@assistant-ui/react-ag-ui 0.0.70 → 0.1.0

package/src/runtime/agUiThreadState.test.ts

@@ -0,0 +1,205 @@
+import { describe, expect, it } from "vitest";
+import {
+  createAgUiThreadState,
+  reduceAgUiThreadState,
+} from "./agUiThreadState";
+
+describe("reduceAgUiThreadState", () => {
+  it("tracks run and load phases without changing state for repeated phases", () => {
+    const initial = createAgUiThreadState();
+    expect(reduceAgUiThreadState(initial, { type: "runFinished" })).toBe(
+      initial,
+    );
+    expect(reduceAgUiThreadState(initial, { type: "loadFinished" })).toBe(
+      initial,
+    );
+
+    const running = reduceAgUiThreadState(initial, { type: "runStarted" });
+    const requested = reduceAgUiThreadState(running, { type: "loadRequested" });
+    const loading = reduceAgUiThreadState(requested, { type: "loadStarted" });
+    expect(loading).toMatchObject({
+      isRunning: true,
+      isLoading: true,
+      loadRequested: true,
+    });
+    expect(reduceAgUiThreadState(running, { type: "runStarted" })).toBe(
+      running,
+    );
+    expect(reduceAgUiThreadState(requested, { type: "loadRequested" })).toBe(
+      requested,
+    );
+    expect(reduceAgUiThreadState(loading, { type: "loadStarted" })).toBe(
+      loading,
+    );
+    expect(initial).toMatchObject({
+      isRunning: false,
+      isLoading: false,
+      loadRequested: false,
+    });
+
+    const stopped = reduceAgUiThreadState(loading, { type: "runFinished" });
+    const loaded = reduceAgUiThreadState(stopped, { type: "loadFinished" });
+    expect(loaded.isRunning).toBe(false);
+    expect(loaded.isLoading).toBe(false);
+    expect(loaded.loadRequested).toBe(true);
+  });
+
+  it("advances history mutation generation across thread resets", () => {
+    const initial = createAgUiThreadState();
+    const first = reduceAgUiThreadState(initial, {
+      type: "historyMutationAdvanced",
+    });
+    const reset = reduceAgUiThreadState(first, { type: "resetThread" });
+    const second = reduceAgUiThreadState(reset, {
+      type: "historyMutationAdvanced",
+    });
+    expect(initial.historyMutationGeneration).toBe(0);
+    expect(reset.historyMutationGeneration).toBe(1);
+    expect(second.historyMutationGeneration).toBe(2);
+  });
+
+  it("replaces the state snapshot", () => {
+    const initial = createAgUiThreadState();
+    const snapshot = { count: 1 };
+    const first = reduceAgUiThreadState(initial, {
+      type: "stateSnapshot",
+      snapshot,
+    });
+    expect(first.stateSnapshot).toBe(snapshot);
+    expect(
+      reduceAgUiThreadState(first, { type: "stateSnapshot", snapshot }),
+    ).toBe(first);
+  });
+
+  it("sets and resets state snapshots", () => {
+    const initial = createAgUiThreadState();
+    const first = reduceAgUiThreadState(initial, {
+      type: "stateSnapshot",
+      snapshot: { count: 1 },
+    });
+    const second = reduceAgUiThreadState(first, {
+      type: "stateSnapshot",
+      snapshot: { count: 2 },
+    });
+    expect(second.stateSnapshot).toEqual({ count: 2 });
+    expect(first.stateSnapshot).toEqual({ count: 1 });
+    expect(
+      reduceAgUiThreadState(second, {
+        type: "stateSnapshot",
+        snapshot: second.stateSnapshot!,
+      }),
+    ).toBe(second);
+    const reset = reduceAgUiThreadState(second, { type: "resetState" });
+    expect(reset.stateSnapshot).toBeUndefined();
+    expect(reduceAgUiThreadState(reset, { type: "resetState" })).toBe(reset);
+  });
+
+  it("records run config and manages deferred continuations by owner", () => {
+    const initial = createAgUiThreadState();
+    const runConfig = {};
+    const configured = reduceAgUiThreadState(initial, {
+      type: "runConfigRecorded",
+      runConfig,
+    });
+    expect(configured.lastRunConfig).toBe(runConfig);
+    expect(
+      reduceAgUiThreadState(configured, {
+        type: "runConfigRecorded",
+        runConfig,
+      }),
+    ).toBe(configured);
+
+    const owner = new AbortController();
+    const withResume = reduceAgUiThreadState(configured, {
+      type: "resumeDeferred",
+      owner,
+      messageId: "assistant",
+    });
+    expect(withResume.pendingResume).toEqual({ owner, messageId: "assistant" });
+    expect(
+      reduceAgUiThreadState(withResume, {
+        type: "resumeDeferred",
+        owner,
+        messageId: "assistant",
+      }),
+    ).toBe(withResume);
+    const withOwner = reduceAgUiThreadState(withResume, {
+      type: "a2uiResumeDeferred",
+      owner,
+    });
+    const action = { value: "submit" };
+    const withAction = reduceAgUiThreadState(withOwner, {
+      type: "a2uiActionDeferred",
+      action,
+    });
+    expect(withAction.pendingA2uiResumeOwner).toBe(owner);
+    expect(withAction.pendingA2uiAction).toBe(action);
+    expect(
+      reduceAgUiThreadState(withOwner, { type: "a2uiResumeDeferred", owner }),
+    ).toBe(withOwner);
+    expect(
+      reduceAgUiThreadState(withAction, { type: "a2uiActionDeferred", action }),
+    ).toBe(withAction);
+
+    const actionCleared = reduceAgUiThreadState(withAction, {
+      type: "a2uiActionCleared",
+    });
+    expect(actionCleared.pendingA2uiAction).toBeUndefined();
+    expect(actionCleared.pendingA2uiResumeOwner).toBe(owner);
+    const ownerCleared = reduceAgUiThreadState(actionCleared, {
+      type: "a2uiResumeCleared",
+    });
+    expect(ownerCleared.pendingA2uiResumeOwner).toBeNull();
+    expect(
+      reduceAgUiThreadState(ownerCleared, { type: "a2uiResumeCleared" }),
+    ).toBe(ownerCleared);
+    expect(
+      reduceAgUiThreadState(ownerCleared, { type: "a2uiActionCleared" }),
+    ).toBe(ownerCleared);
+    const resumeCleared = reduceAgUiThreadState(ownerCleared, {
+      type: "resumeCleared",
+    });
+    expect(resumeCleared.pendingResume).toBeNull();
+    expect(
+      reduceAgUiThreadState(resumeCleared, { type: "resumeCleared" }),
+    ).toBe(resumeCleared);
+  });
+
+  it("clears A2UI continuations and resets thread state without changing run flags", () => {
+    const owner = new AbortController();
+    let state = reduceAgUiThreadState(createAgUiThreadState(), {
+      type: "runStarted",
+    });
+    state = reduceAgUiThreadState(state, { type: "loadStarted" });
+    state = reduceAgUiThreadState(state, {
+      type: "resumeDeferred",
+      owner,
+      messageId: "assistant",
+    });
+    state = reduceAgUiThreadState(state, {
+      type: "a2uiResumeDeferred",
+      owner,
+    });
+    state = reduceAgUiThreadState(state, {
+      type: "a2uiActionDeferred",
+      action: { value: "submit" },
+    });
+    const cleared = reduceAgUiThreadState(state, {
+      type: "a2uiContinuationsCleared",
+    });
+    expect(cleared.pendingA2uiResumeOwner).toBeNull();
+    expect(cleared.pendingA2uiAction).toBeUndefined();
+    expect(cleared.pendingResume).toEqual({ owner, messageId: "assistant" });
+    expect(
+      reduceAgUiThreadState(cleared, { type: "a2uiContinuationsCleared" }),
+    ).toBe(cleared);
+
+    const reset = reduceAgUiThreadState(state, { type: "resetThread" });
+    expect(reset).toMatchObject({ isRunning: true, isLoading: true });
+    expect(reset.stateSnapshot).toBeUndefined();
+    expect(reset.pendingResume).toBeNull();
+    expect(reset.pendingA2uiResumeOwner).toBeNull();
+    expect(reset.pendingA2uiAction).toBeUndefined();
+    expect(reduceAgUiThreadState(reset, { type: "resetThread" })).toBe(reset);
+  });
+});

