Fix CLI follow-up routing after completion asks (#11844)

Fix stdin follow-up routing for completion asks in CLI stream mode

Author: cte

apps/cli/scripts/integration/cases/followup-completion-ask-response.ts

@@ -0,0 +1,143 @@
+import { runStreamCase, StreamEvent } from "../lib/stream-harness"
+
+const START_PROMPT = 'Answer this question and finish: What is 1+1? Reply with only "2", then complete the task.'
+const FOLLOWUP_PROMPT = 'Different question now: what is 3+3? Reply with only "6".'
+
+async function main() {
+	const startRequestId = `start-${Date.now()}`
+	const followupRequestId = `message-${Date.now()}`
+	const shutdownRequestId = `shutdown-${Date.now()}`
+
+	let initSeen = false
+	let sentFollowup = false
+	let sentShutdown = false
+	let startAckCount = 0
+	let sawStartControlAfterFollowup = false
+	let followupDoneCode: string | undefined
+	let sawFollowupUserTurn = false
+	let sawMisroutedToolResult = false
+	let followupResult = ""
+
+	await runStreamCase({
+		onEvent(event: StreamEvent, context) {
+			if (event.type === "system" && event.subtype === "init" && !initSeen) {
+				initSeen = true
+				context.sendCommand({
+					command: "start",
+					requestId: startRequestId,
+					prompt: START_PROMPT,
+				})
+				return
+			}
+
+			if (event.type === "control" && event.subtype === "error") {
+				throw new Error(
+					`received control error for requestId=${event.requestId ?? "unknown"} command=${event.command ?? "unknown"} code=${event.code ?? "unknown"} content=${event.content ?? ""}`,
+				)
+			}
+
+			if (event.type === "control" && event.command === "start" && event.subtype === "ack") {
+				startAckCount += 1
+				if (sentFollowup) {
+					sawStartControlAfterFollowup = true
+				}
+				return
+			}
+
+			if (
+				event.type === "control" &&
+				event.command === "message" &&
+				event.subtype === "done" &&
+				event.requestId === followupRequestId
+			) {
+				followupDoneCode = event.code
+				return
+			}
+
+			if (
+				event.type === "tool_result" &&
+				event.requestId === followupRequestId &&
+				typeof event.content === "string" &&
+				event.content.includes("<user_message>")
+			) {
+				sawMisroutedToolResult = true
+				return
+			}
+
+			if (event.type === "user" && event.requestId === followupRequestId) {
+				sawFollowupUserTurn = typeof event.content === "string" && event.content.includes("3+3")
+				return
+			}
+
+			if (event.type === "result" && event.done === true && event.requestId === startRequestId && !sentFollowup) {
+				context.sendCommand({
+					command: "message",
+					requestId: followupRequestId,
+					prompt: FOLLOWUP_PROMPT,
+				})
+				sentFollowup = true
+				return
+			}
+
+			if (event.type !== "result" || event.done !== true || event.requestId !== followupRequestId) {
+				return
+			}
+
+			followupResult = event.content ?? ""
+			if (followupResult.trim().length === 0) {
+				throw new Error("follow-up produced an empty result")
+			}
+
+			if (followupDoneCode !== "responded") {
+				throw new Error(
+					`follow-up message was not routed as ask response; code="${followupDoneCode ?? "none"}"`,
+				)
+			}
+
+			if (sawMisroutedToolResult) {
+				throw new Error("follow-up message was misrouted into tool_result (<user_message>), old bug reproduced")
+			}
+
+			if (!sawFollowupUserTurn) {
+				throw new Error("follow-up did not appear as a normal user turn in stream output")
+			}
+
+			if (sawStartControlAfterFollowup) {
+				throw new Error("unexpected start control event after follow-up; message should not trigger a new task")
+			}
+
+			if (startAckCount !== 1) {
+				throw new Error(`expected exactly one start ack event, saw ${startAckCount}`)
+			}
+
+			console.log(`[PASS] follow-up control code: "${followupDoneCode}"`)
+			console.log(`[PASS] follow-up user turn observed: ${sawFollowupUserTurn}`)
+			console.log(`[PASS] follow-up result: "${followupResult}"`)
+
+			if (!sentShutdown) {
+				context.sendCommand({
+					command: "shutdown",
+					requestId: shutdownRequestId,
+				})
+				sentShutdown = true
+			}
+		},
+		onTimeoutMessage() {
+			return [
+				"timed out waiting for completion ask-response follow-up validation",
+				`initSeen=${initSeen}`,
+				`sentFollowup=${sentFollowup}`,
+				`startAckCount=${startAckCount}`,
+				`followupDoneCode=${followupDoneCode ?? "none"}`,
+				`sawFollowupUserTurn=${sawFollowupUserTurn}`,
+				`sawMisroutedToolResult=${sawMisroutedToolResult}`,
+				`haveFollowupResult=${Boolean(followupResult)}`,
+			].join(" ")
+		},
+	})
+}
+
+main().catch((error) => {
+	console.error(`[FAIL] ${error instanceof Error ? error.message : String(error)}`)
+	process.exit(1)
+})

apps/cli/scripts/integration/cases/followup-during-streaming.ts

@@ -16,11 +16,9 @@ function looksLikeAttemptCompletionToolUse(event: StreamEvent): boolean {
 	return content.includes('"tool":"attempt_completion"') || content.includes('"name":"attempt_completion"')
 }
 
-function validateFollowupAnswer(text: string): void {
-	const normalized = text.toLowerCase()
-	const hasSix = /\b6\b/.test(normalized) || normalized.includes("six")
-	if (!hasSix) {
-		throw new Error(`follow-up result did not answer follow-up prompt; result="${text}"`)
+function validateFollowupResult(text: string): void {
+	if (text.trim().length === 0) {
+		throw new Error("follow-up produced an empty result")
 	}
 }
 
@@ -117,7 +115,7 @@ async function main() {
 			}
 
 			followupResult = event.content ?? ""
-			validateFollowupAnswer(followupResult)
+			validateFollowupResult(followupResult)
 
 			if (sawMisroutedToolResult) {
 				throw new Error("follow-up message was misrouted into tool_result (<user_message>), old bug reproduced")

apps/cli/src/commands/cli/__tests__/parse-stdin-command.test.ts

@@ -1,4 +1,4 @@
-import { parseStdinStreamCommand } from "../stdin-stream.js"
+import { parseStdinStreamCommand, shouldSendMessageAsAskResponse } from "../stdin-stream.js"
 
 describe("parseStdinStreamCommand", () => {
 	describe("valid commands", () => {
@@ -162,3 +162,22 @@ describe("parseStdinStreamCommand", () => {
 		})
 	})
 })
+
+describe("shouldSendMessageAsAskResponse", () => {
+	it("routes completion_result asks as ask responses", () => {
+		expect(shouldSendMessageAsAskResponse(true, "completion_result")).toBe(true)
+	})
+
+	it("routes followup asks as ask responses", () => {
+		expect(shouldSendMessageAsAskResponse(true, "followup")).toBe(true)
+	})
+
+	it("does not route when not waiting for input", () => {
+		expect(shouldSendMessageAsAskResponse(false, "completion_result")).toBe(false)
+	})
+
+	it("does not route unknown asks", () => {
+		expect(shouldSendMessageAsAskResponse(true, "unknown")).toBe(false)
+		expect(shouldSendMessageAsAskResponse(true, undefined)).toBe(false)
+	})
+})

apps/cli/src/commands/cli/stdin-stream.ts

@@ -214,6 +214,20 @@ const STDIN_EOF_RESUME_WAIT_TIMEOUT_MS = 2_000
 const STDIN_EOF_POLL_INTERVAL_MS = 100
 const STDIN_EOF_IDLE_ASKS = new Set(["completion_result", "resume_completed_task"])
 const STDIN_EOF_IDLE_STABLE_POLLS = 2
+const MESSAGE_AS_ASK_RESPONSE_ASKS = new Set([
+	"followup",
+	"tool",
+	"command",
+	"use_mcp_server",
+	"completion_result",
+	"resume_task",
+	"resume_completed_task",
+	"mistake_limit_reached",
+])
+
+export function shouldSendMessageAsAskResponse(waitingForInput: boolean, currentAsk: string | undefined): boolean {
+	return waitingForInput && typeof currentAsk === "string" && MESSAGE_AS_ASK_RESPONSE_ASKS.has(currentAsk)
+}
 
 function isResumableState(host: ExtensionHost): boolean {
 	const agentState = host.client.getAgentState()
@@ -690,6 +704,8 @@ export async function runStdinStreamMode({ host, jsonEmitter, setStreamRequestId
 					}
 
 					const wasResumable = isResumableState(host)
+					const currentAsk = host.client.getCurrentAsk()
+					const shouldSendAsAskResponse = shouldSendMessageAsAskResponse(host.isWaitingForInput(), currentAsk)
 
 					if (!host.client.hasActiveTask()) {
 						jsonEmitter.emitControl({
@@ -715,6 +731,29 @@ export async function runStdinStreamMode({ host, jsonEmitter, setStreamRequestId
 						success: true,
 					})
 
+					if (shouldSendAsAskResponse) {
+						// Match webview behavior: if there is an active ask, route message directly as an ask response.
+						host.sendToExtension({
+							type: "askResponse",
+							askResponse: "messageResponse",
+							text: stdinCommand.prompt,
+							images: stdinCommand.images,
+						})
+
+						setStreamRequestId(stdinCommand.requestId)
+						jsonEmitter.emitControl({
+							subtype: "done",
+							requestId: stdinCommand.requestId,
+							command: "message",
+							taskId: latestTaskId,
+							content: "message sent to current ask",
+							code: "responded",
+							success: true,
+						})
+						awaitingPostCancelRecovery = false
+						break
+					}
+
 					host.sendToExtension({
 						type: "queueMessage",
 						text: stdinCommand.prompt,