fix(agent-runtime): Prevent fallbacks from triggering on user abort

- Check isAbortError in catch blocks to propagate aborts immediately
- Detect both custom ABORT_ERROR_MESSAGE and native AbortError

Author: brandonkachen

packages/agent-runtime/src/llm-api/gemini-with-fallbacks.ts

@@ -1,4 +1,5 @@
 import { openaiModels, openrouterModels } from '@codebuff/common/old-constants'
+import { isAbortError, unwrapPromptResult } from '@codebuff/common/util/error'
 
 import type {
   FinetunedVertexModel,
@@ -31,7 +32,8 @@ import type { Message } from '@codebuff/common/types/messages/codebuff-message'
  * @param options.costMode - Optional cost mode ('free', 'normal', 'max') influencing fallback model choice.
  * @param options.useGPT4oInsteadOfClaude - Optional flag to use GPT-4o instead of Claude as the final fallback.
  * @returns A promise that resolves to the complete response string from the successful API call.
- * @throws If all API calls (primary and fallbacks) fail.
+ * @throws {Error} If all API calls (primary and fallbacks) fail.
+ * @throws {Error} When the request is aborted by user. Check with `isAbortError()`. Aborts are not retried.
  */
 export async function promptFlashWithFallbacks(
   params: {
@@ -56,12 +58,18 @@ export async function promptFlashWithFallbacks(
   // Try finetuned model first if enabled
   if (useFinetunedModel) {
     try {
-      return await promptAiSdk({
-        ...params,
-        messages,
-        model: useFinetunedModel,
-      })
+      return unwrapPromptResult(
+        await promptAiSdk({
+          ...params,
+          messages,
+          model: useFinetunedModel,
+        }),
+      )
     } catch (error) {
+      // Don't fall back on user-initiated aborts - propagate immediately
+      if (isAbortError(error)) {
+        throw error
+      }
       logger.warn(
         { error },
         'Error calling finetuned model, falling back to Gemini API',
@@ -71,18 +79,24 @@ export async function promptFlashWithFallbacks(
 
   try {
     // First try Gemini
-    return await promptAiSdk({ ...params, messages })
+    return unwrapPromptResult(await promptAiSdk({ ...params, messages }))
   } catch (error) {
+    // Don't fall back on user-initiated aborts - propagate immediately
+    if (isAbortError(error)) {
+      throw error
+    }
     logger.warn(
       { error },
       `Error calling Gemini API, falling back to ${useGPT4oInsteadOfClaude ? 'gpt-4o' : 'Claude'}`,
     )
-    return await promptAiSdk({
-      ...params,
-      messages,
-      model: useGPT4oInsteadOfClaude
-        ? openaiModels.gpt4o
-        : openrouterModels.openrouter_claude_3_5_haiku,
-    })
+    return unwrapPromptResult(
+      await promptAiSdk({
+        ...params,
+        messages,
+        model: useGPT4oInsteadOfClaude
+          ? openaiModels.gpt4o
+          : openrouterModels.openrouter_claude_3_5_haiku,
+      }),
+    )
   }
 }

packages/agent-runtime/src/llm-api/relace-api.ts

@@ -1,12 +1,19 @@
 import { models } from '@codebuff/common/old-constants'
 import { buildArray } from '@codebuff/common/util/array'
+import { isAbortError, unwrapPromptResult } from '@codebuff/common/util/error'
 import { parseMarkdownCodeBlock } from '@codebuff/common/util/file'
 import { assistantMessage, userMessage } from '@codebuff/common/util/messages'
 
 import type { PromptAiSdkFn } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 
+/**
+ * Applies code edits using Relace AI, with fallback to o3-mini on failure.
+ *
+ * @returns The updated code with edits applied.
+ * @throws {Error} When the request is aborted by user. Check with `isAbortError()`. Aborts are not retried.
+ */
 export async function promptRelaceAI(
   params: {
     initialCode: string
@@ -21,24 +28,30 @@ export async function promptRelaceAI(
   try {
     const { tools: _tools, ...rest } = params
     // const model = 'relace-apply-2.5-lite'
-    const content = await promptAiSdk({
-      ...rest,
-      model: 'relace/relace-apply-3',
-      messages: [
-        userMessage(
-          buildArray(
-            instructions && `<instruction>${instructions}</instruction>`,
-            `<code>${initialCode}</code>`,
-            `<update>${editSnippet}</update>`,
-          ).join('\n'),
-        ),
-      ],
-      system: undefined,
-      includeCacheControl: false,
-    })
-
-    return content + '\n'
+    return (
+      unwrapPromptResult(
+        await promptAiSdk({
+          ...rest,
+          model: 'relace/relace-apply-3',
+          messages: [
+            userMessage(
+              buildArray(
+                instructions && `<instruction>${instructions}</instruction>`,
+                `<code>${initialCode}</code>`,
+                `<update>${editSnippet}</update>`,
+              ).join('\n'),
+            ),
+          ],
+          system: undefined,
+          includeCacheControl: false,
+        }),
+      ) + '\n'
+    )
   } catch (error) {
+    // Don't fall back on user-initiated aborts - propagate immediately
+    if (isAbortError(error)) {
+      throw error
+    }
     logger.error(
       {
         error:
@@ -49,7 +62,7 @@ export async function promptRelaceAI(
       'Error calling Relace AI, falling back to o3-mini',
     )
 
-    // Fall back to Gemini
+    // Fall back to o3-mini
     const prompt = `You are an expert programmer. Please rewrite this code file to implement the edit snippet while preserving as much of the original code and behavior as possible.
 
 Initial code:
@@ -70,12 +83,16 @@ Important:
 
 Please output just the complete updated file content with no other text.`
 
-    const content = await promptAiSdk({
-      ...params,
-      messages: [userMessage(prompt), assistantMessage('```\n')],
-      model: models.o3mini,
-    })
-
-    return parseMarkdownCodeBlock(content) + '\n'
+    return (
+      parseMarkdownCodeBlock(
+        unwrapPromptResult(
+          await promptAiSdk({
+            ...params,
+            messages: [userMessage(prompt), assistantMessage('```\n')],
+            model: models.o3mini,
+          }),
+        ),
+      ) + '\n'
+    )
   }
 }