package/src/runtime/agUiThreadState.ts

@@ -0,0 +1,130 @@
+import type { RunConfig } from "@assistant-ui/core";
+import type { ReadonlyJSONValue } from "assistant-stream/utils";
+
+export type AgUiThreadState = {
+  isRunning: boolean;
+  isLoading: boolean;
+  loadRequested: boolean;
+  historyMutationGeneration: number;
+  stateSnapshot: ReadonlyJSONValue | undefined;
+  lastRunConfig: RunConfig | undefined;
+  pendingResume: { owner: AbortController; messageId: string } | null;
+  pendingA2uiResumeOwner: AbortController | null;
+  pendingA2uiAction: Record<string, unknown> | undefined;
+};
+
+export type AgUiThreadAction =
+  | { type: "runStarted" }
+  | { type: "runFinished" }
+  | { type: "loadRequested" }
+  | { type: "loadStarted" }
+  | { type: "loadFinished" }
+  | { type: "historyMutationAdvanced" }
+  | { type: "stateSnapshot"; snapshot: ReadonlyJSONValue }
+  | { type: "resetState" }
+  | { type: "runConfigRecorded"; runConfig: RunConfig }
+  | { type: "resumeDeferred"; owner: AbortController; messageId: string }
+  | { type: "resumeCleared" }
+  | { type: "a2uiResumeDeferred"; owner: AbortController }
+  | { type: "a2uiResumeCleared" }
+  | { type: "a2uiActionDeferred"; action: Record<string, unknown> }
+  | { type: "a2uiActionCleared" }
+  | { type: "a2uiContinuationsCleared" }
+  | { type: "resetThread" };
+
+export const createAgUiThreadState = (): AgUiThreadState => ({
+  isRunning: false,
+  isLoading: false,
+  loadRequested: false,
+  historyMutationGeneration: 0,
+  stateSnapshot: undefined,
+  lastRunConfig: undefined,
+  pendingResume: null,
+  pendingA2uiResumeOwner: null,
+  pendingA2uiAction: undefined,
+});
+
+export const reduceAgUiThreadState = (
+  state: AgUiThreadState,
+  action: AgUiThreadAction,
+): AgUiThreadState => {
+  switch (action.type) {
+    case "runStarted":
+      return state.isRunning ? state : { ...state, isRunning: true };
+    case "runFinished":
+      return !state.isRunning ? state : { ...state, isRunning: false };
+    case "loadRequested":
+      return state.loadRequested ? state : { ...state, loadRequested: true };
+    case "loadStarted":
+      return state.isLoading ? state : { ...state, isLoading: true };
+    case "loadFinished":
+      return !state.isLoading ? state : { ...state, isLoading: false };
+    case "historyMutationAdvanced":
+      return {
+        ...state,
+        historyMutationGeneration: state.historyMutationGeneration + 1,
+      };
+    case "stateSnapshot":
+      return state.stateSnapshot === action.snapshot
+        ? state
+        : { ...state, stateSnapshot: action.snapshot };
+    case "resetState":
+      return state.stateSnapshot === undefined
+        ? state
+        : { ...state, stateSnapshot: undefined };
+    case "runConfigRecorded":
+      return state.lastRunConfig === action.runConfig
+        ? state
+        : { ...state, lastRunConfig: action.runConfig };
+    case "resumeDeferred":
+      return state.pendingResume?.owner === action.owner &&
+        state.pendingResume.messageId === action.messageId
+        ? state
+        : {
+            ...state,
+            pendingResume: { owner: action.owner, messageId: action.messageId },
+          };
+    case "resumeCleared":
+      return state.pendingResume === null
+        ? state
+        : { ...state, pendingResume: null };
+    case "a2uiResumeDeferred":
+      return state.pendingA2uiResumeOwner === action.owner
+        ? state
+        : { ...state, pendingA2uiResumeOwner: action.owner };
+    case "a2uiResumeCleared":
+      return state.pendingA2uiResumeOwner === null
+        ? state
+        : { ...state, pendingA2uiResumeOwner: null };
+    case "a2uiActionDeferred":
+      return state.pendingA2uiAction === action.action
+        ? state
+        : { ...state, pendingA2uiAction: action.action };
+    case "a2uiActionCleared":
+      return state.pendingA2uiAction === undefined
+        ? state
+        : { ...state, pendingA2uiAction: undefined };
+    case "a2uiContinuationsCleared":
+      return state.pendingA2uiResumeOwner === null &&
+        state.pendingA2uiAction === undefined
+        ? state
+        : {
+            ...state,
+            pendingA2uiResumeOwner: null,
+            pendingA2uiAction: undefined,
+          };
+    case "resetThread":
+      return state.stateSnapshot === undefined &&
+        state.pendingResume === null &&
+        state.pendingA2uiResumeOwner === null &&
+        state.pendingA2uiAction === undefined
+        ? state
+        : {
+            ...state,
+            stateSnapshot: undefined,
+            pendingResume: null,
+            pendingA2uiResumeOwner: null,
+            pendingA2uiAction: undefined,
+          };
+  }
+};

