Skip to content

Commit fd5ca92

Browse files
roomoteroomote
committed
fix: resolve message queue race condition during LLM processing
- Add defensive routing in webviewMessageHandler to check for pending ask - Implement active queue monitoring in Task.ask() with event listener - Add queue draining loop in recursivelyMakeClineRequests() - Add UI hardening in ChatView to prefer queuing during transitions - Add comprehensive test coverage for race condition scenarios Fixes #8536
1 parent 6d6b836 commit fd5ca92

File tree

4 files changed

+357
-5
lines changed

4 files changed

+357
-5
lines changed

‎src/core/task/Task.ts‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -862,9 +862,49 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
862862
}
863863
}
864864

865+
// Set up queue listener to actively monitor for newly queued messages while waiting
866+
let queueListener: (() => void) | undefined
867+
if (isBlocking) {
868+
queueListener = () => {
869+
// Check if a message was queued while we're waiting
870+
if (
871+
!this.messageQueueService.isEmpty() &&
872+
this.askResponse === undefined &&
873+
this.lastMessageTs === askTs
874+
) {
875+
console.log("Task#ask detected queued message while waiting")
876+
const message = this.messageQueueService.dequeueMessage()
877+
878+
if (message) {
879+
// Check if this is a tool approval ask that needs to be handled
880+
if (
881+
type === "tool" ||
882+
type === "command" ||
883+
type === "browser_action_launch" ||
884+
type === "use_mcp_server"
885+
) {
886+
// For tool approvals, we need to approve first, then send the message if there's text/images
887+
this.handleWebviewAskResponse("yesButtonClicked", message.text, message.images)
888+
} else {
889+
// For other ask types (like followup), fulfill the ask directly
890+
this.setMessageResponse(message.text, message.images)
891+
}
892+
}
893+
}
894+
}
895+
896+
// Attach the listener
897+
this.messageQueueService.on("stateChanged", queueListener)
898+
}
899+
865900
// Wait for askResponse to be set.
866901
await pWaitFor(() => this.askResponse !== undefined || this.lastMessageTs !== askTs, { interval: 100 })
867902

