hatchery_server/
routing.rs

1//! hq-core router: roles, presets, seats and a limit-aware fallback.
2//!
3//! Config (`<data_dir>/routing.json`, owner-edited, hot-reloaded on each call):
4//! ```json
5//! {
6//!   "seats": [{ "id": "box-claude", "node": "hatchery-box", "provider": "claude" }],
7//!   "roles": {
8//!     "dev":     { "preset": ["claude/opus:high", "codex/gpt-5:high"], "auto": true },
9//!     "advisor": { "preset": ["codex/gpt-5:medium"], "auto": false }
10//!   },
11//!   "limit_cooldown_secs": 3600,
12//!   "advisor": { "min_interval_secs": 300, "dedup_secs": 3600 }
13//! }
14//! ```
15//! A preset entry is `provider/model[:effort]`. A seat is one place a provider can run
16//! (a node + provider, optionally a named account). The router walks the role's preset in
17//! order and, for each entry, the seats of that provider; a seat with a recent `limit`
18//! signal is skipped until a `ready` signal, a parsed reset time, or the cooldown passes.
19//! Signals are hints (screen heuristics): they only reorder and never delete a seat, and
20//! when every seat is limited the router still answers with the soonest-free one.
21//! The router recommends; it picks on its own only in roles with `"auto": true`.
22
23use serde::{Deserialize, Serialize};
24use std::collections::BTreeMap;
25use std::path::Path;
26
27#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
28#[serde(deny_unknown_fields)]
29pub struct Seat {
30    pub id: String,
31    pub node: String,
32    pub provider: String,
33    #[serde(default)]
34    pub account: Option<String>,
35}
36
37#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
38#[serde(deny_unknown_fields)]
39pub struct Role {
40    pub preset: Vec<String>,
41    #[serde(default)]
42    pub auto: bool,
43}
44
45#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
46#[serde(deny_unknown_fields)]
47pub struct AdvisorPolicy {
48    #[serde(default = "d_interval")]
49    pub min_interval_secs: u64,
50    #[serde(default = "d_dedup")]
51    pub dedup_secs: u64,
52}
53fn d_interval() -> u64 { 300 }
54fn d_dedup() -> u64 { 3600 }
55impl Default for AdvisorPolicy {
56    fn default() -> Self { Self { min_interval_secs: d_interval(), dedup_secs: d_dedup() } }
57}
58
59#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
60#[serde(deny_unknown_fields)]
61pub struct RoutingConfig {
62    #[serde(default)]
63    pub seats: Vec<Seat>,
64    #[serde(default)]
65    pub roles: BTreeMap<String, Role>,
66    #[serde(default)]
67    pub limit_cooldown_secs: Option<u64>,
68    #[serde(default)]
69    pub advisor: AdvisorPolicy,
70}
71
72impl RoutingConfig {
73    pub fn load(data_dir: &Path) -> Result<Self, String> {
74        match std::fs::read(data_dir.join("routing.json")) {
75            Ok(b) => serde_json::from_slice(&b).map_err(|e| format!("routing.json: {e}")),
76            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()),
77            Err(e) => Err(format!("routing.json: {e}")),
78        }
79    }
80    fn cooldown_ms(&self) -> u64 {
81        self.limit_cooldown_secs.unwrap_or(3600) * 1000
82    }
83}
84
85/// `provider/model[:effort]`
86#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
87pub struct ModelRef {
88    pub provider: String,
89    pub model: String,
90    pub effort: Option<String>,
91}
92
93pub fn parse_model_ref(s: &str) -> Option<ModelRef> {
94    let (provider, rest) = s.split_once('/')?;
95    let (model, effort) = match rest.split_once(':') {
96        Some((m, e)) => (m, Some(e.to_owned()).filter(|e| !e.is_empty())),
97        None => (rest, None),
98    };
99    let ok = |x: &str| !x.is_empty() && x.bytes().all(|b| b.is_ascii_alphanumeric() || b"-_.".contains(&b));
100    (ok(provider) && ok(model) && effort.as_deref().is_none_or(ok)).then(|| ModelRef { provider: provider.into(), model: model.into(), effort })
101}
102
103/// Latest limit knowledge per (node, provider), built from session signals.
104#[derive(Clone, Debug, Default, Serialize, PartialEq, Eq)]
105pub struct SeatState {
106    /// When the seat is expected to be free again (ms); `None` = free.
107    pub limited_until_ms: Option<u64>,
108    pub hint: Option<String>,
109}
110
111/// Fold signals (oldest first) into per-(node, provider) state.
112pub fn fold_signals<'a>(signals: impl IntoIterator<Item = &'a crate::tasks::SessionSignal>, cfg: &RoutingConfig) -> BTreeMap<(String, String), SeatState> {
113    let mut out: BTreeMap<(String, String), SeatState> = BTreeMap::new();
114    for s in signals {
115        let key = (s.node_id.clone(), s.provider.clone());
116        match s.kind.as_str() {
117            "limit" => {
118                let until = s.hint.as_deref().and_then(|h| parse_reset(h, s.at_ms)).unwrap_or(s.at_ms + cfg.cooldown_ms());
119                out.insert(key, SeatState { limited_until_ms: Some(until), hint: s.hint.clone() });
120            }
121            "ready" => {
122                out.insert(key, SeatState::default());
123            }
124            _ => {}
125        }
126    }
127    out
128}
129
130/// "try again in 2 hours", "in 45 minutes": relative reset, capped at 8 days.
131pub fn parse_reset(hint: &str, at_ms: u64) -> Option<u64> {
132    let lower = hint.to_lowercase();
133    let idx = lower.find(" in ")?;
134    let mut words = lower[idx + 4..].split_whitespace();
135    let n: f64 = words.next()?.parse().ok()?;
136    let unit = words.next()?;
137    let secs = if unit.starts_with("min") { 60.0 } else if unit.starts_with("hour") || unit == "h" || unit.starts_with("hr") { 3600.0 } else if unit.starts_with("day") { 86400.0 } else { return None };
138    let ms = (n * secs * 1000.0) as u64;
139    (ms > 0 && ms <= 8 * 86_400_000).then_some(at_ms + ms)
140}
141
142#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
143pub struct Candidate {
144    pub seat: Seat,
145    pub model: ModelRef,
146    pub available: bool,
147    pub limited_until_ms: Option<u64>,
148    pub hint: Option<String>,
149}
150
151#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
152pub struct Recommendation {
153    pub role: String,
154    /// First available candidate (or the soonest-free one when all are limited).
155    pub pick: Option<Candidate>,
156    /// Whether the router may act on `pick` without the operator (role allows auto).
157    pub auto: bool,
158    /// The preferred candidate when `pick` is a fallback from it.
159    pub fallback_from: Option<Candidate>,
160    pub candidates: Vec<Candidate>,
161}
162
163pub fn recommend(cfg: &RoutingConfig, state: &BTreeMap<(String, String), SeatState>, role: &str, now_ms: u64) -> Result<Recommendation, String> {
164    let r = cfg.roles.get(role).ok_or("unknown role")?;
165    let mut candidates = Vec::new();
166    for entry in &r.preset {
167        let model = parse_model_ref(entry).ok_or_else(|| format!("bad preset entry: {entry}"))?;
168        for seat in cfg.seats.iter().filter(|s| s.provider == model.provider) {
169            let st = state.get(&(seat.node.clone(), seat.provider.clone())).cloned().unwrap_or_default();
170            let limited = st.limited_until_ms.filter(|u| *u > now_ms);
171            candidates.push(Candidate { seat: seat.clone(), model: model.clone(), available: limited.is_none(), limited_until_ms: limited, hint: limited.and(st.hint) });
172        }
173    }
174    let pick = candidates
175        .iter()
176        .find(|c| c.available)
177        .or_else(|| candidates.iter().min_by_key(|c| c.limited_until_ms.unwrap_or(u64::MAX)))
178        .cloned();
179    let fallback_from = match (&pick, candidates.first()) {
180        (Some(p), Some(first)) if p != first => Some(first.clone()),
181        _ => None,
182    };
183    Ok(Recommendation { role: role.into(), pick, auto: r.auto, fallback_from, candidates })
184}
185
186/// Advisor anti-noise gate: one advice per `min_interval` per advisor, and the same topic
187/// not again within `dedup`. Memory only (a restart forgets; that is acceptable noise).
188#[derive(Default)]
189pub struct AdviceGate {
190    last: BTreeMap<String, u64>,
191    topics: BTreeMap<(String, String), u64>,
192}
193
194impl AdviceGate {
195    pub fn check(&mut self, policy: &AdvisorPolicy, advisor: &str, topic: &str, now_ms: u64) -> Result<(), &'static str> {
196        let topic = topic.trim().to_lowercase();
197        if topic.is_empty() {
198            return Err("topic required");
199        }
200        if self.last.get(advisor).is_some_and(|t| now_ms < t + policy.min_interval_secs * 1000) {
201            return Err("too soon");
202        }
203        let key = (advisor.to_owned(), topic);
204        if self.topics.get(&key).is_some_and(|t| now_ms < t + policy.dedup_secs * 1000) {
205            return Err("duplicate topic");
206        }
207        self.last.insert(advisor.to_owned(), now_ms);
208        self.topics.insert(key, now_ms);
209        let horizon = now_ms.saturating_sub(policy.dedup_secs * 1000);
210        self.topics.retain(|_, t| *t >= horizon);
211        Ok(())
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use crate::tasks::SessionSignal;
219
220    fn cfg() -> RoutingConfig {
221        serde_json::from_value(serde_json::json!({
222            "seats": [
223                { "id": "a-claude", "node": "a", "provider": "claude" },
224                { "id": "b-claude", "node": "b", "provider": "claude" },
225                { "id": "a-codex", "node": "a", "provider": "codex" }
226            ],
227            "roles": {
228                "dev": { "preset": ["claude/opus:high", "codex/gpt-5:high"], "auto": true },
229                "advisor": { "preset": ["codex/gpt-5:medium"] }
230            },
231            "limit_cooldown_secs": 600
232        }))
233        .unwrap()
234    }
235
236    fn sig(id: i64, node: &str, provider: &str, kind: &str, hint: Option<&str>, at: u64) -> SessionSignal {
237        SessionSignal { id, node_id: node.into(), workspace_id: "w".into(), instance_id: 1, provider: provider.into(), kind: kind.into(), hint: hint.map(Into::into), at_ms: at }
238    }
239
240    #[test]
241    fn parses_model_refs() {
242        assert_eq!(parse_model_ref("claude/opus:high").unwrap().effort.as_deref(), Some("high"));
243        assert_eq!(parse_model_ref("codex/gpt-5").unwrap().effort, None);
244        assert!(parse_model_ref("bad").is_none());
245        assert!(parse_model_ref("a/b c").is_none());
246    }
247
248    #[test]
249    fn picks_first_free_seat_and_falls_back_on_limit() {
250        let c = cfg();
251        let r = recommend(&c, &BTreeMap::new(), "dev", 0).unwrap();
252        assert_eq!(r.pick.as_ref().unwrap().seat.id, "a-claude");
253        assert!(r.auto && r.fallback_from.is_none());
254        let sigs = [sig(1, "a", "claude", "limit", Some("usage limit, try again in 2 hours"), 1000)];
255        let st = fold_signals(&sigs, &c);
256        let r = recommend(&c, &st, "dev", 5000).unwrap();
257        assert_eq!(r.pick.as_ref().unwrap().seat.id, "b-claude");
258        assert_eq!(r.fallback_from.as_ref().unwrap().seat.id, "a-claude");
259        assert_eq!(r.candidates[0].limited_until_ms, Some(1000 + 7_200_000));
260        // Return after reset.
261        let r = recommend(&c, &st, "dev", 1000 + 7_200_001).unwrap();
262        assert_eq!(r.pick.unwrap().seat.id, "a-claude");
263    }
264
265    #[test]
266    fn ready_clears_and_cooldown_applies_without_hint() {
267        let c = cfg();
268        let sigs = [sig(1, "a", "claude", "limit", None, 0), sig(2, "b", "claude", "limit", None, 0)];
269        let st = fold_signals(&sigs, &c);
270        let r = recommend(&c, &st, "dev", 1).unwrap();
271        assert_eq!(r.pick.as_ref().unwrap().seat.provider, "codex");
272        // All limited: soonest-free is still offered.
273        let sigs2 = [sig(1, "a", "claude", "limit", None, 0), sig(2, "b", "claude", "limit", None, 10), sig(3, "a", "codex", "limit", None, 20)];
274        let r = recommend(&c, &fold_signals(&sigs2, &c), "dev", 30).unwrap();
275        assert_eq!(r.pick.as_ref().unwrap().seat.id, "a-claude");
276        assert!(!r.pick.unwrap().available);
277        // ready clears.
278        let sigs3 = [sig(1, "a", "claude", "limit", None, 0), sig(2, "a", "claude", "ready", None, 5)];
279        assert!(recommend(&c, &fold_signals(&sigs3, &c), "dev", 6).unwrap().pick.unwrap().available);
280        // After cooldown (600 s).
281        assert!(recommend(&c, &st, "dev", 600_001).unwrap().candidates.iter().all(|x| x.available));
282    }
283
284    #[test]
285    fn non_auto_role_and_unknown_role() {
286        let c = cfg();
287        assert!(!recommend(&c, &BTreeMap::new(), "advisor", 0).unwrap().auto);
288        assert!(recommend(&c, &BTreeMap::new(), "nope", 0).is_err());
289    }
290
291    #[test]
292    fn advice_gate_limits_noise() {
293        let p = AdvisorPolicy { min_interval_secs: 60, dedup_secs: 600 };
294        let mut g = AdviceGate::default();
295        assert!(g.check(&p, "adv", "tests failing", 0).is_ok());
296        assert_eq!(g.check(&p, "adv", "other", 1000), Err("too soon"));
297        assert_eq!(g.check(&p, "adv", "Tests failing ", 61_000), Err("duplicate topic"));
298        assert!(g.check(&p, "adv", "other", 61_000).is_ok());
299        assert!(g.check(&p, "adv2", "tests failing", 61_000).is_ok());
300        assert!(g.check(&p, "adv", "tests failing", 700_000).is_ok());
301    }
302}
hatchery_server/routing.rs

hatchery_server/
routing.rs
