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 day = (finished.ended_at_ms / 86_400_000) as i32;
2622 let cfg_key = finished
2623 .station_name
2624 .clone()
2625 .unwrap_or_else(|| "-".to_string());
2626 let provider_key = finished
2627 .provider_id
2628 .clone()
2629 .unwrap_or_else(|| "-".to_string());
2630
2631 let mut rollups = self.usage_rollups.write().await;
2632 let rollup = rollups.entry(finished.service.clone()).or_default();
2633 rollup.loaded.record(
2634 finished.status_code,
2635 finished.duration_ms,
2636 finished.usage.as_ref(),
2637 Some(&finished.cost),
2638 finished.ttfb_ms,
2639 );
2640 rollup.by_day.entry(day).or_default().record(
2641 finished.status_code,
2642 finished.duration_ms,
2643 finished.usage.as_ref(),
2644 Some(&finished.cost),
2645 finished.ttfb_ms,
2646 );
2647 rollup.by_config.entry(cfg_key.clone()).or_default().record(
2648 finished.status_code,
2649 finished.duration_ms,
2650 finished.usage.as_ref(),
2651 Some(&finished.cost),
2652 finished.ttfb_ms,
2653 );
2654 rollup
2655 .by_config_day
2656 .entry(cfg_key)
2657 .or_default()
2658 .entry(day)
2659 .or_default()
2660 .record(
2661 finished.status_code,
2662 finished.duration_ms,
2663 finished.usage.as_ref(),
2664 Some(&finished.cost),
2665 finished.ttfb_ms,
2666 );
2667
2668 rollup
2669 .by_provider
2670 .entry(provider_key.clone())
2671 .or_default()
2672 .record(
2673 finished.status_code,
2674 finished.duration_ms,
2675 finished.usage.as_ref(),
2676 Some(&finished.cost),
2677 finished.ttfb_ms,
2678 );
2679 rollup
2680 .by_provider_day
2681 .entry(provider_key)
2682 .or_default()
2683 .entry(day)
2684 .or_default()
2685 .record(
2686 finished.status_code,
2687 finished.duration_ms,
2688 finished.usage.as_ref(),
2689 Some(&finished.cost),
2690 finished.ttfb_ms,
2691 );
2692 }
2693
2694 if let Some(sid) = finished.session_id.as_deref() {
2695 let mut stats = self.session_stats.write().await;
2696 let entry = stats.entry(sid.to_string()).or_default();
2697 entry.turns_total = entry.turns_total.saturating_add(1);
2698 if finished.session_identity_source.is_some() {
2699 entry.last_session_identity_source = finished.session_identity_source;
2700 }
2701 entry.last_client_name = finished
2702 .client_name
2703 .clone()
2704 .or(entry.last_client_name.clone());
2705 entry.last_client_addr = finished
2706 .client_addr
2707 .clone()
2708 .or(entry.last_client_addr.clone());
2709 entry.last_model = finished.model.clone().or(entry.last_model.clone());
2710 entry.last_reasoning_effort = finished
2711 .reasoning_effort
2712 .clone()
2713 .or(entry.last_reasoning_effort.clone());
2714 entry.last_service_tier = finished
2715 .service_tier
2716 .clone()
2717 .or(entry.last_service_tier.clone());
2718 entry.last_provider_id = finished
2719 .provider_id
2720 .clone()
2721 .or(entry.last_provider_id.clone());
2722 entry.last_station_name = finished
2723 .station_name
2724 .clone()
2725 .or(entry.last_station_name.clone());
2726 if finished.route_decision.is_some() {
2727 entry.last_route_decision = finished.route_decision.clone();
2728 }
2729 if let Some(u) = finished.usage.as_ref() {
2730 entry.last_usage = Some(u.clone());
2731 entry.total_usage.add_assign(u);
2732 entry.turns_with_usage = entry.turns_with_usage.saturating_add(1);
2733 }
2734 entry.last_status = Some(finished.status_code);
2735 entry.last_duration_ms = Some(finished.duration_ms);
2736 entry.last_ended_at_ms = Some(finished.ended_at_ms);
2737 entry.last_seen_ms = finished.ended_at_ms;
2738 }
2739
2740 let mut recent = self.recent_finished.write().await;
2741 recent.push_front(finished);
2742 while recent.len() > recent_finished_max() {
2743 recent.pop_back();
2744 }
2745 self.notify_state_changed();
2746 }
2747
2748 pub async fn list_active_requests(&self) -> Vec<ActiveRequest> {
2749 let guard = self.active_requests.read().await;
2750 let mut vec = guard.values().cloned().collect::<Vec<_>>();
2751 vec.sort_by_key(|r| r.started_at_ms);
2752 vec
2753 }
2754
2755 pub async fn list_recent_finished(&self, limit: usize) -> Vec<FinishedRequest> {
2756 let guard = self.recent_finished.read().await;
2757 guard.iter().take(limit).cloned().collect()
2758 }
2759
2760 pub async fn list_session_stats(&self) -> HashMap<String, SessionStats> {
2761 let guard = self.session_stats.read().await;
2762 guard.clone()
2763 }
2764
2765 pub async fn list_session_identity_cards(
2766 &self,
2767 recent_limit: usize,
2768 ) -> Vec<SessionIdentityCard> {
2769 let recent_limit = recent_limit.clamp(1, recent_finished_max());
2770 let (
2771 active,
2772 recent,
2773 overrides,
2774 station_overrides,
2775 model_overrides,
2776 service_tier_overrides,
2777 bindings,
2778 route_affinities,
2779 global_station_override,
2780 stats,
2781 ) = tokio::join!(
2782 self.list_active_requests(),
2783 self.list_recent_finished(recent_limit),
2784 self.list_session_effort_overrides(),
2785 self.list_session_station_overrides(),
2786 self.list_session_model_overrides(),
2787 self.list_session_service_tier_overrides(),
2788 self.list_session_bindings(),
2789 self.list_session_route_affinities(),
2790 self.get_global_station_override(),
2791 self.list_session_stats(),
2792 );
2793 build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
2794 active: &active,
2795 recent: &recent,
2796 overrides: &overrides,
2797 station_overrides: &station_overrides,
2798 model_overrides: &model_overrides,
2799 service_tier_overrides: &service_tier_overrides,
2800 bindings: &bindings,
2801 route_affinities: &route_affinities,
2802 global_station_override: global_station_override.as_deref(),
2803 stats: &stats,
2804 })
2805 }
2806
2807 async fn resolve_host_transcript_paths_cached(
2808 &self,
2809 session_ids: &[String],
2810 ) -> HashMap<String, Option<String>> {
2811 let mut unique = session_ids
2812 .iter()
2813 .map(|sid| sid.trim())
2814 .filter(|sid| !sid.is_empty())
2815 .map(ToOwned::to_owned)
2816 .collect::<Vec<_>>();
2817 unique.sort();
2818 unique.dedup();
2819 if unique.is_empty() {
2820 return HashMap::new();
2821 }
2822
2823 if self.session_transcript_path_cache_max_entries == 0 {
2824 return sessions::find_codex_session_files_by_ids(&unique)
2825 .await
2826 .unwrap_or_default()
2827 .into_iter()
2828 .map(|(sid, path)| (sid, Some(path.to_string_lossy().to_string())))
2829 .collect();
2830 }
2831
2832 let now_ms = unix_now_ms();
2833 let ttl_ms = self.session_transcript_path_cache_ttl_ms;
2834 let mut resolved = HashMap::<String, Option<String>>::new();
2835 let mut stale_or_missing = Vec::<String>::new();
2836
2837 {
2838 let cache = self.session_transcript_path_cache.read().await;
2839 for sid in &unique {
2840 let fresh = cache.get(sid).filter(|entry| {
2841 ttl_ms == 0 || now_ms.saturating_sub(entry.last_checked_ms) <= ttl_ms
2842 });
2843 if let Some(entry) = fresh {
2844 resolved.insert(sid.clone(), entry.path.clone());
2845 } else {
2846 stale_or_missing.push(sid.clone());
2847 }
2848 }
2849 }
2850
2851 if !stale_or_missing.is_empty() {
2852 let found = sessions::find_codex_session_files_by_ids(&stale_or_missing)
2853 .await
2854 .unwrap_or_default();
2855 let mut cache = self.session_transcript_path_cache.write().await;
2856 for sid in stale_or_missing {
2857 let path = found
2858 .get(&sid)
2859 .map(|path| path.to_string_lossy().to_string());
2860 cache.insert(
2861 sid.clone(),
2862 SessionTranscriptPathCacheEntry {
2863 path: path.clone(),
2864 last_checked_ms: now_ms,
2865 last_seen_ms: now_ms,
2866 },
2867 );
2868 resolved.insert(sid, path);
2869 }
2870 prune_lru_cache(
2871 &mut cache,
2872 self.session_transcript_path_cache_max_entries,
2873 |entry| entry.last_seen_ms,
2874 );
2875 }
2876
2877 {
2878 let mut cache = self.session_transcript_path_cache.write().await;
2879 for sid in &unique {
2880 if let Some(entry) = cache.get_mut(sid) {
2881 entry.last_seen_ms = now_ms;
2882 }
2883 }
2884 }
2885
2886 resolved
2887 }
2888
2889 pub async fn enrich_session_identity_cards_with_cached_host_transcripts(
2890 &self,
2891 cards: &mut [SessionIdentityCard],
2892 ) {
2893 let session_ids = cards
2894 .iter()
2895 .filter_map(|card| card.session_id.as_deref())
2896 .map(str::to_owned)
2897 .collect::<Vec<_>>();
2898 let resolved = self
2899 .resolve_host_transcript_paths_cached(&session_ids)
2900 .await;
2901 for card in cards {
2902 card.host_local_transcript_path = card
2903 .session_id
2904 .as_deref()
2905 .and_then(|sid| resolved.get(sid))
2906 .and_then(Clone::clone);
2907 }
2908 }
2909
2910 pub async fn list_session_identity_cards_with_host_transcripts(
2911 &self,
2912 recent_limit: usize,
2913 ) -> Vec<SessionIdentityCard> {
2914 let mut cards = self.list_session_identity_cards(recent_limit).await;
2915 self.enrich_session_identity_cards_with_cached_host_transcripts(&mut cards)
2916 .await;
2917 cards
2918 }
2919
2920 pub fn spawn_cleanup_task(state: Arc<Self>) {
2921 tokio::spawn(async move {
2923 let mut tick = interval(Duration::from_secs(30));
2924 loop {
2925 tick.tick().await;
2926 state.prune_periodic().await;
2927 }
2928 });
2929 }
2930
2931 async fn prune_periodic(&self) {
2932 let now_ms = std::time::SystemTime::now()
2933 .duration_since(std::time::UNIX_EPOCH)
2934 .map(|d| d.as_millis() as u64)
2935 .unwrap_or(0);
2936
2937 let active = self.active_requests.read().await;
2939 let mut active_sessions: HashMap<String, ()> = HashMap::new();
2940 for req in active.values() {
2941 if let Some(sid) = req.session_id.as_deref() {
2942 active_sessions.insert(sid.to_string(), ());
2943 }
2944 }
2945
2946 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2947 let cutoff_override = now_ms - self.session_override_ttl_ms;
2948 let mut overrides = self.session_effort_overrides.write().await;
2949 overrides.retain(|sid, v| {
2950 if active_sessions.contains_key(sid) {
2951 return true;
2952 }
2953 v.last_seen_ms >= cutoff_override
2954 });
2955 }
2956
2957 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2958 let cutoff_override = now_ms - self.session_override_ttl_ms;
2959 let mut overrides = self.session_station_overrides.write().await;
2960 overrides.retain(|sid, v| {
2961 if active_sessions.contains_key(sid) {
2962 return true;
2963 }
2964 v.last_seen_ms >= cutoff_override
2965 });
2966 }
2967
2968 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2969 let cutoff_override = now_ms - self.session_override_ttl_ms;
2970 let mut overrides = self.session_route_target_overrides.write().await;
2971 overrides.retain(|sid, v| {
2972 if active_sessions.contains_key(sid) {
2973 return true;
2974 }
2975 v.last_seen_ms >= cutoff_override
2976 });
2977 }
2978
2979 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2980 let cutoff_override = now_ms - self.session_override_ttl_ms;
2981 let mut overrides = self.session_model_overrides.write().await;
2982 overrides.retain(|sid, v| {
2983 if active_sessions.contains_key(sid) {
2984 return true;
2985 }
2986 v.last_seen_ms >= cutoff_override
2987 });
2988 }
2989
2990 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
2991 let cutoff_override = now_ms - self.session_override_ttl_ms;
2992 let mut overrides = self.session_service_tier_overrides.write().await;
2993 overrides.retain(|sid, v| {
2994 if active_sessions.contains_key(sid) {
2995 return true;
2996 }
2997 v.last_seen_ms >= cutoff_override
2998 });
2999 }
3000
3001 if self.session_binding_ttl_ms > 0 && now_ms >= self.session_binding_ttl_ms {
3002 let cutoff_binding = now_ms - self.session_binding_ttl_ms;
3003 let mut bindings = self.session_bindings.write().await;
3004 bindings.retain(|sid, entry| {
3005 if active_sessions.contains_key(sid) {
3006 return true;
3007 }
3008 entry.binding.last_seen_ms >= cutoff_binding
3009 });
3010 }
3011 if self.session_binding_max_entries > 0 {
3012 let mut bindings = self.session_bindings.write().await;
3013 if bindings.len() > self.session_binding_max_entries {
3014 let mut removable = bindings
3015 .iter()
3016 .filter(|(sid, _)| !active_sessions.contains_key(*sid))
3017 .map(|(sid, entry)| (sid.clone(), entry.binding.last_seen_ms))
3018 .collect::<Vec<_>>();
3019 removable.sort_by_key(|(_, last_seen_ms)| *last_seen_ms);
3020 let remove_count = bindings
3021 .len()
3022 .saturating_sub(self.session_binding_max_entries)
3023 .min(removable.len());
3024 for (sid, _) in removable.into_iter().take(remove_count) {
3025 bindings.remove(&sid);
3026 }
3027 }
3028 }
3029
3030 {
3031 let mut affinities = self.session_route_affinities.write().await;
3032 if self.session_route_affinity_ttl_ms > 0
3033 && now_ms >= self.session_route_affinity_ttl_ms
3034 {
3035 let cutoff_affinity = now_ms - self.session_route_affinity_ttl_ms;
3036 affinities.retain(|sid, affinity| {
3037 active_sessions.contains_key(sid)
3038 || affinity.last_selected_at_ms >= cutoff_affinity
3039 });
3040 }
3041 prune_lru_cache(
3042 &mut affinities,
3043 self.session_route_affinity_max_entries,
3044 |entry| entry.last_selected_at_ms,
3045 );
3046 }
3047
3048 let keep_days: i32 = std::env::var("CODEX_HELPER_USAGE_ROLLUP_KEEP_DAYS")
3050 .ok()
3051 .and_then(|s| s.trim().parse::<i32>().ok())
3052 .filter(|&n| n > 0)
3053 .unwrap_or(60);
3054 let now_day = (now_ms / 86_400_000) as i32;
3055 let cutoff_day = now_day.saturating_sub(keep_days);
3056 let mut rollups = self.usage_rollups.write().await;
3057 for rollup in rollups.values_mut() {
3058 rollup.by_day.retain(|day, _| *day >= cutoff_day);
3059 rollup.by_config_day.retain(|_, m| {
3060 m.retain(|day, _| *day >= cutoff_day);
3061 !m.is_empty()
3062 });
3063 rollup.by_provider_day.retain(|_, m| {
3064 m.retain(|day, _| *day >= cutoff_day);
3065 !m.is_empty()
3066 });
3067 }
3068
3069 let cutoff_cwd =
3070 if self.session_cwd_cache_ttl_ms == 0 || now_ms < self.session_cwd_cache_ttl_ms {
3071 0
3072 } else {
3073 now_ms - self.session_cwd_cache_ttl_ms
3074 };
3075 self.prune_session_cwd_cache(&active_sessions, cutoff_cwd)
3076 .await;
3077 self.prune_session_transcript_path_cache(now_ms).await;
3078
3079 if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
3080 let cutoff_stats = now_ms - self.session_override_ttl_ms;
3081 let mut stats = self.session_stats.write().await;
3082 stats.retain(|sid, v| {
3083 active_sessions.contains_key(sid) || v.last_seen_ms >= cutoff_stats
3084 });
3085 }
3086 }
3087
3088 async fn prune_session_cwd_cache(&self, active_sessions: &HashMap<String, ()>, cutoff: u64) {
3089 if self.session_cwd_cache_max_entries == 0 {
3090 return;
3091 }
3092 let mut cache = self.session_cwd_cache.write().await;
3093
3094 if self.session_cwd_cache_ttl_ms > 0 {
3095 cache.retain(|sid, v| {
3096 if active_sessions.contains_key(sid) {
3097 return true;
3098 }
3099 v.last_seen_ms >= cutoff
3100 });
3101 }
3102
3103 let max = self.session_cwd_cache_max_entries;
3104 if max == 0 || cache.len() <= max {
3105 return;
3106 }
3107
3108 let mut keys = cache
3110 .iter()
3111 .map(|(sid, v)| (sid.clone(), v.last_seen_ms))
3112 .collect::<Vec<_>>();
3113 keys.sort_by_key(|(_, t)| *t);
3114 let remove_count = keys.len().saturating_sub(max);
3115 for (sid, _) in keys.into_iter().take(remove_count) {
3116 cache.remove(&sid);
3117 }
3118 }
3119
3120 async fn prune_session_transcript_path_cache(&self, now_ms: u64) {
3121 let mut cache = self.session_transcript_path_cache.write().await;
3122 if self.session_transcript_path_cache_max_entries == 0 {
3123 cache.clear();
3124 return;
3125 }
3126
3127 if self.session_transcript_path_cache_ttl_ms > 0 {
3128 let cutoff = now_ms.saturating_sub(self.session_transcript_path_cache_ttl_ms);
3129 cache.retain(|_, entry| entry.last_seen_ms >= cutoff);
3130 }
3131
3132 prune_lru_cache(
3133 &mut cache,
3134 self.session_transcript_path_cache_max_entries,
3135 |entry| entry.last_seen_ms,
3136 );
3137 }
3138}
3139
3140#[cfg(test)]
3141#[derive(Debug)]
3142pub(crate) struct BeginRequestTestBuilder<'a> {
3143 state: &'a ProxyState,
3144 service: &'static str,
3145 method: &'static str,
3146 path: &'static str,
3147 session_id: Option<String>,
3148 session_identity_source: Option<SessionIdentitySource>,
3149 client_name: Option<String>,
3150 client_addr: Option<String>,
3151 cwd: Option<String>,
3152 model: Option<String>,
3153 reasoning_effort: Option<String>,
3154 service_tier: Option<String>,
3155 started_at_ms: u64,
3156}
3157
3158#[cfg(test)]
3159impl<'a> BeginRequestTestBuilder<'a> {
3160 pub(crate) fn new(state: &'a ProxyState) -> Self {
3161 Self {
3162 state,
3163 service: "codex",
3164 method: "POST",
3165 path: "/v1/responses",
3166 session_id: None,
3167 session_identity_source: None,
3168 client_name: None,
3169 client_addr: None,
3170 cwd: None,
3171 model: None,
3172 reasoning_effort: None,
3173 service_tier: None,
3174 started_at_ms: 0,
3175 }
3176 }
3177
3178 pub(crate) fn session_id(mut self, value: impl Into<String>) -> Self {
3179 self.session_id = Some(value.into());
3180 self
3181 }
3182
3183 pub(crate) fn model(mut self, value: impl Into<String>) -> Self {
3184 self.model = Some(value.into());
3185 self
3186 }
3187
3188 pub(crate) fn service_tier(mut self, value: impl Into<String>) -> Self {
3189 self.service_tier = Some(value.into());
3190 self
3191 }
3192
3193 pub(crate) fn started_at_ms(mut self, value: u64) -> Self {
3194 self.started_at_ms = value;
3195 self
3196 }
3197
3198 pub(crate) async fn begin(self) -> u64 {
3199 self.state
3200 .begin_request(
3201 self.service,
3202 self.method,
3203 self.path,
3204 self.session_id,
3205 self.session_identity_source,
3206 self.client_name,
3207 self.client_addr,
3208 self.cwd,
3209 self.model,
3210 self.reasoning_effort,
3211 self.service_tier,
3212 self.started_at_ms,
3213 )
3214 .await
3215 }
3216}
3217
3218#[cfg(test)]
3219mod tests {
3220 use super::*;
3221
3222 use crate::config::{ServiceConfig, ServiceConfigManager, UpstreamAuth, UpstreamConfig};
3223 use crate::runtime_identity::ProviderEndpointKey;
3224 use std::path::Path;
3225 use std::sync::OnceLock;
3226
3227 async fn env_lock() -> tokio::sync::MutexGuard<'static, ()> {
3228 static LOCK: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
3229 LOCK.get_or_init(|| tokio::sync::Mutex::new(()))
3230 .lock()
3231 .await
3232 }
3233
3234 #[derive(Default)]
3235 struct ScopedEnv {
3236 saved: Vec<(String, Option<String>)>,
3237 }
3238
3239 impl ScopedEnv {
3240 unsafe fn set(&mut self, key: &str, value: &str) {
3241 if !self.saved.iter().any(|(saved_key, _)| saved_key == key) {
3242 self.saved.push((key.to_string(), std::env::var(key).ok()));
3243 }
3244 unsafe {
3245 std::env::set_var(key, value);
3246 }
3247 }
3248
3249 unsafe fn set_path(&mut self, key: &str, value: &Path) {
3250 unsafe {
3251 self.set(key, value.to_string_lossy().as_ref());
3252 }
3253 }
3254 }
3255
3256 impl Drop for ScopedEnv {
3257 fn drop(&mut self) {
3258 for (key, value) in self.saved.iter().rev() {
3259 match value {
3260 Some(value) => unsafe {
3261 std::env::set_var(key, value);
3262 },
3263 None => unsafe {
3264 std::env::remove_var(key);
3265 },
3266 }
3267 }
3268 }
3269 }
3270
3271 fn test_runtime_policy(
3272 session_override_ttl_ms: u64,
3273 session_binding_ttl_ms: u64,
3274 session_binding_max_entries: usize,
3275 ) -> RuntimePolicy {
3276 RuntimePolicy {
3277 session_override_ttl_ms,
3278 session_binding_ttl_ms,
3279 session_binding_max_entries,
3280 session_route_affinity_ttl_ms: 0,
3281 session_route_affinity_max_entries: 5_000,
3282 session_route_affinity_store: SessionRouteAffinityStore::from_env(),
3283 session_cwd_cache_ttl_ms: 0,
3284 session_cwd_cache_max_entries: 0,
3285 session_transcript_path_cache_ttl_ms: 30_000,
3286 session_transcript_path_cache_max_entries: 5_000,
3287 }
3288 }
3289
3290 #[test]
3291 fn begin_and_finish_requests_keep_trace_id() {
3292 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3293 runtime.block_on(async {
3294 let state = ProxyState::new();
3295 let request_id = state
3296 .begin_request_for_test()
3297 .model("gpt-5")
3298 .service_tier("priority")
3299 .started_at_ms(100)
3300 .begin()
3301 .await;
3302
3303 let active = state.list_active_requests().await;
3304 assert_eq!(active[0].trace_id.as_deref(), Some("codex-1"));
3305
3306 state
3307 .finish_request(FinishRequestParams {
3308 id: request_id,
3309 status_code: 200,
3310 duration_ms: 10,
3311 ended_at_ms: 110,
3312 observed_service_tier: Some("priority".to_string()),
3313 usage: None,
3314 retry: None,
3315 ttfb_ms: Some(4),
3316 streaming: false,
3317 })
3318 .await;
3319
3320 let recent = state.list_recent_finished(1).await;
3321 assert_eq!(recent[0].trace_id.as_deref(), Some("codex-1"));
3322 assert_eq!(recent[0].observability.trace_id.as_deref(), Some("codex-1"));
3323 assert!(recent[0].observability.fast_mode);
3324 assert_eq!(recent[0].observability.generation_ms, Some(6));
3325 });
3326 }
3327
3328 #[test]
3329 fn state_change_subscription_tracks_dashboard_mutations() {
3330 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3331 runtime.block_on(async {
3332 let state = ProxyState::new();
3333 let mut changes = state.subscribe_state_changes();
3334
3335 state
3336 .record_provider_balance_snapshot(
3337 "codex",
3338 ProviderBalanceSnapshot {
3339 provider_id: "input6".to_string(),
3340 station_name: Some("routing".to_string()),
3341 upstream_index: Some(0),
3342 fetched_at_ms: 100,
3343 stale_after_ms: Some(1_000),
3344 status: BalanceSnapshotStatus::Ok,
3345 total_balance_usd: Some("12.50".to_string()),
3346 ..ProviderBalanceSnapshot::default()
3347 },
3348 )
3349 .await;
3350 changes.changed().await.expect("balance change");
3351 let balance_version = *changes.borrow();
3352 assert!(balance_version > 0);
3353
3354 let request_id = state
3355 .begin_request_for_test()
3356 .session_id("sid-1")
3357 .model("gpt-5")
3358 .started_at_ms(200)
3359 .begin()
3360 .await;
3361 changes.changed().await.expect("begin request change");
3362 let active_version = *changes.borrow();
3363 assert!(active_version > balance_version);
3364
3365 state
3366 .finish_request(FinishRequestParams {
3367 id: request_id,
3368 status_code: 200,
3369 duration_ms: 25,
3370 ended_at_ms: 250,
3371 observed_service_tier: None,
3372 usage: None,
3373 retry: None,
3374 ttfb_ms: Some(5),
3375 streaming: false,
3376 })
3377 .await;
3378 changes.changed().await.expect("finish request change");
3379 assert!(*changes.borrow() > active_version);
3380 });
3381 }
3382
3383 #[test]
3384 fn state_change_subscription_ignores_session_touch_only_updates() {
3385 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3386 runtime.block_on(async {
3387 let state = ProxyState::new();
3388 let mut changes = state.subscribe_state_changes();
3389
3390 state
3391 .set_session_model_override("sid-1".to_string(), "gpt-5".to_string(), 100)
3392 .await;
3393 changes.changed().await.expect("model override change");
3394 let version = *changes.borrow();
3395
3396 state.touch_session_model_override("sid-1", 200).await;
3397 assert_eq!(*changes.borrow(), version);
3398 });
3399 }
3400
3401 #[test]
3402 fn finish_request_estimates_cost_and_rolls_up_cost() {
3403 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3404 runtime.block_on(async {
3405 let state = ProxyState::new();
3406 let request_id = state
3407 .begin_request_for_test()
3408 .model("gpt-5")
3409 .started_at_ms(100)
3410 .begin()
3411 .await;
3412
3413 state
3414 .finish_request(FinishRequestParams {
3415 id: request_id,
3416 status_code: 200,
3417 duration_ms: 10,
3418 ended_at_ms: 110,
3419 observed_service_tier: None,
3420 usage: Some(UsageMetrics {
3421 input_tokens: 1_000,
3422 output_tokens: 500,
3423 cached_input_tokens: 100,
3424 total_tokens: 1_500,
3425 ..UsageMetrics::default()
3426 }),
3427 retry: None,
3428 ttfb_ms: Some(4),
3429 streaming: false,
3430 })
3431 .await;
3432
3433 let recent = state.list_recent_finished(1).await;
3434 assert_eq!(recent[0].cost.total_cost_usd.as_deref(), Some("0.0061375"));
3435
3436 let rollup = state.get_usage_rollup_view("codex", 12, 1).await;
3437 assert_eq!(
3438 rollup.loaded.cost.total_cost_usd.as_deref(),
3439 Some("0.0061375")
3440 );
3441 assert_eq!(rollup.loaded.cost.priced_requests, 1);
3442 assert_eq!(rollup.loaded.cost.unpriced_requests, 0);
3443 });
3444 }
3445
3446 #[test]
3447 fn usage_rollup_view_scores_entities_inside_selected_window() {
3448 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3449 runtime.block_on(async {
3450 let state = ProxyState::new();
3451 let now_ms = std::time::SystemTime::now()
3452 .duration_since(std::time::UNIX_EPOCH)
3453 .map(|d| d.as_millis() as u64)
3454 .unwrap_or(0);
3455 let old_ms = now_ms.saturating_sub(10 * 86_400_000);
3456
3457 let old_id = state
3458 .begin_request_for_test()
3459 .model("gpt-5")
3460 .started_at_ms(old_ms.saturating_sub(1_000))
3461 .begin()
3462 .await;
3463 state
3464 .update_request_route(
3465 old_id,
3466 Some("old-station".to_string()),
3467 Some("old-provider".to_string()),
3468 "https://old.example/v1".to_string(),
3469 None,
3470 )
3471 .await;
3472 state
3473 .finish_request(FinishRequestParams {
3474 id: old_id,
3475 status_code: 200,
3476 duration_ms: 20,
3477 ended_at_ms: old_ms,
3478 observed_service_tier: None,
3479 usage: Some(UsageMetrics {
3480 total_tokens: 100_000,
3481 ..UsageMetrics::default()
3482 }),
3483 retry: None,
3484 ttfb_ms: Some(5),
3485 streaming: false,
3486 })
3487 .await;
3488
3489 let fresh_id = state
3490 .begin_request_for_test()
3491 .model("gpt-5")
3492 .started_at_ms(now_ms.saturating_sub(1_000))
3493 .begin()
3494 .await;
3495 state
3496 .update_request_route(
3497 fresh_id,
3498 Some("fresh-station".to_string()),
3499 Some("fresh-provider".to_string()),
3500 "https://fresh.example/v1".to_string(),
3501 None,
3502 )
3503 .await;
3504 state
3505 .finish_request(FinishRequestParams {
3506 id: fresh_id,
3507 status_code: 200,
3508 duration_ms: 10,
3509 ended_at_ms: now_ms,
3510 observed_service_tier: None,
3511 usage: Some(UsageMetrics {
3512 total_tokens: 10,
3513 ..UsageMetrics::default()
3514 }),
3515 retry: None,
3516 ttfb_ms: Some(3),
3517 streaming: false,
3518 })
3519 .await;
3520
3521 let week = state.get_usage_rollup_view("codex", 10, 7).await;
3522 assert_eq!(week.loaded.requests_total, 2);
3523 assert_eq!(week.window.requests_total, 1);
3524 assert_eq!(week.by_day.len(), 7);
3525 assert_eq!(week.by_config[0].0, "fresh-station");
3526 assert_eq!(week.by_provider[0].0, "fresh-provider");
3527
3528 let loaded = state.get_usage_rollup_view("codex", 10, 0).await;
3529 assert_eq!(loaded.window.requests_total, 2);
3530 assert_eq!(loaded.by_config[0].0, "old-station");
3531 assert_eq!(loaded.by_provider[0].0, "old-provider");
3532 });
3533 }
3534
3535 #[test]
3536 fn build_session_identity_cards_merges_sources_and_sorts_newest_first() {
3537 let active = vec![ActiveRequest {
3538 id: 1,
3539 trace_id: Some("codex-1".to_string()),
3540 session_id: Some("sid-active".to_string()),
3541 session_identity_source: Some(SessionIdentitySource::Header),
3542 client_name: Some("Frank-Laptop".to_string()),
3543 client_addr: Some("100.64.0.8".to_string()),
3544 cwd: Some("G:/codes/project".to_string()),
3545 model: Some("gpt-5.4".to_string()),
3546 reasoning_effort: Some("medium".to_string()),
3547 service_tier: Some("priority".to_string()),
3548 station_name: Some("right".to_string()),
3549 provider_id: Some("right".to_string()),
3550 upstream_base_url: Some("https://right.example/v1".to_string()),
3551 route_decision: None,
3552 service: "codex".to_string(),
3553 method: "POST".to_string(),
3554 path: "/v1/responses".to_string(),
3555 started_at_ms: 500,
3556 }];
3557 let recent = vec![
3558 FinishedRequest {
3559 id: 2,
3560 trace_id: Some("codex-2".to_string()),
3561 session_id: Some("sid-recent".to_string()),
3562 session_identity_source: Some(SessionIdentitySource::PromptCacheKey),
3563 client_name: Some("Studio-Mini".to_string()),
3564 client_addr: Some("100.64.0.9".to_string()),
3565 cwd: Some("G:/codes/other".to_string()),
3566 model: Some("gpt-5.3".to_string()),
3567 reasoning_effort: Some("high".to_string()),
3568 service_tier: Some("default".to_string()),
3569 station_name: Some("vibe".to_string()),
3570 provider_id: Some("vibe".to_string()),
3571 upstream_base_url: Some("https://vibe.example/v1".to_string()),
3572 route_decision: None,
3573 usage: Some(UsageMetrics {
3574 input_tokens: 1,
3575 output_tokens: 2,
3576 reasoning_tokens: 3,
3577 total_tokens: 6,
3578 ..UsageMetrics::default()
3579 }),
3580 cost: CostBreakdown::default(),
3581 retry: None,
3582 observability: RequestObservability::default(),
3583 service: "codex".to_string(),
3584 method: "POST".to_string(),
3585 path: "/v1/responses".to_string(),
3586 status_code: 200,
3587 duration_ms: 1200,
3588 ttfb_ms: Some(100),
3589 streaming: false,
3590 ended_at_ms: 2_000,
3591 },
3592 FinishedRequest {
3593 id: 3,
3594 trace_id: Some("codex-3".to_string()),
3595 session_id: Some("sid-active".to_string()),
3596 session_identity_source: Some(SessionIdentitySource::Header),
3597 client_name: Some("Frank-Laptop".to_string()),
3598 client_addr: Some("100.64.0.8".to_string()),
3599 cwd: Some("G:/codes/project".to_string()),
3600 model: Some("gpt-5.4".to_string()),
3601 reasoning_effort: Some("low".to_string()),
3602 service_tier: Some("flex".to_string()),
3603 station_name: Some("right".to_string()),
3604 provider_id: Some("right".to_string()),
3605 upstream_base_url: Some("https://right.example/v1".to_string()),
3606 route_decision: None,
3607 usage: None,
3608 cost: CostBreakdown::default(),
3609 retry: None,
3610 observability: RequestObservability::default(),
3611 service: "codex".to_string(),
3612 method: "POST".to_string(),
3613 path: "/v1/responses".to_string(),
3614 status_code: 429,
3615 duration_ms: 900,
3616 ttfb_ms: None,
3617 streaming: false,
3618 ended_at_ms: 1_000,
3619 },
3620 ];
3621 let overrides = HashMap::from([("sid-active".to_string(), "xhigh".to_string())]);
3622 let config_overrides = HashMap::from([("sid-active".to_string(), "temp".to_string())]);
3623 let model_overrides =
3624 HashMap::from([("sid-active".to_string(), "gpt-5.4-mini".to_string())]);
3625 let service_tier_overrides =
3626 HashMap::from([("sid-active".to_string(), "priority".to_string())]);
3627 let stats = HashMap::from([(
3628 "sid-active".to_string(),
3629 SessionStats {
3630 turns_total: 3,
3631 last_session_identity_source: Some(SessionIdentitySource::Header),
3632 last_client_name: Some("Frank-Laptop".to_string()),
3633 last_client_addr: Some("100.64.0.8".to_string()),
3634 last_model: Some("gpt-5.4".to_string()),
3635 last_reasoning_effort: Some("low".to_string()),
3636 last_service_tier: Some("flex".to_string()),
3637 last_provider_id: Some("right".to_string()),
3638 last_station_name: Some("right".to_string()),
3639 last_route_decision: None,
3640 last_usage: None,
3641 total_usage: UsageMetrics {
3642 input_tokens: 10,
3643 output_tokens: 20,
3644 reasoning_tokens: 5,
3645 total_tokens: 35,
3646 ..UsageMetrics::default()
3647 },
3648 turns_with_usage: 2,
3649 last_status: Some(429),
3650 last_duration_ms: Some(900),
3651 last_ended_at_ms: Some(1_000),
3652 last_seen_ms: 1_000,
3653 },
3654 )]);
3655
3656 let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
3657 active: &active,
3658 recent: &recent,
3659 overrides: &overrides,
3660 station_overrides: &config_overrides,
3661 model_overrides: &model_overrides,
3662 service_tier_overrides: &service_tier_overrides,
3663 bindings: &HashMap::new(),
3664 route_affinities: &HashMap::new(),
3665 global_station_override: None,
3666 stats: &stats,
3667 });
3668
3669 assert_eq!(cards.len(), 2);
3670 assert_eq!(cards[0].session_id.as_deref(), Some("sid-recent"));
3671 assert_eq!(
3672 cards[0].session_identity_source,
3673 Some(SessionIdentitySource::PromptCacheKey)
3674 );
3675 assert_eq!(
3676 cards[0].observation_scope,
3677 SessionObservationScope::HostLocalEnriched
3678 );
3679 assert_eq!(cards[0].last_client_name.as_deref(), Some("Studio-Mini"));
3680 assert_eq!(cards[0].last_client_addr.as_deref(), Some("100.64.0.9"));
3681 assert_eq!(cards[1].session_id.as_deref(), Some("sid-active"));
3682 assert_eq!(
3683 cards[1].session_identity_source,
3684 Some(SessionIdentitySource::Header)
3685 );
3686 assert_eq!(
3687 cards[1].observation_scope,
3688 SessionObservationScope::HostLocalEnriched
3689 );
3690 assert_eq!(cards[1].active_count, 1);
3691 assert_eq!(cards[1].last_client_name.as_deref(), Some("Frank-Laptop"));
3692 assert_eq!(cards[1].last_client_addr.as_deref(), Some("100.64.0.8"));
3693 assert_eq!(cards[1].last_status, Some(429));
3694 assert_eq!(cards[1].override_effort.as_deref(), Some("xhigh"));
3695 assert_eq!(cards[1].override_station_name.as_deref(), Some("temp"));
3696 assert_eq!(cards[1].override_model.as_deref(), Some("gpt-5.4-mini"));
3697 assert_eq!(cards[1].override_service_tier.as_deref(), Some("priority"));
3698 assert_eq!(
3699 cards[1]
3700 .effective_model
3701 .as_ref()
3702 .map(|value| value.value.as_str()),
3703 Some("gpt-5.4-mini")
3704 );
3705 assert_eq!(
3706 cards[1].effective_model.as_ref().map(|value| value.source),
3707 Some(RouteValueSource::SessionOverride)
3708 );
3709 assert_eq!(
3710 cards[1]
3711 .effective_reasoning_effort
3712 .as_ref()
3713 .map(|value| value.source),
3714 Some(RouteValueSource::SessionOverride)
3715 );
3716 assert_eq!(
3717 cards[1]
3718 .effective_service_tier
3719 .as_ref()
3720 .map(|value| value.source),
3721 Some(RouteValueSource::SessionOverride)
3722 );
3723 assert_eq!(
3724 cards[1]
3725 .effective_station
3726 .as_ref()
3727 .map(|value| value.source),
3728 Some(RouteValueSource::SessionOverride)
3729 );
3730 assert!(cards[1].effective_upstream_base_url.is_none());
3731 assert_eq!(
3732 cards[1].last_upstream_base_url.as_deref(),
3733 Some("https://right.example/v1")
3734 );
3735 assert_eq!(cards[1].turns_total, Some(3));
3736 assert_eq!(cards[1].last_service_tier.as_deref(), Some("flex"));
3737 assert_eq!(
3738 cards[1].total_usage.as_ref().map(|u| u.total_tokens),
3739 Some(35)
3740 );
3741 }
3742
3743 #[test]
3744 fn build_session_identity_cards_default_profile_keeps_request_fields() {
3745 let active = vec![ActiveRequest {
3746 id: 1,
3747 trace_id: Some("codex-1".to_string()),
3748 session_id: Some("sid-bound".to_string()),
3749 session_identity_source: Some(SessionIdentitySource::Header),
3750 client_name: Some("Workstation".to_string()),
3751 client_addr: Some("100.64.0.10".to_string()),
3752 cwd: None,
3753 model: Some("gpt-observed".to_string()),
3754 reasoning_effort: Some("medium".to_string()),
3755 service_tier: Some("default".to_string()),
3756 station_name: Some("right".to_string()),
3757 provider_id: None,
3758 upstream_base_url: None,
3759 route_decision: None,
3760 service: "codex".to_string(),
3761 method: "POST".to_string(),
3762 path: "/v1/responses".to_string(),
3763 started_at_ms: 10,
3764 }];
3765 let bindings = HashMap::from([(
3766 "sid-bound".to_string(),
3767 SessionBinding {
3768 session_id: "sid-bound".to_string(),
3769 profile_name: Some("daily".to_string()),
3770 station_name: Some("vibe".to_string()),
3771 model: Some("gpt-bound".to_string()),
3772 reasoning_effort: Some("high".to_string()),
3773 service_tier: Some("priority".to_string()),
3774 continuity_mode: SessionContinuityMode::DefaultProfile,
3775 created_at_ms: 1,
3776 updated_at_ms: 1,
3777 last_seen_ms: 10,
3778 },
3779 )]);
3780
3781 let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
3782 active: &active,
3783 recent: &[],
3784 overrides: &HashMap::new(),
3785 station_overrides: &HashMap::new(),
3786 model_overrides: &HashMap::new(),
3787 service_tier_overrides: &HashMap::new(),
3788 bindings: &bindings,
3789 route_affinities: &HashMap::new(),
3790 global_station_override: None,
3791 stats: &HashMap::new(),
3792 });
3793
3794 assert_eq!(cards[0].binding_profile_name.as_deref(), Some("daily"));
3795 assert_eq!(
3796 cards[0].observation_scope,
3797 SessionObservationScope::ObservedOnly
3798 );
3799 assert_eq!(
3800 cards[0].binding_continuity_mode,
3801 Some(SessionContinuityMode::DefaultProfile)
3802 );
3803 assert_eq!(
3804 cards[0]
3805 .effective_model
3806 .as_ref()
3807 .map(|value| (value.value.as_str(), value.source)),
3808 Some(("gpt-observed", RouteValueSource::RequestPayload))
3809 );
3810 assert_eq!(
3811 cards[0]
3812 .effective_reasoning_effort
3813 .as_ref()
3814 .map(|value| (value.value.as_str(), value.source)),
3815 Some(("medium", RouteValueSource::RequestPayload))
3816 );
3817 assert_eq!(
3818 cards[0]
3819 .effective_service_tier
3820 .as_ref()
3821 .map(|value| (value.value.as_str(), value.source)),
3822 Some(("default", RouteValueSource::RequestPayload))
3823 );
3824 assert_eq!(
3825 cards[0]
3826 .effective_station
3827 .as_ref()
3828 .map(|value| (value.value.as_str(), value.source)),
3829 Some(("vibe", RouteValueSource::ProfileDefault))
3830 );
3831 }
3832
3833 #[test]
3834 fn build_session_identity_cards_keeps_request_fields_but_allows_global_config_override() {
3835 let active = vec![ActiveRequest {
3836 id: 1,
3837 trace_id: Some("codex-1".to_string()),
3838 session_id: Some("sid-bound".to_string()),
3839 session_identity_source: Some(SessionIdentitySource::Header),
3840 client_name: Some("Workstation".to_string()),
3841 client_addr: Some("100.64.0.10".to_string()),
3842 cwd: None,
3843 model: Some("gpt-observed".to_string()),
3844 reasoning_effort: Some("medium".to_string()),
3845 service_tier: Some("default".to_string()),
3846 station_name: Some("vibe".to_string()),
3847 provider_id: None,
3848 upstream_base_url: Some("https://vibe.example/v1".to_string()),
3849 route_decision: None,
3850 service: "codex".to_string(),
3851 method: "POST".to_string(),
3852 path: "/v1/responses".to_string(),
3853 started_at_ms: 10,
3854 }];
3855 let bindings = HashMap::from([(
3856 "sid-bound".to_string(),
3857 SessionBinding {
3858 session_id: "sid-bound".to_string(),
3859 profile_name: Some("daily".to_string()),
3860 station_name: Some("vibe".to_string()),
3861 model: Some("gpt-bound".to_string()),
3862 reasoning_effort: Some("high".to_string()),
3863 service_tier: Some("priority".to_string()),
3864 continuity_mode: SessionContinuityMode::DefaultProfile,
3865 created_at_ms: 1,
3866 updated_at_ms: 1,
3867 last_seen_ms: 10,
3868 },
3869 )]);
3870
3871 let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
3872 active: &active,
3873 recent: &[],
3874 overrides: &HashMap::new(),
3875 station_overrides: &HashMap::new(),
3876 model_overrides: &HashMap::new(),
3877 service_tier_overrides: &HashMap::new(),
3878 bindings: &bindings,
3879 route_affinities: &HashMap::new(),
3880 global_station_override: Some("right"),
3881 stats: &HashMap::new(),
3882 });
3883
3884 assert_eq!(cards[0].binding_profile_name.as_deref(), Some("daily"));
3885 assert_eq!(
3886 cards[0]
3887 .effective_model
3888 .as_ref()
3889 .map(|value| (value.value.as_str(), value.source)),
3890 Some(("gpt-observed", RouteValueSource::RequestPayload))
3891 );
3892 assert_eq!(
3893 cards[0]
3894 .effective_reasoning_effort
3895 .as_ref()
3896 .map(|value| (value.value.as_str(), value.source)),
3897 Some(("medium", RouteValueSource::RequestPayload))
3898 );
3899 assert_eq!(
3900 cards[0]
3901 .effective_service_tier
3902 .as_ref()
3903 .map(|value| (value.value.as_str(), value.source)),
3904 Some(("default", RouteValueSource::RequestPayload))
3905 );
3906 assert_eq!(
3907 cards[0]
3908 .effective_station
3909 .as_ref()
3910 .map(|value| (value.value.as_str(), value.source)),
3911 Some(("right", RouteValueSource::GlobalOverride))
3912 );
3913 assert!(cards[0].effective_upstream_base_url.is_none());
3914 }
3915
3916 #[test]
3917 fn enrich_session_identity_cards_with_runtime_applies_station_mapping_and_single_upstream() {
3918 let mut cards = vec![SessionIdentityCard {
3919 session_id: Some("sid-1".to_string()),
3920 last_model: Some("gpt-5.4".to_string()),
3921 last_station_name: Some("right".to_string()),
3922 last_upstream_base_url: Some("https://right.example/v1".to_string()),
3923 effective_model: Some(ResolvedRouteValue::new(
3924 "gpt-5.4",
3925 RouteValueSource::RequestPayload,
3926 )),
3927 effective_station: Some(ResolvedRouteValue::new(
3928 "right",
3929 RouteValueSource::RuntimeFallback,
3930 )),
3931 ..SessionIdentityCard::default()
3932 }];
3933
3934 let mut mgr = ServiceConfigManager {
3935 active: Some("right".to_string()),
3936 ..ServiceConfigManager::default()
3937 };
3938 mgr.configs.insert(
3939 "right".to_string(),
3940 ServiceConfig {
3941 name: "right".to_string(),
3942 alias: None,
3943 enabled: true,
3944 level: 1,
3945 upstreams: vec![UpstreamConfig {
3946 base_url: "https://right.example/v1".to_string(),
3947 auth: UpstreamAuth::default(),
3948 tags: HashMap::new(),
3949 supported_models: HashMap::new(),
3950 model_mapping: HashMap::from([(
3951 "gpt-5.4".to_string(),
3952 "gpt-5.4-fast".to_string(),
3953 )]),
3954 }],
3955 },
3956 );
3957
3958 enrich_session_identity_cards_with_runtime(&mut cards, &mgr);
3959
3960 assert_eq!(
3961 cards[0]
3962 .effective_model
3963 .as_ref()
3964 .map(|value| value.value.as_str()),
3965 Some("gpt-5.4-fast")
3966 );
3967 assert_eq!(
3968 cards[0].effective_model.as_ref().map(|value| value.source),
3969 Some(RouteValueSource::StationMapping)
3970 );
3971 assert_eq!(
3972 cards[0]
3973 .effective_upstream_base_url
3974 .as_ref()
3975 .map(|value| value.value.as_str()),
3976 Some("https://right.example/v1")
3977 );
3978 }
3979
3980 #[test]
3981 fn apply_session_profile_binding_sets_binding_and_clears_manual_overrides() {
3982 let runtime = tokio::runtime::Runtime::new().expect("runtime");
3983 runtime.block_on(async {
3984 let state = ProxyState::new();
3985 let now_ms = 42;
3986 let mut mgr = ServiceConfigManager::default();
3987 mgr.configs.insert(
3988 "right".to_string(),
3989 ServiceConfig {
3990 name: "right".to_string(),
3991 alias: None,
3992 enabled: true,
3993 level: 1,
3994 upstreams: vec![UpstreamConfig {
3995 base_url: "https://right.example/v1".to_string(),
3996 auth: UpstreamAuth::default(),
3997 tags: HashMap::from([
3998 ("supports_reasoning_effort".to_string(), "true".to_string()),
3999 ("supports_service_tier".to_string(), "true".to_string()),
4000 ]),
4001 supported_models: HashMap::from([("gpt-5.4".to_string(), true)]),
4002 model_mapping: HashMap::new(),
4003 }],
4004 },
4005 );
4006 mgr.profiles.insert(
4007 "fast".to_string(),
4008 crate::config::ServiceControlProfile {
4009 extends: None,
4010 station: Some("right".to_string()),
4011 model: Some("gpt-5.4".to_string()),
4012 reasoning_effort: Some("low".to_string()),
4013 service_tier: Some("flex".to_string()),
4014 },
4015 );
4016
4017 state
4018 .set_session_station_override("sid-1".to_string(), "other".to_string(), 1)
4019 .await;
4020 state
4021 .set_session_model_override("sid-1".to_string(), "gpt-x".to_string(), 1)
4022 .await;
4023 state
4024 .set_session_effort_override("sid-1".to_string(), "high".to_string(), 1)
4025 .await;
4026 state
4027 .set_session_service_tier_override("sid-1".to_string(), "priority".to_string(), 1)
4028 .await;
4029
4030 state
4031 .apply_session_profile_binding(
4032 "codex",
4033 &mgr,
4034 "sid-1".to_string(),
4035 "fast".to_string(),
4036 now_ms,
4037 )
4038 .await
4039 .expect("apply profile");
4040
4041 let binding = state
4042 .get_session_binding("sid-1")
4043 .await
4044 .expect("binding exists");
4045 assert_eq!(binding.profile_name.as_deref(), Some("fast"));
4046 assert_eq!(binding.station_name.as_deref(), Some("right"));
4047 assert_eq!(binding.model.as_deref(), Some("gpt-5.4"));
4048 assert_eq!(binding.reasoning_effort.as_deref(), Some("low"));
4049 assert_eq!(binding.service_tier.as_deref(), Some("flex"));
4050 assert_eq!(
4051 binding.continuity_mode,
4052 SessionContinuityMode::ManualProfile
4053 );
4054 assert_eq!(binding.updated_at_ms, now_ms);
4055 assert!(state.get_session_station_override("sid-1").await.is_none());
4056 assert!(state.get_session_model_override("sid-1").await.is_none());
4057 assert!(state.get_session_effort_override("sid-1").await.is_none());
4058 assert!(
4059 state
4060 .get_session_service_tier_override("sid-1")
4061 .await
4062 .is_none()
4063 );
4064 });
4065 }
4066
4067 #[test]
4068 fn list_session_manual_overrides_merges_all_dimensions() {
4069 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4070 runtime.block_on(async {
4071 let state = ProxyState::new();
4072 state
4073 .set_session_reasoning_effort_override("sid-1".to_string(), "high".to_string(), 1)
4074 .await;
4075 state
4076 .set_session_station_override("sid-1".to_string(), "right".to_string(), 1)
4077 .await;
4078 state
4079 .set_session_model_override("sid-1".to_string(), "gpt-5.4".to_string(), 1)
4080 .await;
4081 state
4082 .set_session_service_tier_override("sid-1".to_string(), "priority".to_string(), 1)
4083 .await;
4084 state
4085 .set_session_model_override("sid-2".to_string(), "gpt-5.4-mini".to_string(), 2)
4086 .await;
4087
4088 let merged = state.list_session_manual_overrides().await;
4089 assert_eq!(merged.len(), 2);
4090 assert_eq!(
4091 merged
4092 .get("sid-1")
4093 .and_then(|overrides| overrides.reasoning_effort.as_deref()),
4094 Some("high")
4095 );
4096 assert_eq!(
4097 merged
4098 .get("sid-1")
4099 .and_then(|overrides| overrides.station_name.as_deref()),
4100 Some("right")
4101 );
4102 assert_eq!(
4103 merged
4104 .get("sid-1")
4105 .and_then(|overrides| overrides.model.as_deref()),
4106 Some("gpt-5.4")
4107 );
4108 assert_eq!(
4109 merged
4110 .get("sid-1")
4111 .and_then(|overrides| overrides.service_tier.as_deref()),
4112 Some("priority")
4113 );
4114 assert_eq!(
4115 merged
4116 .get("sid-2")
4117 .and_then(|overrides| overrides.model.as_deref()),
4118 Some("gpt-5.4-mini")
4119 );
4120 assert!(
4121 merged
4122 .get("sid-2")
4123 .is_some_and(|overrides| overrides.reasoning_effort.is_none())
4124 );
4125 });
4126 }
4127
4128 #[test]
4129 fn provider_balance_snapshots_are_recorded_and_refreshed() {
4130 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4131 runtime.block_on(async {
4132 let state = ProxyState::new();
4133 state
4134 .record_provider_balance_snapshot(
4135 "codex",
4136 ProviderBalanceSnapshot {
4137 provider_id: "packycode".to_string(),
4138 station_name: Some("right".to_string()),
4139 upstream_index: Some(2),
4140 source: "usage_provider:budget_http_json".to_string(),
4141 fetched_at_ms: 10,
4142 stale_after_ms: Some(0),
4143 stale: false,
4144 status: BalanceSnapshotStatus::Unknown,
4145 exhausted: Some(false),
4146 exhaustion_affects_routing: true,
4147 plan_name: None,
4148 total_balance_usd: Some("3.5".to_string()),
4149 subscription_balance_usd: None,
4150 paygo_balance_usd: None,
4151 monthly_budget_usd: Some("5".to_string()),
4152 monthly_spent_usd: Some("1.5".to_string()),
4153 quota_period: None,
4154 quota_remaining_usd: None,
4155 quota_limit_usd: None,
4156 quota_used_usd: None,
4157 unlimited_quota: None,
4158 total_used_usd: None,
4159 today_used_usd: None,
4160 total_requests: None,
4161 today_requests: None,
4162 total_tokens: None,
4163 today_tokens: None,
4164 error: None,
4165 },
4166 )
4167 .await;
4168
4169 let view = state.get_provider_balance_view("codex").await;
4170 let balances = view.get("right").expect("station balance");
4171 assert_eq!(balances.len(), 1);
4172 assert_eq!(balances[0].provider_id, "packycode");
4173 assert_eq!(balances[0].status, BalanceSnapshotStatus::Stale);
4174 assert_eq!(balances[0].exhausted, Some(false));
4175
4176 let summary = state
4177 .get_provider_balance_summary_view("codex")
4178 .await
4179 .get("right")
4180 .cloned()
4181 .expect("station balance summary");
4182 assert_eq!(summary.snapshots, 1);
4183 assert_eq!(summary.stale, 1);
4184 assert_eq!(summary.routing_snapshots, 1);
4185 });
4186 }
4187
4188 #[test]
4189 fn provider_balance_snapshots_keep_multiple_providers_per_upstream() {
4190 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4191 runtime.block_on(async {
4192 let state = ProxyState::new();
4193 for (provider_id, status) in [
4194 ("general", BalanceSnapshotStatus::Error),
4195 ("newapi", BalanceSnapshotStatus::Ok),
4196 ] {
4197 state
4198 .record_provider_balance_snapshot(
4199 "codex",
4200 ProviderBalanceSnapshot {
4201 provider_id: provider_id.to_string(),
4202 station_name: Some("routing".to_string()),
4203 upstream_index: Some(1),
4204 source: "usage_provider:test".to_string(),
4205 fetched_at_ms: 10,
4206 stale_after_ms: None,
4207 stale: false,
4208 status,
4209 exhausted: if status == BalanceSnapshotStatus::Ok {
4210 Some(false)
4211 } else {
4212 None
4213 },
4214 exhaustion_affects_routing: true,
4215 plan_name: None,
4216 total_balance_usd: if status == BalanceSnapshotStatus::Ok {
4217 Some("3.5".to_string())
4218 } else {
4219 None
4220 },
4221 subscription_balance_usd: None,
4222 paygo_balance_usd: None,
4223 monthly_budget_usd: None,
4224 monthly_spent_usd: None,
4225 quota_period: None,
4226 quota_remaining_usd: None,
4227 quota_limit_usd: None,
4228 quota_used_usd: None,
4229 unlimited_quota: None,
4230 total_used_usd: None,
4231 today_used_usd: None,
4232 total_requests: None,
4233 today_requests: None,
4234 total_tokens: None,
4235 today_tokens: None,
4236 error: if status == BalanceSnapshotStatus::Error {
4237 Some("decode failed".to_string())
4238 } else {
4239 None
4240 },
4241 },
4242 )
4243 .await;
4244 }
4245
4246 let view = state.get_provider_balance_view("codex").await;
4247 let balances = view.get("routing").expect("station balance");
4248 assert_eq!(balances.len(), 2);
4249 assert_eq!(
4250 balances
4251 .iter()
4252 .map(|snapshot| snapshot.provider_id.as_str())
4253 .collect::<Vec<_>>(),
4254 vec!["general", "newapi"]
4255 );
4256 });
4257 }
4258
4259 #[test]
4260 fn provider_balance_error_refresh_preserves_previous_amounts() {
4261 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4262 runtime.block_on(async {
4263 let state = ProxyState::new();
4264 state
4265 .record_provider_balance_snapshot(
4266 "codex",
4267 ProviderBalanceSnapshot {
4268 provider_id: "input".to_string(),
4269 station_name: Some("routing".to_string()),
4270 upstream_index: Some(0),
4271 source: "usage_provider:sub2api_usage".to_string(),
4272 fetched_at_ms: unix_now_ms(),
4273 stale_after_ms: None,
4274 exhausted: Some(false),
4275 quota_period: Some("daily".to_string()),
4276 quota_remaining_usd: Some("263.68".to_string()),
4277 quota_limit_usd: Some("300.00".to_string()),
4278 ..ProviderBalanceSnapshot::default()
4279 },
4280 )
4281 .await;
4282
4283 state
4284 .record_provider_balance_snapshot(
4285 "codex",
4286 ProviderBalanceSnapshot {
4287 provider_id: "input".to_string(),
4288 station_name: Some("routing".to_string()),
4289 upstream_index: Some(0),
4290 source: "usage_provider:openai_balance_http_json".to_string(),
4291 fetched_at_ms: unix_now_ms(),
4292 stale_after_ms: None,
4293 error: Some("usage provider response read failed".to_string()),
4294 ..ProviderBalanceSnapshot::default()
4295 },
4296 )
4297 .await;
4298
4299 let view = state.get_provider_balance_view("codex").await;
4300 let snapshot = view
4301 .get("routing")
4302 .and_then(|snapshots| snapshots.first())
4303 .expect("routing balance snapshot");
4304
4305 assert_eq!(snapshot.status, BalanceSnapshotStatus::Error);
4306 assert_eq!(snapshot.quota_period.as_deref(), Some("daily"));
4307 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("263.68"));
4308 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("300.00"));
4309 assert_eq!(
4310 snapshot.error.as_deref(),
4311 Some("usage provider response read failed")
4312 );
4313 });
4314 }
4315
4316 #[test]
4317 fn prune_runtime_observability_keeps_catalog_provider_balances_on_routing_layout_change() {
4318 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4319 runtime.block_on(async {
4320 let state = ProxyState::new();
4321
4322 let mut initial_mgr = ServiceConfigManager::default();
4323 initial_mgr.configs.insert(
4324 "routing".to_string(),
4325 ServiceConfig {
4326 name: "routing".to_string(),
4327 alias: None,
4328 enabled: true,
4329 level: 1,
4330 upstreams: vec![
4331 UpstreamConfig {
4332 base_url: "https://input.example/v1".to_string(),
4333 auth: UpstreamAuth::default(),
4334 tags: HashMap::from([("provider_id".to_string(), "input".to_string())]),
4335 supported_models: HashMap::new(),
4336 model_mapping: HashMap::new(),
4337 },
4338 UpstreamConfig {
4339 base_url: "https://backup.example/v1".to_string(),
4340 auth: UpstreamAuth::default(),
4341 tags: HashMap::from([(
4342 "provider_id".to_string(),
4343 "backup".to_string(),
4344 )]),
4345 supported_models: HashMap::new(),
4346 model_mapping: HashMap::new(),
4347 },
4348 ],
4349 },
4350 );
4351 state
4352 .prune_runtime_observability_for_service("codex", &initial_mgr)
4353 .await;
4354
4355 state
4356 .record_provider_balance_snapshot(
4357 "codex",
4358 ProviderBalanceSnapshot {
4359 provider_id: "input".to_string(),
4360 station_name: Some("input".to_string()),
4361 upstream_index: Some(0),
4362 source: "usage_provider:test".to_string(),
4363 fetched_at_ms: 10,
4364 stale_after_ms: None,
4365 stale: false,
4366 status: BalanceSnapshotStatus::Ok,
4367 exhausted: Some(false),
4368 exhaustion_affects_routing: true,
4369 total_balance_usd: Some("3.5".to_string()),
4370 ..ProviderBalanceSnapshot::default()
4371 },
4372 )
4373 .await;
4374 state
4375 .record_provider_balance_snapshot(
4376 "codex",
4377 ProviderBalanceSnapshot {
4378 provider_id: "input".to_string(),
4379 station_name: Some("routing".to_string()),
4380 upstream_index: Some(0),
4381 source: "usage_provider:test".to_string(),
4382 fetched_at_ms: 10,
4383 stale_after_ms: None,
4384 stale: false,
4385 status: BalanceSnapshotStatus::Ok,
4386 exhausted: Some(false),
4387 exhaustion_affects_routing: true,
4388 total_balance_usd: Some("3.5".to_string()),
4389 ..ProviderBalanceSnapshot::default()
4390 },
4391 )
4392 .await;
4393
4394 let mut pinned_mgr = ServiceConfigManager::default();
4395 pinned_mgr.configs.insert(
4396 "routing".to_string(),
4397 ServiceConfig {
4398 name: "routing".to_string(),
4399 alias: None,
4400 enabled: true,
4401 level: 1,
4402 upstreams: vec![UpstreamConfig {
4403 base_url: "https://input.example/v1".to_string(),
4404 auth: UpstreamAuth::default(),
4405 tags: HashMap::from([("provider_id".to_string(), "input".to_string())]),
4406 supported_models: HashMap::new(),
4407 model_mapping: HashMap::new(),
4408 }],
4409 },
4410 );
4411 state
4412 .prune_runtime_observability_for_service("codex", &pinned_mgr)
4413 .await;
4414
4415 let view = state.get_provider_balance_view("codex").await;
4416 assert!(view.contains_key("input"));
4417 assert!(!view.contains_key("routing"));
4418 });
4419 }
4420
4421 #[test]
4422 fn prune_runtime_observability_removes_stale_service_keys() {
4423 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4424 runtime.block_on(async {
4425 let state = ProxyState::new();
4426
4427 state
4428 .set_station_enabled_override("codex", "old".to_string(), false, 1)
4429 .await;
4430 state
4431 .set_upstream_enabled_override(
4432 "codex",
4433 "https://old.example/v1".to_string(),
4434 false,
4435 1,
4436 )
4437 .await;
4438 state
4439 .set_provider_endpoint_enabled_override(
4440 "codex",
4441 ProviderEndpointKey::new("codex", "provider-old", "default"),
4442 false,
4443 1,
4444 )
4445 .await;
4446 state
4447 .set_provider_endpoint_runtime_state_override(
4448 "codex",
4449 ProviderEndpointKey::new("codex", "provider-old", "default"),
4450 RuntimeConfigState::BreakerOpen,
4451 1,
4452 )
4453 .await;
4454 state
4455 .record_station_health(
4456 "codex",
4457 "old".to_string(),
4458 StationHealth {
4459 checked_at_ms: 10,
4460 upstreams: vec![UpstreamHealth {
4461 base_url: "https://old.example/v1".to_string(),
4462 ok: Some(false),
4463 status_code: Some(500),
4464 latency_ms: Some(10),
4465 error: Some("boom".to_string()),
4466 passive: None,
4467 }],
4468 },
4469 )
4470 .await;
4471 state
4472 .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
4473 service_name: "codex".to_string(),
4474 station_name: "old".to_string(),
4475 base_url: "https://old.example/v1".to_string(),
4476 status_code: Some(500),
4477 error_class: Some("upstream_transport_error".to_string()),
4478 error: Some("boom".to_string()),
4479 now_ms: 20,
4480 })
4481 .await;
4482 state
4483 .record_provider_balance_snapshot(
4484 "codex",
4485 ProviderBalanceSnapshot {
4486 provider_id: "provider-old".to_string(),
4487 station_name: Some("old".to_string()),
4488 upstream_index: Some(0),
4489 source: "usage_provider:budget_http_json".to_string(),
4490 fetched_at_ms: 10,
4491 stale_after_ms: None,
4492 stale: false,
4493 status: BalanceSnapshotStatus::Ok,
4494 exhausted: Some(false),
4495 exhaustion_affects_routing: true,
4496 plan_name: None,
4497 total_balance_usd: Some("3.5".to_string()),
4498 subscription_balance_usd: None,
4499 paygo_balance_usd: None,
4500 monthly_budget_usd: None,
4501 monthly_spent_usd: None,
4502 quota_period: None,
4503 quota_remaining_usd: None,
4504 quota_limit_usd: None,
4505 quota_used_usd: None,
4506 unlimited_quota: None,
4507 total_used_usd: None,
4508 today_used_usd: None,
4509 total_requests: None,
4510 today_requests: None,
4511 total_tokens: None,
4512 today_tokens: None,
4513 error: None,
4514 },
4515 )
4516 .await;
4517
4518 let request_id = state
4519 .begin_request_for_test()
4520 .session_id("sid-old")
4521 .model("gpt-5")
4522 .started_at_ms(30)
4523 .begin()
4524 .await;
4525 state
4526 .update_request_route(
4527 request_id,
4528 Some("old".to_string()),
4529 Some("provider-old".to_string()),
4530 "https://old.example/v1".to_string(),
4531 None,
4532 )
4533 .await;
4534 state
4535 .finish_request(FinishRequestParams {
4536 id: request_id,
4537 status_code: 200,
4538 duration_ms: 5,
4539 ended_at_ms: 35,
4540 observed_service_tier: None,
4541 usage: None,
4542 retry: None,
4543 ttfb_ms: None,
4544 streaming: false,
4545 })
4546 .await;
4547
4548 let mut mgr = ServiceConfigManager::default();
4549 mgr.configs.insert(
4550 "new".to_string(),
4551 ServiceConfig {
4552 name: "new".to_string(),
4553 alias: None,
4554 enabled: true,
4555 level: 1,
4556 upstreams: vec![UpstreamConfig {
4557 base_url: "https://new.example/v1".to_string(),
4558 auth: UpstreamAuth::default(),
4559 tags: HashMap::from([(
4560 "provider_id".to_string(),
4561 "provider-new".to_string(),
4562 )]),
4563 supported_models: HashMap::new(),
4564 model_mapping: HashMap::new(),
4565 }],
4566 },
4567 );
4568
4569 state
4570 .prune_runtime_observability_for_service("codex", &mgr)
4571 .await;
4572
4573 assert!(state.get_station_meta_overrides("codex").await.is_empty());
4574 assert!(state.get_upstream_meta_overrides("codex").await.is_empty());
4575 assert!(state.get_station_health("codex").await.is_empty());
4576 assert!(state.get_provider_balance_view("codex").await.is_empty());
4577
4578 let rollup = state.get_usage_rollup_view("codex", 10, 10).await;
4579 assert!(rollup.by_config.is_empty());
4580 assert!(rollup.by_provider.is_empty());
4581 });
4582 }
4583
4584 #[test]
4585 fn get_upstream_meta_overrides_merges_endpoint_and_legacy_base_url_entries() {
4586 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4587 runtime.block_on(async {
4588 let state = ProxyState::new();
4589
4590 state
4591 .set_provider_endpoint_enabled_override(
4592 "codex",
4593 ProviderEndpointKey::new("codex", "alpha", "default"),
4594 false,
4595 1,
4596 )
4597 .await;
4598 state
4599 .set_provider_endpoint_runtime_state_override(
4600 "codex",
4601 ProviderEndpointKey::new("codex", "alpha", "default"),
4602 RuntimeConfigState::BreakerOpen,
4603 2,
4604 )
4605 .await;
4606 state
4607 .set_upstream_enabled_override(
4608 "codex",
4609 "https://legacy.example/v1".to_string(),
4610 true,
4611 3,
4612 )
4613 .await;
4614 state
4615 .set_upstream_runtime_state_override(
4616 "codex",
4617 "https://legacy.example/v1".to_string(),
4618 RuntimeConfigState::Draining,
4619 4,
4620 )
4621 .await;
4622
4623 let overrides = state.get_upstream_meta_overrides("codex").await;
4624
4625 assert_eq!(
4626 overrides.get("codex/alpha/default"),
4627 Some(&(Some(false), Some(RuntimeConfigState::BreakerOpen)))
4628 );
4629 assert_eq!(
4630 overrides.get("https://legacy.example/v1"),
4631 Some(&(Some(true), Some(RuntimeConfigState::Draining)))
4632 );
4633 });
4634 }
4635
4636 #[test]
4637 fn provider_endpoint_runtime_health_is_keyed_by_provider_endpoint() {
4638 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4639 runtime.block_on(async {
4640 let state = ProxyState::new();
4641 let monthly = ProviderEndpointKey::new("codex", "monthly", "default");
4642 let fallback = ProviderEndpointKey::new("codex", "fallback", "default");
4643
4644 state
4645 .record_provider_endpoint_attempt_success("codex", fallback.clone(), 10)
4646 .await;
4647 state
4648 .set_provider_endpoint_usage_exhausted("codex", monthly.clone(), true)
4649 .await;
4650 state
4651 .record_provider_endpoint_attempt_failure(
4652 "codex",
4653 monthly.clone(),
4654 0,
4655 crate::lb::CooldownBackoff {
4656 factor: 1,
4657 max_secs: 0,
4658 },
4659 )
4660 .await;
4661 state
4662 .record_provider_endpoint_attempt_failure(
4663 "codex",
4664 monthly.clone(),
4665 0,
4666 crate::lb::CooldownBackoff {
4667 factor: 1,
4668 max_secs: 0,
4669 },
4670 )
4671 .await;
4672 state
4673 .record_provider_endpoint_attempt_failure(
4674 "codex",
4675 monthly.clone(),
4676 30,
4677 crate::lb::CooldownBackoff {
4678 factor: 1,
4679 max_secs: 0,
4680 },
4681 )
4682 .await;
4683
4684 let runtime = state
4685 .route_plan_runtime_state_for_provider_endpoints("codex")
4686 .await;
4687 let monthly_state = runtime.provider_endpoint(&monthly);
4688 assert_eq!(monthly_state.failure_count, crate::lb::FAILURE_THRESHOLD);
4689 assert!(monthly_state.cooldown_active);
4690 assert!(monthly_state.usage_exhausted);
4691 assert_eq!(runtime.affinity_provider_endpoint(), Some(&fallback));
4692
4693 state
4694 .set_provider_endpoint_runtime_state_override(
4695 "codex",
4696 fallback.clone(),
4697 RuntimeConfigState::BreakerOpen,
4698 20,
4699 )
4700 .await;
4701 let runtime = state
4702 .route_plan_runtime_state_for_provider_endpoints("codex")
4703 .await;
4704 assert!(runtime.provider_endpoint(&fallback).runtime_disabled);
4705 });
4706 }
4707
4708 #[test]
4709 fn get_station_health_merges_passive_runtime_observations() {
4710 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4711 runtime.block_on(async {
4712 let state = ProxyState::new();
4713 state
4714 .record_station_health(
4715 "codex",
4716 "right".to_string(),
4717 StationHealth {
4718 checked_at_ms: 10,
4719 upstreams: vec![UpstreamHealth {
4720 base_url: "https://right.example/v1".to_string(),
4721 ok: Some(true),
4722 status_code: Some(200),
4723 latency_ms: Some(120),
4724 error: None,
4725 passive: None,
4726 }],
4727 },
4728 )
4729 .await;
4730 state
4731 .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
4732 service_name: "codex".to_string(),
4733 station_name: "right".to_string(),
4734 base_url: "https://right.example/v1".to_string(),
4735 status_code: Some(500),
4736 error_class: Some("cloudflare_timeout".to_string()),
4737 error: Some("upstream timed out".to_string()),
4738 now_ms: 20,
4739 })
4740 .await;
4741 state
4742 .record_passive_upstream_success(
4743 "codex",
4744 "right",
4745 "https://backup.example/v1",
4746 Some(200),
4747 30,
4748 )
4749 .await;
4750
4751 let health = state.get_station_health("codex").await;
4752 let right = health.get("right").expect("right health");
4753 assert_eq!(right.checked_at_ms, 30);
4754 assert_eq!(right.upstreams.len(), 2);
4755
4756 let primary = right
4757 .upstreams
4758 .iter()
4759 .find(|upstream| upstream.base_url == "https://right.example/v1")
4760 .expect("primary upstream");
4761 assert_eq!(primary.ok, Some(true));
4762 let primary_passive = primary.passive.as_ref().expect("primary passive");
4763 assert_eq!(primary_passive.state, PassiveHealthState::Degraded);
4764 assert_eq!(primary_passive.score, 50);
4765 assert_eq!(primary_passive.last_status_code, Some(500));
4766 assert_eq!(
4767 primary_passive.last_error_class.as_deref(),
4768 Some("cloudflare_timeout")
4769 );
4770
4771 let backup = right
4772 .upstreams
4773 .iter()
4774 .find(|upstream| upstream.base_url == "https://backup.example/v1")
4775 .expect("backup upstream");
4776 assert_eq!(backup.ok, None);
4777 let backup_passive = backup.passive.as_ref().expect("backup passive");
4778 assert_eq!(backup_passive.state, PassiveHealthState::Healthy);
4779 assert_eq!(backup_passive.score, 100);
4780 assert_eq!(backup_passive.last_status_code, Some(200));
4781 });
4782 }
4783
4784 #[test]
4785 fn passive_health_success_recovers_after_failure() {
4786 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4787 runtime.block_on(async {
4788 let state = ProxyState::new();
4789 state
4790 .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
4791 service_name: "codex".to_string(),
4792 station_name: "right".to_string(),
4793 base_url: "https://right.example/v1".to_string(),
4794 status_code: Some(500),
4795 error_class: Some("cloudflare_timeout".to_string()),
4796 error: Some("upstream timed out".to_string()),
4797 now_ms: 10,
4798 })
4799 .await;
4800 state
4801 .record_passive_upstream_success(
4802 "codex",
4803 "right",
4804 "https://right.example/v1",
4805 Some(200),
4806 20,
4807 )
4808 .await;
4809
4810 let health = state.get_station_health("codex").await;
4811 let right = health.get("right").expect("right health");
4812 let upstream = right.upstreams.first().expect("upstream");
4813 let passive = upstream.passive.as_ref().expect("passive");
4814 assert_eq!(passive.state, PassiveHealthState::Healthy);
4815 assert_eq!(passive.score, 100);
4816 assert_eq!(passive.consecutive_failures, 0);
4817 assert_eq!(passive.last_success_at_ms, Some(20));
4818 assert_eq!(passive.last_failure_at_ms, Some(10));
4819 assert_eq!(passive.last_error_class, None);
4820 assert_eq!(passive.last_error, None);
4821 });
4822 }
4823
4824 #[test]
4825 fn apply_session_profile_binding_uses_inherited_profile_values() {
4826 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4827 runtime.block_on(async {
4828 let state = ProxyState::new();
4829 let mut mgr = ServiceConfigManager::default();
4830 mgr.configs.insert(
4831 "right".to_string(),
4832 ServiceConfig {
4833 name: "right".to_string(),
4834 alias: None,
4835 enabled: true,
4836 level: 1,
4837 upstreams: vec![UpstreamConfig {
4838 base_url: "https://right.example/v1".to_string(),
4839 auth: UpstreamAuth::default(),
4840 tags: HashMap::from([
4841 ("supports_reasoning_effort".to_string(), "true".to_string()),
4842 ("supports_service_tier".to_string(), "true".to_string()),
4843 ]),
4844 supported_models: HashMap::from([("gpt-5.4".to_string(), true)]),
4845 model_mapping: HashMap::new(),
4846 }],
4847 },
4848 );
4849 mgr.profiles.insert(
4850 "base".to_string(),
4851 crate::config::ServiceControlProfile {
4852 extends: None,
4853 station: Some("right".to_string()),
4854 model: Some("gpt-5.4".to_string()),
4855 reasoning_effort: None,
4856 service_tier: Some("priority".to_string()),
4857 },
4858 );
4859 mgr.profiles.insert(
4860 "fast".to_string(),
4861 crate::config::ServiceControlProfile {
4862 extends: Some("base".to_string()),
4863 station: None,
4864 model: None,
4865 reasoning_effort: Some("low".to_string()),
4866 service_tier: None,
4867 },
4868 );
4869
4870 state
4871 .apply_session_profile_binding(
4872 "codex",
4873 &mgr,
4874 "sid-inherited".to_string(),
4875 "fast".to_string(),
4876 100,
4877 )
4878 .await
4879 .expect("apply inherited profile");
4880
4881 let binding = state
4882 .get_session_binding("sid-inherited")
4883 .await
4884 .expect("binding exists");
4885 assert_eq!(binding.profile_name.as_deref(), Some("fast"));
4886 assert_eq!(binding.station_name.as_deref(), Some("right"));
4887 assert_eq!(binding.model.as_deref(), Some("gpt-5.4"));
4888 assert_eq!(binding.reasoning_effort.as_deref(), Some("low"));
4889 assert_eq!(binding.service_tier.as_deref(), Some("priority"));
4890 });
4891 }
4892
4893 #[test]
4894 fn prune_periodic_keeps_sticky_binding_after_manual_override_ttl_expires() {
4895 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4896 runtime.block_on(async {
4897 let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(1, 0, 2_000));
4898 state
4899 .set_session_model_override("sid-sticky".to_string(), "gpt-5.4".to_string(), 0)
4900 .await;
4901 state
4902 .set_session_binding(SessionBinding {
4903 session_id: "sid-sticky".to_string(),
4904 profile_name: Some("daily".to_string()),
4905 station_name: Some("right".to_string()),
4906 model: Some("gpt-5.4".to_string()),
4907 reasoning_effort: Some("medium".to_string()),
4908 service_tier: Some("default".to_string()),
4909 continuity_mode: SessionContinuityMode::DefaultProfile,
4910 created_at_ms: 0,
4911 updated_at_ms: 0,
4912 last_seen_ms: 0,
4913 })
4914 .await;
4915
4916 state.prune_periodic().await;
4917
4918 assert!(
4919 state
4920 .get_session_model_override("sid-sticky")
4921 .await
4922 .is_none()
4923 );
4924 assert!(state.get_session_binding("sid-sticky").await.is_some());
4925 });
4926 }
4927
4928 #[test]
4929 fn prune_periodic_honors_opt_in_binding_ttl() {
4930 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4931 runtime.block_on(async {
4932 let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 1, 2_000));
4933 state
4934 .set_session_binding(SessionBinding {
4935 session_id: "sid-expire".to_string(),
4936 profile_name: Some("daily".to_string()),
4937 station_name: Some("right".to_string()),
4938 model: Some("gpt-5.4".to_string()),
4939 reasoning_effort: Some("medium".to_string()),
4940 service_tier: Some("default".to_string()),
4941 continuity_mode: SessionContinuityMode::DefaultProfile,
4942 created_at_ms: 0,
4943 updated_at_ms: 0,
4944 last_seen_ms: 0,
4945 })
4946 .await;
4947
4948 state.prune_periodic().await;
4949
4950 assert!(state.get_session_binding("sid-expire").await.is_none());
4951 });
4952 }
4953
4954 #[test]
4955 fn session_route_affinity_ttl_is_enforced_on_read() {
4956 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4957 runtime.block_on(async {
4958 let mut policy = test_runtime_policy(0, 0, 2_000);
4959 policy.session_route_affinity_ttl_ms = 1;
4960 let state = ProxyState::new_with_runtime_policy(None, policy);
4961 state
4962 .record_session_route_affinity_success(
4963 "sid-expire",
4964 SessionRouteAffinityTarget {
4965 route_graph_key: "graph".to_string(),
4966 session_identity_source: None,
4967 provider_endpoint: ProviderEndpointKey::new("codex", "monthly", "default"),
4968 upstream_base_url: "https://monthly.example/v1".to_string(),
4969 route_path: vec!["monthly_first".to_string(), "monthly".to_string()],
4970 },
4971 Some("first_success".to_string()),
4972 0,
4973 )
4974 .await;
4975
4976 assert!(
4977 state
4978 .get_session_route_affinity("sid-expire")
4979 .await
4980 .is_none()
4981 );
4982 assert!(state.list_session_route_affinities().await.is_empty());
4983 });
4984 }
4985
4986 #[test]
4987 fn session_route_affinity_ledger_does_not_restore_expired_entries() {
4988 let runtime = tokio::runtime::Runtime::new().expect("runtime");
4989 runtime.block_on(async {
4990 let _home = env_lock().await;
4991 let temp_dir = std::env::temp_dir().join(format!(
4992 "codex-helper-route-ledger-test-{}",
4993 uuid::Uuid::new_v4()
4994 ));
4995 std::fs::create_dir_all(&temp_dir).expect("create temp dir");
4996 let mut scoped = ScopedEnv::default();
4997 unsafe {
4998 scoped.set_path("CODEX_HELPER_HOME", temp_dir.as_path());
4999 }
5000
5001 let mut policy = test_runtime_policy(0, 0, 2_000);
5002 policy.session_route_affinity_ttl_ms = 1;
5003 let first_state = ProxyState::new_with_runtime_policy(None, policy.clone());
5004 first_state
5005 .record_session_route_affinity_success(
5006 "sid-expired-ledger",
5007 SessionRouteAffinityTarget {
5008 route_graph_key: "graph".to_string(),
5009 session_identity_source: None,
5010 provider_endpoint: ProviderEndpointKey::new("codex", "monthly", "default"),
5011 upstream_base_url: "https://monthly.example/v1".to_string(),
5012 route_path: vec!["monthly_first".to_string(), "monthly".to_string()],
5013 },
5014 Some("first_success".to_string()),
5015 0,
5016 )
5017 .await;
5018
5019 let second_state = ProxyState::new_with_runtime_policy(None, policy);
5020 assert!(
5021 second_state
5022 .get_session_route_affinity("sid-expired-ledger")
5023 .await
5024 .is_none()
5025 );
5026 });
5027 }
5028
5029 #[test]
5030 fn prune_periodic_caps_sticky_bindings_by_last_seen() {
5031 let runtime = tokio::runtime::Runtime::new().expect("runtime");
5032 runtime.block_on(async {
5033 let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2));
5034 for idx in 0..3 {
5035 state
5036 .set_session_binding(SessionBinding {
5037 session_id: format!("sid-{idx}"),
5038 profile_name: Some("daily".to_string()),
5039 station_name: Some("right".to_string()),
5040 model: Some("gpt-5.4".to_string()),
5041 reasoning_effort: Some("medium".to_string()),
5042 service_tier: Some("default".to_string()),
5043 continuity_mode: SessionContinuityMode::DefaultProfile,
5044 created_at_ms: idx,
5045 updated_at_ms: idx,
5046 last_seen_ms: idx,
5047 })
5048 .await;
5049 }
5050
5051 state.prune_periodic().await;
5052
5053 assert!(state.get_session_binding("sid-0").await.is_none());
5054 assert!(state.get_session_binding("sid-1").await.is_some());
5055 assert!(state.get_session_binding("sid-2").await.is_some());
5056 });
5057 }
5058
5059 #[test]
5060 fn transcript_path_cache_records_negative_lookups() {
5061 let runtime = tokio::runtime::Runtime::new().expect("runtime");
5062 runtime.block_on(async {
5063 let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5064 let paths = state
5065 .resolve_host_transcript_paths_cached(&["missing-session".to_string()])
5066 .await;
5067
5068 assert_eq!(paths.get("missing-session"), Some(&None));
5069 let cache = state.session_transcript_path_cache.read().await;
5070 assert!(cache.contains_key("missing-session"));
5071 });
5072 }
5073}