fix(langgraph): fix default input Zod schema not obtaining langgraph metadata (#1232)

Author: dqbd

libs/langgraph/src/graph/messages_annotation.ts

@@ -93,6 +93,9 @@ export const MessagesZodState = z.object({
       schema: z.custom<Messages>(),
       fn: messagesStateReducer,
     },
+    jsonSchemaExtra: {
+      langgraph_type: "messages",
+    },
     default: () => [],
   }),
 });

libs/langgraph/src/graph/state.ts

@@ -127,6 +127,9 @@ type NodeAction<S, U, C extends SDZod> = RunnableLike<
   LangGraphRunnableConfig<StateType<ToStateDefinition<C>>>
 >;
 
+const PartialStateSchema = Symbol.for("langgraph.state.partial");
+type PartialStateSchema = typeof PartialStateSchema;
+
 /**
  * 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>`.
@@ -213,7 +216,7 @@ export class StateGraph<
   _inputDefinition: I;
 
   /** @internal */
-  _inputRuntimeDefinition: AnyZodObject | undefined;
+  _inputRuntimeDefinition: AnyZodObject | PartialStateSchema | undefined;
 
   /** @internal */
   _outputDefinition: O;
@@ -292,7 +295,7 @@ export class StateGraph<
       this._schemaRuntimeDefinition = fields.state;
 
       this._inputDefinition = inputDef as I;
-      this._inputRuntimeDefinition = fields.input ?? fields.state.partial();
+      this._inputRuntimeDefinition = fields.input ?? PartialStateSchema;
 
       this._outputDefinition = outputDef as O;
       this._outputRuntimeDefinition = fields.output ?? fields.state;
@@ -303,7 +306,7 @@ export class StateGraph<
       this._schemaRuntimeDefinition = fields;
 
       this._inputDefinition = stateDef as I;
-      this._inputRuntimeDefinition = fields.partial();
+      this._inputRuntimeDefinition = PartialStateSchema;
 
       this._outputDefinition = stateDef as O;
       this._outputRuntimeDefinition = fields;
