@mastra/pg 1.28.1 → 1.30.0-alpha.0

package/dist/index.js

@@ -10455,6 +10455,7 @@ function dedupeMessagesForSave(messages) {
 var MemoryPG = class MemoryPG extends MemoryStorage {
 	supportsPartialThreadUpdate = true;
 	supportsObservationalMemory = true;
+	supportsObservationalMemoryHistorySearch = true;
 	/**
 	* Retention-eligible tables. `threads`, `messages`, and `resources` all anchor
 	* on the timezone-aware `createdAtZ` mirror column (kept in sync by triggers),
@@ -12061,6 +12062,10 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			const conditions = [`"lookupKey" = $1`];
 			const params = [lookupKey];
 			let paramIndex = 2;
+			if (options?.recordId !== void 0) {
+				conditions.push(`id = $${paramIndex++}`);
+				params.push(options.recordId);
+			}
 			if (options?.from) {
 				conditions.push(`"createdAtZ" >= $${paramIndex}`);
 				params.push(options.from.toISOString());
@@ -12071,8 +12076,27 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 				params.push(options.to.toISOString());
 				paramIndex++;
 			}
+			if (options?.groupId !== void 0) {
+				conditions.push(`(strpos("activeObservations", $${paramIndex}) > 0 OR EXISTS (
+          SELECT 1 FROM jsonb_array_elements(
+            CASE WHEN jsonb_typeof("bufferedObservationChunks") = 'array' THEN "bufferedObservationChunks" ELSE '[]'::jsonb END
+          ) AS chunk
+          WHERE strpos(chunk->>'observations', $${paramIndex}) > 0
+        ))`);
+				paramIndex++;
+				params.push(`<observation-group id="${options.groupId}"`);
+			}
+			if (options?.beforeGeneration !== void 0) {
+				conditions.push(`"generationCount" < $${paramIndex++}`);
+				params.push(options.beforeGeneration);
+			}
+			if (options?.afterGeneration !== void 0) {
+				conditions.push(`"generationCount" > $${paramIndex++}`);
+				params.push(options.afterGeneration);
+			}
+			const order = options?.sortDirection === "ASC" ? `"generationCount" ASC, "createdAt" ASC, id ASC` : OM_GENERATION_ORDER;
 			params.push(limit);
-			let sql = `SELECT * FROM ${tableName} WHERE ${conditions.join(" AND ")} ORDER BY ${OM_GENERATION_ORDER} LIMIT $${paramIndex}`;
+			let sql = `SELECT * FROM ${tableName} WHERE ${conditions.join(" AND ")} ORDER BY ${order} LIMIT $${paramIndex}`;
 			paramIndex++;
 			if (options?.offset != null) {
 				params.push(options.offset);
@@ -17460,6 +17484,16 @@ function fieldSql(registry, field) {
 	if (sql === void 0) throw new Error(`Unsupported trusted trace-query field: ${field}`);
 	return sql;
 }
+/**
+* Value of trace-root `metadata.<key>` (alias `r`) as filtering, discovery, and aggregate grouping
+* all see it: the normalized search value, falling back to the trimmed raw string.
+*/
+function traceMetadataValueSql(keyParameter) {
+	return `COALESCE(
+      CASE WHEN jsonb_typeof(r."metadataSearch" -> ${keyParameter}) = 'string' THEN r."metadataSearch" ->> ${keyParameter} END,
+      CASE WHEN jsonb_typeof(r."metadataRaw" -> ${keyParameter}) = 'string' THEN NULLIF(btrim(r."metadataRaw" ->> ${keyParameter}), '') END
+    )`;
+}
 function isMetadataField(field) {
 	return field.startsWith("metadata.");
 }
@@ -17489,11 +17523,7 @@ function compileScalarPredicate(predicate, registry, parameterOffset, allowMetad
 	let fieldValues = [];
 	if (isMetadataField(predicate.field)) {
 		if (!allowMetadata) throw new Error(`Unsupported trusted trace-query field: ${predicate.field}`);
-		const keyParameter = `$${parameterOffset++}`;
-		field = `COALESCE(
-      CASE WHEN jsonb_typeof(r."metadataSearch" -> ${keyParameter}) = 'string' THEN r."metadataSearch" ->> ${keyParameter} END,
-      CASE WHEN jsonb_typeof(r."metadataRaw" -> ${keyParameter}) = 'string' THEN NULLIF(btrim(r."metadataRaw" ->> ${keyParameter}), '') END
-    )`;
+		field = traceMetadataValueSql(`$${parameterOffset++}`);
 		fieldValues = [predicate.field.slice(9)];
 	} else field = fieldSql(registry, predicate.field);
 	if (predicate.type === "presence") return {
@@ -17613,6 +17643,54 @@ function collectThreadRelationCollections(predicate, collections) {
 	else collectThreadRelationCollections(predicate.arg, collections);
 	return collections;
 }
+const ROOT_PUSHDOWN_FIELDS = /* @__PURE__ */ new Set([
+	"traceId",
+	"threadId",
+	"resourceId",
+	"runId",
+	"sessionId",
+	"userId",
+	"organizationId",
+	"startedAt",
+	"endedAt",
+	"entityName",
+	"entityType",
+	"environment"
+]);
+function isRootPushdownPredicate(predicate) {
+	if (predicate.type === "relation") return false;
+	if (predicate.type === "boolean") return predicate.args.every(isRootPushdownPredicate);
+	if (predicate.type === "not") return isRootPushdownPredicate(predicate.arg);
+	return ROOT_PUSHDOWN_FIELDS.has(predicate.field);
+}
+/**
+* Splits `where` into top-level AND conjuncts that can filter `root_scope` directly and the rest,
+* which is applied in `candidates`. The latest-root check compares against the unfiltered span
+* table, so filtering a root earlier cannot change which root is current, but it stops long windows
+* from materializing every root (and running the latest-root check on it) before filtering.
+*/
+function splitRootPredicate(predicate) {
+	if (!predicate) return {
+		root: [],
+		rest: []
+	};
+	const conjuncts = predicate.type === "boolean" && predicate.operator === "and" ? predicate.args : [predicate];
+	const root = [];
+	const rest = [];
+	for (const conjunct of conjuncts) (isRootPushdownPredicate(conjunct) ? root : rest).push(conjunct);
+	return {
+		root,
+		rest
+	};
+}
+function compileConjunction(predicates, values) {
+	if (predicates.length === 0) return "TRUE";
+	return predicates.map((predicate) => {
+		const compiled = compilePredicate(predicate, values.length + 1);
+		values.push(...compiled.values);
+		return `(${compiled.sql})`;
+	}).join(" AND ");
+}
 function compilePredicate(predicate, parameterOffset) {
 	if (predicate.type === "relation") {
 		const compiled = predicate.collection === "feedback" ? compileFeedbackScalarPredicate(predicate.predicate, parameterOffset) : compileScalarPredicate(predicate.predicate, predicate.collection === "spans" ? SPAN_FIELDS : SCORE_FIELDS, parameterOffset);
@@ -17677,7 +17755,7 @@ function compileThreadPredicate(predicate, parameterOffset) {
 		values: compiled.values
 	};
 }
-function compilePostgresTraceScope(schema, selection, relationCollections, scope, deltaWindow) {
+function compilePostgresTraceScope(schema, selection, relationCollections, scope, deltaWindow, rootPredicates = []) {
 	const spanTable = qualifiedTable(schema, TABLE_SPAN_EVENTS);
 	const scoreTable = qualifiedTable(schema, TABLE_SCORE_EVENTS);
 	const feedbackTable = qualifiedTable(schema, TABLE_FEEDBACK_EVENTS);
@@ -17697,10 +17775,7 @@ function compilePostgresTraceScope(schema, selection, relationCollections, scope
 		}
 	}
 	const scopeSql = (alias) => scopeConditions.map((condition) => `\n      AND ${alias}.${condition}`).join("");
-	const ctes = [`root_scope AS MATERIALIZED (
-    SELECT *
-    FROM ${spanTable} r
-    WHERE ${[
+	const rootConditions = [
 		`r."parentSpanId" IS NULL`,
 		latestRootPredicate$1(spanTable),
 		`NOT r."isPending"`,
@@ -17709,7 +17784,12 @@ function compilePostgresTraceScope(schema, selection, relationCollections, scope
 		`r."startedAt" < $2`,
 		...deltaConditions,
 		...scopeConditions.map((condition) => `r.${condition}`)
-	].join("\n      AND ")}
+	];
+	if (rootPredicates.length > 0) rootConditions.push(compileConjunction(rootPredicates, values));
+	const ctes = [`root_scope AS MATERIALIZED (
+    SELECT *
+    FROM ${spanTable} r
+    WHERE ${rootConditions.join("\n      AND ")}
   )`];
 	if (relationCollections.has("spans")) ctes.push(`current_spans AS MATERIALIZED (
     SELECT
@@ -17782,8 +17862,27 @@ function compilePostgresTraceScope(schema, selection, relationCollections, scope
 		values
 	};
 }
+/**
+* Builds the CTE chain ending in `candidates`: the completed, current trace roots in the
+* selection's time range and tenant scope that match its `where` predicate. Trace queries and
+* trace aggregates both select from this CTE so they always see the same population.
+*/
+function compilePostgresTraceCandidates(schema, selection, columns, deltaWindow) {
+	const relationCollections = collectRelationCollections(selection.where);
+	const { root, rest } = splitRootPredicate(selection.where);
+	const { ctes, values } = compilePostgresTraceScope(schema, selection, relationCollections, selection.scope, deltaWindow, root);
+	const predicateSql = compileConjunction(rest, values);
+	ctes.push(`candidates AS (
+    SELECT ${columns}
+    FROM root_scope r
+    WHERE ${predicateSql}
+  )`);
+	return {
+		ctes,
+		values
+	};
+}
 function compilePostgresTraceQuery(schema, plan, mode = "data", safeHorizon) {
-	const relationCollections = collectRelationCollections(plan.where);
 	let deltaWindow;
 	if (plan.paginationMode === "delta") {
 		const watermark = coreStorage.getTraceQueryDeltaWatermark(plan, "pg");
@@ -17793,18 +17892,7 @@ function compilePostgresTraceQuery(schema, plan, mode = "data", safeHorizon) {
 			safeHorizon
 		};
 	}
-	const { ctes, values } = compilePostgresTraceScope(schema, plan, relationCollections, plan.scope, deltaWindow);
-	let predicateSql = "TRUE";
-	if (plan.where) {
-		const predicate = compilePredicate(plan.where, values.length + 1);
-		predicateSql = predicate.sql;
-		values.push(...predicate.values);
-	}
-	ctes.push(`candidates AS (
-    SELECT ${TRACE_SELECT}${plan.paginationMode === "delta" ? ", r.\"xactId\", r.\"cursorId\"" : ""}
-    FROM root_scope r
-    WHERE ${predicateSql}
-  )`);
+	const { ctes, values } = compilePostgresTraceCandidates(schema, plan, `${TRACE_SELECT}${plan.paginationMode === "delta" ? ", r.\"xactId\", r.\"cursorId\"" : ""}`, deltaWindow);
 	const candidates = `WITH ${ctes.join(",\n")}`;
 	if (plan.paginationMode === "page" && mode === "count") return {
 		text: `${candidates}
@@ -17957,11 +18045,7 @@ function compilePostgresTraceQueryValues(schema, plan) {
 	let extracted;
 	if (plan.predicateScope === "trace" && plan.path === "tags") extracted = `SELECT value FROM (SELECT DISTINCT r."traceId", UNNEST(r."tags") AS value FROM root_scope r) t`;
 	else if (plan.predicateScope === "trace" && plan.path.startsWith("metadata.")) {
-		const keyParameter = `$${values.length + 1}`;
-		extracted = `SELECT COALESCE(
-      CASE WHEN jsonb_typeof(r."metadataSearch" -> ${keyParameter}) = 'string' THEN r."metadataSearch" ->> ${keyParameter} END,
-      CASE WHEN jsonb_typeof(r."metadataRaw" -> ${keyParameter}) = 'string' THEN NULLIF(btrim(r."metadataRaw" ->> ${keyParameter}), '') END
-    )::text AS value FROM root_scope r`;
+		extracted = `SELECT ${traceMetadataValueSql(`$${values.length + 1}`)}::text AS value FROM root_scope r`;
 		values.push(plan.path.slice(9));
 	} else extracted = `SELECT ${fieldSql(discoveryRegistry(plan.predicateScope), plan.path)}::text AS value FROM ${discoverySource(plan.predicateScope)}`;
 	const searchParameter = values.length + 1;
@@ -18173,6 +18257,187 @@ async function queryThreads(client, schema, plan, timeoutMs) {
 	});
 }
 //#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. */
+const DIMENSION_SQL = {
+	entityType: "r.\"entityType\"",
+	entityName: "r.\"entityName\"",
+	environment: "r.\"environment\"",
+	status: TRACE_STATUS_SQL,
+	serviceName: "r.\"serviceName\"",
+	executionSource: "r.\"executionSource\"",
+	threadId: "r.\"threadId\"",
+	resourceId: "r.\"resourceId\"",
+	userId: "r.\"userId\"",
+	sessionId: "r.\"sessionId\"",
+	organizationId: "r.\"organizationId\"",
+	experimentId: "r.\"experimentId\""
+};
+const PERCENTILES = {
+	"duration.p50": .5,
+	"duration.p90": .9,
+	"duration.p95": .95,
+	"duration.p99": .99
+};
+const COMPARISON_SQL = {
+	eq: "=",
+	ne: "<>",
+	lt: "<",
+	lte: "<=",
+	gt: ">",
+	gte: ">="
+};
+/**
+* Compiles a trusted aggregate plan. The plan's caps are enforced by the planner and trusted here.
+*
+* Every candidate root is first reduced to a `facts` row (dimension values, duration, error flag,
+* countDistinct inputs, bucket) so grouping, measures, and joins only reference plain columns.
+* Without `interval`, groups are filtered, ordered, and limited in one grouped select. With
+* `interval`, a `ranked` CTE ranks groups on whole-window measures, then only the top `limit`
+* groups are expanded into per-bucket rows.
+*/
+function compilePostgresTraceAggregate(schema, plan) {
+	const { ctes, values } = compilePostgresTraceCandidates(schema, plan, "r.*");
+	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 distinctSql = (field) => field === "traceId" ? "r.\"traceId\"" : dimensionSql(field);
+	const dimensionColumns = plan.dimensions.map((_, index) => `"d${index}"`);
+	const distinctColumns = /* @__PURE__ */ new Map();
+	const factColumns = plan.dimensions.map((dimension, index) => `(${dimensionSql(dimension)})::text AS "d${index}"`);
+	for (const measure of plan.measures) {
+		if (measure.type !== "countDistinct" || distinctColumns.has(measure.field)) continue;
+		const column = `"cd${distinctColumns.size}"`;
+		distinctColumns.set(measure.field, column);
+		factColumns.push(`(${distinctSql(measure.field)})::text AS ${column}`);
+	}
+	factColumns.push(`(${durationMsSql("r.\"startedAt\"", "r.\"endedAt\"")})::float8 AS "durationMs"`, `r."error" IS NOT NULL AS "isError"`);
+	if (plan.interval !== void 0) {
+		const intervalMs = coreStorage.TRACE_AGGREGATE_INTERVAL_MS[plan.interval];
+		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"`);
+	}
+	ctes.push(`facts AS (
+    SELECT ${factColumns.join(",\n      ")}
+    FROM candidates r
+  )`);
+	const measureSql = (name, alias = "") => {
+		const column = (name) => `${alias}"${name}"`;
+		if (name === "count") return "COUNT(*)::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")})`;
+		if (name === "duration.min") return `MIN(${column("durationMs")})`;
+		if (name === "duration.max") return `MAX(${column("durationMs")})`;
+		const percentile = PERCENTILES[name];
+		if (percentile !== void 0) return `percentile_cont(${percentile}) WITHIN GROUP (ORDER BY ${column("durationMs")})`;
+		const measure = plan.measures.find((candidate) => candidate.type === "countDistinct" && candidate.name === name);
+		const distinctColumn = measure && distinctColumns.get(measure.field);
+		if (distinctColumn === void 0) throw new Error(`Unsupported trusted trace-aggregate measure: ${name}`);
+		return `(COUNT(DISTINCT ${alias}${distinctColumn}))::float8`;
+	};
+	const compileHaving = (predicate) => {
+		if (predicate.type === "boolean") return predicate.args.map((arg) => `(${compileHaving(arg)})`).join(predicate.operator === "and" ? " AND " : " OR ");
+		if (predicate.type === "not") return `NOT (${compileHaving(predicate.arg)})`;
+		const measure = measureSql(predicate.measure);
+		if (predicate.type === "membership") {
+			if (predicate.values.length === 0) return predicate.operator === "in" ? "FALSE" : "TRUE";
+			const list = predicate.values.map((value) => `${param(value)}::float8`).join(", ");
+			return `${measure} ${predicate.operator === "in" ? "IN" : "NOT IN"} (${list})`;
+		}
+		const operator = COMPARISON_SQL[predicate.operator];
+		if (operator === void 0) throw new Error(`Unsupported trusted trace-aggregate operator: ${predicate.operator}`);
+		return `${measure} ${operator} ${param(predicate.value)}::float8`;
+	};
+	const havingSql = ["COUNT(*) > 0", ...plan.having ? [`(${compileHaving(plan.having)})`] : []].join(" AND ");
+	const direction = plan.orderBy.direction === "asc" ? "ASC" : "DESC";
+	let orderTarget;
+	if (plan.orderBy.target === "measure") orderTarget = measureSql(plan.orderBy.measure);
+	else {
+		const index = plan.dimensions.indexOf(plan.orderBy.dimension);
+		if (index === -1) throw new Error(`Unsupported trusted trace-aggregate order dimension: ${plan.orderBy.dimension}`);
+		orderTarget = `${dimensionColumns[index]} COLLATE "C"`;
+	}
+	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);
+	if (plan.interval === void 0) {
+		const projected = [...dimensionColumns, ...plan.measures.map((measure, index) => `${measureSql(measure.name)} AS "m${index}"`)];
+		return {
+			text: `WITH ${ctes.join(",\n")}
+SELECT ${projected.join(", ")}
+FROM facts
+GROUP BY ${groupSql}
+HAVING ${havingSql}
+ORDER BY ${orderSql}
+LIMIT ${limitParameter}`,
+			values
+		};
+	}
+	const maxRankParameter = param(plan.limit);
+	ctes.push(`ranked AS (
+    SELECT ${[...dimensionColumns, `row_number() OVER (ORDER BY ${orderSql}) AS "rank"`].join(", ")}
+    FROM facts
+    GROUP BY ${groupSql}
+    HAVING ${havingSql}
+    ORDER BY ${orderSql}
+    LIMIT ${limitParameter}
+  )`);
+	const joinSql = dimensionColumns.length > 0 ? dimensionColumns.map((column) => `COALESCE(f.${column}, '') = COALESCE(g.${column}, '') AND (f.${column} IS NULL) = (g.${column} IS NULL)`).join(" AND ") : "TRUE";
+	const projected = [
+		...dimensionColumns.map((column) => `f.${column}`),
+		"f.\"bucket\"",
+		...plan.measures.map((measure, index) => `${measureSql(measure.name, "f.")} AS "m${index}"`),
+		`(SELECT COUNT(*) FROM ranked) > ${maxRankParameter} AS "truncated"`
+	];
+	return {
+		text: `WITH ${ctes.join(",\n")}
+SELECT ${projected.join(", ")}
+FROM facts f
+JOIN ranked g ON ${joinSql}
+WHERE g."rank" <= ${maxRankParameter}
+GROUP BY ${[
+			"g.\"rank\"",
+			...dimensionColumns.map((column) => `f.${column}`),
+			"f.\"bucket\""
+		].join(", ")}
+ORDER BY g."rank" ASC, f."bucket" ASC`,
+		values
+	};
+}
+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);
+	const result = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(compiled.text, compiled.values));
+	let truncated;
+	let resultRows;
+	if (plan.interval === void 0) {
+		truncated = result.length > plan.limit;
+		resultRows = result.slice(0, plan.limit);
+	} else {
+		truncated = result[0]?.truncated === true;
+		resultRows = result;
+	}
+	return {
+		rows: resultRows.map((row) => {
+			const shaped = { measures: Object.fromEntries(plan.measures.map((measure, index) => [measure.name, Number(row[`m${index}`])])) };
+			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;
+		}),
+		truncated
+	};
+}
+//#endregion
 //#region src/storage/domains/observability/v-next/traces.ts
 /**
 * Trace / branch list reads for the v-next Postgres observability domain.
@@ -19074,6 +19339,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"tag-discovery",
 			"metric-discovery",
 			"trace-query",
+			"trace-aggregate",
 			"trace-query-root-duration",
 			"trace-query-discovery",
 			"thread-query",
@@ -19092,6 +19358,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"metric-discovery",
 			"delta-polling",
 			"trace-query",
+			"trace-aggregate",
 			"trace-query-root-duration",
 			"trace-query-discovery",
 			"thread-query",
@@ -19143,6 +19410,9 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async queryTraces(plan) {
 		return this.#run("QUERY_TRACES", () => queryTraces(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}
+	async aggregateTraces(plan) {
+		return this.#run("AGGREGATE_TRACES", () => aggregateTraces(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));
 	}