1use std::collections::{HashMap, HashSet, VecDeque};
2use std::io::{Read, Seek, SeekFrom};
3use std::path::PathBuf;
4use std::sync::Arc;
5use std::sync::Mutex;
6use std::sync::OnceLock;
7use std::sync::atomic::{AtomicU64, Ordering};
8
9use serde_json::Value as JsonValue;
10use tokio::sync::{RwLock, watch};
11use tokio::time::{Duration, interval};
12
13pub use crate::balance::{
14 BalanceSnapshotStatus, ProviderBalanceSnapshot, StationRoutingBalanceSummary,
15};
16use crate::config::ServiceConfigManager;
17use crate::lb::{COOLDOWN_SECS, CooldownBackoff, FAILURE_THRESHOLD, LbState};
18#[cfg(test)]
19use crate::pricing::CostBreakdown;
20use crate::pricing::{CostAdjustments, estimate_request_cost_from_operator_catalog_for_service};
21use crate::routing_ir::{RoutePlanRuntimeState, RoutePlanUpstreamRuntimeState};
22use crate::runtime_identity::ProviderEndpointKey;
23use crate::sessions;
24use crate::usage::UsageMetrics;
25
26mod runtime_types;
27mod session_identity;
28mod session_route_ledger;
29
30use self::runtime_types::{
31 ConfigMetaOverride, RuntimeDefaultProfileOverride, UsageRollup, merge_station_health,
32};
33pub use self::runtime_types::{
34 HealthCheckStatus, LbConfigView, LbUpstreamView, PassiveHealthState, PassiveUpstreamHealth,
35 RuntimeConfigState, StationHealth, UpstreamHealth, UsageBucket, UsageRollupCoverage,
36 UsageRollupView,
37};
38pub use self::session_identity::{
39 ActiveRequest, FinishRequestParams, FinishedRequest, RequestObservability, ResolvedRouteValue,
40 RouteDecisionProvenance, RouteValueSource, SessionBinding, SessionContinuityMode,
41 SessionIdentityCard, SessionIdentityCardBuildInputs, SessionIdentitySource,
42 SessionManualOverrides, SessionObservationScope, SessionRouteAffinity,
43 SessionRouteAffinityTarget, SessionStats, build_session_identity_cards_from_parts,
44 enrich_session_identity_cards_with_host_transcripts,
45 enrich_session_identity_cards_with_runtime,
46};
47use self::session_identity::{
48 SessionBindingEntry, SessionCwdCacheEntry, SessionEffortOverride, SessionModelOverride,
49 SessionRouteTargetOverride, SessionServiceTierOverride, SessionStationOverride,
50};
51use self::session_route_ledger::SessionRouteAffinityStore;
52
53type PassiveStationHealthMap =
54 HashMap<String, HashMap<String, HashMap<String, PassiveUpstreamHealth>>>;
55type ProviderBalanceMap =
56 HashMap<String, HashMap<String, HashMap<usize, HashMap<String, ProviderBalanceSnapshot>>>>;
57type ProviderBalanceSummaryMap = HashMap<String, HashMap<String, StationRoutingBalanceSummary>>;
58type ServiceLayoutSignature = Vec<(String, Vec<String>)>;
59
60#[derive(Debug, Clone, Default)]
61struct ProviderEndpointRuntimeHealth {
62 failure_count: u32,
63 cooldown_until: Option<std::time::Instant>,
64 usage_exhausted: bool,
65 penalty_streak: u32,
66 last_good_at_ms: Option<u64>,
67}
68
69#[derive(Debug, Clone)]
70struct SessionTranscriptPathCacheEntry {
71 path: Option<String>,
72 last_checked_ms: u64,
73 last_seen_ms: u64,
74}
75
76#[derive(Debug, Clone)]
77struct RuntimePolicy {
78 session_override_ttl_ms: u64,
79 session_binding_ttl_ms: u64,
80 session_binding_max_entries: usize,
81 session_route_affinity_ttl_ms: u64,
82 session_route_affinity_max_entries: usize,
83 session_route_affinity_store: SessionRouteAffinityStore,
84 session_cwd_cache_ttl_ms: u64,
85 session_cwd_cache_max_entries: usize,
86 session_transcript_path_cache_ttl_ms: u64,
87 session_transcript_path_cache_max_entries: usize,
88}
89
90pub struct PassiveUpstreamFailureRecord {
91 pub service_name: String,
92 pub station_name: String,
93 pub base_url: String,
94 pub status_code: Option<u16>,
95 pub error_class: Option<String>,
96 pub error: Option<String>,
97 pub now_ms: u64,
98}
99
100pub fn recent_finished_max() -> usize {
101 static MAX: OnceLock<usize> = OnceLock::new();
102 *MAX.get_or_init(|| {
103 std::env::var("CODEX_HELPER_RECENT_FINISHED_MAX")
104 .ok()
105 .and_then(|s| s.trim().parse::<usize>().ok())
106 .filter(|&n| n > 0)
107 .unwrap_or(2_000)
108 .clamp(200, 20_000)
109 })
110}
111
112fn unix_now_ms() -> u64 {
113 std::time::SystemTime::now()
114 .duration_since(std::time::UNIX_EPOCH)
115 .map(|d| d.as_millis() as u64)
116 .unwrap_or(0)
117}
118
119fn prune_lru_cache<T>(
120 cache: &mut HashMap<String, T>,
121 max_entries: usize,
122 last_seen: impl Fn(&T) -> u64,
123) {
124 if max_entries == 0 || cache.len() <= max_entries {
125 return;
126 }
127
128 let mut keys = cache
129 .iter()
130 .map(|(key, value)| (key.clone(), last_seen(value)))
131 .collect::<Vec<_>>();
132 keys.sort_by_key(|(_, seen)| *seen);
133 let remove_count = keys.len().saturating_sub(max_entries);
134 for (key, _) in keys.into_iter().take(remove_count) {
135 cache.remove(&key);
136 }
137}
138
139fn session_route_affinity_is_expired_with_ttl(
140 ttl_ms: u64,
141 affinity: &SessionRouteAffinity,
142 now_ms: u64,
143) -> bool {
144 ttl_ms > 0 && now_ms.saturating_sub(affinity.last_selected_at_ms) >= ttl_ms
145}
146
147fn prune_session_route_affinity_map(
148 affinities: &mut HashMap<String, SessionRouteAffinity>,
149 ttl_ms: u64,
150 max_entries: usize,
151 now_ms: u64,
152) {
153 affinities.retain(|_, affinity| {
154 !session_route_affinity_is_expired_with_ttl(ttl_ms, affinity, now_ms)
155 });
156 prune_lru_cache(affinities, max_entries, |entry| entry.last_selected_at_ms);
157}
158
159fn service_layout_signature(mgr: &ServiceConfigManager) -> ServiceLayoutSignature {
160 let mut entries = mgr
161 .stations()
162 .iter()
163 .map(|(station_name, service)| {
164 (
165 station_name.clone(),
166 service
167 .upstreams
168 .iter()
169 .map(|upstream| upstream.base_url.clone())
170 .collect::<Vec<_>>(),
171 )
172 })
173 .collect::<Vec<_>>();
174 entries.sort_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1)));
175 entries
176}
177
178fn changed_service_layout_stations(
179 previous: &ServiceLayoutSignature,
180 current: &ServiceLayoutSignature,
181) -> HashSet<String> {
182 let previous_by_station = previous
183 .iter()
184 .map(|(station_name, upstreams)| (station_name.as_str(), upstreams.as_slice()))
185 .collect::<HashMap<_, _>>();
186 let current_by_station = current
187 .iter()
188 .map(|(station_name, upstreams)| (station_name.as_str(), upstreams.as_slice()))
189 .collect::<HashMap<_, _>>();
190 let mut changed = HashSet::new();
191
192 for (station_name, upstreams) in previous {
193 if current_by_station
194 .get(station_name.as_str())
195 .is_none_or(|current_upstreams| *current_upstreams != upstreams.as_slice())
196 {
197 changed.insert(station_name.clone());
198 }
199 }
200
201 for (station_name, upstreams) in current {
202 if previous_by_station
203 .get(station_name.as_str())
204 .is_none_or(|previous_upstreams| *previous_upstreams != upstreams.as_slice())
205 {
206 changed.insert(station_name.clone());
207 }
208 }
209
210 changed
211}
212
213#[derive(Debug)]
218pub struct ProxyState {
219 next_request_id: AtomicU64,
220 session_override_ttl_ms: u64,
222 session_binding_ttl_ms: u64,
224 session_binding_max_entries: usize,
225 session_route_affinity_ttl_ms: u64,
226 session_route_affinity_max_entries: usize,
227 session_route_affinity_store: SessionRouteAffinityStore,
228 session_cwd_cache_ttl_ms: u64,
229 session_cwd_cache_max_entries: usize,
230 session_transcript_path_cache_ttl_ms: u64,
231 session_transcript_path_cache_max_entries: usize,
232 session_effort_overrides: RwLock<HashMap<String, SessionEffortOverride>>,
233 session_station_overrides: RwLock<HashMap<String, SessionStationOverride>>,
234 session_route_target_overrides: RwLock<HashMap<String, SessionRouteTargetOverride>>,
235 session_model_overrides: RwLock<HashMap<String, SessionModelOverride>>,
236 session_service_tier_overrides: RwLock<HashMap<String, SessionServiceTierOverride>>,
237 session_bindings: RwLock<HashMap<String, SessionBindingEntry>>,
238 session_route_affinities: RwLock<HashMap<String, SessionRouteAffinity>>,
239 global_station_override: RwLock<Option<String>>,
240 global_route_target_override: RwLock<Option<String>>,
241 runtime_default_profiles: RwLock<HashMap<String, RuntimeDefaultProfileOverride>>,
242 station_meta_overrides: RwLock<HashMap<String, HashMap<String, ConfigMetaOverride>>>,
243 provider_endpoint_meta_overrides:
245 RwLock<HashMap<String, HashMap<ProviderEndpointKey, ConfigMetaOverride>>>,
246 upstream_meta_overrides: RwLock<HashMap<String, HashMap<String, ConfigMetaOverride>>>,
248 session_cwd_cache: RwLock<HashMap<String, SessionCwdCacheEntry>>,
249 session_transcript_path_cache: RwLock<HashMap<String, SessionTranscriptPathCacheEntry>>,
250 session_stats: RwLock<HashMap<String, SessionStats>>,
251 active_requests: RwLock<HashMap<u64, ActiveRequest>>,
252 recent_finished: RwLock<VecDeque<FinishedRequest>>,
253 usage_rollups: RwLock<HashMap<String, UsageRollup>>,
254 station_health: RwLock<HashMap<String, HashMap<String, StationHealth>>>,
255 passive_station_health: RwLock<PassiveStationHealthMap>,
256 provider_balances: RwLock<ProviderBalanceMap>,
257 provider_balance_summaries: RwLock<ProviderBalanceSummaryMap>,
258 provider_endpoint_runtime_health:
259 RwLock<HashMap<String, HashMap<ProviderEndpointKey, ProviderEndpointRuntimeHealth>>>,
260 station_health_checks: RwLock<HashMap<String, HashMap<String, HealthCheckStatus>>>,
261 service_layout_signatures: RwLock<HashMap<String, ServiceLayoutSignature>>,
262 state_version_tx: watch::Sender<u64>,
263 lb_states: Option<Arc<Mutex<HashMap<String, LbState>>>>,
264}
265
266impl ProxyState {
267 const MAX_HEALTH_RECORDS_PER_STATION: usize = 200;
268
269 #[allow(dead_code)]
270 pub fn new() -> Arc<Self> {
271 Self::new_with_lb_states(None)
272 }
273
274 pub fn new_with_lb_states(
275 lb_states: Option<Arc<Mutex<HashMap<String, LbState>>>>,
276 ) -> Arc<Self> {
277 let ttl_secs = std::env::var("CODEX_HELPER_SESSION_OVERRIDE_TTL_SECS")
278 .ok()
279 .and_then(|s| s.trim().parse::<u64>().ok())
280 .filter(|&n| n > 0)
281 .unwrap_or(30 * 60);
282 let ttl_ms = ttl_secs.saturating_mul(1000);
283 let binding_ttl_secs = std::env::var("CODEX_HELPER_SESSION_BINDING_TTL_SECS")
284 .ok()
285 .and_then(|s| s.trim().parse::<u64>().ok())
286 .unwrap_or(0);
287 let binding_ttl_ms = binding_ttl_secs.saturating_mul(1000);
288 let binding_max_entries = std::env::var("CODEX_HELPER_SESSION_BINDING_MAX_ENTRIES")
289 .ok()
290 .and_then(|s| s.trim().parse::<usize>().ok())
291 .unwrap_or(2_000);
292 let route_affinity_ttl_secs = std::env::var("CODEX_HELPER_SESSION_ROUTE_AFFINITY_TTL_SECS")
293 .ok()
294 .and_then(|s| s.trim().parse::<u64>().ok())
295 .unwrap_or(0);
296 let route_affinity_ttl_ms = route_affinity_ttl_secs.saturating_mul(1000);
297 let route_affinity_max_entries =
298 std::env::var("CODEX_HELPER_SESSION_ROUTE_AFFINITY_MAX_ENTRIES")
299 .ok()
300 .and_then(|s| s.trim().parse::<usize>().ok())
301 .unwrap_or(5_000);
302
303 let cwd_cache_ttl_secs = std::env::var("CODEX_HELPER_SESSION_CWD_CACHE_TTL_SECS")
304 .ok()
305 .and_then(|s| s.trim().parse::<u64>().ok())
306 .unwrap_or(12 * 60 * 60);
307 let cwd_cache_ttl_ms = cwd_cache_ttl_secs.saturating_mul(1000);
308 let cwd_cache_max_entries = std::env::var("CODEX_HELPER_SESSION_CWD_CACHE_MAX_ENTRIES")
309 .ok()
310 .and_then(|s| s.trim().parse::<usize>().ok())
311 .unwrap_or(2_000);
312 let transcript_path_cache_ttl_secs =
313 std::env::var("CODEX_HELPER_SESSION_TRANSCRIPT_PATH_CACHE_TTL_SECS")
314 .ok()
315 .and_then(|s| s.trim().parse::<u64>().ok())
316 .unwrap_or(30);
317 let transcript_path_cache_ttl_ms = transcript_path_cache_ttl_secs.saturating_mul(1000);
318 let transcript_path_cache_max_entries =
319 std::env::var("CODEX_HELPER_SESSION_TRANSCRIPT_PATH_CACHE_MAX_ENTRIES")
320 .ok()
321 .and_then(|s| s.trim().parse::<usize>().ok())
322 .unwrap_or(5_000);
323
324 Self::new_with_runtime_policy(
325 lb_states,
326 RuntimePolicy {
327 session_override_ttl_ms: ttl_ms,
328 session_binding_ttl_ms: binding_ttl_ms,
329 session_binding_max_entries: binding_max_entries,
330 session_route_affinity_ttl_ms: route_affinity_ttl_ms,
331 session_route_affinity_max_entries: route_affinity_max_entries,
332 session_route_affinity_store: SessionRouteAffinityStore::from_env(),
333 session_cwd_cache_ttl_ms: cwd_cache_ttl_ms,
334 session_cwd_cache_max_entries: cwd_cache_max_entries,
335 session_transcript_path_cache_ttl_ms: transcript_path_cache_ttl_ms,
336 session_transcript_path_cache_max_entries: transcript_path_cache_max_entries,
337 },
338 )
339 }
340
341 fn new_with_runtime_policy(
342 lb_states: Option<Arc<Mutex<HashMap<String, LbState>>>>,
343 policy: RuntimePolicy,
344 ) -> Arc<Self> {
345 let mut restored_session_route_affinities = policy.session_route_affinity_store.load();
346 prune_session_route_affinity_map(
347 &mut restored_session_route_affinities,
348 policy.session_route_affinity_ttl_ms,
349 policy.session_route_affinity_max_entries,
350 unix_now_ms(),
351 );
352
353 Arc::new(Self {
354 next_request_id: AtomicU64::new(1),
355 session_override_ttl_ms: policy.session_override_ttl_ms,
356 session_binding_ttl_ms: policy.session_binding_ttl_ms,
357 session_binding_max_entries: policy.session_binding_max_entries,
358 session_route_affinity_ttl_ms: policy.session_route_affinity_ttl_ms,
359 session_route_affinity_max_entries: policy.session_route_affinity_max_entries,
360 session_route_affinity_store: policy.session_route_affinity_store,
361 session_cwd_cache_ttl_ms: policy.session_cwd_cache_ttl_ms,
362 session_cwd_cache_max_entries: policy.session_cwd_cache_max_entries,
363 session_transcript_path_cache_ttl_ms: policy.session_transcript_path_cache_ttl_ms,
364 session_transcript_path_cache_max_entries: policy
365 .session_transcript_path_cache_max_entries,
366 session_effort_overrides: RwLock::new(HashMap::new()),
367 session_station_overrides: RwLock::new(HashMap::new()),
368 session_route_target_overrides: RwLock::new(HashMap::new()),
369 session_model_overrides: RwLock::new(HashMap::new()),
370 session_service_tier_overrides: RwLock::new(HashMap::new()),
371 session_bindings: RwLock::new(HashMap::new()),
372 session_route_affinities: RwLock::new(restored_session_route_affinities),
373 global_station_override: RwLock::new(None),
374 global_route_target_override: RwLock::new(None),
375 runtime_default_profiles: RwLock::new(HashMap::new()),
376 station_meta_overrides: RwLock::new(HashMap::new()),
377 provider_endpoint_meta_overrides: RwLock::new(HashMap::new()),
378 upstream_meta_overrides: RwLock::new(HashMap::new()),
379 session_cwd_cache: RwLock::new(HashMap::new()),
380 session_transcript_path_cache: RwLock::new(HashMap::new()),
381 session_stats: RwLock::new(HashMap::new()),
382 active_requests: RwLock::new(HashMap::new()),
383 recent_finished: RwLock::new(VecDeque::new()),
384 usage_rollups: RwLock::new(HashMap::new()),
385 station_health: RwLock::new(HashMap::new()),
386 passive_station_health: RwLock::new(HashMap::new()),
387 provider_balances: RwLock::new(HashMap::new()),
388 provider_balance_summaries: RwLock::new(HashMap::new()),
389 provider_endpoint_runtime_health: RwLock::new(HashMap::new()),
390 station_health_checks: RwLock::new(HashMap::new()),
391 service_layout_signatures: RwLock::new(HashMap::new()),
392 state_version_tx: watch::channel(0).0,
393 lb_states,
394 })
395 }
396
397 pub fn subscribe_state_changes(&self) -> watch::Receiver<u64> {
398 self.state_version_tx.subscribe()
399 }
400
401 fn notify_state_changed(&self) {
402 self.state_version_tx.send_modify(|version| {
403 *version = version.wrapping_add(1);
404 });
405 }
406
407 pub async fn get_session_effort_override(&self, session_id: &str) -> Option<String> {
408 let guard = self.session_effort_overrides.read().await;
409 guard.get(session_id).map(|v| v.effort.clone())
410 }
411
412 pub async fn get_session_reasoning_effort_override(&self, session_id: &str) -> Option<String> {
413 self.get_session_effort_override(session_id).await
414 }
415
416 pub async fn set_session_effort_override(
417 &self,
418 session_id: String,
419 effort: String,
420 now_ms: u64,
421 ) {
422 let mut guard = self.session_effort_overrides.write().await;
423 guard.insert(
424 session_id,
425 SessionEffortOverride {
426 effort,
427 updated_at_ms: now_ms,
428 last_seen_ms: now_ms,
429 },
430 );
431 self.notify_state_changed();
432 }
433
434 pub async fn set_session_reasoning_effort_override(
435 &self,
436 session_id: String,
437 reasoning_effort: String,
438 now_ms: u64,
439 ) {
440 self.set_session_effort_override(session_id, reasoning_effort, now_ms)
441 .await;
442 }
443
444 pub async fn clear_session_effort_override(&self, session_id: &str) {
445 let mut guard = self.session_effort_overrides.write().await;
446 if guard.remove(session_id).is_some() {
447 self.notify_state_changed();
448 }
449 }
450
451 pub async fn clear_session_reasoning_effort_override(&self, session_id: &str) {
452 self.clear_session_effort_override(session_id).await;
453 }
454
455 pub async fn list_session_effort_overrides(&self) -> HashMap<String, String> {
456 let guard = self.session_effort_overrides.read().await;
457 guard
458 .iter()
459 .map(|(k, v)| (k.clone(), v.effort.clone()))
460 .collect()
461 }
462
463 pub async fn list_session_reasoning_effort_overrides(&self) -> HashMap<String, String> {
464 self.list_session_effort_overrides().await
465 }
466
467 pub async fn touch_session_override(&self, session_id: &str, now_ms: u64) {
468 let mut guard = self.session_effort_overrides.write().await;
469 if let Some(v) = guard.get_mut(session_id) {
470 v.last_seen_ms = now_ms;
471 }
472 }
473
474 pub async fn touch_session_reasoning_effort_override(&self, session_id: &str, now_ms: u64) {
475 self.touch_session_override(session_id, now_ms).await;
476 }
477
478 pub async fn get_session_station_override(&self, session_id: &str) -> Option<String> {
479 let guard = self.session_station_overrides.read().await;
480 guard.get(session_id).map(|v| v.station_name.clone())
481 }
482
483 pub async fn get_session_route_target_override(&self, session_id: &str) -> Option<String> {
484 let guard = self.session_route_target_overrides.read().await;
485 guard.get(session_id).map(|v| v.target.clone())
486 }
487
488 pub async fn get_session_model_override(&self, session_id: &str) -> Option<String> {
489 let guard = self.session_model_overrides.read().await;
490 guard.get(session_id).map(|v| v.model.clone())
491 }
492
493 pub async fn set_session_model_override(&self, session_id: String, model: String, now_ms: u64) {
494 let mut guard = self.session_model_overrides.write().await;
495 guard.insert(
496 session_id,
497 SessionModelOverride {
498 model,
499 updated_at_ms: now_ms,
500 last_seen_ms: now_ms,
501 },
502 );
503 self.notify_state_changed();
504 }
505
506 pub async fn clear_session_model_override(&self, session_id: &str) {
507 let mut guard = self.session_model_overrides.write().await;
508 if guard.remove(session_id).is_some() {
509 self.notify_state_changed();
510 }
511 }
512
513 pub async fn list_session_model_overrides(&self) -> HashMap<String, String> {
514 let guard = self.session_model_overrides.read().await;
515 guard
516 .iter()
517 .map(|(k, v)| (k.clone(), v.model.clone()))
518 .collect()
519 }
520
521 pub async fn touch_session_model_override(&self, session_id: &str, now_ms: u64) {
522 let mut guard = self.session_model_overrides.write().await;
523 if let Some(v) = guard.get_mut(session_id) {
524 v.last_seen_ms = now_ms;
525 }
526 }
527
528 pub async fn get_session_service_tier_override(&self, session_id: &str) -> Option<String> {
529 let guard = self.session_service_tier_overrides.read().await;
530 guard.get(session_id).map(|v| v.service_tier.clone())
531 }
532
533 pub async fn set_session_service_tier_override(
534 &self,
535 session_id: String,
536 service_tier: String,
537 now_ms: u64,
538 ) {
539 let mut guard = self.session_service_tier_overrides.write().await;
540 guard.insert(
541 session_id,
542 SessionServiceTierOverride {
543 service_tier,
544 updated_at_ms: now_ms,
545 last_seen_ms: now_ms,
546 },
547 );
548 self.notify_state_changed();
549 }
550
551 pub async fn clear_session_service_tier_override(&self, session_id: &str) {
552 let mut guard = self.session_service_tier_overrides.write().await;
553 if guard.remove(session_id).is_some() {
554 self.notify_state_changed();
555 }
556 }
557
558 pub async fn list_session_service_tier_overrides(&self) -> HashMap<String, String> {
559 let guard = self.session_service_tier_overrides.read().await;
560 guard
561 .iter()
562 .map(|(k, v)| (k.clone(), v.service_tier.clone()))
563 .collect()
564 }
565
566 pub async fn touch_session_service_tier_override(&self, session_id: &str, now_ms: u64) {
567 let mut guard = self.session_service_tier_overrides.write().await;
568 if let Some(v) = guard.get_mut(session_id) {
569 v.last_seen_ms = now_ms;
570 }
571 }
572
573 pub async fn get_session_binding(&self, session_id: &str) -> Option<SessionBinding> {
574 let guard = self.session_bindings.read().await;
575 guard.get(session_id).map(|entry| entry.binding.clone())
576 }
577
578 pub async fn list_session_bindings(&self) -> HashMap<String, SessionBinding> {
579 let guard = self.session_bindings.read().await;
580 guard
581 .iter()
582 .map(|(sid, entry)| (sid.clone(), entry.binding.clone()))
583 .collect()
584 }
585
586 pub async fn get_session_route_affinity(
587 &self,
588 session_id: &str,
589 ) -> Option<SessionRouteAffinity> {
590 let (affinity, snapshot) = {
591 let mut guard = self.session_route_affinities.write().await;
592 let affinity = guard.get(session_id).cloned()?;
593 if self.session_route_affinity_is_expired(&affinity, unix_now_ms()) {
594 guard.remove(session_id);
595 (None, Some(guard.clone()))
596 } else {
597 (Some(affinity), None)
598 }
599 };
600 if let Some(snapshot) = snapshot {
601 self.persist_session_route_affinities(snapshot).await;
602 }
603 affinity
604 }
605
606 pub async fn list_session_route_affinities(&self) -> HashMap<String, SessionRouteAffinity> {
607 let (affinities, snapshot) = {
608 let mut guard = self.session_route_affinities.write().await;
609 let now_ms = unix_now_ms();
610 let before_len = guard.len();
611 guard.retain(|_, affinity| !self.session_route_affinity_is_expired(affinity, now_ms));
612 let affinities = guard.clone();
613 let snapshot = (guard.len() != before_len).then(|| guard.clone());
614 (affinities, snapshot)
615 };
616 if let Some(snapshot) = snapshot {
617 self.persist_session_route_affinities(snapshot).await;
618 }
619 affinities
620 }
621
622 pub async fn record_session_route_affinity_success(
623 &self,
624 session_id: &str,
625 target: SessionRouteAffinityTarget,
626 reason_hint: Option<String>,
627 now_ms: u64,
628 ) -> SessionRouteAffinity {
629 let (affinity, snapshot) = {
630 let mut guard = self.session_route_affinities.write().await;
631 let affinity = match guard.get_mut(session_id) {
632 Some(existing) if target.same_target(existing) => {
633 existing.last_selected_at_ms = now_ms;
634 if target.session_identity_source.is_some() {
635 existing.session_identity_source = target.session_identity_source;
636 }
637 existing.clone()
638 }
639 Some(_) => {
640 let reason = reason_hint.unwrap_or_else(|| "target_changed".to_string());
641 let affinity = SessionRouteAffinity {
642 route_graph_key: target.route_graph_key,
643 session_identity_source: target.session_identity_source,
644 provider_endpoint: target.provider_endpoint,
645 upstream_base_url: target.upstream_base_url,
646 route_path: target.route_path,
647 last_selected_at_ms: now_ms,
648 last_changed_at_ms: now_ms,
649 change_reason: reason,
650 };
651 guard.insert(session_id.to_string(), affinity.clone());
652 affinity
653 }
654 None => {
655 let reason = reason_hint.unwrap_or_else(|| "first_success".to_string());
656 let affinity = SessionRouteAffinity {
657 route_graph_key: target.route_graph_key,
658 session_identity_source: target.session_identity_source,
659 provider_endpoint: target.provider_endpoint,
660 upstream_base_url: target.upstream_base_url,
661 route_path: target.route_path,
662 last_selected_at_ms: now_ms,
663 last_changed_at_ms: now_ms,
664 change_reason: reason,
665 };
666 guard.insert(session_id.to_string(), affinity.clone());
667 affinity
668 }
669 };
670 prune_lru_cache(
671 &mut guard,
672 self.session_route_affinity_max_entries,
673 |entry| entry.last_selected_at_ms,
674 );
675 (affinity, guard.clone())
676 };
677 self.persist_session_route_affinities(snapshot).await;
678 self.notify_state_changed();
679 affinity
680 }
681
682 fn session_route_affinity_is_expired(
683 &self,
684 affinity: &SessionRouteAffinity,
685 now_ms: u64,
686 ) -> bool {
687 session_route_affinity_is_expired_with_ttl(
688 self.session_route_affinity_ttl_ms,
689 affinity,
690 now_ms,
691 )
692 }
693
694 async fn persist_session_route_affinities(
695 &self,
696 snapshot: HashMap<String, SessionRouteAffinity>,
697 ) {
698 if let Err(err) = self
699 .session_route_affinity_store
700 .save(&snapshot, unix_now_ms())
701 .await
702 {
703 tracing::warn!(error = %err, "failed to persist session route affinity ledger");
704 }
705 }
706
707 pub async fn set_session_binding(&self, binding: SessionBinding) {
708 let mut guard = self.session_bindings.write().await;
709 let binding = if let Some(existing) = guard.get(binding.session_id.as_str()) {
710 SessionBinding {
711 created_at_ms: existing.binding.created_at_ms,
712 ..binding
713 }
714 } else {
715 binding
716 };
717 guard.insert(binding.session_id.clone(), SessionBindingEntry { binding });
718 self.notify_state_changed();
719 }
720
721 pub async fn clear_session_binding(&self, session_id: &str) {
722 let mut guard = self.session_bindings.write().await;
723 if guard.remove(session_id).is_some() {
724 self.notify_state_changed();
725 }
726 }
727
728 pub async fn clear_session_manual_overrides(&self, session_id: &str) {
729 self.clear_session_station_override(session_id).await;
730 self.clear_session_route_target_override(session_id).await;
731 self.clear_session_model_override(session_id).await;
732 self.clear_session_effort_override(session_id).await;
733 self.clear_session_service_tier_override(session_id).await;
734 }
735
736 pub async fn get_session_manual_overrides(&self, session_id: &str) -> SessionManualOverrides {
737 let (reasoning_effort, station_name, route_target, model, service_tier) = tokio::join!(
738 self.get_session_reasoning_effort_override(session_id),
739 self.get_session_station_override(session_id),
740 self.get_session_route_target_override(session_id),
741 self.get_session_model_override(session_id),
742 self.get_session_service_tier_override(session_id),
743 );
744
745 SessionManualOverrides {
746 reasoning_effort,
747 station_name,
748 route_target,
749 model,
750 service_tier,
751 }
752 }
753
754 pub async fn list_session_manual_overrides(&self) -> HashMap<String, SessionManualOverrides> {
755 let (reasoning_effort_map, station_map, route_target_map, model_map, service_tier_map) = tokio::join!(
756 self.list_session_reasoning_effort_overrides(),
757 self.list_session_station_overrides(),
758 self.list_session_route_target_overrides(),
759 self.list_session_model_overrides(),
760 self.list_session_service_tier_overrides(),
761 );
762
763 let mut merged = HashMap::<String, SessionManualOverrides>::new();
764 for (session_id, reasoning_effort) in reasoning_effort_map {
765 merged.entry(session_id).or_default().reasoning_effort = Some(reasoning_effort);
766 }
767 for (session_id, station_name) in station_map {
768 merged.entry(session_id).or_default().station_name = Some(station_name);
769 }
770 for (session_id, route_target) in route_target_map {
771 merged.entry(session_id).or_default().route_target = Some(route_target);
772 }
773 for (session_id, model) in model_map {
774 merged.entry(session_id).or_default().model = Some(model);
775 }
776 for (session_id, service_tier) in service_tier_map {
777 merged.entry(session_id).or_default().service_tier = Some(service_tier);
778 }
779 merged.retain(|_, overrides| !overrides.is_empty());
780 merged
781 }
782
783 pub async fn apply_session_profile_binding(
784 &self,
785 service_name: &str,
786 mgr: &ServiceConfigManager,
787 session_id: String,
788 profile_name: String,
789 now_ms: u64,
790 ) -> anyhow::Result<()> {
791 let profile = crate::config::resolve_service_profile(mgr, profile_name.as_str())?;
792 crate::config::validate_profile_station_compatibility(
793 service_name,
794 mgr,
795 profile_name.as_str(),
796 &profile,
797 )?;
798
799 self.set_session_binding(SessionBinding {
800 session_id: session_id.clone(),
801 profile_name: Some(profile_name),
802 station_name: profile.station.clone(),
803 model: profile.model.clone(),
804 reasoning_effort: profile.reasoning_effort.clone(),
805 service_tier: profile.service_tier.clone(),
806 continuity_mode: SessionContinuityMode::ManualProfile,
807 created_at_ms: now_ms,
808 updated_at_ms: now_ms,
809 last_seen_ms: now_ms,
810 })
811 .await;
812 self.clear_session_manual_overrides(session_id.as_str())
813 .await;
814 Ok(())
815 }
816
817 pub async fn touch_session_binding(&self, session_id: &str, now_ms: u64) {
818 let mut guard = self.session_bindings.write().await;
819 if let Some(entry) = guard.get_mut(session_id) {
820 entry.binding.last_seen_ms = now_ms;
821 }
822 }
823
824 pub async fn set_session_station_override(
825 &self,
826 session_id: String,
827 station_name: String,
828 now_ms: u64,
829 ) {
830 let mut guard = self.session_station_overrides.write().await;
831 guard.insert(
832 session_id,
833 SessionStationOverride {
834 station_name,
835 updated_at_ms: now_ms,
836 last_seen_ms: now_ms,
837 },
838 );
839 self.notify_state_changed();
840 }
841
842 pub async fn set_session_route_target_override(
843 &self,
844 session_id: String,
845 target: String,
846 now_ms: u64,
847 ) {
848 let mut guard = self.session_route_target_overrides.write().await;
849 guard.insert(
850 session_id,
851 SessionRouteTargetOverride {
852 target,
853 updated_at_ms: now_ms,
854 last_seen_ms: now_ms,
855 },
856 );
857 self.notify_state_changed();
858 }
859
860 pub async fn clear_session_station_override(&self, session_id: &str) {
861 let mut guard = self.session_station_overrides.write().await;
862 if guard.remove(session_id).is_some() {
863 self.notify_state_changed();
864 }
865 }
866
867 pub async fn clear_session_route_target_override(&self, session_id: &str) {
868 let mut guard = self.session_route_target_overrides.write().await;
869 if guard.remove(session_id).is_some() {
870 self.notify_state_changed();
871 }
872 }
873
874 pub async fn list_session_station_overrides(&self) -> HashMap<String, String> {
875 let guard = self.session_station_overrides.read().await;
876 guard
877 .iter()
878 .map(|(k, v)| (k.clone(), v.station_name.clone()))
879 .collect()
880 }
881
882 pub async fn list_session_route_target_overrides(&self) -> HashMap<String, String> {
883 let guard = self.session_route_target_overrides.read().await;
884 guard
885 .iter()
886 .map(|(k, v)| (k.clone(), v.target.clone()))
887 .collect()
888 }
889
890 pub async fn touch_session_station_override(&self, session_id: &str, now_ms: u64) {
891 let mut guard = self.session_station_overrides.write().await;
892 if let Some(v) = guard.get_mut(session_id) {
893 v.last_seen_ms = now_ms;
894 }
895 }
896
897 pub async fn touch_session_route_target_override(&self, session_id: &str, now_ms: u64) {
898 let mut guard = self.session_route_target_overrides.write().await;
899 if let Some(v) = guard.get_mut(session_id) {
900 v.last_seen_ms = now_ms;
901 }
902 }
903
904 pub async fn get_global_station_override(&self) -> Option<String> {
905 let guard = self.global_station_override.read().await;
906 guard.clone()
907 }
908
909 pub async fn get_global_route_target_override(&self) -> Option<String> {
910 let guard = self.global_route_target_override.read().await;
911 guard.clone()
912 }
913
914 pub async fn set_global_station_override(&self, station_name: String, _now_ms: u64) {
915 let mut guard = self.global_station_override.write().await;
916 *guard = Some(station_name);
917 self.notify_state_changed();
918 }
919
920 pub async fn set_global_route_target_override(&self, target: String, _now_ms: u64) {
921 let mut guard = self.global_route_target_override.write().await;
922 *guard = Some(target);
923 self.notify_state_changed();
924 }
925
926 pub async fn clear_global_station_override(&self) {
927 let mut guard = self.global_station_override.write().await;
928 if guard.take().is_some() {
929 self.notify_state_changed();
930 }
931 }
932
933 pub async fn clear_global_route_target_override(&self) {
934 let mut guard = self.global_route_target_override.write().await;
935 if guard.take().is_some() {
936 self.notify_state_changed();
937 }
938 }
939
940 pub async fn get_runtime_default_profile_override(&self, service_name: &str) -> Option<String> {
941 let guard = self.runtime_default_profiles.read().await;
942 guard
943 .get(service_name)
944 .map(|entry| entry.profile_name.clone())
945 }
946
947 pub async fn set_runtime_default_profile_override(
948 &self,
949 service_name: String,
950 profile_name: String,
951 now_ms: u64,
952 ) {
953 let mut guard = self.runtime_default_profiles.write().await;
954 guard.insert(
955 service_name,
956 RuntimeDefaultProfileOverride {
957 profile_name,
958 updated_at_ms: now_ms,
959 },
960 );
961 self.notify_state_changed();
962 }
963
964 pub async fn clear_runtime_default_profile_override(&self, service_name: &str) {
965 let mut guard = self.runtime_default_profiles.write().await;
966 if guard.remove(service_name).is_some() {
967 self.notify_state_changed();
968 }
969 }
970
971 pub async fn set_station_enabled_override(
972 &self,
973 service_name: &str,
974 station_name: String,
975 enabled: bool,
976 now_ms: u64,
977 ) {
978 let mut guard = self.station_meta_overrides.write().await;
979 let per_service = guard.entry(service_name.to_string()).or_default();
980 let entry = per_service.entry(station_name).or_default();
981 entry.enabled = Some(enabled);
982 entry.updated_at_ms = now_ms;
983 self.notify_state_changed();
984 }
985
986 pub async fn set_station_level_override(
987 &self,
988 service_name: &str,
989 station_name: String,
990 level: u8,
991 now_ms: u64,
992 ) {
993 let mut guard = self.station_meta_overrides.write().await;
994 let per_service = guard.entry(service_name.to_string()).or_default();
995 let entry = per_service.entry(station_name).or_default();
996 entry.level = Some(level.clamp(1, 10));
997 entry.updated_at_ms = now_ms;
998 self.notify_state_changed();
999 }
1000
1001 pub async fn set_station_runtime_state_override(
1002 &self,
1003 service_name: &str,
1004 station_name: String,
1005 state: RuntimeConfigState,
1006 now_ms: u64,
1007 ) {
1008 let mut guard = self.station_meta_overrides.write().await;
1009 let per_service = guard.entry(service_name.to_string()).or_default();
1010 let entry = per_service.entry(station_name).or_default();
1011 entry.state = Some(state);
1012 entry.updated_at_ms = now_ms;
1013 self.notify_state_changed();
1014 }
1015
1016 pub async fn clear_station_enabled_override(&self, service_name: &str, station_name: &str) {
1017 let mut guard = self.station_meta_overrides.write().await;
1018 let Some(per_service) = guard.get_mut(service_name) else {
1019 return;
1020 };
1021 let Some(entry) = per_service.get_mut(station_name) else {
1022 return;
1023 };
1024 if entry.enabled.take().is_none() {
1025 return;
1026 }
1027 if entry.enabled.is_none() && entry.level.is_none() && entry.state.is_none() {
1028 per_service.remove(station_name);
1029 }
1030 if per_service.is_empty() {
1031 guard.remove(service_name);
1032 }
1033 self.notify_state_changed();
1034 }
1035
1036 pub async fn clear_station_level_override(&self, service_name: &str, station_name: &str) {
1037 let mut guard = self.station_meta_overrides.write().await;
1038 let Some(per_service) = guard.get_mut(service_name) else {
1039 return;
1040 };
1041 let Some(entry) = per_service.get_mut(station_name) else {
1042 return;
1043 };
1044 if entry.level.take().is_none() {
1045 return;
1046 }
1047 if entry.enabled.is_none() && entry.level.is_none() && entry.state.is_none() {
1048 per_service.remove(station_name);
1049 }
1050 if per_service.is_empty() {
1051 guard.remove(service_name);
1052 }
1053 self.notify_state_changed();
1054 }
1055
1056 pub async fn clear_station_runtime_state_override(
1057 &self,
1058 service_name: &str,
1059 station_name: &str,
1060 ) {
1061 let mut guard = self.station_meta_overrides.write().await;
1062 let Some(per_service) = guard.get_mut(service_name) else {
1063 return;
1064 };
1065 let Some(entry) = per_service.get_mut(station_name) else {
1066 return;
1067 };
1068 if entry.state.take().is_none() {
1069 return;
1070 }
1071 if entry.enabled.is_none() && entry.level.is_none() && entry.state.is_none() {
1072 per_service.remove(station_name);
1073 }
1074 if per_service.is_empty() {
1075 guard.remove(service_name);
1076 }
1077 self.notify_state_changed();
1078 }
1079
1080 pub async fn get_station_meta_overrides(
1081 &self,
1082 service_name: &str,
1083 ) -> HashMap<String, (Option<bool>, Option<u8>)> {
1084 let guard = self.station_meta_overrides.read().await;
1085 guard
1086 .get(service_name)
1087 .map(|m| {
1088 m.iter()
1089 .map(|(k, v)| (k.clone(), (v.enabled, v.level)))
1090 .collect::<HashMap<_, _>>()
1091 })
1092 .unwrap_or_default()
1093 }
1094
1095 pub async fn get_station_runtime_state_overrides(
1096 &self,
1097 service_name: &str,
1098 ) -> HashMap<String, RuntimeConfigState> {
1099 let guard = self.station_meta_overrides.read().await;
1100 guard
1101 .get(service_name)
1102 .map(|m| {
1103 m.iter()
1104 .filter_map(|(k, v)| v.state.map(|state| (k.clone(), state)))
1105 .collect::<HashMap<_, _>>()
1106 })
1107 .unwrap_or_default()
1108 }
1109
1110 pub async fn set_provider_endpoint_enabled_override(
1111 &self,
1112 service_name: &str,
1113 endpoint_key: ProviderEndpointKey,
1114 enabled: bool,
1115 now_ms: u64,
1116 ) {
1117 let mut guard = self.provider_endpoint_meta_overrides.write().await;
1118 let per_service = guard.entry(service_name.to_string()).or_default();
1119 let entry = per_service.entry(endpoint_key).or_default();
1120 entry.enabled = Some(enabled);
1121 entry.updated_at_ms = now_ms;
1122 self.notify_state_changed();
1123 }
1124
1125 pub async fn clear_provider_endpoint_enabled_override(
1126 &self,
1127 service_name: &str,
1128 endpoint_key: &ProviderEndpointKey,
1129 ) {
1130 let mut guard = self.provider_endpoint_meta_overrides.write().await;
1131 let Some(per_service) = guard.get_mut(service_name) else {
1132 return;
1133 };
1134 let Some(entry) = per_service.get_mut(endpoint_key) else {
1135 return;
1136 };
1137 if entry.enabled.take().is_none() {
1138 return;
1139 }
1140 if entry.enabled.is_none() && entry.level.is_none() && entry.state.is_none() {
1141 per_service.remove(endpoint_key);
1142 }
1143 if per_service.is_empty() {
1144 guard.remove(service_name);
1145 }
1146 self.notify_state_changed();
1147 }
1148
1149 pub async fn set_provider_endpoint_runtime_state_override(
1150 &self,
1151 service_name: &str,
1152 endpoint_key: ProviderEndpointKey,
1153 state: RuntimeConfigState,
1154 now_ms: u64,
1155 ) {
1156 let mut guard = self.provider_endpoint_meta_overrides.write().await;
1157 let per_service = guard.entry(service_name.to_string()).or_default();
1158 let entry = per_service.entry(endpoint_key).or_default();
1159 entry.state = Some(state);
1160 entry.updated_at_ms = now_ms;
1161 self.notify_state_changed();
1162 }
1163
1164 pub async fn clear_provider_endpoint_runtime_state_override(
1165 &self,
1166 service_name: &str,
1167 endpoint_key: &ProviderEndpointKey,
1168 ) {
1169 let mut guard = self.provider_endpoint_meta_overrides.write().await;
1170 let Some(per_service) = guard.get_mut(service_name) else {
1171 return;
1172 };
1173 let Some(entry) = per_service.get_mut(endpoint_key) else {
1174 return;
1175 };
1176 if entry.state.take().is_none() {
1177 return;
1178 }
1179 if entry.enabled.is_none() && entry.level.is_none() && entry.state.is_none() {
1180 per_service.remove(endpoint_key);
1181 }
1182 if per_service.is_empty() {
1183 guard.remove(service_name);
1184 }
1185 self.notify_state_changed();
1186 }
1187
1188 pub async fn set_upstream_enabled_override(
1189 &self,
1190 service_name: &str,
1191 base_url: String,
1192 enabled: bool,
1193 now_ms: u64,
1194 ) {
1195 let mut guard = self.upstream_meta_overrides.write().await;
1196 let per_service = guard.entry(service_name.to_string()).or_default();
1197 let entry = per_service.entry(base_url).or_default();
1198 entry.enabled = Some(enabled);
1199 entry.updated_at_ms = now_ms;
1200 self.notify_state_changed();
1201 }
1202
1203 pub async fn clear_upstream_enabled_override(&self, service_name: &str, base_url: &str) {
1204 let mut guard = self.upstream_meta_overrides.write().await;
1205 let Some(per_service) = guard.get_mut(service_name) else {
1206 return;
1207 };
1208 let Some(entry) = per_service.get_mut(base_url) else {
1209 return;
1210 };
1211 if entry.enabled.take().is_none() {
1212 return;
1213 }
1214 if entry.enabled.is_none() && entry.level.is_none() && entry.state.is_none() {
1215 per_service.remove(base_url);
1216 }
1217 if per_service.is_empty() {
1218 guard.remove(service_name);
1219 }
1220 self.notify_state_changed();
1221 }
1222
1223 pub async fn set_upstream_runtime_state_override(
1224 &self,
1225 service_name: &str,
1226 base_url: String,
1227 state: RuntimeConfigState,
1228 now_ms: u64,
1229 ) {
1230 let mut guard = self.upstream_meta_overrides.write().await;
1231 let per_service = guard.entry(service_name.to_string()).or_default();
1232 let entry = per_service.entry(base_url).or_default();
1233 entry.state = Some(state);
1234 entry.updated_at_ms = now_ms;
1235 self.notify_state_changed();
1236 }
1237
1238 pub async fn clear_upstream_runtime_state_override(&self, service_name: &str, base_url: &str) {
1239 let mut guard = self.upstream_meta_overrides.write().await;
1240 let Some(per_service) = guard.get_mut(service_name) else {
1241 return;
1242 };
1243 let Some(entry) = per_service.get_mut(base_url) else {
1244 return;
1245 };
1246 if entry.state.take().is_none() {
1247 return;
1248 }
1249 if entry.enabled.is_none() && entry.level.is_none() && entry.state.is_none() {
1250 per_service.remove(base_url);
1251 }
1252 if per_service.is_empty() {
1253 guard.remove(service_name);
1254 }
1255 self.notify_state_changed();
1256 }
1257
1258 pub async fn get_upstream_meta_overrides(
1259 &self,
1260 service_name: &str,
1261 ) -> HashMap<String, (Option<bool>, Option<RuntimeConfigState>)> {
1262 let mut overrides = HashMap::new();
1263
1264 {
1265 let guard = self.upstream_meta_overrides.read().await;
1266 if let Some(per_service) = guard.get(service_name) {
1267 overrides.extend(
1268 per_service
1269 .iter()
1270 .map(|(k, v)| (k.clone(), (v.enabled, v.state))),
1271 );
1272 }
1273 }
1274
1275 {
1276 let guard = self.provider_endpoint_meta_overrides.read().await;
1277 if let Some(per_service) = guard.get(service_name) {
1278 overrides.extend(
1279 per_service
1280 .iter()
1281 .map(|(k, v)| (k.stable_key(), (v.enabled, v.state))),
1282 );
1283 }
1284 }
1285
1286 overrides
1287 }
1288
1289 pub async fn route_plan_runtime_state_for_provider_endpoints(
1290 &self,
1291 service_name: &str,
1292 ) -> RoutePlanRuntimeState {
1293 let mut runtime = RoutePlanRuntimeState::default();
1294 let now = std::time::Instant::now();
1295 {
1296 let mut guard = self.provider_endpoint_runtime_health.write().await;
1297 if let Some(per_service) = guard.get_mut(service_name) {
1298 let mut affinity: Option<(ProviderEndpointKey, u64)> = None;
1299 for (endpoint_key, health) in per_service.iter_mut() {
1300 if health.cooldown_until.is_some_and(|until| now >= until) {
1301 health.failure_count = 0;
1302 health.cooldown_until = None;
1303 }
1304 let cooldown_active = health.cooldown_until.is_some_and(|until| now < until);
1305 let cooldown_remaining_secs = health.cooldown_until.and_then(|until| {
1306 if now < until {
1307 Some(until.duration_since(now).as_secs().max(1))
1308 } else {
1309 None
1310 }
1311 });
1312 runtime.set_provider_endpoint(
1313 endpoint_key.clone(),
1314 RoutePlanUpstreamRuntimeState {
1315 runtime_disabled: false,
1316 failure_count: health.failure_count,
1317 cooldown_active,
1318 cooldown_remaining_secs,
1319 usage_exhausted: health.usage_exhausted,
1320 missing_auth: false,
1321 concurrency_saturated: false,
1322 concurrency_active: None,
1323 concurrency_limit: None,
1324 },
1325 );
1326 if let Some(last_good_at_ms) = health.last_good_at_ms
1327 && affinity
1328 .as_ref()
1329 .is_none_or(|(_, current)| last_good_at_ms > *current)
1330 {
1331 affinity = Some((endpoint_key.clone(), last_good_at_ms));
1332 }
1333 }
1334 if let Some((endpoint_key, _)) = affinity {
1335 runtime.set_affinity_provider_endpoint(Some(endpoint_key));
1336 }
1337 }
1338 }
1339
1340 {
1341 let guard = self.provider_endpoint_meta_overrides.read().await;
1342 if let Some(per_service) = guard.get(service_name) {
1343 for (endpoint_key, meta) in per_service {
1344 let mut upstream_state = runtime.provider_endpoint(endpoint_key);
1345 if meta.enabled == Some(false)
1346 || meta
1347 .state
1348 .is_some_and(|state| state != RuntimeConfigState::Normal)
1349 {
1350 upstream_state.runtime_disabled = true;
1351 }
1352 runtime.set_provider_endpoint(endpoint_key.clone(), upstream_state);
1353 }
1354 }
1355 }
1356
1357 runtime
1358 }
1359
1360 pub async fn record_provider_endpoint_attempt_success(
1361 &self,
1362 service_name: &str,
1363 endpoint_key: ProviderEndpointKey,
1364 now_ms: u64,
1365 ) {
1366 let mut guard = self.provider_endpoint_runtime_health.write().await;
1367 let entry = guard
1368 .entry(service_name.to_string())
1369 .or_default()
1370 .entry(endpoint_key)
1371 .or_default();
1372 entry.failure_count = 0;
1373 entry.cooldown_until = None;
1374 entry.penalty_streak = 0;
1375 entry.last_good_at_ms = Some(now_ms);
1376 }
1377
1378 pub async fn record_provider_endpoint_attempt_failure(
1379 &self,
1380 service_name: &str,
1381 endpoint_key: ProviderEndpointKey,
1382 failure_threshold_cooldown_secs: u64,
1383 cooldown_backoff: CooldownBackoff,
1384 ) {
1385 let mut guard = self.provider_endpoint_runtime_health.write().await;
1386 let entry = guard
1387 .entry(service_name.to_string())
1388 .or_default()
1389 .entry(endpoint_key)
1390 .or_default();
1391
1392 entry.failure_count = entry.failure_count.saturating_add(1);
1393 if entry.failure_count >= FAILURE_THRESHOLD {
1394 let base_secs = if failure_threshold_cooldown_secs == 0 {
1395 COOLDOWN_SECS
1396 } else {
1397 failure_threshold_cooldown_secs
1398 };
1399 let effective_secs =
1400 cooldown_backoff.effective_cooldown_secs(base_secs, entry.penalty_streak);
1401 let now = std::time::Instant::now();
1402 let new_until = now + std::time::Duration::from_secs(effective_secs);
1403 if entry
1404 .cooldown_until
1405 .is_none_or(|existing| new_until > existing)
1406 {
1407 entry.cooldown_until = Some(new_until);
1408 }
1409 entry.penalty_streak = entry.penalty_streak.saturating_add(1);
1410 entry.last_good_at_ms = None;
1411 }
1412 }
1413
1414 pub async fn penalize_provider_endpoint_attempt(
1415 &self,
1416 service_name: &str,
1417 endpoint_key: ProviderEndpointKey,
1418 cooldown_secs: u64,
1419 cooldown_backoff: CooldownBackoff,
1420 ) {
1421 let mut guard = self.provider_endpoint_runtime_health.write().await;
1422 let entry = guard
1423 .entry(service_name.to_string())
1424 .or_default()
1425 .entry(endpoint_key)
1426 .or_default();
1427 let effective_secs =
1428 cooldown_backoff.effective_cooldown_secs(cooldown_secs, entry.penalty_streak);
1429 entry.failure_count = FAILURE_THRESHOLD;
1430 entry.cooldown_until =
1431 Some(std::time::Instant::now() + std::time::Duration::from_secs(effective_secs));
1432 entry.penalty_streak = entry.penalty_streak.saturating_add(1);
1433 entry.last_good_at_ms = None;
1434 }
1435
1436 pub async fn set_provider_endpoint_usage_exhausted(
1437 &self,
1438 service_name: &str,
1439 endpoint_key: ProviderEndpointKey,
1440 exhausted: bool,
1441 ) {
1442 let mut guard = self.provider_endpoint_runtime_health.write().await;
1443 let entry = guard
1444 .entry(service_name.to_string())
1445 .or_default()
1446 .entry(endpoint_key)
1447 .or_default();
1448 if entry.usage_exhausted == exhausted {
1449 return;
1450 }
1451 entry.usage_exhausted = exhausted;
1452 }
1453
1454 pub async fn prune_runtime_observability_for_service(
1455 &self,
1456 service_name: &str,
1457 mgr: &ServiceConfigManager,
1458 ) {
1459 let mut changed = false;
1460 let active_stations = mgr.stations().keys().cloned().collect::<HashSet<_>>();
1461 let active_upstreams = mgr
1462 .stations()
1463 .iter()
1464 .map(|(station_name, service)| {
1465 (
1466 station_name.clone(),
1467 service
1468 .upstreams
1469 .iter()
1470 .map(|upstream| upstream.base_url.clone())
1471 .collect::<HashSet<_>>(),
1472 )
1473 })
1474 .collect::<HashMap<_, _>>();
1475 let active_base_urls = active_upstreams
1476 .values()
1477 .flat_map(|upstreams| upstreams.iter().cloned())
1478 .collect::<HashSet<_>>();
1479 let active_provider_endpoint_keys = mgr
1480 .stations()
1481 .iter()
1482 .flat_map(|(station_name, service)| {
1483 service.upstreams.iter().enumerate().map(|(idx, upstream)| {
1484 Self::active_provider_endpoint_key_for_upstream(
1485 service_name,
1486 station_name.as_str(),
1487 idx,
1488 upstream,
1489 )
1490 })
1491 })
1492 .collect::<HashSet<_>>();
1493 let mut active_provider_ids = HashSet::from(["-".to_string()]);
1494 for service in mgr.stations().values() {
1495 for upstream in &service.upstreams {
1496 if let Some(provider_id) = upstream.tags.get("provider_id") {
1497 active_provider_ids.insert(provider_id.clone());
1498 }
1499 }
1500 }
1501
1502 let layout = service_layout_signature(mgr);
1503 let balance_prune_stations = {
1504 let mut signatures = self.service_layout_signatures.write().await;
1505 let changed = signatures.get(service_name).map_or_else(
1506 || {
1507 if active_stations.len() == 1 && active_stations.contains("routing") {
1508 Some(HashSet::from(["routing".to_string()]))
1509 } else {
1510 None
1511 }
1512 },
1513 |previous| Some(changed_service_layout_stations(previous, &layout)),
1514 );
1515 signatures.insert(service_name.to_string(), layout);
1516 changed
1517 };
1518
1519 match balance_prune_stations {
1520 Some(changed_layout_stations) if !changed_layout_stations.is_empty() => {
1521 let mut provider_balances = self.provider_balances.write().await;
1522 if let Some(per_service) = provider_balances.get_mut(service_name) {
1523 let before = per_service.len();
1524 per_service
1525 .retain(|station_name, _| !changed_layout_stations.contains(station_name));
1526 changed |= per_service.len() != before;
1527 if per_service.is_empty() {
1528 provider_balances.remove(service_name);
1529 }
1530 }
1531 let mut provider_balance_summaries = self.provider_balance_summaries.write().await;
1532 if let Some(per_service) = provider_balance_summaries.get_mut(service_name) {
1533 let before = per_service.len();
1534 per_service
1535 .retain(|station_name, _| !changed_layout_stations.contains(station_name));
1536 changed |= per_service.len() != before;
1537 if per_service.is_empty() {
1538 provider_balance_summaries.remove(service_name);
1539 }
1540 }
1541 }
1542 None => {
1543 let mut provider_balances = self.provider_balances.write().await;
1544 changed |= provider_balances.remove(service_name).is_some();
1545 let mut provider_balance_summaries = self.provider_balance_summaries.write().await;
1546 changed |= provider_balance_summaries.remove(service_name).is_some();
1547 }
1548 Some(_) => {}
1549 }
1550
1551 {
1552 let mut guard = self.station_meta_overrides.write().await;
1553 if let Some(per_service) = guard.get_mut(service_name) {
1554 let before = per_service.len();
1555 per_service.retain(|station_name, _| active_stations.contains(station_name));
1556 changed |= per_service.len() != before;
1557 if per_service.is_empty() {
1558 guard.remove(service_name);
1559 }
1560 }
1561 }
1562
1563 {
1564 let mut guard = self.upstream_meta_overrides.write().await;
1565 if let Some(per_service) = guard.get_mut(service_name) {
1566 let before = per_service.len();
1567 per_service.retain(|base_url, _| active_base_urls.contains(base_url));
1568 changed |= per_service.len() != before;
1569 if per_service.is_empty() {
1570 guard.remove(service_name);
1571 }
1572 }
1573 }
1574
1575 {
1576 let mut guard = self.provider_endpoint_meta_overrides.write().await;
1577 if let Some(per_service) = guard.get_mut(service_name) {
1578 let before = per_service.len();
1579 per_service
1580 .retain(|endpoint_key, _| active_provider_endpoint_keys.contains(endpoint_key));
1581 changed |= per_service.len() != before;
1582 if per_service.is_empty() {
1583 guard.remove(service_name);
1584 }
1585 }
1586 }
1587
1588 {
1589 let mut guard = self.provider_endpoint_runtime_health.write().await;
1590 if let Some(per_service) = guard.get_mut(service_name) {
1591 per_service
1592 .retain(|endpoint_key, _| active_provider_endpoint_keys.contains(endpoint_key));
1593 if per_service.is_empty() {
1594 guard.remove(service_name);
1595 }
1596 }
1597 }
1598
1599 {
1600 let mut guard = self.station_health.write().await;
1601 if let Some(per_service) = guard.get_mut(service_name) {
1602 let before_stations = per_service.len();
1603 let before_upstreams = per_service
1604 .values()
1605 .map(|station_health| station_health.upstreams.len())
1606 .sum::<usize>();
1607 per_service.retain(|station_name, station_health| {
1608 if !active_stations.contains(station_name) {
1609 return false;
1610 }
1611 if let Some(allowed_upstreams) = active_upstreams.get(station_name) {
1612 station_health
1613 .upstreams
1614 .retain(|upstream| allowed_upstreams.contains(&upstream.base_url));
1615 }
1616 !station_health.upstreams.is_empty()
1617 });
1618 let after_upstreams = per_service
1619 .values()
1620 .map(|station_health| station_health.upstreams.len())
1621 .sum::<usize>();
1622 changed |=
1623 per_service.len() != before_stations || after_upstreams != before_upstreams;
1624 if per_service.is_empty() {
1625 guard.remove(service_name);
1626 }
1627 }
1628 }
1629
1630 {
1631 let mut guard = self.passive_station_health.write().await;
1632 if let Some(per_service) = guard.get_mut(service_name) {
1633 let before_stations = per_service.len();
1634 let before_upstreams = per_service.values().map(HashMap::len).sum::<usize>();
1635 per_service.retain(|station_name, station_health| {
1636 if !active_stations.contains(station_name) {
1637 return false;
1638 }
1639 if let Some(allowed_upstreams) = active_upstreams.get(station_name) {
1640 station_health.retain(|base_url, _| allowed_upstreams.contains(base_url));
1641 }
1642 !station_health.is_empty()
1643 });
1644 let after_upstreams = per_service.values().map(HashMap::len).sum::<usize>();
1645 changed |=
1646 per_service.len() != before_stations || after_upstreams != before_upstreams;
1647 if per_service.is_empty() {
1648 guard.remove(service_name);
1649 }
1650 }
1651 }
1652
1653 {
1654 let mut guard = self.station_health_checks.write().await;
1655 if let Some(per_service) = guard.get_mut(service_name) {
1656 let before = per_service.len();
1657 per_service.retain(|station_name, _| active_stations.contains(station_name));
1658 changed |= per_service.len() != before;
1659 if per_service.is_empty() {
1660 guard.remove(service_name);
1661 }
1662 }
1663 }
1664
1665 {
1666 let mut guard = self.usage_rollups.write().await;
1667 if let Some(rollup) = guard.get_mut(service_name) {
1668 let before_by_config = rollup.by_config.len();
1669 rollup
1670 .by_config
1671 .retain(|station_name, _| active_stations.contains(station_name));
1672 changed |= rollup.by_config.len() != before_by_config;
1673 let before_by_config_day = rollup.by_config_day.len();
1674 rollup.by_config_day.retain(|station_name, _day_map| {
1675 if !active_stations.contains(station_name) {
1676 return false;
1677 }
1678 true
1679 });
1680 changed |= rollup.by_config_day.len() != before_by_config_day;
1681 let before_by_provider = rollup.by_provider.len();
1682 rollup
1683 .by_provider
1684 .retain(|provider_id, _| active_provider_ids.contains(provider_id));
1685 changed |= rollup.by_provider.len() != before_by_provider;
1686 let before_by_provider_day = rollup.by_provider_day.len();
1687 rollup.by_provider_day.retain(|provider_id, _day_map| {
1688 if !active_provider_ids.contains(provider_id) {
1689 return false;
1690 }
1691 true
1692 });
1693 changed |= rollup.by_provider_day.len() != before_by_provider_day;
1694 }
1695 }
1696 if changed {
1697 self.notify_state_changed();
1698 }
1699 }
1700
1701 fn active_provider_endpoint_key_for_upstream(
1702 service_name: &str,
1703 station_name: &str,
1704 upstream_index: usize,
1705 upstream: &crate::config::UpstreamConfig,
1706 ) -> ProviderEndpointKey {
1707 let provider_id = upstream
1708 .tags
1709 .get("provider_id")
1710 .cloned()
1711 .unwrap_or_else(|| format!("{station_name}#{upstream_index}"));
1712 let endpoint_id = upstream
1713 .tags
1714 .get("endpoint_id")
1715 .cloned()
1716 .unwrap_or_else(|| upstream_index.to_string());
1717
1718 ProviderEndpointKey::new(service_name, provider_id, endpoint_id)
1719 }
1720
1721 pub async fn record_station_health(
1722 &self,
1723 service_name: &str,
1724 station_name: String,
1725 health: StationHealth,
1726 ) {
1727 let mut guard = self.station_health.write().await;
1728 let per_service = guard.entry(service_name.to_string()).or_default();
1729 per_service.insert(station_name, health);
1730 self.notify_state_changed();
1731 }
1732
1733 pub async fn get_station_health(&self, service_name: &str) -> HashMap<String, StationHealth> {
1734 let active = {
1735 let guard = self.station_health.read().await;
1736 guard.get(service_name).cloned().unwrap_or_default()
1737 };
1738 let passive = {
1739 let guard = self.passive_station_health.read().await;
1740 guard.get(service_name).cloned().unwrap_or_default()
1741 };
1742 merge_station_health(active, passive)
1743 }
1744
1745 pub async fn record_provider_balance_snapshot(
1746 &self,
1747 service_name: &str,
1748 mut snapshot: ProviderBalanceSnapshot,
1749 ) {
1750 let (Some(station_name), Some(upstream_index)) =
1751 (snapshot.station_name.clone(), snapshot.upstream_index)
1752 else {
1753 return;
1754 };
1755 let now_ms = unix_now_ms();
1756 snapshot.refresh_status(now_ms);
1757
1758 let station_summary = {
1759 let mut guard = self.provider_balances.write().await;
1760 let station_balances = guard
1761 .entry(service_name.to_string())
1762 .or_default()
1763 .entry(station_name.clone())
1764 .or_default();
1765 if !snapshot.has_amount_data()
1766 && let Some(previous) = station_balances
1767 .get(&upstream_index)
1768 .and_then(|providers| providers.get(&snapshot.provider_id))
1769 {
1770 snapshot.carry_forward_amount_data_from(previous);
1771 snapshot.refresh_status(now_ms);
1772 }
1773 station_balances
1774 .entry(upstream_index)
1775 .or_default()
1776 .insert(snapshot.provider_id.clone(), snapshot);
1777 StationRoutingBalanceSummary::from_snapshot_iter_at(
1778 station_balances
1779 .values()
1780 .flat_map(|providers| providers.values()),
1781 now_ms,
1782 )
1783 };
1784
1785 let mut summaries = self.provider_balance_summaries.write().await;
1786 summaries
1787 .entry(service_name.to_string())
1788 .or_default()
1789 .insert(station_name, station_summary);
1790 self.notify_state_changed();
1791 }
1792
1793 pub async fn get_provider_balance_view(
1794 &self,
1795 service_name: &str,
1796 ) -> HashMap<String, Vec<ProviderBalanceSnapshot>> {
1797 let now_ms = unix_now_ms();
1798 let guard = self.provider_balances.read().await;
1799 let Some(per_service) = guard.get(service_name) else {
1800 return HashMap::new();
1801 };
1802
1803 per_service
1804 .iter()
1805 .map(|(station_name, upstreams)| {
1806 let mut snapshots = upstreams
1807 .values()
1808 .flat_map(|providers| providers.values().cloned())
1809 .collect::<Vec<_>>();
1810 for snapshot in &mut snapshots {
1811 snapshot.refresh_status(now_ms);
1812 }
1813 snapshots.sort_by(|a, b| {
1814 a.upstream_index
1815 .cmp(&b.upstream_index)
1816 .then_with(|| a.provider_id.cmp(&b.provider_id))
1817 });
1818 (station_name.clone(), snapshots)
1819 })
1820 .collect()
1821 }
1822
1823 pub async fn get_provider_balance_summary_view(
1824 &self,
1825 service_name: &str,
1826 ) -> HashMap<String, StationRoutingBalanceSummary> {
1827 let guard = self.provider_balance_summaries.read().await;
1828 let Some(per_service) = guard.get(service_name) else {
1829 return HashMap::new();
1830 };
1831
1832 per_service.clone()
1833 }
1834
1835 pub async fn record_passive_upstream_success(
1836 &self,
1837 service_name: &str,
1838 station_name: &str,
1839 base_url: &str,
1840 status_code: Option<u16>,
1841 now_ms: u64,
1842 ) {
1843 let mut guard = self.passive_station_health.write().await;
1844 let entry = guard
1845 .entry(service_name.to_string())
1846 .or_default()
1847 .entry(station_name.to_string())
1848 .or_default()
1849 .entry(base_url.to_string())
1850 .or_default();
1851 entry.record_success(now_ms, status_code);
1852 self.notify_state_changed();
1853 }
1854
1855 pub async fn record_passive_upstream_failure(&self, params: PassiveUpstreamFailureRecord) {
1856 let PassiveUpstreamFailureRecord {
1857 service_name,
1858 station_name,
1859 base_url,
1860 status_code,
1861 error_class,
1862 error,
1863 now_ms,
1864 } = params;
1865
1866 let mut guard = self.passive_station_health.write().await;
1867 let entry = guard
1868 .entry(service_name)
1869 .or_default()
1870 .entry(station_name)
1871 .or_default()
1872 .entry(base_url)
1873 .or_default();
1874 entry.record_failure(now_ms, status_code, error_class, error);
1875 self.notify_state_changed();
1876 }
1877
1878 pub async fn get_lb_view(&self) -> HashMap<String, LbConfigView> {
1879 let Some(lb_states) = self.lb_states.as_ref() else {
1880 return HashMap::new();
1881 };
1882 let mut map = match lb_states.lock() {
1883 Ok(m) => m,
1884 Err(e) => e.into_inner(),
1885 };
1886
1887 let now = std::time::Instant::now();
1888 let mut out = HashMap::new();
1889 for (cfg_name, st) in map.iter_mut() {
1890 let len = st
1891 .failure_counts
1892 .len()
1893 .max(st.cooldown_until.len())
1894 .max(st.usage_exhausted.len());
1895 if len == 0 {
1896 continue;
1897 }
1898
1899 if st.failure_counts.len() != len {
1901 st.failure_counts.resize(len, 0);
1902 }
1903 if st.cooldown_until.len() != len {
1904 st.cooldown_until.resize(len, None);
1905 }
1906 if st.usage_exhausted.len() != len {
1907 st.usage_exhausted.resize(len, false);
1908 }
1909
1910 let mut upstreams = Vec::with_capacity(len);
1911 for idx in 0..len {
1912 let failure_count = st.failure_counts.get(idx).copied().unwrap_or(0);
1913 let cooldown_remaining_secs = st
1914 .cooldown_until
1915 .get(idx)
1916 .and_then(|v| *v)
1917 .map(|until| until.saturating_duration_since(now).as_secs())
1918 .filter(|&s| s > 0);
1919 let usage_exhausted = st.usage_exhausted.get(idx).copied().unwrap_or(false);
1920 upstreams.push(LbUpstreamView {
1921 failure_count,
1922 cooldown_remaining_secs,
1923 usage_exhausted,
1924 });
1925 }
1926
1927 out.insert(
1928 cfg_name.clone(),
1929 LbConfigView {
1930 last_good_index: st.last_good_index,
1931 upstreams,
1932 },
1933 );
1934 }
1935 out
1936 }
1937
1938 pub async fn list_health_checks(
1939 &self,
1940 service_name: &str,
1941 ) -> HashMap<String, HealthCheckStatus> {
1942 let guard = self.station_health_checks.read().await;
1943 guard.get(service_name).cloned().unwrap_or_default()
1944 }
1945
1946 pub async fn try_begin_station_health_check(
1947 &self,
1948 service_name: &str,
1949 station_name: &str,
1950 total: usize,
1951 now_ms: u64,
1952 ) -> bool {
1953 let mut guard = self.station_health_checks.write().await;
1954 let per_service = guard.entry(service_name.to_string()).or_default();
1955 if let Some(existing) = per_service.get(station_name)
1956 && !existing.done
1957 {
1958 return false;
1959 }
1960 per_service.insert(
1961 station_name.to_string(),
1962 HealthCheckStatus {
1963 started_at_ms: now_ms,
1964 updated_at_ms: now_ms,
1965 total: total.min(u32::MAX as usize) as u32,
1966 completed: 0,
1967 ok: 0,
1968 err: 0,
1969 cancel_requested: false,
1970 canceled: false,
1971 done: false,
1972 last_error: None,
1973 },
1974 );
1975 self.notify_state_changed();
1976 true
1977 }
1978
1979 pub async fn request_cancel_station_health_check(
1980 &self,
1981 service_name: &str,
1982 station_name: &str,
1983 now_ms: u64,
1984 ) -> bool {
1985 let mut guard = self.station_health_checks.write().await;
1986 let Some(per_service) = guard.get_mut(service_name) else {
1987 return false;
1988 };
1989 let Some(st) = per_service.get_mut(station_name) else {
1990 return false;
1991 };
1992 if st.done {
1993 return false;
1994 }
1995 st.cancel_requested = true;
1996 st.updated_at_ms = now_ms;
1997 self.notify_state_changed();
1998 true
1999 }
2000
2001 pub async fn is_station_health_check_cancel_requested(
2002 &self,
2003 service_name: &str,
2004 station_name: &str,
2005 ) -> bool {
2006 let guard = self.station_health_checks.read().await;
2007 guard
2008 .get(service_name)
2009 .and_then(|m| m.get(station_name))
2010 .is_some_and(|s| s.cancel_requested && !s.done)
2011 }
2012
2013 pub async fn record_station_health_check_result(
2014 &self,
2015 service_name: &str,
2016 station_name: &str,
2017 now_ms: u64,
2018 upstream: UpstreamHealth,
2019 ) {
2020 {
2021 let mut guard = self.station_health.write().await;
2022 let per_service = guard.entry(service_name.to_string()).or_default();
2023 let entry = per_service
2024 .entry(station_name.to_string())
2025 .or_insert_with(|| StationHealth {
2026 checked_at_ms: now_ms,
2027 upstreams: Vec::new(),
2028 });
2029 entry.checked_at_ms = entry.checked_at_ms.max(now_ms);
2030 entry.upstreams.push(upstream.clone());
2031 if entry.upstreams.len() > Self::MAX_HEALTH_RECORDS_PER_STATION {
2032 let extra = entry
2033 .upstreams
2034 .len()
2035 .saturating_sub(Self::MAX_HEALTH_RECORDS_PER_STATION);
2036 if extra > 0 {
2037 entry.upstreams.drain(0..extra);
2038 }
2039 }
2040 }
2041
2042 let mut guard = self.station_health_checks.write().await;
2043 let per_service = guard.entry(service_name.to_string()).or_default();
2044 let st = per_service.entry(station_name.to_string()).or_default();
2045 st.updated_at_ms = now_ms;
2046 st.completed = st.completed.saturating_add(1);
2047 match upstream.ok {
2048 Some(true) => st.ok = st.ok.saturating_add(1),
2049 Some(false) => {
2050 st.err = st.err.saturating_add(1);
2051 if st.last_error.is_none() {
2052 st.last_error = upstream.error.clone();
2053 }
2054 }
2055 None => {}
2056 }
2057 self.notify_state_changed();
2058 }
2059
2060 pub async fn finish_station_health_check(
2061 &self,
2062 service_name: &str,
2063 station_name: &str,
2064 now_ms: u64,
2065 canceled: bool,
2066 ) {
2067 let mut guard = self.station_health_checks.write().await;
2068 let per_service = guard.entry(service_name.to_string()).or_default();
2069 let st = per_service.entry(station_name.to_string()).or_default();
2070 st.updated_at_ms = now_ms;
2071 st.canceled = canceled;
2072 st.done = true;
2073 self.notify_state_changed();
2074 }
2075
2076 pub async fn get_usage_rollup_view(
2077 &self,
2078 service_name: &str,
2079 top_n: usize,
2080 days: usize,
2081 ) -> UsageRollupView {
2082 let guard = self.usage_rollups.read().await;
2083 let Some(rollup) = guard.get(service_name) else {
2084 return UsageRollupView::default();
2085 };
2086
2087 fn now_day() -> i32 {
2088 std::time::SystemTime::now()
2089 .duration_since(std::time::UNIX_EPOCH)
2090 .map(|d| (d.as_millis() / 86_400_000) as i32)
2091 .unwrap_or(0)
2092 }
2093
2094 fn sorted_day_series(map: &HashMap<i32, UsageBucket>) -> Vec<(i32, UsageBucket)> {
2095 let mut out = map.iter().map(|(k, v)| (*k, v.clone())).collect::<Vec<_>>();
2096 out.sort_by_key(|(k, _)| *k);
2097 out
2098 }
2099
2100 fn filled_day_series(
2101 map: &HashMap<i32, UsageBucket>,
2102 start_day: i32,
2103 end_day: i32,
2104 ) -> Vec<(i32, UsageBucket)> {
2105 if start_day > end_day {
2106 return Vec::new();
2107 }
2108 (start_day..=end_day)
2109 .map(|day| (day, map.get(&day).cloned().unwrap_or_default()))
2110 .collect()
2111 }
2112
2113 fn sum_series(series: &[(i32, UsageBucket)]) -> UsageBucket {
2114 let mut out = UsageBucket::default();
2115 for (_, bucket) in series {
2116 out.add_assign(bucket);
2117 }
2118 out
2119 }
2120
2121 fn aggregate_entity_window(
2122 source: &HashMap<String, HashMap<i32, UsageBucket>>,
2123 start_day: Option<i32>,
2124 end_day: Option<i32>,
2125 top_n: usize,
2126 ) -> Vec<(String, UsageBucket)> {
2127 let mut out = Vec::new();
2128 for (name, days) in source {
2129 let mut bucket = UsageBucket::default();
2130 for (day, value) in days {
2131 let include = match (start_day, end_day) {
2132 (Some(start), Some(end)) => *day >= start && *day <= end,
2133 _ => true,
2134 };
2135 if include {
2136 bucket.add_assign(value);
2137 }
2138 }
2139 if bucket.requests_total > 0 {
2140 out.push((name.clone(), bucket));
2141 }
2142 }
2143 out.sort_by(|(left_name, left), (right_name, right)| {
2144 right
2145 .usage
2146 .total_tokens
2147 .cmp(&left.usage.total_tokens)
2148 .then_with(|| right.requests_total.cmp(&left.requests_total))
2149 .then_with(|| left_name.cmp(right_name))
2150 });
2151 out.truncate(top_n);
2152 out
2153 }
2154
2155 let all_loaded = days == 0;
2156 let loaded_first_day = rollup.by_day.keys().min().copied();
2157 let loaded_last_day = rollup.by_day.keys().max().copied();
2158 let loaded_days_with_data = rollup
2159 .by_day
2160 .values()
2161 .filter(|bucket| bucket.requests_total > 0)
2162 .count();
2163
2164 let (start_day, end_day) = if all_loaded {
2165 (loaded_first_day, loaded_last_day)
2166 } else {
2167 let end = now_day();
2168 let offset = i32::try_from(days.saturating_sub(1)).unwrap_or(i32::MAX);
2169 (Some(end.saturating_sub(offset)), Some(end))
2170 };
2171
2172 let by_day = match (all_loaded, start_day, end_day) {
2173 (true, _, _) => sorted_day_series(&rollup.by_day),
2174 (false, Some(start), Some(end)) => filled_day_series(&rollup.by_day, start, end),
2175 _ => Vec::new(),
2176 };
2177 let window = if all_loaded {
2178 rollup.loaded.clone()
2179 } else {
2180 sum_series(&by_day)
2181 };
2182
2183 let mut by_config =
2184 aggregate_entity_window(&rollup.by_config_day, start_day, end_day, top_n);
2185 if all_loaded && by_config.is_empty() {
2186 by_config = rollup
2187 .by_config
2188 .iter()
2189 .map(|(k, v)| (k.clone(), v.clone()))
2190 .collect::<Vec<_>>();
2191 by_config.sort_by(|(left_name, left), (right_name, right)| {
2192 right
2193 .usage
2194 .total_tokens
2195 .cmp(&left.usage.total_tokens)
2196 .then_with(|| right.requests_total.cmp(&left.requests_total))
2197 .then_with(|| left_name.cmp(right_name))
2198 });
2199 by_config.truncate(top_n);
2200 }
2201
2202 let mut by_provider =
2203 aggregate_entity_window(&rollup.by_provider_day, start_day, end_day, top_n);
2204 if all_loaded && by_provider.is_empty() {
2205 by_provider = rollup
2206 .by_provider
2207 .iter()
2208 .map(|(k, v)| (k.clone(), v.clone()))
2209 .collect::<Vec<_>>();
2210 by_provider.sort_by(|(left_name, left), (right_name, right)| {
2211 right
2212 .usage
2213 .total_tokens
2214 .cmp(&left.usage.total_tokens)
2215 .then_with(|| right.requests_total.cmp(&left.requests_total))
2216 .then_with(|| left_name.cmp(right_name))
2217 });
2218 by_provider.truncate(top_n);
2219 }
2220
2221 let mut by_config_day = HashMap::new();
2222 for (name, _) in &by_config {
2223 let series = rollup
2224 .by_config_day
2225 .get(name)
2226 .map(|m| match (all_loaded, start_day, end_day) {
2227 (true, _, _) => sorted_day_series(m),
2228 (false, Some(start), Some(end)) => filled_day_series(m, start, end),
2229 _ => Vec::new(),
2230 })
2231 .unwrap_or_default();
2232 by_config_day.insert(name.clone(), series);
2233 }
2234
2235 let mut by_provider_day = HashMap::new();
2236 for (name, _) in &by_provider {
2237 let series = rollup
2238 .by_provider_day
2239 .get(name)
2240 .map(|m| match (all_loaded, start_day, end_day) {
2241 (true, _, _) => sorted_day_series(m),
2242 (false, Some(start), Some(end)) => filled_day_series(m, start, end),
2243 _ => Vec::new(),
2244 })
2245 .unwrap_or_default();
2246 by_provider_day.insert(name.clone(), series);
2247 }
2248
2249 let window_days_with_data = by_day
2250 .iter()
2251 .filter(|(_, bucket)| bucket.requests_total > 0)
2252 .count();
2253 let coverage = UsageRollupCoverage {
2254 requested_days: days,
2255 all_loaded,
2256 loaded_first_day,
2257 loaded_last_day,
2258 loaded_days_with_data,
2259 loaded_requests: rollup.loaded.requests_total,
2260 window_first_day: start_day,
2261 window_last_day: end_day,
2262 window_days_with_data,
2263 window_requests: window.requests_total,
2264 window_exceeds_loaded_start: matches!(
2265 (all_loaded, start_day, loaded_first_day),
2266 (false, Some(start), Some(first)) if start < first
2267 ),
2268 };
2269
2270 UsageRollupView {
2271 loaded: rollup.loaded.clone(),
2272 window,
2273 coverage,
2274 by_day,
2275 by_config,
2276 by_config_day,
2277 by_provider,
2278 by_provider_day,
2279 }
2280 }
2281
2282 pub async fn replay_usage_from_requests_log(
2283 &self,
2284 service_name: &str,
2285 log_path: PathBuf,
2286 base_url_to_provider_id: HashMap<String, String>,
2287 ) -> usize {
2288 let enabled = std::env::var("CODEX_HELPER_USAGE_REPLAY_ON_STARTUP")
2289 .ok()
2290 .map(|v| {
2291 matches!(
2292 v.trim().to_ascii_lowercase().as_str(),
2293 "1" | "true" | "yes" | "y" | "on"
2294 )
2295 })
2296 .unwrap_or(true);
2297 if !enabled {
2298 return 0;
2299 }
2300
2301 let already_has_data = {
2302 let guard = self.usage_rollups.read().await;
2303 guard
2304 .get(service_name)
2305 .is_some_and(|r| r.loaded.requests_total > 0)
2306 };
2307 if already_has_data {
2308 return 0;
2309 }
2310
2311 if !log_path.exists() {
2312 return 0;
2313 }
2314
2315 let max_bytes = std::env::var("CODEX_HELPER_USAGE_REPLAY_MAX_BYTES")
2316 .ok()
2317 .and_then(|s| s.trim().parse::<usize>().ok())
2318 .filter(|&n| n > 0)
2319 .unwrap_or(8 * 1024 * 1024);
2320 let max_lines = std::env::var("CODEX_HELPER_USAGE_REPLAY_MAX_LINES")
2321 .ok()
2322 .and_then(|s| s.trim().parse::<usize>().ok())
2323 .filter(|&n| n > 0)
2324 .unwrap_or(20_000);
2325
2326 let mut file = match std::fs::File::open(&log_path) {
2327 Ok(f) => f,
2328 Err(_) => return 0,
2329 };
2330 let len: u64 = file.metadata().map(|m| m.len()).unwrap_or_default();
2331 let start = len.saturating_sub(max_bytes as u64);
2332 if file.seek(SeekFrom::Start(start)).is_err() {
2333 return 0;
2334 }
2335 let mut buf = Vec::new();
2336 if file.read_to_end(&mut buf).is_err() {
2337 return 0;
2338 }
2339 if start > 0 {
2340 if let Some(pos) = buf.iter().position(|b| *b == b'\n') {
2341 buf = buf[pos + 1..].to_vec();
2342 } else {
2343 return 0;
2344 }
2345 }
2346
2347 let text = match std::str::from_utf8(&buf) {
2348 Ok(s) => s,
2349 Err(_) => return 0,
2350 };
2351 let lines = text
2352 .lines()
2353 .map(|l| l.trim())
2354 .filter(|l| !l.is_empty())
2355 .collect::<Vec<_>>();
2356 let start_idx = lines.len().saturating_sub(max_lines);
2357
2358 let mut events = Vec::new();
2359 for line in &lines[start_idx..] {
2360 let Ok(v) = serde_json::from_str::<JsonValue>(line) else {
2361 continue;
2362 };
2363 let Some(svc) = v.get("service").and_then(|x| x.as_str()) else {
2364 continue;
2365 };
2366 if svc != service_name {
2367 continue;
2368 }
2369
2370 let ended_at_ms = v.get("timestamp_ms").and_then(|x| x.as_u64()).unwrap_or(0);
2371 let status_code = v.get("status_code").and_then(|x| x.as_u64()).unwrap_or(0) as u16;
2372 let duration_ms = v.get("duration_ms").and_then(|x| x.as_u64()).unwrap_or(0);
2373 let station_name = v
2374 .get("station_name")
2375 .and_then(|x| x.as_str())
2376 .unwrap_or("-")
2377 .to_string();
2378 let upstream_base_url = v
2379 .get("upstream_base_url")
2380 .and_then(|x| x.as_str())
2381 .unwrap_or("-")
2382 .to_string();
2383 let provider_id = v
2384 .get("provider_id")
2385 .and_then(|x| x.as_str())
2386 .map(|s| s.to_string())
2387 .or_else(|| base_url_to_provider_id.get(&upstream_base_url).cloned())
2388 .unwrap_or_else(|| "-".to_string());
2389 let usage = v
2390 .get("usage")
2391 .and_then(|u| serde_json::from_value::<UsageMetrics>(u.clone()).ok());
2392 let ttfb_ms = v.get("ttfb_ms").and_then(|x| x.as_u64());
2393
2394 events.push((
2395 ended_at_ms,
2396 status_code,
2397 duration_ms,
2398 station_name,
2399 provider_id,
2400 usage,
2401 ttfb_ms,
2402 ));
2403 }
2404
2405 if events.is_empty() {
2406 return 0;
2407 }
2408
2409 let mut guard = self.usage_rollups.write().await;
2410 let rollup = guard.entry(service_name.to_string()).or_default();
2411 for (ended_at_ms, status_code, duration_ms, cfg_key, provider_key, usage, ttfb_ms) in
2412 events.iter()
2413 {
2414 let day = (*ended_at_ms / 86_400_000) as i32;
2415 rollup
2416 .loaded
2417 .record(*status_code, *duration_ms, usage.as_ref(), None, *ttfb_ms);
2418 rollup.by_day.entry(day).or_default().record(
2419 *status_code,
2420 *duration_ms,
2421 usage.as_ref(),
2422 None,
2423 *ttfb_ms,
2424 );
2425 rollup.by_config.entry(cfg_key.clone()).or_default().record(
2426 *status_code,
2427 *duration_ms,
2428 usage.as_ref(),
2429 None,
2430 *ttfb_ms,
2431 );
2432 rollup
2433 .by_config_day
2434 .entry(cfg_key.clone())
2435 .or_default()
2436 .entry(day)
2437 .or_default()
2438 .record(*status_code, *duration_ms, usage.as_ref(), None, *ttfb_ms);
2439 rollup
2440 .by_provider
2441 .entry(provider_key.clone())
2442 .or_default()
2443 .record(*status_code, *duration_ms, usage.as_ref(), None, *ttfb_ms);
2444 rollup
2445 .by_provider_day
2446 .entry(provider_key.clone())
2447 .or_default()
2448 .entry(day)
2449 .or_default()
2450 .record(*status_code, *duration_ms, usage.as_ref(), None, *ttfb_ms);
2451 }
2452
2453 let replayed = events.len();
2454 self.notify_state_changed();
2455 replayed
2456 }
2457
2458 pub async fn resolve_session_cwd(&self, session_id: &str) -> Option<String> {
2459 if self.session_cwd_cache_max_entries == 0 {
2460 return sessions::find_codex_session_cwd_by_id(session_id)
2461 .await
2462 .ok()
2463 .flatten();
2464 }
2465
2466 let now_ms = std::time::SystemTime::now()
2467 .duration_since(std::time::UNIX_EPOCH)
2468 .map(|d| d.as_millis() as u64)
2469 .unwrap_or(0);
2470
2471 {
2472 let guard = self.session_cwd_cache.read().await;
2473 if let Some(v) = guard.get(session_id) {
2474 let out = v.cwd.clone();
2475 drop(guard);
2476 let mut guard = self.session_cwd_cache.write().await;
2477 if let Some(v) = guard.get_mut(session_id) {
2478 v.last_seen_ms = now_ms;
2479 }
2480 return out;
2481 }
2482 }
2483
2484 let resolved = sessions::find_codex_session_cwd_by_id(session_id)
2487 .await
2488 .ok()
2489 .flatten();
2490
2491 let mut guard = self.session_cwd_cache.write().await;
2492 guard.insert(
2493 session_id.to_string(),
2494 SessionCwdCacheEntry {
2495 cwd: resolved.clone(),
2496 last_seen_ms: now_ms,
2497 },
2498 );
2499 resolved
2500 }
2501
2502 #[allow(clippy::too_many_arguments)]
2503 pub async fn begin_request(
2504 &self,
2505 service: &str,
2506 method: &str,
2507 path: &str,
2508 session_id: Option<String>,
2509 session_identity_source: Option<SessionIdentitySource>,
2510 client_name: Option<String>,
2511 client_addr: Option<String>,
2512 cwd: Option<String>,
2513 model: Option<String>,
2514 reasoning_effort: Option<String>,
2515 service_tier: Option<String>,
2516 started_at_ms: u64,
2517 ) -> u64 {
2518 let id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
2519 let trace_id = Some(crate::logging::request_trace_id(service, id));
2520 let req = ActiveRequest {
2521 id,
2522 trace_id,
2523 session_id,
2524 session_identity_source,
2525 client_name,
2526 client_addr,
2527 cwd,
2528 model,
2529 reasoning_effort,
2530 service_tier,
2531 station_name: None,
2532 provider_id: None,
2533 upstream_base_url: None,
2534 route_decision: None,
2535 service: service.to_string(),
2536 method: method.to_string(),
2537 path: path.to_string(),
2538 started_at_ms,
2539 };
2540 let mut guard = self.active_requests.write().await;
2541 guard.insert(id, req);
2542 self.notify_state_changed();
2543 id
2544 }
2545
2546 #[cfg(test)]
2547 pub(crate) fn begin_request_for_test(&self) -> BeginRequestTestBuilder<'_> {
2548 BeginRequestTestBuilder::new(self)
2549 }
2550
2551 pub async fn update_request_route(
2552 &self,
2553 request_id: u64,
2554 station_name: Option<String>,
2555 provider_id: Option<String>,
2556 upstream_base_url: String,
2557 route_decision: Option<RouteDecisionProvenance>,
2558 ) {
2559 let mut guard = self.active_requests.write().await;
2560 let Some(req) = guard.get_mut(&request_id) else {
2561 return;
2562 };
2563 req.station_name = station_name;
2564 req.provider_id = provider_id;
2565 req.upstream_base_url = Some(upstream_base_url);
2566 req.route_decision = route_decision;
2567 self.notify_state_changed();
2568 }
2569
2570 pub async fn finish_request(&self, params: FinishRequestParams) {
2571 let mut active = self.active_requests.write().await;
2572 let Some(req) = active.remove(¶ms.id) else {
2573 return;
2574 };
2575 drop(active);
2576
2577 let pricing_model = req
2578 .route_decision
2579 .as_ref()
2580 .and_then(|decision| decision.effective_model.as_ref())
2581 .map(|value| value.value.as_str())
2582 .or(req.model.as_deref());
2583 let cost = estimate_request_cost_from_operator_catalog_for_service(
2584 pricing_model,
2585 params.usage.as_ref(),
2586 CostAdjustments::default(),
2587 &req.service,
2588 );
2589
2590 let mut finished = FinishedRequest {
2591 id: params.id,
2592 trace_id: req.trace_id,
2593 session_id: req.session_id,
2594 session_identity_source: req.session_identity_source,
2595 client_name: req.client_name,
2596 client_addr: req.client_addr,
2597 cwd: req.cwd,
2598 model: req.model,
2599 reasoning_effort: req.reasoning_effort,
2600 service_tier: params.observed_service_tier.or(req.service_tier),
2601 station_name: req.station_name,
2602 provider_id: req.provider_id,
2603 upstream_base_url: req.upstream_base_url,
2604 route_decision: req.route_decision,
2605 usage: params.usage.clone(),
2606 cost,
2607 retry: params.retry,
2608 observability: RequestObservability::default(),
2609 service: req.service,
2610 method: req.method,
2611 path: req.path,
2612 status_code: params.status_code,
2613 duration_ms: params.duration_ms,
2614 ttfb_ms: params.ttfb_ms,
2615 streaming: params.streaming,
2616 ended_at_ms: params.ended_at_ms,
2617 };
2618 finished.refresh_observability();
2619
2620 {
2621 let effective_ttfb_ms = finished.observability.ttfb_ms;
2622 let day = (finished.ended_at_ms / 86_400_000) as i32;
2623 let cfg_key = finished
2624 .station_name
2625 .clone()
2626 .unwrap_or_else(|| "-".to_string());
2627 let provider_key = finished
2628 .provider_id
2629 .clone()
2630 .unwrap_or_else(|| "-".to_string());
2631
2632 let mut rollups = self.usage_rollups.write().await;
2633 let rollup = rollups.entry(finished.service.clone()).or_default();
2634 rollup.loaded.record(
2635 finished.status_code,
2636 finished.duration_ms,
2637 finished.usage.as_ref(),
2638 Some(&finished.cost),
2639 effective_ttfb_ms,
2640 );
2641 rollup.by_day.entry(day).or_default().record(
2642 finished.status_code,
2643 finished.duration_ms,
2644 finished.usage.as_ref(),
2645 Some(&finished.cost),
2646 effective_ttfb_ms,
2647 );
2648 rollup.by_config.entry(cfg_key.clone()).or_default().record(
2649 finished.status_code,
2650 finished.duration_ms,
2651 finished.usage.as_ref(),
2652 Some(&finished.cost),
2653 effective_ttfb_ms,
2654 );
2655 rollup
2656 .by_config_day
2657 .entry(cfg_key)
2658 .or_default()
2659 .entry(day)
2660 .or_default()
2661 .record(
2662 finished.status_code,
2663 finished.duration_ms,
2664 finished.usage.as_ref(),
2665 Some(&finished.cost),
2666 effective_ttfb_ms,
2667 );
2668
2669 rollup
2670 .by_provider
2671 .entry(provider_key.clone())
2672 .or_default()
2673 .record(
2674 finished.status_code,
2675 finished.duration_ms,
2676 finished.usage.as_ref(),
2677 Some(&finished.cost),
2678 effective_ttfb_ms,
2679 );
2680 rollup
2681 .by_provider_day
2682 .entry(provider_key)
2683 .or_default()
2684 .entry(day)
2685 .or_default()
2686 .record(
2687 finished.status_code,
2688 finished.duration_ms,
2689 finished.usage.as_ref(),
2690 Some(&finished.cost),
2691 effective_ttfb_ms,
2692 );
2693 }
2694
2695 if let Some(sid) = finished.session_id.as_deref() {
2696 let mut stats = self.session_stats.write().await;
2697 let entry = stats.entry(sid.to_string()).or_default();
2698 entry.turns_total = entry.turns_total.saturating_add(1);
2699 if finished.session_identity_source.is_some() {
2700 entry.last_session_identity_source = finished.session_identity_source;
2701 }
2702 entry.last_client_name = finished
2703 .client_name
2704 .clone()
2705 .or(entry.last_client_name.clone());
2706 entry.last_client_addr = finished
2707 .client_addr
2708 .clone()
2709 .or(entry.last_client_addr.clone());
2710 entry.last_model = finished.model.clone().or(entry.last_model.clone());
2711 entry.last_reasoning_effort = finished
2712 .reasoning_effort
2713 .clone()
2714 .or(entry.last_reasoning_effort.clone());
2715 entry.last_service_tier = finished
2716 .service_tier
2717 .clone()
2718 .or(entry.last_service_tier.clone());
2719 entry.last_provider_id = finished
2720 .provider_id
2721 .clone()
2722 .or(entry.last_provider_id.clone());
2723 entry.last_station_name = finished
2724 .station_name
2725 .clone()
2726 .or(entry.last_station_name.clone());
2727 if finished.route_decision.is_some() {
2728 entry.last_route_decision = finished.route_decision.clone();
2729 }
2730 if let Some(u) = finished.usage.as_ref() {
2731 entry.last_usage = Some(u.clone());
2732 entry.total_usage.add_assign(u);
2733 entry.turns_with_usage = entry.turns_with_usage.saturating_add(1);
2734 }
2735 entry.last_status = Some(finished.status_code);
2736 entry.last_duration_ms = Some(finished.duration_ms);
2737 entry.last_ended_at_ms = Some(finished.ended_at_ms);
2738 entry.last_seen_ms = finished.ended_at_ms;
2739 }
2740
2741 let mut recent = self.recent_finished.write().await;
2742 recent.push_front(finished);
2743 while recent.len() > recent_finished_max() {
2744 recent.pop_back();
2745 }
2746 self.notify_state_changed();
2747 }
2748
2749 pub async fn list_active_requests(&self) -> Vec<ActiveRequest> {
2750 let guard = self.active_requests.read().await;
2751 let mut vec = guard.values().cloned().collect::<Vec<_>>();
2752 vec.sort_by_key(|r| r.started_at_ms);
2753 vec
2754 }
2755
2756 pub async fn list_recent_finished(&self, limit: usize) -> Vec<FinishedRequest> {
2757 let guard = self.recent_finished.read().await;
2758 guard.iter().take(limit).cloned().collect()
2759 }
2760
2761 pub async fn list_session_stats(&self) -> HashMap<String, SessionStats> {
2762 let guard = self.session_stats.read().await;
2763 guard.clone()
2764 }
2765
2766 pub async fn list_session_identity_cards(
2767 &self,
2768 recent_limit: usize,
2769 ) -> Vec<SessionIdentityCard> {
2770 let recent_limit = recent_limit.clamp(1, recent_finished_max());
2771 let (
2772 active,
2773 recent,
2774 overrides,
2775 station_overrides,
2776 model_overrides,
2777 service_tier_overrides,
2778 bindings,
2779 route_affinities,
2780 global_station_override,
2781 stats,
2782 ) = tokio::join!(
2783 self.list_active_requests(),
2784 self.list_recent_finished(recent_limit),
2785 self.list_session_effort_overrides(),
2786 self.list_session_station_overrides(),
2787 self.list_session_model_overrides(),
2788 self.list_session_service_tier_overrides(),
2789 self.list_session_bindings(),
2790 self.list_session_route_affinities(),
2791 self.get_global_station_override(),
2792 self.list_session_stats(),
2793 );
2794 build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
2795 active: &active,
2796 recent: &recent,
2797 overrides: &overrides,
2798 station_overrides: &station_overrides,
2799 model_overrides: &model_overrides,
2800 service_tier_overrides: &service_tier_overrides,
2801 bindings: &bindings,
2802 route_affinities: &route_affinities,
2803 global_station_override: global_station_override.as_deref(),
2804 stats: &stats,
2805 })
2806 }
2807
2808 async fn resolve_host_transcript_paths_cached(
2809 &self,
2810 session_ids: &[String],
2811 ) -> HashMap<String, Option<String>> {
2812 let mut unique = session_ids
2813 .iter()
2814 .map(|sid| sid.trim())
2815 .filter(|sid| !sid.is_empty())
2816 .map(ToOwned::to_owned)
2817 .collect::<Vec<_>>();
2818 unique.sort();
2819 unique.dedup();
2820 if unique.is_empty() {
2821 return HashMap::new();
2822 }
2823
2824 if self.session_transcript_path_cache_max_entries == 0 {
2825 return sessions::find_codex_session_files_by_ids(&unique)
2826 .await
2827 .unwrap_or_default()
2828 .into_iter()
2829 .map(|(sid, path)| (sid, Some(path.to_string_lossy().to_string())))
2830 .collect();
2831 }
2832
2833 let now_ms = unix_now_ms();
2834 let ttl_ms = self.session_transcript_path_cache_ttl_ms;
2835 let mut resolved = HashMap::<String, Option<String>>::new();
2836 let mut stale_or_missing = Vec::<String>::new();
2837
2838 {
2839 let cache = self.session_transcript_path_cache.read().await;
2840 for sid in &unique {
2841 let fresh = cache.get(sid).filter(|entry| {
2842 ttl_ms == 0 || now_ms.saturating_sub(entry.last_checked_ms) <= ttl_ms
2843 });
2844 if let Some(entry) = fresh {
2845 resolved.insert(sid.clone(), entry.path.clone());
2846 } else {
2847 stale_or_missing.push(sid.clone());
2848 }
2849 }
2850 }
2851
2852 if !stale_or_missing.is_empty() {
2853 let found = sessions::find_codex_session_files_by_ids(&stale_or_missing)
2854 .await
2855 .unwrap_or_default();
2856 let mut cache = self.session_transcript_path_cache.write().await;
2857 for sid in stale_or_missing {
2858 let path = found
2859 .get(&sid)
2860 .map(|path| path.to_string_lossy().to_string());
2861 cache.insert(
2862 sid.clone(),
2863 SessionTranscriptPathCacheEntry {
2864 path: path.clone(),
2865 last_checked_ms: now_ms,
2866 last_seen_ms: now_ms,
2867 },
2868 );
2869 resolved.insert(sid, path);
2870 }
2871 prune_lru_cache(
2872 &mut cache,
2873 self.session_transcript_path_cache_max_entries,
2874 |entry| entry.last_seen_ms,
2875 );
2876 }
2877
2878 {
2879 let mut cache = self.session_transcript_path_cache.write().await;
2880 for sid in &unique {
2881 if let Some(entry) = cache.get_mut(sid) {
2882 entry.last_seen_ms = now_ms;
2883 }
2884 }
2885 }
2886
2887 resolved
2888 }
2889
2890 pub async fn enrich_session_identity_cards_with_cached_host_transcripts(
2891 &self,
2892 cards: &mut [SessionIdentityCard],
2893 ) {
2894 let session_ids = cards
2895 .iter()
2896 .filter_map(|card| card.session_id.as_deref())
2897 .map(str::to_owned)
2898 .collect::<Vec<_>>();
2899 let resolved = self
2900 .resolve_host_transcript_paths_cached(&session_ids)
2901 .await;
2902 for card in cards {
2903 card.host_local_transcript_path = card
2904 .session_id
2905 .as_deref()
2906 .and_then(|sid| resolved.get(sid))
2907 .and_then(Clone::clone);
2908 }
2909 }
2910
2911 pub async fn list_session_identity_cards_with_host_transcripts(
2912 &self,
2913 recent_limit: usize,
2914 ) -> Vec<SessionIdentityCard> {
2915 let mut cards = self.list_session_identity_cards(recent_limit).await;
2916 self.enrich_session_identity_cards_with_cached_host_transcripts(&mut cards)
2917 .await;
2918 cards
2919 }
2920
2921 pub fn spawn_cleanup_task(state: Arc<Self>) {
2922 tokio::spawn(async move {
2924 let mut tick = interval(Duration::from_secs(30));
2925 loop {
2926 tick.tick().await;
2927 state.prune_periodic().await;
2928 }
2929 });
2930 }
2931
2932 async fn prune_periodic(&self) {
2933 let now_ms = std::time::SystemTime::now()
2934 .duration_since(std::time::UNIX_EPOCH)
2935 .map(|d| d.as_millis() as u64)
2936 .unwrap_or(0);
2937
2938 let active = self.active_requests.read().await;
2940 let mut active_sessions: HashMap<String, ()> = HashMap::new();
2941 for req in active.values() {
2942 if let Some(sid) = req.session_id.as_deref() {
2943 active_sessions.insert(sid.to_string(), ());
2944 }
2945 }
2946
2947 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2948 let cutoff_override = now_ms - self.session_override_ttl_ms;
2949 let mut overrides = self.session_effort_overrides.write().await;
2950 overrides.retain(|sid, v| {
2951 if active_sessions.contains_key(sid) {
2952 return true;
2953 }
2954 v.last_seen_ms >= cutoff_override
2955 });
2956 }
2957
2958 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2959 let cutoff_override = now_ms - self.session_override_ttl_ms;
2960 let mut overrides = self.session_station_overrides.write().await;
2961 overrides.retain(|sid, v| {
2962 if active_sessions.contains_key(sid) {
2963 return true;
2964 }
2965 v.last_seen_ms >= cutoff_override
2966 });
2967 }
2968
2969 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2970 let cutoff_override = now_ms - self.session_override_ttl_ms;
2971 let mut overrides = self.session_route_target_overrides.write().await;
2972 overrides.retain(|sid, v| {
2973 if active_sessions.contains_key(sid) {
2974 return true;
2975 }
2976 v.last_seen_ms >= cutoff_override
2977 });
2978 }
2979
2980 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2981 let cutoff_override = now_ms - self.session_override_ttl_ms;
2982 let mut overrides = self.session_model_overrides.write().await;
2983 overrides.retain(|sid, v| {
2984 if active_sessions.contains_key(sid) {
2985 return true;
2986 }
2987 v.last_seen_ms >= cutoff_override
2988 });
2989 }
2990
2991 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2992 let cutoff_override = now_ms - self.session_override_ttl_ms;
2993 let mut overrides = self.session_service_tier_overrides.write().await;
2994 overrides.retain(|sid, v| {
2995 if active_sessions.contains_key(sid) {
2996 return true;
2997 }
2998 v.last_seen_ms >= cutoff_override
2999 });
3000 }
3001
3002 if self.session_binding_ttl_ms > 0 && now_ms >= self.session_binding_ttl_ms {
3003 let cutoff_binding = now_ms - self.session_binding_ttl_ms;
3004 let mut bindings = self.session_bindings.write().await;
3005 bindings.retain(|sid, entry| {
3006 if active_sessions.contains_key(sid) {
3007 return true;
3008 }
3009 entry.binding.last_seen_ms >= cutoff_binding
3010 });
3011 }
3012 if self.session_binding_max_entries > 0 {
3013 let mut bindings = self.session_bindings.write().await;
3014 if bindings.len() > self.session_binding_max_entries {
3015 let mut removable = bindings
3016 .iter()
3017 .filter(|(sid, _)| !active_sessions.contains_key(*sid))
3018 .map(|(sid, entry)| (sid.clone(), entry.binding.last_seen_ms))
3019 .collect::<Vec<_>>();
3020 removable.sort_by_key(|(_, last_seen_ms)| *last_seen_ms);
3021 let remove_count = bindings
3022 .len()
3023 .saturating_sub(self.session_binding_max_entries)
3024 .min(removable.len());
3025 for (sid, _) in removable.into_iter().take(remove_count) {
3026 bindings.remove(&sid);
3027 }
3028 }
3029 }
3030
3031 {
3032 let mut affinities = self.session_route_affinities.write().await;
3033 if self.session_route_affinity_ttl_ms > 0
3034 && now_ms >= self.session_route_affinity_ttl_ms
3035 {
3036 let cutoff_affinity = now_ms - self.session_route_affinity_ttl_ms;
3037 affinities.retain(|sid, affinity| {
3038 active_sessions.contains_key(sid)
3039 || affinity.last_selected_at_ms >= cutoff_affinity
3040 });
3041 }
3042 prune_lru_cache(
3043 &mut affinities,
3044 self.session_route_affinity_max_entries,
3045 |entry| entry.last_selected_at_ms,
3046 );
3047 }
3048
3049 let keep_days: i32 = std::env::var("CODEX_HELPER_USAGE_ROLLUP_KEEP_DAYS")
3051 .ok()
3052 .and_then(|s| s.trim().parse::<i32>().ok())
3053 .filter(|&n| n > 0)
3054 .unwrap_or(60);
3055 let now_day = (now_ms / 86_400_000) as i32;
3056 let cutoff_day = now_day.saturating_sub(keep_days);
3057 let mut rollups = self.usage_rollups.write().await;
3058 for rollup in rollups.values_mut() {
3059 rollup.by_day.retain(|day, _| *day >= cutoff_day);
3060 rollup.by_config_day.retain(|_, m| {
3061 m.retain(|day, _| *day >= cutoff_day);
3062 !m.is_empty()
3063 });
3064 rollup.by_provider_day.retain(|_, m| {
3065 m.retain(|day, _| *day >= cutoff_day);
3066 !m.is_empty()
3067 });
3068 }
3069
3070 let cutoff_cwd =
3071 if self.session_cwd_cache_ttl_ms == 0 || now_ms < self.session_cwd_cache_ttl_ms {
3072 0
3073 } else {
3074 now_ms - self.session_cwd_cache_ttl_ms
3075 };
3076 self.prune_session_cwd_cache(&active_sessions, cutoff_cwd)
3077 .await;
3078 self.prune_session_transcript_path_cache(now_ms).await;
3079
3080 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
3081 let cutoff_stats = now_ms - self.session_override_ttl_ms;
3082 let mut stats = self.session_stats.write().await;
3083 stats.retain(|sid, v| {
3084 active_sessions.contains_key(sid) || v.last_seen_ms >= cutoff_stats
3085 });
3086 }
3087 }
3088
3089 async fn prune_session_cwd_cache(&self, active_sessions: &HashMap<String, ()>, cutoff: u64) {
3090 if self.session_cwd_cache_max_entries == 0 {
3091 return;
3092 }
3093 let mut cache = self.session_cwd_cache.write().await;
3094
3095 if self.session_cwd_cache_ttl_ms > 0 {
3096 cache.retain(|sid, v| {
3097 if active_sessions.contains_key(sid) {
3098 return true;
3099 }
3100 v.last_seen_ms >= cutoff
3101 });
3102 }
3103
3104 let max = self.session_cwd_cache_max_entries;
3105 if max == 0 || cache.len() <= max {
3106 return;
3107 }
3108
3109 let mut keys = cache
3111 .iter()
3112 .map(|(sid, v)| (sid.clone(), v.last_seen_ms))
3113 .collect::<Vec<_>>();
3114 keys.sort_by_key(|(_, t)| *t);
3115 let remove_count = keys.len().saturating_sub(max);
3116 for (sid, _) in keys.into_iter().take(remove_count) {
3117 cache.remove(&sid);
3118 }
3119 }
3120
3121 async fn prune_session_transcript_path_cache(&self, now_ms: u64) {
3122 let mut cache = self.session_transcript_path_cache.write().await;
3123 if self.session_transcript_path_cache_max_entries == 0 {
3124 cache.clear();
3125 return;
3126 }
3127
3128 if self.session_transcript_path_cache_ttl_ms > 0 {
3129 let cutoff = now_ms.saturating_sub(self.session_transcript_path_cache_ttl_ms);
3130 cache.retain(|_, entry| entry.last_seen_ms >= cutoff);
3131 }
3132
3133 prune_lru_cache(
3134 &mut cache,
3135 self.session_transcript_path_cache_max_entries,
3136 |entry| entry.last_seen_ms,
3137 );
3138 }
3139}
3140
3141#[cfg(test)]
3142#[derive(Debug)]
3143pub(crate) struct BeginRequestTestBuilder<'a> {
3144 state: &'a ProxyState,
3145 service: &'static str,
3146 method: &'static str,
3147 path: &'static str,
3148 session_id: Option<String>,
3149 session_identity_source: Option<SessionIdentitySource>,
3150 client_name: Option<String>,
3151 client_addr: Option<String>,
3152 cwd: Option<String>,
3153 model: Option<String>,
3154 reasoning_effort: Option<String>,
3155 service_tier: Option<String>,
3156 started_at_ms: u64,
3157}
3158
3159#[cfg(test)]
3160impl<'a> BeginRequestTestBuilder<'a> {
3161 pub(crate) fn new(state: &'a ProxyState) -> Self {
3162 Self {
3163 state,
3164 service: "codex",
3165 method: "POST",
3166 path: "/v1/responses",
3167 session_id: None,
3168 session_identity_source: None,
3169 client_name: None,
3170 client_addr: None,
3171 cwd: None,
3172 model: None,
3173 reasoning_effort: None,
3174 service_tier: None,
3175 started_at_ms: 0,
3176 }
3177 }
3178
3179 pub(crate) fn session_id(mut self, value: impl Into<String>) -> Self {
3180 self.session_id = Some(value.into());
3181 self
3182 }
3183
3184 pub(crate) fn model(mut self, value: impl Into<String>) -> Self {
3185 self.model = Some(value.into());
3186 self
3187 }
3188
3189 pub(crate) fn service_tier(mut self, value: impl Into<String>) -> Self {
3190 self.service_tier = Some(value.into());
3191 self
3192 }
3193
3194 pub(crate) fn started_at_ms(mut self, value: u64) -> Self {
3195 self.started_at_ms = value;
3196 self
3197 }
3198
3199 pub(crate) async fn begin(self) -> u64 {
3200 self.state
3201 .begin_request(
3202 self.service,
3203 self.method,
3204 self.path,
3205 self.session_id,
3206 self.session_identity_source,
3207 self.client_name,
3208 self.client_addr,
3209 self.cwd,
3210 self.model,
3211 self.reasoning_effort,
3212 self.service_tier,
3213 self.started_at_ms,
3214 )
3215 .await
3216 }
3217}
3218
3219#[cfg(test)]
3220mod tests {
3221 use super::*;
3222
3223 use crate::config::{ServiceConfig, ServiceConfigManager, UpstreamAuth, UpstreamConfig};
3224 use crate::runtime_identity::ProviderEndpointKey;
3225 use std::path::Path;
3226 use std::sync::OnceLock;
3227
3228 async fn env_lock() -> tokio::sync::MutexGuard<'static, ()> {
3229 static LOCK: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
3230 LOCK.get_or_init(|| tokio::sync::Mutex::new(()))
3231 .lock()
3232 .await
3233 }
3234
3235 #[derive(Default)]
3236 struct ScopedEnv {
3237 saved: Vec<(String, Option<String>)>,
3238 }
3239
3240 impl ScopedEnv {
3241 unsafe fn set(&mut self, key: &str, value: &str) {
3242 if !self.saved.iter().any(|(saved_key, _)| saved_key == key) {
3243 self.saved.push((key.to_string(), std::env::var(key).ok()));
3244 }
3245 unsafe {
3246 std::env::set_var(key, value);
3247 }
3248 }
3249
3250 unsafe fn set_path(&mut self, key: &str, value: &Path) {
3251 unsafe {
3252 self.set(key, value.to_string_lossy().as_ref());
3253 }
3254 }
3255 }
3256
3257 impl Drop for ScopedEnv {
3258 fn drop(&mut self) {
3259 for (key, value) in self.saved.iter().rev() {
3260 match value {
3261 Some(value) => unsafe {
3262 std::env::set_var(key, value);
3263 },
3264 None => unsafe {
3265 std::env::remove_var(key);
3266 },
3267 }
3268 }
3269 }
3270 }
3271
3272 fn test_runtime_policy(
3273 session_override_ttl_ms: u64,
3274 session_binding_ttl_ms: u64,
3275 session_binding_max_entries: usize,
3276 ) -> RuntimePolicy {
3277 RuntimePolicy {
3278 session_override_ttl_ms,
3279 session_binding_ttl_ms,
3280 session_binding_max_entries,
3281 session_route_affinity_ttl_ms: 0,
3282 session_route_affinity_max_entries: 5_000,
3283 session_route_affinity_store: SessionRouteAffinityStore::from_env(),
3284 session_cwd_cache_ttl_ms: 0,
3285 session_cwd_cache_max_entries: 0,
3286 session_transcript_path_cache_ttl_ms: 30_000,
3287 session_transcript_path_cache_max_entries: 5_000,
3288 }
3289 }
3290
3291 #[test]
3292 fn begin_and_finish_requests_keep_trace_id() {
3293 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3294 runtime.block_on(async {
3295 let state = ProxyState::new();
3296 let request_id = state
3297 .begin_request_for_test()
3298 .model("gpt-5")
3299 .service_tier("priority")
3300 .started_at_ms(100)
3301 .begin()
3302 .await;
3303
3304 let active = state.list_active_requests().await;
3305 assert_eq!(active[0].trace_id.as_deref(), Some("codex-1"));
3306
3307 state
3308 .finish_request(FinishRequestParams {
3309 id: request_id,
3310 status_code: 200,
3311 duration_ms: 10,
3312 ended_at_ms: 110,
3313 observed_service_tier: Some("priority".to_string()),
3314 usage: None,
3315 retry: None,
3316 ttfb_ms: Some(4),
3317 streaming: false,
3318 })
3319 .await;
3320
3321 let recent = state.list_recent_finished(1).await;
3322 assert_eq!(recent[0].trace_id.as_deref(), Some("codex-1"));
3323 assert_eq!(recent[0].observability.trace_id.as_deref(), Some("codex-1"));
3324 assert!(recent[0].observability.fast_mode);
3325 assert_eq!(recent[0].observability.generation_ms, Some(6));
3326 });
3327 }
3328
3329 #[test]
3330 fn state_change_subscription_tracks_dashboard_mutations() {
3331 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3332 runtime.block_on(async {
3333 let state = ProxyState::new();
3334 let mut changes = state.subscribe_state_changes();
3335
3336 state
3337 .record_provider_balance_snapshot(
3338 "codex",
3339 ProviderBalanceSnapshot {
3340 provider_id: "input6".to_string(),
3341 station_name: Some("routing".to_string()),
3342 upstream_index: Some(0),
3343 fetched_at_ms: 100,
3344 stale_after_ms: Some(1_000),
3345 status: BalanceSnapshotStatus::Ok,
3346 total_balance_usd: Some("12.50".to_string()),
3347 ..ProviderBalanceSnapshot::default()
3348 },
3349 )
3350 .await;
3351 changes.changed().await.expect("balance change");
3352 let balance_version = *changes.borrow();
3353 assert!(balance_version > 0);
3354
3355 let request_id = state
3356 .begin_request_for_test()
3357 .session_id("sid-1")
3358 .model("gpt-5")
3359 .started_at_ms(200)
3360 .begin()
3361 .await;
3362 changes.changed().await.expect("begin request change");
3363 let active_version = *changes.borrow();
3364 assert!(active_version > balance_version);
3365
3366 state
3367 .finish_request(FinishRequestParams {
3368 id: request_id,
3369 status_code: 200,
3370 duration_ms: 25,
3371 ended_at_ms: 250,
3372 observed_service_tier: None,
3373 usage: None,
3374 retry: None,
3375 ttfb_ms: Some(5),
3376 streaming: false,
3377 })
3378 .await;
3379 changes.changed().await.expect("finish request change");
3380 assert!(*changes.borrow() > active_version);
3381 });
3382 }
3383
3384 #[test]
3385 fn state_change_subscription_ignores_session_touch_only_updates() {
3386 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3387 runtime.block_on(async {
3388 let state = ProxyState::new();
3389 let mut changes = state.subscribe_state_changes();
3390
3391 state
3392 .set_session_model_override("sid-1".to_string(), "gpt-5".to_string(), 100)
3393 .await;
3394 changes.changed().await.expect("model override change");
3395 let version = *changes.borrow();
3396
3397 state.touch_session_model_override("sid-1", 200).await;
3398 assert_eq!(*changes.borrow(), version);
3399 });
3400 }
3401
3402 #[test]
3403 fn finish_request_estimates_cost_and_rolls_up_cost() {
3404 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3405 runtime.block_on(async {
3406 let state = ProxyState::new();
3407 let request_id = state
3408 .begin_request_for_test()
3409 .model("gpt-5")
3410 .started_at_ms(100)
3411 .begin()
3412 .await;
3413
3414 state
3415 .finish_request(FinishRequestParams {
3416 id: request_id,
3417 status_code: 200,
3418 duration_ms: 10,
3419 ended_at_ms: 110,
3420 observed_service_tier: None,
3421 usage: Some(UsageMetrics {
3422 input_tokens: 1_000,
3423 output_tokens: 500,
3424 cached_input_tokens: 100,
3425 total_tokens: 1_500,
3426 ..UsageMetrics::default()
3427 }),
3428 retry: None,
3429 ttfb_ms: Some(4),
3430 streaming: false,
3431 })
3432 .await;
3433
3434 let recent = state.list_recent_finished(1).await;
3435 assert_eq!(recent[0].cost.total_cost_usd.as_deref(), Some("0.0061375"));
3436
3437 let rollup = state.get_usage_rollup_view("codex", 12, 1).await;
3438 assert_eq!(
3439 rollup.loaded.cost.total_cost_usd.as_deref(),
3440 Some("0.0061375")
3441 );
3442 assert_eq!(rollup.loaded.cost.priced_requests, 1);
3443 assert_eq!(rollup.loaded.cost.unpriced_requests, 0);
3444 });
3445 }
3446
3447 #[test]
3448 fn usage_rollup_view_scores_entities_inside_selected_window() {
3449 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3450 runtime.block_on(async {
3451 let state = ProxyState::new();
3452 let now_ms = std::time::SystemTime::now()
3453 .duration_since(std::time::UNIX_EPOCH)
3454 .map(|d| d.as_millis() as u64)
3455 .unwrap_or(0);
3456 let old_ms = now_ms.saturating_sub(10 * 86_400_000);
3457
3458 let old_id = state
3459 .begin_request_for_test()
3460 .model("gpt-5")
3461 .started_at_ms(old_ms.saturating_sub(1_000))
3462 .begin()
3463 .await;
3464 state
3465 .update_request_route(
3466 old_id,
3467 Some("old-station".to_string()),
3468 Some("old-provider".to_string()),
3469 "https://old.example/v1".to_string(),
3470 None,
3471 )
3472 .await;
3473 state
3474 .finish_request(FinishRequestParams {
3475 id: old_id,
3476 status_code: 200,
3477 duration_ms: 20,
3478 ended_at_ms: old_ms,
3479 observed_service_tier: None,
3480 usage: Some(UsageMetrics {
3481 total_tokens: 100_000,
3482 ..UsageMetrics::default()
3483 }),
3484 retry: None,
3485 ttfb_ms: Some(5),
3486 streaming: false,
3487 })
3488 .await;
3489
3490 let fresh_id = state
3491 .begin_request_for_test()
3492 .model("gpt-5")
3493 .started_at_ms(now_ms.saturating_sub(1_000))
3494 .begin()
3495 .await;
3496 state
3497 .update_request_route(
3498 fresh_id,
3499 Some("fresh-station".to_string()),
3500 Some("fresh-provider".to_string()),
3501 "https://fresh.example/v1".to_string(),
3502 None,
3503 )
3504 .await;
3505 state
3506 .finish_request(FinishRequestParams {
3507 id: fresh_id,
3508 status_code: 200,
3509 duration_ms: 10,
3510 ended_at_ms: now_ms,
3511 observed_service_tier: None,
3512 usage: Some(UsageMetrics {
3513 total_tokens: 10,
3514 ..UsageMetrics::default()
3515 }),
3516 retry: None,
3517 ttfb_ms: Some(3),
3518 streaming: false,
3519 })
3520 .await;
3521
3522 let week = state.get_usage_rollup_view("codex", 10, 7).await;
3523 assert_eq!(week.loaded.requests_total, 2);
3524 assert_eq!(week.window.requests_total, 1);
3525 assert_eq!(week.by_day.len(), 7);
3526 assert_eq!(week.by_config[0].0, "fresh-station");
3527 assert_eq!(week.by_provider[0].0, "fresh-provider");
3528
3529 let loaded = state.get_usage_rollup_view("codex", 10, 0).await;
3530 assert_eq!(loaded.window.requests_total, 2);
3531 assert_eq!(loaded.by_config[0].0, "old-station");
3532 assert_eq!(loaded.by_provider[0].0, "old-provider");
3533 });
3534 }
3535
3536 #[test]
3537 fn build_session_identity_cards_merges_sources_and_sorts_newest_first() {
3538 let active = vec![ActiveRequest {
3539 id: 1,
3540 trace_id: Some("codex-1".to_string()),
3541 session_id: Some("sid-active".to_string()),
3542 session_identity_source: Some(SessionIdentitySource::Header),
3543 client_name: Some("Frank-Laptop".to_string()),
3544 client_addr: Some("100.64.0.8".to_string()),
3545 cwd: Some("G:/codes/project".to_string()),
3546 model: Some("gpt-5.4".to_string()),
3547 reasoning_effort: Some("medium".to_string()),
3548 service_tier: Some("priority".to_string()),
3549 station_name: Some("right".to_string()),
3550 provider_id: Some("right".to_string()),
3551 upstream_base_url: Some("https://right.example/v1".to_string()),
3552 route_decision: None,
3553 service: "codex".to_string(),
3554 method: "POST".to_string(),
3555 path: "/v1/responses".to_string(),
3556 started_at_ms: 500,
3557 }];
3558 let recent = vec![
3559 FinishedRequest {
3560 id: 2,
3561 trace_id: Some("codex-2".to_string()),
3562 session_id: Some("sid-recent".to_string()),
3563 session_identity_source: Some(SessionIdentitySource::PromptCacheKey),
3564 client_name: Some("Studio-Mini".to_string()),
3565 client_addr: Some("100.64.0.9".to_string()),
3566 cwd: Some("G:/codes/other".to_string()),
3567 model: Some("gpt-5.3".to_string()),
3568 reasoning_effort: Some("high".to_string()),
3569 service_tier: Some("default".to_string()),
3570 station_name: Some("vibe".to_string()),
3571 provider_id: Some("vibe".to_string()),
3572 upstream_base_url: Some("https://vibe.example/v1".to_string()),
3573 route_decision: None,
3574 usage: Some(UsageMetrics {
3575 input_tokens: 1,
3576 output_tokens: 2,
3577 reasoning_tokens: 3,
3578 total_tokens: 6,
3579 ..UsageMetrics::default()
3580 }),
3581 cost: CostBreakdown::default(),
3582 retry: None,
3583 observability: RequestObservability::default(),
3584 service: "codex".to_string(),
3585 method: "POST".to_string(),
3586 path: "/v1/responses".to_string(),
3587 status_code: 200,
3588 duration_ms: 1200,
3589 ttfb_ms: Some(100),
3590 streaming: false,
3591 ended_at_ms: 2_000,
3592 },
3593 FinishedRequest {
3594 id: 3,
3595 trace_id: Some("codex-3".to_string()),
3596 session_id: Some("sid-active".to_string()),
3597 session_identity_source: Some(SessionIdentitySource::Header),
3598 client_name: Some("Frank-Laptop".to_string()),
3599 client_addr: Some("100.64.0.8".to_string()),
3600 cwd: Some("G:/codes/project".to_string()),
3601 model: Some("gpt-5.4".to_string()),
3602 reasoning_effort: Some("low".to_string()),
3603 service_tier: Some("flex".to_string()),
3604 station_name: Some("right".to_string()),
3605 provider_id: Some("right".to_string()),
3606 upstream_base_url: Some("https://right.example/v1".to_string()),
3607 route_decision: None,
3608 usage: None,
3609 cost: CostBreakdown::default(),
3610 retry: None,
3611 observability: RequestObservability::default(),
3612 service: "codex".to_string(),
3613 method: "POST".to_string(),
3614 path: "/v1/responses".to_string(),
3615 status_code: 429,
3616 duration_ms: 900,
3617 ttfb_ms: None,
3618 streaming: false,
3619 ended_at_ms: 1_000,
3620 },
3621 ];
3622 let overrides = HashMap::from([("sid-active".to_string(), "xhigh".to_string())]);
3623 let config_overrides = HashMap::from([("sid-active".to_string(), "temp".to_string())]);
3624 let model_overrides =
3625 HashMap::from([("sid-active".to_string(), "gpt-5.4-mini".to_string())]);
3626 let service_tier_overrides =
3627 HashMap::from([("sid-active".to_string(), "priority".to_string())]);
3628 let stats = HashMap::from([(
3629 "sid-active".to_string(),
3630 SessionStats {
3631 turns_total: 3,
3632 last_session_identity_source: Some(SessionIdentitySource::Header),
3633 last_client_name: Some("Frank-Laptop".to_string()),
3634 last_client_addr: Some("100.64.0.8".to_string()),
3635 last_model: Some("gpt-5.4".to_string()),
3636 last_reasoning_effort: Some("low".to_string()),
3637 last_service_tier: Some("flex".to_string()),
3638 last_provider_id: Some("right".to_string()),
3639 last_station_name: Some("right".to_string()),
3640 last_route_decision: None,
3641 last_usage: None,
3642 total_usage: UsageMetrics {
3643 input_tokens: 10,
3644 output_tokens: 20,
3645 reasoning_tokens: 5,
3646 total_tokens: 35,
3647 ..UsageMetrics::default()
3648 },
3649 turns_with_usage: 2,
3650 last_status: Some(429),
3651 last_duration_ms: Some(900),
3652 last_ended_at_ms: Some(1_000),
3653 last_seen_ms: 1_000,
3654 },
3655 )]);
3656
3657 let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
3658 active: &active,
3659 recent: &recent,
3660 overrides: &overrides,
3661 station_overrides: &config_overrides,
3662 model_overrides: &model_overrides,
3663 service_tier_overrides: &service_tier_overrides,
3664 bindings: &HashMap::new(),
3665 route_affinities: &HashMap::new(),
3666 global_station_override: None,
3667 stats: &stats,
3668 });
3669
3670 assert_eq!(cards.len(), 2);
3671 assert_eq!(cards[0].session_id.as_deref(), Some("sid-recent"));
3672 assert_eq!(
3673 cards[0].session_identity_source,
3674 Some(SessionIdentitySource::PromptCacheKey)
3675 );
3676 assert_eq!(
3677 cards[0].observation_scope,
3678 SessionObservationScope::HostLocalEnriched
3679 );
3680 assert_eq!(cards[0].last_client_name.as_deref(), Some("Studio-Mini"));
3681 assert_eq!(cards[0].last_client_addr.as_deref(), Some("100.64.0.9"));
3682 assert_eq!(cards[1].session_id.as_deref(), Some("sid-active"));
3683 assert_eq!(
3684 cards[1].session_identity_source,
3685 Some(SessionIdentitySource::Header)
3686 );
3687 assert_eq!(
3688 cards[1].observation_scope,
3689 SessionObservationScope::HostLocalEnriched
3690 );
3691 assert_eq!(cards[1].active_count, 1);
3692 assert_eq!(cards[1].last_client_name.as_deref(), Some("Frank-Laptop"));
3693 assert_eq!(cards[1].last_client_addr.as_deref(), Some("100.64.0.8"));
3694 assert_eq!(cards[1].last_status, Some(429));
3695 assert_eq!(cards[1].override_effort.as_deref(), Some("xhigh"));
3696 assert_eq!(cards[1].override_station_name.as_deref(), Some("temp"));
3697 assert_eq!(cards[1].override_model.as_deref(), Some("gpt-5.4-mini"));
3698 assert_eq!(cards[1].override_service_tier.as_deref(), Some("priority"));
3699 assert_eq!(
3700 cards[1]
3701 .effective_model
3702 .as_ref()
3703 .map(|value| value.value.as_str()),
3704 Some("gpt-5.4-mini")
3705 );
3706 assert_eq!(
3707 cards[1].effective_model.as_ref().map(|value| value.source),
3708 Some(RouteValueSource::SessionOverride)
3709 );
3710 assert_eq!(
3711 cards[1]
3712 .effective_reasoning_effort
3713 .as_ref()
3714 .map(|value| value.source),
3715 Some(RouteValueSource::SessionOverride)
3716 );
3717 assert_eq!(
3718 cards[1]
3719 .effective_service_tier
3720 .as_ref()
3721 .map(|value| value.source),
3722 Some(RouteValueSource::SessionOverride)
3723 );
3724 assert_eq!(
3725 cards[1]
3726 .effective_station
3727 .as_ref()
3728 .map(|value| value.source),
3729 Some(RouteValueSource::SessionOverride)
3730 );
3731 assert!(cards[1].effective_upstream_base_url.is_none());
3732 assert_eq!(
3733 cards[1].last_upstream_base_url.as_deref(),
3734 Some("https://right.example/v1")
3735 );
3736 assert_eq!(cards[1].turns_total, Some(3));
3737 assert_eq!(cards[1].last_service_tier.as_deref(), Some("flex"));
3738 assert_eq!(
3739 cards[1].total_usage.as_ref().map(|u| u.total_tokens),
3740 Some(35)
3741 );
3742 }
3743
3744 #[test]
3745 fn build_session_identity_cards_default_profile_keeps_request_fields() {
3746 let active = vec![ActiveRequest {
3747 id: 1,
3748 trace_id: Some("codex-1".to_string()),
3749 session_id: Some("sid-bound".to_string()),
3750 session_identity_source: Some(SessionIdentitySource::Header),
3751 client_name: Some("Workstation".to_string()),
3752 client_addr: Some("100.64.0.10".to_string()),
3753 cwd: None,
3754 model: Some("gpt-observed".to_string()),
3755 reasoning_effort: Some("medium".to_string()),
3756 service_tier: Some("default".to_string()),
3757 station_name: Some("right".to_string()),
3758 provider_id: None,
3759 upstream_base_url: None,
3760 route_decision: None,
3761 service: "codex".to_string(),
3762 method: "POST".to_string(),
3763 path: "/v1/responses".to_string(),
3764 started_at_ms: 10,
3765 }];
3766 let bindings = HashMap::from([(
3767 "sid-bound".to_string(),
3768 SessionBinding {
3769 session_id: "sid-bound".to_string(),
3770 profile_name: Some("daily".to_string()),
3771 station_name: Some("vibe".to_string()),
3772 model: Some("gpt-bound".to_string()),
3773 reasoning_effort: Some("high".to_string()),
3774 service_tier: Some("priority".to_string()),
3775 continuity_mode: SessionContinuityMode::DefaultProfile,
3776 created_at_ms: 1,
3777 updated_at_ms: 1,
3778 last_seen_ms: 10,
3779 },
3780 )]);
3781
3782 let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
3783 active: &active,
3784 recent: &[],
3785 overrides: &HashMap::new(),
3786 station_overrides: &HashMap::new(),
3787 model_overrides: &HashMap::new(),
3788 service_tier_overrides: &HashMap::new(),
3789 bindings: &bindings,
3790 route_affinities: &HashMap::new(),
3791 global_station_override: None,
3792 stats: &HashMap::new(),
3793 });
3794
3795 assert_eq!(cards[0].binding_profile_name.as_deref(), Some("daily"));
3796 assert_eq!(
3797 cards[0].observation_scope,
3798 SessionObservationScope::ObservedOnly
3799 );
3800 assert_eq!(
3801 cards[0].binding_continuity_mode,
3802 Some(SessionContinuityMode::DefaultProfile)
3803 );
3804 assert_eq!(
3805 cards[0]
3806 .effective_model
3807 .as_ref()
3808 .map(|value| (value.value.as_str(), value.source)),
3809 Some(("gpt-observed", RouteValueSource::RequestPayload))
3810 );
3811 assert_eq!(
3812 cards[0]
3813 .effective_reasoning_effort
3814 .as_ref()
3815 .map(|value| (value.value.as_str(), value.source)),
3816 Some(("medium", RouteValueSource::RequestPayload))
3817 );
3818 assert_eq!(
3819 cards[0]
3820 .effective_service_tier
3821 .as_ref()
3822 .map(|value| (value.value.as_str(), value.source)),
3823 Some(("default", RouteValueSource::RequestPayload))
3824 );
3825 assert_eq!(
3826 cards[0]
3827 .effective_station
3828 .as_ref()
3829 .map(|value| (value.value.as_str(), value.source)),
3830 Some(("vibe", RouteValueSource::ProfileDefault))
3831 );
3832 }
3833
3834 #[test]
3835 fn build_session_identity_cards_keeps_request_fields_but_allows_global_config_override() {
3836 let active = vec![ActiveRequest {
3837 id: 1,
3838 trace_id: Some("codex-1".to_string()),
3839 session_id: Some("sid-bound".to_string()),
3840 session_identity_source: Some(SessionIdentitySource::Header),
3841 client_name: Some("Workstation".to_string()),
3842 client_addr: Some("100.64.0.10".to_string()),
3843 cwd: None,
3844 model: Some("gpt-observed".to_string()),
3845 reasoning_effort: Some("medium".to_string()),
3846 service_tier: Some("default".to_string()),
3847 station_name: Some("vibe".to_string()),
3848 provider_id: None,
3849 upstream_base_url: Some("https://vibe.example/v1".to_string()),
3850 route_decision: None,
3851 service: "codex".to_string(),
3852 method: "POST".to_string(),
3853 path: "/v1/responses".to_string(),
3854 started_at_ms: 10,
3855 }];
3856 let bindings = HashMap::from([(
3857 "sid-bound".to_string(),
3858 SessionBinding {
3859 session_id: "sid-bound".to_string(),
3860 profile_name: Some("daily".to_string()),
3861 station_name: Some("vibe".to_string()),
3862 model: Some("gpt-bound".to_string()),
3863 reasoning_effort: Some("high".to_string()),
3864 service_tier: Some("priority".to_string()),
3865 continuity_mode: SessionContinuityMode::DefaultProfile,
3866 created_at_ms: 1,
3867 updated_at_ms: 1,
3868 last_seen_ms: 10,
3869 },
3870 )]);
3871
3872 let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
3873 active: &active,
3874 recent: &[],
3875 overrides: &HashMap::new(),
3876 station_overrides: &HashMap::new(),
3877 model_overrides: &HashMap::new(),
3878 service_tier_overrides: &HashMap::new(),
3879 bindings: &bindings,
3880 route_affinities: &HashMap::new(),
3881 global_station_override: Some("right"),
3882 stats: &HashMap::new(),
3883 });
3884
3885 assert_eq!(cards[0].binding_profile_name.as_deref(), Some("daily"));
3886 assert_eq!(
3887 cards[0]
3888 .effective_model
3889 .as_ref()
3890 .map(|value| (value.value.as_str(), value.source)),
3891 Some(("gpt-observed", RouteValueSource::RequestPayload))
3892 );
3893 assert_eq!(
3894 cards[0]
3895 .effective_reasoning_effort
3896 .as_ref()
3897 .map(|value| (value.value.as_str(), value.source)),
3898 Some(("medium", RouteValueSource::RequestPayload))
3899 );
3900 assert_eq!(
3901 cards[0]
3902 .effective_service_tier
3903 .as_ref()
3904 .map(|value| (value.value.as_str(), value.source)),
3905 Some(("default", RouteValueSource::RequestPayload))
3906 );
3907 assert_eq!(
3908 cards[0]
3909 .effective_station
3910 .as_ref()
3911 .map(|value| (value.value.as_str(), value.source)),
3912 Some(("right", RouteValueSource::GlobalOverride))
3913 );
3914 assert!(cards[0].effective_upstream_base_url.is_none());
3915 }
3916
3917 #[test]
3918 fn enrich_session_identity_cards_with_runtime_applies_station_mapping_and_single_upstream() {
3919 let mut cards = vec![SessionIdentityCard {
3920 session_id: Some("sid-1".to_string()),
3921 last_model: Some("gpt-5.4".to_string()),
3922 last_station_name: Some("right".to_string()),
3923 last_upstream_base_url: Some("https://right.example/v1".to_string()),
3924 effective_model: Some(ResolvedRouteValue::new(
3925 "gpt-5.4",
3926 RouteValueSource::RequestPayload,
3927 )),
3928 effective_station: Some(ResolvedRouteValue::new(
3929 "right",
3930 RouteValueSource::RuntimeFallback,
3931 )),
3932 ..SessionIdentityCard::default()
3933 }];
3934
3935 let mut mgr = ServiceConfigManager {
3936 active: Some("right".to_string()),
3937 ..ServiceConfigManager::default()
3938 };
3939 mgr.configs.insert(
3940 "right".to_string(),
3941 ServiceConfig {
3942 name: "right".to_string(),
3943 alias: None,
3944 enabled: true,
3945 level: 1,
3946 upstreams: vec![UpstreamConfig {
3947 base_url: "https://right.example/v1".to_string(),
3948 auth: UpstreamAuth::default(),
3949 tags: HashMap::new(),
3950 supported_models: HashMap::new(),
3951 model_mapping: HashMap::from([(
3952 "gpt-5.4".to_string(),
3953 "gpt-5.4-fast".to_string(),
3954 )]),
3955 }],
3956 },
3957 );
3958
3959 enrich_session_identity_cards_with_runtime(&mut cards, &mgr);
3960
3961 assert_eq!(
3962 cards[0]
3963 .effective_model
3964 .as_ref()
3965 .map(|value| value.value.as_str()),
3966 Some("gpt-5.4-fast")
3967 );
3968 assert_eq!(
3969 cards[0].effective_model.as_ref().map(|value| value.source),
3970 Some(RouteValueSource::StationMapping)
3971 );
3972 assert_eq!(
3973 cards[0]
3974 .effective_upstream_base_url
3975 .as_ref()
3976 .map(|value| value.value.as_str()),
3977 Some("https://right.example/v1")
3978 );
3979 }
3980
3981 #[test]
3982 fn apply_session_profile_binding_sets_binding_and_clears_manual_overrides() {
3983 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3984 runtime.block_on(async {
3985 let state = ProxyState::new();
3986 let now_ms = 42;
3987 let mut mgr = ServiceConfigManager::default();
3988 mgr.configs.insert(
3989 "right".to_string(),
3990 ServiceConfig {
3991 name: "right".to_string(),
3992 alias: None,
3993 enabled: true,
3994 level: 1,
3995 upstreams: vec![UpstreamConfig {
3996 base_url: "https://right.example/v1".to_string(),
3997 auth: UpstreamAuth::default(),
3998 tags: HashMap::from([
3999 ("supports_reasoning_effort".to_string(), "true".to_string()),
4000 ("supports_service_tier".to_string(), "true".to_string()),
4001 ]),
4002 supported_models: HashMap::from([("gpt-5.4".to_string(), true)]),
4003 model_mapping: HashMap::new(),
4004 }],
4005 },
4006 );
4007 mgr.profiles.insert(
4008 "fast".to_string(),
4009 crate::config::ServiceControlProfile {
4010 extends: None,
4011 station: Some("right".to_string()),
4012 model: Some("gpt-5.4".to_string()),
4013 reasoning_effort: Some("low".to_string()),
4014 service_tier: Some("flex".to_string()),
4015 },
4016 );
4017
4018 state
4019 .set_session_station_override("sid-1".to_string(), "other".to_string(), 1)
4020 .await;
4021 state
4022 .set_session_model_override("sid-1".to_string(), "gpt-x".to_string(), 1)
4023 .await;
4024 state
4025 .set_session_effort_override("sid-1".to_string(), "high".to_string(), 1)
4026 .await;
4027 state
4028 .set_session_service_tier_override("sid-1".to_string(), "priority".to_string(), 1)
4029 .await;
4030
4031 state
4032 .apply_session_profile_binding(
4033 "codex",
4034 &mgr,
4035 "sid-1".to_string(),
4036 "fast".to_string(),
4037 now_ms,
4038 )
4039 .await
4040 .expect("apply profile");
4041
4042 let binding = state
4043 .get_session_binding("sid-1")
4044 .await
4045 .expect("binding exists");
4046 assert_eq!(binding.profile_name.as_deref(), Some("fast"));
4047 assert_eq!(binding.station_name.as_deref(), Some("right"));
4048 assert_eq!(binding.model.as_deref(), Some("gpt-5.4"));
4049 assert_eq!(binding.reasoning_effort.as_deref(), Some("low"));
4050 assert_eq!(binding.service_tier.as_deref(), Some("flex"));
4051 assert_eq!(
4052 binding.continuity_mode,
4053 SessionContinuityMode::ManualProfile
4054 );
4055 assert_eq!(binding.updated_at_ms, now_ms);
4056 assert!(state.get_session_station_override("sid-1").await.is_none());
4057 assert!(state.get_session_model_override("sid-1").await.is_none());
4058 assert!(state.get_session_effort_override("sid-1").await.is_none());
4059 assert!(
4060 state
4061 .get_session_service_tier_override("sid-1")
4062 .await
4063 .is_none()
4064 );
4065 });
4066 }
4067
4068 #[test]
4069 fn list_session_manual_overrides_merges_all_dimensions() {
4070 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4071 runtime.block_on(async {
4072 let state = ProxyState::new();
4073 state
4074 .set_session_reasoning_effort_override("sid-1".to_string(), "high".to_string(), 1)
4075 .await;
4076 state
4077 .set_session_station_override("sid-1".to_string(), "right".to_string(), 1)
4078 .await;
4079 state
4080 .set_session_model_override("sid-1".to_string(), "gpt-5.4".to_string(), 1)
4081 .await;
4082 state
4083 .set_session_service_tier_override("sid-1".to_string(), "priority".to_string(), 1)
4084 .await;
4085 state
4086 .set_session_model_override("sid-2".to_string(), "gpt-5.4-mini".to_string(), 2)
4087 .await;
4088
4089 let merged = state.list_session_manual_overrides().await;
4090 assert_eq!(merged.len(), 2);
4091 assert_eq!(
4092 merged
4093 .get("sid-1")
4094 .and_then(|overrides| overrides.reasoning_effort.as_deref()),
4095 Some("high")
4096 );
4097 assert_eq!(
4098 merged
4099 .get("sid-1")
4100 .and_then(|overrides| overrides.station_name.as_deref()),
4101 Some("right")
4102 );
4103 assert_eq!(
4104 merged
4105 .get("sid-1")
4106 .and_then(|overrides| overrides.model.as_deref()),
4107 Some("gpt-5.4")
4108 );
4109 assert_eq!(
4110 merged
4111 .get("sid-1")
4112 .and_then(|overrides| overrides.service_tier.as_deref()),
4113 Some("priority")
4114 );
4115 assert_eq!(
4116 merged
4117 .get("sid-2")
4118 .and_then(|overrides| overrides.model.as_deref()),
4119 Some("gpt-5.4-mini")
4120 );
4121 assert!(
4122 merged
4123 .get("sid-2")
4124 .is_some_and(|overrides| overrides.reasoning_effort.is_none())
4125 );
4126 });
4127 }
4128
4129 #[test]
4130 fn provider_balance_snapshots_are_recorded_and_refreshed() {
4131 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4132 runtime.block_on(async {
4133 let state = ProxyState::new();
4134 state
4135 .record_provider_balance_snapshot(
4136 "codex",
4137 ProviderBalanceSnapshot {
4138 provider_id: "packycode".to_string(),
4139 station_name: Some("right".to_string()),
4140 upstream_index: Some(2),
4141 source: "usage_provider:budget_http_json".to_string(),
4142 fetched_at_ms: 10,
4143 stale_after_ms: Some(0),
4144 stale: false,
4145 status: BalanceSnapshotStatus::Unknown,
4146 exhausted: Some(false),
4147 exhaustion_affects_routing: true,
4148 plan_name: None,
4149 total_balance_usd: Some("3.5".to_string()),
4150 subscription_balance_usd: None,
4151 paygo_balance_usd: None,
4152 monthly_budget_usd: Some("5".to_string()),
4153 monthly_spent_usd: Some("1.5".to_string()),
4154 quota_period: None,
4155 quota_remaining_usd: None,
4156 quota_limit_usd: None,
4157 quota_used_usd: None,
4158 unlimited_quota: None,
4159 total_used_usd: None,
4160 today_used_usd: None,
4161 total_requests: None,
4162 today_requests: None,
4163 total_tokens: None,
4164 today_tokens: None,
4165 error: None,
4166 },
4167 )
4168 .await;
4169
4170 let view = state.get_provider_balance_view("codex").await;
4171 let balances = view.get("right").expect("station balance");
4172 assert_eq!(balances.len(), 1);
4173 assert_eq!(balances[0].provider_id, "packycode");
4174 assert_eq!(balances[0].status, BalanceSnapshotStatus::Stale);
4175 assert_eq!(balances[0].exhausted, Some(false));
4176
4177 let summary = state
4178 .get_provider_balance_summary_view("codex")
4179 .await
4180 .get("right")
4181 .cloned()
4182 .expect("station balance summary");
4183 assert_eq!(summary.snapshots, 1);
4184 assert_eq!(summary.stale, 1);
4185 assert_eq!(summary.routing_snapshots, 1);
4186 });
4187 }
4188
4189 #[test]
4190 fn provider_balance_snapshots_keep_multiple_providers_per_upstream() {
4191 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4192 runtime.block_on(async {
4193 let state = ProxyState::new();
4194 for (provider_id, status) in [
4195 ("general", BalanceSnapshotStatus::Error),
4196 ("newapi", BalanceSnapshotStatus::Ok),
4197 ] {
4198 state
4199 .record_provider_balance_snapshot(
4200 "codex",
4201 ProviderBalanceSnapshot {
4202 provider_id: provider_id.to_string(),
4203 station_name: Some("routing".to_string()),
4204 upstream_index: Some(1),
4205 source: "usage_provider:test".to_string(),
4206 fetched_at_ms: 10,
4207 stale_after_ms: None,
4208 stale: false,
4209 status,
4210 exhausted: if status == BalanceSnapshotStatus::Ok {
4211 Some(false)
4212 } else {
4213 None
4214 },
4215 exhaustion_affects_routing: true,
4216 plan_name: None,
4217 total_balance_usd: if status == BalanceSnapshotStatus::Ok {
4218 Some("3.5".to_string())
4219 } else {
4220 None
4221 },
4222 subscription_balance_usd: None,
4223 paygo_balance_usd: None,
4224 monthly_budget_usd: None,
4225 monthly_spent_usd: None,
4226 quota_period: None,
4227 quota_remaining_usd: None,
4228 quota_limit_usd: None,
4229 quota_used_usd: None,
4230 unlimited_quota: None,
4231 total_used_usd: None,
4232 today_used_usd: None,
4233 total_requests: None,
4234 today_requests: None,
4235 total_tokens: None,
4236 today_tokens: None,
4237 error: if status == BalanceSnapshotStatus::Error {
4238 Some("decode failed".to_string())
4239 } else {
4240 None
4241 },
4242 },
4243 )
4244 .await;
4245 }
4246
4247 let view = state.get_provider_balance_view("codex").await;
4248 let balances = view.get("routing").expect("station balance");
4249 assert_eq!(balances.len(), 2);
4250 assert_eq!(
4251 balances
4252 .iter()
4253 .map(|snapshot| snapshot.provider_id.as_str())
4254 .collect::<Vec<_>>(),
4255 vec!["general", "newapi"]
4256 );
4257 });
4258 }
4259
4260 #[test]
4261 fn provider_balance_error_refresh_preserves_previous_amounts() {
4262 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4263 runtime.block_on(async {
4264 let state = ProxyState::new();
4265 state
4266 .record_provider_balance_snapshot(
4267 "codex",
4268 ProviderBalanceSnapshot {
4269 provider_id: "input".to_string(),
4270 station_name: Some("routing".to_string()),
4271 upstream_index: Some(0),
4272 source: "usage_provider:sub2api_usage".to_string(),
4273 fetched_at_ms: unix_now_ms(),
4274 stale_after_ms: None,
4275 exhausted: Some(false),
4276 quota_period: Some("daily".to_string()),
4277 quota_remaining_usd: Some("263.68".to_string()),
4278 quota_limit_usd: Some("300.00".to_string()),
4279 ..ProviderBalanceSnapshot::default()
4280 },
4281 )
4282 .await;
4283
4284 state
4285 .record_provider_balance_snapshot(
4286 "codex",
4287 ProviderBalanceSnapshot {
4288 provider_id: "input".to_string(),
4289 station_name: Some("routing".to_string()),
4290 upstream_index: Some(0),
4291 source: "usage_provider:openai_balance_http_json".to_string(),
4292 fetched_at_ms: unix_now_ms(),
4293 stale_after_ms: None,
4294 error: Some("usage provider response read failed".to_string()),
4295 ..ProviderBalanceSnapshot::default()
4296 },
4297 )
4298 .await;
4299
4300 let view = state.get_provider_balance_view("codex").await;
4301 let snapshot = view
4302 .get("routing")
4303 .and_then(|snapshots| snapshots.first())
4304 .expect("routing balance snapshot");
4305
4306 assert_eq!(snapshot.status, BalanceSnapshotStatus::Error);
4307 assert_eq!(snapshot.quota_period.as_deref(), Some("daily"));
4308 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("263.68"));
4309 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("300.00"));
4310 assert_eq!(
4311 snapshot.error.as_deref(),
4312 Some("usage provider response read failed")
4313 );
4314 });
4315 }
4316
4317 #[test]
4318 fn prune_runtime_observability_keeps_catalog_provider_balances_on_routing_layout_change() {
4319 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4320 runtime.block_on(async {
4321 let state = ProxyState::new();
4322
4323 let mut initial_mgr = ServiceConfigManager::default();
4324 initial_mgr.configs.insert(
4325 "routing".to_string(),
4326 ServiceConfig {
4327 name: "routing".to_string(),
4328 alias: None,
4329 enabled: true,
4330 level: 1,
4331 upstreams: vec![
4332 UpstreamConfig {
4333 base_url: "https://input.example/v1".to_string(),
4334 auth: UpstreamAuth::default(),
4335 tags: HashMap::from([("provider_id".to_string(), "input".to_string())]),
4336 supported_models: HashMap::new(),
4337 model_mapping: HashMap::new(),
4338 },
4339 UpstreamConfig {
4340 base_url: "https://backup.example/v1".to_string(),
4341 auth: UpstreamAuth::default(),
4342 tags: HashMap::from([(
4343 "provider_id".to_string(),
4344 "backup".to_string(),
4345 )]),
4346 supported_models: HashMap::new(),
4347 model_mapping: HashMap::new(),
4348 },
4349 ],
4350 },
4351 );
4352 state
4353 .prune_runtime_observability_for_service("codex", &initial_mgr)
4354 .await;
4355
4356 state
4357 .record_provider_balance_snapshot(
4358 "codex",
4359 ProviderBalanceSnapshot {
4360 provider_id: "input".to_string(),
4361 station_name: Some("input".to_string()),
4362 upstream_index: Some(0),
4363 source: "usage_provider:test".to_string(),
4364 fetched_at_ms: 10,
4365 stale_after_ms: None,
4366 stale: false,
4367 status: BalanceSnapshotStatus::Ok,
4368 exhausted: Some(false),
4369 exhaustion_affects_routing: true,
4370 total_balance_usd: Some("3.5".to_string()),
4371 ..ProviderBalanceSnapshot::default()
4372 },
4373 )
4374 .await;
4375 state
4376 .record_provider_balance_snapshot(
4377 "codex",
4378 ProviderBalanceSnapshot {
4379 provider_id: "input".to_string(),
4380 station_name: Some("routing".to_string()),
4381 upstream_index: Some(0),
4382 source: "usage_provider:test".to_string(),
4383 fetched_at_ms: 10,
4384 stale_after_ms: None,
4385 stale: false,
4386 status: BalanceSnapshotStatus::Ok,
4387 exhausted: Some(false),
4388 exhaustion_affects_routing: true,
4389 total_balance_usd: Some("3.5".to_string()),
4390 ..ProviderBalanceSnapshot::default()
4391 },
4392 )
4393 .await;
4394
4395 let mut pinned_mgr = ServiceConfigManager::default();
4396 pinned_mgr.configs.insert(
4397 "routing".to_string(),
4398 ServiceConfig {
4399 name: "routing".to_string(),
4400 alias: None,
4401 enabled: true,
4402 level: 1,
4403 upstreams: vec![UpstreamConfig {
4404 base_url: "https://input.example/v1".to_string(),
4405 auth: UpstreamAuth::default(),
4406 tags: HashMap::from([("provider_id".to_string(), "input".to_string())]),
4407 supported_models: HashMap::new(),
4408 model_mapping: HashMap::new(),
4409 }],
4410 },
4411 );
4412 state
4413 .prune_runtime_observability_for_service("codex", &pinned_mgr)
4414 .await;
4415
4416 let view = state.get_provider_balance_view("codex").await;
4417 assert!(view.contains_key("input"));
4418 assert!(!view.contains_key("routing"));
4419 });
4420 }
4421
4422 #[test]
4423 fn prune_runtime_observability_removes_stale_service_keys() {
4424 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4425 runtime.block_on(async {
4426 let state = ProxyState::new();
4427
4428 state
4429 .set_station_enabled_override("codex", "old".to_string(), false, 1)
4430 .await;
4431 state
4432 .set_upstream_enabled_override(
4433 "codex",
4434 "https://old.example/v1".to_string(),
4435 false,
4436 1,
4437 )
4438 .await;
4439 state
4440 .set_provider_endpoint_enabled_override(
4441 "codex",
4442 ProviderEndpointKey::new("codex", "provider-old", "default"),
4443 false,
4444 1,
4445 )
4446 .await;
4447 state
4448 .set_provider_endpoint_runtime_state_override(
4449 "codex",
4450 ProviderEndpointKey::new("codex", "provider-old", "default"),
4451 RuntimeConfigState::BreakerOpen,
4452 1,
4453 )
4454 .await;
4455 state
4456 .record_station_health(
4457 "codex",
4458 "old".to_string(),
4459 StationHealth {
4460 checked_at_ms: 10,
4461 upstreams: vec![UpstreamHealth {
4462 base_url: "https://old.example/v1".to_string(),
4463 ok: Some(false),
4464 status_code: Some(500),
4465 latency_ms: Some(10),
4466 error: Some("boom".to_string()),
4467 passive: None,
4468 }],
4469 },
4470 )
4471 .await;
4472 state
4473 .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
4474 service_name: "codex".to_string(),
4475 station_name: "old".to_string(),
4476 base_url: "https://old.example/v1".to_string(),
4477 status_code: Some(500),
4478 error_class: Some("upstream_transport_error".to_string()),
4479 error: Some("boom".to_string()),
4480 now_ms: 20,
4481 })
4482 .await;
4483 state
4484 .record_provider_balance_snapshot(
4485 "codex",
4486 ProviderBalanceSnapshot {
4487 provider_id: "provider-old".to_string(),
4488 station_name: Some("old".to_string()),
4489 upstream_index: Some(0),
4490 source: "usage_provider:budget_http_json".to_string(),
4491 fetched_at_ms: 10,
4492 stale_after_ms: None,
4493 stale: false,
4494 status: BalanceSnapshotStatus::Ok,
4495 exhausted: Some(false),
4496 exhaustion_affects_routing: true,
4497 plan_name: None,
4498 total_balance_usd: Some("3.5".to_string()),
4499 subscription_balance_usd: None,
4500 paygo_balance_usd: None,
4501 monthly_budget_usd: None,
4502 monthly_spent_usd: None,
4503 quota_period: None,
4504 quota_remaining_usd: None,
4505 quota_limit_usd: None,
4506 quota_used_usd: None,
4507 unlimited_quota: None,
4508 total_used_usd: None,
4509 today_used_usd: None,
4510 total_requests: None,
4511 today_requests: None,
4512 total_tokens: None,
4513 today_tokens: None,
4514 error: None,
4515 },
4516 )
4517 .await;
4518
4519 let request_id = state
4520 .begin_request_for_test()
4521 .session_id("sid-old")
4522 .model("gpt-5")
4523 .started_at_ms(30)
4524 .begin()
4525 .await;
4526 state
4527 .update_request_route(
4528 request_id,
4529 Some("old".to_string()),
4530 Some("provider-old".to_string()),
4531 "https://old.example/v1".to_string(),
4532 None,
4533 )
4534 .await;
4535 state
4536 .finish_request(FinishRequestParams {
4537 id: request_id,
4538 status_code: 200,
4539 duration_ms: 5,
4540 ended_at_ms: 35,
4541 observed_service_tier: None,
4542 usage: None,
4543 retry: None,
4544 ttfb_ms: None,
4545 streaming: false,
4546 })
4547 .await;
4548
4549 let mut mgr = ServiceConfigManager::default();
4550 mgr.configs.insert(
4551 "new".to_string(),
4552 ServiceConfig {
4553 name: "new".to_string(),
4554 alias: None,
4555 enabled: true,
4556 level: 1,
4557 upstreams: vec![UpstreamConfig {
4558 base_url: "https://new.example/v1".to_string(),
4559 auth: UpstreamAuth::default(),
4560 tags: HashMap::from([(
4561 "provider_id".to_string(),
4562 "provider-new".to_string(),
4563 )]),
4564 supported_models: HashMap::new(),
4565 model_mapping: HashMap::new(),
4566 }],
4567 },
4568 );
4569
4570 state
4571 .prune_runtime_observability_for_service("codex", &mgr)
4572 .await;
4573
4574 assert!(state.get_station_meta_overrides("codex").await.is_empty());
4575 assert!(state.get_upstream_meta_overrides("codex").await.is_empty());
4576 assert!(state.get_station_health("codex").await.is_empty());
4577 assert!(state.get_provider_balance_view("codex").await.is_empty());
4578
4579 let rollup = state.get_usage_rollup_view("codex", 10, 10).await;
4580 assert!(rollup.by_config.is_empty());
4581 assert!(rollup.by_provider.is_empty());
4582 });
4583 }
4584
4585 #[test]
4586 fn get_upstream_meta_overrides_merges_endpoint_and_legacy_base_url_entries() {
4587 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4588 runtime.block_on(async {
4589 let state = ProxyState::new();
4590
4591 state
4592 .set_provider_endpoint_enabled_override(
4593 "codex",
4594 ProviderEndpointKey::new("codex", "alpha", "default"),
4595 false,
4596 1,
4597 )
4598 .await;
4599 state
4600 .set_provider_endpoint_runtime_state_override(
4601 "codex",
4602 ProviderEndpointKey::new("codex", "alpha", "default"),
4603 RuntimeConfigState::BreakerOpen,
4604 2,
4605 )
4606 .await;
4607 state
4608 .set_upstream_enabled_override(
4609 "codex",
4610 "https://legacy.example/v1".to_string(),
4611 true,
4612 3,
4613 )
4614 .await;
4615 state
4616 .set_upstream_runtime_state_override(
4617 "codex",
4618 "https://legacy.example/v1".to_string(),
4619 RuntimeConfigState::Draining,
4620 4,
4621 )
4622 .await;
4623
4624 let overrides = state.get_upstream_meta_overrides("codex").await;
4625
4626 assert_eq!(
4627 overrides.get("codex/alpha/default"),
4628 Some(&(Some(false), Some(RuntimeConfigState::BreakerOpen)))
4629 );
4630 assert_eq!(
4631 overrides.get("https://legacy.example/v1"),
4632 Some(&(Some(true), Some(RuntimeConfigState::Draining)))
4633 );
4634 });
4635 }
4636
4637 #[test]
4638 fn provider_endpoint_runtime_health_is_keyed_by_provider_endpoint() {
4639 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4640 runtime.block_on(async {
4641 let state = ProxyState::new();
4642 let monthly = ProviderEndpointKey::new("codex", "monthly", "default");
4643 let fallback = ProviderEndpointKey::new("codex", "fallback", "default");
4644
4645 state
4646 .record_provider_endpoint_attempt_success("codex", fallback.clone(), 10)
4647 .await;
4648 state
4649 .set_provider_endpoint_usage_exhausted("codex", monthly.clone(), true)
4650 .await;
4651 state
4652 .record_provider_endpoint_attempt_failure(
4653 "codex",
4654 monthly.clone(),
4655 0,
4656 crate::lb::CooldownBackoff {
4657 factor: 1,
4658 max_secs: 0,
4659 },
4660 )
4661 .await;
4662 state
4663 .record_provider_endpoint_attempt_failure(
4664 "codex",
4665 monthly.clone(),
4666 0,
4667 crate::lb::CooldownBackoff {
4668 factor: 1,
4669 max_secs: 0,
4670 },
4671 )
4672 .await;
4673 state
4674 .record_provider_endpoint_attempt_failure(
4675 "codex",
4676 monthly.clone(),
4677 30,
4678 crate::lb::CooldownBackoff {
4679 factor: 1,
4680 max_secs: 0,
4681 },
4682 )
4683 .await;
4684
4685 let runtime = state
4686 .route_plan_runtime_state_for_provider_endpoints("codex")
4687 .await;
4688 let monthly_state = runtime.provider_endpoint(&monthly);
4689 assert_eq!(monthly_state.failure_count, crate::lb::FAILURE_THRESHOLD);
4690 assert!(monthly_state.cooldown_active);
4691 assert!(monthly_state.usage_exhausted);
4692 assert_eq!(runtime.affinity_provider_endpoint(), Some(&fallback));
4693
4694 state
4695 .set_provider_endpoint_runtime_state_override(
4696 "codex",
4697 fallback.clone(),
4698 RuntimeConfigState::BreakerOpen,
4699 20,
4700 )
4701 .await;
4702 let runtime = state
4703 .route_plan_runtime_state_for_provider_endpoints("codex")
4704 .await;
4705 assert!(runtime.provider_endpoint(&fallback).runtime_disabled);
4706 });
4707 }
4708
4709 #[test]
4710 fn get_station_health_merges_passive_runtime_observations() {
4711 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4712 runtime.block_on(async {
4713 let state = ProxyState::new();
4714 state
4715 .record_station_health(
4716 "codex",
4717 "right".to_string(),
4718 StationHealth {
4719 checked_at_ms: 10,
4720 upstreams: vec![UpstreamHealth {
4721 base_url: "https://right.example/v1".to_string(),
4722 ok: Some(true),
4723 status_code: Some(200),
4724 latency_ms: Some(120),
4725 error: None,
4726 passive: None,
4727 }],
4728 },
4729 )
4730 .await;
4731 state
4732 .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
4733 service_name: "codex".to_string(),
4734 station_name: "right".to_string(),
4735 base_url: "https://right.example/v1".to_string(),
4736 status_code: Some(500),
4737 error_class: Some("cloudflare_timeout".to_string()),
4738 error: Some("upstream timed out".to_string()),
4739 now_ms: 20,
4740 })
4741 .await;
4742 state
4743 .record_passive_upstream_success(
4744 "codex",
4745 "right",
4746 "https://backup.example/v1",
4747 Some(200),
4748 30,
4749 )
4750 .await;
4751
4752 let health = state.get_station_health("codex").await;
4753 let right = health.get("right").expect("right health");
4754 assert_eq!(right.checked_at_ms, 30);
4755 assert_eq!(right.upstreams.len(), 2);
4756
4757 let primary = right
4758 .upstreams
4759 .iter()
4760 .find(|upstream| upstream.base_url == "https://right.example/v1")
4761 .expect("primary upstream");
4762 assert_eq!(primary.ok, Some(true));
4763 let primary_passive = primary.passive.as_ref().expect("primary passive");
4764 assert_eq!(primary_passive.state, PassiveHealthState::Degraded);
4765 assert_eq!(primary_passive.score, 50);
4766 assert_eq!(primary_passive.last_status_code, Some(500));
4767 assert_eq!(
4768 primary_passive.last_error_class.as_deref(),
4769 Some("cloudflare_timeout")
4770 );
4771
4772 let backup = right
4773 .upstreams
4774 .iter()
4775 .find(|upstream| upstream.base_url == "https://backup.example/v1")
4776 .expect("backup upstream");
4777 assert_eq!(backup.ok, None);
4778 let backup_passive = backup.passive.as_ref().expect("backup passive");
4779 assert_eq!(backup_passive.state, PassiveHealthState::Healthy);
4780 assert_eq!(backup_passive.score, 100);
4781 assert_eq!(backup_passive.last_status_code, Some(200));
4782 });
4783 }
4784
4785 #[test]
4786 fn passive_health_success_recovers_after_failure() {
4787 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4788 runtime.block_on(async {
4789 let state = ProxyState::new();
4790 state
4791 .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
4792 service_name: "codex".to_string(),
4793 station_name: "right".to_string(),
4794 base_url: "https://right.example/v1".to_string(),
4795 status_code: Some(500),
4796 error_class: Some("cloudflare_timeout".to_string()),
4797 error: Some("upstream timed out".to_string()),
4798 now_ms: 10,
4799 })
4800 .await;
4801 state
4802 .record_passive_upstream_success(
4803 "codex",
4804 "right",
4805 "https://right.example/v1",
4806 Some(200),
4807 20,
4808 )
4809 .await;
4810
4811 let health = state.get_station_health("codex").await;
4812 let right = health.get("right").expect("right health");
4813 let upstream = right.upstreams.first().expect("upstream");
4814 let passive = upstream.passive.as_ref().expect("passive");
4815 assert_eq!(passive.state, PassiveHealthState::Healthy);
4816 assert_eq!(passive.score, 100);
4817 assert_eq!(passive.consecutive_failures, 0);
4818 assert_eq!(passive.last_success_at_ms, Some(20));
4819 assert_eq!(passive.last_failure_at_ms, Some(10));
4820 assert_eq!(passive.last_error_class, None);
4821 assert_eq!(passive.last_error, None);
4822 });
4823 }
4824
4825 #[test]
4826 fn apply_session_profile_binding_uses_inherited_profile_values() {
4827 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4828 runtime.block_on(async {
4829 let state = ProxyState::new();
4830 let mut mgr = ServiceConfigManager::default();
4831 mgr.configs.insert(
4832 "right".to_string(),
4833 ServiceConfig {
4834 name: "right".to_string(),
4835 alias: None,
4836 enabled: true,
4837 level: 1,
4838 upstreams: vec![UpstreamConfig {
4839 base_url: "https://right.example/v1".to_string(),
4840 auth: UpstreamAuth::default(),
4841 tags: HashMap::from([
4842 ("supports_reasoning_effort".to_string(), "true".to_string()),
4843 ("supports_service_tier".to_string(), "true".to_string()),
4844 ]),
4845 supported_models: HashMap::from([("gpt-5.4".to_string(), true)]),
4846 model_mapping: HashMap::new(),
4847 }],
4848 },
4849 );
4850 mgr.profiles.insert(
4851 "base".to_string(),
4852 crate::config::ServiceControlProfile {
4853 extends: None,
4854 station: Some("right".to_string()),
4855 model: Some("gpt-5.4".to_string()),
4856 reasoning_effort: None,
4857 service_tier: Some("priority".to_string()),
4858 },
4859 );
4860 mgr.profiles.insert(
4861 "fast".to_string(),
4862 crate::config::ServiceControlProfile {
4863 extends: Some("base".to_string()),
4864 station: None,
4865 model: None,
4866 reasoning_effort: Some("low".to_string()),
4867 service_tier: None,
4868 },
4869 );
4870
4871 state
4872 .apply_session_profile_binding(
4873 "codex",
4874 &mgr,
4875 "sid-inherited".to_string(),
4876 "fast".to_string(),
4877 100,
4878 )
4879 .await
4880 .expect("apply inherited profile");
4881
4882 let binding = state
4883 .get_session_binding("sid-inherited")
4884 .await
4885 .expect("binding exists");
4886 assert_eq!(binding.profile_name.as_deref(), Some("fast"));
4887 assert_eq!(binding.station_name.as_deref(), Some("right"));
4888 assert_eq!(binding.model.as_deref(), Some("gpt-5.4"));
4889 assert_eq!(binding.reasoning_effort.as_deref(), Some("low"));
4890 assert_eq!(binding.service_tier.as_deref(), Some("priority"));
4891 });
4892 }
4893
4894 #[test]
4895 fn prune_periodic_keeps_sticky_binding_after_manual_override_ttl_expires() {
4896 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4897 runtime.block_on(async {
4898 let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(1, 0, 2_000));
4899 state
4900 .set_session_model_override("sid-sticky".to_string(), "gpt-5.4".to_string(), 0)
4901 .await;
4902 state
4903 .set_session_binding(SessionBinding {
4904 session_id: "sid-sticky".to_string(),
4905 profile_name: Some("daily".to_string()),
4906 station_name: Some("right".to_string()),
4907 model: Some("gpt-5.4".to_string()),
4908 reasoning_effort: Some("medium".to_string()),
4909 service_tier: Some("default".to_string()),
4910 continuity_mode: SessionContinuityMode::DefaultProfile,
4911 created_at_ms: 0,
4912 updated_at_ms: 0,
4913 last_seen_ms: 0,
4914 })
4915 .await;
4916
4917 state.prune_periodic().await;
4918
4919 assert!(
4920 state
4921 .get_session_model_override("sid-sticky")
4922 .await
4923 .is_none()
4924 );
4925 assert!(state.get_session_binding("sid-sticky").await.is_some());
4926 });
4927 }
4928
4929 #[test]
4930 fn prune_periodic_honors_opt_in_binding_ttl() {
4931 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4932 runtime.block_on(async {
4933 let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 1, 2_000));
4934 state
4935 .set_session_binding(SessionBinding {
4936 session_id: "sid-expire".to_string(),
4937 profile_name: Some("daily".to_string()),
4938 station_name: Some("right".to_string()),
4939 model: Some("gpt-5.4".to_string()),
4940 reasoning_effort: Some("medium".to_string()),
4941 service_tier: Some("default".to_string()),
4942 continuity_mode: SessionContinuityMode::DefaultProfile,
4943 created_at_ms: 0,
4944 updated_at_ms: 0,
4945 last_seen_ms: 0,
4946 })
4947 .await;
4948
4949 state.prune_periodic().await;
4950
4951 assert!(state.get_session_binding("sid-expire").await.is_none());
4952 });
4953 }
4954
4955 #[test]
4956 fn session_route_affinity_ttl_is_enforced_on_read() {
4957 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4958 runtime.block_on(async {
4959 let mut policy = test_runtime_policy(0, 0, 2_000);
4960 policy.session_route_affinity_ttl_ms = 1;
4961 let state = ProxyState::new_with_runtime_policy(None, policy);
4962 state
4963 .record_session_route_affinity_success(
4964 "sid-expire",
4965 SessionRouteAffinityTarget {
4966 route_graph_key: "graph".to_string(),
4967 session_identity_source: None,
4968 provider_endpoint: ProviderEndpointKey::new("codex", "monthly", "default"),
4969 upstream_base_url: "https://monthly.example/v1".to_string(),
4970 route_path: vec!["monthly_first".to_string(), "monthly".to_string()],
4971 },
4972 Some("first_success".to_string()),
4973 0,
4974 )
4975 .await;
4976
4977 assert!(
4978 state
4979 .get_session_route_affinity("sid-expire")
4980 .await
4981 .is_none()
4982 );
4983 assert!(state.list_session_route_affinities().await.is_empty());
4984 });
4985 }
4986
4987 #[test]
4988 fn session_route_affinity_ledger_does_not_restore_expired_entries() {
4989 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4990 runtime.block_on(async {
4991 let _home = env_lock().await;
4992 let temp_dir = std::env::temp_dir().join(format!(
4993 "codex-helper-route-ledger-test-{}",
4994 uuid::Uuid::new_v4()
4995 ));
4996 std::fs::create_dir_all(&temp_dir).expect("create temp dir");
4997 let mut scoped = ScopedEnv::default();
4998 unsafe {
4999 scoped.set_path("CODEX_HELPER_HOME", temp_dir.as_path());
5000 }
5001
5002 let mut policy = test_runtime_policy(0, 0, 2_000);
5003 policy.session_route_affinity_ttl_ms = 1;
5004 let first_state = ProxyState::new_with_runtime_policy(None, policy.clone());
5005 first_state
5006 .record_session_route_affinity_success(
5007 "sid-expired-ledger",
5008 SessionRouteAffinityTarget {
5009 route_graph_key: "graph".to_string(),
5010 session_identity_source: None,
5011 provider_endpoint: ProviderEndpointKey::new("codex", "monthly", "default"),
5012 upstream_base_url: "https://monthly.example/v1".to_string(),
5013 route_path: vec!["monthly_first".to_string(), "monthly".to_string()],
5014 },
5015 Some("first_success".to_string()),
5016 0,
5017 )
5018 .await;
5019
5020 let second_state = ProxyState::new_with_runtime_policy(None, policy);
5021 assert!(
5022 second_state
5023 .get_session_route_affinity("sid-expired-ledger")
5024 .await
5025 .is_none()
5026 );
5027 });
5028 }
5029
5030 #[test]
5031 fn prune_periodic_caps_sticky_bindings_by_last_seen() {
5032 let runtime = tokio::runtime::Runtime::new().expect("runtime");
5033 runtime.block_on(async {
5034 let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2));
5035 for idx in 0..3 {
5036 state
5037 .set_session_binding(SessionBinding {
5038 session_id: format!("sid-{idx}"),
5039 profile_name: Some("daily".to_string()),
5040 station_name: Some("right".to_string()),
5041 model: Some("gpt-5.4".to_string()),
5042 reasoning_effort: Some("medium".to_string()),
5043 service_tier: Some("default".to_string()),
5044 continuity_mode: SessionContinuityMode::DefaultProfile,
5045 created_at_ms: idx,
5046 updated_at_ms: idx,
5047 last_seen_ms: idx,
5048 })
5049 .await;
5050 }
5051
5052 state.prune_periodic().await;
5053
5054 assert!(state.get_session_binding("sid-0").await.is_none());
5055 assert!(state.get_session_binding("sid-1").await.is_some());
5056 assert!(state.get_session_binding("sid-2").await.is_some());
5057 });
5058 }
5059
5060 #[test]
5061 fn transcript_path_cache_records_negative_lookups() {
5062 let runtime = tokio::runtime::Runtime::new().expect("runtime");
5063 runtime.block_on(async {
5064 let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5065 let paths = state
5066 .resolve_host_transcript_paths_cached(&["missing-session".to_string()])
5067 .await;
5068
5069 assert_eq!(paths.get("missing-session"), Some(&None));
5070 let cache = state.session_transcript_path_cache.read().await;
5071 assert!(cache.contains_key("missing-session"));
5072 });
5073 }
5074}