winx_code_agent/server/
handler.rs

1use rmcp::{
2    model::{
3        CallToolRequestParams, CallToolResponse, CancelTaskParams, GetPromptRequestParams,
4        GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskResult, Icon, Implementation,
5        ListPromptsResult, ListResourcesResult, ListToolsResult, PaginatedRequestParams,
6        PromptMessage, ProtocolVersion, ReadResourceRequestParams, ReadResourceResponse,
7        ReadResourceResult, Resource, ResourceContents, Role, ServerCapabilities, ServerInfo,
8        TaskStatus, Tool, UpdateTaskParams,
9    },
10    service::{NotificationContext, RequestContext, RoleServer},
11    ErrorData as McpError, ServerHandler,
12};
13
14use super::catalog::{server_icon_data_uri, winx_prompts, winx_tools};
15use super::principal::{
16    principal_from_context, scope_tool_request, session_affinity_from_context,
17    task_belongs_to_principal, RequestScope,
18};
19use super::WinxService;
20
21/// One structured `winx::usage` event per tool call. Only metadata is logged —
22/// never tool arguments, command text, or file contents (secrets/PII).
23struct UsageEvent {
24    tool: String,
25    action: String,
26    ws: String,
27    principal: String,
28    thread_id: String,
29    started: std::time::Instant,
30}
31
32impl UsageEvent {
33    fn start(request: &CallToolRequestParams, scope: &RequestScope) -> Self {
34        let arguments = request.arguments.as_ref();
35        let thread_id = arguments
36            .and_then(|arguments| arguments.get("thread_id"))
37            .and_then(serde_json::Value::as_str)
38            .map(|thread_id| scope.unscope_text(thread_id))
39            .unwrap_or_default();
40        let action =
41            if request.name == "BashCommand" { bash_action(arguments) } else { String::new() };
42        let ws = if request.name == "Initialize" {
43            arguments
44                .and_then(|arguments| arguments.get("any_workspace_path"))
45                .and_then(serde_json::Value::as_str)
46                .unwrap_or_default()
47                .to_string()
48        } else {
49            String::new()
50        };
51        Self {
52            tool: request.name.to_string(),
53            action,
54            ws,
55            principal: scope
56                .principal()
57                .map_or_else(|| "local".to_string(), |principal| principal.name().to_string()),
58            thread_id,
59            started: std::time::Instant::now(),
60        }
61    }
62
63    fn emit(&self, outcome: &str) {
64        tracing::info!(
65            target: "winx::usage",
66            event = "tool_call",
67            tool = %self.tool,
68            action = %self.action,
69            ws = %self.ws,
70            principal = %self.principal,
71            thread_id = %self.thread_id,
72            duration_ms = u64::try_from(self.started.elapsed().as_millis()).unwrap_or(u64::MAX),
73            outcome,
74            "tool call"
75        );
76    }
77}
78
79/// Classify a `BashCommand` call by action kind (`command`, `status_check`,
80/// `send_text`, ...) without touching the command text itself. Mirrors the
81/// lenient forms accepted by the `BashCommand` deserializer: a typed
82/// `action_json` object, a JSON-encoded string, a bare command string, or
83/// legacy shorthand keys at the top level.
84fn bash_action(arguments: Option<&rmcp::model::JsonObject>) -> String {
85    use serde_json::Value;
86    const KINDS: [&str; 7] = [
87        "command",
88        "status_check",
89        "send_text",
90        "send_specials",
91        "send_ascii",
92        "screen",
93        "wait_for_turn",
94    ];
95    let Some(arguments) = arguments else {
96        return String::new();
97    };
98    let parsed;
99    let action = match arguments.get("action_json") {
100        Some(Value::Object(object)) => object,
101        Some(Value::String(text)) => {
102            match serde_json::from_str::<Value>(&text.replace('\n', " ")) {
103                Ok(Value::Object(object)) => {
104                    parsed = object;
105                    &parsed
106                }
107                // A non-object string is treated as a bare command downstream.
108                _ => return "command".to_string(),
109            }
110        }
111        // Legacy shorthand: the argument object itself is the action.
112        _ => arguments,
113    };
114    if let Some(kind) = action.get("type").and_then(Value::as_str) {
115        return kind.to_string();
116    }
117    KINDS.iter().find(|kind| action.contains_key(**kind)).map_or("?", |kind| kind).to_string()
118}
119
120impl ServerHandler for WinxService {
121    fn get_info(&self) -> ServerInfo {
122        ServerInfo::new(
123            ServerCapabilities::builder()
124                .enable_tools()
125                .enable_resources()
126                .enable_prompts()
127                .enable_tasks()
128                .build(),
129        )
130        .with_server_info(
131            Implementation::new("winx-mcp-server", self.version.clone())
132                .with_title("Winx High-Performance MCP")
133                .with_icons(vec![Icon::new(server_icon_data_uri())
134                    .with_mime_type("image/png")
135                    .with_sizes(vec!["96x96".to_string()])]),
136        )
137        .with_protocol_version(ProtocolVersion::V_2026_07_28)
138        .with_instructions(
139            "Winx is a high-performance Rust implementation of MCP tools for shell and file management.",
140        )
141    }
142
143    async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
144        self.initialize_from_client_roots(context).await;
145    }
146
147    async fn on_roots_list_changed(&self, context: NotificationContext<RoleServer>) {
148        // A changed Roots list is only used as a retry if Winx is still
149        // uninitialized. Never switch an active shell's workspace implicitly.
150        self.initialize_from_client_roots(context).await;
151    }
152
153    fn get_tool(&self, name: &str) -> Option<Tool> {
154        winx_tools().into_iter().find(|tool| tool.name == name)
155    }
156
157    async fn get_task(
158        &self,
159        request: GetTaskParams,
160        context: RequestContext<RoleServer>,
161    ) -> Result<GetTaskResult, McpError> {
162        let principal = principal_from_context(&context);
163        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
164            return Err(McpError::invalid_request(
165                format!("Unknown or expired task: {}", request.task_id),
166                None,
167            ));
168        }
169        let mut tasks = self.tasks.lock().await;
170        let entry = tasks.get(&request.task_id).ok_or_else(|| {
171            McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
172        })?;
173        Ok(GetTaskResult::new(entry.detailed()))
174    }
175
176    async fn update_task(
177        &self,
178        request: UpdateTaskParams,
179        context: RequestContext<RoleServer>,
180    ) -> Result<(), McpError> {
181        let principal = principal_from_context(&context);
182        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
183            return Err(McpError::invalid_request(
184                format!("Unknown or expired task: {}", request.task_id),
185                None,
186            ));
187        }
188        let mut tasks = self.tasks.lock().await;
189        tasks.get(&request.task_id).ok_or_else(|| {
190            McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
191        })?;
192        Ok(())
193    }
194
195    async fn cancel_task(
196        &self,
197        request: CancelTaskParams,
198        context: RequestContext<RoleServer>,
199    ) -> Result<(), McpError> {
200        let principal = principal_from_context(&context);
201        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
202            return Err(McpError::invalid_request(
203                format!("Unknown or expired task: {}", request.task_id),
204                None,
205            ));
206        }
207        let (abort_handle, thread_id) = {
208            let mut tasks = self.tasks.lock().await;
209            let entry = tasks.get_mut(&request.task_id).ok_or_else(|| {
210                McpError::invalid_request(
211                    format!("Unknown or expired task: {}", request.task_id),
212                    None,
213                )
214            })?;
215            if entry.task.status != TaskStatus::Working
216                && entry.task.status != TaskStatus::InputRequired
217            {
218                return Err(McpError::invalid_request(
219                    format!("Task {} is already terminal", request.task_id),
220                    None,
221                ));
222            }
223            let abort_handle = entry.abort_handle.take();
224            let thread_id = entry.thread_id.clone();
225            entry.finish(TaskStatus::Cancelled, Some("Cancelled by client".to_string()), None);
226            (abort_handle, thread_id)
227        };
228        if let Some(abort_handle) = abort_handle {
229            abort_handle.abort();
230        }
231        self.interrupt_task_thread(&thread_id).await;
232        Ok(())
233    }
234
235    async fn list_tools(
236        &self,
237        _request: Option<PaginatedRequestParams>,
238        _context: RequestContext<RoleServer>,
239    ) -> Result<ListToolsResult, McpError> {
240        Ok(ListToolsResult::with_all_items(winx_tools()))
241    }
242
243    async fn list_resources(
244        &self,
245        _request: Option<PaginatedRequestParams>,
246        _context: RequestContext<RoleServer>,
247    ) -> Result<ListResourcesResult, McpError> {
248        Ok(ListResourcesResult::with_all_items(vec![Resource::new("file://readme", "README")
249            .with_description("Project README documentation")
250            .with_mime_type("text/markdown")]))
251    }
252
253    async fn list_prompts(
254        &self,
255        _request: Option<PaginatedRequestParams>,
256        _context: RequestContext<RoleServer>,
257    ) -> Result<ListPromptsResult, McpError> {
258        Ok(ListPromptsResult::with_all_items(winx_prompts()))
259    }
260
261    async fn get_prompt(
262        &self,
263        request: GetPromptRequestParams,
264        context: RequestContext<RoleServer>,
265    ) -> Result<GetPromptResponse, McpError> {
266        if request.name != "KnowledgeTransfer" {
267            return Err(McpError::invalid_request(
268                format!("Unknown prompt: {}", request.name),
269                None,
270            ));
271        }
272
273        let principal = principal_from_context(&context);
274        let session_prefix = principal.as_ref().map(crate::config::HttpPrincipal::session_prefix);
275        let text = crate::utils::redact::redact(
276            &self.knowledge_transfer_prompt_text(session_prefix.as_deref()).await,
277        )
278        .into_owned();
279        Ok(GetPromptResult::new(vec![PromptMessage::new_text(Role::User, text)])
280            .with_description("Knowledge transfer handoff prompt")
281            .into())
282    }
283
284    async fn read_resource(
285        &self,
286        request: ReadResourceRequestParams,
287        _context: RequestContext<RoleServer>,
288    ) -> Result<ReadResourceResponse, McpError> {
289        let content = match request.uri.as_ref() {
290            "file://readme" => match tokio::fs::read_to_string("README.md").await {
291                Ok(content) => vec![ResourceContents::text(
292                    crate::utils::redact::redact(&content).into_owned(),
293                    request.uri.clone(),
294                )],
295                Err(_) => vec![ResourceContents::text(
296                    "README.md not found".to_string(),
297                    request.uri.clone(),
298                )],
299            },
300            _ => {
301                return Err(McpError::invalid_request(
302                    format!("Unknown resource URI: {}", request.uri),
303                    None,
304                ));
305            }
306        };
307        Ok(ReadResourceResult::new(content).into())
308    }
309
310    async fn call_tool(
311        &self,
312        request: CallToolRequestParams,
313        context: RequestContext<RoleServer>,
314    ) -> Result<CallToolResponse, McpError> {
315        let principal = principal_from_context(&context);
316        let affinity = session_affinity_from_context(&context);
317        let (request, scope) =
318            scope_tool_request(request, principal, affinity).map_err(|error| {
319                McpError::invalid_request(format!("Cannot scope remote request: {error}"), None)
320            })?;
321        let usage = UsageEvent::start(&request, &scope);
322        if context.client_capabilities().is_some_and(|capabilities| capabilities.supports_tasks())
323            && Self::bash_task_is_eligible(&request)
324        {
325            return match self.enqueue_bash_task(request, scope).await {
326                Ok(task) => {
327                    usage.emit("task");
328                    Ok(CallToolResponse::Task(task))
329                }
330                Err(error) => {
331                    usage.emit("error");
332                    Err(error)
333                }
334            };
335        }
336
337        match self.execute_tool_call(request).await {
338            Ok(mut result) => {
339                scope.unscope_result(&mut result);
340                usage.emit("ok");
341                Ok(result.into())
342            }
343            Err(mut error) => {
344                scope.unscope_error(&mut error);
345                usage.emit("error");
346                Err(error)
347            }
348        }
349    }
350}
winx_code_agent/server/handler.rs

winx_code_agent/server/
handler.rs
