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. If the
6//! tokenizer fails to load we fall back to a cheap character/word estimate.
7
8use std::sync::OnceLock;
9use tokenizers::Tokenizer;
10
11/// Embedded `Xenova/claude-tokenizer` definition (Hugging Face `tokenizer.json`).
12static CLAUDE_TOKENIZER_JSON: &[u8] =
13    include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/assets/claude-tokenizer.json"));
14
15fn tokenizer() -> Option<&'static Tokenizer> {
16    static TOKENIZER: OnceLock<Option<Tokenizer>> = OnceLock::new();
17    TOKENIZER
18        .get_or_init(|| match Tokenizer::from_bytes(CLAUDE_TOKENIZER_JSON) {
19            Ok(tokenizer) => Some(tokenizer),
20            Err(error) => {
21                tracing::warn!("Failed to load embedded Claude tokenizer, using estimate: {error}");
22                None
23            }
24        })
25        .as_ref()
26}
27
28/// Count tokens the way Claude does. Falls back to [`estimate_tokens`] on failure.
29pub fn count_tokens(text: &str) -> usize {
30    match encode_ids(text) {
31        Some(ids) => ids.len(),
32        None => estimate_tokens(text),
33    }
34}
35
36/// Encode `text` into Claude token ids. Returns `None` if the tokenizer is
37/// unavailable so callers can pick a byte-based fallback.
38pub fn encode_ids(text: &str) -> Option<Vec<u32>> {
39    let tokenizer = tokenizer()?;
40    tokenizer.encode(text, false).ok().map(|encoding| encoding.get_ids().to_vec())
41}
42
43/// Decode Claude token ids back into text. Returns `None` on failure.
44pub fn decode_ids(ids: &[u32]) -> Option<String> {
45    let tokenizer = tokenizer()?;
46    tokenizer.decode(ids, false).ok()
47}
48
49/// Read a token-budget override from env var `var` (e.g.
50/// `WINX_CODING_TOKEN_BUDGET`), falling back to `default` when unset, zero, or
51/// unparseable. Lets large-context clients tune how much of each file is pulled
52/// into context — and how much saved memory is kept — without a rebuild.
53pub fn budget_from_env(var: &str, default: usize) -> usize {
54    std::env::var(var)
55        .ok()
56        .and_then(|value| value.trim().parse::<usize>().ok())
57        .filter(|&n| n > 0)
58        .unwrap_or(default)
59}
60
61/// Cheap fallback estimate used only when the tokenizer is unavailable.
62pub fn estimate_tokens(text: &str) -> usize {
63    text.chars().count().div_ceil(4).max(text.split_whitespace().count())
64}
65
66#[cfg(test)]
67mod budget_tests {
68    use super::budget_from_env;
69
70    #[test]
71    fn falls_back_to_default_when_unset_or_invalid() {
72        // A name nothing sets -> default. (Avoids mutating process env in tests.)
73        assert_eq!(budget_from_env("WINX_DEFINITELY_UNSET_BUDGET_XYZ", 24_000), 24_000);
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[test]
82    fn counts_tokens_for_simple_text() {
83        // Whatever the backend, a non-empty string must produce at least one token.
84        assert!(count_tokens("hello world") >= 1);
85        assert_eq!(count_tokens(""), 0);
86    }
87
88    #[test]
89    fn estimate_is_nonzero_for_words() {
90        assert!(estimate_tokens("a b c d") >= 4);
91    }
92}
winx_code_agent/utils/encoder.rs

winx_code_agent/utils/
encoder.rs
