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            mutations: super::mutations::MutationCoordinator::default(),
118            version: crate::build_info::display_version().to_string(),
119            isolation,
120        }
121    }
122
123    /// Resolve the session slot for a `thread_id`, creating it if absent.
124    /// Marks the slot as most-recently-used and evicts the LRU idle session when
125    /// adding a new key would exceed [`MAX_SESSIONS`].
126    pub(super) async fn session_for(&self, thread_id: &str) -> (SharedBashState, SessionGuard) {
127        let thread_id = crate::types::normalize_thread_id(thread_id);
128        let mut registry = self.sessions.lock().await;
129        let key = if thread_id.is_empty() {
130            match self.isolation {
131                SessionIsolation::Lenient => {
132                    registry.last_active.clone().unwrap_or_else(|| "default".to_string())
133                }
134                SessionIsolation::Strict => "anonymous".to_string(),
135            }
136        } else {
137            thread_id.clone()
138        };
139
140        if !registry.slots.contains_key(&key) && registry.slots.len() >= MAX_SESSIONS {
141            let victim = registry
142                .last_used
143                .iter()
144                .filter(|(candidate, _)| **candidate != key)
145                .filter(|(candidate, _)| {
146                    registry.in_flight.get(candidate.as_str()).is_none_or(|pin| !pin.is_pinned())
147                })
148                .min_by_key(|(_, last_used)| **last_used)
149                .map(|(candidate, _)| candidate.clone());
150
151            if let Some(victim) = victim {
152                // The daemon runtime owns a process outside this registry. Release
153                // it while the registry lock prevents the same thread_id from being
154                // recreated underneath the termination request.
155                if let Err(error) = self.shell_runtime.terminate_session(&victim).await {
156                    warn!(%error, "failed to terminate LRU shell session '{victim}'");
157                }
158                registry.slots.remove(&victim);
159                registry.last_used.remove(&victim);
160                registry.in_flight.remove(&victim);
161                if registry.last_active.as_deref() == Some(victim.as_str()) {
162                    registry.last_active = None;
163                }
164                warn!("Evicted LRU shell session '{victim}' (session cap {MAX_SESSIONS})");
165            } else {
166                warn!(
167                    "All {MAX_SESSIONS} sessions busy; exceeding the cap rather than evicting an in-flight session"
168                );
169            }
170        }
171
172        let slot =
173            registry.slots.entry(key.clone()).or_insert_with(|| Arc::new(Mutex::new(None))).clone();
174        let pin = registry.in_flight.entry(key.clone()).or_default().clone();
175        registry.last_used.insert(key.clone(), Instant::now());
176        if !thread_id.is_empty() {
177            registry.last_active = Some(key);
178        }
179        (slot, pin.acquire())
180    }
181
182    /// The most recently active session slot, optionally confined to one HTTP
183    /// principal's internal thread-id prefix.
184    pub(super) async fn active_slot(
185        &self,
186        session_prefix: Option<&str>,
187    ) -> Option<SharedBashState> {
188        let registry = self.sessions.lock().await;
189        if let Some(prefix) = session_prefix {
190            return registry
191                .last_used
192                .iter()
193                .filter(|(key, _)| key.starts_with(prefix))
194                .max_by_key(|(_, last_used)| **last_used)
195                .and_then(|(key, _)| registry.slots.get(key).cloned());
196        }
197        registry.last_active.as_ref().and_then(|key| registry.slots.get(key).cloned())
198    }
199
200    /// Read the immutable project identity of an already initialized adapter
201    /// session without creating a slot or changing LRU state. Remote coherence
202    /// checks use this before any tool can touch the shell or filesystem.
203    pub(super) async fn bound_workspace(&self, thread_id: &str) -> Option<std::path::PathBuf> {
204        let thread_id = crate::types::normalize_thread_id(thread_id);
205        if thread_id.is_empty() {
206            return None;
207        }
208        let slot = {
209            let registry = self.sessions.lock().await;
210            registry.slots.get(&thread_id).cloned()
211        }?;
212        let state = slot.lock().await;
213        state.as_ref().filter(|state| state.initialized).map(|state| state.workspace_root.clone())
214    }
215
216    pub(super) async fn has_initialized_session(&self) -> bool {
217        let slots = {
218            let registry = self.sessions.lock().await;
219            registry.slots.values().cloned().collect::<Vec<_>>()
220        };
221        for slot in slots {
222            if slot.lock().await.as_ref().is_some_and(|state| state.initialized) {
223                return true;
224            }
225        }
226        false
227    }
228
229    /// Bootstrap a local single-client session from the client's first MCP Root.
230    /// The negotiated MCP protocol still exposes Roots, although `rmcp` marks this
231    /// compatibility API as deprecated in anticipation of the later SEP-2577 draft.
232    #[allow(deprecated)]
233    pub(super) async fn initialize_from_client_roots(
234        &self,
235        context: NotificationContext<RoleServer>,
236    ) {
237        if self.isolation != SessionIsolation::Lenient {
238            return;
239        }
240        let _bootstrap_guard = self.root_bootstrap.lock().await;
241        if self.has_initialized_session().await {
242            return;
243        }
244        let supports_roots =
245            context.peer.peer_info().is_some_and(|info| info.capabilities.roots.is_some());
246        if !supports_roots {
247            return;
248        }
249        let roots = match context.peer.list_roots().await {
250            Ok(result) => result.roots,
251            Err(error) => {
252                warn!(%error, "client advertised Roots but roots/list failed");
253                return;
254            }
255        };
256        let cwd = std::env::current_dir().ok();
257        let mut paths =
258            roots.iter().filter_map(|root| root_uri_to_path(&root.uri)).collect::<Vec<_>>();
259        let selected = cwd
260            .as_ref()
261            .and_then(|cwd| paths.iter().position(|path| cwd.starts_with(path)))
262            .map(|index| paths.remove(index))
263            .or_else(|| paths.into_iter().next());
264        let Some(workspace) = selected else {
265            warn!("client returned no usable local file Roots; Initialize remains required");
266            return;
267        };
268
269        let initialize = Initialize {
270            init_type: InitializeType::FirstCall,
271            any_workspace_path: workspace.to_string_lossy().into_owned(),
272            initial_files_to_read: Vec::new(),
273            task_id_to_resume: String::new(),
274            mode_name: ModeName::Wcgw,
275            thread_id: generate_thread_id(),
276            code_writer_config: None,
277        };
278        let arguments = match serde_json::to_value(initialize) {
279            Ok(arguments) => arguments,
280            Err(error) => {
281                warn!(%error, "failed to serialize Roots bootstrap request");
282                return;
283            }
284        };
285        match self.handle_initialize(Some(arguments)).await {
286            Ok(_) => info!(workspace = %workspace.display(), "initialized Winx from MCP Roots"),
287            Err(error) => warn!(%error.message, "failed to initialize Winx from MCP Roots"),
288        }
289    }
290}
291
292pub(super) fn root_uri_to_path(uri: &str) -> Option<std::path::PathBuf> {
293    let encoded = uri.strip_prefix("file://")?;
294    let encoded = if let Some(path) = encoded.strip_prefix("localhost/") {
295        format!("/{path}")
296    } else if encoded.starts_with('/') {
297        encoded.to_string()
298    } else {
299        // Refuse non-local file authorities (file://host/path).
300        return None;
301    };
302    let decoded = percent_encoding::percent_decode_str(&encoded).decode_utf8().ok()?;
303    let path = std::path::PathBuf::from(decoded.as_ref());
304    path.is_absolute().then_some(path)
305}
winx_code_agent/server/sessions.rs

winx_code_agent/server/
sessions.rs
