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..83a28892e 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -1139,7 +1139,7 @@ function prepareAISDKChildTracing( modelWrapped: boolean; } { const cleanup: Array<() => void> = []; - const patchedModels = new WeakSet(); + const patchedModels = new WeakMap(); const patchedTools = new WeakSet(); let modelWrapped = false; @@ -1151,22 +1151,34 @@ function prepareAISDKChildTracing( !resolvedModel || typeof resolvedModel !== "object" || typeof resolvedModel.doGenerate !== "function" || - patchedModels.has(resolvedModel) || (resolvedModel as { [AUTO_PATCHED_MODEL]?: boolean })[AUTO_PATCHED_MODEL] ) { 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( + Object.getPrototypeOf(resolvedModel), + ) as AISDKLanguageModel; + Object.defineProperties( + wrappedModel, + Object.getOwnPropertyDescriptors(resolvedModel), + ); + Object.defineProperty(wrappedModel, AUTO_PATCHED_MODEL, { + configurable: true, + value: true, + }); - resolvedModel.doGenerate = async function doGeneratePatched( + wrappedModel.doGenerate = async function doGeneratePatched( options: AISDKCallParams, ) { return parentSpan.traced( @@ -1208,7 +1220,7 @@ function prepareAISDKChildTracing( }; if (originalDoStream) { - resolvedModel.doStream = async function doStreamPatched( + wrappedModel.doStream = async function doStreamPatched( options: AISDKCallParams, ) { const span = parentSpan.startSpan({ @@ -1343,17 +1355,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 => { @@ -1443,13 +1446,13 @@ function prepareAISDKChildTracing( }; if (params && typeof params === "object") { - const patchedParamModel = patchModel(params.model); - if ( - typeof params.model === "string" && - patchedParamModel && - typeof patchedParamModel === "object" - ) { + const originalParamModel = params.model; + const patchedParamModel = patchModel(originalParamModel); + if (patchedParamModel && patchedParamModel !== originalParamModel) { params.model = patchedParamModel; + cleanup.push(() => { + params.model = originalParamModel; + }); } patchTools(params.tools); } @@ -1461,25 +1464,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) {