mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
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:
co-authored by
Illia Polosukhin
Claude Opus 4.6
parent
6330f1b27a
commit
ffb1cc9be8
@@ -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"
|
||||
|
||||
Executable
+56
@@ -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
@@ -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(®istry), Arc::clone(&catalog));
|
||||
(Some(registry), Some(catalog))
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
|
||||
let context_manager = Arc::new(ContextManager::new(self.config.agent.max_parallel_jobs));
|
||||
let hooks = Arc::new(HookRegistry::new());
|
||||
let cost_guard = Arc::new(crate::agent::cost_guard::CostGuard::new(
|
||||
crate::agent::cost_guard::CostGuardConfig {
|
||||
max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents,
|
||||
max_actions_per_hour: self.config.agent.max_actions_per_hour,
|
||||
},
|
||||
));
|
||||
|
||||
tracing::info!(
|
||||
"Tool registry initialized with {} total tools",
|
||||
tools.count()
|
||||
);
|
||||
|
||||
Ok(AppComponents {
|
||||
config: self.config,
|
||||
db: self.db,
|
||||
secrets_store: self.secrets_store,
|
||||
llm,
|
||||
cheap_llm,
|
||||
safety,
|
||||
tools,
|
||||
embeddings,
|
||||
workspace,
|
||||
extension_manager,
|
||||
mcp_session_manager,
|
||||
wasm_tool_runtime,
|
||||
log_broadcaster: self.log_broadcaster,
|
||||
context_manager,
|
||||
hooks,
|
||||
skill_registry,
|
||||
skill_catalog,
|
||||
cost_guard,
|
||||
session: self.session,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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()))),
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}))
|
||||
}
|
||||
@@ -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 }))
|
||||
}
|
||||
@@ -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::*;
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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()))),
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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}"),
|
||||
})?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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()),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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(())
|
||||
}
|
||||
Reference in New Issue
Block a user