@assistant-ui/react-ag-ui 0.0.42 → 0.0.44

package/src/runtime/AgUiThreadRuntimeCore.ts

@@ -4,6 +4,7 @@ import { generateId, fromThreadMessageLike } from "@assistant-ui/core";
 import type {
   AddToolResultOptions,
   AppendMessage,
+  CreateAppendMessage,
   AssistantRuntime,
   ChatModelRunOptions,
   ChatModelRunResult,
@@ -174,6 +175,12 @@ export class AgUiThreadRuntimeCore {
   async append(message: AppendMessage): Promise<void> {
     const startRun = message.startRun ?? message.role === "user";
     if (startRun) this.assertNoPendingInterrupts();
+    const threadMessageId = this.appendEntry(message);
+    if (!startRun) return;
+    await this.startRun(threadMessageId, message.runConfig);
+  }
+
+  private appendEntry(message: AppendMessage): string {
     if (message.sourceId) {
       this.messages = this.messages.filter(
         (entry) => entry.id !== message.sourceId,
@@ -185,9 +192,7 @@ export class AgUiThreadRuntimeCore {
     this.messages = [...this.messages, threadMessage];
     this.notifyUpdate();
     this.recordHistoryEntry(message.parentId ?? null, threadMessage);
-
-    if (!startRun) return;
-    await this.startRun(threadMessage.id, message.runConfig);
+    return threadMessage.id;
   }
 
   async edit(message: AppendMessage): Promise<void> {
@@ -348,6 +353,97 @@ export class AgUiThreadRuntimeCore {
     await this.startRun(pending.messageId, this.lastRunConfig, resume);
   }
 
+  async steerAway(
+    message: CreateAppendMessage,
+    responses?: readonly AgUiResumeEntry[],
+  ): Promise<void> {
+    const pending = this.getPendingInterrupts();
+    if (!pending) {
+      if (responses?.length) {
+        throw new Error(
+          "[agui] steerAway: no pending interrupts on this thread",
+        );
+      }
+      await this.append(this.toAppendMessage(message));
+      return;
+    }
+
+    const resume = this.resolveSteerAwayResume(pending.interrupts, responses);
+
+    if (this.isRunningFlag) {
+      throw new Error("[agui] steerAway: a run is already in progress");
+    }
+
+    const normalized = this.toAppendMessage(message);
+    // Clear before appendEntry so the interrupted assistant is still in
+    // this.messages when its status is flipped; a parentId that points at an
+    // ancestor would otherwise truncate it away before it can be cleared.
+    this.clearPendingInterrupts(pending.messageId);
+    const threadMessageId = this.appendEntry(normalized);
+    await this.startRun(threadMessageId, normalized.runConfig, resume);
+  }
+
+  private resolveSteerAwayResume(
+    interrupts: readonly AgUiInterrupt[],
+    responses: readonly AgUiResumeEntry[] | undefined,
+  ): AgUiResumeEntry[] {
+    const openIds = interrupts.map((interrupt) => interrupt.id);
+    const known = new Set(openIds);
+    const responsesById = new Map<string, AgUiResumeEntry>();
+    for (const entry of responses ?? []) {
+      if (!entry || typeof entry.interruptId !== "string") {
+        throw new Error(
+          "[agui] steerAway: every response must have an interruptId",
+        );
+      }
+      if (entry.status !== "resolved" && entry.status !== "cancelled") {
+        throw new Error(
+          `[agui] steerAway: invalid status "${entry.status}" for interrupt ${entry.interruptId}`,
+        );
+      }
+      if (!known.has(entry.interruptId)) {
+        throw new Error(
+          `[agui] steerAway: unknown interrupt id ${entry.interruptId}`,
+        );
+      }
+      if (responsesById.has(entry.interruptId)) {
+        throw new Error(
+          `[agui] steerAway: duplicate response for interrupt ${entry.interruptId}`,
+        );
+      }
+      responsesById.set(entry.interruptId, entry);
+    }
+    return openIds.map(
+      (id) => responsesById.get(id) ?? { interruptId: id, status: "cancelled" },
+    );
+  }
+
+  private toAppendMessage(message: CreateAppendMessage): AppendMessage {
+    if (typeof message === "string") {
+      return {
+        createdAt: new Date(),
+        parentId: this.messages.at(-1)?.id ?? null,
+        sourceId: null,
+        runConfig: {},
+        role: "user",
+        content: [{ type: "text", text: message }],
+        attachments: [],
+        metadata: { custom: {} },
+      };
+    }
+    return {
+      createdAt: message.createdAt ?? new Date(),
+      parentId: message.parentId ?? this.messages.at(-1)?.id ?? null,
+      sourceId: message.sourceId ?? null,
+      role: message.role ?? "user",
+      content: message.content,
+      attachments: message.attachments ?? [],
+      metadata: message.metadata ?? { custom: {} },
+      runConfig: message.runConfig ?? {},
+      startRun: message.startRun,
+    } as AppendMessage;
+  }
+
   private clearPendingInterrupts(messageId: string): void {
     let touched = false;
     this.messages = this.messages.map((message) => {

package/src/runtime/adapter/conversions.ts

@@ -2,8 +2,14 @@
 
 import type { InputContent } from "@ag-ui/client";
 import type { ThreadMessageLike as CoreThreadMessageLike } from "@assistant-ui/core";
+import { getAutoStatus } from "@assistant-ui/core/internal";
 import { type Tool, toToolsJSONSchema } from "assistant-stream";
 import type { ReadonlyJSONObject } from "assistant-stream/utils";
+import {
+  AG_UI_METADATA_NAMESPACE,
+  type AgUiCustomMetadata,
+} from "./run-aggregator";
+import type { AgUiInterrupt } from "../types";
 
 export type { InputContent };
 
@@ -439,6 +445,28 @@ function extractAssistantToolCalls(
   return parts;
 }
 
+function validateInterrupts(value: unknown): AgUiInterrupt[] | undefined {
+  if (!Array.isArray(value)) return undefined;
+  const valid = value.filter(
+    (entry): entry is AgUiInterrupt =>
+      isObject(entry) &&
+      typeof entry.id === "string" &&
+      typeof entry.reason === "string",
+  );
+  return valid.length > 0 ? valid : undefined;
+}
+
+function readPersistedInterrupts(
+  metadata: unknown,
+): AgUiInterrupt[] | undefined {
+  if (!isObject(metadata)) return undefined;
+  const custom = metadata.custom;
+  if (!isObject(custom)) return undefined;
+  const namespaced = custom[AG_UI_METADATA_NAMESPACE];
+  if (!isObject(namespaced)) return undefined;
+  return validateInterrupts(namespaced.interrupts);
+}
+
 function toAssistantSnapshotMessage(
   rawMessage: Record<string, unknown>,
 ): CoreThreadMessageLike {
@@ -449,11 +477,23 @@ function toAssistantSnapshotMessage(
     ...toolCallParts,
   ];
   const messageName = getString(rawMessage, "name");
+  const interrupts = readPersistedInterrupts(rawMessage.metadata);
   return {
     id: getString(rawMessage, "id") ?? generateId(),
     role: "assistant",
     content: assistantContent.length > 0 ? assistantContent : "",
     ...(messageName !== undefined ? { name: messageName } : {}),
+    ...(interrupts
+      ? {
+          metadata: {
+            custom: {
+              [AG_UI_METADATA_NAMESPACE]: {
+                interrupts,
+              } satisfies AgUiCustomMetadata,
+            },
+          },
+        }
+      : {}),
   };
 }
 
@@ -605,6 +645,35 @@ export function fromAgUiMessages(
     }
   }
 
+  for (let i = 0; i < converted.length; i++) {
+    const message = converted[i]!;
+    if (message.role !== "assistant") continue;
+
+    const hasInterrupt =
+      readPersistedInterrupts(message.metadata) !== undefined;
+    const hasPendingToolCall =
+      Array.isArray(message.content) &&
+      message.content.some(
+        (part) =>
+          isObject(part) &&
+          part.type === "tool-call" &&
+          part.result === undefined,
+      );
+
+    if (hasInterrupt || hasPendingToolCall) {
+      converted[i] = {
+        ...message,
+        status: getAutoStatus(
+          false,
+          false,
+          hasInterrupt,
+          hasPendingToolCall,
+          undefined,
+        ),
+      };
+    }
+  }
+
   return converted;
 }
 

package/src/runtime/adapter/run-aggregator.ts

@@ -334,7 +334,7 @@ export class RunAggregator {
         (part) => part.kind === "tool-call" && part.toolCallId === id,
       )
     ) {
-      this.partOrder.push({ kind: "tool-call", toolCallId: id });
+      this.insertToolPart(id, parentMessageId);
     }
     const state: ToolCallState = {
       toolCallId: id,
@@ -350,6 +350,35 @@ export class RunAggregator {
     this.toolCalls.set(id, state);
   }
 
+  // The message and tool-call channels are unordered on the wire, so anchor a
+  // tool call under its parentMessageId text part instead of appending it.
+  private insertToolPart(id: string, parentMessageId?: string) {
+    const entry = { kind: "tool-call", toolCallId: id } as const;
+    if (parentMessageId) {
+      const parentIndex = this.partOrder.findIndex(
+        (part) => part.kind === "text" && part.key === parentMessageId,
+      );
+      if (parentIndex !== -1) {
+        let insertAt = parentIndex + 1;
+        while (insertAt < this.partOrder.length) {
+          const part = this.partOrder[insertAt];
+          if (
+            part?.kind === "tool-call" &&
+            this.toolCalls.get(part.toolCallId)?.parentMessageId ===
+              parentMessageId
+          ) {
+            insertAt += 1;
+            continue;
+          }
+          break;
+        }
+        this.partOrder.splice(insertAt, 0, entry);
+        return;
+      }
+    }
+    this.partOrder.push(entry);
+  }
+
   private appendToolArgs(id: string | undefined, delta: string) {
     const entry = id ? this.toolCalls.get(id) : undefined;
     if (!entry) return;

package/src/runtime/types.ts

@@ -1,4 +1,5 @@
 import type {
+  CreateAppendMessage,
   AttachmentAdapter,
   DictationAdapter,
   ExternalStoreSharedOptions,
@@ -85,6 +86,17 @@ export type AgUiResumeEntry = {
   payload?: unknown;
 };
 
+export type AgUiRuntimeExtras = {
+  interrupts: readonly AgUiInterrupt[];
+  submitInterruptResponses: (
+    responses: readonly AgUiResumeEntry[],
+  ) => Promise<void>;
+  steerAway: (
+    message: CreateAppendMessage,
+    responses?: readonly AgUiResumeEntry[],
+  ) => Promise<void>;
+};
+
 export type AgUiRunFinishedOutcome =
   | { type: "success" }
   | { type: "interrupt"; interrupts: AgUiInterrupt[] };

package/src/useAgUiRuntime.ts

@@ -21,9 +21,21 @@ import type {
   UseAgUiRuntimeOptions,
 } from "./runtime/types";
 import { AgUiThreadRuntimeCore } from "./runtime/AgUiThreadRuntimeCore";
+import { agUiExtras } from "./agUiExtras";
+
+const EMPTY_INTERRUPTS: readonly AgUiInterrupt[] = [];
 
 export type AgUiAssistantRuntime = AssistantRuntime & {
+  /**
+   * @deprecated Use the `useAgUiInterrupts()` hook instead. This method is kept
+   * for backward compatibility and will be removed in a future major release.
+   */
   unstable_getPendingInterrupts: () => readonly AgUiInterrupt[];
+  /**
+   * @deprecated Use the `useAgUiSubmitInterruptResponses()` hook instead. This
+   * method is kept for backward compatibility and will be removed in a future
+   * major release.
+   */
   unstable_submitInterruptResponses: (
     responses: readonly AgUiResumeEntry[],
   ) => Promise<void>;
@@ -127,6 +139,13 @@ export function useAgUiRuntime(
         messages: core.getMessages(),
         state: core.getState(),
         isRunning: core.isRunning() || hasExecutingTools,
+        extras: agUiExtras.provide({
+          interrupts:
+            core.getPendingInterrupts()?.interrupts ?? EMPTY_INTERRUPTS,
+          submitInterruptResponses: (responses) =>
+            core.submitInterruptResponses(responses),
+          steerAway: (message, responses) => core.steerAway(message, responses),
+        }),
         unstable_enableToolInvocations: true,
         setToolStatuses,
         onNew: (message: AppendMessage) => core.append(message),