Skip to content

Commit 380766c

Browse files
ThomasK33ThomasK33
committed
fix: finalize background task tool outputs
Change-Id: I14a6bd4b2e9076208394515ac6d6bd52afaebcda Signed-off-by: Thomas Kosiewski <tk@coder.com>
1 parent 3bf29da commit 380766c

File tree

2 files changed

+316
-22
lines changed

2 files changed

+316
-22
lines changed

‎src/node/services/taskService.test.ts‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
import { describe, it, expect } from "bun:test";
2+
3+
import type { MuxMessage, MuxToolPart } from "@/common/types/message";
4+
import { Ok } from "@/common/types/result";
5+
import { TaskService } from "./taskService";
6+
7+
function createTaskToolPart(params: {
8+
toolCallId: string;
9+
state: "input-available" | "output-available";
10+
input: unknown;
11+
output?: unknown;
12+
}): MuxToolPart {
13+
if (params.state === "output-available") {
14+
return {
15+
type: "dynamic-tool",
16+
toolCallId: params.toolCallId,
17+
toolName: "task",
18+
state: "output-available",
19+
input: params.input,
20+
output: params.output,
21+
};
22+
}
23+
24+
return {
25+
type: "dynamic-tool",
26+
toolCallId: params.toolCallId,
27+
toolName: "task",
28+
state: "input-available",
29+
input: params.input,
30+
};
31+
}
32+
33+
describe("TaskService", () => {
34+
describe("tryResolveParentTaskToolCall", () => {
35+
it("should finalize background task tool output to completed", async () => {
36+
const toolCallId = "tool-1";
37+
const childWorkspaceId = "child-1";
38+
const reportMarkdown = "done";
39+
40+
const parentMessage: MuxMessage = {
41+
id: "msg-1",
42+
role: "assistant",
43+
parts: [
44+
createTaskToolPart({
45+
toolCallId,
46+
state: "output-available",
47+
input: { agentType: "research", prompt: "hi", runInBackground: true },
48+
output: { status: "started", childWorkspaceId },
49+
}),
50+
],
51+
metadata: {
52+
historySequence: 1,
53+
model: "test-model",
54+
},
55+
};
56+
57+
const reportMessage: MuxMessage = {
58+
id: "msg-2",
59+
role: "assistant",
60+
parts: [{ type: "text", text: "### Subagent report\n\n" + reportMarkdown }],
61+
metadata: {
62+
historySequence: 2,
63+
},
64+
};
65+
66+
const history: MuxMessage[] = [parentMessage, reportMessage];
67+
const emitted: unknown[] = [];
68+
const resumeCalls: unknown[] = [];
69+
70+
const historyService = {
71+
getHistory: () => Ok(history),
72+
updateHistory: (_workspaceId: string, message: MuxMessage) => {
73+
const seq = message.metadata?.historySequence;
74+
expect(seq).toBeDefined();
75+
const index = history.findIndex((m) => m.metadata?.historySequence === seq);
76+
expect(index).toBeGreaterThanOrEqual(0);
77+
history[index] = message;
78+
return Ok(undefined);
79+
},
80+
};
81+
82+
const partialService = {
83+
readPartial: () => null,
84+
writePartial: () => Ok(undefined),
85+
};
86+
87+
const workspaceService = {
88+
emitChatEvent: (_workspaceId: string, event: unknown) => {
89+
emitted.push(event);
90+
},
91+
resumeStream: (_workspaceId: string, options: unknown) => {
92+
resumeCalls.push(options);
93+
return Ok(undefined);
94+
},
95+
};
96+
97+
const config = {
98+
listWorkspaceConfigs: () => [],
99+
};
100+
101+
const aiService = {
102+
on: () => undefined,
103+
};
104+
105+
const service = new TaskService(
106+
config as never,
107+
historyService as never,
108+
partialService as never,
109+
workspaceService as never,
110+
aiService as never
111+
);
112+
113+
await (
114+
service as unknown as { tryResolveParentTaskToolCall: (params: unknown) => Promise<void> }
115+
).tryResolveParentTaskToolCall({
116+
parentWorkspaceId: "parent",
117+
parentToolCallId: toolCallId,
118+
childWorkspaceId,
119+
report: { reportMarkdown },
120+
});
121+
122+
expect(resumeCalls).toHaveLength(0);
123+
124+
const updatedToolPart = history[0].parts[0] as Extract<
125+
MuxToolPart,
126+
{ state: "output-available" }
127+
>;
128+
expect(updatedToolPart.state).toBe("output-available");
129+
expect(updatedToolPart.output).toEqual({
130+
status: "completed",
131+
childWorkspaceId,
132+
reportMarkdown,
133+
});
134+
135+
expect(emitted).toHaveLength(1);
136+
expect(emitted[0]).toMatchObject({
137+
type: "tool-call-end",
138+
toolCallId,
139+
toolName: "task",
140+
result: {
141+
status: "completed",
142+
childWorkspaceId,
143+
reportMarkdown,
144+
},
145+
});
146+
});
147+
148+
it("should auto-resume parent when the tool call was pending", async () => {
149+
const toolCallId = "tool-2";
150+
const childWorkspaceId = "child-2";
151+
const reportMarkdown = "done";
152+
153+
const parentMessage: MuxMessage = {
154+
id: "msg-1",
155+
role: "assistant",
156+
parts: [
157+
createTaskToolPart({
158+
toolCallId,
159+
state: "input-available",
160+
input: { agentType: "research", prompt: "hi" },
161+
}),
162+
],
163+
metadata: {
164+
historySequence: 1,
165+
model: "test-model",
166+
},
167+
};
168+
169+
const reportMessage: MuxMessage = {
170+
id: "msg-2",
171+
role: "assistant",
172+
parts: [{ type: "text", text: "### Subagent report\n\n" + reportMarkdown }],
173+
metadata: {
174+
historySequence: 2,
175+
},
176+
};
177+
178+
const history: MuxMessage[] = [parentMessage, reportMessage];
179+
const emitted: unknown[] = [];
180+
const resumeCalls: unknown[] = [];
181+
182+
const historyService = {
183+
getHistory: () => Ok(history),
184+
updateHistory: (_workspaceId: string, message: MuxMessage) => {
185+
const seq = message.metadata?.historySequence;
186+
expect(seq).toBeDefined();
187+
const index = history.findIndex((m) => m.metadata?.historySequence === seq);
188+
expect(index).toBeGreaterThanOrEqual(0);
189+
history[index] = message;
190+
return Ok(undefined);
191+
},
192+
};
193+
194+
const partialService = {
195+
readPartial: () => null,
196+
writePartial: () => Ok(undefined),
197+
};
198+
199+
const workspaceService = {
200+
emitChatEvent: (_workspaceId: string, event: unknown) => {
201+
emitted.push(event);
202+
},
203+
resumeStream: (_workspaceId: string, options: unknown) => {
204+
resumeCalls.push(options);
205+
return Ok(undefined);
206+
},
207+
};
208+
209+
const config = {
210+
listWorkspaceConfigs: () => [],
211+
};
212+
213+
const aiService = {
214+
on: () => undefined,
215+
};
216+
217+
const service = new TaskService(
218+
config as never,
219+
historyService as never,
220+
partialService as never,
221+
workspaceService as never,
222+
aiService as never
223+
);
224+
225+
await (
226+
service as unknown as { tryResolveParentTaskToolCall: (params: unknown) => Promise<void> }
227+
).tryResolveParentTaskToolCall({
228+
parentWorkspaceId: "parent",
229+
parentToolCallId: toolCallId,
230+
childWorkspaceId,
231+
report: { reportMarkdown },
232+
});
233+
234+
expect(resumeCalls).toHaveLength(1);
235+
expect(resumeCalls[0]).toEqual({
236+
model: "test-model",
237+
mode: undefined,
238+
toolPolicy: undefined,
239+
});
240+
241+
const updatedToolPart = history[0].parts[0] as Extract<
242+
MuxToolPart,
243+
{ state: "output-available" }
244+
>;
245+
expect(updatedToolPart.state).toBe("output-available");
246+
expect(updatedToolPart.output).toEqual({
247+
status: "completed",
248+
childWorkspaceId,
249+
reportMarkdown,
250+
});
251+
252+
expect(emitted).toHaveLength(1);
253+
expect(emitted[0]).toMatchObject({
254+
type: "tool-call-end",
255+
toolCallId,
256+
toolName: "task",
257+
});
258+
});
259+
});
260+
});

