fix: resolve orphaned partial ask messages bug and remove CLI workaround

Author: catrielmuller

cli/src/state/atoms/__tests__/partial-race-condition.test.ts

@@ -98,50 +98,6 @@ describe("Partial Race Condition Bug Fix", () => {
 		expect(messages[0]?.text).toBe("git status")
 	})
 
-	it("should auto-complete orphaned partial ask when subsequent messages arrive (CLI workaround)", () => {
-		// Simulate the extension bug: partial ask message followed by other messages
-		// without ever sending a partial=false update
-		const stateWithOrphanedPartial: ExtensionState = {
-			version: "1.0.0",
-			apiConfiguration: {},
-			chatMessages: [
-				{
-					ts: 1000,
-					type: "ask",
-					ask: "command",
-					text: "git status",
-					partial: true, // Orphaned - never completed by extension
-				},
-				{
-					ts: 2000,
-					type: "say",
-					say: "checkpoint_saved",
-					text: "abc123",
-				},
-				{
-					ts: 3000,
-					type: "say",
-					say: "text",
-					text: "Command executed",
-				},
-			],
-			mode: "code",
-			customModes: [],
-			taskHistoryFullLength: 0,
-			taskHistoryVersion: 0,
-			renderContext: "cli",
-			telemetrySetting: "disabled",
-		}
-
-		store.set(updateExtensionStateAtom, stateWithOrphanedPartial)
-
-		// CLI should auto-complete the orphaned partial ask
-		const messages = store.get(chatMessagesAtom)
-		expect(messages).toHaveLength(3)
-		expect(messages[0]?.partial).toBe(false) // Auto-completed!
-		expect(messages[0]?.ask).toBe("command")
-	})
-
 	it("should handle the race condition for any ask message type", () => {
 		// Test with followup message
 		const initialState: ExtensionState = {

cli/src/state/atoms/extension.ts

@@ -197,10 +197,7 @@ export const updateExtensionStateAtom = atom(null, (get, set, state: ExtensionSt
 		const incomingMessages = state.clineMessages || state.chatMessages || []
 
 		// Reconcile with current messages to preserve streaming state
-		let reconciledMessages = reconcileMessages(currentMessages, incomingMessages, versionMap, streamingSet)
-
-		// Auto-complete orphaned partial ask messages (CLI-only workaround for extension bug)
-		reconciledMessages = autoCompleteOrphanedPartialAsks(reconciledMessages)
+		const reconciledMessages = reconcileMessages(currentMessages, incomingMessages, versionMap, streamingSet)
 
 		set(chatMessagesAtom, reconciledMessages)
 
@@ -459,57 +456,6 @@ function getMessageContentLength(msg: ExtensionChatMessage): number {
 	return length
 }
 
-/**
- * Auto-complete orphaned partial ask messages (CLI-only workaround)
- *
- * This handles the extension bug where ask messages can get stuck with partial=true
- * when other messages (like checkpoint_saved) are added between the partial message
- * and its completion, causing the extension to create a new message instead of updating.
- *
- * Detection logic:
- * - If an ask message has partial=true
- * - AND there's a subsequent message with a later timestamp
- * - AND that subsequent message is NOT command_output (which is expected during command execution)
- * - THEN mark the partial ask as complete (partial=false)
- *
- * This ensures messages don't get stuck in partial state indefinitely.
- */
-function autoCompleteOrphanedPartialAsks(messages: ExtensionChatMessage[]): ExtensionChatMessage[] {
-	const result = [...messages]
-
-	for (let i = 0; i < result.length; i++) {
-		const msg = result[i]
-
-		// Only process partial ask messages
-		if (!msg || msg.type !== "ask" || !msg.partial) {
-			continue
-		}
-
-		// Check if there's a subsequent message (not command_output)
-		let hasSubsequentMessage = false
-		for (let j = i + 1; j < result.length; j++) {
-			const nextMsg = result[j]
-			if (!nextMsg) continue
-
-			// Skip command_output messages as they're expected during command execution
-			if (nextMsg.ask === "command_output") {
-				continue
-			}
-
-			// Found a subsequent non-command_output message
-			hasSubsequentMessage = true
-			break
-		}
-
-		// If there's a subsequent message, this partial ask is orphaned - mark it complete
-		if (hasSubsequentMessage) {
-			result[i] = { ...msg, partial: false }
-		}
-	}
-
-	return result
-}
-
 /**
  * Helper function to reconcile messages from state updates with existing messages
  * Strategy:

src/core/task/Task.ts

@@ -119,6 +119,7 @@ import { ensureLocalKilorulesDirExists } from "../context/instructions/kilo-rule
 import { getMessagesSinceLastSummary, summarizeConversation } from "../condense"
 import { Gpt5Metadata, ClineMessageWithMetadata } from "./types"
 import { MessageQueueService } from "../message-queue/MessageQueueService"
+import { findPartialAskMessage, findPartialSayMessage } from "./message-utils" // kilocode_change
 
 import { AutoApprovalHandler } from "./AutoApprovalHandler"
 import { isAnyRecognizedKiloCodeError, isPaymentRequiredError } from "../../shared/kilocode/errorUtils"
@@ -768,10 +769,14 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		let askTs: number
 
 		if (partial !== undefined) {
-			const lastMessage = this.clineMessages.at(-1)
+			// kilocode_change start: Fix orphaned partial asks by searching backwards
+			// Search for the most recent partial ask of this type, handling cases where
+			// non-interactive messages (like checkpoint_saved) are inserted during streaming
+			const partialResult = findPartialAskMessage(this.clineMessages, type)
+			const lastMessage = partialResult?.message
+			// kilocode_change end
 
-			const isUpdatingPreviousPartial =
-				lastMessage && lastMessage.partial && lastMessage.type === "ask" && lastMessage.ask === type
+			const isUpdatingPreviousPartial = lastMessage !== undefined
 
 			if (partial) {
 				if (isUpdatingPreviousPartial) {
@@ -1157,10 +1162,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		}
 
 		if (partial !== undefined) {
-			const lastMessage = this.clineMessages.at(-1)
+			// kilocode_change start: Fix orphaned partial says by searching backwards
+			// Search for the most recent partial say of this type
+			const partialResult = findPartialSayMessage(this.clineMessages, type)
+			const lastMessage = partialResult?.message
+			// kilocode_change end
 
-			const isUpdatingPreviousPartial =
-				lastMessage && lastMessage.partial && lastMessage.type === "say" && lastMessage.say === type
+			const isUpdatingPreviousPartial = lastMessage !== undefined
 
 			if (partial) {
 				if (isUpdatingPreviousPartial) {
@@ -3146,11 +3154,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		// stream.
 		yield* iterator
 
-        // kilocode_change start
+		// kilocode_change start
 		if (apiConfiguration?.rateLimitAfter) {
 			Task.lastGlobalApiRequestTime = performance.now()
 		}
-        // kilocode_change end
+		// kilocode_change end
 	}
 
 	// Checkpoints

src/core/task/__tests__/message-utils.spec.ts

@@ -0,0 +1,124 @@
+/**
+ * Tests for message utility functions
+ * @kilocode_change - Tests for orphaned partial ask messages bug fix
+ */
+
+import { describe, it, expect } from "vitest"
+import { findPartialAskMessage, findPartialSayMessage } from "../message-utils"
+import type { ClineMessage } from "@roo-code/types"
+
+describe("findPartialAskMessage", () => {
+	it("should find the most recent partial ask message", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "ask", ask: "tool", text: "test", partial: false },
+			{ ts: 2, type: "ask", ask: "tool", text: "test", partial: true },
+			{ ts: 3, type: "say", say: "checkpoint_saved", text: "hash123" },
+		]
+
+		const result = findPartialAskMessage(messages, "tool")
+		expect(result).toBeDefined()
+		expect(result?.message.ts).toBe(2)
+		expect(result?.index).toBe(1)
+	})
+
+	it("should return undefined when no partial ask exists", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "ask", ask: "tool", text: "test", partial: false },
+			{ ts: 2, type: "say", say: "checkpoint_saved", text: "hash123" },
+		]
+
+		const result = findPartialAskMessage(messages, "tool")
+		expect(result).toBeUndefined()
+	})
+
+	it("should find the correct type when multiple ask types exist", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "ask", ask: "tool", text: "test", partial: true },
+			{ ts: 2, type: "ask", ask: "command", text: "test2", partial: true },
+			{ ts: 3, type: "say", say: "checkpoint_saved", text: "hash123" },
+		]
+
+		const result = findPartialAskMessage(messages, "command")
+		expect(result).toBeDefined()
+		expect(result?.message.ask).toBe("command")
+		expect(result?.message.ts).toBe(2)
+	})
+
+	it("should handle multiple checkpoints between partial start and completion", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "ask", ask: "tool", text: "test", partial: true },
+			{ ts: 2, type: "say", say: "checkpoint_saved", text: "hash1" },
+			{ ts: 3, type: "say", say: "checkpoint_saved", text: "hash2" },
+			{ ts: 4, type: "say", say: "checkpoint_saved", text: "hash3" },
+		]
+
+		const result = findPartialAskMessage(messages, "tool")
+		expect(result).toBeDefined()
+		expect(result?.message.ts).toBe(1)
+		expect(result?.index).toBe(0)
+	})
+
+	it("should find the most recent partial when multiple partials of same type exist", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "ask", ask: "tool", text: "old", partial: true },
+			{ ts: 2, type: "ask", ask: "tool", text: "new", partial: true },
+			{ ts: 3, type: "say", say: "checkpoint_saved", text: "hash123" },
+		]
+
+		const result = findPartialAskMessage(messages, "tool")
+		expect(result).toBeDefined()
+		expect(result?.message.text).toBe("new")
+		expect(result?.message.ts).toBe(2)
+	})
+})
+
+describe("findPartialSayMessage", () => {
+	it("should find the most recent partial say message", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "say", say: "text", text: "test", partial: false },
+			{ ts: 2, type: "say", say: "text", text: "test", partial: true },
+			{ ts: 3, type: "say", say: "checkpoint_saved", text: "hash123" },
+		]
+
+		const result = findPartialSayMessage(messages, "text")
+		expect(result).toBeDefined()
+		expect(result?.message.ts).toBe(2)
+		expect(result?.index).toBe(1)
+	})
+
+	it("should return undefined when no partial say exists", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "say", say: "text", text: "test", partial: false },
+			{ ts: 2, type: "say", say: "checkpoint_saved", text: "hash123" },
+		]
+
+		const result = findPartialSayMessage(messages, "text")
+		expect(result).toBeUndefined()
+	})
+
+	it("should find the correct type when multiple say types exist", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "say", say: "text", text: "test", partial: true },
+			{ ts: 2, type: "say", say: "reasoning", text: "thinking", partial: true },
+			{ ts: 3, type: "say", say: "checkpoint_saved", text: "hash123" },
+		]
+
+		const result = findPartialSayMessage(messages, "reasoning")
+		expect(result).toBeDefined()
+		expect(result?.message.say).toBe("reasoning")
+		expect(result?.message.ts).toBe(2)
+	})
+
+	it("should handle multiple checkpoints between partial start and completion", () => {
+		const messages: ClineMessage[] = [
+			{ ts: 1, type: "say", say: "text", text: "test", partial: true },
+			{ ts: 2, type: "say", say: "checkpoint_saved", text: "hash1" },
+			{ ts: 3, type: "say", say: "checkpoint_saved", text: "hash2" },
+		]
+
+		const result = findPartialSayMessage(messages, "text")
+		expect(result).toBeDefined()
+		expect(result?.message.ts).toBe(1)
+		expect(result?.index).toBe(0)
+	})
+})

