feat(langgraph): allow extending state with Zod schema (#1405)

.changeset/slow-camels-sip.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+feat(langgraph): allow extending state with Zod schema

libs/langgraph/src/prebuilt/react_agent_executor.ts

@@ -22,7 +22,10 @@ import {
   type RunnableLike,
 } from "@langchain/core/runnables";
 import { DynamicTool, StructuredToolInterface } from "@langchain/core/tools";
-import { InteropZodType } from "@langchain/core/utils/types";
+import type {
+  InteropZodObject,
+  InteropZodType,
+} from "@langchain/core/utils/types";
 import {
   All,
   BaseCheckpointSaver,
@@ -41,6 +44,7 @@ import { Annotation } from "../graph/annotation.js";
 import { Messages, messagesStateReducer } from "../graph/message.js";
 import { END, START } from "../constants.js";
 import { withAgentName } from "./agentName.js";
+import type { InteropZodToStateDefinition } from "../graph/zod/meta.js";
 
 export interface AgentState<
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -460,9 +464,18 @@ const PreHookAnnotation = Annotation.Root({
 
 type PreHookAnnotation = typeof PreHookAnnotation;
 
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type AnyAnnotationRoot = AnnotationRoot<any>;
+
+type ToAnnotationRoot<A extends AnyAnnotationRoot | InteropZodObject> =
+  A extends AnyAnnotationRoot
+    ? A
+    : A extends InteropZodObject
+    ? AnnotationRoot<InteropZodToStateDefinition<A>>
+    : never;
+
 export type CreateReactAgentParams<
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  A extends AnnotationRoot<any> = AnnotationRoot<any>,
+  A extends AnyAnnotationRoot | InteropZodObject = AnyAnnotationRoot,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   StructuredResponseType = Record<string, any>
 > = {
@@ -550,8 +563,8 @@ export type CreateReactAgentParams<
    * Useful for managing long message histories (e.g., message trimming, summarization, etc.).
    */
   preModelHook?: RunnableLike<
-    A["State"] & PreHookAnnotation["State"],
-    A["Update"] & PreHookAnnotation["Update"],
+    ToAnnotationRoot<A>["State"] & PreHookAnnotation["State"],
+    ToAnnotationRoot<A>["Update"] & PreHookAnnotation["Update"],
     LangGraphRunnableConfig
   >;
 
@@ -560,8 +573,8 @@ export type CreateReactAgentParams<
    * Useful for implementing human-in-the-loop, guardrails, validation, or other post-processing.
    */
   postModelHook?: RunnableLike<
-    A["State"],
-    A["Update"],
+    ToAnnotationRoot<A>["State"],
+    ToAnnotationRoot<A>["Update"],
     LangGraphRunnableConfig
   >;
 };
@@ -608,24 +621,22 @@ export type CreateReactAgentParams<
  * // Returns the messages in the state at each step of execution
  * ```
  */
-
 export function createReactAgent<
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/ban-types
-  A extends AnnotationRoot<any> = typeof MessagesAnnotation,
+  A extends AnyAnnotationRoot | InteropZodObject = typeof MessagesAnnotation,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   StructuredResponseFormat extends Record<string, any> = Record<string, any>
 >(
   params: CreateReactAgentParams<A, StructuredResponseFormat>
 ): CompiledStateGraph<
-  A["State"],
-  A["Update"],
+  ToAnnotationRoot<A>["State"],
+  ToAnnotationRoot<A>["Update"],
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   any,
-  typeof MessagesAnnotation.spec & A["spec"],
+  typeof MessagesAnnotation.spec & ToAnnotationRoot<A>["spec"],
   ReturnType<
     typeof createReactAgentAnnotation<StructuredResponseFormat>
   >["spec"] &
-    A["spec"]
+    ToAnnotationRoot<A>["spec"]
 > {
   const {
     llm,
@@ -757,7 +768,10 @@ export function createReactAgent<
   const schema =
     stateSchema ?? createReactAgentAnnotation<StructuredResponseFormat>();
 
-  const workflow = new StateGraph(schema).addNode("tools", toolNode);
+  const workflow = new StateGraph(schema as AnyAnnotationRoot).addNode(
+    "tools",
+    toolNode
+  );
 
   const allNodeWorkflows = workflow as WithStateGraphNodes<
     | "pre_model_hook"
@@ -774,7 +788,7 @@ export function createReactAgent<
   };
 
   let entrypoint: "agent" | "pre_model_hook" = "agent";
-  let inputSchema: AnnotationRoot<(typeof schema)["spec"]> | undefined;
+  let inputSchema: AnnotationRoot<ToAnnotationRoot<A>["spec"]> | undefined;
   if (preModelHook != null) {
     allNodeWorkflows
       .addNode("pre_model_hook", preModelHook)
@@ -782,7 +796,7 @@ export function createReactAgent<
     entrypoint = "pre_model_hook";
 
     inputSchema = Annotation.Root({
-      ...schema.spec,
+      ...workflow._schemaDefinition,
       ...PreHookAnnotation.spec,
     });
   } else {
@@ -799,7 +813,7 @@ export function createReactAgent<
       .addEdge("agent", "post_model_hook")
       .addConditionalEdges(
         "post_model_hook",
-        (state) => {
+        (state: AgentState<StructuredResponseFormat>) => {
           const { messages } = state;
           const lastMessage = messages[messages.length - 1];
 
@@ -830,7 +844,7 @@ export function createReactAgent<
   if (postModelHook == null) {
     allNodeWorkflows.addConditionalEdges(
       "agent",
-      (state) => {
+      (state: AgentState<StructuredResponseFormat>) => {
         const { messages } = state;
         const lastMessage = messages[messages.length - 1];
 
@@ -855,7 +869,7 @@ export function createReactAgent<
   if (shouldReturnDirect.size > 0) {
     allNodeWorkflows.addConditionalEdges(
       "tools",
-      (state) => {
+      (state: AgentState<StructuredResponseFormat>) => {
         // Check the last consecutive tool calls
         for (let i = state.messages.length - 1; i >= 0; i -= 1) {
           const message = state.messages[i];

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

@@ -12,7 +12,8 @@ import {
   SystemMessage,
   ToolMessage,
 } from "@langchain/core/messages";
-import { z } from "zod";
+import { z } from "zod/v3";
+import { z as z4 } from "zod/v4";
 import {
   Runnable,
   RunnableLambda,
@@ -49,9 +50,12 @@ import {
 } from "../index.js";
 import {
   MessagesAnnotation,
+  MessagesZodMeta,
   MessagesZodState,
 } from "../graph/messages_annotation.js";
 import { gatherIterator } from "../utils.js";
+import { withLangGraph } from "../graph/zod/meta.js";
+import { registry } from "../graph/zod/zod-registry.js";
 
 // Tracing slows down the tests
 beforeAll(() => {
@@ -252,18 +256,10 @@ describe("createReactAgent with prompt/state modifier", () => {
     const agent = createReactAgent({
       llm,
       tools: [
-        tool(
-          async () =>
-            new Command({
-              update: {
-                foo: "baz",
-              },
-            }),
-          {
-            name: "test",
-            schema: z.object({}),
-          }
-        ),
+        tool(async () => new Command({ update: { foo: "baz" } }), {
+          name: "test",
+          schema: z.object({}),
+        }),
       ],
       stateSchema: StateAnnotation,
     });
@@ -283,6 +279,48 @@ describe("createReactAgent with prompt/state modifier", () => {
     expect(result.foo).toEqual("baz");
   });
 
+  it("Allows custom Zod 4 state schema", async () => {
+    const llm = new FakeToolCallingChatModel({
+      responses: [
+        new AIMessage({
+          content: "result1",
+          tool_calls: [{ id: "test1234", args: {}, name: "test" }],
+        }),
+        new AIMessage("result2"),
+      ],
+    });
+
+    const agent = createReactAgent({
+      llm,
+      tools: [
+        tool(async () => new Command({ update: { foo: "baz" } }), {
+          name: "test",
+          schema: z.object({}),
+        }),
+      ],
+      stateSchema: z4.object({
+        messages: z4
+          .custom<BaseMessage[]>()
+          .register(registry, MessagesZodMeta),
+        foo: z4.string(),
+      }),
+    });
+
+    const result = await agent.invoke({
+      messages: [],
+      foo: "bar",
+    });
+    const expected = [
+      new _AnyIdAIMessage({
+        content: "result1",
+        tool_calls: [{ id: "test1234", args: {}, name: "test" }],
+      }),
+      new _AnyIdAIMessage("result2"),
+    ];
+    expect(result.messages).toEqual(expected);
+    expect(result.foo).toEqual("baz");
+  });
+
   it("Should respect a passed signal", async () => {
     const llm = new FakeToolCallingChatModel({
       responses: [
@@ -1114,6 +1152,98 @@ describe("createReactAgent with hooks", () => {
     );
   });
 
+  it("preModelHook + Zod 3 + postModelHook", async () => {
+    const llm = new FakeToolCallingChatModel({
+      responses: [new AIMessage({ id: "0", content: "Hello!" })],
+    });
+
+    const llmSpy = vi.spyOn(llm, "_generate");
+
+    const agent = createReactAgent({
+      llm,
+      tools: [],
+      preModelHook: () => ({
+        llmInputMessages: [
+          new HumanMessage({ id: "human", content: "pre-hook" }),
+        ],
+      }),
+      stateSchema: z.object({
+        messages: withLangGraph(z.custom<BaseMessage[]>(), MessagesZodMeta),
+        flag: withLangGraph(z.boolean(), {
+          reducer: {
+            fn: (a, b) => [a, ...b].reduce((acc, curr) => acc || curr, false),
+            schema: z.array(z.boolean()),
+          },
+          default: () => false,
+        }),
+      }),
+      postModelHook: () => ({ flag: [false, false, true] }),
+    });
+
+    expect(await agent.invoke({ messages: [new HumanMessage("hi?")] })).toEqual(
+      {
+        messages: [
+          new _AnyIdHumanMessage("hi?"),
+          new AIMessage({ id: "0", content: "Hello!" }),
+        ],
+        flag: true,
+      }
+    );
+
+    expect(llmSpy).toHaveBeenCalledWith(
+      [new HumanMessage({ id: "human", content: "pre-hook" })],
+      expect.anything(),
+      undefined
+    );
+  });
+
+  it("preModelHook + Zod 4 + postModelHook", async () => {
+    const llm = new FakeToolCallingChatModel({
+      responses: [new AIMessage({ id: "0", content: "Hello!" })],
+    });
+
+    const llmSpy = vi.spyOn(llm, "_generate");
+
+    const agent = createReactAgent({
+      llm,
+      tools: [],
+      preModelHook: () => ({
+        llmInputMessages: [
+          new HumanMessage({ id: "human", content: "pre-hook" }),
+        ],
+      }),
+      stateSchema: z4.object({
+        messages: z4
+          .custom<BaseMessage[]>()
+          .register(registry, MessagesZodMeta),
+        flag: z4.boolean().register(registry, {
+          reducer: {
+            fn: (a, b) => [a, ...b].reduce((acc, curr) => acc || curr, false),
+            schema: z.array(z.boolean()),
+          },
+          default: () => false,
+        }),
+      }),
+      postModelHook: () => ({ flag: [false, false, true] }),
+    });
+
+    expect(await agent.invoke({ messages: [new HumanMessage("hi?")] })).toEqual(
+      {
+        messages: [
+          new _AnyIdHumanMessage("hi?"),
+          new AIMessage({ id: "0", content: "Hello!" }),
+        ],
+        flag: true,
+      }
+    );
+
+    expect(llmSpy).toHaveBeenCalledWith(
+      [new HumanMessage({ id: "human", content: "pre-hook" })],
+      expect.anything(),
+      undefined
+    );
+  });
+
   it("postModelHook", async () => {
     const FlagAnnotation = Annotation.Root({
       ...MessagesAnnotation.spec,