refactor: Message guard

Author: catrielmuller

src/core/kilocode/task/message-utils.ts

@@ -1,10 +1,62 @@
 /**
  * Utility functions for message handling in Task
- * @kilocode_change - Created to fix orphaned partial ask messages bug
+ * kilocode_change - Created to fix orphaned partial ask messages bug
  */
 
+import pWaitFor from "p-wait-for"
 import type { ClineMessage, ClineAsk, ClineSay } from "@roo-code/types"
 
+const GUARD_TIMEOUT = 30 * 1000 // 30 seconds
+const GUARD_INTERVAL = 50 // 50 milliseconds
+
+/**
+ * Message insertion guard to prevent concurrent message insertions.
+ * This prevents race conditions where checkpoint_saved messages can arrive
+ * during partial message updates, breaking message history order.
+ */
+export class MessageInsertionGuard {
+	private isInserting = false
+	private readonly timeout: number
+	private readonly interval: number
+
+	constructor(timeout = GUARD_TIMEOUT, interval = GUARD_INTERVAL) {
+		this.timeout = timeout
+		this.interval = interval
+	}
+
+	/**
+	 * Wait for any in-flight message insertions to complete before proceeding.
+	 * This ensures messages are inserted sequentially and prevents race conditions.
+	 */
+	async waitForClearance(): Promise<void> {
+		await pWaitFor(() => !this.isInserting, {
+			interval: this.interval,
+			timeout: this.timeout,
+		})
+	}
+
+	/**
+	 * Acquire the insertion lock. Must be released with release() after insertion completes.
+	 */
+	acquire(): void {
+		this.isInserting = true
+	}
+
+	/**
+	 * Release the insertion lock, allowing other insertions to proceed.
+	 */
+	release(): void {
+		this.isInserting = false
+	}
+
+	/**
+	 * Check if a message insertion is currently in progress.
+	 */
+	isLocked(): boolean {
+		return this.isInserting
+	}
+}
+
 /**
  * Search backwards through messages to find the most recent partial ask message
  * of the specified type. This handles cases where non-interactive messages

src/core/task/Task.ts

@@ -119,7 +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 "../kilocode/task/message-utils" // kilocode_change
+import { findPartialAskMessage, findPartialSayMessage, MessageInsertionGuard } from "../kilocode/task/message-utils" // kilocode_change
 
 import { AutoApprovalHandler } from "./AutoApprovalHandler"
 import { isAnyRecognizedKiloCodeError, isPaymentRequiredError } from "../../shared/kilocode/errorUtils"
@@ -221,6 +221,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 	private readonly globalStoragePath: string
 	abort: boolean = false
 
+	// kilocode_change start: Message insertion guard to prevent race conditions with checkpoint messages
+	private readonly messageInsertionGuard = new MessageInsertionGuard()
+	// kilocode_change end
+
 	// TaskStatus
 	idleAsk?: ClineMessage
 	resumableAsk?: ClineMessage
@@ -639,24 +643,32 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		return readTaskMessages({ taskId: this.taskId, globalStoragePath: this.globalStoragePath })
 	}
 
+	// kilocode_change start: Guard against concurrent message insertions to prevent
 	private async addToClineMessages(message: ClineMessage) {
-		this.clineMessages.push(message)
-		const provider = this.providerRef.deref()
-		await provider?.postStateToWebview()
-		this.emit(RooCodeEventName.Message, { action: "created", message })
-		await this.saveClineMessages()
+		await this.messageInsertionGuard.waitForClearance()
+		this.messageInsertionGuard.acquire()
 
-		// kilocode_change start: no cloud service
-		// const shouldCaptureMessage = message.partial !== true && CloudService.isEnabled()
+		try {
+			this.clineMessages.push(message)
+			const provider = this.providerRef.deref()
+			await provider?.postStateToWebview()
+			this.emit(RooCodeEventName.Message, { action: "created", message })
+			await this.saveClineMessages()
 
-		// if (shouldCaptureMessage) {
-		// 	CloudService.instance.captureEvent({
-		// 		event: TelemetryEventName.TASK_MESSAGE,
-		// 		properties: { taskId: this.taskId, message },
-		// 	})
-		// }
-		// kilocode_change end
+			// kilocode_change start: no cloud service
+			// const shouldCaptureMessage = message.partial !== true && CloudService.isEnabled()
+			// if (shouldCaptureMessage) {
+			// 	CloudService.instance.captureEvent({
+			// 		event: TelemetryEventName.TASK_MESSAGE,
+			// 		properties: { taskId: this.taskId, message },
+			// 	})
+			// }
+			// kilocode_change end
+		} finally {
+			this.messageInsertionGuard.release()
+		}
 	}
+	// kilocode_change end
 
 	public async overwriteClineMessages(newMessages: ClineMessage[]) {
 		this.clineMessages = newMessages