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

agent_stream_kit/
context.rs
