@mastra/pg 1.25.1-alpha.0 → 1.26.0-alpha.2

package/dist/index.js

@@ -1,7 +1,7 @@
 import { createHash, randomUUID } from "crypto";
 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 { 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, executeRetentionPrune, getDefaultValue, getSqlType, hasErrorCode, isKnowledgeScopeVisible, knowledgeScopeKey, knowledgeSemanticDocumentId, knowledgeSemanticIdempotencyKey, listBranchesArgsSchema, listFeedbackArgsSchema, listLogsArgsSchema, listMetricsArgsSchema, listScoresArgsSchema, listTracesArgsSchema, matchesExpectedWorkflowStatus, mergeWorkflowStepResult, normalizePerPage, normalizeScheduleTarget, parseDuration, parseKnowledgeNodeCursor, parseKnowledgeWikilinks, resolveRetentionTargets, retentionCutoffMs, runRetentionBatches, safelyParseJSON, storageMessageMatchesMetadataFilter, toTraceSpans, transformScoreRow, validateStorageMetadataFilter } from "@mastra/core/storage";
 import { parseFieldKey, parseSqlIdentifier } from "@mastra/core/utils";
 import { MastraVector, validateTopK, validateUpsertInput } from "@mastra/core/vector";
 import { Mutex } from "async-mutex";
@@ -4903,137 +4903,31 @@ var AgentsPG = class AgentsPG extends AgentsStorage {
 };
 //#endregion
 //#region src/storage/retention.ts
-const DEFAULT_BATCH_SIZE = 1e3;
-async function sleep(ms, signal) {
-	if (ms <= 0 || signal?.aborted) return;
-	await new Promise((resolve) => {
-		const timer = setTimeout(() => {
-			signal?.removeEventListener("abort", onAbort);
-			resolve();
-		}, ms);
-		const onAbort = () => {
-			clearTimeout(timer);
-			resolve();
-		};
-		signal?.addEventListener("abort", onAbort, { once: true });
-	});
-}
-/**
-* Convert a policy's `maxAge` into a cutoff bound matching the anchor's storage
-* type: a `Date` for `timestamptz` columns (pg compares timezone-aware), or a
-* raw millisecond number for `bigint` epoch-ms columns.
-*/
 function cutoffFor(policy, anchorType, now = Date.now()) {
-	const cutoffMs = now - parseDuration(policy.maxAge);
+	const cutoffMs = retentionCutoffMs(policy, now);
 	return anchorType === "epoch-ms" ? cutoffMs : new Date(cutoffMs);
 }
-/**
-* Run the bounded/cancellable batched-delete loop for a single logical target,
-* delegating the actual delete of up to `limit` rows to `deleteBatch`. Returns
-* `{ deleted, done }`; `done: false` means the loop stopped on a bound or the
-* abort signal and eligible rows may remain.
-*/
-async function runBatchedDelete({ deleteBatch, batchSize, options }) {
-	if (!Number.isSafeInteger(batchSize) || batchSize <= 0) throw new Error(`retention batchSize must be a positive integer; received ${batchSize}`);
-	let deleted = 0;
-	let batches = 0;
-	while (true) {
-		if (options?.signal?.aborted) return {
-			deleted,
-			done: false
-		};
-		if (options?.maxBatches !== void 0 && batches >= options.maxBatches) return {
-			deleted,
-			done: false
-		};
-		let limit = batchSize;
-		if (options?.maxRows !== void 0) {
-			const remaining = options.maxRows - deleted;
-			if (remaining <= 0) return {
-				deleted,
-				done: false
-			};
-			limit = Math.min(limit, remaining);
-		}
-		const affected = await deleteBatch(limit);
-		deleted += affected;
-		batches += 1;
-		if (affected < limit) return {
-			deleted,
-			done: true
-		};
-		if (options?.pauseMs) await sleep(options.pauseMs, options.signal);
-	}
-}
-/**
-* Runs the bounded, batched, cancellable delete loop for a set of tables in the
-* given order (callers pass children before parents for cascade-safe pruning),
-* and returns one {@link PruneResult} per table.
-*
-* The loop:
-* - deletes in chunks of `batchSize` (default 1000), each its own statement;
-* - stops a table's loop when a batch deletes fewer rows than requested (drained),
-*   or when `maxBatches`/`maxRows` is hit, or the `signal` aborts — the latter
-*   three leave `done: false` so the caller can resume;
-* - pauses `pauseMs` between batches when set, to avoid starving live traffic.
-*
-* `prune()` only deletes rows; it never reclaims disk. PostgreSQL reuses freed
-* space (dead tuples) via autovacuum on subsequent writes, so tables stop
-* growing. Returning disk to the OS (e.g. `VACUUM FULL`) is left to the operator.
-*/
-async function runPrune({ db, domain, targets, options }) {
-	const results = [];
-	const now = Date.now();
-	for (const target of targets) {
-		if (options?.signal?.aborted) {
-			results.push({
-				domain,
-				table: target.table,
-				deleted: 0,
-				done: false
-			});
-			continue;
-		}
-		const cutoff = cutoffFor(target.policy, target.anchorType, now);
-		const { deleted, done } = await runBatchedDelete({
-			deleteBatch: (limit) => db.pruneBatch({
-				tableName: target.table,
-				column: target.column,
-				cutoff,
-				limit
-			}),
-			batchSize: target.policy.batchSize ?? DEFAULT_BATCH_SIZE,
-			options
-		});
-		results.push({
-			domain,
-			table: target.table,
-			deleted,
-			done
-		});
-	}
-	return results;
+const runBatchedDelete = runRetentionBatches;
+function runPrune({ db, domain, targets, options }) {
+	return executeRetentionPrune({
+		domain,
+		targets,
+		options,
+		cutoffFor: (target, now) => cutoffFor(target.policy, target.anchorType ?? "timestamp", now),
+		deleteBatch: (target, cutoff, limit) => db.pruneBatch({
+			tableName: target.table,
+			column: target.column,
+			cutoff,
+			limit
+		})
+	});
 }
-/**
-* Resolve a domain's `{ tableKey: policy }` map plus its descriptor into an
-* ordered list of {@link PruneTarget}s. `order` lists table keys children-first
-* so cascade-dependent rows are removed before their parents. Table keys not in
-* `policies` are skipped (unset = keep forever).
-*/
 function resolveTargets({ policies, descriptor, order }) {
-	const targets = [];
-	for (const key of order) {
-		const policy = policies[key];
-		const entry = descriptor[key];
-		if (!policy || !entry) continue;
-		targets.push({
-			table: entry.table,
-			column: entry.column,
-			anchorType: entry.anchorType ?? "timestamp",
-			policy
-		});
-	}
-	return targets;
+	return resolveRetentionTargets({
+		policies,
+		descriptor,
+		order
+	});
 }
 //#endregion
 //#region src/storage/domains/background-tasks/index.ts
@@ -17291,6 +17185,11 @@ const FEEDBACK_FIELDS = {
 const TRACE_SELECT = `
   r."traceId" AS "traceId",
   r."spanId" AS "rootSpanId",
+  r."name" AS "name",
+  r."entityId" AS "entityId",
+  r."parentSpanId" AS "parentSpanId",
+  r."metadata" AS "metadata",
+  r."input" AS "input",
   r."threadId" AS "threadId",
   r."resourceId" AS "resourceId",
   r."startedAt" AS "startedAt",
@@ -17692,6 +17591,75 @@ LIMIT $${values.length}`,
 		values
 	};
 }
