| 1 | //! Capacity-controller checkpoints and interventions for the engine loop. |
| 2 | //! |
| 3 | //! Extracted from `core/engine.rs` for issue #74. The main turn loop still |
| 4 | //! decides when checkpoints run; this module owns the guardrail policy side |
| 5 | //! effects, replay verification, canonical-state persistence, and event |
| 6 | //! emission helpers. |
| 7 | |
| 8 | use super::*; |
| 9 | |
| 10 | use crate::models::context_window_for_model; |
| 11 | |
| 12 | impl Engine { |
| 13 | pub(super) async fn run_capacity_pre_request_checkpoint( |
| 14 | &mut self, |
| 15 | turn: &TurnContext, |
| 16 | client: Option<&DeepSeekClient>, |
| 17 | mode: AppMode, |
| 18 | ) -> bool { |
| 19 | let snapshot = self |
| 20 | .capacity_controller |
| 21 | .observe_pre_turn(self.capacity_observation(turn)); |
| 22 | let decision = self |
| 23 | .capacity_controller |
| 24 | .decide(self.turn_counter, snapshot.as_ref()); |
| 25 | self.emit_capacity_decision(turn, snapshot.as_ref(), &decision) |
| 26 | .await; |
| 27 | |
| 28 | if decision.action != GuardrailAction::TargetedContextRefresh { |
| 29 | return false; |
| 30 | } |
| 31 | |
| 32 | self.apply_targeted_context_refresh(turn, client, mode, snapshot.as_ref()) |
| 33 | .await |
| 34 | } |
| 35 | |
| 36 | #[allow(clippy::too_many_arguments)] |
| 37 | pub(super) async fn run_capacity_post_tool_checkpoint( |
| 38 | &mut self, |
| 39 | turn: &TurnContext, |
| 40 | mode: AppMode, |
| 41 | tool_registry: Option<&crate::tools::ToolRegistry>, |
| 42 | tool_exec_lock: Arc<RwLock<()>>, |
| 43 | mcp_pool: Option<Arc<AsyncMutex<McpPool>>>, |
| 44 | _step_error_count: usize, |
| 45 | _consecutive_tool_error_steps: u32, |
| 46 | ) -> bool { |
| 47 | let snapshot = self |
| 48 | .capacity_controller |
| 49 | .observe_post_tool(self.capacity_observation(turn)); |
| 50 | let decision = self |
| 51 | .capacity_controller |
| 52 | .decide(self.turn_counter, snapshot.as_ref()); |
| 53 | self.emit_capacity_decision(turn, snapshot.as_ref(), &decision) |
| 54 | .await; |
| 55 | |
| 56 | match decision.action { |
| 57 | GuardrailAction::VerifyWithToolReplay => { |
| 58 | let _ = self |
| 59 | .apply_verify_with_tool_replay( |
| 60 | turn, |
| 61 | mode, |
| 62 | snapshot.as_ref(), |
| 63 | tool_registry, |
| 64 | tool_exec_lock, |
| 65 | mcp_pool, |
| 66 | ) |
| 67 | .await; |
| 68 | false |
| 69 | } |
| 70 | GuardrailAction::VerifyAndReplan => { |
| 71 | self.apply_verify_and_replan(turn, mode, snapshot.as_ref(), "high_risk_post_tool") |
| 72 | .await |
| 73 | } |
| 74 | GuardrailAction::NoIntervention | GuardrailAction::TargetedContextRefresh => false, |
| 75 | } |
| 76 | } |
| 77 | |
| 78 | pub(super) async fn run_capacity_error_escalation_checkpoint( |
| 79 | &mut self, |
| 80 | turn: &TurnContext, |
| 81 | mode: AppMode, |
| 82 | step_error_count: usize, |
| 83 | consecutive_tool_error_steps: u32, |
| 84 | error_categories: &[ErrorCategory], |
| 85 | ) -> bool { |
| 86 | if step_error_count == 0 && consecutive_tool_error_steps < 2 { |
| 87 | return false; |
| 88 | } |
| 89 | |
| 90 | // Categorize this step's failures by typed `ErrorCategory` rather than |
| 91 | // substring-matching error strings. Context overflow always escalates; |
| 92 | // network / rate-limit / timeout are transient and skip escalation; |
| 93 | // anything else only escalates with consecutive consecutive failures. |
| 94 | let has_context_overflow = error_categories.contains(&ErrorCategory::InvalidInput); |
| 95 | let only_transient = !error_categories.is_empty() |
| 96 | && error_categories.iter().all(|c| { |
| 97 | matches!( |
| 98 | c, |
| 99 | ErrorCategory::Network | ErrorCategory::RateLimit | ErrorCategory::Timeout |
| 100 | ) |
| 101 | }); |
| 102 | if only_transient && !has_context_overflow { |
| 103 | return false; |
| 104 | } |
| 105 | if !has_context_overflow && consecutive_tool_error_steps < 2 { |
| 106 | return false; |
| 107 | } |
| 108 | |
| 109 | let snapshot = self |
| 110 | .capacity_controller |
| 111 | .last_snapshot() |
| 112 | .cloned() |
| 113 | .or_else(|| { |
| 114 | self.capacity_controller |
| 115 | .observe_pre_turn(self.capacity_observation(turn)) |
| 116 | }); |
| 117 | let Some(snapshot) = snapshot else { |
| 118 | return false; |
| 119 | }; |
| 120 | |
| 121 | let repeated_failures = step_error_count >= 2 || consecutive_tool_error_steps >= 2; |
| 122 | let mut forced = snapshot.clone(); |
| 123 | if repeated_failures && !(snapshot.risk_band == RiskBand::High && snapshot.severe) { |
| 124 | forced.risk_band = RiskBand::High; |
| 125 | forced.severe = true; |
| 126 | } |
| 127 | |
| 128 | let decision = self |
| 129 | .capacity_controller |
| 130 | .decide(self.turn_counter, Some(&forced)); |
| 131 | self.emit_capacity_decision(turn, Some(&forced), &decision) |
| 132 | .await; |
| 133 | |
| 134 | if decision.action != GuardrailAction::VerifyAndReplan { |
| 135 | return false; |
| 136 | } |
| 137 | |
| 138 | let category_labels: Vec<String> = error_categories.iter().map(|c| c.to_string()).collect(); |
| 139 | self.apply_verify_and_replan( |
| 140 | turn, |
| 141 | mode, |
| 142 | Some(&forced), |
| 143 | &format!( |
| 144 | "error_escalation: step_errors={}, consecutive_steps={}, categories={}", |
| 145 | step_error_count, |
| 146 | consecutive_tool_error_steps, |
| 147 | category_labels.join(",") |
| 148 | ), |
| 149 | ) |
| 150 | .await |
| 151 | } |
| 152 | |
| 153 | pub(super) fn capacity_observation(&self, turn: &TurnContext) -> CapacityObservationInput { |
| 154 | let message_window = self.config.capacity.profile_window.max(8) * 3; |
| 155 | let action_count_this_turn = usize::try_from(turn.step) |
| 156 | .unwrap_or(usize::MAX) |
| 157 | .saturating_add(turn.tool_calls.len()) |
| 158 | .saturating_add(1); |
| 159 | let tool_calls_recent_window = self.recent_tool_call_count(message_window); |
| 160 | let unique_reference_ids_recent_window = |
| 161 | self.recent_unique_reference_count(message_window, turn); |
| 162 | let context_window = usize::try_from( |
| 163 | context_window_for_model(&self.session.model) |
| 164 | .unwrap_or(LEGACY_DEEPSEEK_CONTEXT_WINDOW_TOKENS), |
| 165 | ) |
| 166 | .unwrap_or(usize::try_from(LEGACY_DEEPSEEK_CONTEXT_WINDOW_TOKENS).unwrap_or(128_000)) |
| 167 | .max(1); |
| 168 | let context_used_ratio = (self.estimated_input_tokens() as f64) / (context_window as f64); |
| 169 | |
| 170 | CapacityObservationInput { |
| 171 | turn_index: self.turn_counter, |
| 172 | model: self.session.model.clone(), |
| 173 | action_count_this_turn, |
| 174 | tool_calls_recent_window, |
| 175 | unique_reference_ids_recent_window, |
| 176 | context_used_ratio, |
| 177 | } |
| 178 | } |
| 179 | |
| 180 | pub(super) fn recent_tool_call_count(&self, message_window: usize) -> usize { |
| 181 | self.session |
| 182 | .messages |
| 183 | .iter() |
| 184 | .rev() |
| 185 | .take(message_window) |
| 186 | .map(|msg| { |
| 187 | msg.content |
| 188 | .iter() |
| 189 | .filter(|block| { |
| 190 | matches!( |
| 191 | block, |
| 192 | ContentBlock::ToolUse { .. } | ContentBlock::ToolResult { .. } |
| 193 | ) |
| 194 | }) |
| 195 | .count() |
| 196 | }) |
| 197 | .sum() |
| 198 | } |
| 199 | |
| 200 | pub(super) fn recent_unique_reference_count( |
| 201 | &self, |
| 202 | message_window: usize, |
| 203 | turn: &TurnContext, |
| 204 | ) -> usize { |
| 205 | let mut refs = std::collections::HashSet::new(); |
| 206 | for msg in self.session.messages.iter().rev().take(message_window) { |
| 207 | for block in &msg.content { |
| 208 | match block { |
| 209 | ContentBlock::ToolUse { id, .. } => { |
| 210 | refs.insert(id.clone()); |
| 211 | } |
| 212 | ContentBlock::ToolResult { tool_use_id, .. } => { |
| 213 | refs.insert(tool_use_id.clone()); |
| 214 | } |
| 215 | ContentBlock::Text { text, .. } => { |
| 216 | for token in text.split_whitespace() { |
| 217 | if token.contains('/') || token.contains('.') { |
| 218 | refs.insert( |
| 219 | token |
| 220 | .trim_matches(|c: char| ",.;:()[]{}".contains(c)) |
| 221 | .to_string(), |
| 222 | ); |
| 223 | } |
| 224 | } |
| 225 | } |
| 226 | ContentBlock::Thinking { .. } |
| 227 | | ContentBlock::ServerToolUse { .. } |
| 228 | | ContentBlock::ToolSearchToolResult { .. } |
| 229 | | ContentBlock::CodeExecutionToolResult { .. } => {} |
| 230 | } |
| 231 | } |
| 232 | } |
| 233 | for tool_call in turn.tool_calls.iter().rev().take(8) { |
| 234 | refs.insert(tool_call.id.clone()); |
| 235 | } |
| 236 | for path in self.session.working_set.top_paths(8) { |
| 237 | refs.insert(path); |
| 238 | } |
| 239 | refs.retain(|item| !item.is_empty()); |
| 240 | refs.len() |
| 241 | } |
| 242 | |
| 243 | pub(super) async fn emit_coherence_signal( |
| 244 | &mut self, |
| 245 | signal: CoherenceSignal, |
| 246 | reason: impl Into<String>, |
| 247 | ) { |
| 248 | let next = next_coherence_state(self.coherence_state, signal); |
| 249 | self.coherence_state = next; |
| 250 | let _ = self |
| 251 | .tx_event |
| 252 | .send(Event::CoherenceState { |
| 253 | state: next, |
| 254 | label: next.label().to_string(), |
| 255 | description: next.description().to_string(), |
| 256 | reason: reason.into(), |
| 257 | }) |
| 258 | .await; |
| 259 | } |
| 260 | |
| 261 | pub(super) async fn emit_compaction_started( |
| 262 | &mut self, |
| 263 | id: String, |
| 264 | auto: bool, |
| 265 | message: String, |
| 266 | ) { |
| 267 | let _ = self |
| 268 | .tx_event |
| 269 | .send(Event::CompactionStarted { |
| 270 | id, |
| 271 | auto, |
| 272 | message: message.clone(), |
| 273 | }) |
| 274 | .await; |
| 275 | self.emit_coherence_signal(CoherenceSignal::CompactionStarted, message) |
| 276 | .await; |
| 277 | } |
| 278 | |
| 279 | pub(super) async fn emit_compaction_completed( |
| 280 | &mut self, |
| 281 | id: String, |
| 282 | auto: bool, |
| 283 | message: String, |
| 284 | messages_before: Option<usize>, |
| 285 | messages_after: Option<usize>, |
| 286 | ) { |
| 287 | let _ = self |
| 288 | .tx_event |
| 289 | .send(Event::CompactionCompleted { |
| 290 | id, |
| 291 | auto, |
| 292 | message: message.clone(), |
| 293 | messages_before, |
| 294 | messages_after, |
| 295 | }) |
| 296 | .await; |
| 297 | self.emit_coherence_signal(CoherenceSignal::CompactionCompleted, message) |
| 298 | .await; |
| 299 | } |
| 300 | |
| 301 | pub(super) async fn emit_compaction_failed(&mut self, id: String, auto: bool, message: String) { |
| 302 | let _ = self |
| 303 | .tx_event |
| 304 | .send(Event::CompactionFailed { |
| 305 | id, |
| 306 | auto, |
| 307 | message: message.clone(), |
| 308 | }) |
| 309 | .await; |
| 310 | self.emit_coherence_signal(CoherenceSignal::CompactionFailed, message) |
| 311 | .await; |
| 312 | } |
| 313 | |
| 314 | pub(super) async fn emit_capacity_decision( |
| 315 | &mut self, |
| 316 | turn: &TurnContext, |
| 317 | snapshot: Option<&CapacitySnapshot>, |
| 318 | decision: &CapacityDecision, |
| 319 | ) { |
| 320 | let Some(snapshot) = snapshot else { |
| 321 | return; |
| 322 | }; |
| 323 | let _ = self |
| 324 | .tx_event |
| 325 | .send(Event::CapacityDecision { |
| 326 | session_id: self.session.id.clone(), |
| 327 | turn_id: turn.id.clone(), |
| 328 | h_hat: snapshot.h_hat, |
| 329 | c_hat: snapshot.c_hat, |
| 330 | slack: snapshot.slack, |
| 331 | min_slack: snapshot.profile.min_slack, |
| 332 | violation_ratio: snapshot.profile.violation_ratio, |
| 333 | p_fail: snapshot.p_fail, |
| 334 | risk_band: snapshot.risk_band.as_str().to_string(), |
| 335 | action: decision.action.as_str().to_string(), |
| 336 | cooldown_blocked: decision.cooldown_blocked, |
| 337 | reason: decision.reason.clone(), |
| 338 | }) |
| 339 | .await; |
| 340 | self.emit_coherence_signal( |
| 341 | CoherenceSignal::CapacityDecision { |
| 342 | risk_band: snapshot.risk_band, |
| 343 | action: decision.action, |
| 344 | cooldown_blocked: decision.cooldown_blocked, |
| 345 | }, |
| 346 | format!( |
| 347 | "capacity_decision: risk={} action={} reason={}", |
| 348 | snapshot.risk_band.as_str(), |
| 349 | decision.action.as_str(), |
| 350 | decision.reason |
| 351 | ), |
| 352 | ) |
| 353 | .await; |
| 354 | } |
| 355 | |
| 356 | pub(super) async fn emit_capacity_intervention( |
| 357 | &mut self, |
| 358 | turn: &TurnContext, |
| 359 | action: GuardrailAction, |
| 360 | before_prompt_tokens: usize, |
| 361 | after_prompt_tokens: usize, |
| 362 | replay_outcome: Option<String>, |
| 363 | replan_performed: bool, |
| 364 | ) { |
| 365 | let _ = self |
| 366 | .tx_event |
| 367 | .send(Event::CapacityIntervention { |
| 368 | session_id: self.session.id.clone(), |
| 369 | turn_id: turn.id.clone(), |
| 370 | action: action.as_str().to_string(), |
| 371 | before_prompt_tokens, |
| 372 | after_prompt_tokens, |
| 373 | compaction_size_reduction: before_prompt_tokens.saturating_sub(after_prompt_tokens), |
| 374 | replay_outcome, |
| 375 | replan_performed, |
| 376 | }) |
| 377 | .await; |
| 378 | self.emit_coherence_signal( |
| 379 | CoherenceSignal::CapacityIntervention { action }, |
| 380 | format!("capacity_intervention: action={}", action.as_str()), |
| 381 | ) |
| 382 | .await; |
| 383 | } |
| 384 | |
| 385 | pub(super) async fn apply_targeted_context_refresh( |
| 386 | &mut self, |
| 387 | turn: &TurnContext, |
| 388 | client: Option<&DeepSeekClient>, |
| 389 | mode: AppMode, |
| 390 | snapshot: Option<&CapacitySnapshot>, |
| 391 | ) -> bool { |
| 392 | let before_tokens = self.estimated_input_tokens(); |
| 393 | let compaction_pins = self |
| 394 | .session |
| 395 | .working_set |
| 396 | .pinned_message_indices(&self.session.messages, &self.session.workspace); |
| 397 | let compaction_paths = self.session.working_set.top_paths(24); |
| 398 | |
| 399 | let mut refreshed = false; |
| 400 | let should_run_summary_compaction = self.config.compaction.enabled |
| 401 | && should_compact( |
| 402 | &self.session.messages, |
| 403 | &self.config.compaction, |
| 404 | Some(&self.session.workspace), |
| 405 | Some(&compaction_pins), |
| 406 | Some(&compaction_paths), |
| 407 | ); |
| 408 | if should_run_summary_compaction && let Some(client) = client { |
| 409 | match compact_messages_safe( |
| 410 | client, |
| 411 | &self.session.messages, |
| 412 | &self.config.compaction, |
| 413 | Some(&self.session.workspace), |
| 414 | Some(&compaction_pins), |
| 415 | Some(&compaction_paths), |
| 416 | ) |
| 417 | .await |
| 418 | { |
| 419 | Ok(result) => { |
| 420 | if !result.messages.is_empty() || self.session.messages.is_empty() { |
| 421 | self.session.messages = result.messages; |
| 422 | self.merge_compaction_summary(result.summary_prompt); |
| 423 | refreshed = true; |
| 424 | } |
| 425 | } |
| 426 | Err(err) => { |
| 427 | let _ = self |
| 428 | .tx_event |
| 429 | .send(Event::status(format!( |
| 430 | "Capacity refresh compaction failed: {err}. Falling back to local trim." |
| 431 | ))) |
| 432 | .await; |
| 433 | } |
| 434 | } |
| 435 | } |
| 436 | |
| 437 | if !refreshed { |
| 438 | let target_budget = context_input_budget(&self.session.model, TURN_MAX_OUTPUT_TOKENS) |
| 439 | .unwrap_or(self.config.compaction.token_threshold.max(1)); |
| 440 | if self.estimated_input_tokens() > target_budget { |
| 441 | let trimmed = self.trim_oldest_messages_to_budget(target_budget); |
| 442 | refreshed = trimmed > 0; |
| 443 | } |
| 444 | } |
| 445 | |
| 446 | if !refreshed { |
| 447 | return false; |
| 448 | } |
| 449 | |
| 450 | let canonical = self.build_canonical_state(turn, None); |
| 451 | let source_message_ids = self.capacity_source_message_ids(turn); |
| 452 | let record = self.build_capacity_record( |
| 453 | turn, |
| 454 | GuardrailAction::TargetedContextRefresh, |
| 455 | snapshot, |
| 456 | canonical.clone(), |
| 457 | source_message_ids, |
| 458 | None, |
| 459 | ); |
| 460 | let pointer = self |
| 461 | .persist_capacity_record(turn, GuardrailAction::TargetedContextRefresh, &record) |
| 462 | .await; |
| 463 | self.merge_compaction_summary(Some(self.canonical_prompt( |
| 464 | &canonical, |
| 465 | &pointer, |
| 466 | GuardrailAction::TargetedContextRefresh, |
| 467 | None, |
| 468 | ))); |
| 469 | self.refresh_system_prompt(mode); |
| 470 | self.emit_session_updated().await; |
| 471 | |
| 472 | let after_tokens = self.estimated_input_tokens(); |
| 473 | self.emit_capacity_intervention( |
| 474 | turn, |
| 475 | GuardrailAction::TargetedContextRefresh, |
| 476 | before_tokens, |
| 477 | after_tokens, |
| 478 | None, |
| 479 | false, |
| 480 | ) |
| 481 | .await; |
| 482 | self.capacity_controller |
| 483 | .mark_intervention_applied(self.turn_counter, GuardrailAction::TargetedContextRefresh); |
| 484 | true |
| 485 | } |
| 486 | |
| 487 | #[allow(clippy::too_many_arguments)] |
| 488 | pub(super) async fn apply_verify_with_tool_replay( |
| 489 | &mut self, |
| 490 | turn: &TurnContext, |
| 491 | mode: AppMode, |
| 492 | snapshot: Option<&CapacitySnapshot>, |
| 493 | tool_registry: Option<&crate::tools::ToolRegistry>, |
| 494 | tool_exec_lock: Arc<RwLock<()>>, |
| 495 | mut mcp_pool: Option<Arc<AsyncMutex<McpPool>>>, |
| 496 | ) -> bool { |
| 497 | let before_tokens = self.estimated_input_tokens(); |
| 498 | let Some(candidate) = self.select_replay_candidate(turn, tool_registry) else { |
| 499 | return false; |
| 500 | }; |
| 501 | |
| 502 | if McpPool::is_mcp_tool(&candidate.name) && mcp_pool.is_none() { |
| 503 | mcp_pool = self.ensure_mcp_pool().await.ok(); |
| 504 | } |
| 505 | |
| 506 | let supports_parallel = if McpPool::is_mcp_tool(&candidate.name) { |
| 507 | mcp_tool_is_parallel_safe(&candidate.name) |
| 508 | } else { |
| 509 | tool_registry |
| 510 | .and_then(|registry| registry.get(&candidate.name)) |
| 511 | .is_some_and(|spec| spec.supports_parallel()) |
| 512 | }; |
| 513 | let interactive = (candidate.name == "exec_shell" |
| 514 | && candidate |
| 515 | .input |
| 516 | .get("interactive") |
| 517 | .and_then(serde_json::Value::as_bool) |
| 518 | == Some(true)) |
| 519 | || candidate.name == REQUEST_USER_INPUT_NAME; |
| 520 | |
| 521 | let replay_result = Self::execute_tool_with_lock( |
| 522 | tool_exec_lock, |
| 523 | supports_parallel, |
| 524 | interactive, |
| 525 | self.tx_event.clone(), |
| 526 | candidate.name.clone(), |
| 527 | candidate.input.clone(), |
| 528 | tool_registry, |
| 529 | mcp_pool.clone(), |
| 530 | None, |
| 531 | ) |
| 532 | .await; |
| 533 | |
| 534 | let (pass, replay_outcome, diff_summary) = match replay_result { |
| 535 | Ok(output) => { |
| 536 | let original = candidate.result.as_deref().unwrap_or_default(); |
| 537 | let replay = output.content.as_str(); |
| 538 | let equal = original.trim() == replay.trim(); |
| 539 | let diff = if equal { |
| 540 | "output_match".to_string() |
| 541 | } else { |
| 542 | format!( |
| 543 | "output_mismatch: original='{}' replay='{}'", |
| 544 | summarize_text(original, 140), |
| 545 | summarize_text(replay, 140) |
| 546 | ) |
| 547 | }; |
| 548 | ( |
| 549 | equal, |
| 550 | if equal { |
| 551 | "pass".to_string() |
| 552 | } else { |
| 553 | "conflict".to_string() |
| 554 | }, |
| 555 | diff, |
| 556 | ) |
| 557 | } |
| 558 | Err(err) => { |
| 559 | self.capacity_controller |
| 560 | .mark_replay_failed(self.turn_counter); |
| 561 | ( |
| 562 | false, |
| 563 | "error".to_string(), |
| 564 | format!("replay_error: {}", summarize_text(&err.to_string(), 180)), |
| 565 | ) |
| 566 | } |
| 567 | }; |
| 568 | |
| 569 | let verification_note = format!( |
| 570 | "[verification replay] tool={} pass={} details={}", |
| 571 | candidate.name, pass, diff_summary |
| 572 | ); |
| 573 | self.add_session_message(Message { |
| 574 | role: "user".to_string(), |
| 575 | content: vec![ContentBlock::ToolResult { |
| 576 | tool_use_id: candidate.id.clone(), |
| 577 | content: verification_note.clone(), |
| 578 | is_error: None, |
| 579 | content_blocks: None, |
| 580 | }], |
| 581 | }) |
| 582 | .await; |
| 583 | |
| 584 | if !pass { |
| 585 | self.capacity_controller |
| 586 | .mark_replay_failed(self.turn_counter); |
| 587 | } |
| 588 | |
| 589 | let canonical = self.build_canonical_state( |
| 590 | turn, |
| 591 | Some(if pass { |
| 592 | "replay verification passed" |
| 593 | } else { |
| 594 | "replay verification failed or conflicted" |
| 595 | }), |
| 596 | ); |
| 597 | let replay_info = Some(ReplayInfo { |
| 598 | tool_id: candidate.id.clone(), |
| 599 | tool_name: candidate.name.clone(), |
| 600 | pass, |
| 601 | diff_summary: diff_summary.clone(), |
| 602 | }); |
| 603 | let source_message_ids = self.capacity_source_message_ids(turn); |
| 604 | let record = self.build_capacity_record( |
| 605 | turn, |
| 606 | GuardrailAction::VerifyWithToolReplay, |
| 607 | snapshot, |
| 608 | canonical.clone(), |
| 609 | source_message_ids, |
| 610 | replay_info, |
| 611 | ); |
| 612 | let pointer = self |
| 613 | .persist_capacity_record(turn, GuardrailAction::VerifyWithToolReplay, &record) |
| 614 | .await; |
| 615 | self.merge_compaction_summary(Some(self.canonical_prompt( |
| 616 | &canonical, |
| 617 | &pointer, |
| 618 | GuardrailAction::VerifyWithToolReplay, |
| 619 | Some(&verification_note), |
| 620 | ))); |
| 621 | self.refresh_system_prompt(mode); |
| 622 | self.emit_session_updated().await; |
| 623 | |
| 624 | let after_tokens = self.estimated_input_tokens(); |
| 625 | self.emit_capacity_intervention( |
| 626 | turn, |
| 627 | GuardrailAction::VerifyWithToolReplay, |
| 628 | before_tokens, |
| 629 | after_tokens, |
| 630 | Some(replay_outcome), |
| 631 | false, |
| 632 | ) |
| 633 | .await; |
| 634 | self.capacity_controller |
| 635 | .mark_intervention_applied(self.turn_counter, GuardrailAction::VerifyWithToolReplay); |
| 636 | true |
| 637 | } |
| 638 | |
| 639 | pub(super) async fn apply_verify_and_replan( |
| 640 | &mut self, |
| 641 | turn: &TurnContext, |
| 642 | mode: AppMode, |
| 643 | snapshot: Option<&CapacitySnapshot>, |
| 644 | reason: &str, |
| 645 | ) -> bool { |
| 646 | let before_tokens = self.estimated_input_tokens(); |
| 647 | let canonical = self.build_canonical_state(turn, Some(reason)); |
| 648 | let source_message_ids = self.capacity_source_message_ids(turn); |
| 649 | let record = self.build_capacity_record( |
| 650 | turn, |
| 651 | GuardrailAction::VerifyAndReplan, |
| 652 | snapshot, |
| 653 | canonical.clone(), |
| 654 | source_message_ids, |
| 655 | None, |
| 656 | ); |
| 657 | let pointer = self |
| 658 | .persist_capacity_record(turn, GuardrailAction::VerifyAndReplan, &record) |
| 659 | .await; |
| 660 | |
| 661 | let latest_user = self |
| 662 | .session |
| 663 | .messages |
| 664 | .iter() |
| 665 | .rev() |
| 666 | .find(|msg| { |
| 667 | msg.role == "user" |
| 668 | && msg |
| 669 | .content |
| 670 | .iter() |
| 671 | .any(|block| matches!(block, ContentBlock::Text { .. })) |
| 672 | }) |
| 673 | .cloned(); |
| 674 | let latest_verified = self |
| 675 | .session |
| 676 | .messages |
| 677 | .iter() |
| 678 | .rev() |
| 679 | .find(|msg| { |
| 680 | msg.role == "user" |
| 681 | && msg.content.iter().any(|block| match block { |
| 682 | ContentBlock::ToolResult { content, .. } => { |
| 683 | content.contains("[verification replay]") |
| 684 | } |
| 685 | _ => false, |
| 686 | }) |
| 687 | }) |
| 688 | .cloned(); |
| 689 | |
| 690 | self.session.messages.clear(); |
| 691 | if let Some(msg) = latest_user { |
| 692 | self.session.messages.push(msg); |
| 693 | } |
| 694 | if let Some(msg) = latest_verified { |
| 695 | self.session.messages.push(msg); |
| 696 | } |
| 697 | |
| 698 | self.merge_compaction_summary(Some(self.canonical_prompt( |
| 699 | &canonical, |
| 700 | &pointer, |
| 701 | GuardrailAction::VerifyAndReplan, |
| 702 | Some("Replan now from canonical state. Keep steps minimal and verifiable."), |
| 703 | ))); |
| 704 | self.refresh_system_prompt(mode); |
| 705 | self.emit_session_updated().await; |
| 706 | |
| 707 | let _ = self |
| 708 | .tx_event |
| 709 | .send(Event::status( |
| 710 | "Capacity guardrail: context reset to canonical state; replanning step." |
| 711 | .to_string(), |
| 712 | )) |
| 713 | .await; |
| 714 | |
| 715 | let after_tokens = self.estimated_input_tokens(); |
| 716 | self.emit_capacity_intervention( |
| 717 | turn, |
| 718 | GuardrailAction::VerifyAndReplan, |
| 719 | before_tokens, |
| 720 | after_tokens, |
| 721 | None, |
| 722 | true, |
| 723 | ) |
| 724 | .await; |
| 725 | self.capacity_controller |
| 726 | .mark_intervention_applied(self.turn_counter, GuardrailAction::VerifyAndReplan); |
| 727 | true |
| 728 | } |
| 729 | |
| 730 | pub(super) fn select_replay_candidate( |
| 731 | &self, |
| 732 | turn: &TurnContext, |
| 733 | tool_registry: Option<&crate::tools::ToolRegistry>, |
| 734 | ) -> Option<TurnToolCall> { |
| 735 | turn.tool_calls |
| 736 | .iter() |
| 737 | .rev() |
| 738 | .find(|call| { |
| 739 | call.error.is_none() |
| 740 | && call.result.is_some() |
| 741 | && self.tool_is_replayable_read_only(&call.name, tool_registry) |
| 742 | }) |
| 743 | .cloned() |
| 744 | } |
| 745 | |
| 746 | pub(super) fn tool_is_replayable_read_only( |
| 747 | &self, |
| 748 | tool_name: &str, |
| 749 | tool_registry: Option<&crate::tools::ToolRegistry>, |
| 750 | ) -> bool { |
| 751 | if tool_name == MULTI_TOOL_PARALLEL_NAME || tool_name == REQUEST_USER_INPUT_NAME { |
| 752 | return false; |
| 753 | } |
| 754 | if McpPool::is_mcp_tool(tool_name) { |
| 755 | return mcp_tool_is_read_only(tool_name); |
| 756 | } |
| 757 | tool_registry |
| 758 | .and_then(|registry| registry.get(tool_name)) |
| 759 | .is_some_and(|spec| spec.is_read_only()) |
| 760 | } |
| 761 | |
| 762 | pub(super) fn build_canonical_state( |
| 763 | &self, |
| 764 | turn: &TurnContext, |
| 765 | note: Option<&str>, |
| 766 | ) -> CanonicalState { |
| 767 | let goal = self |
| 768 | .session |
| 769 | .messages |
| 770 | .iter() |
| 771 | .rev() |
| 772 | .find_map(|msg| { |
| 773 | if msg.role != "user" { |
| 774 | return None; |
| 775 | } |
| 776 | msg.content.iter().find_map(|block| match block { |
| 777 | ContentBlock::Text { text, .. } => Some(summarize_text(text, 220)), |
| 778 | _ => None, |
| 779 | }) |
| 780 | }) |
| 781 | .unwrap_or_else(|| "Continue current task from compact state".to_string()); |
| 782 | |
| 783 | let mut constraints = vec![ |
| 784 | format!("model={}", self.session.model), |
| 785 | format!("workspace={}", self.session.workspace.display()), |
| 786 | ]; |
| 787 | if let Some(note) = note { |
| 788 | constraints.push(summarize_text(note, 180)); |
| 789 | } |
| 790 | |
| 791 | let mut confirmed_facts = Vec::new(); |
| 792 | for msg in self.session.messages.iter().rev() { |
| 793 | for block in &msg.content { |
| 794 | if let ContentBlock::ToolResult { content, .. } = block { |
| 795 | if content.starts_with("Error:") { |
| 796 | continue; |
| 797 | } |
| 798 | confirmed_facts.push(summarize_text(content, 180)); |
| 799 | if confirmed_facts.len() >= 4 { |
| 800 | break; |
| 801 | } |
| 802 | } |
| 803 | } |
| 804 | if confirmed_facts.len() >= 4 { |
| 805 | break; |
| 806 | } |
| 807 | } |
| 808 | |
| 809 | let open_loops: Vec<String> = turn |
| 810 | .tool_calls |
| 811 | .iter() |
| 812 | .rev() |
| 813 | .filter_map(|call| { |
| 814 | call.error |
| 815 | .as_ref() |
| 816 | .map(|error| format!("{}: {}", call.name, summarize_text(error, 180))) |
| 817 | }) |
| 818 | .take(4) |
| 819 | .collect(); |
| 820 | |
| 821 | let pending_actions: Vec<String> = if open_loops.is_empty() { |
| 822 | vec!["Continue with next smallest verifiable step".to_string()] |
| 823 | } else { |
| 824 | vec![ |
| 825 | "Re-evaluate failed tool steps with narrower scope".to_string(), |
| 826 | "Re-derive plan from canonical facts before further edits".to_string(), |
| 827 | ] |
| 828 | }; |
| 829 | |
| 830 | let mut critical_refs = self.session.working_set.top_paths(8); |
| 831 | for tool_call in turn.tool_calls.iter().rev().take(4) { |
| 832 | critical_refs.push(format!("tool:{}", tool_call.id)); |
| 833 | } |
| 834 | critical_refs.dedup(); |
| 835 | |
| 836 | CanonicalState { |
| 837 | goal, |
| 838 | constraints, |
| 839 | confirmed_facts, |
| 840 | open_loops, |
| 841 | pending_actions, |
| 842 | critical_refs, |
| 843 | } |
| 844 | } |
| 845 | |
| 846 | pub(super) fn canonical_prompt( |
| 847 | &self, |
| 848 | canonical: &CanonicalState, |
| 849 | pointer: &str, |
| 850 | action: GuardrailAction, |
| 851 | extra: Option<&str>, |
| 852 | ) -> SystemPrompt { |
| 853 | let mut lines = vec![ |
| 854 | COMPACTION_SUMMARY_MARKER.to_string(), |
| 855 | format!("Capacity Canonical State [{}]", action.as_str()), |
| 856 | format!("Goal: {}", canonical.goal), |
| 857 | "Constraints:".to_string(), |
| 858 | ]; |
| 859 | for item in &canonical.constraints { |
| 860 | lines.push(format!("- {}", summarize_text(item, 200))); |
| 861 | } |
| 862 | lines.push("Confirmed Facts:".to_string()); |
| 863 | for item in &canonical.confirmed_facts { |
| 864 | lines.push(format!("- {}", summarize_text(item, 200))); |
| 865 | } |
| 866 | lines.push("Open Loops:".to_string()); |
| 867 | if canonical.open_loops.is_empty() { |
| 868 | lines.push("- none".to_string()); |
| 869 | } else { |
| 870 | for item in &canonical.open_loops { |
| 871 | lines.push(format!("- {}", summarize_text(item, 200))); |
| 872 | } |
| 873 | } |
| 874 | lines.push("Pending Actions:".to_string()); |
| 875 | for item in &canonical.pending_actions { |
| 876 | lines.push(format!("- {}", summarize_text(item, 200))); |
| 877 | } |
| 878 | lines.push("Critical Refs:".to_string()); |
| 879 | for item in &canonical.critical_refs { |
| 880 | lines.push(format!("- {}", summarize_text(item, 200))); |
| 881 | } |
| 882 | if let Some(extra) = extra { |
| 883 | lines.push(format!("Instruction: {}", summarize_text(extra, 240))); |
| 884 | } |
| 885 | lines.push(format!("Memory Pointer: {pointer}")); |
| 886 | |
| 887 | SystemPrompt::Blocks(vec![crate::models::SystemBlock { |
| 888 | block_type: "text".to_string(), |
| 889 | text: lines.join("\n"), |
| 890 | cache_control: None, |
| 891 | }]) |
| 892 | } |
| 893 | |
| 894 | pub(super) fn capacity_source_message_ids(&self, turn: &TurnContext) -> Vec<String> { |
| 895 | let mut ids: Vec<String> = turn |
| 896 | .tool_calls |
| 897 | .iter() |
| 898 | .rev() |
| 899 | .take(8) |
| 900 | .map(|call| call.id.clone()) |
| 901 | .collect(); |
| 902 | ids.reverse(); |
| 903 | ids |
| 904 | } |
| 905 | |
| 906 | pub(super) fn build_capacity_record( |
| 907 | &self, |
| 908 | turn: &TurnContext, |
| 909 | action: GuardrailAction, |
| 910 | snapshot: Option<&CapacitySnapshot>, |
| 911 | canonical: CanonicalState, |
| 912 | source_message_ids: Vec<String>, |
| 913 | replay_info: Option<ReplayInfo>, |
| 914 | ) -> CapacityMemoryRecord { |
| 915 | let (h_hat, c_hat, slack, risk_band) = snapshot |
| 916 | .map(|s| (s.h_hat, s.c_hat, s.slack, s.risk_band.as_str().to_string())) |
| 917 | .unwrap_or_else(|| (0.0, 0.0, 0.0, "unknown".to_string())); |
| 918 | |
| 919 | CapacityMemoryRecord { |
| 920 | id: new_record_id(), |
| 921 | ts: now_rfc3339(), |
| 922 | turn_index: self.turn_counter, |
| 923 | action_trigger: action.as_str().to_string(), |
| 924 | h_hat, |
| 925 | c_hat, |
| 926 | slack, |
| 927 | risk_band, |
| 928 | canonical_state: canonical, |
| 929 | source_message_ids: if source_message_ids.is_empty() { |
| 930 | vec![turn.id.clone()] |
| 931 | } else { |
| 932 | source_message_ids |
| 933 | }, |
| 934 | replay_info, |
| 935 | } |
| 936 | } |
| 937 | |
| 938 | pub(super) async fn persist_capacity_record( |
| 939 | &mut self, |
| 940 | turn: &TurnContext, |
| 941 | action: GuardrailAction, |
| 942 | record: &CapacityMemoryRecord, |
| 943 | ) -> String { |
| 944 | let pointer = format!("memory://{}/{}", self.session.id, record.id); |
| 945 | if let Err(err) = append_capacity_record(&self.session.id, record) { |
| 946 | let _ = self |
| 947 | .tx_event |
| 948 | .send(Event::CapacityMemoryPersistFailed { |
| 949 | session_id: self.session.id.clone(), |
| 950 | turn_id: turn.id.clone(), |
| 951 | action: action.as_str().to_string(), |
| 952 | error: summarize_text(&err.to_string(), 280), |
| 953 | }) |
| 954 | .await; |
| 955 | return format!("{pointer}?persist=failed"); |
| 956 | } |
| 957 | pointer |
| 958 | } |
| 959 | |
| 960 | pub(super) fn rehydrate_latest_canonical_state(&mut self) { |
| 961 | let Ok(records) = load_last_k_capacity_records(&self.session.id, 1) else { |
| 962 | return; |
| 963 | }; |
| 964 | let Some(last) = records.last() else { |
| 965 | return; |
| 966 | }; |
| 967 | let pointer = format!("memory://{}/{}", self.session.id, last.id); |
| 968 | let prompt = self.canonical_prompt( |
| 969 | &last.canonical_state, |
| 970 | &pointer, |
| 971 | GuardrailAction::NoIntervention, |
| 972 | Some("Rehydrated canonical state from memory."), |
| 973 | ); |
| 974 | self.merge_compaction_summary(Some(prompt)); |
| 975 | } |
| 976 | } |
| 977 |