feat: Add withRetry func and longer timeouts for cloning (#449)

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<string>>()
+      .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<string>>()
+      .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<string>>()
+      .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<string>>()
+      .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<string>>()
+      .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<string>>()
+      .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<string>>()
+      .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<string>>()
+      .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<number>>().mockResolvedValue(42);
+    const objectFn = jest
+      .fn<() => Promise<{ key: string }>>()
+      .mockResolvedValue({ key: "value" });
+    const arrayFn = jest
+      .fn<() => Promise<number[]>>()
+      .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<string>>()
+      .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<string>>()
+      .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<string>>()
+      .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<string>>()
+      .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");
+  });
+});

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,

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", {

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<T>(
+  fn: () => Promise<T>,
+  options: RetryOptions = {},
+): Promise<T> {
+  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<T extends any[], R>(
+  fn: (...args: T) => Promise<R>,
+  options: RetryOptions = {},
+): (...args: T) => Promise<R> {
+  return (...args: T) => withRetry(() => fn(...args), options);
+}