Revert "Fix: Orphaned Partial Ask Messages Bug"

Revert "Fix: Orphaned Partial Ask Messages Bug"

Author: catrielmuller

.changeset/silver-garlics-listen.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Revert orphaned partial ask messages fix

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

@@ -98,6 +98,50 @@ 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,7 +197,10 @@ export const updateExtensionStateAtom = atom(null, (get, set, state: ExtensionSt
 		const incomingMessages = state.clineMessages || state.chatMessages || []
 
 		// Reconcile with current messages to preserve streaming state
-		const reconciledMessages = reconcileMessages(currentMessages, incomingMessages, versionMap, streamingSet)
+		let reconciledMessages = reconcileMessages(currentMessages, incomingMessages, versionMap, streamingSet)
+
+		// Auto-complete orphaned partial ask messages (CLI-only workaround for extension bug)
+		reconciledMessages = autoCompleteOrphanedPartialAsks(reconciledMessages)
 
 		set(chatMessagesAtom, reconciledMessages)
 
@@ -456,6 +459,57 @@ 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/kilocode/task/message-utils.ts

@@ -119,7 +119,6 @@ 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 "../kilocode/task/message-utils" // kilocode_change
 
 import { AutoApprovalHandler } from "./AutoApprovalHandler"
 import { isAnyRecognizedKiloCodeError, isPaymentRequiredError } from "../../shared/kilocode/errorUtils"
@@ -769,13 +768,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		let askTs: number
 
 		if (partial !== undefined) {
-			// 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
-			const isUpdatingPreviousPartial = lastMessage !== undefined
-			// kilocode_change end
+			const lastMessage = this.clineMessages.at(-1)
+
+			const isUpdatingPreviousPartial =
+				lastMessage && lastMessage.partial && lastMessage.type === "ask" && lastMessage.ask === type
 
 			if (partial) {
 				if (isUpdatingPreviousPartial) {
@@ -1161,12 +1157,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		}
 
 		if (partial !== undefined) {
-			// 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
-			const isUpdatingPreviousPartial = lastMessage !== undefined
-			// kilocode_change end
+			const lastMessage = this.clineMessages.at(-1)
+
+			const isUpdatingPreviousPartial =
+				lastMessage && lastMessage.partial && lastMessage.type === "say" && lastMessage.say === type
 
 			if (partial) {
 				if (isUpdatingPreviousPartial) {