fix(a2ui): declare ag-ui state channel + a2ui_params host override (py + sdk-js)

packages/runtime/package.json

@@ -77,11 +77,11 @@
     "attw": "attw --pack . --profile node16"
   },
   "dependencies": {
-    "@ag-ui/a2ui-middleware": "0.0.8",
+    "@ag-ui/a2ui-middleware": "0.0.10",
     "@ag-ui/client": "0.0.57",
     "@ag-ui/core": "0.0.57",
     "@ag-ui/encoder": "0.0.57",
-    "@ag-ui/langgraph": "0.0.41",
+    "@ag-ui/langgraph": "0.0.42",
     "@ag-ui/mcp-apps-middleware": "0.0.3",
     "@ag-ui/mcp-middleware": "0.0.1",
     "@ai-sdk/anthropic": "^3.0.49",

packages/sdk-js/package.json

@@ -59,7 +59,7 @@
     "attw": "attw --pack . --profile node16"
   },
   "dependencies": {
-    "@ag-ui/langgraph": "0.0.41",
+    "@ag-ui/langgraph": "0.0.42",
     "@copilotkit/shared": "workspace:*"
   },
   "devDependencies": {

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

@@ -0,0 +1,134 @@
+/**
+ * Tests for the host `a2uiParams` override on the auto-injected generate_a2ui
+ * tool (the design-guidelines parity gap).
+ *
+ * Before this, the auto-inject path hardwired getA2UITools to the toolkit
+ * defaults: a host serving via copilotkitMiddleware could NOT override the
+ * design/generation guidelines (e.g. to favor a repeating-card layout). The
+ * `a2uiParams` option threads a host override through; the middleware still
+ * injects the bound model and folds the registered catalog in (host wins).
+ *
+ * We mock @ag-ui/langgraph's getA2UITools to capture the params it is handed —
+ * the real factory's bound guidelines are a closure the tool object never
+ * exposes, so capture-at-the-boundary is the only way to assert what reached
+ * the subagent. Isolated in its own file so the module-wide mock does not
+ * affect the behavior tests in middleware.test.ts.
+ */
+
+import { describe, it, expect, beforeEach, vi } from "vitest";
+
+const { captured } = vi.hoisted(() => ({ captured: [] as any[] }));
+
+vi.mock("@ag-ui/langgraph", () => ({
+  getA2UITools: (params: any) => {
+    captured.push(params);
+    return { name: "generate_a2ui" };
+  },
+}));
+
+// Imported AFTER vi.mock so the middleware binds the mocked getA2UITools.
+import { createCopilotkitMiddleware } from "../middleware";
+
+function makeRequest(overrides: any = {}): any {
+  return {
+    model: { _modelType: () => "fake" },
+    messages: [],
+    systemPrompt: undefined,
+    tools: [],
+    state: { messages: [] },
+    runtime: {},
+    ...overrides,
+  };
+}
+
+async function runWrap(middleware: any, request: any) {
+  const handler = async (_req: any) => ({ content: "ok" }) as any;
+  await middleware.wrapModelCall(request, handler);
+}
+
+// Native path: schema present but non-JSON, so no catalogId / compositionGuide.
+const NATIVE_STATE = {
+  messages: [],
+  thread_id: "native",
+  "ag-ui": { a2ui_schema: "<components/>", inject_a2ui_tool: true },
+};
+
+// Runtime-proxy path: the catalog arrives as a context entry, so the middleware
+// derives a compositionGuide + catalogId to fold in.
+const CONTEXT_STATE = {
+  messages: [],
+  thread_id: "context",
+  "ag-ui": { inject_a2ui_tool: true },
+  copilotkit: {
+    context: [
+      {
+        description: "A2UI catalog capabilities",
+        value: "Available A2UI catalog:\n- my-custom-catalog\n  - Card: {...}",
+      },
+    ],
+  },
+};
+
+describe("auto-A2UI host a2uiParams override", () => {
+  beforeEach(() => {
+    captured.length = 0;
+  });
+
+  it("threads host guidelines through to getA2UITools (native path)", async () => {
+    const middleware = createCopilotkitMiddleware({
+      a2uiParams: { guidelines: { designGuidelines: "REPEAT_CARDS_MARK" } },
+    });
+    const request = makeRequest({ state: { ...NATIVE_STATE } });
+
+    await runWrap(middleware, request);
+
+    expect(captured).toHaveLength(1);
+    const params = captured[0];
+    // Host override survives...
+    expect(params.guidelines.designGuidelines).toBe("REPEAT_CARDS_MARK");
+    // ...the middleware still injects the bound model...
+    expect(params.model).toBe(request.model);
+    // ...and the native path contributes no compositionGuide.
+    expect(params.guidelines.compositionGuide).toBeUndefined();
+  });
+
+  it("merges host guidelines with the registered catalog (context path)", async () => {
+    const middleware = createCopilotkitMiddleware({
+      a2uiParams: { guidelines: { designGuidelines: "REPEAT_CARDS_MARK" } },
+    });
+
+    await runWrap(middleware, makeRequest({ state: { ...CONTEXT_STATE } }));
+
+    const params = captured[0];
+    expect(params.guidelines.designGuidelines).toBe("REPEAT_CARDS_MARK");
+    expect(params.guidelines.compositionGuide).toContain("my-custom-catalog");
+    expect(params.defaultCatalogId).toBe("my-custom-catalog");
+  });
+
+  it("host compositionGuide + defaultCatalogId win over the catalog", async () => {
+    const middleware = createCopilotkitMiddleware({
+      a2uiParams: {
+        defaultCatalogId: "host-catalog",
+        guidelines: { compositionGuide: "HOST_COMP" },
+      },
+    });
+
+    await runWrap(middleware, makeRequest({ state: { ...CONTEXT_STATE } }));
+
+    const params = captured[0];
+    expect(params.guidelines.compositionGuide).toBe("HOST_COMP");
+    expect(params.defaultCatalogId).toBe("host-catalog");
+  });
+
+  it("default (no a2uiParams) carries only the inferred model", async () => {
+    const middleware = createCopilotkitMiddleware();
+    const request = makeRequest({ state: { ...NATIVE_STATE } });
+
+    await runWrap(middleware, request);
+
+    const params = captured[0];
+    expect(params.model).toBe(request.model);
+    expect(params.guidelines).toBeUndefined();
+    expect(params.defaultCatalogId).toBeUndefined();
+  });
+});

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

@@ -359,7 +359,10 @@ const copilotKitStateSchema = z.object({
   ),
 });
 
