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            openrouter: None,
177            model_behavior: config.model_behavior.clone(),
178            workspace_root: Some(config.workspace.clone()),
179        },
180    )
181    .with_context(|| format!("Failed to initialize provider '{provider_name}'"))?;
182
183    Ok(Box::new(ProviderClientAdapter::new(provider, config.model.clone())))
184}
185
186fn create_session_info() -> Result<SessionInfo> {
187    let now = SystemTime::now().duration_since(UNIX_EPOCH).map_err(|err| err.duration());
188
189    Ok(build_session_info(now))
190}
191
192fn build_session_info(duration: Result<Duration, Duration>) -> SessionInfo {
193    let (start_time, session_id) = match duration {
194        Ok(duration) => {
195            let secs = duration.as_secs();
196            (secs, format!("session_{secs}"))
197        }
198        Err(delta) => {
199            let fallback = delta.as_secs();
200            warn!(fallback_seconds = fallback, "System time is before UNIX epoch; using fallback session id");
201            (fallback, format!("session_fallback_{fallback}"))
202        }
203    };
204
205    SessionInfo {
206        session_id,
207        start_time,
208        total_turns: 0,
209        total_decisions: 0,
210        error_count: 0,
211    }
212}
213
214fn ensure_workspace_ready(workspace_root: &Path) -> Result<()> {
215    if workspace_root.exists() {
216        let metadata =
217            fs::metadata(workspace_root).with_context(|| ctx_err!(ERR_GET_METADATA, workspace_root.display()))?;
218
219        anyhow::ensure!(metadata.is_dir(), "Workspace path is not a directory: {}", workspace_root.display());
220    } else {
221        fs::create_dir_all(workspace_root).with_context(|| ctx_err!(ERR_CREATE_DIR, workspace_root.display()))?;
222    }
223
224    Ok(())
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use crate::config::constants::models;
231    use crate::config::core::PromptCachingConfig;
232    use crate::config::models::Provider;
233    use crate::config::types::{ModelSelectionSource, ReasoningEffortLevel, UiSurfacePreference};
234    use crate::core::agent::snapshots::{DEFAULT_CHECKPOINTS_ENABLED, DEFAULT_MAX_AGE_DAYS, DEFAULT_MAX_SNAPSHOTS};
235    use std::collections::BTreeMap;
236
237    #[tokio::test]
238    async fn builds_default_component_set() {
239        let temp_dir = tempfile::tempdir().expect("temp dir");
240        let agent_config = AgentConfig {
241            model: models::google::GEMINI_3_FLASH_PREVIEW.to_string(),
242            api_key: "test-api-key".to_owned(),
243            provider: Provider::Gemini.to_string(),
244            api_key_env: Provider::Gemini.default_api_key_env().to_string(),
245            workspace: temp_dir.path().to_path_buf(),
246            verbose: false,
247            theme: "default".to_owned(),
248            reasoning_effort: ReasoningEffortLevel::default(),
249            ui_surface: UiSurfacePreference::Inline,
250            prompt_cache: PromptCachingConfig::default(),
251            model_source: ModelSelectionSource::WorkspaceConfig,
252            custom_api_keys: BTreeMap::new(),
253            checkpointing_enabled: DEFAULT_CHECKPOINTS_ENABLED,
254            checkpointing_storage_dir: None,
255            checkpointing_max_snapshots: DEFAULT_MAX_SNAPSHOTS,
256            checkpointing_max_age_days: Some(DEFAULT_MAX_AGE_DAYS),
257            quiet: false,
258            max_conversation_turns: 1000,
259            model_behavior: None,
260            openai_chatgpt_auth: None,
261        };
262
263        let components = AgentComponentBuilder::new(&agent_config)
264            .build()
265            .await
266            .expect("component build succeeds");
267
268        assert!(components.session_info.session_id.starts_with("session_"));
269        assert_eq!(components.session_info.total_turns, 0);
270        assert!(!components.tool_registry.available_tools().await.is_empty());
271    }
272
273    #[tokio::test]
274    async fn allows_overriding_components() {
275        let temp_dir = tempfile::tempdir().expect("temp dir");
276        let agent_config = AgentConfig {
277            model: models::google::GEMINI_3_FLASH_PREVIEW.to_string(),
278            api_key: "test-api-key".to_owned(),
279            provider: Provider::Gemini.to_string(),
280            api_key_env: Provider::Gemini.default_api_key_env().to_string(),
281            workspace: temp_dir.path().to_path_buf(),
282            verbose: true,
283            theme: "custom".to_owned(),
284            reasoning_effort: ReasoningEffortLevel::High,
285            ui_surface: UiSurfacePreference::Alternate,
286            prompt_cache: PromptCachingConfig::default(),
287            model_source: ModelSelectionSource::WorkspaceConfig,
288            custom_api_keys: BTreeMap::new(),
289            checkpointing_enabled: DEFAULT_CHECKPOINTS_ENABLED,
290            checkpointing_storage_dir: None,
291            checkpointing_max_snapshots: DEFAULT_MAX_SNAPSHOTS,
292            checkpointing_max_age_days: Some(DEFAULT_MAX_AGE_DAYS),
293            quiet: false,
294            max_conversation_turns: 1000,
295            model_behavior: None,
296            openai_chatgpt_auth: None,
297        };
298
299        let custom_session = SessionInfo {
300            session_id: "session_custom".to_owned(),
301            start_time: 42,
302            total_turns: 1,
303            total_decisions: 2,
304            error_count: 3,
305        };
306
307        let registry = Arc::new(ToolRegistry::new(agent_config.workspace.clone()).await);
308        // Install a self-reference so built-in tools are exposed to
309        // `unified_exec` code snippets as callable library functions.
310        registry.set_self_ref(Arc::clone(&registry));
311
312        let components = AgentComponentBuilder::new(&agent_config)
313            .with_session_info(custom_session.clone())
314            .with_tool_registry(Arc::clone(&registry))
315            .build()
316            .await
317            .expect("component build succeeds with overrides");
318
319        assert_eq!(components.session_info.session_id, custom_session.session_id);
320        assert_eq!(components.session_info.start_time, custom_session.start_time);
321        assert_eq!(Arc::as_ptr(&components.tool_registry), Arc::as_ptr(&registry));
322    }
323
324    #[test]
325    fn session_info_uses_fallback_when_clock_is_before_epoch() {
326        let info = build_session_info(Err(Duration::from_secs(42)));
327        assert_eq!(info.session_id, "session_fallback_42");
328        assert_eq!(info.start_time, 42);
329        assert_eq!(info.total_turns, 0);
330    }
331
332    #[tokio::test]
333    async fn rejects_non_directory_workspace() {
334        let temp_dir = tempfile::tempdir().expect("temp dir");
335        let file_path = temp_dir.path().join("not_dir");
336        fs::write(&file_path, "not a dir").expect("write file");
337
338        let agent_config = AgentConfig {
339            workspace: file_path.clone(),
340            ..sample_config(temp_dir.path())
341        };
342
343        let result = AgentComponentBuilder::new(&agent_config).build().await;
344        assert!(result.is_err(), "expected workspace validation to fail");
345    }
346
347    fn sample_config(workspace: &Path) -> AgentConfig {
348        AgentConfig {
349            model: models::google::GEMINI_3_FLASH_PREVIEW.to_string(),
350            api_key: "test-api-key".to_owned(),
351            provider: Provider::Gemini.to_string(),
352            api_key_env: Provider::Gemini.default_api_key_env().to_string(),
353            workspace: workspace.to_path_buf(),
354            verbose: false,
355            theme: "default".to_owned(),
356            reasoning_effort: ReasoningEffortLevel::default(),
357            ui_surface: UiSurfacePreference::Inline,
358            prompt_cache: PromptCachingConfig::default(),
359            model_source: ModelSelectionSource::WorkspaceConfig,
360            custom_api_keys: BTreeMap::new(),
361            checkpointing_enabled: DEFAULT_CHECKPOINTS_ENABLED,
362            checkpointing_storage_dir: None,
363            checkpointing_max_snapshots: DEFAULT_MAX_SNAPSHOTS,
364            checkpointing_max_age_days: Some(DEFAULT_MAX_AGE_DAYS),
365            quiet: false,
366            max_conversation_turns: 1000,
367            model_behavior: None,
368            openai_chatgpt_auth: None,
369        }
370    }
371}
vtcode_core/core/agent/bootstrap.rs

vtcode_core/core/agent/
bootstrap.rs
