From c3ce26278a659161ce1fbce36c5db1cda1928ada Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Sat, 21 Feb 2026 18:54:31 -0800 Subject: [PATCH] refactor: simplify config resolution and consolidate main.rs init (#287) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor: simplify config resolution and consolidate main.rs init into AppBuilder - Add parse_bool_env() and parse_string_env() helpers to eliminate repetitive 5-line optional_env/parse/map_err/unwrap_or boilerplate across 12 config files - Add EmbeddingsConfig::create_provider() to centralize embeddings construction (fixes hardcoded 1536 dimensions and missing Ollama provider in app.rs) - Extract init_cli_tracing(), setup_wasm_channels(), start_tunnel(), run_memory_command(), run_worker(), run_claude_bridge() from main.rs - Replace ~600 lines of inline init in main.rs with AppBuilder::build_all() - Expose catalog_entries from AppComponents for gateway registry entries - Net reduction: ~738 lines across 15 files Co-Authored-By: Claude Opus 4.6 * fix: propagate dev_loaded_tool_names from AppBuilder and add parse_option_env helper Address PR review feedback: - Capture dev_loaded_tool_names from WASM loading in init_extensions() and expose via AppComponents so bootstrap_hooks receives the actual dev tool names instead of an empty slice (fixes silent hook skip) - Add parse_option_env() helper for Option config fields, simplifying max_cost_per_day_cents and max_actions_per_hour in agent.rs Co-Authored-By: Claude Opus 4.6 * fix: fetch real NEAR AI pricing and unify cost calculation path CostGuard was independently looking up pricing via costs::model_cost(), falling back to GPT-4o default rates when NEAR AI model names didn't match the static table — causing ~3x cost overestimates in logs. - Add pricing map to NearAiChatProvider that fetches real rates from /v1/model/list at startup (background, non-blocking) - Update cost_per_token() to check fetched pricing first, then static table, then default - Add cost_per_token parameter to CostGuard::record_llm_call() so the dispatcher passes provider-sourced rates directly Co-Authored-By: Claude Opus 4.6 * chore: update default NEAR AI model to GLM-latest Replace fireworks llama4-maverick-instruct-basic with zai-org/GLM-latest as the default model in config and setup wizard. Co-Authored-By: Claude Opus 4.6 * fix: align wizard default model name with config Change "zai/GLM-latest" to "zai-org/GLM-latest" in wizard.rs to match the default in config/llm.rs. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- Cargo.lock | 18 +- src/agent/cost_guard.rs | 29 +- src/agent/dispatcher.rs | 1 + src/app.rs | 108 +-- src/config/agent.rs | 142 +--- src/config/builder.rs | 20 +- src/config/channels.rs | 37 +- src/config/embeddings.rs | 84 +- src/config/heartbeat.rs | 23 +- src/config/helpers.rs | 42 + src/config/hygiene.rs | 29 +- src/config/llm.rs | 5 +- src/config/routines.rs | 11 +- src/config/safety.rs | 11 +- src/config/sandbox.rs | 36 +- src/config/skills.rs | 11 +- src/config/wasm.rs | 20 +- src/llm/nearai_chat.rs | 269 ++++++- src/llm/reasoning.rs | 2 +- src/main.rs | 1561 ++++++++++++-------------------------- src/setup/wizard.rs | 5 +- 21 files changed, 1004 insertions(+), 1460 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9ee1a335..05b1b238 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4975,15 +4975,6 @@ dependencies = [ "syn 2.0.114", ] -[[package]] -name = "servo_arc" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "170fb83ab34de17dc69aa7c67482b22218ddb85da56546f9bd6b929e32a05930" -dependencies = [ - "stable_deref_trait", -] - [[package]] name = "serde_yml" version = "0.0.12" @@ -4999,6 +4990,15 @@ dependencies = [ "version_check", ] +[[package]] +name = "servo_arc" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "170fb83ab34de17dc69aa7c67482b22218ddb85da56546f9bd6b929e32a05930" +dependencies = [ + "stable_deref_trait", +] + [[package]] name = "sha1" version = "0.10.6" diff --git a/src/agent/cost_guard.rs b/src/agent/cost_guard.rs index 59d4ed85..59d676ca 100644 --- a/src/agent/cost_guard.rs +++ b/src/agent/cost_guard.rs @@ -151,14 +151,19 @@ impl CostGuard { /// Record a completed LLM action: its token costs and the action timestamp. /// /// Call this AFTER an LLM call completes so that costs are tracked. + /// + /// When `cost_per_token` is `Some`, those rates are used directly (provider- + /// sourced pricing). When `None`, falls back to the static `costs::model_cost` + /// lookup table, then `costs::default_cost`. pub async fn record_llm_call( &self, model: &str, input_tokens: u32, output_tokens: u32, + cost_per_token: Option<(Decimal, Decimal)>, ) -> Decimal { - let (input_rate, output_rate) = - costs::model_cost(model).unwrap_or_else(costs::default_cost); + let (input_rate, output_rate) = cost_per_token + .unwrap_or_else(|| costs::model_cost(model).unwrap_or_else(costs::default_cost)); let cost = input_rate * Decimal::from(input_tokens) + output_rate * Decimal::from(output_tokens); @@ -261,7 +266,9 @@ mod tests { assert!(guard.check_allowed().await.is_ok()); // Record a big call, still allowed - guard.record_llm_call("gpt-4o", 100_000, 100_000).await; + guard + .record_llm_call("gpt-4o", 100_000, 100_000, None) + .await; assert!(guard.check_allowed().await.is_ok()); } @@ -278,7 +285,7 @@ mod tests { // Record a call that costs more than $0.01 // gpt-4o: input=$0.0000025/tok, output=$0.00001/tok // 10000 input + 10000 output = $0.025 + $0.10 = $0.125 - guard.record_llm_call("gpt-4o", 10_000, 10_000).await; + guard.record_llm_call("gpt-4o", 10_000, 10_000, None).await; // Now should be blocked let result = guard.check_allowed().await; @@ -301,7 +308,7 @@ mod tests { // First 3 actions allowed for _ in 0..3 { assert!(guard.check_allowed().await.is_ok()); - guard.record_llm_call("gpt-4o", 10, 10).await; + guard.record_llm_call("gpt-4o", 10, 10, None).await; } // 4th should be blocked @@ -322,7 +329,7 @@ mod tests { assert_eq!(guard.daily_spend().await, Decimal::ZERO); - let cost = guard.record_llm_call("gpt-4o", 1000, 500).await; + let cost = guard.record_llm_call("gpt-4o", 1000, 500, None).await; assert!(cost > Decimal::ZERO); assert_eq!(guard.daily_spend().await, cost); } @@ -333,8 +340,8 @@ mod tests { assert_eq!(guard.actions_this_hour().await, 0); - guard.record_llm_call("gpt-4o", 10, 10).await; - guard.record_llm_call("gpt-4o", 10, 10).await; + guard.record_llm_call("gpt-4o", 10, 10, None).await; + guard.record_llm_call("gpt-4o", 10, 10, None).await; assert_eq!(guard.actions_this_hour().await, 2); } @@ -371,10 +378,10 @@ mod tests { assert!(guard.model_usage().await.is_empty()); // Record calls for two different models - guard.record_llm_call("gpt-4o", 1000, 500).await; - guard.record_llm_call("gpt-4o", 2000, 1000).await; + guard.record_llm_call("gpt-4o", 1000, 500, None).await; + guard.record_llm_call("gpt-4o", 2000, 1000, None).await; guard - .record_llm_call("claude-3-5-sonnet-20241022", 500, 200) + .record_llm_call("claude-3-5-sonnet-20241022", 500, 200, None) .await; let usage = guard.model_usage().await; diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 76d7e73a..1fecf803 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -222,6 +222,7 @@ impl Agent { &model_name, output.usage.input_tokens, output.usage.output_tokens, + Some(self.llm().cost_per_token()), ) .await; tracing::debug!( diff --git a/src/app.rs b/src/app.rs index 295d5645..1f9724a0 100644 --- a/src/app.rs +++ b/src/app.rs @@ -22,6 +22,7 @@ use crate::skills::SkillRegistry; use crate::skills::catalog::SkillCatalog; use crate::tools::ToolRegistry; use crate::tools::mcp::McpSessionManager; +use crate::tools::wasm::SharedCredentialRegistry; use crate::tools::wasm::WasmToolRuntime; use crate::workspace::{EmbeddingProvider, Workspace}; @@ -48,6 +49,8 @@ pub struct AppComponents { pub skill_catalog: Option>, pub cost_guard: Arc, pub session: Arc, + pub catalog_entries: Vec, + pub dev_loaded_tool_names: Vec, } /// Options that control optional init phases. @@ -313,54 +316,41 @@ impl AppBuilder { ), anyhow::Error, > { - use crate::workspace::{NearAiEmbeddings, OpenAiEmbeddings}; - let safety = Arc::new(SafetyLayer::new(&self.config.safety)); tracing::info!("Safety layer initialized"); - let tools = Arc::new(ToolRegistry::new()); + // 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(); - // Create embeddings provider if configured - let embeddings: Option> = if self.config.embeddings.enabled { - match self.config.embeddings.provider.as_str() { - "nearai" => { - tracing::info!( - "Embeddings enabled via NEAR AI (model: {})", - self.config.embeddings.model - ); - Some(Arc::new( - NearAiEmbeddings::new( - &self.config.llm.nearai.base_url, - self.session.clone(), - ) - .with_model(&self.config.embeddings.model, 1536), - )) - } - _ => { - if let Some(api_key) = self.config.embeddings.openai_api_key() { - tracing::info!( - "Embeddings enabled via OpenAI (model: {})", - self.config.embeddings.model - ); - Some(Arc::new(OpenAiEmbeddings::with_model( - api_key, - &self.config.embeddings.model, - match self.config.embeddings.model.as_str() { - "text-embedding-3-large" => 3072, - _ => 1536, - }, - ))) - } else { - tracing::warn!("Embeddings configured but OPENAI_API_KEY not set"); - None - } - } - } - } else { - tracing::info!("Embeddings disabled (set OPENAI_API_KEY or EMBEDDING_ENABLED=true)"); - None - }; + // Create embeddings provider using the unified method + let embeddings = self + .config + .embeddings + .create_provider(&self.config.llm.nearai.base_url, self.session.clone()); + + // Warn if libSQL backend is used with non-1536 embedding dimension. + if self.config.database.backend == crate::config::DatabaseBackend::LibSql + && self.config.embeddings.enabled + && self.config.embeddings.dimension != 1536 + { + tracing::warn!( + configured_dimension = self.config.embeddings.dimension, + "Embedding dimension {} is not 1536. The libSQL schema uses \ + F32_BLOB(1536) which requires exactly 1536 dimensions. \ + Embedding storage will fail. Use PostgreSQL or set \ + EMBEDDING_DIMENSION=1536.", + self.config.embeddings.dimension + ); + } // Register memory tools if database is available let workspace = if let Some(ref db) = self.db { @@ -402,6 +392,8 @@ impl AppBuilder { Arc, Option>, Option>, + Vec, + Vec, ), anyhow::Error, > { @@ -431,6 +423,8 @@ impl AppBuilder { let tools = Arc::clone(tools); let wasm_config = self.config.wasm.clone(); async move { + let mut dev_loaded_tool_names: Vec = 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 { @@ -461,10 +455,11 @@ impl AppBuilder { match load_dev_tools(&loader, &wasm_config.tools_dir).await { Ok(results) => { - if !results.loaded.is_empty() { + dev_loaded_tool_names.extend(results.loaded.iter().cloned()); + if !dev_loaded_tool_names.is_empty() { tracing::info!( "Loaded {} dev WASM tools from build artifacts", - results.loaded.len() + dev_loaded_tool_names.len() ); } } @@ -473,6 +468,8 @@ impl AppBuilder { } } } + + dev_loaded_tool_names } }; @@ -577,7 +574,7 @@ impl AppBuilder { } }; - tokio::join!(wasm_tools_future, mcp_servers_future); + let (dev_loaded_tool_names, _) = tokio::join!(wasm_tools_future, mcp_servers_future); // Load registry catalog entries for extension discovery let catalog_entries = match crate::registry::RegistryCatalog::load_or_embedded() { @@ -640,7 +637,13 @@ impl AppBuilder { tools.register_dev_tools(); } - Ok((mcp_session_manager, wasm_tool_runtime, extension_manager)) + Ok(( + mcp_session_manager, + wasm_tool_runtime, + extension_manager, + catalog_entries, + dev_loaded_tool_names, + )) } /// Run all init phases in order and return the assembled components. @@ -654,8 +657,13 @@ impl AppBuilder { // Create hook registry early so runtime extension activation can register hooks. let hooks = Arc::new(HookRegistry::new()); - let (mcp_session_manager, wasm_tool_runtime, extension_manager) = - self.init_extensions(&tools, &hooks).await?; + let ( + mcp_session_manager, + wasm_tool_runtime, + extension_manager, + catalog_entries, + dev_loaded_tool_names, + ) = self.init_extensions(&tools, &hooks).await?; // Seed workspace and backfill embeddings if let Some(ref ws) = workspace { @@ -730,6 +738,8 @@ impl AppBuilder { skill_catalog, cost_guard, session: self.session, + catalog_entries, + dev_loaded_tool_names, }) } } diff --git a/src/config/agent.rs b/src/config/agent.rs index e075d803..22089688 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use crate::config::helpers::optional_env; +use crate::config::helpers::{parse_bool_env, parse_option_env, parse_optional_env}; use crate::error::ConfigError; use crate::settings::Settings; @@ -32,109 +32,43 @@ pub struct AgentConfig { impl AgentConfig { pub(crate) fn resolve(settings: &Settings) -> Result { Ok(Self { - name: optional_env("AGENT_NAME")?.unwrap_or_else(|| settings.agent.name.clone()), - max_parallel_jobs: optional_env("AGENT_MAX_PARALLEL_JOBS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "AGENT_MAX_PARALLEL_JOBS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(settings.agent.max_parallel_jobs as usize), - job_timeout: Duration::from_secs( - optional_env("AGENT_JOB_TIMEOUT_SECS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "AGENT_JOB_TIMEOUT_SECS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(settings.agent.job_timeout_secs), - ), - stuck_threshold: Duration::from_secs( - optional_env("AGENT_STUCK_THRESHOLD_SECS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "AGENT_STUCK_THRESHOLD_SECS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(settings.agent.stuck_threshold_secs), - ), - repair_check_interval: Duration::from_secs( - optional_env("SELF_REPAIR_CHECK_INTERVAL_SECS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "SELF_REPAIR_CHECK_INTERVAL_SECS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(settings.agent.repair_check_interval_secs), - ), - max_repair_attempts: optional_env("SELF_REPAIR_MAX_ATTEMPTS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "SELF_REPAIR_MAX_ATTEMPTS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(settings.agent.max_repair_attempts), - use_planning: optional_env("AGENT_USE_PLANNING")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "AGENT_USE_PLANNING".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(settings.agent.use_planning), - session_idle_timeout: Duration::from_secs( - optional_env("SESSION_IDLE_TIMEOUT_SECS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "SESSION_IDLE_TIMEOUT_SECS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(settings.agent.session_idle_timeout_secs), - ), - allow_local_tools: optional_env("ALLOW_LOCAL_TOOLS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "ALLOW_LOCAL_TOOLS".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(false), - max_cost_per_day_cents: optional_env("MAX_COST_PER_DAY_CENTS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "MAX_COST_PER_DAY_CENTS".to_string(), - message: format!("must be a positive integer: {e}"), - })?, - max_actions_per_hour: optional_env("MAX_ACTIONS_PER_HOUR")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "MAX_ACTIONS_PER_HOUR".to_string(), - message: format!("must be a positive integer: {e}"), - })?, - max_tool_iterations: optional_env("AGENT_MAX_TOOL_ITERATIONS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "AGENT_MAX_TOOL_ITERATIONS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(settings.agent.max_tool_iterations), - auto_approve_tools: optional_env("AGENT_AUTO_APPROVE_TOOLS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "AGENT_AUTO_APPROVE_TOOLS".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(settings.agent.auto_approve_tools), + name: parse_optional_env("AGENT_NAME", settings.agent.name.clone())?, + max_parallel_jobs: parse_optional_env( + "AGENT_MAX_PARALLEL_JOBS", + settings.agent.max_parallel_jobs as usize, + )?, + job_timeout: Duration::from_secs(parse_optional_env( + "AGENT_JOB_TIMEOUT_SECS", + settings.agent.job_timeout_secs, + )?), + stuck_threshold: Duration::from_secs(parse_optional_env( + "AGENT_STUCK_THRESHOLD_SECS", + settings.agent.stuck_threshold_secs, + )?), + repair_check_interval: Duration::from_secs(parse_optional_env( + "SELF_REPAIR_CHECK_INTERVAL_SECS", + settings.agent.repair_check_interval_secs, + )?), + max_repair_attempts: parse_optional_env( + "SELF_REPAIR_MAX_ATTEMPTS", + settings.agent.max_repair_attempts, + )?, + use_planning: parse_bool_env("AGENT_USE_PLANNING", settings.agent.use_planning)?, + session_idle_timeout: Duration::from_secs(parse_optional_env( + "SESSION_IDLE_TIMEOUT_SECS", + settings.agent.session_idle_timeout_secs, + )?), + allow_local_tools: parse_bool_env("ALLOW_LOCAL_TOOLS", false)?, + max_cost_per_day_cents: parse_option_env("MAX_COST_PER_DAY_CENTS")?, + max_actions_per_hour: parse_option_env("MAX_ACTIONS_PER_HOUR")?, + max_tool_iterations: parse_optional_env( + "AGENT_MAX_TOOL_ITERATIONS", + settings.agent.max_tool_iterations, + )?, + auto_approve_tools: parse_bool_env( + "AGENT_AUTO_APPROVE_TOOLS", + settings.agent.auto_approve_tools, + )?, }) } } diff --git a/src/config/builder.rs b/src/config/builder.rs index fede5bce..90bbb185 100644 --- a/src/config/builder.rs +++ b/src/config/builder.rs @@ -1,7 +1,7 @@ use std::path::PathBuf; use std::time::Duration; -use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env}; use crate::error::ConfigError; /// Builder mode configuration. @@ -34,25 +34,11 @@ impl Default for BuilderModeConfig { impl BuilderModeConfig { pub(crate) fn resolve() -> Result { Ok(Self { - enabled: optional_env("BUILDER_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "BUILDER_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), + enabled: parse_bool_env("BUILDER_ENABLED", true)?, build_dir: optional_env("BUILDER_DIR")?.map(PathBuf::from), max_iterations: parse_optional_env("BUILDER_MAX_ITERATIONS", 20)?, timeout_secs: parse_optional_env("BUILDER_TIMEOUT_SECS", 600)?, - auto_register: optional_env("BUILDER_AUTO_REGISTER")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "BUILDER_AUTO_REGISTER".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), + auto_register: parse_bool_env("BUILDER_AUTO_REGISTER", true)?, }) } diff --git a/src/config/channels.rs b/src/config/channels.rs index 31eaffab..ccfdecf3 100644 --- a/src/config/channels.rs +++ b/src/config/channels.rs @@ -2,7 +2,7 @@ use std::path::PathBuf; use secrecy::SecretString; -use crate::config::helpers::optional_env; +use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env}; use crate::error::ConfigError; use crate::settings::Settings; @@ -48,14 +48,7 @@ impl ChannelsConfig { let http = if optional_env("HTTP_PORT")?.is_some() || optional_env("HTTP_HOST")?.is_some() { Some(HttpConfig { host: optional_env("HTTP_HOST")?.unwrap_or_else(|| "0.0.0.0".to_string()), - port: optional_env("HTTP_PORT")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "HTTP_PORT".to_string(), - message: format!("must be a valid port number: {e}"), - })? - .unwrap_or(8080), + port: parse_optional_env("HTTP_PORT", 8080)?, webhook_secret: optional_env("HTTP_WEBHOOK_SECRET")?.map(SecretString::from), user_id: optional_env("HTTP_USER_ID")?.unwrap_or_else(|| "http".to_string()), }) @@ -63,20 +56,11 @@ impl ChannelsConfig { None }; - let gateway = if optional_env("GATEWAY_ENABLED")? - .map(|s| s.to_lowercase() == "true" || s == "1") - .unwrap_or(true) - { + let gateway_enabled = parse_bool_env("GATEWAY_ENABLED", true)?; + let gateway = if gateway_enabled { Some(GatewayConfig { host: optional_env("GATEWAY_HOST")?.unwrap_or_else(|| "127.0.0.1".to_string()), - port: optional_env("GATEWAY_PORT")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "GATEWAY_PORT".to_string(), - message: format!("must be a valid port number: {e}"), - })? - .unwrap_or(3000), + port: parse_optional_env("GATEWAY_PORT", 3000)?, auth_token: optional_env("GATEWAY_AUTH_TOKEN")?, user_id: optional_env("GATEWAY_USER_ID")?.unwrap_or_else(|| "default".to_string()), }) @@ -97,18 +81,11 @@ impl ChannelsConfig { wasm_channels_dir: optional_env("WASM_CHANNELS_DIR")? .map(PathBuf::from) .unwrap_or_else(default_channels_dir), - wasm_channels_enabled: optional_env("WASM_CHANNELS_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "WASM_CHANNELS_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), + wasm_channels_enabled: parse_bool_env("WASM_CHANNELS_ENABLED", true)?, telegram_owner_id: optional_env("TELEGRAM_OWNER_ID")? .map(|s| s.parse()) .transpose() - .map_err(|e| ConfigError::InvalidValue { + .map_err(|e: std::num::ParseIntError| ConfigError::InvalidValue { key: "TELEGRAM_OWNER_ID".to_string(), message: format!("must be an integer: {e}"), })? diff --git a/src/config/embeddings.rs b/src/config/embeddings.rs index 4528aded..501be22c 100644 --- a/src/config/embeddings.rs +++ b/src/config/embeddings.rs @@ -1,8 +1,12 @@ +use std::sync::Arc; + use secrecy::{ExposeSecret, SecretString}; -use crate::config::helpers::optional_env; +use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env}; use crate::error::ConfigError; +use crate::llm::SessionManager; use crate::settings::Settings; +use crate::workspace::EmbeddingProvider; /// Embeddings provider configuration. #[derive(Debug, Clone)] @@ -65,23 +69,10 @@ impl EmbeddingsConfig { .or_else(|| settings.ollama_base_url.clone()) .unwrap_or_else(|| "http://localhost:11434".to_string()); - let dimension = optional_env("EMBEDDING_DIMENSION")? - .map(|s| s.parse::()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "EMBEDDING_DIMENSION".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or_else(|| default_dimension_for_model(&model)); + let dimension = + parse_optional_env("EMBEDDING_DIMENSION", default_dimension_for_model(&model))?; - let enabled = optional_env("EMBEDDING_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "EMBEDDING_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(settings.embeddings.enabled); + let enabled = parse_bool_env("EMBEDDING_ENABLED", settings.embeddings.enabled)?; Ok(Self { enabled, @@ -97,6 +88,65 @@ impl EmbeddingsConfig { pub fn openai_api_key(&self) -> Option<&str> { self.openai_api_key.as_ref().map(|s| s.expose_secret()) } + + /// Create the appropriate embedding provider based on configuration. + /// + /// Returns `None` if embeddings are disabled or the required credentials + /// are missing. The `nearai_base_url` and `session` are needed only for + /// the NEAR AI provider but must be passed unconditionally. + pub fn create_provider( + &self, + nearai_base_url: &str, + session: Arc, + ) -> Option> { + if !self.enabled { + tracing::info!("Embeddings disabled (set EMBEDDING_ENABLED=true to enable)"); + return None; + } + + match self.provider.as_str() { + "nearai" => { + tracing::info!( + "Embeddings enabled via NEAR AI (model: {}, dim: {})", + self.model, + self.dimension, + ); + Some(Arc::new( + crate::workspace::NearAiEmbeddings::new(nearai_base_url, session) + .with_model(&self.model, self.dimension), + )) + } + "ollama" => { + tracing::info!( + "Embeddings enabled via Ollama (model: {}, url: {}, dim: {})", + self.model, + self.ollama_base_url, + self.dimension, + ); + Some(Arc::new( + crate::workspace::OllamaEmbeddings::new(&self.ollama_base_url) + .with_model(&self.model, self.dimension), + )) + } + _ => { + if let Some(api_key) = self.openai_api_key() { + tracing::info!( + "Embeddings enabled via OpenAI (model: {}, dim: {})", + self.model, + self.dimension, + ); + Some(Arc::new(crate::workspace::OpenAiEmbeddings::with_model( + api_key, + &self.model, + self.dimension, + ))) + } else { + tracing::warn!("Embeddings configured but OPENAI_API_KEY not set"); + None + } + } + } + } } #[cfg(test)] diff --git a/src/config/heartbeat.rs b/src/config/heartbeat.rs index 9fe0831b..f2f98071 100644 --- a/src/config/heartbeat.rs +++ b/src/config/heartbeat.rs @@ -1,4 +1,4 @@ -use crate::config::helpers::optional_env; +use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env}; use crate::error::ConfigError; use crate::settings::Settings; @@ -29,22 +29,11 @@ impl Default for HeartbeatConfig { impl HeartbeatConfig { pub(crate) fn resolve(settings: &Settings) -> Result { Ok(Self { - enabled: optional_env("HEARTBEAT_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "HEARTBEAT_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(settings.heartbeat.enabled), - interval_secs: optional_env("HEARTBEAT_INTERVAL_SECS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "HEARTBEAT_INTERVAL_SECS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(settings.heartbeat.interval_secs), + enabled: parse_bool_env("HEARTBEAT_ENABLED", settings.heartbeat.enabled)?, + interval_secs: parse_optional_env( + "HEARTBEAT_INTERVAL_SECS", + settings.heartbeat.interval_secs, + )?, notify_channel: optional_env("HEARTBEAT_NOTIFY_CHANNEL")? .or_else(|| settings.heartbeat.notify_channel.clone()), notify_user: optional_env("HEARTBEAT_NOTIFY_USER")? diff --git a/src/config/helpers.rs b/src/config/helpers.rs index e9e966df..8db271d4 100644 --- a/src/config/helpers.rs +++ b/src/config/helpers.rs @@ -47,3 +47,45 @@ where .transpose() .map(|opt| opt.unwrap_or(default)) } + +/// Parse a boolean from an env var with a default. +/// +/// Accepts "true"/"1" as true, "false"/"0" as false. +pub(crate) fn parse_bool_env(key: &str, default: bool) -> Result { + match optional_env(key)? { + Some(s) => match s.to_lowercase().as_str() { + "true" | "1" => Ok(true), + "false" | "0" => Ok(false), + _ => Err(ConfigError::InvalidValue { + key: key.to_string(), + message: format!("must be 'true' or 'false', got '{s}'"), + }), + }, + None => Ok(default), + } +} + +/// Parse an env var into `Option` — returns `None` when unset, +/// `Some(parsed)` when set to a valid value. +pub(crate) fn parse_option_env(key: &str) -> Result, ConfigError> +where + T: std::str::FromStr, + T::Err: std::fmt::Display, +{ + optional_env(key)? + .map(|s| { + s.parse().map_err(|e| ConfigError::InvalidValue { + key: key.to_string(), + message: format!("{e}"), + }) + }) + .transpose() +} + +/// Parse a string from an env var with a default. +pub(crate) fn parse_string_env( + key: &str, + default: impl Into, +) -> Result { + Ok(optional_env(key)?.unwrap_or_else(|| default.into())) +} diff --git a/src/config/hygiene.rs b/src/config/hygiene.rs index f3d3f414..174ab7c9 100644 --- a/src/config/hygiene.rs +++ b/src/config/hygiene.rs @@ -1,4 +1,4 @@ -use crate::config::helpers::optional_env; +use crate::config::helpers::{parse_bool_env, parse_optional_env}; use crate::error::ConfigError; /// Memory hygiene configuration. @@ -28,30 +28,9 @@ impl Default for HygieneConfig { impl HygieneConfig { pub(crate) fn resolve() -> Result { Ok(Self { - enabled: optional_env("MEMORY_HYGIENE_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "MEMORY_HYGIENE_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), - retention_days: optional_env("MEMORY_HYGIENE_RETENTION_DAYS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "MEMORY_HYGIENE_RETENTION_DAYS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(30), - cadence_hours: optional_env("MEMORY_HYGIENE_CADENCE_HOURS")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "MEMORY_HYGIENE_CADENCE_HOURS".to_string(), - message: format!("must be a positive integer: {e}"), - })? - .unwrap_or(12), + enabled: parse_bool_env("MEMORY_HYGIENE_ENABLED", true)?, + retention_days: parse_optional_env("MEMORY_HYGIENE_RETENTION_DAYS", 30)?, + cadence_hours: parse_optional_env("MEMORY_HYGIENE_CADENCE_HOURS", 12)?, }) } diff --git a/src/config/llm.rs b/src/config/llm.rs index bb49a7b0..60ff9d7f 100644 --- a/src/config/llm.rs +++ b/src/config/llm.rs @@ -206,10 +206,7 @@ impl LlmConfig { let nearai = NearAiConfig { model: optional_env("NEARAI_MODEL")? .or_else(|| settings.selected_model.clone()) - .unwrap_or_else(|| { - "fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic" - .to_string() - }), + .unwrap_or_else(|| "zai-org/GLM-latest".to_string()), cheap_model: optional_env("NEARAI_CHEAP_MODEL")?, base_url: optional_env("NEARAI_BASE_URL")?.unwrap_or_else(|| { if nearai_api_key.is_some() { diff --git a/src/config/routines.rs b/src/config/routines.rs index 03b890de..4357e02b 100644 --- a/src/config/routines.rs +++ b/src/config/routines.rs @@ -1,4 +1,4 @@ -use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::config::helpers::{parse_bool_env, parse_optional_env}; use crate::error::ConfigError; /// Routines configuration. @@ -31,14 +31,7 @@ impl Default for RoutineConfig { impl RoutineConfig { pub(crate) fn resolve() -> Result { Ok(Self { - enabled: optional_env("ROUTINES_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "ROUTINES_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), + enabled: parse_bool_env("ROUTINES_ENABLED", true)?, cron_check_interval_secs: parse_optional_env("ROUTINES_CRON_INTERVAL", 15)?, max_concurrent_routines: parse_optional_env("ROUTINES_MAX_CONCURRENT", 10)?, default_cooldown_secs: parse_optional_env("ROUTINES_DEFAULT_COOLDOWN", 300)?, diff --git a/src/config/safety.rs b/src/config/safety.rs index 21483d73..19c70719 100644 --- a/src/config/safety.rs +++ b/src/config/safety.rs @@ -1,4 +1,4 @@ -use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::config::helpers::{parse_bool_env, parse_optional_env}; use crate::error::ConfigError; /// Safety configuration. @@ -12,14 +12,7 @@ impl SafetyConfig { pub(crate) fn resolve() -> Result { Ok(Self { max_output_length: parse_optional_env("SAFETY_MAX_OUTPUT_LENGTH", 100_000)?, - injection_check_enabled: optional_env("SAFETY_INJECTION_CHECK_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "SAFETY_INJECTION_CHECK_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), + injection_check_enabled: parse_bool_env("SAFETY_INJECTION_CHECK_ENABLED", true)?, }) } } diff --git a/src/config/sandbox.rs b/src/config/sandbox.rs index 57a016fc..85c9c4b2 100644 --- a/src/config/sandbox.rs +++ b/src/config/sandbox.rs @@ -1,4 +1,4 @@ -use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env, parse_string_env}; use crate::error::ConfigError; /// Docker sandbox configuration. @@ -44,28 +44,13 @@ impl SandboxModeConfig { .unwrap_or_default(); Ok(Self { - enabled: optional_env("SANDBOX_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "SANDBOX_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), - policy: optional_env("SANDBOX_POLICY")?.unwrap_or_else(|| "readonly".to_string()), + enabled: parse_bool_env("SANDBOX_ENABLED", true)?, + policy: parse_string_env("SANDBOX_POLICY", "readonly")?, timeout_secs: parse_optional_env("SANDBOX_TIMEOUT_SECS", 120)?, memory_limit_mb: parse_optional_env("SANDBOX_MEMORY_LIMIT_MB", 2048)?, cpu_shares: parse_optional_env("SANDBOX_CPU_SHARES", 1024)?, - image: optional_env("SANDBOX_IMAGE")? - .unwrap_or_else(|| "ironclaw-worker:latest".to_string()), - auto_pull_image: optional_env("SANDBOX_AUTO_PULL")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "SANDBOX_AUTO_PULL".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), + image: parse_string_env("SANDBOX_IMAGE", "ironclaw-worker:latest")?, + auto_pull_image: parse_bool_env("SANDBOX_AUTO_PULL", true)?, extra_allowed_domains: extra_domains, }) } @@ -221,18 +206,11 @@ impl ClaudeCodeConfig { pub(crate) fn resolve() -> Result { let defaults = Self::default(); Ok(Self { - enabled: optional_env("CLAUDE_CODE_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "CLAUDE_CODE_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(defaults.enabled), + enabled: parse_bool_env("CLAUDE_CODE_ENABLED", defaults.enabled)?, config_dir: optional_env("CLAUDE_CONFIG_DIR")? .map(std::path::PathBuf::from) .unwrap_or(defaults.config_dir), - model: optional_env("CLAUDE_CODE_MODEL")?.unwrap_or(defaults.model), + model: parse_string_env("CLAUDE_CODE_MODEL", defaults.model)?, max_turns: parse_optional_env("CLAUDE_CODE_MAX_TURNS", defaults.max_turns)?, memory_limit_mb: parse_optional_env( "CLAUDE_CODE_MEMORY_LIMIT_MB", diff --git a/src/config/skills.rs b/src/config/skills.rs index 71386e74..e58e41b5 100644 --- a/src/config/skills.rs +++ b/src/config/skills.rs @@ -1,6 +1,6 @@ use std::path::PathBuf; -use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env}; use crate::error::ConfigError; /// Skills system configuration. @@ -38,14 +38,7 @@ fn default_skills_dir() -> PathBuf { impl SkillsConfig { pub(crate) fn resolve() -> Result { Ok(Self { - enabled: optional_env("SKILLS_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "SKILLS_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(false), + enabled: parse_bool_env("SKILLS_ENABLED", false)?, local_dir: optional_env("SKILLS_DIR")? .map(PathBuf::from) .unwrap_or_else(default_skills_dir), diff --git a/src/config/wasm.rs b/src/config/wasm.rs index 5d13fa1d..9d069c8b 100644 --- a/src/config/wasm.rs +++ b/src/config/wasm.rs @@ -1,7 +1,7 @@ use std::path::PathBuf; use std::time::Duration; -use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env}; use crate::error::ConfigError; /// WASM sandbox configuration. @@ -48,14 +48,7 @@ fn default_tools_dir() -> PathBuf { impl WasmConfig { pub(crate) fn resolve() -> Result { Ok(Self { - enabled: optional_env("WASM_ENABLED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "WASM_ENABLED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), + enabled: parse_bool_env("WASM_ENABLED", true)?, tools_dir: optional_env("WASM_TOOLS_DIR")? .map(PathBuf::from) .unwrap_or_else(default_tools_dir), @@ -65,14 +58,7 @@ impl WasmConfig { )?, default_timeout_secs: parse_optional_env("WASM_DEFAULT_TIMEOUT_SECS", 60)?, default_fuel_limit: parse_optional_env("WASM_DEFAULT_FUEL_LIMIT", 10_000_000)?, - cache_compiled: optional_env("WASM_CACHE_COMPILED")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "WASM_CACHE_COMPILED".to_string(), - message: format!("must be 'true' or 'false': {e}"), - })? - .unwrap_or(true), + cache_compiled: parse_bool_env("WASM_CACHE_COMPILED", true)?, cache_dir: optional_env("WASM_CACHE_DIR")?.map(PathBuf::from), }) } diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index 8d87dc37..cc4d13fc 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -6,12 +6,13 @@ //! - **Session token auth**: Otherwise, uses `SessionManager` for Bearer session token //! with automatic renewal on 401 errors +use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; use reqwest::Client; use rust_decimal::Decimal; -use rust_decimal_macros::dec; +use rust_decimal::prelude::MathematicalOps; use secrecy::ExposeSecret; use serde::{Deserialize, Serialize}; @@ -21,7 +22,7 @@ use crate::llm::provider::{ ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, }; -use crate::llm::session::SessionManager; +use crate::llm::{costs, session::SessionManager}; /// Information about an available model from NEAR AI API. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -42,6 +43,9 @@ pub struct NearAiChatProvider { session: Arc, active_model: std::sync::RwLock, flatten_tool_messages: bool, + /// Per-model pricing fetched from the NEAR AI `/v1/model/list` endpoint. + /// Maps model ID → (input_cost_per_token, output_cost_per_token). + pricing: Arc>>, } impl NearAiChatProvider { @@ -72,13 +76,49 @@ impl NearAiChatProvider { })?; let active_model = std::sync::RwLock::new(config.model.clone()); - Ok(Self { + let pricing = Arc::new(std::sync::RwLock::new(HashMap::new())); + + let provider = Self { client, config, session, active_model, flatten_tool_messages, - }) + pricing, + }; + + // Fire-and-forget background pricing fetch — don't block startup. + // Only spawns when a tokio runtime is active (skipped in sync tests). + if let Ok(handle) = tokio::runtime::Handle::try_current() { + let client = provider.client.clone(); + let base_url = provider.config.base_url.clone(); + let api_key = provider.config.api_key.clone(); + let session = provider.session.clone(); + let pricing = provider.pricing.clone(); + + handle.spawn(async move { + match fetch_pricing(&client, &base_url, api_key.as_ref(), &session).await { + Ok(map) if !map.is_empty() => { + tracing::info!("Loaded NEAR AI pricing for {} model(s)", map.len()); + match pricing.write() { + Ok(mut guard) => *guard = map, + Err(poisoned) => *poisoned.into_inner() = map, + } + } + Ok(_) => { + tracing::debug!("NEAR AI pricing endpoint returned no pricing data"); + } + Err(e) => { + tracing::debug!( + "Could not fetch NEAR AI pricing (will use fallback): {}", + e + ); + } + } + }); + } + + Ok(provider) } fn api_url(&self, path: &str) -> String { @@ -500,8 +540,14 @@ impl LlmProvider for NearAiChatProvider { } fn cost_per_token(&self) -> (Decimal, Decimal) { - // Default costs - could be model-specific in the future - (dec!(0.000003), dec!(0.000015)) + let model = self.active_model_name(); + // Try fetched pricing first, then static lookup table, then default + if let Ok(guard) = self.pricing.read() + && let Some(&rates) = guard.get(&model) + { + return rates; + } + costs::model_cost(&model).unwrap_or_else(costs::default_cost) } async fn list_models(&self) -> Result, LlmError> { @@ -562,6 +608,143 @@ struct ChatCompletionMessage { tool_calls: Option>, } +// -- Pricing fetch types and logic ----------------------------------------- + +/// Cost amount from the NEAR AI `/v1/model/list` response. +/// +/// Real cost per token = `amount * 10^(-scale)`. +#[derive(Debug, Deserialize)] +struct ModelCost { + amount: f64, + #[serde(default)] + scale: i32, +} + +/// A single model entry from the pricing response. +#[derive(Debug, Deserialize)] +struct PricingModelEntry { + #[serde(default, alias = "modelId", alias = "model_id")] + model_id: Option, + #[serde(default, alias = "inputCostPerToken")] + input_cost_per_token: Option, + #[serde(default, alias = "outputCostPerToken")] + output_cost_per_token: Option, + #[serde(default)] + metadata: Option, +} + +#[derive(Debug, Deserialize)] +struct PricingMetadata { + #[serde(default)] + aliases: Vec, +} + +/// Wrapper for the `/v1/model/list` response body. +#[derive(Debug, Deserialize)] +struct PricingResponse { + #[serde(default)] + models: Option>, + #[serde(default)] + data: Option>, +} + +/// Convert a `ModelCost` to a `Decimal` per-token price. +fn model_cost_to_decimal(mc: &ModelCost) -> Option { + if mc.amount == 0.0 { + return Some(Decimal::ZERO); + } + // amount * 10^(-scale) + let base = Decimal::try_from(mc.amount).ok()?; + let factor = Decimal::TEN.checked_powi(-i64::from(mc.scale))?; + base.checked_mul(factor) +} + +/// Fetch pricing from the NEAR AI `/v1/model/list` endpoint. +/// +/// Returns a map of model_id → (input_cost_per_token, output_cost_per_token). +/// Errors are non-fatal; callers should fall back to the static lookup table. +async fn fetch_pricing( + client: &Client, + base_url: &str, + api_key: Option<&secrecy::SecretString>, + session: &SessionManager, +) -> Result, LlmError> { + let base = base_url.trim_end_matches('/'); + let url = if base.ends_with("/v1") { + format!("{}/model/list", base) + } else { + format!("{}/v1/model/list", base) + }; + + let token = if let Some(key) = api_key { + key.expose_secret().to_string() + } else { + let tok = session.get_token().await?; + tok.expose_secret().to_string() + }; + + let response = client + .get(&url) + .header("Authorization", format!("Bearer {}", token)) + .timeout(std::time::Duration::from_secs(15)) + .send() + .await + .map_err(|e| LlmError::RequestFailed { + provider: "nearai_chat".to_string(), + reason: format!("Failed to fetch pricing: {}", e), + })?; + + if !response.status().is_success() { + return Err(LlmError::RequestFailed { + provider: "nearai_chat".to_string(), + reason: format!("Pricing endpoint returned HTTP {}", response.status()), + }); + } + + let body = response.text().await.map_err(|e| LlmError::RequestFailed { + provider: "nearai_chat".to_string(), + reason: format!("Failed to read pricing response: {}", e), + })?; + + // Parse as {models: [...]} or {data: [...]} or direct array + let entries: Vec = + if let Ok(resp) = serde_json::from_str::(&body) { + resp.models.or(resp.data).unwrap_or_default() + } else if let Ok(arr) = serde_json::from_str::>(&body) { + arr + } else { + return Ok(HashMap::new()); + }; + + let mut map = HashMap::new(); + for entry in &entries { + let (Some(input_mc), Some(output_mc)) = + (&entry.input_cost_per_token, &entry.output_cost_per_token) + else { + continue; + }; + let (Some(input), Some(output)) = ( + model_cost_to_decimal(input_mc), + model_cost_to_decimal(output_mc), + ) else { + continue; + }; + + // Insert under the primary model_id + if let Some(ref id) = entry.model_id { + map.insert(id.clone(), (input, output)); + } + // Also insert under any aliases + if let Some(ref meta) = entry.metadata { + for alias in &meta.aliases { + map.insert(alias.clone(), (input, output)); + } + } + } + + Ok(map) +} + /// Rewrite tool-call / tool-result messages into plain assistant/user text. /// /// NEAR AI cloud-api does not support the OpenAI multi-turn tool-calling @@ -748,6 +931,7 @@ fn parse_usage(usage: Option<&ChatCompletionUsage>) -> (u32, u32) { mod tests { use super::*; use crate::llm::session::SessionConfig; + use rust_decimal_macros::dec; fn test_nearai_config(base_url: &str) -> NearAiConfig { NearAiConfig { @@ -991,4 +1175,77 @@ mod tests { assert!(text.starts_with("Let me check that.")); assert!(text.contains("[Called tool `search`")); } + + #[test] + fn test_model_cost_to_decimal_basic() { + // amount=3, scale=6 → 3 * 10^-6 = 0.000003 + let mc = ModelCost { + amount: 3.0, + scale: 6, + }; + let result = model_cost_to_decimal(&mc).unwrap(); + assert_eq!(result, dec!(0.000003)); + } + + #[test] + fn test_model_cost_to_decimal_zero() { + let mc = ModelCost { + amount: 0.0, + scale: 6, + }; + assert_eq!(model_cost_to_decimal(&mc), Some(Decimal::ZERO)); + } + + #[test] + fn test_model_cost_to_decimal_larger_scale() { + // amount=85, scale=8 → 85 * 10^-8 = 0.00000085 + let mc = ModelCost { + amount: 85.0, + scale: 8, + }; + let result = model_cost_to_decimal(&mc).unwrap(); + assert_eq!(result, dec!(0.00000085)); + } + + #[test] + fn test_cost_per_token_uses_pricing_map() { + let cfg = test_nearai_config("http://127.0.0.1:8318"); + let provider = NearAiChatProvider::new(cfg, test_session()).expect("provider"); + + // Inject pricing directly + { + let mut guard = provider.pricing.write().unwrap(); + guard.insert("test-model".to_string(), (dec!(0.000001), dec!(0.000005))); + } + + let (input, output) = provider.cost_per_token(); + assert_eq!(input, dec!(0.000001)); + assert_eq!(output, dec!(0.000005)); + } + + #[test] + fn test_cost_per_token_falls_back_to_static() { + let mut cfg = test_nearai_config("http://127.0.0.1:8318"); + cfg.model = "gpt-4o".to_string(); + let provider = NearAiChatProvider::new(cfg, test_session()).expect("provider"); + + // No pricing in map, should fall back to static costs::model_cost + let (input, output) = provider.cost_per_token(); + let (expected_in, expected_out) = costs::model_cost("gpt-4o").unwrap(); + assert_eq!(input, expected_in); + assert_eq!(output, expected_out); + } + + #[test] + fn test_cost_per_token_falls_back_to_default() { + let mut cfg = test_nearai_config("http://127.0.0.1:8318"); + cfg.model = "some-unknown-nearai-model".to_string(); + let provider = NearAiChatProvider::new(cfg, test_session()).expect("provider"); + + // No pricing in map, not in static table, should use default_cost + let (input, output) = provider.cost_per_token(); + let (default_in, default_out) = costs::default_cost(); + assert_eq!(input, default_in); + assert_eq!(output, default_out); + } } diff --git a/src/llm/reasoning.rs b/src/llm/reasoning.rs index d78f1b6f..14b8eb89 100644 --- a/src/llm/reasoning.rs +++ b/src/llm/reasoning.rs @@ -614,7 +614,7 @@ Respond with a JSON plan in this format: let group_section = self.build_group_section(); format!( - r#"You are NEAR AI Agent, an autonomous assistant. + r#"You are IronClaw Agent, a secure autonomous assistant. ## Response Format — CRITICAL diff --git a/src/main.rs b/src/main.rs index 0e563ed0..fef4abf7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,7 +6,8 @@ use clap::Parser; use tracing_subscriber::EnvFilter; use ironclaw::{ - agent::{Agent, AgentDeps, SessionManager}, + agent::{Agent, AgentDeps}, + app::{AppBuilder, AppBuilderFlags}, channels::{ ChannelManager, GatewayChannel, HttpChannel, ReplChannel, WebhookServer, WebhookServerConfig, @@ -21,34 +22,28 @@ use ironclaw::{ run_status_command, run_tool_command, }, config::Config, - context::ContextManager, - extensions::ExtensionManager, - hooks::{HookRegistry, bootstrap_hooks}, - llm::{SessionConfig, build_provider_chain, create_session_manager}, + hooks::bootstrap_hooks, + llm::{SessionConfig, create_session_manager}, orchestrator::{ ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore, api::OrchestratorState, }, pairing::PairingStore, - safety::SafetyLayer, secrets::SecretsStore, - tools::{ - ToolRegistry, - mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated}, - wasm::{WasmToolLoader, WasmToolRuntime, load_dev_tools}, - }, - workspace::{ - EmbeddingProvider, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings, Workspace, - }, }; -#[cfg(feature = "libsql")] -use ironclaw::secrets::LibSqlSecretsStore; -#[cfg(feature = "postgres")] -use ironclaw::secrets::PostgresSecretsStore; -use ironclaw::secrets::SecretsCrypto; #[cfg(any(feature = "postgres", feature = "libsql"))] use ironclaw::setup::{SetupConfig, SetupWizard}; + +/// Initialize tracing for simple CLI commands (warn level, no fancy layers). +fn init_cli_tracing() { + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), + ) + .init(); +} + #[tokio::main] async fn main() -> anyhow::Result<()> { let cli = Cli::parse(); @@ -56,161 +51,43 @@ async fn main() -> anyhow::Result<()> { // Handle non-agent commands first (they don't need full setup) match &cli.command { Some(Command::Tool(tool_cmd)) => { - // Simple logging for CLI commands - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - + init_cli_tracing(); return run_tool_command(tool_cmd.clone()).await; } Some(Command::Config(config_cmd)) => { - // Config commands need DB access for settings - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - + init_cli_tracing(); return ironclaw::cli::run_config_command(config_cmd.clone()).await; } Some(Command::Registry(registry_cmd)) => { - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - + init_cli_tracing(); return ironclaw::cli::run_registry_command(registry_cmd.clone()).await; } Some(Command::Mcp(mcp_cmd)) => { - // Simple logging for MCP commands - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - + init_cli_tracing(); return run_mcp_command(mcp_cmd.clone()).await; } Some(Command::Memory(mem_cmd)) => { - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - - // Memory commands need database (and optionally embeddings) - let config = Config::from_env() - .await - .map_err(|e| anyhow::anyhow!("{}", e))?; - - // Set up embeddings if available - let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig { - auth_base_url: config.llm.nearai.auth_base_url.clone(), - session_path: config.llm.nearai.session_path.clone(), - }) - .await; - - let embeddings: Option> = - if config.embeddings.enabled { - match config.embeddings.provider.as_str() { - "nearai" => Some(Arc::new( - ironclaw::workspace::NearAiEmbeddings::new( - &config.llm.nearai.base_url, - session, - ) - .with_model(&config.embeddings.model, config.embeddings.dimension), - )), - "ollama" => Some(Arc::new( - ironclaw::workspace::OllamaEmbeddings::new( - &config.embeddings.ollama_base_url, - ) - .with_model(&config.embeddings.model, config.embeddings.dimension), - )), - _ => { - if let Some(api_key) = config.embeddings.openai_api_key() { - Some(Arc::new(ironclaw::workspace::OpenAiEmbeddings::with_model( - api_key, - &config.embeddings.model, - config.embeddings.dimension, - ))) - } else { - None - } - } - } - } else { - None - }; - - // Warn if libSQL backend is used with non-1536 embedding dimension. - // libSQL schema uses F32_BLOB(1536) which cannot be altered without a - // table rebuild, so non-1536 embeddings will cause storage failures. - if config.database.backend == ironclaw::config::DatabaseBackend::LibSql - && config.embeddings.enabled - && config.embeddings.dimension != 1536 - { - tracing::warn!( - configured_dimension = config.embeddings.dimension, - "Embedding dimension {} is not 1536. The libSQL schema uses \ - F32_BLOB(1536) which requires exactly 1536 dimensions. \ - Embedding storage will fail. Use PostgreSQL or set \ - EMBEDDING_DIMENSION=1536.", - config.embeddings.dimension - ); - } - - // Create a Database-trait-backed workspace for the memory command - let db: Arc = - ironclaw::db::connect_from_config(&config.database) - .await - .map_err(|e| anyhow::anyhow!("{}", e))?; - - return ironclaw::cli::run_memory_command_with_db(mem_cmd.clone(), db, embeddings) - .await; + init_cli_tracing(); + return run_memory_command(mem_cmd).await; } Some(Command::Pairing(pairing_cmd)) => { - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - + init_cli_tracing(); return run_pairing_command(pairing_cmd.clone()).map_err(|e| anyhow::anyhow!("{}", e)); } Some(Command::Service(service_cmd)) => { - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - + init_cli_tracing(); return run_service_command(service_cmd); } Some(Command::Doctor) => { - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - + init_cli_tracing(); let _ = dotenvy::dotenv(); ironclaw::bootstrap::load_ironclaw_env(); - return ironclaw::cli::run_doctor_command().await; } Some(Command::Status) => { - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), - ) - .init(); - + init_cli_tracing(); let _ = dotenvy::dotenv(); ironclaw::bootstrap::load_ironclaw_env(); - return run_status_command().await; } Some(Command::Worker { @@ -218,37 +95,8 @@ async fn main() -> anyhow::Result<()> { orchestrator_url, max_iterations, }) => { - // Worker mode: runs inside a Docker container. - // Simple logging (no TUI, no DB, no channels). - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env() - .unwrap_or_else(|_| EnvFilter::new("ironclaw=info")), - ) - .init(); - - tracing::info!( - "Starting worker for job {} (orchestrator: {})", - job_id, - orchestrator_url - ); - - let config = ironclaw::worker::runtime::WorkerConfig { - job_id: *job_id, - orchestrator_url: orchestrator_url.clone(), - max_iterations: *max_iterations, - timeout: std::time::Duration::from_secs(600), - }; - - let runtime = ironclaw::worker::WorkerRuntime::new(config) - .map_err(|e| anyhow::anyhow!("Worker init failed: {}", e))?; - - runtime - .run() - .await - .map_err(|e| anyhow::anyhow!("Worker failed: {}", e))?; - - return Ok(()); + init_worker_tracing(); + return run_worker(*job_id, orchestrator_url, *max_iterations).await; } Some(Command::ClaudeBridge { job_id, @@ -256,47 +104,13 @@ async fn main() -> anyhow::Result<()> { max_turns, model, }) => { - // Claude Code bridge mode: runs inside a Docker container. - // Spawns the `claude` CLI and streams output to the orchestrator. - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env() - .unwrap_or_else(|_| EnvFilter::new("ironclaw=info")), - ) - .init(); - - tracing::info!( - "Starting Claude Code bridge for job {} (orchestrator: {}, model: {})", - job_id, - orchestrator_url, - model - ); - - let config = ironclaw::worker::claude_bridge::ClaudeBridgeConfig { - job_id: *job_id, - orchestrator_url: orchestrator_url.clone(), - max_turns: *max_turns, - model: model.clone(), - timeout: std::time::Duration::from_secs(1800), - allowed_tools: ironclaw::config::ClaudeCodeConfig::from_env().allowed_tools, - }; - - let runtime = ironclaw::worker::ClaudeBridgeRuntime::new(config) - .map_err(|e| anyhow::anyhow!("Claude bridge init failed: {}", e))?; - - runtime - .run() - .await - .map_err(|e| anyhow::anyhow!("Claude bridge failed: {}", e))?; - - return Ok(()); + init_worker_tracing(); + return run_claude_bridge(*job_id, orchestrator_url, *max_turns, model).await; } Some(Command::Onboard { skip_auth, channels_only, }) => { - // Load .env files before running onboarding wizard. - // Standard ./.env first (higher priority), then ~/.ironclaw/.env. let _ = dotenvy::dotenv(); ironclaw::bootstrap::load_ironclaw_env(); @@ -321,9 +135,10 @@ async fn main() -> anyhow::Result<()> { } } + // ── Agent startup ────────────────────────────────────────────────── + // Load .env files early so DATABASE_URL (and any other vars) are // available to all subsequent env-based config resolution. - // Standard ./.env first (higher priority), then ~/.ironclaw/.env. let _ = dotenvy::dotenv(); ironclaw::bootstrap::load_ironclaw_env(); @@ -340,7 +155,7 @@ async fn main() -> anyhow::Result<()> { // Load initial config from env + disk + optional TOML (before DB is available) let toml_path = cli.config.as_deref(); - let mut config = match Config::from_env_with_toml(toml_path).await { + let config = match Config::from_env_with_toml(toml_path).await { Ok(c) => c, Err(ironclaw::error::ConfigError::MissingRequired { key, hint }) => { eprintln!("Configuration error: Missing required setting '{}'", key); @@ -362,617 +177,45 @@ async fn main() -> anyhow::Result<()> { let session = create_session_manager(session_config).await; // Create log broadcaster before tracing init so the WebLogLayer can capture all events. - // This gets wired to the gateway's /api/logs/events SSE endpoint later. let log_broadcaster = Arc::new(LogBroadcaster::new()); // Initialize tracing with a reloadable EnvFilter so the gateway can switch - // log levels (e.g. ironclaw=debug) at runtime without restarting. + // log levels at runtime without restarting. let log_level_handle = ironclaw::channels::web::log_layer::init_tracing(Arc::clone(&log_broadcaster)); - // Create CLI channel - let repl_channel = if let Some(ref msg) = cli.message { - Some(ReplChannel::with_message(msg.clone())) - } else if config.channels.cli.enabled { - let repl = ReplChannel::new(); - // Suppress the one-liner banner; boot screen will be shown instead. - repl.suppress_banner(); - Some(repl) - } else { - None - }; - tracing::info!("Starting IronClaw..."); tracing::info!("Loaded configuration for agent: {}", config.agent.name); tracing::info!("LLM backend: {}", config.llm.backend); - // Initialize database backend. - // - // Creates an `Arc` that all consumers share. - // Backend is selected by the `DATABASE_BACKEND` env var / config. - // - // NOTE: For simpler call sites (CLI commands, Memory handler) use the shared - // helper `ironclaw::db::connect_from_config()`. This block is kept inline - // because it also captures backend-specific handles (`pg_pool`, `libsql_db`) - // needed by the secrets store. - #[cfg(feature = "postgres")] - let mut pg_pool: Option = None; - #[cfg(feature = "libsql")] - let mut libsql_db: Option> = None; + // ── Phase 1-5: Build all core components via AppBuilder ──────────── - let db: Option> = if cli.no_db { - tracing::warn!("Running without database connection"); - None - } else { - match config.database.backend { - #[cfg(feature = "libsql")] - ironclaw::config::DatabaseBackend::LibSql => { - use ironclaw::db::Database as _; - use ironclaw::db::libsql::LibSqlBackend; - use secrecy::ExposeSecret as _; + let flags = AppBuilderFlags { no_db: cli.no_db }; + let components = AppBuilder::new( + config, + flags, + toml_path.map(std::path::PathBuf::from), + session.clone(), + Arc::clone(&log_broadcaster), + ) + .build_all() + .await?; - let default_path = ironclaw::config::default_libsql_path(); - let db_path = config - .database - .libsql_path - .as_deref() - .unwrap_or(&default_path); - - let backend = if let Some(ref url) = config.database.libsql_url { - let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| { - anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set") - })?; - LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret()).await? - } else { - LibSqlBackend::new_local(db_path).await? - }; - backend.run_migrations().await?; - tracing::info!("libSQL database connected and migrations applied"); - - // Capture the Database handle for SecretsStore (connection-per-op) - libsql_db = Some(backend.shared_db()); - - Some(Arc::new(backend) as Arc) - } - #[cfg(feature = "postgres")] - _ => { - use ironclaw::db::Database as _; - let pg = ironclaw::db::postgres::PgBackend::new(&config.database) - .await - .map_err(|e| anyhow::anyhow!("{}", e))?; - pg.run_migrations() - .await - .map_err(|e| anyhow::anyhow!("{}", e))?; - tracing::info!("PostgreSQL database connected and migrations applied"); - - pg_pool = Some(pg.pool()); - Some(Arc::new(pg) as Arc) - } - #[cfg(not(feature = "postgres"))] - _ => { - anyhow::bail!( - "No database backend available. Enable 'postgres' or 'libsql' feature." - ); - } - } - }; - - // Post-init operations using the database - if let Some(ref db) = db { - // One-time migration: move disk config files into the DB settings table. - if let Err(e) = ironclaw::bootstrap::migrate_disk_to_db(db.as_ref(), "default").await { - tracing::warn!("Disk-to-DB settings migration failed: {}", e); - } - - // Reload config from DB now that we have a connection. - match Config::from_db_with_toml(db.as_ref(), "default", toml_path).await { - Ok(db_config) => { - config = db_config; - tracing::info!("Configuration reloaded from database"); - } - Err(e) => { - tracing::warn!( - "Failed to reload config from DB, keeping env-based config: {}", - e - ); - } - } - - // Attach DB to session manager so tokens save to DB too - session.attach_store(Arc::clone(db), "default").await; - - // Mark any jobs left in "running" or "creating" state as "interrupted". - // Fire-and-forget housekeeping — no need to block startup. - let db_cleanup = Arc::clone(db); - tokio::spawn(async move { - if let Err(e) = db_cleanup.cleanup_stale_sandbox_jobs().await { - tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e); - } - }); - } - - // Create secrets store early: needed for injecting LLM API keys from encrypted - // storage before creating the LLM provider, and later for MCP auth + WASM channels. - // - // When both `postgres` and `libsql` features are compiled, the runtime-selected - // backend determines which store is created: whichever DB init branch ran will - // have set its handle (pg_pool or libsql_db), and the or_else chain picks it up. - let secrets_store: Option> = - if let Some(master_key) = config.secrets.master_key() { - match SecretsCrypto::new(master_key.clone()) { - Ok(crypto) => { - let crypto = Arc::new(crypto); - let store: Option> = None; - - #[cfg(feature = "libsql")] - let store = store.or_else(|| { - libsql_db.take().map(|db| { - Arc::new(LibSqlSecretsStore::new(db, Arc::clone(&crypto))) - as Arc - }) - }); - - #[cfg(feature = "postgres")] - let store = store.or_else(|| { - pg_pool.as_ref().map(|pool| { - Arc::new(PostgresSecretsStore::new(pool.clone(), Arc::clone(&crypto))) - as Arc - }) - }); - - store - } - Err(e) => { - tracing::warn!("Failed to initialize secrets crypto: {}", e); - #[cfg(feature = "libsql")] - let _ = libsql_db.take(); - None - } - } - } else { - #[cfg(feature = "libsql")] - let _ = libsql_db.take(); - None - }; - - // Inject LLM API keys from the encrypted secrets store into a thread-safe - // overlay so that optional_env() (used by LlmConfig::resolve()) picks them - // up. Then re-resolve LlmConfig with the newly available keys (backend may - // have been set during onboarding but the API key is in the secrets store). - if let Some(ref secrets) = secrets_store { - ironclaw::config::inject_llm_keys_from_secrets(secrets.as_ref(), "default").await; - - // Re-resolve LlmConfig now that secrets overlay has been populated - if let Some(ref db_ref) = db { - match Config::from_db_with_toml(db_ref.as_ref(), "default", toml_path).await { - Ok(refreshed) => { - config = refreshed; - tracing::debug!("LlmConfig re-resolved after secret injection"); - } - Err(e) => { - tracing::warn!("Failed to re-resolve config after secret injection: {}", e); - } - } - } - } + let config = components.config; // Session-based auth is only needed for NEAR AI backend without an API key. - // Do this after DB-backed config reload so provider selection from onboarding - // is respected (e.g. OpenAI/OpenAI-compatible should not trigger NEAR auth). if config.llm.backend == ironclaw::config::LlmBackend::NearAi && config.llm.nearai.api_key.is_none() { session.ensure_authenticated().await?; } - // Start managed tunnel if configured and no static URL is already set. - // - // The tunnel process runs in the background, exposing the local gateway - // port to the internet. The resulting public URL is injected into - // config.tunnel.public_url so channels and extensions pick it up. - let active_tunnel: Option> = - if config.tunnel.public_url.is_some() { - tracing::info!( - "Static tunnel URL in use: {}", - config.tunnel.public_url.as_deref().unwrap_or("?") - ); - None - } else if let Some(ref provider_config) = config.tunnel.provider { - let gateway_port = config - .channels - .gateway - .as_ref() - .map(|g| g.port) - .unwrap_or(3000); - let gateway_host = config - .channels - .gateway - .as_ref() - .map(|g| g.host.as_str()) - .unwrap_or("127.0.0.1"); + // ── Tunnel setup ─────────────────────────────────────────────────── - match ironclaw::tunnel::create_tunnel(provider_config) { - Ok(Some(tunnel)) => { - tracing::info!( - "Starting {} tunnel on {}:{}...", - tunnel.name(), - gateway_host, - gateway_port - ); - match tunnel.start(gateway_host, gateway_port).await { - Ok(url) => { - tracing::info!("Tunnel started: {}", url); - config.tunnel.public_url = Some(url); - Some(tunnel) - } - Err(e) => { - tracing::error!("Failed to start tunnel: {}", e); - None - } - } - } - Ok(None) => None, - Err(e) => { - tracing::error!("Failed to create tunnel: {}", e); - None - } - } - } else { - None - }; + let (config, active_tunnel) = start_tunnel(config).await; - // Build the full LLM provider chain (retry → smart routing → failover → circuit breaker → cache) - let (llm, cheap_llm) = build_provider_chain(&config.llm, session.clone())?; + // ── Orchestrator / container job manager ──────────────────────────── - // Initialize safety layer - let safety = Arc::new(SafetyLayer::new(&config.safety)); - tracing::info!("Safety layer initialized"); - - // Initialize tool registry with credential injection support - let credential_registry = Arc::new(ironclaw::tools::wasm::SharedCredentialRegistry::new()); - let tools = if let Some(ref ss) = secrets_store { - Arc::new( - ToolRegistry::new().with_credentials(Arc::clone(&credential_registry), Arc::clone(ss)), - ) - } else { - Arc::new(ToolRegistry::new()) - }; - tools.register_builtin_tools(); - - // Create embeddings provider if configured - let embeddings: Option> = if config.embeddings.enabled { - match config.embeddings.provider.as_str() { - "nearai" => { - tracing::info!( - "Embeddings enabled via NEAR AI (model: {}, dim: {})", - config.embeddings.model, - config.embeddings.dimension, - ); - Some(Arc::new( - NearAiEmbeddings::new(&config.llm.nearai.base_url, session.clone()) - .with_model(&config.embeddings.model, config.embeddings.dimension), - )) - } - "ollama" => { - tracing::info!( - "Embeddings enabled via Ollama (model: {}, url: {}, dim: {})", - config.embeddings.model, - config.embeddings.ollama_base_url, - config.embeddings.dimension, - ); - Some(Arc::new( - OllamaEmbeddings::new(&config.embeddings.ollama_base_url) - .with_model(&config.embeddings.model, config.embeddings.dimension), - )) - } - _ => { - // Default to OpenAI for unknown providers - if let Some(api_key) = config.embeddings.openai_api_key() { - tracing::info!( - "Embeddings enabled via OpenAI (model: {}, dim: {})", - config.embeddings.model, - config.embeddings.dimension, - ); - Some(Arc::new(OpenAiEmbeddings::with_model( - api_key, - &config.embeddings.model, - config.embeddings.dimension, - ))) - } else { - tracing::warn!("Embeddings configured but OPENAI_API_KEY not set"); - None - } - } - } - } else { - tracing::info!("Embeddings disabled (set OPENAI_API_KEY or EMBEDDING_ENABLED=true)"); - None - }; - - // Warn if libSQL backend is used with non-1536 embedding dimension. - if config.database.backend == ironclaw::config::DatabaseBackend::LibSql - && config.embeddings.enabled - && config.embeddings.dimension != 1536 - { - tracing::warn!( - configured_dimension = config.embeddings.dimension, - "Embedding dimension {} is not 1536. The libSQL schema uses \ - F32_BLOB(1536) which requires exactly 1536 dimensions. \ - Embedding storage will fail. Use PostgreSQL or set \ - EMBEDDING_DIMENSION=1536.", - config.embeddings.dimension - ); - } - - // Create workspace once, reused for memory tools and agent - let workspace: Option> = if let Some(ref db) = db { - let mut ws = Workspace::new_with_db("default", Arc::clone(db)); - if let Some(ref emb) = embeddings { - ws = ws.with_embeddings(emb.clone()); - } - Some(Arc::new(ws)) - } else { - None - }; - - // Register memory tools if workspace is available - if let Some(ref ws) = workspace { - tools.register_memory_tools(Arc::clone(ws)); - } - - // Register builder tool if enabled. - // When sandbox is enabled and allow_local_tools is false, skip builder registration - // because register_builder_tool also registers dev tools (shell, file ops) that would - // bypass the sandbox. The builder runs inside containers instead. - if config.builder.enabled && (config.agent.allow_local_tools || !config.sandbox.enabled) { - tools - .register_builder_tool( - llm.clone(), - safety.clone(), - Some(config.builder.to_builder_config()), - ) - .await; - tracing::info!("Builder mode enabled"); - } - - let mcp_session_manager = Arc::new(McpSessionManager::new()); - - // Create hook registry early so runtime extension activation can register hooks. - let hooks = Arc::new(HookRegistry::new()); - - // Create WASM tool runtime (sync, just builds the wasmtime engine) - let wasm_tool_runtime: Option> = - if config.wasm.enabled && config.wasm.tools_dir.exists() { - match WasmToolRuntime::new(config.wasm.to_runtime_config()) { - Ok(runtime) => Some(Arc::new(runtime)), - Err(e) => { - tracing::warn!("Failed to initialize WASM runtime: {}", e); - None - } - } - } else { - None - }; - - // Load WASM tools and MCP servers concurrently. - // Both register into the shared ToolRegistry (RwLock-based) so concurrent writes are safe. - let wasm_tools_future = async { - let mut dev_loaded_tool_names: Vec = 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)); - } - - // Load installed tools from ~/.ironclaw/tools/ - match loader.load_from_dir(&config.wasm.tools_dir).await { - Ok(results) => { - if !results.loaded.is_empty() { - tracing::info!( - "Loaded {} WASM tools from {}", - results.loaded.len(), - config.wasm.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); - } - } - - // Load dev tools from build artifacts (overrides installed if newer) - match load_dev_tools(&loader, &config.wasm.tools_dir).await { - Ok(results) => { - dev_loaded_tool_names.extend(results.loaded.iter().cloned()); - if !results.loaded.is_empty() { - tracing::info!( - "Loaded {} dev WASM tools from build artifacts", - results.loaded.len() - ); - } - } - Err(e) => { - tracing::debug!("No dev WASM tools found: {}", e); - } - } - } - - dev_loaded_tool_names - }; - - let mcp_servers_future = async { - if let Some(ref secrets) = secrets_store { - let servers_result = if let Some(ref d) = db { - load_mcp_servers_from_db(d.as_ref(), "default").await - } else { - ironclaw::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::info!("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_session_manager); - let secrets = Arc::clone(secrets); - let tools = Arc::clone(&tools); - - join_set.spawn(async move { - let server_name = server.name.clone(); - tracing::debug!( - "Checking authentication for MCP server '{}'...", - server_name - ); - let has_tokens = is_authenticated(&server, &secrets, "default").await; - tracing::debug!( - "MCP server '{}' has_tokens={}", - server_name, - has_tokens - ); - - let client = if has_tokens || server.requires_auth() { - McpClient::new_authenticated(server, mcp_sm, secrets, "default") - } else { - McpClient::new_with_name(&server_name, &server.url) - }; - - tracing::debug!("Fetching tools from MCP server '{}'...", server_name); - match client.list_tools().await { - Ok(mcp_tools) => { - let tool_count = mcp_tools.len(); - tracing::debug!( - "Got {} tools from MCP server '{}'", - tool_count, - server_name - ); - match client.create_tools().await { - Ok(tool_impls) => { - for tool in tool_impls { - tools.register(tool).await; - } - tracing::info!( - "Loaded {} tools from MCP server '{}'", - tool_count, - server_name - ); - } - 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 - ); - } - } - } - }); - } - - while let Some(result) = join_set.join_next().await { - if let Err(e) = result { - tracing::warn!("MCP server loading task panicked: {}", e); - } - } - } - Err(e) => { - tracing::debug!("No MCP servers configured ({})", e); - } - } - } - }; - - let (dev_loaded_tool_names, _) = tokio::join!(wasm_tools_future, mcp_servers_future); - - // Load registry catalog entries for in-chat extension discovery - let catalog_entries = match ironclaw::registry::RegistryCatalog::load_or_embedded() { - Ok(catalog) => { - let entries: Vec = catalog - .all() - .iter() - .map(|m| m.to_registry_entry()) - .collect(); - tracing::info!( - count = entries.len(), - "Loaded registry catalog entries for extension discovery" - ); - entries - } - Err(e) => { - tracing::warn!("Failed to load registry catalog: {}", e); - Vec::new() - } - }; - - // Create extension manager for in-chat discovery/install/auth/activate. - // If no persistent secrets store is available, use an ephemeral in-memory store - // so that listing/installing/activating extensions still works (auth won't persist). - let ext_secrets: Arc = if let Some(ref s) = secrets_store { - Arc::clone(s) - } else { - use ironclaw::secrets::{InMemorySecretsStore, SecretsCrypto}; - let ephemeral_key = - secrecy::SecretString::from(ironclaw::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), - ext_secrets, - Arc::clone(&tools), - Some(Arc::clone(&hooks)), - wasm_tool_runtime.clone(), - config.wasm.tools_dir.clone(), - config.channels.wasm_channels_dir.clone(), - config.tunnel.public_url.clone(), - "default".to_string(), - db.clone(), - catalog_entries.clone(), - )); - tools.register_extension_tools(Arc::clone(&manager)); - tracing::info!("Extension manager initialized with in-chat discovery tools"); - Some(manager) - }; - - // Set up orchestrator for sandboxed job execution - // When allow_local_tools is false (default), the LLM uses create_job for FS/shell work. - // When allow_local_tools is true, dev tools are also registered directly (current behavior). - // 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 = - config.builder.enabled && (config.agent.allow_local_tools || !config.sandbox.enabled); - if config.agent.allow_local_tools && !builder_registered_dev_tools { - tools.register_dev_tools(); - } - - // Shared state for job events (used by both orchestrator and web gateway) let job_event_tx: Option< tokio::sync::broadcast::Sender<(uuid::Uuid, ironclaw::channels::web::types::SseEvent)>, > = if config.sandbox.enabled { @@ -1004,13 +247,13 @@ async fn main() -> anyhow::Result<()> { // Start the orchestrator internal API in the background let orchestrator_state = OrchestratorState { - llm: llm.clone(), + llm: components.llm.clone(), job_manager: Arc::clone(&jm), token_store, job_event_tx: job_event_tx.clone(), prompt_queue: Arc::clone(&prompt_queue), - store: db.clone(), - secrets_store: secrets_store.clone(), + store: components.db.clone(), + secrets_store: components.secrets_store.clone(), user_id: "default".to_string(), }; @@ -1032,16 +275,23 @@ async fn main() -> anyhow::Result<()> { None }; - tracing::info!( - "Tool registry initialized with {} total tools", - tools.count() - ); + // ── Channel setup ────────────────────────────────────────────────── - // Initialize channel manager let mut channels = ChannelManager::new(); let mut channel_names: Vec = Vec::new(); let mut loaded_wasm_channel_names: Vec = Vec::new(); + // Create CLI channel + let repl_channel = if let Some(ref msg) = cli.message { + Some(ReplChannel::with_message(msg.clone())) + } else if config.channels.cli.enabled { + let repl = ReplChannel::new(); + repl.suppress_banner(); + Some(repl) + } else { + None + }; + if let Some(repl) = repl_channel { channels.add(Box::new(repl)); if cli.message.is_some() { @@ -1057,172 +307,26 @@ async fn main() -> anyhow::Result<()> { // Load WASM channels and register their webhook routes. if config.channels.wasm_channels_enabled && config.channels.wasm_channels_dir.exists() { - match WasmChannelRuntime::new(WasmChannelRuntimeConfig::default()) { - Ok(runtime) => { - let runtime = Arc::new(runtime); - let pairing_store = Arc::new(PairingStore::new()); - let loader = WasmChannelLoader::new(Arc::clone(&runtime), pairing_store); + let wasm_result = setup_wasm_channels( + &config, + &components.secrets_store, + components.extension_manager.as_ref(), + ) + .await; - match loader - .load_from_dir(&config.channels.wasm_channels_dir) - .await - { - Ok(results) => { - let wasm_router = Arc::new(WasmChannelRouter::new()); - let mut has_webhook_channels = false; - - for loaded in results.loaded { - let channel_name = loaded.name().to_string(); - loaded_wasm_channel_names.push(channel_name.clone()); - tracing::info!("Loaded WASM channel: {}", channel_name); - - let secret_name = loaded.webhook_secret_name(); - - let webhook_secret = if let Some(ref secrets) = secrets_store { - secrets - .get_decrypted("default", &secret_name) - .await - .ok() - .map(|s| s.expose().to_string()) - } else { - None - }; - - let secret_header = - loaded.webhook_secret_header().map(|s| s.to_string()); - - let webhook_path = format!("/webhook/{}", channel_name); - let endpoints = vec![RegisteredEndpoint { - channel_name: channel_name.clone(), - path: webhook_path.clone(), - methods: vec!["POST".to_string()], - require_secret: webhook_secret.is_some(), - }]; - - let channel_arc = Arc::new(loaded.channel); - - { - let mut config_updates = std::collections::HashMap::new(); - - if let Some(ref tunnel_url) = config.tunnel.public_url { - config_updates.insert( - "tunnel_url".to_string(), - serde_json::Value::String(tunnel_url.clone()), - ); - } - - if let Some(ref secret) = webhook_secret { - config_updates.insert( - "webhook_secret".to_string(), - serde_json::Value::String(secret.clone()), - ); - } - - // Inject owner_id for Telegram so the bot only responds - // to the bound user account. - if channel_name == "telegram" - && let Some(owner_id) = config.channels.telegram_owner_id - { - config_updates.insert( - "owner_id".to_string(), - serde_json::json!(owner_id), - ); - } - - if !config_updates.is_empty() { - channel_arc.update_config(config_updates).await; - tracing::info!( - channel = %channel_name, - has_tunnel = config.tunnel.public_url.is_some(), - has_webhook_secret = webhook_secret.is_some(), - "Injected runtime config into channel" - ); - } - } - - tracing::info!( - channel = %channel_name, - has_webhook_secret = webhook_secret.is_some(), - secret_header = ?secret_header, - "Registering channel with router" - ); - - wasm_router - .register( - Arc::clone(&channel_arc), - endpoints, - webhook_secret.clone(), - secret_header, - ) - .await; - has_webhook_channels = true; - - if let Some(ref secrets) = secrets_store { - match inject_channel_credentials( - &channel_arc, - secrets.as_ref(), - &channel_name, - ) - .await - { - Ok(count) => { - if count > 0 { - tracing::info!( - channel = %channel_name, - credentials_injected = count, - "Channel credentials injected" - ); - } - } - Err(e) => { - tracing::error!( - channel = %channel_name, - error = %e, - "Failed to inject channel credentials" - ); - } - } - } - - channel_names.push(channel_name.clone()); - channels.add(Box::new(SharedWasmChannel::new(channel_arc))); - } - - if has_webhook_channels { - webhook_routes.push(create_wasm_channel_router( - wasm_router, - extension_manager.as_ref().map(Arc::clone), - )); - } - - // Tell extension manager which channels are actually loaded - if let Some(ref em) = extension_manager { - em.set_active_channels(loaded_wasm_channel_names.clone()) - .await; - } - - for (path, err) in &results.errors { - tracing::warn!( - "Failed to load WASM channel {}: {}", - path.display(), - err - ); - } - } - Err(e) => { - tracing::warn!("Failed to scan WASM channels directory: {}", e); - } - } + if let Some(result) = wasm_result { + loaded_wasm_channel_names = result.channel_names; + for (name, channel) in result.channels { + channel_names.push(name); + channels.add(channel); } - Err(e) => { - tracing::warn!("Failed to initialize WASM channel runtime: {}", e); + if let Some(routes) = result.webhook_routes { + webhook_routes.push(routes); } } } // Add HTTP channel if configured and not CLI-only mode. - // Extract its routes for the unified server; the channel itself just - // provides the mpsc stream. let mut webhook_server_addr: Option = None; if !cli.cli_only && let Some(ref http_config) = config.channels.http @@ -1265,46 +369,17 @@ async fn main() -> anyhow::Result<()> { None }; - // Seed workspace with core identity files on first boot - if let Some(ref ws) = workspace { - match ws.seed_if_empty().await { - Ok(_) => {} - Err(e) => { - tracing::warn!("Failed to seed workspace: {}", e); - } - } - } - - // Backfill embeddings in background (fire-and-forget housekeeping) - if let (Some(ws), Some(_)) = (&workspace, &embeddings) { - let ws_bg = Arc::clone(ws); - tokio::spawn(async move { - match ws_bg.backfill_embeddings().await { - Ok(count) if count > 0 => { - tracing::info!("Backfilled embeddings for {} chunks", count); - } - Ok(_) => {} - Err(e) => { - tracing::warn!("Failed to backfill embeddings: {}", e); - } - } - }); - } - - // Create context manager (shared between job tools and agent) - let context_manager = Arc::new(ContextManager::new(config.agent.max_parallel_jobs)); - - // Register bundled/plugin/workspace hooks. - let active_tool_names = tools.list().await; + // Register lifecycle hooks. + let active_tool_names = components.tools.list().await; let hook_bootstrap = bootstrap_hooks( - &hooks, - workspace.as_ref(), + &components.hooks, + components.workspace.as_ref(), &config.wasm.tools_dir, &config.channels.wasm_channels_dir, &active_tool_names, &loaded_wasm_channel_names, - &dev_loaded_tool_names, + &components.dev_loaded_tool_names, ) .await; tracing::info!( @@ -1317,13 +392,14 @@ async fn main() -> anyhow::Result<()> { ); // Create session manager (shared between agent and web gateway) - let session_manager = Arc::new(SessionManager::new().with_hooks(hooks.clone())); + let session_manager = + Arc::new(ironclaw::agent::SessionManager::new().with_hooks(components.hooks.clone())); // Register job tools (sandbox deps auto-injected when container_job_manager is available) - tools.register_job_tools( - Arc::clone(&context_manager), + components.tools.register_job_tools( + Arc::clone(&components.context_manager), container_job_manager.clone(), - db.clone(), + components.db.clone(), job_event_tx.clone(), Some(channels.inject_sender()), if config.sandbox.enabled { @@ -1331,69 +407,44 @@ async fn main() -> anyhow::Result<()> { } else { None }, - secrets_store.clone(), + components.secrets_store.clone(), ); - // Initialize skills system (before gateway so we can wire into GatewayState) - let (skill_registry, skill_catalog) = if config.skills.enabled { - let mut registry = ironclaw::skills::SkillRegistry::new(config.skills.local_dir.clone()); - let loaded = registry.discover_all().await; - if !loaded.is_empty() { - tracing::info!("Loaded {} skill(s): {}", loaded.len(), loaded.join(", ")); - } - let registry = Arc::new(std::sync::RwLock::new(registry)); + // ── Gateway channel ──────────────────────────────────────────────── - // Register skill management tools - let catalog = ironclaw::skills::catalog::shared_catalog(); - tools.register_skill_tools(Arc::clone(®istry), Arc::clone(&catalog)); - - (Some(registry), Some(catalog)) - } else { - (None, None) - }; - - // Create cost guard early so gateway can reference it. - let cost_guard = Arc::new(ironclaw::agent::cost_guard::CostGuard::new( - ironclaw::agent::cost_guard::CostGuardConfig { - max_cost_per_day_cents: config.agent.max_cost_per_day_cents, - max_actions_per_hour: config.agent.max_actions_per_hour, - }, - )); - - // Add web gateway channel if configured let mut gateway_url: Option = None; if let Some(ref gw_config) = config.channels.gateway { - let mut gw = GatewayChannel::new(gw_config.clone()).with_llm_provider(Arc::clone(&llm)); - if let Some(ref ws) = workspace { + let mut gw = + GatewayChannel::new(gw_config.clone()).with_llm_provider(Arc::clone(&components.llm)); + if let Some(ref ws) = components.workspace { gw = gw.with_workspace(Arc::clone(ws)); } gw = gw.with_session_manager(Arc::clone(&session_manager)); gw = gw.with_log_broadcaster(Arc::clone(&log_broadcaster)); gw = gw.with_log_level_handle(Arc::clone(&log_level_handle)); - gw = gw.with_tool_registry(Arc::clone(&tools)); - if let Some(ref ext_mgr) = extension_manager { + gw = gw.with_tool_registry(Arc::clone(&components.tools)); + if let Some(ref ext_mgr) = components.extension_manager { gw = gw.with_extension_manager(Arc::clone(ext_mgr)); } - if !catalog_entries.is_empty() { - gw = gw.with_registry_entries(catalog_entries.clone()); + if !components.catalog_entries.is_empty() { + gw = gw.with_registry_entries(components.catalog_entries.clone()); } - if let Some(ref d) = db { + if let Some(ref d) = components.db { gw = gw.with_store(Arc::clone(d)); } if let Some(ref jm) = container_job_manager { gw = gw.with_job_manager(Arc::clone(jm)); } - if let Some(ref sr) = skill_registry { + if let Some(ref sr) = components.skill_registry { gw = gw.with_skill_registry(Arc::clone(sr)); } - if let Some(ref sc) = skill_catalog { + if let Some(ref sc) = components.skill_catalog { gw = gw.with_skill_catalog(Arc::clone(sc)); } - gw = gw.with_cost_guard(Arc::clone(&cost_guard)); + gw = gw.with_cost_guard(Arc::clone(&components.cost_guard)); if config.sandbox.enabled { gw = gw.with_prompt_queue(Arc::clone(&prompt_queue)); - // Spawn a task to forward job events from the broadcast channel to SSE if let Some(ref tx) = job_event_tx { let mut rx = tx.subscribe(); let gw_state = Arc::clone(gw.state()); @@ -1418,37 +469,15 @@ async fn main() -> anyhow::Result<()> { channels.add(Box::new(gw)); } - // Capture boot screen info before moving Arcs into AgentDeps. - let boot_tool_count = tools.count(); - let boot_llm_model = llm.model_name().to_string(); - let boot_cheap_model = cheap_llm.as_ref().map(|c| c.model_name().to_string()); + // ── Boot screen ──────────────────────────────────────────────────── - // Create and run the agent - let deps = AgentDeps { - store: db, - llm, - cheap_llm, - safety, - tools, - workspace, - extension_manager, - skill_registry, - skills_config: config.skills.clone(), - hooks, - cost_guard, - }; - let agent = Agent::new( - config.agent.clone(), - deps, - channels, - Some(config.heartbeat.clone()), - Some(config.hygiene.clone()), - Some(config.routines.clone()), - Some(context_manager), - Some(session_manager), - ); + let boot_tool_count = components.tools.count(); + let boot_llm_model = components.llm.model_name().to_string(); + let boot_cheap_model = components + .cheap_llm + .as_ref() + .map(|c| c.model_name().to_string()); - // Print boot screen for interactive CLI mode (not single-message mode). if config.channels.cli.enabled && cli.message.is_none() { let boot_info = ironclaw::boot_screen::BootInfo { version: env!("CARGO_PKG_VERSION").to_string(), @@ -1485,15 +514,40 @@ async fn main() -> anyhow::Result<()> { ironclaw::boot_screen::print_boot_screen(&boot_info); } - // Run the agent (blocks until shutdown) + // ── Run the agent ────────────────────────────────────────────────── + + let deps = AgentDeps { + store: components.db, + llm: components.llm, + cheap_llm: components.cheap_llm, + safety: components.safety, + tools: components.tools, + workspace: components.workspace, + extension_manager: components.extension_manager, + skill_registry: components.skill_registry, + skills_config: config.skills.clone(), + hooks: components.hooks, + cost_guard: components.cost_guard, + }; + let agent = Agent::new( + config.agent.clone(), + deps, + channels, + Some(config.heartbeat.clone()), + Some(config.hygiene.clone()), + Some(config.routines.clone()), + Some(components.context_manager), + Some(session_manager), + ); + agent.run().await?; - // Shut down the webhook server if one was started + // ── Shutdown ──────────────────────────────────────────────────────── + if let Some(ref mut server) = webhook_server { server.shutdown().await; } - // Stop managed tunnel if one was started if let Some(tunnel) = active_tunnel { tracing::info!("Stopping {} tunnel...", tunnel.name()); if let Err(e) = tunnel.stop().await { @@ -1505,11 +559,341 @@ async fn main() -> anyhow::Result<()> { Ok(()) } +// ── Helper functions ──────────────────────────────────────────────────── + +/// Initialize tracing for worker/bridge processes (info level). +fn init_worker_tracing() { + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("ironclaw=info")), + ) + .init(); +} + +/// Run the Memory CLI subcommand. +async fn run_memory_command(mem_cmd: &ironclaw::cli::MemoryCommand) -> anyhow::Result<()> { + let config = Config::from_env() + .await + .map_err(|e| anyhow::anyhow!("{}", e))?; + + let session = create_session_manager(SessionConfig { + auth_base_url: config.llm.nearai.auth_base_url.clone(), + session_path: config.llm.nearai.session_path.clone(), + }) + .await; + + let embeddings = config + .embeddings + .create_provider(&config.llm.nearai.base_url, session); + + // Warn if libSQL backend is used with non-1536 embedding dimension. + if config.database.backend == ironclaw::config::DatabaseBackend::LibSql + && config.embeddings.enabled + && config.embeddings.dimension != 1536 + { + tracing::warn!( + configured_dimension = config.embeddings.dimension, + "Embedding dimension {} is not 1536. The libSQL schema uses \ + F32_BLOB(1536) which requires exactly 1536 dimensions. \ + Embedding storage will fail. Use PostgreSQL or set \ + EMBEDDING_DIMENSION=1536.", + config.embeddings.dimension + ); + } + + let db: Arc = ironclaw::db::connect_from_config(&config.database) + .await + .map_err(|e| anyhow::anyhow!("{}", e))?; + + ironclaw::cli::run_memory_command_with_db(mem_cmd.clone(), db, embeddings).await +} + +/// Run the Worker subcommand (inside Docker containers). +async fn run_worker( + job_id: uuid::Uuid, + orchestrator_url: &str, + max_iterations: u32, +) -> anyhow::Result<()> { + tracing::info!( + "Starting worker for job {} (orchestrator: {})", + job_id, + orchestrator_url + ); + + let config = ironclaw::worker::runtime::WorkerConfig { + job_id, + orchestrator_url: orchestrator_url.to_string(), + max_iterations, + timeout: std::time::Duration::from_secs(600), + }; + + let runtime = ironclaw::worker::WorkerRuntime::new(config) + .map_err(|e| anyhow::anyhow!("Worker init failed: {}", e))?; + + runtime + .run() + .await + .map_err(|e| anyhow::anyhow!("Worker failed: {}", e)) +} + +/// Run the Claude Code bridge subcommand (inside Docker containers). +async fn run_claude_bridge( + job_id: uuid::Uuid, + orchestrator_url: &str, + max_turns: u32, + model: &str, +) -> anyhow::Result<()> { + tracing::info!( + "Starting Claude Code bridge for job {} (orchestrator: {}, model: {})", + job_id, + orchestrator_url, + model + ); + + let config = ironclaw::worker::claude_bridge::ClaudeBridgeConfig { + job_id, + orchestrator_url: orchestrator_url.to_string(), + max_turns, + model: model.to_string(), + timeout: std::time::Duration::from_secs(1800), + allowed_tools: ironclaw::config::ClaudeCodeConfig::from_env().allowed_tools, + }; + + let runtime = ironclaw::worker::ClaudeBridgeRuntime::new(config) + .map_err(|e| anyhow::anyhow!("Claude bridge init failed: {}", e))?; + + runtime + .run() + .await + .map_err(|e| anyhow::anyhow!("Claude bridge failed: {}", e)) +} + +/// Start managed tunnel if configured and no static URL is already set. +async fn start_tunnel( + mut config: ironclaw::config::Config, +) -> ( + ironclaw::config::Config, + Option>, +) { + if config.tunnel.public_url.is_some() { + tracing::info!( + "Static tunnel URL in use: {}", + config.tunnel.public_url.as_deref().unwrap_or("?") + ); + return (config, None); + } + + let Some(ref provider_config) = config.tunnel.provider else { + return (config, None); + }; + + let gateway_port = config + .channels + .gateway + .as_ref() + .map(|g| g.port) + .unwrap_or(3000); + let gateway_host = config + .channels + .gateway + .as_ref() + .map(|g| g.host.as_str()) + .unwrap_or("127.0.0.1"); + + match ironclaw::tunnel::create_tunnel(provider_config) { + Ok(Some(tunnel)) => { + tracing::info!( + "Starting {} tunnel on {}:{}...", + tunnel.name(), + gateway_host, + gateway_port + ); + match tunnel.start(gateway_host, gateway_port).await { + Ok(url) => { + tracing::info!("Tunnel started: {}", url); + config.tunnel.public_url = Some(url); + (config, Some(tunnel)) + } + Err(e) => { + tracing::error!("Failed to start tunnel: {}", e); + (config, None) + } + } + } + Ok(None) => (config, None), + Err(e) => { + tracing::error!("Failed to create tunnel: {}", e); + (config, None) + } + } +} + +/// Result of WASM channel setup. +struct WasmChannelSetup { + channels: Vec<(String, Box)>, + channel_names: Vec, + webhook_routes: Option, +} + +/// Load WASM channels and register their webhook routes. +async fn setup_wasm_channels( + config: &ironclaw::config::Config, + secrets_store: &Option>, + extension_manager: Option<&Arc>, +) -> Option { + let runtime = match WasmChannelRuntime::new(WasmChannelRuntimeConfig::default()) { + Ok(r) => Arc::new(r), + Err(e) => { + tracing::warn!("Failed to initialize WASM channel runtime: {}", e); + return None; + } + }; + + let pairing_store = Arc::new(PairingStore::new()); + let loader = WasmChannelLoader::new(Arc::clone(&runtime), pairing_store); + + let results = match loader + .load_from_dir(&config.channels.wasm_channels_dir) + .await + { + Ok(r) => r, + Err(e) => { + tracing::warn!("Failed to scan WASM channels directory: {}", e); + return None; + } + }; + + let wasm_router = Arc::new(WasmChannelRouter::new()); + let mut has_webhook_channels = false; + let mut channels: Vec<(String, Box)> = Vec::new(); + let mut channel_names: Vec = Vec::new(); + + for loaded in results.loaded { + let channel_name = loaded.name().to_string(); + channel_names.push(channel_name.clone()); + tracing::info!("Loaded WASM channel: {}", channel_name); + + let secret_name = loaded.webhook_secret_name(); + + let webhook_secret = if let Some(secrets) = secrets_store { + secrets + .get_decrypted("default", &secret_name) + .await + .ok() + .map(|s| s.expose().to_string()) + } else { + None + }; + + let secret_header = loaded.webhook_secret_header().map(|s| s.to_string()); + + let webhook_path = format!("/webhook/{}", channel_name); + let endpoints = vec![RegisteredEndpoint { + channel_name: channel_name.clone(), + path: webhook_path, + methods: vec!["POST".to_string()], + require_secret: webhook_secret.is_some(), + }]; + + let channel_arc = Arc::new(loaded.channel); + + { + let mut config_updates = std::collections::HashMap::new(); + + if let Some(ref tunnel_url) = config.tunnel.public_url { + config_updates.insert( + "tunnel_url".to_string(), + serde_json::Value::String(tunnel_url.clone()), + ); + } + + if let Some(ref secret) = webhook_secret { + config_updates.insert( + "webhook_secret".to_string(), + serde_json::Value::String(secret.clone()), + ); + } + + // Inject owner_id for Telegram so the bot only responds to the bound user. + if channel_name == "telegram" + && let Some(owner_id) = config.channels.telegram_owner_id + { + config_updates.insert("owner_id".to_string(), serde_json::json!(owner_id)); + } + + if !config_updates.is_empty() { + channel_arc.update_config(config_updates).await; + tracing::info!( + channel = %channel_name, + has_tunnel = config.tunnel.public_url.is_some(), + has_webhook_secret = webhook_secret.is_some(), + "Injected runtime config into channel" + ); + } + } + + tracing::info!( + channel = %channel_name, + has_webhook_secret = webhook_secret.is_some(), + secret_header = ?secret_header, + "Registering channel with router" + ); + + wasm_router + .register( + Arc::clone(&channel_arc), + endpoints, + webhook_secret.clone(), + secret_header, + ) + .await; + has_webhook_channels = true; + + if let Some(secrets) = secrets_store { + match inject_channel_credentials(&channel_arc, secrets.as_ref(), &channel_name).await { + Ok(count) => { + if count > 0 { + tracing::info!( + channel = %channel_name, + credentials_injected = count, + "Channel credentials injected" + ); + } + } + Err(e) => { + tracing::error!( + channel = %channel_name, + error = %e, + "Failed to inject channel credentials" + ); + } + } + } + + channels.push((channel_name, Box::new(SharedWasmChannel::new(channel_arc)))); + } + + for (path, err) in &results.errors { + tracing::warn!("Failed to load WASM channel {}: {}", path.display(), err); + } + + let webhook_routes = if has_webhook_channels { + Some(create_wasm_channel_router( + wasm_router, + extension_manager.map(Arc::clone), + )) + } else { + None + }; + + Some(WasmChannelSetup { + channels, + channel_names, + webhook_routes, + }) +} + /// Check if onboarding is needed and return the reason. -/// -/// Returns `Some(reason)` if onboarding should be triggered, `None` otherwise. -/// Called after `load_ironclaw_env()`, so DATABASE_URL from `~/.ironclaw/.env` -/// is already in the environment. #[cfg(any(feature = "postgres", feature = "libsql"))] fn check_onboard_needed() -> Option<&'static str> { let has_db = std::env::var("DATABASE_URL").is_ok() @@ -1520,8 +904,6 @@ fn check_onboard_needed() -> Option<&'static str> { return Some("Database not configured"); } - // The wizard writes ONBOARD_COMPLETED=true to ~/.ironclaw/.env, - // which load_ironclaw_env() loads before this function runs. if std::env::var("ONBOARD_COMPLETED") .map(|v| v == "true") .unwrap_or(false) @@ -1529,9 +911,6 @@ fn check_onboard_needed() -> Option<&'static str> { return None; } - // First run (onboarding never completed and no session). - // Check for a NEAR AI API key or session file as a fallback - // for users who configured credentials manually (no wizard). if std::env::var("NEARAI_API_KEY").is_err() { let session_path = ironclaw::llm::session::default_session_path(); if !session_path.exists() { @@ -1546,14 +925,11 @@ fn check_onboard_needed() -> Option<&'static str> { /// /// Looks for secrets matching the pattern `{channel_name}_*` and injects them /// as credential placeholders (e.g., `telegram_bot_token` -> `{TELEGRAM_BOT_TOKEN}`). -/// -/// Returns the number of credentials injected. async fn inject_channel_credentials( channel: &Arc, secrets: &dyn SecretsStore, channel_name: &str, ) -> anyhow::Result { - // List all secrets for this user and filter by channel prefix let all_secrets = secrets .list("default") .await @@ -1563,12 +939,10 @@ async fn inject_channel_credentials( let mut count = 0; for secret_meta in all_secrets { - // Only process secrets matching the channel prefix if !secret_meta.name.starts_with(&prefix) { continue; } - // Get the decrypted value let decrypted = match secrets.get_decrypted("default", &secret_meta.name).await { Ok(d) => d, Err(e) => { @@ -1581,7 +955,6 @@ async fn inject_channel_credentials( } }; - // Convert secret name to placeholder format (SCREAMING_SNAKE_CASE) let placeholder = secret_meta.name.to_uppercase(); tracing::debug!( diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 34f82915..2000e524 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -1084,9 +1084,8 @@ impl SetupWizard { let fetched = self.fetch_nearai_models().await; let default_models: Vec<(String, String)> = vec![ ( - "fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic" - .into(), - "Llama 4 Maverick (default, fast)".into(), + "zai-org/GLM-latest".into(), + "GLM Latest (default, fast)".into(), ), ( "anthropic::claude-sonnet-4-20250514".into(),