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