fix(langgraph): dispose unused combined signals (#1404)

.changeset/olive-bottles-study.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+fix(langgraph): dispose unused combined signals

libs/langgraph/src/pregel/index.ts

@@ -1840,9 +1840,8 @@ export class Pregel<
     const config = {
       recursionLimit: this.config?.recursionLimit,
       ...options,
-      signal: options?.signal
-        ? combineAbortSignals(options.signal, abortController.signal)
-        : abortController.signal,
+      signal: combineAbortSignals(options?.signal, abortController.signal)
+        .signal,
     };
 
     return new IterableReadableStreamWithAbortSignal(
@@ -1896,9 +1895,8 @@ export class Pregel<
 
       // extend the callbacks with the ones from the config
       callbacks: combineCallbacks(this.config?.callbacks, options?.callbacks),
-      signal: options?.signal
-        ? combineAbortSignals(options.signal, abortController.signal)
-        : abortController.signal,
+      signal: combineAbortSignals(options?.signal, abortController.signal)
+        .signal,
     };
 
     return new IterableReadableStreamWithAbortSignal(

libs/langgraph/src/pregel/runner.ts

@@ -111,21 +111,26 @@ export class PregelRunner {
 
     const nodeErrors: Set<Error> = new Set();
     let graphBubbleUp: GraphBubbleUp | undefined;
+
     const exceptionSignalController = new AbortController();
+    const exceptionSignal = exceptionSignalController.signal;
+    const stepTimeoutSignal = timeout
+      ? AbortSignal.timeout(timeout)
+      : undefined;
 
     // Start task execution
     const pendingTasks = Object.values(this.loop.tasks).filter(
       (t) => t.writes.length === 0
     );
 
-    const currentSignals = this._initializeAbortSignals({
-      exceptionSignalController,
-      timeout,
+    const { signals, disposeCombinedSignal } = this._initializeAbortSignals({
+      exceptionSignal,
+      stepTimeoutSignal,
       signal: options.signal,
     });
 
     const taskStream = this._executeTasksWithRetry(pendingTasks, {
-      signals: currentSignals,
+      signals,
       retryPolicy,
       maxConcurrency,
     });
@@ -156,6 +161,8 @@ export class PregelRunner {
       }
     }
 
+    disposeCombinedSignal?.();
+
     onStepWrite?.(
       this.loop.step,
       Object.values(this.loop.tasks)
@@ -195,65 +202,44 @@ export class PregelRunner {
    * @internal
    */
   private _initializeAbortSignals({
-    exceptionSignalController,
-    timeout,
+    exceptionSignal,
+    stepTimeoutSignal,
     signal,
   }: {
-    exceptionSignalController: AbortController;
-    timeout?: number;
+    exceptionSignal: AbortSignal;
+    stepTimeoutSignal?: AbortSignal;
     signal?: AbortSignal;
-  }): PregelAbortSignals {
-    const previousSignals: PregelAbortSignals =
-      (this.loop.config.configurable?.[
-        CONFIG_KEY_ABORT_SIGNALS
-      ] as PregelAbortSignals) ?? {};
+  }): { signals: PregelAbortSignals; disposeCombinedSignal?: () => void } {
+    const previousSignals = (this.loop.config.configurable?.[
+      CONFIG_KEY_ABORT_SIGNALS
+    ] ?? {}) as PregelAbortSignals;
 
-    // This is true when a node calls a subgraph and, rather than forwarding its own AbortSignal,
-    // it creates a new AbortSignal and passes that along instead.
-    const subgraphCalledWithSignalCreatedByNode =
-      signal &&
-      previousSignals.composedAbortSignal &&
-      signal !== previousSignals.composedAbortSignal;
+    // We always inherit the external abort signal from AsyncLocalStorage,
+    // since that's the only way the signal is inherited by the subgraph calls.
+    const externalAbortSignal = previousSignals.externalAbortSignal ?? signal;
 
-    const externalAbortSignal = subgraphCalledWithSignalCreatedByNode
-      ? // Chain the signals here to make sure that the subgraph receives the external abort signal in
-        // addition to the signal created by the node.
-        combineAbortSignals(previousSignals.externalAbortSignal!, signal!)
-      : // Otherwise, just keep using the external abort signal, or initialize it if it hasn't been
-        // assigned yet
-        previousSignals.externalAbortSignal ?? signal;
+    // inherit the step timeout signal from parent graph
+    const timeoutAbortSignal =
+      stepTimeoutSignal ?? previousSignals.timeoutAbortSignal;
 
-    const errorAbortSignal = previousSignals.errorAbortSignal
-      ? // Chaining here rather than always using a fresh one handles the case where a subgraph is
-        // called in a parallel branch to some other node in the parent graph.
-        combineAbortSignals(
-          previousSignals.errorAbortSignal!,
-          exceptionSignalController.signal
-        )
-      : exceptionSignalController.signal;
+    const { signal: composedAbortSignal, dispose: disposeCombinedSignal } =
+      combineAbortSignals(
+        externalAbortSignal,
+        timeoutAbortSignal,
+        exceptionSignal
+      );
 
-    const timeoutAbortSignal = timeout
-      ? AbortSignal.timeout(timeout)
-      : undefined;
-
-    const composedAbortSignal: AbortSignal = combineAbortSignals(
-      ...(externalAbortSignal ? [externalAbortSignal] : []),
-      ...(timeoutAbortSignal ? [timeoutAbortSignal] : []),
-      errorAbortSignal
-    );
-
-    const currentSignals: PregelAbortSignals = {
+    const signals: PregelAbortSignals = {
       externalAbortSignal,
-      errorAbortSignal,
       timeoutAbortSignal,
       composedAbortSignal,
     };
 
     this.loop.config = patchConfigurable(this.loop.config, {
-      [CONFIG_KEY_ABORT_SIGNALS]: currentSignals,
+      [CONFIG_KEY_ABORT_SIGNALS]: signals,
     });
 
-    return currentSignals;
+    return { signals, disposeCombinedSignal };
   }
 
   /**
@@ -300,21 +286,16 @@ export class PregelRunner {
 
     let startedTasksCount = 0;
 
-    let listener: () => void;
-    const timeoutOrCancelSignal =
-      signals?.externalAbortSignal || signals?.timeoutAbortSignal
-        ? combineAbortSignals(
-            ...(signals.externalAbortSignal
-              ? [signals.externalAbortSignal]
-              : []),
-            ...(signals.timeoutAbortSignal ? [signals.timeoutAbortSignal] : [])
-          )
-        : undefined;
+    let listener: (() => void) | undefined;
+    const timeoutOrCancelSignal = combineAbortSignals(
+      signals?.externalAbortSignal,
+      signals?.timeoutAbortSignal
+    );
 
-    const abortPromise = timeoutOrCancelSignal
+    const abortPromise = timeoutOrCancelSignal.signal
       ? new Promise<never>((_resolve, reject) => {
           listener = () => reject(new Error("Abort"));
-          timeoutOrCancelSignal.addEventListener("abort", listener, {
+          timeoutOrCancelSignal.signal?.addEventListener("abort", listener, {
             once: true,
           });
         })
@@ -357,6 +338,12 @@ export class PregelRunner {
       }
 
       yield settledTask as SettledPregelTask;
+
+      if (listener != null) {
+        timeoutOrCancelSignal.signal?.removeEventListener("abort", listener);
+        timeoutOrCancelSignal.dispose?.();
+      }
+
       delete executingTasksMap[(settledTask as SettledPregelTask).task.id];
     }
   }

libs/langgraph/src/pregel/types.ts

@@ -592,11 +592,6 @@ export type PregelAbortSignals = {
   /** Aborts when the user calls `stream.cancel()` or aborts the {@link AbortSignal} that they passed in via the `signal` option */
   externalAbortSignal?: AbortSignal;
 
-  /**
-   * Aborts when the currently executing task throws any error other than a {@link GraphBubbleUp}
-   */
-  errorAbortSignal?: AbortSignal;
-
   /**
    * Aborts when the currently executing task throws any error other than a {@link GraphBubbleUp}
    */

libs/langgraph/src/pregel/utils/index.ts

@@ -138,32 +138,40 @@ export function patchCheckpointMap(
 /**
  * Combine multiple abort signals into a single abort signal.
  * @param signals - The abort signals to combine.
- * @returns A single abort signal that is aborted if any of the input signals are aborted.
+ * @returns A combined abort signal and a dispose function to remove the abort listener if unused.
  */
-export function combineAbortSignals(...signals: AbortSignal[]): AbortSignal {
-  if (signals.length === 1) {
-    return signals[0];
+export function combineAbortSignals(...x: (AbortSignal | undefined)[]): {
+  signal: AbortSignal | undefined;
+  dispose?: () => void;
+} {
+  const signals = [...new Set(x.filter((s) => s !== undefined))];
+
+  if (signals.length === 0) {
+    return { signal: undefined, dispose: undefined };
+  }
+
+  if (signals.length === 1) {
+    return { signal: signals[0], dispose: undefined };
   }
 
-  // AbortSignal.any() does seem to suffer from memory leaks
-  // @see https://github.com/nodejs/node/issues/55328
-  // if ("any" in AbortSignal) {
-  //   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  //   return (AbortSignal as any).any(signals);
-  // }
   const combinedController = new AbortController();
   const listener = () => {
     combinedController.abort();
     signals.forEach((s) => s.removeEventListener("abort", listener));
   };
 
-  signals.forEach((s) => s.addEventListener("abort", listener));
+  signals.forEach((s) => s.addEventListener("abort", listener, { once: true }));
 
   if (signals.some((s) => s.aborted)) {
     combinedController.abort();
   }
 
-  return combinedController.signal;
+  return {
+    signal: combinedController.signal,
+    dispose: () => {
+      signals.forEach((s) => s.removeEventListener("abort", listener));
+    },
+  };
 }
 
 /**

libs/langgraph/src/tests/pregel/pregel.cancellation.test.ts

@@ -3,7 +3,6 @@ import { v4 as uuidv4 } from "uuid";
 import {
   Annotation,
   Command,
-  CompiledStateGraph,
   END,
   LangGraphRunnableConfig,
   MemorySaver,
@@ -64,82 +63,63 @@ describe("Pregel AbortSignal", () => {
     }),
   });
 
+  function getKey(config: LangGraphRunnableConfig) {
+    return config.configurable?.invokeId ?? "undefined";
+  }
+
   function createGraph({
     mode,
     checkSignal,
   }: {
     mode: TestMode;
     checkSignal: boolean;
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  }): CompiledStateGraph<
-    typeof StateAnnotation.State,
-    typeof StateAnnotation.Update,
-    string,
-    typeof StateAnnotation.spec,
-    typeof StateAnnotation.spec
-  > {
+  }) {
     const graph = new StateGraph(StateAnnotation)
-      .addNode(
-        "one",
-        async (
-          _: typeof StateAnnotation.State,
-          config: LangGraphRunnableConfig
-        ) => {
-          oneCount += 1;
-          if (checkSignal) {
-            const { signal } = config;
-            expect(signal).toBeDefined();
-            return new Promise((resolve, reject) => {
-              const listener = () => {
-                oneRejected = true;
-                reject(new Error("Aborted"));
-              };
+      .addNode("one", async (_, config) => {
+        const key = getKey(config);
 
-              signal!.addEventListener("abort", listener, { once: true });
-              setTimeout(() => {
-                if (!signal!.aborted) {
-                  signal!.removeEventListener("abort", listener);
-                  oneResolved = true;
-                  resolve({
-                    nodeLog: {
-                      [config.configurable!.invokeId]: new Set(["one"]),
-                    },
-                  });
-                }
-              }, 50);
-            });
-          } else {
-            await new Promise((resolve) => {
-              oneResolved = true;
-              setTimeout(resolve, 50);
-            });
-            return {
-              nodeLog: {
-                [config.configurable!.invokeId]: new Set(["one"]),
-              },
+        oneCount += 1;
+
+        if (checkSignal) {
+          const { signal } = config;
+          expect(signal).toBeDefined();
+          return new Promise((resolve, reject) => {
+            const listener = () => {
+              oneRejected = true;
+              reject(new Error("Aborted"));
             };
-          }
+
+            signal!.addEventListener("abort", listener, { once: true });
+            setTimeout(() => {
+              if (!signal!.aborted) {
+                signal!.removeEventListener("abort", listener);
+                oneResolved = true;
+
+                resolve({ nodeLog: { [key]: new Set(["one"]) } });
+              }
+            }, 50);
+          });
         }
-      )
-      .addNode(
-        "two",
-        (
-          state: typeof StateAnnotation.State,
-          config: LangGraphRunnableConfig
-        ) => {
-          twoCount += 1;
-          if (state.shouldThrow) {
-            twoRejected = true;
-            throw new Error("Should not be called!");
-          }
-          twoResolved = true;
-          return {
-            nodeLog: {
-              [config.configurable!.invokeId]: new Set(["two"]),
-            },
-          };
+
+        await new Promise((resolve) => {
+          oneResolved = true;
+          setTimeout(resolve, 50);
+        });
+
+        return { nodeLog: { [key]: new Set(["one"]) } };
+      })
+      .addNode("two", (state, config) => {
+        const key = getKey(config);
+
+        twoCount += 1;
+        if (state.shouldThrow) {
+          twoRejected = true;
+          throw new Error("Should not be called!");
         }
-      )
+
+        twoResolved = true;
+        return { nodeLog: { [key]: new Set(["two"]) } };
+      })
       .addEdge(START, "one")
       .addEdge("one", "two")
       .addEdge("two", END)
@@ -151,30 +131,19 @@ describe("Pregel AbortSignal", () => {
 
     if (mode === "Subgraph called within node without config") {
       return new StateGraph(StateAnnotation)
-        .addNode(
-          "graph",
-          async (
-            { shouldThrow }: typeof StateAnnotation.State,
-            config: LangGraphRunnableConfig
-          ) => {
-            // IMPORTANT: We're explicitly not passing the config here.
-            const result = await graph.invoke({ shouldThrow });
-            // returning two update commands here to make the reducer do the "heavy lifting" of
-            // combining the subgraph result with the parent graph result.
-            return [
-              new Command({
-                update: {
-                  nodeLog: {
-                    [config.configurable!.invokeId]: new Set(["graph"]),
-                  },
-                },
-              }),
-              new Command({
-                update: result,
-              }),
-            ];
-          }
-        )
+        .addNode("graph", async ({ shouldThrow }, config) => {
+          const key = getKey(config);
+
+          // IMPORTANT: We're explicitly not passing the config here.
+          const result = await graph.invoke({ shouldThrow });
+          // returning two update commands here to make the reducer do the "heavy lifting" of
+          // combining the subgraph result with the parent graph result.
+
+          return [
+            new Command({ update: { nodeLog: { [key]: new Set(["graph"]) } } }),
+            new Command({ update: result }),
+          ];
+        })
         .addEdge(START, "graph")
         .addEdge("graph", END)
         .compile({ checkpointer });
@@ -182,30 +151,18 @@ describe("Pregel AbortSignal", () => {
 
     if (mode === "Subgraph called within node with config") {
       return new StateGraph(StateAnnotation)
-        .addNode(
-          "graph",
-          async (
-            { shouldThrow }: typeof StateAnnotation.State,
-            config: LangGraphRunnableConfig
-          ) => {
-            // IMPORTANT: We're explicitly passing the config here, as that's the point of this test case
-            const result = await graph.invoke({ shouldThrow }, config);
-            // returning two update commands here to make the reducer do the "heavy lifting" of
-            // combining the subgraph result with the parent graph result.
-            return [
-              new Command({
-                update: {
-                  nodeLog: {
-                    [config.configurable!.invokeId]: new Set(["graph"]),
-                  },
-                },
-              }),
-              new Command({
-                update: result,
-              }),
-            ];
-          }
-        )
+        .addNode("graph", async ({ shouldThrow }, config) => {
+          const key = getKey(config);
+
+          // IMPORTANT: We're explicitly passing the config here, as that's the point of this test case
+          const result = await graph.invoke({ shouldThrow }, config);
+          // returning two update commands here to make the reducer do the "heavy lifting" of
+          // combining the subgraph result with the parent graph result.
+          return [
+            new Command({ update: { nodeLog: { [key]: new Set(["graph"]) } } }),
+            new Command({ update: result }),
+          ];
+        })
         .addEdge(START, "graph")
         .addEdge("graph", END)
         .compile({ checkpointer });
@@ -233,16 +190,13 @@ describe("Pregel AbortSignal", () => {
       const abortController = new AbortController();
       const config = {
         signal: abortController.signal,
-        configurable: {
-          thread_id: uuidv4(),
-        },
+        configurable: { thread_id: uuidv4() },
       };
 
       setTimeout(() => abortController.abort(), 10);
 
-      await expect(
-        async () =>
-          await createGraph({ mode, checkSignal: false }).invoke({}, config)
+      await expect(() =>
+        createGraph({ mode, checkSignal: false }).invoke({}, config)
       ).rejects.toThrow("Abort");
 
       // Ensure that the `twoCount` has had time to increment before we check it, in case the stream aborted but the graph execution didn't.
@@ -254,6 +208,60 @@ describe("Pregel AbortSignal", () => {
     }
   );
 
+  it.each([
+    "Single layer graph",
+    "Subgraph called within node without config",
+    "Subgraph called within node with config",
+    "Subgraph called as node",
+  ] as TestMode[])(
+    "%s should cancel when step timeout is triggered",
+    async (mode) => {
+      const graph = createGraph({ mode, checkSignal: false });
+      graph.stepTimeout = 10;
+
+      await expect(() =>
+        graph.invoke({}, { configurable: { thread_id: uuidv4() } })
+      ).rejects.toThrow("Abort");
+
+      // Ensure that the `twoCount` has had time to increment before we check it, in case the stream aborted but the graph execution didn't.
+      await new Promise((resolve) => {
+        setTimeout(resolve, 300);
+      });
+      expect(oneCount).toEqual(1);
+      expect(twoCount).toEqual(0);
+    }
+  );
+
+  it.each([
+    "Single layer graph",
+    "Subgraph called within node without config",
+    "Subgraph called within node with config",
+    "Subgraph called as node",
+  ] as TestMode[])(
+    "%s should cancel when external AbortSignal is aborted while step timeout is set",
+    async (mode) => {
+      const graph = createGraph({ mode, checkSignal: false });
+      graph.stepTimeout = 100;
+
+      const abortController = new AbortController();
+      setTimeout(() => abortController.abort(), 10);
+
+      const config = {
+        configurable: { thread_id: uuidv4() },
+        signal: abortController.signal,
+      };
+
+      await expect(() => graph.invoke({}, config)).rejects.toThrow("Abort");
+
+      // Ensure that the `twoCount` has had time to increment before we check it, in case the stream aborted but the graph execution didn't.
+      await new Promise((resolve) => {
+        setTimeout(resolve, 300);
+      });
+      expect(oneCount).toEqual(1);
+      expect(twoCount).toEqual(0);
+    }
+  );
+
   it.each([
     "Single layer graph",
     "Subgraph called within node without config",
@@ -262,17 +270,13 @@ describe("Pregel AbortSignal", () => {
   ] as TestMode[])(
     "%s should pass AbortSignal into nodes via config when timeout is provided but no external signal is given",
     async (mode) => {
-      const config = {
-        timeout: 10,
-        configurable: {
-          thread_id: uuidv4(),
-        },
-      };
+      const config = { timeout: 10, configurable: { thread_id: uuidv4() } };
 
-      await expect(
-        async () =>
-          await createGraph({ mode, checkSignal: true }).invoke({}, config)
-      ).rejects.toThrow("Abort");
+      await expect(() => {
+        const graph = createGraph({ mode, checkSignal: true });
+        graph.stepTimeout = 10;
+        return graph.invoke({}, config);
+      }).rejects.toThrow("Abort");
 
       // Ensure that the `twoCount` has had time to increment before we check it, in case the stream aborted but the graph execution didn't.
       await new Promise((resolve) => {
@@ -297,14 +301,11 @@ describe("Pregel AbortSignal", () => {
       const config = {
         signal: abortController.signal,
         timeout: 10,
-        configurable: {
-          thread_id: uuidv4(),
-        },
+        configurable: { thread_id: uuidv4() },
       };
 
-      await expect(
-        async () =>
-          await createGraph({ mode, checkSignal: true }).invoke({}, config)
+      await expect(() =>
+        createGraph({ mode, checkSignal: true }).invoke({}, config)
       ).rejects.toThrow("Abort");
 
       // Ensure that the `twoCount` has had time to increment before we check it, in case the stream aborted but the graph execution didn't.
@@ -337,9 +338,8 @@ describe("Pregel AbortSignal", () => {
 
       setTimeout(() => abortController.abort(), 10);
 
-      await expect(
-        async () =>
-          await createGraph({ mode, checkSignal: true }).invoke({}, config)
+      await expect(() =>
+        createGraph({ mode, checkSignal: true }).invoke({}, config)
       ).rejects.toThrow("Abort");
 
       // Ensure that the `twoCount` has had time to increment before we check it, in case the stream aborted but the graph execution didn't.
@@ -375,13 +375,12 @@ describe("Pregel AbortSignal", () => {
 
       const thread1Execution1Result = await graph.invoke(
         { shouldThrow: false },
-        {
-          configurable: { thread_id: thread1Id, invokeId: "1" },
-        }
+        { configurable: { thread_id: thread1Id, invokeId: "1" } }
       );
 
       expect(oneCount).toEqual(1);
       expect(twoCount).toEqual(1);
+
       expect(oneRejected).toEqual(false);
       expect(oneResolved).toEqual(true);
       expect(twoRejected).toEqual(false);
@@ -394,9 +393,7 @@ describe("Pregel AbortSignal", () => {
 
       const thread1Execution2Result = await graph.invoke(
         { shouldThrow: false },
-        {
-          configurable: { thread_id: thread1Id, invokeId: "2" },
-        }
+        { configurable: { thread_id: thread1Id, invokeId: "2" } }
       );
 
       expect(oneCount).toEqual(2);
@@ -413,9 +410,7 @@ describe("Pregel AbortSignal", () => {
 
       const thread2Execution1Result = await graph.invoke(
         { shouldThrow: false },
-        {
-          configurable: { thread_id: thread2Id, invokeId: "1" },
-        }
+        { configurable: { thread_id: thread2Id, invokeId: "1" } }
       );
 
       expect(oneCount).toEqual(3);
@@ -471,9 +466,7 @@ describe("Pregel AbortSignal", () => {
 
       const thread2Execution2Attempt2Result = await graph.invoke(
         { shouldThrow: false },
-        {
-          configurable: { thread_id: thread2Id, invokeId: "2" },
-        }
+        { configurable: { thread_id: thread2Id, invokeId: "2" } }
       );
 
       expect(oneCount).toEqual(5);