1#![allow(clippy::unwrap_used)]
2use anyhow::Result;
3use rand::RngExt;
4use serde::{Deserialize, Serialize};
5use sha2::{Digest, Sha256};
6use std::collections::HashMap;
7use std::path::{Path, PathBuf};
8use std::sync::Arc;
9use tokio::sync::Mutex;
10use tracing::info;
11
12use crate::state::persistence::{
13 delete_bash_state as delete_state_file, load_bash_state as load_state_file,
14 save_bash_state as save_state_file, BashStateSnapshot,
15};
16use crate::state::pty::PtyShell;
17use crate::types::{
18 AllowedCommands, AllowedGlobs, BashCommandMode, BashMode, FileEditMode, Modes, WriteIfEmptyMode,
19};
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct FileWhitelistData {
23 pub file_hash: String,
24 pub line_ranges_read: Vec<(usize, usize)>,
25 pub total_lines: usize,
26}
27
28impl FileWhitelistData {
29 pub fn new(
30 file_hash: String,
31 line_ranges_read: Vec<(usize, usize)>,
32 total_lines: usize,
33 ) -> Self {
34 Self { file_hash, line_ranges_read, total_lines }
35 }
36
37 pub fn is_read_enough(&self) -> bool {
38 self.get_percentage_read() >= 99.0
39 }
40
41 pub fn get_percentage_read(&self) -> f64 {
42 if self.total_lines == 0 {
43 return 100.0;
44 }
45 let mut lines_read = std::collections::HashSet::new();
46 for (start, end) in &self.line_ranges_read {
47 for line in *start..=*end {
48 lines_read.insert(line);
49 }
50 }
51 (lines_read.len() as f64 / self.total_lines as f64) * 100.0
52 }
53
54 pub fn get_unread_ranges(&self) -> Vec<(usize, usize)> {
55 if self.total_lines == 0 {
56 return vec![];
57 }
58 let mut lines_read = std::collections::HashSet::new();
59 for (start, end) in &self.line_ranges_read {
60 for line in *start..=*end {
61 lines_read.insert(line);
62 }
63 }
64 let mut unread = vec![];
65 let mut start_range = None;
66 for i in 1..=self.total_lines {
67 if !lines_read.contains(&i) {
68 if start_range.is_none() {
69 start_range = Some(i);
70 }
71 } else if let Some(start) = start_range {
72 unread.push((start, i - 1));
73 start_range = None;
74 }
75 }
76 if let Some(start) = start_range {
77 unread.push((start, self.total_lines));
78 }
79 unread
80 }
81
82 pub fn add_range(&mut self, start: usize, end: usize) {
83 self.line_ranges_read.push((start, end));
84 }
85
86 pub fn get_read_error_message(&self, file_path: &Path) -> String {
87 format!(
88 "File {} needs more reading. Coverage: {:.1}%",
89 file_path.display(),
90 self.get_percentage_read()
91 )
92 }
93
94 pub fn needs_more_reading(&self) -> bool {
95 !self.is_read_enough()
96 }
97}
98
99#[derive(Debug, Clone)]
100pub struct BashState {
101 pub cwd: PathBuf,
102 pub workspace_root: PathBuf,
103 pub current_thread_id: String,
104 pub mode: Modes,
105 pub bash_command_mode: BashCommandMode,
106 pub file_edit_mode: FileEditMode,
107 pub write_if_empty_mode: WriteIfEmptyMode,
108 pub whitelist_for_overwrite: HashMap<String, FileWhitelistData>,
109 pub pty_shell: Arc<Mutex<Option<PtyShell>>>,
110 pub initialized: bool,
111}
112
113impl Default for BashState {
114 fn default() -> Self {
115 Self::new()
116 }
117}
118
119impl BashState {
120 pub fn new() -> Self {
121 let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/tmp"));
122 Self {
123 cwd: cwd.clone(),
124 workspace_root: cwd,
125 current_thread_id: generate_thread_id(),
126 mode: Modes::Wcgw,
127 bash_command_mode: BashCommandMode {
128 bash_mode: BashMode::NormalMode,
129 allowed_commands: AllowedCommands::All("all".to_string()),
130 },
131 file_edit_mode: FileEditMode { allowed_globs: AllowedGlobs::All("all".to_string()) },
132 write_if_empty_mode: WriteIfEmptyMode {
133 allowed_globs: AllowedGlobs::All("all".to_string()),
134 },
135 whitelist_for_overwrite: HashMap::new(),
136 pty_shell: Arc::new(Mutex::new(None)),
137 initialized: false,
138 }
139 }
140
141 pub async fn init_pty_shell(&mut self) -> Result<()> {
142 let shell =
143 PtyShell::new(&self.cwd, self.bash_command_mode.bash_mode == BashMode::RestrictedMode)?;
144 *self.pty_shell.lock().await = Some(shell);
145 Ok(())
146 }
147
148 pub fn update_cwd(&mut self, path: &Path) -> Result<()> {
149 self.cwd = path.to_path_buf();
150 Ok(())
151 }
152
153 pub fn update_workspace_root(&mut self, path: &Path) -> Result<()> {
154 self.workspace_root = path.to_path_buf();
155 Ok(())
156 }
157
158 pub fn is_command_allowed(&self, command: &str) -> bool {
159 self.bash_command_mode.allowed_commands.is_allowed(command)
160 }
161
162 pub fn is_file_edit_allowed(&self, path: &str) -> bool {
163 self.file_edit_mode.allowed_globs.is_allowed(path)
164 }
165
166 pub fn is_file_write_allowed(&self, path: &str) -> bool {
167 self.write_if_empty_mode.allowed_globs.is_allowed(path)
168 }
169 pub fn get_mode_violation_message(&self, op: &str, _target: &str) -> String {
170 format!("Operation {op} not allowed")
171 }
172
173 pub fn save_state_to_disk(&self) -> Result<()> {
174 let snapshot = self.snapshot();
175 save_state_file(&self.current_thread_id, &snapshot)?;
176 Ok(())
177 }
178
179 pub fn snapshot(&self) -> BashStateSnapshot {
180 BashStateSnapshot::from_state(
181 &self.cwd.to_string_lossy(),
182 &self.workspace_root.to_string_lossy(),
183 &self.mode,
184 &self.bash_command_mode,
185 &self.file_edit_mode,
186 &self.write_if_empty_mode,
187 &self.whitelist_for_overwrite,
188 &self.current_thread_id,
189 )
190 }
191
192 pub fn apply_snapshot(&mut self, snapshot: &BashStateSnapshot) {
193 let (cwd, root, mode, bmode, emode, wmode, whitelist, tid) = snapshot.to_state_components();
194
195 self.cwd = PathBuf::from(cwd);
196 self.workspace_root = PathBuf::from(root);
197 self.mode = mode;
198 self.bash_command_mode = bmode;
199 self.file_edit_mode = emode;
200 self.write_if_empty_mode = wmode;
201 self.whitelist_for_overwrite = whitelist;
202 self.current_thread_id = tid;
203 self.initialized = true;
204 }
205
206 pub fn load_state_from_disk(&mut self, thread_id: &str) -> Result<bool> {
207 if let Some(snapshot) = load_state_file(thread_id)? {
208 self.apply_snapshot(&snapshot);
209 Ok(true)
210 } else {
211 Ok(false)
212 }
213 }
214
215 pub fn new_with_thread_id(thread_id: Option<&str>) -> Self {
216 let mut state = Self::new();
217
218 if let Some(tid) = thread_id {
219 if !tid.is_empty() {
220 if let Ok(true) = state.load_state_from_disk(tid) {
221 info!("Loaded state for thread_id '{}'", tid);
222 } else {
223 state.current_thread_id = tid.to_string();
224 }
225 }
226 }
227
228 state
229 }
230}
231
232pub fn generate_thread_id() -> String {
233 let mut rng = rand::rng();
234 format!("tid_{:x}", rng.random::<u64>())
235}