1#![allow(clippy::unwrap_used)]
2use anyhow::Result;
3use rand::RngExt;
4use serde::{Deserialize, Serialize};
5use std::collections::{HashMap, VecDeque};
6use std::path::{Path, PathBuf};
7use std::sync::Arc;
8use tokio::sync::Mutex;
9use tracing::info;
10
11use crate::state::persistence::{
12 load_bash_state as load_state_file, save_bash_state as save_state_file, BashStateSnapshot,
13};
14use crate::state::pty::PtyShell;
15use crate::types::{
16 AllowedCommands, AllowedGlobs, BashCommandMode, BashMode, FileEditMode, Modes, WriteIfEmptyMode,
17};
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct FileWhitelistData {
21 pub file_hash: String,
22 pub line_ranges_read: Vec<(usize, usize)>,
23 pub total_lines: usize,
24}
25
26fn clamped_sorted(ranges: &[(usize, usize)], total_lines: usize) -> Vec<(usize, usize)> {
29 let mut v: Vec<(usize, usize)> = ranges
30 .iter()
31 .map(|&(s, e)| (s.max(1), e.min(total_lines)))
32 .filter(|&(s, e)| s <= e)
33 .collect();
34 v.sort_unstable();
35 v
36}
37
38fn covered_line_count(ranges: &[(usize, usize)], total_lines: usize) -> usize {
41 let mut covered = 0usize;
42 let mut last_end = 0usize; for (s, e) in clamped_sorted(ranges, total_lines) {
44 let s = s.max(last_end + 1);
45 if s <= e {
46 covered += e - s + 1;
47 last_end = e;
48 }
49 }
50 covered
51}
52
53impl FileWhitelistData {
54 pub fn new(
55 file_hash: String,
56 line_ranges_read: Vec<(usize, usize)>,
57 total_lines: usize,
58 ) -> Self {
59 let mut data = Self { file_hash, line_ranges_read: Vec::new(), total_lines };
60 data.merge_ranges(line_ranges_read);
61 data
62 }
63
64 pub fn is_read_enough(&self) -> bool {
65 self.get_percentage_read() >= 99.0
66 }
67
68 #[allow(clippy::cast_precision_loss)]
69 pub fn get_percentage_read(&self) -> f64 {
70 if self.total_lines == 0 {
71 return 100.0;
72 }
73 let covered = covered_line_count(&self.line_ranges_read, self.total_lines);
77 (covered as f64 / self.total_lines as f64) * 100.0
78 }
79
80 pub fn get_unread_ranges(&self) -> Vec<(usize, usize)> {
81 if self.total_lines == 0 {
82 return vec![];
83 }
84 let sorted = clamped_sorted(&self.line_ranges_read, self.total_lines);
85 let mut unread = vec![];
86 let mut next = 1usize; for (s, e) in sorted {
88 if s > next {
89 unread.push((next, s - 1));
90 }
91 next = next.max(e.saturating_add(1));
92 }
93 if next <= self.total_lines {
94 unread.push((next, self.total_lines));
95 }
96 unread
97 }
98
99 pub fn add_range(&mut self, start: usize, end: usize) {
104 self.merge_ranges(std::iter::once((start, end)));
105 }
106
107 pub fn merge_ranges(&mut self, new: impl IntoIterator<Item = (usize, usize)>) {
110 self.line_ranges_read.extend(new);
111 self.line_ranges_read.retain(|(s, e)| s <= e);
112 self.line_ranges_read.sort_unstable();
113 let mut merged: Vec<(usize, usize)> = Vec::with_capacity(self.line_ranges_read.len());
114 for (s, e) in self.line_ranges_read.drain(..) {
115 match merged.last_mut() {
116 Some(last) if s <= last.1.saturating_add(1) => last.1 = last.1.max(e),
117 _ => merged.push((s, e)),
118 }
119 }
120 self.line_ranges_read = merged;
121 }
122
123 pub fn needs_more_reading(&self) -> bool {
124 !self.is_read_enough()
125 }
126
127 pub fn covers_range(&self, start: usize, end: usize) -> bool {
130 start <= end
131 && self
132 .line_ranges_read
133 .iter()
134 .any(|(covered_start, covered_end)| *covered_start <= start && *covered_end >= end)
135 }
136}
137
138const EDIT_CHECKPOINT_CAP: usize = 10;
141
142const EDIT_CHECKPOINT_MAX_CONTENT_BYTES: usize = 1_000_000;
147
148const MAX_WHITELIST_FILES: usize = 1_024;
153
154const IMAGE_DELIVERY_CACHE_CAP: usize = 32;
158
159const EDIT_MUTATION_RECEIPT_CAP: usize = 64;
163
164#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
165pub(crate) struct EditMutationPostcondition {
166 pub path: String,
167 pub sha256: String,
168}
169
170#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
171pub(crate) struct EditMutationReceipt {
172 pub fingerprint: String,
173 pub receipt_id: String,
174 pub tool: String,
175 pub status: String,
176 pub committed_at_unix_ms: u64,
177 pub postconditions: Vec<EditMutationPostcondition>,
178 #[serde(default)]
179 pub persisted: bool,
180 #[serde(default)]
181 pub verification_pending: bool,
182}
183
184#[derive(Debug, Clone)]
188pub struct EditCheckpoint {
189 pub file_path_str: String,
191 pub path: PathBuf,
192 pub prior_content: String,
194 pub prior_whitelist: Option<FileWhitelistData>,
197}
198
199#[derive(Debug, Clone)]
200pub struct BashState {
201 pub cwd: PathBuf,
202 pub workspace_root: PathBuf,
203 pub current_thread_id: String,
204 pub mode: Modes,
205 pub bash_command_mode: BashCommandMode,
206 pub file_edit_mode: FileEditMode,
207 pub write_if_empty_mode: WriteIfEmptyMode,
208 pub whitelist_for_overwrite: HashMap<String, FileWhitelistData>,
209 whitelist_recency: VecDeque<String>,
213 pub pty_shell: Arc<Mutex<Option<PtyShell>>>,
214 pub foreground_command_gate: Arc<Mutex<()>>,
217 pub initialized: bool,
218 pub edit_checkpoints: VecDeque<EditCheckpoint>,
222 pub derived_code_map_usage: crate::utils::agent_temp::DerivedCodeMapUsage,
225 image_deliveries: VecDeque<String>,
226 edit_mutation_receipts: VecDeque<EditMutationReceipt>,
227}
228
229impl Default for BashState {
230 fn default() -> Self {
231 Self::new()
232 }
233}
234
235impl BashState {
236 pub fn new() -> Self {
237 let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/tmp"));
238 Self {
239 cwd: cwd.clone(),
240 workspace_root: cwd,
241 current_thread_id: generate_thread_id(),
242 mode: Modes::Wcgw,
243 bash_command_mode: BashCommandMode {
244 bash_mode: BashMode::NormalMode,
245 allowed_commands: AllowedCommands::All("all".to_string()),
246 },
247 file_edit_mode: FileEditMode { allowed_globs: AllowedGlobs::All("all".to_string()) },
248 write_if_empty_mode: WriteIfEmptyMode {
249 allowed_globs: AllowedGlobs::All("all".to_string()),
250 },
251 whitelist_for_overwrite: HashMap::new(),
252 whitelist_recency: VecDeque::new(),
253 pty_shell: Arc::new(Mutex::new(None)),
254 foreground_command_gate: Arc::new(Mutex::new(())),
255 initialized: false,
256 edit_checkpoints: VecDeque::new(),
257 derived_code_map_usage: crate::utils::agent_temp::DerivedCodeMapUsage::default(),
258 image_deliveries: VecDeque::new(),
259 edit_mutation_receipts: VecDeque::new(),
260 }
261 }
262
263 pub(crate) fn edit_mutation_receipt(
264 &mut self,
265 fingerprint: &str,
266 now_unix_ms: u64,
267 ttl_ms: u64,
268 ) -> Option<EditMutationReceipt> {
269 self.edit_mutation_receipts
270 .retain(|receipt| now_unix_ms.saturating_sub(receipt.committed_at_unix_ms) <= ttl_ms);
271 let index = self
272 .edit_mutation_receipts
273 .iter()
274 .position(|receipt| receipt.fingerprint == fingerprint)?;
275 let receipt = self.edit_mutation_receipts.remove(index)?;
276 self.edit_mutation_receipts.push_back(receipt.clone());
277 Some(receipt)
278 }
279
280 pub(crate) fn record_edit_mutation_receipt(&mut self, receipt: EditMutationReceipt) {
281 if let Some(index) = self
282 .edit_mutation_receipts
283 .iter()
284 .position(|existing| existing.fingerprint == receipt.fingerprint)
285 {
286 self.edit_mutation_receipts.remove(index);
287 }
288 self.edit_mutation_receipts.push_back(receipt);
289 while self.edit_mutation_receipts.len() > EDIT_MUTATION_RECEIPT_CAP {
290 self.edit_mutation_receipts.pop_front();
291 }
292 }
293
294 pub(crate) fn mark_edit_mutation_receipt_volatile(&mut self, fingerprint: &str) {
295 if let Some(receipt) = self
296 .edit_mutation_receipts
297 .iter_mut()
298 .find(|receipt| receipt.fingerprint == fingerprint)
299 {
300 receipt.persisted = false;
301 }
302 }
303
304 pub(crate) fn image_was_delivered(&mut self, fingerprint: &str) -> bool {
308 let Some(index) = self.image_deliveries.iter().position(|item| item == fingerprint) else {
309 return false;
310 };
311 if let Some(item) = self.image_deliveries.remove(index) {
312 self.image_deliveries.push_back(item);
313 }
314 true
315 }
316
317 pub(crate) fn record_image_delivery(&mut self, fingerprint: String) {
319 if let Some(index) = self.image_deliveries.iter().position(|item| item == &fingerprint) {
320 self.image_deliveries.remove(index);
321 }
322 self.image_deliveries.push_back(fingerprint);
323 while self.image_deliveries.len() > IMAGE_DELIVERY_CACHE_CAP {
324 self.image_deliveries.pop_front();
325 }
326 }
327
328 pub fn push_edit_checkpoint(&mut self, checkpoint: EditCheckpoint) {
331 if checkpoint.prior_content.len() > EDIT_CHECKPOINT_MAX_CONTENT_BYTES {
332 info!(
333 file = %checkpoint.file_path_str,
334 "UndoEdit: not checkpointing a file over 1 MB (too large to hold in memory)"
335 );
336 return;
337 }
338 self.edit_checkpoints.push_back(checkpoint);
339 while self.edit_checkpoints.len() > EDIT_CHECKPOINT_CAP {
340 self.edit_checkpoints.pop_front();
341 }
342 }
343
344 pub fn pop_edit_checkpoint_for(&mut self, file_path_str: &str) -> Option<EditCheckpoint> {
348 let index =
349 self.edit_checkpoints.iter().rposition(|cp| cp.file_path_str == file_path_str)?;
350 self.edit_checkpoints.remove(index)
351 }
352
353 pub fn latest_edit_checkpoint_for(&self, file_path_str: &str) -> Option<EditCheckpoint> {
356 self.edit_checkpoints
357 .iter()
358 .rev()
359 .find(|checkpoint| checkpoint.file_path_str == file_path_str)
360 .cloned()
361 }
362
363 pub fn record_read_coverage(
368 &mut self,
369 path: &str,
370 ranges: impl IntoIterator<Item = (usize, usize)>,
371 file_hash: String,
372 total_lines: usize,
373 ) {
374 let ranges = ranges.into_iter().collect::<Vec<_>>();
375 match self.whitelist_for_overwrite.get_mut(path) {
376 Some(existing)
377 if existing.file_hash == file_hash && existing.total_lines == total_lines =>
378 {
379 existing.merge_ranges(ranges);
380 }
381 Some(existing) => {
382 *existing = FileWhitelistData::new(file_hash, ranges, total_lines);
383 }
384 None => {
385 self.whitelist_for_overwrite.insert(
386 path.to_string(),
387 FileWhitelistData::new(file_hash, ranges, total_lines),
388 );
389 }
390 }
391 self.touch_whitelist(path);
392 self.enforce_whitelist_cap();
393 }
394
395 pub fn set_whitelist_entry(&mut self, path: &str, entry: FileWhitelistData) {
397 self.whitelist_for_overwrite.insert(path.to_string(), entry);
398 self.touch_whitelist(path);
399 self.enforce_whitelist_cap();
400 }
401
402 pub fn remove_whitelist_entry(&mut self, path: &str) -> Option<FileWhitelistData> {
404 self.whitelist_recency.retain(|candidate| candidate != path);
405 self.whitelist_for_overwrite.remove(path)
406 }
407
408 fn touch_whitelist(&mut self, path: &str) {
409 self.whitelist_recency.retain(|candidate| candidate != path);
410 self.whitelist_recency.push_back(path.to_string());
411 }
412
413 fn rebuild_whitelist_recency(&mut self) {
414 let mut paths = self.whitelist_for_overwrite.keys().cloned().collect::<Vec<_>>();
415 paths.sort_unstable();
416 self.whitelist_recency = paths.into();
417 self.enforce_whitelist_cap();
418 }
419
420 fn enforce_whitelist_cap(&mut self) {
421 self.whitelist_recency.retain(|path| self.whitelist_for_overwrite.contains_key(path));
422 while self.whitelist_for_overwrite.len() > MAX_WHITELIST_FILES {
423 let victim = self.whitelist_recency.pop_front().or_else(|| {
424 self.whitelist_for_overwrite.keys().min().cloned()
427 });
428 let Some(victim) = victim else { break };
429 self.whitelist_for_overwrite.remove(&victim);
430 }
431 }
432
433 pub async fn init_pty_shell(&mut self) -> Result<()> {
434 let cwd = self.cwd.clone();
435 let workspace_root = self.workspace_root.clone();
436 let temporary_artifact_dir =
437 crate::utils::agent_temp::session_info(&workspace_root, &self.current_thread_id)
438 .directory;
439 let restricted = self.bash_command_mode.bash_mode == BashMode::RestrictedMode;
440 let shell = tokio::task::spawn_blocking(move || {
444 PtyShell::new_with_agent_paths(
445 &cwd,
446 restricted,
447 Some(&workspace_root),
448 Some(&temporary_artifact_dir),
449 )
450 })
451 .await
452 .map_err(|e| {
453 crate::errors::WinxError::ShellInitializationError(format!("PTY init task failed: {e}"))
454 })??;
455 *self.pty_shell.lock().await = Some(shell);
456 Ok(())
457 }
458
459 pub fn update_cwd(&mut self, path: &Path) -> Result<()> {
460 self.cwd = path.to_path_buf();
461 Ok(())
462 }
463
464 pub fn update_workspace_root(&mut self, path: &Path) -> Result<()> {
465 self.workspace_root = path.to_path_buf();
466 Ok(())
467 }
468
469 pub fn is_command_allowed(&self, command: &str) -> bool {
470 if self.mode == Modes::Architect {
471 crate::utils::bash_parser::is_architect_command_allowed(command)
472 } else {
473 self.bash_command_mode.allowed_commands.is_allowed(command)
474 }
475 }
476
477 pub fn is_file_edit_allowed(&self, path: &str) -> bool {
478 self.file_edit_mode.allowed_globs.is_allowed(path)
479 }
480
481 pub fn is_file_write_allowed(&self, path: &str) -> bool {
482 self.write_if_empty_mode.allowed_globs.is_allowed(path)
483 }
484 pub fn save_state_to_disk(&self) -> Result<()> {
485 let snapshot = self.snapshot();
486 save_state_file(&self.current_thread_id, &snapshot)?;
487 Ok(())
488 }
489
490 pub fn snapshot(&self) -> BashStateSnapshot {
491 BashStateSnapshot::from_state(
492 &self.cwd.to_string_lossy(),
493 &self.workspace_root.to_string_lossy(),
494 self.mode,
495 &self.bash_command_mode,
496 &self.file_edit_mode,
497 &self.write_if_empty_mode,
498 &self.whitelist_for_overwrite,
499 &self.current_thread_id,
500 &self.edit_mutation_receipts,
501 )
502 }
503
504 pub fn apply_snapshot(&mut self, snapshot: &BashStateSnapshot) {
505 let (cwd, root, mode, bmode, emode, wmode, whitelist, tid) = snapshot.to_state_components();
506 let root = PathBuf::from(root);
507 let image_identity_changed = self.current_thread_id != tid || self.workspace_root != root;
508
509 self.cwd = PathBuf::from(cwd);
510 self.workspace_root = root;
511 self.mode = mode;
512 self.bash_command_mode = bmode;
513 self.file_edit_mode = emode;
514 self.write_if_empty_mode = wmode;
515 self.whitelist_for_overwrite = whitelist;
516 self.rebuild_whitelist_recency();
517 self.current_thread_id = tid;
518 self.edit_mutation_receipts = snapshot.edit_mutation_receipts.iter().cloned().collect();
519 if image_identity_changed {
520 self.image_deliveries.clear();
521 }
522 self.initialized = true;
523 }
524
525 pub fn load_state_from_disk(&mut self, thread_id: &str) -> Result<bool> {
526 if let Some(snapshot) = load_state_file(thread_id)? {
527 self.apply_snapshot(&snapshot);
528 Ok(true)
529 } else {
530 Ok(false)
531 }
532 }
533
534 pub fn new_with_thread_id(thread_id: Option<&str>) -> Self {
535 let mut state = Self::new();
536
537 if let Some(tid) = thread_id {
538 if !tid.is_empty() {
539 if let Ok(true) = state.load_state_from_disk(tid) {
540 info!("Loaded state for thread_id '{}'", tid);
541 } else {
542 state.current_thread_id = tid.to_string();
543 }
544 }
545 }
546
547 state
548 }
549}
550
551pub fn generate_thread_id() -> String {
552 let mut rng = rand::rng();
553 format!("tid_{:x}", rng.random::<u64>())
554}
555
556#[cfg(test)]
557mod whitelist_range_tests {
558 use super::{
559 BashState, EditMutationPostcondition, EditMutationReceipt, FileWhitelistData,
560 MAX_WHITELIST_FILES,
561 };
562
563 fn wl(ranges: &[(usize, usize)], total: usize) -> FileWhitelistData {
564 FileWhitelistData::new("h".to_string(), ranges.to_vec(), total)
565 }
566
567 #[test]
568 fn merge_collapses_overlap_and_adjacency() {
569 let w = wl(&[(4, 5), (1, 3), (8, 12), (7, 9)], 20);
571 assert_eq!(w.line_ranges_read, vec![(1, 5), (7, 12)]);
572 }
573
574 #[test]
575 fn re_reading_does_not_grow_unbounded() {
576 let mut w = wl(&[(1, 10)], 100);
577 for _ in 0..1000 {
578 w.merge_ranges(std::iter::once((1, 10)));
579 w.merge_ranges(std::iter::once((5, 15)));
580 }
581 assert_eq!(w.line_ranges_read, vec![(1, 15)]);
583 }
584
585 #[test]
586 fn percentage_counts_distinct_lines_with_overlap() {
587 let w = wl(&[(1, 5), (3, 8)], 10);
589 assert!((w.get_percentage_read() - 80.0).abs() < 1e-9);
590 assert!(wl(&[(1, 10)], 10).is_read_enough());
591 }
592
593 #[test]
594 fn unread_ranges_are_the_gaps() {
595 let w = wl(&[(2, 4), (7, 8)], 10);
597 assert_eq!(w.get_unread_ranges(), vec![(1, 1), (5, 6), (9, 10)]);
598 }
599
600 #[test]
601 fn out_of_range_entries_are_clamped() {
602 let w = wl(&[(0, 999)], 10);
604 assert!((w.get_percentage_read() - 100.0).abs() < 1e-9);
605 assert!(w.get_unread_ranges().is_empty());
606 }
607
608 #[test]
609 fn coverage_from_different_file_versions_is_not_combined() {
610 let mut state = BashState::new();
611 state.record_read_coverage("/workspace/file.rs", [(1, 50)], "old-hash".to_string(), 100);
612 state.record_read_coverage("/workspace/file.rs", [(51, 100)], "new-hash".to_string(), 100);
613
614 let coverage = &state.whitelist_for_overwrite["/workspace/file.rs"];
615 assert_eq!(coverage.file_hash, "new-hash");
616 assert_eq!(coverage.line_ranges_read, vec![(51, 100)]);
617 assert_eq!(coverage.get_unread_ranges(), vec![(1, 50)]);
618 }
619
620 #[test]
621 fn whitelist_is_lru_bounded_and_recent_entries_survive() {
622 let mut state = BashState::new();
623 for index in 0..MAX_WHITELIST_FILES {
624 let path = format!("/workspace/file-{index:04}.rs");
625 state.record_read_coverage(&path, [(1, 1)], format!("hash-{index}"), 1);
626 }
627
628 state.record_read_coverage("/workspace/file-0000.rs", [(1, 1)], "hash-0".to_string(), 1);
631 state.record_read_coverage("/workspace/newest.rs", [(1, 1)], "newest-hash".to_string(), 1);
632
633 assert_eq!(state.whitelist_for_overwrite.len(), MAX_WHITELIST_FILES);
634 assert!(state.whitelist_for_overwrite.contains_key("/workspace/file-0000.rs"));
635 assert!(!state.whitelist_for_overwrite.contains_key("/workspace/file-0001.rs"));
636 assert!(state.whitelist_for_overwrite.contains_key("/workspace/newest.rs"));
637 }
638
639 #[test]
640 fn image_delivery_cache_survives_same_session_snapshot_sync_only() {
641 let mut state = BashState::new();
642 state.workspace_root = "/workspace/one".into();
643 state.current_thread_id = "one".to_string();
644 state.record_image_delivery("fingerprint".to_string());
645
646 let same_session = state.snapshot();
647 state.apply_snapshot(&same_session);
648 assert!(state.image_was_delivered("fingerprint"));
649
650 let mut other = BashState::new();
651 other.workspace_root = "/workspace/two".into();
652 other.current_thread_id = "two".to_string();
653 state.apply_snapshot(&other.snapshot());
654 assert!(!state.image_was_delivered("fingerprint"));
655 }
656
657 #[test]
658 fn mutation_receipts_survive_snapshot_and_expire_closed() -> Result<(), &'static str> {
659 let mut state = BashState::new();
660 state.record_edit_mutation_receipt(EditMutationReceipt {
661 fingerprint: "fingerprint".to_string(),
662 receipt_id: "edit_receipt".to_string(),
663 tool: "FileWriteOrEdit".to_string(),
664 status: "completed".to_string(),
665 committed_at_unix_ms: 1_000,
666 postconditions: vec![EditMutationPostcondition {
667 path: "/workspace/file.rs".to_string(),
668 sha256: "hash".to_string(),
669 }],
670 persisted: true,
671 verification_pending: false,
672 });
673
674 let snapshot = state.snapshot();
675 let mut restored = BashState::new();
676 restored.apply_snapshot(&snapshot);
677 let receipt = restored
678 .edit_mutation_receipt("fingerprint", 1_500, 1_000)
679 .ok_or("fresh persisted receipt missing")?;
680 assert_eq!(receipt.receipt_id, "edit_receipt");
681 assert!(receipt.persisted);
682 assert!(restored.edit_mutation_receipt("fingerprint", 3_000, 1_000).is_none());
683 Ok(())
684 }
685}