Gate admission on per-deployment Fireworks health

Replace the fleet-wide worst-of collapse with a per-model map. One
probe per tick still covers every deployment (Fireworks returns them
in a single response), but each model's admission now uses its own
deployment's verdict — a degraded minimax no longer blocks glm.

Models absent from FIREWORKS_DEPLOYMENT_MAP (serverless) default to
'healthy'; TODO for when they move to dedicated deployments.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: jahooma

web/src/server/free-session/__tests__/admission.test.ts

@@ -3,9 +3,10 @@ import { describe, expect, test } from 'bun:test'
 import { runAdmissionTick } from '../admission'
 
 import type { AdmissionDeps } from '../admission'
-import type { FireworksHealth } from '../fireworks-health'
+import type { FireworksHealth, FleetHealth } from '../fireworks-health'
 
 const NOW = new Date('2026-04-17T12:00:00Z')
+const TEST_MODEL = 'test-model'
 
 function makeAdmissionDeps(overrides: Partial<AdmissionDeps> = {}): AdmissionDeps & {
   calls: { admit: number }
@@ -16,10 +17,9 @@ function makeAdmissionDeps(overrides: Partial<AdmissionDeps> = {}): AdmissionDep
     sweepExpired: async () => 0,
     queueDepth: async () => 0,
     activeCount: async () => 0,
-    getFireworksHealth: async () => 'healthy',
-    admitFromQueue: async ({ getFireworksHealth }) => {
+    getFleetHealth: async () => ({}),
+    admitFromQueue: async ({ health }) => {
       calls.admit += 1
-      const health = await getFireworksHealth()
       if (health !== 'healthy') {
         return { admitted: [], skipped: health }
       }
@@ -30,12 +30,16 @@ function makeAdmissionDeps(overrides: Partial<AdmissionDeps> = {}): AdmissionDep
     now: () => NOW,
     // Default to a single model so per-tick assertions (admitted: 1) stay
     // crisp regardless of how many production models are registered.
-    models: ['test-model'],
+    models: [TEST_MODEL],
     ...overrides,
   }
   return deps
 }
 
+function fleet(health: FireworksHealth, model: string = TEST_MODEL): FleetHealth {
+  return { [model]: health }
+}
+
 describe('runAdmissionTick', () => {
   test('admits one user per tick when healthy', async () => {
     const deps = makeAdmissionDeps()
@@ -44,18 +48,18 @@ describe('runAdmissionTick', () => {
     expect(result.skipped).toBeNull()
   })
 
-  test('skips admission when Fireworks is degraded', async () => {
+  test('skips admission when the model deployment is degraded', async () => {
     const deps = makeAdmissionDeps({
-      getFireworksHealth: async () => 'degraded' as FireworksHealth,
+      getFleetHealth: async () => fleet('degraded'),
     })
     const result = await runAdmissionTick(deps)
     expect(result.admitted).toBe(0)
     expect(result.skipped).toBe('degraded')
   })
 
-  test('skips admission when Fireworks is unhealthy', async () => {
+  test('skips admission when the model deployment is unhealthy', async () => {
     const deps = makeAdmissionDeps({
-      getFireworksHealth: async () => 'unhealthy' as FireworksHealth,
+      getFleetHealth: async () => fleet('unhealthy'),
     })
     const result = await runAdmissionTick(deps)
     expect(result.admitted).toBe(0)
@@ -69,13 +73,38 @@ describe('runAdmissionTick', () => {
         swept = 3
         return 3
       },
-      getFireworksHealth: async () => 'unhealthy' as FireworksHealth,
+      getFleetHealth: async () => fleet('unhealthy'),
     })
     const result = await runAdmissionTick(deps)
     expect(swept).toBe(3)
     expect(result.expired).toBe(3)
   })
 
+  test('admits per-model based on per-deployment health', async () => {
+    // Two models: 'good' is healthy, 'bad' is degraded. A single tick should
+    // admit 1 from 'good' and skip 'bad', surfacing the worst skip reason.
+    const deps = makeAdmissionDeps({
+      models: ['good', 'bad'],
+      getFleetHealth: async () => ({ good: 'healthy', bad: 'degraded' }),
+    })
+    const result = await runAdmissionTick(deps)
+    expect(result.admitted).toBe(1)
+    expect(result.skipped).toBe('degraded')
+  })
+
+  test('absent fleet entry defaults to healthy (serverless model)', async () => {
+    // Model isn't in the fleet map (e.g. served via Fireworks serverless).
+    // Admission should proceed rather than stall waiting for a probe that
+    // will never include this deployment.
+    const deps = makeAdmissionDeps({
+      models: ['serverless-model'],
+      getFleetHealth: async () => ({}),
+    })
+    const result = await runAdmissionTick(deps)
+    expect(result.admitted).toBe(1)
+    expect(result.skipped).toBeNull()
+  })
+
   test('propagates expiry count and admit count together', async () => {
     const deps = makeAdmissionDeps({
       sweepExpired: async () => 2,

web/src/server/free-session/__tests__/fireworks-health.test.ts

@@ -4,7 +4,7 @@ import {
   KV_BLOCKS_DEGRADED_FRACTION,
   KV_BLOCKS_UNHEALTHY_FRACTION,
   PREFILL_QUEUE_P90_DEGRADED_MS,
-  classify,
+  classifyOne,
 } from '../fireworks-health'
 
 type PromSample = { name: string; labels: Record<string, string>; value: number }
@@ -57,31 +57,31 @@ function errors(code: string, rate: number): PromSample {
 describe('fireworks health classifier', () => {
   test('healthy when queue well under the threshold', () => {
     const samples: PromSample[] = [kvBlocks(0.5), ...prefillQueueBuckets(150)]
-    expect(classify(samples, [DEPLOY])).toBe('healthy')
+    expect(classifyOne(samples, DEPLOY)).toBe('healthy')
   })
 
   test('degraded when prefill queue p90 exceeds the threshold', () => {
     const samples: PromSample[] = [
       kvBlocks(0.5),
       ...prefillQueueBuckets(PREFILL_QUEUE_P90_DEGRADED_MS + 500),
     ]
-    expect(classify(samples, [DEPLOY])).toBe('degraded')
+    expect(classifyOne(samples, DEPLOY)).toBe('degraded')
   })
 
   test('degraded when KV blocks cross the soft threshold (leading indicator)', () => {
     const samples: PromSample[] = [
       kvBlocks(KV_BLOCKS_DEGRADED_FRACTION + 0.01),
       ...prefillQueueBuckets(300),
     ]
-    expect(classify(samples, [DEPLOY])).toBe('degraded')
+    expect(classifyOne(samples, DEPLOY)).toBe('degraded')
   })
 
   test('unhealthy when KV blocks exceed the backstop', () => {
     const samples: PromSample[] = [
       kvBlocks(KV_BLOCKS_UNHEALTHY_FRACTION + 0.005),
       ...prefillQueueBuckets(300),
     ]
-    expect(classify(samples, [DEPLOY])).toBe('unhealthy')
+    expect(classifyOne(samples, DEPLOY)).toBe('unhealthy')
   })
 
   test('unhealthy when 5xx error fraction exceeds the threshold', () => {
@@ -91,7 +91,7 @@ describe('fireworks health classifier', () => {
       requests(1),
       errors('500', 0.2),
     ]
-    expect(classify(samples, [DEPLOY])).toBe('unhealthy')
+    expect(classifyOne(samples, DEPLOY)).toBe('unhealthy')
   })
 
   test('ignores high error fraction when traffic is too low to be meaningful', () => {
@@ -101,14 +101,17 @@ describe('fireworks health classifier', () => {
       requests(0.05),
       errors('500', 0.05),
     ]
-    expect(classify(samples, [DEPLOY])).toBe('healthy')
+    expect(classifyOne(samples, DEPLOY)).toBe('healthy')
   })
 
   test('healthy with no data yet (new deployment, no events)', () => {
-    expect(classify([], [DEPLOY])).toBe('healthy')
+    expect(classifyOne([], DEPLOY)).toBe('healthy')
   })
 
-  test('worst-of across multiple deployments — unhealthy wins over degraded', () => {
+  test('classifies deployments independently — one bad deployment does not affect another', () => {
+    // The fleet probe builds the result by classifying each deployment
+    // separately, so a saturated 'other' deployment leaves DEPLOY's
+    // (only-degraded) verdict intact.
     const other = 'other123'
     const samples: PromSample[] = [
       kvBlocks(0.5),
@@ -119,6 +122,7 @@ describe('fireworks health classifier', () => {
         value: KV_BLOCKS_UNHEALTHY_FRACTION + 0.005,
       },
     ]
-    expect(classify(samples, [DEPLOY, other])).toBe('unhealthy')
+    expect(classifyOne(samples, DEPLOY)).toBe('degraded')
+    expect(classifyOne(samples, other)).toBe('unhealthy')
   })
 })

web/src/server/free-session/admission.ts

@@ -6,10 +6,10 @@ import {
   getSessionLengthMs,
   isWaitingRoomEnabled,
 } from './config'
-import { getFireworksHealth } from './fireworks-health'
+import { getFleetHealth } from './fireworks-health'
 import { activeCount, admitFromQueue, queueDepth, sweepExpired } from './store'
 
-import type { FireworksHealth } from './fireworks-health'
+import type { FireworksHealth, FleetHealth } from './fireworks-health'
 
 import { logger } from '@/util/logger'
 
@@ -21,9 +21,9 @@ export interface AdmissionDeps {
     model: string
     sessionLengthMs: number
     now: Date
-    getFireworksHealth: () => Promise<FireworksHealth>
+    health: FireworksHealth
   }) => Promise<{ admitted: { user_id: string }[]; skipped: FireworksHealth | null }>
-  getFireworksHealth: () => Promise<FireworksHealth>
+  getFleetHealth: () => Promise<FleetHealth>
   /** Plain values, not thunks — these never change at runtime. */
   sessionLengthMs: number
   graceMs: number
@@ -38,11 +38,13 @@ const defaultDeps: AdmissionDeps = {
   activeCount,
   admitFromQueue,
   // FREEBUFF_DEV_FORCE_ADMIT lets local `dev:freebuff` drive the full
-  // waiting-room → admitted → ended flow without a real upstream.
-  getFireworksHealth:
+  // waiting-room → admitted → ended flow without a real upstream. Returning
+  // an empty fleet means every model resolves to the absence-default of
+  // 'healthy' below.
+  getFleetHealth:
     process.env.FREEBUFF_DEV_FORCE_ADMIT === 'true'
-      ? async () => 'healthy'
-      : getFireworksHealth,
+      ? async () => ({})
+      : getFleetHealth,
   get sessionLengthMs() {
     return getSessionLengthMs()
   },
@@ -81,21 +83,23 @@ export async function runAdmissionTick(
 
   const models = deps.models ?? FREEBUFF_MODELS.map((m) => m.id)
 
-  // Probe upstream health once per tick. Today every model shares a Fireworks
-  // deployment so a single probe gates them all — TODO: when we add a
-  // non-Fireworks model, plumb a model/deploymentId into the probe.
-  const health = await deps.getFireworksHealth()
-  const sharedHealth = async () => health
+  // One probe per tick covers every model — the Fireworks metrics endpoint
+  // returns all deployments in a single response. Models without a dedicated
+  // deployment (e.g. serverless) aren't in the map; treat their absence as
+  // 'healthy' so admission continues. TODO: when those models move to their
+  // own deployments, drop the absence-default and require an explicit entry.
+  const fleet = await deps.getFleetHealth()
 
   // Run per-model admission in parallel — they only contend on independent
   // advisory locks and a single update each.
   const perModel = await Promise.all(
     models.map(async (model) => {
+      const health = fleet[model] ?? 'healthy'
       const { admitted, skipped } = await deps.admitFromQueue({
         model,
         sessionLengthMs: deps.sessionLengthMs,
         now,
-        getFireworksHealth: sharedHealth,
+        health,
       })
       const depth = await deps.queueDepth({ model })
       return { model, admittedCount: admitted.length, depth, skipped }

web/src/server/free-session/fireworks-health.ts

@@ -52,25 +52,35 @@ const HEALTH_CHECK_TIMEOUT_MS = 5_000
  *  pod hits the endpoint at most ~2.4/min. */
 const HEALTH_CACHE_TTL_MS = 25_000
 
-type CacheEntry = { expiresAt: number; health: FireworksHealth }
+/** Map of model id → FireworksHealth. Only includes models that have a
+ *  dedicated Fireworks deployment in `FIREWORKS_DEPLOYMENT_MAP`. Models served
+ *  via the Fireworks serverless API (no deployment id) are not present —
+ *  callers should treat their absence as 'healthy' for now.
+ *  TODO: when serverless models move to dedicated deployments, drop the
+ *        absence-means-healthy fallback at the call site. */
+export type FleetHealth = Record<string, FireworksHealth>
+
+type CacheEntry = { expiresAt: number; fleet: FleetHealth }
 let cache: CacheEntry | null = null
 
 export function __resetFireworksHealthCacheForTests(): void {
   cache = null
 }
 
-export async function getFireworksHealth(): Promise<FireworksHealth> {
+export async function getFleetHealth(): Promise<FleetHealth> {
   const now = Date.now()
-  if (cache && cache.expiresAt > now) return cache.health
+  if (cache && cache.expiresAt > now) return cache.fleet
 
-  const health = await probe()
-  cache = { expiresAt: now + HEALTH_CACHE_TTL_MS, health }
-  return health
+  const fleet = await probe()
+  cache = { expiresAt: now + HEALTH_CACHE_TTL_MS, fleet }
+  return fleet
 }
 
-async function probe(): Promise<FireworksHealth> {
+async function probe(): Promise<FleetHealth> {
   const apiKey = env.FIREWORKS_API_KEY
-  if (!apiKey) return 'unhealthy'
+  // Mark every deployment-mapped model unhealthy when we can't authenticate
+  // the probe. Serverless models (absent from the map) keep their default.
+  if (!apiKey) return allDeploymentsAt('unhealthy')
 
   const controller = new AbortController()
   const timeout = setTimeout(() => controller.abort(), HEALTH_CHECK_TIMEOUT_MS)
@@ -81,18 +91,15 @@ async function probe(): Promise<FireworksHealth> {
       headers: { Authorization: `Bearer ${apiKey}` },
       signal: controller.signal,
     })
-    if (!response.ok) return 'unhealthy'
+    if (!response.ok) return allDeploymentsAt('unhealthy')
     body = await response.text()
   } catch {
-    return 'unhealthy'
+    return allDeploymentsAt('unhealthy')
   } finally {
     clearTimeout(timeout)
   }
 
-  const deploymentIds = Object.values(FIREWORKS_DEPLOYMENT_MAP).map(
-    (name) => name.split('/').pop()!,
-  )
-  if (deploymentIds.length === 0) return 'healthy'
+  if (Object.keys(FIREWORKS_DEPLOYMENT_MAP).length === 0) return {}
 
   const { samples, newestTimestampMs } = parsePrometheus(body)
 
@@ -104,27 +111,26 @@ async function probe(): Promise<FireworksHealth> {
       { ageMs: Date.now() - newestTimestampMs },
       '[FireworksHealth] unhealthy: metrics snapshot is stale',
     )
-    return 'unhealthy'
+    return allDeploymentsAt('unhealthy')
   }
 
-  return classify(samples, deploymentIds)
+  const fleet: FleetHealth = {}
+  for (const [modelId, deploymentName] of Object.entries(FIREWORKS_DEPLOYMENT_MAP)) {
+    const deploymentId = deploymentName.split('/').pop()!
+    fleet[modelId] = classifyOne(samples, deploymentId)
+  }
+  return fleet
 }
 
-/** Treat the whole fleet as degraded/unhealthy if any single deployment is. */
-export function classify(
-  samples: PromSample[],
-  deploymentIds: string[],
-): FireworksHealth {
-  let worst: FireworksHealth = 'healthy'
-  for (const deploymentId of deploymentIds) {
-    const h = classifyOne(samples, deploymentId)
-    if (h === 'unhealthy') return 'unhealthy'
-    if (h === 'degraded') worst = 'degraded'
+function allDeploymentsAt(health: FireworksHealth): FleetHealth {
+  const out: FleetHealth = {}
+  for (const modelId of Object.keys(FIREWORKS_DEPLOYMENT_MAP)) {
+    out[modelId] = health
   }
-  return worst
+  return out
 }
 
-function classifyOne(samples: PromSample[], deploymentId: string): FireworksHealth {
+export function classifyOne(samples: PromSample[], deploymentId: string): FireworksHealth {
   const kvBlocks = scalarFor(
     samples,
     'generator_kv_blocks_fraction:avg_by_deployment',