winx_code_agent/server/
handler.rs

1use rmcp::{
2    model::{
3        CacheScope, CallToolRequestParams, CallToolResponse, CancelTaskParams,
4        GetPromptRequestParams, GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskResult,
5        Icon, Implementation, ListPromptsResult, ListResourcesResult, ListToolsResult,
6        PaginatedRequestParams, PromptMessage, ProtocolVersion, ReadResourceRequestParams,
7        ReadResourceResponse, ReadResourceResult, Resource, ResourceContents, Role,
8        ServerCapabilities, ServerInfo, 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/// Emit the cache fields required by MCP 2026-07-28 without changing the wire
80/// shape for clients that negotiated an older protocol revision.
81fn cache_hints_for_protocol(
82    protocol_version: Option<ProtocolVersion>,
83    cache_scope: CacheScope,
84) -> (Option<u64>, Option<CacheScope>) {
85    if protocol_version.is_some_and(|version| version >= ProtocolVersion::V_2026_07_28) {
86        (Some(0), Some(cache_scope))
87    } else {
88        (None, None)
89    }
90}
91
92/// Classify a `BashCommand` call by action kind (`command`, `status_check`,
93/// `send_text`, ...) without touching the command text itself. Mirrors the
94/// lenient forms accepted by the `BashCommand` deserializer: a typed
95/// `action_json` object, a JSON-encoded string, a bare command string, or
96/// legacy shorthand keys at the top level.
97fn bash_action(arguments: Option<&rmcp::model::JsonObject>) -> String {
98    use serde_json::Value;
99    const KINDS: [&str; 7] = [
100        "command",
101        "status_check",
102        "send_text",
103        "send_specials",
104        "send_ascii",
105        "screen",
106        "wait_for_turn",
107    ];
108    let Some(arguments) = arguments else {
109        return String::new();
110    };
111    let parsed;
112    let action = match arguments.get("action_json") {
113        Some(Value::Object(object)) => object,
114        Some(Value::String(text)) => {
115            match serde_json::from_str::<Value>(&text.replace('\n', " ")) {
116                Ok(Value::Object(object)) => {
117                    parsed = object;
118                    &parsed
119                }
120                // A non-object string is treated as a bare command downstream.
121                _ => return "command".to_string(),
122            }
123        }
124        // Legacy shorthand: the argument object itself is the action.
125        _ => arguments,
126    };
127    if let Some(kind) = action.get("type").and_then(Value::as_str) {
128        return kind.to_string();
129    }
130    KINDS.iter().find(|kind| action.contains_key(**kind)).map_or("?", |kind| kind).to_string()
131}
132
133impl ServerHandler for WinxService {
134    fn get_info(&self) -> ServerInfo {
135        ServerInfo::new(
136            ServerCapabilities::builder()
137                .enable_tools()
138                .enable_resources()
139                .enable_prompts()
140                .enable_tasks()
141                .build(),
142        )
143        .with_server_info(
144            Implementation::new("winx-mcp-server", self.version.clone())
145                .with_title("Winx High-Performance MCP")
146                .with_icons(vec![Icon::new(server_icon_data_uri())
147                    .with_mime_type("image/png")
148                    .with_sizes(vec!["96x96".to_string()])]),
149        )
150        .with_protocol_version(ProtocolVersion::V_2026_07_28)
151        .with_instructions(
152            "Winx is a high-performance Rust implementation of MCP tools for shell and file management.",
153        )
154    }
155
156    async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
157        self.initialize_from_client_roots(context).await;
158    }
159
160    async fn on_roots_list_changed(&self, context: NotificationContext<RoleServer>) {
161        // A changed Roots list is only used as a retry if Winx is still
162        // uninitialized. Never switch an active shell's workspace implicitly.
163        self.initialize_from_client_roots(context).await;
164    }
165
166    fn get_tool(&self, name: &str) -> Option<Tool> {
167        winx_tools().into_iter().find(|tool| tool.name == name)
168    }
169
170    async fn get_task(
171        &self,
172        request: GetTaskParams,
173        context: RequestContext<RoleServer>,
174    ) -> Result<GetTaskResult, McpError> {
175        let principal = principal_from_context(&context);
176        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
177            return Err(McpError::invalid_request(
178                format!("Unknown or expired task: {}", request.task_id),
179                None,
180            ));
181        }
182        let mut tasks = self.tasks.lock().await;
183        let entry = tasks.get(&request.task_id).ok_or_else(|| {
184            McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
185        })?;
186        Ok(GetTaskResult::new(entry.detailed()))
187    }
188
189    async fn update_task(
190        &self,
191        request: UpdateTaskParams,
192        context: RequestContext<RoleServer>,
193    ) -> Result<(), McpError> {
194        let principal = principal_from_context(&context);
195        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
196            return Err(McpError::invalid_request(
197                format!("Unknown or expired task: {}", request.task_id),
198                None,
199            ));
200        }
201        let mut tasks = self.tasks.lock().await;
202        tasks.get(&request.task_id).ok_or_else(|| {
203            McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
204        })?;
205        Ok(())
206    }
207
208    async fn cancel_task(
209        &self,
210        request: CancelTaskParams,
211        context: RequestContext<RoleServer>,
212    ) -> Result<(), McpError> {
213        let principal = principal_from_context(&context);
214        if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
215            return Err(McpError::invalid_request(
216                format!("Unknown or expired task: {}", request.task_id),
217                None,
218            ));
219        }
220        let (abort_handle, thread_id) = {
221            let mut tasks = self.tasks.lock().await;
222            let entry = tasks.get_mut(&request.task_id).ok_or_else(|| {
223                McpError::invalid_request(
224                    format!("Unknown or expired task: {}", request.task_id),
225                    None,
226                )
227            })?;
228            if entry.task.status != TaskStatus::Working
229                && entry.task.status != TaskStatus::InputRequired
230            {
231                return Err(McpError::invalid_request(
232                    format!("Task {} is already terminal", request.task_id),
233                    None,
234                ));
235            }
236            let abort_handle = entry.abort_handle.take();
237            let thread_id = entry.thread_id.clone();
238            entry.finish(TaskStatus::Cancelled, Some("Cancelled by client".to_string()), None);
239            (abort_handle, thread_id)
240        };
241        if let Some(abort_handle) = abort_handle {
242            abort_handle.abort();
243        }
244        self.interrupt_task_thread(&thread_id).await;
245        Ok(())
246    }
247
248    async fn list_tools(
249        &self,
250        _request: Option<PaginatedRequestParams>,
251        ctx: RequestContext<RoleServer>,
252    ) -> Result<ListToolsResult, McpError> {
253        let (ttl_ms, cache_scope) =
254            cache_hints_for_protocol(ctx.protocol_version(), CacheScope::Public);
255        let mut result = ListToolsResult::with_all_items(winx_tools());
256        result.ttl_ms = ttl_ms;
257        result.cache_scope = cache_scope;
258        Ok(result)
259    }
260
261    async fn list_resources(
262        &self,
263        _request: Option<PaginatedRequestParams>,
264        ctx: RequestContext<RoleServer>,
265    ) -> Result<ListResourcesResult, McpError> {
266        let (ttl_ms, cache_scope) =
267            cache_hints_for_protocol(ctx.protocol_version(), CacheScope::Public);
268        let mut result =
269            ListResourcesResult::with_all_items(vec![Resource::new("file://readme", "README")
270                .with_description("Project README documentation")
271                .with_mime_type("text/markdown")]);
272        result.ttl_ms = ttl_ms;
273        result.cache_scope = cache_scope;
274        Ok(result)
275    }
276
277    async fn list_prompts(
278        &self,
279        _request: Option<PaginatedRequestParams>,
280        ctx: RequestContext<RoleServer>,
281    ) -> Result<ListPromptsResult, McpError> {
282        let (ttl_ms, cache_scope) =
283            cache_hints_for_protocol(ctx.protocol_version(), CacheScope::Public);
284        let mut result = ListPromptsResult::with_all_items(winx_prompts());
285        result.ttl_ms = ttl_ms;
286        result.cache_scope = cache_scope;
287        Ok(result)
288    }
289
290    async fn get_prompt(
291        &self,
292        request: GetPromptRequestParams,
293        context: RequestContext<RoleServer>,
294    ) -> Result<GetPromptResponse, McpError> {
295        if request.name != "KnowledgeTransfer" {
296            return Err(McpError::invalid_request(
297                format!("Unknown prompt: {}", request.name),
298                None,
299            ));
300        }
301
302        let principal = principal_from_context(&context);
303        let session_prefix = principal.as_ref().map(crate::config::HttpPrincipal::session_prefix);
304        let text = crate::utils::redact::redact(
305            &self.knowledge_transfer_prompt_text(session_prefix.as_deref()).await,
306        )
307        .into_owned();
308        Ok(GetPromptResult::new(vec![PromptMessage::new_text(Role::User, text)])
309            .with_description("Knowledge transfer handoff prompt")
310            .into())
311    }
312
313    async fn read_resource(
314        &self,
315        request: ReadResourceRequestParams,
316        ctx: RequestContext<RoleServer>,
317    ) -> Result<ReadResourceResponse, McpError> {
318        let content = match request.uri.as_ref() {
319            "file://readme" => match tokio::fs::read_to_string("README.md").await {
320                Ok(content) => vec![ResourceContents::text(
321                    crate::utils::redact::redact(&content).into_owned(),
322                    request.uri.clone(),
323                )],
324                Err(_) => vec![ResourceContents::text(
325                    "README.md not found".to_string(),
326                    request.uri.clone(),
327                )],
328            },
329            _ => {
330                return Err(McpError::invalid_request(
331                    format!("Unknown resource URI: {}", request.uri),
332                    None,
333                ));
334            }
335        };
336        let (ttl_ms, cache_scope) =
337            cache_hints_for_protocol(ctx.protocol_version(), CacheScope::Private);
338        let mut result = ReadResourceResult::new(content);
339        result.ttl_ms = ttl_ms;
340        result.cache_scope = cache_scope;
341        Ok(result.into())
342    }
343
344    async fn call_tool(
345        &self,
346        request: CallToolRequestParams,
347        context: RequestContext<RoleServer>,
348    ) -> Result<CallToolResponse, McpError> {
349        let principal = principal_from_context(&context);
350        let affinity = session_affinity_from_context(&context);
351        let (request, scope) =
352            scope_tool_request(request, principal, affinity).map_err(|error| {
353                McpError::invalid_request(format!("Cannot scope remote request: {error}"), None)
354            })?;
355        let usage = UsageEvent::start(&request, &scope);
356        if context.client_capabilities().is_some_and(|capabilities| capabilities.supports_tasks())
357            && Self::bash_task_is_eligible(&request)
358        {
359            return match self.enqueue_bash_task(request, scope).await {
360                Ok(task) => {
361                    usage.emit("task");
362                    Ok(CallToolResponse::Task(task))
363                }
364                Err(error) => {
365                    usage.emit("error");
366                    Err(error)
367                }
368            };
369        }
370
371        match self.execute_tool_call(request).await {
372            Ok(mut result) => {
373                scope.unscope_result(&mut result);
374                usage.emit("ok");
375                Ok(result.into())
376            }
377            Err(mut error) => {
378                scope.unscope_error(&mut error);
379                usage.emit("error");
380                Err(error)
381            }
382        }
383    }
384}
385
386#[cfg(test)]
387mod tests {
388    use super::*;
389
390    #[test]
391    fn cache_hints_follow_negotiated_protocol_version() {
392        assert_eq!(
393            cache_hints_for_protocol(Some(ProtocolVersion::V_2026_07_28), CacheScope::Private,),
394            (Some(0), Some(CacheScope::Private))
395        );
396        assert_eq!(
397            cache_hints_for_protocol(Some(ProtocolVersion::V_2026_07_28), CacheScope::Public),
398            (Some(0), Some(CacheScope::Public))
399        );
400        assert_eq!(
401            cache_hints_for_protocol(Some(ProtocolVersion::V_2025_11_25), CacheScope::Private,),
402            (None, None)
403        );
404        assert_eq!(cache_hints_for_protocol(None, CacheScope::Private), (None, None));
405    }
406}
winx_code_agent/server/handler.rs

winx_code_agent/server/
handler.rs
