feat(server-utils): Migrate Google GenAI integration to orchestrion

dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts

@@ -22,6 +22,9 @@ import {
   GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
 } from '../../../../../packages/core/src/tracing/ai/gen-ai-attributes';
 import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../utils/runner';
+import { isOrchestrionEnabled } from '../../../utils';
+
+const EXPECTED_ORIGIN = isOrchestrionEnabled() ? 'auto.ai.orchestrion.google_genai' : 'auto.ai.google_genai';
 
 describe('Google GenAI integration', () => {
   afterAll(() => {
@@ -47,7 +50,7 @@ describe('Google GenAI integration', () => {
             expect(chatSpan!.status).toBe('ok');
             expect(chatSpan!.attributes['sentry.op'].value).toBe('gen_ai.chat');
             expect(chatSpan!.attributes[GEN_AI_OPERATION_NAME_ATTRIBUTE].value).toBe('chat');
-            expect(chatSpan!.attributes['sentry.origin'].value).toBe('auto.ai.google_genai');
+            expect(chatSpan!.attributes['sentry.origin'].value).toBe(EXPECTED_ORIGIN);
             expect(chatSpan!.attributes[GEN_AI_SYSTEM_ATTRIBUTE].value).toBe('google_genai');
             expect(chatSpan!.attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE].value).toBe('gemini-1.5-pro');
             expect(chatSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE].value).toBe(8);
@@ -439,7 +442,7 @@ describe('Google GenAI integration', () => {
             for (const span of successfulSpans) {
               expect(span.attributes['sentry.op'].value).toBe('gen_ai.embeddings');
               expect(span.attributes[GEN_AI_OPERATION_NAME_ATTRIBUTE].value).toBe('embeddings');
-              expect(span.attributes['sentry.origin'].value).toBe('auto.ai.google_genai');
+              expect(span.attributes['sentry.origin'].value).toBe(EXPECTED_ORIGIN);
               expect(span.attributes[GEN_AI_SYSTEM_ATTRIBUTE].value).toBe('google_genai');
               expect(span.attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE].value).toBe('text-embedding-004');
               expect(span.attributes[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]).toBeUndefined();

packages/core/src/shared-exports.ts

@@ -205,7 +205,13 @@ export {
 } from './tracing/anthropic-ai';
 export { instrumentAsyncIterableStream, instrumentMessageStream } from './tracing/anthropic-ai/streaming';
 export { ANTHROPIC_AI_INTEGRATION_NAME } from './tracing/anthropic-ai/constants';
-export { instrumentGoogleGenAIClient } from './tracing/google-genai';
+export {
+  instrumentGoogleGenAIClient,
+  extractRequestAttributes as extractGoogleGenAIRequestAttributes,
+  addPrivateRequestAttributes as addGoogleGenAIRequestAttributes,
+  addResponseAttributes as addGoogleGenAIResponseAttributes,
+} from './tracing/google-genai';
+export { instrumentStream as instrumentGoogleGenAIStream } from './tracing/google-genai/streaming';
 export { GOOGLE_GENAI_INTEGRATION_NAME } from './tracing/google-genai/constants';
 export type { GoogleGenAIResponse } from './tracing/google-genai/types';
 export { createLangChainCallbackHandler, instrumentLangChainEmbeddings } from './tracing/langchain';

packages/core/src/tracing/google-genai/index.ts

@@ -101,7 +101,7 @@ function extractConfigAttributes(config: Record<string, unknown>): Record<string
  * Extract request attributes from method arguments
  * Builds the base attributes for span creation including system info, model, and config
  */
-function extractRequestAttributes(
+export function extractRequestAttributes(
   operationName: string,
   params?: Record<string, unknown>,
   context?: unknown,
@@ -140,13 +140,13 @@ function extractRequestAttributes(
  * This is only recorded if recordInputs is true.
  * Handles different parameter formats for different Google GenAI methods.
  */
-function addPrivateRequestAttributes(
+export function addPrivateRequestAttributes(
   span: Span,
   params: Record<string, unknown>,
-  isEmbeddings: boolean,
+  operationName: string,
   enableTruncation: boolean,
 ): void {
-  if (isEmbeddings) {
+  if (operationName === 'embeddings') {
     const contents = params.contents;
     if (contents != null) {
       span.setAttribute(
@@ -206,7 +206,7 @@ function addPrivateRequestAttributes(
  * Add response attributes from the Google GenAI response
  * @see https://github.com/googleapis/js-genai/blob/v1.19.0/src/types.ts#L2313
  */
-function addResponseAttributes(span: Span, response: GoogleGenAIResponse, recordOutputs?: boolean): void {
+export function addResponseAttributes(span: Span, response: GoogleGenAIResponse, recordOutputs?: boolean): void {
   if (!response || typeof response !== 'object') return;
 
   if (response.modelVersion) {
@@ -301,7 +301,7 @@ function instrumentMethod<T extends unknown[], R>(
                 addPrivateRequestAttributes(
                   span,
                   params,
-                  isEmbeddings,
+                  operationName,
                   shouldEnableTruncation(options.enableTruncation),
                 );
               }
@@ -331,7 +331,7 @@ function instrumentMethod<T extends unknown[], R>(
         },
         (span: Span) => {
           if (options.recordInputs && params) {
-            addPrivateRequestAttributes(span, params, isEmbeddings, shouldEnableTruncation(options.enableTruncation));
+            addPrivateRequestAttributes(span, params, operationName, shouldEnableTruncation(options.enableTruncation));
           }
 
           return handleCallbackErrors(

packages/server-utils/src/integrations/tracing-channel/google-genai.ts

@@ -0,0 +1,170 @@
+import * as diagnosticsChannel from 'node:diagnostics_channel';
+import type { GoogleGenAIOptions, GoogleGenAIResponse, IntegrationFn, Span } from '@sentry/core';
+import {
+  _INTERNAL_shouldSkipAiProviderWrapping,
+  addGoogleGenAIRequestAttributes,
+  addGoogleGenAIResponseAttributes,
+  debug,
+  defineIntegration,
+  extractGoogleGenAIRequestAttributes,
+  GEN_AI_REQUEST_MODEL_ATTRIBUTE,
+  getActiveSpan,
+  instrumentGoogleGenAIStream,
+  resolveAIRecordingOptions,
+  SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
+  shouldEnableTruncation,
+  spanToJSON,
+  startInactiveSpan,
+  waitForTracingChannelBinding,
+} from '@sentry/core';
+import { DEBUG_BUILD } from '../../debug-build';
+import { CHANNELS } from '../../orchestrion/channels';
+import { bindTracingChannelToSpan } from '../../tracing-channel';
+
+// Same name as the OTel integration by design: when enabled, the OTel 'Google_GenAI'
+// integration is dropped from the default set (see the Node opt-in loader).
+const INTEGRATION_NAME = 'Google_GenAI' as const;
+
+// Distinct from the proxy's `auto.ai.google_genai` so spans from the orchestrion path
+// are attributable separately from the OTel/proxy one.
+const ORIGIN = 'auto.ai.orchestrion.google_genai';
+
+// Each instrumented method maps to the gen_ai operation its span reports.
+const INSTRUMENTED_CHANNELS = [
+  { channel: CHANNELS.GOOGLE_GENAI_GENERATE_CONTENT, operation: 'generate_content' },
+  { channel: CHANNELS.GOOGLE_GENAI_EMBED_CONTENT, operation: 'embeddings' },
+  { channel: CHANNELS.GOOGLE_GENAI_CHAT, operation: 'chat' },
+] as const;
+
+interface GoogleGenAIChannelContext {
+  arguments: unknown[];
+  // The transform stashes the call's `this` here, which chat methods need since the model lives on
+  // the `Chat` instance (`this.model`/`this.modelVersion`), not in the call arguments.
+  self?: unknown;
+  result?: unknown;
+}
+
+let subscribed = false;
+
+const _googleGenAIChannelIntegration = ((options: GoogleGenAIOptions = {}) => {
+  return {
+    name: INTEGRATION_NAME,
+    setupOnce() {
+      // `tracingChannel` is unavailable before Node 18.19, and a second `init()` would double-subscribe.
+      if (!diagnosticsChannel.tracingChannel || subscribed) {
+        return;
+      }
+      subscribed = true;
+
+      // `bindTracingChannelToSpan` needs the async-context binding that `initOpenTelemetry()` registers
+      // after `setupOnce` runs, so wait for it before subscribing.
+      waitForTracingChannelBinding(() => {
+        for (const { channel, operation } of INSTRUMENTED_CHANNELS) {
+          DEBUG_BUILD && debug.log(`[orchestrion:google-genai] subscribing to channel "${channel}"`);
+          bindTracingChannelToSpan(
+            diagnosticsChannel.tracingChannel<GoogleGenAIChannelContext>(channel),
+            data => createGenAiSpan(data, operation, options),
+            {
+              beforeSpanEnd: (span, data) => {
+                // Embeddings responses carry no content attributes.
+                if (operation !== 'embeddings') {
+                  addGoogleGenAIResponseAttributes(
+                    span,
+                    data.result as GoogleGenAIResponse,
+                    resolveAIRecordingOptions(options).recordOutputs,
+                  );
+                }
+              },
+              deferSpanEnd: ({ span, data }) => wrapStreamResult(span, data, options),
+            },
+          );
+        }
+      });
+    },
+  };
+}) satisfies IntegrationFn;
+
+/**
+ * Build the span for an instrumented call.
+ * Returning `undefined` opts the payload out so no span is opened.
+ */
+function createGenAiSpan(
+  data: GoogleGenAIChannelContext,
+  operation: string,
+  options: GoogleGenAIOptions,
+): Span | undefined {
+  // When another provider (e.g. LangChain) is driving the SDK, it records the spans itself and marks this
+  // provider as skipped; skip here to avoid double spans.
+  if (_INTERNAL_shouldSkipAiProviderWrapping(INTEGRATION_NAME)) {
+    return undefined;
+  }
+
+  // `chat.sendMessage()`/`sendMessageStream()` internally call `Models.generateContent(Stream)`, which
+  // publishes the `generate-content` channel while the chat span is active. Skip that nested event so a
+  // chat call yields a single `gen_ai.chat` span instead of a chat span wrapping a generate_content one.
+  if (operation !== 'chat') {
+    const activeSpan = getActiveSpan();
+    if (activeSpan) {
+      const { op, origin } = spanToJSON(activeSpan);
+      if (origin === ORIGIN && op === 'gen_ai.chat') {
+        return undefined;
+      }
+    }
+  }
+
+  const args = data.arguments ?? [];
+  const params = args[0] as Record<string, unknown> | undefined;
+
+  const { recordInputs } = resolveAIRecordingOptions(options);
+  const enableTruncation = shouldEnableTruncation(options.enableTruncation);
+
+  const attributes = extractGoogleGenAIRequestAttributes(operation, params, data.self);
+  const model = (attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] as string) || 'unknown';
+  attributes[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN] = ORIGIN;
+
+  const span = startInactiveSpan({
+    name: `${operation} ${model}`,
+    op: `gen_ai.${operation}`,
+    attributes,
+  });
+
+  if (recordInputs && params) {
+    addGoogleGenAIRequestAttributes(span, params, operation, enableTruncation);
+  }
+
+  return span;
+}
+
+type AsyncIterableStream = { [Symbol.asyncIterator]: () => AsyncIterator<unknown> };
+
+function isAsyncIterable(value: unknown): value is AsyncIterableStream {
+  return !!value && typeof (value as AsyncIterableStream)[Symbol.asyncIterator] === 'function';
+}
+
+/**
+ * Only the streaming methods (`generateContentStream`/`sendMessageStream`) resolve to an async iterable.
+ * For a stream we patch `result[Symbol.asyncIterator]` in place so `instrumentGoogleGenAIStream` ends the
+ * span when iteration finishes.
+ */
+function wrapStreamResult(span: Span, data: GoogleGenAIChannelContext, options: GoogleGenAIOptions): boolean {
+  const result = data.result;
+  if (!isAsyncIterable(result)) {
+    return false;
+  }
+
+  const { recordOutputs } = resolveAIRecordingOptions(options);
+  const iterate = result[Symbol.asyncIterator].bind(result);
+  const instrumented = instrumentGoogleGenAIStream({ [Symbol.asyncIterator]: iterate }, span, recordOutputs ?? false);
+  result[Symbol.asyncIterator] = () => instrumented;
+
+  return true;
+}
+
+/**
+ * EXPERIMENTAL — orchestrion-driven Google GenAI integration. Subscribes to the
+ * `orchestrion:@google/genai:*` diagnostics_channels injected into the SDK's `Models`
+ * (`generateContent`/`generateContentStream`/`embedContent`) and `Chat`
+ * (`sendMessage`/`sendMessageStream`) methods, so it requires the orchestrion runtime hook or
+ * bundler plugin.
+ */
+export const googleGenAIChannelIntegration = defineIntegration(_googleGenAIChannelIntegration);

packages/server-utils/src/integrations/tracing-channel/openai.ts

@@ -83,7 +83,8 @@ const _openaiChannelIntegration = ((options: OpenAiOptions = {}) => {
  * Returning `undefined` opts the payload out so no span is opened.
  */
 function createGenAiSpan(data: OpenAiChatChannelContext, operation: string, options: OpenAiOptions): Span | undefined {
-  // langchain drives the openai SDK itself and instruments at its own layer; skip here to avoid double spans.
+  // When another provider (e.g. LangChain) is driving the SDK, it records the spans itself and marks this
+  // provider as skipped; skip here to avoid double spans.
   if (_INTERNAL_shouldSkipAiProviderWrapping(INTEGRATION_NAME)) {
     return undefined;
   }

packages/server-utils/src/orchestrion/channels.ts

@@ -4,6 +4,7 @@ import { ioredisChannels } from './config/ioredis';
 import { pgChannels } from './config/pg';
 import { openaiChannels } from './config/openai';
 import { anthropicAiChannels } from './config/anthropic-ai';
+import { googleGenAiChannels } from './config/google-genai';
 import { vercelAiChannels } from './config/vercel-ai';
 import { hapiChannels } from './config/hapi';
 
@@ -27,6 +28,7 @@ export const CHANNELS = {
   ...pgChannels,
   ...openaiChannels,
   ...anthropicAiChannels,
+  ...googleGenAiChannels,
   ...vercelAiChannels,
   ...hapiChannels,
 } as const;

packages/server-utils/src/orchestrion/config/google-genai.ts

@@ -0,0 +1,40 @@
+import type { InstrumentationConfig } from '@apm-js-collab/code-transformer';
+
+// `@google/genai` ships one bundled file per module format and the matcher compares `filePath` exactly,
+// so we list every file the `node` export condition resolves to across the supported range: `index.js`
+// (ESM+CJS for <0.15.0, CJS for <1.1.0), `index.mjs` (ESM for >=0.15.0), and `index.cjs` (CJS for >=1.1.0).
+// A file that doesn't exist in a given version simply never matches, so listing all three is safe.
+const NODE_DIST_FILES = ['dist/node/index.js', 'dist/node/index.mjs', 'dist/node/index.cjs'];
+
+export const googleGenAiConfig = [
+  // `generateContent`/`generateContentStream` are arrow properties assigned in the constructor, not class
+  // methods, so they need `expressionName` rather than `className`/`methodName`.
+  ...NODE_DIST_FILES.flatMap(filePath =>
+    (['generateContent', 'generateContentStream'] as const).map(expressionName => ({
+      channelName: 'generate-content',
+      module: { name: '@google/genai', versionRange: '>=0.10.0 <2', filePath },
+      functionQuery: { expressionName, kind: 'Auto' as const },
+    })),
+  ),
+  // `embedContent` and the `Chat` methods are real class methods.
+  ...NODE_DIST_FILES.map(filePath => ({
+    channelName: 'embed-content',
+    module: { name: '@google/genai', versionRange: '>=0.10.0 <2', filePath },
+    functionQuery: { className: 'Models', methodName: 'embedContent', kind: 'Auto' as const },
+  })),
+  // `sendMessage`/`sendMessageStream` internally delegate to `Models.generateContent(Stream)`; the
+  // subscriber suppresses that nested `generate-content` event so a chat call yields a single span.
+  ...NODE_DIST_FILES.flatMap(filePath =>
+    (['sendMessage', 'sendMessageStream'] as const).map(methodName => ({
+      channelName: 'chat',
+      module: { name: '@google/genai', versionRange: '>=0.10.0 <2', filePath },
+      functionQuery: { className: 'Chat', methodName, kind: 'Auto' as const },
+    })),
+  ),
+] satisfies InstrumentationConfig[];
+
+export const googleGenAiChannels = {
+  GOOGLE_GENAI_GENERATE_CONTENT: 'orchestrion:@google/genai:generate-content',
+  GOOGLE_GENAI_EMBED_CONTENT: 'orchestrion:@google/genai:embed-content',
+  GOOGLE_GENAI_CHAT: 'orchestrion:@google/genai:chat',
+} as const;

packages/server-utils/src/orchestrion/config/index.ts

@@ -5,6 +5,7 @@ import { ioredisConfig } from './ioredis';
 import { openaiConfig } from './openai';
 import { pgConfig } from './pg';
 import { anthropicAiConfig } from './anthropic-ai';
+import { googleGenAiConfig } from './google-genai';
 import { vercelAiConfig } from './vercel-ai';
 import { hapiConfig } from './hapi';
 
@@ -15,6 +16,7 @@ export const SENTRY_INSTRUMENTATIONS: InstrumentationConfig[] = [
   ...openaiConfig,
   ...pgConfig,
   ...anthropicAiConfig,
+  ...googleGenAiConfig,
   ...vercelAiConfig,
   ...hapiConfig,
 ];

packages/server-utils/src/orchestrion/index.ts

@@ -1,4 +1,5 @@
 import { anthropicChannelIntegration } from '../integrations/tracing-channel/anthropic';
+import { googleGenAIChannelIntegration } from '../integrations/tracing-channel/google-genai';
 import { hapiChannelIntegration } from '../integrations/tracing-channel/hapi';
 import { ioredisChannelIntegration } from '../integrations/tracing-channel/ioredis';
 import { lruMemoizerChannelIntegration } from '../integrations/tracing-channel/lru-memoizer';
@@ -10,6 +11,7 @@ import { vercelAiChannelIntegration } from '../integrations/tracing-channel/verc
 export { detectOrchestrionSetup, isOrchestrionInjected } from './detect';
 export {
   anthropicChannelIntegration,
+  googleGenAIChannelIntegration,
   hapiChannelIntegration,
   ioredisChannelIntegration,
   lruMemoizerChannelIntegration,
@@ -39,6 +41,7 @@ export const channelIntegrations = {
   lruMemoizerIntegration: lruMemoizerChannelIntegration,
   openaiIntegration: openaiChannelIntegration,
   anthropicIntegration: anthropicChannelIntegration,
+  googleGenAIIntegration: googleGenAIChannelIntegration,
   vercelAiIntegration: vercelAiChannelIntegration,
   hapiIntegration: hapiChannelIntegration,
 } as const;