| 1 | //! Capacity-aware guardrail controller for context pressure management. |
| 2 | |
| 3 | use std::collections::{HashMap, VecDeque}; |
| 4 | |
| 5 | /// Controller settings. |
| 6 | #[derive(Debug, Clone, PartialEq)] |
| 7 | pub struct CapacityControllerConfig { |
| 8 | pub enabled: bool, |
| 9 | pub low_risk_max: f64, |
| 10 | pub medium_risk_max: f64, |
| 11 | pub severe_min_slack: f64, |
| 12 | pub severe_violation_ratio: f64, |
| 13 | pub refresh_cooldown_turns: u64, |
| 14 | pub replan_cooldown_turns: u64, |
| 15 | pub max_replay_per_turn: usize, |
| 16 | pub min_turns_before_guardrail: u64, |
| 17 | pub profile_window: usize, |
| 18 | pub model_priors: HashMap<String, f64>, |
| 19 | pub fallback_default: f64, |
| 20 | } |
| 21 | |
| 22 | impl Default for CapacityControllerConfig { |
| 23 | fn default() -> Self { |
| 24 | let mut model_priors = HashMap::new(); |
| 25 | model_priors.insert("deepseek_v3_2_chat".to_string(), 3.9); |
| 26 | model_priors.insert("deepseek_v3_2_reasoner".to_string(), 4.1); |
| 27 | model_priors.insert("deepseek_v4_pro".to_string(), 3.5); |
| 28 | model_priors.insert("deepseek_v4_flash".to_string(), 4.2); |
| 29 | |
| 30 | Self { |
| 31 | // OFF BY DEFAULT since v0.8.11. The capacity controller's |
| 32 | // interventions (TargetedContextRefresh, VerifyAndReplan) |
| 33 | // silently rewrite or clear the session message log, which |
| 34 | // surprises the user and destroys V4's prefix cache. v0.8.11 |
| 35 | // committed to "trust the model with the full 1M-token |
| 36 | // context, only compact on explicit user `/compact`." |
| 37 | // Auto-managing the prefix on the user's behalf works against |
| 38 | // that posture. Power users who want the controller can opt |
| 39 | // in via `capacity.enabled = true` in |
| 40 | // `~/.deepseek/config.toml`. |
| 41 | enabled: false, |
| 42 | // Thresholds retained for the opt-in path; tuning notes live |
| 43 | // in git history (#63 follow-up). |
| 44 | low_risk_max: 0.50, |
| 45 | medium_risk_max: 0.62, |
| 46 | severe_min_slack: -0.25, |
| 47 | severe_violation_ratio: 0.40, |
| 48 | refresh_cooldown_turns: 6, |
| 49 | replan_cooldown_turns: 5, |
| 50 | max_replay_per_turn: 1, |
| 51 | min_turns_before_guardrail: 4, |
| 52 | profile_window: 8, |
| 53 | model_priors, |
| 54 | fallback_default: 3.8, |
| 55 | } |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | impl CapacityControllerConfig { |
| 60 | /// Build effective capacity config from app config. |
| 61 | #[must_use] |
| 62 | pub fn from_app_config(config: &crate::config::Config) -> Self { |
| 63 | let mut out = Self::default(); |
| 64 | let Some(capacity) = config.capacity.as_ref() else { |
| 65 | return out; |
| 66 | }; |
| 67 | |
| 68 | if let Some(v) = capacity.enabled { |
| 69 | out.enabled = v; |
| 70 | } |
| 71 | if let Some(v) = capacity.low_risk_max { |
| 72 | out.low_risk_max = v; |
| 73 | } |
| 74 | if let Some(v) = capacity.medium_risk_max { |
| 75 | out.medium_risk_max = v; |
| 76 | } |
| 77 | if let Some(v) = capacity.severe_min_slack { |
| 78 | out.severe_min_slack = v; |
| 79 | } |
| 80 | if let Some(v) = capacity.severe_violation_ratio { |
| 81 | out.severe_violation_ratio = v; |
| 82 | } |
| 83 | if let Some(v) = capacity.refresh_cooldown_turns { |
| 84 | out.refresh_cooldown_turns = v; |
| 85 | } |
| 86 | if let Some(v) = capacity.replan_cooldown_turns { |
| 87 | out.replan_cooldown_turns = v; |
| 88 | } |
| 89 | if let Some(v) = capacity.max_replay_per_turn { |
| 90 | out.max_replay_per_turn = v; |
| 91 | } |
| 92 | if let Some(v) = capacity.min_turns_before_guardrail { |
| 93 | out.min_turns_before_guardrail = v; |
| 94 | } |
| 95 | if let Some(v) = capacity.profile_window { |
| 96 | out.profile_window = v.max(2); |
| 97 | } |
| 98 | |
| 99 | if let Some(v) = capacity.deepseek_v3_2_chat_prior { |
| 100 | out.model_priors.insert("deepseek_v3_2_chat".to_string(), v); |
| 101 | } |
| 102 | if let Some(v) = capacity.deepseek_v3_2_reasoner_prior { |
| 103 | out.model_priors |
| 104 | .insert("deepseek_v3_2_reasoner".to_string(), v); |
| 105 | } |
| 106 | if let Some(v) = capacity.deepseek_v4_pro_prior { |
| 107 | out.model_priors.insert("deepseek_v4_pro".to_string(), v); |
| 108 | } |
| 109 | if let Some(v) = capacity.deepseek_v4_flash_prior { |
| 110 | out.model_priors.insert("deepseek_v4_flash".to_string(), v); |
| 111 | } |
| 112 | if let Some(v) = capacity.fallback_default_prior { |
| 113 | out.fallback_default = v; |
| 114 | } |
| 115 | |
| 116 | out |
| 117 | } |
| 118 | } |
| 119 | |
| 120 | /// Guardrail decision output. |
| 121 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 122 | pub enum GuardrailAction { |
| 123 | NoIntervention, |
| 124 | TargetedContextRefresh, |
| 125 | VerifyWithToolReplay, |
| 126 | VerifyAndReplan, |
| 127 | } |
| 128 | |
| 129 | impl GuardrailAction { |
| 130 | #[must_use] |
| 131 | pub fn as_str(self) -> &'static str { |
| 132 | match self { |
| 133 | GuardrailAction::NoIntervention => "no_intervention", |
| 134 | GuardrailAction::TargetedContextRefresh => "targeted_context_refresh", |
| 135 | GuardrailAction::VerifyWithToolReplay => "verify_with_tool_replay", |
| 136 | GuardrailAction::VerifyAndReplan => "verify_and_replan", |
| 137 | } |
| 138 | } |
| 139 | } |
| 140 | |
| 141 | /// Coarse failure risk band. |
| 142 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 143 | pub enum RiskBand { |
| 144 | Low, |
| 145 | Medium, |
| 146 | High, |
| 147 | } |
| 148 | |
| 149 | impl RiskBand { |
| 150 | #[must_use] |
| 151 | pub fn as_str(self) -> &'static str { |
| 152 | match self { |
| 153 | RiskBand::Low => "low", |
| 154 | RiskBand::Medium => "medium", |
| 155 | RiskBand::High => "high", |
| 156 | } |
| 157 | } |
| 158 | } |
| 159 | |
| 160 | /// Input used to observe current turn pressure. |
| 161 | #[derive(Debug, Clone)] |
| 162 | pub struct CapacityObservationInput { |
| 163 | pub turn_index: u64, |
| 164 | pub model: String, |
| 165 | pub action_count_this_turn: usize, |
| 166 | pub tool_calls_recent_window: usize, |
| 167 | pub unique_reference_ids_recent_window: usize, |
| 168 | pub context_used_ratio: f64, |
| 169 | } |
| 170 | |
| 171 | /// Rolling slack profile. |
| 172 | #[derive(Debug, Clone, Copy, Default)] |
| 173 | pub struct DynamicSlackProfile { |
| 174 | pub final_slack: f64, |
| 175 | pub min_slack: f64, |
| 176 | pub violation_ratio: f64, |
| 177 | pub slack_volatility: f64, |
| 178 | pub slack_drop: f64, |
| 179 | } |
| 180 | |
| 181 | /// Per-checkpoint capacity snapshot. |
| 182 | #[derive(Debug, Clone)] |
| 183 | pub struct CapacitySnapshot { |
| 184 | pub turn_index: u64, |
| 185 | pub h_hat: f64, |
| 186 | pub c_hat: f64, |
| 187 | pub slack: f64, |
| 188 | pub profile: DynamicSlackProfile, |
| 189 | pub p_fail: f64, |
| 190 | pub risk_band: RiskBand, |
| 191 | pub severe: bool, |
| 192 | } |
| 193 | |
| 194 | /// Full controller decision including reason and block flags. |
| 195 | #[derive(Debug, Clone)] |
| 196 | pub struct CapacityDecision { |
| 197 | pub action: GuardrailAction, |
| 198 | pub reason: String, |
| 199 | pub cooldown_blocked: bool, |
| 200 | } |
| 201 | |
| 202 | #[derive(Debug, Clone, Default)] |
| 203 | struct GuardrailRuntimeState { |
| 204 | last_refresh_turn: Option<u64>, |
| 205 | last_replan_turn: Option<u64>, |
| 206 | replay_count_this_turn: usize, |
| 207 | replay_disabled_turn: Option<u64>, |
| 208 | intervention_applied_turn: Option<u64>, |
| 209 | } |
| 210 | |
| 211 | /// Capacity controller. |
| 212 | #[derive(Debug, Clone)] |
| 213 | pub struct CapacityController { |
| 214 | config: CapacityControllerConfig, |
| 215 | slack_window: VecDeque<f64>, |
| 216 | recent_tool_counts: VecDeque<usize>, |
| 217 | recent_ref_counts: VecDeque<usize>, |
| 218 | state: GuardrailRuntimeState, |
| 219 | last_snapshot: Option<CapacitySnapshot>, |
| 220 | } |
| 221 | |
| 222 | impl CapacityController { |
| 223 | #[must_use] |
| 224 | pub fn new(config: CapacityControllerConfig) -> Self { |
| 225 | Self { |
| 226 | config, |
| 227 | slack_window: VecDeque::new(), |
| 228 | recent_tool_counts: VecDeque::new(), |
| 229 | recent_ref_counts: VecDeque::new(), |
| 230 | state: GuardrailRuntimeState::default(), |
| 231 | last_snapshot: None, |
| 232 | } |
| 233 | } |
| 234 | |
| 235 | pub fn observe_pre_turn( |
| 236 | &mut self, |
| 237 | input: CapacityObservationInput, |
| 238 | ) -> Option<CapacitySnapshot> { |
| 239 | self.observe(input) |
| 240 | } |
| 241 | |
| 242 | pub fn observe_post_tool( |
| 243 | &mut self, |
| 244 | input: CapacityObservationInput, |
| 245 | ) -> Option<CapacitySnapshot> { |
| 246 | self.observe(input) |
| 247 | } |
| 248 | |
| 249 | /// Decide intervention from the latest snapshot, with cooldown and safety gates. |
| 250 | #[must_use] |
| 251 | pub fn decide( |
| 252 | &mut self, |
| 253 | turn_index: u64, |
| 254 | snapshot: Option<&CapacitySnapshot>, |
| 255 | ) -> CapacityDecision { |
| 256 | if !self.config.enabled { |
| 257 | return CapacityDecision { |
| 258 | action: GuardrailAction::NoIntervention, |
| 259 | reason: "capacity_controller_disabled".to_string(), |
| 260 | cooldown_blocked: false, |
| 261 | }; |
| 262 | } |
| 263 | |
| 264 | let Some(snapshot) = snapshot else { |
| 265 | return CapacityDecision { |
| 266 | action: GuardrailAction::NoIntervention, |
| 267 | reason: "missing_capacity_data_fail_open".to_string(), |
| 268 | cooldown_blocked: false, |
| 269 | }; |
| 270 | }; |
| 271 | |
| 272 | if turn_index < self.config.min_turns_before_guardrail { |
| 273 | return CapacityDecision { |
| 274 | action: GuardrailAction::NoIntervention, |
| 275 | reason: "min_turns_before_guardrail_not_reached".to_string(), |
| 276 | cooldown_blocked: false, |
| 277 | }; |
| 278 | } |
| 279 | |
| 280 | let proposed = decide_policy(&self.config, snapshot); |
| 281 | if proposed == GuardrailAction::NoIntervention { |
| 282 | return CapacityDecision { |
| 283 | action: proposed, |
| 284 | reason: "low_risk_no_intervention".to_string(), |
| 285 | cooldown_blocked: false, |
| 286 | }; |
| 287 | } |
| 288 | |
| 289 | if self |
| 290 | .state |
| 291 | .intervention_applied_turn |
| 292 | .is_some_and(|t| t == turn_index) |
| 293 | { |
| 294 | return CapacityDecision { |
| 295 | action: GuardrailAction::NoIntervention, |
| 296 | reason: "intervention_already_applied_this_turn".to_string(), |
| 297 | cooldown_blocked: true, |
| 298 | }; |
| 299 | } |
| 300 | |
| 301 | match proposed { |
| 302 | GuardrailAction::TargetedContextRefresh => { |
| 303 | if self |
| 304 | .state |
| 305 | .last_refresh_turn |
| 306 | .is_some_and(|last| turn_index <= last + self.config.refresh_cooldown_turns) |
| 307 | { |
| 308 | return CapacityDecision { |
| 309 | action: GuardrailAction::NoIntervention, |
| 310 | reason: "refresh_cooldown_active".to_string(), |
| 311 | cooldown_blocked: true, |
| 312 | }; |
| 313 | } |
| 314 | } |
| 315 | GuardrailAction::VerifyWithToolReplay => { |
| 316 | if self |
| 317 | .state |
| 318 | .replay_disabled_turn |
| 319 | .is_some_and(|t| t == turn_index) |
| 320 | { |
| 321 | return CapacityDecision { |
| 322 | action: GuardrailAction::NoIntervention, |
| 323 | reason: "replay_disabled_for_turn".to_string(), |
| 324 | cooldown_blocked: true, |
| 325 | }; |
| 326 | } |
| 327 | if self.state.replay_count_this_turn >= self.config.max_replay_per_turn { |
| 328 | return CapacityDecision { |
| 329 | action: GuardrailAction::NoIntervention, |
| 330 | reason: "max_replay_per_turn_reached".to_string(), |
| 331 | cooldown_blocked: true, |
| 332 | }; |
| 333 | } |
| 334 | } |
| 335 | GuardrailAction::VerifyAndReplan => { |
| 336 | if self |
| 337 | .state |
| 338 | .last_replan_turn |
| 339 | .is_some_and(|last| turn_index <= last + self.config.replan_cooldown_turns) |
| 340 | { |
| 341 | return CapacityDecision { |
| 342 | action: GuardrailAction::NoIntervention, |
| 343 | reason: "replan_cooldown_active".to_string(), |
| 344 | cooldown_blocked: true, |
| 345 | }; |
| 346 | } |
| 347 | } |
| 348 | GuardrailAction::NoIntervention => {} |
| 349 | } |
| 350 | |
| 351 | CapacityDecision { |
| 352 | action: proposed, |
| 353 | reason: "policy_selected_action".to_string(), |
| 354 | cooldown_blocked: false, |
| 355 | } |
| 356 | } |
| 357 | |
| 358 | pub fn mark_turn_start(&mut self, turn_index: u64) { |
| 359 | let new_turn = match self.last_snapshot.as_ref() { |
| 360 | None => true, |
| 361 | Some(snapshot) => snapshot.turn_index != turn_index, |
| 362 | }; |
| 363 | if new_turn { |
| 364 | self.state.replay_count_this_turn = 0; |
| 365 | self.state.replay_disabled_turn = None; |
| 366 | self.state.intervention_applied_turn = None; |
| 367 | } |
| 368 | } |
| 369 | |
| 370 | pub fn mark_intervention_applied(&mut self, turn_index: u64, action: GuardrailAction) { |
| 371 | self.state.intervention_applied_turn = Some(turn_index); |
| 372 | match action { |
| 373 | GuardrailAction::TargetedContextRefresh => { |
| 374 | self.state.last_refresh_turn = Some(turn_index); |
| 375 | } |
| 376 | GuardrailAction::VerifyWithToolReplay => { |
| 377 | self.state.replay_count_this_turn = |
| 378 | self.state.replay_count_this_turn.saturating_add(1); |
| 379 | } |
| 380 | GuardrailAction::VerifyAndReplan => { |
| 381 | self.state.last_replan_turn = Some(turn_index); |
| 382 | } |
| 383 | GuardrailAction::NoIntervention => {} |
| 384 | } |
| 385 | } |
| 386 | |
| 387 | pub fn mark_replay_failed(&mut self, turn_index: u64) { |
| 388 | self.state.replay_disabled_turn = Some(turn_index); |
| 389 | } |
| 390 | |
| 391 | #[must_use] |
| 392 | pub fn last_snapshot(&self) -> Option<&CapacitySnapshot> { |
| 393 | self.last_snapshot.as_ref() |
| 394 | } |
| 395 | |
| 396 | fn observe(&mut self, input: CapacityObservationInput) -> Option<CapacitySnapshot> { |
| 397 | if !self.config.enabled { |
| 398 | return None; |
| 399 | } |
| 400 | |
| 401 | let context_used_ratio = input.context_used_ratio.clamp(0.0, 2.0); |
| 402 | let action_complexity_bits = log2_1p(input.action_count_this_turn); |
| 403 | let tool_complexity_bits = log2_1p(input.tool_calls_recent_window); |
| 404 | let ref_complexity_bits = log2_1p(input.unique_reference_ids_recent_window); |
| 405 | let context_pressure_bits = 6.0 * context_used_ratio; |
| 406 | |
| 407 | let h_hat = (0.35 * action_complexity_bits) |
| 408 | + (0.30 * tool_complexity_bits) |
| 409 | + (0.20 * ref_complexity_bits) |
| 410 | + (0.15 * context_pressure_bits); |
| 411 | let c_hat = self.model_prior(&input.model); |
| 412 | let slack = c_hat - h_hat; |
| 413 | |
| 414 | push_window(&mut self.slack_window, slack, self.config.profile_window); |
| 415 | push_window( |
| 416 | &mut self.recent_tool_counts, |
| 417 | input.tool_calls_recent_window, |
| 418 | self.config.profile_window, |
| 419 | ); |
| 420 | push_window( |
| 421 | &mut self.recent_ref_counts, |
| 422 | input.unique_reference_ids_recent_window, |
| 423 | self.config.profile_window, |
| 424 | ); |
| 425 | |
| 426 | let profile = compute_profile(&self.slack_window); |
| 427 | let z = (-1.65 * profile.final_slack) |
| 428 | + (-0.85 * profile.min_slack) |
| 429 | + (1.35 * profile.violation_ratio) |
| 430 | + (0.70 * profile.slack_volatility) |
| 431 | + (0.28 * profile.slack_drop) |
| 432 | - 0.12; |
| 433 | let p_fail = sigmoid(z).clamp(0.0, 1.0); |
| 434 | let risk_band = if p_fail <= self.config.low_risk_max { |
| 435 | RiskBand::Low |
| 436 | } else if p_fail <= self.config.medium_risk_max { |
| 437 | RiskBand::Medium |
| 438 | } else { |
| 439 | RiskBand::High |
| 440 | }; |
| 441 | let severe = profile.min_slack <= self.config.severe_min_slack |
| 442 | || profile.violation_ratio >= self.config.severe_violation_ratio; |
| 443 | |
| 444 | let snapshot = CapacitySnapshot { |
| 445 | turn_index: input.turn_index, |
| 446 | h_hat, |
| 447 | c_hat, |
| 448 | slack, |
| 449 | profile, |
| 450 | p_fail, |
| 451 | risk_band, |
| 452 | severe, |
| 453 | }; |
| 454 | self.last_snapshot = Some(snapshot.clone()); |
| 455 | Some(snapshot) |
| 456 | } |
| 457 | |
| 458 | fn model_prior(&self, model: &str) -> f64 { |
| 459 | let normalized = normalize_model_prior_key(model); |
| 460 | self.config |
| 461 | .model_priors |
| 462 | .get(normalized) |
| 463 | .copied() |
| 464 | .unwrap_or(self.config.fallback_default) |
| 465 | } |
| 466 | } |
| 467 | |
| 468 | /// Pure policy mapping for snapshot -> action. |
| 469 | #[must_use] |
| 470 | pub fn decide_policy( |
| 471 | _config: &CapacityControllerConfig, |
| 472 | snapshot: &CapacitySnapshot, |
| 473 | ) -> GuardrailAction { |
| 474 | match snapshot.risk_band { |
| 475 | RiskBand::Low => GuardrailAction::NoIntervention, |
| 476 | RiskBand::Medium => GuardrailAction::TargetedContextRefresh, |
| 477 | RiskBand::High if snapshot.severe => GuardrailAction::VerifyAndReplan, |
| 478 | RiskBand::High => GuardrailAction::VerifyWithToolReplay, |
| 479 | } |
| 480 | } |
| 481 | |
| 482 | fn normalize_model_prior_key(model: &str) -> &str { |
| 483 | // Strip optional "deepseek-ai/" NIM namespace prefix before pattern matching. |
| 484 | let model = model.strip_prefix("deepseek-ai/").unwrap_or(model); |
| 485 | let lower = model.to_ascii_lowercase(); |
| 486 | // V4 variants must be checked before the generic V3/chat/reasoner branches |
| 487 | // because those branches do not contain "v4" tokens and the ordering prevents |
| 488 | // accidental cross-matches. |
| 489 | if lower.contains("v4-pro") || lower.contains("v4_pro") { |
| 490 | "deepseek_v4_pro" |
| 491 | } else if lower.contains("v4-flash") || lower.contains("v4_flash") { |
| 492 | "deepseek_v4_flash" |
| 493 | } else if lower.contains("reasoner") || lower.contains("r1") { |
| 494 | "deepseek_v3_2_reasoner" |
| 495 | } else if lower.contains("chat") || lower.contains("v3") { |
| 496 | "deepseek_v3_2_chat" |
| 497 | } else { |
| 498 | "fallback_default" |
| 499 | } |
| 500 | } |
| 501 | |
| 502 | fn log2_1p(v: usize) -> f64 { |
| 503 | (1.0 + (v as f64)).log2() |
| 504 | } |
| 505 | |
| 506 | fn push_window<T>(window: &mut VecDeque<T>, value: T, max_len: usize) { |
| 507 | window.push_back(value); |
| 508 | while window.len() > max_len { |
| 509 | window.pop_front(); |
| 510 | } |
| 511 | } |
| 512 | |
| 513 | fn compute_profile(window: &VecDeque<f64>) -> DynamicSlackProfile { |
| 514 | if window.is_empty() { |
| 515 | return DynamicSlackProfile::default(); |
| 516 | } |
| 517 | |
| 518 | let values: Vec<f64> = window.iter().copied().collect(); |
| 519 | let final_slack = *values.last().unwrap_or(&0.0); |
| 520 | let min_slack = values.iter().copied().fold(f64::INFINITY, f64::min); |
| 521 | let violations = values.iter().filter(|v| **v <= 0.0).count() as f64; |
| 522 | let violation_ratio = violations / (values.len() as f64); |
| 523 | |
| 524 | let deltas: Vec<f64> = values.windows(2).map(|w| w[1] - w[0]).collect(); |
| 525 | let slack_drop = if values.len() >= 2 { |
| 526 | (values[values.len() - 2] - values[values.len() - 1]).max(0.0) |
| 527 | } else { |
| 528 | 0.0 |
| 529 | }; |
| 530 | |
| 531 | let slack_volatility = if deltas.is_empty() { |
| 532 | 0.0 |
| 533 | } else { |
| 534 | let mean = deltas.iter().sum::<f64>() / (deltas.len() as f64); |
| 535 | let var = deltas |
| 536 | .iter() |
| 537 | .map(|delta| { |
| 538 | let centered = *delta - mean; |
| 539 | centered * centered |
| 540 | }) |
| 541 | .sum::<f64>() |
| 542 | / (deltas.len() as f64); |
| 543 | var.sqrt() |
| 544 | }; |
| 545 | |
| 546 | DynamicSlackProfile { |
| 547 | final_slack, |
| 548 | min_slack, |
| 549 | violation_ratio, |
| 550 | slack_volatility, |
| 551 | slack_drop, |
| 552 | } |
| 553 | } |
| 554 | |
| 555 | fn sigmoid(z: f64) -> f64 { |
| 556 | if z >= 0.0 { |
| 557 | let ez = (-z).exp(); |
| 558 | 1.0 / (1.0 + ez) |
| 559 | } else { |
| 560 | let ez = z.exp(); |
| 561 | ez / (1.0 + ez) |
| 562 | } |
| 563 | } |
| 564 | |
| 565 | #[cfg(test)] |
| 566 | mod tests { |
| 567 | use super::*; |
| 568 | |
| 569 | fn make_snapshot(p_fail: f64, severe: bool, risk_band: RiskBand) -> CapacitySnapshot { |
| 570 | CapacitySnapshot { |
| 571 | turn_index: 3, |
| 572 | h_hat: 1.0, |
| 573 | c_hat: 3.8, |
| 574 | slack: 2.8, |
| 575 | profile: DynamicSlackProfile { |
| 576 | final_slack: 2.8, |
| 577 | min_slack: if severe { -0.5 } else { 0.2 }, |
| 578 | violation_ratio: if severe { 0.6 } else { 0.1 }, |
| 579 | slack_volatility: 0.2, |
| 580 | slack_drop: 0.1, |
| 581 | }, |
| 582 | p_fail, |
| 583 | risk_band, |
| 584 | severe, |
| 585 | } |
| 586 | } |
| 587 | |
| 588 | #[test] |
| 589 | fn low_risk_maps_to_no_intervention() { |
| 590 | let cfg = CapacityControllerConfig::default(); |
| 591 | let snap = make_snapshot(0.2, false, RiskBand::Low); |
| 592 | assert_eq!(decide_policy(&cfg, &snap), GuardrailAction::NoIntervention); |
| 593 | } |
| 594 | |
| 595 | #[test] |
| 596 | fn medium_risk_maps_to_refresh() { |
| 597 | let cfg = CapacityControllerConfig::default(); |
| 598 | let snap = make_snapshot(0.5, false, RiskBand::Medium); |
| 599 | assert_eq!( |
| 600 | decide_policy(&cfg, &snap), |
| 601 | GuardrailAction::TargetedContextRefresh |
| 602 | ); |
| 603 | } |
| 604 | |
| 605 | #[test] |
| 606 | fn high_non_severe_maps_to_replay() { |
| 607 | let cfg = CapacityControllerConfig::default(); |
| 608 | let snap = make_snapshot(0.8, false, RiskBand::High); |
| 609 | assert_eq!( |
| 610 | decide_policy(&cfg, &snap), |
| 611 | GuardrailAction::VerifyWithToolReplay |
| 612 | ); |
| 613 | } |
| 614 | |
| 615 | #[test] |
| 616 | fn high_severe_maps_to_replan() { |
| 617 | let cfg = CapacityControllerConfig::default(); |
| 618 | let snap = make_snapshot(0.9, true, RiskBand::High); |
| 619 | assert_eq!(decide_policy(&cfg, &snap), GuardrailAction::VerifyAndReplan); |
| 620 | } |
| 621 | |
| 622 | /// v0.8.11 flipped the default to `enabled = false`. The controller's |
| 623 | /// observe / decide methods early-return when disabled — opt-in only. |
| 624 | #[test] |
| 625 | fn default_controller_is_disabled_and_skips_observations() { |
| 626 | let cfg = CapacityControllerConfig::default(); |
| 627 | assert!(!cfg.enabled); |
| 628 | |
| 629 | let mut controller = CapacityController::new(cfg); |
| 630 | let snapshot = controller.observe_pre_turn(CapacityObservationInput { |
| 631 | turn_index: 1, |
| 632 | model: "deepseek-v4-pro".to_string(), |
| 633 | action_count_this_turn: 10, |
| 634 | tool_calls_recent_window: 10, |
| 635 | unique_reference_ids_recent_window: 10, |
| 636 | context_used_ratio: 0.95, |
| 637 | }); |
| 638 | |
| 639 | // With enabled=false, observe_pre_turn returns None. |
| 640 | assert!(snapshot.is_none()); |
| 641 | } |
| 642 | |
| 643 | /// Opting in via `capacity.enabled = true` re-arms the controller — |
| 644 | /// observations produce snapshots, decisions can fire interventions. |
| 645 | #[test] |
| 646 | fn opt_in_controller_observes_and_decides() { |
| 647 | let cfg = CapacityControllerConfig { |
| 648 | enabled: true, |
| 649 | ..Default::default() |
| 650 | }; |
| 651 | |
| 652 | let mut controller = CapacityController::new(cfg); |
| 653 | let snapshot = controller.observe_pre_turn(CapacityObservationInput { |
| 654 | turn_index: 1, |
| 655 | model: "deepseek-v4-pro".to_string(), |
| 656 | action_count_this_turn: 10, |
| 657 | tool_calls_recent_window: 10, |
| 658 | unique_reference_ids_recent_window: 10, |
| 659 | context_used_ratio: 0.95, |
| 660 | }); |
| 661 | |
| 662 | assert!(snapshot.is_some()); |
| 663 | let snap = snapshot.unwrap(); |
| 664 | assert_eq!(snap.turn_index, 1); |
| 665 | assert!(snap.p_fail > 0.0); |
| 666 | } |
| 667 | |
| 668 | #[test] |
| 669 | fn app_config_without_capacity_uses_default_disabled() { |
| 670 | let cfg = CapacityControllerConfig::from_app_config(&crate::config::Config::default()); |
| 671 | // v0.8.11: default is disabled. No capacity section in config |
| 672 | // means the controller stays inert; users opt in deliberately. |
| 673 | assert!(!cfg.enabled); |
| 674 | assert_eq!(cfg.low_risk_max, 0.50); |
| 675 | assert_eq!(cfg.refresh_cooldown_turns, 6); |
| 676 | assert_eq!(cfg.min_turns_before_guardrail, 4); |
| 677 | assert_eq!(cfg.model_priors.get("deepseek_v4_pro"), Some(&3.5)); |
| 678 | assert_eq!(cfg.model_priors.get("deepseek_v4_flash"), Some(&4.2)); |
| 679 | } |
| 680 | |
| 681 | #[test] |
| 682 | fn normalize_v4_pro_variants() { |
| 683 | assert_eq!( |
| 684 | normalize_model_prior_key("deepseek-v4-pro"), |
| 685 | "deepseek_v4_pro" |
| 686 | ); |
| 687 | assert_eq!( |
| 688 | normalize_model_prior_key("deepseek-v4_pro"), |
| 689 | "deepseek_v4_pro" |
| 690 | ); |
| 691 | assert_eq!( |
| 692 | normalize_model_prior_key("deepseek-ai/deepseek-v4-pro"), |
| 693 | "deepseek_v4_pro" |
| 694 | ); |
| 695 | assert_eq!( |
| 696 | normalize_model_prior_key("deepseek-ai/deepseek-v4_pro"), |
| 697 | "deepseek_v4_pro" |
| 698 | ); |
| 699 | } |
| 700 | |
| 701 | #[test] |
| 702 | fn normalize_v4_flash_variants() { |
| 703 | assert_eq!( |
| 704 | normalize_model_prior_key("deepseek-v4-flash"), |
| 705 | "deepseek_v4_flash" |
| 706 | ); |
| 707 | assert_eq!( |
| 708 | normalize_model_prior_key("deepseek-v4_flash"), |
| 709 | "deepseek_v4_flash" |
| 710 | ); |
| 711 | assert_eq!( |
| 712 | normalize_model_prior_key("deepseek-ai/deepseek-v4-flash"), |
| 713 | "deepseek_v4_flash" |
| 714 | ); |
| 715 | assert_eq!( |
| 716 | normalize_model_prior_key("deepseek-ai/deepseek-v4_flash"), |
| 717 | "deepseek_v4_flash" |
| 718 | ); |
| 719 | } |
| 720 | |
| 721 | #[test] |
| 722 | fn normalize_v4_and_fallback_prior_keys() { |
| 723 | assert_eq!( |
| 724 | normalize_model_prior_key("deepseek-v4-pro"), |
| 725 | "deepseek_v4_pro" |
| 726 | ); |
| 727 | assert_eq!( |
| 728 | normalize_model_prior_key("deepseek-v4-flash"), |
| 729 | "deepseek_v4_flash" |
| 730 | ); |
| 731 | assert_eq!( |
| 732 | normalize_model_prior_key("unknown-model"), |
| 733 | "fallback_default" |
| 734 | ); |
| 735 | } |
| 736 | |
| 737 | #[test] |
| 738 | fn v4_priors_loaded_into_default_config() { |
| 739 | let cfg = CapacityControllerConfig::default(); |
| 740 | assert_eq!(cfg.model_priors.get("deepseek_v4_pro").copied(), Some(3.5)); |
| 741 | assert_eq!( |
| 742 | cfg.model_priors.get("deepseek_v4_flash").copied(), |
| 743 | Some(4.2) |
| 744 | ); |
| 745 | } |
| 746 | |
| 747 | #[test] |
| 748 | fn cooldown_blocks_repeated_action() { |
| 749 | // Capacity controller is opt-in (off by default since v0.6.2). This |
| 750 | // test exercises the cooldown logic, so explicitly enable it. |
| 751 | let config = CapacityControllerConfig { |
| 752 | enabled: true, |
| 753 | ..CapacityControllerConfig::default() |
| 754 | }; |
| 755 | let mut controller = CapacityController::new(config); |
| 756 | let turn_index = 5; |
| 757 | controller.mark_turn_start(turn_index); |
| 758 | controller.mark_intervention_applied(turn_index, GuardrailAction::TargetedContextRefresh); |
| 759 | |
| 760 | let snapshot = make_snapshot(0.5, false, RiskBand::Medium); |
| 761 | let decision = controller.decide(turn_index + 1, Some(&snapshot)); |
| 762 | assert_eq!(decision.action, GuardrailAction::NoIntervention); |
| 763 | assert!(decision.cooldown_blocked); |
| 764 | } |
| 765 | |
| 766 | /// Hot-path microbench for `compute_profile`. Run with: |
| 767 | /// |
| 768 | /// ```text |
| 769 | /// cargo test -p deepseek-tui --release capacity::tests::bench_compute_profile -- --ignored --nocapture |
| 770 | /// ``` |
| 771 | /// |
| 772 | /// Establishes a baseline cost so we can detect regressions when the |
| 773 | /// observation cadence is high (50+ message turns × per-step calls). Adds |
| 774 | /// no dev-deps; we measure with `Instant` and print rather than gating CI. |
| 775 | #[test] |
| 776 | #[ignore] |
| 777 | fn bench_compute_profile() { |
| 778 | use std::time::Instant; |
| 779 | |
| 780 | for &window_len in &[16usize, 64, 256, 1024] { |
| 781 | let mut window: VecDeque<f64> = VecDeque::with_capacity(window_len); |
| 782 | for i in 0..window_len { |
| 783 | #[allow(clippy::cast_precision_loss)] |
| 784 | window.push_back((i as f64).sin() * 0.5); |
| 785 | } |
| 786 | |
| 787 | let iters = 100_000usize; |
| 788 | let start = Instant::now(); |
| 789 | for _ in 0..iters { |
| 790 | let profile = compute_profile(&window); |
| 791 | std::hint::black_box(profile); |
| 792 | } |
| 793 | let elapsed = start.elapsed(); |
| 794 | let per_call_ns = elapsed.as_nanos() as f64 / iters as f64; |
| 795 | println!( |
| 796 | "compute_profile window={window_len:>4} total={:?} per-call={per_call_ns:>8.0}ns", |
| 797 | elapsed |
| 798 | ); |
| 799 | } |
| 800 | } |
| 801 | } |
| 802 |