Fix ask_followup_question streaming issue and add missing tool cases (#9561)

Author: daniel-lxs

src/core/assistant-message/NativeToolCallParser.ts

@@ -355,6 +355,117 @@ export class NativeToolCallParser {
 				}
 				break
 
+			case "ask_followup_question":
+				if (partialArgs.question !== undefined || partialArgs.follow_up !== undefined) {
+					nativeArgs = {
+						question: partialArgs.question,
+						follow_up: Array.isArray(partialArgs.follow_up) ? partialArgs.follow_up : undefined,
+					}
+				}
+				break
+
+			case "apply_diff":
+				if (partialArgs.path !== undefined || partialArgs.diff !== undefined) {
+					nativeArgs = {
+						path: partialArgs.path,
+						diff: partialArgs.diff,
+					}
+				}
+				break
+
+			case "browser_action":
+				if (partialArgs.action !== undefined) {
+					nativeArgs = {
+						action: partialArgs.action,
+						url: partialArgs.url,
+						coordinate: partialArgs.coordinate,
+						size: partialArgs.size,
+						text: partialArgs.text,
+					}
+				}
+				break
+
+			case "codebase_search":
+				if (partialArgs.query !== undefined) {
+					nativeArgs = {
+						query: partialArgs.query,
+						path: partialArgs.path,
+					}
+				}
+				break
+
+			case "fetch_instructions":
+				if (partialArgs.task !== undefined) {
+					nativeArgs = {
+						task: partialArgs.task,
+					}
+				}
+				break
+
+			case "generate_image":
+				if (partialArgs.prompt !== undefined || partialArgs.path !== undefined) {
+					nativeArgs = {
+						prompt: partialArgs.prompt,
+						path: partialArgs.path,
+						image: partialArgs.image,
+					}
+				}
+				break
+
+			case "list_code_definition_names":
+				if (partialArgs.path !== undefined) {
+					nativeArgs = {
+						path: partialArgs.path,
+					}
+				}
+				break
+
+			case "run_slash_command":
+				if (partialArgs.command !== undefined) {
+					nativeArgs = {
+						command: partialArgs.command,
+						args: partialArgs.args,
+					}
+				}
+				break
+
+			case "search_files":
+				if (partialArgs.path !== undefined || partialArgs.regex !== undefined) {
+					nativeArgs = {
+						path: partialArgs.path,
+						regex: partialArgs.regex,
+						file_pattern: partialArgs.file_pattern,
+					}
+				}
+				break
+
+			case "switch_mode":
+				if (partialArgs.mode_slug !== undefined || partialArgs.reason !== undefined) {
+					nativeArgs = {
+						mode_slug: partialArgs.mode_slug,
+						reason: partialArgs.reason,
+					}
+				}
+				break
+
+			case "update_todo_list":
+				if (partialArgs.todos !== undefined) {
+					nativeArgs = {
+						todos: partialArgs.todos,
+					}
+				}
+				break
+
+			case "use_mcp_tool":
+				if (partialArgs.server_name !== undefined || partialArgs.tool_name !== undefined) {
+					nativeArgs = {
+						server_name: partialArgs.server_name,
+						tool_name: partialArgs.tool_name,
+						arguments: partialArgs.arguments,
+					}
+				}
+				break
+
 			// Add other tools as needed
 			default:
 				break

src/core/tools/AskFollowupQuestionTool.ts

@@ -91,7 +91,11 @@ export class AskFollowupQuestionTool extends BaseTool<"ask_followup_question"> {
 	}
 
 	override async handlePartial(task: Task, block: ToolUse<"ask_followup_question">): Promise<void> {
-		const question: string | undefined = block.params.question
+		// Get question from params (for XML protocol) or nativeArgs (for native protocol)
+		const question: string | undefined = block.params.question ?? block.nativeArgs?.question
+
+		// During partial streaming, only show the question to avoid displaying raw JSON
+		// The full JSON with suggestions will be sent when the tool call is complete (!block.partial)
 		await task
 			.ask("followup", this.removeClosingTag("question", question, block.partial), block.partial)
 			.catch(() => {})

src/core/tools/__tests__/askFollowupQuestionTool.spec.ts

@@ -1,5 +1,6 @@
 import { askFollowupQuestionTool } from "../AskFollowupQuestionTool"
 import { ToolUse } from "../../../shared/tools"
+import { NativeToolCallParser } from "../../assistant-message/NativeToolCallParser"
 
 describe("askFollowupQuestionTool", () => {
 	let mockCline: any
@@ -101,4 +102,104 @@ describe("askFollowupQuestionTool", () => {
 			false,
 		)
 	})
+
+	describe("handlePartial with native protocol", () => {
+		it("should only send question during partial streaming to avoid raw JSON display", async () => {
+			const block: ToolUse<"ask_followup_question"> = {
+				type: "tool_use",
+				name: "ask_followup_question",
+				params: {
+					question: "What would you like to do?",
+				},
+				partial: true,
+				nativeArgs: {
+					question: "What would you like to do?",
+					follow_up: [{ text: "Option 1", mode: "code" }, { text: "Option 2" }],
+				},
+			}
+
+			await askFollowupQuestionTool.handle(mockCline, block, {
+				askApproval: vi.fn(),
+				handleError: vi.fn(),
+				pushToolResult: mockPushToolResult,
+				removeClosingTag: vi.fn((tag, content) => content || ""),
+				toolProtocol: "native",
+			})
+
+			// During partial streaming, only the question should be sent (not JSON with suggestions)
+			expect(mockCline.ask).toHaveBeenCalledWith("followup", "What would you like to do?", true)
+		})
+
+		it("should handle partial with question from params", async () => {
+			const block: ToolUse<"ask_followup_question"> = {
+				type: "tool_use",
+				name: "ask_followup_question",
+				params: {
+					question: "Choose wisely",
+				},
+				partial: true,
+			}
+
+			await askFollowupQuestionTool.handle(mockCline, block, {
+				askApproval: vi.fn(),
+				handleError: vi.fn(),
+				pushToolResult: mockPushToolResult,
+				removeClosingTag: vi.fn((tag, content) => content || ""),
+				toolProtocol: "xml",
+			})
+
+			expect(mockCline.ask).toHaveBeenCalledWith("followup", "Choose wisely", true)
+		})
+	})
+
+	describe("NativeToolCallParser.createPartialToolUse for ask_followup_question", () => {
+		beforeEach(() => {
+			NativeToolCallParser.clearAllStreamingToolCalls()
+			NativeToolCallParser.clearRawChunkState()
+		})
+
+		it("should build nativeArgs with question and follow_up during streaming", () => {
+			// Start a streaming tool call
+			NativeToolCallParser.startStreamingToolCall("call_123", "ask_followup_question")
+
+			// Simulate streaming JSON chunks
+			const chunk1 = '{"question":"What would you like?","follow_up":[{"text":"Option 1","mode":"code"}'
+			const result1 = NativeToolCallParser.processStreamingChunk("call_123", chunk1)
+
+			expect(result1).not.toBeNull()
+			expect(result1?.name).toBe("ask_followup_question")
+			expect(result1?.params.question).toBe("What would you like?")
+			expect(result1?.nativeArgs).toBeDefined()
+			// Use type assertion to access the specific fields
+			const nativeArgs = result1?.nativeArgs as {
+				question: string
+				follow_up?: Array<{ text: string; mode?: string }>
+			}
+			expect(nativeArgs?.question).toBe("What would you like?")
+			// partial-json should parse the incomplete array
+			expect(nativeArgs?.follow_up).toBeDefined()
+		})
+
+		it("should finalize with complete nativeArgs", () => {
+			NativeToolCallParser.startStreamingToolCall("call_456", "ask_followup_question")
+
+			// Add complete JSON
+			const completeJson =
+				'{"question":"Choose an option","follow_up":[{"text":"Yes","mode":"code"},{"text":"No","mode":null}]}'
+			NativeToolCallParser.processStreamingChunk("call_456", completeJson)
+
+			const result = NativeToolCallParser.finalizeStreamingToolCall("call_456")
+
+			expect(result).not.toBeNull()
+			expect(result?.name).toBe("ask_followup_question")
+			expect(result?.partial).toBe(false)
+			expect(result?.nativeArgs).toEqual({
+				question: "Choose an option",
+				follow_up: [
+					{ text: "Yes", mode: "code" },
+					{ text: "No", mode: null },
+				],
+			})
+		})
+	})
 })