_diffctx/
tokenizer.rs

1use once_cell::sync::Lazy;
2use tiktoken_rs::CoreBPE;
3
4#[derive(Debug, thiserror::Error)]
5pub enum TokenizerError {
6    #[error("failed to load o200k_base BPE tables: {0}")]
7    EncoderInit(String),
8}
9
10// `Lazy<Result<...>>` keeps the static initialization fallible without
11// crashing the host process (PyO3 surface) on sandboxed / proxy-blocked
12// environments where tiktoken-rs may fail to materialize the BPE tables.
13// `try_count_tokens` is the boundary-safe variant that returns the error;
14// `count_tokens` retains the infallible signature used by ~5 internal
15// hot-path call sites and degrades to a conservative byte-length estimate
16// rather than aborting the entire process.
17static ENCODER: Lazy<Result<CoreBPE, TokenizerError>> =
18    Lazy::new(|| tiktoken_rs::o200k_base().map_err(|e| TokenizerError::EncoderInit(e.to_string())));
19
20pub fn try_count_tokens(text: &str) -> Result<u32, TokenizerError> {
21    // Why `encode_ordinary` (not `encode_with_special_tokens`):
22    //
23    // 1. Budget contract (R2-T1 regression): `encode_with_special_tokens`
24    //    collapses literal `<|endoftext|>`-style sequences into a single
25    //    token, breaking byte-accurate accounting against the user budget.
26    //
27    // 2. Prompt-injection safety: a diff is user input. Treating literal
28    //    `<|...|>` sequences as opaque text prevents them from being
29    //    interpreted as model control tokens downstream.
30    match &*ENCODER {
31        Ok(enc) => Ok(enc.encode_ordinary(text).len() as u32),
32        Err(e) => Err(TokenizerError::EncoderInit(e.to_string())),
33    }
34}
35
36pub fn count_tokens(text: &str) -> u32 {
37    // Infallible variant for internal hot-path call sites. On encoder-init
38    // failure, fall back to a conservative byte-length estimate (4 bytes
39    // per token heuristic) so the pipeline degrades gracefully instead of
40    // aborting the host process.
41    match try_count_tokens(text) {
42        Ok(n) => n,
43        Err(_) if text.is_empty() => 0,
44        Err(_) => ((text.len() as u32) / 4).max(1),
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn special_token_literals_are_not_collapsed() {
54        // R2-T1 regression: `encode_with_special_tokens` would collapse
55        // a literal `<|endoftext|>` to a single token, escaping the budget
56        // contract. `encode_ordinary` treats it as plain text.
57        let n = count_tokens("literal <|endoftext|> in code");
58        assert!(
59            n > 1,
60            "tokenizer must not collapse special-token literals; got {n} tokens"
61        );
62    }
63}
_diffctx/tokenizer.rs

_diffctx/
tokenizer.rs
