vtcode_core/core/agent/
bootstrap.rs

1//! Component builder and bootstrap utilities for the core agent.
2//!
3//! This module extracts the initialization logic from [`Agent`](super::core::Agent)
4//! so it can be reused by downstream consumers. The builder pattern makes it easy
5//! to override default components (for example when embedding VTCode in other
6//! applications or exposing a reduced open-source surface area) without relying
7//! on the binary crate's internal setup.
8
9use std::sync::Arc;
10use std::time::{SystemTime, UNIX_EPOCH};
11
12use anyhow::{Context, Result};
13
14use crate::config::models::ModelId;
15use crate::config::types::{AgentConfig, SessionInfo};
16use crate::core::agent::compaction::CompactionEngine;
17use crate::core::conversation_summarizer::ConversationSummarizer;
18use crate::core::decision_tracker::DecisionTracker;
19use crate::core::error_recovery::ErrorRecoveryManager;
20use crate::llm::{AnyClient, make_client};
21use crate::tools::ToolRegistry;
22use crate::tools::tree_sitter::TreeSitterAnalyzer;
23
24/// Collection of dependencies required by the [`Agent`](super::core::Agent).
25///
26/// Consumers that want to reuse the agent loop can either construct this bundle
27/// directly with [`AgentComponentBuilder`] or provide their own specialized
28/// implementation when embedding VTCode.
29pub struct AgentComponentSet {
30    pub client: AnyClient,
31    pub tool_registry: Arc<ToolRegistry>,
32    pub decision_tracker: DecisionTracker,
33    pub error_recovery: ErrorRecoveryManager,
34    pub summarizer: ConversationSummarizer,
35    pub tree_sitter_analyzer: TreeSitterAnalyzer,
36    pub compaction_engine: Arc<CompactionEngine>,
37    pub session_info: SessionInfo,
38}
39
40/// Builder for [`AgentComponentSet`].
41///
42/// The builder exposes hooks for overriding individual components which makes
43/// the agent easier to adapt for open-source scenarios or bespoke deployments.
44pub struct AgentComponentBuilder<'config> {
45    config: &'config AgentConfig,
46    client: Option<AnyClient>,
47    tool_registry: Option<Arc<ToolRegistry>>,
48    decision_tracker: Option<DecisionTracker>,
49    error_recovery: Option<ErrorRecoveryManager>,
50    summarizer: Option<ConversationSummarizer>,
51    tree_sitter_analyzer: Option<TreeSitterAnalyzer>,
52    compaction_engine: Option<Arc<CompactionEngine>>,
53    session_info: Option<SessionInfo>,
54}
55
56impl<'config> AgentComponentBuilder<'config> {
57    /// Create a new builder scoped to the provided configuration.
58    pub fn new(config: &'config AgentConfig) -> Self {
59        Self {
60            config,
61            client: None,
62            tool_registry: None,
63            decision_tracker: None,
64            error_recovery: None,
65            summarizer: None,
66            tree_sitter_analyzer: None,
67            compaction_engine: None,
68            session_info: None,
69        }
70    }
71
72    /// Override the LLM client instance.
73    pub fn with_client(mut self, client: AnyClient) -> Self {
74        self.client = Some(client);
75        self
76    }
77
78    /// Override the tool registry instance.
79    pub fn with_tool_registry(mut self, registry: Arc<ToolRegistry>) -> Self {
80        self.tool_registry = Some(registry);
81        self
82    }
83
84    /// Override the decision tracker instance.
85    pub fn with_decision_tracker(mut self, tracker: DecisionTracker) -> Self {
86        self.decision_tracker = Some(tracker);
87        self
88    }
89
90    /// Override the error recovery manager instance.
91    pub fn with_error_recovery(mut self, manager: ErrorRecoveryManager) -> Self {
92        self.error_recovery = Some(manager);
93        self
94    }
95
96    /// Override the conversation summarizer instance.
97    pub fn with_summarizer(mut self, summarizer: ConversationSummarizer) -> Self {
98        self.summarizer = Some(summarizer);
99        self
100    }
101
102    /// Override the tree-sitter analyzer instance.
103    pub fn with_tree_sitter_analyzer(mut self, analyzer: TreeSitterAnalyzer) -> Self {
104        self.tree_sitter_analyzer = Some(analyzer);
105        self
106    }
107
108    /// Override the compaction engine instance.
109    pub fn with_compaction_engine(mut self, engine: Arc<CompactionEngine>) -> Self {
110        self.compaction_engine = Some(engine);
111        self
112    }
113
114    /// Override the session metadata.
115    pub fn with_session_info(mut self, session_info: SessionInfo) -> Self {
116        self.session_info = Some(session_info);
117        self
118    }
119
120    /// Build the component set, lazily constructing any missing dependencies.
121    pub fn build(mut self) -> Result<AgentComponentSet> {
122        let client = match self.client.take() {
123            Some(client) => client,
124            None => create_llm_client(self.config)?,
125        };
126
127        let tree_sitter_analyzer = match self.tree_sitter_analyzer.take() {
128            Some(analyzer) => analyzer,
129            None => TreeSitterAnalyzer::new()
130                .context("Failed to initialize tree-sitter analyzer for agent components")?,
131        };
132
133        let tool_registry = self
134            .tool_registry
135            .unwrap_or_else(|| Arc::new(ToolRegistry::new(self.config.workspace.clone())));
136
137        let decision_tracker = self.decision_tracker.unwrap_or_else(DecisionTracker::new);
138
139        let error_recovery = self
140            .error_recovery
141            .unwrap_or_else(ErrorRecoveryManager::new);
142
143        let summarizer = self.summarizer.unwrap_or_else(ConversationSummarizer::new);
144
145        let compaction_engine = self
146            .compaction_engine
147            .unwrap_or_else(|| Arc::new(CompactionEngine::new()));
148
149        let session_info = match self.session_info.take() {
150            Some(info) => info,
151            None => create_session_info()
152                .context("Failed to initialize agent session metadata for bootstrap")?,
153        };
154
155        Ok(AgentComponentSet {
156            client,
157            tool_registry,
158            decision_tracker,
159            error_recovery,
160            summarizer,
161            tree_sitter_analyzer,
162            compaction_engine,
163            session_info,
164        })
165    }
166}
167
168fn create_llm_client(config: &AgentConfig) -> Result<AnyClient> {
169    let model_id = config
170        .model
171        .parse::<ModelId>()
172        .with_context(|| format!("Invalid model identifier: {}", config.model))?;
173
174    Ok(make_client(config.api_key.clone(), model_id))
175}
176
177fn create_session_info() -> Result<SessionInfo> {
178    let now = SystemTime::now()
179        .duration_since(UNIX_EPOCH)
180        .context("System time is before the UNIX epoch")?;
181
182    Ok(SessionInfo {
183        session_id: format!("session_{}", now.as_secs()),
184        start_time: now.as_secs(),
185        total_turns: 0,
186        total_decisions: 0,
187        error_count: 0,
188    })
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use crate::config::constants::models;
195    use crate::config::core::PromptCachingConfig;
196    use crate::config::models::Provider;
197    use crate::config::types::{ModelSelectionSource, ReasoningEffortLevel, UiSurfacePreference};
198    use std::collections::BTreeMap;
199
200    #[test]
201    fn builds_default_component_set() {
202        let temp_dir = tempfile::tempdir().expect("temp dir");
203        let agent_config = AgentConfig {
204            model: models::GEMINI_2_5_FLASH_PREVIEW.to_string(),
205            api_key: "test-api-key".to_string(),
206            provider: Provider::Gemini.to_string(),
207            api_key_env: Provider::Gemini.default_api_key_env().to_string(),
208            workspace: temp_dir.path().to_path_buf(),
209            verbose: false,
210            theme: "default".to_string(),
211            reasoning_effort: ReasoningEffortLevel::default(),
212            ui_surface: UiSurfacePreference::Inline,
213            prompt_cache: PromptCachingConfig::default(),
214            model_source: ModelSelectionSource::WorkspaceConfig,
215            custom_api_keys: BTreeMap::new(),
216        };
217
218        let components = AgentComponentBuilder::new(&agent_config)
219            .build()
220            .expect("component build succeeds");
221
222        assert!(components.session_info.session_id.starts_with("session_"));
223        assert_eq!(components.session_info.total_turns, 0);
224        assert!(!components.tool_registry.available_tools().is_empty());
225    }
226
227    #[test]
228    fn allows_overriding_components() {
229        let temp_dir = tempfile::tempdir().expect("temp dir");
230        let agent_config = AgentConfig {
231            model: models::GEMINI_2_5_FLASH_PREVIEW.to_string(),
232            api_key: "test-api-key".to_string(),
233            provider: Provider::Gemini.to_string(),
234            api_key_env: Provider::Gemini.default_api_key_env().to_string(),
235            workspace: temp_dir.path().to_path_buf(),
236            verbose: true,
237            theme: "custom".to_string(),
238            reasoning_effort: ReasoningEffortLevel::High,
239            ui_surface: UiSurfacePreference::Alternate,
240            prompt_cache: PromptCachingConfig::default(),
241            model_source: ModelSelectionSource::WorkspaceConfig,
242            custom_api_keys: BTreeMap::new(),
243        };
244
245        let custom_session = SessionInfo {
246            session_id: "session_custom".to_string(),
247            start_time: 42,
248            total_turns: 1,
249            total_decisions: 2,
250            error_count: 3,
251        };
252
253        let registry = Arc::new(ToolRegistry::new(agent_config.workspace.clone()));
254
255        let components = AgentComponentBuilder::new(&agent_config)
256            .with_session_info(custom_session.clone())
257            .with_tool_registry(Arc::clone(&registry))
258            .build()
259            .expect("component build succeeds with overrides");
260
261        assert_eq!(
262            components.session_info.session_id,
263            custom_session.session_id
264        );
265        assert_eq!(
266            components.session_info.start_time,
267            custom_session.start_time
268        );
269        assert_eq!(
270            Arc::as_ptr(&components.tool_registry),
271            Arc::as_ptr(&registry)
272        );
273    }
274}
vtcode_core/core/agent/bootstrap.rs

vtcode_core/core/agent/
bootstrap.rs