@@ -944,7 +947,11 @@ export class CompiledStateGraph<
   protected async _validateInput(
     input: UpdateType<ToStateDefinition<I>>
   ): Promise<UpdateType<ToStateDefinition<I>>> {
-    const inputSchema = this.builder._inputRuntimeDefinition;
+    let inputSchema = this.builder._inputRuntimeDefinition;
+    if (inputSchema === PartialStateSchema) {
+      inputSchema = this.builder._schemaRuntimeDefinition?.partial();
+    }
+
     if (isCommand(input)) {
       const parsedInput = input;
       if (input.update && isAnyZodObject(inputSchema))

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

@@ -5,6 +5,9 @@ import { getMeta } from "./state.js";
 const TYPE_CACHE: Record<string, WeakMap<z.AnyZodObject, z.AnyZodObject>> = {};
 const DESCRIPTION_PREFIX = "lg:";
 
+const PartialStateSchema = Symbol.for("langgraph.state.partial");
+type PartialStateSchema = typeof PartialStateSchema;
+
 function applyPlugin(
   schema: z.AnyZodObject,
   actions: {
@@ -61,7 +64,7 @@ function applyPlugin(
 interface GraphWithZodLike {
   builder: {
     _schemaRuntimeDefinition: z.AnyZodObject | undefined;
-    _inputRuntimeDefinition: z.AnyZodObject | undefined;
+    _inputRuntimeDefinition: z.AnyZodObject | PartialStateSchema | undefined;
     _outputRuntimeDefinition: z.AnyZodObject | undefined;
     _configRuntimeSchema: z.AnyZodObject | undefined;
   };
@@ -124,7 +127,7 @@ export function getStateTypeSchema(graph: unknown): JsonSchema | undefined {
   if (!isGraphWithZodLike(graph)) return undefined;
   const schemaDef = graph.builder._schemaRuntimeDefinition;
   if (!schemaDef) return undefined;
-  return toJsonSchema(schemaDef);
+  return toJsonSchema(applyPlugin(schemaDef, { jsonSchemaExtra: true }));
 }
 
 /**
@@ -153,7 +156,12 @@ export function getUpdateTypeSchema(graph: unknown): JsonSchema | undefined {
  */
 export function getInputTypeSchema(graph: unknown): JsonSchema | undefined {
   if (!isGraphWithZodLike(graph)) return undefined;
-  const schemaDef = graph.builder._inputRuntimeDefinition;
+  let schemaDef = graph.builder._inputRuntimeDefinition;
+  if (schemaDef === PartialStateSchema) {
+    // No need to pass `.partial()` here, that's being done by `applyPlugin`
+    schemaDef = graph.builder._schemaRuntimeDefinition;
+  }
+
   if (!schemaDef) return undefined;
   return toJsonSchema(
     applyPlugin(schemaDef, {

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

@@ -9,7 +9,7 @@ const META_MAP = new WeakMap<z.ZodType, Meta<any, any>>();
 export interface Meta<ValueType, UpdateType = ValueType> {
   jsonSchemaExtra?: {
     langgraph_nodes?: string[];
-    langgraph_type?: "prompt";
+    langgraph_type?: "prompt" | "messages";
 
     [key: string]: unknown;
   };

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

@@ -2,6 +2,14 @@ import { describe, it, expect } from "vitest";
 import { z } from "zod";
 import { StateGraph } from "../graph/state.js";
 import { END, START } from "../constants.js";
+import { _AnyIdAIMessage, _AnyIdHumanMessage } from "./utils.js";
+import {
+  getOutputTypeSchema,
+  getInputTypeSchema,
+  getUpdateTypeSchema,
+  getStateTypeSchema,
+} from "../graph/zod/schema.js";
+import { MessagesZodState } from "../graph/messages_annotation.js";
 
 describe("StateGraph with Zod schemas", () => {
   it("should accept Zod schema as input in addNode", async () => {
@@ -77,4 +85,63 @@ describe("StateGraph with Zod schemas", () => {
       count: 1,
     });
   });
+
+  it("should accept Zod messages schema & return tagged JSON schema", async () => {
+    const schema = MessagesZodState.extend({ count: z.number() });
+
+    const graph = new StateGraph(schema)
+      .addNode("agent", () => ({
+        messages: [{ type: "ai", content: "agent" }],
+      }))
+      .addNode("tool", () => ({
+        messages: [{ type: "ai", content: "tool" }],
+      }))
+      .addEdge("__start__", "agent")
+      .addEdge("agent", "tool")
+      .compile();
+
+    expect(
+      await graph.invoke({
+        messages: [{ type: "human", content: "hello" }],
+      })
+    ).toMatchObject({
+      messages: [
+        new _AnyIdHumanMessage("hello"),
+        new _AnyIdAIMessage("agent"),
+        new _AnyIdAIMessage("tool"),
+      ],
+    });
+
+    expect.soft(getStateTypeSchema(graph)).toMatchObject({
+      $schema: "http://json-schema.org/draft-07/schema#",
+      properties: {
+        messages: { langgraph_type: "messages" },
+        count: { type: "number" },
+      },
+    });
+
+    expect.soft(getUpdateTypeSchema(graph)).toMatchObject({
+      $schema: "http://json-schema.org/draft-07/schema#",
+      properties: {
+        messages: { langgraph_type: "messages" },
+        count: { type: "number" },
+      },
+    });
+
+    expect.soft(getInputTypeSchema(graph)).toMatchObject({
+      $schema: "http://json-schema.org/draft-07/schema#",
+      properties: {
+        messages: { langgraph_type: "messages" },
+        count: { type: "number" },
+      },
+    });
+
+    expect.soft(getOutputTypeSchema(graph)).toMatchObject({
+      $schema: "http://json-schema.org/draft-07/schema#",
+      properties: {
+        messages: { langgraph_type: "messages" },
+        count: { type: "number" },
+      },
+    });
+  });
 });