use class with helper methods for model catelog

server/lib/aiModelCatalog.ts

@@ -3,13 +3,18 @@ import axios from "axios";
 import config from "@server/lib/config";
 import logger from "@server/logger";
 
-export type CatalogProvider =
+export const CATALOG_PROVIDERS = [
-    | "openai"
+    "openai",
-    | "anthropic"
+    "anthropic",
-    | "gemini"
+    "gemini",
-    | "vertex"
+    "vertex",
-    | "azure"
+    "azure",
-    | "bedrock";
+    "bedrock"
+] as const;
+
+export type CatalogProvider = (typeof CATALOG_PROVIDERS)[number];
+
+const CATALOG_PROVIDER_SET = new Set<string>(CATALOG_PROVIDERS);
 
 export type AiModelCatalogEntry = {
     provider: CatalogProvider;
@@ -22,90 +27,222 @@ export type AiModelCatalogEntry = {
     };
 };
 
-let catalog: AiModelCatalogEntry[] = [];
+type RawCatalogEntry = {
-let refreshTimer: NodeJS.Timeout | null = null;
+    id?: string;
+    name?: string;
+    model?: string;
+    provider: string;
+    input_cost_per_token?: number | null;
+    output_cost_per_token?: number | null;
+    cache_read_input_token_cost?: number | null;
+    output_cost_per_reasoning_token?: number | null;
+    pricing?: {
+        input?: number | null;
+        output?: number | null;
+        cacheRead?: number | null;
+        reasoningOutput?: number | null;
+    };
+};
 
