feat(agents): Improve handling of withLangGraph annotations

libs/langchain/src/agents/nodes/tests/utils.test.ts

@@ -0,0 +1,65 @@
+import { describe, it, expect } from "vitest";
+import { z } from "zod/v3";
+import { withLangGraph } from "@langchain/langgraph/zod";
+import { initializeMiddlewareStates } from "../utils.js";
+import type { AgentMiddleware } from "../../middleware/types.js";
+
+describe("initializeMiddlewareStates", () => {
+  it("should handle withLangGraph defaults", async () => {
+    // Create a middleware with a state schema that has a default
+    const middleware: AgentMiddleware = {
+      name: "TestMiddleware",
+      stateSchema: z.object({
+        testField: withLangGraph(z.string(), {
+          default: () => "default-value",
+        }),
+      }),
+    };
+
+    // Call with empty state (like the test does)
+    const result = await initializeMiddlewareStates([middleware], {});
+
+    // Should apply the default value
+    expect(result).toEqual({
+      testField: "default-value",
+    });
+  });
+
+  it("should handle withLangGraph defaults with reducers", async () => {
+    // Create a middleware like the text editor middleware
+    const middleware: AgentMiddleware = {
+      name: "TestMiddleware",
+      stateSchema: z.object({
+        fieldWithDefault: withLangGraph(z.string(), {
+          reducer: {
+            fn: (_left: string, right: string) => right,
+          },
+          default: () => "default value",
+        }),
+      }),
+    };
+
+    // Call with empty state
+    const result = await initializeMiddlewareStates([middleware], {});
+
+    // Should apply the default value
+    expect(result).toEqual({
+      fieldWithDefault: "default value",
+    });
+  });
+
+  it("should throw error for truly required fields without defaults", async () => {
+    // Create a middleware with a required field (no default)
+    const middleware: AgentMiddleware = {
+      name: "TestMiddleware",
+      stateSchema: z.object({
+        requiredField: z.string(),
+      }),
+    };
+
+    // Should throw because field is required and has no default
+    await expect(initializeMiddlewareStates([middleware], {})).rejects.toThrow(
+      /required/i
+    );
+  });
+});

libs/langchain/src/agents/nodes/utils.ts

@@ -1,22 +1,27 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import { z } from "zod/v3";
 import { type BaseMessage } from "@langchain/core/messages";
 import {
+  getInteropZodObjectShape,
   interopSafeParseAsync,
   interopZodObjectMakeFieldsOptional,
 } from "@langchain/core/utils/types";
-import { type ZodIssue } from "zod/v3";
 import { END } from "@langchain/langgraph";
+import { schemaMetaRegistry } from "@langchain/langgraph/zod";
+import { z } from "zod/v3";
 
-import type { JumpTo } from "../types.js";
 import type { AgentMiddleware } from "../middleware/types.js";
