@assistant-ui/core 0.3.19 → 0.3.21

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();
+        }
       }
     }
   }

package/src/store/clients/external-thread.test.ts

@@ -0,0 +1,145 @@
+import { createTapRoot, useResource } from "@assistant-ui/tap";
+import { describe, expect, it, vi } from "vitest";
+import type { Unstable_RecordToolInteractionOptions } from "../../runtime/interfaces/thread-runtime-core";
+import type { ExternalThreadMessage } from "./external-thread";
+import { ExternalThread } from "./external-thread";
+
+const message: ExternalThreadMessage = {
+  id: "message-1",
+  role: "assistant",
+  createdAt: new Date(0),
+  content: [
+    {
+      type: "tool-call",
+      toolCallId: "call-1",
+      toolName: "weather",
+      args: {},
+      argsText: "{}",
+    },
+  ],
+  status: { type: "complete", reason: "stop" },
+  metadata: {
+    unstable_state: null,
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+};
+
+const createPart = (
+  unstable_onRecordToolInteraction?: (
+    options: Unstable_RecordToolInteractionOptions,
+  ) => void,
+  onResumeToolCall?: (options: {
+    toolCallId: string;
+    payload: unknown;
+  }) => void,
+) => {
+  const root = createTapRoot(function ExternalThreadRoot() {
+    return useResource(
+      ExternalThread({
+        messages: [message],
+        ...(unstable_onRecordToolInteraction
+          ? { unstable_onRecordToolInteraction }
+          : {}),
+        ...(onResumeToolCall ? { onResumeToolCall } : {}),
+      }),
+    );
+  });
+  return {
+    part: root.getValue().message({ index: 0 }).part({ index: 0 }),
+    unmount: () => root.unmount(),
+  };
+};
+
+describe("ExternalThread interaction recording", () => {
+  it("threads records from parts to the callback", async () => {
+    const unstable_onRecordToolInteraction = vi.fn();
+    const { part, unmount } = createPart(unstable_onRecordToolInteraction);
+
+    try {
+      await part.unstable_recordInteraction!({
+        type: "action",
+        payload: { action: "toggle" },
+      });
+
+      expect(unstable_onRecordToolInteraction).toHaveBeenCalledExactlyOnceWith({
+        messageId: "message-1",
+        toolCallId: "call-1",
+        interaction: {
+          type: "action",
+          payload: { action: "toggle" },
+          occurredAt: expect.any(Number),
+        },
+      });
+    } finally {
+      unmount();
+    }
+  });
+
+  it("records an accepted human response after resuming", async () => {
+    const onResumeToolCall = vi.fn();
+    const unstable_onRecordToolInteraction = vi.fn();
+    const { part, unmount } = createPart(
+      unstable_onRecordToolInteraction,
+      onResumeToolCall,
+    );
+    const payload = { answer: "yes" };
+
+    try {
+      part.resumeToolCall(payload);
+
+      expect(onResumeToolCall).toHaveBeenCalledExactlyOnceWith({
+        toolCallId: "call-1",
+        payload,
+      });
+      await vi.waitFor(() =>
+        expect(
+          unstable_onRecordToolInteraction,
+        ).toHaveBeenCalledExactlyOnceWith({
+          messageId: "message-1",
+          toolCallId: "call-1",
+          interaction: {
+            type: "human-response",
+            payload,
+            occurredAt: expect.any(Number),
+          },
+        }),
+      );
+      expect(onResumeToolCall.mock.invocationCallOrder[0]).toBeLessThan(
+        unstable_onRecordToolInteraction.mock.invocationCallOrder[0]!,
+      );
+    } finally {
+      unmount();
+    }
+  });
+
+  it("keeps an accepted resume successful when recording fails", async () => {
+    const onResumeToolCall = vi.fn();
+    const record = vi.fn(() => Promise.reject(new Error("recording failed")));
+    const { part, unmount } = createPart(record, onResumeToolCall);
+
+    try {
+      expect(() => part.resumeToolCall({ answer: "yes" })).not.toThrow();
+      await vi.waitFor(() => expect(record).toHaveBeenCalledOnce());
+      expect(onResumeToolCall).toHaveBeenCalledOnce();
+    } finally {
+      unmount();
+    }
+  });
+
+  it("rejects when no interaction callback is configured", async () => {
+    const { part, unmount } = createPart();
+
+    try {
+      await expect(
+        part.unstable_recordInteraction!({ type: "action", payload: {} }),
+      ).rejects.toThrow(
+        "Runtime does not support recording tool interactions.",
+      );
+    } finally {
+      unmount();
+    }
+  });
+});