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    /// One counter per [`OUTCOMES`] entry (parallel index).
177    requests: Vec<AtomicU64>,
178    /// One counter per [`RL_SCOPES`] entry (parallel index).
179    ratelimit_hits: Vec<AtomicU64>,
180    /// One counter per [`WAF_RULES`] entry (parallel index).
181    waf_hits: Vec<AtomicU64>,
182    /// Cumulative histogram buckets (parallel to [`LATENCY_BUCKETS`]): `bucket[i]` counts
183    /// observations with value <= `LATENCY_BUCKETS[i]`.
184    latency_buckets: Vec<AtomicU64>,
185    latency_sum_micros: AtomicU64,
186    latency_count: AtomicU64,
187    csp_reports: AtomicU64,
188    /// Drainable usage accumulators for managed-mode reporting (requests + bandwidth *since the
189    /// last drain*). Kept separate from the monotonic Prometheus counters above precisely because
190    /// the usage reporter resets these to zero each period — a Prometheus counter must not decrease.
191    usage_requests: AtomicU64,
192    usage_ingress_bytes: AtomicU64,
193    usage_egress_bytes: AtomicU64,
194    /// Drainable count of requests the edge *denied* (auth / WAF-forbidden / rate-limit / quota /
195    /// budget / unpriced-model) since the last drain, for the managed-mode usage report's "blocked"
196    /// figure. A subset of `usage_requests`. Distinct from the monotonic per-outcome counters.
197    usage_blocked: AtomicU64,
198    /// Drainable LLM token/cost usage for managed-mode cost reports (reset each report period,
199    /// like the request/byte accumulators above). Distinct from the monotonic `llm_*` counters.
200    usage_tokens_in: AtomicU64,
201    usage_tokens_out: AtomicU64,
202    usage_cost_micros: AtomicU64,
203    /// Drainable WAF matches by rule class (parallel to [`WAF_RULES`]), for the managed-mode usage
204    /// report's per-category security breakdown. Distinct from the monotonic `waf_hits` counters,
205    /// which the reporter must not reset.
206    usage_waf_hits: Vec<AtomicU64>,
207    /// LLM input (prompt) tokens metered (monotonic).
208    llm_tokens_in: AtomicU64,
209    /// LLM output (completion) tokens metered (monotonic).
210    llm_tokens_out: AtomicU64,
211    /// LLM cached prompt tokens (⊆ input), metered separately (monotonic). The dim the "~7× undercount"
212    /// gap is about — surfaced so cache utilisation and its cost impact are both visible.
213    llm_cached_tokens: AtomicU64,
214    /// LLM reasoning completion tokens (⊆ output), metered separately (monotonic).
215    llm_reasoning_tokens: AtomicU64,
216    /// Accumulated LLM cost in micro-dollars (1e-6 USD), for priced models only.
217    llm_cost_micros: AtomicU64,
218    /// Server-side time-to-first-token histogram (buckets parallel to [`LATENCY_BUCKETS`]) for
219    /// streamed LLM responses, measured in the gateway as frames flow — no client clock, no in-app
220    /// instrumentation. A trace backend only sees span-end duration, so this is a request-path-only
221    /// signal, and a server-measured one rather than a client-clock approximation.
222    llm_ttft_buckets: Vec<AtomicU64>,
223    llm_ttft_sum_micros: AtomicU64,
224    llm_ttft_count: AtomicU64,
225    /// Mean time-per-output-token histogram for streamed LLM responses with >1 output token
226    /// (inter-token latency = span of the output stream / (output_tokens − 1)).
227    llm_tpot_buckets: Vec<AtomicU64>,
228    llm_tpot_sum_micros: AtomicU64,
229    llm_tpot_count: AtomicU64,
230    /// One counter per [`LLM_RESULTS`] entry (parallel index).
231    llm_results: Vec<AtomicU64>,
232    /// Per-model token/cost breakdown (`edgeguard_llm_model_*`), bounded to [`MAX_LLM_MODEL_SERIES`]
233    /// distinct models (overflow → [`LLM_MODEL_OVERFLOW`]) so client-chosen model strings can't blow
234    /// up Prometheus cardinality.
235    per_model: Mutex<BTreeMap<String, PerModelCounters>>,
236    /// Per-team token/cost breakdown (`edgeguard_llm_team_*`), keyed by the `[llm].team_header` value
237    /// (absent → `_none`). Same cardinality bound + overflow bucket as `per_model`, so it answers
238    /// "which team spent this" for chargeback/showback without a Prometheus label explosion.
239    per_team: Mutex<BTreeMap<String, PerModelCounters>>,
240    /// Per-key token/cost breakdown (`edgeguard_llm_key_*`), keyed by the authenticated **principal**
241    /// (API-key id / Basic user / JWT `sub`; unauthenticated → `_anon`) — the OSS identity primitive.
242    /// Same cardinality bound + overflow bucket, so per-user attribution is reachable in the OSS core.
243    per_key: Mutex<BTreeMap<String, PerModelCounters>>,
244    /// One counter per [`BUDGET_SCOPES`] entry (parallel index): requests blocked by a hard budget.
245    budget_blocked: Vec<AtomicU64>,
246    /// Latest observed consumed ratio (`used / limit`, 0.0–1.0+) per budget *name* — the near-limit
247    /// signal. A coarse gauge (one value per budget name, last writer wins across scope keys), which
248    /// is what an "any budget near its cap" alert needs. Bounded by the operator-defined name set.
249    budget_consumed: Mutex<BTreeMap<String, f64>>,
250    /// One counter per [`KEYVAULT_RESULTS`] entry (parallel index).
251    keyvault_results: Vec<AtomicU64>,
252    /// One counter per [`DLP_CATEGORIES`] entry (parallel index): DLP findings by category.
253    dlp_findings: Vec<AtomicU64>,
254    /// Requests blocked (`403`) by DLP `block` mode.
255    dlp_blocked: AtomicU64,
256    /// Budget reconcile/release operations that ultimately FAILED (after retries) against the shared
257    /// store. Each failure means a reserve→settle didn't complete, so the distributed counter has
258    /// drifted (a leaked hold → phantom `BudgetExceededError`, or an uncharged settle → silent
259    /// bypass). Surfaced so this drift is **observable** instead of only logged — alert on it.
260    budget_reconcile_failures: AtomicU64,
261}
262
263impl Default for Metrics {
264    fn default() -> Self {
265        Metrics {
266            requests: OUTCOMES.iter().map(|_| AtomicU64::new(0)).collect(),
267            ratelimit_hits: RL_SCOPES.iter().map(|_| AtomicU64::new(0)).collect(),
268            waf_hits: WAF_RULES.iter().map(|_| AtomicU64::new(0)).collect(),
269            latency_buckets: LATENCY_BUCKETS.iter().map(|_| AtomicU64::new(0)).collect(),
270            latency_sum_micros: AtomicU64::new(0),
271            latency_count: AtomicU64::new(0),
272            csp_reports: AtomicU64::new(0),
273            usage_requests: AtomicU64::new(0),
274            usage_ingress_bytes: AtomicU64::new(0),
275            usage_egress_bytes: AtomicU64::new(0),
276            usage_blocked: AtomicU64::new(0),
277            usage_tokens_in: AtomicU64::new(0),
278            usage_tokens_out: AtomicU64::new(0),
279            usage_cost_micros: AtomicU64::new(0),
280            usage_waf_hits: WAF_RULES.iter().map(|_| AtomicU64::new(0)).collect(),
281            llm_tokens_in: AtomicU64::new(0),
282            llm_tokens_out: AtomicU64::new(0),
283            llm_cached_tokens: AtomicU64::new(0),
284            llm_reasoning_tokens: AtomicU64::new(0),
285            llm_cost_micros: AtomicU64::new(0),
286            llm_ttft_buckets: LATENCY_BUCKETS.iter().map(|_| AtomicU64::new(0)).collect(),
287            llm_ttft_sum_micros: AtomicU64::new(0),
288            llm_ttft_count: AtomicU64::new(0),
289            llm_tpot_buckets: LATENCY_BUCKETS.iter().map(|_| AtomicU64::new(0)).collect(),
290            llm_tpot_sum_micros: AtomicU64::new(0),
291            llm_tpot_count: AtomicU64::new(0),
292            llm_results: LLM_RESULTS.iter().map(|_| AtomicU64::new(0)).collect(),
293            per_model: Mutex::new(BTreeMap::new()),
294            per_team: Mutex::new(BTreeMap::new()),
295            per_key: Mutex::new(BTreeMap::new()),
296            budget_blocked: BUDGET_SCOPES.iter().map(|_| AtomicU64::new(0)).collect(),
297            budget_consumed: Mutex::new(BTreeMap::new()),
298            keyvault_results: KEYVAULT_RESULTS.iter().map(|_| AtomicU64::new(0)).collect(),
299            dlp_findings: DLP_CATEGORIES.iter().map(|_| AtomicU64::new(0)).collect(),
300            dlp_blocked: AtomicU64::new(0),
301            budget_reconcile_failures: AtomicU64::new(0),
302        }
303    }
304}
305
306/// Observe `elapsed` into a cumulative histogram (buckets parallel to [`LATENCY_BUCKETS`]) plus its
307/// running sum (micros) and count. Shared by the request-latency and the LLM TTFT/TPOT histograms so
308/// their bucketing can't drift.
309fn observe_hist(
310    buckets: &[AtomicU64],
311    sum_micros: &AtomicU64,
312    count: &AtomicU64,
313    elapsed: Duration,
314) {
315    let secs = elapsed.as_secs_f64();
316    for (i, bound) in LATENCY_BUCKETS.iter().enumerate() {
317        if secs <= *bound {
318            buckets[i].fetch_add(1, Ordering::Relaxed);
319        }
320    }
321    sum_micros.fetch_add(elapsed.as_micros() as u64, Ordering::Relaxed);
322    count.fetch_add(1, Ordering::Relaxed);
323}
324
325/// Render one Prometheus histogram (`_bucket`/`_sum`/`_count`) to `out`, sharing the exposition
326/// shape with the request-latency histogram above.
327fn render_hist(
328    out: &mut String,
329    name: &str,
330    help: &str,
331    buckets: &[AtomicU64],
332    sum_micros: &AtomicU64,
333    count: &AtomicU64,
334) {
335    out.push_str(&format!("# HELP {name} {help}\n"));
336    out.push_str(&format!("# TYPE {name} histogram\n"));
337    for (i, bound) in LATENCY_BUCKETS.iter().enumerate() {
338        let v = buckets[i].load(Ordering::Relaxed);
339        out.push_str(&format!("{name}_bucket{{le=\"{bound}\"}} {v}\n"));
340    }
341    // The `+Inf` bucket equals the total observation count by definition.
342    let c = count.load(Ordering::Relaxed);
343    out.push_str(&format!("{name}_bucket{{le=\"+Inf\"}} {c}\n"));
344    let sum_secs = sum_micros.load(Ordering::Relaxed) as f64 / 1_000_000.0;
345    out.push_str(&format!("{name}_sum {sum_secs}\n"));
346    out.push_str(&format!("{name}_count {c}\n"));
347}
348
349/// Render one per-dimension token/cost breakdown (`per_model` / `per_team` / `per_key`) to `out`,
350/// sharing the exposition shape across all three so they can't drift from one another.
351fn render_breakdown(
352    out: &mut String,
353    metric_prefix: &str,
354    label_name: &str,
355    dim_desc: &str,
356    map: &BTreeMap<String, PerModelCounters>,
357) {
358    out.push_str(&format!(
359        "# HELP {metric_prefix}_tokens_total LLM tokens metered by {dim_desc} and kind.\n"
360    ));
361    out.push_str(&format!("# TYPE {metric_prefix}_tokens_total counter\n"));
362    for (k, c) in map.iter() {
363        let label = escape_label(k);
364        for kind in LLM_TOKEN_KINDS {
365            let v = match *kind {
366                "input" => c.tokens_in,
367                "output" => c.tokens_out,
368                "cached" => c.cached_tokens,
369                "reasoning" => c.reasoning_tokens,
370                _ => 0,
371            };
372            out.push_str(&format!(
373                "{metric_prefix}_tokens_total{{{label_name}=\"{label}\",kind=\"{kind}\"}} {v}\n"
374            ));
375        }
376    }
377    out.push_str(&format!(
378        "# HELP {metric_prefix}_cost_microdollars_total LLM cost (micro-dollars) by {dim_desc}.\n"
379    ));
380    out.push_str(&format!(
381        "# TYPE {metric_prefix}_cost_microdollars_total counter\n"
382    ));
383    for (k, c) in map.iter() {
384        out.push_str(&format!(
385            "{metric_prefix}_cost_microdollars_total{{{label_name}=\"{}\"}} {}\n",
386            escape_label(k),
387            c.cost_micros
388        ));
389    }
390}
391
392/// Add `s` to the `key` bucket of a bounded per-dimension accumulator (`per_model` / `per_team`): a
393/// new key is inserted only while under [`MAX_LLM_MODEL_SERIES`]; past the cap it folds into
394/// [`LLM_MODEL_OVERFLOW`], so a flood of distinct label values can't grow the map without bound.
395fn accumulate_bounded(map: &Mutex<BTreeMap<String, PerModelCounters>>, key: &str, s: &LlmSample) {
396    let mut map = map.lock().expect("per-dimension mutex poisoned");
397    let k = if map.contains_key(key) || map.len() < MAX_LLM_MODEL_SERIES {
398        key
399    } else {
400        LLM_MODEL_OVERFLOW
401    };
402    let c = map.entry(k.to_string()).or_default();
403    c.tokens_in = c.tokens_in.saturating_add(s.tokens_in);
404    c.tokens_out = c.tokens_out.saturating_add(s.tokens_out);
405    c.cached_tokens = c.cached_tokens.saturating_add(s.cached_tokens);
406    c.reasoning_tokens = c.reasoning_tokens.saturating_add(s.reasoning_tokens);
407    c.cost_micros = c.cost_micros.saturating_add(s.cost_micros.unwrap_or(0));
408}
409
410impl Metrics {
411    pub fn new() -> Self {
412        Self::default()
413    }
414
415    /// Count one finished request under its `outcome` label.
416    pub fn record_request(&self, outcome: &str) {
417        let idx = OUTCOMES
418            .iter()
419            .position(|o| *o == outcome)
420            .unwrap_or(OUTCOMES.len() - 1); // -> "other"
421        self.requests[idx].fetch_add(1, Ordering::Relaxed);
422    }
423
424    /// Observe a request's end-to-end latency into the histogram.
425    pub fn observe_latency(&self, elapsed: Duration) {
426        observe_hist(
427            &self.latency_buckets,
428            &self.latency_sum_micros,
429            &self.latency_count,
430            elapsed,
431        );
432    }
433
434    /// Observe a streamed LLM response's server-side **time-to-first-token** and, when the response
435    /// had more than one output token, its mean **time-per-output-token**. Called once per streamed
436    /// LLM request from the response body's `Drop`, after the terminal `usage` frame is parsed. The
437    /// gateway sits in the token stream, so these are measured with no client clock and no in-app
438    /// instrumentation — the request-path advantage a trace backend (which only sees span-end
439    /// duration) can't offer.
440    pub fn record_llm_latency(&self, ttft: Duration, tpot: Option<Duration>) {
441        observe_hist(
442            &self.llm_ttft_buckets,
443            &self.llm_ttft_sum_micros,
444            &self.llm_ttft_count,
445            ttft,
446        );
447        if let Some(tpot) = tpot {
448            observe_hist(
449                &self.llm_tpot_buckets,
450                &self.llm_tpot_sum_micros,
451                &self.llm_tpot_count,
452                tpot,
453            );
454        }
455    }
456
457    /// Count a rate-limit rejection by which limiter scope tripped (`ip`/`route`/`key`).
458    pub fn record_ratelimit_hit(&self, scope: &str) {
459        if let Some(idx) = RL_SCOPES.iter().position(|s| *s == scope) {
460            self.ratelimit_hits[idx].fetch_add(1, Ordering::Relaxed);
461        }
462    }
463
464    /// Count one WAF rule match by rule class (`sqli`/`xss`/`path_traversal`/`custom`).
465    /// Recorded for both report-only and blocking modes — so a report-first rollout is
466    /// visible — while a *blocked* request is additionally counted under the `forbidden`
467    /// request outcome.
468    pub fn record_waf_hit(&self, class: &str) {
469        if let Some(idx) = WAF_RULES.iter().position(|c| *c == class) {
470            self.waf_hits[idx].fetch_add(1, Ordering::Relaxed);
471            // Parallel drainable accumulator for the managed-mode usage report.
472            self.usage_waf_hits[idx].fetch_add(1, Ordering::Relaxed);
473        }
474    }
475
476    /// Count one received CSP violation report.
477    pub fn record_csp_report(&self) {
478        self.csp_reports.fetch_add(1, Ordering::Relaxed);
479    }
480
481    /// Count one request toward the drainable usage accumulator (managed mode). Called once per
482    /// request from the single `finish` exit, so every request — proxied or rejected — counts.
483    /// `outcome` is the request's outcome label; a denial outcome (see [`outcome_is_blocked`]) also
484    /// bumps the drainable `blocked` accumulator, so the control plane can show what the edge screened.
485    pub fn add_usage_request(&self, outcome: &str) {
486        self.usage_requests.fetch_add(1, Ordering::Relaxed);
487        if outcome_is_blocked(outcome) {
488            self.usage_blocked.fetch_add(1, Ordering::Relaxed);
489        }
490    }
491
492    /// Add request (ingress) + response (egress) bytes to the drainable usage accumulator. Called
493    /// on the proxied path where both bodies are buffered and the counts are known.
494    pub fn add_usage_bytes(&self, ingress: usize, egress: usize) {
495        self.usage_ingress_bytes
496            .fetch_add(ingress as u64, Ordering::Relaxed);
497        self.usage_egress_bytes
498            .fetch_add(egress as u64, Ordering::Relaxed);
499    }
500
501    /// Atomically read-and-zero the usage accumulators — the delta the usage reporter ships to the
502    /// control plane (requests + bandwidth + LLM tokens/cost, gateway L4).
503    pub fn drain_usage(&self) -> DrainedUsage {
504        debug_assert_eq!(WAF_RULES, ["sqli", "xss", "path_traversal", "custom"]);
505        DrainedUsage {
506            requests: self.usage_requests.swap(0, Ordering::Relaxed),
507            ingress_bytes: self.usage_ingress_bytes.swap(0, Ordering::Relaxed),
508            egress_bytes: self.usage_egress_bytes.swap(0, Ordering::Relaxed),
509            tokens_in: self.usage_tokens_in.swap(0, Ordering::Relaxed),
510            tokens_out: self.usage_tokens_out.swap(0, Ordering::Relaxed),
511            cost_micros: self.usage_cost_micros.swap(0, Ordering::Relaxed),
512            blocked: self.usage_blocked.swap(0, Ordering::Relaxed),
513            // Indices parallel to WAF_RULES = [sqli, xss, path_traversal, custom].
514            waf_sqli: self.usage_waf_hits[0].swap(0, Ordering::Relaxed),
515            waf_xss: self.usage_waf_hits[1].swap(0, Ordering::Relaxed),
516            waf_path_traversal: self.usage_waf_hits[2].swap(0, Ordering::Relaxed),
517            waf_custom: self.usage_waf_hits[3].swap(0, Ordering::Relaxed),
518        }
519    }
520
521    /// Add a previously-drained delta back, e.g. when a usage report failed to send — so the
522    /// next period reships it instead of losing billable usage. (New requests that arrived during
523    /// the failed send simply add on top, as intended.)
524    pub fn restore_usage(&self, u: &DrainedUsage) {
525        debug_assert_eq!(WAF_RULES, ["sqli", "xss", "path_traversal", "custom"]);
526        self.usage_requests.fetch_add(u.requests, Ordering::Relaxed);
527        self.usage_ingress_bytes
528            .fetch_add(u.ingress_bytes, Ordering::Relaxed);
529        self.usage_egress_bytes
530            .fetch_add(u.egress_bytes, Ordering::Relaxed);
531        self.usage_tokens_in
532            .fetch_add(u.tokens_in, Ordering::Relaxed);
533        self.usage_tokens_out
534            .fetch_add(u.tokens_out, Ordering::Relaxed);
535        self.usage_cost_micros
536            .fetch_add(u.cost_micros, Ordering::Relaxed);
537        self.usage_blocked.fetch_add(u.blocked, Ordering::Relaxed);
538        // Indices parallel to WAF_RULES = [sqli, xss, path_traversal, custom].
539        self.usage_waf_hits[0].fetch_add(u.waf_sqli, Ordering::Relaxed);
540        self.usage_waf_hits[1].fetch_add(u.waf_xss, Ordering::Relaxed);
541        self.usage_waf_hits[2].fetch_add(u.waf_path_traversal, Ordering::Relaxed);
542        self.usage_waf_hits[3].fetch_add(u.waf_custom, Ordering::Relaxed);
543    }
544
545    /// Record one metered LLM request for `model`: add its four token dimensions and — when the model
546    /// was priced — its cost (micro-dollars). `cost_micros == None` means the model isn't in the price
547    /// book, so tokens are still counted but the request is bucketed `unpriced` rather than `metered`.
548    /// Also updates the bounded per-model breakdown (`edgeguard_llm_model_*`).
549    pub fn record_llm_usage(&self, model: &str, s: LlmSample) {
550        self.llm_tokens_in.fetch_add(s.tokens_in, Ordering::Relaxed);
551        self.llm_tokens_out
552            .fetch_add(s.tokens_out, Ordering::Relaxed);
553        self.llm_cached_tokens
554            .fetch_add(s.cached_tokens, Ordering::Relaxed);
555        self.llm_reasoning_tokens
556            .fetch_add(s.reasoning_tokens, Ordering::Relaxed);
557        // Drainable accumulators for the managed-mode cost report.
558        self.usage_tokens_in
559            .fetch_add(s.tokens_in, Ordering::Relaxed);
560        self.usage_tokens_out
561            .fetch_add(s.tokens_out, Ordering::Relaxed);
562        let result = match s.cost_micros {
563            Some(c) => {
564                self.llm_cost_micros.fetch_add(c, Ordering::Relaxed);
565                self.usage_cost_micros.fetch_add(c, Ordering::Relaxed);
566                "metered"
567            }
568            None => "unpriced",
569        };
570        self.bump_llm_result(result);
571        self.record_per_model(model, &s);
572    }
573
574    /// Update the bounded per-model accumulator. A new model is only inserted while under the cap;
575    /// once full it folds into [`LLM_MODEL_OVERFLOW`], so a flood of distinct model strings can't grow
576    /// the map without bound.
577    fn record_per_model(&self, model: &str, s: &LlmSample) {
578        accumulate_bounded(&self.per_model, model, s);
579    }
580
581    /// Record one metered LLM request against its team (`[llm].team_header` value; absent → `_none`),
582    /// for per-team chargeback/showback (`edgeguard_llm_team_*`). Bounded exactly like the per-model
583    /// breakdown. Called alongside [`Self::record_llm_usage`] from the request path.
584    pub fn record_llm_team_usage(&self, team: &str, s: &LlmSample) {
585        accumulate_bounded(&self.per_team, team, s);
586    }
587
588    /// Record one metered LLM request against its authenticated key/principal (`_anon` when
589    /// unauthenticated), for per-user cost attribution (`edgeguard_llm_key_*`). Bounded exactly like
590    /// the per-model breakdown. Reuses the existing OSS auth principal as the identity — so per-key
591    /// FinOps is reachable without the EE control plane.
592    pub fn record_llm_key_usage(&self, key: &str, s: &LlmSample) {
593        accumulate_bounded(&self.per_key, key, s);
594    }
595
596    /// Record a request blocked by a hard LLM budget, by the budget's `scope` (unknown → `other`).
597    pub fn record_budget_blocked(&self, scope: &str) {
598        let idx = BUDGET_SCOPES
599            .iter()
600            .position(|s| *s == scope)
601            .unwrap_or(BUDGET_SCOPES.len() - 1); // -> "other"
602        self.budget_blocked[idx].fetch_add(1, Ordering::Relaxed);
603    }
604
605    /// Record the latest consumed ratio (`used / limit`) for a budget by `name` — the near-limit
606    /// gauge. Last writer wins per name (a coarse "is any budget near its cap" signal). NaN/negative
607    /// samples are dropped so a divide-by-zero can't poison the gauge.
608    pub fn record_budget_consumed(&self, name: &str, ratio: f64) {
609        if !ratio.is_finite() || ratio < 0.0 {
610            return;
611        }
612        let mut map = self
613            .budget_consumed
614            .lock()
615            .expect("budget_consumed mutex poisoned");
616        // Bound the map the same way as models: operator-defined names are few, but never grow past
617        // the cap if a config churns budget names.
618        if map.contains_key(name) || map.len() < MAX_LLM_MODEL_SERIES {
619            map.insert(name.to_string(), ratio);
620        }
621    }
622
623    /// Record `n` budget reconcile/release failures against the shared store (after retries). A
624    /// non-zero rate here means the distributed budget counter is drifting — the signal to alert on.
625    pub fn record_budget_reconcile_failures(&self, n: usize) {
626        if n > 0 {
627            self.budget_reconcile_failures
628                .fetch_add(n as u64, Ordering::Relaxed);
629        }
630    }
631
632    /// Record an LLM request whose response carried no usage (error, or a stream the client didn't
633    /// opt into usage on). No tokens/cost, but the request is visible as `no_usage`.
634    pub fn record_llm_no_usage(&self) {
635        self.bump_llm_result("no_usage");
636    }
637
638    fn bump_llm_result(&self, result: &str) {
639        if let Some(idx) = LLM_RESULTS.iter().position(|r| *r == result) {
640            self.llm_results[idx].fetch_add(1, Ordering::Relaxed);
641        }
642    }
643
644    /// Record a key-vault decision by `result` (`swapped`/`denied_key`/`denied_model`).
645    pub fn record_keyvault(&self, result: &str) {
646        if let Some(idx) = KEYVAULT_RESULTS.iter().position(|r| *r == result) {
647            self.keyvault_results[idx].fetch_add(1, Ordering::Relaxed);
648        }
649    }
650
651    /// Record one DLP finding under its `category` (unknown → `other`).
652    pub fn record_dlp_finding(&self, category: &str) {
653        let idx = DLP_CATEGORIES
654            .iter()
655            .position(|c| *c == category)
656            .unwrap_or(DLP_CATEGORIES.len() - 1); // -> "other"
657        self.dlp_findings[idx].fetch_add(1, Ordering::Relaxed);
658    }
659
660    /// Record a request blocked by DLP `block` mode.
661    pub fn record_dlp_blocked(&self) {
662        self.dlp_blocked.fetch_add(1, Ordering::Relaxed);
663    }
664
665    /// Render the Prometheus text exposition (format version 0.0.4).
666    pub fn render(&self) -> String {
667        let mut out = String::with_capacity(1024);
668
669        out.push_str("# HELP edgeguard_requests_total Total proxied requests by outcome.\n");
670        out.push_str("# TYPE edgeguard_requests_total counter\n");
671        for (i, label) in OUTCOMES.iter().enumerate() {
672            let v = self.requests[i].load(Ordering::Relaxed);
673            out.push_str(&format!(
674                "edgeguard_requests_total{{outcome=\"{label}\"}} {v}\n"
675            ));
676        }
677
678        out.push_str(
679            "# HELP edgeguard_ratelimit_hits_total Requests rejected by a rate limiter, by scope.\n",
680        );
681        out.push_str("# TYPE edgeguard_ratelimit_hits_total counter\n");
682        for (i, label) in RL_SCOPES.iter().enumerate() {
683            let v = self.ratelimit_hits[i].load(Ordering::Relaxed);
684            out.push_str(&format!(
685                "edgeguard_ratelimit_hits_total{{scope=\"{label}\"}} {v}\n"
686            ));
687        }
688
689        out.push_str(
690            "# HELP edgeguard_waf_hits_total WAF rule matches by class (report-only + blocked).\n",
691        );
692        out.push_str("# TYPE edgeguard_waf_hits_total counter\n");
693        for (i, label) in WAF_RULES.iter().enumerate() {
694            let v = self.waf_hits[i].load(Ordering::Relaxed);
695            out.push_str(&format!(
696                "edgeguard_waf_hits_total{{rule=\"{label}\"}} {v}\n"
697            ));
698        }
699
700        out.push_str("# HELP edgeguard_csp_reports_total CSP violation reports received.\n");
701        out.push_str("# TYPE edgeguard_csp_reports_total counter\n");
702        out.push_str(&format!(
703            "edgeguard_csp_reports_total {}\n",
704            self.csp_reports.load(Ordering::Relaxed)
705        ));
706
707        out.push_str(
708            "# HELP edgeguard_request_duration_seconds Request handling latency in seconds.\n",
709        );
710        out.push_str("# TYPE edgeguard_request_duration_seconds histogram\n");
711        for (i, bound) in LATENCY_BUCKETS.iter().enumerate() {
712            let v = self.latency_buckets[i].load(Ordering::Relaxed);
713            out.push_str(&format!(
714                "edgeguard_request_duration_seconds_bucket{{le=\"{bound}\"}} {v}\n"
715            ));
716        }
717        let count = self.latency_count.load(Ordering::Relaxed);
718        // The `+Inf` bucket equals the total observation count by definition.
719        out.push_str(&format!(
720            "edgeguard_request_duration_seconds_bucket{{le=\"+Inf\"}} {count}\n"
721        ));
722        let sum_secs = self.latency_sum_micros.load(Ordering::Relaxed) as f64 / 1_000_000.0;
723        out.push_str(&format!(
724            "edgeguard_request_duration_seconds_sum {sum_secs}\n"
725        ));
726        out.push_str(&format!(
727            "edgeguard_request_duration_seconds_count {count}\n"
728        ));
729
730        // Streamed-LLM server-side latency: time-to-first-token and mean time-per-output-token,
731        // measured in the request path (the signal a trace backend can't produce). Render at 0 too.
732        render_hist(
733            &mut out,
734            "edgeguard_llm_ttft_seconds",
735            "Server-side time-to-first-token for streamed LLM responses, in seconds.",
736            &self.llm_ttft_buckets,
737            &self.llm_ttft_sum_micros,
738            &self.llm_ttft_count,
739        );
740        render_hist(
741            &mut out,
742            "edgeguard_llm_tpot_seconds",
743            "Mean time-per-output-token for streamed LLM responses (>1 output token), in seconds.",
744            &self.llm_tpot_buckets,
745            &self.llm_tpot_sum_micros,
746            &self.llm_tpot_count,
747        );
748
749        // LLM token metering (gateway L0). All series render even at 0 so dashboards/alerts don't
750        // break on a quiet proxy.
751        out.push_str("# HELP edgeguard_llm_tokens_total LLM tokens metered, by direction.\n");
752        out.push_str("# TYPE edgeguard_llm_tokens_total counter\n");
753        out.push_str(&format!(
754            "edgeguard_llm_tokens_total{{direction=\"input\"}} {}\n",
755            self.llm_tokens_in.load(Ordering::Relaxed)
756        ));
757        out.push_str(&format!(
758            "edgeguard_llm_tokens_total{{direction=\"output\"}} {}\n",
759            self.llm_tokens_out.load(Ordering::Relaxed)
760        ));
761
762        // Cached / reasoning sub-dimensions (⊆ input / ⊆ output). Kept as their own metric so they
763        // are visible for the "~7× undercount" story without being double-summed into the direction
764        // totals above.
765        out.push_str(
766            "# HELP edgeguard_llm_cached_tokens_total Cached prompt tokens metered (subset of input).\n",
767        );
768        out.push_str("# TYPE edgeguard_llm_cached_tokens_total counter\n");
769        out.push_str(&format!(
770            "edgeguard_llm_cached_tokens_total {}\n",
771            self.llm_cached_tokens.load(Ordering::Relaxed)
772        ));
773        out.push_str(
774            "# HELP edgeguard_llm_reasoning_tokens_total Reasoning completion tokens metered (subset of output).\n",
775        );
776        out.push_str("# TYPE edgeguard_llm_reasoning_tokens_total counter\n");
777        out.push_str(&format!(
778            "edgeguard_llm_reasoning_tokens_total {}\n",
779            self.llm_reasoning_tokens.load(Ordering::Relaxed)
780        ));
781
782        out.push_str(
783            "# HELP edgeguard_llm_cost_microdollars_total Accumulated LLM cost in micro-dollars (1e-6 USD).\n",
784        );
785        out.push_str("# TYPE edgeguard_llm_cost_microdollars_total counter\n");
786        out.push_str(&format!(
787            "edgeguard_llm_cost_microdollars_total {}\n",
788            self.llm_cost_micros.load(Ordering::Relaxed)
789        ));
790
791        // Per-model breakdown (bounded cardinality). Tokens carry a `kind` label; cost is a separate
792        // series. Rendered only for models actually seen, so a fresh proxy emits nothing here.
793        {
794            let map = self.per_model.lock().expect("per_model mutex poisoned");
795            render_breakdown(&mut out, "edgeguard_llm_model", "model", "model", &map);
796        }
797
798        // Per-team token/cost breakdown (`edgeguard_llm_team_*`), for chargeback/showback. Same
799        // cardinality bound as per-model; rendered only for teams actually seen.
800        {
801            let map = self.per_team.lock().expect("per_team mutex poisoned");
802            render_breakdown(&mut out, "edgeguard_llm_team", "team", "team", &map);
803        }
804
805        // Per-key (per-principal) token/cost breakdown (`edgeguard_llm_key_*`), for per-user FinOps.
806        // Same cardinality bound as per-model; rendered only for keys actually seen.
807        {
808            let map = self.per_key.lock().expect("per_key mutex poisoned");
809            render_breakdown(
810                &mut out,
811                "edgeguard_llm_key",
812                "key",
813                "authenticated key/principal",
814                &map,
815            );
816        }
817
818        // Hard-budget signals (gateway L1): the near-limit gauge and per-scope block counter.
819        out.push_str(
820            "# HELP edgeguard_llm_budget_blocked_total Requests blocked by a hard LLM budget, by scope.\n",
821        );
822        out.push_str("# TYPE edgeguard_llm_budget_blocked_total counter\n");
823        for (i, label) in BUDGET_SCOPES.iter().enumerate() {
824            let v = self.budget_blocked[i].load(Ordering::Relaxed);
825            out.push_str(&format!(
826                "edgeguard_llm_budget_blocked_total{{scope=\"{label}\"}} {v}\n"
827            ));
828        }
829        {
830            let map = self
831                .budget_consumed
832                .lock()
833                .expect("budget_consumed mutex poisoned");
834            out.push_str(
835                "# HELP edgeguard_llm_budget_consumed_ratio Latest consumed ratio (used/limit) per budget.\n",
836            );
837            out.push_str("# TYPE edgeguard_llm_budget_consumed_ratio gauge\n");
838            for (name, ratio) in map.iter() {
839                out.push_str(&format!(
840                    "edgeguard_llm_budget_consumed_ratio{{budget=\"{}\"}} {ratio}\n",
841                    escape_label(name)
842                ));
843            }
844        }
845        out.push_str(
846            "# HELP edgeguard_llm_budget_reconcile_failures_total Budget reserve->settle reconciles that failed against the shared store (counter drift).\n",
847        );
848        out.push_str("# TYPE edgeguard_llm_budget_reconcile_failures_total counter\n");
849        out.push_str(&format!(
850            "edgeguard_llm_budget_reconcile_failures_total {}\n",
851            self.budget_reconcile_failures.load(Ordering::Relaxed)
852        ));
853
854        out.push_str("# HELP edgeguard_llm_requests_total LLM requests metered, by result.\n");
855        out.push_str("# TYPE edgeguard_llm_requests_total counter\n");
856        for (i, label) in LLM_RESULTS.iter().enumerate() {
857            let v = self.llm_results[i].load(Ordering::Relaxed);
858            out.push_str(&format!(
859                "edgeguard_llm_requests_total{{result=\"{label}\"}} {v}\n"
860            ));
861        }
862
863        out.push_str(
864            "# HELP edgeguard_llm_keyvault_total Key-vault decisions by result (swap / egress denial).\n",
865        );
866        out.push_str("# TYPE edgeguard_llm_keyvault_total counter\n");
867        for (i, label) in KEYVAULT_RESULTS.iter().enumerate() {
868            let v = self.keyvault_results[i].load(Ordering::Relaxed);
869            out.push_str(&format!(
870                "edgeguard_llm_keyvault_total{{result=\"{label}\"}} {v}\n"
871            ));
872        }
873
874        out.push_str(
875            "# HELP edgeguard_llm_dlp_findings_total DLP findings (PII / secrets) by category.\n",
876        );
877        out.push_str("# TYPE edgeguard_llm_dlp_findings_total counter\n");
878        for (i, label) in DLP_CATEGORIES.iter().enumerate() {
879            let v = self.dlp_findings[i].load(Ordering::Relaxed);
880            out.push_str(&format!(
881                "edgeguard_llm_dlp_findings_total{{category=\"{label}\"}} {v}\n"
882            ));
883        }
884        out.push_str(
885            "# HELP edgeguard_llm_dlp_blocked_total Requests blocked by DLP block mode.\n",
886        );
887        out.push_str("# TYPE edgeguard_llm_dlp_blocked_total counter\n");
888        out.push_str(&format!(
889            "edgeguard_llm_dlp_blocked_total {}\n",
890            self.dlp_blocked.load(Ordering::Relaxed)
891        ));
892
893        out
894    }
895}
896
897/// Escape a dynamic label *value* for the Prometheus text format: backslash, double-quote, and
898/// newline must be escaped (per the exposition spec) so a client-chosen `model` or operator-chosen
899/// budget name can't inject a line break or unbalanced quote into the output.
900fn escape_label(s: &str) -> String {
901    let mut out = String::with_capacity(s.len());
902    for ch in s.chars() {
903        match ch {
904            '\\' => out.push_str("\\\\"),
905            '"' => out.push_str("\\\""),
906            '\n' => out.push_str("\\n"),
907            _ => out.push(ch),
908        }
909    }
910    out
911}
912
913#[cfg(test)]
914mod tests {
915    use super::*;
916
917    #[test]
918    fn records_and_renders_request_outcomes() {
919        let m = Metrics::new();
920        m.record_request("ok");
921        m.record_request("ok");
922        m.record_request("rate_limited");
923        // An unknown outcome falls into the `other` bucket, not "ok".
924        m.record_request("totally_unknown");
925
926        let text = m.render();
927        assert!(
928            text.contains("edgeguard_requests_total{outcome=\"ok\"} 2"),
929            "{text}"
930        );
931        assert!(
932            text.contains("edgeguard_requests_total{outcome=\"rate_limited\"} 1"),
933            "{text}"
934        );
935        assert!(
936            text.contains("edgeguard_requests_total{outcome=\"other\"} 1"),
937            "{text}"
938        );
939    }
940
941    #[test]
942    fn latency_histogram_is_cumulative() {
943        let m = Metrics::new();
944        m.observe_latency(Duration::from_millis(3)); // <= 0.005
945        m.observe_latency(Duration::from_millis(40)); // <= 0.05
946        let text = m.render();
947        // 3ms falls under every bucket >= 0.005; 40ms under every bucket >= 0.05.
948        assert!(
949            text.contains("edgeguard_request_duration_seconds_bucket{le=\"0.005\"} 1"),
950            "{text}"
951        );
952        assert!(
953            text.contains("edgeguard_request_duration_seconds_bucket{le=\"0.05\"} 2"),
954            "{text}"
955        );
956        assert!(
957            text.contains("edgeguard_request_duration_seconds_bucket{le=\"+Inf\"} 2"),
958            "{text}"
959        );
960        assert!(
961            text.contains("edgeguard_request_duration_seconds_count 2"),
962            "{text}"
963        );
964    }
965
966    #[test]
967    fn llm_ttft_tpot_histograms_render() {
968        let m = Metrics::new();
969        // TTFT 40ms (<= 0.05), TPOT 8ms (<= 0.01).
970        m.record_llm_latency(Duration::from_millis(40), Some(Duration::from_millis(8)));
971        // A single-output-token response has no defined TPOT — only TTFT is recorded.
972        m.record_llm_latency(Duration::from_millis(3), None);
973        let text = m.render();
974        // Two TTFT observations; both <= 0.05, one <= 0.005.
975        assert!(
976            text.contains("edgeguard_llm_ttft_seconds_bucket{le=\"0.005\"} 1"),
977            "{text}"
978        );
979        assert!(
980            text.contains("edgeguard_llm_ttft_seconds_bucket{le=\"0.05\"} 2"),
981            "{text}"
982        );
983        assert!(
984            text.contains("edgeguard_llm_ttft_seconds_count 2"),
985            "{text}"
986        );
987        // One TPOT observation (8ms), recorded only for the >1-token response.
988        assert!(
989            text.contains("edgeguard_llm_tpot_seconds_bucket{le=\"0.01\"} 1"),
990            "{text}"
991        );
992        assert!(
993            text.contains("edgeguard_llm_tpot_seconds_count 1"),
994            "{text}"
995        );
996    }
997
998    #[test]
999    fn ratelimit_and_csp_counters() {
1000        let m = Metrics::new();
1001        m.record_ratelimit_hit("ip");
1002        m.record_ratelimit_hit("route");
1003        m.record_ratelimit_hit("route");
1004        m.record_csp_report();
1005        let text = m.render();
1006        assert!(
1007            text.contains("edgeguard_ratelimit_hits_total{scope=\"ip\"} 1"),
1008            "{text}"
1009        );
1010        assert!(
1011            text.contains("edgeguard_ratelimit_hits_total{scope=\"route\"} 2"),
1012            "{text}"
1013        );
1014        assert!(text.contains("edgeguard_csp_reports_total 1"), "{text}");
1015    }
1016
1017    #[test]
1018    fn usage_accumulates_drains_and_restores() {
1019        let m = Metrics::new();
1020        m.add_usage_request("ok"); // proxied — not blocked
1021        m.add_usage_request("forbidden"); // edge-denied — also counts toward `blocked`
1022        m.add_usage_bytes(100, 250);
1023        m.add_usage_bytes(0, 50);
1024        // LLM token usage also drains for the cost report (gateway L4).
1025        m.record_llm_usage(
1026            "gpt-4o",
1027            LlmSample {
1028                tokens_in: 1_000,
1029                tokens_out: 400,
1030                cost_micros: Some(2_500),
1031                ..Default::default()
1032            },
1033        );
1034        // Drain returns the accrued delta and zeroes the accumulator.
1035        let drained = m.drain_usage();
1036        assert_eq!(drained.requests, 2);
1037        assert_eq!(drained.blocked, 1); // only the "forbidden" request
1038        assert_eq!(drained.ingress_bytes, 100);
1039        assert_eq!(drained.egress_bytes, 300);
1040        assert_eq!(drained.tokens_in, 1_000);
1041        assert_eq!(drained.tokens_out, 400);
1042        assert_eq!(drained.cost_micros, 2_500);
1043        assert!(m.drain_usage().is_empty());
1044        // Restore (failed-report path) re-adds it for the next period.
1045        m.restore_usage(&drained);
1046        assert_eq!(m.drain_usage(), drained);
1047    }
1048
1049    #[test]
1050    fn llm_token_and_cost_counters() {
1051        let m = Metrics::new();
1052        // Priced model: tokens + cost, bucketed `metered`. Includes cached/reasoning sub-dims.
1053        m.record_llm_usage(
1054            "gpt-4o",
1055            LlmSample {
1056                tokens_in: 100,
1057                tokens_out: 50,
1058                cached_tokens: 40,
1059                reasoning_tokens: 20,
1060                cost_micros: Some(1_250),
1061            },
1062        );
1063        // Unpriced model: tokens counted, no cost, bucketed `unpriced`.
1064        m.record_llm_usage(
1065            "mystery",
1066            LlmSample {
1067                tokens_in: 10,
1068                tokens_out: 5,
1069                cost_micros: None,
1070                ..Default::default()
1071            },
1072        );
1073        // No usage reported.
1074        m.record_llm_no_usage();
1075        let text = m.render();
1076        assert!(
1077            text.contains("edgeguard_llm_tokens_total{direction=\"input\"} 110"),
1078            "{text}"
1079        );
1080        assert!(
1081            text.contains("edgeguard_llm_tokens_total{direction=\"output\"} 55"),
1082            "{text}"
1083        );
1084        assert!(
1085            text.contains("edgeguard_llm_cached_tokens_total 40"),
1086            "{text}"
1087        );
1088        assert!(
1089            text.contains("edgeguard_llm_reasoning_tokens_total 20"),
1090            "{text}"
1091        );
1092        assert!(
1093            text.contains("edgeguard_llm_cost_microdollars_total 1250"),
1094            "{text}"
1095        );
1096        assert!(
1097            text.contains("edgeguard_llm_requests_total{result=\"metered\"} 1"),
1098            "{text}"
1099        );
1100        assert!(
1101            text.contains("edgeguard_llm_requests_total{result=\"unpriced\"} 1"),
1102            "{text}"
1103        );
1104        assert!(
1105            text.contains("edgeguard_llm_requests_total{result=\"no_usage\"} 1"),
1106            "{text}"
1107        );
1108        // Per-model breakdown carries the model + kind labels and the priced model's cost.
1109        assert!(
1110            text.contains("edgeguard_llm_model_tokens_total{model=\"gpt-4o\",kind=\"cached\"} 40"),
1111            "{text}"
1112        );
1113        assert!(
1114            text.contains(
1115                "edgeguard_llm_model_tokens_total{model=\"gpt-4o\",kind=\"reasoning\"} 20"
1116            ),
1117            "{text}"
1118        );
1119        assert!(
1120            text.contains("edgeguard_llm_model_cost_microdollars_total{model=\"gpt-4o\"} 1250"),
1121            "{text}"
1122        );
1123    }
1124
1125    #[test]
1126    fn per_team_tokens_and_cost_are_accumulated_and_rendered() {
1127        let m = Metrics::new();
1128        m.record_llm_team_usage(
1129            "acme",
1130            &LlmSample {
1131                tokens_in: 100,
1132                tokens_out: 40,
1133                cached_tokens: 30,
1134                reasoning_tokens: 10,
1135                cost_micros: Some(77),
1136            },
1137        );
1138        m.record_llm_team_usage(
1139            "acme",
1140            &LlmSample {
1141                tokens_in: 50,
1142                tokens_out: 20,
1143                cost_micros: Some(23),
1144                ..Default::default()
1145            },
1146        );
1147        // A request with no team falls into the shared `_none` bucket.
1148        m.record_llm_team_usage(
1149            "_none",
1150            &LlmSample {
1151                tokens_in: 5,
1152                ..Default::default()
1153            },
1154        );
1155        let text = m.render();
1156        assert!(
1157            text.contains("edgeguard_llm_team_tokens_total{team=\"acme\",kind=\"input\"} 150"),
1158            "{text}"
1159        );
1160        assert!(
1161            text.contains("edgeguard_llm_team_tokens_total{team=\"acme\",kind=\"output\"} 60"),
1162            "{text}"
1163        );
1164        assert!(
1165            text.contains("edgeguard_llm_team_cost_microdollars_total{team=\"acme\"} 100"),
1166            "{text}"
1167        );
1168        assert!(
1169            text.contains("edgeguard_llm_team_tokens_total{team=\"_none\",kind=\"input\"} 5"),
1170            "{text}"
1171        );
1172    }
1173
1174    #[test]
1175    fn budget_reconcile_failures_counter_renders() {
1176        let m = Metrics::new();
1177        m.record_budget_reconcile_failures(0); // a zero is a no-op
1178        m.record_budget_reconcile_failures(2);
1179        m.record_budget_reconcile_failures(1);
1180        assert!(
1181            m.render()
1182                .contains("edgeguard_llm_budget_reconcile_failures_total 3"),
1183            "{}",
1184            m.render()
1185        );
1186    }
1187
1188    #[test]
1189    fn per_key_tokens_and_cost_are_accumulated_and_rendered() {
1190        let m = Metrics::new();
1191        m.record_llm_key_usage(
1192            "key-abc",
1193            &LlmSample {
1194                tokens_in: 100,
1195                tokens_out: 40,
1196                cost_micros: Some(77),
1197                ..Default::default()
1198            },
1199        );
1200        // An unauthenticated request falls into the shared `_anon` bucket.
1201        m.record_llm_key_usage(
1202            "_anon",
1203            &LlmSample {
1204                tokens_in: 5,
1205                ..Default::default()
1206            },
1207        );
1208        let text = m.render();
1209        assert!(
1210            text.contains("edgeguard_llm_key_tokens_total{key=\"key-abc\",kind=\"input\"} 100"),
1211            "{text}"
1212        );
1213        assert!(
1214            text.contains("edgeguard_llm_key_cost_microdollars_total{key=\"key-abc\"} 77"),
1215            "{text}"
1216        );
1217        assert!(
1218            text.contains("edgeguard_llm_key_tokens_total{key=\"_anon\",kind=\"input\"} 5"),
1219            "{text}"
1220        );
1221    }
1222
1223    #[test]
1224    fn per_model_series_are_cardinality_bounded() {
1225        let m = Metrics::new();
1226        // Feed more distinct models than the cap; the overflow bucket absorbs the excess so the map
1227        // never grows past MAX_LLM_MODEL_SERIES + 1 (the overflow key).
1228        for i in 0..(MAX_LLM_MODEL_SERIES + 50) {
1229            m.record_llm_usage(
1230                &format!("model-{i}"),
1231                LlmSample {
1232                    tokens_in: 1,
1233                    ..Default::default()
1234                },
1235            );
1236        }
1237        let map = m.per_model.lock().unwrap();
1238        assert!(map.len() <= MAX_LLM_MODEL_SERIES + 1, "len={}", map.len());
1239        assert!(map.contains_key(LLM_MODEL_OVERFLOW));
1240    }
1241
1242    #[test]
1243    fn budget_blocked_and_consumed_metrics() {
1244        let m = Metrics::new();
1245        m.record_budget_blocked("key");
1246        m.record_budget_blocked("key");
1247        m.record_budget_blocked("team");
1248        m.record_budget_blocked("totally_unknown"); // -> "other"
1249        m.record_budget_consumed("daily-cap", 0.75);
1250        m.record_budget_consumed("daily-cap", 0.92); // last writer wins
1251        m.record_budget_consumed("nan-guard", f64::NAN); // dropped
1252        let text = m.render();
1253        assert!(
1254            text.contains("edgeguard_llm_budget_blocked_total{scope=\"key\"} 2"),
1255            "{text}"
1256        );
1257        assert!(
1258            text.contains("edgeguard_llm_budget_blocked_total{scope=\"team\"} 1"),
1259            "{text}"
1260        );
1261        assert!(
1262            text.contains("edgeguard_llm_budget_blocked_total{scope=\"other\"} 1"),
1263            "{text}"
1264        );
1265        assert!(
1266            text.contains("edgeguard_llm_budget_consumed_ratio{budget=\"daily-cap\"} 0.92"),
1267            "{text}"
1268        );
1269        assert!(
1270            !text.contains("nan-guard"),
1271            "NaN sample must be dropped: {text}"
1272        );
1273    }
1274
1275    #[test]
1276    fn label_values_are_escaped() {
1277        // A client-chosen model with a quote/newline must not break the exposition format.
1278        let m = Metrics::new();
1279        m.record_llm_usage(
1280            "evil\"\nmodel",
1281            LlmSample {
1282                tokens_in: 1,
1283                ..Default::default()
1284            },
1285        );
1286        let text = m.render();
1287        assert!(text.contains("model=\"evil\\\"\\nmodel\""), "{text}");
1288    }
1289
1290    #[test]
1291    fn dlp_finding_and_blocked_counters() {
1292        let m = Metrics::new();
1293        m.record_dlp_finding("email");
1294        m.record_dlp_finding("email");
1295        m.record_dlp_finding("api_key");
1296        m.record_dlp_finding("totally_unknown"); // -> "other"
1297        m.record_dlp_blocked();
1298        let text = m.render();
1299        assert!(
1300            text.contains("edgeguard_llm_dlp_findings_total{category=\"email\"} 2"),
1301            "{text}"
1302        );
1303        assert!(
1304            text.contains("edgeguard_llm_dlp_findings_total{category=\"api_key\"} 1"),
1305            "{text}"
1306        );
1307        assert!(
1308            text.contains("edgeguard_llm_dlp_findings_total{category=\"other\"} 1"),
1309            "{text}"
1310        );
1311        assert!(text.contains("edgeguard_llm_dlp_blocked_total 1"), "{text}");
1312    }
1313
1314    #[test]
1315    fn keyvault_result_counters() {
1316        let m = Metrics::new();
1317        m.record_keyvault("swapped");
1318        m.record_keyvault("swapped");
1319        m.record_keyvault("denied_model");
1320        m.record_keyvault("totally_unknown"); // ignored, not miscounted
1321        let text = m.render();
1322        assert!(
1323            text.contains("edgeguard_llm_keyvault_total{result=\"swapped\"} 2"),
1324            "{text}"
1325        );
1326        assert!(
1327            text.contains("edgeguard_llm_keyvault_total{result=\"denied_model\"} 1"),
1328            "{text}"
1329        );
1330        assert!(
1331            text.contains("edgeguard_llm_keyvault_total{result=\"denied_key\"} 0"),
1332            "{text}"
1333        );
1334    }
1335
1336    #[test]
1337    fn waf_hit_counters_by_class() {
1338        let m = Metrics::new();
1339        m.record_waf_hit("sqli");
1340        m.record_waf_hit("sqli");
1341        m.record_waf_hit("custom");
1342        // An unknown class is ignored rather than miscounted.
1343        m.record_waf_hit("totally_unknown");
1344        let text = m.render();
1345        assert!(
1346            text.contains("edgeguard_waf_hits_total{rule=\"sqli\"} 2"),
1347            "{text}"
1348        );
1349        assert!(
1350            text.contains("edgeguard_waf_hits_total{rule=\"custom\"} 1"),
1351            "{text}"
1352        );
1353        // A class that never fired still renders at 0.
1354        assert!(
1355            text.contains("edgeguard_waf_hits_total{rule=\"xss\"} 0"),
1356            "{text}"
1357        );
1358    }
1359}