Wire database Store into agent loop

Persist jobs and actions to PostgreSQL using fire-and-forget pattern:
- Scheduler passes store to Worker, persists cancellations
- Worker persists job status changes and tool execution actions
- Agent persists new jobs on creation
- All DB writes use tokio::spawn to avoid blocking execution

Store remains optional to preserve --no-db mode.

Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-02 22:53:12 -08:00
co-authored by Claude Opus 4.5
parent aea3f47f8b
commit 45bbfa026d
4 changed files with 137 additions and 17 deletions
+5 -6
View File
@@ -234,12 +234,11 @@ Key test patterns:
1. **Slack/Telegram channels** - Stubs only, need implementation
2. **Tool sandboxing** - `sandbox.rs` is a stub, needs WASM integration
3. **Dynamic tool building** - `builder.rs` placeholder, needs LLM code generation
4. **Database integration** - Store is created but not fully wired into agent loop
5. **Integration tests** - Need testcontainers setup for PostgreSQL
6. **MCP stdio transport** - Only HTTP transport implemented
7. **Workspace integration** - Memory tools need to be registered and workspace passed to workers
8. **Embedding backfill** - Background job to generate embeddings for chunks missing them
9. **Context compaction** - Auto-trigger memory preservation before context window fills
4. **Integration tests** - Need testcontainers setup for PostgreSQL
5. **MCP stdio transport** - Only HTTP transport implemented
6. **Workspace integration** - Memory tools need to be registered and workspace passed to workers
7. **Embedding backfill** - Background job to generate embeddings for chunks missing them
8. **Context compaction** - Auto-trigger memory preservation before context window fills
## Adding a New Tool
+13
View File
@@ -47,6 +47,7 @@ impl Agent {
llm.clone(),
safety.clone(),
tools.clone(),
store.clone(),
));
Self {
@@ -171,6 +172,18 @@ impl Agent {
.await?;
}
// Persist new job to database (fire-and-forget)
if let Some(ref store) = self.store {
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
let store = store.clone();
tokio::spawn(async move {
if let Err(e) = store.save_job(&ctx).await {
tracing::warn!("Failed to persist new job {}: {}", job_id, e);
}
});
}
}
// Schedule for execution
self.scheduler.schedule(job_id).await?;
+22
View File
@@ -11,6 +11,7 @@ use crate::agent::Worker;
use crate::config::AgentConfig;
use crate::context::{ContextManager, JobState};
use crate::error::JobError;
use crate::history::Store;
use crate::llm::LlmProvider;
use crate::safety::SafetyLayer;
use crate::tools::ToolRegistry;
@@ -41,6 +42,7 @@ pub struct Scheduler {
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
/// Running jobs.
jobs: RwLock<HashMap<Uuid, ScheduledJob>>,
}
@@ -53,6 +55,7 @@ impl Scheduler {
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
) -> Self {
Self {
config,
@@ -60,6 +63,7 @@ impl Scheduler {
llm,
safety,
tools,
store,
jobs: RwLock::new(HashMap::new()),
}
}
@@ -103,6 +107,7 @@ impl Scheduler {
self.llm.clone(),
self.safety.clone(),
self.tools.clone(),
self.store.clone(),
self.config.job_timeout,
);
@@ -152,6 +157,23 @@ impl Scheduler {
})
.await?;
// Persist cancellation (fire-and-forget)
if let Some(ref store) = self.store {
let store = store.clone();
tokio::spawn(async move {
if let Err(e) = store
.update_job_status(
job_id,
JobState::Cancelled,
Some("Stopped by scheduler"),
)
.await
{
tracing::warn!("Failed to persist cancellation for job {}: {}", job_id, e);
}
});
}
tracing::info!("Stopped job {}", job_id);
}
+97 -11
View File
@@ -9,6 +9,7 @@ use uuid::Uuid;
use crate::agent::scheduler::WorkerMessage;
use crate::context::{ContextManager, JobState};
use crate::error::Error;
use crate::history::Store;
use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext};
use crate::safety::SafetyLayer;
use crate::tools::ToolRegistry;
@@ -20,6 +21,7 @@ pub struct Worker {
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
timeout: Duration,
}
@@ -31,6 +33,7 @@ impl Worker {
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
timeout: Duration,
) -> Self {
Self {
@@ -39,10 +42,27 @@ impl Worker {
llm,
safety,
tools,
store,
timeout,
}
}
/// Fire-and-forget persistence of job status.
fn persist_status(&self, status: JobState, reason: Option<String>) {
if let Some(ref store) = self.store {
let store = store.clone();
let job_id = self.job_id;
tokio::spawn(async move {
if let Err(e) = store
.update_job_status(job_id, status, reason.as_deref())
.await
{
tracing::warn!("Failed to persist status for job {}: {}", job_id, e);
}
});
}
}
/// Run the worker until the job is complete or stopped.
pub async fn run(self, mut rx: mpsc::Receiver<WorkerMessage>) -> Result<(), Error> {
tracing::info!("Worker starting for job {}", self.job_id);
@@ -241,22 +261,79 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
// Get job context for the tool
let job_ctx = self.context_manager.get_context(self.job_id).await?;
// Execute with timeout
// Execute with timeout and timing
let start = std::time::Instant::now();
let result = tokio::time::timeout(Duration::from_secs(60), async {
tool.execute(params.clone(), &job_ctx).await
})
.await
.map_err(|_| crate::error::ToolError::Timeout {
name: tool_name.to_string(),
timeout: Duration::from_secs(60),
})?
.map_err(|e| crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: e.to_string(),
})?;
.await;
let elapsed = start.elapsed();
// Record action in memory and get the ActionRecord for persistence
let action = match &result {
Ok(Ok(output)) => {
let output_str = serde_json::to_string_pretty(&output.result).ok();
self.context_manager
.update_memory(self.job_id, |mem| {
let rec = mem.create_action(tool_name, params.clone()).succeed(
output_str.clone(),
output.result.clone(),
elapsed,
);
mem.record_action(rec.clone());
rec
})
.await
.ok()
}
Ok(Err(e)) => self
.context_manager
.update_memory(self.job_id, |mem| {
let rec = mem
.create_action(tool_name, params.clone())
.fail(e.to_string(), elapsed);
mem.record_action(rec.clone());
rec
})
.await
.ok(),
Err(_) => self
.context_manager
.update_memory(self.job_id, |mem| {
let rec = mem
.create_action(tool_name, params.clone())
.fail("Execution timeout", elapsed);
mem.record_action(rec.clone());
rec
})
.await
.ok(),
};
// Persist action to database (fire-and-forget)
if let (Some(action), Some(store)) = (action, &self.store) {
let store = store.clone();
let job_id = self.job_id;
tokio::spawn(async move {
if let Err(e) = store.save_action(job_id, &action).await {
tracing::warn!("Failed to persist action for job {}: {}", job_id, e);
}
});
}
// Handle the result
let output = result
.map_err(|_| crate::error::ToolError::Timeout {
name: tool_name.to_string(),
timeout: Duration::from_secs(60),
})?
.map_err(|e| crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: e.to_string(),
})?;
// Return result as string
serde_json::to_string_pretty(&result.result).map_err(|e| {
serde_json::to_string_pretty(&output.result).map_err(|e| {
crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: format!("Failed to serialize result: {}", e),
@@ -278,6 +355,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
id: self.job_id,
reason: s,
})?;
self.persist_status(
JobState::Completed,
Some("Job completed successfully".to_string()),
);
Ok(())
}
@@ -291,6 +373,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
id: self.job_id,
reason: s,
})?;
self.persist_status(JobState::Failed, Some(reason.to_string()));
Ok(())
}
@@ -302,6 +386,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
id: self.job_id,
reason: s,
})?;
self.persist_status(JobState::Stuck, Some(reason.to_string()));
Ok(())
}
}