Merge pull request #236 from cita-777/codex/claude-count-models

Fix Claude count_tokens compatibility and OAuth model discovery
2026-03-21 19:28:18 +08:00
parent c7ba436685 dc93c58582
commit e09fb43b33
5 changed files with 201 additions and 46 deletions
@@ -869,7 +869,17 @@ export async function handleClaudeCountTokensSurfaceRequest(
    }

    excludeChannelIds.push(selected.channel.id);
-    if (String(selected.site.platform || '').trim().toLowerCase() !== 'claude') {
+    const modelName = selected.actualModel || requestedModel;
+    const endpointCandidates = await resolveUpstreamEndpointCandidates(
+      {
+        site: selected.site,
+        account: selected.account,
+      },
+      modelName,
+      'claude',
+      requestedModel,
+    );
+    if (!endpointCandidates.includes('messages')) {
      if (retryCount < MAX_RETRIES) {
        retryCount += 1;
        continue;
@@ -881,8 +891,6 @@ export async function handleClaudeCountTokensSurfaceRequest(
        },
      });
    }
-
-    const modelName = selected.actualModel || requestedModel;
    const oauth = getOauthInfoFromExtraConfig(selected.account.extraConfig);
    const startTime = Date.now();

@@ -156,4 +156,44 @@ describe('claude count_tokens proxy route', () => {
      },
    ]);
  });
+
+  it('supports /v1/messages/count_tokens for openai-platform gateways that expose Claude messages endpoints', async () => {
+    selectChannelMock.mockReturnValue({
+      channel: { id: 12, routeId: 23 },
+      site: { name: 'gateway-site', url: 'https://gateway.example.com', platform: 'openai' },
+      account: { id: 34, username: 'gateway-user@example.com' },
+      tokenName: 'default',
+      tokenValue: 'sk-gateway',
+      actualModel: 'claude-sonnet-4-5-20250929',
+    });
+    fetchMock.mockResolvedValue(new Response(JSON.stringify({
+      input_tokens: 9,
+    }), {
+      status: 200,
+      headers: { 'content-type': 'application/json' },
+    }));
+
+    const response = await app.inject({
+      method: 'POST',
+      url: '/v1/messages/count_tokens',
+      payload: {
+        model: 'claude-sonnet-4-5-20250929',
+        messages: [
+          {
+            role: 'user',
+            content: [{ type: 'text', text: 'count through a compatible gateway' }],
+          },
+        ],
+      },
+    });
+
+    expect(response.statusCode).toBe(200);
+    expect(response.json()).toEqual({ input_tokens: 9 });
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+
+    const [targetUrl, options] = fetchMock.mock.calls[0] as [string, any];
+    expect(targetUrl).toBe('https://gateway.example.com/v1/messages/count_tokens?beta=true');
+    expect(options.headers['x-api-key']).toBe('sk-gateway');
+    expect(options.headers['anthropic-version']).toBe('2023-06-01');
+  });
 });
