feat(sdk): auto-inject A2UI tool in CopilotKitMiddleware

Prebuilt agents get dynamic A2UI with no extra wiring — adding the
middleware is enough. When the frontend registers an A2UI catalog
(surfaced by the runtime into state["ag-ui"].a2ui_schema), the
middleware infers the agent's own model, advertises the generate_a2ui
tool in the model-call hook, and executes it in the tool-call hook.
No catalog → the tool is never advertised.

Covers both @copilotkit/sdk-js and the copilotkit Python SDK. Bumps
the A2UI tool-factory dependency to where get_a2ui_tools ships
(@ag-ui/langgraph 0.0.35, ag-ui-langgraph >=0.0.37).

Author: ranst91

packages/sdk-js/package.json

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

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

@@ -498,6 +498,114 @@ describe("afterAgent", () => {
   });
 });
 
+// ---------------------------------------------------------------------------
+// Auto-A2UI — middleware injects + executes generate_a2ui when the frontend
+// registered a catalog (surfaced into state["ag-ui"].a2ui_schema)
+// ---------------------------------------------------------------------------
+//
+// Contract: the developer passes nothing — using the middleware is enough.
+// generate_a2ui is advertised to the model only when an A2UI catalog is
+// present, is built from the agent's own (inferred) model, and is executed by
+// the middleware itself (it is never in the agent's static tool registry).
+
+async function runWrapTool(middleware: any, request: any) {
+  let received: any = null;
+  const handler = async (req: any) => {
+    received = req;
+    return { content: "tool-ok" } as any;
+  };
+  await middleware.wrapToolCall(request, handler);
+  return received;
+}
+
+describe("auto-A2UI injection", () => {
+  it("does NOT advertise generate_a2ui when there is no A2UI catalog", async () => {
+    const request = makeRequest({
+      state: { messages: [], thread_id: "a2ui-off" },
+      tools: [{ name: "backend" }],
+    });
+
+    const { received } = await runWrap(copilotkitMiddleware, request);
+
+    expect(received.tools.map((t: any) => t.name)).toEqual(["backend"]);
+  });
+
+  it("advertises generate_a2ui (alongside existing tools) when a catalog is present", async () => {
+    const request = makeRequest({
+      state: {
+        messages: [],
+        thread_id: "a2ui-on",
+        "ag-ui": { a2ui_schema: "<components/>" },
+      },
+      tools: [{ name: "backend" }],
+    });
+
+    const { received } = await runWrap(copilotkitMiddleware, request);
+
+    const names = received.tools.map((t: any) => t.name);
+    expect(names).toContain("backend");
+    expect(names).toContain("generate_a2ui");
+  });
+
+  it("executes generate_a2ui via wrapToolCall using the inferred model", async () => {
+    const state = {
+      messages: [],
+      thread_id: "a2ui-exec",
+      "ag-ui": { a2ui_schema: "<components/>" },
+    };
+    // First the model call infers the model + stashes the built tool.
+    await runWrap(copilotkitMiddleware, makeRequest({ state, tools: [] }));
+
+    const received = await runWrapTool(copilotkitMiddleware, {
+      toolCall: { name: "generate_a2ui", id: "1", args: {} },
+      tool: undefined,
+      state,
+      runtime: {},
+    });
+
+    expect(received.tool).toBeDefined();
+    expect(received.tool.name).toBe("generate_a2ui");
+  });
+
+  it("leaves non-A2UI tool calls untouched", async () => {
+    const state = {
+      messages: [],
+      thread_id: "a2ui-other",
+      "ag-ui": { a2ui_schema: "<components/>" },
+    };
+    await runWrap(copilotkitMiddleware, makeRequest({ state, tools: [] }));
+
+    const backendTool = { name: "backend" };
+    const received = await runWrapTool(copilotkitMiddleware, {
+      toolCall: { name: "backend", id: "1", args: {} },
+      tool: backendTool,
+      state,
+      runtime: {},
+    });
+
+    expect(received.tool).toBe(backendTool);
+  });
+
+  it("stops executing generate_a2ui after the run ends (afterAgent clears the bridge)", async () => {
+    const state = {
+      messages: [],
+      thread_id: "a2ui-clean",
+      "ag-ui": { a2ui_schema: "<components/>" },
+    };
+    await runWrap(copilotkitMiddleware, makeRequest({ state, tools: [] }));
+    copilotkitMiddleware.afterAgent(state, {} as any);
+
+    const received = await runWrapTool(copilotkitMiddleware, {
+      toolCall: { name: "generate_a2ui", id: "1", args: {} },
+      tool: undefined,
+      state,
+      runtime: {},
+    });
+
+    expect(received.tool).toBeUndefined();
+  });
+});
+
 // ---------------------------------------------------------------------------
 // zodState — Standard-Schema JSON-schema augmentation
 // ---------------------------------------------------------------------------

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

