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