From d7bcf65968bb1f632bf9ee157c0fe4e3f0e250da Mon Sep 17 00:00:00 2001
From: Askhz <1361267452@qq.com>
Date: Mon, 20 Apr 2026 20:40:59 +0800
Subject: [PATCH] feat: restore CoStrict login option in interactive login

- Added CoStrict option to the login Select component
- Implemented complete OAuth flow with browser opening and token polling
- Added costrict_waiting and costrict_model_select OAuth states
- Fixed TypeScript errors in fetch.ts and index.ts

The CoStrict login option was previously removed during login simplification.
This commit restores the full functionality including model selection after login.
---
 src/components/ConsoleOAuthFlow.tsx | 142 +++++++++++++++++++++++++++-
 src/costrict/provider/fetch.ts      |  39 ++++----
 src/costrict/provider/index.ts      |   7 +-
 3 files changed, 164 insertions(+), 24 deletions(-)

diff --git a/src/components/ConsoleOAuthFlow.tsx b/src/components/ConsoleOAuthFlow.tsx
index b4da18737..512230293 100644
--- a/src/components/ConsoleOAuthFlow.tsx
+++ b/src/components/ConsoleOAuthFlow.tsx
@@ -16,6 +16,7 @@ import { getSettings_DEPRECATED, updateSettingsForSource } from '../utils/settin
 import { Select } from './CustomSelect/select.js'
 import { Spinner } from './Spinner.js'
 import TextInput from './TextInput.js'
