feat(langgraph): support defining interrupt maps in StateGraph (#1671)

.changeset/shaky-items-repair.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+Add support for defining multiple interrupts in StateGraph constructor. Interrupts from the map can be picked from the `Runtime` object, ensuring type-safety across multiple interrupts.

libs/langgraph/src/graph/state.ts

@@ -71,8 +71,11 @@ import {
   InteropZodToStateDefinition,
   schemaMetaRegistry,
 } from "./zod/meta.js";
-import { interrupt } from "../interrupt.js";
-import { writer } from "../writer.js";
+import type {
+  InferInterruptResumeType,
+  InferInterruptInputType,
+} from "../interrupt.js";
+import type { InferWriterType } from "../writer.js";
 
 const ROOT = "__root__";
 
@@ -138,16 +141,6 @@ type ToStateDefinition<T> = T extends InteropZodObject
   ? T
   : never;
 
-type InterruptResumeType<T> = T extends typeof interrupt<unknown, infer R>
-  ? R
-  : unknown;
-
-type InterruptInputType<T> = T extends typeof interrupt<infer I, unknown>
-  ? I
-  : unknown;
-
-type InferWriterType<T> = T extends typeof writer<infer C> ? C : any; // eslint-disable-line @typescript-eslint/no-explicit-any
-
 type NodeAction<
   S,
   U,
@@ -171,7 +164,7 @@ type StrictNodeAction<
   Prettify<S>,
   | U
   | Command<
-      InterruptResumeType<InterruptType>,
+      InferInterruptResumeType<InterruptType>,
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       U & Record<string, any>,
       Nodes
@@ -484,11 +477,12 @@ export class StateGraph<
     }
 
     // Handle runtime config options
-
     if (isOptions(contextSchema)) {
       if (isInteropZodObject(contextSchema.context)) {
         this._configRuntimeSchema = contextSchema.context;
       }
+      this._interrupt = contextSchema.interrupt as InterruptType;
+      this._writer = contextSchema.writer as WriterType;
     } else if (isInteropZodObject(contextSchema)) {
       this._configRuntimeSchema = contextSchema;
     }
@@ -926,8 +920,19 @@ export class StateGraph<
     const streamChannels =
       streamKeys.length === 1 && streamKeys[0] === ROOT ? ROOT : streamKeys;
 
+    const userInterrupt = this._interrupt;
     // create empty compiled graph
