@mastra/pg 1.30.0-alpha.1 → 1.30.0-alpha.3

package/dist/index.js

@@ -18256,6 +18256,166 @@ async function queryThreads(client, schema, plan, timeoutMs) {
 		}) : null }
 	});
 }
+/** Compile a span-row filter with exactly the same rules as trace span predicates. */
+function compileSpanQueryPredicate(predicate, offset) {
+	return compileScalarPredicate(predicate, SPAN_FIELDS, offset);
+}
+//#endregion
+//#region src/storage/domains/observability/v-next/span-query.ts
+const identityFields = [
+	"organizationId",
+	"resourceId",
+	"traceId",
+	"spanId"
+];
+const fields = [
+	...identityFields,
+	"parentSpanId",
+	"name",
+	"spanType",
+	"entityType",
+	"entityId",
+	"entityName",
+	"entityVersionId",
+	"parentEntityVersionId",
+	"rootEntityVersionId",
+	"runId",
+	"sessionId",
+	"userId"
+];
+const identities = identityFields.map((field) => `"${field}"`).join(", ");
+const identityOrder = identityFields.map((field) => `"${field}" COLLATE "C" ASC NULLS LAST`).join(", ");
+function selectedKeys(rows, values) {
+	return rows.map((row) => `(${identityFields.map((field) => {
+		values.push(row[field]);
+		return `"${field}" ${field === "traceId" || field === "spanId" ? "=" : "IS NOT DISTINCT FROM"} $${values.length}::text`;
+	}).join(" AND ")})`).join(" OR ");
+}
+function afterCursor(plan, values) {
+	if (!plan.cursor) return "TRUE";
+	const cursor = plan.cursor;
+	values.push(cursor.sortValue);
+	const timestamp = `$${values.length}::timestamptz`;
+	const field = `s."${plan.orderBy.field}"`;
+	const terms = [`${field} ${plan.orderBy.direction === "asc" ? ">" : "<"} ${timestamp}`];
+	const equal = [`${field} = ${timestamp}`];
+	for (const key of identityFields) {
+		const ref = `s."${key}" COLLATE "C"`;
+		const value = cursor[key];
+		if (value === null) equal.push(`${ref} IS NULL`);
+		else {
+			values.push(value);
+			const parameter = `$${values.length}::text COLLATE "C"`;
+			terms.push(`(${equal.join(" AND ")} AND (${ref} > ${parameter} OR ${ref} IS NULL))`);
+			equal.push(`${ref} = ${parameter}`);
+		}
+	}
+	return `(${terms.join(" OR ")})`;
+}
+function compilePostgresSpanQuery(schema, plan) {
+	const values = [plan.timeRange.from, plan.timeRange.to];
+	const scope = [];
+	if (plan.scope) {
+		values.push(plan.scope.organizationId);
+		scope.push(`"organizationId" = $${values.length}`);
+		if (plan.scope.resourceId !== void 0) {
+			values.push(plan.scope.resourceId);
+			scope.push(`"resourceId" = $${values.length}`);
+		}
+	}
+	const predicate = plan.where ? compileSpanQueryPredicate(plan.where, values.length + 1) : {
+		sql: "TRUE",
+		values: []
+	};
+	values.push(...predicate.values);
+	const after = afterCursor(plan, values);
+	values.push(plan.limit + 1);
+	const table = qualifiedTable(schema, TABLE_SPAN_EVENTS);
+	return {
+		text: `WITH 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 (
+    SELECT DISTINCT ON (${identities})
+      ${fields.map((field) => `"${field}"`).join(", ")}, "startedAt", "endedAt", "cursorId"::text AS version,
+      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,
+      (EXTRACT(EPOCH FROM ("endedAt" - "startedAt")) * 1000)::float8 AS "durationMs",
+      CASE WHEN error IS NULL THEN 'success' ELSE 'error' END AS status,
+      error IS NOT NULL AS "hasError"
+    FROM ${table} e
+    WHERE NOT "isPending" AND "endedAt" >= $1 AND (
+      "organizationId" IS NULL, COALESCE("organizationId", ''), "resourceId" IS NULL, COALESCE("resourceId", ''), "traceId", "spanId"
+    ) IN (
+      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 (
+    SELECT *, CASE WHEN "hasError" THEN 'error' ELSE NULL END AS error FROM current_spans
+  )
+  SELECT ${[
+			...identityFields,
+			"parentSpanId",
+			"name",
+			"spanType",
+			"entityType",
+			"entityId",
+			"entityName",
+			"model",
+			"provider",
+			"startedAt",
+			"endedAt",
+			"durationMs",
+			"status",
+			"version"
+		].map((field) => `s."${field}"`).join(", ")}
+  FROM candidates s WHERE s."startedAt" >= $1 AND s."startedAt" < $2 AND (${predicate.sql}) AND ${after}
+  ORDER BY s."${plan.orderBy.field}" ${plan.orderBy.direction === "asc" ? "ASC" : "DESC"}, ${identityOrder}
+  LIMIT $${values.length}`,
+		values
+	};
+}
+async function querySpans(client, schema, plan, timeoutMs) {
+	const deadline = performance.now() + coreStorage.resolveTraceQueryTimeoutMs(timeoutMs);
+	return runWithPostgresTraceQueryTimeout(client, timeoutMs, async (tx) => {
+		const query = compilePostgresSpanQuery(schema, plan);
+		const rows = await tx.any(query.text, query.values);
+		const selected = rows.map(({ version: _version, startedAt, endedAt, ...row }) => ({
+			...row,
+			startedAt: startedAt.toISOString(),
+			endedAt: endedAt.toISOString()
+		}));
+		const page = rows.slice(0, plan.limit);
+		if (!page.length) return coreStorage.buildSpanQueryResponse(plan, selected, [], []);
+		const remaining = async () => {
+			const ms = Math.floor(deadline - performance.now());
+			if (ms <= 0) throw new coreStorage.TraceQueryExecutionError();
+			await tx.query(`SELECT set_config('statement_timeout', $1, true)`, [`${ms}ms`]);
+		};
+		await remaining();
+		const payloadValues = [];
+		const payloadWhere = page.map((row) => {
+			const key = selectedKeys([row], payloadValues);
+			payloadValues.push(row.version);
+			return `(${key} AND "cursorId" = $${payloadValues.length}::bigint)`;
+		}).join(" OR ");
+		const payloads = await tx.any(`SELECT ${identities},
+      left(input::text, ${coreStorage.SPAN_QUERY_MAX_PREVIEW_CHARACTERS + 1}) AS "inputPreview", left(output::text, ${coreStorage.SPAN_QUERY_MAX_PREVIEW_CHARACTERS + 1}) AS "outputPreview"
+      FROM ${qualifiedTable(schema, TABLE_SPAN_EVENTS)} WHERE ${payloadWhere}`, payloadValues);
+		await remaining();
+		const metricValues = [];
+		const metricWhere = selectedKeys(page, metricValues);
+		const metrics = await tx.any(`SELECT ${identities}, name, "estimatedCost", "costUnit",
+      "costMetadata" ->> 'allocation' AS allocation, COALESCE("costMetadata" ? 'error', false) AS "costError"
+      FROM (SELECT DISTINCT ON (${identities}, "metricId") ${identities}, name, "estimatedCost", "costUnit", "costMetadata" FROM ${qualifiedTable(schema, TABLE_METRIC_EVENTS)}
+        WHERE (${metricWhere}) AND name LIKE 'mastra_model_%'
+        ORDER BY ${identities}, "metricId", "cursorId" DESC) m
+      LIMIT ${coreStorage.SPAN_QUERY_MAX_COST_METRICS + 1}`, metricValues);
+		if (metrics.length > coreStorage.SPAN_QUERY_MAX_COST_METRICS) throw new coreStorage.TraceQueryResourceLimitError();
+		return coreStorage.buildSpanQueryResponse(plan, selected, payloads, metrics);
+	}, { repeatableRead: true });
+}
 //#endregion
 //#region src/storage/domains/observability/v-next/trace-aggregate.ts
 /** Trace-root column expressions (alias `r`) for each groupable dimension; same values `where` sees. */
