Merge pull request #3278 from Kilo-Org/catrielmuller/cli-approval-race-condition

CLI - Fix approval race condition

.changeset/green-teams-pull.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+Improved stability of the approval menu, preventing it from showing when you don't expect it

cli/src/state/atoms/__tests__/approval-delay.test.ts

@@ -0,0 +1,89 @@
+/**
+ * Tests for approval delay functionality
+ */
+
+import { describe, it, expect, afterEach } from "vitest"
+import { createStore } from "jotai"
+import {
+	setPendingApprovalAtom,
+	isApprovalPendingAtom,
+	approvalSetTimestampAtom,
+	clearPendingApprovalAtom,
+} from "../approval.js"
+import type { ExtensionChatMessage } from "../../../types/messages.js"
+
+describe("approval delay", () => {
+	afterEach(() => {
+		// Cleanup if needed
+	})
+
+	it("should set timestamp when pending approval is set", () => {
+		const store = createStore()
+		const message: ExtensionChatMessage = {
+			ts: Date.now(),
+			type: "ask",
+			ask: "tool",
+			text: '{"tool":"readFile"}',
+		}
+
+		store.set(setPendingApprovalAtom, message)
+
+		const timestamp = store.get(approvalSetTimestampAtom)
+		expect(timestamp).not.toBeNull()
+		expect(typeof timestamp).toBe("number")
+	})
+
+	it("should clear timestamp when pending approval is cleared", () => {
+		const store = createStore()
+		const message: ExtensionChatMessage = {
+			ts: Date.now(),
+			type: "ask",
+			ask: "tool",
+			text: '{"tool":"readFile"}',
+		}
+
+		store.set(setPendingApprovalAtom, message)
+		expect(store.get(approvalSetTimestampAtom)).not.toBeNull()
+
+		store.set(clearPendingApprovalAtom)
+		expect(store.get(approvalSetTimestampAtom)).toBeNull()
+	})
+
+	it("isApprovalPendingAtom should return true immediately", () => {
+		const store = createStore()
+		const message: ExtensionChatMessage = {
+			ts: Date.now(),
+			type: "ask",
+			ask: "tool",
+			text: '{"tool":"readFile"}',
+		}
+
+		store.set(setPendingApprovalAtom, message)
+
+		// Immediate check should return true
+		expect(store.get(isApprovalPendingAtom)).toBe(true)
+	})
+
+	it("should handle timestamp and approval state together", () => {
+		const store = createStore()
+		const message: ExtensionChatMessage = {
+			ts: Date.now(),
+			type: "ask",
+			ask: "tool",
+			text: '{"tool":"readFile"}',
+		}
+
+		store.set(setPendingApprovalAtom, message)
+
+		// Both timestamp and approval should be set
+		expect(store.get(approvalSetTimestampAtom)).not.toBeNull()
+		expect(store.get(isApprovalPendingAtom)).toBe(true)
+
+		// Clear approval
+		store.set(clearPendingApprovalAtom)
+
+		// Both should be cleared
+		expect(store.get(approvalSetTimestampAtom)).toBeNull()
+		expect(store.get(isApprovalPendingAtom)).toBe(false)
+	})
+})

cli/src/state/atoms/approval.ts

