@mastra/clickhouse 1.23.0-alpha.2 → 1.23.0

package/dist/index.js

@@ -1689,6 +1689,7 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
 			});
 			const threadIdsToUpdate = /* @__PURE__ */ new Set();
 			const updatePromises = [];
+			const writtenContent = /* @__PURE__ */ new Map();
 			for (const existingMessage of parsedExistingMessages) {
 				const updatePayload = messages.find((m) => m.id === existingMessage.id);
 				if (!updatePayload) continue;
@@ -1715,6 +1716,7 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
 					};
 					setClauses.push(`content = {var_content_${paramIdx}:String}`);
 					values[`var_content_${paramIdx}`] = JSON.stringify(newContent);
+					writtenContent.set(id, values[`var_content_${paramIdx}`]);
 					paramIdx++;
 					delete updatableFields.content;
 				}
@@ -1731,7 +1733,7 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
                 UPDATE ${setClauses.join(", ")}
                 WHERE id = {var_id_${paramIdx}:String}
               `;
-					console.info("Updating message:", id, "with query:", updateQuery, "values:", values);
+					this.logger?.debug?.("Updating message", { id });
 					updatePromises.push(this.client.command({
 						query: updateQuery,
 						query_params: values,
@@ -1767,7 +1769,7 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
 					if (updatedMessage) {
 						let needsRetry = false;
 						for (const [key, value] of Object.entries(fieldsToUpdate)) if (key === "content") {
-							const expectedContent = typeof value === "string" ? value : JSON.stringify(value);
+							const expectedContent = writtenContent.get(id);
 							if ((typeof updatedMessage.content === "string" ? updatedMessage.content : JSON.stringify(updatedMessage.content)) !== expectedContent) {
 								needsRetry = true;
 								break;
@@ -1777,7 +1779,7 @@ var MemoryStorageClickhouse = class extends MemoryStorage {
 							break;
 						}
 						if (needsRetry) {
-							console.info("Update not applied correctly, retrying with DELETE + INSERT for message:", id);
+							this.logger?.warn?.("Update not applied correctly, retrying with DELETE + INSERT", { id });
 							await this.client.command({
 								query: `DELETE FROM ${TABLE_MESSAGES} WHERE id = {messageId:String}`,
 								query_params: { messageId: id },
@@ -2412,7 +2414,7 @@ var ObservabilityStorageClickhouse = class extends ObservabilityStorage {
 				if (filters.hasChildError !== void 0) {
 					const finalClause = (TABLE_ENGINES[TABLE_SPANS] ?? "MergeTree()").startsWith("ReplacingMergeTree") ? "FINAL" : "";
 					conditions.push(`traceId ${filters.hasChildError ? "IN" : "NOT IN"} (
-            SELECT c.traceId FROM ${TABLE_SPANS} ${finalClause} c
+            SELECT c.traceId FROM ${TABLE_SPANS} c ${finalClause}
             WHERE c.error IS NOT NULL AND c.error != ''
           )`);
 				}
@@ -7861,12 +7863,31 @@ function compileClickHouseTraceAggregate(plan) {
 		if (intervalMs === void 0) throw new Error(`Unsupported trusted trace-aggregate interval: ${plan.interval}`);
 		factColumns.push(`fromUnixTimestamp64Milli(intDiv(toUnixTimestamp64Milli(r.startedAt), ${intervalMs}) * ${intervalMs}, 'UTC') AS bucket`);
 	}
+	const measureKinds = new Set(plan.measures.map((measure) => measureRule(measure.name)?.kind));
+	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";
+		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
+    FROM candidates r${usageJoin}
   )`);
