From 2a717a993ee99f9401f5cdf78a23dcecd7cb1a51 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Fri, 22 May 2026 15:39:42 -0700 Subject: [PATCH] =?UTF-8?q?fix(observability):=20G3=20remediation=20?= =?UTF-8?q?=E2=80=94=20bound=20cardinality=20+=20wire=20dead=20metrics?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase G PR-G3 review identified 5 Critical + 4 High + 5 Medium findings. This commit lands all 14 fixes plus the optional nits. CRITICAL * C1 (cardinality DoS): `service_tier` was read from inbound JSON and used verbatim as a metric label. A malicious client could blow up the metric vector unboundedly. Added bounded vocabulary in `metric_names.rs::service_tier` ({auto, default, flex, on_demand, priority, scale, other-sentinel}) + a `validate()` helper. Both request-side (`handlers/responses.rs`) and response-side (`proxy.rs` Responses arm) gate raw values through it. * C2 (dead metric): `proxy_passthrough_bytes_modified_total` had no production emit site. Wired it in `proxy.rs` to fire when a dispatcher arm returning `NoCompression`/`Passthrough` produces a body of a different byte length (a true cache-poisoning regression detector). The check runs BEFORE the PR-E4 prompt_cache_key injector so legitimate injector mutations do not trip the alarm. * C3 (Python/Rust boundary): `proxy_image_generation_call_log_redacted_total` was a dead Rust counter — the redaction happens entirely in the Python proxy's request_logger. Removed the Rust counter; moved the metric to the Python proxy's `/metrics` exporter via the existing `redactions_total()` module-level counter. * C4 (Python/Rust boundary): `wrap_rtk_invocations_total` was a dead Rust counter with no wrap-side bridge. Removed the Rust counter; added new `headroom/cli/wrap_rtk_metrics.py` with `record_rtk_invocation(tool, delta)` + `rtk_invocation_counts()` primitives and surfaced them via the Python proxy's `/metrics` exporter. * C5 (dead metric): `proxy_compression_rejected_by_token_check_total` had no production caller. Wired it in `live_zone_anthropic.rs`, `live_zone_openai.rs`, and `live_zone_responses.rs` to increment on every `BlockAction::RejectedNotSmaller` block in the manifest. The metric now reflects real "compressor ran but kept original" cases. HIGH * H1 (per-strategy ratio garbage): `proxy_compression_ratio_by_strategy` emitted the same aggregate ratio for every strategy in `strategies_applied` when multiple strategies ran on one body. Added `per_strategy_tokens: Vec` to `Outcome::Compressed`; per-strategy `(before, after)` is accumulated from the manifest at the wrapper sites and emitted one sample per strategy in `proxy.rs`. Empty vec → fallback to one aggregate-labelled sample with a debug log (Phase E normalization paths that don't track per-strategy tokens). * H2 (aborted stream): cache_hit_rate observed on client disconnects mid-stream. Added a gate: Anthropic only fires when `state.status == MessageStop`, OpenAI Responses only when `terminal_status().is_some()`. Extracted the gate into the pure function `compute_anthropic_session_hit_rate(state)` so the H2 contract is unit-testable independent of the shared global registry. * H3 (docs lie + alarm contract): docs claimed HELP/TYPE is reachable on fresh boot, then contradicted itself. Force-zero every counter / gauge MetricVec with an `__init__` sentinel label on each scrape so HELP/TYPE + a zero row are visible from boot. Histograms are NOT force-zeroed (a synthetic observe(0.0) would pollute percentiles). PromQL queries in docs filter `{... != "__init__"}` so the sentinel rows are excluded from aggregations. * H4 (crate-version dependency): pinned `prometheus = "=0.13.4"` exactly (no caret) so a future minor bump cannot silently break the H3 force-zero contract that relies on this crate's gather() semantics. Added a clear "retest the alarm contract on bump" paragraph in docs. MEDIUM * M1 (saturate on cached > input): OpenAI Chat + Responses cache- hit-rate computed `non_cached = input.saturating_sub(cached)`, silently clamping to 0 if `cached > input`. Per "no silent fallbacks", log + skip the emit on this wire-format pathology. * M2 (over-fire on non-image base64): Python redactor's "density heuristic" over-fired on encrypted blobs / signed tokens / minified JSON / tool outputs. Tightened: only redact strings inside known image-bearing JSON paths (`data`, `url`, `image_url`, `image`) OR strings starting with `data:image/`. * M3 (NaN clamp): cache_hit_rate::observe used `f64::clamp(0,1)` which returns NaN for NaN input; the `debug_assert!` was compiled out in release. Added `is_finite()` guard with a loud-log + skip before observe. * M4 (PromQL median-only): added p95, p99, mean (sum/count), and Phase H canary-gate query section to docs. Canary fails if ANY of {p50, p95, p99, mean} regresses below the Python baseline. * M5 (label byte vs char): the `` placeholder reported character count, not UTF-8 byte count. Switched to `.encode('utf-8').__len__()` so the label is honest for non-ASCII payloads (ASCII base64 still has byte == char so existing scrapes are unchanged). OPTIONAL * Removed dead `debug_assert_eq!(buffered.len(), buffered.len(), ...)` no-op in proxy.rs. * Normalised `record_response_status` log level from `info` to `debug` to match peer metric helpers. Tests: * Rust: 11 integration_metrics tests (was 6) + 9 cache_hit_rate unit tests (was 4) + 2 compression_ratio (unchanged). New coverage: service_tier known/unknown bucketing, C2 alarm wire, H1 per-strategy ratio, H2 abort gate, M3 NaN/inf skip. * Python: 27 tests (was 13). New coverage: M2 path-gated redaction, M5 byte vs char label, wrap_rtk_metrics primitive thread safety and validation. `cargo fmt --check`, `cargo clippy --workspace -- -D warnings`, `cargo test -p headroom-proxy --lib` (221 passed) and the integration_metrics + integration_compression + integration_volatile_detector + integration_cache_control + integration_cache_drift + integration_responses + integration_bedrock_metrics test files all green. Full `cargo test --workspace` deferred — disk pressure during the agent session left insufficient space for the linker to write the full integration test artifacts; runs that did fit all passed. `make ci-precheck` deferred for the same reason. ruff check + ruff format + mypy headroom/proxy/request_logger.py + headroom/cli/wrap_rtk_metrics.py + headroom/proxy/prometheus_metrics.py green. --- crates/headroom-proxy/Cargo.toml | 11 +- .../src/compression/live_zone_anthropic.rs | 94 +++++- .../src/compression/live_zone_openai.rs | 26 +- .../src/compression/live_zone_responses.rs | 25 +- crates/headroom-proxy/src/compression/mod.rs | 4 +- .../headroom-proxy/src/handlers/responses.rs | 18 +- .../src/observability/cache_hit_rate.rs | 162 ++++++++++ .../src/observability/metric_names.rs | 74 +++-- .../headroom-proxy/src/observability/mod.rs | 5 +- .../src/observability/prometheus.rs | 53 +++- .../src/observability/proxy_metrics.rs | 91 ++---- crates/headroom-proxy/src/proxy.rs | 290 ++++++++++++------ .../headroom-proxy/src/vertex/raw_predict.rs | 1 + .../tests/integration_metrics.rs | 282 +++++++++++++++-- docs/observability.md | 187 ++++++++++- headroom/cli/wrap_rtk_metrics.py | 73 +++++ headroom/proxy/prometheus_metrics.py | 52 ++++ headroom/proxy/request_logger.py | 134 +++++--- tests/test_cli/test_wrap_rtk_metrics.py | 113 +++++++ tests/test_image_log_redaction.py | 77 ++++- 20 files changed, 1454 insertions(+), 318 deletions(-) create mode 100644 headroom/cli/wrap_rtk_metrics.py create mode 100644 tests/test_cli/test_wrap_rtk_metrics.py diff --git a/crates/headroom-proxy/Cargo.toml b/crates/headroom-proxy/Cargo.toml index 88690ce50..d2c889d86 100644 --- a/crates/headroom-proxy/Cargo.toml +++ b/crates/headroom-proxy/Cargo.toml @@ -58,7 +58,16 @@ crc32fast = "1" # don't need (we serve text-format scrapes only), so we disable # defaults and re-enable nothing — pure registry + counter + # histogram + text encoder is sufficient. -prometheus = { version = "0.13", default-features = false } +# +# H4 fix: the H3 force-zero contract (in +# `observability::prometheus::handle_metrics`) relies on this +# crate's v0.13 `gather()` semantics — empty MetricVec families are +# omitted from the scrape, so we force-touch each counter / gauge +# with a sentinel label to surface HELP/TYPE on boot. Pinning the +# exact patch version (no caret, no `~`) so a future minor-version +# bump cannot silently change the alarm contract; the bump must be +# an explicit code review that re-validates the contract. +prometheus = { version = "=0.13.4", default-features = false } # PR-E6: SHA-256 over canonical bytes of the cache hot zone (system, # tools, early messages) for cache-bust drift detection. Already in # the dev-dependencies (and pulled transitively by `aws-sigv4` via diff --git a/crates/headroom-proxy/src/compression/live_zone_anthropic.rs b/crates/headroom-proxy/src/compression/live_zone_anthropic.rs index 4d3dd7e75..c52b6f078 100644 --- a/crates/headroom-proxy/src/compression/live_zone_anthropic.rs +++ b/crates/headroom-proxy/src/compression/live_zone_anthropic.rs @@ -55,6 +55,22 @@ use crate::cache_stabilization::tool_def_normalize::{ use crate::compression::resolve_frozen_count; use crate::config::{CacheControlAutoFrozen, CompressionMode}; +/// Per-strategy aggregate token counts for the +/// `proxy_compression_ratio_by_strategy` metric. One entry per +/// distinct `strategy` tag observed in the manifest's +/// `BlockAction::Compressed` blocks. `H1` remediation: the proxy +/// previously emitted the same aggregate ratio per strategy when +/// multiple strategies ran on one body — meaning Phase H per- +/// strategy dashboards read garbage. This struct surfaces the +/// genuine per-strategy values so the emit loop reports the right +/// numbers. +#[derive(Debug, Clone, Copy)] +pub struct PerStrategyTokens { + pub strategy: &'static str, + pub original_tokens: usize, + pub compressed_tokens: usize, +} + /// What happened. The caller uses the variant to decide whether to /// forward the original bytes (everything PR-B2 lands on) or a /// modified body (PR-B3+). @@ -72,6 +88,13 @@ pub enum Outcome { tokens_after: usize, strategies_applied: Vec<&'static str>, markers_inserted: Vec, + /// H1 remediation: per-strategy `(before, after)` aggregate + /// for the `proxy_compression_ratio_by_strategy` metric. + /// Empty when the proxy compressed via a non-block path + /// (e.g. Phase E normalization that doesn't have per-strategy + /// token accounting) — emit-site falls back to one + /// aggregate-labelled sample. + per_strategy_tokens: Vec, }, /// Dispatcher opted out for a reason we can name. Passthrough { reason: PassthroughReason }, @@ -379,6 +402,12 @@ pub fn compress_anthropic_request( tokens_after: 0, strategies_applied: strategies, markers_inserted: e3_locations, + // H1 remediation: Phase E normalization passes + // (E1 sort, E3 cache_control auto-placement) + // mutate bytes but don't have per-strategy + // token accounting. Empty vec → emit-site falls + // back to one aggregate sample. + per_strategy_tokens: Vec::new(), } } else { Outcome::NoCompression @@ -390,27 +419,63 @@ pub fn compress_anthropic_request( // dispatcher used to gate per-block acceptance — so the // saving the proxy logs is the saving the cache will // actually see. + // + // H1 + C5 remediation: + // - Per-strategy `(before, after)` aggregate populated + // from the manifest so the proxy emits one + // `proxy_compression_ratio_by_strategy` sample with + // the right numbers per strategy (instead of the + // same aggregate ratio per strategy, which was the + // pre-fix behavior). + // - Every `BlockAction::RejectedNotSmaller` increments + // the `proxy_compression_rejected_by_token_check_total` + // counter so dashboards can attribute "compressor ran + // but kept the original" cases. let mut original_bytes_total: usize = 0; let mut compressed_bytes_total: usize = 0; let mut original_tokens_total: usize = 0; let mut compressed_tokens_total: usize = 0; let mut strategies: Vec<&'static str> = Vec::new(); + let mut per_strategy_tokens: Vec = Vec::new(); for entry in &manifest.block_outcomes { - if let BlockAction::Compressed { - strategy, - original_bytes, - compressed_bytes, - original_tokens, - compressed_tokens, - } = entry.action - { - original_bytes_total += original_bytes; - compressed_bytes_total += compressed_bytes; - original_tokens_total += original_tokens; - compressed_tokens_total += compressed_tokens; - if !strategies.contains(&strategy) { - strategies.push(strategy); + match entry.action { + BlockAction::Compressed { + strategy, + original_bytes, + compressed_bytes, + original_tokens, + compressed_tokens, + } => { + original_bytes_total += original_bytes; + compressed_bytes_total += compressed_bytes; + original_tokens_total += original_tokens; + compressed_tokens_total += compressed_tokens; + if !strategies.contains(&strategy) { + strategies.push(strategy); + } + // H1: accumulate per-strategy tokens (one + // entry per strategy; multiple blocks of + // the same strategy sum). + if let Some(slot) = per_strategy_tokens + .iter_mut() + .find(|s| s.strategy == strategy) + { + slot.original_tokens += original_tokens; + slot.compressed_tokens += compressed_tokens; + } else { + per_strategy_tokens.push(PerStrategyTokens { + strategy, + original_tokens, + compressed_tokens, + }); + } } + BlockAction::RejectedNotSmaller { strategy, .. } => { + // C5: surface the tokenizer-validated + // rejection in the dedicated counter. + crate::observability::record_compression_rejected_by_token_check(strategy); + } + _ => {} } } // Stitch in the PR-E1 / PR-E2 / PR-E3 strategy tags so @@ -463,6 +528,7 @@ pub fn compress_anthropic_request( // PR-E3 surfaces tool-slot location(s); PR-B7 will // append CCR retrieval markers when wired. markers_inserted: e3_locations, + per_strategy_tokens, } } Err(LiveZoneError::BodyNotJson(_)) => { diff --git a/crates/headroom-proxy/src/compression/live_zone_openai.rs b/crates/headroom-proxy/src/compression/live_zone_openai.rs index 23271022c..4bdec6571 100644 --- a/crates/headroom-proxy/src/compression/live_zone_openai.rs +++ b/crates/headroom-proxy/src/compression/live_zone_openai.rs @@ -39,7 +39,7 @@ use serde_json::Value; use crate::cache_stabilization::tool_def_normalize::{ any_tool_has_cache_control, sort_schema_keys_recursive, sort_tools_deterministically, }; -use crate::compression::{Outcome, PassthroughReason}; +use crate::compression::{Outcome, PassthroughReason, PerStrategyTokens}; use crate::config::CompressionMode; /// OpenAI Chat Completions live-zone compression entry point. @@ -159,6 +159,7 @@ pub fn compress_openai_chat_request( tokens_after: 0, strategies_applied: normalization_applied.strategies(), markers_inserted: Vec::new(), + per_strategy_tokens: Vec::new(), }; } Outcome::NoCompression @@ -167,11 +168,17 @@ pub fn compress_openai_chat_request( // Aggregate manifest stats. Mirrors the Anthropic // module — same metric shape so dashboards don't need // to special-case the provider. + // + // H1 + C5 remediation: per-strategy token accumulation + // for the proxy's per-strategy compression-ratio metric + + // every rejected-not-smaller block bumps the dedicated + // counter. let mut original_bytes_total: usize = 0; let mut compressed_bytes_total: usize = 0; let mut original_tokens_total: usize = 0; let mut compressed_tokens_total: usize = 0; let mut strategies: Vec<&'static str> = Vec::new(); + let mut per_strategy_tokens: Vec = Vec::new(); let mut had_compressor_error = false; for entry in &manifest.block_outcomes { match entry.action { @@ -189,6 +196,22 @@ pub fn compress_openai_chat_request( if !strategies.contains(&strategy) { strategies.push(strategy); } + if let Some(slot) = per_strategy_tokens + .iter_mut() + .find(|s| s.strategy == strategy) + { + slot.original_tokens += original_tokens; + slot.compressed_tokens += compressed_tokens; + } else { + per_strategy_tokens.push(PerStrategyTokens { + strategy, + original_tokens, + compressed_tokens, + }); + } + } + BlockAction::RejectedNotSmaller { strategy, .. } => { + crate::observability::record_compression_rejected_by_token_check(strategy); } BlockAction::CompressorError { strategy, @@ -246,6 +269,7 @@ pub fn compress_openai_chat_request( tokens_after: compressed_tokens_total, strategies_applied: strategies, markers_inserted: Vec::new(), + per_strategy_tokens, } } Err(LiveZoneError::BodyNotJson(_)) => { diff --git a/crates/headroom-proxy/src/compression/live_zone_responses.rs b/crates/headroom-proxy/src/compression/live_zone_responses.rs index ef4e58378..82380d86a 100644 --- a/crates/headroom-proxy/src/compression/live_zone_responses.rs +++ b/crates/headroom-proxy/src/compression/live_zone_responses.rs @@ -43,7 +43,7 @@ use serde_json::Value; use crate::cache_stabilization::tool_def_normalize::{ any_tool_has_cache_control, sort_schema_keys_recursive, sort_tools_deterministically, }; -use crate::compression::{Outcome, PassthroughReason}; +use crate::compression::{Outcome, PassthroughReason, PerStrategyTokens}; use crate::config::CompressionMode; /// OpenAI Responses live-zone compression entry point. @@ -170,6 +170,7 @@ pub fn compress_openai_responses_request( tokens_after: 0, strategies_applied: normalization_applied.strategies(), markers_inserted: Vec::new(), + per_strategy_tokens: Vec::new(), }; } Outcome::NoCompression @@ -177,12 +178,15 @@ pub fn compress_openai_responses_request( Ok(LiveZoneOutcome::Modified { new_body, manifest }) => { // Aggregate per-block savings for the structured log. // Mirrors the Chat Completions sibling so dashboards - // don't need provider-specific shapes. + // don't need provider-specific shapes. H1 + C5: per- + // strategy token accumulation + rejected-token-check + // counter. let mut original_bytes_total: usize = 0; let mut compressed_bytes_total: usize = 0; let mut original_tokens_total: usize = 0; let mut compressed_tokens_total: usize = 0; let mut strategies: Vec<&'static str> = Vec::new(); + let mut per_strategy_tokens: Vec = Vec::new(); let mut had_compressor_error = false; for entry in &manifest.block_outcomes { match entry.action { @@ -200,6 +204,22 @@ pub fn compress_openai_responses_request( if !strategies.contains(&strategy) { strategies.push(strategy); } + if let Some(slot) = per_strategy_tokens + .iter_mut() + .find(|s| s.strategy == strategy) + { + slot.original_tokens += original_tokens; + slot.compressed_tokens += compressed_tokens; + } else { + per_strategy_tokens.push(PerStrategyTokens { + strategy, + original_tokens, + compressed_tokens, + }); + } + } + BlockAction::RejectedNotSmaller { strategy, .. } => { + crate::observability::record_compression_rejected_by_token_check(strategy); } BlockAction::CompressorError { strategy, @@ -257,6 +277,7 @@ pub fn compress_openai_responses_request( tokens_after: compressed_tokens_total, strategies_applied: strategies, markers_inserted: Vec::new(), + per_strategy_tokens, } } Err(LiveZoneError::BodyNotJson(_)) => { diff --git a/crates/headroom-proxy/src/compression/mod.rs b/crates/headroom-proxy/src/compression/mod.rs index 24bcbb512..af441e39a 100644 --- a/crates/headroom-proxy/src/compression/mod.rs +++ b/crates/headroom-proxy/src/compression/mod.rs @@ -44,7 +44,9 @@ pub mod model_limits; // itself stays through B1 → B2 transition for parallel review; // `compress_anthropic_request` is sourced from the live-zone module. pub use anthropic::resolve_frozen_count; -pub use live_zone_anthropic::{compress_anthropic_request, Outcome, PassthroughReason}; +pub use live_zone_anthropic::{ + compress_anthropic_request, Outcome, PassthroughReason, PerStrategyTokens, +}; pub use live_zone_openai::{ compress_openai_chat_request, should_skip_compression, SkipCompressionReason, }; diff --git a/crates/headroom-proxy/src/handlers/responses.rs b/crates/headroom-proxy/src/handlers/responses.rs index 9376c993f..da18e60c2 100644 --- a/crates/headroom-proxy/src/handlers/responses.rs +++ b/crates/headroom-proxy/src/handlers/responses.rs @@ -120,12 +120,18 @@ pub async fn handle_responses( // bodies do NOT fabricate a tier — per realignment build- // constraint "no silent fallbacks", we just skip the emit and // log at debug. + // + // C1 fix: every raw value is validated against the bounded + // `service_tier` vocabulary BEFORE being used as a label so a + // malicious client cannot blow up label cardinality with + // arbitrary strings. if let Some(tier) = extract_request_service_tier(&body) { let request_id_for_metric = headers .get("x-request-id") .and_then(|v| v.to_str().ok()) .unwrap_or(""); - observability::record_service_tier(&tier, request_id_for_metric); + let bucketed = crate::observability::metric_names::service_tier::validate(&tier); + observability::record_service_tier(bucketed, request_id_for_metric); } else { tracing::debug!( event = "service_tier_skipped", @@ -167,10 +173,12 @@ pub async fn handle_responses( /// Phase G PR-G3: best-effort parse of `service_tier` from the /// inbound request body. Returns `None` when the body is not valid /// JSON, not an object, or lacks the field. The spec defines the -/// field as a string ∈ {auto, default, flex, on_demand, priority}; -/// we do NOT validate against that set here so wire-format drift -/// (e.g. OpenAI adding a tier we haven't enumerated yet) surfaces -/// in the metric rather than getting silently dropped. +/// field as a string ∈ {auto, default, flex, on_demand, priority, +/// scale}; the returned raw string is normalised against the +/// bounded vocabulary at the call site via +/// [`crate::observability::metric_names::service_tier::validate`] +/// so an arbitrary inbound value cannot drive metric-label +/// cardinality unbounded (C1 fix). fn extract_request_service_tier(body: &Bytes) -> Option { let v: serde_json::Value = serde_json::from_slice(body).ok()?; v.get("service_tier") diff --git a/crates/headroom-proxy/src/observability/cache_hit_rate.rs b/crates/headroom-proxy/src/observability/cache_hit_rate.rs index ef152be29..4e72e911a 100644 --- a/crates/headroom-proxy/src/observability/cache_hit_rate.rs +++ b/crates/headroom-proxy/src/observability/cache_hit_rate.rs @@ -118,6 +118,31 @@ pub fn compute_hit_rate( Some(cache_read_input_tokens as f64 / denom as f64) } +/// H2 gate: should we observe a cache-hit-rate sample for this +/// Anthropic session? +/// +/// Returns `Some(rate)` ONLY when the stream completed cleanly +/// (`state.status == MessageStop`) AND the denominator is non-zero. +/// A client disconnect mid-stream closes the channel too — without +/// this gate we'd observe a half-finished session that has only +/// `message_start` usage and pollute the histogram with garbage. +/// +/// Extracted from `proxy.rs::run_sse_state_machine` so the H2 +/// contract is unit-testable independent of the global Prometheus +/// registry (which parallel tests share). +pub fn compute_anthropic_session_hit_rate( + state: &crate::sse::anthropic::AnthropicStreamState, +) -> Option { + if state.status != crate::sse::anthropic::StreamStatus::MessageStop { + return None; + } + compute_hit_rate( + state.usage.input_tokens, + state.usage.cache_read_input_tokens, + state.usage.cache_creation_input_tokens, + ) +} + /// Observe one per-session sample. /// /// `provider` MUST be one of the [`provider`] constants — callers @@ -125,11 +150,31 @@ pub fn compute_hit_rate( /// do not validate the label here because the cardinality is bounded /// by the static call sites; an invalid label would be a bug, not a /// runtime mismatch. +/// +/// M3 fix: NaN / non-finite inputs are loud-skipped instead of +/// silently observed. `f64::clamp(0.0, 1.0)` returns NaN when the +/// input is NaN, so the prior implementation could pollute the +/// histogram with NaN samples in release builds (the debug_assert +/// was compiled out). Per "no silent fallbacks", an unexpected NaN +/// surfaces in the logs rather than getting eaten. pub fn observe(provider: &'static str, request_id: &str, hit_rate: f64) { + if !hit_rate.is_finite() { + tracing::warn!( + event = "cache_hit_rate_non_finite", + metric = METRIC_PROXY_CACHE_HIT_RATE_PER_SESSION, + provider = provider, + request_id = %request_id, + hit_rate = hit_rate, + "refusing to observe a non-finite cache hit rate; this is a caller bug" + ); + return; + } debug_assert!( (0.0..=1.0).contains(&hit_rate), "cache hit rate must be in [0.0, 1.0]; got {hit_rate}" ); + // After the is_finite guard, clamp can only normalise legitimate + // edge-of-range floats (1.0 + epsilon etc.) and never produces NaN. let clamped = hit_rate.clamp(0.0, 1.0); histogram(super::prometheus::registry()) .with_label_values(&[provider]) @@ -175,4 +220,121 @@ mod tests { // observation, not coerce to 0.0. assert!(compute_hit_rate(0, 0, 0).is_none()); } + + #[test] + fn observe_nan_skipped_loudly() { + // M3: a NaN input must NOT reach the histogram. The + // pre-fix code clamped via `f64::clamp(0.0, 1.0)` which + // returns NaN for NaN input — a NaN sample in release + // builds was the bug. After the fix we log + skip. + // The histogram count for this synthetic provider stays + // at whatever it was before the call. + let label = "test_nan_provider_v1"; + // Drive a guaranteed-clean session count via a real + // observation, then push a NaN and assert no count move. + observe(label, "req-nan-baseline", 0.5); + let before = histogram(super::super::prometheus::registry()) + .with_label_values(&[label]) + .get_sample_count(); + observe(label, "req-nan-attempt", f64::NAN); + let after = histogram(super::super::prometheus::registry()) + .with_label_values(&[label]) + .get_sample_count(); + assert_eq!( + before, after, + "NaN must not be observed; expected count unchanged from {before} to {after}" + ); + } + + #[test] + fn h2_aborted_anthropic_stream_returns_none() { + // H2: a stream that closes without `message_stop` (Open + // state) MUST NOT produce a sample, regardless of usage + // values. + use crate::sse::anthropic::{AnthropicStreamState, StreamStatus, UsageBuilder}; + let state = AnthropicStreamState { + status: StreamStatus::Open, + usage: UsageBuilder { + input_tokens: 800, + cache_read_input_tokens: 200, + output_tokens: 50, + cache_creation_input_tokens: 0, + }, + ..Default::default() + }; + assert!( + compute_anthropic_session_hit_rate(&state).is_none(), + "H2 gate must skip emission for non-completed stream" + ); + } + + #[test] + fn h2_errored_anthropic_stream_returns_none() { + // H2: an Errored stream is also not "completed cleanly" — + // skip the observation. + use crate::sse::anthropic::{AnthropicStreamState, StreamStatus, UsageBuilder}; + let state = AnthropicStreamState { + status: StreamStatus::Errored, + usage: UsageBuilder { + input_tokens: 800, + cache_read_input_tokens: 200, + output_tokens: 50, + cache_creation_input_tokens: 0, + }, + ..Default::default() + }; + assert!( + compute_anthropic_session_hit_rate(&state).is_none(), + "H2 gate must skip emission for errored stream" + ); + } + + #[test] + fn h2_completed_anthropic_stream_returns_rate() { + // H2 positive case: MessageStop + non-zero usage → return + // the rate so the caller can observe it. + use crate::sse::anthropic::{AnthropicStreamState, StreamStatus, UsageBuilder}; + let state = AnthropicStreamState { + status: StreamStatus::MessageStop, + usage: UsageBuilder { + input_tokens: 800, + cache_read_input_tokens: 200, + output_tokens: 50, + cache_creation_input_tokens: 0, + }, + ..Default::default() + }; + let rate = compute_anthropic_session_hit_rate(&state).expect("completed stream emits"); + // 200 / (800 + 200 + 0) = 0.2 + assert!((rate - 0.2).abs() < 1e-9); + } + + #[test] + fn h2_completed_but_zero_denominator_returns_none() { + // Even on a completed stream, a zero-token request returns + // None — per "no silent fallbacks", no synthesised 0.0. + use crate::sse::anthropic::{AnthropicStreamState, StreamStatus, UsageBuilder}; + let state = AnthropicStreamState { + status: StreamStatus::MessageStop, + usage: UsageBuilder::default(), + ..Default::default() + }; + assert!(compute_anthropic_session_hit_rate(&state).is_none()); + } + + #[test] + fn observe_infinity_skipped_loudly() { + // Same contract for +/- infinity. + let label = "test_inf_provider_v1"; + observe(label, "req-inf-baseline", 0.5); + let before = histogram(super::super::prometheus::registry()) + .with_label_values(&[label]) + .get_sample_count(); + observe(label, "req-pos-inf", f64::INFINITY); + observe(label, "req-neg-inf", f64::NEG_INFINITY); + let after = histogram(super::super::prometheus::registry()) + .with_label_values(&[label]) + .get_sample_count(); + assert_eq!(before, after); + } } diff --git a/crates/headroom-proxy/src/observability/metric_names.rs b/crates/headroom-proxy/src/observability/metric_names.rs index 4edefb18b..7708f92f5 100644 --- a/crates/headroom-proxy/src/observability/metric_names.rs +++ b/crates/headroom-proxy/src/observability/metric_names.rs @@ -94,29 +94,15 @@ pub const METRIC_PROXY_RESPONSE_STATUS_COUNT_TOTAL_HELP: &str = 'incomplete' detail lands in the structured log paired with each \ increment."; -// ---------- proxy_image_generation_call_log_redacted_total ---------- - -pub const METRIC_PROXY_IMAGE_GENERATION_CALL_LOG_REDACTED_TOTAL: &str = - "proxy_image_generation_call_log_redacted_total"; -pub const METRIC_PROXY_IMAGE_GENERATION_CALL_LOG_REDACTED_TOTAL_HELP: &str = - "Count of base64-encoded image payloads redacted from request logs. \ - Driven from the Python proxy's request logger (large multi-MB \ - base64 strings replaced with size-only placeholders)."; - -// ---------- wrap_rtk_invocations_total ---------- - -pub const METRIC_WRAP_RTK_INVOCATIONS_TOTAL: &str = "wrap_rtk_invocations_total"; -pub const METRIC_WRAP_RTK_INVOCATIONS_TOTAL_HELP: &str = - "Count of RTK invocations observed via the wrap-CLI polling of \ - `rtk gain --format json`, broken down by the rewritten command \ - (`git`, `ls`, `cargo`, …)."; - -// ---------- wrap_rtk_tokens_saved_per_session ---------- - -pub const METRIC_WRAP_RTK_TOKENS_SAVED_PER_SESSION: &str = "wrap_rtk_tokens_saved_per_session"; -pub const METRIC_WRAP_RTK_TOKENS_SAVED_PER_SESSION_HELP: &str = - "Tokens saved by RTK during one wrap session, observed once at \ - session end. Histogram so dashboards can render a distribution."; +// Phase G PR-G3 remediation (C3 + C4): the metric-name constants +// for `proxy_image_generation_call_log_redacted_total`, +// `wrap_rtk_invocations_total`, and `wrap_rtk_tokens_saved_per_session` +// were removed because the underlying counters had no production +// emit site on the Rust side. The same metrics are exported by the +// Python proxy (`headroom/proxy/prometheus_metrics.py`) which is the +// natural owner: image redaction is a Python-proxy operation and RTK +// invocation tracking lives in the wrap CLI, both Python-side +// surfaces. See `docs/observability.md`. // ---------- shared label keys ---------- @@ -126,20 +112,54 @@ pub const LABEL_CONTENT_TYPE: &str = "content_type"; pub const LABEL_PATH: &str = "path"; pub const LABEL_TIER: &str = "tier"; pub const LABEL_STATUS: &str = "status"; -pub const LABEL_TOOL: &str = "tool"; // ---------- bounded label vocabularies ---------- /// OpenAI service-tier values per the Responses API spec -/// (`service_tier` field on the response object). The proxy logs -/// anything outside this set under the literal value so wire-format -/// drift is loud rather than silently bucketed. +/// (`service_tier` field on the response object). The metric label +/// vocabulary is **strictly** this set plus a `"scale"` value +/// (documented in OpenAI's tier-pricing page) and a sentinel +/// `"other"` bucket for anything else, so a malicious client posting +/// `{"service_tier":""}` per request cannot blow up +/// cardinality. pub mod service_tier { pub const AUTO: &str = "auto"; pub const DEFAULT: &str = "default"; pub const FLEX: &str = "flex"; pub const ON_DEMAND: &str = "on_demand"; pub const PRIORITY: &str = "priority"; + pub const SCALE: &str = "scale"; + /// Sentinel for any unknown / unrecognised tier value. Prevents + /// label-cardinality DoS from arbitrary inbound JSON. + pub const OTHER: &str = "other"; + + /// Validate an inbound `service_tier` string against the bounded + /// vocabulary. Returns the matching `&'static` constant or + /// [`OTHER`] for any unrecognised value (with a tracing::warn so + /// wire-format drift is loud rather than silently bucketed). + /// + /// The matching is case-sensitive — the OpenAI spec is + /// case-sensitive on these strings; a case-different value is + /// treated as drift, not as the same tier. + pub fn validate(raw: &str) -> &'static str { + match raw { + AUTO => AUTO, + DEFAULT => DEFAULT, + FLEX => FLEX, + ON_DEMAND => ON_DEMAND, + PRIORITY => PRIORITY, + SCALE => SCALE, + _ => { + tracing::warn!( + event = "service_tier_unknown", + raw = %raw, + bucket = OTHER, + "unknown service_tier value bucketed to 'other' to bound cardinality" + ); + OTHER + } + } + } } /// OpenAI Responses terminal-status vocabulary. `in_progress` is the diff --git a/crates/headroom-proxy/src/observability/mod.rs b/crates/headroom-proxy/src/observability/mod.rs index 295a52b9f..d8ad2c7a7 100644 --- a/crates/headroom-proxy/src/observability/mod.rs +++ b/crates/headroom-proxy/src/observability/mod.rs @@ -63,7 +63,6 @@ pub use compression_ratio::{ record_rejected_by_token_check as record_compression_rejected_by_token_check, }; pub use proxy_metrics::{ - extract_rate_limit_snapshot, record_image_redacted, record_passthrough_bytes_modified, - record_rate_limit_snapshot, record_response_status, record_rtk_invocation, record_service_tier, - RateLimitSnapshot, + extract_rate_limit_snapshot, record_passthrough_bytes_modified, record_rate_limit_snapshot, + record_response_status, record_service_tier, RateLimitSnapshot, }; diff --git a/crates/headroom-proxy/src/observability/prometheus.rs b/crates/headroom-proxy/src/observability/prometheus.rs index ec23a7605..bd6307d85 100644 --- a/crates/headroom-proxy/src/observability/prometheus.rs +++ b/crates/headroom-proxy/src/observability/prometheus.rs @@ -217,21 +217,50 @@ pub async fn handle_metrics() -> Response { // Phase G PR-G3: same idea for the new proxy-wide metric families. // Lazy `OnceLock`-backed singletons; touching each here forces - // registration on first scrape so HELP/TYPE lines appear even - // before traffic has driven a single increment. + // registration on first scrape. + // + // H3 fix: registration alone is NOT enough — the prometheus + // crate's v0.13 `gather()` skips empty MetricVec families + // entirely (no HELP/TYPE lines either). Operators who curl + // /metrics on a fresh boot would otherwise see NO trace of the + // catalogue. We force-touch each counter / gauge MetricVec + // with a sentinel `__init__` label so HELP/TYPE + a zero row + // appear from boot and dashboards/alarms have a predictable + // scrape shape. Counters with the `__init__` label increment + // by 0, so the alarm-able "must stay 0" semantic of + // `proxy_passthrough_bytes_modified_total` is preserved + // (the family becomes visible, the rate stays 0). + // + // Histograms are NOT force-zeroed: a synthetic `observe(0.0)` + // would pollute the per-label distribution. The two histogram + // families (`proxy_cache_hit_rate_per_session` and + // `proxy_compression_ratio_by_strategy`) only surface in the + // scrape after the first real session, by design. let reg = registry(); let _ = super::cache_hit_rate::histogram(reg); let _ = super::compression_ratio::ratio_histogram(reg); - let _ = super::compression_ratio::rejected_counter(reg); - let _ = super::proxy_metrics::passthrough_bytes_modified_counter(reg); - let _ = super::proxy_metrics::rate_limit_remaining_requests_gauge(reg); - let _ = super::proxy_metrics::rate_limit_remaining_tokens_gauge(reg); - let _ = super::proxy_metrics::rate_limit_remaining_input_tokens_gauge(reg); - let _ = super::proxy_metrics::rate_limit_remaining_output_tokens_gauge(reg); - let _ = super::proxy_metrics::service_tier_counter(reg); - let _ = super::proxy_metrics::response_status_counter(reg); - let _ = super::proxy_metrics::image_redacted_counter(reg); - let _ = super::proxy_metrics::rtk_invocations_counter(reg); + let rejected_counter = super::compression_ratio::rejected_counter(reg); + let passthrough_counter = super::proxy_metrics::passthrough_bytes_modified_counter(reg); + let rl_requests_gauge = super::proxy_metrics::rate_limit_remaining_requests_gauge(reg); + let rl_tokens_gauge = super::proxy_metrics::rate_limit_remaining_tokens_gauge(reg); + let rl_input_gauge = super::proxy_metrics::rate_limit_remaining_input_tokens_gauge(reg); + let rl_output_gauge = super::proxy_metrics::rate_limit_remaining_output_tokens_gauge(reg); + let tier_counter = super::proxy_metrics::service_tier_counter(reg); + let status_counter = super::proxy_metrics::response_status_counter(reg); + + const INIT_SENTINEL: &str = "__init__"; + rejected_counter + .with_label_values(&[INIT_SENTINEL]) + .inc_by(0); + passthrough_counter + .with_label_values(&[INIT_SENTINEL]) + .inc_by(0); + rl_requests_gauge.with_label_values(&[INIT_SENTINEL]).set(0); + rl_tokens_gauge.with_label_values(&[INIT_SENTINEL]).set(0); + rl_input_gauge.with_label_values(&[INIT_SENTINEL]).set(0); + rl_output_gauge.with_label_values(&[INIT_SENTINEL]).set(0); + tier_counter.with_label_values(&[INIT_SENTINEL]).inc_by(0); + status_counter.with_label_values(&[INIT_SENTINEL]).inc_by(0); let metric_families = registry().gather(); let mut buffer = Vec::with_capacity(2048); diff --git a/crates/headroom-proxy/src/observability/proxy_metrics.rs b/crates/headroom-proxy/src/observability/proxy_metrics.rs index 527e64ec6..371b2887f 100644 --- a/crates/headroom-proxy/src/observability/proxy_metrics.rs +++ b/crates/headroom-proxy/src/observability/proxy_metrics.rs @@ -13,9 +13,7 @@ use std::sync::OnceLock; use prometheus::{IntCounterVec, IntGaugeVec, Opts, Registry}; use super::metric_names::{ - LABEL_PATH, LABEL_PROVIDER, LABEL_STATUS, LABEL_TIER, LABEL_TOOL, - METRIC_PROXY_IMAGE_GENERATION_CALL_LOG_REDACTED_TOTAL, - METRIC_PROXY_IMAGE_GENERATION_CALL_LOG_REDACTED_TOTAL_HELP, + LABEL_PATH, LABEL_PROVIDER, LABEL_STATUS, LABEL_TIER, METRIC_PROXY_PASSTHROUGH_BYTES_MODIFIED_TOTAL, METRIC_PROXY_PASSTHROUGH_BYTES_MODIFIED_TOTAL_HELP, METRIC_PROXY_RATE_LIMIT_REMAINING_INPUT_TOKENS, @@ -26,7 +24,6 @@ use super::metric_names::{ METRIC_PROXY_RATE_LIMIT_REMAINING_TOKENS, METRIC_PROXY_RATE_LIMIT_REMAINING_TOKENS_HELP, METRIC_PROXY_RESPONSE_STATUS_COUNT_TOTAL, METRIC_PROXY_RESPONSE_STATUS_COUNT_TOTAL_HELP, METRIC_PROXY_SERVICE_TIER_COUNT_TOTAL, METRIC_PROXY_SERVICE_TIER_COUNT_TOTAL_HELP, - METRIC_WRAP_RTK_INVOCATIONS_TOTAL, METRIC_WRAP_RTK_INVOCATIONS_TOTAL_HELP, }; // ---------- proxy_passthrough_bytes_modified_total{path} ---------- @@ -269,7 +266,11 @@ pub fn record_response_status(status: &str, reason: Option<&str>, request_id: &s response_status_counter(super::prometheus::registry()) .with_label_values(&[status]) .inc(); - tracing::info!( + // Optional-3: aligned with the peer `record_*` helpers in this + // module which all use `debug!` for the metric-correlation log + // line. INFO was inconsistent and produced extra log volume + // during normal Responses traffic. + tracing::debug!( event = "metric_recorded", metric = METRIC_PROXY_RESPONSE_STATUS_COUNT_TOTAL, status = %status, @@ -279,74 +280,18 @@ pub fn record_response_status(status: &str, reason: Option<&str>, request_id: &s ); } -// ---------- proxy_image_generation_call_log_redacted_total ---------- - -pub fn image_redacted_counter(registry: &Registry) -> &'static IntCounterVec { - static COUNTER: OnceLock = OnceLock::new(); - COUNTER.get_or_init(|| { - let opts = Opts::new( - METRIC_PROXY_IMAGE_GENERATION_CALL_LOG_REDACTED_TOTAL, - METRIC_PROXY_IMAGE_GENERATION_CALL_LOG_REDACTED_TOTAL_HELP, - ); - // Unlabelled metric — but `IntCounterVec` with zero labels - // is still the right type for symmetry with the rest of the - // module. A bare `IntCounter` would diverge in registration - // shape and read site. - let counter = IntCounterVec::new(opts, &[]) - .expect("proxy_image_generation_call_log_redacted_total descriptor is well-formed"); - registry - .register(Box::new(counter.clone())) - .expect("proxy_image_generation_call_log_redacted_total registers exactly once"); - counter - }) -} - -pub fn record_image_redacted() { - image_redacted_counter(super::prometheus::registry()) - .with_label_values(&[]) - .inc(); - tracing::debug!( - event = "metric_recorded", - metric = METRIC_PROXY_IMAGE_GENERATION_CALL_LOG_REDACTED_TOTAL, - "incremented proxy_image_generation_call_log_redacted_total" - ); -} - -// ---------- wrap_rtk_invocations_total{tool} ---------- -// -// RTK lives on the wrap-CLI side; the proxy never invokes it. The -// counter is registered here so the central /metrics scrape exposes -// it and the wrap-side tail can `inc()` via the registry. The wrap -// side polls `rtk gain --format json` and increments by the delta. - -pub fn rtk_invocations_counter(registry: &Registry) -> &'static IntCounterVec { - static COUNTER: OnceLock = OnceLock::new(); - COUNTER.get_or_init(|| { - let opts = Opts::new( - METRIC_WRAP_RTK_INVOCATIONS_TOTAL, - METRIC_WRAP_RTK_INVOCATIONS_TOTAL_HELP, - ); - let counter = IntCounterVec::new(opts, &[LABEL_TOOL]) - .expect("wrap_rtk_invocations_total descriptor is well-formed"); - registry - .register(Box::new(counter.clone())) - .expect("wrap_rtk_invocations_total registers exactly once"); - counter - }) -} - -pub fn record_rtk_invocation(tool: &str, delta: u64) { - rtk_invocations_counter(super::prometheus::registry()) - .with_label_values(&[tool]) - .inc_by(delta); - tracing::debug!( - event = "metric_recorded", - metric = METRIC_WRAP_RTK_INVOCATIONS_TOTAL, - tool = %tool, - delta = delta, - "incremented wrap_rtk_invocations_total" - ); -} +// Phase G PR-G3 remediation (C3 + C4): the image-redacted counter +// and the wrap_rtk_invocations counter were originally registered +// here but neither had a production emit site that crossed the +// Python/Rust boundary. Both have moved Python-side +// (`headroom.proxy.request_logger::redactions_total` and +// `headroom.cli.wrap_rtk_metrics::rtk_invocation_counts`) and the +// Python proxy's `/metrics` exporter surfaces them — see +// `docs/observability.md` for the placement decision. Keeping a +// dead Rust counter would (a) violate the "no dead metrics +// registered" review finding and (b) mislead Phase H canary +// dashboards into expecting two scrape sources for what is really +// one Python-side counter. #[cfg(test)] mod tests { diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index e81294442..d3fe821cc 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -730,22 +730,24 @@ pub(crate) async fn forward_http( } }; + // C2 fix: snapshot the original buffered byte-length AND the + // dispatcher's "is this a passthrough arm?" decision BEFORE + // `outcome` is consumed by the match below. The + // passthrough-bytes-modified alarm fires when a path that + // promised byte-equal passthrough produces a different + // length downstream. + let original_buffered_len = buffered.len(); + let outcome_is_passthrough_class = matches!( + outcome, + compression::Outcome::NoCompression | compression::Outcome::Passthrough { .. } + ); let body_to_send = match outcome { compression::Outcome::NoCompression => { // PR-B2: forward the *original* buffered bytes. The // cache-safety invariant (bytes-in == bytes-out) // is the whole point of the live-zone architecture // — the dispatcher only mutates body bytes when at - // least one block compressed. PR-B2's no-op - // skeleton always lands here. This assert catches - // accidental future regressions where a compressor - // returns `NoCompression` but already mutated the - // buffer in place. - debug_assert_eq!( - buffered.len(), - buffered.len(), - "buffered bytes length must remain stable on the NoCompression path" - ); + // least one block compressed. buffered } // PR-B3+ produces `Compressed` from the live-zone @@ -758,6 +760,7 @@ pub(crate) async fn forward_http( tokens_after, strategies_applied, markers_inserted, + per_strategy_tokens, } => { tracing::info!( request_id = %request_id, @@ -769,25 +772,50 @@ pub(crate) async fn forward_http( markers = markers_inserted.len(), "compression applied" ); - // Phase G PR-G3: emit one + // Phase G PR-G3 + H1: emit one // `proxy_compression_ratio_by_strategy` sample per - // strategy actually applied. We do NOT have a - // per-block content_type breakdown here — the - // aggregate `tokens_before` / `tokens_after` is - // strategy-wise rather than content-type-wise, so we - // label `content_type=aggregate`. Per-block fanout - // would require threading the manifest through to - // this layer; that's a future PR if dashboards need - // it. - if tokens_before > 0 && tokens_after < tokens_before { - for strategy in &strategies_applied { + // strategy with the *strategy's own* before/after + // token counts. The pre-H1 code emitted the same + // aggregate ratio for every strategy in + // `strategies_applied`, so Phase H per-strategy + // dashboards read garbage when multiple strategies + // ran on one body. We now plumb per-strategy tokens + // from the manifest at the wrapper site + // (`live_zone_anthropic`, `live_zone_openai`, + // `live_zone_responses`). + // + // Fallback: when `per_strategy_tokens` is empty — + // i.e. the Outcome came from a Phase E + // normalization pass that doesn't track per-strategy + // tokens — we emit one aggregate-labelled sample so + // dashboards still see *that* a compression ran. We + // log loudly so this is visible. + if !per_strategy_tokens.is_empty() { + for entry in &per_strategy_tokens { crate::observability::observe_compression_ratio( - strategy, + entry.strategy, "aggregate", - tokens_before, - tokens_after, + entry.original_tokens, + entry.compressed_tokens, ); } + } else if tokens_before > 0 && tokens_after < tokens_before { + tracing::debug!( + event = "compression_ratio_emit_aggregate_only", + request_id = %request_id, + path = %path_for_log, + strategies = ?strategies_applied, + reason = "no_per_strategy_tokens", + "emitting one aggregate-labelled compression_ratio sample because \ + the dispatcher did not surface per-strategy token counts \ + (Phase E normalization paths)" + ); + crate::observability::observe_compression_ratio( + "aggregate", + "aggregate", + tokens_before, + tokens_after, + ); } body } @@ -802,6 +830,24 @@ pub(crate) async fn forward_http( } }; + // C2 fix: cache-safety alarm. When the dispatcher returned + // `NoCompression` or `Passthrough`, the post-dispatcher body + // MUST be byte-length-equal to the original buffered body. + // Any delta is an accidental cache-poisoning regression and + // the alarm metric `proxy_passthrough_bytes_modified_total{path}` + // fires with the byte delta as its increment. We check BEFORE + // the PR-E4 prompt_cache_key injector runs because that + // injector is a legitimate, intentional byte mutation gated + // on PAYG; it must not trip the alarm. + if outcome_is_passthrough_class && body_to_send.len() != original_buffered_len { + let delta = body_to_send.len().abs_diff(original_buffered_len) as u64; + crate::observability::record_passthrough_bytes_modified( + &path_for_log, + delta, + &request_id, + ); + } + // PR-E4: OpenAI `prompt_cache_key` auto-injection. // // Universal safety contract: only mutate when the caller @@ -1242,17 +1288,12 @@ async fn run_sse_state_machine( } } } - // Phase G PR-G3: emit per-session cache-hit-rate from - // the final accumulated `usage`. Anthropic's - // `message_delta` strictly grows token counts, so the - // accumulator's final values are the right denominator. - // `compute_cache_hit_rate` returns `None` for zero - // denominators → log + skip, never synthesise a sample. - match crate::observability::compute_cache_hit_rate( - state.usage.input_tokens, - state.usage.cache_read_input_tokens, - state.usage.cache_creation_input_tokens, - ) { + // Phase G PR-G3 + H2: emit per-session cache-hit-rate + // ONLY when the stream completed cleanly with + // `message_stop`. The gate is encapsulated by the + // pure function `compute_anthropic_session_hit_rate` + // so the H2 contract has a unit-testable surface. + match crate::observability::cache_hit_rate::compute_anthropic_session_hit_rate(&state) { Some(rate) => { crate::observability::observe_cache_hit_rate( crate::observability::cache_hit_rate_provider::ANTHROPIC, @@ -1265,8 +1306,11 @@ async fn run_sse_state_machine( event = "cache_hit_rate_skipped", request_id = %request_id, provider = "anthropic", - reason = "zero_denominator", - "skipping proxy_cache_hit_rate_per_session: no input tokens" + status = ?state.status, + input_tokens = state.usage.input_tokens, + cache_read_input_tokens = state.usage.cache_read_input_tokens, + cache_creation_input_tokens = state.usage.cache_creation_input_tokens, + "skipping proxy_cache_hit_rate_per_session: H2 gate or zero denominator" ); } } @@ -1311,7 +1355,9 @@ async fn run_sse_state_machine( // chunk. OpenAI only emits this when // `stream_options.include_usage = true`; absence is a // signal, not a fallback condition — `usage = None` → - // skip. + // skip. The H2 gate is implicit here: the final usage + // chunk only arrives when the stream completed (it's + // OpenAI's terminal-status equivalent). if let Some(usage) = &state.usage { let input_tokens = usage .get("prompt_tokens") @@ -1322,29 +1368,48 @@ async fn run_sse_state_machine( .and_then(|d| d.get("cached_tokens")) .and_then(|v| v.as_u64()) .unwrap_or(0); - // OpenAI's `prompt_tokens` already INCLUDES cached - // tokens (per Chat Completions API docs), so the - // denominator is `prompt_tokens`, not the sum. The - // numerator is `cached_tokens`; `input_tokens` arg to - // `compute_cache_hit_rate` carries the *non-cached* - // portion (denom-only), so we synthesise that here. - let non_cached = input_tokens.saturating_sub(cached_tokens); - match crate::observability::compute_cache_hit_rate(non_cached, cached_tokens, 0) { - Some(rate) => { - crate::observability::observe_cache_hit_rate( - crate::observability::cache_hit_rate_provider::OPENAI_CHAT, - &request_id, - rate, - ); - } - None => { - tracing::debug!( - event = "cache_hit_rate_skipped", - request_id = %request_id, - provider = "openai_chat", - reason = "zero_denominator", - "skipping proxy_cache_hit_rate_per_session: no input tokens" - ); + // M1: `cached_tokens > input_tokens` is a wire- + // format pathology — log + skip instead of silently + // clamping (saturating_sub would yield 0 → fake 1.0 + // hit-rate sample). + if cached_tokens > input_tokens { + tracing::warn!( + event = "cache_hit_rate_skipped", + request_id = %request_id, + provider = "openai_chat", + reason = "cached_gt_input", + input_tokens = input_tokens, + cached_tokens = cached_tokens, + "skipping proxy_cache_hit_rate_per_session: cached_tokens > prompt_tokens \ + (wire-format pathology; clamping would synthesise a bad sample)" + ); + } else { + // OpenAI's `prompt_tokens` already INCLUDES cached + // tokens (per Chat Completions API docs), so the + // denominator is `prompt_tokens`, not the sum. The + // numerator is `cached_tokens`; `input_tokens` arg + // to `compute_cache_hit_rate` carries the + // *non-cached* portion (denom-only), so we + // synthesise that here. + let non_cached = input_tokens - cached_tokens; + match crate::observability::compute_cache_hit_rate(non_cached, cached_tokens, 0) + { + Some(rate) => { + crate::observability::observe_cache_hit_rate( + crate::observability::cache_hit_rate_provider::OPENAI_CHAT, + &request_id, + rate, + ); + } + None => { + tracing::debug!( + event = "cache_hit_rate_skipped", + request_id = %request_id, + provider = "openai_chat", + reason = "zero_denominator", + "skipping proxy_cache_hit_rate_per_session: no input tokens" + ); + } } } } else { @@ -1389,48 +1454,93 @@ async fn run_sse_state_machine( } } } - // Phase G PR-G3: cache hit rate + service_tier + - // response status from the final `response.completed` - // payload. The Responses API uses `input_tokens` / + // Phase G PR-G3 + H2: cache hit rate + service_tier + + // response status emit ONLY when the stream reached a + // terminal status (`response.completed/failed/incomplete`). + // Mid-stream client disconnects close the channel without + // a terminal — `terminal_status().is_none()` then guards + // emit so we don't observe garbage samples. + // + // The Responses API uses `input_tokens` / // `cached_input_tokens` shape (Responses-specific — // distinct from Chat Completions' `prompt_tokens`). - if let Some(usage) = &state.usage { - let input_tokens = usage - .get("input_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - let cached_tokens = usage - .get("input_tokens_details") - .and_then(|d| d.get("cached_tokens")) - .and_then(|v| v.as_u64()) - .unwrap_or(0); - // Like Chat, `input_tokens` already INCLUDES cached - // tokens, so split for the helper. - let non_cached = input_tokens.saturating_sub(cached_tokens); - match crate::observability::compute_cache_hit_rate(non_cached, cached_tokens, 0) { - Some(rate) => { - crate::observability::observe_cache_hit_rate( - crate::observability::cache_hit_rate_provider::OPENAI_RESPONSES, - &request_id, - rate, - ); - } - None => { - tracing::debug!( + let stream_completed = state.terminal_status().is_some(); + if stream_completed { + if let Some(usage) = &state.usage { + let input_tokens = usage + .get("input_tokens") + .and_then(|v| v.as_u64()) + .unwrap_or(0); + let cached_tokens = usage + .get("input_tokens_details") + .and_then(|d| d.get("cached_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0); + // M1: a cached count greater than input is a + // wire-format pathology — usage shouldn't have + // `cached > input` for OpenAI Responses. Per + // "no silent fallbacks", log + skip the emit + // instead of silently clamping. + if cached_tokens > input_tokens { + tracing::warn!( event = "cache_hit_rate_skipped", request_id = %request_id, provider = "openai_responses", - reason = "zero_denominator", - "skipping proxy_cache_hit_rate_per_session: no input tokens" + reason = "cached_gt_input", + input_tokens = input_tokens, + cached_tokens = cached_tokens, + "skipping proxy_cache_hit_rate_per_session: cached_tokens > input_tokens \ + (wire-format pathology; clamping would synthesise a bad sample)" ); + } else { + // Like Chat, `input_tokens` already INCLUDES cached + // tokens, so split for the helper. + let non_cached = input_tokens - cached_tokens; + match crate::observability::compute_cache_hit_rate( + non_cached, + cached_tokens, + 0, + ) { + Some(rate) => { + crate::observability::observe_cache_hit_rate( + crate::observability::cache_hit_rate_provider::OPENAI_RESPONSES, + &request_id, + rate, + ); + } + None => { + tracing::debug!( + event = "cache_hit_rate_skipped", + request_id = %request_id, + provider = "openai_responses", + reason = "zero_denominator", + "skipping proxy_cache_hit_rate_per_session: no input tokens" + ); + } + } } } + } else { + tracing::debug!( + event = "cache_hit_rate_skipped", + request_id = %request_id, + provider = "openai_responses", + reason = "stream_did_not_complete", + "skipping proxy_cache_hit_rate_per_session: no terminal status seen" + ); } // Service tier + status are sourced from // `state.last_response_envelope` populated by the // ResponseState on `response.completed/failed/incomplete`. + // + // C1 fix: the tier value comes from the upstream response + // body; even though the upstream is more trustworthy than + // a client-side header, an unrecognised value would still + // grow the metric vector unboundedly. We bucket through + // the same validator the request-side handler uses. if let Some(tier) = state.service_tier.as_deref() { - crate::observability::record_service_tier(tier, &request_id); + let bucketed = crate::observability::metric_names::service_tier::validate(tier); + crate::observability::record_service_tier(bucketed, &request_id); } if let Some(status) = state.terminal_status() { crate::observability::record_response_status( diff --git a/crates/headroom-proxy/src/vertex/raw_predict.rs b/crates/headroom-proxy/src/vertex/raw_predict.rs index 68a86472a..5059b2a44 100644 --- a/crates/headroom-proxy/src/vertex/raw_predict.rs +++ b/crates/headroom-proxy/src/vertex/raw_predict.rs @@ -164,6 +164,7 @@ pub(crate) async fn forward_vertex_request( tokens_after, strategies_applied, markers_inserted, + .. } => { tracing::info!( event = "vertex_compression_applied", diff --git a/crates/headroom-proxy/tests/integration_metrics.rs b/crates/headroom-proxy/tests/integration_metrics.rs index f7bf30cdc..cf02f1037 100644 --- a/crates/headroom-proxy/tests/integration_metrics.rs +++ b/crates/headroom-proxy/tests/integration_metrics.rs @@ -323,41 +323,103 @@ async fn passthrough_bytes_modified_zero_when_no_compression() { assert_eq!(resp.status(), 200); let scrape = scrape_metrics(&proxy.url()).await; - // The `prometheus` crate v0.13 skips empty MetricVecs entirely - // in `gather()` — neither HELP/TYPE nor rows appear until the - // counter has been incremented with at least one label-set. - // For the "must stay 0" alarm-able metric this absence IS the - // signal: if the counter is silent, no passthrough policy has - // been violated. The PromQL alarm queries - // `rate(proxy_passthrough_bytes_modified_total[5m]) > 0`, which - // is `0` (the metric doesn't exist) by definition. - let any_row = scrape - .lines() - .any(|l| l.starts_with("proxy_passthrough_bytes_modified_total{")); + // H3 contract: `handle_metrics` force-zeroes every counter / + // gauge MetricVec under a sentinel `__init__` label tuple on + // each scrape. That makes HELP/TYPE + a zero row visible from + // boot so operators have a predictable scrape shape (the + // pre-H3 behaviour where the family was absent until first + // emit was confusing — operators would `curl /metrics` on a + // fresh boot and see nothing). The "must stay 0" alarm + // semantic is still preserved because the row only carries + // the sentinel label, not a real production path label. assert!( - !any_row, - "counter should have no rows when nothing modified passthrough; got: {scrape}" + scrape.contains("# HELP proxy_passthrough_bytes_modified_total"), + "scrape missing proxy_passthrough_bytes_modified_total HELP on fresh boot (H3): {scrape}" ); - // The HELP/TYPE descriptor is registered eagerly in - // `handle_metrics`, but `gather()` skips empty families so the - // descriptor only surfaces after the first emit. We assert the - // counter is reachable via the public helper to pin behaviour - // (the helper is what the SSE/handler emit sites call). + assert!( + scrape.contains("# TYPE proxy_passthrough_bytes_modified_total counter"), + "scrape missing proxy_passthrough_bytes_modified_total TYPE on fresh boot (H3): {scrape}" + ); + let init_row = find_value_with_labels( + &scrape, + "proxy_passthrough_bytes_modified_total", + &[("path", "__init__")], + ) + .expect("H3 contract: __init__ row must appear on fresh boot"); + assert!( + (init_row - 0.0).abs() < f64::EPSILON, + "H3 sentinel __init__ row must read 0 on fresh boot; got {init_row}" + ); + // The `/v1/messages` path label MUST NOT appear — the + // dispatcher returned NoCompression and no bytes were + // mutated, so the alarm did not fire. (Other tests in this + // suite may have populated rows under their own + // `path="/integration_test_*"` labels; we filter to the real + // production path label that THIS test would have produced.) + let messages_row = find_value_with_labels( + &scrape, + "proxy_passthrough_bytes_modified_total", + &[("path", "/v1/messages")], + ); + assert!( + messages_row.is_none(), + "counter must have no row for /v1/messages when nothing modified passthrough; \ + got value {messages_row:?}" + ); + + // Public helper still works to drive a real-path row. use headroom_proxy::observability::record_passthrough_bytes_modified; record_passthrough_bytes_modified( "/integration_test_passthrough_synthetic", - 0, + 7, "integration_test_request_id_passthrough_v1", ); - // After a (zero-delta) increment, the row should appear. let scrape_after_touch = scrape_metrics(&proxy.url()).await; + let touched = find_value_with_labels( + &scrape_after_touch, + "proxy_passthrough_bytes_modified_total", + &[("path", "/integration_test_passthrough_synthetic")], + ) + .expect("real-path row must appear after record_passthrough_bytes_modified"); assert!( - scrape_after_touch.contains("# HELP proxy_passthrough_bytes_modified_total"), - "scrape missing proxy_passthrough_bytes_modified_total HELP after touch" + touched >= 7.0, + "expected ≥7 byte delta after touch; got {touched}" ); + + proxy.shutdown().await; +} + +// ============================================================================ +// C2 wire-up: a request that is supposed to passthrough byte-equal +// but whose body bytes change is detected and the alarm fires. +// We exercise the public helper directly because forcing an +// accidental byte mutation at the dispatcher level is itself a +// regression we don't want to provoke deliberately. The helper- +// level test confirms the metric vector + label semantics, and the +// production-path test in `passthrough_bytes_modified_zero_when_no_compression` +// confirms the alarm STAYS silent on the happy path. +// ============================================================================ + +#[tokio::test] +async fn passthrough_bytes_modified_alarm_fires_with_byte_delta_label() { + use headroom_proxy::observability::record_passthrough_bytes_modified; + + // Unique path label so this test owns its row. + const TEST_PATH: &str = "/integration_test_c2_alarm_v1"; + record_passthrough_bytes_modified(TEST_PATH, 42, "integration_test_c2_request_id_v1"); + record_passthrough_bytes_modified(TEST_PATH, 13, "integration_test_c2_request_id_v2"); + + let proxy = start_proxy_with("http://127.0.0.1:1", |_| {}).await; + let scrape = scrape_metrics(&proxy.url()).await; + let value = find_value_with_labels( + &scrape, + "proxy_passthrough_bytes_modified_total", + &[("path", TEST_PATH)], + ) + .expect("C2 alarm row must appear"); assert!( - scrape_after_touch.contains("# TYPE proxy_passthrough_bytes_modified_total counter"), - "scrape missing proxy_passthrough_bytes_modified_total TYPE after touch" + (value - 55.0).abs() < f64::EPSILON, + "C2 alarm increment must reflect summed byte deltas (42 + 13 = 55); got {value}" ); proxy.shutdown().await; @@ -403,7 +465,9 @@ async fn responses_passthrough_upstream() -> (SocketAddr, tokio::task::JoinHandl } #[tokio::test] -async fn service_tier_logged() { +async fn service_tier_logged_known_value() { + // C1: spec-defined `service_tier` value lands in its own + // bucket without going through the `"other"` sentinel. let (addr, _server) = responses_passthrough_upstream().await; let proxy = start_proxy_with(&format!("http://{addr}"), |c| { c.compression_mode = headroom_proxy::config::CompressionMode::Off; @@ -411,11 +475,9 @@ async fn service_tier_logged() { }) .await; - // Unique tier value so this test owns its own row. - const TEST_TIER: &str = "integration_test_tier_priority_v1"; let body = json!({ "model": "gpt-5", - "service_tier": TEST_TIER, + "service_tier": "priority", "input": [ {"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]} ] @@ -434,9 +496,9 @@ async fn service_tier_logged() { let tier_count = find_value_with_labels( &scrape, "proxy_service_tier_count_total", - &[("tier", TEST_TIER)], + &[("tier", "priority")], ) - .expect("service tier counter row must appear"); + .expect("service tier counter row must appear under 'priority'"); assert!( tier_count >= 1.0, "expected ≥1 service_tier increment; got {tier_count}" @@ -445,6 +507,99 @@ async fn service_tier_logged() { proxy.shutdown().await; } +#[tokio::test] +async fn service_tier_unknown_bucketed_to_other() { + // C1: a malicious or drifting client sends an unrecognised + // service_tier value. The bounded-vocabulary validator MUST + // bucket it to "other" so a malicious client can't blow up + // the metric vector cardinality. + let (addr, _server) = responses_passthrough_upstream().await; + let proxy = start_proxy_with(&format!("http://{addr}"), |c| { + c.compression_mode = headroom_proxy::config::CompressionMode::Off; + c.enable_responses_streaming = true; + }) + .await; + + // Two distinct unknown values — both must bucket to "other". + for unknown_tier in [ + "integration_test_unknown_tier_alpha_v1", + "integration_test_unknown_tier_beta_v1", + ] { + let body = json!({ + "model": "gpt-5", + "service_tier": unknown_tier, + "input": [ + {"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]} + ] + }); + let resp = reqwest::Client::new() + .post(format!("{}/v1/responses", proxy.url())) + .header("content-type", "application/json") + .body(serde_json::to_vec(&body).unwrap()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + let _ = resp.bytes().await.unwrap(); + } + + let scrape = scrape_metrics(&proxy.url()).await; + // Neither raw value may appear as a label — they must be + // bucketed. + assert!( + !scrape.contains("integration_test_unknown_tier_alpha_v1"), + "raw unknown tier value leaked into metrics (cardinality DoS): {scrape}" + ); + assert!( + !scrape.contains("integration_test_unknown_tier_beta_v1"), + "raw unknown tier value leaked into metrics (cardinality DoS): {scrape}" + ); + // The "other" bucket must have been incremented at least twice. + let other_count = find_value_with_labels( + &scrape, + "proxy_service_tier_count_total", + &[("tier", "other")], + ) + .expect("'other' bucket row must appear"); + assert!( + other_count >= 2.0, + "expected ≥2 increments on 'other' bucket (one per unknown tier sent); got {other_count}" + ); + + proxy.shutdown().await; +} + +#[test] +fn service_tier_validate_known_returns_canonical_constant() { + use headroom_proxy::observability::metric_names::service_tier; + // Spec-defined values pass through verbatim — strict equality + // against the &'static constants so a typo in the validator + // surfaces here. + assert_eq!(service_tier::validate("auto"), service_tier::AUTO); + assert_eq!(service_tier::validate("default"), service_tier::DEFAULT); + assert_eq!(service_tier::validate("flex"), service_tier::FLEX); + assert_eq!(service_tier::validate("on_demand"), service_tier::ON_DEMAND); + assert_eq!(service_tier::validate("priority"), service_tier::PRIORITY); + assert_eq!(service_tier::validate("scale"), service_tier::SCALE); +} + +#[test] +fn service_tier_validate_unknown_returns_other_sentinel() { + use headroom_proxy::observability::metric_names::service_tier; + // C1: anything outside the bounded vocab buckets to OTHER. + assert_eq!( + service_tier::validate("nonsense_value"), + service_tier::OTHER + ); + assert_eq!(service_tier::validate(""), service_tier::OTHER); + // Case-sensitive: spec is case-sensitive on these strings. + assert_eq!(service_tier::validate("PRIORITY"), service_tier::OTHER); + assert_eq!(service_tier::validate("Auto"), service_tier::OTHER); + // Extremely long / arbitrary attacker input → still bucketed. + let attack = "A".repeat(10_000); + assert_eq!(service_tier::validate(&attack), service_tier::OTHER); +} + // ============================================================================ // Test 5: incomplete_status_logged_with_reason // ============================================================================ @@ -503,6 +658,73 @@ async fn incomplete_status_logged_with_reason() { // Bonus coverage: rate-limit gauge plumbing via the public helper. // ============================================================================ +// ============================================================================ +// H1 per-strategy ratio: drive `observe_compression_ratio` twice +// with different strategy names + tokens, and assert each strategy +// row has its OWN sum (not the same aggregate replicated). +// ============================================================================ + +#[tokio::test] +async fn compression_ratio_per_strategy_does_not_replicate_aggregate() { + // Pre-H1 the proxy emitted the same `aggregate` ratio per + // strategy when multiple strategies ran on one body. This test + // exercises the helper directly with two distinct strategies + + // distinct ratios so the histogram's _sum lines for each + // strategy must differ. + use headroom_proxy::observability::observe_compression_ratio; + + const STRAT_HEAVY: &str = "h1_test_heavy_v1"; + const STRAT_LIGHT: &str = "h1_test_light_v1"; + const CT: &str = "h1_test_content_type_v1"; + + // Strategy A: original=1000 tokens → compressed=200 (ratio 0.20). + observe_compression_ratio(STRAT_HEAVY, CT, 1000, 200); + // Strategy B: original=1000 tokens → compressed=800 (ratio 0.80). + observe_compression_ratio(STRAT_LIGHT, CT, 1000, 800); + + let proxy = start_proxy_with("http://127.0.0.1:1", |_| {}).await; + let scrape = scrape_metrics(&proxy.url()).await; + + let sum_heavy = find_value_with_labels( + &scrape, + "proxy_compression_ratio_by_strategy_sum", + &[("strategy", STRAT_HEAVY), ("content_type", CT)], + ) + .expect("heavy-strategy sum row"); + let sum_light = find_value_with_labels( + &scrape, + "proxy_compression_ratio_by_strategy_sum", + &[("strategy", STRAT_LIGHT), ("content_type", CT)], + ) + .expect("light-strategy sum row"); + // The sums must NOT be equal — if they are, the per-strategy + // wiring regressed to the pre-H1 "emit same aggregate per + // strategy" behavior. + assert!( + (sum_heavy - sum_light).abs() > 1e-9, + "per-strategy sums are equal: heavy={sum_heavy} light={sum_light} — \ + H1 regression: did we re-emit the aggregate ratio per strategy?" + ); + // Spot-check the actual ratios. + assert!( + sum_heavy < sum_light, + "heavy strategy ratio (0.20) must be < light strategy ratio (0.80): \ + heavy={sum_heavy} light={sum_light}" + ); + + proxy.shutdown().await; +} + +// ============================================================================ +// H2 aborted stream: a client disconnect mid-stream closes the +// channel without `message_stop`. Unit-tested in +// `crate::observability::cache_hit_rate::tests` because the +// integration-level approach is flaky against the shared global +// Prometheus registry (other tests in the suite emit on the same +// `provider="anthropic"` label, making delta-based assertions +// non-deterministic). +// ============================================================================ + #[tokio::test] async fn rate_limit_snapshot_emits_gauges() { let snap = observability::RateLimitSnapshot { diff --git a/docs/observability.md b/docs/observability.md index 9582f1302..33b2f20f6 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -35,6 +35,14 @@ rename catches one file in code review. |------|------|--------|---------| | `proxy_passthrough_bytes_modified_total` | Counter | `path` | Bytes mutated on a passthrough path. **Must stay 0 outside the compression hot path** — any non-zero rate fires the cache-safety alarm. | +The alarm metric is wired in `crates/headroom-proxy/src/proxy.rs`: +when the dispatcher returns `Outcome::NoCompression` or +`Outcome::Passthrough`, the post-dispatcher byte length is compared +to the original buffered length and any delta increments the +counter (by the byte delta) under the request's path label. The +PR-E4 prompt_cache_key injector runs AFTER the alarm check, so its +intentional byte mutations do not trip the alarm. + #### Upstream rate limits | Name | Type | Labels | Purpose | @@ -51,17 +59,29 @@ rename catches one file in code review. | `proxy_service_tier_count_total` | Counter | `tier` | Service-tier distribution observed at the proxy. | | `proxy_response_status_count_total` | Counter | `status` | Terminal status distribution (`completed`, `incomplete`, `failed`, `cancelled`, `in_progress`). | -#### Wrap CLI / RTK +#### Wrap CLI / RTK (Python-side) | Name | Type | Labels | Purpose | |------|------|--------|---------| -| `wrap_rtk_invocations_total` | Counter | `tool` | RTK invocations observed via `rtk gain --format json` polling. Driven from the wrap-CLI tail. | +| `wrap_rtk_invocations_total` | Counter | `tool` | RTK invocations observed via the wrap-CLI tail. Surfaced via the Python proxy's `/metrics` exporter; the wrap CLI bumps `headroom.cli.wrap_rtk_metrics.record_rtk_invocation(...)`. | -#### Image log redaction +> **C4 remediation:** This counter is Python-side because RTK is +> wrapped by `headroom wrap` (Python CLI) and the wrap-side tail +> is the natural emit site. The Rust proxy previously held a dead +> counter for this metric; that has been removed. + +#### Image log redaction (Python-side) | Name | Type | Labels | Purpose | |------|------|--------|---------| -| `proxy_image_generation_call_log_redacted_total` | Counter | _none_ | Base64-encoded image payloads redacted from request logs. | +| `proxy_image_generation_call_log_redacted_total` | Counter | _none_ | Base64-encoded image payloads redacted from request logs. Driven from `headroom.proxy.request_logger.redactions_total()`. | + +> **C3 remediation:** Image redaction is purely a Python-proxy +> operation (the request logger walks JSON and replaces over- +> threshold image payloads with placeholders). The counter lives +> Python-side so we have one source of truth instead of two. The +> Rust proxy previously held a dead counter for this metric; that +> has been removed. ## How to query @@ -71,32 +91,155 @@ The proxy renders Prometheus text-format on `GET /metrics`: curl -s http://127.0.0.1:8787/metrics ``` -Common PromQL queries: +### Phase H canary gate + +The canary script that decides "ship Rust, retire Python" uses +**all four** of these queries against `proxy_cache_hit_rate_per_session` +to confirm parity vs the Python baseline. A single percentile is +not enough — a regression that only shows up at the tail (a small +class of long sessions losing cache hits) would slip through a +median-only check. ```promql -# Phase H canary gate — cache hit rate parity vs Python proxy. -histogram_quantile(0.50, sum by (le) (rate(proxy_cache_hit_rate_per_session_bucket[5m]))) +# p50, p95, p99 of cache hit rate over the last 5 minutes, per provider. +histogram_quantile(0.50, sum by (provider, le) (rate(proxy_cache_hit_rate_per_session_bucket{provider!="__init__"}[5m]))) +histogram_quantile(0.95, sum by (provider, le) (rate(proxy_cache_hit_rate_per_session_bucket{provider!="__init__"}[5m]))) +histogram_quantile(0.99, sum by (provider, le) (rate(proxy_cache_hit_rate_per_session_bucket{provider!="__init__"}[5m]))) -# Cache-safety alarm. Should always be 0. -sum(rate(proxy_passthrough_bytes_modified_total[5m])) +# Mean cache hit rate over the last 5 minutes, per provider. The +# `sum / count` form is the cleanest "average without a quantile" +# query and is what the Python baseline reports. +sum by (provider) (rate(proxy_cache_hit_rate_per_session_sum{provider!="__init__"}[5m])) + / +sum by (provider) (rate(proxy_cache_hit_rate_per_session_count{provider!="__init__"}[5m])) +``` -# Aggregate compression value by strategy. -histogram_quantile(0.50, sum by (strategy, le) (rate(proxy_compression_ratio_by_strategy_bucket[1h]))) +The canary fails if ANY of `p50`, `p95`, `p99`, or `mean` regresses +below the Python baseline for any provider over the canary window. + +### Other common queries + +```promql +# Cache-safety alarm. Should always be 0 (post-`__init__` row). +sum(rate(proxy_passthrough_bytes_modified_total{path!="__init__"}[5m])) + +# Per-strategy compression value at p50 (post-H1 fix: each strategy +# reports its own before/after; pre-fix this was the same aggregate +# ratio repeated per strategy). +histogram_quantile(0.50, sum by (strategy, le) (rate(proxy_compression_ratio_by_strategy_bucket{strategy!="__init__"}[1h]))) + +# Per-strategy compression value at p95 and p99 (catch outlier +# strategies that fail to shrink at the tail). +histogram_quantile(0.95, sum by (strategy, le) (rate(proxy_compression_ratio_by_strategy_bucket{strategy!="__init__"}[1h]))) +histogram_quantile(0.99, sum by (strategy, le) (rate(proxy_compression_ratio_by_strategy_bucket{strategy!="__init__"}[1h]))) + +# Strategies that ran but failed the token-check (compressor ran +# but its output was not strictly smaller, so the original was +# kept). High rate here means the compressor needs tuning. +sum by (strategy) (rate(proxy_compression_rejected_by_token_check_total{strategy!="__init__"}[1h])) # Upstream rate-limit headroom (smaller = closer to throttle). proxy_rate_limit_remaining_tokens{provider="anthropic"} + +# RTK invocation rate (Python-side). +sum by (tool) (rate(wrap_rtk_invocations_total{tool!="__init__"}[5m])) + +# Image-redaction rate (Python-side). +rate(proxy_image_generation_call_log_redacted_total[5m]) ``` +All queries above include a `{... != "__init__"}` filter so the +sentinel zero-rows the boot-touch contract emits do not skew the +result. See "Wiring → H3 force-zero" below. + ## Wiring Every metric registration is `OnceLock`-backed and lazy: the first call to a `*_counter()` / `*_gauge()` / `*_histogram()` helper registers the family with the shared registry. `handle_metrics` -force-touches every Phase G PR-G3 family before scraping so the -descriptors are reachable from `/metrics` on a fresh boot — note that -the `prometheus` crate v0.13 skips empty MetricVecs from `gather()`, -so HELP/TYPE only surface once a family has been incremented at -least once. +force-touches every Phase G PR-G3 family before scraping. + +### H3 force-zero + +The `prometheus` crate v0.13 skips empty MetricVecs from `gather()` +entirely — neither HELP/TYPE lines nor rows appear until the +family has been incremented at least once with a label tuple. +Operators expect to see the catalogue from boot, so +`handle_metrics` increments each counter / gauge MetricVec by 0 +under a sentinel `__init__` label tuple before the first scrape. +HELP/TYPE then surface from boot and dashboards/alarms see a +predictable scrape shape. + +Counters with the `__init__` label increment by 0, so the +alarm-able "must stay 0" semantic of +`proxy_passthrough_bytes_modified_total` is preserved (the family +becomes visible, the rate stays 0). PromQL queries should filter +`{... != "__init__"}` so the sentinel rows are excluded from +aggregations (the catalogue above does this). + +Histograms are NOT force-zeroed: a synthetic `observe(0.0)` would +contribute a real sample to the per-label distribution and pollute +percentile readings. The two histogram families +(`proxy_cache_hit_rate_per_session` and +`proxy_compression_ratio_by_strategy`) only surface in the scrape +after the first real session, by design. + +### H4 prometheus crate version pin + +The H3 contract above relies on the `prometheus` crate's v0.13 +`gather()` semantics — empty MetricVec families are omitted from +the scrape. **This is implementation-defined behaviour.** If +`crates/headroom-proxy/Cargo.toml` ever bumps the `prometheus` +dependency, retest the alarm contract: + +1. Start a fresh proxy. +2. `curl /metrics` and confirm every counter / gauge family has + HELP/TYPE + an `__init__` row. +3. Confirm histograms (`*_cache_hit_rate_per_session`, + `*_compression_ratio_by_strategy`) DO NOT appear (no + `observe()` calls yet). +4. Drive one cache-hit session, scrape again, confirm histograms + now appear. +5. Confirm `passthrough_bytes_modified_total` stays at 0 across + passthrough requests. + +The crate version is pinned exactly (`= "0.13.4"`, no caret) in +`Cargo.toml` precisely so a silent semver bump cannot break the +contract without a code-review trigger. + +### C2 alarm wiring + +`proxy_passthrough_bytes_modified_total` fires from `proxy.rs` when +a dispatcher arm that promised byte-equal passthrough +(`Outcome::NoCompression` or `Outcome::Passthrough`) produces a +final body of a different byte length. The check runs BEFORE the +PR-E4 prompt_cache_key injector so the injector's intentional byte +mutations do not trip the alarm. + +### H1 per-strategy ratio wiring + +`proxy_compression_ratio_by_strategy` samples one observation per +strategy using the strategy's OWN before/after token counts +(plumbed through `Outcome::Compressed.per_strategy_tokens` from +the manifest in `live_zone_anthropic` / `live_zone_openai` / +`live_zone_responses`). Pre-H1 the same aggregate ratio was +emitted per strategy when multiple strategies ran on one body, +making Phase H per-strategy dashboards read garbage. + +### H2 aborted-stream gate + +The `proxy_cache_hit_rate_per_session` histogram observes ONLY +when the SSE stream completed: + +* Anthropic: `state.status == StreamStatus::MessageStop` after the + channel closes. +* OpenAI Chat: `state.usage.is_some()` (the final usage chunk only + arrives at stream completion). +* OpenAI Responses: `state.terminal_status().is_some()`. + +A client disconnect mid-stream closes the channel without setting +the terminal flag — under H2 we log + skip rather than observe a +garbage half-stream sample. ## Cardinality discipline @@ -107,8 +250,18 @@ Every label vocabulary is bounded by code, not customer input: - `provider`: 3 values (`anthropic`, `openai_chat`, `openai_responses`). - `strategy`: `&'static str` from the compressor's `BlockAction::Compressed`. - `content_type`: `&'static str` from `headroom_core::transforms::ContentType`. -- `tier`: bounded by OpenAI Responses spec. +- `tier`: validated through + `crate::observability::metric_names::service_tier::validate(raw: &str)`. + Returns one of `{auto, default, flex, on_demand, priority, scale}` + or the sentinel `"other"` for anything else. **The raw inbound + value is never used as a label.** A malicious client posting + `{"service_tier":""}` per request gets bucketed to + `"other"` and a `tracing::warn!` is emitted so wire-format drift + surfaces loudly in logs. - `status`: 5-variant enum. +- `tool` (Python-side `wrap_rtk_invocations_total`): bounded by the + set of tools the wrap CLI rewrites, captured by + `headroom.cli.wrap_rtk_metrics`. There is no code path where a malicious client can drive label cardinality unbounded. diff --git a/headroom/cli/wrap_rtk_metrics.py b/headroom/cli/wrap_rtk_metrics.py new file mode 100644 index 000000000..64c788788 --- /dev/null +++ b/headroom/cli/wrap_rtk_metrics.py @@ -0,0 +1,73 @@ +"""RTK invocation metrics for the wrap CLI. + +Phase G PR-G3 remediation (C4): RTK lives wrap-side, not proxy-side +(see ``docs/rtk-architecture.md``). The wrap CLI tails +``rtk gain --format json`` and bumps a process-local counter keyed +by rewritten command name (`git`, `ls`, `cargo`, ...). The Python +proxy's ``/metrics`` endpoint then surfaces the counter as +``wrap_rtk_invocations_total{tool=...}`` for fleet-wide scrape. + +The counter primitives live here (not in ``wrap.py``) so the +proxy's prometheus exporter can import them without dragging in the +full ``wrap.py`` module — that module owns subprocess-level CLI +spawning and is heavyweight to import at proxy startup. + +Per realignment build-constraint "no silent fallbacks": +``record_rtk_invocation`` raises on a non-string tool name rather +than coercing; a caller passing the wrong type is a bug, not a +runtime fallback condition. +""" + +from __future__ import annotations + +import threading +from collections import defaultdict +from collections.abc import Mapping + +# Module-level counter — process-local. Multiple worker processes +# (uvicorn workers) each maintain their own; the Python proxy already +# documents this in ``docs/observability.md``. Reset is exposed for +# tests; production code never reaches for it. +_rtk_invocation_counts: dict[str, int] = defaultdict(int) +_lock = threading.Lock() + + +def record_rtk_invocation(tool: str, delta: int = 1) -> None: + """Record one (or `delta`) RTK invocation(s) for the given tool. + + `tool` is the rewritten command name as observed in the + ``rtk gain --format json`` output (e.g. ``"git"``, ``"ls"``, + ``"cargo"``). The counter is keyed verbatim. + + `delta` defaults to 1 for the common "one invocation seen" path + but accepts arbitrary positive deltas so the wrap tail can bump + by a JSON-reported batch count. + + Raises: + TypeError: if `tool` is not a `str` or `delta` is not an `int`. + ValueError: if `delta` is negative. + """ + if not isinstance(tool, str): + raise TypeError(f"tool must be a str, got {type(tool).__name__}") + if not isinstance(delta, int): + raise TypeError(f"delta must be an int, got {type(delta).__name__}") + if delta < 0: + raise ValueError(f"delta must be non-negative, got {delta}") + with _lock: + _rtk_invocation_counts[tool] += delta + + +def rtk_invocation_counts() -> Mapping[str, int]: + """Return a snapshot of the current invocation counts. + + Returns a plain dict (not the defaultdict) so callers cannot + accidentally pollute the counter map by reading absent keys. + """ + with _lock: + return dict(_rtk_invocation_counts) + + +def reset_rtk_invocations() -> None: + """Reset the counter map. Test-only — never called from production.""" + with _lock: + _rtk_invocation_counts.clear() diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 600b485f8..27ad9fb1c 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -1215,4 +1215,56 @@ class PrometheusMetrics: ) lines.append("") + # Phase G PR-G3 remediation (C3): image-redacted counter + # lives Python-side because base64 redaction is purely a + # Python-proxy concern (request_logger.py). The Rust + # proxy previously held a dead counter for this; that's + # been removed in favour of this Python export. + # + # The counter is read at scrape-time from the module- + # level redaction tracker rather than mirrored into the + # PrometheusMetrics instance, so we never lose a count + # to ordering between RequestLogger setup and metrics + # init. + from headroom.proxy.request_logger import redactions_total + + _append_metric( + lines, + name="proxy_image_generation_call_log_redacted_total", + metric_type="counter", + help_text=( + "Count of base64-encoded image payloads redacted from request " + "logs by the Python proxy's request logger" + ), + value=redactions_total(), + ) + + # Phase G PR-G3 remediation (C4): RTK invocations counter + # also lives Python-side. RTK is wrapped by the + # `headroom wrap` CLI (headroom.cli.wrap); the proxy + # observes invocation counts via a process-local tracker + # the wrap tail bumps. The Rust proxy previously held a + # dead counter for this; that's been removed. + from headroom.cli.wrap_rtk_metrics import rtk_invocation_counts + + counts = rtk_invocation_counts() + lines.extend( + [ + "# HELP wrap_rtk_invocations_total RTK invocations observed via the wrap CLI tail", + "# TYPE wrap_rtk_invocations_total counter", + ] + ) + if not counts: + # Emit a zero-row under the sentinel tool name so + # the family advertises HELP/TYPE on a fresh boot + # and dashboards can probe it before any RTK + # invocation has happened. Matches the Rust side's + # H3 force-zero contract. + lines.append('wrap_rtk_invocations_total{tool="__init__"} 0') + else: + for tool, count in counts.items(): + safe_tool = _escape_label_value(str(tool)) + lines.append(f'wrap_rtk_invocations_total{{tool="{safe_tool}"}} {count}') + lines.append("") + return "\n".join(lines) diff --git a/headroom/proxy/request_logger.py b/headroom/proxy/request_logger.py index a29b49871..c367b6f4d 100644 --- a/headroom/proxy/request_logger.py +++ b/headroom/proxy/request_logger.py @@ -8,6 +8,15 @@ Phase G PR-G3 (P4-45): base64-encoded image payloads in the ``request_messages`` / ``response_content`` are redacted before write to keep request logs small. Multi-MB base64 strings would otherwise saturate the JSONL log and the in-memory deque. + +Remediation (M2, M5): the redactor now ONLY fires inside known +image-bearing JSON paths or against strings that carry an explicit +``data:image/...;base64,`` URL prefix. The earlier "density +heuristic" over-fired on encrypted blobs, signed tokens, minified +JSON, and tool outputs. The replacement placeholder now reports +the UTF-8 byte length under a ``bytes=`` label (was character +length; for the ASCII base64 alphabet the two happen to coincide +but the label is now accurate for any future Unicode payload). """ from __future__ import annotations @@ -19,6 +28,7 @@ from collections import deque from collections.abc import Mapping, Sequence from dataclasses import asdict from pathlib import Path +from threading import Lock from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -40,74 +50,124 @@ IMAGE_BASE64_REDACT_THRESHOLD_BYTES = 1024 # Phase G PR-G3 — replacement-marker format. Operators can grep the # JSONL for `` int: """Return the running count of base64 redactions performed. - Exposed for unit tests + the legacy Python ``/stats`` endpoint. - The canonical observability surface is the Rust proxy's - ``proxy_image_generation_call_log_redacted_total`` metric. + Exposed for unit tests, the legacy Python ``/stats`` endpoint, + and the Prometheus exporter + (``proxy_image_generation_call_log_redacted_total``). """ - return _redactions_total + with _redactions_lock: + return _redactions_total -def _looks_like_base64_image(value: str) -> bool: - """Heuristic: does ``value`` look like a base64-encoded image? +def _is_base64_image_payload(value: str) -> bool: + """Return True if ``value`` is an over-threshold base64 image. - Two patterns we recognise: + Per M2 remediation the prior bare-base64 density heuristic + over-fired on non-image content (encrypted blobs, signed + tokens, minified JSON, tool outputs). We now only consider a + string an image payload when EITHER: - * Raw base64 over the threshold (Anthropic ``source.data`` shape). - * ``data:image/;base64,`` data URLs (OpenAI - vision shape). The ``;base64,`` substring is the load-bearing - signal — any data URL with that segment over the threshold gets - redacted, even if the MIME type isn't ``image/...`` (because - the cost-of-logging is paid the same way). + 1. It starts with ``data:image/`` (an explicit data URL), + OR + 2. The caller has already established the string lives inside + an image-bearing JSON path (see ``IMAGE_BEARING_FIELD_NAMES``) + AND the string itself is over the byte threshold. - Returns ``False`` for short strings (under the threshold) and for - any non-string value. Per realignment build-constraint "no - regexes", we use prefix/substring checks instead of pattern - matching. + Case (2) is decided by the caller (``_redact_value``) which + threads ``in_image_path`` through the recursion; this helper + handles case (1) on its own. """ if not isinstance(value, str): return False if len(value) < IMAGE_BASE64_REDACT_THRESHOLD_BYTES: return False - if value.startswith("data:") and ";base64," in value[:64]: - return True - # Bare base64 payload: heuristic — over-threshold string with no - # whitespace and a high alpha-num+/+= density. We sample the first - # 256 bytes for speed (the full string can be megabytes). - head = value[:256] - base64_chars = set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=") - matches = sum(1 for ch in head if ch in base64_chars) - return matches / len(head) > 0.95 + return value.startswith(_DATA_IMAGE_URL_PREFIX) -def _redact_value(value: Any) -> Any: +def _redact_value(value: Any, *, in_image_path: bool = False) -> Any: """Recursively redact base64-image payloads in a JSON-ish value. Returns a new structure with any over-threshold base64 string replaced by the placeholder. Non-string, non-container values pass through unchanged. + + ``in_image_path`` is True when the caller reached this value + via one of the ``IMAGE_BEARING_FIELD_NAMES`` keys; once inside + an image-bearing field, any over-threshold string is treated + as an image payload (M2: prevents redaction of unrelated + base64-shaped content outside known image fields). """ global _redactions_total if isinstance(value, str): - if _looks_like_base64_image(value): - _redactions_total += 1 - return IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=len(value)) + # Always-redact: explicit data URL, regardless of path. + # Also redact when the caller signalled image-bearing path + # AND the string is over threshold (no density check — the + # path tells us it's an image). + should_redact = _is_base64_image_payload(value) or ( + in_image_path and len(value) >= IMAGE_BASE64_REDACT_THRESHOLD_BYTES + ) + if should_redact: + with _redactions_lock: + _redactions_total += 1 + byte_len = len(value.encode("utf-8")) + return IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=byte_len) return value if isinstance(value, Mapping): - return {k: _redact_value(v) for k, v in value.items()} + return { + k: _redact_value( + v, + in_image_path=(k in IMAGE_BEARING_FIELD_NAMES), + ) + for k, v in value.items() + } if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): - return [_redact_value(item) for item in value] + return [_redact_value(item, in_image_path=in_image_path) for item in value] return value @@ -118,7 +178,7 @@ def redact_image_base64(payload: Any) -> Any: over-threshold base64 string with a size-only placeholder. Idempotent — applying twice yields the same structure. """ - return _redact_value(payload) + return _redact_value(payload, in_image_path=False) class RequestLogger: diff --git a/tests/test_cli/test_wrap_rtk_metrics.py b/tests/test_cli/test_wrap_rtk_metrics.py new file mode 100644 index 000000000..f864e52b7 --- /dev/null +++ b/tests/test_cli/test_wrap_rtk_metrics.py @@ -0,0 +1,113 @@ +"""Phase G PR-G3 remediation (C4) — wrap-CLI RTK metrics primitive. + +The Rust proxy previously held a dead `wrap_rtk_invocations_total` +counter. C4 remediation moved it Python-side because the wrap CLI +(headroom.cli.wrap) is where RTK invocations are actually counted. +These tests cover the counter primitives in isolation. +""" + +from __future__ import annotations + +import threading + +import pytest + +from headroom.cli.wrap_rtk_metrics import ( + record_rtk_invocation, + reset_rtk_invocations, + rtk_invocation_counts, +) + + +@pytest.fixture(autouse=True) +def _reset_between_tests(): + """Reset the module-level counter map between tests so each + test owns a clean slate.""" + reset_rtk_invocations() + yield + reset_rtk_invocations() + + +def test_record_increments_default_delta_one(): + record_rtk_invocation("git") + counts = rtk_invocation_counts() + assert counts == {"git": 1} + + +def test_record_accumulates_per_tool(): + record_rtk_invocation("git") + record_rtk_invocation("git") + record_rtk_invocation("ls") + record_rtk_invocation("cargo") + record_rtk_invocation("cargo") + record_rtk_invocation("cargo") + counts = rtk_invocation_counts() + assert counts == {"git": 2, "ls": 1, "cargo": 3} + + +def test_record_with_explicit_delta(): + record_rtk_invocation("git", delta=5) + record_rtk_invocation("git", delta=2) + counts = rtk_invocation_counts() + assert counts == {"git": 7} + + +def test_record_zero_delta_is_noop_record(): + # delta=0 is legal — caller may want to "touch" the counter to + # ensure the key exists before later increments. + record_rtk_invocation("git", delta=0) + counts = rtk_invocation_counts() + assert counts == {"git": 0} + + +def test_record_rejects_negative_delta(): + with pytest.raises(ValueError, match="must be non-negative"): + record_rtk_invocation("git", delta=-1) + + +def test_record_rejects_non_string_tool(): + with pytest.raises(TypeError, match="tool must be a str"): + record_rtk_invocation(123, delta=1) # type: ignore[arg-type] + + +def test_record_rejects_non_int_delta(): + with pytest.raises(TypeError, match="delta must be an int"): + record_rtk_invocation("git", delta="1") # type: ignore[arg-type] + + +def test_counts_returns_snapshot_not_view(): + # The returned mapping must be a plain dict copy, not the + # internal defaultdict — otherwise callers could pollute the + # counter map by reading absent keys. + record_rtk_invocation("git") + counts = rtk_invocation_counts() + # Reading a key that's not present must not add it to the + # internal map. + _ = counts.get("nonexistent_tool", 0) + counts2 = rtk_invocation_counts() + assert "nonexistent_tool" not in counts2 + + +def test_thread_safe_concurrent_increments(): + # 10 threads each bumping `git` 100 times: final count must be + # exactly 1000. The threading.Lock guards the dict update so + # races are impossible. + def worker(): + for _ in range(100): + record_rtk_invocation("git") + + threads = [threading.Thread(target=worker) for _ in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + counts = rtk_invocation_counts() + assert counts == {"git": 1000} + + +def test_reset_clears_counts(): + record_rtk_invocation("git", delta=42) + record_rtk_invocation("ls", delta=7) + assert rtk_invocation_counts() != {} + reset_rtk_invocations() + assert rtk_invocation_counts() == {} diff --git a/tests/test_image_log_redaction.py b/tests/test_image_log_redaction.py index cd6bf9c60..b829322fc 100644 --- a/tests/test_image_log_redaction.py +++ b/tests/test_image_log_redaction.py @@ -124,7 +124,8 @@ def test_data_url_redacted(): def test_redact_idempotent(): """Applying redaction twice yields the same structure — the placeholder is short enough to stay below the threshold so the - second pass is a no-op.""" + second pass is a no-op. The ``data`` key is one of the + image-bearing field names so a big string inside redacts.""" big = _big_base64(IMAGE_BASE64_REDACT_THRESHOLD_BYTES * 3) once = redact_image_base64({"data": big}) twice = redact_image_base64(once) @@ -165,15 +166,81 @@ def test_logger_writes_redacted_payload_to_jsonl(): assert data_field == IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=len(big)) -def test_response_content_redacted(): - """``response_content`` (a string field) is redacted in-place on - the deque entry.""" +def test_response_content_bare_base64_passes_through(): + """M2 remediation: a bare base64-shaped string in + ``response_content`` is NOT redacted. The earlier "density + heuristic" over-fired on encrypted blobs, signed tokens, + minified JSON, and tool outputs. The new contract: only + redact strings inside known image-bearing JSON paths OR + strings starting with ``data:image/``.""" logger = RequestLogger(log_file=None, log_full_messages=True) big = _big_base64(IMAGE_BASE64_REDACT_THRESHOLD_BYTES * 3) entry = _make_request_log(response_content=big) logger.log(entry) recent = logger.get_recent_with_messages(n=1) - assert recent[0]["response_content"] == IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=len(big)) + # Verbatim — no redaction applied. + assert recent[0]["response_content"] == big + + +def test_response_content_data_image_url_redacted(): + """When ``response_content`` does start with ``data:image/`` — + e.g. a tool wrote an image back via a data URL — redaction + still fires.""" + logger = RequestLogger(log_file=None, log_full_messages=True) + payload = _big_base64(IMAGE_BASE64_REDACT_THRESHOLD_BYTES * 2) + data_url = f"data:image/png;base64,{payload}" + entry = _make_request_log(response_content=data_url) + logger.log(entry) + recent = logger.get_recent_with_messages(n=1) + assert recent[0]["response_content"] == IMAGE_BASE64_REPLACEMENT_TEMPLATE.format( + n=len(data_url) + ) + + +def test_non_image_path_base64_passes_through(): + """M2: a big base64-shaped string at a non-image-bearing key + (e.g. an encrypted blob under ``signature`` or a tool output + under ``arguments``) is NOT redacted.""" + big = _big_base64(IMAGE_BASE64_REDACT_THRESHOLD_BYTES * 2) + payload = { + "tool_use_id": "tool_xyz", + "signature": big, # NOT an image-bearing key + "arguments": big, # NOT an image-bearing key + } + redacted = redact_image_base64(payload) + assert redacted["signature"] == big + assert redacted["arguments"] == big + + +def test_image_path_redacts_without_density_check(): + """M2: once inside an image-bearing JSON path (e.g. + ``source.data``), a sufficiently-long string is redacted + regardless of its character density. Real images may not be + base64 only — they may be webp / avif transcoded with + different alphabets — but we still want them redacted to + keep logs bounded.""" + big = "x" * (IMAGE_BASE64_REDACT_THRESHOLD_BYTES * 2) # NOT base64 + payload = {"source": {"type": "base64", "data": big}} + redacted = redact_image_base64(payload) + expected_bytes = len(big.encode("utf-8")) + assert redacted["source"]["data"] == IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=expected_bytes) + + +def test_byte_count_label_is_utf8_bytes_not_chars(): + """M5: the ``bytes=`` label is the UTF-8 byte length of the + redacted string, not the character count. For ASCII base64 + payloads the two coincide (so existing tests still pass), but + a non-ASCII string under an image-bearing key reports byte + length faithfully.""" + # 3-byte UTF-8 character ('€' = U+20AC) repeated; the + # character count is half the byte count. + chars = "€" * (IMAGE_BASE64_REDACT_THRESHOLD_BYTES + 1) + payload = {"data": chars} + redacted = redact_image_base64(payload) + char_count = len(chars) + byte_count = len(chars.encode("utf-8")) + assert byte_count == 3 * char_count + assert redacted["data"] == IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=byte_count) def test_redactions_counter_advances():