agent_framework_core/
session.rs1use std::collections::HashMap;
17use std::sync::{Arc, Mutex};
18
19use serde_json::Value;
20use uuid::Uuid;
21
22use crate::error::{Error, Result};
23use crate::memory::ContextProvider;
24
25#[derive(Clone, Default)]
37pub struct SessionState {
38 inner: Arc<Mutex<HashMap<String, Value>>>,
39}
40
41impl SessionState {
42 pub fn new() -> Self {
44 Self::default()
45 }
46
47 pub fn insert(&self, key: impl Into<String>, value: Value) -> Option<Value> {
49 self.inner.lock().unwrap().insert(key.into(), value)
50 }
51
52 pub fn get(&self, key: &str) -> Option<Value> {
54 self.inner.lock().unwrap().get(key).cloned()
55 }
56
57 pub fn remove(&self, key: &str) -> Option<Value> {
59 self.inner.lock().unwrap().remove(key)
60 }
61
62 pub fn contains_key(&self, key: &str) -> bool {
64 self.inner.lock().unwrap().contains_key(key)
65 }
66
67 pub fn len(&self) -> usize {
69 self.inner.lock().unwrap().len()
70 }
71
72 pub fn is_empty(&self) -> bool {
74 self.inner.lock().unwrap().is_empty()
75 }
76
77 pub fn snapshot(&self) -> HashMap<String, Value> {
79 self.inner.lock().unwrap().clone()
80 }
81
82 pub fn shares_storage_with(&self, other: &SessionState) -> bool {
84 Arc::ptr_eq(&self.inner, &other.inner)
85 }
86}
87
88impl From<HashMap<String, Value>> for SessionState {
89 fn from(map: HashMap<String, Value>) -> Self {
90 Self {
91 inner: Arc::new(Mutex::new(map)),
92 }
93 }
94}
95
96impl std::fmt::Debug for SessionState {
97 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98 f.debug_map().entries(self.snapshot()).finish()
99 }
100}
101
102#[derive(Clone)]
108pub struct AgentSession {
109 session_id: String,
110 service_session_id: Option<String>,
111 pub state: SessionState,
115 pub context_providers: Vec<Arc<dyn ContextProvider>>,
119}
120
121impl std::fmt::Debug for AgentSession {
122 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123 f.debug_struct("AgentSession")
124 .field("session_id", &self.session_id)
125 .field("service_session_id", &self.service_session_id)
126 .field("state", &self.state)
127 .field("context_providers", &self.context_providers.len())
128 .finish()
129 }
130}
131
132impl Default for AgentSession {
133 fn default() -> Self {
134 Self::new()
135 }
136}
137
138impl AgentSession {
139 pub fn new() -> Self {
142 Self {
143 session_id: Uuid::new_v4().to_string(),
144 service_session_id: None,
145 state: SessionState::new(),
146 context_providers: Vec::new(),
147 }
148 }
149
150 pub fn child(&self) -> AgentSession {
166 AgentSession {
167 session_id: self.session_id.clone(),
168 service_session_id: None,
169 state: self.state.clone(),
170 context_providers: Vec::new(),
171 }
172 }
173
174 pub fn service(id: impl Into<String>) -> Self {
176 Self {
177 service_session_id: Some(id.into()),
178 ..Self::new()
179 }
180 }
181
182 pub fn with_context_providers(mut self, providers: Vec<Arc<dyn ContextProvider>>) -> Self {
184 self.context_providers = providers;
185 self
186 }
187
188 pub fn session_id(&self) -> &str {
190 &self.session_id
191 }
192
193 pub fn service_session_id(&self) -> Option<&str> {
195 self.service_session_id.as_deref()
196 }
197
198 pub fn set_service_session_id(&mut self, id: impl Into<String>) {
200 self.service_session_id = Some(id.into());
201 }
202
203 pub fn try_adopt_service_session_id(&mut self, id: &str) -> bool {
210 if self.service_session_id.as_deref() == Some(id) {
211 return false;
212 }
213 self.service_session_id = Some(id.to_string());
214 true
215 }
216
217 pub fn to_dict(&self) -> Value {
223 serde_json::json!({
224 "session_id": self.session_id,
225 "service_session_id": self.service_session_id,
226 "state": self.state.snapshot(),
227 })
228 }
229
230 pub fn from_dict(state: &Value) -> Result<Self> {
236 let session_id = state
237 .get("session_id")
238 .and_then(Value::as_str)
239 .map(str::to_string)
240 .unwrap_or_else(|| Uuid::new_v4().to_string());
241 let service_session_id = state
242 .get("service_session_id")
243 .and_then(Value::as_str)
244 .map(str::to_string);
245 let session_state: HashMap<String, Value> = match state.get("state") {
246 Some(v) if !v.is_null() => serde_json::from_value(v.clone()).map_err(|e| {
247 Error::Serialization(format!("failed to restore session state: {e}"))
248 })?,
249 _ => HashMap::new(),
250 };
251 Ok(Self {
252 session_id,
253 service_session_id,
254 state: SessionState::from(session_state),
255 context_providers: Vec::new(),
256 })
257 }
258}
259
260#[cfg(test)]
261mod tests {
262 use super::*;
263
264 #[test]
265 fn new_session_has_a_generated_id_and_no_service_id() {
266 let session = AgentSession::new();
267 assert!(!session.session_id().is_empty());
268 assert!(session.service_session_id().is_none());
269 assert!(session.state.is_empty());
270 }
271
272 #[test]
273 fn service_session_sets_service_session_id_and_still_has_a_local_id() {
274 let session = AgentSession::service("svc-1");
275 assert_eq!(session.service_session_id(), Some("svc-1"));
276 assert!(!session.session_id().is_empty());
277 }
278
279 #[test]
280 fn try_adopt_service_session_id_reports_whether_it_was_new() {
281 let mut session = AgentSession::new();
282 assert!(session.try_adopt_service_session_id("conv-1"));
283 assert_eq!(session.service_session_id(), Some("conv-1"));
284 assert!(!session.try_adopt_service_session_id("conv-1"));
286 assert!(session.try_adopt_service_session_id("conv-2"));
288 assert_eq!(session.service_session_id(), Some("conv-2"));
289 }
290
291 #[test]
292 fn to_dict_from_dict_round_trips_session_id_service_id_and_state() {
293 let session = AgentSession::service("svc-9");
294 session.state.insert("key", serde_json::json!("value"));
295 let original_id = session.session_id().to_string();
296
297 let state = session.to_dict();
298 assert_eq!(state["session_id"], original_id);
299 assert_eq!(state["service_session_id"], "svc-9");
300 assert_eq!(state["state"]["key"], "value");
301 assert!(state.get("messages").is_none());
303 assert!(state.get("chat_message_store_state").is_none());
304
305 let restored = AgentSession::from_dict(&state).unwrap();
306 assert_eq!(restored.session_id(), original_id);
307 assert_eq!(restored.service_session_id(), Some("svc-9"));
308 assert_eq!(restored.state.get("key"), Some(serde_json::json!("value")));
309 assert!(restored.context_providers.is_empty());
310 }
311
312 #[test]
313 fn clones_share_the_state_bag_by_reference() {
314 let session = AgentSession::new();
315 let clone = session.clone();
316 clone
317 .state
318 .insert("written-via-clone", serde_json::json!(1));
319 assert_eq!(
320 session.state.get("written-via-clone"),
321 Some(serde_json::json!(1)),
322 "a clone must be a view onto the same state bag"
323 );
324 assert!(session.state.shares_storage_with(&clone.state));
325 }
326
327 #[test]
328 fn child_shares_id_and_state_but_isolates_the_service_pointer() {
329 let parent = AgentSession::service("svc-parent");
330 parent.state.insert("k", serde_json::json!("v"));
331
332 let child = parent.child();
333 assert_eq!(child.session_id(), parent.session_id());
334 assert_eq!(
335 child.service_session_id(),
336 None,
337 "the parent's server-side conversation pointer must not leak to the child"
338 );
339 assert!(child.context_providers.is_empty());
340
341 assert_eq!(child.state.get("k"), Some(serde_json::json!("v")));
343 child.state.insert("from-child", serde_json::json!(2));
344 assert_eq!(parent.state.get("from-child"), Some(serde_json::json!(2)));
345
346 let mut child = child;
349 child.set_service_session_id("svc-child");
350 assert_eq!(parent.service_session_id(), Some("svc-parent"));
351 }
352
353 #[test]
354 fn from_dict_generates_a_session_id_when_absent() {
355 let restored = AgentSession::from_dict(&serde_json::json!({})).unwrap();
356 assert!(!restored.session_id().is_empty());
357 assert!(restored.service_session_id().is_none());
358 assert!(restored.state.is_empty());
359 }
360}