winx_code_agent/state/
persistence.rs

1//! Bash state persistence module
2//!
3//! Provides disk persistence for `BashState`, compatible with WCGW Python implementation.
4//! State is stored in `~/.local/share/wcgw/bash_state/` as JSON files.
5
6use anyhow::{Context, Result};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::fs;
10use std::path::{Path, PathBuf};
11use tracing::{debug, info, warn};
12
13use crate::types::{
14    AllowedCommands, AllowedGlobs, BashCommandMode, BashMode, FileEditMode, Modes, WriteIfEmptyMode,
15};
16
17use super::bash_state::FileWhitelistData;
18
19/// Snapshot of `BashState` that can be serialized to disk
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct BashStateSnapshot {
22    pub bash_command_mode: BashCommandModeSnapshot,
23    pub file_edit_mode: FileEditModeSnapshot,
24    pub write_if_empty_mode: WriteIfEmptyModeSnapshot,
25    pub whitelist_for_overwrite: HashMap<String, FileWhitelistDataSnapshot>,
26    pub mode: String,
27    pub workspace_root: String,
28    pub chat_id: String,
29    #[serde(default, skip_serializing_if = "String::is_empty")]
30    pub cwd: String,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct BashCommandModeSnapshot {
35    pub bash_mode: String,
36    pub allowed_commands: AllowedCommandsSnapshot,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
40#[serde(untagged)]
41pub enum AllowedCommandsSnapshot {
42    All(String),
43    List(Vec<String>),
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct FileEditModeSnapshot {
48    pub allowed_globs: AllowedGlobsSnapshot,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct WriteIfEmptyModeSnapshot {
53    pub allowed_globs: AllowedGlobsSnapshot,
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
57#[serde(untagged)]
58pub enum AllowedGlobsSnapshot {
59    All(String),
60    List(Vec<String>),
61}
62
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct FileWhitelistDataSnapshot {
65    pub file_hash: String,
66    pub line_ranges_read: Vec<(usize, usize)>,
67    pub total_lines: usize,
68}
69
70pub fn get_state_dir() -> Result<PathBuf> {
71    let home = home::home_dir().unwrap_or_else(|| PathBuf::from("/tmp"));
72    let bash_state_dir = home.join(".local/share/wcgw/bash_state");
73    if !bash_state_dir.exists() {
74        fs::create_dir_all(&bash_state_dir)?;
75    }
76    Ok(bash_state_dir)
77}
78
79fn get_state_file_path(thread_id: &str) -> Result<PathBuf> {
80    Ok(get_state_dir()?.join(format!("{thread_id}_bash_state.json")))
81}
82
83pub fn save_bash_state(thread_id: &str, state: &BashStateSnapshot) -> Result<()> {
84    if thread_id.is_empty() {
85        return Ok(());
86    }
87    let json = serde_json::to_string_pretty(state)?;
88    fs::write(get_state_file_path(thread_id)?, json)?;
89    Ok(())
90}
91
92pub fn save_bash_state_to_path(path: &Path, state: &BashStateSnapshot) -> Result<()> {
93    let json = serde_json::to_string_pretty(state)?;
94    fs::write(path, json)?;
95    Ok(())
96}
97
98pub fn load_bash_state(thread_id: &str) -> Result<Option<BashStateSnapshot>> {
99    if thread_id.is_empty() {
100        return Ok(None);
101    }
102    let path = get_state_file_path(thread_id)?;
103    if !path.exists() {
104        return Ok(None);
105    }
106    let json = fs::read_to_string(path)?;
107    Ok(Some(serde_json::from_str(&json)?))
108}
109
110pub fn load_bash_state_from_path(path: &Path) -> Result<Option<BashStateSnapshot>> {
111    if !path.exists() {
112        return Ok(None);
113    }
114    let json = fs::read_to_string(path)?;
115    Ok(Some(serde_json::from_str(&json)?))
116}
117
118pub fn delete_bash_state(thread_id: &str) -> Result<()> {
119    if thread_id.is_empty() {
120        return Ok(());
121    }
122    let path = get_state_file_path(thread_id)?;
123    if path.exists() {
124        fs::remove_file(path)?;
125    }
126    Ok(())
127}
128
129impl BashStateSnapshot {
130    pub fn from_state(
131        cwd: &str,
132        workspace_root: &str,
133        mode: &Modes,
134        bash_command_mode: &BashCommandMode,
135        file_edit_mode: &FileEditMode,
136        write_if_empty_mode: &WriteIfEmptyMode,
137        whitelist: &HashMap<String, FileWhitelistData>,
138        thread_id: &str,
139    ) -> Self {
140        Self {
141            bash_command_mode: BashCommandModeSnapshot::from(bash_command_mode),
142            file_edit_mode: FileEditModeSnapshot::from(file_edit_mode),
143            write_if_empty_mode: WriteIfEmptyModeSnapshot::from(write_if_empty_mode),
144            whitelist_for_overwrite: whitelist
145                .iter()
146                .map(|(k, v)| (k.clone(), FileWhitelistDataSnapshot::from(v)))
147                .collect(),
148            mode: mode.to_string(),
149            workspace_root: workspace_root.to_string(),
150            chat_id: thread_id.to_string(),
151            cwd: if cwd == workspace_root { String::new() } else { cwd.to_string() },
152        }
153    }
154
155    pub fn to_state_components(
156        &self,
157    ) -> (
158        String,
159        String,
160        Modes,
161        BashCommandMode,
162        FileEditMode,
163        WriteIfEmptyMode,
164        HashMap<String, FileWhitelistData>,
165        String,
166    ) {
167        let mode = match self.mode.as_str() {
168            "architect" => Modes::Architect,
169            "code_writer" => Modes::CodeWriter,
170            _ => Modes::Wcgw,
171        };
172        let whitelist = self
173            .whitelist_for_overwrite
174            .iter()
175            .map(|(k, v)| (k.clone(), v.to_whitelist_data()))
176            .collect();
177        let cwd = if self.cwd.is_empty() { self.workspace_root.clone() } else { self.cwd.clone() };
178        (
179            cwd,
180            self.workspace_root.clone(),
181            mode,
182            self.bash_command_mode.to_bash_command_mode(),
183            self.file_edit_mode.to_file_edit_mode(),
184            self.write_if_empty_mode.to_write_if_empty_mode(),
185            whitelist,
186            self.chat_id.clone(),
187        )
188    }
189}
190
191impl From<&BashCommandMode> for BashCommandModeSnapshot {
192    fn from(mode: &BashCommandMode) -> Self {
193        Self {
194            bash_mode: match mode.bash_mode {
195                BashMode::RestrictedMode => "restricted_mode".to_string(),
196                BashMode::NormalMode => "normal_mode".to_string(),
197            },
198            allowed_commands: match &mode.allowed_commands {
199                AllowedCommands::All(s) => AllowedCommandsSnapshot::All(s.clone()),
200                AllowedCommands::List(l) => AllowedCommandsSnapshot::List(l.clone()),
201            },
202        }
203    }
204}
205
206impl BashCommandModeSnapshot {
207    fn to_bash_command_mode(&self) -> BashCommandMode {
208        BashCommandMode {
209            bash_mode: if self.bash_mode == "restricted_mode" {
210                BashMode::RestrictedMode
211            } else {
212                BashMode::NormalMode
213            },
214            allowed_commands: match &self.allowed_commands {
215                AllowedCommandsSnapshot::All(s) => AllowedCommands::All(s.clone()),
216                AllowedCommandsSnapshot::List(l) => AllowedCommands::List(l.clone()),
217            },
218        }
219    }
220}
221
222impl From<&FileEditMode> for FileEditModeSnapshot {
223    fn from(mode: &FileEditMode) -> Self {
224        Self {
225            allowed_globs: match &mode.allowed_globs {
226                AllowedGlobs::All(s) => AllowedGlobsSnapshot::All(s.clone()),
227                AllowedGlobs::List(l) => AllowedGlobsSnapshot::List(l.clone()),
228            },
229        }
230    }
231}
232
233impl FileEditModeSnapshot {
234    fn to_file_edit_mode(&self) -> FileEditMode {
235        FileEditMode {
236            allowed_globs: match &self.allowed_globs {
237                AllowedGlobsSnapshot::All(s) => AllowedGlobs::All(s.clone()),
238                AllowedGlobsSnapshot::List(l) => AllowedGlobs::List(l.clone()),
239            },
240        }
241    }
242}
243
244impl From<&WriteIfEmptyMode> for WriteIfEmptyModeSnapshot {
245    fn from(mode: &WriteIfEmptyMode) -> Self {
246        Self {
247            allowed_globs: match &mode.allowed_globs {
248                AllowedGlobs::All(s) => AllowedGlobsSnapshot::All(s.clone()),
249                AllowedGlobs::List(l) => AllowedGlobsSnapshot::List(l.clone()),
250            },
251        }
252    }
253}
254
255impl WriteIfEmptyModeSnapshot {
256    fn to_write_if_empty_mode(&self) -> WriteIfEmptyMode {
257        WriteIfEmptyMode {
258            allowed_globs: match &self.allowed_globs {
259                AllowedGlobsSnapshot::All(s) => AllowedGlobs::All(s.clone()),
260                AllowedGlobsSnapshot::List(l) => AllowedGlobs::List(l.clone()),
261            },
262        }
263    }
264}
265
266impl From<&FileWhitelistData> for FileWhitelistDataSnapshot {
267    fn from(d: &FileWhitelistData) -> Self {
268        Self {
269            file_hash: d.file_hash.clone(),
270            line_ranges_read: d.line_ranges_read.clone(),
271            total_lines: d.total_lines,
272        }
273    }
274}
275
276impl FileWhitelistDataSnapshot {
277    fn to_whitelist_data(&self) -> FileWhitelistData {
278        FileWhitelistData {
279            file_hash: self.file_hash.clone(),
280            line_ranges_read: self.line_ranges_read.clone(),
281            total_lines: self.total_lines,
282        }
283    }
284}
winx_code_agent/state/persistence.rs

winx_code_agent/state/
persistence.rs
