fix connection and release regressions

This commit is contained in:
cita
2026-03-18 18:00:59 +08:00
parent d928e067c1
commit 38e46e165b
16 changed files with 528 additions and 63 deletions
@@ -0,0 +1,178 @@
import Fastify, { type FastifyInstance } from 'fastify';
import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
import { mkdtempSync, rmSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
const verifyTokenMock = vi.fn();
const getApiTokensMock = vi.fn();
const refreshBalanceMock = vi.fn();
const refreshModelsForAccountMock = vi.fn();
const rebuildTokenRoutesFromAvailabilityMock = vi.fn();
const ensureDefaultTokenForAccountMock = vi.fn();
const syncTokensFromUpstreamMock = vi.fn();
vi.mock('../../services/platforms/index.js', () => ({
getAdapter: () => ({
verifyToken: (...args: unknown[]) => verifyTokenMock(...args),
getApiTokens: (...args: unknown[]) => getApiTokensMock(...args),
}),
}));
vi.mock('../../services/balanceService.js', () => ({
refreshBalance: (...args: unknown[]) => refreshBalanceMock(...args),
}));
vi.mock('../../services/modelService.js', () => ({
refreshModelsForAccount: (...args: unknown[]) => refreshModelsForAccountMock(...args),
rebuildTokenRoutesFromAvailability: (...args: unknown[]) => rebuildTokenRoutesFromAvailabilityMock(...args),
}));
vi.mock('../../services/accountTokenService.js', () => ({
ensureDefaultTokenForAccount: (...args: unknown[]) => ensureDefaultTokenForAccountMock(...args),
syncTokensFromUpstream: (...args: unknown[]) => syncTokensFromUpstreamMock(...args),
}));
type DbModule = typeof import('../../db/index.js');
describe('accounts background initialization', () => {
let app: FastifyInstance;
let db: DbModule['db'];
let schema: DbModule['schema'];
let dataDir = '';
let resetBackgroundTasks: (() => void) | null = null;
let getBackgroundTask: ((taskId: string) => { status: string } | null) | null = null;
beforeAll(async () => {
dataDir = mkdtempSync(join(tmpdir(), 'metapi-accounts-background-init-'));
process.env.DATA_DIR = dataDir;
await import('../../db/migrate.js');
const dbModule = await import('../../db/index.js');
const routesModule = await import('./accounts.js');
const backgroundTaskModule = await import('../../services/backgroundTaskService.js');
db = dbModule.db;
schema = dbModule.schema;
resetBackgroundTasks = backgroundTaskModule.__resetBackgroundTasksForTests;
getBackgroundTask = backgroundTaskModule.getBackgroundTask;
app = Fastify();
await app.register(routesModule.accountsRoutes);
});
beforeEach(async () => {
verifyTokenMock.mockReset();
getApiTokensMock.mockReset();
refreshBalanceMock.mockReset();
refreshModelsForAccountMock.mockReset();
rebuildTokenRoutesFromAvailabilityMock.mockReset();
ensureDefaultTokenForAccountMock.mockReset();
syncTokensFromUpstreamMock.mockReset();
resetBackgroundTasks?.();
await db.delete(schema.proxyLogs).run();
await db.delete(schema.checkinLogs).run();
await db.delete(schema.routeChannels).run();
await db.delete(schema.tokenRoutes).run();
await db.delete(schema.tokenModelAvailability).run();
await db.delete(schema.modelAvailability).run();
await db.delete(schema.accountTokens).run();
await db.delete(schema.events).run();
await db.delete(schema.accounts).run();
await db.delete(schema.sites).run();
});
afterAll(async () => {
await app.close();
if (dataDir) {
try {
rmSync(dataDir, { recursive: true, force: true });
} catch {}
}
delete process.env.DATA_DIR;
});
it('returns immediately and queues background initialization when token sync is slow', async () => {
const site = await db.insert(schema.sites).values({
name: 'Session Site',
url: 'https://session.example.com',
platform: 'new-api',
}).returning().get();
verifyTokenMock.mockResolvedValue({
tokenType: 'session',
userInfo: { username: 'demo-user' },
apiToken: 'sk-demo',
});
ensureDefaultTokenForAccountMock.mockResolvedValue(undefined);
refreshBalanceMock.mockResolvedValue({ balance: 1, used: 0, quota: 1 });
refreshModelsForAccountMock.mockResolvedValue(undefined);
rebuildTokenRoutesFromAvailabilityMock.mockResolvedValue(undefined);
syncTokensFromUpstreamMock.mockResolvedValue(undefined);
let releaseTokens: ((value: Array<{ name: string; value: string }>) => void) | null = null;
const pendingTokens = new Promise<Array<{ name: string; value: string }>>((resolve) => {
releaseTokens = resolve;
});
getApiTokensMock.mockReturnValue(pendingTokens);
const responsePromise = app.inject({
method: 'POST',
url: '/api/accounts',
payload: {
siteId: site.id,
accessToken: 'session-token',
},
});
try {
const raceResult = await Promise.race([
responsePromise.then(() => 'response'),
new Promise<'timeout'>((resolve) => setTimeout(() => resolve('timeout'), 25)),
]);
expect(raceResult).toBe('response');
const response = await responsePromise;
expect(response.statusCode).toBe(200);
const body = response.json() as {
id: number;
queued?: boolean;
jobId?: string;
usernameDetected?: boolean;
apiTokenFound?: boolean;
};
expect(body).toMatchObject({
queued: true,
usernameDetected: true,
apiTokenFound: true,
});
expect(body.jobId).toBeTruthy();
const insertedAccounts = await db.select().from(schema.accounts).all();
expect(insertedAccounts).toHaveLength(1);
expect(getBackgroundTask?.(body.jobId!)).toMatchObject({
status: expect.stringMatching(/pending|running/),
});
releaseTokens?.([{ name: 'default', value: 'sk-demo' }]);
for (let attempt = 0; attempt < 20; attempt += 1) {
const task = getBackgroundTask?.(body.jobId!);
if (task?.status === 'succeeded') break;
await new Promise((resolve) => setTimeout(resolve, 0));
}
expect(syncTokensFromUpstreamMock).toHaveBeenCalledTimes(1);
expect(refreshBalanceMock).toHaveBeenCalledTimes(1);
expect(refreshModelsForAccountMock).toHaveBeenCalledTimes(1);
expect(rebuildTokenRoutesFromAvailabilityMock).toHaveBeenCalledTimes(1);
expect(getBackgroundTask?.(body.jobId!)).toMatchObject({ status: 'succeeded' });
} finally {
releaseTokens?.([]);
await responsePromise.catch(() => undefined);
}
});
});
+110 -26
View File
@@ -49,6 +49,17 @@ type AccountCapabilities = {
proxyOnly: boolean;
};
type AccountInitializationParams = {
accountId: number;
site: typeof schema.sites.$inferSelect;
adapter: NonNullable<ReturnType<typeof getAdapter>>;
tokenType: 'session' | 'apikey' | 'unknown';
accessToken: string;
apiToken: string;
platformUserId?: number;
skipModelFetch?: boolean;
};
type VerifyFailureReason = 'needs-user-id' | 'invalid-user-id' | 'shield-blocked' | null;
const limitAccountLogin = createRateLimitGuard({
@@ -123,6 +134,75 @@ function normalizePinnedFlag(input: unknown): boolean | null {
return null;
}
async function initializeAccountInBackground({
accountId,
site,
adapter,
tokenType,
accessToken,
apiToken,
platformUserId,
skipModelFetch,
}: AccountInitializationParams) {
const summary = {
accountId,
syncedTokenCount: 0,
refreshedBalance: false,
refreshedModels: false,
rebuiltRoutes: false,
};
if (tokenType === 'session' && apiToken) {
try {
await ensureDefaultTokenForAccount(accountId, apiToken, { name: 'default', source: 'manual' });
} catch {}
}
if (tokenType === 'session' && accessToken) {
try {
const syncedTokens = await adapter.getApiTokens(site.url, accessToken, platformUserId);
summary.syncedTokenCount = Array.isArray(syncedTokens) ? syncedTokens.length : 0;
if (summary.syncedTokenCount > 0) {
await syncTokensFromUpstream(accountId, syncedTokens);
}
} catch {}
}
if (tokenType === 'session') {
try {
await refreshBalance(accountId);
summary.refreshedBalance = true;
} catch {}
}
if (skipModelFetch !== true) {
try {
await refreshModelsForAccount(accountId);
summary.refreshedModels = true;
await rebuildTokenRoutesFromAvailability();
summary.rebuiltRoutes = true;
} catch {}
}
return summary;
}
function buildQueuedAccountInitializationMessage(
tokenType: 'session' | 'apikey' | 'unknown',
skipModelFetch?: boolean,
) {
if (tokenType === 'session' && skipModelFetch === true) {
return '账号已添加,后台正在同步令牌和余额信息。';
}
if (tokenType === 'session') {
return '账号已添加,后台正在同步令牌、余额和模型信息。';
}
if (skipModelFetch === true) {
return '已添加为 API Key 账号(可用于代理转发)。';
}
return '已添加为 API Key 账号,后台正在同步模型和路由信息。';
}
function normalizeSortOrder(input: unknown): number | null {
if (input === undefined || input === null || input === '') return null;
const parsed = Number.parseInt(String(input), 10);
@@ -1117,32 +1197,33 @@ export async function accountsRoutes(app: FastifyInstance) {
return reply.code(500).send({ success: false, message: '创建账号失败' });
}
if (tokenType === 'session' && apiToken) {
try {
await ensureDefaultTokenForAccount(result.id, apiToken, { name: 'default', source: 'manual' });
} catch { }
}
if (tokenType === 'session' && accessToken) {
try {
const syncedTokens = await adapter.getApiTokens(site.url, accessToken, resolvedPlatformUserId);
if (syncedTokens.length > 0) {
await syncTokensFromUpstream(result.id, syncedTokens);
}
} catch { }
}
// Try to refresh balance
if (tokenType === 'session') {
try { await refreshBalance(result.id); } catch { }
}
// Try to refresh models only if not skipping.
if (body.skipModelFetch !== true) {
try {
await refreshModelsForAccount(result.id);
await rebuildTokenRoutesFromAvailability();
} catch { }
const shouldQueueInitialization = tokenType === 'session' || body.skipModelFetch !== true;
let queuedTaskId: string | undefined;
let queuedMessage: string | undefined;
if (shouldQueueInitialization) {
const taskTitle = `初始化连接 #${result.id}`;
const { task } = startBackgroundTask(
{
type: 'account-init',
title: taskTitle,
dedupeKey: `account-init-${result.id}`,
notifyOnFailure: true,
successMessage: () => `${taskTitle}已完成`,
failureMessage: (currentTask) => `${taskTitle}失败:${currentTask.error || 'unknown error'}`,
},
async () => initializeAccountInBackground({
accountId: result.id,
site,
adapter,
tokenType,
accessToken,
apiToken,
platformUserId: resolvedPlatformUserId,
skipModelFetch: body.skipModelFetch,
}),
);
queuedTaskId = task.id;
queuedMessage = buildQueuedAccountInitializationMessage(tokenType, body.skipModelFetch);
}
const account = await db.select().from(schema.accounts).where(eq(schema.accounts.id, result.id)).get();
@@ -1158,6 +1239,9 @@ export async function accountsRoutes(app: FastifyInstance) {
modelCount: verifiedModels.length,
apiTokenFound: !!apiToken,
usernameDetected: !!(!body.username && username),
queued: !!queuedTaskId,
jobId: queuedTaskId,
message: queuedMessage,
};
});
+6 -6
View File
@@ -16,12 +16,12 @@ import {
saveRouteDecisionSnapshots,
} from '../../services/routeDecisionSnapshotStore.js';
function isExactModelPattern(modelPattern: string): boolean {
const normalized = modelPattern.trim();
if (!normalized) return false;
if (normalized.toLowerCase().startsWith('re:')) return false;
return !/[\*\?\[]/.test(normalized);
}
function isExactModelPattern(modelPattern: string): boolean {
const normalized = modelPattern.trim();
if (!normalized) return false;
if (normalized.toLowerCase().startsWith('re:')) return false;
return !/[\*\?]/.test(normalized);
}
async function getDefaultTokenId(accountId: number): Promise<number | null> {
const token = await db.select().from(schema.accountTokens)
+1 -1
View File
@@ -390,7 +390,7 @@ function isExactModelPattern(modelPattern: string): boolean {
const normalized = modelPattern.trim();
if (!normalized) return false;
if (normalized.toLowerCase().startsWith('re:')) return false;
return !/[\*\?\[]/.test(normalized);
return !/[\*\?]/.test(normalized);
}
async function withTimeout<T>(fn: () => Promise<T>, timeoutMs: number, timeoutMessage: string): Promise<T> {