From fbeaf1395e161b49f822d1d06feba6cfcb647921 Mon Sep 17 00:00:00 2001
From: y574444354 <574444354@qq.com>
Date: Thu, 16 Apr 2026 19:12:01 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=20CoStrict=20provider?=
 =?UTF-8?q?=20=E6=A8=A1=E5=9E=8B=E9=80=89=E6=8B=A9=E4=B8=8E=20token=20?=
 =?UTF-8?q?=E4=BC=B0=E7=AE=97=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- ConsoleOAuthFlow: 替换 ModelPicker 为直接 Select 组件，避免
  getAPIProvider() 竞态导致显示标准 Anthropic 模型列表
- modelOptions: 简化模型选项逻辑，CoStrict 凭证持有者统一走
  CoStrict 模型缓存路径，加载中显示提示而非标准模型列表
- provider/index: 回写 usage 和 stop_reason 到已 yield 的
  assistant messages，与 Anthropic 原生路径保持一致
- tokenEstimation: 非 Anthropic 原生 provider（costrict/openai/
  grok/gemini）统一使用粗略 token 估算
- 品牌文案: Claude → CoStrict 相关措辞修正
---
 src/components/ConsoleOAuthFlow.tsx | 55 +++++++++++++-------
 src/components/ModelPicker.tsx      |  2 +-
 src/constants/prompts.ts            |  2 +-
 src/costrict/provider/index.ts      | 18 +++++++
 src/services/tokenEstimation.ts     | 16 +++++-
 src/utils/betas.ts                  |  2 +-
 src/utils/model/modelOptions.ts     | 80 ++++-------------------------
 7 files changed, 81 insertions(+), 94 deletions(-)

diff --git a/src/components/ConsoleOAuthFlow.tsx b/src/components/ConsoleOAuthFlow.tsx
index 73ba8d0a5..d66cc58f7 100644
--- a/src/components/ConsoleOAuthFlow.tsx
+++ b/src/components/ConsoleOAuthFlow.tsx
@@ -17,7 +17,6 @@ import { getSettings_DEPRECATED, updateSettingsForSource } from '../utils/settin
 import { Select } from './CustomSelect/select.js';
 import { Spinner } from './Spinner.js';
 import TextInput from './TextInput.js';
