@assistant-ui/core 0.3.20 → 0.3.21

package/src/runtimes/tool-invocations/EDGE_CASES.md

@@ -64,7 +64,15 @@ The tracker does **not** restart the stream and does **not** invoke
 `streamCall` a second time. Logs the divergence in non-prod. The host's
 existing `streamCall` keeps its original args view.
 
-### A.5. First resolution (`result` becomes defined)
+The tracked `argsText` advances to the new value, so a snapshot that
+re-observes the same change logs once rather than on every rebuild. This
+is safe only because the stream is already closed; A.2 must keep the
+pre-regression value, since it still drives delta calculation. It also
+means a demoted entry (F.4) carries the changed text, so a pipeline
+restart no longer promotes it and re-fires `streamCall` for a change the
+host's original args view already excluded.
+
+### A.5. First resolution (a final `result` becomes defined)
 The tracker calls `setResponse` on the active controller and closes it.
 The backend result is emitted before the args stream closes, so a stale
 args parse failure cannot replace it. `reader.response.get()` resolves.
@@ -172,7 +180,7 @@ while its `onDelete` path never aborts at all. `reset()` is unaffected
 either way: it clears `_entries` and the recorded ids before aborting,
 because it opens a new execution boundary.
 
-An id is forgotten once a live snapshot observes the call with a result,
+An id is forgotten once a live snapshot observes the call with a final result,
 which bounds the set to the open calls of a discarded turn. That is all
 it does: an id re-emitted inside the same execution boundary keeps its
 existing entry, so A.4 and A.7 govern it rather than this set, and a new
@@ -188,6 +196,11 @@ each snapshot and provider ownership is recomputed, while a rebuilt entry
 would otherwise come back clean. That is what lets F.4 drop a waiting entry
 without a special case.
 
+### A.12. An interim result (`isPreliminary`)
+A result flagged `isPreliminary` marks the entry `skipExecute` like any backend result and reaches the active controller as a preliminary response, which keeps the part open. `entry.hasResult` stays false, so each later snapshot hands the controller its current interim value, while the skip marker keeps the frontend `execute` from running and drops the echo of each value at the result-chunk handler, so `onResult` never fires for it. `reader.response.get()` resolves only when the final result lands, which is A.5 for that call.
+
+A restored or demoted (F.4) entry holding an interim result counts as unresolved, so C.3 promotes it only when its final result lands or its `argsText` changes.
+
 ## B. Tool call disappears from snapshot
 
 ### B.1. Tool call removed entirely (rollback, branch switch)
@@ -211,7 +224,7 @@ Silently kept as restored. Recursion into `content.messages` still
 happens so any nested live tool calls are processed.
 
 ### C.3. Restored entry observed in a live snapshot, signature changed
-A restored entry with a result remains historical regardless of later `argsText` or result changes. An unresolved restored entry is promoted only when a result lands or its `argsText` changes. Complete `argsText` values that parse to equivalent JSON are unchanged. Promotion deletes the restored entry and starts a new active entry via `_startActiveEntry`. `streamCall` fires once, its first and only fire for this `toolCallId`.
+A restored entry with a final result remains historical regardless of later `argsText` or result changes. An unresolved restored entry, with no result or only an interim one (A.12), is promoted only when a final result lands or its `argsText` changes. Complete `argsText` values that parse to equivalent JSON are unchanged. Promotion deletes the restored entry and starts a new active entry via `_startActiveEntry`. `streamCall` fires once, its first and only fire for this `toolCallId`.
 
 ### C.4. `isLoading` transitions `true → false` while messages are stable
 The next `setState` call sees `isLoading === false` and processes
@@ -291,7 +304,9 @@ and a call waiting on the run to settle (A.10) already holds its final
 args, so demoting it would strand it unexecuted.
 
 Starting it over re-fires `streamCall`, which the restart path already
-does for any demoted entry whose signature later changes. The rebuilt
+does for any demoted entry whose signature later changes. A change that
+already happened after completion (A.4) is not such a change: the demoted
+entry carries the changed text, so the restart does not promote it. The rebuilt
 pipeline holds no part for the call, so nothing can reach the executor
 without adding one; a restart is an execution boundary in the same sense
 `reset()` is. Repeated failures

