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 <[email protected]>

* 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 <[email protected]>

* style: fix rustfmt formatting for CI

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* 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 <[email protected]>

---------

Co-authored-by: Illia Polosukhin <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-18 23:05:47 +00:00
committed by GitHub
co-authored by Illia Polosukhin Claude Opus 4.6
parent 6330f1b27a
commit ffb1cc9be8
50 changed files with 8616 additions and 5184 deletions
-4
View File
@@ -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"
+56
View File
@@ -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"
+779
View File
@@ -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<Arc<dyn Database>>,
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
pub llm: Arc<dyn LlmProvider>,
pub cheap_llm: Option<Arc<dyn LlmProvider>>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub embeddings: Option<Arc<dyn EmbeddingProvider>>,
pub workspace: Option<Arc<Workspace>>,
pub extension_manager: Option<Arc<ExtensionManager>>,
pub mcp_session_manager: Arc<McpSessionManager>,
pub wasm_tool_runtime: Option<Arc<WasmToolRuntime>>,
pub log_broadcaster: Arc<LogBroadcaster>,
pub context_manager: Arc<ContextManager>,
pub hooks: Arc<HookRegistry>,
pub skill_registry: Option<Arc<std::sync::RwLock<SkillRegistry>>>,
pub skill_catalog: Option<Arc<SkillCatalog>>,
pub cost_guard: Arc<crate::agent::cost_guard::CostGuard>,
pub session: Arc<SessionManager>,
}
/// Options that control optional init phases.
#[derive(Default)]
pub struct AppBuilderFlags {
pub no_db: bool,
}
/// Builder that orchestrates the 5 mechanical init phases.
pub struct AppBuilder {
config: Config,
flags: AppBuilderFlags,
toml_path: Option<std::path::PathBuf>,
session: Arc<SessionManager>,
log_broadcaster: Arc<LogBroadcaster>,
// Accumulated state
db: Option<Arc<dyn Database>>,
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
// Backend-specific handles needed by secrets store
#[cfg(feature = "postgres")]
pg_pool: Option<deadpool_postgres::Pool>,
#[cfg(feature = "libsql")]
libsql_db: Option<Arc<libsql::Database>>,
}
impl AppBuilder {
/// Create a new builder.
///
/// The `session` and `log_broadcaster` are created before the builder
/// because tracing must be initialized before any init phase runs,
/// and the log broadcaster is part of the tracing layer.
pub fn new(
config: Config,
flags: AppBuilderFlags,
toml_path: Option<std::path::PathBuf>,
session: Arc<SessionManager>,
log_broadcaster: Arc<LogBroadcaster>,
) -> Self {
Self {
config,
flags,
toml_path,
session,
log_broadcaster,
db: None,
secrets_store: None,
#[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<dyn Database> = 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<dyn Database>
}
#[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<dyn Database>
}
#[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<Arc<dyn SecretsStore + Send + Sync>> = 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<dyn SecretsStore + Send + Sync>
})
});
#[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<dyn SecretsStore + Send + Sync>
})
});
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<dyn LlmProvider>, Option<Arc<dyn LlmProvider>>), 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<dyn LlmProvider> = 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<dyn LlmProvider> =
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<dyn LlmProvider> = 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<dyn LlmProvider>,
) -> Result<
(
Arc<SafetyLayer>,
Arc<ToolRegistry>,
Option<Arc<dyn EmbeddingProvider>>,
Option<Arc<Workspace>>,
),
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<Arc<dyn EmbeddingProvider>> = 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<ToolRegistry>,
) -> Result<
(
Arc<McpSessionManager>,
Option<Arc<WasmToolRuntime>>,
Option<Arc<ExtensionManager>>,
),
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<Arc<WasmToolRuntime>> =
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<AppComponents, anyhow::Error> {
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(&registry), Arc::clone(&catalog));
(Some(registry), Some(catalog))
} else {
(None, None)
};
let context_manager = Arc::new(ContextManager::new(self.config.agent.max_parallel_jobs));
let 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,
})
}
}
+633
View File
@@ -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<Arc<GatewayState>>,
Json(req): Json<SendMessageRequest>,
) -> Result<(StatusCode, Json<SendMessageResponse>), (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<Arc<GatewayState>>,
Json(req): Json<ApprovalRequest>,
) -> Result<(StatusCode, Json<SendMessageResponse>), (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<Arc<GatewayState>>,
Json(req): Json<AuthTokenRequest>,
) -> Result<Json<ActionResponse>, (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<Arc<GatewayState>>,
Json(_req): Json<AuthCancelRequest>,
) -> Result<Json<ActionResponse>, (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<Arc<GatewayState>>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
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<Arc<GatewayState>>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
// 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<String>,
pub limit: Option<usize>,
pub before: Option<String>,
}
pub async fn chat_history_handler(
State(state): State<Arc<GatewayState>>,
Query(query): Query<HistoryQuery>,
) -> Result<Json<HistoryResponse>, (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<TurnInfo> = 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<TurnInfo> {
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<Arc<GatewayState>>,
) -> Result<Json<ThreadListResponse>, (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<ThreadInfo> = 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<Arc<GatewayState>>,
) -> Result<Json<ThreadInfo>, (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");
}
}
+153
View File
@@ -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<Arc<GatewayState>>,
) -> Result<Json<ExtensionListResponse>, (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<Arc<GatewayState>>,
) -> Result<Json<ToolListResponse>, (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<Arc<GatewayState>>,
Json(req): Json<InstallExtensionRequest>,
) -> Result<Json<ActionResponse>, (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<Arc<GatewayState>>,
Path(name): Path<String>,
) -> Result<Json<ActionResponse>, (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<Arc<GatewayState>>,
Path(name): Path<String>,
) -> Result<Json<ActionResponse>, (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()))),
}
}
+518
View File
@@ -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<Arc<GatewayState>>,
) -> Result<Json<JobListResponse>, (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<JobInfo> = 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<Arc<GatewayState>>,
) -> Result<Json<JobSummaryResponse>, (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<Arc<GatewayState>>,
Path(id): Path<String>,
) -> Result<Json<JobDetailResponse>, (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<Arc<GatewayState>>,
Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, (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<Arc<GatewayState>>,
Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, (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<crate::orchestrator::auth::CredentialGrant> =
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<Arc<GatewayState>>,
Path(id): Path<String>,
Json(body): Json<serde_json::Value>,
) -> Result<Json<serde_json::Value>, (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<Arc<GatewayState>>,
Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, (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<serde_json::Value> = 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<String>,
}
pub async fn job_files_list_handler(
State(state): State<Arc<GatewayState>>,
Path(id): Path<String>,
Query(query): Query<FilePathQuery>,
) -> Result<Json<ProjectFilesResponse>, (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<Arc<GatewayState>>,
Path(id): Path<String>,
Query(query): Query<FilePathQuery>,
) -> Result<Json<ProjectFileReadResponse>, (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,
}))
}
+171
View File
@@ -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<usize>,
}
pub async fn memory_tree_handler(
State(state): State<Arc<GatewayState>>,
Query(_query): Query<TreeQuery>,
) -> Result<Json<MemoryTreeResponse>, (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<TreeEntry> = Vec::new();
let mut seen_dirs: std::collections::HashSet<String> = 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<String>,
}
pub async fn memory_list_handler(
State(state): State<Arc<GatewayState>>,
Query(query): Query<ListQuery>,
) -> Result<Json<MemoryListResponse>, (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<ListEntry> = 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<Arc<GatewayState>>,
Query(query): Query<ReadQuery>,
) -> Result<Json<MemoryReadResponse>, (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<Arc<GatewayState>>,
Json(req): Json<MemoryWriteRequest>,
) -> Result<Json<MemoryWriteResponse>, (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<Arc<GatewayState>>,
Json(req): Json<MemorySearchRequest>,
) -> Result<Json<MemorySearchResponse>, (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<SearchHit> = 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 }))
}
+23
View File
@@ -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::*;
+330
View File
@@ -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<Arc<GatewayState>>,
) -> Result<Json<RoutineListResponse>, (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<RoutineInfo> = routines.iter().map(routine_to_info).collect();
Ok(Json(RoutineListResponse { routines: items }))
}
pub async fn routines_summary_handler(
State(state): State<Arc<GatewayState>>,
) -> Result<Json<RoutineSummaryResponse>, (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<Arc<GatewayState>>,
Path(id): Path<String>,
) -> Result<Json<RoutineDetailResponse>, (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<RoutineRunInfo> = 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<Arc<GatewayState>>,
Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, (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<bool>,
}
pub async fn routines_toggle_handler(
State(state): State<Arc<GatewayState>>,
Path(id): Path<String>,
body: Option<Json<ToggleRequest>>,
) -> Result<Json<serde_json::Value>, (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<Arc<GatewayState>>,
Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, (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<Arc<GatewayState>>,
Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, (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<RoutineRunInfo> = 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(),
}
}
+133
View File
@@ -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<Arc<GatewayState>>,
) -> Result<Json<SettingsListResponse>, 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<Arc<GatewayState>>,
Path(key): Path<String>,
) -> Result<Json<SettingResponse>, 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<Arc<GatewayState>>,
Path(key): Path<String>,
Json(body): Json<SettingWriteRequest>,
) -> Result<StatusCode, StatusCode> {
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<Arc<GatewayState>>,
Path(key): Path<String>,
) -> Result<StatusCode, StatusCode> {
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<Arc<GatewayState>>,
) -> Result<Json<SettingsExportResponse>, 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<Arc<GatewayState>>,
Json(body): Json<SettingsImportRequest>,
) -> Result<StatusCode, StatusCode> {
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)
}
+257
View File
@@ -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<Arc<GatewayState>>,
) -> Result<Json<SkillListResponse>, (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<SkillInfo> = 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<Arc<GatewayState>>,
Json(req): Json<SkillSearchRequest>,
) -> Result<Json<SkillSearchResponse>, (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<serde_json::Value> = 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<SkillInfo> = {
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<Arc<GatewayState>>,
headers: axum::http::HeaderMap,
Json(req): Json<SkillInstallRequest>,
) -> Result<Json<ActionResponse>, (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<Arc<GatewayState>>,
headers: axum::http::HeaderMap,
Path(name): Path<String>,
) -> Result<Json<ActionResponse>, (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()))),
}
}
+178
View File
@@ -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<HealthResponse> {
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<String>) -> 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<String>) -> 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<Arc<GatewayState>>,
) -> Result<
Sse<impl futures::Stream<Item = Result<Event, Infallible>> + 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<Arc<GatewayState>>,
) -> Json<GatewayStatusResponse> {
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,
}
+1 -1
View File
@@ -519,7 +519,7 @@ async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Syn
#[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();
+1 -1
View File
@@ -737,7 +737,7 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, 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();
-1944
View File
File diff suppressed because it is too large Load Diff
+120
View File
@@ -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<u64>,
/// Maximum LLM/tool actions per hour. None = unlimited.
pub max_actions_per_hour: Option<u64>,
}
impl AgentConfig {
pub(crate) fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
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}"),
})?,
})
}
}
+72
View File
@@ -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<PathBuf>,
/// 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<Self, ConfigError> {
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,
}
}
}
+126
View File
@@ -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<HttpConfig>,
pub gateway: Option<GatewayConfig>,
/// 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<i64>,
}
#[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<SecretString>,
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<String>,
pub user_id: String,
}
impl ChannelsConfig {
pub(crate) fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
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")
}
+130
View File
@@ -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<Self, Self::Err> {
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<PathBuf>,
/// Turso cloud URL for remote sync (optional).
pub libsql_url: Option<String>,
/// Turso auth token (required when libsql_url is set).
pub libsql_auth_token: Option<SecretString>,
}
impl DatabaseConfig {
pub(crate) fn resolve() -> Result<Self, ConfigError> {
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")
}
+165
View File
@@ -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<SecretString>,
/// 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<Self, ConfigError> {
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");
}
}
}
+54
View File
@@ -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<String>,
/// User ID to notify on heartbeat findings.
pub notify_user: Option<String>,
}
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<Self, ConfigError> {
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()),
})
}
}
+40
View File
@@ -0,0 +1,40 @@
use crate::error::ConfigError;
use super::INJECTED_VARS;
pub(crate) fn optional_env(key: &str) -> Result<Option<String>, 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<T>(key: &str, default: T) -> Result<T, ConfigError>
where
T: std::str::FromStr,
T::Err: std::fmt::Display,
{
optional_env(key)?
.map(|s| {
s.parse().map_err(|e| ConfigError::InvalidValue {
key: key.to_string(),
message: format!("{e}"),
})
})
.transpose()
.map(|opt| opt.unwrap_or(default))
}
+426
View File
@@ -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<Self, Self::Err> {
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<SecretString>,
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<OpenAiDirectConfig>,
/// Direct Anthropic config (populated when backend=anthropic)
pub anthropic: Option<AnthropicDirectConfig>,
/// Ollama config (populated when backend=ollama)
pub ollama: Option<OllamaConfig>,
/// OpenAI-compatible config (populated when backend=openai_compatible)
pub openai_compatible: Option<OpenAiCompatibleConfig>,
/// Tinfoil config (populated when backend=tinfoil)
pub tinfoil: Option<TinfoilConfig>,
}
/// 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<Self, Self::Err> {
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<String>,
/// 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<SecretString>,
/// 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<String>,
/// 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<u32>,
/// 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<Self, ConfigError> {
// 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");
}
}
}
+239
View File
@@ -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<HashMap<String, String>> = 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, ConfigError> {
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<Self, ConfigError> {
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, ConfigError> {
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<Self, ConfigError> {
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<Self, ConfigError> {
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);
}
+48
View File
@@ -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<Self, ConfigError> {
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)?,
})
}
}
+25
View File
@@ -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<Self, ConfigError> {
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),
})
}
}
+261
View File
@@ -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<String>,
}
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<Self, ConfigError> {
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<String>,
}
/// 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<String> {
[
// 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<String> {
// 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<Self, ConfigError> {
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<String> {
let creds: serde_json::Value = serde_json::from_str(json).ok()?;
creds["claudeAiOauth"]["accessToken"]
.as_str()
.map(String::from)
}
+70
View File
@@ -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<SecretString>,
/// 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<Self, ConfigError> {
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()
}
}
+56
View File
@@ -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<Self, ConfigError> {
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)?,
})
}
}
+106
View File
@@ -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<String>,
/// Provider configuration for lifecycle-managed tunnels.
/// `None` when using a static URL or no tunnel at all.
pub provider: Option<crate::tunnel::TunnelProviderConfig>,
}
impl TunnelConfig {
pub(crate) fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
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<String> {
self.public_url.as_ref().map(|base| {
let base = base.trim_end_matches('/');
let path = path.trim_start_matches('/');
format!("{}/{}", base, path)
})
}
}
+99
View File
@@ -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<PathBuf>,
}
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<Self, ConfigError> {
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,
}
}
}
+354
View File
@@ -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<Uuid, DatabaseError> {
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<Uuid, DatabaseError> {
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<Vec<ConversationSummary>, 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::<String>(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<Uuid, DatabaseError> {
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<Uuid, DatabaseError> {
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<DateTime<Utc>>,
limit: i64,
) -> Result<(Vec<ConversationMessage>, 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<Option<serde_json::Value>, 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<Vec<ConversationMessage>, 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<bool, DatabaseError> {
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())
}
}
+330
View File
@@ -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<Option<JobContext>, 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<i64> = row.get::<i64>(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<Vec<Uuid>, 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::<String>(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<Vec<ActionRecord>, 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<String> =
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<Uuid, DatabaseError> {
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<Uuid, DatabaseError> {
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<Decimal>,
) -> 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(())
}
}
+460
View File
@@ -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<LibSqlDatabase>,
}
impl LibSqlBackend {
/// Create a new local embedded database.
pub async fn new_local(path: &Path) -> Result<Self, DatabaseError> {
// 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<Self, DatabaseError> {
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<Self, DatabaseError> {
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<LibSqlDatabase> {
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<Connection, DatabaseError> {
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<Utc>.
///
/// 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<DateTime<Utc>, 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<Utc> for SQLite storage (RFC 3339 with millisecond precision).
pub(crate) fn fmt_ts(dt: &DateTime<Utc>) -> String {
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
}
/// Format an optional DateTime<Utc>.
pub(crate) fn fmt_opt_ts(dt: &Option<DateTime<Utc>>) -> 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::<String>(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<String> {
row.get::<String>(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<String>` to a `libsql::Value` (Text or Null).
pub(crate) fn opt_text_owned(s: Option<String>) -> 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::<i64>(idx).unwrap_or(0)
}
/// Extract an optional bool from an integer column.
pub(crate) fn get_opt_bool(row: &libsql::Row, idx: i32) -> Option<bool> {
row.get::<i64>(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::<String>(idx)
.ok()
.and_then(|s| s.parse::<Decimal>().ok())
.unwrap_or_default()
}
/// Parse an optional Decimal from a text column.
pub(crate) fn get_opt_decimal(row: &libsql::Row, idx: i32) -> Option<Decimal> {
row.get::<String>(idx)
.ok()
.and_then(|s| s.parse::<Decimal>().ok())
}
/// Parse a JSON value from a text column.
pub(crate) fn get_json(row: &libsql::Row, idx: i32) -> serde_json::Value {
row.get::<String>(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<Utc> {
match row.get::<String>(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<DateTime<Utc>> {
match row.get::<String>(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<Routine, DatabaseError> {
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<i64> = row.get::<i64>(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<RoutineRun, DatabaseError> {
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::<i64>(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);
}
}
+390
View File
@@ -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<Option<Routine>, 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<Option<Routine>, 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<Vec<Routine>, 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<Vec<Routine>, 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<Vec<Routine>, 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<Utc>,
next_fire_at: Option<DateTime<Utc>>,
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<bool, DatabaseError> {
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<i32>,
) -> 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<Vec<RoutineRun>, 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<i64, DatabaseError> {
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),
}
}
}
+405
View File
@@ -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<Option<SandboxJobRecord>, 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<Vec<SandboxJobRecord>, 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<bool>,
message: Option<&str>,
started_at: Option<DateTime<Utc>>,
completed_at: Option<DateTime<Utc>>,
) -> 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<u64, DatabaseError> {
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<SandboxJobSummary, DatabaseError> {
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<Vec<SandboxJobRecord>, 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<SandboxJobSummary, DatabaseError> {
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<bool, DatabaseError> {
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<Option<String>, 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<i64>,
) -> Result<Vec<JobEventRecord>, 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)
}
}
+208
View File
@@ -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<Option<serde_json::Value>, 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<Option<SettingRow>, 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<bool, DatabaseError> {
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<Vec<SettingRow>, 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<HashMap<String, serde_json::Value>, 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<String, serde_json::Value>,
) -> 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<bool, DatabaseError> {
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),
}
}
}
+97
View File
@@ -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<Vec<BrokenTool>, 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(())
}
}
+607
View File
@@ -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<Uuid>,
path: &str,
) -> Result<MemoryDocument, 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 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<MemoryDocument, WorkspaceError> {
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<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
// 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<Uuid>,
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<Uuid>,
directory: &str,
) -> Result<Vec<WorkspaceEntry>, 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<String, WorkspaceEntry> = 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<WorkspaceEntry> = 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<Uuid>,
) -> Result<Vec<String>, 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<Uuid>,
) -> Result<Vec<MemoryDocument>, 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<Uuid, WorkspaceError> {
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<u8> = 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<u8> = 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<Uuid>,
limit: usize,
) -> Result<Vec<MemoryChunk>, 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<Uuid>,
query: &str,
embedding: Option<&[f32]>,
config: &SearchConfig,
) -> Result<Vec<SearchResult>, 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::<Vec<_>>()
.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))
}
}
File diff suppressed because it is too large Load Diff
+50 -184
View File
@@ -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<dyn Database>` 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<Uuid, DatabaseError>;
/// 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<Uuid, DatabaseError>;
/// 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<Vec<ConversationSummary>, DatabaseError>;
/// Get or create the singleton assistant conversation.
async fn get_or_create_assistant_conversation(
&self,
user_id: &str,
channel: &str,
) -> Result<Uuid, DatabaseError>;
/// Create a conversation with specific metadata.
async fn create_conversation_with_metadata(
&self,
channel: &str,
user_id: &str,
metadata: &serde_json::Value,
) -> Result<Uuid, DatabaseError>;
/// Load messages with cursor-based pagination.
async fn list_conversation_messages_paginated(
&self,
conversation_id: Uuid,
before: Option<DateTime<Utc>>,
limit: i64,
) -> Result<(Vec<ConversationMessage>, 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<Option<serde_json::Value>, DatabaseError>;
/// Load all messages for a conversation.
async fn list_conversation_messages(
&self,
conversation_id: Uuid,
) -> Result<Vec<ConversationMessage>, DatabaseError>;
/// Check if a conversation belongs to a specific user.
async fn conversation_belongs_to_user(
&self,
conversation_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError>;
}
// ==================== 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<Option<JobContext>, 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<Vec<Uuid>, 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<Vec<ActionRecord>, DatabaseError>;
// ==================== LLM Calls ====================
/// Record an LLM call.
async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError>;
// ==================== 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<Uuid, DatabaseError>;
/// 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<Decimal>,
) -> 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<Option<SandboxJobRecord>, DatabaseError>;
/// List all sandbox jobs, most recent first.
async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, 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<DateTime<Utc>>,
completed_at: Option<DateTime<Utc>>,
) -> Result<(), DatabaseError>;
/// Mark stale sandbox jobs as interrupted.
async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError>;
/// Get sandbox job summary.
async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError>;
/// List sandbox jobs for a specific user, most recent first.
async fn list_sandbox_jobs_for_user(
&self,
user_id: &str,
) -> Result<Vec<SandboxJobRecord>, DatabaseError>;
/// Get sandbox job summary for a specific user.
async fn sandbox_job_summary_for_user(
&self,
user_id: &str,
) -> Result<SandboxJobSummary, DatabaseError>;
/// 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<bool, DatabaseError>;
/// 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<Option<String>, 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<i64>,
) -> Result<Vec<JobEventRecord>, 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<Option<Routine>, DatabaseError>;
/// Get a routine by user_id and name.
async fn get_routine_by_name(
&self,
user_id: &str,
name: &str,
) -> Result<Option<Routine>, DatabaseError>;
/// List routines for a user.
async fn list_routines(&self, user_id: &str) -> Result<Vec<Routine>, DatabaseError>;
/// List all enabled event routines.
async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError>;
/// List due cron routines.
async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, 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<bool, DatabaseError>;
// ==================== 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<i32>,
) -> Result<(), DatabaseError>;
/// List recent runs for a routine.
async fn list_routine_runs(
&self,
routine_id: Uuid,
limit: i64,
) -> Result<Vec<RoutineRun>, DatabaseError>;
/// Count currently running runs for a routine.
async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError>;
}
// ==================== 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<Vec<BrokenTool>, 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<Option<serde_json::Value>, DatabaseError>;
/// Get a single setting with metadata.
async fn get_setting_full(
&self,
user_id: &str,
key: &str,
) -> Result<Option<SettingRow>, 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<bool, DatabaseError>;
/// List all settings for a user.
async fn list_settings(&self, user_id: &str) -> Result<Vec<SettingRow>, DatabaseError>;
/// Get all settings as a flat map.
async fn get_all_settings(
&self,
user_id: &str,
) -> Result<HashMap<String, serde_json::Value>, DatabaseError>;
/// Bulk-write settings atomically.
async fn set_all_settings(
&self,
user_id: &str,
settings: &HashMap<String, serde_json::Value>,
) -> Result<(), DatabaseError>;
/// Check if settings exist for a user.
async fn has_settings(&self, user_id: &str) -> Result<bool, DatabaseError>;
}
// ==================== 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<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError>;
/// Get a document by ID.
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError>;
/// Get or create a document by path.
async fn get_or_create_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError>;
/// 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<Uuid>,
path: &str,
) -> Result<(), WorkspaceError>;
/// List files and directories in a directory path.
async fn list_directory(
&self,
user_id: &str,
agent_id: Option<Uuid>,
directory: &str,
) -> Result<Vec<WorkspaceEntry>, WorkspaceError>;
/// List all file paths in the workspace.
async fn list_all_paths(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<String>, WorkspaceError>;
/// List all documents for a user.
async fn list_documents(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<MemoryDocument>, 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<Uuid, WorkspaceError>;
/// 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<Uuid>,
limit: usize,
) -> Result<Vec<MemoryChunk>, 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<Vec<SearchResult>, WorkspaceError>;
}
/// Backend-agnostic database supertrait.
///
/// Combines all sub-traits into one. Existing `Arc<dyn Database>` 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>;
}
+34 -22
View File
@@ -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<Uuid, DatabaseError> {
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<Vec<JobEventRecord>, 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<i64, DatabaseError> {
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<bool, DatabaseError> {
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,
+4
View File
@@ -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};
+12 -126
View File
@@ -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<Self> {
Arc::new(Self {
name: name.to_string(),
should_fail: AtomicBool::new(false),
error_kind: StubError::Transient,
})
}
fn always_fail(name: &str) -> Arc<Self> {
Arc::new(Self {
name: name.to_string(),
should_fail: AtomicBool::new(true),
error_kind: StubError::Transient,
})
}
fn always_fail_non_transient(name: &str) -> Arc<Self> {
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<CompletionResponse, LlmError> {
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<ToolCompletionResponse, LlmError> {
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");
+12 -78
View File
@@ -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<CompletionResponse, LlmError> {
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<ToolCompletionResponse, LlmError> {
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");
}
+1 -1
View File
@@ -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();
+6 -37
View File
@@ -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<CompletionResponse, LlmError> {
Err(LlmError::RequestFailed {
provider: "stub".into(),
reason: "not implemented".into(),
})
}
async fn complete_with_tools(
&self,
_req: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
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),
+3 -3
View File
@@ -79,7 +79,7 @@ pub struct SetupWizard {
db_pool: Option<deadpool_postgres::Pool>,
/// libSQL backend (created during setup, libsql only).
#[cfg(feature = "libsql")]
db_backend: Option<crate::db::libsql_backend::LibSqlBackend>,
db_backend: Option<crate::db::libsql::LibSqlBackend>,
/// Secrets crypto (created during setup).
secrets_crypto: Option<Arc<SecretsCrypto>>,
/// 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
+356
View File
@@ -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<dyn Database>, 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<dyn Database>, 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<String>) -> 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<String>) -> 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<String>) -> 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<String>) -> 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<CompletionResponse, LlmError> {
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<ToolCompletionResponse, LlmError> {
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<dyn Database>`).
pub db: Arc<dyn Database>,
/// 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<Arc<dyn Database>>,
llm: Option<Arc<dyn LlmProvider>>,
tools: Option<Arc<ToolRegistry>>,
}
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<dyn Database>) -> Self {
self.db = Some(db);
self
}
/// Override the LLM provider.
pub fn with_llm(mut self, llm: Arc<dyn LlmProvider>) -> Self {
self.llm = Some(llm);
self
}
/// Override the tool registry.
pub fn with_tools(mut self, tools: Arc<ToolRegistry>) -> 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<dyn LlmProvider> = 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);
}
}
@@ -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(())
}