feat(langgraph): add UpdateType type utility for Zod, improve Zod 4 and Zod 4 mini support (#1402)

.changeset/whole-dots-retire.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+add UpdateType type utility for Zod, improve Zod 4 and Zod 4 mini support

libs/langgraph/src/graph/messages_annotation.ts

@@ -1,10 +1,10 @@
 /* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
 
 import { BaseMessage } from "@langchain/core/messages";
-import { z } from "zod";
+import { z } from "zod/v3";
 import { Annotation } from "./annotation.js";
 import { Messages, messagesStateReducer } from "./message.js";
-import { withLangGraph } from "./zod/meta.js";
+import { SchemaMeta, withLangGraph } from "./zod/meta.js";
 
 /**
  * Prebuilt state annotation that combines returned messages.
@@ -48,6 +48,25 @@ export const MessagesAnnotation = Annotation.Root({
   }),
 });
 
+/**
+ * Prebuilt schema meta for Zod state definition.
+ *
+ * @example
+ * ```ts
+ * import { z } from "zod/v4-mini";
+ * import { MessagesZodState, StateGraph } from "@langchain/langgraph";
+ *
+ * const AgentState = z.object({
+ *   messages: z.custom<BaseMessage[]>().register(registry, MessagesZodMeta),
+ * });
+ * ```
+ */
+export const MessagesZodMeta: SchemaMeta<BaseMessage[], Messages> = {
+  reducer: { fn: messagesStateReducer },
+  jsonSchemaExtra: { langgraph_type: "messages" },
+  default: () => [],
+};
+
 /**
  * Prebuilt state object that uses Zod to combine returned messages.
  * This utility is synonymous with the `MessagesAnnotation` annotation,
@@ -88,14 +107,5 @@ export const MessagesAnnotation = Annotation.Root({
  * ```
  */
 export const MessagesZodState = z.object({
-  messages: withLangGraph(z.custom<BaseMessage[]>(), {
-    reducer: {
-      schema: z.custom<Messages>(),
-      fn: messagesStateReducer,
-    },
-    jsonSchemaExtra: {
-      langgraph_type: "messages",
-    },
-    default: () => [],
-  }),
+  messages: withLangGraph(z.custom<BaseMessage[]>(), MessagesZodMeta),
 });

libs/langgraph/src/graph/zod/meta.ts

@@ -16,12 +16,15 @@ import { LastValue } from "../../channels/last_value.js";
 
 export const META_EXTRAS_DESCRIPTION_PREFIX = "lg:";
 
+/** @internal */
 export type ReducedZodChannel<
   T extends InteropZodType,
   TReducerSchema extends InteropZodType
 > = T & {
   lg_reducer_schema: TReducerSchema;
 };
+
+/** @internal */
 export type InteropZodToStateDefinition<
   T extends InteropZodObject,
   TShape = InteropZodObjectShape<T>
@@ -40,6 +43,24 @@ export type InteropZodToStateDefinition<
     : never;
 };
 
