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    principal: String,
26    thread_id: String,
27    started: std::time::Instant,
28}
29
30impl UsageEvent {
31    fn start(request: &CallToolRequestParams, scope: &RequestScope) -> Self {
32        let thread_id = request
33            .arguments
34            .as_ref()
35            .and_then(|arguments| arguments.get("thread_id"))
36            .and_then(serde_json::Value::as_str)
37            .map(|thread_id| scope.unscope_text(thread_id))
38            .unwrap_or_default();
39        Self {
40            tool: request.name.to_string(),
41            principal: scope
42                .principal()
43                .map_or_else(|| "local".to_string(), |principal| principal.name().to_string()),
44            thread_id,
45            started: std::time::Instant::now(),
46        }
47    }
48
49    fn emit(&self, outcome: &str) {
50        tracing::info!(
51            target: "winx::usage",
52            event = "tool_call",
53            tool = %self.tool,
54            principal = %self.principal,
55            thread_id = %self.thread_id,
56            duration_ms = u64::try_from(self.started.elapsed().as_millis()).unwrap_or(u64::MAX),
57            outcome,
58            "tool call"
59        );
60    }
61}
62
63impl ServerHandler for WinxService {
64    fn get_info(&self) -> ServerInfo {
65        ServerInfo::new(
66            ServerCapabilities::builder()
67                .enable_tools()
68                .enable_resources()
69                .enable_prompts()
70                .enable_tasks()
71                .build(),
72        )
73        .with_server_info(
74            Implementation::new("winx-mcp-server", self.version.clone())
75                .with_title("Winx High-Performance MCP")
76                .with_icons(vec![Icon::new(server_icon_data_uri())
77                    .with_mime_type("image/png")
78                    .with_sizes(vec!["96x96".to_string()])]),
79        )
80        .with_protocol_version(ProtocolVersion::V_2026_07_28)
81        .with_instructions(
82            "Winx is a high-performance Rust implementation of MCP tools for shell and file management.",
83        )
84    }
85
86    async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
87        self.initialize_from_client_roots(context).await;
88    }
89
90    async fn on_roots_list_changed(&self, context: NotificationContext<RoleServer>) {
91        // A changed Roots list is only used as a retry if Winx is still
92        // uninitialized. Never switch an active shell's workspace implicitly.
93        self.initialize_from_client_roots(context).await;
94    }
95
96    fn get_tool(&self, name: &str) -> Option<Tool> {
97        winx_tools().into_iter().find(|tool| tool.name == name)
98    }
99
100    async fn get_task(
101        &self,
102        request: GetTaskParams,
103        context: RequestContext<RoleServer>,
104    ) -> Result<GetTaskResult, McpError> {
105        let principal = principal_from_context(&context);
106        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
107            return Err(McpError::invalid_request(
108                format!("Unknown or expired task: {}", request.task_id),
109                None,
110            ));
111        }
112        let mut tasks = self.tasks.lock().await;
113        let entry = tasks.get(&request.task_id).ok_or_else(|| {
114            McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
115        })?;
116        Ok(GetTaskResult::new(entry.detailed()))
117    }
118
119    async fn update_task(
120        &self,
121        request: UpdateTaskParams,
122        context: RequestContext<RoleServer>,
123    ) -> Result<(), McpError> {
124        let principal = principal_from_context(&context);
125        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
126            return Err(McpError::invalid_request(
127                format!("Unknown or expired task: {}", request.task_id),
128                None,
129            ));
130        }
131        let mut tasks = self.tasks.lock().await;
132        tasks.get(&request.task_id).ok_or_else(|| {
133            McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
134        })?;
135        Ok(())
136    }
137
138    async fn cancel_task(
139        &self,
140        request: CancelTaskParams,
141        context: RequestContext<RoleServer>,
142    ) -> Result<(), McpError> {
143        let principal = principal_from_context(&context);
144        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
145            return Err(McpError::invalid_request(
146                format!("Unknown or expired task: {}", request.task_id),
147                None,
148            ));
149        }
150        let (abort_handle, thread_id) = {
151            let mut tasks = self.tasks.lock().await;
152            let entry = tasks.get_mut(&request.task_id).ok_or_else(|| {
153                McpError::invalid_request(
154                    format!("Unknown or expired task: {}", request.task_id),
155                    None,
156                )
157            })?;
158            if entry.task.status != TaskStatus::Working
159                && entry.task.status != TaskStatus::InputRequired
160            {
161                return Err(McpError::invalid_request(
162                    format!("Task {} is already terminal", request.task_id),
163                    None,
164                ));
165            }
166            let abort_handle = entry.abort_handle.take();
167            let thread_id = entry.thread_id.clone();
168            entry.finish(TaskStatus::Cancelled, Some("Cancelled by client".to_string()), None);
169            (abort_handle, thread_id)
170        };
171        if let Some(abort_handle) = abort_handle {
172            abort_handle.abort();
173        }
174        self.interrupt_task_thread(&thread_id).await;
175        Ok(())
176    }
177
178    async fn list_tools(
179        &self,
180        _request: Option<PaginatedRequestParams>,
181        _context: RequestContext<RoleServer>,
182    ) -> Result<ListToolsResult, McpError> {
183        Ok(ListToolsResult::with_all_items(winx_tools()))
184    }
185
186    async fn list_resources(
187        &self,
188        _request: Option<PaginatedRequestParams>,
189        _context: RequestContext<RoleServer>,
190    ) -> Result<ListResourcesResult, McpError> {
191        Ok(ListResourcesResult::with_all_items(vec![Resource::new("file://readme", "README")
192            .with_description("Project README documentation")
193            .with_mime_type("text/markdown")]))
194    }
195
196    async fn list_prompts(
197        &self,
198        _request: Option<PaginatedRequestParams>,
199        _context: RequestContext<RoleServer>,
200    ) -> Result<ListPromptsResult, McpError> {
201        Ok(ListPromptsResult::with_all_items(winx_prompts()))
202    }
203
204    async fn get_prompt(
205        &self,
206        request: GetPromptRequestParams,
207        context: RequestContext<RoleServer>,
208    ) -> Result<GetPromptResponse, McpError> {
209        if request.name != "KnowledgeTransfer" {
210            return Err(McpError::invalid_request(
211                format!("Unknown prompt: {}", request.name),
212                None,
213            ));
214        }
215
216        let principal = principal_from_context(&context);
217        let session_prefix = principal.as_ref().map(crate::config::HttpPrincipal::session_prefix);
218        let text = crate::utils::redact::redact(
219            &self.knowledge_transfer_prompt_text(session_prefix.as_deref()).await,
220        )
221        .into_owned();
222        Ok(GetPromptResult::new(vec![PromptMessage::new_text(Role::User, text)])
223            .with_description("Knowledge transfer handoff prompt")
224            .into())
225    }
226
227    async fn read_resource(
228        &self,
229        request: ReadResourceRequestParams,
230        _context: RequestContext<RoleServer>,
231    ) -> Result<ReadResourceResponse, McpError> {
232        let content = match request.uri.as_ref() {
233            "file://readme" => match tokio::fs::read_to_string("README.md").await {
234                Ok(content) => vec![ResourceContents::text(
235                    crate::utils::redact::redact(&content).into_owned(),
236                    request.uri.clone(),
237                )],
238                Err(_) => vec![ResourceContents::text(
239                    "README.md not found".to_string(),
240                    request.uri.clone(),
241                )],
242            },
243            _ => {
244                return Err(McpError::invalid_request(
245                    format!("Unknown resource URI: {}", request.uri),
246                    None,
247                ));
248            }
249        };
250        Ok(ReadResourceResult::new(content).into())
251    }
252
253    async fn call_tool(
254        &self,
255        request: CallToolRequestParams,
256        context: RequestContext<RoleServer>,
257    ) -> Result<CallToolResponse, McpError> {
258        let principal = principal_from_context(&context);
259        let affinity = session_affinity_from_context(&context);
260        let (request, scope) =
261            scope_tool_request(request, principal, affinity).map_err(|error| {
262                McpError::invalid_request(format!("Cannot scope remote request: {error}"), None)
263            })?;
264        let usage = UsageEvent::start(&request, &scope);
265        if context.client_capabilities().is_some_and(|capabilities| capabilities.supports_tasks())
266            && Self::bash_task_is_eligible(&request)
267        {
268            return match self.enqueue_bash_task(request, scope).await {
269                Ok(task) => {
270                    usage.emit("task");
271                    Ok(CallToolResponse::Task(task))
272                }
273                Err(error) => {
274                    usage.emit("error");
275                    Err(error)
276                }
277            };
278        }
279
280        match self.execute_tool_call(request).await {
281            Ok(mut result) => {
282                scope.unscope_result(&mut result);
283                usage.emit("ok");
284                Ok(result.into())
285            }
286            Err(mut error) => {
287                scope.unscope_error(&mut error);
288                usage.emit("error");
289                Err(error)
290            }
291        }
292    }
293}
winx_code_agent/server/handler.rs

winx_code_agent/server/
handler.rs
