1use std::collections::HashMap;
2use std::sync::{Arc, Mutex};
3
4use reqwest::Client;
5use tokio::sync::watch;
6
7use crate::config::{
8 PersistedProviderSpec, PersistedProvidersCatalog, PersistedRoutingSpec, ProxyConfig,
9 ProxyConfigV4, ServiceConfigManager,
10};
11use crate::filter::RequestFilter;
12use crate::lb::LbState;
13use crate::routing_explain::RoutingExplainResponse;
14use crate::routing_ir::RouteRequestContext;
15use crate::state::{ProxyState, SessionBinding, SessionContinuityMode};
16
17use super::profile_defaults::effective_default_profile_name;
18use super::{
19 CodexRelayCapabilitiesRequest, CodexRelayCapabilitiesResponse, CodexRelayLiveSmokeRequest,
20 CodexRelayLiveSmokeResponse, PersistedRoutingUpsertRequest, ProfilesResponse,
21 ProviderBalanceRefreshResponse, ProxyControlError, ProxyService, ReloadResult, RuntimeConfig,
22 RuntimeStatusResponse,
23};
24
25impl ProxyService {
26 pub fn new(
27 client: Client,
28 config: Arc<ProxyConfig>,
29 service_name: &'static str,
30 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
31 ) -> Self {
32 Self::new_with_v4_source(client, config, None, service_name, lb_states)
33 }
34
35 pub fn new_with_v4_source(
36 client: Client,
37 config: Arc<ProxyConfig>,
38 v4_source: Option<Arc<ProxyConfigV4>>,
39 service_name: &'static str,
40 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
41 ) -> Self {
42 Self::new_with_v4_source_and_shutdown(
43 client,
44 config,
45 v4_source,
46 service_name,
47 lb_states,
48 None,
49 )
50 }
51
52 pub fn new_with_v4_source_and_shutdown(
53 client: Client,
54 config: Arc<ProxyConfig>,
55 v4_source: Option<Arc<ProxyConfigV4>>,
56 service_name: &'static str,
57 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
58 shutdown_tx: Option<watch::Sender<bool>>,
59 ) -> Self {
60 let state = ProxyState::new_with_lb_states(Some(lb_states.clone()));
61 ProxyState::spawn_cleanup_task(state.clone());
62 if !cfg!(test) {
63 let state = state.clone();
64 let log_path = crate::logging::request_log_path();
65 let mut base_url_to_provider_id = HashMap::new();
66 let mgr = match service_name {
67 "claude" => &config.claude,
68 _ => &config.codex,
69 };
70 for svc in mgr.stations().values() {
71 for up in &svc.upstreams {
72 if let Some(pid) = up.tags.get("provider_id") {
73 base_url_to_provider_id.insert(up.base_url.clone(), pid.clone());
74 }
75 }
76 }
77 tokio::spawn(async move {
78 let _ = state
79 .replay_usage_from_requests_log(service_name, log_path, base_url_to_provider_id)
80 .await;
81 });
82 }
83 Self {
84 client,
85 config: Arc::new(RuntimeConfig::new_with_v4(config, v4_source)),
86 service_name,
87 lb_states,
88 concurrency_limiter: Arc::new(super::concurrency_limits::ConcurrencyLimiter::default()),
89 filter: RequestFilter::new(),
90 state,
91 shutdown_tx,
92 host_local_session_history_mode: crate::host_local::HostLocalSessionHistoryMode::Auto,
93 }
94 }
95
96 pub fn with_host_local_session_history_mode(
97 mut self,
98 mode: crate::host_local::HostLocalSessionHistoryMode,
99 ) -> Self {
100 self.host_local_session_history_mode = mode;
101 self
102 }
103
104 pub(crate) fn host_local_session_history_available(&self) -> bool {
105 crate::host_local::host_local_session_history_available_for_mode(
106 self.host_local_session_history_mode,
107 )
108 }
109
110 pub(super) fn service_manager<'a>(&self, cfg: &'a ProxyConfig) -> &'a ServiceConfigManager {
111 match self.service_name {
112 "codex" => &cfg.codex,
113 "claude" => &cfg.claude,
114 _ => &cfg.codex,
115 }
116 }
117
118 pub(super) async fn ensure_default_session_binding(
119 &self,
120 mgr: &ServiceConfigManager,
121 session_id: &str,
122 now_ms: u64,
123 ) -> Option<SessionBinding> {
124 if let Some(binding) = self.state.get_session_binding(session_id).await {
125 self.state.touch_session_binding(session_id, now_ms).await;
126 return Some(binding);
127 }
128
129 let profile_name =
130 effective_default_profile_name(self.state.as_ref(), self.service_name, mgr).await?;
131 let profile = crate::config::resolve_service_profile(mgr, profile_name.as_str()).ok()?;
132 let binding = SessionBinding {
133 session_id: session_id.to_string(),
134 profile_name: Some(profile_name),
135 station_name: profile.station.clone(),
136 model: profile.model.clone(),
137 reasoning_effort: profile.reasoning_effort.clone(),
138 service_tier: profile.service_tier.clone(),
139 continuity_mode: SessionContinuityMode::DefaultProfile,
140 created_at_ms: now_ms,
141 updated_at_ms: now_ms,
142 last_seen_ms: now_ms,
143 };
144 self.state.set_session_binding(binding.clone()).await;
145 Some(binding)
146 }
147
148 pub fn state_handle(&self) -> Arc<ProxyState> {
149 self.state.clone()
150 }
151
152 pub fn shutdown_available(&self) -> bool {
153 self.shutdown_tx.is_some()
154 }
155
156 pub fn request_shutdown(&self) -> bool {
157 self.shutdown_tx
158 .as_ref()
159 .is_some_and(|shutdown_tx| shutdown_tx.send(true).is_ok())
160 }
161
162 pub fn spawn_initial_balance_refresh(&self) {
163 if cfg!(test) {
164 return;
165 }
166
167 let proxy = self.clone();
168 tokio::spawn(async move {
169 match super::providers_api::refresh_provider_balances_for_proxy(&proxy, None, None)
170 .await
171 {
172 Ok(summary) => {
173 tracing::info!(
174 "initial provider balance refresh finished: attempted={}, refreshed={}, failed={}, missing_token={}, auto_refreshed={}",
175 summary.attempted,
176 summary.refreshed,
177 summary.failed,
178 summary.missing_token,
179 summary.auto_refreshed
180 );
181 }
182 Err((status, message)) => {
183 tracing::warn!(
184 "initial provider balance refresh failed before polling: status={}, {}",
185 status,
186 message
187 );
188 }
189 }
190 });
191 }
192
193 pub async fn refresh_provider_balances(
194 &self,
195 station_name_filter: Option<&str>,
196 provider_id_filter: Option<&str>,
197 ) -> Result<ProviderBalanceRefreshResponse, ProxyControlError> {
198 let refresh = super::providers_api::refresh_provider_balances_for_proxy(
199 self,
200 station_name_filter,
201 provider_id_filter,
202 )
203 .await
204 .map_err(ProxyControlError::from)?;
205 let provider_balances = self
206 .state
207 .get_provider_balance_view(self.service_name)
208 .await;
209
210 Ok(ProviderBalanceRefreshResponse {
211 service_name: self.service_name.to_string(),
212 refresh,
213 provider_balances,
214 })
215 }
216
217 pub async fn runtime_status(&self) -> RuntimeStatusResponse {
218 super::api_responses::build_runtime_status_response(self).await
219 }
220
221 pub async fn codex_relay_capabilities(
222 &self,
223 request: CodexRelayCapabilitiesRequest,
224 ) -> Result<CodexRelayCapabilitiesResponse, ProxyControlError> {
225 super::codex_relay_capabilities::codex_relay_capabilities_for_proxy(self, request).await
226 }
227
228 pub async fn codex_relay_live_smoke(
229 &self,
230 request: CodexRelayLiveSmokeRequest,
231 ) -> Result<CodexRelayLiveSmokeResponse, ProxyControlError> {
232 super::codex_relay_live_smoke::codex_relay_live_smoke_for_proxy(self, request).await
233 }
234
235 pub async fn reload_runtime_config(&self) -> Result<ReloadResult, ProxyControlError> {
236 let changed = self.config.force_reload_from_disk().await.map_err(|err| {
237 ProxyControlError::new(
238 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
239 err.to_string(),
240 )
241 })?;
242 if changed {
243 super::control_plane_service::prune_runtime_observability_after_reload(self).await;
244 }
245 let status = self.runtime_status().await;
246 Ok(super::api_responses::build_reload_result(changed, status))
247 }
248
249 pub async fn profiles(&self) -> ProfilesResponse {
250 super::api_responses::make_profiles_response(self).await
251 }
252
253 pub async fn set_runtime_default_profile(
254 &self,
255 profile_name: Option<String>,
256 ) -> Result<(), ProxyControlError> {
257 let profile_name = normalize_optional_control_name(profile_name);
258
259 if let Some(profile_name) = profile_name {
260 let cfg = self.config.snapshot().await;
261 let mgr = self.service_manager(cfg.as_ref());
262 if mgr.profile(profile_name.as_str()).is_none() {
263 return Err(ProxyControlError::new(
264 axum::http::StatusCode::NOT_FOUND,
265 format!("profile '{}' not found", profile_name),
266 ));
267 }
268 let resolved = crate::config::resolve_service_profile(mgr, profile_name.as_str())
269 .map_err(|err| {
270 ProxyControlError::new(axum::http::StatusCode::BAD_REQUEST, err.to_string())
271 })?;
272 crate::config::validate_profile_station_compatibility(
273 self.service_name,
274 mgr,
275 profile_name.as_str(),
276 &resolved,
277 )
278 .map_err(|err| {
279 ProxyControlError::new(axum::http::StatusCode::BAD_REQUEST, err.to_string())
280 })?;
281 self.state
282 .set_runtime_default_profile_override(
283 self.service_name.to_string(),
284 profile_name,
285 crate::logging::now_ms(),
286 )
287 .await;
288 } else {
289 self.state
290 .clear_runtime_default_profile_override(self.service_name)
291 .await;
292 }
293
294 Ok(())
295 }
296
297 pub async fn set_persisted_default_profile(
298 &self,
299 profile_name: Option<String>,
300 ) -> Result<ProfilesResponse, ProxyControlError> {
301 let profile_name = normalize_optional_control_name(profile_name);
302
303 if let super::control_plane_service::PersistedProxySettingsDocument::V4(mut document) =
304 super::control_plane_service::load_persisted_proxy_settings_document()
305 .await
306 .map_err(ProxyControlError::from)?
307 {
308 if let Some(profile_name) = profile_name.as_deref() {
309 let view =
310 super::control_plane_service::service_view_v4(&document, self.service_name);
311 if !view.profiles.contains_key(profile_name) {
312 return Err(ProxyControlError::new(
313 axum::http::StatusCode::NOT_FOUND,
314 format!("profile '{}' not found", profile_name),
315 ));
316 }
317 let runtime = crate::config::compile_v4_to_runtime(&document).map_err(|err| {
318 ProxyControlError::new(axum::http::StatusCode::BAD_REQUEST, err.to_string())
319 })?;
320 let mgr = super::persisted_registry_api::runtime_service_manager_for_document(
321 &runtime,
322 self.service_name,
323 );
324 let resolved =
325 crate::config::resolve_service_profile(mgr, profile_name).map_err(|err| {
326 ProxyControlError::new(axum::http::StatusCode::BAD_REQUEST, err.to_string())
327 })?;
328 crate::config::validate_profile_station_compatibility(
329 self.service_name,
330 mgr,
331 profile_name,
332 &resolved,
333 )
334 .map_err(|err| {
335 ProxyControlError::new(axum::http::StatusCode::BAD_REQUEST, err.to_string())
336 })?;
337 }
338 let view =
339 super::control_plane_service::service_view_v4_mut(&mut document, self.service_name);
340 view.default_profile = profile_name;
341 super::control_plane_service::save_persisted_proxy_settings_document_and_reload(
342 self,
343 super::control_plane_service::PersistedProxySettingsDocument::V4(document),
344 )
345 .await
346 .map_err(ProxyControlError::from)?;
347 return Ok(self.profiles().await);
348 }
349
350 let cfg_snapshot = self.config.snapshot().await;
351 let mut cfg = cfg_snapshot.as_ref().clone();
352 let mgr =
353 super::control_plane_service::runtime_service_manager_mut(&mut cfg, self.service_name);
354
355 if let Some(profile_name) = profile_name.as_deref() {
356 if mgr.profile(profile_name).is_none() {
357 return Err(ProxyControlError::new(
358 axum::http::StatusCode::NOT_FOUND,
359 format!("profile '{}' not found", profile_name),
360 ));
361 }
362 let resolved =
363 crate::config::resolve_service_profile(mgr, profile_name).map_err(|err| {
364 ProxyControlError::new(axum::http::StatusCode::BAD_REQUEST, err.to_string())
365 })?;
366 crate::config::validate_profile_station_compatibility(
367 self.service_name,
368 mgr,
369 profile_name,
370 &resolved,
371 )
372 .map_err(|err| {
373 ProxyControlError::new(axum::http::StatusCode::BAD_REQUEST, err.to_string())
374 })?;
375 }
376 mgr.default_profile = profile_name;
377
378 super::control_plane_service::save_runtime_profile_settings_and_reload(self, cfg)
379 .await
380 .map_err(ProxyControlError::from)
381 }
382
383 pub async fn persisted_provider_specs(
384 &self,
385 ) -> Result<PersistedProvidersCatalog, ProxyControlError> {
386 let catalog = match super::control_plane_service::load_persisted_proxy_settings_document()
387 .await
388 .map_err(ProxyControlError::from)?
389 {
390 super::control_plane_service::PersistedProxySettingsDocument::V2(cfg) => {
391 crate::config::build_persisted_provider_catalog(
392 super::control_plane_service::service_view_v2(&cfg, self.service_name),
393 )
394 }
395 super::control_plane_service::PersistedProxySettingsDocument::V4(cfg) => {
396 PersistedProvidersCatalog {
397 providers: super::control_plane_service::service_view_v4(
398 &cfg,
399 self.service_name,
400 )
401 .providers
402 .iter()
403 .map(|(name, provider)| {
404 persisted_provider_spec_from_v4_for_service(name, provider)
405 })
406 .collect(),
407 }
408 }
409 };
410 Ok(catalog)
411 }
412
413 pub async fn upsert_persisted_provider_spec(
414 &self,
415 provider_name: String,
416 provider: PersistedProviderSpec,
417 ) -> Result<(), ProxyControlError> {
418 super::persisted_registry_api::upsert_persisted_provider_spec_for_proxy(
419 self,
420 provider_name,
421 provider.into(),
422 )
423 .await
424 .map(|_| ())
425 .map_err(ProxyControlError::from)
426 }
427
428 pub async fn persisted_routing_spec(&self) -> Result<PersistedRoutingSpec, ProxyControlError> {
429 match super::control_plane_service::load_persisted_proxy_settings_document()
430 .await
431 .map_err(ProxyControlError::from)?
432 {
433 super::control_plane_service::PersistedProxySettingsDocument::V4(cfg) => {
434 let view = super::control_plane_service::service_view_v4(&cfg, self.service_name);
435 Ok(persisted_routing_spec_from_v4_for_service(view))
436 }
437 super::control_plane_service::PersistedProxySettingsDocument::V2(_) => {
438 Err(ProxyControlError::new(
439 axum::http::StatusCode::BAD_REQUEST,
440 "routing API requires a version = 5 route graph config",
441 ))
442 }
443 }
444 }
445
446 pub async fn upsert_persisted_routing_spec(
447 &self,
448 payload: PersistedRoutingUpsertRequest,
449 ) -> Result<PersistedRoutingSpec, ProxyControlError> {
450 super::persisted_registry_api::upsert_persisted_routing_spec_for_proxy(self, payload)
451 .await
452 .map_err(ProxyControlError::from)
453 }
454
455 pub async fn routing_explain(
456 &self,
457 request: RouteRequestContext,
458 session_id: Option<String>,
459 ) -> Result<RoutingExplainResponse, ProxyControlError> {
460 super::runtime_admin_api::routing_explain_for_proxy(self, request, session_id)
461 .await
462 .map_err(ProxyControlError::from)
463 }
464}
465
466fn normalize_optional_control_name(value: Option<String>) -> Option<String> {
467 value
468 .as_deref()
469 .map(str::trim)
470 .filter(|value| !value.is_empty())
471 .map(ToOwned::to_owned)
472}
473
474fn persisted_provider_spec_from_v4_for_service(
475 name: &str,
476 provider: &crate::config::ProviderConfigV4,
477) -> PersistedProviderSpec {
478 let mut endpoints = Vec::new();
479 if let Some(base_url) = provider
480 .base_url
481 .as_deref()
482 .map(str::trim)
483 .filter(|value| !value.is_empty())
484 {
485 endpoints.push(crate::config::PersistedProviderEndpointSpec {
486 name: "default".to_string(),
487 base_url: base_url.to_string(),
488 enabled: provider.enabled,
489 priority: 0,
490 tags: std::collections::BTreeMap::new(),
491 limits: crate::config::ProviderConcurrencyLimits::default(),
492 });
493 }
494 endpoints.extend(provider.endpoints.iter().map(|(endpoint_name, endpoint)| {
495 crate::config::PersistedProviderEndpointSpec {
496 name: endpoint_name.clone(),
497 base_url: endpoint.base_url.clone(),
498 enabled: endpoint.enabled,
499 priority: endpoint.priority,
500 tags: endpoint.tags.clone(),
501 limits: endpoint.limits.clone(),
502 }
503 }));
504
505 PersistedProviderSpec {
506 name: name.to_string(),
507 alias: provider.alias.clone(),
508 enabled: provider.enabled,
509 auth_token_env: provider.inline_auth.auth_token_env.clone(),
510 api_key_env: provider.inline_auth.api_key_env.clone(),
511 tags: provider.tags.clone(),
512 limits: provider.limits.clone(),
513 endpoints,
514 }
515}
516
517fn persisted_routing_spec_from_v4_for_service(
518 view: &crate::config::ServiceViewV4,
519) -> PersistedRoutingSpec {
520 let routing = crate::config::effective_v4_routing(view);
521 let order = crate::config::resolved_v4_provider_order("persisted-routing", view)
522 .unwrap_or_else(|_| view.providers.keys().cloned().collect::<Vec<_>>());
523 let entry_node = routing.entry_node();
524 PersistedRoutingSpec {
525 entry: routing.entry.clone(),
526 affinity_policy: routing.affinity_policy,
527 fallback_ttl_ms: routing.fallback_ttl_ms,
528 reprobe_preferred_after_ms: routing.reprobe_preferred_after_ms,
529 routes: routing.routes.clone(),
530 policy: entry_node
531 .map(|node| node.strategy)
532 .unwrap_or(crate::config::RoutingPolicyV4::OrderedFailover),
533 order: order.clone(),
534 target: entry_node.and_then(|node| node.target.clone()),
535 prefer_tags: entry_node
536 .map(|node| node.prefer_tags.clone())
537 .unwrap_or_default(),
538 on_exhausted: entry_node
539 .map(|node| node.on_exhausted)
540 .unwrap_or(crate::config::RoutingExhaustedActionV4::Continue),
541 entry_strategy: entry_node
542 .map(|node| node.strategy)
543 .unwrap_or(crate::config::RoutingPolicyV4::OrderedFailover),
544 expanded_order: order,
545 entry_target: entry_node.and_then(|node| node.target.clone()),
546 providers: view
547 .providers
548 .iter()
549 .map(
550 |(name, provider)| crate::config::PersistedRoutingProviderRef {
551 name: name.clone(),
552 alias: provider.alias.clone(),
553 enabled: provider.enabled,
554 tags: provider.tags.clone(),
555 },
556 )
557 .collect(),
558 }
559}