1use std::fmt::Write as _;
2
3use axum::http::request::Parts;
4use rmcp::{
5 model::{
6 CacheScope, CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams,
7 GetPromptRequestParams, GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskResult,
8 Icon, Implementation, ListPromptsResult, ListResourcesResult, ListToolsResult,
9 PaginatedRequestParams, PromptMessage, ProtocolVersion, ReadResourceRequestParams,
10 ReadResourceResponse, ReadResourceResult, Resource, ResourceContents, Role,
11 ServerCapabilities, ServerInfo, TaskStatus, Tool, UpdateTaskParams,
12 },
13 service::{NotificationContext, RequestContext, RoleServer},
14 ErrorData as McpError, ServerHandler,
15};
16use sha2::{Digest, Sha256};
17
18use super::catalog::{server_icon_data_uri, winx_prompts, winx_tools, winx_tools_for_policy};
19use super::outcomes;
20use super::principal::{
21 conversation_identity_from_context, principal_from_context, scope_tool_request,
22 session_affinity_from_context, task_belongs_to_principal, RequestScope,
23};
24use super::WinxService;
25use crate::runtime::ShellActionOptions;
26use crate::types::{BashCommand, BashWaitPolicy};
27
28pub(crate) const COMPACT_BASH_OUTPUT_EXTENSION: &str = "io.winx/compact-bash-output";
29const ADAPTIVE_TASK_INLINE_SECONDS: f32 = 2.0;
30const RETURN_EARLY_DEFAULT_SECONDS: f32 = 0.25;
31const RETURN_EARLY_MAX_SECONDS: f32 = 5.0;
32const SYNCHRONOUS_MAX_SECONDS: f32 = 60.0;
33
34#[derive(Clone, Copy, Debug, Eq, PartialEq)]
35enum BashTaskRoute {
36 Synchronous,
37 Adaptive,
38 Immediate,
39}
40
41fn bash_task_route(
42 client_tasks: bool,
43 generation_bound_runtime: bool,
44 eligible: bool,
45 policy: BashWaitPolicy,
46) -> BashTaskRoute {
47 if !client_tasks || !generation_bound_runtime || !eligible {
48 return BashTaskRoute::Synchronous;
49 }
50 match policy {
51 BashWaitPolicy::Adaptive => BashTaskRoute::Adaptive,
52 BashWaitPolicy::UntilComplete => BashTaskRoute::Immediate,
53 BashWaitPolicy::ReturnEarly => BashTaskRoute::Synchronous,
54 }
55}
56
57fn adaptive_action_options(compact_output: bool) -> ShellActionOptions {
58 ShellActionOptions {
59 compact_output,
60 require_generation_binding: true,
65 ..ShellActionOptions::default()
66 }
67}
68
69struct UsageEvent {
72 tool: String,
73 action: String,
74 ws: String,
75 principal: String,
76 thread_id: String,
77 request_id: String,
78 client_name: String,
79 client_version: String,
80 protocol: String,
81 client_session: String,
82 conversation_source: String,
83 conversation_id: String,
84 workspace_id: String,
85 workspace_coherence: String,
86 batch_items: usize,
87 worker_limit: usize,
88 started: std::time::Instant,
89}
90
91impl UsageEvent {
92 fn start(
93 request: &CallToolRequestParams,
94 scope: &RequestScope,
95 context: &RequestContext<RoleServer>,
96 ) -> Self {
97 let arguments = request.arguments.as_ref();
98 let thread_id = arguments
99 .and_then(|arguments| arguments.get("thread_id"))
100 .and_then(serde_json::Value::as_str)
101 .map(|thread_id| scope.unscope_text(thread_id))
102 .unwrap_or_default();
103 let action =
104 if request.name == "BashCommand" { bash_action(arguments) } else { String::new() };
105 let ws = if request.name == "Initialize" {
106 arguments
107 .and_then(|arguments| arguments.get("any_workspace_path"))
108 .and_then(serde_json::Value::as_str)
109 .unwrap_or_default()
110 .to_string()
111 } else {
112 String::new()
113 };
114 let workspace = if request.name == "Initialize" {
115 arguments
116 .and_then(|arguments| arguments.get("any_workspace_path"))
117 .and_then(serde_json::Value::as_str)
118 } else {
119 arguments
120 .and_then(|arguments| arguments.get("workspace_root"))
121 .and_then(serde_json::Value::as_str)
122 }
123 .unwrap_or_default();
124 let workspace_id =
125 if workspace.is_empty() { String::new() } else { short_fingerprint("w", workspace) };
126 let client = context.client_info();
127 let client_name =
128 client.as_ref().map_or("unknown", |client| client.name.as_str()).to_string();
129 let client_version =
130 client.as_ref().map_or("unknown", |client| client.version.as_str()).to_string();
131 let protocol = context
132 .protocol_version()
133 .map_or_else(|| "unknown".to_string(), |version| version.to_string());
134 let request_id = context
135 .extensions
136 .get::<Parts>()
137 .and_then(|parts| parts.extensions.get::<crate::http_server::RequestCorrelation>())
138 .map_or_else(
139 || {
140 serde_json::to_string(&context.id)
141 .map_or_else(|_| "unknown".to_string(), |id| short_fingerprint("r", &id))
142 },
143 |correlation| correlation.as_str().to_string(),
144 );
145 let client_session = context
146 .extensions
147 .get::<Parts>()
148 .and_then(|parts| parts.headers.get("mcp-session-id"))
149 .and_then(|value| value.to_str().ok())
150 .map_or_else(|| "stateless".to_string(), |id| short_fingerprint("s", id));
151 let (conversation_source, conversation_id) = conversation_telemetry(context);
152 let batch_items = match request.name.as_ref() {
153 "ReadFiles" => arguments
154 .and_then(|arguments| arguments.get("file_paths"))
155 .and_then(serde_json::Value::as_array)
156 .map_or(0, Vec::len),
157 "MultiFileEdit" => arguments
158 .and_then(|arguments| arguments.get("files"))
159 .and_then(serde_json::Value::as_array)
160 .map_or(0, Vec::len),
161 _ => 0,
162 };
163 let worker_limit = if request.name == "ReadFiles" {
164 crate::tools::read_files::configured_parallelism()
165 } else {
166 0
167 };
168 Self {
169 tool: request.name.to_string(),
170 action,
171 ws,
172 principal: scope
173 .principal()
174 .map_or_else(|| "local".to_string(), |principal| principal.name().to_string()),
175 thread_id,
176 request_id,
177 client_name,
178 client_version,
179 protocol,
180 client_session,
181 conversation_source,
182 conversation_id,
183 workspace_id,
184 workspace_coherence: "unchecked".to_string(),
185 batch_items,
186 worker_limit,
187 started: std::time::Instant::now(),
188 }
189 }
190
191 fn set_workspace_coherence(&mut self, value: &str) {
192 self.workspace_coherence = value.to_string();
193 }
194
195 fn emit(
196 &self,
197 outcome: &str,
198 result_status: &str,
199 response_bytes: usize,
200 result: Option<&CallToolResult>,
201 ) {
202 if self.tool == "Initialize" {
203 let data = result
204 .and_then(|result| result.structured_content.as_ref())
205 .and_then(|structured| structured.get("data"));
206 let initialize_transition = data
207 .and_then(|data| data.get("initialize_transition"))
208 .and_then(serde_json::Value::as_str)
209 .unwrap_or_default();
210 let initialize_reused = data
211 .and_then(|data| data.get("initialize_reused"))
212 .and_then(serde_json::Value::as_bool)
213 .unwrap_or(false);
214 let initialize_response_mode = data
215 .and_then(|data| data.get("initialize_response_mode"))
216 .and_then(serde_json::Value::as_str)
217 .unwrap_or_default();
218 let context_bytes = data
219 .and_then(|data| data.get("context_bytes"))
220 .and_then(serde_json::Value::as_u64)
221 .unwrap_or(0);
222 let guidelines_bytes = data
223 .and_then(|data| data.get("guidelines_bytes"))
224 .and_then(serde_json::Value::as_u64)
225 .unwrap_or(0);
226 let initial_files_count = data
227 .and_then(|data| data.get("initial_files_count"))
228 .and_then(serde_json::Value::as_u64)
229 .unwrap_or(0);
230 let code_writer_policy_strength = data
231 .and_then(|data| data.get("code_writer_policy_strength"))
232 .and_then(serde_json::Value::as_str)
233 .unwrap_or_default();
234 let shell_spawners_present = data
235 .and_then(|data| data.get("shell_spawners_present"))
236 .and_then(serde_json::Value::as_bool)
237 .unwrap_or(false);
238 tracing::info!(
239 target: "winx::usage",
240 event = "tool_call",
241 tool = %self.tool,
242 action = %self.action,
243 ws = %self.ws,
244 principal = %self.principal,
245 thread_id = %self.thread_id,
246 request_id = %self.request_id,
247 client_name = %self.client_name,
248 client_version = %self.client_version,
249 protocol = %self.protocol,
250 client_session = %self.client_session,
251 conversation_source = %self.conversation_source,
252 conversation_id = %self.conversation_id,
253 workspace_id = %self.workspace_id,
254 workspace_coherence = %self.workspace_coherence,
255 batch_items = self.batch_items,
256 worker_limit = self.worker_limit,
257 initialize_transition,
258 initialize_reused,
259 initialize_response_mode,
260 context_bytes,
261 guidelines_bytes,
262 initial_files_count,
263 code_writer_policy_strength,
264 shell_spawners_present,
265 result_status,
266 response_bytes,
267 duration_ms =
268 u64::try_from(self.started.elapsed().as_millis()).unwrap_or(u64::MAX),
269 outcome,
270 "tool call"
271 );
272 return;
273 }
274
275 tracing::info!(
276 target: "winx::usage",
277 event = "tool_call",
278 tool = %self.tool,
279 action = %self.action,
280 ws = %self.ws,
281 principal = %self.principal,
282 thread_id = %self.thread_id,
283 request_id = %self.request_id,
284 client_name = %self.client_name,
285 client_version = %self.client_version,
286 protocol = %self.protocol,
287 client_session = %self.client_session,
288 conversation_source = %self.conversation_source,
289 conversation_id = %self.conversation_id,
290 workspace_id = %self.workspace_id,
291 workspace_coherence = %self.workspace_coherence,
292 batch_items = self.batch_items,
293 worker_limit = self.worker_limit,
294 result_status,
295 response_bytes,
296 duration_ms = u64::try_from(self.started.elapsed().as_millis()).unwrap_or(u64::MAX),
297 outcome,
298 "tool call"
299 );
300 }
301}
302
303fn short_fingerprint(prefix: &str, value: &str) -> String {
304 let digest = Sha256::digest(value.as_bytes());
305 let mut output = format!("{prefix}_");
306 for byte in &digest[..8] {
307 let _ = write!(output, "{byte:02x}");
308 }
309 output
310}
311
312fn conversation_telemetry(context: &RequestContext<RoleServer>) -> (String, String) {
313 let Some(parts) = context.extensions.get::<Parts>() else {
314 return ("none".to_string(), String::new());
315 };
316 for (header, source) in
317 [("mcp-session-id", "mcp_session"), ("x-winx-conversation-id", "gateway_header")]
318 {
319 if let Some(value) = parts
320 .headers
321 .get(header)
322 .and_then(|value| value.to_str().ok())
323 .map(str::trim)
324 .filter(|value| !value.is_empty())
325 {
326 return (source.to_string(), short_fingerprint("c", value));
327 }
328 }
329 ("none".to_string(), String::new())
330}
331
332fn cache_hints_for_protocol(
335 protocol_version: Option<ProtocolVersion>,
336 cache_scope: CacheScope,
337) -> (Option<u64>, Option<CacheScope>) {
338 if protocol_version.is_some_and(|version| version >= ProtocolVersion::V_2026_07_28) {
339 (Some(0), Some(cache_scope))
340 } else {
341 (None, None)
342 }
343}
344
345fn supports_compact_bash_output(context: &RequestContext<RoleServer>) -> bool {
346 context.client_capabilities().is_some_and(|capabilities| {
347 capabilities
348 .extensions
349 .as_ref()
350 .is_some_and(|extensions| extensions.contains_key(COMPACT_BASH_OUTPUT_EXTENSION))
351 })
352}
353
354fn request_requires_bash_capability(request: &CallToolRequestParams) -> bool {
355 matches!(request.name.as_ref(), "FileWriteOrEdit" | "MultiFileEdit")
356 && request
357 .arguments
358 .as_ref()
359 .and_then(|arguments| arguments.get("verify_command"))
360 .is_some_and(|command| !command.is_null())
361}
362
363fn normalized_wait_request(
364 mut request: CallToolRequestParams,
365 policy: BashWaitPolicy,
366 task_inline: bool,
367) -> Result<CallToolRequestParams, McpError> {
368 let mut arguments = request
369 .arguments
370 .clone()
371 .ok_or_else(|| McpError::invalid_request("Missing BashCommand arguments", None))?;
372 let requested = arguments
373 .get("wait_for_seconds")
374 .cloned()
375 .map(serde_json::from_value::<f32>)
376 .transpose()
377 .map_err(|error| {
378 McpError::invalid_request(format!("Invalid wait_for_seconds: {error}"), None)
379 })?;
380 let (default, maximum) = match (policy, task_inline) {
381 (BashWaitPolicy::Adaptive, true) => {
382 (ADAPTIVE_TASK_INLINE_SECONDS, ADAPTIVE_TASK_INLINE_SECONDS)
383 }
384 (BashWaitPolicy::ReturnEarly, _) => {
385 (RETURN_EARLY_DEFAULT_SECONDS, RETURN_EARLY_MAX_SECONDS)
386 }
387 (BashWaitPolicy::Adaptive, false) => (15.0, SYNCHRONOUS_MAX_SECONDS),
388 (BashWaitPolicy::UntilComplete, false) => {
389 (SYNCHRONOUS_MAX_SECONDS, SYNCHRONOUS_MAX_SECONDS)
390 }
391 (BashWaitPolicy::UntilComplete, true) => unreachable!("until_complete starts a Task"),
392 };
393 arguments.insert(
394 "wait_for_seconds".to_string(),
395 serde_json::Value::from(requested.unwrap_or(default).clamp(0.0, maximum)),
396 );
397 request.arguments = Some(arguments);
398 Ok(request)
399}
400
401fn bash_action(arguments: Option<&rmcp::model::JsonObject>) -> String {
407 use serde_json::Value;
408 const KINDS: [&str; 7] = [
409 "command",
410 "status_check",
411 "send_text",
412 "send_specials",
413 "send_ascii",
414 "screen",
415 "wait_for_turn",
416 ];
417 let Some(arguments) = arguments else {
418 return String::new();
419 };
420 let parsed;
421 let action = match arguments.get("action_json") {
422 Some(Value::Object(object)) => object,
423 Some(Value::String(text)) => {
424 match serde_json::from_str::<Value>(&text.replace('\n', " ")) {
425 Ok(Value::Object(object)) => {
426 parsed = object;
427 &parsed
428 }
429 _ => return "command".to_string(),
431 }
432 }
433 _ => arguments,
435 };
436 if let Some(kind) = action.get("type").and_then(Value::as_str) {
437 return kind.to_string();
438 }
439 KINDS.iter().find(|kind| action.contains_key(**kind)).map_or("?", |kind| kind).to_string()
440}
441
442#[allow(clippy::unused_async_trait_impl)] impl ServerHandler for WinxService {
444 fn get_info(&self) -> ServerInfo {
445 let mut capabilities = ServerCapabilities::builder()
446 .enable_tools()
447 .enable_resources()
448 .enable_prompts()
449 .enable_tasks()
450 .build();
451 capabilities
452 .extensions
453 .get_or_insert_default()
454 .insert(COMPACT_BASH_OUTPUT_EXTENSION.to_string(), rmcp::model::JsonObject::new());
455 ServerInfo::new(capabilities)
456 .with_server_info(
457 Implementation::new("winx-mcp-server", self.version.clone())
458 .with_title("Winx High-Performance MCP")
459 .with_icons(vec![Icon::new(server_icon_data_uri())
460 .with_mime_type("image/png")
461 .with_sizes(vec!["96x96".to_string()])]),
462 )
463 .with_protocol_version(ProtocolVersion::V_2026_07_28)
464 .with_instructions(crate::utils::orchestration::server_instructions())
465 }
466
467 async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
468 self.initialize_from_client_roots(context).await;
469 }
470
471 async fn on_roots_list_changed(&self, context: NotificationContext<RoleServer>) {
472 self.initialize_from_client_roots(context).await;
475 }
476
477 fn get_tool(&self, name: &str) -> Option<Tool> {
478 winx_tools().into_iter().find(|tool| tool.name == name)
479 }
480
481 async fn get_task(
482 &self,
483 request: GetTaskParams,
484 context: RequestContext<RoleServer>,
485 ) -> Result<GetTaskResult, McpError> {
486 let principal = principal_from_context(&context);
487 if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
488 return Err(McpError::invalid_request(
489 format!("Unknown or expired task: {}", request.task_id),
490 None,
491 ));
492 }
493 let mut tasks = self.tasks.lock().await;
494 let entry = tasks.get(&request.task_id).ok_or_else(|| {
495 McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
496 })?;
497 Ok(GetTaskResult::new(entry.detailed()))
498 }
499
500 async fn update_task(
501 &self,
502 request: UpdateTaskParams,
503 context: RequestContext<RoleServer>,
504 ) -> Result<(), McpError> {
505 let principal = principal_from_context(&context);
506 if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
507 return Err(McpError::invalid_request(
508 format!("Unknown or expired task: {}", request.task_id),
509 None,
510 ));
511 }
512 let mut tasks = self.tasks.lock().await;
513 tasks.get(&request.task_id).ok_or_else(|| {
514 McpError::invalid_request(format!("Unknown or expired task: {}", request.task_id), None)
515 })?;
516 Ok(())
517 }
518
519 async fn cancel_task(
520 &self,
521 request: CancelTaskParams,
522 context: RequestContext<RoleServer>,
523 ) -> Result<(), McpError> {
524 let principal = principal_from_context(&context);
525 if !task_belongs_to_principal(&request.task_id, principal.as_ref()) {
526 return Err(McpError::invalid_request(
527 format!("Unknown or expired task: {}", request.task_id),
528 None,
529 ));
530 }
531 let (abort_handle, thread_id, mut execution_token, execution_control) = {
532 let mut tasks = self.tasks.lock().await;
533 let entry = tasks.get_mut(&request.task_id).ok_or_else(|| {
534 McpError::invalid_request(
535 format!("Unknown or expired task: {}", request.task_id),
536 None,
537 )
538 })?;
539 if entry.task.status != TaskStatus::Working
540 && entry.task.status != TaskStatus::InputRequired
541 {
542 return Err(McpError::invalid_request(
543 format!("Task {} is already terminal", request.task_id),
544 None,
545 ));
546 }
547 entry.request_cancel();
548 let thread_id = entry.thread_id.clone();
549 let execution_token = entry.execution_token();
550 let execution_control = entry.execution_control();
551 let abort_handle = execution_token.as_ref().and(entry.abort_handle.take());
555 entry.finish(TaskStatus::Cancelled, Some("Cancelled by client".to_string()), None);
556 (abort_handle, thread_id, execution_token, execution_control)
557 };
558 if let Some(abort_handle) = abort_handle {
559 abort_handle.abort();
560 }
561 if execution_token.is_none() {
562 self.cancel_pending_task_action(&thread_id, &request.task_id).await;
563 execution_token = tokio::time::timeout(
564 std::time::Duration::from_millis(500),
565 execution_control.wait_for_execution(),
566 )
567 .await
568 .ok()
569 .flatten();
570 }
571 self.interrupt_task_execution(&thread_id, execution_token).await;
572 Ok(())
573 }
574
575 async fn list_tools(
576 &self,
577 _request: Option<PaginatedRequestParams>,
578 ctx: RequestContext<RoleServer>,
579 ) -> Result<ListToolsResult, McpError> {
580 let principal = principal_from_context(&ctx);
581 let (tools, scope) = principal.as_ref().map_or_else(
582 || (winx_tools(), CacheScope::Public),
583 |principal| {
584 let policy = principal.tool_policy();
585 let scope = if policy == crate::tool_policy::ToolPolicy::default() {
586 CacheScope::Public
587 } else {
588 CacheScope::Private
589 };
590 (winx_tools_for_policy(policy), scope)
591 },
592 );
593 let (ttl_ms, cache_scope) = cache_hints_for_protocol(ctx.protocol_version(), scope);
594 let mut result = ListToolsResult::with_all_items(tools);
595 result.ttl_ms = ttl_ms;
596 result.cache_scope = cache_scope;
597 Ok(result)
598 }
599
600 async fn list_resources(
601 &self,
602 _request: Option<PaginatedRequestParams>,
603 ctx: RequestContext<RoleServer>,
604 ) -> Result<ListResourcesResult, McpError> {
605 let (ttl_ms, cache_scope) =
606 cache_hints_for_protocol(ctx.protocol_version(), CacheScope::Public);
607 let mut result =
608 ListResourcesResult::with_all_items(vec![Resource::new("file://readme", "README")
609 .with_description("Project README documentation")
610 .with_mime_type("text/markdown")]);
611 result.ttl_ms = ttl_ms;
612 result.cache_scope = cache_scope;
613 Ok(result)
614 }
615
616 async fn list_prompts(
617 &self,
618 _request: Option<PaginatedRequestParams>,
619 ctx: RequestContext<RoleServer>,
620 ) -> Result<ListPromptsResult, McpError> {
621 let (ttl_ms, cache_scope) =
622 cache_hints_for_protocol(ctx.protocol_version(), CacheScope::Public);
623 let mut result = ListPromptsResult::with_all_items(winx_prompts());
624 result.ttl_ms = ttl_ms;
625 result.cache_scope = cache_scope;
626 Ok(result)
627 }
628
629 async fn get_prompt(
630 &self,
631 request: GetPromptRequestParams,
632 context: RequestContext<RoleServer>,
633 ) -> Result<GetPromptResponse, McpError> {
634 if request.name != "KnowledgeTransfer" {
635 return Err(McpError::invalid_request(
636 format!("Unknown prompt: {}", request.name),
637 None,
638 ));
639 }
640
641 let principal = principal_from_context(&context);
642 let session_prefix = principal.as_ref().map(crate::config::HttpPrincipal::session_prefix);
643 let text = crate::utils::redact::redact(
644 &self.knowledge_transfer_prompt_text(session_prefix.as_deref()).await,
645 )
646 .into_owned();
647 Ok(GetPromptResult::new(vec![PromptMessage::new_text(Role::User, text)])
648 .with_description("Knowledge transfer handoff prompt")
649 .into())
650 }
651
652 async fn read_resource(
653 &self,
654 request: ReadResourceRequestParams,
655 ctx: RequestContext<RoleServer>,
656 ) -> Result<ReadResourceResponse, McpError> {
657 let content = match request.uri.as_ref() {
658 "file://readme" => match tokio::fs::read_to_string("README.md").await {
659 Ok(content) => vec![ResourceContents::text(
660 crate::utils::redact::redact(&content).into_owned(),
661 request.uri.clone(),
662 )],
663 Err(_) => vec![ResourceContents::text(
664 "README.md not found".to_string(),
665 request.uri.clone(),
666 )],
667 },
668 _ => {
669 return Err(McpError::invalid_request(
670 format!("Unknown resource URI: {}", request.uri),
671 None,
672 ));
673 }
674 };
675 let (ttl_ms, cache_scope) =
676 cache_hints_for_protocol(ctx.protocol_version(), CacheScope::Private);
677 let mut result = ReadResourceResult::new(content);
678 result.ttl_ms = ttl_ms;
679 result.cache_scope = cache_scope;
680 Ok(result.into())
681 }
682
683 #[allow(clippy::too_many_lines)] async fn call_tool(
685 &self,
686 request: CallToolRequestParams,
687 context: RequestContext<RoleServer>,
688 ) -> Result<CallToolResponse, McpError> {
689 let principal = principal_from_context(&context);
690 if principal
691 .as_ref()
692 .is_some_and(|principal| !principal.tool_policy().allows(request.name.as_ref()))
693 {
694 return Err(McpError::invalid_request(
695 format!("Tool is not available for this principal: {}", request.name),
696 None,
697 ));
698 }
699 if request_requires_bash_capability(&request)
700 && principal
701 .as_ref()
702 .is_some_and(|principal| !principal.tool_policy().allows("BashCommand"))
703 {
704 return Err(McpError::invalid_request(
705 "verify_command requires BashCommand in this principal's tool policy",
706 None,
707 ));
708 }
709 let affinity = session_affinity_from_context(&context);
710 let conversation_identity = conversation_identity_from_context(&context);
711 let (request, scope) =
712 scope_tool_request(request, principal, affinity, conversation_identity.as_deref())
713 .map_err(|error| {
714 McpError::invalid_request(format!("Cannot scope remote request: {error}"), None)
715 })?;
716 let mut usage = UsageEvent::start(&request, &scope, &context);
717 match self
718 .validate_workspace_coherence(
719 &request,
720 &scope,
721 affinity,
722 conversation_identity.as_deref(),
723 )
724 .await
725 {
726 Ok(coherence) => usage.set_workspace_coherence(coherence.as_str()),
727 Err(error) => {
728 usage.set_workspace_coherence("rejected");
729 let arguments = request.arguments.clone().map(serde_json::Value::Object);
730 let mut result =
731 outcomes::tool_failure(request.name.as_ref(), &error, arguments.as_ref())?;
732 scope.unscope_result(&mut result);
733 let status = outcomes::result_status(&result);
734 usage.emit(
735 "tool_error",
736 &status,
737 outcomes::result_size_bytes(&result),
738 Some(&result),
739 );
740 return Ok(result.into());
741 }
742 }
743 let compact_bash_output = supports_compact_bash_output(&context);
744 let supports_tasks =
745 context.client_capabilities().is_some_and(|capabilities| capabilities.supports_tasks());
746 let wait_policy = if request.name == "BashCommand" {
747 Some(Self::bash_wait_policy(&request)?)
748 } else {
749 None
750 };
751 if wait_policy == Some(BashWaitPolicy::UntilComplete) {
752 let is_foreground_command = request
753 .arguments
754 .clone()
755 .and_then(|arguments| {
756 serde_json::from_value::<BashCommand>(serde_json::Value::Object(arguments)).ok()
757 })
758 .is_some_and(|bash| {
759 matches!(
760 bash.action_json,
761 crate::types::BashCommandAction::Command { is_background: false, .. }
762 )
763 });
764 if !is_foreground_command {
765 return Err(McpError::invalid_request(
766 "wait_policy=until_complete is valid only for a foreground Command action",
767 None,
768 ));
769 }
770 }
771
772 let eligible = Self::bash_task_is_eligible(&request);
773 let generation_bound_runtime = if supports_tasks && eligible {
774 let thread_id = request
775 .arguments
776 .as_ref()
777 .and_then(|arguments| arguments.get("thread_id"))
778 .and_then(serde_json::Value::as_str)
779 .unwrap_or_default();
780 let (slot, session_guard) = self.session_for(thread_id).await;
781 let supported = self
782 .shell_runtime
783 .supports_generation_bound_actions_for(&slot)
784 .await
785 .unwrap_or(false);
786 drop(session_guard);
787 supported
788 } else {
789 false
790 };
791 let route = wait_policy.map_or(BashTaskRoute::Synchronous, |policy| {
792 bash_task_route(supports_tasks, generation_bound_runtime, eligible, policy)
793 });
794
795 match route {
796 BashTaskRoute::Immediate => {
797 let reservation = self.reserve_bash_task(&request, &scope).await?;
798 return match self
799 .start_reserved_bash_task(
800 reservation,
801 request,
802 scope,
803 None,
804 compact_bash_output,
805 )
806 .await
807 {
808 Ok(task) => {
809 usage.emit("task", "working", 0, None);
810 Ok(CallToolResponse::Task(task))
811 }
812 Err(error) => {
813 usage.emit("protocol_error", "failed", 0, None);
814 Err(error)
815 }
816 };
817 }
818 BashTaskRoute::Adaptive => {
819 let reservation = self.reserve_bash_task(&request, &scope).await?;
820 let task_request = request.clone();
821 let inline_request =
822 normalized_wait_request(request, BashWaitPolicy::Adaptive, true)?;
823 match self
824 .execute_tool_call(inline_request, adaptive_action_options(compact_bash_output))
825 .await
826 {
827 Ok(execution) if outcomes::result_status(&execution.result) == "running" => {
828 return match self
829 .start_reserved_bash_task(
830 reservation,
831 task_request,
832 scope,
833 Some(execution),
834 compact_bash_output,
835 )
836 .await
837 {
838 Ok(task) => {
839 usage.emit("task", "working", 0, None);
840 Ok(CallToolResponse::Task(task))
841 }
842 Err(error) => {
843 usage.emit("protocol_error", "failed", 0, None);
844 Err(error)
845 }
846 };
847 }
848 Ok(mut execution) => {
849 self.release_bash_task(&reservation).await;
850 scope.unscope_result(&mut execution.result);
851 let status = outcomes::result_status(&execution.result);
852 let outcome = if execution.result.is_error == Some(true) {
853 "tool_error"
854 } else {
855 "ok"
856 };
857 usage.emit(
858 outcome,
859 &status,
860 outcomes::result_size_bytes(&execution.result),
861 Some(&execution.result),
862 );
863 return Ok(execution.result.into());
864 }
865 Err(mut error) => {
866 self.release_bash_task(&reservation).await;
867 scope.unscope_error(&mut error);
868 usage.emit("protocol_error", "failed", 0, None);
869 return Err(error);
870 }
871 }
872 }
873 BashTaskRoute::Synchronous => {}
874 }
875
876 let request = match wait_policy {
877 Some(policy) => normalized_wait_request(request, policy, false)?,
878 None => request,
879 };
880 match self
881 .execute_tool_call(
882 request,
883 ShellActionOptions {
884 compact_output: compact_bash_output,
885 ..ShellActionOptions::default()
886 },
887 )
888 .await
889 {
890 Ok(mut execution) => {
891 let result = &mut execution.result;
892 scope.unscope_result(result);
893 let status = outcomes::result_status(result);
894 let outcome = if result.is_error == Some(true) { "tool_error" } else { "ok" };
895 usage.emit(outcome, &status, outcomes::result_size_bytes(result), Some(result));
896 Ok(execution.result.into())
897 }
898 Err(mut error) => {
899 scope.unscope_error(&mut error);
900 usage.emit("protocol_error", "failed", 0, None);
901 Err(error)
902 }
903 }
904 }
905}
906
907#[cfg(test)]
908mod tests {
909 use super::*;
910
911 #[test]
912 fn cache_hints_follow_negotiated_protocol_version() {
913 assert_eq!(
914 cache_hints_for_protocol(Some(ProtocolVersion::V_2026_07_28), CacheScope::Private,),
915 (Some(0), Some(CacheScope::Private))
916 );
917 assert_eq!(
918 cache_hints_for_protocol(Some(ProtocolVersion::V_2026_07_28), CacheScope::Public),
919 (Some(0), Some(CacheScope::Public))
920 );
921 assert_eq!(
922 cache_hints_for_protocol(Some(ProtocolVersion::V_2025_11_25), CacheScope::Private,),
923 (None, None)
924 );
925 assert_eq!(cache_hints_for_protocol(None, CacheScope::Private), (None, None));
926 }
927
928 #[test]
929 fn guardian_1_4_uses_synchronous_policy_fallback_even_for_task_clients() {
930 assert_eq!(
931 bash_task_route(true, false, true, BashWaitPolicy::Adaptive),
932 BashTaskRoute::Synchronous
933 );
934 assert_eq!(
935 bash_task_route(true, false, true, BashWaitPolicy::UntilComplete),
936 BashTaskRoute::Synchronous
937 );
938 }
939
940 #[test]
941 fn task_routing_requires_client_runtime_and_foreground_eligibility() {
942 assert_eq!(
943 bash_task_route(true, true, true, BashWaitPolicy::Adaptive),
944 BashTaskRoute::Adaptive
945 );
946 assert_eq!(
947 bash_task_route(true, true, true, BashWaitPolicy::UntilComplete),
948 BashTaskRoute::Immediate
949 );
950 assert_eq!(
951 bash_task_route(true, true, true, BashWaitPolicy::ReturnEarly),
952 BashTaskRoute::Synchronous
953 );
954 assert_eq!(
955 bash_task_route(false, true, true, BashWaitPolicy::UntilComplete),
956 BashTaskRoute::Synchronous
957 );
958 assert_eq!(
959 bash_task_route(true, true, false, BashWaitPolicy::UntilComplete),
960 BashTaskRoute::Synchronous
961 );
962 }
963
964 #[test]
965 fn adaptive_action_requires_live_generation_binding() {
966 let options = adaptive_action_options(true);
967 assert!(options.compact_output);
968 assert!(options.require_generation_binding);
969 }
970
971 #[test]
972 #[allow(clippy::expect_used, clippy::float_cmp)]
973 fn synchronous_wait_policies_have_explicit_upper_bounds() {
974 let request = |wait_for_seconds: Option<f32>| {
975 let mut arguments = serde_json::Map::new();
976 arguments.insert("command".to_string(), serde_json::Value::String("pwd".to_string()));
977 if let Some(wait) = wait_for_seconds {
978 arguments.insert("wait_for_seconds".to_string(), serde_json::json!(wait));
979 }
980 CallToolRequestParams::new("BashCommand").with_arguments(arguments)
981 };
982 let normalized = |policy, wait| {
983 normalized_wait_request(request(wait), policy, false)
984 .expect("valid wait request")
985 .arguments
986 .and_then(|arguments| arguments.get("wait_for_seconds").cloned())
987 .and_then(|value| value.as_f64())
988 .expect("normalized wait")
989 };
990
991 assert_eq!(normalized(BashWaitPolicy::Adaptive, Some(500.0)), 60.0);
992 assert_eq!(normalized(BashWaitPolicy::UntilComplete, None), 60.0);
993 assert_eq!(normalized(BashWaitPolicy::ReturnEarly, Some(500.0)), 5.0);
994 assert_eq!(normalized(BashWaitPolicy::ReturnEarly, None), 0.25);
995 }
996}