Skip to content

Commit de38fee

Browse files
jaszhixjaszhix
committed
fix(task): recover completed tool turn continuation
1 parent c38db3e commit de38fee

4 files changed

Lines changed: 162 additions & 3 deletions

File tree

‎src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,24 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
209209
})
210210

211211
describe("Custom tool error recording", () => {
212+
it("releases the presentation lock when dispatch throws", async () => {
213+
mockTask.assistantMessageContent = [
214+
{
215+
type: "tool_use",
216+
id: "tool_call_dispatch_failure",
217+
name: "read_file",
218+
params: {},
219+
partial: false,
220+
},
221+
]
222+
mockTask.providerRef.deref = () => ({
223+
getState: vi.fn().mockRejectedValue(new Error("provider state failed")),
224+
})
225+
226+
await expect(presentAssistantMessage(mockTask)).rejects.toThrow("provider state failed")
227+
expect(mockTask.presentAssistantMessageLocked).toBe(false)
228+
})
229+
212230
it("should record custom tool error as 'custom_tool'", async () => {
213231
const toolCallId = "tool_call_custom_error_123"
214232
mockTask.assistantMessageContent = [

‎src/core/assistant-message/presentAssistantMessage.ts‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,16 @@ export async function presentAssistantMessage(cline: Task) {
9696

9797
cline.presentAssistantMessageLocked = true
9898
cline.presentAssistantMessageHasPendingUpdates = false
99+
try {
100+
await presentAssistantMessageBlock(cline)
101+
} finally {
102+
// Tool handlers and provider-state reads can reject. Never strand the
103+
// task behind a dispatch lock after the presenter has unwound.
104+
cline.presentAssistantMessageLocked = false
105+
}
106+
}
99107

108+
async function presentAssistantMessageBlock(cline: Task): Promise<void> {
100109
if (cline.currentStreamingContentIndex >= cline.assistantMessageContent.length) {
101110
// This may happen if the last content block was completed before
102111
// streaming could finish. If streaming is finished, and we're out of
@@ -1063,7 +1072,7 @@ export async function presentAssistantMessage(cline: Task) {
10631072
if (cline.currentStreamingContentIndex < cline.assistantMessageContent.length) {
10641073
// There are already more content blocks to stream, so we'll call
10651074
// this function ourselves.
1066-
return presentAssistantMessage(cline)
1075+
return await presentAssistantMessage(cline)
10671076
} else {
10681077
// CRITICAL FIX: If we're out of bounds and the stream is complete, set userMessageContentReady
10691078
// This handles the case where assistantMessageContent is empty or becomes empty after processing
@@ -1075,7 +1084,7 @@ export async function presentAssistantMessage(cline: Task) {
10751084

10761085
// Block is partial, but the read stream may have finished.
10771086
if (cline.presentAssistantMessageHasPendingUpdates) {
1078-
return presentAssistantMessage(cline)
1087+
return await presentAssistantMessage(cline)
10791088
}
10801089
}
10811090

‎src/core/task/Task.ts‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -487,6 +487,44 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
487487
this.userMessageContent.push(toolResult)
488488
return true
489489
}
490+
491+
/**
492+
* Derives terminal tool-turn readiness from the protocol state instead of
493+
* relying exclusively on the presenter's one-shot boolean latch.
494+
*
495+
* A re-entrant presenter can lose the latch update after every tool has
496+
* already completed. At that point the assistant turn is safe to continue
497+
* when the stream is closed, presentation is idle, every content block is
498+
* final, and every tool call has its matching result.
499+
*/
500+
private hasCompleteToolResultsForCurrentTurn(): boolean {
501+
if (!this.didCompleteReadingStream || this.presentAssistantMessageLocked) {
502+
return false
503+
}
504+
505+
const toolResultIds = new Set(
506+
this.userMessageContent
507+
.filter((block): block is Anthropic.ToolResultBlockParam => block.type === "tool_result")
508+
.map((block) => block.tool_use_id),
509+
)
510+
let toolUseCount = 0
511+
512+
for (const block of this.assistantMessageContent) {
513+
if (block.partial) {
514+
return false
515+
}
516+
if (block.type !== "tool_use" && block.type !== "mcp_tool_use") {
517+
continue
518+
}
519+
520+
toolUseCount++
521+
if (!block.id || !toolResultIds.has(sanitizeToolUseId(block.id))) {
522+
return false
523+
}
524+
}
525+
526+
return toolUseCount > 0
527+
}
490528
didRejectTool = false
491529
didAlreadyUseTool = false
492530
didToolFailInCurrentTurn = false
@@ -4084,7 +4122,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
40844122
// this.userMessageContentReady = true
40854123
// }
40864124

4087-
await pWaitFor(() => this.userMessageContentReady || this.abort || this.abandoned)
4125+
await pWaitFor(
4126+
() =>
4127+
this.userMessageContentReady ||
4128+
this.hasCompleteToolResultsForCurrentTurn() ||
4129+
this.abort ||
4130+
this.abandoned,
4131+
)
40884132

40894133
if (this.abort || this.abandoned) {
40904134
throw new Error(

‎src/core/task/__tests__/Task.spec.ts‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ type TaskTestAccess = {
4343
startTask: (task?: string, images?: string[]) => Promise<void>
4444
resumeTaskFromHistory: () => Promise<void>
4545
presentAssistantMessageSafe: () => void
46+
hasCompleteToolResultsForCurrentTurn: () => boolean
4647
addToClineMessages: (message: import("@roo-code/types").ClineMessage) => Promise<void>
4748
updateClineMessage: (message: import("@roo-code/types").ClineMessage) => Promise<void>
4849
saveClineMessages: () => Promise<boolean>
@@ -576,6 +577,93 @@ describe("Cline", () => {
576577
})
577578

578579
describe("native tool-call request isolation", () => {
580+
it("derives readiness only from a closed, idle, fully paired tool turn", () => {
581+
const task = new Task({
582+
provider: mockProvider,
583+
apiConfiguration: mockApiConfig,
584+
task: "tool turn readiness test",
585+
startTask: false,
586+
})
587+
const readiness = () => getTaskTestAccess(task).hasCompleteToolResultsForCurrentTurn()
588+
task.assistantMessageContent = [
589+
{
590+
type: "tool_use",
591+
id: "call_ready",
592+
name: "read_file",
593+
params: {},
594+
nativeArgs: { path: "README.md" },
595+
partial: false,
596+
},
597+
]
598+
task.userMessageContent = [{ type: "tool_result", tool_use_id: "call_ready", content: "finished" }]
599+
600+
expect(readiness()).toBe(false)
601+
task.didCompleteReadingStream = true
602+
task.presentAssistantMessageLocked = true
603+
expect(readiness()).toBe(false)
604+
task.presentAssistantMessageLocked = false
605+
task.assistantMessageContent[0].partial = true
606+
expect(readiness()).toBe(false)
607+
task.assistantMessageContent[0].partial = false
608+
task.userMessageContent = [{ type: "tool_result", tool_use_id: "different_call", content: "finished" }]
609+
expect(readiness()).toBe(false)
610+
task.userMessageContent = [{ type: "tool_result", tool_use_id: "call_ready", content: "finished" }]
611+
expect(readiness()).toBe(true)
612+
})
613+
614+
it("continues after a complete read_file result when the readiness flag update is lost", async () => {
615+
const task = new Task({
616+
provider: mockProvider,
617+
apiConfiguration: mockApiConfig,
618+
task: "read file continuation test",
619+
startTask: false,
620+
})
621+
622+
vi.spyOn(task.diffViewProvider, "reset").mockResolvedValue(undefined)
623+
vi.spyOn(getTaskTestAccess(task), "safeEnsureModelFetched").mockResolvedValue(stubModelInfo)
624+
vi.spyOn(getTaskTestAccess(task), "presentAssistantMessageSafe").mockImplementation(() => {
625+
const completedRead = task.assistantMessageContent.find(
626+
(block) => block.type === "tool_use" && block.name === "read_file" && !block.partial,
627+
)
628+
if (!completedRead || completedRead.type !== "tool_use" || !completedRead.id) return
629+
630+
task.pushToolResultToUserContent({
631+
type: "tool_result",
632+
tool_use_id: completedRead.id,
633+
content: "File: README.md\nfinished",
634+
})
635+
// Reproduce the persisted field symptom: the tool result exists and
636+
// presentation is no longer running, but the one-shot readiness latch
637+
// never flips, so the old loop waits forever before the next request.
638+
task.userMessageContentReady = false
639+
})
640+
641+
const attemptApiRequestSpy = vi
642+
.spyOn(task, "attemptApiRequest")
643+
.mockImplementationOnce(() =>
644+
asyncStreamFrom<ApiStreamChunk>([
645+
{ type: "tool_call_partial", index: 0, id: "call_read", name: "read_file" },
646+
{ type: "tool_call_partial", index: 0, arguments: '{"path":"README.md"}' },
647+
]),
648+
)
649+
.mockImplementationOnce(() => {
650+
throw new Error("continuation request reached")
651+
})
652+
653+
vi.mocked(pWaitFor).mockImplementation(async (condition) => {
654+
if (!(await condition())) {
655+
throw new Error("read_file continuation would stall")
656+
}
657+
})
658+
659+
try {
660+
await task.recursivelyMakeClineRequests([{ type: "text", text: "read a file, then continue" }])
661+
expect(attemptApiRequestSpy).toHaveBeenCalledTimes(2)
662+
} finally {
663+
vi.mocked(pWaitFor).mockImplementation(async () => {})
664+
}
665+
})
666+
579667
it("keeps overlapping Task parser state scoped to each request", async () => {
580668
const firstTask = new Task({
581669
provider: mockProvider,

0 commit comments

Comments
 (0)