From 8207c547beac52c4b97a6f934efd4dabcbde57ea Mon Sep 17 00:00:00 2001 From: Michael Neale Date: Wed, 10 Jun 2026 18:54:15 +1000 Subject: [PATCH] feat: gate analyze (tree-sitter) and tiktoken behind default-on features MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both are pure data weight in the final binary: the nine tree-sitter grammar state tables are ~14 MB and the tiktoken BPE vocab files ~8 MB (tiktoken-rs include_str!s all four vocabs even though goose only uses o200k_base). Together they are two thirds of a size-optimized no-default-features binary (32 MB -> 14 MB measured on macOS arm64 with opt-level=z, fat LTO, panic=abort, strip). - analyze: gates tree-sitter + grammars, the analyze platform extension, and the analyze_cli bin. Off = the analyze tools are simply not registered. - tiktoken: gates tiktoken-rs. Off = TokenCounter estimates tokens from byte length (len/3, rounding up) — a mild over-estimate, which is the safe direction for the context-budget and compaction decisions TokenCounter feeds. Both features are default-on in goose, goose-cli, and goose-server, so normal builds are unchanged. Size-sensitive embedders (e.g. bundling a stdio ACP agent) build with --no-default-features --features rustls-tls. --- Cargo.toml | 13 ++++++ crates/goose-cli/Cargo.toml | 4 ++ crates/goose-server/Cargo.toml | 4 ++ crates/goose/Cargo.toml | 44 ++++++++++++++----- .../src/agents/platform_extensions/mod.rs | 2 + crates/goose/src/token_counter.rs | 37 +++++++++++++--- 6 files changed, 87 insertions(+), 17 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 54162c1e16..7d9745f51c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -123,3 +123,16 @@ cudaforge = { git = "https://github.com/jbg/cudaforge", rev = "e7c1967340e40673d # full debug info and stay fully debuggable; release builds are unaffected. [profile.dev.package."*"] debug = false + +# Size-optimized release profile for shipping small binaries (e.g. a stdio +# ACP agent bundled into another app). Measured on macOS arm64: goose-cli +# with --no-default-features --features rustls-tls drops from 98 MB +# (release, stripped) to 43 MB. panic=abort is safe for the CLI (smoke +# tested); opt-in only — default release builds are unchanged. +[profile.release-small] +inherits = "release" +opt-level = "z" +lto = "fat" +codegen-units = 1 +panic = "abort" +strip = true diff --git a/crates/goose-cli/Cargo.toml b/crates/goose-cli/Cargo.toml index e47f98bf55..609ab92cd1 100644 --- a/crates/goose-cli/Cargo.toml +++ b/crates/goose-cli/Cargo.toml @@ -83,7 +83,11 @@ default = [ "rustls-tls", "system-keyring", "update", + "analyze", + "tiktoken", ] +analyze = ["goose/analyze"] +tiktoken = ["goose/tiktoken"] code-mode = ["goose/code-mode"] local-inference = ["goose/local-inference"] aws-providers = ["goose/aws-providers"] diff --git a/crates/goose-server/Cargo.toml b/crates/goose-server/Cargo.toml index 50062cbe73..dda5cda88b 100644 --- a/crates/goose-server/Cargo.toml +++ b/crates/goose-server/Cargo.toml @@ -21,7 +21,11 @@ default = [ "otel", "rustls-tls", "system-keyring", + "analyze", + "tiktoken", ] +analyze = ["goose/analyze"] +tiktoken = ["goose/tiktoken"] code-mode = ["goose/code-mode"] local-inference = ["goose/local-inference"] aws-providers = ["goose/aws-providers"] diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 419e8349db..0193e1218a 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -9,7 +9,26 @@ repository.workspace = true description.workspace = true [features] -default = [] +default = ["analyze", "tiktoken"] +# Code-structure analysis platform extension (tree-sitter + per-language +# grammars). The grammar state tables are ~14 MB of the final binary; size- +# sensitive embedders can turn this off and lose only the `analyze` tools. +analyze = [ + "dep:tree-sitter", + "dep:tree-sitter-go", + "dep:tree-sitter-java", + "dep:tree-sitter-javascript", + "dep:tree-sitter-kotlin-ng", + "dep:tree-sitter-python", + "dep:tree-sitter-ruby", + "dep:tree-sitter-rust", + "dep:tree-sitter-swift", + "dep:tree-sitter-typescript", +] +# Exact BPE token counting via tiktoken-rs (~8 MB of embedded vocab data in +# the final binary). Without it, TokenCounter falls back to a conservative +# bytes-per-token estimate — fine for context-budget/compaction decisions. +tiktoken = ["dep:tiktoken-rs"] telemetry = [] otel = [ "dep:tracing-opentelemetry", @@ -106,7 +125,7 @@ async-trait = { workspace = true } async-stream = { workspace = true } minijinja = { version = "2.18", default-features = false, features = ["loader", "multi_template", "serde"] } include_dir = { workspace = true } -tiktoken-rs = { version = "0.12", default-features = false } +tiktoken-rs = { version = "0.12", default-features = false, optional = true } chrono = { workspace = true } clap = { workspace = true } indoc = { workspace = true } @@ -187,16 +206,16 @@ shellexpand = { workspace = true } indexmap = { version = "2.9", default-features = false, features = ["std"] } ignore = { workspace = true } rayon = { workspace = true } -tree-sitter = { workspace = true } -tree-sitter-go = { workspace = true } -tree-sitter-java = { workspace = true } -tree-sitter-javascript = { workspace = true } -tree-sitter-kotlin-ng = { workspace = true } -tree-sitter-python = { workspace = true } -tree-sitter-ruby = { workspace = true } -tree-sitter-rust = { workspace = true } -tree-sitter-swift = { workspace = true } -tree-sitter-typescript = { workspace = true } +tree-sitter = { workspace = true, optional = true } +tree-sitter-go = { workspace = true, optional = true } +tree-sitter-java = { workspace = true, optional = true } +tree-sitter-javascript = { workspace = true, optional = true } +tree-sitter-kotlin-ng = { workspace = true, optional = true } +tree-sitter-python = { workspace = true, optional = true } +tree-sitter-ruby = { workspace = true, optional = true } +tree-sitter-rust = { workspace = true, optional = true } +tree-sitter-swift = { workspace = true, optional = true } +tree-sitter-typescript = { workspace = true, optional = true } which = { workspace = true } pulldown-cmark = { version = "0.13", default-features = false } pastey = { version = "0.2", default-features = false } @@ -275,6 +294,7 @@ required-features = ["local-inference"] [[bin]] name = "analyze_cli" path = "src/bin/analyze_cli.rs" +required-features = ["analyze"] [[bin]] name = "build_canonical_models" diff --git a/crates/goose/src/agents/platform_extensions/mod.rs b/crates/goose/src/agents/platform_extensions/mod.rs index adef19cf99..23333ba547 100644 --- a/crates/goose/src/agents/platform_extensions/mod.rs +++ b/crates/goose/src/agents/platform_extensions/mod.rs @@ -1,3 +1,4 @@ +#[cfg(feature = "analyze")] pub mod analyze; pub mod apps; pub mod chatrecall; @@ -29,6 +30,7 @@ pub static PLATFORM_EXTENSIONS: Lazy || { let mut map = HashMap::new(); + #[cfg(feature = "analyze")] map.insert( analyze::EXTENSION_NAME, PlatformExtensionDef { diff --git a/crates/goose/src/token_counter.rs b/crates/goose/src/token_counter.rs index 6b1f8bb5ff..d9a9057d3a 100644 --- a/crates/goose/src/token_counter.rs +++ b/crates/goose/src/token_counter.rs @@ -1,14 +1,26 @@ use lru::LruCache; use rmcp::model::Tool; use std::num::NonZeroUsize; -use std::sync::{Arc, Mutex}; +#[cfg(feature = "tiktoken")] +use std::sync::Arc; +use std::sync::Mutex; +#[cfg(feature = "tiktoken")] use tiktoken_rs::CoreBPE; +#[cfg(feature = "tiktoken")] use tokio::sync::OnceCell; use crate::conversation::message::Message; +#[cfg(feature = "tiktoken")] static TOKENIZER: OnceCell> = OnceCell::const_new(); +/// Bytes-per-token divisor for the estimator used when the `tiktoken` feature +/// is disabled. Real content sits around 3–4.5 bytes/token (code at the low +/// end), so dividing by 3 gives a mild over-estimate — the safe direction for +/// context-budget and compaction decisions, which is all TokenCounter feeds. +#[cfg(not(feature = "tiktoken"))] +const ESTIMATE_BYTES_PER_TOKEN: usize = 3; + const MAX_TOKEN_CACHE_SIZE: usize = 1_024; // token use for various bits of a tool calls: @@ -20,6 +32,7 @@ const ENUM_ITEM: usize = 3; const FUNC_END: usize = 12; pub struct TokenCounter { + #[cfg(feature = "tiktoken")] tokenizer: Arc, token_cache: Mutex>, } @@ -41,15 +54,29 @@ impl TokenCacheKey { impl TokenCounter { pub async fn new() -> Result { - let tokenizer = get_tokenizer().await?; let cache_capacity = NonZeroUsize::new(MAX_TOKEN_CACHE_SIZE).expect("token cache capacity must be non-zero"); Ok(Self { - tokenizer, + #[cfg(feature = "tiktoken")] + tokenizer: get_tokenizer().await?, token_cache: Mutex::new(LruCache::new(cache_capacity)), }) } + /// Tokenize with the BPE tokenizer (`tiktoken` feature), or estimate from + /// byte length when built without it (mild over-estimate — see + /// `ESTIMATE_BYTES_PER_TOKEN`). + fn count_uncached(&self, text: &str) -> usize { + #[cfg(feature = "tiktoken")] + { + self.tokenizer.encode_with_special_tokens(text).len() + } + #[cfg(not(feature = "tiktoken"))] + { + text.len().div_ceil(ESTIMATE_BYTES_PER_TOKEN) + } + } + pub fn count_tokens(&self, text: &str) -> usize { let cache_key = TokenCacheKey::from_text(text); if let Some(count) = self @@ -62,8 +89,7 @@ impl TokenCounter { return count; } - let tokens = self.tokenizer.encode_with_special_tokens(text); - let count = tokens.len(); + let count = self.count_uncached(text); self.token_cache .lock() @@ -202,6 +228,7 @@ impl TokenCounter { } } +#[cfg(feature = "tiktoken")] async fn get_tokenizer() -> Result, String> { Ok(TOKENIZER .get_or_init(|| async {