feat(followup): add timeout management for follow-up auto-approval (#654)

Author: mini2s

.kilocode/mcp.json

@@ -38,6 +38,7 @@ export type CheckAutoApprovalResult =
 	| {
 			decision: "timeout"
 			timeout: number
+			askType?: string
 			fn: () => { askResponse: ClineAskResponse; text?: string; images?: string[] }
 	  }
 
@@ -73,6 +74,7 @@ export async function checkAutoApproval({
 					return {
 						decision: "timeout",
 						timeout: state.followupAutoApproveTimeoutMs,
+						askType: ask,
 						fn: () => ({ askResponse: "messageResponse", text: suggestion.answer }),
 					}
 				} else {

src/core/auto-approval/index.ts

@@ -161,6 +161,9 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 
 	todoList?: TodoItem[]
 
+	private readonly timeoutMap = new Map<number, NodeJS.Timeout>()
+	private nextTimeoutId = 0
+
 	readonly rootTask: Task | undefined = undefined
 	readonly parentTask: Task | undefined = undefined
 	readonly taskNumber: number
@@ -855,8 +858,6 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 			await this.addToClineMessages({ ts: askTs, type: "ask", ask: type, text, isProtected })
 		}
 
-		let timeouts: NodeJS.Timeout[] = []
-
 		// Automatically approve if the ask according to the user's settings.
 		const provider = this.providerRef.deref()
 		const state = provider ? await provider.getState() : undefined
@@ -867,12 +868,22 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		} else if (approval.decision === "deny") {
 			this.denyAsk()
 		} else if (approval.decision === "timeout") {
-			timeouts.push(
-				setTimeout(() => {
-					const { askResponse, text, images } = approval.fn()
-					this.handleWebviewAskResponse(askResponse, text, images)
-				}, approval.timeout),
-			)
+			const timeoutId = this.nextTimeoutId++
+
+			const timer = setTimeout(() => {
+				const { askResponse, text, images } = approval.fn()
+				this.handleWebviewAskResponse(askResponse, text, images)
+				this.timeoutMap.delete(timeoutId)
+			}, approval.timeout)
+
+			this.timeoutMap.set(timeoutId, timer)
+
+			if (approval.askType === "followup") {
+				provider?.postMessageToWebview({
+					type: "zgsmFollowupClearTimeout",
+					value: timeoutId,
+				})
+			}
 		}
 
 		// The state is mutable if the message is complete and the task will
@@ -891,47 +902,53 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 			const statusMutationTimeout = 2_000
 
 			if (isInteractiveAsk(type)) {
-				timeouts.push(
-					setTimeout(() => {
-						const message = this.findMessageByTimestamp(askTs)
+				const timeoutId = this.nextTimeoutId++
+				const timer = setTimeout(() => {
+					const message = this.findMessageByTimestamp(askTs)
+
+					if (message) {
+						this.interactiveAsk = message
+						this.emit(RooCodeEventName.TaskInteractive, this.taskId)
+						provider?.postMessageToWebview({ type: "interactionRequired" })
+					}
+					this.timeoutMap.delete(timeoutId)
+				}, statusMutationTimeout)
 
-						if (message) {
-							this.interactiveAsk = message
-							this.emit(RooCodeEventName.TaskInteractive, this.taskId)
-							provider?.postMessageToWebview({ type: "interactionRequired" })
-						}
-					}, statusMutationTimeout),
-				)
+				this.timeoutMap.set(timeoutId, timer)
 			} else if (isResumableAsk(type)) {
-				timeouts.push(
-					setTimeout(() => {
-						const message = this.findMessageByTimestamp(askTs)
+				const timeoutId = this.nextTimeoutId++
+				const timer = setTimeout(() => {
+					const message = this.findMessageByTimestamp(askTs)
 
-						if (message) {
-							this.resumableAsk = message
-							this.emit(RooCodeEventName.TaskResumable, this.taskId)
-						}
-					}, statusMutationTimeout),
-				)
+					if (message) {
+						this.resumableAsk = message
+						this.emit(RooCodeEventName.TaskResumable, this.taskId)
+					}
+					this.timeoutMap.delete(timeoutId)
+				}, statusMutationTimeout)
+
+				this.timeoutMap.set(timeoutId, timer)
 			} else if (isIdleAsk(type)) {
-				timeouts.push(
-					setTimeout(() => {
-						const message = this.findMessageByTimestamp(askTs)
+				const timeoutId = this.nextTimeoutId++
+				const timer = setTimeout(() => {
+					const message = this.findMessageByTimestamp(askTs)
 
-						if (message) {
-							this.idleAsk = message
-							this.emit(RooCodeEventName.TaskIdle, this.taskId)
-						}
-					}, statusMutationTimeout),
-				)
+					if (message) {
+						this.idleAsk = message
+						this.emit(RooCodeEventName.TaskIdle, this.taskId)
+					}
+					this.timeoutMap.delete(timeoutId)
+				}, statusMutationTimeout)
+
+				this.timeoutMap.set(timeoutId, timer)
 			}
 		} else if (isMessageQueued) {
 			console.log(`Task#ask: will process message queue -> type: ${type}`)
 
 			const message = this.messageQueueService.dequeueMessage()
 
 			if (message) {
-				// Check if this is a tool approval ask that needs to be handled.
+				// Check if this is a tool approval ask Pthat needs to be handled.
 				if (
 					type === "tool" ||
 					type === "command" ||
@@ -971,7 +988,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		this.askResponseImages = undefined
 
 		// Cancel the timeouts if they are still running.
-		timeouts.forEach((timeout) => clearTimeout(timeout))
+		this.clearAllAutoApprovalTimeouts()
 
 		// Switch back to an active state.
 		if (this.idleAsk || this.resumableAsk || this.interactiveAsk) {
@@ -985,6 +1002,23 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		return result
 	}
 
+	public clearAutoApprovalTimeout(timeoutId: number): boolean {
+		const timer = this.timeoutMap.get(timeoutId)
+		if (timer) {
+			clearTimeout(timer)
+			this.timeoutMap.delete(timeoutId)
+			return true
+		}
+		return false
+	}
+
+	public clearAllAutoApprovalTimeouts(): void {
+		for (const [timeoutId, timer] of this.timeoutMap) {
+			clearTimeout(timer)
+		}
+		this.timeoutMap.clear()
+	}
+
 	public setMessageResponse(text: string, images?: string[]) {
 		this.handleWebviewAskResponse("messageResponse", text, images)
 	}

src/core/task/Task.ts

@@ -3116,6 +3116,16 @@ export const webviewMessageHandler = async (
 			}
 			break
 		}
+		case "zgsmFollowupClearTimeout": {
+			const currentTask = provider?.getCurrentTask()
+			if (currentTask && typeof message.value === "number") {
+				const cleared = currentTask.clearAutoApprovalTimeout(message.value)
+				if (!cleared) {
+					provider?.log(`Failed to clear timeout with ID: ${message.value}`, "info")
+				}
+			}
+			break
+		}
 		case "zgsmRebuildCodebaseIndex": {
 			try {
 				const { apiConfiguration } = await provider.getState()

src/core/webview/webviewMessageHandler.ts

@@ -89,6 +89,7 @@ export interface ExtensionMessage {
 		| "zgsmInviteCode"
 		| "zgsmNotices"
 		| "settingsUpdated"
+		| "zgsmFollowupClearTimeout"
 		// zgsm
 		| "ollamaModels"
 		| "lmStudioModels"

src/shared/ExtensionMessage.ts

@@ -159,6 +159,7 @@ export interface WebviewMessage {
 		| "fetchZgsmQuotaInfo"
 		| "fetchZgsmInviteCode"
 		| "fixCodebase"
+		| "zgsmFollowupClearTimeout"
 		// zgsm
 		| "shareTaskSuccess"
 		| "exportMode"

src/shared/WebviewMessage.ts

@@ -183,8 +183,9 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 			ttl: 1000 * 60 * 5,
 		}),
 	)
-	const autoApproveTimeoutRef = useRef<NodeJS.Timeout | null>(null)
-	const userRespondedRef = useRef<boolean>(false)
+	// const autoApproveTimeoutRef = useRef<NodeJS.Timeout | null>(null)
+	const followUpAutoApproveTimeoutRef = useRef<number | undefined>()
+	// const userRespondedRef = useRef<boolean>(false)
 	const [currentFollowUpTs, setCurrentFollowUpTs] = useState<number | null>(null)
 
 	const clineAskRef = useRef(clineAsk)
@@ -265,7 +266,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 			switch (lastMessage.type) {
 				case "ask":
 					// Reset user response flag when a new ask arrives to allow auto-approval
-					userRespondedRef.current = false
+					// userRespondedRef.current = false
 					const isPartial = lastMessage.partial === true
 					switch (lastMessage.ask) {
 						case "api_req_failed":
@@ -436,12 +437,12 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 		// Note: sendingDisabled is not reset here as it's managed by message effects
 
 		// Clear any pending auto-approval timeout from previous task
-		if (autoApproveTimeoutRef.current) {
-			clearTimeout(autoApproveTimeoutRef.current)
-			autoApproveTimeoutRef.current = null
-		}
+		// if (autoApproveTimeoutRef.current) {
+		// 	clearTimeout(autoApproveTimeoutRef.current)
+		// 	autoApproveTimeoutRef.current = null
+		// }
 		// Reset user response flag for new task
-		userRespondedRef.current = false
+		// userRespondedRef.current = false
 	}, [task?.ts])
 
 	useEffect(() => {
@@ -530,12 +531,12 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 
 	const handleChatReset = useCallback(() => {
 		// Clear any pending auto-approval timeout
-		if (autoApproveTimeoutRef.current) {
-			clearTimeout(autoApproveTimeoutRef.current)
-			autoApproveTimeoutRef.current = null
-		}
+		// if (autoApproveTimeoutRef.current) {
+		// 	clearTimeout(autoApproveTimeoutRef.current)
+		// 	autoApproveTimeoutRef.current = null
+		// }
 		// Reset user response flag for new message
-		userRespondedRef.current = false
+		// userRespondedRef.current = false
 
 		// Only reset message-specific state, preserving mode.
 		setInputValue("")
@@ -578,7 +579,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 				}
 
 				// Mark that user has responded - this prevents any pending auto-approvals.
-				userRespondedRef.current = true
+				// userRespondedRef.current = true
 
 				if (messagesRef.current.length === 0) {
 					vscode.postMessage({ type: "newTask", text, images, values: { chatType } })
@@ -664,7 +665,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	const handlePrimaryButtonClick = useCallback(
 		(text?: string, images?: string[]) => {
 			// Mark that user has responded
-			userRespondedRef.current = true
+			// userRespondedRef.current = true
 
 			const trimmedInput = text?.trim()
 
@@ -714,7 +715,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	const handleSecondaryButtonClick = useCallback(
 		(text?: string, images?: string[]) => {
 			// Mark that user has responded
-			userRespondedRef.current = true
+			// userRespondedRef.current = true
 
 			const trimmedInput = text?.trim()
 
@@ -772,6 +773,10 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 			const message: ExtensionMessage = e.data
 
 			switch (message.type) {
+				case "zgsmFollowupClearTimeout": {
+					followUpAutoApproveTimeoutRef.current = message.value
+					break
+				}
 				case "action":
 					switch (message.action!) {
 						case "didBecomeVisible":
@@ -1223,9 +1228,9 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	const handleSuggestionClickInRow = useCallback(
 		(suggestion: SuggestionItem, event?: React.MouseEvent) => {
 			// Mark that user has responded if this is a manual click (not auto-approval)
-			if (event) {
-				userRespondedRef.current = true
-			}
+			// if (event) {
+			// 	userRespondedRef.current = true
+			// }
 
 			// Mark the current follow-up question as answered when a suggestion is clicked
 			if (clineAsk === "followup" && !event?.shiftKey) {
@@ -1267,7 +1272,14 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	// Handler for when FollowUpSuggest component unmounts
 	const handleFollowUpUnmount = useCallback(() => {
 		// Mark that user has responded
-		userRespondedRef.current = true
+		// userRespondedRef.current = true
+		if (Number.isInteger(followUpAutoApproveTimeoutRef.current)) {
+			vscode.postMessage({
+				type: "zgsmFollowupClearTimeout",
+				value: followUpAutoApproveTimeoutRef.current,
+			})
+		}
+		followUpAutoApproveTimeoutRef.current = undefined
 	}, [])
 	const shouldHighlight = useCallback(
 		(messageOrGroup?: ClineMessage, searchResults: SearchResult[] = [], showSearch?: boolean) => {