903+
// Clean up queue listener
904+
if (queueListener) {
905+
this.messageQueueService.removeListener("stateChanged", queueListener)
906+
}
907+
868908
if (this.lastMessageTs !== askTs) {
869909
// Could happen if we send multiple asks in a row i.e. with
870910
// command_output. It's important that when we know an ask could
@@ -2275,6 +2315,14 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
22752315
// Reset parser after each complete conversation round
22762316
this.assistantMessageParser.reset()
22772317

2318+
// Drain any queued messages after API turn completes
2319+
// This ensures that user messages sent during LLM processing are promptly handled
2320+
// between API steps rather than being stuck in queue
2321+
while (!this.messageQueueService.isEmpty()) {
2322+
console.log("[Task] Draining message queue after API turn")
2323+
this.processQueuedMessages()
2324+
}
2325+
22782326
// Now add to apiConversationHistory.
22792327
// Need to save assistant responses to file before proceeding to
22802328
// tool use since user can exit at any moment and we wouldn't be

‎src/core/webview/__tests__/webviewMessageHandler.race-condition.spec.ts‎

Lines changed: 279 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,279 @@
1+
// npx vitest core/webview/__tests__/webviewMessageHandler.race-condition.spec.ts
2+
3+
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
4+
import * as vscode from "vscode"
5+
import { webviewMessageHandler } from "../webviewMessageHandler"
6+
import { ClineProvider } from "../ClineProvider"
7+
import { Task } from "../../task/Task"
8+
import { WebviewMessage } from "../../../shared/WebviewMessage"
9+
10+
vi.mock("vscode")
11+
12+
describe("webviewMessageHandler - Race Condition Prevention", () => {
13+
let provider: ClineProvider
14+
let task: any
15+
let messageQueueService: any
16+
let postStateToWebviewSpy: any
17+
let addMessageSpy: any
18+
19+
beforeEach(() => {
20+
// Create mock message queue service
21+
messageQueueService = {
22+
addMessage: vi.fn().mockReturnValue(true),
23+
getMessage: vi.fn(),
24+
getAllMessages: vi.fn().mockReturnValue([]),
25+
clearMessages: vi.fn(),
26+
on: vi.fn(),
27+
off: vi.fn(),
28+
emit: vi.fn(),
29+
}
30+
31+
// Create mock task - use 'any' type to allow modifiable taskAsk
32+
task = {
33+
id: "test-task-123",
34+
taskAsk: undefined, // Initially no pending ask
35+
messageQueueService,
36+
handleWebviewAskResponse: vi.fn(),
37+
}
38+
39+
// Create mock provider
40+
provider = {
41+
context: {
42+
globalState: {
43+
get: vi.fn().mockReturnValue(undefined),
44+
update: vi.fn().mockResolvedValue(undefined),
45+
},
46+
workspaceState: {
47+
get: vi.fn().mockReturnValue(undefined),
48+
update: vi.fn().mockResolvedValue(undefined),
49+
},
50+
extensionPath: "/test/path",
51+
},
52+
contextProxy: {
53+
getValue: vi.fn(),
54+
setValue: vi.fn().mockResolvedValue(undefined),
55+
globalStorageUri: { fsPath: "/test/storage" },
56+
},
57+
cwd: "/test/workspace",
58+
postMessageToWebview: vi.fn().mockResolvedValue(undefined),
59+
postStateToWebview: vi.fn().mockResolvedValue(undefined),
60+
getCurrentTask: vi.fn().mockReturnValue(task),
61+
log: vi.fn(),
62+
} as any
63+
64+
// Setup spies - use 'any' type to avoid type issues
65+
postStateToWebviewSpy = vi.spyOn(provider, "postStateToWebview")
66+
addMessageSpy = vi.spyOn(messageQueueService, "addMessage")
67+
})
68+
69+
afterEach(() => {
70+
vi.clearAllMocks()
71+
})
72+
73+
describe("Defensive routing for free user messages", () => {
74+
it('should route message to queue when askResponse is "messageResponse" and no taskAsk exists', async () => {
75+
// Arrange: No pending ask
76+
task.taskAsk = undefined
77+
78+
const message: WebviewMessage = {
79+
type: "askResponse",
80+
askResponse: "messageResponse",
81+
text: "User message during race condition",
82+
images: ["image1.png"],
83+
}
84+
85+
// Act
86+
await webviewMessageHandler(provider, message)
87+
88+
// Assert: Message should be queued instead of processed as askResponse
89+
expect(addMessageSpy).toHaveBeenCalledWith("User message during race condition", ["image1.png"])
90+
expect(task.handleWebviewAskResponse).not.toHaveBeenCalled()
91+
expect(postStateToWebviewSpy).toHaveBeenCalled()
92+
})
93+
94+
it("should process askResponse normally when taskAsk exists", async () => {
95+
// Arrange: Pending ask exists
96+
task.taskAsk = {
97+
type: "followup",
98+
text: "Need user input",
99+
} as any
100+
101+
const message: WebviewMessage = {
102+
type: "askResponse",
103+
askResponse: "messageResponse",
104+
text: "User response to ask",
105+
images: [],
106+
}
107+
108+
// Act
109+
await webviewMessageHandler(provider, message)
110+
111+
// Assert: Message should be processed as askResponse
112+
expect(task.handleWebviewAskResponse).toHaveBeenCalledWith("messageResponse", "User response to ask", [])
113+
expect(addMessageSpy).not.toHaveBeenCalled()
114+
})
115+
116+
it("should not affect other askResponse types", async () => {
117+
// Arrange: No pending ask but different response type
118+
task.taskAsk = undefined
119+
120+
const message: WebviewMessage = {
121+
type: "askResponse",
122+
askResponse: "yesButtonClicked", // Different response type
123+
}
124+
125+
// Act
126+
await webviewMessageHandler(provider, message)
127+
128+
// Assert: Should process normally (not queue)
129+
expect(task.handleWebviewAskResponse).toHaveBeenCalledWith("yesButtonClicked", undefined, undefined)
130+
expect(addMessageSpy).not.toHaveBeenCalled()
131+
})
132+
133+
it("should handle case when no current task exists", async () => {
134+
// Arrange: No current task
135+
provider.getCurrentTask = vi.fn().mockReturnValue(undefined)
136+
137+
const message: WebviewMessage = {
138+
type: "askResponse",
139+
askResponse: "messageResponse",
140+
text: "Message with no task",
141+
images: [],
142+
}
143+
144+
// Act & Assert: Should not throw error
145+
await expect(webviewMessageHandler(provider, message)).resolves.not.toThrow()
146+
expect(addMessageSpy).not.toHaveBeenCalled()
147+
})
148+
})
149+
150+
describe("Queue message handling", () => {
151+
it("should add message to queue when queueMessage type is received", async () => {
152+
// Arrange
153+
const message: WebviewMessage = {
154+
type: "queueMessage",
155+
text: "Queued message",
156+
images: ["image.jpg"],
157+
}
158+
159+
// Act
160+
await webviewMessageHandler(provider, message)
161+
162+
// Assert
163+
expect(addMessageSpy).toHaveBeenCalledWith("Queued message", ["image.jpg"])
164+
// Note: postStateToWebview is not called for queueMessage in actual implementation
165+
})
166+
167+
it("should handle queue full scenario gracefully", async () => {
168+
// Arrange: Queue is full
169+
addMessageSpy.mockReturnValue(false)
170+
171+
const message: WebviewMessage = {
172+
type: "queueMessage",
173+
text: "Message when queue is full",
174+
images: [],
175+
}
176+
177+
// Act
178+
await webviewMessageHandler(provider, message)
179+
180+
// Assert: Should attempt to add to queue even if full
181+
expect(addMessageSpy).toHaveBeenCalledWith("Message when queue is full", [])
182+
// The actual implementation uses message.text ?? "" so it passes the actual text
183+
})
184+
})
185+
186+
describe("Race condition window coverage", () => {
187+
it("should prevent message loss during rapid state transitions", async () => {
188+
// Simulate rapid fire messages during state transition
189+
const messages: WebviewMessage[] = [
190+
{ type: "askResponse", askResponse: "messageResponse", text: "Message 1", images: [] },
191+
{ type: "askResponse", askResponse: "messageResponse", text: "Message 2", images: [] },
192+
{ type: "askResponse", askResponse: "messageResponse", text: "Message 3", images: [] },
193+
]
194+
195+
// Initially no ask
196+
task.taskAsk = undefined
197+
198+
// Process all messages
199+
for (const msg of messages) {
200+
await webviewMessageHandler(provider, msg)
201+
}
202+
203+
// All messages should be queued
204+
expect(addMessageSpy).toHaveBeenCalledTimes(3)
205+
expect(addMessageSpy).toHaveBeenNthCalledWith(1, "Message 1", [])
206+
expect(addMessageSpy).toHaveBeenNthCalledWith(2, "Message 2", [])
207+
expect(addMessageSpy).toHaveBeenNthCalledWith(3, "Message 3", [])
208+
209+
// None should be processed as askResponse
210+
expect(task.handleWebviewAskResponse).not.toHaveBeenCalled()
211+
})
212+
213+
it("should handle mixed message scenarios correctly", async () => {
214+
// Start with no ask
215+
task.taskAsk = undefined
216+
217+
// Message 1: Should be queued
218+
await webviewMessageHandler(provider, {
219+
type: "askResponse",
220+
askResponse: "messageResponse",
221+
text: "Message during no ask",
222+
images: [],
223+
})
224+
225+
expect(addMessageSpy).toHaveBeenCalledWith("Message during no ask", [])
226+
expect(addMessageSpy).toHaveBeenCalledTimes(1)
227+
228+
// Now set taskAsk
229+
task.taskAsk = { type: "followup" } as any
230+
231+
// Message 2: Should be processed
232+
await webviewMessageHandler(provider, {
233+
type: "askResponse",
234+
askResponse: "messageResponse",
235+
text: "Message during ask",
236+
images: [],
237+
})
238+
239+
expect(task.handleWebviewAskResponse).toHaveBeenCalledWith("messageResponse", "Message during ask", [])
240+
241+
// Remove taskAsk again
242+
task.taskAsk = undefined
243+
244+
// Message 3: Should be queued again
245+
await webviewMessageHandler(provider, {
246+
type: "askResponse",
247+
askResponse: "messageResponse",
248+
text: "Another message during no ask",
249+
images: [],
250+
})
251+
252+
expect(addMessageSpy).toHaveBeenCalledTimes(2)
253+
expect(addMessageSpy).toHaveBeenNthCalledWith(2, "Another message during no ask", [])
254+
})
255+
})
256+
257+
describe("Console logging behavior", () => {
258+
it("should log when routing message to queue due to no pending ask", async () => {
259+
// Arrange
260+
const consoleSpy = vi.spyOn(console, "log").mockImplementation(() => {})
261+
task.taskAsk = undefined
262+
263+
const message: WebviewMessage = {
264+
type: "askResponse",
265+
askResponse: "messageResponse",
266+
text: "Test message",
267+
images: [],
268+
}
269+
270+
// Act
271+
await webviewMessageHandler(provider, message)
272+
273+
// Assert
274+
expect(consoleSpy).toHaveBeenCalledWith("[webviewMessageHandler] No pending ask, routing message to queue")
275+
276+
consoleSpy.mockRestore()
277+
})
278+
})
279+
})

‎src/core/webview/webviewMessageHandler.ts‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -577,9 +577,25 @@ export const webviewMessageHandler = async (
577577
await updateGlobalState("alwaysAllowUpdateTodoList", message.bool)
578578
await provider.postStateToWebview()
579579
break
580-
case "askResponse":
581-
provider.getCurrentTask()?.handleWebviewAskResponse(message.askResponse!, message.text, message.images)
580+
case "askResponse": {
581+
// Defensive routing: Check if there's an active ask before processing the response
582+
// This prevents messages from vanishing when sent during the race condition window
583+
const task = provider.getCurrentTask()
584+
if (task) {
585+
// If this is a "messageResponse" and there's no pending ask, queue it instead
586+
if (message.askResponse === "messageResponse" && !task.taskAsk) {
587+
// No ask is pending, route to queue to ensure message is captured
588+
console.log("[webviewMessageHandler] No pending ask, routing message to queue")
589+
task.messageQueueService.addMessage(message.text || "", message.images)
590+
// Post state to update UI with queued message
591+
await provider.postStateToWebview()
592+
} else {
593+
// Normal flow: there's a pending ask or it's not a messageResponse
594+
task.handleWebviewAskResponse(message.askResponse!, message.text, message.images)
595+
}
596+
}
582597
break
598+
}
583599
case "autoCondenseContext":
584600
await updateGlobalState("autoCondenseContext", message.bool)
585601
await provider.postStateToWebview()

0 commit comments

Comments
 (0)