Browse Source

fix: prevent race condition from deleting wrong API messages (#10113)

Co-authored-by: daniel-lxs <[email protected]>
Hannes Rudolph 4 weeks ago
parent
commit
69307ba62e
2 changed files with 159 additions and 5 deletions
  1. 121 0
      src/core/message-manager/index.spec.ts
  2. 38 5
      src/core/message-manager/index.ts

+ 121 - 0
src/core/message-manager/index.spec.ts

@@ -728,4 +728,125 @@ describe("MessageManager", () => {
 			expect(apiCall[0].role).toBe("system")
 		})
 	})
+
+	describe("Race condition handling", () => {
+		it("should preserve assistant message when clineMessage timestamp is earlier due to async execution", async () => {
+			// This test reproduces the bug where deleting a user_feedback clineMessage
+			// incorrectly removes an assistant API message that was added AFTER the
+			// clineMessage (due to async tool execution during streaming).
+			//
+			// Timeline (race condition scenario):
+			// - T1 (100): clineMessage "user_feedback" created during tool execution
+			// - T2 (200): assistant API message added when stream completes
+			// - T3 (300): user API message (tool_result) added after pWaitFor
+			//
+			// When deleting the clineMessage at T1, we should:
+			// - Keep the assistant message at T2
+			// - Remove the user message at T3
+
+			mockTask.clineMessages = [
+				{ ts: 50, say: "user", text: "Initial request" },
+				{ ts: 100, say: "user_feedback", text: "tell me a joke 3" }, // Race: created BEFORE assistant API msg
+			]
+
+			mockTask.apiConversationHistory = [
+				{ ts: 50, role: "user", content: [{ type: "text", text: "Initial request" }] },
+				{
+					ts: 200, // Race: added AFTER clineMessage at ts=100
+					role: "assistant",
+					content: [{ type: "tool_use", id: "tool_1", name: "attempt_completion", input: {} }],
+				},
+				{
+					ts: 300,
+					role: "user",
+					content: [{ type: "tool_result", tool_use_id: "tool_1", content: "tell me a joke 3" }],
+				},
+			]
+
+			// Delete the user_feedback clineMessage at ts=100
+			await manager.rewindToTimestamp(100)
+
+			// The fix ensures we find the first API user message at or after cutoff (ts=300)
+			// and use that as the actual cutoff, preserving the assistant message (ts=200)
+			const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0]
+			expect(apiCall).toHaveLength(2)
+			expect(apiCall[0].ts).toBe(50) // Initial user message preserved
+			expect(apiCall[1].ts).toBe(200) // Assistant message preserved (was incorrectly removed before fix)
+			expect(apiCall[1].role).toBe("assistant")
+		})
+
+		it("should handle normal case where timestamps are properly ordered", async () => {
+			// Normal case: clineMessage timestamp aligns with API message timestamp
+			mockTask.clineMessages = [
+				{ ts: 100, say: "user", text: "First" },
+				{ ts: 200, say: "user_feedback", text: "Feedback" },
+			]
+
+			mockTask.apiConversationHistory = [
+				{ ts: 100, role: "user", content: [{ type: "text", text: "First" }] },
+				{ ts: 150, role: "assistant", content: [{ type: "text", text: "Response" }] },
+				{ ts: 200, role: "user", content: [{ type: "text", text: "Feedback" }] },
+			]
+
+			await manager.rewindToTimestamp(200)
+
+			// Should keep messages before the user message at ts=200
+			const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0]
+			expect(apiCall).toHaveLength(2)
+			expect(apiCall[0].ts).toBe(100)
+			expect(apiCall[1].ts).toBe(150)
+		})
+
+		it("should fall back to original cutoff when no user message found at or after cutoff", async () => {
+			// Edge case: no API user message at or after the cutoff timestamp
+			mockTask.clineMessages = [
+				{ ts: 100, say: "user", text: "First" },
+				{ ts: 200, say: "assistant", text: "Response" },
+				{ ts: 300, say: "assistant", text: "Another response" },
+			]
+
+			mockTask.apiConversationHistory = [
+				{ ts: 100, role: "user", content: [{ type: "text", text: "First" }] },
+				{ ts: 150, role: "assistant", content: [{ type: "text", text: "Response" }] },
+				{ ts: 250, role: "assistant", content: [{ type: "text", text: "Another response" }] },
+				// No user message at or after ts=200
+			]
+
+			await manager.rewindToTimestamp(200)
+
+			// Falls back to original behavior: keep messages with ts < 200
+			// This removes the assistant message at ts=250
+			const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0]
+			expect(apiCall).toHaveLength(2)
+			expect(apiCall[0].ts).toBe(100)
+			expect(apiCall[1].ts).toBe(150)
+		})
+
+		it("should handle multiple assistant messages before the user message in race condition", async () => {
+			// Complex race scenario: multiple assistant messages added before user message
+			mockTask.clineMessages = [
+				{ ts: 50, say: "user", text: "Initial" },
+				{ ts: 100, say: "user_feedback", text: "Feedback" }, // Race condition
+			]
+
+			mockTask.apiConversationHistory = [
+				{ ts: 50, role: "user", content: [{ type: "text", text: "Initial" }] },
+				{ ts: 150, role: "assistant", content: [{ type: "text", text: "First assistant msg" }] }, // After clineMessage
+				{ ts: 200, role: "assistant", content: [{ type: "text", text: "Second assistant msg" }] }, // After clineMessage
+				{ ts: 250, role: "user", content: [{ type: "text", text: "Feedback" }] },
+			]
+
+			await manager.rewindToTimestamp(100)
+
+			// Should preserve both assistant messages (ts=150, ts=200) because the first
+			// user message at or after cutoff is at ts=250
+			const apiCall = mockTask.overwriteApiConversationHistory.mock.calls[0][0]
+			expect(apiCall).toHaveLength(3)
+			expect(apiCall[0].ts).toBe(50)
+			expect(apiCall[1].ts).toBe(150)
+			expect(apiCall[1].role).toBe("assistant")
+			expect(apiCall[2].ts).toBe(200)
+			expect(apiCall[2].role).toBe("assistant")
+		})
+	})
 })

