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