use std::sync::Arc;
use std::path::PathBuf;
use tokio::sync::RwLock;
use bamboo_agent_core::tools::{Tool, ToolExecutor};
use bamboo_engine::{AgentBuilder as EngineAgentBuilder, HookRunner};
use bamboo_llm::{create_provider_with_dir, Config, LLMProvider};
use bamboo_mcp::executor::{CompositeToolExecutor, McpToolExecutor};
use bamboo_mcp::manager::McpServerManager;
use bamboo_mcp::McpServerConfig;
use bamboo_metrics::{MetricsCollector, SqliteMetricsStorage};
use bamboo_skills::{SkillManager, SkillStoreConfig};
use bamboo_storage::{LockedSessionStore, SessionStoreV2};
use bamboo_tools::permission::PermissionChecker;
use bamboo_tools::ToolRegistry;
use super::error::SdkError;
use super::Agent;
const DEFAULT_METRICS_RETENTION_DAYS: u32 = 90;
pub struct AgentBuilder {
inner: EngineAgentBuilder,
system_prompt: Option<String>,
tools: Vec<Arc<dyn Tool>>,
model: Option<String>,
provider_name: Option<String>,
api_key: Option<String>,
mcp_servers: Vec<McpServerConfig>,
permission_checker: Option<Arc<dyn PermissionChecker>>,
session_store: Option<Arc<SessionStoreV2>>,
}
impl AgentBuilder {
pub fn new() -> Self {
Self {
inner: EngineAgentBuilder::new(),
system_prompt: None,
tools: Vec::new(),
model: None,
provider_name: None,
api_key: None,
mcp_servers: Vec::new(),
permission_checker: None,
session_store: None,
}
}
pub fn model(mut self, model: impl Into<String>) -> Self {
self.model = Some(model.into());
self
}
pub fn provider_name(mut self, provider: impl Into<String>) -> Self {
self.provider_name = Some(provider.into().trim().to_ascii_lowercase());
self
}
pub fn instruction(mut self, instruction: impl Into<String>) -> Self {
self.system_prompt = Some(instruction.into());
self
}
pub fn tools<I>(mut self, tools: I) -> Self
where
I: IntoIterator<Item = Arc<dyn Tool>>,
{
self.tools = tools.into_iter().collect();
self
}
pub fn tool<T: Tool + 'static>(mut self, tool: T) -> Self {
self.tools.push(Arc::new(tool));
self
}
pub fn tool_shared(mut self, tool: Arc<dyn Tool>) -> Self {
self.tools.push(tool);
self
}
pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
self.api_key = Some(api_key.into());
self
}
pub fn mcp_server(mut self, config: McpServerConfig) -> Self {
self.mcp_servers.push(config);
self
}
pub fn mcp_servers<I>(mut self, configs: I) -> Self
where
I: IntoIterator<Item = McpServerConfig>,
{
self.mcp_servers.extend(configs);
self
}
pub fn permission_checker(mut self, checker: Arc<dyn PermissionChecker>) -> Self {
self.permission_checker = Some(checker);
self
}
pub fn bypass_permissions(mut self) -> Self {
self.permission_checker = None;
self
}
pub fn hook_runner(mut self, runner: Arc<HookRunner>) -> Self {
self.inner = self.inner.hook_runner(runner);
self
}
pub fn provider(mut self, provider: Arc<dyn LLMProvider>) -> Self {
self.inner = self.inner.provider(provider);
self
}
pub fn default_tools(mut self, tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor>) -> Self {
self.inner = self.inner.default_tools(tools);
self
}
pub fn config(mut self, config: Arc<RwLock<Config>>) -> Self {
self.inner = self.inner.config(config);
self
}
pub async fn with_defaults_for_data_dir(mut self, data_dir: PathBuf) -> Result<Self, SdkError> {
let mut config = Config::from_data_dir(Some(data_dir.clone()));
if let Some(provider) = self.provider_name.clone() {
config.provider = provider;
}
if let Some(api_key) = self.api_key.clone() {
apply_api_key(&mut config, &api_key);
}
let provider = create_provider_with_dir(&config, data_dir.clone())
.await
.map_err(|e| SdkError::ProviderInit(e.to_string()))?;
let config = Arc::new(RwLock::new(config));
let builtin_tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> =
match self.permission_checker.clone() {
Some(checker) => Arc::new(
bamboo_tools::BuiltinToolExecutor::new_with_config_and_permissions(
config.clone(),
checker,
),
),
None => Arc::new(bamboo_tools::BuiltinToolExecutor::new_with_config(
config.clone(),
)),
};
let default_tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> =
if self.mcp_servers.is_empty() {
builtin_tools
} else {
let mcp_manager = Arc::new(McpServerManager::new_with_config(config.clone()));
for server_config in &self.mcp_servers {
let server_id = server_config.id.clone();
mcp_manager
.start_server(server_config.clone())
.await
.map_err(|source| SdkError::McpServerStart { server_id, source })?;
}
let mcp_tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(
McpToolExecutor::new(mcp_manager.clone(), mcp_manager.tool_index()),
);
Arc::new(CompositeToolExecutor::new(builtin_tools, mcp_tools))
};
let store = Arc::new(
SessionStoreV2::new(data_dir.clone())
.await
.map_err(|e| SdkError::StoreInit(e.to_string()))?,
);
let persistence = Arc::new(LockedSessionStore::new(store.clone()));
let skill_manager = Arc::new(SkillManager::with_config(SkillStoreConfig {
skills_dir: data_dir.join("skills"),
project_dir: std::env::current_dir().ok(),
active_mode: None,
}));
skill_manager
.initialize()
.await
.map_err(|e| SdkError::SkillInit(e.to_string()))?;
let metrics_storage: Arc<dyn bamboo_metrics::storage::MetricsStorage> =
Arc::new(SqliteMetricsStorage::new(data_dir.join("metrics.db")));
let metrics_collector =
MetricsCollector::spawn(metrics_storage, DEFAULT_METRICS_RETENTION_DAYS);
self.session_store = Some(store.clone());
self.inner = self
.inner
.storage(store.clone())
.persistence(persistence)
.attachment_reader(store)
.skill_manager(skill_manager)
.metrics_collector(metrics_collector)
.config(config)
.provider(provider)
.default_tools(default_tools);
Ok(self)
}
pub fn build(mut self) -> Result<Agent, SdkError> {
if !self.tools.is_empty() {
let registry = ToolRegistry::new();
for tool in &self.tools {
let _ = registry.register_shared(tool.clone());
}
let executor: Arc<dyn ToolExecutor> =
Arc::new(bamboo_tools::BuiltinToolExecutor::with_registry(registry));
self.inner = self.inner.default_tools(executor);
}
let runtime = self
.inner
.build()
.map_err(|e| SdkError::Build(e.to_string()))?;
Ok(Agent::from_runtime_with_config(
runtime,
self.system_prompt,
self.model,
self.session_store,
self.permission_checker,
))
}
}
impl Default for AgentBuilder {
fn default() -> Self {
Self::new()
}
}
fn apply_api_key(config: &mut Config, api_key: &str) {
let stanza = || serde_json::json!({ "api_key": api_key });
let provider = config.provider.clone();
let providers = config.providers_mut();
let applied = match provider.as_str() {
"openai" => match providers.openai.as_mut() {
Some(c) => {
c.api_key = api_key.to_string();
true
}
None => {
providers.openai = serde_json::from_value(stanza()).ok();
providers.openai.is_some()
}
},
"anthropic" => match providers.anthropic.as_mut() {
Some(c) => {
c.api_key = api_key.to_string();
true
}
None => {
providers.anthropic = serde_json::from_value(stanza()).ok();
providers.anthropic.is_some()
}
},
"gemini" => match providers.gemini.as_mut() {
Some(c) => {
c.api_key = api_key.to_string();
true
}
None => {
providers.gemini = serde_json::from_value(stanza()).ok();
providers.gemini.is_some()
}
},
_ => false,
};
if !applied {
tracing::warn!(
provider = %config.provider,
"AgentBuilder::api_key: key not applied — the active provider either \
takes no plain api_key (e.g. copilot uses cached OAuth) or its config \
could not be built from a key alone"
);
}
}
#[cfg(test)]
mod tests {
use super::apply_api_key;
use bamboo_llm::Config;
#[test]
fn api_key_fabricates_stanza_for_keyed_providers() {
for provider in ["openai", "anthropic", "gemini"] {
let mut config = Config::default();
config.provider = provider.to_string();
config.providers_mut().openai = None;
config.providers_mut().anthropic = None;
config.providers_mut().gemini = None;
apply_api_key(&mut config, "sk-test-123");
let key = match provider {
"openai" => config
.providers()
.openai
.as_ref()
.map(|c| c.api_key.as_str()),
"anthropic" => config
.providers()
.anthropic
.as_ref()
.map(|c| c.api_key.as_str()),
"gemini" => config
.providers()
.gemini
.as_ref()
.map(|c| c.api_key.as_str()),
_ => unreachable!(),
};
assert_eq!(
key,
Some("sk-test-123"),
"expected a fabricated {provider} stanza carrying the api_key"
);
}
}
#[test]
fn provider_selection_before_api_key_routes_key_to_chosen_provider() {
let mut config = Config::default();
config.provider = "anthropic".to_string(); config.providers_mut().openai = None;
config.providers_mut().anthropic = None;
config.provider = "openai".to_string();
apply_api_key(&mut config, "sk-openai-xyz");
assert_eq!(
config
.providers()
.openai
.as_ref()
.map(|c| c.api_key.as_str()),
Some("sk-openai-xyz"),
"api_key must land on the selected provider (openai), not the default"
);
assert!(
config.providers().anthropic.is_none(),
"the default provider must not receive the key"
);
}
}