use super::audit;
use super::config_patch;
use crate::modules::system::logger;
use crate::proxy::admin::ErrorResponse;
use crate::proxy::state::AdminState;
use axum::{
extract::{Json, Path, State},
http::{HeaderMap, StatusCode},
response::IntoResponse,
};
use serde::Deserialize;
use std::collections::{BTreeSet, HashMap};
use std::sync::atomic::Ordering;
fn parse_env_flag(value: &str) -> bool {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on"
)
}
fn admin_stop_shutdown_hook_enabled() -> bool {
std::env::var("ADMIN_STOP_SHUTDOWN")
.ok()
.map(|value| parse_env_flag(&value))
.unwrap_or(false)
}
pub(crate) async fn admin_get_proxy_status(
State(state): State<AdminState>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let active_accounts = state.core.token_manager.len();
let is_running = { *state.runtime.is_running.read().await };
Ok(Json(serde_json::json!({
"running": is_running,
"port": state.runtime.port,
"base_url": format!("http://127.0.0.1:{}", state.runtime.port),
"active_accounts": active_accounts,
})))
}
pub(crate) async fn admin_get_version_routes() -> impl IntoResponse {
let routes = crate::proxy::routes::admin_version_route_capabilities();
Json(serde_json::json!({
"version": env!("CARGO_PKG_VERSION"),
"routes": routes
}))
}
pub(crate) async fn admin_start_proxy_service(
State(state): State<AdminState>,
) -> impl IntoResponse {
if let Ok(mut config) = crate::modules::system::config::load_app_config() {
config.proxy.auto_start = true;
let _ = crate::modules::system::config::save_app_config(&config);
}
if let Err(e) = state.core.token_manager.load_accounts().await {
logger::log_error(&format!(
"[API] Failed to enable service and load accounts: {}",
e
));
}
let mut running = state.runtime.is_running.write().await;
*running = true;
logger::log_info("[API] Proxy service enabled (Persistence synced)");
StatusCode::OK
}
pub(crate) async fn admin_stop_proxy_service(State(state): State<AdminState>) -> impl IntoResponse {
if let Ok(mut config) = crate::modules::system::config::load_app_config() {
config.proxy.auto_start = false;
let _ = crate::modules::system::config::save_app_config(&config);
}
let mut running = state.runtime.is_running.write().await;
*running = false;
logger::log_info("[API] Proxy service disabled (Axum mode / Persistence synced)");
if admin_stop_shutdown_hook_enabled() {
if crate::proxy::server::request_global_shutdown() {
logger::log_warn(
"[API] ADMIN_STOP_SHUTDOWN enabled: requested graceful server shutdown",
);
} else {
logger::log_warn(
"[API] ADMIN_STOP_SHUTDOWN enabled but no active server shutdown hook found",
);
}
}
StatusCode::OK
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct UpdateMappingWrapper {
config: crate::proxy::config::ProxyConfig,
}
fn model_mapping_change_details(
before: &HashMap<String, String>,
after: &HashMap<String, String>,
) -> serde_json::Value {
let before_keys: BTreeSet<String> = before.keys().cloned().collect();
let after_keys: BTreeSet<String> = after.keys().cloned().collect();
let mut added: Vec<String> = after_keys.difference(&before_keys).cloned().collect();
let mut removed: Vec<String> = before_keys.difference(&after_keys).cloned().collect();
let mut updated: Vec<String> = after_keys
.intersection(&before_keys)
.filter_map(|key| {
let before_value = before.get(key)?;
let after_value = after.get(key)?;
if before_value != after_value {
Some(key.clone())
} else {
None
}
})
.collect();
added.sort();
removed.sort();
updated.sort();
serde_json::json!({
"before_entries": before.len(),
"after_entries": after.len(),
"delta": {
"added_count": added.len(),
"removed_count": removed.len(),
"updated_count": updated.len(),
"added_models": added,
"removed_models": removed,
"updated_models": updated
}
})
}
pub(crate) async fn admin_update_model_mapping(
State(state): State<AdminState>,
headers: HeaderMap,
Json(payload): Json<UpdateMappingWrapper>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let actor = audit::resolve_admin_actor(&state, &headers).await;
let config = payload.config;
{
let mut mapping = state.config.custom_mapping.write().await;
*mapping = config.custom_mapping.clone();
}
let mut app_config = crate::modules::system::config::load_app_config().map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse { error: e }),
)
})?;
let before_mapping = app_config.proxy.custom_mapping.clone();
app_config.proxy.custom_mapping = config.custom_mapping;
crate::modules::system::config::save_app_config(&app_config).map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse { error: e }),
)
})?;
logger::log_info("[API] Model mapping hot-reloaded via API and saved");
audit::log_admin_audit(
"update_model_mapping",
&actor,
model_mapping_change_details(&before_mapping, &app_config.proxy.custom_mapping),
);
Ok(StatusCode::OK)
}
pub(crate) async fn admin_generate_api_key() -> impl IntoResponse {
let new_key = format!("sk-{}", uuid::Uuid::new_v4().to_string().replace("-", ""));
Json(new_key)
}
pub(crate) async fn admin_clear_proxy_session_bindings(
State(state): State<AdminState>,
) -> impl IntoResponse {
state.core.token_manager.clear_all_sessions();
logger::log_info("[API] All session bindings cleared");
StatusCode::OK
}
pub(crate) async fn admin_get_proxy_session_bindings(
State(state): State<AdminState>,
) -> impl IntoResponse {
Json(state.core.token_manager.get_sticky_debug_snapshot())
}
pub(crate) async fn admin_get_proxy_sticky_config(
State(state): State<AdminState>,
) -> impl IntoResponse {
let sticky = state.core.token_manager.get_sticky_debug_snapshot();
let preferred_account_id = state.core.token_manager.get_preferred_account().await;
Json(serde_json::json!({
"persist_session_bindings": sticky.persist_session_bindings,
"scheduling": sticky.scheduling,
"preferred_account_id": preferred_account_id
}))
}
#[cfg(test)]
mod tests {
use super::parse_env_flag;
#[test]
fn parse_env_flag_supports_expected_truthy_values() {
assert!(parse_env_flag("true"));
assert!(parse_env_flag("TRUE"));
assert!(parse_env_flag("1"));
assert!(parse_env_flag(" yes "));
assert!(parse_env_flag("On"));
}
#[test]
fn parse_env_flag_rejects_non_truthy_values() {
assert!(!parse_env_flag("false"));
assert!(!parse_env_flag("0"));
assert!(!parse_env_flag("off"));
assert!(!parse_env_flag("no"));
assert!(!parse_env_flag("maybe"));
assert!(!parse_env_flag(""));
}
}
#[derive(Deserialize)]
#[serde(rename_all = "snake_case")]
pub(crate) struct UpdateStickyConfigRequest {
#[serde(default, alias = "persistSessionBindings")]
persist_session_bindings: Option<bool>,
#[serde(default)]
scheduling: Option<crate::proxy::sticky_config::StickySessionConfig>,
}
pub(crate) async fn admin_update_proxy_sticky_config(
State(state): State<AdminState>,
headers: HeaderMap,
Json(payload): Json<UpdateStickyConfigRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let patch = config_patch::patch_proxy_config(
&state,
&headers,
config_patch::RuntimeApplyPolicy::AlwaysHotApplied,
|proxy| {
if payload.persist_session_bindings.is_none() && payload.scheduling.is_none() {
return Err(
"At least one of persist_session_bindings or scheduling must be provided"
.to_string(),
);
}
if let Some(enabled) = payload.persist_session_bindings {
proxy.persist_session_bindings = enabled;
}
if let Some(scheduling) = payload.scheduling.clone() {
proxy.scheduling = scheduling;
}
Ok(())
},
)
.await?;
state
.core
.token_manager
.update_session_binding_persistence(patch.after.persist_session_bindings);
state
.core
.token_manager
.update_sticky_config(patch.after.scheduling.clone())
.await;
let sticky = state.core.token_manager.get_sticky_debug_snapshot();
logger::log_info("[API] Sticky config updated via API and saved");
audit::log_admin_audit(
"update_proxy_sticky",
&patch.actor,
serde_json::json!({
"before": {
"persist_session_bindings": patch.before.persist_session_bindings,
"scheduling": patch.before.scheduling
},
"after": {
"persist_session_bindings": sticky.persist_session_bindings,
"scheduling": sticky.scheduling
}
}),
);
Ok(Json(serde_json::json!({
"ok": true,
"saved": true,
"message": "Sticky config updated",
"runtime_apply": patch.runtime_apply_result(true),
"sticky": {
"persist_session_bindings": sticky.persist_session_bindings,
"scheduling": sticky.scheduling
}
})))
}
pub(crate) async fn admin_get_proxy_request_timeout(
State(state): State<AdminState>,
) -> impl IntoResponse {
let timeout = state.config.request_timeout_secs();
Json(serde_json::json!({
"request_timeout": timeout,
"effective_request_timeout": timeout.max(5)
}))
}
#[derive(Deserialize)]
#[serde(rename_all = "snake_case")]
pub(crate) struct UpdateRequestTimeoutRequest {
#[serde(default, alias = "requestTimeout", alias = "requestTimeoutSeconds")]
request_timeout: Option<u64>,
#[serde(default, alias = "request_timeout_seconds")]
request_timeout_seconds: Option<u64>,
}
pub(crate) async fn admin_update_proxy_request_timeout(
State(state): State<AdminState>,
headers: HeaderMap,
Json(payload): Json<UpdateRequestTimeoutRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let patch = config_patch::patch_proxy_config(
&state,
&headers,
config_patch::RuntimeApplyPolicy::AlwaysHotApplied,
|proxy| {
let request_timeout = payload
.request_timeout
.or(payload.request_timeout_seconds)
.ok_or_else(|| "request_timeout is required".to_string())?;
proxy.request_timeout = request_timeout;
Ok(())
},
)
.await?;
let request_timeout = patch.after.request_timeout;
state
.config
.request_timeout
.store(request_timeout, Ordering::Relaxed);
logger::log_info("[API] Request timeout updated via API and saved");
audit::log_admin_audit(
"update_proxy_request_timeout",
&patch.actor,
serde_json::json!({
"before": {
"request_timeout": patch.before.request_timeout,
"effective_request_timeout": patch.before.request_timeout.max(5)
},
"after": {
"request_timeout": request_timeout,
"effective_request_timeout": request_timeout.max(5)
}
}),
);
Ok(Json(serde_json::json!({
"ok": true,
"saved": true,
"message": "Request timeout updated",
"runtime_apply": patch.runtime_apply_result(true),
"request_timeout": request_timeout,
"effective_request_timeout": request_timeout.max(5)
})))
}
pub(crate) async fn admin_get_proxy_compliance_debug(
State(state): State<AdminState>,
) -> impl IntoResponse {
Json(
state
.core
.token_manager
.get_compliance_debug_snapshot()
.await,
)
}
pub(crate) async fn admin_get_google_outbound_policy(
State(state): State<AdminState>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let policy = state.core.upstream.get_google_policy().await;
let mode = match policy.mode {
crate::proxy::config::GoogleMode::PublicGoogle => "public_google",
crate::proxy::config::GoogleMode::CodeassistCompat => "codeassist_compat",
};
let cfg = crate::modules::system::config::load_app_config().ok();
let google_cfg = cfg
.as_ref()
.map(|c| c.proxy.google.clone())
.unwrap_or_default();
let mimic_profile = match google_cfg.mimic.profile.clone() {
crate::proxy::config::GoogleMimicProfile::StrictMimic => "strict_mimic",
crate::proxy::config::GoogleMimicProfile::Functional => "functional",
};
let userinfo_endpoint = match google_cfg.userinfo_endpoint {
crate::proxy::config::GoogleUserinfoEndpoint::Oauth2V2 => "oauth2_v2",
crate::proxy::config::GoogleUserinfoEndpoint::OpenidconnectV1 => "openidconnect_v1",
crate::proxy::config::GoogleUserinfoEndpoint::DualFallback => "dual_fallback",
};
let cloudcode_host_strategy =
crate::proxy::google::endpoints::cloudcode_host_strategy(google_cfg.mimic.profile.clone());
Ok(Json(serde_json::json!({
"mode": mode,
"inputs": {
"google_source": "runtime_state",
"debug_logging_source": "runtime_state"
},
"headers": {
"send_host_header_configured": policy.send_host_header,
"send_host_header_effective": policy.should_send_host_header(),
"send_x_goog_api_client_configured": policy.send_x_goog_api_client,
"send_x_goog_api_client_effective": policy.send_x_goog_api_client_effective(),
"send_x_goog_api_client_on_cloudcode": policy.send_x_goog_api_client_on_cloudcode,
"x_goog_api_client": policy.x_goog_api_client.clone(),
"x_goog_api_client_ua_guard": "antigravity/* or google-api-nodejs-client/*",
"always_set": [
"authorization",
"user-agent",
"accept-encoding"
],
"conditionally_set": [
"x-goog-api-client"
],
"json_request_header": {
"content-type": "application/json"
},
"passthrough_policy": "deny_by_default",
"allowed_passthrough_headers": [
"anthropic-beta"
],
"blocked_categories": [
"sec-*",
"origin",
"referer",
"cookie",
"x-forwarded-*",
"x-real-ip",
"hop-by-hop",
"host"
]
},
"identity_metadata": {
"ide_type": policy.identity_metadata.ide_type,
"platform": policy.identity_metadata.platform,
"plugin_type": policy.identity_metadata.plugin_type
},
"debug": {
"log_google_outbound_headers": policy.log_google_outbound_headers,
"redaction_applies_to": [
"authorization",
"*token*",
"*api-key*",
"cookie"
]
},
"mimic": {
"profile": mimic_profile,
"trigger_on_auth_events": google_cfg.mimic.trigger_on_auth_events,
"cooldown_seconds": google_cfg.mimic.cooldown_seconds,
"userinfo_endpoint": userinfo_endpoint,
"cloudcode_host_strategy": cloudcode_host_strategy
}
})))
}
pub(crate) async fn admin_get_proxy_metrics(State(state): State<AdminState>) -> impl IntoResponse {
let monitor_stats = state.core.monitor.get_stats().await;
let compliance = state
.core
.token_manager
.get_compliance_debug_snapshot()
.await;
let sticky = state.core.token_manager.get_sticky_debug_snapshot();
let proxy_pool_cfg = { state.runtime.proxy_pool_state.read().await.clone() };
let proxy_pool_observability = state
.runtime
.proxy_pool_manager
.get_observability_snapshot();
let proxy_pool_bindings_count = state
.runtime
.proxy_pool_manager
.get_all_bindings_snapshot()
.len();
let active_accounts = state.core.token_manager.len();
let request_timeout = state.config.request_timeout_secs();
let running = *state.runtime.is_running.read().await;
let refresh_observability = crate::modules::auth::oauth::refresh_observability_snapshot();
let scheduler_observability =
crate::modules::system::scheduler::scheduler_refresh_observability_snapshot();
let total_account_requests_in_last_minute: usize = compliance
.account_requests_in_last_minute
.values()
.copied()
.sum();
let total_account_in_flight: usize = compliance.account_in_flight.values().copied().sum();
Json(serde_json::json!({
"timestamp_unix": chrono::Utc::now().timestamp(),
"runtime": {
"running": running,
"port": state.runtime.port,
"active_accounts": active_accounts,
"request_timeout": request_timeout,
"effective_request_timeout": request_timeout.max(5),
"tls_backend": crate::utils::http::tls_backend_name(),
"tls_requested_backend": crate::utils::http::tls_requested_backend_name(),
"tls_compiled_backends": crate::utils::http::tls_compiled_backends(),
"tls_canary": crate::utils::http::tls_canary_snapshot(),
},
"monitor": {
"enabled": state.core.monitor.is_enabled(),
"total_requests": monitor_stats.total_requests,
"success_count": monitor_stats.success_count,
"error_count": monitor_stats.error_count,
},
"sticky": {
"persist_session_bindings": sticky.persist_session_bindings,
"scheduling_mode": sticky.scheduling.mode,
"scheduling_max_wait_seconds": sticky.scheduling.max_wait_seconds,
"session_bindings_count": sticky.session_bindings.len(),
"recent_events_count": sticky.recent_events.len(),
},
"proxy_pool": {
"enabled": proxy_pool_cfg.enabled,
"auto_failover": proxy_pool_cfg.auto_failover,
"allow_shared_proxy_fallback": proxy_pool_cfg.allow_shared_proxy_fallback,
"require_proxy_for_account_requests": proxy_pool_cfg.require_proxy_for_account_requests,
"strategy": proxy_pool_cfg.strategy,
"configured_proxies": proxy_pool_cfg.proxies.len(),
"account_bindings_count": proxy_pool_bindings_count,
"shared_fallback_selections_total": proxy_pool_observability.shared_fallback_selections_total,
"strict_rejections_total": proxy_pool_observability.strict_rejections_total,
},
"compliance": {
"enabled": compliance.config.enabled,
"max_global_requests_per_minute": compliance.config.max_global_requests_per_minute,
"max_account_requests_per_minute": compliance.config.max_account_requests_per_minute,
"max_account_concurrency": compliance.config.max_account_concurrency,
"risk_cooldown_seconds": compliance.config.risk_cooldown_seconds,
"max_retry_attempts": compliance.config.max_retry_attempts,
"global_requests_in_last_minute": compliance.global_requests_in_last_minute,
"tracked_accounts_last_minute": compliance.account_requests_in_last_minute.len(),
"total_account_requests_in_last_minute": total_account_requests_in_last_minute,
"tracked_accounts_in_flight": compliance.account_in_flight.len(),
"total_account_in_flight": total_account_in_flight,
"accounts_in_cooldown": compliance.account_cooldown_seconds_remaining.len(),
"risk_signals_last_minute": compliance.risk_signals_last_minute,
"account_switches_last_minute": compliance.account_switches_last_minute,
"accounts_with_403_in_last_minute": compliance.account_403_in_last_minute.len(),
"accounts_with_429_in_last_minute": compliance.account_429_in_last_minute.len(),
"account_403_in_last_minute": compliance.account_403_in_last_minute,
"account_429_in_last_minute": compliance.account_429_in_last_minute,
"refresh_attempts_last_minute": refresh_observability.refresh_attempts_last_minute,
"refresh_attempts_by_account_last_minute": refresh_observability.refresh_attempts_by_account_last_minute,
"scheduler_refresh_runs_last_minute": scheduler_observability.scheduler_refresh_runs_last_minute,
"scheduler_refresh_failures_last_minute": scheduler_observability.scheduler_refresh_failures_last_minute,
"scheduler_refresh_accounts_attempted_last_minute": scheduler_observability.scheduler_refresh_accounts_attempted_last_minute,
},
"runtime_apply_policies_supported": config_patch::supported_runtime_apply_policies()
}))
}
pub(crate) async fn admin_run_tls_canary_probe() -> impl IntoResponse {
match crate::utils::http::run_tls_startup_canary_probe().await {
Ok(()) => Json(serde_json::json!({
"ok": true,
"tls_canary": crate::utils::http::tls_canary_snapshot()
}))
.into_response(),
Err(error) => (
StatusCode::BAD_GATEWAY,
Json(serde_json::json!({
"ok": false,
"error": error,
"tls_canary": crate::utils::http::tls_canary_snapshot()
})),
)
.into_response(),
}
}
pub(crate) async fn admin_get_tls_canary_status() -> impl IntoResponse {
Json(crate::utils::http::tls_canary_snapshot())
}
pub(crate) async fn admin_update_proxy_compliance(
State(state): State<AdminState>,
headers: HeaderMap,
Json(compliance): Json<crate::proxy::config::ComplianceConfig>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let patch = config_patch::patch_proxy_config(
&state,
&headers,
config_patch::RuntimeApplyPolicy::AlwaysHotApplied,
|proxy| {
proxy.compliance = compliance.clone();
Ok(())
},
)
.await?;
state
.core
.token_manager
.update_compliance_config(patch.after.compliance.clone())
.await;
logger::log_info("[API] Compliance config updated via API and saved");
audit::log_admin_audit(
"update_proxy_compliance",
&patch.actor,
serde_json::json!({
"before": patch.before.compliance,
"after": patch.after.compliance
}),
);
Ok(Json(serde_json::json!({
"ok": true,
"saved": true,
"message": "Compliance config updated",
"runtime_apply": patch.runtime_apply_result(true),
"compliance": patch.after.compliance
})))
}
pub(crate) async fn admin_clear_all_rate_limits(
State(state): State<AdminState>,
) -> impl IntoResponse {
state.core.token_manager.clear_all_rate_limits();
logger::log_info("[API] All rate limit records cleared");
StatusCode::OK
}
pub(crate) async fn admin_clear_rate_limit(
State(state): State<AdminState>,
Path(account_id): Path<String>,
) -> impl IntoResponse {
let cleared = state.core.token_manager.clear_rate_limit(&account_id);
if cleared {
logger::log_info(&format!(
"[API] Rate limit record for account {} cleared",
account_id
));
StatusCode::OK
} else {
StatusCode::NOT_FOUND
}
}
pub(crate) async fn admin_get_preferred_account(
State(state): State<AdminState>,
) -> impl IntoResponse {
let pref = state.core.token_manager.get_preferred_account().await;
Json(pref)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct SetPreferredAccountRequest {
account_id: Option<String>,
}
pub(crate) async fn admin_set_preferred_account(
State(state): State<AdminState>,
headers: HeaderMap,
Json(payload): Json<SetPreferredAccountRequest>,
) -> impl IntoResponse {
let actor = audit::resolve_admin_actor(&state, &headers).await;
let before = state.core.token_manager.get_preferred_account().await;
state
.core
.token_manager
.set_preferred_account(payload.account_id.clone())
.await;
let after = state.core.token_manager.get_preferred_account().await;
audit::log_admin_audit(
"set_preferred_account",
&actor,
serde_json::json!({
"before": { "preferred_account_id": before },
"after": { "preferred_account_id": after }
}),
);
StatusCode::OK
}
pub(crate) async fn admin_fetch_zai_models(
Json(payload): Json<serde_json::Value>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
let zai_config = payload.get("zai").ok_or_else(|| {
(
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
error: "Missing zai config".to_string(),
}),
)
})?;
let api_key = zai_config
.get("api_key")
.and_then(|v| v.as_str())
.unwrap_or("");
let base_url = zai_config
.get("base_url")
.and_then(|v| v.as_str())
.unwrap_or("https://api.z.ai");
let client = crate::utils::http::get_client();
let resp = client
.get(format!("{}/v1/models", base_url))
.header("Authorization", format!("Bearer {}", api_key))
.send()
.await
.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse {
error: e.to_string(),
}),
)
})?;
let data: serde_json::Value = resp.json().await.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse {
error: e.to_string(),
}),
)
})?;
let models = data
.get("data")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|m| {
m.get("id")
.and_then(|id| id.as_str().map(|s| s.to_string()))
})
.collect::<Vec<String>>()
})
.unwrap_or_default();
Ok(Json(models))
}
pub(crate) async fn admin_set_proxy_monitor_enabled(
State(state): State<AdminState>,
Json(payload): Json<serde_json::Value>,
) -> impl IntoResponse {
let enabled = payload
.get("enabled")
.and_then(|v| v.as_bool())
.unwrap_or(false);
if state.core.monitor.is_enabled() != enabled {
state.core.monitor.set_enabled(enabled);
logger::log_info(&format!("[API] Monitor status set to: {}", enabled));
}
StatusCode::OK
}