@mastra/pg 1.33.0-alpha.2 → 1.33.0-alpha.3

package/dist/index.js

@@ -3996,6 +3996,42 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 			}, error);
 		}
 	}
+	/**
+	* Runs one UPDATE per record in a transaction, changing a row only while every
+	* column in `nullColumns` is still NULL. Never inserts. Returns, in input order,
+	* whether each update changed a row.
+	*/
+	async batchUpdateWhereNull({ tableName, updates, nullColumns }) {
+		const nullConditions = nullColumns.map((column) => ` AND "${parseSqlIdentifier(column, "column name")}" IS NULL`).join("");
+		try {
+			return await this.client.tx(async (tx) => {
+				const applied = [];
+				for (const update of updates) {
+					const query = await this.buildUpdateQuery({
+						tableName,
+						...update
+					});
+					if (!query) {
+						applied.push(false);
+						continue;
+					}
+					const rows = await tx.any(`${query.sql}${nullConditions} RETURNING 1`, query.values);
+					applied.push(rows.length > 0);
+				}
+				return applied;
+			});
+		} catch (error) {
+			throw new MastraError({
+				id: createStorageErrorId("PG", "BATCH_UPDATE_WHERE_NULL", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY,
+				details: {
+					tableName,
+					numberOfRecords: updates.length
+				}
+			}, error);
+		}
+	}
 	async batchDelete({ tableName, keys }) {
 		try {
 			if (keys.length === 0) return;
@@ -4116,9 +4152,14 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 			columns: [column]
 		});
 	}
-	async executeUpdate(client, { tableName, keys, data }) {
+	async executeUpdate(client, args) {
+		const query = await this.buildUpdateQuery(args);
+		if (!query) return;
+		await client.none(query.sql, query.values);
+	}
+	async buildUpdateQuery({ tableName, keys, data }) {
 		const filteredData = await this.filterRecordToKnownColumns(tableName, data);
-		if (Object.keys(filteredData).length === 0) return;
+		if (Object.keys(filteredData).length === 0) return null;
 		const setColumns = [];
 		const setValues = [];
 		let paramIndex = 1;
@@ -4134,12 +4175,13 @@ MIGRATION REQUIRED: Duplicate spans detected in ${duplicateInfo.tableName}\n====
 			whereConditions.push(`"${parsedKey}" = $${paramIndex++}`);
 			whereValues.push(this.prepareValue(value, key, tableName));
 		});
-		const sql = `UPDATE ${getTableName$6({
-			indexName: tableName,
-			schemaName: getSchemaName$6(this.schemaName)
-		})} SET ${setColumns.join(", ")} WHERE ${whereConditions.join(" AND ")}`;
-		const values = [...setValues, ...whereValues];
-		await client.none(sql, values);
+		return {
+			sql: `UPDATE ${getTableName$6({
+				indexName: tableName,
+				schemaName: getSchemaName$6(this.schemaName)
+			})} SET ${setColumns.join(", ")} WHERE ${whereConditions.join(" AND ")}`,
+			values: [...setValues, ...whereValues]
+		};
 	}
 };
 // src/storage/domains/utils.ts
@@ -12015,7 +12057,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 				indexName: OM_TABLE,
 				schemaName: quotedSchemaName
 			});
-			const idxPrefix = schemaPrefix ? `${schemaPrefix}` : "";
+			const idxPrefix = schemaPrefix ? schemaPrefix : "";
 			statements.push(`CREATE INDEX IF NOT EXISTS "${idxPrefix}idx_om_lookup_key" ON ${fullOmTableName} ("lookupKey");`);
 		}
 		for (const idx of MemoryPG.getDefaultIndexDefs(schemaPrefix)) statements.push(generateIndexSQL(idx, schemaName));
@@ -15689,30 +15731,33 @@ var ObservabilityPG = class ObservabilityPG extends ObservabilityStorage {
 			}, error);
 		}
 	}
+	#toSpanUpdates(records) {
+		return records.map((record) => {
+			const data = { ...record.updates };
+			if (data.endedAt instanceof Date) {
+				const endedAt = data.endedAt.toISOString();
+				data.endedAt = endedAt;
+				data.endedAtZ = endedAt;
+			}
+			if (data.startedAt instanceof Date) {
+				const startedAt = data.startedAt.toISOString();
+				data.startedAt = startedAt;
+				data.startedAtZ = startedAt;
+			}
+			return {
+				keys: {
+					spanId: record.spanId,
+					traceId: record.traceId
+				},
+				data
+			};
+		});
+	}
 	async batchUpdateSpans(args) {
 		try {
 			return this.#db.batchUpdate({
 				tableName: TABLE_SPANS,
-				updates: args.records.map((record) => {
-					const data = { ...record.updates };
-					if (data.endedAt instanceof Date) {
-						const endedAt = data.endedAt.toISOString();
-						data.endedAt = endedAt;
-						data.endedAtZ = endedAt;
-					}
-					if (data.startedAt instanceof Date) {
-						const startedAt = data.startedAt.toISOString();
-						data.startedAt = startedAt;
-						data.startedAtZ = startedAt;
-					}
-					return {
-						keys: {
-							spanId: record.spanId,
-							traceId: record.traceId
-						},
-						data
-					};
-				})
+				updates: this.#toSpanUpdates(args.records)
 			});
 		} catch (error) {
 			throw new MastraError({
@@ -15722,6 +15767,25 @@ var ObservabilityPG = class ObservabilityPG extends ObservabilityStorage {
 			}, error);
 		}
 	}
