fix(ai-sdk): isolate shared model tracing

.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

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"
+      }
+    }
+  ]
+}

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"
+            }

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);

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<void> {
+  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);