From 38587769b96f5738a037748e6cce363345d25187 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:40:38 +0000 Subject: [PATCH 1/4] ref: Remove `AsyncLocalStorage.enterWith()` usage --- js/src/global-instrumentation-hooks.test.ts | 2 - js/src/global-instrumentation-hooks.ts | 1 - .../auto-instrumentation-suppression.test.ts | 38 +++---- .../auto-instrumentation-suppression.ts | 24 +---- .../plugins/ai-sdk-plugin.test.ts | 20 +++- .../instrumentation/plugins/ai-sdk-plugin.ts | 44 ++++----- .../plugins/anthropic-sessions-plugin.test.ts | 1 - .../plugins/claude-agent-sdk-channels.ts | 2 +- .../claude-agent-sdk-local-tool-context.ts | 55 ++++------- .../claude-agent-sdk-local-tool-spans.ts | 2 +- .../plugins/claude-agent-sdk-plugin.test.ts | 48 ++++++--- .../plugins/claude-agent-sdk-plugin.ts | 62 ++++++++++-- .../plugins/google-genai-plugin.test.ts | 3 - .../plugins/pi-coding-agent-plugin.test.ts | 13 ++- .../plugins/pi-coding-agent-plugin.ts | 98 ++++++++++++------- .../plugins/strands-agent-sdk-plugin.test.ts | 3 - js/src/instrumentation/registry.test.ts | 1 - js/src/isomorph.ts | 1 - js/src/vendor-sdk-types/pi-coding-agent.ts | 1 + js/src/wrappers/ai-sdk/telemetry.ts | 55 +++-------- .../claude-agent-sdk/claude-agent-sdk.ts | 10 +- js/src/wrappers/vitest/context-manager.ts | 6 -- js/src/wrappers/vitest/wrapper.ts | 8 +- 23 files changed, 274 insertions(+), 224 deletions(-) diff --git a/js/src/global-instrumentation-hooks.test.ts b/js/src/global-instrumentation-hooks.test.ts index f46573e86..687695850 100644 --- a/js/src/global-instrumentation-hooks.test.ts +++ b/js/src/global-instrumentation-hooks.test.ts @@ -446,7 +446,6 @@ describe("global instrumentation hooks", () => { const storeError = new Error("store failed"); const brokenStore = { - enterWith() {}, getStore() { return undefined; }, @@ -475,7 +474,6 @@ describe("global instrumentation hooks", () => { let callback: (() => unknown) | undefined; channel.start.bindStore({ - enterWith() {}, getStore() { return undefined; }, diff --git a/js/src/global-instrumentation-hooks.ts b/js/src/global-instrumentation-hooks.ts index 5dbefc85c..119d1bfd5 100644 --- a/js/src/global-instrumentation-hooks.ts +++ b/js/src/global-instrumentation-hooks.ts @@ -19,7 +19,6 @@ const hookBrand = Symbol.for(GLOBAL_INSTRUMENTATION_HOOK_BRAND); const invocationHookBrand = Symbol.for(GLOBAL_INVOCATION_HOOK_BRAND); export interface GlobalHookAsyncLocalStorage { - enterWith(store: T): void; run(store: T | undefined, callback: () => R): R; getStore(): T | undefined; } diff --git a/js/src/instrumentation/auto-instrumentation-suppression.test.ts b/js/src/instrumentation/auto-instrumentation-suppression.test.ts index d4561deb7..e1672d2ea 100644 --- a/js/src/instrumentation/auto-instrumentation-suppression.test.ts +++ b/js/src/instrumentation/auto-instrumentation-suppression.test.ts @@ -1,8 +1,8 @@ import { beforeAll, describe, expect, it } from "vitest"; import { configureNode } from "../node/config"; import { - enterAutoInstrumentationAllowed, isAutoInstrumentationSuppressed, + runWithAutoInstrumentationAllowed, runWithAutoInstrumentationSuppressed, } from "./auto-instrumentation-suppression"; @@ -19,32 +19,34 @@ describe("auto instrumentation suppression context", () => { await Promise.resolve(); expect(isAutoInstrumentationSuppressed()).toBe(true); - const restoreToolContext = enterAutoInstrumentationAllowed(); - expect(isAutoInstrumentationSuppressed()).toBe(false); + await runWithAutoInstrumentationAllowed(async () => { + expect(isAutoInstrumentationSuppressed()).toBe(false); - await runWithAutoInstrumentationSuppressed(async () => { - expect(isAutoInstrumentationSuppressed()).toBe(true); - await Promise.resolve(); - expect(isAutoInstrumentationSuppressed()).toBe(true); - }); + await runWithAutoInstrumentationSuppressed(async () => { + expect(isAutoInstrumentationSuppressed()).toBe(true); + await Promise.resolve(); + expect(isAutoInstrumentationSuppressed()).toBe(true); + }); - expect(isAutoInstrumentationSuppressed()).toBe(false); - restoreToolContext(); + expect(isAutoInstrumentationSuppressed()).toBe(false); + }); expect(isAutoInstrumentationSuppressed()).toBe(true); }); expect(isAutoInstrumentationSuppressed()).toBe(false); }); - it("keeps instrumentation allowed until every active allow frame exits", async () => { + it("restores nested allow frames at each callback boundary", async () => { await runWithAutoInstrumentationSuppressed(async () => { - const restoreFirstTool = enterAutoInstrumentationAllowed(); - const restoreSecondTool = enterAutoInstrumentationAllowed(); - - expect(isAutoInstrumentationSuppressed()).toBe(false); - restoreFirstTool(); - expect(isAutoInstrumentationSuppressed()).toBe(false); - restoreSecondTool(); + await runWithAutoInstrumentationAllowed(async () => { + expect(isAutoInstrumentationSuppressed()).toBe(false); + await runWithAutoInstrumentationAllowed(async () => { + expect(isAutoInstrumentationSuppressed()).toBe(false); + await Promise.resolve(); + expect(isAutoInstrumentationSuppressed()).toBe(false); + }); + expect(isAutoInstrumentationSuppressed()).toBe(false); + }); expect(isAutoInstrumentationSuppressed()).toBe(true); }); }); diff --git a/js/src/instrumentation/auto-instrumentation-suppression.ts b/js/src/instrumentation/auto-instrumentation-suppression.ts index 2b31044ef..37fc4b4d5 100644 --- a/js/src/instrumentation/auto-instrumentation-suppression.ts +++ b/js/src/instrumentation/auto-instrumentation-suppression.ts @@ -4,7 +4,6 @@ import iso, { } from "../isomorph"; type AutoInstrumentationSuppressionFrame = { - id: symbol; mode: "allow" | "suppress"; }; @@ -34,7 +33,6 @@ export function isAutoInstrumentationSuppressed(): boolean { export function runWithAutoInstrumentationSuppressed(callback: () => R): R { const frame = { - id: Symbol("braintrust.auto-instrumentation-suppress"), mode: "suppress" as const, }; return suppressionStore().run( @@ -56,7 +54,6 @@ export function bindAutoInstrumentationSuppressionToStart( frames: [ ...currentFrames(), { - id: Symbol("braintrust.auto-instrumentation-suppress"), mode: "suppress" as const, }, ], @@ -67,23 +64,12 @@ export function bindAutoInstrumentationSuppressionToStart( }; } -export function enterAutoInstrumentationAllowed(): () => void { +export function runWithAutoInstrumentationAllowed(callback: () => R): R { const frame = { - id: Symbol("braintrust.auto-instrumentation-allow"), mode: "allow" as const, }; - // TODO(luca): Replace ALS.enterWith() with ALS.run() - // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above. - suppressionStore().enterWith({ - frames: [...currentFrames(), frame], - }); - - return () => { - const frames = currentFrames().filter( - (candidate) => candidate.id !== frame.id, - ); - // TODO(luca): Replace ALS.enterWith() with ALS.run() - // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above. - suppressionStore().enterWith(frames.length > 0 ? { frames } : undefined); - }; + return suppressionStore().run( + { frames: [...currentFrames(), frame] }, + callback, + ); } diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts index fcc538deb..0cfd405de 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts @@ -38,6 +38,7 @@ const mockNewTracingChannel = iso.newTracingChannel as ReturnType; type MockTracingChannel = { handlers: any[]; hasSubscribers: boolean; + intercept: ReturnType; subscribe: ReturnType; unsubscribe: ReturnType; }; @@ -66,6 +67,23 @@ describe("AISDKPlugin", () => { const channel: MockTracingChannel = { handlers: [], hasSubscribers: false, + intercept: vi.fn((interceptor: any) => { + const handlers = { + end: (event: any) => + interceptor( + () => event.result, + event.self, + event.arguments ?? [], + {}, + ), + }; + channel.handlers.push(handlers); + return vi.fn(() => { + channel.handlers = channel.handlers.filter( + (candidate) => candidate !== handlers, + ); + }); + }), subscribe: vi.fn((handlers: any) => { channel.handlers.push(handlers); channel.hasSubscribers = true; @@ -243,7 +261,7 @@ describe("AISDKPlugin", () => { const channel = mockChannels.get( "orchestrion:ai:createTelemetryDispatcher", ); - expect(channel?.subscribe).toHaveBeenCalledTimes(1); + expect(channel?.intercept).toHaveBeenCalledTimes(1); channel?.handlers[0]?.end({ arguments: [{ telemetry: {} }], diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index e6cbf92f4..d5526da2f 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -193,7 +193,7 @@ export class AISDKPlugin extends BasePlugin { const denyOutputPaths = this.config.denyOutputPaths || DEFAULT_DENY_OUTPUT_PATHS; - this.unsubscribers.push(subscribeToAISDKV7TelemetryDispatcher()); + this.unsubscribers.push(interceptAISDKV7TelemetryDispatcher()); this.unsubscribers.push(subscribeToHarnessAgentCreateSession()); this.unsubscribers.push( subscribeToHarnessContinuation( @@ -839,31 +839,29 @@ function subscribeToHarnessContinuation( }; } -function subscribeToAISDKV7TelemetryDispatcher(): () => void { - const channel = aiSDKChannels.v7CreateTelemetryDispatcher.tracingChannel(); +function interceptAISDKV7TelemetryDispatcher(): () => void { const telemetry = braintrustAISDKTelemetry(); - const handlers: IsoChannelHandlers< - ChannelMessage - > = { - end: (event) => { - const telemetryOptions = event.arguments?.[0]?.telemetry; - if (telemetryOptions?.isEnabled === false) { - return; + return aiSDKChannels.v7CreateTelemetryDispatcher.intercept( + (target, thisArg, args) => { + const dispatcher = Reflect.apply(target, thisArg, args); + const telemetryOptions = args[0]?.telemetry; + if (telemetryOptions?.isEnabled !== false) { + try { + patchAISDKV7TelemetryDispatcher( + dispatcher, + telemetry, + telemetryOptions, + ); + } catch (error) { + debugLogger.error( + "Error instrumenting AI SDK v7 telemetry dispatcher:", + error, + ); + } } - - patchAISDKV7TelemetryDispatcher( - event.result, - telemetry, - telemetryOptions, - ); + return dispatcher; }, - }; - - channel.subscribe(handlers); - - return () => { - channel.unsubscribe(handlers); - }; + ); } function patchAISDKV7TelemetryDispatcher( diff --git a/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts b/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts index 1d4cb48ac..01cc32592 100644 --- a/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts +++ b/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts @@ -9,7 +9,6 @@ vi.mock("../../isomorph", () => ({ default: { getEnv: vi.fn(), newAsyncLocalStorage: vi.fn(() => ({ - enterWith: vi.fn(), getStore: vi.fn(() => undefined), run: vi.fn((_store: unknown, callback: () => unknown) => callback()), })), diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts b/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts index 959065035..c58e23144 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts @@ -11,7 +11,7 @@ export const claudeAgentSDKChannels = defineChannels( query: channel< [ClaudeAgentSDKQueryParams], AsyncIterable, - Record, + Record, ClaudeAgentSDKMessage >({ channelName: "query", diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts index 3bf15777a..d98223b8b 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts @@ -11,7 +11,6 @@ const LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED = Symbol.for( ); type AsyncLocalStorageLike = { - enterWith: (store: T) => void; getStore: () => T | undefined; run: (store: T, callback: () => R) => R; }; @@ -29,9 +28,6 @@ function createLocalToolContextStore(): AsyncLocalStorageLike(); export function createClaudeLocalToolContext(): ClaudeAgentSDKLocalToolContext { return {}; } -function runWithClaudeLocalToolContext( +export function runWithClaudeLocalToolContext( callback: () => R, context?: ClaudeAgentSDKLocalToolContext, ): R { @@ -64,39 +63,25 @@ function runWithClaudeLocalToolContext( ); } -function ensureClaudeLocalToolContext(): - | ClaudeAgentSDKLocalToolContext - | undefined { - const existing = localToolContextStore.getStore(); - if (existing) { - return existing; - } - - const created: ClaudeAgentSDKLocalToolContext = {}; - // TODO(luca): Replace ALS.enterWith() with ALS.run() - // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above. - localToolContextStore.enterWith(created); - return created; -} - -export function setClaudeLocalToolParentResolver( +export function registerClaudeLocalToolParentResolver( + toolUseId: string, resolver: LocalToolParentResolver, ): void { - fallbackLocalToolParentResolver = resolver; - const context = ensureClaudeLocalToolContext(); - if (!context) { - return; - } - context.resolveLocalToolParent = resolver; + localToolParentResolversByToolUseId.set(toolUseId, resolver); } -export function getClaudeLocalToolParentResolver(): - | LocalToolParentResolver - | undefined { - return ( - localToolContextStore.getStore()?.resolveLocalToolParent ?? - fallbackLocalToolParentResolver - ); +export function getClaudeLocalToolParentResolver( + toolUseId?: string, +): LocalToolParentResolver | undefined { + const currentResolver = + localToolContextStore.getStore()?.resolveLocalToolParent; + if (!toolUseId) { + return currentResolver; + } + + const registeredResolver = localToolParentResolversByToolUseId.get(toolUseId); + localToolParentResolversByToolUseId.delete(toolUseId); + return currentResolver ?? registeredResolver; } function isAsyncIterable(value: unknown): value is AsyncIterable { diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts index 1a1c7db88..843309851 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts @@ -53,7 +53,7 @@ export function wrapLocalClaudeToolHandler( ? `mcp__${metadata.serverName}__${metadata.toolName}` : metadata.toolName; const toolUseId = getToolUseIdFromExtra(handlerArgs[1]); - const localToolParentResolver = getClaudeLocalToolParentResolver(); + const localToolParentResolver = getClaudeLocalToolParentResolver(toolUseId); const spanName = metadata.serverName ? `tool: ${metadata.serverName}/${metadata.toolName}` : `tool: ${metadata.toolName}`; diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts index f1d446bcd..42f9da7a6 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts @@ -123,11 +123,16 @@ describe("ClaudeAgentSDKPlugin", () => { let plugin: ClaudeAgentSDKPlugin; let mockChannel: any; let mockUnsubscribe: any; + let queryInterceptor: any; beforeEach(() => { streamPatcherMock.options = undefined; mockUnsubscribe = vi.fn(); mockChannel = { + intercept: vi.fn((interceptor) => { + queryInterceptor = interceptor; + return mockUnsubscribe; + }), subscribe: vi.fn(), unsubscribe: mockUnsubscribe, hasSubscribers: false, @@ -149,21 +154,15 @@ describe("ClaudeAgentSDKPlugin", () => { expect(mockNewTracingChannel).toHaveBeenCalledWith( "orchestrion:@anthropic-ai/claude-agent-sdk:query", ); - expect(mockChannel.subscribe).toHaveBeenCalledTimes(1); - expect(mockChannel.subscribe).toHaveBeenCalledWith( - expect.objectContaining({ - start: expect.any(Function), - end: expect.any(Function), - error: expect.any(Function), - }), - ); + expect(mockChannel.intercept).toHaveBeenCalledTimes(1); + expect(mockChannel.intercept).toHaveBeenCalledWith(expect.any(Function)); }); it("should not subscribe twice if already enabled", () => { plugin.enable(); plugin.enable(); - expect(mockChannel.subscribe).toHaveBeenCalledTimes(1); + expect(mockChannel.intercept).toHaveBeenCalledTimes(1); }); it("should store unsubscribe function", () => { @@ -202,7 +201,34 @@ describe("ClaudeAgentSDKPlugin", () => { beforeEach(() => { plugin.enable(); - handlers = mockChannel.subscribe.mock.calls[0][0]; + handlers = { + start: (event: any) => + queryInterceptor( + () => ({ + async *[Symbol.asyncIterator]() { + // Keep the query span open so tests can drive stream callbacks. + }, + }), + event.self, + event.arguments ?? [], + {}, + ), + end: () => undefined, + error: (event: any) => { + try { + queryInterceptor( + () => { + throw event.error; + }, + event.self, + event.arguments ?? [], + {}, + ); + } catch { + // The invocation interceptor preserves the target's exception. + } + }, + }; }); describe("start handler", () => { @@ -752,7 +778,7 @@ describe("ClaudeAgentSDKPlugin", () => { plugin.disable(); plugin.enable(); - expect(mockChannel.subscribe).toHaveBeenCalledTimes(2); + expect(mockChannel.intercept).toHaveBeenCalledTimes(2); }); it("should properly clean up on multiple enable/disable cycles", () => { diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts index 96f52c9b6..e9414940b 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts @@ -2,6 +2,7 @@ import { BasePlugin } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import type { IsoChannelHandlers } from "../../isomorph"; +import { debugLogger } from "../../debug-logger"; import { startSpan as startBaseSpan } from "../../logger"; import type { Span } from "../../logger"; import { @@ -25,7 +26,8 @@ import { import { bindClaudeLocalToolContextToAsyncIterable, createClaudeLocalToolContext, - setClaudeLocalToolParentResolver, + registerClaudeLocalToolParentResolver, + runWithClaudeLocalToolContext, type ClaudeAgentSDKLocalToolContext, } from "./claude-agent-sdk-local-tool-context"; import type { @@ -627,6 +629,7 @@ function createToolTracingHooks( (isLocalToolUse(input.tool_name, mcpServers) || localToolHookNames.has(input.tool_name)) ) { + registerClaudeLocalToolParentResolver(toolUseID, resolveParentSpan); return {}; } @@ -1591,7 +1594,6 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { } private subscribeToQuery(): void { - const channel = claudeAgentSDKChannels.query.tracingChannel(); const spans = new WeakMap(); const handlers: IsoChannelHandlers< @@ -1728,7 +1730,6 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { }; localToolContext.resolveLocalToolParent = resolveToolUseParentSpan; - setClaudeLocalToolParentResolver(resolveToolUseParentSpan); const optionsWithHooks = injectTracingHooks( options, resolveToolUseParentSpan, @@ -1853,9 +1854,56 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { }, }; - channel.subscribe(handlers); - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); + this.unsubscribers.push( + claudeAgentSDKChannels.query.intercept( + (target, thisArg, args, additional) => { + const event: ChannelMessage = { + ...additional, + arguments: args, + }; + try { + handlers.start?.(event, claudeAgentSDKChannels.query.channelName); + } catch (error) { + debugLogger.error( + "Error starting Claude Agent SDK instrumentation:", + error, + ); + } + + const state = spans.get(event); + const invokeTarget = () => Reflect.apply(target, thisArg, args); + try { + const result = state + ? runWithClaudeLocalToolContext( + invokeTarget, + state.localToolContext, + ) + : invokeTarget(); + event.result = result; + try { + handlers.end?.(event, claudeAgentSDKChannels.query.channelName); + } catch (error) { + debugLogger.error( + "Error finalizing Claude Agent SDK instrumentation:", + error, + ); + } + return result; + } catch (error) { + event.error = + error instanceof Error ? error : new Error(String(error)); + try { + handlers.error?.(event, claudeAgentSDKChannels.query.channelName); + } catch (instrumentationError) { + debugLogger.error( + "Error handling Claude Agent SDK instrumentation failure:", + instrumentationError, + ); + } + throw error; + } + }, + ), + ); } } diff --git a/js/src/instrumentation/plugins/google-genai-plugin.test.ts b/js/src/instrumentation/plugins/google-genai-plugin.test.ts index f14690e73..4985e0a65 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.test.ts +++ b/js/src/instrumentation/plugins/google-genai-plugin.test.ts @@ -6,9 +6,6 @@ vi.mock("../../isomorph", () => ({ newAsyncLocalStorage: vi.fn(() => { let current: unknown; return { - enterWith: vi.fn((store: unknown) => { - current = store; - }), getStore: vi.fn(() => current), run: vi.fn((store: unknown, callback: () => unknown) => { const previous = current; diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts index 433361974..346901c78 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts +++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts @@ -99,7 +99,8 @@ describe("PiCodingAgentPlugin", () => { return makeStream(finalMessage); }); const agent = makeAgent(originalStreamFn); - agent.state.tools = [bashTool()]; + const tool = bashTool(); + agent.state.tools = [tool]; const session = makeSession(agent); const context = { systemPrompt: "system", @@ -110,7 +111,7 @@ describe("PiCodingAgentPlugin", () => { timestamp: 1, }, ], - tools: [bashTool()], + tools: [tool], }; await interceptor( @@ -127,6 +128,8 @@ describe("PiCodingAgentPlugin", () => { toolName: "bash", type: "tool_execution_start", }); + await tool.execute?.("tool-1", { command: "printf pi_tool_ok" }); + expect(isAutoInstrumentationSuppressed()).toBe(true); await this.agent.emit({ isError: false, result: { stdout: "pi_tool_ok" }, @@ -459,6 +462,12 @@ function bashTool() { return { description: "Run a shell command.", name: "bash", + execute: vi.fn(async (..._args: unknown[]) => { + expect(isAutoInstrumentationSuppressed()).toBe(false); + await Promise.resolve(); + expect(isAutoInstrumentationSuppressed()).toBe(false); + return { stdout: "pi_tool_ok" }; + }), parameters: { type: "object", properties: { diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts index da6c50b30..6d4e442c7 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts +++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts @@ -12,7 +12,7 @@ import { getCurrentUnixTimestamp } from "../../util"; import { SpanTypeAttribute, isObject } from "../../../util/index"; import { processInputAttachments } from "../../wrappers/attachment-utils"; import { - enterAutoInstrumentationAllowed, + runWithAutoInstrumentationAllowed, runWithAutoInstrumentationSuppressed, } from "../auto-instrumentation-suppression"; import { piCodingAgentChannels } from "./pi-coding-agent-channels"; @@ -62,7 +62,6 @@ type PiLlmSpanState = { }; type PiToolSpanState = { - restoreAutoInstrumentation?: () => void; span: Span; }; @@ -73,6 +72,9 @@ type PiAgentPatchState = { const piAgentPatchStates = new WeakMap(); const piAgentEventSubscriptions = new WeakSet(); +const PI_TOOL_EXECUTE_WRAPPED = Symbol.for( + "braintrust.pi_coding_agent.tool_execute_wrapped", +); let piPromptContextStore: | IsoAsyncLocalStorage | undefined; @@ -167,6 +169,7 @@ function startPiPromptRun( return undefined; } installPiAgentInstrumentation(agent); + wrapPiToolExecutors(agent.state?.tools); const metadata = { ...extractSessionMetadata(session), @@ -299,6 +302,7 @@ function makeInstrumentedStreamFn( return invokeOriginal(); } + wrapPiToolExecutors(context.tools); const llmState = await startPiLlmSpan(state, model, context, options); try { const stream = await runWithAutoInstrumentationSuppressed(invokeOriginal); @@ -310,6 +314,41 @@ function makeInstrumentedStreamFn( }; } +function wrapPiToolExecutors(tools: PiTool[] | undefined): void { + if (!tools) { + return; + } + + for (const tool of tools) { + try { + const execute = tool.execute; + if ( + typeof execute !== "function" || + (execute as typeof execute & { [PI_TOOL_EXECUTE_WRAPPED]?: boolean })[ + PI_TOOL_EXECUTE_WRAPPED + ] + ) { + continue; + } + + const wrappedExecute = function (this: unknown, ...args: unknown[]) { + return runWithAutoInstrumentationAllowed(() => + Reflect.apply(execute, this, args), + ); + }; + Object.defineProperty(wrappedExecute, PI_TOOL_EXECUTE_WRAPPED, { + configurable: false, + enumerable: false, + value: true, + writable: false, + }); + tool.execute = wrappedExecute; + } catch (error) { + logInstrumentationError("Pi Coding Agent tool wrapping", error); + } + } +} + async function startPiLlmSpan( state: PiPromptState, model: PiModel, @@ -528,35 +567,26 @@ async function startPiToolSpan( return; } - const restoreAutoInstrumentation = enterAutoInstrumentationAllowed(); const metadata = { "gen_ai.tool.call.id": event.toolCallId, "gen_ai.tool.name": event.toolName, "pi_coding_agent.tool.name": event.toolName, }; - try { - const span = startBaseSpan( - withSpanInstrumentationName( - { - event: { - input: event.args, - metadata, - }, - name: event.toolName || "tool", - parent: await state.span.export(), - spanAttributes: { type: SpanTypeAttribute.TOOL }, + const span = startBaseSpan( + withSpanInstrumentationName( + { + event: { + input: event.args, + metadata, }, - INSTRUMENTATION_NAMES.PI_CODING_AGENT, - ), - ); - state.activeToolSpans.set(event.toolCallId, { - restoreAutoInstrumentation, - span, - }); - } catch (error) { - restoreAutoInstrumentation(); - throw error; - } + name: event.toolName || "tool", + parent: await state.span.export(), + spanAttributes: { type: SpanTypeAttribute.TOOL }, + }, + INSTRUMENTATION_NAMES.PI_CODING_AGENT, + ), + ); + state.activeToolSpans.set(event.toolCallId, { span }); } function finishPiToolSpan( @@ -582,11 +612,7 @@ function finishPiToolSpan( output: event.result, }); } finally { - try { - toolState.span.end(); - } finally { - toolState.restoreAutoInstrumentation?.(); - } + toolState.span.end(); } } @@ -671,14 +697,10 @@ function finishPiLlmSpan( function finishOpenToolSpans(state: PiPromptState, error?: unknown): void { for (const [, toolState] of state.activeToolSpans) { - try { - safeLog(toolState.span, { - error: error ? toLoggedError(error) : "Pi tool did not complete", - }); - toolState.span.end(); - } finally { - toolState.restoreAutoInstrumentation?.(); - } + safeLog(toolState.span, { + error: error ? toLoggedError(error) : "Pi tool did not complete", + }); + toolState.span.end(); } state.activeToolSpans.clear(); } diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts index 54340d1a3..d52d0ce50 100644 --- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts @@ -12,9 +12,6 @@ const { mockNewAsyncLocalStorage: vi.fn(() => { let current: unknown; return { - enterWith: vi.fn((store: unknown) => { - current = store; - }), getStore: vi.fn(() => current), run: vi.fn((store: unknown, callback: () => unknown) => { const previous = current; diff --git a/js/src/instrumentation/registry.test.ts b/js/src/instrumentation/registry.test.ts index fdd49a5c2..9d01c94ef 100644 --- a/js/src/instrumentation/registry.test.ts +++ b/js/src/instrumentation/registry.test.ts @@ -5,7 +5,6 @@ vi.mock("../isomorph", () => ({ default: { newTracingChannel: vi.fn(), newAsyncLocalStorage: vi.fn(() => ({ - enterWith: vi.fn(), getStore: vi.fn(() => undefined), run: vi.fn((_store: unknown, callback: () => unknown) => callback()), })), diff --git a/js/src/isomorph.ts b/js/src/isomorph.ts index d23ddcc7e..a6adf3b48 100644 --- a/js/src/isomorph.ts +++ b/js/src/isomorph.ts @@ -21,7 +21,6 @@ export type IsoAsyncLocalStorage = GlobalHookAsyncLocalStorage; class DefaultAsyncLocalStorage implements IsoAsyncLocalStorage { constructor() {} - enterWith(_: T): void {} run(_: T | undefined, callback: () => R): R { return callback(); } diff --git a/js/src/vendor-sdk-types/pi-coding-agent.ts b/js/src/vendor-sdk-types/pi-coding-agent.ts index 13df2a1b9..0c2810da5 100644 --- a/js/src/vendor-sdk-types/pi-coding-agent.ts +++ b/js/src/vendor-sdk-types/pi-coding-agent.ts @@ -139,6 +139,7 @@ export interface PiContext { export interface PiTool { name: string; description?: string; + execute?: (this: unknown, ...args: unknown[]) => unknown; parameters?: unknown; [key: string]: unknown; } diff --git a/js/src/wrappers/ai-sdk/telemetry.ts b/js/src/wrappers/ai-sdk/telemetry.ts index 1fb2c9812..5bb41f885 100644 --- a/js/src/wrappers/ai-sdk/telemetry.ts +++ b/js/src/wrappers/ai-sdk/telemetry.ts @@ -28,7 +28,6 @@ import type { AISDKResult, AISDKRerankResult, } from "../../vendor-sdk-types/ai-sdk"; -import iso from "../../isomorph"; import type { AISDKV7LanguageModelCallStartEvent, AISDKV7OperationEvent, @@ -74,9 +73,6 @@ type EmbedSpanState = CallSpanState & { export function braintrustAISDKTelemetry(): any { const operations = new Map(); const operationKeysByCallId = new Map(); - const workflowOperationKeyStore = iso.newAsyncLocalStorage< - string | undefined - >(); const modelSpans = new Map(); const objectSpans = new Map(); const embedSpans = new Map(); @@ -135,13 +131,6 @@ export function braintrustAISDKTelemetry(): any { } operations.delete(operationKey); - if (workflowOperationKeyStore.getStore() === operationKey) { - // TODO(luca): Replace ALS.enterWith() with ALS.run() once direct - // telemetry can wrap the full WorkflowAgent callback lifecycle. - // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above. - workflowOperationKeyStore.enterWith(undefined); - } - const keys = operationKeysByCallId.get(state.callId); if (!keys) { return; @@ -208,16 +197,12 @@ export function braintrustAISDKTelemetry(): any { } } - const workflowOperationKey = workflowOperationKeyStore.getStore(); - if (workflowOperationKey && keys.includes(workflowOperationKey)) { - return workflowOperationKey; - } - - if (callId === "workflow-agent") { - return undefined; - } - - return mode === "finish" ? keys[0] : keys[keys.length - 1]; + // Dispatcher instrumentation stamps an explicit key on every callback. + // Direct registerTelemetry() usage has no callback boundary to carry one, + // so overlapping operations with the same callId are best-effort. + return callId === "workflow-agent" || mode === "active" + ? keys[keys.length - 1] + : keys[0]; }; const operationKeyFromEvent = ( @@ -237,20 +222,19 @@ export function braintrustAISDKTelemetry(): any { return operationKey; } - const workflowOperationKey = workflowOperationKeyStore.getStore(); - if (workflowOperationKey && operations.has(workflowOperationKey)) { - return workflowOperationKey; + // Some direct WorkflowAgent telemetry callbacks use a child callId + // instead of the operation's shared `workflow-agent` callId. Without + // a dispatcher key, route these to the newest active workflow as a + // deterministic best-effort fallback. + const workflowAgentKeys = operationKeysByCallId.get("workflow-agent"); + if (workflowAgentKeys?.length) { + return workflowAgentKeys[workflowAgentKeys.length - 1]; } return callId === "workflow-agent" ? undefined : callId; } } - const workflowOperationKey = workflowOperationKeyStore.getStore(); - if (workflowOperationKey && operations.has(workflowOperationKey)) { - return workflowOperationKey; - } - const wrapperSpan = currentWorkflowAgentWrapperSpan(); if (wrapperSpan?.spanId) { for (const [operationKey, state] of operations) { @@ -266,8 +250,8 @@ export function braintrustAISDKTelemetry(): any { // WorkflowAgent uses this callId on the operation, but omits it from // tool start/end callbacks in @ai-sdk/workflow@1.0.x. const workflowAgentKeys = operationKeysByCallId.get("workflow-agent"); - if (workflowAgentKeys?.length === 1) { - return workflowAgentKeys[0]; + if (workflowAgentKeys?.length) { + return workflowAgentKeys[workflowAgentKeys.length - 1]; } if (operations.size === 1) { @@ -521,15 +505,6 @@ export function braintrustAISDKTelemetry(): any { return; } - if (workflowAgent) { - // Direct registerTelemetry() calls do not receive the hidden - // dispatcher operation key used by auto-instrumentation. - // TODO(luca): Replace ALS.enterWith() with ALS.run() once direct - // telemetry can wrap the full WorkflowAgent callback lifecycle. - // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above. - workflowOperationKeyStore.enterWith(operationKey); - } - let metadata = metadataFromEvent(event); const logPayload: { input?: unknown; diff --git a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts index db24f0bcf..ddc6d678e 100644 --- a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts +++ b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts @@ -54,11 +54,11 @@ function wrapClaudeAgentQuery( thisArg === proxy || thisArg === undefined ? (defaultThis ?? thisArg) : thisArg; - return claudeAgentSDKChannels.query.traceSync( - () => Reflect.apply(target, invocationTarget, [wrappedParams]), - // The channel carries no extra context fields, but the generated - // StartOf<> type for Record is overly strict here. - { arguments: [wrappedParams] } as never, + return claudeAgentSDKChannels.query.invoke( + target, + invocationTarget, + [wrappedParams], + {}, ); }, }); diff --git a/js/src/wrappers/vitest/context-manager.ts b/js/src/wrappers/vitest/context-manager.ts index 03cb41799..da25ee17c 100644 --- a/js/src/wrappers/vitest/context-manager.ts +++ b/js/src/wrappers/vitest/context-manager.ts @@ -54,12 +54,6 @@ class VitestContextManager { return this.contextStorage.getStore(); } - setContext(context: VitestExperimentContext): void { - // TODO(luca): Replace ALS.enterWith() with ALS.run() - // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above. - this.contextStorage.enterWith(context); - } - runInContext(context: VitestExperimentContext, callback: () => R): R { return this.contextStorage.run(context, callback); } diff --git a/js/src/wrappers/vitest/wrapper.ts b/js/src/wrappers/vitest/wrapper.ts index 4be2adf5b..b3c7344fa 100644 --- a/js/src/wrappers/vitest/wrapper.ts +++ b/js/src/wrappers/vitest/wrapper.ts @@ -87,8 +87,8 @@ export function wrapTest( // Capture context at registration time (during wrapDescribe factory execution) // as a fallback. Vitest's async test runner creates new async contexts for - // each test, so AsyncLocalStorage.enterWith() set in the describe factory - // doesn't propagate to test execution. The captured context is used when + // each test, so the describe factory's AsyncLocalStorage context doesn't + // propagate to test execution. The captured context is used when // getExperimentContext() returns null at runtime. const registrationContext = getExperimentContext(); @@ -259,9 +259,7 @@ export function wrapDescribe( config.onProgress({ type: "suite_start", suiteName }); } - contextManager.setContext(lazyContext); - - factory(); + contextManager.runInContext(lazyContext, factory); if (afterAll) { afterAll(async () => { From 18613b769d3f6e02398b19afb636ef1312d1f219 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:54:54 +0000 Subject: [PATCH 2/4] Update PR #2401 --- .changeset/remove-enter-with.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/remove-enter-with.md diff --git a/.changeset/remove-enter-with.md b/.changeset/remove-enter-with.md new file mode 100644 index 000000000..d661de675 --- /dev/null +++ b/.changeset/remove-enter-with.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +ref: Remove `AsyncLocalStorage.enterWith()` usage From 4ec703d407c4894cbc2792a6219d16e5e4f0da1a Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 24 Aug 2026 14:55:41 +0000 Subject: [PATCH 3/4] Update PR #2401 --- .../auto-instrumentation-suppression.test.ts | 2 +- .../auto-instrumentation-suppression.ts | 65 +-- .../plugins/ai-sdk-plugin.test.ts | 38 +- .../instrumentation/plugins/ai-sdk-plugin.ts | 61 +- .../plugins/ai-sdk-v7-telemetry.test.ts | 112 ---- .../claude-agent-sdk-local-tool-context.ts | 146 +---- .../plugins/claude-agent-sdk-plugin.test.ts | 2 + .../plugins/claude-agent-sdk-plugin.ts | 520 ++++++++---------- .../plugins/pi-coding-agent-plugin.ts | 10 +- .../plugins/strands-agent-sdk-plugin.test.ts | 113 ++-- .../plugins/strands-agent-sdk-plugin.ts | 274 ++++----- .../vendor-sdk-types/ai-sdk-v7-telemetry.ts | 16 - js/src/wrappers/ai-sdk/telemetry.ts | 22 - js/src/wrappers/strands-agent-sdk.test.ts | 41 +- js/src/wrappers/strands-agent-sdk.ts | 26 +- 15 files changed, 522 insertions(+), 926 deletions(-) diff --git a/js/src/instrumentation/auto-instrumentation-suppression.test.ts b/js/src/instrumentation/auto-instrumentation-suppression.test.ts index e1672d2ea..f3439ac2e 100644 --- a/js/src/instrumentation/auto-instrumentation-suppression.test.ts +++ b/js/src/instrumentation/auto-instrumentation-suppression.test.ts @@ -36,7 +36,7 @@ describe("auto instrumentation suppression context", () => { expect(isAutoInstrumentationSuppressed()).toBe(false); }); - it("restores nested allow frames at each callback boundary", async () => { + it("restores nested allow contexts at each callback boundary", async () => { await runWithAutoInstrumentationSuppressed(async () => { await runWithAutoInstrumentationAllowed(async () => { expect(isAutoInstrumentationSuppressed()).toBe(false); diff --git a/js/src/instrumentation/auto-instrumentation-suppression.ts b/js/src/instrumentation/auto-instrumentation-suppression.ts index 37fc4b4d5..8bb6b4fa7 100644 --- a/js/src/instrumentation/auto-instrumentation-suppression.ts +++ b/js/src/instrumentation/auto-instrumentation-suppression.ts @@ -1,75 +1,22 @@ -import iso, { - type IsoAsyncLocalStorage, - type IsoTracingChannel, -} from "../isomorph"; - -type AutoInstrumentationSuppressionFrame = { - mode: "allow" | "suppress"; -}; - -type AutoInstrumentationSuppressionState = { - frames: AutoInstrumentationSuppressionFrame[]; -}; +import iso, { type IsoAsyncLocalStorage } from "../isomorph"; let autoInstrumentationSuppressionStore: - | IsoAsyncLocalStorage + | IsoAsyncLocalStorage | undefined; function suppressionStore() { - autoInstrumentationSuppressionStore ??= iso.newAsyncLocalStorage< - AutoInstrumentationSuppressionState | undefined - >(); + autoInstrumentationSuppressionStore ??= iso.newAsyncLocalStorage(); return autoInstrumentationSuppressionStore; } -function currentFrames(): AutoInstrumentationSuppressionFrame[] { - return suppressionStore().getStore()?.frames ?? []; -} - export function isAutoInstrumentationSuppressed(): boolean { - const frames = currentFrames(); - return frames[frames.length - 1]?.mode === "suppress"; + return suppressionStore().getStore() === true; } export function runWithAutoInstrumentationSuppressed(callback: () => R): R { - const frame = { - mode: "suppress" as const, - }; - return suppressionStore().run( - { frames: [...currentFrames(), frame] }, - callback, - ); -} - -export function bindAutoInstrumentationSuppressionToStart( - tracingChannel: Pick, "start">, -): (() => void) | undefined { - const startChannel = tracingChannel.start; - if (!startChannel) { - return undefined; - } - - const store = suppressionStore(); - startChannel.bindStore(store, () => ({ - frames: [ - ...currentFrames(), - { - mode: "suppress" as const, - }, - ], - })); - - return () => { - startChannel.unbindStore(store); - }; + return suppressionStore().run(true, callback); } export function runWithAutoInstrumentationAllowed(callback: () => R): R { - const frame = { - mode: "allow" as const, - }; - return suppressionStore().run( - { frames: [...currentFrames(), frame] }, - callback, - ); + return suppressionStore().run(undefined, callback); } diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts index 0cfd405de..cf6522e49 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts @@ -32,7 +32,6 @@ import { } from "./ai-sdk-plugin"; import iso from "../../isomorph"; import { serializeAISDKToolsForLogging } from "../../wrappers/ai-sdk/tool-serialization"; -import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; type MockTracingChannel = { @@ -375,7 +374,20 @@ describe("AISDKPlugin", () => { expect(startEvent).not.toHaveProperty("recordInputs"); }); - it("stamps a stable unique operation key on each dispatcher", async () => { + it("uses an independent telemetry integration for each dispatcher", async () => { + const telemetryA = { + executeTool: vi.fn(({ execute }) => execute()), + onAbort: vi.fn(), + onStart: vi.fn(), + }; + const telemetryB = { + executeTool: vi.fn(({ execute }) => execute()), + onAbort: vi.fn(), + onStart: vi.fn(), + }; + telemetryMocks.braintrustAISDKTelemetry + .mockReturnValueOnce(telemetryA) + .mockReturnValueOnce(telemetryB); const dispatcherA = { executeTool: vi.fn(({ execute }) => execute()), onAbort: vi.fn(), @@ -418,22 +430,12 @@ describe("AISDKPlugin", () => { toolCallId: "tool-b", }); - const runAStart = telemetryMocks.telemetry.onStart?.mock.calls[0]?.[0]; - const runBStart = telemetryMocks.telemetry.onStart?.mock.calls[1]?.[0]; - const runAAbort = telemetryMocks.telemetry.onAbort?.mock.calls[0]?.[0]; - const runBTool = telemetryMocks.telemetry.executeTool?.mock.calls[0]?.[0]; - - expect(runAStart?.[AI_SDK_V7_OPERATION_KEY]).toEqual(expect.any(String)); - expect(runBStart?.[AI_SDK_V7_OPERATION_KEY]).toEqual(expect.any(String)); - expect(runAStart?.[AI_SDK_V7_OPERATION_KEY]).not.toBe( - runBStart?.[AI_SDK_V7_OPERATION_KEY], - ); - expect(runAAbort?.[AI_SDK_V7_OPERATION_KEY]).toBe( - runAStart?.[AI_SDK_V7_OPERATION_KEY], - ); - expect(runBTool?.[AI_SDK_V7_OPERATION_KEY]).toBe( - runBStart?.[AI_SDK_V7_OPERATION_KEY], - ); + expect(telemetryA.onStart).toHaveBeenCalledTimes(1); + expect(telemetryA.onAbort).toHaveBeenCalledTimes(1); + expect(telemetryA.executeTool).not.toHaveBeenCalled(); + expect(telemetryB.onStart).toHaveBeenCalledTimes(1); + expect(telemetryB.onAbort).not.toHaveBeenCalled(); + expect(telemetryB.executeTool).toHaveBeenCalledTimes(1); }); it("preserves existing dispatcher callback return and rejection semantics", async () => { diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index d5526da2f..45ade7d26 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -73,7 +73,6 @@ import type { AISDKV7Telemetry, AISDKV7TelemetryOptions, } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; -import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; interface AISDKPluginConfig { /** @@ -115,7 +114,6 @@ const AUTO_PATCHED_V7_TELEMETRY_DISPATCHER = Symbol.for( const RUNTIME_DENY_OUTPUT_PATHS = Symbol.for( "braintrust.ai-sdk.deny-output-paths", ); -let aiSDKV7TelemetryOperationCounter = 0; const TRANSPORT_PAYLOAD_ROOT_PATHS = [ "rawResponse", "request", @@ -840,18 +838,13 @@ function subscribeToHarnessContinuation( } function interceptAISDKV7TelemetryDispatcher(): () => void { - const telemetry = braintrustAISDKTelemetry(); return aiSDKChannels.v7CreateTelemetryDispatcher.intercept( (target, thisArg, args) => { const dispatcher = Reflect.apply(target, thisArg, args); const telemetryOptions = args[0]?.telemetry; if (telemetryOptions?.isEnabled !== false) { try { - patchAISDKV7TelemetryDispatcher( - dispatcher, - telemetry, - telemetryOptions, - ); + patchAISDKV7TelemetryDispatcher(dispatcher, telemetryOptions); } catch (error) { debugLogger.error( "Error instrumenting AI SDK v7 telemetry dispatcher:", @@ -866,7 +859,6 @@ function interceptAISDKV7TelemetryDispatcher(): () => void { function patchAISDKV7TelemetryDispatcher( dispatcher: unknown, - telemetry: AISDKV7Telemetry, telemetryOptions?: AISDKV7TelemetryOptions, ): void { if (!isObject(dispatcher)) { @@ -878,7 +870,7 @@ function patchAISDKV7TelemetryDispatcher( return; } dispatcherRecord[AUTO_PATCHED_V7_TELEMETRY_DISPATCHER] = true; - let operationKey: string | undefined; + const telemetry = braintrustAISDKTelemetry() as AISDKV7Telemetry; const telemetryEventFields: AISDKV7TelemetryOptions = {}; if (typeof telemetryOptions?.recordInputs === "boolean") { telemetryEventFields.recordInputs = telemetryOptions.recordInputs; @@ -889,49 +881,17 @@ function patchAISDKV7TelemetryDispatcher( if (typeof telemetryOptions?.functionId === "string") { telemetryEventFields.functionId = telemetryOptions.functionId; } + const hasTelemetryEventFields = Object.keys(telemetryEventFields).length > 0; - const eventWithOperationKey = (event: unknown): unknown => { - if (!isObject(event)) { + const eventWithTelemetryFields = (event: unknown): unknown => { + if (!isObject(event) || !hasTelemetryEventFields) { return event; } - const eventRecord = event as Record; - const callId = - typeof eventRecord.callId === "string" ? eventRecord.callId : "unknown"; - operationKey ??= `${callId}:${++aiSDKV7TelemetryOperationCounter}`; - - if (Object.keys(telemetryEventFields).length > 0) { - const augmentedEvent = { - ...telemetryEventFields, - ...(event as Record), - }; - try { - Object.defineProperty(augmentedEvent, AI_SDK_V7_OPERATION_KEY, { - configurable: true, - enumerable: false, - value: operationKey, - }); - } catch { - (augmentedEvent as Record)[ - AI_SDK_V7_OPERATION_KEY - ] = operationKey; - } - return augmentedEvent; - } - - try { - Object.defineProperty(eventRecord, AI_SDK_V7_OPERATION_KEY, { - configurable: true, - enumerable: false, - value: operationKey, - }); - return event; - } catch { - return { - ...(event as Record), - [AI_SDK_V7_OPERATION_KEY]: operationKey, - }; - } + return { + ...telemetryEventFields, + ...(event as Record), + }; }; for (const key of AI_SDK_V7_TELEMETRY_CALLBACKS) { @@ -949,7 +909,7 @@ function patchAISDKV7TelemetryDispatcher( try { const braintrustResult = braintrustCallback.call( telemetry, - eventWithOperationKey(event) as any, + eventWithTelemetryFields(event) as any, ); if (isPromiseLike(braintrustResult)) { void Promise.resolve(braintrustResult).catch(() => undefined); @@ -974,7 +934,6 @@ function patchAISDKV7TelemetryDispatcher( }) => braintrustExecuteTool.call(telemetry, { ...args, - ...(operationKey ? { [AI_SDK_V7_OPERATION_KEY]: operationKey } : {}), execute: () => typeof existingExecuteTool === "function" ? existingExecuteTool.call(dispatcher, args) diff --git a/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts b/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts index 8741a0e1f..517525bec 100644 --- a/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts @@ -12,7 +12,6 @@ import { registerWorkflowAgentWrapperSpan, unregisterWorkflowAgentWrapperSpan, } from "../../wrappers/ai-sdk/workflow-agent-context"; -import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; try { configureNode(); @@ -559,117 +558,6 @@ describe("braintrustAISDKTelemetry", () => { ).toHaveLength(0); }); - it("keeps concurrent WorkflowAgent streams with shared callIds separate", async () => { - const telemetry = braintrustAISDKTelemetry(); - const callId = "workflow-agent"; - const runA = "workflow-agent:run-a"; - const runB = "workflow-agent:run-b"; - - telemetry.onStart?.({ - [AI_SDK_V7_OPERATION_KEY]: runA, - callId, - messages: [{ role: "user", content: "First workflow run" }], - operationId: "ai.workflowAgent.stream", - }); - telemetry.onStart?.({ - [AI_SDK_V7_OPERATION_KEY]: runB, - callId, - messages: [{ role: "user", content: "Second workflow run" }], - operationId: "ai.workflowAgent.stream", - }); - - telemetry.onLanguageModelCallStart?.({ - [AI_SDK_V7_OPERATION_KEY]: runA, - callId, - prompt: [{ role: "user", content: "First workflow run" }], - }); - telemetry.onLanguageModelCallEnd?.({ - [AI_SDK_V7_OPERATION_KEY]: runA, - callId, - text: "First answer", - }); - telemetry.onLanguageModelCallStart?.({ - [AI_SDK_V7_OPERATION_KEY]: runB, - callId, - prompt: [{ role: "user", content: "Second workflow run" }], - }); - telemetry.onLanguageModelCallEnd?.({ - [AI_SDK_V7_OPERATION_KEY]: runB, - callId, - text: "Calling get_weather.", - }); - - telemetry.onToolExecutionStart?.({ - [AI_SDK_V7_OPERATION_KEY]: runB, - toolCall: { - toolCallId: "tool-run-b", - toolName: "get_weather", - input: { location: "Vienna, Austria" }, - }, - }); - telemetry.onToolExecutionEnd?.({ - [AI_SDK_V7_OPERATION_KEY]: runB, - output: { condition: "sunny" }, - success: true, - toolCall: { - toolCallId: "tool-run-b", - toolName: "get_weather", - }, - }); - - telemetry.onEnd?.({ - [AI_SDK_V7_OPERATION_KEY]: runA, - callId, - messages: [{ role: "assistant", content: "First answer" }], - operationId: "ai.workflowAgent.stream", - text: "First answer", - }); - telemetry.onEnd?.({ - [AI_SDK_V7_OPERATION_KEY]: runB, - callId, - messages: [{ role: "assistant", content: "Second answer" }], - operationId: "ai.workflowAgent.stream", - text: "Second answer", - }); - - const spans = (await backgroundLogger.drain()) as Array< - Record - >; - const workflowSpans = spans.filter( - (span) => span.span_attributes?.name === "WorkflowAgent.stream", - ); - const firstWorkflow = workflowSpans.find((span) => - JSON.stringify(span.input).includes("First workflow run"), - ); - const secondWorkflow = workflowSpans.find((span) => - JSON.stringify(span.input).includes("Second workflow run"), - ); - const firstModel = spans.find( - (span) => - span.span_attributes?.name === "doGenerate" && - JSON.stringify(span.input).includes("First workflow run"), - ); - const secondModel = spans.find( - (span) => - span.span_attributes?.name === "doGenerate" && - JSON.stringify(span.input).includes("Second workflow run"), - ); - const tool = spans.find( - (span) => span.span_attributes?.name === "get_weather", - ); - - expect(workflowSpans).toHaveLength(2); - expect(firstWorkflow?.output).toMatchObject({ text: "First answer" }); - expect(secondWorkflow?.output).toMatchObject({ text: "Second answer" }); - expect(firstModel?.span_parents).toEqual([firstWorkflow?.span_id]); - expect(secondModel?.span_parents).toEqual([secondWorkflow?.span_id]); - expect(tool).toMatchObject({ - input: { location: "Vienna, Austria" }, - output: { condition: "sunny" }, - span_parents: [secondWorkflow?.span_id], - }); - }); - it("keeps concurrent direct WorkflowAgent telemetry separated without dispatcher keys", async () => { const telemetry = braintrustAISDKTelemetry(); const callId = "workflow-agent"; diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts index d98223b8b..72ae6b5c8 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts @@ -1,80 +1,34 @@ import iso from "../../isomorph"; -type LocalToolParentResolver = (toolUseId: string) => Promise; - -export type ClaudeAgentSDKLocalToolContext = { - resolveLocalToolParent?: LocalToolParentResolver; -}; - -const LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED = Symbol.for( - "braintrust.claude_agent_sdk.local_tool_context_async_iterator_patched", -); - -type AsyncLocalStorageLike = { - getStore: () => T | undefined; - run: (store: T, callback: () => R) => R; -}; - -function createLocalToolContextStore(): AsyncLocalStorageLike { - const maybeIsoWithAsyncLocalStorage = iso as { - newAsyncLocalStorage?: () => AsyncLocalStorageLike; - }; - - if ( - typeof maybeIsoWithAsyncLocalStorage.newAsyncLocalStorage === "function" - ) { - return maybeIsoWithAsyncLocalStorage.newAsyncLocalStorage(); - } - - let currentStore: ClaudeAgentSDKLocalToolContext | undefined; - return { - getStore() { - return currentStore; - }, - run(store, callback) { - const previousStore = currentStore; - currentStore = store; - try { - return callback(); - } finally { - currentStore = previousStore; - } - }, - }; -} +export type ClaudeLocalToolParentResolver = ( + toolUseId: string, +) => Promise; -const localToolContextStore = createLocalToolContextStore(); +const localToolContextStore = + iso.newAsyncLocalStorage(); const localToolParentResolversByToolUseId = new Map< string, - LocalToolParentResolver + ClaudeLocalToolParentResolver >(); -export function createClaudeLocalToolContext(): ClaudeAgentSDKLocalToolContext { - return {}; -} - export function runWithClaudeLocalToolContext( callback: () => R, - context?: ClaudeAgentSDKLocalToolContext, + resolver: ClaudeLocalToolParentResolver, ): R { - return localToolContextStore.run( - context ?? createClaudeLocalToolContext(), - callback, - ); + return localToolContextStore.run(resolver, callback); } export function registerClaudeLocalToolParentResolver( toolUseId: string, - resolver: LocalToolParentResolver, + resolver: ClaudeLocalToolParentResolver, ): void { localToolParentResolversByToolUseId.set(toolUseId, resolver); } export function getClaudeLocalToolParentResolver( toolUseId?: string, -): LocalToolParentResolver | undefined { - const currentResolver = - localToolContextStore.getStore()?.resolveLocalToolParent; +): ClaudeLocalToolParentResolver | undefined { + const currentResolver = localToolContextStore.getStore(); if (!toolUseId) { return currentResolver; } @@ -83,81 +37,3 @@ export function getClaudeLocalToolParentResolver( localToolParentResolversByToolUseId.delete(toolUseId); return currentResolver ?? registeredResolver; } - -function isAsyncIterable(value: unknown): value is AsyncIterable { - return ( - value !== null && - typeof value === "object" && - Symbol.asyncIterator in value && - typeof value[Symbol.asyncIterator] === "function" - ); -} - -export function bindClaudeLocalToolContextToAsyncIterable( - result: T, - localToolContext: ClaudeAgentSDKLocalToolContext, -): T { - if ( - !isAsyncIterable(result) || - Object.isFrozen(result) || - Object.isSealed(result) - ) { - return result; - } - - const stream = result as AsyncIterable & { - [Symbol.asyncIterator]: (() => AsyncIterator) & { - [LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED]?: boolean; - }; - }; - const originalAsyncIterator = stream[Symbol.asyncIterator]; - if (originalAsyncIterator[LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED]) { - return result; - } - - const patchedAsyncIterator = function (this: unknown) { - return runWithClaudeLocalToolContext(() => { - const iterator = Reflect.apply(originalAsyncIterator, this, []); - if (!iterator || typeof iterator !== "object") { - return iterator; - } - - const patchMethod = (methodName: "next" | "return" | "throw") => { - const originalMethod = Reflect.get(iterator, methodName); - if (typeof originalMethod !== "function") { - return; - } - - Reflect.set(iterator, methodName, (...args: unknown[]) => - runWithClaudeLocalToolContext( - () => - Reflect.apply( - originalMethod as (...methodArgs: unknown[]) => unknown, - iterator, - args, - ), - localToolContext, - ), - ); - }; - - patchMethod("next"); - patchMethod("return"); - patchMethod("throw"); - return iterator; - }, localToolContext); - }; - - Object.defineProperty( - patchedAsyncIterator, - LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED, - { - configurable: false, - enumerable: false, - value: true, - writable: false, - }, - ); - Reflect.set(stream, Symbol.asyncIterator, patchedAsyncIterator); - return result; -} diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts index 42f9da7a6..1c35a4a27 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts @@ -1,4 +1,5 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { AsyncLocalStorage } from "node:async_hooks"; // Mock iso's newTracingChannel - must be before any imports that use it const streamPatcherMock = vi.hoisted(() => ({ @@ -12,6 +13,7 @@ const streamPatcherMock = vi.hoisted(() => ({ vi.mock("../../isomorph", () => ({ default: { + newAsyncLocalStorage: () => new AsyncLocalStorage(), newTracingChannel: vi.fn(), }, })); diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts index e9414940b..86389ee9d 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts @@ -1,7 +1,5 @@ import { BasePlugin } from "../core"; -import type { ChannelMessage } from "../core/channel-definitions"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; -import type { IsoChannelHandlers } from "../../isomorph"; import { debugLogger } from "../../debug-logger"; import { startSpan as startBaseSpan } from "../../logger"; import type { Span } from "../../logger"; @@ -24,11 +22,8 @@ import { wrapLocalMcpServerToolHandlers, } from "./claude-agent-sdk-local-tool-spans"; import { - bindClaudeLocalToolContextToAsyncIterable, - createClaudeLocalToolContext, registerClaudeLocalToolParentResolver, runWithClaudeLocalToolContext, - type ClaudeAgentSDKLocalToolContext, } from "./claude-agent-sdk-local-tool-context"; import type { ClaudeAgentSDKHookCallback, @@ -988,7 +983,7 @@ type QueryState = { latestRootLlmParentRef: { value: string | undefined }; toolUseToParent: Map; usageByMessageId: Map; - localToolContext: ClaudeAgentSDKLocalToolContext; + localToolParentResolver: ParentSpanResolver; }; function setSubAgentPromptMessages( @@ -1583,7 +1578,7 @@ async function finalizeQuerySpan(state: QueryState): Promise { export class ClaudeAgentSDKPlugin extends BasePlugin { protected onEnable(): void { - this.subscribeToQuery(); + this.interceptQuery(); } protected onDisable(): void { @@ -1593,317 +1588,276 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { this.unsubscribers = []; } - private subscribeToQuery(): void { - const spans = new WeakMap(); - - const handlers: IsoChannelHandlers< - ChannelMessage - > = { - start: (event) => { - const params = (event.arguments[0] ?? {}) as ClaudeAgentSDKQueryParams; - const originalPrompt = params.prompt; - const options = params.options ?? {}; - const promptIsAsyncIterable = isAsyncIterable(originalPrompt); - let promptStarted = false; - let capturedPromptMessages: ClaudeAgentSDKMessage[] | undefined; - let resolvePromptDone: (() => void) | undefined; - const promptDone = new Promise((resolve) => { - resolvePromptDone = resolve; - }); + private interceptQuery(): void { + const startQuery = (params: ClaudeAgentSDKQueryParams): QueryState => { + const originalPrompt = params.prompt; + const options = params.options ?? {}; + const promptIsAsyncIterable = isAsyncIterable(originalPrompt); + let promptStarted = false; + let capturedPromptMessages: ClaudeAgentSDKMessage[] | undefined; + let resolvePromptDone: (() => void) | undefined; + const promptDone = new Promise((resolve) => { + resolvePromptDone = resolve; + }); - if (promptIsAsyncIterable) { - capturedPromptMessages = []; - const promptStream = - originalPrompt as AsyncIterable; - params.prompt = (async function* () { - promptStarted = true; - try { - for await (const message of promptStream) { - capturedPromptMessages!.push(message); - yield message; - } - } finally { - resolvePromptDone?.(); + if (promptIsAsyncIterable) { + capturedPromptMessages = []; + const promptStream = + originalPrompt as AsyncIterable; + params.prompt = (async function* () { + promptStarted = true; + try { + for await (const message of promptStream) { + capturedPromptMessages!.push(message); + yield message; } - })(); - } - - const span = startBaseSpan( - withSpanInstrumentationName( - { - name: "Claude Agent", - spanAttributes: { - type: SpanTypeAttribute.TASK, - }, - }, - INSTRUMENTATION_NAMES.CLAUDE_AGENT_SDK, - ), - ); - const startTime = getCurrentUnixTimestamp(); - - try { - span.log({ - input: - typeof originalPrompt === "string" - ? originalPrompt - : promptIsAsyncIterable - ? undefined - : originalPrompt !== undefined - ? String(originalPrompt) - : undefined, - metadata: filterSerializableOptions(options), - }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error("Error extracting input for Claude Agent SDK:", error); - } - - const activeToolSpans = new Map(); - const activeLlmSpansByParentToolUse = new Map(); - const conversationHistoryByParentKey = new Map< - string, - ClaudeConversationMessage[] - >(); - const subAgentSpans = new Map(); - const endedSubAgentSpans = new Set(); - const toolUseToParent = new Map(); - const latestLlmParentBySubAgentToolUse = new Map(); - const latestRootLlmParentRef = { - value: undefined as string | undefined, - }; - const subAgentDetailsByToolUseId = new Map(); - const taskIdToToolUseId = new Map(); - const promptMessagesByParentKey = new Map< - string, - ClaudeConversationMessage[] - >(); - const promptSourcePriorityByParentKey = new Map(); - const localToolContext = createClaudeLocalToolContext(); - const { hasLocalToolHandlers, localToolHookNames } = - prepareLocalToolHandlersInMcpServers(options.mcpServers); - const skipLocalToolHooks = - options[CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION] === true || - hasLocalToolHandlers; - const resolveToolUseParentSpan: ParentSpanResolver = async ( - toolUseID, - context, - ) => { - const trackedParentToolUseId = toolUseToParent.get(toolUseID); - const parentToolUseId = - trackedParentToolUseId ?? - (context?.agentId - ? (taskIdToToolUseId.get(context.agentId) ?? null) - : null); - const parentKey = llmParentKey(parentToolUseId); - const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey); - const latestLlmParent = parentToolUseId - ? latestLlmParentBySubAgentToolUse.get(parentToolUseId) - : latestRootLlmParentRef.value; - - // Tool spans should be siblings of the driving LLM turn, but we still - // materialize that LLM span first so trace ordering reflects that the - // tool call was produced by the model. - if (!activeLlmSpan && !latestLlmParent) { - await ensureActiveLlmSpanForParentToolUse( - span, - activeLlmSpansByParentToolUse, - subAgentDetailsByToolUseId, - activeToolSpans, - subAgentSpans, - parentToolUseId, - getCurrentUnixTimestamp(), - ); + } finally { + resolvePromptDone?.(); } + })(); + } - if (parentToolUseId) { - const subAgentSpan = await ensureSubAgentSpan( - subAgentDetailsByToolUseId, - span, - activeToolSpans, - subAgentSpans, - parentToolUseId, - ); - return subAgentSpan.export(); - } - - return span.export(); - }; - - localToolContext.resolveLocalToolParent = resolveToolUseParentSpan; - const optionsWithHooks = injectTracingHooks( - options, - resolveToolUseParentSpan, - taskIdToToolUseId, - toolUseToParent, - activeToolSpans, - localToolHookNames, - skipLocalToolHooks, - subAgentDetailsByToolUseId, - subAgentSpans, - endedSubAgentSpans, - ); + const span = startBaseSpan( + withSpanInstrumentationName( + { + name: "Claude Agent", + spanAttributes: { + type: SpanTypeAttribute.TASK, + }, + }, + INSTRUMENTATION_NAMES.CLAUDE_AGENT_SDK, + ), + ); + const startTime = getCurrentUnixTimestamp(); - params.options = optionsWithHooks; - event.arguments[0] = params; - - spans.set(event, { - activeLlmSpansByParentToolUse, - activePartialMessageIdByParentKey: new Map(), - activeToolSpans, - conversationHistoryByParentKey, - capturedPromptMessages, - currentMessageId: undefined, - currentMessageStartTime: startTime, - currentMessages: [], - endedSubAgentSpans, - finalOutputUsageMessageIds: new Set(), - finalResults: [], - options: optionsWithHooks, - originalPrompt, - processing: Promise.resolve(), - promptDone, - promptMessagesByParentKey, - promptStarted: () => promptStarted, - promptSourcePriorityByParentKey, - span, - subAgentDetailsByToolUseId, - subAgentSpans, - taskIdToToolUseId, - latestLlmParentBySubAgentToolUse, - latestRootLlmParentRef, - toolUseToParent, - usageByMessageId: new Map(), - localToolContext, + try { + span.log({ + input: + typeof originalPrompt === "string" + ? originalPrompt + : promptIsAsyncIterable + ? undefined + : originalPrompt !== undefined + ? String(originalPrompt) + : undefined, + metadata: filterSerializableOptions(options), }); - }, + } catch (error) { + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.error("Error extracting input for Claude Agent SDK:", error); + } - end: (event) => { - const state = spans.get(event); - if (!state) { - return; + const activeToolSpans = new Map(); + const activeLlmSpansByParentToolUse = new Map(); + const conversationHistoryByParentKey = new Map< + string, + ClaudeConversationMessage[] + >(); + const subAgentSpans = new Map(); + const endedSubAgentSpans = new Set(); + const toolUseToParent = new Map(); + const latestLlmParentBySubAgentToolUse = new Map(); + const latestRootLlmParentRef = { + value: undefined as string | undefined, + }; + const subAgentDetailsByToolUseId = new Map(); + const taskIdToToolUseId = new Map(); + const promptMessagesByParentKey = new Map< + string, + ClaudeConversationMessage[] + >(); + const promptSourcePriorityByParentKey = new Map(); + const { hasLocalToolHandlers, localToolHookNames } = + prepareLocalToolHandlersInMcpServers(options.mcpServers); + const skipLocalToolHooks = + options[CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION] === true || + hasLocalToolHandlers; + const resolveToolUseParentSpan: ParentSpanResolver = async ( + toolUseID, + context, + ) => { + const trackedParentToolUseId = toolUseToParent.get(toolUseID); + const parentToolUseId = + trackedParentToolUseId ?? + (context?.agentId + ? (taskIdToToolUseId.get(context.agentId) ?? null) + : null); + const parentKey = llmParentKey(parentToolUseId); + const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey); + const latestLlmParent = parentToolUseId + ? latestLlmParentBySubAgentToolUse.get(parentToolUseId) + : latestRootLlmParentRef.value; + + // Tool spans should be siblings of the driving LLM turn, but we still + // materialize that LLM span first so trace ordering reflects that the + // tool call was produced by the model. + if (!activeLlmSpan && !latestLlmParent) { + await ensureActiveLlmSpanForParentToolUse( + span, + activeLlmSpansByParentToolUse, + subAgentDetailsByToolUseId, + activeToolSpans, + subAgentSpans, + parentToolUseId, + getCurrentUnixTimestamp(), + ); } - const eventResult = bindClaudeLocalToolContextToAsyncIterable( - event.result, - state.localToolContext, - ); - if (eventResult === undefined) { - state.span.end(); - spans.delete(event); - return; + if (parentToolUseId) { + const subAgentSpan = await ensureSubAgentSpan( + subAgentDetailsByToolUseId, + span, + activeToolSpans, + subAgentSpans, + parentToolUseId, + ); + return subAgentSpan.export(); } - if (isAsyncIterable(eventResult)) { - patchStreamIfNeeded(eventResult, { - onChunk: (message: ClaudeAgentSDKMessage) => { - maybeTrackToolUseContext(state, message); - state.processing = state.processing - .then(() => handleStreamMessage(state, message)) - .catch((error) => { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error( - "Error processing Claude Agent SDK stream chunk:", - error, - ); - }); - }, - onComplete: () => - state.processing - .then(() => finalizeQuerySpan(state)) - .finally(() => { - spans.delete(event); - }), - onError: (error: Error) => - state.processing - .then(() => { - state.span.log({ - error: error.message, - }); - }) - .then(() => finalizeQuerySpan(state)) - .finally(() => { - spans.delete(event); - }), - }); + return span.export(); + }; + + const optionsWithHooks = injectTracingHooks( + options, + resolveToolUseParentSpan, + taskIdToToolUseId, + toolUseToParent, + activeToolSpans, + localToolHookNames, + skipLocalToolHooks, + subAgentDetailsByToolUseId, + subAgentSpans, + endedSubAgentSpans, + ); - return; - } + params.options = optionsWithHooks; + + return { + activeLlmSpansByParentToolUse, + activePartialMessageIdByParentKey: new Map(), + activeToolSpans, + conversationHistoryByParentKey, + capturedPromptMessages, + currentMessageId: undefined, + currentMessageStartTime: startTime, + currentMessages: [], + endedSubAgentSpans, + finalOutputUsageMessageIds: new Set(), + finalResults: [], + options: optionsWithHooks, + originalPrompt, + processing: Promise.resolve(), + promptDone, + promptMessagesByParentKey, + promptStarted: () => promptStarted, + promptSourcePriorityByParentKey, + span, + subAgentDetailsByToolUseId, + subAgentSpans, + taskIdToToolUseId, + latestLlmParentBySubAgentToolUse, + latestRootLlmParentRef, + toolUseToParent, + usageByMessageId: new Map(), + localToolParentResolver: resolveToolUseParentSpan, + }; + }; - try { - state.span.log({ output: eventResult }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error("Error extracting output for Claude Agent SDK:", error); - } finally { - state.span.end(); - spans.delete(event); - } - }, + const finishQuery = ( + state: QueryState, + result: AsyncIterable, + ): void => { + if (isAsyncIterable(result)) { + patchStreamIfNeeded(result, { + aroundNext: (callback) => + runWithClaudeLocalToolContext( + callback, + state.localToolParentResolver, + ), + onChunk: (message: ClaudeAgentSDKMessage) => { + maybeTrackToolUseContext(state, message); + state.processing = state.processing + .then(() => handleStreamMessage(state, message)) + .catch((error) => { + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.error( + "Error processing Claude Agent SDK stream chunk:", + error, + ); + }); + }, + onComplete: () => + state.processing.then(() => finalizeQuerySpan(state)), + onError: (error: Error) => + state.processing + .then(() => { + state.span.log({ error: error.message }); + }) + .then(() => finalizeQuerySpan(state)), + }); - error: (event) => { - const state = spans.get(event); - if (!state || !event.error) { - return; - } + return; + } - state.span.log({ - error: event.error.message, - }); + try { + state.span.log({ output: result }); + } catch (error) { + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.error("Error extracting output for Claude Agent SDK:", error); + } finally { state.span.end(); - spans.delete(event); - }, + } }; this.unsubscribers.push( - claudeAgentSDKChannels.query.intercept( - (target, thisArg, args, additional) => { - const event: ChannelMessage = { - ...additional, - arguments: args, - }; - try { - handlers.start?.(event, claudeAgentSDKChannels.query.channelName); - } catch (error) { - debugLogger.error( - "Error starting Claude Agent SDK instrumentation:", - error, - ); - } + claudeAgentSDKChannels.query.intercept((target, thisArg, args) => { + let state: QueryState | undefined; + try { + args[0] ??= {}; + state = startQuery(args[0]); + } catch (error) { + debugLogger.error( + "Error starting Claude Agent SDK instrumentation:", + error, + ); + } - const state = spans.get(event); - const invokeTarget = () => Reflect.apply(target, thisArg, args); - try { - const result = state - ? runWithClaudeLocalToolContext( - invokeTarget, - state.localToolContext, - ) - : invokeTarget(); - event.result = result; + const invokeTarget = () => Reflect.apply(target, thisArg, args); + try { + const result = state + ? runWithClaudeLocalToolContext( + invokeTarget, + state.localToolParentResolver, + ) + : invokeTarget(); + if (state) { try { - handlers.end?.(event, claudeAgentSDKChannels.query.channelName); + finishQuery(state, result); } catch (error) { debugLogger.error( "Error finalizing Claude Agent SDK instrumentation:", error, ); } - return result; - } catch (error) { - event.error = - error instanceof Error ? error : new Error(String(error)); + } + return result; + } catch (error) { + if (state) { try { - handlers.error?.(event, claudeAgentSDKChannels.query.channelName); + state.span.log({ + error: error instanceof Error ? error.message : String(error), + }); + state.span.end(); } catch (instrumentationError) { debugLogger.error( "Error handling Claude Agent SDK instrumentation failure:", instrumentationError, ); } - throw error; } - }, - ), + throw error; + } + }), ); } } diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts index 6d4e442c7..83381b4ac 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts +++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts @@ -75,9 +75,7 @@ const piAgentEventSubscriptions = new WeakSet(); const PI_TOOL_EXECUTE_WRAPPED = Symbol.for( "braintrust.pi_coding_agent.tool_execute_wrapped", ); -let piPromptContextStore: - | IsoAsyncLocalStorage - | undefined; +let piPromptContextStore: IsoAsyncLocalStorage | undefined; export class PiCodingAgentPlugin extends BasePlugin { private readonly activePromptStates = new Set(); @@ -236,10 +234,8 @@ function isPiAgent(value: unknown): value is PiAgent { ); } -function promptContextStore(): IsoAsyncLocalStorage { - piPromptContextStore ??= iso.newAsyncLocalStorage< - PiPromptState | undefined - >(); +function promptContextStore(): IsoAsyncLocalStorage { + piPromptContextStore ??= iso.newAsyncLocalStorage(); return piPromptContextStore; } diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts index d52d0ce50..a2b3f0527 100644 --- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts @@ -1,32 +1,26 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -const { - mockBindStore, - mockWithCurrent, - mockNewAsyncLocalStorage, - mockStartSpan, - mockUnbindStore, -} = vi.hoisted(() => ({ - mockBindStore: vi.fn(), - mockWithCurrent: vi.fn(), - mockNewAsyncLocalStorage: vi.fn(() => { - let current: unknown; - return { - getStore: vi.fn(() => current), - run: vi.fn((store: unknown, callback: () => unknown) => { - const previous = current; - current = store; - try { - return callback(); - } finally { - current = previous; - } - }), - }; +const { mockWithCurrent, mockNewAsyncLocalStorage, mockStartSpan } = vi.hoisted( + () => ({ + mockWithCurrent: vi.fn(), + mockNewAsyncLocalStorage: vi.fn(() => { + let current: unknown; + return { + getStore: vi.fn(() => current), + run: vi.fn((store: unknown, callback: () => unknown) => { + const previous = current; + current = store; + try { + return callback(); + } finally { + current = previous; + } + }), + }; + }), + mockStartSpan: vi.fn(), }), - mockStartSpan: vi.fn(), - mockUnbindStore: vi.fn(), -})); +); vi.mock("../../isomorph", () => ({ default: { @@ -72,13 +66,48 @@ describe("StrandsAgentSDKPlugin", () => { handlersByName = new Map(); spans = []; mockNewTracingChannel.mockImplementation((name: string) => ({ - start: { - bindStore: mockBindStore, - unbindStore: mockUnbindStore, - }, - subscribe: vi.fn((handlers) => handlersByName.set(name, handlers)), + intercept: vi.fn((interceptor) => { + const handlers = { + end: (event: any) => + interceptor( + () => + typeof event.invoke === "function" + ? event.invoke() + : event.result, + event.self, + event.arguments ?? [], + { + ...(event.agent ? { agent: event.agent } : {}), + ...(event.orchestrator + ? { orchestrator: event.orchestrator } + : {}), + }, + ), + error: (event: any) => { + try { + interceptor( + () => { + throw event.error; + }, + event.self, + event.arguments ?? [], + { + ...(event.agent ? { agent: event.agent } : {}), + ...(event.orchestrator + ? { orchestrator: event.orchestrator } + : {}), + }, + ); + } catch { + // The real interceptor preserves the target error. + } + }, + start: vi.fn(), + }; + handlersByName.set(name, handlers); + return vi.fn(); + }), traceSync: vi.fn((fn) => fn()), - unsubscribe: vi.fn(), })); currentSpan = undefined; mockWithCurrent.mockImplementation((span: any, callback: () => unknown) => { @@ -119,7 +148,7 @@ describe("StrandsAgentSDKPlugin", () => { vi.clearAllMocks(); }); - it("subscribes to Strands stream channels and binds suppression", () => { + it("intercepts Strands stream channels", () => { const plugin = new StrandsAgentSDKPlugin(); plugin.enable(); @@ -132,11 +161,7 @@ describe("StrandsAgentSDKPlugin", () => { expect( handlersByName.has("orchestrion:@strands-agents/sdk:Swarm.stream"), ).toBe(true); - expect(mockBindStore).toHaveBeenCalledTimes(3); - plugin.disable(); - - expect(mockUnbindStore).toHaveBeenCalledTimes(3); }); it("records agent model and tool spans from stream events", async () => { @@ -223,6 +248,10 @@ describe("StrandsAgentSDKPlugin", () => { ); const event = { arguments: ["hello", undefined], + invoke: () => { + suppressionStates.push(isAutoInstrumentationSuppressed()); + return stream; + }, moduleVersion: "1.6.0", result: stream, self: agent, @@ -236,7 +265,15 @@ describe("StrandsAgentSDKPlugin", () => { } expect(chunks).toHaveLength(6); - expect(suppressionStates).toEqual([true, true, true, true, true, true]); + expect(suppressionStates).toEqual([ + true, + true, + true, + true, + true, + true, + true, + ]); const rootSpan = spans.find((span) => span.args.name === "Agent: helper"); const modelSpan = spans.find( (span) => span.args.name === "Strands model: gpt-4o-mini", diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts index 255578a13..13a434cff 100644 --- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts @@ -1,7 +1,5 @@ import { BasePlugin, toLoggedError } from "../core"; -import type { ChannelMessage } from "../core/channel-definitions"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; -import type { IsoChannelHandlers } from "../../isomorph"; import { debugLogger } from "../../debug-logger"; import { Attachment, @@ -18,10 +16,7 @@ import { LRUCache } from "../../lru-cache"; import { getCurrentUnixTimestamp } from "../../util"; import { SpanTypeAttribute, isObject } from "../../../util/index"; import { convertDataToBlob } from "../../wrappers/attachment-utils"; -import { - bindAutoInstrumentationSuppressionToStart, - runWithAutoInstrumentationSuppressed, -} from "../auto-instrumentation-suppression"; +import { runWithAutoInstrumentationSuppressed } from "../auto-instrumentation-suppression"; import { strandsAgentSDKChannels } from "./strands-agent-sdk-channels"; import type { StrandsAfterModelCallEvent, @@ -105,12 +100,12 @@ export class StrandsAgentSDKPlugin extends BasePlugin { private readonly activeChildParents: ActiveChildParents = new WeakMap(); protected onEnable(): void { - this.subscribeToAgentStream(); - this.subscribeToMultiAgentStream( + this.interceptAgentStream(); + this.interceptMultiAgentStream( strandsAgentSDKChannels.graphStream, "Graph.stream", ); - this.subscribeToMultiAgentStream( + this.interceptMultiAgentStream( strandsAgentSDKChannels.swarmStream, "Swarm.stream", ); @@ -123,145 +118,124 @@ export class StrandsAgentSDKPlugin extends BasePlugin { this.unsubscribers = []; } - private subscribeToAgentStream(): void { - const channel = strandsAgentSDKChannels.agentStream.tracingChannel(); - const states = new WeakMap(); - const unbindAutoInstrumentationSuppression = - bindAutoInstrumentationSuppressionToStart(channel); - - const handlers: IsoChannelHandlers< - ChannelMessage - > = { - start: (event) => { - const state = startAgentStream(event, this.activeChildParents); - if (state) { - states.set(event, state); - } - }, - end: (event) => { - const state = states.get(event); - if (!state) { - return; - } - - const result = event.result; - if (isAsyncIterable(result)) { - patchStreamIfNeeded(result, { - aroundNext: (callback) => - runWithAutoInstrumentationSuppressed(callback), - onChunk: (chunk) => handleAgentStreamEvent(state, chunk), - onComplete: () => { - finalizeAgentStream(state); - states.delete(event); - }, - onError: (error) => { - finalizeAgentStream(state, error); - states.delete(event); - }, - }); - return; - } - - finalizeAgentStream(state, undefined, result); - states.delete(event); - }, - error: (event) => { - const state = states.get(event); - if (!state || !event.error) { - return; - } - finalizeAgentStream(state, event.error); - states.delete(event); - }, - }; - - channel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindAutoInstrumentationSuppression?.(); - channel.unsubscribe(handlers); - }); + private interceptAgentStream(): void { + this.unsubscribers.push( + strandsAgentSDKChannels.agentStream.intercept( + (target, thisArg, args, additional) => + instrumentStrandsStreamInvocation< + AgentStreamState, + StrandsAgentStreamEvent, + ReturnType + >({ + finalize: finalizeAgentStream, + handleChunk: handleAgentStreamEvent, + invoke: () => Reflect.apply(target, thisArg, args), + name: "Strands Agent SDK", + start: () => + startAgentStream( + args[0], + extractAgent(additional.agent, thisArg), + this.activeChildParents, + ), + }), + ), + ); } - private subscribeToMultiAgentStream( + private interceptMultiAgentStream( channel: MultiAgentStreamChannel, operation: MultiAgentStreamState["operation"], ): void { - const tracingChannel = channel.tracingChannel(); - const states = new WeakMap(); - const unbindAutoInstrumentationSuppression = - bindAutoInstrumentationSuppressionToStart(tracingChannel); - - const handlers: IsoChannelHandlers> = { - start: (event) => { - const state = startMultiAgentStream( - event, - operation, - this.activeChildParents, - ); - if (state) { - states.set(event, state); - } - }, - end: (event) => { - const state = states.get(event); - if (!state) { - return; - } + this.unsubscribers.push( + channel.intercept((target, thisArg, args, additional) => + instrumentStrandsStreamInvocation< + MultiAgentStreamState, + StrandsMultiAgentStreamEvent, + ReturnType + >({ + finalize: (state, error, output) => + finalizeMultiAgentStream( + state, + this.activeChildParents, + error, + output, + ), + handleChunk: (state, chunk) => + handleMultiAgentStreamEvent(state, chunk, this.activeChildParents), + invoke: () => Reflect.apply(target, thisArg, args), + name: "Strands multi-agent", + start: () => + startMultiAgentStream( + args[0], + extractOrchestrator(additional.orchestrator, thisArg), + operation, + this.activeChildParents, + ), + }), + ), + ); + } +} - const result = event.result; - if (isAsyncIterable(result)) { - patchStreamIfNeeded(result, { - aroundNext: (callback) => - runWithAutoInstrumentationSuppressed(callback), - onChunk: (chunk) => - handleMultiAgentStreamEvent( - state, - chunk, - this.activeChildParents, - ), - onComplete: () => { - finalizeMultiAgentStream(state, this.activeChildParents); - states.delete(event); - }, - onError: (error) => { - finalizeMultiAgentStream(state, this.activeChildParents, error); - states.delete(event); - }, - }); - return; - } +function instrumentStrandsStreamInvocation(options: { + finalize: (state: TState, error?: unknown, output?: unknown) => void; + handleChunk: (state: TState, chunk: TChunk) => void; + invoke: () => TResult; + name: string; + start: () => TState; +}): TResult { + let state: TState | undefined; + try { + state = options.start(); + } catch (error) { + debugLogger.error(`Error starting ${options.name} instrumentation:`, error); + } - finalizeMultiAgentStream( - state, - this.activeChildParents, - undefined, - result, + let result: TResult; + try { + result = runWithAutoInstrumentationSuppressed(options.invoke); + } catch (error) { + if (state) { + try { + options.finalize(state, error); + } catch (instrumentationError) { + debugLogger.error( + `Error handling ${options.name} instrumentation failure:`, + instrumentationError, ); - states.delete(event); - }, - error: (event) => { - const state = states.get(event); - if (!state || !event.error) { - return; - } - finalizeMultiAgentStream(state, this.activeChildParents, event.error); - states.delete(event); - }, - }; - - tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindAutoInstrumentationSuppression?.(); - tracingChannel.unsubscribe(handlers); - }); + } + } + throw error; + } + + if (state) { + try { + if (isAsyncIterable(result)) { + patchStreamIfNeeded(result, { + aroundNext: (callback) => + runWithAutoInstrumentationSuppressed(callback), + onChunk: (chunk) => options.handleChunk(state, chunk), + onComplete: () => options.finalize(state), + onError: (error) => options.finalize(state, error), + }); + } else { + options.finalize(state, undefined, result); + } + } catch (error) { + debugLogger.error( + `Error finalizing ${options.name} instrumentation:`, + error, + ); + } } + return result; } function startAgentStream( - event: ChannelMessage, + input: unknown, + agent: StrandsAgent | undefined, activeChildParents: ActiveChildParents, -): AgentStreamState | undefined { - const agent = extractAgent(event); +): AgentStreamState { const model = agent?.model; const metadata = { ...extractAgentMetadata(agent), @@ -273,17 +247,14 @@ function startAgentStream( ? getOnlyChildParent(activeChildParents, agent) : undefined; const attachmentCache = createStrandsAttachmentCache(); - const input = processStrandsInputAttachments( - event.arguments[0], - attachmentCache, - ); + const processedInput = processStrandsInputAttachments(input, attachmentCache); const span = parentSpan ? withCurrent(parentSpan, () => startBaseSpan( withSpanInstrumentationName( { event: { - input, + input: processedInput, metadata, }, name: formatAgentSpanName(agent), @@ -297,7 +268,7 @@ function startAgentStream( withSpanInstrumentationName( { event: { - input, + input: processedInput, metadata, }, name: formatAgentSpanName(agent), @@ -318,11 +289,11 @@ function startAgentStream( } function startMultiAgentStream( - event: ChannelMessage, + input: unknown, + orchestrator: StrandsMultiAgent | undefined, operation: MultiAgentStreamState["operation"], activeChildParents: ActiveChildParents, ): MultiAgentStreamState { - const orchestrator = extractOrchestrator(event); const metadata = { "strands.operation": operation, provider: "strands", @@ -331,14 +302,14 @@ function startMultiAgentStream( const parentSpan = orchestrator ? getOnlyChildParent(activeChildParents, orchestrator) : undefined; - const input = processStrandsInputAttachments(event.arguments[0]); + const processedInput = processStrandsInputAttachments(input); const span = parentSpan ? withCurrent(parentSpan, () => startBaseSpan( withSpanInstrumentationName( { event: { - input, + input: processedInput, metadata, }, name: @@ -355,7 +326,7 @@ function startMultiAgentStream( withSpanInstrumentationName( { event: { - input, + input: processedInput, metadata, }, name: @@ -847,19 +818,18 @@ function finalizeMultiAgentStream( state.span.end(); } -function extractAgent( - event: ChannelMessage, -): StrandsAgent | undefined { - const candidate = event.agent ?? event.self; +function extractAgent(agent: unknown, self: unknown): StrandsAgent | undefined { + const candidate = agent ?? self; return isObject(candidate) && typeof candidate.stream === "function" ? (candidate as StrandsAgent) : undefined; } function extractOrchestrator( - event: ChannelMessage, + orchestrator: unknown, + self: unknown, ): StrandsMultiAgent | undefined { - const candidate = event.orchestrator ?? event.self; + const candidate = orchestrator ?? self; return isObject(candidate) && typeof candidate.stream === "function" ? (candidate as StrandsMultiAgent) : undefined; diff --git a/js/src/vendor-sdk-types/ai-sdk-v7-telemetry.ts b/js/src/vendor-sdk-types/ai-sdk-v7-telemetry.ts index b47ebd43f..029c01bf6 100644 --- a/js/src/vendor-sdk-types/ai-sdk-v7-telemetry.ts +++ b/js/src/vendor-sdk-types/ai-sdk-v7-telemetry.ts @@ -12,10 +12,6 @@ export interface AISDKV7TelemetryOptions { functionId?: string; } -export const BRAINTRUST_AI_SDK_V7_OPERATION_KEY = Symbol.for( - "braintrust.ai-sdk.v7.telemetry-operation-key", -); - interface AISDKV7ModelInfo { provider?: string; modelId?: string; @@ -25,14 +21,12 @@ export interface AISDKV7OperationEvent extends AISDKV7TelemetryOptions, AISDKV7ModelInfo { callId: string; operationId: string; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } export interface AISDKV7LanguageModelCallStartEvent extends AISDKV7TelemetryOptions, AISDKV7ModelInfo { callId: string; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -43,7 +37,6 @@ export interface AISDKV7LanguageModelCallEndEvent finishReason?: unknown; responseId?: string; usage?: unknown; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -52,7 +45,6 @@ export interface AISDKV7ObjectStepStartEvent callId: string; promptMessages?: unknown; stepNumber?: number; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -67,7 +59,6 @@ export interface AISDKV7ObjectStepEndEvent response?: unknown; usage?: unknown; warnings?: unknown; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -77,7 +68,6 @@ export interface AISDKV7EmbedStartEvent embedCallId: string; operationId: string; values: unknown[]; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -89,7 +79,6 @@ export interface AISDKV7EmbedEndEvent embeddings?: unknown[]; usage?: unknown; values?: unknown[]; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -99,7 +88,6 @@ export interface AISDKV7RerankStartEvent documents?: unknown[]; query?: string; topN?: number; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -107,7 +95,6 @@ export interface AISDKV7RerankEndEvent extends AISDKV7TelemetryOptions, AISDKV7ModelInfo { callId: string; ranking?: Array<{ index?: number; relevanceScore?: number }>; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -122,7 +109,6 @@ export interface AISDKV7ToolExecutionStartEvent extends AISDKV7TelemetryOptions callId?: string; toolCall: AISDKV7ToolCall; toolContext?: unknown; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -141,7 +127,6 @@ export interface AISDKV7ToolExecutionEndEvent extends AISDKV7TelemetryOptions { success?: boolean; toolCall: AISDKV7ToolCall; toolOutput?: AISDKV7ToolOutput; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -188,7 +173,6 @@ export interface AISDKV7Telemetry { callId: string; toolCallId: string; execute: () => PromiseLike; - [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; }) => PromiseLike; } diff --git a/js/src/wrappers/ai-sdk/telemetry.ts b/js/src/wrappers/ai-sdk/telemetry.ts index 5bb41f885..789029490 100644 --- a/js/src/wrappers/ai-sdk/telemetry.ts +++ b/js/src/wrappers/ai-sdk/telemetry.ts @@ -34,7 +34,6 @@ import type { AISDKV7Telemetry, AISDKV7TelemetryOptions, } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; -import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; import { currentWorkflowAgentWrapperSpan } from "./workflow-agent-context"; import { currentHarnessTurnParent, @@ -145,26 +144,10 @@ export function braintrustAISDKTelemetry(): any { } }; - const explicitOperationKey = (event: unknown): string | undefined => { - if (!isObject(event)) { - return undefined; - } - - const key = (event as { [AI_SDK_V7_OPERATION_KEY]?: unknown })[ - AI_SDK_V7_OPERATION_KEY - ]; - return typeof key === "string" ? key : undefined; - }; - const createOperationKey = ( event: AISDKV7OperationEvent, operationName: string, ): string => { - const explicit = explicitOperationKey(event); - if (explicit) { - return explicit; - } - if (operationName === "WorkflowAgent.stream") { workflowAgentOperationCounter += 1; return `${event.callId}:${workflowAgentOperationCounter}`; @@ -209,11 +192,6 @@ export function braintrustAISDKTelemetry(): any { event: { callId?: unknown } | unknown, mode: "active" | "finish" = "active", ): string | undefined => { - const explicit = explicitOperationKey(event); - if (explicit && operations.has(explicit)) { - return explicit; - } - if (isObject(event)) { const callId = (event as { callId?: unknown }).callId; if (typeof callId === "string") { diff --git a/js/src/wrappers/strands-agent-sdk.test.ts b/js/src/wrappers/strands-agent-sdk.test.ts index f811800d4..66d2ad46c 100644 --- a/js/src/wrappers/strands-agent-sdk.test.ts +++ b/js/src/wrappers/strands-agent-sdk.test.ts @@ -1,15 +1,20 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -const { traceSync } = vi.hoisted(() => ({ - traceSync: vi.fn((fn: () => unknown, _event?: unknown) => fn()), +const { invoke } = vi.hoisted(() => ({ + invoke: vi.fn( + ( + target: Function, + thisArg: unknown, + args: unknown[], + _additional?: unknown, + ) => Reflect.apply(target, thisArg, args), + ), })); vi.mock("../isomorph", () => ({ default: { newTracingChannel: vi.fn(() => ({ - subscribe: vi.fn(), - traceSync, - unsubscribe: vi.fn(), + invoke, })), }, })); @@ -61,14 +66,12 @@ describe("wrapStrandsAgentSDK", () => { expect(result).toMatchObject({ lastMessage: { role: "assistant", content: [{ text: "world" }] }, }); - expect(traceSync).toHaveBeenCalledTimes(2); - expect(traceSync.mock.calls[0][1]).toMatchObject({ - arguments: ["hello", undefined], - self: expect.objectContaining({ name: "assistant" }), - }); - expect(traceSync.mock.calls[1][1]).toMatchObject({ - arguments: ["world", undefined], + expect(invoke).toHaveBeenCalledTimes(2); + expect(invoke.mock.calls[0][2]).toEqual(["hello", undefined]); + expect(invoke.mock.calls[0][3]).toMatchObject({ + agent: expect.objectContaining({ name: "assistant" }), }); + expect(invoke.mock.calls[1][2]).toEqual(["world", undefined]); }); it("wraps Graph and Swarm stream/invoke", async () => { @@ -101,12 +104,14 @@ describe("wrapStrandsAgentSDK", () => { await expect(new wrapped.Swarm().invoke("swarm")).resolves.toMatchObject({ status: "COMPLETED", }); - expect(traceSync).toHaveBeenCalledTimes(2); - expect(traceSync.mock.calls[0][1]).toMatchObject({ - arguments: ["graph", undefined], + expect(invoke).toHaveBeenCalledTimes(2); + expect(invoke.mock.calls[0][2]).toEqual(["graph", undefined]); + expect(invoke.mock.calls[0][3]).toMatchObject({ + orchestrator: expect.anything(), }); - expect(traceSync.mock.calls[1][1]).toMatchObject({ - arguments: ["swarm", undefined], + expect(invoke.mock.calls[1][2]).toEqual(["swarm", undefined]); + expect(invoke.mock.calls[1][3]).toMatchObject({ + orchestrator: expect.anything(), }); }); @@ -126,7 +131,7 @@ describe("wrapStrandsAgentSDK", () => { ) as any; await new wrapped.Agent().invoke("hello"); - expect(traceSync).toHaveBeenCalledTimes(1); + expect(invoke).toHaveBeenCalledTimes(1); }); it("preserves private-field-safe method binding", async () => { diff --git a/js/src/wrappers/strands-agent-sdk.ts b/js/src/wrappers/strands-agent-sdk.ts index 6c6173034..5cf339789 100644 --- a/js/src/wrappers/strands-agent-sdk.ts +++ b/js/src/wrappers/strands-agent-sdk.ts @@ -21,8 +21,8 @@ const WRAPPED_INSTANCE = Symbol.for( ); /** - * Wraps the Strands Agent SDK with Braintrust tracing. The wrapper emits - * diagnostics-channel events; the Strands plugin owns span lifecycle. + * Wraps the Strands Agent SDK with Braintrust tracing. The wrapper invokes + * typed instrumentation channels; the Strands plugin owns span lifecycle. */ export function wrapStrandsAgentSDK(sdk: T): T { if (!sdk || typeof sdk !== "object") { @@ -157,13 +157,11 @@ function wrapAgentInstance(agent: StrandsAgent): StrandsAgent { StrandsInvokeArgs, StrandsInvokeOptions | undefined, ]; - return strandsAgentSDKChannels.agentStream.traceSync( - () => Reflect.apply(value, target, callArgs), - { - agent: proxy, - arguments: callArgs, - self: proxy, - } as never, + return strandsAgentSDKChannels.agentStream.invoke( + value as StrandsAgent["stream"], + target, + callArgs, + { agent: proxy }, ); }; } @@ -219,13 +217,13 @@ function wrapMultiAgentInstance( kind === "graph" ? strandsAgentSDKChannels.graphStream : strandsAgentSDKChannels.swarmStream; - return channel.traceSync( - () => Reflect.apply(value, target, callArgs), + return channel.invoke( + value as StrandsMultiAgent["stream"], + target, + callArgs, { - arguments: callArgs, orchestrator: proxy, - self: proxy, - } as never, + }, ); }; } From bb30b929f7fb9bbcfe78fe2ff51715cb0575d1ad Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 24 Aug 2026 15:31:41 +0000 Subject: [PATCH 4/4] Update PR #2401 --- .../plugins/ai-sdk-plugin.test.ts | 38 +++--- .../instrumentation/plugins/ai-sdk-plugin.ts | 62 ++++++++-- .../plugins/ai-sdk-v7-telemetry.test.ts | 112 ++++++++++++++++++ .../claude-agent-sdk-local-tool-context.ts | 4 +- .../vendor-sdk-types/ai-sdk-v7-telemetry.ts | 16 +++ js/src/wrappers/ai-sdk/telemetry.ts | 22 ++++ 6 files changed, 220 insertions(+), 34 deletions(-) diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts index cf6522e49..0cfd405de 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts @@ -32,6 +32,7 @@ import { } from "./ai-sdk-plugin"; import iso from "../../isomorph"; import { serializeAISDKToolsForLogging } from "../../wrappers/ai-sdk/tool-serialization"; +import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; type MockTracingChannel = { @@ -374,20 +375,7 @@ describe("AISDKPlugin", () => { expect(startEvent).not.toHaveProperty("recordInputs"); }); - it("uses an independent telemetry integration for each dispatcher", async () => { - const telemetryA = { - executeTool: vi.fn(({ execute }) => execute()), - onAbort: vi.fn(), - onStart: vi.fn(), - }; - const telemetryB = { - executeTool: vi.fn(({ execute }) => execute()), - onAbort: vi.fn(), - onStart: vi.fn(), - }; - telemetryMocks.braintrustAISDKTelemetry - .mockReturnValueOnce(telemetryA) - .mockReturnValueOnce(telemetryB); + it("stamps a stable unique operation key on each dispatcher", async () => { const dispatcherA = { executeTool: vi.fn(({ execute }) => execute()), onAbort: vi.fn(), @@ -430,12 +418,22 @@ describe("AISDKPlugin", () => { toolCallId: "tool-b", }); - expect(telemetryA.onStart).toHaveBeenCalledTimes(1); - expect(telemetryA.onAbort).toHaveBeenCalledTimes(1); - expect(telemetryA.executeTool).not.toHaveBeenCalled(); - expect(telemetryB.onStart).toHaveBeenCalledTimes(1); - expect(telemetryB.onAbort).not.toHaveBeenCalled(); - expect(telemetryB.executeTool).toHaveBeenCalledTimes(1); + const runAStart = telemetryMocks.telemetry.onStart?.mock.calls[0]?.[0]; + const runBStart = telemetryMocks.telemetry.onStart?.mock.calls[1]?.[0]; + const runAAbort = telemetryMocks.telemetry.onAbort?.mock.calls[0]?.[0]; + const runBTool = telemetryMocks.telemetry.executeTool?.mock.calls[0]?.[0]; + + expect(runAStart?.[AI_SDK_V7_OPERATION_KEY]).toEqual(expect.any(String)); + expect(runBStart?.[AI_SDK_V7_OPERATION_KEY]).toEqual(expect.any(String)); + expect(runAStart?.[AI_SDK_V7_OPERATION_KEY]).not.toBe( + runBStart?.[AI_SDK_V7_OPERATION_KEY], + ); + expect(runAAbort?.[AI_SDK_V7_OPERATION_KEY]).toBe( + runAStart?.[AI_SDK_V7_OPERATION_KEY], + ); + expect(runBTool?.[AI_SDK_V7_OPERATION_KEY]).toBe( + runBStart?.[AI_SDK_V7_OPERATION_KEY], + ); }); it("preserves existing dispatcher callback return and rejection semantics", async () => { diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index 45ade7d26..d23ff888b 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -73,6 +73,7 @@ import type { AISDKV7Telemetry, AISDKV7TelemetryOptions, } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; +import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; interface AISDKPluginConfig { /** @@ -114,6 +115,7 @@ const AUTO_PATCHED_V7_TELEMETRY_DISPATCHER = Symbol.for( const RUNTIME_DENY_OUTPUT_PATHS = Symbol.for( "braintrust.ai-sdk.deny-output-paths", ); +let aiSDKV7TelemetryOperationCounter = 0; const TRANSPORT_PAYLOAD_ROOT_PATHS = [ "rawResponse", "request", @@ -838,13 +840,18 @@ function subscribeToHarnessContinuation( } function interceptAISDKV7TelemetryDispatcher(): () => void { + const telemetry = braintrustAISDKTelemetry(); return aiSDKChannels.v7CreateTelemetryDispatcher.intercept( (target, thisArg, args) => { const dispatcher = Reflect.apply(target, thisArg, args); const telemetryOptions = args[0]?.telemetry; if (telemetryOptions?.isEnabled !== false) { try { - patchAISDKV7TelemetryDispatcher(dispatcher, telemetryOptions); + patchAISDKV7TelemetryDispatcher( + dispatcher, + telemetry, + telemetryOptions, + ); } catch (error) { debugLogger.error( "Error instrumenting AI SDK v7 telemetry dispatcher:", @@ -859,6 +866,7 @@ function interceptAISDKV7TelemetryDispatcher(): () => void { function patchAISDKV7TelemetryDispatcher( dispatcher: unknown, + telemetry: AISDKV7Telemetry, telemetryOptions?: AISDKV7TelemetryOptions, ): void { if (!isObject(dispatcher)) { @@ -870,7 +878,7 @@ function patchAISDKV7TelemetryDispatcher( return; } dispatcherRecord[AUTO_PATCHED_V7_TELEMETRY_DISPATCHER] = true; - const telemetry = braintrustAISDKTelemetry() as AISDKV7Telemetry; + let operationKey: string | undefined; const telemetryEventFields: AISDKV7TelemetryOptions = {}; if (typeof telemetryOptions?.recordInputs === "boolean") { telemetryEventFields.recordInputs = telemetryOptions.recordInputs; @@ -881,17 +889,48 @@ function patchAISDKV7TelemetryDispatcher( if (typeof telemetryOptions?.functionId === "string") { telemetryEventFields.functionId = telemetryOptions.functionId; } - const hasTelemetryEventFields = Object.keys(telemetryEventFields).length > 0; - - const eventWithTelemetryFields = (event: unknown): unknown => { - if (!isObject(event) || !hasTelemetryEventFields) { + const eventWithOperationKey = (event: unknown): unknown => { + if (!isObject(event)) { return event; } - return { - ...telemetryEventFields, - ...(event as Record), - }; + const eventRecord = event as Record; + const callId = + typeof eventRecord.callId === "string" ? eventRecord.callId : "unknown"; + operationKey ??= `${callId}:${++aiSDKV7TelemetryOperationCounter}`; + + if (Object.keys(telemetryEventFields).length > 0) { + const augmentedEvent = { + ...telemetryEventFields, + ...(event as Record), + }; + try { + Object.defineProperty(augmentedEvent, AI_SDK_V7_OPERATION_KEY, { + configurable: true, + enumerable: false, + value: operationKey, + }); + } catch { + (augmentedEvent as Record)[ + AI_SDK_V7_OPERATION_KEY + ] = operationKey; + } + return augmentedEvent; + } + + try { + Object.defineProperty(eventRecord, AI_SDK_V7_OPERATION_KEY, { + configurable: true, + enumerable: false, + value: operationKey, + }); + return event; + } catch { + return { + ...(event as Record), + [AI_SDK_V7_OPERATION_KEY]: operationKey, + }; + } }; for (const key of AI_SDK_V7_TELEMETRY_CALLBACKS) { @@ -909,7 +948,7 @@ function patchAISDKV7TelemetryDispatcher( try { const braintrustResult = braintrustCallback.call( telemetry, - eventWithTelemetryFields(event) as any, + eventWithOperationKey(event) as any, ); if (isPromiseLike(braintrustResult)) { void Promise.resolve(braintrustResult).catch(() => undefined); @@ -934,6 +973,7 @@ function patchAISDKV7TelemetryDispatcher( }) => braintrustExecuteTool.call(telemetry, { ...args, + ...(operationKey ? { [AI_SDK_V7_OPERATION_KEY]: operationKey } : {}), execute: () => typeof existingExecuteTool === "function" ? existingExecuteTool.call(dispatcher, args) diff --git a/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts b/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts index 517525bec..8741a0e1f 100644 --- a/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts @@ -12,6 +12,7 @@ import { registerWorkflowAgentWrapperSpan, unregisterWorkflowAgentWrapperSpan, } from "../../wrappers/ai-sdk/workflow-agent-context"; +import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; try { configureNode(); @@ -558,6 +559,117 @@ describe("braintrustAISDKTelemetry", () => { ).toHaveLength(0); }); + it("keeps concurrent WorkflowAgent streams with shared callIds separate", async () => { + const telemetry = braintrustAISDKTelemetry(); + const callId = "workflow-agent"; + const runA = "workflow-agent:run-a"; + const runB = "workflow-agent:run-b"; + + telemetry.onStart?.({ + [AI_SDK_V7_OPERATION_KEY]: runA, + callId, + messages: [{ role: "user", content: "First workflow run" }], + operationId: "ai.workflowAgent.stream", + }); + telemetry.onStart?.({ + [AI_SDK_V7_OPERATION_KEY]: runB, + callId, + messages: [{ role: "user", content: "Second workflow run" }], + operationId: "ai.workflowAgent.stream", + }); + + telemetry.onLanguageModelCallStart?.({ + [AI_SDK_V7_OPERATION_KEY]: runA, + callId, + prompt: [{ role: "user", content: "First workflow run" }], + }); + telemetry.onLanguageModelCallEnd?.({ + [AI_SDK_V7_OPERATION_KEY]: runA, + callId, + text: "First answer", + }); + telemetry.onLanguageModelCallStart?.({ + [AI_SDK_V7_OPERATION_KEY]: runB, + callId, + prompt: [{ role: "user", content: "Second workflow run" }], + }); + telemetry.onLanguageModelCallEnd?.({ + [AI_SDK_V7_OPERATION_KEY]: runB, + callId, + text: "Calling get_weather.", + }); + + telemetry.onToolExecutionStart?.({ + [AI_SDK_V7_OPERATION_KEY]: runB, + toolCall: { + toolCallId: "tool-run-b", + toolName: "get_weather", + input: { location: "Vienna, Austria" }, + }, + }); + telemetry.onToolExecutionEnd?.({ + [AI_SDK_V7_OPERATION_KEY]: runB, + output: { condition: "sunny" }, + success: true, + toolCall: { + toolCallId: "tool-run-b", + toolName: "get_weather", + }, + }); + + telemetry.onEnd?.({ + [AI_SDK_V7_OPERATION_KEY]: runA, + callId, + messages: [{ role: "assistant", content: "First answer" }], + operationId: "ai.workflowAgent.stream", + text: "First answer", + }); + telemetry.onEnd?.({ + [AI_SDK_V7_OPERATION_KEY]: runB, + callId, + messages: [{ role: "assistant", content: "Second answer" }], + operationId: "ai.workflowAgent.stream", + text: "Second answer", + }); + + const spans = (await backgroundLogger.drain()) as Array< + Record + >; + const workflowSpans = spans.filter( + (span) => span.span_attributes?.name === "WorkflowAgent.stream", + ); + const firstWorkflow = workflowSpans.find((span) => + JSON.stringify(span.input).includes("First workflow run"), + ); + const secondWorkflow = workflowSpans.find((span) => + JSON.stringify(span.input).includes("Second workflow run"), + ); + const firstModel = spans.find( + (span) => + span.span_attributes?.name === "doGenerate" && + JSON.stringify(span.input).includes("First workflow run"), + ); + const secondModel = spans.find( + (span) => + span.span_attributes?.name === "doGenerate" && + JSON.stringify(span.input).includes("Second workflow run"), + ); + const tool = spans.find( + (span) => span.span_attributes?.name === "get_weather", + ); + + expect(workflowSpans).toHaveLength(2); + expect(firstWorkflow?.output).toMatchObject({ text: "First answer" }); + expect(secondWorkflow?.output).toMatchObject({ text: "Second answer" }); + expect(firstModel?.span_parents).toEqual([firstWorkflow?.span_id]); + expect(secondModel?.span_parents).toEqual([secondWorkflow?.span_id]); + expect(tool).toMatchObject({ + input: { location: "Vienna, Austria" }, + output: { condition: "sunny" }, + span_parents: [secondWorkflow?.span_id], + }); + }); + it("keeps concurrent direct WorkflowAgent telemetry separated without dispatcher keys", async () => { const telemetry = braintrustAISDKTelemetry(); const callId = "workflow-agent"; diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts index 72ae6b5c8..e0c2129c3 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts @@ -1,8 +1,6 @@ import iso from "../../isomorph"; -export type ClaudeLocalToolParentResolver = ( - toolUseId: string, -) => Promise; +type ClaudeLocalToolParentResolver = (toolUseId: string) => Promise; const localToolContextStore = iso.newAsyncLocalStorage(); diff --git a/js/src/vendor-sdk-types/ai-sdk-v7-telemetry.ts b/js/src/vendor-sdk-types/ai-sdk-v7-telemetry.ts index 029c01bf6..b47ebd43f 100644 --- a/js/src/vendor-sdk-types/ai-sdk-v7-telemetry.ts +++ b/js/src/vendor-sdk-types/ai-sdk-v7-telemetry.ts @@ -12,6 +12,10 @@ export interface AISDKV7TelemetryOptions { functionId?: string; } +export const BRAINTRUST_AI_SDK_V7_OPERATION_KEY = Symbol.for( + "braintrust.ai-sdk.v7.telemetry-operation-key", +); + interface AISDKV7ModelInfo { provider?: string; modelId?: string; @@ -21,12 +25,14 @@ export interface AISDKV7OperationEvent extends AISDKV7TelemetryOptions, AISDKV7ModelInfo { callId: string; operationId: string; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } export interface AISDKV7LanguageModelCallStartEvent extends AISDKV7TelemetryOptions, AISDKV7ModelInfo { callId: string; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -37,6 +43,7 @@ export interface AISDKV7LanguageModelCallEndEvent finishReason?: unknown; responseId?: string; usage?: unknown; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -45,6 +52,7 @@ export interface AISDKV7ObjectStepStartEvent callId: string; promptMessages?: unknown; stepNumber?: number; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -59,6 +67,7 @@ export interface AISDKV7ObjectStepEndEvent response?: unknown; usage?: unknown; warnings?: unknown; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -68,6 +77,7 @@ export interface AISDKV7EmbedStartEvent embedCallId: string; operationId: string; values: unknown[]; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -79,6 +89,7 @@ export interface AISDKV7EmbedEndEvent embeddings?: unknown[]; usage?: unknown; values?: unknown[]; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -88,6 +99,7 @@ export interface AISDKV7RerankStartEvent documents?: unknown[]; query?: string; topN?: number; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -95,6 +107,7 @@ export interface AISDKV7RerankEndEvent extends AISDKV7TelemetryOptions, AISDKV7ModelInfo { callId: string; ranking?: Array<{ index?: number; relevanceScore?: number }>; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -109,6 +122,7 @@ export interface AISDKV7ToolExecutionStartEvent extends AISDKV7TelemetryOptions callId?: string; toolCall: AISDKV7ToolCall; toolContext?: unknown; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -127,6 +141,7 @@ export interface AISDKV7ToolExecutionEndEvent extends AISDKV7TelemetryOptions { success?: boolean; toolCall: AISDKV7ToolCall; toolOutput?: AISDKV7ToolOutput; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; [key: string]: unknown; } @@ -173,6 +188,7 @@ export interface AISDKV7Telemetry { callId: string; toolCallId: string; execute: () => PromiseLike; + [BRAINTRUST_AI_SDK_V7_OPERATION_KEY]?: string; }) => PromiseLike; } diff --git a/js/src/wrappers/ai-sdk/telemetry.ts b/js/src/wrappers/ai-sdk/telemetry.ts index 789029490..5bb41f885 100644 --- a/js/src/wrappers/ai-sdk/telemetry.ts +++ b/js/src/wrappers/ai-sdk/telemetry.ts @@ -34,6 +34,7 @@ import type { AISDKV7Telemetry, AISDKV7TelemetryOptions, } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; +import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; import { currentWorkflowAgentWrapperSpan } from "./workflow-agent-context"; import { currentHarnessTurnParent, @@ -144,10 +145,26 @@ export function braintrustAISDKTelemetry(): any { } }; + const explicitOperationKey = (event: unknown): string | undefined => { + if (!isObject(event)) { + return undefined; + } + + const key = (event as { [AI_SDK_V7_OPERATION_KEY]?: unknown })[ + AI_SDK_V7_OPERATION_KEY + ]; + return typeof key === "string" ? key : undefined; + }; + const createOperationKey = ( event: AISDKV7OperationEvent, operationName: string, ): string => { + const explicit = explicitOperationKey(event); + if (explicit) { + return explicit; + } + if (operationName === "WorkflowAgent.stream") { workflowAgentOperationCounter += 1; return `${event.callId}:${workflowAgentOperationCounter}`; @@ -192,6 +209,11 @@ export function braintrustAISDKTelemetry(): any { event: { callId?: unknown } | unknown, mode: "active" | "finish" = "active", ): string | undefined => { + const explicit = explicitOperationKey(event); + if (explicit && operations.has(explicit)) { + return explicit; + } + if (isObject(event)) { const callId = (event as { callId?: unknown }).callId; if (typeof callId === "string") {