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#[derive(Debug)]
513pub struct ProxyState {
514 next_request_id: AtomicU64,
515 session_override_ttl_ms: u64,
517 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 provider_endpoint_meta_overrides:
540 RwLock<HashMap<String, HashMap<ProviderEndpointKey, ConfigMetaOverride>>>,
541 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 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 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(¶ms.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 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 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 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 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}