modular_agent_core/
context.rs

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

modular_agent_core/
context.rs
