lean_ctx/core/session/
slot.rs

1// SPDX-License-Identifier: Apache-2.0
2//! The live session of one MCP server and the instance a request is bound to.
3//!
4//! A host may change its workspace roots while a slow call is still running.
5//! Replacing the session *contents* in place let that call's late writes —
6//! receipts, intents, touched files, the batch save — land in the session of
7//! the project the host switched to. A root change therefore installs a new
8//! session instance, and a request keeps the instance it was bound to: late
9//! results belong to the original request and its project (AP01 §3).
10
11use std::cell::RefCell;
12use std::future::Future;
13use std::sync::Arc;
14
15use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};
16
17use super::SessionState;
18
19pub type SharedSession = Arc<RwLock<SessionState>>;
20type LiveCell = Arc<std::sync::RwLock<SharedSession>>;
21
22tokio::task_local! {
23    /// The session instance the current request is bound to, with the slot it
24    /// was taken from. Empty until the request pins it.
25    static BOUND: RefCell<Option<(LiveCell, SharedSession)>>;
26}
27
28/// The session instance the running request is bound to, whichever server's
29/// slot it came from. `None` outside a bound request.
30pub(super) fn bound_session() -> Option<SharedSession> {
31    BOUND
32        .try_with(|bound| {
33            bound
34                .borrow()
35                .as_ref()
36                .map(|(_, session)| Arc::clone(session))
37        })
38        .ok()
39        .flatten()
40}
41
42/// Shared between all clones of one server; the live instance is replaced only
43/// on a project change.
44#[derive(Clone)]
45pub struct SessionSlot {
46    live: LiveCell,
47    /// Per project root, the session this server worked in there before it
48    /// switched away: the one saved session it may take up again. Any other
49    /// saved session of a project is another agent's private state (AP01 §1).
50    own: Arc<std::sync::Mutex<std::collections::HashMap<String, String>>>,
51}
52
53impl SessionSlot {
54    pub fn new(session: SharedSession) -> Self {
55        Self {
56            live: Arc::new(std::sync::RwLock::new(session)),
57            own: Arc::default(),
58        }
59    }
60
61    /// Record that this server's session for `project_root` is `session_id`.
62    pub fn remember_own(&self, project_root: &str, session_id: &str) {
63        self.own
64            .lock()
65            .unwrap_or_else(std::sync::PoisonError::into_inner)
66            .insert(project_root.to_owned(), session_id.to_owned());
67    }
68
69    /// The session this server left `project_root` with, if any.
70    pub fn own_session_id(&self, project_root: &str) -> Option<String> {
71        self.own
72            .lock()
73            .unwrap_or_else(std::sync::PoisonError::into_inner)
74            .get(project_root)
75            .cloned()
76    }
77
78    /// The instance new requests bind to, ignoring any request binding.
79    pub fn live(&self) -> SharedSession {
80        Arc::clone(
81            &self
82                .live
83                .read()
84                .unwrap_or_else(std::sync::PoisonError::into_inner),
85        )
86    }
87
88    /// The instance this request is bound to, or the live one outside a
89    /// bound request.
90    pub fn instance(&self) -> SharedSession {
91        BOUND
92            .try_with(|bound| {
93                bound
94                    .borrow()
95                    .as_ref()
96                    .filter(|(slot, _)| Arc::ptr_eq(slot, &self.live))
97                    .map(|(_, session)| Arc::clone(session))
98            })
99            .ok()
100            .flatten()
101            .unwrap_or_else(|| self.live())
102    }
103
104    pub async fn read(&self) -> OwnedRwLockReadGuard<SessionState> {
105        self.instance().read_owned().await
106    }
107
108    pub async fn write(&self) -> OwnedRwLockWriteGuard<SessionState> {
109        self.instance().write_owned().await
110    }
111
112    /// Install `next` as the live instance if `current` still is it. Requests
113    /// already bound keep the previous one. `false`: another switch won.
114    pub fn replace_if_live(&self, current: &SharedSession, next: SharedSession) -> bool {
115        let mut live = self
116            .live
117            .write()
118            .unwrap_or_else(std::sync::PoisonError::into_inner);
119        if !Arc::ptr_eq(&live, current) {
120            return false;
121        }
122        *live = next;
123        true
124    }
125
126    /// Move the current request's binding to `next`: a request that switched
127    /// the session itself continues in the new one. No-op outside a bound
128    /// request.
129    pub fn rebind(&self, next: SharedSession) {
130        let _ = BOUND.try_with(|bound| {
131            if let Some((slot, session)) = bound.borrow_mut().as_mut()
132                && Arc::ptr_eq(slot, &self.live)
133            {
134                *session = next;
135            }
136        });
137    }
138
139    /// Run one request with an empty binding; [`Self::pin`] fills it.
140    ///
141    /// A request that ends bound to an instance no longer live saves its late
142    /// changes there, under that session's own project: the batch cadence
143    /// that would otherwise pick them up only runs on the live instance.
144    pub async fn scope<F: Future>(&self, request: F) -> F::Output {
145        let (output, bound) = BOUND
146            .scope(RefCell::new(None), async {
147                let output = request.await;
148                let bound = BOUND.with(|bound| bound.borrow_mut().take());
149                (output, bound)
150            })
151            .await;
152        if let Some((_, session)) = bound
153            && !Arc::ptr_eq(&session, &self.live())
154        {
155            let pending = {
156                let state = session.read().await;
157                (state.stats.unsaved_changes > 0).then(|| state.project_root.clone())
158            };
159            if let Some(project) = pending {
160                crate::core::policy::runtime::REQUEST_PROJECT
161                    .scope(
162                        RefCell::new(project.map(std::path::PathBuf::from)),
163                        SessionState::save_shared_logged(session),
164                    )
165                    .await;
166            }
167        }
168        output
169    }
170
171    /// Bind the current request to the live instance once its roots are
172    /// resolved. A request keeps its first binding; outside [`Self::scope`]
173    /// this is a no-op.
174    pub fn pin(&self) {
175        let _ = BOUND.try_with(|bound| {
176            let mut bound = bound.borrow_mut();
177            if bound.is_none() {
178                *bound = Some((Arc::clone(&self.live), self.live()));
179            }
180        });
181    }
182}
lean_ctx/core/session/slot.rs

lean_ctx/core/session/
slot.rs
