diff --git a/src/core/timeout/TimeoutFallbackHandler.ts b/src/core/timeout/TimeoutFallbackHandler.ts index 0fd266c440..b87dd8ca01 100644 --- a/src/core/timeout/TimeoutFallbackHandler.ts +++ b/src/core/timeout/TimeoutFallbackHandler.ts @@ -3,7 +3,6 @@ import { formatResponse } from "../prompts/responses" import type { TimeoutFallbackContext } from "../prompts/instructions/timeout-fallback" import { TimeoutFallbackGenerator, type TimeoutFallbackResult } from "./TimeoutFallbackGenerator" import type { Task } from "../task/Task" -import { parseAssistantMessage } from "../assistant-message/parseAssistantMessage" /** * Generates AI-powered fallback suggestions for timeout scenarios @@ -44,36 +43,27 @@ export class TimeoutFallbackHandler { // Generate AI-powered fallback (with static fallback if AI fails) const aiResult = await TimeoutFallbackGenerator.generateAiFallback(aiContext, task) - if (aiResult.success && aiResult.toolCall && task) { - // Inject the tool call directly into the assistant message content for proper execution - this.injectToolCallIntoMessageContent(aiResult.toolCall, task) - return baseResponse + if (aiResult.success && aiResult.toolCall) { + // Instead of injecting the tool call, we'll return a response that instructs + // the model to ask a follow-up question in its next message + const { question, follow_up } = aiResult.toolCall.params + + // Format the response to explicitly instruct the model to ask the follow-up question + return `${baseResponse} + +The operation timed out. You MUST now use the ask_followup_question tool with the following parameters: + + +${question} + +${follow_up} + + + +This is required to help the user decide how to proceed after the timeout.` } // This should rarely happen since generateAiFallback always provides static fallback return `${baseResponse}\n\nThe operation timed out. Please consider breaking this into smaller steps or trying a different approach.` } - - /** - * Inject a tool call directly into the assistant message content for proper parsing and execution - */ - private static injectToolCallIntoMessageContent(toolCall: TimeoutFallbackResult["toolCall"], task: Task): void { - if (toolCall?.name === "ask_followup_question" && toolCall.params) { - const { question, follow_up } = toolCall.params - - // Create the XML tool call string - const toolCallXml = ` -${question} - -${follow_up} - -` - - // Parse the tool call XML to create proper assistant message content - const parsedContent = parseAssistantMessage(toolCallXml) - - // Add the parsed tool call to the assistant message content - task.assistantMessageContent.push(...parsedContent) - } - } } diff --git a/src/core/timeout/__tests__/timeout-integration.spec.ts b/src/core/timeout/__tests__/timeout-integration.spec.ts index 027677500a..5bd741e51c 100644 --- a/src/core/timeout/__tests__/timeout-integration.spec.ts +++ b/src/core/timeout/__tests__/timeout-integration.spec.ts @@ -123,13 +123,18 @@ describe("Timeout Integration Tests", () => { expect(response).toContain("6s") expect(response.length).toBeGreaterThan(50) - // The AI-generated tool call should be injected into the task's assistant message content - const toolUseBlock = mockTask.assistantMessageContent.find((block) => block.type === "tool_use") - expect(toolUseBlock).toBeDefined() - if (toolUseBlock?.type === "tool_use") { - expect(toolUseBlock.name).toBe("ask_followup_question") - expect(toolUseBlock.params?.question).toContain("timed out") - } + // The response should now contain instructions to use ask_followup_question + expect(response).toContain("You MUST now use the ask_followup_question tool") + expect(response).toContain("") + expect(response).toContain("") + expect(response).toContain("timed out") + expect(response).toContain("") + expect(response).toContain("") + expect(response).toContain("") + expect(response).toContain("") + + // Verify that assistantMessageContent was NOT modified + expect(mockTask.assistantMessageContent).toHaveLength(0) }) test("AbortSignal should be properly handled", async () => { diff --git a/src/core/timeout/__tests__/tool-injection-test.spec.ts b/src/core/timeout/__tests__/tool-injection-test.spec.ts index 1098fd0c73..8208139f03 100644 --- a/src/core/timeout/__tests__/tool-injection-test.spec.ts +++ b/src/core/timeout/__tests__/tool-injection-test.spec.ts @@ -1,8 +1,11 @@ import { describe, test, expect, vi, beforeEach } from "vitest" import { TimeoutFallbackHandler } from "../TimeoutFallbackHandler" +import { TimeoutFallbackGenerator, type TimeoutFallbackResult } from "../TimeoutFallbackGenerator" import { Task } from "../../task/Task" -describe("Tool Call Injection Test", () => { +vi.mock("../TimeoutFallbackGenerator") + +describe("Tool Call Response Test", () => { let mockTask: Task beforeEach(() => { @@ -10,28 +13,28 @@ describe("Tool Call Injection Test", () => { mockTask = { assistantMessageContent: [], cwd: "/test/dir", + say: vi.fn(), } as unknown as Task + + vi.clearAllMocks() }) - test("should inject ask_followup_question tool call into assistant message content", async () => { + test("should return response with ask_followup_question tool instructions", async () => { // Mock the TimeoutFallbackGenerator to return a successful AI result - const mockAiResult = { + const mockAiResult: TimeoutFallbackResult = { success: true, toolCall: { name: "ask_followup_question", params: { - question: "What would you like to do next?", - follow_up: "Try a different approachBreak into smaller steps", + question: "The execute_command operation timed out after 5 seconds. How would you like to proceed?", + follow_up: + "Try a different approach\nBreak into smaller steps", }, }, } // Mock the generateAiFallback method - vi.doMock("../TimeoutFallbackGenerator", () => ({ - TimeoutFallbackGenerator: { - generateAiFallback: vi.fn().mockResolvedValue(mockAiResult), - }, - })) + vi.mocked(TimeoutFallbackGenerator.generateAiFallback).mockResolvedValue(mockAiResult) // Call createTimeoutResponse const response = await TimeoutFallbackHandler.createTimeoutResponse( @@ -42,21 +45,50 @@ describe("Tool Call Injection Test", () => { mockTask, ) - // Check that the response is just the base timeout message + // Check that the response contains the base timeout message expect(response).toContain("timed out after 5 seconds") expect(response).toContain("Execution Time: 6s") - // Check that the tool call was injected into assistantMessageContent - // The XML parser might create multiple blocks (text + tool_use), so find the tool_use block - const toolUseBlock = mockTask.assistantMessageContent.find((block) => block.type === "tool_use") - expect(toolUseBlock).toBeDefined() - expect(toolUseBlock?.type).toBe("tool_use") - expect(toolUseBlock?.name).toBe("ask_followup_question") - expect(toolUseBlock?.params?.question).toBeDefined() - expect(toolUseBlock?.params?.follow_up).toBeDefined() + // Check that the response includes instructions to use ask_followup_question + expect(response).toContain("You MUST now use the ask_followup_question tool") + expect(response).toContain("") + expect(response).toContain(`${mockAiResult.toolCall?.params.question}`) + expect(response).toContain("") + expect(response).toContain(mockAiResult.toolCall?.params.follow_up) + expect(response).toContain("") + expect(response).toContain("") + expect(response).toContain("This is required to help the user decide how to proceed after the timeout.") - // Verify the question contains timeout information - expect(toolUseBlock?.params?.question).toContain("timed out") - expect(toolUseBlock?.params?.question).toContain("5 seconds") + // Verify that assistantMessageContent was NOT modified + expect(mockTask.assistantMessageContent).toHaveLength(0) + }) + + test("should return fallback message when AI generation fails", async () => { + // Mock the generateAiFallback to return a failure + vi.mocked(TimeoutFallbackGenerator.generateAiFallback).mockResolvedValue({ + success: false, + error: "AI generation failed", + }) + + // Call createTimeoutResponse + const response = await TimeoutFallbackHandler.createTimeoutResponse( + "execute_command", + 5000, + 6000, + { command: "npm install" }, + mockTask, + ) + + // Check that the response contains the base timeout message + expect(response).toContain("timed out after 5 seconds") + expect(response).toContain("Execution Time: 6s") + + // Check that the response contains the fallback message + expect(response).toContain( + "The operation timed out. Please consider breaking this into smaller steps or trying a different approach.", + ) + + // Should not contain ask_followup_question instructions + expect(response).not.toContain("ask_followup_question") }) })