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 VT Code 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::fs;
10use std::path::Path;
11use std::sync::Arc;
12use std::time::{Duration, SystemTime, UNIX_EPOCH};
13
14use anyhow::{Context, Result};
15
16use crate::config::models::Provider;
17use crate::config::types::{AgentConfig, SessionInfo};
18use crate::models_manager::ModelsManager;
19use crate::utils::error_messages::{ERR_CREATE_DIR, ERR_GET_METADATA};
20
21use crate::core::decision_tracker::DecisionTracker;
22use crate::core::error_recovery::ErrorRecoveryManager;
23use crate::ctx_err;
24use crate::llm::client::{AnyClient, ProviderClientAdapter};
25use crate::llm::factory::{ProviderConfig, create_provider_with_config, infer_provider_from_model};
26use crate::tools::ToolRegistry;
27use tracing::warn;
28
29/// Collection of dependencies required by the [`Agent`](super::core::Agent).
30///
31/// Consumers that want to reuse the agent loop can either construct this bundle
32/// directly with [`AgentComponentBuilder`] or provide their own specialized
33/// implementation when embedding VT Code.
34pub struct AgentComponentSet {
35    pub client: AnyClient,
36    pub tool_registry: Arc<ToolRegistry>,
37    pub decision_tracker: DecisionTracker,
38    pub error_recovery: ErrorRecoveryManager,
39    pub models_manager: Arc<ModelsManager>,
40
41    pub session_info: SessionInfo,
42}
43
44/// Builder for [`AgentComponentSet`].
45///
46/// The builder exposes hooks for overriding individual components which makes
47/// the agent easier to adapt for open-source scenarios or bespoke deployments.
48pub struct AgentComponentBuilder<'config> {
49    config: &'config AgentConfig,
50    client: Option<AnyClient>,
51    tool_registry: Option<Arc<ToolRegistry>>,
52    decision_tracker: Option<DecisionTracker>,
53    error_recovery: Option<ErrorRecoveryManager>,
54    models_manager: Option<Arc<ModelsManager>>,
55
56    session_info: Option<SessionInfo>,
57}
58
59impl<'config> AgentComponentBuilder<'config> {
60    /// Create a new builder scoped to the provided configuration.
61    pub fn new(config: &'config AgentConfig) -> Self {
62        Self {
63            config,
64            client: None,
65            tool_registry: None,
66            decision_tracker: None,
67            error_recovery: None,
68            models_manager: None,
69            session_info: None,
70        }
71    }
72
73    /// Override the LLM client instance.
74    pub fn with_client(mut self, client: AnyClient) -> Self {
75        self.client = Some(client);
76        self
77    }
78
79    /// Override the tool registry instance.
80    pub fn with_tool_registry(mut self, registry: Arc<ToolRegistry>) -> Self {
81        self.tool_registry = Some(registry);
82        self
83    }
84
85    /// Override the decision tracker instance.
86    pub fn with_decision_tracker(mut self, tracker: DecisionTracker) -> Self {
87        self.decision_tracker = Some(tracker);
88        self
89    }
90
91    /// Override the error recovery manager instance.
92    pub fn with_error_recovery(mut self, manager: ErrorRecoveryManager) -> Self {
93        self.error_recovery = Some(manager);
94        self
95    }
96
97    /// Override the models manager instance.
98    pub fn with_models_manager(mut self, manager: Arc<ModelsManager>) -> Self {
99        self.models_manager = Some(manager);
100        self
101    }
102
103    /// Override the session metadata.
104    pub fn with_session_info(mut self, session_info: SessionInfo) -> Self {
105        self.session_info = Some(session_info);
106        self
107    }
108
109    /// Build the component set, lazily constructing any missing dependencies.
110    pub async fn build(mut self) -> Result<AgentComponentSet> {
111        ensure_workspace_ready(&self.config.workspace)?;
112
113        let client = match self.client.take() {
114            Some(client) => client,
115            None => create_llm_client(self.config)?,
116        };
117
118        let session_info = match self.session_info.take() {
119            Some(info) => info,
120            None => create_session_info().context("Failed to initialize agent session metadata for bootstrap")?,
121        };
122
123        let tool_registry = match self.tool_registry {
124            Some(registry) => registry,
125            None => {
126                let registry = ToolRegistry::new(self.config.workspace.clone()).await;
127                registry.set_harness_session(session_info.session_id.clone());
128                Arc::new(registry)
129            }
130        };
131
132        // Prefer custom manager if provided, otherwise reuse global singleton.
133        // The global singleton is provider-agnostic; provider filtering happens at query time.
134        let models_manager = self.models_manager.take().unwrap_or_else(|| {
135            // Clone Arc from global - this is cheap since ModelsManager is behind LazyLock
136            Arc::new(ModelsManager::with_provider(self.config.provider.parse::<Provider>().ok().unwrap_or_default()))
137        });
138
139        let decision_tracker = self.decision_tracker.unwrap_or_default();
140
141        let error_recovery = self.error_recovery.unwrap_or_default();
142
143        Ok(AgentComponentSet {
144            client,
145            tool_registry,
146            decision_tracker,
147            error_recovery,
148            models_manager,
149
150            session_info,
151        })
152    }
153}
154
155fn create_llm_client(config: &AgentConfig) -> Result<AnyClient> {
156    let provider_name = if config.provider.trim().is_empty() {
157        infer_provider_from_model(&config.model)
158            .map(|provider| provider.to_string())
159            .ok_or_else(|| anyhow::anyhow!("Cannot determine provider for model: {}", config.model))?
160    } else {
161        config.provider.to_lowercase()
162    };
163
164    let provider = create_provider_with_config(
165        &provider_name,
166        ProviderConfig {
167            api_key: Some(config.api_key.clone()),
168            openai_chatgpt_auth: config.openai_chatgpt_auth.clone(),
169            copilot_auth: None,
170            base_url: None,
171            model: Some(config.model.clone()),
172            prompt_cache: Some(config.prompt_cache.clone()),
173            timeouts: None,
174            openai: None,
175            anthropic: None,
176            model_behavior: config.model_behavior.clone(),
177            workspace_root: Some(config.workspace.clone()),
178        },
179    )
180    .with_context(|| format!("Failed to initialize provider '{provider_name}'"))?;
181
182    Ok(Box::new(ProviderClientAdapter::new(provider, config.model.clone())))
183}
184
185fn create_session_info() -> Result<SessionInfo> {
186    let now = SystemTime::now().duration_since(UNIX_EPOCH).map_err(|err| err.duration());
187
188    Ok(build_session_info(now))
189}
190
191fn build_session_info(duration: Result<Duration, Duration>) -> SessionInfo {
192    let (start_time, session_id) = match duration {
193        Ok(duration) => {
194            let secs = duration.as_secs();
195            (secs, format!("session_{secs}"))
196        }
197        Err(delta) => {
198            let fallback = delta.as_secs();
199            warn!(fallback_seconds = fallback, "System time is before UNIX epoch; using fallback session id");
200            (fallback, format!("session_fallback_{fallback}"))
201        }
202    };
203
204    SessionInfo {
205        session_id,
206        start_time,
207        total_turns: 0,
208        total_decisions: 0,
209        error_count: 0,
210    }
211}
212
213fn ensure_workspace_ready(workspace_root: &Path) -> Result<()> {
214    if workspace_root.exists() {
215        let metadata =
216            fs::metadata(workspace_root).with_context(|| ctx_err!(ERR_GET_METADATA, workspace_root.display()))?;
217
218        anyhow::ensure!(metadata.is_dir(), "Workspace path is not a directory: {}", workspace_root.display());
219    } else {
220        fs::create_dir_all(workspace_root).with_context(|| ctx_err!(ERR_CREATE_DIR, workspace_root.display()))?;
221    }
222
223    Ok(())
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use crate::config::constants::models;
230    use crate::config::core::PromptCachingConfig;
231    use crate::config::models::Provider;
232    use crate::config::types::{ModelSelectionSource, ReasoningEffortLevel, UiSurfacePreference};
233    use crate::core::agent::snapshots::{DEFAULT_CHECKPOINTS_ENABLED, DEFAULT_MAX_AGE_DAYS, DEFAULT_MAX_SNAPSHOTS};
234    use std::collections::BTreeMap;
235
236    #[tokio::test]
237    async fn builds_default_component_set() {
238        let temp_dir = tempfile::tempdir().expect("temp dir");
239        let agent_config = AgentConfig {
240            model: models::google::GEMINI_3_FLASH_PREVIEW.to_string(),
241            api_key: "test-api-key".to_owned(),
242            provider: Provider::Gemini.to_string(),
243            api_key_env: Provider::Gemini.default_api_key_env().to_string(),
244            workspace: temp_dir.path().to_path_buf(),
245            verbose: false,
246            theme: "default".to_owned(),
247            reasoning_effort: ReasoningEffortLevel::default(),
248            ui_surface: UiSurfacePreference::Inline,
249            prompt_cache: PromptCachingConfig::default(),
250            model_source: ModelSelectionSource::WorkspaceConfig,
251            custom_api_keys: BTreeMap::new(),
252            checkpointing_enabled: DEFAULT_CHECKPOINTS_ENABLED,
253            checkpointing_storage_dir: None,
254            checkpointing_max_snapshots: DEFAULT_MAX_SNAPSHOTS,
255            checkpointing_max_age_days: Some(DEFAULT_MAX_AGE_DAYS),
256            quiet: false,
257            max_conversation_turns: 1000,
258            model_behavior: None,
259            openai_chatgpt_auth: None,
260        };
261
262        let components = AgentComponentBuilder::new(&agent_config)
263            .build()
264            .await
265            .expect("component build succeeds");
266
267        assert!(components.session_info.session_id.starts_with("session_"));
268        assert_eq!(components.session_info.total_turns, 0);
269        assert!(!components.tool_registry.available_tools().await.is_empty());
270    }
271
272    #[tokio::test]
273    async fn allows_overriding_components() {
274        let temp_dir = tempfile::tempdir().expect("temp dir");
275        let agent_config = AgentConfig {
276            model: models::google::GEMINI_3_FLASH_PREVIEW.to_string(),
277            api_key: "test-api-key".to_owned(),
278            provider: Provider::Gemini.to_string(),
279            api_key_env: Provider::Gemini.default_api_key_env().to_string(),
280            workspace: temp_dir.path().to_path_buf(),
281            verbose: true,
282            theme: "custom".to_owned(),
283            reasoning_effort: ReasoningEffortLevel::High,
284            ui_surface: UiSurfacePreference::Alternate,
285            prompt_cache: PromptCachingConfig::default(),
286            model_source: ModelSelectionSource::WorkspaceConfig,
287            custom_api_keys: BTreeMap::new(),
288            checkpointing_enabled: DEFAULT_CHECKPOINTS_ENABLED,
289            checkpointing_storage_dir: None,
290            checkpointing_max_snapshots: DEFAULT_MAX_SNAPSHOTS,
291            checkpointing_max_age_days: Some(DEFAULT_MAX_AGE_DAYS),
292            quiet: false,
293            max_conversation_turns: 1000,
294            model_behavior: None,
295            openai_chatgpt_auth: None,
296        };
297
298        let custom_session = SessionInfo {
299            session_id: "session_custom".to_owned(),
300            start_time: 42,
301            total_turns: 1,
302            total_decisions: 2,
303            error_count: 3,
304        };
305
306        let registry = Arc::new(ToolRegistry::new(agent_config.workspace.clone()).await);
307        // Install a self-reference so built-in tools are exposed to
308        // `unified_exec` code snippets as callable library functions.
309        registry.set_self_ref(Arc::clone(&registry));
310
311        let components = AgentComponentBuilder::new(&agent_config)
312            .with_session_info(custom_session.clone())
313            .with_tool_registry(Arc::clone(&registry))
314            .build()
315            .await
316            .expect("component build succeeds with overrides");
317
318        assert_eq!(components.session_info.session_id, custom_session.session_id);
319        assert_eq!(components.session_info.start_time, custom_session.start_time);
320        assert_eq!(Arc::as_ptr(&components.tool_registry), Arc::as_ptr(&registry));
321    }
322
323    #[test]
324    fn session_info_uses_fallback_when_clock_is_before_epoch() {
325        let info = build_session_info(Err(Duration::from_secs(42)));
326        assert_eq!(info.session_id, "session_fallback_42");
327        assert_eq!(info.start_time, 42);
328        assert_eq!(info.total_turns, 0);
329    }
330
331    #[tokio::test]
332    async fn rejects_non_directory_workspace() {
333        let temp_dir = tempfile::tempdir().expect("temp dir");
334        let file_path = temp_dir.path().join("not_dir");
335        fs::write(&file_path, "not a dir").expect("write file");
336
337        let agent_config = AgentConfig {
338            workspace: file_path.clone(),
339            ..sample_config(temp_dir.path())
340        };
341
342        let result = AgentComponentBuilder::new(&agent_config).build().await;
343        assert!(result.is_err(), "expected workspace validation to fail");
344    }
345
346    fn sample_config(workspace: &Path) -> AgentConfig {
347        AgentConfig {
348            model: models::google::GEMINI_3_FLASH_PREVIEW.to_string(),
349            api_key: "test-api-key".to_owned(),
350            provider: Provider::Gemini.to_string(),
351            api_key_env: Provider::Gemini.default_api_key_env().to_string(),
352            workspace: workspace.to_path_buf(),
353            verbose: false,
354            theme: "default".to_owned(),
355            reasoning_effort: ReasoningEffortLevel::default(),
356            ui_surface: UiSurfacePreference::Inline,
357            prompt_cache: PromptCachingConfig::default(),
358            model_source: ModelSelectionSource::WorkspaceConfig,
359            custom_api_keys: BTreeMap::new(),
360            checkpointing_enabled: DEFAULT_CHECKPOINTS_ENABLED,
361            checkpointing_storage_dir: None,
362            checkpointing_max_snapshots: DEFAULT_MAX_SNAPSHOTS,
363            checkpointing_max_age_days: Some(DEFAULT_MAX_AGE_DAYS),
364            quiet: false,
365            max_conversation_turns: 1000,
366            model_behavior: None,
367            openai_chatgpt_auth: None,
368        }
369    }
370}
vtcode_core/core/agent/bootstrap.rs

vtcode_core/core/agent/
bootstrap.rs
