winx_code_agent/utils/
encoder.rs

1//! Claude-compatible token counting.
2//!
3//! WCGW counts tokens with the `Xenova/claude-tokenizer` (Hugging Face `tokenizers`).
4//! We embed that same tokenizer definition in the binary and load it lazily, so token
5//! budgets and truncation match the model that actually runs the agent. Small inputs
6//! use a byte-length upper bound and never initialize it; if loading fails for a larger
7//! input we fall back to a cheap character/word estimate.
8
9use std::sync::OnceLock;
10use tokenizers::Tokenizer;
11
12/// Embedded `Xenova/claude-tokenizer` definition (Hugging Face `tokenizer.json`).
13static CLAUDE_TOKENIZER_JSON: &[u8] =
14    include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/assets/claude-tokenizer.json"));
15
16fn tokenizer() -> Option<&'static Tokenizer> {
17    static TOKENIZER: OnceLock<Option<Tokenizer>> = OnceLock::new();
18    TOKENIZER
19        .get_or_init(|| match Tokenizer::from_bytes(CLAUDE_TOKENIZER_JSON) {
20            Ok(tokenizer) => Some(tokenizer),
21            Err(error) => {
22                tracing::warn!("Failed to load embedded Claude tokenizer, using estimate: {error}");
23                None
24            }
25        })
26        .as_ref()
27}
28
29/// Count tokens the way Claude does. Falls back to [`estimate_tokens`] on failure.
30pub fn count_tokens(text: &str) -> usize {
31    match encode_ids(text) {
32        Some(ids) => ids.len(),
33        None => estimate_tokens(text),
34    }
35}
36
37/// Encode `text` into Claude token ids. Returns `None` if the tokenizer is
38/// unavailable so callers can pick a byte-based fallback.
39pub fn encode_ids(text: &str) -> Option<Vec<u32>> {
40    let tokenizer = tokenizer()?;
41    tokenizer.encode(text, false).ok().map(|encoding| encoding.get_ids().to_vec())
42}
43
44/// Decode Claude token ids back into text. Returns `None` on failure.
45pub fn decode_ids(ids: &[u32]) -> Option<String> {
46    let tokenizer = tokenizer()?;
47    tokenizer.decode(ids, false).ok()
48}
49
50/// Return `true` when byte length alone proves `text` fits the token budget.
51///
52/// The embedded tokenizer is byte-level BPE without added special tokens, so it
53/// cannot emit more tokens than the number of UTF-8 input bytes. This conservative
54/// fast path avoids loading or running the tokenizer for small payloads while exact
55/// tokenization remains in charge near and above the budget boundary.
56pub fn definitely_fits_token_budget(text: &str, max_tokens: usize) -> bool {
57    text.len() <= max_tokens
58}
59
60/// Read a token-budget override from env var `var` (e.g.
61/// `WINX_CODING_TOKEN_BUDGET`), falling back to `default` when unset, zero, or
62/// unparseable. Lets large-context clients tune how much of each file is pulled
63/// into context — and how much saved memory is kept — without a rebuild.
64pub fn budget_from_env(var: &str, default: usize) -> usize {
65    std::env::var(var)
66        .ok()
67        .and_then(|value| value.trim().parse::<usize>().ok())
68        .filter(|&n| n > 0)
69        .unwrap_or(default)
70}
71
72/// Cheap fallback estimate used only when the tokenizer is unavailable.
73pub fn estimate_tokens(text: &str) -> usize {
74    text.chars().count().div_ceil(4).max(text.split_whitespace().count())
75}
76
77#[cfg(test)]
78mod budget_tests {
79    use super::budget_from_env;
80
81    #[test]
82    fn falls_back_to_default_when_unset_or_invalid() {
83        // A name nothing sets -> default. (Avoids mutating process env in tests.)
84        assert_eq!(budget_from_env("WINX_DEFINITELY_UNSET_BUDGET_XYZ", 24_000), 24_000);
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn counts_tokens_for_simple_text() {
94        // Whatever the backend, a non-empty string must produce at least one token.
95        assert!(count_tokens("hello world") >= 1);
96        assert_eq!(count_tokens(""), 0);
97    }
98
99    #[test]
100    fn estimate_is_nonzero_for_words() {
101        assert!(estimate_tokens("a b c d") >= 4);
102    }
103
104    #[test]
105    fn byte_length_is_a_safe_token_upper_bound() {
106        for sample in [
107            "plain English prose",
108            "fn main() { println!(\"hello\"); }",
109            "áéíóú 日本語 🚀",
110            "symbols: !@#$%^&*()[]{}",
111            "line one\nline two\nline three",
112        ] {
113            let ids = encode_ids(sample);
114            assert!(
115                ids.as_ref().is_some_and(|ids| ids.len() <= sample.len()),
116                "sample produced more tokens than UTF-8 bytes"
117            );
118        }
119    }
120
121    #[test]
122    fn byte_upper_bound_only_accepts_payloads_proven_to_fit() {
123        assert!(definitely_fits_token_budget("four", 4));
124        assert!(!definitely_fits_token_budget("five!", 4));
125        assert!(!definitely_fits_token_budget("é", 1));
126        assert!(definitely_fits_token_budget("é", 2));
127    }
128}
winx_code_agent/utils/encoder.rs

winx_code_agent/utils/
encoder.rs
