@copilotkit/sdk-js 1.56.4 → 1.56.5

package/src/langgraph/middleware.ts

@@ -1,7 +1,175 @@
 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,
+ * CopilotKit/AG-UI plumbing, or graph-internal scaffolding.
+ */
+const RESERVED_STATE_KEYS: ReadonlySet<string> = new Set([
+  "messages",
+  "copilotkit",
+  "ag-ui",
+  "tools",
+  "structured_response",
+  "thread_id",
+  "remaining_steps",
+]);
+
+/**
+ * Controls how user-defined state keys are surfaced into the LLM prompt
+ * on every model call. Off by default to avoid leaking arbitrary state
+ * into prompts; opt in explicitly.
+ *
+ * - `false` (default) — never surface state.
+ * - `true` — every state key not in the reserved internal set and not
+ *   prefixed with `_` is JSON-serialized into a "Current agent state:"
+ *   note appended to the system prompt.
+ * - `string[]` — only surface the named keys (use this when you want
+ *   explicit control over what the LLM sees, e.g. `["liked", "todos"]`).
+ */
+export type ExposeStateOption = boolean | readonly string[];
+
+const buildStateNote = (
+  state: Record<string, unknown>,
+  expose: ExposeStateOption,
+): string | null => {
+  if (expose === false) return null;
+
+  const allow: ReadonlySet<string> | null = Array.isArray(expose)
+    ? new Set(expose)
+    : null;
+
+  const snapshot: Record<string, unknown> = {};
+  for (const key of Object.keys(state)) {
+    if (
+      allow
+        ? !allow.has(key)
+        : RESERVED_STATE_KEYS.has(key) || key.startsWith("_")
+    ) {
+      continue;
+    }
+    const value = state[key];
+    if (
+      value === undefined ||
+      value === null ||
+      value === "" ||
+      (Array.isArray(value) && value.length === 0) ||
+      (typeof value === "object" &&
+        !Array.isArray(value) &&
+        Object.keys(value as Record<string, unknown>).length === 0)
+    ) {
+      continue;
+    }
+    snapshot[key] = value;
+  }
+
+  if (Object.keys(snapshot).length === 0) return null;
+
+  let body: string;
+  try {
+    body = JSON.stringify(snapshot, null, 2);
+  } catch {
+    body = String(snapshot);
+  }
+  return `Current agent state:\n${body}`;
+};
+
+const applyStateNote = (request: any, expose: ExposeStateOption): any => {
+  const note = buildStateNote(
+    (request.state ?? {}) as Record<string, unknown>,
+    expose,
+  );
+  if (!note) return request;
+
+  const existing = request.systemPrompt;
+  if (existing == null) {
+    return { ...request, systemPrompt: new SystemMessage({ content: note }) };
+  }
+  // existing may be a string OR a SystemMessage
+  const baseText =
+    typeof existing === "string"
+      ? existing
+      : typeof existing.content === "string"
+        ? existing.content
+        : String(existing.content);
+  return {
+    ...request,
+    systemPrompt: new SystemMessage({ content: `${baseText}\n\n${note}` }),
+  };
+};
+
 const createAppContextBeforeAgent = (state, runtime) => {
   const messages = state.messages;
 
@@ -117,23 +285,26 @@ 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 middlewareInput = {
+const buildMiddlewareInput = (exposeState: ExposeStateOption) => ({
   name: "CopilotKitMiddleware",
 
   stateSchema: copilotKitStateSchema as unknown as InteropZodObject,
 
-  // Inject frontend tools before model call
-  wrapModelCall: async (request, handler) => {
+  // Inject frontend tools and surface user state before model call
+  wrapModelCall: async (request: any, handler: (req: any) => Promise<any>) => {
+    request = applyStateNote(request, exposeState);
     const frontendTools = request.state["copilotkit"]?.actions ?? [];
 
     if (frontendTools.length === 0) {
@@ -234,9 +405,33 @@ const middlewareInput = {
       },
     };
   },
-} as any;
-const createCopilotKitMiddleware = () => {
-  return createMiddleware(middlewareInput);
+});
+
+/**
+ * Build a CopilotKit middleware instance with custom options.
+ *
+ * Use this when you want to override the default state-exposure behavior
+ * (for example to hide a sensitive key, or to use an explicit allowlist).
+ *
+ * @example
+ * ```typescript
+ * import { createCopilotkitMiddleware } from "@copilotkit/sdk-js/langgraph";
+ *
+ * const middleware = createCopilotkitMiddleware({
+ *   exposeState: ["liked", "todos"],
+ * });
+ * ```
+ */
+export const createCopilotkitMiddleware = (
+  options: { exposeState?: ExposeStateOption } = {},
+) => {
+  const exposeState = options.exposeState ?? false;
+  return createMiddleware(buildMiddlewareInput(exposeState) as any);
 };
 
-export const copilotkitMiddleware = createCopilotKitMiddleware();
+/**
+ * Default CopilotKit middleware singleton — does NOT surface user state
+ * to the LLM. Pass `exposeState: true` (or an allowlist) to
+ * {@link createCopilotkitMiddleware} to opt in.
+ */
+export const copilotkitMiddleware = createCopilotkitMiddleware();