feat(example): add an assistant-ui chat interface

docs/messages.mdx

@@ -50,9 +50,14 @@ If you want to get MessageDocs, you can use `listMessages` instead.
 
 ## Showing messages in React
 
-See [ChatStreaming.tsx](../example/ui/chat/ChatStreaming.tsx) for a streaming
-example, or [ChatBasic.tsx](../example/ui/chat/ChatBasic.tsx) for a
-non-streaming example.
+Examples:
+
+- [ChatBasic.tsx](../example/ui/chat/ChatBasic.tsx): non-streaming messages.
+- [ChatStreaming.tsx](../example/ui/chat/ChatStreaming.tsx): streaming messages
+  with custom React components.
+- [ChatAssistantUI.tsx](../example/ui/chat/ChatAssistantUI.tsx): streaming
+  messages and tool approvals with
+  [assistant-ui](https://www.assistant-ui.com/).
 
 ### `useUIMessages` hook
 

docs/streaming.mdx

@@ -29,6 +29,8 @@ Example:
 
 - Server: [streaming.ts](../example/convex/chat/streaming.ts)
 - Client: [ChatStreaming.tsx](../example/ui/chat/ChatStreaming.tsx)
+- Client with assistant-ui:
+  [ChatAssistantUI.tsx](../example/ui/chat/ChatAssistantUI.tsx)
 
 ## Streaming message deltas
 

docs/tool-approval.mdx

@@ -12,8 +12,8 @@ action before it happens.
 
 ## Defining tools with approval
 
-Add `needsApproval` to any tool created with `createTool`. It can be a
-boolean or an async function that receives the tool context and input:
+Add `needsApproval` to any tool created with `createTool`. It can be a boolean
+or an async function that receives the tool context and input:
 
 ```ts
 import { createTool } from "@convex-dev/agent";
@@ -56,11 +56,13 @@ The typical approval flow involves four server functions:
 
 1. **Save the user's message** and schedule generation.
 2. **Generate a response.** If the model calls a tool that needs approval,
-   generation pauses and the `tool-approval-request` is persisted in the
-   thread.
+   generation pauses and the `tool-approval-request` is persisted in the thread.
 3. **Submit an approval or denial** for each pending tool call.
    `approveToolCall` and `denyToolCall` work from both mutations and actions.
-4. **Continue generation** from the returned message ID.
+4. **Continue generation** by running the same action with the returned message
+   ID as the prompt. Schedule it from the mutation that saves the last pending
+   decision in the step, so continuation does not depend on the client staying
+   connected.
 
 ```ts
 import { approvalAgent } from "../agents/approval";
@@ -81,7 +83,7 @@ export const sendMessage = mutation({
   },
 });
 
