style: Cleaned up all graph agent code

src/graph/state.rs

@@ -1,5 +1,3 @@
-//! State management and template interpolation for graph execution.
-
 use super::MAX_STATE_SIZE_BYTES;
 use super::types::GraphState;
 use crate::utils::temp_file;
@@ -8,21 +6,13 @@ use fancy_regex::Regex;
 use serde_json::Value;
 use std::collections::HashMap;
 use std::fs;
-use std::fs::{read_to_string, write};
+use std::fs::write;
 use std::path::PathBuf;
 use std::sync::LazyLock;
 
 static TEMPLATE_VAR_RE: LazyLock<Regex> =
     LazyLock::new(|| Regex::new(r"\{\{([a-zA-Z0-9_\.\[\]]+)\}\}").expect("invalid template regex"));
 
-/// Wraps [`GraphState`] with template interpolation, script-output merging,
-/// and a large-state temp-file fallback for use with scripts.
-///
-/// Template syntax: `{{key}}` for top-level keys, `{{a.b.c}}` for nested
-/// JSON paths, and `{{arr[0]}}` / `{{a.b[2].c}}` / `{{matrix[0][1]}}` for
-/// array indices. Use [`StateManager::interpolate`] for strict interpolation
-/// (errors on missing keys) or [`StateManager::interpolate_lenient`] for
-/// best-effort (missing keys become empty strings).
 pub struct StateManager {
     state: GraphState,
     temp_file: Option<PathBuf>,
@@ -44,8 +34,6 @@ impl StateManager {
         &mut self.state
     }
 
-    /// Replace every `{{key}}` in `template` with its state value. Returns
-    /// an error if any referenced key is missing.
     pub fn interpolate(&self, template: &str) -> Result<String> {
         let mut missing = Vec::new();
         let result = self.interpolate_inner(template, |key| {
@@ -65,8 +53,6 @@ impl StateManager {
         Ok(result)
     }
 
-    /// Same as [`Self::interpolate`] but missing keys are silently replaced
-    /// with an empty string.
     pub fn interpolate_lenient(&self, template: &str) -> String {
         self.interpolate_inner(template, |_| String::new())
     }
@@ -96,6 +82,7 @@ impl StateManager {
         for idx in root_indices {
             current = current.get(idx)?;
         }
+
         for part in parts {
             let (segment_key, indices) = split_indices(part)?;
             if !segment_key.is_empty() {
@@ -105,12 +92,10 @@ impl StateManager {
                 current = current.get(idx)?;
             }
         }
+
         Some(current)
     }
 
-    /// Serialize the state for transport to a script. State larger than
-    /// [`MAX_STATE_SIZE_BYTES`] is written to a unique temp file; the file
-    /// is cleaned up when the `StateManager` is dropped.
     pub fn serialize_state(&mut self) -> Result<StateRepresentation> {
         let json = self.state.to_json()?;
         if json.len() > MAX_STATE_SIZE_BYTES {
@@ -119,16 +104,14 @@ impl StateManager {
                 format!("Failed to write state to temp file at '{}'", path.display())
             })?;
             self.temp_file = Some(path.clone());
+
             Ok(StateRepresentation::File(path))
         } else {
             Ok(StateRepresentation::Inline(json))
         }
     }
 
-    pub fn to_json_string(&self) -> Result<String> {
-        self.state.to_json()
-    }
-
+    #[cfg(test)]
     pub fn from_json_string(json: &str) -> Result<Self> {
         let data: HashMap<String, Value> =
             serde_json::from_str(json).context("Failed to parse state JSON")?;
@@ -143,13 +126,11 @@ impl StateManager {
         self.state.size_bytes()
     }
 
+    #[cfg(test)]
     pub fn is_large(&self) -> bool {
         self.size_bytes() > MAX_STATE_SIZE_BYTES
     }
 
-    /// Merge a script's JSON-object stdout into state. The reserved `_next`
-    /// key is extracted (used by the executor for routing) and is not stored
-    /// in state. Errors if the output is not a JSON object.
     pub fn merge_script_output(&mut self, json_output: &str) -> Result<Option<String>> {
         let value: Value =
             serde_json::from_str(json_output).context("Script output must be valid JSON")?;
@@ -170,8 +151,6 @@ impl StateManager {
         Ok(next_node)
     }
 
-    /// Remove the temp file backing this state, if any. Called automatically
-    /// on drop.
     pub fn cleanup(&mut self) {
         if let Some(path) = self.temp_file.take() {
             let _ = fs::remove_file(path);
@@ -185,19 +164,18 @@ impl Drop for StateManager {
     }
 }
 
-/// How serialized state is delivered to a script: inline JSON for small
-/// state, or a file path for state above [`MAX_STATE_SIZE_BYTES`].
 #[derive(Debug, Clone)]
 pub enum StateRepresentation {
     Inline(String),
     File(PathBuf),
 }
 
+#[cfg(test)]
 impl StateRepresentation {
     pub fn as_string(&self) -> Result<String> {
         match self {
             StateRepresentation::Inline(s) => Ok(s.clone()),
-            StateRepresentation::File(path) => read_to_string(path)
+            StateRepresentation::File(path) => fs::read_to_string(path)
                 .with_context(|| format!("Failed to read state file at '{}'", path.display())),
         }
     }
@@ -222,6 +200,7 @@ fn split_indices(segment: &str) -> Option<(&str, Vec<usize>)> {
     };
     let mut indices = Vec::new();
     let mut rest = &segment[bracket_start.unwrap()..];
+
     while !rest.is_empty() {
         if !rest.starts_with('[') {
             return None;
@@ -231,6 +210,7 @@ fn split_indices(segment: &str) -> Option<(&str, Vec<usize>)> {
         indices.push(idx);
         rest = &rest[close + 1..];
     }
+
     Some((key, indices))
 }
 
@@ -261,9 +241,11 @@ mod tests {
     #[test]
     fn simple_interpolation_replaces_top_level_keys() {
         let manager = manager_with(&[("name", json!("Alice")), ("age", json!(30))]);
+
         let result = manager
             .interpolate("Hello {{name}}, you are {{age}} years old")
             .unwrap();
+
         assert_eq!(result, "Hello Alice, you are 30 years old");
     }
 
@@ -271,9 +253,11 @@ mod tests {
     fn nested_interpolation_walks_objects() {
         let manager =
             manager_with(&[("user", json!({ "name": "Bob", "email": "bob@example.com" }))]);
+
         let result = manager
             .interpolate("User: {{user.name}} ({{user.email}})")
             .unwrap();
+
         assert_eq!(result, "User: Bob (bob@example.com)");
     }
 
@@ -285,19 +269,23 @@ mod tests {
                 "api": { "key": "secret123", "endpoint": "https://api.example.com" }
             }),
         )]);
+
         let result = manager
             .interpolate("API: {{config.api.endpoint}} with key {{config.api.key}}")
             .unwrap();
+
         assert_eq!(result, "API: https://api.example.com with key secret123");
     }
 
     #[test]
     fn strict_interpolation_errors_on_missing_keys() {
         let manager = manager_with(&[]);
+
         let err = manager
             .interpolate("Hello {{name}}")
             .unwrap_err()
             .to_string();
+
         assert!(err.contains("not found"), "got: {err}");
         assert!(err.contains("name"), "got: {err}");
     }
@@ -305,24 +293,30 @@ mod tests {
     #[test]
     fn strict_interpolation_collects_all_missing_keys() {
         let manager = manager_with(&[]);
+
         let err = manager
             .interpolate("{{a}} and {{b}}")
             .unwrap_err()
             .to_string();
+
         assert!(err.contains("'a'") && err.contains("'b'"), "got: {err}");
     }
 
     #[test]
     fn lenient_interpolation_substitutes_empty_for_missing() {
         let manager = manager_with(&[("name", json!("Alice"))]);
+
         let result = manager.interpolate_lenient("Hello {{name}}, age: {{age}}");
+
         assert_eq!(result, "Hello Alice, age: ");
     }
 
     #[test]
     fn lenient_interpolation_handles_missing_intermediate() {
         let manager = manager_with(&[("user", json!({ "name": "Bob" }))]);
+
         let result = manager.interpolate_lenient("email: {{user.email}}");
+
         assert_eq!(result, "email: ");
     }
 
@@ -333,6 +327,7 @@ mod tests {
             ("count", json!(42)),
             ("nothing", json!(null)),
         ]);
+
         assert_eq!(manager.interpolate("{{on}}").unwrap(), "true");
         assert_eq!(manager.interpolate("{{count}}").unwrap(), "42");
         assert_eq!(manager.interpolate("{{nothing}}").unwrap(), "null");
@@ -341,20 +336,25 @@ mod tests {
     #[test]
     fn interpolates_arrays_as_json() {
         let manager = manager_with(&[("items", json!(["a", "b", "c"]))]);
+
         let result = manager.interpolate("{{items}}").unwrap();
+
         assert_eq!(result, r#"["a","b","c"]"#);
     }
 
     #[test]
     fn interpolates_objects_as_json() {
         let manager = manager_with(&[("data", json!({ "key": "value" }))]);
+
         let result = manager.interpolate("{{data}}").unwrap();
+
         assert_eq!(result, r#"{"key":"value"}"#);
     }
 
     #[test]
     fn interpolates_array_indices() {
         let manager = manager_with(&[("items", json!(["a", "b", "c"]))]);
+
         assert_eq!(manager.interpolate("{{items[0]}}").unwrap(), "a");
         assert_eq!(manager.interpolate("{{items[2]}}").unwrap(), "c");
     }
@@ -362,24 +362,29 @@ mod tests {
     #[test]
     fn interpolates_array_indices_inside_nested_paths() {
         let manager = manager_with(&[("outer", json!({ "inner": { "arr": ["x", "y", "z"] } }))]);
+
         let result = manager
             .interpolate("first={{outer.inner.arr[0]}} last={{outer.inner.arr[2]}}")
             .unwrap();
+
         assert_eq!(result, "first=x last=z");
     }
 
     #[test]
     fn interpolates_object_fields_after_array_index() {
         let manager = manager_with(&[("users", json!([{ "name": "Alice" }, { "name": "Bob" }]))]);
+
         let result = manager
             .interpolate("{{users[0].name}} and {{users[1].name}}")
             .unwrap();
+
         assert_eq!(result, "Alice and Bob");
     }
 
     #[test]
     fn interpolates_nested_array_indices() {
         let manager = manager_with(&[("matrix", json!([[1, 2], [3, 4]]))]);
+
         assert_eq!(manager.interpolate("{{matrix[0][1]}}").unwrap(), "2");
         assert_eq!(manager.interpolate("{{matrix[1][0]}}").unwrap(), "3");
     }
@@ -387,21 +392,27 @@ mod tests {
     #[test]
     fn out_of_bounds_array_index_is_missing() {
         let manager = manager_with(&[("items", json!(["a", "b"]))]);
+
         let err = manager.interpolate("{{items[5]}}").unwrap_err().to_string();
+
         assert!(err.contains("not found"), "got: {err}");
     }
 
     #[test]
     fn replaces_all_occurrences_of_same_key() {
         let manager = manager_with(&[("n", json!("Alice"))]);
+
         let result = manager.interpolate("{{n}} and {{n}} again").unwrap();
+
         assert_eq!(result, "Alice and Alice again");
     }
 
     #[test]
     fn passes_through_templates_with_no_variables() {
         let manager = manager_with(&[]);
+
         let result = manager.interpolate("No variables here").unwrap();
+
         assert_eq!(result, "No variables here");
     }
 
@@ -409,14 +420,18 @@ mod tests {
     fn from_json_string_round_trips() {
         let json = r#"{"name": "Alice", "age": 30}"#;
         let manager = StateManager::from_json_string(json).unwrap();
+
         let result = manager.interpolate("{{name}} is {{age}}").unwrap();
+
         assert_eq!(result, "Alice is 30");
     }
 
     #[test]
     fn snapshot_clones_state_data() {
         let manager = manager_with(&[("k1", json!("v1")), ("k2", json!(42))]);
+
         let snap = manager.snapshot();
+
         assert_eq!(snap.len(), 2);
         assert_eq!(snap.get("k1"), Some(&json!("v1")));
         assert_eq!(snap.get("k2"), Some(&json!(42)));
@@ -425,7 +440,9 @@ mod tests {
     #[test]
     fn small_state_serializes_inline() {
         let mut manager = manager_with(&[("key", json!("value"))]);
+
         let repr = manager.serialize_state().unwrap();
+
         assert!(matches!(repr, StateRepresentation::Inline(_)));
         assert!(!manager.is_large());
         assert!(!repr.is_file());
@@ -443,7 +460,7 @@ mod tests {
         assert!(path.exists());
 
         let contents = repr.as_string().unwrap();
-        let parsed: serde_json::Value = serde_json::from_str(&contents).unwrap();
+        let parsed: Value = serde_json::from_str(&contents).unwrap();
         assert_eq!(
             parsed.get("blob").unwrap().as_str().unwrap().len(),
             big.len()
@@ -457,7 +474,9 @@ mod tests {
     fn merge_script_output_merges_keys_into_state() {
         let mut manager = manager_with(&[]);
         let output = r#"{"quality_score": 0.85, "issues_found": 3, "status": "complete"}"#;
+
         let next = manager.merge_script_output(output).unwrap();
+
         assert_eq!(next, None);
         assert_eq!(manager.state().get("quality_score"), Some(&json!(0.85)));
         assert_eq!(manager.state().get("issues_found"), Some(&json!(3)));
@@ -468,7 +487,9 @@ mod tests {
     fn merge_script_output_extracts_next_key_for_routing() {
         let mut manager = manager_with(&[]);
         let output = r#"{"_next": "approval_gate", "quality_score": 0.85}"#;
+
         let next = manager.merge_script_output(output).unwrap();
+
         assert_eq!(next.as_deref(), Some("approval_gate"));
         assert_eq!(manager.state().get("quality_score"), Some(&json!(0.85)));
         assert!(
@@ -480,29 +501,35 @@ mod tests {
     #[test]
     fn merge_script_output_rejects_invalid_json() {
         let mut manager = manager_with(&[]);
+
         let err = manager
             .merge_script_output("not json")
             .unwrap_err()
             .to_string();
+
         assert!(err.contains("valid JSON"), "got: {err}");
     }
 
     #[test]
     fn merge_script_output_rejects_non_object() {
         let mut manager = manager_with(&[]);
+
         let err = manager
             .merge_script_output("[1, 2, 3]")
             .unwrap_err()
             .to_string();
+
         assert!(err.contains("must be a JSON object"), "got: {err}");
     }
 
     #[test]
     fn merge_script_output_overwrites_existing_state_keys() {
         let mut manager = manager_with(&[("status", json!("pending"))]);
+
         let _ = manager
             .merge_script_output(r#"{"status": "complete"}"#)
             .unwrap();
+
         assert_eq!(manager.state().get("status"), Some(&json!("complete")));
     }
 }