‎src/node/services/taskService.ts‎

Lines changed: 56 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -493,10 +493,18 @@ export class TaskService {
493493
report: AgentTaskReport;
494494
}): Promise<void> {
495495
let finalized: { updated: MuxMessage; output: unknown } | null = null;
496+
let shouldAutoResumeParent = false;
496497

497498
// 1) Prefer partial.json (most common after restart while waiting).
498499
const parentPartial = await this.partialService.readPartial(params.parentWorkspaceId);
499500
if (parentPartial) {
501+
const hasPendingToolCall = parentPartial.parts.some(
502+
(part) =>
503+
part.type === "dynamic-tool" &&
504+
part.toolCallId === params.parentToolCallId &&
505+
part.state === "input-available"
506+
);
507+
500508
finalized = this.updateTaskToolPartOutput(
501509
parentPartial,
502510
params.parentToolCallId,
@@ -505,6 +513,10 @@ export class TaskService {
505513
);
506514

507515
if (finalized) {
516+
if (hasPendingToolCall) {
517+
shouldAutoResumeParent = true;
518+
}
519+
508520
const writeResult = await this.partialService.writePartial(
509521
params.parentWorkspaceId,
510522
finalized.updated
@@ -560,20 +572,19 @@ export class TaskService {
560572
return;
561573
}
562574

563-
// Guard against finalizing stale tool calls.
575+
const hasPendingToolCall = best.parts.some(
576+
(part) =>
577+
part.type === "dynamic-tool" &&
578+
part.toolCallId === params.parentToolCallId &&
579+
part.state === "input-available"
580+
);
581+
564582
const maxSeq = Math.max(
565583
...historyResult.data
566584
.map((m) => m.metadata?.historySequence)
567585
.filter((n): n is number => typeof n === "number")
568586
);
569-
if (bestSeq !== maxSeq) {
570-
log.error("Refusing to finalize task tool call: not the latest message", {
571-
workspaceId: params.parentWorkspaceId,
572-
toolSeq: bestSeq,
573-
latestSeq: maxSeq,
574-
});
575-
return;
576-
}
587+
const wasLatestOrSecondLatest = bestSeq === maxSeq || bestSeq === maxSeq - 1;
577588

578589
finalized = this.updateTaskToolPartOutput(
579590
best,
@@ -585,6 +596,10 @@ export class TaskService {
585596
return;
586597
}
587598

599+
if (hasPendingToolCall && wasLatestOrSecondLatest) {
600+
shouldAutoResumeParent = true;
601+
}
602+
588603
const updateResult = await this.historyService.updateHistory(
589604
params.parentWorkspaceId,
590605
finalized.updated
@@ -608,7 +623,7 @@ export class TaskService {
608623
} satisfies WorkspaceChatMessage);
609624
}
610625

611-
if (!finalized) {
626+
if (!finalized || !shouldAutoResumeParent) {
612627
return;
613628
}
614629

@@ -636,7 +651,13 @@ export class TaskService {
636651
childWorkspaceId: string,
637652
report: AgentTaskReport
638653
): { updated: MuxMessage; output: unknown } | null {
639-
let output: unknown;
654+
const outputForTool = {
655+
status: "completed",
656+
childWorkspaceId,
657+
reportMarkdown: report.reportMarkdown,
658+
};
659+
660+
let output: unknown = null;
640661
let hasOutput = false;
641662
let changed = false;
642663

@@ -646,19 +667,32 @@ export class TaskService {
646667
}
647668

648669
if (part.state === "output-available") {
649-
return part;
650-
}
670+
const existing = part.output;
671+
const alreadyCompleted =
672+
typeof existing === "object" &&
673+
existing !== null &&
674+
"status" in existing &&
675+
(existing as { status?: unknown }).status === "completed" &&
676+
"childWorkspaceId" in existing &&
677+
(existing as { childWorkspaceId?: unknown }).childWorkspaceId === childWorkspaceId &&
678+
"reportMarkdown" in existing &&
679+
(existing as { reportMarkdown?: unknown }).reportMarkdown === report.reportMarkdown;
680+
681+
if (alreadyCompleted) {
682+
return part;
683+
}
651684

652-
const input = part.input;
653-
const runInBackground =
654-
typeof input === "object" &&
655-
input !== null &&
656-
"runInBackground" in input &&
657-
(input as { runInBackground?: unknown }).runInBackground === true;
685+
output = outputForTool;
686+
hasOutput = true;
687+
changed = true;
658688

659-
const outputForTool = runInBackground
660-
? { status: "started", childWorkspaceId }
661-
: { status: "completed", childWorkspaceId, reportMarkdown: report.reportMarkdown };
689+
const nextToolPart: MuxToolPart = {
690+
...part,
691+
output: outputForTool,
692+
};
693+
694+
return nextToolPart;
695+
}
662696

663697
output = outputForTool;
664698
hasOutput = true;

0 commit comments

Comments
 (0)