@assistant-ui/react
Version:
TypeScript/React library for AI Chat
224 lines • 7.56 kB
JavaScript
// src/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.ts
import { useEffect, useRef, useState } from "react";
import {
createAssistantStreamController,
ToolResponse,
unstable_toolResultStream
} from "assistant-stream";
import {
AssistantMetaTransformStream
} from "assistant-stream/utils";
var isArgsTextComplete = (argsText) => {
try {
JSON.parse(argsText);
return true;
} catch {
return false;
}
};
function useToolInvocations({
state,
getTools,
onResult,
setToolStatuses
}) {
const lastToolStates = useRef({});
const humanInputRef = useRef(/* @__PURE__ */ new Map());
const acRef = useRef(new AbortController());
const executingCountRef = useRef(0);
const settledResolversRef = useRef([]);
const [controller] = useState(() => {
const [stream, controller2] = createAssistantStreamController();
const transform = unstable_toolResultStream(
getTools,
() => acRef.current?.signal ?? new AbortController().signal,
(toolCallId, payload) => {
return new Promise((resolve, reject) => {
const previous = humanInputRef.current.get(toolCallId);
if (previous) {
previous.reject(
new Error("Human input request was superseded by a new request")
);
}
humanInputRef.current.set(toolCallId, { resolve, reject });
setToolStatuses((prev) => ({
...prev,
[toolCallId]: {
type: "interrupt",
payload: { type: "human", payload }
}
}));
});
},
{
onExecutionStart: (toolCallId) => {
executingCountRef.current++;
setToolStatuses((prev) => ({
...prev,
[toolCallId]: { type: "executing" }
}));
},
onExecutionEnd: (toolCallId) => {
executingCountRef.current--;
setToolStatuses((prev) => {
const next = { ...prev };
delete next[toolCallId];
return next;
});
if (executingCountRef.current === 0) {
settledResolversRef.current.forEach((resolve) => resolve());
settledResolversRef.current = [];
}
}
}
);
stream.pipeThrough(transform).pipeThrough(new AssistantMetaTransformStream()).pipeTo(
new WritableStream({
write(chunk) {
if (chunk.type === "result") {
if (lastToolStates.current[chunk.meta.toolCallId]?.hasResult)
return;
onResult({
type: "add-tool-result",
toolCallId: chunk.meta.toolCallId,
toolName: chunk.meta.toolName,
result: chunk.result,
isError: chunk.isError,
...chunk.artifact && { artifact: chunk.artifact }
});
}
}
})
);
return controller2;
});
const ignoredToolIds = useRef(/* @__PURE__ */ new Set());
const isInitialState = useRef(true);
useEffect(() => {
const processMessages = (messages) => {
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) {
const toolCallController = controller.addToolCallPart({
toolName: content.toolName,
toolCallId: content.toolCallId
});
lastState = {
argsText: "",
hasResult: false,
argsComplete: false,
controller: toolCallController
};
lastToolStates.current[content.toolCallId] = lastState;
}
if (content.argsText !== lastState.argsText) {
if (lastState.argsComplete) {
if (process.env["NODE_ENV"] !== "production") {
console.warn(
"argsText updated after controller was closed:",
{
previous: lastState.argsText,
next: content.argsText
}
);
}
} else {
if (!content.argsText.startsWith(lastState.argsText)) {
throw new Error(
`Tool call argsText can only be appended, not updated: ${content.argsText} does not start with ${lastState.argsText}`
);
}
const argsTextDelta = content.argsText.slice(
lastState.argsText.length
);
lastState.controller.argsText.append(argsTextDelta);
const shouldClose = isArgsTextComplete(content.argsText);
if (shouldClose) {
lastState.controller.argsText.close();
}
lastToolStates.current[content.toolCallId] = {
argsText: content.argsText,
hasResult: lastState.hasResult,
argsComplete: shouldClose,
controller: lastState.controller
};
}
}
if (content.result !== void 0 && !lastState.hasResult) {
lastState.controller.setResponse(
new ToolResponse({
result: content.result,
artifact: content.artifact,
isError: content.isError
})
);
lastState.controller.close();
lastToolStates.current[content.toolCallId] = {
hasResult: true,
argsComplete: true,
argsText: lastState.argsText,
controller: lastState.controller
};
}
}
if (content.messages) {
processMessages(content.messages);
}
}
});
});
};
processMessages(state.messages);
if (isInitialState.current) {
isInitialState.current = false;
}
}, [state, controller, onResult]);
const abort = () => {
humanInputRef.current.forEach(({ reject }) => {
reject(new Error("Tool execution aborted"));
});
humanInputRef.current.clear();
acRef.current.abort();
acRef.current = new AbortController();
if (executingCountRef.current === 0) {
return Promise.resolve();
}
return new Promise((resolve) => {
settledResolversRef.current.push(resolve);
});
};
return {
reset: () => {
void abort();
isInitialState.current = true;
},
abort,
resume: (toolCallId, payload) => {
const handlers = humanInputRef.current.get(toolCallId);
if (handlers) {
humanInputRef.current.delete(toolCallId);
setToolStatuses((prev) => ({
...prev,
[toolCallId]: { type: "executing" }
}));
handlers.resolve(payload);
} else {
throw new Error(
`Tool call ${toolCallId} is not waiting for human input`
);
}
}
};
}
export {
useToolInvocations
};
//# sourceMappingURL=useToolInvocations.js.map