fix: tool calls in chat

This commit is contained in:
Richard Tang
2026-04-20 18:10:53 -07:00
parent e7daa59573
commit 2644ab953d
3 changed files with 126 additions and 59 deletions
@@ -515,6 +515,85 @@ describe("replayEventsToMessages", () => {
"stream-session-1-0-t1-research",
]);
});
it("does not carry completed queen tools into a scheduler run", () => {
const events = [
makeEvent({
type: "tool_call_started",
stream_id: "queen",
node_id: "queen",
execution_id: "session-setup",
data: { tool_name: "create_colony", tool_use_id: "tool-create" },
}),
makeEvent({
type: "tool_call_completed",
stream_id: "queen",
node_id: "queen",
execution_id: "session-setup",
data: { tool_name: "create_colony", tool_use_id: "tool-create" },
}),
makeEvent({
type: "llm_turn_complete",
stream_id: "queen",
node_id: "queen",
execution_id: "session-setup",
}),
makeEvent({
type: "node_loop_started",
stream_id: "queen",
node_id: "queen",
execution_id: "session-scheduler",
}),
makeEvent({
type: "tool_call_started",
stream_id: "queen",
node_id: "queen",
execution_id: "session-scheduler",
data: {
tool_name: "list_worker_questions",
tool_use_id: "tool-questions",
},
}),
makeEvent({
type: "tool_call_started",
stream_id: "queen",
node_id: "queen",
execution_id: "session-scheduler",
data: { tool_name: "get_worker_status", tool_use_id: "tool-status" },
}),
makeEvent({
type: "tool_call_completed",
stream_id: "queen",
node_id: "queen",
execution_id: "session-scheduler",
data: {
tool_name: "list_worker_questions",
tool_use_id: "tool-questions",
},
}),
makeEvent({
type: "tool_call_completed",
stream_id: "queen",
node_id: "queen",
execution_id: "session-scheduler",
data: { tool_name: "get_worker_status", tool_use_id: "tool-status" },
}),
];
const restored = replayEventsToMessages(events, "queen-dm", "Alexandra");
const schedulerToolRow = restored.find(
(m) => m.id === "tool-pill-queen-session-scheduler-1",
);
expect(schedulerToolRow).toBeDefined();
expect(JSON.parse(schedulerToolRow!.content)).toEqual({
tools: [
{ name: "list_worker_questions", done: true },
{ name: "get_worker_status", done: true },
],
allDone: true,
});
});
});
// ---------------------------------------------------------------------------
+31 -11
View File
@@ -261,7 +261,10 @@ export interface ReplayState {
string,
{ name: string; done: boolean; streamId: string }
>;
toolUseToPill: Record<string, { msgId: string; name: string }>;
toolUseToPill: Record<
string,
{ msgId: string; name: string; streamId: string }
>;
queenIterText: Record<string, Record<number, string>>;
}
@@ -274,6 +277,20 @@ export function newReplayState(): ReplayState {
};
}
function clearToolStateForStream(state: ReplayState, streamId: string): void {
const activeToolCalls: typeof state.activeToolCalls = {};
for (const [toolUseId, tool] of Object.entries(state.activeToolCalls)) {
if (tool.streamId !== streamId) activeToolCalls[toolUseId] = tool;
}
state.activeToolCalls = activeToolCalls;
const toolUseToPill: typeof state.toolUseToPill = {};
for (const [toolUseId, pill] of Object.entries(state.toolUseToPill)) {
if (pill.streamId !== streamId) toolUseToPill[toolUseId] = pill;
}
state.toolUseToPill = toolUseToPill;
}
/**
* Process a single event and emit zero or more ChatMessage upserts.
*
@@ -317,15 +334,14 @@ export function replayEvent(
switch (event.type) {
case "execution_started":
state.turnCounters[turnKey] = currentTurn + 1;
// New execution for a worker resets its active tools so stale
// tool pills from a previous run cannot bleed into the next run.
if (!isQueen) {
const keepActive: typeof state.activeToolCalls = {};
for (const [k, v] of Object.entries(state.activeToolCalls)) {
if (v.streamId !== streamId) keepActive[k] = v;
}
state.activeToolCalls = keepActive;
}
// New executions reset their active tools so stale completed pills
// from a previous run cannot bleed into the next run.
clearToolStateForStream(state, streamId);
break;
case "node_loop_started":
// Queen-triggered scheduler runs emit node_loop_started rather than
// execution_started, so use it as an execution boundary too.
clearToolStateForStream(state, streamId);
break;
case "llm_turn_complete":
state.turnCounters[turnKey] = currentTurn + 1;
@@ -341,7 +357,11 @@ export function replayEvent(
};
const pillId = `tool-pill-${streamId}-${event.execution_id || "exec"}-${currentTurn}`;
if (toolUseId) {
state.toolUseToPill[toolUseId] = { msgId: pillId, name: toolName };
state.toolUseToPill[toolUseId] = {
msgId: pillId,
name: toolName,
streamId,
};
}
const tools = Object.values(state.activeToolCalls)
.filter((t) => t.streamId === streamId)
+16 -48
View File
@@ -51,7 +51,6 @@ export default function QueenDM() {
null,
);
const [creatingNewSession, setCreatingNewSession] = useState(false);
const [spawning, setSpawning] = useState(false);
const [initialDraft, setInitialDraft] = useState<string | null>(null);
const [cloneDialogOpen, setCloneDialogOpen] = useState(false);
const [cloneColonyName, setCloneColonyName] = useState("");
@@ -422,49 +421,6 @@ export default function QueenDM() {
}
}, [sessionId, compactingAndForking, queenId, setSearchParams]);
const handleColonySpawn = useCallback(async () => {
if (!sessionId || spawning) return;
const colony = cloneColonyName.trim();
if (!colony) return;
setSpawning(true);
try {
const result = await executionApi.colonySpawn(
sessionId,
colony,
cloneTask.trim() || undefined,
);
const msg: ChatMessage = {
id: makeId(),
agent: "System",
agentColor: "",
content: `Forked to colony "${result.colony_name}"${result.is_new ? " (new)" : ""} — session ${result.queen_session_id}`,
timestamp: "",
type: "system",
thread: "queen-dm",
createdAt: Date.now(),
};
setMessages((prev) => [...prev, msg]);
setCloneDialogOpen(false);
setCloneColonyName("");
setCloneTask("");
} catch (err) {
const errMsg = err instanceof Error ? err.message : String(err);
const msg: ChatMessage = {
id: makeId(),
agent: "System",
agentColor: "",
content: `Clone failed: ${errMsg}`,
timestamp: "",
type: "system",
thread: "queen-dm",
createdAt: Date.now(),
};
setMessages((prev) => [...prev, msg]);
} finally {
setSpawning(false);
}
}, [sessionId, spawning, cloneColonyName, cloneTask]);
const handleSelectHistoricalSession = useCallback(
(nextSessionId: string) => {
if (!nextSessionId || nextSessionId === sessionId) return;
@@ -692,6 +648,19 @@ export default function QueenDM() {
[sessionId, awaitingInput, isTyping],
);
const handleColonySpawn = useCallback(() => {
const colony = cloneColonyName.trim();
if (!colony) return;
const task = cloneTask.trim();
const message = task
? `Create a colony named \`${colony}\` for the following task:\n\n${task}`
: `Create a colony named \`${colony}\` from this session.`;
handleSend(message, "queen-dm");
setCloneDialogOpen(false);
setCloneColonyName("");
setCloneTask("");
}, [cloneColonyName, cloneTask, handleSend]);
const handleQuestionAnswer = useCallback(
(answers: Record<string, string>) => {
setAwaitingInput(false);
@@ -788,7 +757,7 @@ export default function QueenDM() {
<div className="fixed inset-0 z-50 flex items-center justify-center">
<div
className="absolute inset-0 bg-black/40 backdrop-blur-sm"
onClick={() => !spawning && setCloneDialogOpen(false)}
onClick={() => setCloneDialogOpen(false)}
/>
<div className="relative bg-card border border-border/60 rounded-xl shadow-2xl w-full max-w-md p-6 space-y-4">
<h2 className="text-sm font-semibold text-foreground">
@@ -837,17 +806,16 @@ export default function QueenDM() {
setCloneColonyName("");
setCloneTask("");
}}
disabled={spawning}
className="px-3 py-1.5 rounded-md text-xs font-medium text-muted-foreground hover:text-foreground hover:bg-muted/50 transition-colors"
>
Cancel
</button>
<button
onClick={handleColonySpawn}
disabled={spawning || !cloneColonyName.trim()}
disabled={!cloneColonyName.trim()}
className="px-3 py-1.5 rounded-md text-xs font-medium bg-primary text-primary-foreground hover:bg-primary/90 transition-colors disabled:opacity-50"
>
{spawning ? "Creating..." : "Create"}
Create
</button>
</div>
</div>