fix(count_tokens): emit the usage event a refusal needs to be findable

crates/aisix-proxy/src/count_tokens.rs

@@ -32,6 +32,20 @@
 //! also keeps the answer honest: `/v1/messages` masks the same spans, so the
 //! count now describes the body the gateway would really send.
 //!
+//! Telemetry: this route is NOT metered and still emits a terminal
+//! `UsageEvent` on every outcome, with `prompt_tokens`/`completion_tokens`
+//! at zero. Those are two different questions, and the route answered only
+//! the first: it generates nothing, so there is nothing to bill — but it
+//! does forward the caller's whole payload to a real upstream, so every
+//! question Logs exists to answer (did this request happen, which key sent
+//! it, how long did it take, did a guardrail refuse it) had no row to read.
+//! A refusal was the sharp end: `/v1/messages/count_tokens` can be blocked
+//! by an input guardrail, and a refusal that emits no event is a 422 the
+//! caller definitely saw and the "Guardrail blocks" view cannot find
+//! (AISIX-Cloud#1435, the same failure mode as AISIX-Cloud#1428).
+//! `guardrail_coverage`'s census asserts the reporting half over the
+//! surfaces it reads out of the router, so this cannot regress quietly.
+//!
 //! Scope: Anthropic-backed models only. `count_tokens` has no upstream
 //! equivalent for OpenAI/Gemini/DeepSeek, so a non-Anthropic Model is
 //! rejected with a 400 at the gateway boundary (parallel to `/v1/rerank`
@@ -107,7 +121,20 @@ pub async fn count_tokens(
     // One snapshot for the whole request (#941) — see `embeddings`.
     let snapshot = state.snapshot.load();
 
-    match dispatch(&state, &snapshot, &auth, &mut body, &request_id, &client).await {
+    // Filled inside `dispatch`, so the failure branch — where a guardrail
+    // block lands — stamps the enforced hits too (AISIX-Cloud#1330 / #1024).
+    let mut screening = InputScreening::default();
+    match dispatch(
+        &state,
+        &snapshot,
+        &auth,
+        &mut body,
+        &request_id,
+        &client,
+        &mut screening,
+    )
+    .await
+    {
         Ok(success) => {
             let elapsed = started.elapsed();
             let status = success.response.status().as_u16();
@@ -136,6 +163,19 @@ pub async fn count_tokens(
                 status,
                 elapsed,
             );
+            emit_usage_event(
+                &state,
+                &snapshot,
+                &pk,
+                &request_id,
+                &success.model_id,
+                &model_name,
+                &api_key_id,
+                status,
+                elapsed,
+                &client,
+                &screening,
+            );
             success.response
         }
         Err(err) => {
@@ -165,28 +205,64 @@ pub async fn count_tokens(
                 status,
                 elapsed,
             );
+            // A failed count_tokens is a request the operator has to be
+            // able to find, and a guardrail refusal is the one that must
+            // carry the flag the "Guardrail blocks" view filters on.
+            crate::usage_attr::emit_error_usage_event(
+                &state,
+                &snapshot,
+                "count_tokens",
+                "anthropic",
+                &request_id,
+                &model_name,
+                &api_key_id,
+                status,
+                err.kind(),
+                err.is_guardrail_block(),
+                &client,
+                crate::usage_attr::enforced_hits(&screening.audit),
+            );
             // Anthropic-shape envelope (#336) — count_tokens callers are
             // the Anthropic SDK, not OpenAI-compatible clients.
             err.into_anthropic_response()
         }
     }
 }
 
+/// What the input hook produced, for the terminal `UsageEvent` to carry.
+///
+/// An out-param rather than part of [`CountTokensSuccess`] because the
+/// failure branch needs it too — a guardrail refusal IS the error, so the
+/// error event is the one that must not drop the audit.
+#[derive(Default)]
+struct InputScreening {
+    /// The request's ENFORCE-mode audit handle (AISIX-Cloud#1330).
+    audit: crate::usage_attr::GuardrailAudit,
+    /// The `{kind, hook}` set of guardrails that governed the request
+    /// (#379 parity) — surfaced on the event so Logs can show them.
+    applied: Vec<aisix_core::AppliedGuardrail>,
+    /// Monitor-mode observations (AISIX-Cloud#562).
+    monitor_hits: Vec<aisix_core::GuardrailMonitorHit>,
+    /// Per-detector PII mask counts (#932/#696). Empty = no redaction.
+    redactions: crate::redact::RedactionCounts,
+}
+
 /// Run the resolved input guardrail chain over the Anthropic-shaped body,
 /// blocking before dispatch and writing mask-action rewrites back into
 /// `body` (which is what `count_tokens_to_target` forwards upstream).
 ///
 /// Deliberately mirrors `messages::dispatch_inner`'s block rather than
-/// sharing a helper with it: that one also threads applied-guardrail,
-/// audit and monitor-hit telemetry into a UsageEvent, and this route emits
-/// none (see [`CountTokensSuccess`]). Keeping the shapes parallel is what
-/// the `guardrail_coverage` census asserts.
+/// sharing a helper with it: that one threads the same telemetry through a
+/// retrying per-attempt emitter, and this route has a single terminal
+/// event. Keeping the shapes parallel is what the `guardrail_coverage`
+/// census asserts.
 async fn screen_input(
     state: &ProxyState,
     auth: &AuthenticatedKey,
     model_entry_id: &str,
     model_name: &str,
     body: &mut Value,
+    screening: &mut InputScreening,
 ) -> Result<(), ProxyError> {
     let chain = state
         .guardrail_index
@@ -197,6 +273,8 @@ async fn screen_input(
             api_key_id: &auth.entry.id,
             team_id: auth.key().team_id.as_deref(),
         });
+    screening.applied = chain.applied().to_vec();
+    screening.audit = chain.audit_log();
     if chain.is_empty() {
         return Ok(());
     }
@@ -218,18 +296,14 @@ async fn screen_input(
             ));
         }
     };
