update with tests

RulaKhaled · RulaKhaled · commit 60eb2d16d154 · 2025-09-18T18:33:58.000+02:00
diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/instrument-with-options.mjs b/dev-packages/node-integration-tests/suites/tracing/google-genai/instrument-with-options.mjs
@@ -7,11 +7,17 @@ Sentry.init({
   tracesSampleRate: 1.0,
   sendDefaultPii: false,
   transport: loggingTransport,
-  registerEsmLoaderHooks: false,
   integrations: [
     Sentry.googleGenAIIntegration({
       recordInputs: true,
       recordOutputs: true,
     }),
   ],
+  beforeSendTransaction: event => {
+    // Filter out mock express server transactions
+    if (event.transaction.includes('/v1beta/')) {
+      return null;
+    }
+    return event;
+  },
 });
diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/instrument-with-pii.mjs b/dev-packages/node-integration-tests/suites/tracing/google-genai/instrument-with-pii.mjs
@@ -7,6 +7,12 @@ Sentry.init({
   tracesSampleRate: 1.0,
   sendDefaultPii: true,
   transport: loggingTransport,
-  registerEsmLoaderHooks: false,
   integrations: [Sentry.googleGenAIIntegration()],
+  beforeSendTransaction: event => {
+    // Filter out mock express server transactions
+    if (event.transaction.includes('/v1beta/')) {
+      return null;
+    }
+    return event;
+  },
 });
diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/instrument.mjs b/dev-packages/node-integration-tests/suites/tracing/google-genai/instrument.mjs
@@ -7,6 +7,12 @@ Sentry.init({
   tracesSampleRate: 1.0,
   sendDefaultPii: false,
   transport: loggingTransport,
-  registerEsmLoaderHooks: false,
   integrations: [Sentry.googleGenAIIntegration()],
+  beforeSendTransaction: event => {
+    // Filter out mock express server transactions
+    if (event.transaction.includes('/v1beta')) {
+      return null;
+    }
+    return event;
+  },
 });
diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario.mjs
@@ -1,89 +1,59 @@
-import { instrumentGoogleGenAIClient } from '@sentry/core';
 import * as Sentry from '@sentry/node';
+import { GoogleGenAI } from '@google/genai';
 
