1
0
forked from boos/BoosAPI

Fix Claude count tokens routing and model discovery

This commit is contained in:
Cita
2026-03-21 18:42:28 +08:00
parent 6a19d809f0
commit dc93c58582
5 changed files with 201 additions and 46 deletions
+11 -3
View File
@@ -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');
});
});
+35 -7
View File
@@ -1521,10 +1521,12 @@ 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,
@@ -1535,10 +1537,7 @@ export function buildClaudeCountTokensUpstreamRequest(input: {
baseHeaders: {
'Content-Type': 'application/json',
},
claudeHeaders: {
...claudeHeaders,
...(betas.length > 0 ? { 'anthropic-beta': betas.join(',') } : {}),
},
claudeHeaders: effectiveClaudeHeaders,
body: sanitizedBody,
action: 'countTokens',
});
@@ -1555,3 +1554,32 @@ export function buildClaudeCountTokensUpstreamRequest(input: {
},
};
}
const anthropicVersion = (
effectiveClaudeHeaders['anthropic-version']
|| '2023-06-01'
);
const isClaudeOauthUpstream = sitePlatform === 'claude' && input.oauthProvider === 'claude';
const headers = buildClaudeRuntimeHeaders({
baseHeaders: {
'Content-Type': 'application/json',
},
claudeHeaders: effectiveClaudeHeaders,
anthropicVersion,
stream: false,
isClaudeOauthUpstream,
tokenValue: input.tokenValue,
});
return {
path: '/v1/messages/count_tokens?beta=true',
headers,
body: sanitizedBody,
runtime: {
executor: 'claude',
modelName: input.modelName,
stream: false,
action: 'countTokens',
},
};
}
@@ -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'));
+32 -25
View File
@@ -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,