From a0b8ea54c891f5ee485dd2e1d6349eba82b103f0 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 15:52:20 -0700 Subject: [PATCH 01/31] feat: add agent tracing lifecycle contract Signed-off-by: Adam Gurary --- packages/appkit/src/beta.ts | 10 +++ .../src/core/agent/consume-adapter-stream.ts | 18 +++- packages/appkit/src/core/agent/run-agent.ts | 24 ++++- .../tests/agent-tracing-public-api.test.ts | 63 ++++++++++++++ .../tests/consume-adapter-stream.test.ts | 87 +++++++++++++++++-- .../src/core/agent/tests/run-agent.test.ts | 56 ++++++++++++ packages/appkit/src/plugins/agents/agents.ts | 7 +- .../src/plugins/agents/event-translator.ts | 4 + .../agents/tests/event-translator.test.ts | 41 +++++++++ .../src/telemetry/agent-tracing/attributes.ts | 26 ++++++ .../src/telemetry/agent-tracing/index.ts | 7 ++ .../telemetry/agent-tracing/serialization.ts | 64 ++++++++++++++ .../agent-tracing/tests/serialization.test.ts | 52 +++++++++++ .../agent-tracing/tests/usage.test.ts | 84 ++++++++++++++++++ .../src/telemetry/agent-tracing/types.ts | 19 ++++ .../src/telemetry/agent-tracing/usage.ts | 46 ++++++++++ packages/shared/src/agent.ts | 47 +++++++++- 17 files changed, 639 insertions(+), 16 deletions(-) create mode 100644 packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/attributes.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/index.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/serialization.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/tests/usage.test.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/types.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/usage.ts diff --git a/packages/appkit/src/beta.ts b/packages/appkit/src/beta.ts index d94b4e0d4..0b0347ec7 100644 --- a/packages/appkit/src/beta.ts +++ b/packages/appkit/src/beta.ts @@ -10,8 +10,12 @@ export type { AgentAdapter, AgentEvent, AgentInput, + AgentModelEndEvent, + AgentModelStartEvent, + AgentRemoteTraceEvent, AgentRunContext, AgentToolDefinition, + AgentUsage, Message, Thread, ThreadStore, @@ -101,3 +105,9 @@ export type { SearchResult, } from "./plugins/ai-search/types"; export * from "./plugins/beta-exports.generated"; +export { + AgentUsageAccumulator, + type CapturedTraceValue, + type CaptureTraceValueOptions, + captureTraceValue, +} from "./telemetry/agent-tracing"; diff --git a/packages/appkit/src/core/agent/consume-adapter-stream.ts b/packages/appkit/src/core/agent/consume-adapter-stream.ts index c4f3d07ed..eb7168558 100644 --- a/packages/appkit/src/core/agent/consume-adapter-stream.ts +++ b/packages/appkit/src/core/agent/consume-adapter-stream.ts @@ -1,4 +1,8 @@ import type { AgentEvent } from "shared"; +import { + AgentUsageAccumulator, + type ConsumedAgentStream, +} from "../../telemetry/agent-tracing"; interface ConsumeAdapterStreamOptions { /** @@ -37,16 +41,26 @@ interface ConsumeAdapterStreamOptions { export async function consumeAdapterStream( stream: AsyncIterable, opts: ConsumeAdapterStreamOptions = {}, -): Promise { +): Promise { let text = ""; + const usage = new AgentUsageAccumulator(); + let remoteTrace: ConsumedAgentStream["remoteTrace"]; for await (const event of stream) { if (opts.signal?.aborted) break; if (event.type === "message_delta") { text += event.content; } else if (event.type === "message") { text = event.content; + } else if (event.type === "model_end") { + usage.add(event.usage); + } else if (event.type === "remote_trace") { + remoteTrace = event; } opts.onEvent?.(event); } - return text; + return { + text, + usage: usage.snapshot(), + ...(remoteTrace ? { remoteTrace } : {}), + }; } diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 5675edd36..1e07368c4 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -1,8 +1,10 @@ import { randomUUID } from "node:crypto"; +import { trace } from "@opentelemetry/api"; import type { AgentAdapter, AgentEvent, AgentToolDefinition, + AgentUsage, Message, PluginConstructor, PluginData, @@ -47,6 +49,10 @@ export interface RunAgentInput { * there is no HTTP request in standalone mode). */ plugins?: PluginData[]; + sessionId?: string; + userId?: string; + requestId?: string; + appName?: string; } export interface RunAgentResult { @@ -54,6 +60,8 @@ export interface RunAgentResult { text: string; /** Every event the adapter yielded, in order. Useful for inspection/tests. */ events: AgentEvent[]; + traceId: string; + usage: AgentUsage; } /** @@ -103,6 +111,9 @@ async function runAgentInternal( input: RunAgentInput, providerCache: Map, ): Promise { + const traceId = + trace.getActiveSpan()?.spanContext().traceId ?? + randomUUID().replaceAll("-", ""); const adapter = await resolveAdapter(def); const messages = normalizeMessages(input.messages, def.instructions); const toolIndex = buildStandaloneToolIndex( @@ -143,6 +154,10 @@ async function runAgentInternal( : JSON.stringify(args), signal, plugins: input.plugins, + sessionId: input.sessionId, + userId: input.userId, + requestId: input.requestId, + appName: input.appName, }; // Reuse the same `providerCache` so sub-agent plugin tools dispatch // through the same instances the parent constructed. @@ -184,14 +199,19 @@ async function runAgentInternal( // Shared accumulation rule (deltas append, `message` replaces). The // `events` array is filled via the `onEvent` side effect so callers that // inspect the raw stream still get the full record. - const text = await consumeAdapterStream(stream, { + const consumed = await consumeAdapterStream(stream, { signal, onEvent: (event) => { events.push(event); }, }); - return { text, events }; + return { + text: consumed.text, + events, + traceId, + usage: consumed.usage, + }; } /** diff --git a/packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts b/packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts new file mode 100644 index 000000000..b7c15d85e --- /dev/null +++ b/packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, test } from "vitest"; +import { + type AgentModelEndEvent, + type AgentModelStartEvent, + type AgentRemoteTraceEvent, + type AgentUsage, + AgentUsageAccumulator, + captureTraceValue, +} from "../../../beta"; + +describe("agent tracing public API", () => { + test("exposes lifecycle, capture, and usage primitives from the beta entrypoint", () => { + const usage: AgentUsage = { + inputTokens: 1, + outputTokens: 2, + totalTokens: 3, + costAvailable: false, + }; + const lifecycle: [ + AgentModelStartEvent, + AgentModelEndEvent, + AgentRemoteTraceEvent, + ] = [ + { + type: "model_start", + stepId: "step-1", + model: "model-a", + provider: "databricks", + input: "hello", + startedAt: 100, + }, + { + type: "model_end", + stepId: "step-1", + model: "model-a", + provider: "databricks", + output: "world", + usage, + streamDurationMs: 20, + endedAt: 120, + }, + { + type: "remote_trace", + traceId: "abcdef0123456789abcdef0123456789", + spanId: "0123456789abcdef", + source: "model-serving", + relation: "linked", + }, + ]; + const accumulator = new AgentUsageAccumulator(); + accumulator.add(usage); + + expect(lifecycle.map((event) => event.type)).toEqual([ + "model_start", + "model_end", + "remote_trace", + ]); + expect(captureTraceValue({ token: "secret" }).value).toBe( + '{"token":"[REDACTED]"}', + ); + expect(accumulator.snapshot()).toEqual(usage); + }); +}); diff --git a/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts b/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts index 98863a62a..e38a4edae 100644 --- a/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts +++ b/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts @@ -12,23 +12,23 @@ async function* streamOf( describe("consumeAdapterStream", () => { test("concatenates message_delta events into the final text", async () => { - const text = await consumeAdapterStream( + const result = await consumeAdapterStream( streamOf([ { type: "message_delta", content: "Hello " }, { type: "message_delta", content: "world" }, ]), ); - expect(text).toBe("Hello world"); + expect(result.text).toBe("Hello world"); }); test("a `message` event replaces whatever deltas arrived so far", async () => { - const text = await consumeAdapterStream( + const result = await consumeAdapterStream( streamOf([ { type: "message_delta", content: "partial" }, { type: "message", content: "final answer" }, ]), ); - expect(text).toBe("final answer"); + expect(result.text).toBe("final answer"); }); test("invokes onEvent once per event, in order, with the raw event", async () => { @@ -68,19 +68,90 @@ describe("consumeAdapterStream", () => { }); test("returns an empty string for a stream with no content events", async () => { - const text = await consumeAdapterStream( + const result = await consumeAdapterStream( streamOf([{ type: "thinking", content: "ā€¦" }]), ); - expect(text).toBe(""); + expect(result.text).toBe(""); }); test("works without a signal (standalone runAgent path)", async () => { - const text = await consumeAdapterStream( + const result = await consumeAdapterStream( streamOf([ { type: "message_delta", content: "x" }, { type: "message_delta", content: "y" }, ]), ); - expect(text).toBe("xy"); + expect(result.text).toBe("xy"); + }); + + test("retains lifecycle usage and the last remote trace without adding user-visible text", async () => { + const result = await consumeAdapterStream( + streamOf([ + { type: "message_delta", content: "answer" }, + { + type: "model_start", + stepId: "step-1", + model: "model-a", + provider: "databricks", + input: { prompt: "hello" }, + startedAt: 100, + }, + { + type: "model_end", + stepId: "step-1", + model: "model-a", + provider: "databricks", + output: { text: "first" }, + usage: { + inputTokens: 10, + outputTokens: 3, + totalTokens: 13, + costUsd: 0.02, + costAvailable: true, + }, + streamDurationMs: 20, + endedAt: 120, + }, + { + type: "remote_trace", + traceId: "abcdef0123456789abcdef0123456789", + spanId: "0123456789abcdef", + source: "model-serving", + relation: "linked", + }, + { + type: "model_end", + stepId: "step-2", + model: "model-b", + provider: "databricks", + output: { text: "second" }, + usage: { + inputTokens: 4, + outputTokens: 2, + totalTokens: 6, + costAvailable: false, + }, + streamDurationMs: 30, + endedAt: 150, + }, + ]), + ); + + expect(result).toEqual({ + text: "answer", + usage: { + inputTokens: 14, + outputTokens: 5, + totalTokens: 19, + costAvailable: false, + }, + remoteTrace: { + type: "remote_trace", + traceId: "abcdef0123456789abcdef0123456789", + spanId: "0123456789abcdef", + source: "model-serving", + relation: "linked", + }, + }); }); }); diff --git a/packages/appkit/src/core/agent/tests/run-agent.test.ts b/packages/appkit/src/core/agent/tests/run-agent.test.ts index efd3e4202..518300f5d 100644 --- a/packages/appkit/src/core/agent/tests/run-agent.test.ts +++ b/packages/appkit/src/core/agent/tests/run-agent.test.ts @@ -1,3 +1,4 @@ +import { type Span, trace } from "@opentelemetry/api"; import type { AgentAdapter, AgentEvent, @@ -42,6 +43,61 @@ describe("runAgent", () => { expect(result.events).toHaveLength(3); }); + test("returns the active trace and aggregate usage for identified runs", async () => { + const traceId = "0123456789abcdef0123456789abcdef"; + const activeSpan = { + spanContext: () => ({ + traceId, + spanId: "0123456789abcdef", + traceFlags: 1, + }), + } as unknown as Span; + const activeSpanSpy = vi + .spyOn(trace, "getActiveSpan") + .mockReturnValue(activeSpan); + const events: AgentEvent[] = [ + { type: "message_delta", content: "done" }, + { + type: "model_end", + stepId: "step-1", + model: "model-a", + provider: "databricks", + output: { text: "done" }, + usage: { + inputTokens: 7, + outputTokens: 2, + totalTokens: 9, + costUsd: 0.01, + costAvailable: true, + }, + streamDurationMs: 10, + endedAt: 110, + }, + ]; + const def = createAgent({ + instructions: "x", + model: scriptedAdapter(events), + }); + + const result = await runAgent(def, { + messages: "hi", + sessionId: "session-1", + userId: "user-1", + requestId: "request-1", + appName: "test-app", + }); + activeSpanSpy.mockRestore(); + + expect(result.traceId).toBe(traceId); + expect(result.usage).toEqual({ + inputTokens: 7, + outputTokens: 2, + totalTokens: 9, + costUsd: 0.01, + costAvailable: true, + }); + }); + test("prefers terminal 'message' event over deltas when present", async () => { const events: AgentEvent[] = [ { type: "message_delta", content: "partial" }, diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 0aefbfe59..a3cf0ea5c 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -1145,7 +1145,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { // The accumulation rule (deltas append, `message` replaces) is shared // with `runAgent` and `runSubAgent`; see `consumeAdapterStream` for // the rationale. - const fullContent = await consumeAdapterStream(stream, { + const { text: fullContent } = await consumeAdapterStream(stream, { signal, onEvent: (event) => { for (const translated of translator.translate(event)) { @@ -1308,7 +1308,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { { executeTool, signal }, ); - fullContent = await consumeAdapterStream(stream, { signal }); + ({ text: fullContent } = await consumeAdapterStream(stream, { signal })); if (fullContent) { await this.threadStore.addMessage(thread.id, userId, { @@ -1563,7 +1563,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { }, ]; - return consumeAdapterStream( + const consumed = await consumeAdapterStream( child.adapter.run( { messages, @@ -1594,6 +1594,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { }, }, ); + return consumed.text; } private async _handleCancel(req: express.Request, res: express.Response) { diff --git a/packages/appkit/src/plugins/agents/event-translator.ts b/packages/appkit/src/plugins/agents/event-translator.ts index 54d749fb0..41ecc7790 100644 --- a/packages/appkit/src/plugins/agents/event-translator.ts +++ b/packages/appkit/src/plugins/agents/event-translator.ts @@ -73,6 +73,10 @@ export class AgentEventTranslator { ]; case "status": return this.handleStatus(event.status, event.error); + case "model_start": + case "model_end": + case "remote_trace": + return []; } } diff --git a/packages/appkit/src/plugins/agents/tests/event-translator.test.ts b/packages/appkit/src/plugins/agents/tests/event-translator.test.ts index 050af001a..0e8d750d5 100644 --- a/packages/appkit/src/plugins/agents/tests/event-translator.test.ts +++ b/packages/appkit/src/plugins/agents/tests/event-translator.test.ts @@ -129,6 +129,47 @@ describe("AgentEventTranslator", () => { } }); + test("keeps model lifecycle and remote trace events off the user-visible SSE stream", () => { + const translator = new AgentEventTranslator(); + + expect( + translator.translate({ + type: "model_start", + stepId: "step-1", + model: "model-a", + provider: "databricks", + input: { prompt: "hello" }, + startedAt: 100, + }), + ).toEqual([]); + expect( + translator.translate({ + type: "model_end", + stepId: "step-1", + model: "model-a", + provider: "databricks", + output: { text: "hello" }, + usage: { + inputTokens: 2, + outputTokens: 1, + totalTokens: 3, + costAvailable: false, + }, + streamDurationMs: 20, + endedAt: 120, + }), + ).toEqual([]); + expect( + translator.translate({ + type: "remote_trace", + traceId: "abcdef0123456789abcdef0123456789", + spanId: "0123456789abcdef", + source: "model-serving", + relation: "linked", + }), + ).toEqual([]); + }); + test("status:complete triggers finalize with response.completed", () => { const translator = new AgentEventTranslator(); translator.translate({ type: "message_delta", content: "Hi" }); diff --git a/packages/appkit/src/telemetry/agent-tracing/attributes.ts b/packages/appkit/src/telemetry/agent-tracing/attributes.ts new file mode 100644 index 000000000..97721f0e7 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/attributes.ts @@ -0,0 +1,26 @@ +export const DEFAULT_TRACE_VALUE_MAX_BYTES = 64 * 1024; +export const REDACTED_TRACE_VALUE = "[REDACTED]"; + +export const DEFAULT_TRACE_REDACT_KEYS = [ + "authorization", + "proxy-authorization", + "cookie", + "set-cookie", + "api-key", + "api_key", + "apikey", + "x-api-key", + "token", + "access-token", + "access_token", + "refresh-token", + "refresh_token", + "databricks-token", + "databricks_token", + "password", + "secret", + "client-secret", + "client_secret", + "credential", + "credentials", +] as const; diff --git a/packages/appkit/src/telemetry/agent-tracing/index.ts b/packages/appkit/src/telemetry/agent-tracing/index.ts new file mode 100644 index 000000000..a6c39ca8b --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/index.ts @@ -0,0 +1,7 @@ +export { captureTraceValue } from "./serialization"; +export type { + CapturedTraceValue, + CaptureTraceValueOptions, + ConsumedAgentStream, +} from "./types"; +export { AgentUsageAccumulator } from "./usage"; diff --git a/packages/appkit/src/telemetry/agent-tracing/serialization.ts b/packages/appkit/src/telemetry/agent-tracing/serialization.ts new file mode 100644 index 000000000..4a48d491a --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/serialization.ts @@ -0,0 +1,64 @@ +import { createHash } from "node:crypto"; +import { + DEFAULT_TRACE_REDACT_KEYS, + DEFAULT_TRACE_VALUE_MAX_BYTES, + REDACTED_TRACE_VALUE, +} from "./attributes"; +import type { CapturedTraceValue, CaptureTraceValueOptions } from "./types"; + +export function captureTraceValue( + value: unknown, + options: CaptureTraceValueOptions = {}, +): CapturedTraceValue { + const redactKeys = new Set( + [...DEFAULT_TRACE_REDACT_KEYS, ...(options.redactKeys ?? [])].map((key) => + key.toLowerCase(), + ), + ); + const serialized = + JSON.stringify(value, (key, current) => { + if (redactKeys.has(key.toLowerCase())) return REDACTED_TRACE_VALUE; + if ( + current !== null && + typeof current === "object" && + !Array.isArray(current) + ) { + return Object.fromEntries( + Object.keys(current) + .sort() + .map((objectKey) => [ + objectKey, + (current as Record)[objectKey], + ]), + ); + } + return current; + }) ?? "null"; + + const encoded = Buffer.from(serialized, "utf8"); + const maxBytes = normalizeMaxBytes(options.maxBytes); + const retained = truncateUtf8(encoded, maxBytes); + + return { + value: retained.toString("utf8"), + originalBytes: encoded.length, + sha256: createHash("sha256").update(encoded).digest("hex"), + truncated: retained.length < encoded.length, + }; +} + +function normalizeMaxBytes(value: number | undefined): number { + if (value === undefined) return DEFAULT_TRACE_VALUE_MAX_BYTES; + if (!Number.isFinite(value) || value < 0) { + throw new RangeError("maxBytes must be a finite, non-negative number"); + } + return Math.floor(value); +} + +function truncateUtf8(value: Buffer, maxBytes: number): Buffer { + if (value.length <= maxBytes) return value; + + let end = maxBytes; + while (end > 0 && (value[end] & 0xc0) === 0x80) end -= 1; + return value.subarray(0, end); +} diff --git a/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts new file mode 100644 index 000000000..9d8d3fddd --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, test } from "vitest"; +import { captureTraceValue } from "../serialization"; + +describe("captureTraceValue", () => { + test("sorts object keys and redacts sensitive keys case-insensitively", () => { + expect( + captureTraceValue({ prompt: "hello", Authorization: "Bearer secret" }), + ).toEqual({ + value: '{"Authorization":"[REDACTED]","prompt":"hello"}', + originalBytes: 47, + sha256: + "d828c3e891ddb30e1c13de236bde9b7833fc29a279769125fc23353d235d9082", + truncated: false, + }); + }); + + test("redacts default and custom sensitive keys recursively", () => { + expect( + captureTraceValue( + { + nested: [{ keep: "c", ToKeN: "a", CustomSecret: "b" }], + }, + { redactKeys: ["customsecret"] }, + ), + ).toEqual({ + value: + '{"nested":[{"CustomSecret":"[REDACTED]","ToKeN":"[REDACTED]","keep":"c"}]}', + originalBytes: 74, + sha256: + "0c5ace5e27c98b82a34f392a38c54ab515ee16ae247ac4b4644ea278eeb1ef49", + truncated: false, + }); + }); + + test("hashes the complete value and truncates only at UTF-8 boundaries", () => { + expect(captureTraceValue("šŸ˜€a", { maxBytes: 5 })).toEqual({ + value: '"šŸ˜€', + originalBytes: 7, + sha256: + "5f83696371a62d8dac1f76186954ebb46376bb3fef359d15a44b5c5db0b51211", + truncated: true, + }); + }); + + test("retains up to 64 KiB by default", () => { + const result = captureTraceValue("a".repeat(70 * 1024)); + + expect(new TextEncoder().encode(result.value)).toHaveLength(64 * 1024); + expect(result.originalBytes).toBe(70 * 1024 + 2); + expect(result.truncated).toBe(true); + }); +}); diff --git a/packages/appkit/src/telemetry/agent-tracing/tests/usage.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/usage.test.ts new file mode 100644 index 000000000..8061ce156 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/tests/usage.test.ts @@ -0,0 +1,84 @@ +import { describe, expect, test } from "vitest"; +import { AgentUsageAccumulator } from "../usage"; + +describe("AgentUsageAccumulator", () => { + test("returns zero usage with unavailable cost before any model steps", () => { + expect(new AgentUsageAccumulator().snapshot()).toEqual({ + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }); + }); + + test("sums token, cache, and cost usage across priced model steps", () => { + const usage = new AgentUsageAccumulator(); + usage.add({ + inputTokens: 10, + outputTokens: 3, + totalTokens: 13, + cacheReadInputTokens: 4, + costUsd: 0.02, + costAvailable: true, + }); + usage.add({ + inputTokens: 4, + outputTokens: 2, + totalTokens: 6, + cacheCreationInputTokens: 2, + costUsd: 0.01, + costAvailable: true, + }); + + expect(usage.snapshot()).toEqual({ + inputTokens: 14, + outputTokens: 5, + totalTokens: 19, + cacheReadInputTokens: 4, + cacheCreationInputTokens: 2, + costUsd: 0.03, + costAvailable: true, + }); + }); + + test("omits partial aggregate cost when any model step is unpriced", () => { + const usage = new AgentUsageAccumulator(); + usage.add({ + inputTokens: 10, + outputTokens: 3, + totalTokens: 13, + costUsd: 0.02, + costAvailable: true, + }); + usage.add({ + inputTokens: 4, + outputTokens: 2, + totalTokens: 6, + costAvailable: false, + }); + + expect(usage.snapshot()).toEqual({ + inputTokens: 14, + outputTokens: 5, + totalTokens: 19, + costAvailable: false, + }); + }); + + test("omits cost when priced steps do not report a value", () => { + const usage = new AgentUsageAccumulator(); + usage.add({ + inputTokens: 1, + outputTokens: 2, + totalTokens: 3, + costAvailable: true, + }); + + expect(usage.snapshot()).toEqual({ + inputTokens: 1, + outputTokens: 2, + totalTokens: 3, + costAvailable: true, + }); + }); +}); diff --git a/packages/appkit/src/telemetry/agent-tracing/types.ts b/packages/appkit/src/telemetry/agent-tracing/types.ts new file mode 100644 index 000000000..1cac5f420 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/types.ts @@ -0,0 +1,19 @@ +import type { AgentRemoteTraceEvent, AgentUsage } from "shared"; + +export interface CapturedTraceValue { + value: string; + originalBytes: number; + sha256: string; + truncated: boolean; +} + +export interface CaptureTraceValueOptions { + maxBytes?: number; + redactKeys?: readonly string[]; +} + +export interface ConsumedAgentStream { + text: string; + usage: AgentUsage; + remoteTrace?: AgentRemoteTraceEvent; +} diff --git a/packages/appkit/src/telemetry/agent-tracing/usage.ts b/packages/appkit/src/telemetry/agent-tracing/usage.ts new file mode 100644 index 000000000..a66540ec5 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/usage.ts @@ -0,0 +1,46 @@ +import type { AgentUsage } from "shared"; + +export class AgentUsageAccumulator { + private modelSteps = 0; + private value: AgentUsage = { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: true, + }; + + add(next: AgentUsage): void { + this.modelSteps += 1; + this.value.inputTokens += next.inputTokens; + this.value.outputTokens += next.outputTokens; + this.value.totalTokens += next.totalTokens; + this.value.cacheReadInputTokens = addOptional( + this.value.cacheReadInputTokens, + next.cacheReadInputTokens, + ); + this.value.cacheCreationInputTokens = addOptional( + this.value.cacheCreationInputTokens, + next.cacheCreationInputTokens, + ); + this.value.costAvailable &&= next.costAvailable; + if (next.costUsd !== undefined) { + this.value.costUsd = (this.value.costUsd ?? 0) + next.costUsd; + } + } + + snapshot(): AgentUsage { + const costAvailable = this.modelSteps > 0 && this.value.costAvailable; + const { costUsd, ...usage } = this.value; + return { + ...usage, + ...(costAvailable && costUsd !== undefined ? { costUsd } : {}), + costAvailable, + }; + } +} + +function addOptional(left?: number, right?: number): number | undefined { + return left === undefined && right === undefined + ? undefined + : (left ?? 0) + (right ?? 0); +} diff --git a/packages/shared/src/agent.ts b/packages/shared/src/agent.ts index 5ec2caf35..97870648c 100644 --- a/packages/shared/src/agent.ts +++ b/packages/shared/src/agent.ts @@ -112,6 +112,48 @@ export interface ThreadStore { // Agent events (SSE protocol) // --------------------------------------------------------------------------- +export interface AgentUsage { + inputTokens: number; + outputTokens: number; + totalTokens: number; + cacheReadInputTokens?: number; + cacheCreationInputTokens?: number; + costUsd?: number; + costAvailable: boolean; +} + +export interface AgentModelStartEvent { + type: "model_start"; + stepId: string; + model: string; + provider: string; + input: unknown; + startedAt: number; +} + +export interface AgentModelEndEvent { + type: "model_end"; + stepId: string; + model: string; + provider: string; + output: unknown; + usage: AgentUsage; + finishReason?: string; + firstTokenAt?: number; + streamDurationMs: number; + endedAt: number; + error?: string; +} + +export interface AgentRemoteTraceEvent { + type: "remote_trace"; + traceId: string; + /** Required when relation is "linked"; omitted only for continued context. */ + spanId?: string; + source: "model-serving" | "supervisor" | "remote-agent"; + relation: "continued" | "linked"; +} + export type AgentEvent = | { type: "message_delta"; content: string } | { type: "message"; content: string } @@ -144,7 +186,10 @@ export type AgentEvent = toolName: string; args: unknown; annotations?: ToolAnnotations; - }; + } + | AgentModelStartEvent + | AgentModelEndEvent + | AgentRemoteTraceEvent; // --------------------------------------------------------------------------- // Responses API types (OpenAI-compatible wire format for HTTP boundary) From 484f7b8b22abb1334dd8e5bc972155a0bde0cdd3 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 16:00:37 -0700 Subject: [PATCH 02/31] fix: require span ID for linked remote traces Signed-off-by: Adam Gurary --- .../tests/agent-tracing-public-api.test.ts | 12 ++++++++++ packages/shared/src/agent.ts | 23 ++++++++++++------- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts b/packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts index b7c15d85e..bbbf41352 100644 --- a/packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts +++ b/packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts @@ -9,6 +9,18 @@ import { } from "../../../beta"; describe("agent tracing public API", () => { + test("requires spanId for linked remote traces", () => { + // @ts-expect-error linked remote traces require a spanId + const invalidLinkedTrace: AgentRemoteTraceEvent = { + type: "remote_trace", + traceId: "abcdef0123456789abcdef0123456789", + source: "model-serving", + relation: "linked", + }; + + expect(invalidLinkedTrace.relation).toBe("linked"); + }); + test("exposes lifecycle, capture, and usage primitives from the beta entrypoint", () => { const usage: AgentUsage = { inputTokens: 1, diff --git a/packages/shared/src/agent.ts b/packages/shared/src/agent.ts index 97870648c..d99516cdd 100644 --- a/packages/shared/src/agent.ts +++ b/packages/shared/src/agent.ts @@ -145,14 +145,21 @@ export interface AgentModelEndEvent { error?: string; } -export interface AgentRemoteTraceEvent { - type: "remote_trace"; - traceId: string; - /** Required when relation is "linked"; omitted only for continued context. */ - spanId?: string; - source: "model-serving" | "supervisor" | "remote-agent"; - relation: "continued" | "linked"; -} +export type AgentRemoteTraceEvent = + | { + type: "remote_trace"; + traceId: string; + spanId?: string; + source: "model-serving" | "supervisor" | "remote-agent"; + relation: "continued"; + } + | { + type: "remote_trace"; + traceId: string; + spanId: string; + source: "model-serving" | "supervisor" | "remote-agent"; + relation: "linked"; + }; export type AgentEvent = | { type: "message_delta"; content: string } From 7c9c767f1657188271416a9d3b4a7ce87fc76448 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 17:08:45 -0700 Subject: [PATCH 03/31] feat: export AppKit agent traces to MLflow UC Signed-off-by: Adam Gurary --- packages/appkit/package.json | 1 + packages/appkit/src/core/appkit.ts | 20 +- packages/appkit/src/telemetry/index.ts | 9 + .../appkit/src/telemetry/mlflow-uc/config.ts | 38 +++ .../src/telemetry/mlflow-uc/exporter.ts | 176 ++++++++++ .../appkit/src/telemetry/mlflow-uc/index.ts | 12 + .../src/telemetry/mlflow-uc/processor.ts | 120 +++++++ .../telemetry/mlflow-uc/tests/config.test.ts | 56 +++ .../mlflow-uc/tests/exporter.test.ts | 318 ++++++++++++++++++ .../mlflow-uc/tests/processor.test.ts | 203 +++++++++++ .../src/telemetry/mlflow-uc/trace-info.ts | 227 +++++++++++++ .../appkit/src/telemetry/telemetry-manager.ts | 112 ++++-- .../telemetry/tests/telemetry-manager.test.ts | 90 ++++- packages/appkit/src/telemetry/types.ts | 3 + pnpm-lock.yaml | 5 +- 15 files changed, 1347 insertions(+), 43 deletions(-) create mode 100644 packages/appkit/src/telemetry/mlflow-uc/config.ts create mode 100644 packages/appkit/src/telemetry/mlflow-uc/exporter.ts create mode 100644 packages/appkit/src/telemetry/mlflow-uc/index.ts create mode 100644 packages/appkit/src/telemetry/mlflow-uc/processor.ts create mode 100644 packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts create mode 100644 packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts create mode 100644 packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts create mode 100644 packages/appkit/src/telemetry/mlflow-uc/trace-info.ts diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 450dd0167..33b9b61a4 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -64,6 +64,7 @@ "@databricks/sdk-experimental": "0.17.0", "@opentelemetry/api": "1.9.0", "@opentelemetry/api-logs": "0.219.0", + "@opentelemetry/core": "2.8.0", "@opentelemetry/auto-instrumentations-node": "0.77.0", "@opentelemetry/exporter-logs-otlp-proto": "0.219.0", "@opentelemetry/exporter-metrics-otlp-proto": "0.219.0", diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 201acf190..051e84f1e 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -191,12 +191,24 @@ export class AppKit { disableInternalTelemetry?: boolean; } = {}, ): Promise> { - // Initialize core services - TelemetryManager.initialize(config?.telemetry); - await CacheManager.getInstance(config?.cache); - const withDefaults = AppKit.withDefaultPlugins(config.plugins as T); const rawPlugins = AppKit.filterDevOnlyPlugins(withDefaults); + const agentsEnabled = rawPlugins.some( + (plugin) => plugin?.name === "agents", + ); + const requestedMlflowUc = config.telemetry?.mlflowUc; + + // Configuration is resolved before plugin construction or server startup. + // The agents plugin enables UC tracing by default; adjacent AppKit projects + // can opt in explicitly without installing that plugin. + await TelemetryManager.initialize( + { + ...config.telemetry, + mlflowUc: agentsEnabled ? requestedMlflowUc || true : requestedMlflowUc, + }, + config.client, + ); + await CacheManager.getInstance(config?.cache); // Collect manifest resources via registry const registry = new ResourceRegistry(); diff --git a/packages/appkit/src/telemetry/index.ts b/packages/appkit/src/telemetry/index.ts index 26877c0df..1500581a5 100644 --- a/packages/appkit/src/telemetry/index.ts +++ b/packages/appkit/src/telemetry/index.ts @@ -7,6 +7,15 @@ export { SpanKind, SpanStatusCode } from "@opentelemetry/api"; export { SeverityNumber } from "@opentelemetry/api-logs"; export { normalizeTelemetryOptions } from "./config"; export { instrumentations } from "./instrumentations"; +export { + constructMlflowV4TraceId, + getMlflowUcTraceId, + type MlflowUcConfig, + MlflowUcSpanExporter, + MlflowUcSpanProcessor, + MlflowUcTraceRegistry, + resolveMlflowUcConfig, +} from "./mlflow-uc"; export { TelemetryManager } from "./telemetry-manager"; export { TelemetryProvider } from "./telemetry-provider"; export type { diff --git a/packages/appkit/src/telemetry/mlflow-uc/config.ts b/packages/appkit/src/telemetry/mlflow-uc/config.ts new file mode 100644 index 000000000..49702578c --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/config.ts @@ -0,0 +1,38 @@ +export interface MlflowUcConfig { + experimentId: string; + catalogName: string; + schemaName: string; + tablePrefix: string; + otelSpansTableName: string; +} + +const ENV_FIELDS = { + experimentId: "MLFLOW_EXPERIMENT_ID", + catalogName: "MLFLOW_UC_CATALOG", + schemaName: "MLFLOW_UC_SCHEMA", + tablePrefix: "MLFLOW_UC_TABLE_PREFIX", + otelSpansTableName: "MLFLOW_OTEL_SPANS_TABLE", +} as const satisfies Record; + +export function resolveMlflowUcConfig( + env: NodeJS.ProcessEnv | Record, + overrides: Partial = {}, +): MlflowUcConfig { + const resolved = Object.fromEntries( + Object.entries(ENV_FIELDS).map(([field, envName]) => [ + field, + overrides[field as keyof MlflowUcConfig] ?? env[envName], + ]), + ) as unknown as MlflowUcConfig; + const missing = Object.entries(ENV_FIELDS) + .filter(([field]) => !resolved[field as keyof MlflowUcConfig]?.trim()) + .map(([, envName]) => envName); + + if (missing.length > 0) { + throw new Error( + `MLflow UC tracing configuration missing: ${missing.join(", ")}`, + ); + } + + return resolved; +} diff --git a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts new file mode 100644 index 000000000..a5a461233 --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts @@ -0,0 +1,176 @@ +import { type ExportResult, ExportResultCode } from "@opentelemetry/core"; +import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto"; +import type { ReadableSpan, SpanExporter } from "@opentelemetry/sdk-trace-base"; +import { createLogger } from "../../logging/logger"; +import { + createWorkspaceClient, + type WorkspaceClient, +} from "../../workspace-client"; +import type { MlflowUcConfig } from "./config"; +import { buildMlflowUcTraceInfo, type MlflowUcTraceInfo } from "./trace-info"; + +const logger = createLogger("telemetry:mlflow-uc"); + +export interface MlflowUcExportBatch { + traceInfo: MlflowUcTraceInfo; + spans: ReadableSpan[]; +} + +export interface MlflowUcTraceExporter { + exportTrace( + batch: MlflowUcExportBatch, + resultCallback: (result: ExportResult) => void, + ): void; + forceFlush(): Promise; + shutdown(): Promise; +} + +interface LoggerLike { + error(message: string, ...args: unknown[]): void; +} + +interface ExporterOptions { + createOtlpExporter?: (options: { + url: string; + headers: Record; + }) => SpanExporter; + logger?: LoggerLike; +} + +export class MlflowUcSpanExporter + implements SpanExporter, MlflowUcTraceExporter +{ + private readonly inFlight = new Set>(); + private readonly createOtlpExporter: NonNullable< + ExporterOptions["createOtlpExporter"] + >; + private readonly logger: LoggerLike; + private closed = false; + + constructor( + private readonly config: MlflowUcConfig, + private readonly client: WorkspaceClient = createWorkspaceClient(), + options: ExporterOptions = {}, + ) { + this.createOtlpExporter = + options.createOtlpExporter ?? + ((exporterOptions) => new OTLPTraceExporter(exporterOptions)); + this.logger = options.logger ?? logger; + } + + export( + spans: ReadableSpan[], + resultCallback: (result: ExportResult) => void, + ): void { + const semanticRoot = spans.find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + if (!semanticRoot) { + resultCallback({ code: ExportResultCode.SUCCESS }); + return; + } + this.exportTrace( + { + traceInfo: buildMlflowUcTraceInfo(this.config, semanticRoot, spans), + spans, + }, + resultCallback, + ); + } + + exportTrace( + batch: MlflowUcExportBatch, + resultCallback: (result: ExportResult) => void, + ): void { + if (this.closed) { + resultCallback({ code: ExportResultCode.SUCCESS }); + return; + } + + let operation!: Promise; + operation = this.exportBatch(batch) + .catch((error) => { + this.logger.error("MLflow UC trace export failed: %O", { + event: "mlflow_uc_trace_export_failed", + traceId: batch.traceInfo.trace_id, + error: error instanceof Error ? error.message : String(error), + }); + }) + .then(() => { + resultCallback({ code: ExportResultCode.SUCCESS }); + }) + .finally(() => { + this.inFlight.delete(operation); + }); + this.inFlight.add(operation); + } + + async forceFlush(): Promise { + while (this.inFlight.size > 0) { + await Promise.allSettled([...this.inFlight]); + } + } + + async shutdown(): Promise { + this.closed = true; + await this.forceFlush(); + } + + private async exportBatch(batch: MlflowUcExportBatch): Promise { + await this.client.config.ensureResolved(); + const configuredHost = this.client.config.host; + if (!configuredHost) { + throw new Error( + "Databricks workspace host is unavailable for MLflow UC export", + ); + } + const host = configuredHost.replace(/\/$/, ""); + const traceInfoHeaders = await this.freshAuthHeaders(); + const { traceInfo } = batch; + const location = `${this.config.catalogName}.${this.config.schemaName}.${this.config.tablePrefix}`; + const otelTraceId = traceInfo.trace_id.slice( + traceInfo.trace_id.lastIndexOf("/") + 1, + ); + const traceInfoResponse = await fetch( + `${host}/api/4.0/mlflow/traces/${encodeURIComponent(location)}/${encodeURIComponent(otelTraceId)}/info`, + { + method: "POST", + headers: { + ...traceInfoHeaders, + "Content-Type": "application/json", + }, + body: JSON.stringify(traceInfo), + }, + ); + if (!traceInfoResponse.ok) { + throw new Error( + `MLflow trace-info request failed with ${traceInfoResponse.status}: ${await traceInfoResponse.text()}`, + ); + } + + const otlpAuthHeaders = await this.freshAuthHeaders(); + const otlpExporter = this.createOtlpExporter({ + url: `${host}/api/2.0/otel/v1/traces`, + headers: { + ...otlpAuthHeaders, + "X-Databricks-UC-Table-Name": this.config.otelSpansTableName, + }, + }); + try { + await new Promise((resolve, reject) => { + otlpExporter.export(batch.spans, (result) => { + if (result.code === ExportResultCode.SUCCESS) resolve(); + else reject(result.error ?? new Error("OTLP trace upload failed")); + }); + }); + } finally { + await otlpExporter.shutdown(); + } + } + + private async freshAuthHeaders(): Promise> { + const headers = new Headers(); + await this.client.config.authenticate(headers); + return Object.fromEntries(headers.entries()); + } +} diff --git a/packages/appkit/src/telemetry/mlflow-uc/index.ts b/packages/appkit/src/telemetry/mlflow-uc/index.ts new file mode 100644 index 000000000..127e705d2 --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/index.ts @@ -0,0 +1,12 @@ +export { + type MlflowUcConfig, + resolveMlflowUcConfig, +} from "./config"; +export { MlflowUcSpanExporter } from "./exporter"; +export { MlflowUcSpanProcessor } from "./processor"; +export { + constructMlflowV4TraceId, + getMlflowUcTraceId, + MlflowUcTraceRegistry, + setActiveMlflowUcTraceRegistry, +} from "./trace-info"; diff --git a/packages/appkit/src/telemetry/mlflow-uc/processor.ts b/packages/appkit/src/telemetry/mlflow-uc/processor.ts new file mode 100644 index 000000000..7e29d9317 --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/processor.ts @@ -0,0 +1,120 @@ +import type { Context } from "@opentelemetry/api"; +import type { + ReadableSpan, + Span, + SpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import type { MlflowUcConfig } from "./config"; +import type { MlflowUcExportBatch, MlflowUcTraceExporter } from "./exporter"; +import { + buildMlflowUcTraceInfo, + MLFLOW_EXPERIMENT_ID_ATTRIBUTE, + MLFLOW_SPAN_TYPE_ATTRIBUTE, + MLFLOW_TRACE_REQUEST_ID_ATTRIBUTE, + MlflowUcTraceRegistry, +} from "./trace-info"; + +interface PendingTrace { + spans: ReadableSpan[]; + memberSpanIds: Set; +} + +export class MlflowUcSpanProcessor implements SpanProcessor { + private readonly pending = new Map(); + private readonly inFlight = new Set>(); + private closed = false; + + constructor( + private readonly config: MlflowUcConfig, + private readonly exporter: MlflowUcTraceExporter, + readonly registry = new MlflowUcTraceRegistry(config), + ) {} + + onStart(span: Span, _parentContext: Context): void { + if (this.closed) return; + const spanContext = span.spanContext(); + const otelTraceId = spanContext.traceId; + const mlflowTraceId = this.registry.ensureTrace(otelTraceId); + span.setAttribute(MLFLOW_TRACE_REQUEST_ID_ATTRIBUTE, mlflowTraceId); + span.setAttribute(MLFLOW_EXPERIMENT_ID_ATTRIBUTE, this.config.experimentId); + + let pending = this.pending.get(otelTraceId); + if (!pending) { + pending = { spans: [], memberSpanIds: new Set() }; + this.pending.set(otelTraceId, pending); + } + + if (span.attributes[MLFLOW_SPAN_TYPE_ATTRIBUTE] === "AGENT") { + if (this.registry.registerSemanticRoot(otelTraceId, spanContext.spanId)) { + pending.memberSpanIds.add(spanContext.spanId); + } else if ( + span.parentSpanContext && + pending.memberSpanIds.has(span.parentSpanContext.spanId) + ) { + pending.memberSpanIds.add(spanContext.spanId); + } + return; + } + + if ( + span.parentSpanContext && + pending.memberSpanIds.has(span.parentSpanContext.spanId) + ) { + pending.memberSpanIds.add(spanContext.spanId); + } + } + + onEnd(span: ReadableSpan): void { + if (this.closed) return; + const spanContext = span.spanContext(); + const otelTraceId = spanContext.traceId; + const pending = this.pending.get(otelTraceId); + if (!pending) return; + + const semanticRootSpanId = this.registry.getSemanticRootSpanId(otelTraceId); + if (!pending.memberSpanIds.has(spanContext.spanId)) { + if (!span.parentSpanContext && !semanticRootSpanId) { + this.pending.delete(otelTraceId); + } + return; + } + + pending.spans.push(span); + if (spanContext.spanId !== semanticRootSpanId) return; + + this.pending.delete(otelTraceId); + const batch: MlflowUcExportBatch = { + traceInfo: buildMlflowUcTraceInfo(this.config, span, pending.spans), + spans: pending.spans, + }; + this.startExport(batch); + } + + async forceFlush(): Promise { + await this.exporter.forceFlush(); + while (this.inFlight.size > 0) { + await Promise.allSettled([...this.inFlight]); + } + } + + async shutdown(): Promise { + if (this.closed) return; + this.closed = true; + await this.forceFlush(); + await this.exporter.shutdown(); + } + + private startExport(batch: MlflowUcExportBatch): void { + let resolveExport!: () => void; + const exportComplete = new Promise((resolve) => { + resolveExport = resolve; + }); + this.inFlight.add(exportComplete); + try { + this.exporter.exportTrace(batch, () => resolveExport()); + } catch { + resolveExport(); + } + void exportComplete.finally(() => this.inFlight.delete(exportComplete)); + } +} diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts new file mode 100644 index 000000000..981b05d4f --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts @@ -0,0 +1,56 @@ +import { describe, expect, test } from "vitest"; +import { constructMlflowV4TraceId, resolveMlflowUcConfig } from "../../index"; + +const completeEnv = { + MLFLOW_EXPERIMENT_ID: "experiment-123", + MLFLOW_UC_CATALOG: "main", + MLFLOW_UC_SCHEMA: "agent_traces", + MLFLOW_UC_TABLE_PREFIX: "appkit", + MLFLOW_OTEL_SPANS_TABLE: "main.agent_traces.appkit_otel_spans", +}; + +describe("resolveMlflowUcConfig", () => { + test("resolves every required field from the environment", () => { + expect(resolveMlflowUcConfig(completeEnv)).toEqual({ + experimentId: "experiment-123", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "appkit", + otelSpansTableName: "main.agent_traces.appkit_otel_spans", + }); + }); + + test("an explicit object overrides only supplied environment-backed fields", () => { + expect( + resolveMlflowUcConfig(completeEnv, { + experimentId: "experiment-override", + tablePrefix: "custom", + }), + ).toEqual({ + experimentId: "experiment-override", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "custom", + otelSpansTableName: "main.agent_traces.appkit_otel_spans", + }); + }); + + test("reports every missing or blank field in one startup error", () => { + expect(() => + resolveMlflowUcConfig({ + MLFLOW_EXPERIMENT_ID: "", + MLFLOW_UC_CATALOG: " ", + }), + ).toThrow( + "MLflow UC tracing configuration missing: MLFLOW_EXPERIMENT_ID, MLFLOW_UC_CATALOG, MLFLOW_UC_SCHEMA, MLFLOW_UC_TABLE_PREFIX, MLFLOW_OTEL_SPANS_TABLE", + ); + }); +}); + +test("constructMlflowV4TraceId uses the UC table-prefix location and OTel ID", () => { + const config = resolveMlflowUcConfig(completeEnv); + + expect( + constructMlflowV4TraceId(config, "0123456789abcdef0123456789abcdef"), + ).toBe("trace:/main.agent_traces.appkit/0123456789abcdef0123456789abcdef"); +}); diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts new file mode 100644 index 000000000..bac5d082d --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts @@ -0,0 +1,318 @@ +import { once } from "node:events"; +import { createServer, type ServerResponse } from "node:http"; +import type { AddressInfo } from "node:net"; +import { SpanStatusCode } from "@opentelemetry/api"; +import { ExportResultCode } from "@opentelemetry/core"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, + type SpanExporter, +} from "@opentelemetry/sdk-trace-base"; +import { afterEach, describe, expect, test, vi } from "vitest"; +import type { WorkspaceClient } from "../../../workspace-client"; +import { type MlflowUcConfig, MlflowUcSpanExporter } from "../../index"; + +const config: MlflowUcConfig = { + experimentId: "experiment-123", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "appkit", + otelSpansTableName: "main.agent_traces.appkit_otel_spans", +}; + +interface ObservedRequest { + path: string; + authorization?: string; + headers: Record; + body: Buffer; +} + +const cleanups: Array<() => Promise> = []; + +afterEach(async () => { + await Promise.all(cleanups.splice(0).map((cleanup) => cleanup())); +}); + +async function startBackend( + handler?: (request: ObservedRequest, response: ServerResponse) => void, +) { + const requests: ObservedRequest[] = []; + const server = createServer(async (request, response) => { + const chunks: Buffer[] = []; + for await (const chunk of request) chunks.push(Buffer.from(chunk)); + const observed = { + path: request.url ?? "", + authorization: request.headers.authorization, + headers: request.headers, + body: Buffer.concat(chunks), + }; + requests.push(observed); + if (handler) { + handler(observed, response); + return; + } + response.statusCode = 200; + response.end(); + }); + server.listen(0, "127.0.0.1"); + await once(server, "listening"); + const { port } = server.address() as AddressInfo; + cleanups.push( + () => + new Promise((resolve, reject) => + server.close((error) => (error ? reject(error) : resolve())), + ), + ); + return { host: `http://127.0.0.1:${port}`, requests }; +} + +function createClient(host: string, tokens: string[]): WorkspaceClient { + let tokenIndex = 0; + return { + config: { + host, + ensureResolved: vi.fn().mockResolvedValue(undefined), + authenticate: vi.fn(async (headers: Headers) => { + headers.set("Authorization", `Bearer ${tokens[tokenIndex++]}`); + }), + }, + } as unknown as WorkspaceClient; +} + +async function createTraceSpans(): Promise { + const inMemory = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(inMemory)], + }); + const tracer = provider.getTracer("mlflow-uc-exporter-test"); + const root = tracer.startSpan("support-agent", { + startTime: 1_700_000_000_000, + attributes: { + "mlflow.spanType": "AGENT", + "mlflow.spanInputs": '{"prompt":"hello"}', + "mlflow.spanOutputs": '{"answer":"world"}', + "mlflow.trace.session": "session-1", + "mlflow.trace.user": "user-1", + "mlflow.trace.tokenUsage": + '{"input_tokens":4,"output_tokens":2,"total_tokens":6}', + "appkit.app.name": "support-console", + "appkit.request.id": "request-1", + "appkit.thread.id": "thread-1", + "appkit.agent.name": "support-agent", + "appkit.route": "chat", + }, + }); + root.setStatus({ code: SpanStatusCode.OK }); + root.end(1_700_000_000_250); + await provider.forceFlush(); + const spans = inMemory.getFinishedSpans(); + await provider.shutdown(); + return spans; +} + +function exportSpans(exporter: MlflowUcSpanExporter, spans: ReadableSpan[]) { + return new Promise<{ code: ExportResultCode; error?: Error }>((resolve) => + exporter.export(spans, resolve), + ); +} + +describe("MlflowUcSpanExporter", () => { + test("authenticates each backend action and registers trace info before protobuf upload", async () => { + const { host, requests } = await startBackend(); + const client = createClient(host, ["token-1", "token-2"]); + const spans = await createTraceSpans(); + const otelTraceId = spans[0].spanContext().traceId; + const exporter = new MlflowUcSpanExporter(config, client); + + await expect(exportSpans(exporter, spans)).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + + expect(requests.map((request) => request.path)).toEqual([ + `/api/4.0/mlflow/traces/main.agent_traces.appkit/${otelTraceId}/info`, + "/api/2.0/otel/v1/traces", + ]); + expect(requests.map((request) => request.authorization)).toEqual([ + "Bearer token-1", + "Bearer token-2", + ]); + expect(requests[1].headers["x-databricks-uc-table-name"]).toBe( + "main.agent_traces.appkit_otel_spans", + ); + + const traceInfo = JSON.parse(requests[0].body.toString("utf8")); + expect(traceInfo).toEqual({ + trace_id: `trace:/main.agent_traces.appkit/${otelTraceId}`, + client_request_id: "request-1", + trace_location: { + type: "UC_TABLE_PREFIX", + uc_table_prefix: { + catalog_name: "main", + schema_name: "agent_traces", + table_prefix: "appkit", + otel_spans_table_name: "main.agent_traces.appkit_otel_spans", + }, + }, + request_preview: '{"prompt":"hello"}', + response_preview: '{"answer":"world"}', + request_time: "2023-11-14T22:13:20.000Z", + execution_duration: "0.25s", + state: "OK", + trace_metadata: { + "mlflow.trace_schema.version": "4", + "mlflow.experimentId": "experiment-123", + "mlflow.traceInputs": '{"prompt":"hello"}', + "mlflow.traceOutputs": '{"answer":"world"}', + "mlflow.trace.session": "session-1", + "mlflow.trace.user": "user-1", + "mlflow.trace.tokenUsage": + '{"input_tokens":4,"output_tokens":2,"total_tokens":6}', + "appkit.app.name": "support-console", + "appkit.request.id": "request-1", + "appkit.thread.id": "thread-1", + "appkit.agent.name": "support-agent", + "appkit.route": "chat", + }, + tags: { "mlflow.traceName": "support-agent" }, + assessments: [], + }); + }); + + test("passes the exact UC header name to a freshly authenticated OTLP exporter", async () => { + const { host, requests } = await startBackend(); + const client = createClient(host, ["token-1", "token-2"]); + const spans = await createTraceSpans(); + const exportCalls: Array<{ + url: string; + headers: Record; + spans: ReadableSpan[]; + }> = []; + const exporter = new MlflowUcSpanExporter(config, client, { + createOtlpExporter(options) { + return { + export(exportedSpans, callback) { + exportCalls.push({ ...options, spans: exportedSpans }); + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await exportSpans(exporter, spans); + + expect(requests.map((request) => request.path)).toEqual([ + expect.stringMatching(/^\/api\/4\.0\/mlflow\/traces\//), + ]); + expect(exportCalls).toHaveLength(1); + expect(exportCalls[0]).toMatchObject({ + url: `${host}/api/2.0/otel/v1/traces`, + headers: { + authorization: "Bearer token-2", + "X-Databricks-UC-Table-Name": "main.agent_traces.appkit_otel_spans", + }, + spans, + }); + }); + + test("isolates a backend rejection, calls back with success, and logs one structured error", async () => { + const { host } = await startBackend((_request, response) => { + response.statusCode = 500; + response.end("backend unavailable"); + }); + const logger = { error: vi.fn() }; + const client = createClient(host, ["token-1"]); + const spans = await createTraceSpans(); + const exporter = new MlflowUcSpanExporter(config, client, { logger }); + + await expect(exportSpans(exporter, spans)).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + expect(logger.error).toHaveBeenCalledTimes(1); + expect(logger.error).toHaveBeenCalledWith( + "MLflow UC trace export failed: %O", + expect.objectContaining({ + event: "mlflow_uc_trace_export_failed", + traceId: expect.stringMatching(/^trace:\/main\.agent_traces\.appkit\//), + error: "MLflow trace-info request failed with 500: backend unavailable", + }), + ); + }); + + test("isolates an unresolved workspace host with an actionable export error", async () => { + const logger = { error: vi.fn() }; + const client = createClient("", ["token-1"]); + const spans = await createTraceSpans(); + const exporter = new MlflowUcSpanExporter(config, client, { logger }); + + await expect(exportSpans(exporter, spans)).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + expect(logger.error).toHaveBeenCalledWith( + "MLflow UC trace export failed: %O", + expect.objectContaining({ + error: "Databricks workspace host is unavailable for MLflow UC export", + }), + ); + }); + + test("shutdown waits for in-flight trace-info and OTLP work", async () => { + let releaseTraceInfo!: () => void; + const traceInfoReleased = new Promise((resolve) => { + releaseTraceInfo = resolve; + }); + let traceInfoStarted!: () => void; + const traceInfoObserved = new Promise((resolve) => { + traceInfoStarted = resolve; + }); + const { host } = await startBackend(async (_request, response) => { + traceInfoStarted(); + await traceInfoReleased; + response.statusCode = 200; + response.end(); + }); + let finishOtlp!: () => void; + const otlpFinished = new Promise((resolve) => { + finishOtlp = resolve; + }); + let otlpStarted!: () => void; + const otlpObserved = new Promise((resolve) => { + otlpStarted = resolve; + }); + const client = createClient(host, ["token-1", "token-2"]); + const spans = await createTraceSpans(); + const exporter = new MlflowUcSpanExporter(config, client, { + createOtlpExporter() { + return { + export(_spans, callback) { + otlpStarted(); + void otlpFinished.then(() => + callback({ code: ExportResultCode.SUCCESS }), + ); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + exporter.export(spans, vi.fn()); + await traceInfoObserved; + let shutdownComplete = false; + const shutdown = exporter.shutdown().then(() => { + shutdownComplete = true; + }); + await Promise.resolve(); + expect(shutdownComplete).toBe(false); + + releaseTraceInfo(); + await otlpObserved; + await Promise.resolve(); + expect(shutdownComplete).toBe(false); + + finishOtlp(); + await shutdown; + expect(shutdownComplete).toBe(true); + }); +}); diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts new file mode 100644 index 000000000..ee291074c --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts @@ -0,0 +1,203 @@ +import { context, SpanStatusCode, trace } from "@opentelemetry/api"; +import { ExportResultCode } from "@opentelemetry/core"; +import { + BasicTracerProvider, + type ReadableSpan, +} from "@opentelemetry/sdk-trace-base"; +import { describe, expect, test, vi } from "vitest"; +import { + getMlflowUcTraceId, + type MlflowUcConfig, + MlflowUcSpanProcessor, + MlflowUcTraceRegistry, +} from "../../index"; +import type { MlflowUcExportBatch, MlflowUcTraceExporter } from "../exporter"; +import { setActiveMlflowUcTraceRegistry } from "../index"; + +const config: MlflowUcConfig = { + experimentId: "experiment-123", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "appkit", + otelSpansTableName: "main.agent_traces.appkit_otel_spans", +}; + +function collectingExporter( + onExport?: ( + batch: MlflowUcExportBatch, + callback: (result: { code: ExportResultCode; error?: Error }) => void, + ) => void, +) { + const batches: MlflowUcExportBatch[] = []; + const exporter: MlflowUcTraceExporter = { + exportTrace(batch, callback) { + batches.push(batch); + if (onExport) onExport(batch, callback); + else callback({ code: ExportResultCode.SUCCESS }); + }, + forceFlush: vi.fn().mockResolvedValue(undefined), + shutdown: vi.fn().mockResolvedValue(undefined), + }; + return { exporter, batches }; +} + +function startTree(processor: MlflowUcSpanProcessor) { + const provider = new BasicTracerProvider({ spanProcessors: [processor] }); + const tracer = provider.getTracer("mlflow-uc-processor-test"); + const http = tracer.startSpan("POST /api/agents/chat"); + const agent = tracer.startSpan( + "support-agent", + { + attributes: { + "mlflow.spanType": "AGENT", + "mlflow.spanInputs": '{"prompt":"hello"}', + "mlflow.spanOutputs": '{"answer":"world"}', + "mlflow.trace.session": "session-1", + "mlflow.trace.user": "user-1", + "appkit.app.name": "support-console", + "appkit.request.id": "request-1", + "appkit.thread.id": "thread-1", + "appkit.agent.name": "support-agent", + "appkit.route": "chat", + }, + }, + trace.setSpan(context.active(), http), + ); + const model = tracer.startSpan( + "model-step", + { + attributes: { + "mlflow.spanType": "CHAT_MODEL", + "mlflow.chat.tokenUsage": + '{"input_tokens":4,"output_tokens":2,"total_tokens":6,"cache_read_input_tokens":1}', + }, + }, + trace.setSpan(context.active(), agent), + ); + return { provider, http, agent, model }; +} + +describe("MlflowUcSpanProcessor", () => { + test("exposes the active registry's V4 trace ID to root tracing", () => { + const registry = new MlflowUcTraceRegistry(config); + const otelTraceId = "0123456789abcdef0123456789abcdef"; + setActiveMlflowUcTraceRegistry(registry); + try { + registry.ensureTrace(otelTraceId); + expect(getMlflowUcTraceId(otelTraceId)).toBe( + `trace:/main.agent_traces.appkit/${otelTraceId}`, + ); + } finally { + setActiveMlflowUcTraceRegistry(undefined); + } + }); + + test("uses the semantic AGENT as MLflow root and exports its complete tree only when it ends", async () => { + const registry = new MlflowUcTraceRegistry(config); + const { exporter, batches } = collectingExporter(); + const processor = new MlflowUcSpanProcessor(config, exporter, registry); + const { provider, http, agent, model } = startTree(processor); + const otelTraceId = agent.spanContext().traceId; + const mlflowTraceId = + `trace:/main.agent_traces.appkit/${otelTraceId}` as const; + + expect(registry.getMlflowTraceId(otelTraceId)).toBe(mlflowTraceId); + + model.end(); + expect(batches).toHaveLength(0); + + agent.setStatus({ code: SpanStatusCode.OK }); + agent.end(); + expect(batches).toHaveLength(1); + expect(batches[0].spans.map((span: ReadableSpan) => span.name)).toEqual([ + "model-step", + "support-agent", + ]); + expect(batches[0].traceInfo).toMatchObject({ + trace_id: mlflowTraceId, + client_request_id: "request-1", + request_preview: '{"prompt":"hello"}', + response_preview: '{"answer":"world"}', + state: "OK", + trace_metadata: { + "mlflow.trace_schema.version": "4", + "mlflow.trace.tokenUsage": + '{"input_tokens":4,"output_tokens":2,"total_tokens":6,"cache_read_input_tokens":1}', + "appkit.app.name": "support-console", + }, + }); + + const exportedRoot = batches[0].spans.find( + (span) => span.name === "support-agent", + ); + expect(exportedRoot?.parentSpanContext?.spanId).toBe( + http.spanContext().spanId, + ); + for (const span of batches[0].spans) { + expect(span.attributes["mlflow.traceRequestId"]).toBe(mlflowTraceId); + expect(span.attributes["mlflow.experimentId"]).toBe("experiment-123"); + } + + http.end(); + expect(batches).toHaveLength(1); + await provider.shutdown(); + }); + + test("forceFlush waits for the root export callback", async () => { + let finishExport!: () => void; + const { exporter } = collectingExporter((_batch, callback) => { + finishExport = () => callback({ code: ExportResultCode.SUCCESS }); + }); + const processor = new MlflowUcSpanProcessor(config, exporter); + const { provider, http, agent, model } = startTree(processor); + model.end(); + agent.end(); + + let flushed = false; + const forceFlush = processor.forceFlush().then(() => { + flushed = true; + }); + await Promise.resolve(); + expect(flushed).toBe(false); + + finishExport(); + await forceFlush; + expect(flushed).toBe(true); + + http.end(); + await provider.shutdown(); + }); + + test("keeps nested AGENT spans inside the first semantic root", async () => { + const { exporter, batches } = collectingExporter(); + const processor = new MlflowUcSpanProcessor(config, exporter); + const { provider, http, agent, model } = startTree(processor); + const tracer = provider.getTracer("nested-agent-test"); + const nestedAgent = tracer.startSpan( + "helper-agent", + { attributes: { "mlflow.spanType": "AGENT" } }, + trace.setSpan(context.active(), agent), + ); + const nestedTool = tracer.startSpan( + "helper-tool", + { attributes: { "mlflow.spanType": "TOOL" } }, + trace.setSpan(context.active(), nestedAgent), + ); + + nestedTool.end(); + nestedAgent.end(); + model.end(); + agent.end(); + + expect(batches).toHaveLength(1); + expect(batches[0].spans.map((span) => span.name)).toEqual([ + "helper-tool", + "helper-agent", + "model-step", + "support-agent", + ]); + + http.end(); + await provider.shutdown(); + }); +}); diff --git a/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts b/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts new file mode 100644 index 000000000..8c1b91ed0 --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts @@ -0,0 +1,227 @@ +import type { Attributes, HrTime } from "@opentelemetry/api"; +import { SpanStatusCode } from "@opentelemetry/api"; +import type { ReadableSpan } from "@opentelemetry/sdk-trace-base"; +import type { MlflowUcConfig } from "./config"; + +export const MLFLOW_TRACE_REQUEST_ID_ATTRIBUTE = "mlflow.traceRequestId"; +export const MLFLOW_EXPERIMENT_ID_ATTRIBUTE = "mlflow.experimentId"; +export const MLFLOW_SPAN_TYPE_ATTRIBUTE = "mlflow.spanType"; + +const TRACE_METADATA_IDENTITIES = [ + "mlflow.trace.session", + "mlflow.trace.user", + "appkit.app.name", + "appkit.request.id", + "appkit.thread.id", + "appkit.agent.name", + "appkit.route", +] as const; + +const MAX_REGISTRY_ENTRIES = 10_000; + +export interface MlflowUcTraceInfo { + trace_id: string; + client_request_id?: string; + trace_location: { + type: "UC_TABLE_PREFIX"; + uc_table_prefix: { + catalog_name: string; + schema_name: string; + table_prefix: string; + otel_spans_table_name: string; + }; + }; + request_preview?: string; + response_preview?: string; + request_time: string; + execution_duration: string; + state: "OK" | "ERROR"; + trace_metadata: Record; + tags: Record; + assessments: []; +} + +interface RegisteredTrace { + mlflowTraceId: string; + semanticRootSpanId?: string; +} + +/** + * Maps the process-local OTel trace identity to MLflow's V4 identity. + * + * The semantic root is registered independently from the OTel root. This is + * intentional: an auto-instrumented HTTP span can be the OTel parent while an + * AGENT span is the record represented by MLflow TraceInfo. Task-level tracing + * helpers can register or query that semantic root without changing provider + * ownership or exporter wiring. + */ +export class MlflowUcTraceRegistry { + private readonly traces = new Map(); + + constructor(private readonly config: MlflowUcConfig) {} + + ensureTrace(otelTraceId: string): string { + const existing = this.traces.get(otelTraceId); + if (existing) return existing.mlflowTraceId; + + if (this.traces.size >= MAX_REGISTRY_ENTRIES) { + const oldest = this.traces.keys().next().value; + if (oldest) this.traces.delete(oldest); + } + const mlflowTraceId = constructMlflowV4TraceId(this.config, otelTraceId); + this.traces.set(otelTraceId, { mlflowTraceId }); + return mlflowTraceId; + } + + registerSemanticRoot(otelTraceId: string, spanId: string): boolean { + const registered = this.getOrCreateTrace(otelTraceId); + if (registered.semanticRootSpanId) { + return registered.semanticRootSpanId === spanId; + } + registered.semanticRootSpanId = spanId; + return true; + } + + getMlflowTraceId(otelTraceId: string): string | undefined { + return this.traces.get(otelTraceId)?.mlflowTraceId; + } + + getSemanticRootSpanId(otelTraceId: string): string | undefined { + return this.traces.get(otelTraceId)?.semanticRootSpanId; + } + + private getOrCreateTrace(otelTraceId: string): RegisteredTrace { + this.ensureTrace(otelTraceId); + return ( + this.traces.get(otelTraceId) ?? { + mlflowTraceId: constructMlflowV4TraceId(this.config, otelTraceId), + } + ); + } +} + +let activeTraceRegistry: MlflowUcTraceRegistry | undefined; + +export function setActiveMlflowUcTraceRegistry( + registry: MlflowUcTraceRegistry | undefined, +): void { + activeTraceRegistry = registry; +} + +export function getMlflowUcTraceId(otelTraceId: string): string | undefined { + return activeTraceRegistry?.getMlflowTraceId(otelTraceId); +} + +export function constructMlflowV4TraceId( + config: MlflowUcConfig, + otelTraceId: string, +): string { + return `trace:/${config.catalogName}.${config.schemaName}.${config.tablePrefix}/${otelTraceId}`; +} + +export function buildMlflowUcTraceInfo( + config: MlflowUcConfig, + semanticRoot: ReadableSpan, + spans: ReadableSpan[], +): MlflowUcTraceInfo { + const otelTraceId = semanticRoot.spanContext().traceId; + const inputs = stringAttribute(semanticRoot.attributes, "mlflow.spanInputs"); + const outputs = stringAttribute( + semanticRoot.attributes, + "mlflow.spanOutputs", + ); + const traceMetadata: Record = { + "mlflow.trace_schema.version": "4", + [MLFLOW_EXPERIMENT_ID_ATTRIBUTE]: config.experimentId, + }; + + if (inputs !== undefined) traceMetadata["mlflow.traceInputs"] = inputs; + if (outputs !== undefined) traceMetadata["mlflow.traceOutputs"] = outputs; + for (const key of TRACE_METADATA_IDENTITIES) { + const value = stringAttribute(semanticRoot.attributes, key); + if (value !== undefined) traceMetadata[key] = value; + } + + const usage = aggregateUsage(semanticRoot, spans); + if (usage) traceMetadata["mlflow.trace.tokenUsage"] = JSON.stringify(usage); + const cost = stringAttribute(semanticRoot.attributes, "mlflow.trace.cost"); + if (cost !== undefined) traceMetadata["mlflow.trace.cost"] = cost; + + const requestTimeMs = hrTimeToMilliseconds(semanticRoot.startTime); + const durationMs = hrTimeToMilliseconds(semanticRoot.duration); + return { + trace_id: constructMlflowV4TraceId(config, otelTraceId), + client_request_id: stringAttribute( + semanticRoot.attributes, + "appkit.request.id", + ), + trace_location: { + type: "UC_TABLE_PREFIX", + uc_table_prefix: { + catalog_name: config.catalogName, + schema_name: config.schemaName, + table_prefix: config.tablePrefix, + otel_spans_table_name: config.otelSpansTableName, + }, + }, + request_preview: inputs, + response_preview: outputs, + request_time: new Date(requestTimeMs).toISOString(), + execution_duration: `${durationMs / 1000}s`, + state: semanticRoot.status.code === SpanStatusCode.ERROR ? "ERROR" : "OK", + trace_metadata: traceMetadata, + tags: { "mlflow.traceName": semanticRoot.name }, + assessments: [], + }; +} + +function stringAttribute( + attributes: Attributes, + key: string, +): string | undefined { + const value = attributes[key]; + return typeof value === "string" ? value : undefined; +} + +function aggregateUsage( + semanticRoot: ReadableSpan, + spans: ReadableSpan[], +): Record | undefined { + const rootUsage = parseUsage( + semanticRoot.attributes["mlflow.trace.tokenUsage"], + ); + if (rootUsage) return rootUsage; + + const total: Record = {}; + let found = false; + for (const span of spans) { + const usage = parseUsage(span.attributes["mlflow.chat.tokenUsage"]); + if (!usage) continue; + found = true; + for (const [key, value] of Object.entries(usage)) { + total[key] = (total[key] ?? 0) + value; + } + } + return found ? total : undefined; +} + +function parseUsage(value: unknown): Record | undefined { + if (typeof value !== "string") return undefined; + try { + const parsed = JSON.parse(value) as unknown; + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + return undefined; + } + const numeric = Object.entries(parsed).filter( + (entry): entry is [string, number] => + typeof entry[1] === "number" && Number.isFinite(entry[1]), + ); + return numeric.length > 0 ? Object.fromEntries(numeric) : undefined; + } catch { + return undefined; + } +} + +function hrTimeToMilliseconds([seconds, nanoseconds]: HrTime): number { + return seconds * 1000 + nanoseconds / 1_000_000; +} diff --git a/packages/appkit/src/telemetry/telemetry-manager.ts b/packages/appkit/src/telemetry/telemetry-manager.ts index a2642f2ab..9d3f00722 100644 --- a/packages/appkit/src/telemetry/telemetry-manager.ts +++ b/packages/appkit/src/telemetry/telemetry-manager.ts @@ -17,12 +17,24 @@ import { import { BatchLogRecordProcessor } from "@opentelemetry/sdk-logs"; import { PeriodicExportingMetricReader } from "@opentelemetry/sdk-metrics"; import { NodeSDK } from "@opentelemetry/sdk-node"; +import { + BatchSpanProcessor, + type SpanProcessor, +} from "@opentelemetry/sdk-trace-base"; import { ATTR_SERVICE_NAME, ATTR_SERVICE_VERSION, } from "@opentelemetry/semantic-conventions"; import type { TelemetryOptions } from "shared"; import { createLogger } from "../logging/logger"; +import type { WorkspaceClient } from "../workspace-client"; +import { + MlflowUcSpanExporter, + MlflowUcSpanProcessor, + MlflowUcTraceRegistry, + resolveMlflowUcConfig, + setActiveMlflowUcTraceRegistry, +} from "./mlflow-uc"; import { TelemetryProvider } from "./telemetry-provider"; import { AppKitSampler } from "./trace-sampler"; import type { TelemetryConfig } from "./types"; @@ -61,45 +73,80 @@ export class TelemetryManager { return TelemetryManager.instance; } - static initialize(config: Partial = {}): void { + static async initialize( + config: Partial = {}, + workspaceClient?: WorkspaceClient, + ): Promise { const instance = TelemetryManager.getInstance(); - instance._initialize(config); + await instance._initialize(config, workspaceClient); } - private _initialize(config: Partial): void { + private async _initialize( + config: Partial, + workspaceClient?: WorkspaceClient, + ): Promise { if (this.sdk) return; - if (!process.env.OTEL_EXPORTER_OTLP_ENDPOINT) { - return; + const genericOtlpEnabled = Boolean(process.env.OTEL_EXPORTER_OTLP_ENDPOINT); + const mlflowUcConfig = config.mlflowUc + ? resolveMlflowUcConfig( + process.env, + config.mlflowUc === true ? {} : config.mlflowUc, + ) + : undefined; + if (!genericOtlpEnabled && !mlflowUcConfig) return; + + const spanProcessors: SpanProcessor[] = []; + if (genericOtlpEnabled) { + spanProcessors.push( + new BatchSpanProcessor( + new OTLPTraceExporter({ headers: config.headers }), + ), + ); } - - try { - this.sdk = new NodeSDK({ - resource: this.createResource(config), - autoDetectResources: false, - sampler: new AppKitSampler(), - traceExporter: new OTLPTraceExporter({ headers: config.headers }), - metricReaders: [ - new PeriodicExportingMetricReader({ - exporter: new OTLPMetricExporter({ headers: config.headers }), - exportIntervalMillis: - config.exportIntervalMs || - TelemetryManager.DEFAULT_EXPORT_INTERVAL_MS, - }), - ], - logRecordProcessors: [ - new BatchLogRecordProcessor( - new OTLPLogExporter({ headers: config.headers }), - ), - ], - instrumentations: this.getDefaultInstrumentations(), - }); - - this.sdk.start(); - logger.debug("Initialized successfully"); - } catch (error) { - logger.error("Failed to initialize: %O", error); + if (mlflowUcConfig) { + const registry = new MlflowUcTraceRegistry(mlflowUcConfig); + const exporter = new MlflowUcSpanExporter( + mlflowUcConfig, + workspaceClient, + ); + spanProcessors.push( + new MlflowUcSpanProcessor(mlflowUcConfig, exporter, registry), + ); + setActiveMlflowUcTraceRegistry(registry); } + + const sdk = new NodeSDK({ + resource: this.createResource(config), + autoDetectResources: false, + sampler: new AppKitSampler(), + spanProcessors, + metricReaders: genericOtlpEnabled + ? [ + new PeriodicExportingMetricReader({ + exporter: new OTLPMetricExporter({ headers: config.headers }), + exportIntervalMillis: + config.exportIntervalMs || + TelemetryManager.DEFAULT_EXPORT_INTERVAL_MS, + }), + ] + : [], + logRecordProcessors: genericOtlpEnabled + ? [ + new BatchLogRecordProcessor( + new OTLPLogExporter({ headers: config.headers }), + ), + ] + : [], + instrumentations: [ + ...this.getDefaultInstrumentations(), + ...(config.instrumentations ?? []), + ], + }); + + sdk.start(); + this.sdk = sdk; + logger.debug("Initialized successfully"); } /** @@ -170,6 +217,7 @@ export class TelemetryManager { if (this.sdk) { const sdk = this.sdk; this.sdk = undefined; + setActiveMlflowUcTraceRegistry(undefined); this.shutdownPromise = (async () => { try { await sdk.shutdown(); diff --git a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts index 11b85d9bf..7eae9c843 100644 --- a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts +++ b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts @@ -1,6 +1,18 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { TelemetryManager } from "../telemetry-manager"; +const { mockNodeSdkStart, mockNodeSdkShutdown } = vi.hoisted(() => ({ + mockNodeSdkStart: vi.fn(), + mockNodeSdkShutdown: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock("@opentelemetry/sdk-node", () => ({ + NodeSDK: vi.fn(function (this: Record) { + this.start = mockNodeSdkStart; + this.shutdown = mockNodeSdkShutdown; + }), +})); + // Mock only exporters to prevent network calls vi.mock("@opentelemetry/exporter-trace-otlp-proto", () => ({ OTLPTraceExporter: vi.fn(() => ({ @@ -52,6 +64,8 @@ describe("TelemetryManager", () => { beforeEach(() => { originalEnv = { ...process.env }; vi.clearAllMocks(); + mockNodeSdkStart.mockImplementation(() => undefined); + mockNodeSdkShutdown.mockResolvedValue(undefined); // @ts-expect-error - accessing private static property for testing TelemetryManager.instance = undefined; // @ts-expect-error - accessing private static property for testing @@ -75,7 +89,7 @@ describe("TelemetryManager", () => { const { detectResources } = await import("@opentelemetry/resources"); vi.clearAllMocks(); - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "test-service-config", }); @@ -85,7 +99,7 @@ describe("TelemetryManager", () => { test("should initialize providers and create telemetry instances", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "integration-test", serviceVersion: "1.0.0", }); @@ -118,7 +132,7 @@ describe("TelemetryManager", () => { test("should support disabled telemetry config", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = ""; - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "disabled-test", serviceVersion: "1.0.0", }); @@ -161,7 +175,7 @@ describe("TelemetryManager", () => { ); vi.clearAllMocks(); - TelemetryManager.initialize({ + await TelemetryManager.initialize({ headers: { Authorization: "Bearer token", "Custom-Header": "value", @@ -182,7 +196,7 @@ describe("TelemetryManager", () => { test("should create and execute spans with real tracer", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "span-test", serviceVersion: "1.0.0", }); @@ -207,7 +221,7 @@ describe("TelemetryManager", () => { test("should handle span errors", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; - TelemetryManager.initialize({ + await TelemetryManager.initialize({ serviceName: "error-test", serviceVersion: "1.0.0", }); @@ -224,4 +238,68 @@ describe("TelemetryManager", () => { ).rejects.toThrow("Test error in span"); }); }); + + test("starts the single AppKit provider for valid UC config without generic OTLP", async () => { + delete process.env.OTEL_EXPORTER_OTLP_ENDPOINT; + process.env.MLFLOW_EXPERIMENT_ID = "experiment-123"; + process.env.MLFLOW_UC_CATALOG = "main"; + process.env.MLFLOW_UC_SCHEMA = "agent_traces"; + process.env.MLFLOW_UC_TABLE_PREFIX = "appkit"; + process.env.MLFLOW_OTEL_SPANS_TABLE = "main.agent_traces.appkit_otel_spans"; + + const { NodeSDK } = await import("@opentelemetry/sdk-node"); + await TelemetryManager.initialize({ mlflowUc: true }); + + expect(NodeSDK).toHaveBeenCalledTimes(1); + expect(vi.mocked(NodeSDK).mock.calls[0][0]).toMatchObject({ + spanProcessors: [expect.anything()], + metricReaders: [], + logRecordProcessors: [], + }); + expect(mockNodeSdkStart).toHaveBeenCalledTimes(1); + }); + + test("coexists with generic OTLP on the same AppKit provider", async () => { + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; + process.env.MLFLOW_EXPERIMENT_ID = "experiment-123"; + process.env.MLFLOW_UC_CATALOG = "main"; + process.env.MLFLOW_UC_SCHEMA = "agent_traces"; + process.env.MLFLOW_UC_TABLE_PREFIX = "appkit"; + process.env.MLFLOW_OTEL_SPANS_TABLE = "main.agent_traces.appkit_otel_spans"; + + const { NodeSDK } = await import("@opentelemetry/sdk-node"); + await TelemetryManager.initialize({ mlflowUc: true }); + + expect(NodeSDK).toHaveBeenCalledTimes(1); + expect(vi.mocked(NodeSDK).mock.calls[0]?.[0]?.spanProcessors).toHaveLength( + 2, + ); + }); + + test("rejects missing UC configuration before starting the provider", async () => { + delete process.env.OTEL_EXPORTER_OTLP_ENDPOINT; + delete process.env.MLFLOW_EXPERIMENT_ID; + delete process.env.MLFLOW_UC_CATALOG; + delete process.env.MLFLOW_UC_SCHEMA; + delete process.env.MLFLOW_UC_TABLE_PREFIX; + delete process.env.MLFLOW_OTEL_SPANS_TABLE; + + const { NodeSDK } = await import("@opentelemetry/sdk-node"); + await expect( + TelemetryManager.initialize({ mlflowUc: true }), + ).rejects.toThrow("MLflow UC tracing configuration missing:"); + expect(NodeSDK).not.toHaveBeenCalled(); + expect(mockNodeSdkStart).not.toHaveBeenCalled(); + }); + + test("propagates provider startup failures", async () => { + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; + mockNodeSdkStart.mockImplementation(() => { + throw new Error("provider start failed"); + }); + + await expect(TelemetryManager.initialize()).rejects.toThrow( + "provider start failed", + ); + }); }); diff --git a/packages/appkit/src/telemetry/types.ts b/packages/appkit/src/telemetry/types.ts index abd73e4c0..008ae95dc 100644 --- a/packages/appkit/src/telemetry/types.ts +++ b/packages/appkit/src/telemetry/types.ts @@ -1,6 +1,7 @@ import type { Meter, Span, SpanOptions, Tracer } from "@opentelemetry/api"; import type { Logger, LogRecord } from "@opentelemetry/api-logs"; import type { Instrumentation } from "@opentelemetry/instrumentation"; +import type { MlflowUcConfig } from "./mlflow-uc"; /** OpenTelemetry configuration for AppKit applications */ export interface TelemetryConfig { @@ -9,6 +10,8 @@ export interface TelemetryConfig { instrumentations?: Instrumentation[]; exportIntervalMs?: number; headers?: Record; + /** Export agent traces to an MLflow experiment backed by Unity Catalog. */ + mlflowUc?: boolean | Partial; } /** diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 2e45b03b5..3e6785a89 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -266,6 +266,9 @@ importers: '@opentelemetry/auto-instrumentations-node': specifier: 0.77.0 version: 0.77.0(@opentelemetry/api@1.9.0)(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.0)) + '@opentelemetry/core': + specifier: 2.8.0 + version: 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/exporter-logs-otlp-proto': specifier: 0.219.0 version: 0.219.0(@opentelemetry/api@1.9.0) @@ -7379,7 +7382,7 @@ packages: git-raw-commits@4.0.0: resolution: {integrity: sha512-ICsMM1Wk8xSGMowkOmPrzo2Fgmfo4bMHLNX6ytHjajRJUqvHOw/TFapQ+QG75c3X/tTDDhOSRPGC52dDbNM8FQ==} engines: {node: '>=16'} - deprecated: This package is no longer maintained. For the JavaScript API, please use @conventional-changelog/git-client instead. + deprecated: Deprecated and no longer maintained. Use @conventional-changelog/git-client instead. hasBin: true git-up@8.1.1: From bda2f9f930bf8160fe290cec390c0934c301f611 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 17:24:14 -0700 Subject: [PATCH 04/31] fix: isolate and release MLflow UC traces Signed-off-by: Adam Gurary --- .../src/telemetry/mlflow-uc/exporter.ts | 65 +++++++++++-- .../src/telemetry/mlflow-uc/processor.ts | 29 ++++-- .../mlflow-uc/tests/exporter.test.ts | 92 ++++++++++++++++++- .../mlflow-uc/tests/processor.test.ts | 48 ++++++++++ .../src/telemetry/mlflow-uc/trace-info.ts | 14 +-- 5 files changed, 220 insertions(+), 28 deletions(-) diff --git a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts index a5a461233..cb2077d75 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts @@ -62,20 +62,36 @@ export class MlflowUcSpanExporter spans: ReadableSpan[], resultCallback: (result: ExportResult) => void, ): void { - const semanticRoot = spans.find( - (span) => span.attributes["mlflow.spanType"] === "AGENT", + const batches = [...groupSpansByTrace(spans).values()].flatMap( + (traceSpans): MlflowUcExportBatch[] => { + const semanticRoot = findSemanticRoot(traceSpans); + return semanticRoot + ? [ + { + traceInfo: buildMlflowUcTraceInfo( + this.config, + semanticRoot, + traceSpans, + ), + spans: traceSpans, + }, + ] + : []; + }, ); - if (!semanticRoot) { + if (batches.length === 0) { resultCallback({ code: ExportResultCode.SUCCESS }); return; } - this.exportTrace( - { - traceInfo: buildMlflowUcTraceInfo(this.config, semanticRoot, spans), - spans, - }, - resultCallback, - ); + + void (async () => { + for (const batch of batches) { + await new Promise((resolve) => { + this.exportTrace(batch, () => resolve()); + }); + } + resultCallback({ code: ExportResultCode.SUCCESS }); + })(); } exportTrace( @@ -174,3 +190,32 @@ export class MlflowUcSpanExporter return Object.fromEntries(headers.entries()); } } + +function groupSpansByTrace(spans: ReadableSpan[]): Map { + const grouped = new Map(); + for (const span of spans) { + const traceId = span.spanContext().traceId; + const traceSpans = grouped.get(traceId) ?? []; + traceSpans.push(span); + grouped.set(traceId, traceSpans); + } + return grouped; +} + +function findSemanticRoot(spans: ReadableSpan[]): ReadableSpan | undefined { + const bySpanId = new Map( + spans.map((span) => [span.spanContext().spanId, span]), + ); + return spans.find((span) => { + if (span.attributes["mlflow.spanType"] !== "AGENT") return false; + + let parent = span.parentSpanContext; + while (parent) { + const parentSpan = bySpanId.get(parent.spanId); + if (!parentSpan) break; + if (parentSpan.attributes["mlflow.spanType"] === "AGENT") return false; + parent = parentSpan.parentSpanContext; + } + return true; + }); +} diff --git a/packages/appkit/src/telemetry/mlflow-uc/processor.ts b/packages/appkit/src/telemetry/mlflow-uc/processor.ts index 7e29d9317..e169fae56 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/processor.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/processor.ts @@ -8,6 +8,7 @@ import type { MlflowUcConfig } from "./config"; import type { MlflowUcExportBatch, MlflowUcTraceExporter } from "./exporter"; import { buildMlflowUcTraceInfo, + constructMlflowV4TraceId, MLFLOW_EXPERIMENT_ID_ATTRIBUTE, MLFLOW_SPAN_TYPE_ATTRIBUTE, MLFLOW_TRACE_REQUEST_ID_ATTRIBUTE, @@ -34,18 +35,25 @@ export class MlflowUcSpanProcessor implements SpanProcessor { if (this.closed) return; const spanContext = span.spanContext(); const otelTraceId = spanContext.traceId; - const mlflowTraceId = this.registry.ensureTrace(otelTraceId); + const mlflowTraceId = + this.registry.getMlflowTraceId(otelTraceId) ?? + constructMlflowV4TraceId(this.config, otelTraceId); span.setAttribute(MLFLOW_TRACE_REQUEST_ID_ATTRIBUTE, mlflowTraceId); span.setAttribute(MLFLOW_EXPERIMENT_ID_ATTRIBUTE, this.config.experimentId); let pending = this.pending.get(otelTraceId); - if (!pending) { - pending = { spans: [], memberSpanIds: new Set() }; - this.pending.set(otelTraceId, pending); - } - if (span.attributes[MLFLOW_SPAN_TYPE_ATTRIBUTE] === "AGENT") { - if (this.registry.registerSemanticRoot(otelTraceId, spanContext.spanId)) { + const registeredRoot = this.registry.registerSemanticRoot( + otelTraceId, + spanContext.spanId, + ); + if (!pending && registeredRoot) { + pending = { spans: [], memberSpanIds: new Set() }; + this.pending.set(otelTraceId, pending); + } + if (!pending) return; + + if (registeredRoot) { pending.memberSpanIds.add(spanContext.spanId); } else if ( span.parentSpanContext && @@ -57,6 +65,7 @@ export class MlflowUcSpanProcessor implements SpanProcessor { } if ( + pending && span.parentSpanContext && pending.memberSpanIds.has(span.parentSpanContext.spanId) ) { @@ -73,9 +82,6 @@ export class MlflowUcSpanProcessor implements SpanProcessor { const semanticRootSpanId = this.registry.getSemanticRootSpanId(otelTraceId); if (!pending.memberSpanIds.has(spanContext.spanId)) { - if (!span.parentSpanContext && !semanticRootSpanId) { - this.pending.delete(otelTraceId); - } return; } @@ -83,6 +89,7 @@ export class MlflowUcSpanProcessor implements SpanProcessor { if (spanContext.spanId !== semanticRootSpanId) return; this.pending.delete(otelTraceId); + this.registry.deleteTrace(otelTraceId); const batch: MlflowUcExportBatch = { traceInfo: buildMlflowUcTraceInfo(this.config, span, pending.spans), spans: pending.spans, @@ -101,6 +108,8 @@ export class MlflowUcSpanProcessor implements SpanProcessor { if (this.closed) return; this.closed = true; await this.forceFlush(); + this.pending.clear(); + this.registry.clear(); await this.exporter.shutdown(); } diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts index bac5d082d..3eb0d408a 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts @@ -1,7 +1,7 @@ import { once } from "node:events"; import { createServer, type ServerResponse } from "node:http"; import type { AddressInfo } from "node:net"; -import { SpanStatusCode } from "@opentelemetry/api"; +import { context, SpanStatusCode, trace } from "@opentelemetry/api"; import { ExportResultCode } from "@opentelemetry/core"; import { BasicTracerProvider, @@ -81,7 +81,9 @@ function createClient(host: string, tokens: string[]): WorkspaceClient { } as unknown as WorkspaceClient; } -async function createTraceSpans(): Promise { +async function createTraceSpans( + options: { nestedAgent?: boolean } = {}, +): Promise { const inMemory = new InMemorySpanExporter(); const provider = new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(inMemory)], @@ -104,6 +106,17 @@ async function createTraceSpans(): Promise { "appkit.route": "chat", }, }); + if (options.nestedAgent) { + const nestedAgent = tracer.startSpan( + "helper-agent", + { + startTime: 1_700_000_000_100, + attributes: { "mlflow.spanType": "AGENT" }, + }, + trace.setSpan(context.active(), root), + ); + nestedAgent.end(1_700_000_000_200); + } root.setStatus({ code: SpanStatusCode.OK }); root.end(1_700_000_000_250); await provider.forceFlush(); @@ -119,6 +132,81 @@ function exportSpans(exporter: MlflowUcSpanExporter, spans: ReadableSpan[]) { } describe("MlflowUcSpanExporter", () => { + test("splits mixed OTel trace batches and registers each semantic root before its isolated upload", async () => { + const events: Array< + | { kind: "trace-info"; traceId: string; rootName: string } + | { kind: "otlp"; traceIds: string[] } + > = []; + const { host } = await startBackend((request, response) => { + const traceId = request.path.match(/\/([0-9a-f]{32})\/info$/)?.[1]; + const traceInfo = JSON.parse(request.body.toString("utf8")); + events.push({ + kind: "trace-info", + traceId: traceId ?? "missing", + rootName: traceInfo.tags["mlflow.traceName"], + }); + response.statusCode = 200; + response.end(); + }); + const client = createClient(host, [ + "token-1", + "token-2", + "token-3", + "token-4", + ]); + const firstTrace = await createTraceSpans({ nestedAgent: true }); + const secondTrace = await createTraceSpans(); + const firstTraceId = firstTrace[0].spanContext().traceId; + const secondTraceId = secondTrace[0].spanContext().traceId; + const exporter = new MlflowUcSpanExporter(config, client, { + createOtlpExporter() { + return { + export(spans, callback) { + events.push({ + kind: "otlp", + traceIds: [ + ...new Set(spans.map((span) => span.spanContext().traceId)), + ], + }); + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await expect( + exportSpans(exporter, [...firstTrace, ...secondTrace]), + ).resolves.toEqual({ code: ExportResultCode.SUCCESS }); + + expect(events.filter((event) => event.kind === "trace-info")).toEqual([ + { + kind: "trace-info", + traceId: firstTraceId, + rootName: "support-agent", + }, + { + kind: "trace-info", + traceId: secondTraceId, + rootName: "support-agent", + }, + ]); + expect(events.filter((event) => event.kind === "otlp")).toEqual([ + { kind: "otlp", traceIds: [firstTraceId] }, + { kind: "otlp", traceIds: [secondTraceId] }, + ]); + for (const traceId of [firstTraceId, secondTraceId]) { + const traceInfoIndex = events.findIndex( + (event) => event.kind === "trace-info" && event.traceId === traceId, + ); + const uploadIndex = events.findIndex( + (event) => event.kind === "otlp" && event.traceIds.includes(traceId), + ); + expect(traceInfoIndex).toBeGreaterThanOrEqual(0); + expect(uploadIndex).toBeGreaterThan(traceInfoIndex); + } + }); + test("authenticates each backend action and registers trace info before protobuf upload", async () => { const { host, requests } = await startBackend(); const client = createClient(host, ["token-1", "token-2"]); diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts index ee291074c..7b6b0dafa 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts @@ -143,6 +143,32 @@ describe("MlflowUcSpanProcessor", () => { await provider.shutdown(); }); + test("releases completed trace state when later non-root work ends", async () => { + const registry = new MlflowUcTraceRegistry(config); + const { exporter, batches } = collectingExporter(); + const processor = new MlflowUcSpanProcessor(config, exporter, registry); + const { provider, http, agent, model } = startTree(processor); + const otelTraceId = agent.spanContext().traceId; + + model.end(); + agent.end(); + + const tracer = provider.getTracer("late-non-root-test"); + const lateSpan = tracer.startSpan( + "late-http-work", + { attributes: { "mlflow.spanType": "CHAIN" } }, + trace.setSpan(context.active(), http), + ); + lateSpan.end(); + + expect(batches).toHaveLength(1); + expect(registry.getSemanticRootSpanId(otelTraceId)).toBeUndefined(); + expect(registry.getMlflowTraceId(otelTraceId)).toBeUndefined(); + + http.end(); + await provider.shutdown(); + }); + test("forceFlush waits for the root export callback", async () => { let finishExport!: () => void; const { exporter } = collectingExporter((_batch, callback) => { @@ -200,4 +226,26 @@ describe("MlflowUcSpanProcessor", () => { http.end(); await provider.shutdown(); }); + + test("exports every concurrently active semantic root beyond the former registry capacity", async () => { + const { exporter, batches } = collectingExporter(); + const processor = new MlflowUcSpanProcessor(config, exporter); + const provider = new BasicTracerProvider({ spanProcessors: [processor] }); + const tracer = provider.getTracer("registry-capacity-test"); + const roots = Array.from({ length: 10_001 }, (_, index) => + tracer.startSpan(`agent-${index}`, { + attributes: { "mlflow.spanType": "AGENT" }, + }), + ); + + for (const root of roots) root.end(); + await processor.forceFlush(); + + expect(batches).toHaveLength(10_001); + expect(new Set(batches.map((batch) => batch.traceInfo.trace_id)).size).toBe( + 10_001, + ); + + await provider.shutdown(); + }, 15_000); }); diff --git a/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts b/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts index 8c1b91ed0..90519796e 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts @@ -17,8 +17,6 @@ const TRACE_METADATA_IDENTITIES = [ "appkit.route", ] as const; -const MAX_REGISTRY_ENTRIES = 10_000; - export interface MlflowUcTraceInfo { trace_id: string; client_request_id?: string; @@ -64,10 +62,6 @@ export class MlflowUcTraceRegistry { const existing = this.traces.get(otelTraceId); if (existing) return existing.mlflowTraceId; - if (this.traces.size >= MAX_REGISTRY_ENTRIES) { - const oldest = this.traces.keys().next().value; - if (oldest) this.traces.delete(oldest); - } const mlflowTraceId = constructMlflowV4TraceId(this.config, otelTraceId); this.traces.set(otelTraceId, { mlflowTraceId }); return mlflowTraceId; @@ -90,6 +84,14 @@ export class MlflowUcTraceRegistry { return this.traces.get(otelTraceId)?.semanticRootSpanId; } + deleteTrace(otelTraceId: string): void { + this.traces.delete(otelTraceId); + } + + clear(): void { + this.traces.clear(); + } + private getOrCreateTrace(otelTraceId: string): RegisteredTrace { this.ensureTrace(otelTraceId); return ( From 605e1d4a337b2ba8332634f7460b9e2487ba5125 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 17:34:46 -0700 Subject: [PATCH 05/31] fix: buffer split MLflow UC export batches Signed-off-by: Adam Gurary --- .../src/telemetry/mlflow-uc/exporter.ts | 109 ++++++++---- .../mlflow-uc/tests/exporter.test.ts | 162 ++++++++++++++++++ 2 files changed, 233 insertions(+), 38 deletions(-) diff --git a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts index cb2077d75..e7cc92a21 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts @@ -41,6 +41,7 @@ export class MlflowUcSpanExporter implements SpanExporter, MlflowUcTraceExporter { private readonly inFlight = new Set>(); + private readonly pendingSpans = new Map>(); private readonly createOtlpExporter: NonNullable< ExporterOptions["createOtlpExporter"] >; @@ -62,36 +63,24 @@ export class MlflowUcSpanExporter spans: ReadableSpan[], resultCallback: (result: ExportResult) => void, ): void { - const batches = [...groupSpansByTrace(spans).values()].flatMap( - (traceSpans): MlflowUcExportBatch[] => { - const semanticRoot = findSemanticRoot(traceSpans); - return semanticRoot - ? [ - { - traceInfo: buildMlflowUcTraceInfo( - this.config, - semanticRoot, - traceSpans, - ), - spans: traceSpans, - }, - ] - : []; - }, - ); + if (this.closed) { + resultCallback({ code: ExportResultCode.SUCCESS }); + return; + } + + const batches = this.accumulateReadyBatches(spans); if (batches.length === 0) { resultCallback({ code: ExportResultCode.SUCCESS }); return; } - void (async () => { + const operation = (async () => { for (const batch of batches) { - await new Promise((resolve) => { - this.exportTrace(batch, () => resolve()); - }); + await this.exportBatchSafely(batch); } resultCallback({ code: ExportResultCode.SUCCESS }); })(); + this.track(operation); } exportTrace( @@ -103,22 +92,10 @@ export class MlflowUcSpanExporter return; } - let operation!: Promise; - operation = this.exportBatch(batch) - .catch((error) => { - this.logger.error("MLflow UC trace export failed: %O", { - event: "mlflow_uc_trace_export_failed", - traceId: batch.traceInfo.trace_id, - error: error instanceof Error ? error.message : String(error), - }); - }) - .then(() => { - resultCallback({ code: ExportResultCode.SUCCESS }); - }) - .finally(() => { - this.inFlight.delete(operation); - }); - this.inFlight.add(operation); + const operation = this.exportBatchSafely(batch).then(() => { + resultCallback({ code: ExportResultCode.SUCCESS }); + }); + this.track(operation); } async forceFlush(): Promise { @@ -130,6 +107,55 @@ export class MlflowUcSpanExporter async shutdown(): Promise { this.closed = true; await this.forceFlush(); + this.pendingSpans.clear(); + } + + private accumulateReadyBatches(spans: ReadableSpan[]): MlflowUcExportBatch[] { + const batches: MlflowUcExportBatch[] = []; + for (const [traceId, newSpans] of groupSpansByTrace(spans)) { + const accumulated = this.pendingSpans.get(traceId) ?? new Map(); + for (const span of newSpans) { + accumulated.set(span.spanContext().spanId, span); + } + this.pendingSpans.set(traceId, accumulated); + + const traceSpans = [...accumulated.values()]; + const semanticRoot = findSemanticRoot(traceSpans); + if (semanticRoot) { + this.pendingSpans.delete(traceId); + batches.push({ + traceInfo: buildMlflowUcTraceInfo( + this.config, + semanticRoot, + traceSpans, + ), + spans: traceSpans, + }); + } else if (isCompletedRootlessTrace(traceSpans)) { + this.pendingSpans.delete(traceId); + } + } + return batches; + } + + private async exportBatchSafely(batch: MlflowUcExportBatch): Promise { + try { + await this.exportBatch(batch); + } catch (error) { + this.logger.error("MLflow UC trace export failed: %O", { + event: "mlflow_uc_trace_export_failed", + traceId: batch.traceInfo.trace_id, + error: error instanceof Error ? error.message : String(error), + }); + } + } + + private track(operation: Promise): void { + this.inFlight.add(operation); + void operation.then( + () => this.inFlight.delete(operation), + () => this.inFlight.delete(operation), + ); } private async exportBatch(batch: MlflowUcExportBatch): Promise { @@ -212,10 +238,17 @@ function findSemanticRoot(spans: ReadableSpan[]): ReadableSpan | undefined { let parent = span.parentSpanContext; while (parent) { const parentSpan = bySpanId.get(parent.spanId); - if (!parentSpan) break; + if (!parentSpan) return false; if (parentSpan.attributes["mlflow.spanType"] === "AGENT") return false; parent = parentSpan.parentSpanContext; } return true; }); } + +function isCompletedRootlessTrace(spans: ReadableSpan[]): boolean { + return ( + spans.some((span) => !span.parentSpanContext) && + spans.every((span) => span.attributes["mlflow.spanType"] !== "AGENT") + ); +} diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts index 3eb0d408a..09120e07e 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts @@ -115,6 +115,15 @@ async function createTraceSpans( }, trace.setSpan(context.active(), root), ); + const nestedTool = tracer.startSpan( + "helper-tool", + { + startTime: 1_700_000_000_150, + attributes: { "mlflow.spanType": "TOOL" }, + }, + trace.setSpan(context.active(), nestedAgent), + ); + nestedTool.end(1_700_000_000_175); nestedAgent.end(1_700_000_000_200); } root.setStatus({ code: SpanStatusCode.OK }); @@ -207,6 +216,65 @@ describe("MlflowUcSpanExporter", () => { } }); + test("buffers split trace fragments until the top semantic AGENT arrives", async () => { + const traceInfoRoots: Array<{ traceId: string; rootName: string }> = []; + const uploads: Array<{ traceId: string; spanNames: string[] }> = []; + const { host } = await startBackend((request, response) => { + const traceId = request.path.match(/\/([0-9a-f]{32})\/info$/)?.[1]; + const traceInfo = JSON.parse(request.body.toString("utf8")); + traceInfoRoots.push({ + traceId: traceId ?? "missing", + rootName: traceInfo.tags["mlflow.traceName"], + }); + response.statusCode = 200; + response.end(); + }); + const client = createClient(host, [ + "token-1", + "token-2", + "token-3", + "token-4", + ]); + const completeTrace = await createTraceSpans({ nestedAgent: true }); + const semanticRoot = completeTrace.find( + (span) => span.name === "support-agent", + ); + const earlierFragments = completeTrace.filter( + (span) => span.name !== "support-agent", + ); + if (!semanticRoot) throw new Error("semantic root fixture is missing"); + const traceId = semanticRoot.spanContext().traceId; + const exporter = new MlflowUcSpanExporter(config, client, { + createOtlpExporter() { + return { + export(spans, callback) { + uploads.push({ + traceId: spans[0].spanContext().traceId, + spanNames: spans.map((span) => span.name), + }); + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await expect(exportSpans(exporter, earlierFragments)).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + await expect(exportSpans(exporter, [semanticRoot])).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + + expect(traceInfoRoots).toEqual([{ traceId, rootName: "support-agent" }]); + expect(uploads).toEqual([ + { + traceId, + spanNames: ["helper-tool", "helper-agent", "support-agent"], + }, + ]); + }); + test("authenticates each backend action and registers trace info before protobuf upload", async () => { const { host, requests } = await startBackend(); const client = createClient(host, ["token-1", "token-2"]); @@ -403,4 +471,98 @@ describe("MlflowUcSpanExporter", () => { await shutdown; expect(shutdownComplete).toBe(true); }); + + test("shutdown drains every trace accepted by an active multi-trace export driver", async () => { + const firstTrace = await createTraceSpans(); + const secondTrace = await createTraceSpans(); + const lateTrace = await createTraceSpans(); + const firstTraceId = firstTrace[0].spanContext().traceId; + const secondTraceId = secondTrace[0].spanContext().traceId; + const lateTraceId = lateTrace[0].spanContext().traceId; + const events: Array<{ kind: "trace-info" | "otlp"; traceId: string }> = []; + + let releaseFirstTraceInfo!: () => void; + const firstTraceInfoReleased = new Promise((resolve) => { + releaseFirstTraceInfo = resolve; + }); + let firstTraceInfoStarted!: () => void; + const firstTraceInfoObserved = new Promise((resolve) => { + firstTraceInfoStarted = resolve; + }); + let finishSecondUpload!: () => void; + const secondUploadFinished = new Promise((resolve) => { + finishSecondUpload = resolve; + }); + let secondUploadStarted!: () => void; + const secondUploadObserved = new Promise((resolve) => { + secondUploadStarted = resolve; + }); + + const { host } = await startBackend(async (request, response) => { + const traceId = request.path.match(/\/([0-9a-f]{32})\/info$/)?.[1]; + if (traceId) events.push({ kind: "trace-info", traceId }); + if (traceId === firstTraceId) { + firstTraceInfoStarted(); + await firstTraceInfoReleased; + } + response.statusCode = 200; + response.end(); + }); + const client = createClient(host, [ + "token-1", + "token-2", + "token-3", + "token-4", + ]); + const exporter = new MlflowUcSpanExporter(config, client, { + createOtlpExporter() { + return { + export(spans, callback) { + const traceId = spans[0].spanContext().traceId; + events.push({ kind: "otlp", traceId }); + if (traceId === secondTraceId) { + secondUploadStarted(); + void secondUploadFinished.then(() => + callback({ code: ExportResultCode.SUCCESS }), + ); + return; + } + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + const activeExport = exportSpans(exporter, [...firstTrace, ...secondTrace]); + await firstTraceInfoObserved; + let shutdownComplete = false; + const shutdown = exporter.shutdown().then(() => { + shutdownComplete = true; + }); + await expect(exportSpans(exporter, lateTrace)).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + await Promise.resolve(); + expect(shutdownComplete).toBe(false); + + releaseFirstTraceInfo(); + const secondAcceptedTraceStarted = await Promise.race([ + secondUploadObserved.then(() => true), + activeExport.then(() => false), + ]); + expect(secondAcceptedTraceStarted).toBe(true); + await Promise.resolve(); + expect(shutdownComplete).toBe(false); + + finishSecondUpload(); + await Promise.all([activeExport, shutdown]); + expect(events).toEqual([ + { kind: "trace-info", traceId: firstTraceId }, + { kind: "otlp", traceId: firstTraceId }, + { kind: "trace-info", traceId: secondTraceId }, + { kind: "otlp", traceId: secondTraceId }, + ]); + expect(events.some((event) => event.traceId === lateTraceId)).toBe(false); + }); }); From 4112f5a5b71aeaeb058ad8da38210260489b2846 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 17:44:36 -0700 Subject: [PATCH 06/31] fix: resolve MLflow UC exporter ancestry Signed-off-by: Adam Gurary --- .../src/telemetry/mlflow-uc/exporter.ts | 40 +++-- .../mlflow-uc/tests/exporter.test.ts | 148 +++++++++++++++++- 2 files changed, 178 insertions(+), 10 deletions(-) diff --git a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts index e7cc92a21..e2e58d530 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts @@ -1,3 +1,4 @@ +import { isSpanContextValid } from "@opentelemetry/api"; import { type ExportResult, ExportResultCode } from "@opentelemetry/core"; import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto"; import type { ReadableSpan, SpanExporter } from "@opentelemetry/sdk-trace-base"; @@ -122,17 +123,16 @@ export class MlflowUcSpanExporter const traceSpans = [...accumulated.values()]; const semanticRoot = findSemanticRoot(traceSpans); if (semanticRoot) { + const semanticSpans = findSemanticSubtree(traceSpans, semanticRoot); this.pendingSpans.delete(traceId); batches.push({ traceInfo: buildMlflowUcTraceInfo( this.config, semanticRoot, - traceSpans, + semanticSpans, ), - spans: traceSpans, + spans: semanticSpans, }); - } else if (isCompletedRootlessTrace(traceSpans)) { - this.pendingSpans.delete(traceId); } } return batches; @@ -237,8 +237,9 @@ function findSemanticRoot(spans: ReadableSpan[]): ReadableSpan | undefined { let parent = span.parentSpanContext; while (parent) { + if (!isSpanContextValid(parent)) return true; const parentSpan = bySpanId.get(parent.spanId); - if (!parentSpan) return false; + if (!parentSpan) return parent.isRemote === true; if (parentSpan.attributes["mlflow.spanType"] === "AGENT") return false; parent = parentSpan.parentSpanContext; } @@ -246,9 +247,30 @@ function findSemanticRoot(spans: ReadableSpan[]): ReadableSpan | undefined { }); } -function isCompletedRootlessTrace(spans: ReadableSpan[]): boolean { - return ( - spans.some((span) => !span.parentSpanContext) && - spans.every((span) => span.attributes["mlflow.spanType"] !== "AGENT") +function findSemanticSubtree( + spans: ReadableSpan[], + semanticRoot: ReadableSpan, +): ReadableSpan[] { + const bySpanId = new Map( + spans.map((span) => [span.spanContext().spanId, span]), ); + const semanticRootId = semanticRoot.spanContext().spanId; + + return spans.filter((span) => { + let current: ReadableSpan | undefined = span; + const visited = new Set(); + while (current) { + const spanId = current.spanContext().spanId; + if (spanId === semanticRootId) return true; + if (visited.has(spanId)) return false; + visited.add(spanId); + + const parentSpanContext: ReadableSpan["parentSpanContext"] = + current.parentSpanContext; + current = parentSpanContext + ? bySpanId.get(parentSpanContext.spanId) + : undefined; + } + return false; + }); } diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts index 09120e07e..b8d560e7d 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts @@ -1,7 +1,7 @@ import { once } from "node:events"; import { createServer, type ServerResponse } from "node:http"; import type { AddressInfo } from "node:net"; -import { context, SpanStatusCode, trace } from "@opentelemetry/api"; +import { context, SpanStatusCode, TraceFlags, trace } from "@opentelemetry/api"; import { ExportResultCode } from "@opentelemetry/core"; import { BasicTracerProvider, @@ -134,6 +134,93 @@ async function createTraceSpans( return spans; } +async function createRemoteParentTraceSpans(): Promise { + const inMemory = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(inMemory)], + }); + const tracer = provider.getTracer("mlflow-uc-remote-parent-test"); + const remoteParent = { + traceId: "11111111111111111111111111111111", + spanId: "2222222222222222", + traceFlags: TraceFlags.SAMPLED, + isRemote: true, + }; + const agent = tracer.startSpan( + "remote-child-agent", + { + startTime: 1_700_000_000_000, + attributes: { + "mlflow.spanType": "AGENT", + "mlflow.spanInputs": '{"prompt":"remote"}', + "mlflow.spanOutputs": '{"answer":"complete"}', + }, + }, + trace.setSpanContext(context.active(), remoteParent), + ); + agent.end(1_700_000_000_100); + await provider.forceFlush(); + const spans = inMemory.getFinishedSpans(); + await provider.shutdown(); + return spans; +} + +async function createHttpWrappedTraceSpans(): Promise<{ + http: ReadableSpan; + agent: ReadableSpan; + model: ReadableSpan; +}> { + const inMemory = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(inMemory)], + }); + const tracer = provider.getTracer("mlflow-uc-http-root-test"); + const http = tracer.startSpan("POST /api/agents/chat", { + startTime: 1_700_000_000_000, + }); + const agent = tracer.startSpan( + "support-agent", + { + startTime: 1_700_000_000_010, + attributes: { + "mlflow.spanType": "AGENT", + "mlflow.spanInputs": '{"prompt":"hello"}', + "mlflow.spanOutputs": '{"answer":"world"}', + }, + }, + trace.setSpan(context.active(), http), + ); + const model = tracer.startSpan( + "model-step", + { + startTime: 1_700_000_000_020, + attributes: { "mlflow.spanType": "CHAT_MODEL" }, + }, + trace.setSpan(context.active(), agent), + ); + + http.end(1_700_000_000_050); + model.end(1_700_000_000_100); + agent.end(1_700_000_000_150); + await provider.forceFlush(); + const spans = inMemory.getFinishedSpans(); + await provider.shutdown(); + + const readableHttp = spans.find( + (span) => span.name === "POST /api/agents/chat", + ); + const readableAgent = spans.find((span) => span.name === "support-agent"); + const readableModel = spans.find((span) => span.name === "model-step"); + if (!readableHttp || !readableAgent || !readableModel) { + throw new Error("HTTP-wrapped trace fixture is incomplete"); + } + return { + http: readableHttp, + agent: readableAgent, + model: readableModel, + }; +} + function exportSpans(exporter: MlflowUcSpanExporter, spans: ReadableSpan[]) { return new Promise<{ code: ExportResultCode; error?: Error }>((resolve) => exporter.export(spans, resolve), @@ -262,6 +349,8 @@ describe("MlflowUcSpanExporter", () => { await expect(exportSpans(exporter, earlierFragments)).resolves.toEqual({ code: ExportResultCode.SUCCESS, }); + expect(traceInfoRoots).toEqual([]); + expect(uploads).toEqual([]); await expect(exportSpans(exporter, [semanticRoot])).resolves.toEqual({ code: ExportResultCode.SUCCESS, }); @@ -275,6 +364,63 @@ describe("MlflowUcSpanExporter", () => { ]); }); + test("accepts a semantic AGENT whose missing parent context is remote", async () => { + const spans = await createRemoteParentTraceSpans(); + const agent = spans[0]; + expect(agent.parentSpanContext?.isRemote).toBe(true); + const { host, requests } = await startBackend(); + const client = createClient(host, ["token-1", "token-2"]); + const exporter = new MlflowUcSpanExporter(config, client); + + await expect(exportSpans(exporter, spans)).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + + expect(requests.map((request) => request.path)).toEqual([ + "/api/4.0/mlflow/traces/main.agent_traces.appkit/11111111111111111111111111111111/info", + "/api/2.0/otel/v1/traces", + ]); + }); + + test("retains an earlier local HTTP root and uploads only its later semantic AGENT subtree", async () => { + const { http, agent, model } = await createHttpWrappedTraceSpans(); + expect(agent.parentSpanContext?.spanId).toBe(http.spanContext().spanId); + expect(agent.parentSpanContext?.isRemote).not.toBe(true); + expect(model.parentSpanContext?.spanId).toBe(agent.spanContext().spanId); + const traceInfoRoots: string[] = []; + const uploads: string[][] = []; + const { host } = await startBackend((request, response) => { + const traceInfo = JSON.parse(request.body.toString("utf8")); + traceInfoRoots.push(traceInfo.tags["mlflow.traceName"]); + response.statusCode = 200; + response.end(); + }); + const client = createClient(host, ["token-1", "token-2"]); + const exporter = new MlflowUcSpanExporter(config, client, { + createOtlpExporter() { + return { + export(spans, callback) { + uploads.push(spans.map((span) => span.name)); + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await expect(exportSpans(exporter, [http])).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + expect(traceInfoRoots).toEqual([]); + expect(uploads).toEqual([]); + await expect(exportSpans(exporter, [model, agent])).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + + expect(traceInfoRoots).toEqual(["support-agent"]); + expect(uploads).toEqual([["model-step", "support-agent"]]); + }); + test("authenticates each backend action and registers trace info before protobuf upload", async () => { const { host, requests } = await startBackend(); const client = createClient(host, ["token-1", "token-2"]); From e7b05966e5eb9355e7ba46b90b85b40798038c57 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 18:02:25 -0700 Subject: [PATCH 07/31] feat: trace every AppKit model step Signed-off-by: Adam Gurary --- packages/appkit/src/agents/databricks.ts | 323 ++++++++-- .../src/agents/tests/databricks.test.ts | 551 +++++++++++++++++- .../appkit/src/connectors/serving/client.ts | 39 +- .../connectors/serving/tests/client.test.ts | 20 +- .../src/core/agent/consume-adapter-stream.ts | 41 +- .../tests/consume-adapter-stream.test.ts | 122 ++++ 6 files changed, 1043 insertions(+), 53 deletions(-) diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 476de0b85..d9332ec83 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -2,10 +2,14 @@ import type { AgentAdapter, AgentEvent, AgentInput, + AgentRemoteTraceEvent, AgentRunContext, AgentToolDefinition, + AgentUsage, } from "shared"; import { + getResponseHeaders, + retainResponseHeaders, type StreamBody, stream as servingStream, } from "../connectors/serving/client"; @@ -31,6 +35,149 @@ function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null; } +function finiteNonNegativeNumber( + record: Record | undefined, + ...keys: string[] +): number | undefined { + for (const key of keys) { + const value = record?.[key]; + if (typeof value === "number" && Number.isFinite(value) && value >= 0) { + return value; + } + } + return undefined; +} + +function normalizeUsage( + parsed: Record, +): AgentUsage | undefined { + const raw = isRecord(parsed.usage) ? parsed.usage : undefined; + if (!raw) return undefined; + + const inputTokens = + finiteNonNegativeNumber(raw, "input_tokens", "prompt_tokens") ?? 0; + const outputTokens = + finiteNonNegativeNumber(raw, "output_tokens", "completion_tokens") ?? 0; + const totalTokens = + finiteNonNegativeNumber(raw, "total_tokens") ?? inputTokens + outputTokens; + const details = isRecord(raw.input_tokens_details) + ? raw.input_tokens_details + : isRecord(raw.prompt_tokens_details) + ? raw.prompt_tokens_details + : undefined; + const cacheReadInputTokens = + finiteNonNegativeNumber(raw, "cache_read_input_tokens", "cached_tokens") ?? + finiteNonNegativeNumber( + details, + "cache_read_input_tokens", + "cached_tokens", + ); + const cacheCreationInputTokens = + finiteNonNegativeNumber(raw, "cache_creation_input_tokens") ?? + finiteNonNegativeNumber( + details, + "cache_creation_input_tokens", + "cache_creation_tokens", + ); + const providerCost = + finiteNonNegativeNumber(raw, "cost_usd", "cost", "total_cost_usd") ?? + finiteNonNegativeNumber(parsed, "cost_usd", "cost"); + + return { + inputTokens, + outputTokens, + totalTokens, + ...(cacheReadInputTokens !== undefined ? { cacheReadInputTokens } : {}), + ...(cacheCreationInputTokens !== undefined + ? { cacheCreationInputTokens } + : {}), + ...(providerCost !== undefined ? { costUsd: providerCost } : {}), + costAvailable: providerCost !== undefined, + }; +} + +function emptyUsage(): AgentUsage { + return { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }; +} + +function firstChoice( + parsed: Record, +): Record | undefined { + const choices = parsed.choices; + if (!Array.isArray(choices) || !isRecord(choices[0])) return undefined; + return choices[0]; +} + +function remoteTraceFromPayload( + parsed: Record, +): AgentRemoteTraceEvent | undefined { + const nested = isRecord(parsed.remote_trace) + ? parsed.remote_trace + : undefined; + const traceIdCandidates = [ + nested?.trace_id, + nested?.traceId, + parsed.mlflow_trace_id, + parsed.databricks_trace_id, + ]; + const traceId = traceIdCandidates.find( + (value): value is string => + typeof value === "string" && value.trim().length > 0, + ); + if (!traceId) return undefined; + + const spanIdCandidates = [ + nested?.span_id, + nested?.spanId, + parsed.mlflow_span_id, + parsed.databricks_span_id, + ]; + const spanId = spanIdCandidates.find( + (value): value is string => + typeof value === "string" && value.trim().length > 0, + ); + return spanId + ? { + type: "remote_trace", + traceId, + spanId, + source: "model-serving", + relation: "linked", + } + : { + type: "remote_trace", + traceId, + source: "model-serving", + relation: "continued", + }; +} + +function sanitizedModelError(error: unknown): string { + const raw = error instanceof Error ? error.message : "Model request failed"; + const withoutControls = Array.from(raw, (character) => { + const code = character.charCodeAt(0); + return code < 32 || code === 127 ? " " : character; + }).join(""); + const singleLine = withoutControls.replace(/\s+/g, " ").trim(); + return (singleLine || "Model request failed").slice(0, 512); +} + +function modelFromEndpointUrl(endpointUrl: string): string { + try { + const parts = new URL(endpointUrl).pathname.split("/"); + const endpointIndex = parts.indexOf("serving-endpoints"); + const encoded = parts[endpointIndex + 1]; + return encoded ? decodeURIComponent(encoded) : endpointUrl; + } catch { + return endpointUrl; + } +} + /** * Optional generation parameters forwarded to the OpenAI-compatible serving * request body. Names match the serving API wire keys. Only keys that are set @@ -80,11 +227,7 @@ function extractLlamaToolJsonSlice(text: string): string | undefined { /** OpenAI SSE payload: `{ choices: [{ delta }] }`. */ function openAiChoicesDelta(parsed: unknown): unknown { if (!isRecord(parsed)) return undefined; - const choices = parsed.choices; - if (!Array.isArray(choices) || choices.length < 1) return undefined; - const first = choices[0]; - if (!isRecord(first)) return undefined; - return first.delta; + return firstChoice(parsed)?.delta; } function isStreamingDeltaToolCall(value: unknown): value is DeltaToolCall { @@ -113,6 +256,8 @@ function throwIfExceedsStreamLimit( interface RawFetchAdapterOptions { endpointUrl: string; authenticate: () => Promise>; + /** Model/endpoint name recorded in lifecycle telemetry. */ + model?: string; maxSteps?: number; maxTokens?: number; /** Optional generation params forwarded to the serving request body. */ @@ -133,6 +278,8 @@ interface RawFetchAdapterOptions { */ interface StreamBodyAdapterOptions { streamBody: StreamBody; + /** Model/endpoint name recorded in lifecycle telemetry. */ + model?: string; maxSteps?: number; maxTokens?: number; generationParams?: GenerationParams; @@ -271,6 +418,7 @@ interface DeltaToolCall { */ export class DatabricksAdapter implements AgentAdapter { private streamBody: StreamBody; + private model: string; private maxSteps: number; private maxTokens: number; private generationParams: GenerationParams; @@ -291,8 +439,10 @@ export class DatabricksAdapter implements AgentAdapter { if (isStreamBodyOptions(options)) { this.streamBody = options.streamBody; + this.model = options.model ?? "databricks-model-serving"; } else { const { endpointUrl, authenticate } = options; + this.model = options.model ?? modelFromEndpointUrl(endpointUrl); this.streamBody = async (body, signal) => { const fetchSignal = signal ?? AbortSignal.timeout(RAW_FETCH_DEFAULT_TIMEOUT_MS); @@ -314,7 +464,7 @@ export class DatabricksAdapter implements AgentAdapter { ); } if (!response.body) throw new Error("No response body"); - return response.body; + return retainResponseHeaders(response.body, response.headers); }; } } @@ -351,6 +501,7 @@ export class DatabricksAdapter implements AgentAdapter { body, signal, ), + model: endpointName, maxSteps, maxTokens, generationParams, @@ -567,20 +718,25 @@ export class DatabricksAdapter implements AgentAdapter { body.tools = tools; } - let responseBody: ReadableStream; - try { - responseBody = await this.streamBody(body, context.signal); - } catch (err) { - const msg = err instanceof Error ? err.message : "Stream request failed"; - yield { type: "status", status: "error", error: msg }; - throw err; - } - - const reader = responseBody.getReader(); + const stepId = globalThis.crypto.randomUUID(); + const startedAt = Date.now(); + const startEvent: AgentEvent = { + type: "model_start", + stepId, + model: this.model, + provider: "databricks", + input: structuredClone(body), + startedAt, + }; - const decoder = new TextDecoder(); - let buffer = ""; let fullText = ""; + let finalUsage = emptyUsage(); + let finishReason: string | undefined; + let firstTokenAt: number | undefined; + let streamStartedAt: number | undefined; + let modelError: string | undefined; + let caughtError: unknown; + let reader: ReadableStreamDefaultReader | undefined; const toolCallAccumulator = new Map< number, { @@ -590,8 +746,37 @@ export class DatabricksAdapter implements AgentAdapter { thoughtSignature?: string; } >(); + const emittedRemoteTraces = new Set(); + const snapshotToolCalls = (): OpenAIToolCall[] => + Array.from(toolCallAccumulator.values()).map((tc) => ({ + id: tc.id, + type: "function" as const, + function: { name: tc.name, arguments: tc.arguments || "{}" }, + ...(tc.thoughtSignature + ? { thoughtSignature: tc.thoughtSignature } + : {}), + })); try { + yield startEvent; + const responseBody = await this.streamBody(body, context.signal); + streamStartedAt = Date.now(); + const headerTraceId = getResponseHeaders(responseBody)?.get( + "x-databricks-trace-id", + ); + if (headerTraceId?.trim()) { + emittedRemoteTraces.add(headerTraceId); + yield { + type: "remote_trace", + traceId: headerTraceId, + source: "model-serving", + relation: "continued", + }; + } + + reader = responseBody.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; while (true) { if (context.signal?.aborted) break; @@ -632,9 +817,38 @@ export class DatabricksAdapter implements AgentAdapter { continue; } + if (!isRecord(parsed)) continue; + + const usage = normalizeUsage(parsed); + if (usage) finalUsage = usage; + + const choice = firstChoice(parsed); + if (typeof choice?.finish_reason === "string") { + finishReason = choice.finish_reason; + } + + const remoteTrace = remoteTraceFromPayload(parsed); + if (remoteTrace) { + const key = `${remoteTrace.traceId}:${remoteTrace.spanId ?? ""}`; + if (!emittedRemoteTraces.has(key)) { + emittedRemoteTraces.add(key); + yield remoteTrace; + } + } + const deltaUnknown = openAiChoicesDelta(parsed); if (!isRecord(deltaUnknown)) continue; + const toolCallsRaw = deltaUnknown.tool_calls; + if ( + firstTokenAt === undefined && + ((typeof deltaUnknown.content === "string" && + deltaUnknown.content.length > 0) || + (Array.isArray(toolCallsRaw) && toolCallsRaw.length > 0)) + ) { + firstTokenAt = Date.now(); + } + if (typeof deltaUnknown.content === "string") { const content = deltaUnknown.content; throwIfExceedsStreamLimit( @@ -647,7 +861,6 @@ export class DatabricksAdapter implements AgentAdapter { yield { type: "message_delta" as const, content }; } - const toolCallsRaw = deltaUnknown.tool_calls; if (!Array.isArray(toolCallsRaw)) continue; for (const tc of toolCallsRaw) { @@ -684,35 +897,59 @@ export class DatabricksAdapter implements AgentAdapter { } } } - } finally { - try { - await reader.cancel(); - } catch (cancelErr) { - console.debug( - "[DatabricksAdapter] reader.cancel() failed during teardown", - cancelErr, - ); + if (context.signal?.aborted && !finishReason) { + finishReason = "cancelled"; } - try { - reader.releaseLock(); - } catch (unlockErr) { - console.debug( - "[DatabricksAdapter] reader.releaseLock() failed during teardown", - unlockErr, - ); + } catch (err) { + if (context.signal?.aborted) { + finishReason ??= "cancelled"; + } else { + modelError = sanitizedModelError(err); + caughtError = err; + yield { type: "status", status: "error", error: modelError }; } + } finally { + if (reader) { + try { + await reader.cancel(); + } catch (cancelErr) { + console.debug( + "[DatabricksAdapter] reader.cancel() failed during teardown", + cancelErr, + ); + } + try { + reader.releaseLock(); + } catch (unlockErr) { + console.debug( + "[DatabricksAdapter] reader.releaseLock() failed during teardown", + unlockErr, + ); + } + } + + const endedAt = Date.now(); + yield { + type: "model_end", + stepId, + model: this.model, + provider: "databricks", + output: { text: fullText, toolCalls: snapshotToolCalls() }, + usage: finalUsage, + ...(finishReason ? { finishReason } : {}), + ...(firstTokenAt !== undefined ? { firstTokenAt } : {}), + streamDurationMs: + streamStartedAt === undefined + ? 0 + : Math.max(0, endedAt - streamStartedAt), + endedAt, + ...(modelError ? { error: modelError } : {}), + }; } - const toolCalls: OpenAIToolCall[] = Array.from( - toolCallAccumulator.values(), - ).map((tc) => ({ - id: tc.id, - type: "function" as const, - function: { name: tc.name, arguments: tc.arguments || "{}" }, - ...(tc.thoughtSignature ? { thoughtSignature: tc.thoughtSignature } : {}), - })); + if (caughtError !== undefined) throw caughtError; - return { text: fullText, toolCalls }; + return { text: fullText, toolCalls: snapshotToolCalls() }; } private async *executeToolCalls( diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index c20f74fa5..0d54565b4 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -66,6 +66,63 @@ function createReadableStream(chunks: string[]): ReadableStream { }); } +function createTimedReadableStream( + reads: Array<{ at: number; chunk?: string; error?: Error }>, +): ReadableStream { + const encoder = new TextEncoder(); + let i = 0; + return { + getReader() { + return { + async read() { + const next = reads[i++]; + if (!next) return { done: true, value: undefined }; + vi.setSystemTime(next.at); + if (next.error) throw next.error; + if (next.chunk === undefined) { + return { done: true, value: undefined }; + } + return { done: false, value: encoder.encode(next.chunk) }; + }, + async cancel() {}, + releaseLock() {}, + }; + }, + } as unknown as ReadableStream; +} + +function completionChunk(options: { + content?: string; + toolCalls?: Array>; + finishReason?: string; + usage?: Record; + mlflowTraceId?: string; + mlflowSpanId?: string; +}): string { + return sseChunk( + JSON.stringify({ + choices: [ + { + delta: { + ...(options.content !== undefined + ? { content: options.content } + : {}), + ...(options.toolCalls ? { tool_calls: options.toolCalls } : {}), + }, + ...(options.finishReason + ? { finish_reason: options.finishReason } + : {}), + }, + ], + ...(options.usage ? { usage: options.usage } : {}), + ...(options.mlflowTraceId + ? { mlflow_trace_id: options.mlflowTraceId } + : {}), + ...(options.mlflowSpanId ? { mlflow_span_id: options.mlflowSpanId } : {}), + }), + ); +} + function mockFetch(chunks: string[]): typeof globalThis.fetch { return vi.fn().mockResolvedValue({ ok: true, @@ -116,6 +173,474 @@ describe("DatabricksAdapter", () => { afterEach(() => { globalThis.fetch = originalFetch; mockAuthenticate.mockClear(); + vi.useRealTimers(); + }); + + test("emits an exact lifecycle pair for every streamed tool-loop model step", async () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + + let request = 0; + globalThis.fetch = vi.fn().mockImplementation(async () => { + request++; + if (request === 1) { + const body = createTimedReadableStream([ + { + at: 1_010, + chunk: completionChunk({ + toolCalls: [ + { + index: 0, + id: "call_1", + type: "function", + function: { + name: "analytics__query", + arguments: '{"query":"SELECT 1"}', + }, + }, + ], + }), + }, + { + at: 1_020, + chunk: completionChunk({ + finishReason: "tool_calls", + usage: { + prompt_tokens: 20, + completion_tokens: 5, + total_tokens: 25, + prompt_tokens_details: { + cached_tokens: 4, + cache_creation_tokens: 2, + }, + cost_usd: 0.04, + }, + }), + }, + { at: 1_025 }, + ]); + return { + ok: true, + body, + headers: new Headers({ + "x-databricks-trace-id": + "trace:/main.agent_traces.appkit/remote-step-1", + }), + text: () => Promise.resolve(""), + }; + } + + return { + ok: true, + body: createTimedReadableStream([ + { + at: 1_040, + chunk: completionChunk({ content: "Final answer" }), + }, + { + at: 1_050, + chunk: completionChunk({ + finishReason: "stop", + usage: { + input_tokens: 10, + output_tokens: 7, + total_tokens: 17, + input_tokens_details: { + cached_tokens: 1, + cache_creation_tokens: 3, + }, + cost: 0.03, + }, + }), + }, + { at: 1_055 }, + ]), + headers: new Headers(), + text: () => Promise.resolve(""), + }; + }); + + const events: AgentEvent[] = []; + const adapter = createAdapter(); + for await (const event of adapter.run( + { + messages: createTestMessages(), + tools: createTestTools(), + threadId: "t1", + }, + { executeTool: vi.fn().mockResolvedValue([{ value: 1 }]) }, + )) { + events.push(event); + } + + expect(events.map((event) => event.type)).toEqual([ + "status", + "model_start", + "remote_trace", + "model_end", + "tool_call", + "tool_result", + "model_start", + "message_delta", + "model_end", + ]); + + const starts = events.filter((event) => event.type === "model_start"); + const ends = events.filter((event) => event.type === "model_end"); + expect(starts).toHaveLength(2); + expect(ends).toHaveLength(2); + expect(starts[0]).toEqual({ + type: "model_start", + stepId: expect.any(String), + model: "my-endpoint", + provider: "databricks", + input: { + messages: [{ role: "user", content: "Hello" }], + stream: true, + max_tokens: 4096, + tools: [ + { + type: "function", + function: { + name: "analytics__query", + description: "Run SQL", + parameters: { + type: "object", + properties: { query: { type: "string" } }, + required: ["query"], + }, + }, + }, + ], + }, + startedAt: 1_000, + }); + expect(starts[1]).toEqual({ + type: "model_start", + stepId: expect.any(String), + model: "my-endpoint", + provider: "databricks", + input: { + messages: [ + { role: "user", content: "Hello" }, + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_1", + type: "function", + function: { + name: "analytics__query", + arguments: '{"query":"SELECT 1"}', + }, + }, + ], + }, + { + role: "tool", + content: '[{"value":1}]', + tool_call_id: "call_1", + }, + ], + stream: true, + max_tokens: 4096, + tools: expect.any(Array), + }, + startedAt: 1_025, + }); + expect(starts[0].stepId).not.toBe(starts[1].stepId); + + expect(ends[0]).toEqual({ + type: "model_end", + stepId: starts[0].stepId, + model: "my-endpoint", + provider: "databricks", + output: { + text: "", + toolCalls: [ + { + id: "call_1", + type: "function", + function: { + name: "analytics__query", + arguments: '{"query":"SELECT 1"}', + }, + }, + ], + }, + usage: { + inputTokens: 20, + outputTokens: 5, + totalTokens: 25, + cacheReadInputTokens: 4, + cacheCreationInputTokens: 2, + costUsd: 0.04, + costAvailable: true, + }, + finishReason: "tool_calls", + firstTokenAt: 1_010, + streamDurationMs: 25, + endedAt: 1_025, + }); + expect(ends[1]).toEqual({ + type: "model_end", + stepId: starts[1].stepId, + model: "my-endpoint", + provider: "databricks", + output: { text: "Final answer", toolCalls: [] }, + usage: { + inputTokens: 10, + outputTokens: 7, + totalTokens: 17, + cacheReadInputTokens: 1, + cacheCreationInputTokens: 3, + costUsd: 0.03, + costAvailable: true, + }, + finishReason: "stop", + firstTokenAt: 1_040, + streamDurationMs: 30, + endedAt: 1_055, + }); + expect(events[2]).toEqual({ + type: "remote_trace", + traceId: "trace:/main.agent_traces.appkit/remote-step-1", + source: "model-serving", + relation: "continued", + }); + }); + + test("preserves a final usage-only frame and marks an unpriced model unavailable", async () => { + vi.useFakeTimers(); + vi.setSystemTime(2_000); + globalThis.fetch = vi.fn().mockResolvedValue({ + ok: true, + body: createTimedReadableStream([ + { + at: 2_015, + chunk: completionChunk({ + finishReason: "stop", + usage: { + prompt_tokens: 8, + completion_tokens: 2, + total_tokens: 10, + }, + }), + }, + { at: 2_020 }, + ]), + headers: new Headers(), + text: () => Promise.resolve(""), + }); + + const adapter = createAdapter({ + endpointUrl: + "https://test.databricks.com/serving-endpoints/unpriced-model/invocations", + }); + const events: AgentEvent[] = []; + for await (const event of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + events.push(event); + } + + expect(events.map((event) => event.type)).toEqual([ + "status", + "model_start", + "model_end", + ]); + expect(events[2]).toEqual({ + type: "model_end", + stepId: (events[1] as Extract) + .stepId, + model: "unpriced-model", + provider: "databricks", + output: { text: "", toolCalls: [] }, + usage: { + inputTokens: 8, + outputTokens: 2, + totalTokens: 10, + costAvailable: false, + }, + finishReason: "stop", + streamDurationMs: 20, + endedAt: 2_020, + }); + expect(events[2]).not.toHaveProperty("usage.costUsd"); + expect(events[2]).not.toHaveProperty("firstTokenAt"); + }); + + test("emits a linked remote trace from a terminal MLflow trace event", async () => { + globalThis.fetch = mockFetch([ + completionChunk({ + content: "ok", + finishReason: "stop", + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, + mlflowTraceId: "trace:/catalog.schema.table/terminal-trace", + mlflowSpanId: "0123456789abcdef", + }), + ]); + + const events: AgentEvent[] = []; + for await (const event of createAdapter().run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + events.push(event); + } + + expect(events).toContainEqual({ + type: "remote_trace", + traceId: "trace:/catalog.schema.table/terminal-trace", + spanId: "0123456789abcdef", + source: "model-serving", + relation: "linked", + }); + }); + + test("finalizes partial output and sanitized error when streaming throws", async () => { + const adapter = new DatabricksAdapter({ + model: "failing-model", + streamBody: async () => + createTimedReadableStream([ + { at: 3_010, chunk: textDelta("partial") }, + { at: 3_020, error: new Error("stream exploded\nwith details") }, + ]), + maxSteps: 1, + }); + + vi.useFakeTimers(); + vi.setSystemTime(3_000); + const events: AgentEvent[] = []; + await expect(async () => { + for await (const event of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + events.push(event); + } + }).rejects.toThrow("stream exploded"); + + expect(events.map((event) => event.type)).toEqual([ + "status", + "model_start", + "message_delta", + "status", + "model_end", + ]); + const start = events[1] as Extract; + expect(events[4]).toEqual({ + type: "model_end", + stepId: start.stepId, + model: "failing-model", + provider: "databricks", + output: { text: "partial", toolCalls: [] }, + usage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }, + firstTokenAt: 3_010, + streamDurationMs: 20, + endedAt: 3_020, + error: "stream exploded with details", + }); + }); + + test("finalizes partial output when the model stream is cancelled", async () => { + vi.useFakeTimers(); + vi.setSystemTime(4_000); + const controller = new AbortController(); + const adapter = new DatabricksAdapter({ + model: "cancelled-model", + streamBody: async () => + createTimedReadableStream([ + { at: 4_010, chunk: textDelta("partial") }, + { at: 4_020, chunk: textDelta("ignored") }, + ]), + maxSteps: 1, + }); + const events: AgentEvent[] = []; + for await (const event of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn(), signal: controller.signal }, + )) { + events.push(event); + if (event.type === "message_delta") controller.abort(); + } + + expect(events.map((event) => event.type)).toEqual([ + "status", + "model_start", + "message_delta", + "model_end", + ]); + expect(events[3]).toEqual( + expect.objectContaining({ + type: "model_end", + stepId: (events[1] as Extract) + .stepId, + model: "cancelled-model", + output: { text: "partial", toolCalls: [] }, + finishReason: "cancelled", + firstTokenAt: 4_010, + streamDurationMs: 10, + endedAt: 4_010, + }), + ); + expect(events[3]).not.toHaveProperty("error"); + }); + + test("finalizes a model step when iteration stops immediately after model_start", async () => { + vi.useFakeTimers(); + vi.setSystemTime(5_000); + const adapter = new DatabricksAdapter({ + model: "early-cancel-model", + streamBody: async () => createReadableStream([]), + maxSteps: 1, + }); + const iterator = adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + ); + + expect((await iterator.next()).value).toEqual({ + type: "status", + status: "running", + }); + const started = await iterator.next(); + expect(started.value).toEqual( + expect.objectContaining({ + type: "model_start", + model: "early-cancel-model", + startedAt: 5_000, + }), + ); + + const finalized = await iterator.return(); + expect(finalized).toEqual({ + done: false, + value: { + type: "model_end", + stepId: (started.value as Extract) + .stepId, + model: "early-cancel-model", + provider: "databricks", + output: { text: "", toolCalls: [] }, + usage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }, + streamDurationMs: 0, + endedAt: 5_000, + }, + }); + await iterator.return(); }); test("streams text deltas from the model", async () => { @@ -136,8 +661,19 @@ describe("DatabricksAdapter", () => { } expect(events[0]).toEqual({ type: "status", status: "running" }); - expect(events[1]).toEqual({ type: "message_delta", content: "Hello" }); - expect(events[2]).toEqual({ type: "message_delta", content: " world" }); + expect(events[1]).toEqual( + expect.objectContaining({ type: "model_start", model: "my-endpoint" }), + ); + expect(events[2]).toEqual({ type: "message_delta", content: "Hello" }); + expect(events[3]).toEqual({ type: "message_delta", content: " world" }); + expect(events[4]).toEqual( + expect.objectContaining({ + type: "model_end", + stepId: (events[1] as Extract) + .stepId, + output: { text: "Hello world", toolCalls: [] }, + }), + ); }); test("calls authenticate() per request for fresh headers", async () => { @@ -813,11 +1349,20 @@ describe("DatabricksAdapter", () => { }).rejects.toThrow("serving_unreachable"); expect(events[0]).toEqual({ type: "status", status: "running" }); - expect(events[1]).toEqual({ + expect(events[1]).toEqual(expect.objectContaining({ type: "model_start" })); + expect(events[2]).toEqual({ type: "status", status: "error", error: "serving_unreachable", }); + expect(events[3]).toEqual( + expect.objectContaining({ + type: "model_end", + stepId: (events[1] as Extract) + .stepId, + error: "serving_unreachable", + }), + ); }); test("yields tool_result with error when executeTool rejects", async () => { diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index de9d0465c..1916b199c 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -18,6 +18,11 @@ export interface ApiClientLike { }; } +const responseHeadersByStream = new WeakMap< + ReadableStream, + Headers +>(); + /** * Transport shim shared by the agent adapters: given a request body, returns * the raw SSE byte stream from a serving / AI-gateway endpoint. Injected at @@ -30,6 +35,33 @@ export type StreamBody = ( signal?: AbortSignal, ) => Promise>; +/** + * Retains response headers without changing or mutating the byte-stream API + * consumed by existing adapters. Weak ownership lets metadata be collected + * with the stream and avoids a discoverable property-name collision. + */ +export function retainResponseHeaders( + stream: ReadableStream, + headers: unknown, +): ReadableStream { + if (headers == null) return stream; + const normalized = + headers instanceof Headers + ? headers + : new Headers( + headers as Headers | Record | [string, string][], + ); + responseHeadersByStream.set(stream, normalized); + return stream; +} + +/** Reads response metadata retained by {@link retainResponseHeaders}. */ +export function getResponseHeaders( + stream: ReadableStream, +): Headers | undefined { + return responseHeadersByStream.get(stream); +} + /** * Invokes a serving endpoint using the SDK's high-level query API. * Returns a typed QueryEndpointResponse. @@ -92,13 +124,16 @@ export async function streamPath( raw: true, }, context, - )) as { contents: ReadableStream | null }; + )) as { + contents: ReadableStream | null; + headers?: unknown; + }; if (!response.contents) { throw new Error("Response body is null ā€” streaming not supported"); } - return response.contents; + return retainResponseHeaders(response.contents, response.headers); } /** diff --git a/packages/appkit/src/connectors/serving/tests/client.test.ts b/packages/appkit/src/connectors/serving/tests/client.test.ts index 34bfd743a..bd75424d7 100644 --- a/packages/appkit/src/connectors/serving/tests/client.test.ts +++ b/packages/appkit/src/connectors/serving/tests/client.test.ts @@ -1,6 +1,6 @@ import { afterEach, describe, expect, test, vi } from "vitest"; import { Context } from "../../../workspace-client"; -import { invoke, stream } from "../client"; +import { getResponseHeaders, invoke, stream } from "../client"; function createMockClient(host = "https://test.databricks.com") { return { @@ -95,6 +95,24 @@ describe("Serving Connector", () => { expect(result).toBeInstanceOf(ReadableStream); }); + test("retains response headers on the returned stream", async () => { + const contents = new ReadableStream(); + const client = createMockClient(); + client.apiClient.request.mockResolvedValue({ + contents, + headers: new Headers({ + "x-databricks-trace-id": + "trace:/main.agent_traces.appkit/remote-trace", + }), + }); + + const result = await stream(client, "my-endpoint", { messages: [] }); + + expect(getResponseHeaders(result)?.get("x-databricks-trace-id")).toBe( + "trace:/main.agent_traces.appkit/remote-trace", + ); + }); + test("sends stream: true in payload via apiClient.request", async () => { const client = createMockClient(); client.apiClient.request.mockResolvedValue({ diff --git a/packages/appkit/src/core/agent/consume-adapter-stream.ts b/packages/appkit/src/core/agent/consume-adapter-stream.ts index eb7168558..66fbceea9 100644 --- a/packages/appkit/src/core/agent/consume-adapter-stream.ts +++ b/packages/appkit/src/core/agent/consume-adapter-stream.ts @@ -1,4 +1,4 @@ -import type { AgentEvent } from "shared"; +import type { AgentEvent, AgentRemoteTraceEvent } from "shared"; import { AgentUsageAccumulator, type ConsumedAgentStream, @@ -20,6 +20,29 @@ interface ConsumeAdapterStreamOptions { onEvent?: (event: AgentEvent) => void; } +function isValidRemoteTrace(event: AgentEvent): event is AgentRemoteTraceEvent { + if ( + event.type !== "remote_trace" || + typeof event.traceId !== "string" || + !event.traceId.trim() + ) { + return false; + } + if ( + event.source !== "model-serving" && + event.source !== "supervisor" && + event.source !== "remote-agent" + ) { + return false; + } + if (event.relation === "continued") return true; + return ( + event.relation === "linked" && + typeof event.spanId === "string" && + event.spanId.trim().length > 0 + ); +} + /** * Consume an adapter's event stream and aggregate the assistant's final text. * @@ -44,16 +67,26 @@ export async function consumeAdapterStream( ): Promise { let text = ""; const usage = new AgentUsageAccumulator(); + const consumedModelSteps = new Set(); let remoteTrace: ConsumedAgentStream["remoteTrace"]; for await (const event of stream) { - if (opts.signal?.aborted) break; + if ( + opts.signal?.aborted && + event.type !== "model_end" && + event.type !== "remote_trace" + ) { + break; + } if (event.type === "message_delta") { text += event.content; } else if (event.type === "message") { text = event.content; } else if (event.type === "model_end") { - usage.add(event.usage); - } else if (event.type === "remote_trace") { + if (!consumedModelSteps.has(event.stepId)) { + consumedModelSteps.add(event.stepId); + usage.add(event.usage); + } + } else if (isValidRemoteTrace(event)) { remoteTrace = event; } opts.onEvent?.(event); diff --git a/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts b/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts index e38a4edae..7e6d82ceb 100644 --- a/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts +++ b/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts @@ -67,6 +67,44 @@ describe("consumeAdapterStream", () => { expect(emitted).toEqual(["first"]); }); + test("still consumes the model finalizer that immediately follows cancellation", async () => { + const controller = new AbortController(); + const result = await consumeAdapterStream( + (async function* () { + yield { type: "message_delta", content: "partial" } as AgentEvent; + controller.abort(); + yield { + type: "model_end", + stepId: "cancelled-step", + model: "model-a", + provider: "databricks", + output: { text: "partial" }, + usage: { + inputTokens: 8, + outputTokens: 2, + totalTokens: 10, + costAvailable: false, + }, + finishReason: "cancelled", + streamDurationMs: 10, + endedAt: 110, + } as AgentEvent; + yield { type: "message_delta", content: "ignored" } as AgentEvent; + })(), + { signal: controller.signal }, + ); + + expect(result).toEqual({ + text: "partial", + usage: { + inputTokens: 8, + outputTokens: 2, + totalTokens: 10, + costAvailable: false, + }, + }); + }); + test("returns an empty string for a stream with no content events", async () => { const result = await consumeAdapterStream( streamOf([{ type: "thinking", content: "ā€¦" }]), @@ -154,4 +192,88 @@ describe("consumeAdapterStream", () => { }, }); }); + + test("aggregates each model step once and retains the last valid remote trace", async () => { + const firstEnd: AgentEvent = { + type: "model_end", + stepId: "step-1", + model: "model-a", + provider: "databricks", + output: { text: "answer" }, + usage: { + inputTokens: 20, + outputTokens: 5, + totalTokens: 25, + cacheReadInputTokens: 4, + cacheCreationInputTokens: 2, + costUsd: 0.04, + costAvailable: true, + }, + streamDurationMs: 20, + endedAt: 120, + }; + const secondEnd: AgentEvent = { + type: "model_end", + stepId: "step-2", + model: "model-a", + provider: "databricks", + output: { text: "answer" }, + usage: { + inputTokens: 10, + outputTokens: 7, + totalTokens: 17, + cacheReadInputTokens: 1, + cacheCreationInputTokens: 3, + costUsd: 0.03, + costAvailable: true, + }, + streamDurationMs: 30, + endedAt: 150, + }; + const result = await consumeAdapterStream( + streamOf([ + { type: "message_delta", content: "answer" }, + firstEnd, + firstEnd, + secondEnd, + { + type: "remote_trace", + traceId: "trace:/catalog.schema.table/valid", + source: "model-serving", + relation: "continued", + }, + { + type: "remote_trace", + traceId: "trace:/catalog.schema.table/invalid-linked", + source: "model-serving", + relation: "linked", + } as AgentEvent, + { + type: "remote_trace", + traceId: null, + source: "model-serving", + relation: "continued", + } as unknown as AgentEvent, + ]), + ); + + expect(result).toEqual({ + text: "answer", + usage: { + inputTokens: 30, + outputTokens: 12, + totalTokens: 42, + cacheReadInputTokens: 5, + cacheCreationInputTokens: 5, + costUsd: 0.07, + costAvailable: true, + }, + remoteTrace: { + type: "remote_trace", + traceId: "trace:/catalog.schema.table/valid", + source: "model-serving", + relation: "continued", + }, + }); + }); }); From 0a0573f328f9f4f2341537ef6e30b187fabe71da Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 18:17:00 -0700 Subject: [PATCH 08/31] fix: harden AppKit model lifecycle tracing Signed-off-by: Adam Gurary --- packages/appkit/src/agents/databricks.ts | 73 +++- .../src/agents/tests/databricks.test.ts | 340 +++++++++++++++++- .../src/core/agent/consume-adapter-stream.ts | 20 +- .../tests/consume-adapter-stream.test.ts | 71 ++++ 4 files changed, 477 insertions(+), 27 deletions(-) diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index d9332ec83..9a062c1ee 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -14,6 +14,10 @@ import { stream as servingStream, } from "../connectors/serving/client"; import { APPKIT_USER_AGENT, getClientOptions } from "../context/client-options"; +import { + DEFAULT_TRACE_REDACT_KEYS, + REDACTED_TRACE_VALUE, +} from "../telemetry/agent-tracing/attributes"; import { createWorkspaceClient } from "../workspace-client"; /** Default cap for a single incomplete SSE line tail (DoS guard). */ @@ -31,6 +35,17 @@ const PYTHON_STYLE_TOOL_PARSE_MAX_INPUT = 64 * 1024; /** Fallback HTTP timeout when the raw fetch adapter path receives no AbortSignal from the runner. */ const RAW_FETCH_DEFAULT_TIMEOUT_MS = 120_000; +const ERROR_SENSITIVE_KEY_PATTERN = new RegExp( + `\\b(${[...DEFAULT_TRACE_REDACT_KEYS] + .sort((left, right) => right.length - left.length) + .map((key) => key.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")) + .join("|")})\\b\\s*[:=]\\s*(?:"[^"]*"|'[^']*'|[^\\s,;&#]+)`, + "gi", +); +const ERROR_AUTHORIZATION_PATTERN = + /\b((?:proxy-)?authorization)\s*[:=]\s*(?:(?:Basic|Bearer)\s+)?[^\s,;]+/gi; +const ERROR_AUTH_SCHEME_PATTERN = /\b(Basic|Bearer)\s+[^\s,;]+/gi; + function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null; } @@ -50,9 +65,20 @@ function finiteNonNegativeNumber( function normalizeUsage( parsed: Record, + previous: AgentUsage, ): AgentUsage | undefined { const raw = isRecord(parsed.usage) ? parsed.usage : undefined; - if (!raw) return undefined; + const providerCost = + finiteNonNegativeNumber(raw, "cost_usd", "cost", "total_cost_usd") ?? + finiteNonNegativeNumber(parsed, "cost_usd", "cost", "total_cost_usd"); + if (!raw) { + if (providerCost === undefined) return undefined; + return { + ...previous, + costUsd: providerCost, + costAvailable: true, + }; + } const inputTokens = finiteNonNegativeNumber(raw, "input_tokens", "prompt_tokens") ?? 0; @@ -79,9 +105,8 @@ function normalizeUsage( "cache_creation_input_tokens", "cache_creation_tokens", ); - const providerCost = - finiteNonNegativeNumber(raw, "cost_usd", "cost", "total_cost_usd") ?? - finiteNonNegativeNumber(parsed, "cost_usd", "cost"); + const retainedCost = + providerCost ?? (previous.costAvailable ? previous.costUsd : undefined); return { inputTokens, @@ -91,8 +116,8 @@ function normalizeUsage( ...(cacheCreationInputTokens !== undefined ? { cacheCreationInputTokens } : {}), - ...(providerCost !== undefined ? { costUsd: providerCost } : {}), - costAvailable: providerCost !== undefined, + ...(retainedCost !== undefined ? { costUsd: retainedCost } : {}), + costAvailable: retainedCost !== undefined, }; } @@ -159,7 +184,20 @@ function remoteTraceFromPayload( function sanitizedModelError(error: unknown): string { const raw = error instanceof Error ? error.message : "Model request failed"; - const withoutControls = Array.from(raw, (character) => { + const redacted = raw + .replace( + ERROR_AUTHORIZATION_PATTERN, + (_match, key: string) => `${key}: ${REDACTED_TRACE_VALUE}`, + ) + .replace( + ERROR_AUTH_SCHEME_PATTERN, + (_match, scheme: string) => `${scheme} ${REDACTED_TRACE_VALUE}`, + ) + .replace( + ERROR_SENSITIVE_KEY_PATTERN, + (_match, key: string) => `${key}=${REDACTED_TRACE_VALUE}`, + ); + const withoutControls = Array.from(redacted, (character) => { const code = character.charCodeAt(0); return code < 32 || code === 127 ? " " : character; }).join(""); @@ -458,10 +496,7 @@ export class DatabricksAdapter implements AgentAdapter { signal: fetchSignal, }); if (!response.ok) { - const errorText = await response.text().catch(() => "Unknown error"); - throw new Error( - `Databricks API error (${response.status}): ${errorText}`, - ); + throw new Error(`Databricks API error (${response.status})`); } if (!response.body) throw new Error("No response body"); return retainResponseHeaders(response.body, response.headers); @@ -747,11 +782,19 @@ export class DatabricksAdapter implements AgentAdapter { } >(); const emittedRemoteTraces = new Set(); - const snapshotToolCalls = (): OpenAIToolCall[] => + const snapshotToolCalls = ( + normalizeCompletedArguments = false, + ): OpenAIToolCall[] => Array.from(toolCallAccumulator.values()).map((tc) => ({ id: tc.id, type: "function" as const, - function: { name: tc.name, arguments: tc.arguments || "{}" }, + function: { + name: tc.name, + arguments: + normalizeCompletedArguments && tc.arguments === "" + ? "{}" + : tc.arguments, + }, ...(tc.thoughtSignature ? { thoughtSignature: tc.thoughtSignature } : {}), @@ -819,7 +862,7 @@ export class DatabricksAdapter implements AgentAdapter { if (!isRecord(parsed)) continue; - const usage = normalizeUsage(parsed); + const usage = normalizeUsage(parsed, finalUsage); if (usage) finalUsage = usage; const choice = firstChoice(parsed); @@ -949,7 +992,7 @@ export class DatabricksAdapter implements AgentAdapter { if (caughtError !== undefined) throw caughtError; - return { text: fullText, toolCalls: snapshotToolCalls() }; + return { text: fullText, toolCalls: snapshotToolCalls(true) }; } private async *executeToolCalls( diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index 0d54565b4..6dbe46a9d 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -1,5 +1,6 @@ import type { AgentEvent, AgentToolDefinition, Message } from "shared"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { consumeAdapterStream } from "../../core/agent/consume-adapter-stream"; import { DatabricksAdapter, type GenerationParams, @@ -67,7 +68,12 @@ function createReadableStream(chunks: string[]): ReadableStream { } function createTimedReadableStream( - reads: Array<{ at: number; chunk?: string; error?: Error }>, + reads: Array<{ + at: number; + chunk?: string; + error?: Error; + onRead?: () => void; + }>, ): ReadableStream { const encoder = new TextEncoder(); let i = 0; @@ -78,6 +84,7 @@ function createTimedReadableStream( const next = reads[i++]; if (!next) return { done: true, value: undefined }; vi.setSystemTime(next.at); + next.onRead?.(); if (next.error) throw next.error; if (next.chunk === undefined) { return { done: true, value: undefined }; @@ -98,6 +105,7 @@ function completionChunk(options: { usage?: Record; mlflowTraceId?: string; mlflowSpanId?: string; + totalCostUsd?: number; }): string { return sseChunk( JSON.stringify({ @@ -119,6 +127,9 @@ function completionChunk(options: { ? { mlflow_trace_id: options.mlflowTraceId } : {}), ...(options.mlflowSpanId ? { mlflow_span_id: options.mlflowSpanId } : {}), + ...(options.totalCostUsd !== undefined + ? { total_cost_usd: options.totalCostUsd } + : {}), }), ); } @@ -472,6 +483,82 @@ describe("DatabricksAdapter", () => { expect(events[2]).not.toHaveProperty("firstTokenAt"); }); + test("applies a cost-only terminal frame without losing prior token usage", async () => { + globalThis.fetch = mockFetch([ + completionChunk({ + content: "answer", + finishReason: "stop", + usage: { + input_tokens: 12, + output_tokens: 4, + total_tokens: 16, + }, + }), + sseChunk( + JSON.stringify({ + choices: [], + total_cost_usd: 0.123, + }), + ), + sseChunk("[DONE]"), + ]); + + const events: AgentEvent[] = []; + for await (const event of createAdapter().run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + events.push(event); + } + + expect(events.find((event) => event.type === "model_end")).toEqual( + expect.objectContaining({ + usage: { + inputTokens: 12, + outputTokens: 4, + totalTokens: 16, + costUsd: 0.123, + costAvailable: true, + }, + }), + ); + }); + + test("maps top-level total_cost_usd when usage shares the terminal frame", async () => { + globalThis.fetch = mockFetch([ + completionChunk({ + finishReason: "stop", + usage: { + prompt_tokens: 6, + completion_tokens: 2, + total_tokens: 8, + }, + totalCostUsd: 0.456, + }), + sseChunk("[DONE]"), + ]); + + const events: AgentEvent[] = []; + for await (const event of createAdapter().run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + events.push(event); + } + + expect(events.find((event) => event.type === "model_end")).toEqual( + expect.objectContaining({ + usage: { + inputTokens: 6, + outputTokens: 2, + totalTokens: 8, + costUsd: 0.456, + costAvailable: true, + }, + }), + ); + }); + test("emits a linked remote trace from a terminal MLflow trace event", async () => { globalThis.fetch = mockFetch([ completionChunk({ @@ -594,6 +681,164 @@ describe("DatabricksAdapter", () => { expect(events[3]).not.toHaveProperty("error"); }); + test("drains buffered deltas through model_end after cancellation without exposing them", async () => { + const controller = new AbortController(); + globalThis.fetch = mockFetch([ + textDelta("visible") + + textDelta("suppressed") + + completionChunk({ + finishReason: "stop", + usage: { + input_tokens: 9, + output_tokens: 3, + total_tokens: 12, + }, + mlflowTraceId: "trace:/catalog.schema.table/cancelled-remote", + mlflowSpanId: "0123456789abcdef", + }) + + sseChunk("[DONE]"), + ]); + const seen: AgentEvent[] = []; + + const result = await consumeAdapterStream( + createAdapter().run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn(), signal: controller.signal }, + ), + { + signal: controller.signal, + onEvent(event) { + seen.push(event); + if (event.type === "message_delta") controller.abort(); + }, + }, + ); + + expect(seen.map((event) => event.type)).toEqual([ + "status", + "model_start", + "message_delta", + "remote_trace", + "model_end", + ]); + expect(result).toEqual({ + text: "visible", + usage: { + inputTokens: 9, + outputTokens: 3, + totalTokens: 12, + costAvailable: false, + }, + remoteTrace: { + type: "remote_trace", + traceId: "trace:/catalog.schema.table/cancelled-remote", + spanId: "0123456789abcdef", + source: "model-serving", + relation: "linked", + }, + }); + }); + + test("retains empty tool arguments when streaming fails after the tool name", async () => { + vi.useFakeTimers(); + vi.setSystemTime(6_000); + const adapter = new DatabricksAdapter({ + model: "partial-tool-model", + streamBody: async () => + createTimedReadableStream([ + { + at: 6_010, + chunk: toolCallDelta(0, "call_partial", "analytics__query", ""), + }, + { at: 6_020, error: new Error("stream failed") }, + ]), + maxSteps: 1, + }); + const events: AgentEvent[] = []; + + await expect(async () => { + for await (const event of adapter.run( + { + messages: createTestMessages(), + tools: createTestTools(), + threadId: "t1", + }, + { executeTool: vi.fn() }, + )) { + events.push(event); + } + }).rejects.toThrow("stream failed"); + + expect(events.find((event) => event.type === "model_end")).toEqual( + expect.objectContaining({ + output: { + text: "", + toolCalls: [ + { + id: "call_partial", + type: "function", + function: { + name: "analytics__query", + arguments: "", + }, + }, + ], + }, + error: "stream failed", + }), + ); + }); + + test("retains empty tool arguments when cancelled after the tool name", async () => { + vi.useFakeTimers(); + vi.setSystemTime(7_000); + const controller = new AbortController(); + const adapter = new DatabricksAdapter({ + model: "cancelled-tool-model", + streamBody: async () => + createTimedReadableStream([ + { + at: 7_010, + chunk: toolCallDelta(0, "call_partial", "analytics__query", ""), + onRead: () => controller.abort(), + }, + { at: 7_020 }, + ]), + maxSteps: 1, + }); + const events: AgentEvent[] = []; + + for await (const event of adapter.run( + { + messages: createTestMessages(), + tools: createTestTools(), + threadId: "t1", + }, + { executeTool: vi.fn(), signal: controller.signal }, + )) { + events.push(event); + } + + expect(events.find((event) => event.type === "model_end")).toEqual( + expect.objectContaining({ + output: { + text: "", + toolCalls: [ + { + id: "call_partial", + type: "function", + function: { + name: "analytics__query", + arguments: "", + }, + }, + ], + }, + finishReason: "cancelled", + }), + ); + }); + test("finalizes a model step when iteration stops immediately after model_start", async () => { vi.useFakeTimers(); vi.setSystemTime(5_000); @@ -1313,23 +1558,106 @@ describe("DatabricksAdapter", () => { }).rejects.toThrow(/tool call arguments exceed/); }); - test("throws on non-ok response", async () => { + test("omits the raw response body from non-ok transport errors", async () => { + const secretValues = [ + "bearer-secret", + "authorization-secret", + "cookie-secret", + "api-key-secret", + "password-secret", + "credential-secret", + "url-token-secret", + "url-api-key-secret", + "url-password-secret", + ]; globalThis.fetch = vi.fn().mockResolvedValue({ ok: false, status: 401, - text: () => Promise.resolve("Unauthorized"), + text: () => + Promise.resolve( + "Unauthorized " + + "Bearer bearer-secret " + + "Authorization: Basic authorization-secret " + + "Cookie: session=cookie-secret " + + "X-API-Key: api-key-secret " + + "password=password-secret " + + "credentials=credential-secret " + + "https://example.test/path?token=url-token-secret&api_key=url-api-key-secret&password=url-password-secret", + ), }); const adapter = createAdapter(); + const events: AgentEvent[] = []; await expect(async () => { - for await (const _ of adapter.run( + for await (const event of adapter.run( { messages: createTestMessages(), tools: [], threadId: "t1" }, { executeTool: vi.fn() }, )) { - // drain + events.push(event); + } + }).rejects.toThrow(/^Databricks API error \(401\)$/); + + const modelEnd = events.find((event) => event.type === "model_end"); + expect(modelEnd).toEqual( + expect.objectContaining({ error: "Databricks API error (401)" }), + ); + for (const secret of secretValues) { + expect( + (modelEnd as Extract).error, + ).not.toContain(secret); + } + }); + + test("redacts secret-bearing patterns from other lifecycle errors", async () => { + const secretValues = [ + "bearer-secret", + "authorization-secret", + "cookie-secret", + "api-key-secret", + "password-secret", + "credential-secret", + "url-token-secret", + "url-api-key-secret", + "url-password-secret", + ]; + const adapter = new DatabricksAdapter({ + model: "secret-error-model", + streamBody: async () => { + throw new Error( + "failed\n" + + "Bearer bearer-secret " + + "Authorization: Basic authorization-secret " + + "Cookie: session=cookie-secret " + + "X-API-Key: api-key-secret " + + "password=password-secret " + + "credentials=credential-secret " + + "https://example.test/path?token=url-token-secret&api_key=url-api-key-secret&password=url-password-secret", + ); + }, + maxSteps: 1, + }); + const events: AgentEvent[] = []; + + await expect(async () => { + for await (const event of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + events.push(event); } - }).rejects.toThrow("Databricks API error (401): Unauthorized"); + }).rejects.toThrow("bearer-secret"); + + const error = ( + events.find( + (event): event is Extract => + event.type === "model_end", + ) as Extract + ).error; + expect(error).toContain("[REDACTED]"); + expect(error?.length).toBeLessThanOrEqual(512); + expect(error).not.toContain("\n"); + for (const secret of secretValues) expect(error).not.toContain(secret); }); test("yields error status then throws when injected streamBody fails", async () => { diff --git a/packages/appkit/src/core/agent/consume-adapter-stream.ts b/packages/appkit/src/core/agent/consume-adapter-stream.ts index 66fbceea9..b6dd4450c 100644 --- a/packages/appkit/src/core/agent/consume-adapter-stream.ts +++ b/packages/appkit/src/core/agent/consume-adapter-stream.ts @@ -67,21 +67,28 @@ export async function consumeAdapterStream( ): Promise { let text = ""; const usage = new AgentUsageAccumulator(); + const activeModelSteps = new Set(); const consumedModelSteps = new Set(); let remoteTrace: ConsumedAgentStream["remoteTrace"]; for await (const event of stream) { - if ( - opts.signal?.aborted && - event.type !== "model_end" && - event.type !== "remote_trace" - ) { - break; + const aborted = opts.signal?.aborted === true; + if (event.type === "model_start") activeModelSteps.add(event.stepId); + if (aborted) { + if (activeModelSteps.size === 0) break; + if ( + event.type !== "model_start" && + event.type !== "model_end" && + event.type !== "remote_trace" + ) { + continue; + } } if (event.type === "message_delta") { text += event.content; } else if (event.type === "message") { text = event.content; } else if (event.type === "model_end") { + activeModelSteps.delete(event.stepId); if (!consumedModelSteps.has(event.stepId)) { consumedModelSteps.add(event.stepId); usage.add(event.usage); @@ -90,6 +97,7 @@ export async function consumeAdapterStream( remoteTrace = event; } opts.onEvent?.(event); + if (aborted && activeModelSteps.size === 0) break; } return { text, diff --git a/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts b/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts index 7e6d82ceb..04a55a795 100644 --- a/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts +++ b/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts @@ -71,6 +71,14 @@ describe("consumeAdapterStream", () => { const controller = new AbortController(); const result = await consumeAdapterStream( (async function* () { + yield { + type: "model_start", + stepId: "cancelled-step", + model: "model-a", + provider: "databricks", + input: { prompt: "hello" }, + startedAt: 100, + } as AgentEvent; yield { type: "message_delta", content: "partial" } as AgentEvent; controller.abort(); yield { @@ -105,6 +113,69 @@ describe("consumeAdapterStream", () => { }); }); + test("tracks a model_start observed in the same turn that cancellation begins", async () => { + const controller = new AbortController(); + const seen: AgentEvent[] = []; + const result = await consumeAdapterStream( + (async function* () { + controller.abort(); + yield { + type: "model_start", + stepId: "racing-step", + model: "model-a", + provider: "databricks", + input: { prompt: "hello" }, + startedAt: 100, + } as AgentEvent; + yield { type: "message_delta", content: "suppressed" } as AgentEvent; + yield { + type: "remote_trace", + traceId: "trace:/catalog.schema.table/racing-step", + source: "model-serving", + relation: "continued", + } as AgentEvent; + yield { + type: "model_end", + stepId: "racing-step", + model: "model-a", + provider: "databricks", + output: { text: "suppressed" }, + usage: { + inputTokens: 5, + outputTokens: 1, + totalTokens: 6, + costAvailable: false, + }, + finishReason: "cancelled", + streamDurationMs: 10, + endedAt: 110, + } as AgentEvent; + })(), + { signal: controller.signal, onEvent: (event) => seen.push(event) }, + ); + + expect(seen.map((event) => event.type)).toEqual([ + "model_start", + "remote_trace", + "model_end", + ]); + expect(result).toEqual({ + text: "", + usage: { + inputTokens: 5, + outputTokens: 1, + totalTokens: 6, + costAvailable: false, + }, + remoteTrace: { + type: "remote_trace", + traceId: "trace:/catalog.schema.table/racing-step", + source: "model-serving", + relation: "continued", + }, + }); + }); + test("returns an empty string for a stream with no content events", async () => { const result = await consumeAdapterStream( streamOf([{ type: "thinking", content: "ā€¦" }]), From 78ecc8e66019d67407d3d5cd8c5898e4b2864780 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 18:28:06 -0700 Subject: [PATCH 09/31] fix: bound cancelled model stream draining Signed-off-by: Adam Gurary --- packages/appkit/src/agents/databricks.ts | 7 + .../src/agents/tests/databricks.test.ts | 26 ++- .../src/core/agent/consume-adapter-stream.ts | 219 ++++++++++++++++-- .../tests/consume-adapter-stream.test.ts | 157 ++++++++++++- 4 files changed, 385 insertions(+), 24 deletions(-) diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 9a062c1ee..a24412075 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -45,6 +45,8 @@ const ERROR_SENSITIVE_KEY_PATTERN = new RegExp( const ERROR_AUTHORIZATION_PATTERN = /\b((?:proxy-)?authorization)\s*[:=]\s*(?:(?:Basic|Bearer)\s+)?[^\s,;]+/gi; const ERROR_AUTH_SCHEME_PATTERN = /\b(Basic|Bearer)\s+[^\s,;]+/gi; +const ERROR_COOKIE_PATTERN = /\b(cookie|set-cookie)\s*[:=]\s*[^\r\n]*/gi; +const ERROR_URL_PATTERN = /\bhttps?:\/\/[^\s]+/gi; function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null; @@ -185,6 +187,11 @@ function remoteTraceFromPayload( function sanitizedModelError(error: unknown): string { const raw = error instanceof Error ? error.message : "Model request failed"; const redacted = raw + .replace( + ERROR_COOKIE_PATTERN, + (_match, key: string) => `${key}: ${REDACTED_TRACE_VALUE}`, + ) + .replace(ERROR_URL_PATTERN, REDACTED_TRACE_VALUE) .replace( ERROR_AUTHORIZATION_PATTERN, (_match, key: string) => `${key}: ${REDACTED_TRACE_VALUE}`, diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index 6dbe46a9d..6dc11b85d 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -1613,26 +1613,32 @@ describe("DatabricksAdapter", () => { const secretValues = [ "bearer-secret", "authorization-secret", - "cookie-secret", + "cookie-one-secret", + "cookie-two-secret", + "set-cookie-one-secret", + "set-cookie-two-secret", "api-key-secret", "password-secret", "credential-secret", "url-token-secret", "url-api-key-secret", "url-password-secret", + "http-url-secret", ]; const adapter = new DatabricksAdapter({ model: "secret-error-model", streamBody: async () => { throw new Error( "failed\n" + - "Bearer bearer-secret " + - "Authorization: Basic authorization-secret " + - "Cookie: session=cookie-secret " + - "X-API-Key: api-key-secret " + - "password=password-secret " + - "credentials=credential-secret " + - "https://example.test/path?token=url-token-secret&api_key=url-api-key-secret&password=url-password-secret", + "Bearer bearer-secret\n" + + "Authorization: Basic authorization-secret\n" + + "Cookie: first=cookie-one-secret; second=cookie-two-secret; theme=public\n" + + "Set-Cookie: session=set-cookie-one-secret; Path=/; preference=set-cookie-two-secret; Secure\n" + + "X-API-Key: api-key-secret\n" + + "password=password-secret\n" + + "credentials=credential-secret\n" + + "https://example.test/path?token=url-token-secret&api_key=url-api-key-secret&password=url-password-secret\n" + + "http://insecure.test/path?secret=http-url-secret", ); }, maxSteps: 1, @@ -1657,6 +1663,10 @@ describe("DatabricksAdapter", () => { expect(error).toContain("[REDACTED]"); expect(error?.length).toBeLessThanOrEqual(512); expect(error).not.toContain("\n"); + expect(error).not.toMatch(/https?:\/\//); + expect(error).not.toContain("theme=public"); + expect(error).not.toContain("Path=/"); + expect(error).not.toContain("Secure"); for (const secret of secretValues) expect(error).not.toContain(secret); }); diff --git a/packages/appkit/src/core/agent/consume-adapter-stream.ts b/packages/appkit/src/core/agent/consume-adapter-stream.ts index b6dd4450c..9a4f0312a 100644 --- a/packages/appkit/src/core/agent/consume-adapter-stream.ts +++ b/packages/appkit/src/core/agent/consume-adapter-stream.ts @@ -18,6 +18,78 @@ interface ConsumeAdapterStreamOptions { * collect a raw event list for tests, or emit telemetry. */ onEvent?: (event: AgentEvent) => void; + /** @internal Bounds cleanup when an adapter violates lifecycle finalization. */ + cancellationDrain?: { + timeoutMs?: number; + maxEvents?: number; + }; +} + +const DEFAULT_CANCELLATION_DRAIN_TIMEOUT_MS = 1_000; +const DEFAULT_CANCELLATION_DRAIN_MAX_EVENTS = 256; +const CANCELLATION_FINALIZER_ERROR = + "Model stream cancellation finalizer unavailable"; + +type ModelStartEvent = Extract; + +function boundedInteger(value: number | undefined, fallback: number): number { + return value === undefined || !Number.isFinite(value) + ? fallback + : Math.max(0, Math.floor(value)); +} + +function requestIteratorCleanup(iterator: AsyncIterator): void { + try { + const cleanup = iterator.return?.(); + if (cleanup) void Promise.resolve(cleanup).catch(() => {}); + } catch { + // Cleanup is best-effort and must never replace the bounded result. + } +} + +async function raceNextWithAbort( + next: Promise>, + signal: AbortSignal, +): Promise< + { type: "next"; result: IteratorResult } | { type: "abort" } +> { + if (signal.aborted) return { type: "abort" }; + return new Promise((resolve, reject) => { + const onAbort = () => { + signal.removeEventListener("abort", onAbort); + resolve({ type: "abort" }); + }; + signal.addEventListener("abort", onAbort, { once: true }); + next.then( + (result) => { + signal.removeEventListener("abort", onAbort); + resolve({ type: "next", result }); + }, + (error) => { + signal.removeEventListener("abort", onAbort); + reject(error); + }, + ); + }); +} + +async function raceNextWithTimeout( + next: Promise>, + timeoutMs: number, +): Promise< + { type: "next"; result: IteratorResult } | { type: "timeout" } +> { + let timeout: ReturnType | undefined; + try { + return await Promise.race([ + next.then((result) => ({ type: "next" as const, result })), + new Promise<{ type: "timeout" }>((resolve) => { + timeout = setTimeout(() => resolve({ type: "timeout" }), timeoutMs); + }), + ]); + } finally { + if (timeout !== undefined) clearTimeout(timeout); + } } function isValidRemoteTrace(event: AgentEvent): event is AgentRemoteTraceEvent { @@ -67,26 +139,28 @@ export async function consumeAdapterStream( ): Promise { let text = ""; const usage = new AgentUsageAccumulator(); - const activeModelSteps = new Set(); + const activeModelSteps = new Map(); const consumedModelSteps = new Set(); let remoteTrace: ConsumedAgentStream["remoteTrace"]; - for await (const event of stream) { - const aborted = opts.signal?.aborted === true; - if (event.type === "model_start") activeModelSteps.add(event.stepId); - if (aborted) { - if (activeModelSteps.size === 0) break; - if ( - event.type !== "model_start" && - event.type !== "model_end" && - event.type !== "remote_trace" - ) { - continue; - } - } + const timeoutMs = boundedInteger( + opts.cancellationDrain?.timeoutMs, + DEFAULT_CANCELLATION_DRAIN_TIMEOUT_MS, + ); + const maxEvents = boundedInteger( + opts.cancellationDrain?.maxEvents, + DEFAULT_CANCELLATION_DRAIN_MAX_EVENTS, + ); + let drainDeadline: number | undefined; + let drainedEvents = 0; + const iterator = stream[Symbol.asyncIterator](); + + const consumeEvent = (event: AgentEvent): void => { if (event.type === "message_delta") { text += event.content; } else if (event.type === "message") { text = event.content; + } else if (event.type === "model_start") { + activeModelSteps.set(event.stepId, event); } else if (event.type === "model_end") { activeModelSteps.delete(event.stepId); if (!consumedModelSteps.has(event.stepId)) { @@ -97,8 +171,123 @@ export async function consumeAdapterStream( remoteTrace = event; } opts.onEvent?.(event); - if (aborted && activeModelSteps.size === 0) break; + }; + + const synthesizeCancellationFinalizers = (): void => { + const endedAt = Date.now(); + for (const start of [...activeModelSteps.values()]) { + consumeEvent({ + type: "model_end", + stepId: start.stepId, + model: start.model, + provider: start.provider, + output: { text: "", toolCalls: [] }, + usage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }, + finishReason: "cancelled", + streamDurationMs: 0, + endedAt, + error: CANCELLATION_FINALIZER_ERROR, + }); + } + }; + + let stopEarly = false; + while (true) { + if (opts.signal?.aborted) { + if (activeModelSteps.size === 0) { + stopEarly = true; + break; + } + drainDeadline ??= Date.now() + timeoutMs; + if (drainedEvents >= maxEvents || Date.now() >= drainDeadline) { + synthesizeCancellationFinalizers(); + stopEarly = true; + break; + } + } + + const pendingNext = Promise.resolve(iterator.next()); + let next: IteratorResult; + if (opts.signal && !opts.signal.aborted) { + const raced = await raceNextWithAbort(pendingNext, opts.signal); + if (raced.type === "abort") { + if (activeModelSteps.size === 0) { + void pendingNext.catch(() => {}); + stopEarly = true; + break; + } + drainDeadline ??= Date.now() + timeoutMs; + const bounded = await raceNextWithTimeout( + pendingNext, + Math.max(0, drainDeadline - Date.now()), + ); + if (bounded.type === "timeout") { + void pendingNext.catch(() => {}); + synthesizeCancellationFinalizers(); + stopEarly = true; + break; + } + next = bounded.result; + } else { + next = raced.result; + } + } else if (opts.signal?.aborted && activeModelSteps.size > 0) { + drainDeadline ??= Date.now() + timeoutMs; + const bounded = await raceNextWithTimeout( + pendingNext, + Math.max(0, drainDeadline - Date.now()), + ); + if (bounded.type === "timeout") { + void pendingNext.catch(() => {}); + synthesizeCancellationFinalizers(); + stopEarly = true; + break; + } + next = bounded.result; + } else { + next = await pendingNext; + } + + if (next.done) break; + const event = next.value; + if (event.type === "model_start") { + activeModelSteps.set(event.stepId, event); + } + const aborted = opts.signal?.aborted === true; + if (!aborted) { + consumeEvent(event); + continue; + } + + drainDeadline ??= Date.now() + timeoutMs; + if (activeModelSteps.size === 0) { + stopEarly = true; + break; + } + drainedEvents += 1; + if ( + event.type === "model_start" || + event.type === "model_end" || + event.type === "remote_trace" + ) { + consumeEvent(event); + } + if (activeModelSteps.size === 0) { + stopEarly = true; + break; + } + if (drainedEvents >= maxEvents || Date.now() >= drainDeadline) { + synthesizeCancellationFinalizers(); + stopEarly = true; + break; + } } + if (stopEarly) requestIteratorCleanup(iterator); return { text, usage: usage.snapshot(), diff --git a/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts b/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts index 04a55a795..4148c752e 100644 --- a/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts +++ b/packages/appkit/src/core/agent/tests/consume-adapter-stream.test.ts @@ -1,5 +1,5 @@ import type { AgentEvent } from "shared"; -import { describe, expect, test } from "vitest"; +import { afterEach, describe, expect, test, vi } from "vitest"; import { consumeAdapterStream } from "../consume-adapter-stream"; async function* streamOf( @@ -11,6 +11,10 @@ async function* streamOf( } describe("consumeAdapterStream", () => { + afterEach(() => { + vi.useRealTimers(); + }); + test("concatenates message_delta events into the final text", async () => { const result = await consumeAdapterStream( streamOf([ @@ -176,6 +180,157 @@ describe("consumeAdapterStream", () => { }); }); + test("synthesizes one finalizer when the post-abort event budget is exhausted", async () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + const controller = new AbortController(); + const seen: AgentEvent[] = []; + let postAbortEvents = 0; + + const result = await consumeAdapterStream( + (async function* () { + yield { + type: "model_start", + stepId: "bounded-step", + model: "model-a", + provider: "databricks", + input: { prompt: "hello" }, + startedAt: 900, + } as AgentEvent; + yield { type: "message_delta", content: "visible" } as AgentEvent; + for (let i = 0; i < 100; i++) { + postAbortEvents++; + yield { + type: "message_delta", + content: `suppressed-${i}`, + } as AgentEvent; + } + })(), + { + signal: controller.signal, + cancellationDrain: { maxEvents: 2, timeoutMs: 10_000 }, + onEvent(event) { + seen.push(event); + if (event.type === "message_delta") controller.abort(); + }, + }, + ); + + expect(postAbortEvents).toBe(2); + expect(seen.map((event) => event.type)).toEqual([ + "model_start", + "message_delta", + "model_end", + ]); + expect(result).toEqual({ + text: "visible", + usage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }, + }); + expect(seen[2]).toEqual({ + type: "model_end", + stepId: "bounded-step", + model: "model-a", + provider: "databricks", + output: { text: "", toolCalls: [] }, + usage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }, + finishReason: "cancelled", + streamDurationMs: 0, + endedAt: 1_000, + error: "Model stream cancellation finalizer unavailable", + }); + }); + + test("times out a stalled post-abort next without awaiting iterator cleanup", async () => { + vi.useFakeTimers(); + vi.setSystemTime(2_000); + const controller = new AbortController(); + const seen: AgentEvent[] = []; + let nextCalls = 0; + let returnCalls = 0; + const never = new Promise>(() => {}); + const iterator: AsyncIterator = { + next() { + nextCalls++; + if (nextCalls === 1) { + return Promise.resolve({ + done: false, + value: { + type: "model_start", + stepId: "stalled-step", + model: "model-b", + provider: "databricks", + input: { prompt: "hello" }, + startedAt: 1_900, + }, + }); + } + if (nextCalls === 2) { + return Promise.resolve({ + done: false, + value: { type: "message_delta", content: "visible" }, + }); + } + return never; + }, + return() { + returnCalls++; + return new Promise>(() => {}); + }, + }; + const stream: AsyncIterable = { + [Symbol.asyncIterator]: () => iterator, + }; + let resolved: Awaited> | undefined; + + void consumeAdapterStream(stream, { + signal: controller.signal, + cancellationDrain: { maxEvents: 100, timeoutMs: 25 }, + onEvent(event) { + seen.push(event); + if (event.type === "message_delta") controller.abort(); + }, + }).then((result) => { + resolved = result; + }); + + await vi.advanceTimersByTimeAsync(25); + + expect(resolved).toEqual({ + text: "visible", + usage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }, + }); + expect(seen.map((event) => event.type)).toEqual([ + "model_start", + "message_delta", + "model_end", + ]); + expect(returnCalls).toBe(1); + expect(seen[2]).toEqual( + expect.objectContaining({ + type: "model_end", + stepId: "stalled-step", + model: "model-b", + endedAt: 2_025, + error: "Model stream cancellation finalizer unavailable", + }), + ); + }); + test("returns an empty string for a stream with no content events", async () => { const result = await consumeAdapterStream( streamOf([{ type: "thinking", content: "ā€¦" }]), From c91bb5cd8653dc174fe9865468f47c816b0d35f2 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 18:53:03 -0700 Subject: [PATCH 10/31] feat: trace every AppKit agent invocation Signed-off-by: Adam Gurary --- packages/appkit/src/core/agent/run-agent.ts | 59 ++- .../src/core/agent/tests/run-agent.test.ts | 84 +++- packages/appkit/src/plugins/agents/agents.ts | 290 +++++++---- .../agents/tests/route-handler-errors.test.ts | 151 +++++- .../src/telemetry/agent-tracing/index.ts | 2 + .../agent-tracing/tests/tracer.test.ts | 463 ++++++++++++++++++ .../src/telemetry/agent-tracing/tracer.ts | 346 +++++++++++++ .../src/telemetry/agent-tracing/types.ts | 26 +- packages/shared/src/agent.ts | 5 +- 9 files changed, 1289 insertions(+), 137 deletions(-) create mode 100644 packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/tracer.ts diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 1e07368c4..76e627568 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -1,5 +1,4 @@ import { randomUUID } from "node:crypto"; -import { trace } from "@opentelemetry/api"; import type { AgentAdapter, AgentEvent, @@ -16,6 +15,11 @@ import { type SupervisorTool, } from "../../agents/supervisor-api"; import { createLogger } from "../../logging/logger"; +import { + type AgentTraceObserver, + resolveAgentTraceAppName, + runWithAgentTrace, +} from "../../telemetry/agent-tracing"; import { consumeAdapterStream } from "./consume-adapter-stream"; import { createPluginsProxy } from "./plugins-map"; import { resolveToolkitFromProvider } from "./toolkit-resolver"; @@ -52,6 +56,7 @@ export interface RunAgentInput { sessionId?: string; userId?: string; requestId?: string; + threadId?: string; appName?: string; } @@ -102,18 +107,42 @@ export async function runAgent( // plugin, and silently diverge in-instance state between parent and child // (e.g. query result caches, connection pools). const providerCache = new Map(); - await initStandalonePlugins(input.plugins ?? [], providerCache); - return runAgentInternal(def, input, providerCache); + const threadId = input.threadId ?? randomUUID(); + const requestId = input.requestId ?? randomUUID(); + const traced = await runWithAgentTrace( + { + appName: resolveAgentTraceAppName(input.appName), + agentName: def.name ?? "agent", + route: "runAgent", + sessionId: input.sessionId ?? threadId, + userId: input.userId ?? "service-principal", + requestId, + threadId, + }, + { messages: input.messages }, + async (observer) => { + await initStandalonePlugins(input.plugins ?? [], providerCache); + return runAgentInternal( + def, + { ...input, requestId, threadId }, + providerCache, + observer, + ); + }, + ); + return { + ...traced.value, + traceId: traced.traceId, + usage: traced.usage, + }; } async function runAgentInternal( def: AgentDefinition, input: RunAgentInput, providerCache: Map, -): Promise { - const traceId = - trace.getActiveSpan()?.spanContext().traceId ?? - randomUUID().replaceAll("-", ""); + observer: AgentTraceObserver, +): Promise> { const adapter = await resolveAdapter(def); const messages = normalizeMessages(input.messages, def.instructions); const toolIndex = buildStandaloneToolIndex( @@ -165,6 +194,7 @@ async function runAgentInternal( entry.agentDef, subInput, providerCache, + observer, ); return res.text; } @@ -189,7 +219,7 @@ async function runAgentInternal( { messages, tools, - threadId: randomUUID(), + threadId: input.threadId ?? randomUUID(), signal, extensions: buildStandaloneExtensions(toolIndex), }, @@ -203,17 +233,26 @@ async function runAgentInternal( signal, onEvent: (event) => { events.push(event); + observer.onEvent(event); }, }); + if (signal?.aborted) { + throw agentAbortError(); + } + return { text: consumed.text, events, - traceId, - usage: consumed.usage, }; } +function agentAbortError(): Error { + const error = new Error("Agent run aborted"); + error.name = "AbortError"; + return error; +} + /** * Eagerly construct every plugin in `input.plugins`, run the standard * AppKit lifecycle (`attachContext({})` + `await setup()`), and populate diff --git a/packages/appkit/src/core/agent/tests/run-agent.test.ts b/packages/appkit/src/core/agent/tests/run-agent.test.ts index 518300f5d..6f8b80240 100644 --- a/packages/appkit/src/core/agent/tests/run-agent.test.ts +++ b/packages/appkit/src/core/agent/tests/run-agent.test.ts @@ -1,4 +1,9 @@ -import { type Span, trace } from "@opentelemetry/api"; +import { trace } from "@opentelemetry/api"; +import { + BasicTracerProvider, + InMemorySpanExporter, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; import type { AgentAdapter, AgentEvent, @@ -43,19 +48,25 @@ describe("runAgent", () => { expect(result.events).toHaveLength(3); }); - test("returns the active trace and aggregate usage for identified runs", async () => { - const traceId = "0123456789abcdef0123456789abcdef"; - const activeSpan = { - spanContext: () => ({ - traceId, - spanId: "0123456789abcdef", - traceFlags: 1, - }), - } as unknown as Span; - const activeSpanSpy = vi - .spyOn(trace, "getActiveSpan") - .mockReturnValue(activeSpan); + test("creates the standalone semantic root and returns its trace and aggregate usage", async () => { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + provider.getTracer(name, version), + ); const events: AgentEvent[] = [ + { + type: "model_start", + stepId: "step-1", + model: "model-a", + provider: "databricks", + input: { messages: [{ role: "user", content: "hi" }] }, + startedAt: Date.now() - 10, + }, { type: "message_delta", content: "done" }, { type: "model_end", @@ -75,20 +86,49 @@ describe("runAgent", () => { }, ]; const def = createAgent({ + name: "planner", instructions: "x", model: scriptedAdapter(events), }); - const result = await runAgent(def, { - messages: "hi", - sessionId: "session-1", - userId: "user-1", - requestId: "request-1", - appName: "test-app", - }); - activeSpanSpy.mockRestore(); + let result!: Awaited>; + let spans = exporter.getFinishedSpans(); + try { + result = await runAgent(def, { + messages: "hi", + sessionId: "session-1", + userId: "user-1", + requestId: "request-1", + threadId: "thread-1", + appName: "test-app", + }); + await provider.forceFlush(); + spans = exporter.getFinishedSpans(); + } finally { + getTracerSpy.mockRestore(); + await provider.shutdown(); + } - expect(result.traceId).toBe(traceId); + const roots = spans.filter( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + const model = spans.find( + (span) => span.attributes["mlflow.spanType"] === "CHAT_MODEL", + ); + expect(roots).toHaveLength(1); + expect(roots[0].attributes).toMatchObject({ + "appkit.app.name": "test-app", + "appkit.agent.name": "planner", + "appkit.route": "runAgent", + "appkit.request.id": "request-1", + "appkit.thread.id": "thread-1", + "mlflow.trace.session": "session-1", + "mlflow.trace.user": "user-1", + }); + expect(model?.parentSpanContext?.spanId).toBe( + roots[0].spanContext().spanId, + ); + expect(result.traceId).toBe(roots[0].spanContext().traceId); expect(result.usage).toEqual({ inputTokens: 7, outputTokens: 2, diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index a3cf0ea5c..5fbf63ee4 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -51,6 +51,11 @@ import { isToolkitEntry } from "../../core/agent/types"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; +import { + type AgentTraceObserver, + resolveAgentTraceAppName, + runWithAgentTrace, +} from "../../telemetry/agent-tracing"; import { agentStreamDefaults } from "./defaults"; import { EventChannel } from "./event-channel"; import { AgentEventTranslator } from "./event-translator"; @@ -128,6 +133,7 @@ interface RunState { }; translator: AgentEventTranslator; outboundEvents: EventChannel; + traceObserver?: AgentTraceObserver; /** Boxed mutable counter shared across parent + all sub-agent dispatches. */ toolCallsUsed: { count: number }; } @@ -1057,7 +1063,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ): Promise { const abortController = new AbortController(); const signal = abortController.signal; - const requestId = randomUUID(); + const requestId = requestTraceId(req) ?? randomUUID(); this.trackStream(requestId, userId, abortController); // `hosted-supervisor` entries are not callable from the Node process @@ -1093,77 +1099,95 @@ export class AgentsPlugin extends Plugin implements ToolProvider { const executeTool = (name: string, args: unknown): Promise => this.dispatchToolCall(runState, registered.toolIndex, name, args, 0); + const pluginNames = this.context + ? this.context + .getPluginNames() + .filter((n) => n !== this.name && n !== "server") + : []; + const fullPrompt = composePromptForAgent( + registered, + this.config.baseSystemPrompt, + { + agentName: registered.name, + pluginNames, + toolNames: tools.map((t) => t.name), + }, + ); + const messagesWithSystem: Message[] = [ + { + id: "system", + role: "system", + content: fullPrompt, + createdAt: new Date(), + }, + ...thread.messages, + ]; + // Drive the adapter and the approval-event side-channel concurrently. // Outbound events from both sources flow through `outboundEvents`; the // generator below drains the channel in order. executeTool pushes // approval-pending events into the same channel before awaiting the gate. const driver = (async () => { try { - for (const evt of translator.translate({ - type: "metadata", - data: { threadId: thread.id }, - })) { - outboundEvents.push(evt); - } - - const pluginNames = this.context - ? this.context - .getPluginNames() - .filter((n) => n !== this.name && n !== "server") - : []; - const fullPrompt = composePromptForAgent( - registered, - this.config.baseSystemPrompt, + await runWithAgentTrace( { + appName: resolveAgentTraceAppName(), agentName: registered.name, - pluginNames, - toolNames: tools.map((t) => t.name), + route: "chat", + sessionId: requestSessionId(req) ?? thread.id, + userId, + requestId, + threadId: thread.id, }, - ); + { messages: messagesWithSystem }, + async (observer) => { + runState.traceObserver = observer; + // Trace discovery must be committed before any SSE event. The + // observer is created synchronously with the semantic root. + res.setHeader("X-MLflow-Trace-Id", observer.traceId); + for (const evt of translator.translate({ + type: "metadata", + data: { threadId: thread.id, traceId: observer.traceId }, + })) { + outboundEvents.push(evt); + } - const messagesWithSystem: Message[] = [ - { - id: "system", - role: "system", - content: fullPrompt, - createdAt: new Date(), - }, - ...thread.messages, - ]; + const stream = registered.adapter.run( + { + messages: messagesWithSystem, + tools, + threadId: thread.id, + signal, + extensions: buildAdapterExtensions(registered.toolIndex), + }, + { executeTool, signal }, + ); - const stream = registered.adapter.run( - { - messages: messagesWithSystem, - tools, - threadId: thread.id, - signal, - extensions: buildAdapterExtensions(registered.toolIndex), - }, - { executeTool, signal }, - ); + const { text: fullContent } = await consumeAdapterStream(stream, { + signal, + onEvent: (event) => { + observer.onEvent(event); + for (const translated of translator.translate(event)) { + outboundEvents.push(translated); + } + }, + }); - // The accumulation rule (deltas append, `message` replaces) is shared - // with `runAgent` and `runSubAgent`; see `consumeAdapterStream` for - // the rationale. - const { text: fullContent } = await consumeAdapterStream(stream, { - signal, - onEvent: (event) => { - for (const translated of translator.translate(event)) { - outboundEvents.push(translated); - } - }, - }); + if (signal.aborted) throw agentRequestAbortError(); - if (fullContent) { - await this.threadStore.addMessage(thread.id, userId, { - id: randomUUID(), - role: "assistant", - content: fullContent, - createdAt: new Date(), - }); - } + if (fullContent) { + await this.threadStore.addMessage(thread.id, userId, { + id: randomUUID(), + role: "assistant", + content: fullContent, + createdAt: new Date(), + }); + } - for (const evt of translator.finalize()) outboundEvents.push(evt); + for (const evt of translator.finalize()) outboundEvents.push(evt); + return { text: fullContent }; + }, + ); } catch (error) { if (signal.aborted) { outboundEvents.close(); @@ -1198,12 +1222,20 @@ export class AgentsPlugin extends Plugin implements ToolProvider { await this.executeStream( res, - async function* () { + async function* (streamSignal?: AbortSignal) { + const abortFromTransport = () => { + if (!signal.aborted) abortController.abort("Stream cancelled"); + }; + if (streamSignal?.aborted) abortFromTransport(); + streamSignal?.addEventListener("abort", abortFromTransport, { + once: true, + }); try { for await (const ev of outboundEvents) { yield ev; } } finally { + streamSignal?.removeEventListener("abort", abortFromTransport); await driver.catch(() => undefined); } }, @@ -1246,7 +1278,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ): Promise { const abortController = new AbortController(); const signal = abortController.signal; - const requestId = randomUUID(); + const requestId = requestTraceId(req) ?? randomUUID(); this.trackStream(requestId, userId, abortController); const tools = Array.from(registered.toolIndex.values()).map((e) => e.def); @@ -1271,56 +1303,78 @@ export class AgentsPlugin extends Plugin implements ToolProvider { const executeTool = (name: string, args: unknown): Promise => this.dispatchToolCall(runState, registered.toolIndex, name, args, 0); + const pluginNames = this.context + ? this.context + .getPluginNames() + .filter((n) => n !== this.name && n !== "server") + : []; + const fullPrompt = composePromptForAgent( + registered, + this.config.baseSystemPrompt, + { + agentName: registered.name, + pluginNames, + toolNames: tools.map((t) => t.name), + }, + ); + const messagesWithSystem: Message[] = [ + { + id: "system", + role: "system", + content: fullPrompt, + createdAt: new Date(), + }, + ...thread.messages, + ]; + let fullContent = ""; + let traceId = ""; try { - const pluginNames = this.context - ? this.context - .getPluginNames() - .filter((n) => n !== this.name && n !== "server") - : []; - const fullPrompt = composePromptForAgent( - registered, - this.config.baseSystemPrompt, + const traced = await runWithAgentTrace( { + appName: resolveAgentTraceAppName(), agentName: registered.name, - pluginNames, - toolNames: tools.map((t) => t.name), - }, - ); - - const messagesWithSystem: Message[] = [ - { - id: "system", - role: "system", - content: fullPrompt, - createdAt: new Date(), - }, - ...thread.messages, - ]; - - const stream = registered.adapter.run( - { - messages: messagesWithSystem, - tools, + route: invokeTraceRoute(req), + sessionId: requestSessionId(req) ?? thread.id, + userId, + requestId, threadId: thread.id, - signal, }, - { executeTool, signal }, + { messages: messagesWithSystem }, + async (observer) => { + runState.traceObserver = observer; + traceId = observer.traceId; + res.setHeader("X-MLflow-Trace-Id", traceId); + const stream = registered.adapter.run( + { + messages: messagesWithSystem, + tools, + threadId: thread.id, + signal, + }, + { executeTool, signal }, + ); + const consumed = await consumeAdapterStream(stream, { + signal, + onEvent: observer.onEvent, + }); + if (signal.aborted) throw agentRequestAbortError(); + if (consumed.text) { + await this.threadStore.addMessage(thread.id, userId, { + id: randomUUID(), + role: "assistant", + content: consumed.text, + createdAt: new Date(), + }); + } + return { text: consumed.text }; + }, ); - - ({ text: fullContent } = await consumeAdapterStream(stream, { signal })); - - if (fullContent) { - await this.threadStore.addMessage(thread.id, userId, { - id: randomUUID(), - role: "assistant", - content: fullContent, - createdAt: new Date(), - }); - } + fullContent = traced.value.text; + traceId = traced.traceId; } catch (error) { if (signal.aborted) { - res.status(499).json({ error: "Request aborted" }); + res.status(499).json({ error: "Request aborted", trace_id: traceId }); return; } logger.error("Agent invoke error: %O", error); @@ -1330,7 +1384,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { : error instanceof Error ? error.message : String(error); - res.status(500).json({ error: message }); + res.status(500).json({ error: message, trace_id: traceId }); return; } finally { this.approvalGate.abortStream(requestId); @@ -1363,6 +1417,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { created_at: Math.floor(Date.now() / 1000), status: "completed", thread_id: thread.id, + trace_id: traceId, output: [message], }); } @@ -1587,6 +1642,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { // Approval-pending events emitted by `dispatchToolCall` already // reach `outboundEvents` directly, so they are not routed here. onEvent: (event) => { + runState.traceObserver?.onEvent(event); if (event.type === "metadata") return; for (const translated of runState.translator.translate(event)) { runState.outboundEvents.push(translated); @@ -1752,6 +1808,36 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } } +function requestTraceId(req: express.Request): string | undefined { + return ( + req.header("x-request-id")?.trim() || + req.header("x-databricks-request-id")?.trim() || + undefined + ); +} + +function requestSessionId(req: express.Request): string | undefined { + return ( + req.header("x-mlflow-session-id")?.trim() || + req.header("x-session-id")?.trim() || + undefined + ); +} + +function invokeTraceRoute(req: express.Request): "invocations" | "responses" { + const requestPath = + req.route?.path ?? req.path ?? req.originalUrl ?? req.url ?? ""; + return String(requestPath).includes("responses") + ? "responses" + : "invocations"; +} + +function agentRequestAbortError(): Error { + const error = new Error("Agent request aborted"); + error.name = "AbortError"; + return error; +} + function normalizeAutoInherit(value: AgentsPluginConfig["autoInheritTools"]): { file: boolean; code: boolean; diff --git a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts index 2fc493ef4..33ebfcaa4 100644 --- a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts +++ b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts @@ -57,6 +57,7 @@ function mockRes() { return statusCode; }, json, + setHeader, }; } @@ -356,7 +357,7 @@ describe("POST /invocations & /responses ā€” successful invoke", () => { delete: vi.fn(), }; - const { res, json } = mockRes(); + const { res, json, setHeader } = mockRes(); await ( plugin as unknown as { _handleInvoke: ( @@ -378,6 +379,7 @@ describe("POST /invocations & /responses ā€” successful invoke", () => { role: string; content: Array<{ type: string; text: string }>; }>; + trace_id: string; }; expect(payload.object).toBe("response"); expect(payload.status).toBe("completed"); @@ -390,6 +392,153 @@ describe("POST /invocations & /responses ā€” successful invoke", () => { type: "output_text", text: "hello world", }); + expect(payload.trace_id).toMatch(/^[0-9a-f]{32}$/); + expect(setHeader).toHaveBeenCalledWith( + "X-MLflow-Trace-Id", + payload.trace_id, + ); + }); + + test("sets trace discovery headers even when the adapter throws", async () => { + const plugin = new AgentsPlugin({ dir: false }); + // biome-ignore lint/suspicious/noExplicitAny: seed private state + (plugin as any).agents.set("default", { + name: "default", + instructions: "hi", + adapter: { + async *run() { + yield { type: "message_delta", content: "partial" }; + throw new Error("adapter failed"); + }, + }, + toolIndex: new Map(), + }); + // biome-ignore lint/suspicious/noExplicitAny: seed private state + (plugin as any).defaultAgentName = "default"; + // biome-ignore lint/suspicious/noExplicitAny: stub + (plugin as any).threadStore = { + create: vi.fn().mockResolvedValue({ id: "t-new", messages: [] }), + addMessage: vi.fn(), + delete: vi.fn(), + }; + + const { res, json, setHeader } = mockRes(); + await ( + plugin as unknown as { + _handleInvoke: ( + r: express.Request, + w: express.Response, + ) => Promise; + } + )._handleInvoke(mockReq({ input: "hi" }), res); + + expect(res.status).toHaveBeenCalledWith(500); + expect(setHeader).toHaveBeenCalledWith( + "X-MLflow-Trace-Id", + expect.stringMatching(/^[0-9a-f]{32}$/), + ); + expect(json).toHaveBeenCalledWith({ + error: "adapter failed", + trace_id: expect.stringMatching(/^[0-9a-f]{32}$/), + }); + }); +}); + +describe("POST /chat ā€” trace discovery ordering", () => { + test("sets the trace header and emits trace metadata before any streamed content", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const order: string[] = []; + const streamed: Array> = []; + // biome-ignore lint/suspicious/noExplicitAny: drive the real stream producer without StreamManager transport + (plugin as any).executeStream = async ( + _res: express.Response, + source: (signal?: AbortSignal) => AsyncIterable>, + ) => { + for await (const event of source(new AbortController().signal)) { + streamed.push(event); + order.push(`body:${String(event.type)}`); + } + }; + // biome-ignore lint/suspicious/noExplicitAny: stub persistence + (plugin as any).threadStore = { + addMessage: vi.fn(), + delete: vi.fn(), + }; + const registered = { + name: "planner", + instructions: "help", + adapter: { + async *run() { + const startedAt = Date.now() - 10; + yield { + type: "model_start" as const, + stepId: "step-1", + model: "model-a", + provider: "databricks", + input: { prompt: "hi" }, + startedAt, + }; + yield { type: "message_delta" as const, content: "hello" }; + yield { + type: "model_end" as const, + stepId: "step-1", + model: "model-a", + provider: "databricks", + output: { text: "hello" }, + usage: { + inputTokens: 1, + outputTokens: 1, + totalTokens: 2, + costAvailable: false, + }, + streamDurationMs: 10, + endedAt: startedAt + 10, + }; + }, + }, + toolIndex: new Map(), + }; + const thread = { + id: "thread-1", + userId: "alice", + messages: [], + createdAt: new Date(), + updatedAt: new Date(), + }; + const req = mockReq({ message: "hi" }); + const { res, setHeader } = mockRes(); + setHeader.mockImplementation((name, value) => { + if (name === "X-MLflow-Trace-Id") order.push(`header:${String(value)}`); + }); + + await ( + plugin as unknown as { + _streamAgent: ( + request: express.Request, + response: express.Response, + agent: unknown, + currentThread: unknown, + userId: string, + ) => Promise; + } + )._streamAgent(req, res, registered, thread, "alice"); + + const metadata = streamed[0] as { + type?: string; + data?: { traceId?: string; threadId?: string }; + }; + expect(order[0]).toMatch(/^header:[0-9a-f]{32}$/); + expect(order[1]).toBe("body:appkit.metadata"); + expect(metadata).toMatchObject({ + type: "appkit.metadata", + data: { threadId: "thread-1", traceId: expect.any(String) }, + }); + expect(setHeader).toHaveBeenCalledWith( + "X-MLflow-Trace-Id", + metadata.data?.traceId, + ); + expect(JSON.stringify(streamed)).not.toContain("model_start"); + expect(JSON.stringify(streamed)).not.toContain("model_end"); }); }); diff --git a/packages/appkit/src/telemetry/agent-tracing/index.ts b/packages/appkit/src/telemetry/agent-tracing/index.ts index a6c39ca8b..4483a96ad 100644 --- a/packages/appkit/src/telemetry/agent-tracing/index.ts +++ b/packages/appkit/src/telemetry/agent-tracing/index.ts @@ -1,5 +1,7 @@ export { captureTraceValue } from "./serialization"; +export { resolveAgentTraceAppName, runWithAgentTrace } from "./tracer"; export type { + AgentTraceObserver, CapturedTraceValue, CaptureTraceValueOptions, ConsumedAgentStream, diff --git a/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts new file mode 100644 index 000000000..40a6c72ac --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts @@ -0,0 +1,463 @@ +import { context, SpanStatusCode, trace } from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, + type SpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import type { AgentEvent, AgentUsage } from "shared"; +import { + afterAll, + afterEach, + beforeAll, + describe, + expect, + test, + vi, +} from "vitest"; +import { + type MlflowUcConfig, + MlflowUcSpanProcessor, + MlflowUcTraceRegistry, + setActiveMlflowUcTraceRegistry, +} from "../../mlflow-uc"; +import * as agentTracing from "../index"; + +type AgentTraceRoute = "chat" | "invocations" | "responses" | "runAgent"; + +interface AgentTraceIdentity { + appName: string; + agentName: string; + route: AgentTraceRoute; + sessionId: string; + userId: string; + requestId: string; + threadId: string; +} + +interface AgentTraceObserver { + readonly traceId: string; + onEvent(event: AgentEvent): void; +} + +type RunWithAgentTrace = ( + identity: AgentTraceIdentity, + inputs: unknown, + operation: (observer: AgentTraceObserver) => Promise, +) => Promise<{ value: T; traceId: string; usage: AgentUsage }>; + +// Intentionally reaches through the existing public module so RED is a +// behavioral failure (the helper is absent), not a module-resolution error. +const runWithAgentTrace = ( + agentTracing as unknown as { + runWithAgentTrace: RunWithAgentTrace; + } +).runWithAgentTrace; + +const mlflowConfig: MlflowUcConfig = { + experimentId: "experiment-123", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "appkit", + otelSpansTableName: "main.agent_traces.appkit_otel_spans", +}; + +let provider: BasicTracerProvider | undefined; +let exporter: InMemorySpanExporter | undefined; +let getTracerSpy: { mockRestore(): void } | undefined; + +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); +}); + +afterEach(async () => { + getTracerSpy?.mockRestore(); + getTracerSpy = undefined; + setActiveMlflowUcTraceRegistry(undefined); + await provider?.shutdown(); + provider = undefined; + exporter = undefined; +}); + +afterAll(() => { + context.disable(); +}); + +function installTracing(extraProcessors: SpanProcessor[] = []): void { + exporter = new InMemorySpanExporter(); + provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter), ...extraProcessors], + }); + const installedProvider = provider; + getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + installedProvider.getTracer(name, version), + ); +} + +function identity(route: AgentTraceRoute): AgentTraceIdentity { + return { + appName: "support-console", + agentName: "planner", + route, + sessionId: "session-1", + userId: "user-1", + requestId: "request-1", + threadId: "thread-1", + }; +} + +function modelEvents( + options: { costAvailable?: boolean; costUsd?: number; error?: string } = {}, +): AgentEvent[] { + const startedAt = Date.now() - 50; + const costAvailable = options.costAvailable ?? true; + return [ + { + type: "model_start", + stepId: "step-1", + model: "dbx-claude-sonnet", + provider: "databricks", + input: { messages: [{ role: "user", content: "Hello" }] }, + startedAt, + }, + { type: "message_delta", content: "Hello " }, + { type: "message_delta", content: "world" }, + { + type: "model_end", + stepId: "step-1", + model: "dbx-claude-sonnet", + provider: "databricks", + output: { text: "Hello world" }, + usage: { + inputTokens: 7, + outputTokens: 2, + totalTokens: 9, + cacheReadInputTokens: 3, + ...(options.costUsd !== undefined + ? { costUsd: options.costUsd } + : costAvailable + ? { costUsd: 0.0125 } + : {}), + costAvailable, + }, + finishReason: options.error ? "error" : "stop", + firstTokenAt: startedAt + 12, + streamDurationMs: 40, + endedAt: startedAt + 50, + ...(options.error ? { error: options.error } : {}), + }, + ]; +} + +async function executeGoldenTrace(route: AgentTraceRoute) { + if (!provider) throw new Error("Tracing is not installed"); + const operation = async (observer: AgentTraceObserver) => { + for (const event of modelEvents()) observer.onEvent(event); + return { text: "Hello world" }; + }; + + if (route === "runAgent") { + return { + traced: await runWithAgentTrace( + identity(route), + { messages: [{ role: "user", content: "Hello" }] }, + operation, + ), + httpSpanId: undefined, + }; + } + + const http = provider.getTracer("http-test").startSpan(`POST /${route}`, { + attributes: { "http.request.method": "POST" }, + }); + try { + const traced = await context.with( + trace.setSpan(context.active(), http), + () => + runWithAgentTrace( + identity(route), + { messages: [{ role: "user", content: "Hello" }] }, + operation, + ), + ); + return { traced, httpSpanId: http.spanContext().spanId }; + } finally { + http.end(); + } +} + +function finishedSpans(): ReadableSpan[] { + return exporter?.getFinishedSpans() ?? []; +} + +describe("runWithAgentTrace golden span trees", () => { + test.each(["chat", "invocations", "responses", "runAgent"])( + "creates one semantic AGENT root and one exact model child for %s", + async (route) => { + installTracing(); + + const { traced, httpSpanId } = await executeGoldenTrace(route); + const spans = finishedSpans(); + const roots = spans.filter( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + const models = spans.filter( + (span) => span.attributes["mlflow.spanType"] === "CHAT_MODEL", + ); + + expect(roots).toHaveLength(1); + expect(models).toHaveLength(1); + const root = roots[0]; + const model = models[0]; + expect(root.attributes).toMatchObject({ + "mlflow.spanType": "AGENT", + "mlflow.spanInputs": '{"messages":[{"content":"Hello","role":"user"}]}', + "mlflow.spanOutputs": '{"text":"Hello world"}', + "mlflow.trace.session": "session-1", + "mlflow.trace.user": "user-1", + "appkit.app.name": "support-console", + "appkit.request.id": "request-1", + "appkit.thread.id": "thread-1", + "appkit.agent.name": "planner", + "appkit.route": route, + "mlflow.trace.tokenUsage": + '{"cache_read_input_tokens":3,"input_tokens":7,"output_tokens":2,"total_tokens":9}', + "appkit.cost.available": true, + "mlflow.llm.cost": 0.0125, + }); + expect(root.status.code).toBe(SpanStatusCode.OK); + expect( + root.duration[0] * 1_000_000_000 + root.duration[1], + ).toBeGreaterThanOrEqual(0); + expect(model.parentSpanContext?.spanId).toBe(root.spanContext().spanId); + expect(model.attributes).toMatchObject({ + "mlflow.spanType": "CHAT_MODEL", + "mlflow.spanInputs": '{"messages":[{"content":"Hello","role":"user"}]}', + "mlflow.spanOutputs": '{"text":"Hello world"}', + "mlflow.chat.model": "dbx-claude-sonnet", + "mlflow.chat.provider": "databricks", + "mlflow.chat.tokenUsage": + '{"cache_read_input_tokens":3,"input_tokens":7,"output_tokens":2,"total_tokens":9}', + "gen_ai.usage.input_tokens": 7, + "gen_ai.usage.output_tokens": 2, + "appkit.cache.read_input_tokens": 3, + "appkit.first_token.duration_ms": 12, + "appkit.stream.duration_ms": 40, + "appkit.cost.available": true, + "mlflow.llm.cost": 0.0125, + }); + expect(model.attributes["gen_ai.response.finish_reasons"]).toEqual([ + "stop", + ]); + expect(model.status.code).toBe(SpanStatusCode.OK); + expect(traced.value).toEqual({ text: "Hello world" }); + expect(traced.usage).toEqual({ + inputTokens: 7, + outputTokens: 2, + totalTokens: 9, + cacheReadInputTokens: 3, + costUsd: 0.0125, + costAvailable: true, + }); + expect(traced.traceId).toMatch(/^[0-9a-f]{32}$/); + expect(traced.traceId).not.toBe(httpSpanId); + expect(root.parentSpanContext?.spanId).toBe(httpSpanId); + }, + ); + + test("exposes the UC V4 root trace ID synchronously before the operation writes", async () => { + const registry = new MlflowUcTraceRegistry(mlflowConfig); + const ucProcessor = new MlflowUcSpanProcessor( + mlflowConfig, + { + exportTrace: (_batch, callback) => callback({ code: 0 }), + forceFlush: vi.fn().mockResolvedValue(undefined), + shutdown: vi.fn().mockResolvedValue(undefined), + }, + registry, + ); + setActiveMlflowUcTraceRegistry(registry); + installTracing([ucProcessor]); + const order: string[] = []; + let observedTraceId = ""; + + const traced = await runWithAgentTrace( + identity("chat"), + { message: "hello" }, + async (observer) => { + observedTraceId = observer.traceId; + order.push(`trace:${observer.traceId}`); + await Promise.resolve(); + order.push("body-write"); + observer.onEvent({ type: "message", content: "done" }); + return { text: "done" }; + }, + ); + + expect(order[0]).toBe(`trace:${observedTraceId}`); + expect(order[1]).toBe("body-write"); + expect(observedTraceId).toMatch( + /^trace:\/main\.agent_traces\.appkit\/[0-9a-f]{32}$/, + ); + expect(traced.traceId).toBe(observedTraceId); + }); + + test("keeps the fallback trace ID active when no tracer provider is installed", async () => { + let activeTraceId: string | undefined; + + const traced = await runWithAgentTrace( + identity("runAgent"), + { message: "local" }, + async (observer) => { + activeTraceId = trace.getActiveSpan()?.spanContext().traceId; + expect(observer.traceId).toMatch(/^[0-9a-f]{32}$/); + observer.onEvent({ type: "message", content: "done" }); + return { text: "done" }; + }, + ); + + expect(activeTraceId).toBe(traced.traceId); + }); + + test("finalizes partial output, model child, exception, and root exactly once on throw", async () => { + installTracing(); + const secretError = new Error("Authorization: Bearer top-secret-token"); + + await expect( + runWithAgentTrace( + identity("responses"), + { password: "secret" }, + async (observer) => { + const [start] = modelEvents(); + observer.onEvent(start); + observer.onEvent({ type: "message_delta", content: "partial" }); + throw secretError; + }, + ), + ).rejects.toBe(secretError); + + const root = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + const model = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "CHAT_MODEL", + ); + expect(root).toBeDefined(); + expect(model).toBeDefined(); + expect(root?.attributes["mlflow.spanInputs"]).toBe( + '{"password":"[REDACTED]"}', + ); + expect(root?.attributes["mlflow.spanOutputs"]).toBe( + '{"error":"[REDACTED]","text":"partial"}', + ); + expect(root?.status.code).toBe(SpanStatusCode.ERROR); + expect(model?.status.code).toBe(SpanStatusCode.ERROR); + expect( + root?.events.filter((event) => event.name === "exception"), + ).toHaveLength(1); + expect(JSON.stringify(root?.events)).not.toContain("top-secret-token"); + expect(JSON.stringify(model?.events)).not.toContain("top-secret-token"); + }); + + test("records an aborted partial stream as error and omits unavailable aggregate cost", async () => { + installTracing(); + const abortError = new DOMException("client cancelled", "AbortError"); + + await expect( + runWithAgentTrace( + identity("chat"), + { message: "hello" }, + async (observer) => { + for (const event of modelEvents({ costAvailable: false })) { + observer.onEvent(event); + if (event.type === "message_delta") break; + } + throw abortError; + }, + ), + ).rejects.toBe(abortError); + + const root = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(root?.attributes["mlflow.spanOutputs"]).toBe( + '{"error":"[REDACTED]","text":"Hello "}', + ); + expect(root?.attributes["appkit.cost.available"]).toBe(false); + expect(root?.attributes["mlflow.llm.cost"]).toBeUndefined(); + expect(root?.status.code).toBe(SpanStatusCode.ERROR); + }); + + test("records a root exception when model lifecycle reports an error without throwing", async () => { + installTracing(); + + const traced = await runWithAgentTrace( + identity("responses"), + { message: "hello" }, + async (observer) => { + for (const event of modelEvents({ error: "token=provider-secret" })) { + observer.onEvent(event); + } + return { text: "Hello world" }; + }, + ); + + const root = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(traced.value).toEqual({ text: "Hello world" }); + expect(root?.status.code).toBe(SpanStatusCode.ERROR); + expect( + root?.events.filter((event) => event.name === "exception"), + ).toHaveLength(1); + expect(JSON.stringify(root?.events)).not.toContain("provider-secret"); + }); + + test("counts each model_end once and withholds partial aggregate cost", async () => { + installTracing(); + + const traced = await runWithAgentTrace( + identity("invocations"), + { message: "hello" }, + async (observer) => { + const priced = modelEvents({ costUsd: 0.01 }); + for (const event of priced) observer.onEvent(event); + const pricedEnd = priced.at(-1); + if (!pricedEnd) throw new Error("Missing priced model_end fixture"); + observer.onEvent(pricedEnd); + for (const event of modelEvents({ costAvailable: false }).map( + (event) => + "stepId" in event ? { ...event, stepId: "step-2" } : event, + )) { + observer.onEvent(event); + } + return { text: "Hello world" }; + }, + ); + + expect(traced.usage).toEqual({ + inputTokens: 14, + outputTokens: 4, + totalTokens: 18, + cacheReadInputTokens: 6, + costAvailable: false, + }); + const root = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(root?.attributes["appkit.cost.available"]).toBe(false); + expect(root?.attributes["mlflow.llm.cost"]).toBeUndefined(); + expect( + finishedSpans().filter( + (span) => span.attributes["mlflow.spanType"] === "CHAT_MODEL", + ), + ).toHaveLength(2); + }); +}); diff --git a/packages/appkit/src/telemetry/agent-tracing/tracer.ts b/packages/appkit/src/telemetry/agent-tracing/tracer.ts new file mode 100644 index 000000000..f5a36a8ca --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/tracer.ts @@ -0,0 +1,346 @@ +import { randomUUID } from "node:crypto"; +import { + context, + isSpanContextValid, + type Span, + SpanStatusCode, + trace, +} from "@opentelemetry/api"; +import type { + AgentModelEndEvent, + AgentModelStartEvent, + AgentUsage, +} from "shared"; +import { getMlflowUcTraceId } from "../mlflow-uc"; +import { captureTraceValue } from "./serialization"; +import type { + AgentTraceIdentity, + AgentTraceObserver, + AgentTraceResult, +} from "./types"; +import { AgentUsageAccumulator } from "./usage"; + +const tracer = () => trace.getTracer("@databricks/appkit-agent-tracing"); + +interface ActiveModelStep { + event: AgentModelStartEvent; + span: Span; +} + +export function resolveAgentTraceAppName(appName?: string): string { + return ( + nonEmpty(appName) ?? + nonEmpty(process.env.DATABRICKS_APP_NAME) ?? + activeServiceName() ?? + nonEmpty(process.env.OTEL_SERVICE_NAME) ?? + "databricks-app" + ); +} + +export async function runWithAgentTrace( + identity: AgentTraceIdentity, + inputs: unknown, + operation: (observer: AgentTraceObserver) => Promise, +): Promise> { + return tracer().startActiveSpan( + `${identity.agentName} agent`, + { + attributes: { + "mlflow.spanType": "AGENT", + "mlflow.trace.session": identity.sessionId, + "mlflow.trace.user": identity.userId, + "appkit.app.name": resolveAgentTraceAppName(identity.appName), + "appkit.request.id": identity.requestId, + "appkit.thread.id": identity.threadId, + "appkit.agent.name": identity.agentName, + "appkit.route": identity.route, + }, + }, + async (root) => { + setCapturedAttribute(root, "mlflow.spanInputs", inputs); + const rootSpanContext = root.spanContext(); + const hasValidRootContext = isSpanContextValid(rootSpanContext); + const otelTraceId = hasValidRootContext + ? rootSpanContext.traceId + : randomUUID().replaceAll("-", ""); + const rootContext = trace.setSpan( + context.active(), + hasValidRootContext + ? root + : trace.wrapSpanContext({ + traceId: otelTraceId, + spanId: randomUUID().replaceAll("-", "").slice(0, 16), + traceFlags: 0, + }), + ); + const traceId = getMlflowUcTraceId(otelTraceId) ?? otelTraceId; + const usage = new AgentUsageAccumulator(); + const activeModels = new Map(); + const completedModels = new Set(); + let outputText = ""; + let lifecycleError = false; + + const observer: AgentTraceObserver = { + traceId, + onEvent(event) { + if (event.type === "message_delta") { + outputText += event.content; + return; + } + if (event.type === "message") { + outputText = event.content; + return; + } + if (event.type === "model_start") { + if ( + activeModels.has(event.stepId) || + completedModels.has(event.stepId) + ) { + return; + } + const span = tracer().startSpan( + `${event.provider} ${event.model}`, + { + startTime: event.startedAt, + attributes: modelStartAttributes(event), + }, + rootContext, + ); + setCapturedAttribute(span, "mlflow.spanInputs", event.input); + activeModels.set(event.stepId, { event, span }); + return; + } + if (event.type === "model_end") { + if (completedModels.has(event.stepId)) return; + completedModels.add(event.stepId); + usage.add(event.usage); + const active = activeModels.get(event.stepId); + activeModels.delete(event.stepId); + if (event.error) lifecycleError = true; + if (active) finalizeModelSpan(active, event); + return; + } + if (event.type === "status" && event.status === "error") { + lifecycleError = true; + } + }, + }; + + let value!: T; + let operationError: unknown; + let failed = false; + try { + value = await context.with(rootContext, () => operation(observer)); + } catch (error) { + failed = true; + operationError = error; + recordSafeException(root, error, "Agent operation failed"); + } finally { + if (activeModels.size > 0) { + lifecycleError = true; + const endedAt = Date.now(); + for (const active of activeModels.values()) { + completedModels.add(active.event.stepId); + const incompleteUsage: AgentUsage = { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }; + usage.add(incompleteUsage); + finalizeModelSpan(active, { + type: "model_end", + stepId: active.event.stepId, + model: active.event.model, + provider: active.event.provider, + output: { text: outputText }, + usage: incompleteUsage, + finishReason: failed ? "error" : "incomplete", + streamDurationMs: Math.max(0, endedAt - active.event.startedAt), + endedAt, + error: failed + ? "Agent operation failed before model completion" + : "Model lifecycle ended without model_end", + }); + } + activeModels.clear(); + } + + if (lifecycleError && !failed) { + recordSafeException( + root, + "Agent lifecycle reported an error", + "Agent operation failed", + ); + } + + const finalUsage = usage.snapshot(); + setRootUsageAttributes(root, finalUsage); + const finalOutputText = outputText || textFromValue(value); + setCapturedAttribute( + root, + "mlflow.spanOutputs", + failed + ? { text: finalOutputText, error: errorValue(operationError) } + : { text: finalOutputText }, + failed ? ["error"] : undefined, + ); + root.setStatus({ + code: + failed || lifecycleError ? SpanStatusCode.ERROR : SpanStatusCode.OK, + ...(failed || lifecycleError + ? { message: "Agent operation failed" } + : {}), + }); + root.end(); + } + + if (failed) throw operationError; + return { value, traceId, usage: usage.snapshot() }; + }, + ); +} + +function modelStartAttributes(event: AgentModelStartEvent) { + return { + "mlflow.spanType": "CHAT_MODEL", + "mlflow.chat.model": event.model, + "mlflow.chat.provider": event.provider, + "gen_ai.request.model": event.model, + "gen_ai.provider.name": event.provider, + "appkit.model.step_id": event.stepId, + }; +} + +function finalizeModelSpan( + active: ActiveModelStep, + event: AgentModelEndEvent, +): void { + const { span } = active; + setCapturedAttribute(span, "mlflow.spanOutputs", event.output); + span.setAttribute( + "mlflow.chat.tokenUsage", + captureTraceValue(mlflowTokenUsage(event.usage)).value, + ); + span.setAttribute("gen_ai.usage.input_tokens", event.usage.inputTokens); + span.setAttribute("gen_ai.usage.output_tokens", event.usage.outputTokens); + if (event.usage.cacheReadInputTokens !== undefined) { + span.setAttribute( + "appkit.cache.read_input_tokens", + event.usage.cacheReadInputTokens, + ); + } + if (event.usage.cacheCreationInputTokens !== undefined) { + span.setAttribute( + "appkit.cache.creation_input_tokens", + event.usage.cacheCreationInputTokens, + ); + } + if (event.finishReason) { + span.setAttribute("gen_ai.response.finish_reasons", [event.finishReason]); + } + if (event.firstTokenAt !== undefined) { + span.setAttribute( + "appkit.first_token.duration_ms", + Math.max(0, event.firstTokenAt - active.event.startedAt), + ); + } + span.setAttribute("appkit.stream.duration_ms", event.streamDurationMs); + setCostAttributes(span, event.usage); + if (event.error) { + recordSafeException(span, event.error, "Model operation failed"); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: "Model operation failed", + }); + } else { + span.setStatus({ code: SpanStatusCode.OK }); + } + span.end(event.endedAt); +} + +function setRootUsageAttributes(span: Span, usage: AgentUsage): void { + span.setAttribute( + "mlflow.trace.tokenUsage", + captureTraceValue(mlflowTokenUsage(usage)).value, + ); + setCostAttributes(span, usage); +} + +function setCostAttributes(span: Span, usage: AgentUsage): void { + span.setAttribute("appkit.cost.available", usage.costAvailable); + if (usage.costAvailable && usage.costUsd !== undefined) { + span.setAttribute("mlflow.llm.cost", usage.costUsd); + } +} + +function mlflowTokenUsage(usage: AgentUsage): Record { + return { + input_tokens: usage.inputTokens, + output_tokens: usage.outputTokens, + total_tokens: usage.totalTokens, + ...(usage.cacheReadInputTokens !== undefined + ? { cache_read_input_tokens: usage.cacheReadInputTokens } + : {}), + ...(usage.cacheCreationInputTokens !== undefined + ? { cache_creation_input_tokens: usage.cacheCreationInputTokens } + : {}), + }; +} + +function setCapturedAttribute( + span: Span, + key: string, + value: unknown, + redactKeys?: readonly string[], +): void { + const captured = captureTraceValue(value, { redactKeys }); + span.setAttribute(key, captured.value); + span.setAttribute(`${key}.original_bytes`, captured.originalBytes); + span.setAttribute(`${key}.sha256`, captured.sha256); + span.setAttribute(`${key}.truncated`, captured.truncated); +} + +function recordSafeException( + span: Span, + error: unknown, + publicMessage: string, +): void { + span.setAttribute( + "appkit.error", + captureTraceValue({ error: errorValue(error) }, { redactKeys: ["error"] }) + .value, + ); + span.recordException({ + name: error instanceof Error ? error.name : "Error", + message: publicMessage, + }); +} + +function errorValue(error: unknown): string { + return error instanceof Error + ? error.message + : String(error ?? "Unknown error"); +} + +function textFromValue(value: unknown): string { + if (typeof value === "string") return value; + if (value && typeof value === "object" && "text" in value) { + const text = (value as { text?: unknown }).text; + if (typeof text === "string") return text; + } + return ""; +} + +function nonEmpty(value: string | undefined): string | undefined { + const trimmed = value?.trim(); + return trimmed ? trimmed : undefined; +} + +function activeServiceName(): string | undefined { + const active = trace.getActiveSpan() as + | (Span & { resource?: { attributes?: Record } }) + | undefined; + const value = active?.resource?.attributes?.["service.name"]; + return typeof value === "string" ? nonEmpty(value) : undefined; +} diff --git a/packages/appkit/src/telemetry/agent-tracing/types.ts b/packages/appkit/src/telemetry/agent-tracing/types.ts index 1cac5f420..fe8796f49 100644 --- a/packages/appkit/src/telemetry/agent-tracing/types.ts +++ b/packages/appkit/src/telemetry/agent-tracing/types.ts @@ -1,4 +1,4 @@ -import type { AgentRemoteTraceEvent, AgentUsage } from "shared"; +import type { AgentEvent, AgentRemoteTraceEvent, AgentUsage } from "shared"; export interface CapturedTraceValue { value: string; @@ -17,3 +17,27 @@ export interface ConsumedAgentStream { usage: AgentUsage; remoteTrace?: AgentRemoteTraceEvent; } + +export type AgentTraceRoute = "chat" | "invocations" | "responses" | "runAgent"; + +export interface AgentTraceIdentity { + appName: string; + agentName: string; + route: AgentTraceRoute; + sessionId: string; + userId: string; + requestId: string; + threadId: string; +} + +export interface AgentTraceObserver { + /** MLflow V4 identity when UC is active; otherwise the 32-hex OTel trace ID. */ + readonly traceId: string; + onEvent(event: AgentEvent): void; +} + +export interface AgentTraceResult { + value: T; + traceId: string; + usage: AgentUsage; +} diff --git a/packages/shared/src/agent.ts b/packages/shared/src/agent.ts index d99516cdd..329fb9659 100644 --- a/packages/shared/src/agent.ts +++ b/packages/shared/src/agent.ts @@ -284,7 +284,10 @@ export interface AppKitThinkingEvent { export interface AppKitMetadataEvent { type: "appkit.metadata"; - data: Record; + data: Record & { + threadId?: string; + traceId?: string; + }; sequence_number: number; } From 53ee53465b1b43f2e8a9f745007d42e78512b233 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 19:15:26 -0700 Subject: [PATCH 11/31] fix: trace early AppKit agent failures Signed-off-by: Adam Gurary --- packages/appkit/src/plugins/agents/agents.ts | 318 +++++++++++------- .../agents/tests/route-handler-errors.test.ts | 317 ++++++++++++++++- .../agent-tracing/tests/tracer.test.ts | 130 +++++++ .../src/telemetry/agent-tracing/tracer.ts | 102 +++++- .../src/telemetry/agent-tracing/types.ts | 3 + 5 files changed, 729 insertions(+), 141 deletions(-) diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 5fbf63ee4..7f8051506 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -871,9 +871,26 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } private async _handleChat(req: express.Request, res: express.Response) { + const requestId = requestTraceId(req) ?? randomUUID(); + await runWithAgentTrace( + provisionalTraceIdentity(req, "chat", requestId), + req.body, + async (observer) => { + res.setHeader("X-MLflow-Trace-Id", observer.traceId); + await this._handleChatWithinTrace(req, res, observer, requestId); + }, + ); + } + + private async _handleChatWithinTrace( + req: express.Request, + res: express.Response, + observer: AgentTraceObserver, + requestId: string, + ): Promise { const parsed = chatRequestSchema.safeParse(req.body); if (!parsed.success) { - res.status(400).json({ + respondWithTraceError(observer, res, 400, { error: "Invalid request", details: parsed.error.flatten().fieldErrors, }); @@ -883,15 +900,17 @@ export class AgentsPlugin extends Plugin implements ToolProvider { const registered = this.resolveAgent(agentName); if (!registered) { - res.status(400).json({ + respondWithTraceError(observer, res, 400, { error: agentName ? `Agent "${agentName}" not found` : "No agent registered", }); return; } + observer.updateIdentity({ agentName: registered.name }); const userId = this.resolveUserId(req); + observer.updateIdentity({ userId }); // Reject early (before allocating a thread) when the user is already at // their concurrent-stream limit. Prevents a misbehaving client from @@ -899,7 +918,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { const limits = this.resolvedLimits; if (this.countUserStreams(userId) >= limits.maxConcurrentStreamsPerUser) { res.setHeader("Retry-After", "5"); - res.status(429).json({ + respondWithTraceError(observer, res, 429, { error: `Too many concurrent streams for this user (limit ${limits.maxConcurrentStreamsPerUser}). Wait for an existing stream to complete before starting another.`, }); return; @@ -916,10 +935,16 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ? await this.threadStore.get(threadId, userId) : null; if (threadId && !existing) { - res.status(404).json({ error: `Thread ${threadId} not found` }); + respondWithTraceError(observer, res, 404, { + error: `Thread ${threadId} not found`, + }); return; } thread = existing ?? (await this.threadStore.create(userId)); + observer.updateIdentity({ + threadId: thread.id, + sessionId: requestSessionId(req) ?? thread.id, + }); const userMessage: Message = { id: randomUUID(), @@ -930,10 +955,20 @@ export class AgentsPlugin extends Plugin implements ToolProvider { await this.threadStore.addMessage(thread.id, userId, userMessage); } catch (err) { logger.error("threadStore failed in /chat: %O", err); - res.status(500).json({ error: "Thread operation failed" }); + respondWithTraceError(observer, res, 500, { + error: "Thread operation failed", + }); return; } - return this._streamAgent(req, res, registered, thread, userId); + return this._streamAgent( + req, + res, + registered, + thread, + userId, + observer, + requestId, + ); } /** @@ -971,9 +1006,27 @@ export class AgentsPlugin extends Plugin implements ToolProvider { * does not provide. See {@link collectApprovalRequiredToolNames}. */ private async _handleInvoke(req: express.Request, res: express.Response) { + const route = invokeTraceRoute(req); + const requestId = requestTraceId(req) ?? randomUUID(); + await runWithAgentTrace( + provisionalTraceIdentity(req, route, requestId), + req.body, + async (observer) => { + res.setHeader("X-MLflow-Trace-Id", observer.traceId); + await this._handleInvokeWithinTrace(req, res, observer, requestId); + }, + ); + } + + private async _handleInvokeWithinTrace( + req: express.Request, + res: express.Response, + observer: AgentTraceObserver, + requestId: string, + ): Promise { const parsed = invocationsRequestSchema.safeParse(req.body); if (!parsed.success) { - res.status(400).json({ + respondWithTraceError(observer, res, 400, { error: "Invalid request", details: parsed.error.flatten().fieldErrors, }); @@ -982,9 +1035,12 @@ export class AgentsPlugin extends Plugin implements ToolProvider { const { input } = parsed.data; const registered = this.resolveAgent(); if (!registered) { - res.status(400).json({ error: "No agent registered" }); + respondWithTraceError(observer, res, 400, { + error: "No agent registered", + }); return; } + observer.updateIdentity({ agentName: registered.name }); // Pre-flight HITL gate. The non-streaming invoke surface has no way to // surface an approval prompt back to the caller and no way to receive @@ -993,7 +1049,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { // with a confusing "denied by user" tool result in the final text). const approvalGated = this.collectApprovalRequiredToolNames(registered); if (approvalGated.length > 0) { - res.status(400).json({ + respondWithTraceError(observer, res, 400, { error: `Agent '${registered.name}' exposes ${approvalGated.length} approval-gated tool(s) ` + `(${approvalGated.join(", ")}); /invocations and /responses are non-streaming and ` + @@ -1004,13 +1060,14 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } const userId = this.resolveUserId(req); + observer.updateIdentity({ userId }); // Match the rate-limit gate on /chat. Without this, a client can bypass // `limits.maxConcurrentStreamsPerUser` by hitting /invocations instead. const limits = this.resolvedLimits; if (this.countUserStreams(userId) >= limits.maxConcurrentStreamsPerUser) { res.setHeader("Retry-After", "5"); - res.status(429).json({ + respondWithTraceError(observer, res, 429, { error: `Too many concurrent streams for this user (limit ${limits.maxConcurrentStreamsPerUser}). Wait for an existing stream to complete before starting another.`, }); return; @@ -1021,6 +1078,10 @@ export class AgentsPlugin extends Plugin implements ToolProvider { let thread: Thread; try { thread = await this.threadStore.create(userId); + observer.updateIdentity({ + threadId: thread.id, + sessionId: requestSessionId(req) ?? thread.id, + }); if (typeof input === "string") { await this.threadStore.addMessage(thread.id, userId, { @@ -1047,11 +1108,21 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } } catch (err) { logger.error("threadStore failed in /invocations: %O", err); - res.status(500).json({ error: "Thread operation failed" }); + respondWithTraceError(observer, res, 500, { + error: "Thread operation failed", + }); return; } - return this._runAgentNonStreaming(req, res, registered, thread, userId); + return this._runAgentNonStreaming( + req, + res, + registered, + thread, + userId, + observer, + requestId, + ); } private async _streamAgent( @@ -1060,10 +1131,11 @@ export class AgentsPlugin extends Plugin implements ToolProvider { registered: RegisteredAgent, thread: Thread, userId: string, + observer: AgentTraceObserver, + requestId: string, ): Promise { const abortController = new AbortController(); const signal = abortController.signal; - const requestId = requestTraceId(req) ?? randomUUID(); this.trackStream(requestId, userId, abortController); // `hosted-supervisor` entries are not callable from the Node process @@ -1094,6 +1166,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { translator, outboundEvents, toolCallsUsed: { count: 0 }, + traceObserver: observer, }; const executeTool = (name: string, args: unknown): Promise => @@ -1123,79 +1196,62 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ...thread.messages, ]; + // Trace discovery is committed by the outer request handler before any + // SSE event. Queue the matching metadata before the adapter can emit. + for (const evt of translator.translate({ + type: "metadata", + data: { threadId: thread.id, traceId: observer.traceId }, + })) { + outboundEvents.push(evt); + } + // Drive the adapter and the approval-event side-channel concurrently. // Outbound events from both sources flow through `outboundEvents`; the // generator below drains the channel in order. executeTool pushes // approval-pending events into the same channel before awaiting the gate. + let driverFailure: { error: unknown } | undefined; const driver = (async () => { try { - await runWithAgentTrace( + const stream = registered.adapter.run( { - appName: resolveAgentTraceAppName(), - agentName: registered.name, - route: "chat", - sessionId: requestSessionId(req) ?? thread.id, - userId, - requestId, + messages: messagesWithSystem, + tools, threadId: thread.id, + signal, + extensions: buildAdapterExtensions(registered.toolIndex), }, - { messages: messagesWithSystem }, - async (observer) => { - runState.traceObserver = observer; - // Trace discovery must be committed before any SSE event. The - // observer is created synchronously with the semantic root. - res.setHeader("X-MLflow-Trace-Id", observer.traceId); - for (const evt of translator.translate({ - type: "metadata", - data: { threadId: thread.id, traceId: observer.traceId }, - })) { - outboundEvents.push(evt); - } - - const stream = registered.adapter.run( - { - messages: messagesWithSystem, - tools, - threadId: thread.id, - signal, - extensions: buildAdapterExtensions(registered.toolIndex), - }, - { executeTool, signal }, - ); + { executeTool, signal }, + ); - const { text: fullContent } = await consumeAdapterStream(stream, { - signal, - onEvent: (event) => { - observer.onEvent(event); - for (const translated of translator.translate(event)) { - outboundEvents.push(translated); - } - }, - }); + const { text: fullContent } = await consumeAdapterStream(stream, { + signal, + onEvent: (event) => { + observer.onEvent(event); + for (const translated of translator.translate(event)) { + outboundEvents.push(translated); + } + }, + }); - if (signal.aborted) throw agentRequestAbortError(); + if (signal.aborted) throw agentRequestAbortError(); - if (fullContent) { - await this.threadStore.addMessage(thread.id, userId, { - id: randomUUID(), - role: "assistant", - content: fullContent, - createdAt: new Date(), - }); - } + if (fullContent) { + await this.threadStore.addMessage(thread.id, userId, { + id: randomUUID(), + role: "assistant", + content: fullContent, + createdAt: new Date(), + }); + } - for (const evt of translator.finalize()) outboundEvents.push(evt); - return { text: fullContent }; - }, - ); + for (const evt of translator.finalize()) outboundEvents.push(evt); + observer.setOutput({ text: fullContent }); } catch (error) { - if (signal.aborted) { - outboundEvents.close(); - return; + driverFailure = { error }; + if (!signal.aborted) { + logger.error("Agent chat error: %O", error); } - logger.error("Agent chat error: %O", error); outboundEvents.close(error); - return; } finally { // Any pending approval gates for this stream are auto-denied so the // adapter can unwind if it was still waiting. @@ -1216,8 +1272,8 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ); } } + outboundEvents.close(); } - outboundEvents.close(); })(); await this.executeStream( @@ -1244,6 +1300,8 @@ export class AgentsPlugin extends Plugin implements ToolProvider { stream: { ...agentStreamDefaults.stream, streamId: requestId }, }, ); + await driver; + if (driverFailure) throw driverFailure.error; } /** @@ -1275,10 +1333,11 @@ export class AgentsPlugin extends Plugin implements ToolProvider { registered: RegisteredAgent, thread: Thread, userId: string, + observer: AgentTraceObserver, + requestId: string, ): Promise { const abortController = new AbortController(); const signal = abortController.signal; - const requestId = requestTraceId(req) ?? randomUUID(); this.trackStream(requestId, userId, abortController); const tools = Array.from(registered.toolIndex.values()).map((e) => e.def); @@ -1298,6 +1357,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { translator: new AgentEventTranslator(), outboundEvents: new EventChannel(), toolCallsUsed: { count: 0 }, + traceObserver: observer, }; const executeTool = (name: string, args: unknown): Promise => @@ -1328,53 +1388,38 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ]; let fullContent = ""; - let traceId = ""; try { - const traced = await runWithAgentTrace( + const stream = registered.adapter.run( { - appName: resolveAgentTraceAppName(), - agentName: registered.name, - route: invokeTraceRoute(req), - sessionId: requestSessionId(req) ?? thread.id, - userId, - requestId, + messages: messagesWithSystem, + tools, threadId: thread.id, + signal, }, - { messages: messagesWithSystem }, - async (observer) => { - runState.traceObserver = observer; - traceId = observer.traceId; - res.setHeader("X-MLflow-Trace-Id", traceId); - const stream = registered.adapter.run( - { - messages: messagesWithSystem, - tools, - threadId: thread.id, - signal, - }, - { executeTool, signal }, - ); - const consumed = await consumeAdapterStream(stream, { - signal, - onEvent: observer.onEvent, - }); - if (signal.aborted) throw agentRequestAbortError(); - if (consumed.text) { - await this.threadStore.addMessage(thread.id, userId, { - id: randomUUID(), - role: "assistant", - content: consumed.text, - createdAt: new Date(), - }); - } - return { text: consumed.text }; - }, + { executeTool, signal }, ); - fullContent = traced.value.text; - traceId = traced.traceId; + const consumed = await consumeAdapterStream(stream, { + signal, + onEvent: observer.onEvent, + }); + if (signal.aborted) throw agentRequestAbortError(); + fullContent = consumed.text; + if (fullContent) { + await this.threadStore.addMessage(thread.id, userId, { + id: randomUUID(), + role: "assistant", + content: fullContent, + createdAt: new Date(), + }); + } } catch (error) { if (signal.aborted) { - res.status(499).json({ error: "Request aborted", trace_id: traceId }); + const body = { + error: "Request aborted", + trace_id: observer.traceId, + }; + observer.recordError(error, body); + res.status(499).json(body); return; } logger.error("Agent invoke error: %O", error); @@ -1384,7 +1429,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { : error instanceof Error ? error.message : String(error); - res.status(500).json({ error: message, trace_id: traceId }); + const body = { error: message, trace_id: observer.traceId }; + observer.recordError(error, body); + res.status(500).json(body); return; } finally { this.approvalGate.abortStream(requestId); @@ -1411,15 +1458,17 @@ export class AgentsPlugin extends Plugin implements ToolProvider { role: "assistant", content: [{ type: "output_text", text: fullContent }], }; - res.json({ + const payload = { id: responseId, object: "response", created_at: Math.floor(Date.now() / 1000), status: "completed", thread_id: thread.id, - trace_id: traceId, + trace_id: observer.traceId, output: [message], - }); + }; + observer.setOutput(payload); + res.json(payload); } /** @@ -1816,6 +1865,43 @@ function requestTraceId(req: express.Request): string | undefined { ); } +function provisionalTraceIdentity( + req: express.Request, + route: "chat" | "invocations" | "responses", + requestId: string, +) { + const body = + req.body && typeof req.body === "object" + ? (req.body as Record) + : undefined; + const threadId = traceIdentityValue(body?.threadId) ?? requestId; + return { + appName: resolveAgentTraceAppName(), + agentName: traceIdentityValue(body?.agent) ?? "unresolved-agent", + route, + sessionId: requestSessionId(req) ?? threadId, + userId: + traceIdentityValue(req.header("x-forwarded-user")) ?? "unresolved-user", + requestId, + threadId, + }; +} + +function traceIdentityValue(value: unknown): string | undefined { + if (typeof value !== "string") return undefined; + return value.trim() || undefined; +} + +function respondWithTraceError( + observer: AgentTraceObserver, + res: express.Response, + status: number, + body: Record, +): void { + observer.recordError(body.error ?? `HTTP ${status}`, body); + res.status(status).json(body); +} + function requestSessionId(req: express.Request): string | undefined { return ( req.header("x-mlflow-session-id")?.trim() || diff --git a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts index 33ebfcaa4..1e9a21a56 100644 --- a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts +++ b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts @@ -1,3 +1,10 @@ +import { SpanStatusCode, trace } from "@opentelemetry/api"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; import type express from "express"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; @@ -35,6 +42,7 @@ function mockReq(body: unknown, userId = "alice"): express.Request { const headers: Record = { "x-forwarded-user": userId, "x-forwarded-access-token": "fake-token", + "x-request-id": "request-early", }; return { body, @@ -43,9 +51,13 @@ function mockReq(body: unknown, userId = "alice"): express.Request { } as unknown as express.Request; } -function mockRes() { - const json = vi.fn(); - const setHeader = vi.fn(); +function mockRes(order?: string[]) { + const json = vi.fn((body: unknown) => { + order?.push(`body:${JSON.stringify(body)}`); + }); + const setHeader = vi.fn((name: string, value: unknown) => { + order?.push(`header:${name}:${String(value)}`); + }); let statusCode = 200; const status = vi.fn((code: number) => { statusCode = code; @@ -61,6 +73,78 @@ function mockRes() { }; } +async function captureRouteSpans( + operation: () => Promise, +): Promise<{ spans: ReadableSpan[]; error?: unknown }> { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + provider.getTracer(name, version), + ); + let error: unknown; + let spans: ReadableSpan[] = []; + try { + await operation(); + } catch (caught) { + error = caught; + } finally { + await provider.forceFlush(); + spans = exporter.getFinishedSpans(); + getTracerSpy.mockRestore(); + await provider.shutdown(); + } + return { spans, ...(error !== undefined ? { error } : {}) }; +} + +function requestForRoute( + route: "chat" | "invocations" | "responses", + body: unknown, + userId = "alice", +): express.Request { + const req = mockReq(body, userId); + Object.assign(req, { + path: `/${route}`, + url: `/${route}`, + originalUrl: `/${route}`, + }); + return req; +} + +function expectEarlyErrorTrace( + spans: ReadableSpan[], + order: string[], + route: "chat" | "invocations" | "responses", + inputKey: "message" | "input", +): void { + const roots = spans.filter( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(roots).toHaveLength(1); + const root = roots[0]; + expect(root.status.code).toBe(SpanStatusCode.ERROR); + expect(root.attributes).toMatchObject({ + "appkit.route": route, + "appkit.request.id": "request-early", + }); + expect(String(root.attributes["mlflow.spanInputs"])).toContain(inputKey); + expect(String(root.attributes["mlflow.spanOutputs"]).length).toBeGreaterThan( + 2, + ); + expect( + spans.filter((span) => span.attributes["mlflow.spanType"] === "AGENT"), + ).toHaveLength(1); + const headerIndex = order.findIndex((entry) => + entry.startsWith("header:X-MLflow-Trace-Id:"), + ); + const bodyIndex = order.findIndex((entry) => entry.startsWith("body:")); + expect(headerIndex).toBeGreaterThanOrEqual(0); + expect(bodyIndex).toBeGreaterThan(headerIndex); +} + function seedPlugin(): AgentsPlugin { const plugin = new AgentsPlugin({ dir: false }); // biome-ignore lint/suspicious/noExplicitAny: seed private state @@ -75,6 +159,208 @@ function seedPlugin(): AgentsPlugin { return plugin; } +describe("early HTTP failures create one semantic root before writing", () => { + test("/chat schema failure traces raw input and sets discovery before the body", async () => { + const plugin = seedPlugin(); + const order: string[] = []; + const { res } = mockRes(order); + const req = requestForRoute("chat", { message: "" }); + + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleChat: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleChat(req, res), + ); + + expectEarlyErrorTrace(observed.spans, order, "chat", "message"); + }); + + test("/chat missing-agent lookup finalizes the provisional root as ERROR", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const order: string[] = []; + const { res } = mockRes(order); + const req = requestForRoute("chat", { + message: "hello", + agent: "missing-agent", + }); + + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleChat: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleChat(req, res), + ); + + expectEarlyErrorTrace(observed.spans, order, "chat", "message"); + }); + + test("/chat thread setup failure updates resolved identity and ends one root", async () => { + const plugin = seedPlugin(); + (plugin as any).threadStore = { + get: vi.fn().mockResolvedValue(null), + create: vi.fn().mockRejectedValue(new Error("DB unavailable")), + addMessage: vi.fn(), + }; + const order: string[] = []; + const { res } = mockRes(order); + const req = requestForRoute("chat", { message: "hello" }); + + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleChat: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleChat(req, res), + ); + + expectEarlyErrorTrace(observed.spans, order, "chat", "message"); + const root = observed.spans.find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(root?.attributes).toMatchObject({ + "appkit.agent.name": "default", + "mlflow.trace.user": "alice", + }); + }); + + test("/chat user-context failure still sets discovery before middleware error body", async () => { + const plugin = seedPlugin(); + const order: string[] = []; + const { res, json } = mockRes(order); + const req = requestForRoute("chat", { message: "hello" }, ""); + const originalNodeEnv = process.env.NODE_ENV; + process.env.NODE_ENV = "production"; + try { + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleChat: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleChat(req, res), + ); + expect(observed.error).toBeDefined(); + json({ error: "Authentication failed" }); + expectEarlyErrorTrace(observed.spans, order, "chat", "message"); + } finally { + process.env.NODE_ENV = originalNodeEnv; + } + }); + + describe.each(["invocations", "responses"] as const)("/%s", (route) => { + test("schema failure creates and finalizes one root", async () => { + const plugin = seedPlugin(); + const order: string[] = []; + const { res } = mockRes(order); + const req = requestForRoute(route, { input: "" }); + + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleInvoke: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleInvoke(req, res), + ); + + expectEarlyErrorTrace(observed.spans, order, route, "input"); + }); + + test("missing-agent lookup creates and finalizes one root", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const order: string[] = []; + const { res } = mockRes(order); + const req = requestForRoute(route, { input: "hello" }); + + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleInvoke: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleInvoke(req, res), + ); + + expectEarlyErrorTrace(observed.spans, order, route, "input"); + }); + + test("thread setup failure creates and finalizes one root", async () => { + const plugin = seedPlugin(); + (plugin as any).threadStore = { + create: vi.fn().mockRejectedValue(new Error("DB unavailable")), + addMessage: vi.fn(), + }; + const order: string[] = []; + const { res } = mockRes(order); + const req = requestForRoute(route, { input: "hello" }); + + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleInvoke: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleInvoke(req, res), + ); + + expectEarlyErrorTrace(observed.spans, order, route, "input"); + const root = observed.spans.find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(root?.attributes).toMatchObject({ + "appkit.agent.name": "default", + "mlflow.trace.user": "alice", + }); + }); + + test("user-context failure sets discovery before middleware error body", async () => { + const plugin = seedPlugin(); + const order: string[] = []; + const { res, json } = mockRes(order); + const req = requestForRoute(route, { input: "hello" }, ""); + const originalNodeEnv = process.env.NODE_ENV; + process.env.NODE_ENV = "production"; + try { + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleInvoke: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleInvoke(req, res), + ); + expect(observed.error).toBeDefined(); + json({ error: "Authentication failed" }); + expectEarlyErrorTrace(observed.spans, order, route, "input"); + } finally { + process.env.NODE_ENV = originalNodeEnv; + } + }); + }); +}); + describe("POST /chat ā€” threadStore failure", () => { test("returns 500 when threadStore.get rejects (existing thread path)", async () => { const plugin = seedPlugin(); @@ -459,11 +745,6 @@ describe("POST /chat ā€” trace discovery ordering", () => { order.push(`body:${String(event.type)}`); } }; - // biome-ignore lint/suspicious/noExplicitAny: stub persistence - (plugin as any).threadStore = { - addMessage: vi.fn(), - delete: vi.fn(), - }; const registered = { name: "planner", instructions: "help", @@ -505,7 +786,18 @@ describe("POST /chat ā€” trace discovery ordering", () => { createdAt: new Date(), updatedAt: new Date(), }; - const req = mockReq({ message: "hi" }); + // biome-ignore lint/suspicious/noExplicitAny: seed private route state + (plugin as any).agents.set("planner", registered); + // biome-ignore lint/suspicious/noExplicitAny: seed private route state + (plugin as any).defaultAgentName = "planner"; + // biome-ignore lint/suspicious/noExplicitAny: stub persistence + (plugin as any).threadStore = { + get: vi.fn(), + create: vi.fn().mockResolvedValue(thread), + addMessage: vi.fn(), + delete: vi.fn(), + }; + const req = mockReq({ message: "hi", agent: "planner" }); const { res, setHeader } = mockRes(); setHeader.mockImplementation((name, value) => { if (name === "X-MLflow-Trace-Id") order.push(`header:${String(value)}`); @@ -513,15 +805,12 @@ describe("POST /chat ā€” trace discovery ordering", () => { await ( plugin as unknown as { - _streamAgent: ( + _handleChat: ( request: express.Request, response: express.Response, - agent: unknown, - currentThread: unknown, - userId: string, ) => Promise; } - )._streamAgent(req, res, registered, thread, "alice"); + )._handleChat(req, res); const metadata = streamed[0] as { type?: string; diff --git a/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts index 40a6c72ac..93c32f93c 100644 --- a/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts +++ b/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts @@ -40,6 +40,9 @@ interface AgentTraceIdentity { interface AgentTraceObserver { readonly traceId: string; onEvent(event: AgentEvent): void; + updateIdentity(identity: Partial>): void; + setOutput(output: unknown): void; + recordError(error: unknown, output?: unknown): void; } type RunWithAgentTrace = ( @@ -366,6 +369,39 @@ describe("runWithAgentTrace golden span trees", () => { expect(JSON.stringify(model?.events)).not.toContain("top-secret-token"); }); + test("never exposes a custom Error name through root or model exception.type", async () => { + installTracing(); + const secret = "adapter-secret-name"; + const customError = new Error("Authorization: Bearer message-secret"); + customError.name = `${secret}-${"x".repeat(2_048)}`; + + await expect( + runWithAgentTrace( + identity("responses"), + { message: "hello" }, + async (observer) => { + for (const event of modelEvents({ error: `${secret}-model` })) { + observer.onEvent(event); + } + throw customError; + }, + ), + ).rejects.toBe(customError); + + const exceptionEvents = finishedSpans().flatMap((span) => + span.events.filter((event) => event.name === "exception"), + ); + expect(exceptionEvents).toHaveLength(2); + for (const event of exceptionEvents) { + expect(event.attributes?.["exception.type"]).toBe("Error"); + expect( + String(event.attributes?.["exception.type"]).length, + ).toBeLessThanOrEqual(64); + } + expect(JSON.stringify(exceptionEvents)).not.toContain(secret); + expect(JSON.stringify(exceptionEvents)).not.toContain("message-secret"); + }); + test("records an aborted partial stream as error and omits unavailable aggregate cost", async () => { installTracing(); const abortError = new DOMException("client cancelled", "AbortError"); @@ -420,6 +456,31 @@ describe("runWithAgentTrace golden span trees", () => { expect(JSON.stringify(root?.events)).not.toContain("provider-secret"); }); + test("redacts explicit handled-error output when the operation does not throw", async () => { + installTracing(); + const secret = "adapter-error-secret"; + + await runWithAgentTrace( + identity("invocations"), + { message: "hello" }, + async (observer) => { + observer.recordError(new Error(secret), { + error: secret, + trace_id: "trace-123", + }); + }, + ); + + const root = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(root?.status.code).toBe(SpanStatusCode.ERROR); + expect(root?.attributes["mlflow.spanOutputs"]).toBe( + '{"error":"[REDACTED]","trace_id":"trace-123"}', + ); + expect(JSON.stringify(root?.events)).not.toContain(secret); + }); + test("counts each model_end once and withholds partial aggregate cost", async () => { installTracing(); @@ -460,4 +521,73 @@ describe("runWithAgentTrace golden span trees", () => { ), ).toHaveLength(2); }); + + test("treats costAvailable without costUsd as unavailable for child and root", async () => { + installTracing(); + + const traced = await runWithAgentTrace( + identity("invocations"), + { message: "hello" }, + async (observer) => { + for (const event of modelEvents({ + costAvailable: true, + costUsd: undefined, + })) { + if (event.type === "model_end") { + observer.onEvent({ + ...event, + usage: { + inputTokens: 7, + outputTokens: 2, + totalTokens: 9, + costAvailable: true, + }, + }); + } else { + observer.onEvent(event); + } + } + return { text: "Hello world" }; + }, + ); + + const root = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + const model = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "CHAT_MODEL", + ); + expect(traced.usage.costAvailable).toBe(false); + expect(root?.attributes["appkit.cost.available"]).toBe(false); + expect(root?.attributes["mlflow.llm.cost"]).toBeUndefined(); + expect(model?.attributes["appkit.cost.available"]).toBe(false); + expect(model?.attributes["mlflow.llm.cost"]).toBeUndefined(); + }); + + test("retains a legitimate zero cost as available for child and root", async () => { + installTracing(); + + const traced = await runWithAgentTrace( + identity("invocations"), + { message: "hello" }, + async (observer) => { + for (const event of modelEvents({ costAvailable: true, costUsd: 0 })) { + observer.onEvent(event); + } + return { text: "Hello world" }; + }, + ); + + const priced = finishedSpans().filter((span) => + ["AGENT", "CHAT_MODEL"].includes( + String(span.attributes["mlflow.spanType"]), + ), + ); + expect(traced.usage).toMatchObject({ costAvailable: true, costUsd: 0 }); + expect(priced).toHaveLength(2); + for (const span of priced) { + expect(span.attributes["appkit.cost.available"]).toBe(true); + expect(span.attributes["mlflow.llm.cost"]).toBe(0); + } + }); }); diff --git a/packages/appkit/src/telemetry/agent-tracing/tracer.ts b/packages/appkit/src/telemetry/agent-tracing/tracer.ts index f5a36a8ca..76beead42 100644 --- a/packages/appkit/src/telemetry/agent-tracing/tracer.ts +++ b/packages/appkit/src/telemetry/agent-tracing/tracer.ts @@ -79,6 +79,9 @@ export async function runWithAgentTrace( const completedModels = new Set(); let outputText = ""; let lifecycleError = false; + let reportedError: unknown; + let explicitOutput: unknown; + let hasExplicitOutput = false; const observer: AgentTraceObserver = { traceId, @@ -113,15 +116,40 @@ export async function runWithAgentTrace( if (event.type === "model_end") { if (completedModels.has(event.stepId)) return; completedModels.add(event.stepId); - usage.add(event.usage); + const normalizedUsage = normalizeUsage(event.usage); + usage.add(normalizedUsage); const active = activeModels.get(event.stepId); activeModels.delete(event.stepId); - if (event.error) lifecycleError = true; - if (active) finalizeModelSpan(active, event); + if (event.error) { + lifecycleError = true; + reportedError ??= event.error; + } + if (active) { + finalizeModelSpan(active, { + ...event, + usage: normalizedUsage, + }); + } return; } if (event.type === "status" && event.status === "error") { lifecycleError = true; + reportedError ??= event.error; + } + }, + updateIdentity(next) { + setIdentityAttributes(root, next); + }, + setOutput(output) { + explicitOutput = output; + hasExplicitOutput = true; + }, + recordError(error, output) { + lifecycleError = true; + reportedError ??= error; + if (output !== undefined) { + explicitOutput = output; + hasExplicitOutput = true; } }, }; @@ -169,7 +197,7 @@ export async function runWithAgentTrace( if (lifecycleError && !failed) { recordSafeException( root, - "Agent lifecycle reported an error", + reportedError ?? "Agent lifecycle reported an error", "Agent operation failed", ); } @@ -177,13 +205,18 @@ export async function runWithAgentTrace( const finalUsage = usage.snapshot(); setRootUsageAttributes(root, finalUsage); const finalOutputText = outputText || textFromValue(value); + const finalOutput = hasExplicitOutput + ? explicitOutput + : failed + ? { text: finalOutputText, error: errorValue(operationError) } + : outputText || textFromValue(value) + ? { text: finalOutputText } + : value; setCapturedAttribute( root, "mlflow.spanOutputs", - failed - ? { text: finalOutputText, error: errorValue(operationError) } - : { text: finalOutputText }, - failed ? ["error"] : undefined, + finalOutput, + failed || lifecycleError ? ["error"] : undefined, ); root.setStatus({ code: @@ -268,12 +301,32 @@ function setRootUsageAttributes(span: Span, usage: AgentUsage): void { } function setCostAttributes(span: Span, usage: AgentUsage): void { - span.setAttribute("appkit.cost.available", usage.costAvailable); - if (usage.costAvailable && usage.costUsd !== undefined) { + const costAvailable = hasCompleteCost(usage); + span.setAttribute("appkit.cost.available", costAvailable); + if (costAvailable && usage.costUsd !== undefined) { span.setAttribute("mlflow.llm.cost", usage.costUsd); } } +function normalizeUsage(usage: AgentUsage): AgentUsage { + const { costUsd, ...withoutCost } = usage; + const costAvailable = hasCompleteCost(usage); + return { + ...withoutCost, + ...(costAvailable ? { costUsd } : {}), + costAvailable, + }; +} + +function hasCompleteCost(usage: AgentUsage): boolean { + return ( + usage.costAvailable && + typeof usage.costUsd === "number" && + Number.isFinite(usage.costUsd) && + usage.costUsd >= 0 + ); +} + function mlflowTokenUsage(usage: AgentUsage): Record { return { input_tokens: usage.inputTokens, @@ -312,11 +365,38 @@ function recordSafeException( .value, ); span.recordException({ - name: error instanceof Error ? error.name : "Error", + name: "Error", message: publicMessage, }); } +function setIdentityAttributes( + span: Span, + identity: Partial>, +): void { + if (identity.appName !== undefined) { + span.setAttribute( + "appkit.app.name", + resolveAgentTraceAppName(identity.appName), + ); + } + if (identity.agentName !== undefined) { + span.setAttribute("appkit.agent.name", identity.agentName); + } + if (identity.sessionId !== undefined) { + span.setAttribute("mlflow.trace.session", identity.sessionId); + } + if (identity.userId !== undefined) { + span.setAttribute("mlflow.trace.user", identity.userId); + } + if (identity.requestId !== undefined) { + span.setAttribute("appkit.request.id", identity.requestId); + } + if (identity.threadId !== undefined) { + span.setAttribute("appkit.thread.id", identity.threadId); + } +} + function errorValue(error: unknown): string { return error instanceof Error ? error.message diff --git a/packages/appkit/src/telemetry/agent-tracing/types.ts b/packages/appkit/src/telemetry/agent-tracing/types.ts index fe8796f49..326522b6a 100644 --- a/packages/appkit/src/telemetry/agent-tracing/types.ts +++ b/packages/appkit/src/telemetry/agent-tracing/types.ts @@ -34,6 +34,9 @@ export interface AgentTraceObserver { /** MLflow V4 identity when UC is active; otherwise the 32-hex OTel trace ID. */ readonly traceId: string; onEvent(event: AgentEvent): void; + updateIdentity(identity: Partial>): void; + setOutput(output: unknown): void; + recordError(error: unknown, output?: unknown): void; } export interface AgentTraceResult { From 7dd81614e3173f10a786c8a6aad13111884c18fd Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 19:53:19 -0700 Subject: [PATCH 12/31] feat: trace AppKit tools approvals and memory Signed-off-by: Adam Gurary --- packages/appkit/src/core/plugin-context.ts | 6 + packages/appkit/src/plugins/agents/agents.ts | 307 +++++++---- .../agents/tests/dispatch-tool-call.test.ts | 489 +++++++++++++++++- .../plugins/agents/tests/thread-store.test.ts | 222 +++++++- .../agents/tests/tool-approval-gate.test.ts | 209 ++++++++ .../appkit/src/plugins/agents/thread-store.ts | 135 +++++ .../src/plugins/agents/tool-approval-gate.ts | 142 ++++- .../telemetry/tests/plugin-telemetry.test.ts | 61 +++ 8 files changed, 1454 insertions(+), 117 deletions(-) diff --git a/packages/appkit/src/core/plugin-context.ts b/packages/appkit/src/core/plugin-context.ts index 4f08dcd91..53abc9772 100644 --- a/packages/appkit/src/core/plugin-context.ts +++ b/packages/appkit/src/core/plugin-context.ts @@ -192,6 +192,10 @@ export class PluginContext { args: unknown, signal?: AbortSignal, timeoutMs: number = 300_000, + traceIdentity: { name: string; source: string } = { + name: `${pluginName}.${toolName}`, + source: "toolkit", + }, ): Promise { const provider = this.toolProviders.get(pluginName); if (!provider) { @@ -204,6 +208,8 @@ export class PluginContext { const operationName = `executeTool:${pluginName}.${toolName}`; return tracer.startActiveSpan(operationName, async (span) => { + span.setAttribute?.("appkit.tool.name", traceIdentity.name); + span.setAttribute?.("appkit.tool.source", traceIdentity.source); const timeoutSignal = AbortSignal.timeout(timeoutMs); const combinedSignal = signal ? AbortSignal.any([signal, timeoutSignal]) diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 7f8051506..a46ee3517 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -1,5 +1,6 @@ import { randomUUID } from "node:crypto"; import path from "node:path"; +import { type Span, SpanStatusCode, trace } from "@opentelemetry/api"; import type express from "express"; import pc from "picocolors"; import type { @@ -13,6 +14,7 @@ import type { ResponseStreamEvent, Thread, ToolAnnotations, + ToolEffect, ToolProvider, } from "shared"; import { @@ -53,6 +55,7 @@ import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; import { type AgentTraceObserver, + captureTraceValue, resolveAgentTraceAppName, runWithAgentTrace, } from "../../telemetry/agent-tracing"; @@ -66,13 +69,93 @@ import { chatRequestSchema, invocationsRequestSchema, } from "./schemas"; -import { InMemoryThreadStore } from "./thread-store"; +import { InMemoryThreadStore, TracedThreadStore } from "./thread-store"; import { ToolApprovalGate } from "./tool-approval-gate"; const logger = createLogger("agents"); const DEFAULT_AGENTS_DIR = "./config/agents"; +const agentOperationTracer = () => + trace.getTracer("@databricks/appkit-agent-tracing"); + +export async function traceToolCall( + input: { + name: string; + source: string; + effect?: ToolEffect; + args: unknown; + }, + operation: (span: Span) => Promise, +): Promise { + return agentOperationTracer().startActiveSpan( + `${input.name} tool`, + { + attributes: { + "mlflow.spanType": "TOOL", + "appkit.tool.name": input.name, + "appkit.tool.source": input.source, + ...(input.effect ? { "appkit.tool.effect": input.effect } : {}), + }, + }, + async (span) => { + const startedAt = Date.now(); + setToolCapturedAttribute(span, "mlflow.spanInputs", input.args); + try { + const result = await operation(span); + setToolCapturedAttribute(span, "mlflow.spanOutputs", result); + span.setAttribute("appkit.tool.state", "completed"); + span.setStatus({ code: SpanStatusCode.OK }); + return result; + } catch (error) { + span.setAttribute("appkit.tool.state", "failed"); + recordSafeToolFailure(span, error); + throw error; + } finally { + span.setAttribute( + "appkit.tool.duration_ms", + Math.max(0, Date.now() - startedAt), + ); + span.end(); + } + }, + ); +} + +function setToolCapturedAttribute( + span: Span, + key: string, + value: unknown, +): void { + const captured = captureTraceValue(value); + span.setAttribute(key, captured.value); + span.setAttribute(`${key}.original_bytes`, captured.originalBytes); + span.setAttribute(`${key}.sha256`, captured.sha256); + span.setAttribute(`${key}.truncated`, captured.truncated); +} + +function recordSafeToolFailure(span: Span, error: unknown): void { + const failure = captureTraceValue( + { + error: + error instanceof Error + ? error.message + : String(error ?? "Unknown error"), + }, + { redactKeys: ["error"] }, + ); + span.setAttribute("appkit.error", failure.value); + span.setAttribute("mlflow.spanOutputs", failure.value); + span.setAttribute("mlflow.spanOutputs.original_bytes", failure.originalBytes); + span.setAttribute("mlflow.spanOutputs.sha256", failure.sha256); + span.setAttribute("mlflow.spanOutputs.truncated", failure.truncated); + span.recordException({ name: "Error", message: "Tool operation failed" }); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: "Tool operation failed", + }); +} + /** * Context flag recorded on the in-memory AgentDefinition to indicate whether * it came from markdown (file) or from user code. Drives the asymmetric @@ -165,9 +248,12 @@ export class AgentsPlugin extends Plugin implements ToolProvider { super(config); this.config = config; if (config.threadStore) { - this.threadStore = config.threadStore; + this.threadStore = + config.threadStore instanceof TracedThreadStore + ? config.threadStore + : new TracedThreadStore(config.threadStore); } else { - this.threadStore = new InMemoryThreadStore(); + this.threadStore = new TracedThreadStore(new InMemoryThreadStore()); if (process.env.NODE_ENV === "production") { logger.warn( "InMemoryThreadStore is in use in a production build (NODE_ENV=production). " + @@ -1489,104 +1575,131 @@ export class AgentsPlugin extends Plugin implements ToolProvider { args: unknown, depth: number, ): Promise { - if (runState.toolCallsUsed.count >= runState.limits.maxToolCalls) { - runState.abortController.abort( - new Error( - `Tool-call budget exhausted (limit ${runState.limits.maxToolCalls}).`, - ), - ); - throw new Error( - `Tool-call budget exhausted (limit ${runState.limits.maxToolCalls}). Raise agents({ limits: { maxToolCalls } }) or review the agent's tool-selection logic.`, - ); - } - runState.toolCallsUsed.count++; - const entry = toolIndex.get(name); - if (!entry) throw new Error(`Unknown tool: ${name}`); - - if ( - runState.approvalPolicy.requireForDestructive && - requiresApproval(entry.def.annotations) - ) { - const approvalId = randomUUID(); - for (const ev of runState.translator.translate({ - type: "approval_pending", - approvalId, - streamId: runState.requestId, - toolName: name, + return traceToolCall( + { + name, + source: entry?.source ?? "unknown", + effect: entry?.def.annotations?.effect, args, - annotations: entry.def.annotations, - })) { - runState.outboundEvents.push(ev); - } - const decision = await this.approvalGate.wait({ - approvalId, - streamId: runState.requestId, - userId: runState.userId, - timeoutMs: runState.approvalPolicy.timeoutMs, - }); - if (decision === "deny") { - return `Tool execution denied by user approval gate (tool: ${name}).`; - } - } + }, + async () => { + if (runState.toolCallsUsed.count >= runState.limits.maxToolCalls) { + runState.abortController.abort( + new Error( + `Tool-call budget exhausted (limit ${runState.limits.maxToolCalls}).`, + ), + ); + throw new Error( + `Tool-call budget exhausted (limit ${runState.limits.maxToolCalls}). Raise agents({ limits: { maxToolCalls } }) or review the agent's tool-selection logic.`, + ); + } + runState.toolCallsUsed.count++; + + if (!entry) throw new Error(`Unknown tool: ${name}`); + + if ( + runState.approvalPolicy.requireForDestructive && + requiresApproval(entry.def.annotations) + ) { + const approvalId = randomUUID(); + for (const ev of runState.translator.translate({ + type: "approval_pending", + approvalId, + streamId: runState.requestId, + toolName: name, + args, + annotations: entry.def.annotations, + })) { + runState.outboundEvents.push(ev); + } + const decision = await this.approvalGate.wait({ + approvalId, + streamId: runState.requestId, + userId: runState.userId, + timeoutMs: runState.approvalPolicy.timeoutMs, + toolName: name, + effect: entry.def.annotations?.effect, + args, + }); + if (decision === "deny") { + return `Tool execution denied by user approval gate (tool: ${name}).`; + } + } - let result: unknown; - if (entry.source === "toolkit") { - if (!this.context) { - throw new Error( - "Plugin tool execution requires PluginContext; this should never happen through createApp", - ); - } - result = await this.context.executeTool( - runState.req, - entry.pluginName, - entry.localName, - args, - runState.signal, - runState.limits.toolCallTimeoutMs, - ); - } else if (entry.source === "function") { - // Function tools declare their parameters as a JSON-object schema, - // so adapters always serialize `args` as an object. A non-object - // value here means the upstream model emitted malformed tool-call - // JSON; surface a clear error rather than silently passing through - // a wrong-shape value the tool will then choke on. - if (typeof args !== "object" || args === null || Array.isArray(args)) { - throw new Error( - `Function tool '${name}' received non-object arguments (got ${args === null ? "null" : Array.isArray(args) ? "array" : typeof args}); expected a JSON object.`, - ); - } - result = await entry.functionTool.execute( - args as Record, - ); - } else if (entry.source === "mcp") { - if (!this.mcpClient) throw new Error("MCP client not connected"); - const oboToken = runState.req.headers["x-forwarded-access-token"]; - const mcpAuth = - typeof oboToken === "string" - ? { Authorization: `Bearer ${oboToken}` } - : undefined; - result = await this.mcpClient.callTool(entry.mcpToolName, args, mcpAuth); - } else if (entry.source === "subagent") { - const childAgent = this.agents.get(entry.agentName); - if (!childAgent) - throw new Error(`Sub-agent not found: ${entry.agentName}`); - result = await this.runSubAgent(runState, childAgent, args, depth + 1); - } else if (entry.source === "hosted-supervisor") { - // Defense-in-depth: should never fire. Hosted-supervisor entries are - // routed via `AgentInput.extensions` and the SA endpoint executes - // them server-side; their `def` is filtered out of the adapter's - // `tools` array, so the model never sees a callable schema for them. - // If we reach here, the agent is paired with a non-SA adapter that - // somehow surfaced the placeholder def to the model ā€” surface a - // clear error rather than crash later in `normalizeToolResult`. - throw new Error( - `Tool '${name}' is a hosted-supervisor tool and cannot be invoked from the Node process. ` + - "It is executed server-side by the Databricks AI Gateway and is only reachable when the agent's model is a Supervisor API adapter.", - ); - } + let result: unknown; + if (entry.source === "toolkit") { + if (!this.context) { + throw new Error( + "Plugin tool execution requires PluginContext; this should never happen through createApp", + ); + } + result = await this.context.executeTool( + runState.req, + entry.pluginName, + entry.localName, + args, + runState.signal, + runState.limits.toolCallTimeoutMs, + { name, source: entry.source }, + ); + } else if (entry.source === "function") { + // Function tools declare their parameters as a JSON-object schema, + // so adapters always serialize `args` as an object. A non-object + // value here means the upstream model emitted malformed tool-call + // JSON; surface a clear error rather than silently passing through + // a wrong-shape value the tool will then choke on. + if ( + typeof args !== "object" || + args === null || + Array.isArray(args) + ) { + throw new Error( + `Function tool '${name}' received non-object arguments (got ${args === null ? "null" : Array.isArray(args) ? "array" : typeof args}); expected a JSON object.`, + ); + } + result = await entry.functionTool.execute( + args as Record, + ); + } else if (entry.source === "mcp") { + if (!this.mcpClient) throw new Error("MCP client not connected"); + const oboToken = runState.req.headers["x-forwarded-access-token"]; + const mcpAuth = + typeof oboToken === "string" + ? { Authorization: `Bearer ${oboToken}` } + : undefined; + result = await this.mcpClient.callTool( + entry.mcpToolName, + args, + mcpAuth, + ); + } else if (entry.source === "subagent") { + const childAgent = this.agents.get(entry.agentName); + if (!childAgent) + throw new Error(`Sub-agent not found: ${entry.agentName}`); + result = await this.runSubAgent( + runState, + childAgent, + args, + depth + 1, + ); + } else if (entry.source === "hosted-supervisor") { + // Defense-in-depth: should never fire. Hosted-supervisor entries are + // routed via `AgentInput.extensions` and the SA endpoint executes + // them server-side; their `def` is filtered out of the adapter's + // `tools` array, so the model never sees a callable schema for them. + // If we reach here, the agent is paired with a non-SA adapter that + // somehow surfaced the placeholder def to the model ā€” surface a + // clear error rather than crash later in `normalizeToolResult`. + throw new Error( + `Tool '${name}' is a hosted-supervisor tool and cannot be invoked from the Node process. ` + + "It is executed server-side by the Databricks AI Gateway and is only reachable when the agent's model is a Supervisor API adapter.", + ); + } - return normalizeToolResult(result); + return normalizeToolResult(result); + }, + ); } /** diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index 2767766c3..a84ffc2fc 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -1,6 +1,23 @@ +import { context, SpanStatusCode, trace } from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; import type express from "express"; -import { beforeEach, describe, expect, test, vi } from "vitest"; +import { + afterAll, + beforeAll, + beforeEach, + describe, + expect, + test, + vi, +} from "vitest"; import { CacheManager } from "../../../cache"; +import { runWithAgentTrace } from "../../../telemetry/agent-tracing"; import { AgentsPlugin } from "../agents"; /** @@ -20,6 +37,17 @@ import { AgentsPlugin } from "../agents"; * common `RunState` object. Tests below pin those guarantees. */ +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); +}); + +afterAll(() => { + context.disable(); +}); + beforeEach(() => { // dispatchToolCall is exercised without going through setup(), so we // need the cache singleton to be initialised before the plugin reads it. @@ -35,6 +63,44 @@ beforeEach(() => { }; }); +async function captureSpans( + operation: () => Promise, +): Promise<{ spans: ReadableSpan[]; error?: unknown }> { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + provider.getTracer(name, version), + ); + let error: unknown; + let spans: ReadableSpan[] = []; + try { + await operation(); + } catch (caught) { + error = caught; + } finally { + await provider.forceFlush(); + spans = exporter.getFinishedSpans(); + getTracerSpy.mockRestore(); + await provider.shutdown(); + } + return { + spans, + ...(error !== undefined ? { error } : {}), + }; +} + +function semanticSpan(spans: ReadableSpan[], spanType: string): ReadableSpan { + const span = spans.find( + (candidate) => candidate.attributes["mlflow.spanType"] === spanType, + ); + expect(span, `missing ${spanType} span`).toBeDefined(); + return span as ReadableSpan; +} + function mockReq(): express.Request { return { body: {}, @@ -99,6 +165,427 @@ function callDispatch( ); } +describe("dispatchToolCall ā€” semantic TOOL spans", () => { + test("creates one TOOL descendant for inline, toolkit, MCP, and local sub-agent dispatch", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const { runState } = makeRunState(plugin); + // These are the slow/external boundaries for toolkit and MCP execution; + // dispatch and span creation remain real. + // biome-ignore lint/suspicious/noExplicitAny: seed private integration seams + (plugin as any).context = { + executeTool: vi.fn().mockResolvedValue({ rows: [1, 2] }), + }; + // biome-ignore lint/suspicious/noExplicitAny: seed private integration seam + (plugin as any).mcpClient = { + callTool: vi.fn().mockResolvedValue({ content: "remote" }), + }; + // biome-ignore lint/suspicious/noExplicitAny: isolate dispatch from adapter streaming + (plugin as any).runSubAgent = vi.fn().mockResolvedValue("child output"); + // biome-ignore lint/suspicious/noExplicitAny: seed the child registry lookup + (plugin as any).agents.set("researcher", { name: "researcher" }); + + const toolIndex = new Map([ + [ + "inline", + { + source: "function", + def: { + name: "inline", + description: "inline", + parameters: { type: "object" }, + annotations: { effect: "read" }, + }, + functionTool: { + execute: vi.fn().mockResolvedValue({ answer: 42 }), + }, + }, + ], + [ + "analytics.query", + { + source: "toolkit", + pluginName: "analytics", + localName: "query", + def: { + name: "analytics.query", + description: "query", + parameters: { type: "object" }, + annotations: { effect: "read" }, + }, + }, + ], + [ + "remote_lookup", + { + source: "mcp", + mcpToolName: "remote_lookup", + def: { + name: "remote_lookup", + description: "remote", + parameters: { type: "object" }, + annotations: { effect: "read" }, + }, + }, + ], + [ + "agent-researcher", + { + source: "subagent", + agentName: "researcher", + def: { + name: "agent-researcher", + description: "delegate", + parameters: { type: "object" }, + }, + }, + ], + ]); + + const observed = await captureSpans(() => + runWithAgentTrace( + { + appName: "trace-test", + agentName: "planner", + route: "chat", + sessionId: "session-1", + userId: "alice", + requestId: "request-1", + threadId: "thread-1", + }, + { message: "run tools" }, + async () => { + await callDispatch(plugin, { + runState, + toolIndex, + name: "inline", + args: { password: "do-not-log", question: "meaning" }, + }); + await callDispatch(plugin, { + runState, + toolIndex, + name: "analytics.query", + args: { sql: "SELECT 1" }, + }); + await callDispatch(plugin, { + runState, + toolIndex, + name: "remote_lookup", + args: { id: 7 }, + }); + await callDispatch(plugin, { + runState, + toolIndex, + name: "agent-researcher", + args: { input: "investigate" }, + }); + return "done"; + }, + ), + ); + + expect(observed.error).toBeUndefined(); + const root = semanticSpan(observed.spans, "AGENT"); + const tools = observed.spans.filter( + (span) => span.attributes["mlflow.spanType"] === "TOOL", + ); + expect(tools).toHaveLength(4); + expect( + tools.map((span) => [ + span.attributes["appkit.tool.name"], + span.attributes["appkit.tool.source"], + span.attributes["appkit.tool.effect"], + ]), + ).toEqual([ + ["inline", "function", "read"], + ["analytics.query", "toolkit", "read"], + ["remote_lookup", "mcp", "read"], + ["agent-researcher", "subagent", undefined], + ]); + expect( + tools.every( + (span) => + span.parentSpanContext?.spanId === root.spanContext().spanId && + span.status.code === SpanStatusCode.OK && + typeof span.attributes["appkit.tool.duration_ms"] === "number", + ), + ).toBe(true); + expect( + JSON.parse(String(tools[0].attributes["mlflow.spanInputs"])), + ).toEqual({ password: "[REDACTED]", question: "meaning" }); + expect(JSON.parse(String(tools[0].attributes["mlflow.spanOutputs"]))).toBe( + '{"answer":42}', + ); + expect(JSON.parse(String(tools[1].attributes["mlflow.spanOutputs"]))).toBe( + '{"rows":[1,2]}', + ); + expect(JSON.parse(String(tools[2].attributes["mlflow.spanOutputs"]))).toBe( + '{"content":"remote"}', + ); + expect(JSON.parse(String(tools[3].attributes["mlflow.spanOutputs"]))).toBe( + "child output", + ); + }); + + test("traces unknown tools as a failed TOOL without exposing the thrown detail", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const { runState } = makeRunState(plugin); + + const observed = await captureSpans(() => + callDispatch(plugin, { + runState, + toolIndex: new Map(), + name: "missing_secret_tool", + args: { apiKey: "sensitive" }, + }), + ); + + expect(observed.error).toEqual( + new Error("Unknown tool: missing_secret_tool"), + ); + const tool = semanticSpan(observed.spans, "TOOL"); + expect(tool.attributes).toMatchObject({ + "appkit.tool.name": "missing_secret_tool", + "appkit.tool.source": "unknown", + "appkit.error": '{"error":"[REDACTED]"}', + }); + expect(tool.status.code).toBe(SpanStatusCode.ERROR); + expect(tool.attributes["appkit.tool.duration_ms"]).toEqual( + expect.any(Number), + ); + expect( + JSON.stringify({ attributes: tool.attributes, events: tool.events }), + ).not.toContain("Unknown tool: missing_secret_tool"); + }); + + test("traces malformed function arguments and never invokes the function body", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const { runState } = makeRunState(plugin); + const execute = vi.fn(); + const toolIndex = new Map([ + [ + "object_only", + { + source: "function", + def: { + name: "object_only", + description: "object", + parameters: { type: "object" }, + }, + functionTool: { execute }, + }, + ], + ]); + + const observed = await captureSpans(() => + callDispatch(plugin, { + runState, + toolIndex, + name: "object_only", + args: ["wrong"], + }), + ); + + expect(observed.error).toEqual( + new Error( + "Function tool 'object_only' received non-object arguments (got array); expected a JSON object.", + ), + ); + expect(execute).not.toHaveBeenCalled(); + const tool = semanticSpan(observed.spans, "TOOL"); + expect(tool.attributes["appkit.tool.source"]).toBe("function"); + expect(tool.status.code).toBe(SpanStatusCode.ERROR); + }); + + test("records a toolkit timeout as a sanitized failed TOOL", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const { runState } = makeRunState(plugin); + // biome-ignore lint/suspicious/noExplicitAny: isolate the PluginContext boundary + (plugin as any).context = { + executeTool: vi + .fn() + .mockRejectedValue( + new DOMException("private timeout detail", "TimeoutError"), + ), + }; + const toolIndex = new Map([ + [ + "analytics.slow", + { + source: "toolkit", + pluginName: "analytics", + localName: "slow", + def: { + name: "analytics.slow", + description: "slow", + parameters: { type: "object" }, + }, + }, + ], + ]); + + const observed = await captureSpans(() => + callDispatch(plugin, { + runState, + toolIndex, + name: "analytics.slow", + args: {}, + }), + ); + + expect(observed.error).toBeInstanceOf(DOMException); + const tool = semanticSpan(observed.spans, "TOOL"); + expect(tool.status.code).toBe(SpanStatusCode.ERROR); + expect(tool.attributes["appkit.error"]).toBe('{"error":"[REDACTED]"}'); + expect( + JSON.stringify({ attributes: tool.attributes, events: tool.events }), + ).not.toContain("private timeout detail"); + }); + + test("records a function failure as a sanitized failed TOOL", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const { runState } = makeRunState(plugin); + const toolIndex = new Map([ + [ + "explode", + { + source: "function", + def: { + name: "explode", + description: "fail", + parameters: { type: "object" }, + }, + functionTool: { + execute: async () => { + throw new Error("database password hunter2"); + }, + }, + }, + ], + ]); + + const observed = await captureSpans(() => + callDispatch(plugin, { + runState, + toolIndex, + name: "explode", + args: {}, + }), + ); + + expect(observed.error).toEqual(new Error("database password hunter2")); + const tool = semanticSpan(observed.spans, "TOOL"); + expect(tool.status.code).toBe(SpanStatusCode.ERROR); + expect(tool.attributes["appkit.error"]).toBe('{"error":"[REDACTED]"}'); + expect( + JSON.stringify({ attributes: tool.attributes, events: tool.events }), + ).not.toContain("hunter2"); + }); +}); + +describe("dispatchToolCall ā€” semantic approval descendants", () => { + function destructiveTool(execute: ReturnType) { + return new Map([ + [ + "delete_user", + { + source: "function", + def: { + name: "delete_user", + description: "delete", + parameters: { type: "object" }, + annotations: { effect: "destructive" }, + }, + functionTool: { execute }, + }, + ], + ]); + } + + test("nests an approved CHAIN under TOOL and runs the body only after approve", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const { runState, pushed } = makeRunState(plugin); + const order: string[] = []; + const execute = vi.fn(async () => { + order.push("tool"); + return "deleted"; + }); + + const observed = await captureSpans(async () => { + const pending = callDispatch(plugin, { + runState, + toolIndex: destructiveTool(execute), + name: "delete_user", + args: { userId: 7, password: "do-not-log" }, + }); + order.push("waiting"); + expect(execute).not.toHaveBeenCalled(); + const approvalId = (pushed[0] as { approvalId: string }).approvalId; + order.push("approved"); + // biome-ignore lint/suspicious/noExplicitAny: exercise the real private gate + (plugin as any).approvalGate.submit({ + approvalId, + userId: "alice", + decision: "approve", + }); + await pending; + }); + + expect(observed.error).toBeUndefined(); + expect(order).toEqual(["waiting", "approved", "tool"]); + const tool = semanticSpan(observed.spans, "TOOL"); + const approval = semanticSpan(observed.spans, "CHAIN"); + expect(approval.parentSpanContext?.spanId).toBe(tool.spanContext().spanId); + expect(approval.attributes).toMatchObject({ + "appkit.approval.decision": "approve", + "appkit.approval.state": "approved", + "appkit.approval.tool_name": "delete_user", + "appkit.approval.duration_ms": expect.any(Number), + }); + expect( + JSON.parse(String(approval.attributes["mlflow.spanInputs"])), + ).toEqual({ + password: "[REDACTED]", + userId: 7, + }); + expect(tool.status.code).toBe(SpanStatusCode.OK); + }); + + test("records denial and leaves the tool body idle", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const { runState, pushed } = makeRunState(plugin); + const execute = vi.fn(); + let result: unknown; + + const observed = await captureSpans(async () => { + const pending = callDispatch(plugin, { + runState, + toolIndex: destructiveTool(execute), + name: "delete_user", + args: { userId: 7 }, + }); + const approvalId = (pushed[0] as { approvalId: string }).approvalId; + // biome-ignore lint/suspicious/noExplicitAny: exercise the real private gate + (plugin as any).approvalGate.submit({ + approvalId, + userId: "alice", + decision: "deny", + }); + result = await pending; + }); + + expect(observed.error).toBeUndefined(); + expect(result).toBe( + "Tool execution denied by user approval gate (tool: delete_user).", + ); + expect(execute).not.toHaveBeenCalled(); + const approval = semanticSpan(observed.spans, "CHAIN"); + expect(approval.attributes).toMatchObject({ + "appkit.approval.decision": "deny", + "appkit.approval.state": "denied", + }); + }); +}); + describe("dispatchToolCall ā€” approval gate honours `effect`", () => { test('fires for `effect: "destructive"` even without legacy `destructive: true`', async () => { // Regression for finding #1 on PR #304: the gate previously checked diff --git a/packages/appkit/src/plugins/agents/tests/thread-store.test.ts b/packages/appkit/src/plugins/agents/tests/thread-store.test.ts index ed4f70bab..c64184bf4 100644 --- a/packages/appkit/src/plugins/agents/tests/thread-store.test.ts +++ b/packages/appkit/src/plugins/agents/tests/thread-store.test.ts @@ -1,6 +1,56 @@ -import { describe, expect, test } from "vitest"; +import { SpanStatusCode, trace } from "@opentelemetry/api"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import type { Message, Thread, ThreadStore } from "shared"; +import { describe, expect, test, vi } from "vitest"; +import { AgentsPlugin } from "../agents"; import { InMemoryThreadStore } from "../thread-store"; +async function captureSpans( + operation: () => Promise, +): Promise<{ spans: ReadableSpan[]; error?: unknown }> { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + provider.getTracer(name, version), + ); + let error: unknown; + let spans: ReadableSpan[] = []; + try { + await operation(); + } catch (caught) { + error = caught; + } finally { + await provider.forceFlush(); + spans = exporter.getFinishedSpans(); + getTracerSpy.mockRestore(); + await provider.shutdown(); + } + return { spans, ...(error !== undefined ? { error } : {}) }; +} + +function configuredStore(backing?: ThreadStore): ThreadStore { + const plugin = new AgentsPlugin({ + dir: false, + ...(backing ? { threadStore: backing } : {}), + }); + return (plugin as unknown as { threadStore: ThreadStore }).threadStore; +} + +function memorySpans(spans: ReadableSpan[]): ReadableSpan[] { + return spans.filter( + (span) => span.attributes["mlflow.spanType"] === "MEMORY", + ); +} + describe("InMemoryThreadStore", () => { test("create() returns a new thread with the given userId", async () => { const store = new InMemoryThreadStore(); @@ -136,3 +186,173 @@ describe("InMemoryThreadStore", () => { expect(user2Threads).toHaveLength(1); }); }); + +describe("TracedThreadStore", () => { + test("wraps configured stores and traces create, get hit/miss, list, add, and delete", async () => { + const store = configuredStore(new InMemoryThreadStore()); + let created: Thread | undefined; + let hit: Thread | null = null; + let miss: Thread | null = null; + let listed: Thread[] = []; + let deleted = false; + const message: Message = { + id: "message-1", + role: "user", + content: "remember this complete message", + createdAt: new Date("2026-08-11T12:00:00.000Z"), + }; + + const observed = await captureSpans(async () => { + created = await store.create("user-7"); + hit = await store.get(created.id, "user-7"); + miss = await store.get("thread-missing", "user-7"); + listed = await store.list("user-7"); + await store.addMessage(created.id, "user-7", message); + deleted = await store.delete(created.id, "user-7"); + }); + + expect(observed.error).toBeUndefined(); + expect((hit as Thread | null)?.id).toBe(created?.id); + expect(miss).toBeNull(); + expect(listed).toHaveLength(1); + expect(deleted).toBe(true); + const spans = memorySpans(observed.spans); + expect(spans).toHaveLength(6); + expect( + spans.map((span) => [ + span.attributes["appkit.memory.operation"], + span.attributes["appkit.memory.state"], + ]), + ).toEqual([ + ["create", "created"], + ["get", "hit"], + ["get", "miss"], + ["list", "completed"], + ["addMessage", "completed"], + ["delete", "deleted"], + ]); + expect( + spans.every( + (span) => + span.status.code === SpanStatusCode.OK && + typeof span.attributes["appkit.memory.duration_ms"] === "number", + ), + ).toBe(true); + expect( + JSON.parse(String(spans[1].attributes["mlflow.spanInputs"])), + ).toEqual({ + threadId: created?.id, + userId: "user-7", + }); + expect( + JSON.parse(String(spans[2].attributes["mlflow.spanOutputs"])), + ).toBeNull(); + expect( + JSON.parse(String(spans[4].attributes["mlflow.spanInputs"])), + ).toEqual({ + message: { + content: "remember this complete message", + createdAt: "2026-08-11T12:00:00.000Z", + id: "message-1", + role: "user", + }, + threadId: created?.id, + userId: "user-7", + }); + expect(JSON.parse(String(spans[5].attributes["mlflow.spanOutputs"]))).toBe( + true, + ); + }); + + test("wraps the default store before any route can use it", async () => { + const store = configuredStore(); + const observed = await captureSpans(() => store.create("default-user")); + + expect(observed.error).toBeUndefined(); + expect(memorySpans(observed.spans)).toHaveLength(1); + expect( + memorySpans(observed.spans)[0].attributes["appkit.memory.operation"], + ).toBe("create"); + }); + + test("records backing-store failures without logging their secret detail", async () => { + const backing: ThreadStore = { + create: async () => { + throw new Error("not used"); + }, + get: async () => { + throw new Error("postgres password super-secret-value"); + }, + list: async () => [], + addMessage: async () => {}, + delete: async () => false, + }; + const store = configuredStore(backing); + + const observed = await captureSpans(() => store.get("thread-7", "user-7")); + + expect(observed.error).toEqual( + new Error("postgres password super-secret-value"), + ); + const [span] = memorySpans(observed.spans); + expect(span.status.code).toBe(SpanStatusCode.ERROR); + expect(span.attributes).toMatchObject({ + "appkit.memory.operation": "get", + "appkit.memory.state": "failed", + "appkit.error": '{"error":"[REDACTED]"}', + "appkit.memory.duration_ms": expect.any(Number), + }); + expect( + JSON.stringify({ attributes: span.attributes, events: span.events }), + ).not.toContain("super-secret-value"); + }); + + test("applies central redaction and truncation to message and thread values", async () => { + const secretThread = { + id: "thread-secret", + userId: "user-secret", + messages: [], + createdAt: new Date("2026-08-11T12:00:00.000Z"), + updatedAt: new Date("2026-08-11T12:00:00.000Z"), + user: { secret: "private-user-value" }, + thread: { secret: "private-thread-value" }, + } as Thread; + const backing: ThreadStore = { + create: async () => secretThread, + get: async () => secretThread, + list: async () => [secretThread], + addMessage: async () => {}, + delete: async () => true, + }; + const store = configuredStore(backing); + + const observed = await captureSpans(async () => { + await store.get("thread-secret", "user-secret"); + await store.addMessage("thread-secret", "user-secret", { + id: "long-message", + role: "user", + content: "x".repeat(70 * 1024), + createdAt: new Date("2026-08-11T12:00:00.000Z"), + }); + }); + + expect(observed.error).toBeUndefined(); + const [getSpan, addSpan] = memorySpans(observed.spans); + expect(JSON.stringify(getSpan.attributes)).not.toContain( + "private-user-value", + ); + expect(JSON.stringify(getSpan.attributes)).not.toContain( + "private-thread-value", + ); + expect(String(getSpan.attributes["mlflow.spanOutputs"])).toContain( + '"secret":"[REDACTED]"', + ); + expect(addSpan.attributes["mlflow.spanInputs.truncated"]).toBe(true); + expect( + addSpan.attributes["mlflow.spanInputs.original_bytes"], + ).toBeGreaterThan(64 * 1024); + expect( + Buffer.byteLength(String(addSpan.attributes["mlflow.spanInputs"])), + ).toBeLessThanOrEqual(64 * 1024); + }); +}); diff --git a/packages/appkit/src/plugins/agents/tests/tool-approval-gate.test.ts b/packages/appkit/src/plugins/agents/tests/tool-approval-gate.test.ts index 1e17ddf63..e32e81145 100644 --- a/packages/appkit/src/plugins/agents/tests/tool-approval-gate.test.ts +++ b/packages/appkit/src/plugins/agents/tests/tool-approval-gate.test.ts @@ -1,6 +1,68 @@ +import { type Span, SpanStatusCode, trace } from "@opentelemetry/api"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import * as approvalGateModule from "../tool-approval-gate"; import { ToolApprovalGate } from "../tool-approval-gate"; +async function captureSpans( + operation: () => Promise, +): Promise<{ spans: ReadableSpan[]; error?: unknown }> { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + provider.getTracer(name, version), + ); + let error: unknown; + let spans: ReadableSpan[] = []; + try { + await operation(); + } catch (caught) { + error = caught; + } finally { + // The SDK's flush/shutdown path uses timers internally. Approval tests use + // fake timers for deterministic wait-state transitions, so restore real + // timers only after the operation (and its measured duration) completes. + vi.useRealTimers(); + await provider.forceFlush(); + spans = exporter.getFinishedSpans(); + getTracerSpy.mockRestore(); + await provider.shutdown(); + } + return { spans, ...(error !== undefined ? { error } : {}) }; +} + +function tracedWait( + gate: ToolApprovalGate, + input: { + approvalId: string; + streamId: string; + userId: string; + timeoutMs: number; + toolName: string; + effect?: "read" | "write" | "update" | "destructive"; + args: unknown; + }, +): Promise<"approve" | "deny"> { + return gate.wait(input as unknown as Parameters[0]); +} + +function approvalSpan(spans: ReadableSpan[]): ReadableSpan { + const span = spans.find( + (candidate) => candidate.attributes["mlflow.spanType"] === "CHAIN", + ); + expect(span, "missing CHAIN span").toBeDefined(); + return span as ReadableSpan; +} + describe("ToolApprovalGate", () => { let gate: ToolApprovalGate; @@ -153,4 +215,151 @@ describe("ToolApprovalGate", () => { }); expect(late).toEqual({ ok: false, reason: "unknown" }); }); + + describe("semantic CHAIN spans", () => { + beforeEach(() => { + // The OpenTelemetry SDK's in-memory processor uses timers internally; + // approval state remains deterministic with a 1ms real timeout here. + vi.useRealTimers(); + }); + + test.each([ + ["approve", "approved"], + ["deny", "denied"], + ] as const)( + "records an explicit %s decision as %s", + async (decision, expectedState) => { + const observed = await captureSpans(async () => { + const waiter = tracedWait(gate, { + approvalId: `explicit-${decision}`, + streamId: "stream-explicit", + userId: "alice", + timeoutMs: 60_000, + toolName: "users.update", + effect: "update", + args: { password: "do-not-log", userId: 7 }, + }); + gate.submit({ + approvalId: `explicit-${decision}`, + userId: "alice", + decision, + }); + await expect(waiter).resolves.toBe(decision); + }); + + expect(observed.error).toBeUndefined(); + const span = approvalSpan(observed.spans); + expect(span.attributes).toMatchObject({ + "appkit.approval.id": `explicit-${decision}`, + "appkit.approval.tool_name": "users.update", + "appkit.approval.effect": "update", + "appkit.approval.decision": decision, + "appkit.approval.state": expectedState, + "appkit.approval.duration_ms": expect.any(Number), + }); + expect( + JSON.parse(String(span.attributes["mlflow.spanInputs"])), + ).toEqual({ + password: "[REDACTED]", + userId: 7, + }); + expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toBe( + decision, + ); + expect(span.status.code).toBe(SpanStatusCode.OK); + }, + ); + + test("records automatic denial as timed_out", async () => { + const observed = await captureSpans(async () => { + const waiter = tracedWait(gate, { + approvalId: "timeout-1", + streamId: "stream-timeout", + userId: "alice", + timeoutMs: 1, + toolName: "users.delete", + effect: "destructive", + args: { userId: 8 }, + }); + await expect(waiter).resolves.toBe("deny"); + }); + + expect(observed.error).toBeUndefined(); + expect(approvalSpan(observed.spans).attributes).toMatchObject({ + "appkit.approval.decision": "deny", + "appkit.approval.state": "timed_out", + "appkit.approval.duration_ms": expect.any(Number), + }); + }); + + test("records stream abort as cancelled", async () => { + const observed = await captureSpans(async () => { + const waiter = tracedWait(gate, { + approvalId: "cancel-1", + streamId: "stream-cancel", + userId: "alice", + timeoutMs: 60_000, + toolName: "users.delete", + effect: "destructive", + args: { userId: 9 }, + }); + gate.abortStream("stream-cancel"); + await expect(waiter).resolves.toBe("deny"); + }); + + expect(observed.error).toBeUndefined(); + expect(approvalSpan(observed.spans).attributes).toMatchObject({ + "appkit.approval.decision": "deny", + "appkit.approval.state": "cancelled", + "appkit.approval.duration_ms": expect.any(Number), + }); + }); + + test("traceApprovalWait records a sanitized failed state", async () => { + type TraceApprovalWait = ( + input: { + approvalId: string; + toolName: string; + effect?: "read" | "write" | "update" | "destructive"; + args: unknown; + }, + operation: (span: Span) => Promise, + ) => Promise; + const traceApprovalWait = ( + approvalGateModule as unknown as { + traceApprovalWait?: TraceApprovalWait; + } + ).traceApprovalWait; + + const observed = await captureSpans(() => + traceApprovalWait + ? traceApprovalWait( + { + approvalId: "failed-1", + toolName: "users.delete", + effect: "destructive", + args: {}, + }, + async () => { + throw new Error("approval backend token secret-token"); + }, + ) + : Promise.reject(new Error("traceApprovalWait is not implemented")), + ); + + expect(observed.error).toBeInstanceOf(Error); + const span = approvalSpan(observed.spans); + expect(span.attributes).toMatchObject({ + "appkit.approval.id": "failed-1", + "appkit.approval.decision": "error", + "appkit.approval.state": "failed", + "appkit.error": '{"error":"[REDACTED]"}', + "appkit.approval.duration_ms": expect.any(Number), + }); + expect(span.status.code).toBe(SpanStatusCode.ERROR); + expect( + JSON.stringify({ attributes: span.attributes, events: span.events }), + ).not.toContain("secret-token"); + }); + }); }); diff --git a/packages/appkit/src/plugins/agents/thread-store.ts b/packages/appkit/src/plugins/agents/thread-store.ts index 7c4622cd3..951109776 100644 --- a/packages/appkit/src/plugins/agents/thread-store.ts +++ b/packages/appkit/src/plugins/agents/thread-store.ts @@ -1,5 +1,78 @@ import { randomUUID } from "node:crypto"; +import { type Span, SpanStatusCode, trace } from "@opentelemetry/api"; import type { Message, Thread, ThreadStore } from "shared"; +import { captureTraceValue } from "../../telemetry/agent-tracing"; + +const tracer = () => trace.getTracer("@databricks/appkit-agent-tracing"); + +type MemoryOperation = "create" | "get" | "list" | "addMessage" | "delete"; + +async function traceMemoryOperation( + operationName: MemoryOperation, + inputs: unknown, + operation: (span: Span) => Promise, +): Promise { + return tracer().startActiveSpan( + `thread.${operationName}`, + { + attributes: { + "mlflow.spanType": "MEMORY", + "appkit.memory.operation": operationName, + }, + }, + async (span) => { + const startedAt = Date.now(); + setCapturedAttribute(span, "mlflow.spanInputs", inputs); + try { + const result = await operation(span); + setCapturedAttribute(span, "mlflow.spanOutputs", result); + span.setStatus({ code: SpanStatusCode.OK }); + return result; + } catch (error) { + span.setAttribute("appkit.memory.state", "failed"); + recordSafeFailure(span, error, "Thread store operation failed"); + throw error; + } finally { + span.setAttribute( + "appkit.memory.duration_ms", + Math.max(0, Date.now() - startedAt), + ); + span.end(); + } + }, + ); +} + +function setCapturedAttribute(span: Span, key: string, value: unknown): void { + const captured = captureTraceValue(value); + span.setAttribute(key, captured.value); + span.setAttribute(`${key}.original_bytes`, captured.originalBytes); + span.setAttribute(`${key}.sha256`, captured.sha256); + span.setAttribute(`${key}.truncated`, captured.truncated); +} + +function recordSafeFailure( + span: Span, + error: unknown, + publicMessage: string, +): void { + const failure = captureTraceValue( + { + error: + error instanceof Error + ? error.message + : String(error ?? "Unknown error"), + }, + { redactKeys: ["error"] }, + ); + span.setAttribute("appkit.error", failure.value); + span.setAttribute("mlflow.spanOutputs", failure.value); + span.setAttribute("mlflow.spanOutputs.original_bytes", failure.originalBytes); + span.setAttribute("mlflow.spanOutputs.sha256", failure.sha256); + span.setAttribute("mlflow.spanOutputs.truncated", failure.truncated); + span.recordException({ name: "Error", message: publicMessage }); + span.setStatus({ code: SpanStatusCode.ERROR, message: publicMessage }); +} /** * In-memory thread store backed by a nested Map. @@ -64,3 +137,65 @@ export class InMemoryThreadStore implements ThreadStore { return map; } } + +/** + * Semantic tracing decorator for any {@link ThreadStore} implementation. + * + * The wrapped store remains the source of truth; this class only adds MEMORY + * descendants using AppKit's active OpenTelemetry provider and central value + * capture policy. + */ +export class TracedThreadStore implements ThreadStore { + constructor(private readonly backing: ThreadStore) {} + + create(userId: string): Promise { + return traceMemoryOperation("create", { userId }, async (span) => { + const thread = await this.backing.create(userId); + span.setAttribute("appkit.memory.state", "created"); + return thread; + }); + } + + get(threadId: string, userId: string): Promise { + return traceMemoryOperation("get", { threadId, userId }, async (span) => { + const thread = await this.backing.get(threadId, userId); + span.setAttribute("appkit.memory.state", thread ? "hit" : "miss"); + return thread; + }); + } + + list(userId: string): Promise { + return traceMemoryOperation("list", { userId }, async (span) => { + const threads = await this.backing.list(userId); + span.setAttribute("appkit.memory.state", "completed"); + return threads; + }); + } + + addMessage( + threadId: string, + userId: string, + message: Message, + ): Promise { + return traceMemoryOperation( + "addMessage", + { message, threadId, userId }, + async (span) => { + await this.backing.addMessage(threadId, userId, message); + span.setAttribute("appkit.memory.state", "completed"); + }, + ); + } + + delete(threadId: string, userId: string): Promise { + return traceMemoryOperation( + "delete", + { threadId, userId }, + async (span) => { + const deleted = await this.backing.delete(threadId, userId); + span.setAttribute("appkit.memory.state", deleted ? "deleted" : "miss"); + return deleted; + }, + ); + } +} diff --git a/packages/appkit/src/plugins/agents/tool-approval-gate.ts b/packages/appkit/src/plugins/agents/tool-approval-gate.ts index 4aeb92925..81ec07f8b 100644 --- a/packages/appkit/src/plugins/agents/tool-approval-gate.ts +++ b/packages/appkit/src/plugins/agents/tool-approval-gate.ts @@ -1,3 +1,90 @@ +import { type Span, SpanStatusCode, trace } from "@opentelemetry/api"; +import type { ToolEffect } from "shared"; +import { captureTraceValue } from "../../telemetry/agent-tracing"; + +const tracer = () => trace.getTracer("@databricks/appkit-agent-tracing"); + +type ApprovalState = + | "approved" + | "denied" + | "timed_out" + | "cancelled" + | "failed"; + +export async function traceApprovalWait( + input: { + approvalId: string; + toolName: string; + effect?: ToolEffect; + args: unknown; + }, + operation: (span: Span) => Promise, +): Promise { + return tracer().startActiveSpan( + `${input.toolName} approval`, + { + attributes: { + "mlflow.spanType": "CHAIN", + "appkit.approval.id": input.approvalId, + "appkit.approval.tool_name": input.toolName, + ...(input.effect ? { "appkit.approval.effect": input.effect } : {}), + }, + }, + async (span) => { + const startedAt = Date.now(); + setCapturedAttribute(span, "mlflow.spanInputs", input.args); + try { + const result = await operation(span); + setCapturedAttribute(span, "mlflow.spanOutputs", result); + span.setStatus({ code: SpanStatusCode.OK }); + return result; + } catch (error) { + span.setAttribute("appkit.approval.decision", "error"); + span.setAttribute("appkit.approval.state", "failed"); + recordSafeFailure(span, error, "Approval wait failed"); + throw error; + } finally { + span.setAttribute( + "appkit.approval.duration_ms", + Math.max(0, Date.now() - startedAt), + ); + span.end(); + } + }, + ); +} + +function setCapturedAttribute(span: Span, key: string, value: unknown): void { + const captured = captureTraceValue(value); + span.setAttribute(key, captured.value); + span.setAttribute(`${key}.original_bytes`, captured.originalBytes); + span.setAttribute(`${key}.sha256`, captured.sha256); + span.setAttribute(`${key}.truncated`, captured.truncated); +} + +function recordSafeFailure( + span: Span, + error: unknown, + publicMessage: string, +): void { + const failure = captureTraceValue( + { + error: + error instanceof Error + ? error.message + : String(error ?? "Unknown error"), + }, + { redactKeys: ["error"] }, + ); + span.setAttribute("appkit.error", failure.value); + span.setAttribute("mlflow.spanOutputs", failure.value); + span.setAttribute("mlflow.spanOutputs.original_bytes", failure.originalBytes); + span.setAttribute("mlflow.spanOutputs.sha256", failure.sha256); + span.setAttribute("mlflow.spanOutputs.truncated", failure.truncated); + span.recordException({ name: "Error", message: publicMessage }); + span.setStatus({ code: SpanStatusCode.ERROR, message: publicMessage }); +} + /** * Server-side state for the human-in-the-loop approval gate on mutating * agent tool calls ā€” tools annotated with `effect: "write" | "update" | @@ -29,7 +116,7 @@ type ApprovalDecision = "approve" | "deny"; interface Pending { - resolve: (decision: ApprovalDecision) => void; + settle: (decision: ApprovalDecision, state: ApprovalState) => void; userId: string; streamId: string; timeout: ReturnType; @@ -52,21 +139,40 @@ export class ToolApprovalGate { streamId: string; userId: string; timeoutMs: number; + toolName?: string; + effect?: ToolEffect; + args?: unknown; }): Promise { - const { approvalId, streamId, userId, timeoutMs } = args; - return new Promise((resolve) => { - const timeout = setTimeout(() => { - if (this.pending.delete(approvalId)) { - resolve("deny"); - } - }, timeoutMs); - this.pending.set(approvalId, { - resolve, - userId, - streamId, - timeout, - }); - }); + const { + approvalId, + streamId, + userId, + timeoutMs, + toolName = "unknown", + effect, + } = args; + return traceApprovalWait( + { approvalId, toolName, effect, args: args.args }, + (span) => + new Promise((resolve) => { + const settle = (decision: ApprovalDecision, state: ApprovalState) => { + span.setAttribute("appkit.approval.decision", decision); + span.setAttribute("appkit.approval.state", state); + resolve(decision); + }; + const timeout = setTimeout(() => { + if (this.pending.delete(approvalId)) { + settle("deny", "timed_out"); + } + }, timeoutMs); + this.pending.set(approvalId, { + settle, + userId, + streamId, + timeout, + }); + }), + ); } /** @@ -88,7 +194,7 @@ export class ToolApprovalGate { if (p.userId !== userId) return { ok: false, reason: "forbidden" }; clearTimeout(p.timeout); this.pending.delete(approvalId); - p.resolve(decision); + p.settle(decision, decision === "approve" ? "approved" : "denied"); return { ok: true }; } @@ -102,7 +208,7 @@ export class ToolApprovalGate { if (p.streamId === streamId) { clearTimeout(p.timeout); this.pending.delete(id); - p.resolve("deny"); + p.settle("deny", "cancelled"); } } } @@ -112,7 +218,7 @@ export class ToolApprovalGate { for (const [id, p] of this.pending) { clearTimeout(p.timeout); this.pending.delete(id); - p.resolve("deny"); + p.settle("deny", "cancelled"); } } diff --git a/packages/appkit/src/telemetry/tests/plugin-telemetry.test.ts b/packages/appkit/src/telemetry/tests/plugin-telemetry.test.ts index 42c310e27..291ac1417 100644 --- a/packages/appkit/src/telemetry/tests/plugin-telemetry.test.ts +++ b/packages/appkit/src/telemetry/tests/plugin-telemetry.test.ts @@ -1,4 +1,5 @@ import { beforeEach, describe, expect, test, vi } from "vitest"; +import { PluginContext } from "../../core/plugin-context"; import { NOOP_LOGGER, NOOP_METER, NOOP_TRACER } from "../noop"; import type { TelemetryManager } from "../telemetry-manager"; import { TelemetryProvider } from "../telemetry-provider"; @@ -190,6 +191,66 @@ describe("TelemetryProvider", () => { expect(result).toBe("result"); }); + test("should tag the plugin execution child with its semantic tool identity", async () => { + const ctx = new PluginContext(); + const span = { + setAttribute: vi.fn(), + setStatus: vi.fn(), + recordException: vi.fn(), + end: vi.fn(), + }; + mockTracer.startActiveSpan.mockImplementation( + async ( + _name: string, + optionsOrFn: unknown, + maybeFn?: (activeSpan: typeof span) => Promise, + ) => { + const fn = + maybeFn ?? + (optionsOrFn as (activeSpan: typeof span) => Promise); + return fn(span); + }, + ); + const provider = { + getAgentTools: vi.fn().mockReturnValue([]), + executeAgentTool: vi.fn().mockResolvedValue("rows"), + asUser: vi.fn().mockReturnThis(), + }; + ctx.registerToolProvider("analytics", provider as any); + + type ExecuteToolWithTraceIdentity = ( + req: unknown, + pluginName: string, + toolName: string, + args: unknown, + signal: AbortSignal | undefined, + timeoutMs: number, + traceIdentity: { name: string; source: string }, + ) => Promise; + await (ctx.executeTool as unknown as ExecuteToolWithTraceIdentity)( + { headers: {} }, + "analytics", + "query", + { sql: "SELECT 1" }, + undefined, + 90_000, + { name: "renamed-query", source: "toolkit" }, + ); + + expect(mockTracer.startActiveSpan).toHaveBeenCalledWith( + "executeTool:analytics.query", + expect.any(Function), + ); + expect(span.setAttribute).toHaveBeenCalledWith( + "appkit.tool.name", + "renamed-query", + ); + expect(span.setAttribute).toHaveBeenCalledWith( + "appkit.tool.source", + "toolkit", + ); + }); + test("should delegate registerInstrumentations to global manager", () => { const telemetry = new TelemetryProvider("test-plugin", mockManager); const instrumentations = [{ name: "test-instrumentation" }] as any; From 8767e077a11daf378e6fa2157aab968630aad03f Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 20:00:41 -0700 Subject: [PATCH 13/31] fix: complete AppKit tracing semantic attributes Signed-off-by: Adam Gurary --- packages/appkit/src/plugins/agents/agents.ts | 2 ++ .../agents/tests/dispatch-tool-call.test.ts | 13 +++++++++- .../plugins/agents/tests/thread-store.test.ts | 14 ++++++----- .../agents/tests/tool-approval-gate.test.ts | 2 +- .../appkit/src/plugins/agents/thread-store.ts | 24 +++++++++++++------ .../src/plugins/agents/tool-approval-gate.ts | 2 +- 6 files changed, 41 insertions(+), 16 deletions(-) diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index a46ee3517..5d734c130 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -93,6 +93,8 @@ export async function traceToolCall( { attributes: { "mlflow.spanType": "TOOL", + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": input.name, "appkit.tool.name": input.name, "appkit.tool.source": input.source, ...(input.effect ? { "appkit.tool.effect": input.effect } : {}), diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index a84ffc2fc..e2a1117ae 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -289,6 +289,17 @@ describe("dispatchToolCall ā€” semantic TOOL spans", () => { (span) => span.attributes["mlflow.spanType"] === "TOOL", ); expect(tools).toHaveLength(4); + expect( + tools.map((span) => [ + span.attributes["gen_ai.operation.name"], + span.attributes["gen_ai.tool.name"], + ]), + ).toEqual([ + ["execute_tool", "inline"], + ["execute_tool", "analytics.query"], + ["execute_tool", "remote_lookup"], + ["execute_tool", "agent-researcher"], + ]); expect( tools.map((span) => [ span.attributes["appkit.tool.name"], @@ -538,7 +549,7 @@ describe("dispatchToolCall ā€” semantic approval descendants", () => { expect(approval.attributes).toMatchObject({ "appkit.approval.decision": "approve", "appkit.approval.state": "approved", - "appkit.approval.tool_name": "delete_user", + "appkit.tool.name": "delete_user", "appkit.approval.duration_ms": expect.any(Number), }); expect( diff --git a/packages/appkit/src/plugins/agents/tests/thread-store.test.ts b/packages/appkit/src/plugins/agents/tests/thread-store.test.ts index c64184bf4..b64b2c499 100644 --- a/packages/appkit/src/plugins/agents/tests/thread-store.test.ts +++ b/packages/appkit/src/plugins/agents/tests/thread-store.test.ts @@ -222,14 +222,16 @@ describe("TracedThreadStore", () => { spans.map((span) => [ span.attributes["appkit.memory.operation"], span.attributes["appkit.memory.state"], + span.attributes["appkit.memory.store"], + span.attributes["appkit.memory.key"], ]), ).toEqual([ - ["create", "created"], - ["get", "hit"], - ["get", "miss"], - ["list", "completed"], - ["addMessage", "completed"], - ["delete", "deleted"], + ["create", "created", "thread", "user-7"], + ["get", "hit", "thread", created?.id], + ["get", "miss", "thread", "thread-missing"], + ["list", "completed", "thread", "user-7"], + ["addMessage", "completed", "thread", created?.id], + ["delete", "deleted", "thread", created?.id], ]); expect( spans.every( diff --git a/packages/appkit/src/plugins/agents/tests/tool-approval-gate.test.ts b/packages/appkit/src/plugins/agents/tests/tool-approval-gate.test.ts index e32e81145..0450e221c 100644 --- a/packages/appkit/src/plugins/agents/tests/tool-approval-gate.test.ts +++ b/packages/appkit/src/plugins/agents/tests/tool-approval-gate.test.ts @@ -251,7 +251,7 @@ describe("ToolApprovalGate", () => { const span = approvalSpan(observed.spans); expect(span.attributes).toMatchObject({ "appkit.approval.id": `explicit-${decision}`, - "appkit.approval.tool_name": "users.update", + "appkit.tool.name": "users.update", "appkit.approval.effect": "update", "appkit.approval.decision": decision, "appkit.approval.state": expectedState, diff --git a/packages/appkit/src/plugins/agents/thread-store.ts b/packages/appkit/src/plugins/agents/thread-store.ts index 951109776..d8c6a905b 100644 --- a/packages/appkit/src/plugins/agents/thread-store.ts +++ b/packages/appkit/src/plugins/agents/thread-store.ts @@ -9,6 +9,7 @@ type MemoryOperation = "create" | "get" | "list" | "addMessage" | "delete"; async function traceMemoryOperation( operationName: MemoryOperation, + key: string, inputs: unknown, operation: (span: Span) => Promise, ): Promise { @@ -18,6 +19,8 @@ async function traceMemoryOperation( attributes: { "mlflow.spanType": "MEMORY", "appkit.memory.operation": operationName, + "appkit.memory.store": "thread", + "appkit.memory.key": key, }, }, async (span) => { @@ -149,7 +152,7 @@ export class TracedThreadStore implements ThreadStore { constructor(private readonly backing: ThreadStore) {} create(userId: string): Promise { - return traceMemoryOperation("create", { userId }, async (span) => { + return traceMemoryOperation("create", userId, { userId }, async (span) => { const thread = await this.backing.create(userId); span.setAttribute("appkit.memory.state", "created"); return thread; @@ -157,15 +160,20 @@ export class TracedThreadStore implements ThreadStore { } get(threadId: string, userId: string): Promise { - return traceMemoryOperation("get", { threadId, userId }, async (span) => { - const thread = await this.backing.get(threadId, userId); - span.setAttribute("appkit.memory.state", thread ? "hit" : "miss"); - return thread; - }); + return traceMemoryOperation( + "get", + threadId, + { threadId, userId }, + async (span) => { + const thread = await this.backing.get(threadId, userId); + span.setAttribute("appkit.memory.state", thread ? "hit" : "miss"); + return thread; + }, + ); } list(userId: string): Promise { - return traceMemoryOperation("list", { userId }, async (span) => { + return traceMemoryOperation("list", userId, { userId }, async (span) => { const threads = await this.backing.list(userId); span.setAttribute("appkit.memory.state", "completed"); return threads; @@ -179,6 +187,7 @@ export class TracedThreadStore implements ThreadStore { ): Promise { return traceMemoryOperation( "addMessage", + threadId, { message, threadId, userId }, async (span) => { await this.backing.addMessage(threadId, userId, message); @@ -190,6 +199,7 @@ export class TracedThreadStore implements ThreadStore { delete(threadId: string, userId: string): Promise { return traceMemoryOperation( "delete", + threadId, { threadId, userId }, async (span) => { const deleted = await this.backing.delete(threadId, userId); diff --git a/packages/appkit/src/plugins/agents/tool-approval-gate.ts b/packages/appkit/src/plugins/agents/tool-approval-gate.ts index 81ec07f8b..6fdf339e0 100644 --- a/packages/appkit/src/plugins/agents/tool-approval-gate.ts +++ b/packages/appkit/src/plugins/agents/tool-approval-gate.ts @@ -26,7 +26,7 @@ export async function traceApprovalWait( attributes: { "mlflow.spanType": "CHAIN", "appkit.approval.id": input.approvalId, - "appkit.approval.tool_name": input.toolName, + "appkit.tool.name": input.toolName, ...(input.effect ? { "appkit.approval.effect": input.effect } : {}), }, }, From 29b377c57592d73dd640f629c2e669c013ad42dc Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 20:25:04 -0700 Subject: [PATCH 14/31] feat: trace retrieval and local sub-agents Signed-off-by: Adam Gurary --- .../appkit/src/connectors/ai-search/client.ts | 169 ++++++++++- .../connectors/ai-search/tests/client.test.ts | 286 ++++++++++++++++++ packages/appkit/src/core/agent/run-agent.ts | 151 +++++---- .../src/core/agent/tests/run-agent.test.ts | 197 +++++++++++- .../appkit/src/core/agent/trace-tool-call.ts | 80 +++++ packages/appkit/src/plugins/agents/agents.ts | 219 ++++++-------- .../agents/tests/dispatch-tool-call.test.ts | 213 ++++++++++++- .../plugins/agents/tests/dos-limits.test.ts | 17 +- .../plugins/ai-search/tests/ai-search.test.ts | 1 + .../src/telemetry/agent-tracing/tracer.ts | 3 + .../src/telemetry/agent-tracing/types.ts | 2 + 11 files changed, 1144 insertions(+), 194 deletions(-) create mode 100644 packages/appkit/src/connectors/ai-search/tests/client.test.ts create mode 100644 packages/appkit/src/core/agent/trace-tool-call.ts diff --git a/packages/appkit/src/connectors/ai-search/client.ts b/packages/appkit/src/connectors/ai-search/client.ts index f4ed44b15..f8793ebc4 100644 --- a/packages/appkit/src/connectors/ai-search/client.ts +++ b/packages/appkit/src/connectors/ai-search/client.ts @@ -6,6 +6,7 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; +import { captureTraceValue } from "../../telemetry/agent-tracing"; import type { WorkspaceClient } from "../../workspace-client"; import { contextFromAbortSignal } from "../context"; import type { @@ -19,6 +20,21 @@ import type { const logger = createLogger("connectors:ai-search"); +const RETRIEVER_SOURCE = "databricks-ai-search"; +const DOCUMENT_ID_COLUMNS = new Set(["id", "doc_id", "document_id"]); + +interface RetrieverDocument { + content: Record; + documentId: string; + score: number | null; +} + +interface RetrieverOutputs { + documents: RetrieverDocument[]; + nextPageToken: string | null; + resultCount: number; +} + export class AiSearchConnector { private readonly telemetry: TelemetryProvider; @@ -70,6 +86,10 @@ export class AiSearchConnector { { kind: SpanKind.CLIENT, attributes: { + "mlflow.spanType": "RETRIEVER", + "appkit.retriever.source": RETRIEVER_SOURCE, + "appkit.retriever.index_name": params.indexName, + "appkit.retriever.query_type": params.queryType, "db.system": "databricks", "vs.index_name": params.indexName, "vs.query_type": params.queryType, @@ -82,6 +102,16 @@ export class AiSearchConnector { }, async (span: Span) => { const startTime = Date.now(); + setCapturedAttribute(span, "mlflow.spanInputs", { + columns: params.columns, + filters: params.filters ?? {}, + indexName: params.indexName, + numResults: params.numResults, + queryText: params.queryText ?? null, + queryType: params.queryType, + queryVector: summarizeVector(params.queryVector), + reranker: params.reranker ?? null, + }); try { const response = (await workspaceClient.apiClient.request( { @@ -96,12 +126,13 @@ export class AiSearchConnector { )) as VsRawResponse; const duration = Date.now() - startTime; + const outputs = retrieverOutputs(response, params.columns); + setRetrieverOutputs(span, outputs); span.setAttribute("vs.result_count", response.result.row_count); span.setAttribute( "vs.query_time_ms", response.debug_info?.response_time ?? 0, ); - span.setAttribute("vs.duration_ms", duration); span.setStatus({ code: SpanStatusCode.OK }); logger.event()?.setContext("ai-search", { @@ -114,12 +145,13 @@ export class AiSearchConnector { return response; } catch (error) { - span.recordException(error as Error); - span.setStatus({ - code: SpanStatusCode.ERROR, - message: error instanceof Error ? error.message : String(error), - }); + recordRetrieverFailure(span, error); throw error; + } finally { + const duration = Math.max(0, Date.now() - startTime); + span.setAttribute("appkit.retriever.latency_ms", duration); + span.setAttribute("vs.duration_ms", duration); + span.end(); } }, { name: "ai-search", includePrefix: true }, @@ -146,12 +178,23 @@ export class AiSearchConnector { { kind: SpanKind.CLIENT, attributes: { + "mlflow.spanType": "RETRIEVER", + "appkit.retriever.source": RETRIEVER_SOURCE, + "appkit.retriever.index_name": params.indexName, + "appkit.retriever.query_type": "next_page", "db.system": "databricks", "vs.index_name": params.indexName, "vs.endpoint_name": params.endpointName, }, }, async (span: Span) => { + const startTime = Date.now(); + setCapturedAttribute(span, "mlflow.spanInputs", { + endpointName: params.endpointName, + indexName: params.indexName, + pageToken: params.pageToken, + queryType: "next_page", + }); try { const response = (await workspaceClient.apiClient.request( { @@ -168,16 +211,23 @@ export class AiSearchConnector { contextFromAbortSignal(signal), )) as VsRawResponse; + const outputs = retrieverOutputs(response); + setRetrieverOutputs(span, outputs); span.setAttribute("vs.result_count", response.result.row_count); + span.setAttribute( + "vs.query_time_ms", + response.debug_info?.response_time ?? 0, + ); span.setStatus({ code: SpanStatusCode.OK }); return response; } catch (error) { - span.recordException(error as Error); - span.setStatus({ - code: SpanStatusCode.ERROR, - message: error instanceof Error ? error.message : String(error), - }); + recordRetrieverFailure(span, error); throw error; + } finally { + const duration = Math.max(0, Date.now() - startTime); + span.setAttribute("appkit.retriever.latency_ms", duration); + span.setAttribute("vs.duration_ms", duration); + span.end(); } }, { name: "ai-search", includePrefix: true }, @@ -227,3 +277,100 @@ export class AiSearchConnector { return (table.columns ?? []).map((c) => c.name); } } + +function summarizeVector( + vector: number[] | undefined, +): { dimensions: number; sha256: string } | null { + if (!vector) return null; + const captured = captureTraceValue(vector); + return { dimensions: vector.length, sha256: captured.sha256 }; +} + +function retrieverOutputs( + response: VsRawResponse, + configuredColumns?: readonly string[], +): RetrieverOutputs { + const columnNames = response.manifest.columns.map((column) => column.name); + const configured = configuredColumns + ? new Set(configuredColumns.map((column) => column.toLowerCase())) + : undefined; + const documentIdIndex = columnNames.findIndex((column) => { + const normalized = column.toLowerCase(); + return ( + DOCUMENT_ID_COLUMNS.has(normalized) && + (configured === undefined || configured.has(normalized)) + ); + }); + const scoreIndex = columnNames.findIndex( + (column) => column.toLowerCase() === "score", + ); + const documents = response.result.data_array.map((row) => { + const content = Object.fromEntries( + columnNames.map((column, index) => [column, row[index] ?? null]), + ); + const capturedRow = captureTraceValue(content); + const returnedId = documentIdIndex >= 0 ? row[documentIdIndex] : undefined; + const score = scoreIndex >= 0 ? row[scoreIndex] : undefined; + return { + content, + documentId: + returnedId === undefined || returnedId === null || returnedId === "" + ? capturedRow.sha256 + : String(returnedId), + score: typeof score === "number" ? score : null, + }; + }); + return { + documents, + nextPageToken: response.next_page_token ?? null, + resultCount: response.result.row_count, + }; +} + +function setRetrieverOutputs(span: Span, outputs: RetrieverOutputs): void { + setCapturedAttribute(span, "mlflow.spanOutputs", outputs); + span.setAttribute("appkit.retriever.result_count", outputs.resultCount); + span.setAttribute( + "appkit.retriever.document_ids", + outputs.documents.map((document) => document.documentId), + ); + const scores = outputs.documents.flatMap((document) => + document.score === null ? [] : [document.score], + ); + if (scores.length > 0) { + span.setAttribute("appkit.retriever.scores", scores); + } +} + +function setCapturedAttribute(span: Span, key: string, value: unknown): void { + const captured = captureTraceValue(value); + span.setAttribute(key, captured.value); + span.setAttribute(`${key}.original_bytes`, captured.originalBytes); + span.setAttribute(`${key}.sha256`, captured.sha256); + span.setAttribute(`${key}.truncated`, captured.truncated); +} + +function recordRetrieverFailure(span: Span, error: unknown): void { + const failure = captureTraceValue( + { + error: + error instanceof Error + ? error.message + : String(error ?? "Unknown error"), + }, + { redactKeys: ["error"] }, + ); + span.setAttribute("appkit.error", failure.value); + span.setAttribute("mlflow.spanOutputs", failure.value); + span.setAttribute("mlflow.spanOutputs.original_bytes", failure.originalBytes); + span.setAttribute("mlflow.spanOutputs.sha256", failure.sha256); + span.setAttribute("mlflow.spanOutputs.truncated", failure.truncated); + span.recordException({ + name: "Error", + message: "Retriever operation failed", + }); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: "Retriever operation failed", + }); +} diff --git a/packages/appkit/src/connectors/ai-search/tests/client.test.ts b/packages/appkit/src/connectors/ai-search/tests/client.test.ts new file mode 100644 index 000000000..f64b5d93a --- /dev/null +++ b/packages/appkit/src/connectors/ai-search/tests/client.test.ts @@ -0,0 +1,286 @@ +import { SpanStatusCode, trace } from "@opentelemetry/api"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import { describe, expect, test, vi } from "vitest"; +import type { WorkspaceClient } from "../../../workspace-client"; +import { AiSearchConnector } from "../client"; +import type { VsRawResponse } from "../types"; + +async function captureSpans( + operation: () => Promise, +): Promise<{ spans: ReadableSpan[]; error?: unknown }> { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + provider.getTracer(name, version), + ); + let error: unknown; + let spans: ReadableSpan[] = []; + try { + await operation(); + } catch (caught) { + error = caught; + } finally { + await provider.forceFlush(); + spans = exporter.getFinishedSpans(); + getTracerSpy.mockRestore(); + await provider.shutdown(); + } + return { spans, ...(error !== undefined ? { error } : {}) }; +} + +function workspaceClient(response: VsRawResponse | Error): WorkspaceClient { + return { + apiClient: { + request: + response instanceof Error + ? vi.fn().mockRejectedValue(response) + : vi.fn().mockResolvedValue(response), + }, + } as unknown as WorkspaceClient; +} + +function retrieverSpan(spans: ReadableSpan[]): ReadableSpan { + const span = spans.find( + (candidate) => candidate.attributes["mlflow.spanType"] === "RETRIEVER", + ); + expect(span, "missing RETRIEVER span").toBeDefined(); + return span as ReadableSpan; +} + +describe("AiSearchConnector semantic retrieval spans", () => { + test("exports complete query inputs, documents, stable IDs, scores, and diagnostics", async () => { + const response: VsRawResponse = { + manifest: { + column_count: 3, + columns: [{ name: "id" }, { name: "text" }, { name: "score" }], + }, + result: { + row_count: 2, + data_array: [ + ["doc-1", "Complete first row", 0.98], + ["doc-2", "Complete second row", 0.87], + ], + }, + next_page_token: "page-2", + debug_info: { response_time: 35 }, + }; + const connector = new AiSearchConnector(); + + const observed = await captureSpans(() => + connector.query(workspaceClient(response), { + columns: ["id", "text", "score"], + filters: { + category: ["observability"], + password: "do-not-export", + }, + indexName: "catalog.schema.docs", + numResults: 2, + queryText: "trace local agents", + queryType: "hybrid", + queryVector: [0.125, -0.5, 1.25], + reranker: { columnsToRerank: ["text"] }, + }), + ); + + expect(observed.error).toBeUndefined(); + const span = retrieverSpan(observed.spans); + expect(span.name).toBe("ai-search.query"); + expect(span.status.code).toBe(SpanStatusCode.OK); + expect(span.attributes).toMatchObject({ + "appkit.retriever.document_ids": ["doc-1", "doc-2"], + "appkit.retriever.index_name": "catalog.schema.docs", + "appkit.retriever.latency_ms": expect.any(Number), + "appkit.retriever.query_type": "hybrid", + "appkit.retriever.result_count": 2, + "appkit.retriever.scores": [0.98, 0.87], + "appkit.retriever.source": "databricks-ai-search", + "db.system": "databricks", + "vs.duration_ms": expect.any(Number), + "vs.has_filters": true, + "vs.has_reranker": true, + "vs.index_name": "catalog.schema.docs", + "vs.num_results": 2, + "vs.query_time_ms": 35, + "vs.query_type": "hybrid", + "vs.result_count": 2, + }); + expect(JSON.parse(String(span.attributes["mlflow.spanInputs"]))).toEqual({ + columns: ["id", "text", "score"], + filters: { + category: ["observability"], + password: "[REDACTED]", + }, + indexName: "catalog.schema.docs", + numResults: 2, + queryText: "trace local agents", + queryType: "hybrid", + queryVector: { + dimensions: 3, + sha256: + "f6b2b238972f104fdfb47a54f079b06eda051734161400b60d305ad49d9b2d31", + }, + reranker: { columnsToRerank: ["text"] }, + }); + expect(span.attributes["mlflow.spanInputs.original_bytes"]).toBe(348); + expect(span.attributes["mlflow.spanInputs.sha256"]).toBe( + "9bed2c6c6453eb7ee53685ae0669929c19fa5f538255c16fb7630d2faa5d4c1d", + ); + expect(span.attributes["mlflow.spanInputs.truncated"]).toBe(false); + expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ + documents: [ + { + content: { + id: "doc-1", + score: 0.98, + text: "Complete first row", + }, + documentId: "doc-1", + score: 0.98, + }, + { + content: { + id: "doc-2", + score: 0.87, + text: "Complete second row", + }, + documentId: "doc-2", + score: 0.87, + }, + ], + nextPageToken: "page-2", + resultCount: 2, + }); + expect(span.attributes["mlflow.spanOutputs.original_bytes"]).toBe(261); + expect(span.attributes["mlflow.spanOutputs.sha256"]).toBe( + "5995c76a74fbcad7875c5d48a303a4d7d8d18317c155c502660938522344cebd", + ); + expect(span.attributes["mlflow.spanOutputs.truncated"]).toBe(false); + expect(span.attributes["mlflow.traceOutputs"]).toBeUndefined(); + }); + + test("exports next-page documents and falls back to the captured-row digest", async () => { + const response: VsRawResponse = { + manifest: { + column_count: 2, + columns: [{ name: "title" }, { name: "score" }], + }, + result: { + row_count: 1, + data_array: [["Fallback row", null]], + }, + next_page_token: null, + }; + const connector = new AiSearchConnector(); + + const observed = await captureSpans(() => + connector.queryNextPage(workspaceClient(response), { + endpointName: "endpoint-a", + indexName: "catalog.schema.docs", + pageToken: "page-2", + }), + ); + + expect(observed.error).toBeUndefined(); + const span = retrieverSpan(observed.spans); + expect(span.name).toBe("ai-search.queryNextPage"); + expect(span.status.code).toBe(SpanStatusCode.OK); + expect(span.attributes).toMatchObject({ + "appkit.retriever.document_ids": [ + "b48a75fb558784bfaaa5305827e9c364bc175160a556c56ff98f33a761980dcb", + ], + "appkit.retriever.index_name": "catalog.schema.docs", + "appkit.retriever.latency_ms": expect.any(Number), + "appkit.retriever.query_type": "next_page", + "appkit.retriever.result_count": 1, + "appkit.retriever.source": "databricks-ai-search", + "vs.endpoint_name": "endpoint-a", + "vs.index_name": "catalog.schema.docs", + "vs.result_count": 1, + }); + expect(JSON.parse(String(span.attributes["mlflow.spanInputs"]))).toEqual({ + endpointName: "endpoint-a", + indexName: "catalog.schema.docs", + pageToken: "page-2", + queryType: "next_page", + }); + expect(span.attributes["mlflow.spanInputs.original_bytes"]).toBe(108); + expect(span.attributes["mlflow.spanInputs.sha256"]).toBe( + "7d0315d01fd5b3d2f6179ff6604466f687cbf9146d3121f5511d7ad5be2e53d8", + ); + expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ + documents: [ + { + content: { score: null, title: "Fallback row" }, + documentId: + "b48a75fb558784bfaaa5305827e9c364bc175160a556c56ff98f33a761980dcb", + score: null, + }, + ], + nextPageToken: null, + resultCount: 1, + }); + expect(span.attributes["mlflow.spanOutputs.original_bytes"]).toBe(195); + expect(span.attributes["mlflow.spanOutputs.sha256"]).toBe( + "16dae38790838bdcc2dc5a62e1b27bbe9f804570cdb467bb111cc454d4c49ffd", + ); + }); + + test("ends a failed retriever with a sanitized exception event", async () => { + const connector = new AiSearchConnector(); + const failure = new Error("vector backend exposed password hunter2"); + + const observed = await captureSpans(() => + connector.query(workspaceClient(failure), { + columns: ["id", "text"], + indexName: "catalog.schema.docs", + numResults: 2, + queryText: "trace failures", + queryType: "ann", + }), + ); + + expect(observed.error).toBe(failure); + const span = retrieverSpan(observed.spans); + expect(span.attributes).toMatchObject({ + "appkit.retriever.latency_ms": expect.any(Number), + "vs.duration_ms": expect.any(Number), + }); + expect(span.status).toEqual({ + code: SpanStatusCode.ERROR, + message: "Retriever operation failed", + }); + expect(span.events).toEqual([ + expect.objectContaining({ + name: "exception", + attributes: expect.objectContaining({ + "exception.message": "Retriever operation failed", + }), + }), + ]); + expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ + error: "[REDACTED]", + }); + expect(JSON.parse(String(span.attributes["mlflow.spanInputs"]))).toEqual({ + columns: ["id", "text"], + filters: {}, + indexName: "catalog.schema.docs", + numResults: 2, + queryText: "trace failures", + queryType: "ann", + queryVector: null, + reranker: null, + }); + expect( + JSON.stringify({ attributes: span.attributes, events: span.events }), + ).not.toContain("hunter2"); + }); +}); diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 76e627568..b5a46e7cc 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -29,6 +29,7 @@ import { isFunctionTool, } from "./tools/function-tool"; import { isHostedTool } from "./tools/hosted-tools"; +import { traceToolCall } from "./trace-tool-call"; import type { AgentDefinition, AgentTool, @@ -69,6 +70,14 @@ export interface RunAgentResult { usage: AgentUsage; } +type ResolvedRunAgentInput = RunAgentInput & { + appName: string; + requestId: string; + sessionId: string; + threadId: string; + userId: string; +}; + /** * Standalone agent execution without `createApp`. Resolves the adapter, binds * inline tools, and drives the adapter's `run()` loop to completion. @@ -109,13 +118,16 @@ export async function runAgent( const providerCache = new Map(); const threadId = input.threadId ?? randomUUID(); const requestId = input.requestId ?? randomUUID(); + const appName = resolveAgentTraceAppName(input.appName); + const sessionId = input.sessionId ?? threadId; + const userId = input.userId ?? "service-principal"; const traced = await runWithAgentTrace( { - appName: resolveAgentTraceAppName(input.appName), + appName, agentName: def.name ?? "agent", route: "runAgent", - sessionId: input.sessionId ?? threadId, - userId: input.userId ?? "service-principal", + sessionId, + userId, requestId, threadId, }, @@ -124,7 +136,14 @@ export async function runAgent( await initStandalonePlugins(input.plugins ?? [], providerCache); return runAgentInternal( def, - { ...input, requestId, threadId }, + { + ...input, + appName, + requestId, + sessionId, + threadId, + userId, + }, providerCache, observer, ); @@ -139,7 +158,7 @@ export async function runAgent( async function runAgentInternal( def: AgentDefinition, - input: RunAgentInput, + input: ResolvedRunAgentInput, providerCache: Map, observer: AgentTraceObserver, ): Promise> { @@ -163,53 +182,81 @@ async function runAgentInternal( const executeTool = async (name: string, args: unknown): Promise => { const entry = toolIndex.get(name); if (!entry) throw new Error(`Unknown tool: ${name}`); - if (entry.kind === "function") { - return entry.tool.execute(args as Record); - } - if (entry.kind === "toolkit") { - return entry.provider.executeAgentTool( - entry.localName, - args as Record, - signal, - ); - } - if (entry.kind === "subagent") { - const subInput: RunAgentInput = { - messages: - typeof args === "object" && - args !== null && - typeof (args as { input?: unknown }).input === "string" - ? (args as { input: string }).input - : JSON.stringify(args), - signal, - plugins: input.plugins, - sessionId: input.sessionId, - userId: input.userId, - requestId: input.requestId, - appName: input.appName, - }; - // Reuse the same `providerCache` so sub-agent plugin tools dispatch - // through the same instances the parent constructed. - const res = await runAgentInternal( - entry.agentDef, - subInput, - providerCache, - observer, - ); - return res.text; - } - if (entry.kind === "hosted-supervisor") { - // Defense-in-depth: should never fire. The placeholder def is - // filtered out of `tools` above, so the model never sees a callable - // schema for hosted-supervisor entries. If we ever reach here, the - // model was somehow handed the def and tried to invoke it directly. - throw new Error( - `runAgent: tool "${name}" is a hosted-supervisor tool, executed server-side by the Databricks AI Gateway. It must not be invoked from the Node process.`, - ); - } - throw new Error( - `runAgent: tool "${name}" is a ${entry.kind} tool. ` + - "Hosted/MCP tools are only usable via createApp({ plugins: [..., agents(...)] }).", + return traceToolCall( + { + name, + source: entry.kind, + effect: entry.def.annotations?.effect, + args, + }, + async () => { + if (entry.kind === "function") { + return entry.tool.execute(args as Record); + } + if (entry.kind === "toolkit") { + return entry.provider.executeAgentTool( + entry.localName, + args as Record, + signal, + ); + } + if (entry.kind === "subagent") { + const childThreadId = randomUUID(); + const subInput: ResolvedRunAgentInput = { + messages: + typeof args === "object" && + args !== null && + typeof (args as { input?: unknown }).input === "string" + ? (args as { input: string }).input + : JSON.stringify(args), + signal, + plugins: input.plugins, + sessionId: input.sessionId, + userId: input.userId, + requestId: input.requestId, + threadId: childThreadId, + appName: input.appName, + }; + const childTrace = await runWithAgentTrace( + { + appName: resolveAgentTraceAppName(input.appName), + agentName: entry.agentDef.name ?? "agent", + route: "runAgent", + sessionId: input.sessionId, + userId: input.userId, + requestId: input.requestId, + threadId: childThreadId, + }, + { messages: subInput.messages }, + (childObserver) => + runAgentInternal( + entry.agentDef, + subInput, + providerCache, + childObserver, + ), + ); + const childResult = { + text: childTrace.value.text, + usage: childTrace.usage, + }; + observer.addChildUsage(childResult.usage); + return childResult.text; + } + if (entry.kind === "hosted-supervisor") { + // Defense-in-depth: should never fire. The placeholder def is + // filtered out of `tools` above, so the model never sees a callable + // schema for hosted-supervisor entries. If we ever reach here, the + // model was somehow handed the def and tried to invoke it directly. + throw new Error( + `runAgent: tool "${name}" is a hosted-supervisor tool, executed server-side by the Databricks AI Gateway. It must not be invoked from the Node process.`, + ); + } + throw new Error( + `runAgent: tool "${name}" is a ${entry.kind} tool. ` + + "Hosted/MCP tools are only usable via createApp({ plugins: [..., agents(...)] }).", + ); + }, ); }; diff --git a/packages/appkit/src/core/agent/tests/run-agent.test.ts b/packages/appkit/src/core/agent/tests/run-agent.test.ts index 6f8b80240..c97d338e2 100644 --- a/packages/appkit/src/core/agent/tests/run-agent.test.ts +++ b/packages/appkit/src/core/agent/tests/run-agent.test.ts @@ -1,7 +1,9 @@ -import { trace } from "@opentelemetry/api"; +import { context, trace } from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; import { BasicTracerProvider, InMemorySpanExporter, + type ReadableSpan, SimpleSpanProcessor, } from "@opentelemetry/sdk-trace-base"; import type { @@ -14,7 +16,7 @@ import type { PluginData, ToolProvider, } from "shared"; -import { describe, expect, test, vi } from "vitest"; +import { afterAll, beforeAll, describe, expect, test, vi } from "vitest"; import { z } from "zod"; import { createAgent } from "../create-agent"; import { runAgent } from "../run-agent"; @@ -31,6 +33,42 @@ function scriptedAdapter(events: AgentEvent[]): AgentAdapter { }; } +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); +}); + +afterAll(() => { + context.disable(); +}); + +async function captureAgentSpans( + operation: () => Promise, +): Promise<{ result: T; spans: ReadableSpan[] }> { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + provider.getTracer(name, version), + ); + let result!: T; + let spans: ReadableSpan[] = []; + try { + result = await operation(); + await provider.forceFlush(); + spans = exporter.getFinishedSpans(); + } finally { + getTracerSpy.mockRestore(); + await provider.shutdown(); + } + return { result, spans }; +} + describe("runAgent", () => { test("drives the adapter and returns aggregated text", async () => { const events: AgentEvent[] = [ @@ -353,6 +391,161 @@ describe("runAgent", () => { expect(result).toBe("child says hi"); }); + test("nests a local agent's model and tool and rolls its usage into the planner once", async () => { + const helperTool = tool({ + name: "helper.tool", + description: "Look up a fact", + schema: z.object({ topic: z.string() }), + execute: async ({ topic }) => `fact:${topic}`, + }); + const helperStartedAt = Date.now(); + const helperAdapter: AgentAdapter = { + async *run(_input, runContext) { + yield { + type: "model_start", + stepId: "helper-step", + model: "helper.model", + provider: "databricks", + input: { messages: [{ role: "user", content: "research" }] }, + startedAt: helperStartedAt, + } satisfies AgentEvent; + const fact = await runContext.executeTool("helper.tool", { + topic: "tracing", + }); + yield { type: "message_delta", content: `helper:${fact}` }; + yield { + type: "model_end", + stepId: "helper-step", + model: "helper.model", + provider: "databricks", + output: { text: `helper:${fact}` }, + usage: { + inputTokens: 6, + outputTokens: 3, + totalTokens: 9, + costAvailable: false, + }, + streamDurationMs: 5, + endedAt: helperStartedAt + 5, + } satisfies AgentEvent; + }, + }; + const plannerStartedAt = Date.now(); + const plannerAdapter: AgentAdapter = { + async *run(_input, runContext) { + yield { + type: "model_start", + stepId: "planner-step", + model: "planner.model", + provider: "databricks", + input: { messages: [{ role: "user", content: "plan" }] }, + startedAt: plannerStartedAt, + } satisfies AgentEvent; + const delegated = await runContext.executeTool("agent-helper", { + input: "research", + }); + yield { type: "message_delta", content: `planner:${delegated}` }; + yield { + type: "model_end", + stepId: "planner-step", + model: "planner.model", + provider: "databricks", + output: { text: `planner:${delegated}` }, + usage: { + inputTokens: 10, + outputTokens: 4, + totalTokens: 14, + costUsd: 0.04, + costAvailable: true, + }, + streamDurationMs: 8, + endedAt: plannerStartedAt + 8, + } satisfies AgentEvent; + }, + }; + const planner = createAgent({ + name: "planner", + instructions: "plan", + model: plannerAdapter, + agents: { + helper: createAgent({ + name: "helper", + instructions: "research", + model: helperAdapter, + tools: { "helper.tool": helperTool }, + }), + }, + }); + + const observed = await captureAgentSpans(() => + runAgent(planner, { + appName: "test-app", + messages: "plan", + requestId: "request-1", + sessionId: "session-1", + threadId: "planner-thread", + userId: "user-1", + }), + ); + + expect(observed.result.usage).toEqual({ + inputTokens: 16, + outputTokens: 7, + totalTokens: 23, + costAvailable: false, + }); + const agentSpans = observed.spans.filter( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + const toolSpans = observed.spans.filter( + (span) => span.attributes["mlflow.spanType"] === "TOOL", + ); + const plannerSpan = agentSpans.find( + (span) => span.attributes["appkit.agent.name"] === "planner", + ); + const helperSpan = agentSpans.find( + (span) => span.attributes["appkit.agent.name"] === "helper", + ); + const delegation = toolSpans.find( + (span) => span.attributes["appkit.tool.name"] === "agent-helper", + ); + const helperToolSpan = toolSpans.find( + (span) => span.attributes["appkit.tool.name"] === "helper.tool", + ); + const helperModel = observed.spans.find( + (span) => span.attributes["mlflow.chat.model"] === "helper.model", + ); + expect(agentSpans).toHaveLength(2); + expect(plannerSpan).toBeDefined(); + expect(delegation?.parentSpanContext?.spanId).toBe( + plannerSpan?.spanContext().spanId, + ); + expect(helperSpan?.parentSpanContext?.spanId).toBe( + delegation?.spanContext().spanId, + ); + expect(helperModel?.parentSpanContext?.spanId).toBe( + helperSpan?.spanContext().spanId, + ); + expect(helperToolSpan?.parentSpanContext?.spanId).toBe( + helperSpan?.spanContext().spanId, + ); + expect(helperSpan?.attributes).toMatchObject({ + "appkit.agent.name": "helper", + "appkit.app.name": "test-app", + "appkit.request.id": "request-1", + "mlflow.trace.session": "session-1", + "mlflow.trace.user": "user-1", + }); + expect(helperSpan?.attributes["appkit.thread.id"]).not.toBe( + "planner-thread", + ); + expect(plannerSpan?.attributes["mlflow.trace.tokenUsage"]).toBe( + '{"input_tokens":16,"output_tokens":7,"total_tokens":23}', + ); + expect(plannerSpan?.attributes["appkit.cost.available"]).toBe(false); + expect(plannerSpan?.attributes["mlflow.llm.cost"]).toBeUndefined(); + }); + test("function-form invoked exactly once per runAgent call", async () => { const toolsFn = vi.fn(() => ({})); const adapter: AgentAdapter = { diff --git a/packages/appkit/src/core/agent/trace-tool-call.ts b/packages/appkit/src/core/agent/trace-tool-call.ts new file mode 100644 index 000000000..4106f71f0 --- /dev/null +++ b/packages/appkit/src/core/agent/trace-tool-call.ts @@ -0,0 +1,80 @@ +import { type Span, SpanStatusCode, trace } from "@opentelemetry/api"; +import type { ToolEffect } from "shared"; +import { captureTraceValue } from "../../telemetry/agent-tracing"; + +const tracer = () => trace.getTracer("@databricks/appkit-agent-tracing"); + +export async function traceToolCall( + input: { + name: string; + source: string; + effect?: ToolEffect; + args: unknown; + }, + operation: (span: Span) => Promise, +): Promise { + return tracer().startActiveSpan( + `${input.name} tool`, + { + attributes: { + "mlflow.spanType": "TOOL", + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": input.name, + "appkit.tool.name": input.name, + "appkit.tool.source": input.source, + ...(input.effect ? { "appkit.tool.effect": input.effect } : {}), + }, + }, + async (span) => { + const startedAt = Date.now(); + setCapturedAttribute(span, "mlflow.spanInputs", input.args); + try { + const result = await operation(span); + setCapturedAttribute(span, "mlflow.spanOutputs", result); + span.setAttribute("appkit.tool.state", "completed"); + span.setStatus({ code: SpanStatusCode.OK }); + return result; + } catch (error) { + span.setAttribute("appkit.tool.state", "failed"); + recordSafeFailure(span, error); + throw error; + } finally { + span.setAttribute( + "appkit.tool.duration_ms", + Math.max(0, Date.now() - startedAt), + ); + span.end(); + } + }, + ); +} + +function setCapturedAttribute(span: Span, key: string, value: unknown): void { + const captured = captureTraceValue(value); + span.setAttribute(key, captured.value); + span.setAttribute(`${key}.original_bytes`, captured.originalBytes); + span.setAttribute(`${key}.sha256`, captured.sha256); + span.setAttribute(`${key}.truncated`, captured.truncated); +} + +function recordSafeFailure(span: Span, error: unknown): void { + const failure = captureTraceValue( + { + error: + error instanceof Error + ? error.message + : String(error ?? "Unknown error"), + }, + { redactKeys: ["error"] }, + ); + span.setAttribute("appkit.error", failure.value); + span.setAttribute("mlflow.spanOutputs", failure.value); + span.setAttribute("mlflow.spanOutputs.original_bytes", failure.originalBytes); + span.setAttribute("mlflow.spanOutputs.sha256", failure.sha256); + span.setAttribute("mlflow.spanOutputs.truncated", failure.truncated); + span.recordException({ name: "Error", message: "Tool operation failed" }); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: "Tool operation failed", + }); +} diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 5d734c130..749b0f920 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -1,12 +1,12 @@ import { randomUUID } from "node:crypto"; import path from "node:path"; -import { type Span, SpanStatusCode, trace } from "@opentelemetry/api"; import type express from "express"; import pc from "picocolors"; import type { AgentAdapter, AgentRunContext, AgentToolDefinition, + AgentUsage, IAppRouter, Message, PluginPhase, @@ -14,7 +14,6 @@ import type { ResponseStreamEvent, Thread, ToolAnnotations, - ToolEffect, ToolProvider, } from "shared"; import { @@ -38,6 +37,7 @@ import { isHostedTool, resolveHostedTools, } from "../../core/agent/tools"; +import { traceToolCall } from "../../core/agent/trace-tool-call"; import type { AgentDefinition, AgentsPluginConfig, @@ -55,7 +55,6 @@ import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; import { type AgentTraceObserver, - captureTraceValue, resolveAgentTraceAppName, runWithAgentTrace, } from "../../telemetry/agent-tracing"; @@ -76,87 +75,7 @@ const logger = createLogger("agents"); const DEFAULT_AGENTS_DIR = "./config/agents"; -const agentOperationTracer = () => - trace.getTracer("@databricks/appkit-agent-tracing"); - -export async function traceToolCall( - input: { - name: string; - source: string; - effect?: ToolEffect; - args: unknown; - }, - operation: (span: Span) => Promise, -): Promise { - return agentOperationTracer().startActiveSpan( - `${input.name} tool`, - { - attributes: { - "mlflow.spanType": "TOOL", - "gen_ai.operation.name": "execute_tool", - "gen_ai.tool.name": input.name, - "appkit.tool.name": input.name, - "appkit.tool.source": input.source, - ...(input.effect ? { "appkit.tool.effect": input.effect } : {}), - }, - }, - async (span) => { - const startedAt = Date.now(); - setToolCapturedAttribute(span, "mlflow.spanInputs", input.args); - try { - const result = await operation(span); - setToolCapturedAttribute(span, "mlflow.spanOutputs", result); - span.setAttribute("appkit.tool.state", "completed"); - span.setStatus({ code: SpanStatusCode.OK }); - return result; - } catch (error) { - span.setAttribute("appkit.tool.state", "failed"); - recordSafeToolFailure(span, error); - throw error; - } finally { - span.setAttribute( - "appkit.tool.duration_ms", - Math.max(0, Date.now() - startedAt), - ); - span.end(); - } - }, - ); -} - -function setToolCapturedAttribute( - span: Span, - key: string, - value: unknown, -): void { - const captured = captureTraceValue(value); - span.setAttribute(key, captured.value); - span.setAttribute(`${key}.original_bytes`, captured.originalBytes); - span.setAttribute(`${key}.sha256`, captured.sha256); - span.setAttribute(`${key}.truncated`, captured.truncated); -} - -function recordSafeToolFailure(span: Span, error: unknown): void { - const failure = captureTraceValue( - { - error: - error instanceof Error - ? error.message - : String(error ?? "Unknown error"), - }, - { redactKeys: ["error"] }, - ); - span.setAttribute("appkit.error", failure.value); - span.setAttribute("mlflow.spanOutputs", failure.value); - span.setAttribute("mlflow.spanOutputs.original_bytes", failure.originalBytes); - span.setAttribute("mlflow.spanOutputs.sha256", failure.sha256); - span.setAttribute("mlflow.spanOutputs.truncated", failure.truncated); - span.recordException({ name: "Error", message: "Tool operation failed" }); - span.setStatus({ - code: SpanStatusCode.ERROR, - message: "Tool operation failed", - }); -} +export { traceToolCall } from "../../core/agent/trace-tool-call"; /** * Context flag recorded on the in-memory AgentDefinition to indicate whether @@ -219,6 +138,13 @@ interface RunState { translator: AgentEventTranslator; outboundEvents: EventChannel; traceObserver?: AgentTraceObserver; + traceIdentity: { + appName: string; + route: "chat" | "invocations" | "responses"; + sessionId: string; + userId: string; + requestId: string; + }; /** Boxed mutable counter shared across parent + all sub-agent dispatches. */ toolCallsUsed: { count: number }; } @@ -1255,6 +1181,13 @@ export class AgentsPlugin extends Plugin implements ToolProvider { outboundEvents, toolCallsUsed: { count: 0 }, traceObserver: observer, + traceIdentity: { + appName: resolveAgentTraceAppName(), + route: "chat", + sessionId: requestSessionId(req) ?? thread.id, + userId, + requestId, + }, }; const executeTool = (name: string, args: unknown): Promise => @@ -1446,6 +1379,13 @@ export class AgentsPlugin extends Plugin implements ToolProvider { outboundEvents: new EventChannel(), toolCallsUsed: { count: 0 }, traceObserver: observer, + traceIdentity: { + appName: resolveAgentTraceAppName(), + route: invokeTraceRoute(req), + sessionId: requestSessionId(req) ?? thread.id, + userId, + requestId, + }, }; const executeTool = (name: string, args: unknown): Promise => @@ -1679,12 +1619,14 @@ export class AgentsPlugin extends Plugin implements ToolProvider { const childAgent = this.agents.get(entry.agentName); if (!childAgent) throw new Error(`Sub-agent not found: ${entry.agentName}`); - result = await this.runSubAgent( + const childResult = await this.runSubAgent( runState, childAgent, args, depth + 1, ); + runState.traceObserver?.addChildUsage(childResult.usage); + result = childResult.text; } else if (entry.source === "hosted-supervisor") { // Defense-in-depth: should never fire. Hosted-supervisor entries are // routed via `AgentInput.extensions` and the SA endpoint executes @@ -1705,9 +1647,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } /** - * Runs a sub-agent in response to an `agent-` tool call. Returns the - * concatenated text output to hand back to the parent adapter as the tool - * result. + * Runs a sub-agent in response to an `agent-` tool call. Returns its + * concatenated text and aggregate usage; dispatch folds that usage into the + * owning trace once, then hands only the text to outer tool normalization. * * `depth` starts at 1 for a top-level sub-agent invocation (i.e. the * outer `_streamAgent` calls `runSubAgent(..., 1)`) and increments on @@ -1723,7 +1665,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { child: RegisteredAgent, args: unknown, depth: number, - ): Promise { + ): Promise<{ text: string; usage: AgentUsage }> { if (depth > runState.limits.maxSubAgentDepth) { throw new Error( `Sub-agent depth exceeded (limit ${runState.limits.maxSubAgentDepth}). ` + @@ -1744,14 +1686,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { .filter((e) => e.source !== "hosted-supervisor") .map((e) => e.def); - const childExecute = (name: string, childArgs: unknown): Promise => - this.dispatchToolCall(runState, child.toolIndex, name, childArgs, depth); - - const runContext: AgentRunContext = { - executeTool: childExecute, - signal: runState.signal, - }; - const pluginNames = this.context ? this.context .getPluginNames() @@ -1782,39 +1716,70 @@ export class AgentsPlugin extends Plugin implements ToolProvider { }, ]; - const consumed = await consumeAdapterStream( - child.adapter.run( - { - messages, - tools: childTools, - threadId: randomUUID(), - signal: runState.signal, - extensions: buildAdapterExtensions(child.toolIndex), - }, - runContext, - ), + const childThreadId = randomUUID(); + const traced = await runWithAgentTrace( { - signal: runState.signal, - // Forward every sub-agent event into the parent's outbound SSE - // stream so the client sees nested tool_call / tool_result events - // (UI-action tools like apply_filter / highlight_period rely on - // this) and the sub-agent's streaming text as it's generated. - // - // `metadata` is the one exception: sub-agents have their own - // threadId, and forwarding it would overwrite the parent's - // thread state on the client and break multi-turn continuity. - // Approval-pending events emitted by `dispatchToolCall` already - // reach `outboundEvents` directly, so they are not routed here. - onEvent: (event) => { - runState.traceObserver?.onEvent(event); - if (event.type === "metadata") return; - for (const translated of runState.translator.translate(event)) { - runState.outboundEvents.push(translated); - } - }, + ...runState.traceIdentity, + agentName: child.name, + threadId: childThreadId, + }, + { messages: input }, + async (childObserver) => { + const childRunState: RunState = { + ...runState, + traceObserver: childObserver, + }; + const childExecute = ( + name: string, + childArgs: unknown, + ): Promise => + this.dispatchToolCall( + childRunState, + child.toolIndex, + name, + childArgs, + depth, + ); + const runContext: AgentRunContext = { + executeTool: childExecute, + signal: runState.signal, + }; + const consumed = await consumeAdapterStream( + child.adapter.run( + { + messages, + tools: childTools, + threadId: childThreadId, + signal: runState.signal, + extensions: buildAdapterExtensions(child.toolIndex), + }, + runContext, + ), + { + signal: runState.signal, + // Forward every sub-agent event into the parent's outbound SSE + // stream so the client sees nested tool_call / tool_result events + // (UI-action tools like apply_filter / highlight_period rely on + // this) and the sub-agent's streaming text as it's generated. + // + // `metadata` is the one exception: sub-agents have their own + // threadId, and forwarding it would overwrite the parent's + // thread state on the client and break multi-turn continuity. + // Approval-pending events emitted by `dispatchToolCall` already + // reach `outboundEvents` directly, so they are not routed here. + onEvent: (event) => { + childObserver.onEvent(event); + if (event.type === "metadata") return; + for (const translated of runState.translator.translate(event)) { + runState.outboundEvents.push(translated); + } + }, + }, + ); + return consumed.text; }, ); - return consumed.text; + return { text: traced.value, usage: traced.usage }; } private async _handleCancel(req: express.Request, res: express.Response) { diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index e2a1117ae..a58c39a77 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -131,6 +131,13 @@ function makeRunState(plugin: AgentsPlugin) { outboundEvents: { push: (event: unknown) => pushed.push(event), }, + traceIdentity: { + appName: "test-app", + route: "chat", + sessionId: "session-1", + userId: "alice", + requestId: "stream-1", + }, toolCallsUsed: { count: 0 }, }; return { runState, pushed, plugin }; @@ -180,7 +187,15 @@ describe("dispatchToolCall ā€” semantic TOOL spans", () => { callTool: vi.fn().mockResolvedValue({ content: "remote" }), }; // biome-ignore lint/suspicious/noExplicitAny: isolate dispatch from adapter streaming - (plugin as any).runSubAgent = vi.fn().mockResolvedValue("child output"); + (plugin as any).runSubAgent = vi.fn().mockResolvedValue({ + text: "child output", + usage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }, + }); // biome-ignore lint/suspicious/noExplicitAny: seed the child registry lookup (plugin as any).agents.set("researcher", { name: "researcher" }); @@ -916,4 +931,200 @@ describe("runSubAgent ā€” sub-agent event forwarding", () => { expect(types).toContain("tool_result"); expect(types).toContain("message_delta"); }); + + test("nests the helper AGENT inside its dispatch TOOL and adds child usage once", async () => { + const plugin = new AgentsPlugin({ dir: false, agents: {} }); + const { runState } = makeRunState(plugin); + Object.assign(runState, { + traceIdentity: { + appName: "test-app", + requestId: "request-1", + route: "chat", + sessionId: "session-1", + userId: "alice", + }, + }); + const helperStartedAt = Date.now(); + let childAdapterThreadId: string | undefined; + const helper = { + name: "helper", + instructions: "research", + adapter: { + async *run(input: any, runContext: any): any { + childAdapterThreadId = input.threadId; + yield { + type: "model_start", + stepId: "helper-step", + model: "helper.model", + provider: "databricks", + input: { messages: [{ role: "user", content: "research" }] }, + startedAt: helperStartedAt, + }; + const fact = await runContext.executeTool("helper.tool", { + topic: "tracing", + }); + yield { type: "message_delta", content: `helper:${fact}` }; + yield { + type: "model_end", + stepId: "helper-step", + model: "helper.model", + provider: "databricks", + output: { text: `helper:${fact}` }, + usage: { + inputTokens: 6, + outputTokens: 3, + totalTokens: 9, + costUsd: 0.02, + costAvailable: true, + }, + streamDurationMs: 5, + endedAt: helperStartedAt + 5, + }; + }, + }, + toolIndex: new Map([ + [ + "helper.tool", + { + source: "function", + def: { + name: "helper.tool", + description: "look up a fact", + parameters: { type: "object" }, + annotations: { effect: "read" }, + }, + functionTool: { + execute: async ({ topic }: { topic: string }) => `fact:${topic}`, + }, + }, + ], + ]), + }; + (plugin as any).agents.set("helper", helper); + const plannerToolIndex = new Map([ + [ + "agent-helper", + { + source: "subagent", + agentName: "helper", + def: { + name: "agent-helper", + description: "delegate", + parameters: { type: "object" }, + }, + }, + ], + ]); + const plannerStartedAt = Date.now(); + let aggregateUsage: unknown; + + const observed = await captureSpans(async () => { + const traced = await runWithAgentTrace( + { + appName: "test-app", + agentName: "planner", + route: "chat", + sessionId: "session-1", + userId: "alice", + requestId: "request-1", + threadId: "planner-thread", + }, + { message: "plan" }, + async (observer) => { + Object.assign(runState, { traceObserver: observer }); + observer.onEvent({ + type: "model_start", + stepId: "planner-step", + model: "planner.model", + provider: "databricks", + input: { messages: [{ role: "user", content: "plan" }] }, + startedAt: plannerStartedAt, + }); + const delegated = await callDispatch(plugin, { + runState, + toolIndex: plannerToolIndex, + name: "agent-helper", + args: { input: "research" }, + }); + observer.onEvent({ + type: "model_end", + stepId: "planner-step", + model: "planner.model", + provider: "databricks", + output: { text: String(delegated) }, + usage: { + inputTokens: 10, + outputTokens: 4, + totalTokens: 14, + costUsd: 0.04, + costAvailable: true, + }, + streamDurationMs: 8, + endedAt: plannerStartedAt + 8, + }); + return delegated; + }, + ); + aggregateUsage = traced.usage; + }); + + expect(observed.error).toBeUndefined(); + expect(aggregateUsage).toEqual({ + inputTokens: 16, + outputTokens: 7, + totalTokens: 23, + costUsd: 0.06, + costAvailable: true, + }); + const agentSpans = observed.spans.filter( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + const toolSpans = observed.spans.filter( + (span) => span.attributes["mlflow.spanType"] === "TOOL", + ); + const planner = agentSpans.find( + (span) => span.attributes["appkit.agent.name"] === "planner", + ); + const child = agentSpans.find( + (span) => span.attributes["appkit.agent.name"] === "helper", + ); + const delegation = toolSpans.find( + (span) => span.attributes["appkit.tool.name"] === "agent-helper", + ); + const helperTool = toolSpans.find( + (span) => span.attributes["appkit.tool.name"] === "helper.tool", + ); + const helperModel = observed.spans.find( + (span) => span.attributes["mlflow.chat.model"] === "helper.model", + ); + expect(agentSpans).toHaveLength(2); + expect(delegation?.parentSpanContext?.spanId).toBe( + planner?.spanContext().spanId, + ); + expect(child?.parentSpanContext?.spanId).toBe( + delegation?.spanContext().spanId, + ); + expect(helperModel?.parentSpanContext?.spanId).toBe( + child?.spanContext().spanId, + ); + expect(helperTool?.parentSpanContext?.spanId).toBe( + child?.spanContext().spanId, + ); + expect(child?.attributes).toMatchObject({ + "appkit.agent.name": "helper", + "appkit.app.name": "test-app", + "appkit.request.id": "request-1", + "appkit.thread.id": childAdapterThreadId, + "mlflow.trace.session": "session-1", + "mlflow.trace.user": "alice", + }); + expect(childAdapterThreadId).not.toBe("planner-thread"); + expect(planner?.attributes["mlflow.trace.tokenUsage"]).toBe( + '{"input_tokens":16,"output_tokens":7,"total_tokens":23}', + ); + expect(planner?.attributes["mlflow.llm.cost"]).toBe(0.06); + expect( + JSON.parse(String(delegation?.attributes["mlflow.spanOutputs"])), + ).toBe("helper:fact:tracing"); + }); }); diff --git a/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts b/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts index 2a4e4a221..0d01fca8e 100644 --- a/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts @@ -338,6 +338,13 @@ describe("runSubAgent ā€” depth guard", () => { outboundEvents: { push: vi.fn(), }, + traceIdentity: { + appName: "test-app", + route: "chat", + sessionId: "session-1", + userId: "alice", + requestId: "test-stream", + }, toolCallsUsed: { count: 0 }, }; } @@ -388,6 +395,14 @@ describe("runSubAgent ā€” depth guard", () => { { input: "test" }, 3, // at the limit, not over ); - expect(result).toBe("hello from depth-3"); + expect(result).toEqual({ + text: "hello from depth-3", + usage: { + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + costAvailable: false, + }, + }); }); }); diff --git a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts index 4160d13d2..e44bf73fe 100644 --- a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts +++ b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts @@ -50,6 +50,7 @@ vi.mock("../../../telemetry", () => ({ _telemetryOpts?: unknown, ) => fn({ + end: vi.fn(), setAttribute: vi.fn(), setStatus: vi.fn(), recordException: vi.fn(), diff --git a/packages/appkit/src/telemetry/agent-tracing/tracer.ts b/packages/appkit/src/telemetry/agent-tracing/tracer.ts index 76beead42..2877d0f10 100644 --- a/packages/appkit/src/telemetry/agent-tracing/tracer.ts +++ b/packages/appkit/src/telemetry/agent-tracing/tracer.ts @@ -137,6 +137,9 @@ export async function runWithAgentTrace( reportedError ??= event.error; } }, + addChildUsage(childUsage) { + usage.add(normalizeUsage(childUsage)); + }, updateIdentity(next) { setIdentityAttributes(root, next); }, diff --git a/packages/appkit/src/telemetry/agent-tracing/types.ts b/packages/appkit/src/telemetry/agent-tracing/types.ts index 326522b6a..2735b8c06 100644 --- a/packages/appkit/src/telemetry/agent-tracing/types.ts +++ b/packages/appkit/src/telemetry/agent-tracing/types.ts @@ -34,6 +34,8 @@ export interface AgentTraceObserver { /** MLflow V4 identity when UC is active; otherwise the 32-hex OTel trace ID. */ readonly traceId: string; onEvent(event: AgentEvent): void; + /** Adds one completed local child trace's aggregate usage to this root. */ + addChildUsage(usage: AgentUsage): void; updateIdentity(identity: Partial>): void; setOutput(output: unknown): void; recordError(error: unknown, output?: unknown): void; From fe802e1e5c3ca179b6b09d26744009c52cd12c2c Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 20:44:53 -0700 Subject: [PATCH 15/31] fix: harden agent and retriever tracing Signed-off-by: Adam Gurary --- .../appkit/src/connectors/ai-search/client.ts | 50 ++-- .../connectors/ai-search/tests/client.test.ts | 228 +++++++++++++++++- packages/appkit/src/core/agent/run-agent.ts | 2 +- .../src/core/agent/tests/run-agent.test.ts | 116 +++++++++ packages/appkit/src/plugins/agents/agents.ts | 2 +- .../agents/tests/dispatch-tool-call.test.ts | 114 +++++++++ .../src/telemetry/agent-tracing/attributes.ts | 1 + .../src/telemetry/agent-tracing/index.ts | 6 +- .../telemetry/agent-tracing/serialization.ts | 8 +- .../agent-tracing/tests/serialization.test.ts | 55 +++++ .../src/telemetry/agent-tracing/tracer.ts | 64 +++-- 11 files changed, 601 insertions(+), 45 deletions(-) diff --git a/packages/appkit/src/connectors/ai-search/client.ts b/packages/appkit/src/connectors/ai-search/client.ts index f8793ebc4..be7a89716 100644 --- a/packages/appkit/src/connectors/ai-search/client.ts +++ b/packages/appkit/src/connectors/ai-search/client.ts @@ -6,7 +6,10 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; -import { captureTraceValue } from "../../telemetry/agent-tracing"; +import { + captureTraceValue, + getActiveAgentTraceIdentity, +} from "../../telemetry/agent-tracing"; import type { WorkspaceClient } from "../../workspace-client"; import { contextFromAbortSignal } from "../context"; import type { @@ -50,10 +53,6 @@ export class AiSearchConnector { params: VsQueryParams, signal?: AbortSignal, ): Promise { - if (signal?.aborted) { - throw new Error("Query cancelled before execution"); - } - const body: Record = { columns: params.columns, num_results: params.numResults, @@ -98,6 +97,7 @@ export class AiSearchConnector { params.filters && Object.keys(params.filters).length > 0 ), "vs.has_reranker": !!params.reranker, + ...activeAgentIdentityAttributes(), }, }, async (span: Span) => { @@ -113,6 +113,9 @@ export class AiSearchConnector { reranker: params.reranker ?? null, }); try { + if (signal?.aborted) { + throw new Error("Query cancelled before execution"); + } const response = (await workspaceClient.apiClient.request( { method: "POST", @@ -126,7 +129,7 @@ export class AiSearchConnector { )) as VsRawResponse; const duration = Date.now() - startTime; - const outputs = retrieverOutputs(response, params.columns); + const outputs = retrieverOutputs(response); setRetrieverOutputs(span, outputs); span.setAttribute("vs.result_count", response.result.row_count); span.setAttribute( @@ -163,10 +166,6 @@ export class AiSearchConnector { params: VsNextPageParams, signal?: AbortSignal, ): Promise { - if (signal?.aborted) { - throw new Error("Query cancelled before execution"); - } - logger.debug( "Fetching next page for index %s (endpoint=%s)", params.indexName, @@ -185,6 +184,7 @@ export class AiSearchConnector { "db.system": "databricks", "vs.index_name": params.indexName, "vs.endpoint_name": params.endpointName, + ...activeAgentIdentityAttributes(), }, }, async (span: Span) => { @@ -196,6 +196,9 @@ export class AiSearchConnector { queryType: "next_page", }); try { + if (signal?.aborted) { + throw new Error("Query cancelled before execution"); + } const response = (await workspaceClient.apiClient.request( { method: "POST", @@ -286,20 +289,11 @@ function summarizeVector( return { dimensions: vector.length, sha256: captured.sha256 }; } -function retrieverOutputs( - response: VsRawResponse, - configuredColumns?: readonly string[], -): RetrieverOutputs { +function retrieverOutputs(response: VsRawResponse): RetrieverOutputs { const columnNames = response.manifest.columns.map((column) => column.name); - const configured = configuredColumns - ? new Set(configuredColumns.map((column) => column.toLowerCase())) - : undefined; const documentIdIndex = columnNames.findIndex((column) => { const normalized = column.toLowerCase(); - return ( - DOCUMENT_ID_COLUMNS.has(normalized) && - (configured === undefined || configured.has(normalized)) - ); + return DOCUMENT_ID_COLUMNS.has(normalized); }); const scoreIndex = columnNames.findIndex( (column) => column.toLowerCase() === "score", @@ -327,6 +321,20 @@ function retrieverOutputs( }; } +function activeAgentIdentityAttributes(): Record { + const identity = getActiveAgentTraceIdentity(); + return identity + ? { + "appkit.agent.name": identity.agentName, + "appkit.app.name": identity.appName, + "appkit.request.id": identity.requestId, + "appkit.thread.id": identity.threadId, + "mlflow.trace.session": identity.sessionId, + "mlflow.trace.user": identity.userId, + } + : {}; +} + function setRetrieverOutputs(span: Span, outputs: RetrieverOutputs): void { setCapturedAttribute(span, "mlflow.spanOutputs", outputs); span.setAttribute("appkit.retriever.result_count", outputs.resultCount); diff --git a/packages/appkit/src/connectors/ai-search/tests/client.test.ts b/packages/appkit/src/connectors/ai-search/tests/client.test.ts index f64b5d93a..930d324d2 100644 --- a/packages/appkit/src/connectors/ai-search/tests/client.test.ts +++ b/packages/appkit/src/connectors/ai-search/tests/client.test.ts @@ -1,11 +1,13 @@ -import { SpanStatusCode, trace } from "@opentelemetry/api"; +import { context, SpanStatusCode, trace } from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; import { BasicTracerProvider, InMemorySpanExporter, type ReadableSpan, SimpleSpanProcessor, } from "@opentelemetry/sdk-trace-base"; -import { describe, expect, test, vi } from "vitest"; +import { afterAll, beforeAll, describe, expect, test, vi } from "vitest"; +import { runWithAgentTrace } from "../../../telemetry/agent-tracing"; import type { WorkspaceClient } from "../../../workspace-client"; import { AiSearchConnector } from "../client"; import type { VsRawResponse } from "../types"; @@ -56,6 +58,17 @@ function retrieverSpan(spans: ReadableSpan[]): ReadableSpan { return span as ReadableSpan; } +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); +}); + +afterAll(() => { + context.disable(); +}); + describe("AiSearchConnector semantic retrieval spans", () => { test("exports complete query inputs, documents, stable IDs, scores, and diagnostics", async () => { const response: VsRawResponse = { @@ -283,4 +296,215 @@ describe("AiSearchConnector semantic retrieval spans", () => { JSON.stringify({ attributes: span.attributes, events: span.events }), ).not.toContain("hunter2"); }); + + test("inherits active agent identity and parentage for query and next-page spans", async () => { + const response: VsRawResponse = { + manifest: { + column_count: 2, + columns: [{ name: "id" }, { name: "text" }], + }, + result: { row_count: 1, data_array: [["doc-1", "content"]] }, + next_page_token: null, + }; + const connector = new AiSearchConnector(); + const client = workspaceClient(response); + + const observed = await captureSpans(() => + runWithAgentTrace( + { + appName: "test-app", + agentName: "planner", + route: "chat", + sessionId: "session-1", + userId: "user-1", + requestId: "request-1", + threadId: "thread-1", + }, + { message: "retrieve" }, + async (observer) => { + observer.updateIdentity({ + agentName: "resolved-agent", + appName: "resolved-app", + requestId: "resolved-request", + sessionId: "resolved-session", + threadId: "resolved-thread", + userId: "resolved-user", + }); + await connector.query(client, { + columns: ["id", "text"], + indexName: "catalog.schema.docs", + numResults: 1, + queryText: "find it", + queryType: "ann", + }); + await connector.queryNextPage(client, { + endpointName: "endpoint-a", + indexName: "catalog.schema.docs", + pageToken: "page-2", + }); + return "done"; + }, + ), + ); + + expect(observed.error).toBeUndefined(); + const root = observed.spans.find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + const retrievers = observed.spans.filter( + (span) => span.attributes["mlflow.spanType"] === "RETRIEVER", + ); + expect(retrievers).toHaveLength(2); + for (const span of retrievers) { + expect(span.parentSpanContext?.spanId).toBe(root?.spanContext().spanId); + expect(span.attributes).toMatchObject({ + "appkit.agent.name": "resolved-agent", + "appkit.app.name": "resolved-app", + "appkit.request.id": "resolved-request", + "appkit.thread.id": "resolved-thread", + "mlflow.trace.session": "resolved-session", + "mlflow.trace.user": "resolved-user", + }); + } + }); + + test("exports a safe failed query span when the signal is already aborted", async () => { + const request = vi.fn(); + const client = { apiClient: { request } } as unknown as WorkspaceClient; + const controller = new AbortController(); + controller.abort(); + const connector = new AiSearchConnector(); + + const observed = await captureSpans(() => + connector.query( + client, + { + columns: ["id", "text"], + indexName: "catalog.schema.docs", + numResults: 2, + queryText: "cancelled query", + queryType: "ann", + }, + controller.signal, + ), + ); + + expect(observed.error).toMatchObject({ + message: "Query cancelled before execution", + }); + expect(request).not.toHaveBeenCalled(); + const span = retrieverSpan(observed.spans); + expect(span.status).toEqual({ + code: SpanStatusCode.ERROR, + message: "Retriever operation failed", + }); + expect(span.attributes).toMatchObject({ + "appkit.retriever.latency_ms": expect.any(Number), + "vs.duration_ms": expect.any(Number), + }); + expect(JSON.parse(String(span.attributes["mlflow.spanInputs"]))).toEqual({ + columns: ["id", "text"], + filters: {}, + indexName: "catalog.schema.docs", + numResults: 2, + queryText: "cancelled query", + queryType: "ann", + queryVector: null, + reranker: null, + }); + expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ + error: "[REDACTED]", + }); + expect(span.events).toEqual([ + expect.objectContaining({ name: "exception" }), + ]); + }); + + test("exports a safe failed next-page span when the signal is already aborted", async () => { + const request = vi.fn(); + const client = { apiClient: { request } } as unknown as WorkspaceClient; + const controller = new AbortController(); + controller.abort(); + const connector = new AiSearchConnector(); + + const observed = await captureSpans(() => + connector.queryNextPage( + client, + { + endpointName: "endpoint-a", + indexName: "catalog.schema.docs", + pageToken: "page-2", + }, + controller.signal, + ), + ); + + expect(observed.error).toMatchObject({ + message: "Query cancelled before execution", + }); + expect(request).not.toHaveBeenCalled(); + const span = retrieverSpan(observed.spans); + expect(span.status).toEqual({ + code: SpanStatusCode.ERROR, + message: "Retriever operation failed", + }); + expect(span.attributes).toMatchObject({ + "appkit.retriever.latency_ms": expect.any(Number), + "vs.duration_ms": expect.any(Number), + }); + expect(JSON.parse(String(span.attributes["mlflow.spanInputs"]))).toEqual({ + endpointName: "endpoint-a", + indexName: "catalog.schema.docs", + pageToken: "page-2", + queryType: "next_page", + }); + expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ + error: "[REDACTED]", + }); + expect(span.events).toEqual([ + expect.objectContaining({ name: "exception" }), + ]); + }); + + test("prefers a returned manifest ID even when requested columns omit it", async () => { + const response: VsRawResponse = { + manifest: { + column_count: 2, + columns: [{ name: "text" }, { name: "id" }], + }, + result: { + row_count: 1, + data_array: [["Complete row", "manifest-doc-id"]], + }, + next_page_token: null, + }; + const connector = new AiSearchConnector(); + + const observed = await captureSpans(() => + connector.query(workspaceClient(response), { + columns: ["text"], + indexName: "catalog.schema.docs", + numResults: 1, + queryText: "find it", + queryType: "ann", + }), + ); + + expect(observed.error).toBeUndefined(); + const span = retrieverSpan(observed.spans); + expect(span.attributes["appkit.retriever.document_ids"]).toEqual([ + "manifest-doc-id", + ]); + expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ + documents: [ + { + content: { id: "manifest-doc-id", text: "Complete row" }, + documentId: "manifest-doc-id", + score: null, + }, + ], + nextPageToken: null, + resultCount: 1, + }); + }); }); diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index b5a46e7cc..3ddeec2e9 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -235,12 +235,12 @@ async function runAgentInternal( providerCache, childObserver, ), + (childUsage) => observer.addChildUsage(childUsage), ); const childResult = { text: childTrace.value.text, usage: childTrace.usage, }; - observer.addChildUsage(childResult.usage); return childResult.text; } if (entry.kind === "hosted-supervisor") { diff --git a/packages/appkit/src/core/agent/tests/run-agent.test.ts b/packages/appkit/src/core/agent/tests/run-agent.test.ts index c97d338e2..33f8ba413 100644 --- a/packages/appkit/src/core/agent/tests/run-agent.test.ts +++ b/packages/appkit/src/core/agent/tests/run-agent.test.ts @@ -69,6 +69,33 @@ async function captureAgentSpans( return { result, spans }; } +async function captureFailedAgentSpans( + operation: () => Promise, +): Promise<{ error: unknown; spans: ReadableSpan[] }> { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracerSpy = vi + .spyOn(trace, "getTracer") + .mockImplementation((name: string, version?: string) => + provider.getTracer(name, version), + ); + let error: unknown; + let spans: ReadableSpan[] = []; + try { + await operation(); + } catch (caught) { + error = caught; + } finally { + await provider.forceFlush(); + spans = exporter.getFinishedSpans(); + getTracerSpy.mockRestore(); + await provider.shutdown(); + } + return { error, spans }; +} + describe("runAgent", () => { test("drives the adapter and returns aggregated text", async () => { const events: AgentEvent[] = [ @@ -546,6 +573,95 @@ describe("runAgent", () => { expect(plannerSpan?.attributes["mlflow.llm.cost"]).toBeUndefined(); }); + test("rolls failed unpriced child usage into the planner once and rethrows the original error", async () => { + const failure = new Error("helper failed after model usage"); + const plannerStartedAt = Date.now(); + const helperStartedAt = plannerStartedAt + 10; + const helperAdapter: AgentAdapter = { + async *run() { + yield { + type: "model_start", + stepId: "helper-step", + model: "helper.model", + provider: "databricks", + input: { messages: [{ role: "user", content: "research" }] }, + startedAt: helperStartedAt, + } satisfies AgentEvent; + yield { + type: "model_end", + stepId: "helper-step", + model: "helper.model", + provider: "databricks", + output: { text: "partial research" }, + usage: { + inputTokens: 6, + outputTokens: 3, + totalTokens: 9, + costAvailable: false, + }, + streamDurationMs: 5, + endedAt: helperStartedAt + 5, + } satisfies AgentEvent; + throw failure; + }, + }; + const plannerAdapter: AgentAdapter = { + async *run(_input, runContext) { + yield { + type: "model_start", + stepId: "planner-step", + model: "planner.model", + provider: "databricks", + input: { messages: [{ role: "user", content: "plan" }] }, + startedAt: plannerStartedAt, + } satisfies AgentEvent; + yield { + type: "model_end", + stepId: "planner-step", + model: "planner.model", + provider: "databricks", + output: { text: "delegating" }, + usage: { + inputTokens: 10, + outputTokens: 4, + totalTokens: 14, + costUsd: 0.04, + costAvailable: true, + }, + streamDurationMs: 8, + endedAt: plannerStartedAt + 8, + } satisfies AgentEvent; + await runContext.executeTool("agent-helper", { input: "research" }); + }, + }; + const planner = createAgent({ + name: "planner", + instructions: "plan", + model: plannerAdapter, + agents: { + helper: createAgent({ + name: "helper", + instructions: "research", + model: helperAdapter, + }), + }, + }); + + const observed = await captureFailedAgentSpans(() => + runAgent(planner, { messages: "plan" }), + ); + + expect(observed.error).toBe(failure); + const plannerSpan = observed.spans.find( + (span) => span.attributes["appkit.agent.name"] === "planner", + ); + expect(plannerSpan?.attributes["mlflow.trace.tokenUsage"]).toBe( + '{"input_tokens":16,"output_tokens":7,"total_tokens":23}', + ); + expect(plannerSpan?.attributes["appkit.cost.available"]).toBe(false); + expect(plannerSpan?.attributes["mlflow.llm.cost"]).toBeUndefined(); + }); + test("function-form invoked exactly once per runAgent call", async () => { const toolsFn = vi.fn(() => ({})); const adapter: AgentAdapter = { diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 749b0f920..0a74febd0 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -1625,7 +1625,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { args, depth + 1, ); - runState.traceObserver?.addChildUsage(childResult.usage); result = childResult.text; } else if (entry.source === "hosted-supervisor") { // Defense-in-depth: should never fire. Hosted-supervisor entries are @@ -1778,6 +1777,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ); return consumed.text; }, + (childUsage) => runState.traceObserver?.addChildUsage(childUsage), ); return { text: traced.value, usage: traced.usage }; } diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index a58c39a77..fafe40b31 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -1127,4 +1127,118 @@ describe("runSubAgent ā€” sub-agent event forwarding", () => { JSON.parse(String(delegation?.attributes["mlflow.spanOutputs"])), ).toBe("helper:fact:tracing"); }); + + test("adds failed unpriced child usage once and rethrows the original error", async () => { + const plugin = new AgentsPlugin({ dir: false, agents: {} }); + const { runState } = makeRunState(plugin); + const failure = new Error("helper failed after model usage"); + const startedAt = Date.now(); + const helper = { + name: "helper", + instructions: "research", + adapter: { + // biome-ignore lint/suspicious/noExplicitAny: stub adapter shape + async *run(): any { + yield { + type: "model_start", + stepId: "helper-step", + model: "helper.model", + provider: "databricks", + input: { messages: [{ role: "user", content: "research" }] }, + startedAt, + }; + yield { + type: "model_end", + stepId: "helper-step", + model: "helper.model", + provider: "databricks", + output: { text: "partial research" }, + usage: { + inputTokens: 6, + outputTokens: 3, + totalTokens: 9, + costAvailable: false, + }, + streamDurationMs: 5, + endedAt: startedAt + 5, + }; + throw failure; + }, + }, + toolIndex: new Map(), + }; + (plugin as any).agents.set("helper", helper); + const plannerToolIndex = new Map([ + [ + "agent-helper", + { + source: "subagent", + agentName: "helper", + def: { + name: "agent-helper", + description: "delegate", + parameters: { type: "object" }, + }, + }, + ], + ]); + + const observed = await captureSpans(() => + runWithAgentTrace( + { + appName: "test-app", + agentName: "planner", + route: "chat", + sessionId: "session-1", + userId: "alice", + requestId: "request-1", + threadId: "planner-thread", + }, + { message: "plan" }, + async (observer) => { + Object.assign(runState, { traceObserver: observer }); + observer.onEvent({ + type: "model_start", + stepId: "planner-step", + model: "planner.model", + provider: "databricks", + input: { messages: [{ role: "user", content: "plan" }] }, + startedAt, + }); + observer.onEvent({ + type: "model_end", + stepId: "planner-step", + model: "planner.model", + provider: "databricks", + output: { text: "delegating" }, + usage: { + inputTokens: 10, + outputTokens: 4, + totalTokens: 14, + costUsd: 0.04, + costAvailable: true, + }, + streamDurationMs: 8, + endedAt: startedAt + 8, + }); + return callDispatch(plugin, { + runState, + toolIndex: plannerToolIndex, + name: "agent-helper", + args: { input: "research" }, + }); + }, + ), + ); + + expect(observed.error).toBe(failure); + const plannerSpan = observed.spans.find( + (span) => span.attributes["appkit.agent.name"] === "planner", + ); + expect(plannerSpan?.attributes["mlflow.trace.tokenUsage"]).toBe( + '{"input_tokens":16,"output_tokens":7,"total_tokens":23}', + ); + expect(plannerSpan?.attributes["appkit.cost.available"]).toBe(false); + expect(plannerSpan?.attributes["mlflow.llm.cost"]).toBeUndefined(); + }); }); diff --git a/packages/appkit/src/telemetry/agent-tracing/attributes.ts b/packages/appkit/src/telemetry/agent-tracing/attributes.ts index 97721f0e7..c4701d9fe 100644 --- a/packages/appkit/src/telemetry/agent-tracing/attributes.ts +++ b/packages/appkit/src/telemetry/agent-tracing/attributes.ts @@ -17,6 +17,7 @@ export const DEFAULT_TRACE_REDACT_KEYS = [ "refresh_token", "databricks-token", "databricks_token", + "sdk-token", "password", "secret", "client-secret", diff --git a/packages/appkit/src/telemetry/agent-tracing/index.ts b/packages/appkit/src/telemetry/agent-tracing/index.ts index 4483a96ad..c55eb8412 100644 --- a/packages/appkit/src/telemetry/agent-tracing/index.ts +++ b/packages/appkit/src/telemetry/agent-tracing/index.ts @@ -1,5 +1,9 @@ export { captureTraceValue } from "./serialization"; -export { resolveAgentTraceAppName, runWithAgentTrace } from "./tracer"; +export { + getActiveAgentTraceIdentity, + resolveAgentTraceAppName, + runWithAgentTrace, +} from "./tracer"; export type { AgentTraceObserver, CapturedTraceValue, diff --git a/packages/appkit/src/telemetry/agent-tracing/serialization.ts b/packages/appkit/src/telemetry/agent-tracing/serialization.ts index 4a48d491a..8fe5ff943 100644 --- a/packages/appkit/src/telemetry/agent-tracing/serialization.ts +++ b/packages/appkit/src/telemetry/agent-tracing/serialization.ts @@ -12,12 +12,12 @@ export function captureTraceValue( ): CapturedTraceValue { const redactKeys = new Set( [...DEFAULT_TRACE_REDACT_KEYS, ...(options.redactKeys ?? [])].map((key) => - key.toLowerCase(), + normalizeRedactKey(key), ), ); const serialized = JSON.stringify(value, (key, current) => { - if (redactKeys.has(key.toLowerCase())) return REDACTED_TRACE_VALUE; + if (redactKeys.has(normalizeRedactKey(key))) return REDACTED_TRACE_VALUE; if ( current !== null && typeof current === "object" && @@ -47,6 +47,10 @@ export function captureTraceValue( }; } +function normalizeRedactKey(value: string): string { + return value.toLowerCase().replaceAll(/[^a-z0-9]/g, ""); +} + function normalizeMaxBytes(value: number | undefined): number { if (value === undefined) return DEFAULT_TRACE_VALUE_MAX_BYTES; if (!Number.isFinite(value) || value < 0) { diff --git a/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts index 9d8d3fddd..c1f222ee5 100644 --- a/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts +++ b/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts @@ -32,6 +32,61 @@ describe("captureTraceValue", () => { }); }); + test("redacts credential keys across camelCase, separator, and case variants", () => { + const captured = captureTraceValue({ + accessToken: "camel-access", + "ACCESS.TOKEN": "dot-access", + access_token: "snake-access", + refreshToken: "camel-refresh", + "refresh token": "space-refresh", + clientSecret: "camel-client", + "CLIENT-SECRET": "kebab-client", + sdkToken: "camel-sdk", + SDK_TOKEN: "snake-sdk", + }); + + expect(JSON.parse(captured.value)).toEqual({ + "ACCESS.TOKEN": "[REDACTED]", + "CLIENT-SECRET": "[REDACTED]", + SDK_TOKEN: "[REDACTED]", + accessToken: "[REDACTED]", + access_token: "[REDACTED]", + clientSecret: "[REDACTED]", + refreshToken: "[REDACTED]", + "refresh token": "[REDACTED]", + sdkToken: "[REDACTED]", + }); + expect(captured.value).not.toContain("camel-access"); + expect(captured.value).not.toContain("camel-refresh"); + expect(captured.value).not.toContain("camel-client"); + expect(captured.value).not.toContain("camel-sdk"); + }); + + test("normalizes custom redaction keys without redacting benign supersets", () => { + const captured = captureTraceValue( + { + customSecret: "private", + CUSTOM_SECRET: "also-private", + accessTokenCount: 4, + clientSecretName: "display-name", + refreshTokenizedAt: "2026-08-11", + sdkTokenizer: "sentencepiece", + secretSauce: "benign", + }, + { redactKeys: ["custom-secret"] }, + ); + + expect(JSON.parse(captured.value)).toEqual({ + CUSTOM_SECRET: "[REDACTED]", + accessTokenCount: 4, + clientSecretName: "display-name", + customSecret: "[REDACTED]", + refreshTokenizedAt: "2026-08-11", + sdkTokenizer: "sentencepiece", + secretSauce: "benign", + }); + }); + test("hashes the complete value and truncates only at UTF-8 boundaries", () => { expect(captureTraceValue("šŸ˜€a", { maxBytes: 5 })).toEqual({ value: '"šŸ˜€', diff --git a/packages/appkit/src/telemetry/agent-tracing/tracer.ts b/packages/appkit/src/telemetry/agent-tracing/tracer.ts index 2877d0f10..af307ae3d 100644 --- a/packages/appkit/src/telemetry/agent-tracing/tracer.ts +++ b/packages/appkit/src/telemetry/agent-tracing/tracer.ts @@ -1,6 +1,7 @@ import { randomUUID } from "node:crypto"; import { context, + createContextKey, isSpanContextValid, type Span, SpanStatusCode, @@ -21,6 +22,9 @@ import type { import { AgentUsageAccumulator } from "./usage"; const tracer = () => trace.getTracer("@databricks/appkit-agent-tracing"); +const ACTIVE_AGENT_TRACE_IDENTITY = createContextKey( + "@databricks/appkit-agent-trace-identity", +); interface ActiveModelStep { event: AgentModelStartEvent; @@ -37,23 +41,35 @@ export function resolveAgentTraceAppName(appName?: string): string { ); } +export function getActiveAgentTraceIdentity(): AgentTraceIdentity | undefined { + const identity = context.active().getValue(ACTIVE_AGENT_TRACE_IDENTITY) as + | AgentTraceIdentity + | undefined; + return identity ? { ...identity } : undefined; +} + export async function runWithAgentTrace( identity: AgentTraceIdentity, inputs: unknown, operation: (observer: AgentTraceObserver) => Promise, + onCompleteUsage?: (usage: AgentUsage) => void, ): Promise> { + const activeIdentity: AgentTraceIdentity = { + ...identity, + appName: resolveAgentTraceAppName(identity.appName), + }; return tracer().startActiveSpan( `${identity.agentName} agent`, { attributes: { "mlflow.spanType": "AGENT", - "mlflow.trace.session": identity.sessionId, - "mlflow.trace.user": identity.userId, - "appkit.app.name": resolveAgentTraceAppName(identity.appName), - "appkit.request.id": identity.requestId, - "appkit.thread.id": identity.threadId, - "appkit.agent.name": identity.agentName, - "appkit.route": identity.route, + "mlflow.trace.session": activeIdentity.sessionId, + "mlflow.trace.user": activeIdentity.userId, + "appkit.app.name": activeIdentity.appName, + "appkit.request.id": activeIdentity.requestId, + "appkit.thread.id": activeIdentity.threadId, + "appkit.agent.name": activeIdentity.agentName, + "appkit.route": activeIdentity.route, }, }, async (root) => { @@ -63,16 +79,18 @@ export async function runWithAgentTrace( const otelTraceId = hasValidRootContext ? rootSpanContext.traceId : randomUUID().replaceAll("-", ""); - const rootContext = trace.setSpan( - context.active(), - hasValidRootContext - ? root - : trace.wrapSpanContext({ - traceId: otelTraceId, - spanId: randomUUID().replaceAll("-", "").slice(0, 16), - traceFlags: 0, - }), - ); + const rootContext = trace + .setSpan( + context.active(), + hasValidRootContext + ? root + : trace.wrapSpanContext({ + traceId: otelTraceId, + spanId: randomUUID().replaceAll("-", "").slice(0, 16), + traceFlags: 0, + }), + ) + .setValue(ACTIVE_AGENT_TRACE_IDENTITY, activeIdentity); const traceId = getMlflowUcTraceId(otelTraceId) ?? otelTraceId; const usage = new AgentUsageAccumulator(); const activeModels = new Map(); @@ -141,6 +159,7 @@ export async function runWithAgentTrace( usage.add(normalizeUsage(childUsage)); }, updateIdentity(next) { + updateActiveIdentity(activeIdentity, next); setIdentityAttributes(root, next); }, setOutput(output) { @@ -207,6 +226,7 @@ export async function runWithAgentTrace( const finalUsage = usage.snapshot(); setRootUsageAttributes(root, finalUsage); + onCompleteUsage?.(finalUsage); const finalOutputText = outputText || textFromValue(value); const finalOutput = hasExplicitOutput ? explicitOutput @@ -237,6 +257,16 @@ export async function runWithAgentTrace( ); } +function updateActiveIdentity( + identity: AgentTraceIdentity, + next: Partial>, +): void { + Object.assign(identity, next); + if (next.appName !== undefined) { + identity.appName = resolveAgentTraceAppName(next.appName); + } +} + function modelStartAttributes(event: AgentModelStartEvent) { return { "mlflow.spanType": "CHAT_MODEL", From 61963cd49af792019cb6655b0644fdc68b5ea099 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 20:58:01 -0700 Subject: [PATCH 16/31] feat: propagate agent trace context Signed-off-by: Adam Gurary --- packages/appkit/src/agents/databricks.ts | 12 +- packages/appkit/src/agents/supervisor-api.ts | 3 +- .../src/agents/tests/databricks.test.ts | 117 ++++++++- .../src/agents/tests/supervisor-api.test.ts | 71 ++++- packages/appkit/src/connectors/mcp/client.ts | 53 +++- .../src/connectors/mcp/tests/client.test.ts | 242 +++++++++++++++++- .../appkit/src/connectors/serving/client.ts | 11 +- .../connectors/serving/tests/client.test.ts | 79 +++++- .../src/telemetry/agent-tracing/index.ts | 5 + .../telemetry/agent-tracing/propagation.ts | 73 ++++++ .../agent-tracing/tests/propagation.test.ts | 216 ++++++++++++++++ 11 files changed, 847 insertions(+), 35 deletions(-) create mode 100644 packages/appkit/src/telemetry/agent-tracing/propagation.ts create mode 100644 packages/appkit/src/telemetry/agent-tracing/tests/propagation.test.ts diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index a24412075..ba6d0c2e5 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -14,6 +14,7 @@ import { stream as servingStream, } from "../connectors/serving/client"; import { APPKIT_USER_AGENT, getClientOptions } from "../context/client-options"; +import { injectActiveTraceContext } from "../telemetry/agent-tracing"; import { DEFAULT_TRACE_REDACT_KEYS, REDACTED_TRACE_VALUE, @@ -492,13 +493,16 @@ export class DatabricksAdapter implements AgentAdapter { const fetchSignal = signal ?? AbortSignal.timeout(RAW_FETCH_DEFAULT_TIMEOUT_MS); const authHeaders = await authenticate(); - const response = await fetch(endpointUrl, { - method: "POST", - headers: { + const headers = injectActiveTraceContext( + new Headers({ "User-Agent": APPKIT_USER_AGENT, "Content-Type": "application/json", ...authHeaders, - }, + }), + ); + const response = await fetch(endpointUrl, { + method: "POST", + headers, body: JSON.stringify(body), signal: fetchSignal, }); diff --git a/packages/appkit/src/agents/supervisor-api.ts b/packages/appkit/src/agents/supervisor-api.ts index 07c4d1a79..abd085fdc 100644 --- a/packages/appkit/src/agents/supervisor-api.ts +++ b/packages/appkit/src/agents/supervisor-api.ts @@ -346,7 +346,8 @@ interface SupervisorApiAdapterCtorOptions { * same mechanism used by `DatabricksAdapter.fromModelServing`. The transport * is injected via {@link SupervisorApiAdapterCtorOptions.streamBody}; the * {@link fromSupervisorApi} factory wires it through the SDK's - * `apiClient.request({ raw: true })`. + * `apiClient.request({ raw: true })`, with active W3C context injected by the + * shared serving transport immediately before each request. * * Set `DEBUG=appkit:agents:supervisor-api` to log the outbound request * shape (model, instructions length, input shape, tool count) and to be diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index 6dc11b85d..0443116a8 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -1,5 +1,23 @@ +import { + context, + createTraceState, + propagation, + TraceFlags, + trace, +} from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { W3CTraceContextPropagator } from "@opentelemetry/core"; import type { AgentEvent, AgentToolDefinition, Message } from "shared"; -import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { + afterAll, + afterEach, + beforeAll, + beforeEach, + describe, + expect, + test, + vi, +} from "vitest"; import { consumeAdapterStream } from "../../core/agent/consume-adapter-stream"; import { DatabricksAdapter, @@ -7,6 +25,34 @@ import { parseTextToolCalls, } from "../databricks"; +const TRACE_ID = "0123456789abcdef0123456789abcdef"; +const SPAN_ID = "0123456789abcdef"; +const TRACEPARENT = `00-${TRACE_ID}-${SPAN_ID}-01`; + +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); + propagation.disable(); + propagation.setGlobalPropagator(new W3CTraceContextPropagator()); +}); + +afterAll(() => { + propagation.disable(); + context.disable(); +}); + +function withActiveTrace(operation: () => T): T { + const span = trace.wrapSpanContext({ + traceId: TRACE_ID, + spanId: SPAN_ID, + traceFlags: TraceFlags.SAMPLED, + traceState: createTraceState("vendor=value"), + }); + return context.with(trace.setSpan(context.active(), span), operation); +} + const mockAuthenticate = vi .fn() .mockResolvedValue({ Authorization: "Bearer test-token" }); @@ -936,7 +982,48 @@ describe("DatabricksAdapter", () => { expect(mockAuthenticate).toHaveBeenCalledTimes(1); const [, init] = (globalThis.fetch as any).mock.calls[0]; - expect(init.headers.Authorization).toBe("Bearer test-token"); + expect(new Headers(init.headers).get("authorization")).toBe( + "Bearer test-token", + ); + }); + + test("injects active W3C context after fresh auth on every raw remote-agent request", async () => { + globalThis.fetch = mockFetch([textDelta("Hi"), sseChunk("[DONE]")]); + const adapter = createAdapter(); + + await withActiveTrace(async () => { + for await (const _ of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + // drain + } + }); + + const [, init] = (globalThis.fetch as any).mock.calls[0]; + const headers = new Headers(init.headers); + expect(headers.get("traceparent")).toBe(TRACEPARENT); + expect(headers.get("tracestate")).toBe("vendor=value"); + expect(headers.get("authorization")).toBe("Bearer test-token"); + expect(headers.get("content-type")).toBe("application/json"); + }); + + test("does not add W3C headers to raw requests without a valid active span", async () => { + globalThis.fetch = mockFetch([textDelta("Hi"), sseChunk("[DONE]")]); + const adapter = createAdapter(); + + for await (const _ of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + // drain + } + + const [, init] = (globalThis.fetch as any).mock.calls[0]; + const headers = new Headers(init.headers); + expect(headers.get("traceparent")).toBeNull(); + expect(headers.get("tracestate")).toBeNull(); + expect(headers.get("authorization")).toBe("Bearer test-token"); }); test("throws when two tool names map to the same wire format", async () => { @@ -1813,6 +1900,32 @@ describe("DatabricksAdapter", () => { }); describe("DatabricksAdapter.fromServingEndpoint", () => { + test("propagates the active W3C context through the SDK-backed adapter route", async () => { + const apiClient = { + request: vi.fn().mockResolvedValue({ + contents: createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), + }), + }; + const adapter = await DatabricksAdapter.fromServingEndpoint({ + workspaceClient: { apiClient }, + endpointName: "remote-agent", + }); + + await withActiveTrace(async () => { + for await (const _ of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + // drain + } + }); + + const [requestArgs] = apiClient.request.mock.calls[0]; + const headers = new Headers(requestArgs.headers); + expect(headers.get("traceparent")).toBe(TRACEPARENT); + expect(headers.get("tracestate")).toBe("vendor=value"); + }); + test("routes tool-free chat through apiClient.request with a streaming payload", async () => { const apiClient = { request: vi.fn().mockResolvedValue({ diff --git a/packages/appkit/src/agents/tests/supervisor-api.test.ts b/packages/appkit/src/agents/tests/supervisor-api.test.ts index 7ff57846c..bc92eb79f 100644 --- a/packages/appkit/src/agents/tests/supervisor-api.test.ts +++ b/packages/appkit/src/agents/tests/supervisor-api.test.ts @@ -1,5 +1,22 @@ +import { + context, + createTraceState, + propagation, + TraceFlags, + trace, +} from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { W3CTraceContextPropagator } from "@opentelemetry/core"; import type { AgentEvent, AgentInput } from "shared"; -import { afterEach, describe, expect, test, vi } from "vitest"; +import { + afterAll, + afterEach, + beforeAll, + describe, + expect, + test, + vi, +} from "vitest"; import { fromSupervisorApi, isSupervisorTool, @@ -10,6 +27,34 @@ import { supervisorTools, } from "../supervisor-api"; +const TRACE_ID = "0123456789abcdef0123456789abcdef"; +const SPAN_ID = "0123456789abcdef"; +const TRACEPARENT = `00-${TRACE_ID}-${SPAN_ID}-01`; + +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); + propagation.disable(); + propagation.setGlobalPropagator(new W3CTraceContextPropagator()); +}); + +afterAll(() => { + propagation.disable(); + context.disable(); +}); + +function withActiveTrace(operation: () => T): T { + const span = trace.wrapSpanContext({ + traceId: TRACE_ID, + spanId: SPAN_ID, + traceFlags: TraceFlags.SAMPLED, + traceState: createTraceState("vendor=value"), + }); + return context.with(trace.setSpan(context.active(), span), operation); +} + function createReadableStream(chunks: string[]): ReadableStream { const encoder = new TextEncoder(); let i = 0; @@ -1109,6 +1154,30 @@ describe("fromSupervisorApi", () => { }); expect(requestArgs.payload).not.toHaveProperty("tools"); }); + + test("injects active W3C context into the Supervisor SDK request", async () => { + const request = vi.fn().mockResolvedValue({ + contents: createReadableStream([sseEvent("response.completed", {})]), + }); + const adapter = await fromSupervisorApi({ + model: "databricks-claude-sonnet-4", + workspaceClient: { + config: { ensureResolved: vi.fn(async () => {}) }, + apiClient: { request }, + }, + }); + + await withActiveTrace(() => + collect(adapter.run(createInput(), { executeTool: vi.fn() })), + ); + + const [requestArgs] = request.mock.calls[0]; + const headers = new Headers(requestArgs.headers); + expect(headers.get("traceparent")).toBe(TRACEPARENT); + expect(headers.get("tracestate")).toBe("vendor=value"); + expect(headers.get("content-type")).toBe("application/json"); + expect(headers.get("accept")).toBe("text/event-stream"); + }); }); describe("DatabricksAdapter.fromSupervisorApi", () => { diff --git a/packages/appkit/src/connectors/mcp/client.ts b/packages/appkit/src/connectors/mcp/client.ts index cbdc82d58..65ec723fa 100644 --- a/packages/appkit/src/connectors/mcp/client.ts +++ b/packages/appkit/src/connectors/mcp/client.ts @@ -22,9 +22,16 @@ * inject our host policy and per-URL auth without fighting the default * transport. */ + +import { trace } from "@opentelemetry/api"; import type { AgentToolDefinition } from "shared"; import { APPKIT_USER_AGENT } from "../../context/client-options"; import { createLogger } from "../../logging/logger"; +import { + attachRemoteTraceLink, + injectActiveTraceContext, + type RemoteTraceReference, +} from "../../telemetry/agent-tracing"; import { assertResolvedHostSafe, checkMcpUrl, @@ -370,6 +377,10 @@ export class AppKitMcpClient { callerSignal, }, ); + const activeSpan = trace.getActiveSpan(); + if (activeSpan && rpcResult.remoteTrace) { + attachRemoteTraceLink(activeSpan, rpcResult.remoteTrace); + } const result = rpcResult.result as McpToolCallResult; // `text` is optional on `McpToolCallResult.content[]` per the MCP @@ -412,7 +423,11 @@ export class AppKitMcpClient { */ callerSignal?: AbortSignal; }, - ): Promise<{ result: unknown; sessionId?: string }> { + ): Promise<{ + result: unknown; + sessionId?: string; + remoteTrace?: RemoteTraceReference; + }> { if (this.closed) throw new Error("MCP client is closed"); const request: JsonRpcRequest = { @@ -423,14 +438,14 @@ export class AppKitMcpClient { }; const authHeaders = await this.resolveAuthHeaders(options); - const headers: Record = { + const headers = new Headers({ "User-Agent": APPKIT_USER_AGENT, "Content-Type": "application/json", Accept: "application/json, text/event-stream", ...authHeaders, - }; + }); if (options?.sessionId) { - headers["Mcp-Session-Id"] = options.sessionId; + headers.set("Mcp-Session-Id", options.sessionId); } const fetchImpl = this.options.fetchImpl ?? fetch; @@ -438,7 +453,7 @@ export class AppKitMcpClient { if (options?.callerSignal) signals.push(options.callerSignal); const response = await fetchImpl(url, { method: "POST", - headers, + headers: injectActiveTraceContext(headers), body: JSON.stringify(request), signal: signals.length > 1 ? AbortSignal.any(signals) : signals[0], }); @@ -484,7 +499,25 @@ export class AppKitMcpClient { } const sid = response.headers.get("mcp-session-id") ?? undefined; - return { result: json.result, sessionId: sid }; + const mlflowTraceId = response.headers.get("x-mlflow-trace-id")?.trim(); + const mlflowSpanId = response.headers.get("x-mlflow-span-id")?.trim(); + const otelTraceId = mlflowTraceId?.slice( + mlflowTraceId.lastIndexOf("/") + 1, + ); + const remoteTrace = + mlflowTraceId && mlflowSpanId && otelTraceId + ? { + traceId: mlflowTraceId, + otelTraceId, + spanId: mlflowSpanId, + source: "mcp" as const, + } + : undefined; + return { + result: json.result, + sessionId: sid, + ...(remoteTrace ? { remoteTrace } : {}), + }; } private async sendNotification( @@ -498,14 +531,14 @@ export class AppKitMcpClient { if (this.closed) return; const authHeaders = await this.resolveAuthHeaders(options); - const headers: Record = { + const headers = new Headers({ "User-Agent": APPKIT_USER_AGENT, "Content-Type": "application/json", Accept: "application/json, text/event-stream", ...authHeaders, - }; + }); if (options?.sessionId) { - headers["Mcp-Session-Id"] = options.sessionId; + headers.set("Mcp-Session-Id", options.sessionId); } const fetchImpl = this.options.fetchImpl ?? fetch; @@ -518,7 +551,7 @@ export class AppKitMcpClient { try { const response = await fetchImpl(url, { method: "POST", - headers, + headers: injectActiveTraceContext(headers), body: JSON.stringify({ jsonrpc: "2.0", method }), signal: AbortSignal.timeout(30_000), }); diff --git a/packages/appkit/src/connectors/mcp/tests/client.test.ts b/packages/appkit/src/connectors/mcp/tests/client.test.ts index 1835d90f3..9809fdd26 100644 --- a/packages/appkit/src/connectors/mcp/tests/client.test.ts +++ b/packages/appkit/src/connectors/mcp/tests/client.test.ts @@ -1,9 +1,67 @@ -import { beforeEach, describe, expect, test, vi } from "vitest"; +import { + context, + createTraceState, + propagation, + TraceFlags, + trace, +} from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { W3CTraceContextPropagator } from "@opentelemetry/core"; +import { + BasicTracerProvider, + InMemorySpanExporter, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import { + afterAll, + afterEach, + beforeAll, + beforeEach, + describe, + expect, + test, + vi, +} from "vitest"; import { APPKIT_USER_AGENT } from "../../../context/client-options"; +import { + MlflowUcTraceRegistry, + setActiveMlflowUcTraceRegistry, +} from "../../../telemetry/mlflow-uc"; import { AppKitMcpClient } from "../client"; import type { DnsLookup, McpHostPolicy } from "../host-policy"; const WORKSPACE = "https://test-workspace.cloud.databricks.com"; +const TRACE_ID = "0123456789abcdef0123456789abcdef"; +const SPAN_ID = "0123456789abcdef"; +const TRACEPARENT = `00-${TRACE_ID}-${SPAN_ID}-01`; + +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); + propagation.disable(); + propagation.setGlobalPropagator(new W3CTraceContextPropagator()); +}); + +afterEach(() => { + setActiveMlflowUcTraceRegistry(undefined); +}); + +afterAll(() => { + propagation.disable(); + context.disable(); +}); + +function withActiveTrace(operation: () => T): T { + const span = trace.wrapSpanContext({ + traceId: TRACE_ID, + spanId: SPAN_ID, + traceFlags: TraceFlags.SAMPLED, + traceState: createTraceState("vendor=value"), + }); + return context.with(trace.setSpan(context.active(), span), operation); +} const workspacePolicy: McpHostPolicy = { workspaceHostname: "test-workspace.cloud.databricks.com", @@ -142,10 +200,10 @@ describe("AppKitMcpClient ā€” host allowlist", () => { `${WORKSPACE}/api/2.0/mcp/genie/abc`, ]); for (const call of calls) { - const headers = call.init.headers as Record; - expect(headers.Authorization).toBe("Bearer SP-TOKEN"); + const headers = new Headers(call.init.headers); + expect(headers.get("authorization")).toBe("Bearer SP-TOKEN"); // Every MCP request is attributed to AppKit via User-Agent. - expect(headers["User-Agent"]).toBe(APPKIT_USER_AGENT); + expect(headers.get("user-agent")).toBe(APPKIT_USER_AGENT); } expect(client.canForwardWorkspaceAuth("genie-1")).toBe(true); }); @@ -177,8 +235,8 @@ describe("AppKitMcpClient ā€” host allowlist", () => { await client.connect({ name: "ext", url: "https://mcp.example.com/mcp" }); for (const call of calls) { - const headers = call.init.headers as Record; - expect(headers.Authorization).toBeUndefined(); + const headers = new Headers(call.init.headers); + expect(headers.get("authorization")).toBeNull(); } expect(authSpy).not.toHaveBeenCalled(); expect(client.canForwardWorkspaceAuth("ext")).toBe(false); @@ -308,6 +366,166 @@ describe("AppKitMcpClient ā€” connectAll partial failures", () => { }); describe("AppKitMcpClient ā€” callTool auth scoping", () => { + test("injects one fresh active context into initialize, notification, tools/list, and tools/call", async () => { + const { fetchImpl, calls } = recordingFetch([ + () => + jsonResponse( + { jsonrpc: "2.0", id: 1, result: {} }, + { "mcp-session-id": "sess-1" }, + ), + () => jsonResponse({ jsonrpc: "2.0", result: null }), + () => + jsonResponse({ + jsonrpc: "2.0", + id: 3, + result: { tools: [{ name: "do" }] }, + }), + () => + jsonResponse({ + jsonrpc: "2.0", + id: 4, + result: { content: [{ type: "text", text: "ok" }] }, + }), + ]); + const client = new AppKitMcpClient( + WORKSPACE, + workspaceAuth, + workspacePolicy, + { fetchImpl, dnsLookup: publicDnsLookup }, + ); + + await withActiveTrace(async () => { + await client.connect({ + name: "genie-1", + url: `${WORKSPACE}/api/2.0/mcp/genie/abc`, + }); + await client.callTool( + "mcp.genie-1.do", + {}, + { Authorization: "Bearer OBO-USER-TOKEN" }, + ); + }); + + expect(calls).toHaveLength(4); + expect(new Set(calls.map((call) => call.init.headers)).size).toBe(4); + for (const call of calls) { + const headers = new Headers(call.init.headers); + expect(headers.get("traceparent")).toBe(TRACEPARENT); + expect(headers.get("tracestate")).toBe("vendor=value"); + expect(headers.get("content-type")).toBe("application/json"); + } + expect(new Headers(calls[3].init.headers).get("authorization")).toBe( + "Bearer OBO-USER-TOKEN", + ); + }); + + test("keeps MCP trace headers internal and links only a different MLflow location on the active TOOL span", async () => { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const registry = new MlflowUcTraceRegistry({ + experimentId: "experiment-1", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "appkit", + otelSpansTableName: "main.agent_traces.appkit_otel_spans", + }); + setActiveMlflowUcTraceRegistry(registry); + let toolCall = 0; + const { fetchImpl } = recordingFetch([ + () => + jsonResponse( + { jsonrpc: "2.0", id: 1, result: {} }, + { "mcp-session-id": "sess-1" }, + ), + () => jsonResponse({ jsonrpc: "2.0", result: null }), + () => + jsonResponse({ + jsonrpc: "2.0", + id: 3, + result: { tools: [{ name: "do" }] }, + }), + (call) => { + const traceparent = new Headers(call.init.headers).get("traceparent"); + const otelTraceId = traceparent?.split("-")[1] ?? "missing"; + toolCall++; + return jsonResponse( + { + jsonrpc: "2.0", + id: 3 + toolCall, + result: { + content: [{ type: "text", text: `content-${toolCall}` }], + }, + }, + { + "X-MLflow-Trace-Id": `trace:/${ + toolCall === 1 ? "main.agent_traces.appkit" : "other.remote.agent" + }/${otelTraceId}`, + "X-MLflow-Span-Id": + toolCall === 1 ? "1111111111111111" : "2222222222222222", + }, + ); + }, + (call) => { + const traceparent = new Headers(call.init.headers).get("traceparent"); + const otelTraceId = traceparent?.split("-")[1] ?? "missing"; + toolCall++; + return jsonResponse( + { + jsonrpc: "2.0", + id: 3 + toolCall, + result: { + content: [{ type: "text", text: `content-${toolCall}` }], + }, + }, + { + "X-MLflow-Trace-Id": `trace:/other.remote.agent/${otelTraceId}`, + "X-MLflow-Span-Id": "2222222222222222", + }, + ); + }, + ]); + const client = new AppKitMcpClient( + WORKSPACE, + workspaceAuth, + workspacePolicy, + { fetchImpl, dnsLookup: publicDnsLookup }, + ); + await client.connect({ + name: "genie-1", + url: `${WORKSPACE}/api/2.0/mcp/genie/abc`, + }); + + const outputs = await provider + .getTracer("mcp-test") + .startActiveSpan( + "mcp.genie-1.do tool", + { attributes: { "mlflow.spanType": "TOOL" } }, + async (span) => { + registry.ensureTrace(span.spanContext().traceId); + const first = await client.callTool("mcp.genie-1.do", {}); + const second = await client.callTool("mcp.genie-1.do", {}); + span.end(); + return [first, second]; + }, + ); + + await provider.forceFlush(); + expect(outputs).toEqual(["content-1", "content-2"]); + const tool = exporter.getFinishedSpans()[0]; + expect(tool.links).toHaveLength(1); + expect(tool.links[0].context).toMatchObject({ + traceId: tool.spanContext().traceId, + spanId: "2222222222222222", + }); + expect(tool.links[0].attributes).toEqual({ + "appkit.remote_trace.source": "mcp", + "mlflow.traceRequestId": `trace:/other.remote.agent/${tool.spanContext().traceId}`, + }); + await provider.shutdown(); + }); + test("drops caller-supplied OBO token when destination is not workspace-origin", async () => { const connectResponders = [ () => @@ -353,8 +571,8 @@ describe("AppKitMcpClient ā€” callTool auth scoping", () => { expect(output).toBe("ok"); const toolCall = calls[calls.length - 1]; - const headers = toolCall.init.headers as Record; - expect(headers.Authorization).toBeUndefined(); + const headers = new Headers(toolCall.init.headers); + expect(headers.get("authorization")).toBeNull(); }); test("forwards caller-supplied OBO token when destination is workspace-origin", async () => { @@ -407,8 +625,8 @@ describe("AppKitMcpClient ā€” callTool auth scoping", () => { ); const toolCall = calls[calls.length - 1]; - const headers = toolCall.init.headers as Record; - expect(headers.Authorization).toBe("Bearer OBO-USER-TOKEN"); + const headers = new Headers(toolCall.init.headers); + expect(headers.get("authorization")).toBe("Bearer OBO-USER-TOKEN"); }); test("falls back to SP auth when no OBO override is provided and destination is workspace", async () => { @@ -451,8 +669,8 @@ describe("AppKitMcpClient ā€” callTool auth scoping", () => { await client.callTool("mcp.genie-1.do", {}, undefined); const toolCall = calls[calls.length - 1]; - const headers = toolCall.init.headers as Record; - expect(headers.Authorization).toBe("Bearer SP-TOKEN"); + const headers = new Headers(toolCall.init.headers); + expect(headers.get("authorization")).toBe("Bearer SP-TOKEN"); }); }); diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index 1916b199c..7fea3af85 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -1,4 +1,5 @@ import { createLogger } from "../../logging/logger"; +import { injectActiveTraceContext } from "../../telemetry/agent-tracing"; import type { serving, WorkspaceClient } from "../../workspace-client"; import { contextFromAbortSignal } from "../context"; @@ -116,10 +117,12 @@ export async function streamPath( { path, method: "POST", - headers: new Headers({ - "Content-Type": "application/json", - Accept: "text/event-stream", - }), + headers: injectActiveTraceContext( + new Headers({ + "Content-Type": "application/json", + Accept: "text/event-stream", + }), + ), payload: body, raw: true, }, diff --git a/packages/appkit/src/connectors/serving/tests/client.test.ts b/packages/appkit/src/connectors/serving/tests/client.test.ts index bd75424d7..eb13c13f4 100644 --- a/packages/appkit/src/connectors/serving/tests/client.test.ts +++ b/packages/appkit/src/connectors/serving/tests/client.test.ts @@ -1,7 +1,52 @@ -import { afterEach, describe, expect, test, vi } from "vitest"; +import { + context, + createTraceState, + propagation, + TraceFlags, + trace, +} from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { W3CTraceContextPropagator } from "@opentelemetry/core"; +import { + afterAll, + afterEach, + beforeAll, + describe, + expect, + test, + vi, +} from "vitest"; import { Context } from "../../../workspace-client"; import { getResponseHeaders, invoke, stream } from "../client"; +const TRACE_ID = "0123456789abcdef0123456789abcdef"; +const SPAN_ID = "0123456789abcdef"; +const TRACEPARENT = `00-${TRACE_ID}-${SPAN_ID}-01`; + +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); + propagation.disable(); + propagation.setGlobalPropagator(new W3CTraceContextPropagator()); +}); + +afterAll(() => { + propagation.disable(); + context.disable(); +}); + +function withActiveTrace(operation: () => T): T { + const span = trace.wrapSpanContext({ + traceId: TRACE_ID, + spanId: SPAN_ID, + traceFlags: TraceFlags.SAMPLED, + traceState: createTraceState("vendor=value"), + }); + return context.with(trace.setSpan(context.active(), span), operation); +} + function createMockClient(host = "https://test.databricks.com") { return { config: { host }, @@ -78,6 +123,38 @@ describe("Serving Connector", () => { }); describe("stream", () => { + test("injects the active W3C context into the actual SDK streaming request", async () => { + const client = createMockClient(); + client.apiClient.request.mockResolvedValue({ + contents: new ReadableStream(), + }); + + await withActiveTrace(() => + stream(client, "my-endpoint", { messages: [] }), + ); + + const [request] = client.apiClient.request.mock.calls[0]; + const headers = new Headers(request.headers); + expect(headers.get("traceparent")).toBe(TRACEPARENT); + expect(headers.get("tracestate")).toBe("vendor=value"); + expect(headers.get("content-type")).toBe("application/json"); + expect(headers.get("accept")).toBe("text/event-stream"); + }); + + test("does not add W3C headers when no valid span is active", async () => { + const client = createMockClient(); + client.apiClient.request.mockResolvedValue({ + contents: new ReadableStream(), + }); + + await stream(client, "my-endpoint", { messages: [] }); + + const [request] = client.apiClient.request.mock.calls[0]; + const headers = new Headers(request.headers); + expect(headers.get("traceparent")).toBeNull(); + expect(headers.get("tracestate")).toBeNull(); + }); + test("returns a ReadableStream from apiClient.request", async () => { const encoder = new TextEncoder(); const mockContents = new ReadableStream({ diff --git a/packages/appkit/src/telemetry/agent-tracing/index.ts b/packages/appkit/src/telemetry/agent-tracing/index.ts index c55eb8412..cf4b9477d 100644 --- a/packages/appkit/src/telemetry/agent-tracing/index.ts +++ b/packages/appkit/src/telemetry/agent-tracing/index.ts @@ -1,3 +1,8 @@ +export { + attachRemoteTraceLink, + injectActiveTraceContext, + type RemoteTraceReference, +} from "./propagation"; export { captureTraceValue } from "./serialization"; export { getActiveAgentTraceIdentity, diff --git a/packages/appkit/src/telemetry/agent-tracing/propagation.ts b/packages/appkit/src/telemetry/agent-tracing/propagation.ts new file mode 100644 index 000000000..2e8bac365 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/propagation.ts @@ -0,0 +1,73 @@ +import { + context, + isSpanContextValid, + propagation, + type Span, + TraceFlags, + trace, +} from "@opentelemetry/api"; +import { getMlflowUcTraceId } from "../mlflow-uc"; + +const MLFLOW_V4_TRACE_ID = /^trace:\/[^/]+\/([0-9a-f]{32})$/; + +export interface RemoteTraceReference { + traceId: string; + otelTraceId: string; + spanId: string; + source: "model-serving" | "supervisor" | "mcp" | "remote-agent"; +} + +/** + * Replaces any stale W3C headers with the currently active sampled context. + * The caller owns the Headers instance and must allocate one per request. + */ +export function injectActiveTraceContext(headers: Headers): Headers { + headers.delete("traceparent"); + headers.delete("tracestate"); + + const activeContext = context.active(); + const activeSpan = trace.getSpan(activeContext); + if (!activeSpan || !isSpanContextValid(activeSpan.spanContext())) { + return headers; + } + + const carrier: Record = {}; + propagation.inject(activeContext, carrier); + for (const [key, value] of Object.entries(carrier)) headers.set(key, value); + return headers; +} + +/** + * Links a valid remote MLflow trace when it is not the same UC trace record + * already represented by the local span. + */ +export function attachRemoteTraceLink( + span: Span, + reference: RemoteTraceReference, +): void { + const mlflowMatch = MLFLOW_V4_TRACE_ID.exec(reference.traceId); + const remoteContext = { + traceId: reference.otelTraceId, + spanId: reference.spanId, + traceFlags: TraceFlags.SAMPLED, + isRemote: true, + }; + if ( + !mlflowMatch || + mlflowMatch[1] !== reference.otelTraceId || + !isSpanContextValid(remoteContext) + ) { + return; + } + + const localContext = span.spanContext(); + if (getMlflowUcTraceId(localContext.traceId) === reference.traceId) return; + + span.addLink({ + context: remoteContext, + attributes: { + "mlflow.traceRequestId": reference.traceId, + "appkit.remote_trace.source": reference.source, + }, + }); +} diff --git a/packages/appkit/src/telemetry/agent-tracing/tests/propagation.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/propagation.test.ts new file mode 100644 index 000000000..0de947d88 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/tests/propagation.test.ts @@ -0,0 +1,216 @@ +import { + context, + createTraceState, + propagation, + type Span, + TraceFlags, + trace, +} from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { W3CTraceContextPropagator } from "@opentelemetry/core"; +import { + BasicTracerProvider, + InMemorySpanExporter, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import { afterAll, afterEach, beforeAll, describe, expect, test } from "vitest"; +import { + MlflowUcTraceRegistry, + setActiveMlflowUcTraceRegistry, +} from "../../mlflow-uc"; +import * as agentTracing from "../index"; + +const TRACE_ID = "0123456789abcdef0123456789abcdef"; +const SPAN_ID = "0123456789abcdef"; +const REMOTE_SPAN_ID = "fedcba9876543210"; +const TRACEPARENT = `00-${TRACE_ID}-${SPAN_ID}-01`; +const TRACESTATE = "vendor=value"; + +interface RemoteTraceReference { + traceId: string; + otelTraceId: string; + spanId: string; + source: "model-serving" | "supervisor" | "mcp" | "remote-agent"; +} + +const injectActiveTraceContext = ( + agentTracing as unknown as { + injectActiveTraceContext: (headers: Headers) => Headers; + } +).injectActiveTraceContext; + +const attachRemoteTraceLink = ( + agentTracing as unknown as { + attachRemoteTraceLink: ( + span: Span, + reference: RemoteTraceReference, + ) => void; + } +).attachRemoteTraceLink; + +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); + propagation.disable(); + propagation.setGlobalPropagator(new W3CTraceContextPropagator()); +}); + +afterEach(() => { + setActiveMlflowUcTraceRegistry(undefined); +}); + +afterAll(() => { + propagation.disable(); + context.disable(); +}); + +function withKnownActiveSpan(operation: () => T): T { + const span = trace.wrapSpanContext({ + traceId: TRACE_ID, + spanId: SPAN_ID, + traceFlags: TraceFlags.SAMPLED, + traceState: createTraceState(TRACESTATE), + }); + return context.with(trace.setSpan(context.active(), span), operation); +} + +describe("agent trace propagation", () => { + test("replaces stale W3C headers with the exact active sampled context and preserves other headers", () => { + withKnownActiveSpan(() => { + const headers = new Headers({ + Authorization: "Bearer secret", + "X-AppKit-Request": "request-1", + traceparent: "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-00", + tracestate: "stale=value", + }); + + const result = injectActiveTraceContext(headers); + + expect(result).toBe(headers); + expect(result.get("traceparent")).toBe(TRACEPARENT); + expect(result.get("tracestate")).toBe(TRACESTATE); + expect(result.get("authorization")).toBe("Bearer secret"); + expect(result.get("x-appkit-request")).toBe("request-1"); + }); + }); + + test("removes stale W3C headers when there is no valid active span", () => { + const headers = new Headers({ + Authorization: "Bearer secret", + traceparent: TRACEPARENT, + tracestate: TRACESTATE, + }); + + const result = injectActiveTraceContext(headers); + + expect(result.get("traceparent")).toBeNull(); + expect(result.get("tracestate")).toBeNull(); + expect(result.get("authorization")).toBe("Bearer secret"); + }); + + test("adds one validated cross-location link and skips same-location continuation", async () => { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const registry = new MlflowUcTraceRegistry({ + experimentId: "experiment-1", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "appkit", + otelSpansTableName: "main.agent_traces.appkit_otel_spans", + }); + setActiveMlflowUcTraceRegistry(registry); + + await provider + .getTracer("propagation-test") + .startActiveSpan( + "remote_lookup tool", + { attributes: { "mlflow.spanType": "TOOL" } }, + async (span) => { + const otelTraceId = span.spanContext().traceId; + const currentTraceId = registry.ensureTrace(otelTraceId); + attachRemoteTraceLink(span, { + traceId: currentTraceId, + otelTraceId, + spanId: SPAN_ID, + source: "mcp", + }); + attachRemoteTraceLink(span, { + traceId: `trace:/other.remote.agent/${otelTraceId}`, + otelTraceId, + spanId: REMOTE_SPAN_ID, + source: "mcp", + }); + span.end(); + }, + ); + + await provider.forceFlush(); + const tool = exporter.getFinishedSpans()[0]; + expect(tool.links).toHaveLength(1); + expect(tool.links[0].context).toMatchObject({ + traceId: tool.spanContext().traceId, + spanId: REMOTE_SPAN_ID, + traceFlags: TraceFlags.SAMPLED, + }); + expect(tool.links[0].attributes).toEqual({ + "appkit.remote_trace.source": "mcp", + "mlflow.traceRequestId": `trace:/other.remote.agent/${tool.spanContext().traceId}`, + }); + + await provider.shutdown(); + }); + + test("ignores malformed MLflow, OTel trace, and span IDs", async () => { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + + await provider + .getTracer("propagation-test") + .startActiveSpan( + "remote_lookup tool", + { attributes: { "mlflow.spanType": "TOOL" } }, + async (span) => { + const references: RemoteTraceReference[] = [ + { + traceId: `trace:/other.remote.agent/${TRACE_ID}`, + otelTraceId: "not-a-trace-id", + spanId: REMOTE_SPAN_ID, + source: "mcp", + }, + { + traceId: `trace:/other.remote.agent/${TRACE_ID}`, + otelTraceId: TRACE_ID, + spanId: "not-a-span-id", + source: "mcp", + }, + { + traceId: `trace:/other.remote.agent/${"a".repeat(32)}`, + otelTraceId: TRACE_ID, + spanId: REMOTE_SPAN_ID, + source: "mcp", + }, + { + traceId: `trace:/other.remote.agent/${TRACE_ID.toUpperCase()}`, + otelTraceId: TRACE_ID, + spanId: REMOTE_SPAN_ID, + source: "mcp", + }, + ]; + for (const reference of references) { + attachRemoteTraceLink(span, reference); + } + span.end(); + }, + ); + + await provider.forceFlush(); + expect(exporter.getFinishedSpans()[0].links).toEqual([]); + await provider.shutdown(); + }); +}); From 9eba2f2f5ce8542d39a99406c28f5dcb157c7050 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 21:09:56 -0700 Subject: [PATCH 17/31] fix: inject trace context after SDK auth Signed-off-by: Adam Gurary --- .../appkit/src/connectors/serving/client.ts | 67 +++++++++-- .../connectors/serving/tests/client.test.ts | 107 +++++++++++++++++- 2 files changed, 163 insertions(+), 11 deletions(-) diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index 7fea3af85..6bb2334e2 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -11,7 +11,9 @@ const logger = createLogger("connectors:serving"); * don't want a hard dependency on the concrete `WorkspaceClient` type. */ export interface ApiClientLike { + config?: object; apiClient: { + config?: object; request( options: Record, context?: unknown, @@ -112,17 +114,17 @@ export async function streamPath( logger.debug("Streaming from path %s", path); const context = contextFromAbortSignal(signal); + const headers = new Headers({ + "Content-Type": "application/json", + Accept: "text/event-stream", + }); - const response = (await client.apiClient.request( + const response = (await requestWithPostAuthTraceContext( + client, { path, method: "POST", - headers: injectActiveTraceContext( - new Headers({ - "Content-Type": "application/json", - Accept: "text/event-stream", - }), - ), + headers, payload: body, raw: true, }, @@ -139,6 +141,57 @@ export async function streamPath( return retainResponseHeaders(response.contents, response.headers); } +async function requestWithPostAuthTraceContext( + client: ApiClientLike, + options: Record & { headers: Headers }, + requestContext?: unknown, +): Promise { + const apiClient = client.apiClient; + type AuthenticatingConfig = { + authenticate(headers: Headers): Promise; + }; + const apiConfig = apiClient.config as AuthenticatingConfig | undefined; + const clientConfig = client.config as AuthenticatingConfig | undefined; + const config = + typeof apiConfig?.authenticate === "function" + ? apiConfig + : typeof clientConfig?.authenticate === "function" + ? clientConfig + : undefined; + if (!config) { + return apiClient.request( + { ...options, headers: injectActiveTraceContext(options.headers) }, + requestContext, + ); + } + + // The SDK owns authentication inside `request()`. A request-scoped receiver + // lets that exact code path run unchanged while decorating only its config's + // authenticate step: propagation occurs after fresh credentials resolve and + // before the SDK builds its fetch options. Shared client/config objects are + // never mutated, so concurrent streams cannot exchange trace contexts. + const traceAwareConfig = new Proxy(config, { + get(target, property) { + if (property === "authenticate") { + return async (headers: Headers) => { + await target.authenticate(headers); + injectActiveTraceContext(headers); + }; + } + const value = Reflect.get(target, property, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + const traceAwareApiClient = new Proxy(apiClient, { + get(target, property, receiver) { + if (property === "config") return traceAwareConfig; + return Reflect.get(target, property, receiver); + }, + }); + + return apiClient.request.call(traceAwareApiClient, options, requestContext); +} + /** * Returns the raw SSE byte stream from a serving endpoint. Thin wrapper over * {@link streamPath} that handles serving-specific URL encoding and forces diff --git a/packages/appkit/src/connectors/serving/tests/client.test.ts b/packages/appkit/src/connectors/serving/tests/client.test.ts index eb13c13f4..d0fe00b11 100644 --- a/packages/appkit/src/connectors/serving/tests/client.test.ts +++ b/packages/appkit/src/connectors/serving/tests/client.test.ts @@ -1,3 +1,4 @@ +import http from "node:http"; import { context, createTraceState, @@ -16,12 +17,13 @@ import { test, vi, } from "vitest"; -import { Context } from "../../../workspace-client"; +import { Context, createWorkspaceClient } from "../../../workspace-client"; import { getResponseHeaders, invoke, stream } from "../client"; const TRACE_ID = "0123456789abcdef0123456789abcdef"; const SPAN_ID = "0123456789abcdef"; const TRACEPARENT = `00-${TRACE_ID}-${SPAN_ID}-01`; +const W3C_PROPAGATOR = new W3CTraceContextPropagator(); beforeAll(() => { context.disable(); @@ -29,7 +31,7 @@ beforeAll(() => { new AsyncLocalStorageContextManager().enable(), ); propagation.disable(); - propagation.setGlobalPropagator(new W3CTraceContextPropagator()); + propagation.setGlobalPropagator(W3C_PROPAGATOR); }); afterAll(() => { @@ -37,10 +39,10 @@ afterAll(() => { context.disable(); }); -function withActiveTrace(operation: () => T): T { +function withActiveTrace(operation: () => T, spanId = SPAN_ID): T { const span = trace.wrapSpanContext({ traceId: TRACE_ID, - spanId: SPAN_ID, + spanId, traceFlags: TraceFlags.SAMPLED, traceState: createTraceState("vendor=value"), }); @@ -123,6 +125,103 @@ describe("Serving Connector", () => { }); describe("stream", () => { + test("injects after SDK authentication on every request and preserves final wire headers", async () => { + const secondSpanId = "fedcba9876543210"; + const order: string[] = []; + const wireHeaders: http.IncomingHttpHeaders[] = []; + let authentication = 0; + const server = http.createServer((request, response) => { + order.push(`wire:${wireHeaders.length + 1}`); + wireHeaders.push(request.headers); + request.resume(); + request.on("end", () => { + response.writeHead(200, { "Content-Type": "text/event-stream" }); + response.end("data: {}\n\n"); + }); + }); + await new Promise((resolve) => + server.listen(0, "127.0.0.1", resolve), + ); + const address = server.address(); + if (!address || typeof address === "string") { + throw new Error("Failed to bind SDK wire-test server"); + } + const client = createWorkspaceClient({ + host: `http://127.0.0.1:${address.port}`, + token: "sdk-test-token", + authType: "pat", + }); + const originalAuthenticate = client.config.authenticate.bind( + client.config, + ); + vi.spyOn(client.config, "authenticate").mockImplementation( + async (headers) => { + authentication++; + order.push(`auth:${authentication}`); + await originalAuthenticate(headers); + headers.set("Authorization", `Bearer fresh-${authentication}`); + headers.set( + "traceparent", + "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-00", + ); + headers.set("tracestate", "auth=stale"); + }, + ); + const originalInject = W3C_PROPAGATOR.inject.bind(W3C_PROPAGATOR); + vi.spyOn(W3C_PROPAGATOR, "inject").mockImplementation( + (activeContext, carrier, setter) => { + order.push(`inject:${authentication}`); + originalInject(activeContext, carrier, setter); + }, + ); + + try { + await withActiveTrace(() => + stream(client, "my-endpoint", { messages: [] }), + ); + await withActiveTrace( + () => stream(client, "my-endpoint", { messages: [] }), + secondSpanId, + ); + } finally { + server.closeAllConnections(); + await new Promise((resolve) => server.close(() => resolve())); + } + + expect(order).toEqual([ + "auth:1", + "inject:1", + "wire:1", + "auth:2", + "inject:2", + "wire:2", + ]); + expect( + wireHeaders.map((headers) => ({ + authorization: headers.authorization, + traceparent: headers.traceparent, + tracestate: headers.tracestate, + contentType: headers["content-type"], + accept: headers.accept, + })), + ).toEqual([ + { + authorization: "Bearer fresh-1", + traceparent: TRACEPARENT, + tracestate: "vendor=value", + contentType: "application/json", + accept: "text/event-stream", + }, + { + authorization: "Bearer fresh-2", + traceparent: `00-${TRACE_ID}-${secondSpanId}-01`, + tracestate: "vendor=value", + contentType: "application/json", + accept: "text/event-stream", + }, + ]); + }); + test("injects the active W3C context into the actual SDK streaming request", async () => { const client = createMockClient(); client.apiClient.request.mockResolvedValue({ From dd10b0fb3eb80a1b05c6a33208387cfc7c548c0b Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 22:57:30 -0700 Subject: [PATCH 18/31] feat: provision and surface MLflow UC tracing Signed-off-by: Adam Gurary --- .../client/src/routes/agent.route.tsx | 25 ++ .../src/routes/smart-dashboard.route.tsx | 27 ++ .../tests/agent-tracing.spec.ts | 37 ++ .../smart-dashboard-agent-tracing.spec.ts | 51 +++ .../api/appkit/Class.AgentUsageAccumulator.md | 43 +++ .../api/appkit/Class.SupervisorApiAdapter.md | 3 +- .../api/appkit/Function.captureTraceValue.md | 16 + .../appkit/Interface.AgentModelEndEvent.md | 89 +++++ .../appkit/Interface.AgentModelStartEvent.md | 49 +++ docs/docs/api/appkit/Interface.AgentUsage.md | 57 +++ .../Interface.CaptureTraceValueOptions.md | 17 + .../appkit/Interface.CapturedTraceValue.md | 33 ++ .../api/appkit/Interface.RunAgentInput.md | 40 ++ .../api/appkit/Interface.RunAgentResult.md | 16 + .../api/appkit/Interface.TelemetryConfig.md | 10 + .../appkit/Interface.WorkspaceClientLike.md | 13 + docs/docs/api/appkit/TypeAlias.AgentEvent.md | 11 +- .../appkit/TypeAlias.AgentRemoteTraceEvent.md | 19 + docs/docs/api/appkit/index.md | 8 + docs/docs/api/appkit/typedoc-sidebar.ts | 40 ++ docs/docs/plugins/agents.md | 22 ++ knip.json | 2 +- packages/appkit/package.json | 1 + .../appkit/scripts/provision-mlflow-uc.py | 231 ++++++++++++ .../scripts/test_provision_mlflow_uc.py | 153 ++++++++ packages/appkit/src/plugins/agents/agents.ts | 14 +- .../appkit/src/plugins/agents/manifest.json | 33 +- .../agents/tests/route-handler-errors.test.ts | 16 +- .../plugin/validate/validate-manifest.test.ts | 53 +++ .../src/cli/commands/setup-mlflow-uc.test.ts | 120 ++++++ packages/shared/src/cli/commands/setup.ts | 342 +++++++++++++++++- packages/shared/src/schemas/manifest.ts | 24 +- template/.env.tmpl | 7 + template/app.yaml.tmpl | 16 +- template/appkit.plugins.json | 35 +- .../client/src/pages/agents/AgentChat.tsx | 21 ++ template/config/agents/planner/agent.md | 6 +- template/databricks.yml.tmpl | 29 +- template/package.json | 2 +- template/server/agents/helper.ts | 2 + template/server/server.ts | 5 + tools/dist-appkit.ts | 11 + tools/generate-app-templates.ts | 65 +++- tools/tests/generate-app-templates.test.ts | 104 ++++++ vitest.config.ts | 8 + 45 files changed, 1904 insertions(+), 22 deletions(-) create mode 100644 apps/dev-playground/tests/agent-tracing.spec.ts create mode 100644 apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts create mode 100644 docs/docs/api/appkit/Class.AgentUsageAccumulator.md create mode 100644 docs/docs/api/appkit/Function.captureTraceValue.md create mode 100644 docs/docs/api/appkit/Interface.AgentModelEndEvent.md create mode 100644 docs/docs/api/appkit/Interface.AgentModelStartEvent.md create mode 100644 docs/docs/api/appkit/Interface.AgentUsage.md create mode 100644 docs/docs/api/appkit/Interface.CaptureTraceValueOptions.md create mode 100644 docs/docs/api/appkit/Interface.CapturedTraceValue.md create mode 100644 docs/docs/api/appkit/TypeAlias.AgentRemoteTraceEvent.md create mode 100644 packages/appkit/scripts/provision-mlflow-uc.py create mode 100644 packages/appkit/scripts/test_provision_mlflow_uc.py create mode 100644 packages/shared/src/cli/commands/setup-mlflow-uc.test.ts create mode 100644 tools/tests/generate-app-templates.test.ts diff --git a/apps/dev-playground/client/src/routes/agent.route.tsx b/apps/dev-playground/client/src/routes/agent.route.tsx index 932176665..bffaac5d5 100644 --- a/apps/dev-playground/client/src/routes/agent.route.tsx +++ b/apps/dev-playground/client/src/routes/agent.route.tsx @@ -155,6 +155,8 @@ function AgentRoute() { const [input, setInput] = useState(""); const [isLoading, setIsLoading] = useState(false); const [threadId, setThreadId] = useState(null); + const [mlflowTraceId, setMlflowTraceId] = useState(null); + const [mlflowTraceUrl, setMlflowTraceUrl] = useState(null); const [agent, setAgent] = useState(AGENT_OPTIONS[0].value); const [pendingApprovals, setPendingApprovals] = useState( [], @@ -217,6 +219,8 @@ function AgentRoute() { { id: ++msgIdCounter.current, role: "user", content: userMessage }, ]); setEvents([]); + setMlflowTraceId(null); + setMlflowTraceUrl(null); setIsLoading(true); try { @@ -287,6 +291,12 @@ function AgentRoute() { if (event.type === "appkit.metadata" && event.data?.threadId) { setThreadId(event.data.threadId as string); } + if (event.type === "appkit.metadata") { + const traceId = event.data?.traceId; + const traceUrl = event.data?.traceUrl; + if (typeof traceId === "string") setMlflowTraceId(traceId); + if (typeof traceUrl === "string") setMlflowTraceUrl(traceUrl); + } if (event.type === "response.output_text.delta" && event.delta) { assistantContent += event.delta; @@ -469,6 +479,21 @@ function AgentRoute() {
+ {mlflowTraceId && mlflowTraceUrl && ( +
+ + {mlflowTraceId} + + + Open trace in MLflow + +
+ )} {hasAutocomplete && (suggestion || isAutocompleting) && (
{isAutocompleting && ( diff --git a/apps/dev-playground/client/src/routes/smart-dashboard.route.tsx b/apps/dev-playground/client/src/routes/smart-dashboard.route.tsx index 3817d70f9..fd6d859e2 100644 --- a/apps/dev-playground/client/src/routes/smart-dashboard.route.tsx +++ b/apps/dev-playground/client/src/routes/smart-dashboard.route.tsx @@ -61,6 +61,8 @@ function SmartDashboardRoute() { ); const [lastAction, setLastAction] = useState(null); const [error, setError] = useState(null); + const [mlflowTraceId, setMlflowTraceId] = useState(null); + const [mlflowTraceUrl, setMlflowTraceUrl] = useState(null); // Multi-turn chat history. Messages accumulate across sends so the user // can scroll back through the conversation rather than having the UI @@ -177,6 +179,13 @@ function SmartDashboardRoute() { (event: SSEEvent) => { handleDispatcherEvent(event); + if (event.type === "appkit.metadata") { + const traceId = event.data?.traceId; + const traceUrl = event.data?.traceUrl; + if (typeof traceId === "string") setMlflowTraceId(traceId); + if (typeof traceUrl === "string") setMlflowTraceUrl(traceUrl); + } + // Capture pending approvals and pin them to the user turn that // triggered them so the ChatDrawer can render the card inline. if (event.type === "appkit.approval_pending") { @@ -250,6 +259,8 @@ function SmartDashboardRoute() { const dispatchToAgent = useCallback( (message: string) => { + setMlflowTraceId(null); + setMlflowTraceUrl(null); const userMsgId = nextMessageId(); const assistantMsgId = nextMessageId(); lastUserMessageIdRef.current = userMsgId; @@ -479,6 +490,22 @@ function SmartDashboardRoute() {
+ {mlflowTraceId && mlflowTraceUrl && ( +
+ + {mlflowTraceId} + + + Open trace in MLflow + +
+ )} + {(error || dataError) && (
diff --git a/apps/dev-playground/tests/agent-tracing.spec.ts b/apps/dev-playground/tests/agent-tracing.spec.ts new file mode 100644 index 000000000..5bfed2421 --- /dev/null +++ b/apps/dev-playground/tests/agent-tracing.spec.ts @@ -0,0 +1,37 @@ +import { expect, test } from "@playwright/test"; + +const traceId = `trace:/main.agent_traces.appkit/${"a".repeat(32)}`; +const traceUrl = `https://example.cloud.databricks.com/ml/experiments/123456789/traces?selectedTraceId=${encodeURIComponent(traceId)}`; + +test("agent invocation surfaces its V4 trace identity and direct MLflow link", async ({ + page, +}) => { + await page.route("**/api/agents/chat", async (route) => { + const body = [ + { + type: "appkit.metadata", + data: { threadId: "thread-1", traceId, traceUrl }, + }, + { type: "response.output_text.delta", delta: "Traced answer" }, + { type: "response.completed", response: {} }, + ] + .map((event) => `data: ${JSON.stringify(event)}\n\n`) + .join(""); + await route.fulfill({ + status: 200, + headers: { "Content-Type": "text/event-stream" }, + body, + }); + }); + + await page.goto("/agent"); + await page.getByPlaceholder("Ask a question...").fill("Trace this request"); + await page.getByRole("button", { name: "Send" }).click(); + + await expect( + page.getByRole("paragraph").filter({ hasText: "Traced answer" }), + ).toBeVisible(); + await expect(page.getByText(traceId)).toBeVisible(); + const link = page.getByRole("link", { name: "Open trace in MLflow" }); + await expect(link).toHaveAttribute("href", traceUrl); +}); diff --git a/apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts b/apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts new file mode 100644 index 000000000..11ee28c64 --- /dev/null +++ b/apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts @@ -0,0 +1,51 @@ +import { expect, test } from "@playwright/test"; +import { setupMockAPI } from "./utils/test-utils"; + +const traceId = `trace:/main.agent_traces.appkit/${"b".repeat(32)}`; +const traceUrl = `https://example.cloud.databricks.com/ml/experiments/123456789/traces?selectedTraceId=${encodeURIComponent(traceId)}`; + +test("smart-dashboard planner action produces one linked semantic trace", async ({ + page, +}) => { + await setupMockAPI(page); + await page.route("**/api/agents/chat", async (route) => { + const body = [ + { + type: "appkit.metadata", + data: { threadId: "dashboard-1", traceId, traceUrl }, + }, + { + type: "response.output_item.done", + item: { + type: "function_call", + call_id: "call-1", + name: "filter_by_date_range", + arguments: JSON.stringify({ start: "2016-11-01", end: "2016-11-30" }), + }, + }, + { + type: "response.output_text.delta", + delta: "Applied the November filter.", + }, + { type: "response.completed", response: {} }, + ] + .map((event) => `data: ${JSON.stringify(event)}\n\n`) + .join(""); + await route.fulfill({ + status: 200, + headers: { "Content-Type": "text/event-stream" }, + body, + }); + }); + + await page.goto("/smart-dashboard"); + await page.getByRole("button", { name: "Toggle chat (āŒ˜J)" }).click(); + await page.getByPlaceholder("Ask the dashboardā€¦").fill("Show November 2016"); + await page.getByPlaceholder("Ask the dashboardā€¦").press("Enter"); + + await expect(page.getByText("Applied the November filter.")).toBeVisible(); + const links = page.getByRole("link", { name: "Open trace in MLflow" }); + await expect(links).toHaveCount(1); + await expect(links).toHaveAttribute("href", traceUrl); + await expect(page.getByText(traceId)).toBeVisible(); +}); diff --git a/docs/docs/api/appkit/Class.AgentUsageAccumulator.md b/docs/docs/api/appkit/Class.AgentUsageAccumulator.md new file mode 100644 index 000000000..8ab546468 --- /dev/null +++ b/docs/docs/api/appkit/Class.AgentUsageAccumulator.md @@ -0,0 +1,43 @@ +# Class: AgentUsageAccumulator + +## Constructors + +### Constructor + +```ts +new AgentUsageAccumulator(): AgentUsageAccumulator; +``` + +#### Returns + +`AgentUsageAccumulator` + +## Methods + +### add() + +```ts +add(next: AgentUsage): void; +``` + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `next` | [`AgentUsage`](Interface.AgentUsage.md) | + +#### Returns + +`void` + +*** + +### snapshot() + +```ts +snapshot(): AgentUsage; +``` + +#### Returns + +[`AgentUsage`](Interface.AgentUsage.md) diff --git a/docs/docs/api/appkit/Class.SupervisorApiAdapter.md b/docs/docs/api/appkit/Class.SupervisorApiAdapter.md index 9d10b484c..006c3b77b 100644 --- a/docs/docs/api/appkit/Class.SupervisorApiAdapter.md +++ b/docs/docs/api/appkit/Class.SupervisorApiAdapter.md @@ -11,7 +11,8 @@ Authentication is handled via the Databricks SDK credential chain ā€” the same mechanism used by `DatabricksAdapter.fromModelServing`. The transport is injected via SupervisorApiAdapterCtorOptions.streamBody; the [fromSupervisorApi](Function.fromSupervisorApi.md) factory wires it through the SDK's -`apiClient.request({ raw: true })`. +`apiClient.request({ raw: true })`, with active W3C context injected by the +shared serving transport immediately before each request. Set `DEBUG=appkit:agents:supervisor-api` to log the outbound request shape (model, instructions length, input shape, tool count) and to be diff --git a/docs/docs/api/appkit/Function.captureTraceValue.md b/docs/docs/api/appkit/Function.captureTraceValue.md new file mode 100644 index 000000000..cc2243365 --- /dev/null +++ b/docs/docs/api/appkit/Function.captureTraceValue.md @@ -0,0 +1,16 @@ +# Function: captureTraceValue() + +```ts +function captureTraceValue(value: unknown, options: CaptureTraceValueOptions): CapturedTraceValue; +``` + +## Parameters + +| Parameter | Type | +| ------ | ------ | +| `value` | `unknown` | +| `options` | [`CaptureTraceValueOptions`](Interface.CaptureTraceValueOptions.md) | + +## Returns + +[`CapturedTraceValue`](Interface.CapturedTraceValue.md) diff --git a/docs/docs/api/appkit/Interface.AgentModelEndEvent.md b/docs/docs/api/appkit/Interface.AgentModelEndEvent.md new file mode 100644 index 000000000..d7ab0f29a --- /dev/null +++ b/docs/docs/api/appkit/Interface.AgentModelEndEvent.md @@ -0,0 +1,89 @@ +# Interface: AgentModelEndEvent + +## Properties + +### endedAt + +```ts +endedAt: number; +``` + +*** + +### error? + +```ts +optional error: string; +``` + +*** + +### finishReason? + +```ts +optional finishReason: string; +``` + +*** + +### firstTokenAt? + +```ts +optional firstTokenAt: number; +``` + +*** + +### model + +```ts +model: string; +``` + +*** + +### output + +```ts +output: unknown; +``` + +*** + +### provider + +```ts +provider: string; +``` + +*** + +### stepId + +```ts +stepId: string; +``` + +*** + +### streamDurationMs + +```ts +streamDurationMs: number; +``` + +*** + +### type + +```ts +type: "model_end"; +``` + +*** + +### usage + +```ts +usage: AgentUsage; +``` diff --git a/docs/docs/api/appkit/Interface.AgentModelStartEvent.md b/docs/docs/api/appkit/Interface.AgentModelStartEvent.md new file mode 100644 index 000000000..dca5405c8 --- /dev/null +++ b/docs/docs/api/appkit/Interface.AgentModelStartEvent.md @@ -0,0 +1,49 @@ +# Interface: AgentModelStartEvent + +## Properties + +### input + +```ts +input: unknown; +``` + +*** + +### model + +```ts +model: string; +``` + +*** + +### provider + +```ts +provider: string; +``` + +*** + +### startedAt + +```ts +startedAt: number; +``` + +*** + +### stepId + +```ts +stepId: string; +``` + +*** + +### type + +```ts +type: "model_start"; +``` diff --git a/docs/docs/api/appkit/Interface.AgentUsage.md b/docs/docs/api/appkit/Interface.AgentUsage.md new file mode 100644 index 000000000..0f505cdfc --- /dev/null +++ b/docs/docs/api/appkit/Interface.AgentUsage.md @@ -0,0 +1,57 @@ +# Interface: AgentUsage + +## Properties + +### cacheCreationInputTokens? + +```ts +optional cacheCreationInputTokens: number; +``` + +*** + +### cacheReadInputTokens? + +```ts +optional cacheReadInputTokens: number; +``` + +*** + +### costAvailable + +```ts +costAvailable: boolean; +``` + +*** + +### costUsd? + +```ts +optional costUsd: number; +``` + +*** + +### inputTokens + +```ts +inputTokens: number; +``` + +*** + +### outputTokens + +```ts +outputTokens: number; +``` + +*** + +### totalTokens + +```ts +totalTokens: number; +``` diff --git a/docs/docs/api/appkit/Interface.CaptureTraceValueOptions.md b/docs/docs/api/appkit/Interface.CaptureTraceValueOptions.md new file mode 100644 index 000000000..1cd955ecb --- /dev/null +++ b/docs/docs/api/appkit/Interface.CaptureTraceValueOptions.md @@ -0,0 +1,17 @@ +# Interface: CaptureTraceValueOptions + +## Properties + +### maxBytes? + +```ts +optional maxBytes: number; +``` + +*** + +### redactKeys? + +```ts +optional redactKeys: readonly string[]; +``` diff --git a/docs/docs/api/appkit/Interface.CapturedTraceValue.md b/docs/docs/api/appkit/Interface.CapturedTraceValue.md new file mode 100644 index 000000000..1b3247f9c --- /dev/null +++ b/docs/docs/api/appkit/Interface.CapturedTraceValue.md @@ -0,0 +1,33 @@ +# Interface: CapturedTraceValue + +## Properties + +### originalBytes + +```ts +originalBytes: number; +``` + +*** + +### sha256 + +```ts +sha256: string; +``` + +*** + +### truncated + +```ts +truncated: boolean; +``` + +*** + +### value + +```ts +value: string; +``` diff --git a/docs/docs/api/appkit/Interface.RunAgentInput.md b/docs/docs/api/appkit/Interface.RunAgentInput.md index b17b4a301..45a842d1a 100644 --- a/docs/docs/api/appkit/Interface.RunAgentInput.md +++ b/docs/docs/api/appkit/Interface.RunAgentInput.md @@ -2,6 +2,14 @@ ## Properties +### appName? + +```ts +optional appName: string; +``` + +*** + ### messages ```ts @@ -26,6 +34,22 @@ there is no HTTP request in standalone mode). *** +### requestId? + +```ts +optional requestId: string; +``` + +*** + +### sessionId? + +```ts +optional sessionId: string; +``` + +*** + ### signal? ```ts @@ -33,3 +57,19 @@ optional signal: AbortSignal; ``` Abort signal for cancellation. + +*** + +### threadId? + +```ts +optional threadId: string; +``` + +*** + +### userId? + +```ts +optional userId: string; +``` diff --git a/docs/docs/api/appkit/Interface.RunAgentResult.md b/docs/docs/api/appkit/Interface.RunAgentResult.md index a9ba258dd..ec54d6739 100644 --- a/docs/docs/api/appkit/Interface.RunAgentResult.md +++ b/docs/docs/api/appkit/Interface.RunAgentResult.md @@ -19,3 +19,19 @@ text: string; ``` Aggregated text output from all `message_delta` events. + +*** + +### traceId + +```ts +traceId: string; +``` + +*** + +### usage + +```ts +usage: AgentUsage; +``` diff --git a/docs/docs/api/appkit/Interface.TelemetryConfig.md b/docs/docs/api/appkit/Interface.TelemetryConfig.md index 132d84040..e4dbeb8c0 100644 --- a/docs/docs/api/appkit/Interface.TelemetryConfig.md +++ b/docs/docs/api/appkit/Interface.TelemetryConfig.md @@ -28,6 +28,16 @@ optional instrumentations: Instrumentation[]; *** +### mlflowUc? + +```ts +optional mlflowUc: boolean | Partial; +``` + +Export agent traces to an MLflow experiment backed by Unity Catalog. + +*** + ### serviceName? ```ts diff --git a/docs/docs/api/appkit/Interface.WorkspaceClientLike.md b/docs/docs/api/appkit/Interface.WorkspaceClientLike.md index 3ee3908d3..661411438 100644 --- a/docs/docs/api/appkit/Interface.WorkspaceClientLike.md +++ b/docs/docs/api/appkit/Interface.WorkspaceClientLike.md @@ -18,10 +18,17 @@ the shape they need to satisfy. ```ts apiClient: { + config?: object; request: Promise; }; ``` +#### config? + +```ts +optional config: object; +``` + #### request() ```ts @@ -64,3 +71,9 @@ ensureResolved(): Promise; ##### Returns `Promise`\<`void`\> + +#### Overrides + +```ts +ApiClientLike.config +``` diff --git a/docs/docs/api/appkit/TypeAlias.AgentEvent.md b/docs/docs/api/appkit/TypeAlias.AgentEvent.md index de9226a2b..d695134e8 100644 --- a/docs/docs/api/appkit/TypeAlias.AgentEvent.md +++ b/docs/docs/api/appkit/TypeAlias.AgentEvent.md @@ -42,7 +42,10 @@ type AgentEvent = streamId: string; toolName: string; type: "approval_pending"; -}; +} + | AgentModelStartEvent + | AgentModelEndEvent + | AgentRemoteTraceEvent; ``` ## Type Declaration @@ -268,3 +271,9 @@ is awaiting human approval ā€” fires for tools annotated with legacy `destructive: true` boolean. Clients should render an approval prompt and POST to `/chat/approve` with the matching `approvalId` and a `decision` of `approve` or `deny`. + +[`AgentModelStartEvent`](Interface.AgentModelStartEvent.md) + +[`AgentModelEndEvent`](Interface.AgentModelEndEvent.md) + +[`AgentRemoteTraceEvent`](TypeAlias.AgentRemoteTraceEvent.md) diff --git a/docs/docs/api/appkit/TypeAlias.AgentRemoteTraceEvent.md b/docs/docs/api/appkit/TypeAlias.AgentRemoteTraceEvent.md new file mode 100644 index 000000000..fd9f5c812 --- /dev/null +++ b/docs/docs/api/appkit/TypeAlias.AgentRemoteTraceEvent.md @@ -0,0 +1,19 @@ +# Type Alias: AgentRemoteTraceEvent + +```ts +type AgentRemoteTraceEvent = + | { + relation: "continued"; + source: "model-serving" | "supervisor" | "remote-agent"; + spanId?: string; + traceId: string; + type: "remote_trace"; +} + | { + relation: "linked"; + source: "model-serving" | "supervisor" | "remote-agent"; + spanId: string; + traceId: string; + type: "remote_trace"; +}; +``` diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index f39a52db2..1b70c71cf 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -14,6 +14,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | Class | Description | | ------ | ------ | +| [AgentUsageAccumulator](Class.AgentUsageAccumulator.md) | - | | [AppKitError](Class.AppKitError.md) | Base error class for all AppKit errors. Provides a consistent structure for error handling across the framework. | | [AppKitMcpClient](Class.AppKitMcpClient.md) | Lightweight MCP client for Databricks-hosted MCP servers. | | [AuthenticationError](Class.AuthenticationError.md) | Error thrown when authentication fails. Use for missing tokens, invalid credentials, or authorization failures. | @@ -37,12 +38,17 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [AgentAdapter](Interface.AgentAdapter.md) | - | | [AgentDefinition](Interface.AgentDefinition.md) | - | | [AgentInput](Interface.AgentInput.md) | - | +| [AgentModelEndEvent](Interface.AgentModelEndEvent.md) | - | +| [AgentModelStartEvent](Interface.AgentModelStartEvent.md) | - | | [AgentRunContext](Interface.AgentRunContext.md) | - | | [AgentsPluginConfig](Interface.AgentsPluginConfig.md) | Base configuration interface for AppKit plugins | | [AgentToolDefinition](Interface.AgentToolDefinition.md) | - | +| [AgentUsage](Interface.AgentUsage.md) | - | | [AutoInheritToolsConfig](Interface.AutoInheritToolsConfig.md) | Auto-inherit configuration. When enabled for a given agent origin, agents with no explicit `tools:` declaration receive every registered ToolProvider plugin tool whose author marked `autoInheritable: true`. Tools without that flag ā€” destructive, state-mutating, or privilege-sensitive ā€” never spread automatically and must be wired via `tools:` (object or function form in code, `plugin:NAME` entries in markdown frontmatter). | | [BasePluginConfig](Interface.BasePluginConfig.md) | Base configuration interface for AppKit plugins | | [CacheConfig](Interface.CacheConfig.md) | Configuration for the CacheInterceptor. Controls TTL, size limits, storage backend, and probabilistic cleanup. | +| [CapturedTraceValue](Interface.CapturedTraceValue.md) | - | +| [CaptureTraceValueOptions](Interface.CaptureTraceValueOptions.md) | - | | [DatabaseCredential](Interface.DatabaseCredential.md) | Database credentials with OAuth token for Postgres connection | | [EndpointConfig](Interface.EndpointConfig.md) | - | | [FilePolicyUser](Interface.FilePolicyUser.md) | Minimal user identity passed to the policy function. | @@ -101,6 +107,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | Type Alias | Description | | ------ | ------ | | [AgentEvent](TypeAlias.AgentEvent.md) | - | +| [AgentRemoteTraceEvent](TypeAlias.AgentRemoteTraceEvent.md) | - | | [AgentTool](TypeAlias.AgentTool.md) | Any tool an agent can invoke: inline function tools (`tool()`), hosted MCP tools (`mcpServer()` / raw hosted), toolkit references from plugins (`analytics().toolkit()`), or adapter-hosted Supervisor-API tools (`supervisorTools.*`). | | [AgentTools](TypeAlias.AgentTools.md) | Per-agent tool record. String keys map to inline tools, toolkit entries, hosted tools, etc. | | [AgentToolsFn](TypeAlias.AgentToolsFn.md) | Function form of `AgentDefinition.tools`. Receives the typed [Plugins](TypeAlias.Plugins.md) map and returns a tool record. Invoked exactly once at setup (or once per `runAgent` call in standalone mode); the result is cached as the agent's resolved tool record. | @@ -142,6 +149,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [agentIdFromMarkdownPath](Function.agentIdFromMarkdownPath.md) | Derives the logical agent id from a markdown path. When the file is named `agent.md`, the id is the parent directory name (folder-based layout); otherwise the id is the file stem (e.g. legacy single-file paths). | | [appKitServingTypesPlugin](Function.appKitServingTypesPlugin.md) | Vite plugin to generate TypeScript types for AppKit serving endpoints. Fetches OpenAPI schemas from Databricks and generates a .d.ts with ServingEndpointRegistry module augmentation. | | [appKitTypesPlugin](Function.appKitTypesPlugin.md) | Vite plugin to generate types for AppKit queries. Calls generateFromEntryPoint under the hood. | +| [captureTraceValue](Function.captureTraceValue.md) | - | | [createAgent](Function.createAgent.md) | Pure factory for agent definitions. Returns the passed-in definition after cycle-detecting the sub-agent graph. Accepts the full `AgentDefinition` shape and is safe to call at module top-level. | | [createApp](Function.createApp.md) | Bootstraps AppKit with the provided configuration. | | [createLakebasePool](Function.createLakebasePool.md) | Create a Lakebase pool with appkit's logger integration. Telemetry automatically uses appkit's OpenTelemetry configuration via global registry. | diff --git a/docs/docs/api/appkit/typedoc-sidebar.ts b/docs/docs/api/appkit/typedoc-sidebar.ts index 18a5333b1..4114cae0d 100644 --- a/docs/docs/api/appkit/typedoc-sidebar.ts +++ b/docs/docs/api/appkit/typedoc-sidebar.ts @@ -21,6 +21,11 @@ const typedocSidebar: SidebarsConfig = { type: "category", label: "Classes", items: [ + { + type: "doc", + id: "api/appkit/Class.AgentUsageAccumulator", + label: "AgentUsageAccumulator" + }, { type: "doc", id: "api/appkit/Class.AppKitError", @@ -117,6 +122,16 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.AgentInput", label: "AgentInput" }, + { + type: "doc", + id: "api/appkit/Interface.AgentModelEndEvent", + label: "AgentModelEndEvent" + }, + { + type: "doc", + id: "api/appkit/Interface.AgentModelStartEvent", + label: "AgentModelStartEvent" + }, { type: "doc", id: "api/appkit/Interface.AgentRunContext", @@ -132,6 +147,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.AgentToolDefinition", label: "AgentToolDefinition" }, + { + type: "doc", + id: "api/appkit/Interface.AgentUsage", + label: "AgentUsage" + }, { type: "doc", id: "api/appkit/Interface.AutoInheritToolsConfig", @@ -147,6 +167,16 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.CacheConfig", label: "CacheConfig" }, + { + type: "doc", + id: "api/appkit/Interface.CapturedTraceValue", + label: "CapturedTraceValue" + }, + { + type: "doc", + id: "api/appkit/Interface.CaptureTraceValueOptions", + label: "CaptureTraceValueOptions" + }, { type: "doc", id: "api/appkit/Interface.DatabaseCredential", @@ -418,6 +448,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/TypeAlias.AgentEvent", label: "AgentEvent" }, + { + type: "doc", + id: "api/appkit/TypeAlias.AgentRemoteTraceEvent", + label: "AgentRemoteTraceEvent" + }, { type: "doc", id: "api/appkit/TypeAlias.AgentTool", @@ -585,6 +620,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Function.appKitTypesPlugin", label: "appKitTypesPlugin" }, + { + type: "doc", + id: "api/appkit/Function.captureTraceValue", + label: "captureTraceValue" + }, { type: "doc", id: "api/appkit/Function.createAgent", diff --git a/docs/docs/plugins/agents.md b/docs/docs/plugins/agents.md index d1b7a79e4..e7cca23df 100644 --- a/docs/docs/plugins/agents.md +++ b/docs/docs/plugins/agents.md @@ -20,6 +20,28 @@ For the non-streaming path against a custom endpoint, use the `serving` plugin's Or skip serving-endpoint setup entirely with the managed [Supervisor API adapter](#managed-agents-the-supervisor-api-adapter) (beta). ::: +## First run: provision MLflow tracing in Unity Catalog + +Agent-enabled AppKit templates turn on MLflow UC tracing through AppKit's existing `TelemetryManager`; they do not initialize another OpenTelemetry provider or an MLflow JavaScript SDK. Run setup once after scaffolding and before starting the app: + +```bash +npm run setup -- \ + --mlflow-catalog main \ + --mlflow-schema agent_traces \ + --mlflow-table-prefix appkit \ + --mlflow-warehouse-id 0123456789abcdef +``` + +Selecting the agents plugin implies `--mlflow-uc`. An adjacent AppKit project without agents can opt in explicitly with `--mlflow-uc`. Setup uses `mlflow[databricks]>=3.14.0,<4`, binds `/Users//appkit-agent-traces` with the supported `UnityCatalog` trace-location API, and persists `MLFLOW_EXPERIMENT_ID`, `MLFLOW_TRACING_SQL_WAREHOUSE_ID`, `MLFLOW_UC_CATALOG`, `MLFLOW_UC_SCHEMA`, `MLFLOW_UC_TABLE_PREFIX`, and `MLFLOW_OTEL_SPANS_TABLE` in the local and deployment configuration. + +The binding is immutable. Re-running setup with the same catalog, schema, and prefix is idempotent; attempting to move an existing experiment to another location fails and prints both locations. Choose the production location deliberately. + +The profile running setup needs permission to create or use the experiment and UC tables. Setup applies `USE CATALOG`, `USE SCHEMA`, and both `MODIFY` and `SELECT` on every discovered trace table to that profile's principal through the selected SQL warehouse. The generated app also binds the MLflow experiment with `CAN_MANAGE` and the tracing warehouse with `CAN_USE`. For production, run setup with the app service-principal profile so those UC grants reach the runtime identity. + +Each invocation emits exactly one semantic `AGENT` root, with `CHAT_MODEL` children for model turns, `TOOL` children for tool and sub-agent dispatch, `CHAIN` spans for approval decisions, and `MEMORY` spans for thread-store operations. The root records redacted, size-bounded inputs and outputs, status/error, agent/request/user/session identity, latency, and descendant token usage. Captured values default to 64 KiB and redact normalized authorization, cookie, API-key, token, password, secret, and credential field names as `[REDACTED]`. Cost is omitted unless every child reports a price; an unpriced model is shown as cost unavailable, never as zero. Runtime export failures are logged and do not fail the user request, while invalid or incomplete startup/provisioning configuration is fatal. + +The initial `appkit.metadata` SSE event carries the MLflow V4 trace ID and direct workspace trace URL. The generated chat page displays both after each invocation so the first runnable planner/helper example can be inspected immediately. + ## Install `agents` is a regular plugin. Add it to `plugins[]` alongside `server()` and any ToolProvider plugins whose tools you want agents to reach. diff --git a/knip.json b/knip.json index 0e96b7df5..647cd6781 100644 --- a/knip.json +++ b/knip.json @@ -29,5 +29,5 @@ "docs/**", ".github/scripts/**" ], - "ignoreBinaries": ["tarball", "appkit"] + "ignoreBinaries": ["tarball", "appkit", "uv"] } diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 33b9b61a4..0569bf577 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -56,6 +56,7 @@ "dist": "tsx ../../tools/dist-appkit.ts", "tarball": "rm -rf tmp && pnpm dist && npm pack ./tmp --pack-destination ./tmp", "tarball:prerelease": "rm -rf tmp && SHORTSHA=$(git rev-parse --short HEAD) && pnpm dist --prerelease $SHORTSHA && npm pack ./tmp --pack-destination ./tmp", + "test:mlflow-uc-provision": "uv run --with 'mlflow[databricks]>=3.14.0,<4' pytest scripts/test_provision_mlflow_uc.py -v", "typecheck": "tsc --noEmit" }, "dependencies": { diff --git a/packages/appkit/scripts/provision-mlflow-uc.py b/packages/appkit/scripts/provision-mlflow-uc.py new file mode 100644 index 000000000..2c35c605b --- /dev/null +++ b/packages/appkit/scripts/provision-mlflow-uc.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python3 +"""Provision an immutable MLflow Unity Catalog trace location for AppKit.""" + +from __future__ import annotations + +import argparse +import json +import os +import tempfile +from pathlib import Path +from typing import Any + + +def _quoted_identifier(value: str) -> str: + return f"`{value.replace('`', '``')}`" + + +def _quoted_string(value: str) -> str: + return "'" + value.replace("'", "''") + "'" + + +def _location_name(location: Any) -> str: + if location is None: + return "" + catalog = getattr(location, "catalog_name", None) + schema = getattr(location, "schema_name", None) + prefix = getattr(location, "table_prefix", None) + if all(isinstance(value, str) and value for value in (catalog, schema, prefix)): + return f"{catalog}.{schema}.{prefix}" + return repr(location) + + +def _execute(workspace: Any, warehouse_id: str, statement: str) -> Any: + return workspace.statement_execution.execute_statement( + statement=statement, + warehouse_id=warehouse_id, + wait_timeout="50s", + ) + + +def _discover_trace_tables( + workspace: Any, + warehouse_id: str, + catalog_name: str, + schema_name: str, + table_prefix: str, +) -> list[str]: + response = _execute( + workspace, + warehouse_id, + " ".join( + [ + "SELECT table_name", + f"FROM {_quoted_identifier(catalog_name)}.information_schema.tables", + f"WHERE table_schema = {_quoted_string(schema_name)}", + f"AND table_name LIKE {_quoted_string(f'{table_prefix}%')}", + "ORDER BY table_name", + ] + ), + ) + rows = getattr(getattr(response, "result", None), "data_array", None) or [] + return [str(row[0]) for row in rows if row and row[0] is not None] + + +def _grant_trace_access( + workspace: Any, + warehouse_id: str, + principal: str, + catalog_name: str, + schema_name: str, + table_names: list[str], +) -> None: + catalog = _quoted_identifier(catalog_name) + schema = _quoted_identifier(schema_name) + grantee = _quoted_identifier(principal) + statements = [ + f"GRANT USE CATALOG ON CATALOG {catalog} TO {grantee}", + f"GRANT USE SCHEMA ON SCHEMA {catalog}.{schema} TO {grantee}", + ] + for table_name in table_names: + table = f"{catalog}.{schema}.{_quoted_identifier(table_name)}" + statements.extend( + [ + f"GRANT MODIFY ON TABLE {table} TO {grantee}", + f"GRANT SELECT ON TABLE {table} TO {grantee}", + ] + ) + for statement in statements: + _execute(workspace, warehouse_id, statement) + + +def provision_mlflow_uc( + *, + profile: str, + experiment_name: str, + catalog_name: str, + schema_name: str, + table_prefix: str, + warehouse_id: str, + mlflow_module: Any | None = None, + workspace: Any | None = None, + unity_catalog_type: type | None = None, +) -> dict[str, str]: + if mlflow_module is None: + import mlflow as mlflow_module + if unity_catalog_type is None: + from mlflow.entities.trace_location import UnityCatalog + + unity_catalog_type = UnityCatalog + if workspace is None: + from databricks.sdk import WorkspaceClient + + workspace = WorkspaceClient(profile=profile) + + tracking_uri = f"databricks://{profile}" + os.environ["MLFLOW_TRACKING_URI"] = tracking_uri + os.environ["MLFLOW_TRACING_SQL_WAREHOUSE_ID"] = warehouse_id + mlflow_module.set_tracking_uri(tracking_uri) + + requested_location = unity_catalog_type( + catalog_name=catalog_name, + schema_name=schema_name, + table_prefix=table_prefix, + ) + experiment = mlflow_module.set_experiment( + experiment_name=experiment_name, + trace_location=requested_location, + ) + existing_location = getattr(experiment, "trace_location", None) + if existing_location != requested_location: + raise ValueError( + "MLflow experiment trace location is immutable: " + f"existing={_location_name(existing_location)}, " + f"requested={_location_name(requested_location)}" + ) + + table_names = _discover_trace_tables( + workspace, + warehouse_id, + catalog_name, + schema_name, + table_prefix, + ) + spans_table = f"{table_prefix}_otel_spans" + if spans_table not in table_names: + raise RuntimeError( + f"Required MLflow trace table {catalog_name}.{schema_name}.{spans_table} " + f"was not created; discovered: {', '.join(table_names) or ''}" + ) + + current_user = workspace.current_user.me() + principal = next( + ( + value + for value in ( + getattr(current_user, "user_name", None), + getattr(current_user, "application_id", None), + getattr(current_user, "id", None), + ) + if isinstance(value, str) and value + ), + None, + ) + if principal is None: + raise RuntimeError("Could not resolve the principal receiving UC trace grants") + + _grant_trace_access( + workspace, + warehouse_id, + principal, + catalog_name, + schema_name, + table_names, + ) + + return { + "MLFLOW_EXPERIMENT_ID": str(experiment.experiment_id), + "MLFLOW_TRACING_SQL_WAREHOUSE_ID": warehouse_id, + "MLFLOW_UC_CATALOG": catalog_name, + "MLFLOW_UC_SCHEMA": schema_name, + "MLFLOW_UC_TABLE_PREFIX": table_prefix, + "MLFLOW_OTEL_SPANS_TABLE": f"{catalog_name}.{schema_name}.{spans_table}", + } + + +def write_output_atomically(output_path: Path, values: dict[str, str]) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + dir=output_path.parent, + prefix=f".{output_path.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + json.dump(values, temporary, indent=2, sort_keys=True) + temporary.write("\n") + temporary.flush() + os.fsync(temporary.fileno()) + temporary_path = Path(temporary.name) + os.replace(temporary_path, output_path) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--profile", required=True) + parser.add_argument("--experiment-name", required=True) + parser.add_argument("--catalog", required=True) + parser.add_argument("--schema", required=True) + parser.add_argument("--table-prefix", required=True) + parser.add_argument("--warehouse-id", required=True) + parser.add_argument("--output-json", type=Path, required=True) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + values = provision_mlflow_uc( + profile=args.profile, + experiment_name=args.experiment_name, + catalog_name=args.catalog, + schema_name=args.schema, + table_prefix=args.table_prefix, + warehouse_id=args.warehouse_id, + ) + write_output_atomically(args.output_json, values) + print(json.dumps(values, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/packages/appkit/scripts/test_provision_mlflow_uc.py b/packages/appkit/scripts/test_provision_mlflow_uc.py new file mode 100644 index 000000000..4899ff61e --- /dev/null +++ b/packages/appkit/scripts/test_provision_mlflow_uc.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from mlflow.entities.trace_location import UnityCatalog + + +SCRIPT_PATH = Path(__file__).with_name("provision-mlflow-uc.py") + + +def load_script(): + spec = importlib.util.spec_from_file_location("provision_mlflow_uc", SCRIPT_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class FakeStatementExecution: + def __init__(self, table_names: list[str]): + self.table_names = table_names + self.statements: list[str] = [] + + def execute_statement(self, statement: str, warehouse_id: str, **_kwargs): + assert warehouse_id == "0123456789abcdef" + self.statements.append(statement) + if "information_schema.tables" in statement: + return SimpleNamespace( + result=SimpleNamespace( + data_array=[[table_name] for table_name in self.table_names] + ) + ) + return SimpleNamespace(result=SimpleNamespace(data_array=[])) + + +class FakeWorkspace: + def __init__(self, table_names: list[str]): + self.statement_execution = FakeStatementExecution(table_names) + self.current_user = SimpleNamespace( + me=lambda: SimpleNamespace(user_name="service-principal") + ) + + +def experiment(location: UnityCatalog): + return SimpleNamespace( + experiment_id="123456789", + name="/Users/user@example.com/appkit-agent-traces", + trace_location=location, + ) + + +def provision(module, *, location: UnityCatalog | None = None, tables=None): + requested = UnityCatalog("main", "agent_traces", "appkit") + mlflow_module = SimpleNamespace( + set_tracking_uri=Mock(), + set_experiment=Mock(return_value=experiment(location or requested)), + ) + workspace = FakeWorkspace( + tables + or ["appkit_otel_spans", "appkit_otel_logs", "appkit_annotations"] + ) + result = module.provision_mlflow_uc( + profile="DEFAULT", + experiment_name="/Users/user@example.com/appkit-agent-traces", + catalog_name="main", + schema_name="agent_traces", + table_prefix="appkit", + warehouse_id="0123456789abcdef", + mlflow_module=mlflow_module, + workspace=workspace, + unity_catalog_type=UnityCatalog, + ) + return result, mlflow_module, workspace + + +def test_provisions_supported_uc_location_and_grants_every_discovered_table( + monkeypatch, +): + module = load_script() + monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False) + monkeypatch.delenv("MLFLOW_TRACING_SQL_WAREHOUSE_ID", raising=False) + + result, mlflow_module, workspace = provision(module) + + mlflow_module.set_tracking_uri.assert_called_once_with("databricks://DEFAULT") + mlflow_module.set_experiment.assert_called_once_with( + experiment_name="/Users/user@example.com/appkit-agent-traces", + trace_location=UnityCatalog( + catalog_name="main", + schema_name="agent_traces", + table_prefix="appkit", + ), + ) + assert result == { + "MLFLOW_EXPERIMENT_ID": "123456789", + "MLFLOW_TRACING_SQL_WAREHOUSE_ID": "0123456789abcdef", + "MLFLOW_UC_CATALOG": "main", + "MLFLOW_UC_SCHEMA": "agent_traces", + "MLFLOW_UC_TABLE_PREFIX": "appkit", + "MLFLOW_OTEL_SPANS_TABLE": "main.agent_traces.appkit_otel_spans", + } + assert workspace.statement_execution.statements[1:] == [ + "GRANT USE CATALOG ON CATALOG `main` TO `service-principal`", + "GRANT USE SCHEMA ON SCHEMA `main`.`agent_traces` TO `service-principal`", + "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_otel_spans` TO `service-principal`", + "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_otel_spans` TO `service-principal`", + "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_otel_logs` TO `service-principal`", + "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_otel_logs` TO `service-principal`", + "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_annotations` TO `service-principal`", + "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_annotations` TO `service-principal`", + ] + + +def test_repeated_setup_is_idempotent(): + module = load_script() + + first, _, _ = provision(module) + second, _, _ = provision(module) + + assert second == first + + +def test_existing_different_uc_location_is_rejected_with_both_locations(): + module = load_script() + + with pytest.raises(ValueError) as error: + provision(module, location=UnityCatalog("other", "traces", "legacy")) + + assert "other.traces.legacy" in str(error.value) + assert "main.agent_traces.appkit" in str(error.value) + + +def test_missing_otel_spans_table_is_fatal(): + module = load_script() + + with pytest.raises(RuntimeError, match="appkit_otel_spans"): + provision(module, tables=["appkit_otel_logs", "appkit_annotations"]) + + +def test_writes_configuration_atomically(tmp_path: Path): + module = load_script() + output = tmp_path / ".databricks" / "mlflow-uc.json" + values = {"MLFLOW_EXPERIMENT_ID": "123456789"} + + module.write_output_atomically(output, values) + + assert json.loads(output.read_text()) == values + assert list(output.parent.glob("*.tmp")) == [] diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 0a74febd0..e05df4240 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -1219,9 +1219,14 @@ export class AgentsPlugin extends Plugin implements ToolProvider { // Trace discovery is committed by the outer request handler before any // SSE event. Queue the matching metadata before the adapter can emit. + const traceUrl = buildMlflowTraceUrl(observer.traceId); for (const evt of translator.translate({ type: "metadata", - data: { threadId: thread.id, traceId: observer.traceId }, + data: { + threadId: thread.id, + traceId: observer.traceId, + ...(traceUrl ? { traceUrl } : {}), + }, })) { outboundEvents.push(evt); } @@ -1972,6 +1977,13 @@ function traceIdentityValue(value: unknown): string | undefined { return value.trim() || undefined; } +function buildMlflowTraceUrl(traceId: string): string | undefined { + const workspaceHost = process.env.DATABRICKS_HOST?.trim().replace(/\/$/, ""); + const experimentId = process.env.MLFLOW_EXPERIMENT_ID?.trim(); + if (!workspaceHost || !experimentId) return undefined; + return `${workspaceHost}/ml/experiments/${encodeURIComponent(experimentId)}/traces?selectedTraceId=${encodeURIComponent(traceId)}`; +} + function respondWithTraceError( observer: AgentTraceObserver, res: express.Response, diff --git a/packages/appkit/src/plugins/agents/manifest.json b/packages/appkit/src/plugins/agents/manifest.json index f2e5420bc..0391c5843 100644 --- a/packages/appkit/src/plugins/agents/manifest.json +++ b/packages/appkit/src/plugins/agents/manifest.json @@ -5,7 +5,38 @@ "stability": "beta", "description": "AI agents driven by markdown configs or code, with auto-tool-discovery from registered plugins", "resources": { - "required": [], + "required": [ + { + "type": "experiment", + "alias": "MLflow agent trace experiment", + "resourceKey": "mlflow-experiment", + "description": "MLflow experiment bound to immutable Unity Catalog trace tables", + "permission": "CAN_MANAGE", + "fields": { + "id": { + "env": "MLFLOW_EXPERIMENT_ID", + "description": "MLflow experiment ID" + } + } + }, + { + "type": "sql_warehouse", + "alias": "MLflow tracing SQL warehouse", + "resourceKey": "mlflow-tracing-warehouse", + "description": "SQL warehouse used to provision and query MLflow UC trace tables", + "permission": "CAN_USE", + "fields": { + "id": { + "env": "MLFLOW_TRACING_SQL_WAREHOUSE_ID", + "description": "SQL warehouse ID for MLflow UC tracing", + "discovery": { + "type": "kind", + "resourceKind": "warehouse" + } + } + } + } + ], "optional": [ { "type": "serving_endpoint", diff --git a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts index 1e9a21a56..7fedb6dab 100644 --- a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts +++ b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts @@ -6,7 +6,7 @@ import { SimpleSpanProcessor, } from "@opentelemetry/sdk-trace-base"; import type express from "express"; -import { beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; import { AgentsPlugin } from "../agents"; @@ -38,6 +38,10 @@ beforeEach(() => { }; }); +afterEach(() => { + vi.unstubAllEnvs(); +}); + function mockReq(body: unknown, userId = "alice"): express.Request { const headers: Record = { "x-forwarded-user": userId, @@ -732,6 +736,8 @@ describe("POST /invocations & /responses ā€” successful invoke", () => { describe("POST /chat ā€” trace discovery ordering", () => { test("sets the trace header and emits trace metadata before any streamed content", async () => { + vi.stubEnv("DATABRICKS_HOST", "https://example.cloud.databricks.com/"); + vi.stubEnv("MLFLOW_EXPERIMENT_ID", "123456789"); const plugin = new AgentsPlugin({ dir: false }); const order: string[] = []; const streamed: Array> = []; @@ -820,7 +826,13 @@ describe("POST /chat ā€” trace discovery ordering", () => { expect(order[1]).toBe("body:appkit.metadata"); expect(metadata).toMatchObject({ type: "appkit.metadata", - data: { threadId: "thread-1", traceId: expect.any(String) }, + data: { + threadId: "thread-1", + traceId: expect.any(String), + traceUrl: expect.stringMatching( + /^https:\/\/example\.cloud\.databricks\.com\/ml\/experiments\/123456789\/traces\?selectedTraceId=/, + ), + }, }); expect(setHeader).toHaveBeenCalledWith( "X-MLflow-Trace-Id", diff --git a/packages/shared/src/cli/commands/plugin/validate/validate-manifest.test.ts b/packages/shared/src/cli/commands/plugin/validate/validate-manifest.test.ts index ee448beea..b9df91249 100644 --- a/packages/shared/src/cli/commands/plugin/validate/validate-manifest.test.ts +++ b/packages/shared/src/cli/commands/plugin/validate/validate-manifest.test.ts @@ -1,3 +1,5 @@ +import { readFileSync } from "node:fs"; +import { resolve } from "node:path"; import { describe, expect, it } from "vitest"; import { RESOURCE_KIND_COMMANDS, @@ -47,6 +49,33 @@ const VALID_MANIFEST_WITH_RESOURCE = { }; describe("validate-manifest", () => { + it("requires MLflow experiment and tracing warehouse resources for agents", () => { + const agentsManifest = JSON.parse( + readFileSync( + resolve( + import.meta.dirname, + "../../../../../../appkit/src/plugins/agents/manifest.json", + ), + "utf8", + ), + ); + + const result = validateManifest(agentsManifest); + expect(result.valid).toBe(true); + expect(result.manifest?.resources.required).toEqual([ + expect.objectContaining({ + type: "experiment", + resourceKey: "mlflow-experiment", + permission: "CAN_MANAGE", + }), + expect.objectContaining({ + type: "sql_warehouse", + resourceKey: "mlflow-tracing-warehouse", + permission: "CAN_USE", + }), + ]); + }); + describe("detectSchemaType", () => { it('returns "plugin-manifest" for plugin manifest $schema', () => { expect( @@ -99,6 +128,30 @@ describe("validate-manifest", () => { expect(result.manifest?.resources.required).toHaveLength(1); }); + it("requires experiment resources to expose an id binding", () => { + const result = validateManifest({ + ...VALID_MANIFEST, + resources: { + required: [ + { + type: "experiment", + alias: "MLflow experiment", + resourceKey: "mlflow-experiment", + description: "MLflow trace destination", + permission: "CAN_MANAGE", + fields: { name: { env: "MLFLOW_EXPERIMENT_NAME" } }, + }, + ], + optional: [], + }, + }); + + expect(result.valid).toBe(false); + expect(formatValidationErrors(result.errors ?? [])).toContain( + "resources.required[0].fields.id", + ); + }); + it("rejects non-object input", () => { expect(validateManifest(null).valid).toBe(false); expect(validateManifest("string").valid).toBe(false); diff --git a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts new file mode 100644 index 000000000..c55cabc89 --- /dev/null +++ b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts @@ -0,0 +1,120 @@ +import { mkdirSync, mkdtempSync, readFileSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { describe, expect, test } from "vitest"; +import { + buildMlflowProvisionCommand, + projectRequiresMlflowUc, + provisionAndPersistMlflowUc, +} from "./setup"; + +const EXPECTED_VALUES = { + MLFLOW_EXPERIMENT_ID: "123456789", + MLFLOW_TRACING_SQL_WAREHOUSE_ID: "0123456789abcdef", + MLFLOW_UC_CATALOG: "main", + MLFLOW_UC_SCHEMA: "agent_traces", + MLFLOW_UC_TABLE_PREFIX: "appkit", + MLFLOW_OTEL_SPANS_TABLE: "main.agent_traces.appkit_otel_spans", +}; + +function createProject(): string { + const cwd = mkdtempSync(join(tmpdir(), "appkit-mlflow-uc-")); + mkdirSync(join(cwd, ".databricks"), { recursive: true }); + writeFileSync( + join(cwd, "appkit.plugins.json"), + JSON.stringify({ + plugins: { agents: { requiredByTemplate: true }, serving: {} }, + }), + ); + writeFileSync(join(cwd, ".env"), "DATABRICKS_CONFIG_PROFILE=DEFAULT\n"); + writeFileSync(join(cwd, "app.yaml"), "command: ['npm', 'run', 'start']\n"); + writeFileSync(join(cwd, "databricks.yml"), "bundle:\n name: traced-app\n"); + return cwd; +} + +describe("MLflow UC setup", () => { + test("agent-enabled projects require UC tracing setup", () => { + const cwd = createProject(); + + expect(projectRequiresMlflowUc(cwd, false)).toBe(true); + }); + + test("builds the supported pinned provisioning invocation", () => { + expect( + buildMlflowProvisionCommand({ + cwd: "/workspace/traced-app", + scriptPath: + "/workspace/traced-app/node_modules/@databricks/appkit/scripts/provision-mlflow-uc.py", + profile: "DEFAULT", + experimentName: "/Users/user@example.com/appkit-agent-traces", + catalog: "main", + schema: "agent_traces", + tablePrefix: "appkit", + warehouseId: "0123456789abcdef", + }), + ).toEqual([ + "uv", + "run", + "--no-project", + "--with", + "mlflow[databricks]>=3.14.0,<4", + "python", + "/workspace/traced-app/node_modules/@databricks/appkit/scripts/provision-mlflow-uc.py", + "--profile", + "DEFAULT", + "--experiment-name", + "/Users/user@example.com/appkit-agent-traces", + "--catalog", + "main", + "--schema", + "agent_traces", + "--table-prefix", + "appkit", + "--warehouse-id", + "0123456789abcdef", + "--output-json", + "/workspace/traced-app/.databricks/mlflow-uc.json", + ]); + }); + + test("persists all tracing values for local and deployed runtimes", async () => { + const cwd = createProject(); + const logged: string[] = []; + + const result = await provisionAndPersistMlflowUc( + { + cwd, + profile: "DEFAULT", + experimentName: "/Users/user@example.com/appkit-agent-traces", + catalog: "main", + schema: "agent_traces", + tablePrefix: "appkit", + warehouseId: "0123456789abcdef", + }, + { + scriptPath: join(cwd, "provision-mlflow-uc.py"), + run(command) { + const outputPath = command[command.indexOf("--output-json") + 1]; + writeFileSync(outputPath, JSON.stringify(EXPECTED_VALUES)); + return 0; + }, + log(message) { + logged.push(message); + }, + workspaceHost: "https://example.cloud.databricks.com", + }, + ); + + expect(result).toEqual(EXPECTED_VALUES); + for (const file of [".env", "app.yaml", "databricks.yml"]) { + const content = readFileSync(join(cwd, file), "utf8"); + for (const [name, value] of Object.entries(EXPECTED_VALUES)) { + expect(content).toContain(name); + expect(content).toContain(value); + } + } + expect(logged).toContain( + "MLflow experiment: https://example.cloud.databricks.com/ml/experiments/123456789/traces", + ); + }); +}); diff --git a/packages/shared/src/cli/commands/setup.ts b/packages/shared/src/cli/commands/setup.ts index c72e661eb..77c1e0060 100644 --- a/packages/shared/src/cli/commands/setup.ts +++ b/packages/shared/src/cli/commands/setup.ts @@ -1,6 +1,8 @@ +import { spawnSync } from "node:child_process"; import fs from "node:fs"; import path from "node:path"; import { Command } from "commander"; +import yaml from "js-yaml"; const PACKAGES = [ { name: "@databricks/appkit", description: "Backend SDK" }, @@ -13,6 +15,231 @@ const PACKAGES = [ const SECTION_START = ""; const SECTION_END = ""; +const MLFLOW_UC_ENV_NAMES = [ + "MLFLOW_EXPERIMENT_ID", + "MLFLOW_TRACING_SQL_WAREHOUSE_ID", + "MLFLOW_UC_CATALOG", + "MLFLOW_UC_SCHEMA", + "MLFLOW_UC_TABLE_PREFIX", + "MLFLOW_OTEL_SPANS_TABLE", +] as const; + +export type MlflowUcValues = Record< + (typeof MLFLOW_UC_ENV_NAMES)[number], + string +>; + +export interface MlflowUcSetupOptions { + cwd: string; + profile: string; + experimentName: string; + catalog: string; + schema: string; + tablePrefix: string; + warehouseId: string; +} + +interface MlflowUcSetupDependencies { + scriptPath?: string; + run?: (command: string[]) => number; + log?: (message: string) => void; + workspaceHost?: string; +} + +function readEnvFile(filePath: string): Record { + if (!fs.existsSync(filePath)) return {}; + return Object.fromEntries( + fs + .readFileSync(filePath, "utf8") + .split(/\r?\n/) + .flatMap((line) => { + const match = /^([A-Z][A-Z0-9_]*)=(.*)$/.exec(line.trim()); + return match ? [[match[1], match[2]]] : []; + }), + ); +} + +export function projectRequiresMlflowUc( + cwd: string, + explicitlySelected: boolean, +): boolean { + if (explicitlySelected) return true; + const manifestPath = path.join(cwd, "appkit.plugins.json"); + if (!fs.existsSync(manifestPath)) return false; + try { + const manifest = JSON.parse(fs.readFileSync(manifestPath, "utf8")) as { + plugins?: { agents?: { requiredByTemplate?: boolean } }; + }; + return manifest.plugins?.agents?.requiredByTemplate === true; + } catch (error) { + throw new Error(`Could not read ${manifestPath}: ${String(error)}`); + } +} + +export function buildMlflowProvisionCommand( + options: MlflowUcSetupOptions & { scriptPath: string }, +): string[] { + return [ + "uv", + "run", + "--no-project", + "--with", + "mlflow[databricks]>=3.14.0,<4", + "python", + options.scriptPath, + "--profile", + options.profile, + "--experiment-name", + options.experimentName, + "--catalog", + options.catalog, + "--schema", + options.schema, + "--table-prefix", + options.tablePrefix, + "--warehouse-id", + options.warehouseId, + "--output-json", + path.join(options.cwd, ".databricks", "mlflow-uc.json"), + ]; +} + +function validateMlflowUcValues(value: unknown): MlflowUcValues { + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw new Error("MLflow UC provisioner returned invalid JSON"); + } + const record = value as Record; + const missing = MLFLOW_UC_ENV_NAMES.filter( + (name) => typeof record[name] !== "string" || !record[name], + ); + if (missing.length > 0) { + throw new Error( + `MLflow UC provisioner output missing: ${missing.join(", ")}`, + ); + } + return Object.fromEntries( + MLFLOW_UC_ENV_NAMES.map((name) => [name, record[name] as string]), + ) as MlflowUcValues; +} + +function persistDotEnv(filePath: string, values: MlflowUcValues): void { + const existing = fs.existsSync(filePath) + ? fs.readFileSync(filePath, "utf8").split(/\r?\n/) + : []; + const remaining = existing.filter( + (line) => !MLFLOW_UC_ENV_NAMES.some((name) => line.startsWith(`${name}=`)), + ); + const lines = [ + ...remaining.filter((line, index) => line || index < remaining.length - 1), + ...MLFLOW_UC_ENV_NAMES.map((name) => `${name}=${values[name]}`), + ]; + fs.writeFileSync(filePath, `${lines.join("\n")}\n`); +} + +function loadYamlObject(filePath: string): Record { + if (!fs.existsSync(filePath)) return {}; + const parsed = yaml.load(fs.readFileSync(filePath, "utf8")); + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new Error(`${filePath} must contain a YAML object`); + } + return parsed as Record; +} + +function writeYamlObject( + filePath: string, + document: Record, +): void { + fs.writeFileSync( + filePath, + yaml.dump(document, { lineWidth: 100, noRefs: true, quotingType: '"' }), + ); +} + +function persistAppYaml(filePath: string, values: MlflowUcValues): void { + const document = loadYamlObject(filePath) as { + env?: Array<{ name?: string; value?: string; valueFrom?: string }>; + }; + const existing = Array.isArray(document.env) ? document.env : []; + document.env = [ + ...existing.filter( + (entry) => + !MLFLOW_UC_ENV_NAMES.includes( + entry.name as (typeof MLFLOW_UC_ENV_NAMES)[number], + ), + ), + ...MLFLOW_UC_ENV_NAMES.map((name) => ({ name, value: values[name] })), + ]; + writeYamlObject(filePath, document as Record); +} + +function persistBundleYaml(filePath: string, values: MlflowUcValues): void { + const document = loadYamlObject(filePath) as { + variables?: Record; + targets?: Record }>; + }; + document.variables ??= {}; + for (const name of MLFLOW_UC_ENV_NAMES) { + document.variables[name] = { + description: `AppKit MLflow UC tracing: ${name}`, + default: values[name], + }; + } + document.targets ??= {}; + document.targets.default ??= {}; + document.targets.default.variables ??= {}; + Object.assign(document.targets.default.variables, values); + writeYamlObject(filePath, document as Record); +} + +export async function provisionAndPersistMlflowUc( + options: MlflowUcSetupOptions, + dependencies: MlflowUcSetupDependencies = {}, +): Promise { + const outputPath = path.join(options.cwd, ".databricks", "mlflow-uc.json"); + fs.mkdirSync(path.dirname(outputPath), { recursive: true }); + const scriptPath = + dependencies.scriptPath ?? + path.join( + options.cwd, + "node_modules", + "@databricks", + "appkit", + "scripts", + "provision-mlflow-uc.py", + ); + if (!fs.existsSync(scriptPath) && !dependencies.run) { + throw new Error(`MLflow UC provisioner not found: ${scriptPath}`); + } + const command = buildMlflowProvisionCommand({ ...options, scriptPath }); + const run = + dependencies.run ?? + ((argv: string[]) => + spawnSync(argv[0], argv.slice(1), { + cwd: options.cwd, + stdio: "inherit", + }).status ?? 1); + const status = run(command); + if (status !== 0) { + throw new Error(`MLflow UC provisioning failed with exit code ${status}`); + } + + const values = validateMlflowUcValues( + JSON.parse(fs.readFileSync(outputPath, "utf8")), + ); + persistDotEnv(path.join(options.cwd, ".env"), values); + persistAppYaml(path.join(options.cwd, "app.yaml"), values); + persistBundleYaml(path.join(options.cwd, "databricks.yml"), values); + + const log = dependencies.log ?? console.log; + const host = dependencies.workspaceHost?.replace(/\/$/, ""); + if (host) { + log( + `MLflow experiment: ${host}/ml/experiments/${encodeURIComponent(values.MLFLOW_EXPERIMENT_ID)}/traces`, + ); + } + return values; +} + /** * Find which AppKit packages are installed by checking for package.json */ @@ -115,10 +342,73 @@ function updateContent(existingContent: string, packages: typeof PACKAGES) { return `${existingContent.trimEnd()}\n\n${newSection}\n`; } +interface SetupCliOptions { + write?: boolean; + mlflowUc?: boolean; + mlflowCatalog: string; + mlflowSchema: string; + mlflowTablePrefix: string; + mlflowWarehouseId?: string; +} + +function databricksJson(args: string[]): unknown { + const result = spawnSync("databricks", args, { encoding: "utf8" }); + if (result.status !== 0) { + throw new Error( + result.stderr.trim() || + `databricks ${args.join(" ")} failed with exit code ${result.status}`, + ); + } + return JSON.parse(result.stdout); +} + +function resolveDatabricksUser(profile: string): string { + const user = databricksJson([ + "current-user", + "me", + "--profile", + profile, + "--output", + "json", + ]) as { userName?: unknown; user_name?: unknown }; + const value = user.userName ?? user.user_name; + if (typeof value !== "string" || !value) { + throw new Error( + "Databricks current user response did not include userName", + ); + } + return value; +} + +function resolveWorkspaceHost( + profile: string, + env: Record, +): string { + const configured = process.env.DATABRICKS_HOST ?? env.DATABRICKS_HOST; + if (configured) return configured; + const response = databricksJson(["auth", "profiles", "--output", "json"]); + const profiles = Array.isArray(response) + ? response + : (response as { profiles?: unknown }).profiles; + const selected = Array.isArray(profiles) + ? profiles.find( + (candidate) => + candidate && + typeof candidate === "object" && + (candidate as { name?: unknown }).name === profile, + ) + : undefined; + const host = (selected as { host?: unknown } | undefined)?.host; + if (typeof host !== "string" || !host) { + throw new Error(`Could not resolve workspace host for profile ${profile}`); + } + return host; +} + /** * Setup command implementation */ -function runSetup(options: { write?: boolean }) { +async function runSetup(options: SetupCliOptions) { const shouldWrite = options.write; // Find installed packages @@ -177,16 +467,62 @@ function runSetup(options: { write?: boolean }) { console.log(generateSection(installed)); console.log("ā”€".repeat(50)); } + + const cwd = process.cwd(); + if (projectRequiresMlflowUc(cwd, options.mlflowUc === true)) { + const env = readEnvFile(path.join(cwd, ".env")); + const profile = + process.env.DATABRICKS_CONFIG_PROFILE ?? + env.DATABRICKS_CONFIG_PROFILE ?? + "DEFAULT"; + const warehouseId = + options.mlflowWarehouseId ?? + process.env.MLFLOW_TRACING_SQL_WAREHOUSE_ID ?? + env.MLFLOW_TRACING_SQL_WAREHOUSE_ID ?? + process.env.DATABRICKS_WAREHOUSE_ID ?? + env.DATABRICKS_WAREHOUSE_ID; + if (!warehouseId || warehouseId === "placeholder") { + throw new Error( + "MLflow UC setup requires --mlflow-warehouse-id (or MLFLOW_TRACING_SQL_WAREHOUSE_ID)", + ); + } + const userName = resolveDatabricksUser(profile); + const workspaceHost = resolveWorkspaceHost(profile, env); + await provisionAndPersistMlflowUc( + { + cwd, + profile, + experimentName: `/Users/${userName}/appkit-agent-traces`, + catalog: options.mlflowCatalog, + schema: options.mlflowSchema, + tablePrefix: options.mlflowTablePrefix, + warehouseId, + }, + { workspaceHost }, + ); + } } export const setupCommand = new Command("setup") - .description("Setup CLAUDE.md with AppKit package references") + .description("Set up AppKit project guidance and optional MLflow UC tracing") .option("-w, --write", "Create or update CLAUDE.md file in current directory") + .option( + "--mlflow-uc", + "Provision MLflow tracing in Unity Catalog (implied by agents)", + ) + .option("--mlflow-catalog ", "Unity Catalog catalog", "main") + .option("--mlflow-schema ", "Unity Catalog schema", "agent_traces") + .option("--mlflow-table-prefix ", "UC trace table prefix", "appkit") + .option( + "--mlflow-warehouse-id ", + "SQL warehouse used to provision and query UC trace tables", + ) .addHelpText( "after", ` Examples: $ appkit setup - $ appkit setup --write`, + $ appkit setup --write + $ appkit setup --write --mlflow-uc --mlflow-warehouse-id 0123456789abcdef`, ) .action(runSetup); diff --git a/packages/shared/src/schemas/manifest.ts b/packages/shared/src/schemas/manifest.ts index bf41293f7..1e068a4c3 100644 --- a/packages/shared/src/schemas/manifest.ts +++ b/packages/shared/src/schemas/manifest.ts @@ -484,6 +484,23 @@ function refineResourceDependsOn( } } +/** + * MLflow experiment bindings are keyed by experiment ID in Databricks Apps. + * Keeping that requirement in the schema prevents a manifest from declaring an + * experiment that passes validation but cannot populate MLFLOW_EXPERIMENT_ID. + */ +function refineExperimentId( + resource: { fields?: Record }, + ctx: z.core.$RefinementCtx, +): void { + if (resource.fields?.id?.env) return; + ctx.addIssue({ + code: "custom", + path: ["fields", "id"], + message: "experiment resources must expose an id environment binding", + }); +} + function makeResourceVariant< TType extends z.ZodLiteral, TPerm extends z.ZodTypeAny, @@ -525,7 +542,10 @@ export const resourceRequirementSchema = z makeResourceVariant(z.literal("database"), databasePermissionSchema), makeResourceVariant(z.literal("postgres"), postgresPermissionSchema), makeResourceVariant(z.literal("genie_space"), genieSpacePermissionSchema), - makeResourceVariant(z.literal("experiment"), experimentPermissionSchema), + makeResourceVariant( + z.literal("experiment"), + experimentPermissionSchema, + ).superRefine(refineExperimentId), makeResourceVariant(z.literal("app"), appPermissionSchema), ]) .describe( @@ -896,7 +916,7 @@ export const templateResourceRequirementSchema = z makeTemplateResourceVariant( z.literal("experiment"), experimentPermissionSchema, - ), + ).superRefine(refineExperimentId), makeTemplateResourceVariant(z.literal("app"), appPermissionSchema), ]) .describe( diff --git a/template/.env.tmpl b/template/.env.tmpl index afa7bb15c..8c78a7dc7 100644 --- a/template/.env.tmpl +++ b/template/.env.tmpl @@ -6,6 +6,13 @@ DATABRICKS_HOST={{.workspaceHost}} {{- if .dotEnv.content}} {{.dotEnv.content}} {{- end}} +{{- if .plugins.agents}} +# `npm run setup` replaces these defaults after provisioning the immutable location. +MLFLOW_UC_CATALOG=main +MLFLOW_UC_SCHEMA=agent_traces +MLFLOW_UC_TABLE_PREFIX=appkit +MLFLOW_OTEL_SPANS_TABLE=main.agent_traces.appkit_otel_spans +{{- end}} DATABRICKS_APP_PORT=8000 DATABRICKS_APP_NAME={{.projectName}} FLASK_RUN_HOST=localhost diff --git a/template/app.yaml.tmpl b/template/app.yaml.tmpl index 66311d121..3649eab07 100644 --- a/template/app.yaml.tmpl +++ b/template/app.yaml.tmpl @@ -1,5 +1,19 @@ command: ['npm', 'run', 'start'] -{{- if .appEnv}} +{{- if or .appEnv .plugins.agents}} env: +{{- if .appEnv}} {{.appEnv}} {{- end}} +{{- if .plugins.agents}} + # UC coordinates are runtime configuration, not fields on the experiment + # resource binding. `npm run setup` replaces these defaults after provisioning. + - name: MLFLOW_UC_CATALOG + value: main + - name: MLFLOW_UC_SCHEMA + value: agent_traces + - name: MLFLOW_UC_TABLE_PREFIX + value: appkit + - name: MLFLOW_OTEL_SPANS_TABLE + value: main.agent_traces.appkit_otel_spans +{{- end}} +{{- end}} diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index 078ab524a..f8784afc5 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -8,7 +8,40 @@ "description": "AI agents driven by markdown configs or code, with auto-tool-discovery from registered plugins", "package": "@databricks/appkit", "resources": { - "required": [], + "required": [ + { + "type": "experiment", + "alias": "MLflow agent trace experiment", + "resourceKey": "mlflow-experiment", + "description": "MLflow experiment bound to immutable Unity Catalog trace tables", + "permission": "CAN_MANAGE", + "fields": { + "id": { + "env": "MLFLOW_EXPERIMENT_ID", + "description": "MLflow experiment ID", + "origin": "user" + } + } + }, + { + "type": "sql_warehouse", + "alias": "MLflow tracing SQL warehouse", + "resourceKey": "mlflow-tracing-warehouse", + "description": "SQL warehouse used to provision and query MLflow UC trace tables", + "permission": "CAN_USE", + "fields": { + "id": { + "env": "MLFLOW_TRACING_SQL_WAREHOUSE_ID", + "description": "SQL warehouse ID for MLflow UC tracing", + "discovery": { + "type": "kind", + "resourceKind": "warehouse" + }, + "origin": "user" + } + } + } + ], "optional": [ { "type": "serving_endpoint", diff --git a/template/client/src/pages/agents/AgentChat.tsx b/template/client/src/pages/agents/AgentChat.tsx index 5d1e5758d..8e0da4f71 100644 --- a/template/client/src/pages/agents/AgentChat.tsx +++ b/template/client/src/pages/agents/AgentChat.tsx @@ -61,12 +61,20 @@ export function AgentChat() { const [pendingAssistantId, setPendingAssistantId] = useState( null, ); + const [mlflowTraceId, setMlflowTraceId] = useState(null); + const [mlflowTraceUrl, setMlflowTraceUrl] = useState(null); const scrollRef = useRef(null); // Surface tool-call events as inline messages. Sub-agent delegations // show up here as `agent-helper` tool calls ā€” same wire shape as a // function tool, so the same row renders both. const handleEvent = (event: AgentChatEvent) => { + if (event.type === 'appkit.metadata') { + const traceId = event.data?.traceId; + const traceUrl = event.data?.traceUrl; + if (typeof traceId === 'string') setMlflowTraceId(traceId); + if (typeof traceUrl === 'string') setMlflowTraceUrl(traceUrl); + } if ( event.type === 'response.output_item.added' && event.item?.type === 'function_call' && @@ -112,6 +120,8 @@ export function AgentChat() { if (!message || isStreaming || !activeAgent) return; setInput(''); + setMlflowTraceId(null); + setMlflowTraceUrl(null); const assistantId = `a-${Date.now()}`; setMessages((prev) => [ @@ -186,6 +196,17 @@ export function AgentChat() { })} + {mlflowTraceId && mlflowTraceUrl && ( +
+ + {mlflowTraceId} + + + Open trace in MLflow + +
+ )} +
= { analytics: "SQL warehouse", + agents: "MLflow experiment, SQL warehouse, Serving Endpoint", files: "Volume", genie: "Genie Space", lakebase: "Database", @@ -61,18 +68,33 @@ const FEATURE_DEPENDENCIES: Record = { const APP_TEMPLATES: AppTemplate[] = [ { name: "appkit-all-in-one", - features: ["analytics", "files", "genie", "lakebase", "serving"], + features: ["agents", "analytics", "files", "genie", "lakebase", "serving"], set: { + "agents.mlflow-experiment.id": "placeholder", + "agents.mlflow-tracing-warehouse.id": "placeholder", + "agents.agents-serving-endpoint.name": "placeholder", "analytics.sql-warehouse.id": "placeholder", "files.files.path": "placeholder", "genie.genie-space.id": "placeholder", "genie.genie-space.name": "placeholder", + "lakebase.postgres.project": "placeholder", "lakebase.postgres.branch": "placeholder", "lakebase.postgres.database": "placeholder", "serving.serving-endpoint.name": "placeholder", }, description: - "Full-stack Node.js app with SQL analytics dashboards, file browser, Genie AI conversations, Lakebase Autoscaling (Postgres) CRUD, and Model Serving", + "Full-stack Node.js app with traced agents, SQL analytics dashboards, file browser, Genie AI conversations, Lakebase Autoscaling (Postgres) CRUD, and Model Serving", + }, + { + name: "appkit-agents", + features: ["agents"], + set: { + "agents.mlflow-experiment.id": "placeholder", + "agents.mlflow-tracing-warehouse.id": "placeholder", + "agents.agents-serving-endpoint.name": "placeholder", + }, + description: + "Node.js agent app with MLflow Unity Catalog tracing and a composed planner/helper example", }, { name: "appkit-analytics", @@ -113,6 +135,7 @@ const APP_TEMPLATES: AppTemplate[] = [ name: "appkit-lakebase", features: ["lakebase"], set: { + "lakebase.postgres.project": "placeholder", "lakebase.postgres.branch": "placeholder", "lakebase.postgres.database": "placeholder", }, @@ -122,7 +145,17 @@ const APP_TEMPLATES: AppTemplate[] = [ ]; function run(cmd: string, args: string[], opts?: { cwd?: string }): number { - const result = spawnSync(cmd, args, { stdio: "inherit", cwd: opts?.cwd }); + const result = spawnSync(cmd, args, { + stdio: "inherit", + cwd: opts?.cwd, + env: { + ...process.env, + DATABRICKS_HOST: + process.env.DATABRICKS_HOST ?? + "https://your-workspace.cloud.databricks.com", + DATABRICKS_TOKEN: process.env.DATABRICKS_TOKEN ?? "appkit-template-only", + }, + }); return result.status ?? 1; } @@ -148,6 +181,7 @@ for (const app of APP_TEMPLATES) { app.features.join(","), "--output-dir", OUTPUT_DIR, + "--skip-install", ]; args.push("--description", app.description); @@ -205,11 +239,34 @@ function postProcess(appDir: string, app: AppTemplate): void { ); writeFileSync(join(appDir, "databricks.yml.tmpl"), databricksYmlTmpl); + // Agent templates depend on the tracing runtime added after 0.58.0. Pin the + // first containing release and remove the stale source lock, which cannot be + // regenerated against 0.59.0 until that release is published. + if (app.features.includes("agents")) { + const packageJsonPath = join(appDir, "package.json"); + const packageJson = JSON.parse(readFileSync(packageJsonPath, "utf8")) as { + dependencies: Record; + }; + packageJson.dependencies["@databricks/appkit"] = + FIRST_MLFLOW_UC_APPKIT_VERSION; + packageJson.dependencies["@databricks/appkit-ui"] = + FIRST_MLFLOW_UC_APPKIT_VERSION; + writeFileSync(packageJsonPath, `${JSON.stringify(packageJson, null, 2)}\n`); + rmSync(join(appDir, "package-lock.json"), { force: true }); + } + // 3. Sync appkit.plugins.json based on server imports (discovers available plugins // and marks the ones used in the plugins array as required). const syncStatus = run( "node", - [join(ROOT, "packages/shared/bin/appkit.js"), "plugin", "sync", "--write"], + [ + join(ROOT, "packages/shared/bin/appkit.js"), + "plugin", + "sync", + "--write", + "--plugins-dir", + join(ROOT, "packages/appkit/src/plugins"), + ], { cwd: appDir }, ); if (syncStatus !== 0) { diff --git a/tools/tests/generate-app-templates.test.ts b/tools/tests/generate-app-templates.test.ts new file mode 100644 index 000000000..ea058c3b7 --- /dev/null +++ b/tools/tests/generate-app-templates.test.ts @@ -0,0 +1,104 @@ +import { execFileSync } from "node:child_process"; +import { existsSync, mkdtempSync, readFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import yaml from "js-yaml"; +import { beforeAll, describe, expect, test } from "vitest"; + +const root = resolve(import.meta.dirname, "../.."); +const outputDir = mkdtempSync(join(tmpdir(), "appkit-traced-templates-")); + +describe("generated AppKit agent templates", () => { + beforeAll(() => { + execFileSync("pnpm", ["generate:app-templates"], { + cwd: root, + env: { ...process.env, APP_TEMPLATES_OUTPUT_DIR: outputDir }, + stdio: "pipe", + }); + }, 120_000); + + test.each(["appkit-agents", "appkit-all-in-one"])( + "%s contains the runnable composed agent example", + (name) => { + const app = join(outputDir, name); + expect( + readFileSync(join(app, "config/agents/planner/agent.md"), "utf8"), + ).toContain("single MLflow trace"); + expect( + readFileSync(join(app, "server/agents/helper.ts"), "utf8"), + ).toContain("AGENT ā†’ TOOL ā†’ AGENT ā†’ TOOL"); + expect(readFileSync(join(app, "server/server.ts"), "utf8")).toContain( + "agents({ agents: { helper } })", + ); + }, + ); + + test.each(["appkit-agents", "appkit-all-in-one"])( + "%s persists all UC tracing configuration and resources", + (name) => { + const app = join(outputDir, name); + const appYaml = yaml.load( + readFileSync(join(app, "app.yaml"), "utf8"), + ) as { env: Array<{ name: string; value?: string; valueFrom?: string }> }; + const generated = [ + readFileSync(join(app, "app.yaml"), "utf8"), + readFileSync(join(app, "databricks.yml"), "utf8"), + readFileSync(join(app, ".env.tmpl"), "utf8"), + ].join("\n"); + for (const variable of [ + "MLFLOW_EXPERIMENT_ID", + "MLFLOW_TRACING_SQL_WAREHOUSE_ID", + "MLFLOW_UC_CATALOG", + "MLFLOW_UC_SCHEMA", + "MLFLOW_UC_TABLE_PREFIX", + "MLFLOW_OTEL_SPANS_TABLE", + ]) { + expect(generated).toContain(variable); + } + const manifest = JSON.parse( + readFileSync(join(app, "appkit.plugins.json"), "utf8"), + ); + expect( + manifest.plugins.agents.resources.required.map( + (resource: { resourceKey: string }) => resource.resourceKey, + ), + ).toEqual(["mlflow-experiment", "mlflow-tracing-warehouse"]); + expect(appYaml.env).toEqual( + expect.arrayContaining([ + { name: "MLFLOW_UC_CATALOG", value: "main" }, + { name: "MLFLOW_UC_SCHEMA", value: "agent_traces" }, + { name: "MLFLOW_UC_TABLE_PREFIX", value: "appkit" }, + { + name: "MLFLOW_OTEL_SPANS_TABLE", + value: "main.agent_traces.appkit_otel_spans", + }, + ]), + ); + for (const entry of appYaml.env.filter( + (item) => + item.name.startsWith("MLFLOW_UC_") || + item.name === "MLFLOW_OTEL_SPANS_TABLE", + )) { + expect(entry.valueFrom).toBeUndefined(); + } + const packageJson = JSON.parse( + readFileSync(join(app, "package.json"), "utf8"), + ); + expect(packageJson.scripts.setup).toBe( + "appkit setup --write --mlflow-uc", + ); + expect(packageJson.dependencies["@databricks/appkit"]).toBe("0.59.0"); + expect(packageJson.dependencies["@databricks/appkit-ui"]).toBe("0.59.0"); + expect(existsSync(join(app, "package-lock.json"))).toBe(false); + }, + ); + + test("generated agent UI exposes the direct MLflow trace link", () => { + const chat = readFileSync( + join(outputDir, "appkit-agents/client/src/pages/agents/AgentChat.tsx"), + "utf8", + ); + expect(chat).toContain("Open trace in MLflow"); + expect(chat).toContain("mlflowTraceUrl"); + }); +}); diff --git a/vitest.config.ts b/vitest.config.ts index 8c2893b01..7c210b359 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -58,6 +58,14 @@ export default defineConfig({ environment: "node", }, }, + { + test: { + name: "tools", + root: ".", + include: ["tools/tests/**/*.test.ts"], + environment: "node", + }, + }, ], }, }); From bad5a7b341f2ea5e144e627aeabfbbf7526c7122 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Tue, 11 Aug 2026 23:40:14 -0700 Subject: [PATCH 19/31] fix: address MLflow UC tracing review Signed-off-by: Adam Gurary --- .../client/src/routes/agent.route.tsx | 20 +-- .../src/routes/smart-dashboard.route.tsx | 20 +-- .../smart-dashboard-agent-tracing.fixture.ts | 147 ++++++++++++++++++ .../smart-dashboard-agent-tracing.test.ts | 40 +++++ .../tests/agent-tracing.spec.ts | 31 ++++ .../smart-dashboard-agent-tracing.spec.ts | 61 +++++--- .../appkit/scripts/provision-mlflow-uc.py | 39 ++++- .../scripts/test_provision_mlflow_uc.py | 109 +++++++++++-- .../appkit/src/plugins/agents/manifest.json | 2 +- .../src/cli/commands/plugin/sync/sync.ts | 18 +-- .../src/cli/commands/setup-mlflow-uc.test.ts | 93 ++++++++++- packages/shared/src/cli/commands/setup.ts | 55 ++++++- template/appkit.plugins.json | 2 +- .../client/src/pages/agents/AgentChat.tsx | 10 +- template/server/server.ts | 5 +- tools/tests/generate-app-templates.test.ts | 24 ++- vitest.config.ts | 9 ++ 17 files changed, 591 insertions(+), 94 deletions(-) create mode 100644 apps/dev-playground/server/tests/smart-dashboard-agent-tracing.fixture.ts create mode 100644 apps/dev-playground/server/tests/smart-dashboard-agent-tracing.test.ts diff --git a/apps/dev-playground/client/src/routes/agent.route.tsx b/apps/dev-playground/client/src/routes/agent.route.tsx index bffaac5d5..a996196ae 100644 --- a/apps/dev-playground/client/src/routes/agent.route.tsx +++ b/apps/dev-playground/client/src/routes/agent.route.tsx @@ -479,19 +479,21 @@ function AgentRoute() {
- {mlflowTraceId && mlflowTraceUrl && ( + {mlflowTraceId && (
{mlflowTraceId} - - Open trace in MLflow - + {mlflowTraceUrl && ( + + Open trace in MLflow + + )}
)} {hasAutocomplete && (suggestion || isAutocompleting) && ( diff --git a/apps/dev-playground/client/src/routes/smart-dashboard.route.tsx b/apps/dev-playground/client/src/routes/smart-dashboard.route.tsx index fd6d859e2..f35b8d324 100644 --- a/apps/dev-playground/client/src/routes/smart-dashboard.route.tsx +++ b/apps/dev-playground/client/src/routes/smart-dashboard.route.tsx @@ -490,19 +490,21 @@ function SmartDashboardRoute() {
- {mlflowTraceId && mlflowTraceUrl && ( + {mlflowTraceId && (
{mlflowTraceId} - - Open trace in MLflow - + {mlflowTraceUrl && ( + + Open trace in MLflow + + )}
)} diff --git a/apps/dev-playground/server/tests/smart-dashboard-agent-tracing.fixture.ts b/apps/dev-playground/server/tests/smart-dashboard-agent-tracing.fixture.ts new file mode 100644 index 000000000..8942d1ed8 --- /dev/null +++ b/apps/dev-playground/server/tests/smart-dashboard-agent-tracing.fixture.ts @@ -0,0 +1,147 @@ +import { context, trace } from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { + BasicTracerProvider, + InMemorySpanExporter, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import type { AgentAdapter } from "shared"; +import { z } from "zod"; +import { createAgent } from "../../../../packages/appkit/src/core/agent/create-agent"; +import { runAgent } from "../../../../packages/appkit/src/core/agent/run-agent"; +import { tool } from "../../../../packages/appkit/src/core/agent/tools/tool"; + +export interface SmartDashboardSpanFixture { + spanId: string; + parentSpanId?: string; + traceId: string; + spanType: string; + agentName?: string; + toolName?: string; +} + +export interface SmartDashboardWireEvent { + type: string; + data?: { threadId: string; traceId: string; traceUrl?: string }; + delta?: string; + response?: Record; +} + +export interface SmartDashboardTracingFixture { + traceId: string; + spans: SmartDashboardSpanFixture[]; + events: SmartDashboardWireEvent[]; +} + +export async function runSmartDashboardTracingFixture(options?: { + includeTraceUrl?: boolean; +}): Promise { + context.disable(); + trace.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); + + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + trace.setGlobalTracerProvider(provider); + + try { + const filterByDateRange = tool({ + name: "filter_by_date_range", + description: "Filter the dashboard to a date range", + schema: z.object({ start: z.string(), end: z.string() }), + execute: async ({ start, end }) => `Filtered ${start} through ${end}.`, + }); + const pilotAdapter: AgentAdapter = { + async *run(_input, runContext) { + const result = await runContext.executeTool("filter_by_date_range", { + start: "2016-11-01", + end: "2016-11-30", + }); + yield { type: "message_delta", content: String(result) }; + }, + }; + const queryAdapter: AgentAdapter = { + async *run(_input, runContext) { + const result = await runContext.executeTool("agent-dashboard_pilot", { + input: "Show November 2016", + }); + yield { type: "message_delta", content: String(result) }; + }, + }; + const query = createAgent({ + name: "query", + instructions: "Delegate dashboard changes to the dashboard pilot.", + model: queryAdapter, + agents: { + dashboard_pilot: createAgent({ + name: "dashboard_pilot", + instructions: "Manipulate the Smart Dashboard.", + model: pilotAdapter, + tools: { filter_by_date_range: filterByDateRange }, + }), + }, + }); + + const result = await runAgent(query, { + appName: "dev-playground", + messages: "Show November 2016", + requestId: "smart-dashboard-request", + sessionId: "smart-dashboard-session", + threadId: "smart-dashboard-thread", + userId: "fixture-user", + }); + await provider.forceFlush(); + + const spans = exporter + .getFinishedSpans() + .filter((span) => { + const spanType = span.attributes["mlflow.spanType"]; + return spanType === "AGENT" || spanType === "TOOL"; + }) + .map((span) => ({ + spanId: span.spanContext().spanId, + ...(span.parentSpanContext?.spanId + ? { parentSpanId: span.parentSpanContext.spanId } + : {}), + traceId: span.spanContext().traceId, + spanType: String(span.attributes["mlflow.spanType"]), + ...(typeof span.attributes["appkit.agent.name"] === "string" + ? { agentName: span.attributes["appkit.agent.name"] } + : {}), + ...(typeof span.attributes["appkit.tool.name"] === "string" + ? { toolName: span.attributes["appkit.tool.name"] } + : {}), + })); + const traceUrl = options?.includeTraceUrl + ? `https://example.cloud.databricks.com/ml/experiments/123456789/traces?selectedTraceId=${encodeURIComponent(result.traceId)}` + : undefined; + + return { + traceId: result.traceId, + spans, + events: [ + { + type: "appkit.metadata", + data: { + threadId: "smart-dashboard-thread", + traceId: result.traceId, + ...(traceUrl ? { traceUrl } : {}), + }, + }, + { + type: "response.output_text.delta", + delta: "Applied the November filter.", + }, + { type: "response.completed", response: {} }, + ], + }; + } finally { + await provider.shutdown(); + trace.disable(); + context.disable(); + } +} diff --git a/apps/dev-playground/server/tests/smart-dashboard-agent-tracing.test.ts b/apps/dev-playground/server/tests/smart-dashboard-agent-tracing.test.ts new file mode 100644 index 000000000..e8737d69a --- /dev/null +++ b/apps/dev-playground/server/tests/smart-dashboard-agent-tracing.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, test } from "vitest"; +import { runSmartDashboardTracingFixture } from "./smart-dashboard-agent-tracing.fixture"; + +describe("Smart Dashboard semantic tracing fixture", () => { + test("emits the exact query ā†’ delegation ā†’ pilot ā†’ action tree", async () => { + const observed = await runSmartDashboardTracingFixture(); + const root = observed.spans.find( + (span) => span.spanType === "AGENT" && span.agentName === "query", + ); + const delegation = observed.spans.find( + (span) => + span.spanType === "TOOL" && span.toolName === "agent-dashboard_pilot", + ); + const pilot = observed.spans.find( + (span) => + span.spanType === "AGENT" && span.agentName === "dashboard_pilot", + ); + const action = observed.spans.find( + (span) => + span.spanType === "TOOL" && span.toolName === "filter_by_date_range", + ); + + expect( + observed.spans.filter( + (span) => span.spanType === "AGENT" && span.parentSpanId === undefined, + ), + ).toHaveLength(1); + expect(root).toBeDefined(); + expect(delegation?.parentSpanId).toBe(root?.spanId); + expect(pilot?.parentSpanId).toBe(delegation?.spanId); + expect(action?.parentSpanId).toBe(pilot?.spanId); + expect(new Set(observed.spans.map((span) => span.traceId))).toEqual( + new Set([observed.traceId]), + ); + expect(observed.events[0]).toMatchObject({ + type: "appkit.metadata", + data: { traceId: observed.traceId }, + }); + }); +}); diff --git a/apps/dev-playground/tests/agent-tracing.spec.ts b/apps/dev-playground/tests/agent-tracing.spec.ts index 5bfed2421..fad2a78e7 100644 --- a/apps/dev-playground/tests/agent-tracing.spec.ts +++ b/apps/dev-playground/tests/agent-tracing.spec.ts @@ -35,3 +35,34 @@ test("agent invocation surfaces its V4 trace identity and direct MLflow link", a const link = page.getByRole("link", { name: "Open trace in MLflow" }); await expect(link).toHaveAttribute("href", traceUrl); }); + +test("agent invocation surfaces its trace ID without a workspace link", async ({ + page, +}) => { + await page.route("**/api/agents/chat", async (route) => { + const body = [ + { + type: "appkit.metadata", + data: { threadId: "thread-2", traceId }, + }, + { type: "response.output_text.delta", delta: "Unlinked trace" }, + { type: "response.completed", response: {} }, + ] + .map((event) => `data: ${JSON.stringify(event)}\n\n`) + .join(""); + await route.fulfill({ + status: 200, + headers: { "Content-Type": "text/event-stream" }, + body, + }); + }); + + await page.goto("/agent"); + await page.getByPlaceholder("Ask a question...").fill("Trace without a URL"); + await page.getByRole("button", { name: "Send" }).click(); + + await expect(page.getByText(traceId)).toBeVisible(); + await expect( + page.getByRole("link", { name: "Open trace in MLflow" }), + ).toHaveCount(0); +}); diff --git a/apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts b/apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts index 11ee28c64..12947e409 100644 --- a/apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts +++ b/apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts @@ -1,34 +1,17 @@ import { expect, test } from "@playwright/test"; +import { runSmartDashboardTracingFixture } from "../server/tests/smart-dashboard-agent-tracing.fixture"; import { setupMockAPI } from "./utils/test-utils"; -const traceId = `trace:/main.agent_traces.appkit/${"b".repeat(32)}`; -const traceUrl = `https://example.cloud.databricks.com/ml/experiments/123456789/traces?selectedTraceId=${encodeURIComponent(traceId)}`; - test("smart-dashboard planner action produces one linked semantic trace", async ({ page, }) => { + const observed = await runSmartDashboardTracingFixture({ + includeTraceUrl: true, + }); + const traceUrl = observed.events[0].data?.traceUrl; await setupMockAPI(page); await page.route("**/api/agents/chat", async (route) => { - const body = [ - { - type: "appkit.metadata", - data: { threadId: "dashboard-1", traceId, traceUrl }, - }, - { - type: "response.output_item.done", - item: { - type: "function_call", - call_id: "call-1", - name: "filter_by_date_range", - arguments: JSON.stringify({ start: "2016-11-01", end: "2016-11-30" }), - }, - }, - { - type: "response.output_text.delta", - delta: "Applied the November filter.", - }, - { type: "response.completed", response: {} }, - ] + const body = observed.events .map((event) => `data: ${JSON.stringify(event)}\n\n`) .join(""); await route.fulfill({ @@ -46,6 +29,34 @@ test("smart-dashboard planner action produces one linked semantic trace", async await expect(page.getByText("Applied the November filter.")).toBeVisible(); const links = page.getByRole("link", { name: "Open trace in MLflow" }); await expect(links).toHaveCount(1); - await expect(links).toHaveAttribute("href", traceUrl); - await expect(page.getByText(traceId)).toBeVisible(); + expect(traceUrl).toBeDefined(); + await expect(links).toHaveAttribute("href", traceUrl as string); + await expect(page.getByText(observed.traceId)).toBeVisible(); +}); + +test("smart-dashboard surfaces its trace ID without a workspace link", async ({ + page, +}) => { + const observed = await runSmartDashboardTracingFixture(); + await setupMockAPI(page); + await page.route("**/api/agents/chat", async (route) => { + const body = observed.events + .map((event) => `data: ${JSON.stringify(event)}\n\n`) + .join(""); + await route.fulfill({ + status: 200, + headers: { "Content-Type": "text/event-stream" }, + body, + }); + }); + + await page.goto("/smart-dashboard"); + await page.getByRole("button", { name: "Toggle chat (āŒ˜J)" }).click(); + await page.getByPlaceholder("Ask the dashboardā€¦").fill("Trace without a URL"); + await page.getByPlaceholder("Ask the dashboardā€¦").press("Enter"); + + await expect(page.getByText(observed.traceId)).toBeVisible(); + await expect( + page.getByRole("link", { name: "Open trace in MLflow" }), + ).toHaveCount(0); }); diff --git a/packages/appkit/scripts/provision-mlflow-uc.py b/packages/appkit/scripts/provision-mlflow-uc.py index 2c35c605b..f0f969668 100644 --- a/packages/appkit/scripts/provision-mlflow-uc.py +++ b/packages/appkit/scripts/provision-mlflow-uc.py @@ -7,6 +7,7 @@ import json import os import tempfile +import time from pathlib import Path from typing import Any @@ -31,11 +32,39 @@ def _location_name(location: Any) -> str: def _execute(workspace: Any, warehouse_id: str, statement: str) -> Any: - return workspace.statement_execution.execute_statement( + response = workspace.statement_execution.execute_statement( statement=statement, warehouse_id=warehouse_id, wait_timeout="50s", ) + for _ in range(120): + status = getattr(response, "status", None) + state = getattr(status, "state", None) + state_value = getattr(state, "value", state) + if state_value == "SUCCEEDED": + return response + if state_value in {"FAILED", "CANCELED", "CLOSED"}: + error = getattr(status, "error", None) + code = getattr(error, "error_code", None) + message = getattr(error, "message", None) + detail = ": ".join(str(value) for value in (code, message) if value) + raise RuntimeError( + f"SQL statement {state_value.lower()}: {detail or statement}" + ) + if state_value not in {"PENDING", "RUNNING"}: + raise RuntimeError( + f"SQL statement returned unknown status {state_value!r}: {statement}" + ) + statement_id = getattr(response, "statement_id", None) + if not isinstance(statement_id, str) or not statement_id: + raise RuntimeError( + f"SQL statement is {state_value.lower()} without a statement ID" + ) + response = workspace.statement_execution.get_statement(statement_id) + next_state = getattr(getattr(response, "status", None), "state", None) + if getattr(next_state, "value", next_state) in {"PENDING", "RUNNING"}: + time.sleep(1) + raise TimeoutError(f"SQL statement did not finish after 120 polls: {statement}") def _discover_trace_tables( @@ -59,7 +88,13 @@ def _discover_trace_tables( ), ) rows = getattr(getattr(response, "result", None), "data_array", None) or [] - return [str(row[0]) for row in rows if row and row[0] is not None] + return [ + str(row[0]) + for row in rows + if row + and row[0] is not None + and str(row[0]).startswith(table_prefix) + ] def _grant_trace_access( diff --git a/packages/appkit/scripts/test_provision_mlflow_uc.py b/packages/appkit/scripts/test_provision_mlflow_uc.py index 4899ff61e..11ddf1143 100644 --- a/packages/appkit/scripts/test_provision_mlflow_uc.py +++ b/packages/appkit/scripts/test_provision_mlflow_uc.py @@ -21,26 +21,65 @@ def load_script(): return module +def statement_response( + state: str, + *, + table_names: list[str] | None = None, + statement_id: str = "statement-1", + error_message: str | None = None, +): + error = ( + SimpleNamespace(message=error_message, error_code="PERMISSION_DENIED") + if error_message + else None + ) + return SimpleNamespace( + statement_id=statement_id, + status=SimpleNamespace(state=state, error=error), + result=SimpleNamespace( + data_array=[[table_name] for table_name in (table_names or [])] + ), + ) + + class FakeStatementExecution: - def __init__(self, table_names: list[str]): + def __init__( + self, + table_names: list[str], + *, + failed_statement: str | None = None, + pending_then_succeeded: bool = False, + ): self.table_names = table_names self.statements: list[str] = [] + self.failed_statement = failed_statement + self.pending_then_succeeded = pending_then_succeeded + self.get_statement_calls: list[str] = [] def execute_statement(self, statement: str, warehouse_id: str, **_kwargs): assert warehouse_id == "0123456789abcdef" self.statements.append(statement) - if "information_schema.tables" in statement: - return SimpleNamespace( - result=SimpleNamespace( - data_array=[[table_name] for table_name in self.table_names] - ) + if self.failed_statement and self.failed_statement in statement: + return statement_response( + "FAILED", + error_message="principal lacks MODIFY", ) - return SimpleNamespace(result=SimpleNamespace(data_array=[])) + if self.pending_then_succeeded: + return statement_response("PENDING") + if "information_schema.tables" in statement: + return statement_response("SUCCEEDED", table_names=self.table_names) + return statement_response("SUCCEEDED") + + def get_statement(self, statement_id: str): + self.get_statement_calls.append(statement_id) + return statement_response("SUCCEEDED", table_names=self.table_names) class FakeWorkspace: - def __init__(self, table_names: list[str]): - self.statement_execution = FakeStatementExecution(table_names) + def __init__(self, table_names: list[str], **statement_options): + self.statement_execution = FakeStatementExecution( + table_names, **statement_options + ) self.current_user = SimpleNamespace( me=lambda: SimpleNamespace(user_name="service-principal") ) @@ -54,22 +93,30 @@ def experiment(location: UnityCatalog): ) -def provision(module, *, location: UnityCatalog | None = None, tables=None): - requested = UnityCatalog("main", "agent_traces", "appkit") +def provision( + module, + *, + location: UnityCatalog | None = None, + tables=None, + statement_options=None, + table_prefix="appkit", +): + requested = UnityCatalog("main", "agent_traces", table_prefix) mlflow_module = SimpleNamespace( set_tracking_uri=Mock(), set_experiment=Mock(return_value=experiment(location or requested)), ) workspace = FakeWorkspace( tables - or ["appkit_otel_spans", "appkit_otel_logs", "appkit_annotations"] + or ["appkit_otel_spans", "appkit_otel_logs", "appkit_annotations"], + **(statement_options or {}), ) result = module.provision_mlflow_uc( profile="DEFAULT", experiment_name="/Users/user@example.com/appkit-agent-traces", catalog_name="main", schema_name="agent_traces", - table_prefix="appkit", + table_prefix=table_prefix, warehouse_id="0123456789abcdef", mlflow_module=mlflow_module, workspace=workspace, @@ -142,6 +189,42 @@ def test_missing_otel_spans_table_is_fatal(): provision(module, tables=["appkit_otel_logs", "appkit_annotations"]) +def test_wildcards_in_prefix_cannot_grant_unrelated_tables(): + module = load_script() + + _, _, workspace = provision( + module, + table_prefix="app%_kit", + tables=["app%_kit_otel_spans", "appXXkit_unrelated"], + ) + + grants = "\n".join(workspace.statement_execution.statements) + assert "app%_kit_otel_spans" in grants + assert "appXXkit_unrelated" not in grants + + +def test_pending_statement_is_polled_to_terminal_success(): + module = load_script() + workspace = FakeWorkspace( + ["appkit_otel_spans"], pending_then_succeeded=True + ) + + response = module._execute(workspace, "0123456789abcdef", "SELECT 1") + + assert response.status.state == "SUCCEEDED" + assert workspace.statement_execution.get_statement_calls == ["statement-1"] + + +def test_failed_grant_response_is_fatal(): + module = load_script() + + with pytest.raises(RuntimeError, match="principal lacks MODIFY"): + provision( + module, + statement_options={"failed_statement": "GRANT MODIFY"}, + ) + + def test_writes_configuration_atomically(tmp_path: Path): module = load_script() output = tmp_path / ".databricks" / "mlflow-uc.json" diff --git a/packages/appkit/src/plugins/agents/manifest.json b/packages/appkit/src/plugins/agents/manifest.json index 0391c5843..fc07cf07a 100644 --- a/packages/appkit/src/plugins/agents/manifest.json +++ b/packages/appkit/src/plugins/agents/manifest.json @@ -46,7 +46,7 @@ "permission": "CAN_QUERY", "fields": { "name": { - "env": "DATABRICKS_SERVING_ENDPOINT_NAME", + "env": "DATABRICKS_AGENT_SERVING_ENDPOINT_NAME", "description": "Default LLM serving endpoint name" } } diff --git a/packages/shared/src/cli/commands/plugin/sync/sync.ts b/packages/shared/src/cli/commands/plugin/sync/sync.ts index 6aab3987b..8cfe13b35 100644 --- a/packages/shared/src/cli/commands/plugin/sync/sync.ts +++ b/packages/shared/src/cli/commands/plugin/sync/sync.ts @@ -770,9 +770,12 @@ async function runPluginsSync(options: { ); }); } else { - // npm import: direct string comparison + // npm import: accept package entrypoints as well as subpath exports. plugin = Object.values(plugins).find( - (p) => p.package === imp.source && p.name === imp.originalName, + (p) => + (p.package === imp.source || + imp.source.startsWith(`${p.package}/`)) && + p.name === imp.originalName, ); } @@ -798,17 +801,6 @@ async function runPluginsSync(options: { } } - // Step 6b: Strip requiredByTemplate for non-GA plugins - for (const plugin of Object.values(plugins)) { - if ( - plugin.requiredByTemplate && - plugin.stability && - plugin.stability !== "ga" - ) { - plugin.requiredByTemplate = undefined; - } - } - if (!options.silent && !options.json) { console.log(`\nFound ${pluginCount} plugin(s):`); for (const [name, manifest] of Object.entries(plugins)) { diff --git a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts index c55cabc89..ef6dadc9e 100644 --- a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts +++ b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts @@ -1,6 +1,7 @@ import { mkdirSync, mkdtempSync, readFileSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import yaml from "js-yaml"; import { describe, expect, test } from "vitest"; import { buildMlflowProvisionCommand, @@ -27,8 +28,38 @@ function createProject(): string { }), ); writeFileSync(join(cwd, ".env"), "DATABRICKS_CONFIG_PROFILE=DEFAULT\n"); - writeFileSync(join(cwd, "app.yaml"), "command: ['npm', 'run', 'start']\n"); - writeFileSync(join(cwd, "databricks.yml"), "bundle:\n name: traced-app\n"); + writeFileSync( + join(cwd, "app.yaml"), + [ + "command: ['npm', 'run', 'start']", + "env:", + " - name: MLFLOW_EXPERIMENT_ID", + " valueFrom: mlflow-experiment", + " - name: MLFLOW_TRACING_SQL_WAREHOUSE_ID", + " valueFrom: mlflow-tracing-warehouse", + " - name: MLFLOW_UC_CATALOG", + " value: main", + "", + ].join("\n"), + ); + writeFileSync( + join(cwd, "databricks.yml"), + [ + "bundle:", + " name: traced-app", + "variables:", + " mlflow_experiment_id:", + " description: MLflow experiment ID", + " mlflow_tracing_warehouse_id:", + " description: MLflow tracing warehouse ID", + "targets:", + " default:", + " variables:", + " mlflow_experiment_id: placeholder", + " mlflow_tracing_warehouse_id: placeholder", + "", + ].join("\n"), + ); return cwd; } @@ -80,6 +111,15 @@ describe("MLflow UC setup", () => { test("persists all tracing values for local and deployed runtimes", async () => { const cwd = createProject(); const logged: string[] = []; + writeFileSync( + join(cwd, ".env"), + [ + "DATABRICKS_CONFIG_PROFILE=DEFAULT", + "DATABRICKS_HOST=https://stale.example.com", + "DATABRICKS_HOST=https://duplicate.example.com", + "", + ].join("\n"), + ); const result = await provisionAndPersistMlflowUc( { @@ -106,13 +146,50 @@ describe("MLflow UC setup", () => { ); expect(result).toEqual(EXPECTED_VALUES); - for (const file of [".env", "app.yaml", "databricks.yml"]) { - const content = readFileSync(join(cwd, file), "utf8"); - for (const [name, value] of Object.entries(EXPECTED_VALUES)) { - expect(content).toContain(name); - expect(content).toContain(value); - } + const dotenv = readFileSync(join(cwd, ".env"), "utf8"); + for (const [name, value] of Object.entries(EXPECTED_VALUES)) { + expect(dotenv).toContain(`${name}=${value}`); } + expect(dotenv.match(/^DATABRICKS_HOST=/gm)).toHaveLength(1); + expect(dotenv).toContain( + "DATABRICKS_HOST=https://example.cloud.databricks.com", + ); + const appYaml = yaml.load(readFileSync(join(cwd, "app.yaml"), "utf8")) as { + env: Array<{ name: string; value?: string; valueFrom?: string }>; + }; + expect(appYaml.env).toEqual( + expect.arrayContaining([ + { + name: "MLFLOW_EXPERIMENT_ID", + valueFrom: "mlflow-experiment", + }, + { + name: "MLFLOW_TRACING_SQL_WAREHOUSE_ID", + valueFrom: "mlflow-tracing-warehouse", + }, + ...Object.entries(EXPECTED_VALUES) + .filter( + ([name]) => + name !== "MLFLOW_EXPERIMENT_ID" && + name !== "MLFLOW_TRACING_SQL_WAREHOUSE_ID", + ) + .map(([name, value]) => ({ name, value })), + ]), + ); + const bundle = yaml.load( + readFileSync(join(cwd, "databricks.yml"), "utf8"), + ) as { + variables: Record; + targets: { default: { variables: Record } }; + }; + expect(bundle.variables.mlflow_experiment_id.default).toBe("123456789"); + expect(bundle.variables.mlflow_tracing_warehouse_id.default).toBe( + "0123456789abcdef", + ); + expect(bundle.targets.default.variables).toMatchObject({ + mlflow_experiment_id: "123456789", + mlflow_tracing_warehouse_id: "0123456789abcdef", + }); expect(logged).toContain( "MLflow experiment: https://example.cloud.databricks.com/ml/experiments/123456789/traces", ); diff --git a/packages/shared/src/cli/commands/setup.ts b/packages/shared/src/cli/commands/setup.ts index 77c1e0060..1a0303d19 100644 --- a/packages/shared/src/cli/commands/setup.ts +++ b/packages/shared/src/cli/commands/setup.ts @@ -24,6 +24,23 @@ const MLFLOW_UC_ENV_NAMES = [ "MLFLOW_OTEL_SPANS_TABLE", ] as const; +const MLFLOW_BUNDLE_VARIABLE_NAMES: Record< + (typeof MLFLOW_UC_ENV_NAMES)[number], + string +> = { + MLFLOW_EXPERIMENT_ID: "mlflow_experiment_id", + MLFLOW_TRACING_SQL_WAREHOUSE_ID: "mlflow_tracing_warehouse_id", + MLFLOW_UC_CATALOG: "MLFLOW_UC_CATALOG", + MLFLOW_UC_SCHEMA: "MLFLOW_UC_SCHEMA", + MLFLOW_UC_TABLE_PREFIX: "MLFLOW_UC_TABLE_PREFIX", + MLFLOW_OTEL_SPANS_TABLE: "MLFLOW_OTEL_SPANS_TABLE", +}; + +const MLFLOW_RESOURCE_BINDING_ENV_NAMES = new Set([ + "MLFLOW_EXPERIMENT_ID", + "MLFLOW_TRACING_SQL_WAREHOUSE_ID", +]); + export type MlflowUcValues = Record< (typeof MLFLOW_UC_ENV_NAMES)[number], string @@ -122,15 +139,22 @@ function validateMlflowUcValues(value: unknown): MlflowUcValues { ) as MlflowUcValues; } -function persistDotEnv(filePath: string, values: MlflowUcValues): void { +function persistDotEnv( + filePath: string, + values: MlflowUcValues, + workspaceHost?: string, +): void { const existing = fs.existsSync(filePath) ? fs.readFileSync(filePath, "utf8").split(/\r?\n/) : []; const remaining = existing.filter( - (line) => !MLFLOW_UC_ENV_NAMES.some((name) => line.startsWith(`${name}=`)), + (line) => + !MLFLOW_UC_ENV_NAMES.some((name) => line.startsWith(`${name}=`)) && + (!workspaceHost || !line.startsWith("DATABRICKS_HOST=")), ); const lines = [ ...remaining.filter((line, index) => line || index < remaining.length - 1), + ...(workspaceHost ? [`DATABRICKS_HOST=${workspaceHost}`] : []), ...MLFLOW_UC_ENV_NAMES.map((name) => `${name}=${values[name]}`), ]; fs.writeFileSync(filePath, `${lines.join("\n")}\n`); @@ -160,6 +184,7 @@ function persistAppYaml(filePath: string, values: MlflowUcValues): void { env?: Array<{ name?: string; value?: string; valueFrom?: string }>; }; const existing = Array.isArray(document.env) ? document.env : []; + const existingByName = new Map(existing.map((entry) => [entry.name, entry])); document.env = [ ...existing.filter( (entry) => @@ -167,7 +192,17 @@ function persistAppYaml(filePath: string, values: MlflowUcValues): void { entry.name as (typeof MLFLOW_UC_ENV_NAMES)[number], ), ), - ...MLFLOW_UC_ENV_NAMES.map((name) => ({ name, value: values[name] })), + ...MLFLOW_UC_ENV_NAMES.map((name) => { + const entry = existingByName.get(name); + if ( + MLFLOW_RESOURCE_BINDING_ENV_NAMES.has(name) && + typeof entry?.valueFrom === "string" && + entry.valueFrom + ) { + return { name, valueFrom: entry.valueFrom }; + } + return { name, value: values[name] }; + }), ]; writeYamlObject(filePath, document as Record); } @@ -179,7 +214,10 @@ function persistBundleYaml(filePath: string, values: MlflowUcValues): void { }; document.variables ??= {}; for (const name of MLFLOW_UC_ENV_NAMES) { - document.variables[name] = { + const variableName = MLFLOW_BUNDLE_VARIABLE_NAMES[name]; + const existing = document.variables[variableName]; + document.variables[variableName] = { + ...(existing && typeof existing === "object" ? existing : {}), description: `AppKit MLflow UC tracing: ${name}`, default: values[name], }; @@ -187,7 +225,10 @@ function persistBundleYaml(filePath: string, values: MlflowUcValues): void { document.targets ??= {}; document.targets.default ??= {}; document.targets.default.variables ??= {}; - Object.assign(document.targets.default.variables, values); + for (const name of MLFLOW_UC_ENV_NAMES) { + document.targets.default.variables[MLFLOW_BUNDLE_VARIABLE_NAMES[name]] = + values[name]; + } writeYamlObject(filePath, document as Record); } @@ -226,12 +267,12 @@ export async function provisionAndPersistMlflowUc( const values = validateMlflowUcValues( JSON.parse(fs.readFileSync(outputPath, "utf8")), ); - persistDotEnv(path.join(options.cwd, ".env"), values); + const host = dependencies.workspaceHost?.replace(/\/+$/, ""); + persistDotEnv(path.join(options.cwd, ".env"), values, host); persistAppYaml(path.join(options.cwd, "app.yaml"), values); persistBundleYaml(path.join(options.cwd, "databricks.yml"), values); const log = dependencies.log ?? console.log; - const host = dependencies.workspaceHost?.replace(/\/$/, ""); if (host) { log( `MLflow experiment: ${host}/ml/experiments/${encodeURIComponent(values.MLFLOW_EXPERIMENT_ID)}/traces`, diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index f8784afc5..241751867 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -51,7 +51,7 @@ "permission": "CAN_QUERY", "fields": { "name": { - "env": "DATABRICKS_SERVING_ENDPOINT_NAME", + "env": "DATABRICKS_AGENT_SERVING_ENDPOINT_NAME", "description": "Default LLM serving endpoint name", "origin": "user" } diff --git a/template/client/src/pages/agents/AgentChat.tsx b/template/client/src/pages/agents/AgentChat.tsx index 8e0da4f71..38b1c69a9 100644 --- a/template/client/src/pages/agents/AgentChat.tsx +++ b/template/client/src/pages/agents/AgentChat.tsx @@ -196,14 +196,16 @@ export function AgentChat() { })} - {mlflowTraceId && mlflowTraceUrl && ( + {mlflowTraceId && (
{mlflowTraceId} - - Open trace in MLflow - + {mlflowTraceUrl && ( + + Open trace in MLflow + + )}
)} diff --git a/template/server/server.ts b/template/server/server.ts index be7dc49d9..357b8a407 100644 --- a/template/server/server.ts +++ b/template/server/server.ts @@ -28,7 +28,10 @@ createApp({ plugins: [ {{- range $name, $_ := .plugins}} {{- if eq $name "agents"}} - agents({ agents: { helper } }), + agents({ + defaultModel: process.env.DATABRICKS_AGENT_SERVING_ENDPOINT_NAME, + agents: { helper }, + }), {{- else}} {{$name}}(), {{- end}} diff --git a/tools/tests/generate-app-templates.test.ts b/tools/tests/generate-app-templates.test.ts index ea058c3b7..37ef13693 100644 --- a/tools/tests/generate-app-templates.test.ts +++ b/tools/tests/generate-app-templates.test.ts @@ -28,7 +28,7 @@ describe("generated AppKit agent templates", () => { readFileSync(join(app, "server/agents/helper.ts"), "utf8"), ).toContain("AGENT ā†’ TOOL ā†’ AGENT ā†’ TOOL"); expect(readFileSync(join(app, "server/server.ts"), "utf8")).toContain( - "agents({ agents: { helper } })", + "agents: { helper }", ); }, ); @@ -58,6 +58,7 @@ describe("generated AppKit agent templates", () => { const manifest = JSON.parse( readFileSync(join(app, "appkit.plugins.json"), "utf8"), ); + expect(manifest.plugins.agents.requiredByTemplate).toBe(true); expect( manifest.plugins.agents.resources.required.map( (resource: { resourceKey: string }) => resource.resourceKey, @@ -100,5 +101,26 @@ describe("generated AppKit agent templates", () => { ); expect(chat).toContain("Open trace in MLflow"); expect(chat).toContain("mlflowTraceUrl"); + expect(chat).toContain("{mlflowTraceId && ("); + expect(chat).toContain("{mlflowTraceUrl && ("); }); + + test.each(["appkit-agents", "appkit-all-in-one"])( + "%s keeps the agent model endpoint distinct from serving", + (name) => { + const app = join(outputDir, name); + const appYaml = yaml.load( + readFileSync(join(app, "app.yaml"), "utf8"), + ) as { env: Array<{ name: string; valueFrom?: string }> }; + const envNames = appYaml.env.map((entry) => entry.name); + expect(new Set(envNames).size).toBe(envNames.length); + expect(appYaml.env).toContainEqual({ + name: "DATABRICKS_AGENT_SERVING_ENDPOINT_NAME", + valueFrom: "agents-serving-endpoint", + }); + expect(readFileSync(join(app, "server/server.ts"), "utf8")).toContain( + "defaultModel: process.env.DATABRICKS_AGENT_SERVING_ENDPOINT_NAME", + ); + }, + ); }); diff --git a/vitest.config.ts b/vitest.config.ts index 7c210b359..a44d3cf9d 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -66,6 +66,15 @@ export default defineConfig({ environment: "node", }, }, + { + plugins: [tsconfigPaths()], + test: { + name: "dev-playground-server", + root: "./apps/dev-playground", + include: ["server/tests/**/*.test.ts"], + environment: "node", + }, + }, ], }, }); From 126a6a85d7794be475dbee2a33fbab11ab490a5f Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Wed, 12 Aug 2026 20:58:40 -0700 Subject: [PATCH 20/31] test: gate AppKit agent trace coverage Signed-off-by: Adam Gurary --- .../trace-conformance.integration.test.ts | 458 ++++++++++++++++++ tools/tests/agent-template-policy.test.ts | 126 +++++ 2 files changed, 584 insertions(+) create mode 100644 packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts create mode 100644 tools/tests/agent-template-policy.test.ts diff --git a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts new file mode 100644 index 000000000..3df323c8d --- /dev/null +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -0,0 +1,458 @@ +import { execFileSync } from "node:child_process"; +import { context, trace } from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import type { AgentAdapter, AgentInput, AgentRunContext } from "shared"; +import { afterAll, beforeAll, describe, expect, test, vi } from "vitest"; +import { z } from "zod"; +import { createAgent } from "../../../core/agent/create-agent"; +import { runAgent } from "../../../core/agent/run-agent"; +import { tool } from "../../../core/agent/tools/tool"; + +interface SpanManifest { + name: string; + spanType: string; + spanId: string; + parentSpanId: string | null; + inputs: unknown; + outputs: unknown; + status: number; + latencyMs: number; + model?: string; + provider?: string; + usage: Record; + costUsd?: number; + costAvailable: boolean; + links: Array<{ traceId: string; spanId: string }>; + attributes: Record; +} + +interface TraceManifest { + template: string; + traceId: string; + spans: SpanManifest[]; +} + +function decoded(value: unknown): unknown { + if (typeof value !== "string") return value; + try { + return JSON.parse(value); + } catch { + return value; + } +} + +function objectValue(value: unknown): Record { + const parsed = decoded(value); + return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) + ? (parsed as Record) + : {}; +} + +function normalize(template: string, spans: ReadableSpan[]): TraceManifest { + const traceIds = new Set(spans.map((span) => span.spanContext().traceId)); + expect(traceIds.size, `${template}: mixed trace IDs`).toBe(1); + return { + template, + traceId: [...traceIds][0], + spans: spans.map((span) => { + const attributes = { ...span.attributes }; + const usage = objectValue( + attributes[ + span.attributes["mlflow.spanType"] === "AGENT" + ? "mlflow.trace.tokenUsage" + : "mlflow.chat.tokenUsage" + ], + ) as Record; + return { + name: span.name, + spanType: String(attributes["mlflow.spanType"] ?? ""), + spanId: span.spanContext().spanId, + parentSpanId: span.parentSpanContext?.spanId ?? null, + inputs: decoded(attributes["mlflow.spanInputs"]), + outputs: decoded(attributes["mlflow.spanOutputs"]), + status: span.status.code, + latencyMs: span.duration[0] * 1_000 + span.duration[1] / 1_000_000, + model: attributes["mlflow.chat.model"] as string | undefined, + provider: attributes["mlflow.chat.provider"] as string | undefined, + usage, + costUsd: attributes["mlflow.llm.cost"] as number | undefined, + costAvailable: attributes["appkit.cost.available"] === true, + links: span.links.map((link) => ({ + traceId: link.context.traceId, + spanId: link.context.spanId, + })), + attributes, + }; + }), + }; +} + +function assertContract(manifest: TraceManifest): void { + const roots = manifest.spans.filter((span) => span.parentSpanId === null); + expect(roots, `${manifest.template}: AGENT roots`).toHaveLength(1); + const root = roots[0]; + expect(root.spanType, `${manifest.template}/${root.name}: span type`).toBe( + "AGENT", + ); + const ids = new Set(manifest.spans.map((span) => span.spanId)); + const models = manifest.spans.filter( + (span) => span.spanType === "CHAT_MODEL", + ); + const tools = manifest.spans.filter((span) => span.spanType === "TOOL"); + expect(models.length, `${manifest.template}: model children`).toBeGreaterThan( + 0, + ); + expect(tools.length, `${manifest.template}: tool children`).toBeGreaterThan( + 0, + ); + for (const span of manifest.spans) { + expect( + span.inputs, + `${manifest.template}/${span.name}: inputs`, + ).toBeDefined(); + expect( + span.outputs, + `${manifest.template}/${span.name}: outputs`, + ).toBeDefined(); + expect(span.status, `${manifest.template}/${span.name}: status`).not.toBe( + 0, + ); + expect( + span.latencyMs, + `${manifest.template}/${span.name}: latency`, + ).toBeGreaterThanOrEqual(0); + if (span.parentSpanId !== null) { + expect( + ids.has(span.parentSpanId), + `${manifest.template}/${span.name}: parent`, + ).toBe(true); + } + } + for (const model of models) { + expect( + model.model, + `${manifest.template}/${model.name}: model`, + ).toBeTruthy(); + expect( + model.provider, + `${manifest.template}/${model.name}: provider`, + ).toBeTruthy(); + expect( + model.usage.total_tokens, + `${manifest.template}/${model.name}: usage`, + ).toBe(model.usage.input_tokens + model.usage.output_tokens); + expect( + model.attributes["appkit.first_token.duration_ms"], + `${manifest.template}/${model.name}: TTFT`, + ).toEqual(expect.any(Number)); + expect( + model.attributes["appkit.stream.duration_ms"], + `${manifest.template}/${model.name}: stream duration`, + ).toEqual(expect.any(Number)); + expect( + model.costAvailable, + `${manifest.template}/${model.name}: cost`, + ).toBe(true); + } + expect( + root.attributes, + `${manifest.template}/${root.name}: identity`, + ).toMatchObject({ + "appkit.app.name": manifest.template, + "mlflow.trace.user": "user-1", + "mlflow.trace.session": "session-1", + }); + const aggregate = models.reduce( + (total, span) => total + span.usage.total_tokens, + 0, + ); + expect( + root.usage.total_tokens, + `${manifest.template}/${root.name}: aggregate usage`, + ).toBe(aggregate); + expect( + root.costUsd, + `${manifest.template}/${root.name}: aggregate cost`, + ).toBe(models.reduce((total, span) => total + (span.costUsd ?? 0), 0)); +} + +async function captureTurn(template: string): Promise { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracer = vi + .spyOn(trace, "getTracer") + .mockImplementation((name, version) => provider.getTracer(name, version)); + const clock = tool({ + name: "clock", + description: "Return UTC time", + schema: z.object({ zone: z.string() }), + execute: async () => ({ time: "12:00", zone: "UTC" }), + }); + const adapter: AgentAdapter = { + async *run(_input: AgentInput, runtime: AgentRunContext) { + const startedAt = Date.now() - 10; + yield { + type: "model_start", + stepId: "step-1", + model: "test-model", + provider: "databricks", + input: { messages: [{ role: "user", content: "Use the clock tool" }] }, + startedAt, + }; + const value = await runtime.executeTool("clock", { zone: "UTC" }); + yield { type: "message_delta", content: JSON.stringify(value) }; + yield { + type: "model_end", + stepId: "step-1", + model: "test-model", + provider: "databricks", + output: value, + usage: { + inputTokens: 7, + outputTokens: 3, + totalTokens: 10, + costUsd: 0.01, + costAvailable: true, + }, + firstTokenAt: startedAt + 2, + streamDurationMs: 10, + endedAt: startedAt + 10, + }; + }, + }; + try { + await runAgent( + createAgent({ + name: "planner", + instructions: "Use tools", + model: adapter, + tools: { clock }, + }), + { + messages: "Use the clock tool", + appName: template, + requestId: "request-1", + sessionId: "session-1", + threadId: "thread-1", + userId: "user-1", + }, + ); + await provider.forceFlush(); + return normalize(template, exporter.getFinishedSpans()); + } finally { + getTracer.mockRestore(); + await provider.shutdown(); + } +} + +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); +}); + +afterAll(() => context.disable()); + +describe.each(["appkit-agents", "appkit-all-in-one"])( + "%s trace conformance", + (template) => { + test("writes, reloads, and validates a deterministic tool turn", async () => { + const manifest = await captureTurn(template); + const reloaded = JSON.parse(JSON.stringify(manifest)) as TraceManifest; + assertContract(reloaded); + }); + }, +); + +const deployedPrerequisites = [ + "APPKIT_TRACE_CONFORMANCE_URL", + "APPKIT_TRACE_CONFORMANCE_PROFILE", + "APPKIT_TRACE_CONFORMANCE_EXPERIMENT_ID", + "APPKIT_TRACE_CONFORMANCE_WAREHOUSE_ID", + "APPKIT_TRACE_CONFORMANCE_OTEL_SPANS_TABLE", + "APPKIT_TRACE_CONFORMANCE_UC_CATALOG", + "APPKIT_TRACE_CONFORMANCE_UC_SCHEMA", + "APPKIT_TRACE_CONFORMANCE_UC_TABLE_PREFIX", +] as const; +const missingDeployed = deployedPrerequisites.filter( + (name) => !process.env[name], +); + +test.skipIf(missingDeployed.length > 0)( + "deployed AppKit agent persists the returned trace at its exact UC location", + async () => { + const profile = process.env.APPKIT_TRACE_CONFORMANCE_PROFILE ?? ""; + const token = JSON.parse( + execFileSync("databricks", ["auth", "token", "-p", profile], { + encoding: "utf8", + }), + ).access_token; + const response = await fetch( + process.env.APPKIT_TRACE_CONFORMANCE_URL ?? "", + { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + "X-MLflow-Return-Trace-Id": "true", + }, + body: JSON.stringify({ + input: [ + { role: "user", content: "What time is it? Use the clock tool." }, + ], + }), + }, + ); + expect(response.ok).toBe(true); + const traceId = response.headers.get("x-mlflow-trace-id"); + expect(traceId).toBeTruthy(); + const experiment = JSON.parse( + execFileSync( + "databricks", + [ + "api", + "get", + `/api/2.0/mlflow/experiments/get?experiment_id=${encodeURIComponent(process.env.APPKIT_TRACE_CONFORMANCE_EXPERIMENT_ID ?? "")}`, + "-p", + profile, + ], + { encoding: "utf8" }, + ), + ); + const location = experiment.experiment?.trace_location; + expect(location).toBeDefined(); + expect(JSON.stringify(location)).toContain( + process.env.APPKIT_TRACE_CONFORMANCE_UC_CATALOG, + ); + expect(JSON.stringify(location)).toContain( + process.env.APPKIT_TRACE_CONFORMANCE_UC_SCHEMA, + ); + expect(JSON.stringify(location)).toContain( + process.env.APPKIT_TRACE_CONFORMANCE_UC_TABLE_PREFIX, + ); + + let storedTrace: Record | undefined; + for (let attempt = 0; attempt < 15 && !storedTrace; attempt += 1) { + try { + storedTrace = JSON.parse( + execFileSync( + "databricks", + [ + "api", + "get", + `/api/3.0/mlflow/traces/${encodeURIComponent(traceId ?? "")}`, + "-p", + profile, + ], + { encoding: "utf8" }, + ), + ); + } catch { + await new Promise((resolve) => setTimeout(resolve, 2_000)); + } + } + expect( + storedTrace, + `MLflow could not retrieve trace ${traceId}`, + ).toBeDefined(); + const statement = + "SELECT trace_id, span_id, parent_span_id, name, attributes\n" + + "FROM IDENTIFIER(:otel_spans_table)\n" + + "WHERE trace_id = :trace_id\n" + + "ORDER BY start_time_unix_nano"; + const sql = JSON.parse( + execFileSync( + "databricks", + [ + "api", + "post", + "/api/2.0/sql/statements", + "-p", + profile, + "--json", + JSON.stringify({ + warehouse_id: process.env.APPKIT_TRACE_CONFORMANCE_WAREHOUSE_ID, + statement, + wait_timeout: "50s", + parameters: [ + { + name: "otel_spans_table", + type: "STRING", + value: process.env.APPKIT_TRACE_CONFORMANCE_OTEL_SPANS_TABLE, + }, + { name: "trace_id", type: "STRING", value: traceId }, + ], + }), + ], + { encoding: "utf8" }, + ), + ); + expect(sql.status?.state).toBe("SUCCEEDED"); + expect(sql.result?.data_array?.length).toBeGreaterThan(0); + const traceRecord = (storedTrace?.trace ?? storedTrace) as { + info?: { trace_id?: string; traceId?: string }; + data?: { spans?: Array> }; + }; + const storedSpans = traceRecord.data?.spans ?? []; + const rows = (sql.result.data_array as unknown[][]).map((row) => ({ + traceId: row[0], + spanId: row[1], + parentSpanId: row[2], + name: row[3], + attributes: objectValue(row[4]), + })); + expect(new Set(rows.map((row) => row.traceId))).toEqual(new Set([traceId])); + const persistedManifest: TraceManifest = { + template: "appkit-agents", + traceId: String(traceRecord.info?.trace_id ?? traceRecord.info?.traceId), + spans: rows.map((row) => { + const source = storedSpans.find( + (span) => (span.span_id ?? span.spanId) === row.spanId, + ); + expect( + source, + `UC span ${String(row.spanId)} missing from MLflow`, + ).toBeDefined(); + const attributes = row.attributes; + const spanType = String(attributes["mlflow.spanType"] ?? ""); + const usage = objectValue( + attributes[ + spanType === "AGENT" + ? "mlflow.trace.tokenUsage" + : "mlflow.chat.tokenUsage" + ], + ) as Record; + return { + name: String(row.name), + spanType, + spanId: String(row.spanId), + parentSpanId: row.parentSpanId ? String(row.parentSpanId) : null, + inputs: decoded(attributes["mlflow.spanInputs"]), + outputs: decoded(attributes["mlflow.spanOutputs"]), + status: Number( + (source?.status as { code?: number } | undefined)?.code ?? 0, + ), + latencyMs: Number(source?.latency_ms ?? source?.latencyMs ?? -1), + model: attributes["mlflow.chat.model"] as string | undefined, + provider: attributes["mlflow.chat.provider"] as string | undefined, + usage, + costUsd: attributes["mlflow.llm.cost"] as number | undefined, + costAvailable: attributes["appkit.cost.available"] === true, + links: [], + attributes, + }; + }), + }; + assertContract(persistedManifest); + }, + 180_000, +); diff --git a/tools/tests/agent-template-policy.test.ts b/tools/tests/agent-template-policy.test.ts new file mode 100644 index 000000000..71035620a --- /dev/null +++ b/tools/tests/agent-template-policy.test.ts @@ -0,0 +1,126 @@ +import { execFileSync } from "node:child_process"; +import { + existsSync, + mkdtempSync, + readdirSync, + readFileSync, + rmSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import yaml from "js-yaml"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; + +const root = resolve(import.meta.dirname, "../.."); +const output = mkdtempSync(join(tmpdir(), "appkit-agent-policy-")); +const requiredTraceEnvironment = [ + "MLFLOW_EXPERIMENT_ID", + "MLFLOW_TRACING_SQL_WAREHOUSE_ID", + "MLFLOW_UC_CATALOG", + "MLFLOW_UC_SCHEMA", + "MLFLOW_UC_TABLE_PREFIX", + "MLFLOW_OTEL_SPANS_TABLE", +]; + +interface GeneratedCandidate { + name: string; + directory: string; +} + +function discoverGeneratedAgentTemplates(): GeneratedCandidate[] { + return readdirSync(output, { withFileTypes: true }) + .filter((entry) => entry.isDirectory()) + .map((entry) => ({ name: entry.name, directory: join(output, entry.name) })) + .filter(({ directory }) => { + const server = readFileSync(join(directory, "server/server.ts"), "utf8"); + return /\bagents\s*\(/.test(server) || /agents:\s*\{/.test(server); + }); +} + +describe("behavior-discovered generated agent template policy", () => { + beforeAll(() => { + const compatibleCli = "/tmp/databricks-cli-1.11.0/databricks"; + execFileSync("pnpm", ["generate:app-templates"], { + cwd: root, + env: { + ...process.env, + APP_TEMPLATES_OUTPUT_DIR: output, + ...(process.env.DATABRICKS_CLI + ? {} + : existsSync(compatibleCli) + ? { DATABRICKS_CLI: compatibleCli } + : {}), + }, + stdio: "pipe", + }); + }, 120_000); + + afterAll(() => rmSync(output, { recursive: true, force: true })); + + test("discovers agent surfaces from generated runtime behavior", () => { + const candidates = discoverGeneratedAgentTemplates(); + expect(candidates.length).toBeGreaterThan(0); + expect(candidates.map(({ name }) => name).sort()).toEqual([ + "appkit-agents", + "appkit-all-in-one", + ]); + }); + + test("every discovered surface declares immutable UC trace resources", () => { + for (const candidate of discoverGeneratedAgentTemplates()) { + const app = yaml.load( + readFileSync(join(candidate.directory, "app.yaml"), "utf8"), + ) as { env: Array<{ name: string; value?: string; valueFrom?: string }> }; + const names = app.env.map((entry) => entry.name); + expect(names, candidate.name).toEqual( + expect.arrayContaining(requiredTraceEnvironment), + ); + const manifest = JSON.parse( + readFileSync(join(candidate.directory, "appkit.plugins.json"), "utf8"), + ); + const resources = manifest.plugins.agents.resources.required.map( + (resource: { type: string; resourceKey: string }) => ({ + type: resource.type, + resourceKey: resource.resourceKey, + }), + ); + expect(resources, candidate.name).toEqual([ + { type: "experiment", resourceKey: "mlflow-experiment" }, + { + type: "sql_warehouse", + resourceKey: "mlflow-tracing-warehouse", + }, + ]); + const staticValues = Object.fromEntries( + app.env + .filter((entry) => entry.name.startsWith("MLFLOW_UC_")) + .map((entry) => [entry.name, entry.value]), + ); + expect(staticValues, candidate.name).toEqual({ + MLFLOW_UC_CATALOG: "main", + MLFLOW_UC_SCHEMA: "agent_traces", + MLFLOW_UC_TABLE_PREFIX: "appkit", + }); + expect( + app.env.find((entry) => entry.name === "MLFLOW_OTEL_SPANS_TABLE") + ?.value, + candidate.name, + ).toBe("main.agent_traces.appkit_otel_spans"); + } + }); + + test("a newly generated behavioral agent cannot bypass the same policy", () => { + for (const candidate of discoverGeneratedAgentTemplates()) { + const packageJson = JSON.parse( + readFileSync(join(candidate.directory, "package.json"), "utf8"), + ); + expect(packageJson.scripts.setup, candidate.name).toBe( + "appkit setup --write --mlflow-uc", + ); + expect( + packageJson.dependencies["@databricks/appkit"], + candidate.name, + ).toBe("0.59.0"); + } + }); +}); From b01b729cbf13f11514cbdd168a7dd22202392523 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Wed, 12 Aug 2026 23:24:39 -0700 Subject: [PATCH 21/31] test: close AppKit trace conformance gaps Signed-off-by: Adam Gurary --- .../trace-conformance.integration.test.ts | 1021 +++++++++++++++-- tools/generate-app-templates.ts | 7 +- tools/tests/agent-template-policy.test.ts | 81 +- tools/tests/generate-app-templates.test.ts | 16 +- 4 files changed, 1009 insertions(+), 116 deletions(-) diff --git a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts index 3df323c8d..9d430f03c 100644 --- a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -1,4 +1,7 @@ import { execFileSync } from "node:child_process"; +import { existsSync, mkdtempSync, rmSync } from "node:fs"; +import { join, resolve } from "node:path"; +import { pathToFileURL } from "node:url"; import { context, trace } from "@opentelemetry/api"; import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; import { @@ -13,6 +16,12 @@ import { z } from "zod"; import { createAgent } from "../../../core/agent/create-agent"; import { runAgent } from "../../../core/agent/run-agent"; import { tool } from "../../../core/agent/tools/tool"; +import type { AgentDefinition } from "../../../core/agent/types"; + +const repositoryRoot = resolve(import.meta.dirname, "../../../../../.."); +const generatedApps = mkdtempSync( + join(repositoryRoot, ".trace-conformance-generated-"), +); interface SpanManifest { name: string; @@ -21,7 +30,7 @@ interface SpanManifest { parentSpanId: string | null; inputs: unknown; outputs: unknown; - status: number; + status: string | number; latencyMs: number; model?: string; provider?: string; @@ -62,6 +71,29 @@ function normalize(template: string, spans: ReadableSpan[]): TraceManifest { traceId: [...traceIds][0], spans: spans.map((span) => { const attributes = { ...span.attributes }; + const firstToken = + attributes.ttft_ms ?? + attributes["appkit.ttft_ms"] ?? + attributes["appkit.first_token.duration_ms"] ?? + attributes["gen_ai.latency.time_to_first_token_ms"]; + const streamDuration = + attributes.stream_duration_ms ?? + attributes["appkit.stream_duration_ms"] ?? + attributes["appkit.stream.duration_ms"] ?? + attributes["gen_ai.latency.stream_ms"]; + if (firstToken !== undefined) attributes.ttft_ms = firstToken; + if (streamDuration !== undefined) { + attributes.stream_duration_ms = streamDuration; + } + if (firstToken !== undefined || streamDuration !== undefined) { + attributes.streaming = true; + } + attributes.app_id ??= + attributes["app.id"] ?? attributes["appkit.app.name"]; + attributes.user_id ??= + attributes["user.id"] ?? attributes["mlflow.trace.user"]; + attributes.session_id ??= + attributes["session.id"] ?? attributes["mlflow.trace.session"]; const usage = objectValue( attributes[ span.attributes["mlflow.spanType"] === "AGENT" @@ -94,92 +126,270 @@ function normalize(template: string, spans: ReadableSpan[]): TraceManifest { } function assertContract(manifest: TraceManifest): void { + const fail = ( + span: SpanManifest | undefined, + field: string, + detail: string, + ): never => { + throw new Error( + `template=${manifest.template || ""} span=${span?.name ?? ""} field=${field}: ${detail}`, + ); + }; + const hasValue = (value: unknown) => + value !== undefined && + value !== null && + value !== "" && + !(Array.isArray(value) && value.length === 0) && + !( + typeof value === "object" && + value !== null && + !Array.isArray(value) && + Object.keys(value).length === 0 + ); + const terminalStatus = (status: string | number) => + status === 1 ? "OK" : status === 2 ? "ERROR" : String(status).toUpperCase(); + const usageFields = [ + "input_tokens", + "output_tokens", + "total_tokens", + ] as const; + const assertUsage = (span: SpanManifest) => { + for (const field of usageFields) { + const value = span.usage[field]; + if (typeof value !== "number" || !Number.isFinite(value) || value < 0) { + fail(span, `usage.${field}`, "must be a non-negative number"); + } + } + if ( + span.usage.total_tokens < + Math.max(span.usage.input_tokens, span.usage.output_tokens) + ) { + fail(span, "usage.total_tokens", "is smaller than a component"); + } + }; + const assertCost = (span: SpanManifest) => { + if (typeof span.costAvailable !== "boolean") { + fail(span, "cost_available", "must explicitly be true or false"); + } + if (span.costAvailable) { + if ( + typeof span.costUsd !== "number" || + !Number.isFinite(span.costUsd) || + span.costUsd < 0 + ) { + fail(span, "cost_usd", "available cost must be non-negative"); + } + } else if (span.costUsd !== undefined) { + fail(span, "cost_usd", "unavailable cost must not be reported as zero"); + } + }; + const secretKeys = new Set([ + "accesstoken", + "apikey", + "authorization", + "clientsecret", + "cookie", + "credential", + "credentials", + "databrickstoken", + "password", + "refreshtoken", + "secret", + "setcookie", + "token", + "xapikey", + ]); + const assertRedacted = (span: SpanManifest, value: unknown): void => { + if (Array.isArray(value)) { + for (const nested of value) assertRedacted(span, nested); + return; + } + if (typeof value === "object" && value !== null) { + for (const [key, nested] of Object.entries(value)) { + if ( + secretKeys.has(key.toLowerCase().replace(/[^a-z0-9]/g, "")) && + nested !== "[REDACTED]" + ) { + fail(span, "credentials", `${key} is not redacted`); + } + assertRedacted(span, nested); + } + return; + } + if ( + typeof value === "string" && + /\b(?:authorization|api[ _-]?key|password|secret|token|credentials?)\b\s*(?::|=|is)?\s+(?!\[REDACTED\])(?:Bearer\s+)?[^\s,;}]+/i.test( + value, + ) + ) { + fail(span, "credentials", "captured text contains an unredacted secret"); + } + }; + if (!manifest.template) fail(undefined, "template", "missing template"); + if (!manifest.traceId) fail(undefined, "trace_id", "missing trace identity"); + if (manifest.spans.length === 0) + fail(undefined, "spans", "trace has no spans"); const roots = manifest.spans.filter((span) => span.parentSpanId === null); - expect(roots, `${manifest.template}: AGENT roots`).toHaveLength(1); + if (roots.length !== 1 || roots[0]?.spanType !== "AGENT") { + fail( + roots.at(-1) ?? manifest.spans[0], + "AGENT root", + "trace must have exactly one parentless AGENT", + ); + } const root = roots[0]; - expect(root.spanType, `${manifest.template}/${root.name}: span type`).toBe( + const supportedTypes = new Set([ "AGENT", - ); - const ids = new Set(manifest.spans.map((span) => span.spanId)); + "CHAIN", + "CHAT_MODEL", + "EMBEDDING", + "LLM", + "MEMORY", + "PARSER", + "RETRIEVER", + "TOOL", + ]); + const spansById = new Map(); + for (const span of manifest.spans) { + if (!span.name) fail(span, "name", "missing span name"); + if (!supportedTypes.has(span.spanType)) { + fail(span, "span_type", `unsupported semantic type ${span.spanType}`); + } + if (!span.spanId) fail(span, "span_id", "missing span identity"); + if (spansById.has(span.spanId)) { + fail(span, "span_id", `duplicate span identity ${span.spanId}`); + } + spansById.set(span.spanId, span); + } const models = manifest.spans.filter( - (span) => span.spanType === "CHAT_MODEL", - ); - const tools = manifest.spans.filter((span) => span.spanType === "TOOL"); - expect(models.length, `${manifest.template}: model children`).toBeGreaterThan( - 0, - ); - expect(tools.length, `${manifest.template}: tool children`).toBeGreaterThan( - 0, + (span) => span.spanType === "CHAT_MODEL" || span.spanType === "LLM", ); + if (models.length === 0) + fail(root, "semantic child", "trace has no model child"); for (const span of manifest.spans) { - expect( - span.inputs, - `${manifest.template}/${span.name}: inputs`, - ).toBeDefined(); - expect( - span.outputs, - `${manifest.template}/${span.name}: outputs`, - ).toBeDefined(); - expect(span.status, `${manifest.template}/${span.name}: status`).not.toBe( - 0, - ); - expect( - span.latencyMs, - `${manifest.template}/${span.name}: latency`, - ).toBeGreaterThanOrEqual(0); + if (!hasValue(span.inputs)) + fail(span, "inputs", "captured inputs are missing"); + if (!hasValue(span.outputs)) + fail(span, "outputs", "captured outputs are missing"); + const status = terminalStatus(span.status); + if (!new Set(["OK", "SUCCESS", "ERROR", "CANCELLED"]).has(status)) { + fail(span, "status", "span is not finalized with a terminal status"); + } + if ( + typeof span.latencyMs !== "number" || + !Number.isFinite(span.latencyMs) || + span.latencyMs < 0 + ) { + fail(span, "latency_ms", "missing or invalid latency"); + } + if ( + status === "ERROR" && + !( + typeof span.outputs === "object" && + span.outputs !== null && + hasValue((span.outputs as Record).partial_output) + ) + ) { + fail( + span, + "outputs.partial_output", + "failed span must retain partial_output", + ); + } + assertCost(span); + assertRedacted(span, span.inputs); + assertRedacted(span, span.outputs); + assertRedacted(span, span.attributes); if (span.parentSpanId !== null) { - expect( - ids.has(span.parentSpanId), - `${manifest.template}/${span.name}: parent`, - ).toBe(true); + if (!spansById.has(span.parentSpanId)) { + fail(span, "parent_span_id", `orphan parent ${span.parentSpanId}`); + } + } + const remoteTraceId = span.attributes.remote_trace_id; + if (remoteTraceId) { + const remoteSpanId = span.attributes.remote_span_id; + if (!remoteSpanId) fail(span, "remote_span_id", "remote root is missing"); + if (span.attributes.remote_lifecycle_complete !== true) { + fail( + span, + "remote_lifecycle_complete", + "remote lifecycle is incomplete", + ); + } + if ( + remoteTraceId !== manifest.traceId && + !span.links.some( + (link) => + link.traceId === remoteTraceId && link.spanId === remoteSpanId, + ) + ) { + fail( + span, + "links", + "orphan remote trace is neither continued nor linked", + ); + } + } + } + for (const span of manifest.spans) { + const ancestry = new Set(); + let current = span; + while (current.parentSpanId !== null) { + if (ancestry.has(current.spanId)) { + fail(span, "parent_span_id", "parent ancestry contains a cycle"); + } + ancestry.add(current.spanId); + current = + spansById.get(current.parentSpanId) ?? + fail(span, "parent_span_id", "span has an orphan parent"); } } for (const model of models) { - expect( - model.model, - `${manifest.template}/${model.name}: model`, - ).toBeTruthy(); - expect( - model.provider, - `${manifest.template}/${model.name}: provider`, - ).toBeTruthy(); - expect( - model.usage.total_tokens, - `${manifest.template}/${model.name}: usage`, - ).toBe(model.usage.input_tokens + model.usage.output_tokens); - expect( - model.attributes["appkit.first_token.duration_ms"], - `${manifest.template}/${model.name}: TTFT`, - ).toEqual(expect.any(Number)); - expect( - model.attributes["appkit.stream.duration_ms"], - `${manifest.template}/${model.name}: stream duration`, - ).toEqual(expect.any(Number)); - expect( - model.costAvailable, - `${manifest.template}/${model.name}: cost`, - ).toBe(true); + if (!model.model) fail(model, "model", "model identity is missing"); + if (!model.provider) + fail(model, "provider", "provider identity is missing"); + assertUsage(model); + if (model.attributes.streaming === true) { + for (const field of ["ttft_ms", "stream_duration_ms"] as const) { + const value = model.attributes[field]; + if (typeof value !== "number" || !Number.isFinite(value) || value < 0) { + fail(model, field, "stream timing is missing or invalid"); + } + } + } + } + for (const field of ["app_id", "user_id", "session_id"] as const) { + if (!hasValue(root.attributes[field])) { + fail(root, field, "request identity is missing"); + } + } + assertUsage(root); + for (const field of usageFields) { + const expected = models.reduce( + (total, span) => total + span.usage[field], + 0, + ); + if (root.usage[field] !== expected) { + fail( + root, + `usage.${field}`, + `aggregate does not equal descendant total ${expected}`, + ); + } + } + const costAvailable = models.every((span) => span.costAvailable); + if (root.costAvailable !== costAvailable) { + fail(root, "cost_available", "does not match descendant availability"); + } + if (costAvailable) { + const expected = models.reduce( + (total, span) => total + (span.costUsd ?? 0), + 0, + ); + if (Math.abs((root.costUsd ?? Number.NaN) - expected) > 1e-12) { + fail(root, "cost_usd", "does not equal descendant cost total"); + } } - expect( - root.attributes, - `${manifest.template}/${root.name}: identity`, - ).toMatchObject({ - "appkit.app.name": manifest.template, - "mlflow.trace.user": "user-1", - "mlflow.trace.session": "session-1", - }); - const aggregate = models.reduce( - (total, span) => total + span.usage.total_tokens, - 0, - ); - expect( - root.usage.total_tokens, - `${manifest.template}/${root.name}: aggregate usage`, - ).toBe(aggregate); - expect( - root.costUsd, - `${manifest.template}/${root.name}: aggregate cost`, - ).toBe(models.reduce((total, span) => total + (span.costUsd ?? 0), 0)); } async function captureTurn(template: string): Promise { @@ -253,14 +463,95 @@ async function captureTurn(template: string): Promise { } } +async function captureGeneratedTurn(template: string): Promise { + const helperPath = join(generatedApps, template, "server/agents/helper.ts"); + const generated = (await import(pathToFileURL(helperPath).href)) as { + helper: AgentDefinition; + }; + const adapter: AgentAdapter = { + async *run(_input: AgentInput, runtime: AgentRunContext) { + const startedAt = Date.now() - 10; + yield { + type: "model_start", + stepId: "generated-step", + model: "generated-test-model", + provider: "databricks", + input: { messages: [{ role: "user", content: "Use count_words" }] }, + startedAt, + }; + const value = await runtime.executeTool("count_words", { + text: "hello traced world", + }); + yield { type: "message_delta", content: JSON.stringify(value) }; + yield { + type: "model_end", + stepId: "generated-step", + model: "generated-test-model", + provider: "databricks", + output: value, + usage: { + inputTokens: 7, + outputTokens: 3, + totalTokens: 10, + costUsd: 0.01, + costAvailable: true, + }, + firstTokenAt: startedAt + 2, + streamDurationMs: 10, + endedAt: startedAt + 10, + }; + }, + }; + generated.helper.model = adapter; + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + const getTracer = vi + .spyOn(trace, "getTracer") + .mockImplementation((name, version) => provider.getTracer(name, version)); + try { + await runAgent(generated.helper, { + messages: "Count the words in hello traced world. Use count_words.", + appName: template, + requestId: "request-1", + sessionId: "session-1", + threadId: "thread-1", + userId: "user-1", + }); + await provider.forceFlush(); + return normalize(template, exporter.getFinishedSpans()); + } finally { + getTracer.mockRestore(); + await provider.shutdown(); + } +} + beforeAll(() => { + const compatibleCli = "/tmp/databricks-cli-1.11.0/databricks"; + execFileSync("pnpm", ["generate:app-templates"], { + cwd: repositoryRoot, + env: { + ...process.env, + APP_TEMPLATES_OUTPUT_DIR: generatedApps, + ...(process.env.DATABRICKS_CLI + ? {} + : existsSync(compatibleCli) + ? { DATABRICKS_CLI: compatibleCli } + : {}), + }, + stdio: "pipe", + }); context.disable(); context.setGlobalContextManager( new AsyncLocalStorageContextManager().enable(), ); }); -afterAll(() => context.disable()); +afterAll(() => { + context.disable(); + rmSync(generatedApps, { recursive: true, force: true }); +}); describe.each(["appkit-agents", "appkit-all-in-one"])( "%s trace conformance", @@ -273,12 +564,488 @@ describe.each(["appkit-agents", "appkit-all-in-one"])( }, ); +describe.each(["appkit-agents", "appkit-all-in-one"])( + "%s generated production package", + (template) => { + test("invokes its own generated agent definition and validates the trace", async () => { + const manifest = await captureGeneratedTurn(template); + expect( + manifest.spans.some( + (span) => + span.spanType === "TOOL" && span.name === "count_words tool", + ), + JSON.stringify( + manifest.spans.map((span) => [span.spanType, span.name]), + ), + ).toBe(true); + assertContract(manifest); + }); + }, +); + +function requiredSpan(manifest: TraceManifest, spanType: string): SpanManifest { + const span = manifest.spans.find( + (candidate) => candidate.spanType === spanType, + ); + if (!span) throw new Error(`fixture is missing ${spanType}`); + return span; +} + +test.each([ + { + name: "duplicate span identity", + mutate(manifest: TraceManifest) { + const model = requiredSpan(manifest, "CHAT_MODEL"); + const toolSpan = requiredSpan(manifest, "TOOL"); + toolSpan.spanId = model.spanId; + }, + expected: /duplicate span identity/, + }, + { + name: "orphan parent", + mutate(manifest: TraceManifest) { + const toolSpan = requiredSpan(manifest, "TOOL"); + toolSpan.parentSpanId = "missing-parent"; + }, + expected: /orphan parent/, + }, + { + name: "parent cycle", + mutate(manifest: TraceManifest) { + const model = requiredSpan(manifest, "CHAT_MODEL"); + const toolSpan = requiredSpan(manifest, "TOOL"); + model.parentSpanId = toolSpan.spanId; + toolSpan.parentSpanId = model.spanId; + }, + expected: /cycle/, + }, +])("rejects $name", async ({ mutate, expected }) => { + const manifest = await captureTurn("topology-fixture"); + mutate(manifest); + + expect(() => assertContract(manifest)).toThrow(expected); +}); + +const fixtureUsage = { + input_tokens: 7, + output_tokens: 3, + total_tokens: 10, +}; + +function fixtureSpan( + name: string, + spanType: string, + spanId: string, + parentSpanId: string | null, +): SpanManifest { + return { + name, + spanType, + spanId, + parentSpanId, + inputs: { value: "input" }, + outputs: { value: "output" }, + status: "OK", + latencyMs: 1, + usage: {}, + costAvailable: false, + links: [], + attributes: {}, + }; +} + +function fixtureModel(spanId = "model", parentSpanId = "root"): SpanManifest { + return { + ...fixtureSpan("model call", "CHAT_MODEL", spanId, parentSpanId), + model: "test-model", + provider: "databricks", + usage: { ...fixtureUsage }, + attributes: { + streaming: true, + ttft_ms: 2, + stream_duration_ms: 8, + }, + }; +} + +function fixtureManifest(children: SpanManifest[]): TraceManifest { + return { + template: "fixture-template", + traceId: "0123456789abcdef0123456789abcdef", + spans: [ + { + ...fixtureSpan("request", "AGENT", "root", null), + usage: { ...fixtureUsage }, + attributes: { + app_id: "fixture-template", + user_id: "user-1", + session_id: "session-1", + }, + }, + ...children, + ], + }; +} + +function validWorkloads(): TraceManifest[] { + const tool = fixtureSpan("clock", "TOOL", "tool", "root"); + tool.inputs = { zone: "UTC" }; + tool.outputs = { time: "12:00" }; + const retriever = fixtureSpan( + "vector search", + "RETRIEVER", + "retriever", + "root", + ); + const remote = fixtureSpan("remote agent", "TOOL", "remote", "root"); + remote.links = [ + { + traceId: "fedcba9876543210fedcba9876543210", + spanId: "0123456789abcdef", + }, + ]; + remote.attributes = { + remote_trace_id: "fedcba9876543210fedcba9876543210", + remote_span_id: "0123456789abcdef", + remote_lifecycle_complete: true, + }; + return [ + fixtureManifest([fixtureModel()]), + fixtureManifest([fixtureModel("plan"), tool]), + fixtureManifest([retriever, fixtureModel("answer", "retriever")]), + fixtureManifest([fixtureModel(), remote]), + ]; +} + +test.each([ + ["simple", 0], + ["tool-using", 1], + ["retrieval-generation", 2], + ["remote-agent", 3], +] as const)("accepts complete %s workload", (_name, index) => { + expect(() => assertContract(validWorkloads()[index])).not.toThrow(); +}); + +test("accepts truthful unavailable cost without inventing zero", () => { + const manifest = fixtureManifest([fixtureModel()]); + expect(manifest.spans[0].costUsd).toBeUndefined(); + expect(manifest.spans[1].costUsd).toBeUndefined(); + expect(() => assertContract(manifest)).not.toThrow(); +}); + +test("accepts and exactly aggregates available model usage and cost", () => { + const first = fixtureModel("model-1"); + first.usage = { input_tokens: 4, output_tokens: 1, total_tokens: 5 }; + first.costAvailable = true; + first.costUsd = 0.01; + const second = fixtureModel("model-2"); + second.usage = { input_tokens: 3, output_tokens: 2, total_tokens: 5 }; + second.costAvailable = true; + second.costUsd = 0.02; + const manifest = fixtureManifest([first, second]); + manifest.spans[0].costAvailable = true; + manifest.spans[0].costUsd = 0.03; + expect(() => assertContract(manifest)).not.toThrow(); +}); + +test.each([ + { + name: "root-only trace", + mutate(manifest: TraceManifest) { + manifest.spans = manifest.spans.slice(0, 1); + }, + expected: /semantic child/, + }, + { + name: "missing model output", + mutate(manifest: TraceManifest) { + manifest.spans[1].outputs = undefined; + }, + expected: /outputs/, + }, + { + name: "missing model usage", + mutate(manifest: TraceManifest) { + manifest.spans[1].usage = {}; + }, + expected: /usage\.input_tokens/, + }, + { + name: "false zero cost", + mutate(manifest: TraceManifest) { + manifest.spans[1].costUsd = 0; + }, + expected: /cost_usd/, + }, + { + name: "orphan remote trace", + mutate(manifest: TraceManifest) { + manifest.spans[2].attributes = { + remote_trace_id: "fedcba9876543210fedcba9876543210", + remote_span_id: "0123456789abcdef", + remote_lifecycle_complete: true, + }; + }, + expected: /links/, + }, + { + name: "incomplete tool", + mutate(manifest: TraceManifest) { + manifest.spans[2].outputs = undefined; + }, + expected: /outputs/, + }, + { + name: "missing identity", + mutate(manifest: TraceManifest) { + delete manifest.spans[0].attributes.user_id; + }, + expected: /user_id/, + }, + { + name: "duplicate roots", + mutate(manifest: TraceManifest) { + manifest.spans.push( + fixtureSpan("second request", "AGENT", "root-2", null), + ); + }, + expected: /AGENT root/, + }, + { + name: "unfinalized span", + mutate(manifest: TraceManifest) { + manifest.spans[1].status = "UNSET"; + }, + expected: /status/, + }, + { + name: "failure without partial output", + mutate(manifest: TraceManifest) { + manifest.spans[1].status = "ERROR"; + manifest.spans[1].outputs = { error: "provider unavailable" }; + }, + expected: /partial_output/, + }, + { + name: "wrong input aggregation", + mutate(manifest: TraceManifest) { + manifest.spans[0].usage.input_tokens = 6; + }, + expected: /usage\.input_tokens/, + }, + { + name: "credential leak", + mutate(manifest: TraceManifest) { + manifest.spans[2].inputs = { Authorization: "Bearer provider-secret" }; + }, + expected: /credentials/, + }, + { + name: "unsupported semantic type", + mutate(manifest: TraceManifest) { + manifest.spans[2].spanType = "HTTP"; + }, + expected: /span_type/, + }, + { + name: "missing provider", + mutate(manifest: TraceManifest) { + manifest.spans[1].provider = undefined; + }, + expected: /provider/, + }, + { + name: "negative latency", + mutate(manifest: TraceManifest) { + manifest.spans[2].latencyMs = -1; + }, + expected: /latency_ms/, + }, +] as const)( + "rejects $name with template and span context", + ({ mutate, expected }) => { + const manifest = fixtureManifest([ + fixtureModel(), + fixtureSpan("clock", "TOOL", "tool", "root"), + ]); + mutate(manifest); + expect(() => assertContract(manifest)).toThrow(expected); + try { + assertContract(manifest); + } catch (error) { + expect(String(error)).toContain("fixture-template"); + } + }, +); + +test.each(["ttft_ms", "stream_duration_ms"] as const)( + "rejects streaming model missing %s", + (field) => { + const manifest = fixtureManifest([fixtureModel()]); + delete manifest.spans[1].attributes[field]; + expect(() => assertContract(manifest)).toThrow(field); + }, +); + +interface StatementResponse { + statement_id?: string; + status?: { state?: string; error?: { message?: string } }; + result?: { data_array?: unknown[][] }; +} + +function deriveUcBinding( + catalog: string, + schema: string, + prefix: string, +): { + location: Record; + spansTable: string; + mlflowTracePrefix: string; +} { + const spansTable = `${catalog}.${schema}.${prefix}_otel_spans`; + return { + location: { + type: "UC_TABLE_PREFIX", + uc_table_prefix: { + catalog_name: catalog, + schema_name: schema, + table_prefix: prefix, + otel_spans_table_name: spansTable, + }, + }, + spansTable, + mlflowTracePrefix: `trace:/${catalog}.${schema}.${prefix}/`, + }; +} + +function otelTraceIdFromReturnedTrace( + traceId: string, + mlflowTracePrefix: string, +): string { + if (!traceId.startsWith(mlflowTracePrefix)) { + throw new Error( + `returned trace ${traceId} is not bound to exact UC location ${mlflowTracePrefix}`, + ); + } + const otelTraceId = traceId.slice(mlflowTracePrefix.length); + if (!/^[0-9a-f]{32}$/i.test(otelTraceId)) { + throw new Error( + `returned trace ${traceId} has no exact OTel trace identity`, + ); + } + return otelTraceId.toLowerCase(); +} + +async function pollStatement( + initial: StatementResponse, + getStatement: (statementId: string) => StatementResponse, + sleep: () => Promise, + maxPolls = 120, +): Promise { + let response = initial; + for (let poll = 0; poll < maxPolls; poll += 1) { + const state = response.status?.state; + if (state === "SUCCEEDED") return response; + if (state === "FAILED" || state === "CANCELED" || state === "CLOSED") { + throw new Error( + `UC trace row query ${state}: ${response.status?.error?.message ?? "unknown error"}`, + ); + } + if (!response.statement_id) { + throw new Error(`UC trace row query is ${state} without a statement ID`); + } + await sleep(); + response = getStatement(response.statement_id); + } + throw new Error(`UC trace row query did not finish after ${maxPolls} polls`); +} + +test("derives the exact immutable UC location, table, and MLflow trace prefix", () => { + expect(deriveUcBinding("main", "agent_traces", "appkit")).toEqual({ + location: { + type: "UC_TABLE_PREFIX", + uc_table_prefix: { + catalog_name: "main", + schema_name: "agent_traces", + table_prefix: "appkit", + otel_spans_table_name: "main.agent_traces.appkit_otel_spans", + }, + }, + spansTable: "main.agent_traces.appkit_otel_spans", + mlflowTracePrefix: "trace:/main.agent_traces.appkit/", + }); +}); + +test("extracts the OTel trace ID only from the exact returned UC trace identity", () => { + expect( + otelTraceIdFromReturnedTrace( + "trace:/main.agent_traces.appkit/0123456789abcdef0123456789abcdef", + "trace:/main.agent_traces.appkit/", + ), + ).toBe("0123456789abcdef0123456789abcdef"); + expect(() => + otelTraceIdFromReturnedTrace( + "trace:/other.schema.prefix/0123456789abcdef0123456789abcdef", + "trace:/main.agent_traces.appkit/", + ), + ).toThrow(/exact UC location/); +}); + +test("polls asynchronous Statement Execution to success", async () => { + const getStatement = vi.fn(() => ({ + statement_id: "statement-1", + status: { state: "SUCCEEDED" }, + result: { data_array: [["trace"]] }, + })); + const result = await pollStatement( + { statement_id: "statement-1", status: { state: "PENDING" } }, + getStatement, + async () => undefined, + ); + expect(getStatement).toHaveBeenCalledWith("statement-1"); + expect(result.result?.data_array).toEqual([["trace"]]); +}); + +test.each(["FAILED", "CANCELED", "CLOSED"])( + "reports terminal Statement Execution state %s", + async (state) => { + await expect( + pollStatement( + { + statement_id: "statement-1", + status: { state, error: { message: "warehouse rejected query" } }, + }, + () => ({ status: { state: "SUCCEEDED" } }), + async () => undefined, + ), + ).rejects.toThrow(new RegExp(`${state}.*warehouse rejected query`, "i")); + }, +); + +test("rejects nonterminal SQL without an identity and bounds polling", async () => { + await expect( + pollStatement( + { status: { state: "PENDING" } }, + () => ({ status: { state: "SUCCEEDED" } }), + async () => undefined, + ), + ).rejects.toThrow(/without a statement ID/); + await expect( + pollStatement( + { statement_id: "statement-1", status: { state: "PENDING" } }, + () => ({ statement_id: "statement-1", status: { state: "RUNNING" } }), + async () => undefined, + 2, + ), + ).rejects.toThrow(/after 2 polls/); +}); + const deployedPrerequisites = [ "APPKIT_TRACE_CONFORMANCE_URL", + "APPKIT_TRACE_CONFORMANCE_APP_NAME", "APPKIT_TRACE_CONFORMANCE_PROFILE", "APPKIT_TRACE_CONFORMANCE_EXPERIMENT_ID", "APPKIT_TRACE_CONFORMANCE_WAREHOUSE_ID", - "APPKIT_TRACE_CONFORMANCE_OTEL_SPANS_TABLE", "APPKIT_TRACE_CONFORMANCE_UC_CATALOG", "APPKIT_TRACE_CONFORMANCE_UC_SCHEMA", "APPKIT_TRACE_CONFORMANCE_UC_TABLE_PREFIX", @@ -291,6 +1058,13 @@ test.skipIf(missingDeployed.length > 0)( "deployed AppKit agent persists the returned trace at its exact UC location", async () => { const profile = process.env.APPKIT_TRACE_CONFORMANCE_PROFILE ?? ""; + const appName = process.env.APPKIT_TRACE_CONFORMANCE_APP_NAME ?? ""; + expect(["appkit-agents", "appkit-all-in-one"]).toContain(appName); + const binding = deriveUcBinding( + process.env.APPKIT_TRACE_CONFORMANCE_UC_CATALOG ?? "", + process.env.APPKIT_TRACE_CONFORMANCE_UC_SCHEMA ?? "", + process.env.APPKIT_TRACE_CONFORMANCE_UC_TABLE_PREFIX ?? "", + ); const token = JSON.parse( execFileSync("databricks", ["auth", "token", "-p", profile], { encoding: "utf8", @@ -315,6 +1089,10 @@ test.skipIf(missingDeployed.length > 0)( expect(response.ok).toBe(true); const traceId = response.headers.get("x-mlflow-trace-id"); expect(traceId).toBeTruthy(); + const otelTraceId = otelTraceIdFromReturnedTrace( + traceId ?? "", + binding.mlflowTracePrefix, + ); const experiment = JSON.parse( execFileSync( "databricks", @@ -329,16 +1107,7 @@ test.skipIf(missingDeployed.length > 0)( ), ); const location = experiment.experiment?.trace_location; - expect(location).toBeDefined(); - expect(JSON.stringify(location)).toContain( - process.env.APPKIT_TRACE_CONFORMANCE_UC_CATALOG, - ); - expect(JSON.stringify(location)).toContain( - process.env.APPKIT_TRACE_CONFORMANCE_UC_SCHEMA, - ); - expect(JSON.stringify(location)).toContain( - process.env.APPKIT_TRACE_CONFORMANCE_UC_TABLE_PREFIX, - ); + expect(location).toEqual(binding.location); let storedTrace: Record | undefined; for (let attempt = 0; attempt < 15 && !storedTrace; attempt += 1) { @@ -364,6 +1133,13 @@ test.skipIf(missingDeployed.length > 0)( storedTrace, `MLflow could not retrieve trace ${traceId}`, ).toBeDefined(); + const traceRecord = (storedTrace?.trace ?? storedTrace) as { + info?: { trace_id?: string; traceId?: string }; + data?: { spans?: Array> }; + }; + expect(traceRecord.info?.trace_id ?? traceRecord.info?.traceId).toBe( + traceId, + ); const statement = "SELECT trace_id, span_id, parent_span_id, name, attributes\n" + "FROM IDENTIFIER(:otel_spans_table)\n" + @@ -387,33 +1163,48 @@ test.skipIf(missingDeployed.length > 0)( { name: "otel_spans_table", type: "STRING", - value: process.env.APPKIT_TRACE_CONFORMANCE_OTEL_SPANS_TABLE, + value: binding.spansTable, }, - { name: "trace_id", type: "STRING", value: traceId }, + { name: "trace_id", type: "STRING", value: otelTraceId }, ], }), ], { encoding: "utf8" }, ), ); - expect(sql.status?.state).toBe("SUCCEEDED"); - expect(sql.result?.data_array?.length).toBeGreaterThan(0); - const traceRecord = (storedTrace?.trace ?? storedTrace) as { - info?: { trace_id?: string; traceId?: string }; - data?: { spans?: Array> }; - }; + const completedSql = await pollStatement( + sql, + (statementId) => + JSON.parse( + execFileSync( + "databricks", + [ + "api", + "get", + `/api/2.0/sql/statements/${encodeURIComponent(statementId)}`, + "-p", + profile, + ], + { encoding: "utf8" }, + ), + ), + () => new Promise((resolve) => setTimeout(resolve, 1_000)), + ); + expect(completedSql.result?.data_array?.length).toBeGreaterThan(0); const storedSpans = traceRecord.data?.spans ?? []; - const rows = (sql.result.data_array as unknown[][]).map((row) => ({ + const rows = (completedSql.result?.data_array ?? []).map((row) => ({ traceId: row[0], spanId: row[1], parentSpanId: row[2], name: row[3], attributes: objectValue(row[4]), })); - expect(new Set(rows.map((row) => row.traceId))).toEqual(new Set([traceId])); + expect(new Set(rows.map((row) => row.traceId))).toEqual( + new Set([otelTraceId]), + ); const persistedManifest: TraceManifest = { - template: "appkit-agents", - traceId: String(traceRecord.info?.trace_id ?? traceRecord.info?.traceId), + template: appName, + traceId: otelTraceId, spans: rows.map((row) => { const source = storedSpans.find( (span) => (span.span_id ?? span.spanId) === row.spanId, @@ -423,6 +1214,24 @@ test.skipIf(missingDeployed.length > 0)( `UC span ${String(row.spanId)} missing from MLflow`, ).toBeDefined(); const attributes = row.attributes; + attributes.app_id ??= + attributes["app.id"] ?? attributes["appkit.app.name"]; + attributes.user_id ??= + attributes["user.id"] ?? attributes["mlflow.trace.user"]; + attributes.session_id ??= + attributes["session.id"] ?? attributes["mlflow.trace.session"]; + attributes.ttft_ms ??= + attributes["appkit.ttft_ms"] ?? + attributes["appkit.first_token.duration_ms"]; + attributes.stream_duration_ms ??= + attributes["appkit.stream_duration_ms"] ?? + attributes["appkit.stream.duration_ms"]; + if ( + attributes.ttft_ms !== undefined || + attributes.stream_duration_ms !== undefined + ) { + attributes.streaming = true; + } const spanType = String(attributes["mlflow.spanType"] ?? ""); const usage = objectValue( attributes[ @@ -438,9 +1247,11 @@ test.skipIf(missingDeployed.length > 0)( parentSpanId: row.parentSpanId ? String(row.parentSpanId) : null, inputs: decoded(attributes["mlflow.spanInputs"]), outputs: decoded(attributes["mlflow.spanOutputs"]), - status: Number( - (source?.status as { code?: number } | undefined)?.code ?? 0, - ), + status: + (source?.status as { code?: string | number } | undefined)?.code ?? + (source?.status as { status_code?: string | number } | undefined) + ?.status_code ?? + String(attributes["mlflow.spanStatus"] ?? "UNSET"), latencyMs: Number(source?.latency_ms ?? source?.latencyMs ?? -1), model: attributes["mlflow.chat.model"] as string | undefined, provider: attributes["mlflow.chat.provider"] as string | undefined, @@ -452,6 +1263,10 @@ test.skipIf(missingDeployed.length > 0)( }; }), }; + const persistedRoot = persistedManifest.spans.find( + (span) => span.parentSpanId === null, + ); + expect(persistedRoot?.attributes.app_id).toBe(appName); assertContract(persistedManifest); }, 180_000, diff --git a/tools/generate-app-templates.ts b/tools/generate-app-templates.ts index 97a43eac2..294c2bf13 100644 --- a/tools/generate-app-templates.ts +++ b/tools/generate-app-templates.ts @@ -43,7 +43,7 @@ const APPKIT_SECTION_END = ""; // variants must never resolve to the pre-feature 0.58.x runtime. The templates // remain publication-gated until this version exists in the registry; local // verification uses packed prerelease tarballs from this checkout. -const FIRST_MLFLOW_UC_APPKIT_VERSION = "0.59.0"; +const FIRST_MLFLOW_UC_APPKIT_VERSION = "0.60.0"; interface AppTemplate { /** Output directory name and --name passed to databricks apps init */ @@ -240,8 +240,9 @@ function postProcess(appDir: string, app: AppTemplate): void { writeFileSync(join(appDir, "databricks.yml.tmpl"), databricksYmlTmpl); // Agent templates depend on the tracing runtime added after 0.58.0. Pin the - // first containing release and remove the stale source lock, which cannot be - // regenerated against 0.59.0 until that release is published. + // first containing release and remove the stale source lock. Published + // 0.59.0 predates this tracing stack, so templates remain intentionally + // installation-gated until 0.60.0 is released. if (app.features.includes("agents")) { const packageJsonPath = join(appDir, "package.json"); const packageJson = JSON.parse(readFileSync(packageJsonPath, "utf8")) as { diff --git a/tools/tests/agent-template-policy.test.ts b/tools/tests/agent-template-policy.test.ts index 71035620a..c73183128 100644 --- a/tools/tests/agent-template-policy.test.ts +++ b/tools/tests/agent-template-policy.test.ts @@ -1,10 +1,12 @@ import { execFileSync } from "node:child_process"; import { existsSync, + mkdirSync, mkdtempSync, readdirSync, readFileSync, rmSync, + writeFileSync, } from "node:fs"; import { tmpdir } from "node:os"; import { join, resolve } from "node:path"; @@ -27,14 +29,39 @@ interface GeneratedCandidate { directory: string; } -function discoverGeneratedAgentTemplates(): GeneratedCandidate[] { - return readdirSync(output, { withFileTypes: true }) +function discoverGeneratedAgentTemplates( + searchRoot = output, +): GeneratedCandidate[] { + const behaviorSignals = [ + /\bAgentServer\b/, + /\b(?:createAgent|agents)\s*\(/, + /agents:\s*\{/, + /\/(?:invocations|responses|api\/agents)\b/, + /(?:for\s+await|while\s*\()[\s\S]*?\bmodel\b[\s\S]*?\b(?:tool|executeTool)\b/i, + /\b(?:retriev|vectorSearch)\w*[\s\S]*?\b(?:generat|model)\w*/i, + ]; + return readdirSync(searchRoot, { withFileTypes: true }) .filter((entry) => entry.isDirectory()) - .map((entry) => ({ name: entry.name, directory: join(output, entry.name) })) + .map((entry) => ({ + name: entry.name, + directory: join(searchRoot, entry.name), + })) .filter(({ directory }) => { - const server = readFileSync(join(directory, "server/server.ts"), "utf8"); - return /\bagents\s*\(/.test(server) || /agents:\s*\{/.test(server); - }); + const sources = readdirSync(directory, { + recursive: true, + encoding: "utf8", + }) + .filter( + (relative) => + !relative.includes("node_modules/") && + /\.(?:ts|tsx|js|jsx|py)$/.test(relative), + ) + .map((relative) => readFileSync(join(directory, relative), "utf8")); + return behaviorSignals.some((signal) => + sources.some((source) => signal.test(source)), + ); + }) + .sort((left, right) => left.name.localeCompare(right.name)); } describe("behavior-discovered generated agent template policy", () => { @@ -66,6 +93,42 @@ describe("behavior-discovered generated agent template policy", () => { ]); }); + test("admits arbitrary candidates from every behavior signal even without proof", () => { + const fixtures = mkdtempSync(join(tmpdir(), "appkit-agent-signals-")); + const signals = new Map([ + ["odd-server", "const app = new AgentServer();"], + ["unusual-constructor", "export const worker = createAgent({});"], + ["endpoint-client", 'fetch("/invocations", { method: "POST" });'], + [ + "loop-surface", + "for await (const event of model.run()) { await executeTool(event); }", + ], + [ + "rag-surface", + "const docs = await vectorSearch.query(); await model.generate(docs);", + ], + ]); + try { + for (const [name, source] of signals) { + const serverDirectory = join(fixtures, name, "server"); + mkdirSync(serverDirectory, { recursive: true }); + writeFileSync(join(serverDirectory, "server.ts"), source); + } + const plainDirectory = join(fixtures, "plain-web", "server"); + mkdirSync(plainDirectory, { recursive: true }); + writeFileSync( + join(plainDirectory, "server.ts"), + "createApp({ plugins: [] });", + ); + + expect( + discoverGeneratedAgentTemplates(fixtures).map(({ name }) => name), + ).toEqual([...signals.keys()].sort()); + } finally { + rmSync(fixtures, { recursive: true, force: true }); + } + }); + test("every discovered surface declares immutable UC trace resources", () => { for (const candidate of discoverGeneratedAgentTemplates()) { const app = yaml.load( @@ -120,7 +183,11 @@ describe("behavior-discovered generated agent template policy", () => { expect( packageJson.dependencies["@databricks/appkit"], candidate.name, - ).toBe("0.59.0"); + ).toBe("0.60.0"); + expect( + packageJson.dependencies["@mlflow/core"], + `${candidate.name} must keep AppKit as its sole tracing provider`, + ).toBeUndefined(); } }); }); diff --git a/tools/tests/generate-app-templates.test.ts b/tools/tests/generate-app-templates.test.ts index 37ef13693..e96c4a2a3 100644 --- a/tools/tests/generate-app-templates.test.ts +++ b/tools/tests/generate-app-templates.test.ts @@ -10,9 +10,18 @@ const outputDir = mkdtempSync(join(tmpdir(), "appkit-traced-templates-")); describe("generated AppKit agent templates", () => { beforeAll(() => { + const compatibleCli = "/tmp/databricks-cli-1.11.0/databricks"; execFileSync("pnpm", ["generate:app-templates"], { cwd: root, - env: { ...process.env, APP_TEMPLATES_OUTPUT_DIR: outputDir }, + env: { + ...process.env, + APP_TEMPLATES_OUTPUT_DIR: outputDir, + ...(process.env.DATABRICKS_CLI + ? {} + : existsSync(compatibleCli) + ? { DATABRICKS_CLI: compatibleCli } + : {}), + }, stdio: "pipe", }); }, 120_000); @@ -88,8 +97,9 @@ describe("generated AppKit agent templates", () => { expect(packageJson.scripts.setup).toBe( "appkit setup --write --mlflow-uc", ); - expect(packageJson.dependencies["@databricks/appkit"]).toBe("0.59.0"); - expect(packageJson.dependencies["@databricks/appkit-ui"]).toBe("0.59.0"); + expect(packageJson.dependencies["@databricks/appkit"]).toBe("0.60.0"); + expect(packageJson.dependencies["@databricks/appkit-ui"]).toBe("0.60.0"); + expect(packageJson.dependencies["@mlflow/core"]).toBeUndefined(); expect(existsSync(join(app, "package-lock.json"))).toBe(false); }, ); From efc8c2b00db7001cfe52b3b5dd237bcee6feb254 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Wed, 12 Aug 2026 23:56:56 -0700 Subject: [PATCH 22/31] fix: harden generated trace conformance proof Signed-off-by: Adam Gurary --- .../trace-conformance.integration.test.ts | 756 +++++++++++++++--- .../appkit/src/plugins/agents/thread-store.ts | 6 +- template/server/server.ts | 6 +- tools/tests/agent-template-policy.test.ts | 23 +- 4 files changed, 662 insertions(+), 129 deletions(-) diff --git a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts index 9d430f03c..0c6d9f361 100644 --- a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -1,5 +1,13 @@ import { execFileSync } from "node:child_process"; -import { existsSync, mkdtempSync, rmSync } from "node:fs"; +import { + existsSync, + mkdtempSync, + readdirSync, + readFileSync, + rmSync, +} from "node:fs"; +import type { Server } from "node:http"; +import { createRequire } from "node:module"; import { join, resolve } from "node:path"; import { pathToFileURL } from "node:url"; import { context, trace } from "@opentelemetry/api"; @@ -10,9 +18,11 @@ import { type ReadableSpan, SimpleSpanProcessor, } from "@opentelemetry/sdk-trace-base"; +import getPort from "get-port"; import type { AgentAdapter, AgentInput, AgentRunContext } from "shared"; -import { afterAll, beforeAll, describe, expect, test, vi } from "vitest"; +import { afterAll, beforeAll, expect, test, vi } from "vitest"; import { z } from "zod"; +import { ServiceContext } from "../../../context/service-context"; import { createAgent } from "../../../core/agent/create-agent"; import { runAgent } from "../../../core/agent/run-agent"; import { tool } from "../../../core/agent/tools/tool"; @@ -23,6 +33,44 @@ const generatedApps = mkdtempSync( join(repositoryRoot, ".trace-conformance-generated-"), ); +interface GeneratedCandidate { + name: string; + directory: string; +} + +function discoverGeneratedAgentTemplates(): GeneratedCandidate[] { + const behaviorSignals = [ + /\bAgentServer\b/, + /\b(?:createAgent|agents)\s*\(/, + /agents:\s*\{/, + /\/(?:invocations|responses|api\/agents)\b/, + /(?:for\s+await|while\s*\()[\s\S]*?\bmodel\b[\s\S]*?\b(?:tool|executeTool)\b/i, + /\b(?:retriev|vectorSearch)\w*[\s\S]*?\b(?:generat|model)\w*/i, + ]; + return readdirSync(generatedApps, { withFileTypes: true }) + .filter((entry) => entry.isDirectory()) + .map((entry) => ({ + name: entry.name, + directory: join(generatedApps, entry.name), + })) + .filter(({ directory }) => { + const sources = readdirSync(directory, { + recursive: true, + encoding: "utf8", + }) + .filter( + (relative) => + !relative.includes("node_modules/") && + /\.(?:ts|tsx|js|jsx|py)$/.test(relative), + ) + .map((relative) => readFileSync(join(directory, relative), "utf8")); + return behaviorSignals.some((signal) => + sources.some((source) => signal.test(source)), + ); + }) + .sort((left, right) => left.name.localeCompare(right.name)); +} + interface SpanManifest { name: string; spanType: string; @@ -65,6 +113,9 @@ function objectValue(value: unknown): Record { function normalize(template: string, spans: ReadableSpan[]): TraceManifest { const traceIds = new Set(spans.map((span) => span.spanContext().traceId)); + const semanticSpanIds = new Set( + spans.map((span) => span.spanContext().spanId), + ); expect(traceIds.size, `${template}: mixed trace IDs`).toBe(1); return { template, @@ -101,11 +152,18 @@ function normalize(template: string, spans: ReadableSpan[]): TraceManifest { : "mlflow.chat.tokenUsage" ], ) as Record; + const spanType = String(attributes["mlflow.spanType"] ?? ""); + const recordedParentSpanId = span.parentSpanContext?.spanId; return { name: span.name, - spanType: String(attributes["mlflow.spanType"] ?? ""), + spanType, spanId: span.spanContext().spanId, - parentSpanId: span.parentSpanContext?.spanId ?? null, + parentSpanId: + spanType === "AGENT" && + recordedParentSpanId !== undefined && + !semanticSpanIds.has(recordedParentSpanId) + ? null + : (recordedParentSpanId ?? null), inputs: decoded(attributes["mlflow.spanInputs"]), outputs: decoded(attributes["mlflow.spanOutputs"]), status: span.status.code, @@ -463,8 +521,22 @@ async function captureTurn(template: string): Promise { } } -async function captureGeneratedTurn(template: string): Promise { - const helperPath = join(generatedApps, template, "server/agents/helper.ts"); +async function closeServer(server: Server | undefined): Promise { + if (!server?.listening) return; + server.closeAllConnections?.(); + await new Promise((resolveClose, rejectClose) => { + server.close((error) => { + if (error) rejectClose(error); + else resolveClose(); + }); + }); +} + +async function captureGeneratedHttpTurn( + candidate: GeneratedCandidate, +): Promise { + const { name: template, directory } = candidate; + const helperPath = join(directory, "server/agents/helper.ts"); const generated = (await import(pathToFileURL(helperPath).href)) as { helper: AgentDefinition; }; @@ -510,18 +582,121 @@ async function captureGeneratedTurn(template: string): Promise { const getTracer = vi .spyOn(trace, "getTracer") .mockImplementation((name, version) => provider.getTracer(name, version)); + const port = await getPort(); + const environment = { + NODE_ENV: "production", + DISABLE_APPKIT_INTERNAL_TELEMETRY: "true", + DATABRICKS_APP_PORT: String(port), + FLASK_RUN_HOST: "127.0.0.1", + DATABRICKS_APP_NAME: template, + DATABRICKS_HOST: "https://test.databricks.com", + DATABRICKS_CLIENT_ID: "test-client-id", + DATABRICKS_AGENT_SERVING_ENDPOINT_NAME: "generated-test-model", + DATABRICKS_WAREHOUSE_ID: "test-warehouse", + DATABRICKS_VOLUME_FILES: "/Volumes/main/default/files", + DATABRICKS_GENIE_SPACE_ID: "test-genie-space", + DATABRICKS_SERVING_ENDPOINT_NAME: "test-serving-endpoint", + LAKEBASE_ENDPOINT: "test-lakebase-endpoint", + PGUSER: "test-client-id", + PGHOST: "localhost", + PGDATABASE: "appkit", + PGPORT: "5432", + PGSSLMODE: "require", + MLFLOW_EXPERIMENT_ID: "test-experiment", + MLFLOW_TRACING_SQL_WAREHOUSE_ID: "test-warehouse", + MLFLOW_UC_CATALOG: "main", + MLFLOW_UC_SCHEMA: "agent_traces", + MLFLOW_UC_TABLE_PREFIX: "appkit", + MLFLOW_OTEL_SPANS_TABLE: "main.agent_traces.appkit_otel_spans", + }; + const previousEnvironment = new Map( + Object.keys(environment).map((name) => [name, process.env[name]]), + ); + Object.assign(process.env, environment); + const serviceContext = { + client: {}, + serviceUserId: "test-client-id", + warehouseId: Promise.resolve("test-warehouse"), + workspaceId: Promise.resolve("test-workspace"), + }; + const initializeServiceContext = vi + .spyOn(ServiceContext, "initialize") + .mockResolvedValue(serviceContext as never); + const getServiceContext = vi + .spyOn(ServiceContext, "get") + .mockReturnValue(serviceContext as never); + let server: Server | undefined; try { - await runAgent(generated.helper, { - messages: "Count the words in hello traced world. Use count_words.", - appName: template, - requestId: "request-1", - sessionId: "session-1", - threadId: "thread-1", - userId: "user-1", - }); + const requireFromGeneratedPackage = createRequire( + join(directory, "package.json"), + ); + expect( + requireFromGeneratedPackage.resolve("@databricks/appkit/package.json"), + `${template} generated package resolution`, + ).toBe(join(repositoryRoot, "packages/appkit/package.json")); + + const serverPath = join(directory, "server/server.ts"); + const generatedServer = (await import(pathToFileURL(serverPath).href)) as { + app?: Promise<{ server: { getServer(): Server } }>; + }; + expect( + generatedServer.app, + `${template} must export its generated createApp execution`, + ).toBeDefined(); + const appkit = await generatedServer.app; + server = appkit?.server.getServer(); + if (!server) { + throw new Error( + `${template} generated createApp did not expose its server`, + ); + } + if (!server.listening) { + await new Promise((resolveListening, rejectListening) => { + server?.once("listening", resolveListening); + server?.once("error", rejectListening); + }); + } + const address = server.address(); + if (!address || typeof address === "string") { + throw new Error(`${template} generated HTTP server did not bind a port`); + } + const response = await fetch( + `http://127.0.0.1:${address.port}/invocations`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + "x-forwarded-user": "user-1", + "x-mlflow-session-id": "session-1", + "x-request-id": "request-1", + }, + body: JSON.stringify({ + input: "Count the words in hello traced world. Use count_words.", + }), + }, + ); + const responseBody = await response.text(); + expect( + response.ok, + `${template} HTTP ${response.status}: ${responseBody}`, + ).toBe(true); + expect( + response.headers.get("x-mlflow-trace-id"), + `${template} generated handler trace identity`, + ).toBeTruthy(); await provider.forceFlush(); - return normalize(template, exporter.getFinishedSpans()); + const semanticSpans = exporter + .getFinishedSpans() + .filter((span) => span.attributes["mlflow.spanType"] !== undefined); + return normalize(template, semanticSpans); } finally { + await closeServer(server); + initializeServiceContext.mockRestore(); + getServiceContext.mockRestore(); + for (const [name, previous] of previousEnvironment) { + if (previous === undefined) delete process.env[name]; + else process.env[name] = previous; + } getTracer.mockRestore(); await provider.shutdown(); } @@ -553,22 +728,19 @@ afterAll(() => { rmSync(generatedApps, { recursive: true, force: true }); }); -describe.each(["appkit-agents", "appkit-all-in-one"])( - "%s trace conformance", - (template) => { - test("writes, reloads, and validates a deterministic tool turn", async () => { - const manifest = await captureTurn(template); - const reloaded = JSON.parse(JSON.stringify(manifest)) as TraceManifest; +test("every behavior-discovered generated surface executes its HTTP trace proof", async () => { + const candidates = discoverGeneratedAgentTemplates(); + expect(candidates.length).toBeGreaterThan(0); + const failures: string[] = []; + for (const candidate of candidates) { + try { + const localManifest = await captureTurn(candidate.name); + const reloaded = JSON.parse( + JSON.stringify(localManifest), + ) as TraceManifest; assertContract(reloaded); - }); - }, -); -describe.each(["appkit-agents", "appkit-all-in-one"])( - "%s generated production package", - (template) => { - test("invokes its own generated agent definition and validates the trace", async () => { - const manifest = await captureGeneratedTurn(template); + const manifest = await captureGeneratedHttpTurn(candidate); expect( manifest.spans.some( (span) => @@ -579,9 +751,17 @@ describe.each(["appkit-agents", "appkit-all-in-one"])( ), ).toBe(true); assertContract(manifest); - }); - }, -); + } catch (error) { + failures.push( + `${candidate.name}: ${error instanceof Error ? error.message : String(error)}`, + ); + } + } + expect( + failures, + `discovered candidates without executable generated HTTP proof:\n${failures.join("\n")}`, + ).toEqual([]); +}, 120_000); function requiredSpan(manifest: TraceManifest, spanType: string): SpanManifest { const span = manifest.spans.find( @@ -893,6 +1073,13 @@ interface StatementResponse { result?: { data_array?: unknown[][] }; } +const persistedSpanStatement = + "SELECT trace_id, span_id, parent_span_id, name, attributes, " + + "status_code, start_time_unix_nano, end_time_unix_nano\n" + + "FROM IDENTIFIER(:otel_spans_table)\n" + + "WHERE trace_id = :trace_id\n" + + "ORDER BY start_time_unix_nano"; + function deriveUcBinding( catalog: string, schema: string, @@ -936,6 +1123,299 @@ function otelTraceIdFromReturnedTrace( return otelTraceId.toLowerCase(); } +interface UcSpanRow { + traceId: unknown; + spanId: unknown; + parentSpanId: unknown; + name: unknown; + attributes: Record; + statusCode?: unknown; + startTimeUnixNano?: unknown; + endTimeUnixNano?: unknown; +} + +interface DeployedTraceProof { + appName: string; + configuredExperimentId: string; + returnedTraceId: string; + binding: ReturnType; + experiment: { + experiment?: { + experiment_id?: string; + trace_location?: Record; + }; + }; + traceRecord: { + info?: { + trace_id?: string; + traceId?: string; + experiment_id?: string; + experimentId?: string; + }; + data?: { spans?: Array> }; + }; + rows: UcSpanRow[]; +} + +function deployedProofFixture(): DeployedTraceProof { + const appName = "appkit-agents"; + const configuredExperimentId = "experiment-123"; + const binding = deriveUcBinding("main", "agent_traces", "appkit"); + const manifest = validWorkloads()[1]; + manifest.template = appName; + manifest.spans[0].attributes.app_id = appName; + const returnedTraceId = `${binding.mlflowTracePrefix}${manifest.traceId}`; + const rows = manifest.spans.map((span, index) => ({ + traceId: manifest.traceId, + spanId: span.spanId, + parentSpanId: span.parentSpanId, + name: span.name, + attributes: { + ...span.attributes, + "mlflow.spanType": span.spanType, + "mlflow.spanInputs": span.inputs, + "mlflow.spanOutputs": span.outputs, + [span.spanType === "AGENT" + ? "mlflow.trace.tokenUsage" + : "mlflow.chat.tokenUsage"]: span.usage, + "mlflow.chat.model": span.model, + "mlflow.chat.provider": span.provider, + "mlflow.llm.cost": span.costUsd, + "appkit.cost.available": span.costAvailable, + "appkit.app.name": span.attributes.app_id, + "mlflow.trace.user": span.attributes.user_id, + "mlflow.trace.session": span.attributes.session_id, + }, + statusCode: "OK", + startTimeUnixNano: String(1_000_000 + index * 2_000_000), + endTimeUnixNano: String(2_000_000 + index * 2_000_000), + })); + return { + appName, + configuredExperimentId, + returnedTraceId, + binding, + experiment: { + experiment: { + experiment_id: configuredExperimentId, + trace_location: binding.location, + }, + }, + traceRecord: { + info: { + trace_id: returnedTraceId, + experiment_id: configuredExperimentId, + }, + data: { + spans: rows.map((row) => ({ + span_id: row.spanId, + parent_span_id: row.parentSpanId, + attributes: row.attributes, + status: { code: "OK" }, + latency_ms: 1, + })), + }, + }, + rows, + }; +} + +function normalizeUcStatus(value: unknown): string | number { + if (typeof value === "number") return value; + if (typeof value !== "string" || !value.trim()) return "UNSET"; + return value + .trim() + .toUpperCase() + .replace(/^STATUS_CODE_/, ""); +} + +function unixNanos(value: unknown): bigint | undefined { + if (typeof value === "bigint") return value; + if (typeof value === "number" && Number.isSafeInteger(value)) { + return BigInt(value); + } + if (typeof value === "string" && /^\d+$/.test(value)) { + return BigInt(value); + } + return undefined; +} + +function ucLatencyMs(row: UcSpanRow): number { + const start = unixNanos(row.startTimeUnixNano); + const end = unixNanos(row.endTimeUnixNano); + if (start === undefined || end === undefined || end < start) return -1; + return Number(end - start) / 1_000_000; +} + +function assertAppIdentity( + attributes: Record, + expected: string, + source: string, +): void { + const values = [ + attributes.app_id, + attributes["app.id"], + attributes["appkit.app.name"], + ].filter((value) => value !== undefined); + if (values.length === 0 || values.some((value) => value !== expected)) { + throw new Error(`${source} app identity does not match configured app`); + } +} + +function normalizeUcRows( + template: string, + traceId: string, + rows: UcSpanRow[], +): TraceManifest { + const semanticRows = rows.filter( + (row) => row.attributes["mlflow.spanType"] !== undefined, + ); + const semanticSpanIds = new Set( + semanticRows.map((row) => String(row.spanId)), + ); + return { + template, + traceId, + spans: semanticRows.map((row) => { + const attributes = { ...row.attributes }; + attributes.app_id ??= + attributes["app.id"] ?? attributes["appkit.app.name"]; + attributes.user_id ??= + attributes["user.id"] ?? attributes["mlflow.trace.user"]; + attributes.session_id ??= + attributes["session.id"] ?? attributes["mlflow.trace.session"]; + attributes.ttft_ms ??= + attributes["appkit.ttft_ms"] ?? + attributes["appkit.first_token.duration_ms"]; + attributes.stream_duration_ms ??= + attributes["appkit.stream_duration_ms"] ?? + attributes["appkit.stream.duration_ms"]; + if ( + attributes.ttft_ms !== undefined || + attributes.stream_duration_ms !== undefined + ) { + attributes.streaming = true; + } + const spanType = String(attributes["mlflow.spanType"] ?? ""); + const usage = objectValue( + attributes[ + spanType === "AGENT" + ? "mlflow.trace.tokenUsage" + : "mlflow.chat.tokenUsage" + ], + ) as Record; + const rawParentSpanId = row.parentSpanId + ? String(row.parentSpanId) + : null; + return { + name: String(row.name), + spanType, + spanId: String(row.spanId), + parentSpanId: + spanType === "AGENT" && + rawParentSpanId !== null && + !semanticSpanIds.has(rawParentSpanId) + ? null + : rawParentSpanId, + inputs: decoded(attributes["mlflow.spanInputs"]), + outputs: decoded(attributes["mlflow.spanOutputs"]), + status: normalizeUcStatus(row.statusCode), + latencyMs: ucLatencyMs(row), + model: attributes["mlflow.chat.model"] as string | undefined, + provider: attributes["mlflow.chat.provider"] as string | undefined, + usage, + costUsd: attributes["mlflow.llm.cost"] as number | undefined, + costAvailable: attributes["appkit.cost.available"] === true, + links: [], + attributes, + }; + }), + }; +} + +function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { + const experiment = proof.experiment.experiment; + if (experiment?.experiment_id !== proof.configuredExperimentId) { + throw new Error( + `experiment API returned ${String(experiment?.experiment_id)} for configured experiment ${proof.configuredExperimentId}`, + ); + } + if ( + JSON.stringify(experiment.trace_location) !== + JSON.stringify(proof.binding.location) + ) { + throw new Error( + "experiment trace location does not match configured UC binding", + ); + } + const traceInfo = proof.traceRecord.info; + const storedTraceId = traceInfo?.trace_id ?? traceInfo?.traceId; + if (storedTraceId !== proof.returnedTraceId) { + throw new Error( + `MLflow trace ${String(storedTraceId)} does not match returned trace ${proof.returnedTraceId}`, + ); + } + const storedExperimentId = + traceInfo?.experiment_id ?? traceInfo?.experimentId; + if (storedExperimentId !== proof.configuredExperimentId) { + throw new Error( + `MLflow trace experiment ${String(storedExperimentId)} does not match configured experiment ${proof.configuredExperimentId}`, + ); + } + const otelTraceId = otelTraceIdFromReturnedTrace( + proof.returnedTraceId, + proof.binding.mlflowTracePrefix, + ); + if (proof.rows.length === 0) { + throw new Error( + `UC table ${proof.binding.spansTable} returned no trace rows`, + ); + } + for (const row of proof.rows) { + if (String(row.traceId).toLowerCase() !== otelTraceId) { + throw new Error( + `UC row trace ${String(row.traceId)} does not match returned OTel trace ${otelTraceId}`, + ); + } + } + + const mlflowSpans = proof.traceRecord.data?.spans ?? []; + const mlflowSpanIds = new Set( + mlflowSpans.map((span) => String(span.span_id ?? span.spanId)), + ); + const semanticRows = proof.rows.filter( + (row) => row.attributes["mlflow.spanType"] !== undefined, + ); + for (const row of semanticRows) { + if (!mlflowSpanIds.has(String(row.spanId))) { + throw new Error( + `UC span ${String(row.spanId)} is not associated with the returned MLflow trace`, + ); + } + } + + const mlflowRoot = mlflowSpans.find((span) => { + const attributes = objectValue(span.attributes); + return attributes["mlflow.spanType"] === "AGENT"; + }); + assertAppIdentity( + objectValue(mlflowRoot?.attributes), + proof.appName, + "MLflow trace", + ); + const ucRoot = semanticRows.find( + (row) => row.attributes["mlflow.spanType"] === "AGENT", + ); + if (!ucRoot) { + throw new Error("UC trace app identity does not match configured app"); + } + assertAppIdentity(ucRoot.attributes, proof.appName, "UC trace"); + + const manifest = normalizeUcRows(proof.appName, otelTraceId, proof.rows); + assertContract(manifest); + return manifest; +} + async function pollStatement( initial: StatementResponse, getStatement: (statementId: string) => StatementResponse, @@ -976,6 +1456,16 @@ test("derives the exact immutable UC location, table, and MLflow trace prefix", }); }); +test("queries UC-native status and timing for the exact returned trace", () => { + expect(persistedSpanStatement).toBe( + "SELECT trace_id, span_id, parent_span_id, name, attributes, " + + "status_code, start_time_unix_nano, end_time_unix_nano\n" + + "FROM IDENTIFIER(:otel_spans_table)\n" + + "WHERE trace_id = :trace_id\n" + + "ORDER BY start_time_unix_nano", + ); +}); + test("extracts the OTel trace ID only from the exact returned UC trace identity", () => { expect( otelTraceIdFromReturnedTrace( @@ -991,6 +1481,87 @@ test("extracts the OTel trace ID only from the exact returned UC trace identity" ).toThrow(/exact UC location/); }); +test.each([ + [ + "experiment API", + (proof: DeployedTraceProof) => { + if (proof.experiment.experiment) { + proof.experiment.experiment.experiment_id = "wrong-experiment"; + } + }, + ], + [ + "MLflow trace", + (proof: DeployedTraceProof) => { + if (proof.traceRecord.info) { + proof.traceRecord.info.experiment_id = "wrong-experiment"; + } + }, + ], +] as const)( + "rejects a wrong configured experiment association from %s despite the matching UC location", + (_name, mutate) => { + const proof = deployedProofFixture(); + mutate(proof); + expect(() => validateDeployedProof(proof)).toThrow(/experiment/i); + }, +); + +test.each([ + [ + "returned trace", + (proof: DeployedTraceProof) => { + if (proof.traceRecord.info) { + proof.traceRecord.info.trace_id = `${proof.binding.mlflowTracePrefix}fedcba9876543210fedcba9876543210`; + } + }, + ], + [ + "MLflow app", + (proof: DeployedTraceProof) => { + const root = proof.traceRecord.data?.spans?.find( + (span) => !span.parent_span_id, + ); + if (root) objectValue(root.attributes)["appkit.app.name"] = "other-app"; + }, + ], + [ + "UC app", + (proof: DeployedTraceProof) => { + const root = proof.rows.find((row) => !row.parentSpanId); + if (root) root.attributes["appkit.app.name"] = "other-app"; + }, + ], +] as const)("rejects a wrong %s association", (_name, mutate) => { + const proof = deployedProofFixture(); + mutate(proof); + expect(() => validateDeployedProof(proof)).toThrow(/trace|app/i); +}); + +test.each([ + ["missing UC status", (row: UcSpanRow) => delete row.statusCode], + [ + "nonterminal UC status", + (row: UcSpanRow) => { + row.statusCode = "UNSET"; + }, + ], + ["missing UC end time", (row: UcSpanRow) => delete row.endTimeUnixNano], + [ + "negative UC duration", + (row: UcSpanRow) => { + row.endTimeUnixNano = "0"; + }, + ], +] as const)( + "rejects %s without borrowing correct MLflow lifecycle values", + (_name, mutate) => { + const proof = deployedProofFixture(); + mutate(proof.rows[0]); + expect(() => validateDeployedProof(proof)).toThrow(/status|latency/i); + }, +); + test("polls asynchronous Statement Execution to success", async () => { const getStatement = vi.fn(() => ({ statement_id: "statement-1", @@ -1059,7 +1630,11 @@ test.skipIf(missingDeployed.length > 0)( async () => { const profile = process.env.APPKIT_TRACE_CONFORMANCE_PROFILE ?? ""; const appName = process.env.APPKIT_TRACE_CONFORMANCE_APP_NAME ?? ""; - expect(["appkit-agents", "appkit-all-in-one"]).toContain(appName); + expect(discoverGeneratedAgentTemplates().map(({ name }) => name)).toContain( + appName, + ); + const configuredExperimentId = + process.env.APPKIT_TRACE_CONFORMANCE_EXPERIMENT_ID ?? ""; const binding = deriveUcBinding( process.env.APPKIT_TRACE_CONFORMANCE_UC_CATALOG ?? "", process.env.APPKIT_TRACE_CONFORMANCE_UC_SCHEMA ?? "", @@ -1099,16 +1674,13 @@ test.skipIf(missingDeployed.length > 0)( [ "api", "get", - `/api/2.0/mlflow/experiments/get?experiment_id=${encodeURIComponent(process.env.APPKIT_TRACE_CONFORMANCE_EXPERIMENT_ID ?? "")}`, + `/api/2.0/mlflow/experiments/get?experiment_id=${encodeURIComponent(configuredExperimentId)}`, "-p", profile, ], { encoding: "utf8" }, ), ); - const location = experiment.experiment?.trace_location; - expect(location).toEqual(binding.location); - let storedTrace: Record | undefined; for (let attempt = 0; attempt < 15 && !storedTrace; attempt += 1) { try { @@ -1133,18 +1705,8 @@ test.skipIf(missingDeployed.length > 0)( storedTrace, `MLflow could not retrieve trace ${traceId}`, ).toBeDefined(); - const traceRecord = (storedTrace?.trace ?? storedTrace) as { - info?: { trace_id?: string; traceId?: string }; - data?: { spans?: Array> }; - }; - expect(traceRecord.info?.trace_id ?? traceRecord.info?.traceId).toBe( - traceId, - ); - const statement = - "SELECT trace_id, span_id, parent_span_id, name, attributes\n" + - "FROM IDENTIFIER(:otel_spans_table)\n" + - "WHERE trace_id = :trace_id\n" + - "ORDER BY start_time_unix_nano"; + const traceRecord = (storedTrace?.trace ?? + storedTrace) as DeployedTraceProof["traceRecord"]; const sql = JSON.parse( execFileSync( "databricks", @@ -1157,7 +1719,7 @@ test.skipIf(missingDeployed.length > 0)( "--json", JSON.stringify({ warehouse_id: process.env.APPKIT_TRACE_CONFORMANCE_WAREHOUSE_ID, - statement, + statement: persistedSpanStatement, wait_timeout: "50s", parameters: [ { @@ -1191,83 +1753,27 @@ test.skipIf(missingDeployed.length > 0)( () => new Promise((resolve) => setTimeout(resolve, 1_000)), ); expect(completedSql.result?.data_array?.length).toBeGreaterThan(0); - const storedSpans = traceRecord.data?.spans ?? []; - const rows = (completedSql.result?.data_array ?? []).map((row) => ({ - traceId: row[0], - spanId: row[1], - parentSpanId: row[2], - name: row[3], - attributes: objectValue(row[4]), - })); - expect(new Set(rows.map((row) => row.traceId))).toEqual( - new Set([otelTraceId]), - ); - const persistedManifest: TraceManifest = { - template: appName, - traceId: otelTraceId, - spans: rows.map((row) => { - const source = storedSpans.find( - (span) => (span.span_id ?? span.spanId) === row.spanId, - ); - expect( - source, - `UC span ${String(row.spanId)} missing from MLflow`, - ).toBeDefined(); - const attributes = row.attributes; - attributes.app_id ??= - attributes["app.id"] ?? attributes["appkit.app.name"]; - attributes.user_id ??= - attributes["user.id"] ?? attributes["mlflow.trace.user"]; - attributes.session_id ??= - attributes["session.id"] ?? attributes["mlflow.trace.session"]; - attributes.ttft_ms ??= - attributes["appkit.ttft_ms"] ?? - attributes["appkit.first_token.duration_ms"]; - attributes.stream_duration_ms ??= - attributes["appkit.stream_duration_ms"] ?? - attributes["appkit.stream.duration_ms"]; - if ( - attributes.ttft_ms !== undefined || - attributes.stream_duration_ms !== undefined - ) { - attributes.streaming = true; - } - const spanType = String(attributes["mlflow.spanType"] ?? ""); - const usage = objectValue( - attributes[ - spanType === "AGENT" - ? "mlflow.trace.tokenUsage" - : "mlflow.chat.tokenUsage" - ], - ) as Record; - return { - name: String(row.name), - spanType, - spanId: String(row.spanId), - parentSpanId: row.parentSpanId ? String(row.parentSpanId) : null, - inputs: decoded(attributes["mlflow.spanInputs"]), - outputs: decoded(attributes["mlflow.spanOutputs"]), - status: - (source?.status as { code?: string | number } | undefined)?.code ?? - (source?.status as { status_code?: string | number } | undefined) - ?.status_code ?? - String(attributes["mlflow.spanStatus"] ?? "UNSET"), - latencyMs: Number(source?.latency_ms ?? source?.latencyMs ?? -1), - model: attributes["mlflow.chat.model"] as string | undefined, - provider: attributes["mlflow.chat.provider"] as string | undefined, - usage, - costUsd: attributes["mlflow.llm.cost"] as number | undefined, - costAvailable: attributes["appkit.cost.available"] === true, - links: [], - attributes, - }; + const rows: UcSpanRow[] = (completedSql.result?.data_array ?? []).map( + (row) => ({ + traceId: row[0], + spanId: row[1], + parentSpanId: row[2], + name: row[3], + attributes: objectValue(row[4]), + statusCode: row[5], + startTimeUnixNano: row[6], + endTimeUnixNano: row[7], }), - }; - const persistedRoot = persistedManifest.spans.find( - (span) => span.parentSpanId === null, ); - expect(persistedRoot?.attributes.app_id).toBe(appName); - assertContract(persistedManifest); + validateDeployedProof({ + appName, + configuredExperimentId, + returnedTraceId: traceId ?? "", + binding, + experiment, + traceRecord, + rows, + }); }, 180_000, ); diff --git a/packages/appkit/src/plugins/agents/thread-store.ts b/packages/appkit/src/plugins/agents/thread-store.ts index d8c6a905b..ac9a6a162 100644 --- a/packages/appkit/src/plugins/agents/thread-store.ts +++ b/packages/appkit/src/plugins/agents/thread-store.ts @@ -28,7 +28,11 @@ async function traceMemoryOperation( setCapturedAttribute(span, "mlflow.spanInputs", inputs); try { const result = await operation(span); - setCapturedAttribute(span, "mlflow.spanOutputs", result); + setCapturedAttribute( + span, + "mlflow.spanOutputs", + result === undefined ? { completed: true } : result, + ); span.setStatus({ code: SpanStatusCode.OK }); return result; } catch (error) { diff --git a/template/server/server.ts b/template/server/server.ts index 357b8a407..5313672e4 100644 --- a/template/server/server.ts +++ b/template/server/server.ts @@ -19,7 +19,7 @@ import { setupSampleLakebaseRoutes } from './routes/lakebase/todo-routes'; import { helper } from './agents/helper'; {{- end}} -createApp({ +export const app = createApp({ {{- if .plugins.agents}} // Uses AppKit's existing TelemetryManager / OTel provider. The setup // command provisions the immutable UC trace location before first run. @@ -42,4 +42,6 @@ createApp({ await setupSampleLakebaseRoutes(appkit); }, {{- end}} -}).catch(console.error); +}); + +app.catch(console.error); diff --git a/tools/tests/agent-template-policy.test.ts b/tools/tests/agent-template-policy.test.ts index c73183128..7cb30384e 100644 --- a/tools/tests/agent-template-policy.test.ts +++ b/tools/tests/agent-template-policy.test.ts @@ -8,13 +8,14 @@ import { rmSync, writeFileSync, } from "node:fs"; +import { createRequire } from "node:module"; import { tmpdir } from "node:os"; import { join, resolve } from "node:path"; import yaml from "js-yaml"; import { afterAll, beforeAll, describe, expect, test } from "vitest"; const root = resolve(import.meta.dirname, "../.."); -const output = mkdtempSync(join(tmpdir(), "appkit-agent-policy-")); +const output = mkdtempSync(join(root, ".appkit-agent-policy-")); const requiredTraceEnvironment = [ "MLFLOW_EXPERIMENT_ID", "MLFLOW_TRACING_SQL_WAREHOUSE_ID", @@ -190,4 +191,24 @@ describe("behavior-discovered generated agent template policy", () => { ).toBeUndefined(); } }); + + test("every discovered surface exposes its generated server as executable proof", () => { + for (const candidate of discoverGeneratedAgentTemplates()) { + const serverSource = readFileSync( + join(candidate.directory, "server/server.ts"), + "utf8", + ); + expect(serverSource, candidate.name).toContain( + "export const app = createApp({", + ); + + const requireFromGeneratedPackage = createRequire( + join(candidate.directory, "package.json"), + ); + expect( + requireFromGeneratedPackage.resolve("@databricks/appkit/package.json"), + `${candidate.name} must execute against its generated package resolution`, + ).toBe(join(root, "packages/appkit/package.json")); + } + }); }); From dab18fe3aa8d2750a1523f217e635e14c1ab4533 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 00:43:30 -0700 Subject: [PATCH 23/31] fix: strengthen deployed trace conformance Signed-off-by: Adam Gurary --- packages/appkit/src/core/appkit.ts | 21 +- packages/appkit/src/core/lifecycle-manager.ts | 51 ++- .../src/core/tests/lifecycle-manager.test.ts | 15 + .../trace-conformance.integration.test.ts | 290 +++++++++++++++++- 4 files changed, 342 insertions(+), 35 deletions(-) diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 051e84f1e..b08b556bc 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -26,10 +26,17 @@ import { isToolProvider, PluginContext } from "./plugin-context"; const logger = createLogger("appkit"); +type AppKitHandle< + TPlugins extends readonly PluginData[], +> = PluginMap & { + shutdown(): Promise; +}; + export class AppKit { #pluginInstances: Record = {}; #setupPromises: Promise[] = []; #context: PluginContext; + #lifecycleManager?: LifecycleManager; private constructor(config: { plugins: TPlugins }) { const { plugins, ...globalConfig } = config; @@ -190,7 +197,7 @@ export class AppKit { onPluginsReady?: (appkit: PluginMap) => void | Promise; disableInternalTelemetry?: boolean; } = {}, - ): Promise> { + ): Promise> { const withDefaults = AppKit.withDefaultPlugins(config.plugins as T); const rawPlugins = AppKit.filterDevOnlyPlugins(withDefaults); const agentsEnabled = rawPlugins.some( @@ -236,7 +243,7 @@ export class AppKit { await Promise.all(instance.#setupPromises); await instance.#context.emitLifecycle("setup:complete"); - const handle = instance as unknown as PluginMap; + const handle = instance as unknown as AppKitHandle; if (config.onPluginsReady) { logger.debug("Running onPluginsReady hook"); @@ -257,11 +264,17 @@ export class AppKit { // plugin has started. Applies uniformly whether or not a server plugin // is present ā€” server-less apps still get their telemetry flushed and // plugin shutdown() hooks run. - new LifecycleManager(instance.#context).installSignalHandlers(); + instance.#lifecycleManager = new LifecycleManager(instance.#context); + instance.#lifecycleManager.installSignalHandlers(); return handle; } + /** Gracefully release plugins, servers, caches, telemetry, and signal hooks. */ + async shutdown(): Promise { + await this.#lifecycleManager?.shutdown({ exitProcess: false }); + } + private static bootstrapInternalTelemetry(): void { const serviceCtx = ServiceContext.get(); const reporter = TelemetryReporter.initialize({ @@ -396,6 +409,6 @@ export async function createApp< onPluginsReady?: (appkit: PluginMap) => void | Promise; disableInternalTelemetry?: boolean; } = {}, -): Promise> { +): Promise> { return AppKit._createApp(config); } diff --git a/packages/appkit/src/core/lifecycle-manager.ts b/packages/appkit/src/core/lifecycle-manager.ts index 8d1811203..da71edddb 100644 --- a/packages/appkit/src/core/lifecycle-manager.ts +++ b/packages/appkit/src/core/lifecycle-manager.ts @@ -56,6 +56,10 @@ export class LifecycleManager { * can say where shutdown got stuck without extra bookkeeping. */ private shutdownPhase = "not started"; + private signalHandlers?: { + SIGTERM: () => void; + SIGINT: () => void; + }; constructor(private readonly context: PluginContext) {} @@ -67,8 +71,20 @@ export class LifecycleManager { * `isShuttingDown` inside {@link shutdown}. */ installSignalHandlers(): void { - process.once("SIGTERM", () => this.shutdown()); - process.once("SIGINT", () => this.shutdown()); + if (this.signalHandlers) return; + this.signalHandlers = { + SIGTERM: () => void this.shutdown(), + SIGINT: () => void this.shutdown(), + }; + process.once("SIGTERM", this.signalHandlers.SIGTERM); + process.once("SIGINT", this.signalHandlers.SIGINT); + } + + private removeSignalHandlers(): void { + if (!this.signalHandlers) return; + process.removeListener("SIGTERM", this.signalHandlers.SIGTERM); + process.removeListener("SIGINT", this.signalHandlers.SIGINT); + this.signalHandlers = undefined; } /** @@ -86,7 +102,11 @@ export class LifecycleManager { * shutdown is not a crash. Exit 1 is reserved for an unexpected error * thrown by the sequence itself. */ - async shutdown(): Promise { + async shutdown({ + exitProcess = true, + }: { + exitProcess?: boolean; + } = {}): Promise { // Must stay synchronous and first: any await before the flag is set // would let a second signal re-enter the shutdown sequence. if (this.isShuttingDown) return; @@ -101,20 +121,22 @@ export class LifecycleManager { // shutdown, not a crash), and orchestrators record nonzero exits on // deploys as crashes. The error log below is the stuck-shutdown // signal instead of the exit code. - const forceExitTimer = setTimeout(() => { - logger.error( - "Graceful shutdown did NOT complete within the %dms budget (phase in flight: %s); force-exiting with code 0.", - LifecycleManager.SHUTDOWN_TIMEOUT_MS, - this.shutdownPhase, - ); - process.exit(0); - }, LifecycleManager.SHUTDOWN_TIMEOUT_MS); + const forceExitTimer = exitProcess + ? setTimeout(() => { + logger.error( + "Graceful shutdown did NOT complete within the %dms budget (phase in flight: %s); force-exiting with code 0.", + LifecycleManager.SHUTDOWN_TIMEOUT_MS, + this.shutdownPhase, + ); + process.exit(0); + }, LifecycleManager.SHUTDOWN_TIMEOUT_MS) + : undefined; // unref so this backstop timer never by itself keeps the process alive. // Any real pending teardown (OTEL export timer, DB pool sockets, the // still-open HTTP listener) is a ref'd handle that holds the loop open // until this fires; if nothing is ref'd, there is nothing left to tear // down and exiting early is correct. - forceExitTimer.unref(); + forceExitTimer?.unref(); try { const plugins = Array.from(this.context.getPlugins().values()); @@ -183,8 +205,9 @@ export class LifecycleManager { exitCode = 1; } - clearTimeout(forceExitTimer); - process.exit(exitCode); + if (forceExitTimer) clearTimeout(forceExitTimer); + this.removeSignalHandlers(); + if (exitProcess) process.exit(exitCode); } /** Close the cache storage, bounded and error-isolated. */ diff --git a/packages/appkit/src/core/tests/lifecycle-manager.test.ts b/packages/appkit/src/core/tests/lifecycle-manager.test.ts index 121e7eb5c..ed9d92f6e 100644 --- a/packages/appkit/src/core/tests/lifecycle-manager.test.ts +++ b/packages/appkit/src/core/tests/lifecycle-manager.test.ts @@ -380,5 +380,20 @@ describe("LifecycleManager", () => { expect(signals).toContain("SIGINT"); onceSpy.mockRestore(); }); + + test("programmatic shutdown removes installed handlers without exiting", async () => { + const baselineSigterm = process.listenerCount("SIGTERM"); + const baselineSigint = process.listenerCount("SIGINT"); + const manager = new LifecycleManager(contextWithPlugins({})); + manager.installSignalHandlers(); + expect(process.listenerCount("SIGTERM")).toBe(baselineSigterm + 1); + expect(process.listenerCount("SIGINT")).toBe(baselineSigint + 1); + + await manager.shutdown({ exitProcess: false }); + + expect(process.listenerCount("SIGTERM")).toBe(baselineSigterm); + expect(process.listenerCount("SIGINT")).toBe(baselineSigint); + expect(exitSpy).not.toHaveBeenCalled(); + }); }); }); diff --git a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts index 0c6d9f361..6f10b2206 100644 --- a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -626,6 +626,9 @@ async function captureGeneratedHttpTurn( .spyOn(ServiceContext, "get") .mockReturnValue(serviceContext as never); let server: Server | undefined; + let appkit: + | { server: { getServer(): Server }; shutdown(): Promise } + | undefined; try { const requireFromGeneratedPackage = createRequire( join(directory, "package.json"), @@ -637,13 +640,16 @@ async function captureGeneratedHttpTurn( const serverPath = join(directory, "server/server.ts"); const generatedServer = (await import(pathToFileURL(serverPath).href)) as { - app?: Promise<{ server: { getServer(): Server } }>; + app?: Promise<{ + server: { getServer(): Server }; + shutdown(): Promise; + }>; }; expect( generatedServer.app, `${template} must export its generated createApp execution`, ).toBeDefined(); - const appkit = await generatedServer.app; + appkit = await generatedServer.app; server = appkit?.server.getServer(); if (!server) { throw new Error( @@ -690,15 +696,19 @@ async function captureGeneratedHttpTurn( .filter((span) => span.attributes["mlflow.spanType"] !== undefined); return normalize(template, semanticSpans); } finally { - await closeServer(server); - initializeServiceContext.mockRestore(); - getServiceContext.mockRestore(); - for (const [name, previous] of previousEnvironment) { - if (previous === undefined) delete process.env[name]; - else process.env[name] = previous; + try { + if (appkit) await appkit.shutdown(); + else await closeServer(server); + } finally { + initializeServiceContext.mockRestore(); + getServiceContext.mockRestore(); + for (const [name, previous] of previousEnvironment) { + if (previous === undefined) delete process.env[name]; + else process.env[name] = previous; + } + getTracer.mockRestore(); + await provider.shutdown(); } - getTracer.mockRestore(); - await provider.shutdown(); } } @@ -734,6 +744,8 @@ test("every behavior-discovered generated surface executes its HTTP trace proof" const failures: string[] = []; for (const candidate of candidates) { try { + const baselineSigterm = process.listenerCount("SIGTERM"); + const baselineSigint = process.listenerCount("SIGINT"); const localManifest = await captureTurn(candidate.name); const reloaded = JSON.parse( JSON.stringify(localManifest), @@ -741,6 +753,12 @@ test("every behavior-discovered generated surface executes its HTTP trace proof" assertContract(reloaded); const manifest = await captureGeneratedHttpTurn(candidate); + expect(process.listenerCount("SIGTERM"), candidate.name).toBe( + baselineSigterm, + ); + expect(process.listenerCount("SIGINT"), candidate.name).toBe( + baselineSigint, + ); expect( manifest.spans.some( (span) => @@ -1137,6 +1155,13 @@ interface UcSpanRow { interface DeployedTraceProof { appName: string; configuredExperimentId: string; + requestBody: unknown; + responseBody: unknown; + expectedTool: { + name: string; + inputs: unknown; + outputs: unknown; + }; returnedTraceId: string; binding: ReturnType; experiment: { @@ -1165,6 +1190,33 @@ function deployedProofFixture(): DeployedTraceProof { manifest.template = appName; manifest.spans[0].attributes.app_id = appName; const returnedTraceId = `${binding.mlflowTracePrefix}${manifest.traceId}`; + const requestBody = { + input: "Count the words in hello traced world. Use count_words.", + }; + const responseBody = { + object: "response", + status: "completed", + trace_id: returnedTraceId, + output: [ + { + type: "message", + status: "completed", + content: [{ type: "output_text", text: "3" }], + }, + ], + }; + const expectedTool = { + name: "count_words tool", + inputs: { text: "hello traced world" }, + outputs: { text: "hello traced world", word_count: 3 }, + }; + manifest.spans[0].inputs = requestBody; + manifest.spans[0].outputs = responseBody; + const toolSpan = manifest.spans.find((span) => span.spanType === "TOOL"); + if (!toolSpan) throw new Error("deployed fixture requires a TOOL span"); + toolSpan.name = expectedTool.name; + toolSpan.inputs = expectedTool.inputs; + toolSpan.outputs = expectedTool.outputs; const rows = manifest.spans.map((span, index) => ({ traceId: manifest.traceId, spanId: span.spanId, @@ -1193,6 +1245,9 @@ function deployedProofFixture(): DeployedTraceProof { return { appName, configuredExperimentId, + requestBody, + responseBody, + expectedTool, returnedTraceId, binding, experiment: { @@ -1210,6 +1265,7 @@ function deployedProofFixture(): DeployedTraceProof { spans: rows.map((row) => ({ span_id: row.spanId, parent_span_id: row.parentSpanId, + name: row.name, attributes: row.attributes, status: { code: "OK" }, latency_ms: 1, @@ -1262,6 +1318,31 @@ function assertAppIdentity( } } +function canonicalTraceValue(value: unknown): unknown { + if (Array.isArray(value)) return value.map(canonicalTraceValue); + if (typeof value === "object" && value !== null) { + return Object.fromEntries( + Object.entries(value) + .sort(([left], [right]) => left.localeCompare(right)) + .map(([key, nested]) => [key, canonicalTraceValue(nested)]), + ); + } + return value; +} + +function assertExactTraceValue( + label: string, + actual: unknown, + expected: unknown, +): void { + if ( + JSON.stringify(canonicalTraceValue(actual)) !== + JSON.stringify(canonicalTraceValue(expected)) + ) { + throw new Error(`${label} does not match the deployed turn`); + } +} + function normalizeUcRows( template: string, traceId: string, @@ -1380,9 +1461,23 @@ function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { } const mlflowSpans = proof.traceRecord.data?.spans ?? []; - const mlflowSpanIds = new Set( - mlflowSpans.map((span) => String(span.span_id ?? span.spanId)), + const mlflowSpanIdentities = mlflowSpans.map((span) => + String(span.span_id ?? span.spanId), ); + const ucSpanIdentities = proof.rows.map((row) => String(row.spanId)); + const mlflowSpanIds = new Set(mlflowSpanIdentities); + const ucSpanIds = new Set(ucSpanIdentities); + // The shared contract does not permit an unpaired provider-created span, + // so there is intentionally no identity exception here. + if ( + mlflowSpanIds.size !== mlflowSpanIdentities.length || + ucSpanIds.size !== ucSpanIdentities.length || + mlflowSpanIds.size !== ucSpanIds.size || + [...mlflowSpanIds].some((spanId) => !ucSpanIds.has(spanId)) || + [...ucSpanIds].some((spanId) => !mlflowSpanIds.has(spanId)) + ) { + throw new Error("MLflow and UC span identity sets do not match exactly"); + } const semanticRows = proof.rows.filter( (row) => row.attributes["mlflow.spanType"] !== undefined, ); @@ -1413,6 +1508,51 @@ function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { const manifest = normalizeUcRows(proof.appName, otelTraceId, proof.rows); assertContract(manifest); + const roots = manifest.spans.filter( + (span) => span.parentSpanId === null && span.spanType === "AGENT", + ); + if (roots.length !== 1) { + throw new Error("deployed trace does not have one exact AGENT root"); + } + const root = roots[0]; + assertExactTraceValue("AGENT request input", root.inputs, proof.requestBody); + if (proof.responseBody === undefined || proof.responseBody === null) { + throw new Error("deployed response output is missing"); + } + const responseTraceId = objectValue(proof.responseBody).trace_id; + if (responseTraceId !== proof.returnedTraceId) { + throw new Error("deployed response trace ID does not match returned trace"); + } + assertExactTraceValue( + "AGENT response output", + root.outputs, + proof.responseBody, + ); + if ( + !manifest.spans.some( + (span) => span.spanType === "CHAT_MODEL" || span.spanType === "LLM", + ) + ) { + throw new Error("deployed trace is missing an LLM span"); + } + const toolSpan = manifest.spans.find( + (span) => span.spanType === "TOOL" && span.name === proof.expectedTool.name, + ); + if (!toolSpan) { + throw new Error( + `deployed trace is missing TOOL ${proof.expectedTool.name}`, + ); + } + assertExactTraceValue( + "TOOL inputs", + toolSpan.inputs, + proof.expectedTool.inputs, + ); + assertExactTraceValue( + "TOOL outputs", + toolSpan.outputs, + proof.expectedTool.outputs, + ); return manifest; } @@ -1562,6 +1702,114 @@ test.each([ }, ); +test.each([ + [ + "missing response output", + (proof: DeployedTraceProof) => { + proof.responseBody = undefined; + }, + ], + [ + "wrong response output", + (proof: DeployedTraceProof) => { + proof.responseBody = { + object: "response", + status: "completed", + trace_id: proof.returnedTraceId, + output: [{ content: [{ type: "output_text", text: "wrong" }] }], + }; + }, + ], +] as const)( + "rejects %s that is not bound to the AGENT root", + (_name, mutate) => { + const proof = deployedProofFixture(); + mutate(proof); + expect(() => validateDeployedProof(proof)).toThrow(/response|output/i); + }, +); + +test.each([ + [ + "absent TOOL span", + (proof: DeployedTraceProof) => { + const toolRow = proof.rows.find( + (row) => row.attributes["mlflow.spanType"] === "TOOL", + ); + proof.rows = proof.rows.filter((row) => row !== toolRow); + if (toolRow && proof.traceRecord.data?.spans) { + proof.traceRecord.data.spans = proof.traceRecord.data.spans.filter( + (span) => + String(span.span_id ?? span.spanId) !== String(toolRow.spanId), + ); + } + }, + ], + [ + "wrong TOOL span", + (proof: DeployedTraceProof) => { + const toolRow = proof.rows.find( + (row) => row.attributes["mlflow.spanType"] === "TOOL", + ); + if (toolRow) toolRow.name = "different tool"; + }, + ], + [ + "wrong TOOL inputs", + (proof: DeployedTraceProof) => { + const toolRow = proof.rows.find( + (row) => row.attributes["mlflow.spanType"] === "TOOL", + ); + if (toolRow) toolRow.attributes["mlflow.spanInputs"] = { text: "wrong" }; + }, + ], + [ + "wrong TOOL outputs", + (proof: DeployedTraceProof) => { + const toolRow = proof.rows.find( + (row) => row.attributes["mlflow.spanType"] === "TOOL", + ); + if (toolRow) + toolRow.attributes["mlflow.spanOutputs"] = { word_count: 99 }; + }, + ], +] as const)( + "rejects %s for the deterministic deployed turn", + (_name, mutate) => { + const proof = deployedProofFixture(); + mutate(proof); + expect(() => validateDeployedProof(proof)).toThrow(/TOOL/i); + }, +); + +test("rejects an extra MLflow span missing from UC", () => { + const proof = deployedProofFixture(); + proof.traceRecord.data?.spans?.push({ + span_id: "mlflow-only", + parent_span_id: null, + name: "provider bookkeeping", + attributes: {}, + }); + + expect(() => validateDeployedProof(proof)).toThrow(/span identity/i); +}); + +test("rejects an extra UC span missing from MLflow", () => { + const proof = deployedProofFixture(); + proof.rows.push({ + traceId: "0123456789abcdef0123456789abcdef", + spanId: "uc-only", + parentSpanId: null, + name: "provider bookkeeping", + attributes: {}, + statusCode: "OK", + startTimeUnixNano: "1000000", + endTimeUnixNano: "2000000", + }); + + expect(() => validateDeployedProof(proof)).toThrow(/span identity/i); +}); + test("polls asynchronous Statement Execution to success", async () => { const getStatement = vi.fn(() => ({ statement_id: "statement-1", @@ -1645,6 +1893,14 @@ test.skipIf(missingDeployed.length > 0)( encoding: "utf8", }), ).access_token; + const requestBody = { + input: "Count the words in hello traced world. Use count_words.", + }; + const expectedTool = { + name: "count_words tool", + inputs: { text: "hello traced world" }, + outputs: { text: "hello traced world", word_count: 3 }, + }; const response = await fetch( process.env.APPKIT_TRACE_CONFORMANCE_URL ?? "", { @@ -1654,14 +1910,11 @@ test.skipIf(missingDeployed.length > 0)( "Content-Type": "application/json", "X-MLflow-Return-Trace-Id": "true", }, - body: JSON.stringify({ - input: [ - { role: "user", content: "What time is it? Use the clock tool." }, - ], - }), + body: JSON.stringify(requestBody), }, ); expect(response.ok).toBe(true); + const responseBody = await response.json(); const traceId = response.headers.get("x-mlflow-trace-id"); expect(traceId).toBeTruthy(); const otelTraceId = otelTraceIdFromReturnedTrace( @@ -1768,6 +2021,9 @@ test.skipIf(missingDeployed.length > 0)( validateDeployedProof({ appName, configuredExperimentId, + requestBody, + responseBody, + expectedTool, returnedTraceId: traceId ?? "", binding, experiment, From 414927704894044d3ec4fda8d3d70be82348a59a Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 01:53:55 -0700 Subject: [PATCH 24/31] fix: harden trace conformance lifecycle Signed-off-by: Adam Gurary --- packages/appkit/src/cache/index.ts | 40 ++ packages/appkit/src/core/lifecycle-manager.ts | 5 +- .../appkit-lifecycle.integration.test.ts | 95 +++++ .../src/core/tests/lifecycle-manager.test.ts | 6 + .../trace-conformance.integration.test.ts | 359 +++++++++++++++++- 5 files changed, 491 insertions(+), 14 deletions(-) create mode 100644 packages/appkit/src/core/tests/appkit-lifecycle.integration.test.ts diff --git a/packages/appkit/src/cache/index.ts b/packages/appkit/src/cache/index.ts index e2b79d04c..f99a3855c 100644 --- a/packages/appkit/src/cache/index.ts +++ b/packages/appkit/src/cache/index.ts @@ -56,6 +56,7 @@ export class CacheManager { private readonly name: string = "cache-manager"; private static instance: CacheManager | null = null; private static initPromise: Promise | null = null; + private static shutdownPromise: Promise | null = null; private storage: CacheStorage; private config: CacheConfig; @@ -119,6 +120,9 @@ export class CacheManager { static async getInstance( userConfig?: Partial, ): Promise { + if (CacheManager.shutdownPromise) { + await CacheManager.shutdownPromise; + } if (CacheManager.instance) { return CacheManager.instance; } @@ -135,6 +139,42 @@ export class CacheManager { return CacheManager.initPromise; } + /** + * Close and retire the process-owned cache singleton. + * + * Concurrent shutdown callers share one close. A later `getInstance()` + * waits for that close and creates fresh storage, so an ended persistent + * pool can never remain reachable through the singleton. + */ + static async shutdown(): Promise { + if (CacheManager.shutdownPromise) return CacheManager.shutdownPromise; + + const shutdown = async () => { + const pendingInitialization = CacheManager.initPromise; + let instance = CacheManager.instance; + if (!instance && pendingInitialization) { + try { + instance = await pendingInitialization; + } catch { + // Failed initialization owns no resource that can be closed. + } + } + try { + await instance?.close(); + } finally { + if (CacheManager.instance === instance) CacheManager.instance = null; + if (CacheManager.initPromise === pendingInitialization) { + CacheManager.initPromise = null; + } + } + }; + + CacheManager.shutdownPromise = shutdown().finally(() => { + CacheManager.shutdownPromise = null; + }); + return CacheManager.shutdownPromise; + } + /** * Create a new cache manager instance * diff --git a/packages/appkit/src/core/lifecycle-manager.ts b/packages/appkit/src/core/lifecycle-manager.ts index da71edddb..1cb99e0d9 100644 --- a/packages/appkit/src/core/lifecycle-manager.ts +++ b/packages/appkit/src/core/lifecycle-manager.ts @@ -212,16 +212,15 @@ export class LifecycleManager { /** Close the cache storage, bounded and error-isolated. */ private async closeCacheStorage(): Promise { - let cache: CacheManager; try { - cache = CacheManager.getInstanceSync(); + CacheManager.getInstanceSync(); } catch { // Cache was never initialized ā€” nothing to close. return; } try { await this.raceWithTimeout( - cache.close(), + CacheManager.shutdown(), LifecycleManager.PHASE_SHUTDOWN_TIMEOUT_MS, "cache storage close", ); diff --git a/packages/appkit/src/core/tests/appkit-lifecycle.integration.test.ts b/packages/appkit/src/core/tests/appkit-lifecycle.integration.test.ts new file mode 100644 index 000000000..6a4cfb205 --- /dev/null +++ b/packages/appkit/src/core/tests/appkit-lifecycle.integration.test.ts @@ -0,0 +1,95 @@ +import { mockServiceContext, setupDatabricksEnv } from "@tools/test-helpers"; +import type { CacheEntry, CacheStorage } from "shared"; +import { afterEach, beforeEach, expect, test } from "vitest"; +import { CacheManager } from "../../cache"; +import { createApp } from "../appkit"; + +class LifecyclePersistentStorage implements CacheStorage { + private readonly entries = new Map>(); + private ended = false; + + private assertOpen(): void { + if (this.ended) throw new Error("persistent pool has ended"); + } + + async get(key: string): Promise | null> { + this.assertOpen(); + return (this.entries.get(key) as CacheEntry | undefined) ?? null; + } + + async set(key: string, entry: CacheEntry): Promise { + this.assertOpen(); + this.entries.set(key, entry as CacheEntry); + } + + async delete(key: string): Promise { + this.assertOpen(); + this.entries.delete(key); + } + + async clear(): Promise { + this.assertOpen(); + this.entries.clear(); + } + + async has(key: string): Promise { + this.assertOpen(); + return this.entries.has(key); + } + + async size(): Promise { + this.assertOpen(); + return this.entries.size; + } + + isPersistent(): boolean { + return true; + } + + async healthCheck(): Promise { + this.assertOpen(); + return true; + } + + async close(): Promise { + if (this.ended) throw new Error("persistent pool ended twice"); + this.ended = true; + } +} + +let serviceContext: Awaited>; + +beforeEach(async () => { + setupDatabricksEnv(); + serviceContext = await mockServiceContext(); +}); + +afterEach(async () => { + serviceContext.restore(); +}); + +test("two sequential app lifecycles recreate persistent cache storage", async () => { + const baselineSigterm = process.listenerCount("SIGTERM"); + const baselineSigint = process.listenerCount("SIGINT"); + const firstStorage = new LifecyclePersistentStorage(); + const first = await createApp({ + plugins: [], + cache: { storage: firstStorage }, + disableInternalTelemetry: true, + }); + await CacheManager.getInstanceSync().set("first", "value"); + await first.shutdown(); + + const secondStorage = new LifecyclePersistentStorage(); + const second = await createApp({ + plugins: [], + cache: { storage: secondStorage }, + disableInternalTelemetry: true, + }); + await CacheManager.getInstanceSync().set("second", "value"); + expect(await CacheManager.getInstanceSync().get("second")).toBe("value"); + await second.shutdown(); + + expect(process.listenerCount("SIGTERM")).toBe(baselineSigterm); + expect(process.listenerCount("SIGINT")).toBe(baselineSigint); +}); diff --git a/packages/appkit/src/core/tests/lifecycle-manager.test.ts b/packages/appkit/src/core/tests/lifecycle-manager.test.ts index ed9d92f6e..735d93b86 100644 --- a/packages/appkit/src/core/tests/lifecycle-manager.test.ts +++ b/packages/appkit/src/core/tests/lifecycle-manager.test.ts @@ -15,6 +15,7 @@ vi.mock("../../cache", () => ({ getInstanceSync: vi.fn().mockReturnValue({ close: vi.fn().mockResolvedValue(undefined), }), + shutdown: vi.fn().mockResolvedValue(undefined), }, })); @@ -227,6 +228,7 @@ describe("LifecycleManager", () => { vi.mocked(CacheManager.getInstanceSync).mockReturnValueOnce({ close, } as any); + vi.mocked(CacheManager.shutdown).mockImplementationOnce(close); vi.mocked(TelemetryManager.getInstance).mockReturnValueOnce({ shutdown: flush, } as any); @@ -266,6 +268,7 @@ describe("LifecycleManager", () => { vi.mocked(CacheManager.getInstanceSync).mockReturnValueOnce({ close: hangingClose, } as any); + vi.mocked(CacheManager.shutdown).mockImplementationOnce(hangingClose); const done = new LifecycleManager(contextWithPlugins({})).shutdown(); await vi.advanceTimersByTimeAsync(2_000); @@ -318,6 +321,9 @@ describe("LifecycleManager", () => { order.push("cache-close"); }), } as any); + vi.mocked(CacheManager.shutdown).mockImplementationOnce(async () => { + order.push("cache-close"); + }); vi.mocked(TelemetryManager.getInstance).mockReturnValueOnce({ shutdown: vi.fn(async () => { order.push("flush"); diff --git a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts index 6f10b2206..e2a3813b4 100644 --- a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -5,6 +5,7 @@ import { readdirSync, readFileSync, rmSync, + writeFileSync, } from "node:fs"; import type { Server } from "node:http"; import { createRequire } from "node:module"; @@ -744,6 +745,25 @@ test("every behavior-discovered generated surface executes its HTTP trace proof" const failures: string[] = []; for (const candidate of candidates) { try { + const requestedCandidate = process.env.APPKIT_TRACE_CONFORMANCE_CANDIDATE; + if (requestedCandidate === candidate.name) { + const sourceDirectory = + process.env.APPKIT_TRACE_CONFORMANCE_SOURCE_DIRECTORY; + if (!sourceDirectory) { + throw new Error( + `${candidate.name} owner proof is missing its generated source directory`, + ); + } + for (const relative of [ + "server/server.ts", + "server/agents/helper.ts", + ]) { + expect( + readFileSync(join(sourceDirectory, relative), "utf8"), + `${candidate.name} source provenance ${relative}`, + ).toBe(readFileSync(join(candidate.directory, relative), "utf8")); + } + } const baselineSigterm = process.listenerCount("SIGTERM"); const baselineSigint = process.listenerCount("SIGINT"); const localManifest = await captureTurn(candidate.name); @@ -769,6 +789,12 @@ test("every behavior-discovered generated surface executes its HTTP trace proof" ), ).toBe(true); assertContract(manifest); + if ( + process.env.APPKIT_TRACE_CONFORMANCE_CANDIDATE === candidate.name && + process.env.TRACE_CONFORMANCE_MANIFEST + ) { + writePythonManifest(process.env.TRACE_CONFORMANCE_MANIFEST, manifest); + } } catch (error) { failures.push( `${candidate.name}: ${error instanceof Error ? error.message : String(error)}`, @@ -885,6 +911,42 @@ function fixtureManifest(children: SpanManifest[]): TraceManifest { }; } +function writePythonManifest(path: string, manifest: TraceManifest): void { + writeFileSync( + path, + `${JSON.stringify( + { + template: manifest.template, + trace_id: manifest.traceId, + spans: manifest.spans.map((span) => ({ + name: span.name, + span_type: span.spanType, + span_id: span.spanId, + parent_span_id: span.parentSpanId, + inputs: span.inputs, + outputs: span.outputs, + status: + span.status === 1 + ? "OK" + : span.status === 2 + ? "ERROR" + : span.status, + latency_ms: span.latencyMs, + model: span.model ?? null, + provider: span.provider ?? null, + usage: span.usage, + cost_usd: span.costUsd ?? null, + cost_available: span.costAvailable, + links: span.links, + attributes: span.attributes, + })), + }, + null, + 2, + )}\n`, + ); +} + function validWorkloads(): TraceManifest[] { const tool = fixtureSpan("clock", "TOOL", "tool", "root"); tool.inputs = { zone: "UTC" }; @@ -1266,7 +1328,7 @@ function deployedProofFixture(): DeployedTraceProof { span_id: row.spanId, parent_span_id: row.parentSpanId, name: row.name, - attributes: row.attributes, + attributes: structuredClone(row.attributes), status: { code: "OK" }, latency_ms: 1, })), @@ -1414,6 +1476,147 @@ function normalizeUcRows( }; } +function normalizeMlflowSpans( + template: string, + traceId: string, + spans: Array>, +): TraceManifest { + const semanticSpans = spans.filter( + (span) => objectValue(span.attributes)["mlflow.spanType"] !== undefined, + ); + const semanticSpanIds = new Set( + semanticSpans.map((span) => String(span.span_id ?? span.spanId)), + ); + return { + template, + traceId, + spans: semanticSpans.map((span) => { + const attributes = { ...objectValue(span.attributes) }; + attributes.app_id ??= + attributes["app.id"] ?? attributes["appkit.app.name"]; + attributes.user_id ??= + attributes["user.id"] ?? attributes["mlflow.trace.user"]; + attributes.session_id ??= + attributes["session.id"] ?? attributes["mlflow.trace.session"]; + attributes.ttft_ms ??= + attributes["appkit.ttft_ms"] ?? + attributes["appkit.first_token.duration_ms"]; + attributes.stream_duration_ms ??= + attributes["appkit.stream_duration_ms"] ?? + attributes["appkit.stream.duration_ms"]; + if ( + attributes.ttft_ms !== undefined || + attributes.stream_duration_ms !== undefined + ) { + attributes.streaming = true; + } + const spanType = String(attributes["mlflow.spanType"] ?? ""); + const usage = objectValue( + attributes[ + spanType === "AGENT" + ? "mlflow.trace.tokenUsage" + : "mlflow.chat.tokenUsage" + ], + ) as Record; + const rawParentSpanId = span.parent_span_id ?? span.parentSpanId; + const parentSpanId = + rawParentSpanId === undefined || rawParentSpanId === null + ? null + : String(rawParentSpanId); + const statusRecord = objectValue(span.status); + return { + name: String(span.name), + spanType, + spanId: String(span.span_id ?? span.spanId), + parentSpanId: + spanType === "AGENT" && + parentSpanId !== null && + !semanticSpanIds.has(parentSpanId) + ? null + : parentSpanId, + inputs: decoded(attributes["mlflow.spanInputs"]), + outputs: decoded(attributes["mlflow.spanOutputs"]), + status: normalizeUcStatus( + statusRecord.code ?? statusRecord.status_code ?? span.status, + ), + latencyMs: Number(span.latency_ms ?? span.latencyMs ?? -1), + model: attributes["mlflow.chat.model"] as string | undefined, + provider: attributes["mlflow.chat.provider"] as string | undefined, + usage, + costUsd: attributes["mlflow.llm.cost"] as number | undefined, + costAvailable: attributes["appkit.cost.available"] === true, + links: Array.isArray(span.links) ? span.links : [], + attributes, + }; + }), + }; +} + +function contractSemantics(span: SpanManifest): Record { + return { + name: span.name, + spanType: span.spanType, + spanId: span.spanId, + parentSpanId: span.parentSpanId, + inputs: span.inputs, + outputs: span.outputs, + status: span.status, + latencyMs: span.latencyMs, + model: span.model, + provider: span.provider, + usage: span.usage, + costUsd: span.costUsd, + costAvailable: span.costAvailable, + links: span.links, + identity: { + app_id: span.attributes.app_id, + user_id: span.attributes.user_id, + session_id: span.attributes.session_id, + }, + streaming: { + enabled: span.attributes.streaming, + ttft_ms: span.attributes.ttft_ms, + stream_duration_ms: span.attributes.stream_duration_ms, + }, + }; +} + +function assertCrossSourceParity( + mlflow: TraceManifest, + uc: TraceManifest, +): void { + if (mlflow.traceId !== uc.traceId) { + throw new Error("MLflow and UC trace identities do not match exactly"); + } + const mlflowById = new Map( + mlflow.spans.map((span) => [span.spanId, span] as const), + ); + const ucById = new Map(uc.spans.map((span) => [span.spanId, span] as const)); + if ( + mlflowById.size !== mlflow.spans.length || + ucById.size !== uc.spans.length || + mlflowById.size !== ucById.size || + [...mlflowById.keys()].some((spanId) => !ucById.has(spanId)) || + [...ucById.keys()].some((spanId) => !mlflowById.has(spanId)) + ) { + throw new Error("MLflow and UC span identity sets do not match exactly"); + } + for (const [spanId, mlflowSpan] of mlflowById) { + const ucSpan = ucById.get(spanId); + if (!ucSpan || mlflowSpan.parentSpanId !== ucSpan.parentSpanId) { + throw new Error( + `MLflow and UC parent identity differs for span ${spanId}`, + ); + } + if ( + JSON.stringify(canonicalTraceValue(contractSemantics(mlflowSpan))) !== + JSON.stringify(canonicalTraceValue(contractSemantics(ucSpan))) + ) { + throw new Error(`MLflow and UC semantics differ for span ${spanId}`); + } + } +} + function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { const experiment = proof.experiment.experiment; if (experiment?.experiment_id !== proof.configuredExperimentId) { @@ -1489,15 +1692,15 @@ function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { } } - const mlflowRoot = mlflowSpans.find((span) => { - const attributes = objectValue(span.attributes); - return attributes["mlflow.spanType"] === "AGENT"; - }); + const mlflowRoot = mlflowSpans.find( + (span) => objectValue(span.attributes)["mlflow.spanType"] === "AGENT", + ); assertAppIdentity( objectValue(mlflowRoot?.attributes), proof.appName, "MLflow trace", ); + const ucRoot = semanticRows.find( (row) => row.attributes["mlflow.spanType"] === "AGENT", ); @@ -1506,9 +1709,16 @@ function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { } assertAppIdentity(ucRoot.attributes, proof.appName, "UC trace"); - const manifest = normalizeUcRows(proof.appName, otelTraceId, proof.rows); - assertContract(manifest); - const roots = manifest.spans.filter( + const mlflowManifest = normalizeMlflowSpans( + proof.appName, + otelTraceId, + mlflowSpans, + ); + const ucManifest = normalizeUcRows(proof.appName, otelTraceId, proof.rows); + assertContract(mlflowManifest); + assertContract(ucManifest); + assertCrossSourceParity(mlflowManifest, ucManifest); + const roots = ucManifest.spans.filter( (span) => span.parentSpanId === null && span.spanType === "AGENT", ); if (roots.length !== 1) { @@ -1529,13 +1739,13 @@ function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { proof.responseBody, ); if ( - !manifest.spans.some( + !ucManifest.spans.some( (span) => span.spanType === "CHAT_MODEL" || span.spanType === "LLM", ) ) { throw new Error("deployed trace is missing an LLM span"); } - const toolSpan = manifest.spans.find( + const toolSpan = ucManifest.spans.find( (span) => span.spanType === "TOOL" && span.name === proof.expectedTool.name, ); if (!toolSpan) { @@ -1553,7 +1763,7 @@ function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { toolSpan.outputs, proof.expectedTool.outputs, ); - return manifest; + return ucManifest; } async function pollStatement( @@ -1810,6 +2020,133 @@ test("rejects an extra UC span missing from MLflow", () => { expect(() => validateDeployedProof(proof)).toThrow(/span identity/i); }); +test.each([ + [ + "missing MLflow AGENT request", + (proof: DeployedTraceProof) => { + const root = proof.traceRecord.data?.spans?.find( + (span) => objectValue(span.attributes)["mlflow.spanType"] === "AGENT", + ); + delete objectValue(root?.attributes)["mlflow.spanInputs"]; + }, + ], + [ + "incorrect MLflow AGENT response", + (proof: DeployedTraceProof) => { + const root = proof.traceRecord.data?.spans?.find( + (span) => objectValue(span.attributes)["mlflow.spanType"] === "AGENT", + ); + objectValue(root?.attributes)["mlflow.spanOutputs"] = { output: "wrong" }; + }, + ], + [ + "incorrect MLflow TOOL name", + (proof: DeployedTraceProof) => { + const tool = proof.traceRecord.data?.spans?.find( + (span) => objectValue(span.attributes)["mlflow.spanType"] === "TOOL", + ); + if (tool) tool.name = "wrong tool"; + }, + ], + [ + "missing MLflow TOOL input", + (proof: DeployedTraceProof) => { + const tool = proof.traceRecord.data?.spans?.find( + (span) => objectValue(span.attributes)["mlflow.spanType"] === "TOOL", + ); + delete objectValue(tool?.attributes)["mlflow.spanInputs"]; + }, + ], + [ + "incorrect MLflow TOOL output", + (proof: DeployedTraceProof) => { + const tool = proof.traceRecord.data?.spans?.find( + (span) => objectValue(span.attributes)["mlflow.spanType"] === "TOOL", + ); + objectValue(tool?.attributes)["mlflow.spanOutputs"] = { word_count: 99 }; + }, + ], + [ + "missing MLflow model input", + (proof: DeployedTraceProof) => { + const model = proof.traceRecord.data?.spans?.find((span) => + ["CHAT_MODEL", "LLM"].includes( + String(objectValue(span.attributes)["mlflow.spanType"]), + ), + ); + delete objectValue(model?.attributes)["mlflow.spanInputs"]; + }, + ], + [ + "incorrect MLflow model output", + (proof: DeployedTraceProof) => { + const model = proof.traceRecord.data?.spans?.find((span) => + ["CHAT_MODEL", "LLM"].includes( + String(objectValue(span.attributes)["mlflow.spanType"]), + ), + ); + objectValue(model?.attributes)["mlflow.spanOutputs"] = { text: "wrong" }; + }, + ], + [ + "missing MLflow model usage", + (proof: DeployedTraceProof) => { + const model = proof.traceRecord.data?.spans?.find((span) => + ["CHAT_MODEL", "LLM"].includes( + String(objectValue(span.attributes)["mlflow.spanType"]), + ), + ); + delete objectValue(model?.attributes)["mlflow.chat.tokenUsage"]; + }, + ], + [ + "incorrect MLflow model cost", + (proof: DeployedTraceProof) => { + const model = proof.traceRecord.data?.spans?.find((span) => + ["CHAT_MODEL", "LLM"].includes( + String(objectValue(span.attributes)["mlflow.spanType"]), + ), + ); + objectValue(model?.attributes)["mlflow.llm.cost"] = 99; + }, + ], + [ + "missing MLflow model timing", + (proof: DeployedTraceProof) => { + const model = proof.traceRecord.data?.spans?.find((span) => + ["CHAT_MODEL", "LLM"].includes( + String(objectValue(span.attributes)["mlflow.spanType"]), + ), + ); + delete objectValue(model?.attributes).ttft_ms; + }, + ], + [ + "incorrect MLflow model status", + (proof: DeployedTraceProof) => { + const model = proof.traceRecord.data?.spans?.find((span) => + ["CHAT_MODEL", "LLM"].includes( + String(objectValue(span.attributes)["mlflow.spanType"]), + ), + ); + if (model) model.status = { code: "UNSET" }; + }, + ], + [ + "incorrect MLflow topology", + (proof: DeployedTraceProof) => { + const tool = proof.traceRecord.data?.spans?.find( + (span) => objectValue(span.attributes)["mlflow.spanType"] === "TOOL", + ); + if (tool) tool.parent_span_id = "model"; + }, + ], +] as const)("rejects %s while UC remains valid", (_name, mutate) => { + const proof = deployedProofFixture(); + mutate(proof); + expect(() => validateDeployedProof(proof)).toThrow(); +}); + test("polls asynchronous Statement Execution to success", async () => { const getStatement = vi.fn(() => ({ statement_id: "statement-1", From c65d4a6b628b5de5dadf169da7559d20e470c713 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 02:16:32 -0700 Subject: [PATCH 25/31] fix: close trace parity and shutdown races Signed-off-by: Adam Gurary --- packages/appkit/src/cache/index.ts | 54 +++--- .../src/cache/tests/cache-manager.test.ts | 55 ++++++ packages/appkit/src/core/lifecycle-manager.ts | 64 +++---- .../src/core/tests/lifecycle-manager.test.ts | 33 ++++ .../trace-conformance.integration.test.ts | 180 ++++++++++++++---- 5 files changed, 292 insertions(+), 94 deletions(-) diff --git a/packages/appkit/src/cache/index.ts b/packages/appkit/src/cache/index.ts index f99a3855c..1aec4ddca 100644 --- a/packages/appkit/src/cache/index.ts +++ b/packages/appkit/src/cache/index.ts @@ -128,12 +128,14 @@ export class CacheManager { } if (!CacheManager.initPromise) { - CacheManager.initPromise = CacheManager.create(userConfig).then( - (instance) => { + let ownedInitialization: Promise; + ownedInitialization = CacheManager.create(userConfig).then((instance) => { + if (CacheManager.initPromise === ownedInitialization) { CacheManager.instance = instance; - return instance; - }, - ); + } + return instance; + }); + CacheManager.initPromise = ownedInitialization; } return CacheManager.initPromise; @@ -146,33 +148,29 @@ export class CacheManager { * waits for that close and creates fresh storage, so an ended persistent * pool can never remain reachable through the singleton. */ - static async shutdown(): Promise { + static shutdown(): Promise { if (CacheManager.shutdownPromise) return CacheManager.shutdownPromise; - const shutdown = async () => { - const pendingInitialization = CacheManager.initPromise; - let instance = CacheManager.instance; - if (!instance && pendingInitialization) { - try { - instance = await pendingInitialization; - } catch { - // Failed initialization owns no resource that can be closed. - } - } - try { - await instance?.close(); - } finally { - if (CacheManager.instance === instance) CacheManager.instance = null; - if (CacheManager.initPromise === pendingInitialization) { - CacheManager.initPromise = null; - } + const detachedInstance = CacheManager.instance; + const detachedInitialization = CacheManager.initPromise; + CacheManager.instance = null; + CacheManager.initPromise = null; + + const detachedOwner = detachedInstance + ? Promise.resolve(detachedInstance) + : detachedInitialization; + let ownedShutdown: Promise; + ownedShutdown = ( + detachedOwner + ? detachedOwner.then((instance) => instance.close()) + : Promise.resolve() + ).finally(() => { + if (CacheManager.shutdownPromise === ownedShutdown) { + CacheManager.shutdownPromise = null; } - }; - - CacheManager.shutdownPromise = shutdown().finally(() => { - CacheManager.shutdownPromise = null; }); - return CacheManager.shutdownPromise; + CacheManager.shutdownPromise = ownedShutdown; + return ownedShutdown; } /** diff --git a/packages/appkit/src/cache/tests/cache-manager.test.ts b/packages/appkit/src/cache/tests/cache-manager.test.ts index 9f6689883..e88be5712 100644 --- a/packages/appkit/src/cache/tests/cache-manager.test.ts +++ b/packages/appkit/src/cache/tests/cache-manager.test.ts @@ -102,6 +102,7 @@ describe("CacheManager", () => { // Access private static fields to reset singleton (CacheManager as any).instance = null; (CacheManager as any).initPromise = null; + (CacheManager as any).shutdownPromise = null; // Default: Lakebase unavailable (most tests pass explicit storage) mockPoolQuery.mockRejectedValue(new Error("Connection failed")); }); @@ -839,6 +840,60 @@ describe("CacheManager", () => { await expect(cache.close()).resolves.not.toThrow(); }); + + test("shutdown detaches closing storage before awaiting close and cannot close its replacement", async () => { + let releaseClose: (() => void) | undefined; + const closingStorage = createMockStorage(true); + closingStorage.close = vi.fn( + () => + new Promise((resolve) => { + releaseClose = resolve; + }), + ); + const closingInstance = await CacheManager.getInstance({ + storage: closingStorage, + }); + + const shutdown = CacheManager.shutdown(); + + await Promise.resolve(); + expect(closingStorage.close).toHaveBeenCalledTimes(1); + expect(() => CacheManager.getInstanceSync()).toThrow( + "CacheManager not initialized", + ); + + const replacementStorage = createMockStorage(true); + let replacementResolved = false; + const replacementPromise = CacheManager.getInstance({ + storage: replacementStorage, + }).then((instance) => { + replacementResolved = true; + return instance; + }); + await Promise.resolve(); + expect(replacementResolved).toBe(false); + expect(() => CacheManager.getInstanceSync()).toThrow( + "CacheManager not initialized", + ); + + releaseClose?.(); + await shutdown; + const replacement = await replacementPromise; + + expect(replacement).not.toBe(closingInstance); + expect(CacheManager.getInstanceSync()).toBe(replacement); + expect(closingStorage.close).toHaveBeenCalledTimes(1); + expect(replacementStorage.close).not.toHaveBeenCalled(); + }); + + test("concurrent shutdown callers close a detached instance only once", async () => { + const storage = createMockStorage(true); + await CacheManager.getInstance({ storage }); + + await Promise.all([CacheManager.shutdown(), CacheManager.shutdown()]); + + expect(storage.close).toHaveBeenCalledTimes(1); + }); }); describe("maybeCleanup", () => { diff --git a/packages/appkit/src/core/lifecycle-manager.ts b/packages/appkit/src/core/lifecycle-manager.ts index 1cb99e0d9..100fb7186 100644 --- a/packages/appkit/src/core/lifecycle-manager.ts +++ b/packages/appkit/src/core/lifecycle-manager.ts @@ -44,13 +44,9 @@ export class LifecycleManager { */ private static readonly PHASE_SHUTDOWN_TIMEOUT_MS = 2_000; - /** - * Guards against re-entrant shutdown (e.g. SIGTERM followed by SIGINT). - * The flag set in `shutdown` must remain synchronous and first ā€” any - * `await` before it would open a window for a second signal to re-enter - * the sequence. - */ - private isShuttingDown = false; + private shutdownPromise?: Promise; + private exitRequested = false; + private forceExitTimer?: ReturnType; /** * Name of the shutdown phase currently in flight, so the force-exit log * can say where shutdown got stuck without extra bookkeeping. @@ -73,8 +69,8 @@ export class LifecycleManager { installSignalHandlers(): void { if (this.signalHandlers) return; this.signalHandlers = { - SIGTERM: () => void this.shutdown(), - SIGINT: () => void this.shutdown(), + SIGTERM: () => void this.shutdown({ exitProcess: true }), + SIGINT: () => void this.shutdown({ exitProcess: true }), }; process.once("SIGTERM", this.signalHandlers.SIGTERM); process.once("SIGINT", this.signalHandlers.SIGINT); @@ -102,16 +98,32 @@ export class LifecycleManager { * shutdown is not a crash. Exit 1 is reserved for an unexpected error * thrown by the sequence itself. */ - async shutdown({ + shutdown({ exitProcess = true, }: { exitProcess?: boolean; } = {}): Promise { - // Must stay synchronous and first: any await before the flag is set - // would let a second signal re-enter the shutdown sequence. - if (this.isShuttingDown) return; - this.isShuttingDown = true; + if (exitProcess) this.requestProcessExit(); + if (this.shutdownPromise) return this.shutdownPromise; + this.shutdownPromise = Promise.resolve().then(() => this.runShutdown()); + return this.shutdownPromise; + } + private requestProcessExit(): void { + this.exitRequested = true; + if (this.forceExitTimer) return; + this.forceExitTimer = setTimeout(() => { + logger.error( + "Graceful shutdown did NOT complete within the %dms budget (phase in flight: %s); force-exiting with code 0.", + LifecycleManager.SHUTDOWN_TIMEOUT_MS, + this.shutdownPhase, + ); + process.exit(0); + }, LifecycleManager.SHUTDOWN_TIMEOUT_MS); + this.forceExitTimer.unref(); + } + + private async runShutdown(): Promise { logger.info("Starting graceful shutdown..."); let exitCode = 0; @@ -121,23 +133,6 @@ export class LifecycleManager { // shutdown, not a crash), and orchestrators record nonzero exits on // deploys as crashes. The error log below is the stuck-shutdown // signal instead of the exit code. - const forceExitTimer = exitProcess - ? setTimeout(() => { - logger.error( - "Graceful shutdown did NOT complete within the %dms budget (phase in flight: %s); force-exiting with code 0.", - LifecycleManager.SHUTDOWN_TIMEOUT_MS, - this.shutdownPhase, - ); - process.exit(0); - }, LifecycleManager.SHUTDOWN_TIMEOUT_MS) - : undefined; - // unref so this backstop timer never by itself keeps the process alive. - // Any real pending teardown (OTEL export timer, DB pool sockets, the - // still-open HTTP listener) is a ref'd handle that holds the loop open - // until this fires; if nothing is ref'd, there is nothing left to tear - // down and exiting early is correct. - forceExitTimer?.unref(); - try { const plugins = Array.from(this.context.getPlugins().values()); @@ -205,9 +200,12 @@ export class LifecycleManager { exitCode = 1; } - if (forceExitTimer) clearTimeout(forceExitTimer); + if (this.forceExitTimer) { + clearTimeout(this.forceExitTimer); + this.forceExitTimer = undefined; + } this.removeSignalHandlers(); - if (exitProcess) process.exit(exitCode); + if (this.exitRequested) process.exit(exitCode); } /** Close the cache storage, bounded and error-isolated. */ diff --git a/packages/appkit/src/core/tests/lifecycle-manager.test.ts b/packages/appkit/src/core/tests/lifecycle-manager.test.ts index 735d93b86..9b6aad8ab 100644 --- a/packages/appkit/src/core/tests/lifecycle-manager.test.ts +++ b/packages/appkit/src/core/tests/lifecycle-manager.test.ts @@ -401,5 +401,38 @@ describe("LifecycleManager", () => { expect(process.listenerCount("SIGINT")).toBe(baselineSigint); expect(exitSpy).not.toHaveBeenCalled(); }); + + test("a signal during programmatic shutdown exits only after the shared teardown completes", async () => { + let releaseShutdown: (() => void) | undefined; + const shutdownHook = vi.fn( + () => + new Promise((resolve) => { + releaseShutdown = resolve; + }), + ); + const manager = new LifecycleManager( + contextWithPlugins({ + agent: { name: "agent", shutdown: shutdownHook }, + }), + ); + manager.installSignalHandlers(); + const sigterm = process.listeners("SIGTERM").at(-1) as () => void; + const sigint = process.listeners("SIGINT").at(-1) as () => void; + + const programmaticShutdown = manager.shutdown({ exitProcess: false }); + await vi.waitFor(() => expect(shutdownHook).toHaveBeenCalledTimes(1)); + + sigterm(); + sigint(); + expect(exitSpy).not.toHaveBeenCalled(); + expect(shutdownHook).toHaveBeenCalledTimes(1); + + releaseShutdown?.(); + await programmaticShutdown; + + expect(exitSpy).toHaveBeenCalledTimes(1); + expect(exitSpy).toHaveBeenCalledWith(0); + expect(shutdownHook).toHaveBeenCalledTimes(1); + }); }); }); diff --git a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts index e2a3813b4..dd1ea67e7 100644 --- a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -1410,16 +1410,10 @@ function normalizeUcRows( traceId: string, rows: UcSpanRow[], ): TraceManifest { - const semanticRows = rows.filter( - (row) => row.attributes["mlflow.spanType"] !== undefined, - ); - const semanticSpanIds = new Set( - semanticRows.map((row) => String(row.spanId)), - ); return { template, traceId, - spans: semanticRows.map((row) => { + spans: rows.map((row) => { const attributes = { ...row.attributes }; attributes.app_id ??= attributes["app.id"] ?? attributes["appkit.app.name"]; @@ -1454,12 +1448,7 @@ function normalizeUcRows( name: String(row.name), spanType, spanId: String(row.spanId), - parentSpanId: - spanType === "AGENT" && - rawParentSpanId !== null && - !semanticSpanIds.has(rawParentSpanId) - ? null - : rawParentSpanId, + parentSpanId: rawParentSpanId, inputs: decoded(attributes["mlflow.spanInputs"]), outputs: decoded(attributes["mlflow.spanOutputs"]), status: normalizeUcStatus(row.statusCode), @@ -1481,16 +1470,10 @@ function normalizeMlflowSpans( traceId: string, spans: Array>, ): TraceManifest { - const semanticSpans = spans.filter( - (span) => objectValue(span.attributes)["mlflow.spanType"] !== undefined, - ); - const semanticSpanIds = new Set( - semanticSpans.map((span) => String(span.span_id ?? span.spanId)), - ); return { template, traceId, - spans: semanticSpans.map((span) => { + spans: spans.map((span) => { const attributes = { ...objectValue(span.attributes) }; attributes.app_id ??= attributes["app.id"] ?? attributes["appkit.app.name"]; @@ -1528,12 +1511,7 @@ function normalizeMlflowSpans( name: String(span.name), spanType, spanId: String(span.span_id ?? span.spanId), - parentSpanId: - spanType === "AGENT" && - parentSpanId !== null && - !semanticSpanIds.has(parentSpanId) - ? null - : parentSpanId, + parentSpanId, inputs: decoded(attributes["mlflow.spanInputs"]), outputs: decoded(attributes["mlflow.spanOutputs"]), status: normalizeUcStatus( @@ -1568,6 +1546,7 @@ function contractSemantics(span: SpanManifest): Record { costUsd: span.costUsd, costAvailable: span.costAvailable, links: span.links, + attributes: span.attributes, identity: { app_id: span.attributes.app_id, user_id: span.attributes.user_id, @@ -1581,6 +1560,23 @@ function contractSemantics(span: SpanManifest): Record { }; } +function semanticProjection(manifest: TraceManifest): TraceManifest { + const spans = manifest.spans.filter((span) => span.spanType !== ""); + const semanticSpanIds = new Set(spans.map((span) => span.spanId)); + return { + ...manifest, + spans: spans.map((span) => ({ + ...span, + parentSpanId: + span.spanType === "AGENT" && + span.parentSpanId !== null && + !semanticSpanIds.has(span.parentSpanId) + ? null + : span.parentSpanId, + })), + }; +} + function assertCrossSourceParity( mlflow: TraceManifest, uc: TraceManifest, @@ -1608,11 +1604,17 @@ function assertCrossSourceParity( `MLflow and UC parent identity differs for span ${spanId}`, ); } - if ( - JSON.stringify(canonicalTraceValue(contractSemantics(mlflowSpan))) !== - JSON.stringify(canonicalTraceValue(contractSemantics(ucSpan))) - ) { - throw new Error(`MLflow and UC semantics differ for span ${spanId}`); + const mlflowSemantics = contractSemantics(mlflowSpan); + const ucSemantics = contractSemantics(ucSpan); + for (const field of Object.keys(mlflowSemantics)) { + if ( + JSON.stringify(canonicalTraceValue(mlflowSemantics[field])) !== + JSON.stringify(canonicalTraceValue(ucSemantics[field])) + ) { + throw new Error( + `MLflow and UC semantics ${field} differs for span ${spanId}`, + ); + } } } } @@ -1709,12 +1711,15 @@ function validateDeployedProof(proof: DeployedTraceProof): TraceManifest { } assertAppIdentity(ucRoot.attributes, proof.appName, "UC trace"); - const mlflowManifest = normalizeMlflowSpans( + const rawMlflowManifest = normalizeMlflowSpans( proof.appName, otelTraceId, mlflowSpans, ); - const ucManifest = normalizeUcRows(proof.appName, otelTraceId, proof.rows); + const rawUcManifest = normalizeUcRows(proof.appName, otelTraceId, proof.rows); + assertCrossSourceParity(rawMlflowManifest, rawUcManifest); + const mlflowManifest = semanticProjection(rawMlflowManifest); + const ucManifest = semanticProjection(rawUcManifest); assertContract(mlflowManifest); assertContract(ucManifest); assertCrossSourceParity(mlflowManifest, ucManifest); @@ -2020,6 +2025,115 @@ test("rejects an extra UC span missing from MLflow", () => { expect(() => validateDeployedProof(proof)).toThrow(/span identity/i); }); +function addPairedNonSemanticSpan( + proof: DeployedTraceProof, + { + spanId = "provider-wrapper", + parentSpanId = null, + }: { spanId?: string; parentSpanId?: string | null } = {}, +): void { + proof.rows.push({ + traceId: "0123456789abcdef0123456789abcdef", + spanId, + parentSpanId, + name: "provider bookkeeping", + attributes: { + "http.request.body": { prompt: "hello" }, + "http.response.body": { answer: "world" }, + "http.route": "/invocations", + }, + statusCode: "OK", + startTimeUnixNano: "11000000", + endTimeUnixNano: "12000000", + }); + proof.traceRecord.data?.spans?.push({ + span_id: spanId, + parent_span_id: parentSpanId, + name: "provider bookkeeping", + attributes: { + "http.request.body": { prompt: "hello" }, + "http.response.body": { answer: "world" }, + "http.route": "/invocations", + }, + status: { code: "OK" }, + latency_ms: 1, + }); +} + +test.each([ + [ + "parent", + (span: Record) => { + span.parent_span_id = "root"; + }, + ], + [ + "name", + (span: Record) => { + span.name = "different provider bookkeeping"; + }, + ], + [ + "status", + (span: Record) => { + span.status = { code: "ERROR" }; + }, + ], + [ + "latency", + (span: Record) => { + span.latency_ms = 2; + }, + ], + [ + "input", + (span: Record) => { + objectValue(span.attributes)["http.request.body"] = { prompt: "wrong" }; + }, + ], + [ + "output", + (span: Record) => { + objectValue(span.attributes)["http.response.body"] = { answer: "wrong" }; + }, + ], + [ + "attribute", + (span: Record) => { + objectValue(span.attributes)["http.route"] = "/wrong"; + }, + ], +] as const)("rejects paired non-semantic span %s mismatch", (_name, mutate) => { + const proof = deployedProofFixture(); + addPairedNonSemanticSpan(proof); + const span = proof.traceRecord.data?.spans?.find( + (candidate) => candidate.span_id === "provider-wrapper", + ); + if (!span) throw new Error("fixture requires provider wrapper"); + mutate(span); + + expect(() => validateDeployedProof(proof)).toThrow( + /parent|parity|semantics/i, + ); +}); + +test("rejects a semantic AGENT child whose filtered parent differs by source", () => { + const proof = deployedProofFixture(); + addPairedNonSemanticSpan(proof); + addPairedNonSemanticSpan(proof, { spanId: "other-provider-wrapper" }); + const mlflowRoot = proof.traceRecord.data?.spans?.find( + (span) => objectValue(span.attributes)["mlflow.spanType"] === "AGENT", + ); + const ucRoot = proof.rows.find( + (row) => row.attributes["mlflow.spanType"] === "AGENT", + ); + if (!mlflowRoot || !ucRoot) throw new Error("fixture requires AGENT roots"); + mlflowRoot.parent_span_id = "provider-wrapper"; + ucRoot.parentSpanId = "other-provider-wrapper"; + + expect(() => validateDeployedProof(proof)).toThrow(/parent|parity/i); +}); + test.each([ [ "missing MLflow AGENT request", From 9df0b9f0f8117ee49faf14b279544f65e9664e18 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 05:18:47 -0700 Subject: [PATCH 26/31] fix: close final tracing conformance gaps Signed-off-by: Adam Gurary --- .../appkit/scripts/provision-mlflow-uc.py | 73 ++++++++--- .../scripts/test_provision_mlflow_uc.py | 56 +++++++-- packages/appkit/src/agents/databricks.ts | 107 +++++++++++----- .../tests/databricks-parser-tracing.test.ts | 73 +++++++++++ .../src/agents/tests/databricks.test.ts | 83 +++++++++++-- .../appkit/src/core/agent/trace-tool-call.ts | 21 ++-- packages/appkit/src/plugins/agents/agents.ts | 5 +- .../agents/tests/dispatch-tool-call.test.ts | 4 + .../agents/tests/route-handler-errors.test.ts | 66 ++++++++++ .../trace-conformance.integration.test.ts | 114 +++++++++++++----- .../src/telemetry/agent-tracing/index.ts | 3 +- .../telemetry/agent-tracing/propagation.ts | 35 +++++- .../telemetry/agent-tracing/serialization.ts | 25 ++++ .../agent-tracing/tests/tracer.test.ts | 57 ++++++++- .../src/telemetry/agent-tracing/tracer.ts | 62 ++++++++-- .../src/cli/commands/setup-mlflow-uc.test.ts | 4 + packages/shared/src/cli/commands/setup.ts | 18 +++ template/README.md | 26 ++++ template/_gitignore | 1 - tools/tests/generate-app-templates.test.ts | 19 +++ 20 files changed, 737 insertions(+), 115 deletions(-) create mode 100644 packages/appkit/src/agents/tests/databricks-parser-tracing.test.ts diff --git a/packages/appkit/scripts/provision-mlflow-uc.py b/packages/appkit/scripts/provision-mlflow-uc.py index f0f969668..205685c03 100644 --- a/packages/appkit/scripts/provision-mlflow-uc.py +++ b/packages/appkit/scripts/provision-mlflow-uc.py @@ -124,6 +124,47 @@ def _grant_trace_access( _execute(workspace, warehouse_id, statement) +def _verify_trace_access( + workspace: Any, + warehouse_id: str, + principal: str, + catalog_name: str, + schema_name: str, + table_names: list[str], +) -> None: + targets = [ + (f"CATALOG {_quoted_identifier(catalog_name)}", {"USE CATALOG"}), + ( + f"SCHEMA {_quoted_identifier(catalog_name)}.{_quoted_identifier(schema_name)}", + {"USE SCHEMA"}, + ), + *[ + ( + "TABLE " + f"{_quoted_identifier(catalog_name)}.{_quoted_identifier(schema_name)}." + f"{_quoted_identifier(table_name)}", + {"MODIFY", "SELECT"}, + ) + for table_name in table_names + ], + ] + for target, required in targets: + response = _execute(workspace, warehouse_id, f"SHOW GRANTS ON {target}") + rows = getattr(getattr(response, "result", None), "data_array", None) or [] + observed = { + str(value).upper() + for row in rows + if any(str(value) == principal for value in row) + for value in row + } + missing = required - observed + if missing: + raise RuntimeError( + f"Runtime principal {principal!r} lacks explicit " + f"{', '.join(sorted(missing))} on {target}" + ) + + def provision_mlflow_uc( *, profile: str, @@ -132,10 +173,14 @@ def provision_mlflow_uc( schema_name: str, table_prefix: str, warehouse_id: str, + runtime_principal: str, mlflow_module: Any | None = None, workspace: Any | None = None, unity_catalog_type: type | None = None, ) -> dict[str, str]: + runtime_principal = runtime_principal.strip() + if not runtime_principal: + raise ValueError("The deployed app runtime principal is required for UC grants") if mlflow_module is None: import mlflow as mlflow_module if unity_catalog_type is None: @@ -183,26 +228,18 @@ def provision_mlflow_uc( f"was not created; discovered: {', '.join(table_names) or ''}" ) - current_user = workspace.current_user.me() - principal = next( - ( - value - for value in ( - getattr(current_user, "user_name", None), - getattr(current_user, "application_id", None), - getattr(current_user, "id", None), - ) - if isinstance(value, str) and value - ), - None, - ) - if principal is None: - raise RuntimeError("Could not resolve the principal receiving UC trace grants") - _grant_trace_access( workspace, warehouse_id, - principal, + runtime_principal, + catalog_name, + schema_name, + table_names, + ) + _verify_trace_access( + workspace, + warehouse_id, + runtime_principal, catalog_name, schema_name, table_names, @@ -244,6 +281,7 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--schema", required=True) parser.add_argument("--table-prefix", required=True) parser.add_argument("--warehouse-id", required=True) + parser.add_argument("--runtime-principal", required=True) parser.add_argument("--output-json", type=Path, required=True) return parser.parse_args() @@ -257,6 +295,7 @@ def main() -> None: schema_name=args.schema, table_prefix=args.table_prefix, warehouse_id=args.warehouse_id, + runtime_principal=args.runtime_principal, ) write_output_atomically(args.output_json, values) print(json.dumps(values, sort_keys=True)) diff --git a/packages/appkit/scripts/test_provision_mlflow_uc.py b/packages/appkit/scripts/test_provision_mlflow_uc.py index 11ddf1143..efe0ce88b 100644 --- a/packages/appkit/scripts/test_provision_mlflow_uc.py +++ b/packages/appkit/scripts/test_provision_mlflow_uc.py @@ -68,6 +68,18 @@ def execute_statement(self, statement: str, warehouse_id: str, **_kwargs): return statement_response("PENDING") if "information_schema.tables" in statement: return statement_response("SUCCEEDED", table_names=self.table_names) + if statement.startswith("SHOW GRANTS ON"): + if " ON CATALOG " in statement: + rows = [["runtime-app-sp", "USE CATALOG"]] + elif " ON SCHEMA " in statement: + rows = [["runtime-app-sp", "USE SCHEMA"]] + else: + rows = [["runtime-app-sp", "MODIFY"], ["runtime-app-sp", "SELECT"]] + return SimpleNamespace( + statement_id="statement-1", + status=SimpleNamespace(state="SUCCEEDED", error=None), + result=SimpleNamespace(data_array=rows), + ) return statement_response("SUCCEEDED") def get_statement(self, statement_id: str): @@ -118,6 +130,7 @@ def provision( schema_name="agent_traces", table_prefix=table_prefix, warehouse_id="0123456789abcdef", + runtime_principal="runtime-app-sp", mlflow_module=mlflow_module, workspace=workspace, unity_catalog_type=UnityCatalog, @@ -151,18 +164,43 @@ def test_provisions_supported_uc_location_and_grants_every_discovered_table( "MLFLOW_UC_TABLE_PREFIX": "appkit", "MLFLOW_OTEL_SPANS_TABLE": "main.agent_traces.appkit_otel_spans", } - assert workspace.statement_execution.statements[1:] == [ - "GRANT USE CATALOG ON CATALOG `main` TO `service-principal`", - "GRANT USE SCHEMA ON SCHEMA `main`.`agent_traces` TO `service-principal`", - "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_otel_spans` TO `service-principal`", - "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_otel_spans` TO `service-principal`", - "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_otel_logs` TO `service-principal`", - "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_otel_logs` TO `service-principal`", - "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_annotations` TO `service-principal`", - "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_annotations` TO `service-principal`", + assert workspace.statement_execution.statements[1:9] == [ + "GRANT USE CATALOG ON CATALOG `main` TO `runtime-app-sp`", + "GRANT USE SCHEMA ON SCHEMA `main`.`agent_traces` TO `runtime-app-sp`", + "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_otel_spans` TO `runtime-app-sp`", + "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_otel_spans` TO `runtime-app-sp`", + "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_otel_logs` TO `runtime-app-sp`", + "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_otel_logs` TO `runtime-app-sp`", + "GRANT MODIFY ON TABLE `main`.`agent_traces`.`appkit_annotations` TO `runtime-app-sp`", + "GRANT SELECT ON TABLE `main`.`agent_traces`.`appkit_annotations` TO `runtime-app-sp`", + ] + verification = workspace.statement_execution.statements[9:] + assert verification == [ + "SHOW GRANTS ON CATALOG `main`", + "SHOW GRANTS ON SCHEMA `main`.`agent_traces`", + "SHOW GRANTS ON TABLE `main`.`agent_traces`.`appkit_otel_spans`", + "SHOW GRANTS ON TABLE `main`.`agent_traces`.`appkit_otel_logs`", + "SHOW GRANTS ON TABLE `main`.`agent_traces`.`appkit_annotations`", ] +def test_runtime_principal_is_required(): + module = load_script() + with pytest.raises(ValueError, match="runtime principal"): + module.provision_mlflow_uc( + profile="DEFAULT", + experiment_name="/Shared/appkit-agent-traces", + catalog_name="main", + schema_name="agent_traces", + table_prefix="appkit", + warehouse_id="0123456789abcdef", + runtime_principal="", + mlflow_module=SimpleNamespace(), + workspace=FakeWorkspace(["appkit_otel_spans"]), + unity_catalog_type=UnityCatalog, + ) + + def test_repeated_setup_is_idempotent(): module = load_script() diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index ba6d0c2e5..f13dd7779 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -1,3 +1,4 @@ +import { SpanStatusCode, trace } from "@opentelemetry/api"; import type { AgentAdapter, AgentEvent, @@ -14,7 +15,12 @@ import { stream as servingStream, } from "../connectors/serving/client"; import { APPKIT_USER_AGENT, getClientOptions } from "../context/client-options"; -import { injectActiveTraceContext } from "../telemetry/agent-tracing"; +import { + captureTraceValue, + injectActiveTraceContext, + normalizeFailureOutput, + verifiedAgentRemoteTrace, +} from "../telemetry/agent-tracing"; import { DEFAULT_TRACE_REDACT_KEYS, REDACTED_TRACE_VALUE, @@ -169,20 +175,7 @@ function remoteTraceFromPayload( (value): value is string => typeof value === "string" && value.trim().length > 0, ); - return spanId - ? { - type: "remote_trace", - traceId, - spanId, - source: "model-serving", - relation: "linked", - } - : { - type: "remote_trace", - traceId, - source: "model-serving", - relation: "continued", - }; + return verifiedAgentRemoteTrace(traceId, spanId, "model-serving"); } function sanitizedModelError(error: unknown): string { @@ -819,13 +812,20 @@ export class DatabricksAdapter implements AgentAdapter { "x-databricks-trace-id", ); if (headerTraceId?.trim()) { - emittedRemoteTraces.add(headerTraceId); - yield { - type: "remote_trace", - traceId: headerTraceId, - source: "model-serving", - relation: "continued", - }; + const headerSpanId = getResponseHeaders(responseBody)?.get( + "x-databricks-span-id", + ); + const remoteTrace = verifiedAgentRemoteTrace( + headerTraceId, + headerSpanId ?? undefined, + "model-serving", + ); + if (remoteTrace) { + emittedRemoteTraces.add( + `${remoteTrace.traceId}:${remoteTrace.spanId ?? ""}`, + ); + yield remoteTrace; + } } reader = responseBody.getReader(); @@ -1119,15 +1119,66 @@ export class DatabricksAdapter implements AgentAdapter { export function parseTextToolCalls( text: string, ): Array<{ name: string; args: unknown }> { + const span = trace + .getTracer("@databricks/appkit-agent-tracing") + .startSpan("databricks text tool-call parser", { + attributes: { + "mlflow.spanType": "PARSER", + "appkit.parser.source": "databricks.text_tool_calls", + }, + }); + const startedAt = Date.now(); + setParserCapturedAttribute(span, "mlflow.spanInputs", { source: text }); const trimmed = text.trim(); + try { + const jsonResult = tryParseLlamaJsonToolCalls(trimmed); + const result = + jsonResult.length > 0 + ? jsonResult + : tryParsePythonStyleToolCalls(trimmed); + const validationFailure = + result.length === 0 && looksLikeTextToolCall(trimmed); + span.setAttribute("appkit.parser.validation_error", validationFailure); + if (validationFailure) { + setParserCapturedAttribute( + span, + "mlflow.spanOutputs", + normalizeFailureOutput([], "Tool-call text failed validation"), + ["error"], + ); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: "Parser validation failed", + }); + } else { + setParserCapturedAttribute(span, "mlflow.spanOutputs", result); + span.setStatus({ code: SpanStatusCode.OK }); + } + return result; + } finally { + span.setAttribute( + "appkit.parser.duration_ms", + Math.max(0, Date.now() - startedAt), + ); + span.end(); + } +} - const jsonResult = tryParseLlamaJsonToolCalls(trimmed); - if (jsonResult.length > 0) return jsonResult; - - const pyResult = tryParsePythonStyleToolCalls(trimmed); - if (pyResult.length > 0) return pyResult; +function looksLikeTextToolCall(text: string): boolean { + return /\[\s*\{|[A-Za-z_][\w.]*\s*\(/.test(text); +} - return []; +function setParserCapturedAttribute( + span: import("@opentelemetry/api").Span, + key: string, + value: unknown, + redactKeys?: readonly string[], +): void { + const captured = captureTraceValue(value, { redactKeys }); + span.setAttribute(key, captured.value); + span.setAttribute(`${key}.original_bytes`, captured.originalBytes); + span.setAttribute(`${key}.sha256`, captured.sha256); + span.setAttribute(`${key}.truncated`, captured.truncated); } function isLlamaToolJsonItem(value: unknown): value is Record< diff --git a/packages/appkit/src/agents/tests/databricks-parser-tracing.test.ts b/packages/appkit/src/agents/tests/databricks-parser-tracing.test.ts new file mode 100644 index 000000000..b0579bac7 --- /dev/null +++ b/packages/appkit/src/agents/tests/databricks-parser-tracing.test.ts @@ -0,0 +1,73 @@ +import { context, SpanStatusCode, trace } from "@opentelemetry/api"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { + BasicTracerProvider, + InMemorySpanExporter, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import { afterAll, afterEach, beforeAll, expect, test, vi } from "vitest"; +import { parseTextToolCalls } from "../databricks"; + +beforeAll(() => { + context.disable(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); +}); + +afterEach(() => vi.restoreAllMocks()); +afterAll(() => context.disable()); + +async function captureParse(source: string) { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], + }); + vi.spyOn(trace, "getTracer").mockImplementation((name, version) => + provider.getTracer(name, version), + ); + const result = parseTextToolCalls(source); + await provider.forceFlush(); + const spans = exporter.getFinishedSpans(); + await provider.shutdown(); + return { result, spans }; +} + +test("traces a valid correctness-changing text tool parse with bounded source and result", async () => { + const { result, spans } = await captureParse( + '[{"name":"analytics.query","parameters":{"sql":"SELECT 1"}}]', + ); + + expect(result).toEqual([ + { name: "analytics.query", args: { sql: "SELECT 1" } }, + ]); + expect(spans).toHaveLength(1); + expect(spans[0].attributes).toMatchObject({ + "mlflow.spanType": "PARSER", + "appkit.parser.source": "databricks.text_tool_calls", + "mlflow.spanInputs": + '{"source":"[{\\"name\\":\\"analytics.query\\",\\"parameters\\":{\\"sql\\":\\"SELECT 1\\"}}]"}', + "mlflow.spanOutputs": + '[{"args":{"sql":"SELECT 1"},"name":"analytics.query"}]', + "appkit.parser.validation_error": false, + }); + expect(spans[0].status.code).toBe(SpanStatusCode.OK); + expect( + spans[0].duration[0] * 1_000 + spans[0].duration[1] / 1_000_000, + ).toBeGreaterThanOrEqual(0); +}); + +test("finalizes malformed tool-like text as a bounded parser failure", async () => { + const source = `[{"name":"analytics.query","parameters":{"secret":"${"x".repeat(80_000)}"}`; + const { result, spans } = await captureParse(source); + + expect(result).toEqual([]); + expect(spans).toHaveLength(1); + const span = spans[0]; + expect(span.status.code).toBe(SpanStatusCode.ERROR); + expect(span.attributes["appkit.parser.validation_error"]).toBe(true); + expect(span.attributes["mlflow.spanInputs.truncated"]).toBe(true); + expect(span.attributes["mlflow.spanOutputs"]).toBe( + '{"error":"[REDACTED]","partial_output":{"available":false,"reason":"no output produced"}}', + ); +}); diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index 0443116a8..cb742085e 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -18,6 +18,7 @@ import { test, vi, } from "vitest"; +import { retainResponseHeaders } from "../../connectors/serving/client"; import { consumeAdapterStream } from "../../core/agent/consume-adapter-stream"; import { DatabricksAdapter, @@ -333,7 +334,6 @@ describe("DatabricksAdapter", () => { expect(events.map((event) => event.type)).toEqual([ "status", "model_start", - "remote_trace", "model_end", "tool_call", "tool_result", @@ -460,14 +460,81 @@ describe("DatabricksAdapter", () => { streamDurationMs: 30, endedAt: 1_055, }); - expect(events[2]).toEqual({ + expect(events.some((event) => event.type === "remote_trace")).toBe(false); + }); + + test("accepts a trace-only response identity only when it proves W3C continuation", async () => { + const adapter = new DatabricksAdapter({ + model: "continued-model", + streamBody: async () => { + const response = await (mockFetch([ + completionChunk({ content: "done", finishReason: "stop" }), + ])("http://test") as unknown as Promise<{ + body: ReadableStream; + }>); + return retainResponseHeaders( + response.body, + new Headers({ "x-databricks-trace-id": TRACE_ID }), + ); + }, + }); + const events: AgentEvent[] = []; + + await withActiveTrace(async () => { + for await (const event of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + events.push(event); + } + }); + + expect(events).toContainEqual({ type: "remote_trace", - traceId: "trace:/main.agent_traces.appkit/remote-step-1", + traceId: TRACE_ID, source: "model-serving", relation: "continued", }); }); + test("turns a foreign response trace and span header into linkable identity", async () => { + const remoteTrace = "11111111111111111111111111111111"; + const remoteSpan = "2222222222222222"; + const adapter = new DatabricksAdapter({ + model: "linked-model", + streamBody: async () => { + const response = await (mockFetch([ + completionChunk({ content: "done", finishReason: "stop" }), + ])("http://test") as unknown as Promise<{ + body: ReadableStream; + }>); + return retainResponseHeaders( + response.body, + new Headers({ + "x-databricks-trace-id": `trace:/main.agent_traces.remote/${remoteTrace}`, + "x-databricks-span-id": remoteSpan, + }), + ); + }, + }); + const events: AgentEvent[] = []; + + for await (const event of adapter.run( + { messages: createTestMessages(), tools: [], threadId: "t1" }, + { executeTool: vi.fn() }, + )) { + events.push(event); + } + + expect(events).toContainEqual({ + type: "remote_trace", + traceId: `trace:/main.agent_traces.remote/${remoteTrace}`, + spanId: remoteSpan, + source: "model-serving", + relation: "linked", + }); + }); + test("preserves a final usage-only frame and marks an unpriced model unavailable", async () => { vi.useFakeTimers(); vi.setSystemTime(2_000); @@ -606,12 +673,13 @@ describe("DatabricksAdapter", () => { }); test("emits a linked remote trace from a terminal MLflow trace event", async () => { + const remoteTrace = "33333333333333333333333333333333"; globalThis.fetch = mockFetch([ completionChunk({ content: "ok", finishReason: "stop", usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, - mlflowTraceId: "trace:/catalog.schema.table/terminal-trace", + mlflowTraceId: `trace:/catalog.schema.table/${remoteTrace}`, mlflowSpanId: "0123456789abcdef", }), ]); @@ -626,7 +694,7 @@ describe("DatabricksAdapter", () => { expect(events).toContainEqual({ type: "remote_trace", - traceId: "trace:/catalog.schema.table/terminal-trace", + traceId: `trace:/catalog.schema.table/${remoteTrace}`, spanId: "0123456789abcdef", source: "model-serving", relation: "linked", @@ -739,7 +807,8 @@ describe("DatabricksAdapter", () => { output_tokens: 3, total_tokens: 12, }, - mlflowTraceId: "trace:/catalog.schema.table/cancelled-remote", + mlflowTraceId: + "trace:/catalog.schema.table/44444444444444444444444444444444", mlflowSpanId: "0123456789abcdef", }) + sseChunk("[DONE]"), @@ -777,7 +846,7 @@ describe("DatabricksAdapter", () => { }, remoteTrace: { type: "remote_trace", - traceId: "trace:/catalog.schema.table/cancelled-remote", + traceId: "trace:/catalog.schema.table/44444444444444444444444444444444", spanId: "0123456789abcdef", source: "model-serving", relation: "linked", diff --git a/packages/appkit/src/core/agent/trace-tool-call.ts b/packages/appkit/src/core/agent/trace-tool-call.ts index 4106f71f0..913f58c0d 100644 --- a/packages/appkit/src/core/agent/trace-tool-call.ts +++ b/packages/appkit/src/core/agent/trace-tool-call.ts @@ -1,6 +1,9 @@ import { type Span, SpanStatusCode, trace } from "@opentelemetry/api"; import type { ToolEffect } from "shared"; -import { captureTraceValue } from "../../telemetry/agent-tracing"; +import { + captureTraceValue, + normalizeFailureOutput, +} from "../../telemetry/agent-tracing"; const tracer = () => trace.getTracer("@databricks/appkit-agent-tracing"); @@ -58,16 +61,16 @@ function setCapturedAttribute(span: Span, key: string, value: unknown): void { } function recordSafeFailure(span: Span, error: unknown): void { - const failure = captureTraceValue( - { - error: - error instanceof Error - ? error.message - : String(error ?? "Unknown error"), - }, + const safeError = + error instanceof Error ? error.message : String(error ?? "Unknown error"); + const errorAttribute = captureTraceValue( + { error: safeError }, { redactKeys: ["error"] }, ); - span.setAttribute("appkit.error", failure.value); + const failure = captureTraceValue(normalizeFailureOutput(undefined, error), { + redactKeys: ["error"], + }); + span.setAttribute("appkit.error", errorAttribute.value); span.setAttribute("mlflow.spanOutputs", failure.value); span.setAttribute("mlflow.spanOutputs.original_bytes", failure.originalBytes); span.setAttribute("mlflow.spanOutputs.sha256", failure.sha256); diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index e05df4240..63c591859 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -1366,7 +1366,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { const signal = abortController.signal; this.trackStream(requestId, userId, abortController); - const tools = Array.from(registered.toolIndex.values()).map((e) => e.def); + const tools = Array.from(registered.toolIndex.values()) + .filter((e) => e.source !== "hosted-supervisor") + .map((e) => e.def); const limits = this.resolvedLimits; const runState: RunState = { @@ -1428,6 +1430,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { tools, threadId: thread.id, signal, + extensions: buildAdapterExtensions(registered.toolIndex), }, { executeTool, signal }, ); diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index fafe40b31..62f07b6b6 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -375,6 +375,10 @@ describe("dispatchToolCall ā€” semantic TOOL spans", () => { "appkit.error": '{"error":"[REDACTED]"}', }); expect(tool.status.code).toBe(SpanStatusCode.ERROR); + expect(JSON.parse(String(tool.attributes["mlflow.spanOutputs"]))).toEqual({ + error: "[REDACTED]", + partial_output: { available: false, reason: "no output produced" }, + }); expect(tool.attributes["appkit.tool.duration_ms"]).toEqual( expect.any(Number), ); diff --git a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts index 7fedb6dab..4cef2657b 100644 --- a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts +++ b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts @@ -624,6 +624,72 @@ describe("POST /invocations & /responses ā€” HITL pre-flight", () => { }); describe("POST /invocations & /responses ā€” successful invoke", () => { + test.each(["invocations", "responses"] as const)( + "passes Supervisor hosted-tool extensions through /%s", + async (route) => { + const observed: unknown[] = []; + const plugin = new AgentsPlugin({ dir: false }); + // biome-ignore lint/suspicious/noExplicitAny: seed private runtime state + (plugin as any).agents.set("default", { + name: "default", + instructions: "hi", + adapter: { + acceptsExtensions: ["databricks.supervisor"], + async *run(input: unknown) { + observed.push(input); + yield { type: "message", content: "done" }; + }, + }, + toolIndex: new Map([ + [ + "genie", + { + source: "hosted-supervisor", + def: { name: "genie", description: "hosted", parameters: {} }, + spec: { + type: "genie_space", + genie_space: { id: "space-1", description: "hosted" }, + }, + }, + ], + ]), + }); + // biome-ignore lint/suspicious/noExplicitAny: seed private runtime state + (plugin as any).defaultAgentName = "default"; + // biome-ignore lint/suspicious/noExplicitAny: stub persistence + (plugin as any).threadStore = { + create: vi.fn().mockResolvedValue({ id: "t-new", messages: [] }), + addMessage: vi.fn(), + delete: vi.fn(), + }; + + const { res } = mockRes(); + await ( + plugin as unknown as { + _handleInvoke: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleInvoke(requestForRoute(route, { input: "hi" }), res); + + expect(observed).toHaveLength(1); + expect(observed[0]).toMatchObject({ + tools: [], + extensions: { + "databricks.supervisor": { + hostedTools: [ + { + type: "genie_space", + genie_space: { id: "space-1", description: "hosted" }, + }, + ], + }, + }, + }); + }, + ); + test("returns OpenAI Responses-shaped JSON with aggregated assistant text", async () => { const plugin = new AgentsPlugin({ dir: false }); // biome-ignore lint/suspicious/noExplicitAny: seed diff --git a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts index dd1ea67e7..8595b23cb 100644 --- a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -533,17 +533,20 @@ async function closeServer(server: Server | undefined): Promise { }); } -async function captureGeneratedHttpTurn( +async function captureGeneratedHttpTurns( candidate: GeneratedCandidate, -): Promise { +): Promise<{ success: TraceManifest; failure: TraceManifest }> { const { name: template, directory } = candidate; const helperPath = join(directory, "server/agents/helper.ts"); const generated = (await import(pathToFileURL(helperPath).href)) as { helper: AgentDefinition; }; const adapter: AgentAdapter = { - async *run(_input: AgentInput, runtime: AgentRunContext) { + async *run(input: AgentInput, runtime: AgentRunContext) { const startedAt = Date.now() - 10; + const injectFailure = JSON.stringify(input.messages).includes( + "INJECT_TRACE_FAILURE", + ); yield { type: "model_start", stepId: "generated-step", @@ -552,6 +555,27 @@ async function captureGeneratedHttpTurn( input: { messages: [{ role: "user", content: "Use count_words" }] }, startedAt, }; + if (injectFailure) { + yield { type: "message_delta", content: "partial" }; + yield { + type: "model_end", + stepId: "generated-step", + model: "generated-test-model", + provider: "databricks", + output: { text: "partial" }, + usage: { + inputTokens: 1, + outputTokens: 1, + totalTokens: 2, + costAvailable: false, + }, + firstTokenAt: startedAt + 2, + streamDurationMs: 10, + endedAt: startedAt + 10, + error: "injected model failure", + }; + return; + } const value = await runtime.executeTool("count_words", { text: "hello traced world", }); @@ -667,35 +691,54 @@ async function captureGeneratedHttpTurn( if (!address || typeof address === "string") { throw new Error(`${template} generated HTTP server did not bind a port`); } - const response = await fetch( - `http://127.0.0.1:${address.port}/invocations`, - { - method: "POST", - headers: { - "Content-Type": "application/json", - "x-forwarded-user": "user-1", - "x-mlflow-session-id": "session-1", - "x-request-id": "request-1", + const invoke = async (input: string) => { + const response = await fetch( + `http://127.0.0.1:${address.port}/invocations`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + "x-forwarded-user": "user-1", + "x-mlflow-session-id": "session-1", + "x-request-id": "request-1", + }, + body: JSON.stringify({ input }), }, - body: JSON.stringify({ - input: "Count the words in hello traced world. Use count_words.", - }), - }, + ); + const responseBody = await response.text(); + expect( + response.ok, + `${template} HTTP ${response.status}: ${responseBody}`, + ).toBe(true); + const traceId = response.headers.get("x-mlflow-trace-id"); + expect( + traceId, + `${template} generated handler trace identity`, + ).toBeTruthy(); + return traceId?.split("/").at(-1) ?? ""; + }; + const successTraceId = await invoke( + "Count the words in hello traced world. Use count_words.", ); - const responseBody = await response.text(); - expect( - response.ok, - `${template} HTTP ${response.status}: ${responseBody}`, - ).toBe(true); - expect( - response.headers.get("x-mlflow-trace-id"), - `${template} generated handler trace identity`, - ).toBeTruthy(); + const failureTraceId = await invoke("INJECT_TRACE_FAILURE"); await provider.forceFlush(); const semanticSpans = exporter .getFinishedSpans() .filter((span) => span.attributes["mlflow.spanType"] !== undefined); - return normalize(template, semanticSpans); + return { + success: normalize( + template, + semanticSpans.filter( + (span) => span.spanContext().traceId === successTraceId, + ), + ), + failure: normalize( + template, + semanticSpans.filter( + (span) => span.spanContext().traceId === failureTraceId, + ), + ), + }; } finally { try { if (appkit) await appkit.shutdown(); @@ -772,7 +815,8 @@ test("every behavior-discovered generated surface executes its HTTP trace proof" ) as TraceManifest; assertContract(reloaded); - const manifest = await captureGeneratedHttpTurn(candidate); + const { success: manifest, failure } = + await captureGeneratedHttpTurns(candidate); expect(process.listenerCount("SIGTERM"), candidate.name).toBe( baselineSigterm, ); @@ -789,11 +833,27 @@ test("every behavior-discovered generated surface executes its HTTP trace proof" ), ).toBe(true); assertContract(manifest); + assertContract(failure); + expect( + failure.spans.some( + (span) => span.status === "ERROR" || span.status === 2, + ), + `${candidate.name} injected failure did not finalize as ERROR: ${JSON.stringify( + failure.spans.map((span) => [span.name, span.status, span.outputs]), + )}`, + ).toBe(true); if ( process.env.APPKIT_TRACE_CONFORMANCE_CANDIDATE === candidate.name && process.env.TRACE_CONFORMANCE_MANIFEST ) { writePythonManifest(process.env.TRACE_CONFORMANCE_MANIFEST, manifest); + const failurePath = process.env.TRACE_CONFORMANCE_FAILURE_MANIFEST; + if (!failurePath) { + throw new Error( + `${candidate.name} owner proof is missing TRACE_CONFORMANCE_FAILURE_MANIFEST`, + ); + } + writePythonManifest(failurePath, failure); } } catch (error) { failures.push( diff --git a/packages/appkit/src/telemetry/agent-tracing/index.ts b/packages/appkit/src/telemetry/agent-tracing/index.ts index cf4b9477d..12e11fc0d 100644 --- a/packages/appkit/src/telemetry/agent-tracing/index.ts +++ b/packages/appkit/src/telemetry/agent-tracing/index.ts @@ -2,8 +2,9 @@ export { attachRemoteTraceLink, injectActiveTraceContext, type RemoteTraceReference, + verifiedAgentRemoteTrace, } from "./propagation"; -export { captureTraceValue } from "./serialization"; +export { captureTraceValue, normalizeFailureOutput } from "./serialization"; export { getActiveAgentTraceIdentity, resolveAgentTraceAppName, diff --git a/packages/appkit/src/telemetry/agent-tracing/propagation.ts b/packages/appkit/src/telemetry/agent-tracing/propagation.ts index 2e8bac365..65b3e9d7a 100644 --- a/packages/appkit/src/telemetry/agent-tracing/propagation.ts +++ b/packages/appkit/src/telemetry/agent-tracing/propagation.ts @@ -9,6 +9,7 @@ import { import { getMlflowUcTraceId } from "../mlflow-uc"; const MLFLOW_V4_TRACE_ID = /^trace:\/[^/]+\/([0-9a-f]{32})$/; +const OTEL_TRACE_ID = /^[0-9a-f]{32}$/; export interface RemoteTraceReference { traceId: string; @@ -17,6 +18,36 @@ export interface RemoteTraceReference { source: "model-serving" | "supervisor" | "mcp" | "remote-agent"; } +export function remoteOtelTraceId(traceId: string): string | undefined { + const normalized = traceId.trim(); + if (OTEL_TRACE_ID.test(normalized)) return normalized; + return MLFLOW_V4_TRACE_ID.exec(normalized)?.[1]; +} + +export function verifiedAgentRemoteTrace( + traceId: string, + spanId: string | undefined, + source: "model-serving" | "supervisor" | "remote-agent", +): import("shared").AgentRemoteTraceEvent | undefined { + const remoteTraceId = remoteOtelTraceId(traceId); + if (!remoteTraceId) return undefined; + const localTraceId = trace.getActiveSpan()?.spanContext().traceId; + if (localTraceId && remoteTraceId === localTraceId) { + return { type: "remote_trace", traceId, source, relation: "continued" }; + } + const normalizedSpanId = spanId?.trim().toLowerCase(); + if (!normalizedSpanId || !/^[0-9a-f]{16}$/.test(normalizedSpanId)) { + return undefined; + } + return { + type: "remote_trace", + traceId, + spanId: normalizedSpanId, + source, + relation: "linked", + }; +} + /** * Replaces any stale W3C headers with the currently active sampled context. * The caller owns the Headers instance and must allocate one per request. @@ -45,7 +76,6 @@ export function attachRemoteTraceLink( span: Span, reference: RemoteTraceReference, ): void { - const mlflowMatch = MLFLOW_V4_TRACE_ID.exec(reference.traceId); const remoteContext = { traceId: reference.otelTraceId, spanId: reference.spanId, @@ -53,8 +83,7 @@ export function attachRemoteTraceLink( isRemote: true, }; if ( - !mlflowMatch || - mlflowMatch[1] !== reference.otelTraceId || + remoteOtelTraceId(reference.traceId) !== reference.otelTraceId || !isSpanContextValid(remoteContext) ) { return; diff --git a/packages/appkit/src/telemetry/agent-tracing/serialization.ts b/packages/appkit/src/telemetry/agent-tracing/serialization.ts index 8fe5ff943..d269d9e51 100644 --- a/packages/appkit/src/telemetry/agent-tracing/serialization.ts +++ b/packages/appkit/src/telemetry/agent-tracing/serialization.ts @@ -47,6 +47,31 @@ export function captureTraceValue( }; } +export function normalizeFailureOutput( + partialOutput: unknown, + error: unknown, +): { partial_output: unknown; error: string } { + const partial = partialOutputValue(partialOutput); + return { + partial_output: + partial === undefined || partial === null || partial === "" + ? { available: false, reason: "no output produced" } + : partial, + error: + error instanceof Error ? error.message : String(error ?? "Unknown error"), + }; +} + +function partialOutputValue(value: unknown): unknown { + if (Array.isArray(value)) return value.length > 0 ? value : undefined; + if (!value || typeof value !== "object" || Array.isArray(value)) return value; + const record = value as Record; + if ("partial_output" in record) return record.partial_output; + if (Object.keys(record).length === 1 && "error" in record) return undefined; + const { error: _error, ...partial } = record; + return Object.keys(partial).length > 0 ? partial : undefined; +} + function normalizeRedactKey(value: string): string { return value.toLowerCase().replaceAll(/[^a-z0-9]/g, ""); } diff --git a/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts index 93c32f93c..a50cd00e4 100644 --- a/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts +++ b/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts @@ -242,6 +242,7 @@ describe("runWithAgentTrace golden span trees", () => { expect(model.parentSpanContext?.spanId).toBe(root.spanContext().spanId); expect(model.attributes).toMatchObject({ "mlflow.spanType": "CHAT_MODEL", + "gen_ai.operation.name": "chat", "mlflow.spanInputs": '{"messages":[{"content":"Hello","role":"user"}]}', "mlflow.spanOutputs": '{"text":"Hello world"}', "mlflow.chat.model": "dbx-claude-sonnet", @@ -250,6 +251,10 @@ describe("runWithAgentTrace golden span trees", () => { '{"cache_read_input_tokens":3,"input_tokens":7,"output_tokens":2,"total_tokens":9}', "gen_ai.usage.input_tokens": 7, "gen_ai.usage.output_tokens": 2, + "gen_ai.usage.cache_read_input_tokens": 3, + "gen_ai.response.model": "dbx-claude-sonnet", + "gen_ai.response.time_to_first_token_ms": 12, + "gen_ai.response.stream_duration_ms": 40, "appkit.cache.read_input_tokens": 3, "appkit.first_token.duration_ms": 12, "appkit.stream.duration_ms": 40, @@ -312,6 +317,49 @@ describe("runWithAgentTrace golden span trees", () => { expect(traced.traceId).toBe(observedTraceId); }); + test("exports a verified linked remote model trace and rejects unverified continuation", async () => { + installTracing(); + const remoteTraceId = + "trace:/main.agent_traces.remote/11111111111111111111111111111111"; + + await runWithAgentTrace( + identity("responses"), + { message: "delegate" }, + async (observer) => { + const [start, ...rest] = modelEvents(); + observer.onEvent(start); + observer.onEvent({ + type: "remote_trace", + traceId: remoteTraceId, + spanId: "2222222222222222", + source: "model-serving", + relation: "linked", + }); + observer.onEvent({ + type: "remote_trace", + traceId: "trace:/main.agent_traces.remote/unverified", + source: "model-serving", + relation: "continued", + }); + for (const event of rest) observer.onEvent(event); + return { text: "done" }; + }, + ); + + const model = finishedSpans().find( + (span) => span.attributes["mlflow.spanType"] === "CHAT_MODEL", + ); + expect(model?.links).toHaveLength(1); + expect(model?.links[0]?.context).toMatchObject({ + traceId: "11111111111111111111111111111111", + spanId: "2222222222222222", + }); + expect(model?.links[0]?.attributes).toMatchObject({ + "mlflow.traceRequestId": remoteTraceId, + "appkit.remote_trace.source": "model-serving", + }); + }); + test("keeps the fallback trace ID active when no tracer provider is installed", async () => { let activeTraceId: string | undefined; @@ -358,7 +406,10 @@ describe("runWithAgentTrace golden span trees", () => { '{"password":"[REDACTED]"}', ); expect(root?.attributes["mlflow.spanOutputs"]).toBe( - '{"error":"[REDACTED]","text":"partial"}', + '{"error":"[REDACTED]","partial_output":"partial"}', + ); + expect(model?.attributes["mlflow.spanOutputs"]).toBe( + '{"error":"[REDACTED]","partial_output":{"text":"partial"}}', ); expect(root?.status.code).toBe(SpanStatusCode.ERROR); expect(model?.status.code).toBe(SpanStatusCode.ERROR); @@ -424,7 +475,7 @@ describe("runWithAgentTrace golden span trees", () => { (span) => span.attributes["mlflow.spanType"] === "AGENT", ); expect(root?.attributes["mlflow.spanOutputs"]).toBe( - '{"error":"[REDACTED]","text":"Hello "}', + '{"error":"[REDACTED]","partial_output":"Hello "}', ); expect(root?.attributes["appkit.cost.available"]).toBe(false); expect(root?.attributes["mlflow.llm.cost"]).toBeUndefined(); @@ -476,7 +527,7 @@ describe("runWithAgentTrace golden span trees", () => { ); expect(root?.status.code).toBe(SpanStatusCode.ERROR); expect(root?.attributes["mlflow.spanOutputs"]).toBe( - '{"error":"[REDACTED]","trace_id":"trace-123"}', + '{"error":"[REDACTED]","partial_output":{"trace_id":"trace-123"}}', ); expect(JSON.stringify(root?.events)).not.toContain(secret); }); diff --git a/packages/appkit/src/telemetry/agent-tracing/tracer.ts b/packages/appkit/src/telemetry/agent-tracing/tracer.ts index af307ae3d..3c360bf1e 100644 --- a/packages/appkit/src/telemetry/agent-tracing/tracer.ts +++ b/packages/appkit/src/telemetry/agent-tracing/tracer.ts @@ -13,7 +13,8 @@ import type { AgentUsage, } from "shared"; import { getMlflowUcTraceId } from "../mlflow-uc"; -import { captureTraceValue } from "./serialization"; +import { attachRemoteTraceLink, remoteOtelTraceId } from "./propagation"; +import { captureTraceValue, normalizeFailureOutput } from "./serialization"; import type { AgentTraceIdentity, AgentTraceObserver, @@ -150,6 +151,20 @@ export async function runWithAgentTrace( } return; } + if (event.type === "remote_trace") { + const active = [...activeModels.values()].at(-1); + if (!active) return; + const otelTraceId = remoteOtelTraceId(event.traceId); + if (event.relation === "linked" && otelTraceId && event.spanId) { + attachRemoteTraceLink(active.span, { + traceId: event.traceId, + otelTraceId, + spanId: event.spanId, + source: event.source, + }); + } + return; + } if (event.type === "status" && event.status === "error") { lifecycleError = true; reportedError ??= event.error; @@ -228,13 +243,17 @@ export async function runWithAgentTrace( setRootUsageAttributes(root, finalUsage); onCompleteUsage?.(finalUsage); const finalOutputText = outputText || textFromValue(value); - const finalOutput = hasExplicitOutput - ? explicitOutput - : failed - ? { text: finalOutputText, error: errorValue(operationError) } - : outputText || textFromValue(value) - ? { text: finalOutputText } - : value; + const finalOutput = + failed || lifecycleError + ? normalizeFailureOutput( + hasExplicitOutput ? explicitOutput : finalOutputText, + operationError ?? reportedError, + ) + : hasExplicitOutput + ? explicitOutput + : outputText || textFromValue(value) + ? { text: finalOutputText } + : value; setCapturedAttribute( root, "mlflow.spanOutputs", @@ -270,6 +289,7 @@ function updateActiveIdentity( function modelStartAttributes(event: AgentModelStartEvent) { return { "mlflow.spanType": "CHAT_MODEL", + "gen_ai.operation.name": "chat", "mlflow.chat.model": event.model, "mlflow.chat.provider": event.provider, "gen_ai.request.model": event.model, @@ -283,7 +303,14 @@ function finalizeModelSpan( event: AgentModelEndEvent, ): void { const { span } = active; - setCapturedAttribute(span, "mlflow.spanOutputs", event.output); + setCapturedAttribute( + span, + "mlflow.spanOutputs", + event.error + ? normalizeFailureOutput(event.output, event.error) + : event.output, + event.error ? ["error"] : undefined, + ); span.setAttribute( "mlflow.chat.tokenUsage", captureTraceValue(mlflowTokenUsage(event.usage)).value, @@ -291,12 +318,20 @@ function finalizeModelSpan( span.setAttribute("gen_ai.usage.input_tokens", event.usage.inputTokens); span.setAttribute("gen_ai.usage.output_tokens", event.usage.outputTokens); if (event.usage.cacheReadInputTokens !== undefined) { + span.setAttribute( + "gen_ai.usage.cache_read_input_tokens", + event.usage.cacheReadInputTokens, + ); span.setAttribute( "appkit.cache.read_input_tokens", event.usage.cacheReadInputTokens, ); } if (event.usage.cacheCreationInputTokens !== undefined) { + span.setAttribute( + "gen_ai.usage.cache_creation_input_tokens", + event.usage.cacheCreationInputTokens, + ); span.setAttribute( "appkit.cache.creation_input_tokens", event.usage.cacheCreationInputTokens, @@ -305,12 +340,21 @@ function finalizeModelSpan( if (event.finishReason) { span.setAttribute("gen_ai.response.finish_reasons", [event.finishReason]); } + span.setAttribute("gen_ai.response.model", event.model); if (event.firstTokenAt !== undefined) { + span.setAttribute( + "gen_ai.response.time_to_first_token_ms", + Math.max(0, event.firstTokenAt - active.event.startedAt), + ); span.setAttribute( "appkit.first_token.duration_ms", Math.max(0, event.firstTokenAt - active.event.startedAt), ); } + span.setAttribute( + "gen_ai.response.stream_duration_ms", + event.streamDurationMs, + ); span.setAttribute("appkit.stream.duration_ms", event.streamDurationMs); setCostAttributes(span, event.usage); if (event.error) { diff --git a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts index ef6dadc9e..0fbe7a910 100644 --- a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts +++ b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts @@ -82,6 +82,7 @@ describe("MLflow UC setup", () => { schema: "agent_traces", tablePrefix: "appkit", warehouseId: "0123456789abcdef", + runtimePrincipal: "runtime-app-sp", }), ).toEqual([ "uv", @@ -103,6 +104,8 @@ describe("MLflow UC setup", () => { "appkit", "--warehouse-id", "0123456789abcdef", + "--runtime-principal", + "runtime-app-sp", "--output-json", "/workspace/traced-app/.databricks/mlflow-uc.json", ]); @@ -130,6 +133,7 @@ describe("MLflow UC setup", () => { schema: "agent_traces", tablePrefix: "appkit", warehouseId: "0123456789abcdef", + runtimePrincipal: "runtime-app-sp", }, { scriptPath: join(cwd, "provision-mlflow-uc.py"), diff --git a/packages/shared/src/cli/commands/setup.ts b/packages/shared/src/cli/commands/setup.ts index 1a0303d19..1b5f511c4 100644 --- a/packages/shared/src/cli/commands/setup.ts +++ b/packages/shared/src/cli/commands/setup.ts @@ -54,6 +54,7 @@ export interface MlflowUcSetupOptions { schema: string; tablePrefix: string; warehouseId: string; + runtimePrincipal: string; } interface MlflowUcSetupDependencies { @@ -116,6 +117,8 @@ export function buildMlflowProvisionCommand( options.tablePrefix, "--warehouse-id", options.warehouseId, + "--runtime-principal", + options.runtimePrincipal, "--output-json", path.join(options.cwd, ".databricks", "mlflow-uc.json"), ]; @@ -390,6 +393,7 @@ interface SetupCliOptions { mlflowSchema: string; mlflowTablePrefix: string; mlflowWarehouseId?: string; + mlflowRuntimePrincipal?: string; } function databricksJson(args: string[]): unknown { @@ -528,6 +532,15 @@ async function runSetup(options: SetupCliOptions) { ); } const userName = resolveDatabricksUser(profile); + const runtimePrincipal = + options.mlflowRuntimePrincipal ?? + process.env.DATABRICKS_APP_SERVICE_PRINCIPAL ?? + env.DATABRICKS_APP_SERVICE_PRINCIPAL; + if (!runtimePrincipal?.trim()) { + throw new Error( + "MLflow UC setup requires --mlflow-runtime-principal (the deployed app service principal application ID)", + ); + } const workspaceHost = resolveWorkspaceHost(profile, env); await provisionAndPersistMlflowUc( { @@ -538,6 +551,7 @@ async function runSetup(options: SetupCliOptions) { schema: options.mlflowSchema, tablePrefix: options.mlflowTablePrefix, warehouseId, + runtimePrincipal: runtimePrincipal.trim(), }, { workspaceHost }, ); @@ -558,6 +572,10 @@ export const setupCommand = new Command("setup") "--mlflow-warehouse-id ", "SQL warehouse used to provision and query UC trace tables", ) + .option( + "--mlflow-runtime-principal ", + "Deployed app service principal receiving explicit UC trace grants", + ) .addHelpText( "after", ` diff --git a/template/README.md b/template/README.md index 621631c13..5e2e309ad 100644 --- a/template/README.md +++ b/template/README.md @@ -12,12 +12,22 @@ A Databricks App powered by [AppKit](https://developers.databricks.com/docs/appk {{- if .plugins.genie}} - **Genie** -- AI/BI Genie conversational interface for natural language data queries {{- end}} +{{- if .plugins.agents}} +- **Agents** -- Composed planner/helper agents with MLflow tracing persisted in Unity Catalog +{{- end}} +{{- if .plugins.files}} +- **Files** -- Governed Databricks Volume browsing +{{- end}} +{{- if .plugins.serving}} +- **Serving** -- Databricks Model Serving invocation +{{- end}} - **Server** -- Express HTTP server with static file serving and Vite dev mode ## Prerequisites - Node.js v22+ and npm - Databricks CLI (for deployment) +- [uv](https://docs.astral.sh/uv/) (for MLflow UC provisioning) - Access to a Databricks workspace ## Databricks Authentication @@ -82,6 +92,22 @@ databricks bundle deploy --profile production ## Getting Started +{{- if .plugins.agents}} +### First-run agent tracing setup + +The generated planner and helper emit one semantic trace per request. Before running the app, identify the deployed app service principal application ID, then provision the experiment, immutable UC trace location, and explicit grants: + +```bash +npm install +npm run setup -- --mlflow-runtime-principal --mlflow-warehouse-id +``` + +Setup grants that runtime identity `USE CATALOG`, `USE SCHEMA`, and explicit `MODIFY` and `SELECT` on every trace table, then verifies each grant. It fails if the runtime principal cannot be resolved or authorized; `ALL PRIVILEGES` is not accepted as proof. + +Run `npm run dev`, open the Agents page, and use the **Open trace in MLflow** link after a turn. Failed root, model, parser, and tool spans retain bounded, redacted `{ partial_output, error }` outputs; runtime export failures do not replace the agent response, while missing startup tracing configuration is fatal. + +{{- end}} + ### Install Dependencies ```bash diff --git a/template/_gitignore b/template/_gitignore index 29895e7c0..23adbc24e 100644 --- a/template/_gitignore +++ b/template/_gitignore @@ -11,4 +11,3 @@ playwright-report/ # Auto-generated types (endpoint-specific, varies per developer) shared/appkit-types/serving.d.ts - diff --git a/tools/tests/generate-app-templates.test.ts b/tools/tests/generate-app-templates.test.ts index e96c4a2a3..a6d1b7654 100644 --- a/tools/tests/generate-app-templates.test.ts +++ b/tools/tests/generate-app-templates.test.ts @@ -42,6 +42,18 @@ describe("generated AppKit agent templates", () => { }, ); + test.each([ + "appkit-analytics", + "appkit-files", + "appkit-genie", + "appkit-lakebase", + "appkit-serving", + ])("%s generates a whitespace-clean gitignore", (name) => { + expect( + readFileSync(join(outputDir, name, ".gitignore"), "utf8"), + ).not.toMatch(/\n\n$/); + }); + test.each(["appkit-agents", "appkit-all-in-one"])( "%s persists all UC tracing configuration and resources", (name) => { @@ -101,6 +113,13 @@ describe("generated AppKit agent templates", () => { expect(packageJson.dependencies["@databricks/appkit-ui"]).toBe("0.60.0"); expect(packageJson.dependencies["@mlflow/core"]).toBeUndefined(); expect(existsSync(join(app, "package-lock.json"))).toBe(false); + const readme = readFileSync(join(app, "README.md"), "utf8"); + expect(readme).toContain("**Agents**"); + expect(readme).toContain("uv"); + expect(readme).toContain("npm run setup -- --mlflow-runtime-principal"); + expect(readme).toContain("USE CATALOG"); + expect(readme).toContain("MODIFY"); + expect(readme).toContain("partial_output"); }, ); From 9c6f1c511f192e6f04d6446f0a057b1055b3e817 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 10:28:24 -0700 Subject: [PATCH 27/31] docs: refresh createApp return type Signed-off-by: Adam Gurary --- docs/docs/api/appkit/Function.createApp.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/docs/api/appkit/Function.createApp.md b/docs/docs/api/appkit/Function.createApp.md index bc656537d..39715ace5 100644 --- a/docs/docs/api/appkit/Function.createApp.md +++ b/docs/docs/api/appkit/Function.createApp.md @@ -8,7 +8,7 @@ function createApp(config: { onPluginsReady?: (appkit: PluginMap) => void | Promise; plugins?: T; telemetry?: TelemetryConfig; -}): Promise>; +}): Promise>; ``` Bootstraps AppKit with the provided configuration. @@ -41,7 +41,7 @@ with an `asUser(req)` method for user-scoped execution. ## Returns -`Promise`\<`PluginMap`\<`T`\>\> +`Promise`\<`AppKitHandle`\<`T`\>\> A `PluginMap` keyed by plugin name with typed exports From c29bee4ffef0a12bcfeb3d09eaf62a9796ba4d17 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 10:50:00 -0700 Subject: [PATCH 28/31] fix: bound tool-call validation time Signed-off-by: Adam Gurary --- packages/appkit/src/agents/databricks.ts | 34 ++++++++++++++++++- .../src/agents/tests/databricks.test.ts | 2 ++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index f13dd7779..b39e72d6b 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -1165,7 +1165,39 @@ export function parseTextToolCalls( } function looksLikeTextToolCall(text: string): boolean { - return /\[\s*\{|[A-Za-z_][\w.]*\s*\(/.test(text); + let arrayStartPending = false; + for (let index = 0; index < text.length; index += 1) { + const character = text[index]; + + if (arrayStartPending) { + if (character === "{") return true; + if (!isWhitespace(character)) { + arrayStartPending = character === "["; + } + } else if (character === "[") { + arrayStartPending = true; + } + + if (!isIdentifierStart(character)) continue; + let cursor = index + 1; + while (cursor < text.length && isIdentifierPart(text[cursor])) cursor += 1; + while (cursor < text.length && isWhitespace(text[cursor])) cursor += 1; + if (text[cursor] === "(") return true; + index = Math.max(index, cursor - 1); + } + return false; +} + +function isIdentifierStart(character: string | undefined): boolean { + return character !== undefined && /[A-Za-z_]/.test(character); +} + +function isIdentifierPart(character: string | undefined): boolean { + return character !== undefined && /[A-Za-z0-9_.]/.test(character); +} + +function isWhitespace(character: string | undefined): boolean { + return character !== undefined && /\s/.test(character); } function setParserCapturedAttribute( diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index cb742085e..182b9f005 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -2185,6 +2185,8 @@ describe("parseTextToolCalls", () => { const cap = 64 * 1024; const filler = "x".repeat(cap); const suffix = "[analytics.query(query='SELECT 1')]"; + const startedAt = performance.now(); expect(parseTextToolCalls(`${filler}${suffix}`)).toEqual([]); + expect(performance.now() - startedAt).toBeLessThan(500); }); }); From 6f14e557b3dd90841e1d20ad666b3c36b10a3921 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 12:44:35 -0700 Subject: [PATCH 29/31] fix: harden MLflow UC tracing integration Signed-off-by: Adam Gurary --- .../api/appkit/Class.DatabricksAdapter.md | 6 +- docs/docs/plugins/stability.md | 5 +- .../appkit/scripts/provision-mlflow-uc.py | 15 +- .../scripts/test_provision_mlflow_uc.py | 13 + packages/appkit/src/agents/databricks.ts | 14 +- .../src/agents/tests/databricks.test.ts | 9 + packages/appkit/src/plugins/agents/agents.ts | 24 +- .../agents/tests/route-handler-errors.test.ts | 66 ++++ .../trace-conformance.integration.test.ts | 6 +- .../telemetry/agent-tracing/serialization.ts | 65 +++- .../agent-tracing/tests/serialization.test.ts | 40 +++ .../appkit/src/telemetry/mlflow-uc/config.ts | 27 +- .../src/telemetry/mlflow-uc/exporter.ts | 340 +++++++++++++++--- .../src/telemetry/mlflow-uc/processor.ts | 82 ++++- .../telemetry/mlflow-uc/tests/config.test.ts | 49 ++- .../mlflow-uc/tests/exporter.test.ts | 297 ++++++++++++++- .../mlflow-uc/tests/processor.test.ts | 36 +- .../src/telemetry/mlflow-uc/trace-info.ts | 9 +- .../telemetry/tests/telemetry-manager.test.ts | 4 +- .../shared/src/cli/commands/doctor/README.md | 12 +- .../cli/commands/doctor/resolve-targets.ts | 5 +- 21 files changed, 988 insertions(+), 136 deletions(-) diff --git a/docs/docs/api/appkit/Class.DatabricksAdapter.md b/docs/docs/api/appkit/Class.DatabricksAdapter.md index 66772a20a..cb779f1d7 100644 --- a/docs/docs/api/appkit/Class.DatabricksAdapter.md +++ b/docs/docs/api/appkit/Class.DatabricksAdapter.md @@ -96,7 +96,8 @@ static fromModelServing(endpointName?: string, options?: ModelServingOptions): P Creates a DatabricksAdapter from a Model Serving endpoint name. Auto-creates a WorkspaceClient internally. Reads the endpoint name -from the argument or the `DATABRICKS_SERVING_ENDPOINT_NAME` env var. +from the argument, the agents resource env var, or the legacy serving +plugin env var (in that order). #### Parameters @@ -112,7 +113,8 @@ from the argument or the `DATABRICKS_SERVING_ENDPOINT_NAME` env var. #### Example ```ts -// Reads endpoint from DATABRICKS_SERVING_ENDPOINT_NAME env var +// Reads DATABRICKS_AGENT_SERVING_ENDPOINT_NAME, falling back to the +// backward-compatible DATABRICKS_SERVING_ENDPOINT_NAME env var const adapter = await DatabricksAdapter.fromModelServing(); // Explicit endpoint diff --git a/docs/docs/plugins/stability.md b/docs/docs/plugins/stability.md index 1a1ba86e4..918d161e7 100644 --- a/docs/docs/plugins/stability.md +++ b/docs/docs/plugins/stability.md @@ -116,7 +116,10 @@ When `plugin sync` discovers non-GA plugins, it includes their stability in the } ``` -Only GA plugins can be marked `requiredByTemplate`. Non-GA plugins always remain optional during init. +Any discovered plugin can be marked `requiredByTemplate` when the template wires +it into `createApp`. Stability still controls how an otherwise optional plugin is +presented, but it does not erase the usage signal for beta plugins that the +generated app requires. ## For Third-Party Plugin Authors diff --git a/packages/appkit/scripts/provision-mlflow-uc.py b/packages/appkit/scripts/provision-mlflow-uc.py index 205685c03..75fbaf23a 100644 --- a/packages/appkit/scripts/provision-mlflow-uc.py +++ b/packages/appkit/scripts/provision-mlflow-uc.py @@ -31,6 +31,19 @@ def _location_name(location: Any) -> str: return repr(location) +def _location_fields(location: Any) -> tuple[str, str, str] | None: + if location is None: + return None + fields = ( + getattr(location, "catalog_name", None), + getattr(location, "schema_name", None), + getattr(location, "table_prefix", None), + ) + if not all(isinstance(value, str) and value for value in fields): + return None + return fields + + def _execute(workspace: Any, warehouse_id: str, statement: str) -> Any: response = workspace.statement_execution.execute_statement( statement=statement, @@ -207,7 +220,7 @@ def provision_mlflow_uc( trace_location=requested_location, ) existing_location = getattr(experiment, "trace_location", None) - if existing_location != requested_location: + if _location_fields(existing_location) != _location_fields(requested_location): raise ValueError( "MLflow experiment trace location is immutable: " f"existing={_location_name(existing_location)}, " diff --git a/packages/appkit/scripts/test_provision_mlflow_uc.py b/packages/appkit/scripts/test_provision_mlflow_uc.py index efe0ce88b..7c76b86a4 100644 --- a/packages/appkit/scripts/test_provision_mlflow_uc.py +++ b/packages/appkit/scripts/test_provision_mlflow_uc.py @@ -210,6 +210,19 @@ def test_repeated_setup_is_idempotent(): assert second == first +def test_accepts_equivalent_trace_location_from_real_sdk_shape(): + module = load_script() + location = SimpleNamespace( + catalog_name="main", + schema_name="agent_traces", + table_prefix="appkit", + ) + + result, _, _ = provision(module, location=location) + + assert result["MLFLOW_EXPERIMENT_ID"] == "123456789" + + def test_existing_different_uc_location_is_rejected_with_both_locations(): module = load_script() diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index b39e72d6b..2c187350a 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -553,11 +553,13 @@ export class DatabricksAdapter implements AgentAdapter { /** * Creates a DatabricksAdapter from a Model Serving endpoint name. * Auto-creates a WorkspaceClient internally. Reads the endpoint name - * from the argument or the `DATABRICKS_SERVING_ENDPOINT_NAME` env var. + * from the argument, the agents resource env var, or the legacy serving + * plugin env var (in that order). * * @example * ```ts - * // Reads endpoint from DATABRICKS_SERVING_ENDPOINT_NAME env var + * // Reads DATABRICKS_AGENT_SERVING_ENDPOINT_NAME, falling back to the + * // backward-compatible DATABRICKS_SERVING_ENDPOINT_NAME env var * const adapter = await DatabricksAdapter.fromModelServing(); * * // Explicit endpoint @@ -575,12 +577,14 @@ export class DatabricksAdapter implements AgentAdapter { options?: ModelServingOptions, ): Promise { const resolvedEndpoint = - endpointName ?? process.env.DATABRICKS_SERVING_ENDPOINT_NAME; + endpointName ?? + process.env.DATABRICKS_AGENT_SERVING_ENDPOINT_NAME ?? + process.env.DATABRICKS_SERVING_ENDPOINT_NAME; if (!resolvedEndpoint) { throw new Error( - "No endpoint name provided and DATABRICKS_SERVING_ENDPOINT_NAME env var is not set. " + - "Pass an endpoint name or set DATABRICKS_SERVING_ENDPOINT_NAME.", + "No endpoint name provided and neither DATABRICKS_AGENT_SERVING_ENDPOINT_NAME nor " + + "DATABRICKS_SERVING_ENDPOINT_NAME is set. Pass an endpoint name or bind an agents serving endpoint.", ); } diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index 182b9f005..02ad215f3 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -2061,6 +2061,15 @@ describe("DatabricksAdapter.fromModelServing", () => { process.env = originalEnv; }); + test("reads endpoint from the agents resource env var", async () => { + delete process.env.DATABRICKS_SERVING_ENDPOINT_NAME; + process.env.DATABRICKS_AGENT_SERVING_ENDPOINT_NAME = "agents-model"; + + const adapter = await DatabricksAdapter.fromModelServing(); + + expect(adapter).toBeInstanceOf(DatabricksAdapter); + }); + test("reads endpoint from DATABRICKS_SERVING_ENDPOINT_NAME env var", async () => { process.env.DATABRICKS_SERVING_ENDPOINT_NAME = "my-model"; diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 42240bf46..6c5c0f24f 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -503,7 +503,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ); } catch (err) { throw new Error( - `Agent '${name}' has no model configured and no DATABRICKS_SERVING_ENDPOINT_NAME default available`, + `Agent '${name}' has no model configured and neither DATABRICKS_AGENT_SERVING_ENDPOINT_NAME nor DATABRICKS_SERVING_ENDPOINT_NAME is available`, { cause: err instanceof Error ? err : undefined }, ); } @@ -1001,13 +1001,23 @@ export class AgentsPlugin extends Plugin implements ToolProvider { registered: RegisteredAgent, ): string[] { if (!this.resolvedApprovalPolicy.requireForDestructive) return []; - const names: string[] = []; - for (const entry of registered.toolIndex.values()) { - if (requiresApproval(entry.def.annotations)) { - names.push(entry.def.name); + const names = new Set(); + const visitedAgents = new Set(); + const visit = (agent: RegisteredAgent): void => { + if (visitedAgents.has(agent.name)) return; + visitedAgents.add(agent.name); + for (const entry of agent.toolIndex.values()) { + if (requiresApproval(entry.def.annotations)) { + names.add(entry.def.name); + } + if (entry.source === "subagent") { + const child = this.agents.get(entry.agentName); + if (child) visit(child); + } } - } - return names; + }; + visit(registered); + return [...names]; } /** diff --git a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts index 092a84ad1..42f98aac2 100644 --- a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts +++ b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts @@ -569,6 +569,72 @@ describe("POST /invocations & /responses ā€” HITL pre-flight", () => { } }); + test("rejects when a nested sub-agent exposes an approval-gated tool", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const childToolIndex = new Map(); + childToolIndex.set("delete_records", { + source: "function", + def: { + name: "delete_records", + description: "deletes records", + parameters: { type: "object", properties: {} }, + annotations: { effect: "destructive" }, + }, + }); + const parentToolIndex = new Map(); + parentToolIndex.set("agent-helper", { + source: "subagent", + agentName: "helper", + def: { + name: "agent-helper", + description: "delegate to helper", + parameters: { type: "object", properties: {} }, + }, + }); + // biome-ignore lint/suspicious/noExplicitAny: seed private state + (plugin as any).agents.set("default", { + name: "default", + instructions: "delegate", + adapter: { async *run() {} }, + toolIndex: parentToolIndex, + }); + // biome-ignore lint/suspicious/noExplicitAny: seed private state + (plugin as any).agents.set("helper", { + name: "helper", + instructions: "delete when asked", + adapter: { async *run() {} }, + toolIndex: childToolIndex, + }); + // biome-ignore lint/suspicious/noExplicitAny: seed private state + (plugin as any).defaultAgentName = "default"; + // biome-ignore lint/suspicious/noExplicitAny: prove the pre-flight rejects before execution + (plugin as any)._runAgentNonStreaming = vi.fn(async () => undefined); + // biome-ignore lint/suspicious/noExplicitAny: stub + (plugin as any).threadStore = { + create: vi.fn().mockResolvedValue({ id: "t-1", messages: [] }), + addMessage: vi.fn(), + }; + + const { res, json } = mockRes(); + await ( + plugin as unknown as { + _handleInvoke: ( + r: express.Request, + w: express.Response, + ) => Promise; + } + )._handleInvoke(mockReq({ input: "hi" }), res); + + expect(res.status).toHaveBeenCalledWith(400); + expect(json).toHaveBeenCalledWith( + expect.objectContaining({ + error: expect.stringMatching(/delete_records/), + }), + ); + // biome-ignore lint/suspicious/noExplicitAny: rejection must happen before adapter execution + expect((plugin as any)._runAgentNonStreaming).not.toHaveBeenCalled(); + }); + test("passes pre-flight when approval.requireForDestructive is disabled", async () => { const plugin = seedPluginWithTools( { effect: "destructive" }, diff --git a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts index 8595b23cb..992580c01 100644 --- a/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -627,7 +627,7 @@ async function captureGeneratedHttpTurns( PGDATABASE: "appkit", PGPORT: "5432", PGSSLMODE: "require", - MLFLOW_EXPERIMENT_ID: "test-experiment", + MLFLOW_EXPERIMENT_ID: "123456789", MLFLOW_TRACING_SQL_WAREHOUSE_ID: "test-warehouse", MLFLOW_UC_CATALOG: "main", MLFLOW_UC_SCHEMA: "agent_traces", @@ -1237,7 +1237,7 @@ function deriveUcBinding( catalog_name: catalog, schema_name: schema, table_prefix: prefix, - otel_spans_table_name: spansTable, + spans_table_name: spansTable, }, }, spansTable, @@ -1863,7 +1863,7 @@ test("derives the exact immutable UC location, table, and MLflow trace prefix", catalog_name: "main", schema_name: "agent_traces", table_prefix: "appkit", - otel_spans_table_name: "main.agent_traces.appkit_otel_spans", + spans_table_name: "main.agent_traces.appkit_otel_spans", }, }, spansTable: "main.agent_traces.appkit_otel_spans", diff --git a/packages/appkit/src/telemetry/agent-tracing/serialization.ts b/packages/appkit/src/telemetry/agent-tracing/serialization.ts index d269d9e51..bd9a73272 100644 --- a/packages/appkit/src/telemetry/agent-tracing/serialization.ts +++ b/packages/appkit/src/telemetry/agent-tracing/serialization.ts @@ -15,25 +15,13 @@ export function captureTraceValue( normalizeRedactKey(key), ), ); - const serialized = - JSON.stringify(value, (key, current) => { - if (redactKeys.has(normalizeRedactKey(key))) return REDACTED_TRACE_VALUE; - if ( - current !== null && - typeof current === "object" && - !Array.isArray(current) - ) { - return Object.fromEntries( - Object.keys(current) - .sort() - .map((objectKey) => [ - objectKey, - (current as Record)[objectKey], - ]), - ); - } - return current; - }) ?? "null"; + let serialized: string; + try { + serialized = + JSON.stringify(canonicalTraceValue(value, redactKeys)) ?? "null"; + } catch { + serialized = '"[Unserializable]"'; + } const encoded = Buffer.from(serialized, "utf8"); const maxBytes = normalizeMaxBytes(options.maxBytes); @@ -47,6 +35,45 @@ export function captureTraceValue( }; } +function canonicalTraceValue( + value: unknown, + redactKeys: Set, + ancestors = new Set(), +): unknown { + if (typeof value === "bigint") return `[BigInt:${value.toString()}]`; + if (value === null || typeof value !== "object") return value; + if (ancestors.has(value)) return "[Circular]"; + + ancestors.add(value); + try { + const toJSON = (value as { toJSON?: unknown }).toJSON; + if (typeof toJSON === "function") { + return canonicalTraceValue(toJSON.call(value), redactKeys, ancestors); + } + if (Array.isArray(value)) { + return value.map((item) => + canonicalTraceValue(item, redactKeys, ancestors), + ); + } + return Object.fromEntries( + Object.keys(value) + .sort() + .map((key) => [ + key, + redactKeys.has(normalizeRedactKey(key)) + ? REDACTED_TRACE_VALUE + : canonicalTraceValue( + (value as Record)[key], + redactKeys, + ancestors, + ), + ]), + ); + } finally { + ancestors.delete(value); + } +} + export function normalizeFailureOutput( partialOutput: unknown, error: unknown, diff --git a/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts index c1f222ee5..78f15724d 100644 --- a/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts +++ b/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts @@ -104,4 +104,44 @@ describe("captureTraceValue", () => { expect(result.originalBytes).toBe(70 * 1024 + 2); expect(result.truncated).toBe(true); }); + + test("captures circular references without failing the agent operation", () => { + const value: Record = { prompt: "hello" }; + value.self = value; + + expect(JSON.parse(captureTraceValue(value).value)).toEqual({ + prompt: "hello", + self: "[Circular]", + }); + }); + + test("does not label a shared acyclic value as circular", () => { + const shared = { value: "reused" }; + + expect( + JSON.parse(captureTraceValue({ first: shared, second: shared }).value), + ).toEqual({ + first: { value: "reused" }, + second: { value: "reused" }, + }); + }); + + test("captures BigInt values without failing the agent operation", () => { + expect(JSON.parse(captureTraceValue({ rows: 42n }).value)).toEqual({ + rows: "[BigInt:42]", + }); + }); + + test("preserves native and custom toJSON semantics", () => { + const custom = { + toJSON: () => ({ kind: "custom", value: 7 }), + }; + + expect( + JSON.parse(captureTraceValue({ at: new Date(0), custom }).value), + ).toEqual({ + at: "1970-01-01T00:00:00.000Z", + custom: { kind: "custom", value: 7 }, + }); + }); }); diff --git a/packages/appkit/src/telemetry/mlflow-uc/config.ts b/packages/appkit/src/telemetry/mlflow-uc/config.ts index 49702578c..9d037146c 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/config.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/config.ts @@ -14,6 +14,8 @@ const ENV_FIELDS = { otelSpansTableName: "MLFLOW_OTEL_SPANS_TABLE", } as const satisfies Record; +const UC_IDENTIFIER = /^[A-Za-z_][A-Za-z0-9_]{0,254}$/; + export function resolveMlflowUcConfig( env: NodeJS.ProcessEnv | Record, overrides: Partial = {}, @@ -21,7 +23,7 @@ export function resolveMlflowUcConfig( const resolved = Object.fromEntries( Object.entries(ENV_FIELDS).map(([field, envName]) => [ field, - overrides[field as keyof MlflowUcConfig] ?? env[envName], + (overrides[field as keyof MlflowUcConfig] ?? env[envName])?.trim(), ]), ) as unknown as MlflowUcConfig; const missing = Object.entries(ENV_FIELDS) @@ -34,5 +36,28 @@ export function resolveMlflowUcConfig( ); } + const invalid: string[] = []; + if (!/^\d+$/.test(resolved.experimentId)) { + invalid.push("MLFLOW_EXPERIMENT_ID must be numeric"); + } + for (const [field, envName] of [ + ["catalogName", "MLFLOW_UC_CATALOG"], + ["schemaName", "MLFLOW_UC_SCHEMA"], + ["tablePrefix", "MLFLOW_UC_TABLE_PREFIX"], + ] as const) { + if (!UC_IDENTIFIER.test(resolved[field])) { + invalid.push(`${envName} must be a simple Unity Catalog identifier`); + } + } + const expectedSpansTable = `${resolved.catalogName}.${resolved.schemaName}.${resolved.tablePrefix}_otel_spans`; + if (resolved.otelSpansTableName !== expectedSpansTable) { + invalid.push(`MLFLOW_OTEL_SPANS_TABLE must equal ${expectedSpansTable}`); + } + if (invalid.length > 0) { + throw new Error( + `Invalid MLflow UC tracing configuration: ${invalid.join("; ")}`, + ); + } + return resolved; } diff --git a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts index e2e58d530..f1e6e08a2 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts @@ -8,7 +8,11 @@ import { type WorkspaceClient, } from "../../workspace-client"; import type { MlflowUcConfig } from "./config"; -import { buildMlflowUcTraceInfo, type MlflowUcTraceInfo } from "./trace-info"; +import { + buildMlflowUcTraceInfo, + constructMlflowV4TraceId, + type MlflowUcTraceInfo, +} from "./trace-info"; const logger = createLogger("telemetry:mlflow-uc"); @@ -33,20 +37,58 @@ interface LoggerLike { interface ExporterOptions { createOtlpExporter?: (options: { url: string; - headers: Record; + headers: Record | (() => Promise>); + timeoutMillis?: number; }) => SpanExporter; logger?: LoggerLike; + maxAttempts?: number; + retryDelayMs?: number; + sleep?: (milliseconds: number) => Promise; + operationTimeoutMs?: number; + maxPendingTraces?: number; + maxSpansPerTrace?: number; + pendingTraceTtlMs?: number; + now?: () => number; +} + +interface PendingSpanTrace { + spans: Map; + lastTouchedMs: number; +} + +const MAX_PENDING_TRACES = 10_000; +const MAX_SPANS_PER_TRACE = 10_000; +const PENDING_TRACE_TTL_MS = 5 * 60_000; + +class TraceInfoRequestError extends Error { + constructor( + readonly status: number, + message: string, + readonly retryAfterMs?: number, + ) { + super(message); + } } export class MlflowUcSpanExporter implements SpanExporter, MlflowUcTraceExporter { private readonly inFlight = new Set>(); - private readonly pendingSpans = new Map>(); + private readonly pendingSpans = new Map(); private readonly createOtlpExporter: NonNullable< ExporterOptions["createOtlpExporter"] >; private readonly logger: LoggerLike; + private readonly maxAttempts: number; + private readonly retryDelayMs: number; + private readonly sleep: (milliseconds: number) => Promise; + private readonly operationTimeoutMs: number; + private readonly maxPendingTraces: number; + private readonly maxSpansPerTrace: number; + private readonly pendingTraceTtlMs: number; + private readonly now: () => number; + private otlpExporter?: SpanExporter; + private shutdownPromise?: Promise; private closed = false; constructor( @@ -58,6 +100,26 @@ export class MlflowUcSpanExporter options.createOtlpExporter ?? ((exporterOptions) => new OTLPTraceExporter(exporterOptions)); this.logger = options.logger ?? logger; + this.maxAttempts = Math.max(1, options.maxAttempts ?? 3); + this.retryDelayMs = Math.max(0, options.retryDelayMs ?? 100); + this.sleep = + options.sleep ?? + ((milliseconds) => + new Promise((resolve) => setTimeout(resolve, milliseconds))); + this.operationTimeoutMs = Math.max(1, options.operationTimeoutMs ?? 10_000); + this.maxPendingTraces = Math.max( + 1, + Math.floor(options.maxPendingTraces ?? MAX_PENDING_TRACES), + ); + this.maxSpansPerTrace = Math.max( + 1, + Math.floor(options.maxSpansPerTrace ?? MAX_SPANS_PER_TRACE), + ); + this.pendingTraceTtlMs = Math.max( + 1, + options.pendingTraceTtlMs ?? PENDING_TRACE_TTL_MS, + ); + this.now = options.now ?? Date.now; } export( @@ -65,7 +127,10 @@ export class MlflowUcSpanExporter resultCallback: (result: ExportResult) => void, ): void { if (this.closed) { - resultCallback({ code: ExportResultCode.SUCCESS }); + resultCallback({ + code: ExportResultCode.FAILED, + error: new Error("MLflow UC trace exporter is shut down"), + }); return; } @@ -76,10 +141,17 @@ export class MlflowUcSpanExporter } const operation = (async () => { - for (const batch of batches) { - await this.exportBatchSafely(batch); + try { + for (const batch of batches) { + await this.exportBatchWithRetry(batch); + } + resultCallback({ code: ExportResultCode.SUCCESS }); + } catch (error) { + resultCallback({ + code: ExportResultCode.FAILED, + error: toError(error), + }); } - resultCallback({ code: ExportResultCode.SUCCESS }); })(); this.track(operation); } @@ -89,13 +161,21 @@ export class MlflowUcSpanExporter resultCallback: (result: ExportResult) => void, ): void { if (this.closed) { - resultCallback({ code: ExportResultCode.SUCCESS }); + resultCallback({ + code: ExportResultCode.FAILED, + error: new Error("MLflow UC trace exporter is shut down"), + }); return; } - const operation = this.exportBatchSafely(batch).then(() => { - resultCallback({ code: ExportResultCode.SUCCESS }); - }); + const operation = this.exportBatchWithRetry(batch).then( + () => resultCallback({ code: ExportResultCode.SUCCESS }), + (error) => + resultCallback({ + code: ExportResultCode.FAILED, + error: toError(error), + }), + ); this.track(operation); } @@ -106,24 +186,37 @@ export class MlflowUcSpanExporter } async shutdown(): Promise { - this.closed = true; - await this.forceFlush(); - this.pendingSpans.clear(); + if (!this.shutdownPromise) { + this.closed = true; + this.shutdownPromise = (async () => { + await this.forceFlush(); + this.pendingSpans.clear(); + const otlpExporter = this.otlpExporter; + this.otlpExporter = undefined; + if (otlpExporter) await otlpExporter.shutdown(); + })(); + } + await this.shutdownPromise; } private accumulateReadyBatches(spans: ReadableSpan[]): MlflowUcExportBatch[] { const batches: MlflowUcExportBatch[] = []; + const now = this.now(); + this.evictExpiredPendingTraces(now); for (const [traceId, newSpans] of groupSpansByTrace(spans)) { - const accumulated = this.pendingSpans.get(traceId) ?? new Map(); + const pending = this.pendingSpans.get(traceId); + const accumulated = new Map(pending?.spans); for (const span of newSpans) { accumulated.set(span.spanContext().spanId, span); } - this.pendingSpans.set(traceId, accumulated); const traceSpans = [...accumulated.values()]; const semanticRoot = findSemanticRoot(traceSpans); if (semanticRoot) { - const semanticSpans = findSemanticSubtree(traceSpans, semanticRoot); + const semanticSpans = this.capCompletedTrace( + findSemanticSubtree(traceSpans, semanticRoot), + semanticRoot, + ); this.pendingSpans.delete(traceId); batches.push({ traceInfo: buildMlflowUcTraceInfo( @@ -133,23 +226,100 @@ export class MlflowUcSpanExporter ), spans: semanticSpans, }); + continue; } + + if (!pending) this.ensurePendingSpanCapacity(); + const retained = pending ?? { spans: new Map(), lastTouchedMs: now }; + retained.lastTouchedMs = now; + retained.spans.clear(); + for (const [spanId, span] of accumulated) { + if (retained.spans.size >= this.maxSpansPerTrace) break; + retained.spans.set(spanId, span); + } + this.pendingSpans.set(traceId, retained); } return batches; } - private async exportBatchSafely(batch: MlflowUcExportBatch): Promise { - try { - await this.exportBatch(batch); - } catch (error) { - this.logger.error("MLflow UC trace export failed: %O", { - event: "mlflow_uc_trace_export_failed", - traceId: batch.traceInfo.trace_id, - error: error instanceof Error ? error.message : String(error), - }); + private capCompletedTrace( + spans: ReadableSpan[], + semanticRoot: ReadableSpan, + ): ReadableSpan[] { + if (spans.length <= this.maxSpansPerTrace) return spans; + const retained = spans.slice(0, this.maxSpansPerTrace); + const semanticRootId = semanticRoot.spanContext().spanId; + if (retained.some((span) => span.spanContext().spanId === semanticRootId)) { + return retained; + } + retained[retained.length - 1] = semanticRoot; + return retained; + } + + private ensurePendingSpanCapacity(): void { + while (this.pendingSpans.size >= this.maxPendingTraces) { + const oldestTraceId = this.pendingSpans.keys().next().value as + | string + | undefined; + if (!oldestTraceId) return; + this.evictPendingTrace(oldestTraceId, "capacity"); } } + private evictExpiredPendingTraces(now: number): void { + for (const [traceId, pending] of this.pendingSpans) { + if (now - pending.lastTouchedMs <= this.pendingTraceTtlMs) continue; + this.evictPendingTrace(traceId, "ttl"); + } + } + + private evictPendingTrace(traceId: string, reason: "capacity" | "ttl"): void { + const pending = this.pendingSpans.get(traceId); + if (!pending) return; + this.pendingSpans.delete(traceId); + this.logger.error("Dropped incomplete MLflow UC trace: %O", { + event: "mlflow_uc_incomplete_trace_dropped", + traceId, + reason, + retainedSpans: pending.spans.size, + }); + } + + private async exportBatchWithRetry( + batch: MlflowUcExportBatch, + ): Promise { + let lastError: Error | undefined; + for (let attempt = 1; attempt <= this.maxAttempts; attempt += 1) { + try { + await this.exportBatch(batch); + return; + } catch (error) { + lastError = toError(error); + if (attempt >= this.maxAttempts || !isRetryableExportError(error)) { + break; + } + const retryAfterMs = + error instanceof TraceInfoRequestError + ? error.retryAfterMs + : undefined; + await this.sleep( + Math.min( + this.operationTimeoutMs, + retryAfterMs ?? this.retryDelayMs * 2 ** (attempt - 1), + ), + ); + } + } + const terminalError = + lastError ?? new Error("MLflow UC trace export failed"); + this.logger.error("MLflow UC trace export failed: %O", { + event: "mlflow_uc_trace_export_failed", + traceId: constructMlflowV4TraceId(this.config, batch.traceInfo.trace_id), + error: terminalError.message, + }); + throw terminalError; + } + private track(operation: Promise): void { this.inFlight.add(operation); void operation.then( @@ -159,7 +329,10 @@ export class MlflowUcSpanExporter } private async exportBatch(batch: MlflowUcExportBatch): Promise { - await this.client.config.ensureResolved(); + await this.withDeadline( + this.client.config.ensureResolved(), + "workspace configuration", + ); const configuredHost = this.client.config.host; if (!configuredHost) { throw new Error( @@ -170,11 +343,8 @@ export class MlflowUcSpanExporter const traceInfoHeaders = await this.freshAuthHeaders(); const { traceInfo } = batch; const location = `${this.config.catalogName}.${this.config.schemaName}.${this.config.tablePrefix}`; - const otelTraceId = traceInfo.trace_id.slice( - traceInfo.trace_id.lastIndexOf("/") + 1, - ); const traceInfoResponse = await fetch( - `${host}/api/4.0/mlflow/traces/${encodeURIComponent(location)}/${encodeURIComponent(otelTraceId)}/info`, + `${host}/api/4.0/mlflow/traces/${encodeURIComponent(location)}/${encodeURIComponent(traceInfo.trace_id)}/info`, { method: "POST", headers: { @@ -182,39 +352,113 @@ export class MlflowUcSpanExporter "Content-Type": "application/json", }, body: JSON.stringify(traceInfo), + signal: AbortSignal.timeout(this.operationTimeoutMs), }, ); if (!traceInfoResponse.ok) { - throw new Error( + throw new TraceInfoRequestError( + traceInfoResponse.status, `MLflow trace-info request failed with ${traceInfoResponse.status}: ${await traceInfoResponse.text()}`, + parseRetryAfter(traceInfoResponse.headers.get("retry-after")), ); } - const otlpAuthHeaders = await this.freshAuthHeaders(); - const otlpExporter = this.createOtlpExporter({ + const otlpExporter = this.getOrCreateOtlpExporter(host); + await new Promise((resolve, reject) => { + otlpExporter.export(batch.spans, (result) => { + if (result.code === ExportResultCode.SUCCESS) resolve(); + else reject(result.error ?? new Error("OTLP trace upload failed")); + }); + }); + } + + private getOrCreateOtlpExporter(host: string): SpanExporter { + if (this.otlpExporter) return this.otlpExporter; + this.otlpExporter = this.createOtlpExporter({ url: `${host}/api/2.0/otel/v1/traces`, - headers: { - ...otlpAuthHeaders, + headers: async () => ({ + ...(await this.freshAuthHeaders()), "X-Databricks-UC-Table-Name": this.config.otelSpansTableName, - }, + }), + timeoutMillis: this.operationTimeoutMs, }); - try { - await new Promise((resolve, reject) => { - otlpExporter.export(batch.spans, (result) => { - if (result.code === ExportResultCode.SUCCESS) resolve(); - else reject(result.error ?? new Error("OTLP trace upload failed")); - }); - }); - } finally { - await otlpExporter.shutdown(); - } + return this.otlpExporter; } private async freshAuthHeaders(): Promise> { const headers = new Headers(); - await this.client.config.authenticate(headers); + await this.withDeadline( + this.client.config.authenticate(headers), + "workspace authentication", + ); return Object.fromEntries(headers.entries()); } + + private async withDeadline( + operation: Promise, + label: string, + ): Promise { + let timeout: ReturnType | undefined; + const deadline = new Promise((_resolve, reject) => { + timeout = setTimeout( + () => + reject( + new Error( + `MLflow UC ${label} timed out after ${this.operationTimeoutMs}ms`, + ), + ), + this.operationTimeoutMs, + ); + }); + try { + return await Promise.race([operation, deadline]); + } finally { + if (timeout) clearTimeout(timeout); + } + } +} + +function isRetryableExportError(error: unknown): boolean { + if (error instanceof TraceInfoRequestError) { + return error.status === 408 || error.status === 429 || error.status >= 500; + } + const candidate = error as { name?: unknown; code?: unknown }; + if ( + candidate?.name === "OTLPExporterError" && + typeof candidate.code === "number" + ) { + return ( + candidate.code === 408 || candidate.code === 429 || candidate.code >= 500 + ); + } + if (typeof candidate?.code === "string") { + return new Set([ + "ECONNRESET", + "ECONNREFUSED", + "EPIPE", + "ETIMEDOUT", + "EAI_AGAIN", + "ENOTFOUND", + "ENETUNREACH", + "EHOSTUNREACH", + ]).has(candidate.code); + } + return !toError(error).message.includes( + "Databricks workspace host is unavailable", + ); +} + +function parseRetryAfter(value: string | null): number | undefined { + if (!value) return undefined; + const seconds = Number(value); + if (Number.isFinite(seconds) && seconds >= 0) return seconds * 1000; + const timestamp = Date.parse(value); + if (!Number.isFinite(timestamp)) return undefined; + return Math.max(0, timestamp - Date.now()); +} + +function toError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)); } function groupSpansByTrace(spans: ReadableSpan[]): Map { diff --git a/packages/appkit/src/telemetry/mlflow-uc/processor.ts b/packages/appkit/src/telemetry/mlflow-uc/processor.ts index e169fae56..50394359d 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/processor.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/processor.ts @@ -4,6 +4,7 @@ import type { Span, SpanProcessor, } from "@opentelemetry/sdk-trace-base"; +import { createLogger } from "../../logging/logger"; import type { MlflowUcConfig } from "./config"; import type { MlflowUcExportBatch, MlflowUcTraceExporter } from "./exporter"; import { @@ -18,12 +19,19 @@ import { interface PendingTrace { spans: ReadableSpan[]; memberSpanIds: Set; + lastTouchedMs: number; } +const MAX_PENDING_TRACES = 10_000; +const MAX_SPANS_PER_TRACE = 10_000; +const PENDING_TRACE_TTL_MS = 5 * 60_000; +const logger = createLogger("telemetry:mlflow-uc:processor"); + export class MlflowUcSpanProcessor implements SpanProcessor { private readonly pending = new Map(); private readonly inFlight = new Set>(); private closed = false; + private shutdownPromise?: Promise; constructor( private readonly config: MlflowUcConfig, @@ -33,6 +41,8 @@ export class MlflowUcSpanProcessor implements SpanProcessor { onStart(span: Span, _parentContext: Context): void { if (this.closed) return; + const now = Date.now(); + this.evictExpired(now); const spanContext = span.spanContext(); const otelTraceId = spanContext.traceId; const mlflowTraceId = @@ -48,16 +58,19 @@ export class MlflowUcSpanProcessor implements SpanProcessor { spanContext.spanId, ); if (!pending && registeredRoot) { - pending = { spans: [], memberSpanIds: new Set() }; + this.ensurePendingCapacity(); + pending = { spans: [], memberSpanIds: new Set(), lastTouchedMs: now }; this.pending.set(otelTraceId, pending); } if (!pending) return; + pending.lastTouchedMs = now; if (registeredRoot) { pending.memberSpanIds.add(spanContext.spanId); } else if ( span.parentSpanContext && - pending.memberSpanIds.has(span.parentSpanContext.spanId) + pending.memberSpanIds.has(span.parentSpanContext.spanId) && + pending.memberSpanIds.size < MAX_SPANS_PER_TRACE ) { pending.memberSpanIds.add(spanContext.spanId); } @@ -67,7 +80,8 @@ export class MlflowUcSpanProcessor implements SpanProcessor { if ( pending && span.parentSpanContext && - pending.memberSpanIds.has(span.parentSpanContext.spanId) + pending.memberSpanIds.has(span.parentSpanContext.spanId) && + pending.memberSpanIds.size < MAX_SPANS_PER_TRACE ) { pending.memberSpanIds.add(spanContext.spanId); } @@ -75,6 +89,8 @@ export class MlflowUcSpanProcessor implements SpanProcessor { onEnd(span: ReadableSpan): void { if (this.closed) return; + const now = Date.now(); + this.evictExpired(now); const spanContext = span.spanContext(); const otelTraceId = spanContext.traceId; const pending = this.pending.get(otelTraceId); @@ -84,9 +100,18 @@ export class MlflowUcSpanProcessor implements SpanProcessor { if (!pending.memberSpanIds.has(spanContext.spanId)) { return; } + pending.lastTouchedMs = now; - pending.spans.push(span); - if (spanContext.spanId !== semanticRootSpanId) return; + const isSemanticRoot = spanContext.spanId === semanticRootSpanId; + if (isSemanticRoot) { + if (pending.spans.length >= MAX_SPANS_PER_TRACE) { + pending.spans.length = MAX_SPANS_PER_TRACE - 1; + } + pending.spans.push(span); + } else if (pending.spans.length < MAX_SPANS_PER_TRACE - 1) { + pending.spans.push(span); + } + if (!isSemanticRoot) return; this.pending.delete(otelTraceId); this.registry.deleteTrace(otelTraceId); @@ -105,12 +130,16 @@ export class MlflowUcSpanProcessor implements SpanProcessor { } async shutdown(): Promise { - if (this.closed) return; - this.closed = true; - await this.forceFlush(); - this.pending.clear(); - this.registry.clear(); - await this.exporter.shutdown(); + if (!this.shutdownPromise) { + this.closed = true; + this.shutdownPromise = (async () => { + await this.forceFlush(); + this.pending.clear(); + this.registry.clear(); + await this.exporter.shutdown(); + })(); + } + await this.shutdownPromise; } private startExport(batch: MlflowUcExportBatch): void { @@ -126,4 +155,35 @@ export class MlflowUcSpanProcessor implements SpanProcessor { } void exportComplete.finally(() => this.inFlight.delete(exportComplete)); } + + private ensurePendingCapacity(): void { + while (this.pending.size >= MAX_PENDING_TRACES) { + const oldestTraceId = this.pending.keys().next().value as + | string + | undefined; + if (!oldestTraceId) return; + this.evictTrace(oldestTraceId, "capacity"); + } + } + + private evictExpired(now: number): void { + for (const [traceId, pending] of this.pending) { + if (now - pending.lastTouchedMs <= PENDING_TRACE_TTL_MS) continue; + this.evictTrace(traceId, "ttl"); + } + } + + private evictTrace(traceId: string, reason: "capacity" | "ttl"): void { + const pending = this.pending.get(traceId); + if (!pending) return; + this.pending.delete(traceId); + this.registry.deleteTrace(traceId); + logger.error("Dropped incomplete MLflow UC trace: %O", { + event: "mlflow_uc_incomplete_trace_dropped", + traceId, + reason, + retainedSpans: pending.spans.length, + memberSpans: pending.memberSpanIds.size, + }); + } } diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts index 981b05d4f..dc04fc5fb 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts @@ -2,7 +2,7 @@ import { describe, expect, test } from "vitest"; import { constructMlflowV4TraceId, resolveMlflowUcConfig } from "../../index"; const completeEnv = { - MLFLOW_EXPERIMENT_ID: "experiment-123", + MLFLOW_EXPERIMENT_ID: "123456789", MLFLOW_UC_CATALOG: "main", MLFLOW_UC_SCHEMA: "agent_traces", MLFLOW_UC_TABLE_PREFIX: "appkit", @@ -12,7 +12,7 @@ const completeEnv = { describe("resolveMlflowUcConfig", () => { test("resolves every required field from the environment", () => { expect(resolveMlflowUcConfig(completeEnv)).toEqual({ - experimentId: "experiment-123", + experimentId: "123456789", catalogName: "main", schemaName: "agent_traces", tablePrefix: "appkit", @@ -23,15 +23,16 @@ describe("resolveMlflowUcConfig", () => { test("an explicit object overrides only supplied environment-backed fields", () => { expect( resolveMlflowUcConfig(completeEnv, { - experimentId: "experiment-override", + experimentId: "987654321", tablePrefix: "custom", + otelSpansTableName: "main.agent_traces.custom_otel_spans", }), ).toEqual({ - experimentId: "experiment-override", + experimentId: "987654321", catalogName: "main", schemaName: "agent_traces", tablePrefix: "custom", - otelSpansTableName: "main.agent_traces.appkit_otel_spans", + otelSpansTableName: "main.agent_traces.custom_otel_spans", }); }); @@ -45,6 +46,44 @@ describe("resolveMlflowUcConfig", () => { "MLflow UC tracing configuration missing: MLFLOW_EXPERIMENT_ID, MLFLOW_UC_CATALOG, MLFLOW_UC_SCHEMA, MLFLOW_UC_TABLE_PREFIX, MLFLOW_OTEL_SPANS_TABLE", ); }); + + test("trims environment and override values before using them in URLs and headers", () => { + expect( + resolveMlflowUcConfig( + { + MLFLOW_EXPERIMENT_ID: " 123456789 ", + MLFLOW_UC_CATALOG: " main ", + MLFLOW_UC_SCHEMA: " agent_traces ", + MLFLOW_UC_TABLE_PREFIX: " appkit ", + MLFLOW_OTEL_SPANS_TABLE: " main.agent_traces.appkit_otel_spans ", + }, + { + tablePrefix: " custom ", + otelSpansTableName: " main.agent_traces.custom_otel_spans ", + }, + ), + ).toEqual({ + experimentId: "123456789", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "custom", + otelSpansTableName: "main.agent_traces.custom_otel_spans", + }); + }); + + test("rejects malformed or internally inconsistent UC routing values together", () => { + expect(() => + resolveMlflowUcConfig({ + MLFLOW_EXPERIMENT_ID: "experiment-123", + MLFLOW_UC_CATALOG: "main.bad", + MLFLOW_UC_SCHEMA: "agent traces", + MLFLOW_UC_TABLE_PREFIX: "appkit/unsafe", + MLFLOW_OTEL_SPANS_TABLE: "main.agent_traces.somewhere_else", + }), + ).toThrow( + /MLFLOW_EXPERIMENT_ID.*MLFLOW_UC_CATALOG.*MLFLOW_UC_SCHEMA.*MLFLOW_UC_TABLE_PREFIX.*MLFLOW_OTEL_SPANS_TABLE/, + ); + }); }); test("constructMlflowV4TraceId uses the UC table-prefix location and OTel ID", () => { diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts index 607decd64..f1aa9405d 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts @@ -365,6 +365,111 @@ describe("MlflowUcSpanExporter", () => { ]); }); + test("evicts the oldest incomplete trace when the pending trace limit is reached", async () => { + const first = await createHttpWrappedTraceSpans(); + const second = await createHttpWrappedTraceSpans(); + const third = await createHttpWrappedTraceSpans(); + const uploadedTraceIds: string[] = []; + const logger = { error: vi.fn() }; + const { host } = await startBackend(); + const client = createClient(host, ["token-1", "token-2"]); + const exporter = new MlflowUcSpanExporter(config, client, { + maxPendingTraces: 2, + logger, + createOtlpExporter() { + return { + export(spans, callback) { + uploadedTraceIds.push(spans[0].spanContext().traceId); + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await exportSpans(exporter, [first.http, second.http, third.http]); + await exportSpans(exporter, [first.agent, first.model]); + await exportSpans(exporter, [third.agent, third.model]); + + expect(uploadedTraceIds).toEqual([third.agent.spanContext().traceId]); + expect(logger.error).toHaveBeenCalledWith( + "Dropped incomplete MLflow UC trace: %O", + expect.objectContaining({ + event: "mlflow_uc_incomplete_trace_dropped", + traceId: first.http.spanContext().traceId, + reason: "capacity", + }), + ); + }); + + test("expires incomplete trace fragments before accepting later spans", async () => { + let now = 1_000; + const first = await createHttpWrappedTraceSpans(); + const second = await createHttpWrappedTraceSpans(); + const uploadedTraceIds: string[] = []; + const logger = { error: vi.fn() }; + const { host } = await startBackend(); + const client = createClient(host, ["token-1", "token-2"]); + const exporter = new MlflowUcSpanExporter(config, client, { + pendingTraceTtlMs: 50, + now: () => now, + logger, + createOtlpExporter() { + return { + export(spans, callback) { + uploadedTraceIds.push(spans[0].spanContext().traceId); + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await exportSpans(exporter, [first.http]); + now += 51; + await exportSpans(exporter, [first.agent, first.model]); + await exportSpans(exporter, [second.http, second.agent, second.model]); + + expect(uploadedTraceIds).toEqual([second.agent.spanContext().traceId]); + expect(logger.error).toHaveBeenCalledWith( + "Dropped incomplete MLflow UC trace: %O", + expect.objectContaining({ + event: "mlflow_uc_incomplete_trace_dropped", + traceId: first.http.spanContext().traceId, + reason: "ttl", + }), + ); + }); + + test("caps retained fragments without dropping a semantic root that completes the trace", async () => { + const spans = await createTraceSpans({ nestedAgent: true }); + const root = spans.find((span) => span.name === "support-agent"); + const fragments = spans.filter((span) => span !== root); + if (!root) throw new Error("semantic root fixture is missing"); + const uploads: string[][] = []; + const { host } = await startBackend(); + const client = createClient(host, ["token-1", "token-2"]); + const exporter = new MlflowUcSpanExporter(config, client, { + maxSpansPerTrace: 2, + createOtlpExporter() { + return { + export(spans, callback) { + uploads.push(spans.map((span) => span.name)); + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await exportSpans(exporter, fragments); + await exportSpans(exporter, [root]); + + expect(uploads).toHaveLength(1); + expect(uploads[0]).toHaveLength(2); + expect(uploads[0]).toContain("support-agent"); + }); + test("accepts a semantic AGENT whose missing parent context is remote", async () => { const spans = await createRemoteParentTraceSpans(); const agent = spans[0]; @@ -447,7 +552,7 @@ describe("MlflowUcSpanExporter", () => { const traceInfo = JSON.parse(requests[0].body.toString("utf8")); expect(traceInfo).toEqual({ - trace_id: `trace:/main.agent_traces.appkit/${otelTraceId}`, + trace_id: otelTraceId, client_request_id: "request-1", trace_location: { type: "UC_TABLE_PREFIX", @@ -455,7 +560,7 @@ describe("MlflowUcSpanExporter", () => { catalog_name: "main", schema_name: "agent_traces", table_prefix: "appkit", - otel_spans_table_name: "main.agent_traces.appkit_otel_spans", + spans_table_name: "main.agent_traces.appkit_otel_spans", }, }, request_preview: '{"prompt":"hello"}', @@ -490,7 +595,7 @@ describe("MlflowUcSpanExporter", () => { const spans = await createTraceSpans(); const exportCalls: Array<{ url: string; - headers: Record; + headers: Record | (() => Promise>); spans: ReadableSpan[]; }> = []; const exporter = new MlflowUcSpanExporter(config, client, { @@ -513,27 +618,142 @@ describe("MlflowUcSpanExporter", () => { expect(exportCalls).toHaveLength(1); expect(exportCalls[0]).toMatchObject({ url: `${host}/api/2.0/otel/v1/traces`, - headers: { - authorization: "Bearer token-2", + spans, + }); + const headers = exportCalls[0].headers; + await expect( + typeof headers === "function" ? headers() : headers, + ).resolves.toEqual({ + authorization: "Bearer token-2", + "X-Databricks-UC-Table-Name": "main.agent_traces.appkit_otel_spans", + }); + }); + + test("reuses one OTLP exporter while resolving fresh auth headers per trace", async () => { + const { host } = await startBackend(); + const client = createClient(host, [ + "trace-1", + "otel-1", + "trace-2", + "otel-2", + ]); + const first = await createTraceSpans(); + const second = await createTraceSpans(); + const observedHeaders: Record[] = []; + const shutdown = vi.fn().mockResolvedValue(undefined); + let exporterCreations = 0; + const exporter = new MlflowUcSpanExporter(config, client, { + createOtlpExporter(options) { + exporterCreations += 1; + return { + export(_spans, callback) { + const headers = options.headers as + | Record + | (() => Promise>); + void Promise.resolve( + typeof headers === "function" ? headers() : headers, + ).then((resolved) => { + observedHeaders.push(resolved); + callback({ code: ExportResultCode.SUCCESS }); + }); + }, + shutdown, + } satisfies SpanExporter; + }, + }); + + await expect(exportSpans(exporter, first)).resolves.toMatchObject({ + code: ExportResultCode.SUCCESS, + }); + await expect(exportSpans(exporter, second)).resolves.toMatchObject({ + code: ExportResultCode.SUCCESS, + }); + await exporter.shutdown(); + + expect(exporterCreations).toBe(1); + expect(observedHeaders).toEqual([ + { + authorization: "Bearer otel-1", "X-Databricks-UC-Table-Name": "main.agent_traces.appkit_otel_spans", }, - spans, + { + authorization: "Bearer otel-2", + "X-Databricks-UC-Table-Name": "main.agent_traces.appkit_otel_spans", + }, + ]); + expect(shutdown).toHaveBeenCalledTimes(1); + }); + + test("retries a transient trace-info rejection and succeeds", async () => { + let attempts = 0; + const { host } = await startBackend((request, response) => { + if (request.path.endsWith("/info")) attempts += 1; + response.statusCode = attempts === 1 ? 500 : 200; + response.end(attempts === 1 ? "backend unavailable" : ""); + }); + const client = createClient(host, ["token-1", "token-2", "token-3"]); + const spans = await createTraceSpans(); + const exporter = new MlflowUcSpanExporter(config, client, { + maxAttempts: 2, + retryDelayMs: 0, }); + + await expect(exportSpans(exporter, spans)).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + expect(attempts).toBe(2); }); - test("isolates a backend rejection, calls back with success, and logs one structured error", async () => { - const { host } = await startBackend((_request, response) => { + test("does not retry a non-retryable OTLP client error", async () => { + const { host } = await startBackend(); + const client = createClient(host, ["token-1"]); + const spans = await createTraceSpans(); + let attempts = 0; + const exporter = new MlflowUcSpanExporter(config, client, { + maxAttempts: 3, + retryDelayMs: 0, + createOtlpExporter() { + return { + export(_spans, callback) { + attempts += 1; + callback({ + code: ExportResultCode.FAILED, + error: Object.assign(new Error("bad request"), { + name: "OTLPExporterError", + code: 400, + }), + }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await expect(exportSpans(exporter, spans)).resolves.toMatchObject({ + code: ExportResultCode.FAILED, + }); + expect(attempts).toBe(1); + }); + + test("reports a persistent backend rejection as failed and logs one structured error", async () => { + const { host, requests } = await startBackend((_request, response) => { response.statusCode = 500; response.end("backend unavailable"); }); const logger = { error: vi.fn() }; const client = createClient(host, ["token-1"]); const spans = await createTraceSpans(); - const exporter = new MlflowUcSpanExporter(config, client, { logger }); + const exporter = new MlflowUcSpanExporter(config, client, { + logger, + maxAttempts: 2, + retryDelayMs: 0, + }); await expect(exportSpans(exporter, spans)).resolves.toEqual({ - code: ExportResultCode.SUCCESS, + code: ExportResultCode.FAILED, + error: expect.any(Error), }); + expect(requests).toHaveLength(2); expect(logger.error).toHaveBeenCalledTimes(1); expect(logger.error).toHaveBeenCalledWith( "MLflow UC trace export failed: %O", @@ -552,7 +772,8 @@ describe("MlflowUcSpanExporter", () => { const exporter = new MlflowUcSpanExporter(config, client, { logger }); await expect(exportSpans(exporter, spans)).resolves.toEqual({ - code: ExportResultCode.SUCCESS, + code: ExportResultCode.FAILED, + error: expect.any(Error), }); expect(logger.error).toHaveBeenCalledWith( "MLflow UC trace export failed: %O", @@ -562,6 +783,57 @@ describe("MlflowUcSpanExporter", () => { ); }); + test("bounds a hung trace-info request and reports failure", async () => { + const { host } = await startBackend(async (request, response) => { + if (request.path.endsWith("/info")) { + await new Promise((resolve) => setTimeout(resolve, 100)); + } + response.statusCode = 200; + response.end(); + }); + const client = createClient(host, ["token-1", "token-2"]); + const spans = await createTraceSpans(); + const exporter = new MlflowUcSpanExporter(config, client, { + maxAttempts: 1, + operationTimeoutMs: 25, + }); + const started = Date.now(); + + await expect(exportSpans(exporter, spans)).resolves.toEqual({ + code: ExportResultCode.FAILED, + error: expect.any(Error), + }); + await exporter.forceFlush(); + expect(Date.now() - started).toBeLessThan(80); + }); + + test("bounds workspace authentication before any trace request begins", async () => { + const client = { + config: { + host: "http://127.0.0.1:1", + ensureResolved: vi.fn().mockResolvedValue(undefined), + authenticate: vi.fn(() => new Promise(() => undefined)), + }, + } as unknown as WorkspaceClient; + const spans = await createTraceSpans(); + const exporter = new MlflowUcSpanExporter(config, client, { + maxAttempts: 1, + operationTimeoutMs: 25, + }); + + const result = await Promise.race([ + exportSpans(exporter, spans), + new Promise<{ timedOut: true }>((resolve) => + setTimeout(() => resolve({ timedOut: true }), 100), + ), + ]); + + expect(result).toMatchObject({ + code: ExportResultCode.FAILED, + error: expect.any(Error), + }); + }); + test("shutdown waits for in-flight trace-info and OTLP work", async () => { let releaseTraceInfo!: () => void; const traceInfoReleased = new Promise((resolve) => { @@ -689,7 +961,8 @@ describe("MlflowUcSpanExporter", () => { shutdownComplete = true; }); await expect(exportSpans(exporter, lateTrace)).resolves.toEqual({ - code: ExportResultCode.SUCCESS, + code: ExportResultCode.FAILED, + error: expect.any(Error), }); await Promise.resolve(); expect(shutdownComplete).toBe(false); diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts index 42eb09fed..d58c98625 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts @@ -115,7 +115,7 @@ describe("MlflowUcSpanProcessor", () => { "support-agent", ]); expect(batches[0].traceInfo).toMatchObject({ - trace_id: mlflowTraceId, + trace_id: otelTraceId, client_request_id: "request-1", request_preview: '{"prompt":"hello"}', response_preview: '{"answer":"world"}', @@ -196,6 +196,31 @@ describe("MlflowUcSpanProcessor", () => { await provider.shutdown(); }); + test("concurrent shutdown callers wait for the same exporter barrier", async () => { + let releaseShutdown!: () => void; + const shutdownBarrier = new Promise((resolve) => { + releaseShutdown = resolve; + }); + const exporter: MlflowUcTraceExporter = { + exportTrace: vi.fn(), + forceFlush: vi.fn().mockResolvedValue(undefined), + shutdown: vi.fn(() => shutdownBarrier), + }; + const processor = new MlflowUcSpanProcessor(config, exporter); + + const first = processor.shutdown(); + let secondComplete = false; + const second = processor.shutdown().then(() => { + secondComplete = true; + }); + await Promise.resolve(); + + expect(secondComplete).toBe(false); + releaseShutdown(); + await Promise.all([first, second]); + expect(exporter.shutdown).toHaveBeenCalledTimes(1); + }); + test("keeps nested AGENT spans inside the first semantic root", async () => { const { exporter, batches } = collectingExporter(); const processor = new MlflowUcSpanProcessor(config, exporter); @@ -229,7 +254,7 @@ describe("MlflowUcSpanProcessor", () => { await provider.shutdown(); }); - test("exports every concurrently active semantic root beyond the former registry capacity", async () => { + test("bounds concurrently active semantic roots and evicts the oldest unfinished trace", async () => { const { exporter, batches } = collectingExporter(); const processor = new MlflowUcSpanProcessor(config, exporter); const provider = new BasicTracerProvider({ spanProcessors: [processor] }); @@ -243,9 +268,12 @@ describe("MlflowUcSpanProcessor", () => { for (const root of roots) root.end(); await processor.forceFlush(); - expect(batches).toHaveLength(10_001); + expect(batches).toHaveLength(10_000); expect(new Set(batches.map((batch) => batch.traceInfo.trace_id)).size).toBe( - 10_001, + 10_000, + ); + expect(batches.some((batch) => batch.spans[0]?.name === "agent-0")).toBe( + false, ); await provider.shutdown(); diff --git a/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts b/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts index 8d4a3fa65..cd8e1ca6e 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts @@ -27,7 +27,7 @@ export interface MlflowUcTraceInfo { catalog_name: string; schema_name: string; table_prefix: string; - otel_spans_table_name: string; + spans_table_name: string; }; }; request_preview?: string; @@ -153,7 +153,10 @@ export function buildMlflowUcTraceInfo( const requestTimeMs = hrTimeToMilliseconds(semanticRoot.startTime); const durationMs = hrTimeToMilliseconds(semanticRoot.duration); return { - trace_id: constructMlflowV4TraceId(config, otelTraceId), + // The Databricks V4 wire proto carries the bare 32-hex OTel identity. + // AppKit's public/span alias remains the fully qualified `trace:/...` ID + // constructed by MlflowUcTraceRegistry. + trace_id: otelTraceId, client_request_id: stringAttribute( semanticRoot.attributes, "appkit.request.id", @@ -164,7 +167,7 @@ export function buildMlflowUcTraceInfo( catalog_name: config.catalogName, schema_name: config.schemaName, table_prefix: config.tablePrefix, - otel_spans_table_name: config.otelSpansTableName, + spans_table_name: config.otelSpansTableName, }, }, request_preview: inputs, diff --git a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts index 7eae9c843..7265bea56 100644 --- a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts +++ b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts @@ -241,7 +241,7 @@ describe("TelemetryManager", () => { test("starts the single AppKit provider for valid UC config without generic OTLP", async () => { delete process.env.OTEL_EXPORTER_OTLP_ENDPOINT; - process.env.MLFLOW_EXPERIMENT_ID = "experiment-123"; + process.env.MLFLOW_EXPERIMENT_ID = "123456789"; process.env.MLFLOW_UC_CATALOG = "main"; process.env.MLFLOW_UC_SCHEMA = "agent_traces"; process.env.MLFLOW_UC_TABLE_PREFIX = "appkit"; @@ -261,7 +261,7 @@ describe("TelemetryManager", () => { test("coexists with generic OTLP on the same AppKit provider", async () => { process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; - process.env.MLFLOW_EXPERIMENT_ID = "experiment-123"; + process.env.MLFLOW_EXPERIMENT_ID = "123456789"; process.env.MLFLOW_UC_CATALOG = "main"; process.env.MLFLOW_UC_SCHEMA = "agent_traces"; process.env.MLFLOW_UC_TABLE_PREFIX = "appkit"; diff --git a/packages/shared/src/cli/commands/doctor/README.md b/packages/shared/src/cli/commands/doctor/README.md index 85a519906..fda89b1a5 100644 --- a/packages/shared/src/cli/commands/doctor/README.md +++ b/packages/shared/src/cli/commands/doctor/README.md @@ -45,15 +45,9 @@ pool's `connectionTimeoutMillis`.) `plugin sync` writes `appkit.plugins.json` cataloguing *every* plugin the installed packages ship (for `apps init`), and marks the ones actually wired into `createApp` with `requiredByTemplate: true`. Doctor checks exactly those, so -an unused built-in doesn't produce phantom "missing env var" errors. - -> **Known limitation ā€” non-GA plugins aren't checked yet.** `plugin sync` strips -> `requiredByTemplate` for beta/experimental plugins (its step 6b), so a *used* -> non-GA plugin (e.g. the agents plugin requiring a serving endpoint) is -> currently skipped. A proper fix needs a usage signal that survives sync -> regardless of stability tier (e.g. a separate `used` marker written before the -> strip); that's a `plugin sync` change tracked as a fast-follow. Until then, -> doctor covers GA plugins wired into your app. +an unused built-in doesn't produce phantom "missing env var" errors. The usage +signal is retained for every stability tier, so a beta plugin wired into the app +(for example, agents) receives the same resource checks as a GA plugin. ## Resource provenance: external vs bundle-managed diff --git a/packages/shared/src/cli/commands/doctor/resolve-targets.ts b/packages/shared/src/cli/commands/doctor/resolve-targets.ts index 35e56cdb8..1c5fd1986 100644 --- a/packages/shared/src/cli/commands/doctor/resolve-targets.ts +++ b/packages/shared/src/cli/commands/doctor/resolve-targets.ts @@ -127,9 +127,8 @@ export function targetsFromManifestFile( // `plugin sync` catalogues every plugin the installed packages ship, marking // only those wired into `createApp` with `requiredByTemplate`. Check exactly // those, else doctor reports phantom "missing env var" errors for unimported - // plugins. - // KNOWN LIMITATION: sync strips `requiredByTemplate` for non-GA plugins, so a - // used beta/experimental plugin is not checked here yet. See the README. + // plugins. Stability does not erase that usage signal: a beta plugin wired + // into the generated app must receive the same resource checks as a GA one. const selected = Object.entries(data.plugins ?? {}).filter( ([, plugin]) => plugin.requiredByTemplate === true, ); From 616ef3319dd81e05dd9dbf5ca17823cf76e5f01f Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 13:00:25 -0700 Subject: [PATCH 30/31] fix: address AppKit AI review findings Signed-off-by: Adam Gurary --- packages/appkit/src/core/appkit.ts | 4 +- .../appkit/src/core/tests/databricks.test.ts | 20 ++ packages/appkit/src/plugins/agents/agents.ts | 222 ++++++++++-------- .../plugins/agents/tests/dos-limits.test.ts | 79 ++++++- .../src/cli/commands/setup-mlflow-uc.test.ts | 55 ++++- packages/shared/src/cli/commands/setup.ts | 24 +- 6 files changed, 303 insertions(+), 101 deletions(-) diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index b08b556bc..3e65cfeaf 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -211,7 +211,9 @@ export class AppKit { await TelemetryManager.initialize( { ...config.telemetry, - mlflowUc: agentsEnabled ? requestedMlflowUc || true : requestedMlflowUc, + mlflowUc: agentsEnabled + ? (requestedMlflowUc ?? true) + : requestedMlflowUc, }, config.client, ); diff --git a/packages/appkit/src/core/tests/databricks.test.ts b/packages/appkit/src/core/tests/databricks.test.ts index b3abc5bea..659eaf2cf 100644 --- a/packages/appkit/src/core/tests/databricks.test.ts +++ b/packages/appkit/src/core/tests/databricks.test.ts @@ -5,6 +5,7 @@ import { ServiceContext } from "../../context/service-context"; import { uiVariants } from "../../plugins/ui-variants"; import type { PluginManifest } from "../../registry/types"; import { ResourceType } from "../../registry/types"; +import { TelemetryManager } from "../../telemetry/telemetry-manager"; import { AppKit, createApp } from "../appkit"; const mockReporter = { @@ -226,6 +227,25 @@ describe("AppKit", () => { expect(instance).toBeInstanceOf(AppKit); }); + test("honors an explicit MLflow UC opt-out when agents are enabled", async () => { + const initialize = vi + .spyOn(TelemetryManager, "initialize") + .mockResolvedValue(undefined); + + const instance = await createApp({ + plugins: [ + { plugin: CoreTestPlugin, config: {}, name: "agents" as const }, + ], + telemetry: { mlflowUc: false }, + }); + + expect(initialize).toHaveBeenCalledWith( + expect.objectContaining({ mlflowUc: false }), + undefined, + ); + await instance.shutdown(); + }); + test("should initialize with single plugin", async () => { const pluginData = [ { diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 6c5c0f24f..35a172cee 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -270,12 +270,31 @@ export class AgentsPlugin extends Plugin implements ToolProvider { requestId: string, userId: string, controller: AbortController, - ): void { + ): boolean { + if (this.activeStreams.has(requestId)) return false; this.activeStreams.set(requestId, { controller, userId }); this.userStreamCounts.set( userId, (this.userStreamCounts.get(userId) ?? 0) + 1, ); + return true; + } + + /** Atomically enforce the per-user limit and reserve a server-owned key. */ + private reserveStream( + userId: string, + maxConcurrentStreams: number, + ): { requestId: string; controller: AbortController } | undefined { + if (this.countUserStreams(userId) >= maxConcurrentStreams) { + return undefined; + } + for (;;) { + const requestId = randomUUID(); + const controller = new AbortController(); + if (this.trackStream(requestId, userId, controller)) { + return { requestId, controller }; + } + } } /** @@ -900,7 +919,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { req: express.Request, res: express.Response, observer: AgentTraceObserver, - requestId: string, + traceRequestId: string, ): Promise { const parsed = chatRequestSchema.safeParse(req.body); if (!parsed.success) { @@ -931,59 +950,68 @@ export class AgentsPlugin extends Plugin implements ToolProvider { // their concurrent-stream limit. Prevents a misbehaving client from // churning thread rows while being denied elsewhere. const limits = this.resolvedLimits; - if (this.countUserStreams(userId) >= limits.maxConcurrentStreamsPerUser) { + const reservation = this.reserveStream( + userId, + limits.maxConcurrentStreamsPerUser, + ); + if (!reservation) { res.setHeader("Retry-After", "5"); respondWithTraceError(observer, res, 429, { error: `Too many concurrent streams for this user (limit ${limits.maxConcurrentStreamsPerUser}). Wait for an existing stream to complete before starting another.`, }); return; } + const { requestId, controller } = reservation; - // ThreadStore can throw on backing-storage failures (DB unreachable, - // permission errors, transient I/O). Without a try/catch the - // `async` Express handler bubbles the rejection without a response and - // the client connection hangs until the proxy times out. Surface the - // failure as a 500 so the SSE client falls back instead of waiting. - let thread: Thread; try { - const existing = threadId - ? await this.threadStore.get(threadId, userId) - : null; - if (threadId && !existing) { - respondWithTraceError(observer, res, 404, { - error: `Thread ${threadId} not found`, + // ThreadStore can throw on backing-storage failures (DB unreachable, + // permission errors, transient I/O). Surface the failure as a 500 so + // the client does not hang until the proxy times out. + let thread: Thread; + try { + const existing = threadId + ? await this.threadStore.get(threadId, userId) + : null; + if (threadId && !existing) { + respondWithTraceError(observer, res, 404, { + error: `Thread ${threadId} not found`, + }); + return; + } + thread = existing ?? (await this.threadStore.create(userId)); + observer.updateIdentity({ + threadId: thread.id, + sessionId: requestSessionId(req) ?? thread.id, + }); + + const userMessage: Message = { + id: randomUUID(), + role: "user", + content: message, + createdAt: new Date(), + }; + await this.threadStore.addMessage(thread.id, userId, userMessage); + } catch (err) { + logger.error("threadStore failed in /chat: %O", err); + respondWithTraceError(observer, res, 500, { + error: "Thread operation failed", }); return; } - thread = existing ?? (await this.threadStore.create(userId)); - observer.updateIdentity({ - threadId: thread.id, - sessionId: requestSessionId(req) ?? thread.id, - }); - - const userMessage: Message = { - id: randomUUID(), - role: "user", - content: message, - createdAt: new Date(), - }; - await this.threadStore.addMessage(thread.id, userId, userMessage); - } catch (err) { - logger.error("threadStore failed in /chat: %O", err); - respondWithTraceError(observer, res, 500, { - error: "Thread operation failed", - }); - return; + return await this._streamAgent( + req, + res, + registered, + thread, + userId, + observer, + requestId, + controller, + traceRequestId, + ); + } finally { + this.untrackStream(requestId); } - return this._streamAgent( - req, - res, - registered, - thread, - userId, - observer, - requestId, - ); } /** @@ -1047,7 +1075,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { req: express.Request, res: express.Response, observer: AgentTraceObserver, - requestId: string, + traceRequestId: string, ): Promise { const parsed = invocationsRequestSchema.safeParse(req.body); if (!parsed.success) { @@ -1091,64 +1119,74 @@ export class AgentsPlugin extends Plugin implements ToolProvider { // Match the rate-limit gate on /chat. Without this, a client can bypass // `limits.maxConcurrentStreamsPerUser` by hitting /invocations instead. const limits = this.resolvedLimits; - if (this.countUserStreams(userId) >= limits.maxConcurrentStreamsPerUser) { + const reservation = this.reserveStream( + userId, + limits.maxConcurrentStreamsPerUser, + ); + if (!reservation) { res.setHeader("Retry-After", "5"); respondWithTraceError(observer, res, 429, { error: `Too many concurrent streams for this user (limit ${limits.maxConcurrentStreamsPerUser}). Wait for an existing stream to complete before starting another.`, }); return; } + const { requestId, controller } = reservation; - // Same rationale as `_handleChat`: surface threadStore failures as a - // 500 instead of letting the async handler hang the client connection. - let thread: Thread; try { - thread = await this.threadStore.create(userId); - observer.updateIdentity({ - threadId: thread.id, - sessionId: requestSessionId(req) ?? thread.id, - }); - - if (typeof input === "string") { - await this.threadStore.addMessage(thread.id, userId, { - id: randomUUID(), - role: "user", - content: input, - createdAt: new Date(), + // Same rationale as `_handleChat`: surface threadStore failures as a + // 500 instead of letting the async handler hang the client connection. + let thread: Thread; + try { + thread = await this.threadStore.create(userId); + observer.updateIdentity({ + threadId: thread.id, + sessionId: requestSessionId(req) ?? thread.id, }); - } else { - for (const item of input) { - const role = (item.role ?? "user") as Message["role"]; - const content = - typeof item.content === "string" - ? item.content - : JSON.stringify(item.content ?? ""); - if (!content) continue; + + if (typeof input === "string") { await this.threadStore.addMessage(thread.id, userId, { id: randomUUID(), - role, - content, + role: "user", + content: input, createdAt: new Date(), }); + } else { + for (const item of input) { + const role = (item.role ?? "user") as Message["role"]; + const content = + typeof item.content === "string" + ? item.content + : JSON.stringify(item.content ?? ""); + if (!content) continue; + await this.threadStore.addMessage(thread.id, userId, { + id: randomUUID(), + role, + content, + createdAt: new Date(), + }); + } } + } catch (err) { + logger.error("threadStore failed in /invocations: %O", err); + respondWithTraceError(observer, res, 500, { + error: "Thread operation failed", + }); + return; } - } catch (err) { - logger.error("threadStore failed in /invocations: %O", err); - respondWithTraceError(observer, res, 500, { - error: "Thread operation failed", - }); - return; + return await this._runAgentNonStreaming( + req, + res, + registered, + thread, + userId, + observer, + requestId, + controller, + traceRequestId, + ); + } finally { + this.untrackStream(requestId); } - - return this._runAgentNonStreaming( - req, - res, - registered, - thread, - userId, - observer, - requestId, - ); } private async _streamAgent( @@ -1159,10 +1197,10 @@ export class AgentsPlugin extends Plugin implements ToolProvider { userId: string, observer: AgentTraceObserver, requestId: string, + abortController: AbortController, + traceRequestId: string, ): Promise { - const abortController = new AbortController(); const signal = abortController.signal; - this.trackStream(requestId, userId, abortController); // `hosted-supervisor` entries are not callable from the Node process // (the SA endpoint executes them server-side). Their `def` is a @@ -1198,7 +1236,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { route: "chat", sessionId: requestSessionId(req) ?? thread.id, userId, - requestId, + requestId: traceRequestId, }, }; @@ -1374,10 +1412,10 @@ export class AgentsPlugin extends Plugin implements ToolProvider { userId: string, observer: AgentTraceObserver, requestId: string, + abortController: AbortController, + traceRequestId: string, ): Promise { - const abortController = new AbortController(); const signal = abortController.signal; - this.trackStream(requestId, userId, abortController); const tools = Array.from(registered.toolIndex.values()) .filter((e) => e.source !== "hosted-supervisor") @@ -1404,7 +1442,7 @@ export class AgentsPlugin extends Plugin implements ToolProvider { route: invokeTraceRoute(req), sessionId: requestSessionId(req) ?? thread.id, userId, - requestId, + requestId: traceRequestId, }, }; diff --git a/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts b/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts index 0d01fca8e..2f9bccc59 100644 --- a/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts @@ -16,8 +16,17 @@ import { chatRequestSchema, invocationsRequestSchema } from "../schemas"; * mocked req/res pattern already used by approval-route.test.ts. */ -function mockReq(body: unknown, userId?: string): express.Request { - const headers: Record = {}; +function mockReq( + body: unknown, + userId?: string, + extraHeaders: Record = {}, +): express.Request { + const headers: Record = Object.fromEntries( + Object.entries(extraHeaders).map(([name, value]) => [ + name.toLowerCase(), + value, + ]), + ); if (userId) { headers["x-forwarded-user"] = userId; headers["x-forwarded-access-token"] = "fake-token"; @@ -192,6 +201,72 @@ describe("POST /chat ā€” per-user concurrent-stream limit", () => { expect(res.status).not.toHaveBeenCalledWith(429); }); + test("uses a unique server-generated stream key when clients reuse x-request-id", async () => { + const plugin = seedPlugin(); + const run = vi.fn(async (..._args: unknown[]) => undefined); + (plugin as any)._streamAgent = run; + + await (plugin as any)._handleChat( + mockReq({ message: "first" }, "alice", { + "x-request-id": "client-controlled-id", + }), + mockRes().res, + ); + await (plugin as any)._handleChat( + mockReq({ message: "second" }, "alice", { + "x-request-id": "client-controlled-id", + }), + mockRes().res, + ); + + const streamIds = run.mock.calls.map((call) => call[6]); + expect(streamIds).toHaveLength(2); + expect(new Set(streamIds).size).toBe(2); + expect(streamIds).not.toContain("client-controlled-id"); + }); + + test("reserves the concurrency slot before awaiting thread storage", async () => { + const plugin = seedPlugin({ + dir: false, + limits: { maxConcurrentStreamsPerUser: 1 }, + }); + let releaseFirst!: () => void; + const firstThread = new Promise<{ id: string; messages: [] }>((resolve) => { + releaseFirst = () => resolve({ id: "thread-1", messages: [] }); + }); + let firstCreateStarted!: () => void; + const firstCreateObserved = new Promise((resolve) => { + firstCreateStarted = resolve; + }); + (plugin as any).threadStore = { + get: vi.fn().mockResolvedValue(null), + create: vi + .fn() + .mockImplementationOnce(() => { + firstCreateStarted(); + return firstThread; + }) + .mockResolvedValue({ id: "thread-2", messages: [] }), + addMessage: vi.fn().mockResolvedValue(undefined), + }; + (plugin as any)._streamAgent = vi.fn(async () => undefined); + + const first = (plugin as any)._handleChat( + mockReq({ message: "first" }, "alice"), + mockRes().res, + ); + await firstCreateObserved; + const secondResponse = mockRes(); + await (plugin as any)._handleChat( + mockReq({ message: "second" }, "alice"), + secondResponse.res, + ); + + expect(secondResponse.res.status).toHaveBeenCalledWith(429); + releaseFirst(); + await first; + }); + test("honours agents({ limits: { maxConcurrentStreamsPerUser } })", async () => { const plugin = seedPlugin({ dir: false, diff --git a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts index 0fbe7a910..ec2108987 100644 --- a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts +++ b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts @@ -2,11 +2,12 @@ import { mkdirSync, mkdtempSync, readFileSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import yaml from "js-yaml"; -import { describe, expect, test } from "vitest"; +import { afterEach, describe, expect, test, vi } from "vitest"; import { buildMlflowProvisionCommand, projectRequiresMlflowUc, provisionAndPersistMlflowUc, + setupCommand, } from "./setup"; const EXPECTED_VALUES = { @@ -64,6 +65,10 @@ function createProject(): string { } describe("MLflow UC setup", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + test("agent-enabled projects require UC tracing setup", () => { const cwd = createProject(); @@ -111,6 +116,54 @@ describe("MLflow UC setup", () => { ]); }); + test("preview mode does not provision agent tracing resources", async () => { + const cwd = createProject(); + const packageDirectory = join(cwd, "node_modules", "@databricks", "appkit"); + mkdirSync(packageDirectory, { recursive: true }); + writeFileSync( + join(packageDirectory, "package.json"), + JSON.stringify({ name: "@databricks/appkit" }), + ); + vi.spyOn(process, "cwd").mockReturnValue(cwd); + vi.spyOn(console, "log").mockImplementation(() => undefined); + + await expect( + setupCommand.parseAsync(["node", "setup"]), + ).resolves.toBeDefined(); + + expect(() => readFileSync(join(cwd, "CLAUDE.md"), "utf8")).toThrow(); + expect(() => + readFileSync(join(cwd, ".databricks", "mlflow-uc.json"), "utf8"), + ).toThrow(); + }); + + test("reports an actionable error when uv is unavailable", async () => { + const cwd = createProject(); + + await expect( + provisionAndPersistMlflowUc( + { + cwd, + profile: "DEFAULT", + experimentName: "/Users/user@example.com/appkit-agent-traces", + catalog: "main", + schema: "agent_traces", + tablePrefix: "appkit", + warehouseId: "0123456789abcdef", + runtimePrincipal: "runtime-app-sp", + }, + { + scriptPath: join(cwd, "provision-mlflow-uc.py"), + run() { + throw Object.assign(new Error("spawn uv ENOENT"), { + code: "ENOENT", + }); + }, + }, + ), + ).rejects.toThrow(/requires uv.*install/i); + }); + test("persists all tracing values for local and deployed runtimes", async () => { const cwd = createProject(); const logged: string[] = []; diff --git a/packages/shared/src/cli/commands/setup.ts b/packages/shared/src/cli/commands/setup.ts index 1b5f511c4..54ab38871 100644 --- a/packages/shared/src/cli/commands/setup.ts +++ b/packages/shared/src/cli/commands/setup.ts @@ -257,12 +257,26 @@ export async function provisionAndPersistMlflowUc( const command = buildMlflowProvisionCommand({ ...options, scriptPath }); const run = dependencies.run ?? - ((argv: string[]) => - spawnSync(argv[0], argv.slice(1), { + ((argv: string[]) => { + const result = spawnSync(argv[0], argv.slice(1), { cwd: options.cwd, stdio: "inherit", - }).status ?? 1); - const status = run(command); + }); + if (result.error) throw result.error; + return result.status ?? 1; + }); + let status: number; + try { + status = run(command); + } catch (error) { + if ((error as { code?: unknown })?.code === "ENOENT") { + throw new Error( + "MLflow UC setup requires uv. Install it from https://docs.astral.sh/uv/getting-started/installation/ and rerun appkit setup --write.", + { cause: error }, + ); + } + throw error; + } if (status !== 0) { throw new Error(`MLflow UC provisioning failed with exit code ${status}`); } @@ -514,7 +528,7 @@ async function runSetup(options: SetupCliOptions) { } const cwd = process.cwd(); - if (projectRequiresMlflowUc(cwd, options.mlflowUc === true)) { + if (shouldWrite && projectRequiresMlflowUc(cwd, options.mlflowUc === true)) { const env = readEnvFile(path.join(cwd, ".env")); const profile = process.env.DATABRICKS_CONFIG_PROFILE ?? From 967dc70365de5fea4e666d7c832935327dbdbe44 Mon Sep 17 00:00:00 2001 From: Adam Gurary Date: Thu, 13 Aug 2026 13:13:55 -0700 Subject: [PATCH 31/31] fix: bound MLflow export pressure and preserve diagnostics Signed-off-by: Adam Gurary --- .../appkit/src/connectors/ai-search/client.ts | 29 +++++--- .../connectors/ai-search/tests/client.test.ts | 14 ++-- .../src/telemetry/mlflow-uc/exporter.ts | 1 + .../src/telemetry/mlflow-uc/processor.ts | 73 +++++++++++++++++-- .../mlflow-uc/tests/exporter.test.ts | 37 ++++++++++ .../mlflow-uc/tests/processor.test.ts | 33 +++++++++ .../src/cli/commands/setup-mlflow-uc.test.ts | 18 +++++ packages/shared/src/cli/commands/setup.ts | 50 +++++++------ 8 files changed, 205 insertions(+), 50 deletions(-) diff --git a/packages/appkit/src/connectors/ai-search/client.ts b/packages/appkit/src/connectors/ai-search/client.ts index be7a89716..114292602 100644 --- a/packages/appkit/src/connectors/ai-search/client.ts +++ b/packages/appkit/src/connectors/ai-search/client.ts @@ -359,26 +359,31 @@ function setCapturedAttribute(span: Span, key: string, value: unknown): void { } function recordRetrieverFailure(span: Span, error: unknown): void { - const failure = captureTraceValue( - { - error: - error instanceof Error - ? error.message - : String(error ?? "Unknown error"), - }, - { redactKeys: ["error"] }, - ); + const message = sanitizeRetrieverErrorMessage(error); + const failure = captureTraceValue({ message }); span.setAttribute("appkit.error", failure.value); span.setAttribute("mlflow.spanOutputs", failure.value); span.setAttribute("mlflow.spanOutputs.original_bytes", failure.originalBytes); span.setAttribute("mlflow.spanOutputs.sha256", failure.sha256); span.setAttribute("mlflow.spanOutputs.truncated", failure.truncated); span.recordException({ - name: "Error", - message: "Retriever operation failed", + name: error instanceof Error ? error.name : "Error", + message, }); span.setStatus({ code: SpanStatusCode.ERROR, - message: "Retriever operation failed", + message, }); } + +function sanitizeRetrieverErrorMessage(error: unknown): string { + return ( + error instanceof Error ? error.message : String(error ?? "Unknown error") + ) + .replace(/\bBearer\s+\S+/gi, "Bearer [REDACTED]") + .replace( + /\b(password|passwd|pwd|secret|token|api[_ -]?key)\b\s*(?:[:=]\s*)?\S+/gi, + "$1 [REDACTED]", + ) + .slice(0, 1024); +} diff --git a/packages/appkit/src/connectors/ai-search/tests/client.test.ts b/packages/appkit/src/connectors/ai-search/tests/client.test.ts index 930d324d2..20b25e504 100644 --- a/packages/appkit/src/connectors/ai-search/tests/client.test.ts +++ b/packages/appkit/src/connectors/ai-search/tests/client.test.ts @@ -269,18 +269,18 @@ describe("AiSearchConnector semantic retrieval spans", () => { }); expect(span.status).toEqual({ code: SpanStatusCode.ERROR, - message: "Retriever operation failed", + message: "vector backend exposed password [REDACTED]", }); expect(span.events).toEqual([ expect.objectContaining({ name: "exception", attributes: expect.objectContaining({ - "exception.message": "Retriever operation failed", + "exception.message": "vector backend exposed password [REDACTED]", }), }), ]); expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ - error: "[REDACTED]", + message: "vector backend exposed password [REDACTED]", }); expect(JSON.parse(String(span.attributes["mlflow.spanInputs"]))).toEqual({ columns: ["id", "text"], @@ -396,7 +396,7 @@ describe("AiSearchConnector semantic retrieval spans", () => { const span = retrieverSpan(observed.spans); expect(span.status).toEqual({ code: SpanStatusCode.ERROR, - message: "Retriever operation failed", + message: "Query cancelled before execution", }); expect(span.attributes).toMatchObject({ "appkit.retriever.latency_ms": expect.any(Number), @@ -413,7 +413,7 @@ describe("AiSearchConnector semantic retrieval spans", () => { reranker: null, }); expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ - error: "[REDACTED]", + message: "Query cancelled before execution", }); expect(span.events).toEqual([ expect.objectContaining({ name: "exception" }), @@ -446,7 +446,7 @@ describe("AiSearchConnector semantic retrieval spans", () => { const span = retrieverSpan(observed.spans); expect(span.status).toEqual({ code: SpanStatusCode.ERROR, - message: "Retriever operation failed", + message: "Query cancelled before execution", }); expect(span.attributes).toMatchObject({ "appkit.retriever.latency_ms": expect.any(Number), @@ -459,7 +459,7 @@ describe("AiSearchConnector semantic retrieval spans", () => { queryType: "next_page", }); expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ - error: "[REDACTED]", + message: "Query cancelled before execution", }); expect(span.events).toEqual([ expect.objectContaining({ name: "exception" }), diff --git a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts index f1e6e08a2..1104b9cff 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/exporter.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts @@ -362,6 +362,7 @@ export class MlflowUcSpanExporter parseRetryAfter(traceInfoResponse.headers.get("retry-after")), ); } + await traceInfoResponse.body?.cancel(); const otlpExporter = this.getOrCreateOtlpExporter(host); await new Promise((resolve, reject) => { diff --git a/packages/appkit/src/telemetry/mlflow-uc/processor.ts b/packages/appkit/src/telemetry/mlflow-uc/processor.ts index 50394359d..d5f2d73ca 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/processor.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/processor.ts @@ -22,14 +22,30 @@ interface PendingTrace { lastTouchedMs: number; } +interface ProcessorOptions { + maxConcurrentExports?: number; + maxQueuedExports?: number; +} + +interface QueuedExport { + batch: MlflowUcExportBatch; + resolve: () => void; +} + const MAX_PENDING_TRACES = 10_000; const MAX_SPANS_PER_TRACE = 10_000; const PENDING_TRACE_TTL_MS = 5 * 60_000; +const MAX_CONCURRENT_EXPORTS = 32; +const MAX_QUEUED_EXPORTS = 10_000; const logger = createLogger("telemetry:mlflow-uc:processor"); export class MlflowUcSpanProcessor implements SpanProcessor { private readonly pending = new Map(); private readonly inFlight = new Set>(); + private readonly exportQueue: QueuedExport[] = []; + private readonly maxConcurrentExports: number; + private readonly maxQueuedExports: number; + private activeExports = 0; private closed = false; private shutdownPromise?: Promise; @@ -37,7 +53,17 @@ export class MlflowUcSpanProcessor implements SpanProcessor { private readonly config: MlflowUcConfig, private readonly exporter: MlflowUcTraceExporter, readonly registry = new MlflowUcTraceRegistry(config), - ) {} + options: ProcessorOptions = {}, + ) { + this.maxConcurrentExports = Math.max( + 1, + Math.floor(options.maxConcurrentExports ?? MAX_CONCURRENT_EXPORTS), + ); + this.maxQueuedExports = Math.max( + 1, + Math.floor(options.maxQueuedExports ?? MAX_QUEUED_EXPORTS), + ); + } onStart(span: Span, _parentContext: Context): void { if (this.closed) return; @@ -123,10 +149,10 @@ export class MlflowUcSpanProcessor implements SpanProcessor { } async forceFlush(): Promise { - await this.exporter.forceFlush(); while (this.inFlight.size > 0) { await Promise.allSettled([...this.inFlight]); } + await this.exporter.forceFlush(); } async shutdown(): Promise { @@ -143,17 +169,50 @@ export class MlflowUcSpanProcessor implements SpanProcessor { } private startExport(batch: MlflowUcExportBatch): void { + if ( + this.activeExports >= this.maxConcurrentExports && + this.exportQueue.length >= this.maxQueuedExports + ) { + logger.error("Dropped completed MLflow UC trace before export: %O", { + event: "mlflow_uc_completed_trace_dropped", + traceId: batch.traceInfo.trace_id, + reason: "export_queue_capacity", + queuedExports: this.exportQueue.length, + }); + return; + } let resolveExport!: () => void; const exportComplete = new Promise((resolve) => { resolveExport = resolve; }); this.inFlight.add(exportComplete); - try { - this.exporter.exportTrace(batch, () => resolveExport()); - } catch { - resolveExport(); - } + this.exportQueue.push({ batch, resolve: resolveExport }); void exportComplete.finally(() => this.inFlight.delete(exportComplete)); + this.drainExports(); + } + + private drainExports(): void { + while ( + this.activeExports < this.maxConcurrentExports && + this.exportQueue.length > 0 + ) { + const queued = this.exportQueue.shift(); + if (!queued) return; + this.activeExports += 1; + let completed = false; + const complete = () => { + if (completed) return; + completed = true; + this.activeExports -= 1; + queued.resolve(); + queueMicrotask(() => this.drainExports()); + }; + try { + this.exporter.exportTrace(queued.batch, complete); + } catch { + complete(); + } + } } private ensurePendingCapacity(): void { diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts index f1aa9405d..6c4f341be 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts @@ -33,6 +33,7 @@ const cleanups: Array<() => Promise> = []; afterEach(async () => { await Promise.all(cleanups.splice(0).map((cleanup) => cleanup())); + vi.unstubAllGlobals(); }); async function startBackend( @@ -589,6 +590,42 @@ describe("MlflowUcSpanExporter", () => { }); }); + test("releases the successful trace-info response body before OTLP upload", async () => { + let traceInfoBodyCancelled = false; + vi.stubGlobal( + "fetch", + vi.fn( + async () => + new Response( + new ReadableStream({ + cancel() { + traceInfoBodyCancelled = true; + }, + }), + { status: 200 }, + ), + ), + ); + const client = createClient("https://example.test", ["token-1"]); + const spans = await createTraceSpans(); + const exporter = new MlflowUcSpanExporter(config, client, { + createOtlpExporter() { + return { + export(_spans, callback) { + callback({ code: ExportResultCode.SUCCESS }); + }, + shutdown: vi.fn().mockResolvedValue(undefined), + } satisfies SpanExporter; + }, + }); + + await expect(exportSpans(exporter, spans)).resolves.toEqual({ + code: ExportResultCode.SUCCESS, + }); + + expect(traceInfoBodyCancelled).toBe(true); + }); + test("passes the exact UC header name to a freshly authenticated OTLP exporter", async () => { const { host, requests } = await startBackend(); const client = createClient(host, ["token-1", "token-2"]); diff --git a/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts index d58c98625..b8eed717f 100644 --- a/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts @@ -196,6 +196,39 @@ describe("MlflowUcSpanProcessor", () => { await provider.shutdown(); }); + test("limits concurrent trace exports and starts queued work as capacity frees", async () => { + const callbacks: Array< + (result: { code: ExportResultCode; error?: Error }) => void + > = []; + const { exporter, batches } = collectingExporter((_batch, callback) => { + callbacks.push(callback); + }); + const processor = new MlflowUcSpanProcessor(config, exporter, undefined, { + maxConcurrentExports: 2, + }); + const trees = [ + startTree(processor), + startTree(processor), + startTree(processor), + ]; + + for (const { model, agent } of trees) { + model.end(); + agent.end(); + } + expect(batches).toHaveLength(2); + + callbacks.shift()?.({ code: ExportResultCode.SUCCESS }); + await Promise.resolve(); + expect(batches).toHaveLength(3); + + for (const callback of callbacks) { + callback({ code: ExportResultCode.SUCCESS }); + } + for (const { http } of trees) http.end(); + await processor.shutdown(); + }); + test("concurrent shutdown callers wait for the same exporter barrier", async () => { let releaseShutdown!: () => void; const shutdownBarrier = new Promise((resolve) => { diff --git a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts index ec2108987..2b2fa09f2 100644 --- a/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts +++ b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts @@ -137,6 +137,24 @@ describe("MLflow UC setup", () => { ).toThrow(); }); + test("does not partially write guidance when tracing prerequisites fail", async () => { + const cwd = createProject(); + const packageDirectory = join(cwd, "node_modules", "@databricks", "appkit"); + mkdirSync(packageDirectory, { recursive: true }); + writeFileSync( + join(packageDirectory, "package.json"), + JSON.stringify({ name: "@databricks/appkit" }), + ); + vi.spyOn(process, "cwd").mockReturnValue(cwd); + vi.spyOn(console, "log").mockImplementation(() => undefined); + + await expect( + setupCommand.parseAsync(["node", "setup", "--write"]), + ).rejects.toThrow(/requires --mlflow-warehouse-id/i); + + expect(() => readFileSync(join(cwd, "CLAUDE.md"), "utf8")).toThrow(); + }); + test("reports an actionable error when uv is unavailable", async () => { const cwd = createProject(); diff --git a/packages/shared/src/cli/commands/setup.ts b/packages/shared/src/cli/commands/setup.ts index 54ab38871..9d8ab619a 100644 --- a/packages/shared/src/cli/commands/setup.ts +++ b/packages/shared/src/cli/commands/setup.ts @@ -503,30 +503,8 @@ async function runSetup(options: SetupCliOptions) { action = "Created"; } - if (shouldWrite) { - fs.writeFileSync(claudePath, finalContent); - console.log(`\nāœ“ ${action} CLAUDE.md`); - console.log(` Path: ${claudePath}`); - } else { - console.log("\nTo create/update CLAUDE.md, run:"); - console.log(" npx appkit setup --write\n"); - - if (existingContent) { - console.log( - `This will ${ - existingContent.includes(SECTION_START) - ? "update the existing" - : "add a new" - } AppKit section.\n`, - ); - } - - console.log("Preview of AppKit section:\n"); - console.log("ā”€".repeat(50)); - console.log(generateSection(installed)); - console.log("ā”€".repeat(50)); - } - + // Validate and provision tracing before writing any local guidance so a + // failed prerequisite cannot leave the project in a partially updated state. const cwd = process.cwd(); if (shouldWrite && projectRequiresMlflowUc(cwd, options.mlflowUc === true)) { const env = readEnvFile(path.join(cwd, ".env")); @@ -570,6 +548,30 @@ async function runSetup(options: SetupCliOptions) { { workspaceHost }, ); } + + if (shouldWrite) { + fs.writeFileSync(claudePath, finalContent); + console.log(`\nāœ“ ${action} CLAUDE.md`); + console.log(` Path: ${claudePath}`); + } else { + console.log("\nTo create/update CLAUDE.md, run:"); + console.log(" npx appkit setup --write\n"); + + if (existingContent) { + console.log( + `This will ${ + existingContent.includes(SECTION_START) + ? "update the existing" + : "add a new" + } AppKit section.\n`, + ); + } + + console.log("Preview of AppKit section:\n"); + console.log("ā”€".repeat(50)); + console.log(generateSection(installed)); + console.log("ā”€".repeat(50)); + } } export const setupCommand = new Command("setup")