@assistant-ui/react-ag-ui 0.0.43 → 0.0.45

package/src/runtime/adapter/conversions.ts

@@ -1,10 +1,18 @@
 "use client";
 
 import type { InputContent } from "@ag-ui/client";
-import type { ThreadMessageLike as CoreThreadMessageLike } from "@assistant-ui/core";
+import type {
+  ThreadMessageLike as CoreThreadMessageLike,
+  ToolModelContentPart,
+} 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 };
 
@@ -54,6 +62,7 @@ type ToolCallPart = {
   args?: ReadonlyJSONObject;
   result?: unknown;
   isError?: boolean;
+  modelContent?: readonly ToolModelContentPart[];
   unstable_toolMessageId?: string;
 };
 
@@ -440,6 +449,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 {
@@ -450,11 +481,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,
+            },
+          },
+        }
+      : {}),
   };
 }
 
@@ -608,19 +651,29 @@ export function fromAgUiMessages(
 
   for (let i = 0; i < converted.length; i++) {
     const message = converted[i]!;
-    if (message.role !== "assistant" || !Array.isArray(message.content)) {
-      continue;
-    }
-    const hasPendingToolCall = message.content.some(
-      (part) =>
-        isObject(part) &&
-        part.type === "tool-call" &&
-        part.result === undefined,
-    );
-    if (hasPendingToolCall) {
+    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, false, true, undefined),
+        status: getAutoStatus(
+          false,
+          false,
+          hasInterrupt,
+          hasPendingToolCall,
+          undefined,
+        ),
       };
     }
   }
