double cache api keys

gateway/src/auth.ts

@@ -1,11 +1,14 @@
 import type { GatewayOptions } from '.'
 import type { ApiKeyInfo } from './types'
-import { ResponseError } from './utils'
+import { ResponseError, runAfter } from './utils'
 
-const CACHE_VERSION = 1
-const CACHE_TTL = 86400
+const CACHE_TTL = 86400 * 30
 
-export async function apiKeyAuth(request: Request, options: GatewayOptions): Promise<ApiKeyInfo> {
+export async function apiKeyAuth(
+  request: Request,
+  ctx: ExecutionContext,
+  options: GatewayOptions,
+): Promise<ApiKeyInfo> {
   const authHeader = request.headers.get('authorization')
 
   let key: string
@@ -23,26 +26,54 @@ export async function apiKeyAuth(request: Request, options: GatewayOptions): Pro
     throw new ResponseError(401, 'Unauthorized - Key too long')
   }
 
-  const cacheKey = apiKeyCacheKey(key, options)
-  const cacheResult = await options.kv.getWithMetadata<ApiKeyInfo, number>(cacheKey, { type: 'json' })
+  const cacheKey = apiKeyCacheKey(key, options.kvVersion)
+  const cacheResult = await options.kv.getWithMetadata<ApiKeyInfo, string>(cacheKey, { type: 'json' })
 
-  let apiKey: ApiKeyInfo | null
-  if (cacheResult && cacheResult.metadata === CACHE_VERSION && cacheResult.value) {
-    apiKey = cacheResult.value
-  } else {
-    apiKey = await options.keysDb.getApiKey(key)
-    if (!apiKey) {
-      throw new ResponseError(401, 'Unauthorized - Key not found')
+  if (cacheResult?.value) {
+    const apiKey = cacheResult.value
+    const projectState = await options.kv.get(projectStateCacheKey(apiKey.project, options.kvVersion))
+    // we only return a cache match if the project state is the same, so updating the project state invalidates the cache
+    // projectState is null if we have never invalidated the cache which will only be true for the first request after a deployment
+    if (projectState === null || projectState === cacheResult.metadata) {
+      return apiKey
     }
-    await options.kv.put(cacheKey, JSON.stringify(apiKey), { metadata: CACHE_VERSION, expirationTtl: CACHE_TTL })
   }
-  // check all key validity in gateway.ts
-  return apiKey
+
+  const apiKey = await options.keysDb.getApiKey(key)
+  if (apiKey) {
+    runAfter(ctx, 'setApiKeyCache', setApiKeyCache(apiKey, options))
+    return apiKey
+  }
+  throw new ResponseError(401, 'Unauthorized - Key not found')
+}
+
+export async function setApiKeyCache(
+  apiKey: ApiKeyInfo,
+  options: Pick<GatewayOptions, 'kv' | 'kvVersion'>,
+  expirationTtl?: number,
+) {
+  const projectState = await options.kv.get(projectStateCacheKey(apiKey.project, options.kvVersion))
+
+  await options.kv.put(apiKeyCacheKey(apiKey.key, options.kvVersion), JSON.stringify(apiKey), {
+    metadata: projectState,
+    // Note: 0 is a valid expirationTtl (for immediate cache expiry if, e.g., the user hits a limit at the end of an interval).
+    // Do not use logical OR (||) for a fallback, as it would treat 0 as false and incorrectly default to CACHE_TTL,
+    // potentially locking out the user much longer than intended.
+    expirationTtl: expirationTtl ?? CACHE_TTL,
+  })
+}
+
+export async function deleteApiKeyCache(
+  apiKey: Pick<ApiKeyInfo, 'key'>,
+  options: Pick<GatewayOptions, 'kv' | 'kvVersion'>,
+) {
+  await options.kv.delete(apiKeyCacheKey(apiKey.key, options.kvVersion))
 }
 
-export async function disableApiKeyAuth(apiKey: ApiKeyInfo, options: GatewayOptions, expirationTtl?: number) {
-  const cacheKey = apiKeyCacheKey(apiKey.key, options)
-  await options.kv.put(cacheKey, JSON.stringify(apiKey), { metadata: CACHE_VERSION, expirationTtl })
+export async function changeProjectState(project: number, options: Pick<GatewayOptions, 'kv' | 'kvVersion'>) {
+  const cacheKey = projectStateCacheKey(project, options.kvVersion)
+  await options.kv.put(cacheKey, crypto.randomUUID(), { expirationTtl: CACHE_TTL })
 }
 
-const apiKeyCacheKey = (key: string, options: GatewayOptions) => `apiKeyAuth:${options.kvVersion}:${key}`
+const apiKeyCacheKey = (key: string, kvVersion: string) => `apiKeyAuth:${kvVersion}:${key}`
+const projectStateCacheKey = (project: number, kvVersion: string) => `projectState:${kvVersion}:${project}`

gateway/src/gateway.ts

@@ -1,5 +1,5 @@
 import type { GatewayOptions } from '.'
-import { apiKeyAuth, disableApiKeyAuth } from './auth'
+import { apiKeyAuth, setApiKeyCache } from './auth'
 import { type ExceededScope, endOfMonth, endOfWeek, type SpendScope, scopeIntervals } from './db'
 import { OtelTrace } from './otel'
 import { genAiOtelAttributes } from './otel/attributes'
@@ -23,7 +23,7 @@ export async function gateway(
     return textResponse(400, `Invalid provider '${provider}', should be one of ${providerIdArray.join(', ')}`)
   }
 
-  const apiKeyInfo = await apiKeyAuth(request, options)
+  const apiKeyInfo = await apiKeyAuth(request, ctx, options)
 
   if (apiKeyInfo.status !== 'active') {
     return textResponse(403, `Unauthorized - Key ${apiKeyInfo.status}`)
@@ -72,7 +72,7 @@ export async function gateway(
   } else if ('error' in result) {
     const { error, disableKey } = result
     if (disableKey) {
-      runAfter(ctx, 'disableApiKey', blockApiKey(apiKeyInfo, options, 'Invalid request'))
+      runAfter(ctx, 'blockApiKey', blockApiKey(apiKeyInfo, options, 'Invalid request'))
       response = textResponse(400, `${error}, API key disabled`)
     } else {
       response = textResponse(400, error)
@@ -97,7 +97,7 @@ export async function disableApiKey(
   expirationTtl?: number,
 ): Promise<void> {
   apiKey.status = newStatus
-  await disableApiKeyAuth(apiKey, options, expirationTtl)
+  await setApiKeyCache(apiKey, options, expirationTtl)
   await options.keysDb.disableKey(apiKey.id, reason, newStatus, expirationTtl)
 }
 

gateway/src/index.ts

@@ -21,6 +21,7 @@ import type { DefaultProviderProxy, Middleware, Next } from './providers/default
 import type { SubFetch } from './types'
 import { ctHeader, ResponseError, response405, textResponse } from './utils'
 
+export { changeProjectState as setProjectState, deleteApiKeyCache, setApiKeyCache } from './auth'
 export type { DefaultProviderProxy, Middleware, Next }
 export * from './db'
 export * from './types'

gateway/test/auth.spec.ts

@@ -0,0 +1,93 @@
+/** biome-ignore-all lint/suspicious/useAwait: don't care in tests */
+import { createExecutionContext, env, waitOnExecutionContext } from 'cloudflare:test'
+import type { KeysDb } from '@pydantic/ai-gateway'
+import { describe, expect } from 'vitest'
+import { apiKeyAuth, changeProjectState } from '../src/auth'
+import type { ApiKeyInfo } from '../src/types'
+import { test } from './setup'
+import { buildGatewayEnv, IDS } from './worker'
+
+class CountingKeysDb implements KeysDb {
+  callCount = 0
+  private wrapped: KeysDb
+
+  constructor(wrapped: KeysDb) {
+    this.wrapped = wrapped
+  }
+
+  async getApiKey(key: string): Promise<ApiKeyInfo | null> {
+    this.callCount++
+    return this.wrapped.getApiKey(key)
+  }
+
+  async disableKey(id: number, reason: string, newStatus: string, expirationTtl?: number): Promise<void> {
+    return this.wrapped.disableKey(id, reason, newStatus, expirationTtl)
+  }
+}
+
+describe('apiKeyAuth cache invalidation', () => {
+  test('caches api key and returns cached value', async () => {
+    const ctx = createExecutionContext()
+    const baseOptions = buildGatewayEnv(env, [], fetch)
+    const countingDb = new CountingKeysDb(baseOptions.keysDb)
+    const options = { ...baseOptions, keysDb: countingDb }
+
+    const request = new Request('https://example.com', { headers: { Authorization: 'healthy' } })
+
+    // First call should fetch from DB
+    const apiKey1 = await apiKeyAuth(request, ctx, options)
+    expect(apiKey1.key).toBe('healthy')
+    // Wait for cache to be set (it's set asynchronously via runAfter)
+    await waitOnExecutionContext(ctx)
+    expect(countingDb.callCount).toBe(1)
+
+    // Verify cache was set
+    const cached = await env.KV.get('apiKeyAuth:test:healthy')
+    expect(cached).toBeTypeOf('string')
+
+    // Second call should use cache, not hit DB
+    const ctx2 = createExecutionContext()
+    const apiKey2 = await apiKeyAuth(request, ctx2, options)
+    expect(apiKey2.key).toBe('healthy')
+
+    expect(countingDb.callCount).toBe(1)
+  })
+
+  test('invalidates cache when project state changes', async () => {
+    const ctx = createExecutionContext()
+    const baseOptions = buildGatewayEnv(env, [], fetch)
+    const countingDb = new CountingKeysDb(baseOptions.keysDb)
+    const options = { ...baseOptions, keysDb: countingDb }
+
+    const request = new Request('https://example.com', { headers: { Authorization: 'healthy' } })
+
+    // First call - fetch from DB and cache
+    await apiKeyAuth(request, ctx, options)
+    await waitOnExecutionContext(ctx)
+    expect(countingDb.callCount).toBe(1)
+
+    const cached1 = await env.KV.getWithMetadata('apiKeyAuth:test:healthy')
+    expect(cached1.value).not.toBeNull()
+    expect(cached1.metadata).toBeNull()
+
+    // Second call - should use cache, not hit DB
+    const ctx2 = createExecutionContext()
+    await apiKeyAuth(request, ctx2, options)
+    await waitOnExecutionContext(ctx2)
+    expect(countingDb.callCount).toBe(1)
+
+    // Change project state - this invalidates the cache
+    await changeProjectState(IDS.projectDefault, options)
+
+    const projectState = await env.KV.get(`projectState:test:${IDS.projectDefault}`)
+    expect(projectState).not.toBeNull()
+
+    // Third call - cache is invalidated, should hit DB again
+    const ctx3 = createExecutionContext()
+    const apiKey3 = await apiKeyAuth(request, ctx3, options)
+    expect(apiKey3.key).toBe('healthy')
+    await waitOnExecutionContext(ctx3)
+
+    expect(countingDb.callCount).toBe(2)
+  })
+})