diff --git a/Cargo.lock b/Cargo.lock index 987259a..15afd3a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -196,6 +196,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5cee35f73844aa3014bb606320a6c1f010249dbdf43342fe54b5a4f6a8ed4b79" dependencies = [ "memchr", + "regex-automata", "serde_core", ] @@ -807,6 +808,17 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "fancy-regex" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" +dependencies = [ + "bit-set 0.8.0", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fast-srgb8" version = "1.0.0" @@ -3521,6 +3533,21 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "tiktoken-rs" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "027853bbf8c7763b77c5c595f1c271c7d536ced7d6f83452911b944621e57fc2" +dependencies = [ + "anyhow", + "base64", + "bstr", + "fancy-regex 0.17.0", + "lazy_static", + "regex", + "rustc-hash", +] + [[package]] name = "time" version = "0.3.53" @@ -4487,6 +4514,7 @@ dependencies = [ "sha2 0.11.0", "similar", "syntect", + "tiktoken-rs", "tokio", "tracing", "tracing-subscriber", diff --git a/Cargo.toml b/Cargo.toml index 8b9bd46..ff824c0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,6 +58,7 @@ chrono = { version = "0.4", features = ["serde"] } tracing-subscriber = { version = "0.3", features = ["env-filter"] } webbrowser = "1" lsp-types = "0.97" +tiktoken-rs = "0.12" [[bin]] name = "zesdex" diff --git a/src/app/runtime/context/mod.rs b/src/app/runtime/context/mod.rs new file mode 100644 index 0000000..c153a92 --- /dev/null +++ b/src/app/runtime/context/mod.rs @@ -0,0 +1,13 @@ +//! Context management: token counting, cross-call tool-result dedup, +//! per-result compression, budget-based shaping, and shared +//! context-window resolution — replaces `runtime::shortsend`. +//! +//! No facade function here: `dedup`, `shaping`, and `tokens` are called +//! directly from each call site (the per-turn auto-compaction loop in +//! `actions::run_agent_turn`, and `Action::Compact`), matching this +//! codebase's "no DI, call modules directly" convention. An orchestration +//! layer would only serve one of the two callers generically — the +//! auto-loop already needs per-stage control to decide when to emit +//! `TurnEvent::Compacted`. + +pub mod tokens; diff --git a/src/app/runtime/context/tokens.rs b/src/app/runtime/context/tokens.rs new file mode 100644 index 0000000..eeb0ae5 --- /dev/null +++ b/src/app/runtime/context/tokens.rs @@ -0,0 +1,69 @@ +//! Unified token-count estimation for context-window budgeting. +//! +//! Flow: text -> `tiktoken_rs::o200k_base_singleton()` (BPE vocab embedded +//! in the binary via `include_str!`, no network access) -> `encode_ordinary` +//! -> token count. +//! +//! Why: replaces three independent char-count heuristics that disagreed +//! with each other (`/3` in the old `shortsend.rs`, `/4` in the turn +//! loop, `/4` again in the status bar) with one real BPE tokenizer. +//! `o200k_base` is an approximation for non-OpenAI providers but is far +//! closer than a flat byte-per-token guess; it's only used for the +//! 85%/95% budget thresholds, not for billing-accurate counts. + +use crate::dto::chat::message::ChatMessage; + +/// Count tokens in a single string under `o200k_base`. +/// +/// Return: the BPE token count for `text`. `encode_ordinary` (not +/// `encode`/`encode_with_special_tokens`) is used deliberately — message +/// content that happens to contain a special-token-shaped substring +/// (e.g. literal text `<|endoftext|>` pasted by a user) must be counted +/// as ordinary text, not interpreted as a control token. +pub fn count_tokens(text: &str) -> usize { + tiktoken_rs::o200k_base_singleton().encode_ordinary(text).len() +} + +/// Count tokens in a `ChatMessage`'s text content. +/// +/// Return: 0 for a message with no `content` (e.g. an assistant message +/// that only carries `tool_calls`). +pub fn count_message_tokens(msg: &ChatMessage) -> usize { + msg.content.as_deref().map_or(0, count_tokens) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dto::chat::message::ChatMessage; + + #[test] + fn empty_string_has_zero_tokens() { + assert_eq!(count_tokens(""), 0); + } + + #[test] + fn known_short_phrase_has_expected_token_count() { + // Verified empirically against tiktoken-rs 0.12's o200k_base: + // "hello world" -> [24912, 2375], i.e. 2 tokens. + assert_eq!(count_tokens("hello world"), 2); + } + + #[test] + fn known_code_snippet_has_expected_token_count() { + // Verified empirically: 9 tokens under o200k_base. + assert_eq!(count_tokens("fn main() { println!(\"hi\"); }"), 9); + } + + #[test] + fn message_with_no_content_counts_zero() { + let msg = ChatMessage::assistant(None); + assert_eq!(count_message_tokens(&msg), 0); + } + + #[test] + fn message_token_count_matches_count_tokens_on_its_content() { + let msg = ChatMessage::user("hello world"); + assert_eq!(count_message_tokens(&msg), count_tokens("hello world")); + } +} diff --git a/src/app/runtime/mod.rs b/src/app/runtime/mod.rs index b0ffa4e..9c4dce8 100644 --- a/src/app/runtime/mod.rs +++ b/src/app/runtime/mod.rs @@ -2,5 +2,6 @@ //! and the LLM streaming pipeline. pub mod actions; pub mod commands; +pub mod context; pub mod shortsend; pub mod stream;