winx_code_agent/server/
sessions.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::Arc;
4use std::time::Instant;
5
6use rmcp::service::{NotificationContext, RoleServer};
7use tokio::sync::Mutex;
8use tracing::{info, warn};
9
10use super::{SharedBashState, WinxService};
11use crate::runtime::{EmbeddedShellRuntime, ShellRuntime};
12use crate::state::bash_state::generate_thread_id;
13use crate::state::task_state::TaskRegistry;
14use crate::types::{Initialize, InitializeType, ModeName};
15
16/// Upper bound on concurrently-live adapter sessions. Each embedded session owns
17/// a PTY; daemon-backed sessions also release their guardian when evicted.
18pub(super) const MAX_SESSIONS: usize = 32;
19
20/// Per-`thread_id` shell sessions. Each `thread_id` gets its own
21/// `BashState`/PTY, so concurrent threads (or HTTP clients sharing the service)
22/// never execute in each other's shell. Tools that don't carry a `thread_id`
23/// (legacy clients) fall back to the most recently active session.
24#[derive(Default)]
25pub(super) struct SessionRegistry {
26    pub(super) slots: HashMap<String, SharedBashState>,
27    /// Last-use timestamps for LRU eviction.
28    pub(super) last_used: HashMap<String, Instant>,
29    /// In-flight operation count per session. A session whose count is > 0 (it
30    /// has a live [`SessionGuard`]) is never chosen as an LRU eviction victim.
31    pub(super) in_flight: HashMap<String, SessionPin>,
32    /// Most recently addressed session, used as the fallback for tool calls
33    /// that omit a `thread_id`.
34    pub(super) last_active: Option<String>,
35}
36
37/// Lock-free pin counter for one session. A live [`SessionGuard`] keeps the
38/// count `> 0`, which marks the session as in-flight so LRU eviction skips it.
39/// Clones share the same counter — the registry hands a clone to each
40/// concurrent call. Keep this discipline in lockstep with the loom model.
41#[derive(Clone)]
42pub(super) struct SessionPin {
43    count: Arc<AtomicUsize>,
44}
45
46impl Default for SessionPin {
47    fn default() -> Self {
48        Self { count: Arc::new(AtomicUsize::new(0)) }
49    }
50}
51
52impl SessionPin {
53    fn acquire(&self) -> SessionGuard {
54        self.count.fetch_add(1, Ordering::SeqCst);
55        SessionGuard { count: self.count.clone() }
56    }
57
58    fn is_pinned(&self) -> bool {
59        self.count.load(Ordering::SeqCst) > 0
60    }
61}
62
63/// RAII marker that a session has an operation in flight. Bumps the session's
64/// in-flight counter on creation and releases it from synchronous `Drop`.
65pub(super) struct SessionGuard {
66    count: Arc<AtomicUsize>,
67}
68
69impl Drop for SessionGuard {
70    fn drop(&mut self) {
71        self.count.fetch_sub(1, Ordering::SeqCst);
72    }
73}
74
75/// How an empty `thread_id` is resolved by the session registry.
76#[derive(Clone, Copy, PartialEq, Eq, Debug)]
77pub enum SessionIsolation {
78    /// Local stdio transport — a single client. An empty `thread_id` falls back
79    /// to the last-active session (legacy single-client convenience).
80    Lenient,
81    /// Remote HTTP transport — multiple clients behind one shared bearer token.
82    /// The last-active fallback is disabled, so an empty `thread_id` from one
83    /// client can never resolve to another client's shell; such calls get a
84    /// dedicated anonymous slot instead.
85    ///
86    /// Residual: two clients that deliberately send the same explicit
87    /// `thread_id` still share a shell. Closing that needs per-client identities.
88    Strict,
89}
90
91impl Default for WinxService {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97impl WinxService {
98    /// Create a new `WinxService` for local stdio transport.
99    pub fn new() -> Self {
100        Self::with_isolation(SessionIsolation::Lenient)
101    }
102
103    /// Create a service with an explicit session-isolation policy. The HTTP
104    /// transport uses [`SessionIsolation::Strict`].
105    pub fn with_isolation(isolation: SessionIsolation) -> Self {
106        Self::with_runtime(isolation, Arc::new(EmbeddedShellRuntime))
107    }
108
109    /// Create a service with an explicit shell-runtime boundary.
110    pub fn with_runtime(isolation: SessionIsolation, shell_runtime: Arc<dyn ShellRuntime>) -> Self {
111        info!(?isolation, "Creating new WinxService instance");
112        Self {
113            sessions: Arc::new(Mutex::new(SessionRegistry::default())),
114            tasks: Arc::new(Mutex::new(TaskRegistry::default())),
115            root_bootstrap: Arc::new(Mutex::new(())),
116            shell_runtime,
117            version: env!("CARGO_PKG_VERSION").to_string(),
118            isolation,
119        }
120    }
121
122    /// Resolve the session slot for a `thread_id`, creating it if absent.
123    /// Marks the slot as most-recently-used and evicts the LRU idle session when
124    /// adding a new key would exceed [`MAX_SESSIONS`].
125    pub(super) async fn session_for(&self, thread_id: &str) -> (SharedBashState, SessionGuard) {
126        let thread_id = crate::types::normalize_thread_id(thread_id);
127        let mut registry = self.sessions.lock().await;
128        let key = if thread_id.is_empty() {
129            match self.isolation {
130                SessionIsolation::Lenient => {
131                    registry.last_active.clone().unwrap_or_else(|| "default".to_string())
132                }
133                SessionIsolation::Strict => "anonymous".to_string(),
134            }
135        } else {
136            thread_id.clone()
137        };
138
139        if !registry.slots.contains_key(&key) && registry.slots.len() >= MAX_SESSIONS {
140            let victim = registry
141                .last_used
142                .iter()
143                .filter(|(candidate, _)| **candidate != key)
144                .filter(|(candidate, _)| {
145                    registry.in_flight.get(candidate.as_str()).is_none_or(|pin| !pin.is_pinned())
146                })
147                .min_by_key(|(_, last_used)| **last_used)
148                .map(|(candidate, _)| candidate.clone());
149
150            if let Some(victim) = victim {
151                // The daemon runtime owns a process outside this registry. Release
152                // it while the registry lock prevents the same thread_id from being
153                // recreated underneath the termination request.
154                if let Err(error) = self.shell_runtime.terminate_session(&victim).await {
155                    warn!(%error, "failed to terminate LRU shell session '{victim}'");
156                }
157                registry.slots.remove(&victim);
158                registry.last_used.remove(&victim);
159                registry.in_flight.remove(&victim);
160                if registry.last_active.as_deref() == Some(victim.as_str()) {
161                    registry.last_active = None;
162                }
163                warn!("Evicted LRU shell session '{victim}' (session cap {MAX_SESSIONS})");
164            } else {
165                warn!(
166                    "All {MAX_SESSIONS} sessions busy; exceeding the cap rather than evicting an in-flight session"
167                );
168            }
169        }
170
171        let slot =
172            registry.slots.entry(key.clone()).or_insert_with(|| Arc::new(Mutex::new(None))).clone();
173        let pin = registry.in_flight.entry(key.clone()).or_default().clone();
174        registry.last_used.insert(key.clone(), Instant::now());
175        if !thread_id.is_empty() {
176            registry.last_active = Some(key);
177        }
178        (slot, pin.acquire())
179    }
180
181    /// The most recently active session slot, optionally confined to one HTTP
182    /// principal's internal thread-id prefix.
183    pub(super) async fn active_slot(
184        &self,
185        session_prefix: Option<&str>,
186    ) -> Option<SharedBashState> {
187        let registry = self.sessions.lock().await;
188        if let Some(prefix) = session_prefix {
189            return registry
190                .last_used
191                .iter()
192                .filter(|(key, _)| key.starts_with(prefix))
193                .max_by_key(|(_, last_used)| **last_used)
194                .and_then(|(key, _)| registry.slots.get(key).cloned());
195        }
196        registry.last_active.as_ref().and_then(|key| registry.slots.get(key).cloned())
197    }
198
199    /// Read the immutable project identity of an already initialized adapter
200    /// session without creating a slot or changing LRU state. Remote coherence
201    /// checks use this before any tool can touch the shell or filesystem.
202    pub(super) async fn bound_workspace(&self, thread_id: &str) -> Option<std::path::PathBuf> {
203        let thread_id = crate::types::normalize_thread_id(thread_id);
204        if thread_id.is_empty() {
205            return None;
206        }
207        let slot = {
208            let registry = self.sessions.lock().await;
209            registry.slots.get(&thread_id).cloned()
210        }?;
211        let state = slot.lock().await;
212        state.as_ref().filter(|state| state.initialized).map(|state| state.workspace_root.clone())
213    }
214
215    pub(super) async fn has_initialized_session(&self) -> bool {
216        let slots = {
217            let registry = self.sessions.lock().await;
218            registry.slots.values().cloned().collect::<Vec<_>>()
219        };
220        for slot in slots {
221            if slot.lock().await.as_ref().is_some_and(|state| state.initialized) {
222                return true;
223            }
224        }
225        false
226    }
227
228    /// Bootstrap a local single-client session from the client's first MCP Root.
229    /// The negotiated MCP protocol still exposes Roots, although `rmcp` marks this
230    /// compatibility API as deprecated in anticipation of the later SEP-2577 draft.
231    #[allow(deprecated)]
232    pub(super) async fn initialize_from_client_roots(
233        &self,
234        context: NotificationContext<RoleServer>,
235    ) {
236        if self.isolation != SessionIsolation::Lenient {
237            return;
238        }
239        let _bootstrap_guard = self.root_bootstrap.lock().await;
240        if self.has_initialized_session().await {
241            return;
242        }
243        let supports_roots =
244            context.peer.peer_info().is_some_and(|info| info.capabilities.roots.is_some());
245        if !supports_roots {
246            return;
247        }
248        let roots = match context.peer.list_roots().await {
249            Ok(result) => result.roots,
250            Err(error) => {
251                warn!(%error, "client advertised Roots but roots/list failed");
252                return;
253            }
254        };
255        let cwd = std::env::current_dir().ok();
256        let mut paths =
257            roots.iter().filter_map(|root| root_uri_to_path(&root.uri)).collect::<Vec<_>>();
258        let selected = cwd
259            .as_ref()
260            .and_then(|cwd| paths.iter().position(|path| cwd.starts_with(path)))
261            .map(|index| paths.remove(index))
262            .or_else(|| paths.into_iter().next());
263        let Some(workspace) = selected else {
264            warn!("client returned no usable local file Roots; Initialize remains required");
265            return;
266        };
267
268        let initialize = Initialize {
269            init_type: InitializeType::FirstCall,
270            any_workspace_path: workspace.to_string_lossy().into_owned(),
271            initial_files_to_read: Vec::new(),
272            task_id_to_resume: String::new(),
273            mode_name: ModeName::Wcgw,
274            thread_id: generate_thread_id(),
275            code_writer_config: None,
276        };
277        let arguments = match serde_json::to_value(initialize) {
278            Ok(arguments) => arguments,
279            Err(error) => {
280                warn!(%error, "failed to serialize Roots bootstrap request");
281                return;
282            }
283        };
284        match self.handle_initialize(Some(arguments)).await {
285            Ok(_) => info!(workspace = %workspace.display(), "initialized Winx from MCP Roots"),
286            Err(error) => warn!(%error.message, "failed to initialize Winx from MCP Roots"),
287        }
288    }
289}
290
291pub(super) fn root_uri_to_path(uri: &str) -> Option<std::path::PathBuf> {
292    let encoded = uri.strip_prefix("file://")?;
293    let encoded = if let Some(path) = encoded.strip_prefix("localhost/") {
294        format!("/{path}")
295    } else if encoded.starts_with('/') {
296        encoded.to_string()
297    } else {
298        // Refuse non-local file authorities (file://host/path).
299        return None;
300    };
301    let decoded = percent_encoding::percent_decode_str(&encoded).decode_utf8().ok()?;
302    let path = std::path::PathBuf::from(decoded.as_ref());
303    path.is_absolute().then_some(path)
304}
winx_code_agent/server/sessions.rs

winx_code_agent/server/
sessions.rs
