diff --git a/e2e/scenarios/ai-sdk-instrumentation/assertions.ts b/e2e/scenarios/ai-sdk-instrumentation/assertions.ts index 8dbef6c8..2bcc5b64 100644 --- a/e2e/scenarios/ai-sdk-instrumentation/assertions.ts +++ b/e2e/scenarios/ai-sdk-instrumentation/assertions.ts @@ -645,6 +645,7 @@ export function defineAISDKInstrumentationAssertions(options: { agentSpanName?: AgentSpanName; name: string; runScenario: RunAISDKScenario; + sdkMajorVersion: number; snapshotName: string; supportsAttachmentScenario: boolean; supportsDenyOutputOverrideScenario: boolean; @@ -894,6 +895,32 @@ export function defineAISDKInstrumentationAssertions(options: { }); } + if (options.sdkMajorVersion >= 4) { + test( + "captures sync streamText()/streamObject() paths in v4+", + testConfig, + () => { + const root = findLatestSpan(events, ROOT_NAME); + const streamTrace = findStreamTrace(events); + + expectOperationParentedByRoot(streamTrace.operation, root); + expectAISDKParentSpan(streamTrace.parent); + expect(operationName(streamTrace.operation)).toBe("stream"); + expect(streamTrace.parent?.span.name).toBe("streamText"); + + if (options.supportsStreamObject) { + const streamObjectTrace = findStreamObjectTrace(events); + expectOperationParentedByRoot(streamObjectTrace.operation, root); + expectAISDKParentSpan(streamObjectTrace.parent); + expect(operationName(streamObjectTrace.operation)).toBe( + "stream-object", + ); + expect(streamObjectTrace.parent?.span.name).toBe("streamObject"); + } + }, + ); + } + if (options.agentSpanName) { test("captures trace for agent.generate()", testConfig, () => { const root = findLatestSpan(events, ROOT_NAME); @@ -922,6 +949,18 @@ export function defineAISDKInstrumentationAssertions(options: { expect(trace.modelChildren.length).toBeGreaterThanOrEqual(1); expect(trace.latestChild?.output).toBeDefined(); }); + + if (options.sdkMajorVersion === 5 && options.agentSpanName === "Agent") { + test("captures Agent.stream() path in v5", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const trace = findAgentStreamTrace(events, "Agent"); + + expectOperationParentedByRoot(trace.operation, root); + expectAISDKParentSpan(trace.parent); + expect(operationName(trace.operation)).toBe("agent-stream"); + expect(trace.parent?.span.name).toBe("Agent.stream"); + }); + } } if (options.supportsDenyOutputOverrideScenario) { diff --git a/e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts b/e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts index 54bdefe8..1c6124ce 100644 --- a/e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts +++ b/e2e/scenarios/ai-sdk-instrumentation/scenario.test.ts @@ -53,6 +53,7 @@ for (const scenario of aiSDKScenarios) { supportsOutputObjectScenario, supportsStreamObject: scenario.supportsStreamObject, supportsToolExecution: scenario.supportsToolExecution, + sdkMajorVersion, testFileUrl: import.meta.url, timeoutMs: AI_SDK_SCENARIO_TIMEOUT_MS, }); @@ -76,6 +77,7 @@ for (const scenario of aiSDKScenarios) { supportsOutputObjectScenario, supportsStreamObject: scenario.supportsStreamObject, supportsToolExecution: scenario.supportsToolExecution, + sdkMajorVersion, testFileUrl: import.meta.url, timeoutMs: AI_SDK_SCENARIO_TIMEOUT_MS, }); diff --git a/js/src/auto-instrumentations/configs/ai-sdk.ts b/js/src/auto-instrumentations/configs/ai-sdk.ts index 5023d65a..2c0ea914 100644 --- a/js/src/auto-instrumentations/configs/ai-sdk.ts +++ b/js/src/auto-instrumentations/configs/ai-sdk.ts @@ -39,12 +39,12 @@ export const aiSDKConfigs: InstrumentationConfig[] = [ }, }, - // streamText - function returning stream + // streamText - async function (v3 only, before the sync refactor in v4) { channelName: aiSDKChannels.streamText.channelName, module: { name: "ai", - versionRange: ">=3.0.0", + versionRange: ">=3.0.0 <4.0.0", filePath: "dist/index.mjs", }, functionQuery: { @@ -52,6 +52,20 @@ export const aiSDKConfigs: InstrumentationConfig[] = [ kind: "Async", }, }, + + // streamText - sync function returning stream (v4+) + { + channelName: aiSDKChannels.streamTextSync.channelName, + module: { + name: "ai", + versionRange: ">=4.0.0", + filePath: "dist/index.mjs", + }, + functionQuery: { + functionName: "streamText", + kind: "Sync", + }, + }, { channelName: aiSDKChannels.streamText.channelName, module: { @@ -91,12 +105,12 @@ export const aiSDKConfigs: InstrumentationConfig[] = [ }, }, - // streamObject - function returning stream + // streamObject - async function (v3 only, before the sync refactor in v4) { channelName: aiSDKChannels.streamObject.channelName, module: { name: "ai", - versionRange: ">=3.0.0", + versionRange: ">=3.0.0 <4.0.0", filePath: "dist/index.mjs", }, functionQuery: { @@ -104,6 +118,20 @@ export const aiSDKConfigs: InstrumentationConfig[] = [ kind: "Async", }, }, + + // streamObject - sync function returning stream (v4+) + { + channelName: aiSDKChannels.streamObjectSync.channelName, + module: { + name: "ai", + versionRange: ">=4.0.0", + filePath: "dist/index.mjs", + }, + functionQuery: { + functionName: "streamObject", + kind: "Sync", + }, + }, { channelName: aiSDKChannels.streamObject.channelName, module: { @@ -147,11 +175,11 @@ export const aiSDKConfigs: InstrumentationConfig[] = [ }, }, - // Agent.stream - async method (v5 only) + // Agent.stream - sync method (v5 only) // The compiled AI SDK bundle emits this as an anonymous class method, so we - // target the first async `stream` method in the file instead of a class name. + // target the first sync `stream` method in the file instead of a class name. { - channelName: aiSDKChannels.agentStream.channelName, + channelName: aiSDKChannels.agentStreamSync.channelName, module: { name: "ai", versionRange: ">=5.0.0 <6.0.0", @@ -159,12 +187,12 @@ export const aiSDKConfigs: InstrumentationConfig[] = [ }, functionQuery: { methodName: "stream", - kind: "Async", + kind: "Sync", index: 0, }, }, { - channelName: aiSDKChannels.agentStream.channelName, + channelName: aiSDKChannels.agentStreamSync.channelName, module: { name: "ai", versionRange: ">=5.0.0 <6.0.0", @@ -172,7 +200,7 @@ export const aiSDKConfigs: InstrumentationConfig[] = [ }, functionQuery: { methodName: "stream", - kind: "Async", + kind: "Sync", index: 0, }, }, diff --git a/js/src/instrumentation/plugins/ai-sdk-channels.ts b/js/src/instrumentation/plugins/ai-sdk-channels.ts index aefe2110..0678982d 100644 --- a/js/src/instrumentation/plugins/ai-sdk-channels.ts +++ b/js/src/instrumentation/plugins/ai-sdk-channels.ts @@ -33,6 +33,15 @@ export const aiSDKChannels = defineChannels("ai", { channelName: "streamText", kind: "async", }), + streamTextSync: channel< + [AISDKCallParams], + AISDKResult, + AISDKChannelContext, + unknown + >({ + channelName: "streamText.sync", + kind: "sync-stream", + }), generateObject: channel< [AISDKCallParams], AISDKStreamResult, @@ -51,6 +60,15 @@ export const aiSDKChannels = defineChannels("ai", { channelName: "streamObject", kind: "async", }), + streamObjectSync: channel< + [AISDKCallParams], + AISDKResult, + AISDKChannelContext, + unknown + >({ + channelName: "streamObject.sync", + kind: "sync-stream", + }), agentGenerate: channel< [AISDKCallParams], AISDKStreamResult, @@ -69,6 +87,15 @@ export const aiSDKChannels = defineChannels("ai", { channelName: "Agent.stream", kind: "async", }), + agentStreamSync: channel< + [AISDKCallParams], + AISDKResult, + AISDKChannelContext, + unknown + >({ + channelName: "Agent.stream.sync", + kind: "sync-stream", + }), toolLoopAgentGenerate: channel< [AISDKCallParams], AISDKStreamResult, diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index 9dfb6a5f..da1c3745 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -1,5 +1,9 @@ import { BasePlugin } from "../core"; -import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing"; +import { + traceStreamingChannel, + traceSyncStreamChannel, + unsubscribeAll, +} from "../core/channel-tracing"; import { SpanTypeAttribute } from "../../../util/index"; import { getCurrentUnixTimestamp } from "../../util"; import { Attachment, type Span, withCurrent } from "../../logger"; @@ -144,6 +148,24 @@ export class AISDKPlugin extends BasePlugin { }), ); + // streamText - sync function returning stream (v4+, used by auto-hook) + this.unsubscribers.push( + traceSyncStreamChannel(aiSDKChannels.streamTextSync, { + name: "streamText", + type: SpanTypeAttribute.LLM, + extractInput: ([params], event, span) => + prepareAISDKInput(params, event, span, denyOutputPaths), + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }), + ); + // generateObject - async function that may return streams this.unsubscribers.push( traceStreamingChannel(aiSDKChannels.generateObject, { @@ -190,6 +212,24 @@ export class AISDKPlugin extends BasePlugin { }), ); + // streamObject - sync function returning stream (v4+, used by auto-hook) + this.unsubscribers.push( + traceSyncStreamChannel(aiSDKChannels.streamObjectSync, { + name: "streamObject", + type: SpanTypeAttribute.LLM, + extractInput: ([params], event, span) => + prepareAISDKInput(params, event, span, denyOutputPaths), + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }), + ); + // Agent.generate - async method this.unsubscribers.push( traceStreamingChannel(aiSDKChannels.agentGenerate, { @@ -210,7 +250,7 @@ export class AISDKPlugin extends BasePlugin { }), ); - // Agent.stream - async method returning stream + // Agent.stream - async method returning stream (v5, used by wrapAISDK) this.unsubscribers.push( traceStreamingChannel(aiSDKChannels.agentStream, { name: "Agent.stream", @@ -236,6 +276,24 @@ export class AISDKPlugin extends BasePlugin { }), ); + // Agent.stream - sync method returning stream (v5, used by auto-hook) + this.unsubscribers.push( + traceSyncStreamChannel(aiSDKChannels.agentStreamSync, { + name: "Agent.stream", + type: SpanTypeAttribute.LLM, + extractInput: ([params], event, span) => + prepareAISDKInput(params, event, span, denyOutputPaths), + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }), + ); + // ToolLoopAgent.generate - async method this.unsubscribers.push( traceStreamingChannel(aiSDKChannels.toolLoopAgentGenerate, {