diff --git a/dev-packages/bun-integration-tests/node-suites/excludes.ts b/dev-packages/bun-integration-tests/node-suites/excludes.ts index 0e868ffeba36..4ba47735d3ef 100644 --- a/dev-packages/bun-integration-tests/node-suites/excludes.ts +++ b/dev-packages/bun-integration-tests/node-suites/excludes.ts @@ -116,6 +116,7 @@ export const NO_AUTO_INSTRUMENTATION = [ 'suites/tracing/openai/v6/test.ts', 'suites/tracing/openai/v7/test.ts', 'suites/tracing/orchestrion-lazy-registration/test.ts', + 'suites/tracing/pi-ai/test.ts', 'suites/tracing/pi-durable/test.ts', 'suites/tracing/postgres/test.ts', 'suites/tracing/postgresjs/test.ts', diff --git a/dev-packages/node-integration-tests/suites/tracing/pi-ai/instrument.mjs b/dev-packages/node-integration-tests/suites/tracing/pi-ai/instrument.mjs new file mode 100644 index 000000000000..e0d424cdae71 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/pi-ai/instrument.mjs @@ -0,0 +1,10 @@ +import * as Sentry from '@sentry/node'; +import { loggingTransport } from '@sentry-internal/node-integration-tests'; + +Sentry.init({ + dsn: 'https://public@dsn.ingest.sentry.io/1337', + release: '1.0', + tracesSampleRate: 1.0, + transport: loggingTransport, + integrations: [Sentry.piAiIntegration()], +}); diff --git a/dev-packages/node-integration-tests/suites/tracing/pi-ai/scenario-compat.mjs b/dev-packages/node-integration-tests/suites/tracing/pi-ai/scenario-compat.mjs new file mode 100644 index 000000000000..c8dd78e9a9ec --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/pi-ai/scenario-compat.mjs @@ -0,0 +1,12 @@ +import * as Sentry from '@sentry/node'; +import { complete, fauxAssistantMessage, registerFauxProvider, streamSimple } from '@earendil-works/pi-ai/compat'; + +// The standalone API dispatches to the providers registered with it. +const compat = registerFauxProvider({ provider: 'compat-faux', models: [{ id: 'compat-model' }] }); +compat.setResponses([fauxAssistantMessage('Hello.'), fauxAssistantMessage('Goodbye.')]); + +await Sentry.startSpan({ name: 'pi-ai-test' }, async () => { + const model = compat.getModel(); + await complete(model, { messages: [{ role: 'user', content: 'Hi.', timestamp: Date.now() }] }); + await streamSimple(model, { messages: [{ role: 'user', content: 'Bye.', timestamp: Date.now() }] }).result(); +}); diff --git a/dev-packages/node-integration-tests/suites/tracing/pi-ai/scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/pi-ai/scenario.mjs new file mode 100644 index 000000000000..2d7742fd6aa8 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/pi-ai/scenario.mjs @@ -0,0 +1,52 @@ +import * as Sentry from '@sentry/node'; +import { Type } from '@earendil-works/pi-ai'; +import { createModels } from '@earendil-works/pi-ai/models'; +import { fauxAssistantMessage, fauxProvider, fauxToolCall } from '@earendil-works/pi-ai/providers/faux'; + +// `pi-ai`'s faux provider scripts model responses in-process. +const faux = fauxProvider({ provider: 'faux', models: [{ id: 'faux-model' }] }); +const models = createModels(); +models.setProvider(faux.provider); +faux.setResponses([ + fauxAssistantMessage(fauxToolCall('get_weather', { city: 'Berlin' }, { id: 'call_1' }), { stopReason: 'toolUse' }), + fauxAssistantMessage('It is sunny in Berlin.'), +]); + +const tools = [ + { + name: 'get_weather', + description: 'Get the current weather for a city.', + parameters: Type.Object({ city: Type.String() }), + }, +]; +const question = { role: 'user', content: 'What is the weather in Berlin?', timestamp: Date.now() }; + +await Sentry.startSpan({ name: 'pi-ai-test' }, async () => { + const model = faux.getModel(); + const events = models.streamSimple(model, { + systemPrompt: 'You are a weather assistant.', + messages: [question], + tools, + }); + for await (const _event of events) { + // Drained like an app consuming the stream. + } + const toolCall = await events.result(); + + await models.complete(model, { + systemPrompt: 'You are a weather assistant.', + messages: [ + question, + toolCall, + { + role: 'toolResult', + toolCallId: 'call_1', + toolName: 'get_weather', + content: [{ type: 'text', text: 'sunny' }], + isError: false, + timestamp: Date.now(), + }, + ], + tools, + }); +}); diff --git a/dev-packages/node-integration-tests/suites/tracing/pi-ai/test.ts b/dev-packages/node-integration-tests/suites/tracing/pi-ai/test.ts new file mode 100644 index 000000000000..2eb62f82eaaa --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/pi-ai/test.ts @@ -0,0 +1,117 @@ +import type { SerializedStreamedSpan } from '@sentry/core'; +import { + GEN_AI_INPUT_MESSAGES, + GEN_AI_OPERATION_NAME, + GEN_AI_OUTPUT_MESSAGES, + GEN_AI_PROVIDER_NAME, + GEN_AI_REQUEST_MODEL, + GEN_AI_RESPONSE_FINISH_REASONS, + GEN_AI_RESPONSE_STREAMING, + GEN_AI_SYSTEM_INSTRUCTIONS, + GEN_AI_TOOL_DEFINITIONS, + GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, + SENTRY_OP, + SENTRY_ORIGIN, +} from '@sentry/conventions/attributes'; +import { GEN_AI_CHAT } from '@sentry/conventions/op'; +import { afterAll, describe, expect } from 'vitest'; +import { supports } from '../../../utils'; +import { cleanupChildProcesses, createEsmTests } from '../../../utils/runner'; + +/** The chat spans under `rootSpanId`, in request order, with what every pi-ai request reports. */ +function chatSpans(spans: SerializedStreamedSpan[], rootSpanId: string): SerializedStreamedSpan[] { + const chats = spans + .filter(span => span.attributes[SENTRY_OP]?.value === GEN_AI_CHAT) + .sort((a, b) => a.start_timestamp - b.start_timestamp); + for (const chat of chats) { + expect(chat.parent_span_id).toBe(rootSpanId); + expect(chat.attributes[SENTRY_ORIGIN]?.value).toBe('auto.ai.pi_ai'); + expect(chat.attributes[GEN_AI_OPERATION_NAME]?.value).toBe('chat'); + expect(chat.attributes[GEN_AI_USAGE_INPUT_TOKENS]?.value).toBeGreaterThan(0); + expect(chat.attributes[GEN_AI_USAGE_OUTPUT_TOKENS]?.value).toBeGreaterThan(0); + } + return chats; +} + +// `@earendil-works/pi-ai` declares `engines.node >= 22.19`. +describe.runIf(supports({ min: 22 }))('pi-ai integration', () => { + afterAll(() => { + cleanupChildProcesses(); + }); + + // `@earendil-works/pi-ai` is ESM-only, so there is no CJS variant of this scenario to run. + createEsmTests(__dirname, 'scenario.mjs', 'instrument.mjs', (createRunner, test) => { + test('traces the requests sent through a `Models` collection', async () => { + await createRunner() + .unordered() + .expect({ + span: container => { + const spans = container.items; + const root = spans.find(span => span.is_segment)!; + expect(root.name).toBe('pi-ai-test'); + + // One span per request: `complete` awaits a `stream` of its own, and no provider SDK + // reports the request a second time. + const chats = chatSpans(spans, root.span_id); + expect(chats.map(chat => chat.name)).toEqual(['chat faux-model', 'chat faux-model']); + + const [toolCall, answer] = chats; + expect(toolCall!.attributes[GEN_AI_PROVIDER_NAME]?.value).toBe('faux'); + expect(toolCall!.attributes[GEN_AI_REQUEST_MODEL]?.value).toBe('faux-model'); + expect(toolCall!.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true); + expect(toolCall!.attributes[GEN_AI_SYSTEM_INSTRUCTIONS]?.value).toBe('You are a weather assistant.'); + expect(toolCall!.attributes[GEN_AI_TOOL_DEFINITIONS]?.value).toContain('get_weather'); + expect(toolCall!.attributes[GEN_AI_RESPONSE_FINISH_REASONS]?.value).toBe('["tool_call"]'); + expect(JSON.parse(String(toolCall!.attributes[GEN_AI_INPUT_MESSAGES]?.value))).toEqual([ + { role: 'user', parts: [{ type: 'text', content: 'What is the weather in Berlin?' }] }, + ]); + expect(JSON.parse(String(toolCall!.attributes[GEN_AI_OUTPUT_MESSAGES]?.value))).toEqual([ + { + role: 'assistant', + parts: [{ type: 'tool_call', id: 'call_1', name: 'get_weather', arguments: '{"city":"Berlin"}' }], + finish_reason: 'tool_call', + }, + ]); + + expect(answer!.attributes[GEN_AI_RESPONSE_STREAMING]).toBeUndefined(); + expect(answer!.attributes[GEN_AI_RESPONSE_FINISH_REASONS]?.value).toBe('["stop"]'); + expect( + JSON.parse(String(answer!.attributes[GEN_AI_INPUT_MESSAGES]?.value)).map( + (message: { role: string }) => message.role, + ), + ).toEqual(['user', 'assistant', 'tool']); + }, + }) + .start() + .completed(); + }); + }); + + createEsmTests(__dirname, 'scenario-compat.mjs', 'instrument.mjs', (createRunner, test) => { + // `compat` loads pi-ai's model catalog with `import ... with { type: 'json' }`. Deno 2.8 drops + // the import attributes once any module hook is registered and compiles the JSON as JavaScript, + // which Deno 2.9 fixed. The SDK registers one for its runtime instrumentation. + test.runIf(supports({ runtimes: ['node', 'bun'] }))( + 'traces the requests sent through the standalone functions', + async () => { + await createRunner() + .unordered() + .expect({ + span: container => { + const root = container.items.find(span => span.is_segment)!; + const [completed, streamed] = chatSpans(container.items, root.span_id); + expect([completed!.name, streamed!.name]).toEqual(['chat compat-model', 'chat compat-model']); + + expect(completed!.attributes[GEN_AI_PROVIDER_NAME]?.value).toBe('compat-faux'); + expect(completed!.attributes[GEN_AI_RESPONSE_STREAMING]).toBeUndefined(); + expect(streamed!.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true); + expect(streamed!.attributes[GEN_AI_OUTPUT_MESSAGES]?.value).toContain('Goodbye.'); + }, + }) + .start() + .completed(); + }, + ); + }); +}); diff --git a/packages/astro/src/index.server.ts b/packages/astro/src/index.server.ts index 6837722ba5b6..d778fef28a90 100644 --- a/packages/astro/src/index.server.ts +++ b/packages/astro/src/index.server.ts @@ -109,6 +109,7 @@ export { createFlueInstrumentation, instrumentPostgresJsSql, mastraIntegration, + piAiIntegration, piDurableIntegration, mcpServerIntegration, SentryMastraExporter, diff --git a/packages/aws-serverless/src/index.ts b/packages/aws-serverless/src/index.ts index 0e684c45dc2d..efb710e97b4a 100644 --- a/packages/aws-serverless/src/index.ts +++ b/packages/aws-serverless/src/index.ts @@ -68,6 +68,7 @@ export { langChainIntegration, langGraphIntegration, mastraIntegration, + piAiIntegration, piDurableIntegration, mcpServerIntegration, SentryMastraExporter, diff --git a/packages/bun/src/index.ts b/packages/bun/src/index.ts index e33986f88bba..7a9ebba13919 100644 --- a/packages/bun/src/index.ts +++ b/packages/bun/src/index.ts @@ -92,6 +92,7 @@ export { langChainIntegration, langGraphIntegration, mastraIntegration, + piAiIntegration, piDurableIntegration, mcpServerIntegration, SentryMastraExporter, diff --git a/packages/cloudflare/src/index.ts b/packages/cloudflare/src/index.ts index 60a32d7c5bbf..dd2f83a2fabb 100644 --- a/packages/cloudflare/src/index.ts +++ b/packages/cloudflare/src/index.ts @@ -193,6 +193,7 @@ export { neonIntegration, nitroIntegration, openAIIntegration, + piAiIntegration, piDurableIntegration, postgresIntegration, postgresJsIntegration, diff --git a/packages/deno/src/index.ts b/packages/deno/src/index.ts index 9f2aac8b30cb..7b6c71244031 100644 --- a/packages/deno/src/index.ts +++ b/packages/deno/src/index.ts @@ -177,6 +177,7 @@ export { langChainIntegration, langGraphIntegration, mastraIntegration, + piAiIntegration, piDurableIntegration, mcpServerIntegration, SentryMastraExporter, diff --git a/packages/elysia/src/index.ts b/packages/elysia/src/index.ts index 014bd9145c4a..2efbdfbc1b3e 100644 --- a/packages/elysia/src/index.ts +++ b/packages/elysia/src/index.ts @@ -69,6 +69,7 @@ export { langGraphIntegration, createFlueInstrumentation, mastraIntegration, + piAiIntegration, piDurableIntegration, mcpServerIntegration, SentryMastraExporter, diff --git a/packages/google-cloud-serverless/src/index.ts b/packages/google-cloud-serverless/src/index.ts index b92984d1b6b6..e0c495683704 100644 --- a/packages/google-cloud-serverless/src/index.ts +++ b/packages/google-cloud-serverless/src/index.ts @@ -68,6 +68,7 @@ export { langChainIntegration, langGraphIntegration, mastraIntegration, + piAiIntegration, piDurableIntegration, mcpServerIntegration, SentryMastraExporter, diff --git a/packages/node/src/index.ts b/packages/node/src/index.ts index 4cf8b6e4b9eb..adf9a8f234a2 100644 --- a/packages/node/src/index.ts +++ b/packages/node/src/index.ts @@ -31,6 +31,7 @@ export { lruMemoizerIntegration, createFlueInstrumentation, mastraIntegration, + piAiIntegration, piDurableIntegration, mcpServerIntegration, SentryMastraExporter, diff --git a/packages/server-utils/src/ai/pi-ai/constants.ts b/packages/server-utils/src/ai/pi-ai/constants.ts new file mode 100644 index 000000000000..18fe024306c6 --- /dev/null +++ b/packages/server-utils/src/ai/pi-ai/constants.ts @@ -0,0 +1,3 @@ +export const PI_AI_INTEGRATION_NAME = 'PiAi' as const; + +export const PI_AI_ORIGIN = 'auto.ai.pi_ai'; diff --git a/packages/server-utils/src/index.ts b/packages/server-utils/src/index.ts index 005e8a4a8b79..0988e68312af 100644 --- a/packages/server-utils/src/index.ts +++ b/packages/server-utils/src/index.ts @@ -55,6 +55,7 @@ export { langChainIntegration } from './integrations/langchain'; export { langGraphIntegration } from './integrations/langgraph'; export { createFlueInstrumentation } from './ai/flue'; export { flueIntegration } from './integrations/flue'; +export { piAiIntegration } from './integrations/pi-ai'; export { piDurableIntegration } from './integrations/pi-durable'; export type { FlueOptions } from './ai/flue'; export { mastraIntegration } from './integrations/mastra'; diff --git a/packages/server-utils/src/integrations/pi-ai.ts b/packages/server-utils/src/integrations/pi-ai.ts new file mode 100644 index 000000000000..7ac0a34fd781 --- /dev/null +++ b/packages/server-utils/src/integrations/pi-ai.ts @@ -0,0 +1,112 @@ +import * as diagnosticsChannel from '../utils/diagnosticsChannel'; +import type { IntegrationFn, Span } from '@sentry/core'; +import { defineIntegration, isThenable, startInactiveSpan } from '@sentry/core'; +import type { GenAiOptions } from '../ai/core/utils'; +import { resolveAIRecordingOptions } from '../ai/core/utils'; +import { getPiAiChatSpanOptions, setPiAiResponseAttributes } from '../ai/pi-ai/chat'; +import { PI_AI_INTEGRATION_NAME, PI_AI_ORIGIN } from '../ai/pi-ai/constants'; +import type { PiAiContext } from '../ai/pi-ai/messages'; +import { skipPiAiProviderIntegrations } from '../ai/pi-ai/providers'; +import type { PiEventStream, PiStreamOptions } from '../ai/pi-ai/types'; +import type { OrchestrionChannelContext } from '../orchestrion/types'; +import { CHANNELS } from '../orchestrion/channels'; +import { piAiModuleNames } from '../orchestrion/config/pi-ai'; +import { invokeOrchestrionInstrumentation } from '../orchestrion/instrumentation'; +import { bindTracingChannelToSpan } from '../tracing-channel'; + +export type PiAiOptions = GenAiOptions; + +const _piAiIntegration = ((options: PiAiOptions = {}) => { + return { + name: PI_AI_INTEGRATION_NAME, + setup(client) { + invokeOrchestrionInstrumentation(client, piAiModuleNames, instrumentPiAi, [options]); + }, + }; +}) satisfies IntegrationFn; + +function instrumentPiAi(options: PiAiOptions): void { + const streamChannel = diagnosticsChannel.tracingChannel(CHANNELS.PI_AI_STREAM); + const completeChannel = diagnosticsChannel.tracingChannel(CHANNELS.PI_AI_COMPLETE); + + // `complete` awaits `stream`, and the `compat` functions hand some models to a `Models` of their + // own, so a traced call can make another one synchronously. Only the outermost gets a span. + let depth = 0; + const startSpan = (data: OrchestrionChannelContext, streaming: boolean): Span | undefined => { + if (depth > 0) { + return undefined; + } + skipPiAiProviderIntegrations(); + const [model, context, callOptions] = data.arguments ?? []; + return startInactiveSpan( + getPiAiChatSpanOptions({ + model: model ?? {}, + // The deferred calls pass a handle here, which carries no request content. + context: context as PiAiContext | undefined, + callOptions: callOptions as PiStreamOptions | undefined, + streaming, + origin: PI_AI_ORIGIN, + recordInputs: resolveAIRecordingOptions(options).recordInputs, + }), + ); + }; + + const setResponse = (span: Span, message: unknown): void => { + if (message && typeof message === 'object') { + setPiAiResponseAttributes(span, message, resolveAIRecordingOptions(options).recordOutputs); + } + }; + + bindTracingChannelToSpan(streamChannel, data => startSpan(data, true), { + // The stream is returned before the request settles; its `result()` settles with the message. + deferSpanEnd: ({ span, data, end }) => { + const events = data.result as Partial | undefined; + const settled = 'error' in data || typeof events?.result !== 'function' ? undefined : events.result(); + if (!isThenable(settled)) { + return false; + } + settled.then(message => { + setResponse(span, message); + end(); + }, end); + return true; + }, + }); + + bindTracingChannelToSpan(completeChannel, data => startSpan(data, false), { + beforeSpanEnd: (span, data) => { + if (!('error' in data)) { + setResponse(span, data.result); + } + }, + }); + + for (const channel of [streamChannel, completeChannel]) { + channel.start.subscribe(() => { + depth++; + }); + channel.end.subscribe(() => { + depth--; + }); + } +} + +/** + * Traces the requests sent with pi-ai (`@earendil-works/pi-ai` >= 0.80.0 < 2) as `gen_ai.chat` spans, + * through `createModels()` or the standalone `stream` and `complete` functions. Needs the Sentry + * runtime hook (`--import`) or bundler plugin. + * + * Not enabled by default yet. Flue and pi-durable send their model calls through pi-ai and still + * report them as chat spans of their own, so enabling this alongside them reports each call twice. + * + * Once pi-ai has sent a request, the `openai`, `@anthropic-ai/sdk`, `@google/genai` and Workers AI + * integrations stop recording requests for the rest of the process, since pi-ai sends through them. + * + * @example + * ```javascript + * Sentry.init({ + * integrations: [Sentry.piAiIntegration()], + * }); + * ``` + */ +export const piAiIntegration = defineIntegration(_piAiIntegration); diff --git a/packages/server-utils/src/orchestrion/channels.ts b/packages/server-utils/src/orchestrion/channels.ts index 67a88bbb3825..b4badf17834d 100644 --- a/packages/server-utils/src/orchestrion/channels.ts +++ b/packages/server-utils/src/orchestrion/channels.ts @@ -30,6 +30,7 @@ import { nodeCronChannels } from './config/node-cron'; import { nodeScheduleChannels } from './config/node-schedule'; import { openaiChannels } from './config/openai'; import { pgChannels } from './config/pg'; +import { piAiChannels } from './config/pi-ai'; import { piDurableChannels } from './config/pi-durable'; import { postgresJsChannels } from './config/postgres'; import { prismaChannels } from './config/prisma'; @@ -89,6 +90,7 @@ export const CHANNELS = { ...nodeScheduleChannels, ...openaiChannels, ...pgChannels, + ...piAiChannels, ...piDurableChannels, ...postgresJsChannels, ...prismaChannels, diff --git a/packages/server-utils/src/orchestrion/config/index.ts b/packages/server-utils/src/orchestrion/config/index.ts index bf2c4b50b178..546984480b6d 100644 --- a/packages/server-utils/src/orchestrion/config/index.ts +++ b/packages/server-utils/src/orchestrion/config/index.ts @@ -23,6 +23,7 @@ import { langgraphConfig } from './langgraph'; import { lruMemoizerConfig } from './lru-memoizer'; import { flueConfig } from './flue'; import { mastraConfig } from './mastra'; +import { piAiConfig } from './pi-ai'; import { piDurableConfig } from './pi-durable'; import { mcpServerConfig } from './mcp-server'; import { mistralConfig } from './mistral'; @@ -97,6 +98,7 @@ export const SENTRY_INSTRUMENTATIONS: InstrumentationConfig[] = [ ...nodeScheduleConfig, ...openaiConfig, ...pgConfig, + ...piAiConfig, ...piDurableConfig, ...postgresJsConfig, ...prismaConfig, diff --git a/packages/server-utils/src/orchestrion/config/pi-ai.ts b/packages/server-utils/src/orchestrion/config/pi-ai.ts new file mode 100644 index 000000000000..c42a4aca2937 --- /dev/null +++ b/packages/server-utils/src/orchestrion/config/pi-ai.ts @@ -0,0 +1,44 @@ +import type { InstrumentationConfig } from '../apmTypes'; + +import { getModuleNames } from './module-names'; + +const module = { name: '@earendil-works/pi-ai', versionRange: '>=0.80.0 <2.0.0' }; + +type Kind = 'Sync' | 'Async'; + +const modelsMethod = (methodName: string, channelName: string, kind: Kind): InstrumentationConfig => ({ + channelName, + module: { ...module, filePath: 'dist/models.js' }, + functionQuery: { className: 'ModelsImpl', methodName, kind }, +}); + +const compatFunction = (functionName: string, channelName: string, kind: Kind): InstrumentationConfig => ({ + channelName, + module: { ...module, filePath: 'dist/compat.js' }, + functionQuery: { functionName, kind }, +}); + +// `ModelsImpl` is what `createModels()` returns, and what Flue and pi-durable send their requests +// through; the `compat` functions are the standalone API. The `complete` functions only await the +// stream of their `stream` counterpart, and the integration leaves out that nested call, so the +// span of a request still tells whether its caller streamed it. `streamDeferred` and `fetchDeferred` +// fetch the answer of a request the provider parked. +export const piAiConfig = [ + modelsMethod('stream', 'stream', 'Sync'), + modelsMethod('streamSimple', 'stream', 'Sync'), + compatFunction('stream', 'stream', 'Sync'), + compatFunction('streamSimple', 'stream', 'Sync'), + modelsMethod('complete', 'complete', 'Async'), + modelsMethod('completeSimple', 'complete', 'Async'), + compatFunction('complete', 'complete', 'Async'), + compatFunction('completeSimple', 'complete', 'Async'), + modelsMethod('streamDeferred', 'stream', 'Sync'), + modelsMethod('fetchDeferred', 'complete', 'Async'), +] satisfies InstrumentationConfig[]; + +export const piAiModuleNames = getModuleNames(piAiConfig); + +export const piAiChannels = { + PI_AI_STREAM: 'orchestrion:@earendil-works/pi-ai:stream', + PI_AI_COMPLETE: 'orchestrion:@earendil-works/pi-ai:complete', +} as const; diff --git a/packages/server-utils/test/integrations/pi-ai.test.ts b/packages/server-utils/test/integrations/pi-ai.test.ts new file mode 100644 index 000000000000..a062e191a6e7 --- /dev/null +++ b/packages/server-utils/test/integrations/pi-ai.test.ts @@ -0,0 +1,380 @@ +import { tracingChannel } from 'node:diagnostics_channel'; +import type { Span } from '@sentry/core'; +import { + _INTERNAL_clearAiProviderSkips, + _INTERNAL_shouldSkipAiProviderWrapping, + _INTERNAL_skipAiProviderWrapping, + GLOBAL_OBJ, + setAsyncContextStrategy, + setCurrentClient, + spanToJSON, + spanToStaticSpanJSON, + startInactiveSpan, + startSpan, + withScope, +} from '@sentry/core'; +import { + GEN_AI_CONVERSATION_ID, + GEN_AI_INPUT_MESSAGES, + GEN_AI_OUTPUT_MESSAGES, + GEN_AI_PROVIDER_NAME, + GEN_AI_REQUEST_MAX_TOKENS, + GEN_AI_REQUEST_MODEL, + GEN_AI_REQUEST_REASONING_LEVEL, + GEN_AI_REQUEST_TEMPERATURE, + GEN_AI_RESPONSE_FINISH_REASONS, + GEN_AI_RESPONSE_ID, + GEN_AI_RESPONSE_MODEL, + GEN_AI_RESPONSE_STREAMING, + GEN_AI_SYSTEM_INSTRUCTIONS, + GEN_AI_TOOL_DEFINITIONS, + GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_USAGE_TOTAL_TOKENS, + SENTRY_OP, + SENTRY_ORIGIN, + SENTRY_STATUS_MESSAGE, + SERVER_ADDRESS, + SERVER_PORT, +} from '@sentry/conventions/attributes'; +import { GEN_AI_CHAT } from '@sentry/conventions/op'; +import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest'; +import { ANTHROPIC_AI_INTEGRATION_NAME } from '../../src/ai/anthropic-ai/constants'; +import { GOOGLE_GENAI_INTEGRATION_NAME } from '../../src/ai/google-genai/constants'; +import { OPENAI_INTEGRATION_NAME } from '../../src/ai/openai/constants'; +import { getPiAiChatSpanOptions, setPiAiResponseAttributes } from '../../src/ai/pi-ai/chat'; +import { PI_AI_ORIGIN } from '../../src/ai/pi-ai/constants'; +import { setAsyncLocalStorageAsyncContextStrategy } from '../../src/async-context'; +import type { PiEventStream } from '../../src/ai/pi-ai/types'; +import { instrumentWorkersAiClient } from '../../src/ai/workers-ai'; +import { piAiIntegration } from '../../src/integrations/pi-ai'; +import { CHANNELS } from '../../src/orchestrion/channels'; +import { getDefaultTestClientOptions, TestClient } from '../mocks/client'; + +const MODEL = { id: 'faux-model', provider: 'faux' }; +const ANSWER = { role: 'assistant', content: [{ type: 'text', text: 'Paris.' }], stopReason: 'stop' }; + +const streamChannel = tracingChannel(CHANNELS.PI_AI_STREAM); +const completeChannel = tracingChannel(CHANNELS.PI_AI_COMPLETE); + +// What the orchestrion transform turns the pi-ai entry points into. +function stream(send: () => unknown, ...args: unknown[]): PiEventStream { + return streamChannel.traceSync(send, { arguments: args }) as PiEventStream; +} + +function complete(send: () => Promise, ...args: unknown[]): Promise { + return completeChannel.tracePromise(send, { arguments: args }); +} + +function settlesWith(message: unknown): () => PiEventStream { + return () => ({ result: async () => message }) as PiEventStream; +} + +describe('piAiIntegration', () => { + let endedSpans: Span[]; + let client: TestClient; + + function useClient(options: Partial[0]> = {}): void { + client = new TestClient( + getDefaultTestClientOptions({ dsn: 'https://public@dsn.ingest.sentry.io/1337', tracesSampleRate: 1, ...options }), + ); + setCurrentClient(client); + client.init(); + client.on('spanEnd', span => endedSpans.push(span)); + } + + beforeAll(() => { + setAsyncLocalStorageAsyncContextStrategy(); + GLOBAL_OBJ.__SENTRY_ORCHESTRION__ = { runtime: ['@earendil-works/pi-ai'] }; + useClient(); + piAiIntegration().setup!(client); + }); + + afterAll(() => { + setAsyncContextStrategy(undefined); + delete GLOBAL_OBJ.__SENTRY_ORCHESTRION__; + }); + + beforeEach(() => { + _INTERNAL_clearAiProviderSkips(); + endedSpans = []; + useClient(); + }); + + it('ends a streamed span only once its result settles', async () => { + let settle!: (message: unknown) => void; + const result = new Promise(resolve => { + settle = resolve; + }); + + const events = stream(() => ({ result: () => result }), MODEL, { messages: [] }); + expect(endedSpans).toHaveLength(0); + + settle(ANSWER); + await events.result(); + expect(endedSpans).toHaveLength(1); + }); + + it.each([true, false])('traces a request with its options and response identity (streaming: %s)', async streaming => { + const message = { + ...ANSWER, + responseModel: 'faux-model-2026', + responseId: 'resp_1', + usage: { input: 10, output: 2, totalTokens: 12 }, + }; + const args = [ + { ...MODEL, baseUrl: 'https://api.example.com:8443/v1' }, + { messages: [{ role: 'user', content: 'Hi.' }] }, + { temperature: 0.2, maxTokens: 50, reasoning: 'low' }, + ]; + + await (streaming ? stream(settlesWith(message), ...args).result() : complete(async () => message, ...args)); + + const chat = spanToJSON(endedSpans[0]!); + expect(chat.name).toBe('chat faux-model'); + expect(chat.attributes).toMatchObject({ + [SENTRY_OP]: GEN_AI_CHAT, + [SENTRY_ORIGIN]: 'auto.ai.pi_ai', + [GEN_AI_PROVIDER_NAME]: 'faux', + [GEN_AI_REQUEST_MODEL]: 'faux-model', + [GEN_AI_RESPONSE_MODEL]: 'faux-model-2026', + [GEN_AI_RESPONSE_ID]: 'resp_1', + [SERVER_ADDRESS]: 'api.example.com', + [SERVER_PORT]: 8443, + [GEN_AI_REQUEST_TEMPERATURE]: 0.2, + [GEN_AI_REQUEST_MAX_TOKENS]: 50, + [GEN_AI_REQUEST_REASONING_LEVEL]: 'low', + [GEN_AI_INPUT_MESSAGES]: '[{"role":"user","parts":[{"type":"text","content":"Hi."}]}]', + [GEN_AI_OUTPUT_MESSAGES]: + '[{"role":"assistant","parts":[{"type":"text","content":"Paris."}],"finish_reason":"stop"}]', + [GEN_AI_RESPONSE_FINISH_REASONS]: '["stop"]', + [GEN_AI_USAGE_INPUT_TOKENS]: 10, + [GEN_AI_USAGE_OUTPUT_TOKENS]: 2, + }); + expect(chat.attributes[GEN_AI_RESPONSE_STREAMING]).toBe(streaming ? true : undefined); + }); + + // `streamDeferred` and `fetchDeferred` pass the handle of a parked request where the others pass it. + it('records no request content for a call that fetches a parked request', async () => { + await stream(settlesWith(ANSWER), MODEL, { id: 'handle_1' }, { signal: new AbortController().signal }).result(); + + const recorded = Object.keys(spanToJSON(endedSpans[0]!).attributes); + expect(recorded).not.toContain(GEN_AI_INPUT_MESSAGES); + expect(recorded).not.toContain(GEN_AI_SYSTEM_INSTRUCTIONS); + expect(recorded).not.toContain(GEN_AI_TOOL_DEFINITIONS); + }); + + // `complete` awaits the stream of `stream`, and the `compat` functions hand some models to a + // `Models` of their own: each is one request. + it('traces a call that makes another traced call as one request', async () => { + await complete(() => stream(settlesWith(ANSWER), MODEL, { messages: [] }).result(), MODEL, { messages: [] }); + await stream(() => stream(settlesWith(ANSWER), MODEL, { messages: [] }), MODEL, { messages: [] }).result(); + + expect(endedSpans.map(span => spanToJSON(span).attributes[GEN_AI_RESPONSE_STREAMING])).toEqual([undefined, true]); + }); + + it('traces requests after a nested one, and after one that threw', async () => { + expect(() => + stream(() => { + throw new Error('sync boom'); + }, MODEL), + ).toThrow('sync boom'); + await complete(() => stream(settlesWith(ANSWER), MODEL).result(), MODEL); + await complete(async () => ANSWER, MODEL); + + expect(endedSpans).toHaveLength(3); + }); + + it('parents the span to the active span and makes it active for the provider call', async () => { + let providerParent: string | undefined; + + await startSpan({ name: 'agent' }, async parent => { + await stream( + () => { + providerParent = startSpan({ name: 'provider' }, span => spanToJSON(span).parent_span_id); + return { result: async () => ANSWER }; + }, + MODEL, + { messages: [] }, + ).result(); + + const chat = endedSpans.map(span => spanToJSON(span)).find(span => span.attributes[SENTRY_OP] === GEN_AI_CHAT)!; + expect(chat.parent_span_id).toBe(parent.spanContext().spanId); + expect(providerParent).toBe(chat.span_id); + }); + }); + + it('sets the conversation id of the scope on the chat span', async () => { + await withScope(async scope => { + scope.setConversationId('harness:7'); + await complete(async () => ANSWER, MODEL, { messages: [] }); + }); + + expect(spanToJSON(endedSpans[0]!).attributes[GEN_AI_CONVERSATION_ID]).toBe('harness:7'); + }); + + // The provider parks the request and the caller fetches the answer later: only the fetch is billed. + it('records no usage for a request the provider parked, and the usage of its answer', async () => { + const parked = { role: 'assistant', content: [], stopReason: 'deferred', usage: { input: 0, totalTokens: 0 } }; + const answer = { ...ANSWER, usage: { input: 8, output: 2, totalTokens: 10 } }; + + await complete(async () => parked, MODEL, { messages: [] }); + await complete(async () => answer, MODEL, { id: 'handle_1' }); + + const [parkedChat, answerChat] = endedSpans.map(span => spanToJSON(span).attributes); + expect(parkedChat![GEN_AI_RESPONSE_FINISH_REASONS]).toBe('["deferred"]'); + expect(parkedChat![GEN_AI_USAGE_TOTAL_TOKENS]).toBeUndefined(); + expect(answerChat![GEN_AI_USAGE_TOTAL_TOKENS]).toBe(10); + }); + + it('marks a failed or aborted request with a fixed status message and maps a tool-calling stop', async () => { + const failed = { + role: 'assistant', + content: [], + stopReason: 'error', + errorMessage: '401 invalid x-api-key for user@example.com', + usage: { input: 0, output: 0, totalTokens: 0 }, + }; + const aborted = { role: 'assistant', content: [], stopReason: 'aborted' }; + const toolCall = { + role: 'assistant', + content: [{ type: 'toolCall', id: 'call_1', name: 'read', arguments: { path: 'a.txt' } }], + stopReason: 'toolUse', + }; + + await complete(async () => failed, MODEL, { messages: [] }); + await stream(settlesWith(aborted), MODEL, { messages: [] }).result(); + await complete(async () => toolCall, MODEL, { messages: [] }); + + const [chatFailed, , chatToolCall] = endedSpans.map(span => spanToJSON(span)); + expect(chatFailed!.status).toBe('error'); + expect(chatFailed!.attributes[SENTRY_STATUS_MESSAGE]).toBe('internal_error'); + expect(JSON.stringify(chatFailed!.attributes)).not.toContain('x-api-key'); + expect(chatFailed!.attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toBeUndefined(); + // Streamed span JSON reports a cancellation as `ok`; the static one keeps it. + expect(spanToStaticSpanJSON(endedSpans[1]!).status).toBe('cancelled'); + expect(chatToolCall!.attributes[GEN_AI_RESPONSE_FINISH_REASONS]).toBe('["tool_call"]'); + expect(chatToolCall!.attributes[GEN_AI_OUTPUT_MESSAGES]).toContain('"finish_reason":"tool_call"'); + }); + + it('marks a rejected, throwing or broken request as errored and keeps the error', async () => { + await expect( + complete(async () => { + throw new Error('boom'); + }, MODEL), + ).rejects.toThrow('boom'); + expect(() => + stream(() => { + throw new Error('sync boom'); + }, MODEL), + ).toThrow('sync boom'); + const broken = stream(() => ({ result: () => Promise.reject(new Error('broken provider')) }), MODEL); + await expect(broken.result()).rejects.toThrow('broken provider'); + + expect(endedSpans.map(span => spanToJSON(span).status)).toEqual(['error', 'error', 'error']); + }); + + it('returns a result of another shape unchanged and ends its span', () => { + const events = { events: [] }; + + expect(stream(() => events, MODEL)).toBe(events); + + expect(endedSpans.map(span => spanToJSON(span).attributes[SENTRY_OP])).toEqual([GEN_AI_CHAT]); + }); + + describe('provider skip', () => { + it('skips the provider integrations on the first request, not at setup', async () => { + expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(false); + + await complete(async () => ANSWER, MODEL, { messages: [] }); + + expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(true); + expect(_INTERNAL_shouldSkipAiProviderWrapping(ANTHROPIC_AI_INTEGRATION_NAME)).toBe(true); + expect(_INTERNAL_shouldSkipAiProviderWrapping(GOOGLE_GENAI_INTEGRATION_NAME)).toBe(true); + }); + + // The registry is reset per client, so a one-shot call would be undone by the next `init()`. + it('applies the skip again after the registry is cleared', async () => { + await complete(async () => ANSWER, MODEL, { messages: [] }); + _INTERNAL_clearAiProviderSkips(); + expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(false); + + await complete(async () => ANSWER, MODEL, { messages: [] }); + + expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(true); + }); + + it('applies the skip to every provider when only some are already registered', async () => { + _INTERNAL_skipAiProviderWrapping([OPENAI_INTEGRATION_NAME]); + + await complete(async () => ANSWER, MODEL, { messages: [] }); + + expect(_INTERNAL_shouldSkipAiProviderWrapping(ANTHROPIC_AI_INTEGRATION_NAME)).toBe(true); + expect(_INTERNAL_shouldSkipAiProviderWrapping(GOOGLE_GENAI_INTEGRATION_NAME)).toBe(true); + }); + + // `createAI()` of `agents/models/pi-ai` sends pi-ai requests through the Workers AI binding. + it('reports a request sent through the Workers AI binding once', async () => { + const ai = instrumentWorkersAiClient({ + run: async (_model: string, _inputs: unknown, _options: unknown) => new Response('{}'), + gateway: () => ({}), + toMarkdown: async () => [], + }); + + await complete( + async () => { + await ai.run('@cf/moonshotai/kimi-k2.7-code', { messages: [] }, { returnRawResponse: true }); + return ANSWER; + }, + MODEL, + { messages: [] }, + ); + + expect(endedSpans.map(span => spanToJSON(span).name)).toEqual(['chat faux-model']); + }); + }); + + describe('content recording', () => { + const message = { + role: 'assistant', + content: [{ type: 'toolCall', id: 'call_1', name: 'get_weather', arguments: { city: 'Berlin' } }], + stopReason: 'toolUse', + }; + const context = { + messages: [ + { role: 'system', content: '', sections: { preamble: 'You help.' }, toolsAdded: [{ name: 'get_weather' }] }, + { role: 'user', content: 'Weather in Berlin?' }, + ], + }; + const CONTENT_ATTRIBUTES = [ + GEN_AI_INPUT_MESSAGES, + GEN_AI_OUTPUT_MESSAGES, + GEN_AI_SYSTEM_INSTRUCTIONS, + GEN_AI_TOOL_DEFINITIONS, + ]; + + it('records nothing when inputs and outputs are off', () => { + const span = startInactiveSpan( + getPiAiChatSpanOptions({ model: MODEL, context, streaming: false, origin: PI_AI_ORIGIN, recordInputs: false }), + ); + setPiAiResponseAttributes(span, message, false); + span.end(); + + const recorded = Object.keys(spanToJSON(span).attributes); + for (const attribute of CONTENT_ATTRIBUTES) { + expect(recorded).not.toContain(attribute); + } + }); + + it('follows the data collection settings of the current client', async () => { + useClient({ dataCollection: { genAI: { inputs: false, outputs: false } } }); + + await complete(async () => message, MODEL, context); + + const recorded = endedSpans.flatMap(span => Object.keys(spanToJSON(span).attributes)); + for (const attribute of CONTENT_ATTRIBUTES) { + expect(recorded).not.toContain(attribute); + } + }); + }); +}); diff --git a/packages/server-utils/test/orchestrion/moduleInjectedTransform.test.ts b/packages/server-utils/test/orchestrion/moduleInjectedTransform.test.ts index 0feac3d58b47..7a4a4e741d8d 100644 --- a/packages/server-utils/test/orchestrion/moduleInjectedTransform.test.ts +++ b/packages/server-utils/test/orchestrion/moduleInjectedTransform.test.ts @@ -44,7 +44,8 @@ describe('channel integration definitions', () => { // `@nestjs/*` and `@remix-run/*` are covered by listeners in their own SDK packages. // Each cron monitor is billed, so bundler-only SDKs must not install crons integrations for every app. - const optIn = new Set(['node-cron', 'node-schedule']); + // pi-ai would duplicate Flue's and pi-durable's own chat spans until they rely on it. + const optIn = new Set(['node-cron', 'node-schedule', '@earendil-works/pi-ai']); const uncovered = [...configured].filter( name => !defined.has(name) && !optIn.has(name) && !name.startsWith('@nestjs/') && !name.startsWith('@remix-run/'), );