-    const compiled = new CompiledStateGraph<S, U, N, I, O, C, NodeReturnType>({
+    const compiled = new CompiledStateGraph<
+      S,
+      U,
+      N,
+      I,
+      O,
+      C,
+      NodeReturnType,
+      InterruptType,
+      WriterType
+    >({
       builder: this,
       checkpointer,
       interruptAfter,
@@ -946,6 +951,7 @@ export class StateGraph<
       cache,
       name,
       description,
+      userInterrupt,
     });
 
     // attach nodes, edges and branches
@@ -1024,7 +1030,7 @@ export class CompiledStateGraph<
   UpdateType<ToStateDefinition<I>>,
   StateType<ToStateDefinition<O>>,
   NodeReturnType,
-  CommandInstance<InterruptResumeType<InterruptType>, Prettify<U>, N>,
+  CommandInstance<InferInterruptResumeType<InterruptType>, Prettify<U>, N>,
   InferWriterType<WriterType>
 > {
   declare builder: StateGraph<unknown, S, U, N, I, O, C, NodeReturnType>;
@@ -1050,7 +1056,7 @@ export class CompiledStateGraph<
       UpdateType<ToStateDefinition<I>>,
       StateType<ToStateDefinition<O>>,
       NodeReturnType,
-      CommandInstance<InterruptResumeType<InterruptType>, Prettify<U>, N>,
+      CommandInstance<InferInterruptResumeType<InterruptType>, Prettify<U>, N>,
       InferWriterType<WriterType>
     >
   >[0]) {
@@ -1311,7 +1317,7 @@ export class CompiledStateGraph<
   }
 
   public isInterrupted(input: unknown): input is {
-    [INTERRUPT]: Interrupt<InterruptInputType<InterruptType>>[];
+    [INTERRUPT]: Interrupt<InferInterruptInputType<InterruptType>>[];
   } {
     return isInterrupted(input);
   }

libs/langgraph/src/interrupt.ts

@@ -111,3 +111,32 @@ export function interrupt<I = unknown, R = any>(value: I): R {
   const id = ns ? XXH3(ns.join(CHECKPOINT_NAMESPACE_SEPARATOR)) : undefined;
   throw new GraphInterrupt([{ id, value }]);
 }
+
+type FilterAny<X> = (<T>() => T extends X ? 1 : 2) extends <
+  T
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+>() => T extends any ? 1 : 2
+  ? never
+  : X;
+
+export type InferInterruptInputType<T> = T extends typeof interrupt<
+  infer I,
+  unknown
+>
+  ? I
+  : // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  T extends { [key: string]: typeof interrupt<any, any> }
+  ? { [K in keyof T]: InferInterruptInputType<T[K]> }[keyof T]
+  : unknown;
+
+export type InferInterruptResumeType<
+  T,
+  TInner = false
+> = T extends typeof interrupt<never, infer R>
+  ? TInner extends true
+    ? FilterAny<R>
+    : R
+  : // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  T extends { [key: string]: typeof interrupt<any, any> }
+  ? { [K in keyof T]: InferInterruptResumeType<T[K], true> }[keyof T]
+  : unknown;

libs/langgraph/src/pregel/index.ts

@@ -399,9 +399,7 @@ export class Pregel<
     OutputType,
     PregelOptions<Nodes, Channels, ContextType>
   >
-  implements
-    PregelInterface<Nodes, Channels, ContextType>,
-    PregelParams<Nodes, Channels>
+  implements PregelInterface<Nodes, Channels, ContextType>
 {
   /**
    * Name of the class when serialized
@@ -497,13 +495,23 @@ export class Pregel<
    */
   store?: BaseStore;
 
-  triggerToNodes: Record<string, string[]> = {};
-
   /**
    * Optional cache for the graph, useful for caching tasks.
    */
   cache?: BaseCache;
 
+  /**
+   * Optional interrupt helper function.
+   * @internal
+   */
+  private userInterrupt?: unknown;
+
+  /**
+   * The trigger to node mapping for the graph run.
+   * @internal
+   */
+  private triggerToNodes: Record<string, string[]> = {};
+
   /**
    * Constructor for Pregel - meant for internal use only.
    *
@@ -549,6 +557,7 @@ export class Pregel<
     this.cache = fields.cache;
     this.name = fields.name;
     this.triggerToNodes = fields.triggerToNodes ?? this.triggerToNodes;
+    this.userInterrupt = fields.userInterrupt;
 
     if (this.autoValidate) {
       this.validate();
@@ -2051,7 +2060,7 @@ export class Pregel<
       stream.push([ns ?? [], "custom", chunk]);
     };
 
-    config.interrupt ??= interrupt;
+    config.interrupt ??= (this.userInterrupt as typeof interrupt) ?? interrupt;
 
     const callbackManager = await getCallbackManagerForConfig(config);
     const runManager = await callbackManager?.handleChainStart(

libs/langgraph/src/pregel/types.ts

@@ -504,6 +504,12 @@ export type PregelParams<
    * @internal
    */
   triggerToNodes?: Record<string, string[]>;
+
+  /**
+   * Interrupt helper function.
+   * @internal
+   */
+  userInterrupt?: unknown;
 };
 
 export interface PregelTaskDescription {

libs/langgraph/src/tests/interrupt.test-d.ts

@@ -0,0 +1,160 @@
+import { expectTypeOf, it } from "vitest";
+import {
+  InferInterruptResumeType,
+  interrupt,
+  type InferInterruptInputType,
+} from "../interrupt.js";
+
+function all<T>(value: T): {
+  InputType: InferInterruptInputType<T>;
+  OutputType: InferInterruptResumeType<T>;
+} {
+  return {
+    InputType: value as InferInterruptInputType<T>,
+    OutputType: value as InferInterruptResumeType<T>,
+  };
+}
+
+it("interrupt single", () => {
+  const actual = all(interrupt<"input:single", "output:single">);
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: "input:single";
+    OutputType: "output:single";
+  }>();
+});
+
+it("interrupt only invoke", () => {
+  const actual = all(interrupt<"input:single">);
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: "input:single";
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    OutputType: any;
+  }>();
+});
+
+it("interrupt map", () => {
+  const actual = all({
+    one: interrupt<"input:one", "output:one">,
+    two: interrupt<"input:two">,
+    three: interrupt<"input:three", "output:three">,
+  });
+
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: "input:one" | "input:two" | "input:three";
+    OutputType: "output:one" | "output:three";
+  }>();
+});
+
+it("interrupt with complex types", () => {
+  interface ComplexInput {
+    id: string;
+    data: { value: number; nested: { flag: boolean } };
+  }
+
+  interface ComplexOutput {
+    result: string;
+    metadata: { processed: boolean };
+  }
+
+  const actual = all(interrupt<ComplexInput, ComplexOutput>);
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: ComplexInput;
+    OutputType: ComplexOutput;
+  }>();
+});
+
+it("interrupt with union types", () => {
+  type InputUnion =
+    | { type: "A"; valueA: string }
+    | { type: "B"; valueB: number };
+  type OutputUnion =
+    | { success: true; data: string }
+    | { success: false; error: string };
+
+  const actual = all(interrupt<InputUnion, OutputUnion>);
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: InputUnion;
+    OutputType: OutputUnion;
+  }>();
+});
+
+it("interrupt with null and undefined", () => {
+  const nullInterrupt = all(interrupt<null, string>);
+  expectTypeOf(nullInterrupt).toEqualTypeOf<{
+    InputType: null;
+    OutputType: string;
+  }>();
+
+  const undefinedInterrupt = all(interrupt<undefined, number>);
+  expectTypeOf(undefinedInterrupt).toEqualTypeOf<{
+    InputType: undefined;
+    OutputType: number;
+  }>();
+
+  const voidInterrupt = all(interrupt<void, boolean>);
+  expectTypeOf(voidInterrupt).toEqualTypeOf<{
+    InputType: void;
+    OutputType: boolean;
+  }>();
+});
+
+it("interrupt with never type", () => {
+  const actual = all(interrupt<never, string>);
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: never;
+    OutputType: string;
+  }>();
+});
+
+it("interrupt mixed map types", () => {
+  const actual = all({
+    stringInterrupt: interrupt<string, number>,
+    objectInterrupt: interrupt<{ id: string }, { result: boolean }>,
+    arrayInterrupt: interrupt<string[], number[]>,
+    primitiveInterrupt: interrupt<boolean>,
+  });
+
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: string | { id: string } | string[] | boolean;
+    OutputType: number | { result: boolean } | number[];
+  }>();
+});
+
+it("interrupt with optional types", () => {
+  const actual = all(
+    interrupt<{ required: string; optional?: number }, string>
+  );
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: { required: string; optional?: number };
+    OutputType: string;
+  }>();
+});
+
+it("interrupt with readonly types", () => {
+  const actual = all(interrupt<readonly string[], Readonly<{ value: number }>>);
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: readonly string[];
+    OutputType: Readonly<{ value: number }>;
+  }>();
+});
+
+it("interrupt with tuple types", () => {
+  const actual = all(interrupt<[string, number, boolean], [number, string]>);
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: [string, number, boolean];
+    OutputType: [number, string];
+  }>();
+});
+
+it("interrupt empty map", () => {
+  const actual = all({});
+  expectTypeOf(actual).toEqualTypeOf<{ InputType: never; OutputType: never }>();
+});
+
+it("interrupt unknown", () => {
+  const actual = all({} as unknown);
+  expectTypeOf(actual).toEqualTypeOf<{
+    InputType: unknown;
+    OutputType: unknown;
+  }>();
+});

