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(compatibility) = identity.compatibility.as_ref() else {
100 return StationRoutingBalanceSummary::default();
101 };
102 let Some(snapshots) =
103 provider_balances.and_then(|balances| balances.get(compatibility.station_name.as_str()))
104 else {
105 return StationRoutingBalanceSummary::default();
106 };
107
108 StationRoutingBalanceSummary::from_snapshot_iter_at(
109 snapshots
110 .iter()
111 .filter(|snapshot| balance_snapshot_matches_candidate(snapshot, identity)),
112 now_ms,
113 )
114}
115
116fn balance_snapshot_matches_candidate(
117 snapshot: &ProviderBalanceSnapshot,
118 identity: &RuntimeUpstreamIdentity,
119) -> bool {
120 let Some(compatibility) = identity.compatibility.as_ref() else {
121 return false;
122 };
123 snapshot.provider_id == identity.provider_endpoint.provider_id
124 && snapshot.station_name.as_deref() == Some(compatibility.station_name.as_str())
125 && snapshot.upstream_index == Some(compatibility.upstream_index)
126}
127
128#[cfg(test)]
129mod tests {
130 use super::*;
131 use crate::balance::BalanceSnapshotStatus;
132 use crate::config::{
133 ProviderConcurrencyLimits, ProviderConfigV4, ProviderEndpointV4, ServiceConfig,
134 ServiceViewV4, UpstreamAuth, UpstreamConfig,
135 };
136 use crate::routing_ir::{compile_legacy_route_plan_template, compile_v4_route_plan_template};
137 use crate::state::{PassiveHealthState, UpstreamHealth};
138 use std::collections::BTreeMap;
139
140 fn provider(base_url: &str) -> ProviderConfigV4 {
141 ProviderConfigV4 {
142 base_url: Some(base_url.to_string()),
143 ..ProviderConfigV4::default()
144 }
145 }
146
147 fn passive_health(state: PassiveHealthState, score: u8) -> PassiveUpstreamHealth {
148 PassiveUpstreamHealth {
149 score,
150 state,
151 observed_at_ms: 100,
152 last_failure_at_ms: Some(100),
153 consecutive_failures: 1,
154 ..PassiveUpstreamHealth::default()
155 }
156 }
157
158 fn balance_snapshot(
159 provider_id: &str,
160 station_name: &str,
161 upstream_index: usize,
162 exhausted: bool,
163 ) -> ProviderBalanceSnapshot {
164 ProviderBalanceSnapshot {
165 provider_id: provider_id.to_string(),
166 station_name: Some(station_name.to_string()),
167 upstream_index: Some(upstream_index),
168 source: "test".to_string(),
169 fetched_at_ms: 100,
170 stale_after_ms: Some(200),
171 status: if exhausted {
172 BalanceSnapshotStatus::Exhausted
173 } else {
174 BalanceSnapshotStatus::Ok
175 },
176 exhausted: Some(exhausted),
177 ..ProviderBalanceSnapshot::default()
178 }
179 }
180
181 #[test]
182 fn route_candidate_runtime_signals_attach_existing_legacy_state() {
183 let view = ServiceViewV4 {
184 providers: BTreeMap::from([(
185 "input".to_string(),
186 provider("https://input.example/v1"),
187 )]),
188 ..ServiceViewV4::default()
189 };
190 let template = compile_v4_route_plan_template("codex", &view).expect("route template");
191
192 let station_health = HashMap::from([(
193 "routing".to_string(),
194 StationHealth {
195 checked_at_ms: 100,
196 upstreams: vec![UpstreamHealth {
197 base_url: "https://input.example/v1".to_string(),
198 passive: Some(passive_health(PassiveHealthState::Failing, 20)),
199 ..UpstreamHealth::default()
200 }],
201 },
202 )]);
203 let load_balancers = HashMap::from([(
204 "routing".to_string(),
205 LbConfigView {
206 last_good_index: None,
207 upstreams: vec![LbUpstreamView {
208 failure_count: 3,
209 cooldown_remaining_secs: Some(11),
210 usage_exhausted: true,
211 }],
212 },
213 )]);
214 let provider_balances = HashMap::from([(
215 "routing".to_string(),
216 vec![balance_snapshot("input", "routing", 0, true)],
217 )]);
218 let inputs = RouteRuntimeSignalInputs {
219 station_health: Some(&station_health),
220 load_balancers: Some(&load_balancers),
221 provider_balances: Some(&provider_balances),
222 now_ms: 150,
223 };
224
225 let signals = template.candidate_runtime_signal_view(&inputs);
226
227 assert_eq!(signals.len(), 1);
228 assert_eq!(
229 signals[0].identity.provider_endpoint.stable_key(),
230 "codex/input/default"
231 );
232 assert!(signals[0].identity.compatibility.is_none());
233 assert!(signals[0].passive_health.is_none());
234 assert!(signals[0].load_balancer.is_none());
235 assert!(signals[0].balance.is_empty());
236 }
237
238 #[test]
239 fn route_candidate_runtime_signals_disambiguate_multi_endpoint_provider_by_legacy_index() {
240 let mut endpoints = BTreeMap::new();
241 endpoints.insert(
242 "slow".to_string(),
243 ProviderEndpointV4 {
244 base_url: "https://slow.example/v1".to_string(),
245 enabled: true,
246 priority: 10,
247 tags: BTreeMap::new(),
248 supported_models: BTreeMap::new(),
249 model_mapping: BTreeMap::new(),
250 limits: ProviderConcurrencyLimits::default(),
251 },
252 );
253 endpoints.insert(
254 "fast".to_string(),
255 ProviderEndpointV4 {
256 base_url: "https://fast.example/v1".to_string(),
257 enabled: true,
258 priority: 0,
259 tags: BTreeMap::new(),
260 supported_models: BTreeMap::new(),
261 model_mapping: BTreeMap::new(),
262 limits: ProviderConcurrencyLimits::default(),
263 },
264 );
265 let view = ServiceViewV4 {
266 providers: BTreeMap::from([(
267 "input".to_string(),
268 ProviderConfigV4 {
269 endpoints,
270 auth: UpstreamAuth::default(),
271 ..ProviderConfigV4::default()
272 },
273 )]),
274 ..ServiceViewV4::default()
275 };
276 let template = compile_v4_route_plan_template("codex", &view).expect("route template");
277 let provider_balances = HashMap::from([(
278 "routing".to_string(),
279 vec![
280 balance_snapshot("input", "routing", 0, false),
281 balance_snapshot("input", "routing", 1, true),
282 ],
283 )]);
284 let load_balancers = HashMap::from([(
285 "routing".to_string(),
286 LbConfigView {
287 last_good_index: Some(0),
288 upstreams: vec![
289 LbUpstreamView {
290 failure_count: 0,
291 cooldown_remaining_secs: None,
292 usage_exhausted: false,
293 },
294 LbUpstreamView {
295 failure_count: 3,
296 cooldown_remaining_secs: Some(30),
297 usage_exhausted: true,
298 },
299 ],
300 },
301 )]);
302 let inputs = RouteRuntimeSignalInputs {
303 load_balancers: Some(&load_balancers),
304 provider_balances: Some(&provider_balances),
305 now_ms: 150,
306 ..RouteRuntimeSignalInputs::default()
307 };
308
309 let signals = template.candidate_runtime_signal_view(&inputs);
310
311 assert_eq!(
312 signals
313 .iter()
314 .map(|signal| signal.identity.provider_endpoint.stable_key())
315 .collect::<Vec<_>>(),
316 vec!["codex/input/fast", "codex/input/slow"]
317 );
318 assert!(signals[0].identity.compatibility.is_none());
319 assert!(signals[1].identity.compatibility.is_none());
320 assert!(signals[0].balance.is_empty());
321 assert!(signals[1].balance.is_empty());
322 assert!(signals[0].load_balancer.is_none());
323 assert!(signals[1].load_balancer.is_none());
324 }
325
326 #[test]
327 fn route_candidate_runtime_signals_keep_legacy_station_compatibility_reads() {
328 let service = ServiceConfig {
329 name: "primary".to_string(),
330 alias: Some("Primary".to_string()),
331 enabled: true,
332 level: 1,
333 upstreams: vec![UpstreamConfig {
334 base_url: "https://legacy.example/v1".to_string(),
335 auth: UpstreamAuth::default(),
336 tags: HashMap::from([
337 ("provider_id".to_string(), "legacy-provider".to_string()),
338 ("endpoint_id".to_string(), "legacy-endpoint".to_string()),
339 ]),
340 supported_models: HashMap::new(),
341 model_mapping: HashMap::new(),
342 }],
343 };
344 let template = compile_legacy_route_plan_template("codex", [&service]);
345
346 let station_health = HashMap::from([(
347 "primary".to_string(),
348 StationHealth {
349 checked_at_ms: 100,
350 upstreams: vec![UpstreamHealth {
351 base_url: "https://legacy.example/v1".to_string(),
352 passive: Some(passive_health(PassiveHealthState::Degraded, 60)),
353 ..UpstreamHealth::default()
354 }],
355 },
356 )]);
357 let load_balancers = HashMap::from([(
358 "primary".to_string(),
359 LbConfigView {
360 last_good_index: Some(0),
361 upstreams: vec![LbUpstreamView {
362 failure_count: 1,
363 cooldown_remaining_secs: None,
364 usage_exhausted: false,
365 }],
366 },
367 )]);
368 let provider_balances = HashMap::from([(
369 "primary".to_string(),
370 vec![balance_snapshot("legacy-provider", "primary", 0, false)],
371 )]);
372 let inputs = RouteRuntimeSignalInputs {
373 station_health: Some(&station_health),
374 load_balancers: Some(&load_balancers),
375 provider_balances: Some(&provider_balances),
376 now_ms: 150,
377 };
378
379 let signals = template.candidate_runtime_signal_view(&inputs);
380
381 assert_eq!(signals.len(), 1);
382 assert_eq!(
383 signals[0].identity.provider_endpoint.stable_key(),
384 "codex/legacy-provider/legacy-endpoint"
385 );
386 assert_eq!(
387 signals[0]
388 .identity
389 .compatibility
390 .as_ref()
391 .map(crate::runtime_identity::LegacyUpstreamKey::stable_key)
392 .as_deref(),
393 Some("codex/primary/0")
394 );
395 assert_eq!(
396 signals[0]
397 .passive_health
398 .as_ref()
399 .map(|health| health.state),
400 Some(PassiveHealthState::Degraded)
401 );
402 assert_eq!(
403 signals[0]
404 .load_balancer
405 .as_ref()
406 .map(|view| view.failure_count),
407 Some(1)
408 );
409 assert_eq!(signals[0].balance.ok, 1);
410 assert_eq!(signals[0].balance.routing_snapshots, 1);
411 }
412}