Skip to main content

vtcode_core/tools/registry/
maintenance.rs

1//! Internal maintenance helpers for ToolRegistry.
2
3use std::time::Duration;
4
5use crate::tool_policy::ToolPolicy;
6use anyhow::Result;
7
8use super::{ToolLatencyStats, ToolRegistry, ToolTimeoutCategory};
9
10impl ToolRegistry {
11    pub(super) async fn sync_policy_catalog(&self) {
12        let lifecycle = {
13            let policy_gateway = self.policy_gateway.clone();
14            policy_gateway.full_auto_catalogue_lifecycle()
15        };
16        let _lifecycle_guard = lifecycle.lock().await;
17        self.sync_policy_catalog_serialized().await;
18    }
19
20    async fn sync_policy_catalog_serialized(&self) {
21        // Include aliases so policy prompts stay in sync with exposed names
22        let mut available = self.available_tools().await;
23        available.extend(
24            self.inventory
25                .registrations_snapshot()
26                .into_iter()
27                .filter(|registration| {
28                    matches!(registration.catalog_source(), super::registration::ToolCatalogSource::Dynamic)
29                })
30                .map(|registration| registration.name().to_string()),
31        );
32        let mcp_keys = self.mcp_policy_keys().await;
33        let full_auto_catalogue_config = self.policy_gateway.full_auto_catalogue_config().await;
34        let full_auto_visible_policy_names = if let Some(config) = &full_auto_catalogue_config {
35            Some(self.visible_policy_names(config.clone()).await)
36        } else {
37            None
38        };
39        #[cfg(test)]
40        {
41            let test_hooks = self.policy_gateway.full_auto_catalogue_test_hooks();
42            test_hooks.pause_after_refresh_snapshot().await;
43        }
44        {
45            let policy_gateway = self.policy_gateway.clone();
46            policy_gateway.sync_available_tools(available, &mcp_keys).await;
47        }
48        if let (Some(config), Some(visible_policy_names)) =
49            (full_auto_catalogue_config.as_ref(), full_auto_visible_policy_names.as_ref())
50        {
51            let policy_gateway = self.policy_gateway.clone();
52            policy_gateway.refresh_full_auto_catalogue(config, visible_policy_names).await;
53        }
54
55        // Seed default permissions from tool metadata when policy manager is present
56        let policy_seeds = {
57            let assembly = self.tool_assembly.read().unwrap_or_else(std::sync::PoisonError::into_inner);
58            assembly
59                .policy_seed_metadata()
60                .iter()
61                .map(|(name, metadata)| (name.clone(), metadata.clone()))
62                .collect::<Vec<_>>()
63        };
64        let policy_arc = self.policy_gateway.tool_policy_arc();
65        let mut seeds_to_apply: Vec<(String, ToolPolicy)> = Vec::new();
66        {
67            let mut policy_lock = policy_arc.lock().await;
68            let policy_guard = match policy_lock.as_mut() {
69                Some(guard) => guard,
70                None => return,
71            };
72            for (name, metadata) in &policy_seeds {
73                if let Some(default_policy) = metadata.default_permission() {
74                    if matches!(policy_guard.get_policy(name), ToolPolicy::Prompt) {
75                        seeds_to_apply.push((name.clone(), default_policy.clone()));
76                    }
77                }
78            }
79        }
80        let mut seeded = 0usize;
81        for (name, policy) in seeds_to_apply {
82            let mut policy_lock = policy_arc.lock().await;
83            if let Some(policy_guard) = policy_lock.as_mut() {
84                if let Err(err) = policy_guard.seed_default_policy(&name, policy).await {
85                    tracing::warn!(
86                        tool = %name,
87                        error = %err,
88                        "Failed to seed default policy from tool metadata"
89                    );
90                } else {
91                    seeded += 1;
92                }
93            }
94        }
95
96        if seeded > 0 {
97            tracing::trace!(seeded, "Seeded default tool policies from tool assembly");
98        }
99    }
100
101    pub(super) fn initialize_resiliency_trackers(&self) {
102        let categories = [
103            ToolTimeoutCategory::Default,
104            ToolTimeoutCategory::Pty,
105            ToolTimeoutCategory::Mcp,
106            ToolTimeoutCategory::LongRunningCommand,
107        ];
108        let mut state = self.resiliency.lock();
109        for category in categories {
110            state.failure_trackers.entry(category).or_default();
111            state.success_trackers.entry(category).or_insert(0);
112            state.latency_stats.entry(category).or_insert_with(|| ToolLatencyStats::new(50));
113            state
114                .adaptive_timeout_ceiling
115                .entry(category)
116                .or_insert_with(|| Duration::from_secs(0));
117        }
118    }
119
120    pub async fn initialize_async(&self) -> Result<()> {
121        // Acquire the MCP client lock once and derive both flags from it,
122        // avoiding a second read-lock acquisition below.
123        let mcp_client_held = self.mcp_client.read().is_some();
124        if self.initialized.load(std::sync::atomic::Ordering::Relaxed)
125            && (!mcp_client_held || !self.mcp_tool_index.read().await.is_empty())
126        {
127            return Ok(());
128        }
129
130        if mcp_client_held
131            && self.mcp_tool_index.read().await.is_empty()
132            && let Err(err) = self.refresh_mcp_tools().await
133        {
134            tracing::warn!(
135                error = %err,
136                "Failed to refresh MCP tools during registry initialization"
137            );
138        }
139
140        self.prewarm_search_runtime();
141        self.sync_policy_catalog().await;
142        self.initialized.store(true, std::sync::atomic::Ordering::Relaxed);
143
144        Ok(())
145    }
146}