@@ -666,10 +719,13 @@ function convertAssistantMessage(
   for (const { id: toolCallId, part } of toolCalls) {
     if (part.result === undefined) continue;
 
+    const modelText = extractText(part.modelContent);
     const resultContent =
-      typeof part.result === "string"
-        ? part.result
-        : JSON.stringify(part.result);
+      modelText.length > 0
+        ? modelText
+        : typeof part.result === "string"
+          ? part.result
+          : JSON.stringify(part.result);
 
     const toolMessage: AgUiMessage = {
       id: part.unstable_toolMessageId ?? `${toolCallId}:tool`,

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

@@ -6,6 +6,7 @@ import {
   type MessageTiming,
   type ThreadAssistantMessagePart,
   type ToolCallMessagePart,
+  type ToolModelContentPart,
 } from "@assistant-ui/core";
 import type { AgUiEvent, AgUiInterrupt } from "../types";
 import type { Logger } from "../logger";
@@ -28,10 +29,16 @@ type ToolCallState = {
   parentMessageId?: string;
   toolMessageId?: string;
   mcpAppResourceUri?: string;
+  mcpAppServerId?: string;
+  modelContent?: ToolModelContentPart[];
+  snapshotResultApplied: boolean;
 };
 
 const MCP_APPS_ACTIVITY_TYPE = "mcp-apps";
 
+const isPlainObject = (value: unknown): value is Record<string, unknown> =>
+  !!value && typeof value === "object" && !Array.isArray(value);
+
 export type RunAggregatorOptions = {
   showThinking: boolean;
   logger: Logger;
@@ -236,8 +243,13 @@ export class RunAggregator {
       }
       case "ACTIVITY_SNAPSHOT": {
         if (event.activityType !== MCP_APPS_ACTIVITY_TYPE) break;
-        const id = this.lastResolvedToolCallId;
-        const entry = id ? this.toolCalls.get(id) : undefined;
+        const toolCallId = event.content.toolCallId;
+        const entry =
+          typeof toolCallId === "string"
+            ? this.toolCalls.get(toolCallId)
+            : this.lastResolvedToolCallId
+              ? this.toolCalls.get(this.lastResolvedToolCallId)
+              : undefined;
         const resourceUri = event.content.resourceUri;
         if (
           entry &&
@@ -245,6 +257,35 @@ export class RunAggregator {
           isMcpAppUri(resourceUri)
         ) {
           entry.mcpAppResourceUri = resourceUri;
+          const id = event.content.serverId;
+          const hash = event.content.serverHash;
+          if (typeof id === "string" && id.length > 0) {
+            entry.mcpAppServerId = id;
+          } else if (typeof hash === "string" && hash.length > 0) {
+            entry.mcpAppServerId = hash;
+          }
+          const result = event.content.result;
+          if (isPlainObject(result)) {
+            if (
+              entry.result !== undefined &&
+              entry.modelContent === undefined
+            ) {
+              entry.modelContent = [
+                {
+                  type: "text",
+                  text:
+                    typeof entry.result === "string"
+                      ? entry.result
+                      : JSON.stringify(entry.result),
+                },
+              ];
+            }
+            entry.result = result;
+            entry.snapshotResultApplied = true;
+            if (typeof result.isError === "boolean") {
+              entry.isError = result.isError;
+            }
+          }
           this.emit();
         }
         break;
@@ -343,6 +384,7 @@ export class RunAggregator {
       parsedArgs: undefined,
       result: undefined,
       isError: undefined,
+      snapshotResultApplied: false,
     };
     if (parentMessageId) {
       state.parentMessageId = parentMessageId;
@@ -411,6 +453,7 @@ export class RunAggregator {
         parsedArgs: undefined,
         result: undefined,
         isError: undefined,
+        snapshotResultApplied: false,
       };
       this.toolCalls.set(id, entry);
     }
@@ -421,8 +464,17 @@ export class RunAggregator {
     ) {
       this.partOrder.push({ kind: "tool-call", toolCallId: id });
     }
-    entry.result = tryParseJSON(content);
-    entry.isError = isError;
+    if (entry.snapshotResultApplied) {
+      if (entry.modelContent === undefined) {
+        entry.modelContent = [{ type: "text", text: content }];
+      }
+      if (entry.isError === undefined) {
+        entry.isError = isError;
+      }
+    } else {
+      entry.result = tryParseJSON(content);
+      entry.isError = isError;
+    }
     if (toolMessageId) {
       entry.toolMessageId = toolMessageId;
     }
@@ -460,9 +512,21 @@ export class RunAggregator {
         args: (entry.parsedArgs ?? {}) as any,
         argsText: entry.argsText,
         ...(entry.result !== undefined ? { result: entry.result } : {}),
+        ...(entry.modelContent !== undefined
+          ? { modelContent: entry.modelContent }
+          : {}),
         ...(entry.isError !== undefined ? { isError: entry.isError } : {}),
         ...(entry.mcpAppResourceUri
-          ? { mcp: { app: { resourceUri: entry.mcpAppResourceUri } } }
+          ? {
+              mcp: {
+                app: {
+                  resourceUri: entry.mcpAppResourceUri,
+                  ...(entry.mcpAppServerId
+                    ? { serverId: entry.mcpAppServerId }
+                    : {}),
+                },
+              },
+            }
           : {}),
         ...(entry.parentMessageId ? { parentId: entry.parentMessageId } : {}),
         ...(entry.toolMessageId

package/src/runtime/types.ts

@@ -59,6 +59,17 @@ export type UseAgUiRuntimeOptions = ExternalStoreSharedOptions & {
   agent: AbstractAgent;
   logger?: Partial<Logger>;
   showThinking?: boolean;
+  /**
+   * When the user sends, edits, or reloads a message while client-side tool
+   * calls are still pending, automatically cancel the unresolved tool calls
+   * with an error result so the agent's tool-call accounting stays
+   * consistent. Pending AG-UI interrupts are exempt: they still reject the
+   * run and must be answered with `useAgUiSubmitInterruptResponses` or
+   * discarded with `useAgUiSteerAway`. When disabled, `useAgUiSteerAway`
+   * remains the explicit way to cancel pending tool calls.
+   * Defaults to `true`.
+   */
+  autoCancelPendingToolCalls?: boolean | undefined;
   onError?: (e: Error) => void;
   onCancel?: () => void;
   adapters?: UseAgUiRuntimeAdapters;
@@ -95,6 +106,12 @@ export type AgUiRuntimeExtras = {
     message: CreateAppendMessage,
     responses?: readonly AgUiResumeEntry[],
   ) => Promise<void>;
+  state: ReadonlyJSONValue | undefined;
+  setState: (
+    next:
+      | ReadonlyJSONValue
+      | ((prev: ReadonlyJSONValue | undefined) => ReadonlyJSONValue),
+  ) => void;
 };
 
 export type AgUiRunFinishedOutcome =

package/src/useAgUiRuntime.ts

@@ -58,6 +58,7 @@ export function useAgUiRuntime(
       agent: options.agent,
       logger,
       showThinking: options.showThinking ?? true,
+      autoCancelPendingToolCalls: options.autoCancelPendingToolCalls,
       ...(options.onError && { onError: options.onError }),
       ...(options.onCancel && { onCancel: options.onCancel }),
       ...(historyAdapter && { history: historyAdapter }),
@@ -70,6 +71,7 @@ export function useAgUiRuntime(
     agent: options.agent,
     logger,
     showThinking: options.showThinking ?? true,
+    autoCancelPendingToolCalls: options.autoCancelPendingToolCalls,
     ...(options.onError && { onError: options.onError }),
     ...(options.onCancel && { onCancel: options.onCancel }),
     ...(historyAdapter && { history: historyAdapter }),
@@ -95,6 +97,7 @@ export function useAgUiRuntime(
         ? async () => {
             await onSwitchToNewThread();
             core.applyExternalMessages([]);
+            core.resetState();
           }
         : undefined,
       onSwitchToThread: onSwitchToThread
@@ -106,6 +109,8 @@ export function useAgUiRuntime(
             core.applyExternalMessages(result.messages);
             if (result.state !== undefined) {
               core.loadExternalState(result.state);
+            } else {
+              core.resetState();
             }
             if (result.unstable_resume) {
               void core.resumeInFlightRun(result.messages);
@@ -136,7 +141,7 @@ export function useAgUiRuntime(
       return {
         ...shared,
         isLoading: core.isLoading,
-        messages: core.getMessages(),
+        messageRepository: core.getMessageRepository(),
         state: core.getState(),
         isRunning: core.isRunning() || hasExecutingTools,
         extras: agUiExtras.provide({
@@ -145,6 +150,8 @@ export function useAgUiRuntime(
           submitInterruptResponses: (responses) =>
             core.submitInterruptResponses(responses),
           steerAway: (message, responses) => core.steerAway(message, responses),
+          state: core.getState(),
+          setState: (next) => core.setState(next),
         }),
         unstable_enableToolInvocations: true,
         setToolStatuses,