fix(showcase/langgraph-typescript): fix beautiful-chat A2UI format + agentId + declarative-gen-ui tool naming

Author: jpr5

showcase/integrations/langgraph-typescript/src/agent/a2ui-dynamic.ts

@@ -5,20 +5,39 @@
  *
  * Pattern:
  * - The agent binds an explicit `generate_a2ui` tool. When called, `generate_a2ui`
- *   invokes a secondary LLM bound to `render_a2ui` (tool_choice forced) using the
- *   registered client catalog injected as `copilotkit.context`.
+ *   invokes a secondary LLM bound to `_design_a2ui_surface` (tool_choice forced)
+ *   using the registered client catalog injected as `copilotkit.context`.
+ *   The internal tool is intentionally NOT named `render_a2ui` because the A2UI
+ *   middleware default-intercepts tool calls by that name from the run's event
+ *   stream and synthesises ACTIVITY_SNAPSHOT events from the LLM's RAW streaming
+ *   args (catalogId + components, before our code can validate). That bypass
+ *   surfaces "Cannot create component root without a type" infinite-loops.
+ *   Renaming sidesteps the middleware's intercept list (`a2uiToolNames`).
  * - The tool result returns an `a2ui_operations` container which the A2UI
  *   middleware detects in the tool-call result and forwards to the frontend
  *   renderer.
  * - The runtime (see `src/app/api/copilotkit-declarative-gen-ui/route.ts`) uses
  *   `injectA2UITool: false` because the tool binding is owned by the agent here.
+ *
+ * State access for `generate_a2ui`:
+ * The tool needs `state.messages` and `state.copilotkit.context` to forward
+ * conversation history and the A2UI catalog schema to the secondary LLM. The
+ * built-in `ToolNode` does not thread graph state through to tool execution
+ * config. In Python, `ToolRuntime` provides this automatically. In TypeScript
+ * we solve it with a state-aware wrapper node that snapshots state into a
+ * module-level variable before delegating to `ToolNode`. This preserves
+ * the standard LangChain tool invocation path (critical for `OnToolEnd`
+ * events that the AG-UI adapter converts to `TOOL_CALL_RESULT` — which the
+ * A2UI middleware needs to detect `a2ui_operations` in the tool result).
  */
 
 import { z } from "zod";
-import { RunnableConfig } from "@langchain/core/runnables";
 import { tool } from "@langchain/core/tools";
+import {
+  AIMessage,
+  SystemMessage,
+} from "@langchain/core/messages";
 import { ToolNode } from "@langchain/langgraph/prebuilt";