@@ -18324,13 +18484,32 @@ function compilePostgresTraceAggregate(schema, plan) {
 		if (intervalMs === void 0) throw new Error(`Unsupported trusted trace-aggregate interval: ${plan.interval}`);
 		factColumns.push(`to_timestamp(floor(EXTRACT(EPOCH FROM r."startedAt") * 1000 / ${intervalMs}) * ${intervalMs} / 1000.0) AS "bucket"`);
 	}
+	const measureKinds = new Set(plan.measures.map((measure) => measureRule(measure.name)?.kind));
+	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\"";
+		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
+    FROM candidates r${usageJoin}
   )`);
+	const mixedUnitsSql = (alias) => `MIN(${alias}"unitMin" COLLATE "C") <> MAX(${alias}"unitMax" COLLATE "C")`;
 	const measureSql = (name, alias = "") => {
 		const column = (name) => `${alias}"${name}"`;
 		if (name === "count") return "COUNT(*)::float8";
+		const rule = measureRule(name);
+		if (rule?.kind === "tokens" && rule.metricNames) {
+			const value = rule.metricNames.map((metricName) => column(`t${usageMetricIndex(metricName)}`)).join(" + ");
+			return `(${rule.statistic === "avg" ? "AVG" : "SUM"}(${value}))::float8`;
+		}
+		if (rule?.kind === "cost") {
+			const aggregate = rule.statistic === "avg" ? "AVG" : "SUM";
+			return `(CASE WHEN ${mixedUnitsSql(alias)} THEN NULL ELSE ${aggregate}(${column("traceCost")}) END)::float8`;
+		}
 		if (name === "errorCount") return `(COUNT(*) FILTER (WHERE ${column("isError")}))::float8`;
 		if (name === "errorRate") return `(COUNT(*) FILTER (WHERE ${column("isError")}))::float8 / COUNT(*)`;
 		if (name === "duration.avg") return `AVG(${column("durationMs")})`;
@@ -18368,8 +18547,17 @@ function compilePostgresTraceAggregate(schema, plan) {
 	const orderSql = [`${orderTarget} ${direction} NULLS LAST`, ...dimensionColumns.map((column) => `${column} COLLATE "C" ASC NULLS LAST`)].join(", ");
 	const groupSql = dimensionColumns.length > 0 ? dimensionColumns.join(", ") : "()";
 	const limitParameter = param(plan.limit + 1);
+	const costColumns = (alias) => {
+		if (!hasCost) return [];
+		const mixedUnit = param(coreStorage.TRACE_AGGREGATE_MIXED_COST_UNIT);
+		return [`(COUNT(*) FILTER (WHERE ${alias}"covered"))::float8 / NULLIF(COUNT(*) FILTER (WHERE ${alias}"usageBearing"), 0) AS "costCoverage"`, `CASE WHEN ${mixedUnitsSql(alias)} THEN ${mixedUnit}::text ELSE MIN(${alias}"unitMin" COLLATE "C") END AS "costUnit"`];
+	};
 	if (plan.interval === void 0) {
-		const projected = [...dimensionColumns, ...plan.measures.map((measure, index) => `${measureSql(measure.name)} AS "m${index}"`)];
+		const projected = [
+			...dimensionColumns,
+			...plan.measures.map((measure, index) => `${measureSql(measure.name)} AS "m${index}"`),
+			...costColumns("")
+		];
 		return {
 			text: `WITH ${ctes.join(",\n")}
 SELECT ${projected.join(", ")}
