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