-// 2. Generate (stops if approval is needed)
+// 2. Generate (stops if approval is needed). Also continues after approvals.
 export const generateResponse = internalAction({
   args: { promptMessageId: v.string(), threadId: v.string() },
   handler: async (ctx, { promptMessageId, threadId }) => {
@@ -106,29 +108,21 @@ export const submitApproval = mutation({
   handler: async (ctx, { threadId, approvalId, approved, reason }) => {
     const { messageId } = approved
       ? await approvalAgent.approveToolCall(ctx, {
-          threadId, approvalId, reason,
+          threadId,
+          approvalId,
+          reason,
         })
       : await approvalAgent.denyToolCall(ctx, {
-          threadId, approvalId, reason,
+          threadId,
+          approvalId,
+          reason,
         });
+    // 4. If no approval request in this step is still undecided, schedule
+    // generateResponse with promptMessageId: messageId. See the example
+    // files below.
     return { messageId };
   },
 });
-
-// 4. Continue generation from the saved approval decision.
-//    Pass the last approval message ID as promptMessageId so the agent
-//    resumes generation from where the approval was issued.
-export const continueAfterApprovals = internalAction({
-  args: { threadId: v.string(), lastApprovalMessageId: v.string() },
-  handler: async (ctx, { threadId, lastApprovalMessageId }) => {
-    const result = await approvalAgent.streamText(
-      ctx,
-      { threadId },
-      { promptMessageId: lastApprovalMessageId },
-    );
-    await result.consumeStream();
-  },
-});
 ```
 
 <Admonition type="tip">
@@ -142,9 +136,9 @@ For example, an action can call `approveToolCall` and then immediately call
 ## Handling multiple tool calls
 
 When the model calls several tools in a single step, some or all of them may
-require approval. Submit decisions from the same assistant request together
-with `respondToToolCallApprovals`. The batch uses one bounded lookup and one
-write, and either every decision is saved or none are.
+require approval. Submit decisions from the same assistant request together with
+`respondToToolCallApprovals`. The batch uses one bounded lookup and one write,
+and either every decision is saved or none are.
 
 ```ts
 const { messageId } = await approvalAgent.respondToToolCallApprovals(ctx, {
@@ -169,12 +163,12 @@ single-decision helpers so the request history is read only once.
 If a partial continuation requests another approval, that newer request also
 remains pending while you resume an earlier sibling.
 
-If an unrelated generation starts while approvals are still unresolved, they
-are **automatically denied** with the reason
-`"auto-denied: new generation started"`. This prevents broken message
-histories where tool calls lack results. These denials are synthetic and only
-added to the model's context for that request; they are not saved, so the
-approvals remain pending in storage.
+If an unrelated generation starts while approvals are still unresolved, they are
+**automatically denied** with the reason
+`"auto-denied: new generation started"`. This prevents broken message histories
+where tool calls lack results. These denials are synthetic and only added to the
+model's context for that request; they are not saved, so the approvals remain
+pending in storage.
 
 This does not apply when resuming approvals: on that continuation, earlier
 approvals that were not part of the resume are omitted from that model request.
@@ -187,24 +181,18 @@ Approve/Deny buttons. Tool parts with `state === "approval-requested"` are
 waiting for a decision.
 
 ```tsx
-import { useEffect, useRef } from "react";
 import { useMutation } from "convex/react";
 import { useUIMessages, type UIMessage } from "@convex-dev/agent/react";
 import type { ToolUIPart } from "ai";
 
 function Chat({ threadId }: { threadId: string }) {
-  const lastApprovalMessageIdRef = useRef<string | null>(null);
-
   const { results: messages } = useUIMessages(
     api.chat.approval.listThreadMessages,
     { threadId },
     { initialNumItems: 10, stream: true },
   );
 
   const submitApproval = useMutation(api.chat.approval.submitApproval);
-  const triggerContinuation = useMutation(
-    api.chat.approval.triggerContinuation,
-  );
 
   const hasPendingApprovals = messages.some((m) =>
     m.parts.some(
@@ -214,36 +202,29 @@ function Chat({ threadId }: { threadId: string }) {
     ),
   );
 
-  // When all approvals are resolved, trigger continuation
-  useEffect(() => {
-    if (!hasPendingApprovals && lastApprovalMessageIdRef.current) {
-      void triggerContinuation({
-        threadId,
-        lastApprovalMessageId: lastApprovalMessageIdRef.current,
-      });
-      lastApprovalMessageIdRef.current = null;
-    }
-  }, [hasPendingApprovals, threadId, triggerContinuation]);
-
   // Render approval buttons for tool parts with state "approval-requested"
+  // and call submitApproval. The server continues generation.
   // ...
 }
 ```
 
 The `ToolUIPart` states relevant to approval are:
 
-| State | Meaning |
-|-------|---------|
-| `approval-requested` | Waiting for the user to approve or deny |
+| State                | Meaning                                              |
+| -------------------- | ---------------------------------------------------- |
+| `approval-requested` | Waiting for the user to approve or deny              |
 | `approval-responded` | User responded; tool is being executed (if approved) |
-| `output-available` | Tool executed successfully |
-| `output-denied` | Tool was denied |
-| `output-error` | Tool execution failed |
+| `output-available`   | Tool executed successfully                           |
+| `output-denied`      | Tool was denied                                      |
+| `output-error`       | Tool execution failed                                |
 
 ## Example files
 
 For a complete working example, see:
 
-- **Agent definition:** [`example/convex/agents/approval.ts`](https://github.com/convex/agent/blob/main/example/convex/agents/approval.ts)
-- **Server functions:** [`example/convex/chat/approval.ts`](https://github.com/convex/agent/blob/main/example/convex/chat/approval.ts)
-- **React UI:** [`example/ui/chat/ChatApproval.tsx`](https://github.com/convex/agent/blob/main/example/ui/chat/ChatApproval.tsx)
+- **Agent definition:**
+  [`example/convex/agents/approval.ts`](https://github.com/convex/agent/blob/main/example/convex/agents/approval.ts)
+- **Server functions:**
+  [`example/convex/chat/approval.ts`](https://github.com/convex/agent/blob/main/example/convex/chat/approval.ts)
+- **React UI:**
+  [`example/ui/chat/ChatApproval.tsx`](https://github.com/convex/agent/blob/main/example/ui/chat/ChatApproval.tsx)

example/README.md

@@ -25,3 +25,23 @@ cd agent
 npm run setup
 npm run dev
 ```
+
+## assistant-ui chat
+
+`/chat-assistant-ui` renders the tool approval agent with
+[assistant-ui](https://www.assistant-ui.com/) primitives. It uses the same
+backend as the Tool Approval example (`convex/chat/approval.ts`).
+
+- `ui/chat/assistant-ui/useConvexRuntime.ts` connects `useUIMessages` to
+  `useExternalStoreRuntime`. Sending runs `sendMessage` with
+  `optimisticallySendMessage`, and approvals go through `submitApproval`, which
+  resumes generation once every approval in the step is answered. Stop calls
+  `abortStreamByOrder`.
+- `ui/chat/assistant-ui/convertMessage.ts` maps `UIMessage` parts to
+  assistant-ui parts, keyed by the message's stable `key`.
+- `ui/chat/assistant-ui/Thread.tsx` composes the UI from assistant-ui
+  primitives.
+
+With no provider key and no Convex AI gateway, the demo model replies with mock
+text and never calls tools, so approvals need a real model (see
+`convex/modelsForDemo.ts`).

example/convex/chat/approval.test.ts

@@ -0,0 +1,91 @@
+import { describe, expect, test } from "vitest";
+import { createThread, saveMessages } from "@convex-dev/agent";
+import { api, components } from "../_generated/api";
+import { initConvexTest } from "../setup.test";
+
+async function setup(approvalIds: string[]) {
+  const t = initConvexTest();
+  const threadId = await t.run(async (ctx) => {
+    const threadId = await createThread(ctx, components.agent, {});
+    await saveMessages(ctx, components.agent, {
+      threadId,
+      messages: [
+        { role: "user", content: "Delete both files" },
+        {
+          role: "assistant",
+          content: approvalIds.flatMap((approvalId) => [
+            {
+              type: "tool-call" as const,
+              toolCallId: `call-${approvalId}`,
+              toolName: "deleteFile",
+              input: { filename: `${approvalId}.txt` },
+            },
+            {
+              type: "tool-approval-request" as const,
+              toolCallId: `call-${approvalId}`,
+              approvalId,
+            },
+          ]),
+        },
+      ],
+    });
+    return threadId;
+  });
+  const decide = (approvalId: string, approved = true) =>
+    t.mutation(api.chat.approval.submitApproval, {
+      threadId,
+      approvalId,
+      approved,
+    });
+  const scheduled = () =>
+    t.run(async (ctx) =>
+      (await ctx.db.system.query("_scheduled_functions").collect()).map(
+        ({ name, args }) => ({ name, args }),
+      ),
+    );
+  return { threadId, decide, scheduled };
+}
+
+const continuation = (threadId: string, promptMessageId: string) => ({
+  name: "chat/approval:generateResponse",
+  args: [{ threadId, promptMessageId }],
+});
+
+describe("submitApproval", () => {
+  test("continues once, after the last of parallel decisions", async () => {
+    const { threadId, decide, scheduled } = await setup(["a", "b"]);
+    await decide("a");
+    expect(await scheduled()).toEqual([]);
+    const { messageId } = await decide("b");
+    expect(await scheduled()).toEqual([continuation(threadId, messageId)]);
+  });
+
+  test("continues when a step's decisions fill the lookup window", async () => {
+    const ids = Array.from({ length: 100 }, (_, i) => `a${i}`);
+    const { threadId, decide, scheduled } = await setup(ids);
+    for (const id of ids.slice(0, -1)) await decide(id);
+    expect(await scheduled()).toEqual([]);
+    const { messageId } = await decide(ids.at(-1)!);
+    expect(await scheduled()).toEqual([continuation(threadId, messageId)]);
+  }, 30_000);
+
+  test("continues after a single decision", async () => {
+    const { threadId, decide, scheduled } = await setup(["a"]);
+    const { messageId } = await decide("a");
+    expect(await scheduled()).toEqual([continuation(threadId, messageId)]);
+  });
+
+  test("counts a denial as a decision", async () => {
+    const { threadId, decide, scheduled } = await setup(["a", "b"]);
+    await decide("a");
+    const { messageId } = await decide("b", false);
+    expect(await scheduled()).toEqual([continuation(threadId, messageId)]);
+  });
+
+  test("rejects a decision for a step that already continued", async () => {
+    const { threadId, decide, scheduled } = await setup(["a"]);
+    const { messageId } = await decide("a");
+    await expect(decide("a", false)).rejects.toThrow(/already handled/);
+    expect(await scheduled()).toEqual([continuation(threadId, messageId)]);
+  });
+});