fix(sdk-js): preserve Zod v4 state fields in LangGraph output_schema (CopilotKit#4358)

- Adds a `zodState` helper that attaches a `~standard.jsonSchema.input`
hook to Standard-Schema–compatible schemas. Without it, Zod v4 fields
carry only `~standard.validate` + `vendor`, so LangGraph's
`isStandardJSONSchema()` returns false and `getJsonSchemaFromSchema`
(called from `StateSchema.getJsonSchema`) silently drops them from the
graph's `output_schema`. The fields then get filtered out of AG-UI
`STATE_SNAPSHOT` payloads even though the underlying thread state has
them, and never reach `useAgent().state.*` on the frontend.
- Wraps the internal `copilotkit` field in `copilotKitStateSchema` with
`zodState(...)`, and exports the helper for user state schemas to do the
same.
- Bumps `@ag-ui/langgraph` to `0.0.30` and `@langchain/{core,langgraph}`
/ `langchain` to the `1.1.41` / `1.2.9` / `1.3.4` line — this is the
dependency upgrade that exposes the Zod v4 behavior the helper works
around.

This recreates the sdk-js portion of CopilotKit#4320 in isolation so it can land
without the integration-demo refresh.

Author: ranst91

packages/runtime/package.json

@@ -92,7 +92,7 @@
     "@ag-ui/client": "0.0.52",
     "@ag-ui/core": "0.0.52",
     "@ag-ui/encoder": "0.0.52",
-    "@ag-ui/langgraph": "0.0.29",
+    "@ag-ui/langgraph": "0.0.31",
     "@ag-ui/mcp-apps-middleware": "0.0.3",
     "@ai-sdk/anthropic": "^3.0.49",
     "@ai-sdk/google": "^3.0.33",

packages/sdk-js/package.json

@@ -59,18 +59,18 @@
     "attw": "attw --pack . --profile node16"
   },
   "dependencies": {
-    "@ag-ui/langgraph": "0.0.29",
+    "@ag-ui/langgraph": "0.0.31",
     "@copilotkit/shared": "workspace:*"
   },
   "devDependencies": {
-    "@langchain/core": "^1.1.8",
-    "@langchain/langgraph": "^1.0.7",
+    "@langchain/core": "^1.1.41",
+    "@langchain/langgraph": "^1.2.9",
     "@swc/core": "1.5.28",
     "@types/express": "^4.17.21",
     "@types/node": "^18.11.17",
     "@whatwg-node/server": "^0.9.34",
     "eslint": "^8.56.0",
-    "langchain": "^1.2.3",
+    "langchain": "^1.3.4",
     "nodemon": "^3.1.3",
     "ts-node": "^10.9.2",
     "tsconfig": "workspace:*",

packages/sdk-js/src/langgraph/__tests__/middleware.test.ts

@@ -18,6 +18,7 @@
  */
 
 import { describe, it, expect } from "vitest";
+import * as z from "zod";
 import {
   AIMessage,
   HumanMessage,
@@ -27,6 +28,7 @@ import {
 import {
   copilotkitMiddleware,
   createCopilotkitMiddleware,
+  zodState,
 } from "../middleware";
 
 // ---------------------------------------------------------------------------
@@ -495,3 +497,115 @@ describe("afterAgent", () => {
     expect(result!.copilotkit.originalAIMessageId).toBeUndefined();
   });
 });
+
+// ---------------------------------------------------------------------------
+// zodState — Standard-Schema JSON-schema augmentation
+// ---------------------------------------------------------------------------
+//
+// Contract: a Zod v4 schema only carries `~standard.validate` + `vendor`, so
+// LangGraph's `isStandardJSONSchema()` returns false and the field is dropped
+// from `output_schema`. `zodState` patches `~standard.jsonSchema.input` onto
+// the schema so the field survives serialization and reaches the frontend
+// via AG-UI `STATE_SNAPSHOT` events.
+
+type JsonSchema = {
+  type?: string;
+  properties?: Record<string, unknown>;
+};
+type StdJson = { jsonSchema?: { input: () => JsonSchema } };
+const standardOf = (s: object): StdJson | undefined =>
+  (s as { "~standard"?: StdJson })["~standard"];
+const hasZodV4ToJsonSchema = (): boolean =>
+  typeof (z as { toJSONSchema?: unknown }).toJSONSchema === "function";
+
+describe("zodState", () => {
+  it("attaches a ~standard.jsonSchema.input hook to a Zod schema", () => {
+    const schema = z.object({ name: z.string() });
+    expect(standardOf(schema)?.jsonSchema).toBeUndefined();
+
+    const wrapped = zodState(schema);
+
+    const std = standardOf(wrapped);
+    expect(std?.jsonSchema).toBeDefined();
+    expect(typeof std?.jsonSchema?.input).toBe("function");
+  });
+
+  it("returns a plain object from input() (real JSON Schema when zod v4 is available, else `{}`)", () => {
+    // The contract langgraph cares about is just "input() returns an
+    // object" — `isStandardJSONSchema()` doesn't introspect the shape.
+    // When zod v4's `toJSONSchema` is present we get the real schema;
+    // otherwise we get `{}`, which is still enough for the field to
+    // appear in `output_schema` (langgraph-api treats it as opaque).
+    const schema = z.object({ todos: z.array(z.string()) });
+    const wrapped = zodState(schema);
+
+    const json = standardOf(wrapped)?.jsonSchema?.input();
+
+    expect(json).toBeTypeOf("object");
+    expect(json).not.toBeNull();
+    if (hasZodV4ToJsonSchema()) {
+      expect(json?.type).toBe("object");
+      expect(json?.properties?.todos).toBeDefined();
+    }
+  });
+
+  it("caches the JSON Schema across calls (input() returns the same object)", () => {
+    const schema = z.object({ x: z.string() });
+    const wrapped = zodState(schema);
+    const std = standardOf(wrapped);
+
+    const a = std?.jsonSchema?.input();
+    const b = std?.jsonSchema?.input();
+
+    expect(a).toBe(b);
+  });
+
+  it("does not overwrite an existing jsonSchema hook", () => {
+    const schema = z.object({ x: z.string() });
+    const preset: StdJson["jsonSchema"] = { input: () => ({}) };
+    const std = standardOf(schema);
+    expect(std).toBeDefined();
+    std!.jsonSchema = preset;
+
+    zodState(schema);
+
+    expect(standardOf(schema)?.jsonSchema).toBe(preset);
+  });
+
+  it("returns the same reference (mutates rather than copies)", () => {
+    const schema = z.object({ x: z.string() });
+    expect(zodState(schema)).toBe(schema);
+  });
+
+  it("is a no-op for a value without ~standard metadata", () => {
+    const plain = { foo: "bar" };
+    expect(() => zodState(plain)).not.toThrow();
+    expect(zodState(plain)).toBe(plain);
+  });
+
+  it("works on optional / array / default-wrapped schemas (state-field shapes)", () => {
+    const todos = zodState(
+      z.array(z.object({ id: z.string(), text: z.string() })).default(() => []),
+    );
+
+    const std = standardOf(todos);
+    expect(std).toBeDefined();
+    expect(std?.jsonSchema).toBeDefined();
+    const json = std?.jsonSchema?.input();
+    expect(json).toBeTypeOf("object");
+    if (hasZodV4ToJsonSchema()) {
+      expect(json?.type).toBe("array");
+    }
+  });
+
+  it("makes the wrapped field pass a StandardJSONSchemaV1-style probe", () => {
+    // Mirrors what LangGraph's isStandardJSONSchema() looks for: an `input`
+    // function on `~standard.jsonSchema` that returns a plain object.
+    const schema = zodState(z.object({ liked: z.array(z.string()) }));
+    const std = standardOf(schema);
+
+    expect(std?.jsonSchema).toBeDefined();
+    expect(typeof std?.jsonSchema?.input).toBe("function");
+    expect(typeof std?.jsonSchema?.input()).toBe("object");
+  });
+});

packages/sdk-js/src/langgraph/middleware.ts

@@ -1,7 +1,77 @@
 import { createMiddleware, AIMessage, SystemMessage } from "langchain";
 import type { InteropZodObject } from "@langchain/core/utils/types";
+import type {
+  StandardJSONSchemaV1,
+  StandardSchemaV1,
+} from "@standard-schema/spec";
 import * as z from "zod";
 
+type WithJsonSchema<T> = T extends { "~standard": infer S }
+  ? Omit<T, "~standard"> & {
+      "~standard": S &
+        StandardJSONSchemaV1.Props<
+          S extends StandardSchemaV1.Props<infer I, any> ? I : unknown,
+          S extends StandardSchemaV1.Props<any, infer O> ? O : unknown
+        >;
+    }
+  : T;
+
+/**
+ * Augment a Standard-Schema–compatible schema (e.g. Zod) with a
+ * `~standard.jsonSchema.input` hook so LangGraph's
+ * `getJsonSchemaFromSchema` (called from `StateSchema.getJsonSchema`)
+ * can serialize the field.
+ *
+ * Without this, Zod v4 fields carry `~standard.validate` + `vendor` only,
+ * and `isStandardJSONSchema()` returns false, so the field is silently
+ * dropped from the graph's `output_schema`. That makes AG-UI
+ * `STATE_SNAPSHOT` events filter the field out of the payload sent to
+ * the frontend even though the underlying thread state has the value.
+ *
+ * Use this on any custom state field you want visible to the frontend
+ * via `useAgent().state.*`.
+ *
+ * @example
+ * ```ts
+ * import { zodState } from "@copilotkit/sdk-js/langgraph";
+ *
+ * const stateSchema = z.object({
+ *   todos: zodState(z.array(TodoSchema).default(() => [])),
+ * });
+ * ```
+ */
+export function zodState<T extends object>(schema: T): WithJsonSchema<T> {
+  const std = (schema as { "~standard"?: { jsonSchema?: unknown } })[
+    "~standard"
+  ];
+  if (std && typeof std === "object" && !("jsonSchema" in std)) {
+    let cached: Record<string, unknown> | undefined;
+    std.jsonSchema = {
+      input: () => {
+        if (cached) return cached;
+        // Prefer zod-v4's native `toJSONSchema` when available. Falls back to
+        // an empty object, which is sufficient for the field to appear in the
+        // graph's output_schema (langgraph-api treats it as an opaque field).
+        try {
+          const maybeV4ToJsonSchema = (
+            z as unknown as {
+              toJSONSchema?: (s: unknown) => Record<string, unknown>;
+            }
+          ).toJSONSchema;
+          cached =
+            typeof maybeV4ToJsonSchema === "function"
+              ? maybeV4ToJsonSchema(schema)
+              : {};
+        } catch {
+          cached = {};
+        }
+        return cached;
+      },
+    };
+  }
+  return schema as WithJsonSchema<T>;
+}
+
 /**
  * Internal/framework state keys that should never be auto-surfaced to the
  * LLM as user-facing state. These are reducer-managed message buckets,
@@ -215,14 +285,16 @@ const createAppContextBeforeAgent = (state, runtime) => {
  * ```
  */
 const copilotKitStateSchema = z.object({
-  copilotkit: z
-    .object({
-      actions: z.array(z.any()),
-      context: z.any().optional(),
-      interceptedToolCalls: z.array(z.any()).optional(),
-      originalAIMessageId: z.string().optional(),
-    })
-    .optional(),
+  copilotkit: zodState(
+    z
+      .object({
+        actions: z.array(z.any()),
+        context: z.any().optional(),
+        interceptedToolCalls: z.array(z.any()).optional(),
+        originalAIMessageId: z.string().optional(),
+      })
+      .optional(),
+  ),
 });
 
 const buildMiddlewareInput = (exposeState: ExposeStateOption) => ({