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
174pub fn required_str<'a>(input: &'a Value, field: &str) -> std::result::Result<&'a str, ToolError> {
176 input.get(field).and_then(Value::as_str).ok_or_else(|| {
177 let provided: Vec<&str> = input
180 .as_object()
181 .map(|obj| obj.keys().map(|k| k.as_str()).collect())
182 .unwrap_or_default();
183 if provided.is_empty() {
184 ToolError::missing_field(field)
185 } else {
186 let hint = format!(
187 "missing required field '{field}'. Input provided: {}",
188 provided.join(", ")
189 );
190 ToolError::invalid_input(hint)
191 }
192 })
193}
194
195#[must_use]
197pub fn optional_str<'a>(input: &'a Value, field: &str) -> Option<&'a str> {
198 input.get(field).and_then(Value::as_str)
199}
200
201pub fn required_u64(input: &Value, field: &str) -> std::result::Result<u64, ToolError> {
203 input
204 .get(field)
205 .and_then(Value::as_u64)
206 .ok_or_else(|| ToolError::missing_field(field))
207}
208
209#[must_use]
211pub fn optional_u64(input: &Value, field: &str, default: u64) -> u64 {
212 input.get(field).and_then(Value::as_u64).unwrap_or(default)
213}
214
215#[must_use]
217pub fn optional_bool(input: &Value, field: &str, default: bool) -> bool {
218 input.get(field).and_then(Value::as_bool).unwrap_or(default)
219}
220
221#[derive(Debug, Clone, Serialize, Deserialize)]
226pub struct ToolDescriptor {
227 pub name: String,
229 pub input_schema: Value,
231 pub output_schema: Value,
233 pub supports_parallel_tool_calls: bool,
235 pub timeout_ms: Option<u64>,
237}
238
239#[derive(Debug, Clone, Serialize, Deserialize)]
244pub struct ConfiguredToolDescriptor {
245 pub spec: ToolDescriptor,
247 pub supports_parallel_tool_calls: bool,
249}
250
251#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
253#[serde(rename_all = "snake_case")]
254pub enum ToolCallSource {
255 Direct,
257 JsRepl,
259}
260
261#[derive(Debug, Clone, Serialize, Deserialize)]
266pub struct ToolCall {
267 pub name: String,
269 pub payload: ToolPayload,
271 pub source: ToolCallSource,
273 pub raw_tool_call_id: Option<String>,
275}
276
277impl ToolCall {
278 pub fn execution_subject(&self, fallback_cwd: &str) -> (String, String, &'static str) {
285 match &self.payload {
286 ToolPayload::LocalShell { params } => (
287 params.command.clone(),
288 params
289 .cwd
290 .clone()
291 .unwrap_or_else(|| fallback_cwd.to_string()),
292 "shell",
293 ),
294 _ => (self.name.clone(), fallback_cwd.to_string(), "tool"),
295 }
296 }
297}
298
299#[derive(Debug, Clone)]
304pub struct ToolInvocation {
305 pub call_id: String,
307 pub tool_name: String,
309 pub payload: ToolPayload,
311 pub source: ToolCallSource,
313}
314
315#[derive(Debug, Clone, Serialize, Deserialize)]
323pub enum FunctionCallError {
324 ToolNotFound { name: String },
326 KindMismatch { expected: ToolKind, got: ToolKind },
328 MutatingToolRejected { name: String },
330 TimedOut { name: String, timeout_ms: u64 },
332 Cancelled { name: String },
334 ExecutionFailed { name: String, error: String },
336}
337
338#[async_trait]
343pub trait ToolHandler: Send + Sync {
344 fn kind(&self) -> ToolKind;
346
347 fn matches_kind(&self, kind: ToolKind) -> bool {
351 self.kind() == kind
352 }
353
354 fn is_mutating(&self) -> bool {
358 false
359 }
360
361 async fn handle(
363 &self,
364 invocation: ToolInvocation,
365 ) -> std::result::Result<ToolOutput, FunctionCallError>;
366}
367
368#[derive(Debug)]
374pub struct ToolCallRuntime {
375 execution_lock: Arc<RwLock<()>>,
376}
377
378impl Default for ToolCallRuntime {
379 fn default() -> Self {
380 Self {
381 execution_lock: Arc::new(RwLock::new(())),
382 }
383 }
384}
385
386#[derive(Debug)]
387enum ToolExecutionGuard {
388 Parallel(#[allow(dead_code)] OwnedRwLockReadGuard<()>),
389 Serial(#[allow(dead_code)] OwnedRwLockWriteGuard<()>),
390 Reentrant,
391}
392
393impl ToolCallRuntime {
394 async fn acquire(&self, supports_parallel: bool) -> ToolExecutionGuard {
395 if TOOL_EXECUTION_LOCK_HELD.try_with(|_| ()).is_ok() {
396 return ToolExecutionGuard::Reentrant;
397 }
398
399 if supports_parallel {
400 ToolExecutionGuard::Parallel(self.execution_lock.clone().read_owned().await)
401 } else {
402 ToolExecutionGuard::Serial(self.execution_lock.clone().write_owned().await)
403 }
404 }
405}
406
407#[derive(Default)]
413pub struct ToolRegistry {
414 handlers: HashMap<String, Arc<dyn ToolHandler>>,
415 specs: HashMap<String, ConfiguredToolDescriptor>,
416 runtime: ToolCallRuntime,
417}
418
419impl ToolRegistry {
420 pub fn register(&mut self, spec: ToolDescriptor, handler: Arc<dyn ToolHandler>) -> Result<()> {
426 let name = spec.name.clone();
427 self.specs.insert(
428 name.clone(),
429 ConfiguredToolDescriptor {
430 supports_parallel_tool_calls: spec.supports_parallel_tool_calls,
431 spec,
432 },
433 );
434 self.handlers.insert(name, handler);
435 Ok(())
436 }
437
438 pub fn list_specs(&self) -> Vec<ConfiguredToolDescriptor> {
440 self.specs.values().cloned().collect()
441 }
442
443 pub async fn dispatch(
451 &self,
452 call: ToolCall,
453 allow_mutating: bool,
454 ) -> std::result::Result<ToolOutput, FunctionCallError> {
455 let handler = self.handlers.get(&call.name).cloned().ok_or_else(|| {
456 FunctionCallError::ToolNotFound {
457 name: call.name.clone(),
458 }
459 })?;
460 let configured =
461 self.specs
462 .get(&call.name)
463 .cloned()
464 .ok_or_else(|| FunctionCallError::ToolNotFound {
465 name: call.name.clone(),
466 })?;
467
468 let payload_kind = tool_payload_kind(&call.payload);
469 let expected = handler.kind();
470 if !handler.matches_kind(payload_kind) {
471 return Err(FunctionCallError::KindMismatch {
472 expected,
473 got: payload_kind,
474 });
475 }
476 if handler.is_mutating() && !allow_mutating {
477 return Err(FunctionCallError::MutatingToolRejected { name: call.name });
478 }
479
480 let invocation = ToolInvocation {
481 call_id: call
482 .raw_tool_call_id
483 .clone()
484 .unwrap_or_else(|| format!("tool-call-{}", uuid::Uuid::new_v4())),
485 tool_name: call.name.clone(),
486 payload: call.payload,
487 source: call.source,
488 };
489
490 let _guard = self
491 .runtime
492 .acquire(configured.supports_parallel_tool_calls)
493 .await;
494
495 TOOL_EXECUTION_LOCK_HELD
496 .scope(
497 (),
498 self.execute_with_timeout(handler, configured.spec.timeout_ms, invocation),
499 )
500 .await
501 }
502
503 async fn execute_with_timeout(
504 &self,
505 handler: Arc<dyn ToolHandler>,
506 timeout_ms: Option<u64>,
507 invocation: ToolInvocation,
508 ) -> std::result::Result<ToolOutput, FunctionCallError> {
509 if let Some(timeout_ms) = timeout_ms {
510 let name = invocation.tool_name.clone();
511 match tokio::time::timeout(
512 Duration::from_millis(timeout_ms),
513 handler.handle(invocation),
514 )
515 .await
516 {
517 Ok(result) => result,
518 Err(_) => Err(FunctionCallError::TimedOut { name, timeout_ms }),
519 }
520 } else {
521 handler.handle(invocation).await
522 }
523 }
524}
525
526fn tool_payload_kind(payload: &ToolPayload) -> ToolKind {
527 match payload {
528 ToolPayload::Mcp { .. } => ToolKind::Mcp,
529 ToolPayload::Function { .. }
530 | ToolPayload::Custom { .. }
531 | ToolPayload::LocalShell { .. } => ToolKind::Function,
532 }
533}
534
535#[cfg(test)]
536mod tests {
537 use serde_json::json;
538
539 use super::*;
540
541 #[test]
542 fn tool_result_success_sets_plain_content() {
543 let content = "operation completed successfully";
544 let result = ToolResult::success(content);
545
546 assert!(result.success);
547 assert_eq!(result.content, content);
548 assert!(result.metadata.is_none());
549 }
550
551 #[test]
552 fn tool_result_json_round_trips_content() {
553 let result = ToolResult::json(&json!({"ok": true})).expect("json");
554 assert!(result.success);
555 let content: serde_json::Value =
556 serde_json::from_str(&result.content).expect("content is valid json");
557 assert_eq!(content, json!({"ok": true}));
558 }
559
560 #[test]
561 fn helper_extractors_validate_shape() {
562 let input = json!({"name": "demo", "count": 7, "enabled": true});
563 assert_eq!(required_str(&input, "name").expect("name"), "demo");
564 assert_eq!(optional_str(&input, "name"), Some("demo"));
565 assert_eq!(optional_str(&input, "missing"), None);
566 assert_eq!(optional_str(&input, "count"), None);
567 assert_eq!(optional_str(&json!({"name": null}), "name"), None);
568 assert_eq!(optional_u64(&input, "count", 0), 7);
569 assert!(optional_bool(&input, "enabled", false));
570 assert!(matches!(
571 required_u64(&input, "name"),
572 Err(ToolError::MissingField { .. })
573 ));
574 }
575
576 #[test]
577 fn required_u64_rejects_missing_or_non_integer_values() {
578 assert!(matches!(
579 required_u64(&json!({}), "count"),
580 Err(ToolError::MissingField { .. })
581 ));
582 assert_eq!(required_u64(&json!({"count": 42}), "count").unwrap(), 42);
583 assert_eq!(
584 required_u64(&json!({"count": u64::MAX}), "count").unwrap(),
585 u64::MAX
586 );
587
588 for value in [json!(-1), json!(2.5), json!("42")] {
589 assert!(matches!(
590 required_u64(&json!({"count": value}), "count"),
591 Err(ToolError::MissingField { .. })
592 ));
593 }
594 }
595
596 #[test]
597 fn required_str_reports_provided_fields_on_missing_required_field() {
598 let input = json!({"path": "src/lib.rs", "content": "new body"});
599 let err = required_str(&input, "replace").expect_err("replace is missing");
600 let message = err.to_string();
601 assert!(message.contains("missing required field 'replace'"));
602 assert!(message.contains("Input provided:"));
603 assert!(message.contains("path"));
604 assert!(message.contains("content"));
605 }
606
607 #[test]
608 fn tool_error_display_matches_legacy_text() {
609 let err = ToolError::missing_field("path");
610 assert_eq!(
611 err.to_string(),
612 "Failed to validate input: missing required field 'path'"
613 );
614 }
615
616 #[test]
617 fn tool_error_missing_field_constructor() {
618 let err = ToolError::missing_field("my_field");
619 assert!(matches!(err, ToolError::MissingField { field } if field == "my_field"));
620 }
621
622 #[test]
623 fn tool_error_not_available_displays_reason() {
624 let err = ToolError::not_available("custom tool not found");
625
626 assert!(matches!(err, ToolError::NotAvailable { .. }));
627 assert_eq!(
628 err.to_string(),
629 "Failed to locate tool: custom tool not found"
630 );
631 }
632
633 #[test]
634 fn tool_error_permission_denied_displays_reason() {
635 let err = ToolError::permission_denied("unauthorized user");
636
637 assert!(matches!(err, ToolError::PermissionDenied { .. }));
638 assert_eq!(
639 err.to_string(),
640 "Failed to authorize tool execution: unauthorized user"
641 );
642 }
643
644 #[test]
645 fn tool_error_execution_failed_displays_reason() {
646 let err = ToolError::execution_failed("process crashed");
647
648 assert!(
649 matches!(err, ToolError::ExecutionFailed { ref message } if message == "process crashed")
650 );
651 assert_eq!(err.to_string(), "Failed to execute tool: process crashed");
652 }
653
654 #[test]
655 fn tool_error_invalid_input_creates_correct_variant() {
656 let err = ToolError::invalid_input("test invalid message");
657 match err {
658 ToolError::InvalidInput { message } => {
659 assert_eq!(message, "test invalid message");
660 }
661 _ => panic!("Expected ToolError::InvalidInput, got {err:?}"),
662 }
663 }
664
665 #[test]
666 fn tool_error_path_escape_display() {
667 let path = std::path::PathBuf::from("../outside");
668 let err = ToolError::path_escape(path);
669 assert_eq!(
670 err.to_string(),
671 "Failed to resolve path '../outside': path escapes workspace"
672 );
673 }
674
675 #[test]
676 fn tool_call_execution_subject_uses_local_shell_command_and_cwd() {
677 let call = ToolCall {
678 name: "shell".to_string(),
679 payload: ToolPayload::LocalShell {
680 params: codewhale_protocol::LocalShellParams {
681 command: "ls -l".to_string(),
682 cwd: Some("/custom/dir".to_string()),
683 timeout_ms: None,
684 },
685 },
686 source: ToolCallSource::Direct,
687 raw_tool_call_id: None,
688 };
689
690 assert_eq!(
691 call.execution_subject("/fallback/dir"),
692 ("ls -l".to_string(), "/custom/dir".to_string(), "shell")
693 );
694 }
695
696 #[test]
697 fn tool_call_execution_subject_falls_back_for_shell_without_cwd() {
698 let call = ToolCall {
699 name: "shell".to_string(),
700 payload: ToolPayload::LocalShell {
701 params: codewhale_protocol::LocalShellParams {
702 command: "echo hello".to_string(),
703 cwd: None,
704 timeout_ms: None,
705 },
706 },
707 source: ToolCallSource::Direct,
708 raw_tool_call_id: None,
709 };
710
711 assert_eq!(
712 call.execution_subject("/fallback/dir"),
713 (
714 "echo hello".to_string(),
715 "/fallback/dir".to_string(),
716 "shell"
717 )
718 );
719 }
720
721 #[test]
722 fn tool_call_execution_subject_uses_tool_name_for_non_shell_payloads() {
723 let call = ToolCall {
724 name: "my_tool".to_string(),
725 payload: ToolPayload::Function {
726 arguments: "{}".to_string(),
727 },
728 source: ToolCallSource::Direct,
729 raw_tool_call_id: None,
730 };
731
732 assert_eq!(
733 call.execution_subject("/fallback/dir"),
734 ("my_tool".to_string(), "/fallback/dir".to_string(), "tool")
735 );
736 }
737}