+ 38 - 5
src/core/message-manager/index.ts

@@ -133,6 +133,14 @@ export class MessageManager {
 	 * 1. Avoids multiple writes to API history
 	 * 2. Only writes if the history actually changed
 	 * 3. Handles both truncation and cleanup atomically
+	 *
+	 * Note on timestamp handling:
+	 * Due to async execution during streaming, clineMessage timestamps may not
+	 * perfectly align with API message timestamps. Specifically, a "user_feedback"
+	 * clineMessage can have a timestamp BEFORE the assistant API message that
+	 * triggered it (because tool execution happens concurrently with stream
+	 * completion). To handle this race condition, we find the first API user
+	 * message at or after the cutoff and use its timestamp as the actual boundary.
 	 */
 	private async truncateApiHistoryWithCleanup(
 		cutoffTs: number,
@@ -142,10 +150,35 @@ export class MessageManager {
 		const originalHistory = this.task.apiConversationHistory
 		let apiHistory = [...originalHistory]
 
-		// Step 1: Filter by timestamp
-		apiHistory = apiHistory.filter((m) => !m.ts || m.ts < cutoffTs)
+		// Step 1: Determine the actual cutoff timestamp
+		// Check if there's an API message with an exact timestamp match
+		const hasExactMatch = apiHistory.some((m) => m.ts === cutoffTs)
+		// Check if there are any messages before the cutoff that would be preserved
+		const hasMessageBeforeCutoff = apiHistory.some((m) => m.ts !== undefined && m.ts < cutoffTs)
+
+		let actualCutoff: number = cutoffTs
+
+		if (!hasExactMatch && hasMessageBeforeCutoff) {
+			// No exact match but there are earlier messages means we might have a race
+			// condition where the clineMessage timestamp is earlier than any API message
+			// due to async execution. In this case, look for the first API user message
+			// at or after the cutoff to use as the actual boundary.
+			// This ensures assistant messages that preceded the user's response are preserved.
+			const firstUserMsgIndexToRemove = apiHistory.findIndex(
+				(m) => m.ts !== undefined && m.ts >= cutoffTs && m.role === "user",
+			)
+
+			if (firstUserMsgIndexToRemove !== -1) {
+				// Use the user message's timestamp as the actual cutoff
+				actualCutoff = apiHistory[firstUserMsgIndexToRemove].ts!
+			}
+			// else: no user message found, use original cutoffTs (fallback)
+		}
+
+		// Step 2: Filter by the actual cutoff timestamp
+		apiHistory = apiHistory.filter((m) => !m.ts || m.ts < actualCutoff)
 
-		// Step 2: Remove Summaries whose condense_context was removed
+		// Step 3: Remove Summaries whose condense_context was removed
 		if (removedIds.condenseIds.size > 0) {
 			apiHistory = apiHistory.filter((msg) => {
 				if (msg.isSummary && msg.condenseId && removedIds.condenseIds.has(msg.condenseId)) {
@@ -156,7 +189,7 @@ export class MessageManager {
 			})
 		}
 
-		// Step 3: Remove truncation markers whose sliding_window_truncation was removed
+		// Step 4: Remove truncation markers whose sliding_window_truncation was removed
 		if (removedIds.truncationIds.size > 0) {
 			apiHistory = apiHistory.filter((msg) => {
 				if (msg.isTruncationMarker && msg.truncationId && removedIds.truncationIds.has(msg.truncationId)) {
@@ -169,7 +202,7 @@ export class MessageManager {
 			})
 		}
 
-		// Step 4: Cleanup orphaned tags (unless skipped)
+		// Step 5: Cleanup orphaned tags (unless skipped)
 		if (!skipCleanup) {
 			apiHistory = cleanupAfterTruncation(apiHistory)
 		}