diff --git a/src/__tests__/ConfigService.test.ts b/src/__tests__/ConfigService.test.ts index 54092880..50505cc9 100644 --- a/src/__tests__/ConfigService.test.ts +++ b/src/__tests__/ConfigService.test.ts @@ -129,6 +129,7 @@ describe("ConfigService", () => { ], }, referenceExamples: {}, // Added referenceExamples + timeoutSeconds: 0, }; (fs.existsSync as jest.Mock).mockReturnValue(true); diff --git a/src/commands/run.ts b/src/commands/run.ts index 6bbb6e21..c51c85ef 100644 --- a/src/commands/run.ts +++ b/src/commands/run.ts @@ -1,5 +1,9 @@ import { Args, Command, Flags } from "@oclif/core"; -import { CrackedAgent, CrackedAgentOptions } from "@services/CrackedAgent"; +import { + CrackedAgent, + CrackedAgentOptions, + ExecutionResult, +} from "@services/CrackedAgent"; import { LLMProviderType } from "@services/LLM/LLMProvider"; import { ModelManager } from "@services/LLM/ModelManager"; import { OpenRouterAPI } from "@services/LLMProviders/OpenRouter/OpenRouterAPI"; @@ -16,6 +20,7 @@ export class Run extends Command { "$ run 'Add error handling'", "$ run --interactive # Start interactive mode", "$ run --init # Initialize configuration", + "$ run 'Add tests' --timeout 300 # Set timeout to 5 minutes", ]; static flags = { @@ -23,8 +28,11 @@ export class Run extends Command { description: "Initialize a default crkdrc.json configuration file", exclusive: ["interactive"], }), + timeout: Flags.integer({ + description: "Set timeout for the operation in seconds", + exclusive: ["init"], + }), }; - static args = { message: Args.string({ description: "Message describing the operation to perform", @@ -110,6 +118,7 @@ export class Run extends Command { ...config, options: this.parseOptions(config.options || ""), provider: config.provider as LLMProviderType, + timeout: (flags.timeout ?? config.timeoutSeconds ?? 0) * 1000, // Convert seconds to milliseconds }; // Validate provider @@ -132,22 +141,40 @@ export class Run extends Command { console.log("Press Enter to start the stream..."); this.rl.once("line", async () => { try { - const result = await agent.execute(args.message!, options); - if (!options.stream && result) { - this.log(result.response); - if (result.actions?.length) { - this.log("\nExecuted Actions:"); - result.actions.forEach(({ action, result }) => { - this.log(`\nAction: ${action}`); - this.log(`Result: ${JSON.stringify(result, null, 2)}`); - }); + const executePromise = agent.execute( + args.message!, + options, + ) as Promise; + + if (options.timeout > 0) { + const timeoutPromise = new Promise((_, reject) => { + setTimeout(() => { + reject( + new Error( + `Operation timed out after ${options.timeout / 1000} seconds`, + ), + ); + }, options.timeout); + }); + + try { + await Promise.race([executePromise, timeoutPromise]); + } catch (error) { + this.sessionManager.cleanup(); + if (error.message.includes("timed out")) { + console.error("\n" + error.message); + process.exit(1); + } + throw error; } + } else { + await executePromise; } - this.sessionManager.cleanup(); + process.exit(0); } catch (error) { - this.sessionManager.cleanup(); - this.error((error as Error).message); + console.error("Error:", error); + process.exit(1); } }); } diff --git a/src/services/ConfigService.ts b/src/services/ConfigService.ts index cf5a9dc1..b38ad394 100644 --- a/src/services/ConfigService.ts +++ b/src/services/ConfigService.ts @@ -97,6 +97,7 @@ const configSchema = z.object({ referenceExamples: z.record(z.string(), z.string()).optional().default({}), projectLanguage: z.string().default("typescript"), packageManager: z.string().default("yarn"), + timeoutSeconds: z.number().optional().default(0), // Add timeout property }); export type Config = z.infer; @@ -209,6 +210,7 @@ export class ConfigService { myService: "src/services/MyService.ts", anotherKey: "path/to/some/other/example.ts", }, + timeoutSeconds: 0, // Add default timeout }; fs.writeFileSync( this.CONFIG_PATH, diff --git a/src/services/CrackedAgent.ts b/src/services/CrackedAgent.ts index 58fb4206..6e77f19c 100644 --- a/src/services/CrackedAgent.ts +++ b/src/services/CrackedAgent.ts @@ -12,6 +12,7 @@ import { PhaseManager } from "./LLM/PhaseManager"; export interface CrackedAgentOptions { root?: string; + timeout: number; instructionsPath?: string; instructions?: string; provider?: LLMProviderType; @@ -98,6 +99,7 @@ export class CrackedAgent { clearContext: false, autoScaler: false, ...options, + timeout: 0, }; this.debugLogger.setDebug(finalOptions.debug); diff --git a/src/services/LLM/__tests__/PhaseManager.test.ts b/src/services/LLM/__tests__/PhaseManager.test.ts index 2a6b65ef..b8bb71ee 100644 --- a/src/services/LLM/__tests__/PhaseManager.test.ts +++ b/src/services/LLM/__tests__/PhaseManager.test.ts @@ -48,6 +48,7 @@ describe("PhaseManager", () => { lockFiles: ["package-lock.json"], }, referenceExamples: {}, + timeoutSeconds: 0, }; beforeAll(() => { @@ -143,6 +144,7 @@ describe("PhaseManager", () => { lockFiles: ["package-lock.json"], }, referenceExamples: {}, + timeoutSeconds: 0, }; jest.spyOn(configService, "getConfig").mockReturnValue(customConfig); @@ -200,6 +202,7 @@ describe("PhaseManager", () => { lockFiles: ["package-lock.json"], }, referenceExamples: {}, + timeoutSeconds: 0, }; jest.spyOn(configService, "getConfig").mockReturnValue(emptyConfig); diff --git a/src/services/LLMProviders/OpenRouter/OpenRouterAPI.ts b/src/services/LLMProviders/OpenRouter/OpenRouterAPI.ts index f77219ff..19bca826 100644 --- a/src/services/LLMProviders/OpenRouter/OpenRouterAPI.ts +++ b/src/services/LLMProviders/OpenRouter/OpenRouterAPI.ts @@ -50,6 +50,7 @@ export class OpenRouterAPI implements ILLMProvider { private retryDelay: number = 1000; private stream: any; private aborted: boolean = false; + private timeout: number = 0; constructor( private htmlEntityDecoder: HtmlEntityDecoder, @@ -389,6 +390,19 @@ export class OpenRouterAPI implements ILLMProvider { currentModel, ); + this.timeout > 0 + ? new Promise((_, reject) => { + setTimeout(() => { + console.error( + "\nOperation timed out in", + this.timeout / 1000, + "seconds", + ); + process.exit(0); + }, this.timeout); + }) + : null; + const streamOperation = async () => { const response = await this.makeRequest( "/chat/completions", @@ -522,4 +536,8 @@ export class OpenRouterAPI implements ILLMProvider { this.stream = null; } } + + updateTimeout(timeout: number) { + this.timeout = timeout; + } } diff --git a/src/services/LLMProviders/OpenRouter/__tests__/OpenRouterAPI.test.ts b/src/services/LLMProviders/OpenRouter/__tests__/OpenRouterAPI.test.ts index ab521dc8..49b3d672 100644 --- a/src/services/LLMProviders/OpenRouter/__tests__/OpenRouterAPI.test.ts +++ b/src/services/LLMProviders/OpenRouter/__tests__/OpenRouterAPI.test.ts @@ -243,6 +243,13 @@ describe("OpenRouterAPI", () => { }); }); describe("Streaming", () => { + // beforeEach(() => { + // jest.useFakeTimers(); + // }); + + // afterEach(() => { + // jest.useRealTimers(); + // }); it("should handle streaming messages correctly", async () => { const mockStreamData = [ 'data: {"choices": [{"delta": {"content": "Hel"}}]}\n', @@ -316,6 +323,35 @@ describe("OpenRouterAPI", () => { expect(callback).not.toHaveBeenCalled(); }); + it("should not timeout when timeout is 0", async () => { + const mockStreamData = [ + 'data: {"choices": [{"delta": {"content": "Hel"}}]}\n', + 'data: {"choices": [{"delta": {"content": "lo"}}]}\n', + 'data: {"choices": [{"delta": {"content": "!"}}]}\n', + "data: [DONE]\n", + ]; + + const mockStream = new Readable({ + read() { + mockStreamData.forEach((chunk) => { + this.push(Buffer.from(chunk)); + }); + this.push(null); + }, + }); + + postSpy.mockResolvedValue({ data: mockStream }); + openRouterAPI.updateTimeout(0); // No timeout + + const callback = jest.fn(); + await openRouterAPI.streamMessage("gpt-4", "Hi", callback); + + expect(callback).toHaveBeenCalledTimes(3); + expect(callback).toHaveBeenNthCalledWith(1, "Hel"); + expect(callback).toHaveBeenNthCalledWith(2, "lo"); + expect(callback).toHaveBeenNthCalledWith(3, "!"); + }); + it("should handle aborted stream", async () => { const mockStreamData = [ 'data: {"choices": [{"delta": {"content": "Hel"}}]}\n', diff --git a/src/services/__tests__/CrackedAgent.test.ts b/src/services/__tests__/CrackedAgent.test.ts index 25ba6d92..71736276 100644 --- a/src/services/__tests__/CrackedAgent.test.ts +++ b/src/services/__tests__/CrackedAgent.test.ts @@ -93,6 +93,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, instructionsPath: "path/to/instructions", + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -113,6 +114,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, instructions: "Custom instructions", + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -130,6 +132,7 @@ describe("CrackedAgent", () => { debug: false, clearContext: false, autoScaler: false, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -146,6 +149,7 @@ describe("CrackedAgent", () => { debug: false, clearContext: true, autoScaler: false, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -160,6 +164,7 @@ describe("CrackedAgent", () => { debug: false, clearContext: false, autoScaler: false, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -174,6 +179,7 @@ describe("CrackedAgent", () => { debug: false, clearContext: false, autoScaler: false, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -189,6 +195,7 @@ describe("CrackedAgent", () => { debug: true, clearContext: false, autoScaler: false, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -203,6 +210,7 @@ describe("CrackedAgent", () => { debug: false, clearContext: false, autoScaler: false, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -217,6 +225,7 @@ describe("CrackedAgent", () => { debug: false, clearContext: false, autoScaler: false, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -234,6 +243,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, root: "/custom/root", + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -253,6 +263,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, options: { key: "value" }, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -271,6 +282,7 @@ describe("CrackedAgent", () => { debug: false, clearContext: false, autoScaler: false, + timeout: 0, }; await crackedAgent.execute("", options); @@ -289,6 +301,7 @@ describe("CrackedAgent", () => { debug: false, clearContext: false, autoScaler: false, + timeout: 0, }; await crackedAgent.execute(undefined as unknown as string, options); @@ -308,6 +321,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, instructionsPath: "", + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -326,6 +340,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, instructionsPath: undefined as unknown as string, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -344,6 +359,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, instructions: "", + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -362,6 +378,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, instructions: undefined as unknown as string, + timeout: 0, }; await crackedAgent.execute("Mock message", options); @@ -380,6 +397,7 @@ describe("CrackedAgent", () => { clearContext: false, autoScaler: false, options: {}, + timeout: 0, }; await crackedAgent.execute("Mock message", options); diff --git a/src/services/streaming/InteractiveSessionManager.ts b/src/services/streaming/InteractiveSessionManager.ts index b7a4e003..a7e081e6 100644 --- a/src/services/streaming/InteractiveSessionManager.ts +++ b/src/services/streaming/InteractiveSessionManager.ts @@ -20,6 +20,7 @@ export class InteractiveSessionManager { this.rl = rl; this.agent = agent; this.options = options; + this.openRouterAPI.updateTimeout(this.options.timeout); } private setupKeypressHandling() { @@ -90,18 +91,33 @@ export class InteractiveSessionManager { public async start() { if (!this.rl) return; + let timeoutId: NodeJS.Timeout | null = null; + + const handleTimeout = () => { + console.error( + `Operation timed out after ${this.options.timeout / 1000} seconds`, + ); + this.cleanup(); + process.exit(1); + }; console.log( 'Interactive mode started. Type "exit" or press Ctrl+C to quit.', ); + this.setupKeypressHandling(); this.rl.prompt(); this.lineHandler = async (input: string) => { + if (timeoutId) clearTimeout(timeoutId); await this.handleInput(input); + if (this.options.timeout > 0) { + timeoutId = setTimeout(handleTimeout, this.options.timeout); + } }; this.closeHandler = () => { + if (timeoutId) clearTimeout(timeoutId); this.cleanup(); process.exit(0); };