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