Skip to main content

winx_code_agent/
server.rs

1//! Winx MCP Server implementation using the official rmcp SDK.
2//! Core MCP tools only - High performance shell and file management
3
4use 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
42/// Map a domain [`WinxError`] to the right JSON-RPC error kind.
43///
44/// Client-caused failures (bad params, security violations, mode restrictions,
45/// state-not-initialized, ambiguous/not-found edits) become `invalid_request`
46/// so the model knows the fix is on its side; genuine server-side failures
47/// (PTY spawn, IO, lock poisoning, persistence) stay `internal_error`.
48fn to_mcp_error(tool: &str, err: &WinxError) -> McpError {
49    let msg = format!("{tool} failed: {err}");
50    // Exhaustive on purpose — NO wildcard arm. A new `WinxError` variant must be
51    // classified by hand (client vs server) or the build breaks. This is the
52    // guard that keeps the JSON-RPC error codes honest over time.
53    match err {
54        // Client-caused: the model can fix its own input or usage. -> invalid_request
55        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        // Server-caused: genuine internal faults the model cannot fix. -> internal_error
78        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
96/// Type alias for the shared bash state - uses `tokio::sync::Mutex` for async safety
97pub type SharedBashState = Arc<Mutex<Option<BashState>>>;
98
99/// Helper function to create JSON schema from schemars Schema
100fn 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    // schemars stamps a redundant `title` (usually just the type/field name) on
104    // every schema node; the LLM pays tokens for it on every tool call for zero
105    // signal. Strip it — context-aware, so a user data field literally named
106    // "title" is never touched.
107    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
115/// Remove integer formats that are valid JSON Schema annotations but rejected
116/// by some MCP clients (notably Claude Code's schema compiler).
117fn 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
136/// Recursively remove `title` keys from JSON-Schema nodes only.
137///
138/// A dict is treated as a schema node when it carries a schema-shaped key
139/// (`type`/`$ref`/`properties`/`items`/`enum`/`const`/`anyOf`/`allOf`/`oneOf`/
140/// `additionalProperties`). This mirrors wcgw's `recursive_purge_dict_key` so a
141/// property whose *name* is "title" keeps its value.
142fn 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
415/// Upper bound on concurrently-live sessions. Each session owns a PTY (a real
416/// bash process), so we evict the least-recently-used one past this to avoid
417/// leaking shells across many short-lived `thread_id`s.
418const MAX_SESSIONS: usize = 32;
419
420/// Per-`thread_id` shell sessions. Each `thread_id` gets its own
421/// `BashState`/PTY, so concurrent threads (or HTTP clients sharing the service)
422/// never execute in each other's shell. Tools that don't carry a `thread_id`
423/// (legacy clients) fall back to the most recently active session.
424#[derive(Default)]
425struct SessionRegistry {
426    slots: HashMap<String, SharedBashState>,
427    /// Last-use timestamps for LRU eviction.
428    last_used: HashMap<String, Instant>,
429    /// In-flight operation count per session. A session whose count is > 0 (it
430    /// has a live [`SessionGuard`]) is never chosen as an LRU eviction victim, so
431    /// a long-running call can't have its shell pulled out from under it.
432    in_flight: HashMap<String, SessionPin>,
433    /// Most recently addressed session, used as the fallback for tool calls
434    /// that omit a `thread_id`.
435    last_active: Option<String>,
436}
437
438/// Lock-free pin counter for one session. A live [`SessionGuard`] keeps the
439/// count `> 0`, which marks the session as in-flight so LRU eviction skips it.
440/// Clones share the same counter — the registry hands a clone to each
441/// concurrent call. The acquire/release discipline (the guard's `Drop` runs
442/// *outside* the registry lock) is loom model-checked by `loom_tests`, which
443/// mirrors this exact logic; keep the two in lockstep.
444#[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    /// Mark a new operation in flight, returning the RAII guard that releases it.
457    fn acquire(&self) -> SessionGuard {
458        self.count.fetch_add(1, Ordering::SeqCst);
459        SessionGuard { count: self.count.clone() }
460    }
461
462    /// Whether any operation is currently in flight on this session.
463    fn is_pinned(&self) -> bool {
464        self.count.load(Ordering::SeqCst) > 0
465    }
466}
467
468/// RAII marker that a session has an operation in flight. Bumps the session's
469/// in-flight counter on creation and drops it on `Drop` — a plain synchronous
470/// `Drop`, so it works even though the registry behind it lives in an async
471/// `tokio::Mutex`. While any guard is alive the session is pinned against LRU
472/// eviction.
473struct 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/// How an empty `thread_id` is resolved by the session registry.
484#[derive(Clone, Copy, PartialEq, Eq, Debug)]
485pub enum SessionIsolation {
486    /// Local stdio transport — a single client. An empty `thread_id` falls back
487    /// to the last-active session (legacy single-client convenience).
488    Lenient,
489    /// Remote HTTP transport — multiple clients behind one shared bearer token.
490    /// The last-active fallback is disabled, so an empty `thread_id` from one
491    /// client can never resolve to another client's shell; such calls get a
492    /// dedicated anonymous slot instead. Clients must use the `thread_id` that
493    /// `Initialize` hands back.
494    ///
495    /// Residual: two clients that deliberately send the *same* explicit
496    /// `thread_id` still share a shell. Closing that needs per-client tokens,
497    /// which the single shared-token model intentionally doesn't provide.
498    Strict,
499}
500
501/// Winx service. Holds the per-thread session registry.
502#[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    /// Version information for the service
509    pub version: String,
510    /// How empty `thread_id`s are resolved (see [`SessionIsolation`]).
511    isolation: SessionIsolation,
512}
513
514impl Default for WinxService {
515    fn default() -> Self {
516        Self::new()
517    }
518}
519
520impl WinxService {
521    /// Create a new `WinxService` for the local stdio transport (lenient,
522    /// single-client session isolation).
523    pub fn new() -> Self {
524        Self::with_isolation(SessionIsolation::Lenient)
525    }
526
527    /// Create a `WinxService` with an explicit session-isolation policy. The
528    /// HTTP transport uses [`SessionIsolation::Strict`].
529    pub fn with_isolation(isolation: SessionIsolation) -> Self {
530        Self::with_runtime(isolation, Arc::new(EmbeddedShellRuntime))
531    }
532
533    /// Create a service with an explicit shell-runtime boundary.
534    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    /// Resolve the session slot for a `thread_id`, creating it if absent.
547    ///
548    /// An empty `thread_id` resolves, under [`SessionIsolation::Lenient`], to the
549    /// most recently active session (the compatibility path for tools — and older
550    /// clients — that don't send one); under [`SessionIsolation::Strict`] it gets
551    /// a dedicated anonymous slot so remote clients can't land in each other's shell.
552    /// Marks the slot as most-recently-used and evicts the LRU session when over
553    /// [`MAX_SESSIONS`].
554    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                // Single local client: reuse the most recently active shell.
559                SessionIsolation::Lenient => {
560                    reg.last_active.clone().unwrap_or_else(|| "default".to_string())
561                }
562                // Many remote clients sharing one token: never resolve to
563                // someone else's active shell. Empty thread_id gets its own slot.
564                SessionIsolation::Strict => "anonymous".to_string(),
565            }
566        } else {
567            thread_id.to_string()
568        };
569
570        // Evict the LRU session if adding a brand-new key would exceed the cap —
571        // but never evict a session with an operation in flight (that would pull
572        // the shell out from under a concurrent long-running call). If every other
573        // session is busy we briefly exceed the cap rather than corrupt one.
574        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    /// The most recently active session slot, without creating one. Used by
608    /// session-agnostic surfaces (e.g. the handoff prompt).
609    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    /// Run one task-augmented foreground shell command through to completion.
615    /// Normal `BashCommand` calls intentionally yield after a short wait; an MCP
616    /// Task keeps polling that same PTY and exposes one final `CallToolResult`.
617    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    /// Bootstrap a local single-client session from the client's first MCP Root.
710    /// Roots are part of stable 2025-11-25, although rmcp 1.8 marks the API as
711    /// deprecated in anticipation of the later SEP-2577 draft.
712    #[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        // Refuse non-local file authorities (file://host/path).
914        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
921/// Server icon (96x96 PNG) as a self-contained data URI, per the 2025-11-25
922/// icons revision. A data URI works over stdio and HTTP alike — no extra
923/// endpoint or auth exemption needed for clients to fetch it.
924fn 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
933/// `ServerHandler` implementation
934impl 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        // A changed Roots list is only used as a retry if Winx is still
963        // uninitialized. Never switch an active shell's workspace implicitly.
964        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        // This prompt embeds git status/diff and file context, so scrub it too —
1067        // it returns a GetPromptResult and bypasses the call_tool choke point.
1068        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(&param)
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        // Audit trail: one structured line per tool call, including the outcome
1125        // and wall-clock. Successes were previously silent — only errors logged —
1126        // which made debugging remote (ChatGPT) sessions guesswork.
1127        // Redact the audit line too: it embeds tool arguments (a BashCommand's
1128        // command line, a file path) that can carry secrets into the logs.
1129        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        // Scrub credentials (provider keys, JWTs, PEM blocks, user:pass URLs)
1147        // from everything leaving the server — shell output, file contents,
1148        // search results — before it reaches the model. `WINX_NO_REDACT` opts out.
1149        let result = match result {
1150            Ok(mut call) => {
1151                redact_result(&mut call);
1152                Ok(call)
1153            }
1154            Err(mut error) => {
1155                // The error message is returned to the client/model verbatim and
1156                // often echoes a path, env value, or command output that can carry
1157                // a secret — scrub it the same as success content.
1158                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            // error.message was already redacted above.
1167            Err(error) => warn!(tool = %tool, ms, "tool call error — {summary}: {}", error.message),
1168        }
1169        result
1170    }
1171}
1172
1173/// Build a short, non-sensitive audit summary of a tool call's arguments.
1174fn 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
1211/// Scrub credential patterns from a tool result's text blocks and structured
1212/// content, in place. Central choke point so every tool is covered (see
1213/// [`crate::utils::redact`]).
1214fn 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        // Sections the ContextSave `description` should contain, tailored to the
1324        // mode: architect produces a plan (no edits), the others produce a status
1325        // + pending-issues handoff (wcgw parity: WCGW_KT vs ARCHITECT_KT).
1326        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    /// Deserialize `args` into `T`, retrying once after JSON-decoding any string
1360    /// field that is itself an encoded object/array. LLMs sometimes send a nested
1361    /// param (e.g. `code_writer_config`) as a JSON string instead of an object;
1362    /// wcgw applies the same leniency in its tool dispatch.
1363    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        // Resolve the session key. A first_call may omit thread_id; generate one
1399        // here and write it back so the handler and the registry agree on the id.
1400        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            // Initialize historically returns a generated thread_id, but local
1437            // clients that rely on MCP Roots may not know it. Resolve an omitted
1438            // id to the session selected by `session_for` before the lower-level
1439            // BashCommand validation runs.
1440            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                // Return a real image content block (not base64 as text) so the
1562                // model can actually see the image. rmcp's `ContentBlock::image`
1563                // takes (data, mime_type).
1564                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        // Read-only: no persist_state (nothing in the session changes).
1577        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
1590/// Create and start the Winx MCP server
1591pub async fn start_winx_server() -> Result<(), Box<dyn std::error::Error>> {
1592    start_winx_server_with_runtime(configured_shell_runtime().await?).await
1593}
1594
1595/// Create and start the stdio MCP server with an explicit shell runtime.
1596pub 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        // Two thread_ids must own two separate slots — the whole point of the
1618        // registry: thread B never executes in thread A's shell.
1619        assert!(!Arc::ptr_eq(&a, &b));
1620        // Same id round-trips to the same slot.
1621        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; // now last-active
1630                                                        // A tool call without a thread_id reuses the most recently active slot.
1631        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; // would be "last-active" under Lenient
1640        let (anon, _) = svc.session_for("").await;
1641        // Strict mode: an empty thread_id must NOT resolve to another client's shell.
1642        assert!(!Arc::ptr_eq(&a, &anon));
1643        // Anonymous calls share one dedicated slot rather than hijacking named ones.
1644        let (anon2, _) = svc.session_for("").await;
1645        assert!(Arc::ptr_eq(&anon, &anon2));
1646        // Explicit thread_ids stay isolated as always.
1647        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        // "keep" is created first (so it's the LRU candidate) and held busy.
1665        let (_keep_slot, _keep_guard) = svc.session_for("keep").await;
1666        // Saturate the cap and then churn well past it; eviction must skip the
1667        // in-flight "keep" and evict idle fillers instead.
1668        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            // A foreground BashCommand may have already returned its documented
1717            // incremental "still running" result before abort wins the race.
1718            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                // a user field literally named "title" must survive as a key,
1759                // and its inner schema's own title must be stripped.
1760                "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        // The Initialize input embeds ModeName. A manual JsonSchema impl used to
1786        // emit a `$ref` to `#/definitions/ModeName` that was never populated —
1787        // strict MCP clients (Cline, and likely ChatGPT) reject the tool call
1788        // with "Reference not found: #/definitions/ModeName". Guard against any
1789        // dangling/definitions-style ref reappearing.
1790        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        // mode_name must be a concrete, self-contained string enum.
1801        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        // A validation suggestion is the model's to fix, not a server fault.
1901        assert_eq!(
1902            code_of(&WinxError::RecoverableSuggestionError {
1903                message: "bad arg".into(),
1904                suggestion: "try x".into(),
1905            }),
1906            ErrorCode::INVALID_REQUEST,
1907        );
1908        // Parse failures come from the model's input.
1909        assert_eq!(
1910            code_of(&WinxError::ParseError("unexpected token".into())),
1911            ErrorCode::INVALID_REQUEST,
1912        );
1913        // A bad path the model handed us is something it can correct.
1914        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/// Loom model-check of the [`SessionPin`] counting discipline — the one piece
1937/// of session state touched off the registry lock (a guard's `Drop` decrements
1938/// while a concurrent eviction may be reading it via `is_pinned`). Loom can't
1939/// model the surrounding `tokio::Mutex`, so we isolate the lock-free counter and
1940/// exhaustively check it under every interleaving.
1941///
1942/// `SessionPin` itself stays on std atomics: instrumenting it would force every
1943/// `WinxService` test (which constructs pins outside `loom::model`) to panic
1944/// under `--all-features`. Instead `PinModel` below mirrors `SessionPin`'s exact
1945/// logic — `acquire` = `fetch_add(SeqCst)`, guard `Drop` = `fetch_sub(SeqCst)`,
1946/// `is_pinned` = `load(SeqCst) > 0`. Keep the two in lockstep.
1947///
1948/// Built only under the `loom` feature. Run with:
1949///   cargo test --features loom --lib loom_
1950#[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    /// Two concurrent in-flight ops on the same session: under *every* thread
1984    /// interleaving each `acquire` is balanced by its guard's `Drop`, so the pin
1985    /// returns to exactly 0 and never underflows. An underflow would wrap the
1986    /// counter to a huge value, making a finished session read as permanently
1987    /// in-flight and leak past LRU eviction forever.
1988    #[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    /// One op stays in flight while another runs concurrently: the observer must
2003    /// never see the live session as unpinned (which the eviction filter would
2004    /// read as "safe to evict", pulling the shell out from under the running op).
2005    #[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            // A second op starts and finishes concurrently; throughout, the
2012            // first guard keeps the session pinned for the observer.
2013            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}