fix(vercel-ai): prevent tool call span map memory leak

packages/core/src/index.ts

@@ -146,7 +146,8 @@ export * as metrics from './metrics/public-api';
 export type { MetricOptions } from './metrics/public-api';
 export { createConsolaReporter } from './integrations/consola';
 export { addVercelAiProcessors } from './tracing/vercel-ai';
-export { _INTERNAL_getSpanForToolCallId, _INTERNAL_cleanupToolCallSpan } from './tracing/vercel-ai/utils';
+export { _INTERNAL_getSpanContextForToolCallId, _INTERNAL_cleanupToolCallSpanContext } from './tracing/vercel-ai/utils';
+export { toolCallSpanContextMap as _INTERNAL_toolCallSpanContextMap } from './tracing/vercel-ai/constants';
 export { instrumentOpenAiClient } from './tracing/openai';
 export { OPENAI_INTEGRATION_NAME } from './tracing/openai/constants';
 export { instrumentAnthropicAiClient } from './tracing/anthropic-ai';

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

@@ -1,8 +1,9 @@
-import type { Span } from '../../types-hoist/span';
+import type { ToolCallSpanContext } from './types';
 
-// Global Map to track tool call IDs to their corresponding spans
+// Global map to track tool call IDs to their corresponding span contexts.
 // This allows us to capture tool errors and link them to the correct span