@@ -18395,6 +18583,7 @@ LIMIT ${limitParameter}`,
 		...dimensionColumns.map((column) => `f.${column}`),
 		"f.\"bucket\"",
 		...plan.measures.map((measure, index) => `${measureSql(measure.name, "f.")} AS "m${index}"`),
+		...costColumns("f."),
 		`(SELECT COUNT(*) FROM ranked) > ${maxRankParameter} AS "truncated"`
 	];
 	return {
@@ -18412,6 +18601,59 @@ ORDER BY g."rank" ASC, f."bucket" ASC`,
 		values
 	};
 }
+const USAGE_METRIC_NAMES = coreStorage.TRACE_AGGREGATE_USAGE_METRIC_NAMES;
+function measureRule(name) {
+	return coreStorage.isTraceAggregateCanonicalMeasure(name) ? coreStorage.TRACE_AGGREGATE_MEASURE_REGISTRY[name] : void 0;
+}
+function usageMetricIndex(name) {
+	const index = USAGE_METRIC_NAMES.indexOf(name);
+	if (index === -1) throw new Error(`Unsupported trusted trace-aggregate token metric: ${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.
+*
+* 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) {
+	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)`,
+		`m."name" = ANY(${param(USAGE_METRIC_NAMES)}::text[])`,
+		`m."timestamp" >= ${param(plan.timeRange.from)}::timestamptz`,
+		...scopeConditions
+	];
+	const costRow = `"name" = ANY(${param(coreStorage.TRACE_AGGREGATE_COST_METRIC_NAMES)}::text[])`;
+	const priced = `${costRow} AND "estimatedCost" IS NOT NULL AND "costUnit" IS NOT NULL AND NOT "hasError"`;
+	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"
+    FROM ${qualifiedTable(schema, TABLE_METRIC_EVENTS)} m
+    WHERE ${rowConditions.join("\n      AND ")}
+    ORDER BY m."metricId", m."timestamp"
+  )`, `usage AS (
+    SELECT "traceId",
+      ${tokenColumns.join(",\n      ")},
+      SUM("estimatedCost") FILTER (WHERE ${priced}) AS "cost",
+      bool_or(${priced}) AS "priced",
+      bool_or(${failed}) AS "pricingFailure",
+      MIN("costUnit" COLLATE "C") FILTER (WHERE ${priced}) AS "unitMin",
+      MAX("costUnit" COLLATE "C") FILTER (WHERE ${priced}) AS "unitMax"
+    FROM usage_rows
+    GROUP BY "traceId"
+  )`];
+}
+function nullableNumber(value) {
+	return value === null || value === void 0 ? null : Number(value);
+}
 function toIsoTimestamp(value) {
 	return value instanceof Date ? value.toISOString() : new Date(value).toISOString();
 }
@@ -18429,7 +18671,11 @@ async function aggregateTraces(client, schema, plan, timeoutMs) {
 	}
 	return {
 		rows: resultRows.map((row) => {
-			const shaped = { measures: Object.fromEntries(plan.measures.map((measure, index) => [measure.name, Number(row[`m${index}`])])) };
+			const shaped = { measures: Object.fromEntries(plan.measures.map((measure, index) => [measure.name, nullableNumber(row[`m${index}`])])) };
+			if ("costCoverage" in row) shaped.cost = {
+				coverage: nullableNumber(row.costCoverage),
+				unit: row.costUnit ?? null
+			};
 			if (plan.dimensions.length > 0) shaped.dimensions = Object.fromEntries(plan.dimensions.map((dimension, index) => [dimension, row[`d${index}`] ?? null]));
 			if (plan.interval !== void 0) shaped.bucket = toIsoTimestamp(row.bucket);
 			return shaped;
@@ -19151,6 +19397,7 @@ async function dangerouslyClearTracing(client, schema) {
 * application database — observability writes will degrade app performance.
 * Use it through `MastraCompositeStore` with a dedicated Postgres connection.
 */
+const spanQueryFeatures = typeof coreStorage.planSpanQuery === "function" ? ["span-query"] : [];
 function wrapError(op, error, details) {
 	if (error instanceof MastraError || error instanceof coreStorage.TraceQueryCursorError || error instanceof coreStorage.TraceQueryExecutionError || error instanceof coreStorage.TraceQueryResourceLimitError) throw error;
 	throw new MastraError({
@@ -19340,6 +19587,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"metric-discovery",
 			"trace-query",
 			"trace-aggregate",
+			...spanQueryFeatures,
 			"trace-query-root-duration",
 			"trace-query-discovery",
 			"thread-query",
@@ -19359,6 +19607,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"delta-polling",
 			"trace-query",
 			"trace-aggregate",
+			...spanQueryFeatures,
 			"trace-query-root-duration",
 			"trace-query-discovery",
 			"thread-query",
@@ -19407,6 +19656,9 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async listTraces(args) {
 		return this.#run("LIST_TRACES", () => listTraces(this.#readClient, this.#schema, args));
 	}
+	async querySpans(plan) {
+		return this.#run("QUERY_SPANS", () => querySpans(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
+	}
 	async queryTraces(plan) {
 		return this.#run("QUERY_TRACES", () => queryTraces(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}