Skip to main content

agent_framework_core/
session.rs

1//! Conversation sessions.
2//!
3//! Rust equivalent of upstream's `_sessions.py`. Upstream deleted
4//! `_threads.py` and `_memory.py`, consolidating a lightweight `AgentSession`
5//! (`{session_id, service_session_id, state}`) with memory/context providers.
6//! Conversation history left the session entirely: it is now injected by a
7//! [`HistoryProvider`](crate::history::HistoryProvider) — just another
8//! [`ContextProvider`] — via `before_run`/`after_run`, instead of being owned
9//! by a message store on the thread/session itself.
10//!
11//! A session is no longer "service-managed XOR locally-stored": it always
12//! carries a `session_id`, optionally a `service_session_id` (when the
13//! underlying service manages the conversation server-side), a bag of
14//! free-form `state`, and the `context_providers` that run on every use.
15
16use 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/// The free-form state bag of an [`AgentSession`], shared **by reference**
26/// across clones.
27///
28/// Upstream's `AgentSession.state` is a plain Python dict, and Python
29/// sessions are reference types: every holder of the session sees the same
30/// dict, including the *child* session `as_tool(propagate_session=True)`
31/// hands a sub-agent (upstream shares the dict by reference while isolating
32/// `service_session_id`). A by-value `HashMap` cannot express that, so the
33/// state bag is an `Arc<Mutex<..>>` handle: cloning an [`AgentSession`]
34/// (or a `SessionState`) yields a view onto the *same* bag, and mutations
35/// through any clone are visible to all of them.
36#[derive(Clone, Default)]
37pub struct SessionState {
38    inner: Arc<Mutex<HashMap<String, Value>>>,
39}
40
41impl SessionState {
42    /// A fresh, empty state bag.
43    pub fn new() -> Self {
44        Self::default()
45    }
46
47    /// Insert `value` at `key`, returning the previous value if any.
48    pub fn insert(&self, key: impl Into<String>, value: Value) -> Option<Value> {
49        self.inner.lock().unwrap().insert(key.into(), value)
50    }
51
52    /// A clone of the value at `key`, if present.
53    pub fn get(&self, key: &str) -> Option<Value> {
54        self.inner.lock().unwrap().get(key).cloned()
55    }
56
57    /// Remove and return the value at `key`, if present.
58    pub fn remove(&self, key: &str) -> Option<Value> {
59        self.inner.lock().unwrap().remove(key)
60    }
61
62    /// Whether `key` is present.
63    pub fn contains_key(&self, key: &str) -> bool {
64        self.inner.lock().unwrap().contains_key(key)
65    }
66
67    /// The number of entries.
68    pub fn len(&self) -> usize {
69        self.inner.lock().unwrap().len()
70    }
71
72    /// Whether the bag is empty.
73    pub fn is_empty(&self) -> bool {
74        self.inner.lock().unwrap().is_empty()
75    }
76
77    /// A point-in-time copy of the whole bag (e.g. for serialization).
78    pub fn snapshot(&self) -> HashMap<String, Value> {
79        self.inner.lock().unwrap().clone()
80    }
81
82    /// Whether `other` is a view onto the same underlying bag.
83    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/// A conversation session: a lightweight identity + state container.
103///
104/// Message history is **not** stored here any more — see
105/// [`crate::history::HistoryProvider`] and
106/// [`crate::history::ensure_history_provider`].
107#[derive(Clone)]
108pub struct AgentSession {
109    session_id: String,
110    service_session_id: Option<String>,
111    /// Free-form session state (for context providers to persist per-session
112    /// data across runs). Shared by reference across clones of this session —
113    /// see [`SessionState`].
114    pub state: SessionState,
115    /// Context providers associated with this session (memory/RAG/history
116    /// injection). Combined with an agent's own providers at request time —
117    /// see [`Agent::combined_providers`](crate::agent::Agent).
118    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    /// A fresh, local (non-service-managed) session with a newly generated
140    /// `session_id`.
141    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    /// A **child** session for delegating part of this conversation to a
151    /// sub-agent (the `as_tool(propagate_session)` path).
152    ///
153    /// The child keeps this session's `session_id` and shares its [`state`]
154    /// bag by reference, but gets an **isolated** (cleared)
155    /// `service_session_id` and no context providers. Isolating the
156    /// server-side conversation pointer matters: after the parent's first
157    /// model call, a service-managed session carries the parent conversation
158    /// id (e.g. an OpenAI Responses `previous_response_id`); a child that
159    /// inherited it would submit a follow-up onto a conversation whose
160    /// tool call is still pending, which the server rejects. Mirrors
161    /// upstream's `_agent_wrapper` child-session construction
162    /// (`_agents.py`, microsoft/agent-framework#5875).
163    ///
164    /// [`state`]: AgentSession::state
165    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    /// A service-managed session identified by a conversation id.
175    pub fn service(id: impl Into<String>) -> Self {
176        Self {
177            service_session_id: Some(id.into()),
178            ..Self::new()
179        }
180    }
181
182    /// Attach context providers to this session, replacing any previously set.
183    pub fn with_context_providers(mut self, providers: Vec<Arc<dyn ContextProvider>>) -> Self {
184        self.context_providers = providers;
185        self
186    }
187
188    /// This session's local identifier.
189    pub fn session_id(&self) -> &str {
190        &self.session_id
191    }
192
193    /// The service-side conversation id, if any.
194    pub fn service_session_id(&self) -> Option<&str> {
195        self.service_session_id.as_deref()
196    }
197
198    /// Set the service-side conversation id explicitly.
199    pub fn set_service_session_id(&mut self, id: impl Into<String>) {
200        self.service_session_id = Some(id.into());
201    }
202
203    /// Adopt a service-managed conversation id returned by the chat service
204    /// (e.g. an OpenAI Responses `previous_response_id` or an Azure AI thread
205    /// id), so follow-up runs continue the same service conversation.
206    ///
207    /// Returns `true` when the id was newly adopted (it differed from what
208    /// the session already carried), `false` when it was already current.
209    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    /// Serialize this session's `{session_id, service_session_id, state}` to
218    /// JSON. Conversation history is deliberately **not** included — it lives
219    /// in whichever [`HistoryProvider`](crate::history::HistoryProvider), if
220    /// any, is attached to `context_providers`; serialize that separately
221    /// (e.g. [`InMemoryHistoryProvider::to_dict`](crate::history::InMemoryHistoryProvider::to_dict)).
222    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    /// Reconstruct a session from state produced by [`AgentSession::to_dict`].
231    ///
232    /// `context_providers` are **not** restored by this call — callers
233    /// reattach their own (including any `HistoryProvider`, whose own state
234    /// is serialized/restored independently).
235    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        // Adopting the same id again is a no-op, reported as such.
285        assert!(!session.try_adopt_service_session_id("conv-1"));
286        // A different id is adopted (and reported as newly adopted).
287        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        // History is deliberately absent from the wire shape.
302        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        // State is shared both ways.
342        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        // And a service id the child adopts during its own run stays local
347        // to the child rather than clobbering the parent's.
348        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}