diff --git a/src/core/timeout/TimeoutFallbackGenerator.ts b/src/core/timeout/TimeoutFallbackGenerator.ts deleted file mode 100644 index ffc1c29390..0000000000 --- a/src/core/timeout/TimeoutFallbackGenerator.ts +++ /dev/null @@ -1,128 +0,0 @@ -import type { ToolName } from "@roo-code/types" -import type { Task } from "../task/Task" -import type { SingleCompletionHandler } from "../../api" -import { - createTimeoutFallbackPrompt, - parseTimeoutFallbackResponse, - type TimeoutFallbackContext, -} from "../prompts/instructions/timeout-fallback" -import { formatResponse } from "../prompts/responses" - -export interface TimeoutFallbackResult { - success: boolean - toolCall?: { - name: "ask_followup_question" - params: { - question: string - follow_up: string - } - } - error?: string -} - -/** - * Generates AI-powered fallback suggestions for timeout scenarios - */ -export class TimeoutFallbackGenerator { - /** - * Generate an AI-powered ask_followup_question tool call for timeout scenarios - */ - public static async generateAiFallback( - context: TimeoutFallbackContext, - task?: Task, - ): Promise { - // Try to use AI to generate contextual suggestions - if (task?.api && "completePrompt" in task.api) { - try { - const aiResult = await this.generateAiSuggestions(context, task.api as SingleCompletionHandler) - if (aiResult.success) { - return aiResult - } - } catch (error) { - // AI failed, fall through to static suggestions - } - } - - // Fallback to static suggestions if AI fails or is unavailable - const toolCall = this.generateStaticToolCall(context) - - return { - success: true, - toolCall, - } - } - - /** - * Generate AI-powered suggestions using the task's API handler - */ - private static async generateAiSuggestions( - context: TimeoutFallbackContext, - apiHandler: SingleCompletionHandler, - ): Promise { - try { - const prompt = createTimeoutFallbackPrompt(context) - const aiResponse = await apiHandler.completePrompt(prompt) - - // Parse the AI response to extract suggestions - const suggestions = parseTimeoutFallbackResponse(aiResponse) - - if (suggestions.length === 0) { - throw new Error("No valid suggestions generated by AI") - } - - const question = `The ${context.toolName} operation timed out after ${Math.round(context.timeoutMs / 1000)} seconds. How would you like to proceed?` - - const followUpXml = suggestions - .map((suggestion) => - suggestion.mode - ? `${suggestion.text}` - : `${suggestion.text}`, - ) - .join("\n") - - return { - success: true, - toolCall: { - name: "ask_followup_question", - params: { - question, - follow_up: followUpXml, - }, - }, - } - } catch (error) { - return { - success: false, - error: error instanceof Error ? error.message : "Unknown error generating AI suggestions", - } - } - } - - /** - * Generate static fallback suggestions when AI is unavailable - */ - private static generateStaticToolCall(context: TimeoutFallbackContext): TimeoutFallbackResult["toolCall"] { - const suggestions = formatResponse.timeoutFallbackSuggestions.generateContextualSuggestions( - context.toolName, - context.toolParams, - ) - - const question = `The ${context.toolName} operation timed out after ${Math.round(context.timeoutMs / 1000)} seconds. How would you like to proceed?` - - const followUpXml = suggestions - .map((suggestion) => - suggestion.mode - ? `${suggestion.text}` - : `${suggestion.text}`, - ) - .join("\n") - - return { - name: "ask_followup_question", - params: { - question, - follow_up: followUpXml, - }, - } - } -} diff --git a/src/core/timeout/TimeoutFallbackHandler.ts b/src/core/timeout/TimeoutFallbackHandler.ts index b87dd8ca01..5aaa2c62f1 100644 --- a/src/core/timeout/TimeoutFallbackHandler.ts +++ b/src/core/timeout/TimeoutFallbackHandler.ts @@ -1,15 +1,32 @@ import type { ToolName } from "@roo-code/types" -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 type { SingleCompletionHandler } from "../../api" +import { + createTimeoutFallbackPrompt, + parseTimeoutFallbackResponse, + type TimeoutFallbackContext, +} from "../prompts/instructions/timeout-fallback" +import { formatResponse } from "../prompts/responses" + +export interface TimeoutFallbackResult { + success: boolean + toolCall?: { + name: "ask_followup_question" + params: { + question: string + follow_up: string + } + } + error?: string +} /** - * Generates AI-powered fallback suggestions for timeout scenarios + * Unified timeout fallback handler that generates AI-powered fallback suggestions + * and creates timeout responses in a single optimized flow */ export class TimeoutFallbackHandler { /** - * Create a timeout response with AI-generated fallback question + * Create a timeout response with AI-generated fallback question in a single optimized query */ public static async createTimeoutResponse( toolName: ToolName, @@ -40,12 +57,11 @@ export class TimeoutFallbackHandler { : undefined, } - // Generate AI-powered fallback (with static fallback if AI fails) - const aiResult = await TimeoutFallbackGenerator.generateAiFallback(aiContext, task) + // Generate AI-powered fallback (with static fallback if AI fails) in a single call + const aiResult = await this.generateAiFallback(aiContext, task) 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 + // Return a response that instructs the model to ask a follow-up question const { question, follow_up } = aiResult.toolCall.params // Format the response to explicitly instruct the model to ask the follow-up question @@ -66,4 +82,106 @@ 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.` } + + /** + * Generate an AI-powered ask_followup_question tool call for timeout scenarios + */ + public static async generateAiFallback( + context: TimeoutFallbackContext, + task?: Task, + ): Promise { + // Try to use AI to generate contextual suggestions + if (task?.api && "completePrompt" in task.api) { + try { + const aiResult = await this.generateAiSuggestions(context, task.api as SingleCompletionHandler) + if (aiResult.success) { + return aiResult + } + } catch (error) { + // AI failed, fall through to static suggestions + } + } + + // Fallback to static suggestions if AI fails or is unavailable + const toolCall = this.generateStaticToolCall(context) + + return { + success: true, + toolCall, + } + } + + /** + * Generate AI-powered suggestions using the task's API handler + */ + private static async generateAiSuggestions( + context: TimeoutFallbackContext, + apiHandler: SingleCompletionHandler, + ): Promise { + try { + const prompt = createTimeoutFallbackPrompt(context) + const aiResponse = await apiHandler.completePrompt(prompt) + + // Parse the AI response to extract suggestions + const suggestions = parseTimeoutFallbackResponse(aiResponse) + + if (suggestions.length === 0) { + throw new Error("No valid suggestions generated by AI") + } + + const question = `The ${context.toolName} operation timed out after ${Math.round(context.timeoutMs / 1000)} seconds. How would you like to proceed?` + + const followUpXml = suggestions + .map((suggestion) => + suggestion.mode + ? `${suggestion.text}` + : `${suggestion.text}`, + ) + .join("\n") + + return { + success: true, + toolCall: { + name: "ask_followup_question", + params: { + question, + follow_up: followUpXml, + }, + }, + } + } catch (error) { + return { + success: false, + error: error instanceof Error ? error.message : "Unknown error generating AI suggestions", + } + } + } + + /** + * Generate static fallback suggestions when AI is unavailable + */ + private static generateStaticToolCall(context: TimeoutFallbackContext): TimeoutFallbackResult["toolCall"] { + const suggestions = formatResponse.timeoutFallbackSuggestions.generateContextualSuggestions( + context.toolName, + context.toolParams, + ) + + const question = `The ${context.toolName} operation timed out after ${Math.round(context.timeoutMs / 1000)} seconds. How would you like to proceed?` + + const followUpXml = suggestions + .map((suggestion) => + suggestion.mode + ? `${suggestion.text}` + : `${suggestion.text}`, + ) + .join("\n") + + return { + name: "ask_followup_question", + params: { + question, + follow_up: followUpXml, + }, + } + } } diff --git a/src/core/timeout/__tests__/ai-fallback-real.spec.ts b/src/core/timeout/__tests__/ai-fallback-real.spec.ts index 021c02eb3f..7b378a2b01 100644 --- a/src/core/timeout/__tests__/ai-fallback-real.spec.ts +++ b/src/core/timeout/__tests__/ai-fallback-real.spec.ts @@ -1,14 +1,14 @@ // npx vitest run src/core/timeout/__tests__/ai-fallback-real.spec.ts import { describe, test, expect, beforeEach, vitest } from "vitest" -import { TimeoutFallbackGenerator } from "../TimeoutFallbackGenerator" +import { TimeoutFallbackHandler } from "../TimeoutFallbackHandler" import type { ApiHandler, SingleCompletionHandler } from "../../../api" import type { Task } from "../../task/Task" // Create a mock API handler that extends ApiHandler and includes completePrompt interface MockApiHandler extends ApiHandler, SingleCompletionHandler {} -describe("TimeoutFallbackGenerator - Real AI Implementation", () => { +describe("TimeoutFallbackHandler - Real AI Implementation", () => { let mockApiHandler: MockApiHandler let mockTask: Partial @@ -47,7 +47,7 @@ describe("TimeoutFallbackGenerator - Real AI Implementation", () => { toolParams: { command: "npm install" }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) expect(result.toolCall).toBeDefined() @@ -80,7 +80,7 @@ describe("TimeoutFallbackGenerator - Real AI Implementation", () => { toolParams: { command: "npm test" }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) expect(result.toolCall).toBeDefined() @@ -105,7 +105,7 @@ describe("TimeoutFallbackGenerator - Real AI Implementation", () => { toolParams: { path: "/large/file.txt" }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, taskWithoutApi as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, taskWithoutApi as Task) expect(result.success).toBe(true) expect(result.toolCall).toBeDefined() @@ -134,7 +134,7 @@ describe("TimeoutFallbackGenerator - Real AI Implementation", () => { toolParams: { command: "build script" }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) const followUp = result.toolCall?.params.follow_up || "" @@ -157,7 +157,7 @@ describe("TimeoutFallbackGenerator - Real AI Implementation", () => { toolParams: { action: "click" }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) const followUp = result.toolCall?.params.follow_up || "" @@ -183,7 +183,7 @@ describe("TimeoutFallbackGenerator - Real AI Implementation", () => { }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) @@ -216,7 +216,7 @@ describe("TimeoutFallbackGenerator - Real AI Implementation", () => { toolParams: { path: "/output.txt" }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) const followUp = result.toolCall?.params.follow_up || "" diff --git a/src/core/timeout/__tests__/e2e-ai-test.spec.ts b/src/core/timeout/__tests__/e2e-ai-test.spec.ts index 7b1fe4d23e..40bd9e5406 100644 --- a/src/core/timeout/__tests__/e2e-ai-test.spec.ts +++ b/src/core/timeout/__tests__/e2e-ai-test.spec.ts @@ -2,14 +2,14 @@ // npx vitest run src/core/timeout/__tests__/e2e-ai-test.spec.ts import { describe, test, expect, beforeEach, vitest } from "vitest" -import { TimeoutFallbackGenerator } from "../TimeoutFallbackGenerator" +import { TimeoutFallbackHandler } from "../TimeoutFallbackHandler" import type { ApiHandler, SingleCompletionHandler } from "../../../api" import type { Task } from "../../task/Task" // Mock API handler that simulates a real AI provider interface TestApiHandler extends ApiHandler, SingleCompletionHandler {} -describe("TimeoutFallbackGenerator - End-to-End AI Test", () => { +describe("TimeoutFallbackHandler - End-to-End AI Test", () => { test("should generate realistic AI suggestions for execute_command timeout", async () => { // Create a realistic mock API handler const mockApiHandler: TestApiHandler = { @@ -42,7 +42,7 @@ Here are some suggestions for the npm install timeout: }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) // Verify the result structure expect(result.success).toBe(true) @@ -99,7 +99,7 @@ Based on the search_files timeout, here are my recommendations: }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) expect(result.toolCall).toBeDefined() @@ -134,7 +134,7 @@ Based on the search_files timeout, here are my recommendations: }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) expect(result.toolCall).toBeDefined() @@ -164,7 +164,7 @@ Based on the search_files timeout, here are my recommendations: }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context, mockTask as Task) + const result = await TimeoutFallbackHandler.generateAiFallback(context, mockTask as Task) expect(result.success).toBe(true) expect(result.toolCall).toBeDefined() diff --git a/src/core/timeout/__tests__/tool-injection-test.spec.ts b/src/core/timeout/__tests__/tool-injection-test.spec.ts index 8208139f03..49c9d0d225 100644 --- a/src/core/timeout/__tests__/tool-injection-test.spec.ts +++ b/src/core/timeout/__tests__/tool-injection-test.spec.ts @@ -1,9 +1,13 @@ import { describe, test, expect, vi, beforeEach } from "vitest" import { TimeoutFallbackHandler } from "../TimeoutFallbackHandler" -import { TimeoutFallbackGenerator, type TimeoutFallbackResult } from "../TimeoutFallbackGenerator" +import { type TimeoutFallbackResult } from "../TimeoutFallbackHandler" import { Task } from "../../task/Task" -vi.mock("../TimeoutFallbackGenerator") +// Import the real module first +import { TimeoutFallbackHandler as RealTimeoutFallbackHandler } from "../TimeoutFallbackHandler" + +// Mock only the generateAiFallback method +vi.spyOn(RealTimeoutFallbackHandler, "generateAiFallback") describe("Tool Call Response Test", () => { let mockTask: Task @@ -20,7 +24,7 @@ describe("Tool Call Response Test", () => { }) test("should return response with ask_followup_question tool instructions", async () => { - // Mock the TimeoutFallbackGenerator to return a successful AI result + // Mock the TimeoutFallbackHandler to return a successful AI result const mockAiResult: TimeoutFallbackResult = { success: true, toolCall: { @@ -34,7 +38,7 @@ describe("Tool Call Response Test", () => { } // Mock the generateAiFallback method - vi.mocked(TimeoutFallbackGenerator.generateAiFallback).mockResolvedValue(mockAiResult) + vi.mocked(RealTimeoutFallbackHandler.generateAiFallback).mockResolvedValue(mockAiResult) // Call createTimeoutResponse const response = await TimeoutFallbackHandler.createTimeoutResponse( @@ -65,7 +69,7 @@ describe("Tool Call Response Test", () => { test("should return fallback message when AI generation fails", async () => { // Mock the generateAiFallback to return a failure - vi.mocked(TimeoutFallbackGenerator.generateAiFallback).mockResolvedValue({ + vi.mocked(RealTimeoutFallbackHandler.generateAiFallback).mockResolvedValue({ success: false, error: "AI generation failed", }) diff --git a/src/core/timeout/__tests__/ui-integration.spec.ts b/src/core/timeout/__tests__/ui-integration.spec.ts index f44610a857..8065aa12ed 100644 --- a/src/core/timeout/__tests__/ui-integration.spec.ts +++ b/src/core/timeout/__tests__/ui-integration.spec.ts @@ -1,5 +1,4 @@ import { describe, it, expect, vi } from "vitest" -import { TimeoutFallbackGenerator } from "../TimeoutFallbackGenerator" import { TimeoutFallbackHandler } from "../TimeoutFallbackHandler" describe("UI Integration - AI Timeout Fallbacks", () => { @@ -11,7 +10,7 @@ describe("UI Integration - AI Timeout Fallbacks", () => { toolParams: { command: "npm install" }, } - const result = await TimeoutFallbackGenerator.generateAiFallback(context) + const result = await TimeoutFallbackHandler.generateAiFallback(context) expect(result.success).toBe(true) expect(result.toolCall).toBeDefined() diff --git a/src/core/timeout/index.ts b/src/core/timeout/index.ts index 8b1f080c71..ad808bc1ca 100644 --- a/src/core/timeout/index.ts +++ b/src/core/timeout/index.ts @@ -1,9 +1,8 @@ export { TimeoutManager, timeoutManager } from "./TimeoutManager" export { ToolExecutionWrapper } from "./ToolExecutionWrapper" export { TimeoutFallbackHandler } from "./TimeoutFallbackHandler" -export { TimeoutFallbackGenerator } from "./TimeoutFallbackGenerator" export type { TimeoutConfig, TimeoutResult, TimeoutEvent } from "./TimeoutManager" export type { ToolExecutionOptions } from "./ToolExecutionWrapper" -export type { TimeoutFallbackResult } from "./TimeoutFallbackGenerator" +export type { TimeoutFallbackResult } from "./TimeoutFallbackHandler" export type { TimeoutFallbackContext } from "../prompts/instructions/timeout-fallback"