@@ -46,6 +46,11 @@ export const approvalProcessingAtom = atom<ApprovalProcessingState>({
  */
 export const selectedApprovalIndexAtom = selectedIndexAtom
 
+/**
+ * Atom to track when approval was set (for delay logic)
+ */
+export const approvalSetTimestampAtom = atom<number | null>(null)
+
 /**
  * Derived atom to check if there's a pending approval
  */
@@ -226,6 +231,11 @@ export const setPendingApprovalAtom = atom(null, (get, set, message: ExtensionCh
 	const messageToSet = message ? { ...message } : null
 	set(pendingApprovalAtom, messageToSet)
 
+	// Set timestamp for delay tracking when setting a new message
+	if (isNewMessage && messageToSet !== null) {
+		set(approvalSetTimestampAtom, Date.now())
+	}
+
 	// Reset selection if this is a new message (different timestamp)
 	if (isNewMessage) {
 		set(selectedIndexAtom, 0)
@@ -241,6 +251,7 @@ export const clearPendingApprovalAtom = atom(null, (get, set) => {
 	const processing = get(approvalProcessingAtom)
 
 	set(pendingApprovalAtom, null)
+	set(approvalSetTimestampAtom, null)
 
 	// Also clear processing state if it matches
 	if (processing.isProcessing && processing.processingTs === current?.ts) {
@@ -284,6 +295,7 @@ export const startApprovalProcessingAtom = atom(null, (get, set, operation: "app
  */
 export const completeApprovalProcessingAtom = atom(null, (get, set) => {
 	set(pendingApprovalAtom, null)
+	set(approvalSetTimestampAtom, null)
 	set(selectedIndexAtom, 0)
 	set(approvalProcessingAtom, { isProcessing: false })
 })

cli/src/state/atoms/ui.ts

@@ -246,7 +246,7 @@ export const lastMessageAtom = atom<CliMessage | null>((get) => {
 
 /**
  * Derived atom to get the last ask message from extension messages
- * Returns the most recent unanswered ask message that requires user approval, or null if none exists
+ * Returns the most recent ask message that requires user approval, or null if none exists
  */
 export const lastAskMessageAtom = atom<ExtensionChatMessage | null>((get) => {
 	const messages = get(chatMessagesAtom)
@@ -254,21 +254,17 @@ export const lastAskMessageAtom = atom<ExtensionChatMessage | null>((get) => {
 	// Ask types that require user approval
 	const approvalAskTypes = ["tool", "command", "browser_action_launch", "use_mcp_server"]
 
-	// Find the last unanswered ask message that requires approval
-	for (let i = messages.length - 1; i >= 0; i--) {
-		const msg = messages[i]
-		if (
-			msg &&
-			msg.type === "ask" &&
-			!msg.isAnswered &&
-			msg.ask &&
-			approvalAskTypes.includes(msg.ask) &&
-			!msg.partial
-		) {
-			return msg
-		}
+	const lastMessage = messages[messages.length - 1]
+	if (
+		lastMessage &&
+		lastMessage.type === "ask" &&
+		!lastMessage.isAnswered &&
+		lastMessage.ask &&
+		approvalAskTypes.includes(lastMessage.ask) &&
+		!lastMessage.partial
+	) {
+		return lastMessage
 	}
-
 	return null
 })
 

cli/src/state/hooks/useApprovalHandler.ts

@@ -1,5 +1,5 @@
 import { useAtomValue, useSetAtom, useStore } from "jotai"
-import { useCallback, useEffect } from "react"
+import { useCallback, useEffect, useState } from "react"
 import {
 	pendingApprovalAtom,
 	approvalOptionsAtom,
@@ -8,6 +8,7 @@ import {
 	selectNextApprovalAtom,
 	selectPreviousApprovalAtom,
 	isApprovalPendingAtom,
+	approvalSetTimestampAtom,
 	startApprovalProcessingAtom,
 	completeApprovalProcessingAtom,
 	approvalProcessingAtom,
@@ -23,6 +24,8 @@ import type { ExtensionChatMessage } from "../../types/messages.js"
 import { logs } from "../../services/logs.js"
 import { useApprovalTelemetry } from "./useApprovalTelemetry.js"
 
+const APPROVAL_MENU_DELAY_MS = 500
+
 /**
  * Options for useApprovalHandler hook
  */
@@ -88,7 +91,14 @@ export function useApprovalHandler(): UseApprovalHandlerReturn {
 	const approvalOptions = useAtomValue(approvalOptionsAtom)
 	const selectedIndex = useAtomValue(selectedApprovalIndexAtom)
 	const selectedOption = useAtomValue(selectedApprovalOptionAtom)
-	const isApprovalPending = useAtomValue(isApprovalPendingAtom)
+	const approvalSetTimestamp = useAtomValue(approvalSetTimestampAtom)
+	const isApprovalPendingImmediate = useAtomValue(isApprovalPendingAtom)
+
+	// Use state to manage delayed visibility
+	const [showApprovalMenu, setShowApprovalMenu] = useState(false)
+
+	// Compute if approval is pending with delay
+	const isApprovalPending = isApprovalPendingImmediate && showApprovalMenu
 
 	const selectNext = useSetAtom(selectNextApprovalAtom)
 	const selectPrevious = useSetAtom(selectPreviousApprovalAtom)
@@ -276,6 +286,28 @@ export function useApprovalHandler(): UseApprovalHandlerReturn {
 		setExecuteSelectedCallback(() => executeSelected)
 	}, [approve, reject, executeSelected, setApproveCallback, setRejectCallback, setExecuteSelectedCallback])
 
+	// Manage delayed visibility of approval menu
+	useEffect(() => {
+		if (approvalSetTimestamp === null) {
+			// No pending approval, hide menu
+			setShowApprovalMenu(false)
+			return
+		}
+
+		// Calculate remaining time for delay
+		const elapsed = Date.now() - approvalSetTimestamp
+		const remaining = Math.max(0, APPROVAL_MENU_DELAY_MS - elapsed)
+
+		// Set timeout to show menu after delay
+		const timeoutId = setTimeout(() => {
+			setShowApprovalMenu(true)
+		}, remaining)
+
+		return () => {
+			clearTimeout(timeoutId)
+		}
+	}, [approvalSetTimestamp, isApprovalPendingImmediate])
+
 	return {
 		pendingApproval,
 		approvalOptions,