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

package/dist/index.js

@@ -18484,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")})`;
@@ -18528,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(", ")}
@@ -18555,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 {
@@ -18572,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();
 }
@@ -18589,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;