@mastra/pg 1.31.3-alpha.0 → 1.32.0-alpha.2

package/dist/index.js

@@ -6303,7 +6303,11 @@ var DatasetsPG = class DatasetsPG extends DatasetsStorage {
 			const queryParams = [];
 			let paramIndex = 1;
 			if (args.filters) {
-				const { organizationId, projectId, candidateKey, candidateId, targetType, targetIds, name } = args.filters;
+				const { organizationId, projectId, candidateKey, candidateId, targetType, targetIds, datasetIds, name } = args.filters;
+				if (datasetIds !== void 0) {
+					conditions.push(`"id" = ANY($${paramIndex++}::text[])`);
+					queryParams.push(datasetIds);
+				}
 				if (organizationId !== void 0) {
 					conditions.push(`"organizationId" = $${paramIndex++}`);
 					queryParams.push(organizationId);
@@ -16006,6 +16010,11 @@ const SPAN_EVENT_COLUMNS = [
 		name: "endedAt",
 		type: "timestamptz"
 	},
+	{
+		name: "timeToFirstChunkMs",
+		type: "double precision",
+		nullable: true
+	},
 	{
 		name: "tags",
 		type: "text[]",
@@ -16740,6 +16749,11 @@ function additiveColumns(schema) {
 			column,
 			ddl: `ALTER TABLE ${qualifiedTable(schema, TABLE_SPAN_EVENTS)} ADD COLUMN IF NOT EXISTS "${column}" text`
 		})),
+		{
+			table: TABLE_SPAN_EVENTS,
+			column: "timeToFirstChunkMs",
+			ddl: `ALTER TABLE ${qualifiedTable(schema, TABLE_SPAN_EVENTS)} ADD COLUMN IF NOT EXISTS "timeToFirstChunkMs" double precision`
+		},
 		{
 			table: TABLE_FEEDBACK_EVENTS,
 			column: "reviewStatus",
@@ -17094,9 +17108,27 @@ function readPreview(preview, payload, jsonEncoded = false) {
 }
 /** Older response builders may shorten previews or turn empty strings into null. */
 function buildSpanQueryResponse(...args) {
-	if (canGeneratePreviews) return coreStorage.buildSpanQueryResponse(...args);
-	const [plan, selected, payloads, metrics] = args;
-	const response = coreStorage.buildSpanQueryResponse(plan, selected, [], metrics);
+	const [plan, selected, payloads, metrics = []] = args;
+	const compatibleSelected = coreStorage.SPAN_QUERY_FIELD_REGISTRY ? selected : selected.map((row) => ({
+		organizationId: row.organizationId,
+		resourceId: row.resourceId,
+		traceId: row.traceId,
+		spanId: row.spanId,
+		parentSpanId: row.parentSpanId,
+		name: row.name,
+		spanType: row.spanType,
+		status: row.status,
+		startedAt: row.startedAt,
+		endedAt: row.endedAt,
+		durationMs: row.durationMs,
+		entityType: row.entityType,
+		entityId: row.entityId,
+		entityName: row.entityName,
+		model: row.model,
+		provider: row.provider
+	}));
+	if (canGeneratePreviews) return coreStorage.buildSpanQueryResponse(plan, compatibleSelected, payloads, metrics);
+	const response = coreStorage.buildSpanQueryResponse(plan, compatibleSelected, [], metrics);
 	const byIdentity = new Map(payloads.map((payload) => [coreStorage.spanQueryIdentityKey(payload), payload]));
 	for (const span of response.spans) {
 		const payload = byIdentity.get(coreStorage.spanQueryIdentityKey(span));
@@ -17443,6 +17475,7 @@ function spanRecordToRow(span) {
 		isPending,
 		startedAt: toIsoOrDate(span.startedAt),
 		endedAt: toIsoOrDate(endedAt),
+		timeToFirstChunkMs: span.timeToFirstChunkMs ?? null,
 		tags: normalizeTags(span.tags),
 		metadataSearch: buildMetadataSearch(metadata),
 		metadataRaw: jsonField(metadata),
@@ -18366,6 +18399,7 @@ async function listLogsDelta(client, table, filters, after, limit) {
 function applyMetricFilters(acc, filters) {
 	applyCommonFilters(acc, filters);
 	applySingleOrArrayFilter(acc, "traceId", filters?.traceIds);
+	applySingleOrArrayFilter(acc, "experimentId", filters?.experimentIds);
 	applySingleOrArrayFilter(acc, "name", filters?.name);
 	applySingleOrArrayFilter(acc, "provider", filters?.provider);
 	applySingleOrArrayFilter(acc, "model", filters?.model);
@@ -19192,6 +19226,25 @@ const SPAN_FIELDS = {
 	userId: "s.\"userId\"",
 	organizationId: "s.\"organizationId\""
 };
+const SPAN_QUERY_FIELDS = {
+	...SPAN_FIELDS,
+	traceId: "s.\"traceId\"",
+	spanId: "s.\"spanId\"",
+	parentSpanId: "s.\"parentSpanId\"",
+	threadId: "s.\"threadId\"",
+	resourceId: "s.\"resourceId\"",
+	requestId: "s.\"requestId\"",
+	serviceName: "s.\"serviceName\"",
+	experimentId: "s.\"experimentId\"",
+	environment: "s.\"environment\"",
+	executionSource: "s.\"executionSource\"",
+	parentEntityType: "s.\"parentEntityType\"",
+	parentEntityId: "s.\"parentEntityId\"",
+	parentEntityName: "s.\"parentEntityName\"",
+	rootEntityType: "s.\"rootEntityType\"",
+	rootEntityId: "s.\"rootEntityId\"",
+	rootEntityName: "s.\"rootEntityName\""
+};
 const SCORE_FIELDS = {
 	scorerId: "s.\"scorerId\"",
 	scorerVersion: "s.\"scorerVersion\"",
@@ -20083,9 +20136,9 @@ async function queryThreads(client, schema, plan, timeoutMs) {
 		}) : null }
 	});
 }
-/** Compile a span-row filter with exactly the same rules as trace span predicates. */
+/** Compile a direct span-query filter using its dedicated field mapping. */
 function compileSpanQueryPredicate(predicate, offset) {
-	return compileScalarPredicate(predicate, SPAN_FIELDS, offset);
+	return compileScalarPredicate(predicate, SPAN_QUERY_FIELDS, offset);
 }
 //#endregion
 //#region src/storage/domains/observability/v-next/span-query.ts
@@ -20104,11 +20157,23 @@ const fields = [
 	"entityId",
 	"entityName",
 	"entityVersionId",
+	"parentEntityType",
+	"parentEntityId",
+	"parentEntityName",
 	"parentEntityVersionId",
+	"rootEntityType",
+	"rootEntityId",
+	"rootEntityName",
 	"rootEntityVersionId",
+	"threadId",
 	"runId",
 	"sessionId",
-	"userId"
+	"userId",
+	"requestId",
+	"serviceName",
+	"experimentId",
+	"environment",
+	"executionSource"
 ];
 const identities = identityFields.map((field) => `"${field}"`).join(", ");
 const identityOrder = identityFields.map((field) => `"${field}" COLLATE "C" ASC NULLS LAST`).join(", ");
@@ -20158,6 +20223,8 @@ function compilePostgresSpanQuery(schema, plan) {
 	const after = afterCursor(plan, values);
 	values.push(plan.limit + 1);
 	const table = qualifiedTable(schema, TABLE_SPAN_EVENTS);
+	const legacyModelFields = coreStorage.SPAN_QUERY_FIELD_REGISTRY ? "" : `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,`;
 	return {
 		text: `WITH candidate_ids AS MATERIALIZED (
     SELECT DISTINCT ${identities} FROM ${table}
@@ -20166,8 +20233,7 @@ function compilePostgresSpanQuery(schema, plan) {
   ), 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,
+      ${legacyModelFields}
       (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"
@@ -20182,15 +20248,7 @@ function compilePostgresSpanQuery(schema, plan) {
     SELECT *, CASE WHEN "hasError" THEN 'error' ELSE NULL END AS error FROM current_spans
   )
   SELECT ${[
-			...identityFields,
-			"parentSpanId",
-			"name",
-			"spanType",
-			"entityType",
-			"entityId",
-			"entityName",
-			"model",
-			"provider",
+			...fields,
 			"startedAt",
 			"endedAt",
 			"durationMs",
@@ -20237,7 +20295,7 @@ async function querySpans(client, schema, plan, timeoutMs) {
 			endedAt: endedAt.toISOString()
 		}));
 		const page = rows.slice(0, plan.limit);
-		if (!page.length) return coreStorage.buildSpanQueryResponse(plan, selected, [], []);
+		if (!page.length) return buildSpanQueryResponse(plan, [], []);
 		const remaining = async () => {
 			const ms = Math.floor(deadline - performance.now());
 			if (ms <= 0) throw new coreStorage.TraceQueryExecutionError();
@@ -20250,20 +20308,32 @@ async function querySpans(client, schema, plan, timeoutMs) {
 			payloadValues.push(row.version);
 			return `(${key} AND "cursorId" = $${payloadValues.length}::bigint)`;
 		}).join(" OR ");
-		const payloads = (await tx.any(`SELECT ${identities},
+		const payloadRows = await tx.any(`SELECT ${identities},
+      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,
       ${PREVIEW_PAYLOADS$1.map(([preview, payload]) => `"${preview}", CASE WHEN ${missingPreviewSql(`"${preview}"`)} THEN "${payload}" END AS "${payload}"`).join(",\n      ")}
-      FROM ${qualifiedTable(schema, TABLE_SPAN_EVENTS)} WHERE ${payloadWhere}`, payloadValues)).map(toSpanQueryPayload);
-		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 buildSpanQueryResponse(plan, selected, payloads, metrics);
+      FROM ${qualifiedTable(schema, TABLE_SPAN_EVENTS)} WHERE ${payloadWhere}`, payloadValues);
+		const payloads = payloadRows.map(toSpanQueryPayload);
+		const metadata = new Map(payloadRows.map((row) => [coreStorage.spanQueryIdentityKey(row), row]));
+		const hydrated = selected.map((row) => ({
+			...row,
+			model: metadata.get(coreStorage.spanQueryIdentityKey(row))?.model ?? null,
+			provider: metadata.get(coreStorage.spanQueryIdentityKey(row))?.provider ?? null
+		}));
+		let metrics = [];
+		if (!coreStorage.SPAN_QUERY_FIELD_REGISTRY) {
+			await remaining();
+			const metricValues = [];
+			const metricWhere = selectedKeys(page, metricValues);
+			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 buildSpanQueryResponse(plan, hydrated, payloads, metrics);
 	}, { repeatableRead: true });
 }
 //#endregion
@@ -21307,7 +21377,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"] : [];
+const spanQueryFeatures = typeof coreStorage.planSpanQuery === "function" ? ["span-query", "span-query-context-fields"] : [];
 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({