1use crate::executors::{
2 file_stat, patch_apply, patch_propose, repo_list, repo_read, repo_search, run_checks,
3};
4use crate::registry::{validate_tool_input_with_canonical_runtime, ToolExecutorV1, ToolRegistryV1};
5use aidens_contracts::{
6 ApprovalRequestV1, CanonicalToolSideEffectClass, PermitUseReportV1, SchemaValidationReportV1,
7 ToolDescriptorV1, ToolInvocationReportV1,
8};
9use aidens_permit_kit::{requires_permit, PermitCheckContextV1, PermitDecisionV1, PermitPolicyV1};
10use serde_json::Value;
11use std::fmt;
12
13#[derive(Debug, Clone)]
14pub struct ToolDispatcher {
15 registry: ToolRegistryV1,
16 permit_policy: PermitPolicyV1,
17}
18
19impl ToolDispatcher {
20 pub fn new(registry: ToolRegistryV1) -> Self {
21 Self {
22 registry,
23 permit_policy: PermitPolicyV1::default(),
24 }
25 }
26
27 pub fn with_permit_policy(mut self, permit_policy: PermitPolicyV1) -> Self {
28 self.permit_policy = permit_policy;
29 self
30 }
31
32 pub async fn invoke(
33 &self,
34 tool_id: &str,
35 input: Value,
36 ) -> anyhow::Result<ToolInvocationOutcome> {
37 let Some(descriptor) = self.registry.descriptor(tool_id) else {
38 let receipt = ToolInvocationReportV1::started(tool_id, input)
39 .complete_failure("tool-not-registered");
40 return Err(ToolInvocationError::new(
41 format!("tool '{tool_id}' is not registered"),
42 receipt,
43 )
44 .into());
45 };
46
47 let mut receipt = ToolInvocationReportV1::started(tool_id, input.clone());
48 let schema_validation = validate_tool_input_with_canonical_runtime(descriptor, &input);
49 if !schema_validation.valid {
50 let failed = receipt
51 .with_schema_validation(schema_validation.receipt_id.clone())
52 .complete_failure("schema-validation-failed");
53 return Err(ToolInvocationError::new(
54 format!("tool '{tool_id}' input failed schema validation"),
55 failed,
56 )
57 .with_schema_validation_receipt(schema_validation)
58 .into());
59 }
60
61 let sandbox_scope = self.sandbox_scope_for(tool_id);
62 let permit_context = PermitCheckContextV1::new(
63 tool_id.to_string(),
64 descriptor.risk_class.clone(),
65 sandbox_scope.clone(),
66 );
67 let mut permit_use_receipt = None;
68 if let Some(grant) = self.permit_policy.grant_for_context(&permit_context) {
69 receipt = receipt.with_permit_grant(grant.permit_id.clone());
70 let use_receipt = PermitUseReportV1::allowed(
71 grant,
72 tool_id.to_string(),
73 sandbox_scope.clone(),
74 None,
75 None,
76 );
77 receipt = receipt.with_permit_use(use_receipt.receipt_id.clone());
78 permit_use_receipt = Some(use_receipt);
79 }
80
81 if requires_permit(&descriptor.risk_class) {
82 match self.permit_policy.decision_for_context(&permit_context) {
83 PermitDecisionV1::Allow => {}
84 PermitDecisionV1::RequiresApproval => {
85 let approval_request = self
86 .permit_policy
87 .approval_request_for_context(&permit_context)
88 .unwrap_or_else(|| {
89 ApprovalRequestV1::scoped(
90 tool_id.to_string(),
91 descriptor.risk_class.clone(),
92 sandbox_scope.clone(),
93 "side-effect tool requires explicit scoped permit",
94 )
95 });
96 let failed = receipt
97 .with_approval_request(approval_request.request_id.clone())
98 .complete_failure(format!("permit-required:{}", descriptor.risk_class));
99 return Err(ToolInvocationError::new(
100 format!(
101 "tool '{tool_id}' requires explicit permit for {}",
102 descriptor.risk_class
103 ),
104 failed,
105 )
106 .with_approval_request(approval_request)
107 .into());
108 }
109 PermitDecisionV1::Deny(reason) => {
110 let failed = receipt.complete_failure(format!("permit-denied:{reason}"));
111 return Err(ToolInvocationError::new(
112 format!("tool '{tool_id}' permit denied: {reason}"),
113 failed,
114 )
115 .into());
116 }
117 }
118 }
119
120 let Some(executor) = self.registry.executors.get(tool_id) else {
121 let failed = receipt.complete_failure("tool-executor-missing");
122 return Err(ToolInvocationError::new(
123 format!("tool '{tool_id}' is registered but not executable this turn"),
124 failed,
125 )
126 .into());
127 };
128
129 match executor {
130 ToolExecutorV1::RepoRead { sandbox_root } => {
131 if !descriptor.read_only
132 || !matches!(
133 &descriptor.risk_class,
134 CanonicalToolSideEffectClass::ReadOnly
135 )
136 {
137 let failed = receipt.complete_failure("read-only-executor-risk-mismatch");
138 return Err(ToolInvocationError::new(
139 format!("tool '{tool_id}' is blocked by read-only executor policy"),
140 failed,
141 )
142 .into());
143 }
144 self.invoke_executor(tool_id, receipt, permit_use_receipt, || {
145 repo_read(sandbox_root, &input)
146 })
147 }
148 ToolExecutorV1::RepoList { sandbox_root } => self.invoke_read_only_executor(
149 tool_id,
150 descriptor,
151 receipt,
152 permit_use_receipt,
153 || repo_list(sandbox_root, &input),
154 ),
155 ToolExecutorV1::FileStat { sandbox_root } => self.invoke_read_only_executor(
156 tool_id,
157 descriptor,
158 receipt,
159 permit_use_receipt,
160 || file_stat(sandbox_root, &input),
161 ),
162 ToolExecutorV1::RepoSearch { sandbox_root } => self.invoke_read_only_executor(
163 tool_id,
164 descriptor,
165 receipt,
166 permit_use_receipt,
167 || repo_search(sandbox_root, &input),
168 ),
169 ToolExecutorV1::PatchPropose { sandbox_root } => self.invoke_read_only_executor(
170 tool_id,
171 descriptor,
172 receipt,
173 permit_use_receipt,
174 || patch_propose(sandbox_root, &input),
175 ),
176 ToolExecutorV1::PatchApply { sandbox_root } => {
177 let permit_grant_id = receipt.permit_grant_id.clone();
178 let permit_use_receipt_id = receipt.permit_use_receipt_id.clone();
179 self.invoke_executor(tool_id, receipt, permit_use_receipt, || {
180 patch_apply(sandbox_root, &input, permit_grant_id, permit_use_receipt_id)
181 })
182 }
183 ToolExecutorV1::RunChecks { sandbox_root } => {
184 let permit_grant_id = receipt.permit_grant_id.clone();
185 let permit_use_receipt_id = receipt.permit_use_receipt_id.clone();
186 self.invoke_executor(tool_id, receipt, permit_use_receipt, || {
187 run_checks(sandbox_root, &input, permit_grant_id, permit_use_receipt_id)
188 })
189 }
190 ToolExecutorV1::Custom(handle) => {
191 let result = handle.inner.execute(tool_id, input.clone()).await;
192 match result {
193 Ok(output_text) => {
194 let output_val = serde_json::json!({ "content": output_text });
195 receipt = receipt.complete_success(output_val.clone());
196 Ok(ToolInvocationOutcome {
197 output: output_val,
198 receipt,
199 permit_use_receipt,
200 })
201 }
202 Err(e) => {
203 let failed =
204 receipt.complete_failure(format!("custom-executor-failed: {e}"));
205 Err(ToolInvocationError::new(
206 format!("custom executor for '{tool_id}' failed: {e}"),
207 failed,
208 )
209 .into())
210 }
211 }
212 }
213 }
214 }
215
216 fn invoke_read_only_executor(
217 &self,
218 tool_id: &str,
219 descriptor: &ToolDescriptorV1,
220 receipt: ToolInvocationReportV1,
221 permit_use_receipt: Option<PermitUseReportV1>,
222 execute: impl FnOnce() -> anyhow::Result<Value>,
223 ) -> anyhow::Result<ToolInvocationOutcome> {
224 if !descriptor.read_only
225 || !matches!(
226 &descriptor.risk_class,
227 CanonicalToolSideEffectClass::ReadOnly
228 )
229 {
230 let failed = receipt.complete_failure("read-only-executor-risk-mismatch");
231 return Err(ToolInvocationError::new(
232 format!("tool '{tool_id}' is blocked by read-only executor policy"),
233 failed,
234 )
235 .into());
236 }
237 self.invoke_executor(tool_id, receipt, permit_use_receipt, execute)
238 }
239
240 fn invoke_executor(
241 &self,
242 tool_id: &str,
243 receipt: ToolInvocationReportV1,
244 permit_use_receipt: Option<PermitUseReportV1>,
245 execute: impl FnOnce() -> anyhow::Result<Value>,
246 ) -> anyhow::Result<ToolInvocationOutcome> {
247 match execute() {
248 Ok(output) => {
249 let receipt = receipt.complete_success(output.clone());
250 Ok(ToolInvocationOutcome {
251 output,
252 receipt,
253 permit_use_receipt,
254 })
255 }
256 Err(error) => {
257 let failed = if let Some(receipt_failure) =
258 error.downcast_ref::<ReceiptBearingToolFailure>()
259 {
260 receipt.complete_failure_with_output(
261 receipt_failure.reason_code.clone(),
262 receipt_failure.output.clone(),
263 )
264 } else {
265 let reason = executor_reason_code(error.to_string());
266 receipt.complete_failure(reason)
267 };
268 Err(ToolInvocationError::new(
269 format!("tool '{tool_id}' execution failed: {error}"),
270 failed,
271 )
272 .into())
273 }
274 }
275 }
276
277 fn sandbox_scope_for(&self, tool_id: &str) -> String {
278 const UNKNOWN_SANDBOX_SCOPE: &str = "unknown-sandbox-root";
279 self.registry
280 .executors
281 .get(tool_id)
282 .map(|executor| match executor {
283 ToolExecutorV1::RepoRead { sandbox_root }
284 | ToolExecutorV1::RepoList { sandbox_root }
285 | ToolExecutorV1::FileStat { sandbox_root }
286 | ToolExecutorV1::RepoSearch { sandbox_root }
287 | ToolExecutorV1::PatchPropose { sandbox_root }
288 | ToolExecutorV1::PatchApply { sandbox_root }
289 | ToolExecutorV1::RunChecks { sandbox_root } => sandbox_root.display().to_string(),
290 ToolExecutorV1::Custom(_) => "custom-executor".to_string(),
291 })
292 .or_else(|| self.registry.sandbox_root_display())
293 .unwrap_or_else(|| UNKNOWN_SANDBOX_SCOPE.into())
294 }
295}
296
297#[derive(Debug, Clone)]
298pub struct ToolInvocationOutcome {
299 pub output: Value,
300 pub receipt: ToolInvocationReportV1,
301 pub permit_use_receipt: Option<PermitUseReportV1>,
302}
303
304impl ToolInvocationOutcome {
305 pub fn output_text(&self) -> String {
306 self.output
307 .get("content")
308 .and_then(Value::as_str)
309 .map(str::to_string)
310 .unwrap_or_else(|| self.output.to_string())
311 }
312}
313
314#[derive(Debug, Clone)]
315pub struct ToolInvocationError {
316 message: String,
317 receipt: ToolInvocationReportV1,
318 approval_request: Option<ApprovalRequestV1>,
319 schema_validation_receipt: Option<SchemaValidationReportV1>,
320}
321
322impl ToolInvocationError {
323 pub fn new(message: String, receipt: ToolInvocationReportV1) -> Self {
324 Self {
325 message,
326 receipt,
327 approval_request: None,
328 schema_validation_receipt: None,
329 }
330 }
331
332 pub fn receipt(&self) -> &ToolInvocationReportV1 {
333 &self.receipt
334 }
335
336 pub fn approval_request(&self) -> Option<&ApprovalRequestV1> {
337 self.approval_request.as_ref()
338 }
339
340 pub fn schema_validation_receipt(&self) -> Option<&SchemaValidationReportV1> {
341 self.schema_validation_receipt.as_ref()
342 }
343
344 pub fn with_approval_request(mut self, approval_request: ApprovalRequestV1) -> Self {
345 self.approval_request = Some(approval_request);
346 self
347 }
348
349 pub fn with_schema_validation_receipt(
350 mut self,
351 schema_validation_receipt: SchemaValidationReportV1,
352 ) -> Self {
353 self.schema_validation_receipt = Some(schema_validation_receipt);
354 self
355 }
356}
357
358impl fmt::Display for ToolInvocationError {
359 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
360 if self.receipt.reason_codes.is_empty() {
361 f.write_str(&self.message)
362 } else {
363 write!(
364 f,
365 "{} [{}]",
366 self.message,
367 self.receipt.reason_codes.join(",")
368 )
369 }
370 }
371}
372
373impl std::error::Error for ToolInvocationError {}
374
375#[derive(Debug, Clone)]
376pub(crate) struct ReceiptBearingToolFailure {
377 pub(crate) message: String,
378 pub(crate) reason_code: String,
379 pub(crate) output: Value,
380}
381
382impl fmt::Display for ReceiptBearingToolFailure {
383 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
384 f.write_str(&self.message)
385 }
386}
387
388impl std::error::Error for ReceiptBearingToolFailure {}
389
390pub(crate) fn executor_reason_code(error: String) -> String {
391 let lower = error.to_ascii_lowercase();
392 if lower.contains("traversal") {
393 "sandbox-path-traversal-denied".into()
394 } else if lower.contains("hardlink") {
395 "sandbox-hardlink-denied".into()
396 } else if lower.contains("sensitive prefix") {
397 "sandbox-sensitive-prefix-denied".into()
398 } else if lower.contains("hidden or sensitive component") {
399 "sandbox-hidden-component-denied".into()
400 } else if lower.contains("escape") || lower.contains("outside sandbox") {
401 "sandbox-escape-denied".into()
402 } else if lower.contains("command-not-allowed") {
403 "command-not-allowed-by-policy".into()
404 } else if lower.contains("patch") {
405 "patch-apply-failed".into()
406 } else {
407 "tool-executor-failed".into()
408 }
409}