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

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

@@ -1,5 +1,7 @@
 export const LANGCHAIN_INTEGRATION_NAME = 'LangChain' as const;
 export const LANGCHAIN_ORIGIN = 'auto.ai.langchain';
+// Matches `MAX_TRACKED_FLUE_SPANS`: LangChain skips the end callbacks when a consumer abandons a stream.
+export const MAX_TRACKED_LANGCHAIN_RUNS = 1000;
 
 export const ROLE_MAP: Record<string, string> = {
   human: 'user',

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

@@ -1,9 +1,11 @@
 /* eslint-disable max-lines */
 import {
+  debug,
   getActiveSpan,
   getClient,
   getRootSpan,
   hasSpanStreamingEnabled,
+  LRUMap,
   SPAN_STATUS_ERROR,
   spanIsIgnored,
   spanToJSON,
@@ -25,8 +27,9 @@ import {
   SENTRY_ORIGIN,
 } from '@sentry/conventions/attributes';
 import { GEN_AI_CHAT, GEN_AI_EXECUTE_TOOL, GEN_AI_INVOKE_AGENT } from '@sentry/conventions/op';
+import { DEBUG_BUILD } from '../../debug-build';
 import { resolveAIRecordingOptions } from '../core/utils';
-import { LANGCHAIN_ORIGIN } from './constants';
+import { LANGCHAIN_ORIGIN, MAX_TRACKED_LANGCHAIN_RUNS } from './constants';
 import {
   addTypeSafeClassifierResponseAttributes,
   isTypeSafeClassifier,
@@ -68,25 +71,53 @@ 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);
+  };
+
+  // `LRUMap.set` would drop the oldest run without ending its span, so it is ended first, like Flue and Mastra.
+  let warnedFull = false;
+  const trackRun = (runId: string, run: LangChainRun): void => {
+    if (runs.size >= MAX_TRACKED_LANGCHAIN_RUNS) {
+      if (!warnedFull) {
+        warnedFull = true;
+        DEBUG_BUILD &&
+          debug.warn(
+            `[LangChain] More than ${MAX_TRACKED_LANGCHAIN_RUNS} runs are open, ending the oldest. Its span may be cut short.`,
+          );
+      }
+      const oldestRunId = runs.keys()[0];
+      if (oldestRunId !== undefined) {
+        exitSpan(oldestRunId);
+      }
+    }
+    runs.set(runId, run);
+  };
+
+  // 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 +129,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);
+    trackRun(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

@@ -1,4 +1,4 @@
-import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import {
   GEN_AI_AGENT_NAME,
   GEN_AI_CONVERSATION_ID,
@@ -10,6 +10,7 @@ import {
 import { GEN_AI_DECIDE } from '@sentry/conventions/op';
 import {
   conversationIdIntegration,
+  debug,
   getClient,
   getMainCarrier,
   setConversationId,
@@ -20,6 +21,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 +322,39 @@ describe('LangChain callback handler run tree', () => {
     });
   });
 
+  it('ends the oldest run once too many runs are open, and warns once', () => {
+    const endedSpans = setupClient();
+    const warn = vi.spyOn(debug, 'warn').mockImplementation(() => undefined);
+    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}`);
+    }
+    expect(endedSpans.map(span => spanToJSON(span).name)).toEqual(['chain abandoned', 'chain run-0']);
+    expect(warn).toHaveBeenCalledTimes(1);
+
+    handler.handleChainEnd?.({}, 'abandoned');
+    expect(endedSpans).toHaveLength(2);
+    warn.mockRestore();
+  });
+
+  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());