From 45e1078a0e220c276d805afde473a876458eeff3 Mon Sep 17 00:00:00 2001 From: Francesco Novy Date: Fri, 30 Oct 2026 11:59:48 +0100 Subject: [PATCH 1/2] feat(server-utils): Trace Flue and pi-durable model calls through pi-ai Enables `piAiIntegration` by default. Flue and pi-durable send their model calls through pi-ai, so they no longer open chat spans of their own: Flue drops its turn spans and hands the conversation id to the pi-ai span through the scope, and pi-durable drops its `models` proxy. Co-Authored-By: Claude Opus 5.5 --- .../tests/flue.test.ts | 2 +- .../tests/pi-durable.test.ts | 4 +- .../node-flue/tests/flue.test.ts | 6 +- .../node-pi-durable/tests/pi-durable.test.ts | 6 +- .../suites/tracing/flue/test.ts | 14 +- .../suites/tracing/pi-ai/instrument.mjs | 1 - .../suites/tracing/pi-durable/test.ts | 14 +- .../server-utils/src/ai/flue/constants.ts | 12 +- packages/server-utils/src/ai/flue/index.ts | 70 +--- packages/server-utils/src/ai/flue/types.ts | 44 +- packages/server-utils/src/ai/flue/utils.ts | 152 +------ packages/server-utils/src/ai/pi-ai/chat.ts | 6 +- .../server-utils/src/ai/pi-ai/providers.ts | 4 +- .../server-utils/src/ai/pi-durable/index.ts | 15 +- .../server-utils/src/ai/pi-durable/models.ts | 113 ------ .../server-utils/src/ai/pi-durable/types.ts | 24 +- .../server-utils/src/integrations/index.ts | 2 + .../server-utils/src/integrations/pi-ai.ts | 11 +- .../src/integrations/pi-durable.ts | 7 +- .../config/channel-integration-definitions.ts | 1 + .../test/ai/lib/tracing/flue.test.ts | 376 ++---------------- .../test/ai/lib/tracing/pi-durable.test.ts | 375 ++--------------- .../test/integrations/pi-ai.test.ts | 3 +- .../moduleInjectedTransform.test.ts | 3 +- 24 files changed, 142 insertions(+), 1123 deletions(-) delete mode 100644 packages/server-utils/src/ai/pi-durable/models.ts diff --git a/dev-packages/e2e-tests/test-applications/cloudflare-flue-send-to-sentry/tests/flue.test.ts b/dev-packages/e2e-tests/test-applications/cloudflare-flue-send-to-sentry/tests/flue.test.ts index ac6c0dab180f..d7842f87eb3f 100644 --- a/dev-packages/e2e-tests/test-applications/cloudflare-flue-send-to-sentry/tests/flue.test.ts +++ b/dev-packages/e2e-tests/test-applications/cloudflare-flue-send-to-sentry/tests/flue.test.ts @@ -48,7 +48,7 @@ test('instruments a Flue agent with no manual instrument() call', async () => { await expect .poll(() => fetchSpanAttributes(traceId!, chat!.event_id!), EVENT_POLLING_OPTIONS) .toMatchObject({ - origin: 'auto.ai.flue', + origin: 'auto.ai.pi_ai', 'gen_ai.provider.name': 'openrouter', 'gen_ai.usage.input_tokens': expect.anything(), 'gen_ai.cost.total_tokens': expect.anything(), diff --git a/dev-packages/e2e-tests/test-applications/cloudflare-pi-durable/tests/pi-durable.test.ts b/dev-packages/e2e-tests/test-applications/cloudflare-pi-durable/tests/pi-durable.test.ts index 0c8f04674b12..092c90f3bdde 100644 --- a/dev-packages/e2e-tests/test-applications/cloudflare-pi-durable/tests/pi-durable.test.ts +++ b/dev-packages/e2e-tests/test-applications/cloudflare-pi-durable/tests/pi-durable.test.ts @@ -53,7 +53,9 @@ test('traces a PiHarness prompt as invoke_agent with chat, execute_tool and prov // pi-ai sends the requests through `@anthropic-ai/sdk`, whose own integration must stay out so // each request is reported once. expect(spans.filter(span => String(span.attributes[SENTRY_ORIGIN]?.value).startsWith('auto.ai.'))).toEqual( - spans.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.pi_durable'), + spans.filter(span => + ['auto.ai.pi_durable', 'auto.ai.pi_ai'].includes(String(span.attributes[SENTRY_ORIGIN]?.value)), + ), ); // One tool-calling response, then the answer. diff --git a/dev-packages/e2e-tests/test-applications/node-flue/tests/flue.test.ts b/dev-packages/e2e-tests/test-applications/node-flue/tests/flue.test.ts index a829a0ba282e..49c3197032d0 100644 --- a/dev-packages/e2e-tests/test-applications/node-flue/tests/flue.test.ts +++ b/dev-packages/e2e-tests/test-applications/node-flue/tests/flue.test.ts @@ -33,11 +33,11 @@ test('captures the invoke_agent / chat / execute_tool hierarchy for a Flue turn' expect(invokeAgent?.attributes?.['gen_ai.operation.name']?.value).toBe('invoke_agent'); expect(invokeAgent?.attributes?.['gen_ai.agent.name']?.value).toBe('Hello'); - expect(chat?.attributes?.['sentry.origin']?.value).toBe('auto.ai.flue'); + expect(chat?.attributes?.['sentry.origin']?.value).toBe('auto.ai.pi_ai'); expect(chat?.attributes?.['gen_ai.provider.name']?.value).toBe('openrouter'); expect(typeof chat?.attributes?.['gen_ai.usage.input_tokens']?.value).toBe('number'); expect(typeof chat?.attributes?.['gen_ai.usage.output_tokens']?.value).toBe('number'); - // Flue computes cost itself; no provider SDK reports it. + // pi-ai computes cost itself; no provider SDK reports it. expect(typeof chat?.attributes?.['gen_ai.cost.total_tokens']?.value).toBe('number'); expect(executeTool?.attributes?.['gen_ai.tool.name']?.value).toBe('get_weather'); @@ -92,7 +92,7 @@ test('nests a manual span raised inside a tool under that tool span', async ({ b expect(manualSpan?.parent_span_id).toBe(executeTool?.span_id); }); -// Flue's `model` operation is wrapped so the turn span is active for it, which is what puts the +// The pi-ai integration makes the chat span active for the request it sends, which is what puts the // provider's HTTP call inside `chat` rather than beside it. test('nests the provider HTTP call inside the chat span', async ({ baseURL }) => { const spansPromise = collectStreamedSpans( diff --git a/dev-packages/e2e-tests/test-applications/node-pi-durable/tests/pi-durable.test.ts b/dev-packages/e2e-tests/test-applications/node-pi-durable/tests/pi-durable.test.ts index 3c18f1bc8cc3..42de516c91bc 100644 --- a/dev-packages/e2e-tests/test-applications/node-pi-durable/tests/pi-durable.test.ts +++ b/dev-packages/e2e-tests/test-applications/node-pi-durable/tests/pi-durable.test.ts @@ -52,14 +52,16 @@ test('traces a run as invoke_agent with chat, execute_tool and provider spans', // pi-ai sends the requests through `@anthropic-ai/sdk`, whose own integration must stay out so // each request is reported once. expect(spans.filter(span => String(span.attributes[SENTRY_ORIGIN]?.value).startsWith('auto.ai.'))).toEqual( - spans.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.pi_durable'), + spans.filter(span => + ['auto.ai.pi_durable', 'auto.ai.pi_ai'].includes(String(span.attributes[SENTRY_ORIGIN]?.value)), + ), ); // One tool-calling response, then the answer. expect(chats.length).toBeGreaterThanOrEqual(2); for (const chat of chats) { expect(chat.parent_span_id).toBe(agent.span_id); - expect(chat.attributes[SENTRY_ORIGIN]?.value).toBe('auto.ai.pi_durable'); + expect(chat.attributes[SENTRY_ORIGIN]?.value).toBe('auto.ai.pi_ai'); expect(chat.attributes[GEN_AI_PROVIDER_NAME]?.value).toBe('openrouter'); expect(chat.attributes[GEN_AI_REQUEST_MODEL]?.value).toBe('anthropic/claude-haiku-4.5'); // pi-durable retries a provider error inside the run; such a request has no usage and no HTTP span diff --git a/dev-packages/node-integration-tests/suites/tracing/flue/test.ts b/dev-packages/node-integration-tests/suites/tracing/flue/test.ts index 83c73526b43e..e64067f3de63 100644 --- a/dev-packages/node-integration-tests/suites/tracing/flue/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/flue/test.ts @@ -40,9 +40,11 @@ describe.runIf(supports({ min: 22 }))('Flue integration', () => { expect(root.is_segment).toBe(true); // Counted rather than looked up: the interceptor skips the submission wrapper - // operation, so one dispatch opens exactly one agent span, and each turn and tool call - // is spanned once. `find` passes just as happily on a duplicate. - expect(spans.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.flue')).toHaveLength(4); + // operation, so one dispatch opens exactly one agent span, and each tool call is spanned + // once. Each turn is spanned once by the pi-ai integration: a provider SDK reporting the + // same request would add a span. `find` passes just as happily on a duplicate. + expect(spans.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.flue')).toHaveLength(2); + expect(spans.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.pi_ai')).toHaveLength(2); const agents = spans.filter(span => span.name === 'invoke_agent Hello'); const chats = spans.filter(span => span.name === 'chat faux-model'); @@ -62,17 +64,17 @@ describe.runIf(supports({ min: 22 }))('Flue integration', () => { expect(conversationId).toEqual(expect.any(String)); // Both turns, not just the first: they leave the provider by different paths (a tool - // call, then a final answer) and resolve their parent through separate tracker lookups. + // call, then a final answer). for (const chat of chats) { expect(chat.attributes[SENTRY_OP]?.value).toBe(GEN_AI_CHAT); - expect(chat.attributes[SENTRY_ORIGIN]?.value).toBe('auto.ai.flue'); + 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_CONVERSATION_ID]?.value).toBe(conversationId); expect(chat.parent_span_id).toBe(agent.span_id); expect(chat.attributes[GEN_AI_USAGE_INPUT_TOKENS]?.value).toBeGreaterThan(0); expect(chat.attributes[GEN_AI_USAGE_OUTPUT_TOKENS]?.value).toBeGreaterThan(0); expect(chat.attributes[GEN_AI_USAGE_TOTAL_TOKENS]?.value).toBeGreaterThan(0); - // Flue computes cost itself; no provider SDK reports it. The faux provider prices + // pi-ai computes cost itself; no provider SDK reports it. The faux provider prices // every model at zero, so this only proves the attribute is mapped. expect(chat.attributes[GEN_AI_COST_TOTAL_TOKENS]?.value).toEqual(expect.any(Number)); } 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 index e0d424cdae71..46a27dd03b74 100644 --- a/dev-packages/node-integration-tests/suites/tracing/pi-ai/instrument.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/pi-ai/instrument.mjs @@ -6,5 +6,4 @@ Sentry.init({ release: '1.0', tracesSampleRate: 1.0, transport: loggingTransport, - integrations: [Sentry.piAiIntegration()], }); diff --git a/dev-packages/node-integration-tests/suites/tracing/pi-durable/test.ts b/dev-packages/node-integration-tests/suites/tracing/pi-durable/test.ts index f26aa34b63d7..38ad1a417e08 100644 --- a/dev-packages/node-integration-tests/suites/tracing/pi-durable/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/pi-durable/test.ts @@ -60,10 +60,12 @@ describe.runIf(supports({ min: 22 }))('pi-durable integration', () => { const agent = spans.find(span => span.is_segment)!; expect(agent.name).toBe('invoke_agent'); expect(spans.filter(span => span.name === 'chat faux-model')).toHaveLength(2); - // One span per run, request and tool call: a tool wrapped twice would add a span. + // One span per run and tool call, and one per request from the pi-ai integration: a tool + // wrapped twice, or a request reported by a provider SDK too, would add a span. expect(spans.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.pi_durable')).toHaveLength( - 5, + 3, ); + expect(spans.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.pi_ai')).toHaveLength(2); expect(agent.attributes[SENTRY_OP]?.value).toBe(GEN_AI_INVOKE_AGENT); expect(agent.attributes[SENTRY_ORIGIN]?.value).toBe('auto.ai.pi_durable'); @@ -77,7 +79,7 @@ describe.runIf(supports({ min: 22 }))('pi-durable integration', () => { for (const chat of chats) { expect(chat.parent_span_id).toBe(agent.span_id); expect(chat.attributes[SENTRY_OP]?.value).toBe(GEN_AI_CHAT); - expect(chat.attributes[SENTRY_ORIGIN]?.value).toBe('auto.ai.pi_durable'); + expect(chat.attributes[SENTRY_ORIGIN]?.value).toBe('auto.ai.pi_ai'); expect(chat.attributes[GEN_AI_PROVIDER_NAME]?.value).toBe('faux'); expect(chat.attributes[GEN_AI_REQUEST_MODEL]?.value).toBe('faux-model'); expect(chat.attributes[GEN_AI_CONVERSATION_ID]?.value).toBe(conversationId); @@ -491,10 +493,8 @@ describe.runIf(supports({ min: 22 }))('pi-durable integration', () => { }) .expect({ span: container => { - // 2. The provider threw: the generation faulted. - const chat = container.items.find(span => span.name === 'chat faux-model')!; - expect(chat.attributes[GEN_AI_RESPONSE_FINISH_REASONS]).toBeUndefined(); - expect(chat.attributes[SENTRY_STATUS_MESSAGE]?.value).toBe('internal_error'); + // 2. The provider threw in front of pi-ai, so there is no request to report. + expect(container.items.find(span => span.name === 'chat faux-model')).toBeUndefined(); const agent = container.items.find(span => span.is_segment)!; expect(agent.attributes[SENTRY_STATUS_MESSAGE]?.value).toBe('internal_error'); }, diff --git a/packages/server-utils/src/ai/flue/constants.ts b/packages/server-utils/src/ai/flue/constants.ts index 6a3bbaaefaa9..2f8a7bff4f6a 100644 --- a/packages/server-utils/src/ai/flue/constants.ts +++ b/packages/server-utils/src/ai/flue/constants.ts @@ -21,11 +21,9 @@ export const FLUE_INSTRUMENTATION_KEY = Symbol.for('sentry.flue.instrumentation' /** * The Flue execution operations we act on. * - * Only `AGENT` is spanned from the interceptor, and only at the outermost depth: Flue drives one - * LLM call through many `MODEL` operations (one per stream read), and its `AGENT` operation nests - * inside itself once per submission. The turn span is driven from the observation stream instead, - * where `turn_start`/`turn` are exactly one-to-one with a model call. `MODEL` and `TOOL` are - * intercepted only to make the already-open span active for the duration of the operation. + * Only `AGENT` is spanned from the interceptor, and only at the outermost depth: its operation nests + * inside itself once per submission. `TOOL` is intercepted to make the tool span `observe` opened + * active, and `MODEL` to hand the conversation to the pi-ai integration, which spans the model call. */ export const FLUE_OPERATION = { AGENT: 'agent', @@ -34,8 +32,8 @@ export const FLUE_OPERATION = { } as const; /** - * Cap on tracked turn and tool spans, matching `MAX_TRACKED_MASTRA_SPANS`. Both maps are keyed off - * an id that is only removed when the matching end observation arrives; a stream that is abandoned + * Cap on tracked tool spans and admissions, matching `MAX_TRACKED_MASTRA_SPANS`. Both maps are keyed + * off an id that is only removed when the matching end observation arrives; a run that is abandoned * mid-turn never emits one, so without a cap the map grows for the lifetime of the process. */ export const MAX_TRACKED_FLUE_SPANS = 1000; diff --git a/packages/server-utils/src/ai/flue/index.ts b/packages/server-utils/src/ai/flue/index.ts index b5f969c8a16e..4edda9f5ccac 100644 --- a/packages/server-utils/src/ai/flue/index.ts +++ b/packages/server-utils/src/ai/flue/index.ts @@ -9,6 +9,7 @@ import { startNewTrace, startSpan, withActiveSpan, + withScope, } from '@sentry/core'; import { GEN_AI_AGENT_NAME, @@ -18,18 +19,10 @@ import { } from '@sentry/conventions/attributes'; import type { GenAiOptions } from '../core/utils'; import { getGenAiSpanOp, resolveAIRecordingOptions } from '../core/utils'; -import { skipPiAiProviderIntegrations } from '../pi-ai/providers'; import { withCleanScopes } from '../../utils/withCleanScopes'; import { FLUE_INSTRUMENTATION_KEY, FLUE_OPERATION, FLUE_ORIGIN, MAX_TRACKED_FLUE_SPANS } from './constants'; import type { SpanTracker } from './utils'; -import { - endToolSpan, - endTurnSpan, - recordRequestContent, - sentryTraceFromTraceparent, - startToolSpan, - startTurnSpan, -} from './utils'; +import { endToolSpan, sentryTraceFromTraceparent, startToolSpan } from './utils'; import type { FlueInstrumentation } from './types'; export type FlueOptions = GenAiOptions; @@ -74,31 +67,21 @@ function isCloudflareWorkers(): boolean { * * - `interceptor` wraps agent execution, so the agent span is *active* for its duration and every * span opened underneath parents correctly. - * - `observe` opens and closes the turn span, because Flue's `turn_start`/`turn` events are the - * only one-to-one signal for a model call and `turn` is what carries usage and cost. + * - `observe` opens and closes the tool spans, and names the agent span. + * + * Model calls get no span here: Flue sends them through pi-ai, whose integration reports each one + * as a `gen_ai.chat` span under the agent span. * * Message content, tool arguments and tool results are gated on `recordInputs`/`recordOutputs`, * which fall back to the current client's `dataCollection.genAI` settings and are read per event. */ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstrumentation { - // Flue drives the providers through `@earendil-works/pi-ai`, which sends its requests through the - // provider SDKs that Sentry instruments too (see `skipPiAiProviderIntegrations`). Left alone, their - // integrations report the same request a second time, as a second `gen_ai.chat` beside the turn's. - // - // The provider integrations are skipped on first use rather than here, for two reasons. - // Constructing the object proves nothing: if `instrument()` rejects it, suppressing the provider - // integrations would leave the app with no `gen_ai.chat` spans at all. And the registry is reset per - // client (`_setupIntegrations` clears it, and Cloudflare calls `init()` per request), so a one-shot - // call at module scope is wiped by the next `init()` and every later request double-reports. - const skipProviders = skipPiAiProviderIntegrations; - // Keyed by the agent operation's own id, which is what the observations carry. That keeps // concurrent runs apart and gives a delegated subagent its own span: Flue nests a second `agent` // operation inside the parent's for `task` delegation, and the nesting is not bounded at two. // A plain map: the entry is removed in a `finally`, so it is bounded by concurrent agent runs. const agentSpans = new Map(); // Capped, unlike the above: these are keyed off ids that only a matching end observation removes. - const turnSpans: SpanTracker = new LRUMap(MAX_TRACKED_FLUE_SPANS); const toolSpans: SpanTracker = new LRUMap(MAX_TRACKED_FLUE_SPANS); // Capped: a submission admitted here but run by another process never settles here. const admissions = new LRUMap(MAX_TRACKED_FLUE_SPANS); @@ -107,17 +90,22 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru key: FLUE_INSTRUMENTATION_KEY, interceptor: async (operation, ctx, next) => { - skipProviders(); - // `observe` has already opened the span for this unit of work; make it active for the // duration so whatever the tool or model call does lands inside it rather than beside it. if (operation?.type === FLUE_OPERATION.TOOL) { const toolSpan = operation.toolCallId ? toolSpans.get(operation.toolCallId) : undefined; return toolSpan ? withActiveSpan(toolSpan, next) : next(); } + // The `gen_ai.chat` span of the model call comes from the pi-ai integration, which reads the + // conversation from the scope. if (operation?.type === FLUE_OPERATION.MODEL) { - const turnSpan = operation.turnId ? turnSpans.get(operation.turnId) : undefined; - return turnSpan ? withActiveSpan(turnSpan, next) : next(); + const { conversationId } = ctx; + return conversationId + ? withScope(scope => { + scope.setConversationId(conversationId); + return next(); + }) + : next(); } if (operation?.type !== FLUE_OPERATION.AGENT) { @@ -172,16 +160,14 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru }, observe: observation => { - skipProviders(); - // Resolved per observation, not once at construction: `resolveAIRecordingOptions` reads the // current client's `dataCollection.genAI`, and Cloudflare replaces the client per request, so // values captured at isolate load would be the wrong ones for every later request. const { recordInputs, recordOutputs } = resolveAIRecordingOptions(options); // Observations are emitted synchronously from inside the agent operation, so the active span - // here is the agent span the interceptor opened — turn and tool spans parent off it without - // any bookkeeping. + // here is the agent span the interceptor opened — tool spans parent off it without any + // bookkeeping. // // The agent operation that gets the span is the one the observations reference, and it knows // neither the agent's name nor the conversation: the submission wrapper carries the name, the @@ -212,17 +198,6 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru admissions.remove(observation.submissionId); } return; - case 'turn_start': - startTurnSpan(observation, turnSpans); - return; - case 'turn_request': - if (recordInputs) { - recordRequestContent(observation, turnSpans); - } - return; - case 'turn': - endTurnSpan(observation, turnSpans, recordOutputs); - return; case 'tool_start': startToolSpan(observation, toolSpans, recordInputs); return; @@ -234,17 +209,16 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru } }, - // Only the turn and tool spans are ended here. They come from `startInactiveSpan`, so nothing - // else will. The agent spans come from `startSpan`, which ends them when its callback settles — - // ending them here would stamp an early end time on work that is still running, and the later - // `end()` would be ignored. + // Only the tool spans are ended here. They come from `startInactiveSpan`, so nothing else will. + // The agent spans come from `startSpan`, which ends them when its callback settles — ending them + // here would stamp an early end time on work that is still running, and the later `end()` would + // be ignored. dispose: () => { - for (const span of [...turnSpans.values(), ...toolSpans.values()]) { + for (const span of toolSpans.values()) { span.end(); } agentSpans.clear(); admissions.clear(); - turnSpans.clear(); toolSpans.clear(); }, }; diff --git a/packages/server-utils/src/ai/flue/types.ts b/packages/server-utils/src/ai/flue/types.ts index 30850dea9949..21f88c26bb49 100644 --- a/packages/server-utils/src/ai/flue/types.ts +++ b/packages/server-utils/src/ai/flue/types.ts @@ -6,44 +6,8 @@ * `FlueExecutionContext` as of `@flue/runtime` 2.x. */ -import type { PiAiUsage } from '../pi-ai/usage'; - -/** Mirrors `ModelRequestInfo`. */ -export interface FlueModelRequestInfo { - requestedModel?: string; - /** Provider slug (`anthropic`, `openrouter`) — what `gen_ai.provider.name` wants. */ - providerId?: string; - temperature?: number; - maxTokens?: number; - reasoningLevel?: string; - serverAddress?: string; - serverPort?: number; -} - -/** Mirrors `ModelRequestInput` — the content half of `ModelRequest`, on `turn_request` only. */ -export interface FlueModelRequestInput { - systemPrompt?: string; - messages?: unknown[]; - tools?: unknown[]; -} - -/** `turn_request` carries `ModelRequest`, which is `ModelRequestInfo` plus the input. */ -export interface FlueModelRequest extends FlueModelRequestInfo { - input?: FlueModelRequestInput; -} - -/** Mirrors `ModelResponse`. */ -export interface FlueModelResponse { - responseId?: string; - responseModel?: string; - output?: unknown; - /** pi-ai's usage, which Flue passes through with its own cost figures. */ - usage?: PiAiUsage; - finishReason?: string; -} - /** - * Mirrors Flue's `errorInfo`, which it attaches to a failed tool or turn observation. Flue catches + * Mirrors Flue's `errorInfo`, which it attaches to a failed tool observation. Flue catches * the throw and turns it into a tool result, so this serialized form is the only trace of it. */ export interface FlueErrorInfo { @@ -62,7 +26,6 @@ export interface FlueObservation { agentName?: string; conversationId?: string; session?: string; - turnId?: string; taskId?: string; submissionId?: string; operationId?: string; @@ -70,14 +33,11 @@ export interface FlueObservation { toolCallId?: string; isError?: boolean; errorInfo?: FlueErrorInfo; - purpose?: string; durationMs?: number; - request?: FlueModelRequest; args?: unknown; result?: unknown; /** The tool result as the model sees it, on successful `tool` events. */ effectiveResult?: unknown; - response?: FlueModelResponse; } /** The execution unit an interceptor wraps. */ @@ -85,7 +45,6 @@ export interface FlueExecutionOperation { type: string; operationId?: string; operationKind?: string; - turnId?: string; toolCallId?: string; } @@ -101,7 +60,6 @@ export interface FlueExecutionContext { submissionId?: string; conversationId?: string; session?: string; - turnId?: string; taskId?: string; } diff --git a/packages/server-utils/src/ai/flue/utils.ts b/packages/server-utils/src/ai/flue/utils.ts index 399a2aec3580..307642e20cb3 100644 --- a/packages/server-utils/src/ai/flue/utils.ts +++ b/packages/server-utils/src/ai/flue/utils.ts @@ -9,36 +9,16 @@ import { } from '@sentry/core'; import { GEN_AI_CONVERSATION_ID, - GEN_AI_INPUT_MESSAGES, GEN_AI_OPERATION_NAME, - 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_SYSTEM_INSTRUCTIONS, GEN_AI_TOOL_CALL_ARGUMENTS, GEN_AI_TOOL_CALL_RESULT, - GEN_AI_TOOL_DEFINITIONS, GEN_AI_TOOL_NAME, - SERVER_ADDRESS, - SERVER_PORT, SENTRY_ORIGIN, } from '@sentry/conventions/attributes'; import { getGenAiSpanOp } from '../core/utils'; -import { - piAiAssistantMessageToGenAiMessage, - piAiContentToString, - piAiFinishReason, - piAiMessagesToGenAiMessages, -} from '../pi-ai/messages'; -import { setPiAiUsageAttributes } from '../pi-ai/usage'; +import { piAiContentToString } from '../pi-ai/messages'; import { FLUE_ORIGIN, MAX_TRACKED_FLUE_SPANS } from './constants'; -import type { FlueErrorInfo, FlueModelRequestInfo, FlueObservation } from './types'; +import type { FlueErrorInfo, FlueObservation } from './types'; /** * Flue persists the incoming W3C `traceparent` at admission and replays it on the agent operation. @@ -108,85 +88,6 @@ function captureToolError(span: Span, errorInfo: FlueErrorInfo | undefined): voi }); } -export function startTurnSpan(observation: FlueObservation, turnSpans: SpanTracker): void { - const { turnId } = observation; - if (!turnId || turnSpans.get(turnId)) { - return; - } - - trackSpan( - turnSpans, - turnId, - startInactiveSpan({ - name: 'chat', - op: getGenAiSpanOp('chat'), - attributes: { - [SENTRY_ORIGIN]: FLUE_ORIGIN, - [GEN_AI_OPERATION_NAME]: 'chat', - [GEN_AI_CONVERSATION_ID]: observation.conversationId, - // No conventional attribute for this; it is the only way to tell a compaction turn from a - // user-facing one. - 'flue.turn.purpose': observation.purpose, - }, - }), - ); -} - -export function endTurnSpan(observation: FlueObservation, turnSpans: SpanTracker, recordOutputs: boolean): void { - const { turnId } = observation; - const span = turnId ? turnSpans.get(turnId) : undefined; - if (!span || !turnId) { - return; - } - turnSpans.remove(turnId); - - const requestedModel = observation.request?.requestedModel; - const responseModel = observation.response?.responseModel; - const model = responseModel ?? requestedModel; - if (model) { - span.updateName(`chat ${model}`); - } - if (requestedModel) { - span.setAttribute(GEN_AI_REQUEST_MODEL, requestedModel); - } - if (responseModel) { - span.setAttribute(GEN_AI_RESPONSE_MODEL, responseModel); - } - - // `providerId` is the slug (`anthropic`, `openrouter`); `providerName` is the display name. - const provider = observation.request?.providerId; - if (provider) { - span.setAttribute(GEN_AI_PROVIDER_NAME, provider); - } - - setRequestAttributes(span, observation.request); - - const responseId = observation.response?.responseId; - if (responseId) { - span.setAttribute(GEN_AI_RESPONSE_ID, responseId); - } - const finishReason = piAiFinishReason(observation.response?.finishReason); - if (finishReason) { - // Serialized, not a raw array: the conventions declare this attribute's value type as `string`, - // and that is what `ai/core`, Mastra, OpenAI and Vercel AI all write. - span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, stringify([finishReason])); - } - - const output = recordOutputs - ? piAiAssistantMessageToGenAiMessage(observation.response?.output, finishReason) - : undefined; - if (output) { - span.setAttribute(GEN_AI_OUTPUT_MESSAGES, stringify([output])); - } - - setPiAiUsageAttributes(span, observation.response?.usage, observation.isError); - - if (observation.isError) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - } - span.end(); -} - /** * Tool spans hang off the agent invocation rather than the turn, matching how Flue's own * OpenTelemetry adapter projects them: siblings of `chat`, correlated to model output by tool call @@ -257,52 +158,3 @@ function toolResultForModel(observation: FlueObservation): string | undefined { } return result === undefined ? undefined : stringify(result); } - -/** - * `turn_request` is the only event carrying the request's content — the settled `turn` reports - * metadata alone — so input messages, system prompt and tool definitions are read from it. - */ -export function recordRequestContent(observation: FlueObservation, turnSpans: SpanTracker): void { - const { turnId } = observation; - const span = turnId ? turnSpans.get(turnId) : undefined; - const input = observation.request?.input; - if (!span || !input) { - return; - } - - if (input.systemPrompt) { - span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, input.systemPrompt); - } - const messages = piAiMessagesToGenAiMessages(input.messages); - if (messages.length) { - span.setAttribute(GEN_AI_INPUT_MESSAGES, stringify(messages)); - } - if (input.tools?.length) { - span.setAttribute(GEN_AI_TOOL_DEFINITIONS, stringify(input.tools)); - } -} - -/** - * Model-call tuning and the provider endpoint, all on the settled turn's `ModelRequestInfo`. These - * are the conventional attributes the other AI integrations in this package set. - */ -export function setRequestAttributes(span: Span, request: FlueModelRequestInfo | undefined): void { - if (!request) { - return; - } - - const attributes: Record = {}; - const set = (key: string, value: string | number | undefined): void => { - if (value !== undefined) { - attributes[key] = value; - } - }; - - set(GEN_AI_REQUEST_TEMPERATURE, request.temperature); - set(GEN_AI_REQUEST_MAX_TOKENS, request.maxTokens); - set(GEN_AI_REQUEST_REASONING_LEVEL, request.reasoningLevel); - set(SERVER_ADDRESS, request.serverAddress); - set(SERVER_PORT, request.serverPort); - - span.setAttributes(attributes); -} diff --git a/packages/server-utils/src/ai/pi-ai/chat.ts b/packages/server-utils/src/ai/pi-ai/chat.ts index 2e401efcba2e..3d35b6c443fe 100644 --- a/packages/server-utils/src/ai/pi-ai/chat.ts +++ b/packages/server-utils/src/ai/pi-ai/chat.ts @@ -28,6 +28,7 @@ import { SERVER_PORT, } from '@sentry/conventions/attributes'; import { GEN_AI_CHAT } from '@sentry/conventions/op'; +import { PI_AI_ORIGIN } from './constants'; import type { PiAiContext } from './messages'; import { piAiAssistantMessageToGenAiMessage, @@ -45,7 +46,6 @@ export interface PiAiChatRequest { context?: PiAiContext; callOptions?: PiStreamOptions; streaming: boolean; - origin: string; recordInputs: boolean; } @@ -54,7 +54,7 @@ export interface PiAiChatRequest { * the caller picks whether the span is active while the request is sent. */ export function getPiAiChatSpanOptions(request: PiAiChatRequest): StartSpanOptions { - const { model, context, callOptions, streaming, origin, recordInputs } = request; + const { model, context, callOptions, streaming, recordInputs } = request; const conversationId = getCurrentScope().getScopeData().conversationId ?? getIsolationScope().getScopeData().conversationId; @@ -62,7 +62,7 @@ export function getPiAiChatSpanOptions(request: PiAiChatRequest): StartSpanOptio name: model.id ? `chat ${model.id}` : 'chat', op: GEN_AI_CHAT, attributes: { - [SENTRY_ORIGIN]: origin, + [SENTRY_ORIGIN]: PI_AI_ORIGIN, [GEN_AI_OPERATION_NAME]: 'chat', [GEN_AI_PROVIDER_NAME]: model.provider, [GEN_AI_REQUEST_MODEL]: model.id, diff --git a/packages/server-utils/src/ai/pi-ai/providers.ts b/packages/server-utils/src/ai/pi-ai/providers.ts index 96bc6149f620..fc2c93427adf 100644 --- a/packages/server-utils/src/ai/pi-ai/providers.ts +++ b/packages/server-utils/src/ai/pi-ai/providers.ts @@ -6,8 +6,8 @@ import { WORKERS_AI_INTEGRATION_NAME } from '../workers-ai/constants'; // pi-ai sends its requests through the `openai`, `@anthropic-ai/sdk` and `@google/genai` clients, and // on Cloudflare through the Workers AI binding (`createAI()` of `agents/models/pi-ai`). Left alone, -// those integrations report the same request a second time beside the `chat` span of the framework -// that sent it. Bedrock requests go through `@aws-sdk/client-bedrock-runtime`, which `awsIntegration` +// those integrations report the same request a second time beside the pi-ai integration's `chat` +// span. Bedrock requests go through `@aws-sdk/client-bedrock-runtime`, which `awsIntegration` // still reports; that one has no skip yet. const PI_AI_PROVIDER_INTEGRATIONS = [ OPENAI_INTEGRATION_NAME, diff --git a/packages/server-utils/src/ai/pi-durable/index.ts b/packages/server-utils/src/ai/pi-durable/index.ts index d2187aae394e..3e9221a0ec31 100644 --- a/packages/server-utils/src/ai/pi-durable/index.ts +++ b/packages/server-utils/src/ai/pi-durable/index.ts @@ -10,9 +10,7 @@ import { } from '@sentry/core'; import { GEN_AI_CONVERSATION_ID } from '@sentry/conventions/attributes'; import type { GenAiOptions } from '../core/utils'; -import { skipPiAiProviderIntegrations } from '../pi-ai/providers'; import { PI_DURABLE_ORIGIN, PI_TASK } from './constants'; -import { instrumentPiModels } from './models'; import type { PiRuns } from './runs'; import { createRuns, endRun, observeRunControl, readProviderSessionId, startRun, toSentryConversationId } from './runs'; import { instrumentTool, instrumentToolPhase } from './tools'; @@ -35,15 +33,15 @@ const INSTRUMENTED = Symbol.for('sentry.pi-durable.instrumented'); const RUNS = Symbol.for('sentry.pi-durable.runs'); /** - * Return pi-durable `HarnessOptions` whose model requests, task phases and tool calls are traced. - * The caller's `models` and `registry` objects are wrapped, never modified. The result has the - * caller's options as prototype: pi-durable reads `settings`, `env` and `conversationCreated` at - * every use, so they stay live, getters included. + * Return pi-durable `HarnessOptions` whose task phases and tool calls are traced. The caller's + * `registry` object is wrapped, never modified. The result has the caller's options as prototype: + * pi-durable reads `settings`, `env` and `conversationCreated` at every use, so they stay live, + * getters included. * * Spans: * - `gen_ai.invoke_agent` for each run, a new trace per run, or a child of the tool call that * created the conversation for a subagent; - * - `gen_ai.chat` for each model request, under its run; + * - `gen_ai.chat` for each model request, under its run, from `piAiIntegration`; * - `gen_ai.execute_tool` for each tool call, under its run. * * The scheduler starts task phases from whichever async context last woke it, often an unrelated @@ -61,9 +59,6 @@ export function instrumentPiDurableHarnessOptions( const runs = createRuns(); const own = (value: unknown): PropertyDescriptor => ({ value, enumerable: true, configurable: true, writable: true }); const instrumented: T = Object.create(harnessOptions, { - ...(harnessOptions.models - ? { models: own(instrumentPiModels(harnessOptions.models, options, skipPiAiProviderIntegrations)) } - : {}), ...(harnessOptions.registry ? { registry: own(instrumentRegistry(harnessOptions.registry, options, runs)) } : {}), // pi-durable passes failures of extension code here, such as a throwing hook, and keeps going. // Only an app that passed its own `onReport` gets a chance to handle them. diff --git a/packages/server-utils/src/ai/pi-durable/models.ts b/packages/server-utils/src/ai/pi-durable/models.ts deleted file mode 100644 index 07cf3a279165..000000000000 --- a/packages/server-utils/src/ai/pi-durable/models.ts +++ /dev/null @@ -1,113 +0,0 @@ -import { isThenable, SPAN_STATUS_ERROR, startSpanManual } from '@sentry/core'; -import type { GenAiOptions } from '../core/utils'; -import { resolveAIRecordingOptions } from '../core/utils'; -import { getPiAiChatSpanOptions, setPiAiResponseAttributes } from '../pi-ai/chat'; -import type { PiAiContext } from '../pi-ai/messages'; -import type { PiAssistantMessage, PiEventStream, PiModel, PiStreamOptions } from '../pi-ai/types'; -import { PI_DURABLE_ORIGIN } from './constants'; -import type { PiModels } from './types'; - -type ModelCall = (model: PiModel, request: unknown, options?: unknown) => unknown; - -const STREAMING_METHODS = new Set(['stream', 'streamSimple', 'streamDeferred']); -const COMPLETE_METHODS = new Set(['complete', 'completeSimple', 'fetchDeferred']); -// These fetch the answer of a parked request; their second argument is a handle, not the request. -const DEFERRED_METHODS = new Set(['streamDeferred', 'fetchDeferred']); - -/** - * Wrap a pi-ai `Models` so every model request it sends gets a `gen_ai.chat` span. - * - * The proxy hands every other member through bound to the original, so a `Models` instance whose - * methods read `this` keeps working, and the caller's own object is never modified. - */ -export function instrumentPiModels(models: T, options: GenAiOptions, onRequest: () => void): T { - const wrapped = new Map(); - - return new Proxy(models, { - get(target, property) { - const value: unknown = Reflect.get(target, property, target); - if (typeof value !== 'function') { - return value; - } - - const cached = wrapped.get(property); - if (cached) { - return cached; - } - - const original = (value as ModelCall).bind(target); - const method = String(property); - const streaming = STREAMING_METHODS.has(method); - const deferred = DEFERRED_METHODS.has(method); - const instrumented = - streaming || COMPLETE_METHODS.has(method) - ? (model: PiModel, request: unknown, callOptions?: unknown): unknown => { - onRequest(); - return traceModelRequest( - model, - deferred ? undefined : (request as PiAiContext), - deferred ? undefined : (callOptions as PiStreamOptions | undefined), - streaming, - options, - () => original(model, request, callOptions), - ); - } - : original; - - wrapped.set(property, instrumented); - return instrumented; - }, - }); -} - -function traceModelRequest( - model: PiModel, - context: PiAiContext | undefined, - callOptions: PiStreamOptions | undefined, - streaming: boolean, - options: GenAiOptions, - send: () => unknown, -): unknown { - // Read per request: `dataCollection.genAI` belongs to the current client, which can change. - const { recordInputs, recordOutputs } = resolveAIRecordingOptions(options); - - return startSpanManual( - getPiAiChatSpanOptions({ model, context, callOptions, streaming, origin: PI_DURABLE_ORIGIN, recordInputs }), - span => { - let result: unknown; - try { - result = send(); - } catch (error) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - span.end(); - throw error; - } - - // pi-ai never rejects a request: a provider failure or an abort settles as a message with - // `stopReason` `error` or `aborted`. The rejection handler only covers a broken provider. - let settled: unknown; - try { - settled = streaming ? (result as PiEventStream).result() : result; - } catch { - settled = undefined; - } - // A result of another shape must reach the caller unchanged, so its span ends without a response. - if (!isThenable(settled)) { - span.end(); - return result; - } - (settled as Promise).then( - message => { - setPiAiResponseAttributes(span, message, recordOutputs); - span.end(); - }, - () => { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - span.end(); - }, - ); - - return result; - }, - ); -} diff --git a/packages/server-utils/src/ai/pi-durable/types.ts b/packages/server-utils/src/ai/pi-durable/types.ts index e4b79f7fcd42..38c4557c8707 100644 --- a/packages/server-utils/src/ai/pi-durable/types.ts +++ b/packages/server-utils/src/ai/pi-durable/types.ts @@ -1,26 +1,9 @@ /** - * Structural subsets of the `@earendil-works/pi-durable` and `@earendil-works/pi-ai` types the - * instrumentation reads. The SDK does not depend on either package, so only the fields used here - * are declared, and every one of them is treated as possibly absent at runtime. + * Structural subsets of the `@earendil-works/pi-durable` types the instrumentation reads. The SDK + * does not depend on the package, so only the fields used here are declared, and every one of them + * is treated as possibly absent at runtime. */ -import type { PiAiContext } from '../pi-ai/messages'; -import type { PiAssistantMessage, PiEventStream, PiModel, PiStreamOptions } from '../pi-ai/types'; - -/** - * The `Models` methods that send a model request. Everything else passes through untouched. The - * deferred methods fetch the answer of a request the provider parked, so they take a handle in - * place of the request context. - */ -export interface PiModels { - stream?(model: PiModel, context: PiAiContext, options?: PiStreamOptions): PiEventStream; - streamSimple?(model: PiModel, context: PiAiContext, options?: PiStreamOptions): PiEventStream; - complete?(model: PiModel, context: PiAiContext, options?: PiStreamOptions): Promise; - completeSimple?(model: PiModel, context: PiAiContext, options?: PiStreamOptions): Promise; - streamDeferred?(model: PiModel, handle: unknown, options?: unknown): PiEventStream; - fetchDeferred?(model: PiModel, handle: unknown, options?: unknown): Promise; -} - export interface PiToolExecutionResult { /** The content blocks the tool returns for the model, in pi-durable 1.2.0 and later. */ output?: unknown[]; @@ -111,7 +94,6 @@ export interface PiRegistryReader { } export interface PiHarnessOptions { - models?: PiModels; registry?: PiRegistryReader; onReport?: (error: unknown) => void; } diff --git a/packages/server-utils/src/integrations/index.ts b/packages/server-utils/src/integrations/index.ts index 9643f95250ff..45eff994348f 100644 --- a/packages/server-utils/src/integrations/index.ts +++ b/packages/server-utils/src/integrations/index.ts @@ -15,6 +15,7 @@ import { mongooseIntegration } from './mongoose'; import { lruMemoizerIntegration } from './lru-memoizer'; import { langChainIntegration } from './langchain'; import { langGraphIntegration } from './langgraph'; +import { piAiIntegration } from './pi-ai'; import { piDurableIntegration } from './pi-durable'; import { mastraIntegration } from './mastra'; import { mcpServerIntegration } from './mcp-server'; @@ -62,6 +63,7 @@ export function getTracingIntegrations(): Integration[] { langChainIntegration(), langGraphIntegration(), mastraIntegration(), + piAiIntegration(), piDurableIntegration(), vercelAIIntegration(), openAIIntegration(), diff --git a/packages/server-utils/src/integrations/pi-ai.ts b/packages/server-utils/src/integrations/pi-ai.ts index 7ac0a34fd781..123ee3f7f89b 100644 --- a/packages/server-utils/src/integrations/pi-ai.ts +++ b/packages/server-utils/src/integrations/pi-ai.ts @@ -4,7 +4,7 @@ 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 { PI_AI_INTEGRATION_NAME } 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'; @@ -45,7 +45,6 @@ function instrumentPiAi(options: PiAiOptions): void { context: context as PiAiContext | undefined, callOptions: callOptions as PiStreamOptions | undefined, streaming, - origin: PI_AI_ORIGIN, recordInputs: resolveAIRecordingOptions(options).recordInputs, }), ); @@ -94,10 +93,8 @@ function instrumentPiAi(options: PiAiOptions): void { /** * 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. + * runtime hook (`--import`) or bundler plugin. This includes the model calls of Flue and pi-durable + * agents, which send them through pi-ai. * * 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. @@ -105,7 +102,7 @@ function instrumentPiAi(options: PiAiOptions): void { * @example * ```javascript * Sentry.init({ - * integrations: [Sentry.piAiIntegration()], + * integrations: [Sentry.piAiIntegration({ recordInputs: true, recordOutputs: true })], * }); * ``` */ diff --git a/packages/server-utils/src/integrations/pi-durable.ts b/packages/server-utils/src/integrations/pi-durable.ts index a9bd9e1b6a99..f04214a7fb17 100644 --- a/packages/server-utils/src/integrations/pi-durable.ts +++ b/packages/server-utils/src/integrations/pi-durable.ts @@ -60,9 +60,8 @@ function instrumentPiDurable(options: PiDurableOptions): void { * channel injected into `Harness.open()` and to the `codingTool` channel injected into the built-in * tool factories, so it requires the Sentry runtime hook or bundler plugin. * - * Traces one `gen_ai.invoke_agent` span per run, with its model requests as `gen_ai.chat` and its - * tool calls as `gen_ai.execute_tool` children. From the first model request on, the `openai`, - * `@anthropic-ai/sdk`, `@google/genai` and Workers AI integrations stop reporting requests in the - * whole process, because pi-ai sends its requests through those clients. + * Traces one `gen_ai.invoke_agent` span per run, with its tool calls as `gen_ai.execute_tool` + * children. Its model requests become `gen_ai.chat` children through `piAiIntegration`, since + * pi-durable sends them through pi-ai. */ export const piDurableIntegration = defineIntegration(_piDurableIntegration); diff --git a/packages/server-utils/src/orchestrion/config/channel-integration-definitions.ts b/packages/server-utils/src/orchestrion/config/channel-integration-definitions.ts index 46cd69d31bf9..635734bf7cdf 100644 --- a/packages/server-utils/src/orchestrion/config/channel-integration-definitions.ts +++ b/packages/server-utils/src/orchestrion/config/channel-integration-definitions.ts @@ -50,6 +50,7 @@ export const CHANNEL_INTEGRATION_DEFINITIONS = [ { exportName: 'langGraphIntegration', modules: ['@langchain/langgraph'] }, { exportName: 'mastraIntegration', modules: ['@mastra/core'] }, { exportName: 'flueIntegration', modules: ['@flue/runtime'] }, + { exportName: 'piAiIntegration', modules: ['@earendil-works/pi-ai'] }, { exportName: 'piDurableIntegration', modules: ['@earendil-works/pi-durable'] }, { exportName: 'mcpServerIntegration', modules: ['@modelcontextprotocol/server', '@modelcontextprotocol/sdk'] }, { exportName: 'awsIntegration', modules: ['@aws-sdk/smithy-client', '@smithy/core', '@smithy/smithy-client'] }, diff --git a/packages/server-utils/test/ai/lib/tracing/flue.test.ts b/packages/server-utils/test/ai/lib/tracing/flue.test.ts index e39ed3fd3bfb..a9e8c8c50743 100644 --- a/packages/server-utils/test/ai/lib/tracing/flue.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/flue.test.ts @@ -1,9 +1,7 @@ import { afterEach, beforeEach, describe, expect, it, onTestFinished, vi } from 'vitest'; import type { Event, Span } from '@sentry/core'; import { - _INTERNAL_clearAiProviderSkips, - _INTERNAL_shouldSkipAiProviderWrapping, - _INTERNAL_skipAiProviderWrapping, + getCurrentScope, getIsolationScope, getMainCarrier, Scope, @@ -15,58 +13,21 @@ import { withActiveSpan, withIsolationScope, } from '@sentry/core'; -import { ANTHROPIC_AI_INTEGRATION_NAME } from '../../../../src/ai/anthropic-ai/constants'; import { setAsyncLocalStorageAsyncContextStrategy } from '../../../../src/async-context'; import { createFlueInstrumentation } from '../../../../src/ai/flue'; -import { MAX_TRACKED_FLUE_SPANS } from '../../../../src/ai/flue/constants'; -import { GOOGLE_GENAI_INTEGRATION_NAME } from '../../../../src/ai/google-genai/constants'; -import type { FlueInstrumentation, FlueObservation } from '../../../../src/ai/flue/types'; -import { OPENAI_INTEGRATION_NAME } from '../../../../src/ai/openai/constants'; +import type { FlueInstrumentation } from '../../../../src/ai/flue/types'; import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; const AGENT_OP = { type: 'agent', operationId: 'op_1' }; const AGENT_CTX = { agentName: 'Hello', submissionId: 'sub_1' }; const INNER_CTX = { conversationId: 'conv_1' }; -/** A settled turn as Flue reports it, with the field names `ModelRequestInfo`/`ModelResponse` use. */ -function turn(overrides: Partial = {}): FlueObservation { - return { - type: 'turn', - turnId: 'turn_1', - submissionId: 'sub_1', - operationId: 'op_1', - request: { - requestedModel: 'claude-haiku-4.5', - providerId: 'anthropic', - temperature: 0.7, - maxTokens: 1024, - reasoningLevel: 'high', - serverAddress: 'api.anthropic.com', - serverPort: 443, - }, - response: { - responseId: 'resp_1', - finishReason: 'stop', - usage: { - input: 924, - output: 57, - totalTokens: 981, - cacheRead: 0, - cacheWrite: 0, - cost: { input: 0.000924, output: 0.000275, total: 0.001199, cacheRead: 0, cacheWrite: 0 }, - }, - }, - ...overrides, - }; -} - describe('createFlueInstrumentation', () => { let endedSpans: Span[]; let client: TestClient; let instrumentation: FlueInstrumentation; beforeEach(() => { - _INTERNAL_clearAiProviderSkips(); getMainCarrier().__SENTRY__ = undefined; client = new TestClient( getDefaultTestClientOptions({ @@ -84,7 +45,6 @@ describe('createFlueInstrumentation', () => { }); afterEach(() => { - _INTERNAL_clearAiProviderSkips(); getMainCarrier().__SENTRY__ = undefined; }); @@ -103,55 +63,6 @@ describe('createFlueInstrumentation', () => { return endedSpans.map(span => spanToStaticSpanJSON(span)).find(json => json.description === description); } - // Flue drives the providers through `pi-ai`, which bundles the `openai` / `@anthropic-ai/sdk` / - // `@google/genai` clients those integrations patch, so their spans duplicate the turn span. - // Not at construction: if `instrument()` rejects the object, suppressing the provider - // integrations would leave the app with no `gen_ai.chat` spans at all. - it('skips raw provider wrapping on first use, not on construction', async () => { - expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(false); - - await withAgent(() => undefined); - - expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(true); - expect(_INTERNAL_shouldSkipAiProviderWrapping(ANTHROPIC_AI_INTEGRATION_NAME)).toBe(true); - }); - - // The registry is reset per client (`_setupIntegrations` clears it, and Cloudflare calls `init()` - // per request), so a one-shot call at construction is wiped by the next reset. - it('re-applies the provider skip after the registry is cleared', async () => { - await withAgent(() => undefined); - expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(true); - - _INTERNAL_clearAiProviderSkips(); - expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(false); - - await withAgent(() => undefined); - expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(true); - }); - - // The guard has to hold for every provider, not just the first: another integration may have - // registered a skip for one of them already, which would otherwise short-circuit the rest. - it('applies the skip to every provider when only some are already registered', async () => { - _INTERNAL_skipAiProviderWrapping([OPENAI_INTEGRATION_NAME]); - - await withAgent(() => undefined); - - expect(_INTERNAL_shouldSkipAiProviderWrapping(ANTHROPIC_AI_INTEGRATION_NAME)).toBe(true); - expect(_INTERNAL_shouldSkipAiProviderWrapping(GOOGLE_GENAI_INTEGRATION_NAME)).toBe(true); - }); - - // A turn whose stream is abandoned never emits the settled `turn` that would remove it, so the - // tracker is capped. Eviction has to end the span it drops, or it is never sent. - it('ends the oldest chat span when the turn tracker overflows', async () => { - await withAgent(() => { - for (let i = 0; i <= MAX_TRACKED_FLUE_SPANS; i++) { - instrumentation.observe({ type: 'turn_start', turnId: `turn_${i}`, operationId: 'op_1' }, {}); - } - }); - - expect(endedSpans.filter(span => spanToStaticSpanJSON(span).data['sentry.op'] === 'gen_ai.chat')).toHaveLength(1); - }); - it('names agent spans `invoke_agent {name}` and sets the gen_ai op', async () => { await withAgent(() => undefined); @@ -166,10 +77,7 @@ describe('createFlueInstrumentation', () => { // the agent, and the conversation arrives on the observations that follow. it('sets the conversation id on the agent span from the observations', async () => { await withAgent(() => { - instrumentation.observe( - { type: 'turn_start', turnId: 'turn_1', operationId: 'op_1', conversationId: 'conv_1' }, - {}, - ); + instrumentation.observe({ type: 'turn_start', operationId: 'op_1', conversationId: 'conv_1' }, {}); }); expect(findSpan('invoke_agent Hello')?.data['gen_ai.conversation.id']).toBe('conv_1'); @@ -384,26 +292,31 @@ describe('createFlueInstrumentation', () => { }); it('does not span operations other than `agent`', async () => { - await instrumentation.interceptor({ type: 'model', turnId: 'turn_1' }, {}, async () => undefined); + await instrumentation.interceptor({ type: 'model' }, {}, async () => undefined); expect(endedSpans).toHaveLength(0); }); - // The `model` and `tool` operations open no span of their own; they make the span `observe` - // already opened active, so the provider's HTTP call and the tool's own work nest inside it - // rather than landing beside it as siblings of the agent invocation. - it('makes the turn span active for the model operation it wraps', async () => { - await withAgent(async () => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); - await instrumentation.interceptor({ type: 'model', turnId: 'turn_1' }, {}, async () => { - startInactiveSpan({ name: 'provider request' }).end(); - }); - instrumentation.observe(turn(), {}); - }); + // The model call's `gen_ai.chat` span comes from the pi-ai integration, which reads the + // conversation from the scope. + it('hands the conversation to the model operation it wraps through the scope', async () => { + let conversationId: string | undefined; - expect(findSpan('provider request')?.parent_span_id).toBe(findSpan('chat claude-haiku-4.5')?.span_id); + await withAgent(() => + instrumentation.interceptor({ type: 'model' }, INNER_CTX, async () => { + conversationId = getCurrentScope().getScopeData().conversationId; + }), + ); + + expect(conversationId).toBe('conv_1'); + expect(getCurrentScope().getScopeData().conversationId).toBeUndefined(); + expect(endedSpans.map(span => spanToStaticSpanJSON(span).description)).toEqual(['invoke_agent Hello']); }); + // The `tool` operation opens no span of its own; it makes the span `observe` already opened + // active, so the tool's own work nests inside it rather than landing beside it as a sibling of + // the agent invocation. + it('makes the tool span active for the tool operation it wraps', async () => { await withAgent(async () => { instrumentation.observe( @@ -419,163 +332,9 @@ describe('createFlueInstrumentation', () => { expect(findSpan('tool work')?.parent_span_id).toBe(findSpan('execute_tool get_weather')?.span_id); }); - it('opens a chat span on turn_start and completes it from the settled turn', async () => { - await withAgent(() => { - instrumentation.observe( - { type: 'turn_start', turnId: 'turn_1', operationId: 'op_1', conversationId: 'conv_1' }, - {}, - ); - instrumentation.observe(turn(), {}); - }); - - const json = findSpan('chat claude-haiku-4.5'); - expect(json?.data['sentry.op']).toBe('gen_ai.chat'); - expect(json?.data['sentry.origin']).toBe('auto.ai.flue'); - expect(json?.data['gen_ai.request.model']).toBe('claude-haiku-4.5'); - expect(json?.data['gen_ai.provider.name']).toBe('anthropic'); - expect(json?.data['gen_ai.response.id']).toBe('resp_1'); - expect(json?.data['gen_ai.response.finish_reasons']).toBe('["stop"]'); - expect(json?.data['gen_ai.conversation.id']).toBe('conv_1'); - }); - - // Flue computes costs itself; the provider SDKs report none. - it('records token usage and Flue-computed cost on the chat span', async () => { - await withAgent(() => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); - instrumentation.observe(turn(), {}); - }); - - const json = findSpan('chat claude-haiku-4.5'); - expect(json?.data['gen_ai.usage.input_tokens']).toBe(924); - expect(json?.data['gen_ai.usage.output_tokens']).toBe(57); - expect(json?.data['gen_ai.usage.total_tokens']).toBe(981); - expect(json?.data['gen_ai.cost.total_tokens']).toBe(0.001199); - }); - - // pi-ai keeps the cached tokens out of `input`; the conventions count them in. - it('counts cached tokens into the input tokens and their cost into the input cost', async () => { - const usage = { - input: 37, - output: 2, - cacheRead: 52, - cacheWrite: 38, - totalTokens: 129, - cost: { input: 0.000037, output: 0.00001, cacheRead: 0.0000052, cacheWrite: 0.0000475, total: 0.0000997 }, - }; - - await withAgent(() => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); - instrumentation.observe(turn({ response: { ...turn().response, usage } }), {}); - }); - - const json = findSpan('chat claude-haiku-4.5'); - expect(json?.data['gen_ai.usage.input_tokens']).toBe(127); - expect(json?.data['gen_ai.usage.cache_read.input_tokens']).toBe(52); - expect(json?.data['gen_ai.usage.cache_creation.input_tokens']).toBe(38); - expect(json?.data['gen_ai.cost.input_tokens']).toBe(0.000037 + 0.0000052 + 0.0000475); - expect(json?.data['gen_ai.cost.total_tokens']).toBe(0.0000997); - }); - - it('reports a tool-calling turn with the finish reason of the conventions', async () => { - await withAgent(() => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); - instrumentation.observe( - turn({ - response: { - ...turn().response, - finishReason: 'toolUse', - output: { - role: 'assistant', - content: [{ type: 'toolCall', id: 'c1', name: 'get_weather', arguments: { city: 'Berlin' } }], - }, - }, - }), - {}, - ); - }); - - const json = findSpan('chat claude-haiku-4.5'); - expect(json?.data['gen_ai.response.finish_reasons']).toBe('["tool_call"]'); - expect(json?.data['gen_ai.output.messages']).toBe( - '[{"role":"assistant","parts":[{"type":"tool_call","id":"c1","name":"get_weather","arguments":"{\\"city\\":\\"Berlin\\"}"}],"finish_reason":"tool_call"}]', - ); - }); - - it('records the model-call tuning and provider endpoint', async () => { - await withAgent(() => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1', purpose: 'agent' }, {}); - instrumentation.observe(turn(), {}); - }); - - const json = findSpan('chat claude-haiku-4.5'); - expect(json?.data['gen_ai.request.temperature']).toBe(0.7); - expect(json?.data['gen_ai.request.max_tokens']).toBe(1024); - expect(json?.data['gen_ai.request.reasoning.level']).toBe('high'); - expect(json?.data['server.address']).toBe('api.anthropic.com'); - expect(json?.data['server.port']).toBe(443); - // Distinguishes a compaction turn from a user-facing one. - expect(json?.data['flue.turn.purpose']).toBe('agent'); - }); - - it('marks a failed turn as errored', async () => { - await withAgent(() => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); - instrumentation.observe(turn({ isError: true }), {}); - }); - - expect(findSpan('chat claude-haiku-4.5')?.status).toBe('internal_error'); - }); - - // A turn that fails before the provider bills anything reports every counter as 0; writing those - // reads as a real zero-cost call. - it('omits usage entirely when a failed turn produced no tokens', async () => { - const empty = { - input: 0, - output: 0, - totalTokens: 0, - cacheRead: 0, - cacheWrite: 0, - cost: { input: 0, output: 0, total: 0, cacheRead: 0, cacheWrite: 0 }, - }; - - await withAgent(() => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); - instrumentation.observe(turn({ isError: true, response: { usage: empty } }), {}); - }); - - const json = findSpan('chat claude-haiku-4.5'); - expect(json).toBeDefined(); - expect(json?.data['gen_ai.usage.total_tokens']).toBeUndefined(); - expect(json?.data['gen_ai.cost.total_tokens']).toBeUndefined(); - }); - describe('content recording', () => { - const requestContent = { - type: 'turn_request', - turnId: 'turn_1', - request: { - requestedModel: 'claude-haiku-4.5', - input: { - systemPrompt: 'You are helpful.', - messages: [{ role: 'user', content: 'hi' }], - tools: [{ name: 'get_weather', description: 'weather', parameters: {} }], - }, - }, - } satisfies FlueObservation; - async function record(instr: FlueInstrumentation): Promise { await instr.interceptor(AGENT_OP, AGENT_CTX, async () => { - instr.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); - instr.observe(requestContent, {}); - instr.observe( - turn({ - response: { - ...turn().response, - output: { role: 'assistant', content: [{ type: 'text', text: 'It is sunny.' }] }, - }, - }), - {}, - ); instr.observe( { type: 'tool_start', @@ -601,18 +360,9 @@ describe('createFlueInstrumentation', () => { }); } - it('records messages, instructions, tool definitions, arguments and results by default', async () => { + it('records tool arguments and results by default', async () => { await record(instrumentation); - const chat = findSpan('chat claude-haiku-4.5'); - expect(chat?.data['gen_ai.system_instructions']).toBe('You are helpful.'); - // Mapped from pi-ai's message shape to the conventions, which is what Sentry renders. - expect(chat?.data['gen_ai.input.messages']).toBe('[{"role":"user","parts":[{"type":"text","content":"hi"}]}]'); - expect(chat?.data['gen_ai.output.messages']).toBe( - '[{"role":"assistant","parts":[{"type":"text","content":"It is sunny."}],"finish_reason":"stop"}]', - ); - expect(chat?.data['gen_ai.tool.definitions']).toContain('get_weather'); - const tool = findSpan('execute_tool get_weather'); expect(tool?.data['gen_ai.tool.call.arguments']).toBe('{"city":"Berlin"}'); expect(tool?.data['gen_ai.tool.call.result']).toBe('sunny'); @@ -692,53 +442,12 @@ describe('createFlueInstrumentation', () => { expect(findSpan('execute_tool get_weather')?.data['gen_ai.tool.call.result']).toBe(''); }); - it('maps a tool result in the request and a tool call in the response', async () => { - await withAgent(() => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); - instrumentation.observe( - { - ...requestContent, - request: { - ...requestContent.request, - input: { - messages: [ - { role: 'user', content: 'Weather in Berlin?' }, - { - role: 'assistant', - content: [{ type: 'toolCall', id: 'c1', name: 'get_weather', arguments: { city: 'Berlin' } }], - }, - { - role: 'toolResult', - toolCallId: 'c1', - toolName: 'get_weather', - content: [{ type: 'text', text: 'sunny' }], - isError: false, - }, - ], - }, - }, - }, - {}, - ); - instrumentation.observe(turn(), {}); - }); - - expect( - JSON.parse(String(findSpan('chat claude-haiku-4.5')?.data['gen_ai.input.messages'])).map( - (message: { role: string }) => message.role, - ), - ).toEqual(['user', 'assistant', 'tool']); - }); - it('omits inputs when recordInputs is false but keeps outputs', async () => { await record(createFlueInstrumentation({ recordInputs: false })); - const chat = findSpan('chat claude-haiku-4.5'); - expect(chat?.data['gen_ai.input.messages']).toBeUndefined(); - expect(chat?.data['gen_ai.system_instructions']).toBeUndefined(); - expect(chat?.data['gen_ai.tool.definitions']).toBeUndefined(); - expect(chat?.data['gen_ai.output.messages']).toBeDefined(); - expect(findSpan('execute_tool get_weather')?.data['gen_ai.tool.call.arguments']).toBeUndefined(); + const tool = findSpan('execute_tool get_weather'); + expect(tool?.data['gen_ai.tool.call.arguments']).toBeUndefined(); + expect(tool?.data['gen_ai.tool.call.result']).toBeDefined(); }); // The client is replaced per request on Cloudflare, so options captured once at construction @@ -746,7 +455,7 @@ describe('createFlueInstrumentation', () => { it('follows the current client when it is replaced', async () => { const instr = createFlueInstrumentation(); await record(instr); - expect(findSpan('chat claude-haiku-4.5')?.data['gen_ai.input.messages']).toBeDefined(); + expect(findSpan('execute_tool get_weather')?.data['gen_ai.tool.call.arguments']).toBeDefined(); endedSpans.length = 0; const strict = new TestClient( @@ -763,18 +472,17 @@ describe('createFlueInstrumentation', () => { await record(instr); - const chat = findSpan('chat claude-haiku-4.5'); - expect(chat).toBeDefined(); - expect(chat?.data['gen_ai.input.messages']).toBeUndefined(); + const tool = findSpan('execute_tool get_weather'); + expect(tool).toBeDefined(); + expect(tool?.data['gen_ai.tool.call.arguments']).toBeUndefined(); }); it('omits outputs when recordOutputs is false but keeps inputs', async () => { await record(createFlueInstrumentation({ recordOutputs: false })); - const chat = findSpan('chat claude-haiku-4.5'); - expect(chat?.data['gen_ai.output.messages']).toBeUndefined(); - expect(chat?.data['gen_ai.input.messages']).toBeDefined(); - expect(findSpan('execute_tool get_weather')?.data['gen_ai.tool.call.result']).toBeUndefined(); + const tool = findSpan('execute_tool get_weather'); + expect(tool?.data['gen_ai.tool.call.result']).toBeUndefined(); + expect(tool?.data['gen_ai.tool.call.arguments']).toBeDefined(); }); }); @@ -878,19 +586,11 @@ describe('createFlueInstrumentation', () => { // for a re-entry of the first: it gets no span, and its conversation id lands on the first's span. it('keeps concurrent agent runs separate', async () => { const runA = instrumentation.interceptor({ type: 'agent', operationId: 'op_a' }, { agentName: 'A' }, async () => { - instrumentation.observe( - { type: 'turn_start', turnId: 'turn_a', operationId: 'op_a', conversationId: 'conv_a' }, - {}, - ); + instrumentation.observe({ type: 'agent_start', operationId: 'op_a', conversationId: 'conv_a' }, {}); // B starts while A is still open. await instrumentation.interceptor({ type: 'agent', operationId: 'op_b' }, { agentName: 'B' }, async () => { - instrumentation.observe( - { type: 'turn_start', turnId: 'turn_b', operationId: 'op_b', conversationId: 'conv_b' }, - {}, - ); - instrumentation.observe(turn({ turnId: 'turn_b', operationId: 'op_b' }), {}); + instrumentation.observe({ type: 'agent_start', operationId: 'op_b', conversationId: 'conv_b' }, {}); }); - instrumentation.observe(turn({ turnId: 'turn_a', operationId: 'op_a' }), {}); }); await runA; @@ -904,9 +604,8 @@ describe('createFlueInstrumentation', () => { expect(agentB?.data['gen_ai.conversation.id']).toBe('conv_b'); }); - it('ignores a settled turn or tool it never opened a span for', async () => { + it('ignores a settled tool it never opened a span for', async () => { await withAgent(() => { - instrumentation.observe(turn({ turnId: 'never_started' }), {}); instrumentation.observe({ type: 'tool', toolCallId: 'never_started', toolName: 'x' }, {}); }); @@ -915,18 +614,15 @@ describe('createFlueInstrumentation', () => { it('ends spans still open at dispose', async () => { await withAgent(() => { - instrumentation.observe({ type: 'turn_start', turnId: 'turn_1', operationId: 'op_1' }, {}); instrumentation.observe( { type: 'tool_start', toolCallId: 'call_1', toolName: 'get_weather', operationId: 'op_1' }, {}, ); }); - // Never settled, so the span keeps the unqualified name it opened with. - expect(findSpan('chat')).toBeUndefined(); + expect(findSpan('execute_tool get_weather')).toBeUndefined(); instrumentation.dispose(); - expect(findSpan('chat')).toBeDefined(); expect(findSpan('execute_tool get_weather')).toBeDefined(); }); }); diff --git a/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts b/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts index ef8c7aeeff31..da6d3b1b1475 100644 --- a/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/pi-durable.test.ts @@ -10,42 +10,23 @@ import { setCurrentClient, spanToJSON, spanToStaticSpanJSON, + startInactiveSpan, 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_CALL_ARGUMENTS, GEN_AI_TOOL_CALL_RESULT, - GEN_AI_TOOL_DEFINITIONS, GEN_AI_TOOL_DESCRIPTION, GEN_AI_TOOL_NAME, - GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, - GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, - GEN_AI_USAGE_INPUT_TOKENS, - GEN_AI_USAGE_OUTPUT_TOKENS, - GEN_AI_USAGE_REASONING_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, GEN_AI_EXECUTE_TOOL, GEN_AI_INVOKE_AGENT } from '@sentry/conventions/op'; import { GEN_AI_TOOL_CALL_ID_ATTRIBUTE } from '../../../../src/ai/core/gen-ai-attributes'; import { endRunsOnClose, instrumentPiDurableHarnessOptions } from '../../../../src/ai/pi-durable'; +import { getPiAiChatSpanOptions } from '../../../../src/ai/pi-ai/chat'; import { MAX_TRACKED_PI_RUNS } from '../../../../src/ai/pi-durable/constants'; import type { PiRuns } from '../../../../src/ai/pi-durable/runs'; import { createRuns, startRun } from '../../../../src/ai/pi-durable/runs'; @@ -54,7 +35,6 @@ import type { PiCommitChange, PiHarnessOptions, PiLiveState, - PiModels, PiRegistryReader, PiSettlement, PiTask, @@ -63,16 +43,21 @@ import type { PiToolExecutionApi, PiToolExecutionResult, } from '../../../../src/ai/pi-durable/types'; -import type { PiEventStream } from '../../../../src/ai/pi-ai/types'; -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 { instrumentWorkersAiClient } from '../../../../src/ai/workers-ai'; import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; const LIVE_DOC = { definition: { kind: 'pi.live' } }; -const MODEL = { id: 'faux-model', provider: 'faux' }; -const ANSWER = { role: 'assistant', content: [{ type: 'text', text: 'Paris.' }], stopReason: 'stop' }; + +// What `piAiIntegration` does for each model request pi-durable sends through pi-ai. +function sendModelRequest(): void { + startInactiveSpan( + getPiAiChatSpanOptions({ + model: { id: 'faux-model', provider: 'faux' }, + context: { messages: [] }, + streaming: false, + recordInputs: false, + }), + ).end(); +} describe('instrumentPiDurableHarnessOptions', () => { let endedSpans: Span[]; @@ -106,14 +91,17 @@ describe('instrumentPiDurableHarnessOptions', () => { }); it('wraps the options once and leaves the caller objects untouched', () => { - const models = { completeSimple: async () => ({}) }; - const options: PiHarnessOptions = { models }; + const registry: PiRegistryReader = { + snapshot: () => ({ task: () => undefined, tasks: () => [] }), + subscribe: () => () => undefined, + }; + const options: PiHarnessOptions = { registry }; const instrumented = instrumentPiDurableHarnessOptions(options); expect(instrumented).not.toBe(options); - expect(instrumented.models).not.toBe(models); - expect(options.models).toBe(models); + expect(instrumented.registry).not.toBe(registry); + expect(options.registry).toBe(registry); expect(instrumentPiDurableHarnessOptions(instrumented)).toBe(instrumented); }); @@ -121,7 +109,7 @@ describe('instrumentPiDurableHarnessOptions', () => { // assignments on the caller's options must reach it. it('keeps the options the caller changes later live', () => { let env = 'first'; - const options = { models: {}, settings: { toolExecution: 'sequential' } } as PiHarnessOptions & { + const options = { settings: { toolExecution: 'sequential' } } as PiHarnessOptions & { settings: unknown; env?: string; }; @@ -135,322 +123,18 @@ describe('instrumentPiDurableHarnessOptions', () => { expect(instrumented.env).toBe('second'); }); - it('hands members other than model requests through, bound to the original', () => { - class Models { - private _provider = 'faux'; - public getProvider(): string { - return this._provider; - } - } - - const instrumented = instrumentPiDurableHarnessOptions({ models: new Models() as PiModels }); - - expect((instrumented.models as unknown as Models).getProvider()).toBe('faux'); - expect(endedSpans).toHaveLength(0); - }); - - it('traces a streamed model request and skips the provider SDK integrations', async () => { - const message = { - ...ANSWER, - model: 'faux-model', - usage: { input: 10, output: 2, totalTokens: 12 }, - }; - const models: PiModels = { streamSimple: () => ({ result: async () => message }) }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - await instrumented.models!.streamSimple!(MODEL, { - messages: [{ role: 'user', content: 'Capital of France?' }], - }).result(); - - expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(true); - const chat = spanToStaticSpanJSON(endedSpans[0]!); - expect(chat.description).toBe('chat faux-model'); - expect(chat.op).toBe(GEN_AI_CHAT); - expect(chat.data).toMatchObject({ - [GEN_AI_PROVIDER_NAME]: 'faux', - [GEN_AI_REQUEST_MODEL]: 'faux-model', - [GEN_AI_RESPONSE_STREAMING]: true, - [GEN_AI_RESPONSE_FINISH_REASONS]: '["stop"]', - [GEN_AI_USAGE_INPUT_TOKENS]: 10, - [GEN_AI_USAGE_OUTPUT_TOKENS]: 2, - }); - expect(chat.data[GEN_AI_INPUT_MESSAGES]).toBe( - '[{"role":"user","parts":[{"type":"text","content":"Capital of France?"}]}]', - ); - expect(chat.data[GEN_AI_OUTPUT_MESSAGES]).toBe( - '[{"role":"assistant","parts":[{"type":"text","content":"Paris."}],"finish_reason":"stop"}]', - ); - }); - - it.each([ - ['stream', true], - ['streamSimple', true], - ['streamDeferred', true], - ['complete', false], - ['completeSimple', false], - ['fetchDeferred', false], - ])('traces %s with the request options and the response identity', async (method, streaming) => { - const message = { ...ANSWER, responseModel: 'faux-model-2026', responseId: 'resp_1' }; - const models = { - [method]: streaming ? () => ({ result: async () => message }) : async () => message, - } as PiModels; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - const result = (instrumented.models as Record unknown>)[method]!( - { ...MODEL, baseUrl: 'https://api.example.com:8443/v1' }, - { messages: [{ role: 'user', content: 'Hi.' }] }, - { temperature: 0.2, maxTokens: 50, reasoning: 'low' }, - ); - await (streaming ? (result as PiEventStream).result() : result); - - const chat = spanToJSON(endedSpans[0]!); - expect(chat.name).toBe('chat faux-model'); - expect(chat.attributes[GEN_AI_RESPONSE_STREAMING]).toBe(streaming ? true : undefined); - expect(chat.attributes[GEN_AI_RESPONSE_MODEL]).toBe('faux-model-2026'); - expect(chat.attributes[GEN_AI_RESPONSE_ID]).toBe('resp_1'); - expect(chat.attributes[SERVER_ADDRESS]).toBe('api.example.com'); - expect(chat.attributes[SERVER_PORT]).toBe(8443); - // The deferred methods take a handle, not a request, so there is nothing to record. - const deferred = method.endsWith('Deferred'); - expect(chat.attributes[GEN_AI_REQUEST_TEMPERATURE]).toBe(deferred ? undefined : 0.2); - expect(chat.attributes[GEN_AI_REQUEST_MAX_TOKENS]).toBe(deferred ? undefined : 50); - expect(chat.attributes[GEN_AI_REQUEST_REASONING_LEVEL]).toBe(deferred ? undefined : 'low'); - expect(chat.attributes[GEN_AI_INPUT_MESSAGES]).toBe( - deferred ? undefined : '[{"role":"user","parts":[{"type":"text","content":"Hi."}]}]', - ); - }); - - it('sets the conversation id of the scope on the chat span itself', async () => { - const models: PiModels = { completeSimple: async () => ANSWER }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - await withScope(async scope => { - scope.setConversationId('harness:7'); - await instrumented.models!.completeSimple!(MODEL, { messages: [] }); - }); - - expect(spanToJSON(endedSpans[0]!).attributes[GEN_AI_CONVERSATION_ID]).toBe('harness:7'); - }); - - it('counts the cached tokens of a response into its input tokens', async () => { - const message = { - ...ANSWER, - usage: { input: 37, output: 2, cacheRead: 52, cacheWrite: 38, reasoning: 1, totalTokens: 129 }, - }; - const models: PiModels = { completeSimple: async () => message }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - await instrumented.models!.completeSimple!(MODEL, { messages: [] }); - - expect(spanToJSON(endedSpans[0]!).attributes).toMatchObject({ - [GEN_AI_USAGE_INPUT_TOKENS]: 127, - [GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]: 52, - [GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]: 38, - [GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]: 1, - [GEN_AI_USAGE_TOTAL_TOKENS]: 129, - }); - }); - - // The provider parks the request and pi-durable 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, output: 0, totalTokens: 0 }, - }; - const answer = { ...ANSWER, usage: { input: 8, output: 2, totalTokens: 10 } }; - const models: PiModels = { completeSimple: async () => parked, fetchDeferred: async () => answer }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - await instrumented.models!.completeSimple!(MODEL, { messages: [] }); - await instrumented.models!.fetchDeferred!(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 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 toolCall = { - role: 'assistant', - content: [{ type: 'toolCall', id: 'call_1', name: 'read', arguments: { path: 'a.txt' } }], - stopReason: 'toolUse', - }; - const models: PiModels = { completeSimple: async () => failed, complete: async () => toolCall }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - await instrumented.models!.completeSimple!(MODEL, { messages: [] }); - await instrumented.models!.complete!(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(); - 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 or throwing request as errored and keeps the error', async () => { - const models: PiModels = { - completeSimple: async () => { - throw new Error('boom'); - }, - stream: () => { - throw new Error('sync boom'); - }, - }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - await expect(instrumented.models!.completeSimple!(MODEL, { messages: [] })).rejects.toThrow('boom'); - expect(() => instrumented.models!.stream!(MODEL, { messages: [] })).toThrow('sync boom'); - - expect(endedSpans.map(span => spanToJSON(span).attributes[SENTRY_STATUS_MESSAGE])).toEqual([ - 'internal_error', - 'internal_error', - ]); - }); - - it('returns a result of another shape unchanged and ends its span', () => { - const stream = { events: [] }; - const models: PiModels = { - stream: () => stream as unknown as PiEventStream, - completeSimple: (() => 'parked') as unknown as PiModels['completeSimple'], - }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - expect(instrumented.models!.stream!(MODEL, { messages: [] })).toBe(stream); - expect(instrumented.models!.completeSimple!(MODEL, { messages: [] })).toBe('parked'); - - expect(endedSpans.map(span => spanToJSON(span).name)).toEqual(['chat faux-model', 'chat faux-model']); - }); - - it('reads the system prompt and the tools from the positional system messages', async () => { - const models: PiModels = { - streamSimple: () => ({ result: async () => ({ role: 'assistant', content: [], stopReason: 'stop' }) }), - }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - // pi-durable never sets `systemPrompt` or `tools`: it sends prompt sections and tool changes as - // system messages, and each request replays them. - await instrumented.models!.streamSimple!(MODEL, { - messages: [ - { - role: 'system', - content: '', - sections: { preamble: 'You help.' }, - toolsAdded: [{ name: 'read' }, { name: 'bash' }], - }, - { role: 'user', content: 'Hi.' }, - { role: 'system', content: '', sections: { mode: 'Plan only.' }, toolsRemoved: [{ name: 'bash' }] }, - ], - }).result(); - - const chat = spanToStaticSpanJSON(endedSpans[0]!); - expect(chat.data[GEN_AI_SYSTEM_INSTRUCTIONS]).toBe('You help.\n\nPlan only.'); - expect(chat.data[GEN_AI_TOOL_DEFINITIONS]).toBe('[{"name":"read"}]'); - expect(chat.data[GEN_AI_INPUT_MESSAGES]).toBe('[{"role":"user","parts":[{"type":"text","content":"Hi."}]}]'); - }); - - describe('provider skip', () => { - const models: PiModels = { completeSimple: async () => ANSWER }; - - it('skips the provider integrations on the first request, not when the options are wrapped', async () => { - const instrumented = instrumentPiDurableHarnessOptions({ models }); - expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(false); - - await instrumented.models!.completeSimple!(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 () => { - const instrumented = instrumentPiDurableHarnessOptions({ models }); - await instrumented.models!.completeSimple!(MODEL, { messages: [] }); - _INTERNAL_clearAiProviderSkips(); - expect(_INTERNAL_shouldSkipAiProviderWrapping(OPENAI_INTEGRATION_NAME)).toBe(false); - - await instrumented.models!.completeSimple!(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]); - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - await instrumented.models!.completeSimple!(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 () => [], - }); - const models: PiModels = { - completeSimple: async () => { - await ai.run('@cf/moonshotai/kimi-k2.7-code', { messages: [] }, { returnRawResponse: true }); - return ANSWER; - }, - }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - - await instrumented.models!.completeSimple!(MODEL, { messages: [] }); - - expect(endedSpans.map(span => spanToStaticSpanJSON(span).description)).toEqual(['chat faux-model']); - }); - }); - describe('content recording', () => { const tool: PiTool = { name: 'get_weather', description: 'Get the weather.', execute: async () => ({ content: [{ type: 'text', text: 'sunny' }] }), }; - 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?' }, - ], - }; - it('records nothing when inputs and outputs are off', async () => { const options = { recordInputs: false, recordOutputs: false }; - const instrumented = instrumentPiDurableHarnessOptions( - { models: { completeSimple: async () => message } as PiModels }, - options, - ); - await instrumented.models!.completeSimple!(MODEL, context); await instrumentTool(tool, createRuns(), options).execute({ city: 'Berlin' }, { callId: 'call_1' }, undefined); - const [chat, execute] = endedSpans.map(span => spanToJSON(span).attributes); - expect(Object.keys(chat!)).not.toContain(GEN_AI_INPUT_MESSAGES); - expect(Object.keys(chat!)).not.toContain(GEN_AI_OUTPUT_MESSAGES); - expect(Object.keys(chat!)).not.toContain(GEN_AI_SYSTEM_INSTRUCTIONS); - expect(Object.keys(chat!)).not.toContain(GEN_AI_TOOL_DEFINITIONS); + const [execute] = endedSpans.map(span => spanToJSON(span).attributes); expect(Object.keys(execute!)).not.toContain(GEN_AI_TOOL_CALL_ARGUMENTS); expect(Object.keys(execute!)).not.toContain(GEN_AI_TOOL_CALL_RESULT); expect(execute![GEN_AI_TOOL_DESCRIPTION]).toBe('Get the weather.'); @@ -467,17 +151,10 @@ describe('instrumentPiDurableHarnessOptions', () => { setCurrentClient(strict); strict.init(); strict.on('spanEnd', span => endedSpans.push(span)); - const models: PiModels = { completeSimple: async () => message }; - const instrumented = instrumentPiDurableHarnessOptions({ models }); - await instrumented.models!.completeSimple!(MODEL, context); await instrumentTool(tool, createRuns(), {}).execute({ city: 'Berlin' }, { callId: 'call_1' }, undefined); const recorded = endedSpans.flatMap(span => Object.keys(spanToJSON(span).attributes)); - expect(recorded).not.toContain(GEN_AI_INPUT_MESSAGES); - expect(recorded).not.toContain(GEN_AI_OUTPUT_MESSAGES); - expect(recorded).not.toContain(GEN_AI_SYSTEM_INSTRUCTIONS); - expect(recorded).not.toContain(GEN_AI_TOOL_DEFINITIONS); expect(recorded).not.toContain(GEN_AI_TOOL_CALL_ARGUMENTS); expect(recorded).not.toContain(GEN_AI_TOOL_CALL_RESULT); }); @@ -958,14 +635,13 @@ describe('instrumentPiDurableHarnessOptions', () => { name: 'pi.generation', phases: { request: async (taskRecord, runtime, context) => { - await harness.models!.completeSimple!(MODEL, { messages: [] }); + sendModelRequest(); await settle.definition.phases.request!(taskRecord, runtime, context); }, }, }, }; const harness: PiHarnessOptions = instrumentPiDurableHarnessOptions({ - models: { completeSimple: async () => ANSWER }, registry: { snapshot: () => ({ task: () => task, tasks: () => [task] }), subscribe: () => () => undefined, @@ -1056,13 +732,12 @@ describe('instrumentPiDurableHarnessOptions', () => { name: 'app.title', phases: { request: async () => { - await harness.models!.completeSimple!(MODEL, { messages: [] }); + sendModelRequest(); }, }, }, }; const harness: PiHarnessOptions = instrumentPiDurableHarnessOptions({ - models: { completeSimple: async () => ANSWER }, registry: { snapshot: () => ({ task: () => task, tasks: () => [task] }), subscribe: () => () => undefined }, }); const runtime: PiTaskRuntime = { diff --git a/packages/server-utils/test/integrations/pi-ai.test.ts b/packages/server-utils/test/integrations/pi-ai.test.ts index a062e191a6e7..5fda26d4e8c4 100644 --- a/packages/server-utils/test/integrations/pi-ai.test.ts +++ b/packages/server-utils/test/integrations/pi-ai.test.ts @@ -43,7 +43,6 @@ import { ANTHROPIC_AI_INTEGRATION_NAME } from '../../src/ai/anthropic-ai/constan 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'; @@ -355,7 +354,7 @@ describe('piAiIntegration', () => { it('records nothing when inputs and outputs are off', () => { const span = startInactiveSpan( - getPiAiChatSpanOptions({ model: MODEL, context, streaming: false, origin: PI_AI_ORIGIN, recordInputs: false }), + getPiAiChatSpanOptions({ model: MODEL, context, streaming: false, recordInputs: false }), ); setPiAiResponseAttributes(span, message, false); span.end(); diff --git a/packages/server-utils/test/orchestrion/moduleInjectedTransform.test.ts b/packages/server-utils/test/orchestrion/moduleInjectedTransform.test.ts index 31565dc9e50a..5dac489035d1 100644 --- a/packages/server-utils/test/orchestrion/moduleInjectedTransform.test.ts +++ b/packages/server-utils/test/orchestrion/moduleInjectedTransform.test.ts @@ -44,8 +44,7 @@ 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. - // pi-ai would duplicate Flue's and pi-durable's own chat spans until they rely on it. - const optIn = new Set(['node-cron', '@earendil-works/pi-ai']); + const optIn = new Set(['node-cron']); const uncovered = [...configured].filter( name => !defined.has(name) && !optIn.has(name) && !name.startsWith('@nestjs/') && !name.startsWith('@remix-run/'), ); From c5e394ff02ba5d50b9abc3ab03e53649a86435c9 Mon Sep 17 00:00:00 2001 From: Francesco Novy Date: Fri, 30 Oct 2026 12:00:06 +0100 Subject: [PATCH 2/2] feat(server-utils): Warn when Flue runs without pi-ai instrumentation Flue's model calls are now traced by the pi-ai integration, which needs the module hook or bundler plugin, while `instrument()` does not. Without them a Flue app would silently lose its `gen_ai.chat` spans; it now gets a warning on its first model call. Co-Authored-By: Claude Opus 5.5 --- packages/server-utils/src/ai/flue/index.ts | 33 ++++++++++++++ .../ai/lib/tracing/flue-pi-ai-warning.test.ts | 43 +++++++++++++++++++ 2 files changed, 76 insertions(+) create mode 100644 packages/server-utils/test/ai/lib/tracing/flue-pi-ai-warning.test.ts diff --git a/packages/server-utils/src/ai/flue/index.ts b/packages/server-utils/src/ai/flue/index.ts index 4edda9f5ccac..6f22e1749d7b 100644 --- a/packages/server-utils/src/ai/flue/index.ts +++ b/packages/server-utils/src/ai/flue/index.ts @@ -1,5 +1,6 @@ import type { Scope, Span } from '@sentry/core'; import { + consoleSandbox, continueTrace, getActiveSpan, getDefaultIsolationScope, @@ -19,6 +20,8 @@ import { } from '@sentry/conventions/attributes'; import type { GenAiOptions } from '../core/utils'; import { getGenAiSpanOp, resolveAIRecordingOptions } from '../core/utils'; +import { getOrchestrionInjectedModules } from '../../orchestrion/detect'; +import { piAiModuleNames } from '../../orchestrion/config/pi-ai'; import { withCleanScopes } from '../../utils/withCleanScopes'; import { FLUE_INSTRUMENTATION_KEY, FLUE_OPERATION, FLUE_ORIGIN, MAX_TRACKED_FLUE_SPANS } from './constants'; import type { SpanTracker } from './utils'; @@ -32,6 +35,35 @@ interface Admission { span: Span | undefined; } +let checkedPiAiInstrumentation = false; + +/** + * `instrument()` works without Sentry's module hook or bundler plugin, but the pi-ai integration that + * spans Flue's model calls doesn't. Checked on the first model call, once pi-ai has loaded. + */ +function warnIfPiAiNotInstrumented(): void { + if (checkedPiAiInstrumentation) { + return; + } + checkedPiAiInstrumentation = true; + + const injected = getOrchestrionInjectedModules(); + if (piAiModuleNames.some(name => injected.includes(name))) { + return; + } + + consoleSandbox(() => { + // eslint-disable-next-line no-console + console.warn( + '[Sentry] Flue model calls are not being traced: `@earendil-works/pi-ai` was not instrumented, so ' + + 'Flue agents record agent and tool spans but no `gen_ai.chat` spans. Model calls are traced by the ' + + "pi-ai integration, which needs Sentry's module hook or bundler plugin. On Node, load Sentry with " + + '`--import` (for example `node --import ./instrument.mjs app.mjs`). On Cloudflare, add ' + + '`sentryCloudflareVitePlugin()`. On Bun, build with `@sentry/bun/plugin`.', + ); + }); +} + /** * Flue's Node coordinator runs every submission from one long-lived claim loop, started by the * first admission and keeping that call's async context for the life of the process. Left alone, @@ -99,6 +131,7 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru // The `gen_ai.chat` span of the model call comes from the pi-ai integration, which reads the // conversation from the scope. if (operation?.type === FLUE_OPERATION.MODEL) { + warnIfPiAiNotInstrumented(); const { conversationId } = ctx; return conversationId ? withScope(scope => { diff --git a/packages/server-utils/test/ai/lib/tracing/flue-pi-ai-warning.test.ts b/packages/server-utils/test/ai/lib/tracing/flue-pi-ai-warning.test.ts new file mode 100644 index 000000000000..5a4cbf2012a8 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/flue-pi-ai-warning.test.ts @@ -0,0 +1,43 @@ +import { GLOBAL_OBJ } from '@sentry/core'; +import { afterEach, describe, expect, it, vi } from 'vitest'; + +const MODEL_OP = { type: 'model' }; + +/** A fresh module per test: the check runs once per process. */ +async function runModelCalls(count: number): Promise { + vi.resetModules(); + const { createFlueInstrumentation } = await import('../../../../src/ai/flue'); + const instrumentation = createFlueInstrumentation(); + for (let i = 0; i < count; i++) { + await instrumentation.interceptor(MODEL_OP, {}, async () => undefined); + } +} + +describe('Flue warning about uninstrumented pi-ai', () => { + afterEach(() => { + delete GLOBAL_OBJ.__SENTRY_ORCHESTRION__; + vi.restoreAllMocks(); + }); + + it('warns once when pi-ai was not instrumented', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + GLOBAL_OBJ.__SENTRY_ORCHESTRION__ = { runtime: ['@flue/runtime'] }; + + await runModelCalls(2); + + expect(warn).toHaveBeenCalledOnce(); + expect(warn).toHaveBeenCalledWith(expect.stringContaining('`@earendil-works/pi-ai` was not instrumented')); + }); + + it.each([ + ['the runtime hook', { runtime: ['@earendil-works/pi-ai'] }], + ['a bundler plugin', { bundler: new Set(['@earendil-works/pi-ai']) }], + ])('stays quiet when %s instrumented pi-ai', async (_label, injected) => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + GLOBAL_OBJ.__SENTRY_ORCHESTRION__ = injected; + + await runModelCalls(1); + + expect(warn).not.toHaveBeenCalled(); + }); +});