From 47fb4fc221b80c63fe82a47c2075524fdcd8c50c Mon Sep 17 00:00:00 2001 From: "kiloconnect[bot]" <240665456+kiloconnect[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 14:08:13 +0000 Subject: [PATCH] refactor(web): read gateway model metadata from postgres --- .../api/openrouter/[...path]/route.test.ts | 4 +-- .../ai-gateway/auto-model/resolution.test.ts | 2 +- .../lib/ai-gateway/auto-model/resolution.ts | 4 +-- apps/web/src/lib/ai-gateway/byok/index.ts | 4 +-- .../providers/gateway-models-cache.ts | 30 +++++++++++-------- .../providers/openrouter/index.test.ts | 2 +- .../ai-gateway/providers/openrouter/index.ts | 4 +-- .../providers/openrouter/sync-providers.ts | 13 +++++++- .../lib/ai-gateway/providers/vercel/index.ts | 4 +-- apps/web/src/lib/cached-fetch.ts | 2 +- apps/web/src/routers/byok-router.ts | 13 ++++---- 11 files changed, 50 insertions(+), 32 deletions(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.test.ts b/apps/web/src/app/api/openrouter/[...path]/route.test.ts index 71bf2941ff..9500c28837 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.test.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.test.ts @@ -5,7 +5,7 @@ import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage import { classifyAbuse } from '@/lib/ai-gateway/abuse-service'; import { getProvider } from '@/lib/ai-gateway/providers/get-provider'; import { upstreamRequest } from '@/lib/ai-gateway/providers/upstream-request'; -import { getOpenRouterModels } from '@/lib/ai-gateway/providers/gateway-models-cache'; +import { getOpenRouterModelsFromRedis } from '@/lib/ai-gateway/providers/gateway-models-cache'; import { emitApiMetricsForResponse } from '@/lib/ai-gateway/o11y/api-metrics.server'; import { accountForMicrodollarUsage } from '@/lib/ai-gateway/llm-proxy-helpers'; import { redisClient } from '@/lib/redis'; @@ -86,7 +86,7 @@ const mockedGetBalanceAndOrgSettings = jest.mocked(getBalanceAndOrgSettings); const mockedClassifyAbuse = jest.mocked(classifyAbuse); const mockedGetProvider = jest.mocked(getProvider); const mockedUpstreamRequest = jest.mocked(upstreamRequest); -const mockedGetOpenRouterModels = jest.mocked(getOpenRouterModels); +const mockedGetOpenRouterModels = jest.mocked(getOpenRouterModelsFromRedis); const mockedEmitApiMetricsForResponse = jest.mocked(emitApiMetricsForResponse); const mockedAccountForMicrodollarUsage = jest.mocked(accountForMicrodollarUsage); const mockedRedisGet = jest.mocked(redisClient.get); diff --git a/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts b/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts index 033272dcd5..3ad34cdebf 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it, jest } from '@jest/globals'; jest.mock('@/lib/ai-gateway/providers/gateway-models-cache', () => ({ - getOpenRouterModels: jest.fn(async () => new Set()), + getOpenRouterModelsFromRedis: jest.fn(async () => new Set()), })); jest.mock('@/lib/kiloclaw/setup-promo', () => ({ diff --git a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts index 3fd63c844f..3ff53046a0 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts @@ -35,7 +35,7 @@ import { findKiloExclusiveModel, isKiloExclusiveFreeModel, } from '@/lib/ai-gateway/models'; -import { getOpenRouterModels } from '@/lib/ai-gateway/providers/gateway-models-cache'; +import { getOpenRouterModelsFromRedis } from '@/lib/ai-gateway/providers/gateway-models-cache'; import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; import type { ProviderId } from '@/lib/ai-gateway/providers/types'; import { @@ -79,7 +79,7 @@ function resolveMode(modeHeader: string | null, featureHeader: FeatureValue | nu export async function getAutoFreeCandidates( apiKind: GatewayRequest['kind'] | null ): Promise> { - const openRouterModels = await getOpenRouterModels(); + const openRouterModels = await getOpenRouterModelsFromRedis(); const candidates = new Set(); for (const model of autoFreeModels) { if (isKiloExclusiveFreeModel(model)) { diff --git a/apps/web/src/lib/ai-gateway/byok/index.ts b/apps/web/src/lib/ai-gateway/byok/index.ts index e0b7bd201b..f049fae0b6 100644 --- a/apps/web/src/lib/ai-gateway/byok/index.ts +++ b/apps/web/src/lib/ai-gateway/byok/index.ts @@ -12,12 +12,12 @@ import { import { isCodestralModel } from '@/lib/ai-gateway/providers/mistral'; import { mapModelIdToVercel } from '@/lib/ai-gateway/providers/vercel/mapModelIdToVercel'; import type { BYOKResult } from '@/lib/ai-gateway/providers/types'; -import { getVercelModelsMetadata } from '@/lib/ai-gateway/providers/gateway-models-cache'; +import { getVercelModelsMetadataFromDatabase } from '@/lib/ai-gateway/providers/gateway-models-cache'; import type { OpenRouterModel } from '@/lib/organizations/organization-types'; import { isKiloExclusiveModel } from '@/lib/ai-gateway/models'; export async function getModelUserByokProviders(modelId: string): Promise { - const vercelModelMetadata = await getVercelModelsMetadata(); + const vercelModelMetadata = await getVercelModelsMetadataFromDatabase(); if (Object.keys(vercelModelMetadata).length === 0) { console.error('[getModelUserByokProviders] no Vercel model metadata for model %s', modelId); return []; diff --git a/apps/web/src/lib/ai-gateway/providers/gateway-models-cache.ts b/apps/web/src/lib/ai-gateway/providers/gateway-models-cache.ts index 6ec6c444c9..21a8a808ad 100644 --- a/apps/web/src/lib/ai-gateway/providers/gateway-models-cache.ts +++ b/apps/web/src/lib/ai-gateway/providers/gateway-models-cache.ts @@ -1,36 +1,42 @@ -import { StoredModelSchema, type StoredModel } from '@kilocode/db'; +import { modelsByProvider, StoredModelSchema, type StoredModel } from '@kilocode/db'; +import { desc } from 'drizzle-orm'; import * as z from 'zod'; import { redisClient } from '@/lib/redis'; import { createCachedFetch } from '@/lib/cached-fetch'; import { GATEWAY_METADATA_REDIS_KEYS, vercelInferenceProvidersRedisKey } from '@/lib/redis-keys'; import type { RedisKey } from '@/lib/redis-keys'; +import { readDb } from '@/lib/drizzle'; export type StoredModelMap = Record; const StoredModelMapSchema = z.record(z.string(), StoredModelSchema); -function createStoredModelsFetcher(redisKey: RedisKey, name: string) { +function createStoredModelsFromDatabaseFetcher(provider: 'openrouter' | 'vercel', name: string) { return createCachedFetch( async () => { - const raw = JSON.parse((await redisClient.get(redisKey)) ?? 'null'); - if (!raw || typeof raw !== 'object' || Object.keys(raw).length === 0) { - console.debug(`[getGatewayModels] no ${name} models found in Redis`); + const [row] = await readDb + .select({ models: modelsByProvider[provider] }) + .from(modelsByProvider) + .orderBy(desc(modelsByProvider.id)) + .limit(1); + if (!row?.models || Object.keys(row.models).length === 0) { + console.debug(`[getGatewayModels] no ${name} models found in the database`); return {}; } - return StoredModelMapSchema.parse(raw); + return StoredModelMapSchema.parse(row.models); }, 600_000, {} ); } -export const getVercelModelsMetadata = createStoredModelsFetcher( - GATEWAY_METADATA_REDIS_KEYS.vercelModels, +export const getVercelModelsMetadataFromDatabase = createStoredModelsFromDatabaseFetcher( + 'vercel', 'Vercel' ); -export const getOpenRouterModelsMetadata = createStoredModelsFetcher( - GATEWAY_METADATA_REDIS_KEYS.openrouterModels, +export const getOpenRouterModelsMetadataFromDatabase = createStoredModelsFromDatabaseFetcher( + 'openrouter', 'OpenRouter' ); @@ -95,12 +101,12 @@ function createModelIdsFetcher(redisKey: RedisKey, name: string) { ); } -export const getVercelModels = createModelIdsFetcher( +export const getVercelModelsFromRedis = createModelIdsFetcher( GATEWAY_METADATA_REDIS_KEYS.vercelModelIds, 'Vercel' ); -export const getOpenRouterModels = createModelIdsFetcher( +export const getOpenRouterModelsFromRedis = createModelIdsFetcher( GATEWAY_METADATA_REDIS_KEYS.openrouterModelIds, 'OpenRouter' ); diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/index.test.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/index.test.ts index a69089f948..76680d767a 100644 --- a/apps/web/src/lib/ai-gateway/providers/openrouter/index.test.ts +++ b/apps/web/src/lib/ai-gateway/providers/openrouter/index.test.ts @@ -21,7 +21,7 @@ import { isFableModel } from '@/lib/ai-gateway/providers/anthropic.constants'; import { KILO_AUTO_EFFICIENT_MODEL } from '@/lib/ai-gateway/auto-model'; jest.mock('@/lib/ai-gateway/providers/gateway-models-cache', () => ({ - getOpenRouterModelsMetadata: jest.fn(() => Promise.resolve({})), + getOpenRouterModelsMetadataFromDatabase: jest.fn(() => Promise.resolve({})), })); const originalFetch = global.fetch; diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts index c558e77e6e..219e30ed79 100644 --- a/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts +++ b/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts @@ -21,7 +21,7 @@ import { isForbiddenFreeModel } from '@/lib/ai-gateway/forbidden-free-models'; import { getGatewayOpenCodeSettings } from '@/lib/ai-gateway/providers/model-settings'; import { AUTO_MODELS, type AutoModel } from '@/lib/ai-gateway/auto-model'; import { ATTRIBUTION_HEADERS } from '@/lib/ai-gateway/providers/openrouter/attribution-headers'; -import { getOpenRouterModelsMetadata } from '@/lib/ai-gateway/providers/gateway-models-cache'; +import { getOpenRouterModelsMetadataFromDatabase } from '@/lib/ai-gateway/providers/gateway-models-cache'; import { getPreferredProviderOrder } from '@/lib/ai-gateway/providers/apply-provider-specific-logic'; import { normalizeInferenceProviderId } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id'; import { getTerminalBenchSummaries, terminalBenchFor } from '@/lib/model-stats/terminal-bench'; @@ -127,7 +127,7 @@ const unavailableModels = [ async function enhancedModelList(models: OpenRouterModel[]) { const autoModels = buildAutoModels(); - const endpointsMetadata = await getOpenRouterModelsMetadata(); + const endpointsMetadata = await getOpenRouterModelsMetadataFromDatabase(); const summaries = await getTerminalBenchSummaries(); const enhancedModels = await Promise.all( models diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/sync-providers.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/sync-providers.ts index ae7f0fa310..839dcd693a 100644 --- a/apps/web/src/lib/ai-gateway/providers/openrouter/sync-providers.ts +++ b/apps/web/src/lib/ai-gateway/providers/openrouter/sync-providers.ts @@ -394,6 +394,8 @@ async function syncProviders( return result; } +const MODEL_METADATA_REDIS_TTL_SECONDS = 7 * 24 * 60 * 60; + async function mirrorToRedis(values: { providers: NormalizedOpenRouterResponse; openrouter: Record; @@ -411,7 +413,16 @@ async function mirrorToRedis(values: { entries.push([GATEWAY_METADATA_REDIS_KEYS.openrouterProviders, values.openrouterProviders]); } await Promise.all([ - ...entries.map(([key, value]) => redisClient.set(key, JSON.stringify(value))), + ...entries.map(([key, value]) => { + const serializedValue = JSON.stringify(value); + if ( + key === GATEWAY_METADATA_REDIS_KEYS.openrouterModels || + key === GATEWAY_METADATA_REDIS_KEYS.vercelModels + ) { + return redisClient.set(key, serializedValue, { ex: MODEL_METADATA_REDIS_TTL_SECONDS }); + } + return redisClient.set(key, serializedValue); + }), mirrorVercelInferenceProvidersToRedis(values.vercel), ]); } diff --git a/apps/web/src/lib/ai-gateway/providers/vercel/index.ts b/apps/web/src/lib/ai-gateway/providers/vercel/index.ts index 4a8064624f..cd30bbc078 100644 --- a/apps/web/src/lib/ai-gateway/providers/vercel/index.ts +++ b/apps/web/src/lib/ai-gateway/providers/vercel/index.ts @@ -23,7 +23,7 @@ import { VERCEL_ROUTING_REDIS_KEY } from '@/lib/redis-keys'; import { getRandomNumber } from '@/lib/ai-gateway/getRandomNumber'; import { getCachedVercelInferenceProviderIdsForModel, - getVercelModels, + getVercelModelsFromRedis, } from '@/lib/ai-gateway/providers/gateway-models-cache'; import type { AnthropicProviderOptions } from '@ai-sdk/anthropic'; import { isDeepseekModel } from '@/lib/ai-gateway/providers/deepseek'; @@ -104,7 +104,7 @@ export async function shouldRouteToVercel( return false; } - const vercelModels = await getVercelModels(); + const vercelModels = await getVercelModelsFromRedis(); const vercelModelId = mapModelIdToVercel(requestedModel, isReasoningExplicitlyDisabled(request)); if (!vercelModels.has(vercelModelId)) { console.debug(`[shouldRouteToVercel] model not found in Vercel model list`); diff --git a/apps/web/src/lib/cached-fetch.ts b/apps/web/src/lib/cached-fetch.ts index 79ae9d4fe3..01a6309304 100644 --- a/apps/web/src/lib/cached-fetch.ts +++ b/apps/web/src/lib/cached-fetch.ts @@ -3,7 +3,7 @@ * * Returns the cached value immediately if it's younger than `ttlMs`, * otherwise calls `fetcher` to refresh. If the fetcher throws (e.g. - * Redis timeout), returns the last-known-good cached value, or + * a network timeout), returns the last-known-good cached value, or * `defaultValue` if nothing has been cached yet. */ export function createCachedFetch(fetcher: () => Promise, ttlMs: number, defaultValue: T) { diff --git a/apps/web/src/routers/byok-router.ts b/apps/web/src/routers/byok-router.ts index f3d55345f5..dac22a9f83 100644 --- a/apps/web/src/routers/byok-router.ts +++ b/apps/web/src/routers/byok-router.ts @@ -29,8 +29,8 @@ import { VercelUserByokInferenceProviderIdSchema, } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id'; import { - getVercelModelsMetadata, - getOpenRouterModelsMetadata, + getVercelModelsMetadataFromDatabase, + getOpenRouterModelsMetadataFromDatabase, } from '@/lib/ai-gateway/providers/gateway-models-cache'; import { createGateway, generateText } from 'ai'; import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; @@ -88,19 +88,20 @@ async function testCodestralApiKey(apiKey: string): Promise<{ success: boolean; async function fetchSupportedModels(): Promise> { const [vercelModelMetadata, openRouterModelMetadata] = await Promise.all([ - getVercelModelsMetadata(), - getOpenRouterModelsMetadata(), + getVercelModelsMetadataFromDatabase(), + getOpenRouterModelsMetadataFromDatabase(), ]); if (Object.keys(vercelModelMetadata).length === 0) { throw new Error( - 'No Vercel model metadata in Redis, use the admin panel at ' + MODELS_BY_PROVIDER_ADMIN_URL + 'No Vercel model metadata in the database, use the admin panel at ' + + MODELS_BY_PROVIDER_ADMIN_URL ); } if (Object.keys(openRouterModelMetadata).length === 0) { throw new Error( - 'No OpenRouter model metadata in Redis, use the admin panel at ' + + 'No OpenRouter model metadata in the database, use the admin panel at ' + MODELS_BY_PROVIDER_ADMIN_URL ); }