Skip to main content

edgeguard/
metrics.rs

1//! Prometheus metrics, hand-rolled.
2//!
3//! A full metrics library (`prometheus`, `metrics`) would be a heavy dependency for the
4//! handful of series EdgeGuard exposes, so — in the same spirit as `parse_host_port` being a
5//! small URL parser rather than a full one — this is a minimal text-exposition renderer over
6//! a few atomics. It emits the Prometheus text format (v0.0.4) at `/__edgeguard/metrics`.
7//!
8//! The registry lives in [`crate::proxy::AppState`] *outside* the hot-swappable runtime, so
9//! counters survive a config hot-reload instead of resetting to zero.
10
11use std::collections::BTreeMap;
12use std::sync::atomic::{AtomicU64, Ordering};
13use std::sync::Mutex;
14use std::time::Duration;
15
16/// Request `outcome` label values. These mirror the `outcome` field already emitted on the
17/// JSON access log in [`crate::proxy`], so a metric series lines up 1:1 with a log line.
18/// Anything not in this list is bucketed under `other` rather than silently dropped.
19const OUTCOMES: &[&str] = &[
20    "ok",
21    "rate_limited",
22    "over_quota",
23    "over_budget",
24    "unpriced_model",
25    "limiter_error",
26    "unauthorized",
27    "forbidden",
28    "method_not_allowed",
29    "not_found",
30    "payload_too_large",
31    "header_too_large",
32    "bad_gateway",
33    "upstream_error",
34    "upstream_timeout",
35    "upstream_body_too_large",
36    "upstream_body_error",
37    "other",
38];
39
40/// Outcomes where the edge itself *denied* the request by policy — auth, WAF-forbidden, rate limit,
41/// quota, hard budget, or an unpriced LLM model. Distinct from protocol/upstream failures
42/// (`not_found`, `bad_gateway`, `upstream_*`), which aren't "blocked by eggrd". Feeds the managed-mode
43/// usage report's `blocked` figure.
44fn outcome_is_blocked(outcome: &str) -> bool {
45    matches!(
46        outcome,
47        "rate_limited"
48            | "over_quota"
49            | "over_budget"
50            | "unpriced_model"
51            | "unauthorized"
52            | "forbidden"
53    )
54}
55
56/// Rate-limit `scope` label values (which limiter rejected the request).
57const RL_SCOPES: &[&str] = &["ip", "route", "key"];
58
59/// WAF `rule` label values (which ruleset class matched). Custom `[[waf.rules]]` all roll up
60/// under `custom`; the specific rule id is in the log line, not the metric.
61const WAF_RULES: &[&str] = &["sqli", "xss", "path_traversal", "custom"];
62
63/// Upper bounds (seconds) for the request-duration histogram, plus an implicit `+Inf`.
64const LATENCY_BUCKETS: &[f64] = &[
65    0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
66];
67
68/// LLM metering `result` label values (per [`crate::llm`]): a request whose response carried usage
69/// for a priced model (`metered`), for an unpriced model (`unpriced`, tokens counted but no cost),
70/// or that reported no usage at all (`no_usage`, e.g. an error or a stream without `include_usage`).
71const LLM_RESULTS: &[&str] = &["metered", "unpriced", "no_usage"];
72
73/// Per-model token `kind` label values on `edgeguard_llm_model_tokens_total`. `input`/`output` are
74/// the prompt/completion totals; `cached` and `reasoning` are the sub-dimensions (⊆ input / ⊆ output)
75/// that providers bill differently — surfaced separately so the "~7× undercount" is visible.
76const LLM_TOKEN_KINDS: &[&str] = &["input", "output", "cached", "reasoning"];
77
78/// Budget `scope` label values (which budget dimension blocked / was consumed). Mirrors
79/// [`crate::budget::BudgetScope`]; a scope not in this list is bucketed under `other`.
80const BUDGET_SCOPES: &[&str] = &["global", "key", "model", "team", "other"];
81
82/// Cap on the number of distinct `model` label values tracked for the per-model token/cost series.
83/// Clients can send an arbitrary `model` string, so the map is bounded to keep Prometheus cardinality
84/// flat; once full, further models fold into the `_over_cap` bucket rather than growing without limit.
85const MAX_LLM_MODEL_SERIES: usize = 128;
86
87/// The overflow bucket a new model folds into once [`MAX_LLM_MODEL_SERIES`] distinct models are seen.
88const LLM_MODEL_OVERFLOW: &str = "_over_cap";
89
90/// Key-vault (gateway L2) `result` label values: a request whose virtual key resolved and was
91/// swapped for the provider key (`swapped`), rejected because the virtual key was unknown
92/// (`denied_key`), or rejected because the requested model was off the key's egress allowlist
93/// (`denied_model`).
94const KEYVAULT_RESULTS: &[&str] = &["swapped", "denied_key", "denied_model"];
95
96/// DLP (gateway L3) finding categories — the `category` label on `edgeguard_llm_dlp_findings_total`.
97/// Mirrors [`crate::dlp::CATEGORIES`]; kept here to avoid a cross-module compile dependency in the
98/// hot render path. A category not in this list is bucketed under `other`.
99const DLP_CATEGORIES: &[&str] = &[
100    "email",
101    "credit_card",
102    "aws_key",
103    "api_key",
104    "private_key",
105    "ssn",
106    "phone",
107    "iban",
108    "high_entropy",
109    "gazetteer",
110    "person",
111    "address",
112    "org",
113    "prompt_injection",
114    "custom",
115    "other",
116];
117
118/// A drained snapshot of the managed-mode usage accumulators (requests + bandwidth + LLM
119/// tokens/cost). Returned by [`Metrics::drain_usage`] and re-applied by [`Metrics::restore_usage`]
120/// if the report fails to send.
121#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
122pub struct DrainedUsage {
123    pub requests: u64,
124    pub ingress_bytes: u64,
125    pub egress_bytes: u64,
126    pub tokens_in: u64,
127    pub tokens_out: u64,
128    pub cost_micros: u64,
129    /// Requests the edge denied since the last drain (a subset of `requests`).
130    pub blocked: u64,
131    /// WAF matches since the last drain, by rule class (indices parallel to [`WAF_RULES`] =
132    /// `[sqli, xss, path_traversal, custom]`). Reported to the control plane for the console's
133    /// per-category security breakdown.
134    pub waf_sqli: u64,
135    pub waf_xss: u64,
136    pub waf_path_traversal: u64,
137    pub waf_custom: u64,
138}
139
140impl DrainedUsage {
141    /// True when nothing accrued — the reporter skips an empty report.
142    pub fn is_empty(&self) -> bool {
143        *self == DrainedUsage::default()
144    }
145}
146
147/// One metered LLM request's token/cost breakdown, passed to [`Metrics::record_llm_usage`]. Carries
148/// the four token dimensions plus the cost (present only when the model was priced). Grouping them in
149/// a struct keeps the call site readable and the four dims impossible to transpose positionally.
150#[derive(Clone, Copy, Debug, Default)]
151pub struct LlmSample {
152    pub tokens_in: u64,
153    pub tokens_out: u64,
154    /// Cached prompt tokens (⊆ `tokens_in`).
155    pub cached_tokens: u64,
156    /// Reasoning completion tokens (⊆ `tokens_out`).
157    pub reasoning_tokens: u64,
158    /// Cost in micro-dollars; `None` when the model is unpriced (tokens still counted).
159    pub cost_micros: Option<u64>,
160}
161
162/// Per-model token/cost accumulators (behind the [`Metrics::per_model`] mutex). Plain integers, not
163/// atomics: the map is small, updated under the lock, and read only at render time.
164#[derive(Clone, Copy, Debug, Default)]
165struct PerModelCounters {
166    tokens_in: u64,
167    tokens_out: u64,
168    cached_tokens: u64,
169    reasoning_tokens: u64,
170    cost_micros: u64,
171}
172
173/// Process-wide metric registry. All methods take `&self` and use relaxed atomics — metrics
174/// are monotonic counters/observations where exact inter-thread ordering doesn't matter.
175pub struct Metrics {
176    /// Off-box access-log shipping, when `[log.ship]` is enabled.
177    ///
178    /// It lives here because `Metrics` is already the telemetry sink threaded to every response
179    /// path — `finish()` in proxy.rs takes it and has 40 call sites. Giving the shipper its own
180    /// parameter would mean touching all forty to carry a value that is `None` in the default
181    /// configuration, for no gain in clarity: an access log IS telemetry, and this is where the
182    /// telemetry object lives.
183    ///
184    /// `OnceLock` because the shipper needs a Tokio runtime to spawn its task, so it cannot be
185    /// built in `Metrics::new()` — it is installed once, at startup, after the runtime exists.
186    log_shipper: std::sync::OnceLock<crate::logship::LogShipper>,
187    /// Request-tracing span shipper, when `[tracing]` is enabled. Same `OnceLock`-installed-at-boot
188    /// shape and the same reason as `log_shipper`: it needs a Tokio runtime to spawn its task, and
189    /// it rides the object the response path already holds rather than a parameter added to
190    /// `finish`'s 38 call sites.
191    span_shipper: std::sync::OnceLock<crate::telemetry::SpanShipper>,
192    /// One counter per [`OUTCOMES`] entry (parallel index).
193    requests: Vec<AtomicU64>,
194    /// One counter per [`RL_SCOPES`] entry (parallel index).
195    ratelimit_hits: Vec<AtomicU64>,
196    /// One counter per [`WAF_RULES`] entry (parallel index).
197    waf_hits: Vec<AtomicU64>,
198    /// Cumulative histogram buckets (parallel to [`LATENCY_BUCKETS`]): `bucket[i]` counts
199    /// observations with value <= `LATENCY_BUCKETS[i]`.
200    latency_buckets: Vec<AtomicU64>,
201    latency_sum_micros: AtomicU64,
202    latency_count: AtomicU64,
203    csp_reports: AtomicU64,
204    /// Drainable usage accumulators for managed-mode reporting (requests + bandwidth *since the
205    /// last drain*). Kept separate from the monotonic Prometheus counters above precisely because
206    /// the usage reporter resets these to zero each period — a Prometheus counter must not decrease.
207    usage_requests: AtomicU64,
208    usage_ingress_bytes: AtomicU64,
209    usage_egress_bytes: AtomicU64,
210    /// Drainable count of requests the edge *denied* (auth / WAF-forbidden / rate-limit / quota /
211    /// budget / unpriced-model) since the last drain, for the managed-mode usage report's "blocked"
212    /// figure. A subset of `usage_requests`. Distinct from the monotonic per-outcome counters.
213    usage_blocked: AtomicU64,
214    /// Drainable LLM token/cost usage for managed-mode cost reports (reset each report period,
215    /// like the request/byte accumulators above). Distinct from the monotonic `llm_*` counters.
216    usage_tokens_in: AtomicU64,
217    usage_tokens_out: AtomicU64,
218    usage_cost_micros: AtomicU64,
219    /// Drainable WAF matches by rule class (parallel to [`WAF_RULES`]), for the managed-mode usage
220    /// report's per-category security breakdown. Distinct from the monotonic `waf_hits` counters,
221    /// which the reporter must not reset.
222    usage_waf_hits: Vec<AtomicU64>,
223    /// LLM input (prompt) tokens metered (monotonic).
224    llm_tokens_in: AtomicU64,
225    /// LLM output (completion) tokens metered (monotonic).
226    llm_tokens_out: AtomicU64,
227    /// LLM cached prompt tokens (⊆ input), metered separately (monotonic). The dim the "~7× undercount"
228    /// gap is about — surfaced so cache utilisation and its cost impact are both visible.
229    llm_cached_tokens: AtomicU64,
230    /// LLM reasoning completion tokens (⊆ output), metered separately (monotonic).
231    llm_reasoning_tokens: AtomicU64,
232    /// Accumulated LLM cost in micro-dollars (1e-6 USD), for priced models only.
233    llm_cost_micros: AtomicU64,
234    /// Server-side time-to-first-token histogram (buckets parallel to [`LATENCY_BUCKETS`]) for
235    /// streamed LLM responses, measured in the gateway as frames flow — no client clock, no in-app
236    /// instrumentation. A trace backend only sees span-end duration, so this is a request-path-only
237    /// signal, and a server-measured one rather than a client-clock approximation.
238    llm_ttft_buckets: Vec<AtomicU64>,
239    llm_ttft_sum_micros: AtomicU64,
240    llm_ttft_count: AtomicU64,
241    /// Mean time-per-output-token histogram for streamed LLM responses with >1 output token
242    /// (inter-token latency = span of the output stream / (output_tokens − 1)).
243    llm_tpot_buckets: Vec<AtomicU64>,
244    llm_tpot_sum_micros: AtomicU64,
245    llm_tpot_count: AtomicU64,
246    /// One counter per [`LLM_RESULTS`] entry (parallel index).
247    llm_results: Vec<AtomicU64>,
248    /// Per-model token/cost breakdown (`edgeguard_llm_model_*`), bounded to [`MAX_LLM_MODEL_SERIES`]
249    /// distinct models (overflow → [`LLM_MODEL_OVERFLOW`]) so client-chosen model strings can't blow
250    /// up Prometheus cardinality.
251    per_model: Mutex<BTreeMap<String, PerModelCounters>>,
252    /// Per-team token/cost breakdown (`edgeguard_llm_team_*`), keyed by the `[llm].team_header` value
253    /// (absent → `_none`). Same cardinality bound + overflow bucket as `per_model`, so it answers
254    /// "which team spent this" for chargeback/showback without a Prometheus label explosion.
255    per_team: Mutex<BTreeMap<String, PerModelCounters>>,
256    /// Per-key token/cost breakdown (`edgeguard_llm_key_*`), keyed by the authenticated **principal**
257    /// (API-key id / Basic user / JWT `sub`; unauthenticated → `_anon`) — the OSS identity primitive.
258    /// Same cardinality bound + overflow bucket, so per-user attribution is reachable in the OSS core.
259    per_key: Mutex<BTreeMap<String, PerModelCounters>>,
260    /// One counter per [`BUDGET_SCOPES`] entry (parallel index): requests blocked by a hard budget.
261    budget_blocked: Vec<AtomicU64>,
262    /// Latest observed consumed ratio (`used / limit`, 0.0–1.0+) per budget *name* — the near-limit
263    /// signal. A coarse gauge (one value per budget name, last writer wins across scope keys), which
264    /// is what an "any budget near its cap" alert needs. Bounded by the operator-defined name set.
265    budget_consumed: Mutex<BTreeMap<String, f64>>,
266    /// One counter per [`KEYVAULT_RESULTS`] entry (parallel index).
267    keyvault_results: Vec<AtomicU64>,
268    /// One counter per [`DLP_CATEGORIES`] entry (parallel index): DLP findings by category.
269    dlp_findings: Vec<AtomicU64>,
270    /// Requests blocked (`403`) by DLP `block` mode.
271    dlp_blocked: AtomicU64,
272    /// Budget reconcile/release operations that ultimately FAILED (after retries) against the shared
273    /// store. Each failure means a reserve→settle didn't complete, so the distributed counter has
274    /// drifted (a leaked hold → phantom `BudgetExceededError`, or an uncharged settle → silent
275    /// bypass). Surfaced so this drift is **observable** instead of only logged — alert on it.
276    budget_reconcile_failures: AtomicU64,
277}
278
279impl Default for Metrics {
280    fn default() -> Self {
281        Metrics {
282            log_shipper: std::sync::OnceLock::new(),
283            span_shipper: std::sync::OnceLock::new(),
284            requests: OUTCOMES.iter().map(|_| AtomicU64::new(0)).collect(),
285            ratelimit_hits: RL_SCOPES.iter().map(|_| AtomicU64::new(0)).collect(),
286            waf_hits: WAF_RULES.iter().map(|_| AtomicU64::new(0)).collect(),
287            latency_buckets: LATENCY_BUCKETS.iter().map(|_| AtomicU64::new(0)).collect(),
288            latency_sum_micros: AtomicU64::new(0),
289            latency_count: AtomicU64::new(0),
290            csp_reports: AtomicU64::new(0),
291            usage_requests: AtomicU64::new(0),
292            usage_ingress_bytes: AtomicU64::new(0),
293            usage_egress_bytes: AtomicU64::new(0),
294            usage_blocked: AtomicU64::new(0),
295            usage_tokens_in: AtomicU64::new(0),
296            usage_tokens_out: AtomicU64::new(0),
297            usage_cost_micros: AtomicU64::new(0),
298            usage_waf_hits: WAF_RULES.iter().map(|_| AtomicU64::new(0)).collect(),
299            llm_tokens_in: AtomicU64::new(0),
300            llm_tokens_out: AtomicU64::new(0),
301            llm_cached_tokens: AtomicU64::new(0),
302            llm_reasoning_tokens: AtomicU64::new(0),
303            llm_cost_micros: AtomicU64::new(0),
304            llm_ttft_buckets: LATENCY_BUCKETS.iter().map(|_| AtomicU64::new(0)).collect(),
305            llm_ttft_sum_micros: AtomicU64::new(0),
306            llm_ttft_count: AtomicU64::new(0),
307            llm_tpot_buckets: LATENCY_BUCKETS.iter().map(|_| AtomicU64::new(0)).collect(),
308            llm_tpot_sum_micros: AtomicU64::new(0),
309            llm_tpot_count: AtomicU64::new(0),
310            llm_results: LLM_RESULTS.iter().map(|_| AtomicU64::new(0)).collect(),
311            per_model: Mutex::new(BTreeMap::new()),
312            per_team: Mutex::new(BTreeMap::new()),
313            per_key: Mutex::new(BTreeMap::new()),
314            budget_blocked: BUDGET_SCOPES.iter().map(|_| AtomicU64::new(0)).collect(),
315            budget_consumed: Mutex::new(BTreeMap::new()),
316            keyvault_results: KEYVAULT_RESULTS.iter().map(|_| AtomicU64::new(0)).collect(),
317            dlp_findings: DLP_CATEGORIES.iter().map(|_| AtomicU64::new(0)).collect(),
318            dlp_blocked: AtomicU64::new(0),
319            budget_reconcile_failures: AtomicU64::new(0),
320        }
321    }
322}
323
324/// Observe `elapsed` into a cumulative histogram (buckets parallel to [`LATENCY_BUCKETS`]) plus its
325/// running sum (micros) and count. Shared by the request-latency and the LLM TTFT/TPOT histograms so
326/// their bucketing can't drift.
327fn observe_hist(
328    buckets: &[AtomicU64],
329    sum_micros: &AtomicU64,
330    count: &AtomicU64,
331    elapsed: Duration,
332) {
333    let secs = elapsed.as_secs_f64();
334    for (i, bound) in LATENCY_BUCKETS.iter().enumerate() {
335        if secs <= *bound {
336            buckets[i].fetch_add(1, Ordering::Relaxed);
337        }
338    }
339    sum_micros.fetch_add(elapsed.as_micros() as u64, Ordering::Relaxed);
340    count.fetch_add(1, Ordering::Relaxed);
341}
342
343/// Render one Prometheus histogram (`_bucket`/`_sum`/`_count`) to `out`, sharing the exposition
344/// shape with the request-latency histogram above.
345fn render_hist(
346    out: &mut String,
347    name: &str,
348    help: &str,
349    buckets: &[AtomicU64],
350    sum_micros: &AtomicU64,
351    count: &AtomicU64,
352) {
353    out.push_str(&format!("# HELP {name} {help}\n"));
354    out.push_str(&format!("# TYPE {name} histogram\n"));
355    for (i, bound) in LATENCY_BUCKETS.iter().enumerate() {
356        let v = buckets[i].load(Ordering::Relaxed);
357        out.push_str(&format!("{name}_bucket{{le=\"{bound}\"}} {v}\n"));
358    }
359    // The `+Inf` bucket equals the total observation count by definition.
360    let c = count.load(Ordering::Relaxed);
361    out.push_str(&format!("{name}_bucket{{le=\"+Inf\"}} {c}\n"));
362    let sum_secs = sum_micros.load(Ordering::Relaxed) as f64 / 1_000_000.0;
363    out.push_str(&format!("{name}_sum {sum_secs}\n"));
364    out.push_str(&format!("{name}_count {c}\n"));
365}
366
367/// Render one per-dimension token/cost breakdown (`per_model` / `per_team` / `per_key`) to `out`,
368/// sharing the exposition shape across all three so they can't drift from one another.
369fn render_breakdown(
370    out: &mut String,
371    metric_prefix: &str,
372    label_name: &str,
373    dim_desc: &str,
374    map: &BTreeMap<String, PerModelCounters>,
375) {
376    out.push_str(&format!(
377        "# HELP {metric_prefix}_tokens_total LLM tokens metered by {dim_desc} and kind.\n"
378    ));
379    out.push_str(&format!("# TYPE {metric_prefix}_tokens_total counter\n"));
380    for (k, c) in map.iter() {
381        let label = escape_label(k);
382        for kind in LLM_TOKEN_KINDS {
383            let v = match *kind {
384                "input" => c.tokens_in,
385                "output" => c.tokens_out,
386                "cached" => c.cached_tokens,
387                "reasoning" => c.reasoning_tokens,
388                _ => 0,
389            };
390            out.push_str(&format!(
391                "{metric_prefix}_tokens_total{{{label_name}=\"{label}\",kind=\"{kind}\"}} {v}\n"
392            ));
393        }
394    }
395    out.push_str(&format!(
396        "# HELP {metric_prefix}_cost_microdollars_total LLM cost (micro-dollars) by {dim_desc}.\n"
397    ));
398    out.push_str(&format!(
399        "# TYPE {metric_prefix}_cost_microdollars_total counter\n"
400    ));
401    for (k, c) in map.iter() {
402        out.push_str(&format!(
403            "{metric_prefix}_cost_microdollars_total{{{label_name}=\"{}\"}} {}\n",
404            escape_label(k),
405            c.cost_micros
406        ));
407    }
408}
409
410/// Add `s` to the `key` bucket of a bounded per-dimension accumulator (`per_model` / `per_team`): a
411/// new key is inserted only while under [`MAX_LLM_MODEL_SERIES`]; past the cap it folds into
412/// [`LLM_MODEL_OVERFLOW`], so a flood of distinct label values can't grow the map without bound.
413fn accumulate_bounded(map: &Mutex<BTreeMap<String, PerModelCounters>>, key: &str, s: &LlmSample) {
414    let mut map = map.lock().expect("per-dimension mutex poisoned");
415    let k = if map.contains_key(key) || map.len() < MAX_LLM_MODEL_SERIES {
416        key
417    } else {
418        LLM_MODEL_OVERFLOW
419    };
420    let c = map.entry(k.to_string()).or_default();
421    c.tokens_in = c.tokens_in.saturating_add(s.tokens_in);
422    c.tokens_out = c.tokens_out.saturating_add(s.tokens_out);
423    c.cached_tokens = c.cached_tokens.saturating_add(s.cached_tokens);
424    c.reasoning_tokens = c.reasoning_tokens.saturating_add(s.reasoning_tokens);
425    c.cost_micros = c.cost_micros.saturating_add(s.cost_micros.unwrap_or(0));
426}
427
428impl Metrics {
429    pub fn new() -> Self {
430        Self::default()
431    }
432
433    /// Install the access-log shipper. Called once at startup, after the Tokio runtime exists.
434    ///
435    /// Returns whether it took: a second call is ignored rather than replacing a live shipper,
436    /// which would leak the first one's background task and its queued records.
437    pub fn set_log_shipper(&self, s: crate::logship::LogShipper) -> bool {
438        self.log_shipper.set(s).is_ok()
439    }
440
441    /// The shipper, if `[log.ship]` is enabled.
442    pub fn log_shipper(&self) -> Option<&crate::logship::LogShipper> {
443        self.log_shipper.get()
444    }
445
446    /// Install the span shipper. Called once at startup, after the Tokio runtime exists.
447    pub fn set_span_shipper(&self, s: crate::telemetry::SpanShipper) -> bool {
448        self.span_shipper.set(s).is_ok()
449    }
450
451    /// The span shipper, if `[tracing]` is enabled. `None` is also the cheap gate on the request
452    /// path: with tracing off, nothing about a span is computed.
453    pub fn span_shipper(&self) -> Option<&crate::telemetry::SpanShipper> {
454        self.span_shipper.get()
455    }
456
457    /// Count one finished request under its `outcome` label.
458    pub fn record_request(&self, outcome: &str) {
459        let idx = OUTCOMES
460            .iter()
461            .position(|o| *o == outcome)
462            .unwrap_or(OUTCOMES.len() - 1); // -> "other"
463        self.requests[idx].fetch_add(1, Ordering::Relaxed);
464    }
465
466    /// Observe a request's end-to-end latency into the histogram.
467    pub fn observe_latency(&self, elapsed: Duration) {
468        observe_hist(
469            &self.latency_buckets,
470            &self.latency_sum_micros,
471            &self.latency_count,
472            elapsed,
473        );
474    }
475
476    /// Observe a streamed LLM response's server-side **time-to-first-token** and, when the response
477    /// had more than one output token, its mean **time-per-output-token**. Called once per streamed
478    /// LLM request from the response body's `Drop`, after the terminal `usage` frame is parsed. The
479    /// gateway sits in the token stream, so these are measured with no client clock and no in-app
480    /// instrumentation — the request-path advantage a trace backend (which only sees span-end
481    /// duration) can't offer.
482    pub fn record_llm_latency(&self, ttft: Duration, tpot: Option<Duration>) {
483        observe_hist(
484            &self.llm_ttft_buckets,
485            &self.llm_ttft_sum_micros,
486            &self.llm_ttft_count,
487            ttft,
488        );
489        if let Some(tpot) = tpot {
490            observe_hist(
491                &self.llm_tpot_buckets,
492                &self.llm_tpot_sum_micros,
493                &self.llm_tpot_count,
494                tpot,
495            );
496        }
497    }
498
499    /// Count a rate-limit rejection by which limiter scope tripped (`ip`/`route`/`key`).
500    pub fn record_ratelimit_hit(&self, scope: &str) {
501        if let Some(idx) = RL_SCOPES.iter().position(|s| *s == scope) {
502            self.ratelimit_hits[idx].fetch_add(1, Ordering::Relaxed);
503        }
504    }
505
506    /// Count one WAF rule match by rule class (`sqli`/`xss`/`path_traversal`/`custom`).
507    /// Recorded for both report-only and blocking modes — so a report-first rollout is
508    /// visible — while a *blocked* request is additionally counted under the `forbidden`
509    /// request outcome.
510    pub fn record_waf_hit(&self, class: &str) {
511        if let Some(idx) = WAF_RULES.iter().position(|c| *c == class) {
512            self.waf_hits[idx].fetch_add(1, Ordering::Relaxed);
513            // Parallel drainable accumulator for the managed-mode usage report.
514            self.usage_waf_hits[idx].fetch_add(1, Ordering::Relaxed);
515        }
516    }
517
518    /// Count one received CSP violation report.
519    pub fn record_csp_report(&self) {
520        self.csp_reports.fetch_add(1, Ordering::Relaxed);
521    }
522
523    /// Count one request toward the drainable usage accumulator (managed mode). Called once per
524    /// request from the single `finish` exit, so every request — proxied or rejected — counts.
525    /// `outcome` is the request's outcome label; a denial outcome (see [`outcome_is_blocked`]) also
526    /// bumps the drainable `blocked` accumulator, so the control plane can show what the edge screened.
527    pub fn add_usage_request(&self, outcome: &str) {
528        self.usage_requests.fetch_add(1, Ordering::Relaxed);
529        if outcome_is_blocked(outcome) {
530            self.usage_blocked.fetch_add(1, Ordering::Relaxed);
531        }
532    }
533
534    /// Add request (ingress) + response (egress) bytes to the drainable usage accumulator. Called
535    /// on the proxied path where both bodies are buffered and the counts are known.
536    pub fn add_usage_bytes(&self, ingress: usize, egress: usize) {
537        self.usage_ingress_bytes
538            .fetch_add(ingress as u64, Ordering::Relaxed);
539        self.usage_egress_bytes
540            .fetch_add(egress as u64, Ordering::Relaxed);
541    }
542
543    /// Atomically read-and-zero the usage accumulators — the delta the usage reporter ships to the
544    /// control plane (requests + bandwidth + LLM tokens/cost, gateway L4).
545    pub fn drain_usage(&self) -> DrainedUsage {
546        debug_assert_eq!(WAF_RULES, ["sqli", "xss", "path_traversal", "custom"]);
547        DrainedUsage {
548            requests: self.usage_requests.swap(0, Ordering::Relaxed),
549            ingress_bytes: self.usage_ingress_bytes.swap(0, Ordering::Relaxed),
550            egress_bytes: self.usage_egress_bytes.swap(0, Ordering::Relaxed),
551            tokens_in: self.usage_tokens_in.swap(0, Ordering::Relaxed),
552            tokens_out: self.usage_tokens_out.swap(0, Ordering::Relaxed),
553            cost_micros: self.usage_cost_micros.swap(0, Ordering::Relaxed),
554            blocked: self.usage_blocked.swap(0, Ordering::Relaxed),
555            // Indices parallel to WAF_RULES = [sqli, xss, path_traversal, custom].
556            waf_sqli: self.usage_waf_hits[0].swap(0, Ordering::Relaxed),
557            waf_xss: self.usage_waf_hits[1].swap(0, Ordering::Relaxed),
558            waf_path_traversal: self.usage_waf_hits[2].swap(0, Ordering::Relaxed),
559            waf_custom: self.usage_waf_hits[3].swap(0, Ordering::Relaxed),
560        }
561    }
562
563    /// Add a previously-drained delta back, e.g. when a usage report failed to send — so the
564    /// next period reships it instead of losing billable usage. (New requests that arrived during
565    /// the failed send simply add on top, as intended.)
566    pub fn restore_usage(&self, u: &DrainedUsage) {
567        debug_assert_eq!(WAF_RULES, ["sqli", "xss", "path_traversal", "custom"]);
568        self.usage_requests.fetch_add(u.requests, Ordering::Relaxed);
569        self.usage_ingress_bytes
570            .fetch_add(u.ingress_bytes, Ordering::Relaxed);
571        self.usage_egress_bytes
572            .fetch_add(u.egress_bytes, Ordering::Relaxed);
573        self.usage_tokens_in
574            .fetch_add(u.tokens_in, Ordering::Relaxed);
575        self.usage_tokens_out
576            .fetch_add(u.tokens_out, Ordering::Relaxed);
577        self.usage_cost_micros
578            .fetch_add(u.cost_micros, Ordering::Relaxed);
579        self.usage_blocked.fetch_add(u.blocked, Ordering::Relaxed);
580        // Indices parallel to WAF_RULES = [sqli, xss, path_traversal, custom].
581        self.usage_waf_hits[0].fetch_add(u.waf_sqli, Ordering::Relaxed);
582        self.usage_waf_hits[1].fetch_add(u.waf_xss, Ordering::Relaxed);
583        self.usage_waf_hits[2].fetch_add(u.waf_path_traversal, Ordering::Relaxed);
584        self.usage_waf_hits[3].fetch_add(u.waf_custom, Ordering::Relaxed);
585    }
586
587    /// Record one metered LLM request for `model`: add its four token dimensions and — when the model
588    /// was priced — its cost (micro-dollars). `cost_micros == None` means the model isn't in the price
589    /// book, so tokens are still counted but the request is bucketed `unpriced` rather than `metered`.
590    /// Also updates the bounded per-model breakdown (`edgeguard_llm_model_*`).
591    pub fn record_llm_usage(&self, model: &str, s: LlmSample) {
592        self.llm_tokens_in.fetch_add(s.tokens_in, Ordering::Relaxed);
593        self.llm_tokens_out
594            .fetch_add(s.tokens_out, Ordering::Relaxed);
595        self.llm_cached_tokens
596            .fetch_add(s.cached_tokens, Ordering::Relaxed);
597        self.llm_reasoning_tokens
598            .fetch_add(s.reasoning_tokens, Ordering::Relaxed);
599        // Drainable accumulators for the managed-mode cost report.
600        self.usage_tokens_in
601            .fetch_add(s.tokens_in, Ordering::Relaxed);
602        self.usage_tokens_out
603            .fetch_add(s.tokens_out, Ordering::Relaxed);
604        let result = match s.cost_micros {
605            Some(c) => {
606                self.llm_cost_micros.fetch_add(c, Ordering::Relaxed);
607                self.usage_cost_micros.fetch_add(c, Ordering::Relaxed);
608                "metered"
609            }
610            None => "unpriced",
611        };
612        self.bump_llm_result(result);
613        self.record_per_model(model, &s);
614    }
615
616    /// Update the bounded per-model accumulator. A new model is only inserted while under the cap;
617    /// once full it folds into [`LLM_MODEL_OVERFLOW`], so a flood of distinct model strings can't grow
618    /// the map without bound.
619    fn record_per_model(&self, model: &str, s: &LlmSample) {
620        accumulate_bounded(&self.per_model, model, s);
621    }
622
623    /// Record one metered LLM request against its team (`[llm].team_header` value; absent → `_none`),
624    /// for per-team chargeback/showback (`edgeguard_llm_team_*`). Bounded exactly like the per-model
625    /// breakdown. Called alongside [`Self::record_llm_usage`] from the request path.
626    pub fn record_llm_team_usage(&self, team: &str, s: &LlmSample) {
627        accumulate_bounded(&self.per_team, team, s);
628    }
629
630    /// Record one metered LLM request against its authenticated key/principal (`_anon` when
631    /// unauthenticated), for per-user cost attribution (`edgeguard_llm_key_*`). Bounded exactly like
632    /// the per-model breakdown. Reuses the existing OSS auth principal as the identity — so per-key
633    /// FinOps is reachable without the EE control plane.
634    pub fn record_llm_key_usage(&self, key: &str, s: &LlmSample) {
635        accumulate_bounded(&self.per_key, key, s);
636    }
637
638    /// Record a request blocked by a hard LLM budget, by the budget's `scope` (unknown → `other`).
639    pub fn record_budget_blocked(&self, scope: &str) {
640        let idx = BUDGET_SCOPES
641            .iter()
642            .position(|s| *s == scope)
643            .unwrap_or(BUDGET_SCOPES.len() - 1); // -> "other"
644        self.budget_blocked[idx].fetch_add(1, Ordering::Relaxed);
645    }
646
647    /// Record the latest consumed ratio (`used / limit`) for a budget by `name` — the near-limit
648    /// gauge. Last writer wins per name (a coarse "is any budget near its cap" signal). NaN/negative
649    /// samples are dropped so a divide-by-zero can't poison the gauge.
650    pub fn record_budget_consumed(&self, name: &str, ratio: f64) {
651        if !ratio.is_finite() || ratio < 0.0 {
652            return;
653        }
654        let mut map = self
655            .budget_consumed
656            .lock()
657            .expect("budget_consumed mutex poisoned");
658        // Bound the map the same way as models: operator-defined names are few, but never grow past
659        // the cap if a config churns budget names.
660        if map.contains_key(name) || map.len() < MAX_LLM_MODEL_SERIES {
661            map.insert(name.to_string(), ratio);
662        }
663    }
664
665    /// Record `n` budget reconcile/release failures against the shared store (after retries). A
666    /// non-zero rate here means the distributed budget counter is drifting — the signal to alert on.
667    pub fn record_budget_reconcile_failures(&self, n: usize) {
668        if n > 0 {
669            self.budget_reconcile_failures
670                .fetch_add(n as u64, Ordering::Relaxed);
671        }
672    }
673
674    /// Record an LLM request whose response carried no usage (error, or a stream the client didn't
675    /// opt into usage on). No tokens/cost, but the request is visible as `no_usage`.
676    pub fn record_llm_no_usage(&self) {
677        self.bump_llm_result("no_usage");
678    }
679
680    fn bump_llm_result(&self, result: &str) {
681        if let Some(idx) = LLM_RESULTS.iter().position(|r| *r == result) {
682            self.llm_results[idx].fetch_add(1, Ordering::Relaxed);
683        }
684    }
685
686    /// Record a key-vault decision by `result` (`swapped`/`denied_key`/`denied_model`).
687    pub fn record_keyvault(&self, result: &str) {
688        if let Some(idx) = KEYVAULT_RESULTS.iter().position(|r| *r == result) {
689            self.keyvault_results[idx].fetch_add(1, Ordering::Relaxed);
690        }
691    }
692
693    /// Record one DLP finding under its `category` (unknown → `other`).
694    pub fn record_dlp_finding(&self, category: &str) {
695        let idx = DLP_CATEGORIES
696            .iter()
697            .position(|c| *c == category)
698            .unwrap_or(DLP_CATEGORIES.len() - 1); // -> "other"
699        self.dlp_findings[idx].fetch_add(1, Ordering::Relaxed);
700    }
701
702    /// Record a request blocked by DLP `block` mode.
703    pub fn record_dlp_blocked(&self) {
704        self.dlp_blocked.fetch_add(1, Ordering::Relaxed);
705    }
706
707    /// Render the Prometheus text exposition (format version 0.0.4).
708    pub fn render(&self) -> String {
709        let mut out = String::with_capacity(1024);
710
711        out.push_str("# HELP edgeguard_requests_total Total proxied requests by outcome.\n");
712        out.push_str("# TYPE edgeguard_requests_total counter\n");
713        for (i, label) in OUTCOMES.iter().enumerate() {
714            let v = self.requests[i].load(Ordering::Relaxed);
715            out.push_str(&format!(
716                "edgeguard_requests_total{{outcome=\"{label}\"}} {v}\n"
717            ));
718        }
719
720        out.push_str(
721            "# HELP edgeguard_ratelimit_hits_total Requests rejected by a rate limiter, by scope.\n",
722        );
723        out.push_str("# TYPE edgeguard_ratelimit_hits_total counter\n");
724        for (i, label) in RL_SCOPES.iter().enumerate() {
725            let v = self.ratelimit_hits[i].load(Ordering::Relaxed);
726            out.push_str(&format!(
727                "edgeguard_ratelimit_hits_total{{scope=\"{label}\"}} {v}\n"
728            ));
729        }
730
731        out.push_str(
732            "# HELP edgeguard_waf_hits_total WAF rule matches by class (report-only + blocked).\n",
733        );
734        out.push_str("# TYPE edgeguard_waf_hits_total counter\n");
735        for (i, label) in WAF_RULES.iter().enumerate() {
736            let v = self.waf_hits[i].load(Ordering::Relaxed);
737            out.push_str(&format!(
738                "edgeguard_waf_hits_total{{rule=\"{label}\"}} {v}\n"
739            ));
740        }
741
742        out.push_str("# HELP edgeguard_csp_reports_total CSP violation reports received.\n");
743        out.push_str("# TYPE edgeguard_csp_reports_total counter\n");
744        out.push_str(&format!(
745            "edgeguard_csp_reports_total {}\n",
746            self.csp_reports.load(Ordering::Relaxed)
747        ));
748
749        // Access-log shipping. `dropped` is the important one: a log pipeline that silently drops
750        // is worse than no pipeline, because the gap is invisible in the destination and the
751        // absence of records reads as an absence of traffic. Alert on it.
752        if let Some(shipper) = self.log_shipper.get() {
753            let (sent, dropped_q, dropped_send, batches, failed) = shipper.stats().snapshot();
754            out.push_str(
755                "# HELP edgeguard_logship_sent_total Access-log records delivered to the collector.\n\
756                 # TYPE edgeguard_logship_sent_total counter\n",
757            );
758            out.push_str(&format!("edgeguard_logship_sent_total {sent}\n"));
759            out.push_str(
760                "# HELP edgeguard_logship_dropped_total Access-log records dropped, by reason.\n\
761                 # TYPE edgeguard_logship_dropped_total counter\n",
762            );
763            // Two reasons, kept apart because they call for different responses: `queue_full`
764            // means the edge is producing faster than the collector accepts (raise queue_size, or
765            // find out why the collector is slow); `send_failed` means the collector rejected or
766            // was unreachable.
767            out.push_str(&format!(
768                "edgeguard_logship_dropped_total{{reason=\"queue_full\"}} {dropped_q}\n"
769            ));
770            out.push_str(&format!(
771                "edgeguard_logship_dropped_total{{reason=\"send_failed\"}} {dropped_send}\n"
772            ));
773            out.push_str(
774                "# HELP edgeguard_logship_batches_total Access-log batches POSTed, by outcome.\n\
775                 # TYPE edgeguard_logship_batches_total counter\n",
776            );
777            out.push_str(&format!(
778                "edgeguard_logship_batches_total{{outcome=\"sent\"}} {batches}\n"
779            ));
780            out.push_str(&format!(
781                "edgeguard_logship_batches_total{{outcome=\"failed\"}} {failed}\n"
782            ));
783        }
784
785        // Span shipping. Same reasoning as the log shipper's counters: a trace pipeline that drops
786        // silently reads, at the destination, as an absence of traffic.
787        if let Some(sh) = self.span_shipper.get() {
788            let st = sh.stats();
789            let sent = st.sent.load(Ordering::Relaxed);
790            let dq = st.dropped_queue_full.load(Ordering::Relaxed);
791            let df = st.dropped_send_failed.load(Ordering::Relaxed);
792            out.push_str(
793                "# HELP edgeguard_spans_sent_total Request spans delivered to the trace collector.\n\
794                 # TYPE edgeguard_spans_sent_total counter\n",
795            );
796            out.push_str(&format!("edgeguard_spans_sent_total {sent}\n"));
797            out.push_str(
798                "# HELP edgeguard_spans_dropped_total Request spans dropped, by reason.\n\
799                 # TYPE edgeguard_spans_dropped_total counter\n",
800            );
801            out.push_str(&format!(
802                "edgeguard_spans_dropped_total{{reason=\"queue_full\"}} {dq}\n"
803            ));
804            out.push_str(&format!(
805                "edgeguard_spans_dropped_total{{reason=\"send_failed\"}} {df}\n"
806            ));
807        }
808
809        out.push_str(
810            "# HELP edgeguard_request_duration_seconds Request handling latency in seconds.\n",
811        );
812        out.push_str("# TYPE edgeguard_request_duration_seconds histogram\n");
813        for (i, bound) in LATENCY_BUCKETS.iter().enumerate() {
814            let v = self.latency_buckets[i].load(Ordering::Relaxed);
815            out.push_str(&format!(
816                "edgeguard_request_duration_seconds_bucket{{le=\"{bound}\"}} {v}\n"
817            ));
818        }
819        let count = self.latency_count.load(Ordering::Relaxed);
820        // The `+Inf` bucket equals the total observation count by definition.
821        out.push_str(&format!(
822            "edgeguard_request_duration_seconds_bucket{{le=\"+Inf\"}} {count}\n"
823        ));
824        let sum_secs = self.latency_sum_micros.load(Ordering::Relaxed) as f64 / 1_000_000.0;
825        out.push_str(&format!(
826            "edgeguard_request_duration_seconds_sum {sum_secs}\n"
827        ));
828        out.push_str(&format!(
829            "edgeguard_request_duration_seconds_count {count}\n"
830        ));
831
832        // Streamed-LLM server-side latency: time-to-first-token and mean time-per-output-token,
833        // measured in the request path (the signal a trace backend can't produce). Render at 0 too.
834        render_hist(
835            &mut out,
836            "edgeguard_llm_ttft_seconds",
837            "Server-side time-to-first-token for streamed LLM responses, in seconds.",
838            &self.llm_ttft_buckets,
839            &self.llm_ttft_sum_micros,
840            &self.llm_ttft_count,
841        );
842        render_hist(
843            &mut out,
844            "edgeguard_llm_tpot_seconds",
845            "Mean time-per-output-token for streamed LLM responses (>1 output token), in seconds.",
846            &self.llm_tpot_buckets,
847            &self.llm_tpot_sum_micros,
848            &self.llm_tpot_count,
849        );
850
851        // LLM token metering (gateway L0). All series render even at 0 so dashboards/alerts don't
852        // break on a quiet proxy.
853        out.push_str("# HELP edgeguard_llm_tokens_total LLM tokens metered, by direction.\n");
854        out.push_str("# TYPE edgeguard_llm_tokens_total counter\n");
855        out.push_str(&format!(
856            "edgeguard_llm_tokens_total{{direction=\"input\"}} {}\n",
857            self.llm_tokens_in.load(Ordering::Relaxed)
858        ));
859        out.push_str(&format!(
860            "edgeguard_llm_tokens_total{{direction=\"output\"}} {}\n",
861            self.llm_tokens_out.load(Ordering::Relaxed)
862        ));
863
864        // Cached / reasoning sub-dimensions (⊆ input / ⊆ output). Kept as their own metric so they
865        // are visible for the "~7× undercount" story without being double-summed into the direction
866        // totals above.
867        out.push_str(
868            "# HELP edgeguard_llm_cached_tokens_total Cached prompt tokens metered (subset of input).\n",
869        );
870        out.push_str("# TYPE edgeguard_llm_cached_tokens_total counter\n");
871        out.push_str(&format!(
872            "edgeguard_llm_cached_tokens_total {}\n",
873            self.llm_cached_tokens.load(Ordering::Relaxed)
874        ));
875        out.push_str(
876            "# HELP edgeguard_llm_reasoning_tokens_total Reasoning completion tokens metered (subset of output).\n",
877        );
878        out.push_str("# TYPE edgeguard_llm_reasoning_tokens_total counter\n");
879        out.push_str(&format!(
880            "edgeguard_llm_reasoning_tokens_total {}\n",
881            self.llm_reasoning_tokens.load(Ordering::Relaxed)
882        ));
883
884        out.push_str(
885            "# HELP edgeguard_llm_cost_microdollars_total Accumulated LLM cost in micro-dollars (1e-6 USD).\n",
886        );
887        out.push_str("# TYPE edgeguard_llm_cost_microdollars_total counter\n");
888        out.push_str(&format!(
889            "edgeguard_llm_cost_microdollars_total {}\n",
890            self.llm_cost_micros.load(Ordering::Relaxed)
891        ));
892
893        // Per-model breakdown (bounded cardinality). Tokens carry a `kind` label; cost is a separate
894        // series. Rendered only for models actually seen, so a fresh proxy emits nothing here.
895        {
896            let map = self.per_model.lock().expect("per_model mutex poisoned");
897            render_breakdown(&mut out, "edgeguard_llm_model", "model", "model", &map);
898        }
899
900        // Per-team token/cost breakdown (`edgeguard_llm_team_*`), for chargeback/showback. Same
901        // cardinality bound as per-model; rendered only for teams actually seen.
902        {
903            let map = self.per_team.lock().expect("per_team mutex poisoned");
904            render_breakdown(&mut out, "edgeguard_llm_team", "team", "team", &map);
905        }
906
907        // Per-key (per-principal) token/cost breakdown (`edgeguard_llm_key_*`), for per-user FinOps.
908        // Same cardinality bound as per-model; rendered only for keys actually seen.
909        {
910            let map = self.per_key.lock().expect("per_key mutex poisoned");
911            render_breakdown(
912                &mut out,
913                "edgeguard_llm_key",
914                "key",
915                "authenticated key/principal",
916                &map,
917            );
918        }
919
920        // Hard-budget signals (gateway L1): the near-limit gauge and per-scope block counter.
921        out.push_str(
922            "# HELP edgeguard_llm_budget_blocked_total Requests blocked by a hard LLM budget, by scope.\n",
923        );
924        out.push_str("# TYPE edgeguard_llm_budget_blocked_total counter\n");
925        for (i, label) in BUDGET_SCOPES.iter().enumerate() {
926            let v = self.budget_blocked[i].load(Ordering::Relaxed);
927            out.push_str(&format!(
928                "edgeguard_llm_budget_blocked_total{{scope=\"{label}\"}} {v}\n"
929            ));
930        }
931        {
932            let map = self
933                .budget_consumed
934                .lock()
935                .expect("budget_consumed mutex poisoned");
936            out.push_str(
937                "# HELP edgeguard_llm_budget_consumed_ratio Latest consumed ratio (used/limit) per budget.\n",
938            );
939            out.push_str("# TYPE edgeguard_llm_budget_consumed_ratio gauge\n");
940            for (name, ratio) in map.iter() {
941                out.push_str(&format!(
942                    "edgeguard_llm_budget_consumed_ratio{{budget=\"{}\"}} {ratio}\n",
943                    escape_label(name)
944                ));
945            }
946        }
947        out.push_str(
948            "# HELP edgeguard_llm_budget_reconcile_failures_total Budget reserve->settle reconciles that failed against the shared store (counter drift).\n",
949        );
950        out.push_str("# TYPE edgeguard_llm_budget_reconcile_failures_total counter\n");
951        out.push_str(&format!(
952            "edgeguard_llm_budget_reconcile_failures_total {}\n",
953            self.budget_reconcile_failures.load(Ordering::Relaxed)
954        ));
955
956        out.push_str("# HELP edgeguard_llm_requests_total LLM requests metered, by result.\n");
957        out.push_str("# TYPE edgeguard_llm_requests_total counter\n");
958        for (i, label) in LLM_RESULTS.iter().enumerate() {
959            let v = self.llm_results[i].load(Ordering::Relaxed);
960            out.push_str(&format!(
961                "edgeguard_llm_requests_total{{result=\"{label}\"}} {v}\n"
962            ));
963        }
964
965        out.push_str(
966            "# HELP edgeguard_llm_keyvault_total Key-vault decisions by result (swap / egress denial).\n",
967        );
968        out.push_str("# TYPE edgeguard_llm_keyvault_total counter\n");
969        for (i, label) in KEYVAULT_RESULTS.iter().enumerate() {
970            let v = self.keyvault_results[i].load(Ordering::Relaxed);
971            out.push_str(&format!(
972                "edgeguard_llm_keyvault_total{{result=\"{label}\"}} {v}\n"
973            ));
974        }
975
976        out.push_str(
977            "# HELP edgeguard_llm_dlp_findings_total DLP findings (PII / secrets) by category.\n",
978        );
979        out.push_str("# TYPE edgeguard_llm_dlp_findings_total counter\n");
980        for (i, label) in DLP_CATEGORIES.iter().enumerate() {
981            let v = self.dlp_findings[i].load(Ordering::Relaxed);
982            out.push_str(&format!(
983                "edgeguard_llm_dlp_findings_total{{category=\"{label}\"}} {v}\n"
984            ));
985        }
986        out.push_str(
987            "# HELP edgeguard_llm_dlp_blocked_total Requests blocked by DLP block mode.\n",
988        );
989        out.push_str("# TYPE edgeguard_llm_dlp_blocked_total counter\n");
990        out.push_str(&format!(
991            "edgeguard_llm_dlp_blocked_total {}\n",
992            self.dlp_blocked.load(Ordering::Relaxed)
993        ));
994
995        out
996    }
997}
998
999/// Escape a dynamic label *value* for the Prometheus text format: backslash, double-quote, and
1000/// newline must be escaped (per the exposition spec) so a client-chosen `model` or operator-chosen
1001/// budget name can't inject a line break or unbalanced quote into the output.
1002fn escape_label(s: &str) -> String {
1003    let mut out = String::with_capacity(s.len());
1004    for ch in s.chars() {
1005        match ch {
1006            '\\' => out.push_str("\\\\"),
1007            '"' => out.push_str("\\\""),
1008            '\n' => out.push_str("\\n"),
1009            _ => out.push(ch),
1010        }
1011    }
1012    out
1013}
1014
1015#[cfg(test)]
1016mod tests {
1017    use super::*;
1018
1019    #[test]
1020    fn records_and_renders_request_outcomes() {
1021        let m = Metrics::new();
1022        m.record_request("ok");
1023        m.record_request("ok");
1024        m.record_request("rate_limited");
1025        // An unknown outcome falls into the `other` bucket, not "ok".
1026        m.record_request("totally_unknown");
1027
1028        let text = m.render();
1029        assert!(
1030            text.contains("edgeguard_requests_total{outcome=\"ok\"} 2"),
1031            "{text}"
1032        );
1033        assert!(
1034            text.contains("edgeguard_requests_total{outcome=\"rate_limited\"} 1"),
1035            "{text}"
1036        );
1037        assert!(
1038            text.contains("edgeguard_requests_total{outcome=\"other\"} 1"),
1039            "{text}"
1040        );
1041    }
1042
1043    #[test]
1044    fn latency_histogram_is_cumulative() {
1045        let m = Metrics::new();
1046        m.observe_latency(Duration::from_millis(3)); // <= 0.005
1047        m.observe_latency(Duration::from_millis(40)); // <= 0.05
1048        let text = m.render();
1049        // 3ms falls under every bucket >= 0.005; 40ms under every bucket >= 0.05.
1050        assert!(
1051            text.contains("edgeguard_request_duration_seconds_bucket{le=\"0.005\"} 1"),
1052            "{text}"
1053        );
1054        assert!(
1055            text.contains("edgeguard_request_duration_seconds_bucket{le=\"0.05\"} 2"),
1056            "{text}"
1057        );
1058        assert!(
1059            text.contains("edgeguard_request_duration_seconds_bucket{le=\"+Inf\"} 2"),
1060            "{text}"
1061        );
1062        assert!(
1063            text.contains("edgeguard_request_duration_seconds_count 2"),
1064            "{text}"
1065        );
1066    }
1067
1068    #[test]
1069    fn llm_ttft_tpot_histograms_render() {
1070        let m = Metrics::new();
1071        // TTFT 40ms (<= 0.05), TPOT 8ms (<= 0.01).
1072        m.record_llm_latency(Duration::from_millis(40), Some(Duration::from_millis(8)));
1073        // A single-output-token response has no defined TPOT — only TTFT is recorded.
1074        m.record_llm_latency(Duration::from_millis(3), None);
1075        let text = m.render();
1076        // Two TTFT observations; both <= 0.05, one <= 0.005.
1077        assert!(
1078            text.contains("edgeguard_llm_ttft_seconds_bucket{le=\"0.005\"} 1"),
1079            "{text}"
1080        );
1081        assert!(
1082            text.contains("edgeguard_llm_ttft_seconds_bucket{le=\"0.05\"} 2"),
1083            "{text}"
1084        );
1085        assert!(
1086            text.contains("edgeguard_llm_ttft_seconds_count 2"),
1087            "{text}"
1088        );
1089        // One TPOT observation (8ms), recorded only for the >1-token response.
1090        assert!(
1091            text.contains("edgeguard_llm_tpot_seconds_bucket{le=\"0.01\"} 1"),
1092            "{text}"
1093        );
1094        assert!(
1095            text.contains("edgeguard_llm_tpot_seconds_count 1"),
1096            "{text}"
1097        );
1098    }
1099
1100    #[test]
1101    fn ratelimit_and_csp_counters() {
1102        let m = Metrics::new();
1103        m.record_ratelimit_hit("ip");
1104        m.record_ratelimit_hit("route");
1105        m.record_ratelimit_hit("route");
1106        m.record_csp_report();
1107        let text = m.render();
1108        assert!(
1109            text.contains("edgeguard_ratelimit_hits_total{scope=\"ip\"} 1"),
1110            "{text}"
1111        );
1112        assert!(
1113            text.contains("edgeguard_ratelimit_hits_total{scope=\"route\"} 2"),
1114            "{text}"
1115        );
1116        assert!(text.contains("edgeguard_csp_reports_total 1"), "{text}");
1117    }
1118
1119    #[test]
1120    fn usage_accumulates_drains_and_restores() {
1121        let m = Metrics::new();
1122        m.add_usage_request("ok"); // proxied — not blocked
1123        m.add_usage_request("forbidden"); // edge-denied — also counts toward `blocked`
1124        m.add_usage_bytes(100, 250);
1125        m.add_usage_bytes(0, 50);
1126        // LLM token usage also drains for the cost report (gateway L4).
1127        m.record_llm_usage(
1128            "gpt-4o",
1129            LlmSample {
1130                tokens_in: 1_000,
1131                tokens_out: 400,
1132                cost_micros: Some(2_500),
1133                ..Default::default()
1134            },
1135        );
1136        // Drain returns the accrued delta and zeroes the accumulator.
1137        let drained = m.drain_usage();
1138        assert_eq!(drained.requests, 2);
1139        assert_eq!(drained.blocked, 1); // only the "forbidden" request
1140        assert_eq!(drained.ingress_bytes, 100);
1141        assert_eq!(drained.egress_bytes, 300);
1142        assert_eq!(drained.tokens_in, 1_000);
1143        assert_eq!(drained.tokens_out, 400);
1144        assert_eq!(drained.cost_micros, 2_500);
1145        assert!(m.drain_usage().is_empty());
1146        // Restore (failed-report path) re-adds it for the next period.
1147        m.restore_usage(&drained);
1148        assert_eq!(m.drain_usage(), drained);
1149    }
1150
1151    #[test]
1152    fn llm_token_and_cost_counters() {
1153        let m = Metrics::new();
1154        // Priced model: tokens + cost, bucketed `metered`. Includes cached/reasoning sub-dims.
1155        m.record_llm_usage(
1156            "gpt-4o",
1157            LlmSample {
1158                tokens_in: 100,
1159                tokens_out: 50,
1160                cached_tokens: 40,
1161                reasoning_tokens: 20,
1162                cost_micros: Some(1_250),
1163            },
1164        );
1165        // Unpriced model: tokens counted, no cost, bucketed `unpriced`.
1166        m.record_llm_usage(
1167            "mystery",
1168            LlmSample {
1169                tokens_in: 10,
1170                tokens_out: 5,
1171                cost_micros: None,
1172                ..Default::default()
1173            },
1174        );
1175        // No usage reported.
1176        m.record_llm_no_usage();
1177        let text = m.render();
1178        assert!(
1179            text.contains("edgeguard_llm_tokens_total{direction=\"input\"} 110"),
1180            "{text}"
1181        );
1182        assert!(
1183            text.contains("edgeguard_llm_tokens_total{direction=\"output\"} 55"),
1184            "{text}"
1185        );
1186        assert!(
1187            text.contains("edgeguard_llm_cached_tokens_total 40"),
1188            "{text}"
1189        );
1190        assert!(
1191            text.contains("edgeguard_llm_reasoning_tokens_total 20"),
1192            "{text}"
1193        );
1194        assert!(
1195            text.contains("edgeguard_llm_cost_microdollars_total 1250"),
1196            "{text}"
1197        );
1198        assert!(
1199            text.contains("edgeguard_llm_requests_total{result=\"metered\"} 1"),
1200            "{text}"
1201        );
1202        assert!(
1203            text.contains("edgeguard_llm_requests_total{result=\"unpriced\"} 1"),
1204            "{text}"
1205        );
1206        assert!(
1207            text.contains("edgeguard_llm_requests_total{result=\"no_usage\"} 1"),
1208            "{text}"
1209        );
1210        // Per-model breakdown carries the model + kind labels and the priced model's cost.
1211        assert!(
1212            text.contains("edgeguard_llm_model_tokens_total{model=\"gpt-4o\",kind=\"cached\"} 40"),
1213            "{text}"
1214        );
1215        assert!(
1216            text.contains(
1217                "edgeguard_llm_model_tokens_total{model=\"gpt-4o\",kind=\"reasoning\"} 20"
1218            ),
1219            "{text}"
1220        );
1221        assert!(
1222            text.contains("edgeguard_llm_model_cost_microdollars_total{model=\"gpt-4o\"} 1250"),
1223            "{text}"
1224        );
1225    }
1226
1227    #[test]
1228    fn per_team_tokens_and_cost_are_accumulated_and_rendered() {
1229        let m = Metrics::new();
1230        m.record_llm_team_usage(
1231            "acme",
1232            &LlmSample {
1233                tokens_in: 100,
1234                tokens_out: 40,
1235                cached_tokens: 30,
1236                reasoning_tokens: 10,
1237                cost_micros: Some(77),
1238            },
1239        );
1240        m.record_llm_team_usage(
1241            "acme",
1242            &LlmSample {
1243                tokens_in: 50,
1244                tokens_out: 20,
1245                cost_micros: Some(23),
1246                ..Default::default()
1247            },
1248        );
1249        // A request with no team falls into the shared `_none` bucket.
1250        m.record_llm_team_usage(
1251            "_none",
1252            &LlmSample {
1253                tokens_in: 5,
1254                ..Default::default()
1255            },
1256        );
1257        let text = m.render();
1258        assert!(
1259            text.contains("edgeguard_llm_team_tokens_total{team=\"acme\",kind=\"input\"} 150"),
1260            "{text}"
1261        );
1262        assert!(
1263            text.contains("edgeguard_llm_team_tokens_total{team=\"acme\",kind=\"output\"} 60"),
1264            "{text}"
1265        );
1266        assert!(
1267            text.contains("edgeguard_llm_team_cost_microdollars_total{team=\"acme\"} 100"),
1268            "{text}"
1269        );
1270        assert!(
1271            text.contains("edgeguard_llm_team_tokens_total{team=\"_none\",kind=\"input\"} 5"),
1272            "{text}"
1273        );
1274    }
1275
1276    #[test]
1277    fn budget_reconcile_failures_counter_renders() {
1278        let m = Metrics::new();
1279        m.record_budget_reconcile_failures(0); // a zero is a no-op
1280        m.record_budget_reconcile_failures(2);
1281        m.record_budget_reconcile_failures(1);
1282        assert!(
1283            m.render()
1284                .contains("edgeguard_llm_budget_reconcile_failures_total 3"),
1285            "{}",
1286            m.render()
1287        );
1288    }
1289
1290    #[test]
1291    fn per_key_tokens_and_cost_are_accumulated_and_rendered() {
1292        let m = Metrics::new();
1293        m.record_llm_key_usage(
1294            "key-abc",
1295            &LlmSample {
1296                tokens_in: 100,
1297                tokens_out: 40,
1298                cost_micros: Some(77),
1299                ..Default::default()
1300            },
1301        );
1302        // An unauthenticated request falls into the shared `_anon` bucket.
1303        m.record_llm_key_usage(
1304            "_anon",
1305            &LlmSample {
1306                tokens_in: 5,
1307                ..Default::default()
1308            },
1309        );
1310        let text = m.render();
1311        assert!(
1312            text.contains("edgeguard_llm_key_tokens_total{key=\"key-abc\",kind=\"input\"} 100"),
1313            "{text}"
1314        );
1315        assert!(
1316            text.contains("edgeguard_llm_key_cost_microdollars_total{key=\"key-abc\"} 77"),
1317            "{text}"
1318        );
1319        assert!(
1320            text.contains("edgeguard_llm_key_tokens_total{key=\"_anon\",kind=\"input\"} 5"),
1321            "{text}"
1322        );
1323    }
1324
1325    #[test]
1326    fn per_model_series_are_cardinality_bounded() {
1327        let m = Metrics::new();
1328        // Feed more distinct models than the cap; the overflow bucket absorbs the excess so the map
1329        // never grows past MAX_LLM_MODEL_SERIES + 1 (the overflow key).
1330        for i in 0..(MAX_LLM_MODEL_SERIES + 50) {
1331            m.record_llm_usage(
1332                &format!("model-{i}"),
1333                LlmSample {
1334                    tokens_in: 1,
1335                    ..Default::default()
1336                },
1337            );
1338        }
1339        let map = m.per_model.lock().unwrap();
1340        assert!(map.len() <= MAX_LLM_MODEL_SERIES + 1, "len={}", map.len());
1341        assert!(map.contains_key(LLM_MODEL_OVERFLOW));
1342    }
1343
1344    #[test]
1345    fn budget_blocked_and_consumed_metrics() {
1346        let m = Metrics::new();
1347        m.record_budget_blocked("key");
1348        m.record_budget_blocked("key");
1349        m.record_budget_blocked("team");
1350        m.record_budget_blocked("totally_unknown"); // -> "other"
1351        m.record_budget_consumed("daily-cap", 0.75);
1352        m.record_budget_consumed("daily-cap", 0.92); // last writer wins
1353        m.record_budget_consumed("nan-guard", f64::NAN); // dropped
1354        let text = m.render();
1355        assert!(
1356            text.contains("edgeguard_llm_budget_blocked_total{scope=\"key\"} 2"),
1357            "{text}"
1358        );
1359        assert!(
1360            text.contains("edgeguard_llm_budget_blocked_total{scope=\"team\"} 1"),
1361            "{text}"
1362        );
1363        assert!(
1364            text.contains("edgeguard_llm_budget_blocked_total{scope=\"other\"} 1"),
1365            "{text}"
1366        );
1367        assert!(
1368            text.contains("edgeguard_llm_budget_consumed_ratio{budget=\"daily-cap\"} 0.92"),
1369            "{text}"
1370        );
1371        assert!(
1372            !text.contains("nan-guard"),
1373            "NaN sample must be dropped: {text}"
1374        );
1375    }
1376
1377    #[test]
1378    fn label_values_are_escaped() {
1379        // A client-chosen model with a quote/newline must not break the exposition format.
1380        let m = Metrics::new();
1381        m.record_llm_usage(
1382            "evil\"\nmodel",
1383            LlmSample {
1384                tokens_in: 1,
1385                ..Default::default()
1386            },
1387        );
1388        let text = m.render();
1389        assert!(text.contains("model=\"evil\\\"\\nmodel\""), "{text}");
1390    }
1391
1392    #[test]
1393    fn dlp_finding_and_blocked_counters() {
1394        let m = Metrics::new();
1395        m.record_dlp_finding("email");
1396        m.record_dlp_finding("email");
1397        m.record_dlp_finding("api_key");
1398        m.record_dlp_finding("totally_unknown"); // -> "other"
1399        m.record_dlp_blocked();
1400        let text = m.render();
1401        assert!(
1402            text.contains("edgeguard_llm_dlp_findings_total{category=\"email\"} 2"),
1403            "{text}"
1404        );
1405        assert!(
1406            text.contains("edgeguard_llm_dlp_findings_total{category=\"api_key\"} 1"),
1407            "{text}"
1408        );
1409        assert!(
1410            text.contains("edgeguard_llm_dlp_findings_total{category=\"other\"} 1"),
1411            "{text}"
1412        );
1413        assert!(text.contains("edgeguard_llm_dlp_blocked_total 1"), "{text}");
1414    }
1415
1416    #[test]
1417    fn keyvault_result_counters() {
1418        let m = Metrics::new();
1419        m.record_keyvault("swapped");
1420        m.record_keyvault("swapped");
1421        m.record_keyvault("denied_model");
1422        m.record_keyvault("totally_unknown"); // ignored, not miscounted
1423        let text = m.render();
1424        assert!(
1425            text.contains("edgeguard_llm_keyvault_total{result=\"swapped\"} 2"),
1426            "{text}"
1427        );
1428        assert!(
1429            text.contains("edgeguard_llm_keyvault_total{result=\"denied_model\"} 1"),
1430            "{text}"
1431        );
1432        assert!(
1433            text.contains("edgeguard_llm_keyvault_total{result=\"denied_key\"} 0"),
1434            "{text}"
1435        );
1436    }
1437
1438    #[test]
1439    fn waf_hit_counters_by_class() {
1440        let m = Metrics::new();
1441        m.record_waf_hit("sqli");
1442        m.record_waf_hit("sqli");
1443        m.record_waf_hit("custom");
1444        // An unknown class is ignored rather than miscounted.
1445        m.record_waf_hit("totally_unknown");
1446        let text = m.render();
1447        assert!(
1448            text.contains("edgeguard_waf_hits_total{rule=\"sqli\"} 2"),
1449            "{text}"
1450        );
1451        assert!(
1452            text.contains("edgeguard_waf_hits_total{rule=\"custom\"} 1"),
1453            "{text}"
1454        );
1455        // A class that never fired still renders at 0.
1456        assert!(
1457            text.contains("edgeguard_waf_hits_total{rule=\"xss\"} 0"),
1458            "{text}"
1459        );
1460    }
1461}