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 continuity_domain: None,
246 enabled: true,
247 priority: 10,
248 tags: BTreeMap::new(),
249 supported_models: BTreeMap::new(),
250 model_mapping: BTreeMap::new(),
251 limits: ProviderConcurrencyLimits::default(),
252 },
253 );
254 endpoints.insert(
255 "fast".to_string(),
256 ProviderEndpointV4 {
257 base_url: "https://fast.example/v1".to_string(),
258 continuity_domain: None,
259 enabled: true,
260 priority: 0,
261 tags: BTreeMap::new(),
262 supported_models: BTreeMap::new(),
263 model_mapping: BTreeMap::new(),
264 limits: ProviderConcurrencyLimits::default(),
265 },
266 );
267 let view = ServiceViewV4 {
268 providers: BTreeMap::from([(
269 "input".to_string(),
270 ProviderConfigV4 {
271 endpoints,
272 auth: UpstreamAuth::default(),
273 ..ProviderConfigV4::default()
274 },
275 )]),
276 ..ServiceViewV4::default()
277 };
278 let template = compile_v4_route_plan_template("codex", &view).expect("route template");
279 let provider_balances = HashMap::from([(
280 "routing".to_string(),
281 vec![
282 balance_snapshot("input", "routing", 0, false),
283 balance_snapshot("input", "routing", 1, true),
284 ],
285 )]);
286 let load_balancers = HashMap::from([(
287 "routing".to_string(),
288 LbConfigView {
289 last_good_index: Some(0),
290 upstreams: vec![
291 LbUpstreamView {
292 failure_count: 0,
293 cooldown_remaining_secs: None,
294 usage_exhausted: false,
295 },
296 LbUpstreamView {
297 failure_count: 3,
298 cooldown_remaining_secs: Some(30),
299 usage_exhausted: true,
300 },
301 ],
302 },
303 )]);
304 let inputs = RouteRuntimeSignalInputs {
305 load_balancers: Some(&load_balancers),
306 provider_balances: Some(&provider_balances),
307 now_ms: 150,
308 ..RouteRuntimeSignalInputs::default()
309 };
310
311 let signals = template.candidate_runtime_signal_view(&inputs);
312
313 assert_eq!(
314 signals
315 .iter()
316 .map(|signal| signal.identity.provider_endpoint.stable_key())
317 .collect::<Vec<_>>(),
318 vec!["codex/input/fast", "codex/input/slow"]
319 );
320 assert!(signals[0].identity.compatibility.is_none());
321 assert!(signals[1].identity.compatibility.is_none());
322 assert!(signals[0].balance.is_empty());
323 assert!(signals[1].balance.is_empty());
324 assert!(signals[0].load_balancer.is_none());
325 assert!(signals[1].load_balancer.is_none());
326 }
327
328 #[test]
329 fn route_candidate_runtime_signals_keep_legacy_station_compatibility_reads() {
330 let service = ServiceConfig {
331 name: "primary".to_string(),
332 alias: Some("Primary".to_string()),
333 enabled: true,
334 level: 1,
335 upstreams: vec![UpstreamConfig {
336 base_url: "https://legacy.example/v1".to_string(),
337 auth: UpstreamAuth::default(),
338 tags: HashMap::from([
339 ("provider_id".to_string(), "legacy-provider".to_string()),
340 ("endpoint_id".to_string(), "legacy-endpoint".to_string()),
341 ]),
342 supported_models: HashMap::new(),
343 model_mapping: HashMap::new(),
344 }],
345 };
346 let template = compile_legacy_route_plan_template("codex", [&service]);
347
348 let station_health = HashMap::from([(
349 "primary".to_string(),
350 StationHealth {
351 checked_at_ms: 100,
352 upstreams: vec![UpstreamHealth {
353 base_url: "https://legacy.example/v1".to_string(),
354 passive: Some(passive_health(PassiveHealthState::Degraded, 60)),
355 ..UpstreamHealth::default()
356 }],
357 },
358 )]);
359 let load_balancers = HashMap::from([(
360 "primary".to_string(),
361 LbConfigView {
362 last_good_index: Some(0),
363 upstreams: vec![LbUpstreamView {
364 failure_count: 1,
365 cooldown_remaining_secs: None,
366 usage_exhausted: false,
367 }],
368 },
369 )]);
370 let provider_balances = HashMap::from([(
371 "primary".to_string(),
372 vec![balance_snapshot("legacy-provider", "primary", 0, false)],
373 )]);
374 let inputs = RouteRuntimeSignalInputs {
375 station_health: Some(&station_health),
376 load_balancers: Some(&load_balancers),
377 provider_balances: Some(&provider_balances),
378 now_ms: 150,
379 };
380
381 let signals = template.candidate_runtime_signal_view(&inputs);
382
383 assert_eq!(signals.len(), 1);
384 assert_eq!(
385 signals[0].identity.provider_endpoint.stable_key(),
386 "codex/legacy-provider/legacy-endpoint"
387 );
388 assert_eq!(
389 signals[0]
390 .identity
391 .compatibility
392 .as_ref()
393 .map(crate::runtime_identity::LegacyUpstreamKey::stable_key)
394 .as_deref(),
395 Some("codex/primary/0")
396 );
397 assert_eq!(
398 signals[0]
399 .passive_health
400 .as_ref()
401 .map(|health| health.state),
402 Some(PassiveHealthState::Degraded)
403 );
404 assert_eq!(
405 signals[0]
406 .load_balancer
407 .as_ref()
408 .map(|view| view.failure_count),
409 Some(1)
410 );
411 assert_eq!(signals[0].balance.ok, 1);
412 assert_eq!(signals[0].balance.routing_snapshots, 1);
413 }
414}