fix(ratelimit): count Anthropic cache tokens in /v1/messages + /v1/responses TPM commits

crates/aisix-proxy/src/messages.rs

@@ -63,6 +63,7 @@ use crate::client_ip::ClientContext;
 use crate::error::ProxyError;
 use crate::request_id::new_request_id;
 use crate::state::ProxyState;
+use crate::usage_attr::total_tokens_with_cache;
 
 /// Anthropic API version header value injected on every forwarded request.
 /// Shared with the `/v1/messages/count_tokens` handler so both Anthropic
@@ -664,8 +665,12 @@ async fn dispatch(
                     // this is skipped.
                     if !outcome.usage_handled_by_stream {
                         if let Some(r) = reservation.take() {
-                            let total = u64::from(outcome.metrics.prompt_tokens)
-                                + u64::from(outcome.metrics.completion_tokens);
+                            let total = total_tokens_with_cache(
+                                outcome.metrics.prompt_tokens,
+                                outcome.metrics.completion_tokens,
+                                outcome.metrics.cache_creation_tokens,
+                                outcome.metrics.cache_read_tokens,
+                            );
                             r.commit_tokens(total).await;
                         }
                     }
@@ -1102,8 +1107,12 @@ async fn anthropic_passthrough_dispatch(
                 // is the sync analog of the reservation's async `commit_tokens`
                 // (this end-of-stream closure can't await); dropping the hold frees
                 // the concurrency slot(s) held for the stream's full lifetime.
-                let streamed_tokens =
-                    u64::from(usage.prompt_tokens) + u64::from(usage.completion_tokens);
+                let streamed_tokens = total_tokens_with_cache(
+                    usage.prompt_tokens,
+                    usage.completion_tokens,
+                    usage.cache_creation_tokens,
+                    usage.cache_read_tokens,
+                );
                 for key in &post_stream_keys {
                     limiter_c.add_tokens_post_stream(key, streamed_tokens);
                 }
@@ -1608,8 +1617,12 @@ async fn cross_provider_dispatch(
                 // concurrency hold now the stream has ended (sync analog of the
                 // reservation's async `commit_tokens`, which this closure can't
                 // await).
-                let streamed_tokens =
-                    u64::from(comp.prompt_tokens) + u64::from(comp.completion_tokens);
+                let streamed_tokens = total_tokens_with_cache(
+                    comp.prompt_tokens,
+                    comp.completion_tokens,
+                    comp.cache_creation_tokens,
+                    comp.cache_read_tokens,
+                );
                 for key in &post_stream_keys {
                     limiter_for_stream.add_tokens_post_stream(key, streamed_tokens);
                 }
@@ -2318,9 +2331,13 @@ fn emit_anthropic_usage_event(
         LlmUsage {
             input_tokens: metrics.prompt_tokens,
             output_tokens: metrics.completion_tokens,
-            total_tokens: metrics
-                .prompt_tokens
-                .saturating_add(metrics.completion_tokens),
+            total_tokens: total_tokens_with_cache(
+                metrics.prompt_tokens,
+                metrics.completion_tokens,
+                metrics.cache_creation_tokens,
+                metrics.cache_read_tokens,
+            )
+            .min(u64::from(u32::MAX)) as u32,
             spend_usd: 0.0,
         },
     );

crates/aisix-proxy/src/responses.rs

@@ -33,7 +33,7 @@ use crate::client_ip::ClientContext;
 use crate::error::ProxyError;
 use crate::request_id::new_request_id;
 use crate::state::ProxyState;
-use crate::usage_attr::provider_telemetry_tags;
+use crate::usage_attr::{provider_telemetry_tags, total_tokens_with_cache};
 
 /// Per-request payload from a successful dispatch — carries the
 /// response + provider label + the bits of usage data needed for
@@ -570,7 +570,12 @@ async fn dispatch(
                                 .usage
                                 .as_ref()
                                 .map(|u| {
-                                    u64::from(u.prompt_tokens) + u64::from(u.completion_tokens)
+                                    total_tokens_with_cache(
+                                        u.prompt_tokens,
+                                        u.completion_tokens,
+                                        u.cache_creation_tokens,
+                                        u.cache_read_tokens,
+                                    )
                                 })
                                 .unwrap_or(0);
                             r.commit_tokens(total).await;
@@ -1145,8 +1150,12 @@ async fn responses_to_target(
                 // #688: apply the terminal token cost to TPM/TPD and release the
                 // concurrency hold now the stream has ended (sync analog of the
                 // reservation's async `commit_tokens`, which this closure can't await).
-                let streamed_tokens =
-                    u64::from(usage.prompt_tokens) + u64::from(usage.completion_tokens);
+                let streamed_tokens = total_tokens_with_cache(
+                    usage.prompt_tokens,
+                    usage.completion_tokens,
+                    usage.cache_creation_tokens,
+                    usage.cache_read_tokens,
+                );
                 for key in &post_stream_keys {
                     limiter_c.add_tokens_post_stream(key, streamed_tokens);
                 }
@@ -1501,8 +1510,12 @@ async fn responses_cross_provider_to_target(
                 // concurrency hold now the stream has ended (sync analog of the
                 // reservation's async `commit_tokens`). Tokens count even on an
                 // output-guardrail block — the upstream still billed them.
-                let streamed_tokens =
-                    u64::from(comp.prompt_tokens) + u64::from(comp.completion_tokens);
+                let streamed_tokens = total_tokens_with_cache(
+                    comp.prompt_tokens,
+                    comp.completion_tokens,
+                    comp.cache_creation_tokens,
+                    comp.cache_read_tokens,
+                );
                 for key in &post_stream_keys {
                     limiter_c.add_tokens_post_stream(key, streamed_tokens);
                 }

crates/aisix-proxy/src/usage_attr.rs

@@ -55,6 +55,27 @@ pub(crate) fn provider_key_metric_name(snap: &AisixSnapshot, provider_key_id: &s
     }
 }
 
+/// Total token cost of a request as committed against TPM/TPD rate limits
+/// (and reported as the prometheus usage total): prompt + completion +
+/// Anthropic cache creation/read. Anthropic reports cache tokens as counters
+/// SEPARATE from `input_tokens`, so a prompt+completion sum silently
+/// undercounts cached traffic — the OpenAI bridge already folds them into
+/// `total_tokens` (#679) and the CP display total includes them (#906); this
+/// keeps the native `/v1/messages` and `/v1/responses` commits consistent
+/// (AISIX-Cloud#995). OpenAI's `cached_tokens` is a subset of
+/// `prompt_tokens` and is deliberately NOT an input here.
+pub(crate) fn total_tokens_with_cache(
+    prompt_tokens: u32,
+    completion_tokens: u32,
+    cache_creation_tokens: u32,
+    cache_read_tokens: u32,
+) -> u64 {
+    u64::from(prompt_tokens)
+        + u64::from(completion_tokens)
+        + u64::from(cache_creation_tokens)
+        + u64::from(cache_read_tokens)
+}
+
 /// The `model` metric label for a request whose client-supplied `model`
 /// field never resolved to a configured model (e.g. model-not-found). See
 /// [`metric_model_label`].