feat(messages): route /v1/messages/count_tokens to Anthropic upstream (#418)#419
Merged
Conversation
…#418) The Anthropic Messages count_tokens sub-route was unregistered, so the DP returned a bare 404 before any auth/bridge logic ran. The Anthropic SDK's `messages.countTokens(...)` — used by Claude Code and most Anthropic-SDK apps to size prompts before a paid /v1/messages call — threw `404 (no body)` through the gateway. Register `/v1/messages/count_tokens` as a passthrough to the Anthropic upstream's `…/v1/messages/count_tokens`, reusing the /v1/messages dispatch helpers (model-alias resolution, x-api-key + anthropic-version auth, build_v1_url) and forwarding the `{input_tokens}` body verbatim. Non-Anthropic models are rejected with a 400 (no upstream equivalent) rather than a misleading 404. Errors use the Anthropic-shape envelope (#336). Shared the JSON-rejection discrimination (malformed JSON vs 413 cap vs transport error) into error::proxy_error_from_json_rejection so the two Anthropic-protocol handlers stay in lockstep, and reused messages::ANTHROPIC_VERSION across both. Refs: Anthropic Count Message Tokens API; LiteLLM exposes the same route and had the identical missing-route bug (BerriAI/litellm#15006).
ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Free Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds a new Changescount_tokens endpoint implementation
🎯 3 (Moderate) | ⏱️ ~25 minutes Note 🎁 Summarized by CodeRabbit FreeYour organization has reached its limit of developer seats under the Pro Plan. For new users, CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please add seats to your subscription by visiting https://app.coderabbit.ai/login.If you believe this is a mistake and have available seats, please assign one to the pull request author through the subscription management page using the link above. Comment |
…unt_tokens Address review feedback on #418: - count_tokens truncated upstream error bodies with `&message[..1024]`, which panics when byte 1024 splits a multibyte codepoint (non-ASCII upstream error body). Truncate on a char boundary instead. Regression test exercises a 3-byte char straddling byte 1024. - Apply the ProviderKey `request.*` override pipeline (param_renames / param_constraints / default_body_fields / default_headers) on count_tokens, identically to /v1/messages — the two routes share the same Anthropic ProviderKey, so operator overrides (notably `anthropic-beta` via default_headers) must reach both. Headers are built via a HeaderMap so the #337 reserved-header blacklist protects x-api-key from operator override. Tests cover default_headers, param_renames, and the x-api-key protection.
Collaborator
Author
|
Follow-up filed for the pre-existing copies of the UTF-8 truncation pattern in |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
crates/aisix-proxy/src/lib.rs:147
inbound_protocol_for_endpointis fed fromrequest.uri().path()(seerecord_in_flight_request), but it only treats the exact string/v1/messagesas Anthropic. The codebase already accounts for/v1/messages/being routed to the Anthropic handler (see the trailing-slash audit note + test aroundenforce_request_body_limit), so in-flight metrics for/v1/messages/will be mislabeled asopenai. Consider including the trailing-slash variant here (and keeping this in sync with the Anthropic-path detection used elsewhere).
fn inbound_protocol_for_endpoint(endpoint: &str) -> &'static str {
if endpoint == "/v1/messages" || endpoint == "/v1/messages/count_tokens" {
"anthropic"
} else {
"openai"
}
Match the /v1/messages handler: insert the response request-id header only when it parses, instead of falling back to an empty value. The request_id is internally generated (always valid ASCII), so this is a consistency/defensiveness fix, not a live bug.
5 tasks
moonming
added a commit
that referenced
this pull request
May 29, 2026
PR #436 audit raised MEDIUM-2 + LOW-1 (both real) plus a HIGH-1 that was based on a factually incorrect premise (see PR comment). Acting on the valid findings: MEDIUM-2 — Unbounded frame buffer. The byte-level SSE parser buffers chunks into a Vec<u8> and only drains on a blank-line terminator. A malformed / hostile upstream that streams bytes WITHOUT a terminator would grow the buffer unboundedly (per-request memory exhaustion). chat.rs doesn't have this risk because it operates on pre-framed ChatChunks, not raw bytes — so this is a genuine new risk the byte-parser introduces. Added a 1 MiB ceiling (MAX_SSE_FRAME_BUF_BYTES); real Anthropic frames are a few KB, so the cap only trips on a non-conformant stream, where we drop the buffer (skipping usage parse for that frame) rather than OOM. The bytes still forward to the client verbatim. LOW-1 — A `message_delta` with a `usage` object but a non-numeric `output_tokens` silently left completion_tokens at the message_start floor (often 1). Added a `tracing::debug!` so a wire-shape drift is visible to operators. Test (audit angle 8c): `sse_frame_parser_tolerates_streams_without_usage` — an error-style stream (ping + error frames, no usage anywhere) drains cleanly leaving the accumulator at zeros, no panic. E2E (DP-level, honours the e2e-first discipline): new `anthropic-streaming-usage-e2e.test.ts` drives a real streaming /v1/messages through the DP binary against a mock Anthropic streaming upstream (input_tokens=37 in message_start, output_tokens=52 in message_delta), then scrapes the DP's /metrics and asserts `aisix_llm_input_tokens_total` / `aisix_llm_output_tokens_total` for /v1/messages are non-zero. Pre-#245 those counters stayed at 0, so this fails red on a regression. Verified locally against the aisix-e2e stack (passes in 2.9s). HIGH-1 (NOT actioned — incorrect premise): the audit claimed `DeliveryCounter` / #419 don't exist and that the PR diverges from chat.rs. Verified false: `grep -rn "DeliveryCounter" crates/` returns chat.rs:2312 (the struct), :2296 (construction), and the #419 cost-leak gate lives at chat.rs:2039-2045 with the `delivered: Arc<AtomicU32>` field at :2010. The PR's "mirrors chat.rs DeliveryCounter / #419 parity" framing is accurate; removing the delivery gate would BREAK parity, not restore it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
/v1/messages/count_tokensreturned a bare 404 through the DP because the route was never registered —count_tokensappeared nowhere in the codebase. The Anthropic SDK'smessages.countTokens(...)(used by Claude Code and most Anthropic-SDK apps to size a prompt before a paid/v1/messagescall) threw404 (no body). Fixes #418.This registers
/v1/messages/count_tokensas a passthrough to the Anthropic upstream, the sibling of the already-working/v1/messagesroute.What changed
count_tokens.rs(new) — handler mirroring thererank.rspassthrough structure. Reuses the existing dispatch helpers (resolve_provider_key,require_secret,require_upstream_model,resolve_base_url,build_v1_url), sends thex-api-key+anthropic-versionauth shape, rewrites the model alias to the upstream id, and forwards the{"input_tokens": <int>}body verbatim.lib.rs— registers the route; adds the path toinbound_protocol_for_endpoint(anthropic label) andis_anthropic_path(so body-limit 413s render the Anthropic envelope).error.rs— extractsproxy_error_from_json_rejection(malformed-JSON vs 413-cap vs transport-error discrimination) into a shared helper.messages.rs— rewired to that helper (removing a ~50-line inline duplicate);ANTHROPIC_VERSIONmadepub(crate)and reused.Scope decision
Anthropic-backed models only — the issue's scope and what the Anthropic SDK targets. A non-Anthropic model returns a clean 400 (
count_tokenshas no OpenAI/Gemini/DeepSeek upstream equivalent) rather than a misleading 404, and without fabricating a local estimate. Errors use the Anthropic-shape envelope (#336).Reference implementations / upstream spec (per CLAUDE.md §7)
POST /v1/messages/count_tokens→{"input_tokens": <int>}, headersx-api-key+anthropic-version: https://platform.claude.com/docs/en/api/messages-count-tokensTesting
count_tokens.rs): 401 (Anthropic envelope), 404 unknown model, 403 forbidden, 400 non-Anthropic provider, happy-path passthrough asserting upstream URL / model rewrite / auth headers.anthropic-count-tokens-e2e.test.ts) against the real backend + etcd:200withinput_tokens, correct upstream sub-route (/v1/messages/count_tokens, not/v1/messages), model rewrite,x-api-key/anthropic-versionforwarded (noAuthorization); plus unknown-model404with the Anthropic error envelope./v1/messages, tools-cross-provider, and body-edges e2e (the refactored paths) 6 passed;cargo fmt --check+clippyclean.The original repro no longer reproduces: the route now returns 200 instead of a bare 404.
Summary by CodeRabbit
New Features
Bug Fixes
Tests