libs/langgraph/src/tests/pregel.test-d.ts

@@ -549,7 +549,10 @@ it("state graph builder", async () => {
       }),
 
       writer: writer<{ custom: string }>,
-      interrupt: interrupt<{ reason: string }, { messages: string[] }>,
+      interrupt: {
+        simple: interrupt<{ reason: string }, { messages: string[] }>,
+        complex: interrupt<{ complex: string }, { messages: string[] }>,
+      },
 
       // Allow optionally specifying all nodes upfront.
       // This is needed to properly type `goto` commands and removes
@@ -565,7 +568,7 @@ it("state graph builder", async () => {
     // @ts-expect-error - Invalid writer input
     if (false) runtime.interrupt();
 
-    const result = runtime.interrupt({ reason: "hello" });
+    const result = runtime.interrupt.simple({ reason: "hello" });
 
     runtime.writer({ custom: "hello" });
 
@@ -582,7 +585,7 @@ it("state graph builder", async () => {
   };
 
   const what: typeof builder.Node = async (_, runtime) => {
-    const result = runtime.interrupt({ reason: "what" });
+    const result = runtime.interrupt.complex({ complex: "what" });
     if (result != null) return { messages: result.messages };
     return {};
   };
@@ -611,7 +614,7 @@ it("state graph builder", async () => {
 
   if (graph.isInterrupted(first)) {
     expectTypeOf(first.__interrupt__).toExtend<
-      { id?: string; value?: { reason: string } }[]
+      { id?: string; value?: { reason: string } | { complex: string } }[]
     >();
   }
 
@@ -621,6 +624,9 @@ it("state graph builder", async () => {
   // @ts-expect-error - Invalid update value
   if (false) await thread.invoke(new Command({ update: { messages: true } }));
 
+  // @ts-expect-error - Invalid resume value
+  if (false) await thread.invoke(new Command({ resume: { messages: true } }));
+
   const second = await gatherIterator(
     thread.stream(new Command({ resume: { messages: ["resume: hello"] } }), {
       streamMode: ["values", "custom"],
@@ -634,7 +640,7 @@ it("state graph builder", async () => {
     [
       "values",
       {
-        __interrupt__: [{ id: expect.any(String), value: { reason: "what" } }],
+        __interrupt__: [{ id: expect.any(String), value: { complex: "what" } }],
       },
     ],
   ]);
@@ -643,7 +649,7 @@ it("state graph builder", async () => {
     if (mode === "values") {
       if (graph.isInterrupted(value)) {
         expectTypeOf(value.__interrupt__).toExtend<
-          { id?: string; value?: { reason: string } }[]
+          { id?: string; value?: { reason: string } | { complex: string } }[]
         >();
       }
     }

libs/langgraph/src/writer.ts

@@ -15,3 +15,5 @@ export function writer<T>(chunk: T): void {
 
   return conf.writer?.(chunk);
 }
+
+export type InferWriterType<T> = T extends typeof writer<infer C> ? C : any; // eslint-disable-line @typescript-eslint/no-explicit-any