src/core/task/message-utils.ts

@@ -0,0 +1,49 @@
+/**
+ * Utility functions for message handling in Task
+ * @kilocode_change - Created to fix orphaned partial ask messages bug
+ */
+
+import type { ClineMessage, ClineAsk, ClineSay } from "@roo-code/types"
+
+/**
+ * Search backwards through messages to find the most recent partial ask message
+ * of the specified type. This handles cases where non-interactive messages
+ * (like checkpoint_saved) are inserted between partial start and completion.
+ *
+ * @param messages - Array of Cline messages to search
+ * @param type - The ask type to search for
+ * @returns The partial message and its index, or undefined if not found
+ */
+export function findPartialAskMessage(
+	messages: ClineMessage[],
+	type: ClineAsk,
+): { message: ClineMessage; index: number } | undefined {
+	for (let i = messages.length - 1; i >= 0; i--) {
+		const msg = messages[i]
+		if (msg.type === "ask" && msg.ask === type && msg.partial === true) {
+			return { message: msg, index: i }
+		}
+	}
+	return undefined
+}
+
+/**
+ * Search backwards through messages to find the most recent partial say message
+ * of the specified type. Similar to findPartialAskMessage but for say messages.
+ *
+ * @param messages - Array of Cline messages to search
+ * @param type - The say type to search for
+ * @returns The partial message and its index, or undefined if not found
+ */
+export function findPartialSayMessage(
+	messages: ClineMessage[],
+	type: ClineSay,
+): { message: ClineMessage; index: number } | undefined {
+	for (let i = messages.length - 1; i >= 0; i--) {
+		const msg = messages[i]
+		if (msg.type === "say" && msg.say === type && msg.partial === true) {
+			return { message: msg, index: i }
+		}
+	}
+	return undefined
+}