From f3759381b1666e26eda55d1f02652c1df1bb2b3d Mon Sep 17 00:00:00 2001
From: anthony bushong <agmsb@users.noreply.github.com>
Date: Mon, 3 Nov 2025 16:04:51 -0800
Subject: [PATCH] feat(core): add timeout to llm edit fix (#12393)

---
 packages/core/src/tools/smart-edit.ts         | 12 ++++
 .../core/src/utils/llm-edit-fixer.test.ts     | 43 +++++++++++++++
 packages/core/src/utils/llm-edit-fixer.ts     | 55 +++++++++++++++----
 3 files changed, 99 insertions(+), 11 deletions(-)

diff --git a/packages/core/src/tools/smart-edit.ts b/packages/core/src/tools/smart-edit.ts
index cb7e21f121..444c893bd0 100644
--- a/packages/core/src/tools/smart-edit.ts
+++ b/packages/core/src/tools/smart-edit.ts
@@ -426,6 +426,18 @@ class EditToolInvocation
       abortSignal,
     );
 
+    // If the self-correction attempt timed out, return the original error.
+    if (fixedEdit === null) {
+      return {
+        currentContent: contentForLlmEditFixer,
+        newContent: currentContent,
+        occurrences: 0,
+        isNewFile: false,
+        error: initialError,
+        originalLineEnding,
+      };
+    }
+
     if (fixedEdit.noChangesRequired) {
       return {
         currentContent,
diff --git a/packages/core/src/utils/llm-edit-fixer.test.ts b/packages/core/src/utils/llm-edit-fixer.test.ts
index 222bf2289b..9caad5ee10 100644
--- a/packages/core/src/utils/llm-edit-fixer.test.ts
+++ b/packages/core/src/utils/llm-edit-fixer.test.ts
@@ -29,6 +29,7 @@ describe('FixLLMEditWithInstruction', () => {
   const abortSignal = abortController.signal;
 
   beforeEach(() => {
+    vi.useFakeTimers();
     vi.clearAllMocks();
     resetLlmEditFixerCaches_TEST_ONLY(); // Ensure cache is cleared before each test
   });
@@ -319,4 +320,46 @@ describe('FixLLMEditWithInstruction', () => {
       });
     });
   });
+
+  it(
+    'should return null if the LLM call times out',
+    { timeout: 60000 },
+    async () => {
+      mockGenerateJson.mockImplementation(
+        async ({ abortSignal }) =>
+          // Simulate a long-running operation that never resolves on its own.
+          // It will only reject when the abort signal is triggered by the timeout.
+          new Promise((_resolve, reject) => {
+            if (abortSignal?.aborted) {
+              return reject(new DOMException('Aborted', 'AbortError'));
+            }
+            abortSignal?.addEventListener('abort', () => {
+              reject(new DOMException('Aborted', 'AbortError'));
+            });
+          }),
+      );
+
+      const testPromptId = 'test-prompt-id-timeout';
+
+      const fixPromise = promptIdContext.run(testPromptId, () =>
+        FixLLMEditWithInstruction(
+          instruction,
+          old_string,
+          new_string,
+          error,
+          current_content,
+          mockBaseLlmClient,
+          abortSignal,
+        ),
+      );
+
+      // Let the timers advance just past the 40000ms default timeout.
+      await vi.advanceTimersByTimeAsync(40001);
+
+      const result = await fixPromise;
+
+      expect(result).toBeNull();
+      expect(mockGenerateJson).toHaveBeenCalledOnce();
+    },
+  );
 });
diff --git a/packages/core/src/utils/llm-edit-fixer.ts b/packages/core/src/utils/llm-edit-fixer.ts
index 93c66a21c5..6c8db6a0d2 100644
--- a/packages/core/src/utils/llm-edit-fixer.ts
+++ b/packages/core/src/utils/llm-edit-fixer.ts
@@ -13,6 +13,7 @@ import { promptIdContext } from './promptIdContext.js';
 import { debugLogger } from './debugLogger.js';
 
 const MAX_CACHE_SIZE = 50;
+const GENERATE_JSON_TIMEOUT_MS = 40000; // 40 seconds
 
 const EDIT_SYS_PROMPT = `
 You are an expert code-editing assistant specializing in debugging and correcting failed search-and-replace operations.
@@ -89,6 +90,32 @@ const editCorrectionWithInstructionCache = new LruCache<
   SearchReplaceEdit
 >(MAX_CACHE_SIZE);
 
+async function generateJsonWithTimeout<T>(
+  client: BaseLlmClient,
+  params: Parameters<BaseLlmClient['generateJson']>[0],
+  timeoutMs: number,
+): Promise<T | null> {
+  try {
+    // Create a signal that aborts automatically after the specified timeout.
+    const timeoutSignal = AbortSignal.timeout(timeoutMs);
+
+    const result = await client.generateJson({
+      ...params,
+      // The operation will be aborted if either the original signal is aborted
+      // or if the timeout is reached.
+      abortSignal: AbortSignal.any([
+        params.abortSignal ?? new AbortController().signal,
+        timeoutSignal,
+      ]),
+    });
+    return result as T;
+  } catch (_err) {
+    // An AbortError will be thrown on timeout.
+    // We catch it and return null to signal that the operation timed out.
+    return null;
+  }
+}
+
 /**
  * Attempts to fix a failed edit by using an LLM to generate a new search and replace pair.
  * @param instruction The instruction for what needs to be done.
@@ -109,7 +136,7 @@ export async function FixLLMEditWithInstruction(
   current_content: string,
   baseLlmClient: BaseLlmClient,
   abortSignal: AbortSignal,
-): Promise<SearchReplaceEdit> {
+): Promise<SearchReplaceEdit | null> {
   let promptId = promptIdContext.getStore();
   if (!promptId) {
     promptId = `llm-fixer-fallback-${Date.now()}-${Math.random().toString(16).slice(2)}`;
@@ -146,17 +173,23 @@ export async function FixLLMEditWithInstruction(
     },
   ];
 
-  const result = (await baseLlmClient.generateJson({
-    contents,
-    schema: SearchReplaceEditSchema,
-    abortSignal,
-    model: DEFAULT_GEMINI_FLASH_MODEL,
-    systemInstruction: EDIT_SYS_PROMPT,
-    promptId,
-    maxAttempts: 1,
-  })) as unknown as SearchReplaceEdit;
+  const result = await generateJsonWithTimeout<SearchReplaceEdit>(
+    baseLlmClient,
+    {
+      contents,
+      schema: SearchReplaceEditSchema,
+      abortSignal,
+      model: DEFAULT_GEMINI_FLASH_MODEL,
+      systemInstruction: EDIT_SYS_PROMPT,
+      promptId,
+      maxAttempts: 1,
+    },
+    GENERATE_JSON_TIMEOUT_MS,
+  );
 
-  editCorrectionWithInstructionCache.set(cacheKey, result);
+  if (result) {
+    editCorrectionWithInstructionCache.set(cacheKey, result);
+  }
   return result;
 }