+import type { JumpTo } from "../types.js";
 
 /**
  * Helper function to initialize middleware state defaults.
  * This is used to ensure all middleware state properties are initialized.
  *
  * Private properties (starting with _) are automatically made optional since
  * users cannot provide them when invoking the agent.
+ *
+ * This function also checks the LangGraph schemaMetaRegistry for fields with
+ * default values defined via withLangGraph(), and applies those defaults when
+ * the fields are omitted from the input state.
  */
 export async function initializeMiddlewareStates(
   middlewareList: readonly AgentMiddleware[],
@@ -41,36 +46,67 @@ export async function initializeMiddlewareStates(
       }
 
       /**
-       * If safeParse fails, there are required public fields missing
+       * If safeParse fails, check if the missing fields have defaults in the
+       * schemaMetaRegistry (from withLangGraph). If they do, apply the defaults.
+       * Only throw an error for truly required fields (no defaults, not optional).
        */
-      const requiredFields = parseResult.error.issues
-        .filter(
-          (issue: ZodIssue) =>
-            issue.code === "invalid_type" && issue.message === "Required"
-        )
-        .map(
-          (issue: ZodIssue) => `  - ${issue.path.join(".")}: ${issue.message}`
-        )
-        .join("\n");
-
-      throw new Error(
-        `Middleware "${middleware.name}" has required state fields that must be initialized:\n` +
-          `${requiredFields}\n\n` +
-          `To fix this, either:\n` +
-          `1. Provide default values in your middleware's state schema using .default():\n` +
-          `   stateSchema: z.object({\n` +
-          `     myField: z.string().default("default value")\n` +
-          `   })\n\n` +
-          `2. Or make the fields optional using .optional():\n` +
-          `   stateSchema: z.object({\n` +
-          `     myField: z.string().optional()\n` +
-          `   })\n\n` +
-          `3. Or ensure you pass these values when invoking the agent:\n` +
-          `   agent.invoke({\n` +
-          `     messages: [...],\n` +
-          `     ${parseResult.error.issues[0]?.path.join(".")}: "value"\n` +
-          `   })`
-      );
+      const shape = getInteropZodObjectShape(middleware.stateSchema);
+      const missingRequiredFields: string[] = [];
+      const fieldsWithDefaults: Record<string, any> = {};
+
+      for (const issue of parseResult.error.issues) {
+        if (issue.code === "invalid_type" && issue.message === "Required") {
+          const fieldName = issue.path[0] as string;
+          const fieldSchema = shape[fieldName];
+
+          if (fieldSchema) {
+            // Check if this field has a default in the registry
+            const meta = schemaMetaRegistry.get(fieldSchema);
+            if (meta?.default) {
+              // Apply the default value
+              fieldsWithDefaults[fieldName] = meta.default();
+            } else {
+              // No default found - this is truly required
+              missingRequiredFields.push(
+                `  - ${issue.path.join(".")}: ${issue.message}`
+              );
+            }
+          }
+        }
+      }
+
+      // If there are truly required fields (no defaults), throw an error
+      if (missingRequiredFields.length > 0) {
+        throw new Error(
+          `Middleware "${middleware.name}" has required state fields that must be initialized:\n` +
+            `${missingRequiredFields.join("\n")}\n\n` +
+            `To fix this, either:\n` +
+            `1. Provide default values in your middleware's state schema using withLangGraph():\n` +
+            `   import { withLangGraph } from "@langchain/langgraph/zod";\n` +
+            `   stateSchema: z.object({\n` +
+            `     myField: withLangGraph(z.string(), { default: () => "default value" })\n` +
+            `   })\n\n` +
+            `2. Or use Zod's .default():\n` +
+            `   stateSchema: z.object({\n` +
+            `     myField: z.string().default("default value")\n` +
+            `   })\n\n` +
+            `3. Or make the fields optional using .optional():\n` +
+            `   stateSchema: z.object({\n` +
+            `     myField: z.string().optional()\n` +
+            `   })\n\n` +
+            `4. Or ensure you pass these values when invoking the agent:\n` +
+            `   agent.invoke({\n` +
+            `     messages: [...],\n` +
+            `     ${missingRequiredFields[0]
+              ?.split(":")[0]
+              ?.trim()
+              .replace("- ", "")}: "value"\n` +
+            `   })`
+        );
+      }
+
+      // Merge the fields with applied defaults into middlewareStates
+      Object.assign(middlewareStates, fieldsWithDefaults);
     }
   }
 

libs/langchain/src/agents/tests/middleware-defaults.test.ts