-export const toolCallSpanMap = new Map<string, Span>();
+// without keeping full Span objects (and their potentially large attributes) alive.
+export const toolCallSpanContextMap = new Map<string, ToolCallSpanContext>();
 
 // Operation sets for efficient mapping to OpenTelemetry semantic convention values
 export const INVOKE_AGENT_OPS = new Set([

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

@@ -1,3 +1,4 @@
+/* eslint-disable max-lines */
 import type { Client } from '../../client';
 import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '../../semanticAttributes';
 import type { Event } from '../../types-hoist/event';
@@ -19,7 +20,13 @@ import {
   GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
   GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
-import { EMBEDDINGS_OPS, GENERATE_CONTENT_OPS, INVOKE_AGENT_OPS, RERANK_OPS, toolCallSpanMap } from './constants';
+import {
+  EMBEDDINGS_OPS,
+  GENERATE_CONTENT_OPS,
+  INVOKE_AGENT_OPS,
+  RERANK_OPS,
+  toolCallSpanContextMap,
+} from './constants';
 import type { TokenSummary } from './types';
 import {
   accumulateTokensForParent,
@@ -232,12 +239,13 @@ function processToolCallSpan(span: Span, attributes: SpanAttributes): void {
   renameAttributeKey(attributes, AI_TOOL_CALL_NAME_ATTRIBUTE, GEN_AI_TOOL_NAME_ATTRIBUTE);
   renameAttributeKey(attributes, AI_TOOL_CALL_ID_ATTRIBUTE, GEN_AI_TOOL_CALL_ID_ATTRIBUTE);
 
-  // Store the span in our global map using the tool call ID
+  // Store the span context in our global map using the tool call ID.
   // This allows us to capture tool errors and link them to the correct span
+  // without retaining the full Span object in memory.
   const toolCallId = attributes[GEN_AI_TOOL_CALL_ID_ATTRIBUTE];
 
   if (typeof toolCallId === 'string') {
-    toolCallSpanMap.set(toolCallId, span);
+    toolCallSpanContextMap.set(toolCallId, span.spanContext());
   }
 
   // https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#gen-ai-tool-type

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

@@ -2,3 +2,8 @@ export interface TokenSummary {
   inputTokens: number;
   outputTokens: number;
 }
+
+export interface ToolCallSpanContext {
+  traceId: string;
+  spanId: string;
+}

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

@@ -17,8 +17,8 @@ import {
   GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
 import { extractSystemInstructions, getTruncatedJsonString } from '../ai/utils';
-import { toolCallSpanMap } from './constants';
-import type { TokenSummary } from './types';
+import { toolCallSpanContextMap } from './constants';
+import type { TokenSummary, ToolCallSpanContext } from './types';
 import { AI_PROMPT_ATTRIBUTE, AI_PROMPT_MESSAGES_ATTRIBUTE } from './vercel-ai-attributes';
 
 /**
@@ -75,17 +75,17 @@ export function applyAccumulatedTokens(
 }
 
 /**
- * Get the span associated with a tool call ID
+ * Get the span context associated with a tool call ID.
  */
-export function _INTERNAL_getSpanForToolCallId(toolCallId: string): Span | undefined {
-  return toolCallSpanMap.get(toolCallId);
+export function _INTERNAL_getSpanContextForToolCallId(toolCallId: string): ToolCallSpanContext | undefined {
+  return toolCallSpanContextMap.get(toolCallId);
 }
 
 /**
  * Clean up the span mapping for a tool call ID
  */
-export function _INTERNAL_cleanupToolCallSpan(toolCallId: string): void {
-  toolCallSpanMap.delete(toolCallId);
+export function _INTERNAL_cleanupToolCallSpanContext(toolCallId: string): void {
+  toolCallSpanContextMap.delete(toolCallId);
 }
 
 /**

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

@@ -1,9 +1,8 @@
 import type { InstrumentationConfig, InstrumentationModuleDefinition } from '@opentelemetry/instrumentation';
 import { InstrumentationBase, InstrumentationNodeModuleDefinition } from '@opentelemetry/instrumentation';
-import type { Span } from '@sentry/core';
 import {
-  _INTERNAL_cleanupToolCallSpan,
-  _INTERNAL_getSpanForToolCallId,
+  _INTERNAL_cleanupToolCallSpanContext,
+  _INTERNAL_getSpanContextForToolCallId,
   addNonEnumerableProperty,
   captureException,
   getActiveSpan,
@@ -71,10 +70,12 @@ function isToolError(obj: unknown): obj is ToolError {
 }
 
 /**
- * Check for tool errors in the result and capture them
- * Tool errors are not rejected in Vercel V5, it is added as metadata to the result content
+ * Process tool call results: capture tool errors and clean up span context mappings.
+ *
+ * Error checking runs first (needs span context for linking), then cleanup removes all entries.
+ * Tool errors are not rejected in Vercel AI V5 — they appear as metadata in the result content.
  */
-function checkResultForToolErrors(result: unknown): void {
+export function processToolCallResults(result: unknown): void {
   if (typeof result !== 'object' || result === null || !('content' in result)) {
     return;
   }
@@ -84,53 +85,68 @@ function checkResultForToolErrors(result: unknown): void {
     return;
   }
 
-  for (const item of resultObj.content) {
-    if (isToolError(item)) {
-      // Try to get the span associated with this tool call ID
-      const associatedSpan = _INTERNAL_getSpanForToolCallId(item.toolCallId) as Span;
+  captureToolErrors(resultObj.content);
+  cleanupToolCallSpanContexts(resultObj.content);
+}
 
-      if (associatedSpan) {
-        // We have the span, so link the error using span and trace IDs from the span
-        const spanContext = associatedSpan.spanContext();
+function captureToolErrors(content: Array<object>): void {
+  for (const item of content) {
+    if (!isToolError(item)) {
+      continue;
+    }
 
-        withScope(scope => {
-          // Set the span and trace context for proper linking
-          scope.setContext('trace', {
-            trace_id: spanContext.traceId,
-            span_id: spanContext.spanId,
-          });
+    // Try to get the span context associated with this tool call ID
+    const spanContext = _INTERNAL_getSpanContextForToolCallId(item.toolCallId);
 
-          scope.setTag('vercel.ai.tool.name', item.toolName);
-          scope.setTag('vercel.ai.tool.callId', item.toolCallId);
+    if (spanContext) {
+      // We have the span context, so link the error using span and trace IDs
+      withScope(scope => {
+        scope.setContext('trace', {
+          trace_id: spanContext.traceId,
+          span_id: spanContext.spanId,
+        });
 
-          scope.setLevel('error');
+        scope.setTag('vercel.ai.tool.name', item.toolName);
+        scope.setTag('vercel.ai.tool.callId', item.toolCallId);
+        scope.setLevel('error');
 
-          captureException(item.error, {
-            mechanism: {
-              type: 'auto.vercelai.otel',
-              handled: false,
-            },
-          });
+        captureException(item.error, {
+          mechanism: {
+            type: 'auto.vercelai.otel',
+            handled: false,
+          },
         });
-
-        // Clean up the span mapping since we've processed this tool error
-        // We won't get multiple { type: 'tool-error' } parts for the same toolCallId.
-        _INTERNAL_cleanupToolCallSpan(item.toolCallId);
-      } else {
-        // Fallback: capture without span linking
-        withScope(scope => {
-          scope.setTag('vercel.ai.tool.name', item.toolName);
-          scope.setTag('vercel.ai.tool.callId', item.toolCallId);
-          scope.setLevel('error');
-
-          captureException(item.error, {
-            mechanism: {
-              type: 'auto.vercelai.otel',
-              handled: false,
-            },
-          });
+      });
+    } else {
+      // Fallback: capture without span linking
+      withScope(scope => {
+        scope.setTag('vercel.ai.tool.name', item.toolName);
+        scope.setTag('vercel.ai.tool.callId', item.toolCallId);
+        scope.setLevel('error');
+
+        captureException(item.error, {
+          mechanism: {
+            type: 'auto.vercelai.otel',
+            handled: false,
+          },
         });
-      }
+      });
+    }
+  }
+}
+
+/**
+ * Remove span context entries for all completed tool calls in the content array.
+ */
+export function cleanupToolCallSpanContexts(content: Array<object>): void {
+  for (const item of content) {
+    if (
+      typeof item === 'object' &&
+      item !== null &&
+      'toolCallId' in item &&
+      typeof (item as Record<string, unknown>).toolCallId === 'string'
+    ) {
+      _INTERNAL_cleanupToolCallSpanContext((item as Record<string, unknown>).toolCallId as string);
     }
   }
 }
@@ -252,7 +268,7 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
             },
             () => {},
             result => {
-              checkResultForToolErrors(result);
+              processToolCallResults(result);
             },
           );
         },

packages/node/test/integrations/tracing/vercelai/instrumentation.test.ts

@@ -1,5 +1,9 @@
-import { describe, expect, test } from 'vitest';
-import { determineRecordingSettings } from '../../../../src/integrations/tracing/vercelai/instrumentation';
+import { _INTERNAL_getSpanContextForToolCallId, _INTERNAL_toolCallSpanContextMap } from '@sentry/core';
+import { beforeEach, describe, expect, test } from 'vitest';
+import {
+  cleanupToolCallSpanContexts,
+  determineRecordingSettings,
+} from '../../../../src/integrations/tracing/vercelai/instrumentation';
 
 describe('determineRecordingSettings', () => {
   test('should use integration recording options when provided (recordInputs: true, recordOutputs: false)', () => {
@@ -212,3 +216,50 @@ describe('determineRecordingSettings', () => {
     });
   });
 });
+
+describe('cleanupToolCallSpanContexts', () => {
+  beforeEach(() => {
+    _INTERNAL_toolCallSpanContextMap.clear();
+  });
+
+  test('cleans up span context for tool-result items', () => {
+    _INTERNAL_toolCallSpanContextMap.set('tool-1', { traceId: 't1', spanId: 's1' });
+    _INTERNAL_toolCallSpanContextMap.set('tool-2', { traceId: 't2', spanId: 's2' });
+
+    cleanupToolCallSpanContexts([{ type: 'tool-result', toolCallId: 'tool-1', toolName: 'bash' }]);
+
+    expect(_INTERNAL_getSpanContextForToolCallId('tool-1')).toBeUndefined();
+    expect(_INTERNAL_getSpanContextForToolCallId('tool-2')).toEqual({ traceId: 't2', spanId: 's2' });
+  });
+
+  test('cleans up span context for tool-error items', () => {
+    _INTERNAL_toolCallSpanContextMap.set('tool-1', { traceId: 't1', spanId: 's1' });
+
+    cleanupToolCallSpanContexts([
+      { type: 'tool-error', toolCallId: 'tool-1', toolName: 'bash', error: new Error('fail') },
+    ]);
+
+    expect(_INTERNAL_getSpanContextForToolCallId('tool-1')).toBeUndefined();
+  });
+
+  test('cleans up mixed tool-result and tool-error in same content array', () => {
+    _INTERNAL_toolCallSpanContextMap.set('tool-1', { traceId: 't1', spanId: 's1' });
+    _INTERNAL_toolCallSpanContextMap.set('tool-2', { traceId: 't2', spanId: 's2' });
+
+    cleanupToolCallSpanContexts([
+      { type: 'tool-result', toolCallId: 'tool-1', toolName: 'bash' },
+      { type: 'tool-error', toolCallId: 'tool-2', toolName: 'bash', error: new Error('fail') },
+    ]);
+
+    expect(_INTERNAL_getSpanContextForToolCallId('tool-1')).toBeUndefined();
+    expect(_INTERNAL_getSpanContextForToolCallId('tool-2')).toBeUndefined();
+  });
+
+  test('ignores items without toolCallId', () => {
+    _INTERNAL_toolCallSpanContextMap.set('tool-1', { traceId: 't1', spanId: 's1' });
+
+    cleanupToolCallSpanContexts([{ type: 'text', text: 'hello' } as unknown as object]);
+
+    expect(_INTERNAL_getSpanContextForToolCallId('tool-1')).toEqual({ traceId: 't1', spanId: 's1' });
+  });
+});