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