From 45bbfa026d551390bd387abb463d26f72f63b056 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Mon, 2 Feb 2026 22:53:12 -0800 Subject: [PATCH] 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 --- CLAUDE.md | 11 ++-- src/agent/agent_loop.rs | 13 +++++ src/agent/scheduler.rs | 22 ++++++++ src/agent/worker.rs | 108 ++++++++++++++++++++++++++++++++++++---- 4 files changed, 137 insertions(+), 17 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 3606c034..724f51d8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 548ce008..75f677ae 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -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?; diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 14d67a88..b92ec923 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -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, safety: Arc, tools: Arc, + store: Option>, /// Running jobs. jobs: RwLock>, } @@ -53,6 +55,7 @@ impl Scheduler { llm: Arc, safety: Arc, tools: Arc, + store: Option>, ) -> 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); } diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 18fbfb34..780bd9f2 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -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, safety: Arc, tools: Arc, + store: Option>, timeout: Duration, } @@ -31,6 +33,7 @@ impl Worker { llm: Arc, safety: Arc, tools: Arc, + store: Option>, 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) { + 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) -> 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(()) } }