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::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 load_bash_state(thread_id: &str) -> Result<Option<BashStateSnapshot>> {
93    if thread_id.is_empty() {
94        return Ok(None);
95    }
96    let path = get_state_file_path(thread_id)?;
97    if !path.exists() {
98        return Ok(None);
99    }
100    let json = fs::read_to_string(path)?;
101    Ok(Some(serde_json::from_str(&json)?))
102}
103
104pub fn delete_bash_state(thread_id: &str) -> Result<()> {
105    if thread_id.is_empty() {
106        return Ok(());
107    }
108    let path = get_state_file_path(thread_id)?;
109    if path.exists() {
110        fs::remove_file(path)?;
111    }
112    Ok(())
113}
114
115impl BashStateSnapshot {
116    pub fn from_state(
117        cwd: &str,
118        workspace_root: &str,
119        mode: &Modes,
120        bash_command_mode: &BashCommandMode,
121        file_edit_mode: &FileEditMode,
122        write_if_empty_mode: &WriteIfEmptyMode,
123        whitelist: &HashMap<String, FileWhitelistData>,
124        thread_id: &str,
125    ) -> Self {
126        Self {
127            bash_command_mode: BashCommandModeSnapshot::from(bash_command_mode),
128            file_edit_mode: FileEditModeSnapshot::from(file_edit_mode),
129            write_if_empty_mode: WriteIfEmptyModeSnapshot::from(write_if_empty_mode),
130            whitelist_for_overwrite: whitelist
131                .iter()
132                .map(|(k, v)| (k.clone(), FileWhitelistDataSnapshot::from(v)))
133                .collect(),
134            mode: mode.to_string(),
135            workspace_root: workspace_root.to_string(),
136            chat_id: thread_id.to_string(),
137            cwd: if cwd == workspace_root { String::new() } else { cwd.to_string() },
138        }
139    }
140
141    pub fn to_state_components(
142        &self,
143    ) -> (
144        String,
145        String,
146        Modes,
147        BashCommandMode,
148        FileEditMode,
149        WriteIfEmptyMode,
150        HashMap<String, FileWhitelistData>,
151        String,
152    ) {
153        let mode = match self.mode.as_str() {
154            "architect" => Modes::Architect,
155            "code_writer" => Modes::CodeWriter,
156            _ => Modes::Wcgw,
157        };
158        let whitelist = self
159            .whitelist_for_overwrite
160            .iter()
161            .map(|(k, v)| (k.clone(), v.to_whitelist_data()))
162            .collect();
163        let cwd = if self.cwd.is_empty() { self.workspace_root.clone() } else { self.cwd.clone() };
164        (
165            cwd,
166            self.workspace_root.clone(),
167            mode,
168            self.bash_command_mode.to_bash_command_mode(),
169            self.file_edit_mode.to_file_edit_mode(),
170            self.write_if_empty_mode.to_write_if_empty_mode(),
171            whitelist,
172            self.chat_id.clone(),
173        )
174    }
175}
176
177impl From<&BashCommandMode> for BashCommandModeSnapshot {
178    fn from(mode: &BashCommandMode) -> Self {
179        Self {
180            bash_mode: match mode.bash_mode {
181                BashMode::RestrictedMode => "restricted_mode".to_string(),
182                BashMode::NormalMode => "normal_mode".to_string(),
183            },
184            allowed_commands: match &mode.allowed_commands {
185                AllowedCommands::All(s) => AllowedCommandsSnapshot::All(s.clone()),
186                AllowedCommands::List(l) => AllowedCommandsSnapshot::List(l.clone()),
187            },
188        }
189    }
190}
191
192impl BashCommandModeSnapshot {
193    fn to_bash_command_mode(&self) -> BashCommandMode {
194        BashCommandMode {
195            bash_mode: if self.bash_mode == "restricted_mode" {
196                BashMode::RestrictedMode
197            } else {
198                BashMode::NormalMode
199            },
200            allowed_commands: match &self.allowed_commands {
201                AllowedCommandsSnapshot::All(s) => AllowedCommands::All(s.clone()),
202                AllowedCommandsSnapshot::List(l) => AllowedCommands::List(l.clone()),
203            },
204        }
205    }
206}
207
208impl From<&FileEditMode> for FileEditModeSnapshot {
209    fn from(mode: &FileEditMode) -> Self {
210        Self {
211            allowed_globs: match &mode.allowed_globs {
212                AllowedGlobs::All(s) => AllowedGlobsSnapshot::All(s.clone()),
213                AllowedGlobs::List(l) => AllowedGlobsSnapshot::List(l.clone()),
214            },
215        }
216    }
217}
218
219impl FileEditModeSnapshot {
220    fn to_file_edit_mode(&self) -> FileEditMode {
221        FileEditMode {
222            allowed_globs: match &self.allowed_globs {
223                AllowedGlobsSnapshot::All(s) => AllowedGlobs::All(s.clone()),
224                AllowedGlobsSnapshot::List(l) => AllowedGlobs::List(l.clone()),
225            },
226        }
227    }
228}
229
230impl From<&WriteIfEmptyMode> for WriteIfEmptyModeSnapshot {
231    fn from(mode: &WriteIfEmptyMode) -> Self {
232        Self {
233            allowed_globs: match &mode.allowed_globs {
234                AllowedGlobs::All(s) => AllowedGlobsSnapshot::All(s.clone()),
235                AllowedGlobs::List(l) => AllowedGlobsSnapshot::List(l.clone()),
236            },
237        }
238    }
239}
240
241impl WriteIfEmptyModeSnapshot {
242    fn to_write_if_empty_mode(&self) -> WriteIfEmptyMode {
243        WriteIfEmptyMode {
244            allowed_globs: match &self.allowed_globs {
245                AllowedGlobsSnapshot::All(s) => AllowedGlobs::All(s.clone()),
246                AllowedGlobsSnapshot::List(l) => AllowedGlobs::List(l.clone()),
247            },
248        }
249    }
250}
251
252impl From<&FileWhitelistData> for FileWhitelistDataSnapshot {
253    fn from(d: &FileWhitelistData) -> Self {
254        Self {
255            file_hash: d.file_hash.clone(),
256            line_ranges_read: d.line_ranges_read.clone(),
257            total_lines: d.total_lines,
258        }
259    }
260}
261
262impl FileWhitelistDataSnapshot {
263    fn to_whitelist_data(&self) -> FileWhitelistData {
264        FileWhitelistData {
265            file_hash: self.file_hash.clone(),
266            line_ranges_read: self.line_ranges_read.clone(),
267            total_lines: self.total_lines,
268        }
269    }
270}
winx_code_agent/state/persistence.rs

winx_code_agent/state/
persistence.rs
