arcan_core/
state.rs

1use crate::protocol::{StatePatch, StatePatchFormat};
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use thiserror::Error;
6
7/// Well-known keys in `AppState.data`. JSON Patch still operates on the raw `Value`.
8#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Default)]
9#[serde(default)]
10pub struct WellKnownState {
11    /// Current working directory for the session.
12    pub cwd: Option<String>,
13    /// Active file paths the agent is aware of.
14    pub open_files: Option<Vec<String>>,
15    /// Session-level metadata.
16    pub session_meta: Option<SessionMeta>,
17    /// Tool execution budget tracking.
18    pub budget: Option<BudgetState>,
19    /// Currently loaded skill names.
20    pub active_skills: Option<Vec<String>>,
21    /// Connected MCP server info.
22    pub mcp_servers: Option<Vec<McpServerInfo>>,
23}
24
25/// Session-level metadata stored in agent state.
26#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
27pub struct SessionMeta {
28    pub session_name: Option<String>,
29    pub user_id: Option<String>,
30    pub created_at: Option<String>,
31}
32
33/// Tool execution budget tracking.
34#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
35pub struct BudgetState {
36    pub total_tokens_used: u64,
37    pub max_tokens_budget: Option<u64>,
38    pub tool_calls_count: u32,
39    pub max_tool_calls: Option<u32>,
40}
41
42/// Info about a connected MCP server.
43#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
44pub struct McpServerInfo {
45    pub name: String,
46    /// Transport type: "stdio" or "http".
47    pub transport: String,
48    pub tool_count: usize,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
52pub struct AppState {
53    pub revision: u64,
54    #[serde(default)]
55    pub data: Value,
56}
57
58impl Default for AppState {
59    fn default() -> Self {
60        Self {
61            revision: 0,
62            data: Value::Object(Default::default()),
63        }
64    }
65}
66
67impl AppState {
68    pub fn new(data: Value) -> Self {
69        Self { revision: 0, data }
70    }
71
72    /// Parse well-known keys from the raw JSON data.
73    pub fn well_known(&self) -> Result<WellKnownState, serde_json::Error> {
74        serde_json::from_value(self.data.clone())
75    }
76
77    /// Get the current working directory from state.
78    pub fn cwd(&self) -> Option<String> {
79        self.data
80            .get("cwd")
81            .and_then(|v| v.as_str())
82            .map(String::from)
83    }
84
85    /// Get the list of open files from state.
86    pub fn open_files(&self) -> Vec<String> {
87        self.data
88            .get("open_files")
89            .and_then(|v| serde_json::from_value(v.clone()).ok())
90            .unwrap_or_default()
91    }
92
93    pub fn apply_patch(&mut self, patch: &StatePatch) -> Result<(), StateError> {
94        match patch.format {
95            StatePatchFormat::JsonPatch => {
96                let parsed_patch: json_patch::Patch = serde_json::from_value(patch.patch.clone())
97                    .map_err(StateError::InvalidJsonPatch)?;
98                json_patch::patch(&mut self.data, &parsed_patch)
99                    .map_err(|e| StateError::PatchApply(e.to_string()))?;
100            }
101            StatePatchFormat::MergePatch => {
102                json_patch::merge(&mut self.data, &patch.patch);
103            }
104        }
105
106        self.revision = self.revision.saturating_add(1);
107        Ok(())
108    }
109}
110
111#[derive(Debug, Error)]
112pub enum StateError {
113    #[error("invalid JSON patch payload: {0}")]
114    InvalidJsonPatch(serde_json::Error),
115    #[error("failed to apply patch: {0}")]
116    PatchApply(String),
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122    use crate::protocol::StatePatchSource;
123    use serde_json::json;
124
125    #[test]
126    fn default_state_is_empty_object() {
127        let state = AppState::default();
128        assert_eq!(state.revision, 0);
129        assert_eq!(state.data, json!({}));
130    }
131
132    #[test]
133    fn merge_patch_adds_fields() {
134        let mut state = AppState::default();
135        let patch = StatePatch {
136            format: StatePatchFormat::MergePatch,
137            patch: json!({"name": "arcan", "version": 1}),
138            source: StatePatchSource::System,
139        };
140        state.apply_patch(&patch).unwrap();
141
142        assert_eq!(state.revision, 1);
143        assert_eq!(state.data["name"], "arcan");
144        assert_eq!(state.data["version"], 1);
145    }
146
147    #[test]
148    fn merge_patch_overwrites_fields() {
149        let mut state = AppState::new(json!({"count": 0}));
150        let patch = StatePatch {
151            format: StatePatchFormat::MergePatch,
152            patch: json!({"count": 42}),
153            source: StatePatchSource::Tool,
154        };
155        state.apply_patch(&patch).unwrap();
156        assert_eq!(state.data["count"], 42);
157    }
158
159    #[test]
160    fn merge_patch_removes_null_fields() {
161        let mut state = AppState::new(json!({"a": 1, "b": 2}));
162        let patch = StatePatch {
163            format: StatePatchFormat::MergePatch,
164            patch: json!({"b": null}),
165            source: StatePatchSource::Model,
166        };
167        state.apply_patch(&patch).unwrap();
168        assert_eq!(state.data, json!({"a": 1}));
169    }
170
171    #[test]
172    fn json_patch_add_operation() {
173        let mut state = AppState::default();
174        let patch = StatePatch {
175            format: StatePatchFormat::JsonPatch,
176            patch: json!([{"op": "add", "path": "/foo", "value": "bar"}]),
177            source: StatePatchSource::System,
178        };
179        state.apply_patch(&patch).unwrap();
180        assert_eq!(state.data["foo"], "bar");
181        assert_eq!(state.revision, 1);
182    }
183
184    #[test]
185    fn json_patch_replace_operation() {
186        let mut state = AppState::new(json!({"x": 10}));
187        let patch = StatePatch {
188            format: StatePatchFormat::JsonPatch,
189            patch: json!([{"op": "replace", "path": "/x", "value": 20}]),
190            source: StatePatchSource::System,
191        };
192        state.apply_patch(&patch).unwrap();
193        assert_eq!(state.data["x"], 20);
194    }
195
196    #[test]
197    fn json_patch_invalid_payload_errors() {
198        let mut state = AppState::default();
199        let patch = StatePatch {
200            format: StatePatchFormat::JsonPatch,
201            patch: json!("not an array"),
202            source: StatePatchSource::System,
203        };
204        assert!(state.apply_patch(&patch).is_err());
205        assert_eq!(state.revision, 0);
206    }
207
208    #[test]
209    fn well_known_parses_populated_state() {
210        let state = AppState::new(json!({
211            "cwd": "/home/user",
212            "open_files": ["main.rs", "lib.rs"],
213            "budget": {
214                "total_tokens_used": 1000,
215                "max_tokens_budget": 10000,
216                "tool_calls_count": 5,
217                "max_tool_calls": 100
218            }
219        }));
220        let wk = state.well_known().unwrap();
221        assert_eq!(wk.cwd.as_deref(), Some("/home/user"));
222        assert_eq!(
223            wk.open_files,
224            Some(vec!["main.rs".to_string(), "lib.rs".to_string()])
225        );
226        assert!(wk.budget.is_some());
227        assert_eq!(wk.budget.unwrap().total_tokens_used, 1000);
228    }
229
230    #[test]
231    fn well_known_defaults_for_empty_state() {
232        let state = AppState::default();
233        let wk = state.well_known().unwrap();
234        assert_eq!(wk.cwd, None);
235        assert_eq!(wk.open_files, None);
236        assert_eq!(wk.session_meta, None);
237        assert_eq!(wk.budget, None);
238        assert_eq!(wk.active_skills, None);
239        assert_eq!(wk.mcp_servers, None);
240    }
241
242    #[test]
243    fn cwd_accessor() {
244        let state = AppState::new(json!({"cwd": "/tmp"}));
245        assert_eq!(state.cwd(), Some("/tmp".to_string()));
246
247        let empty = AppState::default();
248        assert_eq!(empty.cwd(), None);
249    }
250
251    #[test]
252    fn open_files_accessor() {
253        let state = AppState::new(json!({"open_files": ["a.rs", "b.rs"]}));
254        assert_eq!(
255            state.open_files(),
256            vec!["a.rs".to_string(), "b.rs".to_string()]
257        );
258
259        let empty = AppState::default();
260        assert!(empty.open_files().is_empty());
261    }
262
263    #[test]
264    fn json_patch_works_with_well_known_keys() {
265        let mut state = AppState::new(json!({"cwd": "/old"}));
266        let patch = StatePatch {
267            format: StatePatchFormat::JsonPatch,
268            patch: json!([{"op": "replace", "path": "/cwd", "value": "/new"}]),
269            source: StatePatchSource::System,
270        };
271        state.apply_patch(&patch).unwrap();
272        assert_eq!(state.cwd(), Some("/new".to_string()));
273    }
274
275    #[test]
276    fn revision_increments_with_each_patch() {
277        let mut state = AppState::default();
278        for i in 1..=5 {
279            let patch = StatePatch {
280                format: StatePatchFormat::MergePatch,
281                patch: json!({"step": i}),
282                source: StatePatchSource::System,
283            };
284            state.apply_patch(&patch).unwrap();
285            assert_eq!(state.revision, i as u64);
286        }
287    }
288}
arcan_core/state.rs

arcan_core/
state.rs