package/src/runtime/event-parser.ts

@@ -7,6 +7,7 @@ import type {
 import type { Logger } from "./logger";
 import { parseMcpToolCallResult } from "./mcp-tool-result";
 import { withRawResponseSchema } from "./interrupt-internals";
+import type { ContentPart } from "@ag-ui/client";
 
 export type ParseAgUiEventOptions = {
   logger?: Logger;
@@ -79,7 +80,19 @@ const parseRunFinishedOutcome = (
     );
     return undefined;
   }
-  if (raw.type === "success") return { type: "success" };
+  if (raw.type === "success") {
+    return withOptional(
+      { type: "success" as const },
+      {
+        pendingToolCallIds:
+          Array.isArray(raw.pendingToolCallIds) &&
+          raw.pendingToolCallIds.every(isString)
+            ? raw.pendingToolCallIds
+            : undefined,
+      },
+    );
+  }
+  if (raw.type === "cancelled") return { type: "cancelled" };
   if (raw.type === "interrupt") {
     if (!Array.isArray(raw.interrupts)) {
       logRejectedEvent(
@@ -210,74 +223,50 @@ export const parseAgUiEvent = (
           subagentRunId: getString("subagentRunId"),
         },
       );
-    case "TEXT_MESSAGE_CHUNK": {
-      const delta = getString("delta") ?? "";
-      return withOptional(
-        { type: "TEXT_MESSAGE_CHUNK" as const, delta },
-        {
-          messageId: getString("messageId"),
-          subagentRunId: getString("subagentRunId"),
-        },
-      );
-    }
-    case "THINKING_START":
-      return withOptional(
-        { type: "THINKING_START" as const },
-        { title: getString("title"), timestamp },
-      );
-    case "THINKING_TEXT_MESSAGE_START":
-      return withOptional(
-        { type: "THINKING_TEXT_MESSAGE_START" as const },
-        { timestamp },
-      );
-    case "THINKING_TEXT_MESSAGE_CONTENT": {
-      const delta = getString("delta") ?? "";
-      return { type: "THINKING_TEXT_MESSAGE_CONTENT", delta };
-    }
-    case "THINKING_TEXT_MESSAGE_END":
-      return withOptional(
-        { type: "THINKING_TEXT_MESSAGE_END" as const },
-        { timestamp },
-      );
-    case "THINKING_END":
-      return withOptional({ type: "THINKING_END" as const }, { timestamp });
-    case "REASONING_START":
+    case "REASONING_START": {
+      const messageId = getString("messageId");
+      if (!messageId) return reject("missing messageId");
       return withOptional(
-        { type: "REASONING_START" as const },
+        { type: "REASONING_START" as const, messageId },
         {
-          messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
           timestamp,
         },
       );
-    case "REASONING_MESSAGE_START":
+    }
+    case "REASONING_MESSAGE_START": {
+      const messageId = getString("messageId");
+      if (!messageId) return reject("missing messageId");
       return withOptional(
-        { type: "REASONING_MESSAGE_START" as const },
+        { type: "REASONING_MESSAGE_START" as const, messageId },
         {
-          messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
           timestamp,
         },
       );
+    }
     case "REASONING_MESSAGE_CONTENT": {
+      const messageId = getString("messageId");
+      if (!messageId) return reject("missing messageId");
       const delta = getString("delta") ?? "";
       return withOptional(
-        { type: "REASONING_MESSAGE_CONTENT" as const, delta },
+        { type: "REASONING_MESSAGE_CONTENT" as const, messageId, delta },
         {
-          messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
         },
       );
     }
-    case "REASONING_MESSAGE_END":
+    case "REASONING_MESSAGE_END": {
+      const messageId = getString("messageId");
+      if (!messageId) return reject("missing messageId");
       return withOptional(
-        { type: "REASONING_MESSAGE_END" as const },
+        { type: "REASONING_MESSAGE_END" as const, messageId },
         {
-          messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
           timestamp,
         },
       );
+    }
     case "REASONING_ENCRYPTED_VALUE": {
       const entityId = getString("entityId");
       const encryptedValue = getString("encryptedValue");
@@ -299,15 +288,17 @@ export const parseAgUiEvent = (
         ...(subagentRunId !== undefined ? { subagentRunId } : {}),
       };
     }
-    case "REASONING_END":
+    case "REASONING_END": {
+      const messageId = getString("messageId");
+      if (!messageId) return reject("missing messageId");
       return withOptional(
-        { type: "REASONING_END" as const },
+        { type: "REASONING_END" as const, messageId },
         {
-          messageId: getString("messageId"),
           subagentRunId: getString("subagentRunId"),
           timestamp,
         },
       );
+    }
     case "TOOL_CALL_START": {
       const toolCallId = getString("toolCallId");
       if (!toolCallId) return reject("missing toolCallId");
@@ -337,21 +328,12 @@ export const parseAgUiEvent = (
         { subagentRunId: getString("subagentRunId") },
       );
     }
-    case "TOOL_CALL_CHUNK":
-      return withOptional(
-        { type: "TOOL_CALL_CHUNK" as const },
-        {
-          toolCallId: getString("toolCallId"),
-          toolCallName: getString("toolCallName"),
-          parentMessageId: getString("parentMessageId"),
-          delta: getString("delta"),
-          subagentRunId: getString("subagentRunId"),
-        },
-      );
     case "TOOL_CALL_RESULT": {
       const toolCallId = getString("toolCallId");
       if (!toolCallId) return reject("missing toolCallId");
-      const content = getString("content") ?? "";
+      const content = Array.isArray(payload.content)
+        ? (payload.content as ContentPart[])
+        : (getString("content") ?? "");
       return withOptional(
         {
           type: "TOOL_CALL_RESULT" as const,

package/src/runtime/mcp-tool-result.ts

@@ -1,5 +1,6 @@
 import { isRecord } from "@assistant-ui/core/internal";
 import { isMcpAppUri } from "@assistant-ui/core";
+import { contentToText, type ContentPart } from "@ag-ui/client";
 
 export type McpToolCallResult = Record<string, unknown> & {
   content: unknown[];
@@ -8,9 +9,16 @@ export type McpToolCallResult = Record<string, unknown> & {
   isError?: boolean;
 };
 
+export const toolContentText = (content: string | ContentPart[]): string =>
+  contentToText(
+    Array.isArray(content)
+      ? (content.filter(isRecord) as ContentPart[])
+      : content,
+  );
+
 export const parseMcpToolCallResult = (
   source: Record<string, unknown>,
-  modelContent: string,
+  content: string | ContentPart[],
 ): McpToolCallResult | undefined => {
   const structuredContent = isRecord(source.structuredContent)
     ? source.structuredContent
@@ -18,8 +26,14 @@ export const parseMcpToolCallResult = (
   const meta = isRecord(source._meta) ? source._meta : undefined;
   if (!structuredContent && !meta) return undefined;
 
+  const modelContent = toolContentText(content);
+
   return {
-    content: modelContent ? [{ type: "text", text: modelContent }] : [],
+    content: Array.isArray(content)
+      ? content
+      : modelContent
+        ? [{ type: "text", text: modelContent }]
+        : [],
     ...(structuredContent ? { structuredContent } : {}),
     ...(meta ? { _meta: meta } : {}),
     ...(typeof source.isError === "boolean" ? { isError: source.isError } : {}),

package/src/runtime/types.ts

@@ -11,7 +11,13 @@ import type {
   ThreadHistoryAdapter,
   ThreadMessage,
 } from "@assistant-ui/core";
-import type { AbstractAgent, Context, Message } from "@ag-ui/client";
+import type { JoinStrategy } from "@assistant-ui/core/react";
+import type {
+  AbstractAgent,
+  ContentPart,
+  Context,
+  Message,
+} from "@ag-ui/client";
 import type { Logger } from "./logger";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import type { McpToolCallResult } from "./mcp-tool-result";
@@ -32,6 +38,7 @@ export type AgUiCustomMetadata = {
    * represent (a developer record rides as a system message). */
   role?: "developer";
   interrupts?: AgUiInterrupt[];
+  pendingToolCallIds?: string[];
   opaqueReasoning?: AgUiOpaqueReasoning[];
 };
 
@@ -82,6 +89,7 @@ export type UseAgUiRuntimeOptions = ExternalStoreSharedOptions & {
   agent: AbstractAgent;
   logger?: Partial<Logger>;
   showThinking?: boolean;
+  joinStrategy?: JoinStrategy;
   /**
    * Return context entries for each run. `messages` is the exact AG-UI
    * transcript sent in the run input, after transcript selection and message
@@ -187,8 +195,9 @@ export type AgUiRuntimeExtras = {
 };
 
 export type AgUiRunFinishedOutcome =
-  | { type: "success" }
-  | { type: "interrupt"; interrupts: AgUiInterrupt[] };
+  | { type: "success"; pendingToolCallIds?: string[] }
+  | { type: "interrupt"; interrupts: AgUiInterrupt[] }
+  | { type: "cancelled" };
 
 export type AgUiSubagentFinishedOutcome =
   | { type: "success" }
@@ -211,38 +220,27 @@ export type AgUiEvent =
       subagentRunId?: string;
     }
   | { type: "TEXT_MESSAGE_END"; messageId?: string; subagentRunId?: string }
-  | {
-      type: "TEXT_MESSAGE_CHUNK";
-      messageId?: string;
-      delta: string;
-      subagentRunId?: string;
-    }
-  | { type: "THINKING_START"; title?: string; timestamp?: number }
-  | { type: "THINKING_TEXT_MESSAGE_START"; timestamp?: number }
-  | { type: "THINKING_TEXT_MESSAGE_CONTENT"; delta: string }
-  | { type: "THINKING_TEXT_MESSAGE_END"; timestamp?: number }
-  | { type: "THINKING_END"; timestamp?: number }
   | {
       type: "REASONING_START";
-      messageId?: string;
+      messageId: string;
       subagentRunId?: string;
       timestamp?: number;
     }
   | {
       type: "REASONING_MESSAGE_START";
-      messageId?: string;
+      messageId: string;
       subagentRunId?: string;
       timestamp?: number;
     }
   | {
       type: "REASONING_MESSAGE_CONTENT";
-      messageId?: string;
+      messageId: string;
       delta: string;
       subagentRunId?: string;
     }
   | {
       type: "REASONING_MESSAGE_END";
-      messageId?: string;
+      messageId: string;
       subagentRunId?: string;
       timestamp?: number;
     }
@@ -255,7 +253,7 @@ export type AgUiEvent =
     }
   | {
       type: "REASONING_END";
-      messageId?: string;
+      messageId: string;
       subagentRunId?: string;
       timestamp?: number;
     }
@@ -273,19 +271,11 @@ export type AgUiEvent =
       subagentRunId?: string;
     }
   | { type: "TOOL_CALL_END"; toolCallId: string; subagentRunId?: string }
-  | {
-      type: "TOOL_CALL_CHUNK";
-      toolCallId?: string;
-      toolCallName?: string;
-      parentMessageId?: string;
-      delta?: string;
-      subagentRunId?: string;
-    }
   | {
       type: "TOOL_CALL_RESULT";
       messageId?: string;
       toolCallId: string;
-      content: string;
+      content: string | ContentPart[];
       role?: "tool";
       mcpResult?: McpToolCallResult;
       subagentRunId?: string;