-class MockGoogleGenAI {
-  constructor(config) {
-    this.apiKey = config.apiKey;
+import express from 'express';
 
-    this.models = {
-      generateContent: async params => {
-        // Simulate processing time
-        await new Promise(resolve => setTimeout(resolve, 10));
+const PORT = 3333;
 
-        if (params.model === 'error-model') {
-          const error = new Error('Model not found');
-          error.status = 404;
-          throw error;
-        }
+function startMockGoogleGenAIServer() {
+  const app = express();
+  app.use(express.json());
 
-        return {
-          candidates: [
-            {
-              content: {
-                parts: [
-                  {
-                    text: params.contents ? 'The capital of France is Paris.' : 'Mock response from Google GenAI!',
-                  },
-                ],
-                role: 'model',
-              },
-              finishReason: 'stop',
-              index: 0,
-            },
-          ],
-          usageMetadata: {
-            promptTokenCount: 8,
-            candidatesTokenCount: 12,
-            totalTokenCount: 20,
-          },
-        };
-      },
-    };
+  app.post('/v1beta/models/:model\\:generateContent', (req, res) => {
+    const model = req.params.model;
 
-    this.chats = {
-      create: options => {
-        // Return a chat instance with sendMessage method and model info
-        return {
-          model: options?.model || 'unknown', // Include model from create options
-          sendMessage: async () => {
-            // Simulate processing time
-            await new Promise(resolve => setTimeout(resolve, 10));
+    if (model === 'error-model') {
+      res.status(404).set('x-request-id', 'mock-request-123').end('Model not found');
+      return;
+    }
 
-            return {
-              candidates: [
-                {
-                  content: {
-                    parts: [
-                      {
-                        text: 'Mock response from Google GenAI!',
-                      },
-                    ],
-                    role: 'model',
-                  },
-                  finishReason: 'stop',
-                  index: 0,
-                },
-              ],
-              usageMetadata: {
-                promptTokenCount: 10,
-                candidatesTokenCount: 15,
-                totalTokenCount: 25,
+    res.send({
+      candidates: [
+        {
+          content: {
+            parts: [
+              {
+                text: 'Mock response from Google GenAI!',
               },
-            };
+            ],
+            role: 'model',
           },
-        };
+          finishReason: 'stop',
+          index: 0,
+        },
+      ],
+      usageMetadata: {
+        promptTokenCount: 8,
+        candidatesTokenCount: 12,
+        totalTokenCount: 20,
       },
-    };
-  }
+    });
+  });
+
+  return app.listen(PORT);
 }
 
 async function run() {
-  const genAI = new MockGoogleGenAI({ apiKey: 'test-api-key' });
-  const instrumentedClient = instrumentGoogleGenAIClient(genAI);
+  const server = startMockGoogleGenAIServer();
+
+  await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
+    const client = new GoogleGenAI({
+      apiKey: 'mock-api-key',
+      httpOptions: { baseUrl: `http://localhost:${PORT}` }
+    });
 
-  await Sentry.startSpan({ name: 'main', op: 'function' }, async () => {
     // Test 1: chats.create and sendMessage flow
-    const chat = instrumentedClient.chats.create({
+    const chat = client.chats.create({
       model: 'gemini-1.5-pro',
       config: {
         temperature: 0.8,
@@ -103,7 +73,7 @@ async function run() {
     });
 
     // Test 2: models.generateContent
-    await instrumentedClient.models.generateContent({
+    await client.models.generateContent({
       model: 'gemini-1.5-flash',
       config: {
         temperature: 0.7,
@@ -120,7 +90,7 @@ async function run() {
 
     // Test 3: Error handling
     try {
-      await instrumentedClient.models.generateContent({
+      await client.models.generateContent({
         model: 'error-model',
         contents: [
           {
@@ -133,6 +103,8 @@ async function run() {
       // Expected error
     }
   });
+
+  server.close();
 }
 
 run();
diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts b/dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts
@@ -34,9 +34,9 @@ describe('Google GenAI integration', () => {
           'sentry.origin': 'auto.ai.google_genai',
           'gen_ai.system': 'google_genai',
           'gen_ai.request.model': 'gemini-1.5-pro', // Should get from chat context
-          'gen_ai.usage.input_tokens': 10,
-          'gen_ai.usage.output_tokens': 15,
-          'gen_ai.usage.total_tokens': 25,
+          'gen_ai.usage.input_tokens': 8,
+          'gen_ai.usage.output_tokens': 12,
+          'gen_ai.usage.total_tokens': 20,
         },
         description: 'chat gemini-1.5-pro',
         op: 'gen_ai.chat',
@@ -111,9 +111,9 @@ describe('Google GenAI integration', () => {
           'gen_ai.request.model': 'gemini-1.5-pro',
           'gen_ai.request.messages': expect.any(String), // Should include message when recordInputs: true
           'gen_ai.response.text': expect.any(String), // Should include response when recordOutputs: true
-          'gen_ai.usage.input_tokens': 10,
-          'gen_ai.usage.output_tokens': 15,
-          'gen_ai.usage.total_tokens': 25,
+          'gen_ai.usage.input_tokens': 8,
+          'gen_ai.usage.output_tokens': 12,
+          'gen_ai.usage.total_tokens': 20,
         }),
         description: 'chat gemini-1.5-pro',
         op: 'gen_ai.chat',
diff --git a/packages/core/src/utils/google-genai/index.ts b/packages/core/src/utils/google-genai/index.ts
@@ -3,7 +3,6 @@ import { captureException } from '../../exports';
 import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes';
 import { startSpan } from '../../tracing/trace';
 import type { Span, SpanAttributeValue } from '../../types-hoist/span';
-import { handleCallbackErrors } from '../handleCallbackErrors';
 import {
   GEN_AI_OPERATION_NAME_ATTRIBUTE,
   GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE,
@@ -21,11 +20,11 @@ import {
   GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
 import { buildMethodPath, getFinalOperationName, getSpanOperation } from '../ai/utils';
-import { CHAT_PATH, CHATS_CREATE_METHOD, GOOGLE_GENAI_INTEGRATION_NAME, GOOGLE_GENAI_SYSTEM_NAME } from './constants';
+import { handleCallbackErrors } from '../handleCallbackErrors';
+import { CHAT_PATH, CHATS_CREATE_METHOD, GOOGLE_GENAI_SYSTEM_NAME } from './constants';
 import type {
   Candidate,
   ContentPart,
-  GoogleGenAIIntegration,
   GoogleGenAIIstrumentedMethod,
   GoogleGenAIOptions,
   GoogleGenAIResponse,
@@ -189,21 +188,6 @@ function addResponseAttributes(span: Span, response: GoogleGenAIResponse, record
   }
 }
 
-/**
- * Get recording options from the Sentry integration configuration
- * Falls back to sendDefaultPii setting if integration options are not specified
- */
-function getRecordingOptionsFromIntegration(): GoogleGenAIOptions {
-  const client = getClient();
-  const integration = client?.getIntegrationByName<GoogleGenAIIntegration>(GOOGLE_GENAI_INTEGRATION_NAME);
-  const shouldRecordInputsAndOutputs = Boolean(client?.getOptions().sendDefaultPii);
-
-  return {
-    recordInputs: integration?.options?.recordInputs ?? shouldRecordInputsAndOutputs,
-    recordOutputs: integration?.options?.recordOutputs ?? shouldRecordInputsAndOutputs,
-  };
-}
-
 /**
  * Instrument any async or synchronous genai method with Sentry spans
  * Handles operations like models.generateContent and chat.sendMessage and chats.create
@@ -213,12 +197,11 @@ function instrumentMethod<T extends unknown[], R>(
   originalMethod: (...args: T) => R | Promise<R>,
   methodPath: GoogleGenAIIstrumentedMethod,
   context: unknown,
-  options?: GoogleGenAIOptions,
+  options: GoogleGenAIOptions,
 ): (...args: T) => R | Promise<R> {
   const isSyncCreate = methodPath === CHATS_CREATE_METHOD;
 
   const run = (...args: T): R | Promise<R> => {
-    const finalOptions = options || getRecordingOptionsFromIntegration();
     const requestAttributes = extractRequestAttributes(args, methodPath, context);
     const model = requestAttributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] ?? 'unknown';
     const operationName = getFinalOperationName(methodPath);
@@ -231,7 +214,7 @@ function instrumentMethod<T extends unknown[], R>(
         attributes: requestAttributes,
       },
       (span: Span) => {
-        if (finalOptions.recordInputs && args[0] && typeof args[0] === 'object') {
+        if (options.recordInputs && args[0] && typeof args[0] === 'object') {
           addPrivateRequestAttributes(span, args[0] as Record<string, unknown>);
         }
 
@@ -246,7 +229,7 @@ function instrumentMethod<T extends unknown[], R>(
           result => {
             // Only add response attributes for content-producing methods, not for chats.create
             if (!isSyncCreate) {
-              addResponseAttributes(span, result, finalOptions.recordOutputs);
+              addResponseAttributes(span, result, options.recordOutputs);
             }
           },
         );
@@ -261,7 +244,7 @@ function instrumentMethod<T extends unknown[], R>(
  * Create a deep proxy for Google GenAI client instrumentation
  * Recursively instruments methods and handles special cases like chats.create
  */
-function createDeepProxy<T extends object>(target: T, currentPath = '', options?: GoogleGenAIOptions): T {
+function createDeepProxy<T extends object>(target: T, currentPath = '', options: GoogleGenAIOptions): T {
   return new Proxy(target, {
     get: (t, prop, receiver) => {
       const value = Reflect.get(t, prop, receiver);
@@ -321,5 +304,12 @@ function createDeepProxy<T extends object>(target: T, currentPath = '', options?
  * ```
  */
 export function instrumentGoogleGenAIClient<T extends object>(client: T, options?: GoogleGenAIOptions): T {
-  return createDeepProxy(client, '', options);
+  const sendDefaultPii = Boolean(getClient()?.getOptions().sendDefaultPii);
+
+  const _options = {
+    recordInputs: sendDefaultPii,
+    recordOutputs: sendDefaultPii,
+    ...options,
+  };
+  return createDeepProxy(client, '', _options);
 }
diff --git a/packages/core/src/utils/google-genai/types.ts b/packages/core/src/utils/google-genai/types.ts
@@ -1,7 +1,6 @@
-import type { InstrumentationConfig } from '@opentelemetry/instrumentation';
 import type { GOOGLE_GENAI_INSTRUMENTED_METHODS } from './constants';
 
-export interface GoogleGenAIOptions extends InstrumentationConfig {
+export interface GoogleGenAIOptions {
   /**
    * Enable or disable input recording.
    */
diff --git a/packages/node/src/integrations/tracing/google-genai/instrumentation.ts b/packages/node/src/integrations/tracing/google-genai/instrumentation.ts
