From 67492e79db868aa3d905a1e913ed73dac10cdace Mon Sep 17 00:00:00 2001 From: JeremyFunk Date: Wed, 19 Aug 2026 14:29:27 +0200 Subject: [PATCH 1/4] feat(query-engine): AI agent session queries + gen_ai integration layer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Read side for the AI agent session feature. The ingest gateway already stamps `maple_ai.vendor.id` / `.vendor.version` / `.session.id` at decode time; this is what reads them back. Queries (`src/ai/ai-sessions.ts`) `maple_ai.session.id` is sparse — a vendor stamps it on the spans that own the turn (`ai.eve.turn`, `invoke_agent`), never on the sibling `chat`, `execute_tool` or infrastructure spans. So a session resolves at TRACE granularity: any span carrying the id pulls in every span of its trace. Both queries do that fan-out in two stages against two tables. Detection scans `traces`, where `mapContains(SpanAttributes, …)` rides the `mapKeys(SpanAttributes)` bloom index and stays cheap over a week. The fan-out then reads `trace_detail_spans`, where `TraceId` is a sort-key prefix. The same fan-out on raw `traces` times out at 10s on a 7-day window in production — that table is sorted (OrgId, ServiceName, SpanName, Timestamp) and `idx_trace_id` is only a bloom skip index. No new table and no new index: `errorDetailTracesQuery` already splits across the two for the same reason. `vendorId` resolves via argMin over the earliest session-bearing span, not max(): one trace carries several vendors — an eve agent calling through the Vercel AI SDK — and max() picked `vercel_ai_sdk` alphabetically when `eve` was the framework running the turn. Integration layer (`src/ai/ai-span-model.ts`, `ai-integrations.ts`, `ai-vendors.ts`) A default `gen_ai` integration maps each span onto one standardised format covering all 62 GenAI semconv attributes — every one is stability `development`, so there is no stable subset to draw a line at. Fields are declared once in `AI_GENAI_FIELDS`; the Effect schemas and the source-key table are both generated from it. Per-vendor overrides are keyed on `maple_ai.vendor.id` and replace the default's source keys field by field. Three ship: `vercel_ai_sdk` (`ai.*`, read out of the installed SDK), OpenInference (`llm.*` / `input.value`), and `eve`. Keys that could not be verified against a source were dropped rather than guessed; adding a vendor is one table entry. Attribute values arrive as strings from a ClickHouse Map, so decoding is tolerant by design — a value that will not parse yields `undefined`, never a throw, and a span with no AI signal maps to `isAiSpan: false`. Verified end to end against the production warehouse: both queries return correctly over a 7-day window, and real rows fed through `mapAiSpans` produce the expected spans, including the legacy-dialect normalisations. Co-Authored-By: Claude Opus 5 --- .../query-engine-integrations/package.json | 1 + .../src/__sql_baseline__/integrations.sql | 113 ++++++ .../src/ai/ai-integrations.test.ts | 380 ++++++++++++++++++ .../src/ai/ai-integrations.ts | 323 +++++++++++++++ .../src/ai/ai-sessions.test.ts | 210 ++++++++++ .../src/ai/ai-sessions.ts | 308 ++++++++++++++ .../src/ai/ai-span-model.ts | 361 +++++++++++++++++ .../src/ai/ai-vendors.test.ts | 228 +++++++++++ .../src/ai/ai-vendors.ts | 183 +++++++++ .../query-engine-integrations/src/ai/index.ts | 51 +++ .../query-engine-integrations/src/catalog.ts | 27 ++ .../query-engine-integrations/src/index.ts | 1 + 12 files changed, 2186 insertions(+) create mode 100644 packages/query-engine-integrations/src/ai/ai-integrations.test.ts create mode 100644 packages/query-engine-integrations/src/ai/ai-integrations.ts create mode 100644 packages/query-engine-integrations/src/ai/ai-sessions.test.ts create mode 100644 packages/query-engine-integrations/src/ai/ai-sessions.ts create mode 100644 packages/query-engine-integrations/src/ai/ai-span-model.ts create mode 100644 packages/query-engine-integrations/src/ai/ai-vendors.test.ts create mode 100644 packages/query-engine-integrations/src/ai/ai-vendors.ts create mode 100644 packages/query-engine-integrations/src/ai/index.ts diff --git a/packages/query-engine-integrations/package.json b/packages/query-engine-integrations/package.json index 64d645995..5d646c26b 100644 --- a/packages/query-engine-integrations/package.json +++ b/packages/query-engine-integrations/package.json @@ -4,6 +4,7 @@ "type": "module", "exports": { ".": "./src/index.ts", + "./ai": "./src/ai/index.ts", "./cloudflare": "./src/cloudflare/index.ts", "./planetscale": "./src/planetscale/index.ts", "./product": "./src/product/index.ts", diff --git a/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql b/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql index ffcbfbeca..2c172bc56 100644 --- a/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql +++ b/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql @@ -1,3 +1,116 @@ +-- builder:ai-sessions:aiSessionListQuery:default +SELECT + sessionId AS sessionId, + argMin(vendorId, sessionStart) AS vendorId, + argMin(vendorVersion, sessionStart) AS vendorVersion, + uniq(traceId) AS traceCount, + sum(spanCount) AS spanCount, + sum(errorSpanCount) AS errorSpanCount, + groupUniqArrayArray(serviceNames) AS serviceNames, + toString(min(traceStart)) AS startTime, + toString(max(traceEnd)) AS endTime, + intDiv(toUnixTimestamp64Nano(max(traceEnd)) - toUnixTimestamp64Nano(min(traceStart)), 1000000) AS durationMs + FROM (SELECT + TraceId AS traceId, + max(SpanAttributes['maple_ai.session.id']) AS sessionId, + argMin(SpanAttributes['maple_ai.vendor.id'], if((mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != ''), Timestamp, toDateTime('2106-01-01 00:00:00'))) AS vendorId, + argMin(SpanAttributes['maple_ai.vendor.version'], if((mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != ''), Timestamp, toDateTime('2106-01-01 00:00:00'))) AS vendorVersion, + min(if((mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != ''), Timestamp, toDateTime('2106-01-01 00:00:00'))) AS sessionStart, + count() AS spanCount, + countIf(StatusCode = 'Error') AS errorSpanCount, + groupUniqArray(ServiceName) AS serviceNames, + min(Timestamp) AS traceStart, + max(Timestamp) AS traceEnd + FROM trace_detail_spans + WHERE OrgId = 'org_sql_catalog' + AND Timestamp >= '2026-01-01 10:30:00' + AND Timestamp <= '2026-01-03 14:15:00' + AND TraceId IN (SELECT + TraceId AS TraceId + FROM traces + WHERE OrgId = 'org_sql_catalog' + AND Timestamp >= '2026-01-01 10:30:00' + AND Timestamp <= '2026-01-03 14:15:00' + AND (mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != '')) + GROUP BY traceId) AS session_traces + WHERE sessionId != '' + GROUP BY sessionId + ORDER BY startTime DESC + LIMIT 50 + FORMAT JSON + +-- builder:ai-sessions:aiSessionListQuery:filtered +SELECT + sessionId AS sessionId, + argMin(vendorId, sessionStart) AS vendorId, + argMin(vendorVersion, sessionStart) AS vendorVersion, + uniq(traceId) AS traceCount, + sum(spanCount) AS spanCount, + sum(errorSpanCount) AS errorSpanCount, + groupUniqArrayArray(serviceNames) AS serviceNames, + toString(min(traceStart)) AS startTime, + toString(max(traceEnd)) AS endTime, + intDiv(toUnixTimestamp64Nano(max(traceEnd)) - toUnixTimestamp64Nano(min(traceStart)), 1000000) AS durationMs + FROM (SELECT + TraceId AS traceId, + max(SpanAttributes['maple_ai.session.id']) AS sessionId, + argMin(SpanAttributes['maple_ai.vendor.id'], if((mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != ''), Timestamp, toDateTime('2106-01-01 00:00:00'))) AS vendorId, + argMin(SpanAttributes['maple_ai.vendor.version'], if((mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != ''), Timestamp, toDateTime('2106-01-01 00:00:00'))) AS vendorVersion, + min(if((mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != ''), Timestamp, toDateTime('2106-01-01 00:00:00'))) AS sessionStart, + count() AS spanCount, + countIf(StatusCode = 'Error') AS errorSpanCount, + groupUniqArray(ServiceName) AS serviceNames, + min(Timestamp) AS traceStart, + max(Timestamp) AS traceEnd + FROM trace_detail_spans + WHERE OrgId = 'org_sql_catalog' + AND Timestamp >= '2026-01-01 10:30:00' + AND Timestamp <= '2026-01-03 14:15:00' + AND TraceId IN (SELECT + TraceId AS TraceId + FROM traces + WHERE OrgId = 'org_sql_catalog' + AND Timestamp >= '2026-01-01 10:30:00' + AND Timestamp <= '2026-01-03 14:15:00' + AND (mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != '') + AND SpanAttributes['maple_ai.vendor.id'] IN ('eve') + AND ServiceName IN ('maple-slack-agent')) + GROUP BY traceId) AS session_traces + WHERE sessionId != '' + GROUP BY sessionId + ORDER BY startTime DESC + LIMIT 25 + FORMAT JSON + +-- builder:ai-sessions:aiSessionSpansQuery:default +SELECT + TraceId AS traceId, + SpanId AS spanId, + ParentSpanId AS parentSpanId, + SpanName AS spanName, + SpanKind AS spanKind, + ServiceName AS serviceName, + Duration / 1000000 AS durationMs, + StatusCode AS statusCode, + StatusMessage AS statusMessage, + toString(Timestamp) AS timestamp, + SpanAttributes AS spanAttributes, + ResourceAttributes AS resourceAttributes + FROM trace_detail_spans + WHERE OrgId = 'org_sql_catalog' + AND Timestamp >= '2026-01-01 10:30:00' + AND Timestamp <= '2026-01-03 14:15:00' + AND TraceId IN (SELECT + TraceId AS TraceId + FROM traces + WHERE OrgId = 'org_sql_catalog' + AND Timestamp >= '2026-01-01 10:30:00' + AND Timestamp <= '2026-01-03 14:15:00' + AND SpanAttributes['maple_ai.session.id'] = 'wrun_sql_catalog') + ORDER BY timestamp ASC + LIMIT 2000 + FORMAT JSON + -- builder:cloudflare-infra-breakdowns:cloudflareZoneBreakdownTimeseriesSQL:default SELECT formatDateTime(toStartOfInterval(TimeUnix, INTERVAL 300 SECOND), '%Y-%m-%dT%H:%i:%S.%fZ') AS bucket, diff --git a/packages/query-engine-integrations/src/ai/ai-integrations.test.ts b/packages/query-engine-integrations/src/ai/ai-integrations.test.ts new file mode 100644 index 000000000..0ddc53603 --- /dev/null +++ b/packages/query-engine-integrations/src/ai/ai-integrations.test.ts @@ -0,0 +1,380 @@ +import { describe, expect, it } from "vitest" +import { genAiIntegration, mapAiSpan, mapAiSpans, resolveAiIntegration } from "./ai-integrations" +import { AI_GENAI_FIELDS, type AiGenAiField, type AiSessionSpanRow } from "./ai-span-model" + +const row = ( + spanAttributes: Record, + overrides: Partial = {}, +): AiSessionSpanRow => ({ + traceId: "d31eaf1d98a9b26028dfe521f8dbc75c", + spanId: "00233d43ea0d1598", + parentSpanId: "68fc42c0c9f2cf15", + spanName: "invoke_agent openai/gpt-5.6-luna", + spanKind: "Internal", + serviceName: "maple-slack-agent", + durationMs: 2995.573199, + statusCode: "Unset", + statusMessage: "", + timestamp: "2026-08-12 15:18:42.207000000", + spanAttributes, + resourceAttributes: {}, + ...overrides, +}) + +/** A real `invoke_agent` span from this org, canonical `gen_ai.*` throughout. */ +const INVOKE_AGENT_ATTRS = { + "ai.settings.context.eve.session.id": "wrun_01KZAAFFZRHHRYC8MY9MDANASQ", + "ai.settings.context.eve.turn.id": "turn_0", + "gen_ai.agent.name": "slack-agent", + "gen_ai.input.messages": '[{"role":"user","parts":[{"type":"text","content":"hello"}]}]', + "gen_ai.operation.name": "invoke_agent", + "gen_ai.output.messages": '[{"role":"assistant","parts":[{"type":"text","content":"hi"}]}]', + "gen_ai.provider.name": "openrouter", + "gen_ai.request.model": "openai/gpt-5.6-luna", + "gen_ai.response.finish_reasons": '["stop"]', + "gen_ai.system_instructions": '[{"type":"text","content":"You are Maple AI"}]', + "gen_ai.usage.cache_creation.input_tokens": "106", + "gen_ai.usage.cache_read.input_tokens": "4924", + "gen_ai.usage.input_tokens": "5033", + "gen_ai.usage.output_tokens": "38", + "maple_ai.session.id": "wrun_01KZAAFFZRHHRYC8MY9MDANASQ", + "maple_ai.vendor.id": "vercel_ai_sdk", + "maple_ai.vendor.version": "0", +} + +/** A real `execute_tool` span — the tool fields, including two JSON blobs. */ +const EXECUTE_TOOL_ATTRS = { + "gen_ai.execute_tool.duration": "0.5546019470000029", + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.call.arguments": '{"emoji":"wave"}', + "gen_ai.tool.call.id": "call_uKgzomwVJhP3bYZ0fxvwUe86", + "gen_ai.tool.call.result": '{"reacted":true,"emoji":"wave"}', + "gen_ai.tool.name": "add_reaction", + "gen_ai.tool.type": "function", +} + +/** + * A real `workflow.stream.flush` span. Every key is present with an empty value + * because ClickHouse returns `''` for a missing `Map` key — the exact shape that + * would look like "present" to a naive reader. + */ +const NON_AI_ATTRS = { + "cache.name": "", + "db.system": "", + "http.request.method": "", + "http.response.status_code": "", + "server.address": "", + "url.full": "", +} + +const SAMPLE_ATTRIBUTE_VALUE = { + string: "sample", + number: "42", + boolean: "true", + stringArray: '["a"]', + json: '{"a":1}', +} as const + +const SAMPLE_DECODED_VALUE = { + string: "sample", + number: 42, + boolean: true, + stringArray: ["a"], + json: { a: 1 }, +} as const + +describe("the catalog is the contract", () => { + // Driven from `AI_GENAI_FIELDS` itself rather than a hand-written list, so a + // field added to the catalog without a source mapping fails here instead of + // silently returning `undefined` in the UI forever. + for (const [field, def] of Object.entries(AI_GENAI_FIELDS)) { + it(`maps the canonical key ${def.key} to ${field}`, () => { + const mapped = mapAiSpan(row({ [def.key]: SAMPLE_ATTRIBUTE_VALUE[def.type] })) + + expect(mapped.genAi[field as AiGenAiField]).toEqual(SAMPLE_DECODED_VALUE[def.type]) + }) + } + + it("declares a source list for every catalog field", () => { + expect(Object.keys(genAiIntegration.sources).sort()).toEqual(Object.keys(AI_GENAI_FIELDS).sort()) + }) +}) + +describe("value decoding", () => { + it("decodes numbers out of the Map(String, String) wire format", () => { + const mapped = mapAiSpan(row(INVOKE_AGENT_ATTRS)) + + expect(mapped.genAi.usageInputTokens).toBe(5033) + expect(mapped.genAi.usageCacheReadInputTokens).toBe(4924) + }) + + it("rejects a number that is not finite instead of poisoning arithmetic", () => { + // A dashboard that sums NaN token counts shows NaN for the whole session, + // which is strictly worse than showing nothing for one span. + expect( + mapAiSpan(row({ "gen_ai.usage.input_tokens": "not-a-number" })).genAi.usageInputTokens, + ).toBeUndefined() + expect( + mapAiSpan(row({ "gen_ai.usage.input_tokens": "Infinity" })).genAi.usageInputTokens, + ).toBeUndefined() + }) + + it("accepts both the word and the digit form of a boolean", () => { + expect(mapAiSpan(row({ "gen_ai.request.stream": "true" })).genAi.requestStream).toBe(true) + expect(mapAiSpan(row({ "gen_ai.request.stream": "1" })).genAi.requestStream).toBe(true) + expect(mapAiSpan(row({ "gen_ai.request.stream": "false" })).genAi.requestStream).toBe(false) + expect(mapAiSpan(row({ "gen_ai.request.stream": "0" })).genAi.requestStream).toBe(false) + expect(mapAiSpan(row({ "gen_ai.request.stream": "yes" })).genAi.requestStream).toBeUndefined() + }) + + it("accepts both shapes real instrumentation emits for a string array", () => { + // Production carries both for `gen_ai.response.finish_reasons`: the + // serialised array and the bare single value. + expect( + mapAiSpan(row({ "gen_ai.response.finish_reasons": '["stop"]' })).genAi.responseFinishReasons, + ).toEqual(["stop"]) + expect( + mapAiSpan(row({ "gen_ai.response.finish_reasons": "stop" })).genAi.responseFinishReasons, + ).toEqual(["stop"]) + }) + + it("decodes JSON blobs into structured values", () => { + const mapped = mapAiSpan(row(EXECUTE_TOOL_ATTRS)) + + expect(mapped.genAi.toolCallArguments).toEqual({ emoji: "wave" }) + expect(mapped.genAi.toolCallResult).toEqual({ reacted: true, emoji: "wave" }) + expect(mapped.genAi.toolName).toBe("add_reaction") + expect(mapped.genAi.toolCallId).toBe("call_uKgzomwVJhP3bYZ0fxvwUe86") + }) + + it("yields no field for malformed JSON rather than throwing", () => { + // One truncated attribute must not cost the caller the rest of the span. + const mapped = mapAiSpan( + row({ "gen_ai.tool.call.arguments": '{"emoji":', "gen_ai.tool.name": "add_reaction" }), + ) + + expect(mapped.genAi.toolCallArguments).toBeUndefined() + expect(mapped.genAi.toolName).toBe("add_reaction") + }) + + it("treats an empty value as absent, because that is what a missing Map key returns", () => { + expect(mapAiSpan(row({ "gen_ai.request.model": "" })).genAi.requestModel).toBeUndefined() + }) + + it("falls through to the next alias when the first key does not decode", () => { + const mapped = mapAiSpan( + row({ "gen_ai.usage.input_tokens": "n/a", "gen_ai.usage.prompt_tokens": "5033" }), + ) + + expect(mapped.genAi.usageInputTokens).toBe(5033) + }) +}) + +describe("legacy aliases", () => { + // One case per row of the semconv deprecation table. + const cases: ReadonlyArray = [ + ["gen_ai.usage.prompt_tokens", "5033", "usageInputTokens", 5033], + ["gen_ai.usage.completion_tokens", "38", "usageOutputTokens", 38], + ["gen_ai.prompt", '[{"role":"user"}]', "inputMessages", [{ role: "user" }]], + ["gen_ai.completion", '[{"role":"assistant"}]', "outputMessages", [{ role: "assistant" }]], + ["gen_ai.system", "anthropic", "providerName", "anthropic"], + ["gen_ai.openai.request.seed", "7", "requestSeed", 7], + ["gen_ai.openai.request.response_format", "json_object", "outputType", "json_object"], + ["gen_ai.response.finish_reason", "stop", "responseFinishReasons", ["stop"]], + // Not in the deprecation table: the sub-key spelling OpenRouter actually + // emits. Both confirmed present in the warehouse — unlike the plausible + // `gen_ai.usage.reasoning_tokens`, which is not, and so is not mapped. + ["gen_ai.usage.output_tokens.reasoning", "704", "usageReasoningOutputTokens", 704], + ["gen_ai.usage.input_tokens.cached", "2048", "usageCacheReadInputTokens", 2048], + ] + + for (const [key, value, field, expected] of cases) { + it(`reads the deprecated ${key} into ${field}`, () => { + expect(mapAiSpan(row({ [key]: value })).genAi[field]).toEqual(expected) + }) + } + + it("prefers the canonical key when both are present", () => { + const mapped = mapAiSpan( + row({ "gen_ai.usage.input_tokens": "5033", "gen_ai.usage.prompt_tokens": "1" }), + ) + + expect(mapped.genAi.usageInputTokens).toBe(5033) + }) + + it("maps a real legacy OpenRouter span through the aliases alone", () => { + // The dialect an OpenRouter/traceloop-style instrumentor still emits: + // prompt/completion/system, and the singular finish reason. + const mapped = mapAiSpan( + row({ + "gen_ai.system": "openai", + "gen_ai.prompt": '[{"role":"user","content":"hi"}]', + "gen_ai.completion": '[{"role":"assistant","content":"hello"}]', + "gen_ai.response.finish_reason": "stop", + "gen_ai.usage.prompt_tokens": "12", + "gen_ai.usage.completion_tokens": "4", + }), + ) + + expect(mapped.genAi).toEqual({ + providerName: "openai", + inputMessages: [{ role: "user", content: "hi" }], + outputMessages: [{ role: "assistant", content: "hello" }], + responseFinishReasons: ["stop"], + usageInputTokens: 12, + usageOutputTokens: 4, + }) + expect(mapped.isAiSpan).toBe(true) + }) +}) + +describe("value normalisation", () => { + // The `gen_ai.system` enum members that were renamed with the attribute. + const renames: ReadonlyArray = [ + ["vertex_ai", "gcp.vertex_ai"], + ["gemini", "gcp.gemini"], + ["az.ai.inference", "azure.ai.inference"], + ["az.ai.openai", "azure.ai.openai"], + ["xai", "x_ai"], + ] + + for (const [legacy, canonical] of renames) { + it(`rewrites the legacy provider value ${legacy}`, () => { + expect(mapAiSpan(row({ "gen_ai.system": legacy })).genAi.providerName).toBe(canonical) + }) + } + + it("leaves a legacy value alone when it survived the rename", () => { + expect(mapAiSpan(row({ "gen_ai.system": "anthropic" })).genAi.providerName).toBe("anthropic") + }) + + it("never rewrites a value that arrived on the canonical key", () => { + // A span emitting `gen_ai.provider.name` already speaks the new + // vocabulary; a collision with an old enum member is its value to keep. + const mapped = mapAiSpan(row({ "gen_ai.provider.name": "gemini", "gen_ai.system": "vertex_ai" })) + + expect(mapped.genAi.providerName).toBe("gemini") + }) + + it("singularises the old tool_calls finish reason", () => { + const mapped = mapAiSpan(row({ "gen_ai.response.finish_reasons": '["tool_calls","stop"]' })) + + expect(mapped.genAi.responseFinishReasons).toEqual(["tool_call", "stop"]) + }) +}) + +describe("prompt variables", () => { + it("collects the templated gen_ai.prompt.variable. attributes by prefix", () => { + // Templated: the variable name is IN the key, so there is no single key + // the source-list mechanism could look up. + const mapped = mapAiSpan( + row({ + "gen_ai.prompt.name": "triage", + "gen_ai.prompt.variable.service": "maple-api", + "gen_ai.prompt.variable.window": "24h", + }), + ) + + expect(mapped.promptVariables).toEqual({ service: "maple-api", window: "24h" }) + expect(mapped.genAi.promptName).toBe("triage") + }) + + it("omits promptVariables entirely when none are present", () => { + expect(mapAiSpan(row(EXECUTE_TOOL_ATTRS)).promptVariables).toBeUndefined() + }) +}) + +describe("non-AI spans", () => { + it("maps an ordinary infrastructure span to a clean not-an-AI-span result", () => { + const mapped = mapAiSpan(row(NON_AI_ATTRS, { spanName: "workflow.stream.flush", spanKind: "Client" })) + + expect(mapped.isAiSpan).toBe(false) + expect(mapped.genAi).toEqual({}) + expect(mapped.integrationId).toBe("gen_ai") + expect(mapped.vendorId).toBeUndefined() + }) + + it("does not treat a core semconv attribute as AI signal", () => { + // `server.address` is on every HTTP client span in the trace. It is worth + // surfacing next to the AI fields, but it cannot be what decides that a + // span is an AI span. + const mapped = mapAiSpan(row({ "server.address": "openrouter.ai", "server.port": "443" })) + + expect(mapped.genAi.serverAddress).toBe("openrouter.ai") + expect(mapped.isAiSpan).toBe(false) + }) + + it("counts the gateway stamp alone as AI signal", () => { + // The gateway saw evidence the read path cannot (scope, resource SDK + // name, span events), so its stamp outranks the absence of gen_ai keys. + const mapped = mapAiSpan(row({ "maple_ai.vendor.id": "eve", "maple_ai.vendor.version": "0" })) + + expect(mapped.isAiSpan).toBe(true) + }) +}) + +describe("span envelope", () => { + it("carries the warehouse columns and the gateway stamp through untouched", () => { + const mapped = mapAiSpan(row(INVOKE_AGENT_ATTRS)) + + expect(mapped).toMatchObject({ + traceId: "d31eaf1d98a9b26028dfe521f8dbc75c", + spanId: "00233d43ea0d1598", + parentSpanId: "68fc42c0c9f2cf15", + spanName: "invoke_agent openai/gpt-5.6-luna", + spanKind: "Internal", + serviceName: "maple-slack-agent", + timestamp: "2026-08-12 15:18:42.207000000", + durationMs: 2995.573199, + statusCode: "Unset", + statusMessage: "", + sessionId: "wrun_01KZAAFFZRHHRYC8MY9MDANASQ", + vendorId: "vercel_ai_sdk", + vendorVersion: "0", + isAiSpan: true, + }) + }) + + it("lets a span attribute win over a resource attribute of the same name", () => { + const mapped = mapAiSpan( + row( + { "gen_ai.request.model": "openai/gpt-5.6-luna" }, + { resourceAttributes: { "gen_ai.request.model": "stale" } }, + ), + ) + + expect(mapped.genAi.requestModel).toBe("openai/gpt-5.6-luna") + }) + + it("maps a whole trace's worth of spans in order", () => { + const mapped = mapAiSpans([ + row(INVOKE_AGENT_ATTRS), + row(NON_AI_ATTRS, { spanName: "workflow.stream.flush" }), + ]) + + expect(mapped.map((span) => span.isAiSpan)).toEqual([true, false]) + }) +}) + +describe("resolveAiIntegration", () => { + it("falls back to the default integration for a vendor with no override", () => { + // `unknown:other` is a real gateway stamp: AI-shaped span, no recognised + // framework. It must map through the default, not fail. + const integration = resolveAiIntegration("unknown:other") + + expect(integration).toBe(genAiIntegration) + expect( + mapAiSpan(row({ ...INVOKE_AGENT_ATTRS, "maple_ai.vendor.id": "unknown:other" })).integrationId, + ).toBe("gen_ai") + }) + + it("falls back to the default integration for an unstamped span", () => { + expect(resolveAiIntegration(undefined)).toBe(genAiIntegration) + }) + + it("returns the same merged integration on every call for a vendor", () => { + // The merge is memoised per vendor id; re-merging sixty source lists on + // every span of every session is pure waste. + expect(resolveAiIntegration("vercel_ai_sdk")).toBe(resolveAiIntegration("vercel_ai_sdk")) + }) +}) diff --git a/packages/query-engine-integrations/src/ai/ai-integrations.ts b/packages/query-engine-integrations/src/ai/ai-integrations.ts new file mode 100644 index 000000000..b0fac82b5 --- /dev/null +++ b/packages/query-engine-integrations/src/ai/ai-integrations.ts @@ -0,0 +1,323 @@ +// AI span mapping: the default GenAI integration, the vendor registry lookup, +// and the mapper that turns one warehouse row into an `AiAgentSpan`. +// +// The mechanism is deliberately small: an integration is a table of +// `field -> source attribute keys` plus an optional `refine` hook, and a vendor +// integration is the same table applied ON TOP of the default one, per field. +// That covers the two shapes real instrumentation takes — "same meaning, +// different key" (the table) and "same key, different value" (the hook) — +// without a plugin system. +// +// Why keyed on `maple_ai.vendor.id` rather than re-detected here: the ingest +// gateway already ran the detection at decode time over evidence the read path +// no longer has (instrumentation scope, resource SDK name, span events — see +// `SCREEN_KEYS` and the vendor table in `apps/ingest/src/ai_session.rs`) and +// stamped its verdict on the span. Re-deriving a dialect from what survives +// into ClickHouse would be a weaker second copy of that decision, free to +// disagree with the stamp. +// +// Everything here is tolerant by construction. Attributes arrive as +// `Map(String, String)`, so a missing key reads back as `''` and every value is +// a string that may or may not be the shape its field expects. A value that +// fails to decode yields no field — never a throw — because one badly +// serialised attribute must not cost the user the rest of the span. + +import { + AI_GENAI_FIELDS, + AI_PROMPT_VARIABLE_PREFIX, + MAPLE_AI_SESSION_ID_ATTR, + MAPLE_AI_VENDOR_ID_ATTR, + MAPLE_AI_VENDOR_VERSION_ATTR, + type AiAgentSpan, + type AiFieldDef, + type AiGenAiField, + type AiSessionSpanRow, + type MutableAiGenAiValues, +} from "./ai-span-model" +import { AI_VENDOR_INTEGRATIONS } from "./ai-vendors" + +export interface AiRefineContext { + readonly row: AiSessionSpanRow + /** + * Span attributes merged OVER resource attributes — a span-level key wins, + * because resource attributes describe the process, not the operation. Both + * are readable here; the source key lists read from the same merged view. + */ + readonly attributes: Record +} + +export interface AiIntegration { + readonly id: string + /** + * Field → source attribute keys, tried in order; the first key with a value + * that decodes wins. A vendor entry REPLACES the default entry for that + * field, so a vendor that wants the canonical key to keep priority lists it + * first itself. + */ + readonly sources: Partial> + /** + * For what a key list cannot express: normalising a value, or deriving a + * field from something other than a single attribute. + */ + readonly refine?: (values: MutableAiGenAiValues, ctx: AiRefineContext) => void +} + +/** ClickHouse returns `''` for a missing Map key, so empty means absent. */ +const readAttribute = (attributes: Record, key: string): string | undefined => { + const value: string | undefined = attributes[key] + return value === undefined || value === "" ? undefined : value +} + +/** + * What a JSON-typed attribute can decode to. Spelling it out rather than + * returning `unknown` keeps the decoder's contract honest: the value came off a + * `JSON.parse`, so it is JSON, not anything at all. + */ +export type AiJsonValue = + | string + | number + | boolean + | null + | readonly AiJsonValue[] + | { readonly [key: string]: AiJsonValue } + +/** + * Every shape `decodeAttribute` can produce. It is exactly `AiJsonValue`: the + * four scalar field types are JSON scalars and `stringArray` is a JSON array, + * so naming them again would only restate the union. + */ +export type AiDecodedValue = AiJsonValue + +const parseJson = (raw: string): AiJsonValue | undefined => { + try { + return JSON.parse(raw) + } catch { + return undefined + } +} + +const decodeStringArray = (raw: string): readonly string[] => { + // Real data carries both shapes for the same attribute: `'["stop"]'` from + // instrumentation that serialises the array, and a bare `"stop"` from + // instrumentation that emits the single value. Anything that is not a JSON + // array of strings is treated as the bare form rather than discarded. + const parsed = parseJson(raw) + return Array.isArray(parsed) && parsed.every((entry) => typeof entry === "string") ? parsed : [raw] +} + +const decodeAttribute = (type: AiFieldDef["type"], raw: string): AiDecodedValue | undefined => { + switch (type) { + case "string": + return raw + case "number": { + const value = Number(raw) + // `Number("abc")` is NaN and `Number("Infinity")` is Infinity; both + // would poison arithmetic and neither is a token count. + return Number.isFinite(value) ? value : undefined + } + case "boolean": { + const value = raw.toLowerCase() + if (value === "true" || value === "1") return true + if (value === "false" || value === "0") return false + return undefined + } + case "stringArray": + return decodeStringArray(raw) + case "json": + return parseJson(raw) + } +} + +// The catalog correlates each field with its value type, but a loop over the +// union of fields cannot carry that correlation without re-stating the whole +// table. The single unchecked write lives here; `decodeAttribute` is driven by +// the same catalog entry, so the value is right by construction. +const assign = (values: MutableAiGenAiValues, field: AiGenAiField, value: AiDecodedValue): void => { + ;(values as Record)[field] = value +} + +/** + * Deprecated and obsoleted keys the default integration still reads, per field. + * + * `gen_ai.prompt` / `gen_ai.completion` were obsoleted with "no replacement" + * rather than renamed, so mapping them onto the message fields is a pragmatic + * choice, not a spec-blessed rename: in practice they carried exactly that + * content, and OpenRouter instrumentation in production still emits them. + * + * `gen_ai.usage.output_tokens.reasoning` and `gen_ai.usage.input_tokens.cached` + * are likewise absent from the deprecation table — they are the sub-key spelling + * OpenRouter emits in production for what the convention calls + * `gen_ai.usage.reasoning.output_tokens` and `gen_ai.usage.cache_read.input_tokens`. + * Both were confirmed against the warehouse; the plausible-looking + * `gen_ai.usage.reasoning_tokens` was NOT, so it is deliberately not listed. + */ +const GENAI_LEGACY_ALIASES = { + usageInputTokens: ["gen_ai.usage.prompt_tokens"], + usageOutputTokens: ["gen_ai.usage.completion_tokens"], + usageReasoningOutputTokens: ["gen_ai.usage.output_tokens.reasoning"], + usageCacheReadInputTokens: ["gen_ai.usage.input_tokens.cached"], + inputMessages: ["gen_ai.prompt"], + outputMessages: ["gen_ai.completion"], + providerName: ["gen_ai.system"], + requestSeed: ["gen_ai.openai.request.seed"], + outputType: ["gen_ai.openai.request.response_format"], + responseFinishReasons: ["gen_ai.response.finish_reason"], +} satisfies Partial> + +const genAiSources: Partial> = {} +for (const [field, def] of Object.entries(AI_GENAI_FIELDS)) { + genAiSources[field as AiGenAiField] = [def.key] +} +for (const [field, aliases] of Object.entries(GENAI_LEGACY_ALIASES)) { + genAiSources[field as AiGenAiField] = [...(genAiSources[field as AiGenAiField] ?? []), ...aliases] +} + +/** + * `gen_ai.system` enum values that were renamed when the attribute became + * `gen_ai.provider.name`. Values not listed here (`openai`, `anthropic`, …) + * survived the rename unchanged. + */ +const LEGACY_SYSTEM_VALUES = new Map([ + ["vertex_ai", "gcp.vertex_ai"], + ["gemini", "gcp.gemini"], + ["az.ai.inference", "azure.ai.inference"], + ["az.ai.openai", "azure.ai.openai"], + ["xai", "x_ai"], +]) + +const genAiRefine = (values: MutableAiGenAiValues, ctx: AiRefineContext): void => { + // Only a value that actually came from the legacy key gets rewritten. A span + // that emits `gen_ai.provider.name` is already speaking the new vocabulary, + // and its values must be passed through even when they collide with an old + // enum member. + if ( + values.providerName !== undefined && + readAttribute(ctx.attributes, AI_GENAI_FIELDS.providerName.key) === undefined + ) { + const canonical = LEGACY_SYSTEM_VALUES.get(values.providerName) + if (canonical !== undefined) values.providerName = canonical + } + + // The finish reason was singularised in place, so this is a value fix rather + // than a key alias and applies whichever key it arrived on. + if (values.responseFinishReasons !== undefined) { + values.responseFinishReasons = values.responseFinishReasons.map((reason) => + reason === "tool_calls" ? "tool_call" : reason, + ) + } +} + +/** + * The default integration: canonical GenAI keys plus their legacy aliases. It + * is what an unrecognised vendor — and the `unknown:*` buckets the gateway + * stamps — falls back to, and it is the base every vendor override merges onto. + */ +export const genAiIntegration: AiIntegration = { + id: "gen_ai", + sources: genAiSources, + refine: genAiRefine, +} + +const resolvedIntegrations = new Map() + +/** + * The integration for a vendor stamp: the default one, with the vendor's key + * lists replacing the default's per field and both `refine` hooks running — + * default first, so the vendor can correct its work. + */ +export const resolveAiIntegration = (vendorId: string | undefined): AiIntegration => { + if (vendorId === undefined) return genAiIntegration + const vendor = AI_VENDOR_INTEGRATIONS[vendorId] + if (vendor === undefined) return genAiIntegration + const cached = resolvedIntegrations.get(vendorId) + if (cached !== undefined) return cached + const merged: AiIntegration = { + id: vendor.id, + sources: { ...genAiIntegration.sources, ...vendor.sources }, + refine: (values, ctx) => { + genAiIntegration.refine?.(values, ctx) + vendor.refine?.(values, ctx) + }, + } + resolvedIntegrations.set(vendorId, merged) + return merged +} + +const collectPromptVariables = (attributes: Record): Record | undefined => { + // A templated attribute has no fixed key, so it is collected by prefix here + // rather than through the key-list mechanism every other field uses. + let collected: Record | undefined + for (const [key, value] of Object.entries(attributes)) { + if (!key.startsWith(AI_PROMPT_VARIABLE_PREFIX) || value === "") continue + collected ??= {} + collected[key.slice(AI_PROMPT_VARIABLE_PREFIX.length)] = value + } + return collected +} + +/** + * A `core` field is plain OTel semconv that every HTTP client span carries, so + * mapping one is not evidence that this span is an AI span. + */ +const hasAiSignal = (values: MutableAiGenAiValues): boolean => + Object.keys(values).some((field) => AI_GENAI_FIELDS[field as AiGenAiField].group !== "core") + +interface AiSpanOptionalFields { + sessionId?: string + vendorId?: string + vendorVersion?: string + promptVariables?: Record +} + +export const mapAiSpan = (row: AiSessionSpanRow): AiAgentSpan => { + const attributes = { ...row.resourceAttributes, ...row.spanAttributes } + const vendorId = readAttribute(attributes, MAPLE_AI_VENDOR_ID_ATTR) + const integration = resolveAiIntegration(vendorId) + + const genAi: MutableAiGenAiValues = {} + for (const [field, keys] of Object.entries(integration.sources)) { + if (keys === undefined) continue + for (const key of keys) { + const raw = readAttribute(attributes, key) + if (raw === undefined) continue + const value = decodeAttribute(AI_GENAI_FIELDS[field as AiGenAiField].type, raw) + // A key that carries an undecodable value does not consume the + // field: the next alias still gets its turn. + if (value === undefined) continue + assign(genAi, field as AiGenAiField, value) + break + } + } + integration.refine?.(genAi, { row, attributes }) + + const promptVariables = collectPromptVariables(attributes) + const sessionId = readAttribute(attributes, MAPLE_AI_SESSION_ID_ATTR) + const vendorVersion = readAttribute(attributes, MAPLE_AI_VENDOR_VERSION_ATTR) + + // Collected separately so an absent stamp leaves the key off the span + // entirely rather than present-and-undefined. + const optional: AiSpanOptionalFields = {} + if (sessionId !== undefined) optional.sessionId = sessionId + if (vendorId !== undefined) optional.vendorId = vendorId + if (vendorVersion !== undefined) optional.vendorVersion = vendorVersion + if (promptVariables !== undefined) optional.promptVariables = promptVariables + + return { + traceId: row.traceId, + spanId: row.spanId, + parentSpanId: row.parentSpanId, + spanName: row.spanName, + spanKind: row.spanKind, + serviceName: row.serviceName, + timestamp: row.timestamp, + durationMs: row.durationMs, + statusCode: row.statusCode, + statusMessage: row.statusMessage, + ...optional, + integrationId: integration.id, + isAiSpan: vendorId !== undefined || promptVariables !== undefined || hasAiSignal(genAi), + genAi, + } +} + +export const mapAiSpans = (rows: readonly AiSessionSpanRow[]): readonly AiAgentSpan[] => rows.map(mapAiSpan) diff --git a/packages/query-engine-integrations/src/ai/ai-sessions.test.ts b/packages/query-engine-integrations/src/ai/ai-sessions.test.ts new file mode 100644 index 000000000..3fa2056b4 --- /dev/null +++ b/packages/query-engine-integrations/src/ai/ai-sessions.test.ts @@ -0,0 +1,210 @@ +import { describe, expect, it } from "vitest" +import { Effect } from "effect" +import { compileCH, type CompiledQuery } from "@maple-dev/clickhouse-builder" +import { + aiSessionListQuery, + aiSessionListRowSchema, + aiSessionSpansQuery, + aiSessionSpansRowSchema, +} from "./ai-sessions" + +const params = { + orgId: "org_1", + startTime: "2026-08-18 00:00:00", + endTime: "2026-08-19 23:59:59", +} + +const spanParams = { ...params, sessionId: "wrun_01M0CSAEW96BH2W9185XZPRPKH" } + +const decodeRows = (compiled: CompiledQuery, rows: ReadonlyArray>) => + Effect.runSync(compiled.decodeRows(rows)) + +/** `OrgId = 'x'` on the detection level AND on the fan-out level. */ +const orgPredicateCount = (sql: string) => sql.split("OrgId = 'org_1'").length - 1 + +describe("aiSessionListQuery", () => { + it("detects sessions on traces, then fans out over trace_detail_spans", () => { + const { sql } = compileCH(aiSessionListQuery(), params) + + // The detection level is the one the SpanAttributes bloom index serves; + // the fan-out reads the MV whose sort key starts (OrgId, TraceId). + expect(sql).toContain("FROM trace_detail_spans") + expect(sql).toContain("TraceId IN (SELECT") + expect(sql).toContain("FROM traces") + expect(sql).toContain("GROUP BY traceId") + expect(sql).toContain("GROUP BY sessionId") + expect(sql).toContain("ORDER BY startTime DESC") + expect(sql).toContain("LIMIT 50") + }) + + it("repeats the org predicate on every level that reads a table", () => { + const { sql } = compileCH(aiSessionListQuery(), params) + + expect(orgPredicateCount(sql)).toBe(2) + }) + + it("is org-scoped", () => { + expect(compileCH(aiSessionListQuery(), params).tenantScope).toBe("org") + }) + + it("tests session-id presence with mapContains AND a non-empty value", () => { + const { sql } = compileCH(aiSessionListQuery(), params) + + // ClickHouse yields '' for a missing Map key, so mapContains alone would + // admit spans carrying an empty session id. + expect(sql).toContain( + "(mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != '')", + ) + }) + + it("resolves the vendor from the earliest session-bearing span, not max()", () => { + const { sql } = compileCH(aiSessionListQuery(), params) + + // max(vendorId) picked `vercel_ai_sdk` alphabetically over the `eve` that + // actually ran the turn — see the builder's doc comment. + expect(sql).not.toContain("max(SpanAttributes['maple_ai.vendor.id'])") + expect(sql).toContain("argMin(SpanAttributes['maple_ai.vendor.id'], if(") + expect(sql).toContain("argMin(SpanAttributes['maple_ai.vendor.version'], if(") + expect(sql).toContain("argMin(vendorId, sessionStart) AS vendorId") + expect(sql).toContain("argMin(vendorVersion, sessionStart) AS vendorVersion") + // The sentinel must stay inside DateTime's range or toDateTime won't parse. + expect(sql).toContain("toDateTime('2106-01-01 00:00:00')") + }) + + it("escapes an org id carrying a quote", () => { + const { sql } = compileCH(aiSessionListQuery(), { ...params, orgId: "org'evil" }) + + expect(sql).toContain("OrgId = 'org\\'evil'") + }) + + it("omits the optional filters when none are given", () => { + const { sql } = compileCH(aiSessionListQuery(), params) + + expect(sql).not.toContain("SpanAttributes['maple_ai.vendor.id'] IN") + expect(sql).not.toContain("ServiceName IN") + }) + + it("puts both optional filters on the detection level only", () => { + const { sql } = compileCH( + aiSessionListQuery({ limit: 25, vendorIds: ["eve"], serviceNames: ["maple-slack-agent"] }), + params, + ) + + // Filtering the fan-out instead would drop spans and under-count spanCount. + const [fanOut, detection] = sql.split("TraceId IN (SELECT") + expect(detection).toContain("SpanAttributes['maple_ai.vendor.id'] IN ('eve')") + expect(detection).toContain("ServiceName IN ('maple-slack-agent')") + expect(fanOut).not.toContain("IN ('eve')") + expect(sql).toContain("LIMIT 25") + }) + + it("leaves no unresolved param placeholder", () => { + expect(compileCH(aiSessionListQuery(), params).sql).not.toContain("__PARAM_") + }) + + it("decodes quoted 64-bit aggregates and the service-name array", () => { + const compiled = compileCH(aiSessionListQuery(), params, { rowSchema: aiSessionListRowSchema }) + + const [row] = decodeRows(compiled, [ + { + sessionId: "wrun_01M0CSAEW96BH2W9185XZPRPKH", + vendorId: "eve", + vendorVersion: "0", + traceCount: "1", + spanCount: "250", + errorSpanCount: "4", + serviceNames: ["maple-slack-agent", "maple-api"], + startTime: "2026-08-19 10:33:25.825000000", + endTime: "2026-08-19 10:33:36.242000000", + durationMs: "10417", + }, + ]) + + expect(row).toEqual({ + sessionId: "wrun_01M0CSAEW96BH2W9185XZPRPKH", + vendorId: "eve", + vendorVersion: "0", + traceCount: 1, + spanCount: 250, + errorSpanCount: 4, + serviceNames: ["maple-slack-agent", "maple-api"], + startTime: "2026-08-19 10:33:25.825000000", + endTime: "2026-08-19 10:33:36.242000000", + durationMs: 10_417, + }) + }) +}) + +describe("aiSessionSpansQuery", () => { + it("returns every span of every trace in the session, oldest first", () => { + const { sql } = compileCH(aiSessionSpansQuery(), spanParams) + + expect(sql).toContain("FROM trace_detail_spans") + 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("ORDER BY timestamp ASC") + expect(sql).toContain("LIMIT 2000") + }) + + it("repeats the org predicate on every level that reads a table", () => { + const { sql } = compileCH(aiSessionSpansQuery(), spanParams) + + expect(orgPredicateCount(sql)).toBe(2) + }) + + it("is org-scoped", () => { + expect(compileCH(aiSessionSpansQuery(), spanParams).tenantScope).toBe("org") + }) + + it("substitutes and escapes the sessionId param", () => { + const { sql } = compileCH(aiSessionSpansQuery(), spanParams) + expect(sql).toContain("SpanAttributes['maple_ai.session.id'] = 'wrun_01M0CSAEW96BH2W9185XZPRPKH'") + + const escaped = compileCH(aiSessionSpansQuery(), { ...spanParams, sessionId: "sess'evil" }) + expect(escaped.sql).toContain("SpanAttributes['maple_ai.session.id'] = 'sess\\'evil'") + }) + + it("honours a caller-supplied limit", () => { + expect(compileCH(aiSessionSpansQuery({ limit: 100 }), spanParams).sql).toContain("LIMIT 100") + }) + + it("leaves no unresolved param placeholder", () => { + expect(compileCH(aiSessionSpansQuery(), spanParams).sql).not.toContain("__PARAM_") + }) + + it("decodes the raw Map columns as plain objects", () => { + const compiled = compileCH(aiSessionSpansQuery(), spanParams, { + rowSchema: aiSessionSpansRowSchema, + }) + + const [row] = decodeRows(compiled, [ + { + traceId: "6b0c0e0a", + spanId: "aa11", + parentSpanId: "", + spanName: "ai.eve.turn", + spanKind: "Internal", + serviceName: "maple-slack-agent", + durationMs: "250", + statusCode: "Ok", + statusMessage: "", + timestamp: "2026-08-19 10:33:25.825000000", + spanAttributes: { + "maple_ai.vendor.id": "eve", + "maple_ai.session.id": "wrun_01M0CSAEW96BH2W9185XZPRPKH", + }, + resourceAttributes: { "service.name": "maple-slack-agent" }, + }, + ]) + + expect(row?.durationMs).toBe(250) + expect(row?.spanAttributes).toEqual({ + "maple_ai.vendor.id": "eve", + "maple_ai.session.id": "wrun_01M0CSAEW96BH2W9185XZPRPKH", + }) + expect(row?.resourceAttributes).toEqual({ "service.name": "maple-slack-agent" }) + }) +}) diff --git a/packages/query-engine-integrations/src/ai/ai-sessions.ts b/packages/query-engine-integrations/src/ai/ai-sessions.ts new file mode 100644 index 000000000..d5cca112f --- /dev/null +++ b/packages/query-engine-integrations/src/ai/ai-sessions.ts @@ -0,0 +1,308 @@ +// AI agent sessions — read side +// +// The ingest gateway stamps three attributes on AI-agent spans at decode time +// (`apps/ingest/src/ai_session.rs`): `maple_ai.vendor.id`, +// `maple_ai.vendor.version` and `maple_ai.session.id`. Only the last one is +// sparse — a vendor exposes a session key on the spans that own the turn +// (`ai.eve.turn`, `invoke_agent`), never on the sibling `chat`, `execute_tool`, +// `workflow.*` or HTTP-client spans it fans out to. +// +// So a session is resolved at TRACE granularity: a trace belongs to a session +// if ANY of its spans carries that session id, and then EVERY span of that +// trace is part of the session — including the completely non-AI ones. That is +// deliberate; the dashboard shows the full agent context, not just the spans +// the framework happened to label. +// +// Both queries are that fan-out, in two stages against two different tables: +// +// detect — `traces`, filtered on the presence of `maple_ai.session.id`. This +// is the only level that can use the `mapKeys(SpanAttributes)` bloom skip +// index, and with it the scan stays cheap over a week. It yields the +// qualifying trace-id set and nothing else. +// fan out — `trace_detail_spans`, restricted by `TraceId IN (…)`. `TraceId` is +// a sort-key prefix there (`(OrgId, TraceId, SpanId)`), so this is a seek. +// The same fan-out against raw `traces` times out at 10s on a 7-day window +// in production: that table is sorted `(OrgId, ServiceName, SpanName, +// Timestamp)` and `idx_trace_id` is only a bloom skip index, which prunes +// far too little at this org's volume. +// +// That is the "optimise the query, not the storage" answer to the fan-out: no +// new table, no new index, an MV that already exists and that +// `errorDetailTracesQuery` already splits across for exactly this reason. `IN` +// rather than a JOIN for the same reason too — ClickHouse pushes the id set into +// the read, which a JOIN does not do. +// +// The window predicate stays on BOTH levels. On `trace_detail_spans` it prunes +// partitions (`PARTITION BY toDate(Timestamp)`) as well as riding the sort key, +// so it is strictly cheaper than omitting it. +// +// Tenant scoping: a subquery contributes nothing to the outer query's scope, so +// every level that reads a table repeats `OrgId = {orgId}` itself. The outermost +// level of `aiSessionListQuery` reads a derived table rather than a table, and +// inherits `org` scope from it. + +import { Schema } from "effect" +import * as CH from "@maple-dev/clickhouse-builder/expr" +import { + from, + fromQuery, + inSubquery, + param, + type CompiledQueryRowSchema, +} from "@maple-dev/clickhouse-builder" +import { TraceDetailSpans, Traces } from "@maple/query-engine/ch/tables" +import { CHNumber } from "@maple/query-engine/ch/schema" + +const SESSION_ID_ATTR = "maple_ai.session.id" +const VENDOR_ID_ATTR = "maple_ai.vendor.id" +const VENDOR_VERSION_ATTR = "maple_ai.vendor.version" + +/** + * Sorts every span that is NOT session-bearing behind every one that is, so a + * single `argMin`/`min` over it picks the earliest session-bearing span without + * needing an `argMinIf` the DSL does not have. Wrapped in `toDateTime` rather + * than left a bare string, or `if()` would have to reconcile DateTime64(9) with + * String. That wrapper is also why the sentinel is 2106 and not 3000: `DateTime` + * tops out at 2106-02-07 and anything past it fails to parse. + */ +const SESSION_ORDER_SENTINEL = "2106-01-01 00:00:00" + +/** ClickHouse returns `''` for a missing Map key, so presence needs both halves. */ +const hasSessionId = (attrs: CH.Expr>, get: CH.Expr) => + CH.mapContains(attrs, SESSION_ID_ATTR).and(get.neq("")) + +export interface AiSessionListOpts { + /** Sessions returned, most recently started first. */ + readonly limit?: number + readonly vendorIds?: readonly string[] + readonly serviceNames?: readonly string[] +} + +export interface AiSessionListOutput { + readonly sessionId: string + /** Vendor of the earliest session-bearing span — see `aiSessionListQuery`. */ + readonly vendorId: string + readonly vendorVersion: string + readonly traceCount: number + readonly spanCount: number + readonly errorSpanCount: number + readonly serviceNames: readonly string[] + /** ClickHouse datetime literal, e.g. `2026-08-19 10:33:25.825000000`. */ + readonly startTime: string + readonly endTime: string + readonly durationMs: number +} + +export const aiSessionListRowSchema: CompiledQueryRowSchema = Schema.Struct({ + sessionId: Schema.String, + vendorId: Schema.String, + vendorVersion: Schema.String, + traceCount: CHNumber, + spanCount: CHNumber, + errorSpanCount: CHNumber, + serviceNames: Schema.Array(Schema.String), + startTime: Schema.String, + endTime: Schema.String, + durationMs: CHNumber, +}) + +/** + * One row per AI agent session in the window. + * + * `vendorId` is the vendor of the EARLIEST span that carries a session id, not + * `max(vendorId)`. A single trace legitimately carries several vendors — an eve + * agent calls through the Vercel AI SDK — and `max` picked `vercel_ai_sdk` + * alphabetically over `eve` when `eve` was the framework actually running the + * turn. The root-most session-bearing span is the one that names the framework, + * so the two `argMin`s (per trace, then across traces) resolve to it. + * + * The vendor filter goes on the detection subquery: it is the level the bloom + * index serves, and it is the only place `maple_ai.vendor.id` is unambiguous — + * a trace's other spans carry other vendors, or none. + * + * The service filter goes there too, which means "the session-bearing spans came + * from this service" rather than "the trace touched this service". A trace spans + * services by definition, so the alternative — filtering the fan-out — would + * silently drop spans and under-count `spanCount`. The session-bearing spans come + * from the agent's own service, which is the one a user filtering by service means. + */ +export function aiSessionListQuery(opts: AiSessionListOpts = {}) { + const limit = opts.limit ?? 50 + + const sessionTraceIds = from(Traces) + .select(($) => ({ TraceId: $.TraceId })) + .where(($) => [ + $.OrgId.eq(param.string("orgId")), + $.Timestamp.gte(param.dateTime("startTime")), + $.Timestamp.lte(param.dateTime("endTime")), + hasSessionId($.SpanAttributes, $.SpanAttributes.get(SESSION_ID_ATTR)), + opts.vendorIds?.length + ? CH.inList($.SpanAttributes.get(VENDOR_ID_ATTR), opts.vendorIds) + : undefined, + opts.serviceNames?.length ? CH.inList($.ServiceName, opts.serviceNames) : undefined, + ]) + + // Per trace: every span of a qualifying trace, session-bearing or not. + const perTrace = from(TraceDetailSpans) + .select(($) => { + const sessionOrder = CH.if_( + hasSessionId($.SpanAttributes, $.SpanAttributes.get(SESSION_ID_ATTR)), + $.Timestamp, + CH.toDateTime(CH.lit(SESSION_ORDER_SENTINEL)), + ) + return { + traceId: $.TraceId, + // A trace belongs to one session; the non-bearing spans read `''`, + // which `max` discards. + sessionId: CH.max_($.SpanAttributes.get(SESSION_ID_ATTR)), + vendorId: CH.argMin($.SpanAttributes.get(VENDOR_ID_ATTR), sessionOrder), + vendorVersion: CH.argMin($.SpanAttributes.get(VENDOR_VERSION_ATTR), sessionOrder), + // Carried so the outer level can order traces by their first + // session-bearing span rather than by their first span of any kind. + sessionStart: CH.min_(sessionOrder), + spanCount: CH.count(), + errorSpanCount: CH.countIf($.StatusCode.eq("Error")), + serviceNames: CH.groupUniqArray($.ServiceName), + // Named apart from the outer `startTime`/`endTime` on purpose: an + // outer alias shadows the derived table's column of the same name, + // so `min(startTime)` would resolve to the outer `toString(…)` String + // and `toUnixTimestamp64Nano` reject it — verified against production, + // it fails with ILLEGAL_TYPE_OF_ARGUMENT. + traceStart: CH.min_($.Timestamp), + traceEnd: CH.max_($.Timestamp), + } + }) + .where(($) => [ + $.OrgId.eq(param.string("orgId")), + $.Timestamp.gte(param.dateTime("startTime")), + $.Timestamp.lte(param.dateTime("endTime")), + inSubquery($.TraceId, sessionTraceIds), + ]) + .groupBy("traceId") + + return ( + fromQuery(perTrace, "session_traces") + .select(($) => ({ + sessionId: $.sessionId, + vendorId: CH.argMin($.vendorId, $.sessionStart), + vendorVersion: CH.argMin($.vendorVersion, $.sessionStart), + traceCount: CH.uniq($.traceId), + spanCount: CH.sum($.spanCount), + errorSpanCount: CH.sum($.errorSpanCount), + serviceNames: CH.groupUniqArrayArray($.serviceNames), + startTime: CH.toString_(CH.min_($.traceStart)), + endTime: CH.toString_(CH.max_($.traceEnd)), + // Nanoseconds first: `Timestamp` is DateTime64(9), and subtracting two + // of them yields a Decimal whose scale the wire format then quotes. + // Wrapped in `intDiv` because `Expr.sub`/`div` do not parenthesize. + durationMs: CH.intDiv( + CH.toUnixTimestamp64Nano(CH.max_($.traceEnd)).sub( + CH.toUnixTimestamp64Nano(CH.min_($.traceStart)), + ), + 1_000_000, + ), + })) + // A trace can qualify via the IN and still roll up empty if its only + // session-bearing span fell outside the window. + .where(($) => [$.sessionId.neq("")]) + .groupBy("sessionId") + .orderBy(["startTime", "desc"]) + .limit(limit) + .format("JSON") + ) +} + +export interface AiSessionSpansOpts { + readonly limit?: number +} + +export interface AiSessionSpansOutput { + readonly traceId: string + readonly spanId: string + readonly parentSpanId: string + readonly spanName: string + readonly spanKind: string + readonly serviceName: string + readonly durationMs: number + readonly statusCode: string + readonly statusMessage: string + readonly timestamp: string + readonly spanAttributes: Record + readonly resourceAttributes: Record +} + +export const aiSessionSpansRowSchema: CompiledQueryRowSchema = Schema.Struct({ + traceId: Schema.String, + spanId: Schema.String, + parentSpanId: Schema.String, + spanName: Schema.String, + spanKind: Schema.String, + serviceName: Schema.String, + durationMs: CHNumber, + statusCode: Schema.String, + statusMessage: Schema.String, + timestamp: Schema.String, + // A Map column selected directly arrives as a JSON object under FORMAT JSON, + // so this is a plain Record. Not `Schema.fromJsonString(…)` — that is for the + // observability path, which reads maps already serialized to a string. + spanAttributes: Schema.Record(Schema.String, Schema.String), + resourceAttributes: Schema.Record(Schema.String, Schema.String), +}) + +/** + * Every span of every trace belonging to one session, oldest first. + * + * `sessionId` is a compile param rather than an opts field, so one compiled SQL + * string serves every session. + * + * Both attribute Maps come back whole: the integration layer that normalizes + * these into gen_ai form needs keys this query cannot know in advance. Projecting + * only the keys it wants is a later optimisation, and a real one — one production + * trace already carries 250 spans with up to ~17KB of attributes each, so callers + * should expect megabyte-scale payloads at the default limit. + * + * No scope columns: `trace_detail_spans` does not carry `ScopeName`/`ScopeVersion`, + * and the read path does not need them. The ingest gateway already did the + * scope-based vendor detection at write time and encoded its verdict in + * `maple_ai.vendor.id`; re-deriving the dialect here would only second-guess it. + * + * Known v1 limitation: the time window bounds BOTH levels, so a session whose + * traces straddle the window edge returns only the spans inside it. + */ +export function aiSessionSpansQuery(opts: AiSessionSpansOpts = {}) { + const limit = opts.limit ?? 2000 + + const sessionTraceIds = from(Traces) + .select(($) => ({ TraceId: $.TraceId })) + .where(($) => [ + $.OrgId.eq(param.string("orgId")), + $.Timestamp.gte(param.dateTime("startTime")), + $.Timestamp.lte(param.dateTime("endTime")), + $.SpanAttributes.get(SESSION_ID_ATTR).eq(param.string("sessionId")), + ]) + + return from(TraceDetailSpans) + .select(($) => ({ + traceId: $.TraceId, + spanId: $.SpanId, + parentSpanId: $.ParentSpanId, + spanName: $.SpanName, + spanKind: $.SpanKind, + serviceName: $.ServiceName, + durationMs: $.Duration.div(1_000_000), + statusCode: $.StatusCode, + statusMessage: $.StatusMessage, + timestamp: CH.toString_($.Timestamp), + spanAttributes: $.SpanAttributes, + resourceAttributes: $.ResourceAttributes, + })) + .where(($) => [ + $.OrgId.eq(param.string("orgId")), + $.Timestamp.gte(param.dateTime("startTime")), + $.Timestamp.lte(param.dateTime("endTime")), + inSubquery($.TraceId, sessionTraceIds), + ]) + .orderBy(["timestamp", "asc"]) + .limit(limit) + .format("JSON") +} diff --git a/packages/query-engine-integrations/src/ai/ai-span-model.ts b/packages/query-engine-integrations/src/ai/ai-span-model.ts new file mode 100644 index 000000000..b6e7078b5 --- /dev/null +++ b/packages/query-engine-integrations/src/ai/ai-span-model.ts @@ -0,0 +1,361 @@ +// The standardised Maple AI-agent span — field catalog, value types, schema. +// +// Every AI framework speaks a different attribute dialect, but the dashboard +// wants one shape. This module owns that shape: a flat catalog of every OTel +// GenAI semantic-convention attribute (`AI_GENAI_FIELDS`), the value type each +// one decodes to, and the `AiAgentSpan` an integration produces. The mapping +// itself lives in `ai-integrations.ts` — this file is data plus types, so the +// catalog stays the single source of truth for field names, canonical keys and +// value types, and nothing can add a field without also declaring how it +// decodes. +// +// Two facts drive the design: +// +// 1. Every `gen_ai.*` attribute is stability `development` — there are no +// stable ones. So there is no "stable subset" to ship first; the catalog +// carries the whole convention, in-development fields included, and the +// cost of a field the instrumentation never emits is one absent key. +// 2. Warehouse attributes arrive as `Map(String, String)`. `1234`, `true` and +// `["stop"]` are all strings on the wire, and a missing key reads back as +// `''`, not as absent. The `type` tag on each field is what tells the +// decoder how to turn that string back into a value. +// +// There is deliberately NO instrumentation-scope input here. The ingest gateway +// already did scope-based vendor detection at write time +// (`apps/ingest/src/ai_session.rs`) and encoded its verdict in +// `maple_ai.vendor.id`; re-deriving a dialect from the scope on read would be a +// second, weaker copy of that decision that could disagree with the stamp. + +import { Schema } from "effect" + +/** Vendor slug the ingest gateway stamped on the span. */ +export const MAPLE_AI_VENDOR_ID_ATTR = "maple_ai.vendor.id" +/** Version of the vendor detection that produced the stamp, currently `"0"`. */ +export const MAPLE_AI_VENDOR_VERSION_ATTR = "maple_ai.vendor.version" +/** The vendor's own session id, verbatim. */ +export const MAPLE_AI_SESSION_ID_ATTR = "maple_ai.session.id" + +/** + * One span as `aiSessionSpansQuery` returns it. Declared here rather than + * imported from the query module so the mapping layer depends on a shape, not + * on a query — the query satisfies this structurally. + */ +export interface AiSessionSpanRow { + readonly traceId: string + readonly spanId: string + readonly parentSpanId: string + readonly spanName: string + readonly spanKind: string + readonly serviceName: string + readonly durationMs: number + readonly statusCode: string + readonly statusMessage: string + readonly timestamp: string + readonly spanAttributes: Record + readonly resourceAttributes: Record +} + +/** + * Semconv group a field belongs to, for grouping in the UI. + * + * `core` is the odd one out: `error.type`, `server.address` and `server.port` + * are plain core-semconv attributes that AI spans happen to carry, not AI + * signal. Every ordinary HTTP client span in the trace has them too, which is + * why the mapper refuses to treat a `core` field as evidence that a span is an + * AI span. + */ +export type AiFieldGroup = + | "operation" + | "request" + | "response" + | "usage" + | "conversation" + | "agent" + | "tool" + | "content" + | "dataSource" + | "retrieval" + | "memory" + | "embeddings" + | "evaluation" + | "prompt" + | "workflow" + | "core" + +export interface AiFieldDef { + /** Canonical OTel semconv attribute key. */ + readonly key: string + readonly type: "string" | "number" | "boolean" | "stringArray" | "json" + readonly group: AiFieldGroup +} + +/** + * Every GenAI semconv attribute, keyed by a camelCase field name that mirrors + * the semconv path. `key` is the CANONICAL key only — legacy aliases are an + * integration concern and live in `ai-integrations.ts`, because an alias is a + * statement about instrumentation, not about the convention. + */ +export const AI_GENAI_FIELDS = { + // operation + operationName: { key: "gen_ai.operation.name", type: "string", group: "operation" }, + providerName: { key: "gen_ai.provider.name", type: "string", group: "operation" }, + + // request + requestModel: { key: "gen_ai.request.model", type: "string", group: "request" }, + requestMaxTokens: { key: "gen_ai.request.max_tokens", type: "number", group: "request" }, + requestChoiceCount: { key: "gen_ai.request.choice.count", type: "number", group: "request" }, + requestTemperature: { key: "gen_ai.request.temperature", type: "number", group: "request" }, + requestTopP: { key: "gen_ai.request.top_p", type: "number", group: "request" }, + requestTopK: { key: "gen_ai.request.top_k", type: "number", group: "request" }, + requestStopSequences: { key: "gen_ai.request.stop_sequences", type: "stringArray", group: "request" }, + requestFrequencyPenalty: { key: "gen_ai.request.frequency_penalty", type: "number", group: "request" }, + requestPresencePenalty: { key: "gen_ai.request.presence_penalty", type: "number", group: "request" }, + requestEncodingFormats: { key: "gen_ai.request.encoding_formats", type: "stringArray", group: "request" }, + requestSeed: { key: "gen_ai.request.seed", type: "number", group: "request" }, + requestStream: { key: "gen_ai.request.stream", type: "boolean", group: "request" }, + requestReasoningLevel: { key: "gen_ai.request.reasoning.level", type: "string", group: "request" }, + requestPreviousResponseId: { + key: "gen_ai.request.previous_response.id", + type: "string", + group: "request", + }, + requestStreamCursor: { key: "gen_ai.request.stream_cursor", type: "string", group: "request" }, + + // response + responseId: { key: "gen_ai.response.id", type: "string", group: "response" }, + responseModel: { key: "gen_ai.response.model", type: "string", group: "response" }, + responseFinishReasons: { key: "gen_ai.response.finish_reasons", type: "stringArray", group: "response" }, + responseStatus: { key: "gen_ai.response.status", type: "string", group: "response" }, + responseTimeToFirstChunk: { + key: "gen_ai.response.time_to_first_chunk", + type: "number", + group: "response", + }, + outputType: { key: "gen_ai.output.type", type: "string", group: "response" }, + + // usage + usageInputTokens: { key: "gen_ai.usage.input_tokens", type: "number", group: "usage" }, + usageCacheReadInputTokens: { + key: "gen_ai.usage.cache_read.input_tokens", + type: "number", + group: "usage", + }, + usageCacheCreationInputTokens: { + key: "gen_ai.usage.cache_creation.input_tokens", + type: "number", + group: "usage", + }, + usageOutputTokens: { key: "gen_ai.usage.output_tokens", type: "number", group: "usage" }, + usageReasoningOutputTokens: { + key: "gen_ai.usage.reasoning.output_tokens", + type: "number", + group: "usage", + }, + + // conversation + conversationId: { key: "gen_ai.conversation.id", type: "string", group: "conversation" }, + conversationCompacted: { key: "gen_ai.conversation.compacted", type: "boolean", group: "conversation" }, + + // agent + agentId: { key: "gen_ai.agent.id", type: "string", group: "agent" }, + agentName: { key: "gen_ai.agent.name", type: "string", group: "agent" }, + agentDescription: { key: "gen_ai.agent.description", type: "string", group: "agent" }, + agentVersion: { key: "gen_ai.agent.version", type: "string", group: "agent" }, + + // tool + toolName: { key: "gen_ai.tool.name", type: "string", group: "tool" }, + toolCallId: { key: "gen_ai.tool.call.id", type: "string", group: "tool" }, + toolDescription: { key: "gen_ai.tool.description", type: "string", group: "tool" }, + toolType: { key: "gen_ai.tool.type", type: "string", group: "tool" }, + toolCallArguments: { key: "gen_ai.tool.call.arguments", type: "json", group: "tool" }, + toolCallResult: { key: "gen_ai.tool.call.result", type: "json", group: "tool" }, + toolDefinitions: { key: "gen_ai.tool.definitions", type: "json", group: "tool" }, + + // content + systemInstructions: { key: "gen_ai.system_instructions", type: "json", group: "content" }, + inputMessages: { key: "gen_ai.input.messages", type: "json", group: "content" }, + outputMessages: { key: "gen_ai.output.messages", type: "json", group: "content" }, + + // data source / retrieval + dataSourceId: { key: "gen_ai.data_source.id", type: "string", group: "dataSource" }, + retrievalQueryText: { key: "gen_ai.retrieval.query.text", type: "string", group: "retrieval" }, + retrievalTopK: { key: "gen_ai.retrieval.top_k", type: "number", group: "retrieval" }, + retrievalDocuments: { key: "gen_ai.retrieval.documents", type: "json", group: "retrieval" }, + + // memory + memoryStoreId: { key: "gen_ai.memory.store.id", type: "string", group: "memory" }, + memoryRecordId: { key: "gen_ai.memory.record.id", type: "string", group: "memory" }, + memoryRecordCount: { key: "gen_ai.memory.record.count", type: "number", group: "memory" }, + memoryQueryText: { key: "gen_ai.memory.query.text", type: "string", group: "memory" }, + memoryRecords: { key: "gen_ai.memory.records", type: "json", group: "memory" }, + + // embeddings + embeddingsDimensionCount: { + key: "gen_ai.embeddings.dimension.count", + type: "number", + group: "embeddings", + }, + + // evaluation + evaluationName: { key: "gen_ai.evaluation.name", type: "string", group: "evaluation" }, + evaluationScoreValue: { key: "gen_ai.evaluation.score.value", type: "number", group: "evaluation" }, + evaluationScoreLabel: { key: "gen_ai.evaluation.score.label", type: "string", group: "evaluation" }, + evaluationExplanation: { key: "gen_ai.evaluation.explanation", type: "string", group: "evaluation" }, + + // prompt + promptName: { key: "gen_ai.prompt.name", type: "string", group: "prompt" }, + promptVersion: { key: "gen_ai.prompt.version", type: "string", group: "prompt" }, + + // workflow + workflowName: { key: "gen_ai.workflow.name", type: "string", group: "workflow" }, + + // core semconv attributes AI spans carry — see `AiFieldGroup` + errorType: { key: "error.type", type: "string", group: "core" }, + serverAddress: { key: "server.address", type: "string", group: "core" }, + serverPort: { key: "server.port", type: "number", group: "core" }, +} as const satisfies Record + +export type AiGenAiField = keyof typeof AI_GENAI_FIELDS + +/** + * `gen_ai.prompt.variable.` is a TEMPLATED attribute: the key carries the + * variable name, so there is no single key to look up and it cannot live in + * `AI_GENAI_FIELDS` alongside the fixed keys. The mapper collects it by prefix + * into `AiAgentSpan.promptVariables` instead. + */ +export const AI_PROMPT_VARIABLE_PREFIX = "gen_ai.prompt.variable." + +/** + * The documented values of `gen_ai.operation.name`. + * + * NOT a closed enum, and `operationName` is deliberately a plain string: the + * convention explicitly allows system-specific names, and production data + * already carries `agent_step` (Vercel AI SDK). This list exists so the UI can + * group and label the known operations — never to reject an unknown one. + */ +export const AI_KNOWN_OPERATION_NAMES = [ + "chat", + "generate_content", + "text_completion", + "embeddings", + "retrieval", + "fetch_response", + "create_agent", + "invoke_agent", + "execute_tool", + "invoke_workflow", + "plan", + "search_memory", + "create_memory", + "update_memory", + "upsert_memory", + "delete_memory", + "create_memory_store", + "delete_memory_store", +] as const + +export type AiKnownOperationName = (typeof AI_KNOWN_OPERATION_NAMES)[number] + +/** Value a field of the given `type` decodes to. */ +export type AiFieldValue = T extends "string" + ? string + : T extends "number" + ? number + : T extends "boolean" + ? boolean + : T extends "stringArray" + ? readonly string[] + : unknown + +/** Every catalog field, optional, typed from its `type` tag. */ +export type AiGenAiValues = { + readonly [F in AiGenAiField]?: AiFieldValue<(typeof AI_GENAI_FIELDS)[F]["type"]> +} + +/** The same, writable — what an integration's `refine` hook mutates. */ +export type MutableAiGenAiValues = { + -readonly [F in AiGenAiField]?: AiFieldValue<(typeof AI_GENAI_FIELDS)[F]["type"]> +} + +export interface AiAgentSpan { + readonly traceId: string + readonly spanId: string + readonly parentSpanId: string + readonly spanName: string + readonly spanKind: string + readonly serviceName: string + readonly timestamp: string + readonly durationMs: number + readonly statusCode: string + readonly statusMessage: string + /** Maple AI envelope, stamped by the ingest gateway. */ + readonly sessionId?: string + readonly vendorId?: string + readonly vendorVersion?: string + /** Which integration produced `genAi`. */ + readonly integrationId: string + /** + * True when the span carried any recognised AI signal. False for the + * ordinary infrastructure spans that share an agent trace — they are + * returned rather than dropped, because the session view shows the whole + * agent context. + */ + readonly isAiSpan: boolean + readonly genAi: AiGenAiValues + readonly promptVariables?: Record +} + +/** Schema per `type` tag, so the catalog also generates the schema. */ +interface AiFieldValueSchemas { + readonly string: Schema.String + readonly number: Schema.Finite + readonly boolean: Schema.Boolean + readonly stringArray: Schema.$Array + readonly json: Schema.Unknown +} + +const aiFieldValueSchemas: AiFieldValueSchemas = { + string: Schema.String, + number: Schema.Finite, + boolean: Schema.Boolean, + stringArray: Schema.Array(Schema.String), + json: Schema.Unknown, +} + +type AiGenAiFieldSchemas = { + readonly [F in AiGenAiField]: Schema.optionalKey +} + +// Generated from the catalog rather than written out: a hand-written struct of +// sixty optional fields is a second source of truth that silently drifts the +// first time someone adds a field. `Object.fromEntries` erases the key/value +// correlation, so the assertion re-states what the mapped type above already +// spells out. +const aiGenAiFieldSchemas = Object.fromEntries( + Object.entries(AI_GENAI_FIELDS).map(([field, def]) => [ + field, + Schema.optionalKey(aiFieldValueSchemas[def.type]), + ]), +) as AiGenAiFieldSchemas + +export const AiGenAiValuesSchema = Schema.Struct(aiGenAiFieldSchemas) + +export const AiAgentSpanSchema = Schema.Struct({ + traceId: Schema.String, + spanId: Schema.String, + parentSpanId: Schema.String, + spanName: Schema.String, + spanKind: Schema.String, + serviceName: Schema.String, + timestamp: Schema.String, + durationMs: Schema.Finite, + statusCode: Schema.String, + statusMessage: Schema.String, + sessionId: Schema.optionalKey(Schema.String), + vendorId: Schema.optionalKey(Schema.String), + vendorVersion: Schema.optionalKey(Schema.String), + integrationId: Schema.String, + isAiSpan: Schema.Boolean, + genAi: AiGenAiValuesSchema, + promptVariables: Schema.optionalKey(Schema.Record(Schema.String, Schema.String)), +}) diff --git a/packages/query-engine-integrations/src/ai/ai-vendors.test.ts b/packages/query-engine-integrations/src/ai/ai-vendors.test.ts new file mode 100644 index 000000000..585733838 --- /dev/null +++ b/packages/query-engine-integrations/src/ai/ai-vendors.test.ts @@ -0,0 +1,228 @@ +import { describe, expect, it } from "vitest" +import { mapAiSpan, resolveAiIntegration } from "./ai-integrations" +import { AI_VENDOR_INTEGRATIONS } from "./ai-vendors" +import type { AiSessionSpanRow } from "./ai-span-model" + +const row = (vendorId: string, spanAttributes: Record): AiSessionSpanRow => ({ + traceId: "a1e33a5cc671a33952cc0e1117701290", + spanId: "a2d0b69ed027b25b", + parentSpanId: "2bafd07bbfb0eeb3", + spanName: "ai.eve.turn", + spanKind: "Internal", + serviceName: "maple-slack-agent", + durationMs: 3632.565327, + statusCode: "Unset", + statusMessage: "", + timestamp: "2026-08-12 15:19:41.626000000", + spanAttributes: { ...spanAttributes, "maple_ai.vendor.id": vendorId, "maple_ai.vendor.version": "0" }, + resourceAttributes: {}, +}) + +describe("vercel_ai_sdk", () => { + it("reads the older ai.* usage keys the default integration knows nothing about", () => { + const mapped = mapAiSpan( + row("vercel_ai_sdk", { + "ai.usage.promptTokens": "5033", + "ai.usage.completionTokens": "38", + "ai.model.id": "openai/gpt-5.6-luna", + "ai.model.provider": "openrouter", + "ai.response.finishReason": "stop", + }), + ) + + expect(mapped.genAi.usageInputTokens).toBe(5033) + expect(mapped.genAi.usageOutputTokens).toBe(38) + expect(mapped.genAi.requestModel).toBe("openai/gpt-5.6-luna") + expect(mapped.genAi.providerName).toBe("openrouter") + expect(mapped.genAi.responseFinishReasons).toEqual(["stop"]) + expect(mapped.integrationId).toBe("vercel_ai_sdk") + }) + + it("keeps the canonical gen_ai key winning over the ai.* alias", () => { + // Current AI SDK versions emit both dialects on the same span; the + // convention's key has to be the one that lands. + const mapped = mapAiSpan( + row("vercel_ai_sdk", { "gen_ai.usage.input_tokens": "5033", "ai.usage.promptTokens": "1" }), + ) + + expect(mapped.genAi.usageInputTokens).toBe(5033) + }) + + it("maps the tool dialect of an older SDK span", () => { + const mapped = mapAiSpan( + row("vercel_ai_sdk", { + "ai.toolCall.name": "add_reaction", + "ai.toolCall.id": "call_uKgzomwVJhP3bYZ0fxvwUe86", + "ai.toolCall.args": '{"emoji":"wave"}', + "ai.toolCall.result": '{"reacted":true}', + }), + ) + + expect(mapped.genAi.toolName).toBe("add_reaction") + expect(mapped.genAi.toolCallId).toBe("call_uKgzomwVJhP3bYZ0fxvwUe86") + expect(mapped.genAi.toolCallArguments).toEqual({ emoji: "wave" }) + expect(mapped.genAi.toolCallResult).toEqual({ reacted: true }) + }) + + it("falls back to the telemetry function id for the agent name", () => { + // Real spans in this org put the same value in both, and it is the only + // agent identity an older-SDK span carries. + expect( + mapAiSpan(row("vercel_ai_sdk", { "ai.telemetry.functionId": "slack-agent" })).genAi.agentName, + ).toBe("slack-agent") + expect( + mapAiSpan( + row("vercel_ai_sdk", { + "gen_ai.agent.name": "triage", + "ai.telemetry.functionId": "slack-agent", + }), + ).genAi.agentName, + ).toBe("triage") + }) + + it("leaves fields it does not mention on the default source list", () => { + // The merge is per field: `requestSeed` is not in the override, so it + // keeps the default's canonical key AND the default's legacy alias. + const mapped = mapAiSpan(row("vercel_ai_sdk", { "gen_ai.openai.request.seed": "7" })) + + expect(mapped.genAi.requestSeed).toBe(7) + }) + + it("still runs the default refine for a vendor span", () => { + const mapped = mapAiSpan( + row("vercel_ai_sdk", { + "gen_ai.system": "vertex_ai", + "gen_ai.response.finish_reasons": '["tool_calls"]', + }), + ) + + expect(mapped.genAi.providerName).toBe("gcp.vertex_ai") + expect(mapped.genAi.responseFinishReasons).toEqual(["tool_call"]) + }) +}) + +describe("openinference", () => { + it("is registered under both vendor ids the gateway can stamp", () => { + // Same dialect, two detection paths: the OpenAI instrumentor by name, and + // the generic bucket for any other `openinference.instrumentation.*` scope. + expect(AI_VENDOR_INTEGRATIONS["openinference-openai"]).toBe( + AI_VENDOR_INTEGRATIONS["unknown:openinference"], + ) + expect(resolveAiIntegration("unknown:openinference").id).toBe("openinference") + expect(resolveAiIntegration("openinference-openai").id).toBe("openinference") + }) + + it("maps the llm.* dialect", () => { + const mapped = mapAiSpan( + row("openinference-openai", { + "llm.model_name": "gpt-5", + "llm.provider": "openai", + "llm.token_count.prompt": "5033", + "llm.token_count.completion": "38", + "llm.token_count.prompt_details.cache_read": "4924", + "llm.token_count.completion_details.reasoning": "12", + "input.value": '{"messages":[{"role":"user"}]}', + "output.value": '{"messages":[{"role":"assistant"}]}', + "tool.name": "search", + "tool.description": "search the docs", + }), + ) + + expect(mapped.genAi).toMatchObject({ + requestModel: "gpt-5", + providerName: "openai", + usageInputTokens: 5033, + usageOutputTokens: 38, + usageCacheReadInputTokens: 4924, + usageReasoningOutputTokens: 12, + inputMessages: { messages: [{ role: "user" }] }, + outputMessages: { messages: [{ role: "assistant" }] }, + toolName: "search", + toolDescription: "search the docs", + }) + }) + + it("drops the default alias for a field it replaces", () => { + // An override REPLACES the default key list for that field rather than + // extending it, so the default's `gen_ai.completion` alias is gone here. + // This is the whole point of per-field replacement: a dialect gets to say + // which keys are meaningful for it. + const mapped = mapAiSpan( + row("openinference-openai", { "gen_ai.completion": '[{"role":"assistant"}]' }), + ) + + expect(mapped.genAi.outputMessages).toBeUndefined() + }) + + it("translates the span kind into a gen_ai operation name", () => { + expect( + mapAiSpan(row("openinference-openai", { "openinference.span.kind": "LLM" })).genAi.operationName, + ).toBe("chat") + expect( + mapAiSpan(row("openinference-openai", { "openinference.span.kind": "TOOL" })).genAi.operationName, + ).toBe("execute_tool") + expect( + mapAiSpan(row("openinference-openai", { "openinference.span.kind": "AGENT" })).genAi + .operationName, + ).toBe("invoke_agent") + }) + + it("leaves a span kind with no convention equivalent unmapped", () => { + // Better an absent `operationName` than one carrying a value no GenAI + // filter in the product can match. + expect( + mapAiSpan(row("openinference-openai", { "openinference.span.kind": "CHAIN" })).genAi + .operationName, + ).toBeUndefined() + }) + + it("runs after the default refine, so it sees the mapped operation name", () => { + // Hook order is default-then-vendor: the vendor's translation defers to a + // real `gen_ai.operation.name` that the default mapping already produced. + const mapped = mapAiSpan( + row("openinference-openai", { + "gen_ai.operation.name": "chat", + "openinference.span.kind": "TOOL", + }), + ) + + expect(mapped.genAi.operationName).toBe("chat") + }) +}) + +describe("eve", () => { + it("maps a real ai.eve.turn span, which is a session envelope and nothing else", () => { + // eve's own span carries no generation attributes at all — the model call + // happens on Vercel-AI-SDK child spans the gateway stamps separately. + const mapped = mapAiSpan( + row("eve", { + "ai.telemetry.functionId": "slack-agent", + "eve.environment": "production", + "eve.session.id": "wrun_01KZAAFFZRHHRYC8MY9MDANASQ", + "eve.turn.id": "turn_1", + "eve.version": "0.25.3", + "maple_ai.session.id": "wrun_01KZAAFFZRHHRYC8MY9MDANASQ", + }), + ) + + expect(mapped.genAi).toEqual({ conversationId: "turn_1" }) + expect(mapped.sessionId).toBe("wrun_01KZAAFFZRHHRYC8MY9MDANASQ") + expect(mapped.integrationId).toBe("eve") + expect(mapped.isAiSpan).toBe(true) + }) + + it("does not overwrite a conversation id the span already declared", () => { + const mapped = mapAiSpan(row("eve", { "gen_ai.conversation.id": "conv-1", "eve.turn.id": "turn_1" })) + + expect(mapped.genAi.conversationId).toBe("conv-1") + }) + + it("runs both refine hooks, default first", () => { + // Two observable effects on one span: the default's provider rename and + // the vendor's turn-id mapping. + const mapped = mapAiSpan(row("eve", { "gen_ai.system": "xai", "eve.turn.id": "turn_1" })) + + expect(mapped.genAi.providerName).toBe("x_ai") + expect(mapped.genAi.conversationId).toBe("turn_1") + }) +}) diff --git a/packages/query-engine-integrations/src/ai/ai-vendors.ts b/packages/query-engine-integrations/src/ai/ai-vendors.ts new file mode 100644 index 000000000..7aea6eebf --- /dev/null +++ b/packages/query-engine-integrations/src/ai/ai-vendors.ts @@ -0,0 +1,183 @@ +// Per-vendor overrides, keyed by the `maple_ai.vendor.id` the ingest gateway +// stamped on the span. +// +// Only three entries exist, and that is the point: the default GenAI +// integration already maps the twenty-two detected vendors that emit canonical +// `gen_ai.*` attributes, so an override is only worth writing for a framework +// with a genuinely different dialect. Adding one later is a single entry in the +// table at the bottom of this file — no registration step, no new mechanism. +// +// Every key below was verified against the emitting source (the installed AI +// SDK's own telemetry keys, the OpenInference semantic-convention spec, and +// real spans in this org's warehouse). Keys that could not be verified were +// dropped rather than guessed: a wrong key is invisible — it simply never +// matches — which makes guesses uniquely expensive to discover later. + +import type { AiIntegration, AiRefineContext } from "./ai-integrations" +import type { MutableAiGenAiValues } from "./ai-span-model" + +/** + * Vercel AI SDK — the `ai.*` dialect. + * + * Current AI SDK versions emit proper `gen_ai.*` attributes (production spans + * from `apps/slack-agent`, which runs the SDK through eve, carry + * `gen_ai.operation.name`, `gen_ai.usage.*` and friends), so every canonical + * key is listed FIRST and the `ai.*` keys are strictly lower-priority aliases + * for older versions. All `ai.*` keys here appear verbatim in the installed + * `ai` package's telemetry code. + * + * Deliberately not mapped: `ai.response.text` (plain text, not the JSON message + * array `outputMessages` holds), `ai.operationId` (an SDK function id such as + * `ai.generateText.doGenerate`, not a `gen_ai.operation.name` value), and + * `ai.response.msToFirstChunk` (milliseconds, where + * `gen_ai.response.time_to_first_chunk` is seconds — silently mixing units is + * worse than not having the field). + */ +const vercelAiSdkIntegration: AiIntegration = { + id: "vercel_ai_sdk", + sources: { + requestModel: ["gen_ai.request.model", "ai.model.id"], + providerName: ["gen_ai.provider.name", "gen_ai.system", "ai.model.provider"], + responseId: ["gen_ai.response.id", "ai.response.id"], + responseModel: ["gen_ai.response.model", "ai.response.model"], + responseFinishReasons: [ + "gen_ai.response.finish_reasons", + "gen_ai.response.finish_reason", + "ai.response.finishReason", + ], + usageInputTokens: [ + "gen_ai.usage.input_tokens", + "gen_ai.usage.prompt_tokens", + "ai.usage.inputTokens", + "ai.usage.promptTokens", + ], + usageOutputTokens: [ + "gen_ai.usage.output_tokens", + "gen_ai.usage.completion_tokens", + "ai.usage.outputTokens", + "ai.usage.completionTokens", + ], + usageCacheReadInputTokens: [ + "gen_ai.usage.cache_read.input_tokens", + "ai.usage.cachedInputTokens", + "ai.usage.inputTokenDetails.cacheReadTokens", + ], + usageCacheCreationInputTokens: [ + "gen_ai.usage.cache_creation.input_tokens", + "ai.usage.inputTokenDetails.cacheWriteTokens", + ], + usageReasoningOutputTokens: [ + "gen_ai.usage.reasoning.output_tokens", + "gen_ai.usage.output_tokens.reasoning", + "ai.usage.reasoningTokens", + "ai.usage.outputTokenDetails.reasoningTokens", + ], + inputMessages: ["gen_ai.input.messages", "gen_ai.prompt", "ai.prompt.messages", "ai.prompt"], + outputMessages: ["gen_ai.output.messages", "gen_ai.completion"], + toolName: ["gen_ai.tool.name", "ai.toolCall.name"], + toolCallId: ["gen_ai.tool.call.id", "ai.toolCall.id"], + toolCallArguments: ["gen_ai.tool.call.arguments", "ai.toolCall.args"], + toolCallResult: ["gen_ai.tool.call.result", "ai.toolCall.result"], + toolDefinitions: ["gen_ai.tool.definitions", "ai.prompt.tools"], + // `ai.telemetry.functionId` is the name the app gave the traced call. In + // this org's spans it carries the same value the sibling `invoke_agent` + // span puts in `gen_ai.agent.name` (`slack-agent`), which is the only + // agent identity an older-SDK span has. + agentName: ["gen_ai.agent.name", "ai.telemetry.functionId"], + }, +} + +/** + * OpenInference span kinds that have a `gen_ai.operation.name` equivalent. The + * kinds left out (`CHAIN`, `RERANKER`, `GUARDRAIL`, `EVALUATOR`, `PROMPT`, + * `UNKNOWN`) have no counterpart in the convention, and inventing one would put + * a value in `operationName` that no GenAI dashboard filter can match. + */ +const OPENINFERENCE_SPAN_KIND_OPERATIONS = new Map([ + ["LLM", "chat"], + ["TOOL", "execute_tool"], + ["AGENT", "invoke_agent"], + ["EMBEDDING", "embeddings"], + ["RETRIEVER", "retrieval"], +]) + +/** + * OpenInference — the dialect Arize's instrumentors emit. Registered under both + * `openinference-openai` (the gateway's id for the OpenAI instrumentor) and + * `unknown:openinference` (its generic bucket for any other OpenInference + * scope), because the dialect is identical; only the detection path differs. + * + * The integration id is the DIALECT, not the vendor stamp, so both stamps + * report the same `integrationId`. + */ +const openInferenceIntegration: AiIntegration = { + id: "openinference", + sources: { + requestModel: ["gen_ai.request.model", "llm.model_name"], + providerName: ["gen_ai.provider.name", "gen_ai.system", "llm.provider", "llm.system"], + usageInputTokens: ["gen_ai.usage.input_tokens", "llm.token_count.prompt"], + usageOutputTokens: ["gen_ai.usage.output_tokens", "llm.token_count.completion"], + usageCacheReadInputTokens: [ + "gen_ai.usage.cache_read.input_tokens", + "llm.token_count.prompt_details.cache_read", + ], + usageReasoningOutputTokens: [ + "gen_ai.usage.reasoning.output_tokens", + "llm.token_count.completion_details.reasoning", + ], + inputMessages: ["gen_ai.input.messages", "llm.input_messages", "input.value"], + outputMessages: ["gen_ai.output.messages", "llm.output_messages", "output.value"], + toolName: ["gen_ai.tool.name", "tool.name"], + toolDescription: ["gen_ai.tool.description", "tool.description"], + toolCallArguments: ["gen_ai.tool.call.arguments", "tool.parameters"], + toolDefinitions: ["gen_ai.tool.definitions", "llm.tools"], + }, + refine: (values: MutableAiGenAiValues, ctx: AiRefineContext) => { + // `openinference.span.kind` is the dialect's operation classifier, but it + // is an enum of a different vocabulary rather than a differently named + // `gen_ai.operation.name`, so translating it is a refine, not an alias. + if (values.operationName !== undefined) return + const operation = OPENINFERENCE_SPAN_KIND_OPERATIONS.get( + ctx.attributes["openinference.span.kind"] ?? "", + ) + if (operation !== undefined) values.operationName = operation + }, +} + +/** + * eve — a session envelope rather than a GenAI dialect. + * + * eve's own spans (`ai.eve.turn`) carry `eve.session.id`, `eve.turn.id`, + * `eve.environment` and `eve.version`, and nothing else AI-shaped: the model + * call itself is made through the Vercel AI SDK on child spans, which the + * gateway stamps separately. The session id is already lifted into + * `maple_ai.session.id`, and `eve.environment` / `eve.version` describe the + * deployment rather than the generation, so the only mapping worth making is + * the turn id — the conversation-level grouping key inside a session. This + * override is deliberately minimal. + */ +const eveIntegration: AiIntegration = { + id: "eve", + sources: {}, + refine: (values: MutableAiGenAiValues, ctx: AiRefineContext) => { + if (values.conversationId !== undefined) return + const turnId = ctx.attributes["eve.turn.id"] + if (turnId !== undefined && turnId !== "") values.conversationId = turnId + }, +} + +/** + * Vendor id → override. Every id the gateway can stamp that is NOT in here maps + * through the default GenAI integration, which is the right answer for the + * frameworks that emit canonical `gen_ai.*`. + */ +export interface AiVendorRegistry { + readonly [vendorId: string]: AiIntegration | undefined +} + +export const AI_VENDOR_INTEGRATIONS: AiVendorRegistry = { + vercel_ai_sdk: vercelAiSdkIntegration, + "openinference-openai": openInferenceIntegration, + "unknown:openinference": openInferenceIntegration, + eve: eveIntegration, +} diff --git a/packages/query-engine-integrations/src/ai/index.ts b/packages/query-engine-integrations/src/ai/index.ts new file mode 100644 index 000000000..99a2cc7c2 --- /dev/null +++ b/packages/query-engine-integrations/src/ai/index.ts @@ -0,0 +1,51 @@ +// AI agent sessions — the warehouse queries that resolve sessions and their +// spans from the `maple_ai.*` attributes the ingest gateway stamps at decode +// time, plus the integration layer that maps each raw span onto Maple's +// standardised AI agent span format. +// +// The two halves compose: `aiSessionSpansQuery` rows structurally satisfy +// `AiSessionSpanRow`, so `mapAiSpans(rows)` is the whole read path. + +export { + aiSessionListQuery, + aiSessionListRowSchema, + aiSessionSpansQuery, + aiSessionSpansRowSchema, + type AiSessionListOpts, + type AiSessionListOutput, + type AiSessionSpansOpts, + type AiSessionSpansOutput, +} from "./ai-sessions" + +export { + AI_GENAI_FIELDS, + AI_KNOWN_OPERATION_NAMES, + AI_PROMPT_VARIABLE_PREFIX, + AiAgentSpanSchema, + AiGenAiValuesSchema, + MAPLE_AI_SESSION_ID_ATTR, + MAPLE_AI_VENDOR_ID_ATTR, + MAPLE_AI_VENDOR_VERSION_ATTR, + type AiAgentSpan, + type AiFieldDef, + type AiFieldGroup, + type AiFieldValue, + type AiGenAiField, + type AiGenAiValues, + type AiKnownOperationName, + type AiSessionSpanRow, + type MutableAiGenAiValues, +} from "./ai-span-model" + +export { + genAiIntegration, + mapAiSpan, + mapAiSpans, + resolveAiIntegration, + type AiDecodedValue, + type AiIntegration, + type AiJsonValue, + type AiRefineContext, +} from "./ai-integrations" + +export { AI_VENDOR_INTEGRATIONS, type AiVendorRegistry } from "./ai-vendors" diff --git a/packages/query-engine-integrations/src/catalog.ts b/packages/query-engine-integrations/src/catalog.ts index e177d07d7..3dbf55c80 100644 --- a/packages/query-engine-integrations/src/catalog.ts +++ b/packages/query-engine-integrations/src/catalog.ts @@ -26,6 +26,33 @@ const END_TIME = "2026-01-03 14:15:00" const window = { orgId: ORG_ID, startTime: START_TIME, endTime: END_TIME } export const integrationFixtures: ReadonlyArray = [ + { + module: "ai-sessions", + name: "aiSessionListQuery", + label: "default", + compile: () => compile(CH.aiSessionListQuery(), window), + }, + { + // The vendor/service filters the AI sessions list page sends. + module: "ai-sessions", + name: "aiSessionListQuery", + label: "filtered", + compile: () => + compile( + CH.aiSessionListQuery({ + limit: 25, + vendorIds: ["eve"], + serviceNames: ["maple-slack-agent"], + }), + window, + ), + }, + { + module: "ai-sessions", + name: "aiSessionSpansQuery", + label: "default", + compile: () => compile(CH.aiSessionSpansQuery(), { ...window, sessionId: "wrun_sql_catalog" }), + }, { module: "cloudflare-infra", name: "cloudflareZoneLatencySQL", diff --git a/packages/query-engine-integrations/src/index.ts b/packages/query-engine-integrations/src/index.ts index 0598d80d5..9409609e5 100644 --- a/packages/query-engine-integrations/src/index.ts +++ b/packages/query-engine-integrations/src/index.ts @@ -12,6 +12,7 @@ // layer in `@maple/query-engine/observability` builds on those queries, so // moving them would make the two packages circular. +export * from "./ai/index" export * from "./cloudflare/index" export * from "./planetscale/index" export * from "./product/index" From cddfb9fa930f013db63bf43927ffa075453d9989 Mon Sep 17 00:00:00 2001 From: JeremyFunk Date: Wed, 19 Aug 2026 19:53:00 +0200 Subject: [PATCH 2/4] fix(query-engine): harden AI span mapping against hostile vendor stamps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings on the AI agent session read path. The vendor stamp reaches a plain object index, and the ingest gateway strips `maple_ai.*` from span attributes but NOT from resource attributes — so a customer-supplied `maple_ai.vendor.id` of `constructor` or `toString` resolved off `Object.prototype`, passed the `undefined` guard, and produced an integration with `id: undefined`. That violates `AiAgentSpanSchema` at the encode boundary, far from the cause, and was memoised process-wide. Guarded with `Object.hasOwn`, and the attribute merge and prompt-variable accumulator are now null-prototype so the whole class is closed rather than the one reachable instance. `decodeStringArray` wrapped anything that was not a string array into `[raw]`. A parsed object or an array of numbers therefore became a one-element array holding raw JSON text — a value that type-checks, silently consumes the field, and denies the next alias its turn. Only an unparseable value is the bare form now; parsed-but-wrong-shape decodes to nothing, like every other type. `decodeAttribute`'s switch had no exhaustiveness guard, and the package has no `noImplicitReturns`, so a sixth field type would have compiled clean and made every field of that type silently absent. Also: export the spans cap so callers can request `+1` and detect truncation instead of hardcoding it — truncation drops the END of a session, which is where the agent's answer is; attach the row schemas to the catalog fixtures, since the ClickHouse e2e sweep only runs its 64-bit decode assertion for fixtures that carry one; and assert the query row and the mapper input agree, mutually — `extends` alone would accept a query that grew a column the mapper never sees. Co-Authored-By: Claude Opus 5 --- .../src/ai/ai-integrations.test.ts | 70 +++++++++++++++++++ .../src/ai/ai-integrations.ts | 42 +++++++++-- .../src/ai/ai-sessions.ts | 19 ++++- .../query-engine-integrations/src/ai/index.ts | 21 ++++++ .../query-engine-integrations/src/catalog.ts | 13 +++- 5 files changed, 158 insertions(+), 7 deletions(-) diff --git a/packages/query-engine-integrations/src/ai/ai-integrations.test.ts b/packages/query-engine-integrations/src/ai/ai-integrations.test.ts index 0ddc53603..e03bedb39 100644 --- a/packages/query-engine-integrations/src/ai/ai-integrations.test.ts +++ b/packages/query-engine-integrations/src/ai/ai-integrations.test.ts @@ -378,3 +378,73 @@ describe("resolveAiIntegration", () => { expect(resolveAiIntegration("vercel_ai_sdk")).toBe(resolveAiIntegration("vercel_ai_sdk")) }) }) + +// The vendor stamp is customer-reachable: the ingest gateway strips `maple_ai.*` +// from span attributes but NOT from resource attributes, so these names can +// arrive from outside. A plain index reads them off `Object.prototype`. +describe("hostile vendor stamps", () => { + const prototypeKeys = ["constructor", "toString", "valueOf", "hasOwnProperty", "__proto__"] + + for (const stamp of prototypeKeys) { + it(`falls back to the default integration for a stamp of ${stamp}`, () => { + const resolved = resolveAiIntegration(stamp) + + // Before the Object.hasOwn guard this resolved to a merged integration + // with `id: undefined`, which violates AiAgentSpanSchema at the encode + // boundary and was memoised process-wide. + expect(resolved.id).toBe("gen_ai") + expect(resolved).toBe(genAiIntegration) + }) + } + + it("maps a span whose stamp arrives via a resource attribute", () => { + const mapped = mapAiSpan(row({}, { resourceAttributes: { "maple_ai.vendor.id": "constructor" } })) + + expect(mapped.integrationId).toBe("gen_ai") + expect(typeof mapped.integrationId).toBe("string") + }) + + it("keeps a prompt variable literally named __proto__", () => { + const mapped = mapAiSpan(row({ "gen_ai.prompt.variable.__proto__": "kept" })) + + // Asserted through `Object.keys` rather than against an object literal: a + // `{ __proto__: "kept" }` literal collapses to `{}` via the prototype + // setter, so the expectation would silently test nothing. + expect(Object.keys(mapped.promptVariables ?? {})).toEqual(["__proto__"]) + expect(mapped.promptVariables?.["__proto__"]).toBe("kept") + }) + + it("does not mistake an inherited member for an attribute value", () => { + const mapped = mapAiSpan(row({ "gen_ai.request.model": "gpt-5.6-luna" })) + + expect(mapped.genAi.requestModel).toBe("gpt-5.6-luna") + expect(mapped.genAi.toolCallResult).toBeUndefined() + }) +}) + +describe("stringArray decoding rejects malformed arrays", () => { + it("treats a bare value as the single-element form", () => { + expect( + mapAiSpan(row({ "gen_ai.response.finish_reasons": "stop" })).genAi.responseFinishReasons, + ).toEqual(["stop"]) + }) + + it("yields no field for an array of non-strings rather than wrapping the raw JSON", () => { + const mapped = mapAiSpan(row({ "gen_ai.response.finish_reasons": "[1,2]" })) + + // Wrapping to `['[1,2]']` would type-check and silently consume the field + // with a value that never existed. + expect(mapped.genAi.responseFinishReasons).toBeUndefined() + }) + + it("lets the next alias win when the canonical key is malformed", () => { + const mapped = mapAiSpan( + row({ + "gen_ai.response.finish_reasons": '{"reason":"stop"}', + "gen_ai.response.finish_reason": "length", + }), + ) + + expect(mapped.genAi.responseFinishReasons).toEqual(["length"]) + }) +}) diff --git a/packages/query-engine-integrations/src/ai/ai-integrations.ts b/packages/query-engine-integrations/src/ai/ai-integrations.ts index b0fac82b5..420795415 100644 --- a/packages/query-engine-integrations/src/ai/ai-integrations.ts +++ b/packages/query-engine-integrations/src/ai/ai-integrations.ts @@ -96,13 +96,23 @@ const parseJson = (raw: string): AiJsonValue | undefined => { } } -const decodeStringArray = (raw: string): readonly string[] => { +const decodeStringArray = (raw: string): readonly string[] | undefined => { // Real data carries both shapes for the same attribute: `'["stop"]'` from // instrumentation that serialises the array, and a bare `"stop"` from // instrumentation that emits the single value. Anything that is not a JSON // array of strings is treated as the bare form rather than discarded. const parsed = parseJson(raw) - return Array.isArray(parsed) && parsed.every((entry) => typeof entry === "string") ? parsed : [raw] + // Unparseable is the bare form: `stop` and `"stop"`-without-quotes both land + // here, and that is the only case where the raw text IS the value. + if (parsed === undefined) return [raw] + if (typeof parsed === "string") return [parsed] + if (Array.isArray(parsed) && parsed.every((entry) => typeof entry === "string")) return parsed + // Parsed cleanly but into some other shape — an object, a number, an array of + // non-strings. That is structured data in the wrong shape, not a bare value. + // Wrapping the raw JSON text into a one-element array would type-check and + // silently consume the field with a value that never existed, so decode to + // nothing and let the next alias have its turn. + return undefined } const decodeAttribute = (type: AiFieldDef["type"], raw: string): AiDecodedValue | undefined => { @@ -125,6 +135,15 @@ const decodeAttribute = (type: AiFieldDef["type"], raw: string): AiDecodedValue return decodeStringArray(raw) case "json": return parseJson(raw) + default: { + // A sixth field type added to the catalog without a case here would + // otherwise compile clean and make every field of that type silently + // absent from every span — the one failure mode this module is least + // able to surface. `tsconfig` has no `noImplicitReturns`, so the + // never-assignment is what actually enforces it. + const unhandled: never = type + return unhandled + } } } @@ -227,6 +246,12 @@ const resolvedIntegrations = new Map() */ export const resolveAiIntegration = (vendorId: string | undefined): AiIntegration => { if (vendorId === undefined) return genAiIntegration + // `Object.hasOwn`, not a plain index: the vendor stamp is customer-reachable + // (the gateway strips `maple_ai.*` from span attributes but not from RESOURCE + // attributes), and a plain index reads through the prototype chain — a stamp + // of `constructor` or `toString` would resolve truthy, skip the guard below, + // and mint an integration with `id: undefined` into the module-level cache. + if (!Object.hasOwn(AI_VENDOR_INTEGRATIONS, vendorId)) return genAiIntegration const vendor = AI_VENDOR_INTEGRATIONS[vendorId] if (vendor === undefined) return genAiIntegration const cached = resolvedIntegrations.get(vendorId) @@ -249,7 +274,9 @@ const collectPromptVariables = (attributes: Record): Record | undefined for (const [key, value] of Object.entries(attributes)) { if (!key.startsWith(AI_PROMPT_VARIABLE_PREFIX) || value === "") continue - collected ??= {} + // Null-prototype again: a `gen_ai.prompt.variable.__proto__` key would hit + // the prototype setter on a `{}` literal and be dropped without a trace. + collected ??= Object.create(null) as Record collected[key.slice(AI_PROMPT_VARIABLE_PREFIX.length)] = value } return collected @@ -270,7 +297,14 @@ interface AiSpanOptionalFields { } export const mapAiSpan = (row: AiSessionSpanRow): AiAgentSpan => { - const attributes = { ...row.resourceAttributes, ...row.spanAttributes } + // Null-prototype: attribute keys are customer-controlled, and a `Record` + // index would otherwise resolve `toString` or `valueOf` to an inherited + // function that the `=== ""` check below would wave through as a value. + const attributes: Record = Object.assign( + Object.create(null) as Record, + row.resourceAttributes, + row.spanAttributes, + ) const vendorId = readAttribute(attributes, MAPLE_AI_VENDOR_ID_ATTR) const integration = resolveAiIntegration(vendorId) diff --git a/packages/query-engine-integrations/src/ai/ai-sessions.ts b/packages/query-engine-integrations/src/ai/ai-sessions.ts index d5cca112f..ab43030cf 100644 --- a/packages/query-engine-integrations/src/ai/ai-sessions.ts +++ b/packages/query-engine-integrations/src/ai/ai-sessions.ts @@ -67,6 +67,14 @@ const VENDOR_VERSION_ATTR = "maple_ai.vendor.version" */ const SESSION_ORDER_SENTINEL = "2106-01-01 00:00:00" +/** + * Default span cap for `aiSessionSpansQuery`, exported so a caller can request + * `+ 1` and detect truncation instead of hardcoding the number. Mirrors + * `SPAN_HIERARCHY_MAX_SPANS`, which is exported from the core package for the + * same reason. + */ +export const AI_SESSION_SPANS_MAX_SPANS = 2_000 + /** ClickHouse returns `''` for a missing Map key, so presence needs both halves. */ const hasSessionId = (attrs: CH.Expr>, get: CH.Expr) => CH.mapContains(attrs, SESSION_ID_ATTR).and(get.neq("")) @@ -268,9 +276,18 @@ export const aiSessionSpansRowSchema: CompiledQueryRowSchema ({ TraceId: $.TraceId })) diff --git a/packages/query-engine-integrations/src/ai/index.ts b/packages/query-engine-integrations/src/ai/index.ts index 99a2cc7c2..7a5d322d9 100644 --- a/packages/query-engine-integrations/src/ai/index.ts +++ b/packages/query-engine-integrations/src/ai/index.ts @@ -9,6 +9,7 @@ export { aiSessionListQuery, aiSessionListRowSchema, + AI_SESSION_SPANS_MAX_SPANS, aiSessionSpansQuery, aiSessionSpansRowSchema, type AiSessionListOpts, @@ -49,3 +50,23 @@ export { } from "./ai-integrations" export { AI_VENDOR_INTEGRATIONS, type AiVendorRegistry } from "./ai-vendors" + +// The two halves above are joined by structure alone — the query module and the +// span model each declare the row shape independently, and nothing imports the +// other. This assertion is what makes the "structurally satisfy" claim in the +// header enforceable: add a column to the query, retype one, or drop one, and +// the read path fails to compile here rather than quietly mapping fewer fields +// in production. +import type { AiSessionSpansOutput } from "./ai-sessions" +import type { AiSessionSpanRow } from "./ai-span-model" + +type Assert = T +// Mutual, not one-directional: `extends` alone would accept a query that grew a +// column the mapper never sees, which is exactly the drift worth catching. +type _QueryRowSatisfiesMapperInput = Assert< + AiSessionSpansOutput extends AiSessionSpanRow + ? AiSessionSpanRow extends AiSessionSpansOutput + ? true + : false + : false +> diff --git a/packages/query-engine-integrations/src/catalog.ts b/packages/query-engine-integrations/src/catalog.ts index 3dbf55c80..dd16f78b1 100644 --- a/packages/query-engine-integrations/src/catalog.ts +++ b/packages/query-engine-integrations/src/catalog.ts @@ -30,7 +30,10 @@ export const integrationFixtures: ReadonlyArray = [ module: "ai-sessions", name: "aiSessionListQuery", label: "default", - compile: () => compile(CH.aiSessionListQuery(), window), + // Row schemas are attached here, not just in the unit tests: the ClickHouse + // e2e sweep only runs its quoted/unquoted 64-bit decode assertion for + // fixtures whose compiled query carries one. + compile: () => compile(CH.aiSessionListQuery(), window, { rowSchema: CH.aiSessionListRowSchema }), }, { // The vendor/service filters the AI sessions list page sends. @@ -45,13 +48,19 @@ export const integrationFixtures: ReadonlyArray = [ serviceNames: ["maple-slack-agent"], }), window, + { rowSchema: CH.aiSessionListRowSchema }, ), }, { module: "ai-sessions", name: "aiSessionSpansQuery", label: "default", - compile: () => compile(CH.aiSessionSpansQuery(), { ...window, sessionId: "wrun_sql_catalog" }), + compile: () => + compile( + CH.aiSessionSpansQuery(), + { ...window, sessionId: "wrun_sql_catalog" }, + { rowSchema: CH.aiSessionSpansRowSchema }, + ), }, { module: "cloudflare-infra", From 77497e8d5f697358523b76f11c2a4dac4a6dee7f Mon Sep 17 00:00:00 2001 From: JeremyFunk Date: Wed, 19 Aug 2026 19:59:48 +0200 Subject: [PATCH 3/4] fix(query-engine): correct AI session duration and tighten the read path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second round of review findings. `traceEnd` was `max(Timestamp)`, but `Timestamp` is the span's START, so the session's end was the moment its last span BEGAN. Every session under-reported by exactly the last-starting span's duration — 10417 ms against a true 10433 ms on a real production session — and a session whose trace is one long span reported a duration of 0. The error is always negative and always looks like plausible jitter, so it would not have been noticed from the numbers. Now `max(Timestamp + Duration)`, carried through as nanos, the same idiom `tracesDetailQuery` uses. Verified against production: the endpoint is now the nanosecond-exact trace end. `aiSessionSpansQuery` with an empty `sessionId` matched every span that merely LACKS the key, because ClickHouse reads a missing Map key back as `''` — an empty param degraded into a whole-org trace dump. The presence guard the list query already had now covers this one too. The `maple_ai.*` envelope is read from span attributes only. It was read from a merge of span and resource attributes, and the gateway strips that namespace from span attributes but not from resource attributes, so one forged resource attribute marked every span in a service as an AI span and labelled it with a session id the query never matched on. The OpenInference override dropped six of the default's legacy aliases while the Vercel override kept them, so identifying a span as OpenInference LOST its token counts and messages — a recognised vendor mapped strictly worse than an unrecognised one. Both overrides now carry the same policy, and a registry-driven test pins the invariant. The test that asserted the old behaviour as intentional was encoding the bug; it now demonstrates replacement by precedence instead. Also: `count()` rather than approximate `uniq()` for `traceCount`, since the derived table already emits one row per trace; and a `spanId` tiebreaker so truncation is deterministic and cannot orphan children. Co-Authored-By: Claude Opus 5 --- .../src/__sql_baseline__/integrations.sql | 19 +++-- .../src/ai/ai-integrations.ts | 14 +++- .../src/ai/ai-sessions.ts | 84 ++++++++++++------- .../src/ai/ai-vendors.test.ts | 61 ++++++++++++-- .../src/ai/ai-vendors.ts | 29 ++++++- 5 files changed, 154 insertions(+), 53 deletions(-) diff --git a/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql b/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql index 2c172bc56..eb2ce123e 100644 --- a/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql +++ b/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql @@ -3,13 +3,13 @@ SELECT sessionId AS sessionId, argMin(vendorId, sessionStart) AS vendorId, argMin(vendorVersion, sessionStart) AS vendorVersion, - uniq(traceId) AS traceCount, + count() AS traceCount, sum(spanCount) AS spanCount, sum(errorSpanCount) AS errorSpanCount, groupUniqArrayArray(serviceNames) AS serviceNames, toString(min(traceStart)) AS startTime, - toString(max(traceEnd)) AS endTime, - intDiv(toUnixTimestamp64Nano(max(traceEnd)) - toUnixTimestamp64Nano(min(traceStart)), 1000000) AS durationMs + toString(fromUnixTimestamp64Nano(max(traceEndNanos))) AS endTime, + intDiv(max(traceEndNanos) - toUnixTimestamp64Nano(min(traceStart)), 1000000) AS durationMs FROM (SELECT TraceId AS traceId, max(SpanAttributes['maple_ai.session.id']) AS sessionId, @@ -20,7 +20,7 @@ SELECT countIf(StatusCode = 'Error') AS errorSpanCount, groupUniqArray(ServiceName) AS serviceNames, min(Timestamp) AS traceStart, - max(Timestamp) AS traceEnd + max(toUnixTimestamp64Nano(Timestamp) + toInt64(Duration)) AS traceEndNanos FROM trace_detail_spans WHERE OrgId = 'org_sql_catalog' AND Timestamp >= '2026-01-01 10:30:00' @@ -44,13 +44,13 @@ SELECT sessionId AS sessionId, argMin(vendorId, sessionStart) AS vendorId, argMin(vendorVersion, sessionStart) AS vendorVersion, - uniq(traceId) AS traceCount, + count() AS traceCount, sum(spanCount) AS spanCount, sum(errorSpanCount) AS errorSpanCount, groupUniqArrayArray(serviceNames) AS serviceNames, toString(min(traceStart)) AS startTime, - toString(max(traceEnd)) AS endTime, - intDiv(toUnixTimestamp64Nano(max(traceEnd)) - toUnixTimestamp64Nano(min(traceStart)), 1000000) AS durationMs + toString(fromUnixTimestamp64Nano(max(traceEndNanos))) AS endTime, + intDiv(max(traceEndNanos) - toUnixTimestamp64Nano(min(traceStart)), 1000000) AS durationMs FROM (SELECT TraceId AS traceId, max(SpanAttributes['maple_ai.session.id']) AS sessionId, @@ -61,7 +61,7 @@ SELECT countIf(StatusCode = 'Error') AS errorSpanCount, groupUniqArray(ServiceName) AS serviceNames, min(Timestamp) AS traceStart, - max(Timestamp) AS traceEnd + max(toUnixTimestamp64Nano(Timestamp) + toInt64(Duration)) AS traceEndNanos FROM trace_detail_spans WHERE OrgId = 'org_sql_catalog' AND Timestamp >= '2026-01-01 10:30:00' @@ -106,8 +106,9 @@ SELECT WHERE OrgId = 'org_sql_catalog' AND Timestamp >= '2026-01-01 10:30:00' AND Timestamp <= '2026-01-03 14:15:00' + AND (mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != '') AND SpanAttributes['maple_ai.session.id'] = 'wrun_sql_catalog') - ORDER BY timestamp ASC + ORDER BY timestamp ASC, spanId ASC LIMIT 2000 FORMAT JSON diff --git a/packages/query-engine-integrations/src/ai/ai-integrations.ts b/packages/query-engine-integrations/src/ai/ai-integrations.ts index 420795415..6d0bf5f02 100644 --- a/packages/query-engine-integrations/src/ai/ai-integrations.ts +++ b/packages/query-engine-integrations/src/ai/ai-integrations.ts @@ -305,7 +305,15 @@ export const mapAiSpan = (row: AiSessionSpanRow): AiAgentSpan => { row.resourceAttributes, row.spanAttributes, ) - const vendorId = readAttribute(attributes, MAPLE_AI_VENDOR_ID_ATTR) + // The `maple_ai.*` envelope is read from SPAN attributes only, never the + // merged view. The gateway strips this namespace from span attributes before + // stamping its own verdict, so a span-level value is authoritative — but it + // does not touch resource attributes, and `aiSessionSpansQuery` selects + // sessions on `SpanAttributes` alone. Reading the merged view would let one + // forged resource attribute mark every span in a service as an AI span and + // label it with a session id the query never matched on. + const envelope = row.spanAttributes + const vendorId = readAttribute(envelope, MAPLE_AI_VENDOR_ID_ATTR) const integration = resolveAiIntegration(vendorId) const genAi: MutableAiGenAiValues = {} @@ -325,8 +333,8 @@ export const mapAiSpan = (row: AiSessionSpanRow): AiAgentSpan => { integration.refine?.(genAi, { row, attributes }) const promptVariables = collectPromptVariables(attributes) - const sessionId = readAttribute(attributes, MAPLE_AI_SESSION_ID_ATTR) - const vendorVersion = readAttribute(attributes, MAPLE_AI_VENDOR_VERSION_ATTR) + const sessionId = readAttribute(envelope, MAPLE_AI_SESSION_ID_ATTR) + const vendorVersion = readAttribute(envelope, MAPLE_AI_VENDOR_VERSION_ATTR) // Collected separately so an absent stamp leaves the key off the span // entirely rather than present-and-undefined. diff --git a/packages/query-engine-integrations/src/ai/ai-sessions.ts b/packages/query-engine-integrations/src/ai/ai-sessions.ts index ab43030cf..d8450649c 100644 --- a/packages/query-engine-integrations/src/ai/ai-sessions.ts +++ b/packages/query-engine-integrations/src/ai/ai-sessions.ts @@ -44,6 +44,7 @@ import { Schema } from "effect" import * as CH from "@maple-dev/clickhouse-builder/expr" import { + compileFnCall, from, fromQuery, inSubquery, @@ -79,6 +80,10 @@ export const AI_SESSION_SPANS_MAX_SPANS = 2_000 const hasSessionId = (attrs: CH.Expr>, get: CH.Expr) => CH.mapContains(attrs, SESSION_ID_ATTR).and(get.neq("")) +/** Not in the builder's function set; same local helper `tracesDetailQuery` uses. */ +const fromUnixTimestamp64Nano = (nanos: CH.Expr): CH.Expr => + compileFnCall("fromUnixTimestamp64Nano", nanos) + export interface AiSessionListOpts { /** Sessions returned, most recently started first. */ readonly limit?: number @@ -177,7 +182,14 @@ export function aiSessionListQuery(opts: AiSessionListOpts = {}) { // and `toUnixTimestamp64Nano` reject it — verified against production, // it fails with ILLEGAL_TYPE_OF_ARGUMENT. traceStart: CH.min_($.Timestamp), - traceEnd: CH.max_($.Timestamp), + // `Timestamp` is the span's START, so `max(Timestamp)` is when the + // last span BEGAN — the trace end is that span's start plus its own + // duration. Without the `+ Duration` a session whose trace is a + // single long span reports a duration of 0, and every other session + // under-reports by exactly the last-starting span's duration, which + // is invisible because it always looks like plausible jitter. Same + // idiom as `tracesDetailQuery`. + traceEndNanos: CH.max_(CH.toUnixTimestamp64Nano($.Timestamp).add(CH.toInt64($.Duration))), } }) .where(($) => [ @@ -194,19 +206,20 @@ export function aiSessionListQuery(opts: AiSessionListOpts = {}) { sessionId: $.sessionId, vendorId: CH.argMin($.vendorId, $.sessionStart), vendorVersion: CH.argMin($.vendorVersion, $.sessionStart), - traceCount: CH.uniq($.traceId), + // `count()`, not `uniq()`: the derived table already emits exactly one + // row per trace, so this is exact and cheaper — `uniq` is an + // approximate HLL that would start drifting on a very large session. + traceCount: CH.count(), spanCount: CH.sum($.spanCount), errorSpanCount: CH.sum($.errorSpanCount), serviceNames: CH.groupUniqArrayArray($.serviceNames), startTime: CH.toString_(CH.min_($.traceStart)), - endTime: CH.toString_(CH.max_($.traceEnd)), + endTime: CH.toString_(fromUnixTimestamp64Nano(CH.max_($.traceEndNanos))), // Nanoseconds first: `Timestamp` is DateTime64(9), and subtracting two // of them yields a Decimal whose scale the wire format then quotes. // Wrapped in `intDiv` because `Expr.sub`/`div` do not parenthesize. durationMs: CH.intDiv( - CH.toUnixTimestamp64Nano(CH.max_($.traceEnd)).sub( - CH.toUnixTimestamp64Nano(CH.min_($.traceStart)), - ), + CH.max_($.traceEndNanos).sub(CH.toUnixTimestamp64Nano(CH.min_($.traceStart))), 1_000_000, ), })) @@ -295,31 +308,42 @@ export function aiSessionSpansQuery(opts: AiSessionSpansOpts = {}) { $.OrgId.eq(param.string("orgId")), $.Timestamp.gte(param.dateTime("startTime")), $.Timestamp.lte(param.dateTime("endTime")), + // The presence guard is what stops an empty `sessionId` param from + // matching every span that simply LACKS the key — ClickHouse reads a + // missing Map key back as `''`, so equality alone would turn a blank + // session id into a whole-org trace dump. + hasSessionId($.SpanAttributes, $.SpanAttributes.get(SESSION_ID_ATTR)), $.SpanAttributes.get(SESSION_ID_ATTR).eq(param.string("sessionId")), ]) - return from(TraceDetailSpans) - .select(($) => ({ - traceId: $.TraceId, - spanId: $.SpanId, - parentSpanId: $.ParentSpanId, - spanName: $.SpanName, - spanKind: $.SpanKind, - serviceName: $.ServiceName, - durationMs: $.Duration.div(1_000_000), - statusCode: $.StatusCode, - statusMessage: $.StatusMessage, - timestamp: CH.toString_($.Timestamp), - spanAttributes: $.SpanAttributes, - resourceAttributes: $.ResourceAttributes, - })) - .where(($) => [ - $.OrgId.eq(param.string("orgId")), - $.Timestamp.gte(param.dateTime("startTime")), - $.Timestamp.lte(param.dateTime("endTime")), - inSubquery($.TraceId, sessionTraceIds), - ]) - .orderBy(["timestamp", "asc"]) - .limit(limit) - .format("JSON") + return ( + from(TraceDetailSpans) + .select(($) => ({ + traceId: $.TraceId, + spanId: $.SpanId, + parentSpanId: $.ParentSpanId, + spanName: $.SpanName, + spanKind: $.SpanKind, + serviceName: $.ServiceName, + durationMs: $.Duration.div(1_000_000), + statusCode: $.StatusCode, + statusMessage: $.StatusMessage, + timestamp: CH.toString_($.Timestamp), + spanAttributes: $.SpanAttributes, + resourceAttributes: $.ResourceAttributes, + })) + .where(($) => [ + $.OrgId.eq(param.string("orgId")), + $.Timestamp.gte(param.dateTime("startTime")), + $.Timestamp.lte(param.dateTime("endTime")), + inSubquery($.TraceId, sessionTraceIds), + ]) + // `spanId` breaks ties: agent spans routinely share a millisecond, and + // without it the LIMIT cuts an arbitrary subset, so two loads of the same + // truncated session can disagree and a parent can survive while its + // children are dropped. + .orderBy(["timestamp", "asc"], ["spanId", "asc"]) + .limit(limit) + .format("JSON") + ) } diff --git a/packages/query-engine-integrations/src/ai/ai-vendors.test.ts b/packages/query-engine-integrations/src/ai/ai-vendors.test.ts index 585733838..8713119a3 100644 --- a/packages/query-engine-integrations/src/ai/ai-vendors.test.ts +++ b/packages/query-engine-integrations/src/ai/ai-vendors.test.ts @@ -142,16 +142,33 @@ describe("openinference", () => { }) }) - it("drops the default alias for a field it replaces", () => { - // An override REPLACES the default key list for that field rather than - // extending it, so the default's `gen_ai.completion` alias is gone here. - // This is the whole point of per-field replacement: a dialect gets to say - // which keys are meaningful for it. + it("replaces the default key list rather than extending it", () => { + // The replacement is real — an override decides the whole list and its + // order for a field it claims. Demonstrated by precedence rather than by + // loss: the dialect key wins over a lower-priority alias in the same list. const mapped = mapAiSpan( - row("openinference-openai", { "gen_ai.completion": '[{"role":"assistant"}]' }), + row("openinference-openai", { + "llm.output_messages": '[{"role":"assistant","from":"dialect"}]', + "gen_ai.completion": '[{"role":"assistant","from":"legacy"}]', + }), ) - expect(mapped.genAi.outputMessages).toBeUndefined() + expect(mapped.genAi.outputMessages).toEqual([{ role: "assistant", from: "dialect" }]) + }) + + it("still reads the default's legacy aliases it did not supersede", () => { + // Regression: the override used to omit these, so identifying a span as + // OpenInference LOST its token counts and messages — a recognised vendor + // mapped strictly worse than an unrecognised one. + const mapped = mapAiSpan( + row("openinference-openai", { + "gen_ai.usage.prompt_tokens": "120", + "gen_ai.completion": '[{"role":"assistant"}]', + }), + ) + + expect(mapped.genAi.usageInputTokens).toBe(120) + expect(mapped.genAi.outputMessages).toEqual([{ role: "assistant" }]) }) it("translates the span kind into a gen_ai operation name", () => { @@ -226,3 +243,33 @@ describe("eve", () => { expect(mapped.genAi.conversationId).toBe("turn_1") }) }) + +describe("recognising a vendor never maps worse than not recognising it", () => { + // A vendor's key list REPLACES the default's for that field, so an override + // that forgets the default's legacy aliases silently loses fields precisely + // because the span was identified. Driven from the registry so a new override + // inherits the check. + const LEGACY_ONLY_SPAN = { + "gen_ai.usage.prompt_tokens": "120", + "gen_ai.usage.completion_tokens": "34", + "gen_ai.usage.input_tokens.cached": "2048", + "gen_ai.usage.output_tokens.reasoning": "704", + "gen_ai.prompt": '[{"role":"user"}]', + "gen_ai.completion": '[{"role":"assistant"}]', + "gen_ai.system": "anthropic", + "gen_ai.response.finish_reason": "stop", + } + + const mappedFieldCount = (vendorId: string) => + Object.keys(mapAiSpan(row(vendorId, LEGACY_ONLY_SPAN)).genAi).length + + // An unregistered stamp resolves to the default integration, which is the + // baseline every override has to at least match. + const baseline = mappedFieldCount("unknown:other") + + for (const vendorId of Object.keys(AI_VENDOR_INTEGRATIONS)) { + it(`maps at least as many legacy fields under ${vendorId}`, () => { + expect(mappedFieldCount(vendorId)).toBeGreaterThanOrEqual(baseline) + }) + } +}) diff --git a/packages/query-engine-integrations/src/ai/ai-vendors.ts b/packages/query-engine-integrations/src/ai/ai-vendors.ts index 7aea6eebf..e6990cc0d 100644 --- a/packages/query-engine-integrations/src/ai/ai-vendors.ts +++ b/packages/query-engine-integrations/src/ai/ai-vendors.ts @@ -61,6 +61,7 @@ const vercelAiSdkIntegration: AiIntegration = { "gen_ai.usage.cache_read.input_tokens", "ai.usage.cachedInputTokens", "ai.usage.inputTokenDetails.cacheReadTokens", + "gen_ai.usage.input_tokens.cached", ], usageCacheCreationInputTokens: [ "gen_ai.usage.cache_creation.input_tokens", @@ -115,18 +116,38 @@ const openInferenceIntegration: AiIntegration = { sources: { requestModel: ["gen_ai.request.model", "llm.model_name"], providerName: ["gen_ai.provider.name", "gen_ai.system", "llm.provider", "llm.system"], - usageInputTokens: ["gen_ai.usage.input_tokens", "llm.token_count.prompt"], - usageOutputTokens: ["gen_ai.usage.output_tokens", "llm.token_count.completion"], + // Each list re-states the default's legacy aliases after the dialect keys. + // A vendor list REPLACES the default's, so omitting them would make a + // recognised vendor map strictly worse than an unrecognised one — an + // OpenInference span emitting old-semconv `gen_ai.usage.prompt_tokens` + // would lose its token counts precisely because we identified it. + usageInputTokens: [ + "gen_ai.usage.input_tokens", + "llm.token_count.prompt", + "gen_ai.usage.prompt_tokens", + ], + usageOutputTokens: [ + "gen_ai.usage.output_tokens", + "llm.token_count.completion", + "gen_ai.usage.completion_tokens", + ], usageCacheReadInputTokens: [ "gen_ai.usage.cache_read.input_tokens", "llm.token_count.prompt_details.cache_read", + "gen_ai.usage.input_tokens.cached", ], usageReasoningOutputTokens: [ "gen_ai.usage.reasoning.output_tokens", "llm.token_count.completion_details.reasoning", + "gen_ai.usage.output_tokens.reasoning", + ], + inputMessages: ["gen_ai.input.messages", "llm.input_messages", "input.value", "gen_ai.prompt"], + outputMessages: [ + "gen_ai.output.messages", + "llm.output_messages", + "output.value", + "gen_ai.completion", ], - inputMessages: ["gen_ai.input.messages", "llm.input_messages", "input.value"], - outputMessages: ["gen_ai.output.messages", "llm.output_messages", "output.value"], toolName: ["gen_ai.tool.name", "tool.name"], toolDescription: ["gen_ai.tool.description", "tool.description"], toolCallArguments: ["gen_ai.tool.call.arguments", "tool.parameters"], From d7f4212eec8123cd1d8f024e18d235b29285dc6f Mon Sep 17 00:00:00 2001 From: JeremyFunk Date: Wed, 19 Aug 2026 22:31:14 +0200 Subject: [PATCH 4/4] feat(query-engine): AI session facets query for the filter sidebar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `aiSessionFacetsQuery` returns distinct-session counts per vendor id and per service name as a two-branch UNION ALL over the detection scan alone — the `traces` read that finds session-bearing spans. No `trace_detail_spans` fan-out, which is the expensive half of `aiSessionListQuery`. That is exact rather than an approximation because the list applies both of its filters at the same detection level, so a facet describes precisely the population its filter selects. The counts are therefore ANY-span counts: a session belongs to every vendor and every service that any of its session-bearing spans carries, so a multi-vendor session is counted under each and the branch totals exceed the session count. `uniqExact` rather than the builder's approximate `uniq`: a facet count sits next to the list it filters and has to agree with it. Co-Authored-By: Claude Fable 5 --- .../src/__sql_baseline__/integrations.sql | 30 ++++++++ .../src/ai/ai-sessions.test.ts | 71 ++++++++++++++++++- .../src/ai/ai-sessions.ts | 70 ++++++++++++++++++ .../query-engine-integrations/src/ai/index.ts | 3 + .../query-engine-integrations/src/catalog.ts | 7 ++ 5 files changed, 180 insertions(+), 1 deletion(-) diff --git a/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql b/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql index eb2ce123e..bd41f3190 100644 --- a/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql +++ b/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql @@ -1,3 +1,33 @@ +-- builder:ai-sessions:aiSessionFacetsQuery:default +SELECT + SpanAttributes['maple_ai.vendor.id'] AS name, + uniqExact(SpanAttributes['maple_ai.session.id']) AS count, + 'vendor' AS facetType + FROM traces + WHERE OrgId = 'org_sql_catalog' + AND Timestamp >= '2026-01-01 10:30:00' + AND Timestamp <= '2026-01-03 14:15:00' + AND (mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != '') + AND SpanAttributes['maple_ai.vendor.id'] != '' + GROUP BY name + ORDER BY count DESC + LIMIT 50 +UNION ALL +SELECT + ServiceName AS name, + uniqExact(SpanAttributes['maple_ai.session.id']) AS count, + 'service' AS facetType + FROM traces + WHERE OrgId = 'org_sql_catalog' + AND Timestamp >= '2026-01-01 10:30:00' + AND Timestamp <= '2026-01-03 14:15:00' + AND (mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != '') + AND ServiceName != '' + GROUP BY name + ORDER BY count DESC + LIMIT 50 +FORMAT JSON + -- builder:ai-sessions:aiSessionListQuery:default SELECT sessionId AS sessionId, diff --git a/packages/query-engine-integrations/src/ai/ai-sessions.test.ts b/packages/query-engine-integrations/src/ai/ai-sessions.test.ts index 3fa2056b4..bcd57740e 100644 --- a/packages/query-engine-integrations/src/ai/ai-sessions.test.ts +++ b/packages/query-engine-integrations/src/ai/ai-sessions.test.ts @@ -1,7 +1,9 @@ import { describe, expect, it } from "vitest" import { Effect } from "effect" -import { compileCH, type CompiledQuery } from "@maple-dev/clickhouse-builder" +import { compileCH, compileUnion, type CompiledQuery } from "@maple-dev/clickhouse-builder" import { + aiSessionFacetsQuery, + aiSessionFacetsRowSchema, aiSessionListQuery, aiSessionListRowSchema, aiSessionSpansQuery, @@ -135,6 +137,73 @@ describe("aiSessionListQuery", () => { }) }) +describe("aiSessionFacetsQuery", () => { + it("groups the detection scan only — no fan-out over trace_detail_spans", () => { + const { sql } = compileUnion(aiSessionFacetsQuery(), params) + + expect(sql).toContain("FROM traces") + expect(sql).not.toContain("trace_detail_spans") + expect(sql).not.toContain("TraceId IN (SELECT") + expect(sql).toContain("UNION ALL") + }) + + it("counts distinct sessions per vendor and per service", () => { + const { sql } = compileUnion(aiSessionFacetsQuery(), params) + + expect(sql).toContain("SpanAttributes['maple_ai.vendor.id'] AS name") + expect(sql).toContain("ServiceName AS name") + expect(sql).toContain("'vendor' AS facetType") + expect(sql).toContain("'service' AS facetType") + expect(sql.split("uniqExact(SpanAttributes['maple_ai.session.id']) AS count").length - 1).toBe(2) + expect(sql.split("GROUP BY name").length - 1).toBe(2) + expect(sql.split("ORDER BY count DESC").length - 1).toBe(2) + }) + + it("repeats the org and window predicates on every union branch", () => { + const { sql } = compileUnion(aiSessionFacetsQuery(), params) + + expect(orgPredicateCount(sql)).toBe(2) + expect(sql.split(`Timestamp >= '${params.startTime}'`).length - 1).toBe(2) + expect(sql.split(`Timestamp <= '${params.endTime}'`).length - 1).toBe(2) + }) + + it("is org-scoped", () => { + expect(compileUnion(aiSessionFacetsQuery(), params).tenantScope).toBe("org") + }) + + it("counts only session-bearing spans, and drops the blank option", () => { + const { sql } = compileUnion(aiSessionFacetsQuery(), params) + + expect( + sql.split( + "(mapContains(SpanAttributes, 'maple_ai.session.id') AND SpanAttributes['maple_ai.session.id'] != '')", + ).length - 1, + ).toBe(2) + expect(sql).toContain("SpanAttributes['maple_ai.vendor.id'] != ''") + expect(sql).toContain("ServiceName != ''") + }) + + it("leaves no unresolved param placeholder", () => { + expect(compileUnion(aiSessionFacetsQuery(), params).sql).not.toContain("__PARAM_") + }) + + it("decodes the quoted 64-bit uniqExact count", () => { + const compiled = compileUnion(aiSessionFacetsQuery(), params, { + rowSchema: aiSessionFacetsRowSchema, + }) + + expect( + decodeRows(compiled, [ + { name: "eve", count: "12", facetType: "vendor" }, + { name: "maple-slack-agent", count: 9, facetType: "service" }, + ]), + ).toEqual([ + { name: "eve", count: 12, facetType: "vendor" }, + { name: "maple-slack-agent", count: 9, facetType: "service" }, + ]) + }) +}) + describe("aiSessionSpansQuery", () => { it("returns every span of every trace in the session, oldest first", () => { const { sql } = compileCH(aiSessionSpansQuery(), spanParams) diff --git a/packages/query-engine-integrations/src/ai/ai-sessions.ts b/packages/query-engine-integrations/src/ai/ai-sessions.ts index d8450649c..1f46e94d7 100644 --- a/packages/query-engine-integrations/src/ai/ai-sessions.ts +++ b/packages/query-engine-integrations/src/ai/ai-sessions.ts @@ -49,6 +49,9 @@ import { fromQuery, inSubquery, param, + unionAll, + type CHUnionQuery, + type ColumnAccessor, type CompiledQueryRowSchema, } from "@maple-dev/clickhouse-builder" import { TraceDetailSpans, Traces } from "@maple/query-engine/ch/tables" @@ -84,6 +87,13 @@ const hasSessionId = (attrs: CH.Expr>, get: CH.Expr): CH.Expr => compileFnCall("fromUnixTimestamp64Nano", nanos) +/** + * Exact distinct count. Not in the builder's function set, which only carries + * the approximate `uniq`. A facet count sits next to the list it filters, so an + * HLL estimate that disagrees with the visible row count reads as a bug. + */ +const uniqExact = (expr: CH.Expr): CH.Expr => compileFnCall("uniqExact", expr) + export interface AiSessionListOpts { /** Sessions returned, most recently started first. */ readonly limit?: number @@ -233,6 +243,66 @@ export function aiSessionListQuery(opts: AiSessionListOpts = {}) { ) } +// List facets (UNION ALL — vendor / service) + +export interface AiSessionFacetsOutput { + readonly name: string + readonly count: number + readonly facetType: string +} + +export const aiSessionFacetsRowSchema: CompiledQueryRowSchema = Schema.Struct({ + name: Schema.String, + count: CHNumber, + facetType: Schema.String, +}) + +/** + * Distinct sessions per vendor and per service, for the list's filter sidebar. + * + * This is the detection scan of `aiSessionListQuery` and nothing else — no + * `trace_detail_spans` fan-out, which is the expensive half. It can be: both of + * the list's filters are applied at that level, so the population a facet + * describes is exactly the population its filter selects. + * + * That makes the counts ANY-span counts, matching the filter: a session belongs + * to every vendor and every service that ANY of its session-bearing spans + * carries, so a session whose turn spans came from two vendors is counted under + * both and the facet counts sum to more than the number of sessions. Picking one + * value returns exactly the count shown. + * + * `uniqExact` rather than `uniq`: session counts are small enough that the exact + * aggregate costs nothing, and the number has to agree with the list beside it. + */ +export function aiSessionFacetsQuery(): CHUnionQuery { + const facet = (facetType: string, name: ($: ColumnAccessor) => CH.Expr) => + from(Traces) + .select(($) => ({ + name: name($), + count: uniqExact($.SpanAttributes.get(SESSION_ID_ATTR)), + facetType: CH.lit(facetType), + })) + .where(($) => [ + // Every UNION ALL branch reads a table, so every branch carries the org + // predicate itself — see this file's header. + $.OrgId.eq(param.string("orgId")), + $.Timestamp.gte(param.dateTime("startTime")), + $.Timestamp.lte(param.dateTime("endTime")), + hasSessionId($.SpanAttributes, $.SpanAttributes.get(SESSION_ID_ATTR)), + // A span can be session-bearing without a vendor stamp; a blank option + // filters nothing and is not offered. + name($).neq(""), + ]) + .groupBy("name") + .orderBy(["count", "desc"]) + .limit(50) + + return unionAll( + facet("vendor", ($) => $.SpanAttributes.get(VENDOR_ID_ATTR)), + facet("service", ($) => $.ServiceName), + ).format("JSON") +} + export interface AiSessionSpansOpts { readonly limit?: number } diff --git a/packages/query-engine-integrations/src/ai/index.ts b/packages/query-engine-integrations/src/ai/index.ts index 7a5d322d9..d1b445212 100644 --- a/packages/query-engine-integrations/src/ai/index.ts +++ b/packages/query-engine-integrations/src/ai/index.ts @@ -7,11 +7,14 @@ // `AiSessionSpanRow`, so `mapAiSpans(rows)` is the whole read path. export { + aiSessionFacetsQuery, + aiSessionFacetsRowSchema, aiSessionListQuery, aiSessionListRowSchema, AI_SESSION_SPANS_MAX_SPANS, aiSessionSpansQuery, aiSessionSpansRowSchema, + type AiSessionFacetsOutput, type AiSessionListOpts, type AiSessionListOutput, type AiSessionSpansOpts, diff --git a/packages/query-engine-integrations/src/catalog.ts b/packages/query-engine-integrations/src/catalog.ts index dd16f78b1..1f6483313 100644 --- a/packages/query-engine-integrations/src/catalog.ts +++ b/packages/query-engine-integrations/src/catalog.ts @@ -51,6 +51,13 @@ export const integrationFixtures: ReadonlyArray = [ { rowSchema: CH.aiSessionListRowSchema }, ), }, + { + module: "ai-sessions", + name: "aiSessionFacetsQuery", + label: "default", + compile: () => + compileUnion(CH.aiSessionFacetsQuery(), window, { rowSchema: CH.aiSessionFacetsRowSchema }), + }, { module: "ai-sessions", name: "aiSessionSpansQuery",