feat(langgraph): use return types from nodes as update type (#1502)

.changeset/fuzzy-vans-yell.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+Use return type of nodes for streamMode: updates types

.changeset/seven-lamps-matter.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph-cua": minor
+---
+
+feat(cua): Improve update types

.changeset/slimy-singers-trade.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph-swarm": patch
+---
+
+chore: enforce return types for swarm

libs/langgraph-cua/src/index.ts

@@ -169,8 +169,12 @@ export function createCua<
     throw new Error("timeoutHours must be between 0.01 and 24");
   }
 
-  const nodeBefore = nodeBeforeAction ?? (async () => {});
-  const nodeAfter = nodeAfterAction ?? (async () => {});
+  const nodeBefore =
+    nodeBeforeAction ??
+    ((async () => ({})) as (state: CUAState) => Promise<CUAUpdate>);
+  const nodeAfter =
+    nodeAfterAction ??
+    ((async () => ({})) as (state: CUAState) => Promise<CUAUpdate>);
 
   const StateAnnotation = Annotation.Root({
     ...CUAAnnotation.spec,

libs/langgraph-cua/src/nodes/call-model.ts

@@ -10,7 +10,6 @@ import { ChatOpenAI } from "@langchain/openai";
 import {
   CUAEnvironment,
   CUAState,
-  CUAUpdate,
   getConfigurationWithDefaults,
 } from "../types.js";
 import { isComputerCallToolMessage } from "../utils.js";
@@ -102,7 +101,7 @@ const _promptToSysMessage = (prompt: string | SystemMessage | undefined) => {
 export async function callModel(
   state: CUAState,
   config: LangGraphRunnableConfig
-): Promise<CUAUpdate> {
+) {
   const configuration = getConfigurationWithDefaults(config);
 
   const lastMessage = state.messages[state.messages.length - 1];
@@ -132,20 +131,29 @@ export async function callModel(
       previous_response_id: previousResponseId,
     });
 
-  let response: AIMessageChunk;
+  type AIMessageChunkOpenAI = AIMessageChunk & {
+    additional_kwargs: {
+      tool_outputs?: {
+        call_id: string;
+        action: { type: string; display_width: number; display_height: number };
+      }[];
+    };
+  };
+
+  let response: AIMessageChunkOpenAI;
   if (isLastMessageComputerCallOutput && !configuration.zdrEnabled) {
     const formattedMessage =
       await conditionallyUpdateToolMessageContentRunnable.invoke(lastMessage);
-    response = await model.invoke([formattedMessage]);
+    response = (await model.invoke([formattedMessage])) as AIMessageChunkOpenAI;
   } else {
     const formattedMessagesPromise = state.messages.map((m) =>
       conditionallyUpdateToolMessageContentRunnable.invoke(m)
     );
     const prompt = _promptToSysMessage(configuration.prompt);
-    response = await model.invoke([
+    response = (await model.invoke([
       ...(prompt ? [prompt] : []),
       ...(await Promise.all(formattedMessagesPromise)),
-    ]);
+    ])) as AIMessageChunkOpenAI;
   }
 
   return {

libs/langgraph-cua/src/nodes/create-vm-instance.ts

@@ -1,12 +1,12 @@
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
 import { UbuntuInstance, BrowserInstance, WindowsInstance } from "scrapybara";
-import { CUAState, CUAUpdate, getConfigurationWithDefaults } from "../types.js";
+import { CUAState, getConfigurationWithDefaults } from "../types.js";
 import { getScrapybaraClient } from "../utils.js";
 
 export async function createVMInstance(
   state: CUAState,
   config: LangGraphRunnableConfig
-): Promise<CUAUpdate> {
+) {
   const { instanceId } = state;
   if (instanceId) {
     // Instance already exists, no need to initialize
@@ -46,13 +46,8 @@ export async function createVMInstance(
     // If the streamUrl is not yet defined in state, fetch it, then write to the custom stream
     // so that it's made accessible to the client (or whatever is reading the stream) before any actions are taken.
     const { streamUrl } = await instance.getStreamUrl();
-    return {
-      instanceId: instance.id,
-      streamUrl,
-    };
+    return { instanceId: instance.id, streamUrl };
   }
 
-  return {
-    instanceId: instance.id,
-  };
+  return { instanceId: instance.id };
 }

libs/langgraph-cua/src/nodes/take-computer-action.ts

@@ -5,9 +5,9 @@ import {
   Scrapybara,
 } from "scrapybara";
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
-import { BaseMessageLike } from "@langchain/core/messages";
+import { ToolMessage } from "@langchain/core/messages";
 import { RunnableLambda } from "@langchain/core/runnables";
-import { CUAState, CUAUpdate, getConfigurationWithDefaults } from "../types.js";
+import { CUAState, getConfigurationWithDefaults } from "../types.js";
 import { getInstance, getToolOutputs } from "../utils.js";
 
 async function sleep(ms: number): Promise<void> {
@@ -52,7 +52,7 @@ export async function takeComputerAction(
   {
     uploadScreenshot,
   }: { uploadScreenshot?: (screenshot: string) => Promise<string> }
-): Promise<CUAUpdate> {
+) {
   if (!state.instanceId) {
     throw new Error("Can not take computer action without an instance ID.");
   }
@@ -93,7 +93,7 @@ export async function takeComputerAction(
 
   const output = toolOutputs[toolOutputs.length - 1];
   const { action } = output;
-  let computerCallToolMsg: BaseMessageLike | undefined;
+  let computerCallToolMsg: ToolMessage | undefined;
 
   try {
     let computerResponse: Scrapybara.ComputerResponse;
@@ -180,18 +180,14 @@ export async function takeComputerAction(
       );
     }
 
-    computerCallToolMsg = {
-      type: "tool",
+    computerCallToolMsg = new ToolMessage({
+      content: screenshotContent,
       tool_call_id: output.call_id,
       additional_kwargs: { type: "computer_call_output" },
-      content: screenshotContent,
-    };
+    });
   } catch (e) {
     console.error(
-      {
-        error: e,
-        computerCall: output,
-      },
+      { error: e, computerCall: output },
       "Failed to execute computer call."
     );
   }

libs/langgraph-swarm/src/swarm.ts

@@ -130,9 +130,7 @@ const createSwarm = <
     agentNames.add(agent.name);
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  const builder = new StateGraph<any>(stateSchema ?? SwarmState);
-
+  const builder = new StateGraph(stateSchema ?? SwarmState);
   addActiveAgentRouter(builder, {
     routeTo: [...agentNames],
     defaultActiveAgent,
@@ -145,7 +143,14 @@ const createSwarm = <
     });
   }
 
-  return builder;
+  return builder as StateGraph<
+    AnnotationRootT["spec"],
+    AnnotationRootT["State"],
+    AnnotationRootT["Update"],
+    string,
+    AnnotationRootT["spec"],
+    AnnotationRootT["spec"]
+  >;
 };
 
 export { createSwarm, addActiveAgentRouter, SwarmState };

libs/langgraph/src/graph/graph.ts

@@ -222,35 +222,35 @@ export class Graph<
     return this.edges;
   }
 
-  addNode<K extends string>(
+  addNode<K extends string, NodeInput = RunInput, NodeOutput = RunOutput>(
     nodes:
-      | Record<K, NodeAction<RunInput, RunOutput, C>>
+      | Record<K, NodeAction<NodeInput, NodeOutput, C>>
       | [
           key: K,
-          action: NodeAction<RunInput, RunOutput, C>,
+          action: NodeAction<NodeInput, NodeOutput, C>,
           options?: AddNodeOptions
         ][]
   ): Graph<N | K, RunInput, RunOutput>;
 
-  addNode<K extends string, NodeInput = RunInput>(
+  addNode<K extends string, NodeInput = RunInput, NodeOutput = RunOutput>(
     key: K,
-    action: NodeAction<NodeInput, RunOutput, C>,
+    action: NodeAction<NodeInput, NodeOutput, C>,
     options?: AddNodeOptions
   ): Graph<N | K, RunInput, RunOutput>;
 
-  addNode<K extends string, NodeInput = RunInput>(
+  addNode<K extends string, NodeInput = RunInput, NodeOutput = RunOutput>(
     ...args:
       | [
           key: K,
-          action: NodeAction<NodeInput, RunOutput, C>,
+          action: NodeAction<NodeInput, NodeOutput, C>,
           options?: AddNodeOptions
         ]
       | [
           nodes:
-            | Record<K, NodeAction<NodeInput, RunOutput, C>>
+            | Record<K, NodeAction<NodeInput, NodeOutput, C>>
             | [
                 key: K,
-                action: NodeAction<NodeInput, RunOutput, C>,
+                action: NodeAction<NodeInput, NodeOutput, C>,
                 options?: AddNodeOptions
               ][]
         ]
@@ -580,14 +580,18 @@ export class CompiledGraph<
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   InputType = any,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  OutputType = any
+  OutputType = any,
+  NodeReturnType = unknown
 > extends Pregel<
   Record<N | typeof START, PregelNode<State, Update>>,
   Record<N | typeof START | typeof END | string, BaseChannel>,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   ContextType & Record<string, any>,
   InputType,
-  OutputType
+  OutputType,
+  InputType,
+  OutputType,
+  NodeReturnType
 > {
   declare NodeType: N;
 

libs/langgraph/src/graph/state.ts

@@ -140,6 +140,15 @@ type NodeAction<S, U, C extends SDZod> = RunnableLike<
 const PartialStateSchema = Symbol.for("langgraph.state.partial");
 type PartialStateSchema = typeof PartialStateSchema;
 
+type MergeReturnType<Prev, Curr> = Prev & Curr extends infer T
+  ? { [K in keyof T]: T[K] } & unknown
+  : never;
+
+type Prettify<T> = {
+  [K in keyof T]: T[K];
+  // eslint-disable-next-line @typescript-eslint/ban-types
+} & {};
+
 /**
  * A graph whose nodes communicate by reading and writing to a shared state.
  * Each node takes a defined `State` as input and returns a `Partial<State>`.
@@ -209,7 +218,8 @@ export class StateGraph<
   N extends string = typeof START,
   I extends SDZod = SD extends SDZod ? ToStateDefinition<SD> : StateDefinition,
   O extends SDZod = SD extends SDZod ? ToStateDefinition<SD> : StateDefinition,
-  C extends SDZod = StateDefinition
+  C extends SDZod = StateDefinition,
+  NodeReturnType = unknown
 > extends Graph<N, S, U, StateGraphNodeSpec<S, U>, ToStateDefinition<C>> {
   channels: Record<string, BaseChannel> = {};
 
@@ -408,27 +418,73 @@ export class StateGraph<
     }
   }
 
-  override addNode<K extends string>(
+  override addNode<
+    K extends string,
+    NodeMap extends Record<K, NodeAction<S, U, C>>
+  >(
+    nodes: NodeMap
+  ): StateGraph<
+    SD,
+    S,
+    U,
+    N | K,
+    I,
+    O,
+    C,
+    MergeReturnType<
+      NodeReturnType,
+      {
+        [key in keyof NodeMap]: NodeMap[key] extends NodeAction<S, infer U, C>
+          ? U
+          : never;
+      }
+    >
+  >;
+
+  override addNode<K extends string, NodeInput = S, NodeOutput extends U = U>(
     nodes:
-      | Record<K, NodeAction<S, U, C>>
       | [
           key: K,
-          action: NodeAction<S, U, C>,
+          action: NodeAction<NodeInput, NodeOutput, C>,
           options?: StateGraphAddNodeOptions
         ][]
-  ): StateGraph<SD, S, U, N | K, I, O, C>;
+  ): StateGraph<
+    SD,
+    S,
+    U,
+    N | K,
+    I,
+    O,
+    C,
+    MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
+  >;
+
+  override addNode<K extends string, NodeInput = S, NodeOutput extends U = U>(
+    key: K,
+    action: NodeAction<NodeInput, NodeOutput, C>,
+    options?: StateGraphAddNodeOptions
+  ): StateGraph<
+    SD,
+    S,
+    U,
+    N | K,
+    I,
+    O,
+    C,
+    MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
+  >;
 
   override addNode<K extends string, NodeInput = S>(
     key: K,
     action: NodeAction<NodeInput, U, C>,
     options?: StateGraphAddNodeOptions
-  ): StateGraph<SD, S, U, N | K, I, O, C>;
+  ): StateGraph<SD, S, U, N | K, I, O, C, NodeReturnType>;
 
-  override addNode<K extends string, NodeInput = S>(
+  override addNode<K extends string, NodeInput = S, NodeOutput extends U = U>(
     ...args:
       | [
           key: K,
-          action: NodeAction<NodeInput, U, C>,
+          action: NodeAction<NodeInput, NodeOutput, C>,
           options?: StateGraphAddNodeOptions
         ]
       | [
@@ -589,27 +645,61 @@ export class StateGraph<
     return this;
   }
 
-  addSequence<K extends string>(
+  addSequence<K extends string, NodeInput = S, NodeOutput extends U = U>(
     nodes: [
       key: K,
-      action: NodeAction<S, U, C>,
+      action: NodeAction<NodeInput, NodeOutput, C>,
       options?: StateGraphAddNodeOptions
     ][]
-  ): StateGraph<SD, S, U, N | K, I, O, C>;
+  ): StateGraph<
+    SD,
+    S,
+    U,
+    N | K,
+    I,
+    O,
+    C,
+    MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
+  >;
 
-  addSequence<K extends string>(
-    nodes: Record<K, NodeAction<S, U, C>>
-  ): StateGraph<SD, S, U, N | K, I, O, C>;
+  addSequence<K extends string, NodeMap extends Record<K, NodeAction<S, U, C>>>(
+    nodes: NodeMap
+  ): StateGraph<
+    SD,
+    S,
+    U,
+    N | K,
+    I,
+    O,
+    C,
+    MergeReturnType<
+      NodeReturnType,
+      {
+        [key in keyof NodeMap]: NodeMap[key] extends NodeAction<S, infer U, C>
+          ? U
+          : never;
+      }
+    >
+  >;
 
-  addSequence<K extends string>(
+  addSequence<K extends string, NodeInput = S, NodeOutput extends U = U>(
     nodes:
       | [
           key: K,
-          action: NodeAction<S, U, C>,
+          action: NodeAction<NodeInput, NodeOutput, C>,
           options?: StateGraphAddNodeOptions
         ][]
-      | Record<K, NodeAction<S, U, C>>
-  ): StateGraph<SD, S, U, N | K, I, O, C> {
+      | Record<K, NodeAction<NodeInput, NodeOutput, C>>
+  ): StateGraph<
+    SD,
+    S,
+    U,
+    N | K,
+    I,
+    O,
+    C,
+    MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
+  > {
     const parsedNodes = Array.isArray(nodes)
       ? nodes
       : (Object.entries(nodes).map(([key, action]) => [
@@ -632,7 +722,7 @@ export class StateGraph<
       }
 
       const validKey = key as unknown as N;
-      this.addNode(validKey, action, options);
+      this.addNode(validKey, action as NodeAction<S, U, C>, options);
       if (previousNode != null) {
         this.addEdge(previousNode, validKey);
       }
@@ -640,7 +730,16 @@ export class StateGraph<
       previousNode = validKey;
     }
 
-    return this as StateGraph<SD, S, U, N | K, I, O, C>;
+    return this as StateGraph<
+      SD,
+      S,
+      U,
+      N | K,
+      I,
+      O,
+      C,
+      MergeReturnType<NodeReturnType, { [key in K]: NodeOutput }>
+    >;
   }
 
   override compile({
@@ -657,7 +756,15 @@ export class StateGraph<
     interruptBefore?: N[] | All;
     interruptAfter?: N[] | All;
     name?: string;
-  } = {}): CompiledStateGraph<S, U, N, I, O, C> {
+  } = {}): CompiledStateGraph<
+    Prettify<S>,
+    Prettify<U>,
+    N,
+    I,
+    O,
+    C,
+    NodeReturnType
+  > {
     // validate the graph
     this.validate([
       ...(Array.isArray(interruptBefore) ? interruptBefore : []),
@@ -676,7 +783,7 @@ export class StateGraph<
       streamKeys.length === 1 && streamKeys[0] === ROOT ? ROOT : streamKeys;
 
     // create empty compiled graph
-    const compiled = new CompiledStateGraph<S, U, N, I, O, C>({
+    const compiled = new CompiledStateGraph<S, U, N, I, O, C, NodeReturnType>({
       builder: this,
       checkpointer,
       interruptAfter,
@@ -760,16 +867,18 @@ export class CompiledStateGraph<
   N extends string = typeof START,
   I extends SDZod = StateDefinition,
   O extends SDZod = StateDefinition,
-  C extends SDZod = StateDefinition
+  C extends SDZod = StateDefinition,
+  NodeReturnType = unknown
 > extends CompiledGraph<
   N,
   S,
   U,
   StateType<ToStateDefinition<C>>,
   UpdateType<ToStateDefinition<I>>,
-  StateType<ToStateDefinition<O>>
+  StateType<ToStateDefinition<O>>,
+  NodeReturnType
 > {
-  declare builder: StateGraph<unknown, S, U, N, I, O, C>;
+  declare builder: StateGraph<unknown, S, U, N, I, O, C, NodeReturnType>;
 
   /** @internal */
   _metaRegistry: SchemaMetaRegistry = schemaMetaRegistry;

libs/langgraph/src/pregel/index.ts

@@ -385,7 +385,8 @@ export class Pregel<
     InputType = PregelInputType,
     OutputType = PregelOutputType,
     StreamUpdatesType = InputType,
-    StreamValuesType = OutputType
+    StreamValuesType = OutputType,
+    NodeReturnType = unknown
   >
   extends PartialRunnable<
     InputType | CommandInstance | null,
@@ -1824,7 +1825,8 @@ export class Pregel<
         TSubgraphs,
         StreamUpdatesType,
         StreamValuesType,
-        keyof Nodes
+        keyof Nodes,
+        NodeReturnType
       >
     >
   > {
@@ -1849,7 +1851,8 @@ export class Pregel<
           TSubgraphs,
           StreamUpdatesType,
           StreamValuesType,
-          keyof Nodes
+          keyof Nodes,
+          NodeReturnType
         >
       >,
       abortController

libs/langgraph/src/pregel/types.ts

@@ -82,7 +82,8 @@ export type StreamOutputMap<
   TStreamSubgraphs extends boolean,
   StreamUpdates,
   StreamValues,
-  Nodes
+  Nodes,
+  NodeReturnType
 > = (
   undefined extends TStreamMode
     ? []
@@ -100,7 +101,9 @@ export type StreamOutputMap<
         updates: [
           string[],
           "updates",
-          Record<Nodes extends string ? Nodes : string, StreamUpdates>
+          NodeReturnType extends Record<string, unknown>
+            ? { [K in keyof NodeReturnType]?: NodeReturnType[K] }
+            : Record<Nodes extends string ? Nodes : string, StreamUpdates>
         ];
         messages: [string[], "messages", StreamMessageOutput];
         custom: [string[], "custom", StreamCustomOutput];
@@ -120,7 +123,9 @@ export type StreamOutputMap<
         values: ["values", StreamValues];
         updates: [
           "updates",
-          Record<Nodes extends string ? Nodes : string, StreamUpdates>
+          NodeReturnType extends Record<string, unknown>
+            ? { [K in keyof NodeReturnType]?: NodeReturnType[K] }
+            : Record<Nodes extends string ? Nodes : string, StreamUpdates>
         ];
         messages: ["messages", StreamMessageOutput];
         custom: ["custom", StreamCustomOutput];
@@ -136,7 +141,9 @@ export type StreamOutputMap<
         values: [string[], StreamValues];
         updates: [
           string[],
-          Record<Nodes extends string ? Nodes : string, StreamUpdates>
+          NodeReturnType extends Record<string, unknown>
+            ? { [K in keyof NodeReturnType]?: NodeReturnType[K] }
+            : Record<Nodes extends string ? Nodes : string, StreamUpdates>
         ];
         messages: [string[], StreamMessageOutput];
         custom: [string[], StreamCustomOutput];
@@ -149,7 +156,9 @@ export type StreamOutputMap<
       }[Single]
     : {
         values: StreamValues;
-        updates: Record<Nodes extends string ? Nodes : string, StreamUpdates>;
+        updates: NodeReturnType extends Record<string, unknown>
+          ? { [K in keyof NodeReturnType]?: NodeReturnType[K] }
+          : Record<Nodes extends string ? Nodes : string, StreamUpdates>;
         messages: StreamMessageOutput;
         custom: StreamCustomOutput;
         checkpoints: StreamCheckpointsOutput<StreamValues>;

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

@@ -43,7 +43,7 @@ it("StateGraph bad return type", async () => {
 
   const graph = new StateGraph(StateAnnotation)
     // @ts-expect-error Test invalid return value
-    .addNode("foo", () => "hey")
+    .addNode("foo", () => 123)
     .addEdge("__start__", "foo");
 
   const app = graph.compile({ checkpointer: new MemorySaverAssertImmutable() });

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

@@ -33,9 +33,9 @@ it("state graph annotation", async () => {
 
   const graph = new StateGraph(state)
     .addSequence({
-      one: () => ({ foo: "one" }),
-      two: () => ({ foo: "two" }),
-      three: () => ({ foo: "three" }),
+      one: () => ({ foo: "one" as const }),
+      two: () => ({ foo: "two" as const }),
+      three: () => ({ foo: "three" as const }),
     })
     .addEdge("__start__", "one")
     .compile();
@@ -50,7 +50,7 @@ it("state graph annotation", async () => {
   }
 
   expectTypeOf(await gatherIterator(graph.stream(input))).toExtend<
-    Record<"one" | "two" | "three", { foo?: string[] | string }>[]
+    { one?: { foo: "one" }; two?: { foo: "two" }; three?: { foo: "three" } }[]
   >();
 
   expectTypeOf(
@@ -78,12 +78,25 @@ it("state graph annotation", async () => {
 
   expectTypeOf(
     await gatherIterator(graph.stream(input, { streamMode: "updates" }))
-  ).toExtend<Record<"one" | "two" | "three", { foo?: string[] | string }>[]>();
+  ).toExtend<
+    {
+      one?: { foo: "one" };
+      two?: { foo: "two" };
+      three?: { foo: "three" };
+    }[]
+  >();
 
   expectTypeOf(
     await gatherIterator(graph.stream(input, { streamMode: ["updates"] }))
   ).toExtend<
-    ["updates", Record<"one" | "two" | "three", { foo?: string[] | string }>][]
+    [
+      "updates",
+      {
+        one?: { foo: "one" };
+        two?: { foo: "two" };
+        three?: { foo: "three" };
+      }
+    ][]
   >();
 
   expectTypeOf(
@@ -97,7 +110,11 @@ it("state graph annotation", async () => {
     [
       string[],
       "updates",
-      Record<"one" | "two" | "three", { foo?: string[] | string }>
+      {
+        one?: { foo: "one" };
+        two?: { foo: "two" };
+        three?: { foo: "three" };
+      }
     ][]
   >();
 
@@ -109,7 +126,11 @@ it("state graph annotation", async () => {
     (
       | [
           "updates",
-          Record<"one" | "two" | "three", { foo?: string[] | string }>
+          {
+            one?: { foo: "one" };
+            two?: { foo: "two" };
+            three?: { foo: "three" };
+          }
         ]
       | ["values", { foo: string[] }]
     )[]
@@ -127,7 +148,11 @@ it("state graph annotation", async () => {
       | [
           string[],
           "updates",
-          Record<"one" | "two" | "three", { foo?: string[] | string }>
+          {
+            one?: { foo: "one" };
+            two?: { foo: "two" };
+            three?: { foo: "three" };
+          }
         ]
       | [string[], "values", { foo: string[] }]
     )[]
@@ -147,7 +172,11 @@ it("state graph annotation", async () => {
     (
       | [
           "updates",
-          Record<"one" | "two" | "three", { foo?: string[] | string }>
+          {
+            one?: { foo: "one" };
+            two?: { foo: "two" };
+            three?: { foo: "three" };
+          }
         ]
       | ["values", { foo: string[] }]
       | ["debug", Record<string, any>]
@@ -308,9 +337,9 @@ it("state graph zod", async () => {
 
   const graph = new StateGraph(state)
     .addSequence({
-      one: () => ({ foo: "one" }),
-      two: () => ({ foo: "two" }),
-      three: () => ({ foo: "three" }),
+      one: () => ({ foo: "one" as const }),
+      two: () => ({ foo: "two" as const }),
+      three: () => ({ foo: "three" as const }),
     })
     .addEdge("__start__", "one")
     .compile();
@@ -325,7 +354,7 @@ it("state graph zod", async () => {
   }
 
   expectTypeOf(await gatherIterator(graph.stream(input))).toExtend<
-    Record<"one" | "two" | "three", { foo?: string[] | string }>[]
+    { one?: { foo: "one" }; two?: { foo: "two" }; three?: { foo: "three" } }[]
   >();
 
   expectTypeOf(
@@ -350,12 +379,17 @@ it("state graph zod", async () => {
 
   expectTypeOf(
     await gatherIterator(graph.stream(input, { streamMode: "updates" }))
-  ).toExtend<Record<"one" | "two" | "three", { foo?: string[] | string }>[]>();
+  ).toExtend<
+    { one?: { foo: "one" }; two?: { foo: "two" }; three?: { foo: "three" } }[]
+  >();
 
   expectTypeOf(
     await gatherIterator(graph.stream(input, { streamMode: ["updates"] }))
   ).toExtend<
-    ["updates", Record<"one" | "two" | "three", { foo?: string[] | string }>][]
+    [
+      "updates",
+      { one?: { foo: "one" }; two?: { foo: "two" }; three?: { foo: "three" } }
+    ][]
   >();
 
   expectTypeOf(
@@ -366,7 +400,7 @@ it("state graph zod", async () => {
     [
       string[],
       "updates",
-      Record<"one" | "two" | "three", { foo?: string[] | string }>
+      { one?: { foo: "one" }; two?: { foo: "two" }; three?: { foo: "three" } }
     ][]
   >();
 
@@ -378,7 +412,11 @@ it("state graph zod", async () => {
     (
       | [
           "updates",
-          Record<"one" | "two" | "three", { foo?: string[] | string }>
+          {
+            one?: { foo: "one" };
+            two?: { foo: "two" };
+            three?: { foo: "three" };
+          }
         ]
       | ["values", { foo: string[] }]
     )[]
@@ -396,7 +434,11 @@ it("state graph zod", async () => {
       | [
           string[],
           "updates",
-          Record<"one" | "two" | "three", { foo?: string[] | string }>
+          {
+            one?: { foo: "one" };
+            two?: { foo: "two" };
+            three?: { foo: "three" };
+          }
         ]
       | [string[], "values", { foo: string[] }]
     )[]
@@ -416,7 +458,11 @@ it("state graph zod", async () => {
     (
       | [
           "updates",
-          Record<"one" | "two" | "three", { foo?: string[] | string }>
+          {
+            one?: { foo: "one" };
+            two?: { foo: "two" };
+            three?: { foo: "three" };
+          }
         ]
       | ["values", { foo: string[] }]
       | ["debug", Record<string, any>]
@@ -448,7 +494,11 @@ it("state graph zod", async () => {
       | [
           string[],
           "updates",
-          Record<"one" | "two" | "three", { foo?: string[] | string }>
+          {
+            one?: { foo: "one" };
+            two?: { foo: "two" };
+            three?: { foo: "three" };
+          }
         ]
       | [string[], "values", { foo: string[] }]
       | [string[], "debug", Record<string, any>]