provider key cache (#5453)

* provider key cache

* better version

Author: chitalian

worker/src/lib/ai-gateway/AttemptBuilder.ts

@@ -24,16 +24,19 @@ export class AttemptBuilder {
   private readonly pluginHandler: PluginHandler;
   private readonly tracer: DataDogTracer;
   private readonly traceContext: TraceContext | null;
+  private readonly ctx: ExecutionContext;
 
   constructor(
     private readonly providerKeysManager: ProviderKeysManager,
     private readonly env: Env,
     tracer: DataDogTracer,
-    traceContext: TraceContext | null
+    traceContext: TraceContext | null,
+    ctx: ExecutionContext
   ) {
     this.pluginHandler = new PluginHandler();
     this.tracer = tracer;
     this.traceContext = traceContext;
+    this.ctx = ctx;
   }
 
   async buildAttempts(
@@ -210,7 +213,8 @@ export class AttemptBuilder {
       providerData.provider,
       modelSpec.modelName,
       orgId,
-      modelSpec.customUid
+      modelSpec.customUid,
+      this.ctx
     );
 
     this.tracer.finishSpan(keySpan);
@@ -266,7 +270,8 @@ export class AttemptBuilder {
       modelSpec.provider as ModelProviderName,
       modelSpec.modelName,
       orgId,
-      modelSpec.customUid
+      modelSpec.customUid,
+      this.ctx
     );
 
     if (!userKey || !this.isByokEnabled(userKey)) {
@@ -337,7 +342,9 @@ export class AttemptBuilder {
     const heliconeKey = await this.providerKeysManager.getProviderKeyWithFetch(
       providerData.provider,
       providerData.config.providerModelId,
-      this.env.HELICONE_ORG_ID
+      this.env.HELICONE_ORG_ID,
+      undefined,
+      this.ctx
     );
 
     this.tracer.finishSpan(ptbKeySpan);

worker/src/lib/ai-gateway/SimpleAIGateway.ts

@@ -95,7 +95,8 @@ export class SimpleAIGateway {
       providerKeysManager,
       env,
       tracer,
-      traceContext
+      traceContext,
+      ctx
     );
     this.attemptExecutor = new AttemptExecutor(env, ctx, cacheProvider, tracer);
   }

worker/src/lib/managers/ProviderKeysManager.ts

@@ -88,53 +88,109 @@ export class ProviderKeysManager {
     return JSON.parse(keys) as ProviderKey[];
   }
 
+  /**
+   * Get provider key with read-through cache pattern.
+   * Returns cached data immediately, always refreshes in background.
+   */
   async getProviderKeyWithFetch(
     provider: ModelProviderName,
     providerModelId: string,
     orgId: string,
-    keyCuid?: string
+    keyCuid?: string,
+    ctx?: ExecutionContext
   ): Promise<ProviderKey | null> {
+    const cacheKey = `provider_keys_${orgId}`;
+    const ttl = 43200; // 12 hours
+
+    // Get cached keys
     const keys = await this.getProviderKeys(orgId);
+
+    // Try to find key from cache
     const validKey = this.chooseProviderKey(
       keys ?? [],
       provider,
       providerModelId,
       keyCuid
     );
 
-    if (!validKey) {
-      const keys = await this.store.getProviderKeysWithFetch(
+    if (validKey) {
+      // Cache hit - trigger background refresh and return immediately
+      if (ctx) {
+        ctx.waitUntil(
+          this.fetchAndCacheProviderKey(
+            provider,
+            providerModelId,
+            orgId,
+            keyCuid,
+            cacheKey,
+            ttl
+          )
+        );
+      }
+      return validKey;
+    }
+
+    // Cache miss - must wait for fetch
+    return this.fetchAndCacheProviderKey(
+      provider,
+      providerModelId,
+      orgId,
+      keyCuid,
+      cacheKey,
+      ttl
+    );
+  }
+
+  /**
+   * Fetch provider key from Supabase and update cache.
+   * Used both for cache miss (awaited) and background refresh (fire-and-forget).
+   */
+  private async fetchAndCacheProviderKey(
+    provider: ModelProviderName,
+    providerModelId: string,
+    orgId: string,
+    keyCuid: string | undefined,
+    cacheKey: string,
+    ttl: number
+  ): Promise<ProviderKey | null> {
+    try {
+      const fetchedKeys = await this.store.getProviderKeysWithFetch(
         provider,
         orgId,
         keyCuid
       );
 
-      if (!keys) return null;
+      if (!fetchedKeys || fetchedKeys.length === 0) return null;
+
+      // Merge with existing cache
+      const existingCached = await this.getProviderKeys(orgId);
+      const existingKeys = existingCached ?? [];
+
+      // Dedupe by cuid (or provider if no cuid)
+      const keyMap = new Map<string, ProviderKey>();
+      for (const key of existingKeys) {
+        const id = key.cuid ?? `${key.provider}`;
+        keyMap.set(id, key);
+      }
+      for (const key of fetchedKeys) {
+        const id = key.cuid ?? `${key.provider}`;
+        keyMap.set(id, key);
+      }
 
-      const existingKeys = await getFromKVCacheOnly(
-        `provider_keys_${orgId}`,
+      const mergedKeys = Array.from(keyMap.values());
+
+      await storeInCache(
+        cacheKey,
+        JSON.stringify(mergedKeys),
         this.env,
-        43200 // 12 hours
+        ttl,
+        false // Don't use memory cache to avoid test contamination
       );
-      if (existingKeys) {
-        const existingKeysData = JSON.parse(existingKeys) as ProviderKey[];
-        existingKeysData.push(...keys);
-        await storeInCache(
-          `provider_keys_${orgId}`,
-          JSON.stringify(existingKeysData),
-          this.env,
-          43200 // 12 hours
-        );
-      } else {
-        await storeInCache(
-          `provider_keys_${orgId}`,
-          JSON.stringify(keys),
-          this.env,
-          43200 // 12 hours
-        );
-      }
-      return this.chooseProviderKey(keys, provider, providerModelId, keyCuid);
+
+      return this.chooseProviderKey(fetchedKeys, provider, providerModelId, keyCuid);
+    } catch (e) {
+      console.error(`Failed to fetch/cache provider key for ${orgId}:`, e);
+      return null;
     }
-    return validKey;
   }
 }