+import { useSetAppState } from '../state/AppState.js'
 import { fi } from 'zod/v4/locales'
 
 type Props = {
@@ -28,6 +29,14 @@ type Props = {
 type OAuthStatus =
   | { state: 'idle' } // Initial state, waiting to select login method
   | { state: 'platform_setup' } // Show platform setup info (Bedrock/Vertex/Foundry)
+  | {
+      state: 'costrict_waiting'
+      url: string
+    } // CoStrict OAuth: browser opened, waiting for user to login
+  | {
+      state: 'costrict_model_select'
+      models: Array<{ id: string; name?: string }>
+    } // CoStrict: login done, select a model
   | {
       state: 'custom_platform'
       baseUrl: string
@@ -94,6 +103,7 @@ export function ConsoleOAuthFlow({
     }
     return { state: 'idle' }
   })
+  const setAppState = useSetAppState()
 
   const [pastedCode, setPastedCode] = useState('')
   const [cursorOffset, setCursorOffset] = useState(0)
@@ -438,6 +448,7 @@ function OAuthStatusMessage({
   setLoginWithClaudeAi,
   onDone,
 }: OAuthStatusMessageProps): React.ReactNode {
+  const setAppState = useSetAppState()
   switch (oauthStatus.state) {
     case 'idle':
       return (
@@ -453,6 +464,16 @@ function OAuthStatusMessage({
           <Box>
             <Select
               options={[
+                {
+                  label: (
+                    <Text>
+                      CoStrict ·{' '}
+                      <Text dimColor>Sign in with CoStrict account</Text>
+                      {'\n'}
+                    </Text>
+                  ),
+                  value: 'costrict',
+                },
                 {
                   label: (
                     <Text>
@@ -487,7 +508,70 @@ function OAuthStatusMessage({
                 },
               ]}
               onChange={value => {
-                if (value === 'custom_platform') {
+                if (value === 'costrict') {
+                  void (async () => {
+                    try {
+                      const { generateState, getCoStrictBaseURL, buildCoStrictLoginURL, pollLoginToken } = await import('../costrict/provider/auth.js')
+                      const { generateMachineId, saveCoStrictCredentials } = await import('../costrict/provider/credentials.js')
+                      const { extractExpiryFromJWT } = await import('../costrict/provider/token.js')
+                      const { updateSettingsForSource } = await import('../utils/settings/settings.js')
+
+                      const baseUrl = getCoStrictBaseURL()
+                      const state = generateState()
+                      const machineId = generateMachineId()
+                      const loginUrl = buildCoStrictLoginURL(baseUrl, state, machineId)
+
+                      // 打开浏览器
+                      const { openBrowser } = await import('../utils/browser.js')
+                      await openBrowser(loginUrl)
+
+                      setOAuthStatus({ state: 'costrict_waiting', url: loginUrl })
+
+                      // 轮询等待登录完成
+                      const tokens = await pollLoginToken(baseUrl, state, machineId)
+
+                      // 保存凭证
+                      const expiryDate = extractExpiryFromJWT(tokens.access_token)
+                      await saveCoStrictCredentials({
+                        id: 'csc',
+                        name: 'CSC Auth',
+                        access_token: tokens.access_token,
+                        refresh_token: tokens.refresh_token,
+                        state,
+                        machine_id: machineId,
+                        base_url: baseUrl,
+                        expiry_date: expiryDate,
+                        updated_at: new Date().toISOString(),
+                        expired_at: expiryDate ? new Date(expiryDate).toISOString() : undefined,
+                      })
+
+                      // 设置 modelType 为 costrict
+                      updateSettingsForSource('userSettings', { modelType: 'costrict' as any } as any)
+                      process.env.CLAUDE_CODE_USE_COSTRICT = '1'
+
+                      // 预取模型列表，填充同步缓存，并进入模型选择界面
+                      try {
+                        const { fetchCoStrictModels } = await import('../costrict/provider/models.js')
+                        const models = await fetchCoStrictModels(baseUrl, tokens.access_token)
+                        if (models.length > 0) {
+                          setOAuthStatus({ state: 'costrict_model_select', models })
+                          return
+                        }
+                      } catch {
+                        // 预取失败，直接进入 success
+                      }
+
+                      setOAuthStatus({ state: 'success' })
+                      void onDone()
+                    } catch (err: any) {
+                      setOAuthStatus({
+                        state: 'error',
+                        message: err.message || String(err),
+                        toRetry: { state: 'idle' },
+                      })
+                    }
+                  })()
+                } else if (value === 'custom_platform') {
                   logEvent('tengu_custom_platform_selected', {})
                   setOAuthStatus({
                     state: 'custom_platform',
@@ -1220,6 +1304,62 @@ function OAuthStatusMessage({
         )
       }
 
+    case 'costrict_waiting':
+      return (
+        <Box flexDirection="column" gap={1}>
+          <Text>
+            Opening browser for CoStrict login. If it does not open
+            automatically, copy and paste this URL:
+          </Text>
+          <Box marginY={1}>
+            <Text color="suggestion">{oauthStatus.url}</Text>
+          </Box>
+          <Text dimColor>Waiting for authentication...</Text>
+        </Box>
+      )
+
+    case 'costrict_model_select': {
+      const sortedModels = [...oauthStatus.models].sort((a, b) => a.id.localeCompare(b.id));
+      // 直接使用已获取的模型列表渲染选项，避免通用 ModelPicker 因 getAPIProvider() 竞态
+      // 导致显示标准 Anthropic 模型列表而非 CoStrict 动态模型
+      const costrictOptions = sortedModels.map(m => ({
+        value: m.id,
+        label: m.name ?? m.id,
+        description:
+          m.id === 'Auto'
+            ? `${Math.round(((m as any).creditDiscount ?? 0) * 100)}% discount`
+            : `${(m as any).creditConsumption ?? '?'}x credit`,
+      }));
+      return (
+        <Box flexDirection="column">
+          <Box marginBottom={1} flexDirection="column">
+            <Text color="remember" bold>
+              Select model
+            </Text>
+            <Text dimColor>Login successful. Select a CoStrict model to use:</Text>
+          </Box>
+          <Box flexDirection="column" marginBottom={1}>
+            <Select
+              options={costrictOptions}
+              onChange={(value: string) => {
+                process.env.COSTRICT_MODEL = value;
+                setAppState(prev => ({ ...prev, mainLoopModel: value, mainLoopModelForSession: null }));
+                setOAuthStatus({ state: 'success' });
+                void onDone();
+              }}
+              onCancel={() => {
+                const selected = sortedModels[0]?.id ?? '';
+                process.env.COSTRICT_MODEL = selected;
+                setAppState(prev => ({ ...prev, mainLoopModel: selected, mainLoopModelForSession: null }));
+                setOAuthStatus({ state: 'success' });
+                void onDone();
+              }}
+            />
+          </Box>
+        </Box>
+      );
+    }
+
     case 'platform_setup':
       return (
         <Box flexDirection="column" gap={1} marginTop={1}>
diff --git a/src/costrict/provider/fetch.ts b/src/costrict/provider/fetch.ts
index dabbf532a..14dbefc92 100644
--- a/src/costrict/provider/fetch.ts
+++ b/src/costrict/provider/fetch.ts
@@ -34,6 +34,10 @@ function getVersion(): string {
 
 const VERSION = getVersion()
 
+type CoStrictFetch = typeof fetch & {
+  preconnect?: (url: string | URL) => void
+}
+
 /**
  * 创建自定义 fetch 函数，用于 CoStrict API 请求
  *
@@ -43,7 +47,7 @@ const VERSION = getVersion()
  * 3. 注入 Authorization 和 CoStrict 特有 headers
  * 4. 反应性 401 错误恢复（自动重试一次）
  */
-export function createCoStrictFetch() {
+export function createCoStrictFetch(): CoStrictFetch {
   const costrictFetch = async (
     input: RequestInfo | URL,
     init?: RequestInit,
@@ -126,20 +130,21 @@ export function createCoStrictFetch() {
   }
   // Bun 原生支持 fetch.preconnect（共享连接池预热）
   // Node.js 没有 preconnect API，降级为 net.createConnection 做 TCP 预热
-  if (typeof fetch.preconnect === 'function') {
-    costrictFetch.preconnect = fetch.preconnect.bind(fetch)
-  } else {
-    costrictFetch.preconnect = (url: string | URL) => {
-      try {
-        const { hostname, port } = new URL(typeof url === 'string' ? url : url.toString())
-        const { createConnection } = require('node:net') as typeof import('node:net')
-        const sock = createConnection(Number(port) || 443, hostname, () => sock.destroy())
-        sock.setTimeout(3000, () => sock.destroy())
-        sock.on('error', () => sock.destroy())
-      } catch {
-        // 预连接失败不影响正常流程
-      }
-    }
-  }
-  return costrictFetch
+  const costrictFetchWithPreconnect = Object.assign(costrictFetch, {
+    preconnect: typeof fetch.preconnect === 'function'
+      ? fetch.preconnect.bind(fetch)
+      : (url: string | URL) => {
+          try {
+            const { hostname, port } = new URL(typeof url === 'string' ? url : url.toString())
+            const { createConnection } = require('node:net') as typeof import('node:net')
+            const sock = createConnection(Number(port) || 443, hostname, () => sock.destroy())
+            sock.setTimeout(3000, () => sock.destroy())
+            sock.on('error', () => sock.destroy())
+          } catch {
+            // 预连接失败不影响正常流程
+          }
+        }
+  })
+
+  return costrictFetchWithPreconnect as CoStrictFetch
 }
diff --git a/src/costrict/provider/index.ts b/src/costrict/provider/index.ts
index 83fea97e3..7dfe7896e 100644
--- a/src/costrict/provider/index.ts
+++ b/src/costrict/provider/index.ts
@@ -16,12 +16,7 @@ import type { Tools } from '../../Tool.js'
 import type { Options } from '../../services/api/claude.js'
 import OpenAI from 'openai'
 import { getProxyFetchOptions } from '../../utils/proxy.js'
-import { anthropicMessagesToOpenAI } from '../../services/api/openai/convertMessages.js'
-import {
-  anthropicToolsToOpenAI,
-  anthropicToolChoiceToOpenAI,
-} from '../../services/api/openai/convertTools.js'
-import { adaptOpenAIStreamToAnthropic } from '../../services/api/openai/streamAdapter.js'
+import { anthropicMessagesToOpenAI, anthropicToolsToOpenAI, anthropicToolChoiceToOpenAI, adaptOpenAIStreamToAnthropic } from '@ant/model-provider'
 import { normalizeMessagesForAPI } from '../../utils/messages.js'
 import { toolToAPISchema } from '../../utils/api.js'
 import { logForDebugging } from '../../utils/debug.js'