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        self.notify_state_changed();
1623    }
1624
1625    pub async fn prune_runtime_observability_for_service(
1626        &self,
1627        service_name: &str,
1628        mgr: &ServiceConfigManager,
1629    ) {
1630        let mut changed = false;
1631        let active_stations = mgr.stations().keys().cloned().collect::<HashSet<_>>();
1632        let active_upstreams = mgr
1633            .stations()
1634            .iter()
1635            .map(|(station_name, service)| {
1636                (
1637                    station_name.clone(),
1638                    service
1639                        .upstreams
1640                        .iter()
1641                        .map(|upstream| upstream.base_url.clone())
1642                        .collect::<HashSet<_>>(),
1643                )
1644            })
1645            .collect::<HashMap<_, _>>();
1646        let active_base_urls = active_upstreams
1647            .values()
1648            .flat_map(|upstreams| upstreams.iter().cloned())
1649            .collect::<HashSet<_>>();
1650        let active_provider_endpoint_keys = mgr
1651            .stations()
1652            .iter()
1653            .flat_map(|(station_name, service)| {
1654                service.upstreams.iter().enumerate().map(|(idx, upstream)| {
1655                    Self::active_provider_endpoint_key_for_upstream(
1656                        service_name,
1657                        station_name.as_str(),
1658                        idx,
1659                        upstream,
1660                    )
1661                })
1662            })
1663            .collect::<HashSet<_>>();
1664        let mut active_provider_ids = HashSet::from(["-".to_string()]);
1665        for service in mgr.stations().values() {
1666            for upstream in &service.upstreams {
1667                if let Some(provider_id) = upstream.tags.get("provider_id") {
1668                    active_provider_ids.insert(provider_id.clone());
1669                }
1670            }
1671        }
1672
1673        let layout = service_layout_signature(mgr);
1674        let balance_prune_stations = {
1675            let mut signatures = self.service_layout_signatures.write().await;
1676            let changed = signatures.get(service_name).map_or_else(
1677                || {
1678                    if active_stations.len() == 1 && active_stations.contains("routing") {
1679                        Some(HashSet::from(["routing".to_string()]))
1680                    } else {
1681                        None
1682                    }
1683                },
1684                |previous| Some(changed_service_layout_stations(previous, &layout)),
1685            );
1686            signatures.insert(service_name.to_string(), layout);
1687            changed
1688        };
1689
1690        match balance_prune_stations {
1691            Some(changed_layout_stations) if !changed_layout_stations.is_empty() => {
1692                let mut provider_balances = self.provider_balances.write().await;
1693                if let Some(per_service) = provider_balances.get_mut(service_name) {
1694                    let before = per_service.len();
1695                    per_service
1696                        .retain(|station_name, _| !changed_layout_stations.contains(station_name));
1697                    changed |= per_service.len() != before;
1698                    if per_service.is_empty() {
1699                        provider_balances.remove(service_name);
1700                    }
1701                }
1702                let mut provider_balance_history = self.provider_balance_history.write().await;
1703                if let Some(per_service) = provider_balance_history.get_mut(service_name) {
1704                    let before = per_service.len();
1705                    per_service
1706                        .retain(|station_name, _| !changed_layout_stations.contains(station_name));
1707                    changed |= per_service.len() != before;
1708                    if per_service.is_empty() {
1709                        provider_balance_history.remove(service_name);
1710                    }
1711                }
1712                let mut provider_balance_summaries = self.provider_balance_summaries.write().await;
1713                if let Some(per_service) = provider_balance_summaries.get_mut(service_name) {
1714                    let before = per_service.len();
1715                    per_service
1716                        .retain(|station_name, _| !changed_layout_stations.contains(station_name));
1717                    changed |= per_service.len() != before;
1718                    if per_service.is_empty() {
1719                        provider_balance_summaries.remove(service_name);
1720                    }
1721                }
1722            }
1723            None => {
1724                let mut provider_balances = self.provider_balances.write().await;
1725                changed |= provider_balances.remove(service_name).is_some();
1726                let mut provider_balance_history = self.provider_balance_history.write().await;
1727                changed |= provider_balance_history.remove(service_name).is_some();
1728                let mut provider_balance_summaries = self.provider_balance_summaries.write().await;
1729                changed |= provider_balance_summaries.remove(service_name).is_some();
1730            }
1731            Some(_) => {}
1732        }
1733
1734        {
1735            let mut guard = self.station_meta_overrides.write().await;
1736            if let Some(per_service) = guard.get_mut(service_name) {
1737                let before = per_service.len();
1738                per_service.retain(|station_name, _| active_stations.contains(station_name));
1739                changed |= per_service.len() != before;
1740                if per_service.is_empty() {
1741                    guard.remove(service_name);
1742                }
1743            }
1744        }
1745
1746        {
1747            let mut guard = self.upstream_meta_overrides.write().await;
1748            if let Some(per_service) = guard.get_mut(service_name) {
1749                let before = per_service.len();
1750                per_service.retain(|base_url, _| active_base_urls.contains(base_url));
1751                changed |= per_service.len() != before;
1752                if per_service.is_empty() {
1753                    guard.remove(service_name);
1754                }
1755            }
1756        }
1757
1758        {
1759            let mut guard = self.provider_endpoint_meta_overrides.write().await;
1760            if let Some(per_service) = guard.get_mut(service_name) {
1761                let before = per_service.len();
1762                per_service
1763                    .retain(|endpoint_key, _| active_provider_endpoint_keys.contains(endpoint_key));
1764                changed |= per_service.len() != before;
1765                if per_service.is_empty() {
1766                    guard.remove(service_name);
1767                }
1768            }
1769        }
1770
1771        {
1772            let mut guard = self.provider_endpoint_runtime_health.write().await;
1773            if let Some(per_service) = guard.get_mut(service_name) {
1774                per_service
1775                    .retain(|endpoint_key, _| active_provider_endpoint_keys.contains(endpoint_key));
1776                if per_service.is_empty() {
1777                    guard.remove(service_name);
1778                }
1779            }
1780        }
1781
1782        {
1783            let mut guard = self.policy_actions.write().await;
1784            let mut snapshot = guard.clone();
1785            let mut policy_actions_changed = false;
1786            if let Some(per_service) = snapshot.get_mut(service_name) {
1787                let before = per_service.len();
1788                per_service
1789                    .retain(|endpoint_key, _| active_provider_endpoint_keys.contains(endpoint_key));
1790                policy_actions_changed |= per_service.len() != before;
1791                if per_service.is_empty() {
1792                    snapshot.remove(service_name);
1793                }
1794            }
1795            if policy_actions_changed && self.persist_policy_actions(&snapshot).await {
1796                *guard = snapshot;
1797                changed = true;
1798            }
1799        }
1800
1801        {
1802            let mut guard = self.station_health.write().await;
1803            if let Some(per_service) = guard.get_mut(service_name) {
1804                let before_stations = per_service.len();
1805                let before_upstreams = per_service
1806                    .values()
1807                    .map(|station_health| station_health.upstreams.len())
1808                    .sum::<usize>();
1809                per_service.retain(|station_name, station_health| {
1810                    if !active_stations.contains(station_name) {
1811                        return false;
1812                    }
1813                    if let Some(allowed_upstreams) = active_upstreams.get(station_name) {
1814                        station_health
1815                            .upstreams
1816                            .retain(|upstream| allowed_upstreams.contains(&upstream.base_url));
1817                    }
1818                    !station_health.upstreams.is_empty()
1819                });
1820                let after_upstreams = per_service
1821                    .values()
1822                    .map(|station_health| station_health.upstreams.len())
1823                    .sum::<usize>();
1824                changed |=
1825                    per_service.len() != before_stations || after_upstreams != before_upstreams;
1826                if per_service.is_empty() {
1827                    guard.remove(service_name);
1828                }
1829            }
1830        }
1831
1832        {
1833            let mut guard = self.passive_station_health.write().await;
1834            if let Some(per_service) = guard.get_mut(service_name) {
1835                let before_stations = per_service.len();
1836                let before_upstreams = per_service.values().map(HashMap::len).sum::<usize>();
1837                per_service.retain(|station_name, station_health| {
1838                    if !active_stations.contains(station_name) {
1839                        return false;
1840                    }
1841                    if let Some(allowed_upstreams) = active_upstreams.get(station_name) {
1842                        station_health.retain(|base_url, _| allowed_upstreams.contains(base_url));
1843                    }
1844                    !station_health.is_empty()
1845                });
1846                let after_upstreams = per_service.values().map(HashMap::len).sum::<usize>();
1847                changed |=
1848                    per_service.len() != before_stations || after_upstreams != before_upstreams;
1849                if per_service.is_empty() {
1850                    guard.remove(service_name);
1851                }
1852            }
1853        }
1854
1855        {
1856            let mut guard = self.station_health_checks.write().await;
1857            if let Some(per_service) = guard.get_mut(service_name) {
1858                let before = per_service.len();
1859                per_service.retain(|station_name, _| active_stations.contains(station_name));
1860                changed |= per_service.len() != before;
1861                if per_service.is_empty() {
1862                    guard.remove(service_name);
1863                }
1864            }
1865        }
1866
1867        {
1868            let mut guard = self.usage_rollups.write().await;
1869            if let Some(rollup) = guard.get_mut(service_name) {
1870                let before_by_config = rollup.by_config.len();
1871                rollup
1872                    .by_config
1873                    .retain(|station_name, _| active_stations.contains(station_name));
1874                changed |= rollup.by_config.len() != before_by_config;
1875                let before_by_config_day = rollup.by_config_day.len();
1876                rollup.by_config_day.retain(|station_name, _day_map| {
1877                    if !active_stations.contains(station_name) {
1878                        return false;
1879                    }
1880                    true
1881                });
1882                changed |= rollup.by_config_day.len() != before_by_config_day;
1883                let before_by_provider = rollup.by_provider.len();
1884                rollup
1885                    .by_provider
1886                    .retain(|provider_id, _| active_provider_ids.contains(provider_id));
1887                changed |= rollup.by_provider.len() != before_by_provider;
1888                let before_by_provider_day = rollup.by_provider_day.len();
1889                rollup.by_provider_day.retain(|provider_id, _day_map| {
1890                    if !active_provider_ids.contains(provider_id) {
1891                        return false;
1892                    }
1893                    true
1894                });
1895                changed |= rollup.by_provider_day.len() != before_by_provider_day;
1896            }
1897        }
1898        if changed {
1899            self.notify_state_changed();
1900        }
1901    }
1902
1903    fn active_provider_endpoint_key_for_upstream(
1904        service_name: &str,
1905        station_name: &str,
1906        upstream_index: usize,
1907        upstream: &crate::config::UpstreamConfig,
1908    ) -> ProviderEndpointKey {
1909        let provider_id = upstream
1910            .tags
1911            .get("provider_id")
1912            .cloned()
1913            .unwrap_or_else(|| format!("{station_name}#{upstream_index}"));
1914        let endpoint_id = upstream
1915            .tags
1916            .get("endpoint_id")
1917            .cloned()
1918            .unwrap_or_else(|| upstream_index.to_string());
1919
1920        ProviderEndpointKey::new(service_name, provider_id, endpoint_id)
1921    }
1922
1923    pub async fn record_station_health(
1924        &self,
1925        service_name: &str,
1926        station_name: String,
1927        health: StationHealth,
1928    ) {
1929        let mut guard = self.station_health.write().await;
1930        let per_service = guard.entry(service_name.to_string()).or_default();
1931        per_service.insert(station_name, health);
1932        self.notify_state_changed();
1933    }
1934
1935    pub async fn get_station_health(&self, service_name: &str) -> HashMap<String, StationHealth> {
1936        let active = {
1937            let guard = self.station_health.read().await;
1938            guard.get(service_name).cloned().unwrap_or_default()
1939        };
1940        let passive = {
1941            let guard = self.passive_station_health.read().await;
1942            guard.get(service_name).cloned().unwrap_or_default()
1943        };
1944        merge_station_health(active, passive)
1945    }
1946
1947    pub async fn record_provider_balance_snapshot(
1948        &self,
1949        service_name: &str,
1950        mut snapshot: ProviderBalanceSnapshot,
1951    ) {
1952        let (Some(station_name), Some(upstream_index)) =
1953            (snapshot.station_name.clone(), snapshot.upstream_index)
1954        else {
1955            return;
1956        };
1957        let now_ms = unix_now_ms();
1958        snapshot.refresh_status(now_ms);
1959
1960        let station_summary = {
1961            let mut guard = self.provider_balances.write().await;
1962            let station_balances = guard
1963                .entry(service_name.to_string())
1964                .or_default()
1965                .entry(station_name.clone())
1966                .or_default();
1967            if !snapshot.has_amount_data()
1968                && let Some(previous) = station_balances
1969                    .get(&upstream_index)
1970                    .and_then(|providers| providers.get(&snapshot.provider_id))
1971            {
1972                snapshot.carry_forward_amount_data_from(previous);
1973                snapshot.refresh_status(now_ms);
1974            }
1975            station_balances
1976                .entry(upstream_index)
1977                .or_default()
1978                .insert(snapshot.provider_id.clone(), snapshot.clone());
1979            StationRoutingBalanceSummary::from_snapshot_iter_at(
1980                station_balances
1981                    .values()
1982                    .flat_map(|providers| providers.values()),
1983                now_ms,
1984            )
1985        };
1986
1987        {
1988            let mut summaries = self.provider_balance_summaries.write().await;
1989            summaries
1990                .entry(service_name.to_string())
1991                .or_default()
1992                .insert(station_name, station_summary);
1993        }
1994
1995        let history_station_name = snapshot.station_name.clone().unwrap_or_default();
1996        self.record_provider_balance_history_snapshot(
1997            service_name,
1998            &history_station_name,
1999            upstream_index,
2000            snapshot,
2001        )
2002        .await;
2003        self.notify_state_changed();
2004    }
2005
2006    async fn record_provider_balance_history_snapshot(
2007        &self,
2008        service_name: &str,
2009        station_name: &str,
2010        upstream_index: usize,
2011        snapshot: ProviderBalanceSnapshot,
2012    ) {
2013        let mut guard = self.provider_balance_history.write().await;
2014        let history = guard
2015            .entry(service_name.to_string())
2016            .or_default()
2017            .entry(station_name.to_string())
2018            .or_default()
2019            .entry(upstream_index)
2020            .or_default()
2021            .entry(snapshot.provider_id.clone())
2022            .or_default();
2023        if history
2024            .back()
2025            .is_some_and(|previous| previous.fetched_at_ms == snapshot.fetched_at_ms)
2026        {
2027            history.pop_back();
2028        }
2029        history.push_back(snapshot);
2030        let max = provider_balance_history_max();
2031        while history.len() > max {
2032            history.pop_front();
2033        }
2034    }
2035
2036    pub async fn get_provider_balance_view(
2037        &self,
2038        service_name: &str,
2039    ) -> HashMap<String, Vec<ProviderBalanceSnapshot>> {
2040        let now_ms = unix_now_ms();
2041        let guard = self.provider_balances.read().await;
2042        let Some(per_service) = guard.get(service_name) else {
2043            return HashMap::new();
2044        };
2045
2046        per_service
2047            .iter()
2048            .map(|(station_name, upstreams)| {
2049                let mut snapshots = upstreams
2050                    .values()
2051                    .flat_map(|providers| providers.values().cloned())
2052                    .collect::<Vec<_>>();
2053                for snapshot in &mut snapshots {
2054                    snapshot.refresh_status(now_ms);
2055                }
2056                snapshots.sort_by(|a, b| {
2057                    a.upstream_index
2058                        .cmp(&b.upstream_index)
2059                        .then_with(|| a.provider_id.cmp(&b.provider_id))
2060                });
2061                (station_name.clone(), snapshots)
2062            })
2063            .collect()
2064    }
2065
2066    pub async fn get_provider_balance_history_view(
2067        &self,
2068        service_name: &str,
2069    ) -> HashMap<String, Vec<ProviderBalanceSnapshot>> {
2070        let now_ms = unix_now_ms();
2071        let guard = self.provider_balance_history.read().await;
2072        let Some(per_service) = guard.get(service_name) else {
2073            return HashMap::new();
2074        };
2075
2076        per_service
2077            .iter()
2078            .map(|(station_name, upstreams)| {
2079                let mut snapshots = upstreams
2080                    .values()
2081                    .flat_map(|providers| providers.values())
2082                    .flat_map(|history| history.iter().cloned())
2083                    .collect::<Vec<_>>();
2084                for snapshot in &mut snapshots {
2085                    snapshot.refresh_status(now_ms);
2086                }
2087                snapshots.sort_by(|a, b| {
2088                    a.fetched_at_ms
2089                        .cmp(&b.fetched_at_ms)
2090                        .then_with(|| a.upstream_index.cmp(&b.upstream_index))
2091                        .then_with(|| a.provider_id.cmp(&b.provider_id))
2092                });
2093                (station_name.clone(), snapshots)
2094            })
2095            .collect()
2096    }
2097
2098    pub async fn get_provider_balance_summary_view(
2099        &self,
2100        service_name: &str,
2101    ) -> HashMap<String, StationRoutingBalanceSummary> {
2102        let guard = self.provider_balance_summaries.read().await;
2103        let Some(per_service) = guard.get(service_name) else {
2104            return HashMap::new();
2105        };
2106
2107        per_service.clone()
2108    }
2109
2110    pub async fn record_passive_upstream_success(
2111        &self,
2112        service_name: &str,
2113        station_name: &str,
2114        base_url: &str,
2115        status_code: Option<u16>,
2116        now_ms: u64,
2117    ) {
2118        let mut guard = self.passive_station_health.write().await;
2119        let entry = guard
2120            .entry(service_name.to_string())
2121            .or_default()
2122            .entry(station_name.to_string())
2123            .or_default()
2124            .entry(base_url.to_string())
2125            .or_default();
2126        entry.record_success(now_ms, status_code);
2127        self.notify_state_changed();
2128    }
2129
2130    pub async fn record_passive_upstream_failure(&self, params: PassiveUpstreamFailureRecord) {
2131        let PassiveUpstreamFailureRecord {
2132            service_name,
2133            station_name,
2134            base_url,
2135            status_code,
2136            error_class,
2137            error,
2138            now_ms,
2139        } = params;
2140
2141        let mut guard = self.passive_station_health.write().await;
2142        let entry = guard
2143            .entry(service_name)
2144            .or_default()
2145            .entry(station_name)
2146            .or_default()
2147            .entry(base_url)
2148            .or_default();
2149        entry.record_failure(now_ms, status_code, error_class, error);
2150        self.notify_state_changed();
2151    }
2152
2153    pub async fn get_lb_view(&self) -> HashMap<String, LbConfigView> {
2154        let Some(lb_states) = self.lb_states.as_ref() else {
2155            return HashMap::new();
2156        };
2157        let mut map = match lb_states.lock() {
2158            Ok(m) => m,
2159            Err(e) => e.into_inner(),
2160        };
2161
2162        let now = std::time::Instant::now();
2163        let mut out = HashMap::new();
2164        for (cfg_name, st) in map.iter_mut() {
2165            let len = st
2166                .failure_counts
2167                .len()
2168                .max(st.cooldown_until.len())
2169                .max(st.usage_exhausted.len());
2170            if len == 0 {
2171                continue;
2172            }
2173
2174            // 如果结构变化导致长度不一致,做一次对齐,避免 UI 读到越界/脏数据。
2175            if st.failure_counts.len() != len {
2176                st.failure_counts.resize(len, 0);
2177            }
2178            if st.cooldown_until.len() != len {
2179                st.cooldown_until.resize(len, None);
2180            }
2181            if st.usage_exhausted.len() != len {
2182                st.usage_exhausted.resize(len, false);
2183            }
2184
2185            let mut upstreams = Vec::with_capacity(len);
2186            for idx in 0..len {
2187                let failure_count = st.failure_counts.get(idx).copied().unwrap_or(0);
2188                let cooldown_remaining_secs = st
2189                    .cooldown_until
2190                    .get(idx)
2191                    .and_then(|v| *v)
2192                    .map(|until| until.saturating_duration_since(now).as_secs())
2193                    .filter(|&s| s > 0);
2194                let usage_exhausted = st.usage_exhausted.get(idx).copied().unwrap_or(false);
2195                upstreams.push(LbUpstreamView {
2196                    failure_count,
2197                    cooldown_remaining_secs,
2198                    usage_exhausted,
2199                });
2200            }
2201
2202            out.insert(
2203                cfg_name.clone(),
2204                LbConfigView {
2205                    last_good_index: st.last_good_index,
2206                    upstreams,
2207                },
2208            );
2209        }
2210        out
2211    }
2212
2213    pub async fn list_health_checks(
2214        &self,
2215        service_name: &str,
2216    ) -> HashMap<String, HealthCheckStatus> {
2217        let guard = self.station_health_checks.read().await;
2218        guard.get(service_name).cloned().unwrap_or_default()
2219    }
2220
2221    pub async fn try_begin_station_health_check(
2222        &self,
2223        service_name: &str,
2224        station_name: &str,
2225        total: usize,
2226        now_ms: u64,
2227    ) -> bool {
2228        let mut guard = self.station_health_checks.write().await;
2229        let per_service = guard.entry(service_name.to_string()).or_default();
2230        if let Some(existing) = per_service.get(station_name)
2231            && !existing.done
2232        {
2233            return false;
2234        }
2235        per_service.insert(
2236            station_name.to_string(),
2237            HealthCheckStatus {
2238                started_at_ms: now_ms,
2239                updated_at_ms: now_ms,
2240                total: total.min(u32::MAX as usize) as u32,
2241                completed: 0,
2242                ok: 0,
2243                err: 0,
2244                cancel_requested: false,
2245                canceled: false,
2246                done: false,
2247                last_error: None,
2248            },
2249        );
2250        self.notify_state_changed();
2251        true
2252    }
2253
2254    pub async fn request_cancel_station_health_check(
2255        &self,
2256        service_name: &str,
2257        station_name: &str,
2258        now_ms: u64,
2259    ) -> bool {
2260        let mut guard = self.station_health_checks.write().await;
2261        let Some(per_service) = guard.get_mut(service_name) else {
2262            return false;
2263        };
2264        let Some(st) = per_service.get_mut(station_name) else {
2265            return false;
2266        };
2267        if st.done {
2268            return false;
2269        }
2270        st.cancel_requested = true;
2271        st.updated_at_ms = now_ms;
2272        self.notify_state_changed();
2273        true
2274    }
2275
2276    pub async fn is_station_health_check_cancel_requested(
2277        &self,
2278        service_name: &str,
2279        station_name: &str,
2280    ) -> bool {
2281        let guard = self.station_health_checks.read().await;
2282        guard
2283            .get(service_name)
2284            .and_then(|m| m.get(station_name))
2285            .is_some_and(|s| s.cancel_requested && !s.done)
2286    }
2287
2288    pub async fn record_station_health_check_result(
2289        &self,
2290        service_name: &str,
2291        station_name: &str,
2292        now_ms: u64,
2293        upstream: UpstreamHealth,
2294    ) {
2295        {
2296            let mut guard = self.station_health.write().await;
2297            let per_service = guard.entry(service_name.to_string()).or_default();
2298            let entry = per_service
2299                .entry(station_name.to_string())
2300                .or_insert_with(|| StationHealth {
2301                    checked_at_ms: now_ms,
2302                    upstreams: Vec::new(),
2303                });
2304            entry.checked_at_ms = entry.checked_at_ms.max(now_ms);
2305            entry.upstreams.push(upstream.clone());
2306            if entry.upstreams.len() > Self::MAX_HEALTH_RECORDS_PER_STATION {
2307                let extra = entry
2308                    .upstreams
2309                    .len()
2310                    .saturating_sub(Self::MAX_HEALTH_RECORDS_PER_STATION);
2311                if extra > 0 {
2312                    entry.upstreams.drain(0..extra);
2313                }
2314            }
2315        }
2316
2317        let mut guard = self.station_health_checks.write().await;
2318        let per_service = guard.entry(service_name.to_string()).or_default();
2319        let st = per_service.entry(station_name.to_string()).or_default();
2320        st.updated_at_ms = now_ms;
2321        st.completed = st.completed.saturating_add(1);
2322        match upstream.ok {
2323            Some(true) => st.ok = st.ok.saturating_add(1),
2324            Some(false) => {
2325                st.err = st.err.saturating_add(1);
2326                if st.last_error.is_none() {
2327                    st.last_error = upstream.error.clone();
2328                }
2329            }
2330            None => {}
2331        }
2332        self.notify_state_changed();
2333    }
2334
2335    pub async fn finish_station_health_check(
2336        &self,
2337        service_name: &str,
2338        station_name: &str,
2339        now_ms: u64,
2340        canceled: bool,
2341    ) {
2342        let mut guard = self.station_health_checks.write().await;
2343        let per_service = guard.entry(service_name.to_string()).or_default();
2344        let st = per_service.entry(station_name.to_string()).or_default();
2345        st.updated_at_ms = now_ms;
2346        st.canceled = canceled;
2347        st.done = true;
2348        self.notify_state_changed();
2349    }
2350
2351    pub async fn get_usage_rollup_view(
2352        &self,
2353        service_name: &str,
2354        top_n: usize,
2355        days: usize,
2356    ) -> UsageRollupView {
2357        let guard = self.usage_rollups.read().await;
2358        let Some(rollup) = guard.get(service_name) else {
2359            return UsageRollupView::default();
2360        };
2361
2362        fn now_day() -> i32 {
2363            std::time::SystemTime::now()
2364                .duration_since(std::time::UNIX_EPOCH)
2365                .map(|d| (d.as_millis() / 86_400_000) as i32)
2366                .unwrap_or(0)
2367        }
2368
2369        fn sorted_day_series(map: &HashMap<i32, UsageBucket>) -> Vec<(i32, UsageBucket)> {
2370            let mut out = map.iter().map(|(k, v)| (*k, v.clone())).collect::<Vec<_>>();
2371            out.sort_by_key(|(k, _)| *k);
2372            out
2373        }
2374
2375        fn filled_day_series(
2376            map: &HashMap<i32, UsageBucket>,
2377            start_day: i32,
2378            end_day: i32,
2379        ) -> Vec<(i32, UsageBucket)> {
2380            if start_day > end_day {
2381                return Vec::new();
2382            }
2383            (start_day..=end_day)
2384                .map(|day| (day, map.get(&day).cloned().unwrap_or_default()))
2385                .collect()
2386        }
2387
2388        fn sum_series(series: &[(i32, UsageBucket)]) -> UsageBucket {
2389            let mut out = UsageBucket::default();
2390            for (_, bucket) in series {
2391                out.add_assign(bucket);
2392            }
2393            out
2394        }
2395
2396        fn aggregate_entity_window(
2397            source: &HashMap<String, HashMap<i32, UsageBucket>>,
2398            start_day: Option<i32>,
2399            end_day: Option<i32>,
2400            top_n: usize,
2401        ) -> Vec<(String, UsageBucket)> {
2402            let mut out = Vec::new();
2403            for (name, days) in source {
2404                let mut bucket = UsageBucket::default();
2405                for (day, value) in days {
2406                    let include = match (start_day, end_day) {
2407                        (Some(start), Some(end)) => *day >= start && *day <= end,
2408                        _ => true,
2409                    };
2410                    if include {
2411                        bucket.add_assign(value);
2412                    }
2413                }
2414                if bucket.requests_total > 0 {
2415                    out.push((name.clone(), bucket));
2416                }
2417            }
2418            out.sort_by(|(left_name, left), (right_name, right)| {
2419                right
2420                    .usage
2421                    .total_tokens
2422                    .cmp(&left.usage.total_tokens)
2423                    .then_with(|| right.requests_total.cmp(&left.requests_total))
2424                    .then_with(|| left_name.cmp(right_name))
2425            });
2426            out.truncate(top_n);
2427            out
2428        }
2429
2430        let all_loaded = days == 0;
2431        let loaded_first_day = rollup.by_day.keys().min().copied();
2432        let loaded_last_day = rollup.by_day.keys().max().copied();
2433        let loaded_days_with_data = rollup
2434            .by_day
2435            .values()
2436            .filter(|bucket| bucket.requests_total > 0)
2437            .count();
2438
2439        let (start_day, end_day) = if all_loaded {
2440            (loaded_first_day, loaded_last_day)
2441        } else {
2442            let end = now_day();
2443            let offset = i32::try_from(days.saturating_sub(1)).unwrap_or(i32::MAX);
2444            (Some(end.saturating_sub(offset)), Some(end))
2445        };
2446
2447        let by_day = match (all_loaded, start_day, end_day) {
2448            (true, _, _) => sorted_day_series(&rollup.by_day),
2449            (false, Some(start), Some(end)) => filled_day_series(&rollup.by_day, start, end),
2450            _ => Vec::new(),
2451        };
2452        let window = if all_loaded {
2453            rollup.loaded.clone()
2454        } else {
2455            sum_series(&by_day)
2456        };
2457
2458        let mut by_config =
2459            aggregate_entity_window(&rollup.by_config_day, start_day, end_day, top_n);
2460        if all_loaded && by_config.is_empty() {
2461            by_config = rollup
2462                .by_config
2463                .iter()
2464                .map(|(k, v)| (k.clone(), v.clone()))
2465                .collect::<Vec<_>>();
2466            by_config.sort_by(|(left_name, left), (right_name, right)| {
2467                right
2468                    .usage
2469                    .total_tokens
2470                    .cmp(&left.usage.total_tokens)
2471                    .then_with(|| right.requests_total.cmp(&left.requests_total))
2472                    .then_with(|| left_name.cmp(right_name))
2473            });
2474            by_config.truncate(top_n);
2475        }
2476
2477        let mut by_provider =
2478            aggregate_entity_window(&rollup.by_provider_day, start_day, end_day, top_n);
2479        if all_loaded && by_provider.is_empty() {
2480            by_provider = rollup
2481                .by_provider
2482                .iter()
2483                .map(|(k, v)| (k.clone(), v.clone()))
2484                .collect::<Vec<_>>();
2485            by_provider.sort_by(|(left_name, left), (right_name, right)| {
2486                right
2487                    .usage
2488                    .total_tokens
2489                    .cmp(&left.usage.total_tokens)
2490                    .then_with(|| right.requests_total.cmp(&left.requests_total))
2491                    .then_with(|| left_name.cmp(right_name))
2492            });
2493            by_provider.truncate(top_n);
2494        }
2495
2496        let mut by_config_day = HashMap::new();
2497        for (name, _) in &by_config {
2498            let series = rollup
2499                .by_config_day
2500                .get(name)
2501                .map(|m| match (all_loaded, start_day, end_day) {
2502                    (true, _, _) => sorted_day_series(m),
2503                    (false, Some(start), Some(end)) => filled_day_series(m, start, end),
2504                    _ => Vec::new(),
2505                })
2506                .unwrap_or_default();
2507            by_config_day.insert(name.clone(), series);
2508        }
2509
2510        let mut by_provider_day = HashMap::new();
2511        for (name, _) in &by_provider {
2512            let series = rollup
2513                .by_provider_day
2514                .get(name)
2515                .map(|m| match (all_loaded, start_day, end_day) {
2516                    (true, _, _) => sorted_day_series(m),
2517                    (false, Some(start), Some(end)) => filled_day_series(m, start, end),
2518                    _ => Vec::new(),
2519                })
2520                .unwrap_or_default();
2521            by_provider_day.insert(name.clone(), series);
2522        }
2523
2524        let window_days_with_data = by_day
2525            .iter()
2526            .filter(|(_, bucket)| bucket.requests_total > 0)
2527            .count();
2528        let coverage = UsageRollupCoverage {
2529            requested_days: days,
2530            all_loaded,
2531            loaded_first_day,
2532            loaded_last_day,
2533            loaded_days_with_data,
2534            loaded_requests: rollup.loaded.requests_total,
2535            window_first_day: start_day,
2536            window_last_day: end_day,
2537            window_days_with_data,
2538            window_requests: window.requests_total,
2539            window_exceeds_loaded_start: matches!(
2540                (all_loaded, start_day, loaded_first_day),
2541                (false, Some(start), Some(first)) if start < first
2542            ),
2543        };
2544
2545        UsageRollupView {
2546            loaded: rollup.loaded.clone(),
2547            window,
2548            coverage,
2549            by_day,
2550            by_config,
2551            by_config_day,
2552            by_provider,
2553            by_provider_day,
2554        }
2555    }
2556
2557    pub async fn replay_usage_from_requests_log(
2558        &self,
2559        service_name: &str,
2560        log_path: PathBuf,
2561        base_url_to_provider_id: HashMap<String, String>,
2562    ) -> usize {
2563        let enabled = std::env::var("CODEX_HELPER_USAGE_REPLAY_ON_STARTUP")
2564            .ok()
2565            .map(|v| {
2566                matches!(
2567                    v.trim().to_ascii_lowercase().as_str(),
2568                    "1" | "true" | "yes" | "y" | "on"
2569                )
2570            })
2571            .unwrap_or(true);
2572        if !enabled {
2573            return 0;
2574        }
2575
2576        let already_has_data = {
2577            let guard = self.usage_rollups.read().await;
2578            guard
2579                .get(service_name)
2580                .is_some_and(|r| r.loaded.requests_total > 0)
2581        };
2582        if already_has_data {
2583            return 0;
2584        }
2585
2586        if !log_path.exists() {
2587            return 0;
2588        }
2589
2590        let max_bytes = std::env::var("CODEX_HELPER_USAGE_REPLAY_MAX_BYTES")
2591            .ok()
2592            .and_then(|s| s.trim().parse::<usize>().ok())
2593            .filter(|&n| n > 0)
2594            .unwrap_or(8 * 1024 * 1024);
2595        let max_lines = std::env::var("CODEX_HELPER_USAGE_REPLAY_MAX_LINES")
2596            .ok()
2597            .and_then(|s| s.trim().parse::<usize>().ok())
2598            .filter(|&n| n > 0)
2599            .unwrap_or(20_000);
2600
2601        let mut file = match std::fs::File::open(&log_path) {
2602            Ok(f) => f,
2603            Err(_) => return 0,
2604        };
2605        let len: u64 = file.metadata().map(|m| m.len()).unwrap_or_default();
2606        let start = len.saturating_sub(max_bytes as u64);
2607        if file.seek(SeekFrom::Start(start)).is_err() {
2608            return 0;
2609        }
2610        let mut buf = Vec::new();
2611        if file.read_to_end(&mut buf).is_err() {
2612            return 0;
2613        }
2614        if start > 0 {
2615            if let Some(pos) = buf.iter().position(|b| *b == b'\n') {
2616                buf = buf[pos + 1..].to_vec();
2617            } else {
2618                return 0;
2619            }
2620        }
2621
2622        let text = match std::str::from_utf8(&buf) {
2623            Ok(s) => s,
2624            Err(_) => return 0,
2625        };
2626        let lines = text
2627            .lines()
2628            .map(|l| l.trim())
2629            .filter(|l| !l.is_empty())
2630            .collect::<Vec<_>>();
2631        let start_idx = lines.len().saturating_sub(max_lines);
2632
2633        let mut events = Vec::new();
2634        for line in &lines[start_idx..] {
2635            let Ok(v) = serde_json::from_str::<JsonValue>(line) else {
2636                continue;
2637            };
2638            let Some(svc) = v.get("service").and_then(|x| x.as_str()) else {
2639                continue;
2640            };
2641            if svc != service_name {
2642                continue;
2643            }
2644
2645            let ended_at_ms = v.get("timestamp_ms").and_then(|x| x.as_u64()).unwrap_or(0);
2646            let status_code = v.get("status_code").and_then(|x| x.as_u64()).unwrap_or(0) as u16;
2647            let duration_ms = v.get("duration_ms").and_then(|x| x.as_u64()).unwrap_or(0);
2648            let station_name = v
2649                .get("station_name")
2650                .and_then(|x| x.as_str())
2651                .unwrap_or("-")
2652                .to_string();
2653            let upstream_base_url = v
2654                .get("upstream_base_url")
2655                .and_then(|x| x.as_str())
2656                .unwrap_or("-")
2657                .to_string();
2658            let provider_id = v
2659                .get("provider_id")
2660                .and_then(|x| x.as_str())
2661                .map(|s| s.to_string())
2662                .or_else(|| base_url_to_provider_id.get(&upstream_base_url).cloned())
2663                .unwrap_or_else(|| "-".to_string());
2664            let usage = v
2665                .get("usage")
2666                .and_then(|u| serde_json::from_value::<UsageMetrics>(u.clone()).ok());
2667            let ttfb_ms = v.get("ttfb_ms").and_then(|x| x.as_u64());
2668
2669            events.push((
2670                ended_at_ms,
2671                status_code,
2672                duration_ms,
2673                station_name,
2674                provider_id,
2675                usage,
2676                ttfb_ms,
2677            ));
2678        }
2679
2680        if events.is_empty() {
2681            return 0;
2682        }
2683
2684        let mut guard = self.usage_rollups.write().await;
2685        let rollup = guard.entry(service_name.to_string()).or_default();
2686        for (ended_at_ms, status_code, duration_ms, cfg_key, provider_key, usage, ttfb_ms) in
2687            events.iter()
2688        {
2689            let day = (*ended_at_ms / 86_400_000) as i32;
2690            rollup
2691                .loaded
2692                .record(*status_code, *duration_ms, usage.as_ref(), None, *ttfb_ms);
2693            rollup.by_day.entry(day).or_default().record(
2694                *status_code,
2695                *duration_ms,
2696                usage.as_ref(),
2697                None,
2698                *ttfb_ms,
2699            );
2700            rollup.by_config.entry(cfg_key.clone()).or_default().record(
2701                *status_code,
2702                *duration_ms,
2703                usage.as_ref(),
2704                None,
2705                *ttfb_ms,
2706            );
2707            rollup
2708                .by_config_day
2709                .entry(cfg_key.clone())
2710                .or_default()
2711                .entry(day)
2712                .or_default()
2713                .record(*status_code, *duration_ms, usage.as_ref(), None, *ttfb_ms);
2714            rollup
2715                .by_provider
2716                .entry(provider_key.clone())
2717                .or_default()
2718                .record(*status_code, *duration_ms, usage.as_ref(), None, *ttfb_ms);
2719            rollup
2720                .by_provider_day
2721                .entry(provider_key.clone())
2722                .or_default()
2723                .entry(day)
2724                .or_default()
2725                .record(*status_code, *duration_ms, usage.as_ref(), None, *ttfb_ms);
2726        }
2727
2728        let replayed = events.len();
2729        self.notify_state_changed();
2730        replayed
2731    }
2732
2733    pub async fn resolve_session_cwd(&self, session_id: &str) -> Option<String> {
2734        if self.session_cwd_cache_max_entries == 0 {
2735            return sessions::find_codex_session_cwd_by_id(session_id)
2736                .await
2737                .ok()
2738                .flatten();
2739        }
2740
2741        let now_ms = std::time::SystemTime::now()
2742            .duration_since(std::time::UNIX_EPOCH)
2743            .map(|d| d.as_millis() as u64)
2744            .unwrap_or(0);
2745
2746        {
2747            let guard = self.session_cwd_cache.read().await;
2748            if let Some(v) = guard.get(session_id) {
2749                let out = v.cwd.clone();
2750                drop(guard);
2751                let mut guard = self.session_cwd_cache.write().await;
2752                if let Some(v) = guard.get_mut(session_id) {
2753                    v.last_seen_ms = now_ms;
2754                }
2755                return out;
2756            }
2757        }
2758
2759        // Cache miss: resolve from disk and record last_seen.
2760
2761        let resolved = sessions::find_codex_session_cwd_by_id(session_id)
2762            .await
2763            .ok()
2764            .flatten();
2765
2766        let mut guard = self.session_cwd_cache.write().await;
2767        guard.insert(
2768            session_id.to_string(),
2769            SessionCwdCacheEntry {
2770                cwd: resolved.clone(),
2771                last_seen_ms: now_ms,
2772            },
2773        );
2774        resolved
2775    }
2776
2777    #[allow(clippy::too_many_arguments)]
2778    pub async fn begin_request(
2779        &self,
2780        service: &str,
2781        method: &str,
2782        path: &str,
2783        session_id: Option<String>,
2784        session_identity_source: Option<SessionIdentitySource>,
2785        client_name: Option<String>,
2786        client_addr: Option<String>,
2787        cwd: Option<String>,
2788        model: Option<String>,
2789        reasoning_effort: Option<String>,
2790        service_tier: Option<String>,
2791        started_at_ms: u64,
2792    ) -> u64 {
2793        let id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
2794        let trace_id = Some(crate::logging::request_trace_id(service, id));
2795        let req = ActiveRequest {
2796            id,
2797            trace_id,
2798            session_id,
2799            session_identity_source,
2800            client_name,
2801            client_addr,
2802            cwd,
2803            model,
2804            reasoning_effort,
2805            service_tier,
2806            station_name: None,
2807            provider_id: None,
2808            upstream_base_url: None,
2809            route_decision: None,
2810            service: service.to_string(),
2811            method: method.to_string(),
2812            path: path.to_string(),
2813            started_at_ms,
2814        };
2815        let mut guard = self.active_requests.write().await;
2816        guard.insert(id, req);
2817        self.notify_state_changed();
2818        id
2819    }
2820
2821    #[cfg(test)]
2822    pub(crate) fn begin_request_for_test(&self) -> BeginRequestTestBuilder<'_> {
2823        BeginRequestTestBuilder::new(self)
2824    }
2825
2826    pub async fn update_request_route(
2827        &self,
2828        request_id: u64,
2829        station_name: Option<String>,
2830        provider_id: Option<String>,
2831        upstream_base_url: String,
2832        route_decision: Option<RouteDecisionProvenance>,
2833    ) {
2834        let mut guard = self.active_requests.write().await;
2835        let Some(req) = guard.get_mut(&request_id) else {
2836            return;
2837        };
2838        req.station_name = station_name;
2839        req.provider_id = provider_id;
2840        req.upstream_base_url = Some(upstream_base_url);
2841        req.route_decision = route_decision;
2842        self.notify_state_changed();
2843    }
2844
2845    pub async fn finish_request(&self, params: FinishRequestParams) -> bool {
2846        let mut active = self.active_requests.write().await;
2847        let Some(req) = active.remove(&params.id) else {
2848            return false;
2849        };
2850        drop(active);
2851
2852        let pricing_model = req
2853            .route_decision
2854            .as_ref()
2855            .and_then(|decision| decision.effective_model.as_ref())
2856            .map(|value| value.value.as_str())
2857            .or(req.model.as_deref());
2858        let cost = estimate_request_cost_from_operator_catalog_for_service(
2859            pricing_model,
2860            params.usage.as_ref(),
2861            CostAdjustments::default(),
2862            &req.service,
2863        );
2864
2865        let mut finished = FinishedRequest {
2866            id: params.id,
2867            trace_id: req.trace_id,
2868            session_id: req.session_id,
2869            session_identity_source: req.session_identity_source,
2870            client_name: req.client_name,
2871            client_addr: req.client_addr,
2872            cwd: req.cwd,
2873            model: req.model,
2874            reasoning_effort: req.reasoning_effort,
2875            service_tier: params.observed_service_tier.or(req.service_tier),
2876            station_name: req.station_name,
2877            provider_id: req.provider_id,
2878            upstream_base_url: req.upstream_base_url,
2879            route_decision: req.route_decision,
2880            usage: params.usage.clone(),
2881            cost,
2882            retry: params.retry,
2883            provider_signals: Vec::new(),
2884            policy_actions: Vec::new(),
2885            observability: RequestObservability::default(),
2886            service: req.service,
2887            method: req.method,
2888            path: req.path,
2889            status_code: params.status_code,
2890            duration_ms: params.duration_ms,
2891            ttfb_ms: params.ttfb_ms,
2892            streaming: params.streaming,
2893            ended_at_ms: params.ended_at_ms,
2894        };
2895        finished.refresh_observability();
2896
2897        {
2898            let effective_ttfb_ms = finished.observability.ttfb_ms;
2899            let day = (finished.ended_at_ms / 86_400_000) as i32;
2900            let cfg_key = finished
2901                .station_name
2902                .clone()
2903                .unwrap_or_else(|| "-".to_string());
2904            let provider_key = finished
2905                .provider_id
2906                .clone()
2907                .unwrap_or_else(|| "-".to_string());
2908
2909            let mut rollups = self.usage_rollups.write().await;
2910            let rollup = rollups.entry(finished.service.clone()).or_default();
2911            rollup.loaded.record(
2912                finished.status_code,
2913                finished.duration_ms,
2914                finished.usage.as_ref(),
2915                Some(&finished.cost),
2916                effective_ttfb_ms,
2917            );
2918            rollup.by_day.entry(day).or_default().record(
2919                finished.status_code,
2920                finished.duration_ms,
2921                finished.usage.as_ref(),
2922                Some(&finished.cost),
2923                effective_ttfb_ms,
2924            );
2925            rollup.by_config.entry(cfg_key.clone()).or_default().record(
2926                finished.status_code,
2927                finished.duration_ms,
2928                finished.usage.as_ref(),
2929                Some(&finished.cost),
2930                effective_ttfb_ms,
2931            );
2932            rollup
2933                .by_config_day
2934                .entry(cfg_key)
2935                .or_default()
2936                .entry(day)
2937                .or_default()
2938                .record(
2939                    finished.status_code,
2940                    finished.duration_ms,
2941                    finished.usage.as_ref(),
2942                    Some(&finished.cost),
2943                    effective_ttfb_ms,
2944                );
2945
2946            rollup
2947                .by_provider
2948                .entry(provider_key.clone())
2949                .or_default()
2950                .record(
2951                    finished.status_code,
2952                    finished.duration_ms,
2953                    finished.usage.as_ref(),
2954                    Some(&finished.cost),
2955                    effective_ttfb_ms,
2956                );
2957            rollup
2958                .by_provider_day
2959                .entry(provider_key)
2960                .or_default()
2961                .entry(day)
2962                .or_default()
2963                .record(
2964                    finished.status_code,
2965                    finished.duration_ms,
2966                    finished.usage.as_ref(),
2967                    Some(&finished.cost),
2968                    effective_ttfb_ms,
2969                );
2970        }
2971
2972        if let Some(sid) = finished.session_id.as_deref() {
2973            let mut stats = self.session_stats.write().await;
2974            let entry = stats.entry(sid.to_string()).or_default();
2975            entry.turns_total = entry.turns_total.saturating_add(1);
2976            if finished.session_identity_source.is_some() {
2977                entry.last_session_identity_source = finished.session_identity_source;
2978            }
2979            entry.last_client_name = finished
2980                .client_name
2981                .clone()
2982                .or(entry.last_client_name.clone());
2983            entry.last_client_addr = finished
2984                .client_addr
2985                .clone()
2986                .or(entry.last_client_addr.clone());
2987            entry.last_model = finished.model.clone().or(entry.last_model.clone());
2988            entry.last_reasoning_effort = finished
2989                .reasoning_effort
2990                .clone()
2991                .or(entry.last_reasoning_effort.clone());
2992            entry.last_service_tier = finished
2993                .service_tier
2994                .clone()
2995                .or(entry.last_service_tier.clone());
2996            entry.last_provider_id = finished
2997                .provider_id
2998                .clone()
2999                .or(entry.last_provider_id.clone());
3000            entry.last_station_name = finished
3001                .station_name
3002                .clone()
3003                .or(entry.last_station_name.clone());
3004            if finished.route_decision.is_some() {
3005                entry.last_route_decision = finished.route_decision.clone();
3006            }
3007            if let Some(u) = finished.usage.as_ref() {
3008                entry.last_usage = Some(u.clone());
3009                entry.total_usage.add_assign(u);
3010                entry.turns_with_usage = entry.turns_with_usage.saturating_add(1);
3011            }
3012            if finished
3013                .usage
3014                .as_ref()
3015                .is_some_and(|usage| usage.output_tokens > 0)
3016            {
3017                entry.last_output_tokens_per_second =
3018                    finished.observability.output_tokens_per_second;
3019                if let Some(generation_ms) = finished.observability.generation_ms {
3020                    entry.output_generation_ms_total = entry
3021                        .output_generation_ms_total
3022                        .saturating_add(generation_ms);
3023                    entry.avg_output_tokens_per_second =
3024                        self::session_identity::token_weighted_output_tokens_per_second(
3025                            entry.total_usage.output_tokens,
3026                            entry.output_generation_ms_total,
3027                        );
3028                }
3029            }
3030            entry.last_status = Some(finished.status_code);
3031            entry.last_duration_ms = Some(finished.duration_ms);
3032            entry.last_ended_at_ms = Some(finished.ended_at_ms);
3033            entry.last_seen_ms = finished.ended_at_ms;
3034        }
3035
3036        let mut recent = self.recent_finished.write().await;
3037        recent.push_front(finished);
3038        while recent.len() > recent_finished_max() {
3039            recent.pop_back();
3040        }
3041        self.notify_state_changed();
3042        true
3043    }
3044
3045    pub async fn list_active_requests(&self) -> Vec<ActiveRequest> {
3046        let guard = self.active_requests.read().await;
3047        let mut vec = guard.values().cloned().collect::<Vec<_>>();
3048        vec.sort_by_key(|r| r.started_at_ms);
3049        vec
3050    }
3051
3052    pub async fn list_recent_finished(&self, limit: usize) -> Vec<FinishedRequest> {
3053        let guard = self.recent_finished.read().await;
3054        guard.iter().take(limit).cloned().collect()
3055    }
3056
3057    pub async fn list_session_stats(&self) -> HashMap<String, SessionStats> {
3058        let guard = self.session_stats.read().await;
3059        guard.clone()
3060    }
3061
3062    pub async fn list_session_identity_cards(
3063        &self,
3064        recent_limit: usize,
3065    ) -> Vec<SessionIdentityCard> {
3066        let recent_limit = recent_limit.clamp(1, recent_finished_max());
3067        let (
3068            active,
3069            recent,
3070            overrides,
3071            station_overrides,
3072            model_overrides,
3073            service_tier_overrides,
3074            bindings,
3075            route_affinities,
3076            global_station_override,
3077            stats,
3078        ) = tokio::join!(
3079            self.list_active_requests(),
3080            self.list_recent_finished(recent_limit),
3081            self.list_session_effort_overrides(),
3082            self.list_session_station_overrides(),
3083            self.list_session_model_overrides(),
3084            self.list_session_service_tier_overrides(),
3085            self.list_session_bindings(),
3086            self.list_session_route_affinities(),
3087            self.get_global_station_override(),
3088            self.list_session_stats(),
3089        );
3090        build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
3091            active: &active,
3092            recent: &recent,
3093            overrides: &overrides,
3094            station_overrides: &station_overrides,
3095            model_overrides: &model_overrides,
3096            service_tier_overrides: &service_tier_overrides,
3097            bindings: &bindings,
3098            route_affinities: &route_affinities,
3099            global_station_override: global_station_override.as_deref(),
3100            stats: &stats,
3101        })
3102    }
3103
3104    async fn resolve_host_transcript_paths_cached(
3105        &self,
3106        session_ids: &[String],
3107    ) -> HashMap<String, Option<String>> {
3108        let mut unique = session_ids
3109            .iter()
3110            .map(|sid| sid.trim())
3111            .filter(|sid| !sid.is_empty())
3112            .map(ToOwned::to_owned)
3113            .collect::<Vec<_>>();
3114        unique.sort();
3115        unique.dedup();
3116        if unique.is_empty() {
3117            return HashMap::new();
3118        }
3119
3120        if self.session_transcript_path_cache_max_entries == 0 {
3121            return sessions::find_codex_session_files_by_ids(&unique)
3122                .await
3123                .unwrap_or_default()
3124                .into_iter()
3125                .map(|(sid, path)| (sid, Some(path.to_string_lossy().to_string())))
3126                .collect();
3127        }
3128
3129        let now_ms = unix_now_ms();
3130        let ttl_ms = self.session_transcript_path_cache_ttl_ms;
3131        let mut resolved = HashMap::<String, Option<String>>::new();
3132        let mut stale_or_missing = Vec::<String>::new();
3133
3134        {
3135            let cache = self.session_transcript_path_cache.read().await;
3136            for sid in &unique {
3137                let fresh = cache.get(sid).filter(|entry| {
3138                    ttl_ms == 0 || now_ms.saturating_sub(entry.last_checked_ms) <= ttl_ms
3139                });
3140                if let Some(entry) = fresh {
3141                    resolved.insert(sid.clone(), entry.path.clone());
3142                } else {
3143                    stale_or_missing.push(sid.clone());
3144                }
3145            }
3146        }
3147
3148        if !stale_or_missing.is_empty() {
3149            let found = sessions::find_codex_session_files_by_ids(&stale_or_missing)
3150                .await
3151                .unwrap_or_default();
3152            let mut cache = self.session_transcript_path_cache.write().await;
3153            for sid in stale_or_missing {
3154                let path = found
3155                    .get(&sid)
3156                    .map(|path| path.to_string_lossy().to_string());
3157                cache.insert(
3158                    sid.clone(),
3159                    SessionTranscriptPathCacheEntry {
3160                        path: path.clone(),
3161                        last_checked_ms: now_ms,
3162                        last_seen_ms: now_ms,
3163                    },
3164                );
3165                resolved.insert(sid, path);
3166            }
3167            prune_lru_cache(
3168                &mut cache,
3169                self.session_transcript_path_cache_max_entries,
3170                |entry| entry.last_seen_ms,
3171            );
3172        }
3173
3174        {
3175            let mut cache = self.session_transcript_path_cache.write().await;
3176            for sid in &unique {
3177                if let Some(entry) = cache.get_mut(sid) {
3178                    entry.last_seen_ms = now_ms;
3179                }
3180            }
3181        }
3182
3183        resolved
3184    }
3185
3186    pub async fn enrich_session_identity_cards_with_cached_host_transcripts(
3187        &self,
3188        cards: &mut [SessionIdentityCard],
3189    ) {
3190        let session_ids = cards
3191            .iter()
3192            .filter_map(|card| card.session_id.as_deref())
3193            .map(str::to_owned)
3194            .collect::<Vec<_>>();
3195        let resolved = self
3196            .resolve_host_transcript_paths_cached(&session_ids)
3197            .await;
3198        for card in cards {
3199            card.host_local_transcript_path = card
3200                .session_id
3201                .as_deref()
3202                .and_then(|sid| resolved.get(sid))
3203                .and_then(Clone::clone);
3204        }
3205    }
3206
3207    pub async fn list_session_identity_cards_with_host_transcripts(
3208        &self,
3209        recent_limit: usize,
3210    ) -> Vec<SessionIdentityCard> {
3211        let mut cards = self.list_session_identity_cards(recent_limit).await;
3212        self.enrich_session_identity_cards_with_cached_host_transcripts(&mut cards)
3213            .await;
3214        cards
3215    }
3216
3217    pub fn spawn_cleanup_task(state: Arc<Self>) {
3218        // Run periodically; no need to be super frequent.
3219        tokio::spawn(async move {
3220            let mut tick = interval(Duration::from_secs(30));
3221            loop {
3222                tick.tick().await;
3223                state.prune_periodic().await;
3224            }
3225        });
3226    }
3227
3228    async fn prune_periodic(&self) {
3229        let now_ms = std::time::SystemTime::now()
3230            .duration_since(std::time::UNIX_EPOCH)
3231            .map(|d| d.as_millis() as u64)
3232            .unwrap_or(0);
3233
3234        // Collect active session_ids to avoid clearing overrides for currently running requests.
3235        let active = self.active_requests.read().await;
3236        let mut active_sessions: HashMap<String, ()> = HashMap::new();
3237        for req in active.values() {
3238            if let Some(sid) = req.session_id.as_deref() {
3239                active_sessions.insert(sid.to_string(), ());
3240            }
3241        }
3242
3243        if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
3244            let cutoff_override = now_ms - self.session_override_ttl_ms;
3245            let mut overrides = self.session_effort_overrides.write().await;
3246            overrides.retain(|sid, v| {
3247                if active_sessions.contains_key(sid) {
3248                    return true;
3249                }
3250                v.last_seen_ms >= cutoff_override
3251            });
3252        }
3253
3254        if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
3255            let cutoff_override = now_ms - self.session_override_ttl_ms;
3256            let mut overrides = self.session_station_overrides.write().await;
3257            overrides.retain(|sid, v| {
3258                if active_sessions.contains_key(sid) {
3259                    return true;
3260                }
3261                v.last_seen_ms >= cutoff_override
3262            });
3263        }
3264
3265        if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
3266            let cutoff_override = now_ms - self.session_override_ttl_ms;
3267            let mut overrides = self.session_route_target_overrides.write().await;
3268            overrides.retain(|sid, v| {
3269                if active_sessions.contains_key(sid) {
3270                    return true;
3271                }
3272                v.last_seen_ms >= cutoff_override
3273            });
3274        }
3275
3276        if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
3277            let cutoff_override = now_ms - self.session_override_ttl_ms;
3278            let mut overrides = self.session_model_overrides.write().await;
3279            overrides.retain(|sid, v| {
3280                if active_sessions.contains_key(sid) {
3281                    return true;
3282                }
3283                v.last_seen_ms >= cutoff_override
3284            });
3285        }
3286
3287        if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
3288            let cutoff_override = now_ms - self.session_override_ttl_ms;
3289            let mut overrides = self.session_service_tier_overrides.write().await;
3290            overrides.retain(|sid, v| {
3291                if active_sessions.contains_key(sid) {
3292                    return true;
3293                }
3294                v.last_seen_ms >= cutoff_override
3295            });
3296        }
3297
3298        if self.session_binding_ttl_ms > 0 && now_ms >= self.session_binding_ttl_ms {
3299            let cutoff_binding = now_ms - self.session_binding_ttl_ms;
3300            let mut bindings = self.session_bindings.write().await;
3301            bindings.retain(|sid, entry| {
3302                if active_sessions.contains_key(sid) {
3303                    return true;
3304                }
3305                entry.binding.last_seen_ms >= cutoff_binding
3306            });
3307        }
3308        if self.session_binding_max_entries > 0 {
3309            let mut bindings = self.session_bindings.write().await;
3310            if bindings.len() > self.session_binding_max_entries {
3311                let mut removable = bindings
3312                    .iter()
3313                    .filter(|(sid, _)| !active_sessions.contains_key(*sid))
3314                    .map(|(sid, entry)| (sid.clone(), entry.binding.last_seen_ms))
3315                    .collect::<Vec<_>>();
3316                removable.sort_by_key(|(_, last_seen_ms)| *last_seen_ms);
3317                let remove_count = bindings
3318                    .len()
3319                    .saturating_sub(self.session_binding_max_entries)
3320                    .min(removable.len());
3321                for (sid, _) in removable.into_iter().take(remove_count) {
3322                    bindings.remove(&sid);
3323                }
3324            }
3325        }
3326
3327        {
3328            let mut affinities = self.session_route_affinities.write().await;
3329            if self.session_route_affinity_ttl_ms > 0
3330                && now_ms >= self.session_route_affinity_ttl_ms
3331            {
3332                let cutoff_affinity = now_ms - self.session_route_affinity_ttl_ms;
3333                affinities.retain(|sid, affinity| {
3334                    active_sessions.contains_key(sid)
3335                        || affinity.last_selected_at_ms >= cutoff_affinity
3336                });
3337            }
3338            prune_lru_cache(
3339                &mut affinities,
3340                self.session_route_affinity_max_entries,
3341                |entry| entry.last_selected_at_ms,
3342            );
3343        }
3344
3345        // Keep a bounded number of days of rollup data to avoid unbounded growth.
3346        let keep_days: i32 = std::env::var("CODEX_HELPER_USAGE_ROLLUP_KEEP_DAYS")
3347            .ok()
3348            .and_then(|s| s.trim().parse::<i32>().ok())
3349            .filter(|&n| n > 0)
3350            .unwrap_or(60);
3351        let now_day = (now_ms / 86_400_000) as i32;
3352        let cutoff_day = now_day.saturating_sub(keep_days);
3353        let mut rollups = self.usage_rollups.write().await;
3354        for rollup in rollups.values_mut() {
3355            rollup.by_day.retain(|day, _| *day >= cutoff_day);
3356            rollup.by_config_day.retain(|_, m| {
3357                m.retain(|day, _| *day >= cutoff_day);
3358                !m.is_empty()
3359            });
3360            rollup.by_provider_day.retain(|_, m| {
3361                m.retain(|day, _| *day >= cutoff_day);
3362                !m.is_empty()
3363            });
3364        }
3365
3366        let cutoff_cwd =
3367            if self.session_cwd_cache_ttl_ms == 0 || now_ms < self.session_cwd_cache_ttl_ms {
3368                0
3369            } else {
3370                now_ms - self.session_cwd_cache_ttl_ms
3371            };
3372        self.prune_session_cwd_cache(&active_sessions, cutoff_cwd)
3373            .await;
3374        self.prune_session_transcript_path_cache(now_ms).await;
3375
3376        if self.session_override_ttl_ms > 0 && now_ms >= self.session_override_ttl_ms {
3377            let cutoff_stats = now_ms - self.session_override_ttl_ms;
3378            let mut stats = self.session_stats.write().await;
3379            stats.retain(|sid, v| {
3380                active_sessions.contains_key(sid) || v.last_seen_ms >= cutoff_stats
3381            });
3382        }
3383    }
3384
3385    async fn prune_session_cwd_cache(&self, active_sessions: &HashMap<String, ()>, cutoff: u64) {
3386        if self.session_cwd_cache_max_entries == 0 {
3387            return;
3388        }
3389        let mut cache = self.session_cwd_cache.write().await;
3390
3391        if self.session_cwd_cache_ttl_ms > 0 {
3392            cache.retain(|sid, v| {
3393                if active_sessions.contains_key(sid) {
3394                    return true;
3395                }
3396                v.last_seen_ms >= cutoff
3397            });
3398        }
3399
3400        let max = self.session_cwd_cache_max_entries;
3401        if max == 0 || cache.len() <= max {
3402            return;
3403        }
3404
3405        // Drop least-recently-seen entries first.
3406        let mut keys = cache
3407            .iter()
3408            .map(|(sid, v)| (sid.clone(), v.last_seen_ms))
3409            .collect::<Vec<_>>();
3410        keys.sort_by_key(|(_, t)| *t);
3411        let remove_count = keys.len().saturating_sub(max);
3412        for (sid, _) in keys.into_iter().take(remove_count) {
3413            cache.remove(&sid);
3414        }
3415    }
3416
3417    async fn prune_session_transcript_path_cache(&self, now_ms: u64) {
3418        let mut cache = self.session_transcript_path_cache.write().await;
3419        if self.session_transcript_path_cache_max_entries == 0 {
3420            cache.clear();
3421            return;
3422        }
3423
3424        if self.session_transcript_path_cache_ttl_ms > 0 {
3425            let cutoff = now_ms.saturating_sub(self.session_transcript_path_cache_ttl_ms);
3426            cache.retain(|_, entry| entry.last_seen_ms >= cutoff);
3427        }
3428
3429        prune_lru_cache(
3430            &mut cache,
3431            self.session_transcript_path_cache_max_entries,
3432            |entry| entry.last_seen_ms,
3433        );
3434    }
3435}
3436
3437#[cfg(test)]
3438#[derive(Debug)]
3439pub(crate) struct BeginRequestTestBuilder<'a> {
3440    state: &'a ProxyState,
3441    service: &'static str,
3442    method: &'static str,
3443    path: &'static str,
3444    session_id: Option<String>,
3445    session_identity_source: Option<SessionIdentitySource>,
3446    client_name: Option<String>,
3447    client_addr: Option<String>,
3448    cwd: Option<String>,
3449    model: Option<String>,
3450    reasoning_effort: Option<String>,
3451    service_tier: Option<String>,
3452    started_at_ms: u64,
3453}
3454
3455#[cfg(test)]
3456impl<'a> BeginRequestTestBuilder<'a> {
3457    pub(crate) fn new(state: &'a ProxyState) -> Self {
3458        Self {
3459            state,
3460            service: "codex",
3461            method: "POST",
3462            path: "/v1/responses",
3463            session_id: None,
3464            session_identity_source: None,
3465            client_name: None,
3466            client_addr: None,
3467            cwd: None,
3468            model: None,
3469            reasoning_effort: None,
3470            service_tier: None,
3471            started_at_ms: 0,
3472        }
3473    }
3474
3475    pub(crate) fn session_id(mut self, value: impl Into<String>) -> Self {
3476        self.session_id = Some(value.into());
3477        self
3478    }
3479
3480    pub(crate) fn model(mut self, value: impl Into<String>) -> Self {
3481        self.model = Some(value.into());
3482        self
3483    }
3484
3485    pub(crate) fn service_tier(mut self, value: impl Into<String>) -> Self {
3486        self.service_tier = Some(value.into());
3487        self
3488    }
3489
3490    pub(crate) fn started_at_ms(mut self, value: u64) -> Self {
3491        self.started_at_ms = value;
3492        self
3493    }
3494
3495    pub(crate) async fn begin(self) -> u64 {
3496        self.state
3497            .begin_request(
3498                self.service,
3499                self.method,
3500                self.path,
3501                self.session_id,
3502                self.session_identity_source,
3503                self.client_name,
3504                self.client_addr,
3505                self.cwd,
3506                self.model,
3507                self.reasoning_effort,
3508                self.service_tier,
3509                self.started_at_ms,
3510            )
3511            .await
3512    }
3513}
3514
3515#[cfg(test)]
3516mod tests {
3517    use super::*;
3518
3519    use crate::config::{ServiceConfig, ServiceConfigManager, UpstreamAuth, UpstreamConfig};
3520    use crate::runtime_identity::ProviderEndpointKey;
3521    use std::path::Path;
3522    use std::sync::OnceLock;
3523
3524    async fn env_lock() -> tokio::sync::MutexGuard<'static, ()> {
3525        static LOCK: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
3526        LOCK.get_or_init(|| tokio::sync::Mutex::new(()))
3527            .lock()
3528            .await
3529    }
3530
3531    #[derive(Default)]
3532    struct ScopedEnv {
3533        saved: Vec<(String, Option<String>)>,
3534    }
3535
3536    impl ScopedEnv {
3537        unsafe fn set(&mut self, key: &str, value: &str) {
3538            if !self.saved.iter().any(|(saved_key, _)| saved_key == key) {
3539                self.saved.push((key.to_string(), std::env::var(key).ok()));
3540            }
3541            unsafe {
3542                std::env::set_var(key, value);
3543            }
3544        }
3545
3546        unsafe fn set_path(&mut self, key: &str, value: &Path) {
3547            unsafe {
3548                self.set(key, value.to_string_lossy().as_ref());
3549            }
3550        }
3551    }
3552
3553    impl Drop for ScopedEnv {
3554        fn drop(&mut self) {
3555            for (key, value) in self.saved.iter().rev() {
3556                match value {
3557                    Some(value) => unsafe {
3558                        std::env::set_var(key, value);
3559                    },
3560                    None => unsafe {
3561                        std::env::remove_var(key);
3562                    },
3563                }
3564            }
3565        }
3566    }
3567
3568    fn test_runtime_policy(
3569        session_override_ttl_ms: u64,
3570        session_binding_ttl_ms: u64,
3571        session_binding_max_entries: usize,
3572    ) -> RuntimePolicy {
3573        RuntimePolicy {
3574            session_override_ttl_ms,
3575            session_binding_ttl_ms,
3576            session_binding_max_entries,
3577            session_route_affinity_ttl_ms: 0,
3578            session_route_affinity_max_entries: 5_000,
3579            session_route_affinity_store: SessionRouteAffinityStore::from_env(),
3580            session_cwd_cache_ttl_ms: 0,
3581            session_cwd_cache_max_entries: 0,
3582            session_transcript_path_cache_ttl_ms: 30_000,
3583            session_transcript_path_cache_max_entries: 5_000,
3584        }
3585    }
3586
3587    #[test]
3588    fn begin_and_finish_requests_keep_trace_id() {
3589        let runtime = tokio::runtime::Runtime::new().expect("runtime");
3590        runtime.block_on(async {
3591            let state = ProxyState::new();
3592            let request_id = state
3593                .begin_request_for_test()
3594                .model("gpt-5")
3595                .service_tier("priority")
3596                .started_at_ms(100)
3597                .begin()
3598                .await;
3599
3600            let active = state.list_active_requests().await;
3601            assert_eq!(active[0].trace_id.as_deref(), Some("codex-1"));
3602
3603            state
3604                .finish_request(FinishRequestParams {
3605                    id: request_id,
3606                    status_code: 200,
3607                    duration_ms: 10,
3608                    ended_at_ms: 110,
3609                    observed_service_tier: Some("priority".to_string()),
3610                    usage: None,
3611                    retry: None,
3612                    ttfb_ms: Some(4),
3613                    streaming: false,
3614                })
3615                .await;
3616
3617            let recent = state.list_recent_finished(1).await;
3618            assert_eq!(recent[0].trace_id.as_deref(), Some("codex-1"));
3619            assert_eq!(recent[0].observability.trace_id.as_deref(), Some("codex-1"));
3620            assert!(recent[0].observability.fast_mode);
3621            assert_eq!(recent[0].observability.generation_ms, Some(6));
3622        });
3623    }
3624
3625    #[test]
3626    fn recent_finished_max_defaults_to_one_thousand() {
3627        assert_eq!(recent_finished_max_from_env(None), 1_000);
3628    }
3629
3630    #[test]
3631    fn finish_request_reports_exactly_once_publication() {
3632        let runtime = tokio::runtime::Runtime::new().expect("runtime");
3633        runtime.block_on(async {
3634            let state = ProxyState::new();
3635            let request_id = state
3636                .begin_request_for_test()
3637                .started_at_ms(100)
3638                .begin()
3639                .await;
3640
3641            let first = state
3642                .finish_request(FinishRequestParams {
3643                    id: request_id,
3644                    status_code: 200,
3645                    duration_ms: 10,
3646                    ended_at_ms: 110,
3647                    observed_service_tier: None,
3648                    usage: None,
3649                    retry: None,
3650                    ttfb_ms: Some(4),
3651                    streaming: false,
3652                })
3653                .await;
3654            let second = state
3655                .finish_request(FinishRequestParams {
3656                    id: request_id,
3657                    status_code: 500,
3658                    duration_ms: 20,
3659                    ended_at_ms: 120,
3660                    observed_service_tier: None,
3661                    usage: None,
3662                    retry: None,
3663                    ttfb_ms: None,
3664                    streaming: false,
3665                })
3666                .await;
3667
3668            assert!(first);
3669            assert!(!second);
3670
3671            let recent = state.list_recent_finished(10).await;
3672            assert_eq!(recent.len(), 1);
3673            assert_eq!(recent[0].id, request_id);
3674            assert_eq!(recent[0].status_code, 200);
3675        });
3676    }
3677
3678    #[test]
3679    fn session_cards_expose_last_and_average_output_token_speed() {
3680        let runtime = tokio::runtime::Runtime::new().expect("runtime");
3681        runtime.block_on(async {
3682            let state = ProxyState::new();
3683
3684            let first_id = state
3685                .begin_request_for_test()
3686                .session_id("sid-speed")
3687                .started_at_ms(100)
3688                .begin()
3689                .await;
3690            state
3691                .finish_request(FinishRequestParams {
3692                    id: first_id,
3693                    status_code: 200,
3694                    duration_ms: 1_500,
3695                    ended_at_ms: 1_600,
3696                    observed_service_tier: None,
3697                    usage: Some(UsageMetrics {
3698                        output_tokens: 200,
3699                        total_tokens: 200,
3700                        ..UsageMetrics::default()
3701                    }),
3702                    retry: None,
3703                    ttfb_ms: Some(500),
3704                    streaming: true,
3705                })
3706                .await;
3707
3708            let second_id = state
3709                .begin_request_for_test()
3710                .session_id("sid-speed")
3711                .started_at_ms(2_000)
3712                .begin()
3713                .await;
3714            state
3715                .finish_request(FinishRequestParams {
3716                    id: second_id,
3717                    status_code: 200,
3718                    duration_ms: 2_500,
3719                    ended_at_ms: 4_500,
3720                    observed_service_tier: None,
3721                    usage: Some(UsageMetrics {
3722                        output_tokens: 300,
3723                        total_tokens: 300,
3724                        ..UsageMetrics::default()
3725                    }),
3726                    retry: None,
3727                    ttfb_ms: Some(500),
3728                    streaming: true,
3729                })
3730                .await;
3731
3732            let stats = state.list_session_stats().await;
3733            let stats = stats.get("sid-speed").expect("session stats");
3734            assert_eq!(stats.last_output_tokens_per_second, Some(150.0));
3735            assert_eq!(stats.avg_output_tokens_per_second, Some(500.0 / 3.0));
3736
3737            let cards = state.list_session_identity_cards(16).await;
3738            let card = cards
3739                .iter()
3740                .find(|card| card.session_id.as_deref() == Some("sid-speed"))
3741                .expect("session card");
3742            assert_eq!(card.last_output_tokens_per_second, Some(150.0));
3743            assert_eq!(card.avg_output_tokens_per_second, Some(500.0 / 3.0));
3744        });
3745    }
3746
3747    #[test]
3748    fn state_change_subscription_tracks_dashboard_mutations() {
3749        let runtime = tokio::runtime::Runtime::new().expect("runtime");
3750        runtime.block_on(async {
3751            let state = ProxyState::new();
3752            let mut changes = state.subscribe_state_changes();
3753
3754            state
3755                .record_provider_balance_snapshot(
3756                    "codex",
3757                    ProviderBalanceSnapshot {
3758                        provider_id: "input6".to_string(),
3759                        station_name: Some("routing".to_string()),
3760                        upstream_index: Some(0),
3761                        fetched_at_ms: 100,
3762                        stale_after_ms: Some(1_000),
3763                        status: BalanceSnapshotStatus::Ok,
3764                        total_balance_usd: Some("12.50".to_string()),
3765                        ..ProviderBalanceSnapshot::default()
3766                    },
3767                )
3768                .await;
3769            changes.changed().await.expect("balance change");
3770            let balance_version = *changes.borrow();
3771            assert!(balance_version > 0);
3772
3773            state
3774                .set_provider_endpoint_usage_exhausted(
3775                    "codex",
3776                    ProviderEndpointKey::new("codex", "input6", "default"),
3777                    true,
3778                )
3779                .await;
3780            changes.changed().await.expect("usage exhausted change");
3781            let usage_version = *changes.borrow();
3782            assert!(usage_version > balance_version);
3783
3784            let request_id = state
3785                .begin_request_for_test()
3786                .session_id("sid-1")
3787                .model("gpt-5")
3788                .started_at_ms(200)
3789                .begin()
3790                .await;
3791            changes.changed().await.expect("begin request change");
3792            let active_version = *changes.borrow();
3793            assert!(active_version > usage_version);
3794
3795            state
3796                .finish_request(FinishRequestParams {
3797                    id: request_id,
3798                    status_code: 200,
3799                    duration_ms: 25,
3800                    ended_at_ms: 250,
3801                    observed_service_tier: None,
3802                    usage: None,
3803                    retry: None,
3804                    ttfb_ms: Some(5),
3805                    streaming: false,
3806                })
3807                .await;
3808            changes.changed().await.expect("finish request change");
3809            assert!(*changes.borrow() > active_version);
3810        });
3811    }
3812
3813    #[test]
3814    fn state_change_subscription_ignores_session_touch_only_updates() {
3815        let runtime = tokio::runtime::Runtime::new().expect("runtime");
3816        runtime.block_on(async {
3817            let state = ProxyState::new();
3818            let mut changes = state.subscribe_state_changes();
3819
3820            state
3821                .set_session_model_override("sid-1".to_string(), "gpt-5".to_string(), 100)
3822                .await;
3823            changes.changed().await.expect("model override change");
3824            let version = *changes.borrow();
3825
3826            state.touch_session_model_override("sid-1", 200).await;
3827            assert_eq!(*changes.borrow(), version);
3828        });
3829    }
3830
3831    #[test]
3832    fn finish_request_estimates_cost_and_rolls_up_cost() {
3833        let runtime = tokio::runtime::Runtime::new().expect("runtime");
3834        runtime.block_on(async {
3835            let state = ProxyState::new();
3836            let request_id = state
3837                .begin_request_for_test()
3838                .model("gpt-5")
3839                .started_at_ms(100)
3840                .begin()
3841                .await;
3842
3843            state
3844                .finish_request(FinishRequestParams {
3845                    id: request_id,
3846                    status_code: 200,
3847                    duration_ms: 10,
3848                    ended_at_ms: 110,
3849                    observed_service_tier: None,
3850                    usage: Some(UsageMetrics {
3851                        input_tokens: 1_000,
3852                        output_tokens: 500,
3853                        cached_input_tokens: 100,
3854                        total_tokens: 1_500,
3855                        ..UsageMetrics::default()
3856                    }),
3857                    retry: None,
3858                    ttfb_ms: Some(4),
3859                    streaming: false,
3860                })
3861                .await;
3862
3863            let recent = state.list_recent_finished(1).await;
3864            assert_eq!(recent[0].cost.total_cost_usd.as_deref(), Some("0.0061375"));
3865
3866            let rollup = state.get_usage_rollup_view("codex", 12, 1).await;
3867            assert_eq!(
3868                rollup.loaded.cost.total_cost_usd.as_deref(),
3869                Some("0.0061375")
3870            );
3871            assert_eq!(rollup.loaded.cost.priced_requests, 1);
3872            assert_eq!(rollup.loaded.cost.unpriced_requests, 0);
3873        });
3874    }
3875
3876    #[test]
3877    fn usage_rollup_view_scores_entities_inside_selected_window() {
3878        let runtime = tokio::runtime::Runtime::new().expect("runtime");
3879        runtime.block_on(async {
3880            let state = ProxyState::new();
3881            let now_ms = std::time::SystemTime::now()
3882                .duration_since(std::time::UNIX_EPOCH)
3883                .map(|d| d.as_millis() as u64)
3884                .unwrap_or(0);
3885            let old_ms = now_ms.saturating_sub(10 * 86_400_000);
3886
3887            let old_id = state
3888                .begin_request_for_test()
3889                .model("gpt-5")
3890                .started_at_ms(old_ms.saturating_sub(1_000))
3891                .begin()
3892                .await;
3893            state
3894                .update_request_route(
3895                    old_id,
3896                    Some("old-station".to_string()),
3897                    Some("old-provider".to_string()),
3898                    "https://old.example/v1".to_string(),
3899                    None,
3900                )
3901                .await;
3902            state
3903                .finish_request(FinishRequestParams {
3904                    id: old_id,
3905                    status_code: 200,
3906                    duration_ms: 20,
3907                    ended_at_ms: old_ms,
3908                    observed_service_tier: None,
3909                    usage: Some(UsageMetrics {
3910                        total_tokens: 100_000,
3911                        ..UsageMetrics::default()
3912                    }),
3913                    retry: None,
3914                    ttfb_ms: Some(5),
3915                    streaming: false,
3916                })
3917                .await;
3918
3919            let fresh_id = state
3920                .begin_request_for_test()
3921                .model("gpt-5")
3922                .started_at_ms(now_ms.saturating_sub(1_000))
3923                .begin()
3924                .await;
3925            state
3926                .update_request_route(
3927                    fresh_id,
3928                    Some("fresh-station".to_string()),
3929                    Some("fresh-provider".to_string()),
3930                    "https://fresh.example/v1".to_string(),
3931                    None,
3932                )
3933                .await;
3934            state
3935                .finish_request(FinishRequestParams {
3936                    id: fresh_id,
3937                    status_code: 200,
3938                    duration_ms: 10,
3939                    ended_at_ms: now_ms,
3940                    observed_service_tier: None,
3941                    usage: Some(UsageMetrics {
3942                        total_tokens: 10,
3943                        ..UsageMetrics::default()
3944                    }),
3945                    retry: None,
3946                    ttfb_ms: Some(3),
3947                    streaming: false,
3948                })
3949                .await;
3950
3951            let week = state.get_usage_rollup_view("codex", 10, 7).await;
3952            assert_eq!(week.loaded.requests_total, 2);
3953            assert_eq!(week.window.requests_total, 1);
3954            assert_eq!(week.by_day.len(), 7);
3955            assert_eq!(week.by_config[0].0, "fresh-station");
3956            assert_eq!(week.by_provider[0].0, "fresh-provider");
3957
3958            let loaded = state.get_usage_rollup_view("codex", 10, 0).await;
3959            assert_eq!(loaded.window.requests_total, 2);
3960            assert_eq!(loaded.by_config[0].0, "old-station");
3961            assert_eq!(loaded.by_provider[0].0, "old-provider");
3962        });
3963    }
3964
3965    #[test]
3966    fn build_session_identity_cards_merges_sources_and_sorts_newest_first() {
3967        let active = vec![ActiveRequest {
3968            id: 1,
3969            trace_id: Some("codex-1".to_string()),
3970            session_id: Some("sid-active".to_string()),
3971            session_identity_source: Some(SessionIdentitySource::Header),
3972            client_name: Some("Frank-Laptop".to_string()),
3973            client_addr: Some("100.64.0.8".to_string()),
3974            cwd: Some("G:/codes/project".to_string()),
3975            model: Some("gpt-5.4".to_string()),
3976            reasoning_effort: Some("medium".to_string()),
3977            service_tier: Some("priority".to_string()),
3978            station_name: Some("right".to_string()),
3979            provider_id: Some("right".to_string()),
3980            upstream_base_url: Some("https://right.example/v1".to_string()),
3981            route_decision: None,
3982            service: "codex".to_string(),
3983            method: "POST".to_string(),
3984            path: "/v1/responses".to_string(),
3985            started_at_ms: 500,
3986        }];
3987        let recent = vec![
3988            FinishedRequest {
3989                id: 2,
3990                trace_id: Some("codex-2".to_string()),
3991                session_id: Some("sid-recent".to_string()),
3992                session_identity_source: Some(SessionIdentitySource::PromptCacheKey),
3993                client_name: Some("Studio-Mini".to_string()),
3994                client_addr: Some("100.64.0.9".to_string()),
3995                cwd: Some("G:/codes/other".to_string()),
3996                model: Some("gpt-5.3".to_string()),
3997                reasoning_effort: Some("high".to_string()),
3998                service_tier: Some("default".to_string()),
3999                station_name: Some("vibe".to_string()),
4000                provider_id: Some("vibe".to_string()),
4001                upstream_base_url: Some("https://vibe.example/v1".to_string()),
4002                route_decision: None,
4003                usage: Some(UsageMetrics {
4004                    input_tokens: 1,
4005                    output_tokens: 2,
4006                    reasoning_tokens: 3,
4007                    total_tokens: 6,
4008                    ..UsageMetrics::default()
4009                }),
4010                cost: CostBreakdown::default(),
4011                retry: None,
4012                provider_signals: Vec::new(),
4013                policy_actions: Vec::new(),
4014                observability: RequestObservability::default(),
4015                service: "codex".to_string(),
4016                method: "POST".to_string(),
4017                path: "/v1/responses".to_string(),
4018                status_code: 200,
4019                duration_ms: 1200,
4020                ttfb_ms: Some(100),
4021                streaming: false,
4022                ended_at_ms: 2_000,
4023            },
4024            FinishedRequest {
4025                id: 3,
4026                trace_id: Some("codex-3".to_string()),
4027                session_id: Some("sid-active".to_string()),
4028                session_identity_source: Some(SessionIdentitySource::Header),
4029                client_name: Some("Frank-Laptop".to_string()),
4030                client_addr: Some("100.64.0.8".to_string()),
4031                cwd: Some("G:/codes/project".to_string()),
4032                model: Some("gpt-5.4".to_string()),
4033                reasoning_effort: Some("low".to_string()),
4034                service_tier: Some("flex".to_string()),
4035                station_name: Some("right".to_string()),
4036                provider_id: Some("right".to_string()),
4037                upstream_base_url: Some("https://right.example/v1".to_string()),
4038                route_decision: None,
4039                usage: None,
4040                cost: CostBreakdown::default(),
4041                retry: None,
4042                provider_signals: Vec::new(),
4043                policy_actions: Vec::new(),
4044                observability: RequestObservability::default(),
4045                service: "codex".to_string(),
4046                method: "POST".to_string(),
4047                path: "/v1/responses".to_string(),
4048                status_code: 429,
4049                duration_ms: 900,
4050                ttfb_ms: None,
4051                streaming: false,
4052                ended_at_ms: 1_000,
4053            },
4054        ];
4055        let overrides = HashMap::from([("sid-active".to_string(), "xhigh".to_string())]);
4056        let config_overrides = HashMap::from([("sid-active".to_string(), "temp".to_string())]);
4057        let model_overrides =
4058            HashMap::from([("sid-active".to_string(), "gpt-5.4-mini".to_string())]);
4059        let service_tier_overrides =
4060            HashMap::from([("sid-active".to_string(), "priority".to_string())]);
4061        let stats = HashMap::from([(
4062            "sid-active".to_string(),
4063            SessionStats {
4064                turns_total: 3,
4065                last_session_identity_source: Some(SessionIdentitySource::Header),
4066                last_client_name: Some("Frank-Laptop".to_string()),
4067                last_client_addr: Some("100.64.0.8".to_string()),
4068                last_model: Some("gpt-5.4".to_string()),
4069                last_reasoning_effort: Some("low".to_string()),
4070                last_service_tier: Some("flex".to_string()),
4071                last_provider_id: Some("right".to_string()),
4072                last_station_name: Some("right".to_string()),
4073                last_route_decision: None,
4074                last_usage: None,
4075                total_usage: UsageMetrics {
4076                    input_tokens: 10,
4077                    output_tokens: 20,
4078                    reasoning_tokens: 5,
4079                    total_tokens: 35,
4080                    ..UsageMetrics::default()
4081                },
4082                turns_with_usage: 2,
4083                last_output_tokens_per_second: None,
4084                avg_output_tokens_per_second: None,
4085                output_generation_ms_total: 0,
4086                last_status: Some(429),
4087                last_duration_ms: Some(900),
4088                last_ended_at_ms: Some(1_000),
4089                last_seen_ms: 1_000,
4090            },
4091        )]);
4092
4093        let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
4094            active: &active,
4095            recent: &recent,
4096            overrides: &overrides,
4097            station_overrides: &config_overrides,
4098            model_overrides: &model_overrides,
4099            service_tier_overrides: &service_tier_overrides,
4100            bindings: &HashMap::new(),
4101            route_affinities: &HashMap::new(),
4102            global_station_override: None,
4103            stats: &stats,
4104        });
4105
4106        assert_eq!(cards.len(), 2);
4107        assert_eq!(cards[0].session_id.as_deref(), Some("sid-recent"));
4108        assert_eq!(
4109            cards[0].session_identity_source,
4110            Some(SessionIdentitySource::PromptCacheKey)
4111        );
4112        assert_eq!(
4113            cards[0].observation_scope,
4114            SessionObservationScope::HostLocalEnriched
4115        );
4116        assert_eq!(cards[0].last_client_name.as_deref(), Some("Studio-Mini"));
4117        assert_eq!(cards[0].last_client_addr.as_deref(), Some("100.64.0.9"));
4118        assert_eq!(cards[1].session_id.as_deref(), Some("sid-active"));
4119        assert_eq!(
4120            cards[1].session_identity_source,
4121            Some(SessionIdentitySource::Header)
4122        );
4123        assert_eq!(
4124            cards[1].observation_scope,
4125            SessionObservationScope::HostLocalEnriched
4126        );
4127        assert_eq!(cards[1].active_count, 1);
4128        assert_eq!(cards[1].last_client_name.as_deref(), Some("Frank-Laptop"));
4129        assert_eq!(cards[1].last_client_addr.as_deref(), Some("100.64.0.8"));
4130        assert_eq!(cards[1].last_status, Some(429));
4131        assert_eq!(cards[1].override_effort.as_deref(), Some("xhigh"));
4132        assert_eq!(cards[1].override_station_name.as_deref(), Some("temp"));
4133        assert_eq!(cards[1].override_model.as_deref(), Some("gpt-5.4-mini"));
4134        assert_eq!(cards[1].override_service_tier.as_deref(), Some("priority"));
4135        assert_eq!(
4136            cards[1]
4137                .effective_model
4138                .as_ref()
4139                .map(|value| value.value.as_str()),
4140            Some("gpt-5.4-mini")
4141        );
4142        assert_eq!(
4143            cards[1].effective_model.as_ref().map(|value| value.source),
4144            Some(RouteValueSource::SessionOverride)
4145        );
4146        assert_eq!(
4147            cards[1]
4148                .effective_reasoning_effort
4149                .as_ref()
4150                .map(|value| value.source),
4151            Some(RouteValueSource::SessionOverride)
4152        );
4153        assert_eq!(
4154            cards[1]
4155                .effective_service_tier
4156                .as_ref()
4157                .map(|value| value.source),
4158            Some(RouteValueSource::SessionOverride)
4159        );
4160        assert_eq!(
4161            cards[1]
4162                .effective_station
4163                .as_ref()
4164                .map(|value| value.source),
4165            Some(RouteValueSource::SessionOverride)
4166        );
4167        assert!(cards[1].effective_upstream_base_url.is_none());
4168        assert_eq!(
4169            cards[1].last_upstream_base_url.as_deref(),
4170            Some("https://right.example/v1")
4171        );
4172        assert_eq!(cards[1].turns_total, Some(3));
4173        assert_eq!(cards[1].last_service_tier.as_deref(), Some("flex"));
4174        assert_eq!(
4175            cards[1].total_usage.as_ref().map(|u| u.total_tokens),
4176            Some(35)
4177        );
4178    }
4179
4180    #[test]
4181    fn build_session_identity_cards_default_profile_keeps_request_fields() {
4182        let active = vec![ActiveRequest {
4183            id: 1,
4184            trace_id: Some("codex-1".to_string()),
4185            session_id: Some("sid-bound".to_string()),
4186            session_identity_source: Some(SessionIdentitySource::Header),
4187            client_name: Some("Workstation".to_string()),
4188            client_addr: Some("100.64.0.10".to_string()),
4189            cwd: None,
4190            model: Some("gpt-observed".to_string()),
4191            reasoning_effort: Some("medium".to_string()),
4192            service_tier: Some("default".to_string()),
4193            station_name: Some("right".to_string()),
4194            provider_id: None,
4195            upstream_base_url: None,
4196            route_decision: None,
4197            service: "codex".to_string(),
4198            method: "POST".to_string(),
4199            path: "/v1/responses".to_string(),
4200            started_at_ms: 10,
4201        }];
4202        let bindings = HashMap::from([(
4203            "sid-bound".to_string(),
4204            SessionBinding {
4205                session_id: "sid-bound".to_string(),
4206                profile_name: Some("daily".to_string()),
4207                station_name: Some("vibe".to_string()),
4208                model: Some("gpt-bound".to_string()),
4209                reasoning_effort: Some("high".to_string()),
4210                service_tier: Some("priority".to_string()),
4211                continuity_mode: SessionContinuityMode::DefaultProfile,
4212                created_at_ms: 1,
4213                updated_at_ms: 1,
4214                last_seen_ms: 10,
4215            },
4216        )]);
4217
4218        let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
4219            active: &active,
4220            recent: &[],
4221            overrides: &HashMap::new(),
4222            station_overrides: &HashMap::new(),
4223            model_overrides: &HashMap::new(),
4224            service_tier_overrides: &HashMap::new(),
4225            bindings: &bindings,
4226            route_affinities: &HashMap::new(),
4227            global_station_override: None,
4228            stats: &HashMap::new(),
4229        });
4230
4231        assert_eq!(cards[0].binding_profile_name.as_deref(), Some("daily"));
4232        assert_eq!(
4233            cards[0].observation_scope,
4234            SessionObservationScope::ObservedOnly
4235        );
4236        assert_eq!(
4237            cards[0].binding_continuity_mode,
4238            Some(SessionContinuityMode::DefaultProfile)
4239        );
4240        assert_eq!(
4241            cards[0]
4242                .effective_model
4243                .as_ref()
4244                .map(|value| (value.value.as_str(), value.source)),
4245            Some(("gpt-observed", RouteValueSource::RequestPayload))
4246        );
4247        assert_eq!(
4248            cards[0]
4249                .effective_reasoning_effort
4250                .as_ref()
4251                .map(|value| (value.value.as_str(), value.source)),
4252            Some(("medium", RouteValueSource::RequestPayload))
4253        );
4254        assert_eq!(
4255            cards[0]
4256                .effective_service_tier
4257                .as_ref()
4258                .map(|value| (value.value.as_str(), value.source)),
4259            Some(("default", RouteValueSource::RequestPayload))
4260        );
4261        assert_eq!(
4262            cards[0]
4263                .effective_station
4264                .as_ref()
4265                .map(|value| (value.value.as_str(), value.source)),
4266            Some(("vibe", RouteValueSource::ProfileDefault))
4267        );
4268    }
4269
4270    #[test]
4271    fn build_session_identity_cards_keeps_request_fields_but_allows_global_config_override() {
4272        let active = vec![ActiveRequest {
4273            id: 1,
4274            trace_id: Some("codex-1".to_string()),
4275            session_id: Some("sid-bound".to_string()),
4276            session_identity_source: Some(SessionIdentitySource::Header),
4277            client_name: Some("Workstation".to_string()),
4278            client_addr: Some("100.64.0.10".to_string()),
4279            cwd: None,
4280            model: Some("gpt-observed".to_string()),
4281            reasoning_effort: Some("medium".to_string()),
4282            service_tier: Some("default".to_string()),
4283            station_name: Some("vibe".to_string()),
4284            provider_id: None,
4285            upstream_base_url: Some("https://vibe.example/v1".to_string()),
4286            route_decision: None,
4287            service: "codex".to_string(),
4288            method: "POST".to_string(),
4289            path: "/v1/responses".to_string(),
4290            started_at_ms: 10,
4291        }];
4292        let bindings = HashMap::from([(
4293            "sid-bound".to_string(),
4294            SessionBinding {
4295                session_id: "sid-bound".to_string(),
4296                profile_name: Some("daily".to_string()),
4297                station_name: Some("vibe".to_string()),
4298                model: Some("gpt-bound".to_string()),
4299                reasoning_effort: Some("high".to_string()),
4300                service_tier: Some("priority".to_string()),
4301                continuity_mode: SessionContinuityMode::DefaultProfile,
4302                created_at_ms: 1,
4303                updated_at_ms: 1,
4304                last_seen_ms: 10,
4305            },
4306        )]);
4307
4308        let cards = build_session_identity_cards_from_parts(SessionIdentityCardBuildInputs {
4309            active: &active,
4310            recent: &[],
4311            overrides: &HashMap::new(),
4312            station_overrides: &HashMap::new(),
4313            model_overrides: &HashMap::new(),
4314            service_tier_overrides: &HashMap::new(),
4315            bindings: &bindings,
4316            route_affinities: &HashMap::new(),
4317            global_station_override: Some("right"),
4318            stats: &HashMap::new(),
4319        });
4320
4321        assert_eq!(cards[0].binding_profile_name.as_deref(), Some("daily"));
4322        assert_eq!(
4323            cards[0]
4324                .effective_model
4325                .as_ref()
4326                .map(|value| (value.value.as_str(), value.source)),
4327            Some(("gpt-observed", RouteValueSource::RequestPayload))
4328        );
4329        assert_eq!(
4330            cards[0]
4331                .effective_reasoning_effort
4332                .as_ref()
4333                .map(|value| (value.value.as_str(), value.source)),
4334            Some(("medium", RouteValueSource::RequestPayload))
4335        );
4336        assert_eq!(
4337            cards[0]
4338                .effective_service_tier
4339                .as_ref()
4340                .map(|value| (value.value.as_str(), value.source)),
4341            Some(("default", RouteValueSource::RequestPayload))
4342        );
4343        assert_eq!(
4344            cards[0]
4345                .effective_station
4346                .as_ref()
4347                .map(|value| (value.value.as_str(), value.source)),
4348            Some(("right", RouteValueSource::GlobalOverride))
4349        );
4350        assert!(cards[0].effective_upstream_base_url.is_none());
4351    }
4352
4353    #[test]
4354    fn enrich_session_identity_cards_with_runtime_applies_station_mapping_and_single_upstream() {
4355        let mut cards = vec![SessionIdentityCard {
4356            session_id: Some("sid-1".to_string()),
4357            last_model: Some("gpt-5.4".to_string()),
4358            last_station_name: Some("right".to_string()),
4359            last_upstream_base_url: Some("https://right.example/v1".to_string()),
4360            effective_model: Some(ResolvedRouteValue::new(
4361                "gpt-5.4",
4362                RouteValueSource::RequestPayload,
4363            )),
4364            effective_station: Some(ResolvedRouteValue::new(
4365                "right",
4366                RouteValueSource::RuntimeFallback,
4367            )),
4368            ..SessionIdentityCard::default()
4369        }];
4370
4371        let mut mgr = ServiceConfigManager {
4372            active: Some("right".to_string()),
4373            ..ServiceConfigManager::default()
4374        };
4375        mgr.configs.insert(
4376            "right".to_string(),
4377            ServiceConfig {
4378                name: "right".to_string(),
4379                alias: None,
4380                enabled: true,
4381                level: 1,
4382                upstreams: vec![UpstreamConfig {
4383                    base_url: "https://right.example/v1".to_string(),
4384                    auth: UpstreamAuth::default(),
4385                    tags: HashMap::new(),
4386                    supported_models: HashMap::new(),
4387                    model_mapping: HashMap::from([(
4388                        "gpt-5.4".to_string(),
4389                        "gpt-5.4-fast".to_string(),
4390                    )]),
4391                }],
4392            },
4393        );
4394
4395        enrich_session_identity_cards_with_runtime(&mut cards, &mgr);
4396
4397        assert_eq!(
4398            cards[0]
4399                .effective_model
4400                .as_ref()
4401                .map(|value| value.value.as_str()),
4402            Some("gpt-5.4-fast")
4403        );
4404        assert_eq!(
4405            cards[0].effective_model.as_ref().map(|value| value.source),
4406            Some(RouteValueSource::StationMapping)
4407        );
4408        assert_eq!(
4409            cards[0]
4410                .effective_upstream_base_url
4411                .as_ref()
4412                .map(|value| value.value.as_str()),
4413            Some("https://right.example/v1")
4414        );
4415    }
4416
4417    #[test]
4418    fn apply_session_profile_binding_sets_binding_and_clears_manual_overrides() {
4419        let runtime = tokio::runtime::Runtime::new().expect("runtime");
4420        runtime.block_on(async {
4421            let state = ProxyState::new();
4422            let now_ms = 42;
4423            let mut mgr = ServiceConfigManager::default();
4424            mgr.configs.insert(
4425                "right".to_string(),
4426                ServiceConfig {
4427                    name: "right".to_string(),
4428                    alias: None,
4429                    enabled: true,
4430                    level: 1,
4431                    upstreams: vec![UpstreamConfig {
4432                        base_url: "https://right.example/v1".to_string(),
4433                        auth: UpstreamAuth::default(),
4434                        tags: HashMap::from([
4435                            ("supports_reasoning_effort".to_string(), "true".to_string()),
4436                            ("supports_service_tier".to_string(), "true".to_string()),
4437                        ]),
4438                        supported_models: HashMap::from([("gpt-5.4".to_string(), true)]),
4439                        model_mapping: HashMap::new(),
4440                    }],
4441                },
4442            );
4443            mgr.profiles.insert(
4444                "fast".to_string(),
4445                crate::config::ServiceControlProfile {
4446                    extends: None,
4447                    station: Some("right".to_string()),
4448                    model: Some("gpt-5.4".to_string()),
4449                    reasoning_effort: Some("low".to_string()),
4450                    service_tier: Some("flex".to_string()),
4451                },
4452            );
4453
4454            state
4455                .set_session_station_override("sid-1".to_string(), "other".to_string(), 1)
4456                .await;
4457            state
4458                .set_session_model_override("sid-1".to_string(), "gpt-x".to_string(), 1)
4459                .await;
4460            state
4461                .set_session_effort_override("sid-1".to_string(), "high".to_string(), 1)
4462                .await;
4463            state
4464                .set_session_service_tier_override("sid-1".to_string(), "priority".to_string(), 1)
4465                .await;
4466
4467            state
4468                .apply_session_profile_binding(
4469                    "codex",
4470                    &mgr,
4471                    "sid-1".to_string(),
4472                    "fast".to_string(),
4473                    now_ms,
4474                )
4475                .await
4476                .expect("apply profile");
4477
4478            let binding = state
4479                .get_session_binding("sid-1")
4480                .await
4481                .expect("binding exists");
4482            assert_eq!(binding.profile_name.as_deref(), Some("fast"));
4483            assert_eq!(binding.station_name.as_deref(), Some("right"));
4484            assert_eq!(binding.model.as_deref(), Some("gpt-5.4"));
4485            assert_eq!(binding.reasoning_effort.as_deref(), Some("low"));
4486            assert_eq!(binding.service_tier.as_deref(), Some("flex"));
4487            assert_eq!(
4488                binding.continuity_mode,
4489                SessionContinuityMode::ManualProfile
4490            );
4491            assert_eq!(binding.updated_at_ms, now_ms);
4492            assert!(state.get_session_station_override("sid-1").await.is_none());
4493            assert!(state.get_session_model_override("sid-1").await.is_none());
4494            assert!(state.get_session_effort_override("sid-1").await.is_none());
4495            assert!(
4496                state
4497                    .get_session_service_tier_override("sid-1")
4498                    .await
4499                    .is_none()
4500            );
4501        });
4502    }
4503
4504    #[test]
4505    fn list_session_manual_overrides_merges_all_dimensions() {
4506        let runtime = tokio::runtime::Runtime::new().expect("runtime");
4507        runtime.block_on(async {
4508            let state = ProxyState::new();
4509            state
4510                .set_session_reasoning_effort_override("sid-1".to_string(), "high".to_string(), 1)
4511                .await;
4512            state
4513                .set_session_station_override("sid-1".to_string(), "right".to_string(), 1)
4514                .await;
4515            state
4516                .set_session_model_override("sid-1".to_string(), "gpt-5.4".to_string(), 1)
4517                .await;
4518            state
4519                .set_session_service_tier_override("sid-1".to_string(), "priority".to_string(), 1)
4520                .await;
4521            state
4522                .set_session_model_override("sid-2".to_string(), "gpt-5.4-mini".to_string(), 2)
4523                .await;
4524
4525            let merged = state.list_session_manual_overrides().await;
4526            assert_eq!(merged.len(), 2);
4527            assert_eq!(
4528                merged
4529                    .get("sid-1")
4530                    .and_then(|overrides| overrides.reasoning_effort.as_deref()),
4531                Some("high")
4532            );
4533            assert_eq!(
4534                merged
4535                    .get("sid-1")
4536                    .and_then(|overrides| overrides.station_name.as_deref()),
4537                Some("right")
4538            );
4539            assert_eq!(
4540                merged
4541                    .get("sid-1")
4542                    .and_then(|overrides| overrides.model.as_deref()),
4543                Some("gpt-5.4")
4544            );
4545            assert_eq!(
4546                merged
4547                    .get("sid-1")
4548                    .and_then(|overrides| overrides.service_tier.as_deref()),
4549                Some("priority")
4550            );
4551            assert_eq!(
4552                merged
4553                    .get("sid-2")
4554                    .and_then(|overrides| overrides.model.as_deref()),
4555                Some("gpt-5.4-mini")
4556            );
4557            assert!(
4558                merged
4559                    .get("sid-2")
4560                    .is_some_and(|overrides| overrides.reasoning_effort.is_none())
4561            );
4562        });
4563    }
4564
4565    #[test]
4566    fn provider_balance_snapshots_are_recorded_and_refreshed() {
4567        let runtime = tokio::runtime::Runtime::new().expect("runtime");
4568        runtime.block_on(async {
4569            let state = ProxyState::new();
4570            state
4571                .record_provider_balance_snapshot(
4572                    "codex",
4573                    ProviderBalanceSnapshot {
4574                        provider_id: "packycode".to_string(),
4575                        station_name: Some("right".to_string()),
4576                        upstream_index: Some(2),
4577                        source: "usage_provider:budget_http_json".to_string(),
4578                        fetched_at_ms: 10,
4579                        stale_after_ms: Some(0),
4580                        stale: false,
4581                        status: BalanceSnapshotStatus::Unknown,
4582                        exhausted: Some(false),
4583                        exhaustion_affects_routing: true,
4584                        plan_name: None,
4585                        total_balance_usd: Some("3.5".to_string()),
4586                        subscription_balance_usd: None,
4587                        paygo_balance_usd: None,
4588                        monthly_budget_usd: Some("5".to_string()),
4589                        monthly_spent_usd: Some("1.5".to_string()),
4590                        quota_period: None,
4591                        quota_remaining_usd: None,
4592                        quota_limit_usd: None,
4593                        quota_used_usd: None,
4594                        unlimited_quota: None,
4595                        total_used_usd: None,
4596                        today_used_usd: None,
4597                        total_requests: None,
4598                        today_requests: None,
4599                        total_tokens: None,
4600                        today_tokens: None,
4601                        error: None,
4602                        ..ProviderBalanceSnapshot::default()
4603                    },
4604                )
4605                .await;
4606
4607            let view = state.get_provider_balance_view("codex").await;
4608            let balances = view.get("right").expect("station balance");
4609            assert_eq!(balances.len(), 1);
4610            assert_eq!(balances[0].provider_id, "packycode");
4611            assert_eq!(balances[0].status, BalanceSnapshotStatus::Stale);
4612            assert_eq!(balances[0].exhausted, Some(false));
4613
4614            let summary = state
4615                .get_provider_balance_summary_view("codex")
4616                .await
4617                .get("right")
4618                .cloned()
4619                .expect("station balance summary");
4620            assert_eq!(summary.snapshots, 1);
4621            assert_eq!(summary.stale, 1);
4622            assert_eq!(summary.routing_snapshots, 1);
4623        });
4624    }
4625
4626    #[test]
4627    fn provider_balance_history_keeps_recent_refreshes() {
4628        let runtime = tokio::runtime::Runtime::new().expect("runtime");
4629        runtime.block_on(async {
4630            let state = ProxyState::new();
4631            for (fetched_at_ms, remaining) in [(10, "10"), (20, "8")] {
4632                state
4633                    .record_provider_balance_snapshot(
4634                        "codex",
4635                        ProviderBalanceSnapshot {
4636                            provider_id: "packycode".to_string(),
4637                            station_name: Some("right".to_string()),
4638                            upstream_index: Some(2),
4639                            source: "usage_provider:test".to_string(),
4640                            fetched_at_ms,
4641                            quota_period: Some("daily".to_string()),
4642                            quota_remaining_usd: Some(remaining.to_string()),
4643                            quota_limit_usd: Some("20".to_string()),
4644                            exhausted: Some(false),
4645                            ..ProviderBalanceSnapshot::default()
4646                        },
4647                    )
4648                    .await;
4649            }
4650
4651            let history = state.get_provider_balance_history_view("codex").await;
4652            let balances = history.get("right").expect("station history");
4653
4654            assert_eq!(balances.len(), 2);
4655            assert_eq!(balances[0].quota_remaining_usd.as_deref(), Some("10"));
4656            assert_eq!(balances[1].quota_remaining_usd.as_deref(), Some("8"));
4657        });
4658    }
4659
4660    #[test]
4661    fn provider_balance_snapshots_keep_multiple_providers_per_upstream() {
4662        let runtime = tokio::runtime::Runtime::new().expect("runtime");
4663        runtime.block_on(async {
4664            let state = ProxyState::new();
4665            for (provider_id, status) in [
4666                ("general", BalanceSnapshotStatus::Error),
4667                ("newapi", BalanceSnapshotStatus::Ok),
4668            ] {
4669                state
4670                    .record_provider_balance_snapshot(
4671                        "codex",
4672                        ProviderBalanceSnapshot {
4673                            provider_id: provider_id.to_string(),
4674                            station_name: Some("routing".to_string()),
4675                            upstream_index: Some(1),
4676                            source: "usage_provider:test".to_string(),
4677                            fetched_at_ms: 10,
4678                            stale_after_ms: None,
4679                            stale: false,
4680                            status,
4681                            exhausted: if status == BalanceSnapshotStatus::Ok {
4682                                Some(false)
4683                            } else {
4684                                None
4685                            },
4686                            exhaustion_affects_routing: true,
4687                            plan_name: None,
4688                            total_balance_usd: if status == BalanceSnapshotStatus::Ok {
4689                                Some("3.5".to_string())
4690                            } else {
4691                                None
4692                            },
4693                            subscription_balance_usd: None,
4694                            paygo_balance_usd: None,
4695                            monthly_budget_usd: None,
4696                            monthly_spent_usd: None,
4697                            quota_period: None,
4698                            quota_remaining_usd: None,
4699                            quota_limit_usd: None,
4700                            quota_used_usd: None,
4701                            unlimited_quota: None,
4702                            total_used_usd: None,
4703                            today_used_usd: None,
4704                            total_requests: None,
4705                            today_requests: None,
4706                            total_tokens: None,
4707                            today_tokens: None,
4708                            error: if status == BalanceSnapshotStatus::Error {
4709                                Some("decode failed".to_string())
4710                            } else {
4711                                None
4712                            },
4713                            ..ProviderBalanceSnapshot::default()
4714                        },
4715                    )
4716                    .await;
4717            }
4718
4719            let view = state.get_provider_balance_view("codex").await;
4720            let balances = view.get("routing").expect("station balance");
4721            assert_eq!(balances.len(), 2);
4722            assert_eq!(
4723                balances
4724                    .iter()
4725                    .map(|snapshot| snapshot.provider_id.as_str())
4726                    .collect::<Vec<_>>(),
4727                vec!["general", "newapi"]
4728            );
4729        });
4730    }
4731
4732    #[test]
4733    fn provider_balance_error_refresh_preserves_previous_amounts() {
4734        let runtime = tokio::runtime::Runtime::new().expect("runtime");
4735        runtime.block_on(async {
4736            let state = ProxyState::new();
4737            state
4738                .record_provider_balance_snapshot(
4739                    "codex",
4740                    ProviderBalanceSnapshot {
4741                        provider_id: "input".to_string(),
4742                        station_name: Some("routing".to_string()),
4743                        upstream_index: Some(0),
4744                        source: "usage_provider:sub2api_usage".to_string(),
4745                        fetched_at_ms: unix_now_ms(),
4746                        stale_after_ms: None,
4747                        exhausted: Some(false),
4748                        quota_period: Some("daily".to_string()),
4749                        quota_remaining_usd: Some("263.68".to_string()),
4750                        quota_limit_usd: Some("300.00".to_string()),
4751                        ..ProviderBalanceSnapshot::default()
4752                    },
4753                )
4754                .await;
4755
4756            state
4757                .record_provider_balance_snapshot(
4758                    "codex",
4759                    ProviderBalanceSnapshot {
4760                        provider_id: "input".to_string(),
4761                        station_name: Some("routing".to_string()),
4762                        upstream_index: Some(0),
4763                        source: "usage_provider:openai_balance_http_json".to_string(),
4764                        fetched_at_ms: unix_now_ms(),
4765                        stale_after_ms: None,
4766                        error: Some("usage provider response read failed".to_string()),
4767                        ..ProviderBalanceSnapshot::default()
4768                    },
4769                )
4770                .await;
4771
4772            let view = state.get_provider_balance_view("codex").await;
4773            let snapshot = view
4774                .get("routing")
4775                .and_then(|snapshots| snapshots.first())
4776                .expect("routing balance snapshot");
4777
4778            assert_eq!(snapshot.status, BalanceSnapshotStatus::Error);
4779            assert_eq!(snapshot.quota_period.as_deref(), Some("daily"));
4780            assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("263.68"));
4781            assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("300.00"));
4782            assert_eq!(
4783                snapshot.error.as_deref(),
4784                Some("usage provider response read failed")
4785            );
4786        });
4787    }
4788
4789    #[test]
4790    fn prune_runtime_observability_keeps_catalog_provider_balances_on_routing_layout_change() {
4791        let runtime = tokio::runtime::Runtime::new().expect("runtime");
4792        runtime.block_on(async {
4793            let state = ProxyState::new();
4794
4795            let mut initial_mgr = ServiceConfigManager::default();
4796            initial_mgr.configs.insert(
4797                "routing".to_string(),
4798                ServiceConfig {
4799                    name: "routing".to_string(),
4800                    alias: None,
4801                    enabled: true,
4802                    level: 1,
4803                    upstreams: vec![
4804                        UpstreamConfig {
4805                            base_url: "https://input.example/v1".to_string(),
4806                            auth: UpstreamAuth::default(),
4807                            tags: HashMap::from([("provider_id".to_string(), "input".to_string())]),
4808                            supported_models: HashMap::new(),
4809                            model_mapping: HashMap::new(),
4810                        },
4811                        UpstreamConfig {
4812                            base_url: "https://backup.example/v1".to_string(),
4813                            auth: UpstreamAuth::default(),
4814                            tags: HashMap::from([(
4815                                "provider_id".to_string(),
4816                                "backup".to_string(),
4817                            )]),
4818                            supported_models: HashMap::new(),
4819                            model_mapping: HashMap::new(),
4820                        },
4821                    ],
4822                },
4823            );
4824            state
4825                .prune_runtime_observability_for_service("codex", &initial_mgr)
4826                .await;
4827
4828            state
4829                .record_provider_balance_snapshot(
4830                    "codex",
4831                    ProviderBalanceSnapshot {
4832                        provider_id: "input".to_string(),
4833                        station_name: Some("input".to_string()),
4834                        upstream_index: Some(0),
4835                        source: "usage_provider:test".to_string(),
4836                        fetched_at_ms: 10,
4837                        stale_after_ms: None,
4838                        stale: false,
4839                        status: BalanceSnapshotStatus::Ok,
4840                        exhausted: Some(false),
4841                        exhaustion_affects_routing: true,
4842                        total_balance_usd: Some("3.5".to_string()),
4843                        ..ProviderBalanceSnapshot::default()
4844                    },
4845                )
4846                .await;
4847            state
4848                .record_provider_balance_snapshot(
4849                    "codex",
4850                    ProviderBalanceSnapshot {
4851                        provider_id: "input".to_string(),
4852                        station_name: Some("routing".to_string()),
4853                        upstream_index: Some(0),
4854                        source: "usage_provider:test".to_string(),
4855                        fetched_at_ms: 10,
4856                        stale_after_ms: None,
4857                        stale: false,
4858                        status: BalanceSnapshotStatus::Ok,
4859                        exhausted: Some(false),
4860                        exhaustion_affects_routing: true,
4861                        total_balance_usd: Some("3.5".to_string()),
4862                        ..ProviderBalanceSnapshot::default()
4863                    },
4864                )
4865                .await;
4866
4867            let mut pinned_mgr = ServiceConfigManager::default();
4868            pinned_mgr.configs.insert(
4869                "routing".to_string(),
4870                ServiceConfig {
4871                    name: "routing".to_string(),
4872                    alias: None,
4873                    enabled: true,
4874                    level: 1,
4875                    upstreams: vec![UpstreamConfig {
4876                        base_url: "https://input.example/v1".to_string(),
4877                        auth: UpstreamAuth::default(),
4878                        tags: HashMap::from([("provider_id".to_string(), "input".to_string())]),
4879                        supported_models: HashMap::new(),
4880                        model_mapping: HashMap::new(),
4881                    }],
4882                },
4883            );
4884            state
4885                .prune_runtime_observability_for_service("codex", &pinned_mgr)
4886                .await;
4887
4888            let view = state.get_provider_balance_view("codex").await;
4889            assert!(view.contains_key("input"));
4890            assert!(!view.contains_key("routing"));
4891        });
4892    }
4893
4894    #[test]
4895    fn prune_runtime_observability_removes_stale_service_keys() {
4896        let runtime = tokio::runtime::Runtime::new().expect("runtime");
4897        runtime.block_on(async {
4898            let state = ProxyState::new();
4899
4900            state
4901                .set_station_enabled_override("codex", "old".to_string(), false, 1)
4902                .await;
4903            state
4904                .set_upstream_enabled_override(
4905                    "codex",
4906                    "https://old.example/v1".to_string(),
4907                    false,
4908                    1,
4909                )
4910                .await;
4911            state
4912                .set_provider_endpoint_enabled_override(
4913                    "codex",
4914                    ProviderEndpointKey::new("codex", "provider-old", "default"),
4915                    false,
4916                    1,
4917                )
4918                .await;
4919            state
4920                .set_provider_endpoint_runtime_state_override(
4921                    "codex",
4922                    ProviderEndpointKey::new("codex", "provider-old", "default"),
4923                    RuntimeConfigState::BreakerOpen,
4924                    1,
4925                )
4926                .await;
4927            state
4928                .record_station_health(
4929                    "codex",
4930                    "old".to_string(),
4931                    StationHealth {
4932                        checked_at_ms: 10,
4933                        upstreams: vec![UpstreamHealth {
4934                            base_url: "https://old.example/v1".to_string(),
4935                            ok: Some(false),
4936                            status_code: Some(500),
4937                            latency_ms: Some(10),
4938                            error: Some("boom".to_string()),
4939                            passive: None,
4940                        }],
4941                    },
4942                )
4943                .await;
4944            state
4945                .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
4946                    service_name: "codex".to_string(),
4947                    station_name: "old".to_string(),
4948                    base_url: "https://old.example/v1".to_string(),
4949                    status_code: Some(500),
4950                    error_class: Some("upstream_transport_error".to_string()),
4951                    error: Some("boom".to_string()),
4952                    now_ms: 20,
4953                })
4954                .await;
4955            state
4956                .record_provider_balance_snapshot(
4957                    "codex",
4958                    ProviderBalanceSnapshot {
4959                        provider_id: "provider-old".to_string(),
4960                        station_name: Some("old".to_string()),
4961                        upstream_index: Some(0),
4962                        source: "usage_provider:budget_http_json".to_string(),
4963                        fetched_at_ms: 10,
4964                        stale_after_ms: None,
4965                        stale: false,
4966                        status: BalanceSnapshotStatus::Ok,
4967                        exhausted: Some(false),
4968                        exhaustion_affects_routing: true,
4969                        plan_name: None,
4970                        total_balance_usd: Some("3.5".to_string()),
4971                        subscription_balance_usd: None,
4972                        paygo_balance_usd: None,
4973                        monthly_budget_usd: None,
4974                        monthly_spent_usd: None,
4975                        quota_period: None,
4976                        quota_remaining_usd: None,
4977                        quota_limit_usd: None,
4978                        quota_used_usd: None,
4979                        unlimited_quota: None,
4980                        total_used_usd: None,
4981                        today_used_usd: None,
4982                        total_requests: None,
4983                        today_requests: None,
4984                        total_tokens: None,
4985                        today_tokens: None,
4986                        error: None,
4987                        ..ProviderBalanceSnapshot::default()
4988                    },
4989                )
4990                .await;
4991
4992            let request_id = state
4993                .begin_request_for_test()
4994                .session_id("sid-old")
4995                .model("gpt-5")
4996                .started_at_ms(30)
4997                .begin()
4998                .await;
4999            state
5000                .update_request_route(
5001                    request_id,
5002                    Some("old".to_string()),
5003                    Some("provider-old".to_string()),
5004                    "https://old.example/v1".to_string(),
5005                    None,
5006                )
5007                .await;
5008            state
5009                .finish_request(FinishRequestParams {
5010                    id: request_id,
5011                    status_code: 200,
5012                    duration_ms: 5,
5013                    ended_at_ms: 35,
5014                    observed_service_tier: None,
5015                    usage: None,
5016                    retry: None,
5017                    ttfb_ms: None,
5018                    streaming: false,
5019                })
5020                .await;
5021
5022            let mut mgr = ServiceConfigManager::default();
5023            mgr.configs.insert(
5024                "new".to_string(),
5025                ServiceConfig {
5026                    name: "new".to_string(),
5027                    alias: None,
5028                    enabled: true,
5029                    level: 1,
5030                    upstreams: vec![UpstreamConfig {
5031                        base_url: "https://new.example/v1".to_string(),
5032                        auth: UpstreamAuth::default(),
5033                        tags: HashMap::from([(
5034                            "provider_id".to_string(),
5035                            "provider-new".to_string(),
5036                        )]),
5037                        supported_models: HashMap::new(),
5038                        model_mapping: HashMap::new(),
5039                    }],
5040                },
5041            );
5042
5043            state
5044                .prune_runtime_observability_for_service("codex", &mgr)
5045                .await;
5046
5047            assert!(state.get_station_meta_overrides("codex").await.is_empty());
5048            assert!(state.get_upstream_meta_overrides("codex").await.is_empty());
5049            assert!(state.get_station_health("codex").await.is_empty());
5050            assert!(state.get_provider_balance_view("codex").await.is_empty());
5051
5052            let rollup = state.get_usage_rollup_view("codex", 10, 10).await;
5053            assert!(rollup.by_config.is_empty());
5054            assert!(rollup.by_provider.is_empty());
5055        });
5056    }
5057
5058    #[test]
5059    fn get_upstream_meta_overrides_merges_endpoint_and_legacy_base_url_entries() {
5060        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5061        runtime.block_on(async {
5062            let state = ProxyState::new();
5063
5064            state
5065                .set_provider_endpoint_enabled_override(
5066                    "codex",
5067                    ProviderEndpointKey::new("codex", "alpha", "default"),
5068                    false,
5069                    1,
5070                )
5071                .await;
5072            state
5073                .set_provider_endpoint_runtime_state_override(
5074                    "codex",
5075                    ProviderEndpointKey::new("codex", "alpha", "default"),
5076                    RuntimeConfigState::BreakerOpen,
5077                    2,
5078                )
5079                .await;
5080            state
5081                .set_upstream_enabled_override(
5082                    "codex",
5083                    "https://legacy.example/v1".to_string(),
5084                    true,
5085                    3,
5086                )
5087                .await;
5088            state
5089                .set_upstream_runtime_state_override(
5090                    "codex",
5091                    "https://legacy.example/v1".to_string(),
5092                    RuntimeConfigState::Draining,
5093                    4,
5094                )
5095                .await;
5096
5097            let overrides = state.get_upstream_meta_overrides("codex").await;
5098
5099            assert_eq!(
5100                overrides.get("codex/alpha/default"),
5101                Some(&(Some(false), Some(RuntimeConfigState::BreakerOpen)))
5102            );
5103            assert_eq!(
5104                overrides.get("https://legacy.example/v1"),
5105                Some(&(Some(true), Some(RuntimeConfigState::Draining)))
5106            );
5107        });
5108    }
5109
5110    #[test]
5111    fn provider_endpoint_runtime_health_is_keyed_by_provider_endpoint() {
5112        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5113        runtime.block_on(async {
5114            let state = ProxyState::new();
5115            let monthly = ProviderEndpointKey::new("codex", "monthly", "default");
5116            let fallback = ProviderEndpointKey::new("codex", "fallback", "default");
5117
5118            state
5119                .record_provider_endpoint_attempt_success("codex", fallback.clone(), 10)
5120                .await;
5121            state
5122                .set_provider_endpoint_usage_exhausted("codex", monthly.clone(), true)
5123                .await;
5124            state
5125                .record_provider_endpoint_attempt_failure(
5126                    "codex",
5127                    monthly.clone(),
5128                    0,
5129                    crate::lb::CooldownBackoff {
5130                        factor: 1,
5131                        max_secs: 0,
5132                    },
5133                )
5134                .await;
5135            state
5136                .record_provider_endpoint_attempt_failure(
5137                    "codex",
5138                    monthly.clone(),
5139                    0,
5140                    crate::lb::CooldownBackoff {
5141                        factor: 1,
5142                        max_secs: 0,
5143                    },
5144                )
5145                .await;
5146            state
5147                .record_provider_endpoint_attempt_failure(
5148                    "codex",
5149                    monthly.clone(),
5150                    30,
5151                    crate::lb::CooldownBackoff {
5152                        factor: 1,
5153                        max_secs: 0,
5154                    },
5155                )
5156                .await;
5157
5158            let runtime = state
5159                .route_plan_runtime_state_for_provider_endpoints("codex")
5160                .await;
5161            let monthly_state = runtime.provider_endpoint(&monthly);
5162            assert_eq!(monthly_state.failure_count, crate::lb::FAILURE_THRESHOLD);
5163            assert!(monthly_state.cooldown_active);
5164            assert!(monthly_state.usage_exhausted);
5165            assert_eq!(runtime.affinity_provider_endpoint(), Some(&fallback));
5166
5167            state
5168                .set_provider_endpoint_runtime_state_override(
5169                    "codex",
5170                    fallback.clone(),
5171                    RuntimeConfigState::BreakerOpen,
5172                    20,
5173                )
5174                .await;
5175            let runtime = state
5176                .route_plan_runtime_state_for_provider_endpoints("codex")
5177                .await;
5178            assert!(runtime.provider_endpoint(&fallback).runtime_disabled);
5179        });
5180    }
5181
5182    #[test]
5183    fn owned_policy_action_projects_to_runtime_state_below_manual_overrides() {
5184        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5185        runtime.block_on(async {
5186            let state = ProxyState::new();
5187            let endpoint = ProviderEndpointKey::new("codex", "monthly", "default");
5188            let now_ms = unix_now_ms();
5189            let mut signal = crate::provider_signals::ProviderSignal::high_confidence_route_facing(
5190                crate::provider_signals::ProviderSignalKind::Quota,
5191                crate::provider_signals::ProviderSignalSource::UpstreamResponse,
5192                crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5193                    provider_endpoint_key: endpoint.clone(),
5194                },
5195                now_ms,
5196            );
5197            signal.reset_after_secs = Some(60);
5198            signal.reason = Some("usage_limit_reached".to_string());
5199            let action =
5200                crate::policy_actions::PolicyAction::cooldown_from_signal(signal, now_ms, 0, 1)
5201                    .expect("owned cooldown action");
5202
5203            state.upsert_owned_policy_action("codex", action).await;
5204            let runtime = state
5205                .route_plan_runtime_state_for_provider_endpoints("codex")
5206                .await;
5207            let projected = runtime.provider_endpoint(&endpoint);
5208            assert!(projected.cooldown_active);
5209            assert!(projected.usage_exhausted);
5210            assert!(!projected.runtime_disabled);
5211
5212            state
5213                .set_provider_endpoint_runtime_state_override(
5214                    "codex",
5215                    endpoint.clone(),
5216                    RuntimeConfigState::BreakerOpen,
5217                    now_ms.saturating_add(1_000),
5218                )
5219                .await;
5220            let runtime = state
5221                .route_plan_runtime_state_for_provider_endpoints("codex")
5222                .await;
5223            let projected = runtime.provider_endpoint(&endpoint);
5224            assert!(projected.cooldown_active);
5225            assert!(projected.usage_exhausted);
5226            assert!(
5227                projected.runtime_disabled,
5228                "manual runtime override must outrank automatic action projection"
5229            );
5230        });
5231    }
5232
5233    #[test]
5234    fn owned_policy_action_upsert_replaces_existing_action_for_same_source() {
5235        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5236        runtime.block_on(async {
5237            let state = ProxyState::new();
5238            let endpoint = ProviderEndpointKey::new("codex", "monthly", "default");
5239            let signal = crate::provider_signals::ProviderSignal::high_confidence_route_facing(
5240                crate::provider_signals::ProviderSignalKind::Capacity,
5241                crate::provider_signals::ProviderSignalSource::UpstreamResponse,
5242                crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5243                    provider_endpoint_key: endpoint,
5244                },
5245                1_000,
5246            );
5247            let first = crate::policy_actions::PolicyAction::cooldown_from_signal(
5248                signal.clone(),
5249                1_000,
5250                10,
5251                1,
5252            )
5253            .expect("first cooldown");
5254            let second =
5255                crate::policy_actions::PolicyAction::cooldown_from_signal(signal, 2_000, 20, 2)
5256                    .expect("second cooldown");
5257
5258            state.upsert_owned_policy_action("codex", first).await;
5259            state.upsert_owned_policy_action("codex", second).await;
5260
5261            let actions = state.list_policy_actions("codex").await;
5262            assert_eq!(actions.len(), 1);
5263            assert_eq!(actions[0].generation, 2);
5264            assert_eq!(actions[0].expires_at_ms, 22_000);
5265        });
5266    }
5267
5268    #[test]
5269    fn owned_policy_action_clear_is_source_aware() {
5270        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5271        runtime.block_on(async {
5272            let state = ProxyState::new();
5273            let endpoint = ProviderEndpointKey::new("codex", "monthly", "default");
5274            let mut response_signal =
5275                crate::provider_signals::ProviderSignal::high_confidence_route_facing(
5276                    crate::provider_signals::ProviderSignalKind::RateLimit,
5277                    crate::provider_signals::ProviderSignalSource::UpstreamResponse,
5278                    crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5279                        provider_endpoint_key: endpoint.clone(),
5280                    },
5281                    1_000,
5282                );
5283            response_signal.reset_after_secs = Some(30);
5284            let balance_signal = crate::provider_signals::ProviderSignal {
5285                kind: crate::provider_signals::ProviderSignalKind::Balance,
5286                source: crate::provider_signals::ProviderSignalSource::BalanceSnapshot,
5287                target: crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5288                    provider_endpoint_key: endpoint.clone(),
5289                },
5290                confidence: crate::provider_signals::ProviderSignalConfidence::High,
5291                observed_at_ms: 1_000,
5292                route_facing: true,
5293                retry_after_secs: None,
5294                reset_after_secs: Some(60),
5295                reason: Some("balance_exhausted".to_string()),
5296                error_class: None,
5297                trace: Default::default(),
5298            };
5299            let response_action = crate::policy_actions::PolicyAction::cooldown_from_signal(
5300                response_signal,
5301                1_000,
5302                30,
5303                1,
5304            )
5305            .expect("response cooldown");
5306            let balance_action = crate::policy_actions::PolicyAction::cooldown_from_signal(
5307                balance_signal,
5308                1_100,
5309                0,
5310                2,
5311            )
5312            .expect("balance cooldown");
5313
5314            state
5315                .upsert_owned_policy_action("codex", response_action)
5316                .await;
5317            state
5318                .upsert_owned_policy_action("codex", balance_action)
5319                .await;
5320            assert_eq!(state.list_policy_actions("codex").await.len(), 2);
5321
5322            state
5323                .clear_owned_policy_action(
5324                    "codex",
5325                    &endpoint,
5326                    crate::policy_actions::PolicyActionKind::Cooldown,
5327                    crate::provider_signals::ProviderSignalKind::Balance,
5328                    crate::provider_signals::ProviderSignalSource::BalanceSnapshot,
5329                )
5330                .await;
5331
5332            let actions = state.list_policy_actions("codex").await;
5333            assert_eq!(actions.len(), 1);
5334            assert_eq!(
5335                actions[0].source_signal.kind,
5336                crate::provider_signals::ProviderSignalKind::RateLimit
5337            );
5338        });
5339    }
5340
5341    #[test]
5342    fn expired_owned_policy_action_is_not_projected() {
5343        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5344        runtime.block_on(async {
5345            let state = ProxyState::new();
5346            let endpoint = ProviderEndpointKey::new("codex", "monthly", "default");
5347            let signal = crate::provider_signals::ProviderSignal::high_confidence_route_facing(
5348                crate::provider_signals::ProviderSignalKind::Transport,
5349                crate::provider_signals::ProviderSignalSource::UpstreamResponse,
5350                crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5351                    provider_endpoint_key: endpoint.clone(),
5352                },
5353                1_000,
5354            );
5355            let action =
5356                crate::policy_actions::PolicyAction::cooldown_from_signal(signal, 1_000, 1, 1)
5357                    .expect("transport cooldown");
5358
5359            state.upsert_owned_policy_action("codex", action).await;
5360
5361            assert!(
5362                state
5363                    .active_policy_action_projections("codex", 2_001)
5364                    .await
5365                    .is_empty()
5366            );
5367            let runtime = state
5368                .route_plan_runtime_state_for_provider_endpoints("codex")
5369                .await;
5370            let projected = runtime.provider_endpoint(&endpoint);
5371            assert!(!projected.cooldown_active);
5372            assert!(!projected.usage_exhausted);
5373        });
5374    }
5375
5376    #[test]
5377    fn get_station_health_merges_passive_runtime_observations() {
5378        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5379        runtime.block_on(async {
5380            let state = ProxyState::new();
5381            state
5382                .record_station_health(
5383                    "codex",
5384                    "right".to_string(),
5385                    StationHealth {
5386                        checked_at_ms: 10,
5387                        upstreams: vec![UpstreamHealth {
5388                            base_url: "https://right.example/v1".to_string(),
5389                            ok: Some(true),
5390                            status_code: Some(200),
5391                            latency_ms: Some(120),
5392                            error: None,
5393                            passive: None,
5394                        }],
5395                    },
5396                )
5397                .await;
5398            state
5399                .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
5400                    service_name: "codex".to_string(),
5401                    station_name: "right".to_string(),
5402                    base_url: "https://right.example/v1".to_string(),
5403                    status_code: Some(500),
5404                    error_class: Some("cloudflare_timeout".to_string()),
5405                    error: Some("upstream timed out".to_string()),
5406                    now_ms: 20,
5407                })
5408                .await;
5409            state
5410                .record_passive_upstream_success(
5411                    "codex",
5412                    "right",
5413                    "https://backup.example/v1",
5414                    Some(200),
5415                    30,
5416                )
5417                .await;
5418
5419            let health = state.get_station_health("codex").await;
5420            let right = health.get("right").expect("right health");
5421            assert_eq!(right.checked_at_ms, 30);
5422            assert_eq!(right.upstreams.len(), 2);
5423
5424            let primary = right
5425                .upstreams
5426                .iter()
5427                .find(|upstream| upstream.base_url == "https://right.example/v1")
5428                .expect("primary upstream");
5429            assert_eq!(primary.ok, Some(true));
5430            let primary_passive = primary.passive.as_ref().expect("primary passive");
5431            assert_eq!(primary_passive.state, PassiveHealthState::Degraded);
5432            assert_eq!(primary_passive.score, 50);
5433            assert_eq!(primary_passive.last_status_code, Some(500));
5434            assert_eq!(
5435                primary_passive.last_error_class.as_deref(),
5436                Some("cloudflare_timeout")
5437            );
5438
5439            let backup = right
5440                .upstreams
5441                .iter()
5442                .find(|upstream| upstream.base_url == "https://backup.example/v1")
5443                .expect("backup upstream");
5444            assert_eq!(backup.ok, None);
5445            let backup_passive = backup.passive.as_ref().expect("backup passive");
5446            assert_eq!(backup_passive.state, PassiveHealthState::Healthy);
5447            assert_eq!(backup_passive.score, 100);
5448            assert_eq!(backup_passive.last_status_code, Some(200));
5449        });
5450    }
5451
5452    #[test]
5453    fn passive_health_success_recovers_after_failure() {
5454        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5455        runtime.block_on(async {
5456            let state = ProxyState::new();
5457            state
5458                .record_passive_upstream_failure(PassiveUpstreamFailureRecord {
5459                    service_name: "codex".to_string(),
5460                    station_name: "right".to_string(),
5461                    base_url: "https://right.example/v1".to_string(),
5462                    status_code: Some(500),
5463                    error_class: Some("cloudflare_timeout".to_string()),
5464                    error: Some("upstream timed out".to_string()),
5465                    now_ms: 10,
5466                })
5467                .await;
5468            state
5469                .record_passive_upstream_success(
5470                    "codex",
5471                    "right",
5472                    "https://right.example/v1",
5473                    Some(200),
5474                    20,
5475                )
5476                .await;
5477
5478            let health = state.get_station_health("codex").await;
5479            let right = health.get("right").expect("right health");
5480            let upstream = right.upstreams.first().expect("upstream");
5481            let passive = upstream.passive.as_ref().expect("passive");
5482            assert_eq!(passive.state, PassiveHealthState::Healthy);
5483            assert_eq!(passive.score, 100);
5484            assert_eq!(passive.consecutive_failures, 0);
5485            assert_eq!(passive.last_success_at_ms, Some(20));
5486            assert_eq!(passive.last_failure_at_ms, Some(10));
5487            assert_eq!(passive.last_error_class, None);
5488            assert_eq!(passive.last_error, None);
5489        });
5490    }
5491
5492    #[test]
5493    fn apply_session_profile_binding_uses_inherited_profile_values() {
5494        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5495        runtime.block_on(async {
5496            let state = ProxyState::new();
5497            let mut mgr = ServiceConfigManager::default();
5498            mgr.configs.insert(
5499                "right".to_string(),
5500                ServiceConfig {
5501                    name: "right".to_string(),
5502                    alias: None,
5503                    enabled: true,
5504                    level: 1,
5505                    upstreams: vec![UpstreamConfig {
5506                        base_url: "https://right.example/v1".to_string(),
5507                        auth: UpstreamAuth::default(),
5508                        tags: HashMap::from([
5509                            ("supports_reasoning_effort".to_string(), "true".to_string()),
5510                            ("supports_service_tier".to_string(), "true".to_string()),
5511                        ]),
5512                        supported_models: HashMap::from([("gpt-5.4".to_string(), true)]),
5513                        model_mapping: HashMap::new(),
5514                    }],
5515                },
5516            );
5517            mgr.profiles.insert(
5518                "base".to_string(),
5519                crate::config::ServiceControlProfile {
5520                    extends: None,
5521                    station: Some("right".to_string()),
5522                    model: Some("gpt-5.4".to_string()),
5523                    reasoning_effort: None,
5524                    service_tier: Some("priority".to_string()),
5525                },
5526            );
5527            mgr.profiles.insert(
5528                "fast".to_string(),
5529                crate::config::ServiceControlProfile {
5530                    extends: Some("base".to_string()),
5531                    station: None,
5532                    model: None,
5533                    reasoning_effort: Some("low".to_string()),
5534                    service_tier: None,
5535                },
5536            );
5537
5538            state
5539                .apply_session_profile_binding(
5540                    "codex",
5541                    &mgr,
5542                    "sid-inherited".to_string(),
5543                    "fast".to_string(),
5544                    100,
5545                )
5546                .await
5547                .expect("apply inherited profile");
5548
5549            let binding = state
5550                .get_session_binding("sid-inherited")
5551                .await
5552                .expect("binding exists");
5553            assert_eq!(binding.profile_name.as_deref(), Some("fast"));
5554            assert_eq!(binding.station_name.as_deref(), Some("right"));
5555            assert_eq!(binding.model.as_deref(), Some("gpt-5.4"));
5556            assert_eq!(binding.reasoning_effort.as_deref(), Some("low"));
5557            assert_eq!(binding.service_tier.as_deref(), Some("priority"));
5558        });
5559    }
5560
5561    #[test]
5562    fn prune_periodic_keeps_sticky_binding_after_manual_override_ttl_expires() {
5563        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5564        runtime.block_on(async {
5565            let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(1, 0, 2_000));
5566            state
5567                .set_session_model_override("sid-sticky".to_string(), "gpt-5.4".to_string(), 0)
5568                .await;
5569            state
5570                .set_session_binding(SessionBinding {
5571                    session_id: "sid-sticky".to_string(),
5572                    profile_name: Some("daily".to_string()),
5573                    station_name: Some("right".to_string()),
5574                    model: Some("gpt-5.4".to_string()),
5575                    reasoning_effort: Some("medium".to_string()),
5576                    service_tier: Some("default".to_string()),
5577                    continuity_mode: SessionContinuityMode::DefaultProfile,
5578                    created_at_ms: 0,
5579                    updated_at_ms: 0,
5580                    last_seen_ms: 0,
5581                })
5582                .await;
5583
5584            state.prune_periodic().await;
5585
5586            assert!(
5587                state
5588                    .get_session_model_override("sid-sticky")
5589                    .await
5590                    .is_none()
5591            );
5592            assert!(state.get_session_binding("sid-sticky").await.is_some());
5593        });
5594    }
5595
5596    #[test]
5597    fn prune_periodic_honors_opt_in_binding_ttl() {
5598        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5599        runtime.block_on(async {
5600            let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 1, 2_000));
5601            state
5602                .set_session_binding(SessionBinding {
5603                    session_id: "sid-expire".to_string(),
5604                    profile_name: Some("daily".to_string()),
5605                    station_name: Some("right".to_string()),
5606                    model: Some("gpt-5.4".to_string()),
5607                    reasoning_effort: Some("medium".to_string()),
5608                    service_tier: Some("default".to_string()),
5609                    continuity_mode: SessionContinuityMode::DefaultProfile,
5610                    created_at_ms: 0,
5611                    updated_at_ms: 0,
5612                    last_seen_ms: 0,
5613                })
5614                .await;
5615
5616            state.prune_periodic().await;
5617
5618            assert!(state.get_session_binding("sid-expire").await.is_none());
5619        });
5620    }
5621
5622    #[test]
5623    fn session_route_affinity_ttl_is_enforced_on_read() {
5624        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5625        runtime.block_on(async {
5626            let mut policy = test_runtime_policy(0, 0, 2_000);
5627            policy.session_route_affinity_ttl_ms = 1;
5628            let state = ProxyState::new_with_runtime_policy(None, policy);
5629            state
5630                .record_session_route_affinity_success(
5631                    "sid-expire",
5632                    SessionRouteAffinityTarget {
5633                        route_graph_key: "graph".to_string(),
5634                        session_identity_source: None,
5635                        provider_endpoint: ProviderEndpointKey::new("codex", "monthly", "default"),
5636                        upstream_base_url: "https://monthly.example/v1".to_string(),
5637                        route_path: vec!["monthly_first".to_string(), "monthly".to_string()],
5638                    },
5639                    Some("first_success".to_string()),
5640                    0,
5641                )
5642                .await;
5643
5644            assert!(
5645                state
5646                    .get_session_route_affinity("sid-expire")
5647                    .await
5648                    .is_none()
5649            );
5650            assert!(state.list_session_route_affinities().await.is_empty());
5651        });
5652    }
5653
5654    #[test]
5655    fn session_route_affinity_ledger_does_not_restore_expired_entries() {
5656        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5657        runtime.block_on(async {
5658            let _home = env_lock().await;
5659            let temp_dir = std::env::temp_dir().join(format!(
5660                "codex-helper-route-ledger-test-{}",
5661                uuid::Uuid::new_v4()
5662            ));
5663            std::fs::create_dir_all(&temp_dir).expect("create temp dir");
5664            let mut scoped = ScopedEnv::default();
5665            unsafe {
5666                scoped.set_path("CODEX_HELPER_HOME", temp_dir.as_path());
5667            }
5668
5669            let mut policy = test_runtime_policy(0, 0, 2_000);
5670            policy.session_route_affinity_ttl_ms = 1;
5671            let first_state = ProxyState::new_with_runtime_policy(None, policy.clone());
5672            first_state
5673                .record_session_route_affinity_success(
5674                    "sid-expired-ledger",
5675                    SessionRouteAffinityTarget {
5676                        route_graph_key: "graph".to_string(),
5677                        session_identity_source: None,
5678                        provider_endpoint: ProviderEndpointKey::new("codex", "monthly", "default"),
5679                        upstream_base_url: "https://monthly.example/v1".to_string(),
5680                        route_path: vec!["monthly_first".to_string(), "monthly".to_string()],
5681                    },
5682                    Some("first_success".to_string()),
5683                    0,
5684                )
5685                .await;
5686
5687            let second_state = ProxyState::new_with_runtime_policy(None, policy);
5688            assert!(
5689                second_state
5690                    .get_session_route_affinity("sid-expired-ledger")
5691                    .await
5692                    .is_none()
5693            );
5694        });
5695    }
5696
5697    #[test]
5698    fn policy_action_ledger_restores_owned_actions() {
5699        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5700        runtime.block_on(async {
5701            let _home = env_lock().await;
5702            let temp_dir = std::env::temp_dir().join(format!(
5703                "codex-helper-policy-action-ledger-test-{}",
5704                uuid::Uuid::new_v4()
5705            ));
5706            std::fs::create_dir_all(&temp_dir).expect("create temp dir");
5707            let mut scoped = ScopedEnv::default();
5708            unsafe {
5709                scoped.set_path("CODEX_HELPER_HOME", temp_dir.as_path());
5710                scoped.set_path(
5711                    "CODEX_HELPER_POLICY_ACTION_LEDGER",
5712                    temp_dir.join("policy-actions.json").as_path(),
5713                );
5714            }
5715
5716            let first_state =
5717                ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5718            let mut signal = crate::provider_signals::ProviderSignal::high_confidence_route_facing(
5719                crate::provider_signals::ProviderSignalKind::RateLimit,
5720                crate::provider_signals::ProviderSignalSource::UpstreamResponse,
5721                crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5722                    provider_endpoint_key: ProviderEndpointKey::new("codex", "monthly", "default"),
5723                },
5724                unix_now_ms(),
5725            );
5726            signal.reset_after_secs = Some(60);
5727            let action = crate::policy_actions::PolicyAction::cooldown_from_signal(
5728                signal,
5729                unix_now_ms(),
5730                0,
5731                1,
5732            )
5733            .expect("cooldown action");
5734            first_state
5735                .upsert_owned_policy_action("codex", action)
5736                .await;
5737
5738            let second_state =
5739                ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5740            let actions = second_state.list_policy_actions("codex").await;
5741            assert_eq!(actions.len(), 1);
5742            assert_eq!(
5743                actions[0].provider_endpoint_key.stable_key(),
5744                "codex/monthly/default"
5745            );
5746        });
5747    }
5748
5749    #[test]
5750    fn policy_action_upsert_does_not_publish_when_ledger_save_fails() {
5751        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5752        runtime.block_on(async {
5753            let _home = env_lock().await;
5754            let temp_dir = std::env::temp_dir().join(format!(
5755                "codex-helper-policy-action-save-failure-test-{}",
5756                uuid::Uuid::new_v4()
5757            ));
5758            std::fs::create_dir_all(&temp_dir).expect("create temp dir");
5759            let mut scoped = ScopedEnv::default();
5760            unsafe {
5761                scoped.set_path("CODEX_HELPER_HOME", temp_dir.as_path());
5762                scoped.set_path("CODEX_HELPER_POLICY_ACTION_LEDGER", temp_dir.as_path());
5763            }
5764
5765            let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5766            let endpoint = ProviderEndpointKey::new("codex", "monthly", "default");
5767            let mut signal = crate::provider_signals::ProviderSignal::high_confidence_route_facing(
5768                crate::provider_signals::ProviderSignalKind::RateLimit,
5769                crate::provider_signals::ProviderSignalSource::UpstreamResponse,
5770                crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5771                    provider_endpoint_key: endpoint.clone(),
5772                },
5773                unix_now_ms(),
5774            );
5775            signal.reset_after_secs = Some(60);
5776            let action = crate::policy_actions::PolicyAction::cooldown_from_signal(
5777                signal,
5778                unix_now_ms(),
5779                0,
5780                1,
5781            )
5782            .expect("cooldown action");
5783
5784            state.upsert_owned_policy_action("codex", action).await;
5785
5786            assert!(state.list_policy_actions("codex").await.is_empty());
5787            let runtime_state = state
5788                .route_plan_runtime_state_for_provider_endpoints("codex")
5789                .await;
5790            assert!(!runtime_state.provider_endpoint(&endpoint).cooldown_active);
5791        });
5792    }
5793
5794    #[test]
5795    fn policy_action_ledger_drops_expired_actions_on_restore() {
5796        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5797        runtime.block_on(async {
5798            let _home = env_lock().await;
5799            let temp_dir = std::env::temp_dir().join(format!(
5800                "codex-helper-policy-action-expired-ledger-test-{}",
5801                uuid::Uuid::new_v4()
5802            ));
5803            std::fs::create_dir_all(&temp_dir).expect("create temp dir");
5804            let ledger_path = temp_dir.join("policy-actions.json");
5805            let mut scoped = ScopedEnv::default();
5806            unsafe {
5807                scoped.set_path("CODEX_HELPER_HOME", temp_dir.as_path());
5808                scoped.set_path("CODEX_HELPER_POLICY_ACTION_LEDGER", ledger_path.as_path());
5809            }
5810
5811            let endpoint = ProviderEndpointKey::new("codex", "monthly", "default");
5812            let signal = crate::provider_signals::ProviderSignal::high_confidence_route_facing(
5813                crate::provider_signals::ProviderSignalKind::Transport,
5814                crate::provider_signals::ProviderSignalSource::UpstreamResponse,
5815                crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5816                    provider_endpoint_key: endpoint,
5817                },
5818                1_000,
5819            );
5820            let action =
5821                crate::policy_actions::PolicyAction::cooldown_from_signal(signal, 1_000, 1, 1)
5822                    .expect("transport cooldown");
5823            let ledger = serde_json::json!({
5824                "schema_version": 1,
5825                "updated_at_ms": 2_000,
5826                "entries": [
5827                    {
5828                        "service_name": "codex",
5829                        "action": action
5830                    }
5831                ]
5832            });
5833            std::fs::write(
5834                &ledger_path,
5835                serde_json::to_string_pretty(&ledger).expect("ledger json"),
5836            )
5837            .expect("write ledger");
5838
5839            let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5840
5841            assert!(state.list_policy_actions("codex").await.is_empty());
5842            let compacted: serde_json::Value = serde_json::from_str(
5843                &std::fs::read_to_string(&ledger_path).expect("compacted ledger"),
5844            )
5845            .expect("parse compacted ledger");
5846            assert_eq!(
5847                compacted
5848                    .get("entries")
5849                    .and_then(|entries| entries.as_array())
5850                    .map(Vec::len),
5851                Some(0)
5852            );
5853        });
5854    }
5855
5856    #[test]
5857    fn policy_action_prune_persists_removed_endpoint_actions() {
5858        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5859        runtime.block_on(async {
5860            let _home = env_lock().await;
5861            let temp_dir = std::env::temp_dir().join(format!(
5862                "codex-helper-policy-action-prune-ledger-test-{}",
5863                uuid::Uuid::new_v4()
5864            ));
5865            std::fs::create_dir_all(&temp_dir).expect("create temp dir");
5866            let ledger_path = temp_dir.join("policy-actions.json");
5867            let mut scoped = ScopedEnv::default();
5868            unsafe {
5869                scoped.set_path("CODEX_HELPER_HOME", temp_dir.as_path());
5870                scoped.set_path("CODEX_HELPER_POLICY_ACTION_LEDGER", ledger_path.as_path());
5871            }
5872
5873            let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5874            let endpoint = ProviderEndpointKey::new("codex", "monthly", "default");
5875            let mut signal = crate::provider_signals::ProviderSignal::high_confidence_route_facing(
5876                crate::provider_signals::ProviderSignalKind::RateLimit,
5877                crate::provider_signals::ProviderSignalSource::UpstreamResponse,
5878                crate::provider_signals::ProviderSignalTarget::ProviderEndpoint {
5879                    provider_endpoint_key: endpoint,
5880                },
5881                unix_now_ms(),
5882            );
5883            signal.reset_after_secs = Some(60);
5884            let action = crate::policy_actions::PolicyAction::cooldown_from_signal(
5885                signal,
5886                unix_now_ms(),
5887                0,
5888                1,
5889            )
5890            .expect("cooldown action");
5891            state.upsert_owned_policy_action("codex", action).await;
5892            assert_eq!(state.list_policy_actions("codex").await.len(), 1);
5893
5894            let mut mgr = ServiceConfigManager::default();
5895            mgr.configs.insert(
5896                "routing".to_string(),
5897                ServiceConfig {
5898                    name: "routing".to_string(),
5899                    alias: None,
5900                    enabled: true,
5901                    level: 1,
5902                    upstreams: vec![UpstreamConfig {
5903                        base_url: "https://backup.example/v1".to_string(),
5904                        auth: UpstreamAuth::default(),
5905                        tags: HashMap::from([
5906                            ("provider_id".to_string(), "backup".to_string()),
5907                            ("endpoint_id".to_string(), "default".to_string()),
5908                        ]),
5909                        supported_models: HashMap::new(),
5910                        model_mapping: HashMap::new(),
5911                    }],
5912                },
5913            );
5914            state
5915                .prune_runtime_observability_for_service("codex", &mgr)
5916                .await;
5917            assert!(state.list_policy_actions("codex").await.is_empty());
5918
5919            let restored =
5920                ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5921            assert!(restored.list_policy_actions("codex").await.is_empty());
5922        });
5923    }
5924
5925    #[test]
5926    fn prune_periodic_caps_sticky_bindings_by_last_seen() {
5927        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5928        runtime.block_on(async {
5929            let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2));
5930            for idx in 0..3 {
5931                state
5932                    .set_session_binding(SessionBinding {
5933                        session_id: format!("sid-{idx}"),
5934                        profile_name: Some("daily".to_string()),
5935                        station_name: Some("right".to_string()),
5936                        model: Some("gpt-5.4".to_string()),
5937                        reasoning_effort: Some("medium".to_string()),
5938                        service_tier: Some("default".to_string()),
5939                        continuity_mode: SessionContinuityMode::DefaultProfile,
5940                        created_at_ms: idx,
5941                        updated_at_ms: idx,
5942                        last_seen_ms: idx,
5943                    })
5944                    .await;
5945            }
5946
5947            state.prune_periodic().await;
5948
5949            assert!(state.get_session_binding("sid-0").await.is_none());
5950            assert!(state.get_session_binding("sid-1").await.is_some());
5951            assert!(state.get_session_binding("sid-2").await.is_some());
5952        });
5953    }
5954
5955    #[test]
5956    fn transcript_path_cache_records_negative_lookups() {
5957        let runtime = tokio::runtime::Runtime::new().expect("runtime");
5958        runtime.block_on(async {
5959            let state = ProxyState::new_with_runtime_policy(None, test_runtime_policy(0, 0, 2_000));
5960            let paths = state
5961                .resolve_host_transcript_paths_cached(&["missing-session".to_string()])
5962                .await;
5963
5964            assert_eq!(paths.get("missing-session"), Some(&None));
5965            let cache = state.session_transcript_path_cache.read().await;
5966            assert!(cache.contains_key("missing-session"));
5967        });
5968    }
5969}