vtcode_core/tools/registry/
cgp_facade.rs

1//! CGP integration facade for ToolRegistry.
2//!
3//! Provides `enable_cgp_pipeline()` which prefers tool-specific native CGP
4//! facades when available and otherwise wraps registered `TraitObject` tools
5//! through the CGP approval → sandbox → logging/cache/retry pipeline while
6//! preserving registration-sourced metadata.
7
8use std::borrow::Cow;
9use std::path::PathBuf;
10use std::sync::Arc;
11
12use super::ToolRegistry;
13use super::registration::{ToolExecutorFn, ToolHandler, ToolRegistration};
14use crate::components::{wrap_native_tool_ci, wrap_native_tool_interactive, wrap_tool_ci, wrap_tool_interactive};
15use crate::tool_policy::ToolPolicy;
16use crate::tools::result::ToolResult as SplitToolResult;
17use crate::tools::traits::Tool;
18use anyhow::Result;
19use async_trait::async_trait;
20use serde_json::Value;
21
22fn registration_uses_live_task_tracker_metadata(registration: &ToolRegistration) -> bool {
23    registration.name() == crate::config::constants::tools::TASK_TRACKER
24}
25
26fn leak_pattern_str(value: impl Into<String>) -> &'static str {
27    Box::leak(value.into().into_boxed_str())
28}
29
30fn leak_patterns(patterns: &[String]) -> Option<&'static [&'static str]> {
31    if patterns.is_empty() {
32        return None;
33    }
34
35    let leaked_patterns = patterns
36        .iter()
37        .cloned()
38        .map(leak_pattern_str)
39        .collect::<Vec<_>>()
40        .into_boxed_slice();
41    Some(Box::leak(leaked_patterns))
42}
43
44#[derive(Clone)]
45struct RegistrationMetadataSnapshot {
46    name: Arc<str>,
47    description: Arc<str>,
48    parameter_schema: Option<Value>,
49    config_schema: Option<Value>,
50    state_schema: Option<Value>,
51    prompt_path: Option<String>,
52    default_permission: ToolPolicy,
53    allow_patterns: Option<&'static [&'static str]>,
54    deny_patterns: Option<&'static [&'static str]>,
55}
56
57impl RegistrationMetadataSnapshot {
58    fn from_registration(registration: &ToolRegistration) -> Self {
59        Self {
60            name: Arc::<str>::from(registration.name()),
61            description: Arc::<str>::from(registration.metadata().description().unwrap_or_default()),
62            parameter_schema: registration.parameter_schema().cloned(),
63            config_schema: registration.config_schema().cloned(),
64            state_schema: registration.state_schema().cloned(),
65            prompt_path: registration.prompt_path().map(str::to_string),
66            default_permission: registration.default_permission().unwrap_or(ToolPolicy::Prompt),
67            allow_patterns: leak_patterns(registration.metadata().allowlist()),
68            deny_patterns: leak_patterns(registration.metadata().denylist()),
69        }
70    }
71
72    fn from_registration_with_tool<T>(registration: &ToolRegistration, tool: &T) -> Self
73    where
74        T: Tool + ?Sized,
75    {
76        let use_live_task_tracker_metadata = registration_uses_live_task_tracker_metadata(registration);
77        Self {
78            name: Arc::<str>::from(registration.name()),
79            description: if use_live_task_tracker_metadata {
80                Arc::<str>::from(tool.description())
81            } else {
82                registration
83                    .metadata()
84                    .description()
85                    .map(Arc::<str>::from)
86                    .unwrap_or_else(|| Arc::<str>::from(tool.description()))
87            },
88            parameter_schema: if use_live_task_tracker_metadata {
89                tool.parameter_schema()
90            } else {
91                registration.parameter_schema().cloned().or_else(|| tool.parameter_schema())
92            },
93            config_schema: registration.config_schema().cloned().or_else(|| tool.config_schema()),
94            state_schema: registration.state_schema().cloned().or_else(|| tool.state_schema()),
95            prompt_path: registration
96                .prompt_path()
97                .map(str::to_string)
98                .or_else(|| tool.prompt_path().map(Cow::into_owned)),
99            default_permission: registration.default_permission().unwrap_or_else(|| tool.default_permission()),
100            allow_patterns: leak_patterns(registration.metadata().allowlist()).or_else(|| tool.allow_patterns()),
101            deny_patterns: leak_patterns(registration.metadata().denylist()).or_else(|| tool.deny_patterns()),
102        }
103    }
104}
105
106struct RegistryFnTool {
107    registry: ToolRegistry,
108    executor: ToolExecutorFn,
109    metadata: RegistrationMetadataSnapshot,
110}
111
112impl RegistryFnTool {
113    fn from_registration(registry: ToolRegistry, registration: &ToolRegistration) -> Option<Self> {
114        let executor = match registration.handler() {
115            ToolHandler::RegistryFn(executor) => executor,
116            ToolHandler::TraitObject(_) => return None,
117        };
118
119        Some(Self {
120            registry,
121            executor,
122            metadata: RegistrationMetadataSnapshot::from_registration(registration),
123        })
124    }
125}
126
127#[async_trait]
128impl Tool for RegistryFnTool {
129    async fn execute(&self, args: Value) -> Result<Value> {
130        (self.executor)(&self.registry, args).await
131    }
132
133    fn name(&self) -> &str {
134        self.metadata.name.as_ref()
135    }
136
137    fn description(&self) -> &str {
138        self.metadata.description.as_ref()
139    }
140
141    fn parameter_schema(&self) -> Option<Value> {
142        self.metadata.parameter_schema.clone()
143    }
144
145    fn config_schema(&self) -> Option<Value> {
146        self.metadata.config_schema.clone()
147    }
148
149    fn state_schema(&self) -> Option<Value> {
150        self.metadata.state_schema.clone()
151    }
152
153    fn prompt_path(&self) -> Option<Cow<'static, str>> {
154        self.metadata.prompt_path.clone().map(Cow::Owned)
155    }
156
157    fn default_permission(&self) -> ToolPolicy {
158        self.metadata.default_permission.clone()
159    }
160
161    fn allow_patterns(&self) -> Option<&'static [&'static str]> {
162        self.metadata.allow_patterns
163    }
164
165    fn deny_patterns(&self) -> Option<&'static [&'static str]> {
166        self.metadata.deny_patterns
167    }
168}
169
170struct RegistrationBackedTool<T> {
171    inner: T,
172    metadata: RegistrationMetadataSnapshot,
173}
174
175impl<T> RegistrationBackedTool<T>
176where
177    T: Tool + Send + Sync,
178{
179    fn from_registration(inner: T, registration: &ToolRegistration) -> Self {
180        let metadata = RegistrationMetadataSnapshot::from_registration_with_tool(registration, &inner);
181        Self { inner, metadata }
182    }
183}
184
185#[async_trait]
186impl<T> Tool for RegistrationBackedTool<T>
187where
188    T: Tool + Send + Sync,
189{
190    async fn execute(&self, args: Value) -> Result<Value> {
191        self.inner.execute(args).await
192    }
193
194    async fn execute_dual(&self, args: Value) -> Result<SplitToolResult> {
195        let mut result = self.inner.execute_dual(args).await?;
196        result.tool_name = self.name().to_string();
197        Ok(result)
198    }
199
200    fn name(&self) -> &str {
201        self.metadata.name.as_ref()
202    }
203
204    fn description(&self) -> &str {
205        if self.metadata.name.as_ref() == crate::config::constants::tools::TASK_TRACKER {
206            self.inner.description()
207        } else {
208            self.metadata.description.as_ref()
209        }
210    }
211
212    fn validate_args(&self, args: &Value) -> Result<()> {
213        self.inner.validate_args(args)
214    }
215
216    fn parameter_schema(&self) -> Option<Value> {
217        if self.metadata.name.as_ref() == crate::config::constants::tools::TASK_TRACKER {
218            self.inner.parameter_schema()
219        } else {
220            self.metadata.parameter_schema.clone()
221        }
222    }
223
224    fn config_schema(&self) -> Option<Value> {
225        self.metadata.config_schema.clone()
226    }
227
228    fn state_schema(&self) -> Option<Value> {
229        self.metadata.state_schema.clone()
230    }
231
232    fn prompt_path(&self) -> Option<Cow<'static, str>> {
233        self.metadata.prompt_path.clone().map(Cow::Owned)
234    }
235
236    fn default_permission(&self) -> ToolPolicy {
237        self.metadata.default_permission.clone()
238    }
239
240    fn allow_patterns(&self) -> Option<&'static [&'static str]> {
241        self.metadata.allow_patterns
242    }
243
244    fn deny_patterns(&self) -> Option<&'static [&'static str]> {
245        self.metadata.deny_patterns
246    }
247
248    fn is_mutating(&self) -> bool {
249        self.inner.is_mutating()
250    }
251
252    fn is_parallel_safe(&self) -> bool {
253        self.inner.is_parallel_safe()
254    }
255
256    fn kind(&self) -> &'static str {
257        self.inner.kind()
258    }
259
260    fn resource_hints(&self, args: &Value) -> Vec<String> {
261        self.inner.resource_hints(args)
262    }
263
264    fn execution_cost(&self) -> u8 {
265        self.inner.execution_cost()
266    }
267}
268
269struct RegistrationBackedDynTool {
270    inner: Arc<dyn Tool>,
271    metadata: RegistrationMetadataSnapshot,
272}
273
274impl RegistrationBackedDynTool {
275    fn from_registration(inner: Arc<dyn Tool>, registration: &ToolRegistration) -> Self {
276        let metadata = RegistrationMetadataSnapshot::from_registration_with_tool(registration, inner.as_ref());
277        Self { inner, metadata }
278    }
279}
280
281#[async_trait]
282impl Tool for RegistrationBackedDynTool {
283    async fn execute(&self, args: Value) -> Result<Value> {
284        self.inner.execute(args).await
285    }
286
287    async fn execute_dual(&self, args: Value) -> Result<SplitToolResult> {
288        let mut result = self.inner.execute_dual(args).await?;
289        result.tool_name = self.name().to_string();
290        Ok(result)
291    }
292
293    fn name(&self) -> &str {
294        self.metadata.name.as_ref()
295    }
296
297    fn description(&self) -> &str {
298        if self.metadata.name.as_ref() == crate::config::constants::tools::TASK_TRACKER {
299            self.inner.description()
300        } else {
301            self.metadata.description.as_ref()
302        }
303    }
304
305    fn validate_args(&self, args: &Value) -> Result<()> {
306        self.inner.validate_args(args)
307    }
308
309    fn parameter_schema(&self) -> Option<Value> {
310        if self.metadata.name.as_ref() == crate::config::constants::tools::TASK_TRACKER {
311            self.inner.parameter_schema()
312        } else {
313            self.metadata.parameter_schema.clone()
314        }
315    }
316
317    fn config_schema(&self) -> Option<Value> {
318        self.metadata.config_schema.clone()
319    }
320
321    fn state_schema(&self) -> Option<Value> {
322        self.metadata.state_schema.clone()
323    }
324
325    fn prompt_path(&self) -> Option<Cow<'static, str>> {
326        self.metadata.prompt_path.clone().map(Cow::Owned)
327    }
328
329    fn default_permission(&self) -> ToolPolicy {
330        self.metadata.default_permission.clone()
331    }
332
333    fn allow_patterns(&self) -> Option<&'static [&'static str]> {
334        self.metadata.allow_patterns
335    }
336
337    fn deny_patterns(&self) -> Option<&'static [&'static str]> {
338        self.metadata.deny_patterns
339    }
340
341    fn is_mutating(&self) -> bool {
342        self.inner.is_mutating()
343    }
344
345    fn is_parallel_safe(&self) -> bool {
346        self.inner.is_parallel_safe()
347    }
348
349    fn kind(&self) -> &'static str {
350        self.inner.kind()
351    }
352
353    fn resource_hints(&self, args: &Value) -> Vec<String> {
354        self.inner.resource_hints(args)
355    }
356
357    fn execution_cost(&self) -> u8 {
358        self.inner.execution_cost()
359    }
360}
361
362/// Fallback bridge for registrations that already carry shared tool ownership.
363///
364/// Prefer `wrap_registered_native_tool()` when the caller still owns the
365/// concrete tool instance and does not need to preserve an existing
366/// `Arc<dyn Tool>` handle.
367fn wrap_registered_trait_object_tool(
368    registration: &ToolRegistration,
369    tool: Arc<dyn Tool>,
370    workspace_root: PathBuf,
371    mode: CgpRuntimeMode,
372) -> Arc<dyn Tool> {
373    let tool: Arc<dyn Tool> = Arc::new(RegistrationBackedDynTool::from_registration(tool, registration));
374    match mode {
375        CgpRuntimeMode::Interactive => Arc::new(wrap_tool_interactive(tool, workspace_root)),
376        CgpRuntimeMode::Ci => Arc::new(wrap_tool_ci(tool, workspace_root)),
377    }
378}
379
380pub fn wrap_registered_native_tool<T>(
381    registration: &ToolRegistration,
382    tool: T,
383    workspace_root: PathBuf,
384    mode: CgpRuntimeMode,
385) -> Arc<dyn Tool>
386where
387    T: Tool + Send + Sync + 'static,
388{
389    let tool = RegistrationBackedTool::from_registration(tool, registration);
390    match mode {
391        CgpRuntimeMode::Interactive => Arc::new(wrap_native_tool_interactive(tool, workspace_root)),
392        CgpRuntimeMode::Ci => Arc::new(wrap_native_tool_ci(tool, workspace_root)),
393    }
394}
395
396pub fn native_cgp_tool_factory<T, F>(build_tool: F) -> super::registration::NativeCgpToolFactory
397where
398    T: Tool + Send + Sync + 'static,
399    F: Fn() -> T + Send + Sync + 'static,
400{
401    Arc::new(move |registration, workspace_root, mode| {
402        wrap_registered_native_tool(registration, build_tool(), workspace_root, mode)
403    })
404}
405
406/// Runtime mode for CGP pipeline selection.
407#[derive(Debug, Clone, Copy, PartialEq, Eq)]
408pub enum CgpRuntimeMode {
409    /// Interactive TUI sessions: prompt approval + workspace sandbox + tracing.
410    Interactive,
411    /// CI/automation: auto-approval + strict sandbox + no middleware.
412    Ci,
413}
414
415impl ToolRegistry {
416    pub(crate) fn current_cgp_mode(&self) -> Option<CgpRuntimeMode> {
417        *self.cgp_runtime_mode.read().unwrap_or_else(|e| e.into_inner())
418    }
419
420    fn set_cgp_runtime_mode(&self, mode: CgpRuntimeMode) {
421        if let Ok(mut current_mode) = self.cgp_runtime_mode.write() {
422            *current_mode = Some(mode);
423        }
424    }
425
426    pub(crate) fn cgp_handler_for_registration(
427        &self,
428        registration: &ToolRegistration,
429        mode: CgpRuntimeMode,
430    ) -> Option<ToolHandler> {
431        let workspace = self.workspace_root_owned();
432        if let Some(factory) = registration.native_cgp_factory() {
433            return Some(ToolHandler::TraitObject(factory(registration, workspace, mode)));
434        }
435
436        match registration.handler() {
437            ToolHandler::TraitObject(tool) => {
438                Some(ToolHandler::TraitObject(wrap_registered_trait_object_tool(registration, tool, workspace, mode)))
439            }
440            ToolHandler::RegistryFn(_) => {
441                let tool = RegistryFnTool::from_registration(self.clone(), registration)?;
442                Some(ToolHandler::TraitObject(match mode {
443                    CgpRuntimeMode::Interactive => Arc::new(wrap_native_tool_interactive(tool, workspace)),
444                    CgpRuntimeMode::Ci => Arc::new(wrap_native_tool_ci(tool, workspace)),
445                }))
446            }
447        }
448    }
449
450    /// Enable the CGP pipeline for all registered tools.
451    ///
452    /// This replaces each eligible tool's handler with a CGP `ToolFacade`
453    /// determined by the runtime mode. Registrations that provide a native CGP
454    /// factory use that directly; `TraitObject` handlers are wrapped with
455    /// registration-backed metadata before entering the passthrough bridge, and
456    /// `RegistryFn` handlers are projected through a concrete `RegistryFnTool`.
457    pub async fn enable_cgp_pipeline(&self, mode: CgpRuntimeMode) {
458        self.set_cgp_runtime_mode(mode);
459        let snapshot = self.inventory.registrations_snapshot();
460        let mut wrapped_count = 0u32;
461
462        for reg in &snapshot {
463            if reg.is_cgp_wrapped() {
464                continue;
465            }
466
467            let Some(handler) = self.cgp_handler_for_registration(reg, mode) else {
468                continue;
469            };
470
471            if let Err(err) = self.inventory.replace_tool_handler(reg.name(), handler) {
472                tracing::warn!(
473                    tool = %reg.name(),
474                    %err,
475                    "Failed to wrap tool with CGP pipeline"
476                );
477            } else {
478                wrapped_count += 1;
479            }
480        }
481
482        if wrapped_count > 0 {
483            self.rebuild_tool_assembly().await;
484            self.tool_catalog_state.note_explicit_refresh("cgp_pipeline_enable");
485            self.invalidate_hot_cache();
486            tracing::info!(
487                count = wrapped_count,
488                mode = ?mode,
489                "CGP pipeline enabled for registered tools"
490            );
491        }
492    }
493
494    /// Wrap a single tool through the CGP pipeline and register it.
495    ///
496    /// This is the preferred path for new tool registrations that should
497    /// participate in the CGP approval/sandbox/logging/cache/retry pipeline.
498    pub async fn register_cgp_tool(
499        &self,
500        tool: Arc<dyn Tool>,
501        capability: crate::config::types::CapabilityLevel,
502        mode: CgpRuntimeMode,
503    ) -> Result<()> {
504        let workspace = self.workspace_root_owned();
505        let tool_name = Arc::<str>::from(tool.name());
506        let registration = match mode {
507            CgpRuntimeMode::Interactive => {
508                ToolRegistration::from_cgp_tool(tool_name, capability, wrap_tool_interactive(tool, workspace))
509            }
510            CgpRuntimeMode::Ci => ToolRegistration::from_cgp_tool(tool_name, capability, wrap_tool_ci(tool, workspace)),
511        };
512        self.register_tool(registration).await
513    }
514
515    /// Invalidate the hot tool cache after CGP wrapping.
516    fn invalidate_hot_cache(&self) {
517        self.hot_tool_cache.write().clear();
518        *self.cached_available_tools.write() = None;
519    }
520}
521
522#[cfg(test)]
523mod tests {
524    use super::*;
525    use crate::config::constants::tools;
526    use crate::tools::handlers::{PlanningWorkflowState, TaskTrackerTool};
527    use crate::tools::traits::Tool;
528    use futures::future::BoxFuture;
529    use std::path::PathBuf;
530    use tempfile::TempDir;
531
532    struct DummyTool;
533
534    #[async_trait]
535    impl Tool for DummyTool {
536        async fn execute(&self, args: Value) -> Result<Value> {
537            Ok(serde_json::json!({
538                "tool_name": "dummy",
539                "echoed": args,
540            }))
541        }
542
543        fn name(&self) -> &str {
544            "dummy_cgp_test"
545        }
546
547        fn description(&self) -> &str {
548            "A dummy tool for CGP facade tests"
549        }
550    }
551
552    #[tokio::test]
553    async fn enable_cgp_pipeline_wraps_tools() {
554        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
555        let tool: Arc<dyn Tool> = Arc::new(DummyTool);
556        let reg = ToolRegistration::from_tool("dummy_cgp_test", crate::config::types::CapabilityLevel::Basic, tool);
557        registry.register_tool(reg).await.expect("should register");
558
559        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
560
561        let wrapped = registry.get_tool("dummy_cgp_test");
562        assert!(wrapped.is_some(), "tool should still be accessible");
563
564        let result = wrapped
565            .unwrap()
566            .execute(serde_json::json!({"test": true}))
567            .await
568            .expect("should execute");
569        assert_eq!(result.get("echoed").and_then(|v| v.get("test")), Some(&serde_json::json!(true)));
570    }
571
572    #[tokio::test]
573    async fn enable_cgp_pipeline_preserves_registration_metadata_for_trait_object_tools() {
574        struct BridgeTool;
575
576        #[async_trait]
577        impl Tool for BridgeTool {
578            async fn execute(&self, _args: Value) -> Result<Value> {
579                Ok(serde_json::json!({ "path": "bridge" }))
580            }
581
582            async fn execute_dual(&self, _args: Value) -> Result<SplitToolResult> {
583                Ok(SplitToolResult::simple(self.name(), "dual bridge"))
584            }
585
586            fn name(&self) -> &str {
587                "bridge_trait_object"
588            }
589
590            fn description(&self) -> &str {
591                "bridge fallback tool"
592            }
593        }
594
595        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
596        let registration = ToolRegistration::from_tool_with_metadata(
597            "registered_trait_object_cgp_test",
598            crate::config::types::CapabilityLevel::Basic,
599            Arc::new(BridgeTool),
600            crate::tools::registry::ToolMetadata::default()
601                .with_description("registered trait-object tool")
602                .with_parameter_schema(serde_json::json!({
603                    "type": "object",
604                    "properties": {
605                        "query": { "type": "string" }
606                    }
607                }))
608                .with_prompt_path("tools/registered_trait_object.md")
609                .with_permission(ToolPolicy::Allow)
610                .with_allowlist(["tool://allowed"])
611                .with_denylist(["tool://blocked"]),
612        );
613        registry.register_tool(registration).await.expect("should register");
614
615        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
616
617        let tool = registry
618            .get_tool("registered_trait_object_cgp_test")
619            .expect("tool should exist");
620        assert_eq!(tool.name(), "registered_trait_object_cgp_test");
621        assert_eq!(tool.description(), "registered trait-object tool");
622        assert_eq!(tool.prompt_path().as_deref(), Some("tools/registered_trait_object.md"));
623        assert_eq!(tool.default_permission(), ToolPolicy::Allow);
624        assert_eq!(
625            tool.parameter_schema(),
626            Some(serde_json::json!({
627                "type": "object",
628                "properties": {
629                    "query": { "type": "string" }
630                }
631            }))
632        );
633        assert_eq!(tool.allow_patterns(), Some(&["tool://allowed"][..]));
634        assert_eq!(tool.deny_patterns(), Some(&["tool://blocked"][..]));
635
636        let dual = tool
637            .execute_dual(serde_json::json!({ "query": "rust" }))
638            .await
639            .expect("should execute dual");
640        assert_eq!(dual.tool_name, "registered_trait_object_cgp_test");
641    }
642
643    #[tokio::test]
644    async fn enable_cgp_pipeline_ci_mode() {
645        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
646        let tool: Arc<dyn Tool> = Arc::new(DummyTool);
647        let reg = ToolRegistration::from_tool("dummy_cgp_test", crate::config::types::CapabilityLevel::Basic, tool);
648        registry.register_tool(reg).await.expect("should register");
649
650        registry.enable_cgp_pipeline(CgpRuntimeMode::Ci).await;
651
652        let wrapped = registry.get_tool("dummy_cgp_test");
653        assert!(wrapped.is_some());
654
655        let result = wrapped
656            .unwrap()
657            .execute(serde_json::json!({"ci": "mode"}))
658            .await
659            .expect("should execute");
660        assert_eq!(result.get("echoed").and_then(|v| v.get("ci")).and_then(|v| v.as_str()), Some("mode"));
661    }
662
663    #[tokio::test]
664    async fn register_cgp_tool_directly() {
665        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
666        let tool: Arc<dyn Tool> = Arc::new(DummyTool);
667
668        registry
669            .register_cgp_tool(tool, crate::config::types::CapabilityLevel::Basic, CgpRuntimeMode::Interactive)
670            .await
671            .expect("should register");
672
673        let wrapped = registry.get_tool("dummy_cgp_test");
674        assert!(wrapped.is_some());
675    }
676
677    #[tokio::test]
678    async fn enable_cgp_pipeline_skips_already_wrapped_tools() {
679        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
680        let tool: Arc<dyn Tool> = Arc::new(DummyTool);
681
682        registry
683            .register_cgp_tool(tool, crate::config::types::CapabilityLevel::Basic, CgpRuntimeMode::Interactive)
684            .await
685            .expect("should register");
686
687        let before = registry
688            .inventory
689            .registrations_snapshot()
690            .into_iter()
691            .find(|registration| registration.name() == "dummy_cgp_test")
692            .expect("registration should exist");
693        assert!(before.is_cgp_wrapped());
694
695        let before_handler = match before.handler() {
696            ToolHandler::TraitObject(tool) => tool,
697            ToolHandler::RegistryFn(_) => panic!("expected trait object handler"),
698        };
699
700        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
701
702        let after = registry
703            .inventory
704            .registrations_snapshot()
705            .into_iter()
706            .find(|registration| registration.name() == "dummy_cgp_test")
707            .expect("registration should exist");
708        let after_handler = match after.handler() {
709            ToolHandler::TraitObject(tool) => tool,
710            ToolHandler::RegistryFn(_) => panic!("expected trait object handler"),
711        };
712
713        assert!(Arc::ptr_eq(&before_handler, &after_handler));
714    }
715
716    #[tokio::test]
717    async fn enable_cgp_pipeline_prefers_native_cgp_factory() {
718        struct BridgeTool;
719
720        #[async_trait]
721        impl Tool for BridgeTool {
722            async fn execute(&self, _args: Value) -> Result<Value> {
723                Ok(serde_json::json!({ "path": "bridge" }))
724            }
725
726            fn name(&self) -> &str {
727                "native_cgp_factory_test"
728            }
729
730            fn description(&self) -> &str {
731                "bridge fallback tool"
732            }
733        }
734
735        struct NativeTool;
736
737        #[async_trait]
738        impl Tool for NativeTool {
739            async fn execute(&self, _args: Value) -> Result<Value> {
740                Ok(serde_json::json!({ "path": "native" }))
741            }
742
743            fn name(&self) -> &str {
744                "native_cgp_factory_test"
745            }
746
747            fn description(&self) -> &str {
748                "native factory tool"
749            }
750        }
751
752        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
753        let reg = ToolRegistration::from_tool(
754            "native_cgp_factory_test",
755            crate::config::types::CapabilityLevel::Basic,
756            Arc::new(BridgeTool),
757        )
758        .with_description("registered native factory tool")
759        .with_native_cgp_factory(Arc::new(|registration, workspace_root, mode| {
760            wrap_registered_native_tool(registration, NativeTool, workspace_root, mode)
761        }));
762        registry.register_tool(reg).await.expect("should register");
763
764        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
765
766        let tool = registry.get_tool("native_cgp_factory_test").expect("tool should exist");
767        assert_eq!(tool.name(), "native_cgp_factory_test");
768        assert_eq!(tool.description(), "registered native factory tool");
769
770        let result = tool.execute(serde_json::json!({})).await.expect("should execute");
771
772        assert_eq!(result.get("path").and_then(|v| v.as_str()), Some("native"));
773    }
774
775    #[tokio::test]
776    async fn enable_cgp_pipeline_uses_live_task_tracker_metadata_for_native_factory_tools() {
777        let temp_dir = TempDir::new().expect("temp dir");
778        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
779        let plan_state = registry.planning_workflow_state();
780        let registration = registry
781            .inventory
782            .get_registration(tools::TASK_TRACKER)
783            .expect("task_tracker registration exists");
784
785        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
786
787        let tool = registry.get_tool(tools::TASK_TRACKER).expect("wrapped task_tracker tool");
788        assert_eq!(
789            tool.description(),
790            crate::tools::handlers::task_tracker::task_tracker_description_for_workflow(false)
791        );
792        assert_eq!(tool.parameter_schema().expect("standard schema")["properties"]["index"]["minimum"], 0);
793
794        plan_state.enable();
795
796        let tool = registry.get_tool(tools::TASK_TRACKER).expect("wrapped task_tracker tool");
797        assert_eq!(
798            tool.description(),
799            crate::tools::handlers::task_tracker::task_tracker_description_for_workflow(true)
800        );
801        assert_eq!(tool.parameter_schema().expect("planning schema")["properties"]["index"]["minimum"], 1);
802
803        let direct_native = wrap_registered_native_tool(
804            &registration,
805            TaskTrackerTool::new(temp_dir.path().to_path_buf(), plan_state.clone()),
806            temp_dir.path().to_path_buf(),
807            CgpRuntimeMode::Interactive,
808        );
809        assert_eq!(
810            direct_native.description(),
811            crate::tools::handlers::task_tracker::task_tracker_description_for_workflow(true)
812        );
813        assert_eq!(
814            direct_native.parameter_schema().expect("native planning schema")["properties"]["index"]["minimum"],
815            1
816        );
817    }
818
819    #[tokio::test]
820    async fn enable_cgp_pipeline_uses_live_task_tracker_metadata_for_trait_object_tools() {
821        let temp_dir = TempDir::new().expect("temp dir");
822        let registry = ToolRegistry::new(temp_dir.path().to_path_buf()).await;
823        let state = PlanningWorkflowState::new(temp_dir.path().to_path_buf());
824        let registration = ToolRegistration::from_tool(
825            tools::TASK_TRACKER,
826            crate::config::types::CapabilityLevel::Basic,
827            Arc::new(TaskTrackerTool::new(temp_dir.path().to_path_buf(), state.clone())),
828        )
829        .with_description("stale registered task tracker metadata");
830
831        registry.register_tool(registration).await.expect("should register override");
832        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
833
834        let tool = registry.get_tool(tools::TASK_TRACKER).expect("wrapped task_tracker tool");
835        assert_eq!(
836            tool.description(),
837            crate::tools::handlers::task_tracker::task_tracker_description_for_workflow(false)
838        );
839
840        state.enable();
841
842        let tool = registry.get_tool(tools::TASK_TRACKER).expect("wrapped task_tracker tool");
843        assert_eq!(
844            tool.description(),
845            crate::tools::handlers::task_tracker::task_tracker_description_for_workflow(true)
846        );
847        assert_eq!(tool.parameter_schema().expect("planning schema")["properties"]["index"]["minimum"], 1);
848    }
849
850    #[tokio::test]
851    async fn register_tool_after_enabling_cgp_pipeline_wraps_new_tools() {
852        struct LateTool;
853
854        #[async_trait]
855        impl Tool for LateTool {
856            async fn execute(&self, args: Value) -> Result<Value> {
857                Ok(serde_json::json!({
858                    "path": "late-bridge",
859                    "args": args,
860                }))
861            }
862
863            fn name(&self) -> &str {
864                "late_cgp_test"
865            }
866
867            fn description(&self) -> &str {
868                "late registration test"
869            }
870        }
871
872        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
873        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
874
875        registry
876            .register_tool(ToolRegistration::from_tool(
877                "late_cgp_test",
878                crate::config::types::CapabilityLevel::Basic,
879                Arc::new(LateTool),
880            ))
881            .await
882            .expect("should register");
883
884        let registration = registry
885            .inventory
886            .registrations_snapshot()
887            .into_iter()
888            .find(|registration| registration.name() == "late_cgp_test")
889            .expect("registration should exist");
890        assert!(registration.is_cgp_wrapped());
891
892        let result = registry
893            .get_tool("late_cgp_test")
894            .expect("tool should exist")
895            .execute(serde_json::json!({"late": true}))
896            .await
897            .expect("should execute");
898        assert_eq!(result.get("path").and_then(|v| v.as_str()), Some("late-bridge"));
899    }
900
901    #[tokio::test]
902    async fn register_tool_after_enabling_cgp_pipeline_prefers_native_factory() {
903        struct BridgeTool;
904
905        #[async_trait]
906        impl Tool for BridgeTool {
907            async fn execute(&self, _args: Value) -> Result<Value> {
908                Ok(serde_json::json!({ "path": "bridge" }))
909            }
910
911            fn name(&self) -> &str {
912                "late_native_cgp_test"
913            }
914
915            fn description(&self) -> &str {
916                "bridge fallback tool"
917            }
918        }
919
920        struct NativeTool;
921
922        #[async_trait]
923        impl Tool for NativeTool {
924            async fn execute(&self, _args: Value) -> Result<Value> {
925                Ok(serde_json::json!({ "path": "late-native" }))
926            }
927
928            fn name(&self) -> &str {
929                "late_native_cgp_test"
930            }
931
932            fn description(&self) -> &str {
933                "native late tool"
934            }
935        }
936
937        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
938        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
939
940        let registration = ToolRegistration::from_tool(
941            "late_native_cgp_test",
942            crate::config::types::CapabilityLevel::Basic,
943            Arc::new(BridgeTool),
944        )
945        .with_native_cgp_factory(Arc::new(|registration, workspace_root, mode| {
946            wrap_registered_native_tool(registration, NativeTool, workspace_root, mode)
947        }));
948        registry.register_tool(registration).await.expect("should register");
949
950        let result = registry
951            .get_tool("late_native_cgp_test")
952            .expect("tool should exist")
953            .execute(serde_json::json!({}))
954            .await
955            .expect("should execute");
956
957        assert_eq!(result.get("path").and_then(|v| v.as_str()), Some("late-native"));
958    }
959
960    fn registry_fn_test_executor<'a>(_registry: &'a ToolRegistry, args: Value) -> BoxFuture<'a, Result<Value>> {
961        Box::pin(async move {
962            Ok(serde_json::json!({
963                "tool_name": "registry_fn_cgp_test",
964                "echoed": args,
965            }))
966        })
967    }
968
969    #[tokio::test]
970    async fn enable_cgp_pipeline_wraps_registry_fn_tools() {
971        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
972        let registration = ToolRegistration::new(
973            "registry_fn_cgp_test",
974            crate::config::types::CapabilityLevel::Basic,
975            false,
976            registry_fn_test_executor,
977        )
978        .with_description("Registry function CGP test tool")
979        .with_parameter_schema(serde_json::json!({
980            "type": "object",
981            "properties": {
982                "flag": { "type": "boolean" }
983            }
984        }))
985        .with_permission(ToolPolicy::Allow);
986        registry.register_tool(registration).await.expect("should register");
987
988        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
989
990        let wrapped = registry.get_tool("registry_fn_cgp_test").expect("tool exists");
991        assert_eq!(wrapped.name(), "registry_fn_cgp_test");
992        assert_eq!(wrapped.description(), "Registry function CGP test tool");
993        assert!(wrapped.parameter_schema().is_some());
994        assert_eq!(wrapped.default_permission(), ToolPolicy::Allow);
995
996        let result = wrapped
997            .execute(serde_json::json!({"flag": true}))
998            .await
999            .expect("should execute");
1000        assert_eq!(
1001            result
1002                .get("echoed")
1003                .and_then(|value| value.get("flag"))
1004                .and_then(|value| value.as_bool()),
1005            Some(true)
1006        );
1007    }
1008
1009    #[tokio::test]
1010    async fn register_registry_fn_after_enabling_cgp_pipeline_wraps_new_tools() {
1011        let registry = ToolRegistry::new(PathBuf::from("/tmp/test")).await;
1012        registry.enable_cgp_pipeline(CgpRuntimeMode::Interactive).await;
1013
1014        registry
1015            .register_tool(ToolRegistration::new(
1016                "late_registry_fn_cgp_test",
1017                crate::config::types::CapabilityLevel::Basic,
1018                false,
1019                registry_fn_test_executor,
1020            ))
1021            .await
1022            .expect("should register");
1023
1024        let registration = registry
1025            .inventory
1026            .registrations_snapshot()
1027            .into_iter()
1028            .find(|registration| registration.name() == "late_registry_fn_cgp_test")
1029            .expect("registration should exist");
1030        assert!(registration.is_cgp_wrapped());
1031
1032        let result = registry
1033            .get_tool("late_registry_fn_cgp_test")
1034            .expect("tool exists")
1035            .execute(serde_json::json!({"late": true}))
1036            .await
1037            .expect("should execute");
1038        assert_eq!(
1039            result
1040                .get("echoed")
1041                .and_then(|value| value.get("late"))
1042                .and_then(|value| value.as_bool()),
1043            Some(true)
1044        );
1045    }
1046}
vtcode_core/tools/registry/cgp_facade.rs

vtcode_core/tools/registry/
cgp_facade.rs
