mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
feat(engine): self-improving engine via Mission system
Wire the self-improvement loop as a Mission with OnSystemEvent cadence, inspired by karpathy/autoresearch's program.md approach. The mission fires when threads complete with issues, receives trace data as trigger payload, and uses tools directly to diagnose and fix problems. Key changes: Engine self-improvement (Phase A+B from design doc): - Add fire_on_system_event() to MissionManager for OnSystemEvent cadence - Add start_event_listener() that subscribes to thread events and fires matching missions when non-Mission threads complete with trace issues - Add ensure_self_improvement_mission() with autoresearch-style goal prompt (concrete loop steps, not vague instructions) - Add process_self_improvement_output() for structured JSON fallback - Seed fix pattern database with 8 known patterns from debugging - Runtime prompt overlay via MemoryDoc (build_codeact_system_prompt now async + Store-aware, appends learned rules from prompt_overlay docs) - Pass Store to ExecutionLoop for overlay loading Bridge review fixes (P1/P2): - Scope engine v2 SSE events to requesting user (broadcast_for_user) - Per-user pending approvals via HashMap instead of global Option - Reset tool-call limit counter before each thread execution - Only persist auto-approval when user chose "always", not one-off "yes" - Remove dead store/mission_manager fields from EngineState Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
@@ -5,6 +5,7 @@
|
||||
//! - [`PolicyEngine`] — deterministic effect-level allow/deny/approve
|
||||
|
||||
pub mod lease;
|
||||
pub mod planner;
|
||||
pub mod policy;
|
||||
pub mod registry;
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
//! Lease planning for new threads.
|
||||
//!
|
||||
//! Converts capability registry contents plus thread type into explicit
|
||||
//! capability grants so new threads do not receive implicit wildcard leases.
|
||||
|
||||
use crate::capability::registry::CapabilityRegistry;
|
||||
use crate::types::thread::ThreadType;
|
||||
|
||||
/// Explicit grant plan for a single capability.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CapabilityGrantPlan {
|
||||
pub capability_name: String,
|
||||
pub granted_actions: Vec<String>,
|
||||
}
|
||||
|
||||
/// Plans explicit capability leases for new threads.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct LeasePlanner;
|
||||
|
||||
impl LeasePlanner {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
|
||||
/// Build the capability grants for a new thread.
|
||||
///
|
||||
/// Reflection threads are handled by the reflection pipeline's dedicated
|
||||
/// executor, so the default planner grants no host capabilities to them.
|
||||
pub fn plan_for_thread(
|
||||
&self,
|
||||
thread_type: ThreadType,
|
||||
capabilities: &CapabilityRegistry,
|
||||
) -> Vec<CapabilityGrantPlan> {
|
||||
if thread_type == ThreadType::Reflection {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
capabilities
|
||||
.list()
|
||||
.into_iter()
|
||||
.filter_map(|cap| {
|
||||
let granted_actions: Vec<String> = cap
|
||||
.actions
|
||||
.iter()
|
||||
.map(|action| action.name.clone())
|
||||
.collect();
|
||||
if granted_actions.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(CapabilityGrantPlan {
|
||||
capability_name: cap.name.clone(),
|
||||
granted_actions,
|
||||
})
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::types::capability::{ActionDef, Capability, EffectType};
|
||||
|
||||
fn registry() -> CapabilityRegistry {
|
||||
let mut reg = CapabilityRegistry::new();
|
||||
reg.register(Capability {
|
||||
name: "tools".into(),
|
||||
description: "test".into(),
|
||||
actions: vec![ActionDef {
|
||||
name: "read_file".into(),
|
||||
description: "read".into(),
|
||||
parameters_schema: serde_json::json!({}),
|
||||
effects: vec![EffectType::ReadLocal],
|
||||
requires_approval: false,
|
||||
}],
|
||||
knowledge: vec![],
|
||||
policies: vec![],
|
||||
});
|
||||
reg
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn foreground_threads_get_explicit_actions() {
|
||||
let planner = LeasePlanner::new();
|
||||
let plans = planner.plan_for_thread(ThreadType::Foreground, ®istry());
|
||||
assert_eq!(plans.len(), 1);
|
||||
assert_eq!(plans[0].capability_name, "tools");
|
||||
assert_eq!(plans[0].granted_actions, vec!["read_file"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reflection_threads_do_not_get_default_capabilities() {
|
||||
let planner = LeasePlanner::new();
|
||||
let plans = planner.plan_for_thread(ThreadType::Reflection, ®istry());
|
||||
assert!(plans.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -39,6 +39,8 @@ pub struct ExecutionLoop {
|
||||
event_tx: Option<tokio::sync::broadcast::Sender<crate::types::event::ThreadEvent>>,
|
||||
/// Optional retrieval engine for injecting prior knowledge into context.
|
||||
retrieval: Option<crate::memory::RetrievalEngine>,
|
||||
/// Optional Store for runtime prompt overlay loading.
|
||||
store: Option<Arc<dyn crate::traits::store::Store>>,
|
||||
}
|
||||
|
||||
impl ExecutionLoop {
|
||||
@@ -62,6 +64,7 @@ impl ExecutionLoop {
|
||||
capabilities: None,
|
||||
event_tx: None,
|
||||
retrieval: None,
|
||||
store: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -89,6 +92,12 @@ impl ExecutionLoop {
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the Store for runtime prompt overlay loading.
|
||||
pub fn with_store(mut self, store: Arc<dyn crate::traits::store::Store>) -> Self {
|
||||
self.store = Some(store);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add an event to the thread and broadcast it for live status updates.
|
||||
fn emit_event(&mut self, kind: EventKind) {
|
||||
let event = crate::types::event::ThreadEvent::new(self.thread.id, kind);
|
||||
@@ -99,8 +108,32 @@ impl ExecutionLoop {
|
||||
self.thread.updated_at = chrono::Utc::now();
|
||||
}
|
||||
|
||||
async fn persist_runtime_state(
|
||||
&self,
|
||||
step: Option<&Step>,
|
||||
persisted_event_count: &mut usize,
|
||||
) -> Result<(), EngineError> {
|
||||
let Some(store) = self.store.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
if let Some(step) = step {
|
||||
store.save_step(step).await?;
|
||||
}
|
||||
if *persisted_event_count < self.thread.events.len() {
|
||||
store
|
||||
.append_events(&self.thread.events[*persisted_event_count..])
|
||||
.await?;
|
||||
*persisted_event_count = self.thread.events.len();
|
||||
}
|
||||
store.save_thread(&self.thread).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run the execution loop to completion.
|
||||
pub async fn run(&mut self) -> Result<ThreadOutcome, EngineError> {
|
||||
let mut persisted_event_count = 0;
|
||||
|
||||
// Transition to Running
|
||||
self.thread.transition_to(ThreadState::Running, None)?;
|
||||
|
||||
@@ -120,11 +153,18 @@ impl ExecutionLoop {
|
||||
Vec::new()
|
||||
}
|
||||
};
|
||||
let system_prompt = crate::executor::prompt::build_codeact_system_prompt(&actions);
|
||||
let system_prompt = crate::executor::prompt::build_codeact_system_prompt(
|
||||
&actions,
|
||||
self.store.as_ref(),
|
||||
self.thread.project_id,
|
||||
)
|
||||
.await;
|
||||
self.thread
|
||||
.messages
|
||||
.insert(0, ThreadMessage::system(system_prompt));
|
||||
}
|
||||
self.persist_runtime_state(None, &mut persisted_event_count)
|
||||
.await?;
|
||||
|
||||
let max_iterations = self.thread.config.max_iterations;
|
||||
let max_nudges = self.thread.config.max_tool_intent_nudges;
|
||||
@@ -145,10 +185,14 @@ impl ExecutionLoop {
|
||||
SignalAction::Stop => {
|
||||
self.thread
|
||||
.transition_to(ThreadState::Completed, Some("stopped by signal".into()))?;
|
||||
self.persist_runtime_state(None, &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(ThreadOutcome::Stopped);
|
||||
}
|
||||
SignalAction::Inject(msg) => {
|
||||
self.thread.add_message(msg);
|
||||
self.persist_runtime_state(None, &mut persisted_event_count)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -164,6 +208,8 @@ impl ExecutionLoop {
|
||||
);
|
||||
self.thread
|
||||
.transition_to(ThreadState::Completed, Some("token limit exceeded".into()))?;
|
||||
self.persist_runtime_state(None, &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(ThreadOutcome::Failed {
|
||||
error: format!(
|
||||
"Token limit exceeded: {} of {} tokens",
|
||||
@@ -183,6 +229,8 @@ impl ExecutionLoop {
|
||||
);
|
||||
self.thread
|
||||
.transition_to(ThreadState::Completed, Some("timeout".into()))?;
|
||||
self.persist_runtime_state(None, &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(ThreadOutcome::Failed {
|
||||
error: format!("Thread timeout: {elapsed:?} of {max_dur:?}"),
|
||||
});
|
||||
@@ -200,6 +248,8 @@ impl ExecutionLoop {
|
||||
);
|
||||
self.thread
|
||||
.transition_to(ThreadState::Completed, Some("USD budget exceeded".into()))?;
|
||||
self.persist_runtime_state(None, &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(ThreadOutcome::Failed {
|
||||
error: format!(
|
||||
"USD budget exceeded: ${:.4} of ${:.4}",
|
||||
@@ -257,6 +307,8 @@ impl ExecutionLoop {
|
||||
let mut step = Step::new(self.thread.id, iteration + 1);
|
||||
step.status = StepStatus::LlmCalling;
|
||||
self.emit_event(EventKind::StepStarted { step_id: step.id });
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
|
||||
// 7. Call LLM
|
||||
// CodeAct/RLM: send NO structured tool definitions — tools are described
|
||||
@@ -338,6 +390,8 @@ impl ExecutionLoop {
|
||||
ThreadState::Completed,
|
||||
Some("FINAL() in text".into()),
|
||||
)?;
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(ThreadOutcome::Completed {
|
||||
response: Some(answer),
|
||||
});
|
||||
@@ -365,6 +419,8 @@ impl ExecutionLoop {
|
||||
tokens: step.tokens_used,
|
||||
});
|
||||
self.thread.step_count += 1;
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -382,6 +438,8 @@ impl ExecutionLoop {
|
||||
|
||||
self.thread
|
||||
.transition_to(ThreadState::Completed, Some("text response".into()))?;
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(ThreadOutcome::Completed {
|
||||
response: Some(text),
|
||||
});
|
||||
@@ -464,8 +522,12 @@ impl ExecutionLoop {
|
||||
ThreadState::Waiting,
|
||||
Some("awaiting approval".into()),
|
||||
)?;
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(outcome);
|
||||
}
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
}
|
||||
|
||||
LlmResponse::Code { code, content } => {
|
||||
@@ -654,6 +716,8 @@ impl ExecutionLoop {
|
||||
if let Some(answer) = code_result.final_answer {
|
||||
self.thread
|
||||
.transition_to(ThreadState::Completed, Some("FINAL() called".into()))?;
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(ThreadOutcome::Completed {
|
||||
response: Some(answer),
|
||||
});
|
||||
@@ -665,6 +729,8 @@ impl ExecutionLoop {
|
||||
ThreadState::Waiting,
|
||||
Some("awaiting approval".into()),
|
||||
)?;
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(outcome);
|
||||
}
|
||||
|
||||
@@ -674,6 +740,9 @@ impl ExecutionLoop {
|
||||
} else {
|
||||
consecutive_errors = 0;
|
||||
}
|
||||
|
||||
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -693,6 +762,8 @@ impl ExecutionLoop {
|
||||
"consecutive error threshold: {consecutive_errors} errors"
|
||||
)),
|
||||
)?;
|
||||
self.persist_runtime_state(None, &mut persisted_event_count)
|
||||
.await?;
|
||||
return Ok(ThreadOutcome::Failed {
|
||||
error: format!(
|
||||
"Consecutive error threshold exceeded: {consecutive_errors} of {max_errors}"
|
||||
@@ -711,6 +782,8 @@ impl ExecutionLoop {
|
||||
ThreadState::Completed,
|
||||
Some("max iterations reached".into()),
|
||||
)?;
|
||||
self.persist_runtime_state(None, &mut persisted_event_count)
|
||||
.await?;
|
||||
Ok(ThreadOutcome::MaxIterations)
|
||||
}
|
||||
|
||||
|
||||
@@ -5,9 +5,14 @@
|
||||
//!
|
||||
//! Prompt templates live in `crates/ironclaw_engine/prompts/` as plain
|
||||
//! markdown files for easy inspection and iteration. They are embedded
|
||||
//! at compile time via `include_str!`.
|
||||
//! at compile time via `include_str!` and can be extended at runtime with
|
||||
//! prompt overlays stored as MemoryDocs.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::traits::store::Store;
|
||||
use crate::types::capability::ActionDef;
|
||||
use crate::types::project::ProjectId;
|
||||
|
||||
/// The main instruction block (before tool listing).
|
||||
const CODEACT_PREAMBLE: &str = include_str!("../../prompts/codeact_preamble.md");
|
||||
@@ -15,6 +20,15 @@ const CODEACT_PREAMBLE: &str = include_str!("../../prompts/codeact_preamble.md")
|
||||
/// The strategy/closing block (after tool listing).
|
||||
const CODEACT_POSTAMBLE: &str = include_str!("../../prompts/codeact_postamble.md");
|
||||
|
||||
/// Well-known title for the CodeAct preamble overlay.
|
||||
pub const PREAMBLE_OVERLAY_TITLE: &str = "prompt:codeact_preamble";
|
||||
|
||||
/// Well-known tag for prompt overlay docs.
|
||||
pub const PROMPT_OVERLAY_TAG: &str = "prompt_overlay";
|
||||
|
||||
/// Maximum size for a prompt overlay document (in chars).
|
||||
const MAX_PROMPT_OVERLAY_CHARS: usize = 4000;
|
||||
|
||||
/// Build the system prompt for CodeAct/RLM execution.
|
||||
///
|
||||
/// The prompt instructs the LLM to:
|
||||
@@ -23,9 +37,26 @@ const CODEACT_POSTAMBLE: &str = include_str!("../../prompts/codeact_postamble.md
|
||||
/// - Use llm_query(prompt, context) for sub-agent calls
|
||||
/// - Use FINAL(answer) to return the final answer
|
||||
/// - Access thread context via the `context` variable
|
||||
pub fn build_codeact_system_prompt(actions: &[ActionDef]) -> String {
|
||||
///
|
||||
/// If a Store is provided, checks for a runtime prompt overlay (a MemoryDoc
|
||||
/// with tag "prompt_overlay" and title "prompt:codeact_preamble") and appends
|
||||
/// its content after the compiled preamble. This enables the self-improvement
|
||||
/// mission to evolve the system prompt at runtime.
|
||||
pub async fn build_codeact_system_prompt(
|
||||
actions: &[ActionDef],
|
||||
store: Option<&Arc<dyn Store>>,
|
||||
project_id: ProjectId,
|
||||
) -> String {
|
||||
let mut prompt = String::from(CODEACT_PREAMBLE);
|
||||
|
||||
// Append runtime prompt overlay if available
|
||||
if let Some(store) = store
|
||||
&& let Some(overlay) = load_prompt_overlay(store, project_id).await
|
||||
{
|
||||
prompt.push_str("\n\n## Learned Rules (from self-improvement)\n\n");
|
||||
prompt.push_str(&overlay);
|
||||
}
|
||||
|
||||
// Add tool documentation
|
||||
if !actions.is_empty() {
|
||||
prompt.push_str("\n## Available tools (call as Python functions)\n\n");
|
||||
@@ -45,3 +76,109 @@ pub fn build_codeact_system_prompt(actions: &[ActionDef]) -> String {
|
||||
prompt.push_str(CODEACT_POSTAMBLE);
|
||||
prompt
|
||||
}
|
||||
|
||||
/// Load the prompt overlay from the Store, if one exists for this project.
|
||||
async fn load_prompt_overlay(store: &Arc<dyn Store>, project_id: ProjectId) -> Option<String> {
|
||||
let docs = store.list_memory_docs(project_id).await.ok()?;
|
||||
let overlay = docs.iter().find(|d| {
|
||||
d.title == PREAMBLE_OVERLAY_TITLE && d.tags.contains(&PROMPT_OVERLAY_TAG.to_string())
|
||||
})?;
|
||||
|
||||
let content: String = overlay
|
||||
.content
|
||||
.chars()
|
||||
.take(MAX_PROMPT_OVERLAY_CHARS)
|
||||
.collect();
|
||||
if content.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(content)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::types::memory::{DocId, DocType, MemoryDoc};
|
||||
|
||||
#[tokio::test]
|
||||
async fn prompt_without_store_uses_compiled_preamble() {
|
||||
let prompt = build_codeact_system_prompt(&[], None, ProjectId(uuid::Uuid::nil())).await;
|
||||
assert!(prompt.contains("Python REPL environment"));
|
||||
assert!(prompt.contains("Strategy"));
|
||||
assert!(!prompt.contains("Learned Rules"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prompt_with_overlay_appends_rules() {
|
||||
let project_id = ProjectId(uuid::Uuid::new_v4());
|
||||
let overlay = MemoryDoc {
|
||||
id: DocId::new(),
|
||||
project_id,
|
||||
doc_type: DocType::Note,
|
||||
title: PREAMBLE_OVERLAY_TITLE.into(),
|
||||
content: "9. Never call web_fetch — use http() instead.".into(),
|
||||
source_thread_id: None,
|
||||
tags: vec![PROMPT_OVERLAY_TAG.into()],
|
||||
metadata: serde_json::json!({}),
|
||||
created_at: chrono::Utc::now(),
|
||||
updated_at: chrono::Utc::now(),
|
||||
};
|
||||
|
||||
let store = Arc::new(crate::tests::InMemoryStore::with_docs(vec![overlay]));
|
||||
let prompt =
|
||||
build_codeact_system_prompt(&[], Some(&(store as Arc<dyn Store>)), project_id).await;
|
||||
assert!(prompt.contains("Learned Rules"));
|
||||
assert!(prompt.contains("Never call web_fetch"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prompt_overlay_size_is_capped() {
|
||||
let project_id = ProjectId(uuid::Uuid::new_v4());
|
||||
// Create an overlay that exceeds MAX_PROMPT_OVERLAY_CHARS using a char
|
||||
// not found in the compiled preamble/postamble
|
||||
let huge_content = "\u{2603}".repeat(MAX_PROMPT_OVERLAY_CHARS + 1000); // snowman
|
||||
let overlay = MemoryDoc {
|
||||
id: DocId::new(),
|
||||
project_id,
|
||||
doc_type: DocType::Note,
|
||||
title: PREAMBLE_OVERLAY_TITLE.into(),
|
||||
content: huge_content,
|
||||
source_thread_id: None,
|
||||
tags: vec![PROMPT_OVERLAY_TAG.into()],
|
||||
metadata: serde_json::json!({}),
|
||||
created_at: chrono::Utc::now(),
|
||||
updated_at: chrono::Utc::now(),
|
||||
};
|
||||
|
||||
let store = Arc::new(crate::tests::InMemoryStore::with_docs(vec![overlay]));
|
||||
let prompt =
|
||||
build_codeact_system_prompt(&[], Some(&(store as Arc<dyn Store>)), project_id).await;
|
||||
|
||||
let snowman_count = prompt.chars().filter(|c| *c == '\u{2603}').count();
|
||||
assert_eq!(snowman_count, MAX_PROMPT_OVERLAY_CHARS);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prompt_ignores_wrong_project_overlay() {
|
||||
let project_id = ProjectId(uuid::Uuid::new_v4());
|
||||
let other_project = ProjectId(uuid::Uuid::new_v4());
|
||||
let overlay = MemoryDoc {
|
||||
id: DocId::new(),
|
||||
project_id: other_project,
|
||||
doc_type: DocType::Note,
|
||||
title: PREAMBLE_OVERLAY_TITLE.into(),
|
||||
content: "Should not appear".into(),
|
||||
source_thread_id: None,
|
||||
tags: vec![PROMPT_OVERLAY_TAG.into()],
|
||||
metadata: serde_json::json!({}),
|
||||
created_at: chrono::Utc::now(),
|
||||
updated_at: chrono::Utc::now(),
|
||||
};
|
||||
|
||||
let store = Arc::new(crate::tests::InMemoryStore::with_docs(vec![overlay]));
|
||||
let prompt =
|
||||
build_codeact_system_prompt(&[], Some(&(store as Arc<dyn Store>)), project_id).await;
|
||||
assert!(!prompt.contains("Should not appear"));
|
||||
assert!(!prompt.contains("Learned Rules"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,6 +50,7 @@ pub use traits::store::Store;
|
||||
// ── Re-exports: capability ────────────────────────────────────
|
||||
|
||||
pub use capability::lease::LeaseManager;
|
||||
pub use capability::planner::{CapabilityGrantPlan, LeasePlanner};
|
||||
pub use capability::policy::{PolicyDecision, PolicyEngine};
|
||||
pub use capability::registry::CapabilityRegistry;
|
||||
|
||||
@@ -81,3 +82,162 @@ pub use reflection::ReflectionResult;
|
||||
// ── Re-exports: reliability ──────────────────────────────────
|
||||
|
||||
pub use reliability::ReliabilityTracker;
|
||||
|
||||
// ── Test utilities ──────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod tests {
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::traits::store::Store;
|
||||
use crate::types::capability::{CapabilityLease, LeaseId};
|
||||
use crate::types::conversation::{ConversationId, ConversationSurface};
|
||||
use crate::types::error::EngineError;
|
||||
use crate::types::event::ThreadEvent;
|
||||
use crate::types::memory::{DocId, MemoryDoc};
|
||||
use crate::types::mission::{Mission, MissionId, MissionStatus};
|
||||
use crate::types::project::{Project, ProjectId};
|
||||
use crate::types::step::Step;
|
||||
use crate::types::thread::{Thread, ThreadId, ThreadState};
|
||||
|
||||
/// Shared in-memory Store implementation for tests.
|
||||
pub struct InMemoryStore {
|
||||
docs: RwLock<Vec<MemoryDoc>>,
|
||||
missions: RwLock<Vec<Mission>>,
|
||||
}
|
||||
|
||||
impl InMemoryStore {
|
||||
pub fn with_docs(docs: Vec<MemoryDoc>) -> Self {
|
||||
Self {
|
||||
docs: RwLock::new(docs),
|
||||
missions: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Store for InMemoryStore {
|
||||
async fn save_thread(&self, _: &Thread) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn load_thread(&self, _: ThreadId) -> Result<Option<Thread>, EngineError> {
|
||||
Ok(None)
|
||||
}
|
||||
async fn list_threads(&self, _: ProjectId) -> Result<Vec<Thread>, EngineError> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn update_thread_state(
|
||||
&self,
|
||||
_: ThreadId,
|
||||
_: ThreadState,
|
||||
) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn save_step(&self, _: &Step) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn load_steps(&self, _: ThreadId) -> Result<Vec<Step>, EngineError> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn append_events(&self, _: &[ThreadEvent]) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn load_events(&self, _: ThreadId) -> Result<Vec<ThreadEvent>, EngineError> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn save_project(&self, _: &Project) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn load_project(&self, _: ProjectId) -> Result<Option<Project>, EngineError> {
|
||||
Ok(None)
|
||||
}
|
||||
async fn list_projects(&self) -> Result<Vec<Project>, EngineError> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn save_conversation(&self, _: &ConversationSurface) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn load_conversation(
|
||||
&self,
|
||||
_: ConversationId,
|
||||
) -> Result<Option<ConversationSurface>, EngineError> {
|
||||
Ok(None)
|
||||
}
|
||||
async fn list_conversations(
|
||||
&self,
|
||||
_: &str,
|
||||
) -> Result<Vec<ConversationSurface>, EngineError> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn save_memory_doc(&self, doc: &MemoryDoc) -> Result<(), EngineError> {
|
||||
let mut docs = self.docs.write().await;
|
||||
docs.retain(|d| d.id != doc.id);
|
||||
docs.push(doc.clone());
|
||||
Ok(())
|
||||
}
|
||||
async fn load_memory_doc(&self, id: DocId) -> Result<Option<MemoryDoc>, EngineError> {
|
||||
Ok(self.docs.read().await.iter().find(|d| d.id == id).cloned())
|
||||
}
|
||||
async fn list_memory_docs(
|
||||
&self,
|
||||
project_id: ProjectId,
|
||||
) -> Result<Vec<MemoryDoc>, EngineError> {
|
||||
Ok(self
|
||||
.docs
|
||||
.read()
|
||||
.await
|
||||
.iter()
|
||||
.filter(|d| d.project_id == project_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
async fn save_lease(&self, _: &CapabilityLease) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn load_active_leases(
|
||||
&self,
|
||||
_: ThreadId,
|
||||
) -> Result<Vec<CapabilityLease>, EngineError> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn revoke_lease(&self, _: LeaseId, _: &str) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn save_mission(&self, mission: &Mission) -> Result<(), EngineError> {
|
||||
let mut missions = self.missions.write().await;
|
||||
missions.retain(|m| m.id != mission.id);
|
||||
missions.push(mission.clone());
|
||||
Ok(())
|
||||
}
|
||||
async fn load_mission(&self, id: MissionId) -> Result<Option<Mission>, EngineError> {
|
||||
Ok(self
|
||||
.missions
|
||||
.read()
|
||||
.await
|
||||
.iter()
|
||||
.find(|m| m.id == id)
|
||||
.cloned())
|
||||
}
|
||||
async fn list_missions(&self, project_id: ProjectId) -> Result<Vec<Mission>, EngineError> {
|
||||
Ok(self
|
||||
.missions
|
||||
.read()
|
||||
.await
|
||||
.iter()
|
||||
.filter(|m| m.project_id == project_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
async fn update_mission_status(
|
||||
&self,
|
||||
id: MissionId,
|
||||
status: MissionStatus,
|
||||
) -> Result<(), EngineError> {
|
||||
let mut missions = self.missions.write().await;
|
||||
if let Some(m) = missions.iter_mut().find(|m| m.id == id) {
|
||||
m.status = status;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -93,6 +93,8 @@ pub async fn reflect(
|
||||
.grant(refl_thread.id, "reflection_tools", vec![], None, None)
|
||||
.await;
|
||||
refl_thread.capability_leases.push(lease.id);
|
||||
store.save_thread(&refl_thread).await?;
|
||||
store.save_lease(&lease).await?;
|
||||
|
||||
// Run the execution loop
|
||||
let mut exec_loop = ExecutionLoop::new(
|
||||
@@ -103,7 +105,8 @@ pub async fn reflect(
|
||||
policy,
|
||||
signal_rx,
|
||||
"system".to_string(),
|
||||
);
|
||||
)
|
||||
.with_store(Arc::clone(store));
|
||||
|
||||
let outcome = exec_loop.run().await?;
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ use tracing::debug;
|
||||
|
||||
use crate::runtime::manager::ThreadManager;
|
||||
use crate::runtime::messaging::ThreadOutcome;
|
||||
use crate::traits::store::Store;
|
||||
use crate::types::conversation::{ConversationEntry, ConversationId, ConversationSurface};
|
||||
use crate::types::error::EngineError;
|
||||
use crate::types::message::ThreadMessage;
|
||||
@@ -27,42 +28,83 @@ use crate::types::thread::{ThreadConfig, ThreadId, ThreadType};
|
||||
/// 3. Create a new conversation if none exists for this channel+user
|
||||
pub struct ConversationManager {
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
store: Arc<dyn Store>,
|
||||
conversations: RwLock<HashMap<ConversationId, ConversationSurface>>,
|
||||
/// Maps (channel, user_id) → conversation ID for lookup.
|
||||
channel_user_index: RwLock<HashMap<(String, String), ConversationId>>,
|
||||
}
|
||||
|
||||
impl ConversationManager {
|
||||
pub fn new(thread_manager: Arc<ThreadManager>) -> Self {
|
||||
pub fn new(thread_manager: Arc<ThreadManager>, store: Arc<dyn Store>) -> Self {
|
||||
Self {
|
||||
thread_manager,
|
||||
store,
|
||||
conversations: RwLock::new(HashMap::new()),
|
||||
channel_user_index: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Restore persisted conversations for a user into the in-memory index.
|
||||
pub async fn bootstrap_user(&self, user_id: &str) -> Result<usize, EngineError> {
|
||||
let conversations = self.store.list_conversations(user_id).await?;
|
||||
let count = conversations.len();
|
||||
let mut convs = self.conversations.write().await;
|
||||
let mut index = self.channel_user_index.write().await;
|
||||
|
||||
for conversation in conversations {
|
||||
index.insert(
|
||||
(conversation.channel.clone(), conversation.user_id.clone()),
|
||||
conversation.id,
|
||||
);
|
||||
convs.insert(conversation.id, conversation);
|
||||
}
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Get or create a conversation for a channel+user pair.
|
||||
pub async fn get_or_create_conversation(&self, channel: &str, user_id: &str) -> ConversationId {
|
||||
pub async fn get_or_create_conversation(
|
||||
&self,
|
||||
channel: &str,
|
||||
user_id: &str,
|
||||
) -> Result<ConversationId, EngineError> {
|
||||
// Check index first
|
||||
let key = (channel.to_string(), user_id.to_string());
|
||||
{
|
||||
let index = self.channel_user_index.read().await;
|
||||
if let Some(conv_id) = index.get(&key) {
|
||||
return *conv_id;
|
||||
return Ok(*conv_id);
|
||||
}
|
||||
}
|
||||
|
||||
// Check persisted conversations for this user/channel.
|
||||
if let Some(conv) = self
|
||||
.store
|
||||
.list_conversations(user_id)
|
||||
.await?
|
||||
.into_iter()
|
||||
.find(|conv| conv.channel == channel)
|
||||
{
|
||||
let conv_id = conv.id;
|
||||
let mut convs = self.conversations.write().await;
|
||||
let mut index = self.channel_user_index.write().await;
|
||||
convs.insert(conv_id, conv);
|
||||
index.insert(key, conv_id);
|
||||
return Ok(conv_id);
|
||||
}
|
||||
|
||||
// Create new conversation
|
||||
let conv = ConversationSurface::new(channel, user_id);
|
||||
let conv_id = conv.id;
|
||||
|
||||
let mut convs = self.conversations.write().await;
|
||||
let mut index = self.channel_user_index.write().await;
|
||||
convs.insert(conv_id, conv);
|
||||
convs.insert(conv_id, conv.clone());
|
||||
index.insert(key, conv_id);
|
||||
self.store.save_conversation(&conv).await?;
|
||||
|
||||
debug!(conversation_id = %conv_id, channel, user_id, "created conversation");
|
||||
conv_id
|
||||
Ok(conv_id)
|
||||
}
|
||||
|
||||
/// Handle an incoming user message.
|
||||
@@ -101,6 +143,7 @@ impl ConversationManager {
|
||||
self.thread_manager
|
||||
.inject_message(thread_id, ThreadMessage::user(content))
|
||||
.await?;
|
||||
self.store.save_conversation(conv).await?;
|
||||
Ok(thread_id)
|
||||
}
|
||||
None => {
|
||||
@@ -126,6 +169,7 @@ impl ConversationManager {
|
||||
thread_id,
|
||||
"Thread started",
|
||||
));
|
||||
self.store.save_conversation(conv).await?;
|
||||
|
||||
debug!(
|
||||
conversation_id = %conversation_id,
|
||||
@@ -143,7 +187,7 @@ impl ConversationManager {
|
||||
conversation_id: ConversationId,
|
||||
thread_id: ThreadId,
|
||||
outcome: &ThreadOutcome,
|
||||
) {
|
||||
) -> Result<(), EngineError> {
|
||||
let mut convs = self.conversations.write().await;
|
||||
if let Some(conv) = convs.get_mut(&conversation_id) {
|
||||
match outcome {
|
||||
@@ -186,7 +230,9 @@ impl ConversationManager {
|
||||
// Thread stays active — waiting for approval
|
||||
}
|
||||
}
|
||||
self.store.save_conversation(conv).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get a snapshot of a conversation.
|
||||
@@ -259,7 +305,7 @@ mod tests {
|
||||
use crate::traits::llm::{LlmBackend, LlmCallConfig, LlmOutput};
|
||||
use crate::traits::store::Store;
|
||||
use crate::types::capability::{ActionDef, CapabilityLease};
|
||||
use crate::types::conversation::EntrySender;
|
||||
use crate::types::conversation::{ConversationId, ConversationSurface, EntrySender};
|
||||
use crate::types::event::ThreadEvent;
|
||||
use crate::types::memory::{DocId, MemoryDoc};
|
||||
use crate::types::project::Project;
|
||||
@@ -322,7 +368,17 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct MockStore;
|
||||
struct MockStore {
|
||||
conversations: RwLock<HashMap<ConversationId, ConversationSurface>>,
|
||||
}
|
||||
|
||||
impl MockStore {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
conversations: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Store for MockStore {
|
||||
@@ -366,6 +422,35 @@ mod tests {
|
||||
async fn load_project(&self, _: ProjectId) -> Result<Option<Project>, EngineError> {
|
||||
Ok(None)
|
||||
}
|
||||
async fn save_conversation(
|
||||
&self,
|
||||
conversation: &ConversationSurface,
|
||||
) -> Result<(), EngineError> {
|
||||
self.conversations
|
||||
.write()
|
||||
.await
|
||||
.insert(conversation.id, conversation.clone());
|
||||
Ok(())
|
||||
}
|
||||
async fn load_conversation(
|
||||
&self,
|
||||
id: ConversationId,
|
||||
) -> Result<Option<ConversationSurface>, EngineError> {
|
||||
Ok(self.conversations.read().await.get(&id).cloned())
|
||||
}
|
||||
async fn list_conversations(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<ConversationSurface>, EngineError> {
|
||||
Ok(self
|
||||
.conversations
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|conversation| conversation.user_id == user_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
async fn save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
}
|
||||
@@ -419,18 +504,19 @@ mod tests {
|
||||
}
|
||||
|
||||
fn make_conv_manager() -> (Arc<ThreadManager>, ConversationManager) {
|
||||
let store = Arc::new(MockStore::new());
|
||||
let tm = Arc::new(ThreadManager::new(
|
||||
Arc::new(MockLlm(Mutex::new(vec![LlmOutput {
|
||||
response: LlmResponse::Text("Hello!".into()),
|
||||
usage: TokenUsage::default(),
|
||||
}]))),
|
||||
Arc::new(MockEffects),
|
||||
Arc::new(MockStore),
|
||||
store.clone(),
|
||||
Arc::new(CapabilityRegistry::new()),
|
||||
Arc::new(LeaseManager::new()),
|
||||
Arc::new(PolicyEngine::new()),
|
||||
));
|
||||
let cm = ConversationManager::new(Arc::clone(&tm));
|
||||
let cm = ConversationManager::new(Arc::clone(&tm), store);
|
||||
(tm, cm)
|
||||
}
|
||||
|
||||
@@ -439,18 +525,27 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn get_or_create_conversation() {
|
||||
let (_, cm) = make_conv_manager();
|
||||
let c1 = cm.get_or_create_conversation("telegram", "user1").await;
|
||||
let c2 = cm.get_or_create_conversation("telegram", "user1").await;
|
||||
let c1 = cm
|
||||
.get_or_create_conversation("telegram", "user1")
|
||||
.await
|
||||
.unwrap();
|
||||
let c2 = cm
|
||||
.get_or_create_conversation("telegram", "user1")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(c1, c2); // same channel+user returns same conversation
|
||||
|
||||
let c3 = cm.get_or_create_conversation("slack", "user1").await;
|
||||
let c3 = cm
|
||||
.get_or_create_conversation("slack", "user1")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_ne!(c1, c3); // different channel → different conversation
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handle_message_spawns_thread() {
|
||||
let (tm, cm) = make_conv_manager();
|
||||
let conv_id = cm.get_or_create_conversation("web", "user1").await;
|
||||
let conv_id = cm.get_or_create_conversation("web", "user1").await.unwrap();
|
||||
let project = ProjectId::new();
|
||||
|
||||
let tid = cm
|
||||
@@ -471,7 +566,7 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn record_outcome_adds_entry() {
|
||||
let (_, cm) = make_conv_manager();
|
||||
let conv_id = cm.get_or_create_conversation("cli", "user1").await;
|
||||
let conv_id = cm.get_or_create_conversation("cli", "user1").await.unwrap();
|
||||
let tid = ThreadId::new();
|
||||
|
||||
// Manually track a thread
|
||||
@@ -489,7 +584,8 @@ mod tests {
|
||||
response: Some("Done!".into()),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let conv = cm.get_conversation(conv_id).await.unwrap();
|
||||
assert!(conv.active_threads.is_empty());
|
||||
@@ -506,9 +602,11 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn list_conversations_filters_by_user() {
|
||||
let (_, cm) = make_conv_manager();
|
||||
cm.get_or_create_conversation("web", "alice").await;
|
||||
cm.get_or_create_conversation("telegram", "alice").await;
|
||||
cm.get_or_create_conversation("web", "bob").await;
|
||||
cm.get_or_create_conversation("web", "alice").await.unwrap();
|
||||
cm.get_or_create_conversation("telegram", "alice")
|
||||
.await
|
||||
.unwrap();
|
||||
cm.get_or_create_conversation("web", "bob").await.unwrap();
|
||||
|
||||
let alice_convs = cm.list_conversations("alice").await;
|
||||
assert_eq!(alice_convs.len(), 2);
|
||||
@@ -516,4 +614,31 @@ mod tests {
|
||||
let bob_convs = cm.list_conversations("bob").await;
|
||||
assert_eq!(bob_convs.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bootstrap_user_loads_persisted_conversations() {
|
||||
let store = Arc::new(MockStore::new());
|
||||
let mut conv = ConversationSurface::new("web", "user1");
|
||||
conv.add_entry(ConversationEntry::user("persisted"));
|
||||
store.save_conversation(&conv).await.unwrap();
|
||||
|
||||
let tm = Arc::new(ThreadManager::new(
|
||||
Arc::new(MockLlm(Mutex::new(vec![]))),
|
||||
Arc::new(MockEffects),
|
||||
store.clone(),
|
||||
Arc::new(CapabilityRegistry::new()),
|
||||
Arc::new(LeaseManager::new()),
|
||||
Arc::new(PolicyEngine::new()),
|
||||
));
|
||||
let cm = ConversationManager::new(tm, store);
|
||||
|
||||
let loaded = cm.bootstrap_user("user1").await.unwrap();
|
||||
assert_eq!(loaded, 1);
|
||||
|
||||
let conv_id = cm.get_or_create_conversation("web", "user1").await.unwrap();
|
||||
assert_eq!(conv_id, conv.id);
|
||||
let saved = cm.get_conversation(conv.id).await.unwrap();
|
||||
assert_eq!(saved.entries.len(), 1);
|
||||
assert_eq!(saved.entries[0].content, "persisted");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ use tokio::sync::RwLock;
|
||||
use tracing::{debug, error};
|
||||
|
||||
use crate::capability::lease::LeaseManager;
|
||||
use crate::capability::planner::LeasePlanner;
|
||||
use crate::capability::policy::PolicyEngine;
|
||||
use crate::capability::registry::CapabilityRegistry;
|
||||
use crate::executor::ExecutionLoop;
|
||||
@@ -18,7 +19,7 @@ use crate::traits::store::Store;
|
||||
use crate::types::error::EngineError;
|
||||
use crate::types::message::ThreadMessage;
|
||||
use crate::types::project::ProjectId;
|
||||
use crate::types::thread::{Thread, ThreadConfig, ThreadId, ThreadType};
|
||||
use crate::types::thread::{Thread, ThreadConfig, ThreadId, ThreadState, ThreadType};
|
||||
|
||||
/// Handle to a running thread for checking results.
|
||||
struct RunningThread {
|
||||
@@ -36,6 +37,7 @@ pub struct ThreadManager {
|
||||
pub capabilities: Arc<CapabilityRegistry>,
|
||||
pub leases: Arc<LeaseManager>,
|
||||
pub policy: Arc<PolicyEngine>,
|
||||
lease_planner: LeasePlanner,
|
||||
tree: RwLock<ThreadTree>,
|
||||
running: RwLock<HashMap<ThreadId, RunningThread>>,
|
||||
/// Broadcast channel for thread events (for live status updates).
|
||||
@@ -59,6 +61,7 @@ impl ThreadManager {
|
||||
capabilities,
|
||||
leases,
|
||||
policy,
|
||||
lease_planner: LeasePlanner::new(),
|
||||
tree: RwLock::new(ThreadTree::new()),
|
||||
running: RwLock::new(HashMap::new()),
|
||||
event_tx,
|
||||
@@ -124,12 +127,22 @@ impl ThreadManager {
|
||||
self.tree.write().await.add_child(pid, thread_id);
|
||||
}
|
||||
|
||||
// Grant leases for all registered capabilities
|
||||
for cap in self.capabilities.list() {
|
||||
// Grant explicit capability leases based on thread type.
|
||||
for grant in self
|
||||
.lease_planner
|
||||
.plan_for_thread(thread_type, &self.capabilities)
|
||||
{
|
||||
let lease = self
|
||||
.leases
|
||||
.grant(thread_id, &cap.name, vec![], None, None)
|
||||
.grant(
|
||||
thread_id,
|
||||
grant.capability_name,
|
||||
grant.granted_actions,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
self.store.save_lease(&lease).await?;
|
||||
thread.capability_leases.push(lease.id);
|
||||
}
|
||||
|
||||
@@ -159,7 +172,8 @@ impl ThreadManager {
|
||||
let exec_loop = ExecutionLoop::new(thread, llm, effects, leases, policy, rx, user_id)
|
||||
.with_capabilities(Arc::clone(&self.capabilities))
|
||||
.with_event_tx(self.event_tx.clone())
|
||||
.with_retrieval(retrieval);
|
||||
.with_retrieval(retrieval)
|
||||
.with_store(Arc::clone(&self.store));
|
||||
|
||||
// Spawn background task
|
||||
let store_for_task = Arc::clone(&self.store);
|
||||
@@ -268,6 +282,13 @@ impl ThreadManager {
|
||||
crate::executor::trace::write_trace(&trace);
|
||||
}
|
||||
|
||||
if let Err(e) = store_for_task.append_events(&exec.thread.events).await {
|
||||
tracing::warn!(
|
||||
thread_id = %thread_id,
|
||||
"failed to persist thread events: {e}"
|
||||
);
|
||||
}
|
||||
|
||||
// Save final thread state to store
|
||||
if let Err(e) = store_for_task.save_thread(&exec.thread).await {
|
||||
tracing::warn!(
|
||||
@@ -373,6 +394,38 @@ impl ThreadManager {
|
||||
}
|
||||
finished
|
||||
}
|
||||
|
||||
/// Reconcile persisted non-terminal threads after process startup.
|
||||
///
|
||||
/// The current engine does not support mid-thread replay/resume, so any
|
||||
/// thread left in a non-terminal state is marked failed-safe.
|
||||
pub async fn recover_project_threads(
|
||||
&self,
|
||||
project_id: ProjectId,
|
||||
) -> Result<Vec<ThreadId>, EngineError> {
|
||||
let threads = self.store.list_threads(project_id).await?;
|
||||
let mut recovered = Vec::new();
|
||||
|
||||
for mut thread in threads {
|
||||
if thread.state.is_terminal() || thread.state == ThreadState::Completed {
|
||||
continue;
|
||||
}
|
||||
|
||||
if thread
|
||||
.transition_to(
|
||||
ThreadState::Failed,
|
||||
Some("engine restart before thread completion".into()),
|
||||
)
|
||||
.is_ok()
|
||||
{
|
||||
self.store.append_events(&thread.events).await?;
|
||||
self.store.save_thread(&thread).await?;
|
||||
recovered.push(thread.id);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(recovered)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -457,18 +510,38 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct MockStore;
|
||||
struct MockStore {
|
||||
threads: RwLock<HashMap<ThreadId, Thread>>,
|
||||
events: RwLock<HashMap<ThreadId, Vec<ThreadEvent>>>,
|
||||
}
|
||||
|
||||
impl MockStore {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
threads: RwLock::new(HashMap::new()),
|
||||
events: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Store for MockStore {
|
||||
async fn save_thread(&self, _: &Thread) -> Result<(), EngineError> {
|
||||
async fn save_thread(&self, thread: &Thread) -> Result<(), EngineError> {
|
||||
self.threads.write().await.insert(thread.id, thread.clone());
|
||||
Ok(())
|
||||
}
|
||||
async fn load_thread(&self, _: ThreadId) -> Result<Option<Thread>, EngineError> {
|
||||
Ok(None)
|
||||
async fn load_thread(&self, id: ThreadId) -> Result<Option<Thread>, EngineError> {
|
||||
Ok(self.threads.read().await.get(&id).cloned())
|
||||
}
|
||||
async fn list_threads(&self, _: ProjectId) -> Result<Vec<Thread>, EngineError> {
|
||||
Ok(vec![])
|
||||
async fn list_threads(&self, project_id: ProjectId) -> Result<Vec<Thread>, EngineError> {
|
||||
Ok(self
|
||||
.threads
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|thread| thread.project_id == project_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
async fn update_thread_state(
|
||||
&self,
|
||||
@@ -483,11 +556,24 @@ mod tests {
|
||||
async fn load_steps(&self, _: ThreadId) -> Result<Vec<Step>, EngineError> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn append_events(&self, _: &[ThreadEvent]) -> Result<(), EngineError> {
|
||||
async fn append_events(&self, events: &[ThreadEvent]) -> Result<(), EngineError> {
|
||||
let mut stored = self.events.write().await;
|
||||
for event in events {
|
||||
stored
|
||||
.entry(event.thread_id)
|
||||
.or_default()
|
||||
.push(event.clone());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
async fn load_events(&self, _: ThreadId) -> Result<Vec<ThreadEvent>, EngineError> {
|
||||
Ok(vec![])
|
||||
async fn load_events(&self, thread_id: ThreadId) -> Result<Vec<ThreadEvent>, EngineError> {
|
||||
Ok(self
|
||||
.events
|
||||
.read()
|
||||
.await
|
||||
.get(&thread_id)
|
||||
.cloned()
|
||||
.unwrap_or_default())
|
||||
}
|
||||
async fn save_project(&self, _: &Project) -> Result<(), EngineError> {
|
||||
Ok(())
|
||||
@@ -566,7 +652,33 @@ mod tests {
|
||||
ThreadManager::new(
|
||||
llm,
|
||||
Arc::new(MockEffects),
|
||||
Arc::new(MockStore),
|
||||
Arc::new(MockStore::new()),
|
||||
Arc::new(caps),
|
||||
Arc::new(LeaseManager::new()),
|
||||
Arc::new(PolicyEngine::new()),
|
||||
)
|
||||
}
|
||||
|
||||
fn make_manager_with_store(llm: Arc<dyn LlmBackend>, store: Arc<MockStore>) -> ThreadManager {
|
||||
let mut caps = CapabilityRegistry::new();
|
||||
caps.register(Capability {
|
||||
name: "test".into(),
|
||||
description: "Test capability".into(),
|
||||
actions: vec![ActionDef {
|
||||
name: "test_tool".into(),
|
||||
description: "Test".into(),
|
||||
parameters_schema: serde_json::json!({}),
|
||||
effects: vec![EffectType::ReadLocal],
|
||||
requires_approval: false,
|
||||
}],
|
||||
knowledge: vec![],
|
||||
policies: vec![],
|
||||
});
|
||||
|
||||
ThreadManager::new(
|
||||
llm,
|
||||
Arc::new(MockEffects),
|
||||
store,
|
||||
Arc::new(caps),
|
||||
Arc::new(LeaseManager::new()),
|
||||
Arc::new(PolicyEngine::new()),
|
||||
@@ -673,4 +785,39 @@ mod tests {
|
||||
assert_eq!(mgr.parent_of(child).await, Some(parent));
|
||||
assert_eq!(mgr.children_of(parent).await, vec![child]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn recover_project_threads_marks_non_terminal_as_failed() {
|
||||
let store = Arc::new(MockStore::new());
|
||||
let project = ProjectId::new();
|
||||
|
||||
let mut running = Thread::new(
|
||||
"running",
|
||||
ThreadType::Foreground,
|
||||
project,
|
||||
ThreadConfig::default(),
|
||||
);
|
||||
running.transition_to(ThreadState::Running, None).unwrap();
|
||||
store.save_thread(&running).await.unwrap();
|
||||
|
||||
let mut completed = Thread::new(
|
||||
"done",
|
||||
ThreadType::Foreground,
|
||||
project,
|
||||
ThreadConfig::default(),
|
||||
);
|
||||
completed
|
||||
.transition_to(ThreadState::Failed, Some("already terminal".into()))
|
||||
.unwrap();
|
||||
store.save_thread(&completed).await.unwrap();
|
||||
|
||||
let mgr = make_manager_with_store(MockLlm::text("ignored"), Arc::clone(&store));
|
||||
let recovered = mgr.recover_project_threads(project).await.unwrap();
|
||||
|
||||
assert_eq!(recovered, vec![running.id]);
|
||||
let saved = store.load_thread(running.id).await.unwrap().unwrap();
|
||||
assert_eq!(saved.state, ThreadState::Failed);
|
||||
let events = store.load_events(running.id).await.unwrap();
|
||||
assert!(!events.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -36,6 +36,21 @@ impl MissionManager {
|
||||
}
|
||||
}
|
||||
|
||||
/// Populate the active mission index from persisted mission state.
|
||||
pub async fn bootstrap_project(&self, project_id: ProjectId) -> Result<usize, EngineError> {
|
||||
let missions = self.store.list_missions(project_id).await?;
|
||||
let active_ids: Vec<MissionId> = missions
|
||||
.into_iter()
|
||||
.filter(|mission| mission.status == MissionStatus::Active)
|
||||
.map(|mission| mission.id)
|
||||
.collect();
|
||||
|
||||
let count = active_ids.len();
|
||||
*self.active.write().await = active_ids;
|
||||
debug!(project_id = ?project_id, active_missions = count, "bootstrapped active missions");
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Create and persist a new mission. Returns the mission ID.
|
||||
pub async fn create_mission(
|
||||
&self,
|
||||
@@ -57,6 +72,7 @@ impl MissionManager {
|
||||
self.store
|
||||
.update_mission_status(id, MissionStatus::Paused)
|
||||
.await?;
|
||||
self.active.write().await.retain(|mid| *mid != id);
|
||||
debug!(mission_id = %id, "mission paused");
|
||||
Ok(())
|
||||
}
|
||||
@@ -66,6 +82,10 @@ impl MissionManager {
|
||||
self.store
|
||||
.update_mission_status(id, MissionStatus::Active)
|
||||
.await?;
|
||||
let mut active = self.active.write().await;
|
||||
if !active.contains(&id) {
|
||||
active.push(id);
|
||||
}
|
||||
debug!(mission_id = %id, "mission resumed");
|
||||
Ok(())
|
||||
}
|
||||
@@ -196,6 +216,218 @@ impl MissionManager {
|
||||
self.store.load_mission(id).await
|
||||
}
|
||||
|
||||
/// Fire all active `OnSystemEvent` missions whose source and event_type match.
|
||||
///
|
||||
/// The optional `payload` is forwarded as `trigger_payload` to each mission's
|
||||
/// thread, carrying context like trace issues and reflection docs.
|
||||
pub async fn fire_on_system_event(
|
||||
&self,
|
||||
source: &str,
|
||||
event_type: &str,
|
||||
user_id: &str,
|
||||
payload: Option<serde_json::Value>,
|
||||
) -> Result<Vec<ThreadId>, EngineError> {
|
||||
let active_ids = self.active.read().await.clone();
|
||||
let mut spawned = Vec::new();
|
||||
|
||||
for mid in active_ids {
|
||||
let mission = match self.store.load_mission(mid).await? {
|
||||
Some(m) if m.status == MissionStatus::Active => m,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let matches = match &mission.cadence {
|
||||
MissionCadence::OnSystemEvent {
|
||||
source: s,
|
||||
event_type: et,
|
||||
} => s == source && et == event_type,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if matches && let Some(tid) = self.fire_mission(mid, user_id, payload.clone()).await? {
|
||||
spawned.push(tid);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(spawned)
|
||||
}
|
||||
|
||||
/// Start a background event listener that fires `OnSystemEvent` missions
|
||||
/// when threads complete with issues.
|
||||
///
|
||||
/// Subscribes to the ThreadManager's event broadcast channel and watches
|
||||
/// for thread completion events. When a non-Mission, non-Reflection thread
|
||||
/// completes and its trace has issues, fires matching OnSystemEvent missions
|
||||
/// with trace data as the trigger payload.
|
||||
pub fn start_event_listener(self: &Arc<Self>, user_id: String) {
|
||||
let mgr = Arc::clone(self);
|
||||
let mut rx = mgr.thread_manager.subscribe_events();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
match rx.recv().await {
|
||||
Ok(event) => {
|
||||
// React to ReflectionComplete — the thread is done and
|
||||
// we have reflection doc info for the trigger payload.
|
||||
if let crate::types::event::EventKind::ReflectionComplete {
|
||||
docs_produced,
|
||||
ref doc_types,
|
||||
..
|
||||
} = event.kind
|
||||
{
|
||||
// Load the thread to check its type and build the payload
|
||||
let thread = mgr.store.load_thread(event.thread_id).await;
|
||||
let thread = match thread {
|
||||
Ok(Some(t)) => t,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
// Skip Mission and Reflection threads (no recursive self-improvement)
|
||||
if matches!(
|
||||
thread.thread_type,
|
||||
ThreadType::Mission | ThreadType::Reflection
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Build trace to check for issues
|
||||
let trace = crate::executor::trace::build_trace(&thread);
|
||||
if trace.issues.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Build trigger payload with trace issues, error messages,
|
||||
// and reflection summary
|
||||
let issues: Vec<serde_json::Value> = trace
|
||||
.issues
|
||||
.iter()
|
||||
.map(|i| {
|
||||
serde_json::json!({
|
||||
"severity": format!("{:?}", i.severity),
|
||||
"category": i.category,
|
||||
"description": i.description,
|
||||
"step": i.step,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Extract actual error text from ActionFailed events
|
||||
// and system messages (these contain the real diagnostics)
|
||||
let error_messages: Vec<String> = thread
|
||||
.events
|
||||
.iter()
|
||||
.filter_map(|e| {
|
||||
if let crate::types::event::EventKind::ActionFailed {
|
||||
action_name,
|
||||
error,
|
||||
..
|
||||
} = &e.kind
|
||||
{
|
||||
Some(format!("{action_name}: {error}"))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.take(10) // cap to avoid bloating payload
|
||||
.collect();
|
||||
|
||||
let payload = serde_json::json!({
|
||||
"source_thread_id": event.thread_id.0.to_string(),
|
||||
"goal": thread.goal,
|
||||
"issues": issues,
|
||||
"error_messages": error_messages,
|
||||
"reflection": {
|
||||
"docs_produced": docs_produced,
|
||||
"doc_types": doc_types,
|
||||
},
|
||||
});
|
||||
|
||||
if let Err(e) = mgr
|
||||
.fire_on_system_event(
|
||||
"engine",
|
||||
"thread_completed_with_issues",
|
||||
&user_id,
|
||||
Some(payload),
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!("event listener: failed to fire self-improvement: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
|
||||
debug!("event listener: lagged {n} events");
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Ensure a self-improvement mission exists for the given project.
|
||||
///
|
||||
/// Checks if a mission with `"self_improvement": true` in metadata already
|
||||
/// exists. If not, creates one with `OnSystemEvent` cadence that fires
|
||||
/// when threads complete with issues. Also seeds the fix pattern database.
|
||||
///
|
||||
/// Returns the mission ID (existing or newly created).
|
||||
pub async fn ensure_self_improvement_mission(
|
||||
&self,
|
||||
project_id: ProjectId,
|
||||
) -> Result<MissionId, EngineError> {
|
||||
// Check if one already exists
|
||||
let missions = self.store.list_missions(project_id).await?;
|
||||
if let Some(existing) = missions.iter().find(|m| is_self_improvement_mission(m)) {
|
||||
debug!(mission_id = %existing.id, "self-improvement mission already exists");
|
||||
// Make sure it's in the active list
|
||||
let mut active = self.active.write().await;
|
||||
if !active.contains(&existing.id) {
|
||||
active.push(existing.id);
|
||||
}
|
||||
return Ok(existing.id);
|
||||
}
|
||||
|
||||
// Create the self-improvement mission
|
||||
let mut mission = Mission::new(
|
||||
project_id,
|
||||
"self-improvement",
|
||||
SELF_IMPROVEMENT_GOAL,
|
||||
MissionCadence::OnSystemEvent {
|
||||
source: "engine".into(),
|
||||
event_type: "thread_completed_with_issues".into(),
|
||||
},
|
||||
);
|
||||
mission.success_criteria = Some(
|
||||
"Continuously improve system prompts and fix patterns based on execution traces".into(),
|
||||
);
|
||||
mission.metadata = serde_json::json!({"self_improvement": true});
|
||||
mission.max_threads_per_day = 5;
|
||||
|
||||
let id = mission.id;
|
||||
self.store.save_mission(&mission).await?;
|
||||
self.active.write().await.push(id);
|
||||
|
||||
// Seed the fix pattern database if it doesn't exist
|
||||
let docs = self.store.list_memory_docs(project_id).await?;
|
||||
let has_patterns = docs.iter().any(|d| {
|
||||
d.title == FIX_PATTERN_DB_TITLE && d.tags.contains(&FIX_PATTERN_DB_TAG.to_string())
|
||||
});
|
||||
if !has_patterns {
|
||||
use crate::types::memory::{DocType, MemoryDoc};
|
||||
let pattern_doc = MemoryDoc::new(
|
||||
project_id,
|
||||
DocType::Playbook,
|
||||
FIX_PATTERN_DB_TITLE,
|
||||
SEED_FIX_PATTERNS,
|
||||
)
|
||||
.with_tags(vec![FIX_PATTERN_DB_TAG.to_string()]);
|
||||
self.store.save_memory_doc(&pattern_doc).await?;
|
||||
debug!("seeded fix pattern database");
|
||||
}
|
||||
|
||||
debug!(mission_id = %id, "created self-improvement mission");
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
/// Tick — check all active missions and fire any that are due.
|
||||
///
|
||||
/// For `Cron` cadence missions, checks `next_fire_at` against current time.
|
||||
@@ -314,6 +546,8 @@ When done, call FINAL() with your response. Include:\n\
|
||||
/// Process a completed mission thread's outcome.
|
||||
///
|
||||
/// Extracts next_focus from the FINAL() response and updates the mission.
|
||||
/// For self-improvement missions (metadata contains `"self_improvement": true`),
|
||||
/// also processes prompt overlay additions and fix pattern updates.
|
||||
async fn process_mission_outcome(
|
||||
store: &Arc<dyn Store>,
|
||||
mission_id: MissionId,
|
||||
@@ -353,6 +587,16 @@ async fn process_mission_outcome(
|
||||
// Record approach
|
||||
let accomplishment: String = text.chars().take(200).collect();
|
||||
mission.approach_history.push(accomplishment);
|
||||
|
||||
// If this is a self-improvement mission, process structured output
|
||||
if is_self_improvement_mission(&mission)
|
||||
&& let Err(e) = process_self_improvement_output(store, &mission, text).await
|
||||
{
|
||||
warn!(
|
||||
mission_id = %mission_id,
|
||||
"failed to process self-improvement output: {e}"
|
||||
);
|
||||
}
|
||||
}
|
||||
ThreadOutcome::Completed { response: None } => {}
|
||||
ThreadOutcome::Failed { error } => {
|
||||
@@ -370,6 +614,250 @@ async fn process_mission_outcome(
|
||||
store.save_mission(&mission).await
|
||||
}
|
||||
|
||||
/// Check if a mission is the self-improvement mission.
|
||||
fn is_self_improvement_mission(mission: &Mission) -> bool {
|
||||
mission
|
||||
.metadata
|
||||
.get("self_improvement")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Process output from a self-improvement mission thread.
|
||||
///
|
||||
/// Two paths:
|
||||
/// 1. The agent used tools directly (memory_write for prompt overlay, shell for
|
||||
/// code fixes) — in this case the FINAL() response is just a summary and
|
||||
/// there is nothing extra to do here.
|
||||
/// 2. The agent returned structured JSON with `prompt_additions` and/or
|
||||
/// `fix_patterns` — we apply those to the Store.
|
||||
///
|
||||
/// This function handles path 2. Path 1 is handled by the tools themselves.
|
||||
async fn process_self_improvement_output(
|
||||
store: &Arc<dyn Store>,
|
||||
mission: &Mission,
|
||||
response: &str,
|
||||
) -> Result<(), EngineError> {
|
||||
use crate::executor::prompt::{PREAMBLE_OVERLAY_TITLE, PROMPT_OVERLAY_TAG};
|
||||
use crate::types::memory::{DocType, MemoryDoc};
|
||||
|
||||
// Try to extract JSON from the response. If the agent used tools directly
|
||||
// (the preferred autoresearch-style path), there's no JSON and we return
|
||||
// early — the work was already done via tool calls.
|
||||
let json_val = match extract_json_from_response(response) {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
debug!("self-improvement: no structured JSON in response (agent likely used tools directly)");
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
let project_id = mission.project_id;
|
||||
|
||||
// Process prompt additions
|
||||
if let Some(additions) = json_val.get("prompt_additions").and_then(|v| v.as_array())
|
||||
&& !additions.is_empty()
|
||||
{
|
||||
let new_rules: Vec<String> = additions
|
||||
.iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect();
|
||||
|
||||
if !new_rules.is_empty() {
|
||||
// Load or create the prompt overlay doc
|
||||
let docs = store.list_memory_docs(project_id).await?;
|
||||
let existing = docs.iter().find(|d| {
|
||||
d.title == PREAMBLE_OVERLAY_TITLE
|
||||
&& d.tags.contains(&PROMPT_OVERLAY_TAG.to_string())
|
||||
});
|
||||
|
||||
let mut overlay = if let Some(doc) = existing {
|
||||
doc.clone()
|
||||
} else {
|
||||
MemoryDoc::new(project_id, DocType::Note, PREAMBLE_OVERLAY_TITLE, "")
|
||||
.with_tags(vec![PROMPT_OVERLAY_TAG.to_string()])
|
||||
};
|
||||
|
||||
// Append new rules
|
||||
for rule in &new_rules {
|
||||
if !overlay.content.is_empty() {
|
||||
overlay.content.push('\n');
|
||||
}
|
||||
overlay.content.push_str(rule);
|
||||
}
|
||||
overlay.updated_at = chrono::Utc::now();
|
||||
|
||||
store.save_memory_doc(&overlay).await?;
|
||||
debug!(
|
||||
rules_added = new_rules.len(),
|
||||
"self-improvement: updated prompt overlay"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Process fix patterns
|
||||
if let Some(patterns) = json_val.get("fix_patterns").and_then(|v| v.as_array())
|
||||
&& !patterns.is_empty()
|
||||
{
|
||||
let docs = store.list_memory_docs(project_id).await?;
|
||||
let existing = docs.iter().find(|d| {
|
||||
d.title == FIX_PATTERN_DB_TITLE && d.tags.contains(&FIX_PATTERN_DB_TAG.to_string())
|
||||
});
|
||||
|
||||
let mut pattern_doc = if let Some(doc) = existing {
|
||||
doc.clone()
|
||||
} else {
|
||||
MemoryDoc::new(
|
||||
project_id,
|
||||
DocType::Playbook,
|
||||
FIX_PATTERN_DB_TITLE,
|
||||
SEED_FIX_PATTERNS,
|
||||
)
|
||||
.with_tags(vec![FIX_PATTERN_DB_TAG.to_string()])
|
||||
};
|
||||
|
||||
for pattern in patterns {
|
||||
let p = pattern
|
||||
.get("pattern")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
let s = pattern
|
||||
.get("strategy")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
let l = pattern
|
||||
.get("location")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
if !p.is_empty() {
|
||||
pattern_doc
|
||||
.content
|
||||
.push_str(&format!("\n| {p} | {s} | {l} |"));
|
||||
}
|
||||
}
|
||||
pattern_doc.updated_at = chrono::Utc::now();
|
||||
|
||||
store.save_memory_doc(&pattern_doc).await?;
|
||||
debug!(
|
||||
patterns_added = patterns.len(),
|
||||
"self-improvement: updated fix pattern database"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Try to extract a JSON object from a response string.
|
||||
///
|
||||
/// Looks for `{...}` in the text, trying the whole string first,
|
||||
/// then searching for embedded JSON.
|
||||
fn extract_json_from_response(response: &str) -> Option<serde_json::Value> {
|
||||
// Try parsing the whole response as JSON
|
||||
if let Ok(v) = serde_json::from_str::<serde_json::Value>(response)
|
||||
&& v.is_object()
|
||||
{
|
||||
return Some(v);
|
||||
}
|
||||
|
||||
// Search for embedded JSON object
|
||||
let start = response.find('{')?;
|
||||
let end = response.rfind('}')?;
|
||||
if end <= start {
|
||||
return None;
|
||||
}
|
||||
let candidate = &response[start..=end];
|
||||
serde_json::from_str::<serde_json::Value>(candidate)
|
||||
.ok()
|
||||
.filter(|v| v.is_object())
|
||||
}
|
||||
|
||||
/// The goal for the self-improvement mission (autoresearch-style program).
|
||||
///
|
||||
/// This is the "program.md" — a concrete, step-by-step prompt that tells the
|
||||
/// agent exactly what to do. Inspired by karpathy/autoresearch: the entire
|
||||
/// research org is a markdown file with an explicit loop.
|
||||
const SELF_IMPROVEMENT_GOAL: &str = "\
|
||||
You are a self-improvement agent for the IronClaw engine. You receive trigger \
|
||||
payloads containing execution trace issues from completed threads. Your job is \
|
||||
to diagnose root causes and apply fixes so the same issue doesn't recur.
|
||||
|
||||
## What you have access to
|
||||
|
||||
- `state[\"trigger_payload\"]` — JSON with `issues` (list of {severity, category, description, step}), \
|
||||
`error_messages` (actual error text from failed actions), `goal` (what the thread was trying to do), \
|
||||
and `source_thread_id`.
|
||||
- All tools: shell, read_file, write_file, apply_patch, web_search, memory_write, etc.
|
||||
- The codebase at the current working directory.
|
||||
- The fix pattern database in prior knowledge (if loaded).
|
||||
|
||||
## The experiment loop
|
||||
|
||||
For each issue in the trigger payload:
|
||||
|
||||
1. **Diagnose**: Read the error messages and issue descriptions. Classify the root cause:
|
||||
- PROMPT: The LLM made a mistake because the system prompt is missing a rule \
|
||||
(wrong tool name, bad API usage, ignoring tool results)
|
||||
- CONFIG: A default value is wrong (truncation length, iteration limit, timeout)
|
||||
- CODE: There is a bug in the engine or bridge code (crash, type error, missing conversion)
|
||||
|
||||
2. **Check the fix pattern database** in prior knowledge. Has this pattern been seen before? \
|
||||
If yes, apply the known strategy. If no, proceed to step 3.
|
||||
|
||||
3. **Apply the fix** based on the level:
|
||||
|
||||
Level 1 (PROMPT — low risk, apply directly):
|
||||
- Read the current prompt overlay: `memory_search(\"prompt:codeact_preamble\")`
|
||||
- Write an updated overlay with a new rule appended
|
||||
- Use `memory_write` with title=\"prompt:codeact_preamble\" and tags=[\"prompt_overlay\"]
|
||||
- The rule should be specific and actionable (e.g. \"Never call web_fetch — use http() instead\")
|
||||
|
||||
Level 2 (CONFIG — medium risk):
|
||||
- Use `read_file` to find the relevant constant or default
|
||||
- Use `shell` to create a git branch: `git checkout -b self-improve/issue-description`
|
||||
- Apply the change with `apply_patch` or `write_file`
|
||||
- Run tests: `cargo test -p ironclaw_engine`
|
||||
- If tests pass, commit. If not, revert: `git checkout main`
|
||||
|
||||
Level 3 (CODE — high risk, just propose):
|
||||
- Read the relevant source files
|
||||
- Describe the fix needed but DO NOT apply it directly
|
||||
- Log it as a recommendation in your FINAL() response
|
||||
|
||||
4. **Record what you did** — include in your FINAL() response:
|
||||
- What issue you analyzed
|
||||
- What level fix you applied (1/2/3)
|
||||
- What specific change you made
|
||||
- Next focus: what to look for next time
|
||||
|
||||
## Important rules
|
||||
|
||||
- Be specific. \"Never call web_fetch\" is good. \"Be careful with tool names\" is useless.
|
||||
- One fix per issue. Don't try to fix everything at once.
|
||||
- For Level 1 fixes, the rule must be one sentence that can be appended to the prompt.
|
||||
- If the trigger payload has no actionable issues (only Info severity), skip and call FINAL() immediately.
|
||||
- NEVER modify test files to make a fix pass.
|
||||
- NEVER modify security-sensitive code (safety layer, policy engine, leak detection).
|
||||
- If you can't diagnose the root cause after reading the errors, log it and move on.";
|
||||
|
||||
/// Well-known title for the fix pattern database.
|
||||
pub const FIX_PATTERN_DB_TITLE: &str = "fix_pattern_database";
|
||||
|
||||
/// Well-known tag for the fix pattern database.
|
||||
pub const FIX_PATTERN_DB_TAG: &str = "fix_patterns";
|
||||
|
||||
/// Seed content for the fix pattern database.
|
||||
const SEED_FIX_PATTERNS: &str = "\
|
||||
| Trace pattern | Fix strategy | Location pattern |
|
||||
|---|---|---|
|
||||
| Tool X not found | Add name alias or prompt hint about correct name | prompt overlay or effect_adapter |
|
||||
| TypeError: str indices must be integers | Parse JSON before wrapping | Where tool output is converted |
|
||||
| NameError: name 'X' not defined | Add prompt hint about using state dict | prompt overlay |
|
||||
| byte index N is not a char boundary | Replace byte slicing with chars().take(N) | Code that slices strings |
|
||||
| Model calls nonexistent tool | Add prompt rule listing correct tool name | prompt overlay |
|
||||
| Model ignores tool results | Improve output metadata format | prompt overlay |
|
||||
| Excessive steps (>5) for simple task | Add prompt rule or fix tool schema | prompt overlay |
|
||||
| Code error in REPL output | Add prompt hint about correct API usage | prompt overlay |";
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -398,6 +886,7 @@ mod tests {
|
||||
struct TestStore {
|
||||
threads: tokio::sync::RwLock<HashMap<ThreadId, Thread>>,
|
||||
missions: tokio::sync::RwLock<HashMap<MissionId, Mission>>,
|
||||
docs: tokio::sync::RwLock<Vec<MemoryDoc>>,
|
||||
}
|
||||
|
||||
impl TestStore {
|
||||
@@ -405,6 +894,7 @@ mod tests {
|
||||
Self {
|
||||
threads: tokio::sync::RwLock::new(HashMap::new()),
|
||||
missions: tokio::sync::RwLock::new(HashMap::new()),
|
||||
docs: tokio::sync::RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -454,15 +944,28 @@ mod tests {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
// ── MemoryDoc (noop) ──
|
||||
async fn save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> {
|
||||
// ── MemoryDoc ──
|
||||
async fn save_memory_doc(&self, doc: &MemoryDoc) -> Result<(), EngineError> {
|
||||
let mut docs = self.docs.write().await;
|
||||
docs.retain(|d| d.id != doc.id);
|
||||
docs.push(doc.clone());
|
||||
Ok(())
|
||||
}
|
||||
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> {
|
||||
Ok(None)
|
||||
async fn load_memory_doc(&self, id: DocId) -> Result<Option<MemoryDoc>, EngineError> {
|
||||
Ok(self.docs.read().await.iter().find(|d| d.id == id).cloned())
|
||||
}
|
||||
async fn list_memory_docs(&self, _: ProjectId) -> Result<Vec<MemoryDoc>, EngineError> {
|
||||
Ok(vec![])
|
||||
async fn list_memory_docs(
|
||||
&self,
|
||||
project_id: ProjectId,
|
||||
) -> Result<Vec<MemoryDoc>, EngineError> {
|
||||
Ok(self
|
||||
.docs
|
||||
.read()
|
||||
.await
|
||||
.iter()
|
||||
.filter(|d| d.project_id == project_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
// ── Lease (noop) ──
|
||||
@@ -1004,6 +1507,228 @@ mod tests {
|
||||
assert_eq!(mission.threads_today, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fire_on_system_event_matches_cadence() {
|
||||
let store = Arc::new(TestStore::new());
|
||||
let mgr = make_mission_manager_with_response(Arc::clone(&store) as Arc<dyn Store>, "done");
|
||||
let project_id = ProjectId::new();
|
||||
|
||||
// Create an OnSystemEvent mission
|
||||
mgr.create_mission(
|
||||
project_id,
|
||||
"self-improve",
|
||||
"improve prompts",
|
||||
MissionCadence::OnSystemEvent {
|
||||
source: "engine".into(),
|
||||
event_type: "thread_completed_with_issues".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let spawned = mgr
|
||||
.fire_on_system_event(
|
||||
"engine",
|
||||
"thread_completed_with_issues",
|
||||
"test-user",
|
||||
Some(serde_json::json!({"issues": []})),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(spawned.len(), 1, "should fire the matching mission");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fire_on_system_event_ignores_non_matching() {
|
||||
let store = Arc::new(TestStore::new());
|
||||
let mgr = make_mission_manager_with_response(Arc::clone(&store) as Arc<dyn Store>, "done");
|
||||
let project_id = ProjectId::new();
|
||||
|
||||
// Create an OnSystemEvent mission for a different event
|
||||
mgr.create_mission(
|
||||
project_id,
|
||||
"webhook handler",
|
||||
"handle webhooks",
|
||||
MissionCadence::OnSystemEvent {
|
||||
source: "github".into(),
|
||||
event_type: "push".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let spawned = mgr
|
||||
.fire_on_system_event("engine", "thread_completed_with_issues", "test-user", None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(spawned.len(), 0, "should not fire non-matching mission");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fire_on_system_event_skips_manual_and_cron() {
|
||||
let store = Arc::new(TestStore::new());
|
||||
let mgr = make_mission_manager_with_response(Arc::clone(&store) as Arc<dyn Store>, "done");
|
||||
let project_id = ProjectId::new();
|
||||
|
||||
mgr.create_mission(project_id, "manual", "goal", MissionCadence::Manual)
|
||||
.await
|
||||
.unwrap();
|
||||
mgr.create_mission(
|
||||
project_id,
|
||||
"cron",
|
||||
"goal",
|
||||
MissionCadence::Cron {
|
||||
expression: "* * * * *".into(),
|
||||
timezone: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let spawned = mgr
|
||||
.fire_on_system_event("engine", "thread_completed_with_issues", "test-user", None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(spawned.len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn self_improvement_outcome_saves_prompt_overlay() {
|
||||
let store: Arc<dyn Store> = Arc::new(TestStore::new());
|
||||
let project_id = ProjectId::new();
|
||||
|
||||
let mut mission = Mission::new(
|
||||
project_id,
|
||||
"self-improve",
|
||||
"improve prompts",
|
||||
MissionCadence::OnSystemEvent {
|
||||
source: "engine".into(),
|
||||
event_type: "thread_completed_with_issues".into(),
|
||||
},
|
||||
);
|
||||
mission.metadata = serde_json::json!({"self_improvement": true});
|
||||
let id = mission.id;
|
||||
store.save_mission(&mission).await.unwrap();
|
||||
|
||||
let response = r#"{"prompt_additions": ["9. Never call web_fetch — use http() instead."], "fix_patterns": [], "level": 1}"#;
|
||||
let outcome = ThreadOutcome::Completed {
|
||||
response: Some(response.into()),
|
||||
};
|
||||
process_mission_outcome(&store, id, ThreadId::new(), &outcome)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Verify prompt overlay was saved
|
||||
let docs = store.list_memory_docs(project_id).await.unwrap();
|
||||
let overlay = docs
|
||||
.iter()
|
||||
.find(|d| d.title == crate::executor::prompt::PREAMBLE_OVERLAY_TITLE);
|
||||
assert!(overlay.is_some(), "prompt overlay should be saved");
|
||||
assert!(overlay.unwrap().content.contains("Never call web_fetch"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn self_improvement_outcome_saves_fix_patterns() {
|
||||
let store: Arc<dyn Store> = Arc::new(TestStore::new());
|
||||
let project_id = ProjectId::new();
|
||||
|
||||
let mut mission = Mission::new(
|
||||
project_id,
|
||||
"self-improve",
|
||||
"improve prompts",
|
||||
MissionCadence::Manual,
|
||||
);
|
||||
mission.metadata = serde_json::json!({"self_improvement": true});
|
||||
let id = mission.id;
|
||||
store.save_mission(&mission).await.unwrap();
|
||||
|
||||
let response = r#"{"prompt_additions": [], "fix_patterns": [{"pattern": "Tool xyz not found", "strategy": "Add alias xyz -> x-y-z", "location": "effect_adapter"}]}"#;
|
||||
let outcome = ThreadOutcome::Completed {
|
||||
response: Some(response.into()),
|
||||
};
|
||||
process_mission_outcome(&store, id, ThreadId::new(), &outcome)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let docs = store.list_memory_docs(project_id).await.unwrap();
|
||||
let patterns = docs.iter().find(|d| d.title == FIX_PATTERN_DB_TITLE);
|
||||
assert!(patterns.is_some(), "fix patterns should be saved");
|
||||
assert!(patterns.unwrap().content.contains("Tool xyz not found"));
|
||||
// Should also contain seed patterns
|
||||
assert!(patterns.unwrap().content.contains("NameError"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn non_self_improvement_mission_skips_structured_output() {
|
||||
let store: Arc<dyn Store> = Arc::new(TestStore::new());
|
||||
let project_id = ProjectId::new();
|
||||
|
||||
let mission = Mission::new(project_id, "regular", "do stuff", MissionCadence::Manual);
|
||||
let id = mission.id;
|
||||
store.save_mission(&mission).await.unwrap();
|
||||
|
||||
// Even if the response has JSON, it should not create overlays
|
||||
let response = r#"{"prompt_additions": ["should not appear"], "level": 1}"#;
|
||||
let outcome = ThreadOutcome::Completed {
|
||||
response: Some(response.into()),
|
||||
};
|
||||
process_mission_outcome(&store, id, ThreadId::new(), &outcome)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let docs = store.list_memory_docs(project_id).await.unwrap();
|
||||
assert!(docs.is_empty(), "non-SI mission should not create overlay");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ensure_self_improvement_mission_creates_on_first_call() {
|
||||
let store = Arc::new(TestStore::new());
|
||||
let mgr = make_mission_manager(Arc::clone(&store) as Arc<dyn Store>);
|
||||
let project_id = ProjectId::new();
|
||||
|
||||
let id = mgr
|
||||
.ensure_self_improvement_mission(project_id)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mission = mgr.get_mission(id).await.unwrap().unwrap();
|
||||
assert_eq!(mission.name, "self-improvement");
|
||||
assert!(is_self_improvement_mission(&mission));
|
||||
assert!(matches!(
|
||||
mission.cadence,
|
||||
MissionCadence::OnSystemEvent { .. }
|
||||
));
|
||||
assert_eq!(mission.max_threads_per_day, 5);
|
||||
|
||||
// Fix pattern database should be seeded
|
||||
let docs = store.list_memory_docs(project_id).await.unwrap();
|
||||
let patterns = docs.iter().find(|d| d.title == FIX_PATTERN_DB_TITLE);
|
||||
assert!(patterns.is_some(), "fix patterns should be seeded");
|
||||
assert!(patterns.unwrap().content.contains("NameError"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ensure_self_improvement_mission_idempotent() {
|
||||
let store = Arc::new(TestStore::new());
|
||||
let mgr = make_mission_manager(Arc::clone(&store) as Arc<dyn Store>);
|
||||
let project_id = ProjectId::new();
|
||||
|
||||
let id1 = mgr
|
||||
.ensure_self_improvement_mission(project_id)
|
||||
.await
|
||||
.unwrap();
|
||||
let id2 = mgr
|
||||
.ensure_self_improvement_mission(project_id)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(id1, id2, "should return the same mission ID");
|
||||
|
||||
// Should only have one mission
|
||||
let missions = store.list_missions(project_id).await.unwrap();
|
||||
assert_eq!(missions.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn daily_budget_enforced() {
|
||||
let store = Arc::new(TestStore::new());
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
//! this by wrapping its dual-backend `Database` trait (PostgreSQL + libSQL).
|
||||
|
||||
use crate::types::capability::{CapabilityLease, LeaseId};
|
||||
use crate::types::conversation::{ConversationId, ConversationSurface};
|
||||
use crate::types::error::EngineError;
|
||||
use crate::types::event::ThreadEvent;
|
||||
use crate::types::memory::{DocId, MemoryDoc};
|
||||
@@ -40,6 +41,33 @@ pub trait Store: Send + Sync {
|
||||
|
||||
async fn save_project(&self, project: &Project) -> Result<(), EngineError>;
|
||||
async fn load_project(&self, id: ProjectId) -> Result<Option<Project>, EngineError>;
|
||||
async fn list_projects(&self) -> Result<Vec<Project>, EngineError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
// ── Conversation operations ─────────────────────────────
|
||||
|
||||
async fn save_conversation(
|
||||
&self,
|
||||
conversation: &ConversationSurface,
|
||||
) -> Result<(), EngineError> {
|
||||
let _ = conversation;
|
||||
Ok(())
|
||||
}
|
||||
async fn load_conversation(
|
||||
&self,
|
||||
id: ConversationId,
|
||||
) -> Result<Option<ConversationSurface>, EngineError> {
|
||||
let _ = id;
|
||||
Ok(None)
|
||||
}
|
||||
async fn list_conversations(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<ConversationSurface>, EngineError> {
|
||||
let _ = user_id;
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
// ── Memory doc operations ───────────────────────────────
|
||||
|
||||
|
||||
@@ -132,4 +132,14 @@ pub enum EventKind {
|
||||
ReflectionFailed {
|
||||
error: String,
|
||||
},
|
||||
|
||||
// ── Self-improvement ──────────────────────────────────────
|
||||
SelfImprovementStarted,
|
||||
SelfImprovementComplete {
|
||||
prompt_updated: bool,
|
||||
patterns_added: usize,
|
||||
},
|
||||
SelfImprovementFailed {
|
||||
error: String,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -223,8 +223,7 @@ impl EffectBridgeAdapter {
|
||||
})
|
||||
}
|
||||
|
||||
/// Reset the per-step call counter (called between code steps).
|
||||
#[allow(dead_code)]
|
||||
/// Reset the per-step call counter (called between threads/steps).
|
||||
pub fn reset_call_count(&self) {
|
||||
self.call_count
|
||||
.store(0, std::sync::atomic::Ordering::Relaxed);
|
||||
@@ -515,3 +514,60 @@ fn is_v1_only_tool(name: &str) -> bool {
|
||||
| "build-software"
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_adapter() -> EffectBridgeAdapter {
|
||||
use ironclaw_safety::SafetyConfig;
|
||||
let config = SafetyConfig {
|
||||
max_output_length: 10_000,
|
||||
injection_check_enabled: false,
|
||||
};
|
||||
EffectBridgeAdapter::new(
|
||||
Arc::new(ToolRegistry::new()),
|
||||
Arc::new(SafetyLayer::new(&config)),
|
||||
Arc::new(HookRegistry::default()),
|
||||
)
|
||||
}
|
||||
|
||||
/// Verify that reset_call_count resets the counter to zero,
|
||||
/// preventing the "call limit reached" error across threads.
|
||||
#[test]
|
||||
fn call_count_resets_between_threads() {
|
||||
let adapter = make_adapter();
|
||||
|
||||
// Simulate 50 tool calls (the limit)
|
||||
for _ in 0..50 {
|
||||
adapter
|
||||
.call_count
|
||||
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
assert_eq!(
|
||||
adapter
|
||||
.call_count
|
||||
.load(std::sync::atomic::Ordering::Relaxed),
|
||||
50
|
||||
);
|
||||
|
||||
// Reset — simulates what handle_with_engine does before each thread
|
||||
adapter.reset_call_count();
|
||||
assert_eq!(
|
||||
adapter
|
||||
.call_count
|
||||
.load(std::sync::atomic::Ordering::Relaxed),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
/// Verify that auto_approve_tool adds entries and is queryable.
|
||||
#[tokio::test]
|
||||
async fn auto_approve_tracks_tools() {
|
||||
let adapter = make_adapter();
|
||||
|
||||
assert!(!adapter.auto_approved.read().await.contains("shell"));
|
||||
adapter.auto_approve_tool("shell").await;
|
||||
assert!(adapter.auto_approved.read().await.contains("shell"));
|
||||
}
|
||||
}
|
||||
|
||||
+169
-52
@@ -1,5 +1,6 @@
|
||||
//! Engine v2 router — handles user messages via the engine when enabled.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use tokio::sync::RwLock;
|
||||
@@ -40,17 +41,13 @@ struct EngineState {
|
||||
thread_manager: Arc<ThreadManager>,
|
||||
conversation_manager: ConversationManager,
|
||||
effect_adapter: Arc<EffectBridgeAdapter>,
|
||||
#[allow(dead_code)]
|
||||
store: Arc<HybridStore>,
|
||||
default_project_id: ironclaw_engine::ProjectId,
|
||||
/// Currently pending approval (if any).
|
||||
pending_approval: RwLock<Option<PendingApproval>>,
|
||||
/// Per-user pending approvals (keyed by user_id).
|
||||
pending_approvals: RwLock<HashMap<String, PendingApproval>>,
|
||||
/// SSE manager for broadcasting AppEvents to the web gateway.
|
||||
sse: Option<Arc<SseManager>>,
|
||||
/// V1 database for writing conversation messages (gateway reads from here).
|
||||
db: Option<Arc<dyn Database>>,
|
||||
/// Mission manager for long-running goals.
|
||||
mission_manager: Arc<MissionManager>,
|
||||
}
|
||||
|
||||
/// Global engine state, initialized on first use.
|
||||
@@ -85,9 +82,7 @@ async fn get_or_init_engine(agent: &Agent) -> Result<(), Error> {
|
||||
));
|
||||
|
||||
let store = Arc::new(HybridStore::new(agent.workspace().cloned()));
|
||||
|
||||
// Load existing reflection docs from workspace (lessons from prior sessions)
|
||||
store.load_docs_from_workspace().await;
|
||||
store.load_state_from_workspace().await;
|
||||
|
||||
// Build capability registry from available tools
|
||||
let mut capabilities = CapabilityRegistry::new();
|
||||
@@ -114,33 +109,63 @@ async fn get_or_init_engine(agent: &Agent) -> Result<(), Error> {
|
||||
let leases = Arc::new(LeaseManager::new());
|
||||
let policy = Arc::new(PolicyEngine::new());
|
||||
|
||||
let store_dyn: Arc<dyn Store> = store.clone();
|
||||
|
||||
let thread_manager = Arc::new(ThreadManager::new(
|
||||
llm_adapter,
|
||||
effect_adapter.clone(),
|
||||
store.clone(),
|
||||
store_dyn.clone(),
|
||||
Arc::new(capabilities),
|
||||
leases,
|
||||
policy,
|
||||
));
|
||||
|
||||
// Create a default project
|
||||
let project = Project::new("default", "Default project for engine v2");
|
||||
let project_id = project.id;
|
||||
store.save_project(&project).await.map_err(|e| {
|
||||
crate::error::Error::from(crate::error::JobError::ContextError {
|
||||
id: uuid::Uuid::nil(),
|
||||
reason: format!("engine v2 store error: {e}"),
|
||||
})
|
||||
})?;
|
||||
// Reuse the persisted default project when available.
|
||||
let project_id = match store
|
||||
.list_projects()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
crate::error::Error::from(crate::error::JobError::ContextError {
|
||||
id: uuid::Uuid::nil(),
|
||||
reason: format!("engine v2 store error: {e}"),
|
||||
})
|
||||
})?
|
||||
.into_iter()
|
||||
.find(|project| project.name == "default")
|
||||
{
|
||||
Some(project) => project.id,
|
||||
None => {
|
||||
let project = Project::new("default", "Default project for engine v2");
|
||||
let project_id = project.id;
|
||||
store.save_project(&project).await.map_err(|e| {
|
||||
crate::error::Error::from(crate::error::JobError::ContextError {
|
||||
id: uuid::Uuid::nil(),
|
||||
reason: format!("engine v2 store error: {e}"),
|
||||
})
|
||||
})?;
|
||||
project_id
|
||||
}
|
||||
};
|
||||
|
||||
let conversation_manager = ConversationManager::new(Arc::clone(&thread_manager));
|
||||
let conversation_manager = ConversationManager::new(Arc::clone(&thread_manager), store.clone());
|
||||
let _ = conversation_manager
|
||||
.bootstrap_user(&agent.deps.owner_id)
|
||||
.await;
|
||||
|
||||
// Create mission manager and start cron ticker
|
||||
let mission_manager = Arc::new(MissionManager::new(
|
||||
store.clone() as Arc<dyn Store>,
|
||||
Arc::clone(&thread_manager),
|
||||
));
|
||||
let mission_manager = Arc::new(MissionManager::new(store_dyn, Arc::clone(&thread_manager)));
|
||||
let _ = thread_manager.recover_project_threads(project_id).await;
|
||||
let _ = mission_manager.bootstrap_project(project_id).await;
|
||||
mission_manager.start_cron_ticker(agent.deps.owner_id.clone());
|
||||
mission_manager.start_event_listener(agent.deps.owner_id.clone());
|
||||
|
||||
// Ensure self-improvement mission exists for this project
|
||||
if let Err(e) = mission_manager
|
||||
.ensure_self_improvement_mission(project_id)
|
||||
.await
|
||||
{
|
||||
debug!("engine v2: failed to create self-improvement mission: {e}");
|
||||
}
|
||||
|
||||
// Wire mission manager into effect adapter for mission_* function calls
|
||||
effect_adapter
|
||||
@@ -151,12 +176,10 @@ async fn get_or_init_engine(agent: &Agent) -> Result<(), Error> {
|
||||
thread_manager,
|
||||
conversation_manager,
|
||||
effect_adapter,
|
||||
store: store.clone(),
|
||||
default_project_id: project_id,
|
||||
pending_approval: RwLock::new(None),
|
||||
pending_approvals: RwLock::new(HashMap::new()),
|
||||
sse: agent.deps.sse_tx.clone(),
|
||||
db: agent.deps.store.clone(),
|
||||
mission_manager,
|
||||
});
|
||||
|
||||
Ok(())
|
||||
@@ -177,12 +200,16 @@ pub async fn handle_approval(
|
||||
let guard = lock.read().await;
|
||||
let state = guard.as_ref().expect("engine initialized");
|
||||
|
||||
// Take the pending approval
|
||||
let pending = state.pending_approval.write().await.take();
|
||||
// Take the pending approval for this user
|
||||
let pending = state
|
||||
.pending_approvals
|
||||
.write()
|
||||
.await
|
||||
.remove(&message.user_id);
|
||||
let pending = match pending {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
debug!("engine v2: no pending approval, ignoring");
|
||||
debug!(user_id = %message.user_id, "engine v2: no pending approval for user, ignoring");
|
||||
return Ok(Some("No pending approval.".into()));
|
||||
}
|
||||
};
|
||||
@@ -202,22 +229,21 @@ pub async fn handle_approval(
|
||||
)));
|
||||
}
|
||||
|
||||
// Approved — add to auto-approved set
|
||||
// Approved — only persist auto-approval when user chose "always"
|
||||
debug!(
|
||||
tool = %pending.action_name,
|
||||
always,
|
||||
"engine v2: tool approved"
|
||||
);
|
||||
|
||||
// Convert Python name back to registry name for auto-approve
|
||||
let registry_name = pending.action_name.replace('_', "-");
|
||||
state
|
||||
.effect_adapter
|
||||
.auto_approve_tool(&pending.action_name)
|
||||
.await;
|
||||
state.effect_adapter.auto_approve_tool(®istry_name).await;
|
||||
|
||||
if always {
|
||||
// Convert Python name back to registry name for auto-approve
|
||||
let registry_name = pending.action_name.replace('_', "-");
|
||||
state
|
||||
.effect_adapter
|
||||
.auto_approve_tool(&pending.action_name)
|
||||
.await;
|
||||
state.effect_adapter.auto_approve_tool(®istry_name).await;
|
||||
debug!(tool = %pending.action_name, "engine v2: tool auto-approved for session");
|
||||
}
|
||||
|
||||
@@ -263,11 +289,20 @@ pub async fn handle_with_engine(
|
||||
)
|
||||
.await;
|
||||
|
||||
// Reset the per-step call counter so each thread starts fresh
|
||||
state.effect_adapter.reset_call_count();
|
||||
|
||||
// Get or create conversation for this channel+user
|
||||
let conv_id = state
|
||||
.conversation_manager
|
||||
.get_or_create_conversation(&message.channel, &message.user_id)
|
||||
.await;
|
||||
.await
|
||||
.map_err(|e| {
|
||||
crate::error::Error::from(crate::error::JobError::ContextError {
|
||||
id: uuid::Uuid::nil(),
|
||||
reason: format!("engine v2 conversation error: {e}"),
|
||||
})
|
||||
})?;
|
||||
|
||||
// Handle the message — spawns a new thread or injects into active one
|
||||
let thread_id = state
|
||||
@@ -310,7 +345,7 @@ pub async fn handle_with_engine(
|
||||
if let Some(sse) = sse
|
||||
&& let Some(app_event) = thread_event_to_app_event(evt, &tid_str)
|
||||
{
|
||||
sse.broadcast(app_event);
|
||||
sse.broadcast_for_user(&message.user_id, app_event);
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
@@ -341,7 +376,13 @@ pub async fn handle_with_engine(
|
||||
state
|
||||
.conversation_manager
|
||||
.record_thread_outcome(conv_id, thread_id, &outcome)
|
||||
.await;
|
||||
.await
|
||||
.map_err(|e| {
|
||||
crate::error::Error::from(crate::error::JobError::ContextError {
|
||||
id: uuid::Uuid::nil(),
|
||||
reason: format!("engine v2 conversation error: {e}"),
|
||||
})
|
||||
})?;
|
||||
|
||||
// Note: trace recording, retrospective analysis, and LLM reflection
|
||||
// all run automatically inside ThreadManager after the thread completes.
|
||||
@@ -370,16 +411,19 @@ pub async fn handle_with_engine(
|
||||
}
|
||||
}
|
||||
|
||||
// Broadcast final response as AppEvent for web gateway SSE
|
||||
// Broadcast final response as AppEvent for web gateway SSE (scoped to requesting user)
|
||||
if let Some(ref sse) = state.sse
|
||||
&& let ThreadOutcome::Completed {
|
||||
response: Some(ref text),
|
||||
} = outcome
|
||||
{
|
||||
sse.broadcast(AppEvent::Response {
|
||||
content: text.clone(),
|
||||
thread_id: thread_id.to_string(),
|
||||
});
|
||||
sse.broadcast_for_user(
|
||||
&message.user_id,
|
||||
AppEvent::Response {
|
||||
content: text.clone(),
|
||||
thread_id: thread_id.to_string(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// Convert outcome to response
|
||||
@@ -398,11 +442,14 @@ pub async fn handle_with_engine(
|
||||
call_id: _,
|
||||
parameters,
|
||||
} => {
|
||||
// Store pending approval for when the user responds
|
||||
*state.pending_approval.write().await = Some(PendingApproval {
|
||||
action_name: action_name.clone(),
|
||||
original_content: content.to_string(),
|
||||
});
|
||||
// Store pending approval keyed by user so concurrent users don't collide
|
||||
state.pending_approvals.write().await.insert(
|
||||
message.user_id.clone(),
|
||||
PendingApproval {
|
||||
action_name: action_name.clone(),
|
||||
original_content: content.to_string(),
|
||||
},
|
||||
);
|
||||
|
||||
// Send approval request to channel (matches v1 ApprovalNeeded format)
|
||||
let _ = agent
|
||||
@@ -528,3 +575,73 @@ fn thread_event_to_app_event(
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Per-user approval storage: two users' approvals don't collide.
|
||||
#[tokio::test]
|
||||
async fn pending_approvals_are_per_user() {
|
||||
let approvals: RwLock<HashMap<String, PendingApproval>> = RwLock::new(HashMap::new());
|
||||
|
||||
// User A stores an approval
|
||||
approvals.write().await.insert(
|
||||
"alice".into(),
|
||||
PendingApproval {
|
||||
action_name: "shell".into(),
|
||||
original_content: "run ls".into(),
|
||||
},
|
||||
);
|
||||
|
||||
// User B stores a different approval
|
||||
approvals.write().await.insert(
|
||||
"bob".into(),
|
||||
PendingApproval {
|
||||
action_name: "web_fetch".into(),
|
||||
original_content: "fetch example.com".into(),
|
||||
},
|
||||
);
|
||||
|
||||
// Taking Alice's approval doesn't affect Bob's
|
||||
let alice_approval = approvals.write().await.remove("alice");
|
||||
assert_eq!(alice_approval.unwrap().action_name, "shell");
|
||||
|
||||
let bob_approval = approvals.write().await.remove("bob");
|
||||
assert_eq!(bob_approval.unwrap().action_name, "web_fetch");
|
||||
}
|
||||
|
||||
/// A second approval from the same user overwrites their previous one,
|
||||
/// but doesn't affect other users.
|
||||
#[tokio::test]
|
||||
async fn same_user_approval_overwrites() {
|
||||
let approvals: RwLock<HashMap<String, PendingApproval>> = RwLock::new(HashMap::new());
|
||||
|
||||
approvals.write().await.insert(
|
||||
"alice".into(),
|
||||
PendingApproval {
|
||||
action_name: "shell".into(),
|
||||
original_content: "first".into(),
|
||||
},
|
||||
);
|
||||
approvals.write().await.insert(
|
||||
"alice".into(),
|
||||
PendingApproval {
|
||||
action_name: "http".into(),
|
||||
original_content: "second".into(),
|
||||
},
|
||||
);
|
||||
|
||||
let pending = approvals.write().await.remove("alice");
|
||||
assert_eq!(pending.unwrap().action_name, "http");
|
||||
}
|
||||
|
||||
/// No pending approval for an unknown user returns None.
|
||||
#[tokio::test]
|
||||
async fn no_approval_for_unknown_user() {
|
||||
let approvals: RwLock<HashMap<String, PendingApproval>> = RwLock::new(HashMap::new());
|
||||
|
||||
let result = approvals.write().await.remove("nobody");
|
||||
assert!(result.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
+268
-84
@@ -1,40 +1,43 @@
|
||||
//! Hybrid store adapter — in-memory for ephemeral data, workspace for durable knowledge.
|
||||
//! Hybrid store adapter — workspace-backed persistence for engine state.
|
||||
//!
|
||||
//! Threads, steps, events, and leases are ephemeral (per-session).
|
||||
//! MemoryDocs (lessons, specs, playbooks from reflection) persist to the
|
||||
//! workspace so the engine learns across restarts.
|
||||
//! Reflection docs, projects, threads, steps, events, leases, and missions are
|
||||
//! cached in memory and mirrored to the workspace as JSON. This keeps the
|
||||
//! engine restart-safe without introducing dedicated DB tables yet.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::de::DeserializeOwned;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::debug;
|
||||
|
||||
use ironclaw_engine::{
|
||||
CapabilityLease, DocId, DocType, EngineError, LeaseId, MemoryDoc, Project, ProjectId, Step,
|
||||
Store, Thread, ThreadEvent, ThreadId, ThreadState,
|
||||
CapabilityLease, ConversationId, ConversationSurface, DocId, DocType, EngineError, LeaseId,
|
||||
MemoryDoc, Project, ProjectId, Step, Store, Thread, ThreadEvent, ThreadId, ThreadState,
|
||||
types::mission::{Mission, MissionId, MissionStatus},
|
||||
};
|
||||
|
||||
use crate::workspace::Workspace;
|
||||
use crate::workspace::{Workspace, WorkspaceEntry};
|
||||
|
||||
/// Workspace path prefix for engine memory docs.
|
||||
const ENGINE_DOCS_PREFIX: &str = "engine/docs";
|
||||
const PROJECTS_PREFIX: &str = "engine/state/projects";
|
||||
const CONVERSATIONS_PREFIX: &str = "engine/state/conversations";
|
||||
const THREADS_PREFIX: &str = "engine/state/threads";
|
||||
const STEPS_PREFIX: &str = "engine/state/steps";
|
||||
const EVENTS_PREFIX: &str = "engine/state/events";
|
||||
const LEASES_PREFIX: &str = "engine/state/leases";
|
||||
const MISSIONS_PREFIX: &str = "engine/state/missions";
|
||||
|
||||
/// Hybrid store: in-memory for session data, workspace for durable knowledge.
|
||||
/// Workspace-backed engine store.
|
||||
pub struct HybridStore {
|
||||
// ── Ephemeral (in-memory, per-session) ──
|
||||
threads: RwLock<HashMap<ThreadId, Thread>>,
|
||||
steps: RwLock<HashMap<ThreadId, Vec<Step>>>,
|
||||
events: RwLock<HashMap<ThreadId, Vec<ThreadEvent>>>,
|
||||
projects: RwLock<HashMap<ProjectId, Project>>,
|
||||
conversations: RwLock<HashMap<ConversationId, ConversationSurface>>,
|
||||
leases: RwLock<HashMap<LeaseId, CapabilityLease>>,
|
||||
missions: RwLock<HashMap<MissionId, Mission>>,
|
||||
|
||||
// ── Durable (workspace-backed, survives restarts) ──
|
||||
/// In-memory cache of docs (always in sync with workspace).
|
||||
docs: RwLock<HashMap<DocId, MemoryDoc>>,
|
||||
/// Workspace for persistent storage. None if workspace unavailable.
|
||||
workspace: Option<Arc<Workspace>>,
|
||||
}
|
||||
|
||||
@@ -45,6 +48,7 @@ impl HybridStore {
|
||||
steps: RwLock::new(HashMap::new()),
|
||||
events: RwLock::new(HashMap::new()),
|
||||
projects: RwLock::new(HashMap::new()),
|
||||
conversations: RwLock::new(HashMap::new()),
|
||||
leases: RwLock::new(HashMap::new()),
|
||||
missions: RwLock::new(HashMap::new()),
|
||||
docs: RwLock::new(HashMap::new()),
|
||||
@@ -52,66 +56,139 @@ impl HybridStore {
|
||||
}
|
||||
}
|
||||
|
||||
/// Load existing docs from workspace on startup.
|
||||
pub async fn load_docs_from_workspace(&self) {
|
||||
let Some(ref ws) = self.workspace else {
|
||||
/// Load persisted engine state from the workspace on startup.
|
||||
pub async fn load_state_from_workspace(&self) {
|
||||
let Some(ws) = self.workspace.as_ref() else {
|
||||
return;
|
||||
};
|
||||
|
||||
// List all engine doc files
|
||||
let entries = match ws.list(ENGINE_DOCS_PREFIX).await {
|
||||
Ok(entries) => entries,
|
||||
Err(e) => {
|
||||
debug!("no engine docs in workspace: {e}");
|
||||
return;
|
||||
self.load_docs(ws).await;
|
||||
self.load_map(ws, PROJECTS_PREFIX, |project: Project| async {
|
||||
self.projects.write().await.insert(project.id, project);
|
||||
})
|
||||
.await;
|
||||
self.load_map(
|
||||
ws,
|
||||
CONVERSATIONS_PREFIX,
|
||||
|conversation: ConversationSurface| async {
|
||||
self.conversations
|
||||
.write()
|
||||
.await
|
||||
.insert(conversation.id, conversation);
|
||||
},
|
||||
)
|
||||
.await;
|
||||
self.load_map(ws, THREADS_PREFIX, |thread: Thread| async {
|
||||
self.threads.write().await.insert(thread.id, thread);
|
||||
})
|
||||
.await;
|
||||
self.load_map(ws, STEPS_PREFIX, |steps: Vec<Step>| async {
|
||||
if let Some(thread_id) = steps.first().map(|step| step.thread_id) {
|
||||
self.steps.write().await.insert(thread_id, steps);
|
||||
}
|
||||
};
|
||||
})
|
||||
.await;
|
||||
self.load_map(ws, EVENTS_PREFIX, |events: Vec<ThreadEvent>| async {
|
||||
if let Some(thread_id) = events.first().map(|event| event.thread_id) {
|
||||
self.events.write().await.insert(thread_id, events);
|
||||
}
|
||||
})
|
||||
.await;
|
||||
self.load_map(ws, LEASES_PREFIX, |lease: CapabilityLease| async {
|
||||
self.leases.write().await.insert(lease.id, lease);
|
||||
})
|
||||
.await;
|
||||
self.load_map(ws, MISSIONS_PREFIX, |mission: Mission| async {
|
||||
self.missions.write().await.insert(mission.id, mission);
|
||||
})
|
||||
.await;
|
||||
|
||||
let mut loaded = 0;
|
||||
for entry in &entries {
|
||||
if entry.is_directory || !entry.path.ends_with(".json") {
|
||||
continue;
|
||||
}
|
||||
debug!(
|
||||
projects = self.projects.read().await.len(),
|
||||
conversations = self.conversations.read().await.len(),
|
||||
threads = self.threads.read().await.len(),
|
||||
steps = self.steps.read().await.len(),
|
||||
events = self.events.read().await.len(),
|
||||
leases = self.leases.read().await.len(),
|
||||
missions = self.missions.read().await.len(),
|
||||
docs = self.docs.read().await.len(),
|
||||
"loaded engine state from workspace"
|
||||
);
|
||||
}
|
||||
|
||||
async fn load_docs(&self, ws: &Workspace) {
|
||||
for entry in self.json_entries(ws, ENGINE_DOCS_PREFIX).await {
|
||||
match ws.read(&entry.path).await {
|
||||
Ok(ws_doc) => {
|
||||
if let Ok(doc) = serde_json::from_str::<MemoryDoc>(&ws_doc.content) {
|
||||
self.docs.write().await.insert(doc.id, doc);
|
||||
loaded += 1;
|
||||
Ok(doc) => match serde_json::from_str::<MemoryDoc>(&doc.content) {
|
||||
Ok(memory_doc) => {
|
||||
self.docs.write().await.insert(memory_doc.id, memory_doc);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
debug!(path = %entry.path, "failed to read engine doc: {e}");
|
||||
}
|
||||
Err(e) => debug!(path = %entry.path, "failed to parse engine doc: {e}"),
|
||||
},
|
||||
Err(e) => debug!(path = %entry.path, "failed to read engine doc: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
if loaded > 0 {
|
||||
debug!(loaded, "loaded engine docs from workspace");
|
||||
}
|
||||
}
|
||||
|
||||
/// Persist a MemoryDoc to workspace.
|
||||
async fn persist_doc(&self, doc: &MemoryDoc) {
|
||||
let Some(ref ws) = self.workspace else {
|
||||
async fn load_map<T, F, Fut>(&self, ws: &Workspace, directory: &str, on_value: F)
|
||||
where
|
||||
T: DeserializeOwned,
|
||||
F: Fn(T) -> Fut,
|
||||
Fut: std::future::Future<Output = ()>,
|
||||
{
|
||||
for entry in self.json_entries(ws, directory).await {
|
||||
match ws.read(&entry.path).await {
|
||||
Ok(doc) => match serde_json::from_str::<T>(&doc.content) {
|
||||
Ok(value) => on_value(value).await,
|
||||
Err(e) => debug!(path = %entry.path, "failed to parse engine state: {e}"),
|
||||
},
|
||||
Err(e) => debug!(path = %entry.path, "failed to read engine state: {e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn json_entries(&self, ws: &Workspace, directory: &str) -> Vec<WorkspaceEntry> {
|
||||
let top = match ws.list(directory).await {
|
||||
Ok(entries) => entries,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
let mut files = Vec::new();
|
||||
for entry in top {
|
||||
if entry.is_directory {
|
||||
if let Ok(children) = ws.list(&entry.path).await {
|
||||
files.extend(
|
||||
children
|
||||
.into_iter()
|
||||
.filter(|child| !child.is_directory && child.path.ends_with(".json")),
|
||||
);
|
||||
}
|
||||
} else if entry.path.ends_with(".json") {
|
||||
files.push(entry);
|
||||
}
|
||||
}
|
||||
files
|
||||
}
|
||||
|
||||
async fn persist_json<T: serde::Serialize>(&self, path: String, value: &T) {
|
||||
let Some(ws) = self.workspace.as_ref() else {
|
||||
return;
|
||||
};
|
||||
|
||||
let path = doc_workspace_path(doc);
|
||||
let json = match serde_json::to_string_pretty(doc) {
|
||||
Ok(j) => j,
|
||||
let json = match serde_json::to_string_pretty(value) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
debug!("failed to serialize doc: {e}");
|
||||
debug!(path = %path, "failed to serialize engine state: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = ws.write(&path, &json).await {
|
||||
debug!(path = %path, "failed to persist engine doc: {e}");
|
||||
debug!(path = %path, "failed to persist engine state: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Build workspace path for a MemoryDoc.
|
||||
fn doc_workspace_path(doc: &MemoryDoc) -> String {
|
||||
let type_dir = match doc.doc_type {
|
||||
DocType::Summary => "summaries",
|
||||
@@ -124,12 +201,39 @@ fn doc_workspace_path(doc: &MemoryDoc) -> String {
|
||||
format!("{ENGINE_DOCS_PREFIX}/{type_dir}/{}.json", doc.id.0)
|
||||
}
|
||||
|
||||
fn project_path(project_id: ProjectId) -> String {
|
||||
format!("{PROJECTS_PREFIX}/{}.json", project_id.0)
|
||||
}
|
||||
|
||||
fn thread_path(thread_id: ThreadId) -> String {
|
||||
format!("{THREADS_PREFIX}/{}.json", thread_id.0)
|
||||
}
|
||||
|
||||
fn conversation_path(conversation_id: ConversationId) -> String {
|
||||
format!("{CONVERSATIONS_PREFIX}/{}.json", conversation_id.0)
|
||||
}
|
||||
|
||||
fn step_path(thread_id: ThreadId) -> String {
|
||||
format!("{STEPS_PREFIX}/{}.json", thread_id.0)
|
||||
}
|
||||
|
||||
fn event_path(thread_id: ThreadId) -> String {
|
||||
format!("{EVENTS_PREFIX}/{}.json", thread_id.0)
|
||||
}
|
||||
|
||||
fn lease_path(lease_id: LeaseId) -> String {
|
||||
format!("{LEASES_PREFIX}/{}.json", lease_id.0)
|
||||
}
|
||||
|
||||
fn mission_path(mission_id: MissionId) -> String {
|
||||
format!("{MISSIONS_PREFIX}/{}.json", mission_id.0)
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Store for HybridStore {
|
||||
// ── Thread (ephemeral) ──────────────────────────────────
|
||||
|
||||
async fn save_thread(&self, thread: &Thread) -> Result<(), EngineError> {
|
||||
self.threads.write().await.insert(thread.id, thread.clone());
|
||||
self.persist_json(thread_path(thread.id), thread).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -143,7 +247,7 @@ impl Store for HybridStore {
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|t| t.project_id == project_id)
|
||||
.filter(|thread| thread.project_id == project_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
@@ -153,21 +257,38 @@ impl Store for HybridStore {
|
||||
id: ThreadId,
|
||||
state: ThreadState,
|
||||
) -> Result<(), EngineError> {
|
||||
if let Some(thread) = self.threads.write().await.get_mut(&id) {
|
||||
thread.state = state;
|
||||
let updated = {
|
||||
let mut threads = self.threads.write().await;
|
||||
if let Some(thread) = threads.get_mut(&id) {
|
||||
thread.state = state;
|
||||
Some(thread.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
if let Some(thread) = updated.as_ref() {
|
||||
self.persist_json(thread_path(id), thread).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── Step (ephemeral) ────────────────────────────────────
|
||||
|
||||
async fn save_step(&self, step: &Step) -> Result<(), EngineError> {
|
||||
self.steps
|
||||
.write()
|
||||
.await
|
||||
.entry(step.thread_id)
|
||||
.or_default()
|
||||
.push(step.clone());
|
||||
let snapshot = {
|
||||
let mut steps = self.steps.write().await;
|
||||
let thread_steps = steps.entry(step.thread_id).or_default();
|
||||
if let Some(existing) = thread_steps
|
||||
.iter_mut()
|
||||
.find(|existing| existing.id == step.id)
|
||||
{
|
||||
*existing = step.clone();
|
||||
} else {
|
||||
thread_steps.push(step.clone());
|
||||
thread_steps.sort_by_key(|saved| saved.sequence);
|
||||
}
|
||||
thread_steps.clone()
|
||||
};
|
||||
self.persist_json(step_path(step.thread_id), &snapshot)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -181,16 +302,29 @@ impl Store for HybridStore {
|
||||
.unwrap_or_default())
|
||||
}
|
||||
|
||||
// ── Event (ephemeral) ───────────────────────────────────
|
||||
|
||||
async fn append_events(&self, events: &[ThreadEvent]) -> Result<(), EngineError> {
|
||||
let mut store = self.events.write().await;
|
||||
let mut grouped: HashMap<ThreadId, Vec<ThreadEvent>> = HashMap::new();
|
||||
for event in events {
|
||||
store
|
||||
grouped
|
||||
.entry(event.thread_id)
|
||||
.or_default()
|
||||
.push(event.clone());
|
||||
}
|
||||
|
||||
for (thread_id, new_events) in grouped {
|
||||
let snapshot = {
|
||||
let mut stored = self.events.write().await;
|
||||
let thread_events = stored.entry(thread_id).or_default();
|
||||
for event in new_events {
|
||||
if !thread_events.iter().any(|existing| existing.id == event.id) {
|
||||
thread_events.push(event);
|
||||
}
|
||||
}
|
||||
thread_events.sort_by_key(|event| event.timestamp);
|
||||
thread_events.clone()
|
||||
};
|
||||
self.persist_json(event_path(thread_id), &snapshot).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -204,13 +338,12 @@ impl Store for HybridStore {
|
||||
.unwrap_or_default())
|
||||
}
|
||||
|
||||
// ── Project (ephemeral) ─────────────────────────────────
|
||||
|
||||
async fn save_project(&self, project: &Project) -> Result<(), EngineError> {
|
||||
self.projects
|
||||
.write()
|
||||
.await
|
||||
.insert(project.id, project.clone());
|
||||
self.persist_json(project_path(project.id), project).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -218,13 +351,47 @@ impl Store for HybridStore {
|
||||
Ok(self.projects.read().await.get(&id).cloned())
|
||||
}
|
||||
|
||||
// ── MemoryDoc (DURABLE — persisted to workspace) ────────
|
||||
async fn list_projects(&self) -> Result<Vec<Project>, EngineError> {
|
||||
Ok(self.projects.read().await.values().cloned().collect())
|
||||
}
|
||||
|
||||
async fn save_conversation(
|
||||
&self,
|
||||
conversation: &ConversationSurface,
|
||||
) -> Result<(), EngineError> {
|
||||
self.conversations
|
||||
.write()
|
||||
.await
|
||||
.insert(conversation.id, conversation.clone());
|
||||
self.persist_json(conversation_path(conversation.id), conversation)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn load_conversation(
|
||||
&self,
|
||||
id: ConversationId,
|
||||
) -> Result<Option<ConversationSurface>, EngineError> {
|
||||
Ok(self.conversations.read().await.get(&id).cloned())
|
||||
}
|
||||
|
||||
async fn list_conversations(
|
||||
&self,
|
||||
user_id: &str,
|
||||
) -> Result<Vec<ConversationSurface>, EngineError> {
|
||||
Ok(self
|
||||
.conversations
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|conversation| conversation.user_id == user_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn save_memory_doc(&self, doc: &MemoryDoc) -> Result<(), EngineError> {
|
||||
// Save to in-memory cache
|
||||
self.docs.write().await.insert(doc.id, doc.clone());
|
||||
// Persist to workspace
|
||||
self.persist_doc(doc).await;
|
||||
self.persist_json(doc_workspace_path(doc), doc).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -238,15 +405,14 @@ impl Store for HybridStore {
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|d| d.project_id == project_id)
|
||||
.filter(|doc| doc.project_id == project_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
// ── Lease (ephemeral) ───────────────────────────────────
|
||||
|
||||
async fn save_lease(&self, lease: &CapabilityLease) -> Result<(), EngineError> {
|
||||
self.leases.write().await.insert(lease.id, lease.clone());
|
||||
self.persist_json(lease_path(lease.id), lease).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -259,25 +425,33 @@ impl Store for HybridStore {
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|l| l.thread_id == thread_id && l.is_valid())
|
||||
.filter(|lease| lease.thread_id == thread_id && lease.is_valid())
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn revoke_lease(&self, lease_id: LeaseId, _reason: &str) -> Result<(), EngineError> {
|
||||
if let Some(lease) = self.leases.write().await.get_mut(&lease_id) {
|
||||
lease.revoked = true;
|
||||
let updated = {
|
||||
let mut leases = self.leases.write().await;
|
||||
if let Some(lease) = leases.get_mut(&lease_id) {
|
||||
lease.revoked = true;
|
||||
Some(lease.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
if let Some(lease) = updated.as_ref() {
|
||||
self.persist_json(lease_path(lease_id), lease).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── Mission (ephemeral) ──────────────────────────────────
|
||||
|
||||
async fn save_mission(&self, mission: &Mission) -> Result<(), EngineError> {
|
||||
self.missions
|
||||
.write()
|
||||
.await
|
||||
.insert(mission.id, mission.clone());
|
||||
self.persist_json(mission_path(mission.id), mission).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -291,7 +465,7 @@ impl Store for HybridStore {
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|m| m.project_id == project_id)
|
||||
.filter(|mission| mission.project_id == project_id)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
@@ -301,8 +475,18 @@ impl Store for HybridStore {
|
||||
id: MissionId,
|
||||
status: MissionStatus,
|
||||
) -> Result<(), EngineError> {
|
||||
if let Some(mission) = self.missions.write().await.get_mut(&id) {
|
||||
mission.status = status;
|
||||
let updated = {
|
||||
let mut missions = self.missions.write().await;
|
||||
if let Some(mission) = missions.get_mut(&id) {
|
||||
mission.status = status;
|
||||
mission.updated_at = chrono::Utc::now();
|
||||
Some(mission.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
if let Some(mission) = updated.as_ref() {
|
||||
self.persist_json(mission_path(id), mission).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user