@@ -1542,32 +1542,60 @@ export function buildClaudeCountTokensUpstreamRequest(input: {
  delete sanitizedBody.maxTokens;
  delete sanitizedBody.stream;
  const providerProfile = resolveProviderProfile(sitePlatform);
-  if (providerProfile?.id !== 'claude') {
-    throw new Error(`missing claude provider profile for platform: ${sitePlatform || 'unknown'}`);
+  const effectiveClaudeHeaders = {
+    ...claudeHeaders,
+    ...(betas.length > 0 ? { 'anthropic-beta': betas.join(',') } : {}),
+  };
+
+  if (providerProfile?.id === 'claude') {
+    const prepared = providerProfile.prepareRequest({
+      endpoint: 'messages',
+      modelName: input.modelName,
+      stream: false,
+      tokenValue: input.tokenValue,
+      oauthProvider: input.oauthProvider,
+      sitePlatform,
+      baseHeaders: {
+        'Content-Type': 'application/json',
+      },
+      claudeHeaders: effectiveClaudeHeaders,
+      body: sanitizedBody,
+      action: 'countTokens',
+    });
+
+    return {
+      path: prepared.path,
+      headers: prepared.headers,
+      body: prepared.body,
+      runtime: {
+        executor: 'claude',
+        modelName: input.modelName,
+        stream: false,
+        action: 'countTokens',
+      },
+    };
  }

-  const prepared = providerProfile.prepareRequest({
-    endpoint: 'messages',
-    modelName: input.modelName,
-    stream: false,
-    tokenValue: input.tokenValue,
-    oauthProvider: input.oauthProvider,
-    sitePlatform,
+  const anthropicVersion = (
+    effectiveClaudeHeaders['anthropic-version']
+    || '2023-06-01'
+  );
+  const isClaudeOauthUpstream = sitePlatform === 'claude' && input.oauthProvider === 'claude';
+  const headers = buildClaudeRuntimeHeaders({
    baseHeaders: {
      'Content-Type': 'application/json',
    },
-    claudeHeaders: {
-      ...claudeHeaders,
-      ...(betas.length > 0 ? { 'anthropic-beta': betas.join(',') } : {}),
-    },
-    body: sanitizedBody,
-    action: 'countTokens',
+    claudeHeaders: effectiveClaudeHeaders,
+    anthropicVersion,
+    stream: false,
+    isClaudeOauthUpstream,
+    tokenValue: input.tokenValue,
  });

  return {
-    path: prepared.path,
-    headers: prepared.headers,
-    body: prepared.body,
+    path: '/v1/messages/count_tokens?beta=true',
+    headers,
+    body: sanitizedBody,
    runtime: {
      executor: 'claude',
      modelName: input.modelName,
@@ -537,6 +537,78 @@ describe('refreshModelsForAccount credential discovery', () => {
    ]);
  });

+  it('discovers Claude OAuth models from the upstream /v1/models response', async () => {
+    getApiTokenMock.mockResolvedValue(null);
+    getModelsMock.mockRejectedValue(new Error('claude oauth discovery should not call adapter.getModels'));
+    undiciFetchMock.mockResolvedValue({
+      ok: true,
+      status: 200,
+      json: async () => ({
+        data: [
+          { id: 'claude-3-7-sonnet-latest' },
+          { id: 'claude-opus-4-1-20250805' },
+        ],
+      }),
+      text: async () => JSON.stringify({ ok: true }),
+    });
+
+    const site = await db.insert(schema.sites).values({
+      name: 'claude-site',
+      url: 'https://api.anthropic.com',
+      platform: 'claude',
+      status: 'active',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      username: 'claude-user@example.com',
+      accessToken: 'claude-oauth-token',
+      apiToken: null,
+      status: 'active',
+      extraConfig: JSON.stringify({
+        credentialMode: 'session',
+        oauth: {
+          provider: 'claude',
+          email: 'claude-user@example.com',
+          accountKey: 'claude-user@example.com',
+        },
+      }),
+    }).returning().get();
+
+    const result = await refreshModelsForAccount(account.id);
+
+    expect(result).toMatchObject({
+      accountId: account.id,
+      refreshed: true,
+      status: 'success',
+      errorCode: null,
+      tokenScanned: 0,
+      discoveredByCredential: true,
+      discoveredApiToken: false,
+      modelCount: 2,
+      modelsPreview: ['claude-3-7-sonnet-latest', 'claude-opus-4-1-20250805'],
+    });
+    expect(getModelsMock).not.toHaveBeenCalled();
+    expect(undiciFetchMock).toHaveBeenCalledTimes(1);
+    expect(String(undiciFetchMock.mock.calls[0]?.[0] || '')).toBe('https://api.anthropic.com/v1/models');
+    expect(undiciFetchMock.mock.calls[0]?.[1]).toMatchObject({
+      method: 'GET',
+      headers: expect.objectContaining({
+        Authorization: 'Bearer claude-oauth-token',
+        Accept: 'application/json',
+        'anthropic-version': '2023-06-01',
+      }),
+    });
+
+    const rows = await db.select().from(schema.modelAvailability)
+      .where(eq(schema.modelAvailability.accountId, account.id))
+      .all();
+    expect(rows.map((row) => row.modelName).sort()).toEqual([
+      'claude-3-7-sonnet-latest',
+      'claude-opus-4-1-20250805',
+    ]);
+  });
+
  it('marks codex oauth account abnormal when upstream cloud discovery fails', async () => {
    getApiTokenMock.mockResolvedValue(null);
    getModelsMock.mockRejectedValue(new Error('codex plan discovery should not call adapter.getModels'));
@@ -40,18 +40,6 @@ const API_TOKEN_DISCOVERY_TIMEOUT_MS = 8_000;
 const MODEL_DISCOVERY_TIMEOUT_MS = 12_000;
 const MODEL_REFRESH_BATCH_SIZE = 3;
 const CODEX_MODELS_CLIENT_VERSION = '1.0.0';
-const CLAUDE_STATIC_MODELS = [
-  'claude-haiku-4-5-20251001',
-  'claude-sonnet-4-5-20250929',
-  'claude-sonnet-4-6',
-  'claude-opus-4-6',
-  'claude-opus-4-5-20251101',
-  'claude-opus-4-1-20250805',
-  'claude-opus-4-20250514',
-  'claude-sonnet-4-20250514',
-  'claude-3-7-sonnet-20250219',
-  'claude-3-5-haiku-20241022',
-];
 const GEMINI_CLI_STATIC_MODELS = [
  'gemini-2.5-pro',
  'gemini-2.5-flash',
@@ -190,6 +178,21 @@ function extractCodexModelIds(payload: unknown): string[] {
  });
 }

+function extractClaudeModelIds(payload: unknown): string[] {
+  if (!payload || typeof payload !== 'object' || Array.isArray(payload)) return [];
+  const record = payload as { data?: unknown };
+  if (!Array.isArray(record.data)) return [];
+
+  return record.data.flatMap((item) => {
+    if (!item || typeof item !== 'object' || Array.isArray(item)) return [];
+    const value = item as { id?: unknown; name?: unknown };
+    const id = typeof value.id === 'string'
+      ? value.id.trim()
+      : (typeof value.name === 'string' ? value.name.trim() : '');
+    return id ? [id] : [];
+  });
+}
+
 function normalizeBaseUrl(baseUrl: string): string {
  return (baseUrl || '').replace(/\/+$/, '');
 }
@@ -284,10 +287,10 @@ async function discoverCodexModelsFromCloud(input: {
  return normalizeModels(extractCodexModelIds(await response.json()));
 }

-async function validateClaudeOauthConnection(input: {
+async function discoverClaudeModelsFromCloud(input: {
  site: typeof schema.sites.$inferSelect;
  account: typeof schema.accounts.$inferSelect;
-}): Promise<void> {
+}): Promise<string[]> {
  const accessToken = (input.account.accessToken || '').trim();
  if (!accessToken) {
    throw new Error('claude oauth access token missing');
@@ -305,8 +308,9 @@ async function validateClaudeOauthConnection(input: {
  );
  if (!response.ok) {
    const text = await response.text().catch(() => '');
-    throw new Error(`HTTP ${response.status}: ${text || 'claude oauth validation failed'}`);
+    throw new Error(`HTTP ${response.status}: ${text || 'claude oauth model discovery failed'}`);
  }
+  return normalizeModels(extractClaudeModelIds(await response.json()));
 }

 async function validateGeminiCliOauthConnection(input: {
@@ -595,14 +599,17 @@ export async function refreshModelsForAccount(accountId: number): Promise<ModelR
    const checkedAt = new Date().toISOString();
    const startedAt = Date.now();
    try {
-      await withTimeout(
+      const claudeModels = await withTimeout(
        () => withAccountProxyOverride(accountProxyUrl,
-          () => validateClaudeOauthConnection({ site, account })),
+          () => discoverClaudeModelsFromCloud({ site, account })),
        MODEL_DISCOVERY_TIMEOUT_MS,
-        `claude oauth validation timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`,
+        `claude oauth model discovery timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`,
      );
+      if (claudeModels.length === 0) {
+        throw new Error('未获取到可用模型');
+      }
      await db.insert(schema.modelAvailability).values(
-        CLAUDE_STATIC_MODELS.map((modelName) => ({
+        claudeModels.map((modelName) => ({
          accountId,
          modelName,
          available: true,
@@ -614,26 +621,26 @@ export async function refreshModelsForAccount(accountId: number): Promise<ModelR
        account,
        checkedAt,
        status: 'healthy',
-        lastDiscoveredModels: CLAUDE_STATIC_MODELS,
+        lastDiscoveredModels: claudeModels,
      });
      await setAccountRuntimeHealth(accountId, {
        state: 'healthy',
-        reason: 'Claude OAuth 健康探测成功',
+        reason: 'Claude OAuth 模型探测成功',
        source: 'model-discovery',
        checkedAt,
      });
      return buildSuccessfulRefreshResult({
        accountId,
-        modelCount: CLAUDE_STATIC_MODELS.length,
-        modelsPreview: CLAUDE_STATIC_MODELS.slice(0, 10),
+        modelCount: claudeModels.length,
+        modelsPreview: claudeModels.slice(0, 10),
        tokenScanned: 0,
        discoveredByCredential: true,
        discoveredApiToken: false,
      });
    } catch (err) {
-      const rawMessage = (err as { message?: string })?.message || 'claude oauth validation failed';
+      const rawMessage = (err as { message?: string })?.message || 'claude oauth model discovery failed';
      const errorCode = classifyModelDiscoveryError(rawMessage);
-      const errorMessage = `Claude 模型获取失败（${rawMessage}）`;
+      const errorMessage = `Claude OAuth 模型获取失败（${rawMessage}）`;
      await updateOauthModelDiscoveryState({
        account,
        checkedAt,