package/src/runtimes/tool-invocations/ToolInvocationTracker.test.ts

@@ -44,6 +44,7 @@ const createAssistantMessage = (
   options?: {
     result?: ReadonlyJSONValue;
     isError?: boolean;
+    isPreliminary?: boolean;
     toolCallId?: string;
     toolName?: string;
     nestedMessages?: ThreadAssistantMessage[];
@@ -70,6 +71,9 @@ const createAssistantMessage = (
       argsText,
       ...(options?.result !== undefined && { result: options.result }),
       ...(options?.isError !== undefined && { isError: options.isError }),
+      ...(options?.isPreliminary !== undefined && {
+        isPreliminary: options.isPreliminary,
+      }),
       ...(options?.nestedMessages && { messages: options.nestedMessages }),
       ...(options?.approval && { approval: options.approval }),
     },
@@ -292,51 +296,203 @@ describe("ToolInvocationTracker", () => {
     expect(statuses).toEqual({});
   });
 
-  it("marks a fresh execution as executing when an earlier one left a human-input request behind", async () => {
-    const execute = vi
-      .fn()
-      .mockImplementationOnce((_args, { human }) =>
-        human({ request: "approve" }),
-      )
-      .mockImplementationOnce(() => new Promise(() => {}));
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        execute,
-      } satisfies Tool,
-    });
-    let statuses: Record<string, ToolExecutionStatus> = {};
-    const tracker = new ToolInvocationTracker(getTools, {
-      onResult: vi.fn(),
-      onStatusesChange: (s: ReadonlyMap<string, ToolExecutionStatus>) => {
-        statuses = Object.fromEntries(s);
-      },
-    });
-    tracker.setState(createState([], false));
-    tracker.setState(
-      createState(
-        [createAssistantMessage('{"query":"London"}', { query: "London" })],
-        false,
-      ),
-    );
-    await waitFor(() => {
-      expect(statuses["tool-1"]?.type).toBe("interrupt");
+  describe("human-input requests from streamCall", () => {
+    const trackStreamCallHuman = () => {
+      let statuses: Record<string, ToolExecutionStatus> = {};
+      const tracker = new ToolInvocationTracker(
+        () => ({
+          weatherSearch: {
+            parameters: { type: "object", properties: {} },
+            streamCall: async (_reader, { human }) => {
+              await human({ request: "approve" }).catch(() => {});
+            },
+          } satisfies Tool,
+        }),
+        {
+          onResult: vi.fn(),
+          onStatusesChange: (s: ReadonlyMap<string, ToolExecutionStatus>) => {
+            statuses = Object.fromEntries(s);
+          },
+        },
+      );
+      tracker.setState(createState([], false));
+      tracker.setState(
+        createState(
+          [createAssistantMessage('{"query":"London"}', { query: "London" })],
+          false,
+        ),
+      );
+      return { tracker, statuses: () => statuses };
+    };
+
+    it("clears the call's status once the request is resumed", async () => {
+      const { tracker, statuses } = trackStreamCallHuman();
+      await waitFor(() => {
+        expect(statuses()["tool-1"]?.type).toBe("interrupt");
+      });
+
+      expect(tracker.resume("tool-1", true)).toBe(true);
+
+      expect(statuses()).toEqual({});
     });
 
-    killPipeline(tracker);
-    tracker.setState(
-      createState(
-        [createAssistantMessage('{"query":"Paris"}', { query: "Paris" })],
-        false,
-      ),
-    );
+    it("clears the call's interrupt when the tracker aborts", async () => {
+      const { tracker, statuses } = trackStreamCallHuman();
+      await waitFor(() => {
+        expect(statuses()["tool-1"]?.type).toBe("interrupt");
+      });
 
-    await waitFor(() => {
-      expect(execute).toHaveBeenCalledTimes(2);
-      expect(statuses["tool-1"]?.type).toBe("executing");
+      await tracker.abort();
+
+      expect(statuses()).toEqual({});
     });
   });
 
+  it.each(["resume", "abort"] as const)(
+    "marks an execute still running after %s executing until it settles",
+    async (end) => {
+      let finish!: () => void;
+      const execute = vi.fn(async (_args, { human }) => {
+        await human({ request: "approve" }).catch(() => undefined);
+        await new Promise<void>((resolve) => (finish = resolve));
+        return { approved: true };
+      });
+      let statuses: Record<string, ToolExecutionStatus> = {};
+      const tracker = new ToolInvocationTracker(
+        () => ({
+          weatherSearch: {
+            parameters: { type: "object", properties: {} },
+            execute,
+          } satisfies Tool,
+        }),
+        {
+          onResult: vi.fn(),
+          onStatusesChange: (s: ReadonlyMap<string, ToolExecutionStatus>) => {
+            statuses = Object.fromEntries(s);
+          },
+        },
+      );
+      tracker.setState(createState([], false));
+      tracker.setState(
+        createState(
+          [createAssistantMessage('{"query":"London"}', { query: "London" })],
+          false,
+        ),
+      );
+      await waitFor(() => {
+        expect(statuses["tool-1"]?.type).toBe("interrupt");
+      });
+
+      if (end === "resume") tracker.resume("tool-1", true);
+      else void tracker.abort();
+      expect(statuses["tool-1"]?.type).toBe("executing");
+
+      await waitFor(() => expect(finish).toBeDefined());
+      finish();
+      await waitFor(() => expect(statuses).toEqual({}));
+    },
+  );
+
+  it.each(["resume", "abort"] as const)(
+    "marks a fresh execution as executing when an earlier one left a human-input request behind, and keeps it after %s() ends that request",
+    async (ending) => {
+      const execute = vi
+        .fn()
+        .mockImplementationOnce((_args, { human }) =>
+          human({ request: "approve" }),
+        )
+        .mockImplementationOnce(() => new Promise(() => {}));
+      const getTools = () => ({
+        weatherSearch: {
+          parameters: { type: "object", properties: {} },
+          execute,
+        } satisfies Tool,
+      });
+      let statuses: Record<string, ToolExecutionStatus> = {};
+      const tracker = new ToolInvocationTracker(getTools, {
+        onResult: vi.fn(),
+        onStatusesChange: (s: ReadonlyMap<string, ToolExecutionStatus>) => {
+          statuses = Object.fromEntries(s);
+        },
+      });
+      tracker.setState(createState([], false));
+      tracker.setState(
+        createState(
+          [createAssistantMessage('{"query":"London"}', { query: "London" })],
+          false,
+        ),
+      );
+      await waitFor(() => {
+        expect(statuses["tool-1"]?.type).toBe("interrupt");
+      });
+
+      killPipeline(tracker);
+      tracker.setState(
+        createState(
+          [createAssistantMessage('{"query":"Paris"}', { query: "Paris" })],
+          false,
+        ),
+      );
+
+      await waitFor(() => {
+        expect(execute).toHaveBeenCalledTimes(2);
+        expect(statuses["tool-1"]?.type).toBe("executing");
+      });
+
+      if (ending === "resume") tracker.resume("tool-1", true);
+      else void tracker.abort();
+      expect(statuses["tool-1"]?.type).toBe("executing");
+    },
+  );
+
+  it.each(["resume", "abort"] as const)(
+    "clears the status when %s() ends a request no execution owns after a pipeline restart",
+    async (ending) => {
+      const execute = vi.fn((_args, { human }) =>
+        human({ request: "approve" }),
+      );
+      let statuses: Record<string, ToolExecutionStatus> = {};
+      const tracker = new ToolInvocationTracker(
+        () => ({
+          weatherSearch: {
+            parameters: { type: "object", properties: {} },
+            execute,
+          } satisfies Tool,
+        }),
+        {
+          onResult: vi.fn(),
+          onStatusesChange: (s: ReadonlyMap<string, ToolExecutionStatus>) => {
+            statuses = Object.fromEntries(s);
+          },
+        },
+      );
+      tracker.setState(createState([], false));
+      tracker.setState(
+        createState(
+          [createAssistantMessage('{"query":"London"}', { query: "London" })],
+          false,
+        ),
+      );
+      await waitFor(() => {
+        expect(statuses["tool-1"]?.type).toBe("interrupt");
+      });
+
+      killPipeline(tracker);
+      tracker.setState(
+        createState(
+          [createAssistantMessage('{"query":"London"}', { query: "London" })],
+          false,
+        ),
+      );
+      expect(execute).toHaveBeenCalledTimes(1);
+
+      if (ending === "resume")
+        expect(tracker.resume("tool-1", true)).toBe(true);
+      else await tracker.abort();
+      expect(statuses).toEqual({});
+    },
+  );
+
   it("does not auto-submit a parse-error result for a non-executable tool whose divergent argsText closes", async () => {
     // Same close-gating mismatch as the executable case, but for a tool with
     // no frontend execute. Closing on the divergent complete snapshot would
@@ -913,6 +1069,38 @@ describe("ToolInvocationTracker", () => {
     await waitFor(() => expect(onResult).toHaveBeenCalledTimes(1));
   });
 
+  it("does not re-fire streamCall for a call holding a preliminary result when the pipeline restarts", async () => {
+    const streamCall = vi.fn();
+    const tracker = new ToolInvocationTracker(
+      () => ({
+        weatherSearch: {
+          parameters: { type: "object", properties: {} },
+          streamCall,
+        } satisfies Tool,
+      }),
+      { onResult: vi.fn(), onStatusesChange: () => {} },
+    );
+    const interim = () =>
+      createState([
+        createAssistantMessage(
+          '{"city":"London"}',
+          { city: "London" },
+          { result: { forecast: "interim" }, isPreliminary: true },
+        ),
+      ]);
+    tracker.setState(createState([]));
+    tracker.setState(interim());
+    await waitFor(() => {
+      expect(streamCall).toHaveBeenCalledOnce();
+    });
+
+    killPipeline(tracker);
+    tracker.setState(interim());
+    await new Promise((r) => setTimeout(r, 0));
+
+    expect(streamCall).toHaveBeenCalledOnce();
+  });
+
   it("never executes a registered tool the provider gates before its run ends", async () => {
     const execute = vi.fn(async () => ({ deleted: true }));
     const getTools = () => ({
@@ -1598,6 +1786,52 @@ describe("ToolInvocationTracker", () => {
     expect(equivalentStreamCall).not.toHaveBeenCalled();
   });
 
+  it("keeps a tool call restored with a preliminary result historical until its final result lands", async () => {
+    const streamCall = vi.fn();
+    const onResult = vi.fn();
+    const tracker = new ToolInvocationTracker(
+      () => ({
+        weatherSearch: {
+          parameters: { type: "object", properties: {} },
+          streamCall,
+        } satisfies Tool,
+      }),
+      { onResult, onStatusesChange: () => {} },
+    );
+    const interim = () =>
+      createState([
+        createAssistantMessage(
+          '{"city":"London"}',
+          { city: "London" },
+          { result: { forecast: "interim" }, isPreliminary: true },
+        ),
+      ]);
+
+    tracker.setState(interim());
+    tracker.setState(interim());
+    await new Promise((r) => setTimeout(r, 0));
+    expect(streamCall).not.toHaveBeenCalled();
+
+    tracker.setState(
+      createState([
+        createAssistantMessage(
+          '{"city":"London"}',
+          { city: "London" },
+          { result: { forecast: "final" } },
+        ),
+      ]),
+    );
+
+    await waitFor(() => {
+      expect(streamCall).toHaveBeenCalledOnce();
+    });
+    const [reader] = streamCall.mock.calls[0]!;
+    await expect(reader.response.get()).resolves.toMatchObject({
+      result: { forecast: "final" },
+    });
+    expect(onResult).not.toHaveBeenCalled();
+  });
+
   it("promotes an in-progress tool call from the initial snapshot when it changes", async () => {
     const execute = vi.fn(async () => ({ forecast: "ok" }));
     const streamCall = vi.fn();
@@ -2042,6 +2276,78 @@ describe("ToolInvocationTracker", () => {
       warnSpy.mockRestore();
     }
   });
+
+  it("keeps a preliminary backend result open until the final result", async () => {
+    const execute = vi.fn(async () => ({ forecast: "client" }));
+    const streamCall = vi.fn();
+    const onResult = vi.fn();
+    const tracker = new ToolInvocationTracker(
+      () => ({
+        weatherSearch: {
+          parameters: { type: "object", properties: {} },
+          execute,
+          streamCall,
+        } satisfies Tool,
+      }),
+      { onResult, onStatusesChange: () => {} },
+    );
+    tracker.setState(createState([]));
+    tracker.setState(
+      createState([
+        createAssistantMessage('{"city":"London"}', { city: "London" }),
+      ]),
+    );
+
+    await waitFor(() => {
+      expect(streamCall).toHaveBeenCalledOnce();
+    });
+    const [reader] = streamCall.mock.calls[0]!;
+    let resolved = false;
+    void reader.response.get().then(() => {
+      resolved = true;
+    });
+
+    tracker.setState(
+      createState([
+        createAssistantMessage(
+          '{"city":"London"}',
+          { city: "London" },
+          { result: { forecast: "interim" }, isPreliminary: true },
+        ),
+      ]),
+    );
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(resolved).toBe(false);
+
+    tracker.setState(
+      createState([
+        createAssistantMessage(
+          '{"city":"London"}',
+          { city: "London" },
+          { result: { forecast: "interim 2" }, isPreliminary: true },
+        ),
+      ]),
+    );
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(resolved).toBe(false);
+
+    tracker.setState(
+      createState([
+        createAssistantMessage(
+          '{"city":"London"}',
+          { city: "London" },
+          { result: { forecast: "final" } },
+        ),
+      ]),
+    );
+
+    await expect(reader.response.get()).resolves.toMatchObject({
+      result: { forecast: "final" },
+    });
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(execute).not.toHaveBeenCalled();
+    expect(onResult).not.toHaveBeenCalled();
+  });
 });
 
 describe("ToolInvocationTracker reset", () => {
@@ -2301,4 +2607,51 @@ describe("ToolInvocationTracker reset", () => {
 
     tracker.reset();
   });
+
+  it("warns exactly once when a settled tool call's args change is re-observed across renders", async () => {
+    // External-store / AI-SDK runtimes rebuild the messages array on update, so
+    // a settled tool part is re-diffed each snapshot. A post-completion
+    // non-equivalent args change must be recorded so the same change is diffed
+    // once; otherwise it re-warns on every re-observation while retaining both
+    // copies of the payload, which can grow unbounded on large args.
+    const getTools = () => ({
+      weatherSearch: {
+        parameters: { type: "object", properties: {} },
+      } satisfies Tool,
+    });
+    const onResult = vi.fn();
+    const onStatusesChange = () => {};
+    const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
+
+    try {
+      const tracker = new ToolInvocationTracker(getTools, {
+        onResult,
+        onStatusesChange,
+      });
+      tracker.setState(createState([], false));
+
+      // Settle the call with complete args (not running → args stream closes).
+      tracker.setState(
+        createState([createAssistantMessage('{"a":1}', { a: 1 })], false),
+      );
+
+      const afterFirstCompletion = () =>
+        warnSpy.mock.calls.filter((call) =>
+          String(call[0]).includes("changed after first completion"),
+        ).length;
+
+      // A non-equivalent args change after completion, re-observed across three
+      // renders with a fresh messages array each time (identity differs, so the
+      // fast-path skip does not fire).
+      for (let i = 0; i < 3; i++) {
+        tracker.setState(
+          createState([createAssistantMessage('{"a":2}', { a: 2 })], false),
+        );
+      }
+
+      expect(afterFirstCompletion()).toBe(1);
+    } finally {
+      warnSpy.mockRestore();
+    }
+  });
 });

package/src/runtimes/tool-invocations/ToolInvocationTracker.ts

@@ -88,6 +88,9 @@ const isEquivalentCompleteArgsText = (previous: string, next: string) => {
   return isJSONValueEqual(previousValue, nextValue);
 };
 
+const hasFinalResult = (part: ToolCallMessagePart) =>
+  part.result !== undefined && part.isPreliminary !== true;
+
 const getToolExecutionId = (value: object): symbol | undefined =>
   (value as Record<PropertyKey, unknown>)[TOOL_EXECUTION_ID] as
     | symbol
@@ -335,13 +338,14 @@ export class ToolInvocationTracker {
    */
   public abort(options?: { discardPending?: boolean }): Promise<void> {
     try {
-      this._humanInput.forEach(({ reject }) => {
+      this._humanInput.forEach(({ executionId, reject }, toolCallId) => {
         try {
           reject(new Error("Tool execution aborted"));
         } catch {
           // host rejection handler threw — already in the abort path,
           // swallow so we continue cleaning up.
         }
+        this._endHumanRequest(toolCallId, executionId);
       });
       this._humanInput.clear();
 
@@ -370,6 +374,19 @@ export class ToolInvocationTracker {
     }
   }
 
+  // A request from streamCall has no execution whose end would clear the
+  // status, so the call is only marked executing while one runs, and a request
+  // left behind by an earlier execution leaves a newer execution's status alone.
+  private _endHumanRequest(toolCallId: string, executionId: symbol) {
+    if (this._executing.has(executionId)) {
+      this._setStatus(toolCallId, { type: "executing" });
+      return;
+    }
+    const owner = this._entries.get(toolCallId)?.executionId;
+    if (owner === undefined || owner === executionId)
+      this._deleteStatus(toolCallId);
+  }
+
   /**
    * Resolve a pending human-input request for the given tool call. Returns
    * `true` if a pending request was resumed, `false` if the tracker has no
@@ -381,7 +398,7 @@ export class ToolInvocationTracker {
       const handlers = this._humanInput.get(toolCallId);
       if (!handlers) return false;
       this._humanInput.delete(toolCallId);
-      this._setStatus(toolCallId, { type: "executing" });
+      this._endHumanRequest(toolCallId, handlers.executionId);
       handlers.resolve(payload);
       return true;
     } catch (err) {
@@ -727,6 +744,7 @@ export class ToolInvocationTracker {
               },
             );
           }
+          entry.argsText = content.argsText;
           shouldWriteArgsText = false;
         }
       } else if (!content.argsText.startsWith(entry.argsText)) {
@@ -810,7 +828,7 @@ export class ToolInvocationTracker {
           this._entries.set(content.toolCallId, {
             toolName: content.toolName,
             argsText: content.argsText,
-            hasResult: content.result !== undefined,
+            hasResult: hasFinalResult(content),
           });
         }
         continue;
@@ -821,11 +839,11 @@ export class ToolInvocationTracker {
 
       // A discarded id is remembered only until the call is answered, which
       // bounds the set to the open calls of a discarded turn.
-      if (content.result !== undefined)
+      if (hasFinalResult(content))
         this._discardedToolCallIds.delete(content.toolCallId);
 
       if (entry && !entry.controller) {
-        // A restored entry with a result remains historical.
+        // A restored entry with a final result remains historical.
         if (entry.hasResult) continue;
         const argsChanged =
           content.argsText !== entry.argsText &&
@@ -834,7 +852,7 @@ export class ToolInvocationTracker {
             isArgsTextComplete(content.argsText) &&
             isEquivalentCompleteArgsText(entry.argsText, content.argsText)
           );
-        if (!argsChanged && content.result === undefined) continue;
+        if (!argsChanged && !hasFinalResult(content)) continue;
         this._entries.delete(content.toolCallId);
         entry = undefined;
       }
@@ -855,6 +873,8 @@ export class ToolInvocationTracker {
         );
       }
 
+      if (content.result !== undefined) entry.skipExecute = true;
+
       if (content.approval !== undefined) entry.skipExecute = true;
 
       this._processArgsText(entry, content);
@@ -865,19 +885,22 @@ export class ToolInvocationTracker {
         // controller. Narrow once instead of asserting at every use.
         const { controller: activeController } = entry;
         if (!activeController) continue;
-        entry.hasResult = true;
         entry.argsComplete = true;
         activeController.setResponse(
           new ToolResponse({
             result: content.result as ReadonlyJSONValue,
             artifact: content.artifact as ReadonlyJSONValue | undefined,
             isError: content.isError,
+            isPreliminary: content.isPreliminary,
             ...(content.modelContent !== undefined
               ? { modelContent: content.modelContent }
               : {}),
           }),
         );
-        activeController.close();
+        if (hasFinalResult(content)) {
+          entry.hasResult = true;
+          activeController.close();
+        }
       }
     }
   }