-async function fetchFromFile(filePath: string): Promise<AiModelCatalogEntry[] | null> {
+function normalizeCatalogProvider(raw: string): CatalogProvider | null {
-    try {
+    if (CATALOG_PROVIDER_SET.has(raw)) {
-        if (!fs.existsSync(filePath)) {
+        return raw as CatalogProvider;
+    }
+    if (raw.startsWith("bedrock")) {
+        return "bedrock";
+    }
+    if (raw.startsWith("vertex")) {
+        return "vertex";
+    }
+    if (raw.startsWith("azure")) {
+        return "azure";
+    }
+    return null;
+}
+
+function normalizeEntry(raw: RawCatalogEntry): AiModelCatalogEntry | null {
+    const provider = normalizeCatalogProvider(raw.provider);
+    if (!provider) {
+        return null;
+    }
+
+    const model = raw.model ?? raw.name ?? raw.id;
+    if (!model) {
+        return null;
+    }
+
+    return {
+        provider,
+        model,
+        pricing: {
+            input: raw.pricing?.input ?? raw.input_cost_per_token ?? null,
+            output: raw.pricing?.output ?? raw.output_cost_per_token ?? null,
+            cacheRead:
+                raw.pricing?.cacheRead ??
+                raw.cache_read_input_token_cost ??
+                null,
+            reasoningOutput:
+                raw.pricing?.reasoningOutput ??
+                raw.output_cost_per_reasoning_token ??
+                null
+        }
+    };
+}
+
+function providerKey(provider: CatalogProvider, key: string): string {
+    return `${provider}\0${key}`;
+}
+
+export class AiModelCatalog {
+    private entries: AiModelCatalogEntry[] = [];
+    private byProvider = new Map<CatalogProvider, AiModelCatalogEntry[]>();
+    private byProviderAndKey = new Map<string, AiModelCatalogEntry>();
+    private byKey = new Map<string, AiModelCatalogEntry[]>();
+    private refreshTimer: NodeJS.Timeout | null = null;
+
+    /**
+     * Loads the catalog into memory and schedules periodic background refreshes.
+     * Call once at server startup.
+     */
+    async init(): Promise<void> {
+        await this.refresh();
+        this.scheduleNextRefresh();
+    }
+
+    /** Exact lookup by catalog provider and model key. */
+    get(
+        provider: CatalogProvider,
+        key: string
+    ): AiModelCatalogEntry | undefined {
+        return this.byProviderAndKey.get(providerKey(provider, key));
+    }
+
+    /** All models for a catalog provider. */
+    list(provider: CatalogProvider): AiModelCatalogEntry[] {
+        return this.byProvider.get(provider) ?? [];
+    }
+
+    /** All catalog entries that share a model key, across providers. */
+    listByKey(key: string): AiModelCatalogEntry[] {
+        return this.byKey.get(key) ?? [];
+    }
+
+    /** Full in-memory catalog. */
+    getAll(): AiModelCatalogEntry[] {
+        return this.entries;
+    }
+
+    private setEntries(entries: AiModelCatalogEntry[]): void {
+        const byProvider = new Map<CatalogProvider, AiModelCatalogEntry[]>();
+        const byProviderAndKey = new Map<string, AiModelCatalogEntry>();
+        const byKey = new Map<string, AiModelCatalogEntry[]>();
+
+        for (const entry of entries) {
+            const list = byProvider.get(entry.provider) ?? [];
+            list.push(entry);
+            byProvider.set(entry.provider, list);
+
+            const mapKey = providerKey(entry.provider, entry.model);
+            if (!byProviderAndKey.has(mapKey)) {
+                byProviderAndKey.set(mapKey, entry);
+            }
+
+            const keyList = byKey.get(entry.model) ?? [];
+            keyList.push(entry);
+            byKey.set(entry.model, keyList);
+        }
+
+        this.entries = entries;
+        this.byProvider = byProvider;
+        this.byProviderAndKey = byProviderAndKey;
+        this.byKey = byKey;
+    }
+
+    private async fetchFromFile(
+        filePath: string
+    ): Promise<AiModelCatalogEntry[] | null> {
+        try {
+            if (!fs.existsSync(filePath)) {
+                logger.warn(
+                    `AI model catalog file not found at ${filePath}; cost calculation will fall back to unknown pricing`
+                );
+                return null;
+            }
+            const raw = fs.readFileSync(filePath, "utf-8");
+            const parsed = JSON.parse(raw) as { data: RawCatalogEntry[] };
+            return (parsed.data ?? [])
+                .map(normalizeEntry)
+                .filter((e): e is AiModelCatalogEntry => e != null);
+        } catch (error) {
+            logger.warn("Failed to read AI model catalog file", { error });
+            return null;
+        }
+    }
+
+    private async fetchFromUpstream(
+        upstreamUrl: string
+    ): Promise<AiModelCatalogEntry[] | null> {
+        try {
+            const res = await axios.get<{ data: RawCatalogEntry[] }>(
+                upstreamUrl,
+                { timeout: 15_000 }
+            );
+            return (res.data?.data ?? [])
+                .map(normalizeEntry)
+                .filter((e): e is AiModelCatalogEntry => e != null);
+        } catch (error: any) {
             logger.warn(
-                `AI model catalog file not found at ${filePath}; cost calculation will fall back to unknown pricing`
+                `Failed to fetch AI model catalog from ${upstreamUrl}: ${error.message || error}`
             );
             return null;
         }
-        const raw = fs.readFileSync(filePath, "utf-8");
+    }
-        const parsed = JSON.parse(raw) as { data: AiModelCatalogEntry[] };
+
-        return parsed.data ?? [];
+    private async refresh(): Promise<void> {
-    } catch (error) {
+        const { file, upstream_url } = config.getRawConfig().ai.model_catalog;
-        logger.warn("Failed to read AI model catalog file", { error });
+
-        return null;
+        const fetched = file
+            ? await this.fetchFromFile(file)
+            : await this.fetchFromUpstream(upstream_url);
+
+        if (fetched) {
+            this.setEntries(fetched);
+            logger.debug(
+                `AI model catalog refreshed: ${this.entries.length} models loaded`
+            );
+        } else {
+            logger.debug(
+                "AI model catalog refresh failed; keeping previously loaded catalog in memory"
+            );
+        }
+    }
+
+    private scheduleNextRefresh(): void {
+        const { refresh_interval_min_hours, refresh_interval_max_hours } =
+            config.getRawConfig().ai.model_catalog;
+
+        // Jittered rather than fixed so that many self-hosted instances don't
+        // all hit the upstream catalog endpoint at the same moment.
+        const minMs = refresh_interval_min_hours * 60 * 60 * 1000;
+        const maxMs = refresh_interval_max_hours * 60 * 60 * 1000;
+        const delayMs = minMs + Math.random() * Math.max(0, maxMs - minMs);
+
+        if (this.refreshTimer) {
+            clearTimeout(this.refreshTimer);
+        }
+        this.refreshTimer = setTimeout(async () => {
+            await this.refresh();
+            this.scheduleNextRefresh();
+        }, delayMs);
     }
 }
 
-async function fetchFromUpstream(
+export const aiModelCatalog = new AiModelCatalog();
-    upstreamUrl: string
-): Promise<AiModelCatalogEntry[] | null> {
-    try {
-        const res = await axios.get<{ data: AiModelCatalogEntry[] }>(
-            upstreamUrl,
-            { timeout: 15_000 }
-        );
-        return res.data?.data ?? [];
-    } catch (error: any) {
-        logger.warn(
-            `Failed to fetch AI model catalog from ${upstreamUrl}: ${error.message || error}`
-        );
-        return null;
-    }
-}
-
-async function refreshCatalog(): Promise<void> {
-    const { file, upstream_url } = config.getRawConfig().ai.model_catalog;
-
-    const fetched = file
-        ? await fetchFromFile(file)
-        : await fetchFromUpstream(upstream_url);
-
-    if (fetched) {
-        catalog = fetched;
-        logger.debug(
-            `AI model catalog refreshed: ${catalog.length} models loaded`
-        );
-    } else {
-        logger.debug(
-            "AI model catalog refresh failed; keeping previously loaded catalog in memory"
-        );
-    }
-}
-
-function scheduleNextRefresh(): void {
-    const { refresh_interval_min_hours, refresh_interval_max_hours } =
-        config.getRawConfig().ai.model_catalog;
-
-    // Jittered rather than fixed so that many self-hosted instances don't
-    // all hit the upstream catalog endpoint at the same moment.
-    const minMs = refresh_interval_min_hours * 60 * 60 * 1000;
-    const maxMs = refresh_interval_max_hours * 60 * 60 * 1000;
-    const delayMs = minMs + Math.random() * Math.max(0, maxMs - minMs);
-
-    if (refreshTimer) {
-        clearTimeout(refreshTimer);
-    }
-    refreshTimer = setTimeout(async () => {
-        await refreshCatalog();
-        scheduleNextRefresh();
-    }, delayMs);
-}
 
 /**
  * Loads the AI model pricing catalog into memory and schedules periodic
  * background refreshes. Call once at server startup.
  */
 export async function initAiModelCatalog(): Promise<void> {
-    await refreshCatalog();
+    await aiModelCatalog.init();
-    scheduleNextRefresh();
-}
-
-export function getAiModelCatalog(): AiModelCatalogEntry[] {
-    return catalog;
 }

server/lib/aiModelPricing.ts

@@ -1,7 +1,7 @@
 import type { AiProviderType } from "@server/lib/aiProviderDefaults";
 import type { AiUsage } from "@server/lib/aiUsageExtraction";
 import {
-    getAiModelCatalog,
+    aiModelCatalog,
     type AiModelCatalogEntry,
     type CatalogProvider
 } from "@server/lib/aiModelCatalog";
@@ -36,31 +36,6 @@ const PROVIDER_CATALOG_MAP: Record<
     vercelAiGateway: null
 };
 
-// Indexed view over the in-memory catalog, rebuilt only when
-// getAiModelCatalog() returns a different array instance (i.e. after a
-// background refresh swaps it out), not on every lookup.
-let indexedCatalog: AiModelCatalogEntry[] | null = null;
-let indexedByName: Map<string, AiModelCatalogEntry[]> = new Map();
-
-function getIndexedCatalog(): Map<string, AiModelCatalogEntry[]> {
-    const catalog = getAiModelCatalog();
-    if (catalog === indexedCatalog) {
-        return indexedByName;
-    }
-
-    const byName = new Map<string, AiModelCatalogEntry[]>();
-    for (const entry of catalog) {
-        if (!entry.model) continue;
-        const list = byName.get(entry.model) ?? [];
-        list.push(entry);
-        byName.set(entry.model, list);
-    }
-
-    indexedCatalog = catalog;
-    indexedByName = byName;
-    return byName;
-}
-
 function stripVendorPrefix(modelId: string): string | null {
     const idx = modelId.indexOf("/");
     if (idx === -1 || idx === modelId.length - 1) {
@@ -83,7 +58,6 @@ function toPricing(
 }
 
 function findEntry(
-    byName: Map<string, AiModelCatalogEntry[]>,
     modelId: string,
     provider: CatalogProvider | null
 ): AiModelCatalogEntry | null {
@@ -92,11 +66,15 @@ function findEntry(
     );
 
     for (const key of candidates) {
-        const entries = byName.get(key);
+        if (provider) {
-        if (!entries) continue;
+            const match = aiModelCatalog.get(provider, key);
-        const match = provider
+            if (match) {
-            ? entries.find((e) => e.provider === provider)
+                return match;
-            : entries[0];
+            }
+            continue;
+        }
+
+        const match = aiModelCatalog.listByKey(key)[0];
         if (match) {
             return match;
         }
@@ -118,18 +96,17 @@ export function getModelPricing(
         return null;
     }
 
-    const byName = getIndexedCatalog();
     const catalogProvider =
         providerType === "custom" ? null : PROVIDER_CATALOG_MAP[providerType];
 
     if (catalogProvider) {
-        const scoped = findEntry(byName, modelId, catalogProvider);
+        const scoped = findEntry(modelId, catalogProvider);
         if (scoped) {
             return toPricing(scoped, false);
         }
     }
 
-    const fallback = findEntry(byName, modelId, null);
+    const fallback = findEntry(modelId, null);
     if (fallback) {
         return toPricing(fallback, true);
     }

server/routers/site/socketIntegration.ts

@@ -201,7 +201,7 @@ async function checkSocket(
 ): Promise<{ siteId: number; newtId: string }> {
     const { newt } = await getSiteAndNewt(siteId);
 
-    logger.info(
+    logger.debug(
         `Checking Docker socket for site ${siteId} with Newt ${newt.newtId}`
     );