agentguard_core/
context.rs

1//! Request context: tool arguments plus session metadata.
2//!
3//! Tool args are flattened into the top-level context for schema validation;
4//! session metadata is grouped under a `session` record to avoid colliding
5//! with action-specific fields.
6
7use indexmap::IndexMap;
8use serde::{Deserialize, Serialize};
9
10/// Structured context for an agent request.
11#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
12#[non_exhaustive]
13pub struct AgentContext {
14    /// Tool arguments as a JSON object. Flattened to top level on serialization.
15    #[serde(default)]
16    pub args: serde_json::Value,
17    /// Free-form session metadata: `ip`, `user_agent`, `mfa`, `ts`, etc.
18    #[serde(default)]
19    pub session: IndexMap<String, serde_json::Value>,
20}
21
22impl AgentContext {
23    /// Construct an empty context.
24    pub fn new() -> Self {
25        Self::default()
26    }
27
28    /// Add a tool argument and return `self` for chaining.
29    pub fn with_arg(mut self, k: impl Into<String>, v: impl Into<serde_json::Value>) -> Self {
30        if !self.args.is_object() {
31            self.args = serde_json::json!({});
32        }
33        self.args
34            .as_object_mut()
35            .expect("just initialized to object")
36            .insert(k.into(), v.into());
37        self
38    }
39
40    /// Add a session metadata field and return `self` for chaining.
41    pub fn with_session(mut self, k: impl Into<String>, v: impl Into<serde_json::Value>) -> Self {
42        self.session.insert(k.into(), v.into());
43        self
44    }
45
46    /// Serialize to a JSON object suitable for `Context::from_json_str`.
47    ///
48    /// Args are flattened into the top-level so they match the action's
49    /// declared schema fields. Session metadata is grouped under a
50    /// `session` record so it doesn't collide with action-specific fields.
51    pub fn to_json_object(&self) -> serde_json::Map<String, serde_json::Value> {
52        let mut map = serde_json::Map::new();
53        if let serde_json::Value::Object(args) = &self.args {
54            for (k, v) in args {
55                map.insert(k.clone(), v.clone());
56            }
57        }
58        if !self.session.is_empty() {
59            let session_map: serde_json::Map<_, _> = self
60                .session
61                .iter()
62                .map(|(k, v)| (k.clone(), v.clone()))
63                .collect();
64            map.insert("session".into(), serde_json::Value::Object(session_map));
65        }
66        map
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn builder_appends_args_and_session() {
76        let c = AgentContext::new()
77            .with_arg("to", "[email protected]")
78            .with_session("ip", "10.0.0.1");
79        assert_eq!(c.args["to"], "[email protected]");
80        assert_eq!(c.session["ip"], "10.0.0.1");
81    }
82
83    #[test]
84    fn to_json_object_flattens_args() {
85        let c = AgentContext::new()
86            .with_arg("amount", 1000)
87            .with_session("mfa", true);
88        let obj = c.to_json_object();
89        assert_eq!(obj["amount"], 1000);
90        assert!(obj.contains_key("session"));
91    }
92
93    #[test]
94    fn empty_context_serializes_to_empty_object() {
95        let c = AgentContext::new();
96        let obj = c.to_json_object();
97        assert!(obj.is_empty());
98    }
99
100    #[test]
101    fn with_arg_overwrites_previous() {
102        let c = AgentContext::new()
103            .with_arg("to", "[email protected]")
104            .with_arg("to", "[email protected]");
105        assert_eq!(c.args["to"], "[email protected]");
106    }
107
108    #[test]
109    fn context_serde_round_trip() {
110        let c = AgentContext::new()
111            .with_arg("k", "v")
112            .with_session("s", "u");
113        let json = serde_json::to_string(&c).unwrap();
114        let parsed: AgentContext = serde_json::from_str(&json).unwrap();
115        assert_eq!(parsed, c);
116    }
117}
agentguard_core/context.rs

agentguard_core/
context.rs
