@mastra/clickhouse 1.26.0 → 1.27.0-alpha.0

package/dist/index.js

@@ -3648,6 +3648,9 @@ const ALL_MIGRATIONS = [
 	addColumn(TABLE_SPAN_EVENTS, "inputPreview", "Nullable(String)"),
 	addColumn(TABLE_SPAN_EVENTS, "outputPreview", "Nullable(String)"),
 	addColumn(TABLE_SPAN_EVENTS, "errorPreview", "Nullable(String)"),
+	addColumn(TABLE_TRACE_BRANCHES, "entityVersionId", "Nullable(String)"),
+	addColumn(TABLE_TRACE_BRANCHES, "parentEntityVersionId", "Nullable(String)"),
+	addColumn(TABLE_TRACE_BRANCHES, "rootEntityVersionId", "Nullable(String)"),
 	addColumn(TABLE_SPAN_EVENTS, "entityVersionId", "Nullable(String)"),
 	addColumn(TABLE_SPAN_EVENTS, "parentEntityVersionId", "Nullable(String)"),
 	addColumn(TABLE_SPAN_EVENTS, "rootEntityVersionId", "Nullable(String)"),
@@ -7989,8 +7992,7 @@ async function queryThreads(client, plan, timeoutMs) {
 	});
 }
 /** Compile a direct span-query filter using its dedicated field mapping. */
