feat(langgraph): allow explicit type declaration for interrupt and custom stream events (#1612)

.changeset/tame-humans-tap.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+Allow defining types for interrupt and custom events upfront

libs/langgraph/src/constants.ts

@@ -99,8 +99,16 @@ const COMMAND_SYMBOL = Symbol.for("langgraph.command");
  * @see {@link Command}
  * @internal
  */
-export class CommandInstance {
-  [COMMAND_SYMBOL]: true;
+export class CommandInstance<
+  Resume = unknown,
+  Update extends Record<string, unknown> = Record<string, unknown>,
+  Nodes extends string = string
+> {
+  [COMMAND_SYMBOL]: CommandParams<Resume, Update, Nodes>;
+
+  constructor(args: CommandParams<Resume, Update, Nodes>) {
+    this[COMMAND_SYMBOL] = args;
+  }
 }
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -329,7 +337,7 @@ export class Command<
   Resume = unknown,
   Update extends Record<string, unknown> = Record<string, unknown>,
   Nodes extends string = string
-> extends CommandInstance {
+> extends CommandInstance<Resume, Update, Nodes> {
   readonly lg_name = "Command";
 
   lc_direct_tool_output = true;
@@ -364,8 +372,8 @@ export class Command<
 
   static PARENT = "__parent__";
 
-  constructor(args: CommandParams<Resume, Update, Nodes>) {
-    super();
+  constructor(args: Omit<CommandParams<Resume, Update, Nodes>, "lg_name">) {
+    super(args);
     this.resume = args.resume;
     this.graph = args.graph;
     this.update = args.update;

libs/langgraph/src/graph/graph.ts

@@ -12,7 +12,7 @@ import {
   Graph as DrawableGraph,
 } from "@langchain/core/runnables/graph";
 import { All, BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
-import { z } from "zod";
+import { z } from "zod/v4";
 import { validate as isUuid } from "uuid";
 import type {
   RunnableLike,
@@ -581,7 +581,10 @@ export class CompiledGraph<
   InputType = any,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   OutputType = any,
-  NodeReturnType = unknown
+  NodeReturnType = unknown,
+  CommandType = unknown,
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  StreamCustomType = any
 > extends Pregel<
   Record<N | typeof START, PregelNode<State, Update>>,
   Record<N | typeof START | typeof END | string, BaseChannel>,
@@ -591,7 +594,9 @@ export class CompiledGraph<
   OutputType,
   InputType,
   OutputType,
-  NodeReturnType
+  NodeReturnType,
+  CommandType,
+  StreamCustomType
 > {
   declare NodeType: N;
 
@@ -694,12 +699,7 @@ export class CompiledGraph<
     const xray = config?.xray;
     const graph = new DrawableGraph();
     const startNodes: Record<string, DrawableGraphNode> = {
-      [START]: graph.addNode(
-        {
-          schema: z.any(),
-        },
-        START
-      ),
+      [START]: graph.addNode({ schema: z.any() }, START),
     };
     const endNodes: Record<string, DrawableGraphNode> = {};
     // eslint-disable-next-line @typescript-eslint/no-explicit-any

libs/langgraph/src/graph/index.ts

@@ -12,7 +12,6 @@ export {
   type StateGraphArgs,
   StateGraph,
   CompiledStateGraph,
-  typedNode,
 } from "./state.js";
 export {
   MessageGraph,

libs/langgraph/src/graph/state.ts

@@ -15,6 +15,7 @@ import {
 import type {
   RunnableLike,
   LangGraphRunnableConfig,
+  Runtime,
 } from "../pregel/runnable_types.js";
 import { BaseChannel, isBaseChannel } from "../channels/base.js";
 import {
@@ -48,6 +49,10 @@ import {
   Send,
   START,
   TAG_HIDDEN,
+  CommandInstance,
+  isInterrupted,
+  Interrupt,
+  INTERRUPT,
 } from "../constants.js";
 import { InvalidUpdateError, ParentCommand } from "../errors.js";
 import {
@@ -66,6 +71,8 @@ import {
   InteropZodToStateDefinition,
   schemaMetaRegistry,
 } from "./zod/meta.js";
+import { interrupt } from "../interrupt.js";
+import { writer } from "../writer.js";
 
 const ROOT = "__root__";
 
@@ -131,10 +138,49 @@ type ToStateDefinition<T> = T extends InteropZodObject
   ? T
   : never;
 
-type NodeAction<S, U, C extends SDZod> = RunnableLike<
+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,
+  C extends SDZod,
+  InterruptType,
+  WriterType
+> = RunnableLike<
   S,
   U extends object ? U & Record<string, any> : U, // eslint-disable-line @typescript-eslint/no-explicit-any
-  LangGraphRunnableConfig<StateType<ToStateDefinition<C>>>
+  LangGraphRunnableConfig<
+    StateType<ToStateDefinition<C>>,
+    InterruptType,
+    WriterType
+  >
+>;
+
+type StrictNodeAction<
+  S,
+  U,
+  C extends SDZod,
+  Nodes extends string,
+  InterruptType,
+  WriterType
+> = RunnableLike<
+  Prettify<S>,
+  | U
+  | Command<
+      InterruptResumeType<InterruptType>,
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      U & Record<string, any>,
+      Nodes
+    >,
+  Runtime<StateType<ToStateDefinition<C>>, InterruptType, WriterType>
 >;
 
 const PartialStateSchema = Symbol.for("langgraph.state.partial");
@@ -219,7 +265,9 @@ export class StateGraph<
   I extends SDZod = SD extends SDZod ? ToStateDefinition<SD> : StateDefinition,
   O extends SDZod = SD extends SDZod ? ToStateDefinition<SD> : StateDefinition,
   C extends SDZod = StateDefinition,
-  NodeReturnType = unknown
+  NodeReturnType = unknown,
+  InterruptType = unknown,
+  WriterType = unknown
 > extends Graph<N, S, U, StateGraphNodeSpec<S, U>, ToStateDefinition<C>> {
   channels: Record<string, BaseChannel> = {};
 
@@ -259,6 +307,32 @@ export class StateGraph<
   /** @internal */
   _configRuntimeSchema: InteropZodObject | undefined;
 
+  /** @internal */
+  _interrupt: InterruptType;
+
+  /** @internal */
+  _writer: WriterType;
+
+  declare Node: StrictNodeAction<S, U, C, N, InterruptType, WriterType>;
+
+  constructor(
+    state: SD extends InteropZodObject
+      ? SD
+      : SD extends StateDefinition
+      ? AnnotationRoot<SD>
+      : never,
+    options?: {
+      context?: C | AnnotationRoot<ToStateDefinition<C>>;
+      input?: I | AnnotationRoot<ToStateDefinition<I>>;
+      output?: O | AnnotationRoot<ToStateDefinition<O>>;
+
+      interrupt?: InterruptType;
+      writer?: WriterType;
+
+      nodes?: N[];
+    }
+  );
+
   constructor(
     fields: SD extends StateDefinition
       ? StateGraphArgsWithInputOutputSchemas<SD, ToStateDefinition<O>>
@@ -309,7 +383,17 @@ export class StateGraph<
             >
           | StateGraphArgsWithInputOutputSchemas<SD, ToStateDefinition<O>>
       : StateGraphArgs<S>,
-    contextSchema?: C | AnnotationRoot<ToStateDefinition<C>>
+    contextSchema?:
+      | C
+      | AnnotationRoot<ToStateDefinition<C>>
+      | {
+          input?: I | AnnotationRoot<ToStateDefinition<I>>;
+          output?: O | AnnotationRoot<ToStateDefinition<O>>;
+          context?: C | AnnotationRoot<ToStateDefinition<C>>;
+          interrupt?: InterruptType;
+          writer?: WriterType;
+          nodes?: N[];
+        }
   ) {
     super();
 
@@ -377,7 +461,29 @@ export class StateGraph<
     this._addSchema(this._inputDefinition);
     this._addSchema(this._outputDefinition);
 
-    if (isInteropZodObject(contextSchema)) {
+    function isOptions(options: unknown): options is {
+      context?: C | AnnotationRoot<ToStateDefinition<C>>;
+      input?: I | AnnotationRoot<ToStateDefinition<I>>;
+      output?: O | AnnotationRoot<ToStateDefinition<O>>;
+      interrupt?: InterruptType;
+      writer?: WriterType;
+      nodes?: N[];
+    } {
+      return (
+        typeof options === "object" &&
+        options != null &&
+        !("spec" in options) &&
+        !isInteropZodObject(options)
+      );
+    }
+
+    // Handle runtime config options
+
+    if (isOptions(contextSchema)) {
+      if (isInteropZodObject(contextSchema.context)) {
+        this._configRuntimeSchema = contextSchema.context;
+      }
+    } else if (isInteropZodObject(contextSchema)) {
       this._configRuntimeSchema = contextSchema;
     }
   }
@@ -420,7 +526,7 @@ export class StateGraph<
 
   override addNode<
     K extends string,
-    NodeMap extends Record<K, NodeAction<S, U, C>>
+    NodeMap extends Record<K, NodeAction<S, U, C, InterruptType, WriterType>>
   >(
     nodes: NodeMap
   ): StateGraph<
@@ -434,7 +540,13 @@ export class StateGraph<
     MergeReturnType<
       NodeReturnType,
       {
-        [key in keyof NodeMap]: NodeMap[key] extends NodeAction<S, infer U, C>
+        [key in keyof NodeMap]: NodeMap[key] extends NodeAction<
+          S,
+          infer U,
+          C,
+          InterruptType,
+          WriterType
+        >
           ? U
           : never;
       }
@@ -445,7 +557,13 @@ export class StateGraph<
     nodes:
       | [
           key: K,
-          action: NodeAction<NodeInput, NodeOutput, C>,
+          action: NodeAction<
+            NodeInput,
+            NodeOutput,
+            C,
+            InterruptType,
+            WriterType
+          >,
           options?: StateGraphAddNodeOptions
         ][]
   ): StateGraph<
@@ -461,7 +579,7 @@ export class StateGraph<
 
   override addNode<K extends string, NodeInput = S, NodeOutput extends U = U>(
     key: K,
-    action: NodeAction<NodeInput, NodeOutput, C>,
+    action: NodeAction<NodeInput, NodeOutput, C, InterruptType, WriterType>,
     options?: StateGraphAddNodeOptions
   ): StateGraph<
     SD,
@@ -476,7 +594,7 @@ export class StateGraph<
 
   override addNode<K extends string, NodeInput = S>(
     key: K,
-    action: NodeAction<NodeInput, U, C>,
+    action: NodeAction<NodeInput, U, C, InterruptType, WriterType>,
     options?: StateGraphAddNodeOptions
   ): StateGraph<SD, S, U, N | K, I, O, C, NodeReturnType>;
 
@@ -484,15 +602,21 @@ export class StateGraph<
     ...args:
       | [
           key: K,
-          action: NodeAction<NodeInput, NodeOutput, C>,
+          action: NodeAction<
+            NodeInput,
+            NodeOutput,
+            C,
+            InterruptType,
+            WriterType
+          >,
           options?: StateGraphAddNodeOptions
         ]
       | [
           nodes:
-            | Record<K, NodeAction<NodeInput, U, C>>
+            | Record<K, NodeAction<NodeInput, U, C, InterruptType, WriterType>>
             | [
                 key: K,
-                action: NodeAction<NodeInput, U, C>,
+                action: NodeAction<NodeInput, U, C, InterruptType, WriterType>,
                 options?: StateGraphAddNodeOptions
               ][]
         ]
@@ -501,10 +625,10 @@ export class StateGraph<
       args: unknown[]
     ): args is [
       nodes:
-        | Record<K, NodeAction<NodeInput, U, C>>
+        | Record<K, NodeAction<NodeInput, U, C, InterruptType, WriterType>>
         | [
             key: K,
-            action: NodeAction<NodeInput, U, C>,
+            action: NodeAction<NodeInput, U, C, InterruptType, WriterType>,
             options?: AddNodeOptions
           ][]
     ] {
@@ -515,16 +639,11 @@ export class StateGraph<
       isMultipleNodes(args) // eslint-disable-line no-nested-ternary
         ? Array.isArray(args[0])
           ? args[0]
-          : Object.entries(args[0]).map(([key, action]) => [
-              key,
-              action,
-              // eslint-disable-next-line @typescript-eslint/no-explicit-any
-              (action as any)[Symbol.for("langgraph.state.node")] ?? undefined,
-            ])
+          : Object.entries(args[0]).map(([key, action]) => [key, action])
         : [[args[0], args[1], args[2]]]
     ) as [
       K,
-      NodeAction<NodeInput, U, C>,
+      NodeAction<NodeInput, U, C, InterruptType, WriterType>,
       StateGraphAddNodeOptions | undefined
     ][];
 
@@ -648,7 +767,7 @@ export class StateGraph<
   addSequence<K extends string, NodeInput = S, NodeOutput extends U = U>(
     nodes: [
       key: K,
-      action: NodeAction<NodeInput, NodeOutput, C>,
+      action: NodeAction<NodeInput, NodeOutput, C, InterruptType, WriterType>,
       options?: StateGraphAddNodeOptions
     ][]
   ): StateGraph<
@@ -662,7 +781,10 @@ export class StateGraph<
     MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
   >;
 
-  addSequence<K extends string, NodeMap extends Record<K, NodeAction<S, U, C>>>(
+  addSequence<
+    K extends string,
+    NodeMap extends Record<K, NodeAction<S, U, C, InterruptType, WriterType>>
+  >(
     nodes: NodeMap
   ): StateGraph<
     SD,
@@ -675,7 +797,13 @@ export class StateGraph<
     MergeReturnType<
       NodeReturnType,
       {
-        [key in keyof NodeMap]: NodeMap[key] extends NodeAction<S, infer U, C>
+        [key in keyof NodeMap]: NodeMap[key] extends NodeAction<
+          S,
+          infer U,
+          C,
+          InterruptType,
+          WriterType
+        >
           ? U
           : never;
       }
@@ -686,10 +814,19 @@ export class StateGraph<
     nodes:
       | [
           key: K,
-          action: NodeAction<NodeInput, NodeOutput, C>,
+          action: NodeAction<
+            NodeInput,
+            NodeOutput,
+            C,
+            InterruptType,
+            WriterType
+          >,
           options?: StateGraphAddNodeOptions
         ][]
-      | Record<K, NodeAction<NodeInput, NodeOutput, C>>
+      | Record<
+          K,
+          NodeAction<NodeInput, NodeOutput, C, InterruptType, WriterType>
+        >
   ): StateGraph<
     SD,
     S,
@@ -700,14 +837,7 @@ export class StateGraph<
     C,
     MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
   > {
-    const parsedNodes = Array.isArray(nodes)
-      ? nodes
-      : (Object.entries(nodes).map(([key, action]) => [
-          key,
-          action,
-          // eslint-disable-next-line @typescript-eslint/no-explicit-any
-          (action as any)[Symbol.for("langgraph.state.node")] ?? undefined,
-        ]) as [K, NodeAction<S, U, C>, StateGraphAddNodeOptions | undefined][]);
+    const parsedNodes = Array.isArray(nodes) ? nodes : Object.entries(nodes);
 
     if (parsedNodes.length === 0) {
       throw new Error("Sequence requires at least one node.");
@@ -722,7 +852,11 @@ export class StateGraph<
       }
 
       const validKey = key as unknown as N;
-      this.addNode(validKey, action as NodeAction<S, U, C>, options);
+      this.addNode(
+        validKey,
+        action as NodeAction<S, U, C, InterruptType, WriterType>,
+        options
+      );
       if (previousNode != null) {
         this.addEdge(previousNode, validKey);
       }
@@ -765,7 +899,9 @@ export class StateGraph<
     I,
     O,
     C,
-    NodeReturnType
+    NodeReturnType,
+    InterruptType,
+    WriterType
   > {
     // validate the graph
     this.validate([
@@ -871,7 +1007,9 @@ export class CompiledStateGraph<
   I extends SDZod = StateDefinition,
   O extends SDZod = StateDefinition,
   C extends SDZod = StateDefinition,
-  NodeReturnType = unknown
+  NodeReturnType = unknown,
+  InterruptType = unknown,
+  WriterType = unknown
 > extends CompiledGraph<
   N,
   S,
@@ -879,7 +1017,9 @@ export class CompiledStateGraph<
   StateType<ToStateDefinition<C>>,
   UpdateType<ToStateDefinition<I>>,
   StateType<ToStateDefinition<O>>,
-  NodeReturnType
+  NodeReturnType,
+  CommandInstance<InterruptResumeType<InterruptType>, Prettify<U>, N>,
+  InferWriterType<WriterType>
 > {
   declare builder: StateGraph<unknown, S, U, N, I, O, C, NodeReturnType>;
 
@@ -903,7 +1043,9 @@ export class CompiledStateGraph<
       StateType<ToStateDefinition<C>>,
       UpdateType<ToStateDefinition<I>>,
       StateType<ToStateDefinition<O>>,
-      NodeReturnType
+      NodeReturnType,
+      CommandInstance<InterruptResumeType<InterruptType>, Prettify<U>, N>,
+      InferWriterType<WriterType>
     >
   >[0]) {
     super(rest);
@@ -1162,6 +1304,12 @@ export class CompiledStateGraph<
     return input;
   }
 
+  public isInterrupted(input: unknown): input is {
+    [INTERRUPT]: Interrupt<InterruptInputType<InterruptType>>[];
+  } {
+    return isInterrupted(input);
+  }
+
   protected async _validateContext(
     config: Partial<Record<string, unknown>>
   ): Promise<Partial<Record<string, unknown>>> {
@@ -1300,64 +1448,3 @@ function _getControlBranch() {
     path: CONTROL_BRANCH_PATH,
   });
 }
-
-type TypedNodeAction<
-  SD extends StateDefinition,
-  Nodes extends string,
-  C extends StateDefinition = StateDefinition
-> = RunnableLike<
-  StateType<SD>,
-  UpdateType<SD> | Command<unknown, UpdateType<SD>, Nodes>,
-  LangGraphRunnableConfig<StateType<C>>
->;
-
-export function typedNode<
-  SD extends SDZod,
-  Nodes extends string,
-  C extends SDZod = StateDefinition
->(
-  _state: SD extends StateDefinition ? AnnotationRoot<SD> : never,
-  _options?: {
-    nodes?: Nodes[];
-    config?: C extends StateDefinition ? AnnotationRoot<C> : never;
-  }
-): (
-  func: TypedNodeAction<ToStateDefinition<SD>, Nodes, ToStateDefinition<C>>,
-  options?: StateGraphAddNodeOptions<Nodes>
-) => TypedNodeAction<ToStateDefinition<SD>, Nodes, ToStateDefinition<C>>;
-
-export function typedNode<
-  SD extends SDZod,
-  Nodes extends string,
-  C extends SDZod = StateDefinition
->(
-  _state: SD extends InteropZodObject ? SD : never,
-  _options?: {
-    nodes?: Nodes[];
-    config?: C extends InteropZodObject ? C : never;
-  }
-): (
-  func: TypedNodeAction<ToStateDefinition<SD>, Nodes, ToStateDefinition<C>>,
-  options?: StateGraphAddNodeOptions<Nodes>
-) => TypedNodeAction<ToStateDefinition<SD>, Nodes, ToStateDefinition<C>>;
-
-export function typedNode<
-  SD extends SDZod,
-  Nodes extends string,
-  C extends SDZod = StateDefinition
->(
-  _state: SD extends InteropZodObject
-    ? SD
-    : SD extends StateDefinition
-    ? AnnotationRoot<SD>
-    : never,
-  _options?: { nodes?: Nodes[]; config?: AnnotationRoot<ToStateDefinition<C>> }
-) {
-  return (
-    func: TypedNodeAction<ToStateDefinition<SD>, Nodes, ToStateDefinition<C>>,
-    options?: StateGraphAddNodeOptions<Nodes>
-  ) => {
-    Object.assign(func, { [Symbol.for("langgraph.state.node")]: options });
-    return func;
-  };
-}

libs/langgraph/src/index.ts

@@ -8,6 +8,7 @@ initializeAsyncLocalStorageSingleton();
 export * from "./web.js";
 
 export { interrupt } from "./interrupt.js";
+export { writer } from "./writer.js";
 export { getStore, getWriter, getConfig } from "./pregel/utils/config.js";
 export { getPreviousState } from "./func/index.js";
 export { getCurrentTaskInput } from "./pregel/utils/config.js";

libs/langgraph/src/pregel/index.ts

@@ -119,6 +119,7 @@ import { findSubgraphPregel } from "./utils/subgraph.js";
 import { validateGraph, validateKeys } from "./validate.js";
 import { ChannelWrite, ChannelWriteEntry, PASSTHROUGH } from "./write.js";
 import { Topic } from "../channels/topic.js";
+import { interrupt } from "../interrupt.js";
 
 type WriteValue = Runnable | RunnableFunc<unknown, unknown> | unknown;
 type StreamEventsOptions = Parameters<Runnable["streamEvents"]>[2];
@@ -387,10 +388,13 @@ export class Pregel<
     OutputType = PregelOutputType,
     StreamUpdatesType = InputType,
     StreamValuesType = OutputType,
-    NodeReturnType = unknown
+    NodeReturnType = unknown,
+    CommandType = CommandInstance,
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    StreamCustom = any
   >
   extends PartialRunnable<
-    InputType | CommandInstance | null,
+    InputType | CommandType | null,
     OutputType,
     PregelOptions<Nodes, Channels, ContextType>
   >
@@ -1818,7 +1822,7 @@ export class Pregel<
     TStreamMode extends StreamMode | StreamMode[] | undefined,
     TSubgraphs extends boolean
   >(
-    input: InputType | CommandInstance | null,
+    input: InputType | CommandType | null,
     options?: Partial<
       PregelOptions<Nodes, Channels, ContextType, TStreamMode, TSubgraphs>
     >
@@ -1830,7 +1834,8 @@ export class Pregel<
         StreamUpdatesType,
         StreamValuesType,
         keyof Nodes,
-        NodeReturnType
+        NodeReturnType,
+        StreamCustom
       >
     >
   > {
@@ -1856,7 +1861,8 @@ export class Pregel<
           StreamUpdatesType,
           StreamValuesType,
           keyof Nodes,
-          NodeReturnType
+          NodeReturnType,
+          StreamCustom
         >
       >,
       abortController
@@ -1867,7 +1873,7 @@ export class Pregel<
    * @inheritdoc
    */
   override streamEvents(
-    input: InputType | CommandInstance | null,
+    input: InputType | CommandType | null,
     options: Partial<PregelOptions<Nodes, Channels, ContextType>> & {
       version: "v1" | "v2";
     },
@@ -1875,7 +1881,7 @@ export class Pregel<
   ): IterableReadableStream<StreamEvent>;
 
   override streamEvents(
-    input: InputType | CommandInstance | null,
+    input: InputType | CommandType | null,
     options: Partial<PregelOptions<Nodes, Channels, ContextType>> & {
       version: "v1" | "v2";
       encoding: "text/event-stream";
@@ -1884,7 +1890,7 @@ export class Pregel<
   ): IterableReadableStream<Uint8Array>;
 
   override streamEvents(
-    input: InputType | CommandInstance | null,
+    input: InputType | CommandType | null,
     options: Partial<PregelOptions<Nodes, Channels, ContextType>> & {
       version: "v1" | "v2";
     },
@@ -2034,6 +2040,10 @@ export class Pregel<
       };
     }
 
+    if (this.checkpointer != null) {
+      config.interrupt = interrupt;
+    }
+
     const callbackManager = await getCallbackManagerForConfig(config);
     const runManager = await callbackManager?.handleChainStart(
       this.toJSON(), // chain
@@ -2163,7 +2173,7 @@ export class Pregel<
    * @param options The configuration to use for the run.
    */
   override async invoke(
-    input: InputType | CommandInstance | null,
+    input: InputType | CommandType | null,
     options?: Partial<PregelOptions<Nodes, Channels, ContextType>>
   ): Promise<OutputType> {
     const streamMode = options?.streamMode ?? "values";

libs/langgraph/src/pregel/runnable_types.ts

@@ -24,23 +24,43 @@ export type RunnableLike<
   | RunnableFunc<RunInput, RunOutput, CallOptions>
   | RunnableMapLike<RunInput, RunOutput>;
 
+type IsEqual<T, U> = [T] extends [U] ? ([U] extends [T] ? true : false) : false;
+
 export interface LangGraphRunnableConfig<
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  ContextType extends Record<string, any> = Record<string, any>
+  ContextType extends Record<string, any> = Record<string, any>,
+  InterruptType = unknown,
+  WriterType = unknown
 > extends RunnableConfig<ContextType> {
   context?: ContextType;
 
   store?: BaseStore;
 
-  writer?: (chunk: unknown) => void;
+  writer?: IsEqual<WriterType, unknown> extends true
+    ? (chunk: unknown) => void
+    : WriterType;
+
+  interrupt?: IsEqual<InterruptType, unknown> extends true
+    ? (value: unknown) => unknown
+    : InterruptType;
 }
 
-export interface Runtime<ContextType = Record<string, unknown>> {
+export interface Runtime<
+  ContextType = Record<string, unknown>,
+  InterruptType = unknown,
+  WriterType = unknown
+> {
   context?: ContextType;
 
   store?: BaseStore;
 
-  writer?: (chunk: unknown) => void;
+  writer?: IsEqual<WriterType, unknown> extends true
+    ? (chunk: unknown) => void
+    : WriterType;
+
+  interrupt?: IsEqual<InterruptType, unknown> extends true
+    ? (value: unknown) => unknown
+    : InterruptType;
 
   signal?: AbortSignal;
 }

libs/langgraph/src/pregel/types.ts

@@ -40,9 +40,6 @@ export type PregelOutputType = any;
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type StreamMessageOutput = [BaseMessage, Record<string, any>];
 
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-type StreamCustomOutput = any;
-
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type StreamDebugOutput = Record<string, any>;
 
@@ -83,7 +80,8 @@ export type StreamOutputMap<
   StreamUpdates,
   StreamValues,
   Nodes,
-  NodeReturnType
+  NodeReturnType,
+  StreamCustom
 > = (
   undefined extends TStreamMode
     ? []
@@ -106,7 +104,7 @@ export type StreamOutputMap<
             : Record<Nodes extends string ? Nodes : string, StreamUpdates>
         ];
         messages: [string[], "messages", StreamMessageOutput];
-        custom: [string[], "custom", StreamCustomOutput];
+        custom: [string[], "custom", StreamCustom];
         checkpoints: [
           string[],
           "checkpoints",
@@ -128,7 +126,7 @@ export type StreamOutputMap<
             : Record<Nodes extends string ? Nodes : string, StreamUpdates>
         ];
         messages: ["messages", StreamMessageOutput];
-        custom: ["custom", StreamCustomOutput];
+        custom: ["custom", StreamCustom];
         checkpoints: ["checkpoints", StreamCheckpointsOutput<StreamValues>];
         tasks: ["tasks", StreamTasksOutput<StreamUpdates, StreamValues, Nodes>];
         debug: ["debug", StreamDebugOutput];
@@ -146,7 +144,7 @@ export type StreamOutputMap<
             : Record<Nodes extends string ? Nodes : string, StreamUpdates>
         ];
         messages: [string[], StreamMessageOutput];
-        custom: [string[], StreamCustomOutput];
+        custom: [string[], StreamCustom];
         checkpoints: [string[], StreamCheckpointsOutput<StreamValues>];
         tasks: [
           string[],
@@ -160,7 +158,7 @@ export type StreamOutputMap<
           ? { [K in keyof NodeReturnType]?: NodeReturnType[K] }
           : Record<Nodes extends string ? Nodes : string, StreamUpdates>;
         messages: StreamMessageOutput;
-        custom: StreamCustomOutput;
+        custom: StreamCustom;
         checkpoints: StreamCheckpointsOutput<StreamValues>;
         tasks: StreamTasksOutput<StreamUpdates, StreamValues, Nodes>;
         debug: StreamDebugOutput;

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

@@ -23,6 +23,7 @@ const CONFIG_KEYS = [
   "streamMode",
   "store",
   "writer",
+  "interrupt",
   "context",
   "interruptBefore",
   "interruptAfter",

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

@@ -1,6 +1,9 @@
+/* eslint-disable no-constant-condition */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import "../graph/zod/plugin.js";
 import { z } from "zod/v3";
+import { z as z4 } from "zod/v4";
+import { MemorySaver } from "@langchain/langgraph-checkpoint";
 
 import { expectTypeOf, it, beforeAll, expect } from "vitest";
 import type { BaseMessage } from "@langchain/core/messages";
@@ -10,9 +13,13 @@ import { gatherIterator } from "../utils.js";
 import { StreamMode } from "../pregel/types.js";
 import { task, entrypoint } from "../func/index.js";
 import { initializeAsyncLocalStorageSingleton } from "../setup/async_local_storage.js";
-import { INTERRUPT, isInterrupted, START } from "../constants.js";
+import { Command, END, INTERRUPT, isInterrupted, START } from "../constants.js";
 import { withLangGraph } from "../graph/zod/meta.js";
 import { Runtime } from "../pregel/runnable_types.js";
+import { registry } from "../graph/zod/zod-registry.js";
+import { MessagesZodMeta } from "../graph/messages_annotation.js";
+import { interrupt } from "../interrupt.js";
+import { writer } from "../writer.js";
 
 beforeAll(() => {
   // Will occur naturally if user imports from main `@langchain/langgraph` endpoint.
@@ -517,6 +524,142 @@ it("state graph zod", async () => {
   >();
 });
 
+it("state graph builder", async () => {
+  const checkpointer = new MemorySaver();
+  const builder = new StateGraph(
+    z4.object({
+      messages: z4.custom<BaseMessage[]>().register(registry, MessagesZodMeta),
+      innerReason: z4.string(),
+    }),
+    {
+      input: z4.object({
+        messages: z4
+          .custom<BaseMessage[]>()
+          .register(registry, MessagesZodMeta),
+      }),
+
+      writer: writer<{ custom: string }>,
+      interrupt: interrupt<{ reason: string }, { messages: string[] }>,
+
+      // Allow optionally specifying all nodes upfront.
+      // This is needed to properly type `goto` commands and removes
+      // friction when dynamically constructing graphs.
+      nodes: ["hello", "what"],
+    }
+  );
+
+  const hello: typeof builder.Node = async (_, runtime) => {
+    // @ts-expect-error - Invalid interrupt input
+    if (false) runtime.interrupt?.({ reason: false });
+
+    // @ts-expect-error - Invalid writer input
+    if (false) runtime.interrupt?.();
+
+    const result = runtime.interrupt?.({ reason: "hello" });
+
+    if (result != null) {
+      runtime.writer?.({ custom: "hello" });
+
+      // @ts-expect-error - Invalid writer value
+      if (false) runtime.writer?.({ invalid: "hello" });
+
+      // @ts-expect-error - Invalid interrupt value
+      if (false) runtime.writer?.();
+
+      return new Command({
+        update: { messages: result.messages },
+        goto: "what",
+      });
+    }
+
+    return new Command({ goto: "what" });
+  };
+
+  const what: typeof builder.Node = async (_, runtime) => {
+    const result = runtime.interrupt?.({ reason: "what" });
+    if (result != null) return { messages: result.messages };
+    return {};
+  };
+
+  if (false) {
+    // @ts-expect-error - Invalid goto value
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    const invalidGoto: typeof builder.Node = async () => {
+      return new Command({ goto: "invalid" });
+    };
+  }
+
+  builder.addNode("hello", hello, { ends: [END, "what"] });
+  builder.addNode("what", what);
+  builder.addEdge(START, "hello");
+
+  const graph = builder.compile({ checkpointer });
+
+  const thread = graph.withConfig({ configurable: { thread_id: "1" } });
+  const first = await thread.invoke({ messages: "input" });
+
+  expect(first).toMatchObject({
+    messages: [{ text: "input" }],
+    __interrupt__: [{ id: expect.any(String), value: { reason: "hello" } }],
+  });
+
+  if (graph.isInterrupted(first)) {
+    expectTypeOf(first.__interrupt__).toExtend<
+      { id?: string; value?: { reason: string } }[]
+    >();
+  }
+
+  // @ts-expect-error - Invalid goto value
+  if (false) await thread.invoke(new Command({ goto: "xxx" }));
+
+  // @ts-expect-error - Invalid update value
+  if (false) await thread.invoke(new Command({ update: { messages: true } }));
+
+  const second = await gatherIterator(
+    thread.stream(new Command({ resume: { messages: ["resume: hello"] } }), {
+      streamMode: ["values", "custom"],
+    })
+  );
+
+  expect(second).toMatchObject([
+    ["values", { messages: [{ text: "input" }] }],
+    ["custom", { custom: "hello" }],
+    ["values", { messages: [{ text: "input" }, { text: "resume: hello" }] }],
+    [
+      "values",
+      {
+        __interrupt__: [{ id: expect.any(String), value: { reason: "what" } }],
+      },
+    ],
+  ]);
+
+  for await (const [mode, value] of second) {
+    if (mode === "values") {
+      if (graph.isInterrupted(value)) {
+        expectTypeOf(value.__interrupt__).toExtend<
+          { id?: string; value?: { reason: string } }[]
+        >();
+      }
+    }
+
+    if (mode === "custom") {
+      expectTypeOf(value).toExtend<{ custom: string }>();
+    }
+  }
+
+  const third = await thread.invoke(
+    new Command({ resume: { messages: ["resume: what"] } })
+  );
+
+  expect(third).toMatchObject({
+    messages: [
+      { text: "input" },
+      { text: "resume: hello" },
+      { text: "resume: what" },
+    ],
+  });
+});
+
 it("functional", async () => {
   const one = task("one", (input: string) => ({ one: `one(${input})` }));
   const two = task("two", (input: string) => ({ two: `two(${input})` }));

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

@@ -5080,7 +5080,7 @@ graph TD;
             id: "__start__",
             type: "schema",
             data: {
-              $schema: "http://json-schema.org/draft-07/schema#",
+              $schema: expect.any(String),
               title: undefined,
             },
           },
@@ -5104,7 +5104,7 @@ graph TD;
             id: "__end__",
             type: "schema",
             data: {
-              $schema: "http://json-schema.org/draft-07/schema#",
+              $schema: expect.any(String),
               title: undefined,
             },
           },

libs/langgraph/src/tests/typed_state.test.ts

@@ -1,95 +0,0 @@
-import { it, expect } from "vitest";
-import { z } from "zod";
-import { Command, START } from "../constants.js";
-import { Annotation } from "../graph/annotation.js";
-import {
-  MessagesAnnotation,
-  MessagesZodState,
-} from "../graph/messages_annotation.js";
-import { StateGraph, typedNode } from "../graph/state.js";
-import { _AnyIdHumanMessage } from "./utils.js";
-
-it("Annotation.Root", async () => {
-  const StateAnnotation = Annotation.Root({
-    messages: MessagesAnnotation.spec.messages,
-    foo: Annotation<string>,
-  });
-
-  const config = Annotation.Root({
-    random: Annotation<number>,
-  });
-
-  const node = typedNode(StateAnnotation, {
-    nodes: ["nodeA", "nodeB", "nodeC"],
-    config,
-  });
-
-  const nodeA = node(
-    (state) => {
-      const goto = state.foo === "foo" ? "nodeB" : "nodeC";
-      return new Command({
-        update: { messages: [{ type: "user", content: "a" }], foo: "a" },
-        goto,
-      });
-    },
-    { ends: ["nodeB", "nodeC"] }
-  );
-
-  const nodeB = node(
-    () => new Command({ goto: "nodeC", update: { foo: "123" } })
-  );
-  const nodeC = node(async (state) => ({ foo: `${state.foo}|c` }));
-
-  const graph = new StateGraph(StateAnnotation, config)
-    .addNode({ nodeA, nodeB, nodeC })
-    .addEdge(START, "nodeA")
-    .compile();
-
-  expect(await graph.invoke({ foo: "foo" })).toEqual({
-    messages: [new _AnyIdHumanMessage("a")],
-    foo: "123|c",
-  });
-});
-
-it("Zod", async () => {
-  const StateAnnotation = MessagesZodState.extend({
-    foo: z.string(),
-  });
-  const config = z.object({ random: z.number() });
-
-  const node = typedNode(StateAnnotation, {
-    nodes: ["nodeA", "nodeB", "nodeC"],
-    config,
-  });
-
-  const nodeA = node(
-    (state) => {
-      const goto = state.foo === "foo" ? "nodeB" : "nodeC";
-      return new Command({
-        update: { messages: [{ type: "user", content: "a" }], foo: "a" },
-        goto,
-      });
-    },
-    { ends: ["nodeB", "nodeC"] }
-  );
-
-  const nodeB = node(() => {
-    return new Command({
-      goto: "nodeC",
-      update: { foo: "123" },
-    });
-  });
-
-  const nodeC = node((state) => ({ foo: `${state.foo}|c` }));
-  const graph = new StateGraph(StateAnnotation, config)
-    .addNode({ nodeA, nodeB, nodeC })
-    .addEdge(START, "nodeA")
-    .compile();
-
-  expect(
-    await graph.invoke({ foo: "foo" }, { configurable: { random: 123 } })
-  ).toEqual({
-    messages: [new _AnyIdHumanMessage("a")],
-    foo: "123|c",
-  });
-});

libs/langgraph/src/ui/stream.ts

@@ -13,6 +13,7 @@ type StreamCheckpointsOutput<StreamValues> = StreamOutputMap<
   StreamValues,
   unknown,
   string,
+  unknown,
   unknown
 >;
 
@@ -143,6 +144,7 @@ export async function* toLangGraphEventStream<
           unknown,
           unknown,
           string,
+          unknown,
           unknown
         >;
 

libs/langgraph/src/web.ts

@@ -4,7 +4,6 @@ export {
   StateGraph,
   CompiledStateGraph,
   MessageGraph,
-  typedNode,
   messagesStateReducer,
   messagesStateReducer as addMessages,
   REMOVE_ALL_MESSAGES,

libs/langgraph/src/writer.ts

@@ -0,0 +1,17 @@
+import { AsyncLocalStorageProviderSingleton } from "@langchain/core/singletons";
+import { LangGraphRunnableConfig } from "./pregel/runnable_types.js";
+
+export function writer<T>(chunk: T): void {
+  const config: LangGraphRunnableConfig | undefined =
+    AsyncLocalStorageProviderSingleton.getRunnableConfig();
+  if (!config) {
+    throw new Error("Called interrupt() outside the context of a graph.");
+  }
+
+  const conf = config.configurable;
+  if (!conf) {
+    throw new Error("No configurable found in config");
+  }
+
+  return conf.writer?.(chunk);
+}