| 1 | use std::collections::HashMap; |
| 2 | use std::path::PathBuf; |
| 3 | use std::sync::Arc; |
| 4 | use std::time::Duration; |
| 5 | |
| 6 | use anyhow::Result; |
| 7 | use async_trait::async_trait; |
| 8 | use deepseek_protocol::{ToolKind, ToolOutput, ToolPayload}; |
| 9 | use serde::{Deserialize, Serialize}; |
| 10 | use serde_json::Value; |
| 11 | use tokio::sync::RwLock; |
| 12 | |
| 13 | /// Capabilities that a tool may have or require. |
| 14 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] |
| 15 | pub enum ToolCapability { |
| 16 | /// Tool only reads data, never modifies state. |
| 17 | ReadOnly, |
| 18 | /// Tool writes to the filesystem. |
| 19 | WritesFiles, |
| 20 | /// Tool executes arbitrary shell commands. |
| 21 | ExecutesCode, |
| 22 | /// Tool makes network requests. |
| 23 | Network, |
| 24 | /// Tool can be run in a sandbox. |
| 25 | Sandboxable, |
| 26 | /// Tool requires user approval before execution. |
| 27 | RequiresApproval, |
| 28 | } |
| 29 | |
| 30 | /// Approval requirement for a tool. |
| 31 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] |
| 32 | pub enum ApprovalRequirement { |
| 33 | /// Never needs approval: safe read-only operations. |
| 34 | #[default] |
| 35 | Auto, |
| 36 | /// Suggest approval but allow user to skip. |
| 37 | Suggest, |
| 38 | /// Always require explicit user approval. |
| 39 | Required, |
| 40 | } |
| 41 | |
| 42 | /// Errors that can occur during tool execution. |
| 43 | #[derive(Debug, Clone)] |
| 44 | pub enum ToolError { |
| 45 | InvalidInput { message: String }, |
| 46 | MissingField { field: String }, |
| 47 | PathEscape { path: PathBuf }, |
| 48 | ExecutionFailed { message: String }, |
| 49 | Timeout { seconds: u64 }, |
| 50 | NotAvailable { message: String }, |
| 51 | PermissionDenied { message: String }, |
| 52 | } |
| 53 | |
| 54 | impl std::fmt::Display for ToolError { |
| 55 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 56 | match self { |
| 57 | Self::InvalidInput { message } => { |
| 58 | write!(f, "Failed to validate input: {message}") |
| 59 | } |
| 60 | Self::MissingField { field } => { |
| 61 | write!( |
| 62 | f, |
| 63 | "Failed to validate input: missing required field '{field}'" |
| 64 | ) |
| 65 | } |
| 66 | Self::PathEscape { path } => { |
| 67 | write!( |
| 68 | f, |
| 69 | "Failed to resolve path '{}': path escapes workspace", |
| 70 | path.display() |
| 71 | ) |
| 72 | } |
| 73 | Self::ExecutionFailed { message } => { |
| 74 | write!(f, "Failed to execute tool: {message}") |
| 75 | } |
| 76 | Self::Timeout { seconds } => { |
| 77 | write!( |
| 78 | f, |
| 79 | "Failed to execute tool: operation timed out after {seconds}s" |
| 80 | ) |
| 81 | } |
| 82 | Self::NotAvailable { message } => { |
| 83 | write!(f, "Failed to locate tool: {message}") |
| 84 | } |
| 85 | Self::PermissionDenied { message } => { |
| 86 | write!(f, "Failed to authorize tool execution: {message}") |
| 87 | } |
| 88 | } |
| 89 | } |
| 90 | } |
| 91 | |
| 92 | impl std::error::Error for ToolError {} |
| 93 | |
| 94 | impl ToolError { |
| 95 | #[must_use] |
| 96 | pub fn invalid_input(msg: impl Into<String>) -> Self { |
| 97 | Self::InvalidInput { |
| 98 | message: msg.into(), |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | #[must_use] |
| 103 | pub fn missing_field(field: impl Into<String>) -> Self { |
| 104 | Self::MissingField { |
| 105 | field: field.into(), |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | #[must_use] |
| 110 | pub fn execution_failed(msg: impl Into<String>) -> Self { |
| 111 | Self::ExecutionFailed { |
| 112 | message: msg.into(), |
| 113 | } |
| 114 | } |
| 115 | |
| 116 | #[must_use] |
| 117 | pub fn path_escape(path: impl Into<PathBuf>) -> Self { |
| 118 | Self::PathEscape { path: path.into() } |
| 119 | } |
| 120 | |
| 121 | #[must_use] |
| 122 | pub fn not_available(msg: impl Into<String>) -> Self { |
| 123 | Self::NotAvailable { |
| 124 | message: msg.into(), |
| 125 | } |
| 126 | } |
| 127 | |
| 128 | #[must_use] |
| 129 | pub fn permission_denied(msg: impl Into<String>) -> Self { |
| 130 | Self::PermissionDenied { |
| 131 | message: msg.into(), |
| 132 | } |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | /// Result of a tool execution. |
| 137 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 138 | pub struct ToolResult { |
| 139 | /// The output content, which may be JSON or plain text. |
| 140 | pub content: String, |
| 141 | /// Whether the execution was successful. |
| 142 | pub success: bool, |
| 143 | /// Optional structured metadata. |
| 144 | #[serde(skip_serializing_if = "Option::is_none")] |
| 145 | pub metadata: Option<Value>, |
| 146 | } |
| 147 | |
| 148 | impl ToolResult { |
| 149 | /// Create a successful result with content. |
| 150 | #[must_use] |
| 151 | pub fn success(content: impl Into<String>) -> Self { |
| 152 | Self { |
| 153 | content: content.into(), |
| 154 | success: true, |
| 155 | metadata: None, |
| 156 | } |
| 157 | } |
| 158 | |
| 159 | /// Create an error result with message. |
| 160 | #[must_use] |
| 161 | pub fn error(message: impl Into<String>) -> Self { |
| 162 | Self { |
| 163 | content: message.into(), |
| 164 | success: false, |
| 165 | metadata: None, |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | /// Create a successful result from JSON. |
| 170 | pub fn json<T: Serialize>(value: &T) -> std::result::Result<Self, serde_json::Error> { |
| 171 | Ok(Self { |
| 172 | content: serde_json::to_string_pretty(value)?, |
| 173 | success: true, |
| 174 | metadata: None, |
| 175 | }) |
| 176 | } |
| 177 | |
| 178 | /// Add metadata to the result. |
| 179 | #[must_use] |
| 180 | pub fn with_metadata(mut self, metadata: Value) -> Self { |
| 181 | self.metadata = Some(metadata); |
| 182 | self |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | /// Helper to extract a required string field from JSON input. |
| 187 | pub fn required_str<'a>(input: &'a Value, field: &str) -> std::result::Result<&'a str, ToolError> { |
| 188 | input.get(field).and_then(Value::as_str).ok_or_else(|| { |
| 189 | // When the field is missing, list the fields the caller *did* |
| 190 | // supply so the model can spot the mismatch without a retry. |
| 191 | let provided: Vec<&str> = input |
| 192 | .as_object() |
| 193 | .map(|obj| obj.keys().map(|k| k.as_str()).collect()) |
| 194 | .unwrap_or_default(); |
| 195 | if provided.is_empty() { |
| 196 | ToolError::missing_field(field) |
| 197 | } else { |
| 198 | let hint = format!( |
| 199 | "missing required field '{field}'. Input provided: {}", |
| 200 | provided.join(", ") |
| 201 | ); |
| 202 | ToolError::invalid_input(hint) |
| 203 | } |
| 204 | }) |
| 205 | } |
| 206 | |
| 207 | /// Helper to extract an optional string field from JSON input. |
| 208 | #[must_use] |
| 209 | pub fn optional_str<'a>(input: &'a Value, field: &str) -> Option<&'a str> { |
| 210 | input.get(field).and_then(Value::as_str) |
| 211 | } |
| 212 | |
| 213 | /// Helper to extract a required u64 field from JSON input. |
| 214 | pub fn required_u64(input: &Value, field: &str) -> std::result::Result<u64, ToolError> { |
| 215 | input |
| 216 | .get(field) |
| 217 | .and_then(Value::as_u64) |
| 218 | .ok_or_else(|| ToolError::missing_field(field)) |
| 219 | } |
| 220 | |
| 221 | /// Helper to extract an optional u64 field with default. |
| 222 | #[must_use] |
| 223 | pub fn optional_u64(input: &Value, field: &str, default: u64) -> u64 { |
| 224 | input.get(field).and_then(Value::as_u64).unwrap_or(default) |
| 225 | } |
| 226 | |
| 227 | /// Helper to extract an optional bool field with default. |
| 228 | #[must_use] |
| 229 | pub fn optional_bool(input: &Value, field: &str, default: bool) -> bool { |
| 230 | input.get(field).and_then(Value::as_bool).unwrap_or(default) |
| 231 | } |
| 232 | |
| 233 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 234 | pub struct ToolSpec { |
| 235 | pub name: String, |
| 236 | pub input_schema: Value, |
| 237 | pub output_schema: Value, |
| 238 | pub supports_parallel_tool_calls: bool, |
| 239 | pub timeout_ms: Option<u64>, |
| 240 | } |
| 241 | |
| 242 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 243 | pub struct ConfiguredToolSpec { |
| 244 | pub spec: ToolSpec, |
| 245 | pub supports_parallel_tool_calls: bool, |
| 246 | } |
| 247 | |
| 248 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
| 249 | #[serde(rename_all = "snake_case")] |
| 250 | pub enum ToolCallSource { |
| 251 | Direct, |
| 252 | JsRepl, |
| 253 | } |
| 254 | |
| 255 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 256 | pub struct ToolCall { |
| 257 | pub name: String, |
| 258 | pub payload: ToolPayload, |
| 259 | pub source: ToolCallSource, |
| 260 | pub raw_tool_call_id: Option<String>, |
| 261 | } |
| 262 | |
| 263 | impl ToolCall { |
| 264 | pub fn execution_subject(&self, fallback_cwd: &str) -> (String, String, &'static str) { |
| 265 | match &self.payload { |
| 266 | ToolPayload::LocalShell { params } => ( |
| 267 | params.command.clone(), |
| 268 | params |
| 269 | .cwd |
| 270 | .clone() |
| 271 | .unwrap_or_else(|| fallback_cwd.to_string()), |
| 272 | "shell", |
| 273 | ), |
| 274 | _ => (self.name.clone(), fallback_cwd.to_string(), "tool"), |
| 275 | } |
| 276 | } |
| 277 | } |
| 278 | |
| 279 | #[derive(Debug, Clone)] |
| 280 | pub struct ToolInvocation { |
| 281 | pub call_id: String, |
| 282 | pub tool_name: String, |
| 283 | pub payload: ToolPayload, |
| 284 | pub source: ToolCallSource, |
| 285 | } |
| 286 | |
| 287 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 288 | pub enum FunctionCallError { |
| 289 | ToolNotFound { name: String }, |
| 290 | KindMismatch { expected: ToolKind, got: ToolKind }, |
| 291 | MutatingToolRejected { name: String }, |
| 292 | TimedOut { name: String, timeout_ms: u64 }, |
| 293 | Cancelled { name: String }, |
| 294 | ExecutionFailed { name: String, error: String }, |
| 295 | } |
| 296 | |
| 297 | #[async_trait] |
| 298 | pub trait ToolHandler: Send + Sync { |
| 299 | fn kind(&self) -> ToolKind; |
| 300 | fn matches_kind(&self, kind: ToolKind) -> bool { |
| 301 | self.kind() == kind |
| 302 | } |
| 303 | fn is_mutating(&self) -> bool { |
| 304 | false |
| 305 | } |
| 306 | async fn handle( |
| 307 | &self, |
| 308 | invocation: ToolInvocation, |
| 309 | ) -> std::result::Result<ToolOutput, FunctionCallError>; |
| 310 | } |
| 311 | |
| 312 | #[derive(Debug, Default)] |
| 313 | pub struct ToolCallRuntime { |
| 314 | pub parallel_execution: Arc<RwLock<()>>, |
| 315 | } |
| 316 | |
| 317 | #[derive(Default)] |
| 318 | pub struct ToolRegistry { |
| 319 | handlers: HashMap<String, Arc<dyn ToolHandler>>, |
| 320 | specs: HashMap<String, ConfiguredToolSpec>, |
| 321 | runtime: ToolCallRuntime, |
| 322 | } |
| 323 | |
| 324 | impl ToolRegistry { |
| 325 | pub fn register(&mut self, spec: ToolSpec, handler: Arc<dyn ToolHandler>) -> Result<()> { |
| 326 | let name = spec.name.clone(); |
| 327 | self.specs.insert( |
| 328 | name.clone(), |
| 329 | ConfiguredToolSpec { |
| 330 | supports_parallel_tool_calls: spec.supports_parallel_tool_calls, |
| 331 | spec, |
| 332 | }, |
| 333 | ); |
| 334 | self.handlers.insert(name, handler); |
| 335 | Ok(()) |
| 336 | } |
| 337 | |
| 338 | pub fn list_specs(&self) -> Vec<ConfiguredToolSpec> { |
| 339 | self.specs.values().cloned().collect() |
| 340 | } |
| 341 | |
| 342 | pub async fn dispatch( |
| 343 | &self, |
| 344 | call: ToolCall, |
| 345 | allow_mutating: bool, |
| 346 | ) -> std::result::Result<ToolOutput, FunctionCallError> { |
| 347 | let handler = self.handlers.get(&call.name).cloned().ok_or_else(|| { |
| 348 | FunctionCallError::ToolNotFound { |
| 349 | name: call.name.clone(), |
| 350 | } |
| 351 | })?; |
| 352 | let configured = |
| 353 | self.specs |
| 354 | .get(&call.name) |
| 355 | .cloned() |
| 356 | .ok_or_else(|| FunctionCallError::ToolNotFound { |
| 357 | name: call.name.clone(), |
| 358 | })?; |
| 359 | |
| 360 | let payload_kind = tool_payload_kind(&call.payload); |
| 361 | let expected = handler.kind(); |
| 362 | if !handler.matches_kind(payload_kind) { |
| 363 | return Err(FunctionCallError::KindMismatch { |
| 364 | expected, |
| 365 | got: payload_kind, |
| 366 | }); |
| 367 | } |
| 368 | if handler.is_mutating() && !allow_mutating { |
| 369 | return Err(FunctionCallError::MutatingToolRejected { name: call.name }); |
| 370 | } |
| 371 | |
| 372 | let invocation = ToolInvocation { |
| 373 | call_id: call |
| 374 | .raw_tool_call_id |
| 375 | .clone() |
| 376 | .unwrap_or_else(|| format!("tool-call-{}", uuid::Uuid::new_v4())), |
| 377 | tool_name: call.name.clone(), |
| 378 | payload: call.payload, |
| 379 | source: call.source, |
| 380 | }; |
| 381 | |
| 382 | if configured.supports_parallel_tool_calls { |
| 383 | let _guard = self.runtime.parallel_execution.read().await; |
| 384 | self.execute_with_timeout(handler, configured.spec.timeout_ms, invocation) |
| 385 | .await |
| 386 | } else { |
| 387 | let _guard = self.runtime.parallel_execution.write().await; |
| 388 | self.execute_with_timeout(handler, configured.spec.timeout_ms, invocation) |
| 389 | .await |
| 390 | } |
| 391 | } |
| 392 | |
| 393 | async fn execute_with_timeout( |
| 394 | &self, |
| 395 | handler: Arc<dyn ToolHandler>, |
| 396 | timeout_ms: Option<u64>, |
| 397 | invocation: ToolInvocation, |
| 398 | ) -> std::result::Result<ToolOutput, FunctionCallError> { |
| 399 | if let Some(timeout_ms) = timeout_ms { |
| 400 | let name = invocation.tool_name.clone(); |
| 401 | match tokio::time::timeout( |
| 402 | Duration::from_millis(timeout_ms), |
| 403 | handler.handle(invocation), |
| 404 | ) |
| 405 | .await |
| 406 | { |
| 407 | Ok(result) => result, |
| 408 | Err(_) => Err(FunctionCallError::TimedOut { name, timeout_ms }), |
| 409 | } |
| 410 | } else { |
| 411 | handler.handle(invocation).await |
| 412 | } |
| 413 | } |
| 414 | } |
| 415 | |
| 416 | fn tool_payload_kind(payload: &ToolPayload) -> ToolKind { |
| 417 | match payload { |
| 418 | ToolPayload::Mcp { .. } => ToolKind::Mcp, |
| 419 | ToolPayload::Function { .. } |
| 420 | | ToolPayload::Custom { .. } |
| 421 | | ToolPayload::LocalShell { .. } => ToolKind::Function, |
| 422 | } |
| 423 | } |
| 424 | |
| 425 | #[cfg(test)] |
| 426 | mod tests { |
| 427 | use serde_json::json; |
| 428 | |
| 429 | use super::*; |
| 430 | |
| 431 | #[test] |
| 432 | fn tool_result_json_round_trips_content() { |
| 433 | let result = ToolResult::json(&json!({"ok": true})).expect("json"); |
| 434 | assert!(result.success); |
| 435 | assert!(result.content.contains("\"ok\": true")); |
| 436 | } |
| 437 | |
| 438 | #[test] |
| 439 | fn helper_extractors_validate_shape() { |
| 440 | let input = json!({"name": "demo", "count": 7, "enabled": true}); |
| 441 | assert_eq!(required_str(&input, "name").expect("name"), "demo"); |
| 442 | assert_eq!(optional_u64(&input, "count", 0), 7); |
| 443 | assert!(optional_bool(&input, "enabled", false)); |
| 444 | assert!(matches!( |
| 445 | required_u64(&input, "name"), |
| 446 | Err(ToolError::MissingField { .. }) |
| 447 | )); |
| 448 | } |
| 449 | |
| 450 | #[test] |
| 451 | fn required_str_reports_provided_fields_on_missing_required_field() { |
| 452 | let input = json!({"path": "src/lib.rs", "content": "new body"}); |
| 453 | let err = required_str(&input, "replace").expect_err("replace is missing"); |
| 454 | let message = err.to_string(); |
| 455 | assert!(message.contains("missing required field 'replace'")); |
| 456 | assert!(message.contains("Input provided:")); |
| 457 | assert!(message.contains("path")); |
| 458 | assert!(message.contains("content")); |
| 459 | } |
| 460 | |
| 461 | #[test] |
| 462 | fn tool_error_display_matches_legacy_text() { |
| 463 | let err = ToolError::missing_field("path"); |
| 464 | assert_eq!( |
| 465 | err.to_string(), |
| 466 | "Failed to validate input: missing required field 'path'" |
| 467 | ); |
| 468 | } |
| 469 | } |
| 470 |