-    // Monitor-mode hits and redaction counts are collected and dropped:
-    // this route emits no UsageEvent (see [`CountTokensSuccess`]), so there
-    // is nothing to attach them to. Both are out-params of the shared
-    // helpers rather than optional, hence the sinks.
-    let (verdict, _monitor_hits) =
+    let (verdict, monitor_hits) =
         aisix_guardrails::Guardrail::check_input_non_segment_observed(&chain, &chat).await;
-    let mut counts = crate::redact::RedactionCounts::new();
+    screening.monitor_hits = monitor_hits;
     let verdict = crate::redact::moderate_body(
         &chain,
         crate::redact::Direction::Input,
         verdict,
-        &mut counts,
+        &mut screening.redactions,
         &mut Vec::new(),
         |g| crate::redact::redact_anthropic_request(g, body),
     )
@@ -257,24 +331,29 @@ async fn screen_input(
     Ok(())
 }
 
-/// What the winning attempt resolved. `/v1/messages/count_tokens` emits no
-/// UsageEvent, so the only consumer is the request-metric label set — which
-/// still has to match what chat / messages / responses report
-/// (AISIX-Cloud#1234).
+/// What the winning attempt resolved, for the request-metric label set —
+/// which has to match what chat / messages / responses report
+/// (AISIX-Cloud#1234) — and for the terminal `UsageEvent`.
 struct CountTokensSuccess {
     response: Response,
     provider: String,
     upstream_model: String,
     provider_key_id: String,
+    /// The DISPATCHED target's Model row id: a group resolves to one of
+    /// its members, and `UsageEvent::model_id` records that target while
+    /// `requested_model` keeps the alias the caller addressed.
+    model_id: String,
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn dispatch(
     state: &ProxyState,
     snapshot: &aisix_core::AisixSnapshot,
     auth: &AuthenticatedKey,
     body: &mut Value,
     request_id: &str,
     client: &ClientContext,
+    screening: &mut InputScreening,
 ) -> Result<CountTokensSuccess, ProxyError> {
     let model_name = body
         .get("model")
@@ -296,7 +375,7 @@ async fn dispatch(
     // rationale as the `/v1/messages` sibling: before the reservation, so a
     // content-policy refusal doesn't burn an RPM slot. See the module doc
     // for why the input hook applies here and the output hook does not.
-    screen_input(state, auth, &model_entry.id, &model_name, body).await?;
+    screen_input(state, auth, &model_entry.id, &model_name, body, screening).await?;
 
     let model_rl =
         crate::quota::ModelRateLimit::from_model(&model_name, &model_entry.id, &model_entry.value);
@@ -640,9 +719,76 @@ async fn count_tokens_to_target(
         provider: "anthropic".to_string(),
         upstream_model,
         provider_key_id: pk_entry.id.to_string(),
+        model_id: model_id.to_string(),
     })
 }
 
+/// The terminal `UsageEvent` for a served count_tokens.
+///
+/// Token counters stay at zero, deliberately: the `{"input_tokens": N}` the
+/// caller gets back is a MEASUREMENT of a prompt, not tokens any upstream
+/// consumed or billed. Copying it into `prompt_tokens` would put spend on
+/// a request that cost nothing and double-count the prompt once the caller
+/// goes on to issue the real `/v1/messages` call.
+///
+/// No `request_metrics::record_usage` call for the same reason — the
+/// `aisix_llm_*_tokens_total` families are token/spend families, and this
+/// route contributes neither. The request families already carry the call
+/// (`request_metrics::record`, above), and `aisix_usage_events_emitted_total`
+/// counts this event under `handler="count_tokens"`.
+#[allow(clippy::too_many_arguments)]
+fn emit_usage_event(
+    state: &ProxyState,
+    snap: &aisix_core::AisixSnapshot,
+    pk: &crate::usage_attr::ResolvedPk<'_>,
+    request_id: &str,
+    model_id: &str,
+    requested_model: &str,
+    api_key_id: &str,
+    status_code: u16,
+    elapsed: Duration,
+    client: &ClientContext,
+    screening: &InputScreening,
+) {
+    let mut event = aisix_obs::UsageEvent {
+        request_id: request_id.to_string(),
+        occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
+        model_id: model_id.to_string(),
+        api_key_id: api_key_id.to_string(),
+        requested_model: requested_model.to_string(),
+        // Single-attempt route: the attempt spans the whole request, so the
+        // upstream figure and what the caller waited for coincide.
+        upstream_latency_ms: elapsed.as_millis().min(u32::MAX as u128) as u32,
+        downstream_latency_ms: elapsed.as_millis().min(u32::MAX as u128) as u32,
+        status_code,
+        inbound_protocol: "anthropic".to_string(),
+        applied_guardrails: screening.applied.clone(),
+        client_source_ip: client.source_ip.clone(),
+        client_user_agent: client.user_agent.clone(),
+        redacted_entity_counts: screening.redactions.clone(),
+        guardrail_monitor_hits: screening.monitor_hits.clone(),
+        guardrail_enforced_hits: crate::usage_attr::enforced_hits(&screening.audit),
+        ..Default::default()
+    };
+    crate::usage_attr::apply_pk_telemetry(&mut event, pk);
+    crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref());
+    let usage_model =
+        crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned();
+    crate::usage_attr::emit_usage(
+        state,
+        snap,
+        "count_tokens",
+        event,
+        crate::usage_attr::usage_event_labels(&usage_model, pk),
+        // Nothing to capture: the response body is the token count, and the
+        // request body was already screened by the input hook.
+        None,
+        client.trace.as_ref(),
+        /* terminal */ true,
+        /* dispatched */ true,
+    );
+}
+
 fn emit_access_log(
     model: &str,
     provider: &str,

crates/aisix-proxy/src/guardrail_blocked_telemetry.rs

@@ -25,6 +25,17 @@
 //! point at a dead upstream, so a clean request fails too, and a flag that
 //! merely tracked "the request failed" would pass the blocked run and fail
 //! the clean one.
+//!
+//! The list below is hand-written, and AISIX-Cloud#1435 is what a
+//! hand-written list costs: `/v1/messages/count_tokens` gained the chain
+//! and the flag in the same release, was not on it, and emitted no usage
+//! event at all. So the "no surface may be missing" half now lives in
+//! `guardrail_coverage`, whose set is parsed out of the router — add a
+//! route and it is checked whether or not anyone edits a list. What stays
+//! here is what that census cannot express: the CLEAN control above, which
+//! needs a text-dependent guardrail rather than the census's unconditional
+//! script, and `/passthrough/byo`, which needs a configured route prefix
+//! the census snapshot does not carry.
 
 use std::sync::Arc;