1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use incurs::command::RequestContext;
6use incurs::tool::{ToolCallControl, ToolCallOptions, ToolCallOutcome, ToolCatalog};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9
10#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
12#[serde(rename_all = "snake_case")]
13pub enum ReplayPolicy {
14 #[default]
16 Log,
17 Reexecute,
19}
20
21#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
23pub struct ToolAnnotations {
24 pub read_only: Option<bool>,
26 pub destructive: Option<bool>,
28 pub idempotent: Option<bool>,
30 pub open_world: Option<bool>,
32}
33
34#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
36pub struct ToolPolicy {
37 pub requires_approval: bool,
39 pub replay: ReplayPolicy,
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum ToolOrigin {
46 Local,
48 RemoteMcp,
50 OpenApi,
52}
53
54pub trait ToolPolicyResolver: Send + Sync {
56 fn resolve(&self, origin: ToolOrigin, annotations: &ToolAnnotations) -> ToolPolicy;
58}
59
60#[derive(Debug, Clone, Copy, Default)]
62pub struct DefaultToolPolicyResolver;
63
64impl ToolPolicyResolver for DefaultToolPolicyResolver {
65 fn resolve(&self, origin: ToolOrigin, annotations: &ToolAnnotations) -> ToolPolicy {
66 let safe_local_read = origin == ToolOrigin::Local
67 && annotations.read_only == Some(true)
68 && annotations.destructive != Some(true)
69 && annotations.open_world != Some(true);
70 ToolPolicy {
71 requires_approval: !safe_local_read,
72 replay: if safe_local_read {
73 ReplayPolicy::Reexecute
74 } else {
75 ReplayPolicy::Log
76 },
77 }
78 }
79}
80
81#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
83pub struct ConnectorExample {
84 pub command: String,
86 pub description: Option<String>,
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
92pub struct ConnectorTool {
93 pub name: String,
95 pub description: Option<String>,
97 pub input_schema: Value,
99 pub output_schema: Option<Value>,
101 pub instructions: Option<String>,
103 pub examples: Vec<ConnectorExample>,
105 pub annotations: ToolAnnotations,
107 pub policy: ToolPolicy,
109}
110
111#[derive(Debug, Clone, Serialize, Deserialize)]
113pub struct ConnectorDescription {
114 pub name: String,
116 pub instructions: Option<String>,
118 pub tools: Vec<ConnectorTool>,
120}
121
122#[derive(Clone)]
124pub struct ToolContext {
125 pub execution_id: String,
127 pub control: ToolCallControl,
129 pub request: Option<RequestContext>,
131}
132
133#[async_trait]
135pub trait Connector: Send + Sync {
136 fn name(&self) -> &str;
143
144 async fn describe(&self) -> Result<ConnectorDescription, String>;
146
147 async fn execute(
149 &self,
150 method: &str,
151 arguments: Value,
152 context: &ToolContext,
153 ) -> Result<Value, String>;
154
155 async fn revert(
157 &self,
158 _method: &str,
159 _arguments: Value,
160 _result: Value,
161 _context: &ToolContext,
162 ) -> Result<bool, String> {
163 Ok(false)
164 }
165
166 async fn pass_ended(&self, _execution_id: &str, _status: &str) {}
168
169 async fn execution_ended(&self, _execution_id: &str, _status: &str) {}
171}
172
173#[derive(Clone)]
175pub struct IncurConnector {
176 catalog: ToolCatalog,
177 name: String,
178 instructions: Option<String>,
179 options: ToolCallOptions,
180 policy: Arc<dyn ToolPolicyResolver>,
181}
182
183impl IncurConnector {
184 pub fn new(catalog: ToolCatalog) -> Self {
186 let name = sanitize_namespace(catalog.name());
187 Self {
188 catalog,
189 name,
190 instructions: None,
191 options: ToolCallOptions::default(),
192 policy: Arc::new(DefaultToolPolicyResolver),
193 }
194 }
195
196 pub fn with_name(mut self, name: impl Into<String>) -> Self {
198 self.name = name.into();
199 self
200 }
201
202 pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self {
204 self.instructions = Some(instructions.into());
205 self
206 }
207
208 pub fn with_call_options(mut self, options: ToolCallOptions) -> Self {
210 self.options = options;
211 self
212 }
213
214 pub fn with_policy_resolver(mut self, resolver: Arc<dyn ToolPolicyResolver>) -> Self {
216 self.policy = resolver;
217 self
218 }
219}
220
221#[async_trait]
222impl Connector for IncurConnector {
223 fn name(&self) -> &str {
224 &self.name
225 }
226
227 async fn describe(&self) -> Result<ConnectorDescription, String> {
228 Ok(ConnectorDescription {
229 name: self.name.clone(),
230 instructions: self.instructions.clone(),
231 tools: self
232 .catalog
233 .definitions()
234 .into_iter()
235 .map(|tool| {
236 let annotations = ToolAnnotations {
237 read_only: tool
238 .annotations
239 .as_ref()
240 .and_then(|annotations| annotations.read_only_hint),
241 destructive: tool
242 .annotations
243 .as_ref()
244 .and_then(|annotations| annotations.destructive_hint),
245 idempotent: tool
246 .annotations
247 .as_ref()
248 .and_then(|annotations| annotations.idempotent_hint),
249 open_world: tool
250 .annotations
251 .as_ref()
252 .and_then(|annotations| annotations.open_world_hint),
253 };
254 ConnectorTool {
255 name: tool.name,
256 description: (!tool.description.is_empty()).then_some(tool.description),
257 input_schema: tool.input_schema,
258 output_schema: tool.output_schema,
259 instructions: tool.instructions,
260 examples: tool
261 .examples
262 .into_iter()
263 .map(|example| ConnectorExample {
264 command: example.command,
265 description: example.description,
266 })
267 .collect(),
268 policy: self.policy.resolve(ToolOrigin::Local, &annotations),
269 annotations,
270 }
271 })
272 .collect(),
273 })
274 }
275
276 async fn execute(
277 &self,
278 method: &str,
279 arguments: Value,
280 context: &ToolContext,
281 ) -> Result<Value, String> {
282 let arguments = arguments
283 .as_object()
284 .ok_or_else(|| format!("Arguments to {method} must be an object"))?
285 .iter()
286 .map(|(key, value)| (key.clone(), value.clone()))
287 .collect::<BTreeMap<_, _>>();
288 let mut options = self.options.clone();
289 options.control = context.control.clone();
290 if context.request.is_some() {
291 options.request = context.request.clone();
292 }
293 match self.catalog.call(method, arguments, options).await {
294 ToolCallOutcome::Ok { data, cta } => {
295 if cta.is_some() {
296 Ok(serde_json::json!({ "data": data, "cta": cta }))
297 } else {
298 Ok(data)
299 }
300 }
301 ToolCallOutcome::Error {
302 code,
303 message,
304 retryable,
305 field_errors,
306 cta,
307 exit_code,
308 } => Err(serde_json::json!({
309 "code": code,
310 "message": message,
311 "retryable": retryable,
312 "fieldErrors": field_errors,
313 "cta": cta,
314 "exitCode": exit_code,
315 })
316 .to_string()),
317 }
318 }
319}
320
321pub fn sanitize_namespace(value: &str) -> String {
323 let mut result = String::new();
324 for (index, ch) in value.chars().enumerate() {
325 if (index == 0 && !(ch == '_' || ch == '$' || ch.is_ascii_alphabetic()))
326 || (index > 0 && !(ch == '_' || ch == '$' || ch.is_ascii_alphanumeric()))
327 {
328 result.push('_');
329 } else {
330 result.push(ch);
331 }
332 }
333 if result.is_empty() {
334 "tools".to_string()
335 } else {
336 result
337 }
338}
339
340#[derive(Debug, Clone, Serialize, Deserialize)]
342pub struct McpTool {
343 pub name: String,
345 pub description: Option<String>,
347 pub input_schema: Value,
349 pub output_schema: Option<Value>,
351 pub annotations: Option<incurs::command::McpAnnotations>,
353}
354
355#[async_trait]
357pub trait McpClient: Send + Sync {
358 async fn list_tools(&self) -> Result<Vec<McpTool>, String>;
360
361 async fn call_tool(&self, name: &str, arguments: Value) -> Result<Value, String>;
363
364 async fn call_tool_cancellable(
372 &self,
373 name: &str,
374 arguments: Value,
375 cancellation: &tokio_util::sync::CancellationToken,
376 ) -> Result<Value, String> {
377 tokio::select! {
378 biased;
379 () = cancellation.cancelled() => Err("Call cancelled".to_string()),
380 result = self.call_tool(name, arguments) => result,
381 }
382 }
383}
384
385pub struct McpConnector {
387 name: String,
388 instructions: Option<String>,
389 client: Arc<dyn McpClient>,
390 tools: tokio::sync::OnceCell<Vec<(String, McpTool)>>,
391 policy: Arc<dyn ToolPolicyResolver>,
392}
393
394impl McpConnector {
395 pub fn new(name: impl Into<String>, client: Arc<dyn McpClient>) -> Self {
397 Self {
398 name: name.into(),
399 instructions: None,
400 client,
401 tools: tokio::sync::OnceCell::new(),
402 policy: Arc::new(DefaultToolPolicyResolver),
403 }
404 }
405
406 pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self {
408 self.instructions = Some(instructions.into());
409 self
410 }
411
412 pub fn with_policy_resolver(mut self, resolver: Arc<dyn ToolPolicyResolver>) -> Self {
414 self.policy = resolver;
415 self
416 }
417
418 async fn tools(&self) -> Result<&Vec<(String, McpTool)>, String> {
419 self.tools
420 .get_or_try_init(|| async {
421 let mut names = BTreeMap::new();
422 let mut tools = Vec::new();
423 for tool in self.client.list_tools().await? {
424 let name = sanitize_namespace(&tool.name);
425 if let Some(existing) = names.insert(name.clone(), tool.name.clone()) {
426 return Err(format!(
427 "MCP tools \"{existing}\" and \"{}\" both map to \"{name}\"",
428 tool.name
429 ));
430 }
431 tools.push((name, tool));
432 }
433 Ok(tools)
434 })
435 .await
436 }
437}
438
439#[async_trait]
440impl Connector for McpConnector {
441 fn name(&self) -> &str {
442 &self.name
443 }
444
445 async fn describe(&self) -> Result<ConnectorDescription, String> {
446 Ok(ConnectorDescription {
447 name: self.name.clone(),
448 instructions: self.instructions.clone(),
449 tools: self
450 .tools()
451 .await?
452 .iter()
453 .map(|(name, tool)| {
454 let annotations = ToolAnnotations {
455 read_only: tool
456 .annotations
457 .as_ref()
458 .and_then(|annotations| annotations.read_only_hint),
459 destructive: tool
460 .annotations
461 .as_ref()
462 .and_then(|annotations| annotations.destructive_hint),
463 idempotent: tool
464 .annotations
465 .as_ref()
466 .and_then(|annotations| annotations.idempotent_hint),
467 open_world: tool
468 .annotations
469 .as_ref()
470 .and_then(|annotations| annotations.open_world_hint),
471 };
472 ConnectorTool {
473 name: name.clone(),
474 description: tool.description.clone(),
475 input_schema: tool.input_schema.clone(),
476 output_schema: tool.output_schema.clone(),
477 instructions: None,
478 examples: Vec::new(),
479 policy: self.policy.resolve(ToolOrigin::RemoteMcp, &annotations),
480 annotations,
481 }
482 })
483 .collect(),
484 })
485 }
486
487 async fn execute(
488 &self,
489 method: &str,
490 arguments: Value,
491 context: &ToolContext,
492 ) -> Result<Value, String> {
493 let (_, tool) = self
494 .tools()
495 .await?
496 .iter()
497 .find(|(name, _)| name == method)
498 .ok_or_else(|| format!("Tool \"{method}\" not found on {}", self.name))?;
499 self.client
500 .call_tool_cancellable(&tool.name, arguments, &context.control.cancellation)
501 .await
502 }
503}
504
505#[derive(Debug, Clone, Serialize, Deserialize)]
507pub struct OpenApiRequest {
508 pub path: String,
510 pub method: String,
512 pub parameters: BTreeMap<String, Value>,
514 pub body: Option<Value>,
516 pub headers: BTreeMap<String, String>,
518}
519
520#[async_trait]
522pub trait OpenApiClient: Send + Sync {
523 async fn specification(&self) -> Result<Value, String>;
525
526 async fn request(&self, request: OpenApiRequest) -> Result<Value, String>;
528}
529
530#[derive(Clone)]
531struct OpenApiOperation {
532 name: String,
533 method: String,
534 path: String,
535 description: String,
536 input_schema: Value,
537 parameters: Vec<(String, String)>,
538}
539
540pub struct OpenApiConnector {
542 name: String,
543 instructions: Option<String>,
544 client: Arc<dyn OpenApiClient>,
545 operations: tokio::sync::OnceCell<Vec<OpenApiOperation>>,
546 policy: Arc<dyn ToolPolicyResolver>,
547}
548
549impl OpenApiConnector {
550 pub fn new(name: impl Into<String>, client: Arc<dyn OpenApiClient>) -> Self {
552 Self {
553 name: name.into(),
554 instructions: None,
555 client,
556 operations: tokio::sync::OnceCell::new(),
557 policy: Arc::new(DefaultToolPolicyResolver),
558 }
559 }
560
561 pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self {
563 self.instructions = Some(instructions.into());
564 self
565 }
566
567 pub fn with_policy_resolver(mut self, resolver: Arc<dyn ToolPolicyResolver>) -> Self {
569 self.policy = resolver;
570 self
571 }
572
573 async fn operations(&self) -> Result<&Vec<OpenApiOperation>, String> {
574 self.operations
575 .get_or_try_init(|| async {
576 derive_openapi_operations(&self.client.specification().await?)
577 })
578 .await
579 }
580}
581
582#[async_trait]
583impl Connector for OpenApiConnector {
584 fn name(&self) -> &str {
585 &self.name
586 }
587
588 async fn describe(&self) -> Result<ConnectorDescription, String> {
589 let mut tools = vec![ConnectorTool {
590 name: "request".to_string(),
591 description: Some(
592 "Perform an authenticated request when no derived operation fits.".to_string(),
593 ),
594 input_schema: serde_json::json!({
595 "type": "object",
596 "properties": {
597 "path": {"type": "string"},
598 "method": {"type": "string"},
599 "parameters": {"type": "object", "additionalProperties": true},
600 "body": {},
601 "headers": {"type": "object", "additionalProperties": {"type": "string"}}
602 },
603 "required": ["path"]
604 }),
605 output_schema: None,
606 instructions: None,
607 examples: Vec::new(),
608 annotations: ToolAnnotations {
609 open_world: Some(true),
610 ..ToolAnnotations::default()
611 },
612 policy: ToolPolicy {
613 requires_approval: true,
614 replay: ReplayPolicy::Log,
615 },
616 }];
617 tools.extend(self.operations().await?.iter().map(|operation| {
618 let annotations = ToolAnnotations {
619 read_only: Some(operation.method == "get" || operation.method == "head"),
620 open_world: Some(true),
621 ..ToolAnnotations::default()
622 };
623 ConnectorTool {
624 name: operation.name.clone(),
625 description: Some(operation.description.clone()),
626 input_schema: operation.input_schema.clone(),
627 output_schema: None,
628 instructions: None,
629 examples: Vec::new(),
630 policy: self.policy.resolve(ToolOrigin::OpenApi, &annotations),
631 annotations,
632 }
633 }));
634 Ok(ConnectorDescription {
635 name: self.name.clone(),
636 instructions: self.instructions.clone(),
637 tools,
638 })
639 }
640
641 async fn execute(
642 &self,
643 method: &str,
644 arguments: Value,
645 _context: &ToolContext,
646 ) -> Result<Value, String> {
647 if method == "request" {
648 return self.client.request(parse_raw_request(arguments)?).await;
649 }
650 let operation = self
651 .operations()
652 .await?
653 .iter()
654 .find(|operation| operation.name == method)
655 .ok_or_else(|| format!("Tool \"{method}\" not found on {}", self.name))?;
656 self.client
657 .request(operation_request(operation, arguments)?)
658 .await
659 }
660}
661
662fn derive_openapi_operations(document: &Value) -> Result<Vec<OpenApiOperation>, String> {
663 let Some(paths) = document.get("paths").and_then(Value::as_object) else {
664 return Ok(Vec::new());
665 };
666 let mut used = BTreeMap::new();
667 let mut operations = Vec::new();
668 for (path, item) in paths {
669 let Some(item) = item.as_object() else {
670 continue;
671 };
672 for method in ["get", "put", "post", "delete", "patch", "options", "head"] {
673 let Some(operation) = item.get(method).and_then(Value::as_object) else {
674 continue;
675 };
676 let source_name = operation
677 .get("operationId")
678 .and_then(Value::as_str)
679 .map(str::to_string)
680 .unwrap_or_else(|| format!("{method}_{path}"));
681 let name = sanitize_namespace(&source_name);
682 if name == "request" || name == "spec" || used.insert(name.clone(), path).is_some() {
683 continue;
684 }
685 let mut properties = serde_json::Map::new();
686 let mut required = Vec::new();
687 let mut parameters = Vec::new();
688 for parameter in operation
689 .get("parameters")
690 .and_then(Value::as_array)
691 .into_iter()
692 .flatten()
693 {
694 let Some(parameter) = parameter.as_object() else {
695 continue;
696 };
697 let Some(parameter_name) = parameter.get("name").and_then(Value::as_str) else {
698 continue;
699 };
700 let location = parameter
701 .get("in")
702 .and_then(Value::as_str)
703 .unwrap_or("query");
704 properties.insert(
705 parameter_name.to_string(),
706 parameter
707 .get("schema")
708 .cloned()
709 .unwrap_or_else(|| serde_json::json!({})),
710 );
711 parameters.push((parameter_name.to_string(), location.to_string()));
712 if parameter.get("required").and_then(Value::as_bool) == Some(true) {
713 required.push(Value::String(parameter_name.to_string()));
714 }
715 }
716 if let Some(body) = operation
717 .get("requestBody")
718 .and_then(|value| value.get("content"))
719 .and_then(|value| value.get("application/json"))
720 .and_then(|value| value.get("schema"))
721 .cloned()
722 {
723 properties.insert("body".to_string(), body);
724 if operation
725 .get("requestBody")
726 .and_then(|value| value.get("required"))
727 .and_then(Value::as_bool)
728 == Some(true)
729 {
730 required.push(Value::String("body".to_string()));
731 }
732 }
733 operations.push(OpenApiOperation {
734 name,
735 method: method.to_string(),
736 path: path.clone(),
737 description: operation
738 .get("summary")
739 .or_else(|| operation.get("description"))
740 .and_then(Value::as_str)
741 .map(str::to_string)
742 .unwrap_or_else(|| format!("{} {path}", method.to_ascii_uppercase())),
743 input_schema: serde_json::json!({
744 "type": "object",
745 "properties": properties,
746 "required": required,
747 }),
748 parameters,
749 });
750 }
751 }
752 Ok(operations)
753}
754
755fn operation_request(
756 operation: &OpenApiOperation,
757 arguments: Value,
758) -> Result<OpenApiRequest, String> {
759 let input = arguments
760 .as_object()
761 .ok_or_else(|| format!("Arguments to {} must be an object", operation.name))?;
762 let mut path = operation.path.clone();
763 let mut parameters = BTreeMap::new();
764 let mut headers = BTreeMap::new();
765 for (name, location) in &operation.parameters {
766 let Some(value) = input.get(name) else {
767 continue;
768 };
769 match location.as_str() {
770 "path" => {
771 path = path.replace(
772 &format!("{{{name}}}"),
773 value.as_str().unwrap_or(&value.to_string()),
774 )
775 }
776 "header" => {
777 headers.insert(
778 name.clone(),
779 value
780 .as_str()
781 .map(str::to_string)
782 .unwrap_or_else(|| value.to_string()),
783 );
784 }
785 "query" => {
786 parameters.insert(name.clone(), value.clone());
787 }
788 _ => {}
789 }
790 }
791 Ok(OpenApiRequest {
792 path,
793 method: operation.method.clone(),
794 parameters,
795 body: input.get("body").cloned(),
796 headers,
797 })
798}
799
800fn parse_raw_request(arguments: Value) -> Result<OpenApiRequest, String> {
801 let input = arguments
802 .as_object()
803 .ok_or_else(|| "Arguments to request must be an object".to_string())?;
804 Ok(OpenApiRequest {
805 path: input
806 .get("path")
807 .and_then(Value::as_str)
808 .ok_or_else(|| "request.path is required".to_string())?
809 .to_string(),
810 method: input
811 .get("method")
812 .and_then(Value::as_str)
813 .unwrap_or("GET")
814 .to_string(),
815 parameters: object_map(input.get("parameters")),
816 body: input.get("body").cloned(),
817 headers: input
818 .get("headers")
819 .and_then(Value::as_object)
820 .into_iter()
821 .flatten()
822 .map(|(key, value)| {
823 (
824 key.clone(),
825 value
826 .as_str()
827 .map(str::to_string)
828 .unwrap_or_else(|| value.to_string()),
829 )
830 })
831 .collect(),
832 })
833}
834
835fn object_map(value: Option<&Value>) -> BTreeMap<String, Value> {
836 value
837 .and_then(Value::as_object)
838 .into_iter()
839 .flatten()
840 .map(|(key, value)| (key.clone(), value.clone()))
841 .collect()
842}
843
844#[cfg(test)]
845mod policy_tests {
846 use super::*;
847
848 #[test]
849 fn only_safe_local_reads_skip_approval() {
850 let resolver = DefaultToolPolicyResolver;
851 let read = ToolAnnotations {
852 read_only: Some(true),
853 ..ToolAnnotations::default()
854 };
855 assert_eq!(
856 resolver.resolve(ToolOrigin::Local, &read),
857 ToolPolicy {
858 requires_approval: false,
859 replay: ReplayPolicy::Reexecute,
860 }
861 );
862 assert!(
863 resolver
864 .resolve(ToolOrigin::RemoteMcp, &read)
865 .requires_approval
866 );
867 assert!(
868 resolver
869 .resolve(
870 ToolOrigin::Local,
871 &ToolAnnotations {
872 open_world: Some(true),
873 ..read
874 }
875 )
876 .requires_approval
877 );
878 }
879}
880
881#[cfg(test)]
882mod mcp_tests {
883 use std::sync::Arc;
884 use std::sync::atomic::{AtomicUsize, Ordering};
885
886 use async_trait::async_trait;
887 use incurs::tool::ToolCallControl;
888 use serde_json::{Value, json};
889
890 use super::{Connector, McpClient, McpConnector, McpTool, ToolContext};
891
892 struct HangingClient {
898 started: Arc<AtomicUsize>,
899 }
900
901 #[async_trait]
902 impl McpClient for HangingClient {
903 async fn list_tools(&self) -> Result<Vec<McpTool>, String> {
904 Ok(vec![McpTool {
905 name: "wait".to_string(),
906 description: None,
907 input_schema: json!({"type": "object"}),
908 output_schema: None,
909 annotations: None,
910 }])
911 }
912
913 async fn call_tool(&self, _name: &str, _arguments: Value) -> Result<Value, String> {
914 self.started.fetch_add(1, Ordering::SeqCst);
915 std::future::pending::<()>().await;
917 unreachable!("pending future resolved")
918 }
919 }
920
921 fn context(control: ToolCallControl) -> ToolContext {
922 ToolContext {
923 execution_id: "exec_test".to_string(),
924 control,
925 request: None,
926 }
927 }
928
929 #[tokio::test]
930 async fn cancellation_reaches_a_client_that_only_implements_call_tool() {
931 let started = Arc::new(AtomicUsize::new(0));
932 let connector = McpConnector::new(
933 "hang",
934 Arc::new(HangingClient {
935 started: Arc::clone(&started),
936 }),
937 );
938 let control = ToolCallControl::default();
939 let cancellation = control.cancellation.clone();
940
941 let started_probe = Arc::clone(&started);
944 tokio::spawn(async move {
945 while started_probe.load(Ordering::SeqCst) == 0 {
946 tokio::task::yield_now().await;
947 }
948 cancellation.cancel();
949 });
950
951 let result = connector
952 .execute("wait", json!({}), &context(control))
953 .await;
954
955 assert_eq!(started.load(Ordering::SeqCst), 1, "the call never started");
956 assert_eq!(result, Err("Call cancelled".to_string()));
957 }
958
959 #[tokio::test]
960 async fn an_uncancelled_call_is_unaffected() {
961 struct Echo;
962
963 #[async_trait]
964 impl McpClient for Echo {
965 async fn list_tools(&self) -> Result<Vec<McpTool>, String> {
966 Ok(vec![McpTool {
967 name: "echo".to_string(),
968 description: None,
969 input_schema: json!({"type": "object"}),
970 output_schema: None,
971 annotations: None,
972 }])
973 }
974
975 async fn call_tool(&self, _name: &str, arguments: Value) -> Result<Value, String> {
976 Ok(arguments)
977 }
978 }
979
980 let connector = McpConnector::new("echo", Arc::new(Echo));
981 let result = connector
982 .execute(
983 "echo",
984 json!({"v": 1}),
985 &context(ToolCallControl::default()),
986 )
987 .await;
988
989 assert_eq!(result, Ok(json!({"v": 1})));
990 }
991}