@assistant-ui/react-ag-ui 0.1.3 → 0.1.5

package/src/runtime/AgUiThreadController.ts

@@ -300,22 +300,35 @@ const replaceTextParts = (
   return result;
 };
 
+/**
+ * Settles nested sub-agent messages with `status`. Without `runSubagentIds`,
+ * only running messages settle. With it, only that run's sub-agent messages
+ * settle, and an incomplete `status` also settles one still waiting on its own
+ * tool call: the run that could answer it has ended.
+ */
 const settleNestedContent = (
   content: ThreadAssistantMessage["content"],
   status: MessageStatus,
+  runSubagentIds?: ReadonlySet<string>,
 ) =>
-  mapToolCallPartsDeep(content, (part) =>
-    part.messages
-      ? {
-          ...part,
-          messages: part.messages.map((message) =>
-            message.role === "assistant" && message.status?.type === "running"
-              ? { ...message, status }
-              : message,
-          ),
-        }
-      : part,
-  ).content;
+  mapToolCallPartsDeep(content, (part) => {
+    if (!part.messages) return part;
+    let changed = false;
+    const messages = part.messages.map((message) => {
+      if (message.role !== "assistant") return message;
+      const settles = !runSubagentIds
+        ? message.status?.type === "running"
+        : runSubagentIds.has(message.id) &&
+          (message.status?.type === "running" ||
+            (status.type === "incomplete" &&
+              message.status?.type === "requires-action" &&
+              message.status.reason === "tool-calls"));
+      if (!settles) return message;
+      changed = true;
+      return { ...message, status };
+    });
+    return changed ? { ...part, messages } : part;
+  }).content;
 
 type AgUiRuntimeCallbackName = "onError" | "onCancel";
 