-import { AIMessage, SystemMessage } from "@langchain/core/messages";
 import {
   MemorySaver,
   START,
@@ -52,6 +71,47 @@ const SYSTEM_PROMPT =
   "automatically. Keep chat replies to one short sentence; let the UI do " +
   "the talking.";
 
+// Matches Python's _GENERATE_A2UI_PROMPT_HEADER — instructs the secondary LLM
+// on flat component array format, required fields, and catalog constraints.
+const GENERATE_A2UI_PROMPT_HEADER =
+  `You are designing a dynamic A2UI v0.9 surface. Call the \`_design_a2ui_surface\`\n` +
+  `tool with a flat component array.\n\n` +
+  `Hard requirements (failing any of these breaks the renderer — be strict):\n` +
+  `- \`catalogId\` MUST be exactly: "${CUSTOM_CATALOG_ID}"\n` +
+  `- \`surfaceId\` is a short kebab-case identifier (e.g. "kpi-dashboard").\n` +
+  `- \`components\` is a FLAT array. Every entry MUST include both an \`id\` (unique\n` +
+  `  string) AND a \`component\` (string — the catalog component name). The root\n` +
+  `  entry MUST have \`id: "root"\` AND a valid \`component\` field — never emit\n` +
+  `  a root entry without a component type.\n` +
+  `- Container components (Row, Column, Card) reference children by id via their\n` +
+  `  \`children\` (array of strings) or \`child\` (single string) prop. Do NOT inline\n` +
+  `  children objects. Define each child as its own entry in the flat array and\n` +
+  `  reference its id.\n` +
+  `- Use only catalog component names listed in the schema below.`;
+
+/**
+ * Drop component entries that aren't objects or are missing `id`/`component`.
+ * Mirrors Python's `sanitize_a2ui_components`.
+ */
+function sanitizeA2uiComponents(
+  raw: unknown[],
+): Array<Record<string, unknown>> {
+  return (raw ?? []).filter(
+    (c): c is Record<string, unknown> =>
+      typeof c === "object" &&
+      c !== null &&
+      typeof (c as Record<string, unknown>).id === "string" &&
+      (c as Record<string, unknown>).id !== "" &&
+      typeof (c as Record<string, unknown>).component === "string" &&
+      (c as Record<string, unknown>).component !== "",
+  ) as Array<Record<string, unknown>>;
+}
+
+/** True iff `components` contains an entry with `id === "root"`. */
+function hasRootComponent(components: Array<Record<string, unknown>>): boolean {
+  return components.some((c) => c.id === "root");
+}
+
 const AgentStateAnnotation = Annotation.Root({
   ...CopilotKitStateAnnotation.spec,
 });
@@ -90,20 +150,36 @@ function renderA2uiOperations(operations: unknown[]): string {
   return JSON.stringify({ [A2UI_OPERATIONS_KEY]: operations });
 }
 
-const generateA2ui = tool(
-  async (_args, config?: RunnableConfig) => {
-    // The runtime provided by the A2UI middleware injects a `copilotkit.context`
-    // containing the registered client catalog schema + usage guidelines.
-    const configurable = (config?.configurable ?? {}) as Record<
-      string,
-      unknown
-    >;
-    const state = (configurable.state ?? {}) as Record<string, unknown>;
-    const messages = (state.messages ?? []) as unknown[];
-    const copilotkit = (state.copilotkit ?? {}) as Record<string, unknown>;
-    const contextEntries = (copilotkit.context ?? []) as Array<
-      Record<string, unknown>
-    >;
+// ---------------------------------------------------------------------------
+// State snapshot for tool access
+// ---------------------------------------------------------------------------
+// The built-in ToolNode does not forward graph state through config to tools.
+// Python's ToolRuntime provides `runtime.state` automatically; in TypeScript
+// we snapshot the state into a module-level variable before ToolNode runs.
+// This is safe because LangGraph TS runs graph nodes sequentially within a
+// single thread — no concurrent writes to this variable for the same run.
+let _currentState: AgentState | null = null;
+
+/**
+ * `generate_a2ui` — real LangChain tool invoked by ToolNode.
+ *
+ * Reads state from the module-level `_currentState` snapshot (set by the
+ * `stateAwareToolNode` wrapper before ToolNode executes). Returns an
+ * `a2ui_operations` JSON string that the A2UI middleware detects in the
+ * `TOOL_CALL_RESULT` AG-UI event (emitted via the standard OnToolEnd path).
+ */
+const generateA2uiTool = tool(
+  async () => {
+    const state = _currentState;
+    if (!state) {
+      return JSON.stringify({ error: "No state available for generate_a2ui" });
+    }
+
+    const messages = state.messages ?? [];
+    const copilotkit = state.copilotkit ?? {};
+    const contextEntries = (
+      (copilotkit as Record<string, unknown>).context ?? []
+    ) as Array<Record<string, unknown>>;
 
     const contextText = contextEntries
       .map((entry) =>
@@ -114,17 +190,22 @@ const generateA2ui = tool(
       .filter(Boolean)
       .join("\n\n");
 
-    const renderTool = tool(async () => "rendered", {
-      name: "render_a2ui",
-      description: "Render a dynamic A2UI v0.9 surface.",
+    // Internal tool intentionally named `_design_a2ui_surface` (NOT
+    // `render_a2ui`) to avoid the A2UI middleware's default tool-call
+    // intercept. See module docstring.
+    const designTool = tool(async () => "designed", {
+      name: "_design_a2ui_surface",
+      description: "Design a dynamic A2UI v0.9 surface.",
       schema: z.object({
         surfaceId: z.string().describe("Unique surface identifier."),
         catalogId: z
           .string()
           .describe(`The catalog ID (use "${CUSTOM_CATALOG_ID}").`),
         components: z
           .array(z.record(z.unknown()))
-          .describe("A2UI v0.9 component array."),
+          .describe(
+            "A2UI v0.9 component array (flat format). Every entry MUST have `id` and `component`.",
+          ),
         data: z
           .record(z.unknown())
           .optional()
@@ -133,29 +214,94 @@ const generateA2ui = tool(
     });
 
     const model = new ChatOpenAI({ temperature: 0, model: "gpt-4.1" });
-    const modelWithTool = model.bindTools!([renderTool], {
-      tool_choice: { type: "function", function: { name: "render_a2ui" } },
+    const modelWithTool = model.bindTools!([designTool], {
+      tool_choice: {
+        type: "function",
+        function: { name: "_design_a2ui_surface" },
+      },
     });
 
-    // Drop the last message (the tool-call trigger itself) to mirror Python's
-    // `runtime.state["messages"][:-1]`.
-    const priorMessages = messages.slice(0, -1) as any[];
+    // Prepend the explicit instruction header (matching Python's
+    // _GENERATE_A2UI_PROMPT_HEADER) so the LLM knows about flat-array
+    // constraints, required fields, and the canonical catalog ID.
+    const prompt = `${GENERATE_A2UI_PROMPT_HEADER}\n\n${contextText}`.trim();
+
+    // Drop the last message (the tool-call trigger itself) to mirror
+    // Python's `runtime.state["messages"][:-1]`.
+    const rawPrior = (messages as unknown[]).slice(0, -1) as Array<{
+      _getType?: () => string;
+      type?: string;
+      content?: unknown;
+      tool_calls?: unknown[];
+      additional_kwargs?: Record<string, unknown>;
+    }>;
+
+    // Filter out ToolMessages and AIMessages that are pure tool_call
+    // containers (no text content). The secondary LLM only needs
+    // conversational context (human messages + AI text responses);
+    // sending the graph's internal tool_call/tool_result pairs causes
+    // OpenAI to reject with "tool_calls must be followed by tool
+    // messages responding to each tool_call_id".
+    const priorMessages = rawPrior.filter((msg) => {
+      const msgType =
+        typeof msg._getType === "function" ? msg._getType() : msg.type;
+      if (msgType === "tool") return false;
+      if (msgType === "ai") {
+        const hasToolCalls =
+          (msg.tool_calls as unknown[] | undefined)?.length
+            ? (msg.tool_calls as unknown[]).length > 0
+            : false;
+        const addlToolCalls = (
+          msg.additional_kwargs?.tool_calls as unknown[] | undefined
+        );
+        const hasAddlToolCalls = addlToolCalls
+          ? addlToolCalls.length > 0
+          : false;
+        const hasContent =
+          typeof msg.content === "string" &&
+          (msg.content as string).trim().length > 0;
+        if ((hasToolCalls || hasAddlToolCalls) && !hasContent) return false;
+      }
+      return true;
+    });
 
-    const response = (await modelWithTool.invoke([
-      new SystemMessage({ content: contextText }),
-      ...priorMessages,
-    ])) as AIMessage;
+    let response: AIMessage;
+    try {
+      response = (await modelWithTool.invoke([
+        new SystemMessage({ content: prompt }),
+        ...priorMessages,
+      ])) as AIMessage;
+    } catch (err) {
+      console.error("[a2ui-dynamic] Secondary LLM failed:", err);
+      return JSON.stringify({ error: `Secondary LLM failed: ${err}` });
+    }
 
     if (!response.tool_calls?.length) {
-      return JSON.stringify({ error: "LLM did not call render_a2ui" });
+      return JSON.stringify({
+        error: "LLM did not call _design_a2ui_surface",
+      });
     }
 
-    const args = (response.tool_calls[0].args ?? {}) as Record<string, unknown>;
+    const args = (response.tool_calls[0].args ?? {}) as Record<
+      string,
+      unknown
+    >;
     const surfaceId = (args.surfaceId as string) ?? "dynamic-surface";
-    const catalogId = (args.catalogId as string) ?? CUSTOM_CATALOG_ID;
-    const components = (args.components as unknown[]) ?? [];
+    // Force the canonical catalog ID — the secondary LLM has been observed
+    // hallucinating IDs from sibling demos when context is sparse.
+    const catalogId = CUSTOM_CATALOG_ID;
+    const components = sanitizeA2uiComponents(
+      (args.components as unknown[]) ?? [],
+    );
     const data = (args.data as Record<string, unknown>) ?? {};
 
+    if (!hasRootComponent(components)) {
+      return JSON.stringify({
+        error:
+          "LLM produced no valid root component for the A2UI surface.",
+      });
+    }
+
     const ops: unknown[] = [
       createSurfaceOp(surfaceId, catalogId),
       updateComponentsOp(surfaceId, components),
@@ -169,14 +315,21 @@ const generateA2ui = tool(
   {
     name: "generate_a2ui",
     description:
-      "Generate dynamic A2UI components based on the conversation. A secondary LLM designs the UI schema and data. The result is returned as an a2ui_operations container for the A2UI middleware to detect and forward to the frontend renderer.",
+      "Generate dynamic A2UI components based on the conversation. " +
+      "A secondary LLM designs the UI schema and data. The result is " +
+      "returned as an a2ui_operations container for the A2UI middleware " +
+      "to detect and forward to the frontend renderer.",
     schema: z.object({}),
   },
 );
 
-const tools = [generateA2ui];
+const tools = [generateA2uiTool];
+
+// Standard ToolNode — invokes tools via the LangChain runtime so that
+// `OnToolEnd` events fire and the AG-UI adapter emits `TOOL_CALL_RESULT`.
+const _toolNode = new ToolNode(tools);
 
-async function chatNode(state: AgentState, config: RunnableConfig) {
+async function chatNode(state: AgentState) {
   const model = new ChatOpenAI({ temperature: 0, model: "gpt-4.1" });
 
   const modelWithTools = model.bindTools!([
@@ -186,24 +339,36 @@ async function chatNode(state: AgentState, config: RunnableConfig) {
 
   const systemMessage = new SystemMessage({ content: SYSTEM_PROMPT });
 
-  // Forward the full state into the tool execution context via `configurable`
-  // so `generate_a2ui` can access messages + copilotkit.context.
-  const augmentedConfig: RunnableConfig = {
-    ...config,
-    configurable: {
-      ...(config.configurable ?? {}),
-      state,
-    },
-  };
-
-  const response = await modelWithTools.invoke(
-    [systemMessage, ...state.messages],
-    augmentedConfig,
-  );
+  const response = await modelWithTools.invoke([
+    systemMessage,
+    ...state.messages,
+  ]);
 
   return { messages: response };
 }
 
+/**
+ * State-aware tool node wrapper.
+ *
+ * Snapshots the current graph state into the module-level `_currentState`
+ * variable, then delegates to the real `ToolNode`. This gives
+ * `generateA2uiTool` access to `state.messages` and
+ * `state.copilotkit.context` while preserving the standard LangChain tool
+ * invocation path (OnToolEnd events -> TOOL_CALL_RESULT AG-UI events ->
+ * A2UI middleware detection).
+ */
+async function stateAwareToolNode(
+  state: AgentState,
+  config: Record<string, unknown>,
+) {
+  _currentState = state;
+  try {
+    return await _toolNode.invoke(state, config);
+  } finally {
+    _currentState = null;
+  }
+}
+
 function shouldContinue({ messages, copilotkit }: AgentState) {
   const lastMessage = messages[messages.length - 1] as AIMessage;
 
@@ -221,10 +386,16 @@ function shouldContinue({ messages, copilotkit }: AgentState) {
 
 const workflow = new StateGraph(AgentStateAnnotation)
   .addNode("chat_node", chatNode)
-  .addNode("tool_node", new ToolNode(tools))
+  .addNode(
+    "tool_node",
+    stateAwareToolNode as unknown as typeof chatNode,
+  )
   .addEdge(START, "chat_node")
   .addEdge("tool_node", "chat_node")
-  .addConditionalEdges("chat_node", shouldContinue as any);
+  .addConditionalEdges(
+    "chat_node",
+    shouldContinue as unknown as (state: AgentState) => string,
+  );
 
 const memory = new MemorySaver();