| 1 | //! Large-output routing for tool results (issue #548). |
| 2 | //! |
| 3 | //! Any tool result whose estimated token count exceeds the configured threshold |
| 4 | //! is intercepted here before it reaches the parent context. A lightweight |
| 5 | //! V4-Flash synthesis sub-agent condenses the raw output; only the synthesis |
| 6 | //! is returned to the parent. The raw content is stored in the workshop |
| 7 | //! variable `last_tool_result` so the parent agent can call |
| 8 | //! `promote_to_context` later if it needs the full text. |
| 9 | //! |
| 10 | //! Per-tool thresholds can override the global default. Individual tool calls |
| 11 | //! may pass `raw=true` to bypass routing entirely. |
| 12 | |
| 13 | use std::collections::HashMap; |
| 14 | |
| 15 | use serde::{Deserialize, Serialize}; |
| 16 | |
| 17 | use crate::tools::spec::ToolResult; |
| 18 | |
| 19 | // ── Constants ────────────────────────────────────────────────────────────────── |
| 20 | |
| 21 | /// Default token threshold above which a tool result is routed through the |
| 22 | /// workshop. Matches the issue spec of 4 096 tokens. |
| 23 | pub const DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS: usize = 4_096; |
| 24 | |
| 25 | /// Approximate characters-per-token ratio used for the heuristic estimate. |
| 26 | /// We intentionally choose a conservative value (3 chars/token) so we err |
| 27 | /// on the side of routing rather than dumping raw data into the parent. |
| 28 | const CHARS_PER_TOKEN_ESTIMATE: usize = 3; |
| 29 | |
| 30 | /// Workshop variable name where the raw tool output is stored. |
| 31 | pub const WORKSHOP_LAST_TOOL_RESULT_VAR: &str = "last_tool_result"; |
| 32 | |
| 33 | // ── Configuration ───────────────────────────────────────────────────────────── |
| 34 | |
| 35 | /// `[workshop]` section in `config.toml`. |
| 36 | #[derive(Debug, Clone, Deserialize, Default)] |
| 37 | pub struct WorkshopConfig { |
| 38 | /// Token threshold above which tool results are routed through the workshop |
| 39 | /// synthesis sub-agent. Default: [`DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS`]. |
| 40 | #[serde(default)] |
| 41 | pub large_output_threshold_tokens: Option<usize>, |
| 42 | |
| 43 | /// Per-tool threshold overrides (tool name → token limit). A tool whose |
| 44 | /// name appears here uses this limit instead of |
| 45 | /// `large_output_threshold_tokens`. |
| 46 | #[serde(default)] |
| 47 | pub per_tool_thresholds: Option<HashMap<String, usize>>, |
| 48 | } |
| 49 | |
| 50 | impl WorkshopConfig { |
| 51 | /// Resolve the effective threshold for the given tool name. |
| 52 | #[must_use] |
| 53 | pub fn threshold_for(&self, tool_name: &str) -> usize { |
| 54 | if let Some(per_tool) = self.per_tool_thresholds.as_ref() |
| 55 | && let Some(&limit) = per_tool.get(tool_name) |
| 56 | { |
| 57 | return limit; |
| 58 | } |
| 59 | self.large_output_threshold_tokens |
| 60 | .unwrap_or(DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS) |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | // ── Token estimation ────────────────────────────────────────────────────────── |
| 65 | |
| 66 | /// Estimate the number of tokens in `text` using a character-count heuristic. |
| 67 | /// |
| 68 | /// This avoids a real tokeniser dependency; the estimate is deliberately |
| 69 | /// conservative (under-counts tokens) so we route aggressively rather than |
| 70 | /// letting a 5K-token blob slip through. |
| 71 | #[must_use] |
| 72 | pub fn estimate_tokens(text: &str) -> usize { |
| 73 | let chars = text.chars().count(); |
| 74 | // Round up: partial last token still costs a token. |
| 75 | chars.div_ceil(CHARS_PER_TOKEN_ESTIMATE) |
| 76 | } |
| 77 | |
| 78 | // ── Router ──────────────────────────────────────────────────────────────────── |
| 79 | |
| 80 | /// Decision returned by [`LargeOutputRouter::route`]. |
| 81 | #[derive(Debug, Clone, PartialEq)] |
| 82 | pub enum RouteDecision { |
| 83 | /// The output is small enough; pass it through unmodified. |
| 84 | PassThrough, |
| 85 | /// The output exceeded the threshold and was (or should be) synthesised. |
| 86 | Synthesise { |
| 87 | /// Estimated token count of the raw output. |
| 88 | estimated_tokens: usize, |
| 89 | /// The threshold that was breached. |
| 90 | threshold: usize, |
| 91 | }, |
| 92 | } |
| 93 | |
| 94 | /// Intercepts tool results and routes large ones through the workshop. |
| 95 | /// |
| 96 | /// This type is intentionally `Clone` and `Default` so it can be embedded |
| 97 | /// cheaply in [`ToolContext`](crate::tools::spec::ToolContext) without |
| 98 | /// requiring `Arc` wrappers. |
| 99 | #[derive(Debug, Clone, Default)] |
| 100 | pub struct LargeOutputRouter { |
| 101 | config: WorkshopConfig, |
| 102 | } |
| 103 | |
| 104 | impl LargeOutputRouter { |
| 105 | /// Construct a router from the resolved workshop config. |
| 106 | #[must_use] |
| 107 | pub fn new(config: WorkshopConfig) -> Self { |
| 108 | Self { config } |
| 109 | } |
| 110 | |
| 111 | /// Decide whether `result` for `tool_name` should be synthesised. |
| 112 | /// |
| 113 | /// Pass `raw_bypass = true` when the tool call included `raw = true`. |
| 114 | #[must_use] |
| 115 | pub fn route(&self, tool_name: &str, result: &ToolResult, raw_bypass: bool) -> RouteDecision { |
| 116 | if raw_bypass || !result.success { |
| 117 | return RouteDecision::PassThrough; |
| 118 | } |
| 119 | let threshold = self.config.threshold_for(tool_name); |
| 120 | let estimated_tokens = estimate_tokens(&result.content); |
| 121 | if estimated_tokens > threshold { |
| 122 | RouteDecision::Synthesise { |
| 123 | estimated_tokens, |
| 124 | threshold, |
| 125 | } |
| 126 | } else { |
| 127 | RouteDecision::PassThrough |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | /// Build the synthesis prompt sent to the V4-Flash workshop sub-agent. |
| 132 | /// |
| 133 | /// The prompt is intentionally terse — Flash is a fast model and we just |
| 134 | /// want a faithful summary, not deep reasoning. |
| 135 | /// |
| 136 | /// This is the building block for the live LLM synthesis call wired in |
| 137 | /// the follow-up (once the async Flash client is safe to call from the |
| 138 | /// registry layer). The method is public so callers outside this crate |
| 139 | /// can unit-test the prompt shape. |
| 140 | #[must_use] |
| 141 | #[allow(dead_code)] // used by future Flash synthesis call; keep for API stability |
| 142 | pub fn synthesis_prompt(tool_name: &str, raw_output: &str, estimated_tokens: usize) -> String { |
| 143 | format!( |
| 144 | "You are a synthesis assistant. The tool `{tool_name}` produced {estimated_tokens} tokens \ |
| 145 | of output that is too large to include directly in the parent context.\n\n\ |
| 146 | Summarise the output below into a concise, faithful synthesis of ≤ 800 words. \ |
| 147 | Preserve key facts, numbers, file paths, error messages, and any actionable \ |
| 148 | information. Do NOT add commentary or interpretation beyond what is in the source.\n\n\ |
| 149 | <raw_tool_output>\n{raw_output}\n</raw_tool_output>" |
| 150 | ) |
| 151 | } |
| 152 | |
| 153 | /// Wrap a synthesis result with a workshop provenance header and a hint |
| 154 | /// about the stored raw output. |
| 155 | #[must_use] |
| 156 | pub fn wrap_synthesis( |
| 157 | tool_name: &str, |
| 158 | synthesis: &str, |
| 159 | estimated_tokens: usize, |
| 160 | threshold: usize, |
| 161 | ) -> String { |
| 162 | format!( |
| 163 | "[workshop-synthesis: tool={tool_name}, raw_tokens≈{estimated_tokens}, \ |
| 164 | threshold={threshold}, raw_stored_in={WORKSHOP_LAST_TOOL_RESULT_VAR}]\n\n{synthesis}" |
| 165 | ) |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | // ── Workshop variable store ─────────────────────────────────────────────────── |
| 170 | |
| 171 | /// In-process store for workshop variables that persist across tool calls |
| 172 | /// within a session. The only variable exposed today is `last_tool_result` |
| 173 | /// which holds the most recent raw large-tool output for `promote_to_context`. |
| 174 | #[derive(Debug, Clone, Default, Serialize, Deserialize)] |
| 175 | pub struct WorkshopVariables { |
| 176 | /// Raw content of the most recent large tool output that was routed |
| 177 | /// through the workshop. Empty string when no routing has occurred. |
| 178 | #[serde(default)] |
| 179 | pub last_tool_result: String, |
| 180 | |
| 181 | /// Name of the tool that produced `last_tool_result`. |
| 182 | #[serde(default)] |
| 183 | pub last_tool_name: String, |
| 184 | } |
| 185 | |
| 186 | impl WorkshopVariables { |
| 187 | /// Store the raw output from a large-tool routing event. |
| 188 | pub fn store_raw(&mut self, tool_name: &str, raw: &str) { |
| 189 | self.last_tool_result = raw.to_string(); |
| 190 | self.last_tool_name = tool_name.to_string(); |
| 191 | } |
| 192 | |
| 193 | /// Retrieve and clear the stored raw output (consume semantics so the |
| 194 | /// variable is not accidentally promoted twice). |
| 195 | /// |
| 196 | /// Called by the `promote_to_context` tool (not yet wired in this PR). |
| 197 | #[must_use] |
| 198 | #[allow(dead_code)] // consumed by promote_to_context tool in follow-up |
| 199 | pub fn take_raw(&mut self) -> Option<(String, String)> { |
| 200 | if self.last_tool_result.is_empty() { |
| 201 | return None; |
| 202 | } |
| 203 | let content = std::mem::take(&mut self.last_tool_result); |
| 204 | let name = std::mem::take(&mut self.last_tool_name); |
| 205 | Some((name, content)) |
| 206 | } |
| 207 | } |
| 208 | |
| 209 | // ── Unit tests ──────────────────────────────────────────────────────────────── |
| 210 | |
| 211 | #[cfg(test)] |
| 212 | mod tests { |
| 213 | use super::*; |
| 214 | |
| 215 | fn make_result(content: &str) -> ToolResult { |
| 216 | ToolResult::success(content.to_string()) |
| 217 | } |
| 218 | |
| 219 | #[test] |
| 220 | fn pass_through_below_threshold() { |
| 221 | let router = LargeOutputRouter::default(); |
| 222 | let small = "x".repeat(100); |
| 223 | let result = make_result(&small); |
| 224 | assert_eq!( |
| 225 | router.route("read_file", &result, false), |
| 226 | RouteDecision::PassThrough |
| 227 | ); |
| 228 | } |
| 229 | |
| 230 | #[test] |
| 231 | fn synthesise_above_threshold() { |
| 232 | let router = LargeOutputRouter::default(); |
| 233 | // DEFAULT threshold = 4096 tokens; 3 chars/token → 4096*3 = 12288 chars |
| 234 | let big = "a".repeat(13_000); |
| 235 | let result = make_result(&big); |
| 236 | assert!(matches!( |
| 237 | router.route("read_file", &result, false), |
| 238 | RouteDecision::Synthesise { .. } |
| 239 | )); |
| 240 | } |
| 241 | |
| 242 | #[test] |
| 243 | fn raw_bypass_skips_routing() { |
| 244 | let router = LargeOutputRouter::default(); |
| 245 | let big = "a".repeat(13_000); |
| 246 | let result = make_result(&big); |
| 247 | // raw=true → always pass through regardless of size |
| 248 | assert_eq!( |
| 249 | router.route("exec_shell", &result, true), |
| 250 | RouteDecision::PassThrough |
| 251 | ); |
| 252 | } |
| 253 | |
| 254 | #[test] |
| 255 | fn error_results_always_pass_through() { |
| 256 | let router = LargeOutputRouter::default(); |
| 257 | let big = "error: ".repeat(2_000); |
| 258 | let result = ToolResult::error(big); |
| 259 | assert_eq!( |
| 260 | router.route("exec_shell", &result, false), |
| 261 | RouteDecision::PassThrough |
| 262 | ); |
| 263 | } |
| 264 | |
| 265 | #[test] |
| 266 | fn per_tool_threshold_override() { |
| 267 | let mut per_tool = HashMap::new(); |
| 268 | per_tool.insert("grep_files".to_string(), 100); // very low |
| 269 | let config = WorkshopConfig { |
| 270 | large_output_threshold_tokens: Some(4096), |
| 271 | per_tool_thresholds: Some(per_tool), |
| 272 | }; |
| 273 | let router = LargeOutputRouter::new(config); |
| 274 | // 100 tokens * 3 = 300 chars → trigger with 400 chars |
| 275 | let medium = "b".repeat(400); |
| 276 | let result = make_result(&medium); |
| 277 | assert!(matches!( |
| 278 | router.route("grep_files", &result, false), |
| 279 | RouteDecision::Synthesise { .. } |
| 280 | )); |
| 281 | // Other tools still use the global threshold |
| 282 | assert_eq!( |
| 283 | router.route("read_file", &result, false), |
| 284 | RouteDecision::PassThrough |
| 285 | ); |
| 286 | } |
| 287 | |
| 288 | #[test] |
| 289 | fn estimate_tokens_conservative() { |
| 290 | // 9 chars → ceil(9/3) = 3 tokens |
| 291 | assert_eq!(estimate_tokens("123456789"), 3); |
| 292 | // 10 chars → ceil(10/3) = 4 tokens |
| 293 | assert_eq!(estimate_tokens("1234567890"), 4); |
| 294 | // Empty string |
| 295 | assert_eq!(estimate_tokens(""), 0); |
| 296 | } |
| 297 | |
| 298 | #[test] |
| 299 | fn workshop_variables_store_and_take() { |
| 300 | let mut vars = WorkshopVariables::default(); |
| 301 | assert!(vars.take_raw().is_none()); |
| 302 | |
| 303 | vars.store_raw("read_file", "raw content here"); |
| 304 | let taken = vars.take_raw().expect("should have content"); |
| 305 | assert_eq!(taken.0, "read_file"); |
| 306 | assert_eq!(taken.1, "raw content here"); |
| 307 | |
| 308 | // Second take is empty — consume semantics |
| 309 | assert!(vars.take_raw().is_none()); |
| 310 | } |
| 311 | |
| 312 | #[test] |
| 313 | fn wrap_synthesis_includes_provenance_header() { |
| 314 | let wrapped = LargeOutputRouter::wrap_synthesis("web_search", "key facts here", 5000, 4096); |
| 315 | assert!(wrapped.contains("workshop-synthesis")); |
| 316 | assert!(wrapped.contains("web_search")); |
| 317 | assert!(wrapped.contains("5000")); |
| 318 | assert!(wrapped.contains("key facts here")); |
| 319 | } |
| 320 | } |
| 321 |