@assistant-ui/react-ag-ui 0.1.7 → 0.1.9

package/src/runtime/AgUiThreadController.notifications.test.ts

@@ -0,0 +1,126 @@
+import type { HttpAgent } from "@ag-ui/client";
+import { expect, it, vi } from "vitest";
+import { AgUiThreadController } from "./AgUiThreadController";
+import type { AgUiSubscriber } from "./adapter/subscriber";
+import { makeLogger } from "./logger";
+
+it("notifies subscribers synchronously when a delta callback cancels the run", async () => {
+  vi.useFakeTimers();
+  let subscriber!: AgUiSubscriber;
+  let finishRun!: () => void;
+  const agent = {
+    runAgent: (_input: unknown, current: AgUiSubscriber) => {
+      subscriber = current;
+      return new Promise<void>((resolve) => {
+        finishRun = resolve;
+      });
+    },
+    abortRun: vi.fn(),
+  } as unknown as HttpAgent;
+  const runningStates: boolean[] = [];
+  let runningStatesAtCancelReturn: boolean[] = [];
+  const error = vi.fn(() => {
+    void controller.cancel();
+    runningStatesAtCancelReturn = [...runningStates];
+  });
+  const controller = new AgUiThreadController({
+    agent,
+    logger: makeLogger({ error }),
+    showThinking: true,
+  });
+  const pending = controller.reload(null);
+
+  try {
+    subscriber.onTextMessageStartEvent?.({
+      event: { type: "TEXT_MESSAGE_START", messageId: "answer" },
+    });
+    controller.subscribe(() => {
+      runningStates.push(controller.getSnapshot().isRunning);
+    });
+    subscriber.onTextMessageContentEvent?.({
+      event: {
+        type: "TEXT_MESSAGE_CONTENT",
+        messageId: "answer",
+        delta: "partial",
+      },
+    });
+    expect(runningStates).toEqual([]);
+    expect(controller.getSnapshot().isRunning).toBe(true);
+
+    subscriber.onStateDeltaEvent?.({
+      event: {
+        type: "STATE_DELTA",
+        delta: [{ op: "replace", path: "/missing", value: 1 }],
+      },
+    });
+
+    expect(error).toHaveBeenCalledOnce();
+    expect(runningStatesAtCancelReturn.length).toBeGreaterThan(0);
+    expect(runningStatesAtCancelReturn.at(-1)).toBe(false);
+    expect(controller.getSnapshot().isRunning).toBe(false);
+    vi.runAllTimers();
+    expect(runningStates).toEqual(runningStatesAtCancelReturn);
+  } finally {
+    finishRun();
+    await pending;
+    vi.useRealTimers();
+  }
+});
+
+it("notifies once for a text delta burst while the stream stays open", async () => {
+  vi.useFakeTimers();
+  let subscriber!: AgUiSubscriber;
+  let finishRun!: () => void;
+  const agent = {
+    runAgent: (_input: unknown, current: AgUiSubscriber) => {
+      subscriber = current;
+      return new Promise<void>((resolve) => {
+        finishRun = resolve;
+      });
+    },
+    abortRun: vi.fn(),
+  } as unknown as HttpAgent;
+  const controller = new AgUiThreadController({
+    agent,
+    logger: makeLogger(),
+    showThinking: true,
+  });
+  const pending = controller.reload(null);
+
+  try {
+    subscriber.onTextMessageStartEvent?.({
+      event: { type: "TEXT_MESSAGE_START", messageId: "answer" },
+    });
+    const listener = vi.fn(() => controller.getSnapshot());
+    controller.subscribe(listener);
+    for (const delta of ["hello", " ", "world"]) {
+      subscriber.onTextMessageContentEvent?.({
+        event: { type: "TEXT_MESSAGE_CONTENT", messageId: "answer", delta },
+      });
+    }
+
+    expect(listener).not.toHaveBeenCalled();
+    vi.runAllTimers();
+    expect(listener).toHaveBeenCalledOnce();
+    expect(listener.mock.results[0]?.value).toMatchObject({
+      isRunning: true,
+      messageRepository: {
+        messages: [
+          {
+            message: {
+              role: "assistant",
+              content: [{ type: "text", text: "hello world" }],
+              status: { type: "running" },
+            },
+          },
+        ],
+      },
+    });
+    vi.runAllTimers();
+    expect(listener).toHaveBeenCalledOnce();
+  } finally {
+    finishRun();
+    await pending;
+    vi.useRealTimers();
+  }
+});

package/src/runtime/AgUiThreadController.ts

