|
| 1 | +import * as diagnosticsChannel from 'node:diagnostics_channel'; |
| 2 | +import type { GoogleGenAIOptions, GoogleGenAIResponse, IntegrationFn, Span } from '@sentry/core'; |
| 3 | +import { |
| 4 | + _INTERNAL_shouldSkipAiProviderWrapping, |
| 5 | + addGoogleGenAIRequestAttributes, |
| 6 | + addGoogleGenAIResponseAttributes, |
| 7 | + debug, |
| 8 | + defineIntegration, |
| 9 | + extractGoogleGenAIRequestAttributes, |
| 10 | + GEN_AI_REQUEST_MODEL_ATTRIBUTE, |
| 11 | + getActiveSpan, |
| 12 | + instrumentGoogleGenAIStream, |
| 13 | + resolveAIRecordingOptions, |
| 14 | + SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, |
| 15 | + shouldEnableTruncation, |
| 16 | + spanToJSON, |
| 17 | + startInactiveSpan, |
| 18 | + waitForTracingChannelBinding, |
| 19 | +} from '@sentry/core'; |
| 20 | +import { DEBUG_BUILD } from '../../debug-build'; |
| 21 | +import { CHANNELS } from '../../orchestrion/channels'; |
| 22 | +import { bindTracingChannelToSpan } from '../../tracing-channel'; |
| 23 | + |
| 24 | +// Same name as the OTel integration by design: when enabled, the OTel 'Google_GenAI' |
| 25 | +// integration is dropped from the default set (see the Node opt-in loader). |
| 26 | +const INTEGRATION_NAME = 'Google_GenAI' as const; |
| 27 | + |
| 28 | +// Distinct from the proxy's `auto.ai.google_genai` so spans from the orchestrion path |
| 29 | +// are attributable separately from the OTel/proxy one. |
| 30 | +const ORIGIN = 'auto.ai.orchestrion.google_genai'; |
| 31 | + |
| 32 | +const CHAT_OP = 'gen_ai.chat'; |
| 33 | + |
| 34 | +// `stream` determines how the span is ended. `chat` and `generate-content` both carry streaming and |
| 35 | +// non-streaming calls on one channel, so their mode is resolved at runtime from the result shape. |
| 36 | +const INSTRUMENTED_CHANNELS = [ |
| 37 | + { channel: CHANNELS.GOOGLE_GENAI_GENERATE_CONTENT, operation: 'generate_content', stream: 'async-iterable' }, |
| 38 | + { channel: CHANNELS.GOOGLE_GENAI_EMBED_CONTENT, operation: 'embeddings', stream: 'none' }, |
| 39 | + { channel: CHANNELS.GOOGLE_GENAI_CHAT, operation: 'chat', stream: 'async-iterable' }, |
| 40 | +] as const; |
| 41 | + |
| 42 | +type StreamMode = (typeof INSTRUMENTED_CHANNELS)[number]['stream']; |
| 43 | + |
| 44 | +interface GoogleGenAIChannelContext { |
| 45 | + arguments: unknown[]; |
| 46 | + // The transform stashes the call's `this` here, which chat methods need since the model lives on |
| 47 | + // the `Chat` instance (`this.model`/`this.modelVersion`), not in the call arguments. |
| 48 | + self?: unknown; |
| 49 | + result?: unknown; |
| 50 | +} |
| 51 | + |
| 52 | +let subscribed = false; |
| 53 | + |
| 54 | +const _googleGenAIChannelIntegration = ((options: GoogleGenAIOptions = {}) => { |
| 55 | + return { |
| 56 | + name: INTEGRATION_NAME, |
| 57 | + setupOnce() { |
| 58 | + // tracingChannel is unavailable before Node 18.19 and prevent double-subscribe |
| 59 | + if (!diagnosticsChannel.tracingChannel || subscribed) { |
| 60 | + return; |
| 61 | + } |
| 62 | + subscribed = true; |
| 63 | + |
| 64 | + // `bindTracingChannelToSpan` needs the async-context binding that `initOpenTelemetry()` registers |
| 65 | + // after `setupOnce` runs, so wait for it before subscribing. |
| 66 | + waitForTracingChannelBinding(() => { |
| 67 | + for (const { channel, operation, stream } of INSTRUMENTED_CHANNELS) { |
| 68 | + DEBUG_BUILD && debug.log(`[orchestrion:google-genai] subscribing to channel "${channel}"`); |
| 69 | + bindTracingChannelToSpan( |
| 70 | + diagnosticsChannel.tracingChannel<GoogleGenAIChannelContext>(channel), |
| 71 | + data => createGenAiSpan(data, operation, options), |
| 72 | + { |
| 73 | + beforeSpanEnd: (span, data) => { |
| 74 | + // Embeddings responses carry no content attributes, mirroring the OTel integration. |
| 75 | + if (operation !== 'embeddings') { |
| 76 | + addGoogleGenAIResponseAttributes( |
| 77 | + span, |
| 78 | + data.result as GoogleGenAIResponse, |
| 79 | + resolveAIRecordingOptions(options).recordOutputs, |
| 80 | + ); |
| 81 | + } |
| 82 | + }, |
| 83 | + deferSpanEnd: ({ span, data }) => wrapStreamResult(span, data, stream, options), |
| 84 | + }, |
| 85 | + ); |
| 86 | + } |
| 87 | + }); |
| 88 | + }, |
| 89 | + }; |
| 90 | +}) satisfies IntegrationFn; |
| 91 | + |
| 92 | +/** |
| 93 | + * Build the span for an instrumented call. |
| 94 | + * Returning `undefined` opts the payload out so no span is opened. |
| 95 | + */ |
| 96 | +function createGenAiSpan( |
| 97 | + data: GoogleGenAIChannelContext, |
| 98 | + operation: string, |
| 99 | + options: GoogleGenAIOptions, |
| 100 | +): Span | undefined { |
| 101 | + const args = data.arguments ?? []; |
| 102 | + |
| 103 | + // When LangChain (or another provider) is driving the SDK, it records the spans itself and marks this |
| 104 | + // provider as skipped — mirror the OTel integration and don't double-instrument. |
| 105 | + if (_INTERNAL_shouldSkipAiProviderWrapping(INTEGRATION_NAME)) { |
| 106 | + return undefined; |
| 107 | + } |
| 108 | + |
| 109 | + // `chat.sendMessage()`/`sendMessageStream()` internally call `Models.generateContent(Stream)`, which |
| 110 | + // publishes the `generate-content` channel while the chat span is active. Skip that nested event so a |
| 111 | + // chat call yields a single `gen_ai.chat` span instead of a chat span wrapping a generate_content one. |
| 112 | + if (operation !== 'chat') { |
| 113 | + const activeSpan = getActiveSpan(); |
| 114 | + if (activeSpan) { |
| 115 | + const { op, origin } = spanToJSON(activeSpan); |
| 116 | + if (origin === ORIGIN && op === CHAT_OP) { |
| 117 | + return undefined; |
| 118 | + } |
| 119 | + } |
| 120 | + } |
| 121 | + |
| 122 | + const params = typeof args[0] === 'object' && args[0] !== null ? (args[0] as Record<string, unknown>) : undefined; |
| 123 | + |
| 124 | + const { recordInputs } = resolveAIRecordingOptions(options); |
| 125 | + const enableTruncation = shouldEnableTruncation(options.enableTruncation); |
| 126 | + |
| 127 | + const attributes = extractGoogleGenAIRequestAttributes(operation, params, data.self); |
| 128 | + const model = (attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] as string) || 'unknown'; |
| 129 | + attributes[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN] = ORIGIN; |
| 130 | + |
| 131 | + const span = startInactiveSpan({ |
| 132 | + name: `${operation} ${model}`, |
| 133 | + op: `gen_ai.${operation}`, |
| 134 | + attributes, |
| 135 | + }); |
| 136 | + |
| 137 | + if (recordInputs && params) { |
| 138 | + addGoogleGenAIRequestAttributes(span, params, operation === 'embeddings', enableTruncation); |
| 139 | + } |
| 140 | + |
| 141 | + return span; |
| 142 | +} |
| 143 | + |
| 144 | +type AsyncIterableStream = { [Symbol.asyncIterator]: () => AsyncIterator<unknown> }; |
| 145 | + |
| 146 | +function isAsyncIterable(value: unknown): value is AsyncIterableStream { |
| 147 | + return !!value && typeof (value as AsyncIterableStream)[Symbol.asyncIterator] === 'function'; |
| 148 | +} |
| 149 | + |
| 150 | +/** |
| 151 | + * Hand span-ending ownership to a streamed result: returns `true` to skip the normal `beforeSpanEnd`, |
| 152 | + * `false` for non-streaming results (which end via `beforeSpanEnd`). |
| 153 | + * |
| 154 | + * `generateContentStream`/`sendMessageStream` resolve to an async iterable, while `generateContent`/ |
| 155 | + * `sendMessage` resolve to a plain response — so the runtime `isAsyncIterable` check picks the right |
| 156 | + * path even though both share a channel. For streams we patch `result[Symbol.asyncIterator]` in place |
| 157 | + * so `instrumentGoogleGenAIStream` ends the span when iteration finishes. |
| 158 | + */ |
| 159 | +function wrapStreamResult( |
| 160 | + span: Span, |
| 161 | + data: GoogleGenAIChannelContext, |
| 162 | + stream: StreamMode, |
| 163 | + options: GoogleGenAIOptions, |
| 164 | +): boolean { |
| 165 | + const { recordOutputs } = resolveAIRecordingOptions(options); |
| 166 | + const result = data.result; |
| 167 | + |
| 168 | + if (stream === 'async-iterable' && isAsyncIterable(result)) { |
| 169 | + const iterate = result[Symbol.asyncIterator].bind(result); |
| 170 | + const instrumented = instrumentGoogleGenAIStream({ [Symbol.asyncIterator]: iterate }, span, !!recordOutputs); |
| 171 | + result[Symbol.asyncIterator] = () => instrumented; |
| 172 | + return true; |
| 173 | + } |
| 174 | + |
| 175 | + return false; |
| 176 | +} |
| 177 | + |
| 178 | +/** |
| 179 | + * EXPERIMENTAL — orchestrion-driven Google GenAI integration. Subscribes to the |
| 180 | + * `orchestrion:@google/genai:*` diagnostics_channels injected into the SDK's `Models` |
| 181 | + * (`generateContent`/`generateContentStream`/`embedContent`) and `Chat` |
| 182 | + * (`sendMessage`/`sendMessageStream`) methods, so it requires the orchestrion runtime hook or |
| 183 | + * bundler plugin. |
| 184 | + */ |
| 185 | +export const googleGenAIChannelIntegration = defineIntegration(_googleGenAIChannelIntegration); |
0 commit comments