fix(langchain/agents): use channels when defining StateGraph schema (#9358)

Author: hntrl

libs/langchain/src/agents/ReactAgent.ts

@@ -3,7 +3,6 @@
 import { InteropZodObject } from "@langchain/core/utils/types";
 
 import {
-  AnnotationRoot,
   StateGraph,
   END,
   START,
@@ -185,7 +184,7 @@ export class ReactAgent<
     );
 
     const workflow = new StateGraph(
-      schema as unknown as AnnotationRoot<any>,
+      schema as unknown as AnyAnnotationRoot,
       this.options.contextSchema
     );
 

libs/langchain/src/agents/annotation.ts

@@ -7,7 +7,7 @@ import {
   MessagesZodMeta,
   type BinaryOperatorAggregate,
 } from "@langchain/langgraph";
-import { withLangGraph } from "@langchain/langgraph/zod";
+import { withLangGraph, schemaMetaRegistry } from "@langchain/langgraph/zod";
 
 import type { AgentMiddleware, AnyAnnotationRoot } from "./middleware/types.js";
 import { InteropZodObject } from "@langchain/core/utils/types";
@@ -27,7 +27,7 @@ export function createAgentAnnotationConditional<
    * Create Zod schema object to preserve jsonSchemaExtra
    * metadata for LangGraph Studio using v3-compatible withLangGraph
    */
-  const zodSchema: Record<string, any> = {
+  const schemaShape: Record<string, any> = {
     messages: withLangGraph(z.custom<BaseMessage[]>(), MessagesZodMeta),
     jumpTo: z
       .union([
@@ -49,8 +49,8 @@ export function createAgentAnnotationConditional<
         continue;
       }
 
-      if (!(key in zodSchema)) {
-        zodSchema[key] = schema;
+      if (!(key in schemaShape)) {
+        schemaShape[key] = schema;
       }
     }
   };
@@ -68,10 +68,12 @@ export function createAgentAnnotationConditional<
 
   // Only include structuredResponse when responseFormat is defined
   if (hasStructuredResponse) {
-    zodSchema.structuredResponse = z.any().optional();
+    schemaShape.structuredResponse = z.any().optional();
   }
 
-  return z.object(zodSchema);
+  const zodSchema = z.object(schemaShape);
+  const stateDefinition = schemaMetaRegistry.getChannelsForSchema(zodSchema);
+  return stateDefinition;
 }
 
 export const PreHookAnnotation: AnnotationRoot<{

libs/langchain/src/agents/tests/reactAgent.test.ts

@@ -1,5 +1,6 @@
 import { describe, it, expect, beforeEach } from "vitest";
 import { z } from "zod/v3";
+import { z as z4 } from "zod/v4";
 
 import {
   BaseMessage,
@@ -17,7 +18,12 @@ import {
   type BaseCheckpointSaver,
 } from "@langchain/langgraph";
 
-import { providerStrategy, createAgent, createMiddleware } from "../index.js";
+import {
+  providerStrategy,
+  createAgent,
+  createMiddleware,
+  toolStrategy,
+} from "../index.js";
 
 import {
   FakeToolCallingChatModel,
@@ -878,4 +884,128 @@ describe("createAgent", () => {
     const taskInput = JSON.parse(toolMessage.content as string);
     expect(taskInput.customField).toBe("test-value");
   });
+
+  // https://github.com/langchain-ai/langchainjs/issues/9299
+  it("supports zod 3/4 schemas in createAgent and middleware", async () => {
+    // Create middleware with Zod v3 schemas
+    const middleware1 = createMiddleware({
+      name: "middleware1",
+      stateSchema: z.object({
+        middleware1Value: z.string().default("v3-default"),
+      }),
+      contextSchema: z.object({
+        middleware1Context: z.number(),
+      }),
+      beforeModel: (_state, { context }) => {
+        expect(context.middleware1Context).toBe(42);
+        return {
+          middleware1Value: "v3-modified",
+        };
+      },
+    });
+
+    // Create middleware with Zod v4 schemas
+    const middleware2 = createMiddleware({
+      name: "middleware2",
+      stateSchema: z4.object({
+        middleware2Value: z4.string().default("v4-default"),
+      }),
+      contextSchema: z4.object({
+        middleware2Context: z4.boolean(),
+      }),
+      beforeModel: (_state, { context }) => {
+        expect(context.middleware2Context).toBe(true);
+        return {
+          middleware2Value: "v4-modified",
+        };
+      },
+    });
+
+    // Create a tool with Zod v3 schema
+    const testTool = tool(async (input) => `Result: ${input.query}`, {
+      name: "test_tool",
+      description: "A test tool",
+      schema: z.object({
+        query: z.string(),
+      }),
+    });
+
+    const responseFormat = toolStrategy(
+      z.object({
+        result: z.string(),
+        score: z.number(),
+      })
+    );
+
+    const expectedStructuredResponse = { result: "success", score: 95 };
+    const model = new FakeToolCallingChatModel({
+      responses: [
+        new AIMessage({
+          content: "",
+          tool_calls: [
+            {
+              name: "test_tool",
+              id: "tool_1",
+              args: { query: "test" },
+            },
+          ],
+        }),
+        new AIMessage({
+          content: "",
+          tool_calls: [
+            {
+              name: responseFormat[0].tool.function.name,
+              id: "extract",
+              args: expectedStructuredResponse,
+            },
+          ],
+        }),
+      ],
+      structuredResponse: expectedStructuredResponse,
+    });
+
+    const agent = createAgent({
+      model,
+      tools: [testTool],
+      // Use Zod v4 for agent stateSchema
+      stateSchema: z4.object({
+        agentCounter: z4.number().default(0),
+        agentName: z4.string().optional(),
+      }),
+      responseFormat,
+      middleware: [middleware1, middleware2],
+    });
+
+    const result = await agent.invoke(
+      {
+        messages: [new HumanMessage("Test mixed schemas")],
+        agentCounter: 1,
+        agentName: "test-agent",
+      },
+      {
+        context: {
+          middleware1Context: 42,
+          middleware2Context: true,
+        },
+      }
+    );
+
+    // Verify agent state schema (Zod v3)
+    expect(result.agentCounter).toBe(1);
+    expect(result.agentName).toBe("test-agent");
+
+    // Verify middleware1 state (Zod v3)
+    expect(result.middleware1Value).toBe("v3-modified");
+
+    // Verify middleware2 state (Zod v4)
+    expect(result.middleware2Value).toBe("v4-modified");
+
+    // Verify structured response (Zod v4)
+    expect(result.structuredResponse).toEqual(expectedStructuredResponse);
+    expect(result.structuredResponse.result).toBe("success");
+    expect(result.structuredResponse.score).toBe(95);
+
+    // Verify messages were processed correctly
+    expect(result.messages.length).toBeGreaterThan(0);
+  });
 });