diff --git a/docs/configuration.md b/docs/configuration.md index ea0f642..b8ba8a9 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -116,6 +116,12 @@ knowledge_budget_pct = 15.0 # Max graph-traversal hops for knowledge expansion knowledge_max_hops = 2 +# Unit the context budget is enforced in: "line" (count source lines) or +# "token" (estimate tokens per chunk, ~chars/4). Token mode makes injection +# size predictable against the model window. The budget value (e.g. the hook +# `budget`) is then interpreted in this unit. +budget_unit = "line" + [feedback] # Maximum feedback boost multiplier. Actual boost = min(score * boost_weight, boost_max). boost_max = 0.3 @@ -198,6 +204,7 @@ Tunes context assembly — the bridging + knowledge-expansion pipeline that is B | `coupling_threshold` | float | `0.1` | Minimum coupling **score** (0.0–1.0) for a coupled file to enter context. Distinct from `[git].coupling_threshold`, which is an integer co-change **count** applied during indexing. | | `knowledge_budget_pct` | float | `15.0` | Percent of the line budget reserved for knowledge-graph expansion (requires the `knowledge` feature). | | `knowledge_max_hops` | int | `2` | Max graph-traversal hops for knowledge expansion. | +| `budget_unit` | string | `"line"` | Unit the context budget is enforced in: `"line"` (count source lines) or `"token"` (estimate tokens per chunk, ~chars/4). Token mode makes injection size predictable against the model window; the budget value is then interpreted in tokens. | ### `[feedback]` diff --git a/src/cli/context.rs b/src/cli/context.rs index 3933419..a92450a 100644 --- a/src/cli/context.rs +++ b/src/cli/context.rs @@ -187,6 +187,7 @@ pub async fn run(args: ContextArgs, output: OutputConfig) -> Result<()> { let context_config = ContextConfig { budget_lines: args.budget, + budget_unit: config.context.budget_unit, depth: args.depth, max_coupled: args.max_coupled, coupling_threshold: args.coupling_threshold.unwrap_or(config.context.coupling_threshold), diff --git a/src/cli/hook.rs b/src/cli/hook.rs index ca25753..48b8408 100644 --- a/src/cli/hook.rs +++ b/src/cli/hook.rs @@ -2846,6 +2846,7 @@ async fn inject_context_inner(args: InjectContextArgs) -> Result<()> { let context_config = ContextConfig { budget_lines: cal_budget.unwrap_or(budget), + budget_unit: config.context.budget_unit, depth: 1, max_coupled: 2, // Tightened from 3 to reduce coupled noise (matches remote mode) coupling_threshold: adj.coupling_threshold.unwrap_or(config.context.coupling_threshold), diff --git a/src/config.rs b/src/config.rs index 424cfdd..4d2b116 100644 --- a/src/config.rs +++ b/src/config.rs @@ -374,6 +374,10 @@ pub struct ContextTuning { pub knowledge_budget_pct: f32, /// Maximum graph-traversal hops for knowledge expansion. Default: 2. pub knowledge_max_hops: u32, + /// Unit the context budget is enforced in: `"line"` (count source lines) or + /// `"token"` (estimate tokens per chunk, ~chars/4). Token mode makes injection + /// size predictable against the model window. Default: `"line"`. + pub budget_unit: crate::search::context::BudgetUnit, } impl Default for ContextTuning { @@ -385,6 +389,7 @@ impl Default for ContextTuning { coupling_threshold: 0.1, knowledge_budget_pct: 15.0, knowledge_max_hops: 2, + budget_unit: crate::search::context::BudgetUnit::Line, } } } @@ -1206,6 +1211,7 @@ coupling_depth = 500 assert!((c.context.coupling_threshold - 0.1).abs() < f32::EPSILON); assert!((c.context.knowledge_budget_pct - 15.0).abs() < f32::EPSILON); assert_eq!(c.context.knowledge_max_hops, 2); + assert_eq!(c.context.budget_unit, crate::search::context::BudgetUnit::Line); assert!((c.feedback.boost_max - 0.3).abs() < f32::EPSILON); assert!((c.feedback.boost_weight - 0.2).abs() < f32::EPSILON); } @@ -1220,6 +1226,7 @@ max_bridged_chunks_per_file = 3 coupling_threshold = 0.25 knowledge_budget_pct = 20.0 knowledge_max_hops = 3 +budget_unit = "token" [feedback] boost_max = 0.4 @@ -1232,6 +1239,7 @@ boost_weight = 0.15 assert!((config.context.coupling_threshold - 0.25).abs() < f32::EPSILON); assert!((config.context.knowledge_budget_pct - 20.0).abs() < f32::EPSILON); assert_eq!(config.context.knowledge_max_hops, 3); + assert_eq!(config.context.budget_unit, crate::search::context::BudgetUnit::Token); assert!((config.feedback.boost_max - 0.4).abs() < f32::EPSILON); assert!((config.feedback.boost_weight - 0.15).abs() < f32::EPSILON); } diff --git a/src/http/handlers/context.rs b/src/http/handlers/context.rs index c9cd621..66b7d0d 100644 --- a/src/http/handlers/context.rs +++ b/src/http/handlers/context.rs @@ -141,6 +141,7 @@ pub(super) async fn context( let context_config = ContextConfig { budget_lines: params.budget.unwrap_or(500), + budget_unit: state.config.context.budget_unit, depth: params.depth.unwrap_or(1), max_coupled: params.max_coupled.unwrap_or(3), coupling_threshold: params.coupling_threshold.unwrap_or(state.config.context.coupling_threshold), diff --git a/src/search/context.rs b/src/search/context.rs index 93b9033..12fa763 100644 --- a/src/search/context.rs +++ b/src/search/context.rs @@ -52,9 +52,72 @@ impl std::str::FromStr for BridgeMode { } } +/// Unit in which the context-assembly budget is counted. Agents have *token* +/// budgets, but a 50-line dense-code chunk costs far more tokens than 50 lines +/// of comments. `Token` makes injection size predictable against the model +/// window; `Line` preserves the historical line-count behavior. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum BudgetUnit { + /// Count budget in source lines (`end_line - start_line + 1`). Historical default. + #[default] + Line, + /// Count budget in estimated tokens (see [`estimate_tokens`]). + Token, +} + +impl std::fmt::Display for BudgetUnit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BudgetUnit::Line => write!(f, "line"), + BudgetUnit::Token => write!(f, "token"), + } + } +} + +impl std::str::FromStr for BudgetUnit { + type Err = anyhow::Error; + fn from_str(s: &str) -> Result { + match s { + "line" | "lines" => Ok(BudgetUnit::Line), + "token" | "tokens" => Ok(BudgetUnit::Token), + _ => anyhow::bail!("Unknown budget_unit '{}'. Use: line, token", s), + } + } +} + +/// Estimate the token count of a text fragment. +/// +/// Uses the standard `chars / 4` heuristic (≈4 chars per token for English + +/// code) rather than a real tokenizer — it needs no model files, is +/// deterministic, and is accurate enough for budget *accounting* (we only need +/// relative sizing to keep injection under the model window, not exact counts). +/// Non-empty input always costs at least 1 token. +pub fn estimate_tokens(text: &str) -> usize { + let chars = text.chars().count(); + if chars == 0 { + 0 + } else { + chars.div_ceil(4) + } +} + +/// Cost of a chunk against the assembly budget, in the configured unit. +/// `Line` uses the chunk's line span; `Token` estimates from its content. +fn chunk_cost(unit: BudgetUnit, content: &str, start_line: u32, end_line: u32) -> usize { + match unit { + BudgetUnit::Line => (end_line.saturating_sub(start_line) + 1) as usize, + BudgetUnit::Token => estimate_tokens(content), + } +} + /// Configuration for context assembly pub struct ContextConfig { pub budget_lines: usize, + /// Unit the budget is enforced in: `Line` (count source lines) or `Token` + /// (estimate tokens per chunk). When `Token`, `budget_lines` and all derived + /// caps are interpreted as token counts. Default: `Line`. + pub budget_unit: BudgetUnit, pub depth: u32, pub max_coupled: usize, pub coupling_threshold: f32, @@ -127,6 +190,7 @@ impl Default for ContextConfig { fn default() -> Self { Self { budget_lines: 500, + budget_unit: BudgetUnit::Line, depth: 1, max_coupled: 3, coupling_threshold: 0.1, @@ -1375,7 +1439,7 @@ fn assemble_bundle( if seen_chunk_ids.contains(&result.chunk_id) { continue; } - let chunk_lines = (result.end_line - result.start_line + 1) as usize; + let chunk_lines = chunk_cost(config.budget_unit, &result.content, result.start_line, result.end_line); let capped_lines = chunk_lines.min(max_chunk_lines); if used_lines + capped_lines > pin_budget_reserve { break; @@ -1463,7 +1527,7 @@ fn assemble_bundle( continue; } - let chunk_lines = (result.end_line - result.start_line + 1) as usize; + let chunk_lines = chunk_cost(config.budget_unit, &result.content, result.start_line, result.end_line); let capped_lines = chunk_lines.min(max_chunk_lines); if used_lines + capped_lines > budget { @@ -1566,7 +1630,7 @@ fn assemble_bundle( continue; } - let chunk_lines = (chunk.end_line - chunk.start_line + 1) as usize; + let chunk_lines = chunk_cost(config.budget_unit, &chunk.content, chunk.start_line, chunk.end_line); let capped_lines = chunk_lines.min(max_chunk_lines); if used_lines + capped_lines > budget { @@ -1648,7 +1712,7 @@ fn assemble_bundle( continue; } - let chunk_lines = (chunk.end_line - chunk.start_line + 1) as usize; + let chunk_lines = chunk_cost(config.budget_unit, &chunk.content, chunk.start_line, chunk.end_line); let capped_lines = chunk_lines.min(max_chunk_lines); if used_lines + capped_lines > budget || bridged_lines + capped_lines > bridged_budget { @@ -1735,7 +1799,7 @@ fn assemble_bundle( continue; } - let chunk_lines = (chunk.end_line - chunk.start_line + 1) as usize; + let chunk_lines = chunk_cost(config.budget_unit, &chunk.content, chunk.start_line, chunk.end_line); let capped_lines = chunk_lines.min(max_chunk_lines); if used_lines + capped_lines > budget || knowledge_lines + capped_lines > knowledge_budget { @@ -1923,6 +1987,54 @@ mod tests { assert_eq!(bundle.budget.max_lines, 10); } + #[test] + fn test_estimate_tokens() { + assert_eq!(estimate_tokens(""), 0); + assert_eq!(estimate_tokens("a"), 1); // 1 char rounds up to 1 token + assert_eq!(estimate_tokens("abcd"), 1); // 4 chars => 1 + assert_eq!(estimate_tokens("abcde"), 2); // 5 chars => ceil(5/4) = 2 + assert_eq!(estimate_tokens(&"x".repeat(40)), 10); // 40 chars => 10 + } + + #[test] + fn test_budget_unit_roundtrip() { + for unit in [BudgetUnit::Line, BudgetUnit::Token] { + let s = unit.to_string(); + let parsed: BudgetUnit = s.parse().unwrap(); + assert_eq!(unit, parsed); + } + assert_eq!(BudgetUnit::default(), BudgetUnit::Line); + } + + #[test] + fn test_token_budget_enforcement() { + // Token mode counts a chunk's cost from its *content* (~chars/4), not its + // line span. A chunk with a huge line span but tiny content is cheap. + let config = ContextConfig { + budget_lines: 20, // interpreted as 20 tokens here + budget_unit: BudgetUnit::Token, + depth: 0, + bridge_mode: BridgeMode::Inject, + ..ContextConfig::default() + }; + + // a.rs: 199-line span but only 32 chars => 8 tokens (admitted despite span) + let mut a = make_seed("c1", "a.rs", 1, 200, 0.9); + a.content = "x".repeat(32); + // b.rs: 40 chars => 10 tokens; 8 + 10 = 18 <= 20, fits + let mut b = make_seed("c2", "b.rs", 1, 5, 0.8); + b.content = "x".repeat(40); + // c.rs: 16 chars => 4 tokens; 18 + 4 = 22 > 20, dropped + let mut c = make_seed("c3", "c.rs", 1, 3, 0.7); + c.content = "x".repeat(16); + + let bundle = assemble_bundle("test", &config, vec![a, b, c], vec![], vec![], vec![]).unwrap(); + + assert_eq!(bundle.budget.max_lines, 20); + assert_eq!(bundle.budget.used_lines, 18); // tokens, not lines + assert_eq!(bundle.summary.total_chunks, 2); // c dropped on budget + } + #[test] fn test_deduplication() { let config = ContextConfig {