| 1 | //! Adaptive evidence routing for tool results (#4619) — explicit opt-in. |
| 2 | //! |
| 3 | //! Classic bounded spillover is the default: `tools/truncate.rs` keeps results |
| 4 | //! at or under its byte threshold fully inline and gives larger ones a |
| 5 | //! head/tail preview plus a session artifact. Set |
| 6 | //! `CODEWHALE_ADAPTIVE_OUTPUT_ROUTING` to enable the adaptive lane, which |
| 7 | //! classifies results as inline, hybrid, or handle-only by estimated tokens |
| 8 | //! and publishes non-inline results exactly once under their origin session |
| 9 | //! with immutable evidence metadata for bounded retrieval. |
| 10 | |
| 11 | use std::collections::HashMap; |
| 12 | use std::io; |
| 13 | use std::path::PathBuf; |
| 14 | use std::sync::{Mutex, OnceLock}; |
| 15 | |
| 16 | use serde::{Deserialize, Serialize}; |
| 17 | |
| 18 | use crate::tools::spec::ToolResult; |
| 19 | |
| 20 | // ── Constants ────────────────────────────────────────────────────────────────── |
| 21 | |
| 22 | /// Default token threshold separating hybrid from handle-only evidence. |
| 23 | /// |
| 24 | /// 32K tokens (≈96 KiB of text at the 3 chars/token estimate) keeps ordinary |
| 25 | /// tool results — file reads, test runs, build logs up to a few thousand |
| 26 | /// lines — fully inline. Only genuinely large outputs spill to evidence |
| 27 | /// artifacts, where the model-facing preview names the artifact path and how |
| 28 | /// to recover the omitted range. |
| 29 | pub const DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS: usize = 32_768; |
| 30 | |
| 31 | /// Approximate characters-per-token ratio used for the heuristic estimate. |
| 32 | /// We intentionally choose a conservative value (3 chars/token) so we err |
| 33 | /// on the side of routing rather than dumping raw data into the parent. |
| 34 | const CHARS_PER_TOKEN_ESTIMATE: usize = 3; |
| 35 | |
| 36 | static ACTIVE_WORKSHOP: OnceLock<Mutex<WorkshopConfig>> = OnceLock::new(); |
| 37 | |
| 38 | #[cfg(test)] |
| 39 | static ACTIVE_WORKSHOP_TEST_SERIAL: OnceLock<Mutex<()>> = OnceLock::new(); |
| 40 | |
| 41 | #[cfg(test)] |
| 42 | std::thread_local! { |
| 43 | static ACTIVE_WORKSHOP_TEST_SERIAL_HELD: std::cell::Cell<bool> = const { |
| 44 | std::cell::Cell::new(false) |
| 45 | }; |
| 46 | } |
| 47 | |
| 48 | /// Holds every test-side workshop activation behind one process-wide gate. |
| 49 | /// The thread-local marker lets the owning current-thread test call |
| 50 | /// `install_active` without trying to acquire its own non-reentrant lock. |
| 51 | #[cfg(test)] |
| 52 | pub(crate) struct ActiveWorkshopTestGuard { |
| 53 | _serial: std::sync::MutexGuard<'static, ()>, |
| 54 | } |
| 55 | |
| 56 | #[cfg(test)] |
| 57 | impl Drop for ActiveWorkshopTestGuard { |
| 58 | fn drop(&mut self) { |
| 59 | ACTIVE_WORKSHOP_TEST_SERIAL_HELD.with(|held| held.set(false)); |
| 60 | } |
| 61 | } |
| 62 | |
| 63 | #[cfg(test)] |
| 64 | pub(crate) fn active_workshop_test_guard() -> ActiveWorkshopTestGuard { |
| 65 | assert!( |
| 66 | !ACTIVE_WORKSHOP_TEST_SERIAL_HELD.with(std::cell::Cell::get), |
| 67 | "active workshop test guard is not reentrant" |
| 68 | ); |
| 69 | let serial = ACTIVE_WORKSHOP_TEST_SERIAL |
| 70 | .get_or_init(|| Mutex::new(())) |
| 71 | .lock() |
| 72 | .unwrap_or_else(std::sync::PoisonError::into_inner); |
| 73 | ACTIVE_WORKSHOP_TEST_SERIAL_HELD.with(|held| held.set(true)); |
| 74 | ActiveWorkshopTestGuard { _serial: serial } |
| 75 | } |
| 76 | |
| 77 | fn active_workshop_slot() -> &'static Mutex<WorkshopConfig> { |
| 78 | ACTIVE_WORKSHOP.get_or_init(|| Mutex::new(WorkshopConfig::default())) |
| 79 | } |
| 80 | |
| 81 | // ── Configuration ───────────────────────────────────────────────────────────── |
| 82 | |
| 83 | /// Existing `[workshop]` threshold configuration, retained for compatibility. |
| 84 | #[derive(Debug, Clone, Deserialize, Default)] |
| 85 | pub struct WorkshopConfig { |
| 86 | /// Token threshold above which results become handle-only evidence. |
| 87 | #[serde(default)] |
| 88 | pub large_output_threshold_tokens: Option<usize>, |
| 89 | |
| 90 | /// Per-tool threshold overrides (tool name → token limit). A tool whose |
| 91 | /// name appears here uses this limit instead of |
| 92 | /// `large_output_threshold_tokens`. |
| 93 | #[serde(default)] |
| 94 | pub per_tool_thresholds: Option<HashMap<String, usize>>, |
| 95 | |
| 96 | /// Optional model-visible byte budget for a single `read` / `read_file` |
| 97 | /// result. Absent keeps the compile-time default (#5367). |
| 98 | #[serde(default)] |
| 99 | pub read_result_max_bytes: Option<usize>, |
| 100 | |
| 101 | /// Optional model-visible byte budget for a generic tool result after |
| 102 | /// spillover. Absent keeps the compile-time default (#5367). |
| 103 | #[serde(default)] |
| 104 | pub tool_result_max_bytes: Option<usize>, |
| 105 | } |
| 106 | |
| 107 | impl WorkshopConfig { |
| 108 | /// Install the process-wide workshop budgets used by read/tool compactors. |
| 109 | /// |
| 110 | /// The returned immutable receipt is the snapshot written while the |
| 111 | /// singleton lock was held. Callers that need evidence of their own |
| 112 | /// activation can inspect it without racing a later process-wide update. |
| 113 | pub fn install_active(config: Option<&Self>) -> Self { |
| 114 | #[cfg(test)] |
| 115 | let _test_serial = if ACTIVE_WORKSHOP_TEST_SERIAL_HELD.with(std::cell::Cell::get) { |
| 116 | None |
| 117 | } else { |
| 118 | Some( |
| 119 | ACTIVE_WORKSHOP_TEST_SERIAL |
| 120 | .get_or_init(|| Mutex::new(())) |
| 121 | .lock() |
| 122 | .unwrap_or_else(std::sync::PoisonError::into_inner), |
| 123 | ) |
| 124 | }; |
| 125 | let snapshot = config.cloned().unwrap_or_default(); |
| 126 | let mut slot = active_workshop_slot() |
| 127 | .lock() |
| 128 | .unwrap_or_else(std::sync::PoisonError::into_inner); |
| 129 | *slot = snapshot; |
| 130 | slot.clone() |
| 131 | } |
| 132 | |
| 133 | /// Optional model-visible read budget, when the user opted in (#5367). |
| 134 | #[must_use] |
| 135 | pub fn active_read_result_max_bytes() -> Option<usize> { |
| 136 | active_workshop_slot() |
| 137 | .lock() |
| 138 | .ok() |
| 139 | .and_then(|cfg| cfg.read_result_max_bytes.filter(|n| *n > 0)) |
| 140 | } |
| 141 | |
| 142 | /// Optional model-visible tool-result budget, when the user opted in (#5367). |
| 143 | #[must_use] |
| 144 | pub fn active_tool_result_max_bytes() -> Option<usize> { |
| 145 | active_workshop_slot() |
| 146 | .lock() |
| 147 | .ok() |
| 148 | .and_then(|cfg| cfg.tool_result_max_bytes.filter(|n| *n > 0)) |
| 149 | } |
| 150 | |
| 151 | /// Resolve the effective threshold for the given tool name. |
| 152 | #[must_use] |
| 153 | pub fn threshold_for(&self, tool_name: &str) -> usize { |
| 154 | if let Some(per_tool) = self.per_tool_thresholds.as_ref() |
| 155 | && let Some(&limit) = per_tool.get(tool_name) |
| 156 | { |
| 157 | return limit; |
| 158 | } |
| 159 | self.large_output_threshold_tokens |
| 160 | .unwrap_or(DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS) |
| 161 | } |
| 162 | } |
| 163 | |
| 164 | // ── Token estimation ────────────────────────────────────────────────────────── |
| 165 | |
| 166 | /// Estimate the number of tokens in `text` using a character-count heuristic. |
| 167 | /// |
| 168 | /// This avoids a real tokeniser dependency; the estimate is deliberately |
| 169 | /// conservative (under-counts tokens) so we route aggressively rather than |
| 170 | /// letting a 5K-token blob slip through. |
| 171 | #[must_use] |
| 172 | pub fn estimate_tokens(text: &str) -> usize { |
| 173 | let chars = text.chars().count(); |
| 174 | // Round up: partial last token still costs a token. |
| 175 | chars.div_ceil(CHARS_PER_TOKEN_ESTIMATE) |
| 176 | } |
| 177 | |
| 178 | // ── Router ──────────────────────────────────────────────────────────────────── |
| 179 | |
| 180 | /// Classifies tool results for adaptive evidence routing. |
| 181 | /// |
| 182 | /// This type is intentionally `Clone` and `Default` so it can be embedded |
| 183 | /// cheaply in [`ToolContext`](crate::tools::spec::ToolContext) without |
| 184 | /// requiring `Arc` wrappers. |
| 185 | #[derive(Debug, Clone, Default)] |
| 186 | pub struct LargeOutputRouter { |
| 187 | config: WorkshopConfig, |
| 188 | } |
| 189 | |
| 190 | impl LargeOutputRouter { |
| 191 | /// Construct a router from the resolved workshop config. |
| 192 | #[must_use] |
| 193 | pub fn new(config: WorkshopConfig) -> Self { |
| 194 | Self { config } |
| 195 | } |
| 196 | |
| 197 | #[must_use] |
| 198 | pub fn evidence_routing( |
| 199 | &self, |
| 200 | tool_name: &str, |
| 201 | result: &ToolResult, |
| 202 | _raw_bypass: bool, |
| 203 | ) -> (EvidenceRouting, usize, usize) { |
| 204 | let threshold = self.config.threshold_for(tool_name); |
| 205 | let estimated_tokens = estimate_tokens(&result.content); |
| 206 | // `raw=true` no longer bypasses the context bound. Exact bytes remain |
| 207 | // available through the artifact handle, so bypass is unnecessary. |
| 208 | let routing = EvidenceRouting::from_token_estimate(estimated_tokens, threshold); |
| 209 | (routing, estimated_tokens, threshold) |
| 210 | } |
| 211 | } |
| 212 | |
| 213 | // ── Adaptive evidence routing (#4619) ───────────────────────────────────────── |
| 214 | |
| 215 | /// Routing policy for tool results: how much of the output stays inline in the |
| 216 | /// conversation vs. being stored as an external artifact. |
| 217 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 218 | #[serde(rename_all = "snake_case")] |
| 219 | pub enum EvidenceRouting { |
| 220 | /// Full result stays inline in the conversation context. |
| 221 | Inline, |
| 222 | /// A bounded observation (head/tail/summary) stays inline; the exact bytes |
| 223 | /// are stored as an artifact recoverable via handle. |
| 224 | Hybrid, |
| 225 | /// Only a handle/reference stays inline; the full result is artifact-only. |
| 226 | HandleOnly, |
| 227 | } |
| 228 | |
| 229 | impl EvidenceRouting { |
| 230 | /// Determine routing from estimated token count and threshold. |
| 231 | #[must_use] |
| 232 | pub fn from_token_estimate(estimated_tokens: usize, threshold: usize) -> Self { |
| 233 | if estimated_tokens <= threshold / 4 { |
| 234 | Self::Inline |
| 235 | } else if estimated_tokens <= threshold { |
| 236 | Self::Hybrid |
| 237 | } else { |
| 238 | Self::HandleOnly |
| 239 | } |
| 240 | } |
| 241 | } |
| 242 | |
| 243 | /// Immutable metadata for a stored evidence artifact (#4619). |
| 244 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 245 | pub struct EvidenceArtifact { |
| 246 | pub handle: String, |
| 247 | pub digest: String, |
| 248 | pub size_bytes: u64, |
| 249 | pub content_type: String, |
| 250 | pub tool_name: String, |
| 251 | pub call_id: String, |
| 252 | pub origin_session: String, |
| 253 | pub generation: u32, |
| 254 | pub redacted: bool, |
| 255 | pub encoding: String, |
| 256 | pub retention_state: EvidenceRetentionState, |
| 257 | pub created_at_unix_ms: u64, |
| 258 | pub retain_until_unix_ms: u64, |
| 259 | pub storage_path: PathBuf, |
| 260 | } |
| 261 | |
| 262 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 263 | #[serde(rename_all = "snake_case")] |
| 264 | pub enum EvidenceRetentionState { |
| 265 | Live, |
| 266 | Expired, |
| 267 | } |
| 268 | |
| 269 | pub const EVIDENCE_RETENTION_SECS: u64 = 7 * 24 * 60 * 60; |
| 270 | |
| 271 | /// Whether adaptive evidence routing (#4619) is enabled for this process. |
| 272 | /// |
| 273 | /// Off by default — classic bounded spillover owns large results. Set |
| 274 | /// `CODEWHALE_ADAPTIVE_OUTPUT_ROUTING` to opt in. The retired rollback |
| 275 | /// variable is still honored in the negative, so |
| 276 | /// `CODEWHALE_CLASSIC_OUTPUT_ROUTING=0` also selects the adaptive lane. |
| 277 | #[must_use] |
| 278 | pub fn adaptive_output_routing_enabled() -> bool { |
| 279 | if let Ok(value) = std::env::var("CODEWHALE_ADAPTIVE_OUTPUT_ROUTING") { |
| 280 | return matches!(value.trim(), "1" | "true" | "yes" | "on"); |
| 281 | } |
| 282 | matches!( |
| 283 | std::env::var("CODEWHALE_CLASSIC_OUTPUT_ROUTING") |
| 284 | .ok() |
| 285 | .as_deref() |
| 286 | .map(str::trim), |
| 287 | Some("0" | "false" | "no" | "off") |
| 288 | ) |
| 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 file = crate::artifacts::open_session_relative(session_id, &relative, false)?; |
| 312 | read_evidence_metadata_file(&file) |
| 313 | } |
| 314 | |
| 315 | /// Bounded, no-follow read shared by publication/replay and authenticated HTTP |
| 316 | /// retrieval. The caller chooses the existing session-root authority. |
| 317 | pub(crate) fn read_evidence_metadata_file( |
| 318 | file: &crate::fleet::files::WorkspaceFile, |
| 319 | ) -> io::Result<EvidenceArtifact> { |
| 320 | use std::io::Read; |
| 321 | const MAX_MANIFEST_BYTES: u64 = 64 * 1024; |
| 322 | let mut raw = Vec::new(); |
| 323 | file.open_file()? |
| 324 | .take(MAX_MANIFEST_BYTES + 1) |
| 325 | .read_to_end(&mut raw)?; |
| 326 | if raw.len() as u64 > MAX_MANIFEST_BYTES { |
| 327 | return Err(io::Error::new( |
| 328 | io::ErrorKind::InvalidData, |
| 329 | "evidence metadata exceeds limit", |
| 330 | )); |
| 331 | } |
| 332 | serde_json::from_slice(&raw).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) |
| 333 | } |
| 334 | |
| 335 | #[must_use] |
| 336 | pub fn unix_millis_now() -> u64 { |
| 337 | std::time::SystemTime::now() |
| 338 | .duration_since(std::time::UNIX_EPOCH) |
| 339 | .unwrap_or_default() |
| 340 | .as_millis() |
| 341 | .try_into() |
| 342 | .unwrap_or(u64::MAX) |
| 343 | } |
| 344 | |
| 345 | #[must_use] |
| 346 | pub fn evidence_is_expired(artifact: &EvidenceArtifact, now_ms: u64) -> bool { |
| 347 | artifact.retention_state == EvidenceRetentionState::Expired |
| 348 | || now_ms > artifact.retain_until_unix_ms |
| 349 | } |
| 350 | |
| 351 | // ── Unit tests ──────────────────────────────────────────────────────────────── |
| 352 | |
| 353 | #[cfg(test)] |
| 354 | mod tests { |
| 355 | use super::*; |
| 356 | |
| 357 | fn make_result(content: &str) -> ToolResult { |
| 358 | ToolResult::success(content.to_string()) |
| 359 | } |
| 360 | |
| 361 | #[test] |
| 362 | fn default_threshold_is_32k_tokens() { |
| 363 | assert_eq!(DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS, 32_768); |
| 364 | } |
| 365 | |
| 366 | #[test] |
| 367 | fn adaptive_evidence_cannot_bypass_context_bound_with_raw_flag() { |
| 368 | let router = LargeOutputRouter::default(); |
| 369 | let big = make_result(&"a".repeat(100_000)); |
| 370 | let (routing, _, _) = router.evidence_routing("exec_shell", &big, true); |
| 371 | assert_eq!(routing, EvidenceRouting::HandleOnly); |
| 372 | } |
| 373 | |
| 374 | #[test] |
| 375 | fn per_tool_threshold_override() { |
| 376 | let mut per_tool = HashMap::new(); |
| 377 | per_tool.insert("grep_files".to_string(), 100); // very low |
| 378 | let config = WorkshopConfig { |
| 379 | large_output_threshold_tokens: Some(4096), |
| 380 | per_tool_thresholds: Some(per_tool), |
| 381 | read_result_max_bytes: None, |
| 382 | tool_result_max_bytes: None, |
| 383 | }; |
| 384 | assert_eq!(config.threshold_for("grep_files"), 100); |
| 385 | assert_eq!(config.threshold_for("read_file"), 4096); |
| 386 | let default_config = WorkshopConfig::default(); |
| 387 | assert_eq!( |
| 388 | default_config.threshold_for("read_file"), |
| 389 | DEFAULT_LARGE_OUTPUT_THRESHOLD_TOKENS |
| 390 | ); |
| 391 | } |
| 392 | |
| 393 | #[test] |
| 394 | fn workshop_byte_budgets_raise_floor_only() { |
| 395 | let _guard = active_workshop_test_guard(); |
| 396 | let installed = WorkshopConfig::install_active(Some(&WorkshopConfig { |
| 397 | large_output_threshold_tokens: None, |
| 398 | per_tool_thresholds: None, |
| 399 | read_result_max_bytes: Some(102_400), |
| 400 | tool_result_max_bytes: Some(80_000), |
| 401 | })); |
| 402 | assert_eq!(installed.read_result_max_bytes, Some(102_400)); |
| 403 | assert_eq!(installed.tool_result_max_bytes, Some(80_000)); |
| 404 | assert_eq!( |
| 405 | WorkshopConfig::active_read_result_max_bytes(), |
| 406 | Some(102_400) |
| 407 | ); |
| 408 | assert_eq!(WorkshopConfig::active_tool_result_max_bytes(), Some(80_000)); |
| 409 | let cleared = WorkshopConfig::install_active(None); |
| 410 | assert_eq!(cleared.read_result_max_bytes, None); |
| 411 | assert_eq!(cleared.tool_result_max_bytes, None); |
| 412 | assert_eq!(WorkshopConfig::active_read_result_max_bytes(), None); |
| 413 | assert_eq!(WorkshopConfig::active_tool_result_max_bytes(), None); |
| 414 | } |
| 415 | |
| 416 | #[test] |
| 417 | fn estimate_tokens_conservative() { |
| 418 | // 9 chars → ceil(9/3) = 3 tokens |
| 419 | assert_eq!(estimate_tokens("123456789"), 3); |
| 420 | // 10 chars → ceil(10/3) = 4 tokens |
| 421 | assert_eq!(estimate_tokens("1234567890"), 4); |
| 422 | // Empty string |
| 423 | assert_eq!(estimate_tokens(""), 0); |
| 424 | } |
| 425 | } |
| 426 |