vtcode_core/core/agent/
bootstrap.rs1use 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
24pub 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
40pub 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 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 pub fn with_client(mut self, client: AnyClient) -> Self {
74 self.client = Some(client);
75 self
76 }
77
78 pub fn with_tool_registry(mut self, registry: Arc<ToolRegistry>) -> Self {
80 self.tool_registry = Some(registry);
81 self
82 }
83
84 pub fn with_decision_tracker(mut self, tracker: DecisionTracker) -> Self {
86 self.decision_tracker = Some(tracker);
87 self
88 }
89
90 pub fn with_error_recovery(mut self, manager: ErrorRecoveryManager) -> Self {
92 self.error_recovery = Some(manager);
93 self
94 }
95
96 pub fn with_summarizer(mut self, summarizer: ConversationSummarizer) -> Self {
98 self.summarizer = Some(summarizer);
99 self
100 }
101
102 pub fn with_tree_sitter_analyzer(mut self, analyzer: TreeSitterAnalyzer) -> Self {
104 self.tree_sitter_analyzer = Some(analyzer);
105 self
106 }
107
108 pub fn with_compaction_engine(mut self, engine: Arc<CompactionEngine>) -> Self {
110 self.compaction_engine = Some(engine);
111 self
112 }
113
114 pub fn with_session_info(mut self, session_info: SessionInfo) -> Self {
116 self.session_info = Some(session_info);
117 self
118 }
119
120 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::ReasoningEffortLevel;
198
199 #[test]
200 fn builds_default_component_set() {
201 let temp_dir = tempfile::tempdir().expect("temp dir");
202 let agent_config = AgentConfig {
203 model: models::GEMINI_2_5_FLASH_PREVIEW.to_string(),
204 api_key: "test-api-key".to_string(),
205 provider: Provider::Gemini.to_string(),
206 workspace: temp_dir.path().to_path_buf(),
207 verbose: false,
208 theme: "default".to_string(),
209 reasoning_effort: ReasoningEffortLevel::default(),
210 prompt_cache: PromptCachingConfig::default(),
211 };
212
213 let components = AgentComponentBuilder::new(&agent_config)
214 .build()
215 .expect("component build succeeds");
216
217 assert!(components.session_info.session_id.starts_with("session_"));
218 assert_eq!(components.session_info.total_turns, 0);
219 assert!(!components.tool_registry.available_tools().is_empty());
220 }
221
222 #[test]
223 fn allows_overriding_components() {
224 let temp_dir = tempfile::tempdir().expect("temp dir");
225 let agent_config = AgentConfig {
226 model: models::GEMINI_2_5_FLASH_PREVIEW.to_string(),
227 api_key: "test-api-key".to_string(),
228 provider: Provider::Gemini.to_string(),
229 workspace: temp_dir.path().to_path_buf(),
230 verbose: true,
231 theme: "custom".to_string(),
232 reasoning_effort: ReasoningEffortLevel::High,
233 prompt_cache: PromptCachingConfig::default(),
234 };
235
236 let custom_session = SessionInfo {
237 session_id: "session_custom".to_string(),
238 start_time: 42,
239 total_turns: 1,
240 total_decisions: 2,
241 error_count: 3,
242 };
243
244 let registry = Arc::new(ToolRegistry::new(agent_config.workspace.clone()));
245
246 let components = AgentComponentBuilder::new(&agent_config)
247 .with_session_info(custom_session.clone())
248 .with_tool_registry(Arc::clone(®istry))
249 .build()
250 .expect("component build succeeds with overrides");
251
252 assert_eq!(
253 components.session_info.session_id,
254 custom_session.session_id
255 );
256 assert_eq!(
257 components.session_info.start_time,
258 custom_session.start_time
259 );
260 assert_eq!(
261 Arc::as_ptr(&components.tool_registry),
262 Arc::as_ptr(®istry)
263 );
264 }
265}