+function discoveryRegistry(scope) {
+	if (scope === "trace") return TRACE_FIELDS;
+	if (scope === "spans") return SPAN_FIELDS;
+	if (scope === "scores") return SCORE_FIELDS;
+	return FEEDBACK_FIELDS;
+}
+function discoverySource(scope) {
+	if (scope === "trace") return "root_scope r";
+	if (scope === "spans") return "current_spans s";
+	if (scope === "scores") return "current_scores s";
+	return "current_feedback s";
+}
+function discoveryCollections(scope) {
+	return scope === "trace" ? /* @__PURE__ */ new Set() : /* @__PURE__ */ new Set([scope]);
+}
+function compilePostgresTraceQueryObservedFields(schema, plan) {
+	const { ctes, values } = compilePostgresTraceScope(schema, plan, /* @__PURE__ */ new Set());
+	const searchParameter = values.length + 1;
+	const search = plan.search ? `AND strpos(lower('metadata.' || entry.key), lower($${searchParameter})) > 0` : "";
+	if (plan.search) values.push(plan.search);
+	values.push(plan.limit + 1);
+	return {
+		text: `WITH ${ctes.join(",\n")}
+SELECT 'metadata.' || entry.key AS path, count(*)::bigint AS occurrences
+FROM root_scope r
+CROSS JOIN LATERAL jsonb_each(CASE WHEN jsonb_typeof(r."metadataRaw") = 'object' THEN r."metadataRaw" ELSE '{}'::jsonb END) entry
+WHERE jsonb_typeof(entry.value) = 'string'
+  AND btrim(entry.value #>> '{}') <> ''
+  AND entry.key <> ''
+  AND strpos(entry.key, '.') = 0
+  AND octet_length('metadata.' || entry.key) <= ${coreStorage.TRACE_QUERY_MAX_PATH_BYTES}
+  AND octet_length(entry.value #>> '{}') <= ${coreStorage.TRACE_QUERY_MAX_STRING_BYTES}
+  ${search}
+GROUP BY entry.key
+ORDER BY occurrences DESC, ('metadata.' || entry.key) COLLATE "C" ASC
+LIMIT $${values.length}`,
+		values
+	};
+}
+function compilePostgresTraceQueryValues(schema, plan) {
+	const { ctes, values } = compilePostgresTraceScope(schema, plan, discoveryCollections(plan.predicateScope));
+	let field;
+	if (plan.predicateScope === "trace" && plan.path.startsWith("metadata.")) {
+		const keyParameter = `$${values.length + 1}`;
+		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
+    )`;
+		values.push(plan.path.slice(9));
+	} else field = fieldSql(discoveryRegistry(plan.predicateScope), plan.path);
+	const searchParameter = values.length + 1;
+	const search = plan.search ? `AND strpos(lower(value), lower($${searchParameter})) > 0` : "";
+	if (plan.search) values.push(plan.search);
+	values.push(plan.limit + 1);
+	return {
+		text: `WITH ${ctes.join(",\n")}, extracted AS (
+  SELECT ${field}::text AS value FROM ${discoverySource(plan.predicateScope)}
+)
+SELECT value, count(*)::bigint AS count
+FROM extracted
+WHERE value IS NOT NULL
+  AND octet_length(value) <= ${coreStorage.TRACE_QUERY_MAX_STRING_BYTES}
+  ${search}
+GROUP BY value
+ORDER BY count DESC, value COLLATE "C" 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();
@@ -17701,6 +17669,11 @@ function isPostgresStatementTimeout(error) {
 	const candidate = error;
 	return candidate.code === "57014" && String(candidate.message ?? "").includes("statement timeout");
 }
+function isPostgresResourceLimit(error) {
+	if (!error || typeof error !== "object") return false;
+	const candidate = error;
+	return candidate.code === "53200" || candidate.code === "53400";
+}
 async function runWithPostgresTraceQueryTimeout(client, timeoutMs, execute) {
 	const resolvedTimeoutMs = coreStorage.resolveTraceQueryTimeoutMs(timeoutMs);
 	try {
@@ -17710,9 +17683,33 @@ async function runWithPostgresTraceQueryTimeout(client, timeoutMs, execute) {
 		});
 	} catch (error) {
 		if (isPostgresStatementTimeout(error)) throw new coreStorage.TraceQueryExecutionError();
+		if (isPostgresResourceLimit(error)) throw new coreStorage.TraceQueryResourceLimitError();
 		throw error;
 	}
 }
