mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
09e6a7e6d8 | ||
|
|
5814d77b16 | ||
|
|
83773af997 | ||
|
|
7474fd4c52 | ||
|
|
64b6f559fd | ||
|
|
0de6f6aabb |
@@ -151,6 +151,12 @@ src/
|
||||
│ ├── rate_limiter.rs # Per-tool rate limiting
|
||||
│ └── storage.rs # Linear memory persistence
|
||||
│
|
||||
├── db/ # Database abstraction layer
|
||||
│ ├── mod.rs # Database trait (~60 async methods)
|
||||
│ ├── postgres.rs # PostgreSQL backend (delegates to Store + Repository)
|
||||
│ ├── libsql_backend.rs # libSQL/Turso backend (embedded SQLite)
|
||||
│ └── libsql_migrations.rs # SQLite-dialect schema (idempotent)
|
||||
│
|
||||
├── workspace/ # Persistent memory system (OpenClaw-inspired)
|
||||
│ ├── mod.rs # Workspace struct, memory operations
|
||||
│ ├── document.rs # MemoryDocument, MemoryChunk, WorkspaceEntry
|
||||
@@ -201,6 +207,7 @@ When designing new features or systems, always prefer generic/extensible archite
|
||||
- Use `RwLock` for concurrent read/write access
|
||||
|
||||
### Traits for Extensibility
|
||||
- `Database` - Add new database backends (must implement all ~60 methods)
|
||||
- `Channel` - Add new input sources
|
||||
- `Tool` - Add new capabilities
|
||||
- `LlmProvider` - Add new LLM backends
|
||||
@@ -248,7 +255,12 @@ Pending -> InProgress -> Completed -> Submitted -> Accepted
|
||||
|
||||
Environment variables (see `.env.example`):
|
||||
```bash
|
||||
# Database backend (default: postgres)
|
||||
DATABASE_BACKEND=postgres # or "libsql" / "turso"
|
||||
DATABASE_URL=postgres://user:pass@localhost/ironclaw
|
||||
LIBSQL_PATH=~/.ironclaw/ironclaw.db # libSQL local path (default)
|
||||
# LIBSQL_URL=libsql://xxx.turso.io # Turso cloud (optional)
|
||||
# LIBSQL_AUTH_TOKEN=xxx # Required with LIBSQL_URL
|
||||
|
||||
# NEAR AI (required)
|
||||
NEARAI_SESSION_TOKEN=sess_...
|
||||
@@ -308,7 +320,51 @@ Session tokens have the format `sess_xxx` (37 characters). They are authenticate
|
||||
|
||||
## Database
|
||||
|
||||
Single migration in `migrations/V1__initial.sql`. Tables:
|
||||
IronClaw supports two database backends, selected at compile time via Cargo feature flags and at runtime via the `DATABASE_BACKEND` environment variable.
|
||||
|
||||
**IMPORTANT: All new features that touch persistence MUST support both backends.** Implement the operation as a method on the `Database` trait in `src/db/mod.rs`, then add the implementation in both `src/db/postgres.rs` (delegate to Store/Repository) and `src/db/libsql_backend.rs` (native SQL).
|
||||
|
||||
### Backends
|
||||
|
||||
| Backend | Feature Flag | Default | Use Case |
|
||||
|---------|-------------|---------|----------|
|
||||
| PostgreSQL | `postgres` (default) | Yes | Production, existing deployments |
|
||||
| libSQL/Turso | `libsql` | No | Zero-dependency local mode, edge, Turso cloud |
|
||||
|
||||
```bash
|
||||
# Build with PostgreSQL only (default)
|
||||
cargo build
|
||||
|
||||
# Build with libSQL only
|
||||
cargo build --no-default-features --features libsql
|
||||
|
||||
# Build with both backends available
|
||||
cargo build --features "postgres,libsql"
|
||||
```
|
||||
|
||||
### Database Trait
|
||||
|
||||
The `Database` trait (`src/db/mod.rs`) defines ~60 async methods covering all persistence:
|
||||
- Conversations, messages, metadata
|
||||
- Jobs, actions, LLM calls, estimation snapshots
|
||||
- Sandbox jobs, job events
|
||||
- Routines, routine runs
|
||||
- Tool failures, settings
|
||||
- Workspace: documents, chunks, hybrid search
|
||||
|
||||
Both backends implement this trait. PostgreSQL delegates to the existing `Store` + `Repository`. libSQL implements native SQLite-dialect SQL.
|
||||
|
||||
### Schema
|
||||
|
||||
**PostgreSQL:** `migrations/V1__initial.sql` (351 lines). Uses pgvector for embeddings, tsvector for FTS, PL/pgSQL functions. Managed by `refinery`.
|
||||
|
||||
**libSQL:** `src/db/libsql_migrations.rs` (consolidated schema, ~480 lines). Translates PG types:
|
||||
- `UUID` -> `TEXT`, `TIMESTAMPTZ` -> `TEXT` (ISO-8601), `JSONB` -> `TEXT`
|
||||
- `VECTOR(1536)` -> `F32_BLOB(1536)` with `libsql_vector_idx`
|
||||
- `tsvector`/`ts_rank_cd` -> FTS5 virtual table with sync triggers
|
||||
- PL/pgSQL functions -> SQLite triggers
|
||||
|
||||
**Tables (both backends):**
|
||||
|
||||
**Core:**
|
||||
- `conversations` - Multi-channel conversation tracking
|
||||
@@ -320,12 +376,41 @@ Single migration in `migrations/V1__initial.sql`. Tables:
|
||||
|
||||
**Workspace/Memory:**
|
||||
- `memory_documents` - Flexible path-based files (e.g., "context/vision.md", "daily/2024-01-15.md")
|
||||
- `memory_chunks` - Chunked content with FTS (tsvector) and vector (pgvector) indexes
|
||||
- `memory_chunks` - Chunked content with FTS and vector indexes
|
||||
- `heartbeat_state` - Periodic execution tracking
|
||||
|
||||
Requires pgvector extension: `CREATE EXTENSION IF NOT EXISTS vector;`
|
||||
**Other:**
|
||||
- `routines`, `routine_runs` - Scheduled/reactive execution
|
||||
- `settings` - Per-user key-value settings
|
||||
- `tool_failures` - Self-repair tracking
|
||||
- `secrets`, `wasm_tools`, `tool_capabilities` - Extension infrastructure
|
||||
|
||||
Run migrations: `refinery migrate -c refinery.toml`
|
||||
### Configuration
|
||||
|
||||
```bash
|
||||
# Backend selection (default: postgres)
|
||||
DATABASE_BACKEND=libsql
|
||||
|
||||
# PostgreSQL
|
||||
DATABASE_URL=postgres://user:pass@localhost/ironclaw
|
||||
|
||||
# libSQL (embedded)
|
||||
LIBSQL_PATH=~/.ironclaw/ironclaw.db # Default path
|
||||
|
||||
# libSQL (Turso cloud sync)
|
||||
LIBSQL_URL=libsql://your-db.turso.io
|
||||
LIBSQL_AUTH_TOKEN=your-token # Required when LIBSQL_URL is set
|
||||
```
|
||||
|
||||
### Current Limitations (libSQL backend)
|
||||
|
||||
- **Workspace/memory system** not yet wired through Database trait (requires Store migration)
|
||||
- **Secrets store** not yet available (still requires PostgresSecretsStore)
|
||||
- **Hybrid search** uses FTS5 only (vector search via libsql_vector_idx not yet implemented)
|
||||
- **Settings reload from DB** skipped (Config::from_db requires Store)
|
||||
- No incremental migration versioning (schema is CREATE IF NOT EXISTS, no ALTER TABLE support yet)
|
||||
- **No encryption at rest** -- The local SQLite database file stores conversation content, job data, workspace memory, and other application data in plaintext. Only secrets (API tokens, credentials) are encrypted via AES-256-GCM before storage. Users handling sensitive data should use full-disk encryption (FileVault, LUKS, BitLocker) or consider the PostgreSQL backend with TDE/encrypted storage.
|
||||
- **JSON merge patch vs path-targeted update** -- The libSQL backend uses RFC 7396 JSON Merge Patch (`json_patch`) for metadata updates, while PostgreSQL uses path-targeted `jsonb_set`. Merge patch replaces top-level keys entirely, which may drop nested keys not present in the patch. Callers should avoid relying on partial nested object updates in metadata fields.
|
||||
|
||||
## Safety Layer
|
||||
|
||||
@@ -387,6 +472,7 @@ Key test patterns:
|
||||
- ✅ **Claude Code mode** - Delegate jobs to Claude CLI inside containers
|
||||
- ✅ **Routines system** - Cron, event, webhook, and manual triggers with guardrails
|
||||
- ✅ **Extension management** - Install, auth, activate MCP/WASM extensions via CLI and web UI
|
||||
- ✅ **libSQL/Turso backend** - Database trait abstraction (`src/db/`), feature-gated dual backend support (postgres/libsql), embedded SQLite for zero-dependency local mode
|
||||
|
||||
## Adding a New Tool
|
||||
|
||||
@@ -625,7 +711,7 @@ Four tools for LLM use:
|
||||
|
||||
### Hybrid Search (RRF)
|
||||
|
||||
Combines full-text search (PostgreSQL `ts_rank_cd`) and vector similarity (pgvector cosine) using Reciprocal Rank Fusion:
|
||||
Combines full-text search and vector similarity using Reciprocal Rank Fusion:
|
||||
|
||||
```
|
||||
score(d) = Σ 1/(k + rank(d)) for each method where d appears
|
||||
@@ -633,6 +719,10 @@ score(d) = Σ 1/(k + rank(d)) for each method where d appears
|
||||
|
||||
Default k=60. Results from both methods are combined, with documents appearing in both getting boosted scores.
|
||||
|
||||
**Backend differences:**
|
||||
- **PostgreSQL:** `ts_rank_cd` for FTS, pgvector cosine distance for vectors, full RRF
|
||||
- **libSQL:** FTS5 for keyword search only (vector search via `libsql_vector_idx` not yet wired)
|
||||
|
||||
### Heartbeat System
|
||||
|
||||
Proactive periodic execution (default: 30 minutes):
|
||||
|
||||
Generated
+635
-59
File diff suppressed because it is too large
Load Diff
+24
-8
@@ -28,11 +28,14 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "rus
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
|
||||
# Database
|
||||
deadpool-postgres = "0.14"
|
||||
tokio-postgres = { version = "0.7", features = ["with-uuid-1", "with-chrono-0_4", "with-serde_json-1"] }
|
||||
postgres-types = { version = "0.2", features = ["with-serde_json-1"] }
|
||||
refinery = { version = "0.8", features = ["tokio-postgres"] }
|
||||
# Database - PostgreSQL (default, feature-gated)
|
||||
deadpool-postgres = { version = "0.14", optional = true }
|
||||
tokio-postgres = { version = "0.7", features = ["with-uuid-1", "with-chrono-0_4", "with-serde_json-1"], optional = true }
|
||||
postgres-types = { version = "0.2", features = ["with-serde_json-1"], optional = true }
|
||||
refinery = { version = "0.8", features = ["tokio-postgres"], optional = true }
|
||||
|
||||
# Database - libSQL/Turso (optional embedded database)
|
||||
libsql = { version = "0.6", optional = true, default-features = false, features = ["core", "replication"] }
|
||||
|
||||
# Error handling
|
||||
thiserror = "2"
|
||||
@@ -48,7 +51,7 @@ dotenvy = "0.15"
|
||||
# Core types
|
||||
uuid = { version = "1", features = ["v4", "serde"] }
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
rust_decimal = { version = "1", features = ["serde", "serde-with-str", "db-tokio-postgres", "maths"] }
|
||||
rust_decimal = { version = "1", features = ["serde", "serde-with-str", "maths"] }
|
||||
rust_decimal_macros = "1"
|
||||
|
||||
# Async traits
|
||||
@@ -89,7 +92,7 @@ open = "5"
|
||||
|
||||
# Vector embeddings for semantic search
|
||||
# The postgres feature provides ToSql/FromSql for postgres-types (shared by tokio-postgres)
|
||||
pgvector = { version = "0.4", features = ["postgres"] }
|
||||
pgvector = { version = "0.4", features = ["postgres"], optional = true }
|
||||
|
||||
# WASM sandbox for untrusted tool execution
|
||||
wasmtime = { version = "28", features = ["component-model"] }
|
||||
@@ -135,9 +138,22 @@ pretty_assertions = "1"
|
||||
tempfile = "3"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
default = ["postgres"]
|
||||
postgres = [
|
||||
"dep:deadpool-postgres",
|
||||
"dep:tokio-postgres",
|
||||
"dep:postgres-types",
|
||||
"dep:refinery",
|
||||
"dep:pgvector",
|
||||
"rust_decimal/db-tokio-postgres",
|
||||
]
|
||||
libsql = ["dep:libsql"]
|
||||
integration = []
|
||||
|
||||
[[example]]
|
||||
name = "test_heartbeat"
|
||||
required-features = ["postgres"]
|
||||
|
||||
# The profile that 'cargo dist' will build with
|
||||
[profile.dist]
|
||||
inherits = "release"
|
||||
|
||||
@@ -19,9 +19,9 @@ use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusU
|
||||
use crate::config::{AgentConfig, HeartbeatConfig, RoutineConfig};
|
||||
use crate::context::ContextManager;
|
||||
use crate::context::JobContext;
|
||||
use crate::db::Database;
|
||||
use crate::error::Error;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
@@ -59,7 +59,7 @@ enum AgenticLoopResult {
|
||||
///
|
||||
/// Bundles the shared components to reduce argument count.
|
||||
pub struct AgentDeps {
|
||||
pub store: Option<Arc<Store>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
pub llm: Arc<dyn LlmProvider>,
|
||||
pub safety: Arc<SafetyLayer>,
|
||||
pub tools: Arc<ToolRegistry>,
|
||||
@@ -124,7 +124,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Convenience accessors
|
||||
fn store(&self) -> Option<&Arc<Store>> {
|
||||
fn store(&self) -> Option<&Arc<dyn Database>> {
|
||||
self.deps.store.as_ref()
|
||||
}
|
||||
|
||||
|
||||
@@ -24,14 +24,14 @@ use crate::agent::routine::{
|
||||
};
|
||||
use crate::channels::{IncomingMessage, OutgoingResponse};
|
||||
use crate::config::RoutineConfig;
|
||||
use crate::history::Store;
|
||||
use crate::db::Database;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
/// The routine execution engine.
|
||||
pub struct RoutineEngine {
|
||||
config: RoutineConfig,
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
/// Sender for notifications (routed to channel manager).
|
||||
@@ -45,7 +45,7 @@ pub struct RoutineEngine {
|
||||
impl RoutineEngine {
|
||||
pub fn new(
|
||||
config: RoutineConfig,
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
@@ -294,7 +294,7 @@ impl RoutineEngine {
|
||||
|
||||
/// Shared context passed to the execution function.
|
||||
struct EngineContext {
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
|
||||
@@ -12,8 +12,8 @@ use crate::agent::task::{Task, TaskContext, TaskOutput};
|
||||
use crate::agent::worker::{Worker, WorkerDeps};
|
||||
use crate::config::AgentConfig;
|
||||
use crate::context::{ContextManager, JobContext, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::{Error, JobError};
|
||||
use crate::history::Store;
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
@@ -48,7 +48,7 @@ pub struct Scheduler {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
/// Running jobs (main LLM-driven jobs).
|
||||
jobs: Arc<RwLock<HashMap<Uuid, ScheduledJob>>>,
|
||||
/// Running sub-tasks (tool executions, background tasks).
|
||||
@@ -63,7 +63,7 @@ impl Scheduler {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
|
||||
@@ -8,8 +8,8 @@ use chrono::{DateTime, Utc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::context::{ContextManager, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::RepairError;
|
||||
use crate::history::Store;
|
||||
use crate::tools::{BuildRequirement, Language, SoftwareBuilder, SoftwareType, ToolRegistry};
|
||||
|
||||
/// A job that has been detected as stuck.
|
||||
@@ -69,7 +69,7 @@ pub struct DefaultSelfRepair {
|
||||
#[allow(dead_code)] // Will be used for time-based stuck detection
|
||||
stuck_threshold: Duration,
|
||||
max_repair_attempts: u32,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
builder: Option<Arc<dyn SoftwareBuilder>>,
|
||||
#[allow(dead_code)] // Will be used for tool hot-reload after repair
|
||||
tools: Option<Arc<ToolRegistry>>,
|
||||
@@ -94,7 +94,7 @@ impl DefaultSelfRepair {
|
||||
|
||||
/// Add a Store for tool failure tracking.
|
||||
#[allow(dead_code)] // Public API for configuring repair with persistence
|
||||
pub fn with_store(mut self, store: Arc<Store>) -> Self {
|
||||
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
|
||||
self.store = Some(store);
|
||||
self
|
||||
}
|
||||
|
||||
+4
-4
@@ -10,8 +10,8 @@ use uuid::Uuid;
|
||||
use crate::agent::scheduler::WorkerMessage;
|
||||
use crate::agent::task::TaskOutput;
|
||||
use crate::context::{ContextManager, JobState};
|
||||
use crate::db::Database;
|
||||
use crate::error::Error;
|
||||
use crate::history::Store;
|
||||
use crate::llm::{
|
||||
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection,
|
||||
};
|
||||
@@ -28,7 +28,7 @@ pub struct WorkerDeps {
|
||||
pub llm: Arc<dyn LlmProvider>,
|
||||
pub safety: Arc<SafetyLayer>,
|
||||
pub tools: Arc<ToolRegistry>,
|
||||
pub store: Option<Arc<Store>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
pub timeout: Duration,
|
||||
pub use_planning: bool,
|
||||
}
|
||||
@@ -67,7 +67,7 @@ impl Worker {
|
||||
&self.deps.tools
|
||||
}
|
||||
|
||||
fn store(&self) -> Option<&Arc<Store>> {
|
||||
fn store(&self) -> Option<&Arc<dyn Database>> {
|
||||
self.deps.store.as_ref()
|
||||
}
|
||||
|
||||
@@ -381,7 +381,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
tools: Arc<ToolRegistry>,
|
||||
context_manager: Arc<ContextManager>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
job_id: Uuid,
|
||||
tool_name: &str,
|
||||
params: &serde_json::Value,
|
||||
|
||||
+1
-1
@@ -124,7 +124,7 @@ impl BootstrapConfig {
|
||||
/// If both conditions hold, migrates settings, MCP servers, and session data
|
||||
/// to the database, writes `bootstrap.json`, and renames old files to `.migrated`.
|
||||
pub async fn migrate_disk_to_db(
|
||||
store: &crate::history::Store,
|
||||
store: &dyn crate::db::Database,
|
||||
user_id: &str,
|
||||
) -> Result<(), MigrationError> {
|
||||
let legacy_settings_path = BootstrapConfig::legacy_settings_path();
|
||||
|
||||
@@ -32,9 +32,9 @@ use tokio_stream::wrappers::ReceiverStream;
|
||||
use crate::agent::SessionManager;
|
||||
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
|
||||
use crate::config::GatewayConfig;
|
||||
use crate::db::Database;
|
||||
use crate::error::ChannelError;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::workspace::Workspace;
|
||||
@@ -147,7 +147,7 @@ impl GatewayChannel {
|
||||
}
|
||||
|
||||
/// Inject the database store for sandbox job persistence.
|
||||
pub fn with_store(mut self, store: Arc<Store>) -> Self {
|
||||
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
|
||||
self.rebuild_state(|s| s.store = Some(store));
|
||||
self
|
||||
}
|
||||
|
||||
@@ -30,8 +30,8 @@ use crate::channels::web::auth::{AuthState, auth_middleware};
|
||||
use crate::channels::web::log_layer::LogBroadcaster;
|
||||
use crate::channels::web::sse::SseManager;
|
||||
use crate::channels::web::types::*;
|
||||
use crate::db::Database;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::workspace::Workspace;
|
||||
@@ -126,7 +126,7 @@ pub struct GatewayState {
|
||||
/// Tool registry for listing registered tools.
|
||||
pub tool_registry: Option<Arc<ToolRegistry>>,
|
||||
/// Database store for sandbox job persistence.
|
||||
pub store: Option<Arc<Store>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
/// Container job manager for sandbox operations.
|
||||
pub job_manager: Option<Arc<ContainerJobManager>>,
|
||||
/// Prompt queue for Claude Code follow-up prompts.
|
||||
|
||||
+22
-19
@@ -1,7 +1,9 @@
|
||||
//! Configuration management CLI commands.
|
||||
//!
|
||||
//! Commands for viewing and modifying settings.
|
||||
//! Settings are stored in PostgreSQL (env > DB > default).
|
||||
//! Settings are stored in the database (env > DB > default).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::Subcommand;
|
||||
|
||||
@@ -49,8 +51,8 @@ pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
|
||||
let _ = dotenvy::dotenv();
|
||||
|
||||
// Try to connect to the DB for settings access
|
||||
let store = match connect_store().await {
|
||||
Ok(s) => Some(s),
|
||||
let db: Option<Arc<dyn crate::db::Database>> = match connect_db().await {
|
||||
Ok(d) => Some(d),
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"Warning: Could not connect to database ({}), using disk fallback",
|
||||
@@ -60,29 +62,30 @@ pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
|
||||
}
|
||||
};
|
||||
|
||||
let db_ref = db.as_deref();
|
||||
match cmd {
|
||||
ConfigCommand::List { filter } => list_settings(store.as_ref(), filter).await,
|
||||
ConfigCommand::Get { path } => get_setting(store.as_ref(), &path).await,
|
||||
ConfigCommand::Set { path, value } => set_setting(store.as_ref(), &path, &value).await,
|
||||
ConfigCommand::Reset { path } => reset_setting(store.as_ref(), &path).await,
|
||||
ConfigCommand::Path => show_path(store.is_some()),
|
||||
ConfigCommand::List { filter } => list_settings(db_ref, filter).await,
|
||||
ConfigCommand::Get { path } => get_setting(db_ref, &path).await,
|
||||
ConfigCommand::Set { path, value } => set_setting(db_ref, &path, &value).await,
|
||||
ConfigCommand::Reset { path } => reset_setting(db_ref, &path).await,
|
||||
ConfigCommand::Path => show_path(db_ref.is_some()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bootstrap a DB connection for config commands.
|
||||
async fn connect_store() -> anyhow::Result<crate::history::Store> {
|
||||
/// Bootstrap a DB connection for config commands (backend-agnostic).
|
||||
async fn connect_db() -> anyhow::Result<Arc<dyn crate::db::Database>> {
|
||||
let config = crate::config::Config::from_env()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Ok(store)
|
||||
crate::db::connect_from_config(&config.database)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))
|
||||
}
|
||||
|
||||
const DEFAULT_USER_ID: &str = "default";
|
||||
|
||||
/// Load settings: DB if available, else disk.
|
||||
async fn load_settings(store: Option<&crate::history::Store>) -> Settings {
|
||||
async fn load_settings(store: Option<&dyn crate::db::Database>) -> Settings {
|
||||
if let Some(store) = store {
|
||||
match store.get_all_settings(DEFAULT_USER_ID).await {
|
||||
Ok(map) if !map.is_empty() => return Settings::from_db_map(&map),
|
||||
@@ -94,7 +97,7 @@ async fn load_settings(store: Option<&crate::history::Store>) -> Settings {
|
||||
|
||||
/// List all settings.
|
||||
async fn list_settings(
|
||||
store: Option<&crate::history::Store>,
|
||||
store: Option<&dyn crate::db::Database>,
|
||||
filter: Option<String>,
|
||||
) -> anyhow::Result<()> {
|
||||
let settings = load_settings(store).await;
|
||||
@@ -126,7 +129,7 @@ async fn list_settings(
|
||||
}
|
||||
|
||||
/// Get a specific setting.
|
||||
async fn get_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
|
||||
async fn get_setting(store: Option<&dyn crate::db::Database>, path: &str) -> anyhow::Result<()> {
|
||||
let settings = load_settings(store).await;
|
||||
|
||||
match settings.get(path) {
|
||||
@@ -142,7 +145,7 @@ async fn get_setting(store: Option<&crate::history::Store>, path: &str) -> anyho
|
||||
|
||||
/// Set a setting value.
|
||||
async fn set_setting(
|
||||
store: Option<&crate::history::Store>,
|
||||
store: Option<&dyn crate::db::Database>,
|
||||
path: &str,
|
||||
value: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
@@ -171,7 +174,7 @@ async fn set_setting(
|
||||
}
|
||||
|
||||
/// Reset a setting to default.
|
||||
async fn reset_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
|
||||
async fn reset_setting(store: Option<&dyn crate::db::Database>, path: &str) -> anyhow::Result<()> {
|
||||
let default = Settings::default();
|
||||
let default_value = default
|
||||
.get(path)
|
||||
@@ -196,7 +199,7 @@ async fn reset_setting(store: Option<&crate::history::Store>, path: &str) -> any
|
||||
/// Show the settings storage info.
|
||||
fn show_path(has_db: bool) -> anyhow::Result<()> {
|
||||
if has_db {
|
||||
println!("Settings stored in: PostgreSQL (settings table)");
|
||||
println!("Settings stored in: database (settings table)");
|
||||
println!(
|
||||
"Bootstrap config: {}",
|
||||
crate::bootstrap::BootstrapConfig::default_path().display()
|
||||
|
||||
+82
-35
@@ -8,8 +8,10 @@ use std::sync::Arc;
|
||||
use clap::Subcommand;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::history::Store;
|
||||
use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
use crate::db::Database;
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::secrets::PostgresSecretsStore;
|
||||
use crate::secrets::{SecretsCrypto, SecretsStore};
|
||||
use crate::tools::mcp::{
|
||||
McpClient, McpServerConfig, McpSessionManager, OAuthConfig,
|
||||
auth::{authorize_mcp_server, is_authenticated},
|
||||
@@ -172,10 +174,10 @@ async fn add_server(
|
||||
config.validate()?;
|
||||
|
||||
// Save (DB if available, else disk)
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
servers.upsert(config);
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
save_servers(db.as_deref(), &servers).await?;
|
||||
|
||||
println!();
|
||||
println!(" ✓ Added MCP server '{}'", name);
|
||||
@@ -193,12 +195,12 @@ async fn add_server(
|
||||
|
||||
/// Remove an MCP server.
|
||||
async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
if !servers.remove(&name) {
|
||||
anyhow::bail!("Server '{}' not found", name);
|
||||
}
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
save_servers(db.as_deref(), &servers).await?;
|
||||
|
||||
println!();
|
||||
println!(" ✓ Removed MCP server '{}'", name);
|
||||
@@ -209,8 +211,8 @@ async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
|
||||
/// List configured MCP servers.
|
||||
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
|
||||
if servers.servers.is_empty() {
|
||||
println!();
|
||||
@@ -268,8 +270,8 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
/// Authenticate with an MCP server.
|
||||
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
let server = servers
|
||||
.get(&name)
|
||||
.cloned()
|
||||
@@ -341,8 +343,8 @@ async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
/// Test connection to an MCP server.
|
||||
async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let servers = load_servers(db.as_deref()).await?;
|
||||
let server = servers
|
||||
.get(&name)
|
||||
.cloned()
|
||||
@@ -437,8 +439,8 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
|
||||
/// Toggle server enabled/disabled state.
|
||||
async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> {
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
let db = connect_db().await;
|
||||
let mut servers = load_servers(db.as_deref()).await?;
|
||||
|
||||
let server = servers
|
||||
.get_mut(&name)
|
||||
@@ -453,7 +455,7 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
|
||||
};
|
||||
|
||||
server.enabled = new_state;
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
save_servers(db.as_deref(), &servers).await?;
|
||||
|
||||
let status = if new_state { "enabled" } else { "disabled" };
|
||||
println!();
|
||||
@@ -465,18 +467,16 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
|
||||
|
||||
const DEFAULT_USER_ID: &str = "default";
|
||||
|
||||
/// Try to connect to the database store for DB-backed config.
|
||||
async fn connect_store() -> Option<Store> {
|
||||
/// Try to connect to the database (backend-agnostic).
|
||||
async fn connect_db() -> Option<Arc<dyn Database>> {
|
||||
let config = Config::from_env().await.ok()?;
|
||||
let store = Store::new(&config.database).await.ok()?;
|
||||
store.run_migrations().await.ok()?;
|
||||
Some(store)
|
||||
crate::db::connect_from_config(&config.database).await.ok()
|
||||
}
|
||||
|
||||
/// Load MCP servers (DB if available, else disk).
|
||||
async fn load_servers(store: Option<&Store>) -> Result<McpServersFile, config::ConfigError> {
|
||||
if let Some(store) = store {
|
||||
config::load_mcp_servers_from_db(store, DEFAULT_USER_ID).await
|
||||
async fn load_servers(db: Option<&dyn Database>) -> Result<McpServersFile, config::ConfigError> {
|
||||
if let Some(db) = db {
|
||||
config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await
|
||||
} else {
|
||||
config::load_mcp_servers().await
|
||||
}
|
||||
@@ -484,11 +484,11 @@ async fn load_servers(store: Option<&Store>) -> Result<McpServersFile, config::C
|
||||
|
||||
/// Save MCP servers (DB if available, else disk).
|
||||
async fn save_servers(
|
||||
store: Option<&Store>,
|
||||
db: Option<&dyn Database>,
|
||||
servers: &McpServersFile,
|
||||
) -> Result<(), config::ConfigError> {
|
||||
if let Some(store) = store {
|
||||
config::save_mcp_servers_to_db(store, DEFAULT_USER_ID, servers).await
|
||||
if let Some(db) = db {
|
||||
config::save_mcp_servers_to_db(db, DEFAULT_USER_ID, servers).await
|
||||
} else {
|
||||
config::save_mcp_servers(servers).await
|
||||
}
|
||||
@@ -504,14 +504,61 @@ async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Syn
|
||||
)
|
||||
})?;
|
||||
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
Ok(Arc::new(PostgresSecretsStore::new(
|
||||
store.pool(),
|
||||
Arc::new(crypto),
|
||||
)))
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Ok(Arc::new(PostgresSecretsStore::new(
|
||||
store.pool(),
|
||||
Arc::new(crypto),
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
|
||||
{
|
||||
use crate::db::Database as _;
|
||||
use crate::db::libsql_backend::LibSqlBackend;
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config
|
||||
.database
|
||||
.libsql_path
|
||||
.as_deref()
|
||||
.unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.database.libsql_url {
|
||||
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
|
||||
})?;
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
};
|
||||
backend
|
||||
.run_migrations()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
return Ok(Arc::new(crate::secrets::LibSqlSecretsStore::new(
|
||||
backend.shared_db(),
|
||||
Arc::new(crypto),
|
||||
)));
|
||||
}
|
||||
|
||||
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
|
||||
{
|
||||
let _ = crypto;
|
||||
anyhow::bail!(
|
||||
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+26
-1
@@ -9,6 +9,30 @@ use clap::Subcommand;
|
||||
|
||||
use crate::workspace::{EmbeddingProvider, SearchConfig, Workspace};
|
||||
|
||||
/// Run a memory command using the Database trait (works with any backend).
|
||||
pub async fn run_memory_command_with_db(
|
||||
cmd: MemoryCommand,
|
||||
db: std::sync::Arc<dyn crate::db::Database>,
|
||||
embeddings: Option<Arc<dyn EmbeddingProvider>>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut workspace = Workspace::new_with_db("default", db);
|
||||
if let Some(emb) = embeddings {
|
||||
workspace = workspace.with_embeddings(emb);
|
||||
}
|
||||
|
||||
match cmd {
|
||||
MemoryCommand::Search { query, limit } => search(&workspace, &query, limit).await,
|
||||
MemoryCommand::Read { path } => read(&workspace, &path).await,
|
||||
MemoryCommand::Write {
|
||||
path,
|
||||
content,
|
||||
append,
|
||||
} => write(&workspace, &path, content, append).await,
|
||||
MemoryCommand::Tree { path, depth } => tree(&workspace, &path, depth).await,
|
||||
MemoryCommand::Status => status(&workspace).await,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
pub enum MemoryCommand {
|
||||
/// Search workspace memory (hybrid full-text + semantic)
|
||||
@@ -55,7 +79,8 @@ pub enum MemoryCommand {
|
||||
Status,
|
||||
}
|
||||
|
||||
/// Run a memory command.
|
||||
/// Run a memory command (PostgreSQL backend).
|
||||
#[cfg(feature = "postgres")]
|
||||
pub async fn run_memory_command(
|
||||
cmd: MemoryCommand,
|
||||
pool: deadpool_postgres::Pool,
|
||||
|
||||
+4
-1
@@ -18,7 +18,10 @@ mod tool;
|
||||
|
||||
pub use config::{ConfigCommand, run_config_command};
|
||||
pub use mcp::{McpCommand, run_mcp_command};
|
||||
pub use memory::{MemoryCommand, run_memory_command};
|
||||
pub use memory::MemoryCommand;
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use memory::run_memory_command;
|
||||
pub use memory::run_memory_command_with_db;
|
||||
pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store};
|
||||
pub use status::run_status_command;
|
||||
pub use tool::{ToolCommand, run_tool_command};
|
||||
|
||||
@@ -135,6 +135,7 @@ pub async fn run_status_command() -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
async fn check_database() -> anyhow::Result<()> {
|
||||
let _ = dotenvy::dotenv();
|
||||
let settings = Settings::load();
|
||||
@@ -167,6 +168,12 @@ async fn check_database() -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
async fn check_database() -> anyhow::Result<()> {
|
||||
// For non-postgres backends, just report configured
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn count_wasm_files(dir: &std::path::Path) -> usize {
|
||||
std::fs::read_dir(dir)
|
||||
.map(|entries| {
|
||||
|
||||
+63
-13
@@ -11,8 +11,11 @@ use clap::Subcommand;
|
||||
use tokio::fs;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::history::Store;
|
||||
use crate::secrets::{CreateSecretParams, PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
#[allow(unused_imports)]
|
||||
use crate::db::Database;
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::secrets::PostgresSecretsStore;
|
||||
use crate::secrets::{CreateSecretParams, SecretsCrypto, SecretsStore};
|
||||
use crate::tools::wasm::{CapabilitiesFile, compute_binary_hash};
|
||||
|
||||
/// Default tools directory.
|
||||
@@ -722,11 +725,58 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyho
|
||||
)
|
||||
})?;
|
||||
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
let secrets_store = Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)));
|
||||
|
||||
let secrets_store: Arc<dyn SecretsStore + Send + Sync> = {
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)))
|
||||
}
|
||||
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
|
||||
{
|
||||
use crate::db::Database as _;
|
||||
use crate::db::libsql_backend::LibSqlBackend;
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config
|
||||
.database
|
||||
.libsql_path
|
||||
.as_deref()
|
||||
.unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.database.libsql_url {
|
||||
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
|
||||
})?;
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
};
|
||||
backend
|
||||
.run_migrations()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
Arc::new(crate::secrets::LibSqlSecretsStore::new(
|
||||
backend.shared_db(),
|
||||
Arc::new(crypto),
|
||||
))
|
||||
}
|
||||
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
|
||||
{
|
||||
let _ = crypto;
|
||||
anyhow::bail!(
|
||||
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// Check if already configured
|
||||
let already_configured = secrets_store
|
||||
@@ -772,13 +822,13 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyho
|
||||
println!(" Validation failed: {}", e);
|
||||
println!();
|
||||
println!(" Falling back to manual entry...");
|
||||
return auth_tool_manual(&secrets_store, &user_id, &auth).await;
|
||||
return auth_tool_manual(secrets_store.as_ref(), &user_id, &auth).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Save the token
|
||||
save_token(&secrets_store, &user_id, &auth, &token).await?;
|
||||
save_token(secrets_store.as_ref(), &user_id, &auth, &token).await?;
|
||||
print_success(display_name);
|
||||
return Ok(());
|
||||
}
|
||||
@@ -787,16 +837,16 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyho
|
||||
|
||||
// Check for OAuth configuration
|
||||
if let Some(ref oauth) = auth.oauth {
|
||||
return auth_tool_oauth(&secrets_store, &user_id, &auth, oauth).await;
|
||||
return auth_tool_oauth(secrets_store.as_ref(), &user_id, &auth, oauth).await;
|
||||
}
|
||||
|
||||
// Fall back to manual entry
|
||||
auth_tool_manual(&secrets_store, &user_id, &auth).await
|
||||
auth_tool_manual(secrets_store.as_ref(), &user_id, &auth).await
|
||||
}
|
||||
|
||||
/// OAuth browser-based login flow.
|
||||
async fn auth_tool_oauth(
|
||||
store: &PostgresSecretsStore,
|
||||
store: &(dyn SecretsStore + Send + Sync),
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
oauth: &crate::tools::wasm::OAuthConfigSchema,
|
||||
@@ -1044,7 +1094,7 @@ async fn auth_tool_oauth(
|
||||
|
||||
/// Manual token entry flow.
|
||||
async fn auth_tool_manual(
|
||||
store: &PostgresSecretsStore,
|
||||
store: &(dyn SecretsStore + Send + Sync),
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
) -> anyhow::Result<()> {
|
||||
@@ -1217,7 +1267,7 @@ async fn validate_token(
|
||||
|
||||
/// Save token to secrets store.
|
||||
async fn save_token(
|
||||
store: &PostgresSecretsStore,
|
||||
store: &(dyn SecretsStore + Send + Sync),
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
token: &str,
|
||||
|
||||
+86
-1
@@ -38,7 +38,7 @@ impl Config {
|
||||
/// Priority: env var > DB settings > default.
|
||||
/// This is the primary way to load config after DB is connected.
|
||||
pub async fn from_db(
|
||||
store: &crate::history::Store,
|
||||
store: &dyn crate::db::Database,
|
||||
user_id: &str,
|
||||
bootstrap: &crate::bootstrap::BootstrapConfig,
|
||||
) -> Result<Self, ConfigError> {
|
||||
@@ -134,17 +134,72 @@ impl TunnelConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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::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 {
|
||||
fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> 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.
|
||||
let url = optional_env("DATABASE_URL")?
|
||||
.or_else(|| bootstrap.database_url.clone())
|
||||
.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(),
|
||||
@@ -160,9 +215,31 @@ impl DatabaseConfig {
|
||||
.or(bootstrap.database_pool_size)
|
||||
.unwrap_or(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,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -172,6 +249,14 @@ impl DatabaseConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Default libSQL database path (~/.ironclaw/ironclaw.db).
|
||||
pub fn default_libsql_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("ironclaw.db")
|
||||
}
|
||||
|
||||
/// Which LLM backend to use.
|
||||
///
|
||||
/// Defaults to `NearAi` to keep IronClaw close to the NEAR ecosystem.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,549 @@
|
||||
//! SQLite-dialect migrations for the libSQL/Turso backend.
|
||||
//!
|
||||
//! Consolidates all PostgreSQL migrations (V1-V8) into a single SQLite-compatible
|
||||
//! schema. Run once on database creation; idempotent via `IF NOT EXISTS`.
|
||||
|
||||
/// Consolidated schema for libSQL.
|
||||
///
|
||||
/// Translates PostgreSQL types and features:
|
||||
/// - `UUID` -> `TEXT` (store as hex string)
|
||||
/// - `TIMESTAMPTZ` -> `TEXT` (ISO-8601)
|
||||
/// - `JSONB` -> `TEXT` (JSON encoded)
|
||||
/// - `BYTEA` -> `BLOB`
|
||||
/// - `NUMERIC` -> `TEXT` (preserve precision for rust_decimal)
|
||||
/// - `TEXT[]` -> `TEXT` (JSON array)
|
||||
/// - `VECTOR(1536)` -> `F32_BLOB(1536)` (libsql native)
|
||||
/// - `TSVECTOR` -> FTS5 virtual table
|
||||
/// - `BIGSERIAL` -> `INTEGER PRIMARY KEY AUTOINCREMENT`
|
||||
/// - PL/pgSQL functions -> SQLite triggers
|
||||
pub const SCHEMA: &str = r#"
|
||||
|
||||
-- ==================== Migration tracking ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS _migrations (
|
||||
version INTEGER PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
applied_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
-- ==================== Conversations ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
id TEXT PRIMARY KEY,
|
||||
channel TEXT NOT NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
thread_id TEXT,
|
||||
started_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
last_activity TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
metadata TEXT NOT NULL DEFAULT '{}'
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_channel ON conversations(channel);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_last_activity ON conversations(last_activity);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS conversation_messages (
|
||||
id TEXT PRIMARY KEY,
|
||||
conversation_id TEXT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_messages_conversation
|
||||
ON conversation_messages(conversation_id);
|
||||
|
||||
-- ==================== Agent Jobs ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS agent_jobs (
|
||||
id TEXT PRIMARY KEY,
|
||||
marketplace_job_id TEXT,
|
||||
conversation_id TEXT REFERENCES conversations(id),
|
||||
title TEXT NOT NULL,
|
||||
description TEXT NOT NULL,
|
||||
category TEXT,
|
||||
status TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
user_id TEXT NOT NULL DEFAULT 'default',
|
||||
project_dir TEXT,
|
||||
job_mode TEXT NOT NULL DEFAULT 'worker',
|
||||
budget_amount TEXT,
|
||||
budget_token TEXT,
|
||||
bid_amount TEXT,
|
||||
estimated_cost TEXT,
|
||||
estimated_time_secs INTEGER,
|
||||
estimated_value TEXT,
|
||||
actual_cost TEXT,
|
||||
actual_time_secs INTEGER,
|
||||
success INTEGER,
|
||||
failure_reason TEXT,
|
||||
stuck_since TEXT,
|
||||
repair_attempts INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
started_at TEXT,
|
||||
completed_at TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_status ON agent_jobs(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_marketplace ON agent_jobs(marketplace_job_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_conversation ON agent_jobs(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_source ON agent_jobs(source);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_user ON agent_jobs(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_created ON agent_jobs(created_at DESC);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS job_actions (
|
||||
id TEXT PRIMARY KEY,
|
||||
job_id TEXT NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE,
|
||||
sequence_num INTEGER NOT NULL,
|
||||
tool_name TEXT NOT NULL,
|
||||
input TEXT NOT NULL,
|
||||
output_raw TEXT,
|
||||
output_sanitized TEXT,
|
||||
sanitization_warnings TEXT,
|
||||
cost TEXT,
|
||||
duration_ms INTEGER,
|
||||
success INTEGER NOT NULL,
|
||||
error_message TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE(job_id, sequence_num)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_job_actions_job_id ON job_actions(job_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_job_actions_tool ON job_actions(tool_name);
|
||||
|
||||
-- ==================== Dynamic Tools ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS dynamic_tools (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
description TEXT NOT NULL,
|
||||
parameters_schema TEXT NOT NULL,
|
||||
code TEXT NOT NULL,
|
||||
sandbox_config TEXT NOT NULL,
|
||||
created_by_job_id TEXT REFERENCES agent_jobs(id),
|
||||
success_count INTEGER NOT NULL DEFAULT 0,
|
||||
failure_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_error TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_dynamic_tools_status ON dynamic_tools(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_dynamic_tools_name ON dynamic_tools(name);
|
||||
|
||||
-- ==================== LLM Calls ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS llm_calls (
|
||||
id TEXT PRIMARY KEY,
|
||||
job_id TEXT REFERENCES agent_jobs(id) ON DELETE CASCADE,
|
||||
conversation_id TEXT REFERENCES conversations(id),
|
||||
provider TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
input_tokens INTEGER NOT NULL,
|
||||
output_tokens INTEGER NOT NULL,
|
||||
cost TEXT NOT NULL,
|
||||
purpose TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_llm_calls_job ON llm_calls(job_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_llm_calls_conversation ON llm_calls(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_llm_calls_provider ON llm_calls(provider);
|
||||
|
||||
-- ==================== Estimation ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS estimation_snapshots (
|
||||
id TEXT PRIMARY KEY,
|
||||
job_id TEXT NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE,
|
||||
category TEXT NOT NULL,
|
||||
tool_names TEXT NOT NULL DEFAULT '[]',
|
||||
estimated_cost TEXT NOT NULL,
|
||||
actual_cost TEXT,
|
||||
estimated_time_secs INTEGER NOT NULL,
|
||||
actual_time_secs INTEGER,
|
||||
estimated_value TEXT NOT NULL,
|
||||
actual_value TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_estimation_category ON estimation_snapshots(category);
|
||||
CREATE INDEX IF NOT EXISTS idx_estimation_job ON estimation_snapshots(job_id);
|
||||
|
||||
-- ==================== Self Repair ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS repair_attempts (
|
||||
id TEXT PRIMARY KEY,
|
||||
target_type TEXT NOT NULL,
|
||||
target_id TEXT NOT NULL,
|
||||
diagnosis TEXT NOT NULL,
|
||||
action_taken TEXT NOT NULL,
|
||||
success INTEGER NOT NULL,
|
||||
error_message TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_repair_attempts_target ON repair_attempts(target_type, target_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_repair_attempts_created ON repair_attempts(created_at);
|
||||
|
||||
-- ==================== Workspace: Memory Documents ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS memory_documents (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
agent_id TEXT,
|
||||
path TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
metadata TEXT NOT NULL DEFAULT '{}',
|
||||
UNIQUE (user_id, agent_id, path)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_documents_user ON memory_documents(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_documents_path ON memory_documents(user_id, path);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_documents_updated ON memory_documents(updated_at DESC);
|
||||
|
||||
-- Trigger to auto-update updated_at on memory_documents
|
||||
CREATE TRIGGER IF NOT EXISTS update_memory_documents_updated_at
|
||||
AFTER UPDATE ON memory_documents
|
||||
FOR EACH ROW
|
||||
WHEN NEW.updated_at = OLD.updated_at
|
||||
BEGIN
|
||||
UPDATE memory_documents SET updated_at = datetime('now') WHERE id = NEW.id;
|
||||
END;
|
||||
|
||||
-- ==================== Workspace: Memory Chunks ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS memory_chunks (
|
||||
_rowid INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
id TEXT NOT NULL UNIQUE,
|
||||
document_id TEXT NOT NULL REFERENCES memory_documents(id) ON DELETE CASCADE,
|
||||
chunk_index INTEGER NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
embedding F32_BLOB(1536),
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (document_id, chunk_index)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_chunks_document ON memory_chunks(document_id);
|
||||
|
||||
-- Vector index for semantic search (libSQL native)
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_chunks_embedding
|
||||
ON memory_chunks (libsql_vector_idx(embedding));
|
||||
|
||||
-- FTS5 virtual table for full-text search
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS memory_chunks_fts USING fts5(
|
||||
content,
|
||||
content='memory_chunks',
|
||||
content_rowid='_rowid'
|
||||
);
|
||||
|
||||
-- Triggers to keep FTS5 in sync with memory_chunks
|
||||
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_insert AFTER INSERT ON memory_chunks BEGIN
|
||||
INSERT INTO memory_chunks_fts(rowid, content) VALUES (new._rowid, new.content);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_delete AFTER DELETE ON memory_chunks BEGIN
|
||||
INSERT INTO memory_chunks_fts(memory_chunks_fts, rowid, content)
|
||||
VALUES ('delete', old._rowid, old.content);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_update AFTER UPDATE ON memory_chunks BEGIN
|
||||
INSERT INTO memory_chunks_fts(memory_chunks_fts, rowid, content)
|
||||
VALUES ('delete', old._rowid, old.content);
|
||||
INSERT INTO memory_chunks_fts(rowid, content) VALUES (new._rowid, new.content);
|
||||
END;
|
||||
|
||||
-- ==================== Workspace: Heartbeat State ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS heartbeat_state (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
agent_id TEXT,
|
||||
last_run TEXT,
|
||||
next_run TEXT,
|
||||
interval_seconds INTEGER NOT NULL DEFAULT 1800,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
consecutive_failures INTEGER NOT NULL DEFAULT 0,
|
||||
last_checks TEXT NOT NULL DEFAULT '{}',
|
||||
UNIQUE (user_id, agent_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_heartbeat_user ON heartbeat_state(user_id);
|
||||
|
||||
-- ==================== Secrets ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS secrets (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
encrypted_value BLOB NOT NULL,
|
||||
key_salt BLOB NOT NULL,
|
||||
provider TEXT,
|
||||
expires_at TEXT,
|
||||
last_used_at TEXT,
|
||||
usage_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (user_id, name)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_secrets_user ON secrets(user_id);
|
||||
|
||||
-- ==================== WASM Tools ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS wasm_tools (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
version TEXT NOT NULL DEFAULT '1.0.0',
|
||||
description TEXT NOT NULL,
|
||||
wasm_binary BLOB NOT NULL,
|
||||
binary_hash BLOB NOT NULL,
|
||||
parameters_schema TEXT NOT NULL,
|
||||
source_url TEXT,
|
||||
trust_level TEXT NOT NULL DEFAULT 'user',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (user_id, name, version)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_wasm_tools_user ON wasm_tools(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_wasm_tools_name ON wasm_tools(user_id, name);
|
||||
CREATE INDEX IF NOT EXISTS idx_wasm_tools_status ON wasm_tools(status);
|
||||
|
||||
-- ==================== Tool Capabilities ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tool_capabilities (
|
||||
id TEXT PRIMARY KEY,
|
||||
wasm_tool_id TEXT NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
|
||||
http_allowlist TEXT NOT NULL DEFAULT '[]',
|
||||
allowed_secrets TEXT NOT NULL DEFAULT '[]',
|
||||
tool_aliases TEXT NOT NULL DEFAULT '{}',
|
||||
requests_per_minute INTEGER NOT NULL DEFAULT 60,
|
||||
requests_per_hour INTEGER NOT NULL DEFAULT 1000,
|
||||
max_request_body_bytes INTEGER NOT NULL DEFAULT 1048576,
|
||||
max_response_body_bytes INTEGER NOT NULL DEFAULT 10485760,
|
||||
workspace_read_prefixes TEXT NOT NULL DEFAULT '[]',
|
||||
http_timeout_secs INTEGER NOT NULL DEFAULT 30,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (wasm_tool_id)
|
||||
);
|
||||
|
||||
-- ==================== Leak Detection Patterns ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS leak_detection_patterns (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
pattern TEXT NOT NULL,
|
||||
severity TEXT NOT NULL DEFAULT 'high',
|
||||
action TEXT NOT NULL DEFAULT 'block',
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
-- ==================== Rate Limit State ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tool_rate_limit_state (
|
||||
id TEXT PRIMARY KEY,
|
||||
wasm_tool_id TEXT NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
|
||||
user_id TEXT NOT NULL,
|
||||
minute_window_start TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
minute_count INTEGER NOT NULL DEFAULT 0,
|
||||
hour_window_start TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
hour_count INTEGER NOT NULL DEFAULT 0,
|
||||
UNIQUE (wasm_tool_id, user_id)
|
||||
);
|
||||
|
||||
-- ==================== Secret Usage Audit Log ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS secret_usage_log (
|
||||
id TEXT PRIMARY KEY,
|
||||
secret_id TEXT NOT NULL REFERENCES secrets(id) ON DELETE CASCADE,
|
||||
wasm_tool_id TEXT REFERENCES wasm_tools(id) ON DELETE SET NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
target_host TEXT NOT NULL,
|
||||
target_path TEXT,
|
||||
success INTEGER NOT NULL,
|
||||
error_message TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_secret_usage_user ON secret_usage_log(user_id);
|
||||
|
||||
-- ==================== Leak Detection Events ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS leak_detection_events (
|
||||
id TEXT PRIMARY KEY,
|
||||
pattern_id TEXT REFERENCES leak_detection_patterns(id) ON DELETE SET NULL,
|
||||
wasm_tool_id TEXT REFERENCES wasm_tools(id) ON DELETE SET NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
action_taken TEXT NOT NULL,
|
||||
context_preview TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
-- ==================== Tool Failures ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tool_failures (
|
||||
id TEXT PRIMARY KEY,
|
||||
tool_name TEXT NOT NULL UNIQUE,
|
||||
error_message TEXT,
|
||||
error_count INTEGER DEFAULT 1,
|
||||
first_failure TEXT DEFAULT (datetime('now')),
|
||||
last_failure TEXT DEFAULT (datetime('now')),
|
||||
last_build_result TEXT,
|
||||
repaired_at TEXT,
|
||||
repair_attempts INTEGER DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_failures_name ON tool_failures(tool_name);
|
||||
|
||||
-- ==================== Job Events ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS job_events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
job_id TEXT NOT NULL REFERENCES agent_jobs(id),
|
||||
event_type TEXT NOT NULL,
|
||||
data TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_job_events_job ON job_events(job_id, id);
|
||||
|
||||
-- ==================== Routines ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routines (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
user_id TEXT NOT NULL,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
trigger_type TEXT NOT NULL,
|
||||
trigger_config TEXT NOT NULL,
|
||||
action_type TEXT NOT NULL,
|
||||
action_config TEXT NOT NULL,
|
||||
cooldown_secs INTEGER NOT NULL DEFAULT 300,
|
||||
max_concurrent INTEGER NOT NULL DEFAULT 1,
|
||||
dedup_window_secs INTEGER,
|
||||
notify_channel TEXT,
|
||||
notify_user TEXT NOT NULL DEFAULT 'default',
|
||||
notify_on_success INTEGER NOT NULL DEFAULT 0,
|
||||
notify_on_failure INTEGER NOT NULL DEFAULT 1,
|
||||
notify_on_attention INTEGER NOT NULL DEFAULT 1,
|
||||
state TEXT NOT NULL DEFAULT '{}',
|
||||
last_run_at TEXT,
|
||||
next_fire_at TEXT,
|
||||
run_count INTEGER NOT NULL DEFAULT 0,
|
||||
consecutive_failures INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
UNIQUE (user_id, name)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_routines_user ON routines(user_id);
|
||||
|
||||
-- ==================== Routine Runs ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routine_runs (
|
||||
id TEXT PRIMARY KEY,
|
||||
routine_id TEXT NOT NULL REFERENCES routines(id) ON DELETE CASCADE,
|
||||
trigger_type TEXT NOT NULL,
|
||||
trigger_detail TEXT,
|
||||
started_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
completed_at TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'running',
|
||||
result_summary TEXT,
|
||||
tokens_used INTEGER,
|
||||
job_id TEXT REFERENCES agent_jobs(id),
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_routine_runs_routine ON routine_runs(routine_id);
|
||||
|
||||
-- ==================== Settings ====================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS settings (
|
||||
user_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
value TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
PRIMARY KEY (user_id, key)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_settings_user ON settings(user_id);
|
||||
|
||||
-- ==================== Missing indexes (parity with PostgreSQL) ====================
|
||||
|
||||
-- agent_jobs
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_stuck ON agent_jobs(stuck_since);
|
||||
|
||||
-- secrets
|
||||
CREATE INDEX IF NOT EXISTS idx_secrets_provider ON secrets(provider);
|
||||
CREATE INDEX IF NOT EXISTS idx_secrets_expires ON secrets(expires_at);
|
||||
|
||||
-- wasm_tools
|
||||
CREATE INDEX IF NOT EXISTS idx_wasm_tools_trust ON wasm_tools(trust_level);
|
||||
|
||||
-- tool_capabilities
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_capabilities_tool ON tool_capabilities(wasm_tool_id);
|
||||
|
||||
-- leak_detection_patterns
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_patterns_enabled ON leak_detection_patterns(enabled);
|
||||
|
||||
-- tool_rate_limit_state
|
||||
CREATE INDEX IF NOT EXISTS idx_rate_limit_tool ON tool_rate_limit_state(wasm_tool_id);
|
||||
|
||||
-- secret_usage_log
|
||||
CREATE INDEX IF NOT EXISTS idx_secret_usage_secret ON secret_usage_log(secret_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_secret_usage_tool ON secret_usage_log(wasm_tool_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_secret_usage_created ON secret_usage_log(created_at DESC);
|
||||
|
||||
-- leak_detection_events
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_events_pattern ON leak_detection_events(pattern_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_events_tool ON leak_detection_events(wasm_tool_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_events_user ON leak_detection_events(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_leak_events_created ON leak_detection_events(created_at DESC);
|
||||
|
||||
-- tool_failures
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_failures_count ON tool_failures(error_count DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_failures_unrepaired ON tool_failures(tool_name);
|
||||
|
||||
-- routines
|
||||
CREATE INDEX IF NOT EXISTS idx_routines_next_fire ON routines(next_fire_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_routines_event_triggers ON routines(user_id);
|
||||
|
||||
-- routine_runs
|
||||
CREATE INDEX IF NOT EXISTS idx_routine_runs_status ON routine_runs(status);
|
||||
|
||||
-- heartbeat_state
|
||||
CREATE INDEX IF NOT EXISTS idx_heartbeat_next_run ON heartbeat_state(next_run);
|
||||
|
||||
-- ==================== Seed data ====================
|
||||
|
||||
-- Pre-populate leak detection patterns (matches PostgreSQL V2 migration).
|
||||
INSERT OR IGNORE INTO leak_detection_patterns (id, name, pattern, severity, action, enabled, created_at) VALUES
|
||||
('550e8400-e29b-41d4-a716-446655440001', 'openai_api_key', 'sk-(?:proj-)?[a-zA-Z0-9]{20,}(?:T3BlbkFJ[a-zA-Z0-9_-]*)?', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440002', 'anthropic_api_key', 'sk-ant-api[a-zA-Z0-9_-]{90,}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440003', 'aws_access_key', 'AKIA[0-9A-Z]{16}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440004', 'aws_secret_key', '(?<![A-Za-z0-9/+=])[A-Za-z0-9/+=]{40}(?![A-Za-z0-9/+=])', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440005', 'github_token', 'gh[pousr]_[A-Za-z0-9_]{36,}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440006', 'github_fine_grained_pat', 'github_pat_[a-zA-Z0-9]{22}_[a-zA-Z0-9]{59}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440007', 'stripe_api_key', 'sk_(?:live|test)_[a-zA-Z0-9]{24,}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440008', 'nearai_session', 'sess_[a-zA-Z0-9]{32,}', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440009', 'bearer_token', 'Bearer\s+[a-zA-Z0-9_-]{20,}', 'high', 'redact', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000a', 'pem_private_key', '-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000b', 'ssh_private_key', '-----BEGIN\s+(?:OPENSSH|EC|DSA)\s+PRIVATE\s+KEY-----', 'critical', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000c', 'google_api_key', 'AIza[0-9A-Za-z_-]{35}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000d', 'slack_token', 'xox[baprs]-[0-9a-zA-Z-]{10,}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000e', 'discord_token', '[MN][A-Za-z\d]{23,}\.[\w-]{6}\.[\w-]{27}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-44665544000f', 'twilio_api_key', 'SK[a-fA-F0-9]{32}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440010', 'sendgrid_api_key', 'SG\.[a-zA-Z0-9_-]{22}\.[a-zA-Z0-9_-]{43}', 'high', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440011', 'mailchimp_api_key', '[a-f0-9]{32}-us[0-9]{1,2}', 'medium', 'block', 1, datetime('now')),
|
||||
('550e8400-e29b-41d4-a716-446655440012', 'high_entropy_hex', '(?<![a-fA-F0-9])[a-fA-F0-9]{64}(?![a-fA-F0-9])', 'medium', 'warn', 1, datetime('now'));
|
||||
|
||||
"#;
|
||||
+538
@@ -0,0 +1,538 @@
|
||||
//! Database abstraction layer.
|
||||
//!
|
||||
//! Provides a backend-agnostic `Database` trait that unifies all persistence
|
||||
//! operations. Two implementations exist behind feature flags:
|
||||
//!
|
||||
//! - `postgres` (default): Uses `deadpool-postgres` + `tokio-postgres`
|
||||
//! - `libsql`: Uses libSQL (Turso's SQLite fork) for embedded/edge deployment
|
||||
//!
|
||||
//! The existing `Store`, `Repository`, `SecretsStore`, and `WasmToolStore`
|
||||
//! types become thin wrappers that delegate to `Arc<dyn Database>`.
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
pub mod postgres;
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
pub mod libsql_backend;
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
pub mod libsql_migrations;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use rust_decimal::Decimal;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::BrokenTool;
|
||||
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
|
||||
use crate::context::{ActionRecord, JobContext, JobState};
|
||||
use crate::error::DatabaseError;
|
||||
use crate::error::WorkspaceError;
|
||||
use crate::history::{
|
||||
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
|
||||
SandboxJobSummary, SettingRow,
|
||||
};
|
||||
use crate::workspace::{MemoryChunk, MemoryDocument, WorkspaceEntry};
|
||||
use crate::workspace::{SearchConfig, SearchResult};
|
||||
|
||||
/// Create a database backend from configuration, run migrations, and return it.
|
||||
///
|
||||
/// This is the shared helper for CLI commands and other call sites that need
|
||||
/// a simple `Arc<dyn Database>` without retaining backend-specific handles
|
||||
/// (e.g., `pg_pool` or `libsql_conn` for the secrets store). The main agent
|
||||
/// startup in `main.rs` uses its own initialization block because it also
|
||||
/// captures those backend-specific handles.
|
||||
pub async fn connect_from_config(
|
||||
config: &crate::config::DatabaseConfig,
|
||||
) -> Result<Arc<dyn Database>, DatabaseError> {
|
||||
match config.backend {
|
||||
#[cfg(feature = "libsql")]
|
||||
crate::config::DatabaseBackend::LibSql => {
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config.libsql_path.as_deref().unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.libsql_url {
|
||||
let token = config.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
DatabaseError::Pool(
|
||||
"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()))?
|
||||
} else {
|
||||
libsql_backend::LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?
|
||||
};
|
||||
backend.run_migrations().await?;
|
||||
Ok(Arc::new(backend))
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
_ => {
|
||||
let pg = postgres::PgBackend::new(config)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
|
||||
pg.run_migrations().await?;
|
||||
Ok(Arc::new(pg))
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
_ => Err(DatabaseError::Pool(
|
||||
"No database backend available. Enable 'postgres' or 'libsql' feature.".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
#[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.
|
||||
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,
|
||||
channel: &str,
|
||||
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 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,
|
||||
category: &str,
|
||||
tool_names: &[String],
|
||||
estimated_cost: Decimal,
|
||||
estimated_time_secs: i32,
|
||||
estimated_value: Decimal,
|
||||
) -> Result<Uuid, DatabaseError>;
|
||||
|
||||
/// Update estimation snapshot with actual values.
|
||||
async fn update_estimation_actuals(
|
||||
&self,
|
||||
id: Uuid,
|
||||
actual_cost: Decimal,
|
||||
actual_time_secs: i32,
|
||||
actual_value: Option<Decimal>,
|
||||
) -> Result<(), DatabaseError>;
|
||||
|
||||
// ==================== Sandbox Jobs ====================
|
||||
|
||||
/// Insert a new sandbox job.
|
||||
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,
|
||||
status: &str,
|
||||
success: Option<bool>,
|
||||
message: Option<&str>,
|
||||
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 all job events.
|
||||
async fn list_job_events(&self, job_id: Uuid) -> Result<Vec<JobEventRecord>, DatabaseError>;
|
||||
|
||||
// ==================== Routines ====================
|
||||
|
||||
/// Create a new routine.
|
||||
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,
|
||||
last_run_at: DateTime<Utc>,
|
||||
next_fire_at: Option<DateTime<Utc>>,
|
||||
run_count: u64,
|
||||
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,
|
||||
status: RunStatus,
|
||||
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 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 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 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,
|
||||
chunk_index: i32,
|
||||
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,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
config: &SearchConfig,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError>;
|
||||
}
|
||||
@@ -0,0 +1,627 @@
|
||||
//! PostgreSQL backend for the Database trait.
|
||||
//!
|
||||
//! Delegates to the existing `Store` (history) and `Repository` (workspace)
|
||||
//! implementations, avoiding SQL duplication.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use deadpool_postgres::Pool;
|
||||
use rust_decimal::Decimal;
|
||||
use uuid::Uuid;
|
||||
|
||||
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::error::{DatabaseError, WorkspaceError};
|
||||
use crate::history::{
|
||||
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
|
||||
SandboxJobSummary, SettingRow, Store,
|
||||
};
|
||||
use crate::workspace::{
|
||||
MemoryChunk, MemoryDocument, Repository, SearchConfig, SearchResult, WorkspaceEntry,
|
||||
};
|
||||
|
||||
/// PostgreSQL database backend.
|
||||
///
|
||||
/// Wraps the existing `Store` (for history/conversations/jobs/routines/settings)
|
||||
/// and `Repository` (for workspace documents/chunks/search) to implement the
|
||||
/// unified `Database` trait.
|
||||
pub struct PgBackend {
|
||||
store: Store,
|
||||
repo: Repository,
|
||||
}
|
||||
|
||||
impl PgBackend {
|
||||
/// Create a new PostgreSQL backend from configuration.
|
||||
pub async fn new(config: &DatabaseConfig) -> Result<Self, DatabaseError> {
|
||||
let store = Store::new(config).await?;
|
||||
let repo = Repository::new(store.pool());
|
||||
Ok(Self { store, repo })
|
||||
}
|
||||
|
||||
/// Get a clone of the connection pool.
|
||||
///
|
||||
/// Useful for sharing with components that still need raw pool access.
|
||||
pub fn pool(&self) -> Pool {
|
||||
self.store.pool()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Database for PgBackend {
|
||||
async fn run_migrations(&self) -> Result<(), DatabaseError> {
|
||||
self.store.run_migrations().await
|
||||
}
|
||||
|
||||
// ==================== Conversations ====================
|
||||
|
||||
async fn create_conversation(
|
||||
&self,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.create_conversation(channel, user_id, thread_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> {
|
||||
self.store.touch_conversation(id).await
|
||||
}
|
||||
|
||||
async fn add_conversation_message(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
role: &str,
|
||||
content: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.add_conversation_message(conversation_id, role, content)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn ensure_conversation(
|
||||
&self,
|
||||
id: Uuid,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.ensure_conversation(id, channel, user_id, thread_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_conversations_with_preview(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
limit: i64,
|
||||
) -> Result<Vec<ConversationSummary>, DatabaseError> {
|
||||
self.store
|
||||
.list_conversations_with_preview(user_id, channel, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_or_create_assistant_conversation(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.get_or_create_assistant_conversation(user_id, channel)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn create_conversation_with_metadata(
|
||||
&self,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.create_conversation_with_metadata(channel, user_id, metadata)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_conversation_messages_paginated(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
before: Option<DateTime<Utc>>,
|
||||
limit: i64,
|
||||
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError> {
|
||||
self.store
|
||||
.list_conversation_messages_paginated(conversation_id, before, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_conversation_metadata_field(
|
||||
&self,
|
||||
id: Uuid,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_conversation_metadata_field(id, key, value)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_conversation_metadata(
|
||||
&self,
|
||||
id: Uuid,
|
||||
) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||
self.store.get_conversation_metadata(id).await
|
||||
}
|
||||
|
||||
async fn list_conversation_messages(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
) -> Result<Vec<ConversationMessage>, DatabaseError> {
|
||||
self.store.list_conversation_messages(conversation_id).await
|
||||
}
|
||||
|
||||
async fn conversation_belongs_to_user(
|
||||
&self,
|
||||
conversation_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
self.store
|
||||
.conversation_belongs_to_user(conversation_id, user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Jobs ====================
|
||||
|
||||
async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> {
|
||||
self.store.save_job(ctx).await
|
||||
}
|
||||
|
||||
async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
|
||||
self.store.get_job(id).await
|
||||
}
|
||||
|
||||
async fn update_job_status(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: JobState,
|
||||
failure_reason: Option<&str>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_job_status(id, status, failure_reason)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> {
|
||||
self.store.mark_job_stuck(id).await
|
||||
}
|
||||
|
||||
async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError> {
|
||||
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
|
||||
}
|
||||
|
||||
async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError> {
|
||||
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,
|
||||
category: &str,
|
||||
tool_names: &[String],
|
||||
estimated_cost: Decimal,
|
||||
estimated_time_secs: i32,
|
||||
estimated_value: Decimal,
|
||||
) -> Result<Uuid, DatabaseError> {
|
||||
self.store
|
||||
.save_estimation_snapshot(
|
||||
job_id,
|
||||
category,
|
||||
tool_names,
|
||||
estimated_cost,
|
||||
estimated_time_secs,
|
||||
estimated_value,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_estimation_actuals(
|
||||
&self,
|
||||
id: Uuid,
|
||||
actual_cost: Decimal,
|
||||
actual_time_secs: i32,
|
||||
actual_value: Option<Decimal>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Sandbox Jobs ====================
|
||||
|
||||
async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
|
||||
self.store.save_sandbox_job(job).await
|
||||
}
|
||||
|
||||
async fn get_sandbox_job(&self, id: Uuid) -> Result<Option<SandboxJobRecord>, DatabaseError> {
|
||||
self.store.get_sandbox_job(id).await
|
||||
}
|
||||
|
||||
async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
|
||||
self.store.list_sandbox_jobs().await
|
||||
}
|
||||
|
||||
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> {
|
||||
self.store
|
||||
.update_sandbox_job_status(id, status, success, message, started_at, completed_at)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError> {
|
||||
self.store.cleanup_stale_sandbox_jobs().await
|
||||
}
|
||||
|
||||
async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError> {
|
||||
self.store.sandbox_job_summary().await
|
||||
}
|
||||
|
||||
async fn list_sandbox_jobs_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
|
||||
self.store.list_sandbox_jobs_for_user(user_id).await
|
||||
}
|
||||
|
||||
async fn sandbox_job_summary_for_user(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<SandboxJobSummary, DatabaseError> {
|
||||
self.store.sandbox_job_summary_for_user(user_id).await
|
||||
}
|
||||
|
||||
async fn sandbox_job_belongs_to_user(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
user_id: &str,
|
||||
) -> Result<bool, DatabaseError> {
|
||||
self.store
|
||||
.sandbox_job_belongs_to_user(job_id, user_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> {
|
||||
self.store.update_sandbox_job_mode(id, mode).await
|
||||
}
|
||||
|
||||
async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError> {
|
||||
self.store.get_sandbox_job_mode(id).await
|
||||
}
|
||||
|
||||
// ==================== Job Events ====================
|
||||
|
||||
async fn save_job_event(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
event_type: &str,
|
||||
data: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store.save_job_event(job_id, event_type, data).await
|
||||
}
|
||||
|
||||
async fn list_job_events(&self, job_id: Uuid) -> Result<Vec<JobEventRecord>, DatabaseError> {
|
||||
self.store.list_job_events(job_id).await
|
||||
}
|
||||
|
||||
// ==================== Routines ====================
|
||||
|
||||
async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
self.store.create_routine(routine).await
|
||||
}
|
||||
|
||||
async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
|
||||
self.store.get_routine(id).await
|
||||
}
|
||||
|
||||
async fn get_routine_by_name(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<Option<Routine>, DatabaseError> {
|
||||
self.store.get_routine_by_name(user_id, name).await
|
||||
}
|
||||
|
||||
async fn list_routines(&self, user_id: &str) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.store.list_routines(user_id).await
|
||||
}
|
||||
|
||||
async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.store.list_event_routines().await
|
||||
}
|
||||
|
||||
async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
|
||||
self.store.list_due_cron_routines().await
|
||||
}
|
||||
|
||||
async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
self.store.update_routine(routine).await
|
||||
}
|
||||
|
||||
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> {
|
||||
self.store
|
||||
.update_routine_runtime(
|
||||
id,
|
||||
last_run_at,
|
||||
next_fire_at,
|
||||
run_count,
|
||||
consecutive_failures,
|
||||
state,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError> {
|
||||
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
|
||||
}
|
||||
|
||||
async fn complete_routine_run(
|
||||
&self,
|
||||
id: Uuid,
|
||||
status: RunStatus,
|
||||
result_summary: Option<&str>,
|
||||
tokens_used: Option<i32>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.complete_routine_run(id, status, result_summary, tokens_used)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_routine_runs(
|
||||
&self,
|
||||
routine_id: Uuid,
|
||||
limit: i64,
|
||||
) -> Result<Vec<RoutineRun>, DatabaseError> {
|
||||
self.store.list_routine_runs(routine_id, limit).await
|
||||
}
|
||||
|
||||
async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError> {
|
||||
self.store.count_running_routine_runs(routine_id).await
|
||||
}
|
||||
|
||||
// ==================== Tool Failures ====================
|
||||
|
||||
async fn record_tool_failure(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
error_message: &str,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store
|
||||
.record_tool_failure(tool_name, error_message)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError> {
|
||||
self.store.get_broken_tools(threshold).await
|
||||
}
|
||||
|
||||
async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||
self.store.mark_tool_repaired(tool_name).await
|
||||
}
|
||||
|
||||
async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||
self.store.increment_repair_attempts(tool_name).await
|
||||
}
|
||||
|
||||
// ==================== Settings ====================
|
||||
|
||||
async fn get_setting(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
) -> Result<Option<serde_json::Value>, DatabaseError> {
|
||||
self.store.get_setting(user_id, key).await
|
||||
}
|
||||
|
||||
async fn get_setting_full(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
) -> Result<Option<SettingRow>, DatabaseError> {
|
||||
self.store.get_setting_full(user_id, key).await
|
||||
}
|
||||
|
||||
async fn set_setting(
|
||||
&self,
|
||||
user_id: &str,
|
||||
key: &str,
|
||||
value: &serde_json::Value,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store.set_setting(user_id, key, value).await
|
||||
}
|
||||
|
||||
async fn delete_setting(&self, user_id: &str, key: &str) -> Result<bool, DatabaseError> {
|
||||
self.store.delete_setting(user_id, key).await
|
||||
}
|
||||
|
||||
async fn list_settings(&self, user_id: &str) -> Result<Vec<SettingRow>, DatabaseError> {
|
||||
self.store.list_settings(user_id).await
|
||||
}
|
||||
|
||||
async fn get_all_settings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<HashMap<String, serde_json::Value>, DatabaseError> {
|
||||
self.store.get_all_settings(user_id).await
|
||||
}
|
||||
|
||||
async fn set_all_settings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
settings: &HashMap<String, serde_json::Value>,
|
||||
) -> Result<(), DatabaseError> {
|
||||
self.store.set_all_settings(user_id, settings).await
|
||||
}
|
||||
|
||||
async fn has_settings(&self, user_id: &str) -> Result<bool, DatabaseError> {
|
||||
self.store.has_settings(user_id).await
|
||||
}
|
||||
|
||||
// ==================== Workspace: Documents ====================
|
||||
|
||||
async fn get_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
self.repo
|
||||
.get_document_by_path(user_id, agent_id, path)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
|
||||
self.repo.get_document_by_id(id).await
|
||||
}
|
||||
|
||||
async fn get_or_create_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
self.repo
|
||||
.get_or_create_document_by_path(user_id, agent_id, path)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
|
||||
self.repo.update_document(id, content).await
|
||||
}
|
||||
|
||||
async fn delete_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<(), WorkspaceError> {
|
||||
self.repo
|
||||
.delete_document_by_path(user_id, agent_id, path)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_directory(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
directory: &str,
|
||||
) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
|
||||
self.repo.list_directory(user_id, agent_id, directory).await
|
||||
}
|
||||
|
||||
async fn list_all_paths(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
) -> Result<Vec<String>, WorkspaceError> {
|
||||
self.repo.list_all_paths(user_id, agent_id).await
|
||||
}
|
||||
|
||||
async fn list_documents(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
) -> Result<Vec<MemoryDocument>, WorkspaceError> {
|
||||
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
|
||||
}
|
||||
|
||||
async fn insert_chunk(
|
||||
&self,
|
||||
document_id: Uuid,
|
||||
chunk_index: i32,
|
||||
content: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
) -> Result<Uuid, WorkspaceError> {
|
||||
self.repo
|
||||
.insert_chunk(document_id, chunk_index, content, embedding)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update_chunk_embedding(
|
||||
&self,
|
||||
chunk_id: Uuid,
|
||||
embedding: &[f32],
|
||||
) -> Result<(), WorkspaceError> {
|
||||
self.repo.update_chunk_embedding(chunk_id, embedding).await
|
||||
}
|
||||
|
||||
async fn get_chunks_without_embeddings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
|
||||
self.repo
|
||||
.get_chunks_without_embeddings(user_id, agent_id, limit)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Workspace: Search ====================
|
||||
|
||||
async fn hybrid_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
config: &SearchConfig,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError> {
|
||||
self.repo
|
||||
.hybrid_search(user_id, agent_id, query, embedding, config)
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -87,14 +87,21 @@ pub enum DatabaseError {
|
||||
#[error("Serialization error: {0}")]
|
||||
Serialization(String),
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[error("PostgreSQL error: {0}")]
|
||||
Postgres(#[from] tokio_postgres::Error),
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[error("Pool build error: {0}")]
|
||||
PoolBuild(#[from] deadpool_postgres::BuildError),
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[error("Pool runtime error: {0}")]
|
||||
PoolRuntime(#[from] deadpool_postgres::PoolError),
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[error("LibSQL error: {0}")]
|
||||
LibSql(#[from] libsql::Error),
|
||||
}
|
||||
|
||||
/// Channel-related errors.
|
||||
|
||||
@@ -57,7 +57,7 @@ pub struct ExtensionManager {
|
||||
_tunnel_url: Option<String>,
|
||||
user_id: String,
|
||||
/// Optional database store for DB-backed MCP config.
|
||||
store: Option<Arc<crate::history::Store>>,
|
||||
store: Option<Arc<dyn crate::db::Database>>,
|
||||
}
|
||||
|
||||
impl ExtensionManager {
|
||||
@@ -71,7 +71,7 @@ impl ExtensionManager {
|
||||
wasm_channels_dir: PathBuf,
|
||||
tunnel_url: Option<String>,
|
||||
user_id: String,
|
||||
store: Option<Arc<crate::history::Store>>,
|
||||
store: Option<Arc<dyn crate::db::Database>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
registry: ExtensionRegistry::new(),
|
||||
@@ -351,7 +351,7 @@ impl ExtensionManager {
|
||||
) -> Result<crate::tools::mcp::config::McpServersFile, crate::tools::mcp::config::ConfigError>
|
||||
{
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::load_mcp_servers_from_db(store, &self.user_id).await
|
||||
crate::tools::mcp::config::load_mcp_servers_from_db(store.as_ref(), &self.user_id).await
|
||||
} else {
|
||||
crate::tools::mcp::config::load_mcp_servers().await
|
||||
}
|
||||
@@ -375,7 +375,8 @@ impl ExtensionManager {
|
||||
) -> Result<(), crate::tools::mcp::config::ConfigError> {
|
||||
config.validate()?;
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::add_mcp_server_db(store, &self.user_id, config).await
|
||||
crate::tools::mcp::config::add_mcp_server_db(store.as_ref(), &self.user_id, config)
|
||||
.await
|
||||
} else {
|
||||
crate::tools::mcp::config::add_mcp_server(config).await
|
||||
}
|
||||
@@ -386,7 +387,8 @@ impl ExtensionManager {
|
||||
name: &str,
|
||||
) -> Result<(), crate::tools::mcp::config::ConfigError> {
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::remove_mcp_server_db(store, &self.user_id, name).await
|
||||
crate::tools::mcp::config::remove_mcp_server_db(store.as_ref(), &self.user_id, name)
|
||||
.await
|
||||
} else {
|
||||
crate::tools::mcp::config::remove_mcp_server(name).await
|
||||
}
|
||||
|
||||
+5
-1
@@ -5,11 +5,15 @@
|
||||
//! - Learning from past executions
|
||||
//! - Analytics and metrics
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
mod analytics;
|
||||
mod store;
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use analytics::{JobStats, ToolStats};
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use store::Store;
|
||||
pub use store::{
|
||||
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
|
||||
SandboxJobSummary, Store,
|
||||
SandboxJobSummary, SettingRow,
|
||||
};
|
||||
|
||||
@@ -1,13 +1,18 @@
|
||||
//! PostgreSQL store for persisting agent data.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
#[cfg(feature = "postgres")]
|
||||
use deadpool_postgres::{Config, Pool, Runtime};
|
||||
use rust_decimal::Decimal;
|
||||
#[cfg(feature = "postgres")]
|
||||
use tokio_postgres::NoTls;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::config::DatabaseConfig;
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::context::{ActionRecord, JobContext, JobState};
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::error::DatabaseError;
|
||||
|
||||
/// Record for an LLM call to be persisted.
|
||||
@@ -24,10 +29,12 @@ pub struct LlmCallRecord<'a> {
|
||||
}
|
||||
|
||||
/// Database store for the agent.
|
||||
#[cfg(feature = "postgres")]
|
||||
pub struct Store {
|
||||
pool: Pool,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Create a new store and connect to the database.
|
||||
pub async fn new(config: &DatabaseConfig) -> Result<Self, DatabaseError> {
|
||||
@@ -144,7 +151,12 @@ impl Store {
|
||||
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,
|
||||
@@ -466,6 +478,7 @@ pub struct SandboxJobSummary {
|
||||
pub interrupted: usize,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Insert a new sandbox job into `agent_jobs`.
|
||||
pub async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
|
||||
@@ -742,6 +755,7 @@ pub struct JobEventRecord {
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Persist a job event (fire-and-forget from orchestrator handler).
|
||||
pub async fn save_job_event(
|
||||
@@ -814,10 +828,12 @@ impl Store {
|
||||
|
||||
// ==================== Routines ====================
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::agent::routine::{
|
||||
NotifyConfig, Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger,
|
||||
};
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Create a new routine.
|
||||
pub async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
|
||||
@@ -1118,6 +1134,7 @@ impl Store {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_routine(row: &tokio_postgres::Row) -> Result<Routine, DatabaseError> {
|
||||
let trigger_type: String = row.get("trigger_type");
|
||||
let trigger_config: serde_json::Value = row.get("trigger_config");
|
||||
@@ -1162,6 +1179,7 @@ fn row_to_routine(row: &tokio_postgres::Row) -> Result<Routine, DatabaseError> {
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_routine_run(row: &tokio_postgres::Row) -> Result<RoutineRun, DatabaseError> {
|
||||
let status_str: String = row.get("status");
|
||||
let status: RunStatus = status_str
|
||||
@@ -1207,6 +1225,7 @@ pub struct ConversationMessage {
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Ensure a conversation row exists for a given UUID.
|
||||
///
|
||||
@@ -1477,6 +1496,7 @@ impl Store {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn parse_job_state(s: &str) -> JobState {
|
||||
match s {
|
||||
"pending" => JobState::Pending,
|
||||
@@ -1493,8 +1513,10 @@ fn parse_job_state(s: &str) -> JobState {
|
||||
|
||||
// ==================== Tool Failures ====================
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::agent::BrokenTool;
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Record a tool failure (upsert: increment count if exists).
|
||||
pub async fn record_tool_failure(
|
||||
@@ -1588,6 +1610,7 @@ pub struct SettingRow {
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl Store {
|
||||
/// Get a single setting by key.
|
||||
pub async fn get_setting(
|
||||
|
||||
@@ -44,6 +44,7 @@ pub mod channels;
|
||||
pub mod cli;
|
||||
pub mod config;
|
||||
pub mod context;
|
||||
pub mod db;
|
||||
pub mod error;
|
||||
pub mod estimation;
|
||||
pub mod evaluation;
|
||||
|
||||
+2
-2
@@ -62,7 +62,7 @@ pub struct SessionManager {
|
||||
/// Prevents thundering herd during concurrent 401s.
|
||||
renewal_lock: Mutex<()>,
|
||||
/// Optional database store for persisting session to the settings table.
|
||||
store: RwLock<Option<Arc<crate::history::Store>>>,
|
||||
store: RwLock<Option<Arc<dyn crate::db::Database>>>,
|
||||
/// User ID for DB settings (default: "default").
|
||||
user_id: RwLock<String>,
|
||||
}
|
||||
@@ -125,7 +125,7 @@ impl SessionManager {
|
||||
/// When a store is attached, session tokens are saved to the `settings`
|
||||
/// table (key: `nearai.session_token`) in addition to the disk file.
|
||||
/// On load, DB is preferred over disk.
|
||||
pub async fn attach_store(&self, store: Arc<crate::history::Store>, user_id: &str) {
|
||||
pub async fn attach_store(&self, store: Arc<dyn crate::db::Database>, user_id: &str) {
|
||||
*self.store.write().await = Some(store);
|
||||
*self.user_id.write().await = user_id.to_string();
|
||||
|
||||
|
||||
+164
-56
@@ -17,13 +17,11 @@ use ironclaw::{
|
||||
web::log_layer::{LogBroadcaster, WebLogLayer},
|
||||
},
|
||||
cli::{
|
||||
Cli, Command, run_mcp_command, run_memory_command, run_pairing_command, run_status_command,
|
||||
run_tool_command,
|
||||
Cli, Command, run_mcp_command, run_pairing_command, run_status_command, run_tool_command,
|
||||
},
|
||||
config::Config,
|
||||
context::ContextManager,
|
||||
extensions::ExtensionManager,
|
||||
history::Store,
|
||||
llm::{SessionConfig, create_llm_provider, create_session_manager},
|
||||
orchestrator::{
|
||||
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
|
||||
@@ -31,8 +29,7 @@ use ironclaw::{
|
||||
},
|
||||
pairing::PairingStore,
|
||||
safety::SafetyLayer,
|
||||
secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore},
|
||||
setup::{SetupConfig, SetupWizard},
|
||||
secrets::SecretsStore,
|
||||
tools::{
|
||||
ToolRegistry,
|
||||
mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated},
|
||||
@@ -41,6 +38,14 @@ use ironclaw::{
|
||||
workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
|
||||
};
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
use ironclaw::secrets::LibSqlSecretsStore;
|
||||
#[cfg(feature = "postgres")]
|
||||
use ironclaw::secrets::PostgresSecretsStore;
|
||||
use ironclaw::secrets::SecretsCrypto;
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
use ironclaw::setup::{SetupConfig, SetupWizard};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let cli = Cli::parse();
|
||||
@@ -89,8 +94,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
let config = Config::from_env()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
let store = ironclaw::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
// Set up embeddings if available
|
||||
let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig {
|
||||
@@ -130,7 +133,14 @@ async fn main() -> anyhow::Result<()> {
|
||||
None
|
||||
};
|
||||
|
||||
return run_memory_command(mem_cmd.clone(), store.pool(), embeddings).await;
|
||||
// Create a Database-trait-backed workspace for the memory command
|
||||
let db: Arc<dyn ironclaw::db::Database> =
|
||||
ironclaw::db::connect_from_config(&config.database)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
return ironclaw::cli::run_memory_command_with_db(mem_cmd.clone(), db, embeddings)
|
||||
.await;
|
||||
}
|
||||
Some(Command::Pairing(pairing_cmd)) => {
|
||||
tracing_subscriber::fmt()
|
||||
@@ -210,15 +220,13 @@ async fn main() -> anyhow::Result<()> {
|
||||
model
|
||||
);
|
||||
|
||||
// Load allowed tools from config (env var or defaults).
|
||||
let claude_config = ironclaw::config::ClaudeCodeConfig::from_env();
|
||||
let config = ironclaw::worker::claude_bridge::ClaudeBridgeConfig {
|
||||
job_id: *job_id,
|
||||
orchestrator_url: orchestrator_url.clone(),
|
||||
max_turns: *max_turns,
|
||||
model: model.clone(),
|
||||
timeout: std::time::Duration::from_secs(1800),
|
||||
allowed_tools: claude_config.allowed_tools,
|
||||
allowed_tools: Vec::new(),
|
||||
};
|
||||
|
||||
let runtime = ironclaw::worker::ClaudeBridgeRuntime::new(config)
|
||||
@@ -238,12 +246,20 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Load .env before running onboarding wizard
|
||||
let _ = dotenvy::dotenv();
|
||||
|
||||
let config = SetupConfig {
|
||||
skip_auth: *skip_auth,
|
||||
channels_only: *channels_only,
|
||||
};
|
||||
let mut wizard = SetupWizard::with_config(config);
|
||||
wizard.run().await?;
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
{
|
||||
let config = SetupConfig {
|
||||
skip_auth: *skip_auth,
|
||||
channels_only: *channels_only,
|
||||
};
|
||||
let mut wizard = SetupWizard::with_config(config);
|
||||
wizard.run().await?;
|
||||
}
|
||||
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
|
||||
{
|
||||
let _ = (skip_auth, channels_only);
|
||||
eprintln!("Onboarding wizard requires the 'postgres' or 'libsql' feature.");
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
None | Some(Command::Run) => {
|
||||
@@ -255,6 +271,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
let _ = dotenvy::dotenv();
|
||||
|
||||
// Enhanced first-run detection
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
if !cli.no_onboard {
|
||||
if let Some(reason) = check_onboard_needed().await {
|
||||
println!("Onboarding needed: {}", reason);
|
||||
@@ -322,23 +339,86 @@ async fn main() -> anyhow::Result<()> {
|
||||
tracing::info!("Loaded configuration for agent: {}", config.agent.name);
|
||||
tracing::info!("LLM backend: {}", config.llm.backend);
|
||||
|
||||
// Initialize database store (optional for testing)
|
||||
let store = if cli.no_db {
|
||||
// Initialize database backend.
|
||||
//
|
||||
// Creates an `Arc<dyn Database>` that all consumers share.
|
||||
// Backend is selected by the `DATABASE_BACKEND` env var / config.
|
||||
//
|
||||
// NOTE: For simpler call sites (CLI commands, Memory handler) use the shared
|
||||
// helper `ironclaw::db::connect_from_config()`. This block is kept inline
|
||||
// because it also captures backend-specific handles (`pg_pool`, `libsql_db`)
|
||||
// needed by the secrets store.
|
||||
#[cfg(feature = "postgres")]
|
||||
let mut pg_pool: Option<deadpool_postgres::Pool> = None;
|
||||
#[cfg(feature = "libsql")]
|
||||
let mut libsql_db: Option<std::sync::Arc<libsql::Database>> = None;
|
||||
|
||||
let db: Option<Arc<dyn ironclaw::db::Database>> = if cli.no_db {
|
||||
tracing::warn!("Running without database connection");
|
||||
None
|
||||
} else {
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
tracing::info!("Database connected and migrations applied");
|
||||
match config.database.backend {
|
||||
#[cfg(feature = "libsql")]
|
||||
ironclaw::config::DatabaseBackend::LibSql => {
|
||||
use ironclaw::db::Database as _;
|
||||
use ironclaw::db::libsql_backend::LibSqlBackend;
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = ironclaw::config::default_libsql_path();
|
||||
let db_path = config
|
||||
.database
|
||||
.libsql_path
|
||||
.as_deref()
|
||||
.unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.database.libsql_url {
|
||||
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
|
||||
})?;
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret()).await?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path).await?
|
||||
};
|
||||
backend.run_migrations().await?;
|
||||
tracing::info!("libSQL database connected and migrations applied");
|
||||
|
||||
// Capture the Database handle for SecretsStore (connection-per-op)
|
||||
libsql_db = Some(backend.shared_db());
|
||||
|
||||
Some(Arc::new(backend) as Arc<dyn ironclaw::db::Database>)
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
_ => {
|
||||
use ironclaw::db::Database as _;
|
||||
let pg = ironclaw::db::postgres::PgBackend::new(&config.database)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
pg.run_migrations()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
tracing::info!("PostgreSQL database connected and migrations applied");
|
||||
|
||||
pg_pool = Some(pg.pool());
|
||||
Some(Arc::new(pg) as Arc<dyn ironclaw::db::Database>)
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
_ => {
|
||||
anyhow::bail!(
|
||||
"No database backend available. Enable 'postgres' or 'libsql' feature."
|
||||
);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Post-init operations using the database
|
||||
if let Some(ref db) = db {
|
||||
// One-time migration: move disk config files into the DB settings table.
|
||||
if let Err(e) = ironclaw::bootstrap::migrate_disk_to_db(&store, "default").await {
|
||||
if let Err(e) = ironclaw::bootstrap::migrate_disk_to_db(db.as_ref(), "default").await {
|
||||
tracing::warn!("Disk-to-DB settings migration failed: {}", e);
|
||||
}
|
||||
|
||||
// Reload config from DB now that we have a connection.
|
||||
// Priority: env var > DB setting > default.
|
||||
match Config::from_db(&store, "default", &bootstrap).await {
|
||||
match Config::from_db(db.as_ref(), "default", &bootstrap).await {
|
||||
Ok(db_config) => {
|
||||
config = db_config;
|
||||
tracing::info!("Configuration reloaded from database");
|
||||
@@ -351,19 +431,14 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
let store = Arc::new(store);
|
||||
|
||||
// Attach store to session manager so tokens save to DB too
|
||||
session.attach_store(Arc::clone(&store), "default").await;
|
||||
// Attach DB to session manager so tokens save to DB too
|
||||
session.attach_store(Arc::clone(db), "default").await;
|
||||
|
||||
// Mark any jobs left in "running" or "creating" state as "interrupted".
|
||||
if let Err(e) = store.cleanup_stale_sandbox_jobs().await {
|
||||
if let Err(e) = db.cleanup_stale_sandbox_jobs().await {
|
||||
tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e);
|
||||
}
|
||||
|
||||
Some(store)
|
||||
};
|
||||
|
||||
}
|
||||
// Initialize LLM provider (clone session so we can reuse it for embeddings)
|
||||
let llm = create_llm_provider(&config.llm, session.clone())?;
|
||||
tracing::info!("LLM provider initialized: {}", llm.model_name());
|
||||
@@ -417,8 +492,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
};
|
||||
|
||||
// Register memory tools if database is available
|
||||
if let Some(ref store) = store {
|
||||
let mut workspace = Workspace::new("default", store.pool());
|
||||
if let Some(ref db) = db {
|
||||
let mut workspace = Workspace::new_with_db("default", Arc::clone(db));
|
||||
if let Some(ref emb) = embeddings {
|
||||
workspace = workspace.with_embeddings(emb.clone());
|
||||
}
|
||||
@@ -441,20 +516,46 @@ async fn main() -> anyhow::Result<()> {
|
||||
tracing::info!("Builder mode enabled");
|
||||
}
|
||||
|
||||
// Create secrets store if master key is configured (needed for MCP auth and WASM channels)
|
||||
// Create secrets store if master key is configured (needed for MCP auth and WASM channels).
|
||||
//
|
||||
// When both `postgres` and `libsql` features are compiled, the runtime-selected
|
||||
// backend determines which store is created: whichever DB init branch ran will
|
||||
// have set its handle (pg_pool or libsql_db), and the or_else chain picks it up.
|
||||
let secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>> =
|
||||
if let (Some(store), Some(master_key)) = (&store, config.secrets.master_key()) {
|
||||
if let Some(master_key) = config.secrets.master_key() {
|
||||
match SecretsCrypto::new(master_key.clone()) {
|
||||
Ok(crypto) => Some(Arc::new(PostgresSecretsStore::new(
|
||||
store.pool(),
|
||||
Arc::new(crypto),
|
||||
))),
|
||||
Ok(crypto) => {
|
||||
let crypto = Arc::new(crypto);
|
||||
let store: Option<Arc<dyn SecretsStore + Send + Sync>> = None;
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
let store = store.or_else(|| {
|
||||
libsql_db.take().map(|db| {
|
||||
Arc::new(LibSqlSecretsStore::new(db, Arc::clone(&crypto)))
|
||||
as Arc<dyn SecretsStore + Send + Sync>
|
||||
})
|
||||
});
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
let store = store.or_else(|| {
|
||||
pg_pool.as_ref().map(|pool| {
|
||||
Arc::new(PostgresSecretsStore::new(pool.clone(), Arc::clone(&crypto)))
|
||||
as Arc<dyn SecretsStore + Send + Sync>
|
||||
})
|
||||
});
|
||||
|
||||
store
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to initialize secrets crypto: {}", e);
|
||||
#[cfg(feature = "libsql")]
|
||||
let _ = libsql_db.take();
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
#[cfg(feature = "libsql")]
|
||||
let _ = libsql_db.take();
|
||||
None
|
||||
};
|
||||
|
||||
@@ -518,8 +619,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
let mcp_servers_future = async {
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
let servers_result = if let Some(ref s) = store {
|
||||
load_mcp_servers_from_db(s, "default").await
|
||||
let servers_result = if let Some(ref d) = db {
|
||||
load_mcp_servers_from_db(d.as_ref(), "default").await
|
||||
} else {
|
||||
ironclaw::tools::mcp::config::load_mcp_servers().await
|
||||
};
|
||||
@@ -632,7 +733,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
config.channels.wasm_channels_dir.clone(),
|
||||
config.tunnel.public_url.clone(),
|
||||
"default".to_string(),
|
||||
store.clone(),
|
||||
db.clone(),
|
||||
));
|
||||
tools.register_extension_tools(Arc::clone(&manager));
|
||||
tracing::info!("Extension manager initialized with in-chat discovery tools");
|
||||
@@ -695,7 +796,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
token_store,
|
||||
job_event_tx: job_event_tx.clone(),
|
||||
prompt_queue: Arc::clone(&prompt_queue),
|
||||
store: store.clone(),
|
||||
store: db.clone(),
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
@@ -932,13 +1033,15 @@ async fn main() -> anyhow::Result<()> {
|
||||
};
|
||||
|
||||
// Create workspace for agent (shared with memory tools)
|
||||
let workspace = store.as_ref().map(|s| {
|
||||
let mut ws = Workspace::new("default", s.pool());
|
||||
let workspace = if let Some(ref db_ref) = db {
|
||||
let mut ws = Workspace::new_with_db("default", Arc::clone(db_ref));
|
||||
if let Some(ref emb) = embeddings {
|
||||
ws = ws.with_embeddings(emb.clone());
|
||||
}
|
||||
Arc::new(ws)
|
||||
});
|
||||
Some(Arc::new(ws))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Seed workspace with core identity files on first boot
|
||||
if let Some(ref ws) = workspace {
|
||||
@@ -976,7 +1079,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
tools.register_job_tools(
|
||||
Arc::clone(&context_manager),
|
||||
container_job_manager.clone(),
|
||||
store.clone(),
|
||||
db.clone(),
|
||||
);
|
||||
|
||||
// Add web gateway channel if configured
|
||||
@@ -991,13 +1094,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
if let Some(ref ext_mgr) = extension_manager {
|
||||
gw = gw.with_extension_manager(Arc::clone(ext_mgr));
|
||||
}
|
||||
if let Some(ref s) = store {
|
||||
gw = gw.with_store(Arc::clone(s));
|
||||
if let Some(ref d) = db {
|
||||
gw = gw.with_store(Arc::clone(d));
|
||||
}
|
||||
if let Some(ref jm) = container_job_manager {
|
||||
gw = gw.with_job_manager(Arc::clone(jm));
|
||||
}
|
||||
gw = gw.with_llm_provider(Arc::clone(&llm));
|
||||
if config.sandbox.enabled {
|
||||
gw = gw.with_prompt_queue(Arc::clone(&prompt_queue));
|
||||
|
||||
@@ -1030,7 +1132,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
// Create and run the agent
|
||||
let deps = AgentDeps {
|
||||
store,
|
||||
store: db,
|
||||
llm,
|
||||
safety,
|
||||
tools,
|
||||
@@ -1064,11 +1166,17 @@ async fn main() -> anyhow::Result<()> {
|
||||
/// Check if onboarding is needed and return the reason.
|
||||
///
|
||||
/// Returns `Some(reason)` if onboarding should be triggered, `None` otherwise.
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
async fn check_onboard_needed() -> Option<&'static str> {
|
||||
let bootstrap = ironclaw::bootstrap::BootstrapConfig::load();
|
||||
|
||||
// Database not configured (and not in env)
|
||||
if bootstrap.database_url.is_none() && std::env::var("DATABASE_URL").is_err() {
|
||||
let has_db = bootstrap.database_url.is_some()
|
||||
|| std::env::var("DATABASE_URL").is_ok()
|
||||
|| std::env::var("LIBSQL_PATH").is_ok()
|
||||
|| ironclaw::config::default_libsql_path().exists();
|
||||
|
||||
if !has_db {
|
||||
return Some("Database not configured");
|
||||
}
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ use tokio::sync::{Mutex, broadcast};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::types::SseEvent;
|
||||
use crate::history::Store;
|
||||
use crate::db::Database;
|
||||
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
||||
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
@@ -43,7 +43,7 @@ pub struct OrchestratorState {
|
||||
/// Buffered follow-up prompts for sandbox jobs, keyed by job_id.
|
||||
pub prompt_queue: Arc<Mutex<HashMap<Uuid, VecDeque<PendingPrompt>>>>,
|
||||
/// Database handle for persisting job events.
|
||||
pub store: Option<Arc<Store>>,
|
||||
pub store: Option<Arc<dyn Database>>,
|
||||
}
|
||||
|
||||
/// The orchestrator's internal API server.
|
||||
|
||||
+5
-1
@@ -64,7 +64,11 @@ mod store;
|
||||
mod types;
|
||||
|
||||
pub use crypto::SecretsCrypto;
|
||||
pub use store::{PostgresSecretsStore, SecretsStore};
|
||||
#[cfg(feature = "libsql")]
|
||||
pub use store::LibSqlSecretsStore;
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use store::PostgresSecretsStore;
|
||||
pub use store::SecretsStore;
|
||||
pub use types::{
|
||||
CreateSecretParams, CredentialLocation, CredentialMapping, DecryptedSecret, Secret,
|
||||
SecretError, SecretRef,
|
||||
|
||||
@@ -10,6 +10,7 @@ use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::Utc;
|
||||
#[cfg(feature = "postgres")]
|
||||
use deadpool_postgres::Pool;
|
||||
use secrecy::ExposeSecret;
|
||||
use uuid::Uuid;
|
||||
@@ -61,11 +62,13 @@ pub trait SecretsStore: Send + Sync {
|
||||
}
|
||||
|
||||
/// PostgreSQL implementation of SecretsStore.
|
||||
#[cfg(feature = "postgres")]
|
||||
pub struct PostgresSecretsStore {
|
||||
pool: Pool,
|
||||
crypto: Arc<SecretsCrypto>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl PostgresSecretsStore {
|
||||
/// Create a new store with the given database pool and crypto instance.
|
||||
pub fn new(pool: Pool, crypto: Arc<SecretsCrypto>) -> Self {
|
||||
@@ -73,6 +76,7 @@ impl PostgresSecretsStore {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[async_trait]
|
||||
impl SecretsStore for PostgresSecretsStore {
|
||||
async fn create(
|
||||
@@ -283,6 +287,7 @@ impl SecretsStore for PostgresSecretsStore {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_secret(row: &tokio_postgres::Row) -> Secret {
|
||||
Secret {
|
||||
id: row.get("id"),
|
||||
@@ -299,6 +304,332 @@ fn row_to_secret(row: &tokio_postgres::Row) -> Secret {
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== libSQL implementation ====================
|
||||
|
||||
/// libSQL/Turso implementation of SecretsStore.
|
||||
///
|
||||
/// Holds an `Arc<Database>` handle and creates a fresh connection per operation,
|
||||
/// matching the connection-per-request pattern used by the main `LibSqlBackend`.
|
||||
#[cfg(feature = "libsql")]
|
||||
pub struct LibSqlSecretsStore {
|
||||
db: Arc<libsql::Database>,
|
||||
crypto: Arc<SecretsCrypto>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
impl LibSqlSecretsStore {
|
||||
/// Create a new store with the given shared libsql database handle and crypto instance.
|
||||
pub fn new(db: Arc<libsql::Database>, crypto: Arc<SecretsCrypto>) -> Self {
|
||||
Self { db, crypto }
|
||||
}
|
||||
|
||||
fn connect(&self) -> Result<libsql::Connection, SecretError> {
|
||||
self.db
|
||||
.connect()
|
||||
.map_err(|e| SecretError::Database(format!("Connection failed: {}", e)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[async_trait]
|
||||
impl SecretsStore for LibSqlSecretsStore {
|
||||
async fn create(
|
||||
&self,
|
||||
user_id: &str,
|
||||
params: CreateSecretParams,
|
||||
) -> Result<Secret, SecretError> {
|
||||
let plaintext = params.value.expose_secret().as_bytes();
|
||||
let (encrypted_value, key_salt) = self.crypto.encrypt(plaintext)?;
|
||||
|
||||
let id = Uuid::new_v4();
|
||||
let now = Utc::now();
|
||||
let now_str = now.to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
|
||||
let expires_at_str = params
|
||||
.expires_at
|
||||
.map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true));
|
||||
|
||||
// Start transaction for atomic upsert + read-back
|
||||
let conn = self.connect()?;
|
||||
let tx = conn
|
||||
.transaction()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
tx.execute(
|
||||
r#"
|
||||
INSERT INTO secrets (id, user_id, name, encrypted_value, key_salt, provider, expires_at, created_at, updated_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?8)
|
||||
ON CONFLICT (user_id, name) DO UPDATE SET
|
||||
encrypted_value = excluded.encrypted_value,
|
||||
key_salt = excluded.key_salt,
|
||||
provider = excluded.provider,
|
||||
expires_at = excluded.expires_at,
|
||||
updated_at = ?8
|
||||
"#,
|
||||
libsql::params![
|
||||
id.to_string(),
|
||||
user_id,
|
||||
params.name.as_str(),
|
||||
libsql::Value::Blob(encrypted_value.clone()),
|
||||
libsql::Value::Blob(key_salt.clone()),
|
||||
libsql_opt_text(params.provider.as_deref()),
|
||||
libsql_opt_text(expires_at_str.as_deref()),
|
||||
now_str.as_str(),
|
||||
],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
// Read back the row (may have been upserted)
|
||||
let mut rows = tx
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, encrypted_value, key_salt, provider, expires_at,
|
||||
last_used_at, usage_count, created_at, updated_at
|
||||
FROM secrets
|
||||
WHERE user_id = ?1 AND name = ?2
|
||||
"#,
|
||||
libsql::params![user_id, params.name.as_str()],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
let row = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?
|
||||
.ok_or_else(|| SecretError::Database("Insert succeeded but row not found".into()))?;
|
||||
|
||||
let secret = libsql_row_to_secret(&row)?;
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(secret)
|
||||
}
|
||||
|
||||
async fn get(&self, user_id: &str, name: &str) -> Result<Secret, SecretError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, encrypted_value, key_salt, provider, expires_at,
|
||||
last_used_at, usage_count, created_at, updated_at
|
||||
FROM secrets
|
||||
WHERE user_id = ?1 AND name = ?2
|
||||
"#,
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?
|
||||
{
|
||||
Some(row) => {
|
||||
let secret = libsql_row_to_secret(&row)?;
|
||||
|
||||
if let Some(expires_at) = secret.expires_at {
|
||||
if expires_at < Utc::now() {
|
||||
return Err(SecretError::Expired);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(secret)
|
||||
}
|
||||
None => Err(SecretError::NotFound(name.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_decrypted(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<DecryptedSecret, SecretError> {
|
||||
let secret = self.get(user_id, name).await?;
|
||||
self.crypto
|
||||
.decrypt(&secret.encrypted_value, &secret.key_salt)
|
||||
}
|
||||
|
||||
async fn exists(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
"SELECT 1 FROM secrets WHERE user_id = ?1 AND name = ?2",
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?
|
||||
.is_some())
|
||||
}
|
||||
|
||||
async fn list(&self, user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
"SELECT name, provider FROM secrets WHERE user_id = ?1 ORDER BY name",
|
||||
libsql::params![user_id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
let mut refs = Vec::new();
|
||||
while let Some(row) = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?
|
||||
{
|
||||
refs.push(SecretRef {
|
||||
name: row.get::<String>(0).unwrap_or_default(),
|
||||
provider: row.get::<String>(1).ok(),
|
||||
});
|
||||
}
|
||||
Ok(refs)
|
||||
}
|
||||
|
||||
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
|
||||
let conn = self.connect()?;
|
||||
let affected = conn
|
||||
.execute(
|
||||
"DELETE FROM secrets WHERE user_id = ?1 AND name = ?2",
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(affected > 0)
|
||||
}
|
||||
|
||||
async fn record_usage(&self, secret_id: Uuid) -> Result<(), SecretError> {
|
||||
let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
|
||||
let conn = self.connect()?;
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
UPDATE secrets
|
||||
SET last_used_at = ?1, usage_count = usage_count + 1
|
||||
WHERE id = ?2
|
||||
"#,
|
||||
libsql::params![now.as_str(), secret_id.to_string()],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn is_accessible(
|
||||
&self,
|
||||
user_id: &str,
|
||||
secret_name: &str,
|
||||
allowed_secrets: &[String],
|
||||
) -> Result<bool, SecretError> {
|
||||
if !self.exists(user_id, secret_name).await? {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
for pattern in allowed_secrets {
|
||||
if pattern == secret_name {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
if let Some(prefix) = pattern.strip_suffix('*') {
|
||||
if secret_name.starts_with(prefix) {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_opt_text(s: Option<&str>) -> libsql::Value {
|
||||
match s {
|
||||
Some(s) => libsql::Value::Text(s.to_string()),
|
||||
None => libsql::Value::Null,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_parse_timestamp(s: &str) -> Result<chrono::DateTime<Utc>, SecretError> {
|
||||
if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) {
|
||||
return Ok(dt.with_timezone(&Utc));
|
||||
}
|
||||
if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") {
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
Err(SecretError::Database(format!(
|
||||
"unparseable timestamp: {:?}",
|
||||
s
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_row_to_secret(row: &libsql::Row) -> Result<Secret, SecretError> {
|
||||
let id_str: String = row
|
||||
.get(0)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let user_id: String = row
|
||||
.get(1)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let name: String = row
|
||||
.get(2)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let encrypted_value: Vec<u8> = row
|
||||
.get(3)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let key_salt: Vec<u8> = row
|
||||
.get(4)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let provider: Option<String> = row.get::<String>(5).ok().filter(|s| !s.is_empty());
|
||||
let expires_at = row
|
||||
.get::<String>(6)
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.and_then(|s| libsql_parse_timestamp(&s).ok());
|
||||
let last_used_at = row
|
||||
.get::<String>(7)
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.and_then(|s| libsql_parse_timestamp(&s).ok());
|
||||
let usage_count: i64 = row.get::<i64>(8).unwrap_or(0);
|
||||
let created_at_str: String = row
|
||||
.get(9)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
let updated_at_str: String = row
|
||||
.get(10)
|
||||
.map_err(|e| SecretError::Database(e.to_string()))?;
|
||||
|
||||
Ok(Secret {
|
||||
id: id_str
|
||||
.parse()
|
||||
.map_err(|e: uuid::Error| SecretError::Database(e.to_string()))?,
|
||||
user_id,
|
||||
name,
|
||||
encrypted_value,
|
||||
key_salt,
|
||||
provider,
|
||||
expires_at,
|
||||
last_used_at,
|
||||
usage_count,
|
||||
created_at: libsql_parse_timestamp(&created_at_str)?,
|
||||
updated_at: libsql_parse_timestamp(&updated_at_str)?,
|
||||
})
|
||||
}
|
||||
|
||||
/// In-memory implementation for testing.
|
||||
#[cfg(test)]
|
||||
pub mod testing {
|
||||
|
||||
@@ -15,6 +15,10 @@ pub struct Settings {
|
||||
pub onboard_completed: bool,
|
||||
|
||||
// === Step 1: Database ===
|
||||
/// Database backend: "postgres" or "libsql".
|
||||
#[serde(default)]
|
||||
pub database_backend: Option<String>,
|
||||
|
||||
/// Database connection URL (postgres://...).
|
||||
#[serde(default)]
|
||||
pub database_url: Option<String>,
|
||||
@@ -23,6 +27,14 @@ pub struct Settings {
|
||||
#[serde(default)]
|
||||
pub database_pool_size: Option<usize>,
|
||||
|
||||
/// Path to local libSQL database file.
|
||||
#[serde(default)]
|
||||
pub libsql_path: Option<String>,
|
||||
|
||||
/// Turso cloud URL for remote replica sync.
|
||||
#[serde(default)]
|
||||
pub libsql_url: Option<String>,
|
||||
|
||||
// === Step 2: Security ===
|
||||
/// Source for the secrets master key.
|
||||
#[serde(default)]
|
||||
|
||||
+15
-4
@@ -12,7 +12,9 @@ use reqwest::Client;
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::secrets::{CreateSecretParams, PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::secrets::SecretsCrypto;
|
||||
use crate::secrets::{CreateSecretParams, SecretsStore};
|
||||
use crate::settings::Settings;
|
||||
use crate::setup::prompts::{
|
||||
confirm, input, optional_input, print_error, print_info, print_success, secret_input,
|
||||
@@ -20,15 +22,24 @@ use crate::setup::prompts::{
|
||||
|
||||
/// Context for saving secrets during setup.
|
||||
pub struct SecretsContext {
|
||||
store: PostgresSecretsStore,
|
||||
store: Arc<dyn SecretsStore>,
|
||||
user_id: String,
|
||||
}
|
||||
|
||||
impl SecretsContext {
|
||||
/// Create a new secrets context.
|
||||
/// Create a new secrets context from a trait-object store.
|
||||
pub fn from_store(store: Arc<dyn SecretsStore>, user_id: &str) -> Self {
|
||||
Self {
|
||||
store,
|
||||
user_id: user_id.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new secrets context from a PostgreSQL pool and crypto.
|
||||
#[cfg(feature = "postgres")]
|
||||
pub fn new(pool: deadpool_postgres::Pool, crypto: Arc<SecretsCrypto>, user_id: &str) -> Self {
|
||||
Self {
|
||||
store: PostgresSecretsStore::new(pool, crypto),
|
||||
store: Arc::new(crate::secrets::PostgresSecretsStore::new(pool, crypto)),
|
||||
user_id: user_id.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,6 +20,7 @@
|
||||
|
||||
mod channels;
|
||||
mod prompts;
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
mod wizard;
|
||||
|
||||
pub use channels::{
|
||||
@@ -29,4 +30,5 @@ pub use prompts::{
|
||||
confirm, input, optional_input, print_error, print_header, print_info, print_step,
|
||||
print_success, secret_input, select_many, select_one,
|
||||
};
|
||||
#[cfg(any(feature = "postgres", feature = "libsql"))]
|
||||
pub use wizard::{SetupConfig, SetupWizard};
|
||||
|
||||
+297
-37
@@ -12,15 +12,17 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
use deadpool_postgres::{Config as PoolConfig, Runtime};
|
||||
use secrecy::SecretString;
|
||||
#[cfg(feature = "postgres")]
|
||||
use tokio_postgres::NoTls;
|
||||
|
||||
use crate::channels::wasm::{
|
||||
ChannelCapabilitiesFile, available_channel_names, install_bundled_channel,
|
||||
};
|
||||
use crate::llm::{SessionConfig, SessionManager};
|
||||
use crate::secrets::SecretsCrypto;
|
||||
use crate::secrets::{SecretsCrypto, SecretsStore};
|
||||
use crate::settings::{KeySource, Settings};
|
||||
use crate::setup::channels::{
|
||||
SecretsContext, setup_http, setup_telegram, setup_tunnel, setup_wasm_channel,
|
||||
@@ -66,8 +68,12 @@ pub struct SetupWizard {
|
||||
config: SetupConfig,
|
||||
settings: Settings,
|
||||
session_manager: Option<Arc<SessionManager>>,
|
||||
/// Database pool (created during setup).
|
||||
/// Database pool (created during setup, postgres only).
|
||||
#[cfg(feature = "postgres")]
|
||||
db_pool: Option<deadpool_postgres::Pool>,
|
||||
/// libSQL backend (created during setup, libsql only).
|
||||
#[cfg(feature = "libsql")]
|
||||
db_backend: Option<crate::db::libsql_backend::LibSqlBackend>,
|
||||
/// Secrets crypto (created during setup).
|
||||
secrets_crypto: Option<Arc<SecretsCrypto>>,
|
||||
}
|
||||
@@ -79,7 +85,10 @@ impl SetupWizard {
|
||||
config: SetupConfig::default(),
|
||||
settings: Settings::load(),
|
||||
session_manager: None,
|
||||
#[cfg(feature = "postgres")]
|
||||
db_pool: None,
|
||||
#[cfg(feature = "libsql")]
|
||||
db_backend: None,
|
||||
secrets_crypto: None,
|
||||
}
|
||||
}
|
||||
@@ -90,7 +99,10 @@ impl SetupWizard {
|
||||
config,
|
||||
settings: Settings::load(),
|
||||
session_manager: None,
|
||||
#[cfg(feature = "postgres")]
|
||||
db_pool: None,
|
||||
#[cfg(feature = "libsql")]
|
||||
db_backend: None,
|
||||
secrets_crypto: None,
|
||||
}
|
||||
}
|
||||
@@ -153,19 +165,48 @@ impl SetupWizard {
|
||||
|
||||
/// Step 1: Database connection.
|
||||
async fn step_database(&mut self) -> Result<(), SetupError> {
|
||||
// Check if we have an existing URL in env or settings
|
||||
// Determine which backend to use based on compile-time features.
|
||||
// When both features are enabled, prefer the currently configured backend
|
||||
// or default to postgres.
|
||||
#[cfg(all(feature = "postgres", feature = "libsql"))]
|
||||
{
|
||||
let backend = std::env::var("DATABASE_BACKEND")
|
||||
.ok()
|
||||
.or_else(|| self.settings.database_backend.clone())
|
||||
.unwrap_or_else(|| "postgres".to_string());
|
||||
|
||||
if backend == "libsql" || backend == "turso" || backend == "sqlite" {
|
||||
return self.step_database_libsql().await;
|
||||
}
|
||||
return self.step_database_postgres().await;
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "postgres", not(feature = "libsql")))]
|
||||
{
|
||||
return self.step_database_postgres().await;
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
|
||||
{
|
||||
return self.step_database_libsql().await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Step 1 (postgres): Database connection via PostgreSQL URL.
|
||||
#[cfg(feature = "postgres")]
|
||||
async fn step_database_postgres(&mut self) -> Result<(), SetupError> {
|
||||
self.settings.database_backend = Some("postgres".to_string());
|
||||
|
||||
let existing_url = std::env::var("DATABASE_URL")
|
||||
.ok()
|
||||
.or_else(|| self.settings.database_url.clone());
|
||||
|
||||
if let Some(ref url) = existing_url {
|
||||
// Mask the password for display
|
||||
let display_url = mask_password_in_url(url);
|
||||
print_info(&format!("Existing database URL: {}", display_url));
|
||||
|
||||
if confirm("Use this database?", true).map_err(SetupError::Io)? {
|
||||
// Test the connection
|
||||
if let Err(e) = self.test_database_connection(url).await {
|
||||
if let Err(e) = self.test_database_connection_postgres(url).await {
|
||||
print_error(&format!("Connection failed: {}", e));
|
||||
print_info("Let's configure a new database URL.");
|
||||
} else {
|
||||
@@ -176,7 +217,6 @@ impl SetupWizard {
|
||||
}
|
||||
}
|
||||
|
||||
// Prompt for new URL
|
||||
println!();
|
||||
print_info("Enter your PostgreSQL connection URL.");
|
||||
print_info("Format: postgres://user:password@host:port/database");
|
||||
@@ -190,15 +230,13 @@ impl SetupWizard {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Test the connection
|
||||
print_info("Testing connection...");
|
||||
match self.test_database_connection(&url).await {
|
||||
match self.test_database_connection_postgres(&url).await {
|
||||
Ok(()) => {
|
||||
print_success("Database connection successful");
|
||||
|
||||
// Ask if we should run migrations
|
||||
if confirm("Run database migrations?", true).map_err(SetupError::Io)? {
|
||||
self.run_migrations().await?;
|
||||
self.run_migrations_postgres().await?;
|
||||
}
|
||||
|
||||
self.settings.database_url = Some(url);
|
||||
@@ -216,8 +254,115 @@ impl SetupWizard {
|
||||
}
|
||||
}
|
||||
|
||||
/// Test database connection and store the pool.
|
||||
async fn test_database_connection(&mut self, url: &str) -> Result<(), SetupError> {
|
||||
/// Step 1 (libsql): Database connection via local file or Turso remote replica.
|
||||
#[cfg(feature = "libsql")]
|
||||
async fn step_database_libsql(&mut self) -> Result<(), SetupError> {
|
||||
self.settings.database_backend = Some("libsql".to_string());
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let default_path_str = default_path.to_string_lossy().to_string();
|
||||
|
||||
// Check for existing configuration
|
||||
let existing_path = std::env::var("LIBSQL_PATH")
|
||||
.ok()
|
||||
.or_else(|| self.settings.libsql_path.clone());
|
||||
|
||||
if let Some(ref path) = existing_path {
|
||||
print_info(&format!("Existing database path: {}", path));
|
||||
if confirm("Use this database?", true).map_err(SetupError::Io)? {
|
||||
let turso_url = std::env::var("LIBSQL_URL")
|
||||
.ok()
|
||||
.or_else(|| self.settings.libsql_url.clone());
|
||||
let turso_token = std::env::var("LIBSQL_AUTH_TOKEN").ok();
|
||||
|
||||
match self
|
||||
.test_database_connection_libsql(
|
||||
path,
|
||||
turso_url.as_deref(),
|
||||
turso_token.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
print_success("Database connection successful");
|
||||
self.settings.libsql_path = Some(path.clone());
|
||||
if let Some(url) = turso_url {
|
||||
self.settings.libsql_url = Some(url);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
Err(e) => {
|
||||
print_error(&format!("Connection failed: {}", e));
|
||||
print_info("Let's configure a new database path.");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
println!();
|
||||
print_info("IronClaw uses an embedded SQLite database (libSQL).");
|
||||
print_info("No external database server required.");
|
||||
println!();
|
||||
|
||||
let path_input = optional_input(
|
||||
"Database file path",
|
||||
Some(&format!("default: {}", default_path_str)),
|
||||
)
|
||||
.map_err(SetupError::Io)?;
|
||||
|
||||
let db_path = path_input.unwrap_or(default_path_str.clone());
|
||||
|
||||
// Ask about Turso cloud sync
|
||||
println!();
|
||||
let use_turso =
|
||||
confirm("Enable Turso cloud sync (remote replica)?", false).map_err(SetupError::Io)?;
|
||||
|
||||
let (turso_url, turso_token) = if use_turso {
|
||||
print_info("Enter your Turso database URL and auth token.");
|
||||
print_info("Format: libsql://your-db.turso.io");
|
||||
println!();
|
||||
|
||||
let url = input("Turso URL").map_err(SetupError::Io)?;
|
||||
if url.is_empty() {
|
||||
print_error("Turso URL is required for cloud sync.");
|
||||
(None, None)
|
||||
} else {
|
||||
let token = input("Auth token").map_err(SetupError::Io)?;
|
||||
if token.is_empty() {
|
||||
print_error("Auth token is required for cloud sync.");
|
||||
(None, None)
|
||||
} else {
|
||||
(Some(url), Some(token))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
|
||||
print_info("Testing connection...");
|
||||
match self
|
||||
.test_database_connection_libsql(&db_path, turso_url.as_deref(), turso_token.as_deref())
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
print_success("Database connection successful");
|
||||
|
||||
// Always run migrations for libsql (they're idempotent)
|
||||
self.run_migrations_libsql().await?;
|
||||
|
||||
self.settings.libsql_path = Some(db_path);
|
||||
if let Some(url) = turso_url {
|
||||
self.settings.libsql_url = Some(url);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(SetupError::Database(format!("Connection failed: {}", e))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Test PostgreSQL connection and store the pool.
|
||||
#[cfg(feature = "postgres")]
|
||||
async fn test_database_connection_postgres(&mut self, url: &str) -> Result<(), SetupError> {
|
||||
let mut cfg = PoolConfig::new();
|
||||
cfg.url = Some(url.to_string());
|
||||
cfg.pool = Some(deadpool_postgres::PoolConfig {
|
||||
@@ -229,7 +374,6 @@ impl SetupWizard {
|
||||
.create_pool(Some(Runtime::Tokio1), NoTls)
|
||||
.map_err(|e| SetupError::Database(format!("Failed to create pool: {}", e)))?;
|
||||
|
||||
// Test the connection
|
||||
let _ = pool
|
||||
.get()
|
||||
.await
|
||||
@@ -239,8 +383,36 @@ impl SetupWizard {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run database migrations.
|
||||
async fn run_migrations(&self) -> Result<(), SetupError> {
|
||||
/// Test libSQL connection and store the backend.
|
||||
#[cfg(feature = "libsql")]
|
||||
async fn test_database_connection_libsql(
|
||||
&mut self,
|
||||
path: &str,
|
||||
turso_url: Option<&str>,
|
||||
turso_token: Option<&str>,
|
||||
) -> Result<(), SetupError> {
|
||||
use crate::db::libsql_backend::LibSqlBackend;
|
||||
use std::path::Path;
|
||||
|
||||
let db_path = Path::new(path);
|
||||
|
||||
let backend = if let (Some(url), Some(token)) = (turso_url, turso_token) {
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token)
|
||||
.await
|
||||
.map_err(|e| SetupError::Database(format!("Failed to connect: {}", e)))?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| SetupError::Database(format!("Failed to open database: {}", e)))?
|
||||
};
|
||||
|
||||
self.db_backend = Some(backend);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run PostgreSQL migrations.
|
||||
#[cfg(feature = "postgres")]
|
||||
async fn run_migrations_postgres(&self) -> Result<(), SetupError> {
|
||||
if let Some(ref pool) = self.db_pool {
|
||||
use refinery::embed_migrations;
|
||||
embed_migrations!("migrations");
|
||||
@@ -262,6 +434,24 @@ impl SetupWizard {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run libSQL migrations.
|
||||
#[cfg(feature = "libsql")]
|
||||
async fn run_migrations_libsql(&self) -> Result<(), SetupError> {
|
||||
if let Some(ref backend) = self.db_backend {
|
||||
use crate::db::Database;
|
||||
|
||||
print_info("Running migrations...");
|
||||
|
||||
backend
|
||||
.run_migrations()
|
||||
.await
|
||||
.map_err(|e| SetupError::Database(format!("Migration failed: {}", e)))?;
|
||||
|
||||
print_success("Migrations applied");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Step 2: Security (secrets master key).
|
||||
async fn step_security(&mut self) -> Result<(), SetupError> {
|
||||
// Check current configuration
|
||||
@@ -532,24 +722,6 @@ impl SetupWizard {
|
||||
|
||||
/// Initialize secrets context for channel setup.
|
||||
async fn init_secrets_context(&mut self) -> Result<SecretsContext, SetupError> {
|
||||
// Get database pool (should be set from step 1)
|
||||
let pool = if let Some(ref p) = self.db_pool {
|
||||
p.clone()
|
||||
} else {
|
||||
// Fall back to creating one from settings/env
|
||||
let url = self
|
||||
.settings
|
||||
.database_url
|
||||
.clone()
|
||||
.or_else(|| std::env::var("DATABASE_URL").ok())
|
||||
.ok_or_else(|| SetupError::Config("Database URL not configured".to_string()))?;
|
||||
|
||||
self.test_database_connection(&url).await?;
|
||||
// Ensure secrets-related tables exist for channels-only onboarding flows.
|
||||
self.run_migrations().await?;
|
||||
self.db_pool.clone().unwrap()
|
||||
};
|
||||
|
||||
// Get crypto (should be set from step 2, or load from keychain/env)
|
||||
let crypto = if let Some(ref c) = self.secrets_crypto {
|
||||
Arc::clone(c)
|
||||
@@ -571,7 +743,74 @@ impl SetupWizard {
|
||||
Arc::clone(self.secrets_crypto.as_ref().unwrap())
|
||||
};
|
||||
|
||||
Ok(SecretsContext::new(pool, crypto, "default"))
|
||||
// Create backend-appropriate secrets store
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
// Try postgres path first when postgres feature is available
|
||||
if let Some(store) = self.create_postgres_secrets_store(&crypto).await? {
|
||||
return Ok(SecretsContext::from_store(store, "default"));
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
{
|
||||
if let Some(store) = self.create_libsql_secrets_store(&crypto)? {
|
||||
return Ok(SecretsContext::from_store(store, "default"));
|
||||
}
|
||||
}
|
||||
|
||||
Err(SetupError::Config(
|
||||
"No database backend available for secrets storage".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Create a PostgreSQL secrets store from the current pool.
|
||||
#[cfg(feature = "postgres")]
|
||||
async fn create_postgres_secrets_store(
|
||||
&mut self,
|
||||
crypto: &Arc<SecretsCrypto>,
|
||||
) -> Result<Option<Arc<dyn SecretsStore>>, SetupError> {
|
||||
let pool = if let Some(ref p) = self.db_pool {
|
||||
p.clone()
|
||||
} else {
|
||||
// Fall back to creating one from settings/env
|
||||
let url = self
|
||||
.settings
|
||||
.database_url
|
||||
.clone()
|
||||
.or_else(|| std::env::var("DATABASE_URL").ok());
|
||||
|
||||
if let Some(url) = url {
|
||||
self.test_database_connection_postgres(&url).await?;
|
||||
self.run_migrations_postgres().await?;
|
||||
self.db_pool.clone().unwrap()
|
||||
} else {
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let store: Arc<dyn SecretsStore> = Arc::new(crate::secrets::PostgresSecretsStore::new(
|
||||
pool,
|
||||
Arc::clone(crypto),
|
||||
));
|
||||
Ok(Some(store))
|
||||
}
|
||||
|
||||
/// Create a libSQL secrets store from the current backend.
|
||||
#[cfg(feature = "libsql")]
|
||||
fn create_libsql_secrets_store(
|
||||
&self,
|
||||
crypto: &Arc<SecretsCrypto>,
|
||||
) -> Result<Option<Arc<dyn SecretsStore>>, SetupError> {
|
||||
if let Some(ref backend) = self.db_backend {
|
||||
let store: Arc<dyn SecretsStore> = Arc::new(crate::secrets::LibSqlSecretsStore::new(
|
||||
backend.shared_db(),
|
||||
Arc::clone(crypto),
|
||||
));
|
||||
Ok(Some(store))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Step 6: Channel configuration.
|
||||
@@ -794,8 +1033,27 @@ impl SetupWizard {
|
||||
println!("Configuration Summary:");
|
||||
println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
|
||||
|
||||
if self.settings.database_url.is_some() {
|
||||
println!(" Database: configured");
|
||||
let backend = self
|
||||
.settings
|
||||
.database_backend
|
||||
.as_deref()
|
||||
.unwrap_or("postgres");
|
||||
match backend {
|
||||
"libsql" => {
|
||||
if let Some(ref path) = self.settings.libsql_path {
|
||||
println!(" Database: libSQL ({})", path);
|
||||
} else {
|
||||
println!(" Database: libSQL (default path)");
|
||||
}
|
||||
if self.settings.libsql_url.is_some() {
|
||||
println!(" Turso sync: enabled");
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if self.settings.database_url.is_some() {
|
||||
println!(" Database: PostgreSQL (configured)");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match self.settings.secrets_master_key_source {
|
||||
@@ -875,6 +1133,7 @@ impl Default for SetupWizard {
|
||||
}
|
||||
|
||||
/// Mask password in a database URL for display.
|
||||
#[cfg(feature = "postgres")]
|
||||
fn mask_password_in_url(url: &str) -> String {
|
||||
// URL format: scheme://user:password@host/database
|
||||
// Find "://" to locate start of credentials
|
||||
@@ -1065,6 +1324,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "postgres")]
|
||||
fn test_mask_password_in_url() {
|
||||
assert_eq!(
|
||||
mask_password_in_url("postgres://user:secret@localhost/db"),
|
||||
|
||||
@@ -15,7 +15,8 @@ use chrono::Utc;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::context::{ContextManager, JobContext, JobState};
|
||||
use crate::history::{SandboxJobRecord, Store};
|
||||
use crate::db::Database;
|
||||
use crate::history::SandboxJobRecord;
|
||||
use crate::orchestrator::job_manager::{ContainerJobManager, JobMode};
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
|
||||
@@ -27,7 +28,7 @@ use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
pub struct CreateJobTool {
|
||||
context_manager: Arc<ContextManager>,
|
||||
job_manager: Option<Arc<ContainerJobManager>>,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
}
|
||||
|
||||
impl CreateJobTool {
|
||||
@@ -43,7 +44,7 @@ impl CreateJobTool {
|
||||
pub fn with_sandbox(
|
||||
mut self,
|
||||
job_manager: Arc<ContainerJobManager>,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
) -> Self {
|
||||
self.job_manager = Some(job_manager);
|
||||
self.store = store;
|
||||
|
||||
@@ -482,7 +482,7 @@ impl Tool for MemoryTreeTool {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[cfg(all(test, feature = "postgres"))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
|
||||
@@ -19,18 +19,18 @@ use crate::agent::routine::{
|
||||
};
|
||||
use crate::agent::routine_engine::RoutineEngine;
|
||||
use crate::context::JobContext;
|
||||
use crate::history::Store;
|
||||
use crate::db::Database;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
|
||||
// ==================== routine_create ====================
|
||||
|
||||
pub struct RoutineCreateTool {
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
engine: Arc<RoutineEngine>,
|
||||
}
|
||||
|
||||
impl RoutineCreateTool {
|
||||
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
|
||||
pub fn new(store: Arc<dyn Database>, engine: Arc<RoutineEngine>) -> Self {
|
||||
Self { store, engine }
|
||||
}
|
||||
}
|
||||
@@ -277,11 +277,11 @@ impl Tool for RoutineCreateTool {
|
||||
// ==================== routine_list ====================
|
||||
|
||||
pub struct RoutineListTool {
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
}
|
||||
|
||||
impl RoutineListTool {
|
||||
pub fn new(store: Arc<Store>) -> Self {
|
||||
pub fn new(store: Arc<dyn Database>) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
@@ -351,12 +351,12 @@ impl Tool for RoutineListTool {
|
||||
// ==================== routine_update ====================
|
||||
|
||||
pub struct RoutineUpdateTool {
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
engine: Arc<RoutineEngine>,
|
||||
}
|
||||
|
||||
impl RoutineUpdateTool {
|
||||
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
|
||||
pub fn new(store: Arc<dyn Database>, engine: Arc<RoutineEngine>) -> Self {
|
||||
Self { store, engine }
|
||||
}
|
||||
}
|
||||
@@ -474,12 +474,12 @@ impl Tool for RoutineUpdateTool {
|
||||
// ==================== routine_delete ====================
|
||||
|
||||
pub struct RoutineDeleteTool {
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
engine: Arc<RoutineEngine>,
|
||||
}
|
||||
|
||||
impl RoutineDeleteTool {
|
||||
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
|
||||
pub fn new(store: Arc<dyn Database>, engine: Arc<RoutineEngine>) -> Self {
|
||||
Self { store, engine }
|
||||
}
|
||||
}
|
||||
@@ -551,11 +551,11 @@ impl Tool for RoutineDeleteTool {
|
||||
// ==================== routine_history ====================
|
||||
|
||||
pub struct RoutineHistoryTool {
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
}
|
||||
|
||||
impl RoutineHistoryTool {
|
||||
pub fn new(store: Arc<Store>) -> Self {
|
||||
pub fn new(store: Arc<dyn Database>) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
@@ -333,7 +333,7 @@ pub async fn get_mcp_server(name: &str) -> Result<McpServerConfig, ConfigError>
|
||||
///
|
||||
/// Falls back to the disk file if DB has no entry.
|
||||
pub async fn load_mcp_servers_from_db(
|
||||
store: &crate::history::Store,
|
||||
store: &dyn crate::db::Database,
|
||||
user_id: &str,
|
||||
) -> Result<McpServersFile, ConfigError> {
|
||||
match store.get_setting(user_id, "mcp_servers").await {
|
||||
@@ -357,7 +357,7 @@ pub async fn load_mcp_servers_from_db(
|
||||
|
||||
/// Save MCP server configurations to the database settings table.
|
||||
pub async fn save_mcp_servers_to_db(
|
||||
store: &crate::history::Store,
|
||||
store: &dyn crate::db::Database,
|
||||
user_id: &str,
|
||||
config: &McpServersFile,
|
||||
) -> Result<(), ConfigError> {
|
||||
@@ -371,7 +371,7 @@ pub async fn save_mcp_servers_to_db(
|
||||
|
||||
/// Add a new MCP server configuration (DB-backed).
|
||||
pub async fn add_mcp_server_db(
|
||||
store: &crate::history::Store,
|
||||
store: &dyn crate::db::Database,
|
||||
user_id: &str,
|
||||
config: McpServerConfig,
|
||||
) -> Result<(), ConfigError> {
|
||||
@@ -386,7 +386,7 @@ pub async fn add_mcp_server_db(
|
||||
|
||||
/// Remove an MCP server by name (DB-backed).
|
||||
pub async fn remove_mcp_server_db(
|
||||
store: &crate::history::Store,
|
||||
store: &dyn crate::db::Database,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<(), ConfigError> {
|
||||
|
||||
@@ -6,8 +6,8 @@ use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::context::ContextManager;
|
||||
use crate::db::Database;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::llm::{LlmProvider, ToolDefinition};
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
use crate::safety::SafetyLayer;
|
||||
@@ -243,7 +243,7 @@ impl ToolRegistry {
|
||||
&self,
|
||||
context_manager: Arc<ContextManager>,
|
||||
job_manager: Option<Arc<ContainerJobManager>>,
|
||||
store: Option<Arc<Store>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
) {
|
||||
let mut create_tool = CreateJobTool::new(Arc::clone(&context_manager));
|
||||
if let Some(jm) = job_manager {
|
||||
@@ -276,7 +276,7 @@ impl ToolRegistry {
|
||||
/// of routines (scheduled and event-driven tasks).
|
||||
pub fn register_routine_tools(
|
||||
&self,
|
||||
store: Arc<Store>,
|
||||
store: Arc<dyn Database>,
|
||||
engine: Arc<crate::agent::routine_engine::RoutineEngine>,
|
||||
) {
|
||||
use crate::tools::builtin::{
|
||||
|
||||
@@ -108,10 +108,13 @@ pub use credential_injector::{CredentialInjector, InjectedCredentials, Injection
|
||||
pub use rate_limiter::{LimitType, RateLimitError, RateLimitResult, RateLimiter};
|
||||
|
||||
// Storage (V2)
|
||||
#[cfg(feature = "libsql")]
|
||||
pub use storage::LibSqlWasmToolStore;
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use storage::PostgresWasmToolStore;
|
||||
pub use storage::{
|
||||
PostgresWasmToolStore, StoreToolParams, StoredCapabilities, StoredWasmTool,
|
||||
StoredWasmToolWithBinary, ToolStatus, TrustLevel, WasmStorageError, WasmToolStore,
|
||||
compute_binary_hash, verify_binary_integrity,
|
||||
StoreToolParams, StoredCapabilities, StoredWasmTool, StoredWasmToolWithBinary, ToolStatus,
|
||||
TrustLevel, WasmStorageError, WasmToolStore, compute_binary_hash, verify_binary_integrity,
|
||||
};
|
||||
|
||||
// Loader
|
||||
|
||||
@@ -16,6 +16,7 @@ use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
#[cfg(feature = "postgres")]
|
||||
use deadpool_postgres::Pool;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -263,16 +264,19 @@ pub fn verify_binary_integrity(binary: &[u8], expected_hash: &[u8]) -> bool {
|
||||
}
|
||||
|
||||
/// PostgreSQL implementation of WasmToolStore.
|
||||
#[cfg(feature = "postgres")]
|
||||
pub struct PostgresWasmToolStore {
|
||||
pool: Pool,
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
impl PostgresWasmToolStore {
|
||||
pub fn new(pool: Pool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
#[async_trait]
|
||||
impl WasmToolStore for PostgresWasmToolStore {
|
||||
async fn store(&self, params: StoreToolParams) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
@@ -538,6 +542,7 @@ impl WasmToolStore for PostgresWasmToolStore {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
fn row_to_tool(row: &tokio_postgres::Row) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
let trust_level_str: String = row.get("trust_level");
|
||||
let status_str: String = row.get("status");
|
||||
@@ -559,6 +564,458 @@ fn row_to_tool(row: &tokio_postgres::Row) -> Result<StoredWasmTool, WasmStorageE
|
||||
})
|
||||
}
|
||||
|
||||
// ==================== libSQL implementation ====================
|
||||
|
||||
/// libSQL/Turso implementation of WasmToolStore.
|
||||
///
|
||||
/// Holds an `Arc<Database>` handle and creates a fresh connection per operation,
|
||||
/// matching the connection-per-request pattern used by the main `LibSqlBackend`.
|
||||
#[cfg(feature = "libsql")]
|
||||
pub struct LibSqlWasmToolStore {
|
||||
db: std::sync::Arc<libsql::Database>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
impl LibSqlWasmToolStore {
|
||||
pub fn new(db: std::sync::Arc<libsql::Database>) -> Self {
|
||||
Self { db }
|
||||
}
|
||||
|
||||
fn connect(&self) -> Result<libsql::Connection, WasmStorageError> {
|
||||
self.db
|
||||
.connect()
|
||||
.map_err(|e| WasmStorageError::Database(format!("Connection failed: {}", e)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[async_trait]
|
||||
impl WasmToolStore for LibSqlWasmToolStore {
|
||||
async fn store(&self, params: StoreToolParams) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
let binary_hash = compute_binary_hash(¶ms.wasm_binary);
|
||||
let id = Uuid::new_v4();
|
||||
let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
|
||||
let schema_str = serde_json::to_string(¶ms.parameters_schema)
|
||||
.map_err(|e| WasmStorageError::InvalidData(e.to_string()))?;
|
||||
|
||||
// Wrap INSERT + read-back in a transaction to prevent TOCTOU races
|
||||
let conn = self.connect()?;
|
||||
let tx = conn
|
||||
.transaction()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
tx.execute(
|
||||
r#"
|
||||
INSERT INTO wasm_tools (
|
||||
id, user_id, name, version, description, wasm_binary, binary_hash,
|
||||
parameters_schema, source_url, trust_level, status, created_at, updated_at
|
||||
)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, 'active', ?11, ?11)
|
||||
ON CONFLICT (user_id, name, version) DO UPDATE SET
|
||||
description = excluded.description,
|
||||
wasm_binary = excluded.wasm_binary,
|
||||
binary_hash = excluded.binary_hash,
|
||||
parameters_schema = excluded.parameters_schema,
|
||||
source_url = excluded.source_url,
|
||||
updated_at = ?11
|
||||
"#,
|
||||
libsql::params![
|
||||
id.to_string(),
|
||||
params.user_id.as_str(),
|
||||
params.name.as_str(),
|
||||
params.version.as_str(),
|
||||
params.description.as_str(),
|
||||
libsql::Value::Blob(params.wasm_binary),
|
||||
libsql::Value::Blob(binary_hash),
|
||||
schema_str.as_str(),
|
||||
libsql_wasm_opt_text(params.source_url.as_deref()),
|
||||
params.trust_level.to_string(),
|
||||
now.as_str(),
|
||||
],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
// Read back the row within the same transaction
|
||||
let mut rows = tx
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, version, description, parameters_schema,
|
||||
source_url, trust_level, status, created_at, updated_at
|
||||
FROM wasm_tools
|
||||
WHERE user_id = ?1 AND name = ?2
|
||||
ORDER BY version DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
libsql::params![params.user_id.as_str(), params.name.as_str()],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let row = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?
|
||||
.ok_or_else(|| {
|
||||
WasmStorageError::Database("Insert succeeded but row not found".into())
|
||||
})?;
|
||||
|
||||
let tool = libsql_row_to_tool(&row)?;
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
Ok(tool)
|
||||
}
|
||||
|
||||
async fn get(&self, user_id: &str, name: &str) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, version, description, parameters_schema,
|
||||
source_url, trust_level, status, created_at, updated_at
|
||||
FROM wasm_tools
|
||||
WHERE user_id = ?1 AND name = ?2 AND status = 'active'
|
||||
ORDER BY version DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?
|
||||
{
|
||||
Some(row) => {
|
||||
let tool = libsql_row_to_tool(&row)?;
|
||||
match tool.status {
|
||||
ToolStatus::Active => Ok(tool),
|
||||
ToolStatus::Disabled => Err(WasmStorageError::Disabled),
|
||||
ToolStatus::Quarantined => Err(WasmStorageError::Quarantined),
|
||||
}
|
||||
}
|
||||
None => Err(WasmStorageError::NotFound(name.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_with_binary(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
) -> Result<StoredWasmToolWithBinary, WasmStorageError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, version, description, wasm_binary, binary_hash,
|
||||
parameters_schema, source_url, trust_level, status, created_at, updated_at
|
||||
FROM wasm_tools
|
||||
WHERE user_id = ?1 AND name = ?2 AND status = 'active'
|
||||
ORDER BY version DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?
|
||||
{
|
||||
Some(row) => {
|
||||
let wasm_binary: Vec<u8> = row
|
||||
.get(5)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
let binary_hash: Vec<u8> = row
|
||||
.get(6)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
if !verify_binary_integrity(&wasm_binary, &binary_hash) {
|
||||
tracing::error!(
|
||||
user_id = user_id,
|
||||
name = name,
|
||||
"WASM binary integrity check failed"
|
||||
);
|
||||
return Err(WasmStorageError::IntegrityCheckFailed);
|
||||
}
|
||||
|
||||
// Parse metadata from the row (different column offsets due to binary/hash)
|
||||
let tool = libsql_row_to_tool_with_offset(&row)?;
|
||||
|
||||
match tool.status {
|
||||
ToolStatus::Active => Ok(StoredWasmToolWithBinary {
|
||||
tool,
|
||||
wasm_binary,
|
||||
binary_hash,
|
||||
}),
|
||||
ToolStatus::Disabled => Err(WasmStorageError::Disabled),
|
||||
ToolStatus::Quarantined => Err(WasmStorageError::Quarantined),
|
||||
}
|
||||
}
|
||||
None => Err(WasmStorageError::NotFound(name.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_capabilities(
|
||||
&self,
|
||||
tool_id: Uuid,
|
||||
) -> Result<Option<StoredCapabilities>, WasmStorageError> {
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, wasm_tool_id, http_allowlist, allowed_secrets, tool_aliases,
|
||||
requests_per_minute, requests_per_hour, max_request_body_bytes,
|
||||
max_response_body_bytes, workspace_read_prefixes, http_timeout_secs
|
||||
FROM tool_capabilities
|
||||
WHERE wasm_tool_id = ?1
|
||||
"#,
|
||||
libsql::params![tool_id.to_string()],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
match rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?
|
||||
{
|
||||
Some(row) => {
|
||||
let id_str: String = row
|
||||
.get(0)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
let tool_id_str: String = row
|
||||
.get(1)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
let http_allowlist_str: String = row.get::<String>(2).unwrap_or_default();
|
||||
let allowed_secrets_str: String = row.get::<String>(3).unwrap_or_default();
|
||||
let tool_aliases_str: String = row.get::<String>(4).unwrap_or_default();
|
||||
let rpm: i64 = row.get::<i64>(5).unwrap_or(60);
|
||||
let rph: i64 = row.get::<i64>(6).unwrap_or(1000);
|
||||
let max_req: i64 = row.get::<i64>(7).unwrap_or(1048576);
|
||||
let max_resp: i64 = row.get::<i64>(8).unwrap_or(10485760);
|
||||
let ws_prefixes_str: String = row.get::<String>(9).unwrap_or_default();
|
||||
let timeout: i64 = row.get::<i64>(10).unwrap_or(30);
|
||||
|
||||
let http_allowlist: Vec<EndpointPattern> =
|
||||
serde_json::from_str(&http_allowlist_str).unwrap_or_default();
|
||||
let allowed_secrets: Vec<String> =
|
||||
serde_json::from_str(&allowed_secrets_str).unwrap_or_default();
|
||||
let tool_aliases: HashMap<String, String> =
|
||||
serde_json::from_str(&tool_aliases_str).unwrap_or_default();
|
||||
let workspace_read_prefixes: Vec<String> =
|
||||
serde_json::from_str(&ws_prefixes_str).unwrap_or_default();
|
||||
|
||||
Ok(Some(StoredCapabilities {
|
||||
id: id_str
|
||||
.parse()
|
||||
.map_err(|e: uuid::Error| WasmStorageError::InvalidData(e.to_string()))?,
|
||||
wasm_tool_id: tool_id_str
|
||||
.parse()
|
||||
.map_err(|e: uuid::Error| WasmStorageError::InvalidData(e.to_string()))?,
|
||||
http_allowlist,
|
||||
allowed_secrets,
|
||||
tool_aliases,
|
||||
requests_per_minute: rpm as u32,
|
||||
requests_per_hour: rph as u32,
|
||||
max_request_body_bytes: max_req,
|
||||
max_response_body_bytes: max_resp,
|
||||
workspace_read_prefixes,
|
||||
http_timeout_secs: timeout as i32,
|
||||
}))
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list(&self, user_id: &str) -> Result<Vec<StoredWasmTool>, WasmStorageError> {
|
||||
// SQLite doesn't have DISTINCT ON, so we use a subquery to get latest version per name
|
||||
let conn = self.connect()?;
|
||||
let mut rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT id, user_id, name, version, description, parameters_schema,
|
||||
source_url, trust_level, status, created_at, updated_at
|
||||
FROM wasm_tools
|
||||
WHERE user_id = ?1
|
||||
AND rowid IN (
|
||||
SELECT MAX(rowid)
|
||||
FROM wasm_tools
|
||||
WHERE user_id = ?1
|
||||
GROUP BY name
|
||||
)
|
||||
ORDER BY name
|
||||
"#,
|
||||
libsql::params![user_id],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
let mut tools = Vec::new();
|
||||
while let Some(row) = rows
|
||||
.next()
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?
|
||||
{
|
||||
tools.push(libsql_row_to_tool(&row)?);
|
||||
}
|
||||
Ok(tools)
|
||||
}
|
||||
|
||||
async fn update_status(
|
||||
&self,
|
||||
user_id: &str,
|
||||
name: &str,
|
||||
status: ToolStatus,
|
||||
) -> Result<(), WasmStorageError> {
|
||||
let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
|
||||
let conn = self.connect()?;
|
||||
|
||||
let result = conn
|
||||
.execute(
|
||||
"UPDATE wasm_tools SET status = ?1, updated_at = ?2 WHERE user_id = ?3 AND name = ?4",
|
||||
libsql::params![status.to_string(), now.as_str(), user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
if result == 0 {
|
||||
return Err(WasmStorageError::NotFound(name.to_string()));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, WasmStorageError> {
|
||||
let conn = self.connect()?;
|
||||
let result = conn
|
||||
.execute(
|
||||
"DELETE FROM wasm_tools WHERE user_id = ?1 AND name = ?2",
|
||||
libsql::params![user_id, name],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
Ok(result > 0)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_wasm_opt_text(s: Option<&str>) -> libsql::Value {
|
||||
match s {
|
||||
Some(s) => libsql::Value::Text(s.to_string()),
|
||||
None => libsql::Value::Null,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_wasm_parse_ts(s: &str) -> Result<DateTime<Utc>, WasmStorageError> {
|
||||
if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) {
|
||||
return Ok(dt.with_timezone(&Utc));
|
||||
}
|
||||
if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") {
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
Err(WasmStorageError::InvalidData(format!(
|
||||
"unparseable timestamp: {:?}",
|
||||
s
|
||||
)))
|
||||
}
|
||||
|
||||
/// Parse a tool row with standard column order (no binary columns).
|
||||
/// Columns: id(0), user_id(1), name(2), version(3), description(4),
|
||||
/// parameters_schema(5), source_url(6), trust_level(7), status(8),
|
||||
/// created_at(9), updated_at(10)
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_row_to_tool(row: &libsql::Row) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
libsql_row_to_tool_at(row, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
|
||||
}
|
||||
|
||||
/// Parse a tool row when binary columns are present (get_with_binary query).
|
||||
/// Columns: id(0), user_id(1), name(2), version(3), description(4),
|
||||
/// wasm_binary(5), binary_hash(6),
|
||||
/// parameters_schema(7), source_url(8), trust_level(9), status(10),
|
||||
/// created_at(11), updated_at(12)
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_row_to_tool_with_offset(row: &libsql::Row) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
libsql_row_to_tool_at(row, 0, 1, 2, 3, 4, 7, 8, 9, 10, 11, 12)
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
fn libsql_row_to_tool_at(
|
||||
row: &libsql::Row,
|
||||
id_idx: i32,
|
||||
user_id_idx: i32,
|
||||
name_idx: i32,
|
||||
version_idx: i32,
|
||||
description_idx: i32,
|
||||
schema_idx: i32,
|
||||
source_url_idx: i32,
|
||||
trust_level_idx: i32,
|
||||
status_idx: i32,
|
||||
created_at_idx: i32,
|
||||
updated_at_idx: i32,
|
||||
) -> Result<StoredWasmTool, WasmStorageError> {
|
||||
let id_str: String = row
|
||||
.get(id_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
let trust_level_str: String = row
|
||||
.get(trust_level_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
let status_str: String = row
|
||||
.get(status_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
let schema_str: String = row
|
||||
.get(schema_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
let created_at_str: String = row
|
||||
.get(created_at_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
let updated_at_str: String = row
|
||||
.get(updated_at_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?;
|
||||
|
||||
Ok(StoredWasmTool {
|
||||
id: id_str
|
||||
.parse()
|
||||
.map_err(|e: uuid::Error| WasmStorageError::InvalidData(e.to_string()))?,
|
||||
user_id: row
|
||||
.get(user_id_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?,
|
||||
name: row
|
||||
.get(name_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?,
|
||||
version: row
|
||||
.get(version_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?,
|
||||
description: row
|
||||
.get(description_idx)
|
||||
.map_err(|e| WasmStorageError::Database(e.to_string()))?,
|
||||
parameters_schema: serde_json::from_str(&schema_str).unwrap_or_default(),
|
||||
source_url: row
|
||||
.get::<String>(source_url_idx)
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty()),
|
||||
trust_level: trust_level_str
|
||||
.parse()
|
||||
.map_err(WasmStorageError::InvalidData)?,
|
||||
status: status_str.parse().map_err(WasmStorageError::InvalidData)?,
|
||||
created_at: libsql_wasm_parse_ts(&created_at_str)?,
|
||||
updated_at: libsql_wasm_parse_ts(&updated_at_str)?,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tools::wasm::storage::{
|
||||
|
||||
+223
-24
@@ -43,23 +43,206 @@
|
||||
mod chunker;
|
||||
mod document;
|
||||
mod embeddings;
|
||||
#[cfg(feature = "postgres")]
|
||||
mod repository;
|
||||
mod search;
|
||||
|
||||
pub use chunker::{ChunkConfig, chunk_document};
|
||||
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
|
||||
pub use embeddings::{EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OpenAiEmbeddings};
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use repository::Repository;
|
||||
pub use search::{SearchConfig, SearchResult};
|
||||
pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::{NaiveDate, Utc};
|
||||
#[cfg(feature = "postgres")]
|
||||
use deadpool_postgres::Pool;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::WorkspaceError;
|
||||
|
||||
/// Internal storage abstraction for Workspace.
|
||||
///
|
||||
/// Allows Workspace to work with either a PostgreSQL `Repository` (the original
|
||||
/// path) or any `Database` trait implementation (e.g. libSQL backend).
|
||||
enum WorkspaceStorage {
|
||||
/// PostgreSQL-backed repository (uses connection pool directly).
|
||||
#[cfg(feature = "postgres")]
|
||||
Repo(Repository),
|
||||
/// Generic backend implementing the Database trait.
|
||||
Db(Arc<dyn crate::db::Database>),
|
||||
}
|
||||
|
||||
impl WorkspaceStorage {
|
||||
async fn get_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => repo.get_document_by_path(user_id, agent_id, path).await,
|
||||
Self::Db(db) => db.get_document_by_path(user_id, agent_id, path).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => repo.get_document_by_id(id).await,
|
||||
Self::Db(db) => db.get_document_by_id(id).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_or_create_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<MemoryDocument, WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => {
|
||||
repo.get_or_create_document_by_path(user_id, agent_id, path)
|
||||
.await
|
||||
}
|
||||
Self::Db(db) => {
|
||||
db.get_or_create_document_by_path(user_id, agent_id, path)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => repo.update_document(id, content).await,
|
||||
Self::Db(db) => db.update_document(id, content).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_document_by_path(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
path: &str,
|
||||
) -> Result<(), WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => repo.delete_document_by_path(user_id, agent_id, path).await,
|
||||
Self::Db(db) => db.delete_document_by_path(user_id, agent_id, path).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_directory(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
directory: &str,
|
||||
) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => repo.list_directory(user_id, agent_id, directory).await,
|
||||
Self::Db(db) => db.list_directory(user_id, agent_id, directory).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_all_paths(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
) -> Result<Vec<String>, WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => repo.list_all_paths(user_id, agent_id).await,
|
||||
Self::Db(db) => db.list_all_paths(user_id, agent_id).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => repo.delete_chunks(document_id).await,
|
||||
Self::Db(db) => db.delete_chunks(document_id).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn insert_chunk(
|
||||
&self,
|
||||
document_id: Uuid,
|
||||
chunk_index: i32,
|
||||
content: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
) -> Result<Uuid, WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => {
|
||||
repo.insert_chunk(document_id, chunk_index, content, embedding)
|
||||
.await
|
||||
}
|
||||
Self::Db(db) => {
|
||||
db.insert_chunk(document_id, chunk_index, content, embedding)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn update_chunk_embedding(
|
||||
&self,
|
||||
chunk_id: Uuid,
|
||||
embedding: &[f32],
|
||||
) -> Result<(), WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => repo.update_chunk_embedding(chunk_id, embedding).await,
|
||||
Self::Db(db) => db.update_chunk_embedding(chunk_id, embedding).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_chunks_without_embeddings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => {
|
||||
repo.get_chunks_without_embeddings(user_id, agent_id, limit)
|
||||
.await
|
||||
}
|
||||
Self::Db(db) => {
|
||||
db.get_chunks_without_embeddings(user_id, agent_id, limit)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn hybrid_search(
|
||||
&self,
|
||||
user_id: &str,
|
||||
agent_id: Option<Uuid>,
|
||||
query: &str,
|
||||
embedding: Option<&[f32]>,
|
||||
config: &SearchConfig,
|
||||
) -> Result<Vec<SearchResult>, WorkspaceError> {
|
||||
match self {
|
||||
#[cfg(feature = "postgres")]
|
||||
Self::Repo(repo) => {
|
||||
repo.hybrid_search(user_id, agent_id, query, embedding, config)
|
||||
.await
|
||||
}
|
||||
Self::Db(db) => {
|
||||
db.hybrid_search(user_id, agent_id, query, embedding, config)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Default template seeded into HEARTBEAT.md on first access.
|
||||
///
|
||||
/// Intentionally comment-only so the heartbeat runner treats it as
|
||||
@@ -80,25 +263,39 @@ const HEARTBEAT_SEED: &str = "\
|
||||
/// Workspace provides database-backed memory storage for an agent.
|
||||
///
|
||||
/// Each workspace is scoped to a user (and optionally an agent).
|
||||
/// Documents are persisted to PostgreSQL and indexed for search.
|
||||
/// Documents are persisted to the database and indexed for search.
|
||||
/// Supports both PostgreSQL (via Repository) and libSQL (via Database trait).
|
||||
pub struct Workspace {
|
||||
/// User identifier (from channel).
|
||||
user_id: String,
|
||||
/// Optional agent ID for multi-agent isolation.
|
||||
agent_id: Option<Uuid>,
|
||||
/// Database repository.
|
||||
repo: Repository,
|
||||
/// Database storage backend.
|
||||
storage: WorkspaceStorage,
|
||||
/// Embedding provider for semantic search.
|
||||
embeddings: Option<Arc<dyn EmbeddingProvider>>,
|
||||
}
|
||||
|
||||
impl Workspace {
|
||||
/// Create a new workspace for a user.
|
||||
/// Create a new workspace backed by a PostgreSQL connection pool.
|
||||
#[cfg(feature = "postgres")]
|
||||
pub fn new(user_id: impl Into<String>, pool: Pool) -> Self {
|
||||
Self {
|
||||
user_id: user_id.into(),
|
||||
agent_id: None,
|
||||
repo: Repository::new(pool),
|
||||
storage: WorkspaceStorage::Repo(Repository::new(pool)),
|
||||
embeddings: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new workspace backed by any Database implementation.
|
||||
///
|
||||
/// Use this for libSQL or any other backend that implements the Database trait.
|
||||
pub fn new_with_db(user_id: impl Into<String>, db: Arc<dyn crate::db::Database>) -> Self {
|
||||
Self {
|
||||
user_id: user_id.into(),
|
||||
agent_id: None,
|
||||
storage: WorkspaceStorage::Db(db),
|
||||
embeddings: None,
|
||||
}
|
||||
}
|
||||
@@ -138,7 +335,7 @@ impl Workspace {
|
||||
/// ```
|
||||
pub async fn read(&self, path: &str) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let path = normalize_path(path);
|
||||
self.repo
|
||||
self.storage
|
||||
.get_document_by_path(&self.user_id, self.agent_id, &path)
|
||||
.await
|
||||
}
|
||||
@@ -155,14 +352,14 @@ impl Workspace {
|
||||
pub async fn write(&self, path: &str, content: &str) -> Result<MemoryDocument, WorkspaceError> {
|
||||
let path = normalize_path(path);
|
||||
let doc = self
|
||||
.repo
|
||||
.storage
|
||||
.get_or_create_document_by_path(&self.user_id, self.agent_id, &path)
|
||||
.await?;
|
||||
self.repo.update_document(doc.id, content).await?;
|
||||
self.storage.update_document(doc.id, content).await?;
|
||||
self.reindex_document(doc.id).await?;
|
||||
|
||||
// Return updated doc
|
||||
self.repo.get_document_by_id(doc.id).await
|
||||
self.storage.get_document_by_id(doc.id).await
|
||||
}
|
||||
|
||||
/// Append content to a file.
|
||||
@@ -172,7 +369,7 @@ impl Workspace {
|
||||
pub async fn append(&self, path: &str, content: &str) -> Result<(), WorkspaceError> {
|
||||
let path = normalize_path(path);
|
||||
let doc = self
|
||||
.repo
|
||||
.storage
|
||||
.get_or_create_document_by_path(&self.user_id, self.agent_id, &path)
|
||||
.await?;
|
||||
|
||||
@@ -182,7 +379,7 @@ impl Workspace {
|
||||
format!("{}\n{}", doc.content, content)
|
||||
};
|
||||
|
||||
self.repo.update_document(doc.id, &new_content).await?;
|
||||
self.storage.update_document(doc.id, &new_content).await?;
|
||||
self.reindex_document(doc.id).await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -191,7 +388,7 @@ impl Workspace {
|
||||
pub async fn exists(&self, path: &str) -> Result<bool, WorkspaceError> {
|
||||
let path = normalize_path(path);
|
||||
match self
|
||||
.repo
|
||||
.storage
|
||||
.get_document_by_path(&self.user_id, self.agent_id, &path)
|
||||
.await
|
||||
{
|
||||
@@ -206,7 +403,7 @@ impl Workspace {
|
||||
/// Also deletes associated chunks.
|
||||
pub async fn delete(&self, path: &str) -> Result<(), WorkspaceError> {
|
||||
let path = normalize_path(path);
|
||||
self.repo
|
||||
self.storage
|
||||
.delete_document_by_path(&self.user_id, self.agent_id, &path)
|
||||
.await
|
||||
}
|
||||
@@ -229,14 +426,16 @@ impl Workspace {
|
||||
/// ```
|
||||
pub async fn list(&self, directory: &str) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
|
||||
let directory = normalize_directory(directory);
|
||||
self.repo
|
||||
self.storage
|
||||
.list_directory(&self.user_id, self.agent_id, &directory)
|
||||
.await
|
||||
}
|
||||
|
||||
/// List all files recursively (flat list of all paths).
|
||||
pub async fn list_all(&self) -> Result<Vec<String>, WorkspaceError> {
|
||||
self.repo.list_all_paths(&self.user_id, self.agent_id).await
|
||||
self.storage
|
||||
.list_all_paths(&self.user_id, self.agent_id)
|
||||
.await
|
||||
}
|
||||
|
||||
// ==================== Convenience Methods ====================
|
||||
@@ -280,7 +479,7 @@ impl Workspace {
|
||||
|
||||
/// Helper to read or create a file.
|
||||
async fn read_or_create(&self, path: &str) -> Result<MemoryDocument, WorkspaceError> {
|
||||
self.repo
|
||||
self.storage
|
||||
.get_or_create_document_by_path(&self.user_id, self.agent_id, path)
|
||||
.await
|
||||
}
|
||||
@@ -299,7 +498,7 @@ impl Workspace {
|
||||
} else {
|
||||
format!("{}\n\n{}", doc.content, entry)
|
||||
};
|
||||
self.repo.update_document(doc.id, &new_content).await?;
|
||||
self.storage.update_document(doc.id, &new_content).await?;
|
||||
self.reindex_document(doc.id).await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -395,7 +594,7 @@ impl Workspace {
|
||||
None
|
||||
};
|
||||
|
||||
self.repo
|
||||
self.storage
|
||||
.hybrid_search(
|
||||
&self.user_id,
|
||||
self.agent_id,
|
||||
@@ -411,13 +610,13 @@ impl Workspace {
|
||||
/// Re-index a document (chunk and generate embeddings).
|
||||
async fn reindex_document(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
|
||||
// Get the document
|
||||
let doc = self.repo.get_document_by_id(document_id).await?;
|
||||
let doc = self.storage.get_document_by_id(document_id).await?;
|
||||
|
||||
// Chunk the content
|
||||
let chunks = chunk_document(&doc.content, ChunkConfig::default());
|
||||
|
||||
// Delete old chunks
|
||||
self.repo.delete_chunks(document_id).await?;
|
||||
self.storage.delete_chunks(document_id).await?;
|
||||
|
||||
// Insert new chunks
|
||||
for (index, content) in chunks.into_iter().enumerate() {
|
||||
@@ -434,7 +633,7 @@ impl Workspace {
|
||||
None
|
||||
};
|
||||
|
||||
self.repo
|
||||
self.storage
|
||||
.insert_chunk(document_id, index as i32, &content, embedding.as_deref())
|
||||
.await?;
|
||||
}
|
||||
@@ -542,7 +741,7 @@ impl Workspace {
|
||||
};
|
||||
|
||||
let chunks = self
|
||||
.repo
|
||||
.storage
|
||||
.get_chunks_without_embeddings(&self.user_id, self.agent_id, 100)
|
||||
.await?;
|
||||
|
||||
@@ -550,7 +749,7 @@ impl Workspace {
|
||||
for chunk in chunks {
|
||||
match provider.embed(&chunk.content).await {
|
||||
Ok(embedding) => {
|
||||
self.repo
|
||||
self.storage
|
||||
.update_chunk_embedding(chunk.id, &embedding)
|
||||
.await?;
|
||||
count += 1;
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
#![cfg(feature = "postgres")]
|
||||
//! Integration tests for the workspace module.
|
||||
//!
|
||||
//! Requires a running PostgreSQL with pgvector extension.
|
||||
|
||||
Reference in New Issue
Block a user