@@ -0,0 +1,192 @@
+/**
+ * Tests for agent middleware state schemas with default values.
+ *
+ * These tests verify that fields with `withLangGraph` defaults are correctly:
+ * 1. Treated as optional in TypeScript types (no type errors when omitted)
+ * 2. Given their default values at runtime when omitted from invoke()
+ * 3. Still accept explicitly provided values
+ */
+import { describe, it, expect } from "vitest";
+import { z } from "zod/v3";
+import { withLangGraph } from "@langchain/langgraph/zod";
+import { AIMessage, HumanMessage } from "@langchain/core/messages";
+import { createAgent } from "../index.js";
+import { createMiddleware } from "../middleware.js";
+import { FakeToolCallingChatModel } from "./utils.js";
+
+describe("Middleware state with withLangGraph defaults", () => {
+  it("should make fields with withLangGraph defaults optional in invoke parameter", async () => {
+    // Create middleware with a field that has a default
+    const middlewareWithDefaults = createMiddleware({
+      name: "TestMiddleware",
+      stateSchema: z.object({
+        fieldWithDefault: withLangGraph(z.string(), {
+          default: () => "default-value",
+        }),
+      }),
+    });
+
+    const llm = new FakeToolCallingChatModel({
+      responses: [new AIMessage({ id: "0", content: "Response" })],
+    });
+
+    const agent = createAgent({
+      model: llm,
+      middleware: [middlewareWithDefaults],
+    });
+
+    // This should NOT have a type error - fieldWithDefault should be optional
+    // because it has a default value
+    const result1 = await agent.invoke({
+      messages: [new HumanMessage("test")],
+      // fieldWithDefault is omitted - should be OK due to default
+    });
+
+    // Verify the default value was applied
+    expect(result1.fieldWithDefault).toBe("default-value");
+
+    // This should also work when providing the field
+    const result2 = await agent.invoke({
+      messages: [new HumanMessage("test")],
+      fieldWithDefault: "custom-value",
+    });
+
+    // Verify the custom value was used
+    expect(result2.fieldWithDefault).toBe("custom-value");
+  });
+
+  it("should make fields with withLangGraph reducer defaults optional in invoke parameter", async () => {
+    // Create middleware like the text editor middleware
+    const middlewareWithReducerDefaults = createMiddleware({
+      name: "TestMiddleware",
+      stateSchema: z.object({
+        fieldWithDefault: withLangGraph(z.string(), {
+          reducer: {
+            fn: (_left: string, right: string) => right,
+          },
+          default: () => "default value",
+        }),
+      }),
+    });
+
+    const llm = new FakeToolCallingChatModel({
+      responses: [new AIMessage({ id: "0", content: "Response" })],
+    });
+
+    const agent = createAgent({
+      model: llm,
+      middleware: [middlewareWithReducerDefaults],
+    });
+
+    // This should NOT have a type error - fieldWithDefault should be optional
+    const result1 = await agent.invoke({
+      messages: [new HumanMessage("test")],
+      // fieldWithDefault is omitted - should be OK due to default
+    });
+
+    // Verify the default value was applied
+    expect(result1.fieldWithDefault).toEqual("default value");
+
+    // This should also work when providing the field
+    const result2 = await agent.invoke({
+      messages: [new HumanMessage("test")],
+      fieldWithDefault: "nondefault value",
+    });
+
+    // Verify the custom value was used
+    expect(result2.fieldWithDefault).toEqual("nondefault value");
+  });
+
+  it("should require fields without defaults in invoke parameter", async () => {
+    // Create middleware with a required field (no default)
+    const middlewareWithRequired = createMiddleware({
+      name: "TestMiddleware",
+      stateSchema: z.object({
+        requiredField: z.string(),
+      }),
+    });
+
+    const llm = new FakeToolCallingChatModel({
+      responses: [new AIMessage({ id: "0", content: "Response" })],
+    });
+
+    const agent = createAgent({
+      model: llm,
+      middleware: [middlewareWithRequired],
+    });
+
+    await expect(
+      // @ts-expect-error - requiredField should be required
+      agent.invoke({
+        messages: [new HumanMessage("test")],
+        // requiredField is omitted - should be a type error AND runtime error
+      })
+    ).rejects.toThrow(/required/i);
+
+    // This should work when providing the field
+    const result = await agent.invoke({
+      messages: [new HumanMessage("test")],
+      requiredField: "value",
+    });
+
+    // Verify the provided value was used
+    expect(result.requiredField).toBe("value");
+  });
+
+  it("should make fields optional when passed directly via stateSchema", async () => {
+    // When stateSchema is passed directly (not via middleware), it should work the same way
+    const llm = new FakeToolCallingChatModel({
+      responses: [new AIMessage({ id: "0", content: "Response" })],
+    });
+
+    const agent = createAgent({
+      model: llm,
+      stateSchema: z.object({
+        fieldWithDefault: withLangGraph(z.string(), {
+          default: () => "default-value",
+        }),
+      }),
+    });
+
+    // This should NOT have a type error
+    const result = await agent.invoke({
+      messages: [new HumanMessage("test")],
+      // fieldWithDefault is omitted - should be OK due to default
+    });
+
+    // Field should NOT appear in output since no node updated it
+    // (LangGraph only applies defaults when a field is actually written to)
+    expect(result.fieldWithDefault).toBeUndefined();
+    expect(result.messages).toBeDefined();
+  });
+
+  it("should make fields with reducer optional when passed directly via stateSchema", async () => {
+    // When stateSchema is passed directly (not via middleware), it should work the same way
+    const llm = new FakeToolCallingChatModel({
+      responses: [new AIMessage({ id: "0", content: "Response" })],
+    });
+
+    const agent = createAgent({
+      model: llm,
+      stateSchema: z.object({
+        fieldWithDefault: withLangGraph(z.string(), {
+          reducer: {
+            fn: (_left: string, right: string) => right,
+          },
+          default: () => "default-value",
+        }),
+      }),
+    });
+
+    // This should NOT have a type error
+    const result = await agent.invoke({
+      messages: [new HumanMessage("test")],
+      // fieldWithDefault is omitted - should be OK due to default
+    });
+
+    // Field should NOT appear in output since no node updated it
+    // (LangGraph only applies defaults when a field is actually written to)
+    expect(result.fieldWithDefault).toBeUndefined();
+    expect(result.messages).toBeDefined();
+  });
+});