diff --git a/packages/core/src/tracing/openai/index.ts b/packages/core/src/tracing/openai/index.ts index c68e920daf2b..0c5f2e255b86 100644 --- a/packages/core/src/tracing/openai/index.ts +++ b/packages/core/src/tracing/openai/index.ts @@ -2,8 +2,9 @@ import { getClient } from '../../currentScopes'; import { captureException } from '../../exports'; import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes'; import { SPAN_STATUS_ERROR } from '../../tracing'; -import { startSpan, startSpanManual } from '../../tracing/trace'; +import { startInactiveSpan, startSpanManual } from '../../tracing/trace'; import type { Span, SpanAttributeValue } from '../../types-hoist/span'; +import { isThenable } from '../../utils/is'; import { GEN_AI_OPERATION_NAME_ATTRIBUTE, GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE, @@ -126,97 +127,150 @@ function addRequestAttributes(span: Span, params: Record): void } } +/** + * Handle common error catching and reporting for streaming requests + */ +function handleStreamingError(error: unknown, span: Span, methodPath: string): never { + captureException(error, { + mechanism: { handled: false, type: 'auto.ai.openai.stream', data: { function: methodPath } }, + }); + + if (span.isRecording()) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + span.end(); + } + throw error; +} + /** * Instrument a method with Sentry spans * Following Sentry AI Agents Manual Instrumentation conventions * @see https://docs.sentry.io/platforms/javascript/guides/node/tracing/instrumentation/ai-agents-module/#manual-instrumentation + * + * This implementation uses Proxy and startInactiveSpan to preserve the original + * return type (e.g., OpenAI's APIPromise with .withResponse() method). */ function instrumentMethod( - originalMethod: (...args: T) => Promise, + originalMethod: (...args: T) => R | Promise, methodPath: InstrumentedMethod, context: unknown, options: OpenAiOptions, -): (...args: T) => Promise { - return async function instrumentedMethod(...args: T): Promise { - const requestAttributes = extractRequestAttributes(args, methodPath); - const model = (requestAttributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] as string) || 'unknown'; - const operationName = getOperationName(methodPath); +): (...args: T) => R | Promise { + return new Proxy(originalMethod, { + apply(target, _thisArg, args: T): R | Promise { + const requestAttributes = extractRequestAttributes(args, methodPath); + const model = (requestAttributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] as string) || 'unknown'; + const operationName = getOperationName(methodPath); + + const params = args[0] as Record | undefined; + const isStreamRequested = params && typeof params === 'object' && params.stream === true; - const params = args[0] as Record | undefined; - const isStreamRequested = params && typeof params === 'object' && params.stream === true; + if (isStreamRequested) { + // For streaming responses, use manual span management to properly handle the async generator lifecycle + return startSpanManual( + { + name: `${operationName} ${model} stream-response`, + op: getSpanOperation(methodPath), + attributes: requestAttributes as Record, + }, + async (span: Span) => { + try { + if (options.recordInputs && params) { + addRequestAttributes(span, params); + } - if (isStreamRequested) { - // For streaming responses, use manual span management to properly handle the async generator lifecycle - return startSpanManual( - { - name: `${operationName} ${model} stream-response`, - op: getSpanOperation(methodPath), - attributes: requestAttributes as Record, - }, - async (span: Span) => { - try { - if (options.recordInputs && params) { - addRequestAttributes(span, params); + const result = await target.apply(context, args); + + return instrumentStream( + result as OpenAIStream, + span, + options.recordOutputs ?? false, + ) as unknown as R; + } catch (error) { + return handleStreamingError(error, span, methodPath); } + }, + ); + } - const result = await originalMethod.apply(context, args); + // Non-streaming responses: use startInactiveSpan to preserve original return type + // (e.g., OpenAI's APIPromise with .withResponse()) + // + // We use startInactiveSpan instead of startSpan/startSpanManual because those + // internally use handleCallbackErrors which calls .then() on Promises, creating + // a new Promise instance and losing APIPromise's custom methods like .withResponse(). + const span = startInactiveSpan({ + name: `${operationName} ${model}`, + op: getSpanOperation(methodPath), + attributes: requestAttributes as Record, + }); - return instrumentStream( - result as OpenAIStream, - span, - options.recordOutputs ?? false, - ) as unknown as R; - } catch (error) { - // For streaming requests that fail before stream creation, we still want to record - // them as streaming requests but end the span gracefully - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - captureException(error, { - mechanism: { - handled: false, - type: 'auto.ai.openai.stream', - data: { - function: methodPath, + // Handle synchronous exceptions from the API call or request attribute processing + let result: R | Promise; + try { + if (options.recordInputs && params) { + addRequestAttributes(span, params); + } + result = target.apply(context, args); + } catch (err) { + captureException(err, { + mechanism: { + handled: false, + type: 'auto.ai.openai', + data: { + function: methodPath, + }, + }, + }); + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + span.end(); + throw err; + } + + // Attach side-effect handlers without transforming the Promise + // This preserves the original APIPromise type and its methods like .withResponse() + if (isThenable(result)) { + Promise.resolve(result) + .then( + res => { + try { + addResponseAttributes(span, res as OpenAiResponse, options.recordOutputs); + } catch { + // Ignore attribute processing errors - they shouldn't affect the original Promise + // The span will still be ended in finally() + } + }, + err => { + captureException(err, { + mechanism: { + handled: false, + type: 'auto.ai.openai', + data: { + function: methodPath, + }, }, - }, - }); + }); + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + }, + ) + .finally(() => { span.end(); - throw error; - } - }, - ); - } else { - // Non-streaming responses - return startSpan( - { - name: `${operationName} ${model}`, - op: getSpanOperation(methodPath), - attributes: requestAttributes as Record, - }, - async (span: Span) => { - try { - if (options.recordInputs && params) { - addRequestAttributes(span, params); - } + }); + } else { + // Synchronous result (unlikely for OpenAI API but handle it) + try { + addResponseAttributes(span, result as OpenAiResponse, options.recordOutputs); + } catch { + // Ignore attribute processing errors + } finally { + span.end(); + } + } - const result = await originalMethod.apply(context, args); - addResponseAttributes(span, result, options.recordOutputs); - return result; - } catch (error) { - captureException(error, { - mechanism: { - handled: false, - type: 'auto.ai.openai', - data: { - function: methodPath, - }, - }, - }); - throw error; - } - }, - ); - } - }; + // Return the original Promise (APIPromise) with all its methods intact + return result; + }, + }) as (...args: T) => R | Promise; } /** diff --git a/packages/core/test/tracing/openai-integration-functions.test.ts b/packages/core/test/tracing/openai-integration-functions.test.ts index 240ba14d429b..ab76cd11e74b 100644 --- a/packages/core/test/tracing/openai-integration-functions.test.ts +++ b/packages/core/test/tracing/openai-integration-functions.test.ts @@ -2,13 +2,62 @@ import { beforeEach, describe, expect, it } from 'vitest'; import type { OpenAiClient } from '../../src'; import { instrumentOpenAiClient } from '../../src/tracing/openai'; +/** + * Mock APIPromise that simulates OpenAI SDK's APIPromise behavior + * APIPromise extends Promise but has additional methods like withResponse() + * + * IMPORTANT: We do NOT override .then() here, matching the real OpenAI SDK behavior. + * This means calling .then() on an APIPromise returns a standard Promise, losing + * the _response data and withResponse() method. The instrumentation must preserve + * the original APIPromise instance to maintain these methods. + */ +class MockAPIPromise extends Promise { + private _response: { headers: Record }; + + constructor( + executor: (resolve: (value: T) => void, reject: (reason?: unknown) => void) => void, + response?: { headers: Record }, + ) { + super(executor); + this._response = response || { headers: { 'x-request-id': 'test-request-id' } }; + } + + /** + * Simulates OpenAI's APIPromise.withResponse() method + * Returns both the data and the raw response + */ + withResponse(): Promise<{ data: T; response: { headers: Record } }> { + return this.then(data => ({ + data, + response: this._response, + })); + } +} + interface FullOpenAIClient { chat: { completions: { - create: (params: ChatCompletionParams) => Promise; + create: (params: ChatCompletionParams) => MockAPIPromise; parse: (params: ParseCompletionParams) => Promise; }; }; + embeddings: { + create: (params: EmbeddingsParams) => MockAPIPromise; + }; + responses: { + create: (params: { + model: string; + throwSync?: boolean; + rejectAsync?: boolean; + }) => MockAPIPromise; + }; +} + +interface ResponsesResponse { + id: string; + object: string; + model: string; + created_at: number; } interface ChatCompletionParams { model: string; @@ -48,6 +97,18 @@ interface ParseCompletionResponse { parsed: { name: string; age: number }; } +interface EmbeddingsParams { + model: string; + input: string | string[]; +} + +interface EmbeddingsResponse { + object: string; + model: string; + data: Array<{ embedding: number[]; index: number }>; + usage: { prompt_tokens: number; total_tokens: number }; +} + /** * Mock OpenAI client that simulates the private field behavior * that causes the "Cannot read private member" error @@ -59,9 +120,15 @@ class MockOpenAIClient implements FullOpenAIClient { // Simulate instrumented methods chat = { completions: { - create: async (params: ChatCompletionParams): Promise => { + create: (params: ChatCompletionParams): MockAPIPromise => { this.#buildURL('/chat/completions'); - return { id: 'test', model: params.model, choices: [{ message: { content: 'Hello!' } }] }; + return new MockAPIPromise(resolve => { + resolve({ + id: 'test', + model: params.model, + choices: [{ message: { content: 'Hello!' } }], + }); + }); }, // This is NOT instrumented @@ -84,6 +151,50 @@ class MockOpenAIClient implements FullOpenAIClient { }, }; + embeddings = { + create: (params: EmbeddingsParams): MockAPIPromise => { + this.#buildURL('/embeddings'); + return new MockAPIPromise(resolve => { + resolve({ + object: 'list', + model: params.model, + data: [{ embedding: [0.1, 0.2, 0.3], index: 0 }], + usage: { prompt_tokens: 10, total_tokens: 10 }, + }); + }); + }, + }; + + // responses.create is in INSTRUMENTED_METHODS, so we can test error handling here + responses = { + create: (params: { + model: string; + throwSync?: boolean; + rejectAsync?: boolean; + }): MockAPIPromise => { + // Simulate synchronous exception (e.g., validation error before API call) + if (params.throwSync) { + throw new Error('Sync error before API call'); + } + + // Simulate async rejection (e.g., API error response) + if (params.rejectAsync) { + return new MockAPIPromise((_, reject) => { + reject(new Error('Async API error')); + }); + } + + return new MockAPIPromise(resolve => { + resolve({ + id: 'resp_123', + object: 'response', + model: params.model, + created_at: Math.floor(Date.now() / 1000), + }); + }); + }, + }; + constructor() { MockOpenAIClient.#privateData.set(this, { apiKey: 'test-key', @@ -208,3 +319,129 @@ describe('OpenAI Integration Private Field Fix', () => { expect(typeof instrumentedClient.chat.completions.parse).toBe('function'); }); }); + +describe('OpenAI Integration APIPromise Preservation', () => { + let mockClient: MockOpenAIClient; + let instrumentedClient: FullOpenAIClient & OpenAiClient; + + beforeEach(() => { + mockClient = new MockOpenAIClient(); + instrumentedClient = instrumentOpenAiClient(mockClient as unknown as OpenAiClient) as FullOpenAIClient & + OpenAiClient; + }); + + it('should preserve APIPromise.withResponse() method on chat.completions.create', async () => { + const apiPromise = instrumentedClient.chat.completions.create({ + model: 'gpt-4', + messages: [{ role: 'user', content: 'test' }], + }); + + // The key test: withResponse() should exist and work + expect(typeof apiPromise.withResponse).toBe('function'); + + const { data, response } = await apiPromise.withResponse(); + + expect(data.model).toBe('gpt-4'); + expect(data.choices[0]?.message?.content).toBe('Hello!'); + expect(response.headers).toEqual({ 'x-request-id': 'test-request-id' }); + }); + + it('should preserve APIPromise.withResponse() method on embeddings.create', async () => { + const apiPromise = instrumentedClient.embeddings.create({ + model: 'text-embedding-3-small', + input: 'test input', + }); + + // The key test: withResponse() should exist and work + expect(typeof apiPromise.withResponse).toBe('function'); + + const { data, response } = await apiPromise.withResponse(); + + expect(data.model).toBe('text-embedding-3-small'); + expect(data.data[0]?.embedding).toEqual([0.1, 0.2, 0.3]); + expect(response.headers).toEqual({ 'x-request-id': 'test-request-id' }); + }); + + it('should still work with regular await on instrumented methods', async () => { + // Ensure the basic Promise behavior still works + const result = await instrumentedClient.chat.completions.create({ + model: 'gpt-4', + messages: [{ role: 'user', content: 'test' }], + }); + + expect(result.model).toBe('gpt-4'); + expect(result.choices[0]?.message?.content).toBe('Hello!'); + }); + + it('should return the exact same APIPromise instance (not a new Promise)', async () => { + // Create the mock client's APIPromise directly for comparison + const mockClient = new MockOpenAIClient(); + const originalPromise = mockClient.embeddings.create({ + model: 'text-embedding-3-small', + input: 'test', + }); + + // Instrument and call the same method + const instrumentedPromise = instrumentedClient.embeddings.create({ + model: 'text-embedding-3-small', + input: 'test', + }); + + // Both should be MockAPIPromise instances (not converted to regular Promise) + expect(originalPromise.constructor.name).toBe('MockAPIPromise'); + expect(instrumentedPromise.constructor.name).toBe('MockAPIPromise'); + + // The instrumented version should have withResponse available + expect(typeof instrumentedPromise.withResponse).toBe('function'); + + // And it should work correctly + const { data, response } = await instrumentedPromise.withResponse(); + expect(data.model).toBe('text-embedding-3-small'); + expect(response.headers['x-request-id']).toBe('test-request-id'); + }); +}); + +describe('OpenAI Integration Error Handling', () => { + let mockClient: MockOpenAIClient; + let instrumentedClient: FullOpenAIClient & OpenAiClient; + + beforeEach(() => { + mockClient = new MockOpenAIClient(); + instrumentedClient = instrumentOpenAiClient(mockClient as unknown as OpenAiClient) as FullOpenAIClient & + OpenAiClient; + }); + + it('should handle synchronous exceptions and re-throw them', async () => { + // responses.create is instrumented, so this tests the sync error path + expect(() => { + instrumentedClient.responses.create({ + model: 'gpt-4', + throwSync: true, + }); + }).toThrow('Sync error before API call'); + }); + + it('should handle rejected Promises', async () => { + // responses.create is instrumented, so this tests the async error path + const promise = instrumentedClient.responses.create({ + model: 'gpt-4', + rejectAsync: true, + }); + + await expect(promise).rejects.toThrow('Async API error'); + }); + + it('should still preserve APIPromise on success with responses.create', async () => { + const promise = instrumentedClient.responses.create({ + model: 'gpt-4', + }); + + // Should be a MockAPIPromise, not a regular Promise + expect(promise.constructor.name).toBe('MockAPIPromise'); + expect(typeof promise.withResponse).toBe('function'); + + const result = await promise; + expect(result.id).toBe('resp_123'); + expect(result.model).toBe('gpt-4'); + }); +});