+export type UpdateType<
+  T extends InteropZodObject,
+  TShape = InteropZodObjectShape<T>
+> = {
+  [key in keyof TShape]?: TShape[key] extends ReducedZodChannel<
+    infer Schema,
+    infer ReducerSchema
+  >
+    ? Schema extends InteropZodType<unknown>
+      ? ReducerSchema extends InteropZodType<infer U>
+        ? U
+        : never
+      : never
+    : TShape[key] extends InteropZodType<unknown, infer U>
+    ? U
+    : never;
+};
+
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export interface SchemaMeta<TValue = any, TUpdate = TValue> {
   jsonSchemaExtra?: {

libs/langgraph/src/graph/zod/zod-registry.ts

@@ -1,5 +1,9 @@
 // @ts-expect-error If zod/v4 is not imported, the module augmentation will fail in build
 import type { ZodType } from "zod/v4"; // eslint-disable-line @typescript-eslint/no-unused-vars
+
+// @ts-expect-error If zod/v4-mini is not imported, the module augmentation will fail in build
+import type { ZodMiniType } from "zod/v4-mini"; // eslint-disable-line @typescript-eslint/no-unused-vars
+
 import type * as core from "zod/v4/core";
 import { getInteropZodDefaultGetter } from "@langchain/core/utils/types";
 import { $ZodType, $ZodRegistry, $replace } from "zod/v4/core";
@@ -80,4 +84,28 @@ declare module "zod/v4" {
   }
 }
 
+declare module "zod/v4-mini" {
+  export interface ZodMiniType<
+    out Output = unknown,
+    out Input = unknown,
+    out Internals extends core.$ZodTypeInternals<
+      Output,
+      Input
+    > = core.$ZodTypeInternals<Output, Input>
+  > extends core.$ZodType<Output, Input, Internals> {
+    register<
+      R extends LanggraphZodMetaRegistry,
+      TOutput = core.output<this>,
+      TInput = core.input<this>,
+      TInternals extends core.$ZodTypeInternals<
+        TOutput,
+        TInput
+      > = core.$ZodTypeInternals<TOutput, TInput>
+    >(
+      registry: R,
+      meta: SchemaMeta<TOutput, TInput>
+    ): ReducedZodChannel<this, ZodMiniType<TOutput, TInput, TInternals>>;
+  }
+}
+
 export const registry = new LanggraphZodMetaRegistry(schemaMetaRegistry);

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

@@ -0,0 +1,120 @@
+import { BaseMessage, HumanMessage } from "@langchain/core/messages";
+import { describe, expectTypeOf, it } from "vitest";
+
+import { z as z3 } from "zod/v3";
+import { z as z4 } from "zod/v4";
+import { z as z4mini } from "zod/v4-mini";
+
+import { type UpdateType, withLangGraph } from "../graph/zod/meta.js";
+import {
+  MessagesZodMeta,
+  MessagesZodState,
+} from "../graph/messages_annotation.js";
+import type { Messages } from "../graph/message.js";
+import { StateGraph } from "../graph/state.js";
+import { START } from "../constants.js";
+import { registry } from "../graph/zod/zod-registry.js";
+
+describe("zod 3 (prebuilt)", () => {
+  const state = MessagesZodState;
+
+  it("update type", () => {
+    // We're checking for correct union types here, thus we flip the assertion
+    expectTypeOf<{ messages?: Messages | undefined }>().toExtend<
+      UpdateType<typeof state>
+    >();
+  });
+
+  it("state graph", async () => {
+    const builder = new StateGraph(state);
+
+    const node = (state: {
+      messages: BaseMessage[];
+    }): { messages?: Messages } => state;
+
+    // Check for assignability
+    const graph = builder.addNode({ node }).addEdge(START, "node").compile();
+    await graph.invoke({ messages: "input" });
+    await graph.invoke({ messages: [new HumanMessage("input")] });
+  });
+});
+
+describe("zod 3", () => {
+  const state = z3.object({
+    messages: withLangGraph(z3.custom<BaseMessage[]>(), MessagesZodMeta),
+  });
+
+  it("update type", () => {
+    // We're checking for correct union types here, thus we flip the assertion
+    expectTypeOf<{ messages?: Messages | undefined }>().toExtend<
+      UpdateType<typeof state>
+    >();
+  });
+
+  it("state graph", async () => {
+    const builder = new StateGraph(state);
+
+    const node = (state: {
+      messages: BaseMessage[];
+    }): { messages?: Messages } => state;
+
+    // Check for assignability
+    const graph = builder.addNode({ node }).addEdge(START, "node").compile();
+    await graph.invoke({ messages: "input" });
+    await graph.invoke({ messages: [new HumanMessage("input")] });
+  });
+});
+
+describe("zod 4", () => {
+  const state = z4.object({
+    messages: z4.custom<BaseMessage[]>().register(registry, MessagesZodMeta),
+  });
+
+  it("update type", () => {
+    // We're checking for correct union types here, thus we flip the assertion
+    expectTypeOf<{ messages?: Messages | undefined }>().toExtend<
+      UpdateType<typeof state>
+    >();
+  });
+
+  it("state graph", async () => {
+    const builder = new StateGraph(state);
+
+    const node = (state: {
+      messages: BaseMessage[];
+    }): { messages?: Messages } => state;
+
+    // Check for assignability
+    const graph = builder.addNode({ node }).addEdge(START, "node").compile();
+    await graph.invoke({ messages: "input" });
+    await graph.invoke({ messages: [new HumanMessage("input")] });
+  });
+});
+
+describe("zod 4 mini", () => {
+  const state = z4mini.object({
+    messages: z4mini
+      .custom<BaseMessage[]>()
+      .register(registry, MessagesZodMeta),
+  });
+
+  it("update type", () => {
+    // We're checking for correct union types here, thus we flip the assertion
+    expectTypeOf<{ messages?: Messages | undefined }>().toExtend<
+      UpdateType<typeof state>
+    >();
+  });
+
+  it("state graph", async () => {
+    const builder = new StateGraph(state);
+
+    const node = (state: {
+      messages: BaseMessage[];
+    }): { messages?: Messages } => state;
+
+    // Check for assignability
+    const graph = builder.addNode({ node }).addEdge(START, "node").compile();
+    await graph.invoke({ messages: "input" });
+    await graph.invoke({ messages: [new HumanMessage("input")] });
+  });
+});

libs/langgraph/src/web.ts

@@ -89,5 +89,6 @@ export {
 export {
   MessagesAnnotation,
   MessagesZodState,
+  MessagesZodMeta,
 } from "./graph/messages_annotation.js";
 export { type LangGraphRunnableConfig } from "./pregel/runnable_types.js";