winx_code_agent/
server.rs

1//! Winx MCP Server implementation using rmcp 0.12.0
2//! Core MCP tools only - High performance shell and file management
3
4use rmcp::{
5    model::{
6        Annotated, CallToolRequestParams, CallToolResult, Content, Implementation,
7        ListResourcesResult, ListToolsResult, PaginatedRequestParams, ProtocolVersion, RawResource,
8        ReadResourceRequestParams, ReadResourceResult, ResourceContents, ServerCapabilities,
9        ServerInfo, Tool, ToolAnnotations,
10    },
11    service::{RequestContext, RoleServer},
12    transport::stdio,
13    ErrorData as McpError, ServerHandler, ServiceExt,
14};
15use schemars::schema_for;
16use serde_json::Value;
17use std::sync::{Arc, OnceLock};
18use tokio::sync::Mutex;
19use tracing::{info, warn};
20
21use crate::state::BashState;
22use crate::types::{BashCommand, ContextSave, FileWriteOrEdit, Initialize, ReadFiles, ReadImage};
23
24/// Type alias for the shared bash state - uses `tokio::sync::Mutex` for async safety
25pub type SharedBashState = Arc<Mutex<Option<BashState>>>;
26
27/// Helper function to create JSON schema from schemars Schema
28fn schema_to_input_schema<T: schemars::JsonSchema>() -> Arc<serde_json::Map<String, Value>> {
29    let schema = schema_for!(T);
30    let value = serde_json::to_value(schema).unwrap_or(Value::Object(serde_json::Map::new()));
31    match value {
32        Value::Object(map) => Arc::new(map),
33        _ => Arc::new(serde_json::Map::new()),
34    }
35}
36
37fn mcp_tool<T: schemars::JsonSchema>(
38    name: &'static str,
39    description: &'static str,
40    annotations: ToolAnnotations,
41) -> Tool {
42    Tool::new(name, description, schema_to_input_schema::<T>()).with_annotations(annotations)
43}
44
45const INITIALIZE_DESCRIPTION: &str =
46    "- Always call this at the start of the conversation before using any of the shell tools from wcgw. \
47     - Use `any_workspace_path` to initialize the shell in the appropriate project directory. \
48     - If the user has mentioned a workspace or project root or any other file or folder use it to set `any_workspace_path`. \
49     - If user has mentioned any files use `initial_files_to_read` to read, use absolute paths only (~ allowed) \
50     - By default use mode \"wcgw\" \
51     - In \"code-writer\" mode, set the commands and globs which user asked to set, otherwise use 'all'. \
52     - Use type=\"first_call\" if it's the first call to this tool. \
53     - Use type=\"user_asked_mode_change\" if in a conversation user has asked to change mode. \
54     - Use type=\"reset_shell\" if in a conversation shell is not working after multiple tries. \
55     - Use type=\"user_asked_change_workspace\" if in a conversation user asked to change workspace";
56
57const BASH_COMMAND_DESCRIPTION: &str =
58    "- Execute a bash command. This is stateful (beware with subsequent calls). \
59     - Status of the command and the current working directory will always be returned at the end. \
60     - The first or the last line might be `(...truncated)` if the output is too long. \
61     - Always run `pwd` if you get any file or directory not found error to make sure you're not lost. \
62     - Do not run bg commands using \"&\", instead use this tool. \
63     - You must not use echo/cat to read/write files, use ReadFiles/FileWriteOrEdit \
64     - In order to check status of previous command, use `status_check` with empty command argument. \
65     - Only command is allowed to run at a time. You need to wait for any previous command to finish before running a new one. \
66     - 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. \
67     - Do not send Ctrl-c before checking for status till 10 minutes or whatever is appropriate for the program to finish. \
68     - Only run long running commands in background. Each background command is run in a new non-reusable shell. \
69     - On running a bg command you'll get a bg command id that you should use to get status or interact.";
70
71const READ_FILES_DESCRIPTION: &str =
72    "- Read full file content of one or more files. \
73     - Provide absolute paths only (~ allowed) \
74     - Only if the task requires line numbers understanding: \
75     - 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`";
76
77const FILE_WRITE_OR_EDIT_DESCRIPTION: &str =
78    "- Writes or edits a file based on the percentage of changes. \
79     - Use absolute path only (~ allowed). \
80     - First write down percentage of lines that need to be replaced in the file (between 0-100) in percentage_to_change \
81     - 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. \
82     - If percentage_to_change > 50, provide full file content in text_or_search_replace_blocks \
83     - If percentage_to_change <= 50, text_or_search_replace_blocks should be search/replace blocks. \
84     \
85     Instructions for editing files. \
86     # Example \
87     ## Input file \
88     ``` \
89     import numpy as np \
90     from impls import impl1, impl2 \
91     \
92     def hello(): \
93         \"print a greeting\" \
94         print(\"hello\") \
95     \
96     def call_hello(): \
97         \"call hello\" \
98         hello() \
99         print(\"Called\") \
100         impl1() \
101         hello() \
102         impl2() \
103     ``` \
104     ## Edit format on the input file \
105     ``` \
106     <<<<<<< SEARCH \
107     from impls import impl1, impl2 \
108     ======= \
109     from impls import impl1, impl2 \
110     from hello import hello as hello_renamed \
111     >>>>>>> REPLACE \
112     <<<<<<< SEARCH \
113     def hello(): \
114         \"print a greeting\" \
115         print(\"hello\") \
116     ======= \
117     >>>>>>> REPLACE \
118     <<<<<<< SEARCH \
119     def call_hello(): \
120         \"call hello\" \
121         hello() \
122     ======= \
123     def call_hello_renamed(): \
124         \"call hello renamed\" \
125         hello_renamed() \
126     >>>>>>> REPLACE \
127     <<<<<<< SEARCH \
128     impl1() \
129     hello() \
130     impl2() \
131     ======= \
132     impl1() \
133     hello_renamed() \
134     impl2() \
135     >>>>>>> REPLACE \
136     ``` \
137     # *SEARCH/REPLACE block* Rules: \
138     Every \"<<<<<<< SEARCH\" section must *EXACTLY MATCH* the existing file content, character for character, including all comments, docstrings, whitespaces, etc. \
139     Including multiple unique *SEARCH/REPLACE* blocks if needed. \
140     Include enough and only enough lines in each SEARCH section to uniquely match each set of lines that need to change. \
141     Keep *SEARCH/REPLACE* blocks concise. \
142     Break large *SEARCH/REPLACE* blocks into a series of smaller blocks that each change a small portion of the file. \
143     Include just the changing lines, and a few surrounding lines (0-3 lines) if needed for uniqueness. \
144     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. \
145     Preserve leading spaces and indentations in both SEARCH and REPLACE blocks.";
146
147const CONTEXT_SAVE_DESCRIPTION: &str =
148    "Saves provided description and file contents of all the relevant file paths or globs in a single text file. \
149     - Provide random 3 word unqiue id or whatever user provided. \
150     - Leave project path as empty string if no project path";
151
152static WINX_TOOLS: OnceLock<Vec<Tool>> = OnceLock::new();
153
154fn winx_tools() -> Vec<Tool> {
155    WINX_TOOLS.get_or_init(build_winx_tools).clone()
156}
157
158fn build_winx_tools() -> Vec<Tool> {
159    vec![
160        mcp_tool::<Initialize>(
161            "Initialize",
162            INITIALIZE_DESCRIPTION,
163            ToolAnnotations::new().read_only(true).open_world(false),
164        ),
165        mcp_tool::<BashCommand>(
166            "BashCommand",
167            BASH_COMMAND_DESCRIPTION,
168            ToolAnnotations::new().destructive(true).open_world(true),
169        ),
170        mcp_tool::<ReadFiles>(
171            "ReadFiles",
172            READ_FILES_DESCRIPTION,
173            ToolAnnotations::new().read_only(true).open_world(false),
174        ),
175        mcp_tool::<FileWriteOrEdit>(
176            "FileWriteOrEdit",
177            FILE_WRITE_OR_EDIT_DESCRIPTION,
178            ToolAnnotations::new().destructive(true).open_world(false),
179        ),
180        mcp_tool::<ContextSave>(
181            "ContextSave",
182            CONTEXT_SAVE_DESCRIPTION,
183            ToolAnnotations::new().destructive(false).open_world(false),
184        ),
185        mcp_tool::<ReadImage>(
186            "ReadImage",
187            "Read an image from the shell.",
188            ToolAnnotations::new().read_only(true).open_world(false),
189        ),
190    ]
191}
192
193/// Winx service with shared bash state
194#[derive(Clone)]
195pub struct WinxService {
196    /// Shared state for the bash shell environment (async-safe)
197    pub bash_state: SharedBashState,
198    /// Version information for the service
199    pub version: String,
200}
201
202impl Default for WinxService {
203    fn default() -> Self {
204        Self::new()
205    }
206}
207
208impl WinxService {
209    /// Create a new `WinxService` instance
210    pub fn new() -> Self {
211        info!("Creating new WinxService instance");
212        Self {
213            bash_state: Arc::new(Mutex::new(None)),
214            version: env!("CARGO_PKG_VERSION").to_string(),
215        }
216    }
217}
218
219/// `ServerHandler` implementation
220impl ServerHandler for WinxService {
221    fn get_info(&self) -> ServerInfo {
222        ServerInfo::new(
223            ServerCapabilities::builder()
224                .enable_tools()
225                .enable_resources()
226                .build(),
227        )
228        .with_server_info(
229            Implementation::new("winx-mcp-server", self.version.clone())
230                .with_title("Winx High-Performance MCP"),
231        )
232        .with_protocol_version(ProtocolVersion::V_2024_11_05)
233        .with_instructions(
234                "Winx is a high-performance Rust implementation of MCP tools for shell and file management."
235        )
236    }
237
238    async fn list_tools(
239        &self,
240        _request: Option<PaginatedRequestParams>,
241        _context: RequestContext<RoleServer>,
242    ) -> Result<ListToolsResult, McpError> {
243        Ok(ListToolsResult { tools: winx_tools(), next_cursor: None, meta: None })
244    }
245
246    async fn list_resources(
247        &self,
248        _param: Option<PaginatedRequestParams>,
249        _context: RequestContext<RoleServer>,
250    ) -> Result<ListResourcesResult, McpError> {
251        Ok(ListResourcesResult {
252            resources: vec![Annotated {
253                raw: RawResource {
254                    uri: "file://readme".into(),
255                    name: "README".into(),
256                    description: Some("Project README documentation".into()),
257                    mime_type: Some("text/markdown".into()),
258                    size: None,
259                    title: None,
260                    icons: None,
261                    meta: None,
262                },
263                annotations: None,
264            }],
265            next_cursor: None,
266            meta: None,
267        })
268    }
269
270    async fn read_resource(
271        &self,
272        param: ReadResourceRequestParams,
273        _context: RequestContext<RoleServer>,
274    ) -> Result<ReadResourceResult, McpError> {
275        let content = match param.uri.as_ref() {
276            "file://readme" => match tokio::fs::read_to_string("README.md").await {
277                Ok(content) => vec![ResourceContents::text(content, param.uri.clone())],
278                Err(_) => vec![ResourceContents::text(
279                    "README.md not found".to_string(),
280                    param.uri.clone(),
281                )],
282            },
283            _ => {
284                return Err(McpError::invalid_request(
285                    format!("Unknown resource URI: {}", param.uri),
286                    None,
287                ));
288            }
289        };
290
291        Ok(ReadResourceResult::new(content))
292    }
293
294    async fn call_tool(
295        &self,
296        param: CallToolRequestParams,
297        _context: RequestContext<RoleServer>,
298    ) -> Result<CallToolResult, McpError> {
299        let args_value = param.arguments.map(Value::Object);
300        match param.name.as_ref() {
301            "Initialize" => self.handle_initialize(args_value).await,
302            "BashCommand" => self.handle_bash_command(args_value).await,
303            "ReadFiles" => self.handle_read_files(args_value).await,
304            "FileWriteOrEdit" => self.handle_file_write_or_edit(args_value).await,
305            "ContextSave" => self.handle_context_save(args_value).await,
306            "ReadImage" => self.handle_read_image(args_value).await,
307            _ => Err(McpError::invalid_request(format!("Unknown tool: {}", param.name), None)),
308        }
309    }
310}
311
312impl WinxService {
313    async fn handle_initialize(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
314        let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
315        let initialize: Initialize = serde_json::from_value(args).map_err(|e| {
316            McpError::invalid_request(format!("Invalid Initialize parameters: {e}"), None)
317        })?;
318
319        match crate::tools::initialize::handle_tool_call(&self.bash_state, initialize).await {
320            Ok(result) => Ok(CallToolResult::success(vec![Content::text(result)])),
321            Err(e) => Err(McpError::internal_error(format!("Initialize failed: {e}"), None)),
322        }
323    }
324
325    async fn handle_bash_command(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
326        let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
327        let bash_command: BashCommand = serde_json::from_value(args).map_err(|e| {
328            McpError::invalid_request(format!("Invalid BashCommand parameters: {e}"), None)
329        })?;
330
331        match crate::tools::bash_command::handle_tool_call(&self.bash_state, bash_command).await {
332            Ok(output) => Ok(CallToolResult::success(vec![Content::text(output)])),
333            Err(e) => Err(McpError::internal_error(format!("BashCommand failed: {e}"), None)),
334        }
335    }
336
337    async fn handle_read_files(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
338        let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
339        let read_files: ReadFiles = serde_json::from_value(args).map_err(|e| {
340            McpError::invalid_request(format!("Invalid ReadFiles parameters: {e}"), None)
341        })?;
342
343        match crate::tools::read_files::handle_tool_call(&self.bash_state, read_files).await {
344            Ok(result) => Ok(CallToolResult::success(vec![Content::text(result)])),
345            Err(e) => Err(McpError::internal_error(format!("ReadFiles failed: {e}"), None)),
346        }
347    }
348
349    async fn handle_file_write_or_edit(
350        &self,
351        args: Option<Value>,
352    ) -> Result<CallToolResult, McpError> {
353        let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
354        let file_write_or_edit: FileWriteOrEdit = serde_json::from_value(args).map_err(|e| {
355            McpError::invalid_request(format!("Invalid FileWriteOrEdit parameters: {e}"), None)
356        })?;
357
358        match crate::tools::file_write_or_edit::handle_tool_call(
359            &self.bash_state,
360            file_write_or_edit,
361        )
362        .await
363        {
364            Ok(result) => Ok(CallToolResult::success(vec![Content::text(result)])),
365            Err(e) => Err(McpError::internal_error(format!("FileWriteOrEdit failed: {e}"), None)),
366        }
367    }
368
369    async fn handle_context_save(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
370        let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
371        let context_save: ContextSave = serde_json::from_value(args).map_err(|e| {
372            McpError::invalid_request(format!("Invalid ContextSave parameters: {e}"), None)
373        })?;
374
375        match crate::tools::context_save::handle_tool_call(&self.bash_state, context_save).await {
376            Ok(result) => Ok(CallToolResult::success(vec![Content::text(result)])),
377            Err(e) => Err(McpError::internal_error(format!("ContextSave failed: {e}"), None)),
378        }
379    }
380
381    async fn handle_read_image(&self, args: Option<Value>) -> Result<CallToolResult, McpError> {
382        let args = args.ok_or_else(|| McpError::invalid_request("Missing arguments", None))?;
383        let read_image: ReadImage = serde_json::from_value(args).map_err(|e| {
384            McpError::invalid_request(format!("Invalid ReadImage parameters: {e}"), None)
385        })?;
386
387        match crate::tools::read_image::handle_tool_call(&self.bash_state, read_image).await {
388            Ok((mime_type, base64_data)) => {
389                let result_text = format!("MIME: {mime_type}\nData: {base64_data}");
390                Ok(CallToolResult::success(vec![Content::text(result_text)]))
391            }
392            Err(e) => Err(McpError::internal_error(format!("ReadImage failed: {e}"), None)),
393        }
394    }
395}
396
397/// Create and start the Winx MCP server
398pub async fn start_winx_server() -> Result<(), Box<dyn std::error::Error>> {
399    info!("Starting Winx MCP Server");
400    let service = WinxService::new();
401    let server = service.serve(stdio()).await?;
402    server.waiting().await?;
403    Ok(())
404}
winx_code_agent/server.rs

winx_code_agent/
server.rs