+	async batchUpdateUnendedSpans(args) {
+		try {
+			const applied = await this.#db.batchUpdateWhereNull({
+				tableName: TABLE_SPANS,
+				updates: this.#toSpanUpdates(args.records),
+				nullColumns: ["endedAt"]
+			});
+			return { updated: args.records.filter((_, index) => applied[index]).map(({ traceId, spanId }) => ({
+				traceId,
+				spanId
+			})) };
+		} catch (error) {
+			throw new MastraError({
+				id: createStorageErrorId("PG", "BATCH_UPDATE_UNENDED_SPANS", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.USER
+			}, error);
+		}
+	}
 	async batchDeleteTraces(args) {
 		this.assertUnscopedBatchDeleteTraces(args);
 		try {
@@ -20560,7 +20624,7 @@ async function queryThreadRows(client, schema, plan, timeoutMs, strategy = "auto
 			return `"${field}" ${field === "organizationId" || field === "resourceId" ? "IS NOT DISTINCT FROM" : "="} $${values.length}`;
 		}).join(" AND ")})`;
 	}));
-	const payloads = await run(`SELECT "organizationId", "traceId", "cursorId", "spanType", ${contextFields.map((field) => `"${field}"`).join(", ")},
+	const payloads = await run(`SELECT "organizationId", "traceId", "cursorId", ${contextFields.map((field) => `"${field}"`).join(", ")},
     CASE WHEN "spanType" = 'agent_run' THEN "timeToFirstChunkMs" END AS "timeToFirstChunkMs",
     CASE WHEN error IS NULL THEN 'success' ELSE 'error' END AS status,
     ${[
@@ -20575,30 +20639,10 @@ async function queryThreadRows(client, schema, plan, timeoutMs, strategy = "auto
 		String(version)
 	]);
 	const byRoot = new Map(payloads.map((row) => [key(row.organizationId, row.traceId, row.cursorId), row]));
-	const agents = page.filter((row) => byRoot.get(key(row.organizationId, row.firstTraceId, row.firstCursorId))?.spanType === "agent_run");
-	const models = /* @__PURE__ */ new Map();
-	if (agents.length) {
-		const modelValues = [];
-		const rows = await run(`SELECT DISTINCT ON ("organizationId", "traceId") "organizationId", "traceId",
-        CASE WHEN jsonb_typeof(attributes -> 'model') = 'string' THEN attributes ->> 'model' END AS model,
-        CASE WHEN jsonb_typeof(attributes -> 'provider') = 'string' THEN attributes ->> 'provider' END AS provider
-        FROM ${table} WHERE "spanType" = 'model_generation' AND NOT "isPending" AND (${agents.map((row) => `(${Object.entries({
-			organizationId: row.organizationId,
-			resourceId: row.resourceId,
-			traceId: row.firstTraceId,
-			parentSpanId: row.firstSpanId
-		}).map(([field, value]) => {
-			modelValues.push(value);
-			return `"${field}" ${field === "organizationId" || field === "resourceId" ? "IS NOT DISTINCT FROM" : "="} $${modelValues.length}`;
-		}).join(" AND ")})`).join(" OR ")})
-        ORDER BY "organizationId", "traceId", "startedAt", "spanId" COLLATE "C", "cursorId" DESC`, modelValues);
-		for (const row of rows) models.set(JSON.stringify([row.organizationId, row.traceId]), row);
-	}
 	const threads = page.flatMap((row) => {
 		const first = byRoot.get(key(row.organizationId, row.firstTraceId, row.firstCursorId));
 		const last = byRoot.get(key(row.organizationId, row.lastTraceId, row.lastCursorId));
 		if (!first || !last) return [];
-		const model = models.get(JSON.stringify([row.organizationId, row.firstTraceId]));
 		return [coreStorage.threadQueryRowSchema.parse({
 			organizationId: row.organizationId,
 			threadId: row.threadId,
@@ -20610,8 +20654,6 @@ async function queryThreadRows(client, schema, plan, timeoutMs, strategy = "auto
 			endedAt: row.endedAt.toISOString(),
 			traceCount: Number(row.traceCount),
 			timeToFirstChunkMs: first.timeToFirstChunkMs,
-			model: model?.model ?? null,
-			provider: model?.provider ?? null,
 			inputPreview: readPreview(first.inputPreview, first.input),
 			outputPreview: readPreview(last.outputPreview, last.output),
 			errorPreview: readPreview(last.errorPreview, last.error)
@@ -25595,6 +25637,13 @@ var WorkflowsPG = class WorkflowsPG extends WorkflowsStorage {
 	supportsConcurrentUpdates() {
 		return true;
 	}
+	supportsGuardedSnapshotWrites() {
+		return true;
+	}
+	/** SQL for a snapshot's status. Legacy json/text columns can hold escapes that fail a jsonb cast (#11563). */
+	async #statusExpression(column) {
+		return await this.#db.getColumnType(TABLE_WORKFLOW_SNAPSHOT, "snapshot") === "jsonb" ? `${column} ->> 'status'` : `regexp_replace(${column}::text, '\\\\u(0000|[Dd][89A-Fa-f][0-9A-Fa-f]{2})', '', 'g')::jsonb ->> 'status'`;
+	}
 	parseWorkflowRun(row) {
 		let parsedSnapshot = row.snapshot;
 		if (typeof parsedSnapshot === "string") try {
@@ -25990,20 +26039,22 @@ var WorkflowsPG = class WorkflowsPG extends WorkflowsStorage {
 			}, error);
 		}
 	}
-	async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, createdAt, updatedAt, fence }) {
+	async persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot, createdAt, updatedAt, fence, expectedStatus }) {
 		try {
 			const now = /* @__PURE__ */ new Date();
 			const createdAtValue = createdAt ? createdAt : now;
 			const updatedAtValue = updatedAt ? updatedAt : now;
 			const sanitizedSnapshot = toPgJson(snapshot);
-			await withRunFence(this.#db.client, this.#runOwnersTable(), resolveRunFence(this, fence, runId), "persistWorkflowSnapshot", (q) => q.none(`INSERT INTO ${getTableName({
+			const guard = expectedStatus === void 0 ? "" : ` WHERE ${await this.#statusExpression("t.snapshot")} = ANY($9::text[])`;
+			return await withRunFence(this.#db.client, this.#runOwnersTable(), resolveRunFence(this, fence, runId), "persistWorkflowSnapshot", (q) => q.oneOrNone(`INSERT INTO ${getTableName({
 				indexName: TABLE_WORKFLOW_SNAPSHOT,
 				schemaName: getSchemaName(this.#schema)
 			})} AS t
                  (workflow_name, run_id, "resourceId", snapshot, "createdAt", "updatedAt", "createdAtZ", "updatedAtZ")
                  VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                  ON CONFLICT (workflow_name, run_id) DO UPDATE
-                 SET "resourceId" = COALESCE($3, t."resourceId"), snapshot = $4, "updatedAt" = $6, "updatedAtZ" = $8`, [
+                 SET "resourceId" = COALESCE($3, t."resourceId"), snapshot = $4, "updatedAt" = $6, "updatedAtZ" = $8${guard}
+                 RETURNING 1 AS written`, [
 				workflowName,
 				runId,
 				resourceId,
@@ -26011,8 +26062,9 @@ var WorkflowsPG = class WorkflowsPG extends WorkflowsStorage {
 				createdAtValue,
 				updatedAtValue,
 				createdAtValue,
-				updatedAtValue
-			]));
+				updatedAtValue,
+				...expectedStatus === void 0 ? [] : [Array.isArray(expectedStatus) ? expectedStatus : [expectedStatus]]
+			])) !== null;
 		} catch (error) {
 			if (isRunFenceConflictError(error)) throw error;
 			throw new MastraError({
@@ -26022,6 +26074,25 @@ var WorkflowsPG = class WorkflowsPG extends WorkflowsStorage {
 			}, error);
 		}
 	}
+	async loadWorkflowRunStatus({ workflowName, runId }) {
+		try {
+			const statusExpr = await this.#statusExpression("snapshot");
+			return (await this.#db.client.oneOrNone(`SELECT ${statusExpr} AS status FROM ${getTableName({
+				indexName: TABLE_WORKFLOW_SNAPSHOT,
+				schemaName: getSchemaName(this.#schema)
+			})} WHERE "workflow_name" = $1 AND "run_id" = $2 ORDER BY "createdAt" DESC LIMIT 1`, [workflowName, runId]))?.status ?? null;
+		} catch (error) {
+			throw new MastraError({
+				id: createStorageErrorId("PG", "LOAD_WORKFLOW_RUN_STATUS", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY,
+				details: {
+					workflowName,
+					runId
+				}
+			}, error);
+		}
+	}
 	async loadWorkflowSnapshot({ workflowName, runId }) {
 		try {
 			const result = await this.#db.load({