fix(server-utils): Cap the runs the LangChain callback handler keeps open

packages/server-utils/src/ai/langchain/constants.ts

@@ -1,5 +1,8 @@
 export const LANGCHAIN_INTEGRATION_NAME = 'LangChain' as const;
 export const LANGCHAIN_ORIGIN = 'auto.ai.langchain';
+// LangChain skips the end callbacks when a consumer abandons a stream. An evicted run's span is dropped unended,
+// since ending it would send a span that lasted until eviction.
+export const MAX_TRACKED_LANGCHAIN_RUNS = 1000;
 
 export const ROLE_MAP: Record<string, string> = {
   human: 'user',

packages/server-utils/src/ai/langchain/index.ts

@@ -4,6 +4,7 @@ import {
   getClient,
   getRootSpan,
   hasSpanStreamingEnabled,
+  LRUMap,
   SPAN_STATUS_ERROR,
   spanIsIgnored,
   spanToJSON,
@@ -26,7 +27,7 @@ import {
 } from '@sentry/conventions/attributes';
 import { GEN_AI_CHAT, GEN_AI_EXECUTE_TOOL, GEN_AI_INVOKE_AGENT } from '@sentry/conventions/op';
 import { resolveAIRecordingOptions } from '../core/utils';
-import { LANGCHAIN_ORIGIN } from './constants';
+import { LANGCHAIN_ORIGIN, MAX_TRACKED_LANGCHAIN_RUNS } from './constants';
 import {
   addTypeSafeClassifierResponseAttributes,
   isTypeSafeClassifier,
@@ -68,25 +69,34 @@ export function createLangChainCallbackHandler(options: LangChainOptions = {}):
   const { recordInputs, recordOutputs } = resolveAIRecordingOptions(options);
 
   // Internal state - single instance tracks all spans
-  const runs = new Map<string, LangChainRun>();
+  const runs = new LRUMap<string, LangChainRun>(MAX_TRACKED_LANGCHAIN_RUNS);
   const getRun = (runId: string | undefined): LangChainRun | undefined => (runId ? runs.get(runId) : undefined);
 
   /**
    * Exit a span and clean up
    */
   const exitSpan = (runId: string): void => {
-    const span = runs.get(runId)?.span;
+    const span = runs.remove(runId)?.span;
     if (span?.isRecording()) {
       span.end();
     }
-    runs.delete(runId);
+  };
+
+  // Nearest first. `LRUMap.get()` marks each run recently used, so a long-running root run is not evicted while
+  // its subtree is busy.
+  const getAncestors = (parentRunId: string | undefined): LangChainRun[] => {
+    const ancestors: LangChainRun[] = [];
+    for (let run = getRun(parentRunId); run; run = getRun(run.parentRunId)) {
+      ancestors.push(run);
+    }
+    return ancestors;
   };
 
   // LangChain never makes our run spans active, so a run starts under its nearest ancestor run span. Like core,
   // an ignored span only parents when it is a root, which drops its subtree. The active span stays the parent
   // if the ancestor run has no span, or if another span was opened since that run started.
-  const getParentRunSpan = (parentRunId: string | undefined, activeSpan: Span | undefined): Span | undefined => {
-    for (let run = getRun(parentRunId); run; run = getRun(run.parentRunId)) {
+  const getParentRunSpan = (ancestors: LangChainRun[], activeSpan: Span | undefined): Span | undefined => {
+    for (const run of ancestors) {
       if (run.activeSpan !== activeSpan) {
         return undefined;
       }
@@ -98,27 +108,27 @@ export function createLangChainCallbackHandler(options: LangChainOptions = {}):
   };
 
   const startRunSpan = (runId: string, parentRunId: string | undefined, start: () => Span): LangChainRun => {
+    const ancestors = getAncestors(parentRunId);
     const activeSpan = getActiveSpan();
-    const parentSpan = getParentRunSpan(parentRunId, activeSpan);
+    const parentSpan = getParentRunSpan(ancestors, activeSpan);
     const span = parentSpan ? withActiveSpan(parentSpan, start) : start();
 
     // `@langchain/core` >= 1.1.40 keeps `configurable` out of run metadata, so often only the chat span knows
     // the conversation id. The own id is read off the span, since one from `Sentry.setConversationId()` is
     // only set there.
     const ownId = spanToJSON(span).attributes[GEN_AI_CONVERSATION_ID];
-    const conversationId = typeof ownId === 'string' ? ownId : getRun(parentRunId)?.conversationId;
+    const conversationId = typeof ownId === 'string' ? ownId : ancestors[0]?.conversationId;
     if (conversationId && conversationId !== ownId) {
       span.setAttribute(GEN_AI_CONVERSATION_ID, conversationId);
     }
 
     const run: LangChainRun = { span, parentRunId, activeSpan, conversationId };
     runs.set(runId, run);
     if (conversationId) {
-      for (
-        let ancestor = getRun(parentRunId);
-        ancestor && !ancestor.conversationId;
-        ancestor = getRun(ancestor.parentRunId)
-      ) {
+      for (const ancestor of ancestors) {
+        if (ancestor.conversationId) {
+          break;
+        }
         ancestor.conversationId = conversationId;
         ancestor.span.setAttribute(GEN_AI_CONVERSATION_ID, conversationId);
       }

packages/server-utils/test/ai/lib/tracing/langchain-invoke-agent.test.ts

@@ -20,6 +20,7 @@ import {
 } from '@sentry/core';
 import type { ClientOptions, Span } from '@sentry/core';
 import { createLangChainCallbackHandler } from '../../../../src/ai/langchain';
+import { MAX_TRACKED_LANGCHAIN_RUNS } from '../../../../src/ai/langchain/constants';
 import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client';
 
 beforeEach(() => {
@@ -320,6 +321,36 @@ describe('LangChain callback handler run tree', () => {
     });
   });
 
+  it('drops the oldest run without ending it once too many runs are open', () => {
+    const endedSpans = setupClient();
+    const handler = createLangChainCallbackHandler();
+
+    handler.handleChainStart?.({}, {}, 'abandoned', undefined, undefined, undefined, undefined, 'abandoned');
+    for (let i = 0; i < MAX_TRACKED_LANGCHAIN_RUNS; i++) {
+      handler.handleChainStart?.({}, {}, `run-${i}`, undefined, undefined, undefined, undefined, `run-${i}`);
+    }
+    handler.handleChainEnd?.({}, 'abandoned');
+    handler.handleChainEnd?.({}, 'run-0');
+
+    expect(endedSpans.map(span => spanToJSON(span).name)).toEqual(['chain run-0']);
+  });
+
+  it('keeps every ancestor of a busy run open', () => {
+    const endedSpans = setupClient();
+    const handler = createLangChainCallbackHandler();
+
+    handler.handleChainStart?.({}, {}, 'outer', undefined, undefined, undefined, undefined, 'outer');
+    handler.handleChainStart?.({}, {}, 'inner', 'outer', undefined, undefined, undefined, 'inner');
+    for (let i = 0; i < MAX_TRACKED_LANGCHAIN_RUNS; i++) {
+      handler.handleChatModelStart?.(chatModel, [[]], `chat-${i}`, 'inner');
+    }
+    handler.handleChainEnd?.({}, 'inner');
+    handler.handleChainEnd?.({}, 'outer');
+
+    const chainNames = endedSpans.map(span => spanToJSON(span).name).filter(name => name.startsWith('chain'));
+    expect(chainNames).toEqual(['chain inner', 'chain outer']);
+  });
+
   it('lets an id from Sentry.setConversationId() win over one from a child run', () => {
     const endedSpans = setupClient();
     getClient()!.addIntegration(conversationIdIntegration());