@just-every/ensemble
Version:
LLM provider abstraction layer with unified streaming interface
133 lines • 4.23 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.runningToolTracker = exports.RunningToolTracker = void 0;
class RunningToolTracker {
runningTools = new Map();
completionHandlers = [];
addRunningTool(id, toolName, agentName, args) {
const abortController = new AbortController();
const runningTool = {
id,
toolName,
agentName,
args,
startTime: Date.now(),
abortController,
};
this.runningTools.set(id, runningTool);
return runningTool;
}
markTimedOut(id) {
const tool = this.runningTools.get(id);
if (tool) {
tool.timedOut = true;
}
}
async completeRunningTool(id, result, _agent) {
const tool = this.runningTools.get(id);
if (!tool)
return;
tool.completed = true;
tool.result = result;
const duration = Date.now() - tool.startTime;
if (tool.timedOut) {
const event = {
id,
toolName: tool.toolName,
agentName: tool.agentName,
duration,
timedOut: true,
result,
};
this.completionHandlers.forEach(handler => handler(event));
}
this.runningTools.delete(id);
}
async failRunningTool(id, error, _agent) {
const tool = this.runningTools.get(id);
if (!tool)
return;
tool.failed = true;
tool.error = error;
const duration = Date.now() - tool.startTime;
if (tool.timedOut) {
const event = {
id,
toolName: tool.toolName,
agentName: tool.agentName,
duration,
timedOut: true,
error,
};
this.completionHandlers.forEach(handler => handler(event));
}
this.runningTools.delete(id);
}
getRunningTool(id) {
return this.runningTools.get(id);
}
getAllRunningTools() {
return Array.from(this.runningTools.values());
}
getRunningToolsForAgent(agentName) {
return this.getAllRunningTools().filter(tool => tool.agentName === agentName);
}
abortRunningTool(id) {
const tool = this.runningTools.get(id);
if (tool && tool.abortController) {
tool.abortController.abort();
}
}
onCompletion(handler) {
this.completionHandlers.push(handler);
}
clear() {
this.runningTools.forEach(tool => {
if (tool.abortController) {
tool.abortController.abort();
}
});
this.runningTools.clear();
this.completionHandlers = [];
}
isToolRunning(agentName, toolName) {
return this.getAllRunningTools().some(tool => tool.agentName === agentName && tool.toolName === toolName);
}
getRunningToolCount() {
return this.runningTools.size;
}
async waitForTool(id, timeout) {
return new Promise((resolve, reject) => {
const tool = this.runningTools.get(id);
if (!tool) {
resolve(null);
return;
}
let timeoutId;
const cleanup = () => {
if (timeoutId)
clearTimeout(timeoutId);
const index = this.completionHandlers.indexOf(handler);
if (index > -1) {
this.completionHandlers.splice(index, 1);
}
};
const handler = (event) => {
if (event.id === id) {
cleanup();
resolve(event);
}
};
this.onCompletion(handler);
if (timeout) {
timeoutId = setTimeout(() => {
cleanup();
reject(new Error(`Timeout waiting for tool ${id}`));
}, timeout);
}
});
}
}
exports.RunningToolTracker = RunningToolTracker;
exports.runningToolTracker = new RunningToolTracker();
//# sourceMappingURL=running_tool_tracker.js.map