agent_shadow_mode/
lib.rs

1/*!
2agent-shadow-mode: toggleable shadow mode for agent tool calls.
3
4When rolling out a new agent, you don't want it to actually charge cards,
5send emails, or delete files on day one. Wrap your tool calls through
6`ShadowMode`: when active, it records the intended call and returns a safe
7placeholder without running the real function.
8
9```rust
10use agent_shadow_mode::ShadowMode;
11use serde_json::json;
12
13let mut shadow = ShadowMode::new(true);
14
15// Shadow mode active: real closure never runs
16let result = shadow.intercept("charge_card", json!([]), json!({}), json!({"status": "shadowed"}), || {
17    json!({"status": "charged"}) // would run the real call
18});
19assert_eq!(result, json!({"status": "shadowed"}));
20assert_eq!(shadow.records().len(), 1);
21assert_eq!(shadow.records()[0].tool_name, "charge_card");
22
23// Deactivate: real closure runs
24shadow.deactivate();
25let real = shadow.intercept("charge_card", json!([]), json!({}), json!({"status": "shadowed"}), || {
26    json!({"status": "charged"})
27});
28assert_eq!(real, json!({"status": "charged"}));
29```
30*/
31
32use serde_json::{Map, Value};
33use std::path::Path;
34
35// ---- ShadowRecord ---------------------------------------------------------
36
37/// One recorded would-be tool call.
38#[derive(Debug, Clone)]
39pub struct ShadowRecord {
40    pub tool_name: String,
41    /// Positional args passed to the tool (serialized as JSON array items).
42    pub args: Vec<Value>,
43    /// Keyword args passed to the tool.
44    pub kwargs: Map<String, Value>,
45    /// Unix timestamp when the call was intercepted.
46    pub ts: f64,
47    /// The value that was returned in shadow mode (the placeholder).
48    pub returned: Value,
49}
50
51impl ShadowRecord {
52    /// Serialize to a single-line JSON string for JSONL output.
53    pub fn to_json_line(&self) -> String {
54        let obj = serde_json::json!({
55            "tool_name": self.tool_name,
56            "args": self.args,
57            "kwargs": self.kwargs,
58            "ts": self.ts,
59            "returned": self.returned,
60        });
61        serde_json::to_string(&obj).unwrap_or_default()
62    }
63
64    /// Deserialize from a JSONL line.
65    pub fn from_json_line(line: &str) -> Option<Self> {
66        let v: Value = serde_json::from_str(line.trim()).ok()?;
67        let obj = v.as_object()?;
68        Some(Self {
69            tool_name: obj.get("tool_name")?.as_str()?.to_owned(),
70            args: obj
71                .get("args")
72                .and_then(|a| a.as_array().cloned())
73                .unwrap_or_default(),
74            kwargs: obj
75                .get("kwargs")
76                .and_then(|k| k.as_object().cloned())
77                .unwrap_or_default(),
78            ts: obj.get("ts").and_then(|t| t.as_f64()).unwrap_or(0.0),
79            returned: obj.get("returned").cloned().unwrap_or(Value::Null),
80        })
81    }
82}
83
84// ---- ShadowMode -----------------------------------------------------------
85
86/// Toggleable shadow-mode manager for agent tool calls.
87pub struct ShadowMode {
88    active: bool,
89    records: Vec<ShadowRecord>,
90    clock: Box<dyn Fn() -> f64 + Send>,
91}
92
93impl ShadowMode {
94    /// Create a new instance. Pass `true` to start in shadow mode.
95    pub fn new(active: bool) -> Self {
96        Self {
97            active,
98            records: Vec::new(),
99            clock: Box::new(system_time_secs),
100        }
101    }
102
103    /// Create with a custom clock (useful for tests).
104    pub fn with_clock(active: bool, clock: impl Fn() -> f64 + Send + 'static) -> Self {
105        Self {
106            active,
107            records: Vec::new(),
108            clock: Box::new(clock),
109        }
110    }
111
112    // ---- state -------------------------------------------------------
113
114    pub fn is_active(&self) -> bool {
115        self.active
116    }
117
118    pub fn activate(&mut self) {
119        self.active = true;
120    }
121
122    pub fn deactivate(&mut self) {
123        self.active = false;
124    }
125
126    // ---- records -----------------------------------------------------
127
128    pub fn records(&self) -> &[ShadowRecord] {
129        &self.records
130    }
131
132    pub fn clear_records(&mut self) {
133        self.records.clear();
134    }
135
136    // ---- core API ----------------------------------------------------
137
138    /// Intercept a tool call.
139    ///
140    /// If active: record the call, return `when_shadow` without calling `f`.
141    /// If inactive: call `f` and return its result.
142    ///
143    /// - `args`: positional args (usually `json!([arg1, arg2])`)
144    /// - `kwargs`: keyword args (usually `json!({"key": val})`)
145    /// - `when_shadow`: the placeholder to return in shadow mode
146    /// - `f`: the real implementation (only called when inactive)
147    pub fn intercept<F>(
148        &mut self,
149        tool_name: &str,
150        args: Value,
151        kwargs: Value,
152        when_shadow: Value,
153        f: F,
154    ) -> Value
155    where
156        F: FnOnce() -> Value,
157    {
158        if !self.active {
159            return f();
160        }
161        let shadow_val = when_shadow;
162        let rec = ShadowRecord {
163            tool_name: tool_name.to_owned(),
164            args: match &args {
165                Value::Array(v) => v.clone(),
166                _ => vec![args],
167            },
168            kwargs: match kwargs {
169                Value::Object(m) => m,
170                _ => Map::new(),
171            },
172            ts: (self.clock)(),
173            returned: shadow_val.clone(),
174        };
175        self.records.push(rec);
176        shadow_val
177    }
178
179    // ---- persistence -------------------------------------------------
180
181    /// Save all records to `path` as JSONL (overwrites). Returns count.
182    pub fn to_jsonl(&self, path: impl AsRef<Path>) -> std::io::Result<usize> {
183        let path = path.as_ref();
184        if let Some(parent) = path.parent() {
185            if !parent.as_os_str().is_empty() {
186                std::fs::create_dir_all(parent)?;
187            }
188        }
189        let mut out = String::new();
190        for rec in &self.records {
191            out.push_str(&rec.to_json_line());
192            out.push('\n');
193        }
194        std::fs::write(path, out)?;
195        Ok(self.records.len())
196    }
197
198    /// Load records from a JSONL file.
199    pub fn load_jsonl(path: impl AsRef<Path>) -> std::io::Result<Vec<ShadowRecord>> {
200        let text = std::fs::read_to_string(path)?;
201        let recs = text
202            .lines()
203            .filter(|l| !l.trim().is_empty())
204            .filter_map(ShadowRecord::from_json_line)
205            .collect();
206        Ok(recs)
207    }
208}
209
210fn system_time_secs() -> f64 {
211    use std::time::{SystemTime, UNIX_EPOCH};
212    SystemTime::now()
213        .duration_since(UNIX_EPOCH)
214        .map(|d| d.as_secs_f64())
215        .unwrap_or(0.0)
216}
217
218// ---- tests ----------------------------------------------------------------
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use serde_json::json;
224
225    fn shadow_ts() -> ShadowMode {
226        ShadowMode::with_clock(true, || 1_700_000_000.0)
227    }
228
229    #[test]
230    fn active_returns_shadow_value() {
231        let mut s = shadow_ts();
232        let r = s.intercept("charge", json!([]), json!({}), json!("shadowed"), || json!("real"));
233        assert_eq!(r, json!("shadowed"));
234    }
235
236    #[test]
237    fn inactive_returns_real_value() {
238        let mut s = shadow_ts();
239        s.deactivate();
240        let r = s.intercept("charge", json!([]), json!({}), json!("shadowed"), || json!("real"));
241        assert_eq!(r, json!("real"));
242    }
243
244    #[test]
245    fn records_captured() {
246        let mut s = shadow_ts();
247        s.intercept("t1", json!([1, 2]), json!({"k": "v"}), json!(null), || json!(null));
248        assert_eq!(s.records().len(), 1);
249        let rec = &s.records()[0];
250        assert_eq!(rec.tool_name, "t1");
251        assert_eq!(rec.args, vec![json!(1), json!(2)]);
252        assert_eq!(rec.kwargs["k"], json!("v"));
253        assert_eq!(rec.ts, 1_700_000_000.0);
254        assert_eq!(rec.returned, json!(null));
255    }
256
257    #[test]
258    fn inactive_does_not_record() {
259        let mut s = shadow_ts();
260        s.deactivate();
261        s.intercept("t1", json!([]), json!({}), json!(null), || json!(null));
262        assert!(s.records().is_empty());
263    }
264
265    #[test]
266    fn toggle_activate_deactivate() {
267        let mut s = ShadowMode::new(false);
268        assert!(!s.is_active());
269        s.activate();
270        assert!(s.is_active());
271        s.deactivate();
272        assert!(!s.is_active());
273    }
274
275    #[test]
276    fn clear_records() {
277        let mut s = shadow_ts();
278        s.intercept("a", json!([]), json!({}), json!(1), || json!(1));
279        s.intercept("b", json!([]), json!({}), json!(2), || json!(2));
280        assert_eq!(s.records().len(), 2);
281        s.clear_records();
282        assert!(s.records().is_empty());
283    }
284
285    #[test]
286    fn multiple_intercepts() {
287        let mut s = shadow_ts();
288        for i in 0..5u64 {
289            s.intercept("t", json!([i]), json!({}), json!(i), || json!(i));
290        }
291        assert_eq!(s.records().len(), 5);
292    }
293
294    #[test]
295    fn shadow_value_captured_in_record() {
296        let mut s = shadow_ts();
297        s.intercept("send_email", json!([]), json!({"to": "x@y.com"}), json!({"ok": true}), || unreachable!());
298        assert_eq!(s.records()[0].returned, json!({"ok": true}));
299    }
300
301    #[test]
302    fn kwargs_captured() {
303        let mut s = shadow_ts();
304        s.intercept("f", json!([]), json!({"a": 1, "b": "c"}), json!(null), || json!(null));
305        let kw = &s.records()[0].kwargs;
306        assert_eq!(kw["a"], json!(1));
307        assert_eq!(kw["b"], json!("c"));
308    }
309
310    #[test]
311    fn jsonl_round_trip() {
312        let mut s = shadow_ts();
313        s.intercept("tool_a", json!([42]), json!({"x": "y"}), json!({"status": "ok"}), || unreachable!());
314        s.intercept("tool_b", json!([]), json!({}), json!(null), || unreachable!());
315
316        let path = std::env::temp_dir().join("shadow_test_roundtrip.jsonl");
317        let n = s.to_jsonl(&path).unwrap();
318        assert_eq!(n, 2);
319
320        let loaded = ShadowMode::load_jsonl(&path).unwrap();
321        assert_eq!(loaded.len(), 2);
322        assert_eq!(loaded[0].tool_name, "tool_a");
323        assert_eq!(loaded[0].args[0], json!(42));
324        assert_eq!(loaded[1].tool_name, "tool_b");
325        assert_eq!(loaded[1].returned, Value::Null);
326
327        std::fs::remove_file(path).ok();
328    }
329
330    #[test]
331    fn to_json_line_roundtrip() {
332        let rec = ShadowRecord {
333            tool_name: "foo".to_owned(),
334            args: vec![json!(1), json!(2)],
335            kwargs: serde_json::from_str(r#"{"k": "v"}"#).unwrap(),
336            ts: 12345.6,
337            returned: json!("ok"),
338        };
339        let line = rec.to_json_line();
340        let back = ShadowRecord::from_json_line(&line).unwrap();
341        assert_eq!(back.tool_name, "foo");
342        assert_eq!(back.args, vec![json!(1), json!(2)]);
343        assert_eq!(back.ts, 12345.6);
344    }
345
346    #[test]
347    fn from_json_line_blank_returns_none() {
348        assert!(ShadowRecord::from_json_line("").is_none());
349        assert!(ShadowRecord::from_json_line("   ").is_none());
350        assert!(ShadowRecord::from_json_line("not json").is_none());
351    }
352
353    #[test]
354    fn non_array_args_wrapped() {
355        let mut s = shadow_ts();
356        // When args is not an array, wrap it in a vec
357        s.intercept("t", json!("single_arg"), json!({}), json!(null), || json!(null));
358        assert_eq!(s.records()[0].args, vec![json!("single_arg")]);
359    }
360}
agent_shadow_mode/lib.rs

agent_shadow_mode/
lib.rs