-function compileSpanQueryPredicate(predicate) {
-	const parameters = new ParameterBuilder();
+function compileSpanQueryPredicate(predicate, parameters = new ParameterBuilder()) {
 	return {
 		sql: compileScalarPredicate(predicate, SPAN_QUERY_FIELDS, parameters),
 		params: parameters.params
@@ -8151,6 +8153,131 @@ function compileClickHouseSpanQuery(plan) {
 		}
 	};
 }
+const branchSpanTypes = new Set(BRANCH_SPAN_TYPE_VALUES);
+const sameColumn = (field) => [field];
+/**
+* Span fields `trace_branches` can serve, as the `trace_branches` columns each one reads. The two
+* tables are expected to diverge, so this list is explicit: a selection whose fields are not all
+* listed reads `span_events`. A new span-query field has to be classified here deliberately, and a
+* missing column then costs speed, not correctness.
+*/
+const TRACE_BRANCHES_SPAN_FIELD_COLUMNS = {
+	...Object.fromEntries([
+		"organizationId",
+		"resourceId",
+		"traceId",
+		"spanId",
+		"parentSpanId",
+		"name",
+		"spanType",
+		"startedAt",
+		"endedAt",
+		"timeToFirstChunkMs",
+		"threadId",
+		"runId",
+		"sessionId",
+		"userId",
+		"requestId",
+		"serviceName",
+		"experimentId",
+		"environment",
+		"executionSource",
+		"entityType",
+		"entityId",
+		"entityName",
+		"entityVersionId",
+		"parentEntityType",
+		"parentEntityId",
+		"parentEntityName",
+		"parentEntityVersionId",
+		"rootEntityType",
+		"rootEntityId",
+		"rootEntityName",
+		"rootEntityVersionId"
+	].map((field) => [field, sameColumn(field)])),
+	status: ["error"],
+	error: ["error"],
+	durationMs: ["startedAt", "endedAt"]
+};
+/**
+* Whether a span selection reads `trace_branches`: `where` must restrict `spanType` to branch
+* types, and `trace_branches` must serve every selected field (`served`, for tests).
+*/
+function readsTraceBranches(where, fields, served = TRACE_BRANCHES_SPAN_FIELD_COLUMNS) {
+	return positiveSpanConjuncts(where).some(({ field, values }) => field === "spanType" && values.length > 0 && values.every((value) => branchSpanTypes.has(String(value)))) && [...fields].every((field) => Object.hasOwn(served, field));
+}
+/**
+* Top-level `spanType` / `name` equality or `in` conditions. Every span the predicate can match
+* satisfies them, and copies of a span never differ in either field, so they can be applied
+* before collapsing copies.
+*/
+function positiveSpanConjuncts(where) {
+	return (where?.type === "boolean" && where.operator === "and" ? where.args : where ? [where] : []).flatMap((conjunct) => {
+		if (conjunct.type !== "comparison" && conjunct.type !== "membership") return [];
+		if (conjunct.field !== "spanType" && conjunct.field !== "name") return [];
+		if (conjunct.type === "comparison" && conjunct.operator === "eq") return [{
+			field: conjunct.field,
+			values: [conjunct.value]
+		}];
+		if (conjunct.type === "membership" && conjunct.operator === "in") return [{
+			field: conjunct.field,
+			values: conjunct.values
+		}];
+		return [];
+	});
+}
+/**
+* The spans `querySpans()` selects, as CTEs ending in `selected` (alias each row as you like):
+* one row per current completed span with `startedAt` in the window, matching `where`, carrying
+* the identity, `startedAt`, `endedAt`, `status`, and `fields`.
+*
+* Copies of a span collapse to the copy with the latest `endedAt`. When `where` requires branch
+* span types (agent and workflow runs, tool calls, …) and `trace_branches` serves every selected
+* field, the spans are read from `trace_branches`, whose sort key starts with `spanType`, instead of
+* every span in the window. That path can differ from `querySpans()` in two ways:
+* - `trace_branches` is filled only by its materialized view, with no backfill, so an install that
+*   upgraded from @mastra/clickhouse 1.4–1.6 has no branch rows for spans written before then.
+* - It is a `ReplacingMergeTree` without a version column, so once a background merge collapses a
+*   retried span's copies it keeps the last-inserted copy, not the one with the latest `endedAt`.
+*   Counts are unaffected; `status` and duration of an out-of-order retry can differ.
+*
+* Like `querySpans()`, the read has no upper bound on `endedAt` (`PREWHERE endedAt >= from` only,
+* since a span that starts in the window can end any time later), so a historical window reads
+* every span that ended since `from`.
+*/
+function compileClickHouseSpanSelection(plan, fields, parameters) {
+	const from = parameters.add(plan.timeRange.from, "DateTime64(3, 'UTC')");
+	const to = parameters.add(plan.timeRange.to, "DateTime64(3, 'UTC')");
+	const prefilter = positiveSpanConjuncts(plan.where).map(({ field, values }) => values.length ? `\n      AND e.${field} IN (${values.map((value) => parameters.add(String(value), "String")).join(", ")})` : "\n      AND 0").join("");
+	const predicate = plan.where ? compileSpanQueryPredicate(plan.where, parameters).sql : "1";
+	let tenant = "";
+	if (plan.scope) {
+		tenant += `\n      AND e.organizationId = ${parameters.add(plan.scope.organizationId, "String")}`;
+		if (plan.scope.resourceId !== void 0) tenant += `\n      AND e.resourceId = ${parameters.add(plan.scope.resourceId, "String")}`;
+	}
+	const selectionFields = [.../* @__PURE__ */ new Set([
+		"startedAt",
+		"endedAt",
+		"status",
+		...fields,
+		...predicateFields$1(plan.where)
+	])].filter((field) => !new Set(identityFields).has(field));
+	const branches = readsTraceBranches(plan.where, selectionFields);
+	const scan = `FROM ${branches ? TABLE_TRACE_BRANCHES : TABLE_SPAN_EVENTS} e
+    PREWHERE e.endedAt >= ${from}
+    WHERE e.startedAt >= ${from} AND e.startedAt < ${to}${tenant}${prefilter}`;
+	return [`span_candidates AS (
+    ${branches ? `SELECT ${identityFields.map((field) => `e.${field} AS ${field}`).join(", ")}, ${selectionFields.map((field) => field === "endedAt" ? "max(e.endedAt) AS endedAt" : `argMax(tuple(${fieldExpression(field)}), e.endedAt).1 AS ${field}`).join(", ")}
+    ${scan}
+    GROUP BY ${identityFields.map((field) => `e.${field}`).join(", ")}` : `SELECT ${identities}, ${selectionFields.map((field) => `${fieldExpression(field)} AS ${field}`).join(", ")}
+    ${scan}
+    ORDER BY ${identities}, e.endedAt DESC
+    LIMIT 1 BY ${identities}`}
+  )`, `selected AS (
+    SELECT * FROM span_candidates s
+    WHERE s.startedAt >= ${from} AND s.startedAt < ${to} AND (${predicate})
+  )`];
+}
 const LEGACY_PAYLOAD_BATCH_ROWS = 32;
 function iso$1(value) {
 	return new Date(value.includes("T") ? value : `${value.replace(" ", "T")}Z`).toISOString();
@@ -8554,6 +8681,28 @@ const DIMENSION_SQL = {
 	organizationId: "r.organizationId",
 	experimentId: "r.experimentId"
 };
+/** Span column expressions (alias `r`, a `selected` span row) for each groupable span dimension. */
+const SPAN_DIMENSION_SQL = Object.fromEntries([
+	"name",
+	"spanType",
+	"entityType",
+	"entityId",
+	"entityName",
+	"parentEntityType",
+	"parentEntityName",
+	"rootEntityType",
+	"rootEntityName",
+	"environment",
+	"status",
+	"serviceName",
+	"executionSource",
+	"threadId",
+	"resourceId",
+	"userId",
+	"sessionId",
+	"organizationId",
+	"experimentId"
+].map((field) => [field, `r.${field}`]));
 const PERCENTILES = {
 	"duration.p50": .5,
 	"duration.p90": .9,
@@ -8580,13 +8729,41 @@ const COMPARISON_SQL = {
 */
 function compileClickHouseTraceAggregate(plan) {
 	const parameters = new ParameterBuilder();
-	const ctes = compileClickHouseTraceCandidates(plan, "*", parameters, { oneRootPerTrace: true });
-	const dimensionSql = (field) => {
-		if (field.startsWith("metadata.")) return traceMetadataValueSql(parameters.add(field.slice(9), "String"));
-		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;
-	};
+	return compileClickHouseAggregate(plan, parameters, {
+		ctes: compileClickHouseTraceCandidates(plan, "*", parameters, { oneRootPerTrace: true }),
+		from: "candidates",
+		dimensionSql: (field) => {
+			if (field.startsWith("metadata.")) return traceMetadataValueSql(parameters.add(field.slice(9), "String"));
+			return lookupDimensionSql(DIMENSION_SQL, field);
+		},
+		isErrorSql: "isNotNull(r.error)",
+		usageKey: ["traceId"],
+		usageBound: "traceId IN (SELECT traceId FROM candidates)"
+	});
+}
+/**
+* Compiles a trusted span aggregate plan over the spans `querySpans()` selects. Usage is
+* attributed by `(traceId, spanId)`, so only the span that recorded a metric row carries it.
+*/
+function compileClickHouseSpanAggregate(plan) {
+	const parameters = new ParameterBuilder();
+	return compileClickHouseAggregate(plan, parameters, {
+		ctes: compileClickHouseSpanSelection(plan, [...plan.dimensions, ...plan.measures.flatMap((measure) => measure.type === "countDistinct" && measure.field !== "traceId" ? [measure.field] : [])], parameters),
+		from: "selected",
+		dimensionSql: (field) => lookupDimensionSql(SPAN_DIMENSION_SQL, field),
+		isErrorSql: "r.status = 'error'",
+		usageKey: ["traceId", "spanId"],
+		usageBound: "isNotNull(spanId)"
+	});
+}
+function lookupDimensionSql(columns, field) {
+	const sql = Object.hasOwn(columns, field) ? columns[field] : void 0;
+	if (sql === void 0) throw new Error(`Unsupported trusted aggregate field: ${field}`);
+	return sql;
+}
+function compileClickHouseAggregate(plan, parameters, source) {
+	const { ctes } = source;
+	const dimensionSql = source.dimensionSql;
 	const distinctSql = (field) => field === "traceId" ? "r.traceId" : dimensionSql(field);
 	const dimensionColumns = plan.dimensions.map((_, index) => `d${index}`);
 	const distinctColumns = /* @__PURE__ */ new Map();
@@ -8597,7 +8774,7 @@ function compileClickHouseTraceAggregate(plan) {
 		distinctColumns.set(measure.field, column);
 		factColumns.push(`CAST(${distinctSql(measure.field)} AS Nullable(String)) AS ${column}`);
 	}
-	factColumns.push(`toFloat64(${durationMsSql("r.startedAt", "r.endedAt")}) AS durationMs`, `isNotNull(r.error) AS isError`, `cityHash64(r.traceId) AS traceSeed`);
+	factColumns.push(`toFloat64(${durationMsSql("r.startedAt", "r.endedAt")}) AS durationMs`, `${source.isErrorSql} AS isError`, `cityHash64(${source.usageKey.map((column) => `r.${column}`).join(", ")}) AS traceSeed`);
 	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}`);
@@ -8607,14 +8784,14 @@ function compileClickHouseTraceAggregate(plan) {
 	const hasCost = measureKinds.has("cost");
 	let usageJoin = "";
 	if (hasCost || measureKinds.has("tokens")) {
-		ctes.push(compileUsageCte(plan, parameters));
-		usageJoin = "\n    LEFT JOIN usage u ON u.traceId = r.traceId";
+		ctes.push(compileUsageCte(plan, parameters, source));
+		usageJoin = `\n    LEFT JOIN usage u ON ${source.usageKey.map((column) => `u.${column} = r.${column}`).join(" AND ")}`;
 		USAGE_METRIC_NAMES.forEach((_, index) => factColumns.push(`if(u.hasUsage = 1, u.t${index}, NULL) AS t${index}`));
 		factColumns.push(`u.hasUsage = 1 AS usageBearing`, `if(u.hasUsage = 1 AND u.priced = 1, u.cost, NULL) AS traceCost`, `u.hasUsage = 1 AND u.priced = 1 AND u.pricingFailure = 0 AS covered`, `if(u.hasUsage = 1 AND u.priced = 1, u.unitMin, NULL) AS unitMin`, `if(u.hasUsage = 1 AND u.priced = 1, u.unitMax, NULL) AS unitMax`);
 	}
 	ctes.push(`facts AS (
     SELECT ${factColumns.join(",\n      ")}
-    FROM candidates r${usageJoin}
+    FROM ${source.from} r${usageJoin}
   )`);
 	const mixedUnitsSql = "minOrNull(unitMin) != maxOrNull(unitMax)";
 	const measureSql = (name) => {
@@ -8729,26 +8906,28 @@ function usageMetricIndex(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.
+* Builds `usage`: one row per usage key (a trace, or a span) that has at least one token metric
+* row, holding per-name token sums (`t<i>`, indexed by `TRACE_AGGREGATE_USAGE_METRIC_NAMES`) and
+* its pricing state. Metric rows are matched by the source's usage bound, usage metric name,
+* tenant scope, and `timestamp >= from` only: a trace or span that starts in the window can emit
+* metrics after `to`. For spans that makes the read cost grow with the age of `from` (see
+* `usageBound`).
 *
 * Rows are deduplicated with `LIMIT 1 BY metricId` at query time. `ReplacingMergeTree` only
 * collapses rows with equal sort keys `(name, timestamp, metricId)`, so retried copies with
 * different timestamps never merge, and copies with equal keys stay separate until a background
 * merge runs.
 */
-function compileUsageCte(plan, parameters) {
+function compileUsageCte(plan, parameters, source) {
+	const key = source.usageKey.join(", ");
 	const names = USAGE_METRIC_NAMES.map((name) => parameters.add(name, "String"));
 	const costRow = `name IN (${coreStorage.TRACE_AGGREGATE_COST_METRIC_NAMES.map((name) => names[usageMetricIndex(name)]).join(", ")})`;
 	const from = parameters.add(plan.timeRange.from, "DateTime64(3, 'UTC')");
 	const tenant = compileTenantScope(plan.scope, parameters);
-	const priced = `${costRow} AND isNotNull(estimatedCost) AND isNotNull(costUnit) AND NOT hasError`;
+	const priced = `${costRow} AND isNotNull(estimatedCost) AND isNotNull(costUnit) AND (NOT hasError OR partialCost)`;
 	const failed = `${costRow} AND (hasError OR (isNotNull(estimatedCost) AND isNull(costUnit)))`;
 	return `usage AS (
-    SELECT traceId,
+    SELECT ${key},
       toUInt8(1) AS hasUsage,
       ${names.map((name, index) => `sumIf(value, name = ${name}) AS t${index}`).join(",\n      ")},
       sumIf(assumeNotNull(estimatedCost), ${priced}) AS cost,
@@ -8757,34 +8936,42 @@ function compileUsageCte(plan, parameters) {
       minIf(assumeNotNull(costUnit), ${priced}) AS unitMin,
       maxIf(assumeNotNull(costUnit), ${priced}) AS unitMax
     FROM (
-      SELECT traceId,
+      SELECT ${key},
         latest.1 AS name,
         latest.2 AS value,
         latest.3 AS estimatedCost,
         latest.4 AS costUnit,
-        latest.5 AS hasError
+        latest.5 AS hasError,
+        latest.6 AS partialCost
       FROM (
-        SELECT traceId,
-          argMax(tuple(name, value, estimatedCost, costUnit, hasError), timestamp) AS latest
+        SELECT ${key},
+          argMax(tuple(name, value, estimatedCost, costUnit, hasError, partialCost), timestamp) AS latest
         FROM (
-          SELECT traceId, metricId, timestamp, name, value, estimatedCost, costUnit,
-            if(${costRow}, ifNull(JSONHas(costMetadata, 'error') AND JSONType(costMetadata, 'error') != 'Null', 0), 0) AS hasError
+          SELECT ${key}, metricId, timestamp, name, value, estimatedCost, costUnit,
+            if(${costRow}, ifNull(JSONHas(costMetadata, 'error') AND JSONType(costMetadata, 'error') != 'Null', 0), 0) AS hasError,
+            if(${costRow}, ifNull(JSONExtractString(costMetadata, 'error') = 'partial_cost', 0), 0) AS partialCost
           FROM ${TABLE_METRIC_EVENTS}
-          WHERE traceId IN (SELECT traceId FROM candidates)
+          WHERE ${source.usageBound}
             AND name IN (${names.join(", ")})
             AND timestamp >= ${from}${tenant}
         )
-        GROUP BY traceId, metricId
+        GROUP BY ${key}, metricId
       )
     )
-    GROUP BY traceId
+    GROUP BY ${key}
   )`;
 }
 function nullableNumber(value) {
 	return value === null || value === void 0 ? null : Number(value);
 }
-async function aggregateTraces(client, plan, timeoutMs) {
-	const result = await runWithClickHouseTraceQueryTimeout(client, { timeoutMs }, compileClickHouseTraceAggregate(plan));
+function aggregateTraces(client, plan, timeoutMs) {
+	return runAggregate(client, plan, compileClickHouseTraceAggregate(plan), timeoutMs);
+}
+async function aggregateSpans(client, plan, timeoutMs) {
+	return await runAggregate(client, plan, compileClickHouseSpanAggregate(plan), timeoutMs);
+}
+async function runAggregate(client, plan, compiled, timeoutMs) {
+	const result = await runWithClickHouseTraceQueryTimeout(client, { timeoutMs }, compiled);
 	let truncated;
 	let resultRows;
 	if (plan.interval === void 0) {
@@ -9680,6 +9867,7 @@ function buildBranchCursor(row) {
 //#endregion
 //#region src/storage/domains/observability/v-next/index.ts
 const spanQueryFeatures = typeof coreStorage.planSpanQuery === "function" ? ["span-query", "span-query-context-fields"] : [];
+const spanAggregateFeatures = typeof coreStorage.planSpanAggregate === "function" ? ["span-aggregate"] : [];
 const threadQueryFeatures = typeof coreStorage.planThreadRowsQuery === "function" ? ["thread-query-rows"] : [];
 function buildSignalMigrationRequiredMessage(args) {
 	const tableList = args.tables.map((table) => `  - ${table.table} (${table.engine})`).join("\n");
@@ -10123,6 +10311,7 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 			"trace-query-root-entity-id",
 			"trace-aggregate",
 			...spanQueryFeatures,
+			...spanAggregateFeatures,
 			...threadQueryFeatures
 		];
 		return [
@@ -10147,6 +10336,7 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 			"trace-query-root-entity-id",
 			"trace-aggregate",
 			...spanQueryFeatures,
+			...spanAggregateFeatures,
 			...threadQueryFeatures
 		];
 	}
@@ -10298,6 +10488,18 @@ var ObservabilityStorageClickhouseVNext = class extends ObservabilityStorage {
 			}, error);
 		}
 	}
+	async aggregateSpans(plan) {
+		try {
+			return await aggregateSpans(this.#client, plan, this.#traceQueryTimeoutMs);
+		} catch (error) {
+			if (error instanceof MastraError || error instanceof coreStorage.TraceQueryExecutionError || error instanceof coreStorage.TraceQueryResourceLimitError) throw error;
+			throw new MastraError({
+				id: createStorageErrorId("CLICKHOUSE", "AGGREGATE_SPANS", "FAILED"),
+				domain: ErrorDomain.STORAGE,
+				category: ErrorCategory.THIRD_PARTY
+			}, error);
+		}
+	}
 	async getTraceQueryObservedFields(plan) {
 		try {
 			return await getTraceQueryObservedFields(this.#client, plan, this.#traceQueryDiscoveryLimits);