diff --git a/crates/aisix-admin/src/lib.rs b/crates/aisix-admin/src/lib.rs index 4ff36f83..a6c0a0c9 100644 --- a/crates/aisix-admin/src/lib.rs +++ b/crates/aisix-admin/src/lib.rs @@ -63,9 +63,8 @@ pub fn build_router(state: AdminState) -> Router { // subsequent handler call is a free lookup. let _ = openapi::merged_openapi(); - Router::new() + let mut router = Router::new() .route("/livez", get(livez)) - .route("/metrics", get(metrics_handler)) // OpenAPI scalar UI is unauthenticated like /metrics — admin // listener is private in production. .route("/admin/openapi.json", get(openapi::openapi_json)) @@ -150,8 +149,28 @@ pub fn build_router(state: AdminState) -> Router { .route( "/playground/chat/completions", post(playground_handler::playground_chat_completions), - ) - .with_state(state) + ); + + if state.prometheus.enabled { + router = router.route( + &normalized_prometheus_path(&state.prometheus.path), + get(metrics_handler), + ); + } + + router.with_state(state) +} + +fn normalized_prometheus_path(path: &str) -> String { + let path = path.trim(); + if path.is_empty() { + return "/metrics".to_string(); + } + if path.starts_with('/') { + path.to_string() + } else { + format!("/{path}") + } } async fn livez( @@ -335,6 +354,84 @@ mod tests { assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); } + #[tokio::test] + async fn metrics_endpoint_uses_configured_path() { + use aisix_core::config::PrometheusConfig; + use aisix_obs::Metrics; + + let state = build_state() + .with_metrics(Arc::new(Metrics::new(false))) + .with_prometheus_config(PrometheusConfig { + enabled: true, + path: "/internal/prom".into(), + }); + let app = build_router(state); + + let resp = run( + app.clone(), + Request::builder() + .uri("/metrics") + .body(Body::empty()) + .unwrap(), + ) + .await; + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + let resp = run( + app, + Request::builder() + .uri("/internal/prom") + .body(Body::empty()) + .unwrap(), + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + } + + #[tokio::test] + async fn metrics_endpoint_normalizes_configured_path() { + use aisix_core::config::PrometheusConfig; + use aisix_obs::Metrics; + + let state = build_state() + .with_metrics(Arc::new(Metrics::new(false))) + .with_prometheus_config(PrometheusConfig { + enabled: true, + path: "internal/prom".into(), + }); + let app = build_router(state); + + let resp = run( + app, + Request::builder() + .uri("/internal/prom") + .body(Body::empty()) + .unwrap(), + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + } + + #[tokio::test] + async fn metrics_endpoint_can_be_disabled() { + use aisix_core::config::PrometheusConfig; + use aisix_obs::Metrics; + + let state = build_state() + .with_metrics(Arc::new(Metrics::new(false))) + .with_prometheus_config(PrometheusConfig { + enabled: false, + path: "/metrics".into(), + }); + let app = build_router(state); + let req = Request::builder() + .uri("/metrics") + .body(Body::empty()) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + } + #[tokio::test] async fn livez_reports_plain_ok_by_default() { let app = build_router(build_state()); diff --git a/crates/aisix-admin/src/state.rs b/crates/aisix-admin/src/state.rs index c24d7ef6..0e199409 100644 --- a/crates/aisix-admin/src/state.rs +++ b/crates/aisix-admin/src/state.rs @@ -11,6 +11,7 @@ //! wire an etcd-backed impl and tests can use `InMemoryStore` via the //! same type. +use aisix_core::config::PrometheusConfig; use aisix_core::snapshot::SnapshotHandle; use aisix_core::{AdminConfig, AisixSnapshot}; use aisix_etcd::WatchStatus; @@ -27,6 +28,7 @@ pub struct AdminState { pub admin_keys: Arc<[String]>, pub store: Arc, pub metrics: Option>, + pub prometheus: PrometheusConfig, /// Shared in-process health tracker from the proxy. Used by the /// `/admin/v1/health` endpoint to report per-model health status. pub health_tracker: Option>, @@ -59,6 +61,7 @@ impl AdminState { admin_keys: Arc::from(cfg.admin_keys.clone()), store, metrics: None, + prometheus: PrometheusConfig::default(), health_tracker: None, runtime_status_tracker: None, watch_status: None, @@ -83,6 +86,11 @@ impl AdminState { self } + pub fn with_prometheus_config(mut self, prometheus: PrometheusConfig) -> Self { + self.prometheus = prometheus; + self + } + /// Attach the in-process health tracker from the proxy. When set, /// `GET /admin/v1/health` reflects per-model upstream health. pub fn with_health_tracker(mut self, tracker: Arc) -> Self { diff --git a/crates/aisix-obs/src/lib.rs b/crates/aisix-obs/src/lib.rs index dffb9d46..3590d67e 100644 --- a/crates/aisix-obs/src/lib.rs +++ b/crates/aisix-obs/src/lib.rs @@ -22,7 +22,10 @@ use aisix_core::ObservabilityConfig; use tracing_subscriber::{fmt, prelude::*, EnvFilter}; pub use access_log::AccessLog; -pub use metrics::{Metrics, RequestOutcome}; +pub use metrics::{ + BudgetGauges, BudgetLabels, DeploymentLabels, DeploymentState, LlmUsage, Metrics, + RequestLabels, RequestOutcome, UsageLabels, +}; pub use otlp::{install_otlp_tracer, shutdown_otlp, OtlpError, OtlpHandle}; pub use otlp_http_sink::OtlpHttpFanOut; pub use usage::{UsageEvent, UsageSink}; diff --git a/crates/aisix-obs/src/metrics.rs b/crates/aisix-obs/src/metrics.rs index a43ae109..23efb1dd 100644 --- a/crates/aisix-obs/src/metrics.rs +++ b/crates/aisix-obs/src/metrics.rs @@ -1,7 +1,7 @@ //! Prometheus metrics registry shared across the proxy middleware and //! the admin `/metrics` endpoint. //! -//! Four series cover spec §7: +//! Existing compatibility series cover spec §7: //! - `aisix_requests_total{provider,model,status,outcome}` — counter //! incremented at the end of every proxy request. //! - `aisix_request_duration_seconds{provider,model,status}` — histogram @@ -10,13 +10,19 @@ //! - `aisix_tokens_consumed_total{provider,model}` — counter of //! `usage.total_tokens` summed across completed non-streaming calls. //! +//! Newer AISIX-native series use `aisix_proxy_*` and `aisix_llm_*` +//! names with bounded, DP-stable labels. They intentionally do not +//! copy LiteLLM label names that the data plane does not have. +//! //! A single [`Metrics`] instance is held `Arc`'d inside `ObsState` and //! cloned into axum state. The exposition format is emitted via //! `metrics-exporter-prometheus`'s text renderer; no global recorder is //! installed, so tests can spin up isolated instances per case. use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle, PrometheusRecorder}; +use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; /// Metric names (public so the admin `/metrics` handler and tests can @@ -25,6 +31,36 @@ pub const M_REQUESTS_TOTAL: &str = "aisix_requests_total"; pub const M_REQUEST_DURATION: &str = "aisix_request_duration_seconds"; pub const M_RATELIMIT_REJECTIONS: &str = "aisix_ratelimit_rejections_total"; pub const M_TOKENS_CONSUMED: &str = "aisix_tokens_consumed_total"; +pub const M_LLM_SPEND_MICRO_USD_TOTAL: &str = "aisix_llm_spend_micro_usd_total"; +pub const M_LLM_INPUT_TOKENS_TOTAL: &str = "aisix_llm_input_tokens_total"; +pub const M_LLM_OUTPUT_TOKENS_TOTAL: &str = "aisix_llm_output_tokens_total"; +pub const M_LLM_TOTAL_TOKENS_TOTAL: &str = "aisix_llm_total_tokens_total"; +pub const M_LLM_REQUESTS_TOTAL: &str = "aisix_llm_requests_total"; +pub const M_LLM_REQUEST_DURATION: &str = "aisix_llm_request_duration_seconds"; +pub const M_LLM_API_LATENCY: &str = "aisix_llm_api_latency_seconds"; +pub const M_LLM_TTFT: &str = "aisix_llm_time_to_first_token_seconds"; +pub const M_PROXY_IN_FLIGHT: &str = "aisix_proxy_in_flight_requests"; +pub const M_PROXY_REQUESTS_TOTAL: &str = "aisix_proxy_requests_total"; +pub const M_PROXY_FAILED_REQUESTS_TOTAL: &str = "aisix_proxy_failed_requests_total"; +pub const M_PROXY_REQUEST_DURATION: &str = "aisix_proxy_request_duration_seconds"; +pub const M_DEPLOYMENT_REQUESTS_TOTAL: &str = "aisix_deployment_requests_total"; +pub const M_DEPLOYMENT_SUCCESS_TOTAL: &str = "aisix_deployment_success_responses_total"; +pub const M_DEPLOYMENT_FAILURE_TOTAL: &str = "aisix_deployment_failure_responses_total"; +pub const M_DEPLOYMENT_STATE: &str = "aisix_deployment_state"; +pub const M_DEPLOYMENT_COOLED_DOWN_TOTAL: &str = "aisix_deployment_cooled_down_total"; +pub const M_ROUTING_SUCCESSFUL_FALLBACKS_TOTAL: &str = "aisix_routing_successful_fallbacks_total"; +pub const M_ROUTING_FAILED_FALLBACKS_TOTAL: &str = "aisix_routing_failed_fallbacks_total"; +pub const M_RATELIMIT_REMAINING_REQUESTS: &str = "aisix_ratelimit_remaining_requests"; +pub const M_RATELIMIT_REMAINING_TOKENS: &str = "aisix_ratelimit_remaining_tokens"; +pub const M_BUDGET_LIMIT_USD: &str = "aisix_budget_limit_usd"; +pub const M_BUDGET_SPENT_USD: &str = "aisix_budget_spent_usd"; +pub const M_BUDGET_REMAINING_USD: &str = "aisix_budget_remaining_usd"; +pub const M_BUDGET_RESET_SECONDS: &str = "aisix_budget_reset_seconds"; +pub const M_BUDGET_DETAILS_PRESENT: &str = "aisix_budget_details_present"; +pub const M_REDIS_FAILURES_TOTAL: &str = "aisix_redis_failures_total"; +pub const M_USAGE_EVENT_DROPS_TOTAL: &str = "aisix_usage_event_drops_total"; +pub const M_OTLP_FANOUT_DROPS_TOTAL: &str = "aisix_otlp_fanout_drops_total"; +pub const M_OTLP_FANOUT_FAILURES_TOTAL: &str = "aisix_otlp_fanout_failures_total"; /// Holds an isolated `PrometheusRecorder` plus its render handle. /// `metrics::*` macros talk to whatever recorder is in scope; we use @@ -38,6 +74,7 @@ pub struct Metrics { struct MetricsInner { recorder: PrometheusRecorder, handle: PrometheusHandle, + proxy_in_flight: Mutex>, } impl std::fmt::Debug for Metrics { @@ -54,7 +91,11 @@ impl Metrics { let recorder = PrometheusBuilder::new().build_recorder(); let handle = recorder.handle(); Self { - inner: Arc::new(MetricsInner { recorder, handle }), + inner: Arc::new(MetricsInner { + recorder, + handle, + proxy_in_flight: Mutex::new(HashMap::new()), + }), } } @@ -114,6 +155,482 @@ impl Metrics { .increment(total_tokens); }); } + + pub fn increment_proxy_in_flight(&self, endpoint: &str, inbound_protocol: &str) { + let value = { + let mut counters = self.inner.proxy_in_flight.lock().expect("lock in-flight"); + let value = counters + .entry((endpoint.to_string(), inbound_protocol.to_string())) + .or_insert(0); + *value += 1; + *value + }; + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::gauge!( + M_PROXY_IN_FLIGHT, + "endpoint" => endpoint.to_string(), + "inbound_protocol" => inbound_protocol.to_string(), + ) + .set(value as f64); + }); + } + + pub fn decrement_proxy_in_flight(&self, endpoint: &str, inbound_protocol: &str) { + let value = { + let mut counters = self.inner.proxy_in_flight.lock().expect("lock in-flight"); + let key = (endpoint.to_string(), inbound_protocol.to_string()); + let value = counters.entry(key.clone()).or_insert(0); + *value = (*value - 1).max(0); + let current = *value; + if current == 0 { + counters.remove(&key); + } + current + }; + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::gauge!( + M_PROXY_IN_FLIGHT, + "endpoint" => endpoint.to_string(), + "inbound_protocol" => inbound_protocol.to_string(), + ) + .set(value as f64); + }); + } + + pub fn record_proxy_request(&self, labels: RequestLabels<'_>, duration: Duration) { + metrics::with_local_recorder(&self.inner.recorder, || { + labels.record_request_counter(M_PROXY_REQUESTS_TOTAL); + metrics::histogram!( + M_PROXY_REQUEST_DURATION, + "endpoint" => labels.endpoint.to_string(), + "inbound_protocol" => labels.inbound_protocol.to_string(), + "provider" => labels.provider.to_string(), + "model" => labels.model.to_string(), + "upstream_model" => labels.upstream_model.to_string(), + "provider_key_id" => labels.provider_key_id.to_string(), + "api_key_id" => labels.api_key_id.to_string(), + "team_id" => labels.team_id.to_string(), + "owner_id" => labels.owner_id.to_string(), + "status" => labels.status.to_string(), + "outcome" => labels.outcome.as_str().to_string(), + ) + .record(duration.as_secs_f64()); + if labels.outcome != RequestOutcome::Success { + labels.record_request_counter(M_PROXY_FAILED_REQUESTS_TOTAL); + } + }); + } + + pub fn record_llm_request(&self, labels: RequestLabels<'_>, duration: Duration) { + metrics::with_local_recorder(&self.inner.recorder, || { + labels.record_request_counter(M_LLM_REQUESTS_TOTAL); + metrics::histogram!( + M_LLM_REQUEST_DURATION, + "endpoint" => labels.endpoint.to_string(), + "inbound_protocol" => labels.inbound_protocol.to_string(), + "provider" => labels.provider.to_string(), + "model" => labels.model.to_string(), + "upstream_model" => labels.upstream_model.to_string(), + "provider_key_id" => labels.provider_key_id.to_string(), + "api_key_id" => labels.api_key_id.to_string(), + "team_id" => labels.team_id.to_string(), + "owner_id" => labels.owner_id.to_string(), + "status" => labels.status.to_string(), + "outcome" => labels.outcome.as_str().to_string(), + ) + .record(duration.as_secs_f64()); + }); + } + + pub fn record_llm_usage(&self, labels: UsageLabels<'_>, usage: LlmUsage) { + if usage.is_empty() { + return; + } + metrics::with_local_recorder(&self.inner.recorder, || { + if usage.input_tokens > 0 { + labels.record_counter(M_LLM_INPUT_TOKENS_TOTAL, u64::from(usage.input_tokens)); + } + if usage.output_tokens > 0 { + labels.record_counter(M_LLM_OUTPUT_TOKENS_TOTAL, u64::from(usage.output_tokens)); + } + if usage.total_tokens > 0 { + labels.record_counter(M_LLM_TOTAL_TOKENS_TOTAL, u64::from(usage.total_tokens)); + } + if usage.spend_usd > 0.0 { + labels.record_spend_usd(usage.spend_usd); + } + }); + } + + pub fn record_time_to_first_token(&self, labels: UsageLabels<'_>, ttft: Duration) { + if ttft.is_zero() { + return; + } + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::histogram!( + M_LLM_TTFT, + "endpoint" => labels.endpoint.to_string(), + "inbound_protocol" => labels.inbound_protocol.to_string(), + "provider" => labels.provider.to_string(), + "model" => labels.model.to_string(), + "upstream_model" => labels.upstream_model.to_string(), + "provider_key_id" => labels.provider_key_id.to_string(), + "api_key_id" => labels.api_key_id.to_string(), + "team_id" => labels.team_id.to_string(), + "owner_id" => labels.owner_id.to_string(), + ) + .record(ttft.as_secs_f64()); + }); + } + + pub fn record_deployment_request(&self, labels: DeploymentLabels<'_>, outcome: RequestOutcome) { + metrics::with_local_recorder(&self.inner.recorder, || { + labels.record_counter(M_DEPLOYMENT_REQUESTS_TOTAL); + match outcome { + RequestOutcome::Success => labels.record_counter(M_DEPLOYMENT_SUCCESS_TOTAL), + _ => labels.record_counter(M_DEPLOYMENT_FAILURE_TOTAL), + } + }); + } + + pub fn set_deployment_state(&self, labels: DeploymentLabels<'_>, state: DeploymentState) { + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::gauge!( + M_DEPLOYMENT_STATE, + "provider" => labels.provider.to_string(), + "model" => labels.model.to_string(), + "upstream_model" => labels.upstream_model.to_string(), + "provider_key_id" => labels.provider_key_id.to_string(), + ) + .set(state.as_f64()); + }); + } + + pub fn record_deployment_cooldown(&self, labels: DeploymentLabels<'_>) { + metrics::with_local_recorder(&self.inner.recorder, || { + labels.record_counter(M_DEPLOYMENT_COOLED_DOWN_TOTAL); + }); + } + + pub fn record_routing_fallback(&self, success: bool, model: &str) { + let metric = if success { + M_ROUTING_SUCCESSFUL_FALLBACKS_TOTAL + } else { + M_ROUTING_FAILED_FALLBACKS_TOTAL + }; + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::counter!(metric, "model" => model.to_string()).increment(1); + }); + } + + pub fn set_rate_limit_remaining( + &self, + api_key_id: &str, + model: &str, + requests: Option, + tokens: Option, + ) { + metrics::with_local_recorder(&self.inner.recorder, || { + if let Some(value) = requests { + metrics::gauge!( + M_RATELIMIT_REMAINING_REQUESTS, + "api_key_id" => api_key_id.to_string(), + "model" => model.to_string(), + ) + .set(value as f64); + } + if let Some(value) = tokens { + metrics::gauge!( + M_RATELIMIT_REMAINING_TOKENS, + "api_key_id" => api_key_id.to_string(), + "model" => model.to_string(), + ) + .set(value as f64); + } + }); + } + + pub fn set_budget_gauges(&self, labels: BudgetLabels<'_>, budget: BudgetGauges) { + metrics::with_local_recorder(&self.inner.recorder, || { + labels.record_gauge(M_BUDGET_DETAILS_PRESENT, 1.0); + if let Some(value) = budget.limit_usd { + labels.record_gauge(M_BUDGET_LIMIT_USD, value); + } + if let Some(value) = budget.spent_usd { + labels.record_gauge(M_BUDGET_SPENT_USD, value); + } + if let Some(value) = budget.remaining_usd { + labels.record_gauge(M_BUDGET_REMAINING_USD, value); + } + if let Some(value) = budget.reset_seconds { + labels.record_gauge(M_BUDGET_RESET_SECONDS, value as f64); + } + }); + } + + pub fn clear_budget_gauges(&self, labels: BudgetLabels<'_>) { + metrics::with_local_recorder(&self.inner.recorder, || { + labels.record_gauge(M_BUDGET_DETAILS_PRESENT, 0.0); + }); + } + + pub fn record_redis_failure(&self, operation: &str) { + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::counter!(M_REDIS_FAILURES_TOTAL, "operation" => operation.to_string()) + .increment(1); + }); + } + + pub fn record_usage_event_drop(&self, reason: &str) { + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::counter!(M_USAGE_EVENT_DROPS_TOTAL, "reason" => reason.to_string()) + .increment(1); + }); + } + + pub fn record_otlp_fanout_drop(&self, exporter: &str, reason: &str) { + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::counter!( + M_OTLP_FANOUT_DROPS_TOTAL, + "exporter" => exporter.to_string(), + "reason" => reason.to_string(), + ) + .increment(1); + }); + } + + pub fn record_otlp_fanout_failure(&self, exporter: &str) { + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::counter!(M_OTLP_FANOUT_FAILURES_TOTAL, "exporter" => exporter.to_string()) + .increment(1); + }); + } +} + +#[derive(Debug, Clone, Copy)] +pub struct RequestLabels<'a> { + pub endpoint: &'a str, + pub inbound_protocol: &'a str, + pub provider: &'a str, + pub model: &'a str, + pub upstream_model: &'a str, + pub provider_key_id: &'a str, + pub api_key_id: &'a str, + pub team_id: &'a str, + pub owner_id: &'a str, + pub status: u16, + pub outcome: RequestOutcome, +} + +impl Default for RequestLabels<'_> { + fn default() -> Self { + Self { + endpoint: "unknown", + inbound_protocol: "openai", + provider: "unknown", + model: "unknown", + upstream_model: "unknown", + provider_key_id: "unknown", + api_key_id: "unknown", + team_id: "unknown", + owner_id: "unknown", + status: 0, + outcome: RequestOutcome::UpstreamError, + } + } +} + +impl RequestLabels<'_> { + fn record_request_counter(&self, metric: &'static str) { + metrics::counter!( + metric, + "endpoint" => self.endpoint.to_string(), + "inbound_protocol" => self.inbound_protocol.to_string(), + "provider" => self.provider.to_string(), + "model" => self.model.to_string(), + "upstream_model" => self.upstream_model.to_string(), + "provider_key_id" => self.provider_key_id.to_string(), + "api_key_id" => self.api_key_id.to_string(), + "team_id" => self.team_id.to_string(), + "owner_id" => self.owner_id.to_string(), + "status" => self.status.to_string(), + "outcome" => self.outcome.as_str().to_string(), + ) + .increment(1); + } +} + +#[derive(Debug, Clone, Copy)] +pub struct UsageLabels<'a> { + pub endpoint: &'a str, + pub inbound_protocol: &'a str, + pub provider: &'a str, + pub model: &'a str, + pub upstream_model: &'a str, + pub provider_key_id: &'a str, + pub api_key_id: &'a str, + pub team_id: &'a str, + pub owner_id: &'a str, +} + +impl Default for UsageLabels<'_> { + fn default() -> Self { + Self { + endpoint: "unknown", + inbound_protocol: "openai", + provider: "unknown", + model: "unknown", + upstream_model: "unknown", + provider_key_id: "unknown", + api_key_id: "unknown", + team_id: "unknown", + owner_id: "unknown", + } + } +} + +impl UsageLabels<'_> { + fn record_counter(&self, metric: &'static str, value: u64) { + metrics::counter!( + metric, + "endpoint" => self.endpoint.to_string(), + "inbound_protocol" => self.inbound_protocol.to_string(), + "provider" => self.provider.to_string(), + "model" => self.model.to_string(), + "upstream_model" => self.upstream_model.to_string(), + "provider_key_id" => self.provider_key_id.to_string(), + "api_key_id" => self.api_key_id.to_string(), + "team_id" => self.team_id.to_string(), + "owner_id" => self.owner_id.to_string(), + ) + .increment(value); + } + + fn record_spend_usd(&self, value: f64) { + if !value.is_finite() || value <= 0.0 { + return; + } + let micro_usd = (value * 1_000_000.0).round(); + if micro_usd <= 0.0 { + return; + } + metrics::counter!( + M_LLM_SPEND_MICRO_USD_TOTAL, + "endpoint" => self.endpoint.to_string(), + "inbound_protocol" => self.inbound_protocol.to_string(), + "provider" => self.provider.to_string(), + "model" => self.model.to_string(), + "upstream_model" => self.upstream_model.to_string(), + "provider_key_id" => self.provider_key_id.to_string(), + "api_key_id" => self.api_key_id.to_string(), + "team_id" => self.team_id.to_string(), + "owner_id" => self.owner_id.to_string(), + ) + .increment(micro_usd as u64); + } +} + +#[derive(Debug, Clone, Copy, Default)] +pub struct LlmUsage { + pub input_tokens: u32, + pub output_tokens: u32, + pub total_tokens: u32, + pub spend_usd: f64, +} + +impl LlmUsage { + fn is_empty(self) -> bool { + self.input_tokens == 0 + && self.output_tokens == 0 + && self.total_tokens == 0 + && self.spend_usd <= 0.0 + } +} + +#[derive(Debug, Clone, Copy)] +pub struct DeploymentLabels<'a> { + pub provider: &'a str, + pub model: &'a str, + pub upstream_model: &'a str, + pub provider_key_id: &'a str, +} + +impl Default for DeploymentLabels<'_> { + fn default() -> Self { + Self { + provider: "unknown", + model: "unknown", + upstream_model: "unknown", + provider_key_id: "unknown", + } + } +} + +impl DeploymentLabels<'_> { + fn record_counter(&self, metric: &'static str) { + metrics::counter!( + metric, + "provider" => self.provider.to_string(), + "model" => self.model.to_string(), + "upstream_model" => self.upstream_model.to_string(), + "provider_key_id" => self.provider_key_id.to_string(), + ) + .increment(1); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeploymentState { + Healthy, + PartialFailure, + Down, +} + +impl DeploymentState { + fn as_f64(self) -> f64 { + match self { + Self::Healthy => 0.0, + Self::PartialFailure => 1.0, + Self::Down => 2.0, + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct BudgetLabels<'a> { + pub api_key_id: &'a str, + pub team_id: &'a str, + pub owner_id: &'a str, +} + +impl Default for BudgetLabels<'_> { + fn default() -> Self { + Self { + api_key_id: "unknown", + team_id: "unknown", + owner_id: "unknown", + } + } +} + +impl BudgetLabels<'_> { + fn record_gauge(&self, metric: &'static str, value: f64) { + metrics::gauge!( + metric, + "api_key_id" => self.api_key_id.to_string(), + "team_id" => self.team_id.to_string(), + "owner_id" => self.owner_id.to_string(), + ) + .set(value); + } +} + +#[derive(Debug, Clone, Copy, Default)] +pub struct BudgetGauges { + pub limit_usd: Option, + pub spent_usd: Option, + pub remaining_usd: Option, + pub reset_seconds: Option, } /// Canonical outcome label for [`Metrics::record_request`]. Keeps the @@ -217,4 +734,91 @@ mod tests { "expected total 42 in exposition, got:\n{rendered}" ); } + + #[test] + fn aisix_native_request_usage_and_latency_metrics_render() { + let m = Metrics::new(false); + let labels = RequestLabels { + endpoint: "/v1/chat/completions", + inbound_protocol: "openai", + provider: "openai", + model: "gpt", + upstream_model: "gpt-4o", + provider_key_id: "pk-1", + api_key_id: "ak-1", + team_id: "team-1", + owner_id: "owner-1", + status: 200, + outcome: RequestOutcome::Success, + }; + m.record_proxy_request(labels, Duration::from_millis(25)); + m.record_llm_request(labels, Duration::from_millis(20)); + m.record_llm_usage( + UsageLabels { + endpoint: "/v1/chat/completions", + inbound_protocol: "openai", + provider: "openai", + model: "gpt", + upstream_model: "gpt-4o", + provider_key_id: "pk-1", + api_key_id: "ak-1", + team_id: "team-1", + owner_id: "owner-1", + }, + LlmUsage { + input_tokens: 5, + output_tokens: 7, + total_tokens: 12, + spend_usd: 0.001, + }, + ); + m.record_time_to_first_token( + UsageLabels { + endpoint: "/v1/chat/completions", + inbound_protocol: "openai", + provider: "openai", + model: "gpt", + upstream_model: "gpt-4o", + provider_key_id: "pk-1", + api_key_id: "ak-1", + team_id: "team-1", + owner_id: "owner-1", + }, + Duration::from_millis(42), + ); + + let rendered = m.render(); + assert!(rendered.contains(M_PROXY_REQUESTS_TOTAL)); + assert!(rendered.contains(M_LLM_REQUESTS_TOTAL)); + assert!(rendered.contains(M_LLM_INPUT_TOKENS_TOTAL)); + assert!(rendered.contains(M_LLM_OUTPUT_TOKENS_TOTAL)); + assert!(rendered.contains(M_LLM_TOTAL_TOKENS_TOTAL)); + assert!(rendered.contains(M_LLM_SPEND_MICRO_USD_TOTAL)); + assert!(rendered.contains(M_LLM_REQUEST_DURATION)); + assert!(rendered.contains(M_LLM_TTFT)); + assert!(rendered.contains("endpoint=\"/v1/chat/completions\"")); + assert!(rendered.contains("team_id=\"team-1\"")); + } + + #[test] + fn zero_llm_usage_does_not_emit_samples() { + let m = Metrics::new(false); + m.record_llm_usage(UsageLabels::default(), LlmUsage::default()); + let rendered = m.render(); + assert!(!rendered.contains(M_LLM_INPUT_TOKENS_TOTAL)); + assert!(!rendered.contains(M_LLM_TOTAL_TOKENS_TOTAL)); + } + + #[test] + fn in_flight_gauge_returns_to_zero() { + let m = Metrics::new(false); + m.increment_proxy_in_flight("/v1/chat/completions", "openai"); + m.decrement_proxy_in_flight("/v1/chat/completions", "openai"); + let rendered = m.render(); + assert!(rendered.contains(M_PROXY_IN_FLIGHT)); + assert!( + rendered.contains(" 0"), + "expected gauge to return to zero:\n{rendered}" + ); + } } diff --git a/crates/aisix-proxy/src/budget.rs b/crates/aisix-proxy/src/budget.rs index 9cf367df..d9ff510d 100644 --- a/crates/aisix-proxy/src/budget.rs +++ b/crates/aisix-proxy/src/budget.rs @@ -15,6 +15,7 @@ use dashmap::DashMap; use serde::Deserialize; +use serde_json::Value; use std::time::{Duration, Instant}; const CACHE_TTL: Duration = Duration::from_secs(5); @@ -43,6 +44,15 @@ pub struct Decision { pub allowed: bool, pub fail_mode: FailMode, pub reason: Option, + pub budget: Option, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct BudgetDetails { + pub limit_usd: Option, + pub spent_usd: Option, + pub remaining_usd: Option, + pub reset_seconds: Option, } impl Decision { @@ -51,6 +61,7 @@ impl Decision { allowed: true, fail_mode: FailMode::Open, reason: None, + budget: None, } } } @@ -170,6 +181,7 @@ impl BudgetClient { allowed: false, fail_mode: FailMode::Sticky, reason: Some("cp-api unreachable and no cached decision".to_string()), + budget: None, } } @@ -204,16 +216,19 @@ fn apply_fail_mode(prev: &Decision) -> Decision { allowed: true, fail_mode: FailMode::Open, reason: None, + budget: prev.budget.clone(), }, FailMode::Closed => Decision { allowed: false, fail_mode: FailMode::Closed, reason: Some("cp-api unreachable; fail_mode=closed".to_string()), + budget: prev.budget.clone(), }, FailMode::Sticky => Decision { allowed: false, fail_mode: FailMode::Sticky, reason: Some("cp-api unreachable; cached decision stale".to_string()), + budget: prev.budget.clone(), }, } } @@ -244,12 +259,34 @@ struct WireDecision { fail_mode: String, #[serde(default)] reason: Option, + #[serde(default)] + budget: Option, } #[derive(Debug, Deserialize)] struct WireReason { #[serde(default)] message: String, + #[serde(default)] + limit_usd: Option, + #[serde(default)] + spent_usd: Option, + #[serde(default)] + remaining_usd: Option, + #[serde(default, alias = "period_resets_at")] + reset_seconds: Option, +} + +#[derive(Debug, Deserialize)] +struct WireBudget { + #[serde(default, alias = "max_usd")] + limit_usd: Option, + #[serde(default)] + spent_usd: Option, + #[serde(default)] + remaining_usd: Option, + #[serde(default, alias = "period_resets_at")] + reset_seconds: Option, } async fn fetch_decision( @@ -265,6 +302,18 @@ async fn fetch_decision( .await? .error_for_status()?; let wire: WireDecision = resp.json().await?; + let reason_budget = wire.reason.as_ref().map(|r| BudgetDetails { + limit_usd: value_as_f64(r.limit_usd.as_ref()), + spent_usd: value_as_f64(r.spent_usd.as_ref()), + remaining_usd: value_as_f64(r.remaining_usd.as_ref()), + reset_seconds: value_as_u64(r.reset_seconds.as_ref()), + }); + let top_budget = wire.budget.as_ref().map(|b| BudgetDetails { + limit_usd: value_as_f64(b.limit_usd.as_ref()), + spent_usd: value_as_f64(b.spent_usd.as_ref()), + remaining_usd: value_as_f64(b.remaining_usd.as_ref()), + reset_seconds: value_as_u64(b.reset_seconds.as_ref()), + }); let reason = wire.reason.and_then(|r| { if r.message.is_empty() { None @@ -276,9 +325,26 @@ async fn fetch_decision( allowed: wire.allow, fail_mode: FailMode::parse(&wire.fail_mode), reason, + budget: top_budget.or(reason_budget), }) } +fn value_as_f64(value: Option<&Value>) -> Option { + match value? { + Value::Number(n) => n.as_f64(), + Value::String(s) => s.parse().ok(), + _ => None, + } +} + +fn value_as_u64(value: Option<&Value>) -> Option { + match value? { + Value::Number(n) => n.as_u64(), + Value::String(s) => s.parse().ok(), + _ => None, + } +} + #[cfg(test)] mod tests { use super::*; @@ -309,6 +375,37 @@ mod tests { assert_eq!(d.fail_mode, FailMode::Sticky); } + #[tokio::test] + async fn live_client_parses_optional_budget_details() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/dp/budget_check")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "allow": true, + "fail_mode": "sticky", + "budget": { + "limit_usd": "10.5", + "spent_usd": 4.25, + "remaining_usd": "6.25", + "reset_seconds": 3600 + } + }))) + .mount(&server) + .await; + + let c = BudgetClient::new(server.uri(), reqwest::Client::new()); + let d = c.check("k-1").await; + assert_eq!( + d.budget, + Some(BudgetDetails { + limit_usd: Some(10.5), + spent_usd: Some(4.25), + remaining_usd: Some(6.25), + reset_seconds: Some(3600), + }) + ); + } + #[tokio::test] async fn live_client_returns_deny_when_cp_says_no() { let server = MockServer::start().await; diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index ad10d12b..351f9025 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -19,7 +19,9 @@ use aisix_cache::CacheKey; use aisix_gateway::{BridgeContext, BridgeError, ChatFormat}; use aisix_guardrails::GuardrailVerdict; -use aisix_obs::{AccessLog, Metrics, RequestOutcome, UsageEvent}; +use aisix_obs::{ + AccessLog, LlmUsage, Metrics, RequestLabels, RequestOutcome, UsageEvent, UsageLabels, +}; use axum::extract::State; use axum::http::HeaderValue; use axum::response::sse::{Event, KeepAlive, Sse}; @@ -86,6 +88,9 @@ pub async fn chat_completions( &state.metrics, &success.provider, &model_name, + &api_key_id, + auth.key().team_id.as_deref(), + auth.key().owner_id.as_deref(), status, &success, elapsed, @@ -143,6 +148,12 @@ pub async fn chat_completions( let rl_limits = auth.key().rate_limit.clone().unwrap_or_default(); if let Some(rl_status) = state.limiter.peek(&api_key_id, &rl_limits) { crate::render::inject_ratelimit_headers(&mut success.response, &rl_status); + state.metrics.set_rate_limit_remaining( + &api_key_id, + &model_name, + rl_status.rpm_remaining(), + rl_status.tpm_remaining(), + ); } // Correlation / routing headers. if let Ok(v) = axum::http::HeaderValue::try_from(request_id.as_str()) { @@ -262,6 +273,8 @@ struct Success { provider_request_id: String, /// Resolved model the provider actually billed. provider_model_version: String, + provider_key_id: String, + upstream_model: String, /// finish_reason / stop_reason as the upstream returned it. Empty /// for streaming (no terminal event yet) and cache hits. finish_reason: String, @@ -437,6 +450,11 @@ async fn dispatch( // Budget pre-check via cp-api. The DP no longer owns budget state; // cp-api returns a cached/live decision per api_key. let decision = state.budgets.check(&auth.entry.id).await; + if let Some(budget) = decision.budget.as_ref() { + record_budget_gauges(&state.metrics, auth, Some(budget)); + } else { + record_budget_gauges(&state.metrics, auth, None); + } if !decision.allowed { return Err(with_model(ProxyError::BudgetExceeded( decision.reason.unwrap_or_else(|| auth.entry.id.clone()), @@ -543,9 +561,16 @@ async fn dispatch( // return (the non-streaming path's spot) would record zeros. let limiter = Arc::clone(&state.limiter); let state_for_telem = state.clone(); + let metrics_for_stream = state.metrics.clone(); let request_id_for_telem = request_id.to_string(); let model_id_for_telem = model_id.clone(); let api_key_id_for_telem = auth.entry.id.clone(); + let team_id_for_metrics = auth.key().team_id.clone(); + let owner_id_for_metrics = auth.key().owner_id.clone(); + let provider_for_metrics = format!("{provider:?}").to_lowercase(); + let model_for_metrics = req.model.clone(); + let provider_key_id_for_metrics = pk_entry.id.clone(); + let upstream_model_for_metrics = model.upstream_model().unwrap_or("unknown").to_string(); let bypass_reason_for_telem = bypass_reason.clone().unwrap_or_default(); // Per #204: pass the gateway's guardrail chain so the // streaming path can run output guardrails at end-of-stream @@ -623,6 +648,39 @@ async fn dispatch( /* cost_usd */ 0.0, comp.guardrail_blocked, ); + metrics_for_stream.record_llm_usage( + UsageLabels { + endpoint: "/v1/chat/completions", + inbound_protocol: "openai", + provider: &provider_for_metrics, + model: &model_for_metrics, + upstream_model: &upstream_model_for_metrics, + provider_key_id: &provider_key_id_for_metrics, + api_key_id: &api_key_id_for_telem, + team_id: team_id_for_metrics.as_deref().unwrap_or("unknown"), + owner_id: owner_id_for_metrics.as_deref().unwrap_or("unknown"), + }, + LlmUsage { + input_tokens: comp.prompt_tokens, + output_tokens: comp.completion_tokens, + total_tokens: comp.total_tokens.min(u64::from(u32::MAX)) as u32, + spend_usd: 0.0, + }, + ); + metrics_for_stream.record_time_to_first_token( + UsageLabels { + endpoint: "/v1/chat/completions", + inbound_protocol: "openai", + provider: &provider_for_metrics, + model: &model_for_metrics, + upstream_model: &upstream_model_for_metrics, + provider_key_id: &provider_key_id_for_metrics, + api_key_id: &api_key_id_for_telem, + team_id: team_id_for_metrics.as_deref().unwrap_or("unknown"), + owner_id: owner_id_for_metrics.as_deref().unwrap_or("unknown"), + }, + Duration::from_millis(u64::from(comp.ttft_ms)), + ); }, ); let response = @@ -645,6 +703,8 @@ async fn dispatch( cache_read_tokens: 0, provider_request_id: String::new(), provider_model_version: String::new(), + provider_key_id: pk_entry.id.clone(), + upstream_model: model.upstream_model().unwrap_or("unknown").to_string(), finish_reason: String::new(), bypass_reason: bypass_reason.clone(), // Streaming responses aren't cached at this layer — see @@ -733,6 +793,16 @@ async fn dispatch( .provider .map(|p| format!("{p:?}").to_lowercase()) .unwrap_or_else(|| "unknown".into()); + let provider_key_id = attempt_models[0] + .model + .provider_key_id + .clone() + .unwrap_or_else(|| "unknown".into()); + let upstream_model = attempt_models[0] + .model + .upstream_model() + .unwrap_or("unknown") + .to_string(); let mut response = Json(render_response(now, cached)).into_response(); response .headers_mut() @@ -754,6 +824,8 @@ async fn dispatch( // so we leave these blank deliberately. provider_request_id: String::new(), provider_model_version: String::new(), + provider_key_id, + upstream_model, finish_reason: String::new(), // Cache hits don't burn cost on our side (we already // paid the upstream price the first time around). @@ -783,6 +855,8 @@ async fn dispatch( // (non-429 4xx) errors stop immediately. let mut last_err: Option = None; let mut chosen_provider: Option = None; + let mut chosen_provider_key_id: Option = None; + let mut chosen_upstream_model: Option = None; let mut upstream: Option = None; let retries = virtual_entry .value @@ -833,6 +907,9 @@ async fn dispatch( state.health.record_success(&model.display_name); state.runtime_status.mark_healthy(&attempt.id); chosen_provider = Some(format!("{provider:?}").to_lowercase()); + chosen_provider_key_id = Some(pk_entry.id.clone()); + chosen_upstream_model = + Some(model.upstream_model().unwrap_or("unknown").to_string()); upstream = Some(resp); break; } @@ -886,6 +963,8 @@ async fn dispatch( return Err(with_model(ProxyError::Bridge(err))); }; let provider_name = chosen_provider.unwrap_or_else(|| "unknown".into()); + let provider_key_id = chosen_provider_key_id.unwrap_or_else(|| "unknown".into()); + let upstream_model = chosen_upstream_model.unwrap_or_else(|| "unknown".into()); // Output guardrail. Tokens still count against quota — the upstream // already burned them — so commit before the check, and refuse the @@ -1008,6 +1087,8 @@ async fn dispatch( cache_read_tokens, provider_request_id, provider_model_version, + provider_key_id, + upstream_model, finish_reason, cost_usd, bypass_reason, @@ -1116,19 +1197,82 @@ fn finish_reason_label(reason: &aisix_gateway::FinishReason) -> String { } } +#[allow(clippy::too_many_arguments)] fn record_success( metrics: &Metrics, provider: &str, model: &str, + api_key_id: &str, + team_id: Option<&str>, + owner_id: Option<&str>, status: u16, s: &Success, elapsed: Duration, ) { let outcome = RequestOutcome::from_status(status); metrics.record_request(provider, model, status, outcome, elapsed); + let request_labels = RequestLabels { + endpoint: "/v1/chat/completions", + inbound_protocol: "openai", + provider, + model, + upstream_model: &s.upstream_model, + provider_key_id: &s.provider_key_id, + api_key_id, + team_id: team_id.unwrap_or("unknown"), + owner_id: owner_id.unwrap_or("unknown"), + status, + outcome, + }; + metrics.record_proxy_request(request_labels, elapsed); + metrics.record_llm_request(request_labels, elapsed); if let Some(total) = s.total_tokens { metrics.record_tokens(provider, model, total); } + metrics.record_llm_usage( + UsageLabels { + endpoint: "/v1/chat/completions", + inbound_protocol: "openai", + provider, + model, + upstream_model: &s.upstream_model, + provider_key_id: &s.provider_key_id, + api_key_id, + team_id: team_id.unwrap_or("unknown"), + owner_id: owner_id.unwrap_or("unknown"), + }, + LlmUsage { + input_tokens: s.prompt_tokens.unwrap_or(0).min(u64::from(u32::MAX)) as u32, + output_tokens: s.completion_tokens.unwrap_or(0).min(u64::from(u32::MAX)) as u32, + total_tokens: s.total_tokens.unwrap_or(0).min(u64::from(u32::MAX)) as u32, + spend_usd: s.cost_usd, + }, + ); +} + +fn record_budget_gauges( + metrics: &Metrics, + auth: &AuthenticatedKey, + budget: Option<&crate::budget::BudgetDetails>, +) { + let labels = aisix_obs::BudgetLabels { + api_key_id: &auth.entry.id, + team_id: auth.key().team_id.as_deref().unwrap_or("unknown"), + owner_id: auth.key().owner_id.as_deref().unwrap_or("unknown"), + }; + if let Some(budget) = budget { + metrics.set_budget_gauges( + labels, + aisix_obs::BudgetGauges { + limit_usd: budget.limit_usd, + spent_usd: budget.spent_usd, + remaining_usd: budget.remaining_usd, + reset_seconds: budget.reset_seconds, + }, + ); + } else { + metrics.clear_budget_gauges(labels); + } } /// Push one telemetry event onto the CP-side sink **and** fan it out diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index f654b302..42924730 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -97,9 +97,60 @@ pub fn build_router(state: ProxyState) -> Router { state.clone(), enforce_request_body_limit, )) + .layer(middleware::from_fn_with_state( + state.clone(), + record_in_flight_request, + )) .with_state(state) } +async fn record_in_flight_request( + State(state): State, + request: Request, + next: Next, +) -> Response { + let endpoint = request.uri().path().to_string(); + let inbound_protocol = inbound_protocol_for_endpoint(&endpoint).to_string(); + let _guard = InFlightGuard::new(state.metrics.clone(), endpoint, inbound_protocol); + next.run(request).await +} + +fn inbound_protocol_for_endpoint(endpoint: &str) -> &'static str { + if endpoint == "/v1/messages" { + "anthropic" + } else { + "openai" + } +} + +struct InFlightGuard { + metrics: std::sync::Arc, + endpoint: String, + inbound_protocol: String, +} + +impl InFlightGuard { + fn new( + metrics: std::sync::Arc, + endpoint: String, + inbound_protocol: String, + ) -> Self { + metrics.increment_proxy_in_flight(&endpoint, &inbound_protocol); + Self { + metrics, + endpoint, + inbound_protocol, + } + } +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + self.metrics + .decrement_proxy_in_flight(&self.endpoint, &self.inbound_protocol); + } +} + /// Per RFC 9110 §15.5.14, a request body that exceeds the gateway's /// configured `request_body_limit_bytes` must surface as a clean /// `413 Content Too Large` response — NOT an `ECONNRESET` from a diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 99c43469..992c6aa3 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -26,7 +26,7 @@ //! side can handle them consistently regardless of which endpoint was used. use aisix_core::models::Provider; -use aisix_obs::{AccessLog, RequestOutcome, UsageEvent}; +use aisix_obs::{AccessLog, LlmUsage, RequestLabels, RequestOutcome, UsageEvent, UsageLabels}; use axum::extract::State; use axum::http::{HeaderName, HeaderValue}; use axum::response::{IntoResponse, Response}; @@ -69,6 +69,8 @@ pub async fn messages( Ok(DispatchOutcome { response, provider_label, + provider_key_id, + upstream_model, metrics, }) => { let elapsed = started.elapsed(); @@ -88,11 +90,33 @@ pub async fn messages( RequestOutcome::from_status(status), elapsed, ); + let outcome = RequestOutcome::from_status(status); + let labels = RequestLabels { + endpoint: "/v1/messages", + inbound_protocol: "anthropic", + provider: &provider_label, + model: &model_name, + upstream_model: &upstream_model, + provider_key_id: &provider_key_id, + api_key_id: &api_key_id, + team_id: auth.key().team_id.as_deref().unwrap_or("unknown"), + owner_id: auth.key().owner_id.as_deref().unwrap_or("unknown"), + status, + outcome, + }; + state.metrics.record_proxy_request(labels, elapsed); + state.metrics.record_llm_request(labels, elapsed); emit_anthropic_usage_event( &state, &request_id, &model_id, &api_key_id, + &provider_label, + &model_name, + &provider_key_id, + &upstream_model, + auth.key().team_id.as_deref(), + auth.key().owner_id.as_deref(), status, elapsed, metrics, @@ -126,6 +150,12 @@ pub async fn messages( &request_id, &model_id, &api_key_id, + "unknown", + &model_name, + "unknown", + "unknown", + auth.key().team_id.as_deref(), + auth.key().owner_id.as_deref(), status, elapsed, AnthropicUsageMetrics::default(), @@ -310,6 +340,8 @@ async fn dispatch( Ok(DispatchOutcome { response, provider_label, + provider_key_id: pk_entry.id.clone(), + upstream_model: upstream_model.clone(), metrics: AnthropicUsageMetrics::default(), }) } else { @@ -344,6 +376,8 @@ async fn dispatch( Ok(DispatchOutcome { response: Json(json_body).into_response(), provider_label, + provider_key_id: pk_entry.id.clone(), + upstream_model, metrics, }) } @@ -455,6 +489,8 @@ async fn cross_provider_dispatch( let pk_arc = Arc::new(provider_key.clone()); let ctx = BridgeContext::new(request_id, model_arc, pk_arc); let provider_label = format!("{provider:?}").to_lowercase(); + let provider_key_id = model.provider_key_id.as_deref().unwrap_or("unknown"); + let upstream_model = model.upstream_model().unwrap_or("unknown").to_string(); if is_stream { let upstream = bridge.chat_stream(&chat, &ctx).await.map_err(|err| { @@ -495,6 +531,8 @@ async fn cross_provider_dispatch( return Ok(DispatchOutcome { response, provider_label, + provider_key_id: provider_key_id.to_string(), + upstream_model, metrics: AnthropicUsageMetrics::default(), }); } @@ -523,6 +561,8 @@ async fn cross_provider_dispatch( Ok(DispatchOutcome { response: Json(json).into_response(), provider_label, + provider_key_id: provider_key_id.to_string(), + upstream_model, metrics, }) } @@ -577,6 +617,8 @@ fn build_anthropic_sse_stream( struct DispatchOutcome { response: Response, provider_label: String, + provider_key_id: String, + upstream_model: String, metrics: AnthropicUsageMetrics, } @@ -604,11 +646,18 @@ struct AnthropicUsageMetrics { /// to an Anthropic upstream skips the call — the upstream byte stream /// isn't parsed in-flight, so token counts aren't available; that /// path's UsageEvent emission is tracked as follow-up work. +#[allow(clippy::too_many_arguments)] fn emit_anthropic_usage_event( state: &ProxyState, request_id: &str, model_id: &str, api_key_id: &str, + provider: &str, + model: &str, + provider_key_id: &str, + upstream_model: &str, + team_id: Option<&str>, + owner_id: Option<&str>, status_code: u16, elapsed: Duration, metrics: AnthropicUsageMetrics, @@ -636,6 +685,27 @@ fn emit_anthropic_usage_event( state .otlp_fan_out .fan_out(&event, exporters.iter().map(|e| &e.value)); + state.metrics.record_llm_usage( + UsageLabels { + endpoint: "/v1/messages", + inbound_protocol: "anthropic", + provider, + model, + upstream_model, + provider_key_id, + api_key_id, + team_id: team_id.unwrap_or("unknown"), + owner_id: owner_id.unwrap_or("unknown"), + }, + LlmUsage { + input_tokens: metrics.prompt_tokens, + output_tokens: metrics.completion_tokens, + total_tokens: metrics + .prompt_tokens + .saturating_add(metrics.completion_tokens), + spend_usd: 0.0, + }, + ); } fn emit_access_log( diff --git a/crates/aisix-proxy/src/quota.rs b/crates/aisix-proxy/src/quota.rs index 236bb68c..62f3baf7 100644 --- a/crates/aisix-proxy/src/quota.rs +++ b/crates/aisix-proxy/src/quota.rs @@ -134,6 +134,24 @@ pub(crate) async fn enforce<'a>( model_rl: Option<&ModelRateLimit>, ) -> Result, ProxyError> { let decision = state.budgets.check(&auth.entry.id).await; + let budget_labels = aisix_obs::BudgetLabels { + api_key_id: &auth.entry.id, + team_id: auth.key().team_id.as_deref().unwrap_or("unknown"), + owner_id: auth.key().owner_id.as_deref().unwrap_or("unknown"), + }; + if let Some(budget) = decision.budget.as_ref() { + state.metrics.set_budget_gauges( + budget_labels, + aisix_obs::BudgetGauges { + limit_usd: budget.limit_usd, + spent_usd: budget.spent_usd, + remaining_usd: budget.remaining_usd, + reset_seconds: budget.reset_seconds, + }, + ); + } else { + state.metrics.clear_budget_gauges(budget_labels); + } if !decision.allowed { return Err(ProxyError::BudgetExceeded( decision.reason.unwrap_or_else(|| auth.entry.id.clone()), diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index f53cc185..68762afa 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -452,6 +452,7 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { Arc::new(EtcdConfigStore::new(admin_client, etcd_prefix.clone())); let admin_state = AdminState::new(snapshot_handle.clone(), admin_store, &cfg.admin) .with_metrics(metrics.clone()) + .with_prometheus_config(cfg.observability.metrics.prometheus.clone()) // Share the health tracker so /admin/v1/health reflects live // per-model upstream failure counts. .with_health_tracker(health_tracker) diff --git a/docs/configuration/bootstrap-config.md b/docs/configuration/bootstrap-config.md index 0f1f5e1a..7ad14fa4 100644 --- a/docs/configuration/bootstrap-config.md +++ b/docs/configuration/bootstrap-config.md @@ -70,6 +70,10 @@ observability: service_name: "aisix" log_level: "info" access_log: true + metrics: + prometheus: + enabled: true + path: "/metrics" cache: backend: "memory" @@ -157,6 +161,8 @@ Use `observability` to configure: Bootstrap observability settings are process-wide. They are different from dynamic `ObservabilityExporter` rows, which control data-plane telemetry fan-out for request events. +`observability.metrics.prometheus.enabled` controls whether the admin listener mounts the Prometheus scrape endpoint. `observability.metrics.prometheus.path` controls the mounted path and defaults to `/metrics`. + ## `cache` Use `cache` to choose the bootstrap cache backend. diff --git a/docs/operations/metrics-and-logs.md b/docs/operations/metrics-and-logs.md index 865d6437..3396af5d 100644 --- a/docs/operations/metrics-and-logs.md +++ b/docs/operations/metrics-and-logs.md @@ -10,12 +10,30 @@ Use them together. No single signal tells the whole story. ## Metrics -`GET /metrics` on the admin listener is the Prometheus scrape endpoint. +`GET /metrics` on the admin listener is the default Prometheus scrape endpoint. Operators can change it with `observability.metrics.prometheus.path`, or disable the endpoint with `observability.metrics.prometheus.enabled: false`. This endpoint is unauthenticated by design on the private admin listener. Treat `/metrics` as infrastructure-facing, not as a public diagnostics surface. +AISIX exposes native metric names with the `aisix_` prefix. Existing compatibility series remain: + +- `aisix_requests_total` +- `aisix_request_duration_seconds` +- `aisix_ratelimit_rejections_total` +- `aisix_tokens_consumed_total` + +The Prometheus integration also emits LiteLLM-category equivalents under AISIX-native names: + +- usage and cost: `aisix_llm_input_tokens_total`, `aisix_llm_output_tokens_total`, `aisix_llm_total_tokens_total`, `aisix_llm_spend_micro_usd_total` +- request volume and latency: `aisix_llm_requests_total`, `aisix_llm_request_duration_seconds`, `aisix_llm_time_to_first_token_seconds` +- proxy health: `aisix_proxy_requests_total`, `aisix_proxy_failed_requests_total`, `aisix_proxy_request_duration_seconds`, `aisix_proxy_in_flight_requests` +- quotas and budgets: `aisix_ratelimit_remaining_requests`, `aisix_ratelimit_remaining_tokens`, and budget gauges when the control plane returns budget detail fields; `aisix_budget_details_present` tells scrapers whether the current budget response carried those optional fields +- deployment and routing: `aisix_deployment_*` and `aisix_routing_*` metric families when the request path has those events +- exporter/cache health: Redis, usage-event drop, and OTLP fan-out drop/failure counters + +Labels are limited to values the data plane has reliably: `endpoint`, `inbound_protocol`, `provider`, `model`, `upstream_model`, `provider_key_id`, `api_key_id`, `team_id`, `owner_id`, `status`, and `outcome`. User email, team alias, and end-user labels are not fabricated by the data plane. + ## Access Logs And Usage Signals Current proxy behavior emits: diff --git a/tests/e2e/src/cases/prometheus-metrics-e2e.test.ts b/tests/e2e/src/cases/prometheus-metrics-e2e.test.ts new file mode 100644 index 00000000..a44b712d --- /dev/null +++ b/tests/e2e/src/cases/prometheus-metrics-e2e.test.ts @@ -0,0 +1,176 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + ProxyClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +const CALLER_PLAINTEXT = "sk-prometheus-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +describe("prometheus metrics e2e", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream({ + nonStreamBody: responseBody(), + }); + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + await configureOpenAi(admin, upstream, "prometheus-gpt"); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + test("scrape contains AISIX-native request and token metrics", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + + const proxy = new ProxyClient(app.proxyUrl, CALLER_PLAINTEXT); + await waitConfigPropagation(async () => { + const probe = await proxy.chat({ + model: "prometheus-gpt", + messages: [{ role: "user", content: "ready" }], + }); + return probe.status === 200; + }); + + const { status, body } = await proxy.chat({ + model: "prometheus-gpt", + messages: [{ role: "user", content: "metrics" }], + }); + expect(status, JSON.stringify(body)).toBe(200); + + const scrape = await fetch(`${app.adminUrl}/metrics`); + expect(scrape.status).toBe(200); + const text = await scrape.text(); + + expect(text).toContain("aisix_proxy_requests_total"); + expect(text).toContain("aisix_llm_requests_total"); + expect(text).toContain("aisix_llm_input_tokens_total"); + expect(text).toContain("aisix_llm_output_tokens_total"); + expect(text).toContain("aisix_llm_total_tokens_total"); + expect(text).toContain("aisix_proxy_in_flight_requests"); + expect(text).toMatch( + /aisix_proxy_requests_total\{[^}]*endpoint="\/v1\/chat\/completions"[^}]*model="prometheus-gpt"[^}]*status="200"/, + ); + expect(text).toMatch( + /aisix_llm_requests_total\{[^}]*endpoint="\/v1\/chat\/completions"[^}]*model="prometheus-gpt"[^}]*status="200"/, + ); + expect(text).toContain('team_id="unknown"'); + expect(text).toContain('owner_id="unknown"'); + }); + + test("custom prometheus path is used for scrapes", async (ctx) => { + if (!etcdReachable) { + ctx.skip(); + return; + } + + const customUpstream = await startOpenAiUpstream({ + nonStreamBody: responseBody(), + }); + const customApp = await spawnApp({ prometheusPath: "/custom-metrics" }); + try { + const customAdmin = new AdminClient(customApp.adminUrl, customApp.adminKey); + await configureOpenAi(customAdmin, customUpstream, "prometheus-custom-gpt"); + const proxy = new ProxyClient(customApp.proxyUrl, CALLER_PLAINTEXT); + await waitConfigPropagation(async () => { + const probe = await proxy.chat({ + model: "prometheus-custom-gpt", + messages: [{ role: "user", content: "ready" }], + }); + return probe.status === 200; + }); + + const defaultScrape = await fetch(`${customApp.adminUrl}/metrics`); + expect(defaultScrape.status).toBe(404); + + const scrape = await fetch(`${customApp.adminUrl}/custom-metrics`); + expect(scrape.status).toBe(200); + const text = await scrape.text(); + expect(text).toMatch( + /aisix_proxy_requests_total\{[^}]*endpoint="\/v1\/chat\/completions"[^}]*model="prometheus-custom-gpt"/, + ); + expect(text).toContain("aisix_llm_total_tokens_total"); + } finally { + await customApp.exit(); + await customUpstream.close(); + } + }); + + test("disabled prometheus endpoint is not mounted", async (ctx) => { + if (!etcdReachable) { + ctx.skip(); + return; + } + + const disabledApp = await spawnApp({ prometheus: false }); + try { + const scrape = await fetch(`${disabledApp.adminUrl}/metrics`); + expect(scrape.status).toBe(404); + expect(await scrape.text()).not.toContain("aisix_"); + } finally { + await disabledApp.exit(); + } + }); +}); + +function responseBody() { + return { + id: "chatcmpl-prom-1", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "hello" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 11, completion_tokens: 13, total_tokens: 24 }, + }; +} + +async function configureOpenAi( + admin: AdminClient, + upstream: OpenAiUpstream, + modelName: string, +) { + const pk = await admin.createProviderKey({ + display_name: `${modelName}-pk`, + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: modelName, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: [modelName], + }); +} diff --git a/tests/e2e/src/harness/app.ts b/tests/e2e/src/harness/app.ts index 5b78ecf7..ae27373d 100644 --- a/tests/e2e/src/harness/app.ts +++ b/tests/e2e/src/harness/app.ts @@ -12,8 +12,10 @@ import { harnessRequest } from "./http.js"; export interface AppOverrides { /** Inserted into `admin.admin_keys`. Defaults to a fresh random key. */ adminKey?: string; - /** Whether to enable Prometheus on `/metrics`. Defaults to true. */ + /** Whether to enable the Prometheus scrape endpoint. Defaults to true. */ prometheus?: boolean; + /** Prometheus scrape path. Defaults to `/metrics`. */ + prometheusPath?: string; /** Extra raw config keys merged into the YAML at the top level. */ extra?: Record; } @@ -66,7 +68,10 @@ export async function spawnApp(overrides: AppOverrides = {}): Promise