From 88ea54da4846105c85563db831cfe638caf5c830 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 3 Jul 2026 15:26:03 -0400 Subject: [PATCH 1/2] fix(ai-sdk): isolate shared model tracing --- .changeset/calm-llamas-trace.md | 5 ++ .../concurrent-shared-model.span-tree.json | 59 +++++++++++++ .../concurrent-shared-model.span-tree.txt | 33 +++++++ .../scenario.concurrent-shared-model.test.ts | 87 +++++++++++++++++++ .../scenario.concurrent-shared-model.ts | 74 ++++++++++++++++ .../instrumentation/plugins/ai-sdk-plugin.ts | 62 ++++++------- 6 files changed, 285 insertions(+), 35 deletions(-) create mode 100644 .changeset/calm-llamas-trace.md create mode 100644 e2e/scenarios/ai-sdk-instrumentation/__snapshots__/concurrent-shared-model.span-tree.json create mode 100644 e2e/scenarios/ai-sdk-instrumentation/__snapshots__/concurrent-shared-model.span-tree.txt create mode 100644 e2e/scenarios/ai-sdk-instrumentation/scenario.concurrent-shared-model.test.ts create mode 100644 e2e/scenarios/ai-sdk-instrumentation/scenario.concurrent-shared-model.ts diff --git a/.changeset/calm-llamas-trace.md b/.changeset/calm-llamas-trace.md new file mode 100644 index 000000000..2114d88a5 --- /dev/null +++ b/.changeset/calm-llamas-trace.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(ai-sdk): Keep `wrapAISDK` model child spans correctly parented for concurrent calls sharing a model instance diff --git a/e2e/scenarios/ai-sdk-instrumentation/__snapshots__/concurrent-shared-model.span-tree.json b/e2e/scenarios/ai-sdk-instrumentation/__snapshots__/concurrent-shared-model.span-tree.json new file mode 100644 index 000000000..d231512bf --- /dev/null +++ b/e2e/scenarios/ai-sdk-instrumentation/__snapshots__/concurrent-shared-model.span-tree.json @@ -0,0 +1,59 @@ +{ + "span_tree": [ + { + "name": "ai-sdk-concurrent-shared-model-root", + "type": "task", + "children": [ + { + "name": "generateText", + "type": "function", + "children": [ + { + "name": "doGenerate", + "type": "llm", + "children": [], + "input": { + "marker": "MARKER_A" + }, + "output": { + "marker": "MARKER_A" + } + } + ], + "input": { + "marker": "MARKER_A" + }, + "output": { + "marker": "MARKER_A" + } + }, + { + "name": "generateText", + "type": "function", + "children": [ + { + "name": "doGenerate", + "type": "llm", + "children": [], + "input": { + "marker": "MARKER_B" + }, + "output": { + "marker": "MARKER_B" + } + } + ], + "input": { + "marker": "MARKER_B" + }, + "output": { + "marker": "MARKER_B" + } + } + ], + "metadata": { + "scenario": "ai-sdk-concurrent-shared-model" + } + } + ] +} diff --git a/e2e/scenarios/ai-sdk-instrumentation/__snapshots__/concurrent-shared-model.span-tree.txt b/e2e/scenarios/ai-sdk-instrumentation/__snapshots__/concurrent-shared-model.span-tree.txt new file mode 100644 index 000000000..e8953a8c9 --- /dev/null +++ b/e2e/scenarios/ai-sdk-instrumentation/__snapshots__/concurrent-shared-model.span-tree.txt @@ -0,0 +1,33 @@ +span_tree: +└── ai-sdk-concurrent-shared-model-root [task] + metadata: { + "scenario": "ai-sdk-concurrent-shared-model" + } + ├── generateText [function] + │ input: { + │ "marker": "MARKER_A" + │ } + │ output: { + │ "marker": "MARKER_A" + │ } + │ └── doGenerate [llm] + │ input: { + │ "marker": "MARKER_A" + │ } + │ output: { + │ "marker": "MARKER_A" + │ } + └── generateText [function] + input: { + "marker": "MARKER_B" + } + output: { + "marker": "MARKER_B" + } + └── doGenerate [llm] + input: { + "marker": "MARKER_B" + } + output: { + "marker": "MARKER_B" + } diff --git a/e2e/scenarios/ai-sdk-instrumentation/scenario.concurrent-shared-model.test.ts b/e2e/scenarios/ai-sdk-instrumentation/scenario.concurrent-shared-model.test.ts new file mode 100644 index 000000000..25f494da9 --- /dev/null +++ b/e2e/scenarios/ai-sdk-instrumentation/scenario.concurrent-shared-model.test.ts @@ -0,0 +1,87 @@ +import { expect, test } from "vitest"; +import { resolveFileSnapshotPath } from "../../helpers/file-snapshot"; +import type { CapturedLogEvent } from "../../helpers/mock-braintrust-server"; +import { + prepareScenarioDir, + resolveScenarioDir, + withScenarioHarness, +} from "../../helpers/scenario-harness"; +import { matchSpanTreeSnapshot } from "../../helpers/span-tree"; +import { findChildSpans, findLatestSpan } from "../../helpers/trace-selectors"; + +const originalScenarioDir = resolveScenarioDir(import.meta.url); +const scenarioDir = await prepareScenarioDir({ + scenarioDir: originalScenarioDir, +}); +const ROOT_NAME = "ai-sdk-concurrent-shared-model-root"; +const MARKERS = ["MARKER_A", "MARKER_B"] as const; +const spanSnapshotPath = resolveFileSnapshotPath( + import.meta.url, + "concurrent-shared-model.span-tree.json", +); + +function stringified(value: unknown): string { + return JSON.stringify(value) ?? ""; +} + +function markerFrom(value: unknown): string | undefined { + const text = stringified(value); + return MARKERS.find((marker) => text.includes(marker)); +} + +function summarizedEntry(event: CapturedLogEvent) { + const marker = markerFrom(event.input) ?? markerFrom(event.output); + const fields = + marker === undefined + ? { metadata: { scenario: event.metadata?.scenario } } + : { + input: { marker }, + output: markerFrom(event.output) + ? { marker: markerFrom(event.output) } + : undefined, + }; + + return { event, fields }; +} + +test("wrapAISDK keeps doGenerate spans under each concurrent generateText parent when sharing a model", async () => { + await withScenarioHarness(async (harness) => { + await harness.runScenarioDir({ + entry: "scenario.concurrent-shared-model.ts", + scenarioDir, + timeoutMs: 30_000, + }); + + const events = harness.events(); + const root = findLatestSpan(events, ROOT_NAME); + expect(root).toBeDefined(); + + const parents = findChildSpans(events, "generateText", root?.span.id); + expect(parents).toHaveLength(MARKERS.length); + const children = parents.flatMap((parent) => + findChildSpans(events, "doGenerate", parent.span.id), + ); + + await matchSpanTreeSnapshot( + [root, ...parents, ...children].map(summarizedEntry), + spanSnapshotPath, + ); + + for (const marker of MARKERS) { + const parent = parents.find((event) => + stringified(event.input).includes(marker), + ); + expect(parent, `missing generateText parent for ${marker}`).toBeDefined(); + + const markerChildren = findChildSpans( + events, + "doGenerate", + parent?.span.id, + ); + expect(markerChildren, `doGenerate children for ${marker}`).toHaveLength( + 1, + ); + expect(stringified(markerChildren[0]?.input)).toContain(marker); + } + }); +}, 30_000); diff --git a/e2e/scenarios/ai-sdk-instrumentation/scenario.concurrent-shared-model.ts b/e2e/scenarios/ai-sdk-instrumentation/scenario.concurrent-shared-model.ts new file mode 100644 index 000000000..7472c150f --- /dev/null +++ b/e2e/scenarios/ai-sdk-instrumentation/scenario.concurrent-shared-model.ts @@ -0,0 +1,74 @@ +import { wrapAISDK } from "braintrust"; +import * as ai from "ai-sdk-v6"; +import { MockLanguageModelV3 } from "ai-sdk-v6/test"; +import { runMain, runTracedScenario } from "../../helpers/provider-runtime.mjs"; + +const ROOT_NAME = "ai-sdk-concurrent-shared-model-root"; +const SCENARIO_NAME = "ai-sdk-concurrent-shared-model"; +const MARKERS = ["MARKER_A", "MARKER_B"] as const; + +function markerFrom(value: unknown): string { + const text = JSON.stringify(value); + return MARKERS.find((marker) => text.includes(marker)) ?? "MARKER_UNKNOWN"; +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +async function main() { + const { generateText } = wrapAISDK(ai); + const sharedModel = new MockLanguageModelV3({ + provider: "mock-provider", + modelId: "mock-concurrent-model", + doGenerate: async (options) => { + const marker = markerFrom(options); + await sleep(marker === "MARKER_A" ? 75 : 25); + + return { + content: [{ type: "text", text: `response for ${marker}` }], + finishReason: { raw: "stop", unified: "stop" }, + usage: { + inputTokens: { + cacheRead: undefined, + cacheWrite: undefined, + noCache: 5, + total: 5, + }, + outputTokens: { + reasoning: undefined, + text: 3, + total: 3, + }, + }, + }; + }, + }); + + await runTracedScenario({ + callback: async () => { + await Promise.all( + MARKERS.map((marker) => + generateText({ + model: sharedModel, + messages: [ + { + role: "user", + content: `${marker}: write one short sentence about cats.`, + }, + ], + temperature: 0, + maxOutputTokens: 16, + }), + ), + ); + }, + metadata: { + scenario: SCENARIO_NAME, + }, + projectNameBase: "e2e-ai-sdk-concurrent-shared-model", + rootName: ROOT_NAME, + }); +} + +runMain(main); diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index 2c53704ed..90ed57a8e 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -70,7 +70,6 @@ export const DEFAULT_DENY_OUTPUT_PATHS: string[] = [ "steps[].response.headers", ]; -const AUTO_PATCHED_MODEL = Symbol.for("braintrust.ai-sdk.auto-patched-model"); const AUTO_PATCHED_TOOL = Symbol.for("braintrust.ai-sdk.auto-patched-tool"); const AUTO_PATCHED_V7_TELEMETRY_DISPATCHER = Symbol.for( "braintrust.ai-sdk.v7.auto-patched-telemetry-dispatcher", @@ -1139,7 +1138,7 @@ function prepareAISDKChildTracing( modelWrapped: boolean; } { const cleanup: Array<() => void> = []; - const patchedModels = new WeakSet(); + const patchedModels = new WeakMap(); const patchedTools = new WeakSet(); let modelWrapped = false; @@ -1150,23 +1149,24 @@ function prepareAISDKChildTracing( if ( !resolvedModel || typeof resolvedModel !== "object" || - typeof resolvedModel.doGenerate !== "function" || - patchedModels.has(resolvedModel) || - (resolvedModel as { [AUTO_PATCHED_MODEL]?: boolean })[AUTO_PATCHED_MODEL] + typeof resolvedModel.doGenerate !== "function" ) { return resolvedModel; } - patchedModels.add(resolvedModel); - (resolvedModel as { [AUTO_PATCHED_MODEL]?: boolean })[AUTO_PATCHED_MODEL] = - true; + const existingWrappedModel = patchedModels.get(resolvedModel); + if (existingWrappedModel) { + return existingWrappedModel; + } + modelWrapped = true; const originalDoGenerate = resolvedModel.doGenerate; const originalDoStream = resolvedModel.doStream; const baseMetadata = buildAISDKChildMetadata(resolvedModel); + const wrappedModel = Object.create(resolvedModel) as AISDKLanguageModel; - resolvedModel.doGenerate = async function doGeneratePatched( + wrappedModel.doGenerate = async function doGeneratePatched( options: AISDKCallParams, ) { return parentSpan.traced( @@ -1208,7 +1208,7 @@ function prepareAISDKChildTracing( }; if (originalDoStream) { - resolvedModel.doStream = async function doStreamPatched( + wrappedModel.doStream = async function doStreamPatched( options: AISDKCallParams, ) { const span = parentSpan.startSpan({ @@ -1343,17 +1343,8 @@ function prepareAISDKChildTracing( }; } - cleanup.push(() => { - resolvedModel.doGenerate = originalDoGenerate; - if (originalDoStream) { - resolvedModel.doStream = originalDoStream; - } - delete (resolvedModel as { [AUTO_PATCHED_MODEL]?: boolean })[ - AUTO_PATCHED_MODEL - ]; - }); - - return resolvedModel; + patchedModels.set(resolvedModel, wrappedModel); + return wrappedModel; }; const patchTool = (tool: AISDKTool, name: string): void => { @@ -1444,11 +1435,7 @@ function prepareAISDKChildTracing( if (params && typeof params === "object") { const patchedParamModel = patchModel(params.model); - if ( - typeof params.model === "string" && - patchedParamModel && - typeof patchedParamModel === "object" - ) { + if (patchedParamModel && patchedParamModel !== params.model) { params.model = patchedParamModel; } patchTools(params.tools); @@ -1461,25 +1448,30 @@ function prepareAISDKChildTracing( }; if (selfRecord.model !== undefined) { - const patchedSelfModel = patchModel(selfRecord.model); - if ( - typeof selfRecord.model === "string" && - patchedSelfModel && - typeof patchedSelfModel === "object" - ) { + const originalSelfModel = selfRecord.model; + const patchedSelfModel = patchModel(originalSelfModel); + if (patchedSelfModel && patchedSelfModel !== originalSelfModel) { selfRecord.model = patchedSelfModel; + cleanup.push(() => { + selfRecord.model = originalSelfModel; + }); } } if (selfRecord.settings && typeof selfRecord.settings === "object") { if (selfRecord.settings.model !== undefined) { - const patchedSettingsModel = patchModel(selfRecord.settings.model); + const originalSettingsModel = selfRecord.settings.model; + const patchedSettingsModel = patchModel(originalSettingsModel); if ( - typeof selfRecord.settings.model === "string" && patchedSettingsModel && - typeof patchedSettingsModel === "object" + patchedSettingsModel !== originalSettingsModel ) { selfRecord.settings.model = patchedSettingsModel; + cleanup.push(() => { + if (selfRecord.settings) { + selfRecord.settings.model = originalSettingsModel; + } + }); } } if (selfRecord.settings.tools !== undefined) { From c6220869de76a1747fa4f4a3e1a1639fad98fe86 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 6 Jul 2026 12:24:59 +0200 Subject: [PATCH 2/2] fix double wrapping of models --- .../instrumentation/plugins/ai-sdk-plugin.ts | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index 90ed57a8e..83a28892e 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -70,6 +70,7 @@ export const DEFAULT_DENY_OUTPUT_PATHS: string[] = [ "steps[].response.headers", ]; +const AUTO_PATCHED_MODEL = Symbol.for("braintrust.ai-sdk.auto-patched-model"); const AUTO_PATCHED_TOOL = Symbol.for("braintrust.ai-sdk.auto-patched-tool"); const AUTO_PATCHED_V7_TELEMETRY_DISPATCHER = Symbol.for( "braintrust.ai-sdk.v7.auto-patched-telemetry-dispatcher", @@ -1149,7 +1150,8 @@ function prepareAISDKChildTracing( if ( !resolvedModel || typeof resolvedModel !== "object" || - typeof resolvedModel.doGenerate !== "function" + typeof resolvedModel.doGenerate !== "function" || + (resolvedModel as { [AUTO_PATCHED_MODEL]?: boolean })[AUTO_PATCHED_MODEL] ) { return resolvedModel; } @@ -1164,7 +1166,17 @@ function prepareAISDKChildTracing( const originalDoGenerate = resolvedModel.doGenerate; const originalDoStream = resolvedModel.doStream; const baseMetadata = buildAISDKChildMetadata(resolvedModel); - const wrappedModel = Object.create(resolvedModel) as AISDKLanguageModel; + const wrappedModel = Object.create( + Object.getPrototypeOf(resolvedModel), + ) as AISDKLanguageModel; + Object.defineProperties( + wrappedModel, + Object.getOwnPropertyDescriptors(resolvedModel), + ); + Object.defineProperty(wrappedModel, AUTO_PATCHED_MODEL, { + configurable: true, + value: true, + }); wrappedModel.doGenerate = async function doGeneratePatched( options: AISDKCallParams, @@ -1434,9 +1446,13 @@ function prepareAISDKChildTracing( }; if (params && typeof params === "object") { - const patchedParamModel = patchModel(params.model); - if (patchedParamModel && patchedParamModel !== params.model) { + const originalParamModel = params.model; + const patchedParamModel = patchModel(originalParamModel); + if (patchedParamModel && patchedParamModel !== originalParamModel) { params.model = patchedParamModel; + cleanup.push(() => { + params.model = originalParamModel; + }); } patchTools(params.tools); }