1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::balance::{ProviderBalanceSnapshot, StationRoutingBalanceSummary};
6use crate::routing_ir::{RouteCandidate, RoutePlanTemplate};
7use crate::runtime_identity::RuntimeUpstreamIdentity;
8use crate::state::{LbConfigView, LbUpstreamView, PassiveUpstreamHealth, StationHealth};
9
10#[derive(Debug, Clone, Copy, Default)]
11pub struct RouteRuntimeSignalInputs<'a> {
12 pub station_health: Option<&'a HashMap<String, StationHealth>>,
13 pub load_balancers: Option<&'a HashMap<String, LbConfigView>>,
14 pub provider_balances: Option<&'a HashMap<String, Vec<ProviderBalanceSnapshot>>>,
15 pub now_ms: u64,
16}
17
18impl<'a> RouteRuntimeSignalInputs<'a> {
19 pub fn empty(now_ms: u64) -> Self {
20 Self {
21 station_health: None,
22 load_balancers: None,
23 provider_balances: None,
24 now_ms,
25 }
26 }
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
30pub struct RouteCandidateRuntimeSignals {
31 pub identity: RuntimeUpstreamIdentity,
32 #[serde(default, skip_serializing_if = "Option::is_none")]
33 pub passive_health: Option<PassiveUpstreamHealth>,
34 #[serde(default, skip_serializing_if = "Option::is_none")]
35 pub load_balancer: Option<LbUpstreamView>,
36 #[serde(
37 default,
38 skip_serializing_if = "StationRoutingBalanceSummary::is_empty"
39 )]
40 pub balance: StationRoutingBalanceSummary,
41}
42
43impl RoutePlanTemplate {
44 pub fn candidate_runtime_signals(
45 &self,
46 candidate: &RouteCandidate,
47 inputs: &RouteRuntimeSignalInputs<'_>,
48 ) -> RouteCandidateRuntimeSignals {
49 let identity = self.candidate_identity(candidate);
50 RouteCandidateRuntimeSignals {
51 passive_health: candidate_passive_health(&identity, inputs.station_health),
52 load_balancer: candidate_load_balancer_state(&identity, inputs.load_balancers),
53 balance: candidate_balance_summary(&identity, inputs.provider_balances, inputs.now_ms),
54 identity,
55 }
56 }
57
58 pub fn candidate_runtime_signal_view(
59 &self,
60 inputs: &RouteRuntimeSignalInputs<'_>,
61 ) -> Vec<RouteCandidateRuntimeSignals> {
62 self.candidates
63 .iter()
64 .map(|candidate| self.candidate_runtime_signals(candidate, inputs))
65 .collect()
66 }
67}
68
69fn candidate_passive_health(
70 identity: &RuntimeUpstreamIdentity,
71 station_health: Option<&HashMap<String, StationHealth>>,
72) -> Option<PassiveUpstreamHealth> {
73 let compatibility = identity.compatibility.as_ref()?;
74 station_health?
75 .get(compatibility.station_name.as_str())?
76 .upstreams
77 .iter()
78 .find(|upstream| upstream.base_url == identity.base_url)
79 .and_then(|upstream| upstream.passive.clone())
80}
81
82fn candidate_load_balancer_state(
83 identity: &RuntimeUpstreamIdentity,
84 load_balancers: Option<&HashMap<String, LbConfigView>>,
85) -> Option<LbUpstreamView> {
86 let compatibility = identity.compatibility.as_ref()?;
87 load_balancers?
88 .get(compatibility.station_name.as_str())?
89 .upstreams
90 .get(compatibility.upstream_index)
91 .cloned()
92}
93
94fn candidate_balance_summary(
95 identity: &RuntimeUpstreamIdentity,
96 provider_balances: Option<&HashMap<String, Vec<ProviderBalanceSnapshot>>>,
97 now_ms: u64,
98) -> StationRoutingBalanceSummary {
99 let Some(provider_balances) = provider_balances else {
100 return StationRoutingBalanceSummary::default();
101 };
102
103 let snapshots = if let Some(compatibility) = identity.compatibility.as_ref() {
104 provider_balances
105 .get(compatibility.station_name.as_str())
106 .into_iter()
107 .flat_map(|snapshots| snapshots.iter())
108 .filter(|snapshot| balance_snapshot_matches_candidate(snapshot, identity))
109 .collect::<Vec<_>>()
110 } else {
111 let provider_endpoint_key = identity.provider_endpoint.stable_key();
112 provider_balances
113 .values()
114 .flat_map(|snapshots| snapshots.iter())
115 .filter(|snapshot| {
116 snapshot
117 .provider_endpoint_key
118 .as_deref()
119 .is_some_and(|key| key == provider_endpoint_key)
120 })
121 .collect::<Vec<_>>()
122 };
123
124 StationRoutingBalanceSummary::from_snapshot_iter_at(snapshots, now_ms)
125}
126
127fn balance_snapshot_matches_candidate(
128 snapshot: &ProviderBalanceSnapshot,
129 identity: &RuntimeUpstreamIdentity,
130) -> bool {
131 let Some(compatibility) = identity.compatibility.as_ref() else {
132 return false;
133 };
134 snapshot.provider_id == identity.provider_endpoint.provider_id
135 && snapshot.station_name.as_deref() == Some(compatibility.station_name.as_str())
136 && snapshot.upstream_index == Some(compatibility.upstream_index)
137}
138
139#[cfg(test)]
140mod tests {
141 use super::*;
142 use crate::balance::BalanceSnapshotStatus;
143 use crate::config::{
144 ProviderConcurrencyLimits, ProviderConfigV4, ProviderEndpointV4, ServiceConfig,
145 ServiceViewV4, UpstreamAuth, UpstreamConfig,
146 };
147 use crate::routing_ir::{compile_legacy_route_plan_template, compile_v4_route_plan_template};
148 use crate::state::{PassiveHealthState, UpstreamHealth};
149 use std::collections::BTreeMap;
150
151 fn provider(base_url: &str) -> ProviderConfigV4 {
152 ProviderConfigV4 {
153 base_url: Some(base_url.to_string()),
154 ..ProviderConfigV4::default()
155 }
156 }
157
158 fn passive_health(state: PassiveHealthState, score: u8) -> PassiveUpstreamHealth {
159 PassiveUpstreamHealth {
160 score,
161 state,
162 observed_at_ms: 100,
163 last_failure_at_ms: Some(100),
164 consecutive_failures: 1,
165 ..PassiveUpstreamHealth::default()
166 }
167 }
168
169 fn balance_snapshot(
170 provider_id: &str,
171 station_name: &str,
172 upstream_index: usize,
173 exhausted: bool,
174 ) -> ProviderBalanceSnapshot {
175 ProviderBalanceSnapshot {
176 provider_id: provider_id.to_string(),
177 station_name: Some(station_name.to_string()),
178 upstream_index: Some(upstream_index),
179 source: "test".to_string(),
180 fetched_at_ms: 100,
181 stale_after_ms: Some(200),
182 status: if exhausted {
183 BalanceSnapshotStatus::Exhausted
184 } else {
185 BalanceSnapshotStatus::Ok
186 },
187 exhausted: Some(exhausted),
188 ..ProviderBalanceSnapshot::default()
189 }
190 }
191
192 #[test]
193 fn route_candidate_runtime_signals_attach_existing_legacy_state() {
194 let view = ServiceViewV4 {
195 providers: BTreeMap::from([(
196 "input".to_string(),
197 provider("https://input.example/v1"),
198 )]),
199 ..ServiceViewV4::default()
200 };
201 let template = compile_v4_route_plan_template("codex", &view).expect("route template");
202
203 let station_health = HashMap::from([(
204 "routing".to_string(),
205 StationHealth {
206 checked_at_ms: 100,
207 upstreams: vec![UpstreamHealth {
208 base_url: "https://input.example/v1".to_string(),
209 passive: Some(passive_health(PassiveHealthState::Failing, 20)),
210 ..UpstreamHealth::default()
211 }],
212 },
213 )]);
214 let load_balancers = HashMap::from([(
215 "routing".to_string(),
216 LbConfigView {
217 last_good_index: None,
218 upstreams: vec![LbUpstreamView {
219 failure_count: 3,
220 cooldown_remaining_secs: Some(11),
221 usage_exhausted: true,
222 }],
223 },
224 )]);
225 let provider_balances = HashMap::from([(
226 "routing".to_string(),
227 vec![balance_snapshot("input", "routing", 0, true)],
228 )]);
229 let inputs = RouteRuntimeSignalInputs {
230 station_health: Some(&station_health),
231 load_balancers: Some(&load_balancers),
232 provider_balances: Some(&provider_balances),
233 now_ms: 150,
234 };
235
236 let signals = template.candidate_runtime_signal_view(&inputs);
237
238 assert_eq!(signals.len(), 1);
239 assert_eq!(
240 signals[0].identity.provider_endpoint.stable_key(),
241 "codex/input/default"
242 );
243 assert!(signals[0].identity.compatibility.is_none());
244 assert!(signals[0].passive_health.is_none());
245 assert!(signals[0].load_balancer.is_none());
246 assert!(signals[0].balance.is_empty());
247 }
248
249 #[test]
250 fn route_candidate_runtime_signals_disambiguate_multi_endpoint_provider_by_endpoint_key() {
251 let mut endpoints = BTreeMap::new();
252 endpoints.insert(
253 "slow".to_string(),
254 ProviderEndpointV4 {
255 base_url: "https://slow.example/v1".to_string(),
256 continuity_domain: None,
257 enabled: true,
258 priority: 10,
259 tags: BTreeMap::new(),
260 supported_models: BTreeMap::new(),
261 model_mapping: BTreeMap::new(),
262 limits: ProviderConcurrencyLimits::default(),
263 },
264 );
265 endpoints.insert(
266 "fast".to_string(),
267 ProviderEndpointV4 {
268 base_url: "https://fast.example/v1".to_string(),
269 continuity_domain: None,
270 enabled: true,
271 priority: 0,
272 tags: BTreeMap::new(),
273 supported_models: BTreeMap::new(),
274 model_mapping: BTreeMap::new(),
275 limits: ProviderConcurrencyLimits::default(),
276 },
277 );
278 let view = ServiceViewV4 {
279 providers: BTreeMap::from([(
280 "input".to_string(),
281 ProviderConfigV4 {
282 endpoints,
283 auth: UpstreamAuth::default(),
284 ..ProviderConfigV4::default()
285 },
286 )]),
287 ..ServiceViewV4::default()
288 };
289 let template = compile_v4_route_plan_template("codex", &view).expect("route template");
290 let provider_balances = HashMap::from([(
291 "routing".to_string(),
292 vec![
293 balance_snapshot("input", "routing", 0, false)
294 .with_provider_endpoint_key("codex/input/fast"),
295 balance_snapshot("input", "routing", 1, true)
296 .with_provider_endpoint_key("codex/input/slow"),
297 ],
298 )]);
299 let load_balancers = HashMap::from([(
300 "routing".to_string(),
301 LbConfigView {
302 last_good_index: Some(0),
303 upstreams: vec![
304 LbUpstreamView {
305 failure_count: 0,
306 cooldown_remaining_secs: None,
307 usage_exhausted: false,
308 },
309 LbUpstreamView {
310 failure_count: 3,
311 cooldown_remaining_secs: Some(30),
312 usage_exhausted: true,
313 },
314 ],
315 },
316 )]);
317 let inputs = RouteRuntimeSignalInputs {
318 load_balancers: Some(&load_balancers),
319 provider_balances: Some(&provider_balances),
320 now_ms: 150,
321 ..RouteRuntimeSignalInputs::default()
322 };
323
324 let signals = template.candidate_runtime_signal_view(&inputs);
325
326 assert_eq!(
327 signals
328 .iter()
329 .map(|signal| signal.identity.provider_endpoint.stable_key())
330 .collect::<Vec<_>>(),
331 vec!["codex/input/fast", "codex/input/slow"]
332 );
333 assert!(signals[0].identity.compatibility.is_none());
334 assert!(signals[1].identity.compatibility.is_none());
335 assert_eq!(signals[0].balance.ok, 1);
336 assert_eq!(signals[0].balance.routing_snapshots, 1);
337 assert_eq!(signals[1].balance.exhausted, 1);
338 assert_eq!(signals[1].balance.routing_exhausted, 1);
339 assert!(signals[0].load_balancer.is_none());
340 assert!(signals[1].load_balancer.is_none());
341 }
342
343 #[test]
344 fn route_candidate_runtime_signals_keep_legacy_station_compatibility_reads() {
345 let service = ServiceConfig {
346 name: "primary".to_string(),
347 alias: Some("Primary".to_string()),
348 enabled: true,
349 level: 1,
350 upstreams: vec![UpstreamConfig {
351 base_url: "https://legacy.example/v1".to_string(),
352 auth: UpstreamAuth::default(),
353 tags: HashMap::from([
354 ("provider_id".to_string(), "legacy-provider".to_string()),
355 ("endpoint_id".to_string(), "legacy-endpoint".to_string()),
356 ]),
357 supported_models: HashMap::new(),
358 model_mapping: HashMap::new(),
359 }],
360 };
361 let template = compile_legacy_route_plan_template("codex", [&service]);
362
363 let station_health = HashMap::from([(
364 "primary".to_string(),
365 StationHealth {
366 checked_at_ms: 100,
367 upstreams: vec![UpstreamHealth {
368 base_url: "https://legacy.example/v1".to_string(),
369 passive: Some(passive_health(PassiveHealthState::Degraded, 60)),
370 ..UpstreamHealth::default()
371 }],
372 },
373 )]);
374 let load_balancers = HashMap::from([(
375 "primary".to_string(),
376 LbConfigView {
377 last_good_index: Some(0),
378 upstreams: vec![LbUpstreamView {
379 failure_count: 1,
380 cooldown_remaining_secs: None,
381 usage_exhausted: false,
382 }],
383 },
384 )]);
385 let provider_balances = HashMap::from([(
386 "primary".to_string(),
387 vec![balance_snapshot("legacy-provider", "primary", 0, false)],
388 )]);
389 let inputs = RouteRuntimeSignalInputs {
390 station_health: Some(&station_health),
391 load_balancers: Some(&load_balancers),
392 provider_balances: Some(&provider_balances),
393 now_ms: 150,
394 };
395
396 let signals = template.candidate_runtime_signal_view(&inputs);
397
398 assert_eq!(signals.len(), 1);
399 assert_eq!(
400 signals[0].identity.provider_endpoint.stable_key(),
401 "codex/legacy-provider/legacy-endpoint"
402 );
403 assert_eq!(
404 signals[0]
405 .identity
406 .compatibility
407 .as_ref()
408 .map(crate::runtime_identity::LegacyUpstreamKey::stable_key)
409 .as_deref(),
410 Some("codex/primary/0")
411 );
412 assert_eq!(
413 signals[0]
414 .passive_health
415 .as_ref()
416 .map(|health| health.state),
417 Some(PassiveHealthState::Degraded)
418 );
419 assert_eq!(
420 signals[0]
421 .load_balancer
422 .as_ref()
423 .map(|view| view.failure_count),
424 Some(1)
425 );
426 assert_eq!(signals[0].balance.ok, 1);
427 assert_eq!(signals[0].balance.routing_snapshots, 1);
428 }
429}