perf(agent-sessions): project the span read down to the keys the mapper reads

apps/api/src/routes/internal/ai-sessions.http.test.ts

@@ -68,11 +68,7 @@ const spanRow = (index: number) => ({
 	statusCode: "Unset",
 	statusMessage: "",
 	timestamp: "2026-08-19 10:00:00.000000000",
-	spanAttributes: {
-		"gen_ai.operation.name": "chat",
-		"maple_ai.session.id": SESSION_ID,
-	},
-	resourceAttributes: {},
+	spanAttributes: { "gen_ai.operation.name": "chat", "maple_ai.session.id": SESSION_ID },
 })
 
 const makeHarness = (overrides: Partial<WarehouseQueryServiceApi>) => {
@@ -249,7 +245,8 @@ describe("POST /internal/ai-sessions/spans", () => {
 			expect(windowSql).toContain(`TraceId = '${TRACE_ID}'`)
 			expect(windowSql).not.toContain("maple_ai.session.id")
 			expect(spansSql).toContain(`TraceId = '${TRACE_ID}'`)
-			expect(spansSql).not.toContain("maple_ai.session.id")
+			// The projection names the key; the predicate is what must be absent.
+			expect(spansSql).not.toContain("SpanAttributes['maple_ai.session.id']")
 			// The bounds the window read handed back still prune the span read.
 			expect(spansSql).toContain(`Timestamp >= '${resolved.startTime}'`)
 			expect(spansSql).not.toContain("__PARAM_")

lib/clickhouse-builder/src/ch/core-dsl.test.ts

@@ -88,6 +88,14 @@ describe("expression functions", () => {
 		expect(sql).toContain("map('key1', Name, 'key2', 'val') AS m")
 	})
 
+	it("compiles mapFilterKeys with the DSL's own conditions on the key", () => {
+		const q = CH.from(TestTable).select(($) => ({
+			m: CH.mapFilterKeys($.Attrs, (k) => k.in_("a", "b").or(k.like("x.%"))),
+		}))
+		const { sql } = compileCHUnsafe(q, {})
+		expect(sql).toContain("mapFilter((k, v) -> (k IN ('a', 'b') OR k LIKE 'x.%'), Attrs) AS m")
+	})
+
 	it("compiles empty mapLiteral", () => {
 		const q = CH.from(TestTable).select(() => ({ m: CH.mapLiteral() }))
 		const { sql } = compileCHUnsafe(q, {})

lib/clickhouse-builder/src/ch/functions/index.ts

@@ -90,7 +90,7 @@ export {
 	has,
 } from "./array"
 
-export { mapContains, mapGet, mapKeys, mapValues, mapLiteral } from "./map"
+export { mapContains, mapFilterKeys, mapGet, mapKeys, mapValues, mapLiteral } from "./map"
 
 export { toJSONString } from "./json"
 

lib/clickhouse-builder/src/ch/functions/map.ts

@@ -24,6 +24,23 @@ export function mapValues(mapExpr: Expr<Record<string, string>>): Expr<ReadonlyA
 	return makeExpr(raw(`mapValues(${compile(mapExpr.toFragment())})`), STRINGS)
 }
 
+/**
+ * `mapFilter((k, v) -> <predicate>, map)` — the entries whose KEY passes.
+ *
+ * The predicate is built from the lambda's key parameter, so it can use every
+ * condition the DSL has (`in_`, `like`, `or`, …). Values are not inspected.
+ */
+export function mapFilterKeys(
+	mapExpr: Expr<Record<string, string>>,
+	predicate: (key: Expr<string>) => Condition,
+): Expr<Record<string, string>> {
+	const key = makeExpr(raw("k"), T.string.schema)
+	return makeExpr(
+		raw(`mapFilter((k, v) -> ${compile(predicate(key).toFragment())}, ${compile(mapExpr.toFragment())})`),
+		STRING_MAP,
+	)
+}
+
 export function mapLiteral(...pairs: Array<[string, Expr<string>]>): Expr<Record<string, string>> {
 	if (pairs.length === 0) return makeExpr(raw("map()"), STRING_MAP)
 	const args = pairs.map(([k, v]) => `${compile(str(k))}, ${compile(v.toFragment())}`).join(", ")

lib/clickhouse-builder/src/ch/index.ts

@@ -203,6 +203,7 @@ export {
 	mapGet,
 	mapKeys,
 	mapValues,
+	mapFilterKeys,
 	mapLiteral,
 	// JSON
 	toJSONString,

packages/query-engine-integrations/src/__sql_baseline__/integrations.sql

@@ -517,8 +517,7 @@ SELECT
           StatusCode AS statusCode,
           StatusMessage AS statusMessage,
           toString(Timestamp) AS timestamp,
-          SpanAttributes AS spanAttributes,
-          ResourceAttributes AS resourceAttributes
+          mapFilter((k, v) -> (k IN ('maple_ai.session.id', 'maple_ai.vendor.id', 'maple_ai.vendor.version', 'gen_ai.operation.name', 'gen_ai.provider.name', 'gen_ai.system', 'gen_ai.request.model', 'gen_ai.request.max_tokens', 'gen_ai.request.choice.count', 'gen_ai.request.temperature', 'gen_ai.request.top_p', 'gen_ai.request.top_k', 'gen_ai.request.stop_sequences', 'gen_ai.request.frequency_penalty', 'gen_ai.request.presence_penalty', 'gen_ai.request.encoding_formats', 'gen_ai.request.seed', 'gen_ai.openai.request.seed', 'gen_ai.request.stream', 'gen_ai.request.reasoning.level', 'gen_ai.request.previous_response.id', 'gen_ai.request.stream_cursor', 'gen_ai.response.id', 'gen_ai.response.model', 'gen_ai.response.finish_reasons', 'gen_ai.response.finish_reason', 'gen_ai.response.status', 'gen_ai.response.time_to_first_chunk', 'gen_ai.output.type', 'gen_ai.usage.input_tokens', 'gen_ai.usage.prompt_tokens', 'gen_ai.usage.cache_read.input_tokens', 'gen_ai.usage.input_tokens.cached', 'gen_ai.usage.cache_creation.input_tokens', 'gen_ai.usage.cache_write.input_tokens', 'gen_ai.usage.output_tokens', 'gen_ai.usage.completion_tokens', 'gen_ai.usage.reasoning.output_tokens', 'gen_ai.usage.output_tokens.reasoning', 'gen_ai.usage.cost', 'gen_ai.usage.total_cost', 'gen_ai.conversation.id', 'gen_ai.conversation.compacted', 'gen_ai.agent.id', 'gen_ai.agent.name', 'gen_ai.agent.description', 'gen_ai.agent.version', 'gen_ai.tool.name', 'gen_ai.tool.call.id', 'gen_ai.tool.description', 'gen_ai.tool.type', 'gen_ai.tool.call.arguments', 'gen_ai.tool.call.result', 'gen_ai.tool.definitions', 'gen_ai.system_instructions', 'gen_ai.input.messages', 'gen_ai.prompt', 'gen_ai.output.messages', 'gen_ai.completion', 'gen_ai.data_source.id', 'gen_ai.retrieval.query.text', 'gen_ai.retrieval.top_k', 'gen_ai.retrieval.documents', 'gen_ai.memory.store.id', 'gen_ai.memory.record.id', 'gen_ai.memory.record.count', 'gen_ai.memory.query.text', 'gen_ai.memory.records', 'gen_ai.embeddings.dimension.count', 'gen_ai.evaluation.name', 'gen_ai.evaluation.score.value', 'gen_ai.evaluation.score.label', 'gen_ai.evaluation.explanation', 'gen_ai.prompt.name', 'gen_ai.prompt.version', 'gen_ai.workflow.name', 'error.type', 'server.address', 'server.port', 'ai.model.provider', 'ai.model.id', 'ai.response.id', 'ai.response.model', 'ai.response.finishReason', 'gen_ai.client.operation.time_to_first_chunk', 'ai.usage.inputTokens', 'ai.usage.promptTokens', 'ai.usage.cachedInputTokens', 'ai.usage.inputTokenDetails.cacheReadTokens', 'ai.usage.inputTokenDetails.cacheWriteTokens', 'ai.usage.outputTokens', 'ai.usage.completionTokens', 'ai.usage.reasoningTokens', 'ai.usage.outputTokenDetails.reasoningTokens', 'ai.telemetry.functionId', 'ai.toolCall.name', 'ai.toolCall.id', 'ai.toolCall.args', 'ai.toolCall.result', 'ai.prompt.tools', 'ai.prompt.messages', 'ai.prompt', 'llm.provider', 'llm.system', 'llm.model_name', 'llm.token_count.prompt', 'llm.token_count.prompt_details.cache_read', 'llm.token_count.completion', 'llm.token_count.completion_details.reasoning', 'llm.cost.total', 'tool.name', 'tool.description', 'llm.tools', 'llm.input_messages', 'input.value', 'llm.output_messages', 'output.value', 'openinference.span.kind', 'eve.turn.id', 'maple_ai.turn.id') OR k LIKE 'gen_ai.prompt.variable.%'), SpanAttributes) AS spanAttributes
         FROM trace_detail_spans
         WHERE OrgId = 'org_sql_catalog'
           AND Timestamp >= '2026-01-01 10:30:00'
@@ -558,8 +557,7 @@ SELECT
           StatusCode AS statusCode,
           StatusMessage AS statusMessage,
           toString(Timestamp) AS timestamp,
-          SpanAttributes AS spanAttributes,
-          ResourceAttributes AS resourceAttributes
+          mapFilter((k, v) -> (k IN ('maple_ai.session.id', 'maple_ai.vendor.id', 'maple_ai.vendor.version', 'gen_ai.operation.name', 'gen_ai.provider.name', 'gen_ai.system', 'gen_ai.request.model', 'gen_ai.request.max_tokens', 'gen_ai.request.choice.count', 'gen_ai.request.temperature', 'gen_ai.request.top_p', 'gen_ai.request.top_k', 'gen_ai.request.stop_sequences', 'gen_ai.request.frequency_penalty', 'gen_ai.request.presence_penalty', 'gen_ai.request.encoding_formats', 'gen_ai.request.seed', 'gen_ai.openai.request.seed', 'gen_ai.request.stream', 'gen_ai.request.reasoning.level', 'gen_ai.request.previous_response.id', 'gen_ai.request.stream_cursor', 'gen_ai.response.id', 'gen_ai.response.model', 'gen_ai.response.finish_reasons', 'gen_ai.response.finish_reason', 'gen_ai.response.status', 'gen_ai.response.time_to_first_chunk', 'gen_ai.output.type', 'gen_ai.usage.input_tokens', 'gen_ai.usage.prompt_tokens', 'gen_ai.usage.cache_read.input_tokens', 'gen_ai.usage.input_tokens.cached', 'gen_ai.usage.cache_creation.input_tokens', 'gen_ai.usage.cache_write.input_tokens', 'gen_ai.usage.output_tokens', 'gen_ai.usage.completion_tokens', 'gen_ai.usage.reasoning.output_tokens', 'gen_ai.usage.output_tokens.reasoning', 'gen_ai.usage.cost', 'gen_ai.usage.total_cost', 'gen_ai.conversation.id', 'gen_ai.conversation.compacted', 'gen_ai.agent.id', 'gen_ai.agent.name', 'gen_ai.agent.description', 'gen_ai.agent.version', 'gen_ai.tool.name', 'gen_ai.tool.call.id', 'gen_ai.tool.description', 'gen_ai.tool.type', 'gen_ai.tool.call.arguments', 'gen_ai.tool.call.result', 'gen_ai.tool.definitions', 'gen_ai.system_instructions', 'gen_ai.input.messages', 'gen_ai.prompt', 'gen_ai.output.messages', 'gen_ai.completion', 'gen_ai.data_source.id', 'gen_ai.retrieval.query.text', 'gen_ai.retrieval.top_k', 'gen_ai.retrieval.documents', 'gen_ai.memory.store.id', 'gen_ai.memory.record.id', 'gen_ai.memory.record.count', 'gen_ai.memory.query.text', 'gen_ai.memory.records', 'gen_ai.embeddings.dimension.count', 'gen_ai.evaluation.name', 'gen_ai.evaluation.score.value', 'gen_ai.evaluation.score.label', 'gen_ai.evaluation.explanation', 'gen_ai.prompt.name', 'gen_ai.prompt.version', 'gen_ai.workflow.name', 'error.type', 'server.address', 'server.port', 'ai.model.provider', 'ai.model.id', 'ai.response.id', 'ai.response.model', 'ai.response.finishReason', 'gen_ai.client.operation.time_to_first_chunk', 'ai.usage.inputTokens', 'ai.usage.promptTokens', 'ai.usage.cachedInputTokens', 'ai.usage.inputTokenDetails.cacheReadTokens', 'ai.usage.inputTokenDetails.cacheWriteTokens', 'ai.usage.outputTokens', 'ai.usage.completionTokens', 'ai.usage.reasoningTokens', 'ai.usage.outputTokenDetails.reasoningTokens', 'ai.telemetry.functionId', 'ai.toolCall.name', 'ai.toolCall.id', 'ai.toolCall.args', 'ai.toolCall.result', 'ai.prompt.tools', 'ai.prompt.messages', 'ai.prompt', 'llm.provider', 'llm.system', 'llm.model_name', 'llm.token_count.prompt', 'llm.token_count.prompt_details.cache_read', 'llm.token_count.completion', 'llm.token_count.completion_details.reasoning', 'llm.cost.total', 'tool.name', 'tool.description', 'llm.tools', 'llm.input_messages', 'input.value', 'llm.output_messages', 'output.value', 'openinference.span.kind', 'eve.turn.id', 'maple_ai.turn.id') OR k LIKE 'gen_ai.prompt.variable.%'), SpanAttributes) AS spanAttributes
         FROM trace_detail_spans
         WHERE OrgId = 'org_sql_catalog'
           AND Timestamp >= '2026-01-01 10:30:00'

packages/query-engine-integrations/src/ai/ai-integrations.test.ts

@@ -18,7 +18,6 @@ const row = (
 	statusMessage: "",
 	timestamp: "2026-08-12 15:18:42.207000000",
 	spanAttributes,
-	resourceAttributes: {},
 	...overrides,
 })
 
@@ -351,18 +350,6 @@ describe("span envelope", () => {
 		})
 	})
 
-	it("reads gen_ai keys from span attributes alone", () => {
-		// A resource-level `gen_ai.*` key describes the process, not the
-		// operation: honouring it would stamp every span of that service —
-		// Postgres, HTTP, everything — as an AI span.
-		const mapped = mapAiSpan(
-			row({}, { resourceAttributes: { "gen_ai.request.model": "resource-level" } }),
-		)
-
-		expect(mapped.genAi.requestModel).toBeUndefined()
-		expect(mapped.isAiSpan).toBe(false)
-	})
-
 	it("maps a whole trace's worth of spans in order", () => {
 		const mapped = mapAiSpans([
 			row(INVOKE_AGENT_ATTRS),
@@ -397,15 +384,6 @@ describe("resolveAiIntegration", () => {
 })
 
 describe("untrusted attribute keys", () => {
-	it("ignores a vendor stamp that arrives via a resource attribute", () => {
-		// The envelope is read from span attributes alone, so a resource-level
-		// stamp neither selects an integration nor marks the span.
-		const mapped = mapAiSpan(row({}, { resourceAttributes: { "maple_ai.vendor.id": "eve" } }))
-
-		expect(mapped.vendorId).toBeUndefined()
-		expect(mapped.isAiSpan).toBe(false)
-	})
-
 	it("keeps a prompt variable literally named __proto__ as AI signal", () => {
 		expect(mapAiSpan(row({ "gen_ai.prompt.variable.__proto__": "kept" })).isAiSpan).toBe(true)
 	})

packages/query-engine-integrations/src/ai/ai-integrations.ts

@@ -50,6 +50,12 @@ export interface AiIntegration {
 	 * field from something other than a single attribute.
 	 */
 	readonly refine?: (values: MutableAiGenAiValues, ctx: AiRefineContext) => void
+	/**
+	 * Attribute keys `refine` reads that no source list names. The span read
+	 * projects only the keys the mapper is known to read (`aiSpanAttributeKeys`),
+	 * so a key missing here is a key `refine` never sees.
+	 */
+	readonly refineKeys?: readonly string[]
 }
 
 /** An integration carrying a source list for every catalog field. */
@@ -239,6 +245,26 @@ const resolvedIntegrations = new Map<string, ResolvedAiIntegration>(
 	]),
 )
 
+/**
+ * Every attribute key the mapper can read off a span, across every integration:
+ * the envelope, each field's source keys, and what the refine hooks read. The
+ * span read projects the attribute map down to these (plus the
+ * `AI_PROMPT_VARIABLE_PREFIX` family, which has no fixed key), so a key not in
+ * this list never reaches `mapAiSpan` — in production the map's bulk is
+ * `db.query.text` and friends, which the mapper never looked at.
+ */
+export const aiSpanAttributeKeys: readonly string[] = [
+	...new Set([
+		MAPLE_AI_SESSION_ID_ATTR,
+		MAPLE_AI_VENDOR_ID_ATTR,
+		MAPLE_AI_VENDOR_VERSION_ATTR,
+		...[genAiIntegration, ...resolvedIntegrations.values()].flatMap((integration) =>
+			Object.values(integration.sources).flat(),
+		),
+		...Object.values(AI_VENDOR_INTEGRATIONS).flatMap((vendor) => vendor.refineKeys ?? []),
+	]),
+]
+
 /** The integration for a vendor stamp, or the default for a stamp with no entry. */
 export const resolveAiIntegration = (vendorId: string | undefined): ResolvedAiIntegration => {
 	const resolved = vendorId === undefined ? undefined : resolvedIntegrations.get(vendorId)

packages/query-engine-integrations/src/ai/ai-sessions.test.ts

@@ -813,12 +813,31 @@ describe("aiSessionSpansQuery", () => {
 		expect(sql).toContain("TraceId IN (SELECT")
 		expect(sql).toContain("FROM traces")
 		expect(sql).toContain("Duration / 1000000 AS durationMs")
-		expect(sql).toContain("SpanAttributes AS spanAttributes")
-		expect(sql).toContain("ResourceAttributes AS resourceAttributes")
+		expect(sql).toContain("mapFilter((k, v) -> (k IN ('maple_ai.session.id', ")
+		expect(sql).toContain("OR k LIKE 'gen_ai.prompt.variable.%'), SpanAttributes) AS spanAttributes")
+		expect(sql).not.toContain("ResourceAttributes")
 		expect(sql).toContain("ORDER BY timestamp ASC")
 		expect(sql).toContain("LIMIT 2000")
 	})
 
+	it("projects every key the mapper reads, across vendors", () => {
+		const { sql } = compileUnsafe(aiSessionSpansQuery(), spanParams)
+
+		for (const key of [
+			"maple_ai.vendor.id",
+			"gen_ai.input.messages",
+			"gen_ai.usage.prompt_tokens", // legacy alias
+			"ai.usage.inputTokens", // vercel_ai_sdk
+			"llm.token_count.prompt", // openinference
+			"openinference.span.kind", // read by a refine hook, not a source list
+			"eve.turn.id",
+			"maple_ai.turn.id",
+			"error.type",
+		]) {
+			expect(sql, key).toContain(`'${key}'`)
+		}
+	})
+
 	it("repeats the org predicate on every level that reads a table", () => {
 		const { sql } = compileUnsafe(aiSessionSpansQuery(), spanParams)
 
@@ -878,7 +897,6 @@ describe("aiSessionSpansQuery", () => {
 					"maple_ai.vendor.id": "eve",
 					"maple_ai.session.id": "wrun_01M0CSAEW96BH2W9185XZPRPKH",
 				},
-				resourceAttributes: { "service.name": "maple-slack-agent" },
 			},
 		])
 
@@ -887,9 +905,6 @@ describe("aiSessionSpansQuery", () => {
 			"maple_ai.vendor.id": "eve",
 			"maple_ai.session.id": "wrun_01M0CSAEW96BH2W9185XZPRPKH",
 		})
-		expect(row?.resourceAttributes).toEqual({
-			"service.name": "maple-slack-agent",
-		})
 	})
 })
 
@@ -1033,16 +1048,17 @@ describe("aiTraceSpansQuery", () => {
 		expect(sql).toContain(`TraceId = '${TRACE_ID}'`)
 		expect(sql).not.toContain("TraceId IN (SELECT")
 		expect(sql).not.toContain("FROM traces")
-		expect(sql).not.toContain("maple_ai.session.id")
+		// The projection still names the key; only the predicate is gone.
+		expect(sql).not.toContain("SpanAttributes['maple_ai.session.id']")
 	})
 
 	it("keeps the projection and the order of the session form", () => {
 		const { sql } = compileUnsafe(aiTraceSpansQuery(), traceParams)
 
 		// One shape whichever kind of session the detail page opened.
 		expect(sql).toContain("Duration / 1000000 AS durationMs")
-		expect(sql).toContain("SpanAttributes AS spanAttributes")
-		expect(sql).toContain("ResourceAttributes AS resourceAttributes")
+		expect(sql).toContain("SpanAttributes) AS spanAttributes")
+		expect(sql).not.toContain("ResourceAttributes")
 		expect(sql).toContain("ORDER BY timestamp ASC, spanId ASC")
 		expect(sql).toContain("LIMIT 2000")
 		expect(compileUnsafe(aiTraceSpansQuery({ limit: 100 }), traceParams).sql).toContain("LIMIT 100")
@@ -1092,7 +1108,6 @@ describe("aiTraceSpansQuery", () => {
 				timestamp: "2026-08-19 10:33:25.825000000",
 				// A sessionless vendor: the stamp is there, the session key is not.
 				spanAttributes: { "maple_ai.vendor.id": "llamaindex" },
-				resourceAttributes: { "service.name": "rag-service" },
 			},
 		])