@@ -392,6 +405,8 @@ export class AgUiThreadController {
   // The agent that started the active run. updateOptions can swap this.agent
   // mid-run, and cancelling has to reach the agent holding the live request.
   private activeRunAgent: AbstractAgent | null = null;
+  // Counts runs started, so a finished run can tell a newer one began.
+  private runCount = 0;
   private history: ThreadHistoryAdapter | undefined;
   private readonly assistantHistoryParents = new Map<string, string | null>();
   private readonly subagentParentToolCallIds = new Map<
@@ -400,6 +415,9 @@ export class AgUiThreadController {
   >();
   private readonly snapshotHistoryIds = new Set<string>();
   private readonly persistedHistoryIds = new Set<string>();
+  // Ids whose append has not succeeded yet. A write queued behind that append
+  // runs as an append too, since an update may not move the stored head.
+  private readonly unappendedHistoryIds = new Set<string>();
   private readonly historyWrites = createHistoryWriteChain();
   private _historyImportPromise: Promise<void> | undefined;
   private _loadPromise: Promise<void> | undefined;
@@ -525,6 +543,7 @@ export class AgUiThreadController {
       this.subagentParentToolCallIds.clear();
       this.snapshotHistoryIds.clear();
       this.persistedHistoryIds.clear();
+      this.unappendedHistoryIds.clear();
       for (const { message } of repo.messages) {
         this.persistedHistoryIds.add(message.id);
       }
@@ -683,6 +702,17 @@ export class AgUiThreadController {
     await this.startRun(parentId, config.runConfig);
   }
 
+  /** Deletes a message that has no replies under it, and its stored history entry. */
+  async deleteMessage(messageId: string): Promise<void> {
+    const item = this.session.tryGetMessage(messageId);
+    if (!item) return;
+    const parentId = item.parentId ?? null;
+    this.session.deleteMessage(messageId, parentId);
+    this.snapshotHistoryIds.delete(messageId);
+    this.publishUpdate();
+    await this.history?.delete?.([{ parentId, message: item.message }]);
+  }
+
   async cancel(): Promise<void> {
     this.dispatch({
       type: "historyMutationAdvanced",
@@ -1721,6 +1751,8 @@ export class AgUiThreadController {
     const assistantParentId = parent ? parentId : this.session.headId;
     let assistantMessageId: string | undefined;
     let snapshotAssistantId: string | undefined;
+    const importedRunAssistants = new Set<string>();
+    const runSubagentIds = new Set<string>();
     let replayInterruptMessageId: string | undefined;
     // A snapshot the preserve gate declines still evicts the in-flight
     // assistant; recreating under the cached id on the next content-bearing
@@ -2263,9 +2295,38 @@ export class AgUiThreadController {
         adoptServerMessageId(serverId, true, frameOpenIds);
       },
     });
+    const settleRunSubagents = (status: MessageStatus) => {
+      if (runSubagentIds.size === 0) return;
+      for (const message of this.session.getMessages()) {
+        if (message.role !== "assistant") continue;
+        const content = settleNestedContent(
+          message.content,
+          status,
+          runSubagentIds,
+        );
+        if (content === message.content) continue;
+        this.updateAssistantMessage(message.id, { content });
+        this.persistSettledAssistant(message.id);
+      }
+    };
+    // Writes this run's other assistants in branch order: before the answer
+    // settles, the ones above it; after, the ones below. An adapter that
+    // moves its head on each first write then ends on the branch's tail.
+    const persistRunAssistants = (aboveAnswer: boolean) => {
+      for (const message of this.session.getMessages()) {
+        if (message.id === assistantMessageId) {
+          if (aboveAnswer) return;
+          continue;
+        }
+        if (message.role === "assistant" && !runStartMessageIds.has(message.id))
+          this.persistAssistantHistory(message.id);
+      }
+    };
     let runFinished = false;
     const dispatch = (event: AgUiEvent) => {
       if (this.abortController !== abortController) return;
+      if (event.type === "SUBAGENT_STARTED")
+        runSubagentIds.add(event.subagentRunId);
       if (event.type === "RUN_FINISHED") {
         runFinished = true;
         if (event.outcome?.type === "interrupt")
@@ -2288,6 +2349,7 @@ export class AgUiThreadController {
         event.type === "RUN_FINISHED" ||
         event.type === "RUN_ERROR" ||
         event.type === "RUN_CANCELLED";
+      if (settlingRun) persistRunAssistants(true);
       const beforeSnapshot =
         event.type === "MESSAGES_SNAPSHOT" ? this.session.getMessages() : [];
       try {
@@ -2301,6 +2363,31 @@ export class AgUiThreadController {
           candidateContexts.size ? reconcileReplaySnapshot : undefined,
         );
         if (event.type === "MESSAGES_SNAPSHOT") {
+          // Assistants a snapshot adds belong to the run's history even when
+          // the run never streams into them. Every import resets the pending
+          // marks, so each snapshot marks them again. The import also routes
+          // every message to update; a run assistant this client never wrote
+          // gets an append first, which is what moves the stored head.
+          const previousIds = new Set(beforeSnapshot.map(({ id }) => id));
+          for (const message of this.session.getMessages()) {
+            if (
+              message.role !== "assistant" ||
+              runStartMessageIds.has(message.id)
+            )
+              continue;
+            if (!previousIds.has(message.id))
+              importedRunAssistants.add(message.id);
+            if (
+              !this.persistedHistoryIds.has(message.id) &&
+              this.historyWrites.pending(message.id) === undefined
+            )
+              this.snapshotHistoryIds.delete(message.id);
+          }
+          for (const messageId of importedRunAssistants) {
+            const item = this.session.tryGetMessage(messageId);
+            if (item)
+              this.markPendingAssistantHistory(messageId, item.parentId);
+          }
           for (const message of beforeSnapshot) {
             const context = replayContexts.get(message.id);
             const current = this.session.tryGetMessage(message.id)?.message;
@@ -2332,6 +2419,16 @@ export class AgUiThreadController {
               tail?.role === "assistant" ? tail.id : undefined;
           }
         }
+        if (event.type === "RUN_ERROR") {
+          settleRunSubagents({
+            type: "incomplete",
+            reason: "error",
+            ...(event.message !== undefined ? { error: event.message } : {}),
+          });
+        } else if (event.type === "RUN_CANCELLED") {
+          settleRunSubagents({ type: "incomplete", reason: "cancelled" });
+        }
+        if (settlingRun) persistRunAssistants(false);
       } finally {
         settlingRun = false;
       }
@@ -2340,6 +2437,7 @@ export class AgUiThreadController {
     this.abortController = abortController;
     const runAgentInstance = this.agent;
     this.activeRunAgent = runAgentInstance;
+    const runNumber = ++this.runCount;
 
     let cancelRun = () => dispatch({ type: "RUN_CANCELLED" });
     abortSignal.addEventListener(
@@ -2423,6 +2521,7 @@ export class AgUiThreadController {
         current?.role === "assistant" &&
         current.status?.type === "running"
       ) {
+        persistRunAssistants(true);
         // Finalization is transport cleanup, not proof of a terminal event.
         const status: MessageStatus = {
           type: "incomplete",
@@ -2435,6 +2534,25 @@ export class AgUiThreadController {
           content: settleNestedContent(current.content, status),
         });
       }
+      if (this.abortController === abortController && !runFinished) {
+        // The run ended without RUN_FINISHED, so no sub-agent call it left
+        // open can still be answered.
+        const settled = assistantMessageId
+          ? this.session.tryGetMessage(assistantMessageId)?.message
+          : undefined;
+        settleRunSubagents(
+          settled?.status?.type === "incomplete"
+            ? settled.status
+            : {
+                type: "incomplete",
+                reason: "error",
+                error: "The connection ended before the answer finished.",
+              },
+        );
+      }
+      // Stop may have finished this run already; a newer run owns its own
+      // writes.
+      if (this.runCount === runNumber) persistRunAssistants(false);
       this.finishRun(abortController);
     }
 
@@ -2773,11 +2891,14 @@ export class AgUiThreadController {
       oldId,
       collidesWithExisting ? undefined : newId,
     );
+    if (this.unappendedHistoryIds.delete(oldId) && !collidesWithExisting)
+      this.unappendedHistoryIds.add(newId);
     if (pendingWrite && !collidesWithExisting) {
       void pendingWrite.then(
         () => {
           this.persistedHistoryIds.delete(oldId);
           this.persistedHistoryIds.add(newId);
+          this.unappendedHistoryIds.delete(newId);
         },
         () => {},
       );
@@ -3859,7 +3980,9 @@ export class AgUiThreadController {
 
     if (shouldUpdate) {
       const write = this.historyWrites.chain(messageId, () =>
-        update.call(history, { parentId, message }),
+        this.unappendedHistoryIds.has(messageId)
+          ? this.appendOnce(history, messageId, { parentId, message })
+          : update.call(history, { parentId, message }),
       );
       this.assistantHistoryParents.delete(messageId);
       void write.then(
@@ -3897,6 +4020,15 @@ export class AgUiThreadController {
     );
   }
 
+  private async appendOnce(
+    history: ThreadHistoryAdapter,
+    messageId: string,
+    item: { parentId: string | null; message: ThreadMessage },
+  ): Promise<void> {
+    await history.append(item);
+    this.unappendedHistoryIds.delete(messageId);
+  }
+
   private appendHistoryItem(
     parentId: string | null,
     message: ThreadMessage,
@@ -3905,9 +4037,10 @@ export class AgUiThreadController {
     const pending = this.historyWrites.pending(message.id);
     if (pending) return pending;
 
-    const append = this.history.append.bind(this.history);
+    const history = this.history;
+    this.unappendedHistoryIds.add(message.id);
     const write = this.historyWrites.chain(message.id, () =>
-      append({ parentId, message }),
+      this.appendOnce(history, message.id, { parentId, message }),
     );
     void write.then(
       () => {

package/src/useAgUiRuntime.cloud-run-leases.test.tsx

@@ -4,11 +4,13 @@ import { act, cleanup, render, waitFor } from "@testing-library/react";
 import { HttpAgent, type RunAgentInput } from "@ag-ui/client";
 import {
   AssistantRuntimeProvider,
+  MessageByIndexProvider,
+  useActionBarResume,
   useCloudThreadListAdapter,
   useRemoteThreadListRuntime,
 } from "@assistant-ui/core/react";
 import type { AssistantRuntime } from "@assistant-ui/core";
-import { useAui, type AssistantClient } from "@assistant-ui/store";
+import { useAui, useAuiState, type AssistantClient } from "@assistant-ui/store";
 import { AssistantCloud } from "assistant-cloud";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { useAgUiRuntime } from "./useAgUiRuntime";
@@ -107,6 +109,13 @@ const createServer = (
         if (row) row.content = body.content;
         return json({});
       }
+      if (method === "DELETE" && path === "/threads/remote-1/messages") {
+        for (const id of body.message_ids as string[]) {
+          const index = rows.findIndex((r) => r.id === id);
+          if (index !== -1) rows.splice(index, 1);
+        }
+        return json({});
+      }
       if (/^\/threads\/remote-1\/runs\/[^/]+\/heartbeat$/.test(path))
         return json({ heartbeat_interval_ms: heartbeatIntervalMs });
       if (/^\/threads\/remote-1\/runs\/[^/]+\/end$/.test(path)) return json({});
@@ -192,6 +201,20 @@ const renderApp = async (staleAfter?: number) => {
     aui = useAui();
     return null;
   };
+  const resume: { current: ReturnType<typeof useActionBarResume> | undefined } =
+    { current: undefined };
+  const CaptureResume = () => {
+    resume.current = useActionBarResume();
+    return null;
+  };
+  const ReplyResume = () => {
+    const hasReply = useAuiState((s) => s.thread.messages.length >= 2);
+    return hasReply ? (
+      <MessageByIndexProvider index={1}>
+        <CaptureResume />
+      </MessageByIndexProvider>
+    ) : null;
+  };
   const App = () => {
     const adapter = useCloudThreadListAdapter({
       cloud,
@@ -206,6 +229,7 @@ const renderApp = async (staleAfter?: number) => {
     return (
       <AssistantRuntimeProvider runtime={runtime}>
         <Capture />
+        <ReplyResume />
       </AssistantRuntimeProvider>
     );
   };
@@ -213,19 +237,25 @@ const renderApp = async (staleAfter?: number) => {
   await act(async () => {
     await aui.threads.getLoadThreadsPromise();
   });
-  return { agent, getAui: () => aui, getRuntime: () => runtime };
+  return {
+    agent,
+    getAui: () => aui,
+    getRuntime: () => runtime,
+    getResume: () => resume.current,
+  };
 };
 
 const assistantRow = (
   last_heartbeat_ms_ago: number,
   status: "running" | "ended" = "running",
+  content: readonly unknown[] = [{ type: "text", text: "half an answer" }],
 ): Row => ({
   id: "a1",
   parent_id: "u1",
   format: "aui/v0",
   content: {
     role: "assistant",
-    content: [{ type: "text", text: "half an answer" }],
+    content: [...content],
     status: { type: "incomplete", reason: "cancelled" },
   },
   external_id: "a1",
@@ -304,7 +334,50 @@ describe("useAgUiRuntime with Assistant Cloud run leases", () => {
     expect(reply.id).toBe("agent-msg-1");
   });
 
-  it("shows a reloaded run quiet for longer than staleAfter as interrupted; Resume posts a sibling with its own lease and runs the agent again", async () => {
+  it("shows a reloaded empty run quiet for longer than staleAfter as stale; Resume deletes it and runs the agent again from the user message with no sibling", async () => {
+    useServer([userRow, assistantRow(45_000, "running", [])]);
+    const { agent, getAui, getRuntime, getResume } = await renderApp();
+    await act(async () => {
+      await getRuntime().threads.switchToThread("remote-1");
+    });
+    await waitFor(() =>
+      expect(getAui().thread.message({ index: 1 }).getState().runState).toBe(
+        "stale",
+      ),
+    );
+    await waitFor(() => expect(getResume()?.disabled).toBe(false));
+
+    await act(async () => {
+      getResume()!.resume();
+    });
+    await waitFor(() => expect(agent.runs).toHaveLength(1));
+    await waitFor(() =>
+      expect(createsWith(server.calls, true)).toHaveLength(1),
+    );
+    expect(createsWith(server.calls, true)[0]!.body).toMatchObject({
+      parent_id: "u1",
+      open_run: true,
+    });
+    await act(async () => {
+      agent.runs[0]!.finish();
+    });
+    await waitFor(() =>
+      expect(
+        server.calls.some((c) => c.path.endsWith("/runs/lease_1/end")),
+      ).toBe(true),
+    );
+    const deletes = server.calls.filter((c) => c.method === "DELETE");
+    expect(deletes.map((c) => c.body)).toEqual([{ message_ids: ["a1"] }]);
+    expect(server.rows.map((r) => r.id)).toEqual(["u1", "msg_1"]);
+    const messages = getAui().thread.getState().messages;
+    expect(messages.map((m) => m.role)).toEqual(["user", "assistant"]);
+    expect(messages[1]!.id).not.toBe("a1");
+    expect(getAui().thread.message({ index: 1 }).getState().branchCount).toBe(
+      1,
+    );
+  });
+
+  it("shows a reloaded run quiet for longer than staleAfter as interrupted; Regenerate posts a sibling with its own lease and runs the agent again", async () => {
     useServer([userRow, assistantRow(45_000)]);
     const { agent, getAui, getRuntime } = await renderApp();
     await act(async () => {
@@ -314,7 +387,12 @@ describe("useAgUiRuntime with Assistant Cloud run leases", () => {
       expect(getAui().thread.getState().messages).toHaveLength(2),
     );
     const stale = getAui().thread.message({ index: 1 });
-    await waitFor(() => expect(stale.getState().runState).toBe("stale"));
+    await waitFor(() =>
+      expect(stale.getState().status).toMatchObject({
+        error: { code: "run_stale" },
+      }),
+    );
+    expect(stale.getState().runState).toBeUndefined();
 
     await act(async () => {
       stale.reload();
@@ -363,8 +441,8 @@ describe("useAgUiRuntime with Assistant Cloud run leases", () => {
     });
     await waitFor(() =>
       expect(
-        short.getAui().thread.message({ index: 1 }).getState().runState,
-      ).toBe("stale"),
+        short.getAui().thread.message({ index: 1 }).getState().status,
+      ).toMatchObject({ error: { code: "run_stale" } }),
     );
   });
 

package/src/useAgUiRuntime.ts

@@ -345,6 +345,7 @@ export function useAgUiRuntime(
           queueController?.clear();
           return core.reload(parentId, config);
         },
+        onDelete: (messageId: string) => core.deleteMessage(messageId),
         onCancel: async () => {
           queueController?.notifyCancelled();
           core.cancel();