-import { ModelPicker } from './ModelPicker.js';
 import { useSetAppState } from '../state/AppState.js';
 
 type Props = {
@@ -1260,25 +1259,43 @@ function OAuthStatusMessage({
 
     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 (
-        <ModelPicker
-          initial={sortedModels[0]?.id ?? null}
-          headerText="Login successful. Select a CoStrict model to use:"
-          onSelect={model => {
-            const selected = model ?? sortedModels[0]?.id ?? '';
-            process.env.COSTRICT_MODEL = selected;
-            setAppState(prev => ({ ...prev, mainLoopModel: selected, 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 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>
       );
     }
 
diff --git a/src/components/ModelPicker.tsx b/src/components/ModelPicker.tsx
index 2c32535df..4a4ae7589 100644
--- a/src/components/ModelPicker.tsx
+++ b/src/components/ModelPicker.tsx
@@ -224,7 +224,7 @@ export function ModelPicker({
           </Text>
           <Text dimColor>
             {headerText ??
-              'Switch between Claude models. Applies to this session and future CoStrict sessions. For other/previous model names, specify with --model.'}
+              'Switch between CoStrict models. Applies to this session and future CoStrict sessions. For other/previous model names, specify with --model.'}
           </Text>
           {sessionModel && (
             <Text dimColor>
diff --git a/src/constants/prompts.ts b/src/constants/prompts.ts
index 95fe50396..62063eded 100644
--- a/src/constants/prompts.ts
+++ b/src/constants/prompts.ts
@@ -694,7 +694,7 @@ export async function computeSimpleEnvInfo(
     knowledgeCutoffMessage,
     process.env.USER_TYPE === 'ant' && isUndercover()
       ? null
-      : `The most recent Claude model family is Claude 4.5/4.6. Model IDs — Opus 4.6: '${CLAUDE_4_5_OR_4_6_MODEL_IDS.opus}', Sonnet 4.6: '${CLAUDE_4_5_OR_4_6_MODEL_IDS.sonnet}', Haiku 4.5: '${CLAUDE_4_5_OR_4_6_MODEL_IDS.haiku}'. When building AI applications, default to the latest and most capable Claude models.`,
+      : `The most recent Claude model family is Claude 4.5/4.6. Model IDs — Opus 4.6: '${CLAUDE_4_5_OR_4_6_MODEL_IDS.opus}', Sonnet 4.6: '${CLAUDE_4_5_OR_4_6_MODEL_IDS.sonnet}', Haiku 4.5: '${CLAUDE_4_5_OR_4_6_MODEL_IDS.haiku}'. When building AI applications, default to the latest and most capable CoStrict models.`,
     process.env.USER_TYPE === 'ant' && isUndercover()
       ? null
       : `Claude Code is available as a CLI in the terminal, desktop app (Mac/Windows), web app (costrict.ai), and IDE extensions (VS Code, JetBrains).`,
diff --git a/src/costrict/provider/index.ts b/src/costrict/provider/index.ts
index 00bfdda81..83fea97e3 100644
--- a/src/costrict/provider/index.ts
+++ b/src/costrict/provider/index.ts
@@ -136,7 +136,10 @@ export async function* queryModelCoStrict(
     const adaptedStream = adaptOpenAIStreamToAnthropic(stream, costrictModel)
 
     const contentBlocks: Record<number, any> = {}
+    // 跟踪已 yield 的 assistant messages，用于 message_delta 时回写 usage
+    const yieldedMessages: AssistantMessage[] = []
     let partialMessage: any
+    let stopReason: string | null = null
     let usage = {
       input_tokens: 0,
       output_tokens: 0,
@@ -194,18 +197,33 @@ export async function* queryModelCoStrict(
             message: {
               ...partialMessage,
               content: normalizeContentFromAPI([block], tools, options.agentId),
+              usage,
             },
             requestId: undefined,
             type: 'assistant',
             uuid: randomUUID(),
             timestamp: new Date().toISOString(),
           }
+          yieldedMessages.push(m)
           yield m
           break
         }
         case 'message_delta': {
           const deltaUsage = (event as any).usage
           if (deltaUsage) usage = { ...usage, ...deltaUsage }
+          // 回写 usage 到已 yield 的 assistant messages
+          // 与 Anthropic 原生路径 claude.ts:2298 保持一致
+          for (const msg of yieldedMessages) {
+            msg.message.usage = usage
+          }
+          // 记录 stop_reason，回写到最后的 message
+          if ((event as any).delta?.stop_reason != null) {
+            stopReason = (event as any).delta.stop_reason
+            const lastMsg = yieldedMessages[yieldedMessages.length - 1]
+            if (lastMsg) {
+              lastMsg.message.stop_reason = stopReason
+            }
+          }
           break
         }
         case 'message_stop':
diff --git a/src/services/tokenEstimation.ts b/src/services/tokenEstimation.ts
index c59d53a3a..0502f53f9 100644
--- a/src/services/tokenEstimation.ts
+++ b/src/services/tokenEstimation.ts
@@ -144,7 +144,13 @@ export async function countMessagesTokensWithAPI(
   return withTokenCountVCR(messages, tools, async () => {
     try {
       const provider = getAPIProvider()
-      if (provider === 'gemini') {
+      // 非 Anthropic 原生 provider 不支持 countTokens API，直接使用粗略估算
+      if (
+        provider === 'gemini' ||
+        provider === 'openai' ||
+        provider === 'grok' ||
+        provider === 'costrict'
+      ) {
         return roughTokenCountEstimationForAPIRequest(messages, tools)
       }
 
@@ -258,7 +264,13 @@ export async function countTokensViaHaikuFallback(
   tools: Anthropic.Beta.Messages.BetaToolUnion[],
 ): Promise<number | null> {
   const provider = getAPIProvider()
-  if (provider === 'gemini') {
+  // 非 Anthropic 原生 provider 不支持 countTokens API，直接使用粗略估算
+  if (
+    provider === 'gemini' ||
+    provider === 'openai' ||
+    provider === 'grok' ||
+    provider === 'costrict'
+  ) {
     return roughTokenCountEstimationForAPIRequest(messages, tools)
   }
 
diff --git a/src/utils/betas.ts b/src/utils/betas.ts
index 3a48a943a..326d49c96 100644
--- a/src/utils/betas.ts
+++ b/src/utils/betas.ts
@@ -181,7 +181,7 @@ export function modelSupportsAutoMode(model: string): boolean {
       return true
     }
     if (process.env.USER_TYPE === 'ant') {
-      // Denylist: block known-unsupported claude models, allow everything else (ant-internal models etc.)
+      // Denylist: block known-unsupported CoStrict models, allow everything else (ant-internal models etc.)
       if (m.includes('claude-3-')) return false
       // claude-*-4 not followed by -[6-9]: blocks bare -4, -4-YYYYMMDD, -4@, -4-0 thru -4-5
       if (/claude-(opus|sonnet|haiku)-4(?!-[6-9])/.test(m)) return false
diff --git a/src/utils/model/modelOptions.ts b/src/utils/model/modelOptions.ts
index 5ecf35367..16798ef84 100644
--- a/src/utils/model/modelOptions.ts
+++ b/src/utils/model/modelOptions.ts
@@ -380,7 +380,10 @@ function getModelOptionsBase(fastMode = false): ModelOption[] {
   }
 
   // CoStrict provider: 从缓存中动态展示 CoStrict 可用模型
-  if (getAPIProvider() === 'costrict') {
+  // 也覆盖其他 provider 但用户已有 CoStrict 凭证的场景（modelType 丢失/竞态）
+  const isCoStrictProvider = getAPIProvider() === 'costrict'
+  const hasCoStrictCreds = hasCoStrictCredentialsSync()
+  if (isCoStrictProvider || hasCoStrictCreds) {
     // 延迟 require 避免循环依赖，且该模块只在 costrict provider 下加载
     // eslint-disable-next-line @typescript-eslint/no-require-imports
     const { getCachedCoStrictModels } =
@@ -400,82 +403,19 @@ function getModelOptionsBase(fastMode = false): ModelOption[] {
         }
       })
     }
-    // 缓存尚未加载（凭证过期或首次加载中），提示用户重新登录
+    // 缓存尚未加载（启动时预取未完成），显示提示而非标准模型列表
     return [
       {
         value: null,
-        label: 'CoStrict (not logged in)',
-        description: 'Run /login to sign in with CoStrict again',
+        label: 'CoStrict (loading models...)',
+        description:
+          'CoStrict models are being loaded, please try again in a moment',
       },
     ]
   }
 
-  // PAYG 1P API: Default (Sonnet) + Sonnet 1M + Opus 4.6 + Opus 1M + Haiku
-  if (getAPIProvider() === 'firstParty') {
-    // 当用户未登录 CoStrict 时，只显示推荐登录选项
-    if (!hasCoStrictCredentialsSync()) {
-      return [getCoStrictLoginOption()]
-    }
-
-    const payg1POptions: ModelOption[] = []
-    payg1POptions.push(getDefaultOptionForUser(fastMode))
-
-    if (checkSonnet1mAccess()) {
-      payg1POptions.push(getSonnet46_1MOption())
-    }
-    if (isOpus1mMergeEnabled()) {
-      payg1POptions.push(getMergedOpus1MOption(fastMode))
-    } else {
-      payg1POptions.push(getOpus46Option(fastMode))
-      if (checkOpus1mAccess()) {
-        payg1POptions.push(getOpus46_1MOption(fastMode))
-      }
-    }
-    payg1POptions.push(getHaiku45Option())
-
-    return payg1POptions
-  }
-
-  // PAYG 3P: Default (Sonnet 4.5) + Sonnet (3P custom) or Sonnet 4.6/1M + Opus (3P custom) or Opus 4.1/Opus 4.6/Opus1M + Haiku + Opus 4.1
-
-  // 当用户未登录 CoStrict 时，只显示推荐登录选项
-  if (getAPIProvider() !== 'costrict' && !hasCoStrictCredentialsSync()) {
-    return [getCoStrictLoginOption()]
-  }
-
-  const payg3pOptions: ModelOption[] = []
-  payg3pOptions.push(getDefaultOptionForUser(fastMode))
-
-  const customSonnet = getCustomSonnetOption()
-  if (customSonnet !== undefined) {
-    payg3pOptions.push(customSonnet)
-  } else {
-    // Add Sonnet 4.6 since Sonnet 4.5 is the default
-    payg3pOptions.push(getSonnet46Option())
-    if (checkSonnet1mAccess()) {
-      payg3pOptions.push(getSonnet46_1MOption())
-    }
-  }
-
-  const customOpus = getCustomOpusOption()
-  if (customOpus !== undefined) {
-    payg3pOptions.push(customOpus)
-  } else {
-    // Add Opus 4.1, Opus 4.6 and Opus 4.6 1M
-    payg3pOptions.push(getOpus41Option()) // This is the default opus
-    payg3pOptions.push(getOpus46Option(fastMode))
-    if (checkOpus1mAccess()) {
-      payg3pOptions.push(getOpus46_1MOption(fastMode))
-    }
-  }
-  const customHaiku = getCustomHaikuOption()
-  if (customHaiku !== undefined) {
-    payg3pOptions.push(customHaiku)
-  } else {
-    payg3pOptions.push(getHaikuOption())
-  }
-
-  return payg3pOptions
+  // 用户没有 CoStrict 凭证，显示推荐登录选项
+  return [getCoStrictLoginOption()]
 }
 
 /**