Files
optimclaw/src/app.rs
T
4c043bf057 feat: complete multi-tenant isolation — phases 2–4 (#1614)
* feat: complete multi-tenant isolation — per-user budgets, model selection, heartbeat cycling

Finishes the remaining isolation work from phases 2–4 of #59:

Phase 2 (DB scoping): Fix /status and /list commands to use _for_user
DB variants instead of global queries that leaked cross-user job data.

Phase 3 (Runtime isolation): Per-user workspace in routine engine's
spawn_fire so lightweight routines run in the correct user context.
Per-user daily cost tracking in CostGuard with configurable budget via
MAX_COST_PER_USER_PER_DAY_CENTS. Multi-user heartbeat that cycles
through all users with routines, auto-detected from GATEWAY_USER_TOKENS.

Phase 4 (Provider/tools): Per-user model selection via preferred_model
setting — looked up from SettingsStore on first iteration, threaded
through ReasoningContext.model_override to CompletionRequest. Works
with providers that support per-request model overrides (NearAI).

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: use selected_model setting key to match /model command persistence

The dispatcher was reading "preferred_model" but the /model command
(merged from staging) persists to "selected_model". Since set_setting
is already per-user scoped, using the same key makes /model work as
the per-user model override in multi-tenant mode.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: heartbeat hygiene, /model multi-tenant guard, RigAdapter model override

Three follow-up fixes for multi-tenant isolation:

1. Multi-user heartbeat now runs memory hygiene per user before each
   heartbeat check, matching single-user heartbeat behavior.

2. /model command in multi-tenant mode only persists to per-user
   settings (selected_model) without calling set_model() on the shared
   LlmProvider. The per-request model_override in the dispatcher reads
   from the same setting. Added multi_tenant flag to AgentConfig
   (auto-detected from GATEWAY_USER_TOKENS).

3. RigAdapter now supports per-request model overrides by injecting the
   model name into rig-core's additional_params. OpenAI/Anthropic/Ollama
   API servers use last-key-wins for duplicate JSON keys, so the override
   takes effect via serde's flatten serialization order.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address PR review — cost model attribution, heartbeat concurrency, pruning

Fixes from review comments on #1614:

- Cost tracking now uses the override model name (not active_model_name)
  when a per-user model override is active, for accurate attribution.
- Multi-user heartbeat runs per-user checks concurrently via JoinSet
  instead of sequentially, preventing one slow user from blocking others.
- Per-user failure counts tracked independently; users exceeding
  max_failures are skipped (matching single-user semantics).
- per_user_daily_cost HashMap pruned on day rollover to prevent
  unbounded growth in long-lived deployments.
- Doc comment fixed: says "routines" not "active routines".

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: /status ownership, model persistence scoping, heartbeat robustness

Addresses second round of PR review on #1614:

- /status <job_id> DB path now validates job.user_id == requesting user
  before returning data (was missing ownership check, security fix).

- persist_selected_model takes user_id param instead of owner_id, and
  skips .env/TOML writes in multi-tenant mode (these are shared global
  files). handle_system_command now receives user_id from caller.

- JoinSet collection handles Err(JoinError) explicitly instead of
  silently dropping panicked tasks.

- Notification forwarder extracts owner_id from response metadata in
  multi-tenant mode for per-user routing instead of broadcasting to
  the agent owner.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: cost pricing, fire_manual workspace, heartbeat concurrency cap

Round 3 review fixes:

- Cost tracking passes None for cost_per_token when model override is
  active, letting CostGuard look up pricing by model name instead of
  using the default provider's rates (serrrfirat).

- fire_manual() now uses per-user workspace, matching spawn_fire()
  pattern (serrrfirat).

- Removed MULTI_TENANT env var — multi-tenant mode is auto-detected
  solely from GATEWAY_USER_TOKENS presence (serrrfirat + Copilot).

- Multi-user heartbeat capped at 8 concurrent tasks to avoid flooding
  the LLM provider (serrrfirat + Copilot).

- Fixed inject_model_override doc comment accuracy (Copilot).

- Added comment explaining multi-tenant notification routing priority
  (Copilot).

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* feat: user-scoped webhook endpoint for multi-tenant isolation

Adds POST /api/webhooks/u/{user_id}/{path} — a user-scoped webhook
endpoint that filters the routine lookup by user_id, preventing
cross-user webhook triggering when paths collide.

The existing /api/webhooks/{path} endpoint remains unchanged for
backward compatibility in single-user deployments.

Changes:
- get_webhook_routine_by_path gains user_id: Option<&str> param
- Both postgres and libsql implementations add AND user_id = ? filter
  when user_id is provided
- New webhook_trigger_user_scoped_handler extracts (user_id, path)
  from URL and passes to shared fire_webhook_inner logic
- Route registered on public router (webhooks are called by external
  services that can't send bearer tokens)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* feat: add TenantCtx for compile-time tenant isolation

Implements zmanian's architectural proposal from #1614 review:
two-tier scoped database access (TenantScope/AdminScope) so handler
code cannot accidentally bypass tenant scoping.

TenantScope (default): wraps user_id + Arc<dyn Database>, auto-binds
user_id on every operation. ID-based lookups return None for cross-
tenant resources. No escape hatch — forgetting to scope is a compile
error.

AdminScope (explicit opt-in): cross-tenant access for system-level
components (heartbeat, routine engine, self-repair, scheduler, worker).

TenantCtx bundles TenantScope + workspace + cost guard + per-user
rate limiting. Constructed once per request in handle_message, threaded
through all command handlers and ChatDelegate.

Key changes:
- New src/tenant.rs (~920 lines): TenantScope, AdminScope, TenantCtx,
  TenantRateState, TenantRateRegistry
- All command handlers: user_id: &str → ctx: &TenantCtx
- ChatDelegate: cost check/record/settings via self.tenant
- System components: store field changed to AdminScope
- Config: TENANT_MAX_LLM_CONCURRENT, TENANT_MAX_JOBS_CONCURRENT env vars
- Fixes bug: /status <job_id> cross-tenant leak (now auto-filtered)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-25 17:24:48 -07:00

987 lines
39 KiB
Rust

//! Application builder for initializing core IronClaw components.
//!
//! Extracts the mechanical initialization phases from `main.rs` into a
//! reusable builder so that:
//!
//! - Tests can construct a full `AppComponents` without wiring channels
//! - Main stays focused on CLI dispatch and channel setup
//! - Each init phase is independently testable
use std::sync::Arc;
use crate::agent::SessionManager as AgentSessionManager;
use crate::channels::web::log_layer::LogBroadcaster;
use crate::config::Config;
use crate::context::ContextManager;
use crate::db::Database;
use crate::extensions::ExtensionManager;
use crate::hooks::HookRegistry;
use crate::llm::{LlmProvider, RecordingLlm, SessionManager};
use crate::safety::SafetyLayer;
use crate::secrets::SecretsStore;
use crate::skills::SkillRegistry;
use crate::skills::catalog::SkillCatalog;
use crate::tools::ToolRegistry;
use crate::tools::mcp::{McpProcessManager, McpSessionManager};
use crate::tools::wasm::SharedCredentialRegistry;
use crate::tools::wasm::WasmToolRuntime;
use crate::workspace::{EmbeddingCacheConfig, EmbeddingProvider, Workspace};
/// Fully initialized application components, ready for channel wiring
/// and agent construction.
pub struct AppComponents {
/// The (potentially mutated) config after DB reload and secret injection.
pub config: Config,
pub db: Option<Arc<dyn Database>>,
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
pub llm: Arc<dyn LlmProvider>,
pub cheap_llm: Option<Arc<dyn LlmProvider>>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub embeddings: Option<Arc<dyn EmbeddingProvider>>,
pub workspace: Option<Arc<Workspace>>,
pub extension_manager: Option<Arc<ExtensionManager>>,
pub mcp_session_manager: Arc<McpSessionManager>,
pub mcp_process_manager: Arc<McpProcessManager>,
pub wasm_tool_runtime: Option<Arc<WasmToolRuntime>>,
pub log_broadcaster: Arc<LogBroadcaster>,
pub context_manager: Arc<ContextManager>,
pub hooks: Arc<HookRegistry>,
/// Shared thread/session manager used by the standard agent runtime.
pub agent_session_manager: Arc<AgentSessionManager>,
pub skill_registry: Option<Arc<std::sync::RwLock<SkillRegistry>>>,
pub skill_catalog: Option<Arc<SkillCatalog>>,
pub cost_guard: Arc<crate::agent::cost_guard::CostGuard>,
pub recording_handle: Option<Arc<RecordingLlm>>,
pub session: Arc<SessionManager>,
pub catalog_entries: Vec<crate::extensions::RegistryEntry>,
pub dev_loaded_tool_names: Vec<String>,
pub builder: Option<Arc<dyn crate::tools::SoftwareBuilder>>,
}
/// Options that control optional init phases.
#[derive(Default)]
pub struct AppBuilderFlags {
pub no_db: bool,
}
/// Builder that orchestrates the 5 mechanical init phases.
pub struct AppBuilder {
config: Config,
flags: AppBuilderFlags,
toml_path: Option<std::path::PathBuf>,
session: Arc<SessionManager>,
log_broadcaster: Arc<LogBroadcaster>,
// Accumulated state
db: Option<Arc<dyn Database>>,
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
// Test overrides
llm_override: Option<Arc<dyn LlmProvider>>,
// Backend-specific handles needed by secrets store
handles: Option<crate::db::DatabaseHandles>,
}
impl AppBuilder {
/// Create a new builder.
///
/// The `session` and `log_broadcaster` are created before the builder
/// because tracing must be initialized before any init phase runs,
/// and the log broadcaster is part of the tracing layer.
pub fn new(
config: Config,
flags: AppBuilderFlags,
toml_path: Option<std::path::PathBuf>,
session: Arc<SessionManager>,
log_broadcaster: Arc<LogBroadcaster>,
) -> Self {
Self {
config,
flags,
toml_path,
session,
log_broadcaster,
db: None,
secrets_store: None,
llm_override: None,
handles: None,
}
}
/// Inject a pre-created database, skipping `init_database()`.
pub fn with_database(&mut self, db: Arc<dyn Database>) {
self.db = Some(db);
}
/// Inject a pre-created LLM provider, skipping `init_llm()`.
pub fn with_llm(&mut self, llm: Arc<dyn LlmProvider>) {
self.llm_override = Some(llm);
}
/// Phase 1: Initialize database backend.
///
/// Creates the database connection, runs migrations, reloads config
/// from DB, attaches DB to session manager, and cleans up stale jobs.
pub async fn init_database(&mut self) -> Result<(), anyhow::Error> {
if self.db.is_some() {
tracing::debug!("Database already provided, skipping init_database()");
return Ok(());
}
if self.flags.no_db {
tracing::warn!("Running without database connection");
return Ok(());
}
let (db, handles) = crate::db::connect_with_handles(&self.config.database)
.await
.map_err(|e| anyhow::anyhow!("{}", e))?;
self.handles = Some(handles);
// Post-init: migrate disk config, reload config from DB, attach session, cleanup
if let Err(e) =
crate::bootstrap::migrate_disk_to_db(db.as_ref(), &self.config.owner_id).await
{
tracing::warn!("Disk-to-DB settings migration failed: {}", e);
}
let toml_path = self.toml_path.as_deref();
match Config::from_db_with_toml(db.as_ref(), &self.config.owner_id, toml_path).await {
Ok(db_config) => {
self.config = db_config;
tracing::debug!("Configuration reloaded from database");
}
Err(e) => {
tracing::warn!(
"Failed to reload config from DB, keeping env-based config: {}",
e
);
}
}
self.session
.attach_store(db.clone(), &self.config.owner_id)
.await;
// Fire-and-forget housekeeping — no need to block startup.
let db_cleanup = db.clone();
tokio::spawn(async move {
if let Err(e) = db_cleanup.cleanup_stale_sandbox_jobs().await {
tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e);
}
});
self.db = Some(db);
Ok(())
}
/// Phase 2: Create secrets store.
///
/// Requires a master key and a backend-specific DB handle. After creating
/// the store, injects any encrypted LLM API keys into the config overlay
/// and re-resolves config.
pub async fn init_secrets(&mut self) -> Result<(), anyhow::Error> {
let master_key = match self.config.secrets.master_key() {
Some(k) => k,
None => {
// No secrets DB available, but we can still load tokens from
// OS credential stores (e.g., Anthropic OAuth via Claude Code's
// macOS Keychain / Linux ~/.claude/.credentials.json).
crate::config::inject_os_credentials();
// Consume unused handles
self.handles.take();
// Re-resolve only the LLM config with OS credentials.
let store: Option<&(dyn crate::db::SettingsStore + Sync)> =
self.db.as_ref().map(|db| db.as_ref() as _);
let toml_path = self.toml_path.as_deref();
let owner_id = self.config.owner_id.clone();
if let Err(e) = self
.config
.re_resolve_llm(store, &owner_id, toml_path)
.await
{
tracing::warn!(
"Failed to re-resolve LLM config after OS credential injection: {e}"
);
}
return Ok(());
}
};
let crypto = match crate::secrets::SecretsCrypto::new(master_key.clone()) {
Ok(c) => Arc::new(c),
Err(e) => {
tracing::warn!("Failed to initialize secrets crypto: {}", e);
self.handles.take();
return Ok(());
}
};
// Fallback covers the no-database path where `init_database` returned
// early before populating `self.handles`.
let empty_handles = crate::db::DatabaseHandles::default();
let handles = self.handles.as_ref().unwrap_or(&empty_handles);
let store = crate::secrets::create_secrets_store(crypto, handles);
if let Some(ref secrets) = store {
// Inject LLM API keys from encrypted storage
crate::config::inject_llm_keys_from_secrets(secrets.as_ref(), &self.config.owner_id)
.await;
// Re-resolve only the LLM config with newly available keys.
let store: Option<&(dyn crate::db::SettingsStore + Sync)> =
self.db.as_ref().map(|db| db.as_ref() as _);
let toml_path = self.toml_path.as_deref();
let owner_id = self.config.owner_id.clone();
if let Err(e) = self
.config
.re_resolve_llm(store, &owner_id, toml_path)
.await
{
tracing::warn!("Failed to re-resolve LLM config after secret injection: {e}");
}
}
self.secrets_store = store;
Ok(())
}
/// Phase 3: Initialize LLM provider chain.
///
/// Delegates to `build_provider_chain` which applies all decorators
/// (retry, smart routing, failover, circuit breaker, response cache).
#[allow(clippy::type_complexity)]
pub async fn init_llm(
&self,
) -> Result<
(
Arc<dyn LlmProvider>,
Option<Arc<dyn LlmProvider>>,
Option<Arc<RecordingLlm>>,
),
anyhow::Error,
> {
let (llm, cheap_llm, recording_handle) =
crate::llm::build_provider_chain(&self.config.llm, self.session.clone()).await?;
Ok((llm, cheap_llm, recording_handle))
}
/// Phase 4: Initialize safety, tools, embeddings, and workspace.
pub async fn init_tools(
&self,
llm: &Arc<dyn LlmProvider>,
) -> Result<
(
Arc<SafetyLayer>,
Arc<ToolRegistry>,
Option<Arc<dyn EmbeddingProvider>>,
Option<Arc<Workspace>>,
Option<Arc<dyn crate::tools::SoftwareBuilder>>,
),
anyhow::Error,
> {
let safety = Arc::new(SafetyLayer::new(&self.config.safety));
tracing::debug!("Safety layer initialized");
// Initialize tool registry with credential injection support
let credential_registry = Arc::new(SharedCredentialRegistry::new());
let tools = if let Some(ref ss) = self.secrets_store {
Arc::new(
ToolRegistry::new()
.with_credentials(Arc::clone(&credential_registry), Arc::clone(ss)),
)
} else {
Arc::new(ToolRegistry::new())
};
tools.register_builtin_tools();
tools.register_tool_info();
if let Some(ref ss) = self.secrets_store {
tools.register_secrets_tools(Arc::clone(ss));
}
// Create embeddings provider using the unified method
let embeddings = self
.config
.embeddings
.create_provider(&self.config.llm.nearai.base_url, self.session.clone());
// Register memory tools if database is available
let workspace_user_id = self.config.owner_id.as_str();
let workspace = if let Some(ref db) = self.db {
let emb_cache_config = EmbeddingCacheConfig {
max_entries: self.config.embeddings.cache_size,
};
let mut ws = Workspace::new_with_db(workspace_user_id, db.clone())
.with_search_config(&self.config.search);
if let Some(ref emb) = embeddings {
ws = ws.with_embeddings_cached(emb.clone(), emb_cache_config.clone());
}
// Wire workspace-level settings (read scopes, memory layers)
if !self.config.workspace.read_scopes.is_empty() {
ws = ws.with_additional_read_scopes(self.config.workspace.read_scopes.clone());
tracing::info!(
user_id = workspace_user_id,
read_scopes = ?ws.read_user_ids(),
"Workspace configured with multi-scope reads"
);
}
ws = ws.with_memory_layers(self.config.workspace.memory_layers.clone());
let ws = Arc::new(ws);
// Detect multi-tenant mode: when GATEWAY_USER_TOKENS is configured,
// each authenticated user needs their own workspace scope. Use
// WorkspacePool (which implements WorkspaceResolver) to create
// per-user workspaces on demand instead of sharing the startup
// workspace across all users.
let is_multi_tenant = self
.config
.channels
.gateway
.as_ref()
.is_some_and(|gw| gw.user_tokens.is_some());
if is_multi_tenant {
let pool = Arc::new(crate::channels::web::server::WorkspacePool::new(
Arc::clone(db),
embeddings.clone(),
emb_cache_config,
self.config.search.clone(),
self.config.workspace.clone(),
));
tools.register_memory_tools_with_resolver(pool);
tracing::info!(
"Memory tools configured with per-user workspace resolver (multi-tenant mode)"
);
} else {
tools.register_memory_tools(Arc::clone(&ws));
}
Some(ws)
} else {
None
};
// Register image/vision tools if we have a workspace and LLM API credentials
if workspace.is_some() {
let (api_base, api_key_opt) = if let Some(ref provider) = self.config.llm.provider {
(
provider.base_url.clone(),
provider.api_key.as_ref().map(|s| {
use secrecy::ExposeSecret;
s.expose_secret().to_string()
}),
)
} else {
(
self.config.llm.nearai.base_url.clone(),
self.config.llm.nearai.api_key.as_ref().map(|s| {
use secrecy::ExposeSecret;
s.expose_secret().to_string()
}),
)
};
if let Some(api_key) = api_key_opt {
// Check for image generation models
let model_name = self
.config
.llm
.provider
.as_ref()
.map(|p| p.model.clone())
.unwrap_or_else(|| self.config.llm.nearai.model.clone());
let models = vec![model_name.clone()];
let gen_model = crate::llm::image_models::suggest_image_model(&models)
.unwrap_or("flux-1.1-pro")
.to_string();
tools.register_image_tools(api_base.clone(), api_key.clone(), gen_model, None);
// Check for vision models
let vision_model = crate::llm::vision_models::suggest_vision_model(&models)
.unwrap_or(&model_name)
.to_string();
tools.register_vision_tools(api_base, api_key, vision_model, None);
}
}
// Register builder tool if enabled
let builder = if self.config.builder.enabled
&& (self.config.agent.allow_local_tools || !self.config.sandbox.enabled)
{
let b = tools
.register_builder_tool(llm.clone(), Some(self.config.builder.to_builder_config()))
.await;
tracing::debug!("Builder mode enabled");
Some(b)
} else {
None
};
Ok((safety, tools, embeddings, workspace, builder))
}
/// Phase 5: Load WASM tools, MCP servers, and create extension manager.
pub async fn init_extensions(
&self,
tools: &Arc<ToolRegistry>,
hooks: &Arc<HookRegistry>,
) -> Result<
(
Arc<McpSessionManager>,
Arc<McpProcessManager>,
Option<Arc<WasmToolRuntime>>,
Option<Arc<ExtensionManager>>,
Vec<crate::extensions::RegistryEntry>,
Vec<String>,
),
anyhow::Error,
> {
use crate::tools::mcp::config::load_mcp_servers_from_db;
use crate::tools::wasm::{WasmToolLoader, load_dev_tools};
let mcp_session_manager = Arc::new(McpSessionManager::new());
let mcp_process_manager = Arc::new(McpProcessManager::new());
// Create WASM tool runtime eagerly so extensions installed after startup
// (e.g. via the web UI) can still be activated. The tools directory is only
// needed when loading modules, not for engine initialisation.
let wasm_tool_runtime: Option<Arc<WasmToolRuntime>> = if self.config.wasm.enabled {
WasmToolRuntime::new(self.config.wasm.to_runtime_config())
.map(Arc::new)
.map_err(|e| tracing::warn!("Failed to initialize WASM runtime: {}", e))
.ok()
} else {
None
};
// Load WASM tools and MCP servers concurrently
let wasm_tools_future = {
let wasm_tool_runtime = wasm_tool_runtime.clone();
let secrets_store = self.secrets_store.clone();
let tools = Arc::clone(tools);
let wasm_config = self.config.wasm.clone();
async move {
let mut dev_loaded_tool_names: Vec<String> = Vec::new();
if let Some(ref runtime) = wasm_tool_runtime {
let mut loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
if let Some(ref secrets) = secrets_store {
loader = loader.with_secrets_store(Arc::clone(secrets));
}
match loader.load_from_dir(&wasm_config.tools_dir).await {
Ok(results) => {
if !results.loaded.is_empty() {
tracing::debug!(
"Loaded {} WASM tools from {}",
results.loaded.len(),
wasm_config.tools_dir.display()
);
}
for (path, err) in &results.errors {
tracing::warn!(
"Failed to load WASM tool {}: {}",
path.display(),
err
);
}
}
Err(e) => {
tracing::warn!("Failed to scan WASM tools directory: {}", e);
}
}
match load_dev_tools(&loader, &wasm_config.tools_dir).await {
Ok(results) => {
dev_loaded_tool_names.extend(results.loaded.iter().cloned());
if !dev_loaded_tool_names.is_empty() {
tracing::debug!(
"Loaded {} dev WASM tools from build artifacts",
dev_loaded_tool_names.len()
);
}
}
Err(e) => {
tracing::debug!("No dev WASM tools found: {}", e);
}
}
}
dev_loaded_tool_names
}
};
let mcp_servers_future = {
let secrets_store = self.secrets_store.clone();
let db = self.db.clone();
let tools = Arc::clone(tools);
let mcp_sm = Arc::clone(&mcp_session_manager);
let pm = Arc::clone(&mcp_process_manager);
let owner_id = self.config.owner_id.clone();
async move {
let servers_result = if let Some(ref d) = db {
load_mcp_servers_from_db(d.as_ref(), &owner_id).await
} else {
crate::tools::mcp::config::load_mcp_servers().await
};
match servers_result {
Ok(servers) => {
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
if !enabled.is_empty() {
tracing::debug!(
"Loading {} configured MCP server(s)...",
enabled.len()
);
}
let mut join_set = tokio::task::JoinSet::new();
for server in enabled {
let mcp_sm = Arc::clone(&mcp_sm);
let secrets = secrets_store.clone();
let tools = Arc::clone(&tools);
let pm = Arc::clone(&pm);
let owner_id = owner_id.clone();
join_set.spawn(async move {
let server_name = server.name.clone();
let client = match crate::tools::mcp::create_client_from_config(
server,
&mcp_sm,
&pm,
secrets,
&owner_id,
)
.await
{
Ok(c) => c,
Err(e) => {
tracing::warn!(
"Failed to create MCP client for '{}': {}",
server_name,
e
);
return None;
}
};
match client.list_tools().await {
Ok(mcp_tools) => {
let tool_count = mcp_tools.len();
match client.create_tools().await {
Ok(tool_impls) => {
for tool in tool_impls {
tools.register(tool).await;
}
tracing::debug!(
"Loaded {} tools from MCP server '{}'",
tool_count,
server_name
);
return Some((
server_name,
Arc::new(client),
));
}
Err(e) => {
tracing::warn!(
"Failed to create tools from MCP server '{}': {}",
server_name,
e
);
}
}
}
Err(e) => {
let err_str = e.to_string();
if err_str.contains("401")
|| err_str.contains("authentication")
{
tracing::warn!(
"MCP server '{}' requires authentication. \
Run: ironclaw mcp auth {}",
server_name,
server_name
);
} else {
tracing::warn!(
"Failed to connect to MCP server '{}': {}",
server_name,
e
);
}
}
}
None
});
}
let mut startup_clients = Vec::new();
while let Some(result) = join_set.join_next().await {
match result {
Ok(Some(client_pair)) => {
startup_clients.push(client_pair);
}
Ok(None) => {}
Err(e) => {
if e.is_panic() {
tracing::error!("MCP server loading task panicked: {}", e);
} else {
tracing::warn!("MCP server loading task failed: {}", e);
}
}
}
}
return startup_clients;
}
Err(e) => {
if matches!(
e,
crate::tools::mcp::config::ConfigError::InvalidConfig { .. }
| crate::tools::mcp::config::ConfigError::Json(_)
) {
tracing::warn!(
"MCP server configuration is invalid: {}. \
Fix or remove the corrupted config.",
e
);
} else {
tracing::debug!("No MCP servers configured ({})", e);
}
}
}
Vec::new()
}
};
let (dev_loaded_tool_names, startup_mcp_clients) =
tokio::join!(wasm_tools_future, mcp_servers_future);
// Load registry catalog entries for extension discovery
let mut catalog_entries = match crate::registry::RegistryCatalog::load_or_embedded() {
Ok(catalog) => {
let entries: Vec<_> = catalog
.all()
.iter()
.filter_map(|m| m.to_registry_entry())
.collect();
tracing::debug!(
count = entries.len(),
"Loaded registry catalog entries for extension discovery"
);
entries
}
Err(e) => {
tracing::warn!("Failed to load registry catalog: {}", e);
Vec::new()
}
};
// Append builtin entries (e.g. channel-relay integrations) so they appear
// in the web UI's available extensions list.
let builtin = crate::extensions::registry::builtin_entries();
for entry in builtin {
if !catalog_entries.iter().any(|e| e.name == entry.name) {
catalog_entries.push(entry);
}
}
// Create extension manager. Use ephemeral in-memory secrets if no
// persistent store is configured (listing/install/activate still work).
let ext_secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> = if let Some(ref s) =
self.secrets_store
{
Arc::clone(s)
} else {
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
let ephemeral_key =
secrecy::SecretString::from(crate::secrets::keychain::generate_master_key_hex());
let crypto = Arc::new(SecretsCrypto::new(ephemeral_key).expect("ephemeral crypto"));
tracing::debug!("Using ephemeral in-memory secrets store for extension manager");
Arc::new(InMemorySecretsStore::new(crypto))
};
let extension_manager = {
let manager = Arc::new(ExtensionManager::new(
Arc::clone(&mcp_session_manager),
Arc::clone(&mcp_process_manager),
ext_secrets,
Arc::clone(tools),
Some(Arc::clone(hooks)),
wasm_tool_runtime.clone(),
self.config.wasm.tools_dir.clone(),
self.config.channels.wasm_channels_dir.clone(),
self.config.tunnel.public_url.clone(),
self.config.owner_id.clone(),
self.db.clone(),
catalog_entries.clone(),
));
tools.register_extension_tools(Arc::clone(&manager));
tracing::debug!("Extension manager initialized with in-chat discovery tools");
if !startup_mcp_clients.is_empty() {
tracing::info!(
count = startup_mcp_clients.len(),
"Injecting startup MCP clients into extension manager"
);
for (name, client) in startup_mcp_clients {
manager.inject_mcp_client(name, client).await;
}
}
Some(manager)
};
// register_builder_tool() already calls register_dev_tools() internally,
// so only register them here when the builder didn't already do it.
let builder_registered_dev_tools = self.config.builder.enabled
&& (self.config.agent.allow_local_tools || !self.config.sandbox.enabled);
if self.config.agent.allow_local_tools && !builder_registered_dev_tools {
tools.register_dev_tools();
}
Ok((
mcp_session_manager,
mcp_process_manager,
wasm_tool_runtime,
extension_manager,
catalog_entries,
dev_loaded_tool_names,
))
}
/// Run all init phases in order and return the assembled components.
pub async fn build_all(mut self) -> Result<AppComponents, anyhow::Error> {
self.init_database().await?;
self.init_secrets().await?;
// Post-init validation: backends with dedicated config (nearai, gemini_oauth,
// bedrock, openai_codex) handle their own credential resolution. For registry-based
// backends, fail early if no provider config was resolved.
if !matches!(
self.config.llm.backend.as_str(),
"nearai" | "gemini_oauth" | "bedrock" | "openai_codex"
) && self.config.llm.provider.is_none()
{
let backend = &self.config.llm.backend;
anyhow::bail!(
"LLM_BACKEND={backend} is configured but no credentials were found. \
Set the appropriate API key environment variable or run the setup wizard."
);
}
let (llm, cheap_llm, recording_handle) = if let Some(llm) = self.llm_override.take() {
(llm, None, None)
} else {
self.init_llm().await?
};
let (safety, tools, embeddings, workspace, builder) = self.init_tools(&llm).await?;
// Create hook registry early so runtime extension activation can register hooks.
let hooks = Arc::new(HookRegistry::new());
let agent_session_manager =
Arc::new(AgentSessionManager::new().with_hooks(Arc::clone(&hooks)));
let (
mcp_session_manager,
mcp_process_manager,
wasm_tool_runtime,
extension_manager,
catalog_entries,
dev_loaded_tool_names,
) = self.init_extensions(&tools, &hooks).await?;
// Load bootstrap-completed flag from settings so that existing users
// who already completed onboarding don't re-get bootstrap injection.
if let Some(ref ws) = workspace {
let toml_path = crate::settings::Settings::default_toml_path();
if let Ok(Some(settings)) = crate::settings::Settings::load_toml(&toml_path)
&& settings.profile_onboarding_completed
{
ws.mark_bootstrap_completed();
}
}
// Seed workspace and backfill embeddings
if let Some(ref ws) = workspace {
// Import workspace files from disk FIRST if WORKSPACE_IMPORT_DIR is set.
// This lets Docker images / deployment scripts ship customized
// workspace templates (e.g., AGENTS.md, TOOLS.md) that override
// the generic seeds. Only imports files that don't already exist
// in the database — never overwrites user edits.
//
// Runs before seed_if_empty() so that custom templates take priority
// over generic seeds. seed_if_empty() then fills any remaining gaps.
if let Ok(import_dir) = std::env::var("WORKSPACE_IMPORT_DIR") {
let import_path = std::path::Path::new(&import_dir);
match ws.import_from_directory(import_path).await {
Ok(count) if count > 0 => {
tracing::debug!("Imported {} workspace file(s) from {}", count, import_dir);
}
Ok(_) => {}
Err(e) => {
tracing::warn!(
"Failed to import workspace files from {}: {}",
import_dir,
e
);
}
}
}
match ws.seed_if_empty().await {
Ok(_) => {}
Err(e) => {
tracing::warn!("Failed to seed workspace: {}", e);
}
}
if embeddings.is_some() {
let ws_bg = Arc::clone(ws);
tokio::spawn(async move {
match ws_bg.backfill_embeddings().await {
Ok(count) if count > 0 => {
tracing::debug!("Backfilled embeddings for {} chunks", count);
}
Ok(_) => {}
Err(e) => {
tracing::warn!("Failed to backfill embeddings: {}", e);
}
}
});
}
}
// Skills system
let (skill_registry, skill_catalog) = if self.config.skills.enabled {
let mut registry = SkillRegistry::new(self.config.skills.local_dir.clone())
.with_installed_dir(self.config.skills.installed_dir.clone());
let loaded = registry.discover_all().await;
if !loaded.is_empty() {
tracing::debug!("Loaded {} skill(s): {}", loaded.len(), loaded.join(", "));
}
let registry = Arc::new(std::sync::RwLock::new(registry));
let catalog = crate::skills::catalog::shared_catalog();
tools.register_skill_tools(Arc::clone(&registry), Arc::clone(&catalog));
(Some(registry), Some(catalog))
} else {
(None, None)
};
let context_manager = Arc::new(ContextManager::new(self.config.agent.max_parallel_jobs));
let cost_guard = Arc::new(crate::agent::cost_guard::CostGuard::new(
crate::agent::cost_guard::CostGuardConfig {
max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents,
max_actions_per_hour: self.config.agent.max_actions_per_hour,
max_cost_per_user_per_day_cents: self.config.agent.max_cost_per_user_per_day_cents,
},
));
tracing::debug!(
"Tool registry initialized with {} total tools",
tools.count()
);
Ok(AppComponents {
config: self.config,
db: self.db,
secrets_store: self.secrets_store,
llm,
cheap_llm,
safety,
tools,
embeddings,
workspace,
extension_manager,
mcp_session_manager,
mcp_process_manager,
wasm_tool_runtime,
log_broadcaster: self.log_broadcaster,
context_manager,
hooks,
agent_session_manager,
skill_registry,
skill_catalog,
cost_guard,
recording_handle,
session: self.session,
catalog_entries,
dev_loaded_tool_names,
builder,
})
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use async_trait::async_trait;
use tokio::sync::mpsc;
use crate::agent::SessionManager as AgentSessionManager;
use crate::hooks::{
Hook, HookContext, HookError, HookEvent, HookOutcome, HookPoint, HookRegistry,
};
struct SessionStartHook {
tx: mpsc::UnboundedSender<(String, String)>,
}
#[async_trait]
impl Hook for SessionStartHook {
fn name(&self) -> &str {
"session-start-test"
}
fn hook_points(&self) -> &[HookPoint] {
&[HookPoint::OnSessionStart]
}
async fn execute(
&self,
event: &HookEvent,
_ctx: &HookContext,
) -> Result<HookOutcome, HookError> {
if let HookEvent::SessionStart {
user_id,
session_id,
} = event
{
self.tx
.send((user_id.clone(), session_id.clone()))
.expect("test channel receiver should be alive");
} else {
panic!("SessionStartHook received an unexpected event: {event:?}");
}
Ok(HookOutcome::ok())
}
}
#[tokio::test]
async fn agent_session_manager_runs_session_start_hooks() {
let hooks = Arc::new(HookRegistry::new());
let (tx, mut rx) = mpsc::unbounded_channel();
hooks.register(Arc::new(SessionStartHook { tx })).await;
let manager = AgentSessionManager::new().with_hooks(Arc::clone(&hooks));
manager.get_or_create_session("user-123").await;
let (user_id, session_id) =
tokio::time::timeout(std::time::Duration::from_secs(1), rx.recv())
.await
.expect("session start hook should fire")
.expect("session start payload should be present");
assert_eq!(user_id, "user-123");
assert!(!session_id.is_empty());
}
}