Skip to content

Commit b47e911

Browse files
lithdewlithdew
committed
fix(vercel-ai): prevent tool call span map memory leak
Tool calls were stored in a global map and only cleaned up on tool errors, causing unbounded retention in tool-heavy apps (and potential OOMs when inputs/outputs were recorded). Store only span context in a bounded LRU cache and clean up on successful tool results; add tests for caching/eviction.
1 parent 6fb1ee1 commit b47e911

File tree

6 files changed

+214
-65
lines changed

6 files changed

+214
-65
lines changed

‎packages/core/src/tracing/vercel-ai/constants.ts‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
1-
import type { Span } from '../../types-hoist/span';
1+
import { LRUMap } from '../../utils/lru';
2+
import type { ToolCallSpanContext } from './types';
23

3-
// Global Map to track tool call IDs to their corresponding spans
4+
export const TOOL_CALL_SPAN_MAP_MAX_SIZE = 10_000;
5+
6+
// Global LRU map to track tool call IDs to their corresponding span contexts.
47
// This allows us to capture tool errors and link them to the correct span
5-
export const toolCallSpanMap = new Map<string, Span>();
8+
// without keeping full Span objects (and their potentially large attributes) alive.
9+
export const toolCallSpanMap = new LRUMap<string, ToolCallSpanContext>(TOOL_CALL_SPAN_MAP_MAX_SIZE);
610

