diff --git a/apps/dev-playground/client/src/routes/agent.route.tsx b/apps/dev-playground/client/src/routes/agent.route.tsx index 932176665..a996196ae 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,23 @@ function AgentRoute() {
+ {mlflowTraceId && ( +
+ + {mlflowTraceId} + + {mlflowTraceUrl && ( + + 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..f35b8d324 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,24 @@ function SmartDashboardRoute() {
+ {mlflowTraceId && ( +
+ + {mlflowTraceId} + + {mlflowTraceUrl && ( + + Open trace in MLflow + + )} +
+ )} + {(error || dataError) && (
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 new file mode 100644 index 000000000..fad2a78e7 --- /dev/null +++ b/apps/dev-playground/tests/agent-tracing.spec.ts @@ -0,0 +1,68 @@ +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); +}); + +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 new file mode 100644 index 000000000..12947e409 --- /dev/null +++ b/apps/dev-playground/tests/smart-dashboard-agent-tracing.spec.ts @@ -0,0 +1,62 @@ +import { expect, test } from "@playwright/test"; +import { runSmartDashboardTracingFixture } from "../server/tests/smart-dashboard-agent-tracing.fixture"; +import { setupMockAPI } from "./utils/test-utils"; + +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 = 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("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); + 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/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.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/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/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 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/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/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 7bc8e25a1..c4da4039e 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": { @@ -64,6 +65,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", @@ -85,7 +87,6 @@ "get-port": "7.2.0", "js-yaml": "4.2.0", "magic-string": "0.30.21", - "mlflow-tracing": "0.1.3", "obug": "2.1.1", "pg": "8.18.0", "picocolors": "1.1.1", diff --git a/packages/appkit/scripts/provision-mlflow-uc.py b/packages/appkit/scripts/provision-mlflow-uc.py new file mode 100644 index 000000000..75fbaf23a --- /dev/null +++ b/packages/appkit/scripts/provision-mlflow-uc.py @@ -0,0 +1,318 @@ +#!/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 +import time +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 _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, + 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( + 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 + and str(row[0]).startswith(table_prefix) + ] + + +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 _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, + experiment_name: str, + catalog_name: str, + 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: + 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 _location_fields(existing_location) != _location_fields(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 ''}" + ) + + _grant_trace_access( + workspace, + warehouse_id, + runtime_principal, + catalog_name, + schema_name, + table_names, + ) + _verify_trace_access( + workspace, + warehouse_id, + runtime_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("--runtime-principal", 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, + runtime_principal=args.runtime_principal, + ) + 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..7c76b86a4 --- /dev/null +++ b/packages/appkit/scripts/test_provision_mlflow_uc.py @@ -0,0 +1,287 @@ +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 + + +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], + *, + 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 self.failed_statement and self.failed_statement in statement: + return statement_response( + "FAILED", + error_message="principal lacks MODIFY", + ) + if self.pending_then_succeeded: + 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): + 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], **statement_options): + self.statement_execution = FakeStatementExecution( + table_names, **statement_options + ) + 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, + 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"], + **(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=table_prefix, + warehouse_id="0123456789abcdef", + runtime_principal="runtime-app-sp", + 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: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() + + first, _, _ = provision(module) + second, _, _ = provision(module) + + 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() + + 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_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" + 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/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 476de0b85..2c187350a 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -1,15 +1,30 @@ +import { SpanStatusCode, trace } from "@opentelemetry/api"; import type { AgentAdapter, AgentEvent, AgentInput, + AgentRemoteTraceEvent, AgentRunContext, AgentToolDefinition, + AgentUsage, } from "shared"; import { + getResponseHeaders, + retainResponseHeaders, type StreamBody, stream as servingStream, } from "../connectors/serving/client"; import { APPKIT_USER_AGENT, getClientOptions } from "../context/client-options"; +import { + captureTraceValue, + injectActiveTraceContext, + normalizeFailureOutput, + verifiedAgentRemoteTrace, +} from "../telemetry/agent-tracing"; +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). */ @@ -27,10 +42,181 @@ 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; +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; } +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, + previous: AgentUsage, +): AgentUsage | undefined { + const raw = isRecord(parsed.usage) ? parsed.usage : 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; + 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 retainedCost = + providerCost ?? (previous.costAvailable ? previous.costUsd : undefined); + + return { + inputTokens, + outputTokens, + totalTokens, + ...(cacheReadInputTokens !== undefined ? { cacheReadInputTokens } : {}), + ...(cacheCreationInputTokens !== undefined + ? { cacheCreationInputTokens } + : {}), + ...(retainedCost !== undefined ? { costUsd: retainedCost } : {}), + costAvailable: retainedCost !== 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 verifiedAgentRemoteTrace(traceId, spanId, "model-serving"); +} + +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}`, + ) + .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(""); + 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 +266,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 +295,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 +317,8 @@ interface RawFetchAdapterOptions { */ interface StreamBodyAdapterOptions { streamBody: StreamBody; + /** Model/endpoint name recorded in lifecycle telemetry. */ + model?: string; maxSteps?: number; maxTokens?: number; generationParams?: GenerationParams; @@ -271,6 +457,7 @@ interface DeltaToolCall { */ export class DatabricksAdapter implements AgentAdapter { private streamBody: StreamBody; + private model: string; private maxSteps: number; private maxTokens: number; private generationParams: GenerationParams; @@ -291,30 +478,32 @@ 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); 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, }); 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 response.body; + return retainResponseHeaders(response.body, response.headers); }; } } @@ -351,6 +540,7 @@ export class DatabricksAdapter implements AgentAdapter { body, signal, ), + model: endpointName, maxSteps, maxTokens, generationParams, @@ -363,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 @@ -385,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.", ); } @@ -567,20 +761,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 +789,52 @@ export class DatabricksAdapter implements AgentAdapter { thoughtSignature?: string; } >(); + const emittedRemoteTraces = new Set(); + const snapshotToolCalls = ( + normalizeCompletedArguments = false, + ): OpenAIToolCall[] => + Array.from(toolCallAccumulator.values()).map((tc) => ({ + id: tc.id, + type: "function" as const, + function: { + name: tc.name, + arguments: + normalizeCompletedArguments && tc.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()) { + 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(); + const decoder = new TextDecoder(); + let buffer = ""; while (true) { if (context.signal?.aborted) break; @@ -632,9 +875,38 @@ export class DatabricksAdapter implements AgentAdapter { continue; } + if (!isRecord(parsed)) continue; + + const usage = normalizeUsage(parsed, finalUsage); + 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 +919,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 +955,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(true) }; } private async *executeToolCalls( @@ -828,15 +1123,98 @@ 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; +function looksLikeTextToolCall(text: string): boolean { + let arrayStartPending = false; + for (let index = 0; index < text.length; index += 1) { + const character = text[index]; - const pyResult = tryParsePythonStyleToolCalls(trimmed); - if (pyResult.length > 0) return pyResult; + 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); +} - return []; +function isWhitespace(character: string | undefined): boolean { + return character !== undefined && /\s/.test(character); +} + +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/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-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 c20f74fa5..02ad215f3 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -1,11 +1,59 @@ +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 { retainResponseHeaders } from "../../connectors/serving/client"; +import { consumeAdapterStream } from "../../core/agent/consume-adapter-stream"; import { DatabricksAdapter, type GenerationParams, 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" }); @@ -66,6 +114,73 @@ function createReadableStream(chunks: string[]): ReadableStream { }); } +function createTimedReadableStream( + reads: Array<{ + at: number; + chunk?: string; + error?: Error; + onRead?: () => void; + }>, +): 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); + next.onRead?.(); + 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; + totalCostUsd?: number; +}): 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 } : {}), + ...(options.totalCostUsd !== undefined + ? { total_cost_usd: options.totalCostUsd } + : {}), + }), + ); +} + function mockFetch(chunks: string[]): typeof globalThis.fetch { return vi.fn().mockResolvedValue({ ok: true, @@ -116,6 +231,776 @@ 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", + "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.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_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); + 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("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 () => { + 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/${remoteTrace}`, + 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/${remoteTrace}`, + 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("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/44444444444444444444444444444444", + 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/44444444444444444444444444444444", + 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); + 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 +1021,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 () => { @@ -155,7 +1051,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 () => { @@ -777,23 +1714,116 @@ 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-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\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, + }); + 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"); + 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); }); test("yields error status then throws when injected streamBody fails", async () => { @@ -813,11 +1843,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 () => { @@ -930,6 +1969,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({ @@ -996,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"; @@ -1120,6 +2194,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); }); }); 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/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/cache/index.ts b/packages/appkit/src/cache/index.ts index e2b79d04c..1aec4ddca 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,22 +120,59 @@ export class CacheManager { static async getInstance( userConfig?: Partial, ): Promise { + if (CacheManager.shutdownPromise) { + await CacheManager.shutdownPromise; + } if (CacheManager.instance) { return CacheManager.instance; } 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; } + /** + * 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 shutdown(): Promise { + if (CacheManager.shutdownPromise) return CacheManager.shutdownPromise; + + 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 = ownedShutdown; + return ownedShutdown; + } + /** * Create a new cache manager instance * 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/connectors/ai-search/client.ts b/packages/appkit/src/connectors/ai-search/client.ts index f4ed44b15..114292602 100644 --- a/packages/appkit/src/connectors/ai-search/client.ts +++ b/packages/appkit/src/connectors/ai-search/client.ts @@ -6,6 +6,10 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; +import { + captureTraceValue, + getActiveAgentTraceIdentity, +} from "../../telemetry/agent-tracing"; import type { WorkspaceClient } from "../../workspace-client"; import { contextFromAbortSignal } from "../context"; import type { @@ -19,6 +23,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; @@ -34,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, @@ -70,6 +85,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, @@ -78,11 +97,25 @@ export class AiSearchConnector { params.filters && Object.keys(params.filters).length > 0 ), "vs.has_reranker": !!params.reranker, + ...activeAgentIdentityAttributes(), }, }, 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 { + if (signal?.aborted) { + throw new Error("Query cancelled before execution"); + } const response = (await workspaceClient.apiClient.request( { method: "POST", @@ -96,12 +129,13 @@ export class AiSearchConnector { )) as VsRawResponse; const duration = Date.now() - startTime; + 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.setAttribute("vs.duration_ms", duration); span.setStatus({ code: SpanStatusCode.OK }); logger.event()?.setContext("ai-search", { @@ -114,12 +148,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 }, @@ -131,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, @@ -146,13 +177,28 @@ 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, + ...activeAgentIdentityAttributes(), }, }, async (span: Span) => { + const startTime = Date.now(); + setCapturedAttribute(span, "mlflow.spanInputs", { + endpointName: params.endpointName, + indexName: params.indexName, + pageToken: params.pageToken, + queryType: "next_page", + }); try { + if (signal?.aborted) { + throw new Error("Query cancelled before execution"); + } const response = (await workspaceClient.apiClient.request( { method: "POST", @@ -168,16 +214,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 +280,110 @@ 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): RetrieverOutputs { + const columnNames = response.manifest.columns.map((column) => column.name); + const documentIdIndex = columnNames.findIndex((column) => { + const normalized = column.toLowerCase(); + return DOCUMENT_ID_COLUMNS.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 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); + 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 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 instanceof Error ? error.name : "Error", + message, + }); + span.setStatus({ + code: SpanStatusCode.ERROR, + 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 new file mode 100644 index 000000000..20b25e504 --- /dev/null +++ b/packages/appkit/src/connectors/ai-search/tests/client.test.ts @@ -0,0 +1,510 @@ +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 { 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"; + +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; +} + +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 = { + 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: "vector backend exposed password [REDACTED]", + }); + expect(span.events).toEqual([ + expect.objectContaining({ + name: "exception", + attributes: expect.objectContaining({ + "exception.message": "vector backend exposed password [REDACTED]", + }), + }), + ]); + expect(JSON.parse(String(span.attributes["mlflow.spanOutputs"]))).toEqual({ + message: "vector backend exposed password [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"); + }); + + 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: "Query cancelled before execution", + }); + 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({ + message: "Query cancelled before execution", + }); + 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: "Query cancelled before execution", + }); + 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({ + message: "Query cancelled before execution", + }); + 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/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 de9d0465c..6bb2334e2 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"; @@ -10,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, @@ -18,6 +21,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 +38,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. @@ -79,26 +114,82 @@ 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: new Headers({ - "Content-Type": "application/json", - Accept: "text/event-stream", - }), + headers, payload: body, 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); +} + +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); } /** diff --git a/packages/appkit/src/connectors/serving/tests/client.test.ts b/packages/appkit/src/connectors/serving/tests/client.test.ts index 34bfd743a..d0fe00b11 100644 --- a/packages/appkit/src/connectors/serving/tests/client.test.ts +++ b/packages/appkit/src/connectors/serving/tests/client.test.ts @@ -1,6 +1,53 @@ -import { afterEach, describe, expect, test, vi } from "vitest"; -import { Context } from "../../../workspace-client"; -import { invoke, stream } from "../client"; +import http from "node:http"; +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, 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(); + context.setGlobalContextManager( + new AsyncLocalStorageContextManager().enable(), + ); + propagation.disable(); + propagation.setGlobalPropagator(W3C_PROPAGATOR); +}); + +afterAll(() => { + propagation.disable(); + context.disable(); +}); + +function withActiveTrace(operation: () => T, spanId = SPAN_ID): T { + const span = trace.wrapSpanContext({ + traceId: TRACE_ID, + spanId, + traceFlags: TraceFlags.SAMPLED, + traceState: createTraceState("vendor=value"), + }); + return context.with(trace.setSpan(context.active(), span), operation); +} function createMockClient(host = "https://test.databricks.com") { return { @@ -78,6 +125,135 @@ 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({ + 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({ @@ -95,6 +271,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 c4f3d07ed..9a4f0312a 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 type { AgentEvent, AgentRemoteTraceEvent } from "shared"; +import { + AgentUsageAccumulator, + type ConsumedAgentStream, +} from "../../telemetry/agent-tracing"; interface ConsumeAdapterStreamOptions { /** @@ -14,6 +18,101 @@ 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 { + 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 + ); } /** @@ -37,16 +136,161 @@ interface ConsumeAdapterStreamOptions { export async function consumeAdapterStream( stream: AsyncIterable, opts: ConsumeAdapterStreamOptions = {}, -): Promise { +): Promise { let text = ""; - for await (const event of stream) { - if (opts.signal?.aborted) break; + const usage = new AgentUsageAccumulator(); + const activeModelSteps = new Map(); + const consumedModelSteps = new Set(); + let remoteTrace: ConsumedAgentStream["remoteTrace"]; + 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)) { + consumedModelSteps.add(event.stepId); + usage.add(event.usage); + } + } else if (isValidRemoteTrace(event)) { + remoteTrace = event; } opts.onEvent?.(event); + }; + + 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; + } } - return text; + if (stopEarly) requestIteratorCleanup(iterator); + 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..3ddeec2e9 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -3,6 +3,7 @@ import type { AgentAdapter, AgentEvent, AgentToolDefinition, + AgentUsage, Message, PluginConstructor, PluginData, @@ -14,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"; @@ -23,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, @@ -47,6 +54,11 @@ export interface RunAgentInput { * there is no HTTP request in standalone mode). */ plugins?: PluginData[]; + sessionId?: string; + userId?: string; + requestId?: string; + threadId?: string; + appName?: string; } export interface RunAgentResult { @@ -54,8 +66,18 @@ export interface RunAgentResult { text: string; /** Every event the adapter yielded, in order. Useful for inspection/tests. */ events: AgentEvent[]; + traceId: string; + 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. @@ -94,15 +116,52 @@ 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 appName = resolveAgentTraceAppName(input.appName); + const sessionId = input.sessionId ?? threadId; + const userId = input.userId ?? "service-principal"; + const traced = await runWithAgentTrace( + { + appName, + agentName: def.name ?? "agent", + route: "runAgent", + sessionId, + userId, + requestId, + threadId, + }, + { messages: input.messages }, + async (observer) => { + await initStandalonePlugins(input.plugins ?? [], providerCache); + return runAgentInternal( + def, + { + ...input, + appName, + requestId, + sessionId, + threadId, + userId, + }, + providerCache, + observer, + ); + }, + ); + return { + ...traced.value, + traceId: traced.traceId, + usage: traced.usage, + }; } async function runAgentInternal( def: AgentDefinition, - input: RunAgentInput, + input: ResolvedRunAgentInput, providerCache: Map, -): Promise { + observer: AgentTraceObserver, +): Promise> { const adapter = await resolveAdapter(def); const messages = normalizeMessages(input.messages, def.instructions); const toolIndex = buildStandaloneToolIndex( @@ -123,48 +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, - }; - // 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, - ); - 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, + ), + (childUsage) => observer.addChildUsage(childUsage), + ); + const childResult = { + text: childTrace.value.text, + usage: childTrace.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(...)] }).", + ); + }, ); }; @@ -174,7 +266,7 @@ async function runAgentInternal( { messages, tools, - threadId: randomUUID(), + threadId: input.threadId ?? randomUUID(), signal, extensions: buildStandaloneExtensions(toolIndex), }, @@ -184,14 +276,28 @@ 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); + observer.onEvent(event); }, }); - return { text, events }; + if (signal?.aborted) { + throw agentAbortError(); + } + + return { + text: consumed.text, + events, + }; +} + +function agentAbortError(): Error { + const error = new Error("Agent run aborted"); + error.name = "AbortError"; + return error; } /** 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..bbbf41352 --- /dev/null +++ b/packages/appkit/src/core/agent/tests/agent-tracing-public-api.test.ts @@ -0,0 +1,75 @@ +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("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, + 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..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,24 +11,28 @@ async function* streamOf( } describe("consumeAdapterStream", () => { + afterEach(() => { + vi.useRealTimers(); + }); + 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 () => { @@ -67,20 +71,435 @@ 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: "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 { + 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("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("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 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", + }, + }); + }); + + 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", + }, + }); }); }); 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..33f8ba413 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,11 @@ +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, AgentEvent, @@ -8,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"; @@ -25,6 +33,69 @@ 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 }; +} + +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[] = [ @@ -42,6 +113,96 @@ describe("runAgent", () => { expect(result.events).toHaveLength(3); }); + 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", + 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({ + name: "planner", + instructions: "x", + model: scriptedAdapter(events), + }); + + 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(); + } + + 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, + 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" }, @@ -257,6 +418,250 @@ 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("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/core/agent/trace-tool-call.ts b/packages/appkit/src/core/agent/trace-tool-call.ts new file mode 100644 index 000000000..913f58c0d --- /dev/null +++ b/packages/appkit/src/core/agent/trace-tool-call.ts @@ -0,0 +1,83 @@ +import { type Span, SpanStatusCode, trace } from "@opentelemetry/api"; +import type { ToolEffect } from "shared"; +import { + captureTraceValue, + normalizeFailureOutput, +} 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 safeError = + error instanceof Error ? error.message : String(error ?? "Unknown error"); + const errorAttribute = captureTraceValue( + { error: safeError }, + { redactKeys: ["error"] }, + ); + 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); + 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/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 201acf190..3e65cfeaf 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,13 +197,27 @@ export class AppKit { onPluginsReady?: (appkit: PluginMap) => void | Promise; disableInternalTelemetry?: boolean; } = {}, - ): Promise> { - // Initialize core services - TelemetryManager.initialize(config?.telemetry); - await CacheManager.getInstance(config?.cache); - + ): Promise> { 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(); @@ -224,7 +245,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"); @@ -245,11 +266,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({ @@ -384,6 +411,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..100fb7186 100644 --- a/packages/appkit/src/core/lifecycle-manager.ts +++ b/packages/appkit/src/core/lifecycle-manager.ts @@ -44,18 +44,18 @@ 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. */ private shutdownPhase = "not started"; + private signalHandlers?: { + SIGTERM: () => void; + SIGINT: () => void; + }; constructor(private readonly context: PluginContext) {} @@ -67,8 +67,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({ exitProcess: true }), + SIGINT: () => void this.shutdown({ exitProcess: true }), + }; + 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,12 +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(): 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; + shutdown({ + exitProcess = true, + }: { + exitProcess?: boolean; + } = {}): Promise { + 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; @@ -101,21 +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 = 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); - // 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()); @@ -183,22 +200,25 @@ export class LifecycleManager { exitCode = 1; } - clearTimeout(forceExitTimer); - process.exit(exitCode); + if (this.forceExitTimer) { + clearTimeout(this.forceExitTimer); + this.forceExitTimer = undefined; + } + this.removeSignalHandlers(); + if (this.exitRequested) process.exit(exitCode); } /** 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/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/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/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/core/tests/lifecycle-manager.test.ts b/packages/appkit/src/core/tests/lifecycle-manager.test.ts index 121e7eb5c..9b6aad8ab 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"); @@ -380,5 +386,53 @@ 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(); + }); + + 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/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 1d9162a96..35a172cee 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -6,6 +6,7 @@ import type { AgentAdapter, AgentRunContext, AgentToolDefinition, + AgentUsage, IAppRouter, Message, PluginPhase, @@ -36,6 +37,7 @@ import { isHostedTool, resolveHostedTools, } from "../../core/agent/tools"; +import { traceToolCall } from "../../core/agent/trace-tool-call"; import type { AgentDefinition, AgentsPluginConfig, @@ -51,30 +53,30 @@ 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"; import manifest from "./manifest.json"; -import { - currentTraceId, - initAgentTracing, - linkTraceToRun, - traceAgent, - traceTool, -} from "./mlflow"; import { approvalRequestSchema, cancelRequestSchema, 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"; +export { traceToolCall } from "../../core/agent/trace-tool-call"; + /** * Context flag recorded on the in-memory AgentDefinition to indicate whether * it came from markdown (file) or from user code. Drives the asymmetric @@ -135,16 +137,20 @@ 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 }; } export class AgentsPlugin extends Plugin implements ToolProvider { - // Routed through `unknown`: the optional resources have differing `fields` - // keys (serving `name`, experiment `experimentId`), which TS widens to an - // incompatible union on the JSON import. The shape is validated at runtime - // against the plugin-manifest schema. - static manifest = manifest as unknown as PluginManifest; + static manifest = manifest as PluginManifest; static phase: PluginPhase = "deferred"; protected declare config: AgentsPluginConfig; @@ -170,9 +176,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). " + @@ -261,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 }; + } + } } /** @@ -287,7 +315,6 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } async setup() { - await initAgentTracing(); const { agents, defaultAgentName } = await this.buildAgentRegistry(); this.agents = agents; this.defaultAgentName = defaultAgentName; @@ -495,7 +522,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 }, ); } @@ -877,69 +904,114 @@ 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, + traceRequestId: 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, }); return; } const { message, threadId, agent: agentName, mlflowRunId } = parsed.data; + if (mlflowRunId) observer.linkToRun(mlflowRunId); 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 // 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"); - 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; } + 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) { - res.status(404).json({ 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)); - - 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); - res.status(500).json({ 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, mlflowRunId); } /** @@ -957,13 +1029,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]; } /** @@ -977,20 +1059,42 @@ 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, + traceRequestId: 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, }); return; } const { input, mlflowRunId } = parsed.data; + if (mlflowRunId) observer.linkToRun(mlflowRunId); 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 @@ -999,7 +1103,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 ` + @@ -1010,61 +1114,79 @@ 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) { + const reservation = this.reserveStream( + userId, + limits.maxConcurrentStreamsPerUser, + ); + if (!reservation) { 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; } + 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); - - 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); - res.status(500).json({ 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, - mlflowRunId, - ); } private async _streamAgent( @@ -1073,12 +1195,12 @@ export class AgentsPlugin extends Plugin implements ToolProvider { registered: RegisteredAgent, thread: Thread, userId: string, - mlflowRunId?: string, + observer: AgentTraceObserver, + requestId: string, + abortController: AbortController, + traceRequestId: string, ): Promise { - const abortController = new AbortController(); const signal = abortController.signal; - const requestId = randomUUID(); - 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 @@ -1108,119 +1230,105 @@ export class AgentsPlugin extends Plugin implements ToolProvider { translator, outboundEvents, toolCallsUsed: { count: 0 }, + traceObserver: observer, + traceIdentity: { + appName: resolveAgentTraceAppName(), + route: "chat", + sessionId: requestSessionId(req) ?? thread.id, + userId, + requestId: traceRequestId, + }, }; 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, + ]; + + // 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, + mlflowTraceId: observer.traceId, + ...(traceUrl ? { traceUrl } : {}), + }, + })) { + 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 { - for (const evt of translator.translate({ - type: "metadata", - data: { threadId: thread.id }, - })) { - outboundEvents.push(evt); - } - - // Root MLflow span for the turn; tool-call spans nest under it. - await traceAgent( - registered.name ?? "agent", + const stream = registered.adapter.run( { - messages: thread.messages.map((m) => ({ - role: m.role, - content: m.content, - })), + messages: messagesWithSystem, + tools, + threadId: thread.id, + signal, + extensions: buildAdapterExtensions(registered.toolIndex), }, - async (span) => { - // Link this turn's trace to an eval run when the eval runner - // supplied one, so the trace shows under the MLflow evaluation run. - if (mlflowRunId) linkTraceToRun(mlflowRunId); - - 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, - ]; - - const stream = registered.adapter.run( - { - messages: messagesWithSystem, - tools, - threadId: thread.id, - signal, - extensions: buildAdapterExtensions(registered.toolIndex), - }, - { executeTool, signal }, - ); - - // The accumulation rule (deltas append, `message` replaces) is - // shared with `runAgent` and `runSubAgent`; see - // `consumeAdapterStream` for the rationale. - const fullContent = await consumeAdapterStream(stream, { - signal, - onEvent: (event) => { - for (const translated of translator.translate(event)) { - outboundEvents.push(translated); - } - }, - }); - - if (fullContent) { - span.setOutputs({ role: "assistant", content: fullContent }); - await this.threadStore.addMessage(thread.id, userId, { - id: randomUUID(), - role: "assistant", - content: fullContent, - createdAt: new Date(), - }); - } + { executeTool, signal }, + ); - // Surface the MLflow trace id so eval runs can attach assessments - // to this turn's trace. No-op when tracing is disabled. - const mlflowTraceId = currentTraceId(); - if (mlflowTraceId) { - for (const evt of translator.translate({ - type: "metadata", - data: { mlflowTraceId }, - })) { - outboundEvents.push(evt); - } + 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 (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); + 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. @@ -1241,18 +1349,26 @@ export class AgentsPlugin extends Plugin implements ToolProvider { ); } } + outboundEvents.close(); } - outboundEvents.close(); })(); 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); } }, @@ -1261,6 +1377,8 @@ export class AgentsPlugin extends Plugin implements ToolProvider { stream: { ...agentStreamDefaults.stream, streamId: requestId }, }, ); + await driver; + if (driverFailure) throw driverFailure.error; } /** @@ -1292,20 +1410,18 @@ export class AgentsPlugin extends Plugin implements ToolProvider { registered: RegisteredAgent, thread: Thread, userId: string, - mlflowRunId?: string, + observer: AgentTraceObserver, + requestId: string, + abortController: AbortController, + traceRequestId: string, ): Promise { - const abortController = new AbortController(); const signal = abortController.signal; - const requestId = randomUUID(); - 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; - // Assigned inside the span below (the only place the active trace id - // resolves), read into the response envelope after. - let mlflowTraceId: string | undefined; - const runState: RunState = { req, userId, @@ -1320,82 +1436,77 @@ export class AgentsPlugin extends Plugin implements ToolProvider { translator: new AgentEventTranslator(), outboundEvents: new EventChannel(), toolCallsUsed: { count: 0 }, + traceObserver: observer, + traceIdentity: { + appName: resolveAgentTraceAppName(), + route: invokeTraceRoute(req), + sessionId: requestSessionId(req) ?? thread.id, + userId, + requestId: traceRequestId, + }, }; 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 = ""; try { - // Root MLflow span for the turn; tool-call spans nest under it. Mirrors - // the streaming path so the invoke surface produces the same trace shape - // instead of orphan root TOOL spans. - await traceAgent( - registered.name ?? "agent", + const stream = registered.adapter.run( { - messages: thread.messages.map((m) => ({ - role: m.role, - content: m.content, - })), - }, - async (span) => { - // Link this turn's trace to an eval run when the eval runner - // supplied one, so the trace shows under the MLflow evaluation run. - if (mlflowRunId) linkTraceToRun(mlflowRunId); - - 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, - ]; - - const stream = registered.adapter.run( - { - messages: messagesWithSystem, - tools, - threadId: thread.id, - signal, - }, - { executeTool, signal }, - ); - - fullContent = await consumeAdapterStream(stream, { signal }); - - if (fullContent) { - span.setOutputs({ role: "assistant", content: fullContent }); - await this.threadStore.addMessage(thread.id, userId, { - id: randomUUID(), - role: "assistant", - content: fullContent, - createdAt: new Date(), - }); - } - - mlflowTraceId = currentTraceId(); + messages: messagesWithSystem, + tools, + threadId: thread.id, + signal, + extensions: buildAdapterExtensions(registered.toolIndex), }, + { executeTool, signal }, ); + 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" }); + 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); @@ -1405,7 +1516,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { : error instanceof Error ? error.message : String(error); - res.status(500).json({ error: message }); + const body = { error: message, trace_id: observer.traceId }; + observer.recordError(error, body); + res.status(500).json(body); return; } finally { this.approvalGate.abortStream(requestId); @@ -1432,17 +1545,18 @@ 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, - // Lets an eval runner attach assessments to this turn's trace; absent - // when tracing is disabled. - ...(mlflowTraceId ? { mlflow_trace_id: mlflowTraceId } : {}), + trace_id: observer.traceId, + mlflow_trace_id: observer.traceId, output: [message], - }); + }; + observer.setOutput(payload); + res.json(payload); } /** @@ -1463,120 +1577,138 @@ 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}).`; - } - } - - // Traced from here so the span covers execution only, not the approval - // wait above (which is human latency). - const toolResult = await traceTool(name, args, async () => { - let result: unknown; - if (entry.source === "toolkit") { - if (!this.context) { + }, + async () => { + if (runState.toolCallsUsed.count >= runState.limits.maxToolCalls) { + runState.abortController.abort( + new Error( + `Tool-call budget exhausted (limit ${runState.limits.maxToolCalls}).`, + ), + ); throw new Error( - "Plugin tool execution requires PluginContext; this should never happen through createApp", + `Tool-call budget exhausted (limit ${runState.limits.maxToolCalls}). Raise agents({ limits: { maxToolCalls } }) or review the agent's tool-selection logic.`, ); } - 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)) { + 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, + { 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}`); + const childResult = await this.runSubAgent( + runState, + childAgent, + args, + depth + 1, + ); + 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 + // 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( - `Function tool '${name}' received non-object arguments (got ${args === null ? "null" : Array.isArray(args) ? "array" : typeof args}); expected a JSON object.`, + `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.", ); } - 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 result; - }); - return normalizeToolResult(toolResult); + return normalizeToolResult(result); + }, + ); } /** - * 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 @@ -1592,7 +1724,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}). ` + @@ -1613,14 +1745,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() @@ -1651,37 +1775,71 @@ export class AgentsPlugin extends Plugin implements ToolProvider { }, ]; - return 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) => { - 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; + }, + (childUsage) => runState.traceObserver?.addChildUsage(childUsage), ); + return { text: traced.value, usage: traced.usage }; } private async _handleCancel(req: express.Request, res: express.Response) { @@ -1839,6 +1997,80 @@ 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 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 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, + 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() || + 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/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/manifest.json b/packages/appkit/src/plugins/agents/manifest.json index 4d6f52485..fc07cf07a 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", @@ -15,23 +46,10 @@ "permission": "CAN_QUERY", "fields": { "name": { - "env": "DATABRICKS_SERVING_ENDPOINT_NAME", + "env": "DATABRICKS_AGENT_SERVING_ENDPOINT_NAME", "description": "Default LLM serving endpoint name" } } - }, - { - "type": "experiment", - "alias": "MLflow experiment for agent traces", - "resourceKey": "agents-mlflow-experiment", - "description": "When bound, agent turns and tool calls are traced to this MLflow experiment via OpenTelemetry. Tracing is a no-op when unset.", - "permission": "CAN_EDIT", - "fields": { - "experimentId": { - "env": "MLFLOW_EXPERIMENT_ID", - "description": "MLflow experiment id traces are logged to" - } - } } ] } diff --git a/packages/appkit/src/plugins/agents/mlflow.ts b/packages/appkit/src/plugins/agents/mlflow.ts deleted file mode 100644 index cd48b8190..000000000 --- a/packages/appkit/src/plugins/agents/mlflow.ts +++ /dev/null @@ -1,157 +0,0 @@ -import { createLogger } from "../../logging/logger"; - -const logger = createLogger("agents"); - -type MlflowModule = typeof import("mlflow-tracing"); - -let mlflow: MlflowModule | undefined; -let enabled = false; -let initStarted = false; - -/** The bound MLflow experiment id, from the optional `experiment` resource. */ -function experimentId(): string | undefined { - const id = process.env.MLFLOW_EXPERIMENT_ID?.trim(); - return id || undefined; -} - -/** - * Databricks host with a scheme. The mlflow-tracing SDK uses `DATABRICKS_HOST` - * verbatim to build request URLs and doesn't add `https://`, so a bare host - * (`workspace.cloud.databricks.com`) makes `new URL()` throw. Pass an explicit - * normalized host when the env var is set; when it isn't (profile-based auth), - * return undefined and let the SDK read the host from `~/.databrickscfg`. - */ -function normalizedDatabricksHost(): string | undefined { - const raw = process.env.DATABRICKS_HOST?.trim(); - if (!raw) return undefined; - return /^https?:\/\//i.test(raw) ? raw : `https://${raw}`; -} - -/** - * Initialize MLflow agent tracing once, when an experiment is bound — i.e. the - * agents plugin's optional `experiment` resource is set (`MLFLOW_EXPERIMENT_ID`). - * - * Auth is resolved by the `mlflow-tracing` SDK from the app's own Databricks - * credentials — `DATABRICKS_HOST`/`DATABRICKS_TOKEN` or a `~/.databrickscfg` - * profile (`MLFLOW_TRACKING_URI=databricks://profile`) — so no tokens or OTLP - * headers are wired by hand. A failure (missing creds, bad experiment) logs and - * leaves tracing disabled rather than breaking the agent. - * - * Safe to call repeatedly; only the first call does work. - */ -export async function initAgentTracing(): Promise { - if (initStarted) return; - initStarted = true; - - const id = experimentId(); - if (!id) return; - - try { - mlflow = await import("mlflow-tracing"); - const host = normalizedDatabricksHost(); - mlflow.init({ - trackingUri: process.env.MLFLOW_TRACKING_URI?.trim() || "databricks", - experimentId: id, - ...(host ? { host } : {}), - }); - enabled = true; - logger.info("MLflow agent tracing enabled (experiment %s)", id); - } catch (err) { - logger.warn("MLflow agent tracing disabled: %O", err); - } -} - -/** - * Records a span's outputs. Callers get one from `traceAgent`/`traceTool`; - * it's a no-op when tracing is disabled, so call sites never branch on it. - */ -export interface SpanRecorder { - setOutputs(outputs: unknown): void; -} - -const noopRecorder: SpanRecorder = { setOutputs() {} }; - -/** - * Run `fn` inside an MLflow span of `spanType` when tracing is enabled, - * otherwise just run it (zero overhead). Spans auto-nest via the SDK's active - * context, so a TOOL span opened inside an AGENT span's callback becomes its - * child. The callback's resolved value is recorded as the span's outputs unless - * it called `setOutputs` first; return `undefined` (or set outputs explicitly) - * when the return value isn't the output you want traced. - */ -async function trace( - spanType: "AGENT" | "TOOL", - name: string, - inputs: unknown, - fn: (span: SpanRecorder) => Promise, -): Promise { - if (!enabled || !mlflow) return fn(noopRecorder); - const type = - spanType === "AGENT" ? mlflow.SpanType.AGENT : mlflow.SpanType.TOOL; - return await mlflow.withSpan( - async (span) => { - if (inputs !== undefined) span.setInputs(inputs); - let outputsSet = false; - const result = await fn({ - setOutputs(outputs) { - outputsSet = true; - span.setOutputs(outputs); - }, - }); - if (!outputsSet && result !== undefined) span.setOutputs(result); - return result; - }, - { name, spanType: type }, - ); -} - -/** Trace a turn's root AGENT span. See {@link trace}. */ -export function traceAgent( - name: string, - inputs: unknown, - fn: (span: SpanRecorder) => Promise, -): Promise { - return trace("AGENT", name, inputs, fn); -} - -/** Trace a TOOL span, nested under the active AGENT span. See {@link trace}. */ -export function traceTool( - name: string, - inputs: unknown, - fn: (span: SpanRecorder) => Promise, -): Promise { - return trace("TOOL", name, inputs, fn); -} - -/** - * The MLflow trace id for the active turn, when tracing is enabled. Must be - * read inside an agent span so eval runs can correlate the turn to its trace - * and attach assessments. Returns undefined when tracing is off. - * - * Reads the context-active span rather than `getLastActiveTraceId()`: the - * latter is only populated when a root span *ends* (on export), so mid-turn it - * returns the previous turn's id — or, under concurrent turns, another turn's. - */ -export function currentTraceId(): string | undefined { - if (!enabled || !mlflow) return undefined; - try { - return mlflow.getCurrentActiveSpan()?.traceId; - } catch { - return undefined; - } -} - -/** - * Link the active turn's trace to an MLflow run by id, via the `mlflow.sourceRun` - * trace metadata. Used by eval runs so each case's trace shows under the run. - * Must be called while a trace is active (inside an agent span). No-op when - * tracing is disabled. - */ -export function linkTraceToRun(runId: string): void { - if (!enabled || !mlflow) return; - try { - mlflow.updateCurrentTrace({ metadata: { "mlflow.sourceRun": runId } }); - } catch (err) { - logger.warn("Failed to link trace to run %s: %O", runId, err); - } -} 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..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 @@ -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: {}, @@ -65,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 }; @@ -99,6 +172,450 @@ 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({ + 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" }); + + 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["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"], + 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(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), + ); + 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.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 @@ -418,4 +935,314 @@ 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"); + }); + + 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/plugins/agents/tests/dos-limits.test.ts b/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts index 2a4e4a221..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, @@ -338,6 +413,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 +470,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/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/plugins/agents/tests/mlflow.test.ts b/packages/appkit/src/plugins/agents/tests/mlflow.test.ts deleted file mode 100644 index a43c29cec..000000000 --- a/packages/appkit/src/plugins/agents/tests/mlflow.test.ts +++ /dev/null @@ -1,90 +0,0 @@ -import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; - -/** - * The tracing module keeps module-level singleton state (`enabled`, - * `initStarted`) and lazily `import()`s `mlflow-tracing`. Each test resets the - * module registry and re-mocks the SDK so init runs fresh. - */ - -function stubSdk(overrides: Record = {}) { - const setInputs = vi.fn(); - const setOutputs = vi.fn(); - const span = { setInputs, setOutputs }; - const sdk = { - init: vi.fn(), - SpanType: { AGENT: "AGENT", TOOL: "TOOL" }, - withSpan: vi.fn(async (fn: (s: unknown) => unknown) => fn(span)), - getCurrentActiveSpan: vi.fn(() => ({ traceId: "tr-active" })), - // Must NOT be consulted by currentTraceId — it only reflects the last - // root span that *ended*, i.e. a previous/other turn. - getLastActiveTraceId: vi.fn(() => "tr-STALE"), - updateCurrentTrace: vi.fn(), - ...overrides, - }; - vi.doMock("mlflow-tracing", () => sdk); - return { sdk, span, setInputs, setOutputs }; -} - -describe("agent tracing (mlflow)", () => { - beforeEach(() => { - vi.resetModules(); - delete process.env.MLFLOW_EXPERIMENT_ID; - }); - - afterEach(() => { - vi.doUnmock("mlflow-tracing"); - delete process.env.MLFLOW_EXPERIMENT_ID; - }); - - test("disabled (no experiment bound): still runs fn and returns its value", async () => { - const mod = await import("../mlflow"); - await mod.initAgentTracing(); - - const fn = vi.fn(async () => "result"); - await expect(mod.traceTool("t", { a: 1 }, fn)).resolves.toBe("result"); - expect(fn).toHaveBeenCalledOnce(); - expect(mod.currentTraceId()).toBeUndefined(); - }); - - test("currentTraceId reads the context-active span, not getLastActiveTraceId", async () => { - process.env.MLFLOW_EXPERIMENT_ID = "exp-123"; - const { sdk } = stubSdk(); - const mod = await import("../mlflow"); - await mod.initAgentTracing(); - - let seen: string | undefined; - await mod.traceAgent("agent", { messages: [] }, async () => { - seen = mod.currentTraceId(); - }); - - expect(seen).toBe("tr-active"); - expect(sdk.getCurrentActiveSpan).toHaveBeenCalled(); - expect(sdk.getLastActiveTraceId).not.toHaveBeenCalled(); - }); - - test("auto-captures the callback's return value as span outputs", async () => { - process.env.MLFLOW_EXPERIMENT_ID = "exp-123"; - const { setOutputs } = stubSdk(); - const mod = await import("../mlflow"); - await mod.initAgentTracing(); - - await mod.traceTool("t", { a: 1 }, async () => ({ ok: true })); - expect(setOutputs).toHaveBeenCalledExactlyOnceWith({ ok: true }); - }); - - test("explicit setOutputs wins over auto-capture (no double-set)", async () => { - process.env.MLFLOW_EXPERIMENT_ID = "exp-123"; - const { setOutputs } = stubSdk(); - const mod = await import("../mlflow"); - await mod.initAgentTracing(); - - await mod.traceAgent("agent", { messages: [] }, async (span) => { - span.setOutputs({ role: "assistant", content: "hi" }); - return "ignored-return"; - }); - expect(setOutputs).toHaveBeenCalledExactlyOnceWith({ - role: "assistant", - content: "hi", - }); - }); -}); 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 0e31af4bf..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 @@ -1,28 +1,15 @@ +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 { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; import { AgentsPlugin } from "../agents"; -// Partial-mock the tracing module: traceAgent/traceTool still run their -// callbacks, but the trace id is deterministic and run-linking is a spy. -const linkTraceToRun = vi.hoisted(() => vi.fn()); -let mockTraceId: string | undefined; -vi.mock("../mlflow", () => ({ - initAgentTracing: vi.fn(async () => {}), - traceAgent: ( - _name: string, - _inputs: unknown, - fn: (span: { setOutputs: () => void }) => Promise, - ) => fn({ setOutputs: () => {} }), - traceTool: ( - _name: string, - _inputs: unknown, - fn: (span: { setOutputs: () => void }) => Promise, - ) => fn({ setOutputs: () => {} }), - currentTraceId: () => mockTraceId, - linkTraceToRun, -})); - /** * Surface-level guarantees on the agents plugin's HTTP route handlers when * downstream dependencies fail. Prior to PR #305 review finding #1+#2, @@ -39,8 +26,6 @@ vi.mock("../mlflow", () => ({ */ beforeEach(() => { - linkTraceToRun.mockClear(); - mockTraceId = undefined; // biome-ignore lint/suspicious/noExplicitAny: test seam, mirrors other suites (CacheManager as any).instance = { get: vi.fn(), @@ -53,10 +38,15 @@ beforeEach(() => { }; }); +afterEach(() => { + vi.unstubAllEnvs(); +}); + 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, @@ -65,9 +55,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; @@ -79,16 +73,89 @@ function mockRes() { return statusCode; }, json, + setHeader, }; } -function seedPlugin(adapter: unknown = { async *run() {} }): AgentsPlugin { +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 (plugin as any).agents.set("default", { name: "default", instructions: "hi", - adapter, + adapter: { async *run() {} }, toolIndex: new Map(), }); // biome-ignore lint/suspicious/noExplicitAny: seed private state @@ -96,6 +163,208 @@ function seedPlugin(adapter: unknown = { async *run() {} }): 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(); @@ -300,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" }, @@ -355,6 +690,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 @@ -378,15 +779,17 @@ describe("POST /invocations & /responses — successful invoke", () => { delete: vi.fn(), }; - const { res, json } = mockRes(); - await ( - plugin as unknown as { - _handleInvoke: ( - r: express.Request, - w: express.Response, - ) => Promise; - } - )._handleInvoke(mockReq({ input: "hi" }), res); + const { res, json, setHeader } = mockRes(); + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleInvoke: ( + r: express.Request, + w: express.Response, + ) => Promise; + } + )._handleInvoke(mockReq({ input: "hi", mlflowRunId: "run-99" }), res), + ); expect(res.status).not.toHaveBeenCalledWith(500); expect(json).toHaveBeenCalledTimes(1); @@ -400,6 +803,8 @@ describe("POST /invocations & /responses — successful invoke", () => { role: string; content: Array<{ type: string; text: string }>; }>; + trace_id: string; + mlflow_trace_id: string; }; expect(payload.object).toBe("response"); expect(payload.status).toBe("completed"); @@ -412,28 +817,42 @@ describe("POST /invocations & /responses — successful invoke", () => { type: "output_text", text: "hello world", }); + expect(payload.trace_id).toMatch(/^[0-9a-f]{32}$/); + expect(payload.mlflow_trace_id).toBe(payload.trace_id); + const root = observed.spans.find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(root?.attributes["mlflow.sourceRun"]).toBe("run-99"); + expect(setHeader).toHaveBeenCalledWith( + "X-MLflow-Trace-Id", + payload.trace_id, + ); }); - function seedEchoPlugin(): AgentsPlugin { - const plugin = seedPlugin({ - async *run() { - yield { type: "message_delta", content: "ok" }; + 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(), }; - return plugin; - } - async function invoke( - plugin: AgentsPlugin, - body: unknown, - ): Promise> { - const { res, json } = mockRes(); + const { res, json, setHeader } = mockRes(); await ( plugin as unknown as { _handleInvoke: ( @@ -441,42 +860,142 @@ describe("POST /invocations & /responses — successful invoke", () => { w: express.Response, ) => Promise; } - )._handleInvoke(mockReq(body), res); - return json.mock.calls[0]?.[0] as Record; - } - - test("links the trace to the run and echoes mlflow_trace_id when tracing is on", async () => { - mockTraceId = "tr-abc123"; - const plugin = seedEchoPlugin(); + )._handleInvoke(mockReq({ input: "hi" }), res); - const payload = await invoke(plugin, { - input: "hi", - mlflowRunId: "run-99", + 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}$/), }); - - expect(linkTraceToRun).toHaveBeenCalledWith("run-99"); - expect(payload.mlflow_trace_id).toBe("tr-abc123"); - }); - - test("omits mlflow_trace_id and does not link when tracing is off", async () => { - mockTraceId = undefined; // currentTraceId() no-ops when disabled - const plugin = seedEchoPlugin(); - - const payload = await invoke(plugin, { input: "hi" }); - - expect(linkTraceToRun).not.toHaveBeenCalled(); - expect(payload).not.toHaveProperty("mlflow_trace_id"); }); +}); - test("does not link when no run id is supplied even if tracing is on", async () => { - mockTraceId = "tr-standalone"; - const plugin = seedEchoPlugin(); +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> = []; + // 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)}`); + } + }; + 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(), + }; + // 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", + mlflowRunId: "run-chat-99", + }); + const { res, setHeader } = mockRes(); + setHeader.mockImplementation((name, value) => { + if (name === "X-MLflow-Trace-Id") order.push(`header:${String(value)}`); + }); - const payload = await invoke(plugin, { input: "hi" }); + const observed = await captureRouteSpans(() => + ( + plugin as unknown as { + _handleChat: ( + request: express.Request, + response: express.Response, + ) => Promise; + } + )._handleChat(req, res), + ); - expect(linkTraceToRun).not.toHaveBeenCalled(); - // Trace still exists and its id is surfaced — just not linked to a run. - expect(payload.mlflow_trace_id).toBe("tr-standalone"); + const metadata = streamed[0] as { + type?: string; + data?: { + traceId?: string; + mlflowTraceId?: 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), + mlflowTraceId: expect.any(String), + traceUrl: expect.stringMatching( + /^https:\/\/example\.cloud\.databricks\.com\/ml\/experiments\/123456789\/traces\?selectedTraceId=/, + ), + }, + }); + expect(setHeader).toHaveBeenCalledWith( + "X-MLflow-Trace-Id", + metadata.data?.traceId, + ); + expect(metadata.data?.mlflowTraceId).toBe(metadata.data?.traceId); + const root = observed.spans.find( + (span) => span.attributes["mlflow.spanType"] === "AGENT", + ); + expect(root?.attributes["mlflow.sourceRun"]).toBe("run-chat-99"); + expect(JSON.stringify(streamed)).not.toContain("model_start"); + expect(JSON.stringify(streamed)).not.toContain("model_end"); }); }); 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..b64b2c499 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,175 @@ 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"], + span.attributes["appkit.memory.store"], + span.attributes["appkit.memory.key"], + ]), + ).toEqual([ + ["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( + (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..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 @@ -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.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/tests/trace-conformance.integration.test.ts b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts new file mode 100644 index 000000000..992580c01 --- /dev/null +++ b/packages/appkit/src/plugins/agents/tests/trace-conformance.integration.test.ts @@ -0,0 +1,2546 @@ +import { execFileSync } from "node:child_process"; +import { + existsSync, + mkdtempSync, + readdirSync, + readFileSync, + rmSync, + writeFileSync, +} 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"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { + BasicTracerProvider, + InMemorySpanExporter, + type ReadableSpan, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import getPort from "get-port"; +import type { AgentAdapter, AgentInput, AgentRunContext } from "shared"; +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"; +import type { AgentDefinition } from "../../../core/agent/types"; + +const repositoryRoot = resolve(import.meta.dirname, "../../../../../.."); +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; + spanId: string; + parentSpanId: string | null; + inputs: unknown; + outputs: unknown; + status: string | 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)); + const semanticSpanIds = new Set( + spans.map((span) => span.spanContext().spanId), + ); + expect(traceIds.size, `${template}: mixed trace IDs`).toBe(1); + return { + template, + 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" + ? "mlflow.trace.tokenUsage" + : "mlflow.chat.tokenUsage" + ], + ) as Record; + const spanType = String(attributes["mlflow.spanType"] ?? ""); + const recordedParentSpanId = span.parentSpanContext?.spanId; + return { + name: span.name, + spanType, + spanId: span.spanContext().spanId, + 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, + 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 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); + 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]; + const supportedTypes = new Set([ + "AGENT", + "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" || span.spanType === "LLM", + ); + if (models.length === 0) + fail(root, "semantic child", "trace has no model child"); + for (const span of manifest.spans) { + 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) { + 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) { + 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"); + } + } +} + +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(); + } +} + +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 captureGeneratedHttpTurns( + candidate: GeneratedCandidate, +): 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) { + const startedAt = Date.now() - 10; + const injectFailure = JSON.stringify(input.messages).includes( + "INJECT_TRACE_FAILURE", + ); + yield { + type: "model_start", + stepId: "generated-step", + model: "generated-test-model", + provider: "databricks", + 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", + }); + 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)); + 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: "123456789", + 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; + let appkit: + | { server: { getServer(): Server }; shutdown(): Promise } + | undefined; + try { + 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 }; + shutdown(): Promise; + }>; + }; + expect( + generatedServer.app, + `${template} must export its generated createApp execution`, + ).toBeDefined(); + 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 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 }), + }, + ); + 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 failureTraceId = await invoke("INJECT_TRACE_FAILURE"); + await provider.forceFlush(); + const semanticSpans = exporter + .getFinishedSpans() + .filter((span) => span.attributes["mlflow.spanType"] !== undefined); + 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(); + 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(); + } + } +} + +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(); + rmSync(generatedApps, { recursive: true, force: true }); +}); + +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 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); + const reloaded = JSON.parse( + JSON.stringify(localManifest), + ) as TraceManifest; + assertContract(reloaded); + + const { success: manifest, failure } = + await captureGeneratedHttpTurns(candidate); + expect(process.listenerCount("SIGTERM"), candidate.name).toBe( + baselineSigterm, + ); + expect(process.listenerCount("SIGINT"), candidate.name).toBe( + baselineSigint, + ); + 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); + 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( + `${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( + (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 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" }; + 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[][] }; +} + +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, + 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, + 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(); +} + +interface UcSpanRow { + traceId: unknown; + spanId: unknown; + parentSpanId: unknown; + name: unknown; + attributes: Record; + statusCode?: unknown; + startTimeUnixNano?: unknown; + endTimeUnixNano?: unknown; +} + +interface DeployedTraceProof { + appName: string; + configuredExperimentId: string; + requestBody: unknown; + responseBody: unknown; + expectedTool: { + name: string; + inputs: unknown; + outputs: unknown; + }; + 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 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, + 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, + requestBody, + responseBody, + expectedTool, + 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, + name: row.name, + attributes: structuredClone(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 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, + rows: UcSpanRow[], +): TraceManifest { + return { + template, + traceId, + spans: rows.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: 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 normalizeMlflowSpans( + template: string, + traceId: string, + spans: Array>, +): TraceManifest { + return { + template, + traceId, + spans: spans.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, + 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, + attributes: span.attributes, + 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 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, +): 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}`, + ); + } + 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}`, + ); + } + } + } +} + +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 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, + ); + 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) => objectValue(span.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 rawMlflowManifest = normalizeMlflowSpans( + proof.appName, + otelTraceId, + mlflowSpans, + ); + 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); + const roots = ucManifest.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 ( + !ucManifest.spans.some( + (span) => span.spanType === "CHAT_MODEL" || span.spanType === "LLM", + ) + ) { + throw new Error("deployed trace is missing an LLM span"); + } + const toolSpan = ucManifest.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 ucManifest; +} + +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", + spans_table_name: "main.agent_traces.appkit_otel_spans", + }, + }, + spansTable: "main.agent_traces.appkit_otel_spans", + mlflowTracePrefix: "trace:/main.agent_traces.appkit/", + }); +}); + +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( + "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.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.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); +}); + +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", + (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", + 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_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 appName = process.env.APPKIT_TRACE_CONFORMANCE_APP_NAME ?? ""; + 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 ?? "", + process.env.APPKIT_TRACE_CONFORMANCE_UC_TABLE_PREFIX ?? "", + ); + const token = JSON.parse( + execFileSync("databricks", ["auth", "token", "-p", profile], { + 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 ?? "", + { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + "X-MLflow-Return-Trace-Id": "true", + }, + 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( + traceId ?? "", + binding.mlflowTracePrefix, + ); + const experiment = JSON.parse( + execFileSync( + "databricks", + [ + "api", + "get", + `/api/2.0/mlflow/experiments/get?experiment_id=${encodeURIComponent(configuredExperimentId)}`, + "-p", + profile, + ], + { encoding: "utf8" }, + ), + ); + 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 traceRecord = (storedTrace?.trace ?? + storedTrace) as DeployedTraceProof["traceRecord"]; + 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: persistedSpanStatement, + wait_timeout: "50s", + parameters: [ + { + name: "otel_spans_table", + type: "STRING", + value: binding.spansTable, + }, + { name: "trace_id", type: "STRING", value: otelTraceId }, + ], + }), + ], + { encoding: "utf8" }, + ), + ); + 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 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], + }), + ); + validateDeployedProof({ + appName, + configuredExperimentId, + requestBody, + responseBody, + expectedTool, + 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 7c4622cd3..ac9a6a162 100644 --- a/packages/appkit/src/plugins/agents/thread-store.ts +++ b/packages/appkit/src/plugins/agents/thread-store.ts @@ -1,5 +1,85 @@ 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, + key: string, + inputs: unknown, + operation: (span: Span) => Promise, +): Promise { + return tracer().startActiveSpan( + `thread.${operationName}`, + { + attributes: { + "mlflow.spanType": "MEMORY", + "appkit.memory.operation": operationName, + "appkit.memory.store": "thread", + "appkit.memory.key": key, + }, + }, + async (span) => { + const startedAt = Date.now(); + setCapturedAttribute(span, "mlflow.spanInputs", inputs); + try { + const result = await operation(span); + setCapturedAttribute( + span, + "mlflow.spanOutputs", + result === undefined ? { completed: true } : 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 +144,72 @@ 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, { 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, + { 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, { 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", + threadId, + { 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, + { 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..6fdf339e0 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.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/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/attributes.ts b/packages/appkit/src/telemetry/agent-tracing/attributes.ts new file mode 100644 index 000000000..c4701d9fe --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/attributes.ts @@ -0,0 +1,27 @@ +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", + "sdk-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..12e11fc0d --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/index.ts @@ -0,0 +1,19 @@ +export { + attachRemoteTraceLink, + injectActiveTraceContext, + type RemoteTraceReference, + verifiedAgentRemoteTrace, +} from "./propagation"; +export { captureTraceValue, normalizeFailureOutput } from "./serialization"; +export { + getActiveAgentTraceIdentity, + resolveAgentTraceAppName, + runWithAgentTrace, +} from "./tracer"; +export type { + AgentTraceObserver, + CapturedTraceValue, + CaptureTraceValueOptions, + ConsumedAgentStream, +} from "./types"; +export { AgentUsageAccumulator } from "./usage"; 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..65b3e9d7a --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/propagation.ts @@ -0,0 +1,102 @@ +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})$/; +const OTEL_TRACE_ID = /^[0-9a-f]{32}$/; + +export interface RemoteTraceReference { + traceId: string; + otelTraceId: string; + spanId: string; + 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. + */ +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 remoteContext = { + traceId: reference.otelTraceId, + spanId: reference.spanId, + traceFlags: TraceFlags.SAMPLED, + isRemote: true, + }; + if ( + remoteOtelTraceId(reference.traceId) !== 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/serialization.ts b/packages/appkit/src/telemetry/agent-tracing/serialization.ts new file mode 100644 index 000000000..bd9a73272 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/serialization.ts @@ -0,0 +1,120 @@ +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) => + normalizeRedactKey(key), + ), + ); + 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); + 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 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, +): { 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, ""); +} + +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/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(); + }); +}); 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..78f15724d --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/tests/serialization.test.ts @@ -0,0 +1,147 @@ +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("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: '"😀', + 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); + }); + + 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/agent-tracing/tests/tracer.test.ts b/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts new file mode 100644 index 000000000..a50cd00e4 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/tests/tracer.test.ts @@ -0,0 +1,644 @@ +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; + updateIdentity(identity: Partial>): void; + setOutput(output: unknown): void; + recordError(error: unknown, output?: unknown): 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", + "gen_ai.operation.name": "chat", + "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, + "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, + "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("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; + + 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]","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); + 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("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"); + + 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]","partial_output":"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("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]","partial_output":{"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(); + + 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); + }); + + 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/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/tracer.ts b/packages/appkit/src/telemetry/agent-tracing/tracer.ts new file mode 100644 index 000000000..4280515a1 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/tracer.ts @@ -0,0 +1,506 @@ +import { randomUUID } from "node:crypto"; +import { + context, + createContextKey, + isSpanContextValid, + type Span, + SpanStatusCode, + trace, +} from "@opentelemetry/api"; +import type { + AgentModelEndEvent, + AgentModelStartEvent, + AgentUsage, +} from "shared"; +import { getMlflowUcTraceId } from "../mlflow-uc"; +import { attachRemoteTraceLink, remoteOtelTraceId } from "./propagation"; +import { captureTraceValue, normalizeFailureOutput } from "./serialization"; +import type { + AgentTraceIdentity, + AgentTraceObserver, + AgentTraceResult, +} from "./types"; +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; + 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 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": 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) => { + 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, + }), + ) + .setValue(ACTIVE_AGENT_TRACE_IDENTITY, activeIdentity); + const traceId = getMlflowUcTraceId(otelTraceId) ?? otelTraceId; + const usage = new AgentUsageAccumulator(); + const activeModels = new Map(); + const completedModels = new Set(); + let outputText = ""; + let lifecycleError = false; + let reportedError: unknown; + let explicitOutput: unknown; + let hasExplicitOutput = 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); + const normalizedUsage = normalizeUsage(event.usage); + usage.add(normalizedUsage); + const active = activeModels.get(event.stepId); + activeModels.delete(event.stepId); + if (event.error) { + lifecycleError = true; + reportedError ??= event.error; + } + if (active) { + finalizeModelSpan(active, { + ...event, + usage: normalizedUsage, + }); + } + 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; + } + }, + addChildUsage(childUsage) { + usage.add(normalizeUsage(childUsage)); + }, + linkToRun(runId) { + root.setAttribute("mlflow.sourceRun", runId); + }, + updateIdentity(next) { + updateActiveIdentity(activeIdentity, next); + setIdentityAttributes(root, next); + }, + setOutput(output) { + explicitOutput = output; + hasExplicitOutput = true; + }, + recordError(error, output) { + lifecycleError = true; + reportedError ??= error; + if (output !== undefined) { + explicitOutput = output; + hasExplicitOutput = 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, + reportedError ?? "Agent lifecycle reported an error", + "Agent operation failed", + ); + } + + const finalUsage = usage.snapshot(); + setRootUsageAttributes(root, finalUsage); + onCompleteUsage?.(finalUsage); + const finalOutputText = outputText || textFromValue(value); + const finalOutput = + failed || lifecycleError + ? normalizeFailureOutput( + hasExplicitOutput ? explicitOutput : finalOutputText, + operationError ?? reportedError, + ) + : hasExplicitOutput + ? explicitOutput + : outputText || textFromValue(value) + ? { text: finalOutputText } + : value; + setCapturedAttribute( + root, + "mlflow.spanOutputs", + finalOutput, + failed || lifecycleError ? ["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 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", + "gen_ai.operation.name": "chat", + "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.error + ? normalizeFailureOutput(event.output, event.error) + : event.output, + event.error ? ["error"] : undefined, + ); + 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( + "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, + ); + } + 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) { + 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 { + 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, + 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", + 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 + : 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 new file mode 100644 index 000000000..17cb43138 --- /dev/null +++ b/packages/appkit/src/telemetry/agent-tracing/types.ts @@ -0,0 +1,50 @@ +import type { AgentEvent, 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; +} + +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; + /** Adds one completed local child trace's aggregate usage to this root. */ + addChildUsage(usage: AgentUsage): void; + /** Links this trace to the MLflow run that launched the agent invocation. */ + linkToRun(runId: string): void; + updateIdentity(identity: Partial>): void; + setOutput(output: unknown): void; + recordError(error: unknown, output?: unknown): void; +} + +export interface AgentTraceResult { + value: T; + traceId: string; + usage: AgentUsage; +} 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/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..9d037146c --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/config.ts @@ -0,0 +1,63 @@ +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; + +const UC_IDENTIFIER = /^[A-Za-z_][A-Za-z0-9_]{0,254}$/; + +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])?.trim(), + ]), + ) 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(", ")}`, + ); + } + + 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 new file mode 100644 index 000000000..1104b9cff --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/exporter.ts @@ -0,0 +1,521 @@ +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"; +import { createLogger } from "../../logging/logger"; +import { + createWorkspaceClient, + type WorkspaceClient, +} from "../../workspace-client"; +import type { MlflowUcConfig } from "./config"; +import { + buildMlflowUcTraceInfo, + constructMlflowV4TraceId, + 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 | (() => 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 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( + private readonly config: MlflowUcConfig, + private readonly client: WorkspaceClient = createWorkspaceClient(), + options: ExporterOptions = {}, + ) { + this.createOtlpExporter = + 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( + spans: ReadableSpan[], + resultCallback: (result: ExportResult) => void, + ): void { + if (this.closed) { + resultCallback({ + code: ExportResultCode.FAILED, + error: new Error("MLflow UC trace exporter is shut down"), + }); + return; + } + + const batches = this.accumulateReadyBatches(spans); + if (batches.length === 0) { + resultCallback({ code: ExportResultCode.SUCCESS }); + return; + } + + const operation = (async () => { + try { + for (const batch of batches) { + await this.exportBatchWithRetry(batch); + } + resultCallback({ code: ExportResultCode.SUCCESS }); + } catch (error) { + resultCallback({ + code: ExportResultCode.FAILED, + error: toError(error), + }); + } + })(); + this.track(operation); + } + + exportTrace( + batch: MlflowUcExportBatch, + resultCallback: (result: ExportResult) => void, + ): void { + if (this.closed) { + resultCallback({ + code: ExportResultCode.FAILED, + error: new Error("MLflow UC trace exporter is shut down"), + }); + return; + } + + const operation = this.exportBatchWithRetry(batch).then( + () => resultCallback({ code: ExportResultCode.SUCCESS }), + (error) => + resultCallback({ + code: ExportResultCode.FAILED, + error: toError(error), + }), + ); + this.track(operation); + } + + async forceFlush(): Promise { + while (this.inFlight.size > 0) { + await Promise.allSettled([...this.inFlight]); + } + } + + async shutdown(): Promise { + 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 pending = this.pendingSpans.get(traceId); + const accumulated = new Map(pending?.spans); + for (const span of newSpans) { + accumulated.set(span.spanContext().spanId, span); + } + + const traceSpans = [...accumulated.values()]; + const semanticRoot = findSemanticRoot(traceSpans); + if (semanticRoot) { + const semanticSpans = this.capCompletedTrace( + findSemanticSubtree(traceSpans, semanticRoot), + semanticRoot, + ); + this.pendingSpans.delete(traceId); + batches.push({ + traceInfo: buildMlflowUcTraceInfo( + this.config, + semanticRoot, + semanticSpans, + ), + 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 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( + () => this.inFlight.delete(operation), + () => this.inFlight.delete(operation), + ); + } + + private async exportBatch(batch: MlflowUcExportBatch): Promise { + await this.withDeadline( + this.client.config.ensureResolved(), + "workspace configuration", + ); + 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 traceInfoResponse = await fetch( + `${host}/api/4.0/mlflow/traces/${encodeURIComponent(location)}/${encodeURIComponent(traceInfo.trace_id)}/info`, + { + method: "POST", + headers: { + ...traceInfoHeaders, + "Content-Type": "application/json", + }, + body: JSON.stringify(traceInfo), + signal: AbortSignal.timeout(this.operationTimeoutMs), + }, + ); + if (!traceInfoResponse.ok) { + throw new TraceInfoRequestError( + traceInfoResponse.status, + `MLflow trace-info request failed with ${traceInfoResponse.status}: ${await traceInfoResponse.text()}`, + parseRetryAfter(traceInfoResponse.headers.get("retry-after")), + ); + } + await traceInfoResponse.body?.cancel(); + + 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: async () => ({ + ...(await this.freshAuthHeaders()), + "X-Databricks-UC-Table-Name": this.config.otelSpansTableName, + }), + timeoutMillis: this.operationTimeoutMs, + }); + return this.otlpExporter; + } + + private async freshAuthHeaders(): Promise> { + const headers = new 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 { + 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) { + if (!isSpanContextValid(parent)) return true; + const parentSpan = bySpanId.get(parent.spanId); + if (!parentSpan) return parent.isRemote === true; + if (parentSpan.attributes["mlflow.spanType"] === "AGENT") return false; + parent = parentSpan.parentSpanContext; + } + return true; + }); +} + +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/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..d5f2d73ca --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/processor.ts @@ -0,0 +1,248 @@ +import type { Context } from "@opentelemetry/api"; +import type { + ReadableSpan, + Span, + SpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import { createLogger } from "../../logging/logger"; +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, + MlflowUcTraceRegistry, +} from "./trace-info"; + +interface PendingTrace { + spans: ReadableSpan[]; + memberSpanIds: Set; + 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; + + constructor( + 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; + const now = Date.now(); + this.evictExpired(now); + const spanContext = span.spanContext(); + const otelTraceId = spanContext.traceId; + 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 (span.attributes[MLFLOW_SPAN_TYPE_ATTRIBUTE] === "AGENT") { + const registeredRoot = this.registry.registerSemanticRoot( + otelTraceId, + spanContext.spanId, + ); + if (!pending && registeredRoot) { + 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.size < MAX_SPANS_PER_TRACE + ) { + pending.memberSpanIds.add(spanContext.spanId); + } + return; + } + + if ( + pending && + span.parentSpanContext && + pending.memberSpanIds.has(span.parentSpanContext.spanId) && + pending.memberSpanIds.size < MAX_SPANS_PER_TRACE + ) { + pending.memberSpanIds.add(spanContext.spanId); + } + } + + 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); + if (!pending) return; + + const semanticRootSpanId = this.registry.getSemanticRootSpanId(otelTraceId); + if (!pending.memberSpanIds.has(spanContext.spanId)) { + return; + } + pending.lastTouchedMs = now; + + 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); + const batch: MlflowUcExportBatch = { + traceInfo: buildMlflowUcTraceInfo(this.config, span, pending.spans), + spans: pending.spans, + }; + this.startExport(batch); + } + + async forceFlush(): Promise { + while (this.inFlight.size > 0) { + await Promise.allSettled([...this.inFlight]); + } + await this.exporter.forceFlush(); + } + + async shutdown(): Promise { + 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 { + 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); + 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 { + 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 new file mode 100644 index 000000000..dc04fc5fb --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/config.test.ts @@ -0,0 +1,95 @@ +import { describe, expect, test } from "vitest"; +import { constructMlflowV4TraceId, resolveMlflowUcConfig } from "../../index"; + +const completeEnv = { + 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", +}; + +describe("resolveMlflowUcConfig", () => { + test("resolves every required field from the environment", () => { + expect(resolveMlflowUcConfig(completeEnv)).toEqual({ + experimentId: "123456789", + 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: "987654321", + tablePrefix: "custom", + otelSpansTableName: "main.agent_traces.custom_otel_spans", + }), + ).toEqual({ + experimentId: "987654321", + catalogName: "main", + schemaName: "agent_traces", + tablePrefix: "custom", + otelSpansTableName: "main.agent_traces.custom_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("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", () => { + 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..6c4f341be --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/exporter.test.ts @@ -0,0 +1,1026 @@ +import { once } from "node:events"; +import { createServer, type ServerResponse } from "node:http"; +import type { AddressInfo } from "node:net"; +import { context, SpanStatusCode, TraceFlags, trace } 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())); + vi.unstubAllGlobals(); +}); + +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( + options: { nestedAgent?: boolean } = {}, +): 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.sourceRun": "run-99", + "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", + }, + }); + if (options.nestedAgent) { + const nestedAgent = tracer.startSpan( + "helper-agent", + { + startTime: 1_700_000_000_100, + attributes: { "mlflow.spanType": "AGENT" }, + }, + 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 }); + root.end(1_700_000_000_250); + await provider.forceFlush(); + const spans = inMemory.getFinishedSpans(); + await provider.shutdown(); + 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), + ); +} + +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("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, + }); + expect(traceInfoRoots).toEqual([]); + expect(uploads).toEqual([]); + 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("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]; + 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"]); + 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: 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", + 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.sourceRun": "run-99", + "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("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"]); + const spans = await createTraceSpans(); + const exportCalls: Array<{ + url: string; + headers: Record | (() => Promise>); + 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`, + 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", + }, + { + 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("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, + maxAttempts: 2, + retryDelayMs: 0, + }); + + await expect(exportSpans(exporter, spans)).resolves.toEqual({ + 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", + 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.FAILED, + error: expect.any(Error), + }); + expect(logger.error).toHaveBeenCalledWith( + "MLflow UC trace export failed: %O", + expect.objectContaining({ + error: "Databricks workspace host is unavailable for MLflow UC export", + }), + ); + }); + + 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) => { + 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); + }); + + 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.FAILED, + error: expect.any(Error), + }); + 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); + }); +}); 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..b8eed717f --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/tests/processor.test.ts @@ -0,0 +1,314 @@ +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", + "mlflow.sourceRun": "run-99", + "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: otelTraceId, + 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}', + "mlflow.sourceRun": "run-99", + "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("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) => { + 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("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) => { + 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); + 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(); + }); + + 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] }); + 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_000); + expect(new Set(batches.map((batch) => batch.traceInfo.trace_id)).size).toBe( + 10_000, + ); + expect(batches.some((batch) => batch.spans[0]?.name === "agent-0")).toBe( + false, + ); + + 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 new file mode 100644 index 000000000..cd8e1ca6e --- /dev/null +++ b/packages/appkit/src/telemetry/mlflow-uc/trace-info.ts @@ -0,0 +1,233 @@ +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", + "mlflow.sourceRun", + "appkit.app.name", + "appkit.request.id", + "appkit.thread.id", + "appkit.agent.name", + "appkit.route", +] as const; + +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; + 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; + + 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; + } + + deleteTrace(otelTraceId: string): void { + this.traces.delete(otelTraceId); + } + + clear(): void { + this.traces.clear(); + } + + 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 { + // 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", + ), + trace_location: { + type: "UC_TABLE_PREFIX", + uc_table_prefix: { + catalog_name: config.catalogName, + schema_name: config.schemaName, + table_prefix: config.tablePrefix, + 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/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; diff --git a/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts b/packages/appkit/src/telemetry/tests/telemetry-manager.test.ts index 11b85d9bf..7265bea56 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 = "123456789"; + 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 = "123456789"; + 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/packages/shared/src/agent.ts b/packages/shared/src/agent.ts index 5ec2caf35..329fb9659 100644 --- a/packages/shared/src/agent.ts +++ b/packages/shared/src/agent.ts @@ -112,6 +112,55 @@ 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 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 } | { type: "message"; content: string } @@ -144,7 +193,10 @@ export type AgentEvent = toolName: string; args: unknown; annotations?: ToolAnnotations; - }; + } + | AgentModelStartEvent + | AgentModelEndEvent + | AgentRemoteTraceEvent; // --------------------------------------------------------------------------- // Responses API types (OpenAI-compatible wire format for HTTP boundary) @@ -232,7 +284,10 @@ export interface AppKitThinkingEvent { export interface AppKitMetadataEvent { type: "appkit.metadata"; - data: Record; + data: Record & { + threadId?: string; + traceId?: string; + }; sequence_number: number; } 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, ); 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/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..2b2fa09f2 --- /dev/null +++ b/packages/shared/src/cli/commands/setup-mlflow-uc.test.ts @@ -0,0 +1,272 @@ +import { mkdirSync, mkdtempSync, readFileSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import yaml from "js-yaml"; +import { afterEach, describe, expect, test, vi } from "vitest"; +import { + buildMlflowProvisionCommand, + projectRequiresMlflowUc, + provisionAndPersistMlflowUc, + setupCommand, +} 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']", + "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; +} + +describe("MLflow UC setup", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + 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", + runtimePrincipal: "runtime-app-sp", + }), + ).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", + "--runtime-principal", + "runtime-app-sp", + "--output-json", + "/workspace/traced-app/.databricks/mlflow-uc.json", + ]); + }); + + 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("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(); + + 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[] = []; + 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( + { + 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(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); + 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 c72e661eb..9d8ab619a 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,289 @@ 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; + +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 +>; + +export interface MlflowUcSetupOptions { + cwd: string; + profile: string; + experimentName: string; + catalog: string; + schema: string; + tablePrefix: string; + warehouseId: string; + runtimePrincipal: 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, + "--runtime-principal", + options.runtimePrincipal, + "--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, + 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}=`)) && + (!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`); +} + +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 : []; + const existingByName = new Map(existing.map((entry) => [entry.name, entry])); + 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) => { + 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); +} + +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) { + 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], + }; + } + document.targets ??= {}; + document.targets.default ??= {}; + document.targets.default.variables ??= {}; + for (const name of MLFLOW_UC_ENV_NAMES) { + document.targets.default.variables[MLFLOW_BUNDLE_VARIABLE_NAMES[name]] = + values[name]; + } + 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[]) => { + const result = spawnSync(argv[0], argv.slice(1), { + cwd: options.cwd, + stdio: "inherit", + }); + 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}`); + } + + const values = validateMlflowUcValues( + JSON.parse(fs.readFileSync(outputPath, "utf8")), + ); + 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; + 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 +400,74 @@ 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; + mlflowRuntimePrincipal?: 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 @@ -154,6 +503,52 @@ function runSetup(options: { write?: boolean }) { action = "Created"; } + // 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")); + 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 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( + { + cwd, + profile, + experimentName: `/Users/${userName}/appkit-agent-traces`, + catalog: options.mlflowCatalog, + schema: options.mlflowSchema, + tablePrefix: options.mlflowTablePrefix, + warehouseId, + runtimePrincipal: runtimePrincipal.trim(), + }, + { workspaceHost }, + ); + } + if (shouldWrite) { fs.writeFileSync(claudePath, finalContent); console.log(`\n✓ ${action} CLAUDE.md`); @@ -180,13 +575,29 @@ function runSetup(options: { write?: boolean }) { } 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", + ) + .option( + "--mlflow-runtime-principal ", + "Deployed app service principal receiving explicit UC trace grants", + ) .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/pnpm-lock.yaml b/pnpm-lock.yaml index 4c2eeca07..a55c39687 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -267,6 +267,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) @@ -327,9 +330,6 @@ importers: magic-string: specifier: 0.30.21 version: 0.30.21 - mlflow-tracing: - specifier: 0.1.3 - version: 0.1.3 obug: specifier: 2.1.1 version: 2.1.1 @@ -1974,10 +1974,6 @@ packages: engines: {node: ^20 || ^22 || ^24 || ^25, pnpm: '>=10'} hasBin: true - '@databricks/sdk-experimental@0.15.0': - resolution: {integrity: sha512-HkoMiF7dNDt6WRW0xhi7oPlBJQfxJ9suJhEZRFt08VwLMaWcw2PiF8monfHlkD4lkufEYV6CTxi5njQkciqiHA==} - engines: {node: '>=22.0', npm: '>=10.0.0'} - '@databricks/sdk-experimental@0.17.0': resolution: {integrity: sha512-dOJIt4F2nBk6HKObnv7Xbmy/qLYTy2835qhXSuW0Qw1QAXui9plmCet1KqG3yeQcMTyncWGbnhjGdQi8GEGQSA==} engines: {node: '>=22.0', npm: '>=10.0.0'} @@ -2810,10 +2806,6 @@ packages: '@octokit/types@16.0.0': resolution: {integrity: sha512-sKq+9r1Mm4efXW1FCk7hFSeJo4QKreL/tTbR0rz/qx/r1Oa2VV83LTA/H/MuCOX7uCIJmQVRKBcbmWoySjAnSg==} - '@opentelemetry/api-logs@0.205.0': - resolution: {integrity: sha512-wBlPk1nFB37Hsm+3Qy73yQSobVn28F4isnWIBvKpd5IUH/eat8bwcL02H9yzmHyyPmukeccSl2mbN5sDQZYnPg==} - engines: {node: '>=8.0.0'} - '@opentelemetry/api-logs@0.219.0': resolution: {integrity: sha512-FFx7YnaYJlIjqWW/AG/yAZ0L/NEY724PipXXXQLdtZPbLwBGbUMTGL1i/esI56TWfTUXxhLfpgrnWJCG8aUJyg==} engines: {node: '>=8.0.0'} @@ -2835,12 +2827,6 @@ packages: peerDependencies: '@opentelemetry/api': ^1.9.0 - '@opentelemetry/context-async-hooks@2.1.0': - resolution: {integrity: sha512-zOyetmZppnwTyPrt4S7jMfXiSX9yyfF0hxlA8B5oo2TtKl+/RGCy7fi4DrBfIf3lCPrkKsRBWZZD7RFojK7FDg==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.0.0 <1.10.0' - '@opentelemetry/context-async-hooks@2.8.0': resolution: {integrity: sha512-/3FIraneMcng67SUJCxvyInk/oxzwsxyadufk0wwfOBLf5wqtAGX4MoQASwSbndBPeARzBryUM9Azr5kHIdWLw==} engines: {node: ^18.19.0 || >=20.6.0} @@ -2853,132 +2839,66 @@ packages: peerDependencies: '@opentelemetry/api': '>=1.0.0 <1.10.0' - '@opentelemetry/exporter-logs-otlp-grpc@0.205.0': - resolution: {integrity: sha512-jQlw7OHbqZ8zPt+pOrW2KGN7T55P50e3NXBMr4ckPOF+DWDwSy4W7mkG09GpYWlQAQ5C9BXg5gfUlv5ldTgWsw==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-logs-otlp-grpc@0.219.0': resolution: {integrity: sha512-7SvzDCIclHWAcCwZ1MTOLcwn4BVNPGI3QxS/DJraPNe1TTL+4TvUBq5zeQV8tsnYvtDN7wKW2qocVmaCP2l7sQ==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-logs-otlp-http@0.205.0': - resolution: {integrity: sha512-5JteMyVWiro4ghF0tHQjfE6OJcF7UBUcoEqX3UIQ5jutKP1H+fxFdyhqjjpmeHMFxzOHaYuLlNR1Bn7FOjGyJg==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-logs-otlp-http@0.219.0': resolution: {integrity: sha512-mhl2HL6GmZI8b8PwPfqMws/5ovJfbRTxwc9Y5agVVHiQ+e5SL1btsFr/kJDgt7YCexDtsUn5HAreHQO9szFS0A==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-logs-otlp-proto@0.205.0': - resolution: {integrity: sha512-q3VS9wS+lpZ01txKxiDGBtBpTNge3YhbVEFDgem9ZQR9eI3EZ68+9tVZH9zJcSxI37nZPJ6lEEZO58yEjYZsVA==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-logs-otlp-proto@0.219.0': resolution: {integrity: sha512-Ayw4Gf71PS9jhBVaYywa4WsajnqfDehMkTdVH3TSAVHqPcsAv/AhH/wTNRYNt99szeYr6Gbd/D6RjZD77wAxHg==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-metrics-otlp-grpc@0.205.0': - resolution: {integrity: sha512-1Vxlo4lUwqSKYX+phFkXHKYR3DolFHxCku6lVMP1H8sVE3oj4wwmwxMzDsJ7zF+sXd8M0FCr+ckK4SnNNKkV+w==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-metrics-otlp-grpc@0.219.0': resolution: {integrity: sha512-6LaaSrPxK5L55bXevWajvOMxGOpNm0n12tG53TeZaUeNzXwLPg6d2KCC1zAlGsojan+xRG71mA4Qqs9K2VVrKQ==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-metrics-otlp-http@0.205.0': - resolution: {integrity: sha512-fFxNQ/HbbpLmh1pgU6HUVbFD1kNIjrkoluoKJkh88+gnmpFD92kMQ8WFNjPnSbjg2mNVnEkeKXgCYEowNW+p1w==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-metrics-otlp-http@0.219.0': resolution: {integrity: sha512-6CaDRbMVHZSDWzNXwrR8y/H4B/Z1eMNnkHiPQlTx3Ojz2OHY4X/aff/UC4P/3pHUQSuTfi3oh2UsPPZppw+Vrg==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-metrics-otlp-proto@0.205.0': - resolution: {integrity: sha512-qIbNnedw9QfFjwpx4NQvdgjK3j3R2kWH/2T+7WXAm1IfMFe9fwatYxE61i7li4CIJKf8HgUC3GS8Du0C3D+AuQ==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-metrics-otlp-proto@0.219.0': resolution: {integrity: sha512-DUS7XyIiEnoeccQUvuKy0G2/YqeKhpN8FVIrGbrLNIVMj10yeIFLRzRv0tibCI2kXXvlTTABVexGAk78wHk2ug==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-prometheus@0.205.0': - resolution: {integrity: sha512-xsot/Qm9VLDTag4GEwAunD1XR1U8eBHTLAgO7IZNo2JuD/c/vL7xmDP7mQIUr6Lk3gtj/yGGIR2h3vhTeVzv4w==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-prometheus@0.219.0': resolution: {integrity: sha512-TxOnJ85eWJY5JyOJsNMXiRTYlkDcOv0u3KbXEzWCc+tUS9sjL/BC6BcdxZ0B9r2OFVqsrZFXUzSD2sZUy42Ucw==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-trace-otlp-grpc@0.205.0': - resolution: {integrity: sha512-ZBksUk84CcQOuDJB65yu5A4PORkC4qEsskNwCrPZxDLeWjPOFZNSWt0E0jQxKCY8PskLhjNXJYo12YaqsYvGFA==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-trace-otlp-grpc@0.219.0': resolution: {integrity: sha512-BkDNv1UD6BscW19MxbAxVmSYSSFuyeqR6buV2/HTYqA7GrR0EbTFzqG6h86T3PtXmpdbsWjMGLDdjG2rikG27Q==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-trace-otlp-http@0.205.0': - resolution: {integrity: sha512-vr2bwwPCSc9u7rbKc74jR+DXFvyMFQo9o5zs+H/fgbK672Whw/1izUKVf+xfWOdJOvuwTnfWxy+VAY+4TSo74Q==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-trace-otlp-http@0.219.0': resolution: {integrity: sha512-9t6SvBXXBEjOBcIzgozvBbd3jWrv3Gt3ngGhl1fhdZ/zRc7oZDVOFEqbi2zlBpW9BXhgDMKv422J0DL/3iQWfw==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-trace-otlp-proto@0.205.0': - resolution: {integrity: sha512-bGtFzqiENO2GpJk988mOBMe0MfeNpTQjbLm/LBijas6VRyEDQarUzdBHpFlu89A25k1+BCntdWGsWTa9Ai4FyA==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-trace-otlp-proto@0.219.0': resolution: {integrity: sha512-lF/LUBfhOFmxJa+SQsLN7ziV4MHa2pyKgOM6JNehSOfU+npjM4gwm9oIKEJrzrWcexMcqydiyoFy0XCb1Ql3wQ==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/exporter-zipkin@2.1.0': - resolution: {integrity: sha512-0mEI0VDZrrX9t5RE1FhAyGz+jAGt96HSuXu73leswtY3L5YZD11gtcpARY2KAx/s6Z2+rj5Mhj566JsI2C7mfA==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.0.0 - '@opentelemetry/exporter-zipkin@2.8.0': resolution: {integrity: sha512-Mj84UkEa17BK2o903VTXW3wM8CrSZexGs4tRGVZVIMM9ni1T6TuGx5IrRfoWKAbshx42D5/kc7YV+axypLPYyA==} engines: {node: ^18.19.0 || >=20.6.0} @@ -3231,48 +3151,24 @@ packages: peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/instrumentation@0.205.0': - resolution: {integrity: sha512-cgvm7tvQdu9Qo7VurJP84wJ7ZV9F6WqDDGZpUc6rUEXwjV7/bXWs0kaYp9v+1Vh1+3TZCD3i6j/lUBcPhu8NhA==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/instrumentation@0.219.0': resolution: {integrity: sha512-X5t7I8GyIO9rmGHwoedZLREpQqrF1WW2nxzNNym6HOKpFiE+rvqV3ngC0xcZVO2YwIGf3KKmRdWrYwdwz3H9RQ==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/otlp-exporter-base@0.205.0': - resolution: {integrity: sha512-2MN0C1IiKyo34M6NZzD6P9Nv9Dfuz3OJ3rkZwzFmF6xzjDfqqCTatc9v1EpNfaP55iDOCLHFyYNCgs61FFgtUQ==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/otlp-exporter-base@0.219.0': resolution: {integrity: sha512-zvIxQX/AZUVKDU+hCuYx+7UkiP7GRdnk1ZbFQRYzHvYp47cAWR4j3IhoPhV9KaeXEv2xdGq3IA6PnpzDmLcmSA==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/otlp-grpc-exporter-base@0.205.0': - resolution: {integrity: sha512-AeuLfrciGYffqsp4EUTdYYc6Ee2BQS+hr08mHZk1C524SFWx0WnfcTnV0NFXbVURUNU6DZu1DhS89zRRrcx/hg==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/otlp-grpc-exporter-base@0.219.0': resolution: {integrity: sha512-iIk/s8QQu39zpTrRRmsW/Eg3SE2+Hg8tLWepr2FLRgmwUpNd0IpCTLJEHJ77hpt4hgIS8MAh44UYI4xQPZwWlw==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': ^1.3.0 - '@opentelemetry/otlp-transformer@0.205.0': - resolution: {integrity: sha512-KmObgqPtk9k/XTlWPJHdMbGCylRAmMJNXIRh6VYJmvlRDMfe+DonH41G7eenG8t4FXn3fxOGh14o/WiMRR6vPg==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': ^1.3.0 - '@opentelemetry/otlp-transformer@0.219.0': resolution: {integrity: sha512-aaYKAyXhw9VchKZVGOopD3Gw/kPsyrX2c6IQ0AW32mTjqmZOh5Y6Gf5OYqTNqVktAeBjmFinhyFaCwW6GYK9YQ==} engines: {node: ^18.19.0 || >=20.6.0} @@ -3285,24 +3181,12 @@ packages: peerDependencies: '@opentelemetry/api': '>=1.0.0 <1.10.0' - '@opentelemetry/propagator-b3@2.1.0': - resolution: {integrity: sha512-yOdHmFseIChYanddMMz0mJIFQHyjwbNhoxc65fEAA8yanxcBPwoFDoh1+WBUWAO/Z0NRgk+k87d+aFIzAZhcBw==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.0.0 <1.10.0' - '@opentelemetry/propagator-b3@2.8.0': resolution: {integrity: sha512-SazlvuSKi5533rPHTW2TwBwdMakhjZST4SYs0YauuvfGDkT13KbG1gJS75hV0uWVeevhtVP9sAIlaZLTHdSbMg==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': '>=1.0.0 <1.10.0' - '@opentelemetry/propagator-jaeger@2.1.0': - resolution: {integrity: sha512-QYo7vLyMjrBCUTpwQBF/e+rvP7oGskrSELGxhSvLj5gpM0az9oJnu/0O4l2Nm7LEhAff80ntRYKkAcSwVgvSVQ==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.0.0 <1.10.0' - '@opentelemetry/propagator-jaeger@2.8.0': resolution: {integrity: sha512-Xnz9zZvvQzUw+9DrOn0MomR7BxFCkA2pcfXBQuHC28ndJpSbjLs7knzYb05kw5SyCjSsEWombkZMgGcJSk8JVg==} engines: {node: ^18.19.0 || >=20.6.0} @@ -3343,72 +3227,36 @@ packages: peerDependencies: '@opentelemetry/api': ^1.0.0 - '@opentelemetry/resources@2.1.0': - resolution: {integrity: sha512-1CJjf3LCvoefUOgegxi8h6r4B/wLSzInyhGP2UmIBYNlo4Qk5CZ73e1eEyWmfXvFtm1ybkmfb2DqWvspsYLrWw==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.3.0 <1.10.0' - '@opentelemetry/resources@2.8.0': resolution: {integrity: sha512-qmXQ27ilDbUK/vGMqwL8D4/rhn76C+sherM4wTbjlfknR8Nvfc/hCxjRJPhkzZzUsPiNg16SA31NxMabwttRjg==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': '>=1.3.0 <1.10.0' - '@opentelemetry/sdk-logs@0.205.0': - resolution: {integrity: sha512-nyqhNQ6eEzPWQU60Nc7+A5LIq8fz3UeIzdEVBQYefB4+msJZ2vuVtRuk9KxPMw1uHoHDtYEwkr2Ct0iG29jU8w==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.4.0 <1.10.0' - '@opentelemetry/sdk-logs@0.219.0': resolution: {integrity: sha512-s6lTKRakaPClvKoWHRChxnXjDMkM/TQ30ff78jN6EBGf7MI7VzANE5PU3f4z9qDUudWjvZjOLHG0rBnBKYvoXA==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': '>=1.4.0 <1.10.0' - '@opentelemetry/sdk-metrics@2.1.0': - resolution: {integrity: sha512-J9QX459mzqHLL9Y6FZ4wQPRZG4TOpMCyPOh6mkr/humxE1W2S3Bvf4i75yiMW9uyed2Kf5rxmLhTm/UK8vNkAw==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.9.0 <1.10.0' - '@opentelemetry/sdk-metrics@2.8.0': resolution: {integrity: sha512-UDBGaj6W0Rgy5rTTaoxs8gVGF/aGkAKyjurJv7se6wjRxJu7FoquTLT/vt54DZfo4crbprYfhX/SOK9+BPw1qg==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': '>=1.9.0 <1.10.0' - '@opentelemetry/sdk-node@0.205.0': - resolution: {integrity: sha512-Y4Wcs8scj/Wy1u61pX1ggqPXPtCsGaqx/UnFu7BtRQE1zCQR+b0h56K7I0jz7U2bRlPUZIFdnNLtoaJSMNzz2g==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.3.0 <1.10.0' - '@opentelemetry/sdk-node@0.219.0': resolution: {integrity: sha512-NWLpWLEb8gV3+JBHYoIrktbM385wyHpRJoh3J/4Q52d4PR+AlPMNGJT3DzBUrDSUEVbKAXoHR+EDAPxtiNcj8g==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': '>=1.3.0 <1.10.0' - '@opentelemetry/sdk-trace-base@2.1.0': - resolution: {integrity: sha512-uTX9FBlVQm4S2gVQO1sb5qyBLq/FPjbp+tmGoxu4tIgtYGmBYB44+KX/725RFDe30yBSaA9Ml9fqphe1hbUyLQ==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.3.0 <1.10.0' - '@opentelemetry/sdk-trace-base@2.8.0': resolution: {integrity: sha512-mhU4jp+vW0mGbFRd+GeXHvmfA4aDqWjBjLC3pE5XMpLs0IE2ryYb019Ts2AQrOq67gaTF25D91+fgvEHDZEnuQ==} engines: {node: ^18.19.0 || >=20.6.0} peerDependencies: '@opentelemetry/api': '>=1.3.0 <1.10.0' - '@opentelemetry/sdk-trace-node@2.1.0': - resolution: {integrity: sha512-SvVlBFc/jI96u/mmlKm86n9BbTCbQ35nsPoOohqJX6DXH92K0kTe73zGY5r8xoI1QkjR9PizszVJLzMC966y9Q==} - engines: {node: ^18.19.0 || >=20.6.0} - peerDependencies: - '@opentelemetry/api': '>=1.0.0 <1.10.0' - '@opentelemetry/sdk-trace-node@2.8.0': resolution: {integrity: sha512-nZt9OGufioAc3AfoLTqA9bsAeaMJAictYDdI2VcNQ+PmT+3rfKjAZDZvgPfd8VPX0O5Bw1hdQF6kDK8VSpZiWg==} engines: {node: ^18.19.0 || >=20.6.0} @@ -7285,9 +7133,6 @@ packages: fast-levenshtein@2.0.6: resolution: {integrity: sha512-DCXu6Ifhqcks7TZKY3Hxp3y6qphY5SJZmrWMDrKcERSOXWQdMhU9Ig/PYrzyw/ul9jOIyh0N4M0tbC5hodg8dw==} - fast-safe-stringify@2.1.1: - resolution: {integrity: sha512-W+KJc2dmILlPplD/H4K9l9LcAHAfPtP6BY84uVLXQ6Evcz9Lcg33Y2z1IVblT6xdY54PXYVHEv+0Wpq8Io6zkA==} - fast-uri@3.1.0: resolution: {integrity: sha512-iPeeDKJSWf4IEOasVVrknXpaBV0IApz/gp7S2bb7Z4Lljbl2MGJRqInZiUrQwV16cpzw/D3S5j5Julj/gT52AA==} @@ -7538,7 +7383,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: @@ -7966,9 +7811,6 @@ packages: resolution: {integrity: sha512-TR3KfrTZTYLPB6jUjfx6MF9WcWrHL9su5TObK4ZkYgBdWKPOFoSoQIdEuTuR82pmtxH2spWG9h6etwfr1pLBqQ==} engines: {node: '>=6'} - import-in-the-middle@1.15.0: - resolution: {integrity: sha512-bpQy+CrsRmYmoPMAE/0G33iwRqwW4ouqdRg8jgbH3aKuCtOc8lxgmYXg2dMM92CRiGP660EtBcymH/eVUpCSaA==} - import-in-the-middle@3.0.1: resolution: {integrity: sha512-pYkiyXVL2Mf3pozdlDGV6NAObxQx13Ae8knZk1UJRJ6uRW/ZRmTGHlQYtrsSl7ubuE5F8CD1z+s1n4RHNuTtuA==} engines: {node: '>=18'} @@ -8017,10 +7859,6 @@ packages: resolution: {integrity: sha512-QQnnxNyfvmHFIsj7gkPcYymR8Jdw/o7mp5ZFihxn6h8Ci6fh3Dx4E1gPjpQEpIuPo9XVNY/ZUwh4BPMjGyL01g==} engines: {node: ^14.17.0 || ^16.13.0 || >=18.0.0} - ini@5.0.0: - resolution: {integrity: sha512-+N0ngpO3e7cRUWOJAS7qw0IZIVc6XPrW4MlFBdD066F2L4k1L6ker3hLqSq7iXxU5tgS4WGkIUElWn5vogAEnw==} - engines: {node: ^18.17.0 || >=20.5.0} - ini@6.0.0: resolution: {integrity: sha512-IBTdIkzZNOpqm7q3dRqJvMaldXjDHWkEDfrwGEQTs5eaQMWV+djAhR+wahyNNMAa+qpbDUhBMVt4ZKNwpPm7xQ==} engines: {node: ^20.17.0 || >=22.9.0} @@ -9069,10 +8907,6 @@ packages: engines: {node: '>=10'} hasBin: true - mlflow-tracing@0.1.3: - resolution: {integrity: sha512-Koqkwaid5ubGHuLprBP6J7Su70WddlD11f2vgzgxbFFHYKsAsJatMGvjIck5CkyhT/gMUyBqpA3Lkl+zC3W3uQ==} - engines: {node: '>=18'} - mlly@1.8.0: resolution: {integrity: sha512-l8D9ODSRWLe2KHJSifWGwBqpTZXIXTeo8mlKjY+E2HAakaTeNpqAyBZ8GSqLzHgw4XmHmC8whvpjJNMbFZN7/g==} @@ -10381,10 +10215,6 @@ packages: resolution: {integrity: sha512-Xf0nWe6RseziFMu+Ap9biiUbmplq6S9/p+7w7YXP/JBHhrUDDUhwa+vANyubuqfZWTveU//DYVGsDG7RKL/vEw==} engines: {node: '>=0.10.0'} - require-in-the-middle@7.5.2: - resolution: {integrity: sha512-gAZ+kLqBdHarXB64XpAe2VCjB7rIRv+mU8tfRWziHRJ5umKsIHN2tLLv6EtMw7WCdP19S0ERVMldNvxYCHnhSQ==} - engines: {node: '>=8.6.0'} - require-in-the-middle@8.0.1: resolution: {integrity: sha512-QT7FVMXfWOYFbeRBF6nu+I6tr2Tf3u0q8RIEjNob/heKY/nh7drD/k7eeMFmSQgnTtCzLDcCu/XEnpW2wk4xCQ==} engines: {node: '>=9.3.0 || >=8.10.0 <9.0.0'} @@ -13765,15 +13595,6 @@ snapshots: transitivePeerDependencies: - supports-color - '@databricks/sdk-experimental@0.15.0': - dependencies: - google-auth-library: 10.5.0 - ini: 6.0.0 - reflect-metadata: 0.2.2 - semver: 7.7.3 - transitivePeerDependencies: - - supports-color - '@databricks/sdk-experimental@0.17.0': dependencies: google-auth-library: 10.5.0 @@ -15208,10 +15029,6 @@ snapshots: dependencies: '@octokit/openapi-types': 27.0.0 - '@opentelemetry/api-logs@0.205.0': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/api-logs@0.219.0': dependencies: '@opentelemetry/api': 1.9.0 @@ -15280,10 +15097,6 @@ snapshots: '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) yaml: 2.8.2 - '@opentelemetry/context-async-hooks@2.1.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/context-async-hooks@2.8.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15293,16 +15106,6 @@ snapshots: '@opentelemetry/api': 1.9.0 '@opentelemetry/semantic-conventions': 1.38.0 - '@opentelemetry/exporter-logs-otlp-grpc@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@grpc/grpc-js': 1.14.4 - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-grpc-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-logs': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-logs-otlp-grpc@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@grpc/grpc-js': 1.14.4 @@ -15313,15 +15116,6 @@ snapshots: '@opentelemetry/otlp-transformer': 0.219.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-logs': 0.219.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-logs-otlp-http@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/api-logs': 0.205.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-logs': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-logs-otlp-http@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15331,17 +15125,6 @@ snapshots: '@opentelemetry/otlp-transformer': 0.219.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-logs': 0.219.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-logs-otlp-proto@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/api-logs': 0.205.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-logs': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-logs-otlp-proto@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15353,18 +15136,6 @@ snapshots: '@opentelemetry/sdk-logs': 0.219.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-trace-base': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-grpc@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@grpc/grpc-js': 1.14.4 - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-http': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-grpc-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-metrics': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-grpc@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@grpc/grpc-js': 1.14.4 @@ -15377,15 +15148,6 @@ snapshots: '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-metrics': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-http@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-metrics': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-http@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15395,16 +15157,6 @@ snapshots: '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-metrics': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-proto@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-http': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-metrics': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-proto@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15415,13 +15167,6 @@ snapshots: '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-metrics': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-prometheus@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-metrics': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-prometheus@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15430,17 +15175,6 @@ snapshots: '@opentelemetry/sdk-metrics': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/semantic-conventions': 1.38.0 - '@opentelemetry/exporter-trace-otlp-grpc@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@grpc/grpc-js': 1.14.4 - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-grpc-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-trace-otlp-grpc@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@grpc/grpc-js': 1.14.4 @@ -15452,15 +15186,6 @@ snapshots: '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-trace-base': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-trace-otlp-http@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-trace-otlp-http@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15470,15 +15195,6 @@ snapshots: '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-trace-base': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-trace-otlp-proto@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-trace-otlp-proto@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15488,14 +15204,6 @@ snapshots: '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-trace-base': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-zipkin@2.1.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/semantic-conventions': 1.38.0 - '@opentelemetry/exporter-zipkin@2.8.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15855,15 +15563,6 @@ snapshots: transitivePeerDependencies: - supports-color - '@opentelemetry/instrumentation@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/api-logs': 0.205.0 - import-in-the-middle: 1.15.0 - require-in-the-middle: 7.5.2 - transitivePeerDependencies: - - supports-color - '@opentelemetry/instrumentation@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15873,26 +15572,12 @@ snapshots: transitivePeerDependencies: - supports-color - '@opentelemetry/otlp-exporter-base@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/otlp-transformer': 0.219.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-grpc-exporter-base@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@grpc/grpc-js': 1.14.4 - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-exporter-base': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-grpc-exporter-base@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@grpc/grpc-js': 1.14.4 @@ -15901,17 +15586,6 @@ snapshots: '@opentelemetry/otlp-exporter-base': 0.219.0(@opentelemetry/api@1.9.0) '@opentelemetry/otlp-transformer': 0.219.0(@opentelemetry/api@1.9.0) - '@opentelemetry/otlp-transformer@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/api-logs': 0.205.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-logs': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-metrics': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 2.1.0(@opentelemetry/api@1.9.0) - protobufjs: 7.6.2 - '@opentelemetry/otlp-transformer@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15926,21 +15600,11 @@ snapshots: dependencies: '@opentelemetry/api': 1.9.0 - '@opentelemetry/propagator-b3@2.1.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/propagator-b3@2.8.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/propagator-jaeger@2.1.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/propagator-jaeger@2.8.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -15983,25 +15647,12 @@ snapshots: transitivePeerDependencies: - supports-color - '@opentelemetry/resources@2.1.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/semantic-conventions': 1.38.0 - '@opentelemetry/resources@2.8.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/semantic-conventions': 1.38.0 - '@opentelemetry/sdk-logs@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/api-logs': 0.205.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-logs@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -16010,46 +15661,12 @@ snapshots: '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/semantic-conventions': 1.38.0 - '@opentelemetry/sdk-metrics@2.1.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-metrics@2.8.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-node@0.205.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/api-logs': 0.205.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-logs-otlp-grpc': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-logs-otlp-http': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-logs-otlp-proto': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-grpc': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-http': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-metrics-otlp-proto': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-prometheus': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-trace-otlp-grpc': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-trace-otlp-http': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-trace-otlp-proto': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/exporter-zipkin': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/instrumentation': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/propagator-b3': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/propagator-jaeger': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-logs': 0.205.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-metrics': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-node': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/semantic-conventions': 1.38.0 - transitivePeerDependencies: - - supports-color - '@opentelemetry/sdk-node@0.219.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -16082,13 +15699,6 @@ snapshots: transitivePeerDependencies: - supports-color - '@opentelemetry/sdk-trace-base@2.1.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/semantic-conventions': 1.38.0 - '@opentelemetry/sdk-trace-base@2.8.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -16096,13 +15706,6 @@ snapshots: '@opentelemetry/resources': 2.8.0(@opentelemetry/api@1.9.0) '@opentelemetry/semantic-conventions': 1.38.0 - '@opentelemetry/sdk-trace-node@2.1.0(@opentelemetry/api@1.9.0)': - dependencies: - '@opentelemetry/api': 1.9.0 - '@opentelemetry/context-async-hooks': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/core': 2.8.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 2.1.0(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-node@2.8.0(@opentelemetry/api@1.9.0)': dependencies: '@opentelemetry/api': 1.9.0 @@ -20053,8 +19656,6 @@ snapshots: fast-levenshtein@2.0.6: {} - fast-safe-stringify@2.1.1: {} - fast-uri@3.1.0: {} fastq@1.19.1: @@ -20969,13 +20570,6 @@ snapshots: parent-module: 1.0.1 resolve-from: 4.0.0 - import-in-the-middle@1.15.0: - dependencies: - acorn: 8.15.0 - acorn-import-attributes: 1.9.5(acorn@8.15.0) - cjs-module-lexer: 1.4.3 - module-details-from-path: 1.0.4 - import-in-the-middle@3.0.1: dependencies: acorn: 8.15.0 @@ -21007,8 +20601,6 @@ snapshots: ini@4.1.1: {} - ini@5.0.0: {} - ini@6.0.0: {} inline-style-parser@0.2.7: {} @@ -22299,17 +21891,6 @@ snapshots: mkdirp@3.0.1: {} - mlflow-tracing@0.1.3: - dependencies: - '@databricks/sdk-experimental': 0.15.0 - '@opentelemetry/api': 1.9.0 - '@opentelemetry/sdk-node': 0.205.0(@opentelemetry/api@1.9.0) - bignumber.js: 9.3.1 - fast-safe-stringify: 2.1.1 - ini: 5.0.0 - transitivePeerDependencies: - - supports-color - mlly@1.8.0: dependencies: acorn: 8.15.0 @@ -23828,14 +23409,6 @@ snapshots: require-from-string@2.0.2: {} - require-in-the-middle@7.5.2: - dependencies: - debug: 4.4.3 - module-details-from-path: 1.0.4 - resolve: 1.22.10 - transitivePeerDependencies: - - supports-color - require-in-the-middle@8.0.1: dependencies: debug: 4.4.3 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/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/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 503ee7983..241751867 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", @@ -18,25 +51,11 @@ "permission": "CAN_QUERY", "fields": { "name": { - "env": "DATABRICKS_SERVING_ENDPOINT_NAME", + "env": "DATABRICKS_AGENT_SERVING_ENDPOINT_NAME", "description": "Default LLM serving endpoint name", "origin": "user" } } - }, - { - "type": "experiment", - "alias": "MLflow experiment for agent traces", - "resourceKey": "agents-mlflow-experiment", - "description": "When bound, agent turns and tool calls are traced to this MLflow experiment via OpenTelemetry. Tracing is a no-op when unset.", - "permission": "CAN_EDIT", - "fields": { - "experimentId": { - "env": "MLFLOW_EXPERIMENT_ID", - "description": "MLflow experiment id traces are logged to", - "origin": "user" - } - } } ] }, diff --git a/template/client/src/pages/agents/AgentChat.tsx b/template/client/src/pages/agents/AgentChat.tsx index 5d1e5758d..38b1c69a9 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,19 @@ export function AgentChat() { })} + {mlflowTraceId && ( +
+ + {mlflowTraceId} + + {mlflowTraceUrl && ( + + 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,35 @@ 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. 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 { + 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/agent-template-policy.test.ts b/tools/tests/agent-template-policy.test.ts new file mode 100644 index 000000000..7cb30384e --- /dev/null +++ b/tools/tests/agent-template-policy.test.ts @@ -0,0 +1,214 @@ +import { execFileSync } from "node:child_process"; +import { + existsSync, + mkdirSync, + mkdtempSync, + readdirSync, + readFileSync, + 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(root, ".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( + 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(searchRoot, 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)); +} + +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("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( + 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.60.0"); + expect( + packageJson.dependencies["@mlflow/core"], + `${candidate.name} must keep AppKit as its sole tracing provider`, + ).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")); + } + }); +}); diff --git a/tools/tests/generate-app-templates.test.ts b/tools/tests/generate-app-templates.test.ts new file mode 100644 index 000000000..a6d1b7654 --- /dev/null +++ b/tools/tests/generate-app-templates.test.ts @@ -0,0 +1,155 @@ +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(() => { + const compatibleCli = "/tmp/databricks-cli-1.11.0/databricks"; + execFileSync("pnpm", ["generate:app-templates"], { + cwd: root, + env: { + ...process.env, + APP_TEMPLATES_OUTPUT_DIR: outputDir, + ...(process.env.DATABRICKS_CLI + ? {} + : existsSync(compatibleCli) + ? { DATABRICKS_CLI: compatibleCli } + : {}), + }, + 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: { helper }", + ); + }, + ); + + 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) => { + 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.requiredByTemplate).toBe(true); + 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.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); + 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"); + }, + ); + + 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"); + 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 8c2893b01..a44d3cf9d 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -58,6 +58,23 @@ export default defineConfig({ environment: "node", }, }, + { + test: { + name: "tools", + root: ".", + include: ["tools/tests/**/*.test.ts"], + environment: "node", + }, + }, + { + plugins: [tsconfigPaths()], + test: { + name: "dev-playground-server", + root: "./apps/dev-playground", + include: ["server/tests/**/*.test.ts"], + environment: "node", + }, + }, ], }, });