feat(sdk): source A2UI catalog wherever the frontend registered it

The auto-injected generate_a2ui tool now resolves the A2UI catalog from
both delivery paths: the CopilotKit runtime proxy (a copilotkit.context
entry) and the AG-UI native endpoint (state["ag-ui"].a2ui_schema). The
registered catalog id is extracted and bound to generated surfaces so
BYOC custom catalogs render their own components instead of the basic
catalog. Covers @copilotkit/sdk-js and the copilotkit Python SDK.

Author: ranst91

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

@@ -547,6 +547,32 @@ describe("auto-A2UI injection", () => {
     expect(names).toContain("generate_a2ui");
   });
 
+  it("advertises generate_a2ui when the catalog arrives via copilotkit.context (runtime-proxy path)", async () => {
+    const request = makeRequest({
+      state: {
+        messages: [],
+        thread_id: "a2ui-ctx",
+        copilotkit: {
+          context: [
+            {
+              description:
+                "A2UI catalog capabilities: available catalog IDs and custom component definitions.",
+              value:
+                "Available A2UI catalog:\n- declarative-gen-ui-catalog\n  - Card: {...}\n  - Metric: {...}",
+            },
+          ],
+        },
+      },
+      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: [],

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

@@ -19,6 +19,45 @@ const A2UI_DEFAULT_THREAD_KEY = "__copilotkit_a2ui_default__";
 const a2uiThreadKey = (state: any): string =>
   (state?.thread_id as string) || A2UI_DEFAULT_THREAD_KEY;
 
+/**
+ * Find the frontend-registered A2UI catalog wherever it was passed. Returns
+ * `{ compositionGuide?, catalogId? }` when a catalog is present, else `null`
+ * (so the tool is never advertised when the client can't render A2UI). Two
+ * delivery paths, depending on how the agent is served:
+ *  - AG-UI native endpoint → `state["ag-ui"].a2ui_schema` (JSON
+ *    `{ catalogId, components }`); the toolkit reads it from state itself.
+ *  - CopilotKit runtime proxy → a `state.copilotkit.context` entry describing
+ *    the A2UI catalog (catalog id + component schemas as text), passed to the
+ *    subagent via `compositionGuide`.
+ * `catalogId` binds generated surfaces to the frontend's catalog so BYOC
+ * custom catalogs render their own components (not the basic one).
+ */
+const resolveA2uiCatalog = (
+  state: any,
+): { compositionGuide?: string; catalogId?: string } | null => {
+  const a2uiSchema = state?.["ag-ui"]?.a2ui_schema;
+  if (a2uiSchema) {
+    let catalogId: string | undefined;
+    try {
+      const parsed =
+        typeof a2uiSchema === "string" ? JSON.parse(a2uiSchema) : a2uiSchema;
+      catalogId = parsed?.catalogId;
+    } catch {
+      // non-JSON schema — fall back to the toolkit's basic catalog
+    }
+    return { catalogId };
+  }
+  const context = state?.copilotkit?.context;
+  for (const entry of Array.isArray(context) ? context : []) {
+    const description = entry?.description ?? "";
+    const value = entry?.value ?? "";
+    if (!description.includes("A2UI catalog") || !value) continue;
+    const match = /^\s*-\s+(\S+)/m.exec(value);
+    return { compositionGuide: value, catalogId: match?.[1] };
+  }
+  return null;
+};
+
 type WithJsonSchema<T> = T extends { "~standard": infer S }
   ? Omit<T, "~standard"> & {
       "~standard": S &
@@ -335,16 +374,24 @@ 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.
+    // catalog — sourced wherever the FE passed it (CopilotKit runtime proxy via
+    // copilotkit.context, or AG-UI native via ag-ui.a2ui_schema). The catalog's
+    // presence is the signal that the client can render A2UI surfaces, so this
+    // never advertises a tool that would render nowhere while staying fully
+    // zero-config. The model is inferred from request.model; the catalog id
+    // binds surfaces to the FE's catalog; 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);
+    const a2uiCatalog =
+      typeof getA2UITools === "function"
+        ? resolveA2uiCatalog(request.state)
+        : null;
+    if (a2uiCatalog) {
+      const opts: { defaultCatalogId?: string; compositionGuide?: string } = {};
+      if (a2uiCatalog.catalogId) opts.defaultCatalogId = a2uiCatalog.catalogId;
+      if (a2uiCatalog.compositionGuide)
+        opts.compositionGuide = a2uiCatalog.compositionGuide;
+      a2uiTool = getA2UITools(request.model, opts);
       a2uiToolsByThread.set(a2uiThreadKey(request.state), a2uiTool);
     }
 

sdk-python/copilotkit/copilotkit_lg_middleware.py

@@ -318,28 +318,92 @@ def _apply_state_note(self, request: ModelRequest) -> ModelRequest:
     # Auto-A2UI tool injection
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _resolve_a2ui_catalog(state: dict) -> "tuple[str | None, str | None] | None":
+        """Find the frontend-registered A2UI catalog wherever it was passed.
+
+        Returns ``(component_schema, catalog_id)`` when a catalog is present,
+        else ``None`` (so the tool is never advertised when the client can't
+        render A2UI). Two delivery paths are supported, because the catalog
+        lands in different places depending on how the agent is served:
+
+        - **AG-UI native endpoint** → ``state["ag-ui"]["a2ui_schema"]``, a JSON
+          string ``{"catalogId": ..., "components": [...]}``.
+        - **CopilotKit runtime proxy** → a ``state["copilotkit"]["context"]``
+          entry describing the A2UI catalog (catalog id + component schemas as
+          text).
+
+        ``component_schema`` is the text/JSON the subagent should compose from;
+        ``catalog_id`` binds generated surfaces to the frontend's catalog (so
+        BYOC custom catalogs render their own components, not the basic one).
+        """
+        # AG-UI native path.
+        ag_ui = state.get("ag-ui") or {}
+        a2ui_schema = ag_ui.get("a2ui_schema")
+        if a2ui_schema:
+            catalog_id = None
+            try:
+                parsed = (
+                    json.loads(a2ui_schema)
+                    if isinstance(a2ui_schema, str)
+                    else a2ui_schema
+                )
+                if isinstance(parsed, dict):
+                    catalog_id = parsed.get("catalogId")
+            except (TypeError, ValueError):
+                pass
+            # Native path: the toolkit reads ``a2ui_schema`` from state itself,
+            # so no composition_guide is needed — just surface the catalog id.
+            return None, catalog_id
+
+        # CopilotKit runtime-proxy path: the catalog arrives as a context entry.
+        context = (state.get("copilotkit") or {}).get("context") or []
+        for entry in context:
+            if not isinstance(entry, dict):
+                continue
+            description = entry.get("description") or ""
+            value = entry.get("value") or ""
+            if "A2UI catalog" not in description or not value:
+                continue
+            # The value lists catalogs as "- <catalogId>" lines; the first is
+            # the custom catalog the client registered.
+            match = re.search(r"(?m)^\s*-\s+(\S+)", value)
+            catalog_id = match.group(1) if match else None
+            return value, catalog_id
+
+        return None
+
     def _maybe_build_a2ui_tool(self, request: ModelRequest) -> Any | None:
         """Build a ``generate_a2ui`` tool bound to the agent's own model when
-        the frontend has registered an A2UI catalog.
+        the frontend has registered an A2UI catalog — wherever it was passed.
 
-        The catalog is surfaced by the AG-UI runtime into
-        ``state["ag-ui"]["a2ui_schema"]``; its presence is the signal that the
-        client can actually render A2UI surfaces. Gating on it means agents
-        whose frontend has no catalog never see the tool — so we never advertise
-        a tool that would render nowhere — while keeping the feature fully
-        zero-config: the developer only uses the middleware, nothing else.
+        The catalog's presence is the signal that the client can actually render
+        A2UI surfaces, so agents whose frontend has no catalog never see the
+        tool (we never advertise a tool that would render nowhere) while the
+        feature stays fully zero-config: the developer only uses the middleware.
 
         The model is inferred from ``request.model`` (the bound agent model);
-        the built tool is stashed for the tool-call hook to execute. Returns the
-        tool (also appended to the advertised tools) or ``None`` when A2UI is
-        not applicable.
+        the component schema and catalog id come from the registered catalog so
+        the subagent composes the right components and surfaces bind to the
+        frontend's catalog. The built tool is stashed for the tool-call hook to
+        execute. Returns the tool or ``None`` when A2UI is not applicable.
         """
         if get_a2ui_tools is None:
             return None
-        ag_ui = request.state.get("ag-ui") or {}
-        if not ag_ui.get("a2ui_schema"):
+        resolved = self._resolve_a2ui_catalog(request.state or {})
+        if resolved is None:
             return None
-        tool = get_a2ui_tools(request.model)
+        component_schema, catalog_id = resolved
+
+        kwargs: dict[str, Any] = {}
+        if catalog_id:
+            kwargs["default_catalog_id"] = catalog_id
+        # Feed the registered component schema to the subagent so it composes
+        # only catalog components (the toolkit appends this to its prompt).
+        if component_schema:
+            kwargs["composition_guide"] = component_schema
+
+        tool = get_a2ui_tools(request.model, **kwargs)
         _a2ui_tools_by_thread[_current_thread_id() or _DEFAULT_THREAD_KEY] = tool
         return tool
 

sdk-python/tests/test_copilotkit_lg_middleware.py

@@ -1172,3 +1172,60 @@ def test_bridge_cleared_after_agent_stops_execution(self):
         )
 
         assert received.tool is None
+
+    # --- catalog sourced from wherever the frontend passed it ----------------
+
+    def test_injected_from_copilotkit_context(self):
+        """CopilotKit runtime-proxy path: catalog arrives as a context entry."""
+        middleware = CopilotKitMiddleware()
+        state = {
+            "messages": [],
+            "copilotkit": {
+                "context": [
+                    {
+                        "description": "A2UI catalog capabilities: available "
+                        "catalog IDs and custom component definitions.",
+                        "value": "Available A2UI catalog:\n- my-custom-catalog\n"
+                        "  - Card: {...}\n  - Metric: {...}",
+                    }
+                ]
+            },
+        }
+        request = _make_request(state=state, tools=[{"name": "backend"}])
+
+        seen, _ = _run_wrap(middleware, request)
+
+        names = [getattr(t, "name", None) or t.get("name") for t in seen.tools]
+        assert "generate_a2ui" in names
+        assert "backend" in names
+
+    def test_resolve_catalog_from_context_extracts_catalog_id(self):
+        schema, catalog_id = CopilotKitMiddleware._resolve_a2ui_catalog(
+            {
+                "copilotkit": {
+                    "context": [
+                        {
+                            "description": "A2UI catalog capabilities",
+                            "value": "Available A2UI catalog:\n- declarative-gen-ui-catalog\n  ...",
+                        }
+                    ]
+                }
+            }
+        )
+        assert catalog_id == "declarative-gen-ui-catalog"
+        assert schema and "declarative-gen-ui-catalog" in schema
+
+    def test_resolve_catalog_from_native_schema_extracts_catalog_id(self):
+        schema, catalog_id = CopilotKitMiddleware._resolve_a2ui_catalog(
+            {
+                "ag-ui": {
+                    "a2ui_schema": json.dumps({"catalogId": "cat-x", "components": []})
+                }
+            }
+        )
+        # Native path: toolkit reads a2ui_schema from state, so no guide needed.
+        assert schema is None
+        assert catalog_id == "cat-x"
+
+    def test_resolve_catalog_none_when_absent(self):
+        assert CopilotKitMiddleware._resolve_a2ui_catalog({"messages": []}) is None