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        // Include aliases so policy prompts stay in sync with exposed names
13        let available = self.available_tools().await;
14        let mcp_keys = self.mcp_policy_keys().await;
15        self.policy_gateway
16            .lock()
17            .await
18            .sync_available_tools(available, &mcp_keys)
19            .await;
20
21        // Seed default permissions from tool metadata when policy manager is present
22        let policy_seeds = {
23            let assembly = self
24                .tool_assembly
25                .read()
26                .unwrap_or_else(std::sync::PoisonError::into_inner);
27            assembly
28                .policy_seed_metadata()
29                .iter()
30                .map(|(name, metadata)| (name.clone(), metadata.clone()))
31                .collect::<Vec<_>>()
32        };
33        let mut policy_gateway = self.policy_gateway.lock().await;
34        if let Ok(policy) = policy_gateway.policy_manager_mut() {
35            let mut seeded = 0usize;
36            for (name, metadata) in policy_seeds {
37                if let Some(default_policy) = metadata.default_permission() {
38                    let current = policy.get_policy(&name);
39                    if matches!(current, ToolPolicy::Prompt) {
40                        if let Err(err) = policy
41                            .seed_default_policy(&name, default_policy.clone())
42                            .await
43                        {
44                            tracing::warn!(
45                                tool = %name,
46                                error = %err,
47                                "Failed to seed default policy from tool metadata"
48                            );
49                        } else {
50                            seeded += 1;
51                        }
52                    }
53                }
54            }
55
56            if seeded > 0 {
57                tracing::trace!(seeded, "Seeded default tool policies from tool assembly");
58            }
59        }
60    }
61
62    pub(super) fn initialize_resiliency_trackers(&self) {
63        let categories = [
64            ToolTimeoutCategory::Default,
65            ToolTimeoutCategory::Pty,
66            ToolTimeoutCategory::Mcp,
67        ];
68        let mut state = self.resiliency.lock();
69        for category in categories {
70            state.failure_trackers.entry(category).or_default();
71            state.success_trackers.entry(category).or_insert(0);
72            state
73                .latency_stats
74                .entry(category)
75                .or_insert_with(|| ToolLatencyStats::new(50));
76            state
77                .adaptive_timeout_ceiling
78                .entry(category)
79                .or_insert_with(|| Duration::from_secs(0));
80        }
81    }
82
83    pub async fn initialize_async(&self) -> Result<()> {
84        let mcp_client_is_none = {
85            self.mcp_client
86                .read()
87                .ok()
88                .map(|g| g.is_none())
89                .unwrap_or(true)
90        };
91        if self.initialized.load(std::sync::atomic::Ordering::Relaxed)
92            && (mcp_client_is_none || !self.mcp_tool_index.read().await.is_empty())
93        {
94            return Ok(());
95        }
96
97        let mcp_client_is_some = {
98            self.mcp_client
99                .read()
100                .ok()
101                .map(|g| g.is_some())
102                .unwrap_or(false)
103        };
104        if mcp_client_is_some
105            && self.mcp_tool_index.read().await.is_empty()
106            && let Err(err) = self.refresh_mcp_tools().await
107        {
108            tracing::warn!(
109                error = %err,
110                "Failed to refresh MCP tools during registry initialization"
111            );
112        }
113
114        self.prewarm_search_runtime();
115        self.sync_policy_catalog().await;
116        self.initialized
117            .store(true, std::sync::atomic::Ordering::Relaxed);
118
119        Ok(())
120    }
121}