@@ -5,8 +5,20 @@ import type {
   StandardSchemaV1,
 } from "@standard-schema/spec";
 import * as z from "zod";
+import { getA2UITools } from "@ag-ui/langgraph";
 import { getForwardedHeaders } from "../header-propagation";
 
+// ---------------------------------------------------------------------------
+// Auto-A2UI: bridge the inferred model's generate_a2ui tool from wrapModelCall
+// (the only hook that exposes the bound model) to wrapToolCall (where the tool
+// actually executes but the model is absent). Keyed by the run's thread id so
+// concurrent runs don't clobber each other.
+// ---------------------------------------------------------------------------
+const a2uiToolsByThread = new Map<string, any>();
+const A2UI_DEFAULT_THREAD_KEY = "__copilotkit_a2ui_default__";
+const a2uiThreadKey = (state: any): string =>
+  (state?.thread_id as string) || A2UI_DEFAULT_THREAD_KEY;
+
 type WithJsonSchema<T> = T extends { "~standard": infer S }
   ? Omit<T, "~standard"> & {
       "~standard": S &
@@ -322,25 +334,60 @@ const buildMiddlewareInput = (exposeState: ExposeStateOption) => ({
       };
     }
 
+    // Auto-inject generate_a2ui when the frontend has registered an A2UI
+    // catalog. The AG-UI runtime surfaces that catalog into
+    // state["ag-ui"].a2ui_schema; its presence is the signal that the client
+    // can actually render A2UI surfaces. Gating on it keeps the feature fully
+    // zero-config (the developer only uses the middleware) while never
+    // advertising a tool that would render nowhere. The model is inferred from
+    // request.model; the built tool is stashed for wrapToolCall to execute.
+    let a2uiTool: any = null;
+    const a2uiSchema = (request.state["ag-ui"] as any)?.a2ui_schema;
+    if (a2uiSchema && typeof getA2UITools === "function") {
+      a2uiTool = getA2UITools(request.model);
+      a2uiToolsByThread.set(a2uiThreadKey(request.state), a2uiTool);
+    }
+
     const frontendTools = request.state["copilotkit"]?.actions ?? [];
 
-    if (frontendTools.length === 0) {
+    if (frontendTools.length === 0 && !a2uiTool) {
       return handler(request);
     }
 
     const existingTools = request.tools || [];
-    const mergedTools = [...existingTools, ...frontendTools];
+    const mergedTools = [
+      ...existingTools,
+      ...(a2uiTool ? [a2uiTool] : []),
+      ...frontendTools,
+    ];
 
     return handler({
       ...request,
       tools: mergedTools,
     });
   },
 
+  // Execute the dynamically-advertised generate_a2ui tool. It is not in the
+  // agent's static tool registry, so the tool node cannot run it on its own;
+  // we supply the implementation (built with the inferred model) for that one
+  // tool. This hook's presence also disables createAgent's "unknown tool"
+  // guard for dynamically-advertised tools.
+  wrapToolCall: async (request: any, handler: (req: any) => Promise<any>) => {
+    const tool = a2uiToolsByThread.get(a2uiThreadKey(request.state));
+    if (tool && !request.tool && request.toolCall?.name === tool.name) {
+      return handler({ ...request, tool });
+    }
+    return handler(request);
+  },
+
   beforeAgent: createAppContextBeforeAgent,
 
   // Restore frontend tool calls to AIMessage before agent exits
   afterAgent: (state) => {
+    // Drop the bridged A2UI tool for this run — all tool calls for the turn
+    // have executed by now; the next model call re-stashes if needed.
+    a2uiToolsByThread.delete(a2uiThreadKey(state));
+
     const interceptedToolCalls = state["copilotkit"]?.interceptedToolCalls;
     const originalMessageId = state["copilotkit"]?.originalAIMessageId;