1use rand::RngExt;
5use rmcp::{
6 model::{
7 CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, ContentBlock,
8 CreateTaskResult, GetPromptRequestParams, GetPromptResponse, GetPromptResult,
9 GetTaskParams, GetTaskResult, Icon, Implementation, ListPromptsResult, ListResourcesResult,
10 ListToolsResult, PaginatedRequestParams, Prompt, PromptMessage, ProtocolVersion,
11 ReadResourceRequestParams, ReadResourceResponse, ReadResourceResult, Resource,
12 ResourceContents, Role, ServerCapabilities, ServerInfo, Task, TaskStatus, Tool,
13 ToolAnnotations, UpdateTaskParams,
14 },
15 service::{NotificationContext, RequestContext, RoleServer},
16 transport::{async_rw::AsyncRwTransport, stdio},
17 ErrorData as McpError, ServerHandler, ServiceExt,
18};
19use schemars::schema_for;
20use serde_json::Value;
21use std::collections::HashMap;
22use std::fmt::Write as FmtWrite;
23use std::path::Path;
24use std::process::Command;
25use std::sync::atomic::{AtomicUsize, Ordering};
26use std::sync::{Arc, OnceLock};
27use std::time::{Duration, Instant};
28use tokio::sync::Mutex;
29use tracing::{info, warn};
30
31use crate::errors::WinxError;
32use crate::runtime::{configured_shell_runtime, EmbeddedShellRuntime, ShellRuntime};
33use crate::state::bash_state::generate_thread_id;
34use crate::state::task_state::{TaskEntry, TaskRegistry, DEFAULT_TASK_TTL_MS};
35use crate::state::BashState;
36use crate::types::{
37 normalize_thread_id, BashCommand, BashCommandAction, CodeMap, CodeMapStructuredOutput,
38 ContextSave, FileWriteOrEdit, Initialize, InitializeType, ModeName, MultiFileEdit, ReadFiles,
39 ReadImage, UndoEdit,
40};
41
42fn to_mcp_error(tool: &str, err: &WinxError) -> McpError {
49 let msg = format!("{tool} failed: {err}");
50 match err {
54 WinxError::BashStateNotInitialized
56 | WinxError::CommandNotAllowed(_)
57 | WinxError::PathSecurityError { .. }
58 | WinxError::ThreadIdMismatch(_)
59 | WinxError::ParameterValidationError { .. }
60 | WinxError::MissingParameterError { .. }
61 | WinxError::NullValueError { .. }
62 | WinxError::ArgumentParseError(_)
63 | WinxError::JsonParseError(_)
64 | WinxError::DeserializationError(_)
65 | WinxError::WorkspacePathError(_)
66 | WinxError::InvalidInput(_)
67 | WinxError::ParseError(_)
68 | WinxError::FileAccessError { .. }
69 | WinxError::RecoverableSuggestionError { .. }
70 | WinxError::SearchReplaceSyntaxError(_)
71 | WinxError::SearchReplaceSyntaxErrorDetailed { .. }
72 | WinxError::SearchBlockNotFound(_)
73 | WinxError::SearchBlockAmbiguous { .. }
74 | WinxError::FileTooLarge { .. }
75 | WinxError::InteractiveCommandDetected { .. }
76 | WinxError::CommandAlreadyRunning { .. } => McpError::invalid_request(msg, None),
77 WinxError::ShellInitializationError(_)
79 | WinxError::BashStateLockError(_)
80 | WinxError::CommandExecutionError(_)
81 | WinxError::SerializationError(_)
82 | WinxError::FileWriteError { .. }
83 | WinxError::DataLoadingError(_)
84 | WinxError::ContextSaveError(_)
85 | WinxError::CommandTimeout { .. }
86 | WinxError::ProcessCleanupError { .. }
87 | WinxError::BufferOverflow { .. }
88 | WinxError::SessionRecoveryError { .. }
89 | WinxError::ResourceAllocationError { .. }
90 | WinxError::IoError(_)
91 | WinxError::ConfigurationError(_)
92 | WinxError::FileError(_) => McpError::internal_error(msg, None),
93 }
94}
95
96pub type SharedBashState = Arc<Mutex<Option<BashState>>>;
98
99fn schema_to_input_schema<T: schemars::JsonSchema>() -> Arc<serde_json::Map<String, Value>> {
101 let schema = schema_for!(T);
102 let mut value = serde_json::to_value(schema).unwrap_or(Value::Object(serde_json::Map::new()));
103 strip_schema_titles(&mut value);
108 strip_unsupported_schema_formats(&mut value);
109 match value {
110 Value::Object(map) => Arc::new(map),
111 _ => Arc::new(serde_json::Map::new()),
112 }
113}
114
115fn strip_unsupported_schema_formats(value: &mut Value) {
118 match value {
119 Value::Object(map) => {
120 if map.get("format").and_then(Value::as_str) == Some("uint") {
121 map.remove("format");
122 }
123 for child in map.values_mut() {
124 strip_unsupported_schema_formats(child);
125 }
126 }
127 Value::Array(items) => {
128 for item in items {
129 strip_unsupported_schema_formats(item);
130 }
131 }
132 _ => {}
133 }
134}
135
136fn strip_schema_titles(value: &mut Value) {
143 match value {
144 Value::Object(map) => {
145 const SCHEMA_KEYS: &[&str] = &[
146 "type",
147 "$ref",
148 "properties",
149 "items",
150 "additionalProperties",
151 "enum",
152 "const",
153 "anyOf",
154 "allOf",
155 "oneOf",
156 ];
157 if SCHEMA_KEYS.iter().any(|k| map.contains_key(*k)) {
158 map.remove("title");
159 }
160 for child in map.values_mut() {
161 strip_schema_titles(child);
162 }
163 }
164 Value::Array(items) => {
165 for item in items {
166 strip_schema_titles(item);
167 }
168 }
169 _ => {}
170 }
171}
172
173fn mcp_tool<T: schemars::JsonSchema>(
174 name: &'static str,
175 description: &'static str,
176 annotations: ToolAnnotations,
177) -> Tool {
178 Tool::new(name, description, schema_to_input_schema::<T>()).with_annotations(annotations)
179}
180
181fn with_output_schema<T: schemars::JsonSchema>(mut tool: Tool) -> Tool {
182 tool.output_schema = Some(schema_to_input_schema::<T>());
183 tool
184}
185
186const INITIALIZE_DESCRIPTION: &str =
187 "- Call this at the start of the conversation before using shell tools, unless a local MCP client supplied Roots and Winx initialized that workspace automatically. \
188 - Use `any_workspace_path` to initialize the shell in the appropriate project directory. \
189 - If the user has mentioned a workspace or project root or any other file or folder use it to set `any_workspace_path`. \
190 - If user has mentioned any files use `initial_files_to_read` to read, use absolute paths only (~ allowed) \
191 - By default use mode \"wcgw\" \
192 - In \"code-writer\" mode, set the commands and globs which user asked to set, otherwise use 'all'. \
193 - Use type=\"first_call\" if it's the first call to this tool. \
194 - Use type=\"user_asked_mode_change\" if in a conversation user has asked to change mode. \
195 - Use type=\"reset_shell\" if in a conversation shell is not working after multiple tries. \
196 - Use type=\"user_asked_change_workspace\" if in a conversation user asked to change workspace";
197
198const BASH_COMMAND_DESCRIPTION: &str =
199 "- Execute a bash command. This is stateful (beware with subsequent calls). \
200 - Accepted payloads include action_json with an explicit type, action_json shorthand such as {\"command\":\"pwd\"}, or top-level shorthand such as {\"command\":\"pwd\"}. \
201 - Status of the command and the current working directory will always be returned at the end. \
202 - The first or the last line might be `(...truncated)` if the output is too long. \
203 - Always run `pwd` if you get any file or directory not found error to make sure you're not lost. \
204 - Do not run bg commands using \"&\", instead use this tool. \
205 - You must not use echo/cat to read/write files, use ReadFiles/FileWriteOrEdit \
206 - In order to check status of previous command, use `status_check` with empty command argument. \
207 - Only command is allowed to run at a time. You need to wait for any previous command to finish before running a new one. \
208 - Programs don't hang easily, so most likely explanation for no output is usually that the program is still running, and you need to check status again. \
209 - Do not send Ctrl-c before checking for status till 10 minutes or whatever is appropriate for the program to finish. \
210 - Only run long running commands in background. Each background command is run in a new non-reusable shell. \
211 - On running a bg command you'll get a bg command id that you should use to get status or interact. \
212 - MCP Tasks are supported for a single foreground `command` when the client declares the `io.modelcontextprotocol/tasks` extension. Use an explicit `thread_id`; poll with tasks/get, whose completed state contains the final CallToolResult. Do not combine an MCP Task with `is_background=true`. \
213 - Piloting an interactive full-screen TUI (the `claude` CLI, vim, htop, fzf, a REPL)? Run it in the background, then drive it with these two actions: \
214 - `screen` ({\"screen\":true,\"bg_command_id\":\"...\",\"lines\":N,\"diff\":true}) returns a STABLE snapshot of the live terminal screen (cursor moves, redraws, alternate-screen and synchronized-output already applied; ANSI stripped), with the cursor position in the header. Use this to read the current frame — unlike `status_check`, it never stacks redraw generations and never waits. Pass \"diff\":true to get back ONLY the lines that changed since your last `screen` look (large token savings when polling a TUI frame-by-frame; first look or a big change still returns the full frame). \
215 - `wait_for_turn` ({\"wait_for_turn\":true,\"bg_command_id\":\"...\",\"recognizer\":\"auto|claude|codex|antigravity|generic\",\"quiet_ms\":600,\"timeout_seconds\":30}) waits for the TUI's turn and returns the stable snapshot plus the detected state (busy / awaiting_input / awaiting_approval). By default it returns as soon as it confirms the app is actively working (state=busy) so a long-running child never pins you for the whole timeout — poll again to keep watching; pass \"wait_through_busy\":true to instead block through busy until it is ready for input (or the timeout fires). Typical REPL loop: run the app in bg -> wait_for_turn until awaiting_input -> send_text(submit:true) -> wait_for_turn -> screen, repeat.";
216
217const READ_FILES_DESCRIPTION: &str =
218 "- Read full file content of one or more files. \
219 - Prefer this over reading files with BashCommand (cat/head/tail): the output is token-budgeted and the read is recorded so FileWriteOrEdit can edit the file afterward. \
220 - Do NOT use this for binary files or images — use ReadImage for images. \
221 - Provide absolute paths only (~ allowed) \
222 - Only if the task requires line numbers understanding: \
223 - You may extract a range of lines. E.g., `/path/to/file:1-10` for lines 1-10. You can drop start or end like `/path/to/file:1-` or `/path/to/file:-10`";
224
225const FILE_WRITE_OR_EDIT_DESCRIPTION: &str =
226 "- Writes or edits a file based on the percentage of changes. \
227 - Prefer this over writing/editing files with BashCommand (echo/sed/redirects/heredocs). \
228 - For an edit, the file must have been read first with ReadFiles, otherwise the edit is rejected. \
229 - Use absolute path only (~ allowed). \
230 - First write down percentage of lines that need to be replaced in the file (between 0-100) in percentage_to_change \
231 - percentage_to_change should be low if mostly new code is to be added. It should be high if a lot of things are to be replaced. \
232 - If percentage_to_change > 50, provide full file content in text_or_search_replace_blocks \
233 - If percentage_to_change <= 50, text_or_search_replace_blocks should be search/replace blocks. \
234 \
235 Instructions for editing files. \
236 # Example \
237 ## Input file \
238 ``` \
239 import numpy as np \
240 from impls import impl1, impl2 \
241 \
242 def hello(): \
243 \"print a greeting\" \
244 print(\"hello\") \
245 \
246 def call_hello(): \
247 \"call hello\" \
248 hello() \
249 print(\"Called\") \
250 impl1() \
251 hello() \
252 impl2() \
253 ``` \
254 ## Edit format on the input file \
255 ``` \
256 <<<<<<< SEARCH \
257 from impls import impl1, impl2 \
258 ======= \
259 from impls import impl1, impl2 \
260 from hello import hello as hello_renamed \
261 >>>>>>> REPLACE \
262 <<<<<<< SEARCH \
263 def hello(): \
264 \"print a greeting\" \
265 print(\"hello\") \
266 ======= \
267 >>>>>>> REPLACE \
268 <<<<<<< SEARCH \
269 def call_hello(): \
270 \"call hello\" \
271 hello() \
272 ======= \
273 def call_hello_renamed(): \
274 \"call hello renamed\" \
275 hello_renamed() \
276 >>>>>>> REPLACE \
277 <<<<<<< SEARCH \
278 impl1() \
279 hello() \
280 impl2() \
281 ======= \
282 impl1() \
283 hello_renamed() \
284 impl2() \
285 >>>>>>> REPLACE \
286 ``` \
287 # *SEARCH/REPLACE block* Rules: \
288 Every \"<<<<<<< SEARCH\" section must *EXACTLY MATCH* the existing file content, character for character, including all comments, docstrings, whitespaces, etc. \
289 Including multiple unique *SEARCH/REPLACE* blocks if needed. \
290 Include enough and only enough lines in each SEARCH section to uniquely match each set of lines that need to change. \
291 Keep *SEARCH/REPLACE* blocks concise. \
292 Break large *SEARCH/REPLACE* blocks into a series of smaller blocks that each change a small portion of the file. \
293 Include just the changing lines, and a few surrounding lines (0-3 lines) if needed for uniqueness. \
294 Other than for uniqueness, avoid including those lines which do not change in search (and replace) blocks. Target 0-3 non trivial extra lines per block. \
295 Preserve leading spaces and indentations in both SEARCH and REPLACE blocks. \
296 If a short block would match multiple places, anchor it to a line number from ReadFiles instead of padding with context: write the marker as \"<<<<<<< SEARCH @42\" (or a range \"@42-50\") to target that 1-based location. A stale anchor falls back to the normal search, so it never makes a valid edit fail.";
297
298const MULTI_FILE_EDIT_DESCRIPTION: &str =
299 "- Edits SEVERAL files together, all-or-nothing. Use this over multiple FileWriteOrEdit calls when a change spans files and a partial apply would be bad (e.g. rename a symbol across files). \
300 - Every file's edit is validated and computed in memory FIRST; only if ALL succeed is anything written, so a SEARCH that fails to match in the last file leaves the earlier files untouched. \
301 - Each entry has the same fields as FileWriteOrEdit: file_path (absolute, ~ allowed), percentage_to_change, and text_or_search_replace_blocks. Each file must have been read with ReadFiles first. \
302 - Provide 2+ files; for a single file use FileWriteOrEdit. Do not list the same file twice. \
303 - If a write fails mid-batch (rare: disk/permissions), it stops and reports which files were already written; those are not rolled back.";
304
305const UNDO_EDIT_DESCRIPTION: &str =
306 "- Reverts a file to the content it had before the last FileWriteOrEdit/MultiFileEdit you made to it THIS session. \
307 - Use this to back out a wrong edit instead of re-typing the old content. \
308 - Per-file: call it again on the same file to walk further back through its edits (the last ~10 edits per session are kept, in memory only). \
309 - Refused if the file changed on disk since your edit (so an undo never discards newer changes), and a brand-new file's creation cannot be undone (no prior content) - use BashCommand rm for that. \
310 - Provide file_path (absolute, ~ allowed).";
311
312const CONTEXT_SAVE_DESCRIPTION: &str =
313 "Saves provided description and file contents of all the relevant file paths or globs in a single text file. \
314 - Provide random 3 word unqiue id or whatever user provided. \
315 - Leave project path as empty string if no project path";
316
317const CODE_MAP_DESCRIPTION: &str =
318 "- Navigate code structure via tree-sitter - the semantic layer plain grep/rg can't give you. Pick an `operation`: \
319 - operation=\"outline\": map symbols (functions, types, methods, classes, ...). `path` to a FILE returns that file's definitions; `path` to a DIRECTORY (or empty = the whole workspace) returns a relevance-ranked, token-budgeted symbol map across files. Use it instead of reading whole files just to learn their shape. \
320 - operation=\"references\": find where a symbol is defined and referenced (called/used), by name. `name` is required (exact identifier). Counts only real symbol occurrences, never matches inside strings or comments. Output lists definitions first, then references, as `def|ref file:line kind name`. \
321 - Scope either operation with `path` (file or directory; empty = whole workspace); cap with `max_results`. gitignore-aware, workspace-confined, works in every mode. \
322 - 11 languages (rust, js/ts, go, c, c++, java, ruby, c#, php, lua); other files return no symbols. Note: C/C++ grammars tag definitions only, so references reads 0 for `.c`/`.h`/`.cpp`. \
323 - For plain-text/regex search or file discovery, use rg / grep / fd / find via BashCommand.";
324
325static WINX_TOOLS: OnceLock<Vec<Tool>> = OnceLock::new();
326static WINX_PROMPTS: OnceLock<Vec<Prompt>> = OnceLock::new();
327
328fn winx_tools() -> Vec<Tool> {
329 WINX_TOOLS.get_or_init(build_winx_tools).clone()
330}
331
332fn build_winx_tools() -> Vec<Tool> {
333 vec![
334 mcp_tool::<Initialize>(
335 "Initialize",
336 INITIALIZE_DESCRIPTION,
337 ToolAnnotations::new().read_only(true).open_world(false),
338 ),
339 mcp_tool::<BashCommand>(
340 "BashCommand",
341 BASH_COMMAND_DESCRIPTION,
342 ToolAnnotations::new().destructive(true).open_world(true),
343 ),
344 mcp_tool::<ReadFiles>(
345 "ReadFiles",
346 READ_FILES_DESCRIPTION,
347 ToolAnnotations::new().read_only(true).open_world(false),
348 ),
349 mcp_tool::<FileWriteOrEdit>(
350 "FileWriteOrEdit",
351 FILE_WRITE_OR_EDIT_DESCRIPTION,
352 ToolAnnotations::new().destructive(true).open_world(false),
353 ),
354 mcp_tool::<MultiFileEdit>(
355 "MultiFileEdit",
356 MULTI_FILE_EDIT_DESCRIPTION,
357 ToolAnnotations::new().destructive(true).open_world(false),
358 ),
359 mcp_tool::<UndoEdit>(
360 "UndoEdit",
361 UNDO_EDIT_DESCRIPTION,
362 ToolAnnotations::new().destructive(true).open_world(false),
363 ),
364 mcp_tool::<ContextSave>(
365 "ContextSave",
366 CONTEXT_SAVE_DESCRIPTION,
367 ToolAnnotations::new().destructive(false).open_world(false),
368 ),
369 mcp_tool::<ReadImage>(
370 "ReadImage",
371 "Read an image from the shell.",
372 ToolAnnotations::new().read_only(true).open_world(false),
373 ),
374 with_output_schema::<CodeMapStructuredOutput>(mcp_tool::<CodeMap>(
375 "CodeMap",
376 CODE_MAP_DESCRIPTION,
377 ToolAnnotations::new().read_only(true).open_world(false),
378 )),
379 ]
380}
381
382fn winx_prompts() -> Vec<Prompt> {
383 WINX_PROMPTS
384 .get_or_init(|| {
385 vec![Prompt::new(
386 "KnowledgeTransfer",
387 Some("Summarize current Winx state, workspace context, and handoff notes."),
388 None,
389 )]
390 })
391 .clone()
392}
393
394fn append_command_section<const N: usize>(
395 output: &mut String,
396 title: &str,
397 cwd: &Path,
398 args: [&str; N],
399) {
400 let Ok(command_output) = Command::new("git").args(["-C"]).arg(cwd).args(args).output() else {
401 return;
402 };
403 if !command_output.status.success() {
404 return;
405 }
406
407 let content = String::from_utf8_lossy(&command_output.stdout);
408 if content.trim().is_empty() {
409 return;
410 }
411
412 let _ = writeln!(output, "\n# {title}\n{}", content.trim_end());
413}
414
415const MAX_SESSIONS: usize = 32;
419
420#[derive(Default)]
425struct SessionRegistry {
426 slots: HashMap<String, SharedBashState>,
427 last_used: HashMap<String, Instant>,
429 in_flight: HashMap<String, SessionPin>,
433 last_active: Option<String>,
436}
437
438#[derive(Clone)]
445struct SessionPin {
446 count: Arc<AtomicUsize>,
447}
448
449impl Default for SessionPin {
450 fn default() -> Self {
451 Self { count: Arc::new(AtomicUsize::new(0)) }
452 }
453}
454
455impl SessionPin {
456 fn acquire(&self) -> SessionGuard {
458 self.count.fetch_add(1, Ordering::SeqCst);
459 SessionGuard { count: self.count.clone() }
460 }
461
462 fn is_pinned(&self) -> bool {
464 self.count.load(Ordering::SeqCst) > 0
465 }
466}
467
468struct SessionGuard {
474 count: Arc<AtomicUsize>,
475}
476
477impl Drop for SessionGuard {
478 fn drop(&mut self) {
479 self.count.fetch_sub(1, Ordering::SeqCst);
480 }
481}
482
483#[derive(Clone, Copy, PartialEq, Eq, Debug)]
485pub enum SessionIsolation {
486 Lenient,
489 Strict,
499}
500
501#[derive(Clone)]
503pub struct WinxService {
504 sessions: Arc<Mutex<SessionRegistry>>,
505 tasks: Arc<Mutex<TaskRegistry>>,
506 root_bootstrap: Arc<Mutex<()>>,
507 shell_runtime: Arc<dyn ShellRuntime>,
508 pub version: String,
510 isolation: SessionIsolation,
512}
513
514impl Default for WinxService {
515 fn default() -> Self {
516 Self::new()
517 }
518}
519
520impl WinxService {
521 pub fn new() -> Self {
524 Self::with_isolation(SessionIsolation::Lenient)
525 }
526
527 pub fn with_isolation(isolation: SessionIsolation) -> Self {
530 Self::with_runtime(isolation, Arc::new(EmbeddedShellRuntime))
531 }
532
533 pub fn with_runtime(isolation: SessionIsolation, shell_runtime: Arc<dyn ShellRuntime>) -> Self {
535 info!(?isolation, "Creating new WinxService instance");
536 Self {
537 sessions: Arc::new(Mutex::new(SessionRegistry::default())),
538 tasks: Arc::new(Mutex::new(TaskRegistry::default())),
539 root_bootstrap: Arc::new(Mutex::new(())),
540 shell_runtime,
541 version: env!("CARGO_PKG_VERSION").to_string(),
542 isolation,
543 }
544 }
545
546 async fn session_for(&self, thread_id: &str) -> (SharedBashState, SessionGuard) {
555 let mut reg = self.sessions.lock().await;
556 let key = if thread_id.is_empty() {
557 match self.isolation {
558 SessionIsolation::Lenient => {
560 reg.last_active.clone().unwrap_or_else(|| "default".to_string())
561 }
562 SessionIsolation::Strict => "anonymous".to_string(),
565 }
566 } else {
567 thread_id.to_string()
568 };
569
570 if !reg.slots.contains_key(&key) && reg.slots.len() >= MAX_SESSIONS {
575 let victim = reg
576 .last_used
577 .iter()
578 .filter(|(k, _)| **k != key)
579 .filter(|(k, _)| reg.in_flight.get(k.as_str()).map_or(true, |p| !p.is_pinned()))
580 .min_by_key(|(_, t)| **t)
581 .map(|(k, _)| k.clone());
582 if let Some(victim) = victim {
583 reg.slots.remove(&victim);
584 reg.last_used.remove(&victim);
585 reg.in_flight.remove(&victim);
586 if reg.last_active.as_deref() == Some(victim.as_str()) {
587 reg.last_active = None;
588 }
589 warn!("Evicted LRU shell session '{victim}' (session cap {MAX_SESSIONS})");
590 } else {
591 warn!(
592 "All {MAX_SESSIONS} sessions busy; exceeding the cap rather than evicting an in-flight session"
593 );
594 }
595 }
596
597 let slot =
598 reg.slots.entry(key.clone()).or_insert_with(|| Arc::new(Mutex::new(None))).clone();
599 let pin = reg.in_flight.entry(key.clone()).or_default().clone();
600 reg.last_used.insert(key.clone(), Instant::now());
601 if !thread_id.is_empty() {
602 reg.last_active = Some(key);
603 }
604 (slot, pin.acquire())
605 }
606
607 async fn active_slot(&self) -> Option<SharedBashState> {
610 let reg = self.sessions.lock().await;
611 reg.last_active.as_ref().and_then(|key| reg.slots.get(key).cloned())
612 }
613
614 async fn run_bash_task(
618 &self,
619 mut request: CallToolRequestParams,
620 ) -> Result<CallToolResult, McpError> {
621 let args = request
622 .arguments
623 .clone()
624 .ok_or_else(|| McpError::invalid_request("Missing BashCommand arguments", None))?;
625 let mut bash: BashCommand =
626 serde_json::from_value(Value::Object(args)).map_err(|error| {
627 McpError::invalid_request(
628 format!("Invalid BashCommand task parameters: {error}"),
629 None,
630 )
631 })?;
632 let BashCommandAction::Command { is_background, .. } = &bash.action_json else {
633 return Err(McpError::invalid_request(
634 "MCP Tasks support BashCommand's foreground command action only",
635 None,
636 ));
637 };
638 if *is_background {
639 return Err(McpError::invalid_request(
640 "Do not combine an MCP Task with BashCommand is_background=true",
641 None,
642 ));
643 }
644
645 let thread_id = normalize_thread_id(&bash.thread_id);
646 if thread_id.is_empty() {
647 return Err(McpError::invalid_request(
648 "Task-augmented BashCommand requires an explicit thread_id",
649 None,
650 ));
651 }
652 bash.thread_id.clone_from(&thread_id);
653 bash.wait_for_seconds = Some(20.0);
654 request.arguments = Some(task_arguments(&bash)?);
655
656 let mut output = String::new();
657 loop {
658 let result = self.execute_tool_call(request).await?;
659 let chunk = result
660 .content
661 .iter()
662 .filter_map(|content| content.as_text().map(|text| text.text.as_str()))
663 .collect::<Vec<_>>()
664 .join("\n");
665 append_task_output(&mut output, &chunk);
666 if !chunk.contains("status = still running") {
667 return Ok(CallToolResult::success(vec![ContentBlock::text(output)]));
668 }
669
670 let status = BashCommand {
671 action_json: BashCommandAction::StatusCheck {
672 status_check: true,
673 bg_command_id: None,
674 scrollback_lines: None,
675 verbose: false,
676 },
677 wait_for_seconds: Some(20.0),
678 thread_id: thread_id.clone(),
679 };
680 request =
681 CallToolRequestParams::new("BashCommand").with_arguments(task_arguments(&status)?);
682 }
683 }
684
685 async fn interrupt_task_thread(&self, thread_id: &str) {
686 let slot = {
687 let registry = self.sessions.lock().await;
688 registry.slots.get(thread_id).cloned()
689 };
690 let Some(slot) = slot else { return };
691 if let Err(error) = self.shell_runtime.interrupt(&slot).await {
692 warn!(thread_id, %error, "failed to interrupt cancelled MCP task");
693 }
694 }
695
696 async fn has_initialized_session(&self) -> bool {
697 let slots = {
698 let registry = self.sessions.lock().await;
699 registry.slots.values().cloned().collect::<Vec<_>>()
700 };
701 for slot in slots {
702 if slot.lock().await.as_ref().is_some_and(|state| state.initialized) {
703 return true;
704 }
705 }
706 false
707 }
708
709 #[allow(deprecated)]
713 async fn initialize_from_client_roots(&self, context: NotificationContext<RoleServer>) {
714 if self.isolation != SessionIsolation::Lenient {
715 return;
716 }
717 let _bootstrap_guard = self.root_bootstrap.lock().await;
718 if self.has_initialized_session().await {
719 return;
720 }
721 let supports_roots =
722 context.peer.peer_info().is_some_and(|info| info.capabilities.roots.is_some());
723 if !supports_roots {
724 return;
725 }
726 let roots = match context.peer.list_roots().await {
727 Ok(result) => result.roots,
728 Err(error) => {
729 warn!(%error, "client advertised Roots but roots/list failed");
730 return;
731 }
732 };
733 let cwd = std::env::current_dir().ok();
734 let mut paths =
735 roots.iter().filter_map(|root| root_uri_to_path(&root.uri)).collect::<Vec<_>>();
736 let selected = cwd
737 .as_ref()
738 .and_then(|cwd| paths.iter().position(|path| cwd.starts_with(path)))
739 .map(|index| paths.remove(index))
740 .or_else(|| paths.into_iter().next());
741 let Some(workspace) = selected else {
742 warn!("client returned no usable local file Roots; Initialize remains required");
743 return;
744 };
745
746 let initialize = Initialize {
747 init_type: InitializeType::FirstCall,
748 any_workspace_path: workspace.to_string_lossy().into_owned(),
749 initial_files_to_read: Vec::new(),
750 task_id_to_resume: String::new(),
751 mode_name: ModeName::Wcgw,
752 thread_id: generate_thread_id(),
753 code_writer_config: None,
754 };
755 let arguments = match serde_json::to_value(initialize) {
756 Ok(arguments) => arguments,
757 Err(error) => {
758 warn!(%error, "failed to serialize Roots bootstrap request");
759 return;
760 }
761 };
762 match self.handle_initialize(Some(arguments)).await {
763 Ok(_) => info!(workspace = %workspace.display(), "initialized Winx from MCP Roots"),
764 Err(error) => warn!(%error.message, "failed to initialize Winx from MCP Roots"),
765 }
766 }
767
768 fn bash_task_is_eligible(request: &CallToolRequestParams) -> bool {
769 if request.name != "BashCommand" {
770 return false;
771 }
772 let Some(arguments) = request.arguments.clone() else {
773 return false;
774 };
775 let Ok(bash) = serde_json::from_value::<BashCommand>(Value::Object(arguments)) else {
776 return false;
777 };
778 matches!(bash.action_json, BashCommandAction::Command { is_background: false, .. })
779 && !normalize_thread_id(&bash.thread_id).is_empty()
780 }
781
782 async fn enqueue_bash_task(
783 &self,
784 request: CallToolRequestParams,
785 ) -> Result<CreateTaskResult, McpError> {
786 let args = request
787 .arguments
788 .clone()
789 .ok_or_else(|| McpError::invalid_request("Missing BashCommand arguments", None))?;
790 let bash: BashCommand = serde_json::from_value(Value::Object(args)).map_err(|error| {
791 McpError::invalid_request(format!("Invalid BashCommand task parameters: {error}"), None)
792 })?;
793 let BashCommandAction::Command { is_background, .. } = bash.action_json else {
794 return Err(McpError::invalid_request(
795 "MCP Tasks support BashCommand's foreground command action only",
796 None,
797 ));
798 };
799 if is_background {
800 return Err(McpError::invalid_request(
801 "Do not combine an MCP Task with BashCommand is_background=true",
802 None,
803 ));
804 }
805 let thread_id = normalize_thread_id(&bash.thread_id);
806 if thread_id.is_empty() {
807 return Err(McpError::invalid_request(
808 "Task-augmented BashCommand requires an explicit thread_id",
809 None,
810 ));
811 }
812
813 let task_id = format!("task_{:032x}", rand::rng().random::<u128>());
814 let now = rmcp::task_manager::current_timestamp();
815 let task = Task::new(task_id.clone(), TaskStatus::Working, now.clone(), now)
816 .with_status_message("Running BashCommand")
817 .with_ttl_ms(DEFAULT_TASK_TTL_MS)
818 .with_poll_interval_ms(1_000);
819 self.tasks
820 .lock()
821 .await
822 .insert(task_id.clone(), TaskEntry::working(task.clone(), thread_id.clone()))
823 .map_err(|message| McpError::internal_error(message, None))?;
824
825 let service = self.clone();
826 let worker_task_id = task_id.clone();
827 let handle = tokio::spawn(async move {
828 let outcome =
829 tokio::time::timeout(MAX_TASK_RUNTIME, service.run_bash_task(request)).await;
830
831 let (status, message, result) = match outcome {
832 Ok(Ok(result)) => match serde_json::to_value(result) {
833 Ok(result) => (
834 TaskStatus::Completed,
835 Some("BashCommand completed".to_string()),
836 Some(result),
837 ),
838 Err(error) => (
839 TaskStatus::Failed,
840 Some(format!("Failed to serialize task result: {error}")),
841 None,
842 ),
843 },
844 Ok(Err(error)) => (
845 TaskStatus::Failed,
846 Some(crate::utils::redact::redact(&error.message).into_owned()),
847 None,
848 ),
849 Err(_) => {
850 service.interrupt_task_thread(&thread_id).await;
851 (
852 TaskStatus::Failed,
853 Some("BashCommand task exceeded the one-hour runtime limit".to_string()),
854 None,
855 )
856 }
857 };
858
859 let mut tasks = service.tasks.lock().await;
860 if let Some(entry) = tasks.get_mut(&worker_task_id) {
861 if entry.task.status == TaskStatus::Working {
862 entry.finish(status, message, result);
863 }
864 }
865 });
866 let abort_handle = handle.abort_handle();
867 drop(handle);
868 if let Some(entry) = self.tasks.lock().await.get_mut(&task_id) {
869 if entry.task.status == TaskStatus::Working {
870 entry.abort_handle = Some(abort_handle);
871 }
872 }
873
874 Ok(CreateTaskResult::new(task))
875 }
876}
877
878const MAX_TASK_RUNTIME: Duration = Duration::from_secs(60 * 60);
879const MAX_TASK_OUTPUT_BYTES: usize = 1_000_000;
880
881fn task_arguments<T: serde::Serialize>(
882 value: &T,
883) -> Result<serde_json::Map<String, Value>, McpError> {
884 match serde_json::to_value(value) {
885 Ok(Value::Object(arguments)) => Ok(arguments),
886 Ok(_) => Err(McpError::internal_error("Task arguments were not an object", None)),
887 Err(error) => Err(McpError::internal_error(
888 format!("Failed to serialize task arguments: {error}"),
889 None,
890 )),
891 }
892}
893
894fn append_task_output(output: &mut String, chunk: &str) {
895 if !output.is_empty() && !chunk.is_empty() {
896 output.push('\n');
897 }
898 output.push_str(chunk);
899 if output.len() <= MAX_TASK_OUTPUT_BYTES {
900 return;
901 }
902 let cut = crate::utils::floor_char_boundary(output, output.len() - MAX_TASK_OUTPUT_BYTES);
903 output.replace_range(..cut, "(...earlier task output truncated...)\n");
904}
905
906fn root_uri_to_path(uri: &str) -> Option<std::path::PathBuf> {
907 let encoded = uri.strip_prefix("file://")?;
908 let encoded = if let Some(path) = encoded.strip_prefix("localhost/") {
909 format!("/{path}")
910 } else if encoded.starts_with('/') {
911 encoded.to_string()
912 } else {
913 return None;
915 };
916 let decoded = percent_encoding::percent_decode_str(&encoded).decode_utf8().ok()?;
917 let path = std::path::PathBuf::from(decoded.as_ref());
918 path.is_absolute().then_some(path)
919}
920
921fn server_icon_data_uri() -> &'static str {
925 static URI: std::sync::OnceLock<String> = std::sync::OnceLock::new();
926 URI.get_or_init(|| {
927 use base64::Engine as _;
928 let png = include_bytes!("../.github/assets/icon-96.png");
929 format!("data:image/png;base64,{}", base64::engine::general_purpose::STANDARD.encode(png))
930 })
931}
932
933impl ServerHandler for WinxService {
935 fn get_info(&self) -> ServerInfo {
936 ServerInfo::new(
937 ServerCapabilities::builder()
938 .enable_tools()
939 .enable_resources()
940 .enable_prompts()
941 .enable_tasks()
942 .build(),
943 )
944 .with_server_info(
945 Implementation::new("winx-mcp-server", self.version.clone())
946 .with_title("Winx High-Performance MCP")
947 .with_icons(vec![Icon::new(server_icon_data_uri())
948 .with_mime_type("image/png")
949 .with_sizes(vec!["96x96".to_string()])]),
950 )
951 .with_protocol_version(ProtocolVersion::V_2026_07_28)
952 .with_instructions(
953 "Winx is a high-performance Rust implementation of MCP tools for shell and file management."
954 )
955 }
956
957 async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
958 self.initialize_from_client_roots(context).await;
959 }
960
961 async fn on_roots_list_changed(&self, context: NotificationContext<RoleServer>) {
962 self.initialize_from_client_roots(context).await;
965 }
966
967 fn get_tool(&self, name: &str) -> Option<Tool> {
968 winx_tools().into_iter().find(|tool| tool.name == name)
969 }
970
971 async fn get_task(
972 &self,
973 request: GetTaskParams,
974 _context: RequestContext<RoleServer>,
975 ) -> Result<GetTaskResult, McpError> {
976 let mut tasks = self.tasks.lock().await;
977 let entry = tasks.get(&request.task_id).ok_or_else(|| {
978 McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
979 })?;
980 Ok(GetTaskResult::new(entry.detailed()))
981 }
982
983 async fn update_task(
984 &self,
985 request: UpdateTaskParams,
986 _context: RequestContext<RoleServer>,
987 ) -> Result<(), McpError> {
988 let mut tasks = self.tasks.lock().await;
989 tasks.get(&request.task_id).ok_or_else(|| {
990 McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
991 })?;
992 Ok(())
993 }
994
995 async fn cancel_task(
996 &self,
997 request: CancelTaskParams,
998 _context: RequestContext<RoleServer>,
999 ) -> Result<(), McpError> {
1000 let (abort_handle, thread_id) = {
1001 let mut tasks = self.tasks.lock().await;
1002 let entry = tasks.get_mut(&request.task_id).ok_or_else(|| {
1003 McpError::invalid_request(
1004 format!("Unknown or expired task: {}", request.task_id),
1005 None,
1006 )
1007 })?;
1008 if entry.task.status != TaskStatus::Working
1009 && entry.task.status != TaskStatus::InputRequired
1010 {
1011 return Err(McpError::invalid_request(
1012 format!("Task {} is already terminal", request.task_id),
1013 None,
1014 ));
1015 }
1016 let abort_handle = entry.abort_handle.take();
1017 let thread_id = entry.thread_id.clone();
1018 entry.finish(TaskStatus::Cancelled, Some("Cancelled by client".to_string()), None);
1019 (abort_handle, thread_id)
1020 };
1021 if let Some(abort_handle) = abort_handle {
1022 abort_handle.abort();
1023 }
1024 self.interrupt_task_thread(&thread_id).await;
1025 Ok(())
1026 }
1027
1028 async fn list_tools(
1029 &self,
1030 _request: Option<PaginatedRequestParams>,
1031 _context: RequestContext<RoleServer>,
1032 ) -> Result<ListToolsResult, McpError> {
1033 Ok(ListToolsResult::with_all_items(winx_tools()))
1034 }
1035
1036 async fn list_resources(
1037 &self,
1038 _param: Option<PaginatedRequestParams>,
1039 _context: RequestContext<RoleServer>,
1040 ) -> Result<ListResourcesResult, McpError> {
1041 Ok(ListResourcesResult::with_all_items(vec![Resource::new("file://readme", "README")
1042 .with_description("Project README documentation")
1043 .with_mime_type("text/markdown")]))
1044 }
1045
1046 async fn list_prompts(
1047 &self,
1048 _request: Option<PaginatedRequestParams>,
1049 _context: RequestContext<RoleServer>,
1050 ) -> Result<ListPromptsResult, McpError> {
1051 Ok(ListPromptsResult::with_all_items(winx_prompts()))
1052 }
1053
1054 async fn get_prompt(
1055 &self,
1056 request: GetPromptRequestParams,
1057 _context: RequestContext<RoleServer>,
1058 ) -> Result<GetPromptResponse, McpError> {
1059 if request.name != "KnowledgeTransfer" {
1060 return Err(McpError::invalid_request(
1061 format!("Unknown prompt: {}", request.name),
1062 None,
1063 ));
1064 }
1065
1066 let text =
1069 crate::utils::redact::redact(&self.knowledge_transfer_prompt_text().await).into_owned();
1070
1071 Ok(GetPromptResult::new(vec![PromptMessage::new_text(Role::User, text)])
1072 .with_description("Knowledge transfer handoff prompt")
1073 .into())
1074 }
1075
1076 async fn read_resource(
1077 &self,
1078 param: ReadResourceRequestParams,
1079 _context: RequestContext<RoleServer>,
1080 ) -> Result<ReadResourceResponse, McpError> {
1081 let content = match param.uri.as_ref() {
1082 "file://readme" => match tokio::fs::read_to_string("README.md").await {
1083 Ok(content) => vec![ResourceContents::text(
1084 crate::utils::redact::redact(&content).into_owned(),
1085 param.uri.clone(),
1086 )],
1087 Err(_) => vec![ResourceContents::text(
1088 "README.md not found".to_string(),
1089 param.uri.clone(),
1090 )],
1091 },
1092 _ => {
1093 return Err(McpError::invalid_request(
1094 format!("Unknown resource URI: {}", param.uri),
1095 None,
1096 ));
1097 }
1098 };
1099
1100 Ok(ReadResourceResult::new(content).into())
1101 }
1102
1103 async fn call_tool(
1104 &self,
1105 param: CallToolRequestParams,
1106 context: RequestContext<RoleServer>,
1107 ) -> Result<CallToolResponse, McpError> {
1108 if context.client_capabilities().is_some_and(|capabilities| capabilities.supports_tasks())
1109 && Self::bash_task_is_eligible(¶m)
1110 {
1111 return Ok(CallToolResponse::Task(self.enqueue_bash_task(param).await?));
1112 }
1113 Ok(self.execute_tool_call(param).await?.into())
1114 }
1115}
1116
1117impl WinxService {
1118 async fn execute_tool_call(
1119 &self,
1120 param: CallToolRequestParams,
1121 ) -> Result<CallToolResult, McpError> {
1122 let tool = param.name.to_string();
1123 let args_value = param.arguments.map(Value::Object);
1124 let summary =
1130 crate::utils::redact::redact(&audit_summary(&tool, args_value.as_ref())).into_owned();
1131 let started = std::time::Instant::now();
1132
1133 let result = match tool.as_str() {
1134 "Initialize" => self.handle_initialize(args_value).await,
1135 "BashCommand" => self.handle_bash_command(args_value).await,
1136 "ReadFiles" => self.handle_read_files(args_value).await,
1137 "FileWriteOrEdit" => self.handle_file_write_or_edit(args_value).await,
1138 "MultiFileEdit" => self.handle_multi_file_edit(args_value).await,
1139 "UndoEdit" => self.handle_undo_edit(args_value).await,
1140 "ContextSave" => self.handle_context_save(args_value).await,
1141 "ReadImage" => self.handle_read_image(args_value).await,
1142 "CodeMap" => self.handle_code_map(args_value).await,
1143 _ => Err(McpError::invalid_request(format!("Unknown tool: {tool}"), None)),
1144 };
1145
1146 let result = match result {
1150 Ok(mut call) => {
1151 redact_result(&mut call);
1152 Ok(call)
1153 }
1154 Err(mut error) => {
1155 error.message = crate::utils::redact::redact(&error.message).into_owned().into();
1159 Err(error)
1160 }
1161 };
1162
1163 let ms = started.elapsed().as_millis();
1164 match &result {
1165 Ok(_) => info!(tool = %tool, ms, "tool call ok — {summary}"),
1166 Err(error) => warn!(tool = %tool, ms, "tool call error — {summary}: {}", error.message),
1168 }
1169 result
1170 }
1171}
1172
1173fn audit_summary(tool: &str, args: Option<&Value>) -> String {
1175 let Some(args) = args else {
1176 return "(no args)".to_string();
1177 };
1178 let s = |key: &str| args.get(key).and_then(Value::as_str).unwrap_or("").to_string();
1179 match tool {
1180 "BashCommand" => {
1181 let action = args.get("action_json");
1182 let cmd = action
1183 .and_then(|a| a.get("command"))
1184 .and_then(Value::as_str)
1185 .or_else(|| args.get("command").and_then(Value::as_str));
1186 if let Some(cmd) = cmd {
1187 format!("command bytes={}", cmd.len())
1188 } else {
1189 let kind =
1190 action.and_then(|a| a.get("type")).and_then(Value::as_str).unwrap_or("?");
1191 format!("action={kind}")
1192 }
1193 }
1194 "FileWriteOrEdit" | "ReadImage" | "UndoEdit" => format!("path={}", s("file_path")),
1195 "MultiFileEdit" => {
1196 format!("files={}", args.get("files").and_then(Value::as_array).map_or(0, Vec::len))
1197 }
1198 "ReadFiles" => {
1199 format!(
1200 "files={}",
1201 args.get("file_paths").and_then(Value::as_array).map_or(0, Vec::len)
1202 )
1203 }
1204 "Initialize" => format!("ws={} mode={}", s("any_workspace_path"), s("mode_name")),
1205 "ContextSave" => format!("id={}", s("id")),
1206 "CodeMap" => format!("op={} path={} name={}", s("operation"), s("path"), s("name")),
1207 _ => String::new(),
1208 }
1209}
1210
1211fn redact_result(result: &mut CallToolResult) {
1215 for content in &mut result.content {
1216 if let ContentBlock::Text(text) = content {
1217 if let std::borrow::Cow::Owned(scrubbed) = crate::utils::redact::redact(&text.text) {
1218 text.text = scrubbed;
1219 }
1220 }
1221 }
1222 if let Some(structured) = result.structured_content.as_mut() {
1223 crate::utils::redact::redact_json(structured);
1224 }
1225}
1226
1227#[cfg(test)]
1228mod audit_tests {
1229 use super::audit_summary;
1230 use serde_json::json;
1231
1232 #[test]
1233 fn bash_audit_summary_never_contains_command_content() {
1234 let args = json!({
1235 "action_json": {
1236 "type": "command",
1237 "command": "curl -H 'Authorization: secret-value' https://example.invalid"
1238 }
1239 });
1240 let summary = audit_summary("BashCommand", Some(&args));
1241 assert!(summary.starts_with("command bytes="));
1242 assert!(!summary.contains("secret-value"));
1243 assert!(!summary.contains("example.invalid"));
1244 }
1245}
1246
1247impl WinxService {
1248 async fn knowledge_transfer_prompt_text(&self) -> String {
1249 let mut text = String::from(
1250 "Prepare a concise handoff for another agent. Include active objective, current state, important files, changed files, blockers, validation already run, and exact next commands.\n",
1251 );
1252
1253 let state_snapshot = if let Some(slot) = self.active_slot().await {
1254 let guard = slot.lock().await;
1255 guard.as_ref().map(|state| {
1256 let whitelist = state
1257 .whitelist_for_overwrite
1258 .iter()
1259 .take(12)
1260 .map(|(path, data)| {
1261 format!(
1262 "- {} ({:.1}% read, {} lines)",
1263 path,
1264 data.get_percentage_read(),
1265 data.total_lines
1266 )
1267 })
1268 .collect::<Vec<_>>();
1269 (
1270 state.current_thread_id.clone(),
1271 state.workspace_root.clone(),
1272 state.cwd.clone(),
1273 state.mode.to_string(),
1274 whitelist,
1275 state.whitelist_for_overwrite.len(),
1276 )
1277 })
1278 } else {
1279 None
1280 };
1281
1282 let Some((thread_id, workspace_root, cwd, mode, whitelist, whitelist_count)) =
1283 state_snapshot
1284 else {
1285 text.push_str("\n# Current Winx state\nWinx is not initialized.\n");
1286 return text;
1287 };
1288
1289 let _ = writeln!(
1290 text,
1291 "\n# Current Winx state\nThread: {thread_id}\nWorkspace: {}\nCwd: {}\nMode: {mode}\nWhitelisted files: {whitelist_count}",
1292 workspace_root.display(),
1293 cwd.display()
1294 );
1295
1296 if !whitelist.is_empty() {
1297 text.push_str("\n# Recently readable files\n");
1298 text.push_str(&whitelist.join("\n"));
1299 text.push('\n');
1300 }
1301
1302 let active_files = crate::utils::workspace_stats::active_files(&workspace_root);
1303 if !active_files.is_empty() {
1304 text.push_str("\n# Active files by Winx usage\n");
1305 for file in active_files.iter().take(12) {
1306 let _ = writeln!(text, "- {file}");
1307 }
1308 }
1309
1310 if let Ok((repo_context, _)) = crate::utils::repo::get_repo_context(&workspace_root) {
1311 let repo_excerpt = repo_context.lines().take(80).collect::<Vec<_>>().join("\n");
1312 let _ = writeln!(text, "\n# Workspace context\n{repo_excerpt}");
1313 }
1314
1315 append_command_section(&mut text, "Git status", &workspace_root, ["status", "--short"]);
1316 append_command_section(
1317 &mut text,
1318 "Git diff stat",
1319 &workspace_root,
1320 ["diff", "--stat", "HEAD"],
1321 );
1322
1323 let sections = if mode == "architect" {
1327 "\n# Sections for the ContextSave description (architect mode)\n\
1328 - `# Objective` — project and task objective.\n\
1329 - `# All user instructions` — everything the user asked, verbatim.\n\
1330 - `# Designed plan` — the plan you designed, in detail.\n\
1331 - Provide all relevant file paths so the next agent can resume; err toward more.\n"
1332 } else {
1333 "\n# Sections for the ContextSave description\n\
1334 - `# Objective` — project and task objective.\n\
1335 - `# All user instructions` — everything the user asked, verbatim.\n\
1336 - `# Current status` — what's already done (not what's left).\n\
1337 - `# Pending issues with snippets` — verbatim errors/tracebacks/commands; be verbose.\n\
1338 - `# Build and development instructions` — how to build/run/test; leave empty if unknown.\n\
1339 - Provide all relevant file paths so the next agent can resume; err toward more.\n"
1340 };
1341 text.push_str(sections);
1342
1343 text.push_str(
1344 "\n# Handoff checklist\n- State what changed and why.\n- Include files touched and any user-owned dirty work to preserve.\n- Include validation commands already run and their result.\n- Include the next safest command to continue.\n",
1345 );
1346
1347 text
1348 }
1349
1350 async fn persist_state(&self, slot: &SharedBashState) {
1351 let guard = slot.lock().await;
1352 if let Some(state) = guard.as_ref() {
1353 if let Err(error) = state.save_state_to_disk() {
1354 warn!("Failed to persist bash state: {}", error);
1355 }
1356 }
1357 }
1358
1359 fn lenient_from_value<T: serde::de::DeserializeOwned>(
1364 args: Value,
1365 ) -> Result<T, serde_json::Error> {
1366 match serde_json::from_value::<T>(args.clone()) {
1367 Ok(value) => Ok(value),
1368 Err(first_err) => {
1369 let Value::Object(mut map) = args else {
1370 return Err(first_err);
1371 };
1372 let mut changed = false;
1373 for value in map.values_mut() {
1374 if let Value::String(text) = value {
1375 if let Ok(parsed) = serde_json::from_str::<Value>(text) {
1376 if parsed.is_object() || parsed.is_array() {
1377 *value = parsed;
1378 changed = true;
1379 }
1380 }
1381 }
1382 }
1383 if changed {
1384 serde_json::from_value::<T>(Value::Object(map))
1385 } else {
1386 Err(first_err)
1387 }
1388 }
1389 }
1390 }
1391
1392 async fn handle_initialize(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
1393 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1394 let mut initialize: Initialize = Self::lenient_from_value(args).map_err(|e| {
1395 McpError::invalid_request(format!("Invalid Initialize parameters: {e}"), None)
1396 })?;
1397
1398 let mut thread_id = normalize_thread_id(&initialize.thread_id);
1401 if thread_id.is_empty() {
1402 thread_id = generate_thread_id();
1403 initialize.thread_id.clone_from(&thread_id);
1404 }
1405 let (slot, _session_guard) = self.session_for(&thread_id).await;
1406
1407 match crate::tools::initialize::handle_tool_call_with_runtime(
1408 self.shell_runtime.as_ref(),
1409 &slot,
1410 initialize,
1411 )
1412 .await
1413 {
1414 Ok(result) => {
1415 self.persist_state(&slot).await;
1416 Ok(CallToolResult::success(vec![ContentBlock::text(result)]))
1417 }
1418 Err(e) => Err(to_mcp_error("Initialize", &e)),
1419 }
1420 }
1421
1422 async fn handle_bash_command(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
1423 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1424 let mut bash_command: BashCommand = serde_json::from_value(args).map_err(|e| {
1425 McpError::invalid_request(
1426 format!(
1427 "Invalid BashCommand parameters: {e}. Accepted forms include {{\"action_json\": {{\"command\": \"pwd\"}}}}, {{\"command\": \"pwd\"}}, or {{\"action_json\": {{\"type\": \"status_check\", \"status_check\": true}}}}."
1428 ),
1429 None,
1430 )
1431 })?;
1432
1433 let requested_thread_id = normalize_thread_id(&bash_command.thread_id);
1434 let (slot, _session_guard) = self.session_for(&requested_thread_id).await;
1435 if requested_thread_id.is_empty() {
1436 if let Some(thread_id) =
1441 slot.lock().await.as_ref().map(|state| state.current_thread_id.clone())
1442 {
1443 bash_command.thread_id = thread_id;
1444 }
1445 }
1446 match crate::tools::bash_command::handle_tool_call_with_runtime(
1447 self.shell_runtime.as_ref(),
1448 &slot,
1449 bash_command,
1450 )
1451 .await
1452 {
1453 Ok(output) => {
1454 self.persist_state(&slot).await;
1455 Ok(CallToolResult::success(vec![ContentBlock::text(output)]))
1456 }
1457 Err(e) => Err(to_mcp_error("BashCommand", &e)),
1458 }
1459 }
1460
1461 async fn handle_read_files(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
1462 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1463 let read_files: ReadFiles = Self::lenient_from_value(args).map_err(|e| {
1464 McpError::invalid_request(format!("Invalid ReadFiles parameters: {e}"), None)
1465 })?;
1466
1467 let (slot, _session_guard) =
1468 self.session_for(&normalize_thread_id(&read_files.thread_id)).await;
1469 match crate::tools::read_files::handle_tool_call(&slot, read_files).await {
1470 Ok(result) => {
1471 self.persist_state(&slot).await;
1472 Ok(CallToolResult::success(vec![ContentBlock::text(result)]))
1473 }
1474 Err(e) => Err(to_mcp_error("ReadFiles", &e)),
1475 }
1476 }
1477
1478 async fn handle_file_write_or_edit(
1479 &self,
1480 args: Option<Value>,
1481 ) -> Result<CallToolResult, McpError> {
1482 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1483 let file_write_or_edit: FileWriteOrEdit = Self::lenient_from_value(args).map_err(|e| {
1484 McpError::invalid_request(format!("Invalid FileWriteOrEdit parameters: {e}"), None)
1485 })?;
1486
1487 let (slot, _session_guard) =
1488 self.session_for(&normalize_thread_id(&file_write_or_edit.thread_id)).await;
1489 match crate::tools::file_write_or_edit::handle_tool_call(&slot, file_write_or_edit).await {
1490 Ok(result) => {
1491 self.persist_state(&slot).await;
1492 Ok(CallToolResult::success(vec![ContentBlock::text(result)]))
1493 }
1494 Err(e) => Err(to_mcp_error("FileWriteOrEdit", &e)),
1495 }
1496 }
1497
1498 async fn handle_multi_file_edit(
1499 &self,
1500 args: Option<Value>,
1501 ) -> Result<CallToolResult, McpError> {
1502 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1503 let multi: MultiFileEdit = Self::lenient_from_value(args).map_err(|e| {
1504 McpError::invalid_request(format!("Invalid MultiFileEdit parameters: {e}"), None)
1505 })?;
1506
1507 let (slot, _session_guard) = self.session_for(&normalize_thread_id(&multi.thread_id)).await;
1508 match crate::tools::multi_file_edit::handle_tool_call(&slot, multi).await {
1509 Ok(result) => {
1510 self.persist_state(&slot).await;
1511 Ok(CallToolResult::success(vec![ContentBlock::text(result)]))
1512 }
1513 Err(e) => Err(to_mcp_error("MultiFileEdit", &e)),
1514 }
1515 }
1516
1517 async fn handle_undo_edit(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
1518 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1519 let undo: UndoEdit = Self::lenient_from_value(args).map_err(|e| {
1520 McpError::invalid_request(format!("Invalid UndoEdit parameters: {e}"), None)
1521 })?;
1522
1523 let (slot, _session_guard) = self.session_for(&normalize_thread_id(&undo.thread_id)).await;
1524 match crate::tools::undo_edit::handle_tool_call(&slot, undo).await {
1525 Ok(result) => {
1526 self.persist_state(&slot).await;
1527 Ok(CallToolResult::success(vec![ContentBlock::text(result)]))
1528 }
1529 Err(e) => Err(to_mcp_error("UndoEdit", &e)),
1530 }
1531 }
1532
1533 async fn handle_context_save(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
1534 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1535 let context_save: ContextSave = Self::lenient_from_value(args).map_err(|e| {
1536 McpError::invalid_request(format!("Invalid ContextSave parameters: {e}"), None)
1537 })?;
1538
1539 let (slot, _session_guard) =
1540 self.session_for(&normalize_thread_id(&context_save.thread_id)).await;
1541 match crate::tools::context_save::handle_tool_call(&slot, context_save).await {
1542 Ok(result) => {
1543 self.persist_state(&slot).await;
1544 Ok(CallToolResult::success(vec![ContentBlock::text(result)]))
1545 }
1546 Err(e) => Err(to_mcp_error("ContextSave", &e)),
1547 }
1548 }
1549
1550 async fn handle_read_image(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
1551 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1552 let read_image: ReadImage = Self::lenient_from_value(args).map_err(|e| {
1553 McpError::invalid_request(format!("Invalid ReadImage parameters: {e}"), None)
1554 })?;
1555
1556 let (slot, _session_guard) =
1557 self.session_for(&normalize_thread_id(&read_image.thread_id)).await;
1558 match crate::tools::read_image::handle_tool_call(&slot, read_image).await {
1559 Ok((mime_type, base64_data)) => {
1560 self.persist_state(&slot).await;
1561 Ok(CallToolResult::success(vec![ContentBlock::image(base64_data, mime_type)]))
1565 }
1566 Err(e) => Err(to_mcp_error("ReadImage", &e)),
1567 }
1568 }
1569
1570 async fn handle_code_map(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
1571 let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
1572 let code_map: CodeMap = Self::lenient_from_value(args).map_err(|e| {
1573 McpError::invalid_request(format!("Invalid CodeMap parameters: {e}"), None)
1574 })?;
1575
1576 let (slot, _session_guard) =
1578 self.session_for(&normalize_thread_id(&code_map.thread_id)).await;
1579 match crate::tools::code_map::handle_tool_call(&slot, code_map).await {
1580 Ok((text, structured)) => {
1581 let mut result = CallToolResult::success(vec![ContentBlock::text(text)]);
1582 result.structured_content = Some(structured);
1583 Ok(result)
1584 }
1585 Err(e) => Err(to_mcp_error("CodeMap", &e)),
1586 }
1587 }
1588}
1589
1590pub async fn start_winx_server() -> Result<(), Box<dyn std::error::Error>> {
1592 start_winx_server_with_runtime(configured_shell_runtime().await?).await
1593}
1594
1595pub async fn start_winx_server_with_runtime(
1597 runtime: Arc<dyn ShellRuntime>,
1598) -> Result<(), Box<dyn std::error::Error>> {
1599 info!("Starting Winx MCP Server");
1600 let service = WinxService::with_runtime(SessionIsolation::Lenient, runtime);
1601 let (stdin, stdout) = stdio();
1602 let transport = AsyncRwTransport::new_server(stdin, stdout);
1603 let server = service.serve(transport).await?;
1604 server.waiting().await?;
1605 Ok(())
1606}
1607
1608#[cfg(test)]
1609mod session_registry_tests {
1610 use super::*;
1611
1612 #[tokio::test]
1613 async fn distinct_threads_get_distinct_sessions() {
1614 let svc = WinxService::new();
1615 let (a, _) = svc.session_for("thread_a").await;
1616 let (b, _) = svc.session_for("thread_b").await;
1617 assert!(!Arc::ptr_eq(&a, &b));
1620 let (a2, _) = svc.session_for("thread_a").await;
1622 assert!(Arc::ptr_eq(&a, &a2));
1623 }
1624
1625 #[tokio::test]
1626 async fn empty_thread_id_falls_back_to_last_active() {
1627 let svc = WinxService::new();
1628 let (_a, _) = svc.session_for("thread_a").await;
1629 let (b, _) = svc.session_for("thread_b").await; let (fallback, _) = svc.session_for("").await;
1632 assert!(Arc::ptr_eq(&b, &fallback));
1633 assert!(svc.active_slot().await.is_some());
1634 }
1635
1636 #[tokio::test]
1637 async fn strict_isolation_empty_thread_id_does_not_steal_active_session() {
1638 let svc = WinxService::with_isolation(SessionIsolation::Strict);
1639 let (a, _) = svc.session_for("thread_a").await; let (anon, _) = svc.session_for("").await;
1641 assert!(!Arc::ptr_eq(&a, &anon));
1643 let (anon2, _) = svc.session_for("").await;
1645 assert!(Arc::ptr_eq(&anon, &anon2));
1646 let (b, _) = svc.session_for("thread_b").await;
1648 assert!(!Arc::ptr_eq(&a, &b));
1649 }
1650
1651 #[tokio::test]
1652 async fn lru_eviction_caps_live_sessions() {
1653 let svc = WinxService::new();
1654 for i in 0..(MAX_SESSIONS + 5) {
1655 let (_, _) = svc.session_for(&format!("t{i}")).await;
1656 }
1657 let reg = svc.sessions.lock().await;
1658 assert!(reg.slots.len() <= MAX_SESSIONS, "session count {} over cap", reg.slots.len());
1659 }
1660
1661 #[tokio::test]
1662 async fn in_flight_session_is_not_evicted() {
1663 let svc = WinxService::new();
1664 let (_keep_slot, _keep_guard) = svc.session_for("keep").await;
1666 for i in 0..(MAX_SESSIONS + 10) {
1669 let (_, _) = svc.session_for(&format!("filler{i}")).await;
1670 }
1671 let reg = svc.sessions.lock().await;
1672 assert!(
1673 reg.slots.contains_key("keep"),
1674 "an in-flight session must survive LRU eviction churn"
1675 );
1676 }
1677}
1678
1679#[cfg(test)]
1680mod task_lifecycle_tests {
1681 use super::*;
1682
1683 #[tokio::test]
1684 async fn interrupt_waits_until_main_shell_is_reusable() {
1685 let service = WinxService::new();
1686 let initialized = service
1687 .handle_initialize(Some(serde_json::json!({
1688 "type": "first_call",
1689 "any_workspace_path": std::env::temp_dir(),
1690 "mode_name": "wcgw",
1691 "thread_id": "task_cancel_regression"
1692 })))
1693 .await;
1694 assert!(initialized.is_ok(), "test session failed to initialize: {initialized:?}");
1695
1696 let worker = {
1697 let service = service.clone();
1698 tokio::spawn(async move {
1699 service
1700 .handle_bash_command(Some(serde_json::json!({
1701 "action_json": {
1702 "type": "command",
1703 "command": "sleep 30",
1704 "is_background": false
1705 },
1706 "thread_id": "task_cancel_regression"
1707 })))
1708 .await
1709 })
1710 };
1711 tokio::time::sleep(Duration::from_millis(300)).await;
1712 worker.abort();
1713 service.interrupt_task_thread("task_cancel_regression").await;
1714 let interrupted = tokio::time::timeout(Duration::from_secs(5), worker).await;
1715 match interrupted {
1716 Ok(Ok(_)) => {}
1719 Ok(Err(error)) if error.is_cancelled() => {}
1720 Ok(Err(error)) => panic!("aborted task worker failed instead of cancelling: {error}"),
1721 Err(_) => panic!("aborted task worker did not settle"),
1722 }
1723
1724 let recovery = service
1725 .handle_bash_command(Some(serde_json::json!({
1726 "action_json": {
1727 "type": "command",
1728 "command": "printf task-shell-recovered",
1729 "is_background": false
1730 },
1731 "thread_id": "task_cancel_regression"
1732 })))
1733 .await;
1734 let rendered = format!("{recovery:?}");
1735 assert!(
1736 recovery.is_ok() && rendered.contains("task-shell-recovered"),
1737 "shell was not reusable after cancellation: {rendered}"
1738 );
1739 }
1740}
1741
1742#[cfg(test)]
1743mod schema_tests {
1744 #![allow(clippy::expect_used)]
1745
1746 use super::{root_uri_to_path, schema_to_input_schema, strip_schema_titles, winx_tools};
1747 use rmcp::model::ProtocolVersion;
1748 use rmcp::ServerHandler;
1749 use serde_json::json;
1750 use std::path::PathBuf;
1751
1752 #[test]
1753 fn strips_titles_from_schema_nodes_only() {
1754 let mut v = json!({
1755 "type": "object",
1756 "title": "ShouldGo",
1757 "properties": {
1758 "title": { "type": "string", "title": "InnerGoes" }
1761 }
1762 });
1763 strip_schema_titles(&mut v);
1764 assert!(v.get("title").is_none(), "schema-node title not stripped");
1765 let props = v.get("properties").and_then(serde_json::Value::as_object);
1766 assert!(
1767 props.is_some_and(|p| p.contains_key("title")),
1768 "property key named 'title' must be preserved"
1769 );
1770 assert!(
1771 props.and_then(|p| p.get("title")).and_then(|t| t.get("title")).is_none(),
1772 "inner schema title not stripped"
1773 );
1774 }
1775
1776 #[test]
1777 fn real_tool_schema_carries_no_titles() {
1778 let schema = schema_to_input_schema::<crate::types::Initialize>();
1779 let blob = serde_json::to_string(&*schema).unwrap_or_default();
1780 assert!(!blob.contains("\"title\""), "tool schema still contains titles: {blob}");
1781 }
1782
1783 #[test]
1784 fn initialize_schema_has_no_dangling_refs() {
1785 let schema = schema_to_input_schema::<crate::types::Initialize>();
1791 let blob = serde_json::to_string(&*schema).unwrap_or_default();
1792 assert!(
1793 !blob.contains("#/definitions/"),
1794 "schema has a draft-07 #/definitions ref (dangling for 2020-12 clients): {blob}"
1795 );
1796 assert!(
1797 !blob.contains("definitions/ModeName"),
1798 "ModeName is referenced but never defined: {blob}"
1799 );
1800 assert!(
1802 blob.contains("wcgw") && blob.contains("architect") && blob.contains("code_writer"),
1803 "mode_name enum not inlined: {blob}"
1804 );
1805 }
1806
1807 #[test]
1808 fn advertises_latest_protocol_and_tasks_extension() {
1809 let info = ServerHandler::get_info(&super::WinxService::new());
1810 assert_eq!(info.protocol_version, ProtocolVersion::V_2026_07_28);
1811 assert!(info.capabilities.supports_tasks());
1812 }
1813
1814 #[test]
1815 fn bash_command_and_code_map_output_schema_are_advertised() {
1816 let tools = winx_tools();
1817 assert!(tools.iter().any(|tool| tool.name == "BashCommand"));
1818
1819 let code_map = tools.iter().find(|tool| tool.name == "CodeMap").expect("CodeMap");
1820 let output = code_map.output_schema.as_ref().expect("CodeMap outputSchema");
1821 assert!(output.get("properties").is_some_and(|properties| {
1822 properties.as_object().is_some_and(|properties| properties.contains_key("truncated"))
1823 }));
1824 let blob = serde_json::to_string(output).unwrap_or_default();
1825 assert!(!blob.contains("\"format\":\"uint\""), "unsupported uint format: {blob}");
1826 }
1827
1828 #[test]
1829 fn local_root_uri_decoding_rejects_remote_authorities() {
1830 assert_eq!(
1831 root_uri_to_path("file:///tmp/a%20project"),
1832 Some(PathBuf::from("/tmp/a project"))
1833 );
1834 assert_eq!(
1835 root_uri_to_path("file://localhost/tmp/project"),
1836 Some(PathBuf::from("/tmp/project"))
1837 );
1838 assert_eq!(root_uri_to_path("file://remote.example/tmp/project"), None);
1839 assert_eq!(root_uri_to_path("https://example.com/project"), None);
1840 }
1841}
1842
1843#[cfg(test)]
1844mod discovery_protocol_tests {
1845 #![allow(clippy::expect_used)]
1846
1847 use super::WinxService;
1848 use rmcp::{transport::async_rw::AsyncRwTransport, ServiceExt};
1849 use serde_json::{json, Value};
1850 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
1851
1852 #[tokio::test]
1853 async fn discover_probe_negotiates_modern_stateless_lifecycle() {
1854 let (client, server) = tokio::io::duplex(32 * 1024);
1855 let (server_read, server_write) = tokio::io::split(server);
1856 let transport = AsyncRwTransport::new_server(server_read, server_write);
1857 let task = tokio::spawn(async move {
1858 let running = WinxService::new()
1859 .serve(transport)
1860 .await
1861 .expect("server should accept modern discovery");
1862 running.waiting().await.expect("server should shut down cleanly");
1863 });
1864
1865 let (client_read, mut client_write) = tokio::io::split(client);
1866 let mut client_read = BufReader::new(client_read);
1867
1868 client_write
1869 .write_all(
1870 b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"server/discover\",\"params\":{\"_meta\":{\"io.modelcontextprotocol/protocolVersion\":\"2026-07-28\",\"io.modelcontextprotocol/clientInfo\":{\"name\":\"discovery-test\",\"version\":\"1\"},\"io.modelcontextprotocol/clientCapabilities\":{}}}}\n",
1871 )
1872 .await
1873 .expect("write discovery probe");
1874 let mut line = String::new();
1875 client_read.read_line(&mut line).await.expect("read discovery response");
1876 let response: Value = serde_json::from_str(&line).expect("valid discovery response");
1877 assert_eq!(response["id"], 1);
1878 assert!(response["result"]["supportedVersions"]
1879 .as_array()
1880 .is_some_and(|versions| versions.iter().any(|version| version == "2026-07-28")));
1881
1882 drop(client_read);
1883 drop(client_write);
1884 task.await.expect("server task");
1885 }
1886}
1887
1888#[cfg(test)]
1889mod error_mapping_tests {
1890 use super::*;
1891 use rmcp::model::ErrorCode;
1892 use std::path::PathBuf;
1893
1894 fn code_of(err: &WinxError) -> ErrorCode {
1895 to_mcp_error("Tool", err).code
1896 }
1897
1898 #[test]
1899 fn client_caused_errors_map_to_invalid_request() {
1900 assert_eq!(
1902 code_of(&WinxError::RecoverableSuggestionError {
1903 message: "bad arg".into(),
1904 suggestion: "try x".into(),
1905 }),
1906 ErrorCode::INVALID_REQUEST,
1907 );
1908 assert_eq!(
1910 code_of(&WinxError::ParseError("unexpected token".into())),
1911 ErrorCode::INVALID_REQUEST,
1912 );
1913 assert_eq!(
1915 code_of(&WinxError::FileAccessError {
1916 path: PathBuf::from("/nope"),
1917 message: "no such file".into(),
1918 }),
1919 ErrorCode::INVALID_REQUEST,
1920 );
1921 }
1922
1923 #[test]
1924 fn server_caused_errors_stay_internal_error() {
1925 assert_eq!(
1926 code_of(&WinxError::IoError(std::io::Error::other("disk gone"))),
1927 ErrorCode::INTERNAL_ERROR,
1928 );
1929 assert_eq!(
1930 code_of(&WinxError::BashStateLockError("poisoned".into())),
1931 ErrorCode::INTERNAL_ERROR,
1932 );
1933 }
1934}
1935
1936#[cfg(all(test, feature = "loom"))]
1951mod loom_tests {
1952 use loom::sync::atomic::{AtomicUsize, Ordering};
1953 use loom::sync::Arc;
1954
1955 #[derive(Clone)]
1956 struct PinModel {
1957 count: Arc<AtomicUsize>,
1958 }
1959
1960 struct GuardModel {
1961 count: Arc<AtomicUsize>,
1962 }
1963
1964 impl PinModel {
1965 fn new() -> Self {
1966 Self { count: Arc::new(AtomicUsize::new(0)) }
1967 }
1968 fn acquire(&self) -> GuardModel {
1969 self.count.fetch_add(1, Ordering::SeqCst);
1970 GuardModel { count: self.count.clone() }
1971 }
1972 fn is_pinned(&self) -> bool {
1973 self.count.load(Ordering::SeqCst) > 0
1974 }
1975 }
1976
1977 impl Drop for GuardModel {
1978 fn drop(&mut self) {
1979 self.count.fetch_sub(1, Ordering::SeqCst);
1980 }
1981 }
1982
1983 #[test]
1989 fn loom_concurrent_guards_balance_to_zero() {
1990 loom::model(|| {
1991 let pin = PinModel::new();
1992 let p1 = pin.clone();
1993 let p2 = pin.clone();
1994 let h1 = loom::thread::spawn(move || drop(p1.acquire()));
1995 let h2 = loom::thread::spawn(move || drop(p2.acquire()));
1996 assert!(h1.join().is_ok());
1997 assert!(h2.join().is_ok());
1998 assert_eq!(pin.count.load(Ordering::SeqCst), 0, "pin must settle back to 0");
1999 });
2000 }
2001
2002 #[test]
2006 fn loom_live_guard_always_reads_pinned() {
2007 loom::model(|| {
2008 let pin = PinModel::new();
2009 let observer = pin.clone();
2010 let held = pin.acquire();
2011 let worker = pin.clone();
2014 let h = loom::thread::spawn(move || drop(worker.acquire()));
2015 assert!(observer.is_pinned(), "session with a live guard must read pinned");
2016 assert!(h.join().is_ok());
2017 assert!(observer.is_pinned(), "still pinned while the first guard lives");
2018 drop(held);
2019 assert!(!observer.is_pinned(), "all guards gone -> unpinned");
2020 });
2021 }
2022}