From ffb1cc9be8712106213f0837025f4cd6f31e0203 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Wed, 18 Feb 2026 15:05:47 -0800 Subject: [PATCH] refactor: architecture improvements for contributor velocity (#198) * refactor: split large files and consolidate test stubs for contributor velocity - Extract 7 Database sub-traits (ConversationStore, JobStore, SandboxStore, RoutineStore, ToolFailureStore, SettingsStore, WorkspaceStore) with Database as a supertrait combining them all - Split libsql_backend.rs (2769 lines) into src/db/libsql/ directory with one file per sub-trait implementation - Split config.rs (1753 lines) into src/config/ directory with 16 domain files - Consolidate 3 duplicate test LLM stubs into shared StubLlm in src/testing.rs - Split server.rs handlers into src/channels/web/handlers/ directory - Extract main.rs init phases into AppBuilder (src/app.rs) - Add developer setup script (scripts/dev-setup.sh) Co-Authored-By: Claude Opus 4.6 * refactor: move heartbeat test from examples/ to tests/ Convert standalone example binary into a proper #[ignore] integration test, matching the convention of the other integration tests. Co-Authored-By: Claude Opus 4.6 * style: fix rustfmt formatting for CI Co-Authored-By: Claude Opus 4.6 * fix: address PR review comments from Copilot - tunnel.rs: replace .ok().flatten() with ? to propagate env var errors - secrets.rs: remove misleading "process-wide cache" comment - database.rs: use uppercase "DATABASE_URL" in error key - testing.rs: gate harness tests with #[cfg(feature = "libsql")] Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Illia Polosukhin Co-authored-by: Claude Opus 4.6 --- Cargo.toml | 4 - scripts/dev-setup.sh | 56 + src/app.rs | 779 +++++ src/channels/web/handlers/chat.rs | 633 ++++ src/channels/web/handlers/extensions.rs | 153 + src/channels/web/handlers/jobs.rs | 518 +++ src/channels/web/handlers/memory.rs | 171 + src/channels/web/handlers/mod.rs | 23 + src/channels/web/handlers/routines.rs | 330 ++ src/channels/web/handlers/settings.rs | 133 + src/channels/web/handlers/skills.rs | 257 ++ src/channels/web/handlers/static_files.rs | 178 ++ src/cli/mcp.rs | 2 +- src/cli/tool.rs | 2 +- src/config.rs | 1944 ------------ src/config/agent.rs | 120 + src/config/builder.rs | 72 + src/config/channels.rs | 126 + src/config/database.rs | 130 + src/config/embeddings.rs | 165 + src/config/heartbeat.rs | 54 + src/config/helpers.rs | 40 + src/config/llm.rs | 426 +++ src/config/mod.rs | 239 ++ src/config/routines.rs | 48 + src/config/safety.rs | 25 + src/config/sandbox.rs | 261 ++ src/config/secrets.rs | 70 + src/config/skills.rs | 56 + src/config/tunnel.rs | 106 + src/config/wasm.rs | 99 + src/db/libsql/conversations.rs | 354 +++ src/db/libsql/jobs.rs | 330 ++ src/db/libsql/mod.rs | 460 +++ src/db/libsql/routines.rs | 390 +++ src/db/libsql/sandbox.rs | 405 +++ src/db/libsql/settings.rs | 208 ++ src/db/libsql/tool_failures.rs | 97 + src/db/libsql/workspace.rs | 607 ++++ src/db/libsql_backend.rs | 2769 ----------------- src/db/mod.rs | 234 +- src/db/postgres.rs | 56 +- src/lib.rs | 4 + src/llm/circuit_breaker.rs | 138 +- src/llm/response_cache.rs | 90 +- src/main.rs | 2 +- src/orchestrator/api.rs | 43 +- src/setup/wizard.rs | 6 +- src/testing.rs | 356 +++ .../heartbeat_integration.rs | 31 +- 50 files changed, 8616 insertions(+), 5184 deletions(-) create mode 100755 scripts/dev-setup.sh create mode 100644 src/app.rs create mode 100644 src/channels/web/handlers/chat.rs create mode 100644 src/channels/web/handlers/extensions.rs create mode 100644 src/channels/web/handlers/jobs.rs create mode 100644 src/channels/web/handlers/memory.rs create mode 100644 src/channels/web/handlers/mod.rs create mode 100644 src/channels/web/handlers/routines.rs create mode 100644 src/channels/web/handlers/settings.rs create mode 100644 src/channels/web/handlers/skills.rs create mode 100644 src/channels/web/handlers/static_files.rs delete mode 100644 src/config.rs create mode 100644 src/config/agent.rs create mode 100644 src/config/builder.rs create mode 100644 src/config/channels.rs create mode 100644 src/config/database.rs create mode 100644 src/config/embeddings.rs create mode 100644 src/config/heartbeat.rs create mode 100644 src/config/helpers.rs create mode 100644 src/config/llm.rs create mode 100644 src/config/mod.rs create mode 100644 src/config/routines.rs create mode 100644 src/config/safety.rs create mode 100644 src/config/sandbox.rs create mode 100644 src/config/secrets.rs create mode 100644 src/config/skills.rs create mode 100644 src/config/tunnel.rs create mode 100644 src/config/wasm.rs create mode 100644 src/db/libsql/conversations.rs create mode 100644 src/db/libsql/jobs.rs create mode 100644 src/db/libsql/mod.rs create mode 100644 src/db/libsql/routines.rs create mode 100644 src/db/libsql/sandbox.rs create mode 100644 src/db/libsql/settings.rs create mode 100644 src/db/libsql/tool_failures.rs create mode 100644 src/db/libsql/workspace.rs delete mode 100644 src/db/libsql_backend.rs create mode 100644 src/testing.rs rename examples/test_heartbeat.rs => tests/heartbeat_integration.rs (84%) diff --git a/Cargo.toml b/Cargo.toml index e25b0e5e..2deef38a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -164,10 +164,6 @@ postgres = [ libsql = ["dep:libsql"] integration = [] -[[example]] -name = "test_heartbeat" -required-features = ["postgres"] - # The profile that 'cargo dist' will build with [profile.dist] inherits = "release" diff --git a/scripts/dev-setup.sh b/scripts/dev-setup.sh new file mode 100755 index 00000000..d052c9d1 --- /dev/null +++ b/scripts/dev-setup.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Developer setup script for IronClaw. +# +# Gets a fresh checkout ready for development without requiring +# Docker, PostgreSQL, or any external services. +# +# Usage: +# ./scripts/dev-setup.sh +# +# After running, you can: +# cargo check # default features (postgres + libsql) +# cargo test # default test suite (uses libsql temp DB) +# cargo test --all-features # full test suite + +set -euo pipefail + +cd "$(dirname "$0")/.." + +echo "=== IronClaw Developer Setup ===" +echo "" + +# 1. Check rustup +if ! command -v rustup &>/dev/null; then + echo "ERROR: rustup not found. Install from https://rustup.rs" + exit 1 +fi +echo "[1/5] rustup found: $(rustup --version 2>/dev/null | head -1)" + +# 2. Add WASM target (required by build.rs for channel compilation) +echo "[2/5] Adding wasm32-wasip2 target..." +rustup target add wasm32-wasip2 + +# 3. Install wasm-tools (required by build.rs for WASM component model) +echo "[3/5] Installing wasm-tools..." +if command -v wasm-tools &>/dev/null; then + echo " wasm-tools already installed: $(wasm-tools --version)" +else + cargo install wasm-tools --locked +fi + +# 4. Verify the project compiles +echo "[4/5] Running cargo check..." +cargo check + +# 5. Run tests using libsql temp DB (no Docker/external DB needed) +echo "[5/5] Running tests (no external DB required)..." +cargo test + +echo "" +echo "=== Setup complete ===" +echo "" +echo "Quick start:" +echo " cargo run # Run with default features" +echo " cargo test # Test suite (libsql temp DB)" +echo " cargo test --all-features # Full test suite" +echo " cargo clippy --all-features # Lint all code" diff --git a/src/app.rs b/src/app.rs new file mode 100644 index 00000000..0209fb75 --- /dev/null +++ b/src/app.rs @@ -0,0 +1,779 @@ +//! Application builder for initializing core IronClaw components. +//! +//! Extracts the mechanical initialization phases from `main.rs` into a +//! reusable builder so that: +//! +//! - Tests can construct a full `AppComponents` without wiring channels +//! - Main stays focused on CLI dispatch and channel setup +//! - Each init phase is independently testable + +use std::sync::Arc; + +use crate::channels::web::log_layer::LogBroadcaster; +use crate::config::Config; +use crate::context::ContextManager; +use crate::db::Database; +use crate::extensions::ExtensionManager; +use crate::hooks::HookRegistry; +use crate::llm::{LlmProvider, SessionManager}; +use crate::safety::SafetyLayer; +use crate::secrets::SecretsStore; +use crate::skills::SkillRegistry; +use crate::skills::catalog::SkillCatalog; +use crate::tools::ToolRegistry; +use crate::tools::mcp::McpSessionManager; +use crate::tools::wasm::WasmToolRuntime; +use crate::workspace::{EmbeddingProvider, Workspace}; + +/// Fully initialized application components, ready for channel wiring +/// and agent construction. +pub struct AppComponents { + /// The (potentially mutated) config after DB reload and secret injection. + pub config: Config, + pub db: Option>, + pub secrets_store: Option>, + pub llm: Arc, + pub cheap_llm: Option>, + pub safety: Arc, + pub tools: Arc, + pub embeddings: Option>, + pub workspace: Option>, + pub extension_manager: Option>, + pub mcp_session_manager: Arc, + pub wasm_tool_runtime: Option>, + pub log_broadcaster: Arc, + pub context_manager: Arc, + pub hooks: Arc, + pub skill_registry: Option>>, + pub skill_catalog: Option>, + pub cost_guard: Arc, + pub session: Arc, +} + +/// Options that control optional init phases. +#[derive(Default)] +pub struct AppBuilderFlags { + pub no_db: bool, +} + +/// Builder that orchestrates the 5 mechanical init phases. +pub struct AppBuilder { + config: Config, + flags: AppBuilderFlags, + toml_path: Option, + session: Arc, + log_broadcaster: Arc, + + // Accumulated state + db: Option>, + secrets_store: Option>, + + // Backend-specific handles needed by secrets store + #[cfg(feature = "postgres")] + pg_pool: Option, + #[cfg(feature = "libsql")] + libsql_db: Option>, +} + +impl AppBuilder { + /// Create a new builder. + /// + /// The `session` and `log_broadcaster` are created before the builder + /// because tracing must be initialized before any init phase runs, + /// and the log broadcaster is part of the tracing layer. + pub fn new( + config: Config, + flags: AppBuilderFlags, + toml_path: Option, + session: Arc, + log_broadcaster: Arc, + ) -> Self { + Self { + config, + flags, + toml_path, + session, + log_broadcaster, + db: None, + secrets_store: None, + #[cfg(feature = "postgres")] + pg_pool: None, + #[cfg(feature = "libsql")] + libsql_db: None, + } + } + + /// Phase 1: Initialize database backend. + /// + /// Creates the database connection, runs migrations, reloads config + /// from DB, attaches DB to session manager, and cleans up stale jobs. + pub async fn init_database(&mut self) -> Result<(), anyhow::Error> { + if self.flags.no_db { + tracing::warn!("Running without database connection"); + return Ok(()); + } + + let db: Arc = match self.config.database.backend { + #[cfg(feature = "libsql")] + crate::config::DatabaseBackend::LibSql => { + use crate::db::Database as _; + use crate::db::libsql::LibSqlBackend; + use secrecy::ExposeSecret as _; + + let default_path = crate::config::default_libsql_path(); + let db_path = self + .config + .database + .libsql_path + .as_deref() + .unwrap_or(&default_path); + + let backend = if let Some(ref url) = self.config.database.libsql_url { + let token = + self.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"); + + #[cfg(feature = "libsql")] + { + self.libsql_db = Some(backend.shared_db()); + } + + Arc::new(backend) as Arc + } + #[cfg(feature = "postgres")] + _ => { + use crate::db::Database as _; + let pg = crate::db::postgres::PgBackend::new(&self.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"); + + #[cfg(feature = "postgres")] + { + self.pg_pool = Some(pg.pool()); + } + + Arc::new(pg) as Arc + } + #[cfg(not(feature = "postgres"))] + _ => { + anyhow::bail!( + "No database backend available. Enable 'postgres' or 'libsql' feature." + ); + } + }; + + // Post-init: migrate disk config, reload config from DB, attach session, cleanup + if let Err(e) = crate::bootstrap::migrate_disk_to_db(db.as_ref(), "default").await { + tracing::warn!("Disk-to-DB settings migration failed: {}", e); + } + + let toml_path = self.toml_path.as_deref(); + match Config::from_db_with_toml(db.as_ref(), "default", toml_path).await { + Ok(db_config) => { + self.config = db_config; + tracing::info!("Configuration reloaded from database"); + } + Err(e) => { + tracing::warn!( + "Failed to reload config from DB, keeping env-based config: {}", + e + ); + } + } + + self.session.attach_store(db.clone(), "default").await; + + if let Err(e) = db.cleanup_stale_sandbox_jobs().await { + tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e); + } + + self.db = Some(db); + Ok(()) + } + + /// Phase 2: Create secrets store. + /// + /// Requires a master key and a backend-specific DB handle. After creating + /// the store, injects any encrypted LLM API keys into the config overlay + /// and re-resolves config. + pub async fn init_secrets(&mut self) -> Result<(), anyhow::Error> { + let master_key = match self.config.secrets.master_key() { + Some(k) => k, + None => { + // Consume unused handles + #[cfg(feature = "libsql")] + { + self.libsql_db.take(); + } + return Ok(()); + } + }; + + let crypto = match crate::secrets::SecretsCrypto::new(master_key.clone()) { + Ok(c) => Arc::new(c), + Err(e) => { + tracing::warn!("Failed to initialize secrets crypto: {}", e); + #[cfg(feature = "libsql")] + { + self.libsql_db.take(); + } + return Ok(()); + } + }; + + let store: Option> = None; + + #[cfg(feature = "libsql")] + let store = store.or_else(|| { + self.libsql_db.take().map(|db| { + Arc::new(crate::secrets::LibSqlSecretsStore::new( + db, + Arc::clone(&crypto), + )) as Arc + }) + }); + + #[cfg(feature = "postgres")] + let store = store.or_else(|| { + self.pg_pool.as_ref().map(|pool| { + Arc::new(crate::secrets::PostgresSecretsStore::new( + pool.clone(), + Arc::clone(&crypto), + )) as Arc + }) + }); + + if let Some(ref secrets) = store { + // Inject LLM API keys from encrypted storage + crate::config::inject_llm_keys_from_secrets(secrets.as_ref(), "default").await; + + // Re-resolve config with newly available keys + if let Some(ref db) = self.db { + let toml_path = self.toml_path.as_deref(); + match Config::from_db_with_toml(db.as_ref(), "default", toml_path).await { + Ok(refreshed) => { + self.config = refreshed; + tracing::debug!("LlmConfig re-resolved after secret injection"); + } + Err(e) => { + tracing::warn!("Failed to re-resolve config after secret injection: {}", e); + } + } + } + } + + self.secrets_store = store; + Ok(()) + } + + /// Phase 3: Initialize LLM provider chain. + /// + /// Creates the primary provider, then wraps with failover, circuit + /// breaker, and response cache as configured. + #[allow(clippy::type_complexity)] + pub fn init_llm( + &self, + ) -> Result<(Arc, Option>), anyhow::Error> { + use crate::llm::{ + CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig, + FailoverProvider, ResponseCacheConfig, create_cheap_llm_provider, create_llm_provider, + create_llm_provider_with_config, + }; + + let llm = create_llm_provider(&self.config.llm, self.session.clone())?; + tracing::info!("LLM provider initialized: {}", llm.model_name()); + + // Wrap in failover if a fallback model is configured + let llm: Arc = if let Some(fallback_model) = + self.config.llm.nearai.fallback_model.as_ref() + { + if fallback_model == &self.config.llm.nearai.model { + tracing::warn!( + "fallback_model is the same as primary model, failover may not be effective" + ); + } + let mut fallback_config = self.config.llm.nearai.clone(); + fallback_config.model = fallback_model.clone(); + let fallback = create_llm_provider_with_config(&fallback_config, self.session.clone())?; + tracing::info!( + primary = %llm.model_name(), + fallback = %fallback.model_name(), + "LLM failover enabled" + ); + let cooldown_config = CooldownConfig { + cooldown_duration: std::time::Duration::from_secs( + self.config.llm.nearai.failover_cooldown_secs, + ), + failure_threshold: self.config.llm.nearai.failover_cooldown_threshold, + }; + Arc::new(FailoverProvider::with_cooldown( + vec![llm, fallback], + cooldown_config, + )?) + } else { + llm + }; + + // Wrap in circuit breaker if configured + let llm: Arc = + if let Some(threshold) = self.config.llm.nearai.circuit_breaker_threshold { + let cb_config = CircuitBreakerConfig { + failure_threshold: threshold, + recovery_timeout: std::time::Duration::from_secs( + self.config.llm.nearai.circuit_breaker_recovery_secs, + ), + ..CircuitBreakerConfig::default() + }; + tracing::info!( + threshold, + recovery_secs = self.config.llm.nearai.circuit_breaker_recovery_secs, + "LLM circuit breaker enabled" + ); + Arc::new(CircuitBreakerProvider::new(llm, cb_config)) + } else { + llm + }; + + // Wrap in response cache if configured + let llm: Arc = if self.config.llm.nearai.response_cache_enabled { + let rc_config = ResponseCacheConfig { + ttl: std::time::Duration::from_secs(self.config.llm.nearai.response_cache_ttl_secs), + max_entries: self.config.llm.nearai.response_cache_max_entries, + }; + tracing::info!( + ttl_secs = self.config.llm.nearai.response_cache_ttl_secs, + max_entries = self.config.llm.nearai.response_cache_max_entries, + "LLM response cache enabled" + ); + Arc::new(CachedProvider::new(llm, rc_config)) + } else { + llm + }; + + // Cheap LLM for lightweight tasks + let cheap_llm = create_cheap_llm_provider(&self.config.llm, self.session.clone())?; + if let Some(ref cheap) = cheap_llm { + tracing::info!("Cheap LLM provider initialized: {}", cheap.model_name()); + } + + Ok((llm, cheap_llm)) + } + + /// Phase 4: Initialize safety, tools, embeddings, and workspace. + pub async fn init_tools( + &self, + llm: &Arc, + ) -> Result< + ( + Arc, + Arc, + Option>, + Option>, + ), + 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()); + tools.register_builtin_tools(); + tracing::info!("Registered {} built-in tools", tools.count()); + + // 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 + }; + + // Register memory tools if database is available + let workspace = if let Some(ref db) = self.db { + let mut ws = Workspace::new_with_db("default", db.clone()); + if let Some(ref emb) = embeddings { + ws = ws.with_embeddings(emb.clone()); + } + let ws = Arc::new(ws); + tools.register_memory_tools(Arc::clone(&ws)); + Some(ws) + } else { + None + }; + + // Register builder tool if enabled + if self.config.builder.enabled + && (self.config.agent.allow_local_tools || !self.config.sandbox.enabled) + { + tools + .register_builder_tool( + llm.clone(), + safety.clone(), + Some(self.config.builder.to_builder_config()), + ) + .await; + tracing::info!("Builder mode enabled"); + } + + Ok((safety, tools, embeddings, workspace)) + } + + /// Phase 5: Load WASM tools, MCP servers, and create extension manager. + pub async fn init_extensions( + &self, + tools: &Arc, + ) -> Result< + ( + Arc, + Option>, + Option>, + ), + anyhow::Error, + > { + use crate::tools::mcp::{McpClient, config::load_mcp_servers_from_db, is_authenticated}; + use crate::tools::wasm::{WasmToolLoader, load_dev_tools}; + + let mcp_session_manager = Arc::new(McpSessionManager::new()); + + // Create WASM tool runtime + let wasm_tool_runtime: Option> = + if self.config.wasm.enabled && self.config.wasm.tools_dir.exists() { + match WasmToolRuntime::new(self.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 + let wasm_tools_future = { + let wasm_tool_runtime = wasm_tool_runtime.clone(); + let secrets_store = self.secrets_store.clone(); + let tools = Arc::clone(tools); + let wasm_config = self.config.wasm.clone(); + async move { + if let Some(ref runtime) = wasm_tool_runtime { + let mut loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools)); + if let Some(ref secrets) = secrets_store { + loader = loader.with_secrets_store(Arc::clone(secrets)); + } + + match loader.load_from_dir(&wasm_config.tools_dir).await { + Ok(results) => { + if !results.loaded.is_empty() { + tracing::info!( + "Loaded {} WASM tools from {}", + results.loaded.len(), + wasm_config.tools_dir.display() + ); + } + for (path, err) in &results.errors { + tracing::warn!( + "Failed to load WASM tool {}: {}", + path.display(), + err + ); + } + } + Err(e) => { + tracing::warn!("Failed to scan WASM tools directory: {}", e); + } + } + + match load_dev_tools(&loader, &wasm_config.tools_dir).await { + Ok(results) => { + 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); + } + } + } + } + }; + + let mcp_servers_future = { + let secrets_store = self.secrets_store.clone(); + let db = self.db.clone(); + let tools = Arc::clone(tools); + let mcp_sm = Arc::clone(&mcp_session_manager); + async move { + 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 { + crate::tools::mcp::config::load_mcp_servers().await + }; + match servers_result { + Ok(servers) => { + let enabled: Vec<_> = servers.enabled_servers().cloned().collect(); + if !enabled.is_empty() { + tracing::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_sm); + let secrets = Arc::clone(secrets); + let tools = Arc::clone(&tools); + + join_set.spawn(async move { + let server_name = server.name.clone(); + let has_tokens = + is_authenticated(&server, &secrets, "default").await; + + 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) + }; + + match client.list_tools().await { + Ok(mcp_tools) => { + let tool_count = mcp_tools.len(); + match client.create_tools().await { + Ok(tool_impls) => { + for tool in tool_impls { + tools.register(tool).await; + } + tracing::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); + } + } + } + } + }; + + tokio::join!(wasm_tools_future, mcp_servers_future); + + // Create extension manager + let extension_manager = if let Some(ref secrets) = self.secrets_store { + let manager = Arc::new(ExtensionManager::new( + Arc::clone(&mcp_session_manager), + Arc::clone(secrets), + Arc::clone(tools), + wasm_tool_runtime.clone(), + self.config.wasm.tools_dir.clone(), + self.config.channels.wasm_channels_dir.clone(), + self.config.tunnel.public_url.clone(), + "default".to_string(), + self.db.clone(), + )); + tools.register_extension_tools(Arc::clone(&manager)); + tracing::info!("Extension manager initialized with in-chat discovery tools"); + Some(manager) + } else { + tracing::debug!( + "Extension manager not available (no secrets store). \ + Extension tools won't be registered." + ); + None + }; + + // Register dev tools if local tools are enabled + if self.config.agent.allow_local_tools { + tools.register_dev_tools(); + tracing::info!( + "Local tools enabled (allow_local_tools=true), dev tools registered directly" + ); + } + + Ok((mcp_session_manager, wasm_tool_runtime, extension_manager)) + } + + /// Run all init phases in order and return the assembled components. + pub async fn build_all(mut self) -> Result { + self.init_database().await?; + self.init_secrets().await?; + + let (llm, cheap_llm) = self.init_llm()?; + let (safety, tools, embeddings, workspace) = self.init_tools(&llm).await?; + let (mcp_session_manager, wasm_tool_runtime, extension_manager) = + self.init_extensions(&tools).await?; + + // Seed workspace and backfill embeddings + if let Some(ref ws) = workspace { + match ws.seed_if_empty().await { + Ok(count) if count > 0 => { + tracing::info!("Workspace seeded with {} core files", count); + } + Ok(_) => {} + Err(e) => { + tracing::warn!("Failed to seed workspace: {}", e); + } + } + + if embeddings.is_some() { + match ws.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); + } + } + } + } + + // Skills system + let (skill_registry, skill_catalog) = if self.config.skills.enabled { + let mut registry = SkillRegistry::new(self.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)); + let catalog = crate::skills::catalog::shared_catalog(); + tools.register_skill_tools(Arc::clone(®istry), Arc::clone(&catalog)); + (Some(registry), Some(catalog)) + } else { + (None, None) + }; + + let context_manager = Arc::new(ContextManager::new(self.config.agent.max_parallel_jobs)); + let hooks = Arc::new(HookRegistry::new()); + let cost_guard = Arc::new(crate::agent::cost_guard::CostGuard::new( + crate::agent::cost_guard::CostGuardConfig { + max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents, + max_actions_per_hour: self.config.agent.max_actions_per_hour, + }, + )); + + tracing::info!( + "Tool registry initialized with {} total tools", + tools.count() + ); + + Ok(AppComponents { + config: self.config, + db: self.db, + secrets_store: self.secrets_store, + llm, + cheap_llm, + safety, + tools, + embeddings, + workspace, + extension_manager, + mcp_session_manager, + wasm_tool_runtime, + log_broadcaster: self.log_broadcaster, + context_manager, + hooks, + skill_registry, + skill_catalog, + cost_guard, + session: self.session, + }) + } +} diff --git a/src/channels/web/handlers/chat.rs b/src/channels/web/handlers/chat.rs new file mode 100644 index 00000000..753b6d99 --- /dev/null +++ b/src/channels/web/handlers/chat.rs @@ -0,0 +1,633 @@ +//! Chat handlers: send, approval, auth, SSE events, WebSocket, history, threads. + +use std::sync::Arc; + +use axum::{ + Json, + extract::{Query, State, WebSocketUpgrade}, + http::StatusCode, + response::IntoResponse, +}; +use serde::Deserialize; +use uuid::Uuid; + +use crate::channels::IncomingMessage; +use crate::channels::web::server::GatewayState; +use crate::channels::web::types::*; + +pub async fn chat_send_handler( + State(state): State>, + Json(req): Json, +) -> Result<(StatusCode, Json), (StatusCode, String)> { + if !state.chat_rate_limiter.check() { + return Err(( + StatusCode::TOO_MANY_REQUESTS, + "Rate limit exceeded. Try again shortly.".to_string(), + )); + } + + let mut msg = IncomingMessage::new("gateway", &state.user_id, &req.content); + + if let Some(ref thread_id) = req.thread_id { + msg = msg.with_thread(thread_id); + msg = msg.with_metadata(serde_json::json!({"thread_id": thread_id})); + } + + let msg_id = msg.id; + + let tx_guard = state.msg_tx.read().await; + let tx = tx_guard.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Channel not started".to_string(), + ))?; + + tx.send(msg).await.map_err(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + "Channel closed".to_string(), + ) + })?; + + Ok(( + StatusCode::ACCEPTED, + Json(SendMessageResponse { + message_id: msg_id, + status: "accepted", + }), + )) +} + +pub async fn chat_approval_handler( + State(state): State>, + Json(req): Json, +) -> Result<(StatusCode, Json), (StatusCode, String)> { + let (approved, always) = match req.action.as_str() { + "approve" => (true, false), + "always" => (true, true), + "deny" => (false, false), + other => { + return Err(( + StatusCode::BAD_REQUEST, + format!("Unknown action: {}", other), + )); + } + }; + + let request_id = Uuid::parse_str(&req.request_id).map_err(|_| { + ( + StatusCode::BAD_REQUEST, + "Invalid request_id (expected UUID)".to_string(), + ) + })?; + + // Build a structured ExecApproval submission as JSON, sent through the + // existing message pipeline so the agent loop picks it up. + let approval = crate::agent::submission::Submission::ExecApproval { + request_id, + approved, + always, + }; + let content = serde_json::to_string(&approval).map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to serialize approval: {}", e), + ) + })?; + + let mut msg = IncomingMessage::new("gateway", &state.user_id, content); + + if let Some(ref thread_id) = req.thread_id { + msg = msg.with_thread(thread_id); + } + + let msg_id = msg.id; + + let tx_guard = state.msg_tx.read().await; + let tx = tx_guard.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Channel not started".to_string(), + ))?; + + tx.send(msg).await.map_err(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + "Channel closed".to_string(), + ) + })?; + + Ok(( + StatusCode::ACCEPTED, + Json(SendMessageResponse { + message_id: msg_id, + status: "accepted", + }), + )) +} + +/// Submit an auth token directly to the extension manager, bypassing the message pipeline. +/// +/// The token never touches the LLM, chat history, or SSE stream. +pub async fn chat_auth_token_handler( + State(state): State>, + Json(req): Json, +) -> Result, (StatusCode, String)> { + let ext_mgr = state.extension_manager.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Extension manager not available".to_string(), + ))?; + + let result = ext_mgr + .auth(&req.extension_name, Some(&req.token)) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + if result.status == "authenticated" { + // Auto-activate so tools are available immediately + let msg = match ext_mgr.activate(&req.extension_name).await { + Ok(r) => format!( + "{} authenticated ({} tools loaded)", + req.extension_name, + r.tools_loaded.len() + ), + Err(e) => format!( + "{} authenticated but activation failed: {}", + req.extension_name, e + ), + }; + + // Clear auth mode on the active thread + clear_auth_mode(&state).await; + + state.sse.broadcast(SseEvent::AuthCompleted { + extension_name: req.extension_name, + success: true, + message: msg.clone(), + }); + + Ok(Json(ActionResponse::ok(msg))) + } else { + // Re-emit auth_required for retry + state.sse.broadcast(SseEvent::AuthRequired { + extension_name: req.extension_name.clone(), + instructions: result.instructions.clone(), + auth_url: result.auth_url.clone(), + setup_url: result.setup_url.clone(), + }); + Ok(Json(ActionResponse::fail( + result + .instructions + .unwrap_or_else(|| "Invalid token".to_string()), + ))) + } +} + +/// Cancel an in-progress auth flow. +pub async fn chat_auth_cancel_handler( + State(state): State>, + Json(_req): Json, +) -> Result, (StatusCode, String)> { + clear_auth_mode(&state).await; + Ok(Json(ActionResponse::ok("Auth cancelled"))) +} + +/// Clear pending auth mode on the active thread. +pub async fn clear_auth_mode(state: &GatewayState) { + if let Some(ref sm) = state.session_manager { + let session = sm.get_or_create_session(&state.user_id).await; + let mut sess = session.lock().await; + if let Some(thread_id) = sess.active_thread + && let Some(thread) = sess.threads.get_mut(&thread_id) + { + thread.pending_auth = None; + } + } +} + +pub async fn chat_events_handler( + State(state): State>, +) -> Result { + state.sse.subscribe().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Too many connections".to_string(), + )) +} + +pub async fn chat_ws_handler( + headers: axum::http::HeaderMap, + ws: WebSocketUpgrade, + State(state): State>, +) -> Result { + // Validate Origin header to prevent cross-site WebSocket hijacking. + let origin = headers + .get("origin") + .and_then(|v| v.to_str().ok()) + .ok_or_else(|| { + ( + StatusCode::FORBIDDEN, + "WebSocket Origin header required".to_string(), + ) + })?; + + let host = origin + .strip_prefix("http://") + .or_else(|| origin.strip_prefix("https://")) + .and_then(|rest| rest.split(':').next()?.split('/').next()) + .unwrap_or(""); + + let is_local = matches!(host, "localhost" | "127.0.0.1" | "[::1]"); + if !is_local { + return Err(( + StatusCode::FORBIDDEN, + "WebSocket origin not allowed".to_string(), + )); + } + Ok(ws.on_upgrade(move |socket| crate::channels::web::ws::handle_ws_connection(socket, state))) +} + +#[derive(Deserialize)] +pub struct HistoryQuery { + pub thread_id: Option, + pub limit: Option, + pub before: Option, +} + +pub async fn chat_history_handler( + State(state): State>, + Query(query): Query, +) -> Result, (StatusCode, String)> { + let session_manager = state.session_manager.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Session manager not available".to_string(), + ))?; + + let session = session_manager.get_or_create_session(&state.user_id).await; + let sess = session.lock().await; + + let limit = query.limit.unwrap_or(50); + let before_cursor = query + .before + .as_deref() + .map(|s| { + chrono::DateTime::parse_from_rfc3339(s) + .map(|dt| dt.with_timezone(&chrono::Utc)) + .map_err(|_| { + ( + StatusCode::BAD_REQUEST, + "Invalid 'before' timestamp".to_string(), + ) + }) + }) + .transpose()?; + + // Find the thread + let thread_id = if let Some(ref tid) = query.thread_id { + Uuid::parse_str(tid) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid thread_id".to_string()))? + } else { + sess.active_thread + .ok_or((StatusCode::NOT_FOUND, "No active thread".to_string()))? + }; + + // Verify the thread belongs to the authenticated user before returning any data. + if query.thread_id.is_some() + && let Some(ref store) = state.store + { + let owned = store + .conversation_belongs_to_user(thread_id, &state.user_id) + .await + .unwrap_or(false); + if !owned && !sess.threads.contains_key(&thread_id) { + return Err((StatusCode::NOT_FOUND, "Thread not found".to_string())); + } + } + + // For paginated requests (before cursor set), always go to DB + if before_cursor.is_some() + && let Some(ref store) = state.store + { + let (messages, has_more) = store + .list_conversation_messages_paginated(thread_id, before_cursor, limit as i64) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339()); + let turns = build_turns_from_db_messages(&messages); + return Ok(Json(HistoryResponse { + thread_id, + turns, + has_more, + oldest_timestamp, + })); + } + + // Try in-memory first (freshest data for active threads) + if let Some(thread) = sess.threads.get(&thread_id) + && !thread.turns.is_empty() + { + let turns: Vec = thread + .turns + .iter() + .map(|t| TurnInfo { + turn_number: t.turn_number, + user_input: t.user_input.clone(), + response: t.response.clone(), + state: format!("{:?}", t.state), + started_at: t.started_at.to_rfc3339(), + completed_at: t.completed_at.map(|dt| dt.to_rfc3339()), + tool_calls: t + .tool_calls + .iter() + .map(|tc| ToolCallInfo { + name: tc.name.clone(), + has_result: tc.result.is_some(), + has_error: tc.error.is_some(), + }) + .collect(), + }) + .collect(); + + return Ok(Json(HistoryResponse { + thread_id, + turns, + has_more: false, + oldest_timestamp: None, + })); + } + + // Fall back to DB for historical threads not in memory (paginated) + if let Some(ref store) = state.store { + let (messages, has_more) = store + .list_conversation_messages_paginated(thread_id, None, limit as i64) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + if !messages.is_empty() { + let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339()); + let turns = build_turns_from_db_messages(&messages); + return Ok(Json(HistoryResponse { + thread_id, + turns, + has_more, + oldest_timestamp, + })); + } + } + + // Empty thread (just created, no messages yet) + Ok(Json(HistoryResponse { + thread_id, + turns: Vec::new(), + has_more: false, + oldest_timestamp: None, + })) +} + +/// Build TurnInfo pairs from flat DB messages (alternating user/assistant). +pub fn build_turns_from_db_messages( + messages: &[crate::history::ConversationMessage], +) -> Vec { + let mut turns = Vec::new(); + let mut turn_number = 0; + let mut iter = messages.iter().peekable(); + + while let Some(msg) = iter.next() { + if msg.role == "user" { + let mut turn = TurnInfo { + turn_number, + user_input: msg.content.clone(), + response: None, + state: "Completed".to_string(), + started_at: msg.created_at.to_rfc3339(), + completed_at: None, + tool_calls: Vec::new(), + }; + + // Check if next message is an assistant response + if let Some(next) = iter.peek() + && next.role == "assistant" + { + let assistant_msg = iter.next().expect("peeked"); + turn.response = Some(assistant_msg.content.clone()); + turn.completed_at = Some(assistant_msg.created_at.to_rfc3339()); + } + + // Incomplete turn (user message without response) + if turn.response.is_none() { + turn.state = "Failed".to_string(); + } + + turns.push(turn); + turn_number += 1; + } + } + + turns +} + +pub async fn chat_threads_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let session_manager = state.session_manager.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Session manager not available".to_string(), + ))?; + + let session = session_manager.get_or_create_session(&state.user_id).await; + let sess = session.lock().await; + + // Try DB first for persistent thread list + if let Some(ref store) = state.store { + // Auto-create assistant thread if it doesn't exist + let assistant_id = store + .get_or_create_assistant_conversation(&state.user_id, "gateway") + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + if let Ok(summaries) = store + .list_conversations_with_preview(&state.user_id, "gateway", 50) + .await + { + let mut assistant_thread = None; + let mut threads = Vec::new(); + + for s in &summaries { + let info = ThreadInfo { + id: s.id, + state: "Idle".to_string(), + turn_count: (s.message_count / 2).max(0) as usize, + created_at: s.started_at.to_rfc3339(), + updated_at: s.last_activity.to_rfc3339(), + title: s.title.clone(), + thread_type: s.thread_type.clone(), + }; + + if s.id == assistant_id { + assistant_thread = Some(info); + } else { + threads.push(info); + } + } + + // If assistant wasn't in the list (0 messages), synthesize it + if assistant_thread.is_none() { + assistant_thread = Some(ThreadInfo { + id: assistant_id, + state: "Idle".to_string(), + turn_count: 0, + created_at: chrono::Utc::now().to_rfc3339(), + updated_at: chrono::Utc::now().to_rfc3339(), + title: None, + thread_type: Some("assistant".to_string()), + }); + } + + return Ok(Json(ThreadListResponse { + assistant_thread, + threads, + active_thread: sess.active_thread, + })); + } + } + + // Fallback: in-memory only (no assistant thread without DB) + let threads: Vec = sess + .threads + .values() + .map(|t| ThreadInfo { + id: t.id, + state: format!("{:?}", t.state), + turn_count: t.turns.len(), + created_at: t.created_at.to_rfc3339(), + updated_at: t.updated_at.to_rfc3339(), + title: None, + thread_type: None, + }) + .collect(); + + Ok(Json(ThreadListResponse { + assistant_thread: None, + threads, + active_thread: sess.active_thread, + })) +} + +pub async fn chat_new_thread_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let session_manager = state.session_manager.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Session manager not available".to_string(), + ))?; + + let session = session_manager.get_or_create_session(&state.user_id).await; + let mut sess = session.lock().await; + let thread = sess.create_thread(); + let thread_id = thread.id; + let info = ThreadInfo { + id: thread.id, + state: format!("{:?}", thread.state), + turn_count: thread.turns.len(), + created_at: thread.created_at.to_rfc3339(), + updated_at: thread.updated_at.to_rfc3339(), + title: None, + thread_type: Some("thread".to_string()), + }; + + // Persist the empty conversation row with thread_type metadata + if let Some(ref store) = state.store { + let store = Arc::clone(store); + let user_id = state.user_id.clone(); + tokio::spawn(async move { + if let Err(e) = store + .ensure_conversation(thread_id, "gateway", &user_id, None) + .await + { + tracing::warn!("Failed to persist new thread: {}", e); + } + let metadata_val = serde_json::json!("thread"); + if let Err(e) = store + .update_conversation_metadata_field(thread_id, "thread_type", &metadata_val) + .await + { + tracing::warn!("Failed to set thread_type metadata: {}", e); + } + }); + } + + Ok(Json(info)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_build_turns_from_db_messages_complete() { + let now = chrono::Utc::now(); + let messages = vec![ + crate::history::ConversationMessage { + id: Uuid::new_v4(), + role: "user".to_string(), + content: "Hello".to_string(), + created_at: now, + }, + crate::history::ConversationMessage { + id: Uuid::new_v4(), + role: "assistant".to_string(), + content: "Hi there!".to_string(), + created_at: now + chrono::TimeDelta::seconds(1), + }, + crate::history::ConversationMessage { + id: Uuid::new_v4(), + role: "user".to_string(), + content: "How are you?".to_string(), + created_at: now + chrono::TimeDelta::seconds(2), + }, + crate::history::ConversationMessage { + id: Uuid::new_v4(), + role: "assistant".to_string(), + content: "Doing well!".to_string(), + created_at: now + chrono::TimeDelta::seconds(3), + }, + ]; + + let turns = build_turns_from_db_messages(&messages); + assert_eq!(turns.len(), 2); + assert_eq!(turns[0].user_input, "Hello"); + assert_eq!(turns[0].response.as_deref(), Some("Hi there!")); + assert_eq!(turns[0].state, "Completed"); + assert_eq!(turns[1].user_input, "How are you?"); + assert_eq!(turns[1].response.as_deref(), Some("Doing well!")); + } + + #[test] + fn test_build_turns_from_db_messages_incomplete_last() { + let now = chrono::Utc::now(); + let messages = vec![ + crate::history::ConversationMessage { + id: Uuid::new_v4(), + role: "user".to_string(), + content: "Hello".to_string(), + created_at: now, + }, + crate::history::ConversationMessage { + id: Uuid::new_v4(), + role: "assistant".to_string(), + content: "Hi!".to_string(), + created_at: now + chrono::TimeDelta::seconds(1), + }, + crate::history::ConversationMessage { + id: Uuid::new_v4(), + role: "user".to_string(), + content: "Lost message".to_string(), + created_at: now + chrono::TimeDelta::seconds(2), + }, + ]; + + let turns = build_turns_from_db_messages(&messages); + assert_eq!(turns.len(), 2); + assert_eq!(turns[1].user_input, "Lost message"); + assert!(turns[1].response.is_none()); + assert_eq!(turns[1].state, "Failed"); + } +} diff --git a/src/channels/web/handlers/extensions.rs b/src/channels/web/handlers/extensions.rs new file mode 100644 index 00000000..7860184a --- /dev/null +++ b/src/channels/web/handlers/extensions.rs @@ -0,0 +1,153 @@ +//! Extension management API handlers. + +use std::sync::Arc; + +use axum::{ + Json, + extract::{Path, State}, + http::StatusCode, +}; + +use crate::channels::web::server::GatewayState; +use crate::channels::web::types::*; + +pub async fn extensions_list_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let ext_mgr = state.extension_manager.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Extension manager not available (secrets store required)".to_string(), + ))?; + + let installed = ext_mgr + .list(None) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let extensions = installed + .into_iter() + .map(|ext| ExtensionInfo { + name: ext.name, + kind: ext.kind.to_string(), + description: ext.description, + url: ext.url, + authenticated: ext.authenticated, + active: ext.active, + tools: ext.tools, + }) + .collect(); + + Ok(Json(ExtensionListResponse { extensions })) +} + +pub async fn extensions_tools_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let registry = state.tool_registry.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Tool registry not available".to_string(), + ))?; + + let definitions = registry.tool_definitions().await; + let tools = definitions + .into_iter() + .map(|td| ToolInfo { + name: td.name, + description: td.description, + }) + .collect(); + + Ok(Json(ToolListResponse { tools })) +} + +pub async fn extensions_install_handler( + State(state): State>, + Json(req): Json, +) -> Result, (StatusCode, String)> { + let ext_mgr = state.extension_manager.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Extension manager not available (secrets store required)".to_string(), + ))?; + + let kind_hint = req.kind.as_deref().and_then(|k| match k { + "mcp_server" => Some(crate::extensions::ExtensionKind::McpServer), + "wasm_tool" => Some(crate::extensions::ExtensionKind::WasmTool), + "wasm_channel" => Some(crate::extensions::ExtensionKind::WasmChannel), + _ => None, + }); + + match ext_mgr + .install(&req.name, req.url.as_deref(), kind_hint) + .await + { + Ok(result) => Ok(Json(ActionResponse::ok(result.message))), + Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))), + } +} + +pub async fn extensions_activate_handler( + State(state): State>, + Path(name): Path, +) -> Result, (StatusCode, String)> { + let ext_mgr = state.extension_manager.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Extension manager not available (secrets store required)".to_string(), + ))?; + + match ext_mgr.activate(&name).await { + Ok(result) => Ok(Json(ActionResponse::ok(result.message))), + Err(activate_err) => { + let err_str = activate_err.to_string(); + let needs_auth = err_str.contains("authentication") + || err_str.contains("401") + || err_str.contains("Unauthorized"); + + if !needs_auth { + return Ok(Json(ActionResponse::fail(err_str))); + } + + // Activation failed due to auth; try authenticating first. + match ext_mgr.auth(&name, None).await { + Ok(auth_result) if auth_result.status == "authenticated" => { + // Auth succeeded, retry activation. + match ext_mgr.activate(&name).await { + Ok(result) => Ok(Json(ActionResponse::ok(result.message))), + Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))), + } + } + Ok(auth_result) => { + // Auth in progress (OAuth URL or awaiting manual token). + let mut resp = ActionResponse::fail( + auth_result + .instructions + .clone() + .unwrap_or_else(|| format!("'{}' requires authentication.", name)), + ); + resp.auth_url = auth_result.auth_url; + resp.awaiting_token = Some(auth_result.awaiting_token); + resp.instructions = auth_result.instructions; + Ok(Json(resp)) + } + Err(auth_err) => Ok(Json(ActionResponse::fail(format!( + "Authentication failed: {}", + auth_err + )))), + } + } + } +} + +pub async fn extensions_remove_handler( + State(state): State>, + Path(name): Path, +) -> Result, (StatusCode, String)> { + let ext_mgr = state.extension_manager.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Extension manager not available (secrets store required)".to_string(), + ))?; + + match ext_mgr.remove(&name).await { + Ok(message) => Ok(Json(ActionResponse::ok(message))), + Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))), + } +} diff --git a/src/channels/web/handlers/jobs.rs b/src/channels/web/handlers/jobs.rs new file mode 100644 index 00000000..567acc7a --- /dev/null +++ b/src/channels/web/handlers/jobs.rs @@ -0,0 +1,518 @@ +//! Job and sandbox API handlers. + +use std::sync::Arc; + +use axum::{ + Json, + extract::{Path, Query, State}, + http::StatusCode, +}; +use serde::Deserialize; +use uuid::Uuid; + +use crate::channels::web::server::GatewayState; +use crate::channels::web::types::*; + +pub async fn jobs_list_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + // Fetch sandbox jobs scoped to the authenticated user. + let sandbox_jobs = store + .list_sandbox_jobs_for_user(&state.user_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + // Scope jobs to the authenticated user. + let mut jobs: Vec = sandbox_jobs + .iter() + .filter(|j| j.user_id == state.user_id) + .map(|j| { + let ui_state = match j.status.as_str() { + "creating" => "pending", + "running" => "in_progress", + s => s, + }; + JobInfo { + id: j.id, + title: j.task.clone(), + state: ui_state.to_string(), + user_id: j.user_id.clone(), + created_at: j.created_at.to_rfc3339(), + started_at: j.started_at.map(|dt| dt.to_rfc3339()), + } + }) + .collect(); + + // Most recent first. + jobs.sort_by(|a, b| b.created_at.cmp(&a.created_at)); + + Ok(Json(JobListResponse { jobs })) +} + +pub async fn jobs_summary_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let s = store + .sandbox_job_summary_for_user(&state.user_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + Ok(Json(JobSummaryResponse { + total: s.total, + pending: s.creating, + in_progress: s.running, + completed: s.completed, + failed: s.failed + s.interrupted, + stuck: 0, + })) +} + +pub async fn jobs_detail_handler( + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let job_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; + + // Try sandbox job from DB first, scoped to the authenticated user. + if let Some(ref store) = state.store + && let Ok(Some(job)) = store.get_sandbox_job(job_id).await + { + if job.user_id != state.user_id { + return Err((StatusCode::NOT_FOUND, "Job not found".to_string())); + } + let browse_id = std::path::Path::new(&job.project_dir) + .file_name() + .map(|n| n.to_string_lossy().to_string()) + .unwrap_or_else(|| job.id.to_string()); + + let ui_state = match job.status.as_str() { + "creating" => "pending", + "running" => "in_progress", + s => s, + }; + + let elapsed_secs = job.started_at.map(|start| { + let end = job.completed_at.unwrap_or_else(chrono::Utc::now); + (end - start).num_seconds().max(0) as u64 + }); + + // Synthesize transitions from timestamps. + let mut transitions = Vec::new(); + if let Some(started) = job.started_at { + transitions.push(TransitionInfo { + from: "creating".to_string(), + to: "running".to_string(), + timestamp: started.to_rfc3339(), + reason: None, + }); + } + if let Some(completed) = job.completed_at { + transitions.push(TransitionInfo { + from: "running".to_string(), + to: job.status.clone(), + timestamp: completed.to_rfc3339(), + reason: job.failure_reason.clone(), + }); + } + + return Ok(Json(JobDetailResponse { + id: job.id, + title: job.task.clone(), + description: String::new(), + state: ui_state.to_string(), + user_id: job.user_id.clone(), + created_at: job.created_at.to_rfc3339(), + started_at: job.started_at.map(|dt| dt.to_rfc3339()), + completed_at: job.completed_at.map(|dt| dt.to_rfc3339()), + elapsed_secs, + project_dir: Some(job.project_dir.clone()), + browse_url: Some(format!("/projects/{}/", browse_id)), + job_mode: { + let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten(); + mode.filter(|m| m != "worker") + }, + transitions, + })); + } + + Err((StatusCode::NOT_FOUND, "Job not found".to_string())) +} + +pub async fn jobs_cancel_handler( + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let job_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; + + // Try sandbox job cancellation, scoped to the authenticated user. + if let Some(ref store) = state.store + && let Ok(Some(job)) = store.get_sandbox_job(job_id).await + { + if job.user_id != state.user_id { + return Err((StatusCode::NOT_FOUND, "Job not found".to_string())); + } + if job.status == "running" || job.status == "creating" { + // Stop the container if we have a job manager. + if let Some(ref jm) = state.job_manager + && let Err(e) = jm.stop_job(job_id).await + { + tracing::warn!(job_id = %job_id, error = %e, "Failed to stop container during cancellation"); + } + store + .update_sandbox_job_status( + job_id, + "failed", + Some(false), + Some("Cancelled by user"), + None, + Some(chrono::Utc::now()), + ) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + } + return Ok(Json(serde_json::json!({ + "status": "cancelled", + "job_id": job_id, + }))); + } + + Err((StatusCode::NOT_FOUND, "Job not found".to_string())) +} + +pub async fn jobs_restart_handler( + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + let jm = state.job_manager.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Sandbox not enabled".to_string(), + ))?; + + let old_job_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; + + let old_job = store + .get_sandbox_job(old_job_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?; + + // Scope to the authenticated user. + if old_job.user_id != state.user_id { + return Err((StatusCode::NOT_FOUND, "Job not found".to_string())); + } + + if old_job.status != "interrupted" && old_job.status != "failed" { + return Err(( + StatusCode::CONFLICT, + format!("Cannot restart job in state '{}'", old_job.status), + )); + } + + // Create a new job with the same task and project_dir. + let new_job_id = Uuid::new_v4(); + let now = chrono::Utc::now(); + + let record = crate::history::SandboxJobRecord { + id: new_job_id, + task: old_job.task.clone(), + status: "creating".to_string(), + user_id: old_job.user_id.clone(), + project_dir: old_job.project_dir.clone(), + success: None, + failure_reason: None, + created_at: now, + started_at: None, + completed_at: None, + credential_grants_json: old_job.credential_grants_json.clone(), + }; + store + .save_sandbox_job(&record) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + // Look up the original job's mode so the restart uses the same mode. + let mode = match store.get_sandbox_job_mode(old_job_id).await { + Ok(Some(m)) if m == "claude_code" => crate::orchestrator::job_manager::JobMode::ClaudeCode, + _ => crate::orchestrator::job_manager::JobMode::Worker, + }; + + // Restore credential grants from the original job so the restarted container + // has access to the same secrets. + let credential_grants: Vec = + serde_json::from_str(&old_job.credential_grants_json).unwrap_or_else(|e| { + tracing::warn!( + job_id = %old_job.id, + "Failed to deserialize credential grants from stored job: {}. \ + Restarted job will have no credentials.", + e + ); + vec![] + }); + + let project_dir = std::path::PathBuf::from(&old_job.project_dir); + let _token = jm + .create_job( + new_job_id, + &old_job.task, + Some(project_dir), + mode, + credential_grants, + ) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to create container: {}", e), + ) + })?; + + store + .update_sandbox_job_status(new_job_id, "running", None, None, Some(now), None) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + Ok(Json(serde_json::json!({ + "status": "restarted", + "old_job_id": old_job_id, + "new_job_id": new_job_id, + }))) +} + +/// Submit a follow-up prompt to a running Claude Code sandbox job. +pub async fn jobs_prompt_handler( + State(state): State>, + Path(id): Path, + Json(body): Json, +) -> Result, (StatusCode, String)> { + let prompt_queue = state.prompt_queue.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Claude Code not configured".to_string(), + ))?; + + let job_id: uuid::Uuid = id + .parse() + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; + + // Verify user owns this job. + if let Some(ref store) = state.store + && !store + .sandbox_job_belongs_to_user(job_id, &state.user_id) + .await + .unwrap_or(false) + { + return Err((StatusCode::NOT_FOUND, "Job not found".to_string())); + } + + let content = body + .get("content") + .and_then(|v| v.as_str()) + .ok_or(( + StatusCode::BAD_REQUEST, + "Missing 'content' field".to_string(), + ))? + .to_string(); + + let done = body.get("done").and_then(|v| v.as_bool()).unwrap_or(false); + + let prompt = crate::orchestrator::api::PendingPrompt { content, done }; + + { + let mut queue = prompt_queue.lock().await; + queue.entry(job_id).or_default().push_back(prompt); + } + + Ok(Json(serde_json::json!({ + "status": "queued", + "job_id": job_id.to_string(), + }))) +} + +/// Load persisted job events for a job (for history replay on page open). +pub async fn jobs_events_handler( + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Database not available".to_string(), + ))?; + + let job_id: uuid::Uuid = id + .parse() + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; + + // Verify user owns this job. + if !store + .sandbox_job_belongs_to_user(job_id, &state.user_id) + .await + .unwrap_or(false) + { + return Err((StatusCode::NOT_FOUND, "Job not found".to_string())); + } + + let events = store + .list_job_events(job_id, None) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let events_json: Vec = events + .into_iter() + .map(|e| { + serde_json::json!({ + "id": e.id, + "event_type": e.event_type, + "data": e.data, + "created_at": e.created_at.to_rfc3339(), + }) + }) + .collect(); + + Ok(Json(serde_json::json!({ + "job_id": job_id.to_string(), + "events": events_json, + }))) +} + +// --- Project file handlers for sandbox jobs --- + +#[derive(Deserialize)] +pub struct FilePathQuery { + pub path: Option, +} + +pub async fn job_files_list_handler( + State(state): State>, + Path(id): Path, + Query(query): Query, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let job_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; + + let job = store + .get_sandbox_job(job_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?; + + // Verify user owns this job. + if job.user_id != state.user_id { + return Err((StatusCode::NOT_FOUND, "Job not found".to_string())); + } + + let base = std::path::PathBuf::from(&job.project_dir); + let rel_path = query.path.as_deref().unwrap_or(""); + let target = base.join(rel_path); + + // Path traversal guard. + let canonical = target + .canonicalize() + .map_err(|_| (StatusCode::NOT_FOUND, "Path not found".to_string()))?; + let base_canonical = base + .canonicalize() + .map_err(|_| (StatusCode::NOT_FOUND, "Project dir not found".to_string()))?; + if !canonical.starts_with(&base_canonical) { + return Err((StatusCode::FORBIDDEN, "Forbidden".to_string())); + } + + let mut entries = Vec::new(); + let mut read_dir = tokio::fs::read_dir(&canonical) + .await + .map_err(|_| (StatusCode::NOT_FOUND, "Cannot read directory".to_string()))?; + + while let Ok(Some(entry)) = read_dir.next_entry().await { + let name = entry.file_name().to_string_lossy().to_string(); + let is_dir = entry + .file_type() + .await + .map(|ft| ft.is_dir()) + .unwrap_or(false); + let rel = if rel_path.is_empty() { + name.clone() + } else { + format!("{}/{}", rel_path, name) + }; + entries.push(ProjectFileEntry { + name, + path: rel, + is_dir, + }); + } + + entries.sort_by(|a, b| b.is_dir.cmp(&a.is_dir).then_with(|| a.name.cmp(&b.name))); + + Ok(Json(ProjectFilesResponse { entries })) +} + +pub async fn job_files_read_handler( + State(state): State>, + Path(id): Path, + Query(query): Query, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let job_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; + + let job = store + .get_sandbox_job(job_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?; + + // Verify user owns this job. + if job.user_id != state.user_id { + return Err((StatusCode::NOT_FOUND, "Job not found".to_string())); + } + + let path = query.path.as_deref().ok_or(( + StatusCode::BAD_REQUEST, + "path parameter required".to_string(), + ))?; + + let base = std::path::PathBuf::from(&job.project_dir); + let file_path = base.join(path); + + let canonical = file_path + .canonicalize() + .map_err(|_| (StatusCode::NOT_FOUND, "File not found".to_string()))?; + let base_canonical = base + .canonicalize() + .map_err(|_| (StatusCode::NOT_FOUND, "Project dir not found".to_string()))?; + if !canonical.starts_with(&base_canonical) { + return Err((StatusCode::FORBIDDEN, "Forbidden".to_string())); + } + + let content = tokio::fs::read_to_string(&canonical) + .await + .map_err(|_| (StatusCode::NOT_FOUND, "Cannot read file".to_string()))?; + + Ok(Json(ProjectFileReadResponse { + path: path.to_string(), + content, + })) +} diff --git a/src/channels/web/handlers/memory.rs b/src/channels/web/handlers/memory.rs new file mode 100644 index 00000000..59655d51 --- /dev/null +++ b/src/channels/web/handlers/memory.rs @@ -0,0 +1,171 @@ +//! Memory/workspace API handlers. + +use std::sync::Arc; + +use axum::{ + Json, + extract::{Query, State}, + http::StatusCode, +}; +use serde::Deserialize; + +use crate::channels::web::server::GatewayState; +use crate::channels::web::types::*; + +#[derive(Deserialize)] +pub struct TreeQuery { + #[allow(dead_code)] + pub depth: Option, +} + +pub async fn memory_tree_handler( + State(state): State>, + Query(_query): Query, +) -> Result, (StatusCode, String)> { + let workspace = state.workspace.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Workspace not available".to_string(), + ))?; + + // Build tree from list_all (flat list of all paths) + let all_paths = workspace + .list_all() + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + // Collect unique directories and files + let mut entries: Vec = Vec::new(); + let mut seen_dirs: std::collections::HashSet = std::collections::HashSet::new(); + + for path in &all_paths { + // Add parent directories + let parts: Vec<&str> = path.split('/').collect(); + for i in 0..parts.len().saturating_sub(1) { + let dir_path = parts[..=i].join("/"); + if seen_dirs.insert(dir_path.clone()) { + entries.push(TreeEntry { + path: dir_path, + is_dir: true, + }); + } + } + // Add the file itself + entries.push(TreeEntry { + path: path.clone(), + is_dir: false, + }); + } + + entries.sort_by(|a, b| a.path.cmp(&b.path)); + + Ok(Json(MemoryTreeResponse { entries })) +} + +#[derive(Deserialize)] +pub struct ListQuery { + pub path: Option, +} + +pub async fn memory_list_handler( + State(state): State>, + Query(query): Query, +) -> Result, (StatusCode, String)> { + let workspace = state.workspace.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Workspace not available".to_string(), + ))?; + + let path = query.path.as_deref().unwrap_or(""); + let entries = workspace + .list(path) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let list_entries: Vec = entries + .iter() + .map(|e| ListEntry { + name: e.path.rsplit('/').next().unwrap_or(&e.path).to_string(), + path: e.path.clone(), + is_dir: e.is_directory, + updated_at: e.updated_at.map(|dt| dt.to_rfc3339()), + }) + .collect(); + + Ok(Json(MemoryListResponse { + path: path.to_string(), + entries: list_entries, + })) +} + +#[derive(Deserialize)] +pub struct ReadQuery { + pub path: String, +} + +pub async fn memory_read_handler( + State(state): State>, + Query(query): Query, +) -> Result, (StatusCode, String)> { + let workspace = state.workspace.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Workspace not available".to_string(), + ))?; + + let doc = workspace + .read(&query.path) + .await + .map_err(|e| (StatusCode::NOT_FOUND, e.to_string()))?; + + Ok(Json(MemoryReadResponse { + path: query.path, + content: doc.content, + updated_at: Some(doc.updated_at.to_rfc3339()), + })) +} + +pub async fn memory_write_handler( + State(state): State>, + Json(req): Json, +) -> Result, (StatusCode, String)> { + let workspace = state.workspace.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Workspace not available".to_string(), + ))?; + + workspace + .write(&req.path, &req.content) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + Ok(Json(MemoryWriteResponse { + path: req.path, + status: "written", + })) +} + +pub async fn memory_search_handler( + State(state): State>, + Json(req): Json, +) -> Result, (StatusCode, String)> { + let workspace = state.workspace.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Workspace not available".to_string(), + ))?; + + let limit = req.limit.unwrap_or(10); + let results = workspace + .search(&req.query, limit) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let hits: Vec = results + .iter() + .map(|r| SearchHit { + path: r.document_id.to_string(), + content: r.content.clone(), + score: r.score as f64, + }) + .collect(); + + Ok(Json(MemorySearchResponse { results: hits })) +} diff --git a/src/channels/web/handlers/mod.rs b/src/channels/web/handlers/mod.rs new file mode 100644 index 00000000..88cd3d91 --- /dev/null +++ b/src/channels/web/handlers/mod.rs @@ -0,0 +1,23 @@ +//! Handler modules for the web gateway API. +//! +//! Each module groups related endpoint handlers by domain. + +pub mod chat; +pub mod extensions; +pub mod jobs; +pub mod memory; +pub mod routines; +pub mod settings; +pub mod skills; +pub mod static_files; + +// Re-export all handler functions so `server.rs` can reference them +// as `handlers::chat_send_handler`, etc. +pub use chat::*; +pub use extensions::*; +pub use jobs::*; +pub use memory::*; +pub use routines::*; +pub use settings::*; +pub use skills::*; +pub use static_files::*; diff --git a/src/channels/web/handlers/routines.rs b/src/channels/web/handlers/routines.rs new file mode 100644 index 00000000..be25681b --- /dev/null +++ b/src/channels/web/handlers/routines.rs @@ -0,0 +1,330 @@ +//! Routine management API handlers. + +use std::sync::Arc; + +use axum::{ + Json, + extract::{Path, State}, + http::StatusCode, +}; +use serde::Deserialize; +use uuid::Uuid; + +use crate::channels::IncomingMessage; +use crate::channels::web::server::GatewayState; +use crate::channels::web::types::*; + +pub async fn routines_list_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let routines = store + .list_routines(&state.user_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let items: Vec = routines.iter().map(routine_to_info).collect(); + + Ok(Json(RoutineListResponse { routines: items })) +} + +pub async fn routines_summary_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let routines = store + .list_routines(&state.user_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let total = routines.len() as u64; + let enabled = routines.iter().filter(|r| r.enabled).count() as u64; + let disabled = total - enabled; + let failing = routines + .iter() + .filter(|r| r.consecutive_failures > 0) + .count() as u64; + + let today_start = chrono::Utc::now() + .date_naive() + .and_hms_opt(0, 0, 0) + .map(|dt| dt.and_utc()); + let runs_today = if let Some(start) = today_start { + routines + .iter() + .filter(|r| r.last_run_at.is_some_and(|ts| ts >= start)) + .count() as u64 + } else { + 0 + }; + + Ok(Json(RoutineSummaryResponse { + total, + enabled, + disabled, + failing, + runs_today, + })) +} + +pub async fn routines_detail_handler( + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let routine_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?; + + let routine = store + .get_routine(routine_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?; + + let runs = store + .list_routine_runs(routine_id, 20) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let recent_runs: Vec = runs + .iter() + .map(|run| RoutineRunInfo { + id: run.id, + trigger_type: run.trigger_type.clone(), + started_at: run.started_at.to_rfc3339(), + completed_at: run.completed_at.map(|dt| dt.to_rfc3339()), + status: format!("{:?}", run.status), + result_summary: run.result_summary.clone(), + tokens_used: run.tokens_used, + }) + .collect(); + + Ok(Json(RoutineDetailResponse { + id: routine.id, + name: routine.name.clone(), + description: routine.description.clone(), + enabled: routine.enabled, + trigger: serde_json::to_value(&routine.trigger).unwrap_or_default(), + action: serde_json::to_value(&routine.action).unwrap_or_default(), + guardrails: serde_json::to_value(&routine.guardrails).unwrap_or_default(), + notify: serde_json::to_value(&routine.notify).unwrap_or_default(), + last_run_at: routine.last_run_at.map(|dt| dt.to_rfc3339()), + next_fire_at: routine.next_fire_at.map(|dt| dt.to_rfc3339()), + run_count: routine.run_count, + consecutive_failures: routine.consecutive_failures, + created_at: routine.created_at.to_rfc3339(), + recent_runs, + })) +} + +pub async fn routines_trigger_handler( + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let routine_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?; + + let routine = store + .get_routine(routine_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?; + + // Send the routine prompt through the message pipeline as a manual trigger. + let prompt = match &routine.action { + crate::agent::routine::RoutineAction::Lightweight { prompt, .. } => prompt.clone(), + crate::agent::routine::RoutineAction::FullJob { + title, description, .. + } => format!("{}: {}", title, description), + }; + + let content = format!("[routine:{}] {}", routine.name, prompt); + let msg = IncomingMessage::new("gateway", &state.user_id, content); + + let tx_guard = state.msg_tx.read().await; + let tx = tx_guard.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Channel not started".to_string(), + ))?; + + tx.send(msg).await.map_err(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + "Channel closed".to_string(), + ) + })?; + + Ok(Json(serde_json::json!({ + "status": "triggered", + "routine_id": routine_id, + }))) +} + +#[derive(Deserialize)] +pub struct ToggleRequest { + pub enabled: Option, +} + +pub async fn routines_toggle_handler( + State(state): State>, + Path(id): Path, + body: Option>, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let routine_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?; + + let mut routine = store + .get_routine(routine_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .ok_or((StatusCode::NOT_FOUND, "Routine not found".to_string()))?; + + // If a specific value was provided, use it; otherwise toggle. + routine.enabled = match body { + Some(Json(req)) => req.enabled.unwrap_or(!routine.enabled), + None => !routine.enabled, + }; + + store + .update_routine(&routine) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + Ok(Json(serde_json::json!({ + "status": if routine.enabled { "enabled" } else { "disabled" }, + "routine_id": routine_id, + }))) +} + +pub async fn routines_delete_handler( + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let routine_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?; + + let deleted = store + .delete_routine(routine_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + if deleted { + Ok(Json(serde_json::json!({ + "status": "deleted", + "routine_id": routine_id, + }))) + } else { + Err((StatusCode::NOT_FOUND, "Routine not found".to_string())) + } +} + +pub async fn routines_runs_handler( + State(state): State>, + Path(id): Path, +) -> Result, (StatusCode, String)> { + let store = state.store.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Database not available".to_string(), + ))?; + + let routine_id = Uuid::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid routine ID".to_string()))?; + + let runs = store + .list_routine_runs(routine_id, 50) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let run_infos: Vec = runs + .iter() + .map(|run| RoutineRunInfo { + id: run.id, + trigger_type: run.trigger_type.clone(), + started_at: run.started_at.to_rfc3339(), + completed_at: run.completed_at.map(|dt| dt.to_rfc3339()), + status: format!("{:?}", run.status), + result_summary: run.result_summary.clone(), + tokens_used: run.tokens_used, + }) + .collect(); + + Ok(Json(serde_json::json!({ + "routine_id": routine_id, + "runs": run_infos, + }))) +} + +/// Convert a Routine to the trimmed RoutineInfo for list display. +fn routine_to_info(r: &crate::agent::routine::Routine) -> RoutineInfo { + let (trigger_type, trigger_summary) = match &r.trigger { + crate::agent::routine::Trigger::Cron { schedule } => { + ("cron".to_string(), format!("cron: {}", schedule)) + } + crate::agent::routine::Trigger::Event { + pattern, channel, .. + } => { + let ch = channel.as_deref().unwrap_or("any"); + ("event".to_string(), format!("on {} /{}/", ch, pattern)) + } + crate::agent::routine::Trigger::Webhook { path, .. } => { + let p = path.as_deref().unwrap_or("/"); + ("webhook".to_string(), format!("webhook: {}", p)) + } + crate::agent::routine::Trigger::Manual => ("manual".to_string(), "manual only".to_string()), + }; + + let action_type = match &r.action { + crate::agent::routine::RoutineAction::Lightweight { .. } => "lightweight", + crate::agent::routine::RoutineAction::FullJob { .. } => "full_job", + }; + + let status = if !r.enabled { + "disabled" + } else if r.consecutive_failures > 0 { + "failing" + } else { + "active" + }; + + RoutineInfo { + id: r.id, + name: r.name.clone(), + description: r.description.clone(), + enabled: r.enabled, + trigger_type, + trigger_summary, + action_type: action_type.to_string(), + last_run_at: r.last_run_at.map(|dt| dt.to_rfc3339()), + next_fire_at: r.next_fire_at.map(|dt| dt.to_rfc3339()), + run_count: r.run_count, + consecutive_failures: r.consecutive_failures, + status: status.to_string(), + } +} diff --git a/src/channels/web/handlers/settings.rs b/src/channels/web/handlers/settings.rs new file mode 100644 index 00000000..dd66027b --- /dev/null +++ b/src/channels/web/handlers/settings.rs @@ -0,0 +1,133 @@ +//! Settings API handlers. + +use std::sync::Arc; + +use axum::{ + Json, + extract::{Path, State}, + http::StatusCode, +}; + +use crate::channels::web::server::GatewayState; +use crate::channels::web::types::*; + +pub async fn settings_list_handler( + State(state): State>, +) -> Result, StatusCode> { + let store = state + .store + .as_ref() + .ok_or(StatusCode::SERVICE_UNAVAILABLE)?; + let rows = store.list_settings(&state.user_id).await.map_err(|e| { + tracing::error!("Failed to list settings: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let settings = rows + .into_iter() + .map(|r| SettingResponse { + key: r.key, + value: r.value, + updated_at: r.updated_at.to_rfc3339(), + }) + .collect(); + + Ok(Json(SettingsListResponse { settings })) +} + +pub async fn settings_get_handler( + State(state): State>, + Path(key): Path, +) -> Result, StatusCode> { + let store = state + .store + .as_ref() + .ok_or(StatusCode::SERVICE_UNAVAILABLE)?; + let row = store + .get_setting_full(&state.user_id, &key) + .await + .map_err(|e| { + tracing::error!("Failed to get setting '{}': {}", key, e); + StatusCode::INTERNAL_SERVER_ERROR + })? + .ok_or(StatusCode::NOT_FOUND)?; + + Ok(Json(SettingResponse { + key: row.key, + value: row.value, + updated_at: row.updated_at.to_rfc3339(), + })) +} + +pub async fn settings_set_handler( + State(state): State>, + Path(key): Path, + Json(body): Json, +) -> Result { + let store = state + .store + .as_ref() + .ok_or(StatusCode::SERVICE_UNAVAILABLE)?; + store + .set_setting(&state.user_id, &key, &body.value) + .await + .map_err(|e| { + tracing::error!("Failed to set setting '{}': {}", key, e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + Ok(StatusCode::NO_CONTENT) +} + +pub async fn settings_delete_handler( + State(state): State>, + Path(key): Path, +) -> Result { + let store = state + .store + .as_ref() + .ok_or(StatusCode::SERVICE_UNAVAILABLE)?; + store + .delete_setting(&state.user_id, &key) + .await + .map_err(|e| { + tracing::error!("Failed to delete setting '{}': {}", key, e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + Ok(StatusCode::NO_CONTENT) +} + +pub async fn settings_export_handler( + State(state): State>, +) -> Result, StatusCode> { + let store = state + .store + .as_ref() + .ok_or(StatusCode::SERVICE_UNAVAILABLE)?; + let settings = store.get_all_settings(&state.user_id).await.map_err(|e| { + tracing::error!("Failed to export settings: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + Ok(Json(SettingsExportResponse { settings })) +} + +pub async fn settings_import_handler( + State(state): State>, + Json(body): Json, +) -> Result { + let store = state + .store + .as_ref() + .ok_or(StatusCode::SERVICE_UNAVAILABLE)?; + store + .set_all_settings(&state.user_id, &body.settings) + .await + .map_err(|e| { + tracing::error!("Failed to import settings: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + Ok(StatusCode::NO_CONTENT) +} diff --git a/src/channels/web/handlers/skills.rs b/src/channels/web/handlers/skills.rs new file mode 100644 index 00000000..dc281e40 --- /dev/null +++ b/src/channels/web/handlers/skills.rs @@ -0,0 +1,257 @@ +//! Skills management API handlers. + +use std::sync::Arc; + +use axum::{ + Json, + extract::{Path, State}, + http::StatusCode, +}; + +use crate::channels::web::server::GatewayState; +use crate::channels::web::types::*; + +pub async fn skills_list_handler( + State(state): State>, +) -> Result, (StatusCode, String)> { + let registry = state.skill_registry.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Skills system not enabled".to_string(), + ))?; + + let guard = registry.read().map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Skill registry lock poisoned: {}", e), + ) + })?; + + let skills: Vec = guard + .skills() + .iter() + .map(|s| SkillInfo { + name: s.manifest.name.clone(), + description: s.manifest.description.clone(), + version: s.manifest.version.clone(), + trust: s.trust.to_string(), + source: format!("{:?}", s.source), + keywords: s.manifest.activation.keywords.clone(), + }) + .collect(); + + let count = skills.len(); + Ok(Json(SkillListResponse { skills, count })) +} + +pub async fn skills_search_handler( + State(state): State>, + Json(req): Json, +) -> Result, (StatusCode, String)> { + let registry = state.skill_registry.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Skills system not enabled".to_string(), + ))?; + + let catalog = state.skill_catalog.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Skill catalog not available".to_string(), + ))?; + + // Search ClawHub catalog + let catalog_results = catalog.search(&req.query).await; + let catalog_json: Vec = catalog_results + .into_iter() + .map(|e| { + serde_json::json!({ + "slug": e.slug, + "name": e.name, + "description": e.description, + "version": e.version, + "score": e.score, + }) + }) + .collect(); + + // Search local skills + let query_lower = req.query.to_lowercase(); + let installed: Vec = { + let guard = registry.read().map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Skill registry lock poisoned: {}", e), + ) + })?; + guard + .skills() + .iter() + .filter(|s| { + s.manifest.name.to_lowercase().contains(&query_lower) + || s.manifest.description.to_lowercase().contains(&query_lower) + }) + .map(|s| SkillInfo { + name: s.manifest.name.clone(), + description: s.manifest.description.clone(), + version: s.manifest.version.clone(), + trust: s.trust.to_string(), + source: format!("{:?}", s.source), + keywords: s.manifest.activation.keywords.clone(), + }) + .collect() + }; + + Ok(Json(SkillSearchResponse { + catalog: catalog_json, + installed, + registry_url: catalog.registry_url().to_string(), + })) +} + +pub async fn skills_install_handler( + State(state): State>, + headers: axum::http::HeaderMap, + Json(req): Json, +) -> Result, (StatusCode, String)> { + // Require explicit confirmation header to prevent accidental installs. + // Chat tools have requires_approval(); this is the equivalent for the web API. + if headers + .get("x-confirm-action") + .and_then(|v| v.to_str().ok()) + != Some("true") + { + return Err(( + StatusCode::BAD_REQUEST, + "Skill install requires X-Confirm-Action: true header".to_string(), + )); + } + + let registry = state.skill_registry.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Skills system not enabled".to_string(), + ))?; + + let content = if let Some(ref raw) = req.content { + raw.clone() + } else if let Some(ref url) = req.url { + // Fetch from explicit URL (with SSRF protection) + crate::tools::builtin::skill_tools::fetch_skill_content(url) + .await + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))? + } else if let Some(ref catalog) = state.skill_catalog { + let url = crate::skills::catalog::skill_download_url(catalog.registry_url(), &req.name); + crate::tools::builtin::skill_tools::fetch_skill_content(&url) + .await + .map_err(|e| (StatusCode::BAD_GATEWAY, e.to_string()))? + } else { + return Ok(Json(ActionResponse::fail( + "Provide 'content' or 'url' to install a skill".to_string(), + ))); + }; + + // Parse, check duplicates, and get user_dir under a brief read lock. + let (user_dir, skill_name_from_parse) = { + let guard = registry.read().map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Skill registry lock poisoned: {}", e), + ) + })?; + + let normalized = crate::skills::normalize_line_endings(&content); + let parsed = crate::skills::parser::parse_skill_md(&normalized) + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; + let skill_name = parsed.manifest.name.clone(); + + if guard.has(&skill_name) { + return Ok(Json(ActionResponse::fail(format!( + "Skill '{}' already exists", + skill_name + )))); + } + + (guard.user_dir().to_path_buf(), skill_name) + }; + + // Perform async I/O (write to disk, load) with no lock held. + let normalized = crate::skills::normalize_line_endings(&content); + let (skill_name, loaded_skill) = + crate::skills::registry::SkillRegistry::prepare_install_to_disk( + &user_dir, + &skill_name_from_parse, + &normalized, + ) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + // Commit: brief write lock for in-memory addition + let mut guard = registry.write().map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Skill registry lock poisoned: {}", e), + ) + })?; + + match guard.commit_install(&skill_name, loaded_skill) { + Ok(()) => Ok(Json(ActionResponse::ok(format!( + "Skill '{}' installed", + skill_name + )))), + Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))), + } +} + +pub async fn skills_remove_handler( + State(state): State>, + headers: axum::http::HeaderMap, + Path(name): Path, +) -> Result, (StatusCode, String)> { + // Require explicit confirmation header to prevent accidental removals. + if headers + .get("x-confirm-action") + .and_then(|v| v.to_str().ok()) + != Some("true") + { + return Err(( + StatusCode::BAD_REQUEST, + "Skill removal requires X-Confirm-Action: true header".to_string(), + )); + } + + let registry = state.skill_registry.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Skills system not enabled".to_string(), + ))?; + + // Validate removal under a brief read lock + let skill_path = { + let guard = registry.read().map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Skill registry lock poisoned: {}", e), + ) + })?; + guard + .validate_remove(&name) + .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))? + }; + + // Delete files from disk (async I/O, no lock held) + crate::skills::registry::SkillRegistry::delete_skill_files(&skill_path) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + // Remove from in-memory registry under a brief write lock + let mut guard = registry.write().map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Skill registry lock poisoned: {}", e), + ) + })?; + + match guard.commit_remove(&name) { + Ok(()) => Ok(Json(ActionResponse::ok(format!( + "Skill '{}' removed", + name + )))), + Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))), + } +} diff --git a/src/channels/web/handlers/static_files.rs b/src/channels/web/handlers/static_files.rs new file mode 100644 index 00000000..cd0eeece --- /dev/null +++ b/src/channels/web/handlers/static_files.rs @@ -0,0 +1,178 @@ +//! Static file and health handlers. + +use axum::{ + Json, + http::{StatusCode, header}, + response::{Html, IntoResponse}, +}; + +use crate::channels::web::types::*; + +// --- Static file handlers --- + +pub async fn index_handler() -> Html<&'static str> { + Html(include_str!("../static/index.html")) +} + +pub async fn css_handler() -> impl IntoResponse { + ( + [(header::CONTENT_TYPE, "text/css")], + include_str!("../static/style.css"), + ) +} + +pub async fn js_handler() -> impl IntoResponse { + ( + [(header::CONTENT_TYPE, "application/javascript")], + include_str!("../static/app.js"), + ) +} + +// --- Health --- + +pub async fn health_handler() -> Json { + Json(HealthResponse { + status: "healthy", + channel: "gateway", + }) +} + +// --- Project file serving handlers --- + +use axum::extract::Path; + +/// Redirect `/projects/{id}` to `/projects/{id}/` so relative paths in +/// the served HTML resolve within the project namespace. +pub async fn project_redirect_handler(Path(project_id): Path) -> impl IntoResponse { + axum::response::Redirect::permanent(&format!("/projects/{project_id}/")) +} + +/// Serve `index.html` when hitting `/projects/{project_id}/`. +pub async fn project_index_handler(Path(project_id): Path) -> impl IntoResponse { + serve_project_file(&project_id, "index.html").await +} + +/// Serve any file under `/projects/{project_id}/{path}`. +pub async fn project_file_handler( + Path((project_id, path)): Path<(String, String)>, +) -> impl IntoResponse { + serve_project_file(&project_id, &path).await +} + +/// Shared logic: resolve the file inside `~/.ironclaw/projects/{project_id}/`, +/// guard against path traversal, and stream the content with the right MIME type. +async fn serve_project_file(project_id: &str, path: &str) -> axum::response::Response { + // Reject project_id values that could escape the projects directory. + if project_id.contains('/') + || project_id.contains('\\') + || project_id.contains("..") + || project_id.is_empty() + { + return (StatusCode::BAD_REQUEST, "Invalid project ID").into_response(); + } + + let base = dirs::home_dir() + .unwrap_or_else(|| std::path::PathBuf::from(".")) + .join(".ironclaw") + .join("projects") + .join(project_id); + + let file_path = base.join(path); + + // Path traversal guard + let canonical = match file_path.canonicalize() { + Ok(p) => p, + Err(_) => return (StatusCode::NOT_FOUND, "Not found").into_response(), + }; + let base_canonical = match base.canonicalize() { + Ok(p) => p, + Err(_) => return (StatusCode::NOT_FOUND, "Not found").into_response(), + }; + if !canonical.starts_with(&base_canonical) { + return (StatusCode::FORBIDDEN, "Forbidden").into_response(); + } + + match tokio::fs::read(&canonical).await { + Ok(contents) => { + let mime = mime_guess::from_path(&canonical) + .first_or_octet_stream() + .to_string(); + ([(header::CONTENT_TYPE, mime)], contents).into_response() + } + Err(_) => (StatusCode::NOT_FOUND, "Not found").into_response(), + } +} + +// --- Logs --- + +use std::convert::Infallible; +use std::sync::Arc; + +use axum::extract::State; +use axum::response::sse::{Event, KeepAlive, Sse}; +use tokio_stream::StreamExt; + +use crate::channels::web::server::GatewayState; + +pub async fn logs_events_handler( + State(state): State>, +) -> Result< + Sse> + Send + 'static>, + (StatusCode, String), +> { + let broadcaster = state.log_broadcaster.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Log broadcaster not available".to_string(), + ))?; + + // Replay recent history so late-joining browsers see startup logs. + // Subscribe BEFORE snapshotting to avoid a gap between history and live. + let rx = broadcaster.subscribe(); + let history = broadcaster.recent_entries(); + + let history_stream = futures::stream::iter(history).map(|entry| { + let data = serde_json::to_string(&entry).unwrap_or_default(); + Ok(Event::default().event("log").data(data)) + }); + + let live_stream = tokio_stream::wrappers::BroadcastStream::new(rx) + .filter_map(|result| result.ok()) + .map(|entry| { + let data = serde_json::to_string(&entry).unwrap_or_default(); + Ok(Event::default().event("log").data(data)) + }); + + let stream = history_stream.chain(live_stream); + + Ok(Sse::new(stream).keep_alive( + KeepAlive::new() + .interval(std::time::Duration::from_secs(30)) + .text(""), + )) +} + +// --- Gateway status --- + +pub async fn gateway_status_handler( + State(state): State>, +) -> Json { + let sse_connections = state.sse.connection_count(); + let ws_connections = state + .ws_tracker + .as_ref() + .map(|t| t.connection_count()) + .unwrap_or(0); + + Json(GatewayStatusResponse { + sse_connections, + ws_connections, + total_connections: sse_connections + ws_connections, + }) +} + +#[derive(serde::Serialize)] +pub struct GatewayStatusResponse { + pub sse_connections: u64, + pub ws_connections: u64, + pub total_connections: u64, +} diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index e61b65de..dc9cb99e 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -519,7 +519,7 @@ async fn get_secrets_store() -> anyhow::Result, user_id: String) -> anyho #[cfg(all(feature = "libsql", not(feature = "postgres")))] { use crate::db::Database as _; - use crate::db::libsql_backend::LibSqlBackend; + use crate::db::libsql::LibSqlBackend; use secrecy::ExposeSecret as _; let default_path = crate::config::default_libsql_path(); diff --git a/src/config.rs b/src/config.rs deleted file mode 100644 index 68951ffa..00000000 --- a/src/config.rs +++ /dev/null @@ -1,1944 +0,0 @@ -//! Configuration for IronClaw. -//! -//! Settings are loaded with priority: env var > database > default. -//! `DATABASE_URL` lives in `~/.ironclaw/.env` (loaded via dotenvy early -//! in startup). Everything else comes from env vars, the DB settings -//! table, or auto-detection. - -use std::collections::HashMap; -use std::path::PathBuf; -use std::sync::OnceLock; -use std::time::Duration; - -use secrecy::{ExposeSecret, SecretString}; - -use crate::error::ConfigError; -use crate::settings::Settings; - -/// Thread-safe overlay for injected env vars (secrets loaded from DB). -/// -/// Used by `inject_llm_keys_from_secrets()` to make API keys available to -/// `optional_env()` without unsafe `set_var` calls. `optional_env()` checks -/// real env vars first, then falls back to this overlay. -static INJECTED_VARS: OnceLock> = OnceLock::new(); - -/// Main configuration for the agent. -#[derive(Debug, Clone)] -pub struct Config { - pub database: DatabaseConfig, - pub llm: LlmConfig, - pub embeddings: EmbeddingsConfig, - pub tunnel: TunnelConfig, - pub channels: ChannelsConfig, - pub agent: AgentConfig, - pub safety: SafetyConfig, - pub wasm: WasmConfig, - pub secrets: SecretsConfig, - pub builder: BuilderModeConfig, - pub heartbeat: HeartbeatConfig, - pub routines: RoutineConfig, - pub sandbox: SandboxModeConfig, - pub claude_code: ClaudeCodeConfig, - pub skills: SkillsConfig, - pub observability: crate::observability::ObservabilityConfig, -} - -impl Config { - /// Load configuration from environment variables and the database. - /// - /// Priority: env var > TOML config file > DB settings > default. - /// This is the primary way to load config after DB is connected. - pub async fn from_db( - store: &dyn crate::db::Database, - user_id: &str, - ) -> Result { - Self::from_db_with_toml(store, user_id, None).await - } - - /// Load from DB with an optional TOML config file overlay. - pub async fn from_db_with_toml( - store: &dyn crate::db::Database, - user_id: &str, - toml_path: Option<&std::path::Path>, - ) -> Result { - let _ = dotenvy::dotenv(); - crate::bootstrap::load_ironclaw_env(); - - // Load all settings from DB into a Settings struct - let mut db_settings = match store.get_all_settings(user_id).await { - Ok(map) => Settings::from_db_map(&map), - Err(e) => { - tracing::warn!("Failed to load settings from DB, using defaults: {}", e); - Settings::default() - } - }; - - // Overlay TOML config file (values win over DB settings) - Self::apply_toml_overlay(&mut db_settings, toml_path)?; - - Self::build(&db_settings).await - } - - /// Load configuration from environment variables only (no database). - /// - /// Used during early startup before the database is connected, - /// and by CLI commands that don't have DB access. - /// Falls back to legacy `settings.json` on disk if present. - /// - /// Loads both `./.env` (standard, higher priority) and `~/.ironclaw/.env` - /// (lower priority) via dotenvy, which never overwrites existing vars. - pub async fn from_env() -> Result { - Self::from_env_with_toml(None).await - } - - /// Load from env with an optional TOML config file overlay. - pub async fn from_env_with_toml( - toml_path: Option<&std::path::Path>, - ) -> Result { - let _ = dotenvy::dotenv(); - crate::bootstrap::load_ironclaw_env(); - let mut settings = Settings::load(); - - // Overlay TOML config file (values win over JSON settings) - Self::apply_toml_overlay(&mut settings, toml_path)?; - - Self::build(&settings).await - } - - /// Load and merge a TOML config file into settings. - /// - /// If `explicit_path` is `Some`, loads from that path (errors are fatal). - /// If `None`, tries the default path `~/.ironclaw/config.toml` (missing - /// file is silently ignored). - fn apply_toml_overlay( - settings: &mut Settings, - explicit_path: Option<&std::path::Path>, - ) -> Result<(), ConfigError> { - let path = explicit_path - .map(std::path::PathBuf::from) - .unwrap_or_else(Settings::default_toml_path); - - match Settings::load_toml(&path) { - Ok(Some(toml_settings)) => { - settings.merge_from(&toml_settings); - tracing::debug!("Loaded TOML config from {}", path.display()); - } - Ok(None) => { - if explicit_path.is_some() { - return Err(ConfigError::ParseError(format!( - "Config file not found: {}", - path.display() - ))); - } - } - Err(e) => { - if explicit_path.is_some() { - return Err(ConfigError::ParseError(format!( - "Failed to load config file {}: {}", - path.display(), - e - ))); - } - tracing::warn!("Failed to load default config file: {}", e); - } - } - Ok(()) - } - - /// Build config from settings (shared by from_env and from_db). - async fn build(settings: &Settings) -> Result { - Ok(Self { - database: DatabaseConfig::resolve()?, - llm: LlmConfig::resolve(settings)?, - embeddings: EmbeddingsConfig::resolve(settings)?, - tunnel: TunnelConfig::resolve(settings)?, - channels: ChannelsConfig::resolve(settings)?, - agent: AgentConfig::resolve(settings)?, - safety: SafetyConfig::resolve()?, - wasm: WasmConfig::resolve()?, - secrets: SecretsConfig::resolve().await?, - builder: BuilderModeConfig::resolve()?, - heartbeat: HeartbeatConfig::resolve(settings)?, - routines: RoutineConfig::resolve()?, - sandbox: SandboxModeConfig::resolve()?, - claude_code: ClaudeCodeConfig::resolve()?, - skills: SkillsConfig::resolve()?, - observability: crate::observability::ObservabilityConfig { - backend: std::env::var("OBSERVABILITY_BACKEND").unwrap_or_else(|_| "none".into()), - }, - }) - } -} - -/// Tunnel configuration for exposing the agent to the internet. -/// -/// Used by channels and tools that need public webhook endpoints. -/// The tunnel URL is shared across all channels (Telegram, Slack, etc.). -/// -/// Two modes: -/// - **Static URL** (`TUNNEL_URL`): set the public URL directly (manual tunnel) -/// - **Managed provider** (`TUNNEL_PROVIDER`): lifecycle-managed tunnel process -/// -/// When a managed provider is configured _and_ no static URL is set, -/// the gateway starts the tunnel on boot and populates `public_url`. -#[derive(Debug, Clone, Default)] -pub struct TunnelConfig { - /// Public URL from tunnel provider (e.g., "https://abc123.ngrok.io"). - /// Set statically via `TUNNEL_URL` or populated at runtime by a managed tunnel. - pub public_url: Option, - /// Provider configuration for lifecycle-managed tunnels. - /// `None` when using a static URL or no tunnel at all. - pub provider: Option, -} - -impl TunnelConfig { - fn resolve(settings: &Settings) -> Result { - let public_url = optional_env("TUNNEL_URL")? - .or_else(|| settings.tunnel.public_url.clone().filter(|s| !s.is_empty())); - - if let Some(ref url) = public_url - && !url.starts_with("https://") - { - return Err(ConfigError::InvalidValue { - key: "TUNNEL_URL".to_string(), - message: "must start with https:// (webhooks require HTTPS)".to_string(), - }); - } - - // Resolve managed tunnel provider config. - // Priority: env var > settings > default (none). - let provider_name = optional_env("TUNNEL_PROVIDER")? - .or_else(|| settings.tunnel.provider.clone()) - .unwrap_or_default(); - - let provider = if provider_name.is_empty() || provider_name == "none" { - None - } else { - Some(crate::tunnel::TunnelProviderConfig { - provider: provider_name.clone(), - cloudflare: optional_env("TUNNEL_CF_TOKEN")? - .or_else(|| settings.tunnel.cf_token.clone()) - .map(|token| crate::tunnel::CloudflareTunnelConfig { token }), - tailscale: Some(crate::tunnel::TailscaleTunnelConfig { - funnel: optional_env("TUNNEL_TS_FUNNEL") - .ok() - .flatten() - .map(|s| s == "true" || s == "1") - .unwrap_or(settings.tunnel.ts_funnel), - hostname: optional_env("TUNNEL_TS_HOSTNAME") - .ok() - .flatten() - .or_else(|| settings.tunnel.ts_hostname.clone()), - }), - ngrok: optional_env("TUNNEL_NGROK_TOKEN")? - .or_else(|| settings.tunnel.ngrok_token.clone()) - .map(|auth_token| crate::tunnel::NgrokTunnelConfig { - auth_token, - domain: optional_env("TUNNEL_NGROK_DOMAIN") - .ok() - .flatten() - .or_else(|| settings.tunnel.ngrok_domain.clone()), - }), - custom: optional_env("TUNNEL_CUSTOM_COMMAND")? - .or_else(|| settings.tunnel.custom_command.clone()) - .map(|start_command| crate::tunnel::CustomTunnelConfig { - start_command, - health_url: optional_env("TUNNEL_CUSTOM_HEALTH_URL") - .ok() - .flatten() - .or_else(|| settings.tunnel.custom_health_url.clone()), - url_pattern: optional_env("TUNNEL_CUSTOM_URL_PATTERN") - .ok() - .flatten() - .or_else(|| settings.tunnel.custom_url_pattern.clone()), - }), - }) - }; - - Ok(Self { - public_url, - provider, - }) - } - - /// Check if a tunnel is configured (static URL or managed provider). - pub fn is_enabled(&self) -> bool { - self.public_url.is_some() || self.provider.is_some() - } - - /// Get the webhook URL for a given path. - pub fn webhook_url(&self, path: &str) -> Option { - self.public_url.as_ref().map(|base| { - let base = base.trim_end_matches('/'); - let path = path.trim_start_matches('/'); - format!("{}/{}", base, path) - }) - } -} - -/// Which database backend to use. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum DatabaseBackend { - /// PostgreSQL via deadpool-postgres (default). - #[default] - Postgres, - /// libSQL/Turso embedded database. - LibSql, -} - -impl std::fmt::Display for DatabaseBackend { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Postgres => write!(f, "postgres"), - Self::LibSql => write!(f, "libsql"), - } - } -} - -impl std::str::FromStr for DatabaseBackend { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "postgres" | "postgresql" | "pg" => Ok(Self::Postgres), - "libsql" | "turso" | "sqlite" => Ok(Self::LibSql), - _ => Err(format!( - "invalid database backend '{}', expected 'postgres' or 'libsql'", - s - )), - } - } -} - -/// Database configuration. -#[derive(Debug, Clone)] -pub struct DatabaseConfig { - /// Which backend to use (default: Postgres). - pub backend: DatabaseBackend, - - // -- PostgreSQL fields -- - pub url: SecretString, - pub pool_size: usize, - - // -- libSQL fields -- - /// Path to local libSQL database file (default: ~/.ironclaw/ironclaw.db). - pub libsql_path: Option, - /// Turso cloud URL for remote sync (optional). - pub libsql_url: Option, - /// Turso auth token (required when libsql_url is set). - pub libsql_auth_token: Option, -} - -impl DatabaseConfig { - fn resolve() -> Result { - let backend: DatabaseBackend = if let Some(b) = optional_env("DATABASE_BACKEND")? { - b.parse().map_err(|e| ConfigError::InvalidValue { - key: "DATABASE_BACKEND".to_string(), - message: e, - })? - } else { - DatabaseBackend::default() - }; - - // PostgreSQL URL is required only when using the postgres backend. - // For libsql backend, default to an empty placeholder. - // DATABASE_URL is loaded from ~/.ironclaw/.env via dotenvy early in startup. - let url = optional_env("DATABASE_URL")? - .or_else(|| { - if backend == DatabaseBackend::LibSql { - Some("unused://libsql".to_string()) - } else { - None - } - }) - .ok_or_else(|| ConfigError::MissingRequired { - key: "database_url".to_string(), - hint: "Run 'ironclaw onboard' or set DATABASE_URL environment variable".to_string(), - })?; - - let pool_size = parse_optional_env("DATABASE_POOL_SIZE", 10)?; - - let libsql_path = optional_env("LIBSQL_PATH")?.map(PathBuf::from).or_else(|| { - if backend == DatabaseBackend::LibSql { - Some(default_libsql_path()) - } else { - None - } - }); - - let libsql_url = optional_env("LIBSQL_URL")?; - let libsql_auth_token = optional_env("LIBSQL_AUTH_TOKEN")?.map(SecretString::from); - - if libsql_url.is_some() && libsql_auth_token.is_none() { - return Err(ConfigError::MissingRequired { - key: "LIBSQL_AUTH_TOKEN".to_string(), - hint: "LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set".to_string(), - }); - } - - Ok(Self { - backend, - url: SecretString::from(url), - pool_size, - libsql_path, - libsql_url, - libsql_auth_token, - }) - } - - /// Get the database URL (exposes the secret). - pub fn url(&self) -> &str { - self.url.expose_secret() - } -} - -/// Default libSQL database path (~/.ironclaw/ironclaw.db). -pub fn default_libsql_path() -> PathBuf { - dirs::home_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join(".ironclaw") - .join("ironclaw.db") -} - -/// Which LLM backend to use. -/// -/// Defaults to `NearAi` to keep IronClaw close to the NEAR ecosystem. -/// Users can override with `LLM_BACKEND` env var to use their own API keys. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum LlmBackend { - /// NEAR AI proxy (default) -- session or API key auth - #[default] - NearAi, - /// Direct OpenAI API - OpenAi, - /// Direct Anthropic API - Anthropic, - /// Local Ollama instance - Ollama, - /// Any OpenAI-compatible endpoint (e.g. vLLM, LiteLLM, Together) - OpenAiCompatible, - /// Tinfoil private inference - Tinfoil, -} - -impl std::str::FromStr for LlmBackend { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "nearai" | "near_ai" | "near" => Ok(Self::NearAi), - "openai" | "open_ai" => Ok(Self::OpenAi), - "anthropic" | "claude" => Ok(Self::Anthropic), - "ollama" => Ok(Self::Ollama), - "openai_compatible" | "openai-compatible" | "compatible" => Ok(Self::OpenAiCompatible), - "tinfoil" => Ok(Self::Tinfoil), - _ => Err(format!( - "invalid LLM backend '{}', expected one of: nearai, openai, anthropic, ollama, openai_compatible, tinfoil", - s - )), - } - } -} - -impl std::fmt::Display for LlmBackend { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::NearAi => write!(f, "nearai"), - Self::OpenAi => write!(f, "openai"), - Self::Anthropic => write!(f, "anthropic"), - Self::Ollama => write!(f, "ollama"), - Self::OpenAiCompatible => write!(f, "openai_compatible"), - Self::Tinfoil => write!(f, "tinfoil"), - } - } -} - -/// Configuration for direct OpenAI API access. -#[derive(Debug, Clone)] -pub struct OpenAiDirectConfig { - pub api_key: SecretString, - pub model: String, -} - -/// Configuration for direct Anthropic API access. -#[derive(Debug, Clone)] -pub struct AnthropicDirectConfig { - pub api_key: SecretString, - pub model: String, -} - -/// Configuration for local Ollama. -#[derive(Debug, Clone)] -pub struct OllamaConfig { - pub base_url: String, - pub model: String, -} - -/// Configuration for any OpenAI-compatible endpoint. -#[derive(Debug, Clone)] -pub struct OpenAiCompatibleConfig { - pub base_url: String, - pub api_key: Option, - pub model: String, -} - -/// Configuration for Tinfoil private inference. -#[derive(Debug, Clone)] -pub struct TinfoilConfig { - pub api_key: SecretString, - pub model: String, -} - -/// LLM provider configuration. -/// -/// NEAR AI remains the default backend. Users can switch to other providers -/// by setting `LLM_BACKEND` (e.g. `openai`, `anthropic`, `ollama`). -#[derive(Debug, Clone)] -pub struct LlmConfig { - /// Which backend to use (default: NearAi) - pub backend: LlmBackend, - /// NEAR AI config (always populated for NEAR AI embeddings, etc.) - pub nearai: NearAiConfig, - /// Direct OpenAI config (populated when backend=openai) - pub openai: Option, - /// Direct Anthropic config (populated when backend=anthropic) - pub anthropic: Option, - /// Ollama config (populated when backend=ollama) - pub ollama: Option, - /// OpenAI-compatible config (populated when backend=openai_compatible) - pub openai_compatible: Option, - /// Tinfoil config (populated when backend=tinfoil) - pub tinfoil: Option, -} - -/// API mode for NEAR AI. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum NearAiApiMode { - /// Use the Responses API (chat-api proxy) - session-based auth - #[default] - Responses, - /// Use the Chat Completions API (cloud-api) - API key auth - ChatCompletions, -} - -impl std::str::FromStr for NearAiApiMode { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "responses" | "response" => Ok(Self::Responses), - "chat_completions" | "chatcompletions" | "chat" | "completions" => { - Ok(Self::ChatCompletions) - } - _ => Err(format!( - "invalid API mode '{}', expected 'responses' or 'chat_completions'", - s - )), - } - } -} - -/// NEAR AI chat-api configuration. -#[derive(Debug, Clone)] -pub struct NearAiConfig { - /// Model to use (e.g., "claude-3-5-sonnet-20241022", "gpt-4o") - pub model: String, - /// Cheap/fast model for lightweight tasks (heartbeat, routing, evaluation). - /// Falls back to the main model if not set. - pub cheap_model: Option, - /// Base URL for the NEAR AI API (default: https://private.near.ai). - pub base_url: String, - /// Base URL for auth/refresh endpoints (default: https://private.near.ai) - pub auth_base_url: String, - /// Path to session file (default: ~/.ironclaw/session.json) - pub session_path: PathBuf, - /// API mode: "responses" (chat-api) or "chat_completions" (cloud-api) - pub api_mode: NearAiApiMode, - /// API key for cloud-api (required for chat_completions mode) - pub api_key: Option, - /// Optional fallback model for failover (default: None). - /// When set, a secondary provider is created with this model and wrapped - /// in a `FailoverProvider` so transient errors on the primary model - /// automatically fall through to the fallback. - pub fallback_model: Option, - /// Maximum number of retries for transient errors (default: 3). - /// With the default of 3, the provider makes up to 4 total attempts - /// (1 initial + 3 retries) before giving up. - pub max_retries: u32, - /// Consecutive transient failures before the circuit breaker opens. - /// None = disabled (default). E.g. 5 means after 5 consecutive failures - /// all requests are rejected until recovery timeout elapses. - pub circuit_breaker_threshold: Option, - /// How long (seconds) the circuit stays open before allowing a probe (default: 30). - pub circuit_breaker_recovery_secs: u64, - /// Enable in-memory response caching for `complete()` calls. - /// Saves tokens on repeated prompts within a session. Default: false. - pub response_cache_enabled: bool, - /// TTL in seconds for cached responses (default: 3600 = 1 hour). - pub response_cache_ttl_secs: u64, - /// Max cached responses before LRU eviction (default: 1000). - pub response_cache_max_entries: usize, - /// Cooldown duration in seconds for the failover provider (default: 300). - /// When a provider accumulates enough consecutive failures it is skipped - /// for this many seconds. - pub failover_cooldown_secs: u64, - /// Number of consecutive retryable failures before a provider enters - /// cooldown (default: 3). - pub failover_cooldown_threshold: u32, -} - -impl LlmConfig { - fn resolve(settings: &Settings) -> Result { - // Determine backend: env var > settings > default (NearAi) - let backend: LlmBackend = if let Some(b) = optional_env("LLM_BACKEND")? { - b.parse().map_err(|e| ConfigError::InvalidValue { - key: "LLM_BACKEND".to_string(), - message: e, - })? - } else if let Some(ref b) = settings.llm_backend { - match b.parse() { - Ok(backend) => backend, - Err(e) => { - tracing::warn!( - "Invalid llm_backend '{}' in settings: {}. Using default NearAi.", - b, - e - ); - LlmBackend::NearAi - } - } - } else { - LlmBackend::NearAi - }; - - // Resolve NEAR AI config only when backend is NearAi (or when explicitly configured) - let nearai_api_key = optional_env("NEARAI_API_KEY")?.map(SecretString::from); - - let api_mode = if let Some(mode_str) = optional_env("NEARAI_API_MODE")? { - mode_str.parse().map_err(|e| ConfigError::InvalidValue { - key: "NEARAI_API_MODE".to_string(), - message: e, - })? - } else if nearai_api_key.is_some() { - NearAiApiMode::ChatCompletions - } else { - NearAiApiMode::Responses - }; - - 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() - }), - cheap_model: optional_env("NEARAI_CHEAP_MODEL")?, - base_url: optional_env("NEARAI_BASE_URL")? - .unwrap_or_else(|| "https://private.near.ai".to_string()), - auth_base_url: optional_env("NEARAI_AUTH_URL")? - .unwrap_or_else(|| "https://private.near.ai".to_string()), - session_path: optional_env("NEARAI_SESSION_PATH")? - .map(PathBuf::from) - .unwrap_or_else(default_session_path), - api_mode, - api_key: nearai_api_key, - fallback_model: optional_env("NEARAI_FALLBACK_MODEL")?, - max_retries: parse_optional_env("NEARAI_MAX_RETRIES", 3)?, - circuit_breaker_threshold: optional_env("CIRCUIT_BREAKER_THRESHOLD")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "CIRCUIT_BREAKER_THRESHOLD".to_string(), - message: format!("must be a positive integer: {e}"), - })?, - circuit_breaker_recovery_secs: parse_optional_env("CIRCUIT_BREAKER_RECOVERY_SECS", 30)?, - response_cache_enabled: parse_optional_env("RESPONSE_CACHE_ENABLED", false)?, - response_cache_ttl_secs: parse_optional_env("RESPONSE_CACHE_TTL_SECS", 3600)?, - response_cache_max_entries: parse_optional_env("RESPONSE_CACHE_MAX_ENTRIES", 1000)?, - failover_cooldown_secs: parse_optional_env("LLM_FAILOVER_COOLDOWN_SECS", 300)?, - failover_cooldown_threshold: parse_optional_env("LLM_FAILOVER_THRESHOLD", 3)?, - }; - - // Resolve provider-specific configs based on backend - let openai = if backend == LlmBackend::OpenAi { - let api_key = optional_env("OPENAI_API_KEY")? - .map(SecretString::from) - .ok_or_else(|| ConfigError::MissingRequired { - key: "OPENAI_API_KEY".to_string(), - hint: "Set OPENAI_API_KEY when LLM_BACKEND=openai".to_string(), - })?; - let model = optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4o".to_string()); - Some(OpenAiDirectConfig { api_key, model }) - } else { - None - }; - - let anthropic = if backend == LlmBackend::Anthropic { - let api_key = optional_env("ANTHROPIC_API_KEY")? - .map(SecretString::from) - .ok_or_else(|| ConfigError::MissingRequired { - key: "ANTHROPIC_API_KEY".to_string(), - hint: "Set ANTHROPIC_API_KEY when LLM_BACKEND=anthropic".to_string(), - })?; - let model = optional_env("ANTHROPIC_MODEL")? - .unwrap_or_else(|| "claude-sonnet-4-20250514".to_string()); - Some(AnthropicDirectConfig { api_key, model }) - } else { - None - }; - - let ollama = if backend == LlmBackend::Ollama { - let base_url = optional_env("OLLAMA_BASE_URL")? - .or_else(|| settings.ollama_base_url.clone()) - .unwrap_or_else(|| "http://localhost:11434".to_string()); - let model = optional_env("OLLAMA_MODEL")?.unwrap_or_else(|| "llama3".to_string()); - Some(OllamaConfig { base_url, model }) - } else { - None - }; - - let openai_compatible = if backend == LlmBackend::OpenAiCompatible { - let base_url = optional_env("LLM_BASE_URL")? - .or_else(|| settings.openai_compatible_base_url.clone()) - .ok_or_else(|| ConfigError::MissingRequired { - key: "LLM_BASE_URL".to_string(), - hint: "Set LLM_BASE_URL when LLM_BACKEND=openai_compatible".to_string(), - })?; - let api_key = optional_env("LLM_API_KEY")?.map(SecretString::from); - let model = optional_env("LLM_MODEL")? - .or_else(|| settings.selected_model.clone()) - .unwrap_or_else(|| "default".to_string()); - Some(OpenAiCompatibleConfig { - base_url, - api_key, - model, - }) - } else { - None - }; - - let tinfoil = if backend == LlmBackend::Tinfoil { - let api_key = optional_env("TINFOIL_API_KEY")? - .map(SecretString::from) - .ok_or_else(|| ConfigError::MissingRequired { - key: "TINFOIL_API_KEY".to_string(), - hint: "Set TINFOIL_API_KEY when LLM_BACKEND=tinfoil".to_string(), - })?; - let model = optional_env("TINFOIL_MODEL")?.unwrap_or_else(|| "kimi-k2-5".to_string()); - Some(TinfoilConfig { api_key, model }) - } else { - None - }; - - Ok(Self { - backend, - nearai, - openai, - anthropic, - ollama, - openai_compatible, - tinfoil, - }) - } -} - -/// Embeddings provider configuration. -#[derive(Debug, Clone)] -pub struct EmbeddingsConfig { - /// Whether embeddings are enabled. - pub enabled: bool, - /// Provider to use: "openai" or "nearai" - pub provider: String, - /// OpenAI API key (for OpenAI provider). - pub openai_api_key: Option, - /// Model to use for embeddings. - pub model: String, -} - -impl Default for EmbeddingsConfig { - fn default() -> Self { - Self { - enabled: false, - provider: "openai".to_string(), - openai_api_key: None, - model: "text-embedding-3-small".to_string(), - } - } -} - -impl EmbeddingsConfig { - fn resolve(settings: &Settings) -> Result { - let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from); - - let provider = optional_env("EMBEDDING_PROVIDER")? - .unwrap_or_else(|| settings.embeddings.provider.clone()); - - let model = - optional_env("EMBEDDING_MODEL")?.unwrap_or_else(|| settings.embeddings.model.clone()); - - 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); - - Ok(Self { - enabled, - provider, - openai_api_key, - model, - }) - } - - /// Get the OpenAI API key if configured. - pub fn openai_api_key(&self) -> Option<&str> { - self.openai_api_key.as_ref().map(|s| s.expose_secret()) - } -} - -/// Get the default session file path (~/.ironclaw/session.json). -fn default_session_path() -> PathBuf { - dirs::home_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join(".ironclaw") - .join("session.json") -} - -/// Channel configurations. -#[derive(Debug, Clone)] -pub struct ChannelsConfig { - pub cli: CliConfig, - pub http: Option, - pub gateway: Option, - /// Directory containing WASM channel modules (default: ~/.ironclaw/channels/). - pub wasm_channels_dir: std::path::PathBuf, - /// Whether WASM channels are enabled. - pub wasm_channels_enabled: bool, - /// Telegram owner user ID. When set, the bot only responds to this user. - pub telegram_owner_id: Option, -} - -#[derive(Debug, Clone)] -pub struct CliConfig { - pub enabled: bool, -} - -#[derive(Debug, Clone)] -pub struct HttpConfig { - pub host: String, - pub port: u16, - pub webhook_secret: Option, - pub user_id: String, -} - -/// Web gateway configuration. -#[derive(Debug, Clone)] -pub struct GatewayConfig { - pub host: String, - pub port: u16, - /// Bearer token for authentication. Random hex generated at startup if unset. - pub auth_token: Option, - pub user_id: String, -} - -impl ChannelsConfig { - fn resolve(settings: &Settings) -> Result { - 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), - webhook_secret: optional_env("HTTP_WEBHOOK_SECRET")?.map(SecretString::from), - user_id: optional_env("HTTP_USER_ID")?.unwrap_or_else(|| "http".to_string()), - }) - } else { - None - }; - - let gateway = if optional_env("GATEWAY_ENABLED")? - .map(|s| s.to_lowercase() == "true" || s == "1") - .unwrap_or(true) - { - 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), - auth_token: optional_env("GATEWAY_AUTH_TOKEN")?, - user_id: optional_env("GATEWAY_USER_ID")?.unwrap_or_else(|| "default".to_string()), - }) - } else { - None - }; - - let cli_enabled = optional_env("CLI_ENABLED")? - .map(|s| s.to_lowercase() != "false" && s != "0") - .unwrap_or(true); - - Ok(Self { - cli: CliConfig { - enabled: cli_enabled, - }, - http, - gateway, - 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), - telegram_owner_id: optional_env("TELEGRAM_OWNER_ID")? - .map(|s| s.parse()) - .transpose() - .map_err(|e| ConfigError::InvalidValue { - key: "TELEGRAM_OWNER_ID".to_string(), - message: format!("must be an integer: {e}"), - })? - .or(settings.channels.telegram_owner_id), - }) - } -} - -/// Get the default channels directory (~/.ironclaw/channels/). -fn default_channels_dir() -> PathBuf { - dirs::home_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join(".ironclaw") - .join("channels") -} - -/// Agent behavior configuration. -#[derive(Debug, Clone)] -pub struct AgentConfig { - pub name: String, - pub max_parallel_jobs: usize, - pub job_timeout: Duration, - pub stuck_threshold: Duration, - pub repair_check_interval: Duration, - pub max_repair_attempts: u32, - /// Whether to use planning before tool execution. - pub use_planning: bool, - /// Session idle timeout. Sessions inactive longer than this are pruned. - pub session_idle_timeout: Duration, - /// Allow chat to use filesystem/shell tools directly (bypass sandbox). - pub allow_local_tools: bool, - /// Maximum daily LLM spend in cents (e.g. 10000 = $100). None = unlimited. - pub max_cost_per_day_cents: Option, - /// Maximum LLM/tool actions per hour. None = unlimited. - pub max_actions_per_hour: Option, -} - -impl AgentConfig { - 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}"), - })?, - }) - } -} - -/// Safety configuration. -#[derive(Debug, Clone)] -pub struct SafetyConfig { - pub max_output_length: usize, - pub injection_check_enabled: bool, -} - -impl SafetyConfig { - 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), - }) - } -} - -/// WASM sandbox configuration. -#[derive(Debug, Clone)] -pub struct WasmConfig { - /// Whether WASM tool execution is enabled. - pub enabled: bool, - /// Directory containing installed WASM tools (default: ~/.ironclaw/tools/). - pub tools_dir: PathBuf, - /// Default memory limit in bytes (default: 10 MB). - pub default_memory_limit: u64, - /// Default execution timeout in seconds (default: 60). - pub default_timeout_secs: u64, - /// Default fuel limit for CPU metering (default: 10M). - pub default_fuel_limit: u64, - /// Whether to cache compiled modules. - pub cache_compiled: bool, - /// Directory for compiled module cache. - pub cache_dir: Option, -} - -/// Secrets management configuration. -#[derive(Clone, Default)] -pub struct SecretsConfig { - /// Master key for encrypting secrets. - pub master_key: Option, - /// Whether secrets management is enabled. - pub enabled: bool, - /// Source of the master key. - pub source: crate::settings::KeySource, -} - -impl std::fmt::Debug for SecretsConfig { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("SecretsConfig") - .field("master_key", &self.master_key.is_some()) - .field("enabled", &self.enabled) - .field("source", &self.source) - .finish() - } -} - -/// Process-wide cache for the keychain master key. -/// -/// Avoids re-prompting the OS keychain on every `SecretsConfig::resolve()` call -/// (e.g. `Config::from_env()` then `Config::from_db()`). Thread-safe alternative -/// to caching in a process env var. -impl SecretsConfig { - /// Auto-detect secrets master key from env var, then OS keychain. - /// - /// Sequential probe: SECRETS_MASTER_KEY env var first, then OS keychain. - /// No saved "source" needed; just try each source in order. - async fn resolve() -> Result { - use crate::settings::KeySource; - - let (master_key, source) = if let Some(env_key) = optional_env("SECRETS_MASTER_KEY")? { - (Some(SecretString::from(env_key)), KeySource::Env) - } else { - // Probe the OS keychain; if a key is stored, use it - match crate::secrets::keychain::get_master_key().await { - Ok(key_bytes) => { - let key_hex: String = key_bytes.iter().map(|b| format!("{:02x}", b)).collect(); - (Some(SecretString::from(key_hex)), KeySource::Keychain) - } - Err(_) => (None, KeySource::None), - } - }; - - let enabled = master_key.is_some(); - - if let Some(ref key) = master_key - && key.expose_secret().len() < 32 - { - return Err(ConfigError::InvalidValue { - key: "SECRETS_MASTER_KEY".to_string(), - message: "must be at least 32 bytes for AES-256-GCM".to_string(), - }); - } - - Ok(Self { - master_key, - enabled, - source, - }) - } - - /// Get the master key if configured. - pub fn master_key(&self) -> Option<&SecretString> { - self.master_key.as_ref() - } -} - -impl Default for WasmConfig { - fn default() -> Self { - Self { - enabled: true, - tools_dir: default_tools_dir(), - default_memory_limit: 10 * 1024 * 1024, // 10 MB - default_timeout_secs: 60, - default_fuel_limit: 10_000_000, - cache_compiled: true, - cache_dir: None, - } - } -} - -/// Get the default tools directory (~/.ironclaw/tools/). -fn default_tools_dir() -> PathBuf { - dirs::home_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join(".ironclaw") - .join("tools") -} - -impl WasmConfig { - 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), - tools_dir: optional_env("WASM_TOOLS_DIR")? - .map(PathBuf::from) - .unwrap_or_else(default_tools_dir), - default_memory_limit: parse_optional_env( - "WASM_DEFAULT_MEMORY_LIMIT", - 10 * 1024 * 1024, - )?, - 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_dir: optional_env("WASM_CACHE_DIR")?.map(PathBuf::from), - }) - } - - /// Convert to WasmRuntimeConfig. - pub fn to_runtime_config(&self) -> crate::tools::wasm::WasmRuntimeConfig { - use crate::tools::wasm::{FuelConfig, ResourceLimits, WasmRuntimeConfig}; - use std::time::Duration; - - WasmRuntimeConfig { - default_limits: ResourceLimits { - memory_bytes: self.default_memory_limit, - fuel: self.default_fuel_limit, - timeout: Duration::from_secs(self.default_timeout_secs), - }, - fuel_config: FuelConfig { - initial_fuel: self.default_fuel_limit, - enabled: true, - }, - cache_compiled: self.cache_compiled, - cache_dir: self.cache_dir.clone(), - optimization_level: wasmtime::OptLevel::Speed, - } - } -} - -/// Builder mode configuration. -#[derive(Debug, Clone)] -pub struct BuilderModeConfig { - /// Whether the software builder tool is enabled. - pub enabled: bool, - /// Directory for build artifacts (default: temp dir). - pub build_dir: Option, - /// Maximum iterations for the build loop. - pub max_iterations: u32, - /// Build timeout in seconds. - pub timeout_secs: u64, - /// Whether to automatically register built WASM tools. - pub auto_register: bool, -} - -impl Default for BuilderModeConfig { - fn default() -> Self { - Self { - enabled: true, - build_dir: None, - max_iterations: 20, - timeout_secs: 600, - auto_register: true, - } - } -} - -impl BuilderModeConfig { - 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), - 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), - }) - } - - /// Convert to BuilderConfig for the builder tool. - pub fn to_builder_config(&self) -> crate::tools::BuilderConfig { - crate::tools::BuilderConfig { - build_dir: self.build_dir.clone().unwrap_or_else(std::env::temp_dir), - max_iterations: self.max_iterations, - timeout: Duration::from_secs(self.timeout_secs), - cleanup_on_failure: true, - validate_wasm: true, - run_tests: true, - auto_register: self.auto_register, - wasm_output_dir: None, - } - } -} - -/// Heartbeat configuration. -#[derive(Debug, Clone)] -pub struct HeartbeatConfig { - /// Whether heartbeat is enabled. - pub enabled: bool, - /// Interval between heartbeat checks in seconds. - pub interval_secs: u64, - /// Channel to notify on heartbeat findings. - pub notify_channel: Option, - /// User ID to notify on heartbeat findings. - pub notify_user: Option, -} - -impl Default for HeartbeatConfig { - fn default() -> Self { - Self { - enabled: false, - interval_secs: 1800, // 30 minutes - notify_channel: None, - notify_user: None, - } - } -} - -impl HeartbeatConfig { - 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), - notify_channel: optional_env("HEARTBEAT_NOTIFY_CHANNEL")? - .or_else(|| settings.heartbeat.notify_channel.clone()), - notify_user: optional_env("HEARTBEAT_NOTIFY_USER")? - .or_else(|| settings.heartbeat.notify_user.clone()), - }) - } -} - -/// Routines configuration. -#[derive(Debug, Clone)] -pub struct RoutineConfig { - /// Whether the routines system is enabled. - pub enabled: bool, - /// How often (seconds) to poll for cron routines that need firing. - pub cron_check_interval_secs: u64, - /// Max routines executing concurrently across all users. - pub max_concurrent_routines: usize, - /// Default cooldown between fires (seconds). - pub default_cooldown_secs: u64, - /// Max output tokens for lightweight routine LLM calls. - pub max_lightweight_tokens: u32, -} - -impl Default for RoutineConfig { - fn default() -> Self { - Self { - enabled: true, - cron_check_interval_secs: 15, - max_concurrent_routines: 10, - default_cooldown_secs: 300, - max_lightweight_tokens: 4096, - } - } -} - -impl RoutineConfig { - 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), - 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)?, - max_lightweight_tokens: parse_optional_env("ROUTINES_MAX_TOKENS", 4096)?, - }) - } -} - -/// Docker sandbox configuration. -#[derive(Debug, Clone)] -pub struct SandboxModeConfig { - /// Whether the Docker sandbox is enabled. - pub enabled: bool, - /// Sandbox policy: "readonly", "workspace_write", or "full_access". - pub policy: String, - /// Command timeout in seconds. - pub timeout_secs: u64, - /// Memory limit in megabytes. - pub memory_limit_mb: u64, - /// CPU shares (relative weight). - pub cpu_shares: u32, - /// Docker image for the sandbox. - pub image: String, - /// Whether to auto-pull the image if not found. - pub auto_pull_image: bool, - /// Additional domains to allow through the network proxy. - pub extra_allowed_domains: Vec, -} - -impl Default for SandboxModeConfig { - fn default() -> Self { - Self { - enabled: true, - policy: "readonly".to_string(), - timeout_secs: 120, - memory_limit_mb: 2048, - cpu_shares: 1024, - image: "ghcr.io/nearai/sandbox:latest".to_string(), - auto_pull_image: true, - extra_allowed_domains: Vec::new(), - } - } -} - -impl SandboxModeConfig { - fn resolve() -> Result { - let extra_domains = optional_env("SANDBOX_EXTRA_DOMAINS")? - .map(|s| s.split(',').map(|d| d.trim().to_string()).collect()) - .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()), - 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(|| "ghcr.io/nearai/sandbox: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), - extra_allowed_domains: extra_domains, - }) - } - - /// Convert to SandboxConfig for the sandbox module. - pub fn to_sandbox_config(&self) -> crate::sandbox::SandboxConfig { - use crate::sandbox::SandboxPolicy; - use std::time::Duration; - - let policy = self.policy.parse().unwrap_or(SandboxPolicy::ReadOnly); - - let mut allowlist = crate::sandbox::default_allowlist(); - allowlist.extend(self.extra_allowed_domains.clone()); - - crate::sandbox::SandboxConfig { - enabled: self.enabled, - policy, - timeout: Duration::from_secs(self.timeout_secs), - memory_limit_mb: self.memory_limit_mb, - cpu_shares: self.cpu_shares, - network_allowlist: allowlist, - image: self.image.clone(), - auto_pull_image: self.auto_pull_image, - proxy_port: 0, // Auto-assign - } - } -} - -/// Claude Code sandbox configuration. -#[derive(Debug, Clone)] -pub struct ClaudeCodeConfig { - /// Whether Claude Code sandbox mode is available. - pub enabled: bool, - /// Host directory containing Claude auth config (not mounted into containers; - /// auth is handled via ANTHROPIC_API_KEY env var instead). - pub config_dir: std::path::PathBuf, - /// Claude model to use (e.g. "sonnet", "opus"). - pub model: String, - /// Maximum agentic turns before stopping. - pub max_turns: u32, - /// Memory limit in MB for Claude Code containers (heavier than workers). - pub memory_limit_mb: u64, - /// Allowed tool patterns for Claude Code permission settings. - /// - /// Written to `/workspace/.claude/settings.json` before spawning the CLI. - /// Provides defense-in-depth: only explicitly listed tools are auto-approved. - /// Any new/unknown tools would require interactive approval (which times out - /// in the non-interactive container, failing safely). - /// - /// Patterns follow Claude Code syntax: `"Bash(*)"`, `"Read"`, `"Edit(*)"`, etc. - pub allowed_tools: Vec, -} - -/// Default allowed tools for Claude Code inside containers. -/// -/// These cover all standard Claude Code tools needed for autonomous operation. -/// The Docker container provides the primary security boundary; this allowlist -/// provides defense-in-depth by preventing any future unknown tools from being -/// silently auto-approved. -fn default_claude_code_allowed_tools() -> Vec { - [ - // File system -- glob patterns match Claude Code's settings.json format - "Read(*)", - "Write(*)", - "Edit(*)", - "Glob(*)", - "Grep(*)", - "NotebookEdit(*)", - // Execution - "Bash(*)", - "Task(*)", - // Network - "WebFetch(*)", - "WebSearch(*)", - ] - .into_iter() - .map(String::from) - .collect() -} - -impl Default for ClaudeCodeConfig { - fn default() -> Self { - Self { - enabled: false, - config_dir: dirs::home_dir() - .unwrap_or_else(|| std::path::PathBuf::from(".")) - .join(".claude"), - model: "sonnet".to_string(), - max_turns: 50, - memory_limit_mb: 4096, - allowed_tools: default_claude_code_allowed_tools(), - } - } -} - -impl ClaudeCodeConfig { - /// Load from environment variables only (used inside containers where - /// there is no database or full config). - pub fn from_env() -> Self { - match Self::resolve() { - Ok(c) => c, - Err(e) => { - tracing::warn!("Failed to resolve ClaudeCodeConfig: {e}, using defaults"); - Self::default() - } - } - } - - /// Extract the OAuth access token from the host's credential store. - /// - /// On macOS: reads from Keychain (`Claude Code-credentials` service). - /// On Linux: reads from `~/.claude/.credentials.json`. - /// - /// Returns the access token if found. The token typically expires in - /// 8-12 hours, which is sufficient for any single container job. - pub fn extract_oauth_token() -> Option { - // macOS: extract from Keychain - if cfg!(target_os = "macos") { - match std::process::Command::new("security") - .args([ - "find-generic-password", - "-s", - "Claude Code-credentials", - "-w", - ]) - .output() - { - Ok(output) if output.status.success() => { - if let Ok(json) = String::from_utf8(output.stdout) { - return parse_oauth_access_token(json.trim()); - } - } - Ok(_) => { - tracing::debug!("No Claude Code credentials in macOS Keychain"); - } - Err(e) => { - tracing::debug!("Failed to query macOS Keychain: {e}"); - } - } - } - - // Linux / fallback: read from ~/.claude/.credentials.json - if let Some(home) = dirs::home_dir() { - let creds_path = home.join(".claude").join(".credentials.json"); - if let Ok(json) = std::fs::read_to_string(&creds_path) { - return parse_oauth_access_token(&json); - } - } - - None - } - - 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), - 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), - max_turns: parse_optional_env("CLAUDE_CODE_MAX_TURNS", defaults.max_turns)?, - memory_limit_mb: parse_optional_env( - "CLAUDE_CODE_MEMORY_LIMIT_MB", - defaults.memory_limit_mb, - )?, - allowed_tools: optional_env("CLAUDE_CODE_ALLOWED_TOOLS")? - .map(|s| { - s.split(',') - .map(|t| t.trim().to_string()) - .filter(|t| !t.is_empty()) - .collect() - }) - .unwrap_or(defaults.allowed_tools), - }) - } -} - -/// Parse the OAuth access token from a Claude Code credentials JSON blob. -/// -/// Expected shape: `{"claudeAiOauth": {"accessToken": "sk-ant-oat01-..."}}` -fn parse_oauth_access_token(json: &str) -> Option { - let creds: serde_json::Value = serde_json::from_str(json).ok()?; - creds["claudeAiOauth"]["accessToken"] - .as_str() - .map(String::from) -} - -/// Skills system configuration. -#[derive(Debug, Clone)] -pub struct SkillsConfig { - /// Whether the skills system is enabled. - pub enabled: bool, - /// Directory containing local skills (default: ~/.ironclaw/skills/). - pub local_dir: PathBuf, - /// Maximum number of skills that can be active simultaneously. - pub max_active_skills: usize, - /// Maximum total context tokens allocated to skill prompts. - pub max_context_tokens: usize, -} - -impl Default for SkillsConfig { - fn default() -> Self { - Self { - enabled: false, - local_dir: default_skills_dir(), - max_active_skills: 3, - max_context_tokens: 4000, - } - } -} - -/// Get the default skills directory (~/.ironclaw/skills/). -fn default_skills_dir() -> PathBuf { - dirs::home_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join(".ironclaw") - .join("skills") -} - -impl SkillsConfig { - 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), - local_dir: optional_env("SKILLS_DIR")? - .map(PathBuf::from) - .unwrap_or_else(default_skills_dir), - max_active_skills: parse_optional_env("SKILLS_MAX_ACTIVE", 3)?, - max_context_tokens: parse_optional_env("SKILLS_MAX_CONTEXT_TOKENS", 4000)?, - }) - } -} - -/// Load API keys from the encrypted secrets store into a thread-safe overlay. -/// -/// This bridges the gap between secrets stored during onboarding and the -/// env-var-first resolution in `LlmConfig::resolve()`. Keys in the overlay -/// are read by `optional_env()` before falling back to `std::env::var()`, -/// so explicit env vars always win. -pub async fn inject_llm_keys_from_secrets( - secrets: &dyn crate::secrets::SecretsStore, - user_id: &str, -) { - let mappings = [ - ("llm_openai_api_key", "OPENAI_API_KEY"), - ("llm_anthropic_api_key", "ANTHROPIC_API_KEY"), - ("llm_compatible_api_key", "LLM_API_KEY"), - ]; - - let mut injected = HashMap::new(); - - for (secret_name, env_var) in mappings { - match std::env::var(env_var) { - Ok(val) if !val.is_empty() => continue, - _ => {} - } - match secrets.get_decrypted(user_id, secret_name).await { - Ok(decrypted) => { - injected.insert(env_var.to_string(), decrypted.expose().to_string()); - tracing::debug!("Loaded secret '{}' for env var '{}'", secret_name, env_var); - } - Err(_) => { - // Secret doesn't exist, that's fine - } - } - } - - let _ = INJECTED_VARS.set(injected); -} - -// Helper functions - -fn optional_env(key: &str) -> Result, ConfigError> { - // Check real env vars first (always win over injected secrets) - match std::env::var(key) { - Ok(val) if val.is_empty() => {} - Ok(val) => return Ok(Some(val)), - Err(std::env::VarError::NotPresent) => {} - Err(e) => { - return Err(ConfigError::ParseError(format!( - "failed to read {key}: {e}" - ))); - } - } - - // Fall back to thread-safe overlay (secrets injected from DB) - if let Some(val) = INJECTED_VARS.get().and_then(|map| map.get(key)) { - return Ok(Some(val.clone())); - } - - Ok(None) -} - -fn parse_optional_env(key: &str, default: T) -> Result -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() - .map(|opt| opt.unwrap_or(default)) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::settings::{EmbeddingsSettings, Settings}; - use std::sync::Mutex; - - /// Serializes env-mutating tests to prevent parallel races. - static ENV_MUTEX: Mutex<()> = Mutex::new(()); - - /// Clear all embedding-related env vars. - fn clear_embedding_env() { - // SAFETY: Only called under ENV_MUTEX in tests. No other threads - // observe these vars while the lock is held. - unsafe { - std::env::remove_var("EMBEDDING_ENABLED"); - std::env::remove_var("EMBEDDING_PROVIDER"); - std::env::remove_var("EMBEDDING_MODEL"); - std::env::remove_var("OPENAI_API_KEY"); - } - } - - /// Clear all openai-compatible-related env vars. - fn clear_openai_compatible_env() { - // SAFETY: Only called under ENV_MUTEX in tests. - unsafe { - std::env::remove_var("LLM_BACKEND"); - std::env::remove_var("LLM_BASE_URL"); - std::env::remove_var("LLM_MODEL"); - } - } - - #[test] - fn embeddings_disabled_not_overridden_by_openai_key() { - let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); - - clear_embedding_env(); - // SAFETY: Under ENV_MUTEX, no concurrent env access. - unsafe { - std::env::set_var("OPENAI_API_KEY", "sk-test-key-for-issue-129"); - } - - let settings = Settings { - embeddings: EmbeddingsSettings { - enabled: false, - ..Default::default() - }, - ..Default::default() - }; - - let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); - assert!( - !config.enabled, - "embeddings should remain disabled when settings.embeddings.enabled=false, \ - even when OPENAI_API_KEY is set (issue #129)" - ); - - // SAFETY: Under ENV_MUTEX. - unsafe { - std::env::remove_var("OPENAI_API_KEY"); - } - } - - #[test] - fn embeddings_enabled_from_settings() { - let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); - clear_embedding_env(); - - let settings = Settings { - embeddings: EmbeddingsSettings { - enabled: true, - ..Default::default() - }, - ..Default::default() - }; - - let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); - assert!( - config.enabled, - "embeddings should be enabled when settings say so" - ); - } - - #[test] - fn embeddings_env_override_takes_precedence() { - let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); - - clear_embedding_env(); - // SAFETY: Under ENV_MUTEX. - unsafe { - std::env::set_var("EMBEDDING_ENABLED", "true"); - } - - let settings = Settings { - embeddings: EmbeddingsSettings { - enabled: false, - ..Default::default() - }, - ..Default::default() - }; - - let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); - assert!( - config.enabled, - "EMBEDDING_ENABLED=true env var should override settings" - ); - - // SAFETY: Under ENV_MUTEX. - unsafe { - std::env::remove_var("EMBEDDING_ENABLED"); - } - } - - #[test] - fn openai_compatible_uses_selected_model_when_llm_model_unset() { - let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); - clear_openai_compatible_env(); - - let settings = Settings { - llm_backend: Some("openai_compatible".to_string()), - openai_compatible_base_url: Some("https://openrouter.ai/api/v1".to_string()), - selected_model: Some("openai/gpt-5.1-codex".to_string()), - ..Default::default() - }; - - let cfg = LlmConfig::resolve(&settings).expect("resolve should succeed"); - let compat = cfg - .openai_compatible - .expect("openai-compatible config should be present"); - - assert_eq!(compat.model, "openai/gpt-5.1-codex"); - } - - #[test] - fn openai_compatible_llm_model_env_overrides_selected_model() { - let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); - clear_openai_compatible_env(); - // SAFETY: Under ENV_MUTEX. - unsafe { - std::env::set_var("LLM_MODEL", "openai/gpt-5-codex"); - } - - let settings = Settings { - llm_backend: Some("openai_compatible".to_string()), - openai_compatible_base_url: Some("https://openrouter.ai/api/v1".to_string()), - selected_model: Some("openai/gpt-5.1-codex".to_string()), - ..Default::default() - }; - - let cfg = LlmConfig::resolve(&settings).expect("resolve should succeed"); - let compat = cfg - .openai_compatible - .expect("openai-compatible config should be present"); - - assert_eq!(compat.model, "openai/gpt-5-codex"); - - // SAFETY: Under ENV_MUTEX. - unsafe { - std::env::remove_var("LLM_MODEL"); - } - } -} diff --git a/src/config/agent.rs b/src/config/agent.rs new file mode 100644 index 00000000..25b02f1d --- /dev/null +++ b/src/config/agent.rs @@ -0,0 +1,120 @@ +use std::time::Duration; + +use crate::config::helpers::optional_env; +use crate::error::ConfigError; +use crate::settings::Settings; + +/// Agent behavior configuration. +#[derive(Debug, Clone)] +pub struct AgentConfig { + pub name: String, + pub max_parallel_jobs: usize, + pub job_timeout: Duration, + pub stuck_threshold: Duration, + pub repair_check_interval: Duration, + pub max_repair_attempts: u32, + /// Whether to use planning before tool execution. + pub use_planning: bool, + /// Session idle timeout. Sessions inactive longer than this are pruned. + pub session_idle_timeout: Duration, + /// Allow chat to use filesystem/shell tools directly (bypass sandbox). + pub allow_local_tools: bool, + /// Maximum daily LLM spend in cents (e.g. 10000 = $100). None = unlimited. + pub max_cost_per_day_cents: Option, + /// Maximum LLM/tool actions per hour. None = unlimited. + pub max_actions_per_hour: Option, +} + +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}"), + })?, + }) + } +} diff --git a/src/config/builder.rs b/src/config/builder.rs new file mode 100644 index 00000000..fede5bce --- /dev/null +++ b/src/config/builder.rs @@ -0,0 +1,72 @@ +use std::path::PathBuf; +use std::time::Duration; + +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; + +/// Builder mode configuration. +#[derive(Debug, Clone)] +pub struct BuilderModeConfig { + /// Whether the software builder tool is enabled. + pub enabled: bool, + /// Directory for build artifacts (default: temp dir). + pub build_dir: Option, + /// Maximum iterations for the build loop. + pub max_iterations: u32, + /// Build timeout in seconds. + pub timeout_secs: u64, + /// Whether to automatically register built WASM tools. + pub auto_register: bool, +} + +impl Default for BuilderModeConfig { + fn default() -> Self { + Self { + enabled: true, + build_dir: None, + max_iterations: 20, + timeout_secs: 600, + auto_register: true, + } + } +} + +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), + 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), + }) + } + + /// Convert to BuilderConfig for the builder tool. + pub fn to_builder_config(&self) -> crate::tools::BuilderConfig { + crate::tools::BuilderConfig { + build_dir: self.build_dir.clone().unwrap_or_else(std::env::temp_dir), + max_iterations: self.max_iterations, + timeout: Duration::from_secs(self.timeout_secs), + cleanup_on_failure: true, + validate_wasm: true, + run_tests: true, + auto_register: self.auto_register, + wasm_output_dir: None, + } + } +} diff --git a/src/config/channels.rs b/src/config/channels.rs new file mode 100644 index 00000000..31eaffab --- /dev/null +++ b/src/config/channels.rs @@ -0,0 +1,126 @@ +use std::path::PathBuf; + +use secrecy::SecretString; + +use crate::config::helpers::optional_env; +use crate::error::ConfigError; +use crate::settings::Settings; + +/// Channel configurations. +#[derive(Debug, Clone)] +pub struct ChannelsConfig { + pub cli: CliConfig, + pub http: Option, + pub gateway: Option, + /// Directory containing WASM channel modules (default: ~/.ironclaw/channels/). + pub wasm_channels_dir: std::path::PathBuf, + /// Whether WASM channels are enabled. + pub wasm_channels_enabled: bool, + /// Telegram owner user ID. When set, the bot only responds to this user. + pub telegram_owner_id: Option, +} + +#[derive(Debug, Clone)] +pub struct CliConfig { + pub enabled: bool, +} + +#[derive(Debug, Clone)] +pub struct HttpConfig { + pub host: String, + pub port: u16, + pub webhook_secret: Option, + pub user_id: String, +} + +/// Web gateway configuration. +#[derive(Debug, Clone)] +pub struct GatewayConfig { + pub host: String, + pub port: u16, + /// Bearer token for authentication. Random hex generated at startup if unset. + pub auth_token: Option, + pub user_id: String, +} + +impl ChannelsConfig { + pub(crate) fn resolve(settings: &Settings) -> Result { + 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), + webhook_secret: optional_env("HTTP_WEBHOOK_SECRET")?.map(SecretString::from), + user_id: optional_env("HTTP_USER_ID")?.unwrap_or_else(|| "http".to_string()), + }) + } else { + None + }; + + let gateway = if optional_env("GATEWAY_ENABLED")? + .map(|s| s.to_lowercase() == "true" || s == "1") + .unwrap_or(true) + { + 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), + auth_token: optional_env("GATEWAY_AUTH_TOKEN")?, + user_id: optional_env("GATEWAY_USER_ID")?.unwrap_or_else(|| "default".to_string()), + }) + } else { + None + }; + + let cli_enabled = optional_env("CLI_ENABLED")? + .map(|s| s.to_lowercase() != "false" && s != "0") + .unwrap_or(true); + + Ok(Self { + cli: CliConfig { + enabled: cli_enabled, + }, + http, + gateway, + 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), + telegram_owner_id: optional_env("TELEGRAM_OWNER_ID")? + .map(|s| s.parse()) + .transpose() + .map_err(|e| ConfigError::InvalidValue { + key: "TELEGRAM_OWNER_ID".to_string(), + message: format!("must be an integer: {e}"), + })? + .or(settings.channels.telegram_owner_id), + }) + } +} + +/// Get the default channels directory (~/.ironclaw/channels/). +fn default_channels_dir() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("channels") +} diff --git a/src/config/database.rs b/src/config/database.rs new file mode 100644 index 00000000..12b176c0 --- /dev/null +++ b/src/config/database.rs @@ -0,0 +1,130 @@ +use std::path::PathBuf; + +use secrecy::{ExposeSecret, SecretString}; + +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; + +/// Which database backend to use. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum DatabaseBackend { + /// PostgreSQL via deadpool-postgres (default). + #[default] + Postgres, + /// libSQL/Turso embedded database. + LibSql, +} + +impl std::fmt::Display for DatabaseBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Postgres => write!(f, "postgres"), + Self::LibSql => write!(f, "libsql"), + } + } +} + +impl std::str::FromStr for DatabaseBackend { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "postgres" | "postgresql" | "pg" => Ok(Self::Postgres), + "libsql" | "turso" | "sqlite" => Ok(Self::LibSql), + _ => Err(format!( + "invalid database backend '{}', expected 'postgres' or 'libsql'", + s + )), + } + } +} + +/// Database configuration. +#[derive(Debug, Clone)] +pub struct DatabaseConfig { + /// Which backend to use (default: Postgres). + pub backend: DatabaseBackend, + + // -- PostgreSQL fields -- + pub url: SecretString, + pub pool_size: usize, + + // -- libSQL fields -- + /// Path to local libSQL database file (default: ~/.ironclaw/ironclaw.db). + pub libsql_path: Option, + /// Turso cloud URL for remote sync (optional). + pub libsql_url: Option, + /// Turso auth token (required when libsql_url is set). + pub libsql_auth_token: Option, +} + +impl DatabaseConfig { + pub(crate) fn resolve() -> Result { + let backend: DatabaseBackend = if let Some(b) = optional_env("DATABASE_BACKEND")? { + b.parse().map_err(|e| ConfigError::InvalidValue { + key: "DATABASE_BACKEND".to_string(), + message: e, + })? + } else { + DatabaseBackend::default() + }; + + // PostgreSQL URL is required only when using the postgres backend. + // For libsql backend, default to an empty placeholder. + // DATABASE_URL is loaded from ~/.ironclaw/.env via dotenvy early in startup. + let url = optional_env("DATABASE_URL")? + .or_else(|| { + if backend == DatabaseBackend::LibSql { + Some("unused://libsql".to_string()) + } else { + None + } + }) + .ok_or_else(|| ConfigError::MissingRequired { + key: "DATABASE_URL".to_string(), + hint: "Run 'ironclaw onboard' or set DATABASE_URL environment variable".to_string(), + })?; + + let pool_size = parse_optional_env("DATABASE_POOL_SIZE", 10)?; + + let libsql_path = optional_env("LIBSQL_PATH")?.map(PathBuf::from).or_else(|| { + if backend == DatabaseBackend::LibSql { + Some(default_libsql_path()) + } else { + None + } + }); + + let libsql_url = optional_env("LIBSQL_URL")?; + let libsql_auth_token = optional_env("LIBSQL_AUTH_TOKEN")?.map(SecretString::from); + + if libsql_url.is_some() && libsql_auth_token.is_none() { + return Err(ConfigError::MissingRequired { + key: "LIBSQL_AUTH_TOKEN".to_string(), + hint: "LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set".to_string(), + }); + } + + Ok(Self { + backend, + url: SecretString::from(url), + pool_size, + libsql_path, + libsql_url, + libsql_auth_token, + }) + } + + /// Get the database URL (exposes the secret). + pub fn url(&self) -> &str { + self.url.expose_secret() + } +} + +/// Default libSQL database path (~/.ironclaw/ironclaw.db). +pub fn default_libsql_path() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("ironclaw.db") +} diff --git a/src/config/embeddings.rs b/src/config/embeddings.rs new file mode 100644 index 00000000..466f11c5 --- /dev/null +++ b/src/config/embeddings.rs @@ -0,0 +1,165 @@ +use secrecy::{ExposeSecret, SecretString}; + +use crate::config::helpers::optional_env; +use crate::error::ConfigError; +use crate::settings::Settings; + +/// Embeddings provider configuration. +#[derive(Debug, Clone)] +pub struct EmbeddingsConfig { + /// Whether embeddings are enabled. + pub enabled: bool, + /// Provider to use: "openai" or "nearai" + pub provider: String, + /// OpenAI API key (for OpenAI provider). + pub openai_api_key: Option, + /// Model to use for embeddings. + pub model: String, +} + +impl Default for EmbeddingsConfig { + fn default() -> Self { + Self { + enabled: false, + provider: "openai".to_string(), + openai_api_key: None, + model: "text-embedding-3-small".to_string(), + } + } +} + +impl EmbeddingsConfig { + pub(crate) fn resolve(settings: &Settings) -> Result { + let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from); + + let provider = optional_env("EMBEDDING_PROVIDER")? + .unwrap_or_else(|| settings.embeddings.provider.clone()); + + let model = + optional_env("EMBEDDING_MODEL")?.unwrap_or_else(|| settings.embeddings.model.clone()); + + 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); + + Ok(Self { + enabled, + provider, + openai_api_key, + model, + }) + } + + /// Get the OpenAI API key if configured. + pub fn openai_api_key(&self) -> Option<&str> { + self.openai_api_key.as_ref().map(|s| s.expose_secret()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::settings::{EmbeddingsSettings, Settings}; + use std::sync::Mutex; + + /// Serializes env-mutating tests to prevent parallel races. + static ENV_MUTEX: Mutex<()> = Mutex::new(()); + + /// Clear all embedding-related env vars. + fn clear_embedding_env() { + // SAFETY: Only called under ENV_MUTEX in tests. No other threads + // observe these vars while the lock is held. + unsafe { + std::env::remove_var("EMBEDDING_ENABLED"); + std::env::remove_var("EMBEDDING_PROVIDER"); + std::env::remove_var("EMBEDDING_MODEL"); + std::env::remove_var("OPENAI_API_KEY"); + } + } + + #[test] + fn embeddings_disabled_not_overridden_by_openai_key() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + + clear_embedding_env(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::set_var("OPENAI_API_KEY", "sk-test-key-for-issue-129"); + } + + let settings = Settings { + embeddings: EmbeddingsSettings { + enabled: false, + ..Default::default() + }, + ..Default::default() + }; + + let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); + assert!( + !config.enabled, + "embeddings should remain disabled when settings.embeddings.enabled=false, \ + even when OPENAI_API_KEY is set (issue #129)" + ); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("OPENAI_API_KEY"); + } + } + + #[test] + fn embeddings_enabled_from_settings() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_embedding_env(); + + let settings = Settings { + embeddings: EmbeddingsSettings { + enabled: true, + ..Default::default() + }, + ..Default::default() + }; + + let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); + assert!( + config.enabled, + "embeddings should be enabled when settings say so" + ); + } + + #[test] + fn embeddings_env_override_takes_precedence() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + + clear_embedding_env(); + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("EMBEDDING_ENABLED", "true"); + } + + let settings = Settings { + embeddings: EmbeddingsSettings { + enabled: false, + ..Default::default() + }, + ..Default::default() + }; + + let config = EmbeddingsConfig::resolve(&settings).expect("resolve should succeed"); + assert!( + config.enabled, + "EMBEDDING_ENABLED=true env var should override settings" + ); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("EMBEDDING_ENABLED"); + } + } +} diff --git a/src/config/heartbeat.rs b/src/config/heartbeat.rs new file mode 100644 index 00000000..9fe0831b --- /dev/null +++ b/src/config/heartbeat.rs @@ -0,0 +1,54 @@ +use crate::config::helpers::optional_env; +use crate::error::ConfigError; +use crate::settings::Settings; + +/// Heartbeat configuration. +#[derive(Debug, Clone)] +pub struct HeartbeatConfig { + /// Whether heartbeat is enabled. + pub enabled: bool, + /// Interval between heartbeat checks in seconds. + pub interval_secs: u64, + /// Channel to notify on heartbeat findings. + pub notify_channel: Option, + /// User ID to notify on heartbeat findings. + pub notify_user: Option, +} + +impl Default for HeartbeatConfig { + fn default() -> Self { + Self { + enabled: false, + interval_secs: 1800, // 30 minutes + notify_channel: None, + notify_user: None, + } + } +} + +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), + notify_channel: optional_env("HEARTBEAT_NOTIFY_CHANNEL")? + .or_else(|| settings.heartbeat.notify_channel.clone()), + notify_user: optional_env("HEARTBEAT_NOTIFY_USER")? + .or_else(|| settings.heartbeat.notify_user.clone()), + }) + } +} diff --git a/src/config/helpers.rs b/src/config/helpers.rs new file mode 100644 index 00000000..b463bb50 --- /dev/null +++ b/src/config/helpers.rs @@ -0,0 +1,40 @@ +use crate::error::ConfigError; + +use super::INJECTED_VARS; + +pub(crate) fn optional_env(key: &str) -> Result, ConfigError> { + // Check real env vars first (always win over injected secrets) + match std::env::var(key) { + Ok(val) if val.is_empty() => {} + Ok(val) => return Ok(Some(val)), + Err(std::env::VarError::NotPresent) => {} + Err(e) => { + return Err(ConfigError::ParseError(format!( + "failed to read {key}: {e}" + ))); + } + } + + // Fall back to thread-safe overlay (secrets injected from DB) + if let Some(val) = INJECTED_VARS.get().and_then(|map| map.get(key)) { + return Ok(Some(val.clone())); + } + + Ok(None) +} + +pub(crate) fn parse_optional_env(key: &str, default: T) -> Result +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() + .map(|opt| opt.unwrap_or(default)) +} diff --git a/src/config/llm.rs b/src/config/llm.rs new file mode 100644 index 00000000..53b21892 --- /dev/null +++ b/src/config/llm.rs @@ -0,0 +1,426 @@ +use std::path::PathBuf; + +use secrecy::SecretString; + +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; +use crate::settings::Settings; + +/// Which LLM backend to use. +/// +/// Defaults to `NearAi` to keep IronClaw close to the NEAR ecosystem. +/// Users can override with `LLM_BACKEND` env var to use their own API keys. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum LlmBackend { + /// NEAR AI proxy (default) -- session or API key auth + #[default] + NearAi, + /// Direct OpenAI API + OpenAi, + /// Direct Anthropic API + Anthropic, + /// Local Ollama instance + Ollama, + /// Any OpenAI-compatible endpoint (e.g. vLLM, LiteLLM, Together) + OpenAiCompatible, + /// Tinfoil private inference + Tinfoil, +} + +impl std::str::FromStr for LlmBackend { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "nearai" | "near_ai" | "near" => Ok(Self::NearAi), + "openai" | "open_ai" => Ok(Self::OpenAi), + "anthropic" | "claude" => Ok(Self::Anthropic), + "ollama" => Ok(Self::Ollama), + "openai_compatible" | "openai-compatible" | "compatible" => Ok(Self::OpenAiCompatible), + "tinfoil" => Ok(Self::Tinfoil), + _ => Err(format!( + "invalid LLM backend '{}', expected one of: nearai, openai, anthropic, ollama, openai_compatible, tinfoil", + s + )), + } + } +} + +impl std::fmt::Display for LlmBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::NearAi => write!(f, "nearai"), + Self::OpenAi => write!(f, "openai"), + Self::Anthropic => write!(f, "anthropic"), + Self::Ollama => write!(f, "ollama"), + Self::OpenAiCompatible => write!(f, "openai_compatible"), + Self::Tinfoil => write!(f, "tinfoil"), + } + } +} + +/// Configuration for direct OpenAI API access. +#[derive(Debug, Clone)] +pub struct OpenAiDirectConfig { + pub api_key: SecretString, + pub model: String, +} + +/// Configuration for direct Anthropic API access. +#[derive(Debug, Clone)] +pub struct AnthropicDirectConfig { + pub api_key: SecretString, + pub model: String, +} + +/// Configuration for local Ollama. +#[derive(Debug, Clone)] +pub struct OllamaConfig { + pub base_url: String, + pub model: String, +} + +/// Configuration for any OpenAI-compatible endpoint. +#[derive(Debug, Clone)] +pub struct OpenAiCompatibleConfig { + pub base_url: String, + pub api_key: Option, + pub model: String, +} + +/// Configuration for Tinfoil private inference. +#[derive(Debug, Clone)] +pub struct TinfoilConfig { + pub api_key: SecretString, + pub model: String, +} + +/// LLM provider configuration. +/// +/// NEAR AI remains the default backend. Users can switch to other providers +/// by setting `LLM_BACKEND` (e.g. `openai`, `anthropic`, `ollama`). +#[derive(Debug, Clone)] +pub struct LlmConfig { + /// Which backend to use (default: NearAi) + pub backend: LlmBackend, + /// NEAR AI config (always populated for NEAR AI embeddings, etc.) + pub nearai: NearAiConfig, + /// Direct OpenAI config (populated when backend=openai) + pub openai: Option, + /// Direct Anthropic config (populated when backend=anthropic) + pub anthropic: Option, + /// Ollama config (populated when backend=ollama) + pub ollama: Option, + /// OpenAI-compatible config (populated when backend=openai_compatible) + pub openai_compatible: Option, + /// Tinfoil config (populated when backend=tinfoil) + pub tinfoil: Option, +} + +/// API mode for NEAR AI. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum NearAiApiMode { + /// Use the Responses API (chat-api proxy) - session-based auth + #[default] + Responses, + /// Use the Chat Completions API (cloud-api) - API key auth + ChatCompletions, +} + +impl std::str::FromStr for NearAiApiMode { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "responses" | "response" => Ok(Self::Responses), + "chat_completions" | "chatcompletions" | "chat" | "completions" => { + Ok(Self::ChatCompletions) + } + _ => Err(format!( + "invalid API mode '{}', expected 'responses' or 'chat_completions'", + s + )), + } + } +} + +/// NEAR AI chat-api configuration. +#[derive(Debug, Clone)] +pub struct NearAiConfig { + /// Model to use (e.g., "claude-3-5-sonnet-20241022", "gpt-4o") + pub model: String, + /// Cheap/fast model for lightweight tasks (heartbeat, routing, evaluation). + /// Falls back to the main model if not set. + pub cheap_model: Option, + /// Base URL for the NEAR AI API (default: https://private.near.ai). + pub base_url: String, + /// Base URL for auth/refresh endpoints (default: https://private.near.ai) + pub auth_base_url: String, + /// Path to session file (default: ~/.ironclaw/session.json) + pub session_path: PathBuf, + /// API mode: "responses" (chat-api) or "chat_completions" (cloud-api) + pub api_mode: NearAiApiMode, + /// API key for cloud-api (required for chat_completions mode) + pub api_key: Option, + /// Optional fallback model for failover (default: None). + /// When set, a secondary provider is created with this model and wrapped + /// in a `FailoverProvider` so transient errors on the primary model + /// automatically fall through to the fallback. + pub fallback_model: Option, + /// Maximum number of retries for transient errors (default: 3). + /// With the default of 3, the provider makes up to 4 total attempts + /// (1 initial + 3 retries) before giving up. + pub max_retries: u32, + /// Consecutive transient failures before the circuit breaker opens. + /// None = disabled (default). E.g. 5 means after 5 consecutive failures + /// all requests are rejected until recovery timeout elapses. + pub circuit_breaker_threshold: Option, + /// How long (seconds) the circuit stays open before allowing a probe (default: 30). + pub circuit_breaker_recovery_secs: u64, + /// Enable in-memory response caching for `complete()` calls. + /// Saves tokens on repeated prompts within a session. Default: false. + pub response_cache_enabled: bool, + /// TTL in seconds for cached responses (default: 3600 = 1 hour). + pub response_cache_ttl_secs: u64, + /// Max cached responses before LRU eviction (default: 1000). + pub response_cache_max_entries: usize, + /// Cooldown duration in seconds for the failover provider (default: 300). + /// When a provider accumulates enough consecutive failures it is skipped + /// for this many seconds. + pub failover_cooldown_secs: u64, + /// Number of consecutive retryable failures before a provider enters + /// cooldown (default: 3). + pub failover_cooldown_threshold: u32, +} + +impl LlmConfig { + pub(crate) fn resolve(settings: &Settings) -> Result { + // Determine backend: env var > settings > default (NearAi) + let backend: LlmBackend = if let Some(b) = optional_env("LLM_BACKEND")? { + b.parse().map_err(|e| ConfigError::InvalidValue { + key: "LLM_BACKEND".to_string(), + message: e, + })? + } else if let Some(ref b) = settings.llm_backend { + match b.parse() { + Ok(backend) => backend, + Err(e) => { + tracing::warn!( + "Invalid llm_backend '{}' in settings: {}. Using default NearAi.", + b, + e + ); + LlmBackend::NearAi + } + } + } else { + LlmBackend::NearAi + }; + + // Resolve NEAR AI config only when backend is NearAi (or when explicitly configured) + let nearai_api_key = optional_env("NEARAI_API_KEY")?.map(SecretString::from); + + let api_mode = if let Some(mode_str) = optional_env("NEARAI_API_MODE")? { + mode_str.parse().map_err(|e| ConfigError::InvalidValue { + key: "NEARAI_API_MODE".to_string(), + message: e, + })? + } else if nearai_api_key.is_some() { + NearAiApiMode::ChatCompletions + } else { + NearAiApiMode::Responses + }; + + 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() + }), + cheap_model: optional_env("NEARAI_CHEAP_MODEL")?, + base_url: optional_env("NEARAI_BASE_URL")? + .unwrap_or_else(|| "https://private.near.ai".to_string()), + auth_base_url: optional_env("NEARAI_AUTH_URL")? + .unwrap_or_else(|| "https://private.near.ai".to_string()), + session_path: optional_env("NEARAI_SESSION_PATH")? + .map(PathBuf::from) + .unwrap_or_else(default_session_path), + api_mode, + api_key: nearai_api_key, + fallback_model: optional_env("NEARAI_FALLBACK_MODEL")?, + max_retries: parse_optional_env("NEARAI_MAX_RETRIES", 3)?, + circuit_breaker_threshold: optional_env("CIRCUIT_BREAKER_THRESHOLD")? + .map(|s| s.parse()) + .transpose() + .map_err(|e| ConfigError::InvalidValue { + key: "CIRCUIT_BREAKER_THRESHOLD".to_string(), + message: format!("must be a positive integer: {e}"), + })?, + circuit_breaker_recovery_secs: parse_optional_env("CIRCUIT_BREAKER_RECOVERY_SECS", 30)?, + response_cache_enabled: parse_optional_env("RESPONSE_CACHE_ENABLED", false)?, + response_cache_ttl_secs: parse_optional_env("RESPONSE_CACHE_TTL_SECS", 3600)?, + response_cache_max_entries: parse_optional_env("RESPONSE_CACHE_MAX_ENTRIES", 1000)?, + failover_cooldown_secs: parse_optional_env("LLM_FAILOVER_COOLDOWN_SECS", 300)?, + failover_cooldown_threshold: parse_optional_env("LLM_FAILOVER_THRESHOLD", 3)?, + }; + + // Resolve provider-specific configs based on backend + let openai = if backend == LlmBackend::OpenAi { + let api_key = optional_env("OPENAI_API_KEY")? + .map(SecretString::from) + .ok_or_else(|| ConfigError::MissingRequired { + key: "OPENAI_API_KEY".to_string(), + hint: "Set OPENAI_API_KEY when LLM_BACKEND=openai".to_string(), + })?; + let model = optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4o".to_string()); + Some(OpenAiDirectConfig { api_key, model }) + } else { + None + }; + + let anthropic = if backend == LlmBackend::Anthropic { + let api_key = optional_env("ANTHROPIC_API_KEY")? + .map(SecretString::from) + .ok_or_else(|| ConfigError::MissingRequired { + key: "ANTHROPIC_API_KEY".to_string(), + hint: "Set ANTHROPIC_API_KEY when LLM_BACKEND=anthropic".to_string(), + })?; + let model = optional_env("ANTHROPIC_MODEL")? + .unwrap_or_else(|| "claude-sonnet-4-20250514".to_string()); + Some(AnthropicDirectConfig { api_key, model }) + } else { + None + }; + + let ollama = if backend == LlmBackend::Ollama { + let base_url = optional_env("OLLAMA_BASE_URL")? + .or_else(|| settings.ollama_base_url.clone()) + .unwrap_or_else(|| "http://localhost:11434".to_string()); + let model = optional_env("OLLAMA_MODEL")?.unwrap_or_else(|| "llama3".to_string()); + Some(OllamaConfig { base_url, model }) + } else { + None + }; + + let openai_compatible = if backend == LlmBackend::OpenAiCompatible { + let base_url = optional_env("LLM_BASE_URL")? + .or_else(|| settings.openai_compatible_base_url.clone()) + .ok_or_else(|| ConfigError::MissingRequired { + key: "LLM_BASE_URL".to_string(), + hint: "Set LLM_BASE_URL when LLM_BACKEND=openai_compatible".to_string(), + })?; + let api_key = optional_env("LLM_API_KEY")?.map(SecretString::from); + let model = optional_env("LLM_MODEL")? + .or_else(|| settings.selected_model.clone()) + .unwrap_or_else(|| "default".to_string()); + Some(OpenAiCompatibleConfig { + base_url, + api_key, + model, + }) + } else { + None + }; + + let tinfoil = if backend == LlmBackend::Tinfoil { + let api_key = optional_env("TINFOIL_API_KEY")? + .map(SecretString::from) + .ok_or_else(|| ConfigError::MissingRequired { + key: "TINFOIL_API_KEY".to_string(), + hint: "Set TINFOIL_API_KEY when LLM_BACKEND=tinfoil".to_string(), + })?; + let model = optional_env("TINFOIL_MODEL")?.unwrap_or_else(|| "kimi-k2-5".to_string()); + Some(TinfoilConfig { api_key, model }) + } else { + None + }; + + Ok(Self { + backend, + nearai, + openai, + anthropic, + ollama, + openai_compatible, + tinfoil, + }) + } +} + +/// Get the default session file path (~/.ironclaw/session.json). +fn default_session_path() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("session.json") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::settings::Settings; + use std::sync::Mutex; + + /// Serializes env-mutating tests to prevent parallel races. + static ENV_MUTEX: Mutex<()> = Mutex::new(()); + + /// Clear all openai-compatible-related env vars. + fn clear_openai_compatible_env() { + // SAFETY: Only called under ENV_MUTEX in tests. + unsafe { + std::env::remove_var("LLM_BACKEND"); + std::env::remove_var("LLM_BASE_URL"); + std::env::remove_var("LLM_MODEL"); + } + } + + #[test] + fn openai_compatible_uses_selected_model_when_llm_model_unset() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_openai_compatible_env(); + + let settings = Settings { + llm_backend: Some("openai_compatible".to_string()), + openai_compatible_base_url: Some("https://openrouter.ai/api/v1".to_string()), + selected_model: Some("openai/gpt-5.1-codex".to_string()), + ..Default::default() + }; + + let cfg = LlmConfig::resolve(&settings).expect("resolve should succeed"); + let compat = cfg + .openai_compatible + .expect("openai-compatible config should be present"); + + assert_eq!(compat.model, "openai/gpt-5.1-codex"); + } + + #[test] + fn openai_compatible_llm_model_env_overrides_selected_model() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + clear_openai_compatible_env(); + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::set_var("LLM_MODEL", "openai/gpt-5-codex"); + } + + let settings = Settings { + llm_backend: Some("openai_compatible".to_string()), + openai_compatible_base_url: Some("https://openrouter.ai/api/v1".to_string()), + selected_model: Some("openai/gpt-5.1-codex".to_string()), + ..Default::default() + }; + + let cfg = LlmConfig::resolve(&settings).expect("resolve should succeed"); + let compat = cfg + .openai_compatible + .expect("openai-compatible config should be present"); + + assert_eq!(compat.model, "openai/gpt-5-codex"); + + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("LLM_MODEL"); + } + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs new file mode 100644 index 00000000..2d227723 --- /dev/null +++ b/src/config/mod.rs @@ -0,0 +1,239 @@ +//! Configuration for IronClaw. +//! +//! Settings are loaded with priority: env var > database > default. +//! `DATABASE_URL` lives in `~/.ironclaw/.env` (loaded via dotenvy early +//! in startup). Everything else comes from env vars, the DB settings +//! table, or auto-detection. + +mod agent; +mod builder; +mod channels; +mod database; +mod embeddings; +mod heartbeat; +pub(crate) mod helpers; +mod llm; +mod routines; +mod safety; +mod sandbox; +mod secrets; +mod skills; +mod tunnel; +mod wasm; + +use std::collections::HashMap; +use std::sync::OnceLock; + +use crate::error::ConfigError; +use crate::settings::Settings; + +// Re-export all public types so `crate::config::FooConfig` continues to work. +pub use self::agent::AgentConfig; +pub use self::builder::BuilderModeConfig; +pub use self::channels::{ChannelsConfig, CliConfig, GatewayConfig, HttpConfig}; +pub use self::database::{DatabaseBackend, DatabaseConfig, default_libsql_path}; +pub use self::embeddings::EmbeddingsConfig; +pub use self::heartbeat::HeartbeatConfig; +pub use self::llm::{ + AnthropicDirectConfig, LlmBackend, LlmConfig, NearAiApiMode, NearAiConfig, OllamaConfig, + OpenAiCompatibleConfig, OpenAiDirectConfig, TinfoilConfig, +}; +pub use self::routines::RoutineConfig; +pub use self::safety::SafetyConfig; +pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig}; +pub use self::secrets::SecretsConfig; +pub use self::skills::SkillsConfig; +pub use self::tunnel::TunnelConfig; +pub use self::wasm::WasmConfig; + +/// Thread-safe overlay for injected env vars (secrets loaded from DB). +/// +/// Used by `inject_llm_keys_from_secrets()` to make API keys available to +/// `optional_env()` without unsafe `set_var` calls. `optional_env()` checks +/// real env vars first, then falls back to this overlay. +static INJECTED_VARS: OnceLock> = OnceLock::new(); + +/// Main configuration for the agent. +#[derive(Debug, Clone)] +pub struct Config { + pub database: DatabaseConfig, + pub llm: LlmConfig, + pub embeddings: EmbeddingsConfig, + pub tunnel: TunnelConfig, + pub channels: ChannelsConfig, + pub agent: AgentConfig, + pub safety: SafetyConfig, + pub wasm: WasmConfig, + pub secrets: SecretsConfig, + pub builder: BuilderModeConfig, + pub heartbeat: HeartbeatConfig, + pub routines: RoutineConfig, + pub sandbox: SandboxModeConfig, + pub claude_code: ClaudeCodeConfig, + pub skills: SkillsConfig, + pub observability: crate::observability::ObservabilityConfig, +} + +impl Config { + /// Load configuration from environment variables and the database. + /// + /// Priority: env var > TOML config file > DB settings > default. + /// This is the primary way to load config after DB is connected. + pub async fn from_db( + store: &(dyn crate::db::SettingsStore + Sync), + user_id: &str, + ) -> Result { + Self::from_db_with_toml(store, user_id, None).await + } + + /// Load from DB with an optional TOML config file overlay. + pub async fn from_db_with_toml( + store: &(dyn crate::db::SettingsStore + Sync), + user_id: &str, + toml_path: Option<&std::path::Path>, + ) -> Result { + let _ = dotenvy::dotenv(); + crate::bootstrap::load_ironclaw_env(); + + // Load all settings from DB into a Settings struct + let mut db_settings = match store.get_all_settings(user_id).await { + Ok(map) => Settings::from_db_map(&map), + Err(e) => { + tracing::warn!("Failed to load settings from DB, using defaults: {}", e); + Settings::default() + } + }; + + // Overlay TOML config file (values win over DB settings) + Self::apply_toml_overlay(&mut db_settings, toml_path)?; + + Self::build(&db_settings).await + } + + /// Load configuration from environment variables only (no database). + /// + /// Used during early startup before the database is connected, + /// and by CLI commands that don't have DB access. + /// Falls back to legacy `settings.json` on disk if present. + /// + /// Loads both `./.env` (standard, higher priority) and `~/.ironclaw/.env` + /// (lower priority) via dotenvy, which never overwrites existing vars. + pub async fn from_env() -> Result { + Self::from_env_with_toml(None).await + } + + /// Load from env with an optional TOML config file overlay. + pub async fn from_env_with_toml( + toml_path: Option<&std::path::Path>, + ) -> Result { + let _ = dotenvy::dotenv(); + crate::bootstrap::load_ironclaw_env(); + let mut settings = Settings::load(); + + // Overlay TOML config file (values win over JSON settings) + Self::apply_toml_overlay(&mut settings, toml_path)?; + + Self::build(&settings).await + } + + /// Load and merge a TOML config file into settings. + /// + /// If `explicit_path` is `Some`, loads from that path (errors are fatal). + /// If `None`, tries the default path `~/.ironclaw/config.toml` (missing + /// file is silently ignored). + fn apply_toml_overlay( + settings: &mut Settings, + explicit_path: Option<&std::path::Path>, + ) -> Result<(), ConfigError> { + let path = explicit_path + .map(std::path::PathBuf::from) + .unwrap_or_else(Settings::default_toml_path); + + match Settings::load_toml(&path) { + Ok(Some(toml_settings)) => { + settings.merge_from(&toml_settings); + tracing::debug!("Loaded TOML config from {}", path.display()); + } + Ok(None) => { + if explicit_path.is_some() { + return Err(ConfigError::ParseError(format!( + "Config file not found: {}", + path.display() + ))); + } + } + Err(e) => { + if explicit_path.is_some() { + return Err(ConfigError::ParseError(format!( + "Failed to load config file {}: {}", + path.display(), + e + ))); + } + tracing::warn!("Failed to load default config file: {}", e); + } + } + Ok(()) + } + + /// Build config from settings (shared by from_env and from_db). + async fn build(settings: &Settings) -> Result { + Ok(Self { + database: DatabaseConfig::resolve()?, + llm: LlmConfig::resolve(settings)?, + embeddings: EmbeddingsConfig::resolve(settings)?, + tunnel: TunnelConfig::resolve(settings)?, + channels: ChannelsConfig::resolve(settings)?, + agent: AgentConfig::resolve(settings)?, + safety: SafetyConfig::resolve()?, + wasm: WasmConfig::resolve()?, + secrets: SecretsConfig::resolve().await?, + builder: BuilderModeConfig::resolve()?, + heartbeat: HeartbeatConfig::resolve(settings)?, + routines: RoutineConfig::resolve()?, + sandbox: SandboxModeConfig::resolve()?, + claude_code: ClaudeCodeConfig::resolve()?, + skills: SkillsConfig::resolve()?, + observability: crate::observability::ObservabilityConfig { + backend: std::env::var("OBSERVABILITY_BACKEND").unwrap_or_else(|_| "none".into()), + }, + }) + } +} + +/// Load API keys from the encrypted secrets store into a thread-safe overlay. +/// +/// This bridges the gap between secrets stored during onboarding and the +/// env-var-first resolution in `LlmConfig::resolve()`. Keys in the overlay +/// are read by `optional_env()` before falling back to `std::env::var()`, +/// so explicit env vars always win. +pub async fn inject_llm_keys_from_secrets( + secrets: &dyn crate::secrets::SecretsStore, + user_id: &str, +) { + let mappings = [ + ("llm_openai_api_key", "OPENAI_API_KEY"), + ("llm_anthropic_api_key", "ANTHROPIC_API_KEY"), + ("llm_compatible_api_key", "LLM_API_KEY"), + ]; + + let mut injected = HashMap::new(); + + for (secret_name, env_var) in mappings { + match std::env::var(env_var) { + Ok(val) if !val.is_empty() => continue, + _ => {} + } + match secrets.get_decrypted(user_id, secret_name).await { + Ok(decrypted) => { + injected.insert(env_var.to_string(), decrypted.expose().to_string()); + tracing::debug!("Loaded secret '{}' for env var '{}'", secret_name, env_var); + } + Err(_) => { + // Secret doesn't exist, that's fine + } + } + } + + let _ = INJECTED_VARS.set(injected); +} diff --git a/src/config/routines.rs b/src/config/routines.rs new file mode 100644 index 00000000..03b890de --- /dev/null +++ b/src/config/routines.rs @@ -0,0 +1,48 @@ +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; + +/// Routines configuration. +#[derive(Debug, Clone)] +pub struct RoutineConfig { + /// Whether the routines system is enabled. + pub enabled: bool, + /// How often (seconds) to poll for cron routines that need firing. + pub cron_check_interval_secs: u64, + /// Max routines executing concurrently across all users. + pub max_concurrent_routines: usize, + /// Default cooldown between fires (seconds). + pub default_cooldown_secs: u64, + /// Max output tokens for lightweight routine LLM calls. + pub max_lightweight_tokens: u32, +} + +impl Default for RoutineConfig { + fn default() -> Self { + Self { + enabled: true, + cron_check_interval_secs: 15, + max_concurrent_routines: 10, + default_cooldown_secs: 300, + max_lightweight_tokens: 4096, + } + } +} + +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), + 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)?, + max_lightweight_tokens: parse_optional_env("ROUTINES_MAX_TOKENS", 4096)?, + }) + } +} diff --git a/src/config/safety.rs b/src/config/safety.rs new file mode 100644 index 00000000..21483d73 --- /dev/null +++ b/src/config/safety.rs @@ -0,0 +1,25 @@ +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; + +/// Safety configuration. +#[derive(Debug, Clone)] +pub struct SafetyConfig { + pub max_output_length: usize, + pub injection_check_enabled: bool, +} + +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), + }) + } +} diff --git a/src/config/sandbox.rs b/src/config/sandbox.rs new file mode 100644 index 00000000..1473507e --- /dev/null +++ b/src/config/sandbox.rs @@ -0,0 +1,261 @@ +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; + +/// Docker sandbox configuration. +#[derive(Debug, Clone)] +pub struct SandboxModeConfig { + /// Whether the Docker sandbox is enabled. + pub enabled: bool, + /// Sandbox policy: "readonly", "workspace_write", or "full_access". + pub policy: String, + /// Command timeout in seconds. + pub timeout_secs: u64, + /// Memory limit in megabytes. + pub memory_limit_mb: u64, + /// CPU shares (relative weight). + pub cpu_shares: u32, + /// Docker image for the sandbox. + pub image: String, + /// Whether to auto-pull the image if not found. + pub auto_pull_image: bool, + /// Additional domains to allow through the network proxy. + pub extra_allowed_domains: Vec, +} + +impl Default for SandboxModeConfig { + fn default() -> Self { + Self { + enabled: true, + policy: "readonly".to_string(), + timeout_secs: 120, + memory_limit_mb: 2048, + cpu_shares: 1024, + image: "ghcr.io/nearai/sandbox:latest".to_string(), + auto_pull_image: true, + extra_allowed_domains: Vec::new(), + } + } +} + +impl SandboxModeConfig { + pub(crate) fn resolve() -> Result { + let extra_domains = optional_env("SANDBOX_EXTRA_DOMAINS")? + .map(|s| s.split(',').map(|d| d.trim().to_string()).collect()) + .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()), + 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(|| "ghcr.io/nearai/sandbox: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), + extra_allowed_domains: extra_domains, + }) + } + + /// Convert to SandboxConfig for the sandbox module. + pub fn to_sandbox_config(&self) -> crate::sandbox::SandboxConfig { + use crate::sandbox::SandboxPolicy; + use std::time::Duration; + + let policy = self.policy.parse().unwrap_or(SandboxPolicy::ReadOnly); + + let mut allowlist = crate::sandbox::default_allowlist(); + allowlist.extend(self.extra_allowed_domains.clone()); + + crate::sandbox::SandboxConfig { + enabled: self.enabled, + policy, + timeout: Duration::from_secs(self.timeout_secs), + memory_limit_mb: self.memory_limit_mb, + cpu_shares: self.cpu_shares, + network_allowlist: allowlist, + image: self.image.clone(), + auto_pull_image: self.auto_pull_image, + proxy_port: 0, // Auto-assign + } + } +} + +/// Claude Code sandbox configuration. +#[derive(Debug, Clone)] +pub struct ClaudeCodeConfig { + /// Whether Claude Code sandbox mode is available. + pub enabled: bool, + /// Host directory containing Claude auth config (not mounted into containers; + /// auth is handled via ANTHROPIC_API_KEY env var instead). + pub config_dir: std::path::PathBuf, + /// Claude model to use (e.g. "sonnet", "opus"). + pub model: String, + /// Maximum agentic turns before stopping. + pub max_turns: u32, + /// Memory limit in MB for Claude Code containers (heavier than workers). + pub memory_limit_mb: u64, + /// Allowed tool patterns for Claude Code permission settings. + /// + /// Written to `/workspace/.claude/settings.json` before spawning the CLI. + /// Provides defense-in-depth: only explicitly listed tools are auto-approved. + /// Any new/unknown tools would require interactive approval (which times out + /// in the non-interactive container, failing safely). + /// + /// Patterns follow Claude Code syntax: `"Bash(*)"`, `"Read"`, `"Edit(*)"`, etc. + pub allowed_tools: Vec, +} + +/// Default allowed tools for Claude Code inside containers. +/// +/// These cover all standard Claude Code tools needed for autonomous operation. +/// The Docker container provides the primary security boundary; this allowlist +/// provides defense-in-depth by preventing any future unknown tools from being +/// silently auto-approved. +fn default_claude_code_allowed_tools() -> Vec { + [ + // File system -- glob patterns match Claude Code's settings.json format + "Read(*)", + "Write(*)", + "Edit(*)", + "Glob(*)", + "Grep(*)", + "NotebookEdit(*)", + // Execution + "Bash(*)", + "Task(*)", + // Network + "WebFetch(*)", + "WebSearch(*)", + ] + .into_iter() + .map(String::from) + .collect() +} + +impl Default for ClaudeCodeConfig { + fn default() -> Self { + Self { + enabled: false, + config_dir: dirs::home_dir() + .unwrap_or_else(|| std::path::PathBuf::from(".")) + .join(".claude"), + model: "sonnet".to_string(), + max_turns: 50, + memory_limit_mb: 4096, + allowed_tools: default_claude_code_allowed_tools(), + } + } +} + +impl ClaudeCodeConfig { + /// Load from environment variables only (used inside containers where + /// there is no database or full config). + pub fn from_env() -> Self { + match Self::resolve() { + Ok(c) => c, + Err(e) => { + tracing::warn!("Failed to resolve ClaudeCodeConfig: {e}, using defaults"); + Self::default() + } + } + } + + /// Extract the OAuth access token from the host's credential store. + /// + /// On macOS: reads from Keychain (`Claude Code-credentials` service). + /// On Linux: reads from `~/.claude/.credentials.json`. + /// + /// Returns the access token if found. The token typically expires in + /// 8-12 hours, which is sufficient for any single container job. + pub fn extract_oauth_token() -> Option { + // macOS: extract from Keychain + if cfg!(target_os = "macos") { + match std::process::Command::new("security") + .args([ + "find-generic-password", + "-s", + "Claude Code-credentials", + "-w", + ]) + .output() + { + Ok(output) if output.status.success() => { + if let Ok(json) = String::from_utf8(output.stdout) { + return parse_oauth_access_token(json.trim()); + } + } + Ok(_) => { + tracing::debug!("No Claude Code credentials in macOS Keychain"); + } + Err(e) => { + tracing::debug!("Failed to query macOS Keychain: {e}"); + } + } + } + + // Linux / fallback: read from ~/.claude/.credentials.json + if let Some(home) = dirs::home_dir() { + let creds_path = home.join(".claude").join(".credentials.json"); + if let Ok(json) = std::fs::read_to_string(&creds_path) { + return parse_oauth_access_token(&json); + } + } + + None + } + + 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), + 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), + max_turns: parse_optional_env("CLAUDE_CODE_MAX_TURNS", defaults.max_turns)?, + memory_limit_mb: parse_optional_env( + "CLAUDE_CODE_MEMORY_LIMIT_MB", + defaults.memory_limit_mb, + )?, + allowed_tools: optional_env("CLAUDE_CODE_ALLOWED_TOOLS")? + .map(|s| { + s.split(',') + .map(|t| t.trim().to_string()) + .filter(|t| !t.is_empty()) + .collect() + }) + .unwrap_or(defaults.allowed_tools), + }) + } +} + +/// Parse the OAuth access token from a Claude Code credentials JSON blob. +/// +/// Expected shape: `{"claudeAiOauth": {"accessToken": "sk-ant-oat01-..."}}` +fn parse_oauth_access_token(json: &str) -> Option { + let creds: serde_json::Value = serde_json::from_str(json).ok()?; + creds["claudeAiOauth"]["accessToken"] + .as_str() + .map(String::from) +} diff --git a/src/config/secrets.rs b/src/config/secrets.rs new file mode 100644 index 00000000..863fb07d --- /dev/null +++ b/src/config/secrets.rs @@ -0,0 +1,70 @@ +use secrecy::{ExposeSecret, SecretString}; + +use crate::config::helpers::optional_env; +use crate::error::ConfigError; + +/// Secrets management configuration. +#[derive(Clone, Default)] +pub struct SecretsConfig { + /// Master key for encrypting secrets. + pub master_key: Option, + /// Whether secrets management is enabled. + pub enabled: bool, + /// Source of the master key. + pub source: crate::settings::KeySource, +} + +impl std::fmt::Debug for SecretsConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SecretsConfig") + .field("master_key", &self.master_key.is_some()) + .field("enabled", &self.enabled) + .field("source", &self.source) + .finish() + } +} + +impl SecretsConfig { + /// Auto-detect secrets master key from env var, then OS keychain. + /// + /// Sequential probe: SECRETS_MASTER_KEY env var first, then OS keychain. + /// No saved "source" needed; just try each source in order. + pub(crate) async fn resolve() -> Result { + use crate::settings::KeySource; + + let (master_key, source) = if let Some(env_key) = optional_env("SECRETS_MASTER_KEY")? { + (Some(SecretString::from(env_key)), KeySource::Env) + } else { + // Probe the OS keychain; if a key is stored, use it + match crate::secrets::keychain::get_master_key().await { + Ok(key_bytes) => { + let key_hex: String = key_bytes.iter().map(|b| format!("{:02x}", b)).collect(); + (Some(SecretString::from(key_hex)), KeySource::Keychain) + } + Err(_) => (None, KeySource::None), + } + }; + + let enabled = master_key.is_some(); + + if let Some(ref key) = master_key + && key.expose_secret().len() < 32 + { + return Err(ConfigError::InvalidValue { + key: "SECRETS_MASTER_KEY".to_string(), + message: "must be at least 32 bytes for AES-256-GCM".to_string(), + }); + } + + Ok(Self { + master_key, + enabled, + source, + }) + } + + /// Get the master key if configured. + pub fn master_key(&self) -> Option<&SecretString> { + self.master_key.as_ref() + } +} diff --git a/src/config/skills.rs b/src/config/skills.rs new file mode 100644 index 00000000..71386e74 --- /dev/null +++ b/src/config/skills.rs @@ -0,0 +1,56 @@ +use std::path::PathBuf; + +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; + +/// Skills system configuration. +#[derive(Debug, Clone)] +pub struct SkillsConfig { + /// Whether the skills system is enabled. + pub enabled: bool, + /// Directory containing local skills (default: ~/.ironclaw/skills/). + pub local_dir: PathBuf, + /// Maximum number of skills that can be active simultaneously. + pub max_active_skills: usize, + /// Maximum total context tokens allocated to skill prompts. + pub max_context_tokens: usize, +} + +impl Default for SkillsConfig { + fn default() -> Self { + Self { + enabled: false, + local_dir: default_skills_dir(), + max_active_skills: 3, + max_context_tokens: 4000, + } + } +} + +/// Get the default skills directory (~/.ironclaw/skills/). +fn default_skills_dir() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("skills") +} + +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), + local_dir: optional_env("SKILLS_DIR")? + .map(PathBuf::from) + .unwrap_or_else(default_skills_dir), + max_active_skills: parse_optional_env("SKILLS_MAX_ACTIVE", 3)?, + max_context_tokens: parse_optional_env("SKILLS_MAX_CONTEXT_TOKENS", 4000)?, + }) + } +} diff --git a/src/config/tunnel.rs b/src/config/tunnel.rs new file mode 100644 index 00000000..7c175753 --- /dev/null +++ b/src/config/tunnel.rs @@ -0,0 +1,106 @@ +use crate::config::helpers::optional_env; +use crate::error::ConfigError; +use crate::settings::Settings; + +/// Tunnel configuration for exposing the agent to the internet. +/// +/// Used by channels and tools that need public webhook endpoints. +/// The tunnel URL is shared across all channels (Telegram, Slack, etc.). +/// +/// Two modes: +/// - **Static URL** (`TUNNEL_URL`): set the public URL directly (manual tunnel) +/// - **Managed provider** (`TUNNEL_PROVIDER`): lifecycle-managed tunnel process +/// +/// When a managed provider is configured _and_ no static URL is set, +/// the gateway starts the tunnel on boot and populates `public_url`. +#[derive(Debug, Clone, Default)] +pub struct TunnelConfig { + /// Public URL from tunnel provider (e.g., "https://abc123.ngrok.io"). + /// Set statically via `TUNNEL_URL` or populated at runtime by a managed tunnel. + pub public_url: Option, + /// Provider configuration for lifecycle-managed tunnels. + /// `None` when using a static URL or no tunnel at all. + pub provider: Option, +} + +impl TunnelConfig { + pub(crate) fn resolve(settings: &Settings) -> Result { + let public_url = optional_env("TUNNEL_URL")? + .or_else(|| settings.tunnel.public_url.clone().filter(|s| !s.is_empty())); + + if let Some(ref url) = public_url + && !url.starts_with("https://") + { + return Err(ConfigError::InvalidValue { + key: "TUNNEL_URL".to_string(), + message: "must start with https:// (webhooks require HTTPS)".to_string(), + }); + } + + // Resolve managed tunnel provider config. + // Priority: env var > settings > default (none). + let provider_name = optional_env("TUNNEL_PROVIDER")? + .or_else(|| settings.tunnel.provider.clone()) + .unwrap_or_default(); + + let provider = if provider_name.is_empty() || provider_name == "none" { + None + } else { + Some(crate::tunnel::TunnelProviderConfig { + provider: provider_name.clone(), + cloudflare: optional_env("TUNNEL_CF_TOKEN")? + .or_else(|| settings.tunnel.cf_token.clone()) + .map(|token| crate::tunnel::CloudflareTunnelConfig { token }), + tailscale: Some(crate::tunnel::TailscaleTunnelConfig { + funnel: optional_env("TUNNEL_TS_FUNNEL")? + .map(|s| s == "true" || s == "1") + .unwrap_or(settings.tunnel.ts_funnel), + hostname: optional_env("TUNNEL_TS_HOSTNAME")? + .or_else(|| settings.tunnel.ts_hostname.clone()), + }), + ngrok: { + let ngrok_domain = optional_env("TUNNEL_NGROK_DOMAIN")? + .or_else(|| settings.tunnel.ngrok_domain.clone()); + optional_env("TUNNEL_NGROK_TOKEN")? + .or_else(|| settings.tunnel.ngrok_token.clone()) + .map(|auth_token| crate::tunnel::NgrokTunnelConfig { + auth_token, + domain: ngrok_domain, + }) + }, + custom: { + let health_url = optional_env("TUNNEL_CUSTOM_HEALTH_URL")? + .or_else(|| settings.tunnel.custom_health_url.clone()); + let url_pattern = optional_env("TUNNEL_CUSTOM_URL_PATTERN")? + .or_else(|| settings.tunnel.custom_url_pattern.clone()); + optional_env("TUNNEL_CUSTOM_COMMAND")? + .or_else(|| settings.tunnel.custom_command.clone()) + .map(|start_command| crate::tunnel::CustomTunnelConfig { + start_command, + health_url, + url_pattern, + }) + }, + }) + }; + + Ok(Self { + public_url, + provider, + }) + } + + /// Check if a tunnel is configured (static URL or managed provider). + pub fn is_enabled(&self) -> bool { + self.public_url.is_some() || self.provider.is_some() + } + + /// Get the webhook URL for a given path. + pub fn webhook_url(&self, path: &str) -> Option { + self.public_url.as_ref().map(|base| { + let base = base.trim_end_matches('/'); + let path = path.trim_start_matches('/'); + format!("{}/{}", base, path) + }) + } +} diff --git a/src/config/wasm.rs b/src/config/wasm.rs new file mode 100644 index 00000000..5d13fa1d --- /dev/null +++ b/src/config/wasm.rs @@ -0,0 +1,99 @@ +use std::path::PathBuf; +use std::time::Duration; + +use crate::config::helpers::{optional_env, parse_optional_env}; +use crate::error::ConfigError; + +/// WASM sandbox configuration. +#[derive(Debug, Clone)] +pub struct WasmConfig { + /// Whether WASM tool execution is enabled. + pub enabled: bool, + /// Directory containing installed WASM tools (default: ~/.ironclaw/tools/). + pub tools_dir: PathBuf, + /// Default memory limit in bytes (default: 10 MB). + pub default_memory_limit: u64, + /// Default execution timeout in seconds (default: 60). + pub default_timeout_secs: u64, + /// Default fuel limit for CPU metering (default: 10M). + pub default_fuel_limit: u64, + /// Whether to cache compiled modules. + pub cache_compiled: bool, + /// Directory for compiled module cache. + pub cache_dir: Option, +} + +impl Default for WasmConfig { + fn default() -> Self { + Self { + enabled: true, + tools_dir: default_tools_dir(), + default_memory_limit: 10 * 1024 * 1024, // 10 MB + default_timeout_secs: 60, + default_fuel_limit: 10_000_000, + cache_compiled: true, + cache_dir: None, + } + } +} + +/// Get the default tools directory (~/.ironclaw/tools/). +fn default_tools_dir() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("tools") +} + +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), + tools_dir: optional_env("WASM_TOOLS_DIR")? + .map(PathBuf::from) + .unwrap_or_else(default_tools_dir), + default_memory_limit: parse_optional_env( + "WASM_DEFAULT_MEMORY_LIMIT", + 10 * 1024 * 1024, + )?, + 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_dir: optional_env("WASM_CACHE_DIR")?.map(PathBuf::from), + }) + } + + /// Convert to WasmRuntimeConfig. + pub fn to_runtime_config(&self) -> crate::tools::wasm::WasmRuntimeConfig { + use crate::tools::wasm::{FuelConfig, ResourceLimits, WasmRuntimeConfig}; + + WasmRuntimeConfig { + default_limits: ResourceLimits { + memory_bytes: self.default_memory_limit, + fuel: self.default_fuel_limit, + timeout: Duration::from_secs(self.default_timeout_secs), + }, + fuel_config: FuelConfig { + initial_fuel: self.default_fuel_limit, + enabled: true, + }, + cache_compiled: self.cache_compiled, + cache_dir: self.cache_dir.clone(), + optimization_level: wasmtime::OptLevel::Speed, + } + } +} diff --git a/src/db/libsql/conversations.rs b/src/db/libsql/conversations.rs new file mode 100644 index 00000000..c78e5c11 --- /dev/null +++ b/src/db/libsql/conversations.rs @@ -0,0 +1,354 @@ +//! Conversation-related ConversationStore implementation for LibSqlBackend. + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use libsql::params; +use uuid::Uuid; + +use super::{LibSqlBackend, fmt_ts, get_i64, get_json, get_opt_text, get_text, get_ts, opt_text}; +use crate::db::ConversationStore; +use crate::error::DatabaseError; +use crate::history::{ConversationMessage, ConversationSummary}; + +#[async_trait] +impl ConversationStore for LibSqlBackend { + async fn create_conversation( + &self, + channel: &str, + user_id: &str, + thread_id: Option<&str>, + ) -> Result { + let conn = self.connect().await?; + let id = Uuid::new_v4(); + conn.execute( + "INSERT INTO conversations (id, channel, user_id, thread_id) VALUES (?1, ?2, ?3, ?4)", + params![id.to_string(), channel, user_id, opt_text(thread_id)], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(id) + } + + async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute( + "UPDATE conversations SET last_activity = ?2 WHERE id = ?1", + params![id.to_string(), now], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn add_conversation_message( + &self, + conversation_id: Uuid, + role: &str, + content: &str, + ) -> Result { + let conn = self.connect().await?; + let id = Uuid::new_v4(); + conn.execute( + "INSERT INTO conversation_messages (id, conversation_id, role, content) VALUES (?1, ?2, ?3, ?4)", + params![id.to_string(), conversation_id.to_string(), role, content], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + self.touch_conversation(conversation_id).await?; + Ok(id) + } + + async fn ensure_conversation( + &self, + id: Uuid, + channel: &str, + user_id: &str, + thread_id: Option<&str>, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute( + r#" + INSERT INTO conversations (id, channel, user_id, thread_id) + VALUES (?1, ?2, ?3, ?4) + ON CONFLICT (id) DO UPDATE SET last_activity = ?5 + "#, + params![id.to_string(), channel, user_id, opt_text(thread_id), now], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn list_conversations_with_preview( + &self, + user_id: &str, + channel: &str, + limit: i64, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + r#" + SELECT + c.id, + c.started_at, + c.last_activity, + c.metadata, + (SELECT COUNT(*) FROM conversation_messages m WHERE m.conversation_id = c.id) AS message_count, + (SELECT substr(m2.content, 1, 100) + FROM conversation_messages m2 + WHERE m2.conversation_id = c.id AND m2.role = 'user' + ORDER BY m2.created_at ASC + LIMIT 1 + ) AS title + FROM conversations c + WHERE c.user_id = ?1 AND c.channel = ?2 + ORDER BY c.last_activity DESC + LIMIT ?3 + "#, + params![user_id, channel, limit], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut results = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + let metadata = get_json(&row, 3); + let thread_type = metadata + .get("thread_type") + .and_then(|v| v.as_str()) + .map(String::from); + results.push(ConversationSummary { + id: row + .get::(0) + .unwrap_or_default() + .parse() + .unwrap_or_default(), + started_at: get_ts(&row, 1), + last_activity: get_ts(&row, 2), + message_count: get_i64(&row, 4), + title: get_opt_text(&row, 5), + thread_type, + }); + } + Ok(results) + } + + async fn get_or_create_assistant_conversation( + &self, + user_id: &str, + channel: &str, + ) -> Result { + let conn = self.connect().await?; + // Try to find existing + let mut rows = conn + .query( + r#" + SELECT id FROM conversations + WHERE user_id = ?1 AND channel = ?2 + AND json_extract(metadata, '$.thread_type') = 'assistant' + LIMIT 1 + "#, + params![user_id, channel], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + if let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + let id_str: String = row.get(0).unwrap_or_default(); + return id_str + .parse() + .map_err(|_| DatabaseError::Serialization("Invalid UUID".to_string())); + } + + // Create new + let id = Uuid::new_v4(); + let metadata = serde_json::json!({"thread_type": "assistant", "title": "Assistant"}); + conn.execute( + "INSERT INTO conversations (id, channel, user_id, metadata) VALUES (?1, ?2, ?3, ?4)", + params![id.to_string(), channel, user_id, metadata.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(id) + } + + async fn create_conversation_with_metadata( + &self, + channel: &str, + user_id: &str, + metadata: &serde_json::Value, + ) -> Result { + let conn = self.connect().await?; + let id = Uuid::new_v4(); + conn.execute( + "INSERT INTO conversations (id, channel, user_id, metadata) VALUES (?1, ?2, ?3, ?4)", + params![id.to_string(), channel, user_id, metadata.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(id) + } + + async fn list_conversation_messages_paginated( + &self, + conversation_id: Uuid, + before: Option>, + limit: i64, + ) -> Result<(Vec, bool), DatabaseError> { + let conn = self.connect().await?; + let fetch_limit = limit + 1; + let cid = conversation_id.to_string(); + + let mut rows = if let Some(before_ts) = before { + conn.query( + r#" + SELECT id, role, content, created_at + FROM conversation_messages + WHERE conversation_id = ?1 AND created_at < ?2 + ORDER BY created_at DESC + LIMIT ?3 + "#, + params![cid, fmt_ts(&before_ts), fetch_limit], + ) + .await + } else { + conn.query( + r#" + SELECT id, role, content, created_at + FROM conversation_messages + WHERE conversation_id = ?1 + ORDER BY created_at DESC + LIMIT ?2 + "#, + params![cid, fetch_limit], + ) + .await + } + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut all = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + all.push(ConversationMessage { + id: get_text(&row, 0).parse().unwrap_or_default(), + role: get_text(&row, 1), + content: get_text(&row, 2), + created_at: get_ts(&row, 3), + }); + } + + let has_more = all.len() as i64 > limit; + all.truncate(limit as usize); + all.reverse(); // oldest first + Ok((all, has_more)) + } + + async fn update_conversation_metadata_field( + &self, + id: Uuid, + key: &str, + value: &serde_json::Value, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + // SQLite: use json_patch to merge the key + let patch = serde_json::json!({ key: value }); + conn.execute( + "UPDATE conversations SET metadata = json_patch(metadata, ?2) WHERE id = ?1", + params![id.to_string(), patch.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn get_conversation_metadata( + &self, + id: Uuid, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT metadata FROM conversations WHERE id = ?1", + params![id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(Some(get_json(&row, 0))), + None => Ok(None), + } + } + + async fn list_conversation_messages( + &self, + conversation_id: Uuid, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + r#" + SELECT id, role, content, created_at + FROM conversation_messages + WHERE conversation_id = ?1 + ORDER BY created_at ASC + "#, + params![conversation_id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut messages = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + messages.push(ConversationMessage { + id: get_text(&row, 0).parse().unwrap_or_default(), + role: get_text(&row, 1), + content: get_text(&row, 2), + created_at: get_ts(&row, 3), + }); + } + Ok(messages) + } + + async fn conversation_belongs_to_user( + &self, + conversation_id: Uuid, + user_id: &str, + ) -> Result { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT 1 FROM conversations WHERE id = ?1 AND user_id = ?2", + libsql::params![conversation_id.to_string(), user_id], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + let found = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(found.is_some()) + } +} diff --git a/src/db/libsql/jobs.rs b/src/db/libsql/jobs.rs new file mode 100644 index 00000000..0c46b231 --- /dev/null +++ b/src/db/libsql/jobs.rs @@ -0,0 +1,330 @@ +//! Job-related JobStore implementation for LibSqlBackend. + +use async_trait::async_trait; +use libsql::params; +use rust_decimal::Decimal; +use uuid::Uuid; + +use super::{ + LibSqlBackend, fmt_opt_ts, fmt_ts, get_decimal, get_i64, get_json, get_opt_decimal, + get_opt_text, get_opt_ts, get_text, get_ts, opt_text, opt_text_owned, parse_job_state, +}; +use crate::context::{ActionRecord, JobContext, JobState}; +use crate::db::JobStore; +use crate::error::DatabaseError; +use crate::history::LlmCallRecord; + +use chrono::Utc; + +#[async_trait] +impl JobStore for LibSqlBackend { + async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let status = ctx.state.to_string(); + let estimated_time_secs = ctx.estimated_duration.map(|d| d.as_secs() as i64); + + conn + .execute( + r#" + INSERT INTO agent_jobs ( + id, conversation_id, title, description, category, status, source, + budget_amount, budget_token, bid_amount, estimated_cost, estimated_time_secs, + actual_cost, repair_attempts, created_at, started_at, completed_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17) + ON CONFLICT (id) DO UPDATE SET + title = excluded.title, + description = excluded.description, + category = excluded.category, + status = excluded.status, + estimated_cost = excluded.estimated_cost, + estimated_time_secs = excluded.estimated_time_secs, + actual_cost = excluded.actual_cost, + repair_attempts = excluded.repair_attempts, + started_at = excluded.started_at, + completed_at = excluded.completed_at + "#, + params![ + ctx.job_id.to_string(), + opt_text_owned(ctx.conversation_id.map(|id| id.to_string())), + ctx.title.as_str(), + ctx.description.as_str(), + opt_text(ctx.category.as_deref()), + status, + "direct", + opt_text_owned(ctx.budget.map(|d| d.to_string())), + opt_text(ctx.budget_token.as_deref()), + opt_text_owned(ctx.bid_amount.map(|d| d.to_string())), + opt_text_owned(ctx.estimated_cost.map(|d| d.to_string())), + estimated_time_secs, + ctx.actual_cost.to_string(), + ctx.repair_attempts as i64, + fmt_ts(&ctx.created_at), + fmt_opt_ts(&ctx.started_at), + fmt_opt_ts(&ctx.completed_at), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn get_job(&self, id: Uuid) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + r#" + SELECT id, conversation_id, title, description, category, status, user_id, + budget_amount, budget_token, bid_amount, estimated_cost, estimated_time_secs, + actual_cost, repair_attempts, created_at, started_at, completed_at + FROM agent_jobs WHERE id = ?1 + "#, + params![id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => { + let status_str = get_text(&row, 5); + let state = parse_job_state(&status_str); + let estimated_time_secs: Option = row.get::(11).ok(); + + Ok(Some(JobContext { + job_id: get_text(&row, 0).parse().unwrap_or_default(), + state, + user_id: get_text(&row, 6), + conversation_id: get_opt_text(&row, 1).and_then(|s| s.parse().ok()), + title: get_text(&row, 2), + description: get_text(&row, 3), + category: get_opt_text(&row, 4), + budget: get_opt_decimal(&row, 7), + budget_token: get_opt_text(&row, 8), + bid_amount: get_opt_decimal(&row, 9), + estimated_cost: get_opt_decimal(&row, 10), + estimated_duration: estimated_time_secs + .map(|s| std::time::Duration::from_secs(s as u64)), + actual_cost: get_decimal(&row, 12), + total_tokens_used: 0, + max_tokens: 0, + repair_attempts: get_i64(&row, 13) as u32, + created_at: get_ts(&row, 14), + started_at: get_opt_ts(&row, 15), + completed_at: get_opt_ts(&row, 16), + transitions: Vec::new(), + metadata: serde_json::Value::Null, + extra_env: std::sync::Arc::new(std::collections::HashMap::new()), + })) + } + None => Ok(None), + } + } + + async fn update_job_status( + &self, + id: Uuid, + status: JobState, + failure_reason: Option<&str>, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + conn.execute( + "UPDATE agent_jobs SET status = ?2, failure_reason = ?3 WHERE id = ?1", + params![id.to_string(), status.to_string(), opt_text(failure_reason)], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute( + "UPDATE agent_jobs SET status = 'stuck', stuck_since = ?2 WHERE id = ?1", + params![id.to_string(), now], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn get_stuck_jobs(&self) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query("SELECT id FROM agent_jobs WHERE status = 'stuck'", ()) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut ids = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + if let Ok(id_str) = row.get::(0) + && let Ok(id) = id_str.parse() + { + ids.push(id); + } + } + Ok(ids) + } + + async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let duration_ms = action.duration.as_millis() as i64; + let warnings_json = serde_json::to_string(&action.sanitization_warnings) + .map_err(|e| DatabaseError::Serialization(e.to_string()))?; + + conn.execute( + r#" + INSERT INTO job_actions ( + id, job_id, sequence_num, tool_name, input, output_raw, output_sanitized, + sanitization_warnings, cost, duration_ms, success, error_message, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13) + "#, + params![ + action.id.to_string(), + job_id.to_string(), + action.sequence as i64, + action.tool_name.as_str(), + action.input.to_string(), + opt_text(action.output_raw.as_deref()), + opt_text_owned(action.output_sanitized.as_ref().map(|v| v.to_string())), + warnings_json, + opt_text_owned(action.cost.map(|d| d.to_string())), + duration_ms, + action.success as i64, + opt_text(action.error.as_deref()), + fmt_ts(&action.executed_at), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn get_job_actions(&self, job_id: Uuid) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + r#" + SELECT id, sequence_num, tool_name, input, output_raw, output_sanitized, + sanitization_warnings, cost, duration_ms, success, error_message, created_at + FROM job_actions WHERE job_id = ?1 ORDER BY sequence_num + "#, + params![job_id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut actions = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + let warnings: Vec = + serde_json::from_str(&get_text(&row, 6)).unwrap_or_default(); + actions.push(ActionRecord { + id: get_text(&row, 0).parse().unwrap_or_default(), + sequence: get_i64(&row, 1) as u32, + tool_name: get_text(&row, 2), + input: get_json(&row, 3), + output_raw: get_opt_text(&row, 4), + output_sanitized: get_opt_text(&row, 5).and_then(|s| serde_json::from_str(&s).ok()), + sanitization_warnings: warnings, + cost: get_opt_decimal(&row, 7), + duration: std::time::Duration::from_millis(get_i64(&row, 8) as u64), + success: get_i64(&row, 9) != 0, + error: get_opt_text(&row, 10), + executed_at: get_ts(&row, 11), + }); + } + Ok(actions) + } + + async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result { + let conn = self.connect().await?; + let id = Uuid::new_v4(); + conn.execute( + r#" + INSERT INTO llm_calls (id, job_id, conversation_id, provider, model, input_tokens, output_tokens, cost, purpose) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) + "#, + params![ + id.to_string(), + opt_text_owned(record.job_id.map(|id| id.to_string())), + opt_text_owned(record.conversation_id.map(|id| id.to_string())), + record.provider, + record.model, + record.input_tokens as i64, + record.output_tokens as i64, + record.cost.to_string(), + opt_text(record.purpose), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(id) + } + + async fn save_estimation_snapshot( + &self, + job_id: Uuid, + category: &str, + tool_names: &[String], + estimated_cost: Decimal, + estimated_time_secs: i32, + estimated_value: Decimal, + ) -> Result { + let conn = self.connect().await?; + let id = Uuid::new_v4(); + let tools_json = serde_json::to_string(tool_names) + .map_err(|e| DatabaseError::Serialization(e.to_string()))?; + + conn.execute( + r#" + INSERT INTO estimation_snapshots (id, job_id, category, tool_names, estimated_cost, estimated_time_secs, estimated_value) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) + "#, + params![ + id.to_string(), + job_id.to_string(), + category, + tools_json, + estimated_cost.to_string(), + estimated_time_secs as i64, + estimated_value.to_string(), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(id) + } + + async fn update_estimation_actuals( + &self, + id: Uuid, + actual_cost: Decimal, + actual_time_secs: i32, + actual_value: Option, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + conn.execute( + "UPDATE estimation_snapshots SET actual_cost = ?2, actual_time_secs = ?3, actual_value = ?4 WHERE id = ?1", + params![ + id.to_string(), + actual_cost.to_string(), + actual_time_secs as i64, + actual_value.map(|d| d.to_string()).unwrap_or_default(), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } +} diff --git a/src/db/libsql/mod.rs b/src/db/libsql/mod.rs new file mode 100644 index 00000000..d83fdbe8 --- /dev/null +++ b/src/db/libsql/mod.rs @@ -0,0 +1,460 @@ +//! libSQL/Turso backend for the Database trait. +//! +//! Provides an embedded SQLite-compatible database using Turso's libSQL fork. +//! Supports three modes: +//! - Local embedded (file-based, no server needed) +//! - Turso cloud with embedded replica (sync to cloud) +//! - In-memory (for testing) + +mod conversations; +mod jobs; +mod routines; +mod sandbox; +mod settings; +mod tool_failures; +mod workspace; + +use std::path::Path; +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::{DateTime, NaiveDateTime, Utc}; +use libsql::{Connection, Database as LibSqlDatabase}; +use rust_decimal::Decimal; + +use crate::agent::routine::{ + NotifyConfig, Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger, +}; +use crate::context::JobState; +use crate::db::Database; +use crate::error::DatabaseError; +use crate::workspace::MemoryDocument; + +use crate::db::libsql_migrations; + +/// Explicit column list for routines table (matches positional access in `row_to_routine_libsql`). +pub(crate) const ROUTINE_COLUMNS: &str = "\ + id, name, description, user_id, enabled, \ + trigger_type, trigger_config, action_type, action_config, \ + cooldown_secs, max_concurrent, dedup_window_secs, \ + notify_channel, notify_user, notify_on_success, notify_on_failure, notify_on_attention, \ + state, last_run_at, next_fire_at, run_count, consecutive_failures, \ + created_at, updated_at"; + +/// Explicit column list for routine_runs table (matches positional access in `row_to_routine_run_libsql`). +pub(crate) const ROUTINE_RUN_COLUMNS: &str = "\ + id, routine_id, trigger_type, trigger_detail, started_at, \ + status, completed_at, result_summary, tokens_used, job_id, created_at"; + +/// libSQL/Turso database backend. +/// +/// Stores the `Database` handle in an `Arc` so that the same underlying +/// database can be shared with stores (SecretsStore, WasmToolStore) that +/// create their own connections per-operation. +pub struct LibSqlBackend { + db: Arc, +} + +impl LibSqlBackend { + /// Create a new local embedded database. + pub async fn new_local(path: &Path) -> Result { + // Ensure parent directory exists + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).map_err(|e| { + DatabaseError::Pool(format!("Failed to create database directory: {}", e)) + })?; + } + + let db = libsql::Builder::new_local(path) + .build() + .await + .map_err(|e| DatabaseError::Pool(format!("Failed to open libSQL database: {}", e)))?; + + Ok(Self { db: Arc::new(db) }) + } + + /// Create a new in-memory database (for testing). + pub async fn new_memory() -> Result { + let db = libsql::Builder::new_local(":memory:") + .build() + .await + .map_err(|e| { + DatabaseError::Pool(format!("Failed to create in-memory database: {}", e)) + })?; + + Ok(Self { db: Arc::new(db) }) + } + + /// Create with Turso cloud sync (embedded replica). + pub async fn new_remote_replica( + path: &Path, + url: &str, + auth_token: &str, + ) -> Result { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).map_err(|e| { + DatabaseError::Pool(format!("Failed to create database directory: {}", e)) + })?; + } + + let db = libsql::Builder::new_remote_replica(path, url.to_string(), auth_token.to_string()) + .build() + .await + .map_err(|e| DatabaseError::Pool(format!("Failed to open remote replica: {}", e)))?; + + Ok(Self { db: Arc::new(db) }) + } + + /// Get a shared reference to the underlying database handle. + /// + /// Use this to pass the database to stores (SecretsStore, WasmToolStore) + /// that need to create their own connections per-operation. + pub fn shared_db(&self) -> Arc { + Arc::clone(&self.db) + } + + /// Create a new connection to the database. + /// + /// Sets `PRAGMA busy_timeout = 5000` on every connection so concurrent + /// writers wait up to 5 seconds instead of failing instantly with + /// "database is locked". + pub async fn connect(&self) -> Result { + let conn = self + .db + .connect() + .map_err(|e| DatabaseError::Pool(format!("Failed to create connection: {}", e)))?; + conn.query("PRAGMA busy_timeout = 5000", ()) + .await + .map_err(|e| DatabaseError::Pool(format!("Failed to set busy_timeout: {}", e)))?; + Ok(conn) + } +} + +// ==================== Helper functions ==================== + +/// Parse an ISO-8601 timestamp string from SQLite into DateTime. +/// +/// Tries multiple formats in order: +/// 1. RFC 3339 with timezone (e.g. `2024-01-15T10:30:00.123Z`) +/// 2. Naive datetime with fractional seconds (e.g. `2024-01-15 10:30:00.123`) +/// 3. Naive datetime without fractional seconds (e.g. `2024-01-15 10:30:00`) +/// +/// Returns an error if none of the formats match. +pub(crate) fn parse_timestamp(s: &str) -> Result, String> { + // RFC 3339 (our canonical write format) + if let Ok(dt) = DateTime::parse_from_rfc3339(s) { + return Ok(dt.with_timezone(&Utc)); + } + // Naive with fractional seconds (legacy or SQLite datetime() output) + if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") { + return Ok(ndt.and_utc()); + } + // Naive without fractional seconds (legacy format) + if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") { + return Ok(ndt.and_utc()); + } + Err(format!("unparseable timestamp: {:?}", s)) +} + +/// Format a DateTime for SQLite storage (RFC 3339 with millisecond precision). +pub(crate) fn fmt_ts(dt: &DateTime) -> String { + dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true) +} + +/// Format an optional DateTime. +pub(crate) fn fmt_opt_ts(dt: &Option>) -> libsql::Value { + match dt { + Some(dt) => libsql::Value::Text(fmt_ts(dt)), + None => libsql::Value::Null, + } +} + +pub(crate) fn parse_job_state(s: &str) -> JobState { + match s { + "pending" => JobState::Pending, + "in_progress" => JobState::InProgress, + "completed" => JobState::Completed, + "submitted" => JobState::Submitted, + "accepted" => JobState::Accepted, + "failed" => JobState::Failed, + "stuck" => JobState::Stuck, + "cancelled" => JobState::Cancelled, + _ => JobState::Pending, + } +} + +/// Extract a text column from a libsql Row, returning empty string for NULL. +pub(crate) fn get_text(row: &libsql::Row, idx: i32) -> String { + row.get::(idx).unwrap_or_default() +} + +/// Extract an optional text column. +/// Returns None for SQL NULL, preserves empty strings as Some(""). +pub(crate) fn get_opt_text(row: &libsql::Row, idx: i32) -> Option { + row.get::(idx).ok() +} + +/// Convert an `Option<&str>` to a `libsql::Value` (Text or Null). +/// Use this instead of `.unwrap_or("")` to preserve NULL semantics. +pub(crate) fn opt_text(s: Option<&str>) -> libsql::Value { + match s { + Some(s) => libsql::Value::Text(s.to_string()), + None => libsql::Value::Null, + } +} + +/// Convert an `Option` to a `libsql::Value` (Text or Null). +pub(crate) fn opt_text_owned(s: Option) -> libsql::Value { + match s { + Some(s) => libsql::Value::Text(s), + None => libsql::Value::Null, + } +} + +/// Extract an i64 column, defaulting to 0. +pub(crate) fn get_i64(row: &libsql::Row, idx: i32) -> i64 { + row.get::(idx).unwrap_or(0) +} + +/// Extract an optional bool from an integer column. +pub(crate) fn get_opt_bool(row: &libsql::Row, idx: i32) -> Option { + row.get::(idx).ok().map(|v| v != 0) +} + +/// Parse a Decimal from a text column. +pub(crate) fn get_decimal(row: &libsql::Row, idx: i32) -> Decimal { + row.get::(idx) + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or_default() +} + +/// Parse an optional Decimal from a text column. +pub(crate) fn get_opt_decimal(row: &libsql::Row, idx: i32) -> Option { + row.get::(idx) + .ok() + .and_then(|s| s.parse::().ok()) +} + +/// Parse a JSON value from a text column. +pub(crate) fn get_json(row: &libsql::Row, idx: i32) -> serde_json::Value { + row.get::(idx) + .ok() + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or(serde_json::Value::Null) +} + +/// Parse a timestamp from a text column. +/// +/// If the column is NULL or the value cannot be parsed, logs a warning and +/// returns the Unix epoch (1970-01-01T00:00:00Z) so the error is detectable +/// rather than silently replaced by the current time. +pub(crate) fn get_ts(row: &libsql::Row, idx: i32) -> DateTime { + match row.get::(idx) { + Ok(s) => match parse_timestamp(&s) { + Ok(dt) => dt, + Err(e) => { + tracing::warn!("Timestamp parse failure at column {}: {}", idx, e); + DateTime::UNIX_EPOCH + } + }, + Err(_) => DateTime::UNIX_EPOCH, + } +} + +/// Parse an optional timestamp from a text column. +/// +/// Returns None if the column is NULL. Logs a warning and returns None if the +/// value is present but cannot be parsed. +pub(crate) fn get_opt_ts(row: &libsql::Row, idx: i32) -> Option> { + match row.get::(idx) { + Ok(s) if s.is_empty() => None, + Ok(s) => match parse_timestamp(&s) { + Ok(dt) => Some(dt), + Err(e) => { + tracing::warn!("Timestamp parse failure at column {}: {}", idx, e); + None + } + }, + Err(_) => None, + } +} + +#[async_trait] +impl Database for LibSqlBackend { + async fn run_migrations(&self) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + // WAL mode persists in the database file: all future connections benefit. + // Readers no longer block writers and vice versa. + conn.query("PRAGMA journal_mode=WAL", ()) + .await + .map_err(|e| DatabaseError::Migration(format!("Failed to enable WAL mode: {}", e)))?; + conn.execute_batch(libsql_migrations::SCHEMA) + .await + .map_err(|e| DatabaseError::Migration(format!("libSQL migration failed: {}", e)))?; + Ok(()) + } +} + +// ==================== Row conversion helpers ==================== + +pub(crate) fn row_to_memory_document(row: &libsql::Row) -> MemoryDocument { + MemoryDocument { + id: get_text(row, 0).parse().unwrap_or_default(), + user_id: get_text(row, 1), + agent_id: get_opt_text(row, 2).and_then(|s| s.parse().ok()), + path: get_text(row, 3), + content: get_text(row, 4), + created_at: get_ts(row, 5), + updated_at: get_ts(row, 6), + metadata: get_json(row, 7), + } +} + +pub(crate) fn row_to_routine_libsql(row: &libsql::Row) -> Result { + let trigger_type = get_text(row, 5); + let trigger_config = get_json(row, 6); + let action_type = get_text(row, 7); + let action_config = get_json(row, 8); + let cooldown_secs = get_i64(row, 9); + let max_concurrent = get_i64(row, 10); + let dedup_window_secs: Option = row.get::(11).ok(); + + let trigger = + Trigger::from_db(&trigger_type, trigger_config).map_err(DatabaseError::Serialization)?; + let action = RoutineAction::from_db(&action_type, action_config) + .map_err(DatabaseError::Serialization)?; + + Ok(Routine { + id: get_text(row, 0).parse().unwrap_or_default(), + name: get_text(row, 1), + description: get_text(row, 2), + user_id: get_text(row, 3), + enabled: get_i64(row, 4) != 0, + trigger, + action, + guardrails: RoutineGuardrails { + cooldown: std::time::Duration::from_secs(cooldown_secs as u64), + max_concurrent: max_concurrent as u32, + dedup_window: dedup_window_secs.map(|s| std::time::Duration::from_secs(s as u64)), + }, + notify: NotifyConfig { + channel: get_opt_text(row, 12), + user: get_text(row, 13), + on_success: get_i64(row, 14) != 0, + on_failure: get_i64(row, 15) != 0, + on_attention: get_i64(row, 16) != 0, + }, + state: get_json(row, 17), + last_run_at: get_opt_ts(row, 18), + next_fire_at: get_opt_ts(row, 19), + run_count: get_i64(row, 20) as u64, + consecutive_failures: get_i64(row, 21) as u32, + created_at: get_ts(row, 22), + updated_at: get_ts(row, 23), + }) +} + +pub(crate) fn row_to_routine_run_libsql(row: &libsql::Row) -> Result { + let status_str = get_text(row, 5); + let status: RunStatus = status_str + .parse() + .map_err(|e: String| DatabaseError::Serialization(e))?; + + Ok(RoutineRun { + id: get_text(row, 0).parse().unwrap_or_default(), + routine_id: get_text(row, 1).parse().unwrap_or_default(), + trigger_type: get_text(row, 2), + trigger_detail: get_opt_text(row, 3), + started_at: get_ts(row, 4), + completed_at: get_opt_ts(row, 6), + status, + result_summary: get_opt_text(row, 7), + tokens_used: row.get::(8).ok().map(|v| v as i32), + job_id: get_opt_text(row, 9).and_then(|s| s.parse().ok()), + created_at: get_ts(row, 10), + }) +} + +#[cfg(test)] +mod tests { + use crate::db::Database; + use crate::db::libsql::LibSqlBackend; + + #[tokio::test] + async fn test_wal_mode_after_migrations() { + let backend = LibSqlBackend::new_memory().await.unwrap(); + backend.run_migrations().await.unwrap(); + + let conn = backend.connect().await.unwrap(); + let mut rows = conn.query("PRAGMA journal_mode", ()).await.unwrap(); + let row = rows.next().await.unwrap().unwrap(); + let mode: String = row.get(0).unwrap(); + // In-memory databases use "memory" journal mode (WAL doesn't apply to :memory:), + // but the PRAGMA still executes without error. For file-based databases it returns "wal". + assert!( + mode == "wal" || mode == "memory", + "expected wal or memory, got: {}", + mode, + ); + } + + #[tokio::test] + async fn test_busy_timeout_set_on_connect() { + let backend = LibSqlBackend::new_memory().await.unwrap(); + backend.run_migrations().await.unwrap(); + + let conn = backend.connect().await.unwrap(); + let mut rows = conn.query("PRAGMA busy_timeout", ()).await.unwrap(); + let row = rows.next().await.unwrap().unwrap(); + let timeout: i64 = row.get(0).unwrap(); + assert_eq!(timeout, 5000); + } + + #[tokio::test] + async fn test_concurrent_writes_succeed() { + // Use a temp file so connections share state (in-memory DBs are connection-local) + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("test_concurrent.db"); + let backend = LibSqlBackend::new_local(&db_path).await.unwrap(); + backend.run_migrations().await.unwrap(); + + // Spawn 20 concurrent inserts into the conversations table + let mut handles = Vec::new(); + for i in 0..20 { + let conn = backend.connect().await.unwrap(); + let handle = tokio::spawn(async move { + let id = uuid::Uuid::new_v4().to_string(); + let val = format!("ch_{}", i); + conn.execute( + "INSERT INTO conversations (id, channel, user_id) VALUES (?1, ?2, ?3)", + libsql::params![id, val, "test_user"], + ) + .await + }); + handles.push(handle); + } + + for handle in handles { + let result = handle.await.unwrap(); + assert!( + result.is_ok(), + "concurrent write failed: {:?}", + result.err() + ); + } + + // Verify all 20 rows landed + let conn = backend.connect().await.unwrap(); + let mut rows = conn + .query( + "SELECT COUNT(*) FROM conversations WHERE user_id = ?1", + libsql::params!["test_user"], + ) + .await + .unwrap(); + let row = rows.next().await.unwrap().unwrap(); + let count: i64 = row.get(0).unwrap(); + assert_eq!(count, 20); + } +} diff --git a/src/db/libsql/routines.rs b/src/db/libsql/routines.rs new file mode 100644 index 00000000..48028c13 --- /dev/null +++ b/src/db/libsql/routines.rs @@ -0,0 +1,390 @@ +//! Routine-related RoutineStore implementation for LibSqlBackend. + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use libsql::params; +use uuid::Uuid; + +use super::{ + LibSqlBackend, ROUTINE_COLUMNS, ROUTINE_RUN_COLUMNS, fmt_opt_ts, fmt_ts, get_i64, opt_text, + opt_text_owned, row_to_routine_libsql, row_to_routine_run_libsql, +}; +use crate::agent::routine::{Routine, RoutineRun, RunStatus}; +use crate::db::RoutineStore; +use crate::error::DatabaseError; + +#[async_trait] +impl RoutineStore for LibSqlBackend { + async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let trigger_type = routine.trigger.type_tag(); + let trigger_config = routine.trigger.to_config_json(); + let action_type = routine.action.type_tag(); + let action_config = routine.action.to_config_json(); + let cooldown_secs = routine.guardrails.cooldown.as_secs() as i64; + let max_concurrent = routine.guardrails.max_concurrent as i64; + let dedup_window_secs = routine.guardrails.dedup_window.map(|d| d.as_secs() as i64); + + conn.execute( + r#" + INSERT INTO routines ( + id, name, description, user_id, enabled, + trigger_type, trigger_config, action_type, action_config, + cooldown_secs, max_concurrent, dedup_window_secs, + notify_channel, notify_user, notify_on_success, notify_on_failure, notify_on_attention, + state, next_fire_at, created_at, updated_at + ) VALUES ( + ?1, ?2, ?3, ?4, ?5, + ?6, ?7, ?8, ?9, + ?10, ?11, ?12, + ?13, ?14, ?15, ?16, ?17, + ?18, ?19, ?20, ?21 + ) + "#, + params![ + routine.id.to_string(), + routine.name.as_str(), + routine.description.as_str(), + routine.user_id.as_str(), + routine.enabled as i64, + trigger_type, + trigger_config.to_string(), + action_type, + action_config.to_string(), + cooldown_secs, + max_concurrent, + dedup_window_secs, + opt_text(routine.notify.channel.as_deref()), + routine.notify.user.as_str(), + routine.notify.on_success as i64, + routine.notify.on_failure as i64, + routine.notify.on_attention as i64, + routine.state.to_string(), + fmt_opt_ts(&routine.next_fire_at), + fmt_ts(&routine.created_at), + fmt_ts(&routine.updated_at), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn get_routine(&self, id: Uuid) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + &format!("SELECT {} FROM routines WHERE id = ?1", ROUTINE_COLUMNS), + params![id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(Some(row_to_routine_libsql(&row)?)), + None => Ok(None), + } + } + + async fn get_routine_by_name( + &self, + user_id: &str, + name: &str, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + &format!( + "SELECT {} FROM routines WHERE user_id = ?1 AND name = ?2", + ROUTINE_COLUMNS + ), + params![user_id, name], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(Some(row_to_routine_libsql(&row)?)), + None => Ok(None), + } + } + + async fn list_routines(&self, user_id: &str) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + &format!( + "SELECT {} FROM routines WHERE user_id = ?1 ORDER BY name", + ROUTINE_COLUMNS + ), + params![user_id], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut routines = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + routines.push(row_to_routine_libsql(&row)?); + } + Ok(routines) + } + + async fn list_event_routines(&self) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + &format!( + "SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'event'", + ROUTINE_COLUMNS + ), + (), + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut routines = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + routines.push(row_to_routine_libsql(&row)?); + } + Ok(routines) + } + + async fn list_due_cron_routines(&self) -> Result, DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + let mut rows = conn + .query( + &format!( + "SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'cron' AND next_fire_at IS NOT NULL AND next_fire_at <= ?1", + ROUTINE_COLUMNS + ), + params![now], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut routines = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + routines.push(row_to_routine_libsql(&row)?); + } + Ok(routines) + } + + async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let trigger_type = routine.trigger.type_tag(); + let trigger_config = routine.trigger.to_config_json(); + let action_type = routine.action.type_tag(); + let action_config = routine.action.to_config_json(); + let cooldown_secs = routine.guardrails.cooldown.as_secs() as i64; + let max_concurrent = routine.guardrails.max_concurrent as i64; + let dedup_window_secs = routine.guardrails.dedup_window.map(|d| d.as_secs() as i64); + let now = fmt_ts(&Utc::now()); + + conn.execute( + r#" + UPDATE routines SET + name = ?2, description = ?3, enabled = ?4, + trigger_type = ?5, trigger_config = ?6, + action_type = ?7, action_config = ?8, + cooldown_secs = ?9, max_concurrent = ?10, dedup_window_secs = ?11, + notify_channel = ?12, notify_user = ?13, + notify_on_success = ?14, notify_on_failure = ?15, notify_on_attention = ?16, + state = ?17, next_fire_at = ?18, + updated_at = ?19 + WHERE id = ?1 + "#, + params![ + routine.id.to_string(), + routine.name.as_str(), + routine.description.as_str(), + routine.enabled as i64, + trigger_type, + trigger_config.to_string(), + action_type, + action_config.to_string(), + cooldown_secs, + max_concurrent, + dedup_window_secs, + opt_text(routine.notify.channel.as_deref()), + routine.notify.user.as_str(), + routine.notify.on_success as i64, + routine.notify.on_failure as i64, + routine.notify.on_attention as i64, + routine.state.to_string(), + fmt_opt_ts(&routine.next_fire_at), + now, + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn update_routine_runtime( + &self, + id: Uuid, + last_run_at: DateTime, + next_fire_at: Option>, + run_count: u64, + consecutive_failures: u32, + state: &serde_json::Value, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute( + r#" + UPDATE routines SET + last_run_at = ?2, next_fire_at = ?3, + run_count = ?4, consecutive_failures = ?5, + state = ?6, updated_at = ?7 + WHERE id = ?1 + "#, + params![ + id.to_string(), + fmt_ts(&last_run_at), + fmt_opt_ts(&next_fire_at), + run_count as i64, + consecutive_failures as i64, + state.to_string(), + now, + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn delete_routine(&self, id: Uuid) -> Result { + let conn = self.connect().await?; + let count = conn + .execute( + "DELETE FROM routines WHERE id = ?1", + params![id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(count > 0) + } + + async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + conn.execute( + r#" + INSERT INTO routine_runs ( + id, routine_id, trigger_type, trigger_detail, + started_at, status, job_id + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) + "#, + params![ + run.id.to_string(), + run.routine_id.to_string(), + run.trigger_type.as_str(), + opt_text(run.trigger_detail.as_deref()), + fmt_ts(&run.started_at), + run.status.to_string(), + opt_text_owned(run.job_id.map(|id| id.to_string())), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn complete_routine_run( + &self, + id: Uuid, + status: RunStatus, + result_summary: Option<&str>, + tokens_used: Option, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute( + r#" + UPDATE routine_runs SET + completed_at = ?5, status = ?2, + result_summary = ?3, tokens_used = ?4 + WHERE id = ?1 + "#, + params![ + id.to_string(), + status.to_string(), + opt_text(result_summary), + tokens_used.map(|t| t as i64), + now, + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn list_routine_runs( + &self, + routine_id: Uuid, + limit: i64, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + &format!( + "SELECT {} FROM routine_runs WHERE routine_id = ?1 ORDER BY started_at DESC LIMIT ?2", + ROUTINE_RUN_COLUMNS + ), + params![routine_id.to_string(), limit], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut runs = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + runs.push(row_to_routine_run_libsql(&row)?); + } + Ok(runs) + } + + async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT COUNT(*) as cnt FROM routine_runs WHERE routine_id = ?1 AND status = 'running'", + params![routine_id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(get_i64(&row, 0)), + None => Ok(0), + } + } +} diff --git a/src/db/libsql/sandbox.rs b/src/db/libsql/sandbox.rs new file mode 100644 index 00000000..f7449608 --- /dev/null +++ b/src/db/libsql/sandbox.rs @@ -0,0 +1,405 @@ +//! Sandbox-related SandboxStore implementation for LibSqlBackend. + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use libsql::params; +use uuid::Uuid; + +use super::{ + LibSqlBackend, fmt_opt_ts, fmt_ts, get_i64, get_json, get_opt_bool, get_opt_text, get_opt_ts, + get_text, get_ts, opt_text, +}; +use crate::db::SandboxStore; +use crate::error::DatabaseError; +use crate::history::{JobEventRecord, SandboxJobRecord, SandboxJobSummary}; + +#[async_trait] +impl SandboxStore for LibSqlBackend { + async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + conn.execute( + r#" + INSERT INTO agent_jobs ( + id, title, description, status, source, user_id, project_dir, + success, failure_reason, created_at, started_at, completed_at + ) VALUES (?1, ?2, ?3, ?4, 'sandbox', ?5, ?6, ?7, ?8, ?9, ?10, ?11) + ON CONFLICT (id) DO UPDATE SET + status = excluded.status, + success = excluded.success, + failure_reason = excluded.failure_reason, + started_at = excluded.started_at, + completed_at = excluded.completed_at + "#, + params![ + job.id.to_string(), + job.task.as_str(), + job.credential_grants_json.as_str(), + job.status.as_str(), + job.user_id.as_str(), + job.project_dir.as_str(), + job.success.map(|b| b as i64), + opt_text(job.failure_reason.as_deref()), + fmt_ts(&job.created_at), + fmt_opt_ts(&job.started_at), + fmt_opt_ts(&job.completed_at), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn get_sandbox_job(&self, id: Uuid) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + r#" + SELECT id, title, description, status, user_id, project_dir, + success, failure_reason, created_at, started_at, completed_at + FROM agent_jobs WHERE id = ?1 AND source = 'sandbox' + "#, + params![id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(Some(SandboxJobRecord { + id: get_text(&row, 0).parse().unwrap_or_default(), + task: get_text(&row, 1), + credential_grants_json: get_text(&row, 2), + status: get_text(&row, 3), + user_id: get_text(&row, 4), + project_dir: get_text(&row, 5), + success: get_opt_bool(&row, 6), + failure_reason: get_opt_text(&row, 7), + created_at: get_ts(&row, 8), + started_at: get_opt_ts(&row, 9), + completed_at: get_opt_ts(&row, 10), + })), + None => Ok(None), + } + } + + async fn list_sandbox_jobs(&self) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + r#" + SELECT id, title, description, status, user_id, project_dir, + success, failure_reason, created_at, started_at, completed_at + FROM agent_jobs WHERE source = 'sandbox' + ORDER BY created_at DESC + "#, + (), + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut jobs = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + jobs.push(SandboxJobRecord { + id: get_text(&row, 0).parse().unwrap_or_default(), + task: get_text(&row, 1), + credential_grants_json: get_text(&row, 2), + status: get_text(&row, 3), + user_id: get_text(&row, 4), + project_dir: get_text(&row, 5), + success: get_opt_bool(&row, 6), + failure_reason: get_opt_text(&row, 7), + created_at: get_ts(&row, 8), + started_at: get_opt_ts(&row, 9), + completed_at: get_opt_ts(&row, 10), + }); + } + Ok(jobs) + } + + async fn update_sandbox_job_status( + &self, + id: Uuid, + status: &str, + success: Option, + message: Option<&str>, + started_at: Option>, + completed_at: Option>, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + conn.execute( + r#" + UPDATE agent_jobs SET + status = ?2, + success = COALESCE(?3, success), + failure_reason = COALESCE(?4, failure_reason), + started_at = COALESCE(?5, started_at), + completed_at = COALESCE(?6, completed_at) + WHERE id = ?1 AND source = 'sandbox' + "#, + params![ + id.to_string(), + status, + success.map(|b| b as i64), + message, + fmt_opt_ts(&started_at), + fmt_opt_ts(&completed_at), + ], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn cleanup_stale_sandbox_jobs(&self) -> Result { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + let count = conn + .execute( + r#" + UPDATE agent_jobs SET + status = 'interrupted', + failure_reason = 'Process restarted', + completed_at = ?1 + WHERE source = 'sandbox' AND status IN ('running', 'creating') + "#, + params![now], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + if count > 0 { + tracing::info!("Marked {} stale sandbox jobs as interrupted", count); + } + Ok(count) + } + + async fn sandbox_job_summary(&self) -> Result { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'sandbox' GROUP BY status", + (), + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut summary = SandboxJobSummary::default(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + let status = get_text(&row, 0); + let count = get_i64(&row, 1) as usize; + summary.total += count; + match status.as_str() { + "creating" => summary.creating += count, + "running" => summary.running += count, + "completed" => summary.completed += count, + "failed" => summary.failed += count, + "interrupted" => summary.interrupted += count, + _ => {} + } + } + Ok(summary) + } + + async fn list_sandbox_jobs_for_user( + &self, + user_id: &str, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + r#" + SELECT id, title, description, status, user_id, project_dir, + success, failure_reason, created_at, started_at, completed_at + FROM agent_jobs WHERE source = 'sandbox' AND user_id = ?1 + ORDER BY created_at DESC + "#, + libsql::params![user_id], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut jobs = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + jobs.push(SandboxJobRecord { + id: get_text(&row, 0).parse().unwrap_or_default(), + task: get_text(&row, 1), + credential_grants_json: get_text(&row, 2), + status: get_text(&row, 3), + user_id: get_text(&row, 4), + project_dir: get_text(&row, 5), + success: get_opt_bool(&row, 6), + failure_reason: get_opt_text(&row, 7), + created_at: get_ts(&row, 8), + started_at: get_opt_ts(&row, 9), + completed_at: get_opt_ts(&row, 10), + }); + } + Ok(jobs) + } + + async fn sandbox_job_summary_for_user( + &self, + user_id: &str, + ) -> Result { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'sandbox' AND user_id = ?1 GROUP BY status", + libsql::params![user_id], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut summary = SandboxJobSummary::default(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + let status = get_text(&row, 0); + let count = get_i64(&row, 1) as usize; + summary.total += count; + match status.as_str() { + "creating" => summary.creating += count, + "running" => summary.running += count, + "completed" => summary.completed += count, + "failed" => summary.failed += count, + "interrupted" => summary.interrupted += count, + _ => {} + } + } + Ok(summary) + } + + async fn sandbox_job_belongs_to_user( + &self, + job_id: Uuid, + user_id: &str, + ) -> Result { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT 1 FROM agent_jobs WHERE id = ?1 AND user_id = ?2 AND source = 'sandbox'", + libsql::params![job_id.to_string(), user_id], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + let found = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(found.is_some()) + } + + async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + conn.execute( + "UPDATE agent_jobs SET job_mode = ?2 WHERE id = ?1", + params![id.to_string(), mode], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn get_sandbox_job_mode(&self, id: Uuid) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT job_mode FROM agent_jobs WHERE id = ?1", + params![id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(Some(get_text(&row, 0))), + None => Ok(None), + } + } + + async fn save_job_event( + &self, + job_id: Uuid, + event_type: &str, + data: &serde_json::Value, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + conn.execute( + "INSERT INTO job_events (job_id, event_type, data) VALUES (?1, ?2, ?3)", + params![job_id.to_string(), event_type, data.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn list_job_events( + &self, + job_id: Uuid, + limit: Option, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = if let Some(n) = limit { + conn.query( + r#" + SELECT id, job_id, event_type, data, created_at + FROM ( + SELECT id, job_id, event_type, data, created_at + FROM job_events WHERE job_id = ?1 + ORDER BY id DESC + LIMIT ?2 + ) + ORDER BY id ASC + "#, + params![job_id.to_string(), n], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + } else { + conn.query( + r#" + SELECT id, job_id, event_type, data, created_at + FROM job_events WHERE job_id = ?1 ORDER BY id ASC + "#, + params![job_id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + }; + + let mut events = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + events.push(JobEventRecord { + id: get_i64(&row, 0), + job_id: get_text(&row, 1).parse().unwrap_or_default(), + event_type: get_text(&row, 2), + data: get_json(&row, 3), + created_at: get_ts(&row, 4), + }); + } + Ok(events) + } +} diff --git a/src/db/libsql/settings.rs b/src/db/libsql/settings.rs new file mode 100644 index 00000000..bf2703e9 --- /dev/null +++ b/src/db/libsql/settings.rs @@ -0,0 +1,208 @@ +//! Settings-related SettingsStore implementation for LibSqlBackend. + +use std::collections::HashMap; + +use async_trait::async_trait; +use libsql::params; + +use super::{LibSqlBackend, fmt_ts, get_i64, get_json, get_text, get_ts}; +use crate::db::SettingsStore; +use crate::error::DatabaseError; +use crate::history::SettingRow; + +use chrono::Utc; + +#[async_trait] +impl SettingsStore for LibSqlBackend { + async fn get_setting( + &self, + user_id: &str, + key: &str, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT value FROM settings WHERE user_id = ?1 AND key = ?2", + params![user_id, key], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(Some(get_json(&row, 0))), + None => Ok(None), + } + } + + async fn get_setting_full( + &self, + user_id: &str, + key: &str, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT key, value, updated_at FROM settings WHERE user_id = ?1 AND key = ?2", + params![user_id, key], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(Some(SettingRow { + key: get_text(&row, 0), + value: get_json(&row, 1), + updated_at: get_ts(&row, 2), + })), + None => Ok(None), + } + } + + async fn set_setting( + &self, + user_id: &str, + key: &str, + value: &serde_json::Value, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute( + r#" + INSERT INTO settings (user_id, key, value, updated_at) + VALUES (?1, ?2, ?3, ?4) + ON CONFLICT (user_id, key) DO UPDATE SET + value = excluded.value, + updated_at = ?4 + "#, + params![user_id, key, value.to_string(), now], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn delete_setting(&self, user_id: &str, key: &str) -> Result { + let conn = self.connect().await?; + let count = conn + .execute( + "DELETE FROM settings WHERE user_id = ?1 AND key = ?2", + params![user_id, key], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(count > 0) + } + + async fn list_settings(&self, user_id: &str) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT key, value, updated_at FROM settings WHERE user_id = ?1 ORDER BY key", + params![user_id], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut settings = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + settings.push(SettingRow { + key: get_text(&row, 0), + value: get_json(&row, 1), + updated_at: get_ts(&row, 2), + }); + } + Ok(settings) + } + + async fn get_all_settings( + &self, + user_id: &str, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT key, value FROM settings WHERE user_id = ?1", + params![user_id], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut map = HashMap::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + map.insert(get_text(&row, 0), get_json(&row, 1)); + } + Ok(map) + } + + async fn set_all_settings( + &self, + user_id: &str, + settings: &HashMap, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute("BEGIN", ()) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + for (key, value) in settings { + if let Err(e) = conn + .execute( + r#" + INSERT INTO settings (user_id, key, value, updated_at) + VALUES (?1, ?2, ?3, ?4) + ON CONFLICT (user_id, key) DO UPDATE SET + value = excluded.value, + updated_at = ?4 + "#, + params![user_id, key.as_str(), value.to_string(), now.as_str()], + ) + .await + { + let _ = conn.execute("ROLLBACK", ()).await; + return Err(DatabaseError::Query(e.to_string())); + } + } + + conn.execute("COMMIT", ()) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn has_settings(&self, user_id: &str) -> Result { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT COUNT(*) as cnt FROM settings WHERE user_id = ?1", + params![user_id], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + match rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Some(row) => Ok(get_i64(&row, 0) > 0), + None => Ok(false), + } + } +} diff --git a/src/db/libsql/tool_failures.rs b/src/db/libsql/tool_failures.rs new file mode 100644 index 00000000..80ccb354 --- /dev/null +++ b/src/db/libsql/tool_failures.rs @@ -0,0 +1,97 @@ +//! Tool failure-related ToolFailureStore implementation for LibSqlBackend. + +use async_trait::async_trait; +use libsql::params; +use uuid::Uuid; + +use super::{LibSqlBackend, fmt_ts, get_i64, get_opt_text, get_text, get_ts}; +use crate::agent::BrokenTool; +use crate::db::ToolFailureStore; +use crate::error::DatabaseError; + +use chrono::Utc; + +#[async_trait] +impl ToolFailureStore for LibSqlBackend { + async fn record_tool_failure( + &self, + tool_name: &str, + error_message: &str, + ) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute( + r#" + INSERT INTO tool_failures (id, tool_name, error_message, error_count, last_failure) + VALUES (?1, ?2, ?3, 1, ?4) + ON CONFLICT (tool_name) DO UPDATE SET + error_message = ?3, + error_count = tool_failures.error_count + 1, + last_failure = ?4 + "#, + params![Uuid::new_v4().to_string(), tool_name, error_message, now], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn get_broken_tools(&self, threshold: i32) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + r#" + SELECT tool_name, error_message, error_count, first_failure, last_failure, + last_build_result, repair_attempts + FROM tool_failures + WHERE error_count >= ?1 AND repaired_at IS NULL + ORDER BY error_count DESC + "#, + params![threshold as i64], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + let mut tools = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + tools.push(BrokenTool { + name: get_text(&row, 0), + last_error: get_opt_text(&row, 1), + failure_count: get_i64(&row, 2) as u32, + first_failure: get_ts(&row, 3), + last_failure: get_ts(&row, 4), + last_build_result: get_opt_text(&row, 5) + .and_then(|s| serde_json::from_str(&s).ok()), + repair_attempts: get_i64(&row, 6) as u32, + }); + } + Ok(tools) + } + + async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + let now = fmt_ts(&Utc::now()); + conn.execute( + "UPDATE tool_failures SET repaired_at = ?2, error_count = 0 WHERE tool_name = ?1", + params![tool_name, now], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } + + async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> { + let conn = self.connect().await?; + conn.execute( + "UPDATE tool_failures SET repair_attempts = repair_attempts + 1 WHERE tool_name = ?1", + params![tool_name], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + Ok(()) + } +} diff --git a/src/db/libsql/workspace.rs b/src/db/libsql/workspace.rs new file mode 100644 index 00000000..31c9da17 --- /dev/null +++ b/src/db/libsql/workspace.rs @@ -0,0 +1,607 @@ +//! Workspace-related WorkspaceStore implementation for LibSqlBackend. + +use std::collections::HashMap; + +use async_trait::async_trait; +use libsql::params; +use uuid::Uuid; + +use super::{ + LibSqlBackend, fmt_ts, get_i64, get_opt_text, get_opt_ts, get_text, get_ts, + row_to_memory_document, +}; +use crate::db::WorkspaceStore; +use crate::error::WorkspaceError; +use crate::workspace::{ + MemoryChunk, MemoryDocument, RankedResult, SearchConfig, SearchResult, WorkspaceEntry, + reciprocal_rank_fusion, +}; + +use chrono::Utc; + +#[async_trait] +impl WorkspaceStore for LibSqlBackend { + async fn get_document_by_path( + &self, + user_id: &str, + agent_id: Option, + path: &str, + ) -> Result { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let agent_id_str = agent_id.map(|id| id.to_string()); + let mut rows = conn + .query( + r#" + SELECT id, user_id, agent_id, path, content, + created_at, updated_at, metadata + FROM memory_documents + WHERE user_id = ?1 AND agent_id IS ?2 AND path = ?3 + "#, + params![user_id, agent_id_str.as_deref(), path], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })?; + + match rows + .next() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })? { + Some(row) => Ok(row_to_memory_document(&row)), + None => Err(WorkspaceError::DocumentNotFound { + doc_type: path.to_string(), + user_id: user_id.to_string(), + }), + } + } + + async fn get_document_by_id(&self, id: Uuid) -> Result { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let mut rows = conn + .query( + r#" + SELECT id, user_id, agent_id, path, content, + created_at, updated_at, metadata + FROM memory_documents WHERE id = ?1 + "#, + params![id.to_string()], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })?; + + match rows + .next() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })? { + Some(row) => Ok(row_to_memory_document(&row)), + None => Err(WorkspaceError::DocumentNotFound { + doc_type: "unknown".to_string(), + user_id: "unknown".to_string(), + }), + } + } + + async fn get_or_create_document_by_path( + &self, + user_id: &str, + agent_id: Option, + path: &str, + ) -> Result { + // Try get + match self.get_document_by_path(user_id, agent_id, path).await { + Ok(doc) => return Ok(doc), + Err(WorkspaceError::DocumentNotFound { .. }) => {} + Err(e) => return Err(e), + } + + // Create + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let id = Uuid::new_v4(); + let agent_id_str = agent_id.map(|id| id.to_string()); + conn.execute( + r#" + INSERT INTO memory_documents (id, user_id, agent_id, path, content, metadata) + VALUES (?1, ?2, ?3, ?4, '', '{}') + ON CONFLICT (user_id, agent_id, path) DO NOTHING + "#, + params![id.to_string(), user_id, agent_id_str.as_deref(), path], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Insert failed: {}", e), + })?; + + self.get_document_by_path(user_id, agent_id, path).await + } + + async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let now = fmt_ts(&Utc::now()); + conn.execute( + "UPDATE memory_documents SET content = ?2, updated_at = ?3 WHERE id = ?1", + params![id.to_string(), content, now], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Update failed: {}", e), + })?; + Ok(()) + } + + async fn delete_document_by_path( + &self, + user_id: &str, + agent_id: Option, + path: &str, + ) -> Result<(), WorkspaceError> { + let doc = self.get_document_by_path(user_id, agent_id, path).await?; + self.delete_chunks(doc.id).await?; + + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let agent_id_str = agent_id.map(|id| id.to_string()); + conn.execute( + "DELETE FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 AND path = ?3", + params![user_id, agent_id_str.as_deref(), path], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Delete failed: {}", e), + })?; + Ok(()) + } + + async fn list_directory( + &self, + user_id: &str, + agent_id: Option, + directory: &str, + ) -> Result, WorkspaceError> { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let dir = if !directory.is_empty() && !directory.ends_with('/') { + format!("{}/", directory) + } else { + directory.to_string() + }; + + let agent_id_str = agent_id.map(|id| id.to_string()); + let pattern = if dir.is_empty() { + "%".to_string() + } else { + format!("{}%", dir) + }; + + let mut rows = conn + .query( + r#" + SELECT path, updated_at, substr(content, 1, 200) as content_preview + FROM memory_documents + WHERE user_id = ?1 AND agent_id IS ?2 + AND (?3 = '%' OR path LIKE ?3) + ORDER BY path + "#, + params![user_id, agent_id_str.as_deref(), pattern], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("List directory failed: {}", e), + })?; + + let mut entries_map: HashMap = HashMap::new(); + + while let Some(row) = rows + .next() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })? + { + let full_path = get_text(&row, 0); + let updated_at = get_opt_ts(&row, 1); + let content_preview = get_opt_text(&row, 2); + + let relative = if dir.is_empty() { + &full_path + } else if let Some(stripped) = full_path.strip_prefix(&dir) { + stripped + } else { + continue; + }; + + let child_name = if let Some(slash_pos) = relative.find('/') { + &relative[..slash_pos] + } else { + relative + }; + + if child_name.is_empty() { + continue; + } + + let is_dir = relative.contains('/'); + let entry_path = if dir.is_empty() { + child_name.to_string() + } else { + format!("{}{}", dir, child_name) + }; + + entries_map + .entry(child_name.to_string()) + .and_modify(|e| { + if is_dir { + e.is_directory = true; + e.content_preview = None; + } + if let (Some(existing), Some(new)) = (&e.updated_at, &updated_at) + && new > existing + { + e.updated_at = Some(*new); + } + }) + .or_insert(WorkspaceEntry { + path: entry_path, + is_directory: is_dir, + updated_at, + content_preview: if is_dir { None } else { content_preview }, + }); + } + + let mut entries: Vec = entries_map.into_values().collect(); + entries.sort_by(|a, b| a.path.cmp(&b.path)); + Ok(entries) + } + + async fn list_all_paths( + &self, + user_id: &str, + agent_id: Option, + ) -> Result, WorkspaceError> { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let agent_id_str = agent_id.map(|id| id.to_string()); + let mut rows = conn + .query( + "SELECT path FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 ORDER BY path", + params![user_id, agent_id_str.as_deref()], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("List paths failed: {}", e), + })?; + + let mut paths = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })? + { + paths.push(get_text(&row, 0)); + } + Ok(paths) + } + + async fn list_documents( + &self, + user_id: &str, + agent_id: Option, + ) -> Result, WorkspaceError> { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let agent_id_str = agent_id.map(|id| id.to_string()); + let mut rows = conn + .query( + r#" + SELECT id, user_id, agent_id, path, content, + created_at, updated_at, metadata + FROM memory_documents + WHERE user_id = ?1 AND agent_id IS ?2 + ORDER BY updated_at DESC + "#, + params![user_id, agent_id_str.as_deref()], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })?; + + let mut docs = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })? + { + docs.push(row_to_memory_document(&row)); + } + Ok(docs) + } + + async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::ChunkingFailed { + reason: e.to_string(), + })?; + conn.execute( + "DELETE FROM memory_chunks WHERE document_id = ?1", + params![document_id.to_string()], + ) + .await + .map_err(|e| WorkspaceError::ChunkingFailed { + reason: format!("Delete failed: {}", e), + })?; + Ok(()) + } + + async fn insert_chunk( + &self, + document_id: Uuid, + chunk_index: i32, + content: &str, + embedding: Option<&[f32]>, + ) -> Result { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::ChunkingFailed { + reason: e.to_string(), + })?; + let id = Uuid::new_v4(); + let embedding_blob = embedding.map(|e| { + let bytes: Vec = e.iter().flat_map(|f| f.to_le_bytes()).collect(); + bytes + }); + + conn.execute( + r#" + INSERT INTO memory_chunks (id, document_id, chunk_index, content, embedding) + VALUES (?1, ?2, ?3, ?4, ?5) + "#, + params![ + id.to_string(), + document_id.to_string(), + chunk_index as i64, + content, + embedding_blob.map(libsql::Value::Blob), + ], + ) + .await + .map_err(|e| WorkspaceError::ChunkingFailed { + reason: format!("Insert failed: {}", e), + })?; + Ok(id) + } + + async fn update_chunk_embedding( + &self, + chunk_id: Uuid, + embedding: &[f32], + ) -> Result<(), WorkspaceError> { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::EmbeddingFailed { + reason: e.to_string(), + })?; + let bytes: Vec = embedding.iter().flat_map(|f| f.to_le_bytes()).collect(); + + conn.execute( + "UPDATE memory_chunks SET embedding = ?2 WHERE id = ?1", + params![chunk_id.to_string(), libsql::Value::Blob(bytes)], + ) + .await + .map_err(|e| WorkspaceError::EmbeddingFailed { + reason: format!("Update failed: {}", e), + })?; + Ok(()) + } + + async fn get_chunks_without_embeddings( + &self, + user_id: &str, + agent_id: Option, + limit: usize, + ) -> Result, WorkspaceError> { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let agent_id_str = agent_id.map(|id| id.to_string()); + let mut rows = conn + .query( + r#" + SELECT c.id, c.document_id, c.chunk_index, c.content, c.created_at + FROM memory_chunks c + JOIN memory_documents d ON d.id = c.document_id + WHERE d.user_id = ?1 AND d.agent_id IS ?2 + AND c.embedding IS NULL + LIMIT ?3 + "#, + params![user_id, agent_id_str.as_deref(), limit as i64], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })?; + + let mut chunks = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Query failed: {}", e), + })? + { + chunks.push(MemoryChunk { + id: get_text(&row, 0).parse().unwrap_or_default(), + document_id: get_text(&row, 1).parse().unwrap_or_default(), + chunk_index: get_i64(&row, 2) as i32, + content: get_text(&row, 3), + embedding: None, + created_at: get_ts(&row, 4), + }); + } + Ok(chunks) + } + + async fn hybrid_search( + &self, + user_id: &str, + agent_id: Option, + query: &str, + embedding: Option<&[f32]>, + config: &SearchConfig, + ) -> Result, WorkspaceError> { + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; + let agent_id_str = agent_id.map(|id| id.to_string()); + let pre_limit = config.pre_fusion_limit as i64; + + let fts_results = if config.use_fts { + let mut rows = conn + .query( + r#" + SELECT c.id, c.document_id, c.content + FROM memory_chunks_fts fts + JOIN memory_chunks c ON c._rowid = fts.rowid + JOIN memory_documents d ON d.id = c.document_id + WHERE d.user_id = ?1 AND d.agent_id IS ?2 + AND memory_chunks_fts MATCH ?3 + ORDER BY rank + LIMIT ?4 + "#, + params![user_id, agent_id_str.as_deref(), query, pre_limit], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("FTS query failed: {}", e), + })?; + + let mut results = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("FTS row fetch failed: {}", e), + })? + { + results.push(RankedResult { + chunk_id: get_text(&row, 0).parse().unwrap_or_default(), + document_id: get_text(&row, 1).parse().unwrap_or_default(), + content: get_text(&row, 2), + rank: results.len() as u32 + 1, + }); + } + results + } else { + Vec::new() + }; + + let vector_results = if let (true, Some(emb)) = (config.use_vector, embedding) { + let vector_json = format!( + "[{}]", + emb.iter() + .map(|f| f.to_string()) + .collect::>() + .join(",") + ); + + let mut rows = conn + .query( + r#" + SELECT c.id, c.document_id, c.content + FROM vector_top_k('idx_memory_chunks_embedding', vector(?1), ?2) AS top_k + JOIN memory_chunks c ON c._rowid = top_k.id + JOIN memory_documents d ON d.id = c.document_id + WHERE d.user_id = ?3 AND d.agent_id IS ?4 + "#, + params![vector_json, pre_limit, user_id, agent_id_str.as_deref()], + ) + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Vector query failed: {}", e), + })?; + + let mut results = Vec::new(); + while let Some(row) = rows + .next() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: format!("Vector row fetch failed: {}", e), + })? + { + results.push(RankedResult { + chunk_id: get_text(&row, 0).parse().unwrap_or_default(), + document_id: get_text(&row, 1).parse().unwrap_or_default(), + content: get_text(&row, 2), + rank: results.len() as u32 + 1, + }); + } + results + } else { + Vec::new() + }; + + if embedding.is_some() && !config.use_vector { + tracing::warn!( + "Embedding provided but vector search is disabled in config; using FTS-only results" + ); + } + + Ok(reciprocal_rank_fusion(fts_results, vector_results, config)) + } +} diff --git a/src/db/libsql_backend.rs b/src/db/libsql_backend.rs deleted file mode 100644 index 31f79e77..00000000 --- a/src/db/libsql_backend.rs +++ /dev/null @@ -1,2769 +0,0 @@ -//! libSQL/Turso backend for the Database trait. -//! -//! Provides an embedded SQLite-compatible database using Turso's libSQL fork. -//! Supports three modes: -//! - Local embedded (file-based, no server needed) -//! - Turso cloud with embedded replica (sync to cloud) -//! - In-memory (for testing) - -use std::collections::HashMap; -use std::path::Path; -use std::sync::Arc; - -use async_trait::async_trait; -use chrono::{DateTime, NaiveDateTime, Utc}; -use libsql::{Connection, Database as LibSqlDatabase, params}; -use rust_decimal::Decimal; -use uuid::Uuid; - -use crate::agent::BrokenTool; -use crate::agent::routine::{ - NotifyConfig, Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger, -}; -use crate::context::{ActionRecord, JobContext, JobState}; -use crate::db::Database; -use crate::error::{DatabaseError, WorkspaceError}; -use crate::history::{ - ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord, - SandboxJobSummary, SettingRow, -}; -use crate::workspace::{ - MemoryChunk, MemoryDocument, RankedResult, SearchConfig, SearchResult, WorkspaceEntry, - reciprocal_rank_fusion, -}; - -use crate::db::libsql_migrations; - -/// Explicit column list for routines table (matches positional access in `row_to_routine_libsql`). -const ROUTINE_COLUMNS: &str = "\ - id, name, description, user_id, enabled, \ - trigger_type, trigger_config, action_type, action_config, \ - cooldown_secs, max_concurrent, dedup_window_secs, \ - notify_channel, notify_user, notify_on_success, notify_on_failure, notify_on_attention, \ - state, last_run_at, next_fire_at, run_count, consecutive_failures, \ - created_at, updated_at"; - -/// Explicit column list for routine_runs table (matches positional access in `row_to_routine_run_libsql`). -const ROUTINE_RUN_COLUMNS: &str = "\ - id, routine_id, trigger_type, trigger_detail, started_at, \ - status, completed_at, result_summary, tokens_used, job_id, created_at"; - -/// libSQL/Turso database backend. -/// -/// Stores the `Database` handle in an `Arc` so that the same underlying -/// database can be shared with stores (SecretsStore, WasmToolStore) that -/// create their own connections per-operation. -pub struct LibSqlBackend { - db: Arc, -} - -impl LibSqlBackend { - /// Create a new local embedded database. - pub async fn new_local(path: &Path) -> Result { - // Ensure parent directory exists - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent).map_err(|e| { - DatabaseError::Pool(format!("Failed to create database directory: {}", e)) - })?; - } - - let db = libsql::Builder::new_local(path) - .build() - .await - .map_err(|e| DatabaseError::Pool(format!("Failed to open libSQL database: {}", e)))?; - - Ok(Self { db: Arc::new(db) }) - } - - /// Create a new in-memory database (for testing). - pub async fn new_memory() -> Result { - let db = libsql::Builder::new_local(":memory:") - .build() - .await - .map_err(|e| { - DatabaseError::Pool(format!("Failed to create in-memory database: {}", e)) - })?; - - Ok(Self { db: Arc::new(db) }) - } - - /// Create with Turso cloud sync (embedded replica). - pub async fn new_remote_replica( - path: &Path, - url: &str, - auth_token: &str, - ) -> Result { - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent).map_err(|e| { - DatabaseError::Pool(format!("Failed to create database directory: {}", e)) - })?; - } - - let db = libsql::Builder::new_remote_replica(path, url.to_string(), auth_token.to_string()) - .build() - .await - .map_err(|e| DatabaseError::Pool(format!("Failed to open remote replica: {}", e)))?; - - Ok(Self { db: Arc::new(db) }) - } - - /// Get a shared reference to the underlying database handle. - /// - /// Use this to pass the database to stores (SecretsStore, WasmToolStore) - /// that need to create their own connections per-operation. - pub fn shared_db(&self) -> Arc { - Arc::clone(&self.db) - } - - /// Create a new connection to the database. - /// - /// Sets `PRAGMA busy_timeout = 5000` on every connection so concurrent - /// writers wait up to 5 seconds instead of failing instantly with - /// "database is locked". - pub async fn connect(&self) -> Result { - let conn = self - .db - .connect() - .map_err(|e| DatabaseError::Pool(format!("Failed to create connection: {}", e)))?; - conn.query("PRAGMA busy_timeout = 5000", ()) - .await - .map_err(|e| DatabaseError::Pool(format!("Failed to set busy_timeout: {}", e)))?; - Ok(conn) - } -} - -// ==================== Helper functions ==================== - -/// Parse an ISO-8601 timestamp string from SQLite into DateTime. -/// -/// Tries multiple formats in order: -/// 1. RFC 3339 with timezone (e.g. `2024-01-15T10:30:00.123Z`) -/// 2. Naive datetime with fractional seconds (e.g. `2024-01-15 10:30:00.123`) -/// 3. Naive datetime without fractional seconds (e.g. `2024-01-15 10:30:00`) -/// -/// Returns an error if none of the formats match. -fn parse_timestamp(s: &str) -> Result, String> { - // RFC 3339 (our canonical write format) - if let Ok(dt) = DateTime::parse_from_rfc3339(s) { - return Ok(dt.with_timezone(&Utc)); - } - // Naive with fractional seconds (legacy or SQLite datetime() output) - if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") { - return Ok(ndt.and_utc()); - } - // Naive without fractional seconds (legacy format) - if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") { - return Ok(ndt.and_utc()); - } - Err(format!("unparseable timestamp: {:?}", s)) -} - -/// Format a DateTime for SQLite storage (RFC 3339 with millisecond precision). -fn fmt_ts(dt: &DateTime) -> String { - dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true) -} - -/// Format an optional DateTime. -fn fmt_opt_ts(dt: &Option>) -> libsql::Value { - match dt { - Some(dt) => libsql::Value::Text(fmt_ts(dt)), - None => libsql::Value::Null, - } -} - -fn parse_job_state(s: &str) -> JobState { - match s { - "pending" => JobState::Pending, - "in_progress" => JobState::InProgress, - "completed" => JobState::Completed, - "submitted" => JobState::Submitted, - "accepted" => JobState::Accepted, - "failed" => JobState::Failed, - "stuck" => JobState::Stuck, - "cancelled" => JobState::Cancelled, - _ => JobState::Pending, - } -} - -/// Extract a text column from a libsql Row, returning empty string for NULL. -fn get_text(row: &libsql::Row, idx: i32) -> String { - row.get::(idx).unwrap_or_default() -} - -/// Extract an optional text column. -/// Returns None for SQL NULL, preserves empty strings as Some(""). -fn get_opt_text(row: &libsql::Row, idx: i32) -> Option { - row.get::(idx).ok() -} - -/// Convert an `Option<&str>` to a `libsql::Value` (Text or Null). -/// Use this instead of `.unwrap_or("")` to preserve NULL semantics. -fn opt_text(s: Option<&str>) -> libsql::Value { - match s { - Some(s) => libsql::Value::Text(s.to_string()), - None => libsql::Value::Null, - } -} - -/// Convert an `Option` to a `libsql::Value` (Text or Null). -fn opt_text_owned(s: Option) -> libsql::Value { - match s { - Some(s) => libsql::Value::Text(s), - None => libsql::Value::Null, - } -} - -/// Extract an i64 column, defaulting to 0. -fn get_i64(row: &libsql::Row, idx: i32) -> i64 { - row.get::(idx).unwrap_or(0) -} - -/// Extract an optional bool from an integer column. -fn get_opt_bool(row: &libsql::Row, idx: i32) -> Option { - row.get::(idx).ok().map(|v| v != 0) -} - -/// Parse a Decimal from a text column. -fn get_decimal(row: &libsql::Row, idx: i32) -> Decimal { - row.get::(idx) - .ok() - .and_then(|s| s.parse::().ok()) - .unwrap_or_default() -} - -/// Parse an optional Decimal from a text column. -fn get_opt_decimal(row: &libsql::Row, idx: i32) -> Option { - row.get::(idx) - .ok() - .and_then(|s| s.parse::().ok()) -} - -/// Parse a JSON value from a text column. -fn get_json(row: &libsql::Row, idx: i32) -> serde_json::Value { - row.get::(idx) - .ok() - .and_then(|s| serde_json::from_str(&s).ok()) - .unwrap_or(serde_json::Value::Null) -} - -/// Parse a timestamp from a text column. -/// -/// If the column is NULL or the value cannot be parsed, logs a warning and -/// returns the Unix epoch (1970-01-01T00:00:00Z) so the error is detectable -/// rather than silently replaced by the current time. -fn get_ts(row: &libsql::Row, idx: i32) -> DateTime { - match row.get::(idx) { - Ok(s) => match parse_timestamp(&s) { - Ok(dt) => dt, - Err(e) => { - tracing::warn!("Timestamp parse failure at column {}: {}", idx, e); - DateTime::UNIX_EPOCH - } - }, - Err(_) => DateTime::UNIX_EPOCH, - } -} - -/// Parse an optional timestamp from a text column. -/// -/// Returns None if the column is NULL. Logs a warning and returns None if the -/// value is present but cannot be parsed. -fn get_opt_ts(row: &libsql::Row, idx: i32) -> Option> { - match row.get::(idx) { - Ok(s) if s.is_empty() => None, - Ok(s) => match parse_timestamp(&s) { - Ok(dt) => Some(dt), - Err(e) => { - tracing::warn!("Timestamp parse failure at column {}: {}", idx, e); - None - } - }, - Err(_) => None, - } -} - -#[async_trait] -impl Database for LibSqlBackend { - async fn run_migrations(&self) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - // WAL mode persists in the database file: all future connections benefit. - // Readers no longer block writers and vice versa. - conn.query("PRAGMA journal_mode=WAL", ()) - .await - .map_err(|e| DatabaseError::Migration(format!("Failed to enable WAL mode: {}", e)))?; - conn.execute_batch(libsql_migrations::SCHEMA) - .await - .map_err(|e| DatabaseError::Migration(format!("libSQL migration failed: {}", e)))?; - Ok(()) - } - - // ==================== Conversations ==================== - - async fn create_conversation( - &self, - channel: &str, - user_id: &str, - thread_id: Option<&str>, - ) -> Result { - let conn = self.connect().await?; - let id = Uuid::new_v4(); - conn.execute( - "INSERT INTO conversations (id, channel, user_id, thread_id) VALUES (?1, ?2, ?3, ?4)", - params![id.to_string(), channel, user_id, opt_text(thread_id)], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(id) - } - - async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute( - "UPDATE conversations SET last_activity = ?2 WHERE id = ?1", - params![id.to_string(), now], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn add_conversation_message( - &self, - conversation_id: Uuid, - role: &str, - content: &str, - ) -> Result { - let conn = self.connect().await?; - let id = Uuid::new_v4(); - conn.execute( - "INSERT INTO conversation_messages (id, conversation_id, role, content) VALUES (?1, ?2, ?3, ?4)", - params![id.to_string(), conversation_id.to_string(), role, content], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - self.touch_conversation(conversation_id).await?; - Ok(id) - } - - async fn ensure_conversation( - &self, - id: Uuid, - channel: &str, - user_id: &str, - thread_id: Option<&str>, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute( - r#" - INSERT INTO conversations (id, channel, user_id, thread_id) - VALUES (?1, ?2, ?3, ?4) - ON CONFLICT (id) DO UPDATE SET last_activity = ?5 - "#, - params![id.to_string(), channel, user_id, opt_text(thread_id), now], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn list_conversations_with_preview( - &self, - user_id: &str, - channel: &str, - limit: i64, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - r#" - SELECT - c.id, - c.started_at, - c.last_activity, - c.metadata, - (SELECT COUNT(*) FROM conversation_messages m WHERE m.conversation_id = c.id) AS message_count, - (SELECT substr(m2.content, 1, 100) - FROM conversation_messages m2 - WHERE m2.conversation_id = c.id AND m2.role = 'user' - ORDER BY m2.created_at ASC - LIMIT 1 - ) AS title - FROM conversations c - WHERE c.user_id = ?1 AND c.channel = ?2 - ORDER BY c.last_activity DESC - LIMIT ?3 - "#, - params![user_id, channel, limit], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut results = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - let metadata = get_json(&row, 3); - let thread_type = metadata - .get("thread_type") - .and_then(|v| v.as_str()) - .map(String::from); - results.push(ConversationSummary { - id: row - .get::(0) - .unwrap_or_default() - .parse() - .unwrap_or_default(), - started_at: get_ts(&row, 1), - last_activity: get_ts(&row, 2), - message_count: get_i64(&row, 4), - title: get_opt_text(&row, 5), - thread_type, - }); - } - Ok(results) - } - - async fn get_or_create_assistant_conversation( - &self, - user_id: &str, - channel: &str, - ) -> Result { - let conn = self.connect().await?; - // Try to find existing - let mut rows = conn - .query( - r#" - SELECT id FROM conversations - WHERE user_id = ?1 AND channel = ?2 - AND json_extract(metadata, '$.thread_type') = 'assistant' - LIMIT 1 - "#, - params![user_id, channel], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - if let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - let id_str: String = row.get(0).unwrap_or_default(); - return id_str - .parse() - .map_err(|_| DatabaseError::Serialization("Invalid UUID".to_string())); - } - - // Create new - let id = Uuid::new_v4(); - let metadata = serde_json::json!({"thread_type": "assistant", "title": "Assistant"}); - conn.execute( - "INSERT INTO conversations (id, channel, user_id, metadata) VALUES (?1, ?2, ?3, ?4)", - params![id.to_string(), channel, user_id, metadata.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(id) - } - - async fn create_conversation_with_metadata( - &self, - channel: &str, - user_id: &str, - metadata: &serde_json::Value, - ) -> Result { - let conn = self.connect().await?; - let id = Uuid::new_v4(); - conn.execute( - "INSERT INTO conversations (id, channel, user_id, metadata) VALUES (?1, ?2, ?3, ?4)", - params![id.to_string(), channel, user_id, metadata.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(id) - } - - async fn list_conversation_messages_paginated( - &self, - conversation_id: Uuid, - before: Option>, - limit: i64, - ) -> Result<(Vec, bool), DatabaseError> { - let conn = self.connect().await?; - let fetch_limit = limit + 1; - let cid = conversation_id.to_string(); - - let mut rows = if let Some(before_ts) = before { - conn.query( - r#" - SELECT id, role, content, created_at - FROM conversation_messages - WHERE conversation_id = ?1 AND created_at < ?2 - ORDER BY created_at DESC - LIMIT ?3 - "#, - params![cid, fmt_ts(&before_ts), fetch_limit], - ) - .await - } else { - conn.query( - r#" - SELECT id, role, content, created_at - FROM conversation_messages - WHERE conversation_id = ?1 - ORDER BY created_at DESC - LIMIT ?2 - "#, - params![cid, fetch_limit], - ) - .await - } - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut all = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - all.push(ConversationMessage { - id: get_text(&row, 0).parse().unwrap_or_default(), - role: get_text(&row, 1), - content: get_text(&row, 2), - created_at: get_ts(&row, 3), - }); - } - - let has_more = all.len() as i64 > limit; - all.truncate(limit as usize); - all.reverse(); // oldest first - Ok((all, has_more)) - } - - async fn update_conversation_metadata_field( - &self, - id: Uuid, - key: &str, - value: &serde_json::Value, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - // SQLite: use json_patch to merge the key - let patch = serde_json::json!({ key: value }); - conn.execute( - "UPDATE conversations SET metadata = json_patch(metadata, ?2) WHERE id = ?1", - params![id.to_string(), patch.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn get_conversation_metadata( - &self, - id: Uuid, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT metadata FROM conversations WHERE id = ?1", - params![id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(Some(get_json(&row, 0))), - None => Ok(None), - } - } - - async fn list_conversation_messages( - &self, - conversation_id: Uuid, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - r#" - SELECT id, role, content, created_at - FROM conversation_messages - WHERE conversation_id = ?1 - ORDER BY created_at ASC - "#, - params![conversation_id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut messages = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - messages.push(ConversationMessage { - id: get_text(&row, 0).parse().unwrap_or_default(), - role: get_text(&row, 1), - content: get_text(&row, 2), - created_at: get_ts(&row, 3), - }); - } - Ok(messages) - } - - async fn conversation_belongs_to_user( - &self, - conversation_id: Uuid, - user_id: &str, - ) -> Result { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT 1 FROM conversations WHERE id = ?1 AND user_id = ?2", - libsql::params![conversation_id.to_string(), user_id], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - let found = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(found.is_some()) - } - - // ==================== Jobs ==================== - - async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let status = ctx.state.to_string(); - let estimated_time_secs = ctx.estimated_duration.map(|d| d.as_secs() as i64); - - conn - .execute( - r#" - INSERT INTO agent_jobs ( - id, conversation_id, title, description, category, status, source, - budget_amount, budget_token, bid_amount, estimated_cost, estimated_time_secs, - actual_cost, repair_attempts, created_at, started_at, completed_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17) - ON CONFLICT (id) DO UPDATE SET - title = excluded.title, - description = excluded.description, - category = excluded.category, - status = excluded.status, - estimated_cost = excluded.estimated_cost, - estimated_time_secs = excluded.estimated_time_secs, - actual_cost = excluded.actual_cost, - repair_attempts = excluded.repair_attempts, - started_at = excluded.started_at, - completed_at = excluded.completed_at - "#, - params![ - ctx.job_id.to_string(), - opt_text_owned(ctx.conversation_id.map(|id| id.to_string())), - ctx.title.as_str(), - ctx.description.as_str(), - opt_text(ctx.category.as_deref()), - status, - "direct", - opt_text_owned(ctx.budget.map(|d| d.to_string())), - opt_text(ctx.budget_token.as_deref()), - opt_text_owned(ctx.bid_amount.map(|d| d.to_string())), - opt_text_owned(ctx.estimated_cost.map(|d| d.to_string())), - estimated_time_secs, - ctx.actual_cost.to_string(), - ctx.repair_attempts as i64, - fmt_ts(&ctx.created_at), - fmt_opt_ts(&ctx.started_at), - fmt_opt_ts(&ctx.completed_at), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn get_job(&self, id: Uuid) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - r#" - SELECT id, conversation_id, title, description, category, status, user_id, - budget_amount, budget_token, bid_amount, estimated_cost, estimated_time_secs, - actual_cost, repair_attempts, created_at, started_at, completed_at - FROM agent_jobs WHERE id = ?1 - "#, - params![id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => { - let status_str = get_text(&row, 5); - let state = parse_job_state(&status_str); - let estimated_time_secs: Option = row.get::(11).ok(); - - Ok(Some(JobContext { - job_id: get_text(&row, 0).parse().unwrap_or_default(), - state, - user_id: get_text(&row, 6), - conversation_id: get_opt_text(&row, 1).and_then(|s| s.parse().ok()), - title: get_text(&row, 2), - description: get_text(&row, 3), - category: get_opt_text(&row, 4), - budget: get_opt_decimal(&row, 7), - budget_token: get_opt_text(&row, 8), - bid_amount: get_opt_decimal(&row, 9), - estimated_cost: get_opt_decimal(&row, 10), - estimated_duration: estimated_time_secs - .map(|s| std::time::Duration::from_secs(s as u64)), - actual_cost: get_decimal(&row, 12), - total_tokens_used: 0, - max_tokens: 0, - repair_attempts: get_i64(&row, 13) as u32, - created_at: get_ts(&row, 14), - started_at: get_opt_ts(&row, 15), - completed_at: get_opt_ts(&row, 16), - transitions: Vec::new(), - metadata: serde_json::Value::Null, - extra_env: std::sync::Arc::new(std::collections::HashMap::new()), - })) - } - None => Ok(None), - } - } - - async fn update_job_status( - &self, - id: Uuid, - status: JobState, - failure_reason: Option<&str>, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - conn.execute( - "UPDATE agent_jobs SET status = ?2, failure_reason = ?3 WHERE id = ?1", - params![id.to_string(), status.to_string(), opt_text(failure_reason)], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute( - "UPDATE agent_jobs SET status = 'stuck', stuck_since = ?2 WHERE id = ?1", - params![id.to_string(), now], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn get_stuck_jobs(&self) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query("SELECT id FROM agent_jobs WHERE status = 'stuck'", ()) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut ids = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - if let Ok(id_str) = row.get::(0) - && let Ok(id) = id_str.parse() - { - ids.push(id); - } - } - Ok(ids) - } - - // ==================== Actions ==================== - - async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let duration_ms = action.duration.as_millis() as i64; - let warnings_json = serde_json::to_string(&action.sanitization_warnings) - .map_err(|e| DatabaseError::Serialization(e.to_string()))?; - - conn.execute( - r#" - INSERT INTO job_actions ( - id, job_id, sequence_num, tool_name, input, output_raw, output_sanitized, - sanitization_warnings, cost, duration_ms, success, error_message, created_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13) - "#, - params![ - action.id.to_string(), - job_id.to_string(), - action.sequence as i64, - action.tool_name.as_str(), - action.input.to_string(), - opt_text(action.output_raw.as_deref()), - opt_text_owned(action.output_sanitized.as_ref().map(|v| v.to_string())), - warnings_json, - opt_text_owned(action.cost.map(|d| d.to_string())), - duration_ms, - action.success as i64, - opt_text(action.error.as_deref()), - fmt_ts(&action.executed_at), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn get_job_actions(&self, job_id: Uuid) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - r#" - SELECT id, sequence_num, tool_name, input, output_raw, output_sanitized, - sanitization_warnings, cost, duration_ms, success, error_message, created_at - FROM job_actions WHERE job_id = ?1 ORDER BY sequence_num - "#, - params![job_id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut actions = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - let warnings: Vec = - serde_json::from_str(&get_text(&row, 6)).unwrap_or_default(); - actions.push(ActionRecord { - id: get_text(&row, 0).parse().unwrap_or_default(), - sequence: get_i64(&row, 1) as u32, - tool_name: get_text(&row, 2), - input: get_json(&row, 3), - output_raw: get_opt_text(&row, 4), - output_sanitized: get_opt_text(&row, 5).and_then(|s| serde_json::from_str(&s).ok()), - sanitization_warnings: warnings, - cost: get_opt_decimal(&row, 7), - duration: std::time::Duration::from_millis(get_i64(&row, 8) as u64), - success: get_i64(&row, 9) != 0, - error: get_opt_text(&row, 10), - executed_at: get_ts(&row, 11), - }); - } - Ok(actions) - } - - // ==================== LLM Calls ==================== - - async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result { - let conn = self.connect().await?; - let id = Uuid::new_v4(); - conn.execute( - r#" - INSERT INTO llm_calls (id, job_id, conversation_id, provider, model, input_tokens, output_tokens, cost, purpose) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) - "#, - params![ - id.to_string(), - opt_text_owned(record.job_id.map(|id| id.to_string())), - opt_text_owned(record.conversation_id.map(|id| id.to_string())), - record.provider, - record.model, - record.input_tokens as i64, - record.output_tokens as i64, - record.cost.to_string(), - opt_text(record.purpose), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(id) - } - - // ==================== Estimation Snapshots ==================== - - async fn save_estimation_snapshot( - &self, - job_id: Uuid, - category: &str, - tool_names: &[String], - estimated_cost: Decimal, - estimated_time_secs: i32, - estimated_value: Decimal, - ) -> Result { - let conn = self.connect().await?; - let id = Uuid::new_v4(); - let tools_json = serde_json::to_string(tool_names) - .map_err(|e| DatabaseError::Serialization(e.to_string()))?; - - conn.execute( - r#" - INSERT INTO estimation_snapshots (id, job_id, category, tool_names, estimated_cost, estimated_time_secs, estimated_value) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) - "#, - params![ - id.to_string(), - job_id.to_string(), - category, - tools_json, - estimated_cost.to_string(), - estimated_time_secs as i64, - estimated_value.to_string(), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(id) - } - - async fn update_estimation_actuals( - &self, - id: Uuid, - actual_cost: Decimal, - actual_time_secs: i32, - actual_value: Option, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - conn.execute( - "UPDATE estimation_snapshots SET actual_cost = ?2, actual_time_secs = ?3, actual_value = ?4 WHERE id = ?1", - params![ - id.to_string(), - actual_cost.to_string(), - actual_time_secs as i64, - actual_value.map(|d| d.to_string()).unwrap_or_default(), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - // ==================== Sandbox Jobs ==================== - - async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - conn.execute( - r#" - INSERT INTO agent_jobs ( - id, title, description, status, source, user_id, project_dir, - success, failure_reason, created_at, started_at, completed_at - ) VALUES (?1, ?2, ?3, ?4, 'sandbox', ?5, ?6, ?7, ?8, ?9, ?10, ?11) - ON CONFLICT (id) DO UPDATE SET - status = excluded.status, - success = excluded.success, - failure_reason = excluded.failure_reason, - started_at = excluded.started_at, - completed_at = excluded.completed_at - "#, - params![ - job.id.to_string(), - job.task.as_str(), - job.credential_grants_json.as_str(), - job.status.as_str(), - job.user_id.as_str(), - job.project_dir.as_str(), - job.success.map(|b| b as i64), - opt_text(job.failure_reason.as_deref()), - fmt_ts(&job.created_at), - fmt_opt_ts(&job.started_at), - fmt_opt_ts(&job.completed_at), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn get_sandbox_job(&self, id: Uuid) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - r#" - SELECT id, title, description, status, user_id, project_dir, - success, failure_reason, created_at, started_at, completed_at - FROM agent_jobs WHERE id = ?1 AND source = 'sandbox' - "#, - params![id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(Some(SandboxJobRecord { - id: get_text(&row, 0).parse().unwrap_or_default(), - task: get_text(&row, 1), - credential_grants_json: get_text(&row, 2), - status: get_text(&row, 3), - user_id: get_text(&row, 4), - project_dir: get_text(&row, 5), - success: get_opt_bool(&row, 6), - failure_reason: get_opt_text(&row, 7), - created_at: get_ts(&row, 8), - started_at: get_opt_ts(&row, 9), - completed_at: get_opt_ts(&row, 10), - })), - None => Ok(None), - } - } - - async fn list_sandbox_jobs(&self) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - r#" - SELECT id, title, description, status, user_id, project_dir, - success, failure_reason, created_at, started_at, completed_at - FROM agent_jobs WHERE source = 'sandbox' - ORDER BY created_at DESC - "#, - (), - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut jobs = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - jobs.push(SandboxJobRecord { - id: get_text(&row, 0).parse().unwrap_or_default(), - task: get_text(&row, 1), - credential_grants_json: get_text(&row, 2), - status: get_text(&row, 3), - user_id: get_text(&row, 4), - project_dir: get_text(&row, 5), - success: get_opt_bool(&row, 6), - failure_reason: get_opt_text(&row, 7), - created_at: get_ts(&row, 8), - started_at: get_opt_ts(&row, 9), - completed_at: get_opt_ts(&row, 10), - }); - } - Ok(jobs) - } - - async fn update_sandbox_job_status( - &self, - id: Uuid, - status: &str, - success: Option, - message: Option<&str>, - started_at: Option>, - completed_at: Option>, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - conn.execute( - r#" - UPDATE agent_jobs SET - status = ?2, - success = COALESCE(?3, success), - failure_reason = COALESCE(?4, failure_reason), - started_at = COALESCE(?5, started_at), - completed_at = COALESCE(?6, completed_at) - WHERE id = ?1 AND source = 'sandbox' - "#, - params![ - id.to_string(), - status, - success.map(|b| b as i64), - message, - fmt_opt_ts(&started_at), - fmt_opt_ts(&completed_at), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn cleanup_stale_sandbox_jobs(&self) -> Result { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - let count = conn - .execute( - r#" - UPDATE agent_jobs SET - status = 'interrupted', - failure_reason = 'Process restarted', - completed_at = ?1 - WHERE source = 'sandbox' AND status IN ('running', 'creating') - "#, - params![now], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - if count > 0 { - tracing::info!("Marked {} stale sandbox jobs as interrupted", count); - } - Ok(count) - } - - async fn sandbox_job_summary(&self) -> Result { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'sandbox' GROUP BY status", - (), - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut summary = SandboxJobSummary::default(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - let status = get_text(&row, 0); - let count = get_i64(&row, 1) as usize; - summary.total += count; - match status.as_str() { - "creating" => summary.creating += count, - "running" => summary.running += count, - "completed" => summary.completed += count, - "failed" => summary.failed += count, - "interrupted" => summary.interrupted += count, - _ => {} - } - } - Ok(summary) - } - - async fn list_sandbox_jobs_for_user( - &self, - user_id: &str, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - r#" - SELECT id, title, description, status, user_id, project_dir, - success, failure_reason, created_at, started_at, completed_at - FROM agent_jobs WHERE source = 'sandbox' AND user_id = ?1 - ORDER BY created_at DESC - "#, - libsql::params![user_id], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut jobs = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - jobs.push(SandboxJobRecord { - id: get_text(&row, 0).parse().unwrap_or_default(), - task: get_text(&row, 1), - credential_grants_json: get_text(&row, 2), - status: get_text(&row, 3), - user_id: get_text(&row, 4), - project_dir: get_text(&row, 5), - success: get_opt_bool(&row, 6), - failure_reason: get_opt_text(&row, 7), - created_at: get_ts(&row, 8), - started_at: get_opt_ts(&row, 9), - completed_at: get_opt_ts(&row, 10), - }); - } - Ok(jobs) - } - - async fn sandbox_job_summary_for_user( - &self, - user_id: &str, - ) -> Result { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'sandbox' AND user_id = ?1 GROUP BY status", - libsql::params![user_id], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut summary = SandboxJobSummary::default(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - let status = get_text(&row, 0); - let count = get_i64(&row, 1) as usize; - summary.total += count; - match status.as_str() { - "creating" => summary.creating += count, - "running" => summary.running += count, - "completed" => summary.completed += count, - "failed" => summary.failed += count, - "interrupted" => summary.interrupted += count, - _ => {} - } - } - Ok(summary) - } - - async fn sandbox_job_belongs_to_user( - &self, - job_id: Uuid, - user_id: &str, - ) -> Result { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT 1 FROM agent_jobs WHERE id = ?1 AND user_id = ?2 AND source = 'sandbox'", - libsql::params![job_id.to_string(), user_id], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - let found = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(found.is_some()) - } - - async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - conn.execute( - "UPDATE agent_jobs SET job_mode = ?2 WHERE id = ?1", - params![id.to_string(), mode], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn get_sandbox_job_mode(&self, id: Uuid) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT job_mode FROM agent_jobs WHERE id = ?1", - params![id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(Some(get_text(&row, 0))), - None => Ok(None), - } - } - - // ==================== Job Events ==================== - - async fn save_job_event( - &self, - job_id: Uuid, - event_type: &str, - data: &serde_json::Value, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - conn.execute( - "INSERT INTO job_events (job_id, event_type, data) VALUES (?1, ?2, ?3)", - params![job_id.to_string(), event_type, data.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn list_job_events( - &self, - job_id: Uuid, - limit: Option, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = if let Some(n) = limit { - conn.query( - r#" - SELECT id, job_id, event_type, data, created_at - FROM ( - SELECT id, job_id, event_type, data, created_at - FROM job_events WHERE job_id = ?1 - ORDER BY id DESC - LIMIT ?2 - ) - ORDER BY id ASC - "#, - params![job_id.to_string(), n], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - } else { - conn.query( - r#" - SELECT id, job_id, event_type, data, created_at - FROM job_events WHERE job_id = ?1 ORDER BY id ASC - "#, - params![job_id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - }; - - let mut events = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - events.push(JobEventRecord { - id: get_i64(&row, 0), - job_id: get_text(&row, 1).parse().unwrap_or_default(), - event_type: get_text(&row, 2), - data: get_json(&row, 3), - created_at: get_ts(&row, 4), - }); - } - Ok(events) - } - - // ==================== Routines ==================== - - async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let trigger_type = routine.trigger.type_tag(); - let trigger_config = routine.trigger.to_config_json(); - let action_type = routine.action.type_tag(); - let action_config = routine.action.to_config_json(); - let cooldown_secs = routine.guardrails.cooldown.as_secs() as i64; - let max_concurrent = routine.guardrails.max_concurrent as i64; - let dedup_window_secs = routine.guardrails.dedup_window.map(|d| d.as_secs() as i64); - - conn.execute( - r#" - INSERT INTO routines ( - id, name, description, user_id, enabled, - trigger_type, trigger_config, action_type, action_config, - cooldown_secs, max_concurrent, dedup_window_secs, - notify_channel, notify_user, notify_on_success, notify_on_failure, notify_on_attention, - state, next_fire_at, created_at, updated_at - ) VALUES ( - ?1, ?2, ?3, ?4, ?5, - ?6, ?7, ?8, ?9, - ?10, ?11, ?12, - ?13, ?14, ?15, ?16, ?17, - ?18, ?19, ?20, ?21 - ) - "#, - params![ - routine.id.to_string(), - routine.name.as_str(), - routine.description.as_str(), - routine.user_id.as_str(), - routine.enabled as i64, - trigger_type, - trigger_config.to_string(), - action_type, - action_config.to_string(), - cooldown_secs, - max_concurrent, - dedup_window_secs, - opt_text(routine.notify.channel.as_deref()), - routine.notify.user.as_str(), - routine.notify.on_success as i64, - routine.notify.on_failure as i64, - routine.notify.on_attention as i64, - routine.state.to_string(), - fmt_opt_ts(&routine.next_fire_at), - fmt_ts(&routine.created_at), - fmt_ts(&routine.updated_at), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn get_routine(&self, id: Uuid) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - &format!("SELECT {} FROM routines WHERE id = ?1", ROUTINE_COLUMNS), - params![id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(Some(row_to_routine_libsql(&row)?)), - None => Ok(None), - } - } - - async fn get_routine_by_name( - &self, - user_id: &str, - name: &str, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - &format!( - "SELECT {} FROM routines WHERE user_id = ?1 AND name = ?2", - ROUTINE_COLUMNS - ), - params![user_id, name], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(Some(row_to_routine_libsql(&row)?)), - None => Ok(None), - } - } - - async fn list_routines(&self, user_id: &str) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - &format!( - "SELECT {} FROM routines WHERE user_id = ?1 ORDER BY name", - ROUTINE_COLUMNS - ), - params![user_id], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut routines = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - routines.push(row_to_routine_libsql(&row)?); - } - Ok(routines) - } - - async fn list_event_routines(&self) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - &format!( - "SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'event'", - ROUTINE_COLUMNS - ), - (), - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut routines = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - routines.push(row_to_routine_libsql(&row)?); - } - Ok(routines) - } - - async fn list_due_cron_routines(&self) -> Result, DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - let mut rows = conn - .query( - &format!( - "SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'cron' AND next_fire_at IS NOT NULL AND next_fire_at <= ?1", - ROUTINE_COLUMNS - ), - params![now], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut routines = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - routines.push(row_to_routine_libsql(&row)?); - } - Ok(routines) - } - - async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let trigger_type = routine.trigger.type_tag(); - let trigger_config = routine.trigger.to_config_json(); - let action_type = routine.action.type_tag(); - let action_config = routine.action.to_config_json(); - let cooldown_secs = routine.guardrails.cooldown.as_secs() as i64; - let max_concurrent = routine.guardrails.max_concurrent as i64; - let dedup_window_secs = routine.guardrails.dedup_window.map(|d| d.as_secs() as i64); - let now = fmt_ts(&Utc::now()); - - conn.execute( - r#" - UPDATE routines SET - name = ?2, description = ?3, enabled = ?4, - trigger_type = ?5, trigger_config = ?6, - action_type = ?7, action_config = ?8, - cooldown_secs = ?9, max_concurrent = ?10, dedup_window_secs = ?11, - notify_channel = ?12, notify_user = ?13, - notify_on_success = ?14, notify_on_failure = ?15, notify_on_attention = ?16, - state = ?17, next_fire_at = ?18, - updated_at = ?19 - WHERE id = ?1 - "#, - params![ - routine.id.to_string(), - routine.name.as_str(), - routine.description.as_str(), - routine.enabled as i64, - trigger_type, - trigger_config.to_string(), - action_type, - action_config.to_string(), - cooldown_secs, - max_concurrent, - dedup_window_secs, - opt_text(routine.notify.channel.as_deref()), - routine.notify.user.as_str(), - routine.notify.on_success as i64, - routine.notify.on_failure as i64, - routine.notify.on_attention as i64, - routine.state.to_string(), - fmt_opt_ts(&routine.next_fire_at), - now, - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn update_routine_runtime( - &self, - id: Uuid, - last_run_at: DateTime, - next_fire_at: Option>, - run_count: u64, - consecutive_failures: u32, - state: &serde_json::Value, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute( - r#" - UPDATE routines SET - last_run_at = ?2, next_fire_at = ?3, - run_count = ?4, consecutive_failures = ?5, - state = ?6, updated_at = ?7 - WHERE id = ?1 - "#, - params![ - id.to_string(), - fmt_ts(&last_run_at), - fmt_opt_ts(&next_fire_at), - run_count as i64, - consecutive_failures as i64, - state.to_string(), - now, - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn delete_routine(&self, id: Uuid) -> Result { - let conn = self.connect().await?; - let count = conn - .execute( - "DELETE FROM routines WHERE id = ?1", - params![id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(count > 0) - } - - // ==================== Routine Runs ==================== - - async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - conn.execute( - r#" - INSERT INTO routine_runs ( - id, routine_id, trigger_type, trigger_detail, - started_at, status, job_id - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) - "#, - params![ - run.id.to_string(), - run.routine_id.to_string(), - run.trigger_type.as_str(), - opt_text(run.trigger_detail.as_deref()), - fmt_ts(&run.started_at), - run.status.to_string(), - opt_text_owned(run.job_id.map(|id| id.to_string())), - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn complete_routine_run( - &self, - id: Uuid, - status: RunStatus, - result_summary: Option<&str>, - tokens_used: Option, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute( - r#" - UPDATE routine_runs SET - completed_at = ?5, status = ?2, - result_summary = ?3, tokens_used = ?4 - WHERE id = ?1 - "#, - params![ - id.to_string(), - status.to_string(), - opt_text(result_summary), - tokens_used.map(|t| t as i64), - now, - ], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn list_routine_runs( - &self, - routine_id: Uuid, - limit: i64, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - &format!( - "SELECT {} FROM routine_runs WHERE routine_id = ?1 ORDER BY started_at DESC LIMIT ?2", - ROUTINE_RUN_COLUMNS - ), - params![routine_id.to_string(), limit], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut runs = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - runs.push(row_to_routine_run_libsql(&row)?); - } - Ok(runs) - } - - async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT COUNT(*) as cnt FROM routine_runs WHERE routine_id = ?1 AND status = 'running'", - params![routine_id.to_string()], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(get_i64(&row, 0)), - None => Ok(0), - } - } - - // ==================== Tool Failures ==================== - - async fn record_tool_failure( - &self, - tool_name: &str, - error_message: &str, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute( - r#" - INSERT INTO tool_failures (id, tool_name, error_message, error_count, last_failure) - VALUES (?1, ?2, ?3, 1, ?4) - ON CONFLICT (tool_name) DO UPDATE SET - error_message = ?3, - error_count = tool_failures.error_count + 1, - last_failure = ?4 - "#, - params![Uuid::new_v4().to_string(), tool_name, error_message, now], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn get_broken_tools(&self, threshold: i32) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - r#" - SELECT tool_name, error_message, error_count, first_failure, last_failure, - last_build_result, repair_attempts - FROM tool_failures - WHERE error_count >= ?1 AND repaired_at IS NULL - ORDER BY error_count DESC - "#, - params![threshold as i64], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut tools = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - tools.push(BrokenTool { - name: get_text(&row, 0), - last_error: get_opt_text(&row, 1), - failure_count: get_i64(&row, 2) as u32, - first_failure: get_ts(&row, 3), - last_failure: get_ts(&row, 4), - last_build_result: get_opt_text(&row, 5) - .and_then(|s| serde_json::from_str(&s).ok()), - repair_attempts: get_i64(&row, 6) as u32, - }); - } - Ok(tools) - } - - async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute( - "UPDATE tool_failures SET repaired_at = ?2, error_count = 0 WHERE tool_name = ?1", - params![tool_name, now], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - conn.execute( - "UPDATE tool_failures SET repair_attempts = repair_attempts + 1 WHERE tool_name = ?1", - params![tool_name], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - // ==================== Settings ==================== - - async fn get_setting( - &self, - user_id: &str, - key: &str, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT value FROM settings WHERE user_id = ?1 AND key = ?2", - params![user_id, key], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(Some(get_json(&row, 0))), - None => Ok(None), - } - } - - async fn get_setting_full( - &self, - user_id: &str, - key: &str, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT key, value, updated_at FROM settings WHERE user_id = ?1 AND key = ?2", - params![user_id, key], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(Some(SettingRow { - key: get_text(&row, 0), - value: get_json(&row, 1), - updated_at: get_ts(&row, 2), - })), - None => Ok(None), - } - } - - async fn set_setting( - &self, - user_id: &str, - key: &str, - value: &serde_json::Value, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute( - r#" - INSERT INTO settings (user_id, key, value, updated_at) - VALUES (?1, ?2, ?3, ?4) - ON CONFLICT (user_id, key) DO UPDATE SET - value = excluded.value, - updated_at = ?4 - "#, - params![user_id, key, value.to_string(), now], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn delete_setting(&self, user_id: &str, key: &str) -> Result { - let conn = self.connect().await?; - let count = conn - .execute( - "DELETE FROM settings WHERE user_id = ?1 AND key = ?2", - params![user_id, key], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(count > 0) - } - - async fn list_settings(&self, user_id: &str) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT key, value, updated_at FROM settings WHERE user_id = ?1 ORDER BY key", - params![user_id], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut settings = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - settings.push(SettingRow { - key: get_text(&row, 0), - value: get_json(&row, 1), - updated_at: get_ts(&row, 2), - }); - } - Ok(settings) - } - - async fn get_all_settings( - &self, - user_id: &str, - ) -> Result, DatabaseError> { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT key, value FROM settings WHERE user_id = ?1", - params![user_id], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - let mut map = HashMap::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - map.insert(get_text(&row, 0), get_json(&row, 1)); - } - Ok(map) - } - - async fn set_all_settings( - &self, - user_id: &str, - settings: &HashMap, - ) -> Result<(), DatabaseError> { - let conn = self.connect().await?; - let now = fmt_ts(&Utc::now()); - conn.execute("BEGIN", ()) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - for (key, value) in settings { - if let Err(e) = conn - .execute( - r#" - INSERT INTO settings (user_id, key, value, updated_at) - VALUES (?1, ?2, ?3, ?4) - ON CONFLICT (user_id, key) DO UPDATE SET - value = excluded.value, - updated_at = ?4 - "#, - params![user_id, key.as_str(), value.to_string(), now.as_str()], - ) - .await - { - let _ = conn.execute("ROLLBACK", ()).await; - return Err(DatabaseError::Query(e.to_string())); - } - } - - conn.execute("COMMIT", ()) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - Ok(()) - } - - async fn has_settings(&self, user_id: &str) -> Result { - let conn = self.connect().await?; - let mut rows = conn - .query( - "SELECT COUNT(*) as cnt FROM settings WHERE user_id = ?1", - params![user_id], - ) - .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; - - match rows - .next() - .await - .map_err(|e| DatabaseError::Query(e.to_string()))? - { - Some(row) => Ok(get_i64(&row, 0) > 0), - None => Ok(false), - } - } - - // ==================== Workspace: Documents ==================== - - async fn get_document_by_path( - &self, - user_id: &str, - agent_id: Option, - path: &str, - ) -> Result { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let agent_id_str = agent_id.map(|id| id.to_string()); - let mut rows = conn - .query( - r#" - SELECT id, user_id, agent_id, path, content, - created_at, updated_at, metadata - FROM memory_documents - WHERE user_id = ?1 AND agent_id IS ?2 AND path = ?3 - "#, - params![user_id, agent_id_str.as_deref(), path], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })?; - - match rows - .next() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })? { - Some(row) => Ok(row_to_memory_document(&row)), - None => Err(WorkspaceError::DocumentNotFound { - doc_type: path.to_string(), - user_id: user_id.to_string(), - }), - } - } - - async fn get_document_by_id(&self, id: Uuid) -> Result { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let mut rows = conn - .query( - r#" - SELECT id, user_id, agent_id, path, content, - created_at, updated_at, metadata - FROM memory_documents WHERE id = ?1 - "#, - params![id.to_string()], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })?; - - match rows - .next() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })? { - Some(row) => Ok(row_to_memory_document(&row)), - None => Err(WorkspaceError::DocumentNotFound { - doc_type: "unknown".to_string(), - user_id: "unknown".to_string(), - }), - } - } - - async fn get_or_create_document_by_path( - &self, - user_id: &str, - agent_id: Option, - path: &str, - ) -> Result { - // Try get - match self.get_document_by_path(user_id, agent_id, path).await { - Ok(doc) => return Ok(doc), - Err(WorkspaceError::DocumentNotFound { .. }) => {} - Err(e) => return Err(e), - } - - // Create - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let id = Uuid::new_v4(); - let agent_id_str = agent_id.map(|id| id.to_string()); - conn.execute( - r#" - INSERT INTO memory_documents (id, user_id, agent_id, path, content, metadata) - VALUES (?1, ?2, ?3, ?4, '', '{}') - ON CONFLICT (user_id, agent_id, path) DO NOTHING - "#, - params![id.to_string(), user_id, agent_id_str.as_deref(), path], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Insert failed: {}", e), - })?; - - self.get_document_by_path(user_id, agent_id, path).await - } - - async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let now = fmt_ts(&Utc::now()); - conn.execute( - "UPDATE memory_documents SET content = ?2, updated_at = ?3 WHERE id = ?1", - params![id.to_string(), content, now], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Update failed: {}", e), - })?; - Ok(()) - } - - async fn delete_document_by_path( - &self, - user_id: &str, - agent_id: Option, - path: &str, - ) -> Result<(), WorkspaceError> { - let doc = self.get_document_by_path(user_id, agent_id, path).await?; - self.delete_chunks(doc.id).await?; - - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let agent_id_str = agent_id.map(|id| id.to_string()); - conn.execute( - "DELETE FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 AND path = ?3", - params![user_id, agent_id_str.as_deref(), path], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Delete failed: {}", e), - })?; - Ok(()) - } - - async fn list_directory( - &self, - user_id: &str, - agent_id: Option, - directory: &str, - ) -> Result, WorkspaceError> { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - // Implement the list_workspace_files logic in Rust instead of PL/pgSQL. - let dir = if !directory.is_empty() && !directory.ends_with('/') { - format!("{}/", directory) - } else { - directory.to_string() - }; - - let agent_id_str = agent_id.map(|id| id.to_string()); - let pattern = if dir.is_empty() { - "%".to_string() - } else { - format!("{}%", dir) - }; - - let mut rows = conn - .query( - r#" - SELECT path, updated_at, substr(content, 1, 200) as content_preview - FROM memory_documents - WHERE user_id = ?1 AND agent_id IS ?2 - AND (?3 = '%' OR path LIKE ?3) - ORDER BY path - "#, - params![user_id, agent_id_str.as_deref(), pattern], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("List directory failed: {}", e), - })?; - - let mut entries_map: HashMap = HashMap::new(); - - while let Some(row) = rows - .next() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })? - { - let full_path = get_text(&row, 0); - let updated_at = get_opt_ts(&row, 1); - let content_preview = get_opt_text(&row, 2); - - // Extract the immediate child name relative to directory - let relative = if dir.is_empty() { - &full_path - } else if let Some(stripped) = full_path.strip_prefix(&dir) { - stripped - } else { - continue; - }; - - let child_name = if let Some(slash_pos) = relative.find('/') { - &relative[..slash_pos] - } else { - relative - }; - - if child_name.is_empty() { - continue; - } - - let is_dir = relative.contains('/'); - let entry_path = if dir.is_empty() { - child_name.to_string() - } else { - format!("{}{}", dir, child_name) - }; - - entries_map - .entry(child_name.to_string()) - .and_modify(|e| { - // Mark as directory if any sub-paths exist - if is_dir { - e.is_directory = true; - e.content_preview = None; - } - // Update to latest timestamp - if let (Some(existing), Some(new)) = (&e.updated_at, &updated_at) - && new > existing - { - e.updated_at = Some(*new); - } - }) - .or_insert(WorkspaceEntry { - path: entry_path, - is_directory: is_dir, - updated_at, - content_preview: if is_dir { None } else { content_preview }, - }); - } - - let mut entries: Vec = entries_map.into_values().collect(); - entries.sort_by(|a, b| a.path.cmp(&b.path)); - Ok(entries) - } - - async fn list_all_paths( - &self, - user_id: &str, - agent_id: Option, - ) -> Result, WorkspaceError> { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let agent_id_str = agent_id.map(|id| id.to_string()); - let mut rows = conn - .query( - "SELECT path FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 ORDER BY path", - params![user_id, agent_id_str.as_deref()], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("List paths failed: {}", e), - })?; - - let mut paths = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })? - { - paths.push(get_text(&row, 0)); - } - Ok(paths) - } - - async fn list_documents( - &self, - user_id: &str, - agent_id: Option, - ) -> Result, WorkspaceError> { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let agent_id_str = agent_id.map(|id| id.to_string()); - let mut rows = conn - .query( - r#" - SELECT id, user_id, agent_id, path, content, - created_at, updated_at, metadata - FROM memory_documents - WHERE user_id = ?1 AND agent_id IS ?2 - ORDER BY updated_at DESC - "#, - params![user_id, agent_id_str.as_deref()], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })?; - - let mut docs = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })? - { - docs.push(row_to_memory_document(&row)); - } - Ok(docs) - } - - // ==================== Workspace: Chunks ==================== - - async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::ChunkingFailed { - reason: e.to_string(), - })?; - conn.execute( - "DELETE FROM memory_chunks WHERE document_id = ?1", - params![document_id.to_string()], - ) - .await - .map_err(|e| WorkspaceError::ChunkingFailed { - reason: format!("Delete failed: {}", e), - })?; - Ok(()) - } - - async fn insert_chunk( - &self, - document_id: Uuid, - chunk_index: i32, - content: &str, - embedding: Option<&[f32]>, - ) -> Result { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::ChunkingFailed { - reason: e.to_string(), - })?; - let id = Uuid::new_v4(); - let embedding_blob = embedding.map(|e| { - // Convert f32 slice to bytes for F32_BLOB - let bytes: Vec = e.iter().flat_map(|f| f.to_le_bytes()).collect(); - bytes - }); - - conn.execute( - r#" - INSERT INTO memory_chunks (id, document_id, chunk_index, content, embedding) - VALUES (?1, ?2, ?3, ?4, ?5) - "#, - params![ - id.to_string(), - document_id.to_string(), - chunk_index as i64, - content, - embedding_blob.map(libsql::Value::Blob), - ], - ) - .await - .map_err(|e| WorkspaceError::ChunkingFailed { - reason: format!("Insert failed: {}", e), - })?; - Ok(id) - } - - async fn update_chunk_embedding( - &self, - chunk_id: Uuid, - embedding: &[f32], - ) -> Result<(), WorkspaceError> { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::EmbeddingFailed { - reason: e.to_string(), - })?; - let bytes: Vec = embedding.iter().flat_map(|f| f.to_le_bytes()).collect(); - - conn.execute( - "UPDATE memory_chunks SET embedding = ?2 WHERE id = ?1", - params![chunk_id.to_string(), libsql::Value::Blob(bytes)], - ) - .await - .map_err(|e| WorkspaceError::EmbeddingFailed { - reason: format!("Update failed: {}", e), - })?; - Ok(()) - } - - async fn get_chunks_without_embeddings( - &self, - user_id: &str, - agent_id: Option, - limit: usize, - ) -> Result, WorkspaceError> { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let agent_id_str = agent_id.map(|id| id.to_string()); - let mut rows = conn - .query( - r#" - SELECT c.id, c.document_id, c.chunk_index, c.content, c.created_at - FROM memory_chunks c - JOIN memory_documents d ON d.id = c.document_id - WHERE d.user_id = ?1 AND d.agent_id IS ?2 - AND c.embedding IS NULL - LIMIT ?3 - "#, - params![user_id, agent_id_str.as_deref(), limit as i64], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })?; - - let mut chunks = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Query failed: {}", e), - })? - { - chunks.push(MemoryChunk { - id: get_text(&row, 0).parse().unwrap_or_default(), - document_id: get_text(&row, 1).parse().unwrap_or_default(), - chunk_index: get_i64(&row, 2) as i32, - content: get_text(&row, 3), - embedding: None, - created_at: get_ts(&row, 4), - }); - } - Ok(chunks) - } - - // ==================== Workspace: Search ==================== - - async fn hybrid_search( - &self, - user_id: &str, - agent_id: Option, - query: &str, - embedding: Option<&[f32]>, - config: &SearchConfig, - ) -> Result, WorkspaceError> { - let conn = self - .connect() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; - let agent_id_str = agent_id.map(|id| id.to_string()); - let pre_limit = config.pre_fusion_limit as i64; - - // FTS search using FTS5 - let fts_results = if config.use_fts { - let mut rows = conn - .query( - r#" - SELECT c.id, c.document_id, c.content - FROM memory_chunks_fts fts - JOIN memory_chunks c ON c._rowid = fts.rowid - JOIN memory_documents d ON d.id = c.document_id - WHERE d.user_id = ?1 AND d.agent_id IS ?2 - AND memory_chunks_fts MATCH ?3 - ORDER BY rank - LIMIT ?4 - "#, - params![user_id, agent_id_str.as_deref(), query, pre_limit], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("FTS query failed: {}", e), - })?; - - let mut results = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("FTS row fetch failed: {}", e), - })? - { - results.push(RankedResult { - chunk_id: get_text(&row, 0).parse().unwrap_or_default(), - document_id: get_text(&row, 1).parse().unwrap_or_default(), - content: get_text(&row, 2), - rank: results.len() as u32 + 1, - }); - } - results - } else { - Vec::new() - }; - - // Vector search using libsql_vector_idx - let vector_results = if let (true, Some(emb)) = (config.use_vector, embedding) { - // Format as JSON array string for vector() SQL function - let vector_json = format!( - "[{}]", - emb.iter() - .map(|f| f.to_string()) - .collect::>() - .join(",") - ); - - // vector_top_k returns rowids from the vector index. - // We join back to memory_chunks and filter by user/agent. - let mut rows = conn - .query( - r#" - SELECT c.id, c.document_id, c.content - FROM vector_top_k('idx_memory_chunks_embedding', vector(?1), ?2) AS top_k - JOIN memory_chunks c ON c._rowid = top_k.id - JOIN memory_documents d ON d.id = c.document_id - WHERE d.user_id = ?3 AND d.agent_id IS ?4 - "#, - params![vector_json, pre_limit, user_id, agent_id_str.as_deref()], - ) - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Vector query failed: {}", e), - })?; - - let mut results = Vec::new(); - while let Some(row) = rows - .next() - .await - .map_err(|e| WorkspaceError::SearchFailed { - reason: format!("Vector row fetch failed: {}", e), - })? - { - results.push(RankedResult { - chunk_id: get_text(&row, 0).parse().unwrap_or_default(), - document_id: get_text(&row, 1).parse().unwrap_or_default(), - content: get_text(&row, 2), - rank: results.len() as u32 + 1, - }); - } - results - } else { - Vec::new() - }; - - if embedding.is_some() && !config.use_vector { - tracing::warn!( - "Embedding provided but vector search is disabled in config; using FTS-only results" - ); - } - - Ok(reciprocal_rank_fusion(fts_results, vector_results, config)) - } -} - -// ==================== Row conversion helpers ==================== - -fn row_to_memory_document(row: &libsql::Row) -> MemoryDocument { - MemoryDocument { - id: get_text(row, 0).parse().unwrap_or_default(), - user_id: get_text(row, 1), - agent_id: get_opt_text(row, 2).and_then(|s| s.parse().ok()), - path: get_text(row, 3), - content: get_text(row, 4), - created_at: get_ts(row, 5), - updated_at: get_ts(row, 6), - metadata: get_json(row, 7), - } -} - -fn row_to_routine_libsql(row: &libsql::Row) -> Result { - let trigger_type = get_text(row, 5); - let trigger_config = get_json(row, 6); - let action_type = get_text(row, 7); - let action_config = get_json(row, 8); - let cooldown_secs = get_i64(row, 9); - let max_concurrent = get_i64(row, 10); - let dedup_window_secs: Option = row.get::(11).ok(); - - let trigger = - Trigger::from_db(&trigger_type, trigger_config).map_err(DatabaseError::Serialization)?; - let action = RoutineAction::from_db(&action_type, action_config) - .map_err(DatabaseError::Serialization)?; - - Ok(Routine { - id: get_text(row, 0).parse().unwrap_or_default(), - name: get_text(row, 1), - description: get_text(row, 2), - user_id: get_text(row, 3), - enabled: get_i64(row, 4) != 0, - trigger, - action, - guardrails: RoutineGuardrails { - cooldown: std::time::Duration::from_secs(cooldown_secs as u64), - max_concurrent: max_concurrent as u32, - dedup_window: dedup_window_secs.map(|s| std::time::Duration::from_secs(s as u64)), - }, - notify: NotifyConfig { - channel: get_opt_text(row, 12), - user: get_text(row, 13), - on_success: get_i64(row, 14) != 0, - on_failure: get_i64(row, 15) != 0, - on_attention: get_i64(row, 16) != 0, - }, - state: get_json(row, 17), - last_run_at: get_opt_ts(row, 18), - next_fire_at: get_opt_ts(row, 19), - run_count: get_i64(row, 20) as u64, - consecutive_failures: get_i64(row, 21) as u32, - created_at: get_ts(row, 22), - updated_at: get_ts(row, 23), - }) -} - -fn row_to_routine_run_libsql(row: &libsql::Row) -> Result { - let status_str = get_text(row, 5); - let status: RunStatus = status_str - .parse() - .map_err(|e: String| DatabaseError::Serialization(e))?; - - Ok(RoutineRun { - id: get_text(row, 0).parse().unwrap_or_default(), - routine_id: get_text(row, 1).parse().unwrap_or_default(), - trigger_type: get_text(row, 2), - trigger_detail: get_opt_text(row, 3), - started_at: get_ts(row, 4), - completed_at: get_opt_ts(row, 6), - status, - result_summary: get_opt_text(row, 7), - tokens_used: row.get::(8).ok().map(|v| v as i32), - job_id: get_opt_text(row, 9).and_then(|s| s.parse().ok()), - created_at: get_ts(row, 10), - }) -} - -#[cfg(test)] -mod tests { - use crate::db::Database; - use crate::db::libsql_backend::LibSqlBackend; - - #[tokio::test] - async fn test_wal_mode_after_migrations() { - let backend = LibSqlBackend::new_memory().await.unwrap(); - backend.run_migrations().await.unwrap(); - - let conn = backend.connect().await.unwrap(); - let mut rows = conn.query("PRAGMA journal_mode", ()).await.unwrap(); - let row = rows.next().await.unwrap().unwrap(); - let mode: String = row.get(0).unwrap(); - // In-memory databases use "memory" journal mode (WAL doesn't apply to :memory:), - // but the PRAGMA still executes without error. For file-based databases it returns "wal". - assert!( - mode == "wal" || mode == "memory", - "expected wal or memory, got: {}", - mode, - ); - } - - #[tokio::test] - async fn test_busy_timeout_set_on_connect() { - let backend = LibSqlBackend::new_memory().await.unwrap(); - backend.run_migrations().await.unwrap(); - - let conn = backend.connect().await.unwrap(); - let mut rows = conn.query("PRAGMA busy_timeout", ()).await.unwrap(); - let row = rows.next().await.unwrap().unwrap(); - let timeout: i64 = row.get(0).unwrap(); - assert_eq!(timeout, 5000); - } - - #[tokio::test] - async fn test_concurrent_writes_succeed() { - // Use a temp file so connections share state (in-memory DBs are connection-local) - let dir = tempfile::tempdir().unwrap(); - let db_path = dir.path().join("test_concurrent.db"); - let backend = LibSqlBackend::new_local(&db_path).await.unwrap(); - backend.run_migrations().await.unwrap(); - - // Spawn 20 concurrent inserts into the conversations table - let mut handles = Vec::new(); - for i in 0..20 { - let conn = backend.connect().await.unwrap(); - let handle = tokio::spawn(async move { - let id = uuid::Uuid::new_v4().to_string(); - let val = format!("ch_{}", i); - conn.execute( - "INSERT INTO conversations (id, channel, user_id) VALUES (?1, ?2, ?3)", - libsql::params![id, val, "test_user"], - ) - .await - }); - handles.push(handle); - } - - for handle in handles { - let result = handle.await.unwrap(); - assert!( - result.is_ok(), - "concurrent write failed: {:?}", - result.err() - ); - } - - // Verify all 20 rows landed - let conn = backend.connect().await.unwrap(); - let mut rows = conn - .query( - "SELECT COUNT(*) FROM conversations WHERE user_id = ?1", - libsql::params!["test_user"], - ) - .await - .unwrap(); - let row = rows.next().await.unwrap().unwrap(); - let count: i64 = row.get(0).unwrap(); - assert_eq!(count, 20); - } -} diff --git a/src/db/mod.rs b/src/db/mod.rs index 40e85cae..86c9d568 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -13,7 +13,7 @@ pub mod postgres; #[cfg(feature = "libsql")] -pub mod libsql_backend; +pub mod libsql; #[cfg(feature = "libsql")] pub mod libsql_migrations; @@ -62,15 +62,11 @@ pub async fn connect_from_config( "LIBSQL_AUTH_TOKEN required when LIBSQL_URL is set".to_string(), ) })?; - libsql_backend::LibSqlBackend::new_remote_replica( - db_path, - url, - token.expose_secret(), - ) - .await - .map_err(|e| DatabaseError::Pool(e.to_string()))? + libsql::LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret()) + .await + .map_err(|e| DatabaseError::Pool(e.to_string()))? } else { - libsql_backend::LibSqlBackend::new_local(db_path) + libsql::LibSqlBackend::new_local(db_path) .await .map_err(|e| DatabaseError::Pool(e.to_string()))? }; @@ -92,37 +88,27 @@ pub async fn connect_from_config( } } -/// Backend-agnostic database trait. -/// -/// Combines all persistence operations from Store, Repository, and related -/// stores into a single trait that can be implemented for different backends. +// ==================== Sub-traits ==================== +// +// Each sub-trait groups related persistence methods. The `Database` supertrait +// combines them all, so existing `Arc` consumers keep working. +// Leaf consumers can depend on a specific sub-trait instead. + #[async_trait] -pub trait Database: Send + Sync { - /// Run schema migrations for this backend. - async fn run_migrations(&self) -> Result<(), DatabaseError>; - - // ==================== Conversations ==================== - - /// Create a new conversation. +pub trait ConversationStore: Send + Sync { async fn create_conversation( &self, channel: &str, user_id: &str, thread_id: Option<&str>, ) -> Result; - - /// Update conversation last activity. async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError>; - - /// Add a message to a conversation. async fn add_conversation_message( &self, conversation_id: Uuid, role: &str, content: &str, ) -> Result; - - /// Ensure a conversation row exists (upsert). async fn ensure_conversation( &self, id: Uuid, @@ -130,103 +116,65 @@ pub trait Database: Send + Sync { user_id: &str, thread_id: Option<&str>, ) -> Result<(), DatabaseError>; - - /// List conversations with a title preview. async fn list_conversations_with_preview( &self, user_id: &str, channel: &str, limit: i64, ) -> Result, DatabaseError>; - - /// Get or create the singleton assistant conversation. async fn get_or_create_assistant_conversation( &self, user_id: &str, channel: &str, ) -> Result; - - /// Create a conversation with specific metadata. async fn create_conversation_with_metadata( &self, channel: &str, user_id: &str, metadata: &serde_json::Value, ) -> Result; - - /// Load messages with cursor-based pagination. async fn list_conversation_messages_paginated( &self, conversation_id: Uuid, before: Option>, limit: i64, ) -> Result<(Vec, bool), DatabaseError>; - - /// Merge a single key into conversation metadata. async fn update_conversation_metadata_field( &self, id: Uuid, key: &str, value: &serde_json::Value, ) -> Result<(), DatabaseError>; - - /// Read conversation metadata. async fn get_conversation_metadata( &self, id: Uuid, ) -> Result, DatabaseError>; - - /// Load all messages for a conversation. async fn list_conversation_messages( &self, conversation_id: Uuid, ) -> Result, DatabaseError>; - - /// Check if a conversation belongs to a specific user. async fn conversation_belongs_to_user( &self, conversation_id: Uuid, user_id: &str, ) -> Result; +} - // ==================== Jobs ==================== - - /// Save a job context. +#[async_trait] +pub trait JobStore: Send + Sync { async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError>; - - /// Get a job by ID. async fn get_job(&self, id: Uuid) -> Result, DatabaseError>; - - /// Update job status. async fn update_job_status( &self, id: Uuid, status: JobState, failure_reason: Option<&str>, ) -> Result<(), DatabaseError>; - - /// Mark job as stuck. async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError>; - - /// Get stuck jobs. async fn get_stuck_jobs(&self) -> Result, DatabaseError>; - - // ==================== Actions ==================== - - /// Save a job action. async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError>; - - /// Get actions for a job. async fn get_job_actions(&self, job_id: Uuid) -> Result, DatabaseError>; - - // ==================== LLM Calls ==================== - - /// Record an LLM call. async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result; - - // ==================== Estimation Snapshots ==================== - - /// Save an estimation snapshot. async fn save_estimation_snapshot( &self, job_id: Uuid, @@ -236,8 +184,6 @@ pub trait Database: Send + Sync { estimated_time_secs: i32, estimated_value: Decimal, ) -> Result; - - /// Update estimation snapshot with actual values. async fn update_estimation_actuals( &self, id: Uuid, @@ -245,19 +191,13 @@ pub trait Database: Send + Sync { actual_time_secs: i32, actual_value: Option, ) -> Result<(), DatabaseError>; +} - // ==================== Sandbox Jobs ==================== - - /// Insert a new sandbox job. +#[async_trait] +pub trait SandboxStore: Send + Sync { async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError>; - - /// Get a sandbox job by ID. async fn get_sandbox_job(&self, id: Uuid) -> Result, DatabaseError>; - - /// List all sandbox jobs, most recent first. async fn list_sandbox_jobs(&self) -> Result, DatabaseError>; - - /// Update sandbox job status. async fn update_sandbox_job_status( &self, id: Uuid, @@ -267,83 +207,49 @@ pub trait Database: Send + Sync { started_at: Option>, completed_at: Option>, ) -> Result<(), DatabaseError>; - - /// Mark stale sandbox jobs as interrupted. async fn cleanup_stale_sandbox_jobs(&self) -> Result; - - /// Get sandbox job summary. async fn sandbox_job_summary(&self) -> Result; - - /// List sandbox jobs for a specific user, most recent first. async fn list_sandbox_jobs_for_user( &self, user_id: &str, ) -> Result, DatabaseError>; - - /// Get sandbox job summary for a specific user. async fn sandbox_job_summary_for_user( &self, user_id: &str, ) -> Result; - - /// Check if a sandbox job belongs to a specific user. async fn sandbox_job_belongs_to_user( &self, job_id: Uuid, user_id: &str, ) -> Result; - - /// Update sandbox job mode. async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError>; - - /// Get sandbox job mode. async fn get_sandbox_job_mode(&self, id: Uuid) -> Result, DatabaseError>; - - // ==================== Job Events ==================== - - /// Persist a job event. async fn save_job_event( &self, job_id: Uuid, event_type: &str, data: &serde_json::Value, ) -> Result<(), DatabaseError>; - - /// Load job events, returning the most recent `limit` entries (or all if `None`). async fn list_job_events( &self, job_id: Uuid, limit: Option, ) -> Result, DatabaseError>; +} - // ==================== Routines ==================== - - /// Create a new routine. +#[async_trait] +pub trait RoutineStore: Send + Sync { async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError>; - - /// Get a routine by ID. async fn get_routine(&self, id: Uuid) -> Result, DatabaseError>; - - /// Get a routine by user_id and name. async fn get_routine_by_name( &self, user_id: &str, name: &str, ) -> Result, DatabaseError>; - - /// List routines for a user. async fn list_routines(&self, user_id: &str) -> Result, DatabaseError>; - - /// List all enabled event routines. async fn list_event_routines(&self) -> Result, DatabaseError>; - - /// List due cron routines. async fn list_due_cron_routines(&self) -> Result, DatabaseError>; - - /// Update a routine. async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError>; - - /// Update runtime state after a routine fires. async fn update_routine_runtime( &self, id: Uuid, @@ -353,16 +259,8 @@ pub trait Database: Send + Sync { consecutive_failures: u32, state: &serde_json::Value, ) -> Result<(), DatabaseError>; - - /// Delete a routine. async fn delete_routine(&self, id: Uuid) -> Result; - - // ==================== Routine Runs ==================== - - /// Record a routine run starting. async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError>; - - /// Complete a routine run. async fn complete_routine_run( &self, id: Uuid, @@ -370,141 +268,97 @@ pub trait Database: Send + Sync { result_summary: Option<&str>, tokens_used: Option, ) -> Result<(), DatabaseError>; - - /// List recent runs for a routine. async fn list_routine_runs( &self, routine_id: Uuid, limit: i64, ) -> Result, DatabaseError>; - - /// Count currently running runs for a routine. async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result; +} - // ==================== Tool Failures ==================== - - /// Record a tool failure (upsert). +#[async_trait] +pub trait ToolFailureStore: Send + Sync { async fn record_tool_failure( &self, tool_name: &str, error_message: &str, ) -> Result<(), DatabaseError>; - - /// Get broken tools exceeding threshold. async fn get_broken_tools(&self, threshold: i32) -> Result, DatabaseError>; - - /// Mark a tool as repaired. async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError>; - - /// Increment repair attempts. async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError>; +} - // ==================== Settings ==================== - - /// Get a single setting. +#[async_trait] +pub trait SettingsStore: Send + Sync { async fn get_setting( &self, user_id: &str, key: &str, ) -> Result, DatabaseError>; - - /// Get a single setting with metadata. async fn get_setting_full( &self, user_id: &str, key: &str, ) -> Result, DatabaseError>; - - /// Set a single setting (upsert). async fn set_setting( &self, user_id: &str, key: &str, value: &serde_json::Value, ) -> Result<(), DatabaseError>; - - /// Delete a single setting. async fn delete_setting(&self, user_id: &str, key: &str) -> Result; - - /// List all settings for a user. async fn list_settings(&self, user_id: &str) -> Result, DatabaseError>; - - /// Get all settings as a flat map. async fn get_all_settings( &self, user_id: &str, ) -> Result, DatabaseError>; - - /// Bulk-write settings atomically. async fn set_all_settings( &self, user_id: &str, settings: &HashMap, ) -> Result<(), DatabaseError>; - - /// Check if settings exist for a user. async fn has_settings(&self, user_id: &str) -> Result; +} - // ==================== Workspace: Documents ==================== - - /// Get a document by path. +#[async_trait] +pub trait WorkspaceStore: Send + Sync { async fn get_document_by_path( &self, user_id: &str, agent_id: Option, path: &str, ) -> Result; - - /// Get a document by ID. async fn get_document_by_id(&self, id: Uuid) -> Result; - - /// Get or create a document by path. async fn get_or_create_document_by_path( &self, user_id: &str, agent_id: Option, path: &str, ) -> Result; - - /// Update a document's content. async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError>; - - /// Delete a document by path. async fn delete_document_by_path( &self, user_id: &str, agent_id: Option, path: &str, ) -> Result<(), WorkspaceError>; - - /// List files and directories in a directory path. async fn list_directory( &self, user_id: &str, agent_id: Option, directory: &str, ) -> Result, WorkspaceError>; - - /// List all file paths in the workspace. async fn list_all_paths( &self, user_id: &str, agent_id: Option, ) -> Result, WorkspaceError>; - - /// List all documents for a user. async fn list_documents( &self, user_id: &str, agent_id: Option, ) -> Result, WorkspaceError>; - - // ==================== Workspace: Chunks ==================== - - /// Delete all chunks for a document. async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError>; - - /// Insert a chunk. async fn insert_chunk( &self, document_id: Uuid, @@ -512,25 +366,17 @@ pub trait Database: Send + Sync { content: &str, embedding: Option<&[f32]>, ) -> Result; - - /// Update a chunk's embedding. async fn update_chunk_embedding( &self, chunk_id: Uuid, embedding: &[f32], ) -> Result<(), WorkspaceError>; - - /// Get chunks without embeddings for backfilling. async fn get_chunks_without_embeddings( &self, user_id: &str, agent_id: Option, limit: usize, ) -> Result, WorkspaceError>; - - // ==================== Workspace: Search ==================== - - /// Perform hybrid search combining FTS and vector similarity. async fn hybrid_search( &self, user_id: &str, @@ -540,3 +386,23 @@ pub trait Database: Send + Sync { config: &SearchConfig, ) -> Result, WorkspaceError>; } + +/// Backend-agnostic database supertrait. +/// +/// Combines all sub-traits into one. Existing `Arc` consumers +/// continue to work; leaf consumers can depend on a specific sub-trait instead. +#[async_trait] +pub trait Database: + ConversationStore + + JobStore + + SandboxStore + + RoutineStore + + ToolFailureStore + + SettingsStore + + WorkspaceStore + + Send + + Sync +{ + /// Run schema migrations for this backend. + async fn run_migrations(&self) -> Result<(), DatabaseError>; +} diff --git a/src/db/postgres.rs b/src/db/postgres.rs index 096f3a95..9404dc7e 100644 --- a/src/db/postgres.rs +++ b/src/db/postgres.rs @@ -15,7 +15,10 @@ use crate::agent::BrokenTool; use crate::agent::routine::{Routine, RoutineRun, RunStatus}; use crate::config::DatabaseConfig; use crate::context::{ActionRecord, JobContext, JobState}; -use crate::db::Database; +use crate::db::{ + ConversationStore, Database, JobStore, RoutineStore, SandboxStore, SettingsStore, + ToolFailureStore, WorkspaceStore, +}; use crate::error::{DatabaseError, WorkspaceError}; use crate::history::{ ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord, @@ -51,14 +54,19 @@ impl PgBackend { } } +// ==================== Database (supertrait) ==================== + #[async_trait] impl Database for PgBackend { async fn run_migrations(&self) -> Result<(), DatabaseError> { self.store.run_migrations().await } +} - // ==================== Conversations ==================== +// ==================== ConversationStore ==================== +#[async_trait] +impl ConversationStore for PgBackend { async fn create_conversation( &self, channel: &str, @@ -174,9 +182,12 @@ impl Database for PgBackend { .conversation_belongs_to_user(conversation_id, user_id) .await } +} - // ==================== Jobs ==================== +// ==================== JobStore ==================== +#[async_trait] +impl JobStore for PgBackend { async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> { self.store.save_job(ctx).await } @@ -204,8 +215,6 @@ impl Database for PgBackend { self.store.get_stuck_jobs().await } - // ==================== Actions ==================== - async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> { self.store.save_action(job_id, action).await } @@ -214,14 +223,10 @@ impl Database for PgBackend { self.store.get_job_actions(job_id).await } - // ==================== LLM Calls ==================== - async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result { self.store.record_llm_call(record).await } - // ==================== Estimation Snapshots ==================== - async fn save_estimation_snapshot( &self, job_id: Uuid, @@ -254,9 +259,12 @@ impl Database for PgBackend { .update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value) .await } +} - // ==================== Sandbox Jobs ==================== +// ==================== SandboxStore ==================== +#[async_trait] +impl SandboxStore for PgBackend { async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> { self.store.save_sandbox_job(job).await } @@ -323,8 +331,6 @@ impl Database for PgBackend { self.store.get_sandbox_job_mode(id).await } - // ==================== Job Events ==================== - async fn save_job_event( &self, job_id: Uuid, @@ -341,9 +347,12 @@ impl Database for PgBackend { ) -> Result, DatabaseError> { self.store.list_job_events(job_id, limit).await } +} - // ==================== Routines ==================== +// ==================== RoutineStore ==================== +#[async_trait] +impl RoutineStore for PgBackend { async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { self.store.create_routine(routine).await } @@ -401,8 +410,6 @@ impl Database for PgBackend { self.store.delete_routine(id).await } - // ==================== Routine Runs ==================== - async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> { self.store.create_routine_run(run).await } @@ -430,9 +437,12 @@ impl Database for PgBackend { async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result { self.store.count_running_routine_runs(routine_id).await } +} - // ==================== Tool Failures ==================== +// ==================== ToolFailureStore ==================== +#[async_trait] +impl ToolFailureStore for PgBackend { async fn record_tool_failure( &self, tool_name: &str, @@ -454,9 +464,12 @@ impl Database for PgBackend { async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> { self.store.increment_repair_attempts(tool_name).await } +} - // ==================== Settings ==================== +// ==================== SettingsStore ==================== +#[async_trait] +impl SettingsStore for PgBackend { async fn get_setting( &self, user_id: &str, @@ -508,9 +521,12 @@ impl Database for PgBackend { async fn has_settings(&self, user_id: &str) -> Result { self.store.has_settings(user_id).await } +} - // ==================== Workspace: Documents ==================== +// ==================== WorkspaceStore ==================== +#[async_trait] +impl WorkspaceStore for PgBackend { async fn get_document_by_path( &self, user_id: &str, @@ -577,8 +593,6 @@ impl Database for PgBackend { self.repo.list_documents(user_id, agent_id).await } - // ==================== Workspace: Chunks ==================== - async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> { self.repo.delete_chunks(document_id).await } @@ -614,8 +628,6 @@ impl Database for PgBackend { .await } - // ==================== Workspace: Search ==================== - async fn hybrid_search( &self, user_id: &str, diff --git a/src/lib.rs b/src/lib.rs index 69fa2f61..202fcbb0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -39,6 +39,7 @@ //! - **Continuous learning** - Improve estimates from historical data pub mod agent; +pub mod app; pub mod boot_screen; pub mod bootstrap; pub mod channels; @@ -70,6 +71,9 @@ pub mod util; pub mod worker; pub mod workspace; +#[cfg(test)] +pub mod testing; + pub use config::Config; pub use error::{Error, Result}; diff --git a/src/llm/circuit_breaker.rs b/src/llm/circuit_breaker.rs index 7ac39970..8f7718fd 100644 --- a/src/llm/circuit_breaker.rs +++ b/src/llm/circuit_breaker.rs @@ -298,121 +298,7 @@ impl LlmProvider for CircuitBreakerProvider { mod tests { use super::*; - use std::sync::atomic::{AtomicBool, Ordering}; - - use crate::llm::provider::{CompletionResponse, FinishReason, ToolCompletionResponse}; - - /// A test stub that either always succeeds or always fails with a - /// configurable error. The `should_fail` flag can be flipped at - /// runtime for half-open recovery tests. - struct StubProvider { - name: String, - should_fail: AtomicBool, - error_kind: StubError, - } - - #[derive(Clone)] - enum StubError { - Transient, - NonTransient, - } - - impl StubProvider { - fn always_ok(name: &str) -> Arc { - Arc::new(Self { - name: name.to_string(), - should_fail: AtomicBool::new(false), - error_kind: StubError::Transient, - }) - } - - fn always_fail(name: &str) -> Arc { - Arc::new(Self { - name: name.to_string(), - should_fail: AtomicBool::new(true), - error_kind: StubError::Transient, - }) - } - - fn always_fail_non_transient(name: &str) -> Arc { - Arc::new(Self { - name: name.to_string(), - should_fail: AtomicBool::new(true), - error_kind: StubError::NonTransient, - }) - } - - fn set_failing(&self, fail: bool) { - self.should_fail.store(fail, Ordering::Relaxed); - } - - fn make_error(&self) -> LlmError { - match self.error_kind { - StubError::Transient => LlmError::RequestFailed { - provider: self.name.clone(), - reason: "server error".to_string(), - }, - StubError::NonTransient => LlmError::ContextLengthExceeded { - used: 100_000, - limit: 50_000, - }, - } - } - - fn ok_response() -> CompletionResponse { - CompletionResponse { - content: "ok".to_string(), - input_tokens: 10, - output_tokens: 5, - finish_reason: FinishReason::Stop, - response_id: None, - } - } - - fn ok_tool_response() -> ToolCompletionResponse { - ToolCompletionResponse { - content: Some("ok".to_string()), - tool_calls: vec![], - input_tokens: 10, - output_tokens: 5, - finish_reason: FinishReason::Stop, - response_id: None, - } - } - } - - #[async_trait] - impl LlmProvider for StubProvider { - fn model_name(&self) -> &str { - &self.name - } - - fn cost_per_token(&self) -> (Decimal, Decimal) { - (Decimal::ZERO, Decimal::ZERO) - } - - async fn complete( - &self, - _request: CompletionRequest, - ) -> Result { - if self.should_fail.load(Ordering::Relaxed) { - Err(self.make_error()) - } else { - Ok(Self::ok_response()) - } - } - - async fn complete_with_tools( - &self, - _request: ToolCompletionRequest, - ) -> Result { - if self.should_fail.load(Ordering::Relaxed) { - Err(self.make_error()) - } else { - Ok(Self::ok_tool_response()) - } - } - } + use crate::testing::StubLlm; fn make_request() -> CompletionRequest { CompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")]) @@ -434,7 +320,7 @@ mod tests { #[tokio::test] async fn closed_allows_calls_and_resets_on_success() { - let stub = StubProvider::always_ok("test"); + let stub = Arc::new(StubLlm::new("ok").with_model_name("test")); let cb = CircuitBreakerProvider::new(stub, fast_config(3)); let resp = cb.complete(make_request()).await; @@ -445,7 +331,7 @@ mod tests { #[tokio::test] async fn failures_accumulate_then_trip_to_open() { - let stub = StubProvider::always_fail("test"); + let stub = Arc::new(StubLlm::failing("test")); let cb = CircuitBreakerProvider::new(stub, fast_config(3)); // First 2 failures: still closed @@ -462,7 +348,7 @@ mod tests { #[tokio::test] async fn open_rejects_immediately() { - let stub = StubProvider::always_fail("test"); + let stub = Arc::new(StubLlm::failing("test")); let cb = CircuitBreakerProvider::new( stub, CircuitBreakerConfig { @@ -492,7 +378,7 @@ mod tests { #[tokio::test] async fn recovery_timeout_transitions_to_half_open() { - let stub = StubProvider::always_fail("test"); + let stub = Arc::new(StubLlm::failing("test")); let cb = CircuitBreakerProvider::new(stub, fast_config(1)); // Trip to open @@ -510,7 +396,7 @@ mod tests { #[tokio::test] async fn half_open_success_closes_circuit() { - let stub = StubProvider::always_fail("test"); + let stub = Arc::new(StubLlm::failing("test")); let cb = CircuitBreakerProvider::new(stub.clone(), fast_config(1)); // Trip to open @@ -530,7 +416,7 @@ mod tests { #[tokio::test] async fn half_open_failure_reopens_circuit() { - let stub = StubProvider::always_fail("test"); + let stub = Arc::new(StubLlm::failing("test")); let cb = CircuitBreakerProvider::new(stub, fast_config(1)); // Trip to open @@ -546,7 +432,7 @@ mod tests { #[tokio::test] async fn non_transient_errors_do_not_trip_breaker() { - let stub = StubProvider::always_fail_non_transient("test"); + let stub = Arc::new(StubLlm::failing_non_transient("test")); let cb = CircuitBreakerProvider::new(stub, fast_config(1)); // ContextLengthExceeded is not transient; breaker should stay closed @@ -559,7 +445,7 @@ mod tests { #[tokio::test] async fn success_resets_failure_count() { - let stub = StubProvider::always_fail("test"); + let stub = Arc::new(StubLlm::failing("test")); let cb = CircuitBreakerProvider::new(stub.clone(), fast_config(3)); // Accumulate 2 failures @@ -576,7 +462,7 @@ mod tests { #[tokio::test] async fn complete_with_tools_uses_same_breaker_logic() { - let stub = StubProvider::always_fail("test"); + let stub = Arc::new(StubLlm::failing("test")); let cb = CircuitBreakerProvider::new(stub, fast_config(2)); let _ = cb.complete_with_tools(make_tool_request()).await; @@ -586,7 +472,7 @@ mod tests { #[tokio::test] async fn multiple_half_open_successes_needed() { - let stub = StubProvider::always_fail("test"); + let stub = Arc::new(StubLlm::failing("test")); let cb = CircuitBreakerProvider::new( stub.clone(), CircuitBreakerConfig { @@ -663,7 +549,7 @@ mod tests { #[tokio::test] async fn passthrough_methods_delegate_to_inner() { - let stub = StubProvider::always_ok("my-model"); + let stub = Arc::new(StubLlm::new("ok").with_model_name("my-model")); let cb = CircuitBreakerProvider::new(stub, fast_config(3)); assert_eq!(cb.model_name(), "my-model"); diff --git a/src/llm/response_cache.rs b/src/llm/response_cache.rs index cd91cabf..0f8468df 100644 --- a/src/llm/response_cache.rs +++ b/src/llm/response_cache.rs @@ -235,75 +235,9 @@ impl LlmProvider for CachedProvider { #[cfg(test)] mod tests { - use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; - - use crate::llm::provider::{ChatMessage, FinishReason}; + use crate::llm::provider::ChatMessage; use crate::llm::response_cache::*; - - /// Controllable stub provider for testing cache behavior. - struct StubProvider { - call_count: AtomicU32, - should_fail: AtomicBool, - } - - impl StubProvider { - fn new() -> Self { - Self { - call_count: AtomicU32::new(0), - should_fail: AtomicBool::new(false), - } - } - - fn calls(&self) -> u32 { - self.call_count.load(Ordering::Relaxed) - } - } - - #[async_trait] - impl LlmProvider for StubProvider { - fn model_name(&self) -> &str { - "stub-model" - } - - fn cost_per_token(&self) -> (Decimal, Decimal) { - (Decimal::ZERO, Decimal::ZERO) - } - - async fn complete( - &self, - _request: CompletionRequest, - ) -> Result { - self.call_count.fetch_add(1, Ordering::Relaxed); - if self.should_fail.load(Ordering::Relaxed) { - return Err(LlmError::RequestFailed { - provider: "stub".into(), - reason: "forced failure".into(), - }); - } - Ok(CompletionResponse { - content: "cached response".into(), - input_tokens: 10, - output_tokens: 5, - finish_reason: FinishReason::Stop, - response_id: None, - }) - } - - async fn complete_with_tools( - &self, - _request: ToolCompletionRequest, - ) -> Result { - self.call_count.fetch_add(1, Ordering::Relaxed); - Ok(ToolCompletionResponse { - content: Some("tool response".into()), - tool_calls: vec![], - input_tokens: 10, - output_tokens: 5, - finish_reason: FinishReason::Stop, - response_id: None, - }) - } - } + use crate::testing::StubLlm; fn simple_request() -> CompletionRequest { CompletionRequest { @@ -369,7 +303,7 @@ mod tests { #[tokio::test] async fn cache_hit_avoids_provider_call() { - let stub = Arc::new(StubProvider::new()); + let stub = Arc::new(StubLlm::new("cached response")); let cached = CachedProvider::new( stub.clone(), ResponseCacheConfig { @@ -393,7 +327,7 @@ mod tests { #[tokio::test] async fn different_messages_get_different_entries() { - let stub = Arc::new(StubProvider::new()); + let stub = Arc::new(StubLlm::new("cached response")); let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default()); cached.complete(simple_request()).await.unwrap(); @@ -405,7 +339,7 @@ mod tests { #[tokio::test] async fn expired_entries_are_evicted() { - let stub = Arc::new(StubProvider::new()); + let stub = Arc::new(StubLlm::new("cached response")); let cached = CachedProvider::new( stub.clone(), ResponseCacheConfig { @@ -427,7 +361,7 @@ mod tests { #[tokio::test] async fn lru_eviction_removes_oldest() { - let stub = Arc::new(StubProvider::new()); + let stub = Arc::new(StubLlm::new("cached response")); let cached = CachedProvider::new( stub.clone(), ResponseCacheConfig { @@ -456,7 +390,7 @@ mod tests { #[tokio::test] async fn tool_calls_are_never_cached() { - let stub = Arc::new(StubProvider::new()); + let stub = Arc::new(StubLlm::new("cached response")); let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default()); let req = ToolCompletionRequest { @@ -478,7 +412,7 @@ mod tests { #[tokio::test] async fn provider_errors_are_not_cached() { - let stub = Arc::new(StubProvider::new()); + let stub = Arc::new(StubLlm::new("cached response")); let cached = CachedProvider::new( stub.clone(), ResponseCacheConfig { @@ -487,20 +421,20 @@ mod tests { }, ); - stub.should_fail.store(true, Ordering::Relaxed); + stub.set_failing(true); let result = cached.complete(simple_request()).await; assert!(result.is_err()); assert!(cached.is_empty().await); // After fixing the provider, should succeed and cache - stub.should_fail.store(false, Ordering::Relaxed); + stub.set_failing(false); cached.complete(simple_request()).await.unwrap(); assert_eq!(cached.len().await, 1); } #[tokio::test] async fn clear_empties_cache() { - let stub = Arc::new(StubProvider::new()); + let stub = Arc::new(StubLlm::new("cached response")); let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default()); cached.complete(simple_request()).await.unwrap(); @@ -519,7 +453,7 @@ mod tests { #[tokio::test] async fn delegates_model_name() { - let stub = Arc::new(StubProvider::new()); + let stub = Arc::new(StubLlm::new("cached response")); let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default()); assert_eq!(cached.model_name(), "stub-model"); } diff --git a/src/main.rs b/src/main.rs index 97e69f90..a7b3e2af 100644 --- a/src/main.rs +++ b/src/main.rs @@ -386,7 +386,7 @@ async fn main() -> anyhow::Result<()> { #[cfg(feature = "libsql")] ironclaw::config::DatabaseBackend::LibSql => { use ironclaw::db::Database as _; - use ironclaw::db::libsql_backend::LibSqlBackend; + use ironclaw::db::libsql::LibSqlBackend; use secrecy::ExposeSecret as _; let default_path = ironclaw::config::default_libsql_path(); diff --git a/src/orchestrator/api.rs b/src/orchestrator/api.rs index 6ffd411f..8138dc1c 100644 --- a/src/orchestrator/api.rs +++ b/src/orchestrator/api.rs @@ -449,48 +449,17 @@ mod tests { use tower::ServiceExt; use uuid::Uuid; - use crate::error::LlmError; - use crate::llm::{ - CompletionRequest, CompletionResponse, ToolCompletionRequest, ToolCompletionResponse, - }; use crate::orchestrator::auth::TokenStore; use crate::orchestrator::job_manager::{ContainerJobConfig, ContainerJobManager}; + use crate::testing::StubLlm; use super::*; - /// Stub LLM provider that panics if called (tests only exercise routing/auth). - struct StubLlm; - - #[async_trait::async_trait] - impl crate::llm::LlmProvider for StubLlm { - fn model_name(&self) -> &str { - "stub" - } - fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) { - (rust_decimal::Decimal::ZERO, rust_decimal::Decimal::ZERO) - } - async fn complete(&self, _req: CompletionRequest) -> Result { - Err(LlmError::RequestFailed { - provider: "stub".into(), - reason: "not implemented".into(), - }) - } - async fn complete_with_tools( - &self, - _req: ToolCompletionRequest, - ) -> Result { - Err(LlmError::RequestFailed { - provider: "stub".into(), - reason: "not implemented".into(), - }) - } - } - fn test_state() -> OrchestratorState { let token_store = TokenStore::new(); let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone()); OrchestratorState { - llm: Arc::new(StubLlm), + llm: Arc::new(StubLlm::default()), job_manager: Arc::new(jm), token_store, job_event_tx: None, @@ -722,7 +691,7 @@ mod tests { .await; let state = OrchestratorState { - llm: Arc::new(StubLlm), + llm: Arc::new(StubLlm::default()), job_manager: Arc::new(jm), token_store, job_event_tx: None, @@ -757,7 +726,7 @@ mod tests { let token_store = TokenStore::new(); let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone()); let state = OrchestratorState { - llm: Arc::new(StubLlm), + llm: Arc::new(StubLlm::default()), job_manager: Arc::new(jm), token_store: token_store.clone(), job_event_tx: Some(tx), @@ -812,7 +781,7 @@ mod tests { let token_store = TokenStore::new(); let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone()); let state = OrchestratorState { - llm: Arc::new(StubLlm), + llm: Arc::new(StubLlm::default()), job_manager: Arc::new(jm), token_store: token_store.clone(), job_event_tx: Some(tx), @@ -860,7 +829,7 @@ mod tests { let token_store = TokenStore::new(); let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone()); let state = OrchestratorState { - llm: Arc::new(StubLlm), + llm: Arc::new(StubLlm::default()), job_manager: Arc::new(jm), token_store: token_store.clone(), job_event_tx: Some(tx), diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 4988abaf..e204b17c 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -79,7 +79,7 @@ pub struct SetupWizard { db_pool: Option, /// libSQL backend (created during setup, libsql only). #[cfg(feature = "libsql")] - db_backend: Option, + db_backend: Option, /// Secrets crypto (created during setup). secrets_crypto: Option>, /// Cached API key from provider setup (used by model fetcher without env mutation). @@ -438,7 +438,7 @@ impl SetupWizard { turso_url: Option<&str>, turso_token: Option<&str>, ) -> Result<(), SetupError> { - use crate::db::libsql_backend::LibSqlBackend; + use crate::db::libsql::LibSqlBackend; use std::path::Path; let db_path = Path::new(path); @@ -1486,7 +1486,7 @@ impl SetupWizard { #[cfg(feature = "libsql")] let saved = if !saved { if let Some(ref backend) = self.db_backend { - use crate::db::Database as _; + use crate::db::SettingsStore as _; backend .set_all_settings("default", &db_map) .await diff --git a/src/testing.rs b/src/testing.rs new file mode 100644 index 00000000..2bdc74ad --- /dev/null +++ b/src/testing.rs @@ -0,0 +1,356 @@ +//! Test harness for constructing `AgentDeps` with sensible defaults. +//! +//! Provides: +//! - [`StubLlm`]: A configurable LLM provider that returns a fixed response +//! - [`TestHarnessBuilder`]: Builder for wiring `AgentDeps` with defaults +//! - [`TestHarness`]: The assembled components ready for use in tests +//! +//! # Usage +//! +//! ```rust,no_run +//! use ironclaw::testing::TestHarnessBuilder; +//! +//! #[tokio::test] +//! async fn test_something() { +//! let harness = TestHarnessBuilder::new().build().await; +//! // use harness.deps, harness.db, etc. +//! } +//! ``` + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; + +use async_trait::async_trait; +use rust_decimal::Decimal; + +use crate::agent::AgentDeps; +use crate::db::Database; +use crate::error::LlmError; +use crate::llm::{ + CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ToolCompletionRequest, + ToolCompletionResponse, +}; +use crate::tools::ToolRegistry; + +/// Create a libSQL-backed test database in a temporary directory. +/// +/// Returns the database and a `TempDir` guard — the database file is +/// deleted when the guard is dropped. +#[cfg(feature = "libsql")] +pub async fn test_db() -> (Arc, tempfile::TempDir) { + use crate::db::libsql::LibSqlBackend; + + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let path = dir.path().join("test.db"); + let backend = LibSqlBackend::new_local(&path) + .await + .expect("failed to create test LibSqlBackend"); + backend + .run_migrations() + .await + .expect("failed to run migrations"); + (Arc::new(backend) as Arc, dir) +} + +/// What kind of error the stub should produce when failing. +#[derive(Clone, Copy, Debug)] +pub enum StubErrorKind { + /// Transient/retryable error (`LlmError::RequestFailed`). + Transient, + /// Non-transient error (`LlmError::ContextLengthExceeded`). + NonTransient, +} + +/// A configurable LLM provider stub for tests. +/// +/// Supports: +/// - Fixed response content +/// - Call counting via [`calls()`](Self::calls) +/// - Runtime failure toggling via [`set_failing()`](Self::set_failing) +/// - Configurable error kinds (transient vs non-transient) +/// +/// Use this in tests instead of creating ad-hoc stub implementations. +pub struct StubLlm { + model_name: String, + response: String, + call_count: AtomicU32, + should_fail: AtomicBool, + error_kind: StubErrorKind, +} + +impl StubLlm { + /// Create a new stub that returns the given response. + pub fn new(response: impl Into) -> Self { + Self { + model_name: "stub-model".to_string(), + response: response.into(), + call_count: AtomicU32::new(0), + should_fail: AtomicBool::new(false), + error_kind: StubErrorKind::Transient, + } + } + + /// Create a stub that always fails with a transient error. + pub fn failing(name: impl Into) -> Self { + Self { + model_name: name.into(), + response: String::new(), + call_count: AtomicU32::new(0), + should_fail: AtomicBool::new(true), + error_kind: StubErrorKind::Transient, + } + } + + /// Create a stub that always fails with a non-transient error. + pub fn failing_non_transient(name: impl Into) -> Self { + Self { + model_name: name.into(), + response: String::new(), + call_count: AtomicU32::new(0), + should_fail: AtomicBool::new(true), + error_kind: StubErrorKind::NonTransient, + } + } + + /// Set the model name. + pub fn with_model_name(mut self, name: impl Into) -> Self { + self.model_name = name.into(); + self + } + + /// Get the number of times `complete` or `complete_with_tools` was called. + pub fn calls(&self) -> u32 { + self.call_count.load(Ordering::Relaxed) + } + + /// Toggle whether calls should fail at runtime. + pub fn set_failing(&self, fail: bool) { + self.should_fail.store(fail, Ordering::Relaxed); + } + + fn make_error(&self) -> LlmError { + match self.error_kind { + StubErrorKind::Transient => LlmError::RequestFailed { + provider: self.model_name.clone(), + reason: "server error".to_string(), + }, + StubErrorKind::NonTransient => LlmError::ContextLengthExceeded { + used: 100_000, + limit: 50_000, + }, + } + } +} + +impl Default for StubLlm { + fn default() -> Self { + Self::new("OK") + } +} + +#[async_trait] +impl LlmProvider for StubLlm { + fn model_name(&self) -> &str { + &self.model_name + } + + fn cost_per_token(&self) -> (Decimal, Decimal) { + (Decimal::ZERO, Decimal::ZERO) + } + + async fn complete(&self, _request: CompletionRequest) -> Result { + self.call_count.fetch_add(1, Ordering::Relaxed); + if self.should_fail.load(Ordering::Relaxed) { + return Err(self.make_error()); + } + Ok(CompletionResponse { + content: self.response.clone(), + input_tokens: 10, + output_tokens: 5, + finish_reason: FinishReason::Stop, + response_id: None, + }) + } + + async fn complete_with_tools( + &self, + _request: ToolCompletionRequest, + ) -> Result { + self.call_count.fetch_add(1, Ordering::Relaxed); + if self.should_fail.load(Ordering::Relaxed) { + return Err(self.make_error()); + } + Ok(ToolCompletionResponse { + content: Some(self.response.clone()), + tool_calls: Vec::new(), + input_tokens: 10, + output_tokens: 5, + finish_reason: FinishReason::Stop, + response_id: None, + }) + } +} + +/// Assembled test components. +pub struct TestHarness { + /// The agent dependencies, ready for use. + pub deps: AgentDeps, + /// Direct reference to the database (as `Arc`). + pub db: Arc, + /// Temp directory guard — keeps the test database alive. Dropped + /// automatically when the harness goes out of scope. + #[cfg(feature = "libsql")] + _temp_dir: tempfile::TempDir, +} + +/// Builder for constructing a [`TestHarness`] with sensible defaults. +/// +/// All defaults are designed to work without any external services: +/// - Database: libSQL in a temp directory (real SQL, FTS5, no network) +/// - LLM: `StubLlm` returning "OK" +/// - Safety: permissive config +/// - Tools: builtin tools registered +/// - Hooks: empty registry +/// - Cost guard: no limits +pub struct TestHarnessBuilder { + db: Option>, + llm: Option>, + tools: Option>, +} + +impl TestHarnessBuilder { + /// Create a new builder with all defaults. + pub fn new() -> Self { + Self { + db: None, + llm: None, + tools: None, + } + } + + /// Override the database backend. + pub fn with_db(mut self, db: Arc) -> Self { + self.db = Some(db); + self + } + + /// Override the LLM provider. + pub fn with_llm(mut self, llm: Arc) -> Self { + self.llm = Some(llm); + self + } + + /// Override the tool registry. + pub fn with_tools(mut self, tools: Arc) -> Self { + self.tools = Some(tools); + self + } + + /// Build the harness with defaults applied. + #[cfg(feature = "libsql")] + pub async fn build(self) -> TestHarness { + use crate::agent::cost_guard::{CostGuard, CostGuardConfig}; + use crate::config::{SafetyConfig, SkillsConfig}; + use crate::hooks::HookRegistry; + use crate::safety::SafetyLayer; + + let (db, temp_dir) = if let Some(db) = self.db { + // Caller provided a DB; create a dummy temp dir to satisfy the struct. + let dir = tempfile::tempdir().expect("failed to create temp dir"); + (db, dir) + } else { + test_db().await + }; + + let llm: Arc = self.llm.unwrap_or_else(|| Arc::new(StubLlm::default())); + + let tools = self.tools.unwrap_or_else(|| { + let t = Arc::new(ToolRegistry::new()); + t.register_builtin_tools(); + t + }); + + let safety = Arc::new(SafetyLayer::new(&SafetyConfig { + max_output_length: 100_000, + injection_check_enabled: false, + })); + + let hooks = Arc::new(HookRegistry::new()); + + let cost_guard = Arc::new(CostGuard::new(CostGuardConfig { + max_cost_per_day_cents: None, + max_actions_per_hour: None, + })); + + let deps = AgentDeps { + store: Some(Arc::clone(&db)), + llm, + cheap_llm: None, + safety, + tools, + workspace: None, + extension_manager: None, + skill_registry: None, + skills_config: SkillsConfig::default(), + hooks, + cost_guard, + }; + + TestHarness { + deps, + db, + _temp_dir: temp_dir, + } + } +} + +impl Default for TestHarnessBuilder { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(feature = "libsql")] + #[tokio::test] + async fn test_harness_builds_with_defaults() { + let harness = TestHarnessBuilder::new().build().await; + assert!(harness.deps.store.is_some()); + assert_eq!(harness.deps.llm.model_name(), "stub-model"); + } + + #[cfg(feature = "libsql")] + #[tokio::test] + async fn test_harness_custom_llm() { + let custom_llm = Arc::new(StubLlm::new("custom response").with_model_name("my-model")); + let harness = TestHarnessBuilder::new().with_llm(custom_llm).build().await; + assert_eq!(harness.deps.llm.model_name(), "my-model"); + } + + #[cfg(feature = "libsql")] + #[tokio::test] + async fn test_harness_db_works() { + let harness = TestHarnessBuilder::new().build().await; + + let id = harness + .db + .create_conversation("test", "user1", None) + .await + .expect("create conversation"); + assert!(!id.is_nil()); + } + + #[tokio::test] + async fn test_stub_llm_complete() { + let llm = StubLlm::new("hello world"); + let response = llm + .complete(CompletionRequest::new(vec![])) + .await + .expect("complete"); + assert_eq!(response.content, "hello world"); + assert_eq!(response.finish_reason, FinishReason::Stop); + } +} diff --git a/examples/test_heartbeat.rs b/tests/heartbeat_integration.rs similarity index 84% rename from examples/test_heartbeat.rs rename to tests/heartbeat_integration.rs index fcb9333d..a4c07357 100644 --- a/examples/test_heartbeat.rs +++ b/tests/heartbeat_integration.rs @@ -1,11 +1,12 @@ -//! Standalone heartbeat test. +#![cfg(feature = "postgres")] +//! Heartbeat integration test. //! //! Exercises the heartbeat system in isolation: connects to the real //! database, reads the real HEARTBEAT.md, calls the real LLM, and prints //! every step so you can see exactly where it breaks. //! //! Usage: -//! cargo run --example test_heartbeat +//! cargo test --test heartbeat_integration -- --ignored --nocapture use std::sync::Arc; @@ -17,20 +18,19 @@ use ironclaw::{ workspace::Workspace, }; -#[tokio::main] -async fn main() -> anyhow::Result<()> { +#[tokio::test] +#[ignore] // Requires running database and LLM credentials +async fn test_heartbeat_end_to_end() { // Load .env and set up logging let _ = dotenvy::dotenv(); - tracing_subscriber::fmt() + let _ = tracing_subscriber::fmt() .with_env_filter("ironclaw=debug") - .init(); + .try_init(); println!("=== Heartbeat Integration Test ===\n"); // 1. Load config - let config = Config::from_env() - .await - .map_err(|e| anyhow::anyhow!("Config: {}", e))?; + let config = Config::from_env().await.expect("Failed to load config"); println!("[1/6] Config loaded"); println!(" heartbeat.enabled = {}", config.heartbeat.enabled); println!( @@ -47,8 +47,13 @@ async fn main() -> anyhow::Result<()> { ); // 2. Connect to database - let store = Store::new(&config.database).await?; - store.run_migrations().await?; + let store = Store::new(&config.database) + .await + .expect("Failed to connect to database"); + store + .run_migrations() + .await + .expect("Failed to run migrations"); println!("[2/6] Database connected"); // 3. Create workspace @@ -83,7 +88,7 @@ async fn main() -> anyhow::Result<()> { session_path: config.llm.nearai.session_path.clone(), }) .await; - let llm = create_llm_provider(&config.llm, session)?; + let llm = create_llm_provider(&config.llm, session).expect("Failed to create LLM provider"); println!("[5/6] LLM provider created (model: {})", llm.model_name()); // 6. Run heartbeat check @@ -116,6 +121,4 @@ async fn main() -> anyhow::Result<()> { println!(" Error: {}", err); } } - - Ok(()) }