@mastra/pg 1.32.0 → 1.33.0-alpha.1

package/dist/docs/references/reference-tools-vector-query-tool.md

@@ -148,7 +148,7 @@ const queryTool = createVectorQueryTool({
   indexName: 'documentation',
   model: new ModelRouterEmbeddingModel('openai/text-embedding-3-small'),
   reranker: {
-    model: 'openai/gpt-5.6-sol',
+    model: 'openai/gpt-6.1-sol',
     options: {
       weights: {
         semantic: 0.5, // Semantic relevance weight

package/dist/docs/references/reference-vectors-pg.md

@@ -412,7 +412,7 @@ export const pgAgent = new Agent({
   name: 'PG Agent',
   instructions:
     'You are an AI agent with the ability to automatically recall memories from previous interactions.',
-  model: 'openai/gpt-5.6-sol',
+  model: 'openai/gpt-6.1-sol',
   memory: new Memory({
     storage: new PostgresStore({
       id: 'pg-agent-storage',

package/dist/index.cjs

@@ -20280,7 +20280,12 @@ function afterCursor(plan, values) {
 	}
 	return `(${terms.join(" OR ")})`;
 }
-function compilePostgresSpanQuery(schema, plan) {
+/**
+* The spans `querySpans()` pages through, as a `selected` CTE: the latest completed copy of each
+* span that starts in the window and matches `where`. `aggregateSpans()` groups the same CTE.
+* `values` holds the CTEs' parameters; callers append their own.
+*/
+function compilePostgresSpanSelection(schema, plan) {
 	const values = [plan.timeRange.from, plan.timeRange.to];
 	const scope = [];
 	if (plan.scope) {
@@ -20296,17 +20301,17 @@ function compilePostgresSpanQuery(schema, plan) {
 		values: []
 	};
 	values.push(...predicate.values);
-	const after = afterCursor(plan, values);
-	values.push(plan.limit + 1);
 	const table = qualifiedTable(schema, TABLE_SPAN_EVENTS);
 	const legacyModelFields = _mastra_core_storage.SPAN_QUERY_FIELD_REGISTRY ? "" : `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,`;
 	return {
-		text: `WITH candidate_ids AS MATERIALIZED (
+		ctes: [
+			`candidate_ids AS MATERIALIZED (
     SELECT DISTINCT ${identities} FROM ${table}
     WHERE "startedAt" >= $1 AND "startedAt" < $2 AND "endedAt" >= $1 AND NOT "isPending"
       ${scope.length ? `AND ${scope.join(" AND ")}` : ""}
-  ), current_spans AS MATERIALIZED (
+  )`,
+			`current_spans AS MATERIALIZED (
     SELECT DISTINCT ON (${identities})
       ${fields.map((field) => `"${field}"`).join(", ")}, "startedAt", "endedAt", "cursorId"::text AS version,
       ${legacyModelFields}
@@ -20320,9 +20325,23 @@ function compilePostgresSpanQuery(schema, plan) {
       SELECT "organizationId" IS NULL, COALESCE("organizationId", ''), "resourceId" IS NULL, COALESCE("resourceId", ''), "traceId", "spanId" FROM candidate_ids
     ) ${scope.length ? `AND ${scope.join(" AND ")}` : ""}
     ORDER BY ${identities}, "endedAt" DESC, "cursorId" DESC
-  ), candidates AS (
+  )`,
+			`candidates AS (
     SELECT *, CASE WHEN "hasError" THEN 'error' ELSE NULL END AS error FROM current_spans
-  )
+  )`,
+			`selected AS (
+    SELECT * FROM candidates s WHERE s."startedAt" >= $1 AND s."startedAt" < $2 AND (${predicate.sql})
+  )`
+		],
+		values
+	};
+}
+function compilePostgresSpanQuery(schema, plan) {
+	const { ctes, values } = compilePostgresSpanSelection(schema, plan);
+	const after = afterCursor(plan, values);
+	values.push(plan.limit + 1);
+	return {
+		text: `WITH ${ctes.join(", ")}
   SELECT ${[
 			...fields,
 			"startedAt",
@@ -20331,7 +20350,7 @@ function compilePostgresSpanQuery(schema, plan) {
 			"status",
 			"version"
 		].map((field) => `s."${field}"`).join(", ")}
-  FROM candidates s WHERE s."startedAt" >= $1 AND s."startedAt" < $2 AND (${predicate.sql}) AND ${after}
+  FROM selected s WHERE ${after}
   ORDER BY s."${plan.orderBy.field}" ${plan.orderBy.direction === "asc" ? "ASC" : "DESC"} NULLS LAST, ${identityOrder}
   LIMIT $${values.length}`,
 		values
@@ -20665,6 +20684,28 @@ const DIMENSION_SQL = {
 	organizationId: "r.\"organizationId\"",
 	experimentId: "r.\"experimentId\""
 };
+/** Span column expressions (alias `r`, a `selected` span row) for each groupable span dimension. */
+const SPAN_DIMENSION_SQL = Object.fromEntries([
+	"name",
+	"spanType",
+	"entityType",
+	"entityId",
+	"entityName",
+	"parentEntityType",
+	"parentEntityName",
+	"rootEntityType",
+	"rootEntityName",
+	"environment",
+	"status",
+	"serviceName",
+	"executionSource",
+	"threadId",
+	"resourceId",
+	"userId",
+	"sessionId",
+	"organizationId",
+	"experimentId"
+].map((field) => [field, `r."${field}"`]));
 const PERCENTILES = {
 	"duration.p50": .5,
 	"duration.p90": .9,
@@ -20690,16 +20731,43 @@ const COMPARISON_SQL = {
 */
 function compilePostgresTraceAggregate(schema, plan) {
 	const { ctes, values } = compilePostgresTraceCandidates(schema, plan, "r.*");
+	return compilePostgresAggregate(schema, plan, {
+		ctes,
+		values,
+		from: "candidates",
+		dimensionSql: (field, param) => {
+			if (field.startsWith("metadata.")) return traceMetadataValueSql(param(field.slice(9)));
+			return lookupDimensionSql(DIMENSION_SQL, field);
+		},
+		usageKey: ["traceId"]
+	});
+}
+/**
+* Compiles a trusted span aggregate plan over the spans `querySpans()` selects. Usage is
+* attributed by `(traceId, spanId)`, so only the span that recorded a metric row carries it.
+*/
+function compilePostgresSpanAggregate(schema, plan) {
+	const { ctes, values } = compilePostgresSpanSelection(schema, plan);
+	return compilePostgresAggregate(schema, plan, {
+		ctes,
+		values,
+		from: "selected",
+		dimensionSql: (field) => lookupDimensionSql(SPAN_DIMENSION_SQL, field),
+		usageKey: ["traceId", "spanId"]
+	});
+}
+function lookupDimensionSql(columns, field) {
+	const sql = Object.hasOwn(columns, field) ? columns[field] : void 0;
+	if (sql === void 0) throw new Error(`Unsupported trusted aggregate field: ${field}`);
+	return sql;
+}
+function compilePostgresAggregate(schema, plan, source) {
+	const { ctes, values } = source;
 	const param = (value) => {
 		values.push(value);
 		return `$${values.length}`;
 	};
-	const dimensionSql = (field) => {
-		if (field.startsWith("metadata.")) return traceMetadataValueSql(param(field.slice(9)));
-		const sql = Object.hasOwn(DIMENSION_SQL, field) ? DIMENSION_SQL[field] : void 0;
-		if (sql === void 0) throw new Error(`Unsupported trusted trace-aggregate field: ${field}`);
-		return sql;
-	};
+	const dimensionSql = (field) => source.dimensionSql(field, param);
 	const distinctSql = (field) => field === "traceId" ? "r.\"traceId\"" : dimensionSql(field);
 	const dimensionColumns = plan.dimensions.map((_, index) => `"d${index}"`);
 	const distinctColumns = /* @__PURE__ */ new Map();
@@ -20720,14 +20788,14 @@ function compilePostgresTraceAggregate(schema, plan) {
 	const hasCost = measureKinds.has("cost");
 	let usageJoin = "";
 	if (hasCost || measureKinds.has("tokens")) {
-		ctes.push(...compileUsageCtes(schema, plan, param));
-		usageJoin = "\n    LEFT JOIN usage u ON u.\"traceId\" = r.\"traceId\"";
+		ctes.push(...compileUsageCtes(schema, plan, source, param));
+		usageJoin = `\n    LEFT JOIN usage u ON ${source.usageKey.map((key) => `u."${key}" = r."${key}"`).join(" AND ")}`;
 		USAGE_METRIC_NAMES.forEach((_, index) => factColumns.push(`CASE WHEN u."traceId" IS NOT NULL THEN COALESCE(u."t${index}", 0) END AS "t${index}"`));
 		factColumns.push(`u."traceId" IS NOT NULL AS "usageBearing"`, `CASE WHEN u."priced" THEN u."cost" END AS "traceCost"`, `COALESCE(u."priced" AND NOT u."pricingFailure", FALSE) AS "covered"`, `u."unitMin"`, `u."unitMax"`);
 	}
 	ctes.push(`facts AS (
     SELECT ${factColumns.join(",\n      ")}
-    FROM candidates r${usageJoin}
+    FROM ${source.from} r${usageJoin}
   )`);
 	const mixedUnitsSql = (alias) => `MIN(${alias}"unitMin" COLLATE "C") <> MAX(${alias}"unitMax" COLLATE "C")`;
 	const measureSql = (name, alias = "") => {
@@ -20843,36 +20911,39 @@ function usageMetricIndex(name) {
 	return index;
 }
 /**
-* Builds `usage`: one row per candidate trace that has at least one token metric row, holding
-* per-name token sums (`t<i>`, indexed by `TRACE_AGGREGATE_USAGE_METRIC_NAMES`) and the trace's
-* pricing state. Metric rows are matched by traceId, usage metric name, tenant scope, and
-* `timestamp >= from` only: a trace that starts in the window can emit metrics after `to`, and
-* the candidates semi-join already bounds the scan.
+* Builds `usage`: one row per population row (trace, or span) that has at least one token metric
+* row, holding per-name token sums (`t<i>`, indexed by `TRACE_AGGREGATE_USAGE_METRIC_NAMES`) and
+* its pricing state. Metric rows are matched by the source's usage key, usage metric name, tenant
+* scope, and `timestamp >= from` only: a trace or span that starts in the window can emit metrics
+* after `to`, and the population semi-join already bounds the scan.
 *
 * Rows are deduplicated on `metricId` alone. The primary key is (`metricId`, `timestamp`)
 * because the table is partitioned by timestamp, so retried copies with different timestamps
 * are both stored.
 */
-function compileUsageCtes(schema, plan, param) {
+function compileUsageCtes(schema, plan, source, param) {
+	const key = source.usageKey.map((column) => `"${column}"`).join(", ");
+	const metricKey = source.usageKey.map((column) => `m."${column}"`).join(", ");
 	const scopeConditions = plan.scope ? [`m."organizationId" = ${param(plan.scope.organizationId)}`, ...plan.scope.resourceId === void 0 ? [] : [`m."resourceId" = ${param(plan.scope.resourceId)}`]] : [];
 	const rowConditions = [
-		`m."traceId" IN (SELECT "traceId" FROM candidates)`,
+		`${source.usageKey.length === 1 ? metricKey : `(${metricKey})`} IN (SELECT ${key} FROM ${source.from})`,
 		`m."name" = ANY(${param(USAGE_METRIC_NAMES)}::text[])`,
 		`m."timestamp" >= ${param(plan.timeRange.from)}::timestamptz`,
 		...scopeConditions
 	];
 	const costRow = `"name" = ANY(${param(_mastra_core_storage.TRACE_AGGREGATE_COST_METRIC_NAMES)}::text[])`;
-	const priced = `${costRow} AND "estimatedCost" IS NOT NULL AND "costUnit" IS NOT NULL AND NOT "hasError"`;
+	const priced = `${costRow} AND "estimatedCost" IS NOT NULL AND "costUnit" IS NOT NULL AND (NOT "hasError" OR "partialCost")`;
 	const failed = `${costRow} AND ("hasError" OR ("estimatedCost" IS NOT NULL AND "costUnit" IS NULL))`;
 	const tokenColumns = USAGE_METRIC_NAMES.map((name, index) => `SUM("value") FILTER (WHERE "name" = ${param(name)}) AS "t${index}"`);
 	return [`usage_rows AS (
-    SELECT DISTINCT ON (m."metricId") m."traceId", m."name", m."value", m."estimatedCost", m."costUnit",
-      COALESCE(jsonb_typeof(m."costMetadata" -> 'error'), 'null') <> 'null' AS "hasError"
+    SELECT DISTINCT ON (m."metricId") ${metricKey}, m."name", m."value", m."estimatedCost", m."costUnit",
+      COALESCE(jsonb_typeof(m."costMetadata" -> 'error'), 'null') <> 'null' AS "hasError",
+      COALESCE(m."costMetadata" ->> 'error' = 'partial_cost', FALSE) AS "partialCost"
     FROM ${qualifiedTable(schema, TABLE_METRIC_EVENTS)} m
     WHERE ${rowConditions.join("\n      AND ")}
     ORDER BY m."metricId", m."timestamp"
   )`, `usage AS (
-    SELECT "traceId",
+    SELECT ${key},
       ${tokenColumns.join(",\n      ")},
       SUM("estimatedCost") FILTER (WHERE ${priced}) AS "cost",
       bool_or(${priced}) AS "priced",
@@ -20880,7 +20951,7 @@ function compileUsageCtes(schema, plan, param) {
       MIN("costUnit" COLLATE "C") FILTER (WHERE ${priced}) AS "unitMin",
       MAX("costUnit" COLLATE "C") FILTER (WHERE ${priced}) AS "unitMax"
     FROM usage_rows
-    GROUP BY "traceId"
+    GROUP BY ${key}
   )`];
 }
 function nullableNumber(value) {
@@ -20890,7 +20961,12 @@ function toIsoTimestamp(value) {
 	return value instanceof Date ? value.toISOString() : new Date(value).toISOString();
 }
 async function aggregateTraces(client, schema, plan, timeoutMs) {
-	const compiled = compilePostgresTraceAggregate(schema, plan);
+	return runAggregate(client, plan, compilePostgresTraceAggregate(schema, plan), timeoutMs);
+}
+async function aggregateSpans(client, schema, plan, timeoutMs) {
+	return runAggregate(client, plan, compilePostgresSpanAggregate(schema, plan), timeoutMs);
+}
+async function runAggregate(client, plan, compiled, timeoutMs) {
 	const result = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(compiled.text, compiled.values));
 	let truncated;
 	let resultRows;
@@ -21686,6 +21762,7 @@ async function dangerouslyClearTracing(client, schema) {
 * Use it through `MastraCompositeStore` with a dedicated Postgres connection.
 */
 const spanQueryFeatures = typeof _mastra_core_storage.planSpanQuery === "function" ? ["span-query", "span-query-context-fields"] : [];
+const spanAggregateFeatures = typeof _mastra_core_storage.planSpanAggregate === "function" ? ["span-aggregate"] : [];
 function wrapError(op, error, details) {
 	if (error instanceof _mastra_core_error.MastraError || error instanceof _mastra_core_storage.TraceQueryCursorError || error instanceof _mastra_core_storage.TraceQueryExecutionError || error instanceof _mastra_core_storage.TraceQueryResourceLimitError) throw error;
 	throw new _mastra_core_error.MastraError({
@@ -21876,6 +21953,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"trace-query",
 			"trace-aggregate",
 			...spanQueryFeatures,
+			...spanAggregateFeatures,
 			...typeof _mastra_core_storage.planThreadRowsQuery === "function" ? ["thread-query-rows"] : [],
 			"trace-query-root-duration",
 			"trace-query-discovery",
@@ -21900,6 +21978,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"trace-query",
 			"trace-aggregate",
 			...spanQueryFeatures,
+			...spanAggregateFeatures,
 			...typeof _mastra_core_storage.planThreadRowsQuery === "function" ? ["thread-query-rows"] : [],
 			"trace-query-root-duration",
 			"trace-query-discovery",
@@ -21964,6 +22043,9 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async aggregateTraces(plan) {
 		return this.#run("AGGREGATE_TRACES", () => aggregateTraces(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}
+	async aggregateSpans(plan) {
+		return this.#run("AGGREGATE_SPANS", () => aggregateSpans(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
+	}
 	async getTraceQueryObservedFields(plan) {
 		return this.#run("GET_TRACE_QUERY_OBSERVED_FIELDS", () => getTraceQueryObservedFields(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}
@@ -25826,7 +25908,10 @@ var WorkflowsPG = class WorkflowsPG extends _mastra_core_storage.WorkflowsStorag
 			throw this.#ownershipError("GET_RUN_OWNERSHIP", runId, error);
 		}
 	}
-	async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state, fence: explicitFence }) {
+	supportsWorkflowStepGuards() {
+		return true;
+	}
+	async updateWorkflowResults({ workflowName, runId, stepId, result, requestContext, state, fence: explicitFence, startGuard, completionGuard }) {
 		const fence = (0, _mastra_core_storage.resolveRunFence)(this, explicitFence, runId);
 		try {
 			return await this.#db.client.tx(async (t) => {
@@ -25836,16 +25921,48 @@ var WorkflowsPG = class WorkflowsPG extends _mastra_core_storage.WorkflowsStorag
 					schemaName: getSchemaName(this.#schema)
 				});
 				const existingSnapshotResult = await t.oneOrNone(`SELECT snapshot FROM ${tableName} WHERE workflow_name = $1 AND run_id = $2 FOR UPDATE`, [workflowName, runId]);
-				if (!existingSnapshotResult) return {};
+				if (!existingSnapshotResult) {
+					(0, _mastra_core_storage.assertWorkflowStepGuard)({
+						snapshot: null,
+						stepId,
+						startGuard,
+						completionGuard
+					});
+					return {};
+				}
 				const existingSnapshot = existingSnapshotResult.snapshot;
 				const snapshot = typeof existingSnapshot === "string" ? JSON.parse(existingSnapshot) : existingSnapshot;
-				(0, _mastra_core_storage.mergeWorkflowStepResult)({
+				(0, _mastra_core_storage.assertWorkflowStepGuard)({
+					snapshot,
+					stepId,
+					startGuard,
+					completionGuard
+				});
+				if (completionGuard?.pruneOnly) snapshot.context[stepId] = {
+					...result,
+					metadata: {
+						...snapshot.context[stepId]?.metadata,
+						...result.metadata,
+						__continuationStarted: true
+					}
+				};
+				else (0, _mastra_core_storage.mergeWorkflowStepResult)({
 					snapshot,
 					stepId,
 					result,
 					requestContext,
 					state
 				});
+				const previous = startGuard?.previousStep;
+				const previousResult = previous && snapshot.context[previous.stepId];
+				if (previous && previous.stepId !== stepId && previousResult?.status === "success" && previousResult.metadata?.startedBy === previous.eventId) snapshot.context[previous.stepId] = {
+					...previous.result,
+					metadata: {
+						...previousResult.metadata,
+						...previous.result.metadata,
+						__continuationStarted: true
+					}
+				};
 				const now = /* @__PURE__ */ new Date();
 				const sanitizedSnapshot = toPgJson(snapshot);
 				await t.none(`INSERT INTO ${tableName}
@@ -25864,7 +25981,7 @@ var WorkflowsPG = class WorkflowsPG extends _mastra_core_storage.WorkflowsStorag
 				return snapshot.context;
 			});
 		} catch (error) {
-			if ((0, _mastra_core_storage.isRunFenceConflictError)(error)) throw error;
+			if ((0, _mastra_core_storage.isRunFenceConflictError)(error) || (0, _mastra_core_storage.isWorkflowStepConflictError)(error)) throw error;
 			throw new _mastra_core_error.MastraError({
 				id: (0, _mastra_core_storage.createStorageErrorId)("PG", "UPDATE_WORKFLOW_RESULTS", "FAILED"),
 				domain: _mastra_core_error.ErrorDomain.STORAGE,