fix: cool an account down per served model, not for every model it serves

crates/dag/src/nodes.rs

@@ -308,9 +308,8 @@ impl DagNode for SelectAccount {
             .protocol()
             .ok_or_else(|| GatewayError::internal("select_account before resolve_model"))?;
         // post-variant: the served model's knobs ride the provider lookup
-        let conf = ctx
-            .cfg
-            .find_model(served_model(ctx.request.model_param_v2.as_ref()));
+        let model = served_model(ctx.request.model_param_v2.as_ref());
+        let conf = ctx.cfg.find_model(model);
         ctx.request.prompt_cache = conf.is_some_and(|m| m.prompt_cache);
         let provider = conf.and_then(|m| m.provider.as_deref());
         let latency = ctx
@@ -321,7 +320,7 @@ impl DagNode for SelectAccount {
         let account = ctx
             .state
             .pool
-            .select_healthy(mt, provider, &[], ctx.state.health.as_ref(), latency)
+            .select_healthy(mt, model, provider, &[], ctx.state.health.as_ref(), latency)
             .await;
         let Some(account) = account else {
             // unsampled, an exhausted pool would read no_data forever
@@ -516,6 +515,7 @@ impl DagNode for CallEngine {
                     .pool
                     .select_healthy(
                         mt,
+                        served_model(ctx.request.model_param_v2.as_ref()),
                         provider,
                         std::slice::from_ref(&failed.name),
                         ctx.state.health.as_ref(),
@@ -599,7 +599,11 @@ async fn note_engine_outcome(
         .avail
         .record(requested_model(ctx.request.model_param_v2.as_ref()), true);
     if let Some(account) = ctx.request.account.as_ref() {
-        ctx.state.health.record_success(&account.name).await;
+        let model = served_model(ctx.request.model_param_v2.as_ref());
+        ctx.state
+            .health
+            .record_success(&gw_state::health_key(&account.name, model))
+            .await;
         if let Some(started) = started {
             ctx.state.latency.record(&account.name, started.elapsed());
         }
@@ -630,23 +634,24 @@ fn account_fault(class: gw_consts::ErrClass) -> bool {
 async fn note_failure(ctx: &mut DagContext, account: &str) {
     let threshold = ctx.cfg.stability.failure_threshold;
     let secs = ctx.cfg.stability.cooldown_seconds;
+    let key = gw_state::health_key(account, served_model(ctx.request.model_param_v2.as_ref()));
     if !ctx
         .state
         .health
-        .record_failure(account, threshold, std::time::Duration::from_secs(secs))
+        .record_failure(&key, threshold, std::time::Duration::from_secs(secs))
         .await
     {
         return;
     }
+    ctx.decide(
+        "account_health",
+        format!("{key} entered cooldown ({secs}s)"),
+    );
     ctx.state.alerts.emit(
         "account_cooldown",
-        account.to_owned(),
+        key,
         format!("{threshold} consecutive failures; cooling {secs}s"),
     );
-    ctx.decide(
-        "account_health",
-        format!("{account} entered cooldown ({secs}s)"),
-    );
 }
 
 /// post_process/common_usage: RawUsageJSON -> CommonUsage.

crates/handler/src/lib.rs

@@ -2361,6 +2361,44 @@ mod tests {
         }
     }
 
+    #[derive(Debug)]
+    struct RefusesOneModel {
+        model: &'static str,
+        in_band: bool,
+    }
+
+    #[async_trait::async_trait]
+    impl gw_engines::transport::Transport for RefusesOneModel {
+        async fn send(
+            &self,
+            req: gw_engines::transport::UpstreamRequest,
+        ) -> GResult<gw_engines::transport::UpstreamResponse> {
+            let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap_or_default();
+            if body["model"] != self.model {
+                return gw_engines::MockTransport.send(req).await;
+            }
+            let refusal = r#"{"error":{"message":"high demand: exceeds the maximum usage size allowed during peak load","code":503}}"#;
+            let (status, body) = if self.in_band {
+                (
+                    200,
+                    gw_engines::transport::UpstreamBody::Sse(
+                        format!("data: {refusal}\n\n").into_bytes(),
+                    ),
+                )
+            } else {
+                (
+                    503,
+                    gw_engines::transport::UpstreamBody::Json(bytes::Bytes::from(refusal)),
+                )
+            };
+            Ok(gw_engines::transport::UpstreamResponse {
+                status,
+                body,
+                headers: Default::default(),
+            })
+        }
+    }
+
     #[derive(Debug)]
     struct PausedStream {
         release: Arc<tokio::sync::Notify>,
@@ -2485,7 +2523,10 @@ mod tests {
             let outcome = ctx.outcome.expect("outcome");
             assert!(outcome.terminal_error.is_some(), "{code} native={native}");
             assert_eq!(
-                h.state().health.available("a1").await,
+                h.state()
+                    .health
+                    .available(&gw_state::health_key("a1", "m"))
+                    .await,
                 !faults,
                 "{code} native={native}"
             );
@@ -2526,21 +2567,75 @@ mod tests {
                 .map(|error| error.class),
             Some(gw_consts::ErrClass::ModelStreamError)
         );
-        assert!(!h.state().health.available("a-down").await);
-        assert!(!h.state().health.available("a-up").await);
+        let health = &h.state().health;
+        assert!(!health.available(&gw_state::health_key("a-down", "m")).await);
+        assert!(!health.available(&gw_state::health_key("a-up", "m")).await);
         let first = alerts.try_recv().expect("first account cooldown alert");
         let second = alerts.try_recv().expect("retry account cooldown alert");
         assert_eq!(
             (first.kind, first.subject.as_str()),
-            ("account_cooldown", "a-down")
+            ("account_cooldown", "a-down/m")
         );
         assert_eq!(
             (second.kind, second.subject.as_str()),
-            ("account_cooldown", "a-up")
+            ("account_cooldown", "a-up/m")
         );
         assert!(alerts.try_recv().is_err(), "one alert per failed account");
     }
 
+    #[tokio::test]
+    async fn a_model_cooling_down_leaves_its_only_account_serving_other_models() {
+        let yaml = "listen: {host: h, port: 1}\nmodels: [{name: pub, protocol: openai-chat, variants: [{model: busy, weight: 1}]}, {name: busy, protocol: openai-chat}, {name: calm, protocol: openai-chat}]\naccounts: [{name: a1, provider: p, protocols: [openai-chat]}]\nstability: {failure_threshold: 1, cooldown_seconds: 30}\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]";
+        for (in_band, requested) in [(false, "busy"), (true, "pub")] {
+            let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap());
+            let state = Arc::new(GatewayState::from_config(&cfg));
+            let h = OnlineHandler::new(
+                gw_state::SharedConfig::new(cfg, state),
+                Arc::new(RefusesOneModel {
+                    model: "busy",
+                    in_band,
+                }),
+            );
+            let key = h.state().auth.get(&access_key_id("k1")).await.unwrap();
+            let refused = h
+                .run(drained_stream_req(requested, "hi"), key.clone())
+                .await;
+            assert!(
+                refused.map_or(true, |ctx| ctx
+                    .outcome
+                    .is_some_and(|o| o.terminal_error.is_some())),
+                "in_band={in_band}"
+            );
+            let health = &h.state().health;
+            assert!(
+                !health.available(&gw_state::health_key("a1", "busy")).await,
+                "in_band={in_band}"
+            );
+            assert!(
+                health.available(&gw_state::health_key("a1", "calm")).await,
+                "in_band={in_band}"
+            );
+
+            let ctx = h.run(chat_req("calm", "hi"), key.clone()).await.unwrap();
+            let outcome = ctx.outcome.expect("calm is served");
+            assert!(outcome.terminal_error.is_none(), "in_band={in_band}");
+            assert_eq!(
+                ctx.request.account.as_ref().map(|a| a.name.as_str()),
+                Some("a1")
+            );
+
+            let cooling = h
+                .run(chat_req(requested, "hi"), key)
+                .await
+                .err()
+                .expect("busy is cooling");
+            assert!(
+                cooling.message.contains("no healthy upstream"),
+                "in_band={in_band}: {cooling:?}"
+            );
+        }
+    }
+
     #[tokio::test]
     async fn client_disconnect_after_commit_stays_health_neutral() {
         let yaml = "listen: {host: h, port: 1}\nsecurity: {dlp_redact: false}\nmodels: [{name: m, protocol: openai-chat}]\naccounts: [{name: a1, provider: p, protocols: [openai-chat]}]\nstability: {failure_threshold: 1, cooldown_seconds: 30}\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]";
@@ -2575,7 +2670,12 @@ mod tests {
         avail.flush().await;
         let minute = gw_state::epoch_secs() / 60;
         assert_eq!(avail.window("m", minute - 5, minute).await, (0, 0));
-        assert!(h.state().health.available("a1").await);
+        assert!(
+            h.state()
+                .health
+                .available(&gw_state::health_key("a1", "m"))
+                .await
+        );
         assert!(alerts.try_recv().is_err());
     }
 

crates/server/tests/e2e.rs

@@ -4224,12 +4224,12 @@ async fn account_cooldown_and_recovery() {
             .unwrap()
             .iter()
             .find(|a| a["name"] == "mock-spark-down")
-            .unwrap()["health"]
-            .as_str()
             .unwrap()
-            .to_owned()
+            .clone()
     };
-    assert_eq!(spark_health(&app).await, "cooling");
+    let spark = spark_health(&app).await;
+    assert_eq!(spark["health"], "cooling");
+    assert_eq!(spark["cooling_models"], serde_json::json!(["spark-lite"]));
     let r = app
         .clone()
         .oneshot(post("/v1/chat/completions", Some("ak-demo-123"), body))
@@ -4245,7 +4245,7 @@ async fn account_cooldown_and_recovery() {
     let mut recovered = false;
     for _ in 0..40 {
         tokio::time::sleep(std::time::Duration::from_millis(200)).await;
-        if spark_health(&app).await == "ok" {
+        if spark_health(&app).await["health"] == "ok" {
             recovered = true;
             break;
         }

crates/state/src/lib.rs

@@ -456,20 +456,22 @@ impl AccountPool {
             .map(Arc::clone)
     }
 
-    /// PTU accounts are preferred over paygo; `excluded` (failed-over) and
-    /// cooldown accounts are skipped; within a tier, pick by priority then
-    /// round-robin. A model bound to a `provider` only uses that provider's
-    /// accounts, and only those are health-checked.
+    /// PTU accounts are preferred over paygo; `excluded` (failed-over) accounts
+    /// and accounts cooling down for `model` are skipped; within a tier, pick by
+    /// priority then round-robin. A model bound to a `provider` only uses that
+    /// provider's accounts, and only those are health-checked.
     pub async fn select_healthy(
         &self,
         p: Protocol,
+        model: &str,
         provider: Option<&str>,
         excluded: &[String],
         health: &dyn HealthStore,
         latency: Option<&latency::Latency>,
     ) -> Option<Arc<Account>> {
         let serving = || self.accounts.iter().filter(|a| serves(a, p, provider));
-        let checks = futures::future::join_all(serving().map(|a| health.available(&a.name))).await;
+        let keys: Vec<String> = serving().map(|a| health_key(&a.name, model)).collect();
+        let checks = futures::future::join_all(keys.iter().map(|k| health.available(k))).await;
         let unhealthy: Vec<&str> = serving()
             .zip(checks)
             .filter(|(_, ok)| !ok)
@@ -988,6 +990,10 @@ pub fn epoch_millis() -> i64 {
         .unwrap_or(0)
 }
 
+pub fn health_key(account: &str, model: &str) -> String {
+    format!("{account}/{model}")
+}
+
 /// Recovers a poisoned lock instead of panicking: every critical section here is infallible.
 pub(crate) fn lock<T>(m: &std::sync::Mutex<T>) -> std::sync::MutexGuard<'_, T> {
     m.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
@@ -1133,6 +1139,7 @@ mod tests {
             s.pool
                 .select_healthy(
                     Protocol::OpenaiChat,
+                    "m",
                     Some("p"),
                     &[],
                     s.health.as_ref(),
@@ -1525,32 +1532,59 @@ mod tests {
         let h = s.health.as_ref();
         let a = s
             .pool
-            .select_healthy(Protocol::OpenaiChat, Some("openai"), &[], h, None)
+            .select_healthy(Protocol::OpenaiChat, "m", Some("openai"), &[], h, None)
             .await
             .unwrap();
         let b = s
             .pool
-            .select_healthy(Protocol::OpenaiChat, Some("openai"), &[], h, None)
+            .select_healthy(Protocol::OpenaiChat, "m", Some("openai"), &[], h, None)
             .await
             .unwrap();
         assert_eq!(a.name, "mock-openai-1");
         assert_eq!(b.name, "mock-openai-1");
         assert_eq!(
             s.pool
-                .select_healthy(Protocol::AnthropicMessages, None, &[], h, None)
+                .select_healthy(Protocol::AnthropicMessages, "m", None, &[], h, None)
                 .await
                 .unwrap()
                 .name,
             "mock-anthropic-1"
         );
         assert!(
             s.pool
-                .select_healthy(Protocol::Video, Some("nonexistent"), &[], h, None)
+                .select_healthy(Protocol::Video, "m", Some("nonexistent"), &[], h, None)
                 .await
                 .is_none()
         );
     }
 
+    #[tokio::test]
+    async fn a_cooldown_benches_the_account_for_one_model_only() {
+        let yaml = "listen: {host: h, port: 1}\naccounts: [{name: only, provider: p, protocols: ['responses']}]";
+        let s = GatewayState::from_config(&GatewayConfig::from_yaml(yaml).unwrap());
+        let cd = Duration::from_secs(30);
+        assert!(
+            s.health
+                .record_failure(&health_key("only", "busy"), 1, cd)
+                .await
+        );
+        let pick = |model| {
+            s.pool.select_healthy(
+                Protocol::Responses,
+                model,
+                None,
+                &[],
+                s.health.as_ref(),
+                None,
+            )
+        };
+        assert!(pick("busy").await.is_none());
+        assert_eq!(
+            pick("calm").await.map(|a| a.name.clone()).as_deref(),
+            Some("only")
+        );
+    }
+
     #[test]
     fn failure_streak_decays_when_idle() {
         let h = AccountHealth::default();
@@ -1585,7 +1619,7 @@ mod tests {
         let h = s.health.as_ref();
         let first = s
             .pool
-            .select_healthy(Protocol::OpenaiChat, Some("tencent"), &[], h, None)
+            .select_healthy(Protocol::OpenaiChat, "m", Some("tencent"), &[], h, None)
             .await
             .unwrap();
         assert_eq!(first.name, "mock-hunyuan-ptu-down");
@@ -1594,6 +1628,7 @@ mod tests {
             .pool
             .select_healthy(
                 Protocol::OpenaiChat,
+                "m",
                 Some("tencent"),
                 &["mock-hunyuan-ptu-down".into()],
                 h,