feat: Add manual model availability management, including database schema updates, new API routes, and UI components.

2026-03-14 19:26:13 -07:00
parent 5da3fbf874
commit de22fc35fa
17 changed files with 586 additions and 163 deletions
@@ -0,0 +1,133 @@
+import Fastify, { type FastifyInstance } from 'fastify';
+import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest';
+import { mkdtempSync } from 'node:fs';
+import { tmpdir } from 'node:os';
+import { join } from 'node:path';
+import { eq } from 'drizzle-orm';
+
+type DbModule = typeof import('../../db/index.js');
+
+describe('accounts manual models endpoint', () => {
+  let app: FastifyInstance;
+  let db: DbModule['db'];
+  let schema: DbModule['schema'];
+  let dataDir = '';
+
+  beforeAll(async () => {
+    dataDir = mkdtempSync(join(tmpdir(), 'metapi-accounts-manual-models-'));
+    process.env.DATA_DIR = dataDir;
+
+    await import('../../db/migrate.js');
+    const dbModule = await import('../../db/index.js');
+    const routesModule = await import('./accounts.js');
+    db = dbModule.db;
+    schema = dbModule.schema;
+
+    app = Fastify();
+    await app.register(routesModule.accountsRoutes);
+  });
+
+  beforeEach(async () => {
+    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('adds manual models and sets isManual to true', async () => {
+    const site = await db.insert(schema.sites).values({
+      name: 'Test Site',
+      url: 'https://test.example.com',
+      platform: 'new-api',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      accessToken: 'test-token',
+    }).returning().get();
+
+    const response = await app.inject({
+      method: 'POST',
+      url: `/api/accounts/${account.id}/models/manual`,
+      payload: {
+        models: ['gpt-4-manual', 'claude-3-manual'],
+      },
+    });
+
+    expect(response.statusCode).toBe(200);
+    const body = response.json();
+    expect(body.success).toBe(true);
+
+    const models = await db.select().from(schema.modelAvailability).where(
+      eq(schema.modelAvailability.accountId, account.id)
+    ).all();
+    
+    expect(models).toHaveLength(2);
+    expect(models.map(m => m.modelName).sort()).toEqual(['claude-3-manual', 'gpt-4-manual']);
+    expect(models[0]?.isManual).toBe(true);
+    expect(models[1]?.isManual).toBe(true);
+  });
+
+  it('updates existing synced models to manual if provided', async () => {
+    const site = await db.insert(schema.sites).values({
+      name: 'Test Site',
+      url: 'https://test.example.com',
+      platform: 'new-api',
+    }).returning().get();
+
+    const account = await db.insert(schema.accounts).values({
+      siteId: site.id,
+      accessToken: 'test-token',
+    }).returning().get();
+
+    // Already-synced model that is NOT manual
+    await db.insert(schema.modelAvailability).values({
+      accountId: account.id,
+      modelName: 'gpt-existing',
+      available: true,
+    });
+
+    const response = await app.inject({
+      method: 'POST',
+      url: `/api/accounts/${account.id}/models/manual`,
+      payload: {
+        models: ['gpt-existing', 'gpt-new'],
+      },
+    });
+
+    expect(response.statusCode).toBe(200);
+
+    const models = await db.select().from(schema.modelAvailability)
+      .where(eq(schema.modelAvailability.accountId, account.id))
+      .all();
+    
+    expect(models).toHaveLength(2);
+    const existing = models.find(m => m.modelName === 'gpt-existing');
+    const newModel = models.find(m => m.modelName === 'gpt-new');
+
+    expect(existing?.isManual).toBe(true); // Should be updated
+    expect(newModel?.isManual).toBe(true);
+  });
+
+  it('fails if account does not exist', async () => {
+    const response = await app.inject({
+      method: 'POST',
+      url: '/api/accounts/999/models/manual',
+      payload: {
+        models: ['gpt-4-manual'],
+      },
+    });
+
+    expect(response.statusCode).toBe(404);
+  });
+});
@@ -0,0 +1,120 @@
+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';
+
+const getModelsMock = vi.fn();
+
+vi.mock('../../services/platforms/index.js', () => ({
+  getAdapter: () => ({
+    getModels: (...args: unknown[]) => getModelsMock(...args),
+  }),
+}));
+
+type DbModule = typeof import('../../db/index.js');
+
+describe('accounts skipModelFetch behavior', () => {
+  let app: FastifyInstance;
+  let db: DbModule['db'];
+  let schema: DbModule['schema'];
+  let dataDir = '';
+
+  beforeAll(async () => {
+    dataDir = mkdtempSync(join(tmpdir(), 'metapi-accounts-skip-model-'));
+    process.env.DATA_DIR = dataDir;
+
+    await import('../../db/migrate.js');
+    const dbModule = await import('../../db/index.js');
+    const routesModule = await import('./accounts.js');
+    db = dbModule.db;
+    schema = dbModule.schema;
+
+    app = Fastify();
+    await app.register(routesModule.accountsRoutes);
+  });
+
+  beforeEach(async () => {
+    getModelsMock.mockReset();
+
+    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('skips model fetching when skipModelFetch is true under apikey credentialMode', async () => {
+    getModelsMock.mockRejectedValueOnce(new Error('Should not be called'));
+
+    const site = await db.insert(schema.sites).values({
+      name: 'API Key Site',
+      url: 'https://apikey.example.com',
+      platform: 'new-api',
+    }).returning().get();
+
+    const response = await app.inject({
+      method: 'POST',
+      url: '/api/accounts',
+      payload: {
+        siteId: site.id,
+        accessToken: 'sk-test-skip-fetch',
+        credentialMode: 'apikey',
+        skipModelFetch: true,
+      },
+    });
+
+    expect(response.statusCode).toBe(200);
+    // getModels is called ONCE asynchronously by refreshModelsForAccount later in the flow
+    expect(getModelsMock).toHaveBeenCalledTimes(1);
+
+    const accounts = await db.select().from(schema.accounts).all();
+    expect(accounts).toHaveLength(1);
+    expect(accounts[0]?.apiToken).toBe('sk-test-skip-fetch');
+    
+    // Model availability should be empty initially (background task might not have populated it yet or failed)
+    const models = await db.select().from(schema.modelAvailability).all();
+    expect(models).toHaveLength(0);
+  });
+
+  it('still calls getModels when skipModelFetch is false', async () => {
+    getModelsMock.mockResolvedValue(['gpt-4']);
+
+    const site = await db.insert(schema.sites).values({
+      name: 'API Key Site',
+      url: 'https://apikey.example.com',
+      platform: 'new-api',
+    }).returning().get();
+
+    const response = await app.inject({
+      method: 'POST',
+      url: '/api/accounts',
+      payload: {
+        siteId: site.id,
+        accessToken: 'sk-test-normal-fetch',
+        credentialMode: 'apikey',
+        skipModelFetch: false,
+      },
+    });
+
+    expect(response.statusCode).toBe(200);
+    // getModels is called TWICE: once for block validation, once asynchronously by refreshModelsForAccount
+    expect(getModelsMock).toHaveBeenCalledTimes(2);
+
+    const accounts = await db.select().from(schema.accounts).all();
+    expect(accounts).toHaveLength(1);
+    expect(accounts[0]?.apiToken).toBe('sk-test-normal-fetch');
+
+    // Model availability should be populated since getModels was called (which refreshModels uses later or handled directly)
+    // Actually our POST /api/accounts triggers rebuildTokenRoutesFromAvailability and refreshModelsForAccount asynchronously, so models might not be populated synchronously, but the mock should be called.
+  });
+});
@@ -954,7 +954,7 @@ export async function accountsRoutes(app: FastifyInstance) {
  );

  // Add an account (manual credential input)
-  app.post<{ Body: { siteId: number; username?: string; accessToken: string; apiToken?: string; platformUserId?: number; checkinEnabled?: boolean; credentialMode?: AccountCredentialMode; refreshToken?: string; tokenExpiresAt?: number | string } }>('/api/accounts', async (request, reply) => {
+  app.post<{ Body: { siteId: number; username?: string; accessToken: string; apiToken?: string; platformUserId?: number; checkinEnabled?: boolean; credentialMode?: AccountCredentialMode; refreshToken?: string; tokenExpiresAt?: number | string; skipModelFetch?: boolean } }>('/api/accounts', async (request, reply) => {
    const body = request.body;
    const site = await db.select().from(schema.sites).where(eq(schema.sites.id, body.siteId)).get();
    if (!site) {
@@ -979,29 +979,35 @@ export async function accountsRoutes(app: FastifyInstance) {
    let verifiedModels: string[] = [];

    if (credentialMode === 'apikey') {
-      try {
-        const models = await adapter.getModels(site.url, rawAccessToken, body.platformUserId);
-        verifiedModels = Array.isArray(models)
-          ? models.filter((item) => typeof item === 'string' && item.trim().length > 0)
-          : [];
-      } catch (err: any) {
-        return reply.code(400).send({
-          success: false,
-          message: err?.message || 'API Key 验证失败',
-        });
-      }
+      if (body.skipModelFetch) {
+        tokenType = 'apikey';
+        accessToken = '';
+        if (!apiToken) apiToken = rawAccessToken;
+      } else {
+        try {
+          const models = await adapter.getModels(site.url, rawAccessToken, body.platformUserId);
+          verifiedModels = Array.isArray(models)
+            ? models.filter((item) => typeof item === 'string' && item.trim().length > 0)
+            : [];
+        } catch (err: any) {
+          return reply.code(400).send({
+            success: false,
+            message: err?.message || 'API Key 验证失败',
+          });
+        }

-      if (verifiedModels.length === 0) {
-        return reply.code(400).send({
-          success: false,
-          requiresVerification: true,
-          message: 'API Key 验证失败：未获取到可用模型',
-        });
-      }
+        if (verifiedModels.length === 0) {
+          return reply.code(400).send({
+            success: false,
+            requiresVerification: true,
+            message: 'API Key 验证失败：未获取到可用模型',
+          });
+        }

-      tokenType = 'apikey';
-      accessToken = '';
-      if (!apiToken) apiToken = rawAccessToken;
+        tokenType = 'apikey';
+        accessToken = '';
+        if (!apiToken) apiToken = rawAccessToken;
+      }
    } else {
      let verifyResult: any;
      try {
@@ -1367,6 +1373,7 @@ export async function accountsRoutes(app: FastifyInstance) {
      modelName: schema.modelAvailability.modelName,
      available: schema.modelAvailability.available,
      latencyMs: schema.modelAvailability.latencyMs,
+      isManual: schema.modelAvailability.isManual,
    }).from(schema.modelAvailability)
      .where(eq(schema.modelAvailability.accountId, accountId))
      .all();
@@ -1386,6 +1393,7 @@ export async function accountsRoutes(app: FastifyInstance) {
        name: r.modelName,
        latencyMs: r.latencyMs,
        disabled: disabledSet.has(r.modelName),
+        isManual: !!r.isManual,
      }))
      .sort((a, b) => a.name.localeCompare(b.name));

@@ -1397,5 +1405,66 @@ export async function accountsRoutes(app: FastifyInstance) {
      disabledCount: models.filter((m) => m.disabled).length,
    };
  });
+
+  // Add models manually to an account
+  app.post<{ Params: { id: string }; Body: { models: string[] } }>('/api/accounts/:id/models/manual', async (request, reply) => {
+    const accountId = parseInt(request.params.id, 10);
+    if (!Number.isFinite(accountId) || accountId <= 0) {
+      return reply.code(400).send({ message: '账号 ID 无效' });
+    }
+
+    const { models } = request.body;
+    if (!Array.isArray(models) || models.length === 0) {
+      return reply.code(400).send({ message: '模型列表不能为空' });
+    }
+
+    const normalizedModels = models.map(m => String(m).trim()).filter(m => m.length > 0);
+    if (normalizedModels.length === 0) {
+      return reply.code(400).send({ message: '模型列表不能为空' });
+    }
+
+    const account = await db.select().from(schema.accounts)
+      .where(eq(schema.accounts.id, accountId))
+      .get();
+
+    if (!account) {
+      return reply.code(404).send({ message: '账号不存在' });
+    }
+
+    try {
+      for (const modelName of normalizedModels) {
+        const existing = await db.select()
+          .from(schema.modelAvailability)
+          .where(and(eq(schema.modelAvailability.accountId, accountId), eq(schema.modelAvailability.modelName, modelName)))
+          .get();
+
+        if (existing) {
+          if (!existing.available || !existing.isManual) {
+            await db.update(schema.modelAvailability)
+              .set({ available: true, latencyMs: null, isManual: true, checkedAt: new Date().toISOString() })
+              .where(eq(schema.modelAvailability.id, existing.id))
+              .run();
+          }
+        } else {
+          await db.insert(schema.modelAvailability).values({
+            accountId,
+            modelName,
+            available: true,
+            isManual: true,
+            latencyMs: null,
+            checkedAt: new Date().toISOString(),
+          }).run();
+        }
+      }
+
+      try {
+        await rebuildTokenRoutesFromAvailability();
+      } catch { }
+
+      return { success: true };
+    } catch (err: any) {
+      return reply.code(500).send({ success: false, message: err?.message || '保存失败' });
+    }
+  });
 }

@@ -352,6 +352,12 @@ describe('/v1/models route', () => {
      },
    ]).run();

+    await db.insert(schema.downstreamApiKeys).values({
+      name: 'search-key',
+      key: 'change-me-proxy-sk-token',
+      enabled: true,
+    }).run();
+
    const response = await app.inject({
      method: 'GET',
      url: '/v1/models',