fix(server): count expanded media tokens in input_tokens
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem
Found during the real-checkpoint run for PR #1777. The prompt-inspection routes document their count as "the number of tokens the same body would actually have been prefilled with", but `count_prompt_tokens` encodes the rendered template text only. Image and video soft-token expansion happens later on the model worker, which these routes never reach. On `gemma-3-4b-it-4bit` with two clips of four frames each, `/v1/chat/completions/input_tokens` answers 57 while the same body's `usage.prompt_tokens` is 2,129. A client that budgets context with these routes undercounts every image or video request by the size of its media.
## Evidence
- Contract: `src/server/routes/prompt_inspection.rs:28-35`, repeated in `docs/llama-server-compat.md:129`.
- `count_prompt_tokens` (`prompt_inspection.rs:159-176`) encodes `rendered.prompt`, produced by `render_chat_prompt` (`:95-138`), which stops at `prepare_chat_request_with_cache`.
- Both handlers in the file use it: `chat_input_tokens` (`:215-236`) and `responses_input_tokens` (`:244-275`), so `/v1/responses/input_tokens` has the same gap.
- `/v1/messages/count_tokens` (`src/server/routes/anthropic.rs:736-799`) repeats the same encode of `prepared.prompt` (`:780-791`), and `src/server/anthropic_translator.rs` turns `image` blocks into `image_url` parts (module doc, `:31`), so it undercounts image bodies too. Its missing media gate is #1748, a separate issue.
- The expansion: `prepare_request_vlm_embeddings` (`src/server/model_worker.rs:1440`) rewrites `prompt_tokens` in place, in the same per-family call that computes the vision embeddings (`src/multimodal/vlm_runtime.rs:1126`).
## Proposed fix
Count through the same media expansion the generation path runs, or return the prepared request's post-expansion token ids, without loading media twice where avoidable. Cover all three counting routes, with the `count_tokens` handler in `routes/anthropic.rs` calling the shared counter instead of keeping its own copy.
## Acceptance criteria
- [ ] Text-only bodies return the same count as today on all three routes.
- [ ] For an image body and for a video-fallback body, `input_tokens` equals `usage.prompt_tokens` of the same body; `/v1/messages/count_tokens` matches `/v1/messages` usage for an image body (tests).
- [ ] The same equality holds on a real checkpoint (`gemma-3-4b-it-4bit` with the PR #1777 two-clip body).
## Verification
```bash
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --all -- --check
```
Manual: serve `models/mlx/gemma-3-4b-it-4bit`, post one image body to `/v1/chat/completions/input_tokens` and to `/v1/chat/completions`, and confirm `input_tokens` equals `usage.prompt_tokens`.
Contributor guide
Research direction
Start with src/server/routes/prompt_inspection.rs, especially count_prompt_tokens and the chat_input_tokens and responses_input_tokens handlers, then compare them with src/server/routes/anthropic.rs and the model-worker expansion path. Trace how media reaches prepare_request_vlm_embeddings and the shared counting behavior. Done means all three routes preserve text-only counts and match usage.prompt_tokens for image and video-fallback bodies, with tests passing.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, api, backend
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100