aidens_tool_kit/
registry.rs

1use crate::descriptors::safe_coding_tool_plan;
2use crate::sandbox::canonical_sandbox_root;
3use crate::{canonical_stack, ToolExposurePolicyV1};
4use aidens_contracts::{
5    generated_artifact_id_from_material, ApprovalRequestV1, ArtifactId, CanonicalBackpointerV1,
6    CapabilityGateDecisionDraftV1, CapabilityGateDecisionV1, CapabilityGateOutcomeV1,
7    PermitUseReportV1, SchemaValidationReportV1, ToolDescriptorV1, ToolExposureSetV1,
8    ToolLifecycleStateV1, ToolProviderSchemaV1,
9};
10use aidens_permit_kit::{requires_permit, PermitCheckContextV1, PermitDecisionV1};
11use serde_json::Value;
12use std::collections::BTreeMap;
13use std::path::{Path, PathBuf};
14
15#[derive(Debug, Clone, Default)]
16pub struct ToolRegistryV1 {
17    pub(crate) tools: BTreeMap<String, ToolDescriptorV1>,
18    pub(crate) executors: BTreeMap<String, ToolExecutorV1>,
19    pub(crate) sandbox_root: Option<PathBuf>,
20    pub(crate) construction_degradation_reasons: Vec<String>,
21}
22
23#[derive(Debug, Clone)]
24pub(crate) enum ToolExecutorV1 {
25    RepoRead { sandbox_root: PathBuf },
26    RepoList { sandbox_root: PathBuf },
27    FileStat { sandbox_root: PathBuf },
28    RepoSearch { sandbox_root: PathBuf },
29    PatchPropose { sandbox_root: PathBuf },
30    PatchApply { sandbox_root: PathBuf },
31    RunChecks { sandbox_root: PathBuf },
32    Custom(crate::custom::CustomExecutorHandle),
33}
34
35#[derive(Debug, Clone)]
36pub(crate) struct GateDescriptorOutcome {
37    outcome: CapabilityGateOutcomeV1,
38    reason_codes: Vec<String>,
39    approval_request: Option<ApprovalRequestV1>,
40    permit_grant_id: Option<ArtifactId>,
41    permit_use_receipt: Option<PermitUseReportV1>,
42}
43
44impl GateDescriptorOutcome {
45    fn exposed() -> Self {
46        Self {
47            outcome: CapabilityGateOutcomeV1::Exposed,
48            reason_codes: Vec::new(),
49            approval_request: None,
50            permit_grant_id: None,
51            permit_use_receipt: None,
52        }
53    }
54
55    fn exposed_with_permit(
56        permit_use_receipt: PermitUseReportV1,
57        permit_grant_id: ArtifactId,
58    ) -> Self {
59        Self {
60            outcome: CapabilityGateOutcomeV1::Exposed,
61            reason_codes: vec!["permit-scope-matched".into()],
62            approval_request: None,
63            permit_grant_id: Some(permit_grant_id),
64            permit_use_receipt: Some(permit_use_receipt),
65        }
66    }
67
68    fn hidden(reason_codes: Vec<String>) -> Self {
69        Self {
70            outcome: CapabilityGateOutcomeV1::Hidden,
71            reason_codes,
72            approval_request: None,
73            permit_grant_id: None,
74            permit_use_receipt: None,
75        }
76    }
77
78    fn blocked(reason_codes: Vec<String>) -> Self {
79        Self {
80            outcome: CapabilityGateOutcomeV1::Blocked,
81            reason_codes,
82            approval_request: None,
83            permit_grant_id: None,
84            permit_use_receipt: None,
85        }
86    }
87
88    fn blocked_with_approval(
89        reason_codes: Vec<String>,
90        approval_request: ApprovalRequestV1,
91    ) -> Self {
92        Self {
93            outcome: CapabilityGateOutcomeV1::Blocked,
94            reason_codes,
95            approval_request: Some(approval_request),
96            permit_grant_id: None,
97            permit_use_receipt: None,
98        }
99    }
100}
101
102pub(crate) fn canonical_descriptor_from_aidens(
103    descriptor: &ToolDescriptorV1,
104) -> canonical_stack::CanonicalToolDescriptor {
105    canonical_stack::CanonicalToolDescriptor {
106        name: descriptor.tool_id(),
107        version: descriptor.version.clone(),
108        description: Some(descriptor.description.clone()),
109        backend_kind: canonical_stack::ToolBackendKind::LocalFunction,
110        input_schema: descriptor.schema.input_schema.clone(),
111        output_mode: canonical_stack::ToolOutputMode::StructuredJson,
112        read_only: descriptor.read_only,
113        side_effect_class: descriptor.risk_class.clone(),
114        idempotency_class: if descriptor.read_only {
115            canonical_stack::ToolIdempotencyClass::Idempotent
116        } else {
117            canonical_stack::ToolIdempotencyClass::BestEffort
118        },
119        approval_kind: if requires_permit(&descriptor.risk_class) {
120            canonical_stack::ToolApprovalKind::PolicyRequired
121        } else {
122            canonical_stack::ToolApprovalKind::None
123        },
124        timeout_ms: 30_000,
125        concurrency_key: None,
126        cache_ttl_ms: None,
127        exposure_mode: if descriptor.hidden {
128            canonical_stack::ToolExposureMode::Hidden
129        } else {
130            canonical_stack::ToolExposureMode::Auto
131        },
132        mcp_surface_kind: canonical_stack::McpSurfaceKind::Tool,
133        exposure_policy: Default::default(),
134        receipt_persistence: if descriptor.read_only || requires_permit(&descriptor.risk_class) {
135            canonical_stack::ToolReceiptPersistence::ForgeRaw
136        } else {
137            canonical_stack::ToolReceiptPersistence::Ephemeral
138        },
139        output_size_limit_bytes: None,
140        provider_payload: Some(serde_json::json!({
141            "aidens_namespace": descriptor.namespace,
142            "aidens_name": descriptor.name,
143            "canonical_owner": "llm-tool-runtime",
144        })),
145    }
146}
147
148pub(crate) fn validate_tool_input_with_canonical_runtime(
149    descriptor: &ToolDescriptorV1,
150    input: &Value,
151) -> SchemaValidationReportV1 {
152    let canonical_descriptor = canonical_descriptor_from_aidens(descriptor);
153    let errors = canonical_stack::validate_canonical_arguments(&canonical_descriptor, input)
154        .err()
155        .map(|error| vec![error.message])
156        .unwrap_or_default();
157    SchemaValidationReportV1::new(Some(&descriptor.schema.input_schema), input, errors)
158        .with_tool_id(descriptor.tool_id())
159}
160
161impl ToolRegistryV1 {
162    pub fn register_enabled(&mut self, descriptor: ToolDescriptorV1, enabled: bool) -> bool {
163        if enabled {
164            self.tools.insert(descriptor.tool_id(), descriptor);
165            true
166        } else {
167            false
168        }
169    }
170
171    pub fn register_enabled_with_repo_read_dispatcher(
172        &mut self,
173        descriptor: ToolDescriptorV1,
174        enabled: bool,
175        sandbox_root: impl AsRef<Path>,
176    ) -> anyhow::Result<bool> {
177        let tool_id = descriptor.tool_id();
178        if !self.register_enabled(descriptor, enabled) {
179            return Ok(false);
180        }
181        let sandbox_root = canonical_sandbox_root(sandbox_root.as_ref())?;
182        self.sandbox_root = Some(sandbox_root.clone());
183        self.executors
184            .insert(tool_id, ToolExecutorV1::RepoRead { sandbox_root });
185        Ok(true)
186    }
187
188    /// Register a tool with a custom async executor.
189    /// The executor implements `CustomToolExecutor` and will be called
190    /// when the tool is dispatched during the agent loop.
191    pub fn register_enabled_with_custom_executor(
192        &mut self,
193        descriptor: ToolDescriptorV1,
194        enabled: bool,
195        executor: std::sync::Arc<dyn crate::custom::CustomToolExecutor>,
196    ) -> bool {
197        let tool_id = descriptor.tool_id();
198        if !self.register_enabled(descriptor, enabled) {
199            return false;
200        }
201        self.executors.insert(
202            tool_id,
203            ToolExecutorV1::Custom(crate::custom::CustomExecutorHandle::new(executor)),
204        );
205        true
206    }
207
208    fn register_enabled_with_executor(
209        &mut self,
210        descriptor: ToolDescriptorV1,
211        enabled: bool,
212        sandbox_root: impl AsRef<Path>,
213        executor: fn(PathBuf) -> ToolExecutorV1,
214    ) -> anyhow::Result<bool> {
215        let tool_id = descriptor.tool_id();
216        if !self.register_enabled(descriptor, enabled) {
217            return Ok(false);
218        }
219        let sandbox_root = canonical_sandbox_root(sandbox_root.as_ref())?;
220        self.sandbox_root = Some(sandbox_root.clone());
221        self.executors.insert(tool_id, executor(sandbox_root));
222        Ok(true)
223    }
224
225    pub fn with_sandbox_root(mut self, sandbox_root: impl AsRef<Path>) -> anyhow::Result<Self> {
226        self.sandbox_root = Some(canonical_sandbox_root(sandbox_root.as_ref())?);
227        Ok(self)
228    }
229
230    pub fn sandbox_root(&self) -> Option<&Path> {
231        self.sandbox_root.as_deref()
232    }
233
234    pub fn contains_tool_id(&self, tool_id: &str) -> bool {
235        self.tools.contains_key(tool_id)
236    }
237
238    pub fn descriptor(&self, tool_id: &str) -> Option<&ToolDescriptorV1> {
239        self.tools.get(tool_id)
240    }
241
242    pub fn descriptors(&self) -> Vec<ToolDescriptorV1> {
243        self.tools.values().cloned().collect()
244    }
245
246    pub fn tool_ids(&self) -> Vec<String> {
247        self.tools.keys().cloned().collect()
248    }
249
250    pub fn executable_tool_ids(&self) -> Vec<String> {
251        self.tools
252            .keys()
253            .filter(|tool_id| self.can_execute(tool_id))
254            .cloned()
255            .collect()
256    }
257
258    pub fn len(&self) -> usize {
259        self.tools.len()
260    }
261
262    pub fn is_empty(&self) -> bool {
263        self.tools.is_empty()
264    }
265
266    pub fn expose_read_only(&self) -> ToolExposureSetV1 {
267        self.plan_exposure(&ToolExposurePolicyV1::read_only_default())
268    }
269
270    pub fn can_execute(&self, tool_id: &str) -> bool {
271        self.executors.contains_key(tool_id)
272    }
273
274    pub fn plan_exposure(&self, policy: &ToolExposurePolicyV1) -> ToolExposureSetV1 {
275        self.plan_exposure_with_declarations(policy, self.descriptors())
276    }
277
278    pub fn plan_exposure_with_declarations(
279        &self,
280        policy: &ToolExposurePolicyV1,
281        declarations: Vec<ToolDescriptorV1>,
282    ) -> ToolExposureSetV1 {
283        let mut declarations = declarations;
284        declarations.sort_by_key(ToolDescriptorV1::tool_id);
285        let mut exposed = Vec::new();
286        let mut hidden = Vec::new();
287        let mut blocked = Vec::new();
288        let mut decisions = Vec::new();
289        let mut approval_requests = Vec::new();
290        let mut permit_use_receipts = Vec::new();
291        let mut provider_tool_schemas = Vec::new();
292        let mut reason_codes = self.construction_degradation_reasons.clone();
293        let declared_tool_ids = declarations
294            .iter()
295            .map(ToolDescriptorV1::tool_id)
296            .collect::<Vec<_>>();
297        let registered_tool_ids = self.tool_ids();
298        let executable_tool_ids = self.executable_tool_ids();
299        let sandbox_root = policy
300            .sandbox_root
301            .clone()
302            .or_else(|| self.sandbox_root_display());
303
304        for descriptor in declarations {
305            let tool_id = descriptor.tool_id();
306            let registered = self.contains_tool_id(&tool_id);
307            let executable = self.can_execute(&tool_id);
308            let mut lifecycle = vec![ToolLifecycleStateV1::Declared];
309            if registered {
310                lifecycle.push(ToolLifecycleStateV1::Registered);
311            }
312            if executable {
313                lifecycle.push(ToolLifecycleStateV1::Executable);
314            }
315
316            let gate = self.gate_descriptor(
317                &descriptor,
318                policy,
319                exposed.len(),
320                registered,
321                executable,
322                sandbox_root.as_deref(),
323            );
324
325            match gate.outcome {
326                CapabilityGateOutcomeV1::Exposed => {
327                    lifecycle.push(ToolLifecycleStateV1::ExposedThisTurn);
328                    exposed.push(tool_id.clone());
329                    provider_tool_schemas.push(ToolProviderSchemaV1::from_descriptor(&descriptor));
330                }
331                CapabilityGateOutcomeV1::Hidden => {
332                    lifecycle.push(ToolLifecycleStateV1::Hidden);
333                    hidden.push(tool_id.clone());
334                }
335                CapabilityGateOutcomeV1::Blocked => {
336                    lifecycle.push(ToolLifecycleStateV1::Blocked);
337                    blocked.push(tool_id.clone());
338                }
339            }
340
341            if let Some(request) = gate.approval_request.clone() {
342                approval_requests.push(request);
343            }
344            if let Some(permit_use_receipt) = gate.permit_use_receipt.clone() {
345                permit_use_receipts.push(permit_use_receipt);
346            }
347            reason_codes.extend(gate.reason_codes.iter().cloned());
348            decisions.push(CapabilityGateDecisionV1::for_tool(
349                CapabilityGateDecisionDraftV1 {
350                    tool_id,
351                    outcome: gate.outcome,
352                    lifecycle,
353                    risk_class: descriptor.risk_class.clone(),
354                    permit_required: requires_permit(&descriptor.risk_class),
355                    executable_this_turn: executable,
356                    sandbox_root: sandbox_root.clone(),
357                    approval_request: gate.approval_request,
358                    permit_grant_id: gate.permit_grant_id,
359                    permit_use_receipt_id: gate
360                        .permit_use_receipt
361                        .map(|receipt| receipt.receipt_id),
362                    reason_codes: gate.reason_codes,
363                },
364            ));
365        }
366
367        reason_codes.sort();
368        reason_codes.dedup();
369        let exposure_material = format!(
370            "declared={declared:?}|registered={registered:?}|executable={executable_ids:?}|exposed={exposed:?}|hidden={hidden:?}|blocked={blocked:?}|sandbox={sandbox}",
371            declared = &declared_tool_ids,
372            registered = &registered_tool_ids,
373            executable_ids = &executable_tool_ids,
374            sandbox = sandbox_root
375                .as_ref()
376                .cloned()
377                .unwrap_or_else(|| "<none>".into()),
378        );
379
380        ToolExposureSetV1 {
381            exposure_id: generated_artifact_id_from_material("tool-exposure", &exposure_material),
382            declared_tool_ids,
383            registered_tool_ids,
384            executable_tool_ids,
385            exposed_tool_ids: exposed,
386            hidden_tool_ids: hidden,
387            blocked_tool_ids: blocked,
388            decisions,
389            approval_requests,
390            permit_use_receipts,
391            provider_tool_schemas,
392            sandbox_root,
393            degraded: !self.construction_degradation_reasons.is_empty(),
394            reason_codes,
395            canonical_backpointers: vec![CanonicalBackpointerV1::owner_type(
396                "llm-tool-runtime",
397                "ToolExposurePlan",
398                "canonical-tool-exposure-owner",
399            )],
400            reason: Some("policy-filtered exposure".into()),
401        }
402    }
403
404    pub fn declared_not_registered_tool_ids(
405        &self,
406        declarations: &[ToolDescriptorV1],
407    ) -> Vec<String> {
408        declarations
409            .iter()
410            .map(ToolDescriptorV1::tool_id)
411            .filter(|tool_id| !self.contains_tool_id(tool_id))
412            .collect()
413    }
414
415    pub(crate) fn sandbox_root_display(&self) -> Option<String> {
416        self.sandbox_root
417            .as_ref()
418            .map(|path| path.display().to_string())
419    }
420
421    fn gate_descriptor(
422        &self,
423        descriptor: &ToolDescriptorV1,
424        policy: &ToolExposurePolicyV1,
425        exposed_count: usize,
426        registered: bool,
427        executable: bool,
428        sandbox_root: Option<&str>,
429    ) -> GateDescriptorOutcome {
430        let tool_id = descriptor.tool_id();
431        if !registered {
432            return GateDescriptorOutcome::hidden(vec![format!(
433                "tool-declared-not-registered:{tool_id}"
434            )]);
435        }
436        if descriptor.hidden {
437            return GateDescriptorOutcome::hidden(vec![format!("tool-hidden:{tool_id}")]);
438        }
439        if policy
440            .allowed_tool_ids
441            .as_ref()
442            .is_some_and(|allowed| !allowed.contains(&tool_id))
443        {
444            return GateDescriptorOutcome::hidden(vec![format!(
445                "tool-not-allowed-this-turn:{tool_id}"
446            )]);
447        }
448        if !policy.allowed_risk_classes.contains(&descriptor.risk_class) {
449            return GateDescriptorOutcome::blocked(vec![format!(
450                "risk-blocked:{}",
451                descriptor.risk_class
452            )]);
453        }
454        let mut permit_grant_id = None;
455        let mut permit_use_receipt = None;
456        if requires_permit(&descriptor.risk_class) {
457            let Some(permit_scope) = sandbox_root else {
458                return GateDescriptorOutcome::blocked(vec![
459                    "permit-scope-missing-sandbox-root".into(),
460                    format!("permit-required:{}", descriptor.risk_class),
461                ]);
462            };
463            let context = PermitCheckContextV1::new(
464                tool_id.clone(),
465                descriptor.risk_class.clone(),
466                permit_scope,
467            );
468            match policy.permit_policy.decision_for_context(&context) {
469                PermitDecisionV1::Allow => {
470                    if let Some(receipt) = policy
471                        .permit_policy
472                        .permit_use_receipt_for_context(&context)
473                    {
474                        permit_grant_id = Some(receipt.permit_id.clone());
475                        permit_use_receipt = Some(receipt);
476                    }
477                }
478                PermitDecisionV1::RequiresApproval => {
479                    let approval_request = policy
480                        .permit_policy
481                        .approval_request_for_context(&context)
482                        .unwrap_or_else(|| {
483                            ApprovalRequestV1::scoped(
484                                tool_id.clone(),
485                                descriptor.risk_class.clone(),
486                                permit_scope,
487                                "side-effect tool requires explicit scoped permit",
488                            )
489                        });
490                    return GateDescriptorOutcome::blocked_with_approval(
491                        vec![format!("permit-required:{}", descriptor.risk_class)],
492                        approval_request,
493                    );
494                }
495                PermitDecisionV1::Deny(reason) => {
496                    return GateDescriptorOutcome::blocked(vec![format!("permit-denied:{reason}")]);
497                }
498            }
499        }
500        if descriptor.requires_native_tool_loop && !policy.native_tool_loop_available {
501            return GateDescriptorOutcome::blocked(vec![format!(
502                "route-requires-native-tool-loop:{tool_id}"
503            )]);
504        }
505        if !executable {
506            return GateDescriptorOutcome::hidden(vec![format!("tool-executor-missing:{tool_id}")]);
507        }
508        if policy.max_tools.is_some_and(|max| exposed_count >= max) {
509            return GateDescriptorOutcome::hidden(vec!["max-tools-reached".into()]);
510        }
511        match (permit_use_receipt, permit_grant_id) {
512            (Some(receipt), Some(grant_id)) => {
513                GateDescriptorOutcome::exposed_with_permit(receipt, grant_id)
514            }
515            _ => GateDescriptorOutcome::exposed(),
516        }
517    }
518
519    pub fn safe_coding_with_dispatchers(sandbox_root: impl AsRef<Path>) -> anyhow::Result<Self> {
520        let mut registry = ToolRegistryV1::default();
521        for (descriptor, enabled) in safe_coding_tool_plan() {
522            match descriptor.tool_id().as_str() {
523                "aidens:repo-read:1" => {
524                    registry.register_enabled_with_repo_read_dispatcher(
525                        descriptor,
526                        enabled,
527                        sandbox_root.as_ref(),
528                    )?;
529                }
530                "aidens:repo-list:1" => {
531                    registry.register_enabled_with_executor(
532                        descriptor,
533                        enabled,
534                        sandbox_root.as_ref(),
535                        |sandbox_root| ToolExecutorV1::RepoList { sandbox_root },
536                    )?;
537                }
538                "aidens:file-stat:1" => {
539                    registry.register_enabled_with_executor(
540                        descriptor,
541                        enabled,
542                        sandbox_root.as_ref(),
543                        |sandbox_root| ToolExecutorV1::FileStat { sandbox_root },
544                    )?;
545                }
546                "aidens:repo-search:1" => {
547                    registry.register_enabled_with_executor(
548                        descriptor,
549                        enabled,
550                        sandbox_root.as_ref(),
551                        |sandbox_root| ToolExecutorV1::RepoSearch { sandbox_root },
552                    )?;
553                }
554                "aidens:patch-propose:1" => {
555                    registry.register_enabled_with_executor(
556                        descriptor,
557                        enabled,
558                        sandbox_root.as_ref(),
559                        |sandbox_root| ToolExecutorV1::PatchPropose { sandbox_root },
560                    )?;
561                }
562                "aidens:patch-apply:1" => {
563                    registry.register_enabled_with_executor(
564                        descriptor,
565                        enabled,
566                        sandbox_root.as_ref(),
567                        |sandbox_root| ToolExecutorV1::PatchApply { sandbox_root },
568                    )?;
569                }
570                "aidens:run-checks:1" => {
571                    registry.register_enabled_with_executor(
572                        descriptor,
573                        enabled,
574                        sandbox_root.as_ref(),
575                        |sandbox_root| ToolExecutorV1::RunChecks { sandbox_root },
576                    )?;
577                }
578                _ => {
579                    registry.register_enabled(descriptor, enabled);
580                }
581            }
582        }
583        Ok(registry)
584    }
585}
586
587pub fn safe_coding_registry() -> ToolRegistryV1 {
588    let mut registry = ToolRegistryV1::default();
589    for (descriptor, enabled) in safe_coding_tool_plan() {
590        registry.register_enabled(descriptor, enabled);
591    }
592    registry
593}
594
595pub fn safe_coding_registry_for_current_dir() -> ToolRegistryV1 {
596    safe_coding_registry_for_sandbox_root(".")
597}
598
599pub fn safe_coding_registry_for_sandbox_root(root: impl AsRef<Path>) -> ToolRegistryV1 {
600    ToolRegistryV1::safe_coding_with_dispatchers(root.as_ref()).unwrap_or_else(|error| {
601        let mut registry = safe_coding_registry();
602        registry.construction_degradation_reasons.push(format!(
603            "safe-coding dispatcher construction failed for sandbox root {}: {error}",
604            root.as_ref().display()
605        ));
606        registry
607    })
608}
609
610pub fn registry_from_enabled_bundles(
611    enabled_bundles: &[String],
612    sandbox_root: Option<&str>,
613) -> ToolRegistryV1 {
614    if !enabled_bundles.iter().any(|bundle| {
615        matches!(
616            bundle.as_str(),
617            "safe-coding"
618                | "repo-read"
619                | "repo-list"
620                | "file-stat"
621                | "repo-search"
622                | "patch-propose"
623                | "patch-apply"
624                | "run-checks"
625        )
626    }) {
627        return ToolRegistryV1::default();
628    }
629
630    sandbox_root
631        .and_then(|root| ToolRegistryV1::safe_coding_with_dispatchers(root).ok())
632        .unwrap_or_else(safe_coding_registry)
633}
aidens_tool_kit/registry.rs

aidens_tool_kit/
registry.rs
