modular_agent_core/
context.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use serde::{Deserialize, Serialize};
4
5use crate::error::AgentError;
6use crate::value::AgentValue;
7
8/// Event-scoped context that identifies a single flow across agents and carries auxiliary metadata.
9///
10/// A context is created per externally triggered event (user input, timer, webhook, etc.) so that
11/// agents connected through connections can recognize they are handling the same flow. It can carry
12/// auxiliary metadata useful for processing without altering the primary payload.
13///
14/// When a single datum fans out into multiple derived items (e.g., a `map` operation), frames track
15/// the branching lineage. Because mapping can nest, frames behave like a stack to preserve ancestry.
16/// Instances are cheap to clone and return new copies instead of mutating in place.
17#[derive(Clone, Debug, Default, Serialize, Deserialize)]
18pub struct AgentContext {
19    /// Unique identifier assigned when the context is created.
20    id: usize,
21
22    /// Variables stored in this context.
23    #[serde(skip_serializing_if = "Option::is_none")]
24    vars: Option<im::HashMap<String, AgentValue>>,
25
26    /// Frame stack for tracking branching (e.g., map operations).
27    #[serde(default, skip_serializing_if = "Option::is_none")]
28    frames: Option<im::Vector<Frame>>,
29}
30
31/// Frame type for map operations.
32pub const FRAME_MAP: &str = "map";
33
34/// Key for the index value in map frames.
35pub const FRAME_KEY_INDEX: &str = "index";
36
37/// Key for the length value in map frames.
38pub const FRAME_KEY_LENGTH: &str = "length";
39
40impl AgentContext {
41    /// Creates a new context with a unique identifier and no state.
42    pub fn new() -> Self {
43        Self {
44            id: new_id(),
45            vars: None,
46            frames: None,
47        }
48    }
49
50    /// Returns the unique identifier for this context.
51    pub fn id(&self) -> usize {
52        self.id
53    }
54
55    // Variables
56
57    /// Retrieves an immutable reference to a stored variable, if present.
58    pub fn get_var(&self, key: &str) -> Option<&AgentValue> {
59        self.vars.as_ref().and_then(|vars| vars.get(key))
60    }
61
62    /// Returns a new context with the provided variable inserted while keeping the current context unchanged.
63    pub fn with_var(&self, key: String, value: AgentValue) -> Self {
64        let mut vars = if let Some(vars) = &self.vars {
65            vars.clone()
66        } else {
67            im::HashMap::new()
68        };
69        vars.insert(key, value);
70        Self {
71            id: self.id,
72            vars: Some(vars),
73            frames: self.frames.clone(),
74        }
75    }
76}
77
78// ID generation
79static CONTEXT_ID_COUNTER: AtomicUsize = AtomicUsize::new(1);
80
81/// Generates a monotonically increasing identifier for contexts.
82fn new_id() -> usize {
83    CONTEXT_ID_COUNTER.fetch_add(1, Ordering::Relaxed)
84}
85
86// Frame stack
87
88/// Describes a single stack frame captured during agent execution.
89#[derive(Clone, Debug, Serialize, Deserialize)]
90pub struct Frame {
91    /// The frame type name (e.g., "map").
92    pub name: String,
93
94    /// Data associated with this frame.
95    pub data: AgentValue,
96}
97
98fn map_frame_data(index: usize, len: usize) -> AgentValue {
99    let mut data = AgentValue::object_default();
100    let _ = data.set(
101        FRAME_KEY_INDEX.to_string(),
102        AgentValue::integer(index as i64),
103    );
104    let _ = data.set(
105        FRAME_KEY_LENGTH.to_string(),
106        AgentValue::integer(len as i64),
107    );
108    data
109}
110
111fn read_map_frame(frame: &Frame) -> Result<(usize, usize), AgentError> {
112    let idx = frame
113        .data
114        .get(FRAME_KEY_INDEX)
115        .and_then(|v| v.as_i64())
116        .ok_or_else(|| AgentError::InvalidValue("map frame missing integer index".into()))?;
117    let len = frame
118        .data
119        .get(FRAME_KEY_LENGTH)
120        .and_then(|v| v.as_i64())
121        .ok_or_else(|| AgentError::InvalidValue("map frame missing integer length".into()))?;
122    if idx < 0 || len < 1 {
123        return Err(AgentError::InvalidValue("Invalid map frame values".into()));
124    }
125    let (idx, len) = (idx as usize, len as usize);
126    if idx >= len {
127        return Err(AgentError::InvalidValue(
128            "map frame index is out of bounds".into(),
129        ));
130    }
131    Ok((idx, len))
132}
133
134impl AgentContext {
135    /// Returns the current frame stack, if any frames have been pushed.
136    pub fn frames(&self) -> Option<&im::Vector<Frame>> {
137        self.frames.as_ref()
138    }
139
140    /// Appends a new frame to the end of the stack and returns the updated context.
141    pub fn push_frame(&self, name: String, data: AgentValue) -> Self {
142        let mut frames = if let Some(frames) = &self.frames {
143            frames.clone()
144        } else {
145            im::Vector::new()
146        };
147        frames.push_back(Frame { name, data });
148        Self {
149            id: self.id,
150            vars: self.vars.clone(),
151            frames: Some(frames),
152        }
153    }
154
155    /// Pushes a map frame with index/length metadata after validating bounds.
156    pub fn push_map_frame(&self, index: usize, len: usize) -> Result<Self, AgentError> {
157        if len == 0 {
158            return Err(AgentError::InvalidValue(
159                "map frame length must be positive".into(),
160            ));
161        }
162        if index >= len {
163            return Err(AgentError::InvalidValue(
164                "map frame index is out of bounds".into(),
165            ));
166        }
167        Ok(self.push_frame(FRAME_MAP.to_string(), map_frame_data(index, len)))
168    }
169
170    /// Returns the most recent map frame's (index, length) if present at the top of the stack.
171    pub fn current_map_frame(&self) -> Result<Option<(usize, usize)>, AgentError> {
172        let frames = match self.frames() {
173            Some(frames) => frames,
174            None => return Ok(None),
175        };
176        let Some(last_index) = frames.len().checked_sub(1) else {
177            return Ok(None);
178        };
179        let Some(frame) = frames.get(last_index) else {
180            return Ok(None);
181        };
182        if frame.name != FRAME_MAP {
183            return Ok(None);
184        }
185        read_map_frame(frame).map(Some)
186    }
187
188    /// Removes the most recent map frame, erroring if the top frame is missing or not a map frame.
189    pub fn pop_map_frame(&self) -> Result<AgentContext, AgentError> {
190        let (frame, next_ctx) = self.pop_frame();
191        match frame {
192            Some(f) if f.name == FRAME_MAP => Ok(next_ctx),
193            Some(f) => Err(AgentError::InvalidValue(format!(
194                "Unexpected frame '{}', expected map",
195                f.name
196            ))),
197            None => Err(AgentError::InvalidValue(
198                "Missing map frame in context".into(),
199            )),
200        }
201    }
202
203    /// Collects all map frame (index, length) tuples in order, validating each entry.
204    pub fn map_frame_indices(&self) -> Result<Vec<(usize, usize)>, AgentError> {
205        let mut indices = Vec::new();
206        let Some(frames) = self.frames() else {
207            return Ok(indices);
208        };
209        for frame in frames.iter() {
210            if frame.name != FRAME_MAP {
211                continue;
212            }
213            let (idx, len) = read_map_frame(frame)?;
214            indices.push((idx, len));
215        }
216        Ok(indices)
217    }
218
219    /// Returns a stable key combining the context id with all map frame indices, if present.
220    pub fn ctx_key(&self) -> Result<String, AgentError> {
221        let map_frames = self.map_frame_indices()?;
222        if map_frames.is_empty() {
223            return Ok(self.id().to_string());
224        }
225        let parts: Vec<String> = map_frames
226            .iter()
227            .map(|(idx, len)| format!("{}:{}", idx, len))
228            .collect();
229        Ok(format!("{}:{}", self.id(), parts.join(",")))
230    }
231
232    /// Removes the most recently pushed frame and returns it together with the updated context.
233    /// If the stack is empty, `None` is returned alongside an unchanged context.
234    pub fn pop_frame(&self) -> (Option<Frame>, Self) {
235        if let Some(frames) = &self.frames {
236            if frames.is_empty() {
237                return (None, self.clone());
238            }
239            let mut frames = frames.clone();
240            let last = frames.pop_back().unwrap(); // safe unwrap after is_empty check
241
242            let new_frames = if frames.is_empty() {
243                None
244            } else {
245                Some(frames)
246            };
247            return (
248                Some(last),
249                Self {
250                    id: self.id,
251                    vars: self.vars.clone(),
252                    frames: new_frames,
253                },
254            );
255        }
256        (None, self.clone())
257    }
258}
259
260// Tests
261#[cfg(test)]
262mod tests {
263    use super::*;
264    use serde_json::json;
265
266    #[test]
267    fn new_assigns_unique_ids() {
268        let ctx1 = AgentContext::new();
269        let ctx2 = AgentContext::new();
270
271        assert_ne!(ctx1.id(), 0);
272        assert_ne!(ctx2.id(), 0);
273        assert_ne!(ctx1.id(), ctx2.id());
274        assert_eq!(ctx1.id(), ctx1.clone().id());
275    }
276
277    #[test]
278    fn with_var_sets_value_without_mutating_original() {
279        let ctx = AgentContext::new();
280        assert!(ctx.get_var("answer").is_none());
281
282        let updated = ctx.with_var("answer".into(), AgentValue::integer(42));
283
284        assert!(ctx.get_var("answer").is_none());
285        assert_eq!(updated.get_var("answer"), Some(&AgentValue::integer(42)));
286        assert_eq!(ctx.id(), updated.id());
287    }
288
289    #[test]
290    fn push_and_pop_frames() {
291        let ctx = AgentContext::new();
292        assert!(ctx.frames().is_none());
293
294        let ctx = ctx
295            .push_frame("first".into(), AgentValue::string("a"))
296            .push_frame("second".into(), AgentValue::integer(2));
297
298        let frames = ctx.frames().expect("frames should be present");
299        assert_eq!(frames.len(), 2);
300        assert_eq!(frames[0].name, "first");
301        assert_eq!(frames[1].name, "second");
302        assert_eq!(frames[1].data, AgentValue::integer(2));
303
304        let (popped_second, ctx) = ctx.pop_frame();
305        let popped_second = popped_second.expect("second frame should exist");
306        assert_eq!(popped_second.name, "second");
307        assert_eq!(ctx.frames().unwrap().len(), 1);
308        assert_eq!(ctx.frames().unwrap()[0].name, "first");
309
310        let (popped_first, ctx) = ctx.pop_frame();
311        assert_eq!(popped_first.unwrap().name, "first");
312        assert!(ctx.frames().is_none());
313
314        let (no_frame, ctx_after_empty) = ctx.pop_frame();
315        assert!(no_frame.is_none());
316        assert!(ctx_after_empty.frames().is_none());
317    }
318
319    #[test]
320    fn clone_preserves_vars() {
321        let ctx = AgentContext::new().with_var("key".into(), AgentValue::integer(1));
322        let cloned = ctx.clone();
323
324        assert_eq!(cloned.get_var("key"), Some(&AgentValue::integer(1)));
325        assert_eq!(cloned.id(), ctx.id());
326    }
327
328    #[test]
329    fn clone_preserves_frames() {
330        let ctx = AgentContext::new().push_frame("frame".into(), AgentValue::string("data"));
331        let cloned = ctx.clone();
332
333        let frames = cloned.frames().expect("cloned frames should exist");
334        assert_eq!(frames.len(), 1);
335        assert_eq!(frames[0].name, "frame");
336        assert_eq!(frames[0].data, AgentValue::string("data"));
337        assert_eq!(cloned.id(), ctx.id());
338    }
339
340    #[test]
341    fn serialization_skips_empty_optional_fields() {
342        let ctx = AgentContext::new();
343        let json_ctx = serde_json::to_value(&ctx).unwrap();
344
345        assert!(json_ctx.get("id").and_then(|v| v.as_u64()).is_some());
346        assert!(json_ctx.get("vars").is_none());
347        assert!(json_ctx.get("frames").is_none());
348
349        let populated = ctx
350            .with_var("key".into(), AgentValue::string("value"))
351            .push_frame("frame".into(), AgentValue::integer(1));
352        let json_populated = serde_json::to_value(&populated).unwrap();
353
354        assert_eq!(json_populated["vars"]["key"], json!("value"));
355        let frames = json_populated["frames"]
356            .as_array()
357            .expect("frames should serialize as array");
358        assert_eq!(frames.len(), 1);
359        assert_eq!(frames[0]["name"], json!("frame"));
360        assert_eq!(frames[0]["data"], json!(1));
361    }
362
363    #[test]
364    fn map_frame_helpers_validate_and_track_indices() -> Result<(), AgentError> {
365        let ctx = AgentContext::new();
366        let ctx = ctx.push_map_frame(0, 2)?;
367        let ctx = ctx.push_map_frame(1, 3)?;
368
369        let indices = ctx.map_frame_indices()?;
370        assert_eq!(indices, vec![(0, 2), (1, 3)]);
371
372        let current = ctx.current_map_frame()?.expect("map frame should exist");
373        assert_eq!(current, (1, 3));
374
375        let key = ctx.ctx_key()?;
376        assert_eq!(key, format!("{}:0:2,1:3", ctx.id()));
377
378        let ctx = ctx.pop_map_frame()?;
379        let current_after_pop = ctx.current_map_frame()?.expect("map frame should remain");
380        assert_eq!(current_after_pop, (0, 2));
381
382        Ok(())
383    }
384
385    #[test]
386    fn pop_map_frame_errors_when_missing_or_wrong_kind() {
387        let ctx = AgentContext::new();
388        assert!(ctx.pop_map_frame().is_err());
389
390        let ctx = ctx.push_frame("other".into(), AgentValue::unit());
391        assert!(ctx.pop_map_frame().is_err());
392    }
393
394    #[test]
395    fn push_map_frame_rejects_invalid_bounds() {
396        let ctx = AgentContext::new();
397        assert!(ctx.push_map_frame(0, 0).is_err());
398        assert!(ctx.push_map_frame(2, 1).is_err());
399    }
400}
modular_agent_core/context.rs

modular_agent_core/
context.rs
