refactor(auth): save authType after successfully authenticated (#1036)

2025-12-20 16:57:46 +00:00 · 2025-11-19 11:21:46 +08:00
parent 3ed93d5b5d
commit d0e76c76a8
30 changed files with 822 additions and 518 deletions
--- a/packages/cli/src/ui/components/DialogManager.tsx
+++ b/packages/cli/src/ui/components/DialogManager.tsx
@@ -12,9 +12,9 @@ import { ShellConfirmationDialog } from './ShellConfirmationDialog.js';
 import { ConsentPrompt } from './ConsentPrompt.js';
 import { ThemeDialog } from './ThemeDialog.js';
 import { SettingsDialog } from './SettingsDialog.js';
-import { AuthInProgress } from '../auth/AuthInProgress.js';
 import { QwenOAuthProgress } from './QwenOAuthProgress.js';
 import { AuthDialog } from '../auth/AuthDialog.js';
+import { OpenAIKeyPrompt } from './OpenAIKeyPrompt.js';
 import { EditorSettingsDialog } from './EditorSettingsDialog.js';
 import { WorkspaceMigrationDialog } from './WorkspaceMigrationDialog.js';
 import { ProQuotaDialog } from './ProQuotaDialog.js';
@@ -26,6 +26,9 @@ import { useUIState } from '../contexts/UIStateContext.js';
 import { useUIActions } from '../contexts/UIActionsContext.js';
 import { useConfig } from '../contexts/ConfigContext.js';
 import { useSettings } from '../contexts/SettingsContext.js';
+import { SettingScope } from '../../config/settings.js';
+import { AuthState } from '../types.js';
+import { AuthType } from '@qwen-code/qwen-code-core';
 import process from 'node:process';
 import { type UseHistoryManagerReturn } from '../hooks/useHistoryManager.js';
 import { IdeTrustChangeDialog } from './IdeTrustChangeDialog.js';
@@ -56,6 +59,16 @@ export const DialogManager = ({
  const { constrainHeight, terminalHeight, staticExtraHeight, mainAreaWidth } =
    uiState;

+  const getDefaultOpenAIConfig = () => {
+    const fromSettings = settings.merged.security?.auth;
+    const modelSettings = settings.merged.model;
+    return {
+      apiKey: fromSettings?.apiKey || process.env['OPENAI_API_KEY'] || '',
+      baseUrl: fromSettings?.baseUrl || process.env['OPENAI_BASE_URL'] || '',
+      model: modelSettings?.name || process.env['OPENAI_MODEL'] || '',
+    };
+  };
+
  if (uiState.showWelcomeBackDialog && uiState.welcomeBackInfo?.hasHistory) {
    return (
      <WelcomeBackDialog
@@ -207,39 +220,56 @@ export const DialogManager = ({
  if (uiState.isVisionSwitchDialogOpen) {
    return <ModelSwitchDialog onSelect={uiActions.handleVisionSwitchSelect} />;
  }
+
+  if (uiState.isAuthDialogOpen || uiState.authError) {
+    return (
+      <Box flexDirection="column">
+        <AuthDialog />
+      </Box>
+    );
+  }
+
  if (uiState.isAuthenticating) {
-    // Show Qwen OAuth progress if it's Qwen auth and OAuth is active
-    if (uiState.isQwenAuth && uiState.isQwenAuthenticating) {
+    if (uiState.pendingAuthType === AuthType.USE_OPENAI) {
+      const defaults = getDefaultOpenAIConfig();
      return (
-        <QwenOAuthProgress
-          deviceAuth={uiState.deviceAuth || undefined}
-          authStatus={uiState.authStatus}
-          authMessage={uiState.authMessage}
-          onTimeout={uiActions.handleQwenAuthTimeout}
-          onCancel={uiActions.handleQwenAuthCancel}
+        <OpenAIKeyPrompt
+          onSubmit={(apiKey, baseUrl, model) => {
+            uiActions.handleAuthSelect(AuthType.USE_OPENAI, SettingScope.User, {
+              apiKey,
+              baseUrl,
+              model,
+            });
+          }}
+          onCancel={() => {
+            uiActions.cancelAuthentication();
+            uiActions.setAuthState(AuthState.Updating);
+          }}
+          defaultApiKey={defaults.apiKey}
+          defaultBaseUrl={defaults.baseUrl}
+          defaultModel={defaults.model}
        />
      );
    }

-    // Default auth progress for other auth types
-    return (
-      <AuthInProgress
-        onTimeout={() => {
-          uiActions.onAuthError('Authentication cancelled.');
-        }}
-      />
-    );
-  }
-  if (uiState.isAuthDialogOpen) {
-    return (
-      <Box flexDirection="column">
-        <AuthDialog
-          onSelect={uiActions.handleAuthSelect}
-          settings={settings}
-          initialErrorMessage={uiState.authError}
+    if (uiState.pendingAuthType === AuthType.QWEN_OAUTH) {
+      return (
+        <QwenOAuthProgress
+          deviceAuth={uiState.qwenAuthState.deviceAuth || undefined}
+          authStatus={uiState.qwenAuthState.authStatus}
+          authMessage={uiState.qwenAuthState.authMessage}
+          onTimeout={() => {
+            uiActions.onAuthError('Qwen OAuth authentication timed out.');
+            uiActions.cancelAuthentication();
+            uiActions.setAuthState(AuthState.Updating);
+          }}
+          onCancel={() => {
+            uiActions.cancelAuthentication();
+            uiActions.setAuthState(AuthState.Updating);
+          }}
        />
-      </Box>
-    );
+      );
+    }
  }
  if (uiState.isEditorDialogOpen) {
    return (
--- a/packages/cli/src/ui/components/OpenAIKeyPrompt.tsx
+++ b/packages/cli/src/ui/components/OpenAIKeyPrompt.tsx
@@ -6,6 +6,7 @@

 import type React from 'react';
 import { useState } from 'react';
+import { z } from 'zod';
 import { Box, Text } from 'ink';
 import { Colors } from '../colors.js';
 import { useKeypress } from '../hooks/useKeypress.js';
@@ -18,6 +19,16 @@ interface OpenAIKeyPromptProps {
  defaultModel?: string;
 }

+export const credentialSchema = z.object({
+  apiKey: z.string().min(1, 'API key is required'),
+  baseUrl: z
+    .union([z.string().url('Base URL must be a valid URL'), z.literal('')])
+    .optional(),
+  model: z.string().min(1, 'Model must be a non-empty string').optional(),
+});
+
+export type OpenAICredentials = z.infer<typeof credentialSchema>;
+
 export function OpenAIKeyPrompt({
  onSubmit,
  onCancel,
@@ -31,6 +42,34 @@ export function OpenAIKeyPrompt({
  const [currentField, setCurrentField] = useState<
    'apiKey' | 'baseUrl' | 'model'
  >('apiKey');
+  const [validationError, setValidationError] = useState<string | null>(null);
+
+  const validateAndSubmit = () => {
+    setValidationError(null);
+
+    try {
+      const validated = credentialSchema.parse({
+        apiKey: apiKey.trim(),
+        baseUrl: baseUrl.trim() || undefined,
+        model: model.trim() || undefined,
+      });
+
+      onSubmit(
+        validated.apiKey,
+        validated.baseUrl === '' ? '' : validated.baseUrl || '',
+        validated.model || '',
+      );
+    } catch (error) {
+      if (error instanceof z.ZodError) {
+        const errorMessage = error.errors
+          .map((e) => `${e.path.join('.')}: ${e.message}`)
+          .join(', ');
+        setValidationError(`Invalid credentials: ${errorMessage}`);
+      } else {
+        setValidationError('Failed to validate credentials');
+      }
+    }
+  };

  useKeypress(
    (key) => {
@@ -52,7 +91,7 @@ export function OpenAIKeyPrompt({
        } else if (currentField === 'model') {
          // 只有在提交时才检查 API key 是否为空
          if (apiKey.trim()) {
-            onSubmit(apiKey.trim(), baseUrl.trim(), model.trim());
+            validateAndSubmit();
          } else {
            // 如果 API key 为空，回到 API key 字段
            setCurrentField('apiKey');
@@ -168,6 +207,11 @@ export function OpenAIKeyPrompt({
      <Text bold color={Colors.AccentBlue}>
        OpenAI Configuration Required
      </Text>
+      {validationError && (
+        <Box marginTop={1}>
+          <Text color={Colors.AccentRed}>{validationError}</Text>
+        </Box>
+      )}
      <Box marginTop={1}>
        <Text>
          Please enter your OpenAI configuration. You can get an API key from{' '}
--- a/packages/cli/src/ui/components/QwenOAuthProgress.test.tsx
+++ b/packages/cli/src/ui/components/QwenOAuthProgress.test.tsx
@@ -8,7 +8,7 @@
 import { render } from 'ink-testing-library';
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { QwenOAuthProgress } from './QwenOAuthProgress.js';
-import type { DeviceAuthorizationInfo } from '../hooks/useQwenAuth.js';
+import type { DeviceAuthorizationData } from '@qwen-code/qwen-code-core';
 import { useKeypress } from '../hooks/useKeypress.js';
 import type { Key } from '../contexts/KeypressContext.js';

@@ -42,12 +42,13 @@ describe('QwenOAuthProgress', () => {
  let keypressHandler: ((key: Key) => void) | null = null;

  const createMockDeviceAuth = (
-    overrides: Partial<DeviceAuthorizationInfo> = {},
-  ): DeviceAuthorizationInfo => ({
+    overrides: Partial<DeviceAuthorizationData> = {},
+  ): DeviceAuthorizationData => ({
    verification_uri: 'https://example.com/device',
    verification_uri_complete: 'https://example.com/device?user_code=ABC123',
    user_code: 'ABC123',
    expires_in: 300,
+    device_code: 'test-device-code',
    ...overrides,
  });

@@ -55,7 +56,7 @@ describe('QwenOAuthProgress', () => {

  const renderComponent = (
    props: Partial<{
-      deviceAuth: DeviceAuthorizationInfo;
+      deviceAuth: DeviceAuthorizationData;
      authStatus:
        | 'idle'
        | 'polling'
@@ -158,7 +159,7 @@ describe('QwenOAuthProgress', () => {
    });

    it('should format time correctly', () => {
-      const deviceAuthWithCustomTime: DeviceAuthorizationInfo = {
+      const deviceAuthWithCustomTime: DeviceAuthorizationData = {
        ...mockDeviceAuth,
        expires_in: 125, // 2 minutes and 5 seconds
      };
@@ -176,7 +177,7 @@ describe('QwenOAuthProgress', () => {
    });

    it('should format single digit seconds with leading zero', () => {
-      const deviceAuthWithCustomTime: DeviceAuthorizationInfo = {
+      const deviceAuthWithCustomTime: DeviceAuthorizationData = {
        ...mockDeviceAuth,
        expires_in: 67, // 1 minute and 7 seconds
      };
@@ -196,7 +197,7 @@ describe('QwenOAuthProgress', () => {

  describe('Timer functionality', () => {
    it('should countdown and call onTimeout when timer expires', async () => {
-      const deviceAuthWithShortTime: DeviceAuthorizationInfo = {
+      const deviceAuthWithShortTime: DeviceAuthorizationData = {
        ...mockDeviceAuth,
        expires_in: 2, // 2 seconds
      };
@@ -520,7 +521,7 @@ describe('QwenOAuthProgress', () => {

  describe('Props changes', () => {
    it('should display initial timer value from deviceAuth', () => {
-      const deviceAuthWith10Min: DeviceAuthorizationInfo = {
+      const deviceAuthWith10Min: DeviceAuthorizationData = {
        ...mockDeviceAuth,
        expires_in: 600, // 10 minutes
      };
--- a/packages/cli/src/ui/components/QwenOAuthProgress.tsx
+++ b/packages/cli/src/ui/components/QwenOAuthProgress.tsx
@@ -11,13 +11,13 @@ import Spinner from 'ink-spinner';
 import Link from 'ink-link';
 import qrcode from 'qrcode-terminal';
 import { Colors } from '../colors.js';
-import type { DeviceAuthorizationInfo } from '../hooks/useQwenAuth.js';
+import type { DeviceAuthorizationData } from '@qwen-code/qwen-code-core';
 import { useKeypress } from '../hooks/useKeypress.js';

 interface QwenOAuthProgressProps {
  onTimeout: () => void;
  onCancel: () => void;
-  deviceAuth?: DeviceAuthorizationInfo;
+  deviceAuth?: DeviceAuthorizationData;
  authStatus?:
    | 'idle'
    | 'polling'
@@ -131,8 +131,8 @@ export function QwenOAuthProgress({

  useKeypress(
    (key) => {
-      if (authStatus === 'timeout') {
-        // Any key press in timeout state should trigger cancel to return to auth dialog
+      if (authStatus === 'timeout' || authStatus === 'error') {
+        // Any key press in timeout or error state should trigger cancel to return to auth dialog
        onCancel();
      } else if (key.name === 'escape' || (key.ctrl && key.name === 'c')) {
        onCancel();
@@ -234,6 +234,35 @@ export function QwenOAuthProgress({
    );
  }

+  if (authStatus === 'error') {
+    return (
+      <Box
+        borderStyle="round"
+        borderColor={Colors.AccentRed}
+        flexDirection="column"
+        padding={1}
+        width="100%"
+      >
+        <Text bold color={Colors.AccentRed}>
+          Qwen OAuth Authentication Error
+        </Text>
+
+        <Box marginTop={1}>
+          <Text>
+            {authMessage ||
+              'An error occurred during authentication. Please try again.'}
+          </Text>
+        </Box>
+
+        <Box marginTop={1}>
+          <Text color={Colors.Gray}>
+            Press any key to return to authentication type selection.
+          </Text>
+        </Box>
+      </Box>
+    );
+  }
+
  // Show loading state when no device auth is available yet
  if (!deviceAuth) {
    return (
--- a/packages/cli/src/ui/components/SettingsDialog.test.tsx
+++ b/packages/cli/src/ui/components/SettingsDialog.test.tsx
@@ -487,8 +487,11 @@ describe('SettingsDialog', () => {
      it('loops back when reaching the end of an enum', async () => {
        vi.mocked(saveModifiedSettings).mockClear();
        vi.mocked(getSettingsSchema).mockReturnValue(FAKE_SCHEMA);
-        const settings = createMockSettings();
-        settings.setValue(SettingScope.User, 'ui.theme', StringEnum.BAZ);
+        const settings = createMockSettings({
+          ui: {
+            theme: StringEnum.BAZ,
+          },
+        });
        const onSelect = vi.fn();
        const component = (
          <KeypressProvider kittyProtocolEnabled={false}>