|
| 1 | +import type { Span } from '@sentry/core'; |
| 2 | +import { |
| 3 | + _INTERNAL_shouldSkipAiProviderWrapping, |
| 4 | + _INTERNAL_skipAiProviderWrapping, |
| 5 | + continueTrace, |
| 6 | + getActiveSpan, |
| 7 | + LRUMap, |
| 8 | + SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, |
| 9 | + startSpan, |
| 10 | + withActiveSpan, |
| 11 | +} from '@sentry/core'; |
| 12 | +import { GEN_AI_AGENT_NAME, GEN_AI_CONVERSATION_ID, GEN_AI_OPERATION_NAME } from '@sentry/conventions/attributes'; |
| 13 | +import { ANTHROPIC_AI_INTEGRATION_NAME } from '../anthropic-ai/constants'; |
| 14 | +import type { GenAiOptions } from '../core/utils'; |
| 15 | +import { getGenAiSpanOp, resolveAIRecordingOptions } from '../core/utils'; |
| 16 | +import { GOOGLE_GENAI_INTEGRATION_NAME } from '../google-genai/constants'; |
| 17 | +import { OPENAI_INTEGRATION_NAME } from '../openai/constants'; |
| 18 | +import { FLUE_INSTRUMENTATION_KEY, FLUE_OPERATION, FLUE_ORIGIN, MAX_TRACKED_FLUE_SPANS } from './constants'; |
| 19 | +import type { SpanTracker } from './utils'; |
| 20 | +import { |
| 21 | + endToolSpan, |
| 22 | + endTurnSpan, |
| 23 | + recordRequestContent, |
| 24 | + sentryTraceFromTraceparent, |
| 25 | + startToolSpan, |
| 26 | + startTurnSpan, |
| 27 | +} from './utils'; |
| 28 | +import type { FlueInstrumentation } from './types'; |
| 29 | + |
| 30 | +export type FlueOptions = GenAiOptions; |
| 31 | + |
| 32 | +const SKIPPED_PROVIDERS = [OPENAI_INTEGRATION_NAME, ANTHROPIC_AI_INTEGRATION_NAME, GOOGLE_GENAI_INTEGRATION_NAME]; |
| 33 | + |
| 34 | +/** |
| 35 | + * Build the object to hand to `instrument()` from `@flue/runtime`. |
| 36 | + * |
| 37 | + * The two callbacks own different halves of the result: |
| 38 | + * |
| 39 | + * - `interceptor` wraps agent execution, so the agent span is *active* for its duration and every |
| 40 | + * span opened underneath parents correctly. |
| 41 | + * - `observe` opens and closes the turn span, because Flue's `turn_start`/`turn` events are the |
| 42 | + * only one-to-one signal for a model call and `turn` is what carries usage and cost. |
| 43 | + * |
| 44 | + * Message content, tool arguments and tool results are gated on `recordInputs`/`recordOutputs`, |
| 45 | + * which fall back to the current client's `dataCollection.genAI` settings and are read per event. |
| 46 | + */ |
| 47 | +export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstrumentation { |
| 48 | + // Flue drives the providers through `@earendil-works/pi-ai`, which bundles the `openai`, |
| 49 | + // `@anthropic-ai/sdk` and `@google/genai` clients. Left alone they instrument the same call this |
| 50 | + // reports as a turn, emitting a second `gen_ai.chat` beside ours. |
| 51 | + // |
| 52 | + // Applied on first use rather than here, for two reasons. Constructing the object proves nothing |
| 53 | + // — if `instrument()` rejects it, suppressing the provider integrations would leave the app with |
| 54 | + // no `gen_ai.chat` spans at all. And the registry is reset per client (`_setupIntegrations` |
| 55 | + // clears it, and Cloudflare calls `init()` per request), so a one-shot call at module scope is |
| 56 | + // wiped by the next `init()` and every later request double-reports. |
| 57 | + const skipProviders = (): void => { |
| 58 | + if (!SKIPPED_PROVIDERS.every(provider => _INTERNAL_shouldSkipAiProviderWrapping(provider))) { |
| 59 | + _INTERNAL_skipAiProviderWrapping(SKIPPED_PROVIDERS); |
| 60 | + } |
| 61 | + }; |
| 62 | + |
| 63 | + // Keyed by the agent operation's own id, which is what the observations carry. That keeps |
| 64 | + // concurrent runs apart and gives a delegated subagent its own span: Flue nests a second `agent` |
| 65 | + // operation inside the parent's for `task` delegation, and the nesting is not bounded at two. |
| 66 | + // A plain map: the entry is removed in a `finally`, so it is bounded by concurrent agent runs. |
| 67 | + const agentSpans = new Map<string, Span>(); |
| 68 | + // Capped, unlike the above: these are keyed off ids that only a matching end observation removes. |
| 69 | + const turnSpans: SpanTracker = new LRUMap(MAX_TRACKED_FLUE_SPANS); |
| 70 | + const toolSpans: SpanTracker = new LRUMap(MAX_TRACKED_FLUE_SPANS); |
| 71 | + |
| 72 | + return { |
| 73 | + key: FLUE_INSTRUMENTATION_KEY, |
| 74 | + |
| 75 | + interceptor: async (operation, ctx, next) => { |
| 76 | + skipProviders(); |
| 77 | + |
| 78 | + // `observe` has already opened the span for this unit of work; make it active for the |
| 79 | + // duration so whatever the tool or model call does lands inside it rather than beside it. |
| 80 | + if (operation?.type === FLUE_OPERATION.TOOL) { |
| 81 | + const toolSpan = operation.toolCallId ? toolSpans.get(operation.toolCallId) : undefined; |
| 82 | + return toolSpan ? withActiveSpan(toolSpan, next) : next(); |
| 83 | + } |
| 84 | + if (operation?.type === FLUE_OPERATION.MODEL) { |
| 85 | + const turnSpan = operation.turnId ? turnSpans.get(operation.turnId) : undefined; |
| 86 | + return turnSpan ? withActiveSpan(turnSpan, next) : next(); |
| 87 | + } |
| 88 | + |
| 89 | + if (operation?.type !== FLUE_OPERATION.AGENT) { |
| 90 | + return next(); |
| 91 | + } |
| 92 | + |
| 93 | + // A submission opens with a wrapper operation whose id *is* the submission id; the run it |
| 94 | + // wraps gets its own operation id, and that is the one the observations reference. Spanning |
| 95 | + // the wrapper too would double-count every agent invocation. |
| 96 | + const operationId = operation.operationId; |
| 97 | + if (!operationId || operationId === ctx.submissionId || agentSpans.has(operationId)) { |
| 98 | + return next(); |
| 99 | + } |
| 100 | + |
| 101 | + const openAgentSpan = () => |
| 102 | + startSpan( |
| 103 | + { |
| 104 | + name: `invoke_agent ${ctx.agentName ?? 'agent'}`, |
| 105 | + op: getGenAiSpanOp('invoke_agent'), |
| 106 | + attributes: { |
| 107 | + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: FLUE_ORIGIN, |
| 108 | + [GEN_AI_OPERATION_NAME]: 'invoke_agent', |
| 109 | + ...(ctx.agentName ? { [GEN_AI_AGENT_NAME]: ctx.agentName } : {}), |
| 110 | + ...(ctx.conversationId ? { [GEN_AI_CONVERSATION_ID]: ctx.conversationId } : {}), |
| 111 | + }, |
| 112 | + }, |
| 113 | + async (span: Span) => { |
| 114 | + agentSpans.set(operationId, span); |
| 115 | + try { |
| 116 | + return await next(); |
| 117 | + } finally { |
| 118 | + agentSpans.delete(operationId); |
| 119 | + } |
| 120 | + }, |
| 121 | + ); |
| 122 | + |
| 123 | + // A durable or dispatched submission is run later by the coordinator, possibly in another |
| 124 | + // isolate or process, where nothing links it back to the request that enqueued it. Flue |
| 125 | + // replays that request's `traceparent` here, so continue from it — but only when no trace is |
| 126 | + // already active, so an in-process dispatch keeps the trace it is genuinely part of. |
| 127 | + const sentryTrace = ctx.traceCarrier?.traceparent |
| 128 | + ? sentryTraceFromTraceparent(ctx.traceCarrier.traceparent) |
| 129 | + : undefined; |
| 130 | + |
| 131 | + return sentryTrace && !getActiveSpan() |
| 132 | + ? continueTrace({ sentryTrace, baggage: undefined }, openAgentSpan) |
| 133 | + : openAgentSpan(); |
| 134 | + }, |
| 135 | + |
| 136 | + observe: observation => { |
| 137 | + skipProviders(); |
| 138 | + |
| 139 | + // Resolved per observation, not once at construction: `resolveAIRecordingOptions` reads the |
| 140 | + // current client's `dataCollection.genAI`, and Cloudflare replaces the client per request, so |
| 141 | + // values captured at isolate load would be the wrong ones for every later request. |
| 142 | + const { recordInputs, recordOutputs } = resolveAIRecordingOptions(options); |
| 143 | + |
| 144 | + // Observations are emitted synchronously from inside the agent operation, so the active span |
| 145 | + // here is the agent span the interceptor opened — turn and tool spans parent off it without |
| 146 | + // any bookkeeping. |
| 147 | + // |
| 148 | + // The agent operation that gets the span is the one the observations reference, and it knows |
| 149 | + // neither the agent's name nor the conversation: the submission wrapper carries the name, the |
| 150 | + // re-entry carries the conversation, and neither opens a span. Both arrive here instead. |
| 151 | + const agentSpan = observation.operationId ? agentSpans.get(observation.operationId) : undefined; |
| 152 | + if (agentSpan) { |
| 153 | + if (observation.conversationId) { |
| 154 | + agentSpan.setAttribute(GEN_AI_CONVERSATION_ID, observation.conversationId); |
| 155 | + } |
| 156 | + if (observation.agentName) { |
| 157 | + agentSpan.setAttribute(GEN_AI_AGENT_NAME, observation.agentName); |
| 158 | + agentSpan.updateName(`invoke_agent ${observation.agentName}`); |
| 159 | + } |
| 160 | + } |
| 161 | + |
| 162 | + switch (observation.type) { |
| 163 | + case 'turn_start': |
| 164 | + startTurnSpan(observation, turnSpans); |
| 165 | + return; |
| 166 | + case 'turn_request': |
| 167 | + if (recordInputs) { |
| 168 | + recordRequestContent(observation, turnSpans); |
| 169 | + } |
| 170 | + return; |
| 171 | + case 'turn': |
| 172 | + endTurnSpan(observation, turnSpans, recordOutputs); |
| 173 | + return; |
| 174 | + case 'tool_start': |
| 175 | + startToolSpan(observation, toolSpans, recordInputs); |
| 176 | + return; |
| 177 | + case 'tool': |
| 178 | + endToolSpan(observation, toolSpans, recordOutputs); |
| 179 | + return; |
| 180 | + default: |
| 181 | + return; |
| 182 | + } |
| 183 | + }, |
| 184 | + |
| 185 | + // Only the turn and tool spans are ended here. They come from `startInactiveSpan`, so nothing |
| 186 | + // else will. The agent spans come from `startSpan`, which ends them when its callback settles — |
| 187 | + // ending them here would stamp an early end time on work that is still running, and the later |
| 188 | + // `end()` would be ignored. |
| 189 | + dispose: () => { |
| 190 | + for (const span of [...turnSpans.values(), ...toolSpans.values()]) { |
| 191 | + span.end(); |
| 192 | + } |
| 193 | + agentSpans.clear(); |
| 194 | + turnSpans.clear(); |
| 195 | + toolSpans.clear(); |
| 196 | + }, |
| 197 | + }; |
| 198 | +} |
0 commit comments