fix(auth): surface OAuth setup failures

Amp-Thread-ID: https://ampcode.com/threads/T-019fb8cf-a980-7568-80f1-fe34faaaecae
Co-authored-by: Pierre Bruno <pierrebruno@hotmail.ch>

Author: ampagent

internal/cli/provider_onboarding.go

@@ -509,13 +509,16 @@ func runProvidersRemove(args []string, stdout io.Writer, stderr io.Writer, deps
 	// non-default config path cannot leave the encrypted key behind.
 	keyRemoved := false
 	var keyErr error
+	markerTransferFailed := false
 	if survivor, ok := caseVariantProviderName(cfg, name); ok {
 		if hadBefore && before.APIKeyStored {
 			if survivorRow, foundSurvivor, err := config.ProviderRow(configPath, survivor); err != nil {
 				keyErr = err
+				markerTransferFailed = true
 			} else if foundSurvivor && !survivorRow.APIKeyStored {
 				if _, err := config.TransferProviderAPIKeyStoredMarker(configPath, survivor); err != nil {
 					keyErr = err
+					markerTransferFailed = true
 				}
 			}
 		}
@@ -542,7 +545,11 @@ func runProvidersRemove(args []string, stdout io.Writer, stderr io.Writer, deps
 		return exitCrash
 	}
 	if keyErr != nil {
-		if _, err := fmt.Fprintf(stderr, "warning: its stored API key could not be deleted and remains in the credential store: %v\n", keyErr); err != nil {
+		warning := "its stored API key could not be deleted and remains in the credential store"
+		if markerTransferFailed {
+			warning = "the stored API key marker could not be transferred to the surviving case-variant provider, so the shared key may be unreachable"
+		}
+		if _, err := fmt.Fprintf(stderr, "warning: %s: %v\n", warning, keyErr); err != nil {
 			return exitCrash
 		}
 	} else if keyRemoved {

internal/tui/onboarding.go

@@ -575,8 +575,11 @@ type setupOAuthDeviceMsg struct {
 	err        error
 }
 
-func setupDevicePrepareCmd(name string) tea.Cmd {
+func setupDevicePrepareCmd(name string, configPath ...string) tea.Cmd {
 	return func() tea.Msg {
+		if err := preflightOAuthUserConfig(firstString(configPath)); err != nil {
+			return setupOAuthDeviceMsg{providerID: name, err: err}
+		}
 		auth, cfg, err := oauthDevicePrepare(name)
 		if err != nil {
 			return setupOAuthDeviceMsg{providerID: name, err: err}
@@ -612,7 +615,7 @@ func (m model) startSetupDeviceLogin(descriptor providercatalog.Descriptor) (tea
 	m.setup.oauthErr = ""
 	m.setup.deviceUserCode = ""
 	m.setup.deviceVerificationURI = ""
-	return m, setupDevicePrepareCmd(descriptor.ID)
+	return m, setupDevicePrepareCmd(descriptor.ID, m.setup.configPath)
 }
 
 // applySetupOAuthDeviceCode handles phase 1 of device-code login: show the code,
@@ -660,9 +663,12 @@ func (m model) applySetupOAuth(msg setupOAuthMsg) (tea.Model, tea.Cmd) {
 	if msg.tokenLogin {
 		// Early persist on the Update goroutine (see persistOAuthLoginProvider's
 		// threading contract) so quitting setup after the login doesn't lose it;
-		// completeSetup persists the full profile with the chosen model anyway,
-		// so a failure here is recoverable and not fatal to setup.
-		_ = persistOAuthLoginProvider(m.setup.configPath, msg.providerID)
+		// do not advance when that write fails or the stored token would have no
+		// reachable provider profile after setup exits.
+		if err := persistOAuthLoginProvider(m.setup.configPath, msg.providerID); err != nil {
+			m.setup.oauthErr = "Signed in, but the provider profile could not be saved: " + redaction.ErrorMessage(err, redaction.Options{})
+			return m, nil
+		}
 	}
 	m.setup.oauthErr = ""
 	m.setup.err = ""
@@ -882,6 +888,9 @@ func (m model) completeSetup() (tea.Model, tea.Cmd) {
 	})
 	if err != nil {
 		m.setup.err = err.Error()
+		if m.setup.oauthMode && m.setupProviderDescriptor().OAuthMintsKey && strings.TrimSpace(apiKey) != "" {
+			m.setup.err = oauthMintedKeySaveError(m.setup.err, apiKey)
+		}
 		return m, nil
 	}
 
@@ -911,6 +920,10 @@ func (m model) completeSetup() (tea.Model, tea.Cmd) {
 	return m.exitSetupToChat()
 }
 
+func oauthMintedKeySaveError(message, apiKey string) string {
+	return fmt.Sprintf("%s\nThe minted API key was not saved. Copy it now before leaving Zero:\n%s", message, apiKey)
+}
+
 func (m *model) resetSetupModels() {
 	option := m.setupProvider()
 	provider := setupProviderDescriptor(option)

internal/tui/onboarding_test.go

@@ -2066,6 +2066,61 @@ func TestApplySetupOAuthSuccessAdvancesToModel(t *testing.T) {
 	}
 }
 
+func TestApplySetupOAuthTokenPersistFailureStaysOnProvider(t *testing.T) {
+	m := newModel(context.Background(), Options{
+		Setup: SetupOptions{Visible: true, Providers: []SetupProviderOption{
+			{ID: "xai", Name: "xAI", RequiresAuth: true},
+		}},
+	})
+	m.setup.stage = setupStageProvider
+	m.setup.oauthPending = true
+	m.setup.oauthMode = true
+	m.setup.configPath = t.TempDir() // A directory cannot be read as config.json.
+
+	updated, cmd := m.applySetupOAuth(setupOAuthMsg{tokenLogin: true, providerID: "xai"})
+	next := updated.(model)
+	if cmd != nil {
+		t.Fatal("failed profile persistence should not start model discovery")
+	}
+	if next.setup.stage != setupStageProvider {
+		t.Fatalf("stage = %v, want provider", next.setup.stage)
+	}
+	if next.setup.oauthErr == "" || !strings.Contains(next.setup.oauthErr, "could not be saved") {
+		t.Fatalf("oauthErr = %q, want profile-save failure", next.setup.oauthErr)
+	}
+}
+
+func TestCompleteSetupSurfacesMintedOpenRouterKeyWhenSaveFails(t *testing.T) {
+	const mintedKey = "sk-or-minted-recovery"
+	m := newModel(context.Background(), Options{
+		Setup: SetupOptions{
+			Visible: true,
+			Providers: []SetupProviderOption{
+				{ID: "openrouter", Name: "OpenRouter", RequiresAuth: true},
+			},
+			Save: func(SetupSelection) (SetupResult, error) {
+				return SetupResult{}, errors.New("config write failed")
+			},
+		},
+	})
+	m.setup.oauthMode = true
+	m.setup.stage = setupStageReady
+	m.setup.apiKey.SetValue(mintedKey)
+
+	updated, _ := m.completeSetup()
+	next := updated.(model)
+	if !strings.Contains(next.setup.err, mintedKey) || !strings.Contains(next.setup.err, "was not saved") {
+		t.Fatalf("setup error = %q, want minted-key recovery", next.setup.err)
+	}
+}
+
+func TestSetupDevicePreparePreflightsConfigBeforeRequestingCode(t *testing.T) {
+	msg := setupDevicePrepareCmd("xai", t.TempDir())().(setupOAuthDeviceMsg)
+	if msg.err == nil {
+		t.Fatal("device-code preparation should reject an unreadable config before requesting a code")
+	}
+}
+
 func pressSetupContinue(m model) model {
 	m = pressSetupContinueOnce(m)
 	// Transparently skip the connect-method chooser via the API-key/browse path so

internal/tui/provider_wizard.go

@@ -327,8 +327,11 @@ type providerWizardDeviceCodeMsg struct {
 
 // providerWizardDevicePrepareCmd runs phase 1 of the device-code login off the UI
 // goroutine and reports the code to display (or an error).
-func providerWizardDevicePrepareCmd(name string, attemptID int) tea.Cmd {
+func providerWizardDevicePrepareCmd(name string, attemptID int, configPath ...string) tea.Cmd {
 	return func() tea.Msg {
+		if err := preflightOAuthUserConfig(firstString(configPath)); err != nil {
+			return providerWizardDeviceCodeMsg{providerID: name, attemptID: attemptID, err: err}
+		}
 		auth, cfg, err := oauthDevicePrepare(name)
 		if err != nil {
 			return providerWizardDeviceCodeMsg{providerID: name, attemptID: attemptID, err: err}
@@ -364,7 +367,7 @@ func (m model) startProviderDeviceLogin() (model, tea.Cmd) {
 		return m, nil
 	}
 	attemptID := m.providerWizard.beginOAuthAttempt(true)
-	return m, providerWizardDevicePrepareCmd(provider.ID, attemptID)
+	return m, providerWizardDevicePrepareCmd(provider.ID, attemptID, m.userConfigPath)
 }
 
 const maxProviderWizardProvidersVisible = 10
@@ -1300,11 +1303,17 @@ func (m model) applyProviderWizard() (model, tea.Cmd) {
 		adopted, adoptErr := config.AdoptPersistedCatalogProviderName(m.userConfigPath, profile)
 		if adoptErr != nil {
 			wizard.err = adoptErr.Error()
+			if provider.OAuthMintsKey && strings.TrimSpace(wizard.apiKey) != "" {
+				wizard.err = oauthMintedKeySaveError(wizard.err, wizard.apiKey)
+			}
 			return m, nil
 		}
 		profile = adopted
 		if err := config.PreflightProviderWrite(m.userConfigPath, profile.Name); err != nil {
 			wizard.err = err.Error()
+			if provider.OAuthMintsKey && strings.TrimSpace(wizard.apiKey) != "" {
+				wizard.err = oauthMintedKeySaveError(wizard.err, wizard.apiKey)
+			}
 			return m, nil
 		}
 		// Capture flip: move the freshly entered key into the encrypted credential
@@ -1316,6 +1325,9 @@ func (m model) applyProviderWizard() (model, tea.Cmd) {
 		}
 		if _, err := config.UpsertProvider(m.userConfigPath, profile, true); err != nil {
 			wizard.err = redaction.RedactString(err.Error(), redaction.Options{ExtraSecretValues: []string{secret, profile.APIKey}})
+			if provider.OAuthMintsKey && strings.TrimSpace(secret) != "" {
+				wizard.err = oauthMintedKeySaveError(wizard.err, secret)
+			}
 			return m, nil // nothing committed to live state yet
 		}
 	}

internal/tui/provider_wizard_oauth_test.go

@@ -364,6 +364,29 @@ func TestApplyProviderWizardOAuthIgnoresStaleAttempt(t *testing.T) {
 	}
 }
 
+func TestProviderWizardSurfacesMintedOpenRouterKeyWhenSaveFails(t *testing.T) {
+	const mintedKey = "sk-or-wizard-recovery"
+	m := wizardModelAt(t, "openrouter", providerWizardStepDone)
+	m.providerWizard.oauthMode = true
+	m.providerWizard.apiKey = mintedKey
+	m.userConfigPath = t.TempDir() // A directory cannot be read as config.json.
+
+	next, _ := m.applyProviderWizard()
+	if !strings.Contains(next.providerWizard.err, mintedKey) || !strings.Contains(next.providerWizard.err, "was not saved") {
+		t.Fatalf("wizard error = %q, want minted-key recovery", next.providerWizard.err)
+	}
+}
+
+func TestProviderWizardDevicePreparePreflightsConfigBeforeRequestingCode(t *testing.T) {
+	msg := providerWizardDevicePrepareCmd("xai", 42, t.TempDir())().(providerWizardDeviceCodeMsg)
+	if msg.err == nil {
+		t.Fatal("device-code preparation should reject an unreadable config before requesting a code")
+	}
+	if msg.attemptID != 42 {
+		t.Fatalf("attemptID = %d, want 42", msg.attemptID)
+	}
+}
+
 func TestProviderWizardDeviceCodeIgnoresStaleAttempt(t *testing.T) {
 	m := mouseTestModel()
 	m.providerWizard = m.newProviderWizard()