+	const mixedUnitsSql = "minOrNull(unitMin) != maxOrNull(unitMax)";
 	const measureSql = (name) => {
 		if (name === "count") return "toFloat64(count())";
+		const rule = measureRule(name);
+		if (rule?.kind === "tokens" && rule.metricNames) {
+			const value = rule.metricNames.map((metricName) => `t${usageMetricIndex(metricName)}`).join(" + ");
+			return `toFloat64(${rule.statistic === "avg" ? "avgOrNull" : "sumOrNull"}(${value}))`;
+		}
+		if (rule?.kind === "cost") {
+			const aggregate = rule.statistic === "avg" ? "avgOrNull" : "sumOrNull";
+			return `if(${mixedUnitsSql}, NULL, toFloat64(${aggregate}(traceCost)))`;
+		}
 		if (name === "errorCount") return "toFloat64(countIf(isError))";
 		if (name === "errorRate") return "countIf(isError) / count()";
 		if (name === "duration.avg") return "avg(durationMs)";
@@ -7904,6 +7925,11 @@ function compileClickHouseTraceAggregate(plan) {
 	const tiebreakSql = dimensionColumns.map((column) => `${column} ASC NULLS LAST`);
 	const measureColumns = plan.measures.map((measure, index) => `${measureSql(measure.name)} AS m${index}`);
 	const measureNames = plan.measures.map((_, index) => `m${index}`);
+	if (hasCost) {
+		const mixedUnit = parameters.add(coreStorage.TRACE_AGGREGATE_MIXED_COST_UNIT, "String");
+		measureColumns.push(`countIf(covered) / nullIf(countIf(usageBearing), 0) AS costCoverage`, `if(${mixedUnitsSql}, ${mixedUnit}, minOrNull(unitMin)) AS costUnit`);
+		measureNames.push("costCoverage", "costUnit");
+	}
 	if (plan.interval === void 0) {
 		const limit = parameters.add(plan.limit + 1, "UInt64");
 		return {
@@ -7925,7 +7951,7 @@ LIMIT ${limit}`,
 		"bucket",
 		"grouping(bucket) = 1 AS __isGroup",
 		...measureColumns,
-		`(${havingSql}) AS __keep`,
+		`ifNull(${havingSql}, 0) AS __keep`,
 		`${orderTarget} AS __order`
 	].join(", ")}
     FROM facts
@@ -7953,6 +7979,70 @@ ORDER BY __rank ASC, bucket ASC`,
 		query_params: parameters.params
 	};
 }
+const USAGE_METRIC_NAMES = coreStorage.TRACE_AGGREGATE_USAGE_METRIC_NAMES;
+function measureRule(name) {
+	return coreStorage.isTraceAggregateCanonicalMeasure(name) ? coreStorage.TRACE_AGGREGATE_MEASURE_REGISTRY[name] : void 0;
+}
+function usageMetricIndex(name) {
+	const index = USAGE_METRIC_NAMES.indexOf(name);
+	if (index === -1) throw new Error(`Unsupported trusted trace-aggregate token metric: ${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.
+*
+* 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) {
+	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 failed = `${costRow} AND (hasError OR (isNotNull(estimatedCost) AND isNull(costUnit)))`;
+	return `usage AS (
+    SELECT traceId,
+      toUInt8(1) AS hasUsage,
+      ${names.map((name, index) => `sumIf(value, name = ${name}) AS t${index}`).join(",\n      ")},
+      sumIf(assumeNotNull(estimatedCost), ${priced}) AS cost,
+      toUInt8(countIf(${priced}) > 0) AS priced,
+      toUInt8(countIf(${failed}) > 0) AS pricingFailure,
+      minIf(assumeNotNull(costUnit), ${priced}) AS unitMin,
+      maxIf(assumeNotNull(costUnit), ${priced}) AS unitMax
+    FROM (
+      SELECT traceId,
+        latest.1 AS name,
+        latest.2 AS value,
+        latest.3 AS estimatedCost,
+        latest.4 AS costUnit,
+        latest.5 AS hasError
+      FROM (
+        SELECT traceId,
+          argMax(tuple(name, value, estimatedCost, costUnit, hasError), 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
+          FROM ${TABLE_METRIC_EVENTS}
+          WHERE traceId IN (SELECT traceId FROM candidates)
+            AND name IN (${names.join(", ")})
+            AND timestamp >= ${from}${tenant}
+        )
+        GROUP BY traceId, metricId
+      )
+    )
+    GROUP BY traceId
+  )`;
+}
+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));
 	let truncated;
@@ -7966,7 +8056,11 @@ async function aggregateTraces(client, plan, timeoutMs) {
 	}
 	return {
 		rows: resultRows.map((row) => {
-			const shaped = { measures: Object.fromEntries(plan.measures.map((measure, index) => [measure.name, Number(row[`m${index}`])])) };
+			const shaped = { measures: Object.fromEntries(plan.measures.map((measure, index) => [measure.name, nullableNumber(row[`m${index}`])])) };
+			if ("costCoverage" in row) shaped.cost = {
+				coverage: nullableNumber(row.costCoverage),
+				unit: row.costUnit ?? null
+			};
 			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 = new Date(row.bucket).toISOString();
 			return shaped;