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