From 95c269a35522da66974ccf6c33b9cfbfafcf2bec Mon Sep 17 00:00:00 2001 From: Jarvis Date: Fri, 28 Aug 2026 10:35:15 +0000 Subject: [PATCH 1/5] fix(count_tokens): emit the usage event a refusal needs to be findable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `/v1/messages/count_tokens` emitted no UsageEvent on any outcome. #1064 made it refusable by an input guardrail and #1065 put the flag the Logs "Guardrail blocks" view filters on onto the event a refusal produces — on the nine families that emit one. This route had nothing to put it on, so it refused correctly and the refusal was unfindable. The route stays unmetered: token counters are zero on both paths. What it gains is a row. Also adds the reporting half to `guardrail_coverage`'s census, whose surface set is parsed out of the router — the census that already existed asserted refusals against a hand-written list, which is how this surface was missing from it. --- crates/aisix-proxy/src/count_tokens.rs | 180 ++++++++++++++-- .../src/guardrail_blocked_telemetry.rs | 11 + crates/aisix-proxy/src/guardrail_coverage.rs | 102 +++++++++ .../count-tokens-usage-event-e2e.test.ts | 197 ++++++++++++++++++ 4 files changed, 473 insertions(+), 17 deletions(-) create mode 100644 tests/e2e/src/cases/count-tokens-usage-event-e2e.test.ts diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index 772b8ce1..91e9b36c 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/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,6 +205,23 @@ 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() @@ -172,21 +229,40 @@ pub async fn count_tokens( } } +/// 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, + /// Monitor-mode observations (AISIX-Cloud#562). + monitor_hits: Vec, + /// 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,17 +331,21 @@ 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, @@ -275,6 +353,7 @@ async fn dispatch( body: &mut Value, request_id: &str, client: &ClientContext, + screening: &mut InputScreening, ) -> Result { 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, diff --git a/crates/aisix-proxy/src/guardrail_blocked_telemetry.rs b/crates/aisix-proxy/src/guardrail_blocked_telemetry.rs index 460b24bc..264c536c 100644 --- a/crates/aisix-proxy/src/guardrail_blocked_telemetry.rs +++ b/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; diff --git a/crates/aisix-proxy/src/guardrail_coverage.rs b/crates/aisix-proxy/src/guardrail_coverage.rs index 7fab0499..c2599361 100644 --- a/crates/aisix-proxy/src/guardrail_coverage.rs +++ b/crates/aisix-proxy/src/guardrail_coverage.rs @@ -28,6 +28,7 @@ use std::sync::Arc; use aisix_core::snapshot::SnapshotHandle; use aisix_core::{AisixSnapshot, ApiKey, ProxyConfig, ResourceEntry}; +use aisix_obs::{UsageEvent, UsageSink}; use axum::body::Body; use axum::http::Request; use tower::ServiceExt; @@ -412,6 +413,23 @@ fn census_router() -> axum::Router { crate::build_router(state) } +/// [`census_router`] plus the receiver its usage events land in. +/// +/// One router per surface rather than a shared one: the sink is a single +/// channel, so surfaces driven through the same router would interleave +/// their events and "which surface emitted nothing" would stop being +/// answerable — which is the whole question below. +fn census_router_with_usage() -> (axum::Router, tokio::sync::mpsc::Receiver) { + let handle = SnapshotHandle::new(census_snapshot()); + let index = aisix_guardrails::LiveGuardrailIndex::new(handle.clone(), None); + let (tx, rx) = tokio::sync::mpsc::channel(32); + let state = crate::ProxyState::new(handle, census_hub(), &cfg()) + .without_cache() + .with_guardrail_index(index) + .with_usage_sink(UsageSink::new(tx)); + (crate::build_router(state), rx) +} + const MULTIPART_BOUNDARY: &str = "censusboundary"; /// A multipart body with the parts named, in order. Values are inline @@ -633,6 +651,90 @@ async fn enforced_surfaces_refuse_a_blocking_guardrail() { ); } +/// AISIX-Cloud#1435: refusing is half the job — the refusal also has to be +/// REPORTED, and on the same router-derived set. +/// +/// `guardrail_blocked_telemetry` already pins this invariant, but against a +/// hand-written list of surfaces, and that is exactly how the gap it was +/// written for came back: `/v1/messages/count_tokens` gained the chain +/// (#1064) and the flag (#1065) in the same release, was absent from the +/// list, and emitted no usage event at all — so it refused correctly and +/// the refusal was unfindable, on a route whose whole job is to ship the +/// caller's entire payload to a provider. Here the set comes out of the +/// router, so a surface cannot be missing from it. +/// +/// The counters are asserted alongside because a refusal that BILLS is the +/// other way to get this wrong: nothing ran upstream, so nothing is owed. +#[tokio::test] +async fn an_enforced_surface_reports_the_refusal_it_makes() { + let mut wrong = Vec::new(); + + for (surface, posture) in POSTURE { + if !matches!(posture, Posture::Enforced) { + continue; + } + let Some(request) = fixture(surface) else { + // `enforced_surfaces_refuse_a_blocking_guardrail` owns the + // missing-fixture complaint; do not duplicate it here. + continue; + }; + let (router, mut rx) = census_router_with_usage(); + let response = router.oneshot(request).await.expect("router must answer"); + let bytes = axum::body::to_bytes(response.into_body(), 1 << 20) + .await + .expect("body must read"); + if !refused_by_guardrail(&String::from_utf8_lossy(&bytes)) { + // Ditto: a surface that did not refuse is the sibling test's + // finding, and reporting it twice buries the new one. + continue; + } + + // Drained on a short timeout rather than counted — how many events + // a surface emits is its own business, and pinning it here would + // make this file fail for reasons that are not the flag. + let mut events = Vec::new(); + while let Ok(Some(event)) = + tokio::time::timeout(std::time::Duration::from_millis(300), rx.recv()).await + { + events.push(event); + } + + if events.is_empty() { + wrong.push(format!("{surface}: refused but emitted no usage event")); + continue; + } + if !events.iter().any(|e| e.guardrail_blocked) { + wrong.push(format!( + "{surface}: emitted {} usage event(s), none marked guardrail_blocked", + events.len(), + )); + continue; + } + // `/a2a` is exempt: its counters are the gateway's own reading of + // the words, flagged `usage_estimated` and never charged — they are + // filled from the request before the chain even runs. + if *surface == "/a2a/:agent" { + continue; + } + for event in &events { + if event.prompt_tokens != 0 || event.completion_tokens != 0 { + wrong.push(format!( + "{surface}: refused request billed {}+{} tokens", + event.prompt_tokens, event.completion_tokens, + )); + } + } + } + + assert!( + wrong.is_empty(), + "a guardrail refusal must reach the Logs \"Guardrail blocks\" view \ + (usage_events.guardrail_blocked = true) on every enforced surface, and cost \ + the caller nothing:\n {}", + wrong.join("\n "), + ); +} + /// Drive every enforced surface with `guardrails` holding `row` — or /// nothing at all when `row` is `None` — and return the surfaces that /// answered with a guardrail refusal. diff --git a/tests/e2e/src/cases/count-tokens-usage-event-e2e.test.ts b/tests/e2e/src/cases/count-tokens-usage-event-e2e.test.ts new file mode 100644 index 00000000..5408f78b --- /dev/null +++ b/tests/e2e/src/cases/count-tokens-usage-event-e2e.test.ts @@ -0,0 +1,197 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + slsLogsFor, + spawnApp, + startMockSls, + startOpenAiUpstream, + waitConfigPropagation, + waitForSlsLog, + type MockSls, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for AISIX-Cloud#1435: `/v1/messages/count_tokens` records the +// requests it serves and the ones it refuses. +// +// The route shipped deliberately unmetered — it generates nothing, so +// there is nothing to bill — and that was read as "emits nothing", which +// is a different statement. It forwards the caller's entire `system` + +// `messages` + `tools` payload to a real provider, and #1064 made it +// refusable by an input guardrail. So a refusal was a 422 the caller +// definitely saw and the Logs "Guardrail blocks" view could not find: +// nine handler families gained the flag in #1065, this one had no event +// to put it on. +// +// Both halves are asserted here because only the pair pins the design. +// The refusal must be findable; the SUCCESS must be findable too and must +// still bill nothing — the `{"input_tokens": N}` a caller gets back is a +// measurement of a prompt, not tokens an upstream consumed, and copying it +// into `prompt_tokens` would charge for a free call and double-count the +// prompt once the caller issues the real `/v1/messages`. +// +// Read back off a real Aliyun-SLS export from a real `aisix` binary, so +// what is asserted is the row a consumer receives. + +const CALLER_PLAINTEXT = "sk-ct-usage-e2e-caller"; +const CALLER_KEY_HASH = createHash("sha256").update(CALLER_PLAINTEXT).digest("hex"); + +const CREDENTIAL_REF = "mock"; +const MOCK_AK_ID = "LTAI_mock_ak"; +const MOCK_AK_SECRET = "mock_ak_secret"; +const SLS_PROJECT = "aisix-e2e-obs"; +const LOGSTORE = "count-tokens-usage"; + +const FORBIDDEN_WORD = "counttokensentinel"; +const MODEL_ALIAS = "ctu-e2e"; +const UPSTREAM_MODEL_ID = "claude-haiku-4-5-20251001"; + +// The route emits no captured content, so rows are identified by the +// outcome rather than by a marker planted in the prompt. One probe of each +// kind therefore has to mean one row of each status. +const rowsWithStatus = (sls: MockSls, status: string) => + slsLogsFor(sls, LOGSTORE).filter((l) => l.get("status_code") === status); + +describe("count_tokens usage e2e: served and refused requests both leave a row (#1435)", () => { + let upstream: OpenAiUpstream | undefined; + let sls: MockSls | undefined; + let app: SpawnedApp | undefined; + let etcdReachable = false; + + async function countTokens(text: string): Promise { + const res = await fetch(`${app!.proxyUrl}/v1/messages/count_tokens`, { + method: "POST", + headers: { + "x-api-key": CALLER_PLAINTEXT, + "anthropic-version": "2023-06-01", + "content-type": "application/json", + }, + body: JSON.stringify({ + model: MODEL_ALIAS, + messages: [{ role: "user", content: text }], + }), + }); + await res.text(); + return res; + } + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + sls = await startMockSls(); + // Anthropic's documented count_tokens response shape. The mock is + // path-agnostic, so it stands in for the upstream's own sub-route. + upstream = await startOpenAiUpstream({ nonStreamBody: { input_tokens: 42 } }); + + app = await spawnApp({ + extraEnv: { + [`SLS_CRED_${CREDENTIAL_REF.toUpperCase()}_AK_ID`]: MOCK_AK_ID, + [`SLS_CRED_${CREDENTIAL_REF.toUpperCase()}_AK_SECRET`]: MOCK_AK_SECRET, + }, + }); + const seed = new SeedClient(etcd, app.etcdPrefix); + + await seed.createObservabilityExporter({ + name: "ctu-sls", + enabled: true, + kind: "aliyun_sls", + endpoint: sls.url, + project: SLS_PROJECT, + logstore: LOGSTORE, + credential_ref: CREDENTIAL_REF, + }); + + // Anthropic bridge appends the path to the bare host (no `/v1`), and + // count_tokens only dispatches to Anthropic-protocol targets. + const pk = await seed.createProviderKey({ + display_name: "ctu-pk", + secret: "sk-ant-mock", + api_base: upstream.baseUrl, + }); + await seed.createModel({ + display_name: MODEL_ALIAS, + provider: "anthropic", + model_name: UPSTREAM_MODEL_ID, + provider_key_id: pk.id, + }); + await seed.createGuardrail({ + name: "ctu-guard", + enabled: true, + hook_point: "input", + kind: "keyword", + patterns: [{ kind: "literal", value: FORBIDDEN_WORD }], + }); + // Written last: one etcd watch applies events in revision order, so + // the moment this key authenticates everything above is in the + // snapshot. + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: [MODEL_ALIAS], + }); + + await waitConfigPropagation(async () => (await countTokens("readiness")).status === 200); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + await sls?.close(); + }); + + test("a served count_tokens leaves a row, and bills nothing", async (ctx) => { + if (!etcdReachable || !app || !sls) { + ctx.skip(); + return; + } + const before = slsLogsFor(sls, LOGSTORE).length; + const res = await countTokens("how long is this prompt"); + expect(res.status).toBe(200); + + const log = await waitForSlsLog( + sls, + LOGSTORE, + (l) => l.get("status_code") === "200", + "served count_tokens usage row", + 15_000, + ); + expect(log.get("requested_model")).toBe(MODEL_ALIAS); + expect(log.get("inbound_protocol")).toBe("anthropic"); + expect(log.get("guardrail_blocked")).not.toBe("true"); + // The upstream answered `input_tokens: 42`; that is a measurement of + // the prompt, not consumption, so it must not become spend. + expect(log.get("prompt_tokens") ?? "0").toBe("0"); + expect(log.get("completion_tokens") ?? "0").toBe("0"); + expect(slsLogsFor(sls, LOGSTORE).length).toBeGreaterThan(before); + }); + + test("a refused count_tokens reaches the Blocked view", async (ctx) => { + if (!etcdReachable || !app || !sls) { + ctx.skip(); + return; + } + const res = await countTokens(`please ${FORBIDDEN_WORD} now`); + expect(res.status).toBe(422); + + const log = await waitForSlsLog( + sls, + LOGSTORE, + (l) => l.get("status_code") === "422", + "refused count_tokens usage row", + 15_000, + ); + // The predicate the dashboard's "Guardrail blocks" view filters on. + expect(log.get("guardrail_blocked")).toBe("true"); + expect(log.get("requested_model")).toBe(MODEL_ALIAS); + expect(log.get("inbound_protocol")).toBe("anthropic"); + // Refused before dispatch, so nothing was sent and nothing is owed. + expect(log.get("prompt_tokens") ?? "0").toBe("0"); + expect(log.get("completion_tokens") ?? "0").toBe("0"); + // One refusal, one row — not one per resolved target. + expect(rowsWithStatus(sls, "422")).toHaveLength(1); + }); +}); From f2f793d2f6c0653ff7f8b34b9a083f249178d703 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Fri, 28 Aug 2026 10:43:27 +0000 Subject: [PATCH 2/5] fix(count_tokens): follow apply_jwt_identity to apply_caller_identity #1066 renamed the family so a new emit site cannot be added without considering member attribution. This is that site. --- crates/aisix-proxy/src/count_tokens.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index 91e9b36c..7b92d5cf 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -771,7 +771,11 @@ fn emit_usage_event( ..Default::default() }; crate::usage_attr::apply_pk_telemetry(&mut event, pk); - crate::usage_attr::apply_jwt_identity(&mut event, client.jwt.as_ref()); + crate::usage_attr::apply_caller_identity( + &mut event, + client.jwt.as_ref(), + client.caller.user_id.as_deref(), + ); let usage_model = crate::usage_attr::usage_event_model_label(snap, &event.requested_model).into_owned(); crate::usage_attr::emit_usage( From 65af6bfe173af40b5ab6e632879ce43b774be682 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Fri, 28 Aug 2026 10:46:44 +0000 Subject: [PATCH 3/5] fix(count_tokens): attempt-scope the upstream latency; drop a vacuous census exemption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The emitter reported the handler's own elapsed as upstream_latency_ms. This route fails over across a group's Anthropic targets and retries within one, so on a group that figure includes every attempt that lost — while upstream_latency_ms is attempt-scoped everywhere else in Logs. The census skipped /a2a for the bills-nothing assertion, copying an exemption guardrail_blocked_telemetry needs because its fixtures carry real text. The census's are contentless, so a2a bills zero there anyway: the skip protected nothing and would have hidden a surface that started billing on a refusal. --- crates/aisix-proxy/src/count_tokens.rs | 22 +++++++++++++++----- crates/aisix-proxy/src/guardrail_coverage.rs | 13 ++++++------ 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index 7b92d5cf..ffd39613 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -172,6 +172,7 @@ pub async fn count_tokens( &model_name, &api_key_id, status, + success.upstream_elapsed, elapsed, &client, &screening, @@ -343,6 +344,12 @@ struct CountTokensSuccess { /// its members, and `UsageEvent::model_id` records that target while /// `requested_model` keeps the alias the caller addressed. model_id: String, + /// How long the WINNING attempt took. Not the handler's own elapsed: + /// this route fails over across a group's Anthropic targets and + /// retries within one, so on a group the two diverge by every attempt + /// that lost — and `upstream_latency_ms` is attempt-scoped everywhere + /// else in Logs (`downstream_latency_ms` is the request-scoped one). + upstream_elapsed: Duration, } #[allow(clippy::too_many_arguments)] @@ -552,6 +559,7 @@ async fn count_tokens_to_target( request_id: &str, client: &ClientContext, ) -> Result { + let attempt_started = Instant::now(); let mut body = body.clone(); let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; let api_key = crate::dispatch::require_api_key(&pk_entry.value, model)?; @@ -720,6 +728,7 @@ async fn count_tokens_to_target( upstream_model, provider_key_id: pk_entry.id.to_string(), model_id: model_id.to_string(), + upstream_elapsed: attempt_started.elapsed(), }) } @@ -746,6 +755,10 @@ fn emit_usage_event( requested_model: &str, api_key_id: &str, status_code: u16, + // Attempt-scoped, from the winning attempt; see `CountTokensSuccess`. + upstream_elapsed: Duration, + // Request-scoped: what the caller actually waited for, guardrails and + // any lost attempts included. elapsed: Duration, client: &ClientContext, screening: &InputScreening, @@ -756,9 +769,7 @@ fn emit_usage_event( 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, + upstream_latency_ms: upstream_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(), @@ -784,8 +795,9 @@ fn emit_usage_event( "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. + // Content capture (#700) is not wired on this route — it is a + // separate, per-exporter opt-in capability, and #1435 is about the + // event existing at all. None, client.trace.as_ref(), /* terminal */ true, diff --git a/crates/aisix-proxy/src/guardrail_coverage.rs b/crates/aisix-proxy/src/guardrail_coverage.rs index c2599361..e7aa4ccb 100644 --- a/crates/aisix-proxy/src/guardrail_coverage.rs +++ b/crates/aisix-proxy/src/guardrail_coverage.rs @@ -710,12 +710,13 @@ async fn an_enforced_surface_reports_the_refusal_it_makes() { )); continue; } - // `/a2a` is exempt: its counters are the gateway's own reading of - // the words, flagged `usage_estimated` and never charged — they are - // filled from the request before the chain even runs. - if *surface == "/a2a/:agent" { - continue; - } + // No surface is exempt here, `/a2a` included — + // `guardrail_blocked_telemetry` has to exempt it because its + // counters are the gateway's own reading of the request text, + // filled before the chain runs, and that file's fixtures carry + // real text. These are contentless by construction, so zero is + // the honest answer on every one of them and an exemption would + // protect nothing while hiding a surface that started billing. for event in &events { if event.prompt_tokens != 0 || event.completion_tokens != 0 { wrong.push(format!( From 58cdb03512f561e607611fdb9ec387265f308c1d Mon Sep 17 00:00:00 2001 From: Jarvis Date: Fri, 28 Aug 2026 10:50:39 +0000 Subject: [PATCH 4/5] fix(count_tokens): report the masks it applied, like /v1/messages does MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The post-block-check masking pass was run and its counts discarded. Now that this route emits an event, redacted_entity_counts must match what /v1/messages reports for the same body and chain — otherwise the sibling route reads as masking more of the same payload. Invisible from the audit side: the enforced hit carries its own copy of the count, so the field cp-api persists was empty while a reader checking guardrail_enforced_hits would have seen the mask recorded. --- crates/aisix-proxy/src/count_tokens.rs | 92 +++++++++++++++++++++++++- 1 file changed, 91 insertions(+), 1 deletion(-) diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index ffd39613..7e5a6a38 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -328,7 +328,14 @@ async fn screen_input( )); } // Mask-action rules rewrite the body that is about to be forwarded. - crate::redact::redact_anthropic_request(&chain, body); + // Merged, not discarded: `/v1/messages` merges the same pass into the + // counts its event reports (#932), and the two routes screen the same + // body with the same chain — a mask this side under-reported would + // read as the sibling route masking more of the same payload. + crate::redact::merge_counts( + &mut screening.redactions, + crate::redact::redact_anthropic_request(&chain, body), + ); Ok(()) } @@ -931,6 +938,89 @@ mod tests { .unwrap() } + /// AISIX-Cloud#1435: the served request leaves a row, and that row + /// carries what the guardrail chain did to the body. + /// + /// The mask count is the part worth pinning. `/v1/messages` merges the + /// post-block-check masking pass into the counts its event reports + /// (#932), and this route screens the same body with the same chain — + /// so a mask counted on one and not the other reads as the sibling + /// route masking more of the same payload. It is also invisible from + /// the audit side: the enforced hit carries its own copy, so a reader + /// checking only that would see the mask recorded while the field + /// cp-api persists stayed empty. The upstream answers + /// `input_tokens: 42`, which must NOT become spend: it measures a + /// prompt, it does not consume one. + #[tokio::test] + async fn a_served_request_emits_a_zero_token_row_carrying_the_mask_it_applied() { + use aisix_obs::UsageSink; + + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages/count_tokens")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"input_tokens": 42})), + ) + .mount(&upstream) + .await; + + let snap = new_snap(&upstream.uri()); + snap.models.insert(anthropic_model("ct-mask")); + snap.apikeys.insert(apikey_entry(&["ct-mask"])); + let row: aisix_core::models::Guardrail = serde_json::from_str( + r#"{ + "name": "eda-mask", + "kind": "pii", + "hook_point": "input", + "detectors": [], + "custom_patterns": [ + {"name": "eda_version", "regex": "version\\s*:\\s*(\\d+(?:\\.\\d+)+)", "action": "mask", "replacement": "***"} + ] + }"#, + ) + .unwrap(); + snap.guardrails.insert(ResourceEntry::new("g-mask", row, 1)); + + let hub = Arc::new(Hub::new()); + hub.register_specialized( + "anthropic", + Arc::new(aisix_provider_anthropic::AnthropicBridge::new()), + ); + let handle = SnapshotHandle::new(snap); + let index = aisix_guardrails::LiveGuardrailIndex::new(handle.clone(), None); + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let app = crate::build_router( + crate::ProxyState::new(handle, hub, &cfg()) + .without_cache() + .with_guardrail_index(index) + .with_usage_sink(UsageSink::new(tx)), + ); + + let res = app + .oneshot(make_req(serde_json::json!({ + "model": "ct-mask", + "messages": [{ "role": "user", "content": "version: 9.9.9" }], + }))) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + + let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("count_tokens must emit a usage event") + .expect("channel open"); + assert_eq!(event.status_code, 200); + assert_eq!(event.inbound_protocol, "anthropic"); + assert_eq!(event.requested_model, "ct-mask"); + assert_eq!(event.prompt_tokens, 0); + assert_eq!(event.completion_tokens, 0); + assert_eq!( + event.redacted_entity_counts.get("eda_version").copied(), + Some(1), + "the mask this route applied is missing from its own row: {event:?}", + ); + } + /// Mixed group [anthropic, openai]: the openai target is `continue`d /// past (count_tokens has no upstream there), so it is NOT a usable /// fallback — the default retry budget must apply on the anthropic From eebc8fda833ca14c2e6c03e1c220e11f1694c2e3 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Fri, 28 Aug 2026 10:53:32 +0000 Subject: [PATCH 5/5] test(count_tokens): gate readiness off the route under test, anchor rows by request id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both from the CodeRabbit review on #1068, both correct. The gate sent successful count_tokens requests, so a regression in the route would have surfaced as a propagation timeout in beforeAll instead of a failed assertion — and every gate attempt planted a served row the first test could match instead of its own. tests/e2e/AGENTS.md says a gate must not exercise the behavior under test; the sibling spec this one was modelled on already gates on GET /v1/models and explains why. Rows are now selected by x-aisix-request-id rather than status_code. ensure_request_id is the outermost layer and stamps it on every response, short-circuited 4xx included, and it is the telemetry request_id. --- .../count-tokens-usage-event-e2e.test.ts | 45 +++++++++++++------ 1 file changed, 32 insertions(+), 13 deletions(-) diff --git a/tests/e2e/src/cases/count-tokens-usage-event-e2e.test.ts b/tests/e2e/src/cases/count-tokens-usage-event-e2e.test.ts index 5408f78b..26b5068e 100644 --- a/tests/e2e/src/cases/count-tokens-usage-event-e2e.test.ts +++ b/tests/e2e/src/cases/count-tokens-usage-event-e2e.test.ts @@ -2,6 +2,7 @@ import { createHash } from "node:crypto"; import { afterAll, beforeAll, describe, expect, test } from "vitest"; import { EtcdClient, + ProxyClient, SeedClient, slsLogsFor, spawnApp, @@ -49,11 +50,14 @@ const FORBIDDEN_WORD = "counttokensentinel"; const MODEL_ALIAS = "ctu-e2e"; const UPSTREAM_MODEL_ID = "claude-haiku-4-5-20251001"; -// The route emits no captured content, so rows are identified by the -// outcome rather than by a marker planted in the prompt. One probe of each -// kind therefore has to mean one row of each status. -const rowsWithStatus = (sls: MockSls, status: string) => - slsLogsFor(sls, LOGSTORE).filter((l) => l.get("status_code") === status); +// The route emits no captured content, so a row cannot be found by a marker +// planted in the prompt the way the sibling specs do. `x-aisix-request-id` +// is the anchor instead: `ensure_request_id` is the outermost layer and +// stamps it onto every response, short-circuited 4xx included, and it IS +// the telemetry `request_id` — so each assertion pins the row its own +// request produced rather than any row that happens to share a status. +const rowsForRequest = (sls: MockSls, requestId: string) => + slsLogsFor(sls, LOGSTORE).filter((l) => l.get("request_id") === requestId); describe("count_tokens usage e2e: served and refused requests both leave a row (#1435)", () => { let upstream: OpenAiUpstream | undefined; @@ -61,7 +65,9 @@ describe("count_tokens usage e2e: served and refused requests both leave a row ( let app: SpawnedApp | undefined; let etcdReachable = false; - async function countTokens(text: string): Promise { + async function countTokens( + text: string, + ): Promise<{ status: number; requestId: string }> { const res = await fetch(`${app!.proxyUrl}/v1/messages/count_tokens`, { method: "POST", headers: { @@ -75,7 +81,10 @@ describe("count_tokens usage e2e: served and refused requests both leave a row ( }), }); await res.text(); - return res; + return { + status: res.status, + requestId: res.headers.get("x-aisix-request-id") ?? "", + }; } beforeAll(async () => { @@ -134,7 +143,13 @@ describe("count_tokens usage e2e: served and refused requests both leave a row ( allowed_models: [MODEL_ALIAS], }); - await waitConfigPropagation(async () => (await countTokens("readiness")).status === 200); + // Independent of the route under test (tests/e2e/AGENTS.md). Gating on + // count_tokens itself would make a regression in it surface as a + // propagation timeout in `beforeAll` instead of a failed assertion, and + // every gate attempt would plant a served row the first test could then + // match instead of its own. + const proxy = new ProxyClient(app.proxyUrl, CALLER_PLAINTEXT); + await waitConfigPropagation(async () => (await proxy.listModels()).status === 200); }); afterAll(async () => { @@ -148,17 +163,18 @@ describe("count_tokens usage e2e: served and refused requests both leave a row ( ctx.skip(); return; } - const before = slsLogsFor(sls, LOGSTORE).length; const res = await countTokens("how long is this prompt"); expect(res.status).toBe(200); + expect(res.requestId).not.toBe(""); const log = await waitForSlsLog( sls, LOGSTORE, - (l) => l.get("status_code") === "200", + (l) => l.get("request_id") === res.requestId, "served count_tokens usage row", 15_000, ); + expect(log.get("status_code")).toBe("200"); expect(log.get("requested_model")).toBe(MODEL_ALIAS); expect(log.get("inbound_protocol")).toBe("anthropic"); expect(log.get("guardrail_blocked")).not.toBe("true"); @@ -166,7 +182,8 @@ describe("count_tokens usage e2e: served and refused requests both leave a row ( // the prompt, not consumption, so it must not become spend. expect(log.get("prompt_tokens") ?? "0").toBe("0"); expect(log.get("completion_tokens") ?? "0").toBe("0"); - expect(slsLogsFor(sls, LOGSTORE).length).toBeGreaterThan(before); + // One request, one row. + expect(rowsForRequest(sls, res.requestId)).toHaveLength(1); }); test("a refused count_tokens reaches the Blocked view", async (ctx) => { @@ -176,22 +193,24 @@ describe("count_tokens usage e2e: served and refused requests both leave a row ( } const res = await countTokens(`please ${FORBIDDEN_WORD} now`); expect(res.status).toBe(422); + expect(res.requestId).not.toBe(""); const log = await waitForSlsLog( sls, LOGSTORE, - (l) => l.get("status_code") === "422", + (l) => l.get("request_id") === res.requestId, "refused count_tokens usage row", 15_000, ); // The predicate the dashboard's "Guardrail blocks" view filters on. expect(log.get("guardrail_blocked")).toBe("true"); + expect(log.get("status_code")).toBe("422"); expect(log.get("requested_model")).toBe(MODEL_ALIAS); expect(log.get("inbound_protocol")).toBe("anthropic"); // Refused before dispatch, so nothing was sent and nothing is owed. expect(log.get("prompt_tokens") ?? "0").toBe("0"); expect(log.get("completion_tokens") ?? "0").toBe("0"); // One refusal, one row — not one per resolved target. - expect(rowsWithStatus(sls, "422")).toHaveLength(1); + expect(rowsForRequest(sls, res.requestId)).toHaveLength(1); }); });