@@ -218,6 +218,28 @@ const isActivityPart = (part: ThreadAssistantMessage["content"][number]) =>
 type ResumeStream = (
   options: ChatModelRunOptions,
 ) => AsyncGenerator<ChatModelRunResult, void, unknown>;
+/** A run already in flight on the agent, joined through `agent.connectAgent()`. */
+type ResumeSource = ResumeStream | "connect";
+
+// connectAgent() swallows the base class's not-implemented error and resolves
+// with no events, so only an agent that overrides connect() can rejoin a run.
+// AbstractAgent extends nothing, so the prototype owning connect() is the base
+// exactly when its own prototype is Object.prototype. Checking ownership instead
+// of function identity holds when the app loads its own copy of @ag-ui/client.
+const implementsConnect = (agent: AbstractAgent): boolean => {
+  if (typeof agent.connectAgent !== "function") return false;
+  if (Object.prototype.hasOwnProperty.call(agent, "connect")) return true;
+  for (
+    let proto: object | null = Object.getPrototypeOf(agent);
+    proto !== null;
+    proto = Object.getPrototypeOf(proto)
+  ) {
+    if (Object.prototype.hasOwnProperty.call(proto, "connect")) {
+      return Object.getPrototypeOf(proto) !== Object.prototype;
+    }
+  }
+  return false;
+};
 type ResumeRunConfig = {
   parentId: string | null;
   sourceId: string | null;
@@ -240,6 +262,7 @@ type CoreOptions = {
   logger: Logger;
   showThinking: boolean;
   joinStrategy?: UseAgUiRuntimeOptions["joinStrategy"];
+  joinAdjacentAssistantMessages?: boolean | undefined;
   getRunContext?: UseAgUiRuntimeOptions["getRunContext"];
   resumeTranscript?: AgUiResumeTranscript | undefined;
   autoCancelPendingToolCalls?: boolean | undefined;
@@ -377,6 +400,7 @@ export class AgUiThreadController {
   private agent: AbstractAgent;
   private logger: Logger;
   private showThinking: boolean;
+  private joinAdjacentAssistantMessages: boolean;
   private joinStrategy: UseAgUiRuntimeOptions["joinStrategy"];
   private getRunContext: UseAgUiRuntimeOptions["getRunContext"];
   private resumeTranscript: AgUiResumeTranscript;
@@ -384,6 +408,7 @@ export class AgUiThreadController {
   private onError: ((error: Error) => void) | undefined;
   private onCancel: (() => void) | undefined;
   private readonly listeners = new Set<() => void>();
+  private pendingUpdateNotification: ReturnType<typeof setTimeout> | undefined;
   private snapshot:
     | {
         isLoading: boolean;
@@ -445,6 +470,8 @@ export class AgUiThreadController {
     }
     this.logger = options.logger;
     this.showThinking = options.showThinking;
+    this.joinAdjacentAssistantMessages =
+      options.joinAdjacentAssistantMessages ?? false;
     this.joinStrategy = options.joinStrategy;
     this.getRunContext = options.getRunContext;
     this.resumeTranscript = options.resumeTranscript ?? "full";
@@ -472,8 +499,20 @@ export class AgUiThreadController {
     this.state = reduceAgUiThreadState(this.state, action);
   }
 
-  private publishUpdate(): void {
+  private publishUpdate(deferNotification = false): void {
     this.snapshot = undefined;
+    if (deferNotification) {
+      if (this.pendingUpdateNotification !== undefined) return;
+      this.pendingUpdateNotification = setTimeout(() => {
+        this.pendingUpdateNotification = undefined;
+        for (const listener of this.listeners) listener();
+      }, 0);
+      return;
+    }
+    if (this.pendingUpdateNotification !== undefined) {
+      clearTimeout(this.pendingUpdateNotification);
+      this.pendingUpdateNotification = undefined;
+    }
     for (const listener of this.listeners) listener();
   }
 
@@ -492,6 +531,8 @@ export class AgUiThreadController {
     }
     this.logger = options.logger;
     this.showThinking = options.showThinking;
+    this.joinAdjacentAssistantMessages =
+      options.joinAdjacentAssistantMessages ?? false;
     this.joinStrategy = options.joinStrategy;
     this.getRunContext = options.getRunContext;
     this.resumeTranscript = options.resumeTranscript ?? "full";
@@ -517,9 +558,22 @@ export class AgUiThreadController {
 
   detachRuntime() {
     this.runtime = undefined;
-    void this.cancel().catch((error) => {
+    try {
+      this.leave();
+    } catch (error) {
       this.logger.error("[agui] failed to cancel run during teardown", error);
+    }
+  }
+
+  /**
+   * Stops following the active run because the thread is being left (a thread
+   * switch or teardown), as opposed to cancel(), which the user asked for.
+   */
+  leave(): void {
+    this.dispatch({
+      type: "historyMutationAdvanced",
     });
+    this.abortActiveRun(true);
   }
 
   getMessages(): readonly ThreadMessage[] {
@@ -582,12 +636,11 @@ export class AgUiThreadController {
       .then(async (repo) => {
         if (repo?.unstable_resume) {
           const parentId = repo.headId ?? this.session.headId;
-          const resumeStream = this.history?.resume?.bind(this.history);
           await this.startRun(
             parentId,
             this.state.lastRunConfig,
             undefined,
-            resumeStream,
+            this.getResumeSource(),
           );
         }
       })
@@ -739,16 +792,26 @@ export class AgUiThreadController {
   // synchronously, so a callback that starts another run must keep the
   // controller it installed, and a throw from a subclass's abortRun would
   // otherwise strand the thread as running.
-  private abortActiveRun(): void {
+  private abortActiveRun(leaveRunning = false): void {
     const controller = this.abortController;
     if (!controller) return;
     try {
-      (this.activeRunAgent ?? this.agent).abortRun();
+      this.stopAgent(this.activeRunAgent ?? this.agent, leaveRunning);
     } finally {
       controller.abort();
     }
   }
 
+  private stopAgent(agent: AbstractAgent, leaveRunning: boolean): void {
+    if (leaveRunning && implementsConnect(agent)) {
+      void agent.detachActiveRun().catch((error: unknown) => {
+        this.logger.error?.("[agui] agent detachActiveRun failed", error);
+      });
+      return;
+    }
+    agent.abortRun();
+  }
+
   // Starting a run supersedes the active one, and reports whether the caller
   // still owns the thread afterwards: the local abort runs onCancel
   // synchronously, so a callback that starts its own run keeps it. abortRun is a
@@ -785,11 +848,11 @@ export class AgUiThreadController {
   async resumeInFlightRun(messages: readonly ThreadMessage[]): Promise<void> {
     const historyImport = this.waitForHistoryImport();
     if (historyImport && !(await historyImport)) return;
-    // Without a resume stream startRun would re-run the agent from scratch.
-    const resumeStream = this.history?.resume?.bind(this.history);
-    if (!resumeStream) {
+    // Without a resume source startRun would re-run the agent from scratch.
+    const resumeSource = this.getResumeSource();
+    if (!resumeSource) {
       const error = new Error(
-        "[agui] unstable_resume requires a ThreadHistoryAdapter with a resume() method; skipping resume after thread switch",
+        "[agui] unstable_resume requires a ThreadHistoryAdapter with a resume() method or an agent that implements connect(); skipping resume after thread switch",
       );
       this.logger.error?.(error.message);
       invokeRuntimeCallback("onError", this.onError, error);
@@ -801,13 +864,19 @@ export class AgUiThreadController {
         parentId,
         this.state.lastRunConfig,
         undefined,
-        resumeStream,
+        resumeSource,
       );
     } catch {
       // startRun already reported via onError; don't reject the switch.
     }
   }
 
+  private getResumeSource(): ResumeSource | undefined {
+    const resumeStream = this.history?.resume?.bind(this.history);
+    if (resumeStream) return resumeStream;
+    return implementsConnect(this.agent) ? "connect" : undefined;
+  }
+
   private assertNoPendingInterrupts(): void {
     if (!this.getPendingInterrupts()) return;
     throw new Error(
@@ -1664,7 +1733,7 @@ export class AgUiThreadController {
       this.activeRunAgent = null;
       this.setRunning(false);
       try {
-        (activeRunAgent ?? this.agent).abortRun();
+        this.stopAgent(activeRunAgent ?? this.agent, true);
       } catch (error) {
         this.logger.error?.("[agui] agent abortRun failed", error);
       } finally {
@@ -1679,10 +1748,12 @@ export class AgUiThreadController {
     parentId: string | null,
     runConfig?: RunConfig,
     resume?: ResumeDispatch,
-    resumeStream?: ResumeStream,
+    resumeSource?: ResumeSource,
     replayExistingMessages = false,
   ): Promise<void> {
     if (this.isThreadSwitching?.()) return;
+    const resumeStream =
+      typeof resumeSource === "function" ? resumeSource : undefined;
     // A default AG-UI run supersedes the active run; the hook's opt-in message
     // queue serializes sends instead. append supersedes earlier, before it links
     // its message; this covers the entry points that start a run without one.
@@ -1778,6 +1849,7 @@ export class AgUiThreadController {
     let frameOwnerIds = new Set<string>();
     let restrictedPendingToolCallIds: ReadonlySet<string> | undefined;
     let settlingRun = false;
+    let deferUpdateNotifications = false;
     let latestSnapshot: ThreadAssistantMessage | undefined;
     let snapshotToolCallIds = new Set<string>();
     const snapshotResultIds = new Set<string>();
@@ -1813,7 +1885,11 @@ export class AgUiThreadController {
         this.session.hasMessage(assistantParentId)
           ? assistantParentId
           : repositoryHeadId;
-      const created = this.insertAssistantPlaceholder(parentId, cached);
+      const created = this.insertAssistantPlaceholder(
+        parentId,
+        cached,
+        deferUpdateNotifications,
+      );
       assistantMessageId = created;
       this.markPendingAssistantHistory(created, parentId);
       return created;
@@ -2061,6 +2137,7 @@ export class AgUiThreadController {
             },
             replayContexts.get(ownerId),
             interrupt ? new Set<string>() : restrictedPendingToolCallIds,
+            deferUpdateNotifications,
           );
         }
         const { content: _content, ...controlUpdate } = update;
@@ -2078,6 +2155,7 @@ export class AgUiThreadController {
         primaryUpdate,
         replayContexts.get(messageId),
         restrictedPendingToolCallIds,
+        deferUpdateNotifications,
       );
       if (messageId === frameMessageId && resolved !== assistantMessageId) {
         assistantMessageId = resolved;
@@ -2151,7 +2229,11 @@ export class AgUiThreadController {
         const previousId = assistantMessageId;
         assistantMessageId = undefined;
         if (!adoptsReplay && this.session.tryGetMessage(previousId)) {
-          const created = this.insertAssistantPlaceholder(previousId);
+          const created = this.insertAssistantPlaceholder(
+            previousId,
+            undefined,
+            deferUpdateNotifications,
+          );
           assistantMessageId = created;
           this.markPendingAssistantHistory(created, previousId);
         }
@@ -2165,7 +2247,11 @@ export class AgUiThreadController {
       replayEnded = true;
       const placeholder = ensureAssistant(true);
       if (placeholder === serverId) return;
-      const reassigned = this.reassignAssistantId(placeholder, serverId);
+      const reassigned = this.reassignAssistantId(
+        placeholder,
+        serverId,
+        deferUpdateNotifications,
+      );
       // A collision drops the placeholder before revealing the existing
       // server message as the current head. Only messages introduced during
       // this run can replace the placeholder; regeneration must not rewrite
@@ -2238,13 +2324,19 @@ export class AgUiThreadController {
         (settlingRun && heldMessages.has(messageId));
       if (held) heldMessages.set(messageId, new Set(unresolved));
       else heldMessages.delete(messageId);
-      this.updateAssistantMessage(messageId, {
-        ...update,
-        ...(content !== undefined ? { content } : {}),
-        status: held
-          ? { type: "running" }
-          : { type: "complete", reason: "unknown" },
-      });
+      this.updateAssistantMessage(
+        messageId,
+        {
+          ...update,
+          ...(content !== undefined ? { content } : {}),
+          status: held
+            ? { type: "running" }
+            : { type: "complete", reason: "unknown" },
+        },
+        undefined,
+        undefined,
+        deferUpdateNotifications,
+      );
       const current = this.session.tryGetMessage(messageId)?.message;
       if (current?.role === "assistant")
         retainedSnapshots.set(messageId, { ...snapshot, message: current });
@@ -2262,7 +2354,13 @@ export class AgUiThreadController {
         !runStartMessageIds.has(head.id) &&
         !aggregator.isRetained(head.id)
       )
-        this.updateAssistantMessage(head.id, { status: { type: "running" } });
+        this.updateAssistantMessage(
+          head.id,
+          { status: { type: "running" } },
+          undefined,
+          undefined,
+          deferUpdateNotifications,
+        );
     };
 
     const abortController = new AbortController();
@@ -2477,6 +2575,13 @@ export class AgUiThreadController {
       if (settlingRun) persistRunAssistants(true);
       const beforeSnapshot =
         event.type === "MESSAGES_SNAPSHOT" ? this.session.getMessages() : [];
+      const previousDeferredNotifications = deferUpdateNotifications;
+      deferUpdateNotifications =
+        event.type === "TEXT_MESSAGE_CONTENT" ||
+        event.type === "REASONING_MESSAGE_CONTENT" ||
+        event.type === "TOOL_CALL_ARGS" ||
+        event.type === "ACTIVITY_DELTA" ||
+        event.type === "STATE_DELTA";
       try {
         const nextAssistantMessageId = this.handleEvent(
           aggregator,
@@ -2569,6 +2674,7 @@ export class AgUiThreadController {
         }
         if (settlingRun) persistRunAssistants(false);
       } finally {
+        deferUpdateNotifications = previousDeferredNotifications;
         settlingRun = false;
       }
     };
@@ -2638,9 +2744,13 @@ export class AgUiThreadController {
         } catch {
           // ignore
         }
-        const runAgent: RunAgentWithRunOptions =
-          runAgentInstance.runAgent.bind(runAgentInstance);
-        await runAgent(input, subscriber, { signal: abortSignal });
+        if (resumeSource === "connect") {
+          await runAgentInstance.connectAgent(input, subscriber);
+        } else {
+          const runAgent: RunAgentWithRunOptions =
+            runAgentInstance.runAgent.bind(runAgentInstance);
+          await runAgent(input, subscriber, { signal: abortSignal });
+        }
       }
     } catch (error) {
       // HttpAgent rethrows a failure it already passed to the subscriber's
@@ -2962,6 +3072,7 @@ export class AgUiThreadController {
   private insertAssistantPlaceholder(
     parentId: string | null,
     id: string = generateOptimisticId(),
+    deferNotification = false,
   ): string {
     const assistant: ThreadAssistantMessage = {
       id,
@@ -2980,11 +3091,15 @@ export class AgUiThreadController {
     };
     this.session.addOrUpdateMessage(parentId ?? this.session.headId, assistant);
     this.session.switchToBranch(id);
-    this.publishUpdate();
+    this.publishUpdate(deferNotification);
     return id;
   }
 
-  private reassignAssistantId(oldId: string, newId: string): boolean {
+  private reassignAssistantId(
+    oldId: string,
+    newId: string,
+    deferNotification = false,
+  ): boolean {
     if (oldId === newId) return true;
     const oldItem = this.session.tryGetMessage(oldId);
     if (!oldItem) return false;
@@ -3045,7 +3160,7 @@ export class AgUiThreadController {
       );
     }
 
-    this.publishUpdate();
+    this.publishUpdate(deferNotification);
     return !collidesWithExisting;
   }
 
@@ -3054,6 +3169,7 @@ export class AgUiThreadController {
     update: ChatModelRunResult,
     replay?: ReplayContext,
     pendingToolCallIds?: ReadonlySet<string>,
+    deferNotification = false,
   ): string {
     let latestStatus: MessageStatus | undefined;
     const touched = this.session.updateMessage(messageId, (message) => {
@@ -3177,10 +3293,10 @@ export class AgUiThreadController {
       latestStatus !== undefined && latestStatus.type !== "running";
     if (isSettled && isOptimisticId(messageId)) {
       const stableId = generateId();
-      this.reassignAssistantId(messageId, stableId);
+      this.reassignAssistantId(messageId, stableId, deferNotification);
       resolvedMessageId = stableId;
     } else {
-      this.publishUpdate();
+      this.publishUpdate(deferNotification);
     }
     if (this.isPersistableStatus(latestStatus)) {
       this.persistAssistantHistory(resolvedMessageId);
@@ -3546,7 +3662,7 @@ export class AgUiThreadController {
             type: "stateSnapshot",
             snapshot: result.newDocument as ReadonlyJSONValue,
           });
-          this.updateActiveAssistantState(activeAssistantId);
+          this.updateActiveAssistantState(activeAssistantId, true);
         } catch (error) {
           this.logger.error?.("[agui] failed to apply state delta", error);
         }
@@ -3658,7 +3774,10 @@ export class AgUiThreadController {
     return activeAssistantId;
   }
 
-  private updateActiveAssistantState(messageId: string | undefined): void {
+  private updateActiveAssistantState(
+    messageId: string | undefined,
+    deferNotification = false,
+  ): void {
     const current =
       messageId === undefined
         ? undefined
@@ -3680,7 +3799,7 @@ export class AgUiThreadController {
         };
       });
     }
-    this.publishUpdate();
+    this.publishUpdate(deferNotification);
   }
 
   private applyCrossRunToolResult(
@@ -3860,7 +3979,10 @@ export class AgUiThreadController {
         activeMessage?.role === "assistant" ? activeMessage : undefined;
       const normalized = fromAgUiSnapshotMessages(
         rawMessages,
-        { showThinking: this.showThinking },
+        {
+          showThinking: this.showThinking,
+          joinAdjacentAssistantMessages: this.joinAdjacentAssistantMessages,
+        },
         this.subagentParentToolCallIds,
       );
       const converted: ThreadMessage[] = [];