From e45a4a0cf80b70bdaa9cfea057db1bca25524609 Mon Sep 17 00:00:00 2001 From: Brace Sproul Date: Fri, 18 Jul 2025 10:19:33 -0700 Subject: [PATCH] feat: Add withRetry func and longer timeouts for cloning (#449) --- apps/open-swe/src/__tests__/retry.test.ts | 184 ++++++++++++++++++ .../src/graphs/shared/initialize-sandbox.ts | 17 +- apps/open-swe/src/utils/github/git.ts | 6 + apps/open-swe/src/utils/retry.ts | 50 +++++ 4 files changed, 253 insertions(+), 4 deletions(-) create mode 100644 apps/open-swe/src/__tests__/retry.test.ts create mode 100644 apps/open-swe/src/utils/retry.ts diff --git a/apps/open-swe/src/__tests__/retry.test.ts b/apps/open-swe/src/__tests__/retry.test.ts new file mode 100644 index 00000000..f4954ea6 --- /dev/null +++ b/apps/open-swe/src/__tests__/retry.test.ts @@ -0,0 +1,184 @@ +import { describe, it, expect, jest } from "@jest/globals"; +import { withRetry, createRetryWrapper } from "../utils/retry.js"; + +describe("withRetry", () => { + it("should return result on first success", async () => { + const mockFn = jest + .fn<() => Promise>() + .mockResolvedValue("success"); + + const result = await withRetry(mockFn); + + expect(result).toBe("success"); + expect(mockFn).toHaveBeenCalledTimes(1); + }); + + it("should retry on failure and eventually succeed", async () => { + const mockFn = jest + .fn<() => Promise>() + .mockRejectedValueOnce(new Error("fail 1")) + .mockRejectedValueOnce(new Error("fail 2")) + .mockResolvedValue("success"); + + const result = await withRetry(mockFn); + + expect(result).toBe("success"); + expect(mockFn).toHaveBeenCalledTimes(3); + }); + + it("should use default retries of 3", async () => { + const mockFn = jest + .fn<() => Promise>() + .mockRejectedValue(new Error("always fails")); + + await expect(withRetry(mockFn)).rejects.toThrow("always fails"); + expect(mockFn).toHaveBeenCalledTimes(4); // 1 initial + 3 retries + }); + + it("should respect custom retry count", async () => { + const mockFn = jest + .fn<() => Promise>() + .mockRejectedValue(new Error("always fails")); + + await expect(withRetry(mockFn, { retries: 2 })).rejects.toThrow( + "always fails", + ); + expect(mockFn).toHaveBeenCalledTimes(3); // 1 initial + 2 retries + }); + + it("should respect custom delay", async () => { + const mockFn = jest + .fn<() => Promise>() + .mockRejectedValue(new Error("always fails")); + const startTime = Date.now(); + + await expect(withRetry(mockFn, { retries: 2, delay: 100 })).rejects.toThrow( + "always fails", + ); + const endTime = Date.now(); + + expect(mockFn).toHaveBeenCalledTimes(3); + expect(endTime - startTime).toBeGreaterThanOrEqual(200); // 2 delays of 100ms each + }); + + it("should not delay with default delay of 0", async () => { + const mockFn = jest + .fn<() => Promise>() + .mockRejectedValue(new Error("always fails")); + const startTime = Date.now(); + + await expect(withRetry(mockFn, { retries: 2 })).rejects.toThrow( + "always fails", + ); + const endTime = Date.now(); + + expect(mockFn).toHaveBeenCalledTimes(3); + expect(endTime - startTime).toBeLessThan(50); // Should be very fast with no delay + }); + + it("should handle non-Error objects", async () => { + const mockFn = jest + .fn<() => Promise>() + .mockRejectedValue("string error"); + + await expect(withRetry(mockFn, { retries: 1 })).rejects.toThrow( + "string error", + ); + expect(mockFn).toHaveBeenCalledTimes(2); + }); + + it("should throw the last error after all retries", async () => { + const error1 = new Error("first error"); + const error2 = new Error("second error"); + const lastError = new Error("last error"); + + const mockFn = jest + .fn<() => Promise>() + .mockRejectedValueOnce(error1) + .mockRejectedValueOnce(error2) + .mockRejectedValue(lastError); + + await expect(withRetry(mockFn, { retries: 2 })).rejects.toThrow( + "last error", + ); + expect(mockFn).toHaveBeenCalledTimes(3); + }); + + it("should work with async functions that return different types", async () => { + const numberFn = jest.fn<() => Promise>().mockResolvedValue(42); + const objectFn = jest + .fn<() => Promise<{ key: string }>>() + .mockResolvedValue({ key: "value" }); + const arrayFn = jest + .fn<() => Promise>() + .mockResolvedValue([1, 2, 3]); + + expect(await withRetry(numberFn)).toBe(42); + expect(await withRetry(objectFn)).toEqual({ key: "value" }); + expect(await withRetry(arrayFn)).toEqual([1, 2, 3]); + }); +}); + +describe("createRetryWrapper", () => { + it("should create a wrapper that retries with default options", async () => { + const originalFn = jest + .fn<() => Promise>() + .mockRejectedValueOnce(new Error("fail")) + .mockResolvedValue("success"); + + const wrappedFn = createRetryWrapper(originalFn); + const result = await wrappedFn(); + + expect(result).toBe("success"); + expect(originalFn).toHaveBeenCalledTimes(2); + }); + + it("should create a wrapper that retries with custom options", async () => { + const originalFn = jest + .fn<() => Promise>() + .mockRejectedValue(new Error("always fails")); + + const wrappedFn = createRetryWrapper(originalFn, { retries: 1 }); + + await expect(wrappedFn()).rejects.toThrow("always fails"); + expect(originalFn).toHaveBeenCalledTimes(2); // 1 initial + 1 retry + }); + + it("should preserve function arguments", async () => { + const originalFn = jest + .fn<(a: string, b: string, c: number) => Promise>() + .mockResolvedValue("success"); + + const wrappedFn = createRetryWrapper(originalFn); + const result = await wrappedFn("arg1", "arg2", 123); + + expect(result).toBe("success"); + expect(originalFn).toHaveBeenCalledWith("arg1", "arg2", 123); + }); + + it("should work with functions that have multiple parameters", async () => { + const originalFn = jest.fn((a: string, b: number, c: boolean) => + Promise.resolve(`${a}-${b}-${c}`), + ); + + const wrappedFn = createRetryWrapper(originalFn); + const result = await wrappedFn("test", 42, true); + + expect(result).toBe("test-42-true"); + expect(originalFn).toHaveBeenCalledWith("test", 42, true); + }); + + it("should retry with the same arguments on each attempt", async () => { + const originalFn = jest + .fn<(a: string, b: string) => Promise>() + .mockRejectedValueOnce(new Error("fail")) + .mockResolvedValue("success"); + + const wrappedFn = createRetryWrapper(originalFn); + await wrappedFn("arg1", "arg2"); + + expect(originalFn).toHaveBeenCalledTimes(2); + expect(originalFn).toHaveBeenNthCalledWith(1, "arg1", "arg2"); + expect(originalFn).toHaveBeenNthCalledWith(2, "arg1", "arg2"); + }); +}); diff --git a/apps/open-swe/src/graphs/shared/initialize-sandbox.ts b/apps/open-swe/src/graphs/shared/initialize-sandbox.ts index ea8c6199..8fe5785b 100644 --- a/apps/open-swe/src/graphs/shared/initialize-sandbox.ts +++ b/apps/open-swe/src/graphs/shared/initialize-sandbox.ts @@ -27,6 +27,7 @@ import { Sandbox } from "@daytonaio/sdk"; import { AIMessage, BaseMessage } from "@langchain/core/messages"; import { DEFAULT_SANDBOX_CREATE_PARAMS } from "../../constants.js"; import { getCustomRules } from "../../utils/custom-rules.js"; +import { withRetry } from "../../utils/retry.js"; const logger = createLogger(LogLevel.INFO, "InitializeSandbox"); @@ -258,10 +259,18 @@ export async function initializeSandbox( }, }; emitStepEvent(baseCloneRepoAction, "pending"); - const cloneRepoRes = await cloneRepo(sandbox, targetRepository, { - githubInstallationToken, - stateBranchName: branchName, - }); + + // Retry the clone command up to 3 times. Sometimes, it can timeout if the repo is large. + const cloneRepoRes = await withRetry( + async () => { + return await cloneRepo(sandbox, targetRepository, { + githubInstallationToken, + stateBranchName: branchName, + }); + }, + { retries: 3, delay: 0 }, + ); + if (cloneRepoRes.exitCode !== 0) { emitStepEvent( baseCloneRepoAction, diff --git a/apps/open-swe/src/utils/github/git.ts b/apps/open-swe/src/utils/github/git.ts index 73df3ef8..ad55bfe4 100644 --- a/apps/open-swe/src/utils/github/git.ts +++ b/apps/open-swe/src/utils/github/git.ts @@ -492,6 +492,9 @@ export async function cloneRepo( cloneResult = await sandbox.process.executeCommand( gitCloneCommand.join(" "), + undefined, + undefined, + TIMEOUT_SEC * 2, // two min timeout since large repos can take a while to clone ); if (!targetRepository.baseCommit) { @@ -514,6 +517,9 @@ export async function cloneRepo( ); const cloneDefaultBranchResult = await sandbox.process.executeCommand( cloneDefaultBranchCommand.join(" "), + undefined, + undefined, + TIMEOUT_SEC * 2, // two min timeout since large repos can take a while to clone ); if (cloneDefaultBranchResult.exitCode !== 0) { logger.error("Failed to clone default branch", { diff --git a/apps/open-swe/src/utils/retry.ts b/apps/open-swe/src/utils/retry.ts new file mode 100644 index 00000000..d5156d53 --- /dev/null +++ b/apps/open-swe/src/utils/retry.ts @@ -0,0 +1,50 @@ +interface RetryOptions { + retries?: number; + delay?: number; +} + +/** + * Executes an async function with retry logic + * @param fn - The async function to execute + * @param options - Configuration options for retry behavior + * @returns Promise that resolves with the function result or rejects with the last error + */ +export async function withRetry( + fn: () => Promise, + options: RetryOptions = {}, +): Promise { + const { retries = 3, delay = 0 } = options; + + let lastError: Error; + + for (let attempt = 0; attempt <= retries; attempt++) { + try { + return await fn(); + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); + + if (attempt === retries) { + throw lastError; + } + + if (delay > 0) { + await new Promise((resolve) => setTimeout(resolve, delay)); + } + } + } + + throw lastError!; +} + +/** + * Creates a retry wrapper for a specific function with predefined options + * @param fn - The async function to wrap + * @param options - Configuration options for retry behavior + * @returns A new function that will retry on failure + */ +export function createRetryWrapper( + fn: (...args: T) => Promise, + options: RetryOptions = {}, +): (...args: T) => Promise { + return (...args: T) => withRetry(() => fn(...args), options); +}