From ce8fd694aeda228d28bed154ce88ef533ca9fada Mon Sep 17 00:00:00 2001
From: Charles-Henri ROBICHE <chrobiche@gmail.com>
Date: Thu, 29 Jan 2026 10:59:30 +0100
Subject: [PATCH] fix(litellm): handle AUTH_CHOICE_CANCELLED to retry auth
 selection

- Add retry loops in configure.gateway-auth and agents.commands.add
- Catch AUTH_CHOICE_CANCELLED error and re-prompt for auth method
- Users can now go back to auth method selection when LiteLLM setup fails
- Import AuthChoice type for proper typing

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/commands/agents.commands.add.ts    | 52 ++++++++++++--------
 src/commands/configure.gateway-auth.ts | 67 ++++++++++++++++----------
 2 files changed, 75 insertions(+), 44 deletions(-)

diff --git a/src/commands/agents.commands.add.ts b/src/commands/agents.commands.add.ts
index ee99fcdfe..ea977d057 100644
--- a/src/commands/agents.commands.add.ts
+++ b/src/commands/agents.commands.add.ts
@@ -254,27 +254,41 @@ export async function agentsAddCommand(
       const authStore = ensureAuthProfileStore(agentDir, {
         allowKeychainPrompt: false,
       });
-      const authChoice = await promptAuthChoiceGrouped({
-        prompter,
-        store: authStore,
-        includeSkip: true,
-      });
 
-      const authResult = await applyAuthChoice({
-        authChoice,
-        config: nextConfig,
-        prompter,
-        runtime,
-        agentDir,
-        setDefaultModel: false,
-        agentId,
-      });
-      nextConfig = authResult.config;
-      if (authResult.agentModelOverride) {
-        nextConfig = applyAgentConfig(nextConfig, {
-          agentId,
-          model: authResult.agentModelOverride,
+      // Loop to allow retrying auth choice if user cancels during configuration
+      while (true) {
+        const authChoice = await promptAuthChoiceGrouped({
+          prompter,
+          store: authStore,
+          includeSkip: true,
         });
+
+        try {
+          const authResult = await applyAuthChoice({
+            authChoice,
+            config: nextConfig,
+            prompter,
+            runtime,
+            agentDir,
+            setDefaultModel: false,
+            agentId,
+          });
+          nextConfig = authResult.config;
+          if (authResult.agentModelOverride) {
+            nextConfig = applyAgentConfig(nextConfig, {
+              agentId,
+              model: authResult.agentModelOverride,
+            });
+          }
+          break; // Success - exit the loop
+        } catch (error) {
+          // If user cancelled to go back to auth selection, loop again
+          if (error instanceof Error && error.message === "AUTH_CHOICE_CANCELLED") {
+            continue;
+          }
+          // Re-throw other errors
+          throw error;
+        }
       }
     }
 
diff --git a/src/commands/configure.gateway-auth.ts b/src/commands/configure.gateway-auth.ts
index e4e72d5ec..d0cc6d56d 100644
--- a/src/commands/configure.gateway-auth.ts
+++ b/src/commands/configure.gateway-auth.ts
@@ -4,6 +4,7 @@ import type { RuntimeEnv } from "../runtime.js";
 import type { WizardPrompter } from "../wizard/prompts.js";
 import { applyAuthChoice, resolvePreferredProviderForAuthChoice } from "./auth-choice.js";
 import { promptAuthChoiceGrouped } from "./auth-choice-prompt.js";
+import type { AuthChoice } from "./onboard-types.js";
 import {
   applyModelAllowlist,
   applyModelFallbacksFromSelection,
@@ -41,34 +42,50 @@ export async function promptAuthConfig(
   runtime: RuntimeEnv,
   prompter: WizardPrompter,
 ): Promise<MoltbotConfig> {
-  const authChoice = await promptAuthChoiceGrouped({
-    prompter,
-    store: ensureAuthProfileStore(undefined, {
-      allowKeychainPrompt: false,
-    }),
-    includeSkip: true,
-  });
-
   let next = cfg;
-  if (authChoice !== "skip") {
-    const applied = await applyAuthChoice({
-      authChoice,
-      config: next,
+  let authChoice: AuthChoice;
+
+  // Loop to allow retrying auth choice if user cancels during configuration
+  while (true) {
+    authChoice = await promptAuthChoiceGrouped({
       prompter,
-      runtime,
-      setDefaultModel: true,
+      store: ensureAuthProfileStore(undefined, {
+        allowKeychainPrompt: false,
+      }),
+      includeSkip: true,
     });
-    next = applied.config;
-  } else {
-    const modelSelection = await promptDefaultModel({
-      config: next,
-      prompter,
-      allowKeep: true,
-      ignoreAllowlist: true,
-      preferredProvider: resolvePreferredProviderForAuthChoice(authChoice),
-    });
-    if (modelSelection.model) {
-      next = applyPrimaryModel(next, modelSelection.model);
+
+    if (authChoice !== "skip") {
+      try {
+        const applied = await applyAuthChoice({
+          authChoice,
+          config: next,
+          prompter,
+          runtime,
+          setDefaultModel: true,
+        });
+        next = applied.config;
+        break; // Success - exit the loop
+      } catch (error) {
+        // If user cancelled to go back to auth selection, loop again
+        if (error instanceof Error && error.message === "AUTH_CHOICE_CANCELLED") {
+          continue;
+        }
+        // Re-throw other errors
+        throw error;
+      }
+    } else {
+      const modelSelection = await promptDefaultModel({
+        config: next,
+        prompter,
+        allowKeep: true,
+        ignoreAllowlist: true,
+        preferredProvider: resolvePreferredProviderForAuthChoice(authChoice),
+      });
+      if (modelSelection.model) {
+        next = applyPrimaryModel(next, modelSelection.model);
+      }
+      break; // Skip selected - exit the loop
     }
   }