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