agent_stream_kit/
context.rs

1use std::{
2    collections::BTreeMap,
3    sync::{
4        Arc,
5        atomic::{AtomicUsize, Ordering},
6    },
7};
8
9use serde::{Deserialize, Serialize};
10
11use super::data::AgentValue;
12
13#[derive(Clone, Debug, Default, Serialize, Deserialize)]
14pub struct AgentContext {
15    id: usize,
16
17    port: String,
18
19    #[serde(skip_serializing_if = "Option::is_none")]
20    vars: Option<Arc<BTreeMap<String, AgentValue>>>,
21}
22
23impl AgentContext {
24    pub fn new() -> Self {
25        Self {
26            id: new_id(),
27            port: "".to_string(),
28            vars: None,
29        }
30    }
31
32    // Port
33
34    pub fn new_with_port(port: impl Into<String>) -> Self {
35        Self {
36            id: new_id(),
37            port: port.into(),
38            vars: None,
39        }
40    }
41
42    pub fn with_port(&self, port: impl Into<String>) -> Self {
43        Self {
44            id: self.id,
45            port: port.into(),
46            vars: self.vars.clone(),
47        }
48    }
49
50    pub fn id(&self) -> usize {
51        self.id
52    }
53
54    pub fn port(&self) -> &str {
55        &self.port
56    }
57
58    // Variables
59
60    pub fn get_var(&self, key: &str) -> Option<&AgentValue> {
61        self.vars.as_ref().and_then(|vars| vars.get(key))
62    }
63
64    pub fn with_var(&self, key: String, value: AgentValue) -> Self {
65        let mut vars = if let Some(vars) = &self.vars {
66            vars.as_ref().clone()
67        } else {
68            BTreeMap::new()
69        };
70        vars.insert(key, value);
71        Self {
72            id: self.id,
73            port: self.port.clone(),
74            vars: Some(Arc::new(vars)),
75        }
76    }
77}
78
79static CONTEXT_ID_COUNTER: AtomicUsize = AtomicUsize::new(1);
80
81fn new_id() -> usize {
82    CONTEXT_ID_COUNTER.fetch_add(1, Ordering::Relaxed)
83}
agent_stream_kit/context.rs

agent_stream_kit/
context.rs