+async function getTraceQueryObservedFields(client, schema, plan, timeoutMs) {
+	if (plan.predicateScope !== "trace") return {
+		observedFields: [],
+		observedFieldsTruncated: false
+	};
+	const query = compilePostgresTraceQueryObservedFields(schema, plan);
+	const rows = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(query.text, query.values));
+	return {
+		observedFields: rows.slice(0, plan.limit).map((row) => coreStorage.createTraceQueryObservedFieldDescriptor(String(row.path), Number(row.occurrences))),
+		observedFieldsTruncated: rows.length > plan.limit
+	};
+}
+async function getTraceQueryValues(client, schema, plan, timeoutMs) {
+	const query = compilePostgresTraceQueryValues(schema, plan);
+	const rows = await runWithPostgresTraceQueryTimeout(client, timeoutMs, (transaction) => transaction.any(query.text, query.values));
+	return coreStorage.getTraceQueryValuesResponseSchema.parse({
+		values: rows.slice(0, plan.limit).map((row) => ({
+			value: String(row.value),
+			count: Number(row.count)
+		})),
+		valuesTruncated: rows.length > plan.limit
+	});
+}
 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));
@@ -17731,6 +17728,12 @@ async function queryTraces(client, schema, plan, timeoutMs) {
 	const traces = visibleRows.map((row) => ({
 		traceId: String(row.traceId),
 		rootSpanId: String(row.rootSpanId),
+		name: row.name,
+		entityId: row.entityId ?? null,
+		parentSpanId: row.parentSpanId ?? null,
+		createdAt: asIsoTimestamp$1(row.startedAt),
+		metadata: row.metadata ?? null,
+		inputPreview: coreStorage.buildInputPreview(row.input) ?? null,
 		threadId: row.threadId == null ? null : String(row.threadId),
 		resourceId: row.resourceId == null ? null : String(row.resourceId),
 		startedAt: asIsoTimestamp$1(row.startedAt),
@@ -18478,7 +18481,7 @@ async function dangerouslyClearTracing(client, schema) {
 * Use it through `MastraCompositeStore` with a dedicated Postgres connection.
 */
 function wrapError(op, error, details) {
-	if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError) throw error;
+	if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError || error instanceof coreStorage.TraceQueryResourceLimitError) throw error;
 	throw new MastraError({
 		id: createStorageErrorId("PG", op, "FAILED"),
 		domain: ErrorDomain.STORAGE,
@@ -18659,6 +18662,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"metrics",
 			"logs",
 			"trace-query",
+			"trace-query-discovery",
 			"thread-query"
 		];
 		return [
@@ -18666,6 +18670,7 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 			"logs",
 			"delta-polling",
 			"trace-query",
+			"trace-query-discovery",
 			"thread-query"
 		];
 	}
@@ -18712,6 +18717,12 @@ var ObservabilityStoragePostgresVNext = class ObservabilityStoragePostgresVNext
 	async queryTraces(plan) {
 		return this.#run("QUERY_TRACES", () => queryTraces(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));
+	}
+	async getTraceQueryValues(plan) {
+		return this.#run("GET_TRACE_QUERY_VALUES", () => getTraceQueryValues(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
+	}
 	async queryThreads(plan) {
 		return this.#run("QUERY_THREADS", () => queryThreads(this.#readClient, this.#schema, plan, this.#traceQueryTimeoutMs));
 	}