fix(vercel): nested runs being unflattened, incorrect inheritance beh…

…avior of user runId / traceable (#1383)

Fixes vercel/ai#4223

js/package.json

@@ -109,20 +109,21 @@
     "uuid": "^10.0.0"
   },
   "devDependencies": {
-    "@ai-sdk/openai": "^0.0.68",
+    "@ai-sdk/openai": "^1.0.13",
     "@babel/preset-env": "^7.22.4",
     "@faker-js/faker": "^8.4.1",
     "@jest/globals": "^29.5.0",
     "@langchain/core": "^0.3.14",
     "@langchain/langgraph": "^0.2.20",
     "@langchain/openai": "^0.3.11",
+    "@opentelemetry/api": "^1.9.0",
     "@opentelemetry/sdk-trace-base": "^1.26.0",
     "@opentelemetry/sdk-trace-node": "^1.26.0",
     "@tsconfig/recommended": "^1.0.2",
     "@types/jest": "^29.5.1",
     "@typescript-eslint/eslint-plugin": "^5.59.8",
     "@typescript-eslint/parser": "^5.59.8",
-    "ai": "^3.4.17",
+    "ai": "^4.0.27",
     "babel-jest": "^29.5.0",
     "cross-env": "^7.0.3",
     "dotenv": "^16.1.3",

js/src/tests/vercel.int.test.ts

@@ -102,7 +102,7 @@ test("generateText with image", async () => {
   expect(storedRun.id).toEqual(runId);
 });
 
-test.skip("streamText", async () => {
+test("streamText", async () => {
   const runId = uuid();
   const result = await streamText({
     model: openai("gpt-4o-mini"),
@@ -234,9 +234,7 @@ test("traceable", async () => {
         async () => {
           return "bar";
         },
-        {
-          name: "foo",
-        }
+        { name: "foo" }
       );
 
       await foo();
@@ -254,6 +252,64 @@ test("traceable", async () => {
   expect(storedRun.outputs).toEqual(result);
 });
 
+test("nested generateText", async () => {
+  const runId = uuid();
+  const childRunId = uuid();
+
+  await generateText({
+    model: openai("gpt-4o-mini"),
+    messages: [
+      {
+        role: "user",
+        content: "What are my orders and where are they? My user ID is 123",
+      },
+    ],
+    tools: {
+      listOrders: tool({
+        description: "list all orders",
+        parameters: z.object({ userId: z.string() }),
+        execute: async ({ userId }) =>
+          `User ${userId} has the following orders: 1`,
+      }),
+      viewTrackingInformation: tool({
+        description: "view tracking information for a specific order",
+        parameters: z.object({ orderId: z.string() }),
+        execute: async ({ orderId }) =>
+          await generateText({
+            model: openai("gpt-4o-mini"),
+            experimental_telemetry: AISDKExporter.getSettings({
+              isEnabled: true,
+              runId: childRunId,
+            }),
+            messages: [
+              {
+                role: "user",
+                content: `Generate a random tracking information, include order ID ${orderId}`,
+              },
+            ],
+          }),
+      }),
+    },
+    experimental_telemetry: AISDKExporter.getSettings({
+      isEnabled: true,
+      runId,
+      functionId: "functionId",
+      metadata: { userId: "123", language: "english" },
+    }),
+    maxSteps: 10,
+  });
+
+  await provider.forceFlush();
+  await waitUntilRunFound(client, runId, true);
+
+  const storedRun = await client.readRun(runId);
+  expect(storedRun.id).toEqual(runId);
+
+  await waitUntilRunFound(client, childRunId, true);
+  const storedChildRun = await client.readRun(childRunId);
+  expect(storedChildRun.id).toEqual(childRunId);
+});
+
 afterAll(async () => {
   await provider.shutdown();
 });

js/src/tests/vercel.test.ts

@@ -1,6 +1,7 @@
 import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
 import { BatchSpanProcessor } from "@opentelemetry/sdk-trace-base";
-
+import { context, trace } from "@opentelemetry/api";
+import { v4 as uuidv4 } from "uuid";
 import {
   generateText,
   streamText,
@@ -940,3 +941,117 @@ test("traceable", async () => {
     },
   });
 });
+
+test("unrelated spans around", async () => {
+  const tracer = provider.getTracer("test");
+
+  const inner = async () => {
+    const span = tracer.startSpan("inner-unrelated");
+    const ctx = trace.setSpan(context.active(), span);
+
+    await context.with(ctx, async () => {
+      await generateText({
+        model: new MockMultiStepLanguageModelV1({
+          provider: "inner-model",
+          doGenerate: async () => {
+            return {
+              rawCall: { rawPrompt: null, rawSettings: {} },
+              finishReason: "stop",
+              usage: { promptTokens: 10, completionTokens: 20 },
+              text: `Hello, world!`,
+            };
+          },
+        }),
+        messages: [{ role: "user", content: "Hello" }],
+        experimental_telemetry: AISDKExporter.getSettings({
+          isEnabled: true,
+          metadata: { userId: "123" },
+        }),
+      });
+    });
+
+    span.end();
+  };
+
+  const rootRunId = uuidv4();
+
+  const outer = async () => {
+    const span = tracer.startSpan("outer-unrelated");
+    const ctx = trace.setSpan(context.active(), span);
+
+    const model = new MockMultiStepLanguageModelV1({
+      provider: "outer-model",
+      doGenerate: async () => {
+        if (model.generateStep === 0) {
+          return {
+            rawCall: { rawPrompt: null, rawSettings: {} },
+            finishReason: "stop",
+            usage: { promptTokens: 10, completionTokens: 20 },
+            toolCalls: [
+              {
+                toolCallType: "function",
+                toolName: "callInner",
+                toolCallId: "tool-id",
+                args: JSON.stringify({}),
+              },
+            ],
+          };
+        }
+
+        return {
+          rawCall: { rawPrompt: null, rawSettings: {} },
+          finishReason: "stop",
+          usage: { promptTokens: 10, completionTokens: 20 },
+          text: `Hello, world!`,
+        };
+      },
+    });
+
+    await context.with(ctx, async () => {
+      await generateText({
+        model,
+        tools: {
+          callInner: tool({
+            description: "call inner",
+            parameters: z.object({}),
+            execute: inner,
+          }),
+        },
+        messages: [{ role: "user", content: "Nested call" }],
+        experimental_telemetry: AISDKExporter.getSettings({
+          isEnabled: true,
+          runName: "outer",
+          runId: rootRunId,
+        }),
+        maxSteps: 10,
+      });
+    });
+
+    span.end();
+  };
+
+  await outer();
+  await provider.forceFlush();
+
+  const actual = getAssumedTreeFromCalls(callSpy.mock.calls);
+  expect(actual).toMatchObject({
+    nodes: [
+      "outer:0",
+      "outer-model:1",
+      "callInner:2",
+      "outer-model:3",
+      "inner-model:4",
+      "inner-model:5",
+    ],
+    edges: [
+      ["outer:0", "outer-model:1"],
+      ["outer:0", "callInner:2"],
+      ["outer:0", "outer-model:3"],
+      ["callInner:2", "inner-model:4"],
+      ["inner-model:4", "inner-model:5"],
+    ],
+    data: {
+      "outer:0": { id: rootRunId },
+    },
+  });
+});

js/src/tests/wrapped_ai_sdk.int.test.ts

@@ -35,7 +35,7 @@ test("AI SDK generateText with a tool", async () => {
           JSON.stringify(["pasta", "tomato", "cheese", "onions"]),
       }),
     },
-    maxToolRoundtrips: 2,
+    maxSteps: 2,
   });
   DEBUG && console.log(text);
 });