fix: address gemini native and route coverage regressions

2026-03-15 21:47:00 +08:00
parent a64a51aa0c
commit 4b42751953
15 changed files with 1006 additions and 74 deletions
@@ -0,0 +1,203 @@
+import Fastify, { type FastifyInstance } from 'fastify';
+import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+import { mkdtempSync } from 'node:fs';
+import { tmpdir } from 'node:os';
+import { join } from 'node:path';
+import { eq } from 'drizzle-orm';
+
+const getApiTokensMock = vi.fn();
+const getApiTokenMock = vi.fn();
+const createApiTokenMock = vi.fn();
+const getUserGroupsMock = vi.fn();
+const deleteApiTokenMock = vi.fn();
+const getModelsMock = vi.fn();
+
+vi.mock('../../services/platforms/index.js', () => ({
+  getAdapter: () => ({
+    getApiTokens: (...args: unknown[]) => getApiTokensMock(...args),
+    getApiToken: (...args: unknown[]) => getApiTokenMock(...args),
+    createApiToken: (...args: unknown[]) => createApiTokenMock(...args),
+    getUserGroups: (...args: unknown[]) => getUserGroupsMock(...args),
+    deleteApiToken: (...args: unknown[]) => deleteApiTokenMock(...args),
+    getModels: (...args: unknown[]) => getModelsMock(...args),
+  }),
+}));
+
+type DbModule = typeof import('../../db/index.js');
+
+describe('account token coverage refresh', () => {
+  let app: FastifyInstance;
+  let db: DbModule['db'];
+  let schema: DbModule['schema'];
+  let dataDir = '';
+  let seedId = 0;
+
+  const nextSeed = () => {
+    seedId += 1;
+    return seedId;
+  };
+
+  const seedAccount = async (modelName: string) => {
+    const id = nextSeed();
+    const site = await db.insert(schema.sites).values({
+      name: `site-${id}`,
+      url: `https://site-${id}.example.com`,
+      platform: 'new-api',
+      status: 'active',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      username: `user-${id}`,
+      accessToken: `acc-token-${id}`,
+      status: 'active',
+    }).returning().get();
+
+    await db.insert(schema.modelAvailability).values({
+      accountId: account.id,
+      modelName,
+      available: true,
+    }).run();
+
+    return { site, account };
+  };
+
+  const readTokenCandidates = async () => {
+    const response = await app.inject({
+      method: 'GET',
+      url: '/api/models/token-candidates',
+    });
+    expect(response.statusCode).toBe(200);
+    return response.json() as {
+      models: Record<string, Array<{ tokenId: number; accountId: number }>>;
+      modelsWithoutToken: Record<string, Array<{ accountId: number }>>;
+    };
+  };
+
+  beforeAll(async () => {
+    dataDir = mkdtempSync(join(tmpdir(), 'metapi-account-token-coverage-'));
+    process.env.DATA_DIR = dataDir;
+
+    await import('../../db/migrate.js');
+    const dbModule = await import('../../db/index.js');
+    const accountTokenRoutesModule = await import('./accountTokens.js');
+    const statsRoutesModule = await import('./stats.js');
+    db = dbModule.db;
+    schema = dbModule.schema;
+
+    app = Fastify();
+    await app.register(accountTokenRoutesModule.accountTokensRoutes);
+    await app.register(statsRoutesModule.statsRoutes);
+  });
+
+  beforeEach(async () => {
+    getApiTokensMock.mockReset();
+    getApiTokenMock.mockReset();
+    createApiTokenMock.mockReset();
+    getUserGroupsMock.mockReset();
+    deleteApiTokenMock.mockReset();
+    getModelsMock.mockReset();
+    seedId = 0;
+
+    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.accounts).run();
+    await db.delete(schema.sites).run();
+  });
+
+  afterAll(async () => {
+    await app.close();
+    delete process.env.DATA_DIR;
+  });
+
+  it('refreshes token coverage after manually adding an account token', async () => {
+    const modelName = 'gpt-4o-mini';
+    const { account } = await seedAccount(modelName);
+    getModelsMock.mockImplementation(async (_url: string, credential: string) => {
+      if (credential === account.accessToken || credential === 'sk-manual-token') {
+        return [modelName];
+      }
+      return [];
+    });
+
+    const response = await app.inject({
+      method: 'POST',
+      url: '/api/account-tokens',
+      payload: {
+        accountId: account.id,
+        name: 'manual-token',
+        token: 'sk-manual-token',
+      },
+    });
+
+    expect(response.statusCode).toBe(200);
+
+    const token = await db.select()
+      .from(schema.accountTokens)
+      .where(eq(schema.accountTokens.accountId, account.id))
+      .get();
+    expect(token).toBeDefined();
+
+    const tokenAvailability = await db.select()
+      .from(schema.tokenModelAvailability)
+      .where(eq(schema.tokenModelAvailability.tokenId, token!.id))
+      .all();
+    expect(tokenAvailability.map((row) => row.modelName)).toContain(modelName);
+
+    const candidates = await readTokenCandidates();
+    expect(candidates.modelsWithoutToken[modelName]).toBeUndefined();
+    expect(candidates.models[modelName]).toEqual([
+      expect.objectContaining({
+        accountId: account.id,
+        tokenId: token!.id,
+      }),
+    ]);
+  });
+
+  it('refreshes token coverage after syncing account tokens from upstream', async () => {
+    const modelName = 'claude-sonnet-4-5-20250929';
+    const { account } = await seedAccount(modelName);
+    getApiTokensMock.mockResolvedValue([
+      { name: 'default', key: 'sk-synced-token', enabled: true },
+    ]);
+    getModelsMock.mockImplementation(async (_url: string, credential: string) => {
+      if (credential === account.accessToken || credential === 'sk-synced-token') {
+        return [modelName];
+      }
+      return [];
+    });
+
+    const response = await app.inject({
+      method: 'POST',
+      url: `/api/account-tokens/sync/${account.id}`,
+    });
+
+    expect(response.statusCode).toBe(200);
+
+    const token = await db.select()
+      .from(schema.accountTokens)
+      .where(eq(schema.accountTokens.accountId, account.id))
+      .get();
+    expect(token).toBeDefined();
+
+    const tokenAvailability = await db.select()
+      .from(schema.tokenModelAvailability)
+      .where(eq(schema.tokenModelAvailability.tokenId, token!.id))
+      .all();
+    expect(tokenAvailability.map((row) => row.modelName)).toContain(modelName);
+
+    const candidates = await readTokenCandidates();
+    expect(candidates.modelsWithoutToken[modelName]).toBeUndefined();
+    expect(candidates.models[modelName]).toEqual([
+      expect.objectContaining({
+        accountId: account.id,
+        tokenId: token!.id,
+      }),
+    ]);
+  });
+});
@@ -14,6 +14,7 @@ import {
 import { getAdapter } from '../../services/platforms/index.js';
 import { getCredentialModeFromExtraConfig, resolvePlatformUserId } from '../../services/accountExtraConfig.js';
 import { startBackgroundTask } from '../../services/backgroundTaskService.js';
+import { rebuildTokenRoutesFromAvailability, refreshModelsForAccount } from '../../services/modelService.js';

 type AccountWithSiteRow = {
  accounts: typeof schema.accounts.$inferSelect;
@@ -337,6 +338,12 @@ async function executeSyncAllAccountTokens() {
    results.push(...batchResults);
  }

+  const coverageRefresh = await refreshCoverageForAccounts(
+    results
+      .filter((item) => item.status === 'synced')
+      .map((item) => item.accountId),
+  );
+
  const summary = {
    total: results.length,
    synced: results.filter((item) => item.status === 'synced').length,
@@ -346,7 +353,23 @@ async function executeSyncAllAccountTokens() {
    updated: results.reduce((acc, item) => acc + item.updated, 0),
  };

-  return { summary, results };
+  return { summary, results, coverageRefresh };
+}
+
+async function refreshCoverageForAccounts(accountIds: number[]) {
+  const uniqueAccountIds = Array.from(new Set(
+    accountIds.filter((id) => Number.isFinite(id) && id > 0),
+  ));
+
+  if (uniqueAccountIds.length === 0) {
+    return { refresh: [], rebuild: null };
+  }
+
+  const refresh = await Promise.all(
+    uniqueAccountIds.map(async (accountId) => refreshModelsForAccount(accountId)),
+  );
+  const rebuild = await rebuildTokenRoutesFromAvailability();
+  return { refresh, rebuild };
 }

 export async function accountTokensRoutes(app: FastifyInstance) {
@@ -416,8 +439,8 @@ export async function accountTokensRoutes(app: FastifyInstance) {
      } else if (existing.every((token) => !token.isDefault) && (body.enabled ?? true)) {
        await setDefaultToken(created.id);
      }
-
-      return { success: true, token: created };
+      const coverageRefresh = await refreshCoverageForAccounts([body.accountId]);
+      return { success: true, token: created, coverageRefresh };
    }

    const account = row.accounts;
@@ -496,6 +519,7 @@ export async function accountTokensRoutes(app: FastifyInstance) {
    if (syncResult.status === 'skipped') {
      return reply.code(502).send({ success: false, message: syncResult.message || '站点未返回可用令牌' });
    }
+    const coverageRefresh = await refreshCoverageForAccounts([account.id]);

    const preferred = await db.select().from(schema.accountTokens)
      .where(and(eq(schema.accountTokens.accountId, account.id), eq(schema.accountTokens.isDefault, true)))
@@ -509,6 +533,7 @@ export async function accountTokensRoutes(app: FastifyInstance) {
      success: true,
      createdViaUpstream: true,
      ...syncResult,
+      coverageRefresh,
      token,
    };
  });
@@ -823,7 +848,8 @@ export async function accountTokensRoutes(app: FastifyInstance) {
    if (result.status === 'failed') {
      return reply.code(502).send({ success: false, message: result.message || '同步失败' });
    }
-    return { success: true, ...result };
+    const coverageRefresh = await refreshCoverageForAccounts([accountId]);
+    return { success: true, ...result, coverageRefresh };
  });

  app.post<{ Body?: { wait?: boolean } }>('/api/account-tokens/sync-all', async (request, reply) => {