From f901884913540a6eb811773e2233789c90319c8d Mon Sep 17 00:00:00 2001
From: Charles-Henri ROBICHE <chrobiche@gmail.com>
Date: Thu, 29 Jan 2026 10:31:11 +0100
Subject: [PATCH] fix(litellm): improve retry flow for API key and base URL

- Extract promptForApiKey and promptForBaseUrl as helper functions
- Use recursive retry instead of throwing errors
- Re-enter API key now retries the entire flow with new key
- Re-enter base URL now properly prompts for new URL and retries
- Both options maintain full onboarding flow instead of crashing

Co-Authored-By: Claude <noreply@anthropic.com>
---
 .../auth-choice.apply.api-providers.ts        | 55 ++++++++++++-------
 1 file changed, 35 insertions(+), 20 deletions(-)

diff --git a/src/commands/auth-choice.apply.api-providers.ts b/src/commands/auth-choice.apply.api-providers.ts
index f5f61bb93..ef3965915 100644
--- a/src/commands/auth-choice.apply.api-providers.ts
+++ b/src/commands/auth-choice.apply.api-providers.ts
@@ -627,20 +627,17 @@ export async function applyAuthChoiceApiProviders(
       }
     }
 
-    if (!hasCredential) {
+    // Helper function to prompt for API key
+    const promptForApiKey = async () => {
       const key = await params.prompter.text({
         message: "Enter LiteLLM API key",
         validate: validateApiKeyInput,
       });
-      apiKey = normalizeApiKeyInput(String(key));
-      await setLitellmApiKey(apiKey, params.agentDir);
-    }
+      return normalizeApiKeyInput(String(key));
+    };
 
-    // Check for pre-provided base URL via CLI option (--litellm-base-url)
-    let normalizedBaseUrl: string;
-    if (params.opts?.litellmBaseUrl) {
-      normalizedBaseUrl = params.opts.litellmBaseUrl.trim();
-    } else {
+    // Helper function to prompt for base URL
+    const promptForBaseUrl = async () => {
       const defaultBaseUrl = process.env.LITELLM_BASE_URL ?? "http://localhost:4000";
       const baseUrl = await params.prompter.text({
         message: "Enter LiteLLM base URL",
@@ -656,7 +653,20 @@ export async function applyAuthChoiceApiProviders(
           }
         },
       });
-      normalizedBaseUrl = String(baseUrl).trim();
+      return String(baseUrl).trim();
+    };
+
+    if (!hasCredential) {
+      apiKey = await promptForApiKey();
+      await setLitellmApiKey(apiKey, params.agentDir);
+    }
+
+    // Check for pre-provided base URL via CLI option (--litellm-base-url)
+    let normalizedBaseUrl: string;
+    if (params.opts?.litellmBaseUrl) {
+      normalizedBaseUrl = params.opts.litellmBaseUrl.trim();
+    } else {
+      normalizedBaseUrl = await promptForBaseUrl();
     }
 
     // Try to fetch available models from LiteLLM
@@ -794,20 +804,25 @@ export async function applyAuthChoiceApiProviders(
       }
 
       if (action === "retry-apikey") {
-        // Re-prompt for API key
-        const key = await params.prompter.text({
-          message: "Enter LiteLLM API key",
-          validate: validateApiKeyInput,
-        });
-        apiKey = normalizeApiKeyInput(String(key));
+        // Re-prompt for API key and retry
+        apiKey = await promptForApiKey();
         await setLitellmApiKey(apiKey, params.agentDir);
-
-        // Retry the entire LiteLLM flow by returning null and letting the caller retry
-        throw new Error("Please try the LiteLLM setup again with the new API key");
+        // Retry fetch by recursively calling this function
+        return await applyAuthChoiceApiProviders({ ...params, authChoice: "litellm-api-key" });
       }
 
       if (action === "retry-baseurl") {
-        throw new Error("Please restart the setup and provide the correct base URL");
+        // Re-prompt for base URL and retry the entire flow
+        // This ensures we go through the full fetch process again with the new URL
+        const newParams = {
+          ...params,
+          authChoice: "litellm-api-key" as const,
+          opts: {
+            ...params.opts,
+            litellmBaseUrl: undefined, // Clear the CLI-provided URL so we can prompt
+          },
+        };
+        return await applyAuthChoiceApiProviders(newParams);
       }
 
       // Manual model entry