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:
2026-03-25 00:15:50 -07:00
co-authored by Claude Opus 4.6
parent de71cd2ca7
commit 8180a417ff
14 changed files with 2046 additions and 182 deletions
@@ -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, &registry());
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, &registry());
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)
}
+139 -2
View File
@@ -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"));
}
}
+160
View File
@@ -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");
}
}
+162 -15
View File
@@ -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());
}
}
+731 -6
View File
@@ -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 ───────────────────────────────
+10
View File
@@ -132,4 +132,14 @@ pub enum EventKind {
ReflectionFailed {
error: String,
},
// ── Self-improvement ──────────────────────────────────────
SelfImprovementStarted,
SelfImprovementComplete {
prompt_updated: bool,
patterns_added: usize,
},
SelfImprovementFailed {
error: String,
},
}
+58 -2
View File
@@ -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
View File
@@ -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(&registry_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(&registry_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
View File
@@ -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(())
}