@assistant-ui/core 0.2.0 → 0.2.3

package/dist/react/runtimes/useToolInvocations.js

@@ -27,25 +27,53 @@ const isEquivalentCompleteArgsText = (previous, next) => {
     return isJSONValueEqual(previousValue, nextValue);
 };
 export function useToolInvocations({ state, getTools, onResult, setToolStatuses, }) {
-    const lastToolStates = useRef({});
+    /**
+     * Single source of truth for per-tool-call lifecycle. Keyed by *logical*
+     * toolCallId (the id the host knows). Restored entries have no controller;
+     * active entries carry their stream id and rewrite/execution bookkeeping.
+     */
+    const entriesRef = useRef(new Map());
+    /**
+     * Reverse alias map populated only when a rewrite assigns a synthetic stream
+     * id to an entry. Identity mappings are implicit via the fallback in
+     * `getLogicalToolCallId`.
+     */
+    const streamToLogicalRef = useRef(new Map());
+    /**
+     * Stream ids whose `result` chunks must be dropped before reaching `onResult`.
+     * Populated when:
+     *   - an argsText rewrite supersedes a stream (the old stream's result, if
+     *     any, is no longer authoritative)
+     *   - `reset()` is called while a pre-resolved tool call has a never-settling
+     *     Promise pending in the executor — the eventual cancellation chunk
+     *     would otherwise be forwarded to a host that has already moved on.
+     */
+    const abandonedStreamIdsRef = useRef(new Set());
+    /**
+     * Stream ids whose `execute` should be short-circuited in the tool wrapper.
+     * Tracked by physical stream id (not logical id) so cleanup is keyed off
+     * the same id the wrapper sees in its context.
+     */
+    const skipExecuteStreamIdsRef = useRef(new Set());
     const humanInputRef = useRef(new Map());
+    /**
+     * In-flight `execute` invocations keyed by physical stream id. Lives outside
+     * `entriesRef` so `reset()` can drop tool-call state without orphaning the
+     * cleanup the cancellation `onExecutionEnd` still needs.
+     */
+    const executingRef = useRef(new Map());
     const acRef = useRef(new AbortController());
     const executingCountRef = useRef(0);
-    const startedExecutionToolCallIdsRef = useRef(new Set());
     const settledResolversRef = useRef([]);
-    const toolCallIdAliasesRef = useRef(new Map());
-    const ignoredResultToolCallIdsRef = useRef(new Set());
     const rewriteCounterRef = useRef(0);
-    const getLogicalToolCallId = (toolCallId) => {
-        return toolCallIdAliasesRef.current.get(toolCallId) ?? toolCallId;
-    };
-    const shouldIgnoreAndCleanupResult = (toolCallId) => {
-        if (!ignoredResultToolCallIdsRef.current.has(toolCallId))
-            return false;
-        ignoredResultToolCallIdsRef.current.delete(toolCallId);
-        toolCallIdAliasesRef.current.delete(toolCallId);
-        return true;
-    };
+    /**
+     * `true` until the first snapshot has been processed; `reset()` flips it
+     * back to `true`. Snapshots observed while this is `true` are treated as
+     * historical: their tool calls are recorded in `entriesRef` as restored
+     * but no streamCall/execute fires. The next snapshot is processed as live.
+     */
+    const pendingRestoreRef = useRef(true);
+    const getLogicalToolCallId = (streamId) => streamToLogicalRef.current.get(streamId) ?? streamId;
     const getWrappedTools = () => {
         const tools = getTools();
         if (!tools)
@@ -53,13 +81,27 @@ export function useToolInvocations({ state, getTools, onResult, setToolStatuses,
         return Object.fromEntries(Object.entries(tools).map(([name, tool]) => {
             const execute = tool.execute;
             const streamCall = tool.streamCall;
+            const toModelOutput = tool.toModelOutput;
             const wrappedTool = {
                 ...tool,
                 ...(execute !== undefined && {
-                    execute: (...[args, context]) => execute(args, {
-                        ...context,
-                        toolCallId: getLogicalToolCallId(context.toolCallId),
-                    }),
+                    execute: (...[args, context]) => {
+                        if (skipExecuteStreamIdsRef.current.has(context.toolCallId)) {
+                            // Pre-resolved on first live observation: never invoke the
+                            // host's execute fn. Returning a never-settling Promise keeps
+                            // the executor's pending entry alive but enqueues nothing.
+                            // The membership in skipExecuteStreamIdsRef must outlive the
+                            // wrapper call so `reset()`'s seeding loop (which reads this
+                            // Set to identify pre-resolved entries needing cancellation
+                            // suppression) sees the entry. Growth is bounded by the
+                            // number of pre-resolved tool calls observed in the session.
+                            return new Promise(() => { });
+                        }
+                        return execute(args, {
+                            ...context,
+                            toolCallId: getLogicalToolCallId(context.toolCallId),
+                        });
+                    },
                 }),
                 ...(streamCall !== undefined && {
                     streamCall: (...[reader, context]) => streamCall(reader, {
@@ -67,16 +109,27 @@ export function useToolInvocations({ state, getTools, onResult, setToolStatuses,
                         toolCallId: getLogicalToolCallId(context.toolCallId),
                     }),
                 }),
+                ...(toModelOutput !== undefined && {
+                    toModelOutput: (options) => toModelOutput({
+                        ...options,
+                        toolCallId: getLogicalToolCallId(options.toolCallId),
+                    }),
+                }),
             };
             return [name, wrappedTool];
         }));
     };
+    const resolveAllSettledResolvers = () => {
+        const resolvers = settledResolversRef.current;
+        settledResolversRef.current = [];
+        // biome-ignore lint/suspicious/useIterableCallbackReturn: forEach callback intentionally has no return
+        resolvers.forEach((resolve) => resolve());
+    };
     const [controller] = useState(() => {
         const [stream, controller] = createAssistantStreamController();
         const transform = unstable_toolResultStream(getWrappedTools, () => acRef.current?.signal ?? new AbortController().signal, (toolCallId, payload) => {
             const logicalToolCallId = getLogicalToolCallId(toolCallId);
             return new Promise((resolve, reject) => {
-                // Reject previous human input request if it exists
                 const previous = humanInputRef.current.get(logicalToolCallId);
                 if (previous) {
                     previous.reject(new Error("Human input request was superseded by a new request"));
@@ -91,46 +144,38 @@ export function useToolInvocations({ state, getTools, onResult, setToolStatuses,
                 }));
             });
         }, {
-            onExecutionStart: (toolCallId) => {
-                if (ignoredResultToolCallIdsRef.current.has(toolCallId)) {
+            onExecutionStart: (streamId) => {
+                if (skipExecuteStreamIdsRef.current.has(streamId))
                     return;
-                }
-                startedExecutionToolCallIdsRef.current.add(toolCallId);
-                const logicalToolCallId = getLogicalToolCallId(toolCallId);
+                const logicalToolCallId = getLogicalToolCallId(streamId);
+                const abandoned = abandonedStreamIdsRef.current.has(streamId);
+                executingRef.current.set(streamId, {
+                    logicalToolCallId,
+                    abandoned,
+                });
                 executingCountRef.current++;
-                setToolStatuses((prev) => ({
-                    ...prev,
-                    [logicalToolCallId]: { type: "executing" },
-                }));
-            },
-            onExecutionEnd: (toolCallId) => {
-                const wasStarted = startedExecutionToolCallIdsRef.current.delete(toolCallId);
-                if (ignoredResultToolCallIdsRef.current.has(toolCallId)) {
-                    if (wasStarted) {
-                        executingCountRef.current--;
-                        if (executingCountRef.current === 0) {
-                            // biome-ignore lint/suspicious/useIterableCallbackReturn: forEach callback intentionally has no return
-                            settledResolversRef.current.forEach((resolve) => resolve());
-                            settledResolversRef.current = [];
-                        }
-                    }
-                    return;
+                if (!abandoned) {
+                    setToolStatuses((prev) => ({
+                        ...prev,
+                        [logicalToolCallId]: { type: "executing" },
+                    }));
                 }
-                if (!wasStarted) {
+            },
+            onExecutionEnd: (streamId) => {
+                const info = executingRef.current.get(streamId);
+                if (!info)
                     return;
-                }
-                const logicalToolCallId = getLogicalToolCallId(toolCallId);
+                executingRef.current.delete(streamId);
                 executingCountRef.current--;
-                setToolStatuses((prev) => {
-                    const next = { ...prev };
-                    delete next[logicalToolCallId];
-                    return next;
-                });
-                // Resolve any waiting abort promises when all tools have settled
+                if (!info.abandoned) {
+                    setToolStatuses((prev) => {
+                        const next = { ...prev };
+                        delete next[info.logicalToolCallId];
+                        return next;
+                    });
+                }
                 if (executingCountRef.current === 0) {
-                    // biome-ignore lint/suspicious/useIterableCallbackReturn: forEach callback intentionally has no return
-                    settledResolversRef.current.forEach((resolve) => resolve());
-                    settledResolversRef.current = [];
+                    resolveAllSettledResolvers();
                 }
             },
         });
@@ -139,47 +184,49 @@ export function useToolInvocations({ state, getTools, onResult, setToolStatuses,
             .pipeThrough(new AssistantMetaTransformStream())
             .pipeTo(new WritableStream({
             write(chunk) {
-                if (chunk.type === "result") {
-                    if (shouldIgnoreAndCleanupResult(chunk.meta.toolCallId)) {
-                        return;
-                    }
-                    const logicalToolCallId = getLogicalToolCallId(chunk.meta.toolCallId);
-                    if (logicalToolCallId !== chunk.meta.toolCallId) {
-                        toolCallIdAliasesRef.current.delete(chunk.meta.toolCallId);
-                    }
-                    // the tool call result was already set by the backend
-                    if (lastToolStates.current[logicalToolCallId]?.hasResult)
-                        return;
-                    onResult({
-                        type: "add-tool-result",
-                        toolCallId: logicalToolCallId,
-                        toolName: chunk.meta.toolName,
-                        result: chunk.result,
-                        isError: chunk.isError,
-                        ...(chunk.artifact && { artifact: chunk.artifact }),
-                    });
+                if (chunk.type !== "result")
+                    return;
+                const streamId = chunk.meta.toolCallId;
+                const logicalToolCallId = getLogicalToolCallId(streamId);
+                const entry = entriesRef.current.get(logicalToolCallId);
+                // Result chunk from a rewrite-superseded stream: drop and clean
+                // up the alias.
+                if (abandonedStreamIdsRef.current.delete(streamId)) {
+                    streamToLogicalRef.current.delete(streamId);
+                    return;
+                }
+                // Pre-resolved tool call whose entry has been cleared by
+                // `reset()`. Both the real result chunk and the post-abort
+                // cancellation chunk can land here in either order; suppress
+                // both via the long-lived `skipExecuteStreamIdsRef` marker.
+                if (!entry && skipExecuteStreamIdsRef.current.has(streamId)) {
+                    return;
                 }
+                // The host already set the result (via the live snapshot's
+                // `setResponse` path). Suppress the executor's redundant emit.
+                if (entry?.hasResult)
+                    return;
+                if (streamId !== logicalToolCallId) {
+                    streamToLogicalRef.current.delete(streamId);
+                }
+                onResult({
+                    type: "add-tool-result",
+                    toolCallId: logicalToolCallId,
+                    toolName: chunk.meta.toolName,
+                    result: chunk.result,
+                    isError: chunk.isError,
+                    ...(chunk.artifact !== undefined && {
+                        artifact: chunk.artifact,
+                    }),
+                    ...(chunk.modelContent !== undefined && {
+                        modelContent: chunk.modelContent,
+                    }),
+                });
             },
         }));
         return controller;
     });
-    const ignoredToolIds = useRef(new Set());
-    const isInitialState = useRef(true);
     useEffect(() => {
-        const createToolState = ({ controller, streamToolCallId, }) => ({
-            argsText: "",
-            hasResult: false,
-            argsComplete: false,
-            streamToolCallId,
-            controller,
-        });
-        const setToolState = (toolCallId, state) => {
-            lastToolStates.current[toolCallId] = state;
-            return state;
-        };
-        const patchToolState = (toolCallId, state, patch) => {
-            return setToolState(toolCallId, { ...state, ...patch });
-        };
         const hasExecutableTool = (toolName) => {
             const tool = getTools()?.[toolName];
             return tool?.execute !== undefined || tool?.streamCall !== undefined;
@@ -188,191 +235,210 @@ export function useToolInvocations({ state, getTools, onResult, setToolStatuses,
             if (hasResult)
                 return true;
             if (!hasExecutableTool(toolName)) {
-                // Non-executable tools can emit parseable snapshots mid-stream.
-                // Wait until the run settles before closing the args stream.
+                // Non-executable tools can emit parseable JSON mid-stream; wait for
+                // the run to settle before closing.
                 return !state.isRunning && isArgsTextComplete(argsText);
             }
             return isArgsTextComplete(argsText);
         };
-        const restartToolArgsStream = ({ toolCallId, toolName, state, }) => {
-            ignoredResultToolCallIdsRef.current.add(state.streamToolCallId);
-            state.controller.argsText.close();
-            const streamToolCallId = `${toolCallId}:rewrite:${rewriteCounterRef.current++}`;
-            toolCallIdAliasesRef.current.set(streamToolCallId, toolCallId);
+        const startActiveEntry = (toolCallId, toolName, skipExecute) => {
             const toolCallController = controller.addToolCallPart({
                 toolName,
-                toolCallId: streamToolCallId,
+                toolCallId,
+            });
+            if (skipExecute) {
+                skipExecuteStreamIdsRef.current.add(toolCallId);
+            }
+            const entry = {
+                toolName,
+                controller: toolCallController,
+                streamId: toolCallId,
+                argsText: "",
+                hasResult: false,
+                argsComplete: false,
+            };
+            entriesRef.current.set(toolCallId, entry);
+            return entry;
+        };
+        const restartArgsStream = (entry, toolCallId) => {
+            if (!entry.controller)
+                return;
+            abandonedStreamIdsRef.current.add(entry.streamId);
+            // The wrapper's execute short-circuit follows the current stream id;
+            // the abandoned id stays in `skipExecuteStreamIdsRef` if it was there,
+            // which is harmless and keeps in-flight chunks consistent.
+            const wasSkipExecute = skipExecuteStreamIdsRef.current.has(entry.streamId);
+            entry.controller.argsText.close();
+            const newStreamId = `${toolCallId}:rewrite:${rewriteCounterRef.current++}`;
+            streamToLogicalRef.current.set(newStreamId, toolCallId);
+            const newController = controller.addToolCallPart({
+                toolName: entry.toolName,
+                toolCallId: newStreamId,
             });
+            if (wasSkipExecute) {
+                skipExecuteStreamIdsRef.current.add(newStreamId);
+            }
             if (process.env.NODE_ENV !== "production") {
                 console.warn("started replacement stream tool call", {
                     toolCallId,
-                    streamToolCallId,
+                    streamToolCallId: newStreamId,
                 });
             }
-            return setToolState(toolCallId, {
-                ...createToolState({
-                    controller: toolCallController,
-                    streamToolCallId,
-                }),
-                hasResult: state.hasResult,
-            });
+            entry.controller = newController;
+            entry.streamId = newStreamId;
+            entry.argsText = "";
+            entry.argsComplete = false;
+        };
+        const processArgsText = (entry, content) => {
+            if (!entry.controller)
+                return;
+            const hasResult = content.result !== undefined;
+            if (content.argsText !== entry.argsText) {
+                let shouldWriteArgsText = true;
+                if (entry.argsComplete) {
+                    if (isEquivalentCompleteArgsText(entry.argsText, content.argsText)) {
+                        entry.argsText = content.argsText;
+                        shouldWriteArgsText = false;
+                    }
+                    else {
+                        const canRestart = !entry.hasResult && !executingRef.current.has(entry.streamId);
+                        if (process.env.NODE_ENV !== "production") {
+                            console.warn(canRestart
+                                ? "argsText updated after controller was closed, restarting tool args stream:"
+                                : "argsText updated after controller was closed:", { previous: entry.argsText, next: content.argsText });
+                        }
+                        if (!canRestart) {
+                            entry.argsText = content.argsText;
+                            shouldWriteArgsText = false;
+                        }
+                        else {
+                            restartArgsStream(entry, content.toolCallId);
+                        }
+                    }
+                }
+                else if (!content.argsText.startsWith(entry.argsText)) {
+                    // Mid-stream rewrite. If both texts parse to equivalent JSON it's a
+                    // key-reorder snapshot — accept silently. Otherwise restart.
+                    if (isArgsTextComplete(entry.argsText) &&
+                        isArgsTextComplete(content.argsText) &&
+                        isEquivalentCompleteArgsText(entry.argsText, content.argsText)) {
+                        const shouldClose = shouldCloseArgsStream({
+                            toolName: content.toolName,
+                            argsText: content.argsText,
+                            hasResult,
+                        });
+                        if (shouldClose)
+                            entry.controller.argsText.close();
+                        entry.argsText = content.argsText;
+                        entry.argsComplete = shouldClose;
+                        shouldWriteArgsText = false;
+                    }
+                    else {
+                        if (process.env.NODE_ENV !== "production") {
+                            console.warn("argsText rewrote previous snapshot, restarting tool args stream:", {
+                                previous: entry.argsText,
+                                next: content.argsText,
+                                toolCallId: content.toolCallId,
+                            });
+                        }
+                        restartArgsStream(entry, content.toolCallId);
+                    }
+                }
+                if (shouldWriteArgsText) {
+                    const delta = content.argsText.slice(entry.argsText.length);
+                    entry.controller.argsText.append(delta);
+                    const shouldClose = shouldCloseArgsStream({
+                        toolName: content.toolName,
+                        argsText: content.argsText,
+                        hasResult,
+                    });
+                    if (shouldClose)
+                        entry.controller.argsText.close();
+                    entry.argsText = content.argsText;
+                    entry.argsComplete = shouldClose;
+                }
+            }
+            if (!entry.argsComplete) {
+                const shouldClose = shouldCloseArgsStream({
+                    toolName: content.toolName,
+                    argsText: content.argsText,
+                    hasResult,
+                });
+                if (shouldClose) {
+                    entry.controller.argsText.close();
+                    entry.argsText = content.argsText;
+                    entry.argsComplete = true;
+                }
+            }
         };
         const processMessages = (messages) => {
+            const isRestore = pendingRestoreRef.current;
             messages.forEach((message) => {
                 message.content.forEach((content) => {
-                    if (content.type === "tool-call") {
-                        if (isInitialState.current) {
-                            ignoredToolIds.current.add(content.toolCallId);
-                        }
-                        else {
-                            if (ignoredToolIds.current.has(content.toolCallId)) {
-                                return;
-                            }
-                            let lastState = lastToolStates.current[content.toolCallId];
-                            if (!lastState) {
-                                if (content.result !== undefined) {
-                                    if (content.messages) {
-                                        processMessages(content.messages);
-                                    }
-                                    return;
-                                }
-                                toolCallIdAliasesRef.current.set(content.toolCallId, content.toolCallId);
-                                const toolCallController = controller.addToolCallPart({
-                                    toolName: content.toolName,
-                                    toolCallId: content.toolCallId,
-                                });
-                                lastState = setToolState(content.toolCallId, createToolState({
-                                    controller: toolCallController,
-                                    streamToolCallId: content.toolCallId,
-                                }));
-                            }
-                            if (content.argsText !== lastState.argsText) {
-                                let shouldWriteArgsText = true;
-                                if (lastState.argsComplete) {
-                                    if (isEquivalentCompleteArgsText(lastState.argsText, content.argsText)) {
-                                        lastState = patchToolState(content.toolCallId, lastState, {
-                                            argsText: content.argsText,
-                                        });
-                                        shouldWriteArgsText = false;
-                                    }
-                                    if (shouldWriteArgsText) {
-                                        const canRestartClosedArgsStream = !lastState.hasResult &&
-                                            !startedExecutionToolCallIdsRef.current.has(lastState.streamToolCallId);
-                                        if (process.env.NODE_ENV !== "production") {
-                                            console.warn(canRestartClosedArgsStream
-                                                ? "argsText updated after controller was closed, restarting tool args stream:"
-                                                : "argsText updated after controller was closed:", {
-                                                previous: lastState.argsText,
-                                                next: content.argsText,
-                                            });
-                                        }
-                                        if (!canRestartClosedArgsStream) {
-                                            lastState = patchToolState(content.toolCallId, lastState, {
-                                                argsText: content.argsText,
-                                            });
-                                            shouldWriteArgsText = false;
-                                        }
-                                    }
-                                    if (shouldWriteArgsText) {
-                                        lastState = restartToolArgsStream({
-                                            toolCallId: content.toolCallId,
-                                            toolName: content.toolName,
-                                            state: lastState,
-                                        });
-                                    }
-                                }
-                                else if (!content.argsText.startsWith(lastState.argsText)) {
-                                    // Check if this is key reordering (both are complete JSON)
-                                    // This happens when transitioning from streaming to complete state
-                                    // and the provider returns keys in a different order
-                                    if (isArgsTextComplete(lastState.argsText) &&
-                                        isArgsTextComplete(content.argsText) &&
-                                        isEquivalentCompleteArgsText(lastState.argsText, content.argsText)) {
-                                        const shouldClose = shouldCloseArgsStream({
-                                            toolName: content.toolName,
-                                            argsText: content.argsText,
-                                            hasResult: content.result !== undefined,
-                                        });
-                                        if (shouldClose) {
-                                            lastState.controller.argsText.close();
-                                        }
-                                        lastState = patchToolState(content.toolCallId, lastState, {
-                                            argsText: content.argsText,
-                                            argsComplete: shouldClose,
-                                        });
-                                        shouldWriteArgsText = false;
-                                    }
-                                    if (shouldWriteArgsText) {
-                                        if (process.env.NODE_ENV !== "production") {
-                                            console.warn("argsText rewrote previous snapshot, restarting tool args stream:", {
-                                                previous: lastState.argsText,
-                                                next: content.argsText,
-                                                toolCallId: content.toolCallId,
-                                            });
-                                        }
-                                        lastState = restartToolArgsStream({
-                                            toolCallId: content.toolCallId,
-                                            toolName: content.toolName,
-                                            state: lastState,
-                                        });
-                                    }
-                                }
-                                if (shouldWriteArgsText) {
-                                    const argsTextDelta = content.argsText.slice(lastState.argsText.length);
-                                    lastState.controller.argsText.append(argsTextDelta);
-                                    const shouldClose = shouldCloseArgsStream({
-                                        toolName: content.toolName,
-                                        argsText: content.argsText,
-                                        hasResult: content.result !== undefined,
-                                    });
-                                    if (shouldClose) {
-                                        lastState.controller.argsText.close();
-                                    }
-                                    lastState = patchToolState(content.toolCallId, lastState, {
-                                        argsText: content.argsText,
-                                        argsComplete: shouldClose,
-                                    });
-                                }
-                            }
-                            if (!lastState.argsComplete) {
-                                const shouldClose = shouldCloseArgsStream({
-                                    toolName: content.toolName,
-                                    argsText: content.argsText,
-                                    hasResult: content.result !== undefined,
-                                });
-                                if (shouldClose) {
-                                    lastState.controller.argsText.close();
-                                    lastState = patchToolState(content.toolCallId, lastState, {
-                                        argsText: content.argsText,
-                                        argsComplete: true,
-                                    });
-                                }
-                            }
-                            if (content.result !== undefined && !lastState.hasResult) {
-                                patchToolState(content.toolCallId, lastState, {
-                                    hasResult: true,
-                                    argsComplete: true,
-                                });
-                                lastState.controller.setResponse(new ToolResponse({
-                                    result: content.result,
-                                    artifact: content.artifact,
-                                    isError: content.isError,
-                                }));
-                                lastState.controller.close();
-                            }
+                    if (content.type !== "tool-call")
+                        return;
+                    const existing = entriesRef.current.get(content.toolCallId);
+                    if (isRestore) {
+                        // Don't overwrite an already-active entry (e.g. live tool-call
+                        // observed before this restore snapshot landed). Restore can only
+                        // seed entries the runtime has never seen.
+                        if (!existing?.controller) {
+                            entriesRef.current.set(content.toolCallId, {
+                                toolName: content.toolName,
+                                argsText: content.argsText,
+                                hasResult: content.result !== undefined,
+                            });
                         }
-                        // Recursively process nested messages
-                        if (content.messages) {
+                        if (content.messages)
                             processMessages(content.messages);
+                        return;
+                    }
+                    // Live snapshot.
+                    let entry = existing;
+                    if (entry && !entry.controller) {
+                        // Restored entry observed in a live snapshot. Promote if its
+                        // signature has changed; otherwise treat as still-historical.
+                        const signatureChanged = content.argsText !== entry.argsText ||
+                            (content.result !== undefined) !== entry.hasResult;
+                        if (!signatureChanged) {
+                            if (content.messages)
+                                processMessages(content.messages);
+                            return;
                         }
+                        entriesRef.current.delete(content.toolCallId);
+                        entry = undefined;
+                    }
+                    if (!entry) {
+                        entry = startActiveEntry(content.toolCallId, content.toolName, content.result !== undefined);
                     }
+                    processArgsText(entry, content);
+                    if (content.result !== undefined && !entry.hasResult) {
+                        // `entry` is in active phase from this point — either it was
+                        // just created by `startActiveEntry` above, or it pre-existed
+                        // and `processArgsText` preserved (or replaced via rewrite) its
+                        // controller. Narrow once instead of asserting at every use.
+                        const { controller: activeController } = entry;
+                        if (!activeController)
+                            return;
+                        entry.hasResult = true;
+                        entry.argsComplete = true;
+                        activeController.setResponse(new ToolResponse({
+                            result: content.result,
+                            artifact: content.artifact,
+                            isError: content.isError,
+                            ...(content.modelContent !== undefined
+                                ? { modelContent: content.modelContent }
+                                : {}),
+                        }));
+                        activeController.close();
+                    }
+                    if (content.messages)
+                        processMessages(content.messages);
                 });
             });
         };
         processMessages(state.messages);
-        if (isInitialState.current) {
-            isInitialState.current = false;
-        }
+        pendingRestoreRef.current = false;
     }, [state, controller, getTools]);
     const abort = () => {
         humanInputRef.current.forEach(({ reject }) => {
@@ -381,7 +447,6 @@ export function useToolInvocations({ state, getTools, onResult, setToolStatuses,
         humanInputRef.current.clear();
         acRef.current.abort();
         acRef.current = new AbortController();
-        // Return a promise that resolves when all executing tools have settled
         if (executingCountRef.current === 0) {
             return Promise.resolve();
         }
@@ -391,13 +456,17 @@ export function useToolInvocations({ state, getTools, onResult, setToolStatuses,
     };
     return {
         reset: () => {
-            isInitialState.current = true;
-            ignoredToolIds.current.clear();
-            lastToolStates.current = {};
+            pendingRestoreRef.current = true;
+            entriesRef.current.clear();
+            // `skipExecuteStreamIdsRef` is not cleared: it has to outlive `reset()`
+            // so (a) any wrapper call still inbound through the stream pipeline
+            // continues to short-circuit `execute`, and (b) the consumer can
+            // recognize and drop any post-abort cancellation `result` chunks for
+            // pre-resolved streams whose entries have been cleared. Membership
+            // grows by one per pre-resolved tool call observed in the session.
             void abort().finally(() => {
-                startedExecutionToolCallIdsRef.current.clear();
-                toolCallIdAliasesRef.current.clear();
-                ignoredResultToolCallIdsRef.current.clear();
+                executingRef.current.clear();
+                streamToLogicalRef.current.clear();
                 rewriteCounterRef.current = 0;
             });
         },