@mastra/pg 1.22.3 → 1.23.0-alpha.0

package/dist/index.js

@@ -1,4 +1,5 @@
 import { ErrorCategory, ErrorDomain, MastraError } from "@mastra/core/error";
+import * as coreStorage from "@mastra/core/storage";
 import { AgentsStorage, BRANCH_SPAN_TYPES, BackgroundTasksStorage, BlobStore, ChannelsStorage, DATASETS_SCHEMA, DATASET_ITEMS_SCHEMA, DATASET_VERSIONS_SCHEMA, DatasetsStorage, EXPERIMENTS_SCHEMA, EXPERIMENT_RESULTS_SCHEMA, EntityType, ExperimentsStorage, FactoryStorage, FavoritesStorage, KNOWLEDGE_ACTIVITY_SCHEMA, KNOWLEDGE_CURSORS_SCHEMA, KNOWLEDGE_MENTIONS_SCHEMA, KNOWLEDGE_NODES_SCHEMA, KNOWLEDGE_RECORDS_SCHEMA, KNOWLEDGE_SEMANTIC_OUTBOX_SCHEMA, KnowledgeConflictError, KnowledgeNotFoundError, KnowledgeStorage, MCPClientsStorage, MCPServersStorage, METRIC_DISTINCT_COLUMNS, MastraCompositeStore, MemoryStorage, NotificationsStorage, OBSERVATIONAL_MEMORY_TABLE_SCHEMA, ObservabilityStorage, PromptBlocksStorage, SchedulesStorage, ScorerDefinitionsStorage, ScoresStorage, SkillsStorage, TABLE_AGENTS, TABLE_AGENT_VERSIONS, TABLE_BACKGROUND_TASKS, TABLE_CHANNEL_CONFIG, TABLE_CHANNEL_INSTALLATIONS, TABLE_CONFIGS, TABLE_DATASETS, TABLE_DATASET_ITEMS, TABLE_DATASET_VERSIONS, TABLE_EXPERIMENTS, TABLE_EXPERIMENT_RESULTS, TABLE_FAVORITES, TABLE_KNOWLEDGE_ACTIVITY, TABLE_KNOWLEDGE_CURSORS, TABLE_KNOWLEDGE_MENTIONS, TABLE_KNOWLEDGE_NODES, TABLE_KNOWLEDGE_RECORDS, TABLE_KNOWLEDGE_SEMANTIC_OUTBOX, TABLE_MCP_CLIENTS, TABLE_MCP_CLIENT_VERSIONS, TABLE_MCP_SERVERS, TABLE_MCP_SERVER_VERSIONS, TABLE_MESSAGES, TABLE_NOTIFICATIONS, TABLE_PROMPT_BLOCKS, TABLE_PROMPT_BLOCK_VERSIONS, TABLE_RESOURCES, TABLE_SCHEDULES, TABLE_SCHEDULE_TRIGGERS, TABLE_SCHEMAS, TABLE_SCORERS, TABLE_SCORER_DEFINITIONS, TABLE_SCORER_DEFINITION_VERSIONS, TABLE_SKILLS, TABLE_SKILL_BLOBS, TABLE_SKILL_VERSIONS, TABLE_SPANS, TABLE_THREADS, TABLE_THREAD_STATE, TABLE_TOOL_PROVIDER_CONNECTIONS, TABLE_WORKFLOW_DEFINITIONS, TABLE_WORKFLOW_SNAPSHOT, TABLE_WORKSPACES, TABLE_WORKSPACE_VERSIONS, ThreadStateStorage, ToolProviderConnectionsStorage, TraceStatus, UniqueViolationError, WorkflowDefinitionsStorage, WorkflowsStorage, WorkspacesStorage, assertKnowledgeCeilingRaised, assertKnowledgeScopeWithinCeiling, calculatePagination, canonicalizeKnowledgeScope, createKnowledgeUlid, createStorageErrorId, createVectorErrorId, ensureDate, getDefaultValue, getSqlType, hasErrorCode, isKnowledgeScopeVisible, knowledgeScopeKey, knowledgeSemanticDocumentId, knowledgeSemanticIdempotencyKey, listBranchesArgsSchema, listFeedbackArgsSchema, listLogsArgsSchema, listMetricsArgsSchema, listScoresArgsSchema, listTracesArgsSchema, matchesExpectedWorkflowStatus, mergeWorkflowStepResult, normalizePerPage, normalizeScheduleTarget, parseDuration, parseKnowledgeNodeCursor, parseKnowledgeWikilinks, safelyParseJSON, storageMessageMatchesMetadataFilter, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
 import { parseFieldKey, parseSqlIdentifier } from "@mastra/core/utils";
 import { MastraVector, validateTopK, validateUpsertInput } from "@mastra/core/vector";
@@ -12145,12 +12146,20 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			};
 			const lastBufferedAtTime = input.lastBufferedAtTime ? input.lastBufferedAtTime.toISOString() : null;
 			if ((await this.#db.client.query(`UPDATE ${tableName} SET
-          "bufferedObservationChunks" = COALESCE("bufferedObservationChunks", '[]'::jsonb) || $1::jsonb,
-          "lastBufferedAtTime" = COALESCE($2, "lastBufferedAtTime"),
-          "updatedAt" = $3,
-          "updatedAtZ" = $4
-        WHERE id = $5`, [
+          "bufferedObservationChunks" = CASE
+            WHEN EXISTS (
+              SELECT 1
+              FROM jsonb_array_elements(COALESCE("bufferedObservationChunks", '[]'::jsonb)) AS chunk
+              WHERE chunk->>'cycleId' = $2
+            ) THEN COALESCE("bufferedObservationChunks", '[]'::jsonb)
+            ELSE COALESCE("bufferedObservationChunks", '[]'::jsonb) || $1::jsonb
+          END,
+          "lastBufferedAtTime" = COALESCE($3, "lastBufferedAtTime"),
+          "updatedAt" = $4,
+          "updatedAtZ" = $5
+        WHERE id = $6`, [
 				JSON.stringify([newChunk]),
+				input.chunk.cycleId,
 				lastBufferedAtTime,
 				nowStr,
 				nowStr,
@@ -16715,6 +16724,304 @@ async function getScorePercentiles(client, schema, args) {
 	return { series: percentileSeriesFromRows(await client.manyOrNone(sql, acc.params), args.percentiles) };
 }
 //#endregion
+//#region src/storage/domains/observability/v-next/trace-query.ts
+const TRACE_STATUS_SQL = `CASE WHEN r."error" IS NOT NULL THEN 'error' ELSE 'success' END`;
+const TRACE_FIELDS = {
+	traceId: "r.\"traceId\"",
+	threadId: "r.\"threadId\"",
+	resourceId: "r.\"resourceId\"",
+	startedAt: "r.\"startedAt\"",
+	endedAt: "r.\"endedAt\"",
+	entityName: "r.\"entityName\"",
+	entityType: "r.\"entityType\"",
+	environment: "r.\"environment\"",
+	status: TRACE_STATUS_SQL
+};
+const SPAN_FIELDS = {
+	spanType: "s.\"spanType\"",
+	error: "s.\"error\""
+};
+const SCORE_FIELDS = {
+	scorerId: "s.\"scorerId\"",
+	score: "s.\"score\""
+};
+const TRACE_SELECT = `
+  r."traceId" AS "traceId",
+  r."spanId" AS "rootSpanId",
+  r."threadId" AS "threadId",
+  r."resourceId" AS "resourceId",
+  r."startedAt" AS "startedAt",
+  r."endedAt" AS "endedAt",
+  r."entityName" AS "entityName",
+  r."entityType" AS "entityType",
+  r."environment" AS "environment",
+  ${TRACE_STATUS_SQL} AS "status"`;
+function fieldSql(registry, field) {
+	const sql = registry[field];
+	if (sql === void 0) throw new Error(`Unsupported trusted trace-query field: ${field}`);
+	return sql;
+}
+function placeholders(values, offset) {
+	return values.map((_, index) => `$${offset + index}`).join(", ");
+}
+function compileScalarPredicate(predicate, registry, parameterOffset) {
+	if (predicate.type === "boolean") {
+		const values = [];
+		return {
+			sql: predicate.args.map((arg) => {
+				const compiled = compileScalarPredicate(arg, registry, parameterOffset + values.length);
+				values.push(...compiled.values);
+				return `(${compiled.sql})`;
+			}).join(predicate.operator === "and" ? " AND " : " OR "),
+			values
+		};
+	}
+	if (predicate.type === "not") {
+		const compiled = compileScalarPredicate(predicate.arg, registry, parameterOffset);
+		return {
+			sql: `NOT (${compiled.sql})`,
+			values: compiled.values
+		};
+	}
+	const field = fieldSql(registry, predicate.field);
+	if (predicate.type === "presence") return {
+		sql: `${field} IS ${predicate.operator === "exists" ? "NOT " : ""}NULL`,
+		values: []
+	};
+	if (predicate.type === "membership") {
+		const list = placeholders(predicate.values, parameterOffset);
+		if (predicate.operator === "in") return {
+			sql: `${field} IS NOT NULL AND ${field} IN (${list})`,
+			values: predicate.values
+		};
+		return {
+			sql: `${field} IS NULL OR ${field} NOT IN (${list})`,
+			values: predicate.values
+		};
+	}
+	const parameter = `$${parameterOffset}`;
+	const operators = {
+		lt: "<",
+		lte: "<=",
+		gt: ">",
+		gte: ">="
+	};
+	if (predicate.operator === "eq") return {
+		sql: `${field} IS NOT DISTINCT FROM ${parameter}`,
+		values: [predicate.value]
+	};
+	if (predicate.operator === "ne") return {
+		sql: `${field} IS DISTINCT FROM ${parameter}`,
+		values: [predicate.value]
+	};
+	const operator = operators[predicate.operator];
+	if (operator === void 0) throw new Error(`Unsupported trusted trace-query operator: ${predicate.operator}`);
+	return {
+		sql: `${field} IS NOT NULL AND ${field} ${operator} ${parameter}`,
+		values: [predicate.value]
+	};
+}
+function latestRootPredicate$1(spanTable) {
+	return `NOT EXISTS (
+    SELECT 1 FROM ${spanTable} newer
+    WHERE newer."traceId" = r."traceId"
+      AND newer."parentSpanId" IS NULL
+      AND newer."cursorId" > r."cursorId"
+  )`;
+}
+function latestSpanPredicate$1(spanTable) {
+	return `NOT EXISTS (
+    SELECT 1 FROM ${spanTable} newer
+    WHERE newer."traceId" = s."traceId"
+      AND newer."spanId" = s."spanId"
+      AND (newer."isPending" < s."isPending" OR (newer."isPending" = s."isPending" AND newer."cursorId" > s."cursorId"))
+  )`;
+}
+function latestScorePredicate(scoreTable) {
+	return `NOT EXISTS (
+    SELECT 1 FROM ${scoreTable} newer
+    WHERE newer."scoreId" = s."scoreId"
+      AND newer."cursorId" > s."cursorId"
+  )`;
+}
+function collectRelationCollections(predicate, collections = /* @__PURE__ */ new Set()) {
+	if (!predicate) return collections;
+	if (predicate.type === "relation") collections.add(predicate.collection);
+	else if (predicate.type === "boolean") for (const arg of predicate.args) collectRelationCollections(arg, collections);
+	else if (predicate.type === "not") collectRelationCollections(predicate.arg, collections);
+	return collections;
+}
+function compilePredicate(predicate, parameterOffset) {
+	if (predicate.type === "relation") {
+		const registry = predicate.collection === "spans" ? SPAN_FIELDS : SCORE_FIELDS;
+		const compiled = compileScalarPredicate(predicate.predicate, registry, parameterOffset);
+		const existence = `EXISTS (
+      SELECT 1 FROM ${predicate.collection === "spans" ? "current_spans" : "current_scores"} s
+      WHERE s."traceId" = r."traceId"
+        AND (${compiled.sql})
+    )`;
+		return {
+			sql: predicate.quantifier === "some" ? existence : `NOT ${existence}`,
+			values: compiled.values
+		};
+	}
+	if (predicate.type === "boolean") {
+		const values = [];
+		return {
+			sql: predicate.args.map((arg) => {
+				const compiled = compilePredicate(arg, parameterOffset + values.length);
+				values.push(...compiled.values);
+				return `(${compiled.sql})`;
+			}).join(predicate.operator === "and" ? " AND " : " OR "),
+			values
+		};
+	}
+	if (predicate.type === "not") {
+		const compiled = compilePredicate(predicate.arg, parameterOffset);
+		return {
+			sql: `NOT (${compiled.sql})`,
+			values: compiled.values
+		};
+	}
+	return compileScalarPredicate(predicate, TRACE_FIELDS, parameterOffset);
+}
+function compilePostgresTraceQuery(schema, plan) {
+	const spanTable = qualifiedTable(schema, TABLE_SPAN_EVENTS);
+	const scoreTable = qualifiedTable(schema, TABLE_SCORE_EVENTS);
+	const values = [plan.timeRange.from, plan.timeRange.to];
+	const relationCollections = collectRelationCollections(plan.where);
+	const ctes = [`root_scope AS MATERIALIZED (
+    SELECT *
+    FROM ${spanTable} r
+    WHERE ${[
+		`r."parentSpanId" IS NULL`,
+		latestRootPredicate$1(spanTable),
+		`NOT r."isPending"`,
+		`r."endedAt" IS NOT NULL`,
+		`r."startedAt" >= $1`,
+		`r."startedAt" < $2`
+	].join("\n      AND ")}
+  )`];
+	if (relationCollections.has("spans")) ctes.push(`current_spans AS MATERIALIZED (
+    SELECT s."traceId", s."spanType", s."error"
+    FROM ${spanTable} s
+    WHERE s."traceId" IS NOT NULL
+      AND s."traceId" IN (SELECT "traceId" FROM root_scope)
+      AND ${latestSpanPredicate$1(spanTable)}
+  )`);
+	if (relationCollections.has("scores")) ctes.push(`current_scores AS MATERIALIZED (
+    SELECT s."traceId", s."scorerId", s."score"
+    FROM ${scoreTable} s
+    WHERE s."traceId" IS NOT NULL
+      AND s."traceId" IN (SELECT "traceId" FROM root_scope)
+      AND ${latestScorePredicate(scoreTable)}
+  )`);
+	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}
+    FROM root_scope r
+    WHERE ${predicateSql}
+  )`);
+	const candidates = `WITH ${ctes.join(",\n")}`;
+	if (plan.result === "groups") {
+		const pageCondition = plan.cursor ? `AND "threadId" > $${values.length + 1}` : "";
+		if (plan.cursor) values.push(plan.cursor.threadId);
+		values.push(plan.limit + 1);
+		return {
+			text: `${candidates}
+SELECT "threadId"
+FROM candidates
+WHERE "threadId" IS NOT NULL ${pageCondition}
+GROUP BY "threadId"
+ORDER BY "threadId" ASC
+LIMIT $${values.length}`,
+			values
+		};
+	}
+	const orderField = plan.orderBy.field === "startedAt" ? "\"startedAt\"" : "\"endedAt\"";
+	const direction = plan.orderBy.direction === "asc" ? "ASC" : "DESC";
+	let pageCondition = "";
+	if (plan.cursor) {
+		const comparison = plan.orderBy.direction === "asc" ? ">" : "<";
+		const sortParameter = `$${values.length + 1}`;
+		pageCondition = `WHERE (${orderField} ${comparison} ${sortParameter} OR (${orderField} = ${sortParameter} AND "traceId" > ${`$${values.length + 2}`}))`;
+		values.push(plan.cursor.sortValue, plan.cursor.traceId);
+	}
+	values.push(plan.limit + 1);
+	return {
+		text: `${candidates}
+SELECT *
+FROM candidates
+${pageCondition}
+ORDER BY ${orderField} ${direction}, "traceId" ASC
+LIMIT $${values.length}`,
+		values
+	};
+}
+function asIsoTimestamp$1(value) {
+	if (value === null || value === void 0) throw new Error("Trace query returned a null timestamp");
+	return value instanceof Date ? value.toISOString() : new Date(value).toISOString();
+}
+function isPostgresStatementTimeout(error) {
+	if (!error || typeof error !== "object") return false;
+	const candidate = error;
+	return candidate.code === "57014" && String(candidate.message ?? "").includes("statement timeout");
+}
+async function runWithPostgresTraceQueryTimeout(client, timeoutMs, execute) {
+	const resolvedTimeoutMs = coreStorage.resolveTraceQueryTimeoutMs(timeoutMs);
+	try {
+		return await client.tx(async (transaction) => {
+			await transaction.query(`SELECT set_config('statement_timeout', $1, true)`, [`${resolvedTimeoutMs}ms`]);
+			return execute(transaction);
+		});
+	} catch (error) {
+		if (isPostgresStatementTimeout(error)) throw new coreStorage.TraceQueryExecutionError();
+		throw error;
+	}
+}
+async function queryTraces(client, schema, plan, timeoutMs) {
+	const query = compilePostgresTraceQuery(schema, plan);
+	const rows = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(query.text, query.values));
+	const visibleRows = rows.slice(0, plan.limit);
+	if (plan.result === "groups") {
+		const groups = visibleRows.map((row) => ({ threadId: String(row.threadId) }));
+		const last = groups.at(-1);
+		return coreStorage.traceQueryResponseSchema.parse({
+			groups,
+			page: { next: rows.length > plan.limit && last ? coreStorage.encodeTraceQueryCursor(plan, {
+				result: "groups",
+				threadId: last.threadId
+			}) : null }
+		});
+	}
+	const traces = visibleRows.map((row) => ({
+		traceId: String(row.traceId),
+		rootSpanId: String(row.rootSpanId),
+		threadId: row.threadId == null ? null : String(row.threadId),
+		resourceId: row.resourceId == null ? null : String(row.resourceId),
+		startedAt: asIsoTimestamp$1(row.startedAt),
+		endedAt: asIsoTimestamp$1(row.endedAt),
+		entityName: row.entityName == null ? null : String(row.entityName),
+		entityType: row.entityType == null ? null : String(row.entityType),
+		environment: row.environment == null ? null : String(row.environment),
+		status: row.status
+	}));
+	const last = traces.at(-1);
+	return coreStorage.traceQueryResponseSchema.parse({
+		traces,
+		page: { next: rows.length > plan.limit && last ? coreStorage.encodeTraceQueryCursor(plan, {
+			result: "traces",
+			sortValue: last[plan.orderBy.field],
+			traceId: last.traceId
+		}) : null }
+	});
+}
+//#endregion
 //#region src/storage/domains/observability/v-next/traces.ts
 /**
 * Trace / branch list reads for the v-next Postgres observability domain.
@@ -17401,7 +17708,7 @@ async function dangerouslyClearTracing(client, schema) {
 * Use it through `MastraCompositeStore` with a dedicated Postgres connection.
 */
 function wrapError(op, error, details) {
-	if (error instanceof MastraError) throw error;
+	if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError) throw error;
 	throw new MastraError({
 		id: createStorageErrorId("PG", op, "FAILED"),
 		domain: ErrorDomain.STORAGE,
@@ -17414,6 +17721,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	#schema;
 	#partitioning;
 	#discoveryConfig;
+	#traceQueryTimeoutMs;
 	#partitionMode;
 	constructor(config) {
 		super();
@@ -17422,6 +17730,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 		this.#schema = schemaName ?? "public";
 		this.#partitioning = config.partitioning ?? {};
 		this.#discoveryConfig = config.discovery ?? {};
+		this.#traceQueryTimeoutMs = coreStorage.resolveTraceQueryTimeoutMs(config.traceQueryTimeoutMs);
 	}
 	/**
 	* Build the discovery config used at each call site, with the framework
@@ -17573,11 +17882,16 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 		};
 	}
 	getFeatures() {
-		if (!deltaPollingFeatureEnabled()) return ["metrics", "logs"];
+		if (!deltaPollingFeatureEnabled()) return [
+			"metrics",
+			"logs",
+			"trace-query"
+		];
 		return [
 			"metrics",
 			"logs",
-			"delta-polling"
+			"delta-polling",
+			"trace-query"
 		];
 	}
 	async #run(op, fn, details) {
@@ -17620,6 +17934,9 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async listTraces(args) {
 		return this.#run("LIST_TRACES", () => listTraces(this.#client, this.#schema, args));
 	}
+	async queryTraces(plan) {
+		return this.#run("QUERY_TRACES", () => queryTraces(this.#client, this.#schema, plan, this.#traceQueryTimeoutMs));
+	}
 	async listBranches(args) {
 		return this.#run("LIST_BRANCHES", () => listBranches(this.#client, this.#schema, args));
 	}
@@ -22819,7 +23136,8 @@ var PostgresStoreVNext = class extends PostgresStore {
 			client: observabilityClient,
 			schemaName: obsConfig.schemaName ?? config.schemaName,
 			partitioning: obsConfig.partitioning,
-			discovery: obsConfig.discovery
+			discovery: obsConfig.discovery,
+			traceQueryTimeoutMs: obsConfig.traceQueryTimeoutMs
 		});
 		this.stores = {
 			...this.stores,