-const buildMiddlewareInput = (exposeState: ExposeStateOption) => ({
+const buildMiddlewareInput = (
+  exposeState: ExposeStateOption,
+  a2uiParams?: Omit<A2UIToolParams, "model">,
+) => ({
   name: "CopilotKitMiddleware",
 
   stateSchema: copilotKitStateSchema as unknown as InteropZodObject,
@@ -395,12 +398,24 @@ const buildMiddlewareInput = (exposeState: ExposeStateOption) => ({
     if (typeof getA2UITools === "function" && decision) {
       const catalog = resolveA2uiCatalog(request.state);
       // Shared A2UIToolParams: a single params object owned by the toolkit.
-      // `model` lives inside it; `compositionGuide` is folded into the
-      // `guidelines` bag alongside generation/design overrides.
-      const params: A2UIToolParams = { model: request.model };
-      if (catalog?.catalogId) params.defaultCatalogId = catalog.catalogId;
-      if (catalog?.compositionGuide)
-        params.guidelines = { compositionGuide: catalog.compositionGuide };
+      // Start from the host overrides (guidelines / catalog id / tool name /
+      // recovery) so a host can steer the subagent, then layer in only what the
+      // host cannot know — the bound model, and the registered catalog id +
+      // compositionGuide — without clobbering any host-set value.
+      const params: A2UIToolParams = {
+        ...(a2uiParams ?? {}),
+        model: request.model,
+      };
+      if (catalog?.catalogId && params.defaultCatalogId == null)
+        params.defaultCatalogId = catalog.catalogId;
+      // Merge the registered catalog schema into any host `guidelines` bag; a
+      // host-set compositionGuide wins, host generation/design overrides stay.
+      if (catalog?.compositionGuide) {
+        const guidelines = { ...(params.guidelines ?? {}) };
+        if (guidelines.compositionGuide == null)
+          guidelines.compositionGuide = catalog.compositionGuide;
+        params.guidelines = guidelines;
+      }
       const candidate = getA2UITools(params);
       const existingNames = new Set(
         (request.tools || []).map((t: any) => t?.name),
@@ -546,22 +561,36 @@ const buildMiddlewareInput = (exposeState: ExposeStateOption) => ({
  * 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).
+ * (for example to hide a sensitive key, or to use an explicit allowlist), or
+ * to steer the auto-injected `generate_a2ui` subagent via `a2uiParams`.
+ *
+ * `a2uiParams` is an `A2UIToolParams` without `model` (the middleware always
+ * injects the bound model). Use it to override the subagent guidelines
+ * (`generationGuidelines` / `designGuidelines` / `compositionGuide`),
+ * `defaultCatalogId`, `toolName`, `recovery`, etc. on the auto-inject path —
+ * which otherwise only ever uses the toolkit defaults. The registered catalog
+ * is still folded in, but host-set values win.
  *
  * @example
  * ```typescript
  * import { createCopilotkitMiddleware } from "@copilotkit/sdk-js/langgraph";
  *
  * const middleware = createCopilotkitMiddleware({
  *   exposeState: ["liked", "todos"],
+ *   a2uiParams: { guidelines: { designGuidelines: "...repeating-card layout..." } },
  * });
  * ```
  */
 export const createCopilotkitMiddleware = (
-  options: { exposeState?: ExposeStateOption } = {},
+  options: {
+    exposeState?: ExposeStateOption;
+    a2uiParams?: Omit<A2UIToolParams, "model">;
+  } = {},
 ) => {
   const exposeState = options.exposeState ?? false;
-  return createMiddleware(buildMiddlewareInput(exposeState) as any);
+  return createMiddleware(
+    buildMiddlewareInput(exposeState, options.a2uiParams) as any,
+  );
 };
 
 /**