Skip to main content

codex_helper_core/
state.rs

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