From 84e94a85612053d4f9c1842563c3c7792f00f984 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 15 Jun 2026 16:30:00 +0200 Subject: [PATCH 1/6] fix: Fix TTFT in AI SDK v6 --- .../plugins/ai-sdk-plugin.streaming.test.ts | 147 ++++++++++++++++++ .../instrumentation/plugins/ai-sdk-plugin.ts | 126 ++++++++++++++- 2 files changed, 267 insertions(+), 6 deletions(-) create mode 100644 js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts new file mode 100644 index 000000000..e74d2e19d --- /dev/null +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts @@ -0,0 +1,147 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +/* eslint-disable @typescript-eslint/consistent-type-assertions */ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + test, +} from "vitest"; +import * as ai from "ai"; +import { configureNode } from "../../node/config"; +import { + _exportsForTestingOnly, + initLogger, + Logger, + TestBackgroundLogger, +} from "../../logger"; +import { wrapAISDK } from "../../wrappers/ai-sdk"; + +try { + configureNode(); +} catch {} + +describe("AI SDK streaming instrumentation", () => { + let backgroundLogger: TestBackgroundLogger; + let _logger: Logger; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + _logger = initLogger({ + projectName: "ai-sdk-plugin.streaming.test.ts", + projectId: "test-project-id", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + test("streamText time_to_first_token ignores AI SDK v6 framing chunks", async () => { + expect(await backgroundLogger.drain()).toHaveLength(0); + + const requestDelayMs = 40; + const contentDelayMs = 80; + const sleep = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + let sentContent = false; + const model = { + specificationVersion: "v3", + provider: "mock-delayed-provider", + modelId: "mock-delayed-model", + supportedUrls: {}, + doGenerate: async () => { + throw new Error("doGenerate should not be called"); + }, + doStream: async () => { + await sleep(requestDelayMs); + + return { + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "stream-start", warnings: [] }); + controller.enqueue({ + type: "response-metadata", + id: "response-id", + modelId: "mock-delayed-model", + timestamp: new Date(0), + }); + }, + async pull(controller) { + if (sentContent) { + controller.close(); + return; + } + + sentContent = true; + await sleep(contentDelayMs); + controller.enqueue({ type: "text-start", id: "delayed-text" }); + controller.enqueue({ + type: "text-delta", + id: "delayed-text", + delta: "DELAYED", + }); + controller.enqueue({ type: "text-end", id: "delayed-text" }); + controller.enqueue({ + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { + total: 1, + noCache: 1, + cacheRead: 0, + cacheWrite: 0, + }, + outputTokens: { + total: 1, + text: 1, + reasoning: 0, + }, + }, + }); + }, + }), + warnings: [], + }; + }, + } as any; + + const wrappedAI = wrapAISDK(ai); + const result = wrappedAI.streamText({ + model, + prompt: "Reply with exactly DELAYED.", + maxOutputTokens: 16, + }); + + let fullText = ""; + for await (const chunk of result.textStream) { + fullText += chunk; + } + await result.text; + + expect(fullText).toBe("DELAYED"); + + const spans = (await backgroundLogger.drain()) as any[]; + const streamTextSpan = spans.find( + (s) => s?.span_attributes?.name === "streamText", + ); + const doStreamSpan = spans.find( + (s) => s?.span_attributes?.name === "doStream", + ); + const minimumExpectedTTFT = (requestDelayMs + contentDelayMs) / 1000 / 2; + + expect(streamTextSpan?.metrics?.time_to_first_token).toBeGreaterThanOrEqual( + minimumExpectedTTFT, + ); + expect(doStreamSpan?.metrics?.time_to_first_token).toBeGreaterThanOrEqual( + minimumExpectedTTFT, + ); + expect(streamTextSpan?.output?.text).toBe("DELAYED"); + expect(doStreamSpan?.output?.text).toBe("DELAYED"); + }); +}); diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index 63f2fc0d6..4552edd62 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -6,7 +6,11 @@ import { unsubscribeAll, } from "../core/channel-tracing"; import { isAsyncIterable } from "../core/stream-patcher"; -import { SpanTypeAttribute, isPromiseLike } from "../../../util/index"; +import { + SpanTypeAttribute, + isObject, + isPromiseLike, +} from "../../../util/index"; import { getCurrentUnixTimestamp } from "../../util"; import { Attachment, type Span, withCurrent } from "../../logger"; import { @@ -1103,10 +1107,10 @@ function prepareAISDKChildTracing( }, }); + const streamStartTime = getCurrentUnixTimestamp(); const result = await withCurrent(span, () => Reflect.apply(originalDoStream, resolvedModel, [options]), ); - const streamStartTime = getCurrentUnixTimestamp(); let firstChunkTime: number | undefined; const output: Record = {}; let text = ""; @@ -1116,7 +1120,10 @@ function prepareAISDKChildTracing( const transformStream = new TransformStream({ transform(chunk: AISDKModelStreamChunk, controller) { - if (firstChunkTime === undefined) { + if ( + firstChunkTime === undefined && + isAISDKContentStreamChunk(chunk) + ) { firstChunkTime = getCurrentUnixTimestamp(); } @@ -1365,6 +1372,107 @@ function finalizeAISDKChildTracing(event?: { [key: string]: unknown }): void { } } +function extractAISDKStreamPart(chunk: unknown): unknown { + if (!isObject(chunk) || !isObject(chunk.part)) { + return chunk; + } + + return chunk.part; +} + +function stringContent(value: unknown): string | undefined { + return typeof value === "string" && value.length > 0 ? value : undefined; +} + +function rawValueHasAISDKContent(value: unknown): boolean { + if (value === undefined || value === null) { + return false; + } + + if (typeof value === "string") { + return value.length > 0; + } + + if (!isObject(value)) { + return true; + } + + if ( + stringContent(value.content) || + stringContent(value.text) || + (isObject(value.delta) && stringContent(value.delta.content)) || + (Array.isArray(value.choices) && + value.choices.some( + (choice) => + isObject(choice) && + isObject(choice.delta) && + stringContent(choice.delta.content), + )) + ) { + return true; + } + + const type = value.type; + return ( + typeof type === "string" && + type !== "start" && + type !== "stream-start" && + type !== "response-metadata" + ); +} + +function isAISDKContentStreamChunk(chunk: unknown): boolean { + const part = extractAISDKStreamPart(chunk); + if (typeof part === "string") { + return part.length > 0; + } + + if (!isObject(part) || typeof part.type !== "string") { + return false; + } + + switch (part.type) { + case "text-delta": + return ( + stringContent(part.textDelta) !== undefined || + stringContent(part.delta) !== undefined || + stringContent(part.text) !== undefined || + stringContent(part.content) !== undefined + ); + case "reasoning-delta": + return ( + stringContent(part.delta) !== undefined || + stringContent(part.text) !== undefined || + stringContent(part.content) !== undefined + ); + case "tool-input-start": + case "tool-call": + case "tool-result": + case "object": + return true; + case "tool-input-delta": + case "tool-call-delta": + return ( + stringContent(part.delta) !== undefined || + stringContent(part.text) !== undefined || + stringContent(part.content) !== undefined + ); + case "raw": + return rawValueHasAISDKContent(part.rawValue); + default: + return false; + } +} + +function isAISDKContentAsyncIterableChunk(chunk: unknown): boolean { + if (isAISDKContentStreamChunk(chunk)) { + return true; + } + + const part = extractAISDKStreamPart(chunk); + return isObject(part) && typeof part.type !== "string"; +} + function patchAISDKStreamingResult(args: { defaultDenyOutputPaths: string[]; endEvent: { denyOutputPaths?: string[]; [key: string]: unknown }; @@ -1387,7 +1495,10 @@ function patchAISDKStreamingResult(args: { const wrappedBaseStream = resultRecord.baseStream.pipeThrough( new TransformStream({ transform(chunk, controller) { - if (firstChunkTime === undefined) { + if ( + firstChunkTime === undefined && + isAISDKContentStreamChunk(chunk) + ) { firstChunkTime = getCurrentUnixTimestamp(); } controller.enqueue(chunk); @@ -1441,8 +1552,11 @@ function patchAISDKStreamingResult(args: { let firstChunkTime: number | undefined; const wrappedStream = createPatchedAsyncIterable(streamField.stream, { - onChunk: () => { - if (firstChunkTime === undefined) { + onChunk: (chunk) => { + if ( + firstChunkTime === undefined && + isAISDKContentAsyncIterableChunk(chunk) + ) { firstChunkTime = getCurrentUnixTimestamp(); } }, From de4e33b0278476ad4332d5a97ea495c7b9b7a8f8 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 15 Jun 2026 16:37:05 +0200 Subject: [PATCH 2/6] cs --- .changeset/hungry-hounds-yell.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/hungry-hounds-yell.md diff --git a/.changeset/hungry-hounds-yell.md b/.changeset/hungry-hounds-yell.md new file mode 100644 index 000000000..1621c8191 --- /dev/null +++ b/.changeset/hungry-hounds-yell.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix: Fix TTFT in AI SDK v6 From 3edfd3931a823d220673a8916e897752bcb2416c Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 15 Jun 2026 16:38:42 +0200 Subject: [PATCH 3/6] fix --- .../plugins/ai-sdk-plugin.streaming.test.ts | 8 +++ .../instrumentation/plugins/ai-sdk-plugin.ts | 59 ++++++++++++++----- 2 files changed, 53 insertions(+), 14 deletions(-) diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts index e74d2e19d..6c5fe0139 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts @@ -71,6 +71,13 @@ describe("AI SDK streaming instrumentation", () => { modelId: "mock-delayed-model", timestamp: new Date(0), }); + controller.enqueue({ + type: "raw", + rawValue: { + type: "response.created", + response: { id: "response-id" }, + }, + }); }, async pull(controller) { if (sentContent) { @@ -115,6 +122,7 @@ describe("AI SDK streaming instrumentation", () => { const result = wrappedAI.streamText({ model, prompt: "Reply with exactly DELAYED.", + includeRawChunks: true, maxOutputTokens: 16, }); diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index 4552edd62..a482cbbd2 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -1111,6 +1111,8 @@ function prepareAISDKChildTracing( const result = await withCurrent(span, () => Reflect.apply(originalDoStream, resolvedModel, [options]), ); + // firstChunkTime !== streamStartTime because the first few chunks may be actual bookkeeping stuff for the SDK + // instead of actual LLM output that can be streamed to users. let firstChunkTime: number | undefined; const output: Record = {}; let text = ""; @@ -1397,28 +1399,56 @@ function rawValueHasAISDKContent(value: unknown): boolean { return true; } + // Raw streams also include provider lifecycle events; only content shapes count. + const delta = value.delta; if ( stringContent(value.content) || stringContent(value.text) || - (isObject(value.delta) && stringContent(value.delta.content)) || + stringContent(value.delta) || + stringContent(value.arguments) || + stringContent(value.partial_json) || + (isObject(delta) && + (stringContent(delta.content) || + stringContent(delta.text) || + stringContent(delta.arguments) || + stringContent(delta.partial_json) || + stringContent(delta.thinking))) || (Array.isArray(value.choices) && - value.choices.some( - (choice) => - isObject(choice) && - isObject(choice.delta) && - stringContent(choice.delta.content), - )) + value.choices.some((choice) => { + if (!isObject(choice) || !isObject(choice.delta)) { + return false; + } + + const choiceDelta = choice.delta; + if ( + stringContent(choiceDelta.content) || + stringContent(choiceDelta.text) + ) { + return true; + } + + if ( + isObject(choiceDelta.function_call) && + stringContent(choiceDelta.function_call.arguments) + ) { + return true; + } + + return ( + Array.isArray(choiceDelta.tool_calls) && + choiceDelta.tool_calls.some( + (toolCall) => + isObject(toolCall) && + isObject(toolCall.function) && + stringContent(toolCall.function.arguments), + ) + ); + })) ) { return true; } - const type = value.type; - return ( - typeof type === "string" && - type !== "start" && - type !== "stream-start" && - type !== "response-metadata" - ); + return false; } function isAISDKContentStreamChunk(chunk: unknown): boolean { @@ -1431,6 +1461,7 @@ function isAISDKContentStreamChunk(chunk: unknown): boolean { return false; } + // TTFT should ignore framing like stream-start, metadata, and text-start. switch (part.type) { case "text-delta": return ( From 36eeab6af2e1c99793c47aaaea5da4e6f68637e2 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 16 Jun 2026 14:25:52 +0200 Subject: [PATCH 4/6] fix --- .../plugins/ai-sdk-plugin.streaming.test.ts | 65 ++++++++++++++++++- .../instrumentation/plugins/ai-sdk-plugin.ts | 2 + 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts index 6c5fe0139..3d54bd2dc 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts @@ -17,11 +17,15 @@ import { TestBackgroundLogger, } from "../../logger"; import { wrapAISDK } from "../../wrappers/ai-sdk"; +import { aiSDKChannels } from "./ai-sdk-channels"; try { configureNode(); } catch {} +const sleep = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + describe("AI SDK streaming instrumentation", () => { let backgroundLogger: TestBackgroundLogger; let _logger: Logger; @@ -47,8 +51,6 @@ describe("AI SDK streaming instrumentation", () => { const requestDelayMs = 40; const contentDelayMs = 80; - const sleep = (ms: number) => - new Promise((resolve) => setTimeout(resolve, ms)); let sentContent = false; const model = { specificationVersion: "v3", @@ -152,4 +154,63 @@ describe("AI SDK streaming instrumentation", () => { expect(streamTextSpan?.output?.text).toBe("DELAYED"); expect(doStreamSpan?.output?.text).toBe("DELAYED"); }); + + test("streamText time_to_first_token counts streamed tool-call arguments", async () => { + expect(await backgroundLogger.drain()).toHaveLength(0); + + const contentDelayMs = 80; + let sentContent = false; + const result = (await aiSDKChannels.streamText.tracePromise( + async () => ({ + baseStream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "stream-start", warnings: [] }); + controller.enqueue({ type: "text-start", id: "ignored-text" }); + }, + async pull(controller) { + if (sentContent) { + controller.close(); + return; + } + + sentContent = true; + await sleep(contentDelayMs); + controller.enqueue({ + type: "tool-call-delta", + toolCallType: "function", + toolCallId: "call-1", + toolName: "lookup", + argsTextDelta: '{"query"', + }); + controller.close(); + }, + }), + }), + { + arguments: [ + { + model: "mock-tool-model", + prompt: "Call the lookup tool.", + }, + ], + } as any, + )) as any; + + const reader = result.baseStream.getReader(); + while (true) { + const { done } = await reader.read(); + if (done) { + break; + } + } + + const spans = (await backgroundLogger.drain()) as any[]; + const streamTextSpan = spans.find( + (s) => s?.span_attributes?.name === "streamText", + ); + + expect(streamTextSpan?.metrics?.time_to_first_token).toBeGreaterThanOrEqual( + contentDelayMs / 1000 / 2, + ); + }); }); diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index a482cbbd2..a27bd5f55 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -1484,6 +1484,8 @@ function isAISDKContentStreamChunk(chunk: unknown): boolean { case "tool-input-delta": case "tool-call-delta": return ( + stringContent(part.argsTextDelta) !== undefined || + stringContent(part.inputTextDelta) !== undefined || stringContent(part.delta) !== undefined || stringContent(part.text) !== undefined || stringContent(part.content) !== undefined From b4232d4bf0944a8bf77ed88b18454b8a229eb101 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 16 Jun 2026 14:49:32 +0200 Subject: [PATCH 5/6] fix --- .../plugins/ai-sdk-plugin.streaming.test.ts | 15 +++++++++------ js/src/instrumentation/plugins/ai-sdk-plugin.ts | 3 +-- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts index 3d54bd2dc..90497a2dc 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts @@ -155,7 +155,7 @@ describe("AI SDK streaming instrumentation", () => { expect(doStreamSpan?.output?.text).toBe("DELAYED"); }); - test("streamText time_to_first_token counts streamed tool-call arguments", async () => { + test("streamText time_to_first_token counts streamed tool input arguments", async () => { expect(await backgroundLogger.drain()).toHaveLength(0); const contentDelayMs = 80; @@ -166,6 +166,11 @@ describe("AI SDK streaming instrumentation", () => { start(controller) { controller.enqueue({ type: "stream-start", warnings: [] }); controller.enqueue({ type: "text-start", id: "ignored-text" }); + controller.enqueue({ + type: "tool-input-start", + id: "call-1", + toolName: "lookup", + }); }, async pull(controller) { if (sentContent) { @@ -176,11 +181,9 @@ describe("AI SDK streaming instrumentation", () => { sentContent = true; await sleep(contentDelayMs); controller.enqueue({ - type: "tool-call-delta", - toolCallType: "function", - toolCallId: "call-1", - toolName: "lookup", - argsTextDelta: '{"query"', + type: "tool-input-delta", + id: "call-1", + inputTextDelta: '{"query"', }); controller.close(); }, diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index a27bd5f55..07c25a663 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -1461,7 +1461,7 @@ function isAISDKContentStreamChunk(chunk: unknown): boolean { return false; } - // TTFT should ignore framing like stream-start, metadata, and text-start. + // TTFT should ignore framing like stream-start, metadata, and text/tool starts. switch (part.type) { case "text-delta": return ( @@ -1476,7 +1476,6 @@ function isAISDKContentStreamChunk(chunk: unknown): boolean { stringContent(part.text) !== undefined || stringContent(part.content) !== undefined ); - case "tool-input-start": case "tool-call": case "tool-result": case "object": From 6a6838998cfe00060cfe15b2b1a099b07e88cc96 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 16 Jun 2026 18:51:44 +0200 Subject: [PATCH 6/6] file, but no tool result --- js/src/instrumentation/plugins/ai-sdk-plugin.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index 07c25a663..343f25222 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -1477,8 +1477,8 @@ function isAISDKContentStreamChunk(chunk: unknown): boolean { stringContent(part.content) !== undefined ); case "tool-call": - case "tool-result": case "object": + case "file": return true; case "tool-input-delta": case "tool-call-delta":