Add checkpoint-based engine thread recovery

This commit is contained in:
2026-03-25 10:02:45 -07:00
parent 8180a417ff
commit d7a3e04cec
10 changed files with 1310 additions and 133 deletions
@@ -24,6 +24,16 @@ use crate::types::message::ThreadMessage;
use crate::types::step::{ExecutionTier, LlmResponse, Step, StepStatus};
use crate::types::thread::{Thread, ThreadState};
const RUNTIME_CHECKPOINT_METADATA_KEY: &str = "runtime_checkpoint";
#[derive(Default)]
struct RuntimeCheckpoint {
persisted_state: serde_json::Value,
nudge_count: u32,
consecutive_errors: u32,
compaction_count: u32,
}
/// The core execution loop for a thread.
pub struct ExecutionLoop {
pub thread: Thread,
@@ -108,6 +118,67 @@ impl ExecutionLoop {
self.thread.updated_at = chrono::Utc::now();
}
fn load_runtime_checkpoint(&self) -> RuntimeCheckpoint {
let Some(checkpoint) = self
.thread
.metadata
.get(RUNTIME_CHECKPOINT_METADATA_KEY)
.and_then(|value| value.as_object())
else {
return RuntimeCheckpoint {
persisted_state: serde_json::json!({}),
..RuntimeCheckpoint::default()
};
};
RuntimeCheckpoint {
persisted_state: checkpoint
.get("persisted_state")
.cloned()
.unwrap_or_else(|| serde_json::json!({})),
nudge_count: checkpoint
.get("nudge_count")
.and_then(|value| value.as_u64())
.unwrap_or(0) as u32,
consecutive_errors: checkpoint
.get("consecutive_errors")
.and_then(|value| value.as_u64())
.unwrap_or(0) as u32,
compaction_count: checkpoint
.get("compaction_count")
.and_then(|value| value.as_u64())
.unwrap_or(0) as u32,
}
}
fn save_runtime_checkpoint(
&mut self,
persisted_state: &serde_json::Value,
nudge_count: u32,
consecutive_errors: u32,
compaction_count: u32,
) {
if let Some(metadata) = self.thread.metadata.as_object_mut() {
metadata.insert(
RUNTIME_CHECKPOINT_METADATA_KEY.into(),
serde_json::json!({
"persisted_state": persisted_state,
"nudge_count": nudge_count,
"consecutive_errors": consecutive_errors,
"compaction_count": compaction_count,
}),
);
}
self.thread.updated_at = chrono::Utc::now();
}
fn clear_runtime_checkpoint(&mut self) {
if let Some(metadata) = self.thread.metadata.as_object_mut() {
metadata.remove(RUNTIME_CHECKPOINT_METADATA_KEY);
}
self.thread.updated_at = chrono::Utc::now();
}
async fn persist_runtime_state(
&self,
step: Option<&Step>,
@@ -132,10 +203,13 @@ impl ExecutionLoop {
/// Run the execution loop to completion.
pub async fn run(&mut self) -> Result<ThreadOutcome, EngineError> {
let mut persisted_event_count = 0;
let mut persisted_event_count = self.thread.events.len();
let checkpoint = self.load_runtime_checkpoint();
// Transition to Running
self.thread.transition_to(ThreadState::Running, None)?;
// Transition to Running if this is a fresh start or restart from a resumable state.
if self.thread.state != ThreadState::Running {
self.thread.transition_to(ThreadState::Running, None)?;
}
// Inject CodeAct/RLM system prompt if none exists
if !self
@@ -173,18 +247,19 @@ impl ExecutionLoop {
// Persisted state across code steps — accumulates return values
// and tool results so the next step can access them via `state`.
let mut persisted_state = serde_json::json!({});
let mut nudge_count: u32 = 0;
let mut consecutive_errors: u32 = 0;
let mut compaction_count: u32 = 0;
let mut persisted_state = checkpoint.persisted_state;
let mut nudge_count = checkpoint.nudge_count;
let mut consecutive_errors = checkpoint.consecutive_errors;
let mut compaction_count = checkpoint.compaction_count;
for iteration in 0..max_iterations {
for iteration in self.thread.step_count..max_iterations {
// 1. Check signals
match self.check_signals() {
SignalAction::Continue => {}
SignalAction::Stop => {
self.thread
.transition_to(ThreadState::Completed, Some("stopped by signal".into()))?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(None, &mut persisted_event_count)
.await?;
return Ok(ThreadOutcome::Stopped);
@@ -208,6 +283,7 @@ impl ExecutionLoop {
);
self.thread
.transition_to(ThreadState::Completed, Some("token limit exceeded".into()))?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(None, &mut persisted_event_count)
.await?;
return Ok(ThreadOutcome::Failed {
@@ -229,6 +305,7 @@ impl ExecutionLoop {
);
self.thread
.transition_to(ThreadState::Completed, Some("timeout".into()))?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(None, &mut persisted_event_count)
.await?;
return Ok(ThreadOutcome::Failed {
@@ -248,6 +325,7 @@ impl ExecutionLoop {
);
self.thread
.transition_to(ThreadState::Completed, Some("USD budget exceeded".into()))?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(None, &mut persisted_event_count)
.await?;
return Ok(ThreadOutcome::Failed {
@@ -288,7 +366,7 @@ impl ExecutionLoop {
let active_leases = self.leases.active_for_thread(self.thread.id).await;
// 5. Build context (inject prior knowledge on first iteration only)
let retrieval_ref = if iteration == 0 {
let retrieval_ref = if self.thread.step_count == 0 {
self.retrieval.as_ref()
} else {
None
@@ -304,7 +382,7 @@ impl ExecutionLoop {
.await?;
// 6. Create step
let mut step = Step::new(self.thread.id, iteration + 1);
let mut step = Step::new(self.thread.id, self.thread.step_count + 1);
step.status = StepStatus::LlmCalling;
self.emit_event(EventKind::StepStarted { step_id: step.id });
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
@@ -315,7 +393,7 @@ impl ExecutionLoop {
// in the system prompt as Python functions. The LLM produces text with
// ```repl code blocks that the bridge detects and converts to LlmResponse::Code.
// This avoids the LLM using structured tool calls instead of writing code.
let force_text = iteration >= max_iterations.saturating_sub(1);
let force_text = self.thread.step_count >= max_iterations.saturating_sub(1);
let config = LlmCallConfig {
force_text,
depth: self.thread.config.depth,
@@ -390,6 +468,7 @@ impl ExecutionLoop {
ThreadState::Completed,
Some("FINAL() in text".into()),
)?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
.await?;
return Ok(ThreadOutcome::Completed {
@@ -419,6 +498,12 @@ impl ExecutionLoop {
tokens: step.tokens_used,
});
self.thread.step_count += 1;
self.save_runtime_checkpoint(
&persisted_state,
nudge_count,
consecutive_errors,
compaction_count,
);
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
.await?;
continue;
@@ -438,6 +523,7 @@ impl ExecutionLoop {
self.thread
.transition_to(ThreadState::Completed, Some("text response".into()))?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
.await?;
return Ok(ThreadOutcome::Completed {
@@ -522,10 +608,22 @@ impl ExecutionLoop {
ThreadState::Waiting,
Some("awaiting approval".into()),
)?;
self.save_runtime_checkpoint(
&persisted_state,
nudge_count,
consecutive_errors,
compaction_count,
);
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
.await?;
return Ok(outcome);
}
self.save_runtime_checkpoint(
&persisted_state,
nudge_count,
consecutive_errors,
compaction_count,
);
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
.await?;
}
@@ -716,6 +814,7 @@ impl ExecutionLoop {
if let Some(answer) = code_result.final_answer {
self.thread
.transition_to(ThreadState::Completed, Some("FINAL() called".into()))?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
.await?;
return Ok(ThreadOutcome::Completed {
@@ -729,6 +828,12 @@ impl ExecutionLoop {
ThreadState::Waiting,
Some("awaiting approval".into()),
)?;
self.save_runtime_checkpoint(
&persisted_state,
nudge_count,
consecutive_errors,
compaction_count,
);
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
.await?;
return Ok(outcome);
@@ -741,6 +846,12 @@ impl ExecutionLoop {
consecutive_errors = 0;
}
self.save_runtime_checkpoint(
&persisted_state,
nudge_count,
consecutive_errors,
compaction_count,
);
self.persist_runtime_state(Some(&step), &mut persisted_event_count)
.await?;
}
@@ -762,6 +873,7 @@ impl ExecutionLoop {
"consecutive error threshold: {consecutive_errors} errors"
)),
)?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(None, &mut persisted_event_count)
.await?;
return Ok(ThreadOutcome::Failed {
@@ -782,6 +894,7 @@ impl ExecutionLoop {
ThreadState::Completed,
Some("max iterations reached".into()),
)?;
self.clear_runtime_checkpoint();
self.persist_runtime_state(None, &mut persisted_event_count)
.await?;
Ok(ThreadOutcome::MaxIterations)
@@ -18,7 +18,12 @@ use crate::types::conversation::{ConversationEntry, ConversationId, Conversation
use crate::types::error::EngineError;
use crate::types::message::ThreadMessage;
use crate::types::project::ProjectId;
use crate::types::thread::{ThreadConfig, ThreadId, ThreadType};
use crate::types::thread::{ThreadConfig, ThreadId, ThreadState, ThreadType};
enum ActiveForeground {
Running(ThreadId),
Resumable(ThreadId),
}
/// Manages conversation surfaces and routes messages to threads.
///
@@ -133,8 +138,7 @@ impl ConversationManager {
let active_foreground = self.find_active_foreground(conv).await;
match active_foreground {
Some(thread_id) => {
// Inject into existing thread
Some(ActiveForeground::Running(thread_id)) => {
debug!(
conversation_id = %conversation_id,
thread_id = %thread_id,
@@ -146,6 +150,22 @@ impl ConversationManager {
self.store.save_conversation(conv).await?;
Ok(thread_id)
}
Some(ActiveForeground::Resumable(thread_id)) => {
debug!(
conversation_id = %conversation_id,
thread_id = %thread_id,
"resuming suspended foreground thread"
);
self.thread_manager
.resume_thread(thread_id, user_id, Some(ThreadMessage::user(content)), None)
.await?;
conv.add_entry(ConversationEntry::system_for_thread(
thread_id,
"Thread resumed",
));
self.store.save_conversation(conv).await?;
Ok(thread_id)
}
None => {
// Build conversation history from prior entries for context continuity
let history = build_history_from_entries(&conv.entries);
@@ -255,10 +275,16 @@ impl ConversationManager {
}
/// Find an active foreground thread in a conversation.
async fn find_active_foreground(&self, conv: &ConversationSurface) -> Option<ThreadId> {
async fn find_active_foreground(&self, conv: &ConversationSurface) -> Option<ActiveForeground> {
for &tid in &conv.active_threads {
if self.thread_manager.is_running(tid).await {
return Some(tid);
return Some(ActiveForeground::Running(tid));
}
if let Ok(Some(thread)) = self.store.load_thread(tid).await
&& thread.thread_type == ThreadType::Foreground
&& thread.state == ThreadState::Suspended
{
return Some(ActiveForeground::Resumable(tid));
}
}
None
@@ -370,32 +396,45 @@ mod tests {
struct MockStore {
conversations: RwLock<HashMap<ConversationId, ConversationSurface>>,
threads: RwLock<HashMap<ThreadId, crate::types::thread::Thread>>,
}
impl MockStore {
fn new() -> Self {
Self {
conversations: RwLock::new(HashMap::new()),
threads: RwLock::new(HashMap::new()),
}
}
}
#[async_trait::async_trait]
impl Store for MockStore {
async fn save_thread(&self, _: &crate::types::thread::Thread) -> Result<(), EngineError> {
async fn save_thread(
&self,
thread: &crate::types::thread::Thread,
) -> Result<(), EngineError> {
self.threads.write().await.insert(thread.id, thread.clone());
Ok(())
}
async fn load_thread(
&self,
_: ThreadId,
id: ThreadId,
) -> Result<Option<crate::types::thread::Thread>, EngineError> {
Ok(None)
Ok(self.threads.read().await.get(&id).cloned())
}
async fn list_threads(
&self,
_: ProjectId,
project_id: ProjectId,
) -> Result<Vec<crate::types::thread::Thread>, EngineError> {
Ok(vec![])
Ok(self
.threads
.read()
.await
.values()
.filter(|thread| thread.project_id == project_id)
.cloned()
.collect())
}
async fn update_thread_state(
&self,
@@ -563,6 +602,71 @@ mod tests {
assert!(matches!(outcome, ThreadOutcome::Completed { .. }));
}
#[tokio::test]
async fn handle_message_resumes_suspended_thread() {
let store = Arc::new(MockStore::new());
let tm = Arc::new(ThreadManager::new(
Arc::new(MockLlm(Mutex::new(vec![LlmOutput {
response: LlmResponse::Text("Recovered".into()),
usage: TokenUsage::default(),
}]))),
Arc::new(MockEffects),
store.clone(),
Arc::new(CapabilityRegistry::new()),
Arc::new(LeaseManager::new()),
Arc::new(PolicyEngine::new()),
));
let cm = ConversationManager::new(Arc::clone(&tm), store.clone());
let conv_id = cm.get_or_create_conversation("web", "user1").await.unwrap();
let project = ProjectId::new();
let mut thread = crate::types::thread::Thread::new(
"resume",
ThreadType::Foreground,
project,
ThreadConfig::default(),
);
thread.transition_to(ThreadState::Running, None).unwrap();
thread.add_message(ThreadMessage::user("earlier"));
thread.step_count = 1;
thread.metadata = serde_json::json!({
"runtime_checkpoint": {
"persisted_state": {"last_return": 7},
"nudge_count": 0,
"consecutive_errors": 0,
"compaction_count": 0
}
});
thread
.transition_to(
ThreadState::Suspended,
Some("engine restart; resumable from checkpoint".into()),
)
.unwrap();
store.save_thread(&thread).await.unwrap();
{
let mut convs = cm.conversations.write().await;
let conv = convs.get_mut(&conv_id).unwrap();
conv.track_thread(thread.id);
}
let resumed = cm
.handle_user_message(
conv_id,
"continue from there",
project,
"user1",
ThreadConfig::default(),
)
.await
.unwrap();
assert_eq!(resumed, thread.id);
let outcome = tm.join_thread(thread.id).await.unwrap();
assert!(matches!(outcome, ThreadOutcome::Completed { .. }));
}
#[tokio::test]
async fn record_outcome_adds_entry() {
let (_, cm) = make_conv_manager();
+251 -4
View File
@@ -39,7 +39,8 @@ pub struct ThreadManager {
pub policy: Arc<PolicyEngine>,
lease_planner: LeasePlanner,
tree: RwLock<ThreadTree>,
running: RwLock<HashMap<ThreadId, RunningThread>>,
running: Arc<RwLock<HashMap<ThreadId, RunningThread>>>,
completed: Arc<RwLock<HashMap<ThreadId, ThreadOutcome>>>,
/// Broadcast channel for thread events (for live status updates).
event_tx: tokio::sync::broadcast::Sender<crate::types::event::ThreadEvent>,
}
@@ -63,7 +64,8 @@ impl ThreadManager {
policy,
lease_planner: LeasePlanner::new(),
tree: RwLock::new(ThreadTree::new()),
running: RwLock::new(HashMap::new()),
running: Arc::new(RwLock::new(HashMap::new())),
completed: Arc::new(RwLock::new(HashMap::new())),
event_tx,
}
}
@@ -121,6 +123,9 @@ impl ThreadManager {
}
let thread_id = thread.id;
let user_id = user_id.into();
if let Some(metadata) = thread.metadata.as_object_mut() {
metadata.insert("user_id".into(), serde_json::Value::String(user_id.clone()));
}
// Register in tree
if let Some(pid) = parent_id {
@@ -157,6 +162,69 @@ impl ThreadManager {
// Persist
self.store.save_thread(&thread).await?;
self.start_thread(thread, user_id, false).await
}
/// Resume a persisted waiting or suspended thread.
pub async fn resume_thread(
&self,
thread_id: ThreadId,
user_id: impl Into<String>,
injected_message: Option<ThreadMessage>,
approval_event: Option<(String, bool)>,
) -> Result<(), EngineError> {
if self.is_running(thread_id).await {
return Err(EngineError::Thread(
crate::types::error::ThreadError::AlreadyRunning(thread_id),
));
}
let mut thread = self
.store
.load_thread(thread_id)
.await?
.ok_or(EngineError::ThreadNotFound(thread_id))?;
if !matches!(
thread.state,
crate::types::thread::ThreadState::Waiting
| crate::types::thread::ThreadState::Suspended
) {
return Err(EngineError::Store {
reason: format!(
"thread {thread_id} is not resumable from {:?}",
thread.state
),
});
}
if let Some((call_id, approved)) = approval_event {
let event = crate::types::event::ThreadEvent::new(
thread_id,
crate::types::event::EventKind::ApprovalReceived { call_id, approved },
);
let _ = self.event_tx.send(event.clone());
thread.events.push(event);
thread.updated_at = chrono::Utc::now();
}
if let Some(message) = injected_message {
thread.add_message(message);
}
self.store.save_thread(&thread).await?;
self.start_thread(thread, user_id.into(), true).await?;
Ok(())
}
async fn start_thread(
&self,
thread: Thread,
user_id: String,
is_resume: bool,
) -> Result<ThreadId, EngineError> {
let thread_id = thread.id;
// Create signal channel
let (tx, rx) = messaging::signal_channel(32);
@@ -180,6 +248,8 @@ impl ThreadManager {
let llm_for_reflection = Arc::clone(&self.llm);
let caps_for_reflection = Arc::clone(&self.capabilities);
let event_tx = self.event_tx.clone();
let running = Arc::clone(&self.running);
let completed = Arc::clone(&self.completed);
let handle = tokio::spawn(async move {
let mut exec = exec_loop;
let result = exec.run().await;
@@ -296,7 +366,16 @@ impl ThreadManager {
"failed to save final thread state: {e}"
);
}
result
let outcome = match result {
Ok(outcome) => outcome,
Err(error) => ThreadOutcome::Failed {
error: error.to_string(),
},
};
completed.write().await.insert(thread_id, outcome.clone());
running.write().await.remove(&thread_id);
Ok(outcome)
});
self.running.write().await.insert(
@@ -307,6 +386,10 @@ impl ThreadManager {
},
);
if is_resume {
debug!(thread_id = %thread_id, "resumed thread");
}
Ok(thread_id)
}
@@ -350,6 +433,10 @@ impl ThreadManager {
/// Wait for a thread to finish and return its outcome.
/// Removes the thread from the running set.
pub async fn join_thread(&self, thread_id: ThreadId) -> Result<ThreadOutcome, EngineError> {
if let Some(outcome) = self.completed.write().await.remove(&thread_id) {
return Ok(outcome);
}
let rt = {
let mut running = self.running.write().await;
running.remove(&thread_id)
@@ -395,6 +482,41 @@ impl ThreadManager {
finished
}
/// Automatically resume checkpointed non-foreground threads.
pub async fn resume_background_threads(
&self,
project_id: ProjectId,
) -> Result<Vec<ThreadId>, EngineError> {
let threads = self.store.list_threads(project_id).await?;
let mut resumed = Vec::new();
for thread in threads {
if thread.state != ThreadState::Suspended {
continue;
}
if thread.thread_type != ThreadType::Research {
continue;
}
if thread.metadata.get("runtime_checkpoint").is_none() {
continue;
}
let Some(user_id) = thread
.metadata
.get("user_id")
.and_then(|value| value.as_str())
.filter(|user_id| !user_id.is_empty())
else {
continue;
};
self.resume_thread(thread.id, user_id.to_string(), None, None)
.await?;
resumed.push(thread.id);
}
Ok(resumed)
}
/// Reconcile persisted non-terminal threads after process startup.
///
/// The current engine does not support mid-thread replay/resume, so any
@@ -403,6 +525,8 @@ impl ThreadManager {
&self,
project_id: ProjectId,
) -> Result<Vec<ThreadId>, EngineError> {
const PENDING_APPROVAL_METADATA_KEY: &str = "pending_approval";
const RUNTIME_CHECKPOINT_METADATA_KEY: &str = "runtime_checkpoint";
let threads = self.store.list_threads(project_id).await?;
let mut recovered = Vec::new();
@@ -411,6 +535,30 @@ impl ThreadManager {
continue;
}
if thread.state == ThreadState::Waiting
&& thread.metadata.get(PENDING_APPROVAL_METADATA_KEY).is_some()
{
continue;
}
if thread
.metadata
.get(RUNTIME_CHECKPOINT_METADATA_KEY)
.is_some()
&& matches!(thread.state, ThreadState::Running | ThreadState::Suspended)
{
if thread.state == ThreadState::Running {
thread.transition_to(
ThreadState::Suspended,
Some("engine restart; resumable from checkpoint".into()),
)?;
}
self.store.append_events(&thread.events).await?;
self.store.save_thread(&thread).await?;
recovered.push(thread.id);
continue;
}
if thread
.transition_to(
ThreadState::Failed,
@@ -744,7 +892,7 @@ mod tests {
// Give it a moment to start, then stop
tokio::time::sleep(Duration::from_millis(10)).await;
mgr.stop_thread(tid).await.unwrap();
let _ = mgr.stop_thread(tid).await;
let outcome = mgr.join_thread(tid).await.unwrap();
assert!(matches!(
@@ -820,4 +968,103 @@ mod tests {
let events = store.load_events(running.id).await.unwrap();
assert!(!events.is_empty());
}
#[tokio::test]
async fn recover_project_threads_preserves_waiting_approval_threads() {
let store = Arc::new(MockStore::new());
let project = ProjectId::new();
let mut waiting = Thread::new(
"awaiting approval",
ThreadType::Foreground,
project,
ThreadConfig::default(),
);
waiting.transition_to(ThreadState::Running, None).unwrap();
waiting
.transition_to(ThreadState::Waiting, Some("approval".into()))
.unwrap();
waiting.metadata = serde_json::json!({
"pending_approval": {
"request_id": "req-1",
"action_name": "shell",
"call_id": "call-1"
}
});
store.save_thread(&waiting).await.unwrap();
let mgr = make_manager_with_store(MockLlm::text("ignored"), Arc::clone(&store));
let recovered = mgr.recover_project_threads(project).await.unwrap();
assert!(recovered.is_empty());
let saved = store.load_thread(waiting.id).await.unwrap().unwrap();
assert_eq!(saved.state, ThreadState::Waiting);
}
#[tokio::test]
async fn recover_project_threads_suspends_checkpointed_threads() {
let store = Arc::new(MockStore::new());
let project = ProjectId::new();
let mut running = Thread::new(
"resume me",
ThreadType::Foreground,
project,
ThreadConfig::default(),
);
running.transition_to(ThreadState::Running, None).unwrap();
running.metadata = serde_json::json!({
"runtime_checkpoint": {
"persisted_state": {"last_return": 7},
"nudge_count": 0,
"consecutive_errors": 0,
"compaction_count": 0
}
});
store.save_thread(&running).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::Suspended);
}
#[tokio::test]
async fn resume_background_threads_restarts_suspended_research_threads() {
let store = Arc::new(MockStore::new());
let project = ProjectId::new();
let mut research = Thread::new(
"background research",
ThreadType::Research,
project,
ThreadConfig::default(),
);
research.transition_to(ThreadState::Running, None).unwrap();
research.metadata = serde_json::json!({
"user_id": "owner",
"runtime_checkpoint": {
"persisted_state": {},
"nudge_count": 0,
"consecutive_errors": 0,
"compaction_count": 0
}
});
research
.transition_to(
ThreadState::Suspended,
Some("engine restart; resumable from checkpoint".into()),
)
.unwrap();
store.save_thread(&research).await.unwrap();
let mgr = make_manager_with_store(MockLlm::text("done"), Arc::clone(&store));
let resumed = mgr.resume_background_threads(project).await.unwrap();
assert_eq!(resumed, vec![research.id]);
let outcome = mgr.join_thread(research.id).await.unwrap();
assert!(matches!(outcome, ThreadOutcome::Completed { .. }));
}
}
+59 -20
View File
@@ -163,29 +163,47 @@ impl MissionManager {
self.store.save_mission(&updated).await?;
debug!(mission_id = %id, thread_id = %thread_id, "mission fired");
// Wait for thread completion and process the outcome
let tm = Arc::clone(&self.thread_manager);
let store = Arc::clone(&self.store);
let mission_id = id;
tokio::spawn(async move {
match tm.join_thread(thread_id).await {
Ok(outcome) => {
if let Err(e) =
process_mission_outcome(&store, mission_id, thread_id, &outcome).await
{
warn!(mission_id = %mission_id, "failed to process outcome: {e}");
}
}
Err(e) => {
warn!(mission_id = %mission_id, "thread join failed: {e}");
}
}
});
self.spawn_mission_outcome_watcher(id, thread_id);
Ok(Some(thread_id))
}
/// Resume suspended checkpointed mission threads after restart.
pub async fn resume_recoverable_threads(
&self,
user_id: &str,
) -> Result<Vec<ThreadId>, EngineError> {
let mut resumed = Vec::new();
for mission_id in self.active.read().await.clone() {
let Some(mission) = self.store.load_mission(mission_id).await? else {
continue;
};
for &thread_id in mission.thread_history.iter().rev() {
let Some(thread) = self.store.load_thread(thread_id).await? else {
continue;
};
if thread.thread_type != ThreadType::Mission
|| thread.state != crate::types::thread::ThreadState::Suspended
{
continue;
}
if thread.metadata.get("runtime_checkpoint").is_none() {
continue;
}
self.thread_manager
.resume_thread(thread_id, user_id.to_string(), None, None)
.await?;
self.spawn_mission_outcome_watcher(mission_id, thread_id);
resumed.push(thread_id);
}
}
Ok(resumed)
}
/// Start a background cron ticker that fires due missions every 60 seconds.
pub fn start_cron_ticker(self: &Arc<Self>, user_id: String) {
let mgr = Arc::clone(self);
@@ -462,6 +480,25 @@ impl MissionManager {
Ok(spawned)
}
fn spawn_mission_outcome_watcher(&self, mission_id: MissionId, thread_id: ThreadId) {
let tm = Arc::clone(&self.thread_manager);
let store = Arc::clone(&self.store);
tokio::spawn(async move {
match tm.join_thread(thread_id).await {
Ok(outcome) => {
if let Err(e) =
process_mission_outcome(&store, mission_id, thread_id, &outcome).await
{
warn!(mission_id = %mission_id, "failed to process outcome: {e}");
}
}
Err(e) => {
warn!(mission_id = %mission_id, "thread join failed: {e}");
}
}
});
}
}
// ── Meta-prompt generation ───────────────────────────────────
@@ -647,7 +684,9 @@ async fn process_self_improvement_output(
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)");
debug!(
"self-improvement: no structured JSON in response (agent likely used tools directly)"
);
return Ok(());
}
};
+3 -1
View File
@@ -9,4 +9,6 @@ mod llm_adapter;
mod router;
mod store_adapter;
pub use router::{handle_approval, handle_with_engine, is_engine_v2_enabled};
pub use router::{
handle_approval, handle_with_engine, is_engine_v2_enabled, pending_approval_for_user_thread,
};
+712 -73
View File
@@ -30,10 +30,23 @@ pub fn is_engine_v2_enabled() -> bool {
}
/// Pending approval info stored between the NeedApproval outcome and the user's response.
#[derive(Clone)]
struct PendingApproval {
request_id: String,
action_name: String,
/// The user message that triggered this (for re-submission after approval).
original_content: String,
thread_id: ironclaw_engine::ThreadId,
conversation_id: ironclaw_engine::ConversationId,
call_id: String,
description: String,
parameters: serde_json::Value,
}
#[derive(Debug, Clone)]
pub struct PendingApprovalView {
pub request_id: String,
pub tool_name: String,
pub description: String,
pub parameters: String,
}
/// Persistent engine state that lives across messages.
@@ -41,6 +54,7 @@ struct EngineState {
thread_manager: Arc<ThreadManager>,
conversation_manager: ConversationManager,
effect_adapter: Arc<EffectBridgeAdapter>,
store: Arc<dyn Store>,
default_project_id: ironclaw_engine::ProjectId,
/// Per-user pending approvals (keyed by user_id).
pending_approvals: RwLock<HashMap<String, PendingApproval>>,
@@ -53,6 +67,14 @@ struct EngineState {
/// Global engine state, initialized on first use.
static ENGINE_STATE: OnceLock<RwLock<Option<EngineState>>> = OnceLock::new();
const PENDING_APPROVAL_METADATA_KEY: &str = "pending_approval";
enum PendingApprovalResolution {
None,
Resolved(PendingApproval),
Ambiguous,
}
/// Get or initialize the engine state using the agent's dependencies.
async fn get_or_init_engine(agent: &Agent) -> Result<(), Error> {
let lock = ENGINE_STATE.get_or_init(|| RwLock::new(None));
@@ -156,6 +178,10 @@ async fn get_or_init_engine(agent: &Agent) -> Result<(), Error> {
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;
let _ = mission_manager
.resume_recoverable_threads(&agent.deps.owner_id)
.await;
let _ = thread_manager.resume_background_threads(project_id).await;
mission_manager.start_cron_ticker(agent.deps.owner_id.clone());
mission_manager.start_event_listener(agent.deps.owner_id.clone());
@@ -176,6 +202,7 @@ 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_approvals: RwLock::new(HashMap::new()),
sse: agent.deps.sse_tx.clone(),
@@ -185,6 +212,252 @@ async fn get_or_init_engine(agent: &Agent) -> Result<(), Error> {
Ok(())
}
async fn persist_pending_approval(
store: &Arc<dyn Store>,
pending: &PendingApproval,
) -> Result<(), Error> {
let mut thread = store
.load_thread(pending.thread_id)
.await
.map_err(|e| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 store error: {e}"),
})
})?
.ok_or_else(|| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 thread {} not found", pending.thread_id),
})
})?;
let metadata = thread.metadata.as_object_mut().ok_or_else(|| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: "engine v2 thread metadata must be an object".into(),
})
})?;
metadata.insert(
PENDING_APPROVAL_METADATA_KEY.into(),
serde_json::json!({
"request_id": pending.request_id,
"action_name": pending.action_name,
"thread_id": pending.thread_id.to_string(),
"conversation_id": pending.conversation_id.to_string(),
"call_id": pending.call_id,
"description": pending.description,
"parameters": pending.parameters,
}),
);
thread.updated_at = chrono::Utc::now();
store.save_thread(&thread).await.map_err(|e| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 store error: {e}"),
})
})
}
async fn load_pending_approval_from_thread(
store: &Arc<dyn Store>,
conversation_id: ironclaw_engine::ConversationId,
thread_id: ironclaw_engine::ThreadId,
) -> Result<Option<PendingApproval>, Error> {
let Some(thread) = store.load_thread(thread_id).await.map_err(|e| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 store error: {e}"),
})
})?
else {
return Ok(None);
};
if thread.state != ironclaw_engine::ThreadState::Waiting {
return Ok(None);
}
let Some(pending) = thread
.metadata
.get(PENDING_APPROVAL_METADATA_KEY)
.and_then(|value| value.as_object())
else {
return Ok(None);
};
let Some(request_id) = pending.get("request_id").and_then(|value| value.as_str()) else {
return Ok(None);
};
let Some(action_name) = pending.get("action_name").and_then(|value| value.as_str()) else {
return Ok(None);
};
let Some(call_id) = pending.get("call_id").and_then(|value| value.as_str()) else {
return Ok(None);
};
let description = pending
.get("description")
.and_then(|value| value.as_str())
.map(str::to_string)
.unwrap_or_else(|| format!("Tool '{}' requires approval to execute.", action_name));
let parameters = pending
.get("parameters")
.cloned()
.unwrap_or(serde_json::Value::Null);
Ok(Some(PendingApproval {
request_id: request_id.to_string(),
action_name: action_name.to_string(),
thread_id,
conversation_id,
call_id: call_id.to_string(),
description,
parameters,
}))
}
async fn clear_pending_approval_metadata(
store: &Arc<dyn Store>,
thread_id: ironclaw_engine::ThreadId,
) -> Result<(), Error> {
let Some(mut thread) = store.load_thread(thread_id).await.map_err(|e| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 store error: {e}"),
})
})?
else {
return Ok(());
};
if let Some(metadata) = thread.metadata.as_object_mut() {
metadata.remove(PENDING_APPROVAL_METADATA_KEY);
thread.updated_at = chrono::Utc::now();
store.save_thread(&thread).await.map_err(|e| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 store error: {e}"),
})
})?;
}
Ok(())
}
async fn resolve_pending_approval_for_thread(
store: &Arc<dyn Store>,
pending_approvals: &RwLock<HashMap<String, PendingApproval>>,
user_id: &str,
thread_id_hint: Option<&str>,
) -> Result<PendingApprovalResolution, Error> {
let hinted_thread_id = thread_id_hint.and_then(|id| uuid::Uuid::parse_str(id).ok());
if let Some(cached) = pending_approvals.read().await.get(user_id).cloned() {
let hint_matches = hinted_thread_id
.map(|id| cached.thread_id.0 == id)
.unwrap_or(true);
if hint_matches {
if let Some(pending) =
load_pending_approval_from_thread(store, cached.conversation_id, cached.thread_id)
.await?
{
return Ok(PendingApprovalResolution::Resolved(pending));
}
let mut approvals = pending_approvals.write().await;
if approvals
.get(user_id)
.is_some_and(|pending| pending.thread_id == cached.thread_id)
{
approvals.remove(user_id);
}
}
}
let conversations = store.list_conversations(user_id).await.map_err(|e| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 store error: {e}"),
})
})?;
let mut candidates = Vec::new();
for conversation in conversations {
for thread_id in conversation.active_threads {
if hinted_thread_id.is_some_and(|hint| thread_id.0 != hint) {
continue;
}
let Some(thread) = store.load_thread(thread_id).await.map_err(|e| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 store error: {e}"),
})
})?
else {
continue;
};
let Some(pending) =
load_pending_approval_from_thread(store, conversation.id, thread_id).await?
else {
continue;
};
candidates.push((thread.updated_at, pending));
}
}
if hinted_thread_id.is_none() && candidates.len() > 1 {
return Ok(PendingApprovalResolution::Ambiguous);
}
candidates.sort_by_key(|(updated_at, _)| *updated_at);
let resolved = candidates.pop().map(|(_, pending)| pending);
if let Some(ref pending) = resolved {
pending_approvals
.write()
.await
.insert(user_id.to_string(), pending.clone());
}
Ok(match resolved {
Some(pending) => PendingApprovalResolution::Resolved(pending),
None => PendingApprovalResolution::None,
})
}
pub async fn pending_approval_for_user_thread(
user_id: &str,
thread_id: Option<&str>,
) -> Result<Option<PendingApprovalView>, Error> {
let Some(lock) = ENGINE_STATE.get() else {
return Ok(None);
};
let guard = lock.read().await;
let Some(state) = guard.as_ref() else {
return Ok(None);
};
match resolve_pending_approval_for_thread(
&state.store,
&state.pending_approvals,
user_id,
thread_id,
)
.await?
{
PendingApprovalResolution::Resolved(pending) => Ok(Some(PendingApprovalView {
request_id: pending.request_id,
tool_name: pending.action_name,
description: pending.description,
parameters: serde_json::to_string_pretty(&pending.parameters)
.unwrap_or_else(|_| pending.parameters.to_string()),
})),
PendingApprovalResolution::None | PendingApprovalResolution::Ambiguous => Ok(None),
}
}
/// Handle an approval response (yes/no/always) for engine v2.
///
/// Called from `handle_message` when the user responds to an approval request.
@@ -200,17 +473,23 @@ pub async fn handle_approval(
let guard = lock.read().await;
let state = guard.as_ref().expect("engine initialized");
// 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 => {
let pending = match resolve_pending_approval_for_thread(
&state.store,
&state.pending_approvals,
&message.user_id,
message.thread_id.as_deref(),
)
.await?
{
PendingApprovalResolution::Resolved(p) => p,
PendingApprovalResolution::None => {
debug!(user_id = %message.user_id, "engine v2: no pending approval for user, ignoring");
return Ok(Some("No pending approval.".into()));
return Ok(Some("No pending approval for this thread.".into()));
}
PendingApprovalResolution::Ambiguous => {
return Ok(Some(
"Multiple pending approvals are waiting. Approve from the original thread or retry with that thread selected.".into(),
));
}
};
@@ -223,41 +502,83 @@ pub async fn handle_approval(
&message.metadata,
)
.await;
return Ok(Some(format!(
"Denied: tool '{}' was not executed.",
pending.action_name
)));
}
// Approved — only persist auto-approval when user chose "always"
// Approved — persist auto-approval when user chose "always"
debug!(
tool = %pending.action_name,
always,
"engine v2: tool approved"
approved,
"engine v2: tool approval received"
);
if always {
// Convert Python name back to registry name for auto-approve
if approved && always {
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");
debug!(
tool = %pending.action_name,
"engine v2: tool auto-approved for session"
);
}
// Re-process the original message — the tool will now pass approval
let _ = agent
.channels
.send_status(
&message.channel,
StatusUpdate::Thinking("Re-executing with approval...".into()),
StatusUpdate::Thinking("Resuming pending thread...".into()),
&message.metadata,
)
.await;
handle_with_engine(agent, message, &pending.original_content).await
let resume_message = if approved {
ironclaw_engine::ThreadMessage::user(format!(
"User approved action '{}'. Continue from the pending step and reuse the approved action if still needed.",
pending.action_name
))
} else {
ironclaw_engine::ThreadMessage::user(format!(
"User denied action '{}'. Do not execute it; choose an alternative approach.",
pending.action_name
))
};
state.effect_adapter.reset_call_count();
state
.thread_manager
.resume_thread(
pending.thread_id,
message.user_id.clone(),
Some(resume_message),
Some((pending.call_id.clone(), approved)),
)
.await
.map_err(|e| {
crate::error::Error::from(crate::error::JobError::ContextError {
id: uuid::Uuid::nil(),
reason: format!("engine v2 resume error: {e}"),
})
})?;
clear_pending_approval_metadata(&state.store, pending.thread_id).await?;
let mut approvals = state.pending_approvals.write().await;
if approvals
.get(&message.user_id)
.is_some_and(|cached| cached.thread_id == pending.thread_id)
{
approvals.remove(&message.user_id);
}
await_thread_outcome(
agent,
state,
message,
pending.conversation_id,
pending.thread_id,
)
.await
}
/// Handle a user message through the engine v2 pipeline.
@@ -325,9 +646,27 @@ pub async fn handle_with_engine(
})
})?;
debug!(thread_id = %thread_id, "engine v2: thread spawned");
if let Some(ref db) = state.db
&& let Ok(conv_id_v1) = db
.get_or_create_assistant_conversation(&message.user_id, &message.channel)
.await
{
let _ = db
.add_conversation_message(conv_id_v1, "user", content)
.await;
}
// Subscribe to live events for progress updates
debug!(thread_id = %thread_id, "engine v2: thread spawned");
await_thread_outcome(agent, state, message, conv_id, thread_id).await
}
async fn await_thread_outcome(
agent: &Agent,
state: &EngineState,
message: &IncomingMessage,
conv_id: ironclaw_engine::ConversationId,
thread_id: ironclaw_engine::ThreadId,
) -> Result<Option<String>, Error> {
let mut event_rx = state.thread_manager.subscribe_events();
let channels = &agent.channels;
let channel_name = &message.channel;
@@ -335,7 +674,6 @@ pub async fn handle_with_engine(
let sse = state.sse.as_ref();
let tid_str = thread_id.to_string();
// Forward events to both the channel (REPL) and SSE (web gateway)
loop {
tokio::select! {
event = event_rx.recv() => {
@@ -360,7 +698,6 @@ pub async fn handle_with_engine(
}
}
// Join the thread to get the outcome
let outcome = state
.thread_manager
.join_thread(thread_id)
@@ -372,7 +709,6 @@ pub async fn handle_with_engine(
})
})?;
// Record outcome in conversation
state
.conversation_manager
.record_thread_outcome(conv_id, thread_id, &outcome)
@@ -384,34 +720,19 @@ pub async fn handle_with_engine(
})
})?;
// Note: trace recording, retrospective analysis, and LLM reflection
// all run automatically inside ThreadManager after the thread completes.
// Persist to v1 conversation DB so web gateway can display messages
if let Some(ref db) = state.db {
// get_or_create_assistant_conversation gives us a per-user, per-channel conversation
if let Ok(conv_id_v1) = db
if let Some(ref db) = state.db
&& let Ok(conv_id_v1) = db
.get_or_create_assistant_conversation(&message.user_id, &message.channel)
.await
{
// Write user message
let _ = db
.add_conversation_message(conv_id_v1, "user", content)
.await;
// Write agent response
if let ThreadOutcome::Completed {
response: Some(ref text),
} = outcome
{
let _ = db
.add_conversation_message(conv_id_v1, "assistant", text)
.await;
}
}
&& let ThreadOutcome::Completed {
response: Some(ref text),
} = outcome
{
let _ = db
.add_conversation_message(conv_id_v1, "assistant", text)
.await;
}
// 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),
@@ -426,7 +747,6 @@ pub async fn handle_with_engine(
);
}
// Convert outcome to response
match outcome {
ThreadOutcome::Completed { response } => {
debug!(thread_id = %thread_id, "engine v2: completed");
@@ -439,17 +759,26 @@ pub async fn handle_with_engine(
ThreadOutcome::Failed { error } => Ok(Some(format!("Error: {error}"))),
ThreadOutcome::NeedApproval {
action_name,
call_id: _,
call_id,
parameters,
} => {
// 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(),
},
);
let request_id = uuid::Uuid::new_v4().to_string();
let description = format!("Tool '{}' requires approval to execute.", action_name);
let pending = PendingApproval {
request_id: request_id.clone(),
action_name: action_name.clone(),
thread_id,
conversation_id: conv_id,
call_id,
description: description.clone(),
parameters: parameters.clone(),
};
state
.pending_approvals
.write()
.await
.insert(message.user_id.clone(), pending.clone());
persist_pending_approval(&state.store, &pending).await?;
// Send approval request to channel (matches v1 ApprovalNeeded format)
let _ = agent
@@ -457,12 +786,9 @@ pub async fn handle_with_engine(
.send_status(
&message.channel,
StatusUpdate::ApprovalNeeded {
request_id: uuid::Uuid::new_v4().to_string(),
request_id,
tool_name: action_name.clone(),
description: format!(
"Tool '{}' requires approval to execute.",
action_name
),
description,
parameters,
allow_always: true,
},
@@ -471,7 +797,7 @@ pub async fn handle_with_engine(
.await;
Ok(Some(format!(
"Tool '{}' requires approval. Reply 'yes' to approve, 'always' to auto-approve, or 'no' to deny.",
"Tool '{}' requires approval. Reply 'yes' to approve, 'always' to auto-approve future uses of this tool, or 'no' to deny.",
action_name
)))
}
@@ -579,6 +905,190 @@ fn thread_event_to_app_event(
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::RwLock as TokioRwLock;
struct TestStore {
conversations: TokioRwLock<Vec<ironclaw_engine::ConversationSurface>>,
threads: TokioRwLock<HashMap<ironclaw_engine::ThreadId, ironclaw_engine::Thread>>,
}
impl TestStore {
fn new() -> Self {
Self {
conversations: TokioRwLock::new(Vec::new()),
threads: TokioRwLock::new(HashMap::new()),
}
}
}
#[async_trait::async_trait]
impl Store for TestStore {
async fn save_thread(
&self,
thread: &ironclaw_engine::Thread,
) -> Result<(), ironclaw_engine::EngineError> {
self.threads.write().await.insert(thread.id, thread.clone());
Ok(())
}
async fn load_thread(
&self,
id: ironclaw_engine::ThreadId,
) -> Result<Option<ironclaw_engine::Thread>, ironclaw_engine::EngineError> {
Ok(self.threads.read().await.get(&id).cloned())
}
async fn list_threads(
&self,
_project_id: ironclaw_engine::ProjectId,
) -> Result<Vec<ironclaw_engine::Thread>, ironclaw_engine::EngineError> {
Ok(self.threads.read().await.values().cloned().collect())
}
async fn update_thread_state(
&self,
_id: ironclaw_engine::ThreadId,
_state: ironclaw_engine::ThreadState,
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
async fn save_step(
&self,
_: &ironclaw_engine::Step,
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
async fn load_steps(
&self,
_: ironclaw_engine::ThreadId,
) -> Result<Vec<ironclaw_engine::Step>, ironclaw_engine::EngineError> {
Ok(vec![])
}
async fn append_events(
&self,
_: &[ironclaw_engine::ThreadEvent],
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
async fn load_events(
&self,
_: ironclaw_engine::ThreadId,
) -> Result<Vec<ironclaw_engine::ThreadEvent>, ironclaw_engine::EngineError> {
Ok(vec![])
}
async fn save_project(
&self,
_: &ironclaw_engine::Project,
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
async fn load_project(
&self,
_: ironclaw_engine::ProjectId,
) -> Result<Option<ironclaw_engine::Project>, ironclaw_engine::EngineError> {
Ok(None)
}
async fn list_projects(
&self,
) -> Result<Vec<ironclaw_engine::Project>, ironclaw_engine::EngineError> {
Ok(vec![])
}
async fn save_conversation(
&self,
conversation: &ironclaw_engine::ConversationSurface,
) -> Result<(), ironclaw_engine::EngineError> {
let mut conversations = self.conversations.write().await;
conversations.retain(|existing| existing.id != conversation.id);
conversations.push(conversation.clone());
Ok(())
}
async fn load_conversation(
&self,
id: ironclaw_engine::ConversationId,
) -> Result<Option<ironclaw_engine::ConversationSurface>, ironclaw_engine::EngineError>
{
Ok(self
.conversations
.read()
.await
.iter()
.find(|conversation| conversation.id == id)
.cloned())
}
async fn list_conversations(
&self,
user_id: &str,
) -> Result<Vec<ironclaw_engine::ConversationSurface>, ironclaw_engine::EngineError>
{
Ok(self
.conversations
.read()
.await
.iter()
.filter(|conversation| conversation.user_id == user_id)
.cloned()
.collect())
}
async fn save_memory_doc(
&self,
_: &ironclaw_engine::MemoryDoc,
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
async fn load_memory_doc(
&self,
_: ironclaw_engine::DocId,
) -> Result<Option<ironclaw_engine::MemoryDoc>, ironclaw_engine::EngineError> {
Ok(None)
}
async fn list_memory_docs(
&self,
_: ironclaw_engine::ProjectId,
) -> Result<Vec<ironclaw_engine::MemoryDoc>, ironclaw_engine::EngineError> {
Ok(vec![])
}
async fn save_lease(
&self,
_: &ironclaw_engine::CapabilityLease,
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
async fn load_active_leases(
&self,
_: ironclaw_engine::ThreadId,
) -> Result<Vec<ironclaw_engine::CapabilityLease>, ironclaw_engine::EngineError> {
Ok(vec![])
}
async fn revoke_lease(
&self,
_: ironclaw_engine::LeaseId,
_: &str,
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
async fn save_mission(
&self,
_: &ironclaw_engine::Mission,
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
async fn load_mission(
&self,
_: ironclaw_engine::MissionId,
) -> Result<Option<ironclaw_engine::Mission>, ironclaw_engine::EngineError> {
Ok(None)
}
async fn list_missions(
&self,
_: ironclaw_engine::ProjectId,
) -> Result<Vec<ironclaw_engine::Mission>, ironclaw_engine::EngineError> {
Ok(vec![])
}
async fn update_mission_status(
&self,
_: ironclaw_engine::MissionId,
_: ironclaw_engine::MissionStatus,
) -> Result<(), ironclaw_engine::EngineError> {
Ok(())
}
}
/// Per-user approval storage: two users' approvals don't collide.
#[tokio::test]
@@ -589,8 +1099,13 @@ mod tests {
approvals.write().await.insert(
"alice".into(),
PendingApproval {
request_id: "req-a".into(),
action_name: "shell".into(),
original_content: "run ls".into(),
thread_id: ironclaw_engine::ThreadId::new(),
conversation_id: ironclaw_engine::ConversationId::new(),
call_id: "call-a".into(),
description: "desc".into(),
parameters: serde_json::json!({}),
},
);
@@ -598,8 +1113,13 @@ mod tests {
approvals.write().await.insert(
"bob".into(),
PendingApproval {
request_id: "req-b".into(),
action_name: "web_fetch".into(),
original_content: "fetch example.com".into(),
thread_id: ironclaw_engine::ThreadId::new(),
conversation_id: ironclaw_engine::ConversationId::new(),
call_id: "call-b".into(),
description: "desc".into(),
parameters: serde_json::json!({}),
},
);
@@ -620,15 +1140,25 @@ mod tests {
approvals.write().await.insert(
"alice".into(),
PendingApproval {
request_id: "req-1".into(),
action_name: "shell".into(),
original_content: "first".into(),
thread_id: ironclaw_engine::ThreadId::new(),
conversation_id: ironclaw_engine::ConversationId::new(),
call_id: "call-1".into(),
description: "desc".into(),
parameters: serde_json::json!({}),
},
);
approvals.write().await.insert(
"alice".into(),
PendingApproval {
request_id: "req-2".into(),
action_name: "http".into(),
original_content: "second".into(),
thread_id: ironclaw_engine::ThreadId::new(),
conversation_id: ironclaw_engine::ConversationId::new(),
call_id: "call-2".into(),
description: "desc".into(),
parameters: serde_json::json!({}),
},
);
@@ -644,4 +1174,113 @@ mod tests {
let result = approvals.write().await.remove("nobody");
assert!(result.is_none());
}
#[tokio::test]
async fn persist_and_resolve_pending_approval_from_thread_metadata() {
let store: Arc<dyn Store> = Arc::new(TestStore::new());
let thread_id = ironclaw_engine::ThreadId::new();
let conversation_id = ironclaw_engine::ConversationId::new();
let pending_approvals = RwLock::new(HashMap::new());
let mut thread = ironclaw_engine::Thread::new(
"goal",
ironclaw_engine::ThreadType::Foreground,
ironclaw_engine::ProjectId::new(),
ironclaw_engine::ThreadConfig::default(),
);
thread.id = thread_id;
thread
.transition_to(ironclaw_engine::ThreadState::Running, None)
.unwrap();
thread
.transition_to(
ironclaw_engine::ThreadState::Waiting,
Some("approval".into()),
)
.unwrap();
store.save_thread(&thread).await.unwrap();
let mut conversation = ironclaw_engine::ConversationSurface::new("web", "user1");
conversation.id = conversation_id;
conversation.track_thread(thread_id);
store.save_conversation(&conversation).await.unwrap();
let pending = PendingApproval {
request_id: "req-123".into(),
action_name: "shell".into(),
thread_id,
conversation_id,
call_id: "call-123".into(),
description: "Tool 'shell' requires approval to execute.".into(),
parameters: serde_json::json!({"cmd": "ls"}),
};
persist_pending_approval(&store, &pending).await.unwrap();
let resolved =
resolve_pending_approval_for_thread(&store, &pending_approvals, "user1", None)
.await
.unwrap();
let PendingApprovalResolution::Resolved(resolved) = resolved else {
panic!("expected resolved pending approval");
};
assert_eq!(resolved.action_name, "shell");
assert_eq!(resolved.thread_id, thread_id);
assert_eq!(resolved.request_id, "req-123");
assert_eq!(resolved.parameters["cmd"], "ls");
clear_pending_approval_metadata(&store, thread_id)
.await
.unwrap();
let thread = store.load_thread(thread_id).await.unwrap().unwrap();
assert!(thread.metadata.get(PENDING_APPROVAL_METADATA_KEY).is_none());
}
#[tokio::test]
async fn resolve_pending_approval_detects_ambiguity_without_thread_hint() {
let store: Arc<dyn Store> = Arc::new(TestStore::new());
let pending_approvals = RwLock::new(HashMap::new());
for call_id in ["call-1", "call-2"] {
let thread_id = ironclaw_engine::ThreadId::new();
let mut thread = ironclaw_engine::Thread::new(
"goal",
ironclaw_engine::ThreadType::Foreground,
ironclaw_engine::ProjectId::new(),
ironclaw_engine::ThreadConfig::default(),
);
thread.id = thread_id;
thread
.transition_to(ironclaw_engine::ThreadState::Running, None)
.unwrap();
thread
.transition_to(
ironclaw_engine::ThreadState::Waiting,
Some("approval".into()),
)
.unwrap();
store.save_thread(&thread).await.unwrap();
let mut conversation = ironclaw_engine::ConversationSurface::new("web", "user1");
conversation.track_thread(thread_id);
let conversation_id = conversation.id;
store.save_conversation(&conversation).await.unwrap();
let pending = PendingApproval {
request_id: format!("req-{call_id}"),
action_name: "shell".into(),
thread_id,
conversation_id,
call_id: call_id.into(),
description: "Tool 'shell' requires approval to execute.".into(),
parameters: serde_json::json!({}),
};
persist_pending_approval(&store, &pending).await.unwrap();
}
let resolved =
resolve_pending_approval_for_thread(&store, &pending_approvals, "user1", None)
.await
.unwrap();
assert!(matches!(resolved, PendingApprovalResolution::Ambiguous));
}
}
+17 -8
View File
@@ -103,15 +103,24 @@ impl HybridStore {
})
.await;
let projects = self.projects.read().await.len();
let conversations = self.conversations.read().await.len();
let threads = self.threads.read().await.len();
let steps = self.steps.read().await.len();
let events = self.events.read().await.len();
let leases = self.leases.read().await.len();
let missions = self.missions.read().await.len();
let docs = self.docs.read().await.len();
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(),
projects,
conversations,
threads,
steps,
events,
leases,
missions,
docs,
"loaded engine state from workspace"
);
}
+1 -1
View File
@@ -201,7 +201,7 @@ All responses include:
## Pending Approvals
Tool approval state is **in-memory only** (not persisted to DB). Server restart clears all pending approvals. The `pending_approval` field in `HistoryResponse` is re-populated on thread switch from in-memory state.
Classic agent approvals are in-memory, but engine v2 approvals are mirrored into durable engine metadata. `HistoryResponse.pending_approval` should rehydrate from live session state first and then from engine metadata so approval cards survive thread switches and restart recovery.
## Adding a New API Endpoint
+27 -3
View File
@@ -290,6 +290,23 @@ pub struct HistoryQuery {
pub before: Option<String>,
}
async fn engine_pending_approval(user_id: &str, thread_id: Uuid) -> Option<PendingApprovalInfo> {
if !crate::bridge::is_engine_v2_enabled() {
return None;
}
crate::bridge::pending_approval_for_user_thread(user_id, Some(&thread_id.to_string()))
.await
.ok()
.flatten()
.map(|pending| PendingApprovalInfo {
request_id: pending.request_id,
tool_name: pending.tool_name,
description: pending.description,
parameters: pending.parameters,
})
}
pub async fn chat_history_handler(
State(state): State<Arc<GatewayState>>,
AuthenticatedUser(identity): AuthenticatedUser,
@@ -357,12 +374,13 @@ pub async fn chat_history_handler(
let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339());
let turns = build_turns_from_db_messages(&messages);
let pending_approval = engine_pending_approval(&identity.user_id, thread_id).await;
return Ok(Json(HistoryResponse {
thread_id,
turns,
has_more,
oldest_timestamp,
pending_approval: None,
pending_approval,
}));
}
@@ -412,6 +430,10 @@ pub async fn chat_history_handler(
description: pa.description.clone(),
parameters: serde_json::to_string_pretty(&pa.parameters).unwrap_or_default(),
});
let pending_approval = match pending_approval {
Some(pending) => Some(pending),
None => engine_pending_approval(&identity.user_id, thread_id).await,
};
return Ok(Json(HistoryResponse {
thread_id,
@@ -433,23 +455,25 @@ pub async fn chat_history_handler(
if !messages.is_empty() {
let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339());
let turns = build_turns_from_db_messages(&messages);
let pending_approval = engine_pending_approval(&identity.user_id, thread_id).await;
return Ok(Json(HistoryResponse {
thread_id,
turns,
has_more,
oldest_timestamp,
pending_approval: None,
pending_approval,
}));
}
}
// Empty thread (just created, no messages yet)
let pending_approval = engine_pending_approval(&identity.user_id, thread_id).await;
Ok(Json(HistoryResponse {
thread_id,
turns: Vec::new(),
has_more: false,
oldest_timestamp: None,
pending_approval: None,
pending_approval,
}))
}
+2 -2
View File
@@ -88,8 +88,8 @@ pub struct HistoryResponse {
pub oldest_timestamp: Option<String>,
/// Pending tool approval that needs user action (re-rendered on thread switch).
///
/// Only populated from in-memory state; not persisted to DB.
/// Server restart clears pending approvals.
/// Populated from live session state and, when engine v2 is active,
/// durable engine metadata so approval cards survive thread switches and restarts.
#[serde(skip_serializing_if = "Option::is_none")]
pub pending_approval: Option<PendingApprovalInfo>,
}