some refacotor

Author: RulaKhaled

dev-packages/node-integration-tests/suites/tracing/google-genai/scenario.mjs

@@ -22,9 +22,7 @@ class MockGoogleGenAI {
               content: {
                 parts: [
                   {
-                    text: params.contents
-                      ? 'The capital of France is Paris.'
-                      : 'Mock response from Google GenAI!',
+                    text: params.contents ? 'The capital of France is Paris.' : 'Mock response from Google GenAI!',
                   },
                 ],
                 role: 'model',

packages/core/src/utils/ai/gen-ai-attributes.ts

@@ -178,15 +178,3 @@ export const OPENAI_OPERATIONS = {
  * The response timestamp from Anthropic AI (ISO string)
  */
 export const ANTHROPIC_AI_RESPONSE_TIMESTAMP_ATTRIBUTE = 'anthropic.response.timestamp';
-
-// =============================================================================
-// GOOGLE GENAI OPERATIONS
-// =============================================================================
-
-/**
- * Google GenAI API operations
- */
-export const GOOGLE_GENAI_OPERATIONS = {
-  GENERATE_CONTENT: 'generateContent',
-  STREAM_GENERATE_CONTENT: 'streamGenerateContent',
-} as const;

packages/core/src/utils/google-genai/constants.ts

@@ -5,7 +5,9 @@ export const GOOGLE_GENAI_INTEGRATION_NAME = 'Google_GenAI';
 export const GOOGLE_GENAI_INSTRUMENTED_METHODS = ['models.generateContent', 'chats.create', 'sendMessage'] as const;
 
 // Constants for internal use
-export const GOOGLE_GENAI_MODEL_PROPERTY = '_sentryGoogleGenAIModel';
 export const GOOGLE_GENAI_SYSTEM_NAME = 'google_genai';
 export const CHATS_CREATE_METHOD = 'chats.create';
 export const CHAT_PATH = 'chat';
+
+// Span attribute key for storing chat context
+export const GOOGLE_GENAI_SCOPE_MODEL_KEY = 'google_genai.chat.model';

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

@@ -3,6 +3,7 @@ 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 { getActiveSpan, getRootSpan, spanToJSON } from '../../utils/spanUtils';
 import {
   GEN_AI_OPERATION_NAME_ATTRIBUTE,
   GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE,
@@ -25,7 +26,7 @@ import {
   CHAT_PATH,
   CHATS_CREATE_METHOD,
   GOOGLE_GENAI_INTEGRATION_NAME,
-  GOOGLE_GENAI_MODEL_PROPERTY,
+  GOOGLE_GENAI_SCOPE_MODEL_KEY,
   GOOGLE_GENAI_SYSTEM_NAME,
 } from './constants';
 import type {
@@ -39,21 +40,35 @@ import type {
 import { shouldInstrument } from './utils';
 
 /**
- * Extract model from parameters or context
- * For chat instances, the model is stored during chat creation and retrieved from context
+ * Store model information on the root span for later retrieval by chat.sendMessage
  */
-export function extractModel(params: Record<string, unknown>, context?: unknown): string {
+function storeModelOnRootSpan(model: string): void {
+  const activeSpan = getActiveSpan();
+  if (activeSpan) {
+    const rootSpan = getRootSpan(activeSpan);
+    rootSpan.setAttributes({
+      [GOOGLE_GENAI_SCOPE_MODEL_KEY]: model,
+    });
+  }
+}
+
+/**
+ * Extract model from parameters or root span attributes
+ * For chat instances, the model is stored on the root span during chat creation
+ */
+export function extractModel(params: Record<string, unknown>, _context?: unknown): string {
   if ('model' in params && typeof params.model === 'string') {
     return params.model;
   }
 
-  // For chat instances, try to get the model from the chat context
-  // This is because the model is set during chat creation
-  // and not passed as a parameter to the chat.sendMessage method
-  if (context && typeof context === 'object') {
-    const chatObj = context as Record<string, unknown>;
-    if (chatObj[GOOGLE_GENAI_MODEL_PROPERTY] && typeof chatObj[GOOGLE_GENAI_MODEL_PROPERTY] === 'string') {
-      return chatObj[GOOGLE_GENAI_MODEL_PROPERTY] as string;
+  // Try to get model from root span attributes
+  const activeSpan = getActiveSpan();
+  if (activeSpan) {
+    const rootSpan = getRootSpan(activeSpan);
+    const spanData = spanToJSON(rootSpan);
+    const modelFromSpan = spanData.data?.[GOOGLE_GENAI_SCOPE_MODEL_KEY];
+    if (typeof modelFromSpan === 'string') {
+      return modelFromSpan;
     }
   }
 
@@ -238,12 +253,11 @@ function instrumentMethod<T extends unknown[], R>(
             if (finalOptions.recordInputs && args[0] && typeof args[0] === 'object') {
               addPrivateRequestAttributes(span, args[0] as Record<string, unknown>);
             }
-            const result = (originalMethod as (...args: T) => R).apply(context, args) as R;
+            const result = (originalMethod as (...args: T) => R).apply(context, args);
 
-            if (typeof model === 'string' && model !== 'unknown' && typeof result === 'object') {
-              // We store the model in the result object so that it can be accessed later
-              // This is because the model is not passed as a parameter to the chat.sendMessage method
-              (result as Record<string, unknown>)[GOOGLE_GENAI_MODEL_PROPERTY] = model;
+            if (typeof model === 'string' && model !== 'unknown') {
+              // Store model information on root span for later retrieval by chat.sendMessage
+              storeModelOnRootSpan(model);
             }
 
             // No response attributes for create (returns object of chat instance, not generated content)
@@ -255,7 +269,7 @@ function instrumentMethod<T extends unknown[], R>(
             throw error;
           }
         },
-      ) as R;
+      );
     }
 
     // Async/content-producing path
@@ -273,15 +287,15 @@ function instrumentMethod<T extends unknown[], R>(
 
           const result = await Promise.resolve((originalMethod as (...args: T) => Promise<R>).apply(context, args));
           addResponseAttributes(span, result as GoogleGenAIResponse, finalOptions.recordOutputs);
-          return result as R;
+          return result;
         } catch (error) {
           captureException(error, {
             mechanism: { handled: false, type: 'auto.ai.google_genai', data: { function: methodPath } },
           });
           throw error;
         }
       },
-    ) as Promise<R>;
+    );
   };
 
   return run;

packages/node/src/integrations/tracing/google-genai/instrumentation.ts

@@ -4,11 +4,16 @@ import {
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
 } from '@opentelemetry/instrumentation';
-import type { GoogleGenAIClient, GoogleGenAIOptions, Integration } from '@sentry/core';
+import type { GoogleGenAIClient, Integration } from '@sentry/core';
 import { getCurrentScope, GOOGLE_GENAI_INTEGRATION_NAME, instrumentGoogleGenAIClient, SDK_VERSION } from '@sentry/core';
 
 const supportedVersions = ['>=0.10.0 <2'];
 
+interface GoogleGenAIOptions {
+  recordInputs?: boolean;
+  recordOutputs?: boolean;
+}
+
 export interface GoogleGenAIIntegration extends Integration {
   options: GoogleGenAIOptions;
 }
@@ -28,9 +33,11 @@ function determineRecordingSettings(
   integrationOptions: GoogleGenAIOptions | undefined,
   defaultEnabled: boolean,
 ): { recordInputs: boolean; recordOutputs: boolean } {
-  const recordInputs = integrationOptions?.recordInputs ?? defaultEnabled;
-  const recordOutputs = integrationOptions?.recordOutputs ?? defaultEnabled;
-  return { recordInputs, recordOutputs };
+  if (!integrationOptions?.recordInputs || !integrationOptions.recordOutputs) {
+    return { recordInputs: defaultEnabled, recordOutputs: defaultEnabled };
+  }
+
+  return { recordInputs: integrationOptions.recordInputs, recordOutputs: integrationOptions.recordOutputs };
 }
 
 /**