give the llm model registry a short window to become ready when the server first tries to enrich spans so we don't miss enriching the first Xms of spans every time a server boots

Author: ericallam

apps/webapp/app/env.server.ts

@@ -1287,6 +1287,7 @@ const EnvironmentSchema = z
     LLM_COST_TRACKING_ENABLED: BoolEnv.default(true),
     LLM_PRICING_RELOAD_INTERVAL_MS: z.coerce.number().int().default(5 * 60 * 1000), // 5 minutes
     LLM_PRICING_SEED_ON_STARTUP: BoolEnv.default(false),
+    LLM_PRICING_READY_TIMEOUT_MS: z.coerce.number().int().default(500),
 
     // Bootstrap
     TRIGGER_BOOTSTRAP_ENABLED: z.string().default("0"),

apps/webapp/app/v3/llmPricingRegistry.server.ts

@@ -1,12 +1,13 @@
 import { ModelPricingRegistry, seedLlmPricing } from "@internal/llm-pricing";
 import { prisma, $replica } from "~/db.server";
 import { env } from "~/env.server";
+import { signalsEmitter } from "~/services/signals.server";
 import { singleton } from "~/utils/singleton";
 import { setLlmPricingRegistry } from "./utils/enrichCreatableEvents.server";
 
 async function initRegistry(registry: ModelPricingRegistry) {
   if (env.LLM_PRICING_SEED_ON_STARTUP) {
-    const result = await seedLlmPricing(prisma);
+    await seedLlmPricing(prisma);
   }
 
   await registry.loadFromDatabase();
@@ -28,15 +29,34 @@ export const llmPricingRegistry = singleton("llmPricingRegistry", () => {
 
   // Periodic reload
   const reloadInterval = env.LLM_PRICING_RELOAD_INTERVAL_MS;
-  setInterval(() => {
-    registry
-      .reload()
-      .then(() => {
-      })
-      .catch((err) => {
-        console.error("Failed to reload LLM pricing registry", err);
-      });
+  const interval = setInterval(() => {
+    registry.reload().catch((err) => {
+      console.error("Failed to reload LLM pricing registry", err);
+    });
   }, reloadInterval);
 
+  signalsEmitter.on("SIGTERM", () => {
+    clearInterval(interval);
+  });
+  signalsEmitter.on("SIGINT", () => {
+    clearInterval(interval);
+  });
+
   return registry;
 });
+
+/**
+ * Wait for the LLM pricing registry to finish its initial load, with a timeout.
+ * After the first call resolves (or times out), subsequent calls are no-ops.
+ */
+export async function waitForLlmPricingReady(): Promise<void> {
+  if (!llmPricingRegistry || llmPricingRegistry.isLoaded) return;
+
+  const timeoutMs = env.LLM_PRICING_READY_TIMEOUT_MS;
+  if (timeoutMs <= 0) return;
+
+  await Promise.race([
+    llmPricingRegistry.isReady,
+    new Promise<void>((resolve) => setTimeout(resolve, timeoutMs)),
+  ]);
+}

apps/webapp/app/v3/otlpExporter.server.ts

@@ -37,7 +37,7 @@ import type {
 } from "./eventRepository/eventRepository.types";
 import { startSpan } from "./tracing.server";
 import { enrichCreatableEvents } from "./utils/enrichCreatableEvents.server";
-import "./llmPricingRegistry.server"; // Initialize LLM pricing registry on startup
+import { waitForLlmPricingReady } from "./llmPricingRegistry.server";
 import { env } from "~/env.server";
 import { detectBadJsonStrings } from "~/utils/detectBadJsonStrings";
 import { singleton } from "~/utils/singleton";
@@ -129,6 +129,7 @@ class OTLPExporter {
     for (const [store, events] of Object.entries(eventsGroupedByStore)) {
       const eventRepository = this.#getEventRepositoryForStore(store);
 
+      await waitForLlmPricingReady();
       const enrichedEvents = enrichCreatableEvents(events);
 
       this.#logEventsVerbose(enrichedEvents, `exportEvents ${store}`);

apps/webapp/app/v3/utils/enrichCreatableEvents.server.ts

@@ -20,6 +20,8 @@ type CostRegistry = {
 
 let _registry: CostRegistry | undefined;
 
+const ENRICHABLE_KINDS = new Set(["INTERNAL", "SERVER", "CLIENT", "CONSUMER", "PRODUCER"]);
+
 export function setLlmPricingRegistry(registry: CostRegistry): void {
   _registry = registry;
 }
@@ -46,8 +48,7 @@ function enrichLlmMetrics(event: CreateEventInput): void {
   if (!props) return;
 
   // Only enrich span-like events (INTERNAL, SERVER, CLIENT, CONSUMER, PRODUCER — not LOG, UNSPECIFIED)
-  const enrichableKinds = new Set(["INTERNAL", "SERVER", "CLIENT", "CONSUMER", "PRODUCER"]);
-  if (!enrichableKinds.has(event.kind as string)) return;
+  if (!ENRICHABLE_KINDS.has(event.kind as string)) return;
 
   // Skip partial spans (they don't have final token counts)
   if (event.isPartial) return;

internal-packages/llm-pricing/src/registry.ts

@@ -24,9 +24,16 @@ export class ModelPricingRegistry {
   private _patterns: CompiledPattern[] = [];
   private _exactMatchCache: Map<string, LlmModelWithPricing | null> = new Map();
   private _loaded = false;
+  private _readyResolve!: () => void;
+
+  /** Resolves once the initial `loadFromDatabase()` completes successfully. */
+  readonly isReady: Promise<void>;
 
   constructor(prisma: PrismaClient | PrismaReplicaClient) {
     this._prisma = prisma;
+    this.isReady = new Promise<void>((resolve) => {
+      this._readyResolve = resolve;
+    });
   }
 
   get isLoaded(): boolean {
@@ -81,7 +88,11 @@ export class ModelPricingRegistry {
 
     this._patterns = compiled;
     this._exactMatchCache.clear();
-    this._loaded = true;
+
+    if (!this._loaded) {
+      this._loaded = true;
+      this._readyResolve();
+    }
   }
 
   async reload(): Promise<void> {