711
// Operation sets for efficient mapping to OpenTelemetry semantic convention values
812
export const INVOKE_AGENT_OPS = new Set([

‎packages/core/src/tracing/vercel-ai/index.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -232,12 +232,13 @@ function processToolCallSpan(span: Span, attributes: SpanAttributes): void {
232232
renameAttributeKey(attributes, AI_TOOL_CALL_NAME_ATTRIBUTE, GEN_AI_TOOL_NAME_ATTRIBUTE);
233233
renameAttributeKey(attributes, AI_TOOL_CALL_ID_ATTRIBUTE, GEN_AI_TOOL_CALL_ID_ATTRIBUTE);
234234

235-
// Store the span in our global map using the tool call ID
235+
// Store the span context in our global map using the tool call ID.
236236
// This allows us to capture tool errors and link them to the correct span
237+
// without retaining the full Span object in memory.
237238
const toolCallId = attributes[GEN_AI_TOOL_CALL_ID_ATTRIBUTE];
238239

239240
if (typeof toolCallId === 'string') {
240-
toolCallSpanMap.set(toolCallId, span);
241+
toolCallSpanMap.set(toolCallId, span.spanContext());
241242
}
242243

243244
// https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#gen-ai-tool-type

‎packages/core/src/tracing/vercel-ai/types.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,8 @@ export interface TokenSummary {
22
inputTokens: number;
33
outputTokens: number;
44
}
5+
6+
export interface ToolCallSpanContext {
7+
traceId: string;
8+
spanId: string;
9+
}

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import {
1818
} from '../ai/gen-ai-attributes';
1919
import { extractSystemInstructions, getTruncatedJsonString } from '../ai/utils';
2020
import { toolCallSpanMap } from './constants';
21-
import type { TokenSummary } from './types';
21+
import type { TokenSummary, ToolCallSpanContext } from './types';
2222
import { AI_PROMPT_ATTRIBUTE, AI_PROMPT_MESSAGES_ATTRIBUTE } from './vercel-ai-attributes';
2323

2424
/**
@@ -75,17 +75,17 @@ export function applyAccumulatedTokens(
7575
}
7676

7777
/**
78-
* Get the span associated with a tool call ID
78+
* Get the span context associated with a tool call ID.
7979
*/
80-
export function _INTERNAL_getSpanForToolCallId(toolCallId: string): Span | undefined {
80+
export function _INTERNAL_getSpanForToolCallId(toolCallId: string): ToolCallSpanContext | undefined {
8181
return toolCallSpanMap.get(toolCallId);
8282
}
8383

8484
/**
8585
* Clean up the span mapping for a tool call ID
8686
*/
8787
export function _INTERNAL_cleanupToolCallSpan(toolCallId: string): void {
88-
toolCallSpanMap.delete(toolCallId);
88+
toolCallSpanMap.remove(toolCallId);
8989
}
9090

9191
/**

‎packages/core/test/lib/tracing/vercel-ai-tool-call-span-map.test.ts‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
import { beforeEach, describe, expect, it } from 'vitest';
2+
import { addVercelAiProcessors } from '../../../src/tracing/vercel-ai';
3+
import { TOOL_CALL_SPAN_MAP_MAX_SIZE, toolCallSpanMap } from '../../../src/tracing/vercel-ai/constants';
4+
import { _INTERNAL_cleanupToolCallSpan, _INTERNAL_getSpanForToolCallId } from '../../../src/tracing/vercel-ai/utils';
5+
import {
6+
AI_TOOL_CALL_ID_ATTRIBUTE,
7+
AI_TOOL_CALL_NAME_ATTRIBUTE,
8+
} from '../../../src/tracing/vercel-ai/vercel-ai-attributes';
9+
import type { SpanAttributes, SpanAttributeValue, SpanTimeInput } from '../../../src/types-hoist/span';
10+
import type { SpanStatus } from '../../../src/types-hoist/spanStatus';
11+
import type { OpenTelemetrySdkTraceBaseSpan } from '../../../src/utils/spanUtils';
12+
import { getDefaultTestClientOptions, TestClient } from '../../mocks/client';
13+
14+
function createToolCallSpan(params: {
15+
toolCallId: string;
16+
toolName: string;
17+
traceId: string;
18+
spanId: string;
19+
}): OpenTelemetrySdkTraceBaseSpan {
20+
const attributes: SpanAttributes = {
21+
[AI_TOOL_CALL_ID_ATTRIBUTE]: params.toolCallId,
22+
[AI_TOOL_CALL_NAME_ATTRIBUTE]: params.toolName,
23+
};
24+
25+
const startTime: SpanTimeInput = [0, 0];
26+
const endTime: SpanTimeInput = [0, 0];
27+
const status: SpanStatus = { code: 0 };
28+
29+
const span: OpenTelemetrySdkTraceBaseSpan = {
30+
attributes,
31+
startTime,
32+
endTime,
33+
name: 'ai.toolCall',
34+
status,
35+
spanContext: () => ({
36+
traceId: params.traceId,
37+
spanId: params.spanId,
38+
traceFlags: 1,
39+
}),
40+
end: () => undefined,
41+
setAttribute: (key: string, value: SpanAttributeValue | undefined) => {
42+
if (value === undefined) {
43+
// eslint-disable-next-line @typescript-eslint/no-dynamic-delete
44+
delete attributes[key];
45+
} else {
46+
attributes[key] = value;
47+
}
48+
return span;
49+
},
50+
setAttributes: (nextAttributes: SpanAttributes) => {
51+
for (const key of Object.keys(nextAttributes)) {
52+
const value = nextAttributes[key];
53+
if (value === undefined) {
54+
// eslint-disable-next-line @typescript-eslint/no-dynamic-delete
55+
delete attributes[key];
56+
} else {
57+
attributes[key] = value;
58+
}
59+
}
60+
return span;
61+
},
62+
setStatus: (nextStatus: SpanStatus) => {
63+
span.status = nextStatus;
64+
return span;
65+
},
66+
updateName: (name: string) => {
67+
span.name = name;
68+
return span;
69+
},
70+
isRecording: () => true,
71+
addEvent: () => span,
72+
addLink: () => span,
73+
addLinks: () => span,
74+
recordException: () => undefined,
75+
};
76+
77+
return span;
78+
}
79+
80+
describe('vercel-ai tool call span context map', () => {
81+
beforeEach(() => {
82+
toolCallSpanMap.clear();
83+
});
84+
85+
it('stores toolCallId -> span context on spanStart', () => {
86+
const options = getDefaultTestClientOptions({ tracesSampleRate: 1.0 });
87+
const client = new TestClient(options);
88+
client.init();
89+
addVercelAiProcessors(client);
90+
91+
const span = createToolCallSpan({
92+
toolCallId: 'tool-call-1',
93+
toolName: 'bash',
94+
traceId: 'trace-id-1',
95+
spanId: 'span-id-1',
96+
});
97+
98+
client.emit('spanStart', span);
99+
100+
expect(_INTERNAL_getSpanForToolCallId('tool-call-1')).toMatchObject({
101+
traceId: 'trace-id-1',
102+
spanId: 'span-id-1',
103+
});
104+
105+
_INTERNAL_cleanupToolCallSpan('tool-call-1');
106+
expect(_INTERNAL_getSpanForToolCallId('tool-call-1')).toBeUndefined();
107+
});
108+
109+
it('evicts old entries when the map exceeds max size', () => {
110+
for (let i = 0; i < TOOL_CALL_SPAN_MAP_MAX_SIZE + 1; i++) {
111+
toolCallSpanMap.set(`tool-call-${i}`, { traceId: `trace-${i}`, spanId: `span-${i}` });
112+
}
113+
114+
expect(toolCallSpanMap.size).toBe(TOOL_CALL_SPAN_MAP_MAX_SIZE);
115+
expect(toolCallSpanMap.get('tool-call-0')).toBeUndefined();
116+
expect(toolCallSpanMap.get(`tool-call-${TOOL_CALL_SPAN_MAP_MAX_SIZE}`)).toEqual({
117+
traceId: `trace-${TOOL_CALL_SPAN_MAP_MAX_SIZE}`,
118+
spanId: `span-${TOOL_CALL_SPAN_MAP_MAX_SIZE}`,
119+
});
120+
});
121+
});

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

Lines changed: 74 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import type { InstrumentationConfig, InstrumentationModuleDefinition } from '@opentelemetry/instrumentation';
22
import { InstrumentationBase, InstrumentationNodeModuleDefinition } from '@opentelemetry/instrumentation';
3-
import type { Span } from '@sentry/core';
43
import {
54
_INTERNAL_cleanupToolCallSpan,
65
_INTERNAL_getSpanForToolCallId,
@@ -43,33 +42,46 @@ interface RecordingOptions {
4342
recordOutputs?: boolean;
4443
}
4544

46-
interface ToolError {
47-
type: 'tool-error' | 'tool-result' | 'tool-call';
45+
interface ToolErrorPart {
46+
type: 'tool-error';
4847
toolCallId: string;
4948
toolName: string;
50-
input?: {
51-
[key: string]: unknown;
52-
};
5349
error: Error;
54-
dynamic?: boolean;
5550
}
5651

57-
function isToolError(obj: unknown): obj is ToolError {
52+
interface ToolResultPart {
53+
type: 'tool-result';
54+
toolCallId: string;
55+
toolName: string;
56+
}
57+
58+
function isToolErrorPart(obj: unknown): obj is ToolErrorPart {
5859
if (typeof obj !== 'object' || obj === null) {
5960
return false;
6061
}
6162

6263
const candidate = obj as Record<string, unknown>;
6364
return (
64-
'type' in candidate &&
65-
'error' in candidate &&
66-
'toolName' in candidate &&
67-
'toolCallId' in candidate &&
6865
candidate.type === 'tool-error' &&
66+
typeof candidate.toolName === 'string' &&
67+
typeof candidate.toolCallId === 'string' &&
6968
candidate.error instanceof Error
7069
);
7170
}
7271

72+
function isToolResultPart(obj: unknown): obj is ToolResultPart {
73+
if (typeof obj !== 'object' || obj === null) {
74+
return false;
75+
}
76+
77+
const candidate = obj as Record<string, unknown>;
78+
return (
79+
candidate.type === 'tool-result' &&
80+
typeof candidate.toolName === 'string' &&
81+
typeof candidate.toolCallId === 'string'
82+
);
83+
}
84+
7385
/**
7486
* Check for tool errors in the result and capture them
7587
* Tool errors are not rejected in Vercel V5, it is added as metadata to the result content
@@ -79,59 +91,65 @@ function checkResultForToolErrors(result: unknown): void {
7991
return;
8092
}
8193

82-
const resultObj = result as { content: Array<object> };
94+
const resultObj = result as { content: unknown };
8395
if (!Array.isArray(resultObj.content)) {
8496
return;
8597
}
8698

8799
for (const item of resultObj.content) {
88-
if (isToolError(item)) {
89-
// Try to get the span associated with this tool call ID
90-
const associatedSpan = _INTERNAL_getSpanForToolCallId(item.toolCallId) as Span;
91-
92-
if (associatedSpan) {
93-
// We have the span, so link the error using span and trace IDs from the span
94-
const spanContext = associatedSpan.spanContext();
95-
96-
withScope(scope => {
97-
// Set the span and trace context for proper linking
98-
scope.setContext('trace', {
99-
trace_id: spanContext.traceId,
100-
span_id: spanContext.spanId,
101-
});
102-
103-
scope.setTag('vercel.ai.tool.name', item.toolName);
104-
scope.setTag('vercel.ai.tool.callId', item.toolCallId);
105-
106-
scope.setLevel('error');
107-
108-
captureException(item.error, {
109-
mechanism: {
110-
type: 'auto.vercelai.otel',
111-
handled: false,
112-
},
113-
});
100+
// Successful tool calls should not keep toolCallId -> span context mappings alive.
101+
if (isToolResultPart(item)) {
102+
_INTERNAL_cleanupToolCallSpan(item.toolCallId);
103+
continue;
104+
}
105+
106+
if (!isToolErrorPart(item)) {
107+
continue;
108+
}
109+
110+
// Try to get the span context associated with this tool call ID
111+
const spanContext = _INTERNAL_getSpanForToolCallId(item.toolCallId);
112+
113+
if (spanContext) {
114+
// We have a span context, so link the error using span and trace IDs from the span
115+
withScope(scope => {
116+
// Set the span and trace context for proper linking
117+
scope.setContext('trace', {
118+
trace_id: spanContext.traceId,
119+
span_id: spanContext.spanId,
114120
});
115121

116-
// Clean up the span mapping since we've processed this tool error
117-
// We won't get multiple { type: 'tool-error' } parts for the same toolCallId.
118-
_INTERNAL_cleanupToolCallSpan(item.toolCallId);
119-
} else {
120-
// Fallback: capture without span linking
121-
withScope(scope => {
122-
scope.setTag('vercel.ai.tool.name', item.toolName);
123-
scope.setTag('vercel.ai.tool.callId', item.toolCallId);
124-
scope.setLevel('error');
125-
126-
captureException(item.error, {
127-
mechanism: {
128-
type: 'auto.vercelai.otel',
129-
handled: false,
130-
},
131-
});
122+
scope.setTag('vercel.ai.tool.name', item.toolName);
123+
scope.setTag('vercel.ai.tool.callId', item.toolCallId);
124+
125+
scope.setLevel('error');
126+
127+
captureException(item.error, {
128+
mechanism: {
129+
type: 'auto.vercelai.otel',
130+
handled: false,
131+
},
132132
});
133-
}
133+
});
134+
} else {
135+
// Fallback: capture without span linking
136+
withScope(scope => {
137+
scope.setTag('vercel.ai.tool.name', item.toolName);
138+
scope.setTag('vercel.ai.tool.callId', item.toolCallId);
139+
scope.setLevel('error');
140+
141+
captureException(item.error, {
142+
mechanism: {
143+
type: 'auto.vercelai.otel',
144+
handled: false,
145+
},
146+
});
147+
});
134148
}
149+
150+
// Clean up the span mapping since we've processed this tool error
151+
// We won't get multiple { type: 'tool-error' } parts for the same toolCallId.
152+
_INTERNAL_cleanupToolCallSpan(item.toolCallId);
135153
}
136154
}
137155

0 commit comments

Comments
 (0)