Skip to content

Commit 6b09503

Browse files
mydeamydea
committed
refactor to fix stuff
1 parent e0997f3 commit 6b09503

File tree

5 files changed

+61
-38
lines changed

5 files changed

+61
-38
lines changed

‎packages/cloudflare/src/integrations/tracing/vercelai.ts

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
*/
1010

1111
import type { Client, IntegrationFn } from '@sentry/core';
12-
import { defineIntegration, processVercelAiSpan } from '@sentry/core';
12+
import { addVercelAiProcessors, defineIntegration } from '@sentry/core';
1313
import type { modulesIntegration } from '../modules';
1414

1515
interface VercelAiOptions {
@@ -36,12 +36,8 @@ const _vercelAIIntegration = ((options: VercelAiOptions = {}) => {
3636
name: INTEGRATION_NAME,
3737
options,
3838
setup(client) {
39-
function registerProcessors(): void {
40-
client.on('spanEnd', processVercelAiSpan);
41-
}
42-
4339
if (options.force || shouldRunIntegration(client)) {
44-
registerProcessors();
40+
addVercelAiProcessors(client);
4541
}
4642
},
4743
};

‎packages/core/src/client.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,8 @@ export abstract class Client<O extends ClientOptions = ClientOptions> {
498498
): void;
499499

500500
/**
501-
* Register a callback for whenever a span is ended.
501+
* Register a callback for after a span is ended.
502+
* NOTE: The span cannot be mutated anymore in this callback.
502503
* Receives the span as argument.
503504
* @returns {() => void} A function that, when executed, removes the registered callback.
504505
*/

‎packages/core/src/index.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ export { captureFeedback } from './feedback';
123123
export type { ReportDialogOptions } from './report-dialog';
124124
export { _INTERNAL_captureLog, _INTERNAL_flushLogsBuffer, _INTERNAL_captureSerializedLog } from './logs/exports';
125125
export { consoleLoggingIntegration } from './logs/console-integration';
126-
export { processVercelAiSpan } from './utils/vercel-ai';
126+
export { addVercelAiProcessors } from './utils/vercel-ai';
127127

128128
export type { FeatureFlag } from './utils/featureFlags';
129129
export {

‎packages/core/src/utils/vercel-ai.ts

Lines changed: 53 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
1+
import type { Client } from '../client';
12
import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../semanticAttributes';
2-
import type { Span, SpanAttributes, SpanOrigin } from '../types-hoist/span';
3+
import type { Event } from '../types-hoist/event';
4+
import type { Span, SpanAttributes, SpanJSON, SpanOrigin } from '../types-hoist/span';
35
import { spanToJSON } from './spanUtils';
46
import {
57
AI_MODEL_ID_ATTRIBUTE,
@@ -25,9 +27,9 @@ function addOriginToSpan(span: Span, origin: SpanOrigin): void {
2527

2628
/**
2729
* Post-process spans emitted by the Vercel AI SDK.
28-
* This is supposed to be used in `client.on('spanEnd', ...)`, to ensure all data is already finished.
30+
* This is supposed to be used in `client.on('spanStart', ...)
2931
*/
30-
export function processVercelAiSpan(span: Span): void {
32+
function onVercelAiSpanStart(span: Span): void {
3133
const { data: attributes, description: name } = spanToJSON(span);
3234

3335
if (!name) {
@@ -38,7 +40,6 @@ export function processVercelAiSpan(span: Span): void {
3840
// https://ai-sdk.dev/docs/ai-sdk-core/telemetry#tool-call-spans
3941
if (attributes[AI_TOOL_CALL_NAME_ATTRIBUTE] && attributes[AI_TOOL_CALL_ID_ATTRIBUTE] && name === 'ai.toolCall') {
4042
processToolCallSpan(span, attributes);
41-
sharedProcessSpan(span, attributes);
4243
return;
4344
}
4445

@@ -52,7 +53,47 @@ export function processVercelAiSpan(span: Span): void {
5253
}
5354

5455
processGenerateSpan(span, name, attributes);
55-
sharedProcessSpan(span, attributes);
56+
}
57+
58+
const vercelAiEventProcessor = Object.assign(
59+
(event: Event): Event => {
60+
if (event.type === 'transaction' && event.spans) {
61+
for (const span of event.spans) {
62+
// this mutates spans in-place
63+
processEndedVercelAiSpan(span);
64+
}
65+
}
66+
return event;
67+
},
68+
{ id: 'VercelAiEventProcessor' },
69+
);
70+
71+
/**
72+
* Post-process spans emitted by the Vercel AI SDK.
73+
*/
74+
function processEndedVercelAiSpan(span: SpanJSON): void {
75+
const { data: attributes, origin } = span;
76+
77+
if (origin !== 'auto.vercelai.otel') {
78+
return;
79+
}
80+
81+
renameAttributeKey(attributes, AI_USAGE_COMPLETION_TOKENS_ATTRIBUTE, GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE);
82+
renameAttributeKey(attributes, AI_USAGE_PROMPT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE);
83+
84+
if (
85+
typeof attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] === 'number' &&
86+
typeof attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] === 'number'
87+
) {
88+
attributes['gen_ai.usage.total_tokens'] =
89+
attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] + attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE];
90+
}
91+
92+
// Rename AI SDK attributes to standardized gen_ai attributes
93+
renameAttributeKey(attributes, AI_PROMPT_MESSAGES_ATTRIBUTE, 'gen_ai.request.messages');
94+
renameAttributeKey(attributes, AI_RESPONSE_TEXT_ATTRIBUTE, 'gen_ai.response.text');
95+
renameAttributeKey(attributes, AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, 'gen_ai.response.tool_calls');
96+
renameAttributeKey(attributes, AI_PROMPT_TOOLS_ATTRIBUTE, 'gen_ai.request.available_tools');
5697
}
5798

5899
/**
@@ -170,22 +211,11 @@ function processGenerateSpan(span: Span, name: string, attributes: SpanAttribute
170211
}
171212
}
172213

173-
// Processing for both tool call and non-tool call spans
174-
function sharedProcessSpan(span: Span, attributes: SpanAttributes): void {
175-
renameAttributeKey(attributes, AI_USAGE_COMPLETION_TOKENS_ATTRIBUTE, GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE);
176-
renameAttributeKey(attributes, AI_USAGE_PROMPT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE);
177-
178-
if (
179-
typeof attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] === 'number' &&
180-
typeof attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] === 'number'
181-
) {
182-
attributes['gen_ai.usage.total_tokens'] =
183-
attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] + attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE];
184-
}
185-
186-
// Rename AI SDK attributes to standardized gen_ai attributes
187-
renameAttributeKey(attributes, AI_PROMPT_MESSAGES_ATTRIBUTE, 'gen_ai.request.messages');
188-
renameAttributeKey(attributes, AI_RESPONSE_TEXT_ATTRIBUTE, 'gen_ai.response.text');
189-
renameAttributeKey(attributes, AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, 'gen_ai.response.tool_calls');
190-
renameAttributeKey(attributes, AI_PROMPT_TOOLS_ATTRIBUTE, 'gen_ai.request.available_tools');
214+
/**
215+
* Add event processors to the given client to process Vercel AI spans.
216+
*/
217+
export function addVercelAiProcessors(client: Client): void {
218+
client.on('spanStart', onVercelAiSpanStart);
219+
// Note: We cannot do this on `spanEnd`, because the span cannot be mutated anymore at this point
220+
client.addEventProcessor(vercelAiEventProcessor);
191221
}

‎packages/node/src/integrations/tracing/vercelai/index.ts

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { Client, IntegrationFn } from '@sentry/core';
2-
import { defineIntegration, processVercelAiSpan } from '@sentry/core';
2+
import { addVercelAiProcessors, defineIntegration } from '@sentry/core';
33
import { generateInstrumentOnce } from '../../../otel/instrument';
44
import type { modulesIntegration } from '../../modules';
55
import { INTEGRATION_NAME } from './constants';
@@ -27,18 +27,14 @@ const _vercelAIIntegration = ((options: VercelAiOptions = {}) => {
2727
instrumentation = instrumentVercelAi();
2828
},
2929
afterAllSetup(client) {
30-
function registerProcessors(): void {
31-
client.on('spanEnd', processVercelAiSpan);
32-
}
33-
3430
// Auto-detect if we should force the integration when running with 'ai' package available
3531
// Note that this can only be detected if the 'Modules' integration is available, and running in CJS mode
3632
const shouldForce = options.force ?? shouldForceIntegration(client);
3733

3834
if (shouldForce) {
39-
registerProcessors();
35+
addVercelAiProcessors(client);
4036
} else {
41-
instrumentation?.callWhenPatched(registerProcessors);
37+
instrumentation?.callWhenPatched(() => addVercelAiProcessors(client));
4238
}
4339
},
4440
};

0 commit comments

Comments
 (0)