| 1 | //! Context budgeting and prompt-shaping helpers for the engine. |
| 2 | //! |
| 3 | //! These functions are shared by the streaming turn loop, capacity flow, and |
| 4 | //! engine session maintenance code. Keeping them here prevents the top-level |
| 5 | //! engine module from accumulating unrelated context-policy details. |
| 6 | |
| 7 | use crate::compaction::estimate_tokens; |
| 8 | use crate::error_taxonomy::ErrorCategory; |
| 9 | use crate::models::{Message, SystemPrompt, context_window_for_model}; |
| 10 | use crate::tools::spec::ToolResult; |
| 11 | |
| 12 | /// Max output tokens requested for normal agent turns. Generous on purpose: |
| 13 | /// V4 thinking models can produce tens of thousands of reasoning tokens on |
| 14 | /// hard prompts before the visible reply, and DeepSeek V4 ships with a 1M |
| 15 | /// context window. v0.7.5 keeps this cap fixed instead of silently lowering |
| 16 | /// `max_tokens` near pressure; hard-cycle/preflight checks reserve this budget |
| 17 | /// plus safety headroom before sending the next request. |
| 18 | pub(super) const TURN_MAX_OUTPUT_TOKENS: u32 = 262_144; |
| 19 | |
| 20 | /// Safe max output tokens sent in the API request. This must be low enough to |
| 21 | /// work with providers that have smaller context limits than the model's native |
| 22 | /// window (e.g., self-hosted vLLM/SGLang with `--max-model-len 131072`). |
| 23 | /// DeepSeek's API will still produce as many tokens as needed for thinking; |
| 24 | /// this cap just prevents HTTP 400 from providers with tight limits. |
| 25 | const API_MAX_OUTPUT_TOKENS: u32 = 65_536; |
| 26 | |
| 27 | /// Compute the effective `max_tokens` to send in the API request for a given |
| 28 | /// model. Uses `API_MAX_OUTPUT_TOKENS` (64K) which fits within common provider |
| 29 | /// limits (128K+ total). For non-V4 models with smaller context windows, caps |
| 30 | /// at half the context window. |
| 31 | pub(super) fn effective_max_output_tokens(model: &str) -> u32 { |
| 32 | let window = context_window_for_model(model).unwrap_or(128_000); |
| 33 | if window >= 500_000 { |
| 34 | // V4-class models on large-context providers: use 64K which is safe |
| 35 | // for most deployments while still allowing substantial output. |
| 36 | API_MAX_OUTPUT_TOKENS |
| 37 | } else { |
| 38 | // Smaller models: cap at half the context window (leave room for input) |
| 39 | let capped = window / 2; |
| 40 | capped.min(API_MAX_OUTPUT_TOKENS) |
| 41 | } |
| 42 | } |
| 43 | /// Keep this many most recent messages when emergency trimming is required. |
| 44 | pub(super) const MIN_RECENT_MESSAGES_TO_KEEP: usize = 4; |
| 45 | /// Allow a few emergency recovery attempts before failing the turn. |
| 46 | pub(super) const MAX_CONTEXT_RECOVERY_ATTEMPTS: u8 = 2; |
| 47 | /// Reserve additional headroom to avoid hitting provider hard limits. |
| 48 | const CONTEXT_HEADROOM_TOKENS: usize = 1024; |
| 49 | /// Hard cap for any tool output inserted into model context. |
| 50 | const TOOL_RESULT_CONTEXT_HARD_LIMIT_CHARS: usize = 12_000; |
| 51 | /// Soft cap for known noisy tools inserted into model context. |
| 52 | const TOOL_RESULT_CONTEXT_SOFT_LIMIT_CHARS: usize = 2_000; |
| 53 | /// Snippet length kept when compacting tool output for model context. |
| 54 | const TOOL_RESULT_CONTEXT_SNIPPET_CHARS: usize = 900; |
| 55 | /// Hard cap for tool output inserted into a large-context model. |
| 56 | const LARGE_CONTEXT_TOOL_RESULT_HARD_LIMIT_CHARS: usize = 180_000; |
| 57 | /// Soft cap for known noisy tools inserted into a large-context model. |
| 58 | const LARGE_CONTEXT_TOOL_RESULT_SOFT_LIMIT_CHARS: usize = 60_000; |
| 59 | /// Snippet length kept when compacting large-context tool output. |
| 60 | const LARGE_CONTEXT_TOOL_RESULT_SNIPPET_CHARS: usize = 40_000; |
| 61 | /// Context window size at which tool output limits can be relaxed. |
| 62 | const LARGE_CONTEXT_WINDOW_TOKENS: u32 = 500_000; |
| 63 | /// Max chars to keep from metadata-provided output summaries. |
| 64 | const TOOL_RESULT_METADATA_SUMMARY_CHARS: usize = 320; |
| 65 | |
| 66 | pub(super) const COMPACTION_SUMMARY_MARKER: &str = "Conversation Summary (Auto-Generated)"; |
| 67 | |
| 68 | #[derive(Debug, Clone, Copy)] |
| 69 | struct ToolResultContextLimits { |
| 70 | hard_limit_chars: usize, |
| 71 | noisy_soft_limit_chars: usize, |
| 72 | snippet_chars: usize, |
| 73 | } |
| 74 | |
| 75 | pub(super) fn summarize_text(text: &str, limit: usize) -> String { |
| 76 | if text.chars().count() <= limit { |
| 77 | return text.to_string(); |
| 78 | } |
| 79 | let take = limit.saturating_sub(3); |
| 80 | let mut out: String = text.chars().take(take).collect(); |
| 81 | out.push_str("..."); |
| 82 | out |
| 83 | } |
| 84 | |
| 85 | fn summarize_text_head_tail(text: &str, limit: usize) -> String { |
| 86 | let total = text.chars().count(); |
| 87 | if total <= limit { |
| 88 | return text.to_string(); |
| 89 | } |
| 90 | if limit <= 20 { |
| 91 | return summarize_text(text, limit); |
| 92 | } |
| 93 | |
| 94 | let marker = "\n\n[... output truncated for context ...]\n\n"; |
| 95 | let marker_len = marker.chars().count(); |
| 96 | if limit <= marker_len + 20 { |
| 97 | return summarize_text(text, limit); |
| 98 | } |
| 99 | |
| 100 | let remaining = limit - marker_len; |
| 101 | let head_len = remaining.saturating_mul(2) / 3; |
| 102 | let tail_len = remaining.saturating_sub(head_len); |
| 103 | let head: String = text.chars().take(head_len).collect(); |
| 104 | let tail_vec: Vec<char> = text.chars().rev().take(tail_len).collect(); |
| 105 | let tail: String = tail_vec.into_iter().rev().collect(); |
| 106 | format!("{head}{marker}{tail}") |
| 107 | } |
| 108 | |
| 109 | fn tool_result_is_noisy(tool_name: &str) -> bool { |
| 110 | matches!( |
| 111 | tool_name, |
| 112 | "exec_shell" |
| 113 | | "exec_shell_wait" |
| 114 | | "exec_shell_interact" |
| 115 | | "multi_tool_use.parallel" |
| 116 | | "web_search" |
| 117 | ) |
| 118 | } |
| 119 | |
| 120 | fn tool_result_metadata_summary(metadata: Option<&serde_json::Value>) -> Option<String> { |
| 121 | let obj = metadata?.as_object()?; |
| 122 | for key in ["summary", "stdout_summary", "stderr_summary", "message"] { |
| 123 | if let Some(text) = obj.get(key).and_then(serde_json::Value::as_str) { |
| 124 | let trimmed = text.trim(); |
| 125 | if !trimmed.is_empty() { |
| 126 | return Some(summarize_text(trimmed, TOOL_RESULT_METADATA_SUMMARY_CHARS)); |
| 127 | } |
| 128 | } |
| 129 | } |
| 130 | None |
| 131 | } |
| 132 | |
| 133 | fn summarize_subagent_status(status: &serde_json::Value) -> String { |
| 134 | if let Some(raw) = status.as_str() { |
| 135 | return raw.to_string(); |
| 136 | } |
| 137 | if let Some(obj) = status.as_object() |
| 138 | && let Some((kind, value)) = obj.iter().next() |
| 139 | { |
| 140 | if let Some(reason) = value.as_str().filter(|s| !s.trim().is_empty()) { |
| 141 | return format!("{kind}({})", summarize_text(reason.trim(), 120)); |
| 142 | } |
| 143 | return kind.to_string(); |
| 144 | } |
| 145 | status.to_string() |
| 146 | } |
| 147 | |
| 148 | fn summarize_subagent_snapshot(snapshot: &serde_json::Value, index: usize) -> String { |
| 149 | let Some(obj) = snapshot.as_object() else { |
| 150 | return format!( |
| 151 | "- item {index}: {}", |
| 152 | summarize_text(&snapshot.to_string(), 240) |
| 153 | ); |
| 154 | }; |
| 155 | |
| 156 | let agent_id = obj |
| 157 | .get("agent_id") |
| 158 | .and_then(serde_json::Value::as_str) |
| 159 | .unwrap_or("unknown"); |
| 160 | let agent_type = obj |
| 161 | .get("agent_type") |
| 162 | .and_then(serde_json::Value::as_str) |
| 163 | .unwrap_or("agent"); |
| 164 | let status = obj |
| 165 | .get("status") |
| 166 | .map(summarize_subagent_status) |
| 167 | .unwrap_or_else(|| "unknown".to_string()); |
| 168 | let objective = obj |
| 169 | .get("assignment") |
| 170 | .and_then(|assignment| assignment.get("objective")) |
| 171 | .and_then(serde_json::Value::as_str) |
| 172 | .map(str::trim) |
| 173 | .filter(|s| !s.is_empty()) |
| 174 | .map(|s| summarize_text(s, 220)); |
| 175 | let result = obj |
| 176 | .get("result") |
| 177 | .and_then(serde_json::Value::as_str) |
| 178 | .map(str::trim) |
| 179 | .filter(|s| !s.is_empty()) |
| 180 | .map(|s| summarize_text(s, 1_600)); |
| 181 | let steps = obj.get("steps_taken").and_then(serde_json::Value::as_u64); |
| 182 | let duration_ms = obj.get("duration_ms").and_then(serde_json::Value::as_u64); |
| 183 | |
| 184 | let mut lines = vec![format!("- {agent_id} ({agent_type}) status={status}")]; |
| 185 | if let Some(objective) = objective { |
| 186 | lines.push(format!(" objective: {objective}")); |
| 187 | } |
| 188 | match result { |
| 189 | Some(result) => lines.push(format!(" result: {result}")), |
| 190 | None => lines.push(" result: not available yet".to_string()), |
| 191 | } |
| 192 | if steps.is_some() || duration_ms.is_some() { |
| 193 | let steps = steps |
| 194 | .map(|n| n.to_string()) |
| 195 | .unwrap_or_else(|| "?".to_string()); |
| 196 | let duration_ms = duration_ms |
| 197 | .map(|n| n.to_string()) |
| 198 | .unwrap_or_else(|| "?".to_string()); |
| 199 | lines.push(format!(" stats: steps={steps}, duration_ms={duration_ms}")); |
| 200 | } |
| 201 | lines.join("\n") |
| 202 | } |
| 203 | |
| 204 | fn compact_subagent_tool_result_for_context(tool_name: &str, raw: &str) -> Option<String> { |
| 205 | if !matches!(tool_name, "agent_result" | "agent_wait" | "wait") { |
| 206 | return None; |
| 207 | } |
| 208 | |
| 209 | let parsed: serde_json::Value = serde_json::from_str(raw).ok()?; |
| 210 | let snapshots: Vec<&serde_json::Value> = match &parsed { |
| 211 | serde_json::Value::Array(items) => items.iter().collect(), |
| 212 | serde_json::Value::Object(_) => vec![&parsed], |
| 213 | _ => return None, |
| 214 | }; |
| 215 | |
| 216 | let mut out = String::from("[sub-agent result summarized for parent context]\n"); |
| 217 | out.push_str("Use `agent_result` again only if you need the full raw payload.\n"); |
| 218 | for (idx, snapshot) in snapshots.iter().enumerate() { |
| 219 | if idx >= 8 { |
| 220 | out.push_str(&format!( |
| 221 | "- ... {} more sub-agent result(s) omitted from context summary\n", |
| 222 | snapshots.len().saturating_sub(idx) |
| 223 | )); |
| 224 | break; |
| 225 | } |
| 226 | out.push_str(&summarize_subagent_snapshot(snapshot, idx + 1)); |
| 227 | out.push('\n'); |
| 228 | } |
| 229 | Some(out.trim_end().to_string()) |
| 230 | } |
| 231 | |
| 232 | fn tool_result_context_limits_for_model(model: &str) -> ToolResultContextLimits { |
| 233 | let is_large_context = |
| 234 | context_window_for_model(model).is_some_and(|window| window >= LARGE_CONTEXT_WINDOW_TOKENS); |
| 235 | |
| 236 | if is_large_context { |
| 237 | ToolResultContextLimits { |
| 238 | hard_limit_chars: LARGE_CONTEXT_TOOL_RESULT_HARD_LIMIT_CHARS, |
| 239 | noisy_soft_limit_chars: LARGE_CONTEXT_TOOL_RESULT_SOFT_LIMIT_CHARS, |
| 240 | snippet_chars: LARGE_CONTEXT_TOOL_RESULT_SNIPPET_CHARS, |
| 241 | } |
| 242 | } else { |
| 243 | ToolResultContextLimits { |
| 244 | hard_limit_chars: TOOL_RESULT_CONTEXT_HARD_LIMIT_CHARS, |
| 245 | noisy_soft_limit_chars: TOOL_RESULT_CONTEXT_SOFT_LIMIT_CHARS, |
| 246 | snippet_chars: TOOL_RESULT_CONTEXT_SNIPPET_CHARS, |
| 247 | } |
| 248 | } |
| 249 | } |
| 250 | |
| 251 | pub(crate) fn compact_tool_result_for_context( |
| 252 | model: &str, |
| 253 | tool_name: &str, |
| 254 | output: &ToolResult, |
| 255 | ) -> String { |
| 256 | let raw = output.content.trim(); |
| 257 | if raw.is_empty() { |
| 258 | return String::new(); |
| 259 | } |
| 260 | |
| 261 | if let Some(summary) = compact_subagent_tool_result_for_context(tool_name, raw) { |
| 262 | return summary; |
| 263 | } |
| 264 | |
| 265 | let limits = tool_result_context_limits_for_model(model); |
| 266 | let raw_chars = raw.chars().count(); |
| 267 | let should_compact = raw_chars > limits.hard_limit_chars |
| 268 | || (tool_result_is_noisy(tool_name) && raw_chars > limits.noisy_soft_limit_chars); |
| 269 | if !should_compact { |
| 270 | return raw.to_string(); |
| 271 | } |
| 272 | |
| 273 | let snippet = summarize_text_head_tail(raw, limits.snippet_chars); |
| 274 | let omitted = raw_chars.saturating_sub(snippet.chars().count()); |
| 275 | let summary = tool_result_metadata_summary(output.metadata.as_ref()); |
| 276 | |
| 277 | if let Some(summary) = summary { |
| 278 | format!( |
| 279 | "[{tool_name} output compacted to protect context]\nSummary: {summary}\nSnippet: {snippet}\n(Original: {raw_chars} chars, omitted: {omitted} chars.)" |
| 280 | ) |
| 281 | } else { |
| 282 | format!( |
| 283 | "[{tool_name} output compacted to protect context]\nSnippet: {snippet}\n(Original: {raw_chars} chars, omitted: {omitted} chars.)" |
| 284 | ) |
| 285 | } |
| 286 | } |
| 287 | |
| 288 | pub(super) fn extract_compaction_summary_prompt( |
| 289 | prompt: Option<SystemPrompt>, |
| 290 | ) -> Option<SystemPrompt> { |
| 291 | match prompt { |
| 292 | Some(SystemPrompt::Blocks(blocks)) => { |
| 293 | let summary_blocks: Vec<_> = blocks |
| 294 | .into_iter() |
| 295 | .filter(|block| block.text.contains(COMPACTION_SUMMARY_MARKER)) |
| 296 | .collect(); |
| 297 | if summary_blocks.is_empty() { |
| 298 | None |
| 299 | } else { |
| 300 | Some(SystemPrompt::Blocks(summary_blocks)) |
| 301 | } |
| 302 | } |
| 303 | Some(SystemPrompt::Text(text)) => { |
| 304 | if text.contains(COMPACTION_SUMMARY_MARKER) { |
| 305 | Some(SystemPrompt::Text(text)) |
| 306 | } else { |
| 307 | None |
| 308 | } |
| 309 | } |
| 310 | None => None, |
| 311 | } |
| 312 | } |
| 313 | |
| 314 | fn estimate_text_tokens_conservative(text: &str) -> usize { |
| 315 | text.chars().count().div_ceil(3) |
| 316 | } |
| 317 | |
| 318 | fn estimate_system_tokens_conservative(system: Option<&SystemPrompt>) -> usize { |
| 319 | match system { |
| 320 | Some(SystemPrompt::Text(text)) => estimate_text_tokens_conservative(text), |
| 321 | Some(SystemPrompt::Blocks(blocks)) => blocks |
| 322 | .iter() |
| 323 | .map(|block| estimate_text_tokens_conservative(&block.text)) |
| 324 | .sum(), |
| 325 | None => 0, |
| 326 | } |
| 327 | } |
| 328 | |
| 329 | pub(super) fn estimate_input_tokens_conservative( |
| 330 | messages: &[Message], |
| 331 | system: Option<&SystemPrompt>, |
| 332 | ) -> usize { |
| 333 | let message_tokens = estimate_tokens(messages).saturating_mul(3).div_ceil(2); |
| 334 | let system_tokens = estimate_system_tokens_conservative(system); |
| 335 | let framing_overhead = messages.len().saturating_mul(12).saturating_add(48); |
| 336 | message_tokens |
| 337 | .saturating_add(system_tokens) |
| 338 | .saturating_add(framing_overhead) |
| 339 | } |
| 340 | |
| 341 | pub(super) fn context_input_budget(model: &str, requested_output_tokens: u32) -> Option<usize> { |
| 342 | let window = usize::try_from(context_window_for_model(model)?).ok()?; |
| 343 | let output = usize::try_from(requested_output_tokens).ok()?; |
| 344 | window |
| 345 | .checked_sub(output) |
| 346 | .and_then(|v| v.checked_sub(CONTEXT_HEADROOM_TOKENS)) |
| 347 | } |
| 348 | |
| 349 | pub(super) fn turn_response_headroom_tokens() -> u64 { |
| 350 | u64::from(TURN_MAX_OUTPUT_TOKENS).saturating_add(CONTEXT_HEADROOM_TOKENS as u64) |
| 351 | } |
| 352 | |
| 353 | pub(super) fn is_context_length_error_message(message: &str) -> bool { |
| 354 | crate::error_taxonomy::classify_error_message(message) == ErrorCategory::InvalidInput |
| 355 | } |
| 356 |