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
16pub(super) const MAX_SESSIONS: usize = 32;
19
20#[derive(Default)]
25pub(super) struct SessionRegistry {
26 pub(super) slots: HashMap<String, SharedBashState>,
27 pub(super) last_used: HashMap<String, Instant>,
29 pub(super) in_flight: HashMap<String, SessionPin>,
32 pub(super) last_active: Option<String>,
35}
36
37#[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
63pub(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#[derive(Clone, Copy, PartialEq, Eq, Debug)]
77pub enum SessionIsolation {
78 Lenient,
81 Strict,
89}
90
91impl Default for WinxService {
92 fn default() -> Self {
93 Self::new()
94 }
95}
96
97impl WinxService {
98 pub fn new() -> Self {
100 Self::with_isolation(SessionIsolation::Lenient)
101 }
102
103 pub fn with_isolation(isolation: SessionIsolation) -> Self {
106 Self::with_runtime(isolation, Arc::new(EmbeddedShellRuntime))
107 }
108
109 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 execution_history: Arc::new(super::executions::history::WorkflowHistory::default()),
119 version: crate::build_info::display_version().to_string(),
120 isolation,
121 }
122 }
123
124 pub(super) async fn session_for(&self, thread_id: &str) -> (SharedBashState, SessionGuard) {
128 let thread_id = crate::types::normalize_thread_id(thread_id);
129 let mut registry = self.sessions.lock().await;
130 let key = if thread_id.is_empty() {
131 match self.isolation {
132 SessionIsolation::Lenient => {
133 registry.last_active.clone().unwrap_or_else(|| "default".to_string())
134 }
135 SessionIsolation::Strict => "anonymous".to_string(),
136 }
137 } else {
138 thread_id.clone()
139 };
140
141 if !registry.slots.contains_key(&key) && registry.slots.len() >= MAX_SESSIONS {
142 let victim = registry
143 .last_used
144 .iter()
145 .filter(|(candidate, _)| **candidate != key)
146 .filter(|(candidate, _)| {
147 registry.in_flight.get(candidate.as_str()).is_none_or(|pin| !pin.is_pinned())
148 })
149 .min_by_key(|(_, last_used)| **last_used)
150 .map(|(candidate, _)| candidate.clone());
151
152 if let Some(victim) = victim {
153 if let Err(error) = self.shell_runtime.terminate_session(&victim).await {
157 warn!(%error, "failed to terminate LRU shell session '{victim}'");
158 }
159 registry.slots.remove(&victim);
160 registry.last_used.remove(&victim);
161 registry.in_flight.remove(&victim);
162 if registry.last_active.as_deref() == Some(victim.as_str()) {
163 registry.last_active = None;
164 }
165 warn!("Evicted LRU shell session '{victim}' (session cap {MAX_SESSIONS})");
166 } else {
167 warn!(
168 "All {MAX_SESSIONS} sessions busy; exceeding the cap rather than evicting an in-flight session"
169 );
170 }
171 }
172
173 let slot =
174 registry.slots.entry(key.clone()).or_insert_with(|| Arc::new(Mutex::new(None))).clone();
175 let pin = registry.in_flight.entry(key.clone()).or_default().clone();
176 registry.last_used.insert(key.clone(), Instant::now());
177 if !thread_id.is_empty() {
178 registry.last_active = Some(key);
179 }
180 (slot, pin.acquire())
181 }
182
183 pub(super) async fn tool_session_for(
195 &self,
196 thread_id: &str,
197 ) -> (SharedBashState, SessionGuard) {
198 let (slot, guard) = self.session_for(thread_id).await;
199 self.rehydrate_slot_if_persisted(&slot, thread_id).await;
200 (slot, guard)
201 }
202
203 async fn rehydrate_slot_if_persisted(&self, slot: &SharedBashState, thread_id: &str) {
208 use crate::runtime::ShellSessionTransition;
209
210 let thread_id = crate::types::normalize_thread_id(thread_id);
211 if thread_id.is_empty() {
212 return;
213 }
214 let mut state_guard = slot.lock().await;
215 if state_guard.is_some() {
216 return;
217 }
218 let snapshot = match crate::state::persistence::load_bash_state(&thread_id) {
219 Ok(Some(snapshot)) => snapshot,
220 Ok(None) => return,
221 Err(error) => {
222 warn!(%error, thread_id, "failed to read persisted session snapshot");
223 return;
224 }
225 };
226
227 let mut state = crate::state::bash_state::BashState::new();
228 state.apply_snapshot(&snapshot);
229 state.current_thread_id.clone_from(&thread_id);
231 if !state.workspace_root.exists() {
232 warn!(
233 thread_id,
234 workspace_root = %state.workspace_root.display(),
235 "not rehydrating session: persisted workspace no longer exists"
236 );
237 return;
238 }
239
240 match self
245 .shell_runtime
246 .configure_session(&mut state, ShellSessionTransition::FirstCall)
247 .await
248 {
249 Ok(configured) => {
250 info!(
251 thread_id,
252 attached_existing = configured.attached_existing,
253 "rehydrated session from persisted state after adapter restart"
254 );
255 if !configured.attached_existing {
256 state.recovery_note = Some(lost_shell_recovery_note(&thread_id, &state));
257 }
258 *state_guard = Some(state);
259 }
260 Err(error) => {
261 warn!(%error, thread_id, "failed to rehydrate persisted session");
262 }
263 }
264 }
265
266 pub(super) async fn active_slot(
269 &self,
270 session_prefix: Option<&str>,
271 ) -> Option<SharedBashState> {
272 let registry = self.sessions.lock().await;
273 if let Some(prefix) = session_prefix {
274 return registry
275 .last_used
276 .iter()
277 .filter(|(key, _)| key.starts_with(prefix))
278 .max_by_key(|(_, last_used)| **last_used)
279 .and_then(|(key, _)| registry.slots.get(key).cloned());
280 }
281 registry.last_active.as_ref().and_then(|key| registry.slots.get(key).cloned())
282 }
283
284 pub(super) async fn bound_workspace(&self, thread_id: &str) -> Option<std::path::PathBuf> {
288 let thread_id = crate::types::normalize_thread_id(thread_id);
289 if thread_id.is_empty() {
290 return None;
291 }
292 let slot = {
293 let registry = self.sessions.lock().await;
294 registry.slots.get(&thread_id).cloned()
295 }?;
296 let state = slot.lock().await;
297 state.as_ref().filter(|state| state.initialized).map(|state| state.workspace_root.clone())
298 }
299
300 pub(super) async fn has_initialized_session(&self) -> bool {
301 let slots = {
302 let registry = self.sessions.lock().await;
303 registry.slots.values().cloned().collect::<Vec<_>>()
304 };
305 for slot in slots {
306 if slot.lock().await.as_ref().is_some_and(|state| state.initialized) {
307 return true;
308 }
309 }
310 false
311 }
312
313 #[allow(deprecated)]
317 pub(super) async fn initialize_from_client_roots(
318 &self,
319 context: NotificationContext<RoleServer>,
320 ) {
321 if self.isolation != SessionIsolation::Lenient {
322 return;
323 }
324 let _bootstrap_guard = self.root_bootstrap.lock().await;
325 if self.has_initialized_session().await {
326 return;
327 }
328 let supports_roots =
329 context.peer.peer_info().is_some_and(|info| info.capabilities.roots.is_some());
330 if !supports_roots {
331 return;
332 }
333 let roots = match context.peer.list_roots().await {
334 Ok(result) => result.roots,
335 Err(error) => {
336 warn!(%error, "client advertised Roots but roots/list failed");
337 return;
338 }
339 };
340 let cwd = std::env::current_dir().ok();
341 let mut paths =
342 roots.iter().filter_map(|root| root_uri_to_path(&root.uri)).collect::<Vec<_>>();
343 let selected = cwd
344 .as_ref()
345 .and_then(|cwd| paths.iter().position(|path| cwd.starts_with(path)))
346 .map(|index| paths.remove(index))
347 .or_else(|| paths.into_iter().next());
348 let Some(workspace) = selected else {
349 warn!("client returned no usable local file Roots; Initialize remains required");
350 return;
351 };
352
353 let initialize = Initialize {
354 init_type: InitializeType::FirstCall,
355 any_workspace_path: workspace.to_string_lossy().into_owned(),
356 initial_files_to_read: Vec::new(),
357 task_id_to_resume: String::new(),
358 mode_name: ModeName::Wcgw,
359 thread_id: generate_thread_id(),
360 code_writer_config: None,
361 };
362 let arguments = match serde_json::to_value(initialize) {
363 Ok(arguments) => arguments,
364 Err(error) => {
365 warn!(%error, "failed to serialize Roots bootstrap request");
366 return;
367 }
368 };
369 match self.handle_initialize(Some(arguments)).await {
370 Ok(_) => info!(workspace = %workspace.display(), "initialized Winx from MCP Roots"),
371 Err(error) => warn!(%error.message, "failed to initialize Winx from MCP Roots"),
372 }
373 }
374}
375
376fn lost_shell_recovery_note(
380 thread_id: &str,
381 state: &crate::state::bash_state::BashState,
382) -> String {
383 let base = format!(
384 "Recovered session {thread_id} from its persisted snapshot after a restart: mode, \
385 permissions, and edit receipts were restored, and a fresh shell was started at {}.",
386 state.cwd.display()
387 );
388 match crate::state::agent_resume::load_agent_session(thread_id) {
389 Ok(Some(record)) => {
390 let advice = record.recovery_advice_for(thread_id, &state.workspace_root);
391 let classification = advice.classification.as_str();
392 let detail = match advice.resume_argv {
393 Some(argv) => format!(
394 "The recorded `{}` launch explicitly selected a conversation target. \
395 Recovery classification={classification}, source=launch_argument, \
396 current_conversation_verified=false. To target that recorded ID, use `{}` \
397 (recorded cwd {}). Winx has not verified backend acceptance or the \
398 conversation currently selected after an in-TUI switch; no recovery \
399 command was executed.",
400 record.agent,
401 argv.join(" "),
402 record.cwd
403 ),
404 None => format!(
405 "Recorded agent: `{}`. Recovery classification={classification}, \
406 current_conversation_verified=false. No supported, scope-bound \
407 conversation selector is available. Directory-most-recent continuation \
408 may select a different conversation; no recovery command was executed.",
409 record.agent
410 ),
411 };
412 format!("{base} {detail}")
413 }
414 Ok(None) => base,
415 Err(error) => {
416 warn!(%error, thread_id, "failed to read agent-session record");
417 base
418 }
419 }
420}
421
422pub(super) fn root_uri_to_path(uri: &str) -> Option<std::path::PathBuf> {
423 let encoded = uri.strip_prefix("file://")?;
424 let encoded = if let Some(path) = encoded.strip_prefix("localhost/") {
425 format!("/{path}")
426 } else if encoded.starts_with('/') {
427 encoded.to_string()
428 } else {
429 return None;
431 };
432 let decoded = percent_encoding::percent_decode_str(&encoded).decode_utf8().ok()?;
433 let decoded: &str = if cfg!(windows) {
435 decoded.strip_prefix('/').filter(|rest| is_drive_prefixed(rest)).unwrap_or(&decoded)
436 } else {
437 &decoded
438 };
439 let path = std::path::PathBuf::from(decoded);
440 path.is_absolute().then_some(path)
441}
442
443fn is_drive_prefixed(path: &str) -> bool {
444 let mut chars = path.chars();
445 matches!((chars.next(), chars.next()), (Some(drive), Some(':')) if drive.is_ascii_alphabetic())
446}
447
448#[cfg(test)]
449#[allow(clippy::expect_used)]
450mod rehydration_tests {
451 use std::sync::Arc;
452
453 use super::SessionIsolation;
454 use crate::errors::WinxError;
455 use crate::runtime::{
456 ShellRuntime, ShellRuntimeConfigureFuture, ShellRuntimeFuture, ShellRuntimeUnitFuture,
457 ShellSessionConfiguration, ShellSessionTransition,
458 };
459 use crate::server::WinxService;
460 use crate::state::agent_resume::{save_agent_session, AgentSessionRecord};
461 use crate::state::bash_state::BashState;
462 use crate::state::persistence::{delete_bash_state, save_bash_state};
463
464 struct StubRuntime {
467 attached_existing: bool,
468 }
469
470 impl ShellRuntime for StubRuntime {
471 fn configure_session<'a>(
472 &'a self,
473 _bash_state: &'a mut BashState,
474 _transition: ShellSessionTransition,
475 ) -> ShellRuntimeConfigureFuture<'a> {
476 let attached_existing = self.attached_existing;
477 Box::pin(async move {
478 Ok(ShellSessionConfiguration { attach_hint: None, attached_existing })
479 })
480 }
481
482 fn run_action<'a>(
483 &'a self,
484 _bash_state: &'a Arc<tokio::sync::Mutex<Option<BashState>>>,
485 _command: crate::types::BashCommand,
486 ) -> ShellRuntimeFuture<'a> {
487 Box::pin(async { Err(WinxError::CommandExecutionError("stub".to_string())) })
488 }
489
490 fn interrupt<'a>(
491 &'a self,
492 _bash_state: &'a Arc<tokio::sync::Mutex<Option<BashState>>>,
493 ) -> ShellRuntimeUnitFuture<'a> {
494 Box::pin(async { Ok(()) })
495 }
496
497 fn terminate_session<'a>(&'a self, _thread_id: &'a str) -> ShellRuntimeUnitFuture<'a> {
498 Box::pin(async { Ok(()) })
499 }
500 }
501
502 fn service(attached_existing: bool) -> WinxService {
503 WinxService::with_runtime(
504 SessionIsolation::Strict,
505 Arc::new(StubRuntime { attached_existing }),
506 )
507 }
508
509 fn unique_thread_id(tag: &str) -> String {
510 let nanos = std::time::SystemTime::now()
511 .duration_since(std::time::UNIX_EPOCH)
512 .expect("system time after epoch")
513 .as_nanos();
514 format!("winxtest_rehydrate_{tag}_{}_{nanos}", std::process::id())
517 }
518
519 fn persist_snapshot(thread_id: &str, workspace: &std::path::Path) {
520 let mut state = BashState::new();
521 state.current_thread_id = thread_id.to_string();
522 state.workspace_root = workspace.to_path_buf();
523 state.cwd = workspace.to_path_buf();
524 state.initialized = true;
525 save_bash_state(thread_id, &state.snapshot()).expect("persist snapshot");
526 }
527
528 #[tokio::test]
529 async fn tool_call_rehydrates_a_persisted_session_transparently() {
530 let workspace = tempfile::tempdir().expect("workspace");
531 let thread_id = unique_thread_id("attach");
532 persist_snapshot(&thread_id, workspace.path());
533
534 let (slot, _guard) = service(true).tool_session_for(&thread_id).await;
536 {
537 let state = slot.lock().await;
538 let state = state.as_ref().expect("session should rehydrate from disk");
539 assert!(state.initialized);
540 assert_eq!(state.current_thread_id, thread_id);
541 assert_eq!(state.workspace_root, workspace.path());
542 assert!(state.recovery_note.is_none());
544 }
545 delete_bash_state(&thread_id).ok();
546 }
547
548 #[tokio::test]
549 async fn lost_shell_rehydration_surfaces_the_agent_resume_hint() {
550 let workspace = tempfile::tempdir().expect("workspace");
551 let thread_id = unique_thread_id("resume");
552 persist_snapshot(&thread_id, workspace.path());
553 save_agent_session(
554 &thread_id,
555 &AgentSessionRecord {
556 agent: "claude".to_string(),
557 command: "claude".to_string(),
558 cwd: workspace.path().to_string_lossy().into_owned(),
559 launched_at_unix_ms: 1,
560 workspace_root: None,
561 scoped_thread_id: None,
562 execution_token: None,
563 conversation_selector: None,
564 },
565 )
566 .expect("persist agent record");
567
568 let (slot, _guard) = service(false).tool_session_for(&thread_id).await;
571 {
572 let state = slot.lock().await;
573 let note = state
574 .as_ref()
575 .expect("session should rehydrate")
576 .recovery_note
577 .clone()
578 .expect("lost shell must set a recovery note");
579 assert!(note.contains("best_effort"), "missing recovery classification: {note}");
580 assert!(!note.contains("claude --continue"), "unbound hint selected latest: {note}");
581 assert!(!note.contains("resume its conversation"), "unverified identity claim: {note}");
582 }
583 delete_bash_state(&thread_id).ok();
584 crate::state::agent_resume::clear_agent_session(&thread_id).ok();
585 }
586
587 #[tokio::test]
588 async fn without_a_persisted_snapshot_the_slot_stays_empty() {
589 let thread_id = unique_thread_id("missing");
590 let (slot, _guard) = service(true).tool_session_for(&thread_id).await;
591 assert!(slot.lock().await.is_none(), "no snapshot on disk must not fabricate state");
592 }
593
594 #[tokio::test]
595 async fn empty_thread_id_never_rehydrates() {
596 let (slot, _guard) = service(true).tool_session_for("").await;
597 assert!(slot.lock().await.is_none());
598 }
599
600 fn status_probe(thread_id: &str, workspace: &std::path::Path) -> serde_json::Value {
601 serde_json::json!({
602 "type": "status",
603 "any_workspace_path": workspace,
604 "mode_name": "wcgw",
605 "thread_id": thread_id
606 })
607 }
608
609 fn tool_request(arguments: &serde_json::Value) -> rmcp::model::CallToolRequestParams {
610 rmcp::model::CallToolRequestParams::new("Initialize")
611 .with_arguments(arguments.as_object().expect("request object").clone())
612 }
613
614 #[tokio::test]
618 async fn status_probe_rehydrates_a_persisted_session_after_restart() {
619 let workspace = tempfile::tempdir().expect("workspace");
620 let thread_id = unique_thread_id("status");
621 persist_snapshot(&thread_id, workspace.path());
622 let arguments = status_probe(&thread_id, workspace.path());
623
624 let service = service(false);
625 let coherence = service
626 .validate_workspace_coherence(
627 &tool_request(&arguments),
628 &crate::server::principal::RequestScope::default(),
629 crate::config::HttpSessionAffinity::Thread,
630 None,
631 )
632 .await
633 .expect("known pair must validate after restart");
634 assert_eq!(coherence, crate::server::coherence::WorkspaceCoherence::Validated);
635
636 let result = service.handle_initialize(Some(arguments)).await.expect("status result");
637 assert_ne!(result.is_error, Some(true), "{result:?}");
638 let structured = result.structured_content.clone().expect("structured status");
639 assert_eq!(structured["initialize_transition"], "status", "{structured}");
640 assert_eq!(structured["initialize_recovered_missing_session"], true, "{structured}");
641 let text = &result.content[0].as_text().expect("status text").text;
642 assert!(text.contains("[session recovery]"), "lost shell must be reported: {text}");
643 let state = service.bound_workspace(&thread_id).await.expect("rehydrated binding");
644 assert_eq!(state, workspace.path());
645 delete_bash_state(&thread_id).ok();
646 }
647
648 #[tokio::test]
649 async fn non_strict_status_probe_rehydrates_a_persisted_session() {
650 let workspace = tempfile::tempdir().expect("workspace");
651 let thread_id = unique_thread_id("status_local");
652 persist_snapshot(&thread_id, workspace.path());
653 let service = WinxService::with_runtime(
654 SessionIsolation::Lenient,
655 Arc::new(StubRuntime { attached_existing: true }),
656 );
657
658 let result = service
659 .handle_initialize(Some(status_probe(&thread_id, workspace.path())))
660 .await
661 .expect("status result");
662 assert_ne!(result.is_error, Some(true), "{result:?}");
663 let structured = result.structured_content.clone().expect("structured status");
664 assert_eq!(structured["initialize_transition"], "status", "{structured}");
665 assert_eq!(structured["initialize_recovered_missing_session"], false, "{structured}");
667 delete_bash_state(&thread_id).ok();
668 }
669
670 #[tokio::test]
671 async fn status_probe_keeps_cross_workspace_pairs_fail_closed_after_restart() {
672 let bound = tempfile::tempdir().expect("bound workspace");
673 let other = tempfile::tempdir().expect("other workspace");
674 let thread_id = unique_thread_id("status_mismatch");
675 persist_snapshot(&thread_id, bound.path());
676 let arguments = status_probe(&thread_id, other.path());
677 let scope = crate::server::principal::RequestScope::default();
678
679 let service = service(true);
680 let error = service
681 .validate_workspace_coherence(
682 &tool_request(&arguments),
683 &scope,
684 crate::config::HttpSessionAffinity::Thread,
685 None,
686 )
687 .await
688 .expect_err("persisted pair belongs to another workspace");
689 assert!(matches!(error, WinxError::WorkspaceBindingMismatch { .. }), "{error:?}");
690
691 let fresh = WinxService::with_runtime(
694 SessionIsolation::Strict,
695 Arc::new(StubRuntime { attached_existing: true }),
696 );
697 let error = fresh
698 .validate_workspace_coherence(
699 &tool_request(&arguments),
700 &scope,
701 crate::config::HttpSessionAffinity::Workspace,
702 None,
703 )
704 .await
705 .expect_err("invented workspace id must fail");
706 assert!(matches!(error, WinxError::BashStateNotInitialized), "{error:?}");
707 assert!(fresh.bound_workspace(&thread_id).await.is_none(), "must not rehydrate");
708 delete_bash_state(&thread_id).ok();
709 }
710}