@assistant-ui/core 0.3.14 → 0.3.16

package/src/runtimes/local/local-thread-runtime-core.test.ts

@@ -7,7 +7,11 @@ import type {
 } from "../../runtime/utils/chat-model-adapter";
 import type { AppendMessage } from "../../types/message";
 import type { LocalRuntimeOptionsBase } from "./local-runtime-options";
-import type { ExportedMessageRepositoryItem } from "../../runtime/utils/message-repository";
+import {
+  ExportedMessageRepository,
+  type ExportedMessageRepositoryItem,
+} from "../../runtime/utils/message-repository";
+import type { ThreadMessageLike } from "../../runtime/utils/thread-message-like";
 import type { ThreadSuggestion } from "../../runtime/interfaces/thread-runtime-core";
 import { isMessageNotSentError } from "../../types/error";
 
@@ -673,6 +677,325 @@ describe("LocalThreadRuntimeCore tool approvals", () => {
 });
 
 describe("LocalThreadRuntimeCore cancellation", () => {
+  it("settles a superseded tool-call result without aborting its replacement", async () => {
+    let resolveFirst!: (result: ChatModelRunResult) => void;
+    let resolveSecond!: (result: ChatModelRunResult) => void;
+    const firstResult = new Promise<ChatModelRunResult>((resolve) => {
+      resolveFirst = resolve;
+    });
+    const secondResult = new Promise<ChatModelRunResult>((resolve) => {
+      resolveSecond = resolve;
+    });
+    const runs: ChatModelRunOptions[] = [];
+    const thread = createThread({
+      run(options) {
+        runs.push(options);
+        return runs.length === 1 ? firstResult : secondResult;
+      },
+    });
+
+    const firstAppend = thread.append(userMessage("first"));
+    await flush();
+    const supersededMessageId = thread.messages.at(-1)!.id;
+    const secondAppend = thread.append({
+      ...userMessage("second"),
+      parentId: supersededMessageId,
+    });
+    await flush();
+
+    expect(runs).toHaveLength(2);
+    expect(runs[0]?.abortSignal.aborted).toBe(true);
+    expect(runs[1]?.abortSignal.aborted).toBe(false);
+
+    resolveFirst(toolCallResult("lookup_weather"));
+    await firstAppend;
+
+    expect(runs).toHaveLength(2);
+    expect(runs[1]?.abortSignal.aborted).toBe(false);
+    expect(
+      thread.messages.find((item) => item.id === supersededMessageId)?.status,
+    ).toEqual({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+
+    resolveSecond({ content: [{ type: "text", text: "replacement" }] });
+    await secondAppend;
+  });
+
+  it("cancels a superseded approval pause", async () => {
+    let resolveFirst!: (result: ChatModelRunResult) => void;
+    let resolveSecond!: (result: ChatModelRunResult) => void;
+    const firstResult = new Promise<ChatModelRunResult>((resolve) => {
+      resolveFirst = resolve;
+    });
+    const secondResult = new Promise<ChatModelRunResult>((resolve) => {
+      resolveSecond = resolve;
+    });
+    const runs: ChatModelRunOptions[] = [];
+    const thread = createThread({
+      run(options) {
+        runs.push(options);
+        return runs.length === 1 ? firstResult : secondResult;
+      },
+    });
+
+    const firstAppend = thread.append(userMessage("first"));
+    await flush();
+    const supersededMessageId = thread.messages.at(-1)!.id;
+    const secondAppend = thread.append({
+      ...userMessage("second"),
+      parentId: supersededMessageId,
+    });
+    await flush();
+
+    resolveFirst(toolCallResult("deploy", { id: "approval-1" }));
+    await firstAppend;
+
+    expect(
+      thread.messages.find((item) => item.id === supersededMessageId)?.status,
+    ).toEqual({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+
+    resolveSecond({ content: [{ type: "text", text: "replacement" }] });
+    await secondAppend;
+  });
+
+  it("ignores a stale result after addToolResult replaces the same message", async () => {
+    let releaseFirst!: () => void;
+    let resolveSecond!: (result: ChatModelRunResult) => void;
+    const firstTeardown = new Promise<void>((resolve) => {
+      releaseFirst = resolve;
+    });
+    const secondResult = new Promise<ChatModelRunResult>((resolve) => {
+      resolveSecond = resolve;
+    });
+    const runs: ChatModelRunOptions[] = [];
+    const thread = createThread({
+      run(options) {
+        runs.push(options);
+        if (runs.length === 1) {
+          return (async function* () {
+            yield toolCallResult("send_email");
+            await firstTeardown;
+            yield toolCallResult("stale_tool");
+          })();
+        }
+        return secondResult;
+      },
+    });
+
+    const firstAppend = thread.append(userMessage("send an email"));
+    await flush();
+    const messageId = thread.messages.at(-1)!.id;
+
+    thread.addToolResult({
+      messageId,
+      toolCallId: "call-send_email",
+      toolName: "send_email",
+      result: { approved: true },
+      isError: false,
+    });
+    await flush();
+
+    expect(runs).toHaveLength(2);
+    expect(runs[1]?.abortSignal.aborted).toBe(false);
+
+    releaseFirst();
+    await firstAppend;
+
+    const toolCall = thread.messages
+      .find((item) => item.id === messageId)
+      ?.content.find(
+        (part) =>
+          part.type === "tool-call" && part.toolCallId === "call-send_email",
+      );
+    expect(toolCall?.result).toEqual({ approved: true });
+    expect(runs[1]?.abortSignal.aborted).toBe(false);
+
+    resolveSecond({ content: [{ type: "text", text: "replacement" }] });
+    await vi.waitFor(() => {
+      expect(
+        thread.messages.find((item) => item.id === messageId)?.status.type,
+      ).toBe("complete");
+    });
+  });
+
+  it("ignores stale chunks after a partial tool result updates the active message", async () => {
+    let releaseTeardown!: () => void;
+    const teardown = new Promise<void>((resolve) => {
+      releaseTeardown = resolve;
+    });
+    const runs: ChatModelRunOptions[] = [];
+    const thread = createThread({
+      run(options) {
+        runs.push(options);
+        return (async function* () {
+          yield {
+            content: [
+              toolCallPart("send_email"),
+              toolCallPart("deploy", { id: "approval-1" }),
+            ],
+            status: {
+              type: "requires-action",
+              reason: "tool-calls",
+            },
+          } satisfies ChatModelRunResult;
+          await teardown;
+          yield { content: [{ type: "text", text: "stale" }] };
+        })();
+      },
+    });
+
+    const appendPromise = thread.append(userMessage("send and deploy"));
+    await flush();
+    const messageId = thread.messages.at(-1)!.id;
+
+    thread.addToolResult({
+      messageId,
+      toolCallId: "call-send_email",
+      toolName: "send_email",
+      result: { approved: true },
+      isError: false,
+    });
+    await flush();
+
+    expect(runs).toHaveLength(1);
+
+    releaseTeardown();
+    await appendPromise;
+
+    const message = thread.messages.find((item) => item.id === messageId);
+    const sendEmail = message?.content.find(
+      (part) =>
+        part.type === "tool-call" && part.toolCallId === "call-send_email",
+    );
+    const deploy = message?.content.find(
+      (part) => part.type === "tool-call" && part.toolCallId === "call-deploy",
+    );
+    expect(sendEmail?.result).toEqual({ approved: true });
+    expect(deploy?.approval).toEqual({ id: "approval-1" });
+    expect(message?.status.type).toBe("requires-action");
+  });
+
+  it("ignores a superseded result after its message is removed", async () => {
+    let resolveFirst!: (result: ChatModelRunResult) => void;
+    const firstResult = new Promise<ChatModelRunResult>((resolve) => {
+      resolveFirst = resolve;
+    });
+    let runCount = 0;
+    const thread = createThread({
+      run() {
+        runCount++;
+        if (runCount === 1) return firstResult;
+        return Promise.resolve({
+          content: [{ type: "text", text: "replacement" }],
+        });
+      },
+    });
+
+    const firstAppend = thread.append(userMessage("first"));
+    await flush();
+    await thread.append(userMessage("second"));
+    thread.reset();
+
+    resolveFirst(toolCallResult("lookup_weather"));
+
+    await expect(firstAppend).resolves.toBeUndefined();
+    expect(thread.messages).toHaveLength(0);
+  });
+
+  it("does not continue after cancelRun when a delayed tool call resolves", async () => {
+    let resolveFirst!: (result: ChatModelRunResult) => void;
+    const firstResult = new Promise<ChatModelRunResult>((resolve) => {
+      resolveFirst = resolve;
+    });
+    const runs: ChatModelRunOptions[] = [];
+    const thread = createThread({
+      run(options) {
+        runs.push(options);
+        if (runs.length === 1) return firstResult;
+        return Promise.resolve({
+          content: [{ type: "text", text: "unexpected continuation" }],
+          status: { type: "complete", reason: "stop" },
+        });
+      },
+    });
+
+    const appendPromise = thread.append(userMessage("first"));
+    await flush();
+
+    thread.cancelRun();
+    resolveFirst(toolCallResult("lookup_weather"));
+    await appendPromise;
+
+    expect(runs).toHaveLength(1);
+    expect(thread.messages.at(-1)?.status).toEqual({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+  });
+
+  it("preserves a terminal adapter status when cancelled during teardown", async () => {
+    let releaseTeardown!: () => void;
+    const teardown = new Promise<void>((resolve) => {
+      releaseTeardown = resolve;
+    });
+    const thread = createThread({
+      async *run() {
+        yield {
+          content: [{ type: "text", text: "done" }],
+          status: { type: "complete", reason: "stop" },
+        };
+        await teardown;
+      },
+    });
+
+    const appendPromise = thread.append(userMessage("first"));
+    await flush();
+
+    expect(thread.messages.at(-1)?.status).toEqual({
+      type: "complete",
+      reason: "stop",
+    });
+
+    thread.cancelRun();
+    releaseTeardown();
+    await appendPromise;
+
+    expect(thread.messages.at(-1)?.status).toEqual({
+      type: "complete",
+      reason: "stop",
+    });
+  });
+
+  it("preserves an approval pause when the adapter yields after cancellation", async () => {
+    let releaseTeardown!: () => void;
+    const teardown = new Promise<void>((resolve) => {
+      releaseTeardown = resolve;
+    });
+    const thread = createThread({
+      async *run() {
+        yield toolCallResult("deploy", { id: "approval-1" });
+        await teardown;
+        yield { content: [{ type: "text", text: "late" }] };
+      },
+    });
+
+    const appendPromise = thread.append(userMessage("deploy"));
+    await flush();
+
+    expect(thread.messages.at(-1)?.status?.type).toBe("requires-action");
+
+    thread.cancelRun();
+    releaseTeardown();
+    await appendPromise;
+
+    expect(thread.messages.at(-1)?.status?.type).toBe("requires-action");
+  });
+
   it("keeps a replacement run cancellable after the previous run settles", async () => {
     let releaseFirst!: () => void;
     let releaseSecond!: () => void;
@@ -1683,3 +2006,170 @@ describe("LocalRuntimeCore composer.canCancel", () => {
     thread.cancelRun();
   });
 });
+
+describe("LocalThreadRuntimeCore imported approvals", () => {
+  const pausedOnApproval = (
+    approval: { id: string } | undefined,
+    extra: Record<string, unknown> = {},
+  ): ThreadMessageLike[] => [
+    { role: "user", content: [{ type: "text", text: "send an email" }] },
+    {
+      role: "assistant",
+      content: [
+        {
+          type: "tool-call",
+          toolCallId: "call-send_email",
+          toolName: "send_email",
+          args: {},
+          argsText: "{}",
+          ...(approval !== undefined ? { approval } : {}),
+          ...extra,
+        },
+      ],
+    },
+  ];
+
+  const createImportedThread = (
+    initialMessages: ThreadMessageLike[],
+    history?: LocalRuntimeOptionsBase["adapters"]["history"],
+  ) => {
+    const runs: ChatModelRunOptions[] = [];
+    const core = new LocalRuntimeCore(
+      {
+        adapters: {
+          chatModel: {
+            async run(options) {
+              runs.push(options);
+              return { content: [{ type: "text", text: "done" }] };
+            },
+          },
+          ...(history !== undefined && { history }),
+        },
+        unstable_humanToolNames: ["send_email"],
+      },
+      initialMessages,
+    );
+    return { thread: core.threads.getMainThreadRuntimeCore(), runs };
+  };
+
+  it("resumes the run after approving an imported pending approval", async () => {
+    const { thread, runs } = createImportedThread(
+      pausedOnApproval({ id: "a1" }),
+    );
+
+    expect(thread.messages.at(-1)?.status).toMatchObject({
+      type: "requires-action",
+      reason: "tool-calls",
+    });
+
+    thread.respondToToolApproval({ approvalId: "a1", approved: true });
+    await flush();
+
+    expect(runs).toHaveLength(1);
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+  });
+
+  it("keeps the interrupt reason for an imported interrupted tool call", () => {
+    const { thread } = createImportedThread(
+      pausedOnApproval(undefined, {
+        interrupt: { type: "human", payload: { question: "which one?" } },
+      }),
+    );
+
+    expect(thread.messages.at(-1)?.status).toMatchObject({
+      type: "requires-action",
+      reason: "interrupt",
+    });
+  });
+
+  it("resumes the run after approving an approval restored by the history adapter", async () => {
+    const { thread, runs } = createImportedThread([], {
+      async load() {
+        return ExportedMessageRepository.fromArray(
+          pausedOnApproval({ id: "a1" }),
+        );
+      },
+      async append() {},
+      async update() {},
+    });
+
+    thread.__internal_load();
+    await flush();
+
+    expect(thread.messages.at(-1)?.status).toMatchObject({
+      type: "requires-action",
+      reason: "tool-calls",
+    });
+
+    thread.respondToToolApproval({ approvalId: "a1", approved: true });
+    await flush();
+
+    expect(runs).toHaveLength(1);
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+  });
+
+  it("normalizes a serialized approval pause that lost its auto-status marker", async () => {
+    const { thread, runs } = createImportedThread([]);
+    thread.import(
+      JSON.parse(
+        JSON.stringify(
+          ExportedMessageRepository.fromArray(pausedOnApproval({ id: "a1" })),
+        ),
+      ),
+    );
+
+    expect(thread.messages.at(-1)?.status).toMatchObject({
+      type: "requires-action",
+      reason: "tool-calls",
+    });
+
+    thread.respondToToolApproval({ approvalId: "a1", approved: true });
+    await flush();
+
+    expect(runs).toHaveLength(1);
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+  });
+
+  it("completes an imported interrupt that has nothing left to act on", () => {
+    const { thread } = createImportedThread([
+      { role: "user", content: [{ type: "text", text: "send an email" }] },
+      {
+        role: "assistant",
+        status: { type: "requires-action", reason: "interrupt" },
+        content: [
+          {
+            type: "tool-call",
+            toolCallId: "call-send_email",
+            toolName: "send_email",
+            args: {},
+            argsText: "{}",
+            result: { sent: true },
+          },
+        ],
+      } as ThreadMessageLike,
+    ]);
+
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+  });
+
+  it("resumes the run after adding a result to an imported resultless tool call", async () => {
+    const { thread, runs } = createImportedThread(pausedOnApproval(undefined));
+
+    expect(thread.messages.at(-1)?.status).toMatchObject({
+      type: "requires-action",
+      reason: "tool-calls",
+    });
+
+    thread.addToolResult({
+      messageId: thread.messages.at(-1)!.id,
+      toolCallId: "call-send_email",
+      toolName: "send_email",
+      result: { sent: true },
+      isError: false,
+    });
+    await flush();
+
+    expect(runs).toHaveLength(1);
+    expect(thread.messages.at(-1)?.status?.type).toBe("complete");
+  });
+});

package/src/runtimes/local/local-thread-runtime-core.ts

@@ -5,6 +5,8 @@ import type {
   ChatModelRunResult,
 } from "../../runtime/utils/chat-model-adapter";
 import { shouldContinue } from "./should-continue";
+import { getAutoStatus } from "../../runtime/utils/auto-status";
+import type { ExportedMessageRepository } from "../../runtime/utils/message-repository";
 import type { LocalRuntimeOptionsBase } from "./local-runtime-options";
 import { consumeSuggestionResult } from "../../adapters/suggestion";
 import type {
@@ -21,7 +23,7 @@ import type {
   AppendMessage,
   ThreadAssistantMessage,
 } from "../../types/message";
-import type { RunConfig } from "../../types/message";
+import type { RunConfig, ThreadMessage } from "../../types/message";
 import { MessageNotSentError, toAssistantError } from "../../types/error";
 import type { ModelContextProvider } from "../../model-context/types";
 import {
@@ -32,7 +34,7 @@ import type { QueuePlacement } from "../../runtime/queue/external-thread-queue-a
 import {
   EMPTY_QUEUE_ITEMS,
   type QueueItemState,
-} from "../../store/scopes/queue-item";
+} from "../../runtime/queue/queue-item";
 import {
   captureThreadRuntimeGeneration,
   invalidateThreadRuntime,
@@ -49,6 +51,45 @@ class AbortError extends Error {
   }
 }
 
+// `shouldContinue` resumes only on `tool-calls` and `resumeToolCall` throws, so
+// an imported `interrupt` with no interrupt payload would be stranded here.
+// Provenance cannot gate it: a repository that went through JSON has no marker.
+// The replacement stays content-derived, so a message with nothing resultless
+// left to act on lands on `complete` rather than an unactionable pause.
+const withLocalPauseReason = (message: ThreadMessage): ThreadMessage => {
+  if (
+    message.role !== "assistant" ||
+    message.status.type !== "requires-action" ||
+    message.status.reason !== "interrupt" ||
+    message.content.some(
+      (c) =>
+        c.type === "tool-call" && c.result === undefined && c.interrupt != null,
+    )
+  )
+    return message;
+  return {
+    ...message,
+    status: getAutoStatus(
+      false,
+      false,
+      false,
+      message.content.some(
+        (c) => c.type === "tool-call" && c.result === undefined,
+      ),
+    ),
+  };
+};
+
+const withLocalPauseReasons = (
+  data: ExportedMessageRepository,
+): ExportedMessageRepository => ({
+  ...data,
+  messages: data.messages.map((item) => ({
+    ...item,
+    message: withLocalPauseReason(item.message),
+  })),
+});
+
 export class LocalThreadRuntimeCore
   extends BaseThreadRuntimeCore
   implements ThreadRuntimeCore
@@ -272,7 +313,7 @@ export class LocalThreadRuntimeCore
     this._loadPromise = promise
       .then((repo) => {
         if (!repo) return;
-        this.repository.import(repo);
+        this.repository.import(withLocalPauseReasons(repo));
         if (repo.messages.length > 0) {
           this.ensureInitialized();
         }
@@ -434,6 +475,10 @@ export class LocalThreadRuntimeCore
     throw new Error("Runtime does not support exporting external states.");
   }
 
+  public override import(data: ExportedMessageRepository) {
+    super.import(withLocalPauseReasons(data));
+  }
+
   public importExternalState(): void {
     throw new Error("Runtime does not support importing external states.");
   }
@@ -501,9 +546,11 @@ export class LocalThreadRuntimeCore
           message,
           runConfig,
           alreadyPersisted,
+          run,
           runCallback,
         );
         runCallback = undefined;
+        if (this._activeRun !== run) break;
       } while (shouldContinue(message, this._options.unstable_humanToolNames));
     } finally {
       this._notifyEventSubscribers("runEnd", {});
@@ -548,6 +595,7 @@ export class LocalThreadRuntimeCore
     message: ThreadAssistantMessage,
     runConfig: RunConfig | undefined,
     alreadyPersisted: boolean,
+    run: { cancelled: boolean },
     runCallback?: ChatModelAdapter["run"],
   ) {
     const messages = parentId ? this.repository.getMessages(parentId) : [];
@@ -562,7 +610,23 @@ export class LocalThreadRuntimeCore
     const initialData = message.metadata?.unstable_data;
     const initialSteps = message.metadata?.steps;
     const initialCustom = message.metadata?.custom;
+    let hasStoredMessage = true;
+    try {
+      this.repository.getMessage(message.id);
+    } catch {
+      hasStoredMessage = false;
+    }
+    // Other writers replace the stored message object, so identity distinguishes this run from a newer owner.
+    const ownsMessage = () => {
+      if (!hasStoredMessage) return this._activeRun === run;
+      try {
+        return this.repository.getMessage(message.id).message === message;
+      } catch {
+        return false;
+      }
+    };
     const updateMessage = (m: Partial<ChatModelRunResult>) => {
+      if (!ownsMessage()) return;
       const newSteps = m.metadata?.steps;
       const steps = newSteps
         ? [...(initialSteps ?? []), ...newSteps]
@@ -609,6 +673,7 @@ export class LocalThreadRuntimeCore
           : undefined),
       };
       this.repository.addOrUpdateMessage(parentId, message);
+      hasStoredMessage = true;
       this._notifySubscribers();
     };
 
@@ -647,6 +712,12 @@ export class LocalThreadRuntimeCore
         this.adapters.chatModel.run.bind(this.adapters.chatModel);
 
       const abortSignal = abortController.signal;
+      const shouldCancelMessage = () =>
+        abortSignal.aborted &&
+        (message.status.type === "running" ||
+          (message.status.type === "requires-action" &&
+            (this._activeRun !== run ||
+              shouldContinue(message, this._options.unstable_humanToolNames))));
       const threadId = this._getThreadId?.();
       const promiseOrGenerator = runCallback({
         messages,
@@ -665,9 +736,11 @@ export class LocalThreadRuntimeCore
       if (Symbol.asyncIterator in promiseOrGenerator) {
         for await (const r of promiseOrGenerator) {
           if (abortSignal.aborted) {
-            updateMessage({
-              status: { type: "incomplete", reason: "cancelled" },
-            });
+            if (shouldCancelMessage()) {
+              updateMessage({
+                status: { type: "incomplete", reason: "cancelled" },
+              });
+            }
             break;
           }
 
@@ -677,11 +750,13 @@ export class LocalThreadRuntimeCore
         updateMessage(await promiseOrGenerator);
       }
 
-      if (message.status.type === "running") {
+      if (shouldCancelMessage()) {
+        updateMessage({
+          status: { type: "incomplete", reason: "cancelled" },
+        });
+      } else if (message.status.type === "running") {
         updateMessage({
-          status: abortSignal.aborted
-            ? { type: "incomplete", reason: "cancelled" }
-            : { type: "complete", reason: "unknown" },
+          status: { type: "complete", reason: "unknown" },
         });
       }
     } catch (e) {
@@ -724,7 +799,7 @@ export class LocalThreadRuntimeCore
 
       // Pauses are written only for adapters that can rewrite the entry later;
       // an append-only adapter would strand a half-finished run in history.
-      if (isTerminal || (isPausing && history?.update)) {
+      if (ownsMessage() && (isTerminal || (isPausing && history?.update))) {
         const write =
           alreadyPersisted && history?.update
             ? history.update.bind(history)