| 1 | //! Adaptive evidence routing for tool results (#4619). |
| 2 | //! |
| 3 | //! Results are classified as inline, hybrid, or handle-only before they enter |
| 4 | //! model context. Non-inline results are published exactly once under their |
| 5 | //! origin session and remain available through bounded retrieval. The earlier |
| 6 | //! workshop preview behavior remains only behind the explicit classic-output |
| 7 | //! rollback switch. |
| 8 | |
| 9 | use std::collections::HashMap; |
| 10 | use std::io; |
| 11 | use std::path::PathBuf; |
| 12 | |
| 13 | use serde::{Deserialize, Serialize}; |
| 14 | |
| 15 | use crate::tools::spec::ToolResult; |
| 16 | |
| 17 | // ── Constants ────────────────────────────────────────────────────────────────── |
| 18 | |
| 19 | /// Default token threshold separating hybrid from handle-only evidence. |
| 20 | /// |
| 21 | /// 32K tokens (≈96 KiB of text at the 3 chars/token estimate) keeps ordinary |
| 22 | /// tool results — file reads, test runs, build logs up to a few thousand |
| 23 | /// lines — fully inline. Only genuinely large outputs spill to evidence |
| 24 | /// artifacts, where the model-facing preview names the artifact path and how |
| 25 | /// to recover the omitted range. |
| 26 | pub const DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS: usize = 32_768; |
| 27 | |
| 28 | /// Approximate characters-per-token ratio used for the heuristic estimate. |
| 29 | /// We intentionally choose a conservative value (3 chars/token) so we err |
| 30 | /// on the side of routing rather than dumping raw data into the parent. |
| 31 | const CHARS_PER_TOKEN_ESTIMATE: usize = 3; |
| 32 | |
| 33 | /// Workshop variable name where the raw tool output is stored. |
| 34 | pub const WORKSHOP_LAST_TOOL_RESULT_VAR: &str = "last_tool_result"; |
| 35 | |
| 36 | // ── Configuration ───────────────────────────────────────────────────────────── |
| 37 | |
| 38 | /// Existing `[workshop]` threshold configuration, retained for compatibility. |
| 39 | #[derive(Debug, Clone, Deserialize, Default)] |
| 40 | pub struct WorkshopConfig { |
| 41 | /// Token threshold above which results become handle-only evidence. |
| 42 | #[serde(default)] |
| 43 | pub large_output_threshold_tokens: Option<usize>, |
| 44 | |
| 45 | /// Per-tool threshold overrides (tool name → token limit). A tool whose |
| 46 | /// name appears here uses this limit instead of |
| 47 | /// `large_output_threshold_tokens`. |
| 48 | #[serde(default)] |
| 49 | pub per_tool_thresholds: Option<HashMap<String, usize>>, |
| 50 | } |
| 51 | |
| 52 | impl WorkshopConfig { |
| 53 | /// Resolve the effective threshold for the given tool name. |
| 54 | #[must_use] |
| 55 | pub fn threshold_for(&self, tool_name: &str) -> usize { |
| 56 | if let Some(per_tool) = self.per_tool_thresholds.as_ref() |
| 57 | && let Some(&limit) = per_tool.get(tool_name) |
| 58 | { |
| 59 | return limit; |
| 60 | } |
| 61 | self.large_output_threshold_tokens |
| 62 | .unwrap_or(DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS) |
| 63 | } |
| 64 | } |
| 65 | |
| 66 | // ── Token estimation ────────────────────────────────────────────────────────── |
| 67 | |
| 68 | /// Estimate the number of tokens in `text` using a character-count heuristic. |
| 69 | /// |
| 70 | /// This avoids a real tokeniser dependency; the estimate is deliberately |
| 71 | /// conservative (under-counts tokens) so we route aggressively rather than |
| 72 | /// letting a 5K-token blob slip through. |
| 73 | #[must_use] |
| 74 | pub fn estimate_tokens(text: &str) -> usize { |
| 75 | let chars = text.chars().count(); |
| 76 | // Round up: partial last token still costs a token. |
| 77 | chars.div_ceil(CHARS_PER_TOKEN_ESTIMATE) |
| 78 | } |
| 79 | |
| 80 | // ── Router ──────────────────────────────────────────────────────────────────── |
| 81 | |
| 82 | /// Decision returned by [`LargeOutputRouter::route`]. |
| 83 | #[derive(Debug, Clone, PartialEq)] |
| 84 | pub enum RouteDecision { |
| 85 | /// The output is small enough; pass it through unmodified. |
| 86 | PassThrough, |
| 87 | /// The output exceeded the threshold and was (or should be) synthesised. |
| 88 | Synthesise { |
| 89 | /// Estimated token count of the raw output. |
| 90 | estimated_tokens: usize, |
| 91 | /// The threshold that was breached. |
| 92 | threshold: usize, |
| 93 | }, |
| 94 | } |
| 95 | |
| 96 | /// Intercepts tool results and routes large ones through the workshop. |
| 97 | /// |
| 98 | /// This type is intentionally `Clone` and `Default` so it can be embedded |
| 99 | /// cheaply in [`ToolContext`](crate::tools::spec::ToolContext) without |
| 100 | /// requiring `Arc` wrappers. |
| 101 | #[derive(Debug, Clone, Default)] |
| 102 | pub struct LargeOutputRouter { |
| 103 | config: WorkshopConfig, |
| 104 | } |
| 105 | |
| 106 | impl LargeOutputRouter { |
| 107 | /// Construct a router from the resolved workshop config. |
| 108 | #[must_use] |
| 109 | pub fn new(config: WorkshopConfig) -> Self { |
| 110 | Self { config } |
| 111 | } |
| 112 | |
| 113 | /// Decide whether classic routing would synthesize `result`. |
| 114 | /// |
| 115 | /// This is used only by the rollback implementation. |
| 116 | #[must_use] |
| 117 | pub fn route(&self, tool_name: &str, result: &ToolResult, raw_bypass: bool) -> RouteDecision { |
| 118 | if raw_bypass || !result.success { |
| 119 | return RouteDecision::PassThrough; |
| 120 | } |
| 121 | let threshold = self.config.threshold_for(tool_name); |
| 122 | let estimated_tokens = estimate_tokens(&result.content); |
| 123 | if estimated_tokens > threshold { |
| 124 | RouteDecision::Synthesise { |
| 125 | estimated_tokens, |
| 126 | threshold, |
| 127 | } |
| 128 | } else { |
| 129 | RouteDecision::PassThrough |
| 130 | } |
| 131 | } |
| 132 | |
| 133 | #[must_use] |
| 134 | pub fn evidence_routing( |
| 135 | &self, |
| 136 | tool_name: &str, |
| 137 | result: &ToolResult, |
| 138 | _raw_bypass: bool, |
| 139 | ) -> (EvidenceRouting, usize, usize) { |
| 140 | let threshold = self.config.threshold_for(tool_name); |
| 141 | let estimated_tokens = estimate_tokens(&result.content); |
| 142 | // `raw=true` no longer bypasses the context bound. Exact bytes remain |
| 143 | // available through the artifact handle, so bypass is unnecessary. |
| 144 | let routing = EvidenceRouting::from_token_estimate(estimated_tokens, threshold); |
| 145 | (routing, estimated_tokens, threshold) |
| 146 | } |
| 147 | |
| 148 | /// Build the synthesis prompt sent to the V4-Flash workshop sub-agent. |
| 149 | /// |
| 150 | /// The prompt is intentionally terse — Flash is a fast model and we just |
| 151 | /// want a faithful summary, not deep reasoning. |
| 152 | /// |
| 153 | /// This is the building block for the live LLM synthesis call wired in |
| 154 | /// the follow-up (once the async Flash client is safe to call from the |
| 155 | /// registry layer). The method is public so callers outside this crate |
| 156 | /// can unit-test the prompt shape. |
| 157 | #[must_use] |
| 158 | #[allow(dead_code)] // used by future Flash synthesis call; keep for API stability |
| 159 | pub fn synthesis_prompt(tool_name: &str, raw_output: &str, estimated_tokens: usize) -> String { |
| 160 | format!( |
| 161 | "You are a synthesis assistant. The tool `{tool_name}` produced {estimated_tokens} tokens \ |
| 162 | of output that is too large to include directly in the parent context.\n\n\ |
| 163 | Summarise the output below into a concise, faithful synthesis of ≤ 800 words. \ |
| 164 | Preserve key facts, numbers, file paths, error messages, and any actionable \ |
| 165 | information. Do NOT add commentary or interpretation beyond what is in the source.\n\n\ |
| 166 | <raw_tool_output>\n{raw_output}\n</raw_tool_output>" |
| 167 | ) |
| 168 | } |
| 169 | |
| 170 | /// Wrap a synthesis result with a workshop provenance header and a hint |
| 171 | /// about the stored raw output. |
| 172 | #[must_use] |
| 173 | pub fn wrap_synthesis( |
| 174 | tool_name: &str, |
| 175 | synthesis: &str, |
| 176 | estimated_tokens: usize, |
| 177 | threshold: usize, |
| 178 | ) -> String { |
| 179 | format!( |
| 180 | "[workshop-synthesis: tool={tool_name}, raw_tokens≈{estimated_tokens}, \ |
| 181 | threshold={threshold}, raw_stored_in={WORKSHOP_LAST_TOOL_RESULT_VAR}]\n\n{synthesis}" |
| 182 | ) |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | // ── Workshop variable store ─────────────────────────────────────────────────── |
| 187 | |
| 188 | /// In-process store for workshop variables that persist across tool calls |
| 189 | /// within a session. The only variable exposed today is `last_tool_result` |
| 190 | /// which holds the most recent raw large-tool output for `promote_to_context`. |
| 191 | #[derive(Debug, Clone, Default, Serialize, Deserialize)] |
| 192 | pub struct WorkshopVariables { |
| 193 | /// Raw content of the most recent large tool output that was routed |
| 194 | /// through the workshop. Empty string when no routing has occurred. |
| 195 | #[serde(default)] |
| 196 | pub last_tool_result: String, |
| 197 | |
| 198 | /// Name of the tool that produced `last_tool_result`. |
| 199 | #[serde(default)] |
| 200 | pub last_tool_name: String, |
| 201 | } |
| 202 | |
| 203 | impl WorkshopVariables { |
| 204 | /// Store the raw output from a large-tool routing event. |
| 205 | pub fn store_raw(&mut self, tool_name: &str, raw: &str) { |
| 206 | self.last_tool_result = raw.to_string(); |
| 207 | self.last_tool_name = tool_name.to_string(); |
| 208 | } |
| 209 | |
| 210 | /// Retrieve and clear the stored raw output (consume semantics so the |
| 211 | /// variable is not accidentally promoted twice). |
| 212 | /// |
| 213 | /// Called by the `promote_to_context` tool (not yet wired in this PR). |
| 214 | #[must_use] |
| 215 | #[allow(dead_code)] // consumed by promote_to_context tool in follow-up |
| 216 | pub fn take_raw(&mut self) -> Option<(String, String)> { |
| 217 | if self.last_tool_result.is_empty() { |
| 218 | return None; |
| 219 | } |
| 220 | let content = std::mem::take(&mut self.last_tool_result); |
| 221 | let name = std::mem::take(&mut self.last_tool_name); |
| 222 | Some((name, content)) |
| 223 | } |
| 224 | } |
| 225 | |
| 226 | // ── Adaptive evidence routing (#4619) ───────────────────────────────────────── |
| 227 | |
| 228 | /// Routing policy for tool results: how much of the output stays inline in the |
| 229 | /// conversation vs. being stored as an external artifact. |
| 230 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 231 | #[serde(rename_all = "snake_case")] |
| 232 | pub enum EvidenceRouting { |
| 233 | /// Full result stays inline in the conversation context. |
| 234 | Inline, |
| 235 | /// A bounded observation (head/tail/summary) stays inline; the exact bytes |
| 236 | /// are stored as an artifact recoverable via handle. |
| 237 | Hybrid, |
| 238 | /// Only a handle/reference stays inline; the full result is artifact-only. |
| 239 | HandleOnly, |
| 240 | } |
| 241 | |
| 242 | impl EvidenceRouting { |
| 243 | /// Determine routing from estimated token count and threshold. |
| 244 | #[must_use] |
| 245 | pub fn from_token_estimate(estimated_tokens: usize, threshold: usize) -> Self { |
| 246 | if estimated_tokens <= threshold / 4 { |
| 247 | Self::Inline |
| 248 | } else if estimated_tokens <= threshold { |
| 249 | Self::Hybrid |
| 250 | } else { |
| 251 | Self::HandleOnly |
| 252 | } |
| 253 | } |
| 254 | } |
| 255 | |
| 256 | /// Immutable metadata for a stored evidence artifact (#4619). |
| 257 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 258 | pub struct EvidenceArtifact { |
| 259 | pub handle: String, |
| 260 | pub digest: String, |
| 261 | pub size_bytes: u64, |
| 262 | pub content_type: String, |
| 263 | pub tool_name: String, |
| 264 | pub call_id: String, |
| 265 | pub origin_session: String, |
| 266 | pub generation: u32, |
| 267 | pub redacted: bool, |
| 268 | pub encoding: String, |
| 269 | pub retention_state: EvidenceRetentionState, |
| 270 | pub created_at_unix_ms: u64, |
| 271 | pub retain_until_unix_ms: u64, |
| 272 | pub storage_path: PathBuf, |
| 273 | } |
| 274 | |
| 275 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 276 | #[serde(rename_all = "snake_case")] |
| 277 | pub enum EvidenceRetentionState { |
| 278 | Live, |
| 279 | Expired, |
| 280 | } |
| 281 | |
| 282 | pub const EVIDENCE_RETENTION_SECS: u64 = 7 * 24 * 60 * 60; |
| 283 | |
| 284 | #[must_use] |
| 285 | pub fn classic_output_routing_enabled() -> bool { |
| 286 | std::env::var("CODEWHALE_CLASSIC_OUTPUT_ROUTING") |
| 287 | .ok() |
| 288 | .is_some_and(|value| matches!(value.trim(), "1" | "true" | "yes" | "on")) |
| 289 | } |
| 290 | |
| 291 | #[must_use] |
| 292 | pub fn evidence_metadata_relative_path(handle: &str) -> PathBuf { |
| 293 | PathBuf::from(crate::artifacts::ARTIFACTS_DIR_NAME).join(format!("{handle}.evidence.json")) |
| 294 | } |
| 295 | |
| 296 | pub fn publish_evidence_metadata( |
| 297 | session_id: &str, |
| 298 | artifact: &EvidenceArtifact, |
| 299 | ) -> io::Result<PathBuf> { |
| 300 | let bytes = serde_json::to_vec_pretty(artifact) |
| 301 | .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; |
| 302 | crate::artifacts::write_session_relative_immutable( |
| 303 | session_id, |
| 304 | &evidence_metadata_relative_path(&artifact.handle), |
| 305 | &bytes, |
| 306 | ) |
| 307 | } |
| 308 | |
| 309 | pub fn read_evidence_metadata(session_id: &str, handle: &str) -> io::Result<EvidenceArtifact> { |
| 310 | let relative = evidence_metadata_relative_path(handle); |
| 311 | let path = crate::artifacts::session_artifact_absolute_path(session_id, &relative) |
| 312 | .ok_or_else(|| io::Error::new(io::ErrorKind::PermissionDenied, "invalid evidence owner"))?; |
| 313 | let raw = std::fs::read(path)?; |
| 314 | serde_json::from_slice(&raw).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) |
| 315 | } |
| 316 | |
| 317 | #[must_use] |
| 318 | pub fn unix_millis_now() -> u64 { |
| 319 | std::time::SystemTime::now() |
| 320 | .duration_since(std::time::UNIX_EPOCH) |
| 321 | .unwrap_or_default() |
| 322 | .as_millis() |
| 323 | .try_into() |
| 324 | .unwrap_or(u64::MAX) |
| 325 | } |
| 326 | |
| 327 | #[must_use] |
| 328 | pub fn evidence_is_expired(artifact: &EvidenceArtifact, now_ms: u64) -> bool { |
| 329 | artifact.retention_state == EvidenceRetentionState::Expired |
| 330 | || now_ms > artifact.retain_until_unix_ms |
| 331 | } |
| 332 | |
| 333 | // ── Unit tests ──────────────────────────────────────────────────────────────── |
| 334 | |
| 335 | #[cfg(test)] |
| 336 | mod tests { |
| 337 | use super::*; |
| 338 | |
| 339 | fn make_result(content: &str) -> ToolResult { |
| 340 | ToolResult::success(content.to_string()) |
| 341 | } |
| 342 | |
| 343 | #[test] |
| 344 | fn pass_through_below_threshold() { |
| 345 | let router = LargeOutputRouter::default(); |
| 346 | let small = "x".repeat(100); |
| 347 | let result = make_result(&small); |
| 348 | assert_eq!( |
| 349 | router.route("read_file", &result, false), |
| 350 | RouteDecision::PassThrough |
| 351 | ); |
| 352 | } |
| 353 | |
| 354 | #[test] |
| 355 | fn default_threshold_is_32k_tokens() { |
| 356 | assert_eq!(DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS, 32_768); |
| 357 | } |
| 358 | |
| 359 | #[test] |
| 360 | fn synthesise_above_threshold() { |
| 361 | let router = LargeOutputRouter::default(); |
| 362 | // DEFAULT threshold = 32768 tokens; 3 chars/token → 32768*3 = 98304 chars |
| 363 | let big = "a".repeat(100_000); |
| 364 | let result = make_result(&big); |
| 365 | assert!(matches!( |
| 366 | router.route("read_file", &result, false), |
| 367 | RouteDecision::Synthesise { .. } |
| 368 | )); |
| 369 | } |
| 370 | |
| 371 | #[test] |
| 372 | fn raw_bypass_skips_routing() { |
| 373 | let router = LargeOutputRouter::default(); |
| 374 | let big = "a".repeat(100_000); |
| 375 | let result = make_result(&big); |
| 376 | // raw=true → always pass through regardless of size |
| 377 | assert_eq!( |
| 378 | router.route("exec_shell", &result, true), |
| 379 | RouteDecision::PassThrough |
| 380 | ); |
| 381 | } |
| 382 | |
| 383 | #[test] |
| 384 | fn adaptive_evidence_cannot_bypass_context_bound_with_raw_flag() { |
| 385 | let router = LargeOutputRouter::default(); |
| 386 | let big = make_result(&"a".repeat(100_000)); |
| 387 | let (routing, _, _) = router.evidence_routing("exec_shell", &big, true); |
| 388 | assert_eq!(routing, EvidenceRouting::HandleOnly); |
| 389 | } |
| 390 | |
| 391 | #[test] |
| 392 | fn error_results_always_pass_through() { |
| 393 | let router = LargeOutputRouter::default(); |
| 394 | let big = "error: ".repeat(2_000); |
| 395 | let result = ToolResult::error(big); |
| 396 | assert_eq!( |
| 397 | router.route("exec_shell", &result, false), |
| 398 | RouteDecision::PassThrough |
| 399 | ); |
| 400 | } |
| 401 | |
| 402 | #[test] |
| 403 | fn per_tool_threshold_override() { |
| 404 | let mut per_tool = HashMap::new(); |
| 405 | per_tool.insert("grep_files".to_string(), 100); // very low |
| 406 | let config = WorkshopConfig { |
| 407 | large_output_threshold_tokens: Some(4096), |
| 408 | per_tool_thresholds: Some(per_tool), |
| 409 | }; |
| 410 | let router = LargeOutputRouter::new(config); |
| 411 | // 100 tokens * 3 = 300 chars → trigger with 400 chars |
| 412 | let medium = "b".repeat(400); |
| 413 | let result = make_result(&medium); |
| 414 | assert!(matches!( |
| 415 | router.route("grep_files", &result, false), |
| 416 | RouteDecision::Synthesise { .. } |
| 417 | )); |
| 418 | // Other tools still use the global threshold |
| 419 | assert_eq!( |
| 420 | router.route("read_file", &result, false), |
| 421 | RouteDecision::PassThrough |
| 422 | ); |
| 423 | } |
| 424 | |
| 425 | #[test] |
| 426 | fn estimate_tokens_conservative() { |
| 427 | // 9 chars → ceil(9/3) = 3 tokens |
| 428 | assert_eq!(estimate_tokens("123456789"), 3); |
| 429 | // 10 chars → ceil(10/3) = 4 tokens |
| 430 | assert_eq!(estimate_tokens("1234567890"), 4); |
| 431 | // Empty string |
| 432 | assert_eq!(estimate_tokens(""), 0); |
| 433 | } |
| 434 | |
| 435 | #[test] |
| 436 | fn workshop_variables_store_and_take() { |
| 437 | let mut vars = WorkshopVariables::default(); |
| 438 | assert!(vars.take_raw().is_none()); |
| 439 | |
| 440 | vars.store_raw("read_file", "raw content here"); |
| 441 | let taken = vars.take_raw().expect("should have content"); |
| 442 | assert_eq!(taken.0, "read_file"); |
| 443 | assert_eq!(taken.1, "raw content here"); |
| 444 | |
| 445 | // Second take is empty — consume semantics |
| 446 | assert!(vars.take_raw().is_none()); |
| 447 | } |
| 448 | |
| 449 | #[test] |
| 450 | fn wrap_synthesis_includes_provenance_header() { |
| 451 | let wrapped = LargeOutputRouter::wrap_synthesis("web_search", "key facts here", 5000, 4096); |
| 452 | assert!(wrapped.contains("workshop-synthesis")); |
| 453 | assert!(wrapped.contains("web_search")); |
| 454 | assert!(wrapped.contains("5000")); |
| 455 | assert!(wrapped.contains("key facts here")); |
| 456 | } |
| 457 | } |
| 458 |