1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::sync::Arc;
4use std::time::Duration;
5
6use anyhow::Result;
7use async_trait::async_trait;
8use codewhale_protocol::{ToolKind, ToolOutput, ToolPayload};
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};
12
13mod outcome;
14mod prepared;
15mod resources;
16
17pub use outcome::{ToolExecutionOutcome, ToolTerminalStatus};
18pub use prepared::PreparedToolCall;
19pub use resources::{ResourceClaim, schedule_non_conflicting};
20
21tokio::task_local! {
22 static TOOL_EXECUTION_LOCK_HELD: ();
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
27pub enum ToolCapability {
28 ReadOnly,
30 WritesFiles,
32 ExecutesCode,
34 Network,
36 Sandboxable,
38 RequiresApproval,
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
44pub enum ApprovalRequirement {
45 #[default]
47 Auto,
48 Suggest,
50 Required,
52}
53
54#[derive(Debug, Clone, thiserror::Error)]
56pub enum ToolError {
57 #[error("Failed to validate input: {message}")]
58 InvalidInput { message: String },
59 #[error("Failed to validate input: missing required field '{field}'")]
60 MissingField { field: String },
61 #[error("Failed to resolve path '{}': path escapes workspace", path.display())]
62 PathEscape { path: PathBuf },
63 #[error("Failed to execute tool: {message}")]
64 ExecutionFailed { message: String },
65 #[error("Failed to execute tool: operation timed out after {seconds}s")]
66 Timeout { seconds: u64 },
67 #[error("Tool execution cancelled: {message}")]
68 Cancelled { message: String },
69 #[error("Failed to locate tool: {message}")]
70 NotAvailable { message: String },
71 #[error("Failed to authorize tool execution: {message}")]
72 PermissionDenied { message: String },
73}
74
75impl ToolError {
76 #[must_use]
77 pub fn invalid_input(msg: impl Into<String>) -> Self {
78 Self::InvalidInput {
79 message: msg.into(),
80 }
81 }
82
83 #[must_use]
84 pub fn missing_field(field: impl Into<String>) -> Self {
85 Self::MissingField {
86 field: field.into(),
87 }
88 }
89
90 #[must_use]
91 pub fn execution_failed(msg: impl Into<String>) -> Self {
92 Self::ExecutionFailed {
93 message: msg.into(),
94 }
95 }
96
97 #[must_use]
98 pub fn cancelled(msg: impl Into<String>) -> Self {
99 Self::Cancelled {
100 message: msg.into(),
101 }
102 }
103
104 #[must_use]
105 pub fn path_escape(path: impl Into<PathBuf>) -> Self {
106 Self::PathEscape { path: path.into() }
107 }
108
109 #[must_use]
110 pub fn not_available(msg: impl Into<String>) -> Self {
111 Self::NotAvailable {
112 message: msg.into(),
113 }
114 }
115
116 #[must_use]
117 pub fn permission_denied(msg: impl Into<String>) -> Self {
118 Self::PermissionDenied {
119 message: msg.into(),
120 }
121 }
122}
123
124#[derive(Debug, Clone, Serialize, Deserialize)]
126pub struct ToolResult {
127 pub content: String,
129 pub success: bool,
131 #[serde(skip_serializing_if = "Option::is_none")]
133 pub metadata: Option<Value>,
134}
135
136impl ToolResult {
137 #[must_use]
139 pub fn success(content: impl Into<String>) -> Self {
140 Self {
141 content: content.into(),
142 success: true,
143 metadata: None,
144 }
145 }
146
147 #[must_use]
149 pub fn error(message: impl Into<String>) -> Self {
150 Self {
151 content: message.into(),
152 success: false,
153 metadata: None,
154 }
155 }
156
157 pub fn json<T: Serialize>(value: &T) -> std::result::Result<Self, serde_json::Error> {
159 Ok(Self {
160 content: serde_json::to_string(value)?,
161 success: true,
162 metadata: None,
163 })
164 }
165
166 #[must_use]
168 pub fn with_metadata(mut self, metadata: Value) -> Self {
169 self.metadata = Some(metadata);
170 self
171 }
172}
173
174#[must_use]
176pub fn json_type_name(value: &Value) -> &'static str {
177 match value {
178 Value::Null => "null",
179 Value::Bool(_) => "boolean",
180 Value::Number(_) => "number",
181 Value::String(_) => "string",
182 Value::Array(_) => "array",
183 Value::Object(_) => "object",
184 }
185}
186
187#[must_use]
190pub fn value_preview(value: &Value) -> String {
191 let preview = value.to_string();
192 if preview.chars().count() > 120 {
193 preview.chars().take(117).collect::<String>() + "..."
194 } else {
195 preview
196 }
197}
198
199#[must_use]
205pub fn type_mismatch(field: &str, value: &Value, expected: &str) -> ToolError {
206 ToolError::invalid_input(format!(
207 "field '{field}' must be {expected}; got {}. Received: {}",
208 json_type_name(value),
209 value_preview(value)
210 ))
211}
212
213fn is_absent(value: Option<&Value>) -> bool {
221 matches!(value, None | Some(Value::Null))
222}
223
224pub fn required_str<'a>(input: &'a Value, field: &str) -> std::result::Result<&'a str, ToolError> {
226 if let Some(value) = input.get(field) {
227 if let Some(string_value) = value.as_str() {
228 return Ok(string_value);
229 }
230
231 return Err(type_mismatch(field, value, "a string"));
232 }
233
234 let provided: Vec<&str> = input
237 .as_object()
238 .map(|obj| obj.keys().map(|k| k.as_str()).collect())
239 .unwrap_or_default();
240 if provided.is_empty() {
241 Err(ToolError::missing_field(field))
242 } else {
243 let hint = format!(
244 "missing required field '{field}'. Input provided: {}",
245 provided.join(", ")
246 );
247 Err(ToolError::invalid_input(hint))
248 }
249}
250
251pub fn optional_str<'a>(
256 input: &'a Value,
257 field: &str,
258) -> std::result::Result<Option<&'a str>, ToolError> {
259 let value = input.get(field);
260 if is_absent(value) {
261 return Ok(None);
262 }
263 let value = value.expect("is_absent covers the None case");
264 value
265 .as_str()
266 .map(Some)
267 .ok_or_else(|| type_mismatch(field, value, "a string"))
268}
269
270pub fn required_u64(input: &Value, field: &str) -> std::result::Result<u64, ToolError> {
277 let value = input.get(field);
278 if is_absent(value) {
279 return Err(ToolError::missing_field(field));
280 }
281 let value = value.expect("is_absent covers the None case");
282 value
283 .as_u64()
284 .ok_or_else(|| type_mismatch(field, value, "a non-negative integer"))
285}
286
287pub fn optional_u64(
291 input: &Value,
292 field: &str,
293 default: u64,
294) -> std::result::Result<u64, ToolError> {
295 let value = input.get(field);
296 if is_absent(value) {
297 return Ok(default);
298 }
299 let value = value.expect("is_absent covers the None case");
300 value
301 .as_u64()
302 .ok_or_else(|| type_mismatch(field, value, "a non-negative integer"))
303}
304
305pub fn optional_bool(
313 input: &Value,
314 field: &str,
315 default: bool,
316) -> std::result::Result<bool, ToolError> {
317 Ok(optional_bool_opt(input, field)?.unwrap_or(default))
318}
319
320pub fn optional_bool_opt(
327 input: &Value,
328 field: &str,
329) -> std::result::Result<Option<bool>, ToolError> {
330 let value = input.get(field);
331 if is_absent(value) {
332 return Ok(None);
333 }
334 let value = value.expect("is_absent covers the None case");
335 value
336 .as_bool()
337 .map(Some)
338 .ok_or_else(|| type_mismatch(field, value, "a boolean"))
339}
340
341#[derive(Debug, Clone, Serialize, Deserialize)]
346pub struct ToolDescriptor {
347 pub name: String,
349 pub input_schema: Value,
351 pub output_schema: Value,
353 pub supports_parallel_tool_calls: bool,
355 pub timeout_ms: Option<u64>,
357}
358
359#[derive(Debug, Clone, Serialize, Deserialize)]
364pub struct ConfiguredToolDescriptor {
365 pub spec: ToolDescriptor,
367 pub supports_parallel_tool_calls: bool,
369}
370
371#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
373#[serde(rename_all = "snake_case")]
374pub enum ToolCallSource {
375 Direct,
377 JsRepl,
379}
380
381#[derive(Debug, Clone, Serialize, Deserialize)]
386pub struct ToolCall {
387 pub name: String,
389 pub payload: ToolPayload,
391 pub source: ToolCallSource,
393 pub raw_tool_call_id: Option<String>,
395}
396
397impl ToolCall {
398 pub fn execution_subject(&self, fallback_cwd: &str) -> (String, String, &'static str) {
405 match &self.payload {
406 ToolPayload::LocalShell { params } => (
407 params.command.clone(),
408 params
409 .cwd
410 .clone()
411 .unwrap_or_else(|| fallback_cwd.to_string()),
412 "shell",
413 ),
414 _ => (self.name.clone(), fallback_cwd.to_string(), "tool"),
415 }
416 }
417}
418
419#[derive(Debug, Clone)]
424pub struct ToolInvocation {
425 pub call_id: String,
427 pub tool_name: String,
429 pub payload: ToolPayload,
431 pub source: ToolCallSource,
433}
434
435#[derive(Debug, Clone, Serialize, Deserialize)]
443pub enum FunctionCallError {
444 ToolNotFound { name: String },
446 KindMismatch { expected: ToolKind, got: ToolKind },
448 MutatingToolRejected { name: String },
450 TimedOut { name: String, timeout_ms: u64 },
452 Cancelled { name: String },
454 ExecutionFailed { name: String, error: String },
456}
457
458#[async_trait]
463pub trait ToolHandler: Send + Sync {
464 fn kind(&self) -> ToolKind;
466
467 fn matches_kind(&self, kind: ToolKind) -> bool {
471 self.kind() == kind
472 }
473
474 fn is_mutating(&self) -> bool {
478 false
479 }
480
481 async fn handle(
483 &self,
484 invocation: ToolInvocation,
485 ) -> std::result::Result<ToolOutput, FunctionCallError>;
486}
487
488#[derive(Debug)]
494pub struct ToolCallRuntime {
495 execution_lock: Arc<RwLock<()>>,
496}
497
498impl Default for ToolCallRuntime {
499 fn default() -> Self {
500 Self {
501 execution_lock: Arc::new(RwLock::new(())),
502 }
503 }
504}
505
506#[derive(Debug)]
507enum ToolExecutionGuard {
508 Parallel(#[allow(dead_code)] OwnedRwLockReadGuard<()>),
509 Serial(#[allow(dead_code)] OwnedRwLockWriteGuard<()>),
510 Reentrant,
511}
512
513impl ToolCallRuntime {
514 async fn acquire(&self, supports_parallel: bool) -> ToolExecutionGuard {
515 if TOOL_EXECUTION_LOCK_HELD.try_with(|_| ()).is_ok() {
516 return ToolExecutionGuard::Reentrant;
517 }
518
519 if supports_parallel {
520 ToolExecutionGuard::Parallel(self.execution_lock.clone().read_owned().await)
521 } else {
522 ToolExecutionGuard::Serial(self.execution_lock.clone().write_owned().await)
523 }
524 }
525}
526
527#[derive(Default)]
533pub struct ToolRegistry {
534 handlers: HashMap<String, Arc<dyn ToolHandler>>,
535 specs: HashMap<String, ConfiguredToolDescriptor>,
536 runtime: ToolCallRuntime,
537}
538
539impl ToolRegistry {
540 pub fn register(&mut self, spec: ToolDescriptor, handler: Arc<dyn ToolHandler>) -> Result<()> {
546 let name = spec.name.clone();
547 self.specs.insert(
548 name.clone(),
549 ConfiguredToolDescriptor {
550 supports_parallel_tool_calls: spec.supports_parallel_tool_calls,
551 spec,
552 },
553 );
554 self.handlers.insert(name, handler);
555 Ok(())
556 }
557
558 pub fn list_specs(&self) -> Vec<ConfiguredToolDescriptor> {
560 self.specs.values().cloned().collect()
561 }
562
563 pub async fn dispatch(
571 &self,
572 call: ToolCall,
573 allow_mutating: bool,
574 ) -> std::result::Result<ToolOutput, FunctionCallError> {
575 let handler = self.handlers.get(&call.name).cloned().ok_or_else(|| {
576 FunctionCallError::ToolNotFound {
577 name: call.name.clone(),
578 }
579 })?;
580 let configured =
581 self.specs
582 .get(&call.name)
583 .cloned()
584 .ok_or_else(|| FunctionCallError::ToolNotFound {
585 name: call.name.clone(),
586 })?;
587
588 let payload_kind = tool_payload_kind(&call.payload);
589 let expected = handler.kind();
590 if !handler.matches_kind(payload_kind) {
591 return Err(FunctionCallError::KindMismatch {
592 expected,
593 got: payload_kind,
594 });
595 }
596 if handler.is_mutating() && !allow_mutating {
597 return Err(FunctionCallError::MutatingToolRejected { name: call.name });
598 }
599
600 let invocation = ToolInvocation {
601 call_id: call
602 .raw_tool_call_id
603 .clone()
604 .unwrap_or_else(|| format!("tool-call-{}", uuid::Uuid::new_v4())),
605 tool_name: call.name.clone(),
606 payload: call.payload,
607 source: call.source,
608 };
609
610 let _guard = self
611 .runtime
612 .acquire(configured.supports_parallel_tool_calls)
613 .await;
614
615 TOOL_EXECUTION_LOCK_HELD
616 .scope(
617 (),
618 self.execute_with_timeout(handler, configured.spec.timeout_ms, invocation),
619 )
620 .await
621 }
622
623 async fn execute_with_timeout(
624 &self,
625 handler: Arc<dyn ToolHandler>,
626 timeout_ms: Option<u64>,
627 invocation: ToolInvocation,
628 ) -> std::result::Result<ToolOutput, FunctionCallError> {
629 if let Some(timeout_ms) = timeout_ms {
630 let name = invocation.tool_name.clone();
631 match tokio::time::timeout(
632 Duration::from_millis(timeout_ms),
633 handler.handle(invocation),
634 )
635 .await
636 {
637 Ok(result) => result,
638 Err(_) => Err(FunctionCallError::TimedOut { name, timeout_ms }),
639 }
640 } else {
641 handler.handle(invocation).await
642 }
643 }
644}
645
646fn tool_payload_kind(payload: &ToolPayload) -> ToolKind {
647 match payload {
648 ToolPayload::Mcp { .. } => ToolKind::Mcp,
649 ToolPayload::Function { .. }
650 | ToolPayload::Custom { .. }
651 | ToolPayload::LocalShell { .. } => ToolKind::Function,
652 }
653}
654
655#[cfg(test)]
656mod tests {
657 use serde_json::json;
658
659 use super::*;
660
661 #[test]
662 fn tool_result_success_sets_plain_content() {
663 let content = "operation completed successfully";
664 let result = ToolResult::success(content);
665
666 assert!(result.success);
667 assert_eq!(result.content, content);
668 assert!(result.metadata.is_none());
669 }
670
671 #[test]
672 fn tool_result_json_round_trips_content() {
673 let result = ToolResult::json(&json!({"ok": true})).expect("json");
674 assert!(result.success);
675 let content: serde_json::Value =
676 serde_json::from_str(&result.content).expect("content is valid json");
677 assert_eq!(content, json!({"ok": true}));
678 }
679
680 #[test]
681 fn helper_extractors_validate_shape() {
682 let input = json!({"name": "demo", "count": 7, "enabled": true});
683 assert_eq!(required_str(&input, "name").expect("name"), "demo");
684 assert_eq!(optional_str(&input, "name").unwrap(), Some("demo"));
685 assert_eq!(optional_str(&input, "missing").unwrap(), None);
686 assert_eq!(optional_str(&json!({"name": null}), "name").unwrap(), None);
687 assert_eq!(optional_u64(&input, "count", 0).unwrap(), 7);
688 assert!(optional_bool(&input, "enabled", false).unwrap());
689 let err = required_u64(&input, "name")
692 .expect_err("a present string is not a missing u64")
693 .to_string();
694 assert!(
695 err.contains("field 'name' must be a non-negative integer"),
696 "{err}"
697 );
698 }
699
700 #[test]
704 fn optional_extractors_refuse_type_mismatches_instead_of_defaulting() {
705 let err = optional_bool(&json!({"dry_run": "true"}), "dry_run", false)
709 .expect_err("a stringy bool must not become the default")
710 .to_string();
711 assert!(err.contains("dry_run"), "{err}");
712 assert!(err.contains("must be a boolean"), "{err}");
713 assert!(err.contains("got string"), "{err}");
714 assert!(err.contains("\"true\""), "{err}");
715
716 for bad in [json!("true"), json!(1), json!(0), json!([]), json!({})] {
717 assert!(
718 optional_bool(&json!({"flag": bad}), "flag", false).is_err(),
719 "optional_bool accepted {bad}"
720 );
721 }
722 for bad in [json!("7"), json!(-1), json!(1.5), json!(true), json!([7])] {
723 assert!(
724 optional_u64(&json!({"n": bad}), "n", 42).is_err(),
725 "optional_u64 accepted {bad}"
726 );
727 }
728 for bad in [json!(7), json!(true), json!(["a"]), json!({"a": 1})] {
729 assert!(
730 optional_str(&json!({"s": bad}), "s").is_err(),
731 "optional_str accepted {bad}"
732 );
733 }
734
735 assert!(optional_bool(&json!({"flag": null}), "flag", true).unwrap());
737 assert_eq!(optional_u64(&json!({"n": null}), "n", 42).unwrap(), 42);
738 assert_eq!(optional_str(&json!({"s": null}), "s").unwrap(), None);
739 }
740
741 #[test]
742 fn type_mismatch_truncates_a_huge_offending_value() {
743 let big = Value::String("x".repeat(500));
744 let err = type_mismatch("body", &big, "a boolean").to_string();
745 assert!(err.contains("body"), "{err}");
746 assert!(err.ends_with("..."), "{err}");
747 assert!(err.chars().count() < 250, "{err}");
748 }
749
750 #[test]
751 fn required_u64_distinguishes_missing_from_type_mismatch() {
752 assert!(matches!(
754 required_u64(&json!({}), "count"),
755 Err(ToolError::MissingField { .. })
756 ));
757 assert!(matches!(
758 required_u64(&json!({"count": null}), "count"),
759 Err(ToolError::MissingField { .. })
760 ));
761
762 assert_eq!(required_u64(&json!({"count": 42}), "count").unwrap(), 42);
764 assert_eq!(
765 required_u64(&json!({"count": u64::MAX}), "count").unwrap(),
766 u64::MAX
767 );
768
769 for value in [json!(-1), json!(2.5), json!("42")] {
772 let err = required_u64(&json!({"count": value}), "count")
773 .expect_err("wrong type must not look missing")
774 .to_string();
775 assert!(
776 err.contains("field 'count' must be a non-negative integer"),
777 "{err}"
778 );
779 }
780 }
781
782 #[test]
783 fn required_str_reports_provided_fields_on_missing_required_field() {
784 let input = json!({"path": "src/lib.rs", "content": "new body"});
785 let err = required_str(&input, "replace").expect_err("replace is missing");
786 let message = err.to_string();
787 assert!(message.contains("missing required field 'replace'"));
788 assert!(message.contains("Input provided:"));
789 assert!(message.contains("path"));
790 assert!(message.contains("content"));
791 }
792
793 #[test]
794 fn required_str_reports_wrong_type_when_field_exists() {
795 let input = json!({"replace": [{"path": "src/lib.rs", "content": "new body"}]});
796 let err = required_str(&input, "replace").expect_err("replace has wrong type");
797 let message = err.to_string();
798 assert!(message.contains("field 'replace' must be a string"));
799 assert!(message.contains("got array"));
800 assert!(message.contains(r#""content":"new body""#));
801 assert!(message.contains(r#""path":"src/lib.rs""#));
802 }
803
804 #[test]
805 fn tool_error_display_matches_legacy_text() {
806 let err = ToolError::missing_field("path");
807 assert_eq!(
808 err.to_string(),
809 "Failed to validate input: missing required field 'path'"
810 );
811 }
812
813 #[test]
814 fn tool_error_missing_field_constructor() {
815 let err = ToolError::missing_field("my_field");
816 assert!(matches!(err, ToolError::MissingField { field } if field == "my_field"));
817 }
818
819 #[test]
820 fn tool_error_not_available_displays_reason() {
821 let err = ToolError::not_available("custom tool not found");
822
823 assert!(matches!(err, ToolError::NotAvailable { .. }));
824 assert_eq!(
825 err.to_string(),
826 "Failed to locate tool: custom tool not found"
827 );
828 }
829
830 #[test]
831 fn tool_error_permission_denied_displays_reason() {
832 let err = ToolError::permission_denied("unauthorized user");
833
834 assert!(matches!(err, ToolError::PermissionDenied { .. }));
835 assert_eq!(
836 err.to_string(),
837 "Failed to authorize tool execution: unauthorized user"
838 );
839 }
840
841 #[test]
842 fn tool_error_execution_failed_displays_reason() {
843 let err = ToolError::execution_failed("process crashed");
844
845 assert!(
846 matches!(err, ToolError::ExecutionFailed { ref message } if message == "process crashed")
847 );
848 assert_eq!(err.to_string(), "Failed to execute tool: process crashed");
849 }
850
851 #[test]
852 fn tool_error_invalid_input_creates_correct_variant() {
853 let err = ToolError::invalid_input("test invalid message");
854 match err {
855 ToolError::InvalidInput { message } => {
856 assert_eq!(message, "test invalid message");
857 }
858 _ => panic!("Expected ToolError::InvalidInput, got {err:?}"),
859 }
860 }
861
862 #[test]
863 fn tool_error_path_escape_display() {
864 let path = std::path::PathBuf::from("../outside");
865 let err = ToolError::path_escape(path);
866 assert_eq!(
867 err.to_string(),
868 "Failed to resolve path '../outside': path escapes workspace"
869 );
870 }
871
872 #[test]
873 fn tool_call_execution_subject_uses_local_shell_command_and_cwd() {
874 let call = ToolCall {
875 name: "shell".to_string(),
876 payload: ToolPayload::LocalShell {
877 params: codewhale_protocol::LocalShellParams {
878 command: "ls -l".to_string(),
879 cwd: Some("/custom/dir".to_string()),
880 timeout_ms: None,
881 },
882 },
883 source: ToolCallSource::Direct,
884 raw_tool_call_id: None,
885 };
886
887 assert_eq!(
888 call.execution_subject("/fallback/dir"),
889 ("ls -l".to_string(), "/custom/dir".to_string(), "shell")
890 );
891 }
892
893 #[test]
894 fn tool_call_execution_subject_falls_back_for_shell_without_cwd() {
895 let call = ToolCall {
896 name: "shell".to_string(),
897 payload: ToolPayload::LocalShell {
898 params: codewhale_protocol::LocalShellParams {
899 command: "echo hello".to_string(),
900 cwd: None,
901 timeout_ms: None,
902 },
903 },
904 source: ToolCallSource::Direct,
905 raw_tool_call_id: None,
906 };
907
908 assert_eq!(
909 call.execution_subject("/fallback/dir"),
910 (
911 "echo hello".to_string(),
912 "/fallback/dir".to_string(),
913 "shell"
914 )
915 );
916 }
917
918 #[test]
919 fn tool_call_execution_subject_uses_tool_name_for_non_shell_payloads() {
920 let call = ToolCall {
921 name: "my_tool".to_string(),
922 payload: ToolPayload::Function {
923 arguments: "{}".to_string(),
924 },
925 source: ToolCallSource::Direct,
926 raw_tool_call_id: None,
927 };
928
929 assert_eq!(
930 call.execution_subject("/fallback/dir"),
931 ("my_tool".to_string(), "/fallback/dir".to_string(), "tool")
932 );
933 }
934}