feat(bridge): implement tool approval flow for engine v2

Adds a complete approval flow that mirrors v1 behavior, using the
existing v1 security controls (Tool::requires_approval, auto-approve
sets, StatusUpdate::ApprovalNeeded).

## How it works

### Step 1: Tool blocked at execution
When the LLM's code calls a tool (e.g., `shell("ls")`):
1. EffectBridgeAdapter.execute_action() looks up the Tool object
2. Calls tool.requires_approval(&params) — returns ApprovalRequirement
3. If Always → EngineError::LeaseDenied (always blocks)
4. If UnlessAutoApproved → checks auto_approved HashSet → if not in set,
   returns EngineError::LeaseDenied
5. If Never → proceeds to execution

### Step 2: Engine returns NeedApproval
The LeaseDenied error propagates through:
- CodeAct path: becomes Python RuntimeError, code halts, thread returns
  NeedApproval with action_name + parameters
- Structured path: same via ActionResult.is_error

### Step 3: Router stores pending approval
- PendingApproval { action_name, original_content } stored on EngineState
- StatusUpdate::ApprovalNeeded sent to channel (shows approval card in
  CLI/web with tool name, parameters, yes/always/no buttons)
- Returns text: "Tool 'shell' requires approval. Reply yes/always/no."

### Step 4: User responds
handle_message() intercepts Submission::ApprovalResponse when ENGINE_V2:
- 'yes' → auto_approve_tool(name) on EffectBridgeAdapter, re-processes
  original message (tool now passes the approval check on second run)
- 'always' → same + logs for session persistence
- 'no' → returns "Denied: tool was not executed."

### Key design choice
Instead of pausing/resuming mid-execution (which needs engine changes
to freeze/restore the Monty VM state), we auto-approve the tool and
re-run the full message. The EffectBridgeAdapter's auto_approved set
persists across runs, so the second execution passes immediately.

This trades one extra LLM call for zero engine modifications.

## Files changed
- src/bridge/router.rs: PendingApproval struct, handle_approval(),
  NeedApproval → StatusUpdate::ApprovalNeeded conversion
- src/bridge/mod.rs: export handle_approval
- src/agent/agent_loop.rs: intercept ApprovalResponse for engine v2
- src/bridge/effect_adapter.rs: fmt fixes

151 tests passing, clippy + fmt clean.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
2026-03-23 21:10:40 -07:00
co-authored by Claude Opus 4.6
parent 38fc870223
commit e82dcbd5e6
31 changed files with 2125 additions and 436 deletions
+18 -8
View File
@@ -35,32 +35,42 @@ src/
│ ├── capability.rs # Capability, ActionDef, EffectType, CapabilityLease, PolicyRule
│ ├── memory.rs # MemoryDoc, DocId, DocType (Summary/Lesson/Playbook/Issue/Spec/Note)
│ ├── project.rs # Project, ProjectId
│ ├── event.rs # ThreadEvent, EventKind (16 variants for event sourcing)
│ ├── event.rs # ThreadEvent, EventKind (18 variants for event sourcing)
│ ├── message.rs # ThreadMessage, MessageRole
│ ├── provenance.rs # Provenance enum (User/System/ToolOutput/LlmGenerated/etc.)
│ ├── conversation.rs # ConversationSurface, ConversationEntry, EntrySender
│ ├── mission.rs # Mission, MissionId, MissionCadence, MissionStatus
│ └── error.rs # EngineError, ThreadError, StepError, CapabilityError
├── traits/ # External dependency abstractions (host implements these)
│ ├── llm.rs # LlmBackend trait
│ ├── store.rs # Store trait (18 CRUD methods)
│ ├── store.rs # Store trait (20 CRUD methods)
│ └── effect.rs # EffectExecutor trait
├── capability/ # Capability management
│ ├── registry.rs # CapabilityRegistry — register/get/list capabilities
│ ├── lease.rs # LeaseManager — grant/check/consume/revoke/expire leases
│ └── policy.rs # PolicyEngine — deterministic effect-level allow/deny/approve
│ └── policy.rs # PolicyEngine — deterministic effect-level allow/deny/approve + provenance taint
├── runtime/ # Thread lifecycle management
│ ├── manager.rs # ThreadManager — spawn, stop, inject messages, join threads
│ ├── conversation.rs # ConversationManager — routes UI messages to threads
│ ├── mission.rs # MissionManager — long-running goals that spawn threads on cadence
│ ├── tree.rs # ThreadTree — parent-child relationships
│ └── messaging.rs # ThreadSignal, ThreadOutcome, signal channels
├── executor/ # Step execution
│ ├── loop_engine.rs # ExecutionLoop — core loop replacing run_agentic_loop()
│ ├── structured.rs # Tier 0: structured tool call execution
│ ├── context.rs # Context builder (messages + actions from leases)
── intent.rs # Tool intent nudge detection
│ ├── scripting.rs # Tier 1: embedded Python via Monty (CodeAct/RLM)
── context.rs # Context builder (messages + actions from leases + memory docs)
│ ├── compaction.rs # Context compaction when approaching model context limit
│ ├── prompt.rs # System prompt construction (CodeAct preamble/postamble)
│ ├── intent.rs # Tool intent nudge detection
│ └── trace.rs # Execution trace recording and retrospective analysis
├── memory/ # Memory document system
│ ├── store.rs # MemoryStore — project-scoped doc CRUD
│ └── retrieval.rs # RetrievalEngine — context building (stub, Phase 4)
── reflection/ # Post-thread reflection (stub, Phase 4)
── mod.rs
│ └── retrieval.rs # RetrievalEngine — keyword-based context retrieval from project docs
── reflection/ # Post-thread reflection pipeline
── pipeline.rs # reflect() (CodeAct) + reflect_simple() (direct LLM) + output parsing
│ └── executor.rs # ReflectionExecutor — read-only tools for reflection threads
└── reliability.rs # ReliabilityTracker — per-action success rate and latency via EMA
```
## Thread State Machine
+20 -8
View File
@@ -53,9 +53,11 @@ impl LeaseManager {
/// Check whether a lease is still valid. Returns the lease if valid.
pub async fn check(&self, lease_id: LeaseId) -> Result<CapabilityLease, EngineError> {
let leases = self.active.read().await;
let lease = leases.get(&lease_id).ok_or_else(|| EngineError::LeaseExpired {
capability_name: format!("lease {lease_id:?} not found"),
})?;
let lease = leases
.get(&lease_id)
.ok_or_else(|| EngineError::LeaseExpired {
capability_name: format!("lease {lease_id:?} not found"),
})?;
if !lease.is_valid() {
return Err(EngineError::LeaseExpired {
capability_name: lease.capability_name.clone(),
@@ -67,9 +69,11 @@ impl LeaseManager {
/// Consume one use of a lease. Returns error if the lease is invalid or exhausted.
pub async fn consume_use(&self, lease_id: LeaseId) -> Result<(), EngineError> {
let mut leases = self.active.write().await;
let lease = leases.get_mut(&lease_id).ok_or_else(|| EngineError::LeaseExpired {
capability_name: format!("lease {lease_id:?} not found"),
})?;
let lease = leases
.get_mut(&lease_id)
.ok_or_else(|| EngineError::LeaseExpired {
capability_name: format!("lease {lease_id:?} not found"),
})?;
if !lease.is_valid() {
return Err(EngineError::LeaseExpired {
capability_name: lease.capability_name.clone(),
@@ -202,8 +206,16 @@ mod tests {
None,
)
.await;
assert!(mgr.find_lease_for_action(tid, "create_issue").await.is_some());
assert!(mgr.find_lease_for_action(tid, "delete_repo").await.is_none());
assert!(
mgr.find_lease_for_action(tid, "create_issue")
.await
.is_some()
);
assert!(
mgr.find_lease_for_action(tid, "delete_repo")
.await
.is_none()
);
}
#[tokio::test]
@@ -314,12 +314,8 @@ mod tests {
let engine = PolicyEngine::new();
let action = make_action("transfer_funds", vec![EffectType::Financial], false);
let lease = make_lease();
let decision = engine.evaluate_with_provenance(
&action,
&lease,
&[],
&Provenance::LlmGenerated,
);
let decision =
engine.evaluate_with_provenance(&action, &lease, &[], &Provenance::LlmGenerated);
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
}
@@ -328,12 +324,8 @@ mod tests {
let engine = PolicyEngine::new();
let action = make_action("post_message", vec![EffectType::WriteExternal], false);
let lease = make_lease();
let decision = engine.evaluate_with_provenance(
&action,
&lease,
&[],
&Provenance::LlmGenerated,
);
let decision =
engine.evaluate_with_provenance(&action, &lease, &[], &Provenance::LlmGenerated);
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
}
@@ -342,12 +334,7 @@ mod tests {
let engine = PolicyEngine::new();
let action = make_action("transfer_funds", vec![EffectType::Financial], false);
let lease = make_lease();
let decision = engine.evaluate_with_provenance(
&action,
&lease,
&[],
&Provenance::User,
);
let decision = engine.evaluate_with_provenance(&action, &lease, &[], &Provenance::User);
assert_eq!(decision, PolicyDecision::Allow);
}
@@ -360,7 +347,9 @@ mod tests {
&action,
&lease,
&[],
&Provenance::ToolOutput { action_name: "scrape_invoices".into() },
&Provenance::ToolOutput {
action_name: "scrape_invoices".into(),
},
);
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
}
@@ -20,7 +20,8 @@ impl CapabilityRegistry {
/// Register a capability. Overwrites any existing capability with the same name.
pub fn register(&mut self, capability: Capability) {
self.capabilities.insert(capability.name.clone(), capability);
self.capabilities
.insert(capability.name.clone(), capability);
}
/// Look up a capability by name.
@@ -26,9 +26,7 @@ pub fn estimate_tokens(messages: &[ThreadMessage]) -> usize {
let total_chars: usize = messages
.iter()
.map(|m| {
m.content.len()
+ m.action_name.as_ref().map_or(0, |n| n.len())
+ 4 // overhead per message (role token, delimiters)
m.content.len() + m.action_name.as_ref().map_or(0, |n| n.len()) + 4 // overhead per message (role token, delimiters)
})
.sum();
total_chars.div_ceil(CHARS_PER_TOKEN)
+103 -45
View File
@@ -6,12 +6,12 @@
use std::sync::Arc;
use crate::memory::RetrievalEngine;
use crate::traits::effect::EffectExecutor;
use crate::types::capability::{ActionDef, CapabilityLease};
use crate::types::error::EngineError;
use crate::types::memory::MemoryDoc;
use crate::types::message::ThreadMessage;
use crate::types::project::ProjectId;
use crate::traits::effect::EffectExecutor;
/// Maximum number of memory docs to inject into context.
const MAX_CONTEXT_DOCS: usize = 5;
@@ -74,7 +74,11 @@ fn format_docs_as_context(docs: &[MemoryDoc]) -> String {
};
// Truncate long docs to avoid context bloat
let content: String = doc.content.chars().take(500).collect();
let truncated = if doc.content.len() > 500 { "..." } else { "" };
let truncated = if doc.content.chars().count() > 500 {
"..."
} else {
""
};
parts.push(format!(
"### [{type_label}] {}\n{content}{truncated}\n",
doc.title
@@ -126,36 +130,102 @@ mod tests {
#[async_trait::async_trait]
impl crate::traits::store::Store for DocStore {
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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> { Ok(()) }
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> { Ok(None) }
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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> {
Ok(())
}
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> {
Ok(None)
}
async fn list_memory_docs(&self, pid: ProjectId) -> Result<Vec<MemoryDoc>, EngineError> {
Ok(self.0.iter().filter(|d| d.project_id == pid).cloned().collect())
Ok(self
.0
.iter()
.filter(|d| d.project_id == pid)
.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,
_: &crate::types::mission::Mission,
) -> Result<(), EngineError> {
Ok(())
}
async fn load_mission(
&self,
_: crate::types::mission::MissionId,
) -> Result<Option<crate::types::mission::Mission>, EngineError> {
Ok(None)
}
async fn list_missions(
&self,
_: ProjectId,
) -> Result<Vec<crate::types::mission::Mission>, EngineError> {
Ok(vec![])
}
async fn update_mission_status(
&self,
_: crate::types::mission::MissionId,
_: crate::types::mission::MissionStatus,
) -> Result<(), EngineError> {
Ok(())
}
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, _: &crate::types::mission::Mission) -> Result<(), EngineError> { Ok(()) }
async fn load_mission(&self, _: crate::types::mission::MissionId) -> Result<Option<crate::types::mission::Mission>, EngineError> { Ok(None) }
async fn list_missions(&self, _: ProjectId) -> Result<Vec<crate::types::mission::Mission>, EngineError> { Ok(vec![]) }
async fn update_mission_status(&self, _: crate::types::mission::MissionId, _: crate::types::mission::MissionStatus) -> Result<(), EngineError> { Ok(()) }
}
#[tokio::test]
async fn context_injects_docs_after_system_prompt() {
let project = ProjectId::new();
let store: Arc<dyn crate::traits::store::Store> = Arc::new(DocStore(vec![
MemoryDoc::new(project, DocType::Lesson, "web tool alias", "Use web-search not web_search"),
]));
let store: Arc<dyn crate::traits::store::Store> = Arc::new(DocStore(vec![MemoryDoc::new(
project,
DocType::Lesson,
"web tool alias",
"Use web-search not web_search",
)]));
let retrieval = RetrievalEngine::new(store);
let effects: Arc<dyn EffectExecutor> = Arc::new(MockEffects);
@@ -193,16 +263,10 @@ mod tests {
ThreadMessage::user("hello"),
];
let (ctx_msgs, _) = build_step_context(
&messages,
&[],
&effects,
None,
ProjectId::new(),
"hello",
)
.await
.unwrap();
let (ctx_msgs, _) =
build_step_context(&messages, &[], &effects, None, ProjectId::new(), "hello")
.await
.unwrap();
// No injection — same number of messages
assert_eq!(ctx_msgs.len(), 2);
@@ -217,16 +281,10 @@ mod tests {
let messages = vec![ThreadMessage::user("hello")];
let (ctx_msgs, _) = build_step_context(
&messages,
&[],
&effects,
Some(&retrieval),
project,
"hello",
)
.await
.unwrap();
let (ctx_msgs, _) =
build_step_context(&messages, &[], &effects, Some(&retrieval), project, "hello")
.await
.unwrap();
assert_eq!(ctx_msgs.len(), 1);
}
@@ -58,8 +58,7 @@ pub fn signals_tool_intent(response: &str) -> bool {
}
/// The nudge message injected into context when tool intent is detected.
pub const TOOL_INTENT_NUDGE: &str =
"You expressed intent to use a tool but didn't make an action call. \
pub const TOOL_INTENT_NUDGE: &str = "You expressed intent to use a tool but didn't make an action call. \
Please go ahead and call the appropriate action.";
#[cfg(test)]
@@ -33,6 +33,8 @@ pub struct ExecutionLoop {
policy: Arc<PolicyEngine>,
signal_rx: SignalReceiver,
user_id: String,
/// Optional capability registry for resolving capability-level policies.
capabilities: Option<Arc<crate::capability::registry::CapabilityRegistry>>,
/// Optional broadcast sender for live event streaming.
event_tx: Option<tokio::sync::broadcast::Sender<crate::types::event::ThreadEvent>>,
/// Optional retrieval engine for injecting prior knowledge into context.
@@ -57,6 +59,7 @@ impl ExecutionLoop {
policy,
signal_rx,
user_id,
capabilities: None,
event_tx: None,
retrieval: None,
}
@@ -71,6 +74,15 @@ impl ExecutionLoop {
self
}
/// Set the capability registry for resolving capability-level policies.
pub fn with_capabilities(
mut self,
capabilities: Arc<crate::capability::registry::CapabilityRegistry>,
) -> Self {
self.capabilities = Some(capabilities);
self
}
/// Set the retrieval engine for injecting prior knowledge into context.
pub fn with_retrieval(mut self, retrieval: crate::memory::RetrievalEngine) -> Self {
self.retrieval = Some(retrieval);
@@ -93,13 +105,25 @@ impl ExecutionLoop {
self.thread.transition_to(ThreadState::Running, None)?;
// Inject CodeAct/RLM system prompt if none exists
if !self.thread.messages.iter().any(|m| m.role == crate::types::message::MessageRole::System) {
if !self
.thread
.messages
.iter()
.any(|m| m.role == crate::types::message::MessageRole::System)
{
// Get available actions for the prompt
let active_leases = self.leases.active_for_thread(self.thread.id).await;
let actions = self.effects.available_actions(&active_leases).await
.unwrap_or_default();
let actions = match self.effects.available_actions(&active_leases).await {
Ok(a) => a,
Err(e) => {
warn!(thread_id = %self.thread.id, "failed to load actions for system prompt: {e}");
Vec::new()
}
};
let system_prompt = crate::executor::prompt::build_codeact_system_prompt(&actions);
self.thread.messages.insert(0, ThreadMessage::system(system_prompt));
self.thread
.messages
.insert(0, ThreadMessage::system(system_prompt));
}
let max_iterations = self.thread.config.max_iterations;
@@ -138,10 +162,8 @@ impl ExecutionLoop {
limit = max_tokens,
"token limit exceeded"
);
self.thread.transition_to(
ThreadState::Completed,
Some("token limit exceeded".into()),
)?;
self.thread
.transition_to(ThreadState::Completed, Some("token limit exceeded".into()))?;
return Ok(ThreadOutcome::Failed {
error: format!(
"Token limit exceeded: {} of {} tokens",
@@ -176,10 +198,8 @@ impl ExecutionLoop {
limit = max_usd,
"USD budget exceeded"
);
self.thread.transition_to(
ThreadState::Completed,
Some("USD budget exceeded".into()),
)?;
self.thread
.transition_to(ThreadState::Completed, Some("USD budget exceeded".into()))?;
return Ok(ThreadOutcome::Failed {
error: format!(
"USD budget exceeded: ${:.4} of ${:.4}",
@@ -236,9 +256,7 @@ impl ExecutionLoop {
// 6. Create step
let mut step = Step::new(self.thread.id, iteration + 1);
step.status = StepStatus::LlmCalling;
self.emit_event(EventKind::StepStarted {
step_id: step.id,
});
self.emit_event(EventKind::StepStarted { step_id: step.id });
// 7. Call LLM
// CodeAct/RLM: send NO structured tool definitions — tools are described
@@ -308,8 +326,7 @@ impl ExecutionLoop {
// sometimes write FINAL() outside code blocks)
if let Some(answer) = extract_final_from_text(&text) {
debug!(thread_id = %self.thread.id, "FINAL() detected in text response");
self.thread
.add_message(ThreadMessage::assistant(text));
self.thread.add_message(ThreadMessage::assistant(text));
step.status = StepStatus::Completed;
step.completed_at = Some(chrono::Utc::now());
self.emit_event(EventKind::StepCompleted {
@@ -337,8 +354,7 @@ impl ExecutionLoop {
nudge_count,
"tool intent detected, injecting nudge"
);
self.thread
.add_message(ThreadMessage::assistant(text));
self.thread.add_message(ThreadMessage::assistant(text));
self.thread
.add_message(ThreadMessage::user(intent::TOOL_INTENT_NUDGE));
@@ -364,10 +380,8 @@ impl ExecutionLoop {
});
self.thread.step_count += 1;
self.thread.transition_to(
ThreadState::Completed,
Some("text response".into()),
)?;
self.thread
.transition_to(ThreadState::Completed, Some("text response".into()))?;
return Ok(ThreadOutcome::Completed {
response: Some(text),
});
@@ -378,7 +392,10 @@ impl ExecutionLoop {
// Record assistant message with action calls
self.thread
.add_message(ThreadMessage::assistant_with_actions(content, calls.clone()));
.add_message(ThreadMessage::assistant_with_actions(
content,
calls.clone(),
));
step.status = StepStatus::Executing;
@@ -391,6 +408,21 @@ impl ExecutionLoop {
step_id: step.id,
};
// Collect capability-level policies from the registry.
// Gathers policies from all registered capabilities — they're
// cheap and deterministic. Per-action scoping happens inside
// the policy engine via condition matching.
let cap_policies: Vec<crate::types::capability::PolicyRule> = self
.capabilities
.as_ref()
.map(|caps| {
caps.list()
.iter()
.flat_map(|c| c.policies.iter().cloned())
.collect()
})
.unwrap_or_default();
// Execute actions
let batch = execute_action_calls(
&calls,
@@ -399,7 +431,7 @@ impl ExecutionLoop {
&self.leases,
&self.policy,
&exec_ctx,
&[], // capability-level policies (TODO: resolve from registry)
&cap_policies,
)
.await?;
@@ -428,8 +460,10 @@ impl ExecutionLoop {
// Check if approval is needed
if let Some(outcome) = batch.need_approval {
self.thread
.transition_to(ThreadState::Waiting, Some("awaiting approval".into()))?;
self.thread.transition_to(
ThreadState::Waiting,
Some("awaiting approval".into()),
)?;
return Ok(outcome);
}
}
@@ -559,15 +593,18 @@ impl ExecutionLoop {
output_str
};
if result.is_error {
output_parts.push(format!("[{} error] {}", result.action_name, truncated));
output_parts
.push(format!("[{} error] {}", result.action_name, truncated));
} else {
output_parts.push(format!("[{} result] {}", result.action_name, truncated));
output_parts
.push(format!("[{} result] {}", result.action_name, truncated));
}
}
if code_result.return_value != serde_json::Value::Null {
output_parts.push(format!(
"[return] {}",
serde_json::to_string_pretty(&code_result.return_value).unwrap_or_default()
serde_json::to_string_pretty(&code_result.return_value)
.unwrap_or_default()
));
}
let output_text = if output_parts.is_empty() {
@@ -577,14 +614,22 @@ impl ExecutionLoop {
};
// Truncate total output to prevent context bloat
let mut metadata = if output_text.chars().count() > 8000 {
let tail: String = output_text.chars().skip(output_text.chars().count() - 8000).collect();
format!("[TRUNCATED: last 8000 of {} chars]\n{tail}", output_text.chars().count())
let tail: String = output_text
.chars()
.skip(output_text.chars().count() - 8000)
.collect();
format!(
"[TRUNCATED: last 8000 of {} chars]\n{tail}",
output_text.chars().count()
)
} else {
output_text
};
// If code had errors, remind the model about `state`
if code_result.had_error && !persisted_state.as_object().is_some_and(|m| m.is_empty()) {
if code_result.had_error
&& !persisted_state.as_object().is_some_and(|m| m.is_empty())
{
let keys: Vec<&str> = persisted_state
.as_object()
.map(|m| m.keys().map(String::as_str).collect())
@@ -607,10 +652,8 @@ impl ExecutionLoop {
// Check FINAL() termination
if let Some(answer) = code_result.final_answer {
self.thread.transition_to(
ThreadState::Completed,
Some("FINAL() called".into()),
)?;
self.thread
.transition_to(ThreadState::Completed, Some("FINAL() called".into()))?;
return Ok(ThreadOutcome::Completed {
response: Some(answer),
});
@@ -618,8 +661,10 @@ impl ExecutionLoop {
// Check if approval is needed
if let Some(outcome) = code_result.need_approval {
self.thread
.transition_to(ThreadState::Waiting, Some("awaiting approval".into()))?;
self.thread.transition_to(
ThreadState::Waiting,
Some("awaiting approval".into()),
)?;
return Ok(outcome);
}
@@ -763,11 +808,11 @@ fn extract_final_from_text(text: &str) -> Option<String> {
#[cfg(test)]
mod tests {
use super::*;
use crate::traits::llm::{LlmCallConfig, LlmOutput};
use crate::types::capability::{ActionDef, CapabilityLease, EffectType};
use crate::types::project::ProjectId;
use crate::types::step::{ActionResult, TokenUsage};
use crate::types::thread::{ThreadConfig, ThreadType};
use crate::traits::llm::{LlmCallConfig, LlmOutput};
use std::sync::Mutex;
use std::time::Duration;
@@ -913,21 +958,11 @@ mod tests {
let policy = Arc::new(PolicyEngine::new());
// Grant a default lease
leases
.grant(tid, "test_cap", vec![], None, None)
.await;
leases.grant(tid, "test_cap", vec![], None, None).await;
let (tx, rx) = crate::runtime::messaging::signal_channel(16);
let exec = ExecutionLoop::new(
thread,
llm,
effects,
leases,
policy,
rx,
"test-user".into(),
);
let exec = ExecutionLoop::new(thread, llm, effects, leases, policy, rx, "test-user".into());
(exec, tx)
}
@@ -1059,11 +1094,12 @@ mod tests {
let outcome = exec.run().await.unwrap();
assert!(matches!(outcome, ThreadOutcome::Completed { .. }));
assert!(exec
.thread
.messages
.iter()
.any(|m| m.content == "injected!"));
assert!(
exec.thread
.messages
.iter()
.any(|m| m.content == "injected!")
);
}
#[tokio::test]
@@ -1088,11 +1124,12 @@ mod tests {
);
assert_eq!(exec.thread.step_count, 2);
// Should have nudge system message
assert!(exec
.thread
.messages
.iter()
.any(|m| m.content.contains("didn't make an action call")));
assert!(
exec.thread
.messages
.iter()
.any(|m| m.content.contains("didn't make an action call"))
);
}
#[tokio::test]
@@ -1165,9 +1202,9 @@ mod tests {
async fn codeact_tool_call_then_final() {
// LLM outputs code that calls a tool, then uses the result
let (mut exec, _tx) = make_loop(
vec![
code_response("result = test_tool()\nprint(result)\nFINAL('got result')"),
],
vec![code_response(
"result = test_tool()\nprint(result)\nFINAL('got result')",
)],
vec![Ok(ActionResult {
call_id: "code_call_1".into(),
action_name: "test_tool".into(),
@@ -1225,7 +1262,12 @@ mod tests {
);
assert_eq!(exec.thread.step_count, 2);
// The output metadata from first step should be in messages
assert!(exec.thread.messages.iter().any(|m| m.content.contains("x = 30")));
assert!(
exec.thread
.messages
.iter()
.any(|m| m.content.contains("x = 30"))
);
}
#[tokio::test]
@@ -1243,12 +1285,17 @@ mod tests {
.await;
let outcome = exec.run().await.unwrap();
assert!(matches!(outcome, ThreadOutcome::Completed { response: Some(r) } if r == "recovered"));
assert!(
matches!(outcome, ThreadOutcome::Completed { response: Some(r) } if r == "recovered")
);
assert_eq!(exec.thread.step_count, 2);
// First step should have error in output metadata
assert!(exec.thread.messages.iter().any(|m| {
m.content.contains("NameError") || m.content.contains("Error")
}));
assert!(
exec.thread
.messages
.iter()
.any(|m| { m.content.contains("NameError") || m.content.contains("Error") })
);
}
#[tokio::test]
@@ -1321,7 +1368,9 @@ mod tests {
let (mut exec, _tx) = make_loop(
vec![
// First response: code that calls llm_query
code_response("answer = llm_query('What is 2+2?')\nFINAL(f'Sub-agent said: {answer}')"),
code_response(
"answer = llm_query('What is 2+2?')\nFINAL(f'Sub-agent said: {answer}')",
),
// This text response will be consumed by the llm_query sub-call
// (MockLlm pops from the same queue)
],
@@ -79,7 +79,10 @@ pub fn compact_output_metadata(stdout: &str, return_value: &serde_json::Value) -
if !stdout.is_empty() {
if stdout.chars().count() > OUTPUT_TRUNCATE_LEN {
let truncated: String = stdout.chars().skip(stdout.chars().count() - OUTPUT_TRUNCATE_LEN).collect();
let truncated: String = stdout
.chars()
.skip(stdout.chars().count() - OUTPUT_TRUNCATE_LEN)
.collect();
parts.push(format!(
"[TRUNCATED: last {OUTPUT_TRUNCATE_LEN} of {} chars shown]\n{truncated}",
stdout.len()
@@ -130,7 +133,7 @@ pub fn build_orientation_preamble(thread: &Thread) -> String {
.find(|m| m.role == MessageRole::User)
{
let content_preview: String = last_user.content.chars().take(500).collect();
let truncated = if last_user.content.len() > 500 {
let truncated = if last_user.content.chars().count() > 500 {
"..."
} else {
""
@@ -332,11 +335,7 @@ pub async fn execute_code(
let ext_result = match action_name.as_str() {
// FINAL(answer) — explicit termination
"FINAL" => {
let answer = call
.args
.first()
.map(monty_to_string)
.unwrap_or_default();
let answer = call.args.first().map(monty_to_string).unwrap_or_default();
final_answer = Some(answer);
ExtFunctionResult::Return(MontyObject::None)
}
@@ -359,8 +358,7 @@ pub async fn execute_code(
// llm_query(prompt, context) — recursive sub-call
"llm_query" => {
handle_llm_query(&call.args, &call.kwargs, llm, &mut recursive_tokens)
.await
handle_llm_query(&call.args, &call.kwargs, llm, &mut recursive_tokens).await
}
// llm_query_batched(prompts) — parallel sub-calls
@@ -741,12 +739,14 @@ async fn handle_rlm_query(
max_iterations: parent_thread.config.max_iterations.min(20), // cap child iterations
enable_reflection: false,
enable_tool_intent_nudge: false,
max_tokens_total: parent_thread.config.max_tokens_total.map(|max| {
max.saturating_sub(parent_thread.total_tokens_used)
}),
max_budget_usd: parent_thread.config.max_budget_usd.map(|max| {
(max - parent_thread.total_cost_usd).max(0.0)
}),
max_tokens_total: parent_thread
.config
.max_tokens_total
.map(|max| max.saturating_sub(parent_thread.total_tokens_used)),
max_budget_usd: parent_thread
.config
.max_budget_usd
.map(|max| (max - parent_thread.total_cost_usd).max(0.0)),
max_duration: parent_thread.config.max_duration,
depth: current_depth + 1,
max_depth,
@@ -1031,9 +1031,7 @@ fn json_to_monty(val: &serde_json::Value) -> MontyObject {
}
}
serde_json::Value::String(s) => MontyObject::String(s.clone()),
serde_json::Value::Array(arr) => {
MontyObject::List(arr.iter().map(json_to_monty).collect())
}
serde_json::Value::Array(arr) => MontyObject::List(arr.iter().map(json_to_monty).collect()),
serde_json::Value::Object(map) => MontyObject::dict(
map.iter()
.map(|(k, v)| (MontyObject::String(k.clone()), json_to_monty(v)))
@@ -48,7 +48,10 @@ pub async fn execute_action_calls(
for call in calls {
// 1. Find the lease for this action
let lease = match leases.find_lease_for_action(thread.id, &call.action_name).await {
let lease = match leases
.find_lease_for_action(thread.id, &call.action_name)
.await
{
Some(l) => l,
None => {
let error_result = ActionResult {
+16 -18
View File
@@ -88,8 +88,8 @@ pub fn build_trace(thread: &Thread) -> ExecutionTrace {
let preview: String = m.content.chars().take(300).collect();
MessageRecord {
role: format!("{:?}", m.role),
content_length: m.content.len(),
content_preview: if m.content.len() > 300 {
content_length: m.content.chars().count(),
content_preview: if m.content.chars().count() > 300 {
format!("{preview}...")
} else {
preview
@@ -119,10 +119,7 @@ pub fn build_trace(thread: &Thread) -> ExecutionTrace {
/// Write a trace to a JSON file.
pub fn write_trace(trace: &ExecutionTrace) -> Option<PathBuf> {
let filename = format!(
"engine_trace_{}.json",
Utc::now().format("%Y%m%dT%H%M%S")
);
let filename = format!("engine_trace_{}.json", Utc::now().format("%Y%m%dT%H%M%S"));
let path = PathBuf::from(&filename);
match serde_json::to_string_pretty(trace) {
@@ -205,7 +202,11 @@ pub fn log_trace_summary(trace: &ExecutionTrace) {
);
for doc in &refl.docs {
let preview: String = doc.content.chars().take(200).collect();
let truncated = if doc.content.len() > 200 { "..." } else { "" };
let truncated = if doc.content.chars().count() > 200 {
"..."
} else {
""
};
info!(
doc_type = %doc.doc_type,
title = %doc.title,
@@ -240,7 +241,8 @@ fn analyze_trace(thread: &Thread) -> Vec<TraceIssue> {
issues.push(TraceIssue {
severity: IssueSeverity::Warning,
category: "no_response".into(),
description: "No assistant message in thread — model may not have generated output".into(),
description: "No assistant message in thread — model may not have generated output"
.into(),
step: None,
});
}
@@ -254,9 +256,7 @@ fn analyze_trace(thread: &Thread) -> Vec<TraceIssue> {
if !tool_errors.is_empty() {
for event in &tool_errors {
if let crate::types::event::EventKind::ActionFailed {
action_name,
error,
..
action_name, error, ..
} = &event.kind
{
issues.push(TraceIssue {
@@ -313,7 +313,10 @@ fn analyze_trace(thread: &Thread) -> Vec<TraceIssue> {
issues.push(TraceIssue {
severity: IssueSeverity::Warning,
category: "excessive_steps".into(),
description: format!("Thread took {} steps — may be stuck in a loop", thread.step_count),
description: format!(
"Thread took {} steps — may be stuck in a loop",
thread.step_count
),
step: None,
});
}
@@ -338,12 +341,7 @@ fn analyze_trace(thread: &Thread) -> Vec<TraceIssue> {
let code_steps = thread
.events
.iter()
.filter(|e| {
matches!(
e.kind,
crate::types::event::EventKind::StepStarted { .. }
)
})
.filter(|e| matches!(e.kind, crate::types::event::EventKind::StepStarted { .. }))
.count();
let text_responses_without_code = thread
.messages
+2 -2
View File
@@ -33,12 +33,12 @@ pub use types::error::{CapabilityError, EngineError, StepError, ThreadError};
pub use types::event::{EventId, EventKind, ThreadEvent};
pub use types::memory::{DocId, DocType, MemoryDoc};
pub use types::message::{MessageRole, ThreadMessage};
pub use types::mission::{Mission, MissionCadence, MissionId, MissionStatus};
pub use types::project::{Project, ProjectId};
pub use types::provenance::Provenance;
pub use types::step::{
ActionCall, ActionResult, ExecutionTier, LlmResponse, Step, StepId, StepStatus, TokenUsage,
};
pub use types::mission::{Mission, MissionCadence, MissionId, MissionStatus};
pub use types::thread::{Thread, ThreadConfig, ThreadId, ThreadState, ThreadType};
// ── Re-exports: traits ──────────────────────────────────────
@@ -49,9 +49,9 @@ pub use traits::store::Store;
// ── Re-exports: capability ────────────────────────────────────
pub use capability::registry::CapabilityRegistry;
pub use capability::lease::LeaseManager;
pub use capability::policy::{PolicyDecision, PolicyEngine};
pub use capability::registry::CapabilityRegistry;
// ── Re-exports: runtime ───────────────────────────────────────
+150 -38
View File
@@ -74,12 +74,12 @@ impl RetrievalEngine {
/// Extract lowercase keywords from a query, filtering out stop words.
fn extract_keywords(query: &str) -> Vec<String> {
const STOP_WORDS: &[&str] = &[
"a", "an", "the", "is", "are", "was", "were", "be", "been", "being", "have", "has",
"had", "do", "does", "did", "will", "would", "could", "should", "may", "might", "shall",
"can", "to", "of", "in", "for", "on", "with", "at", "by", "from", "as", "into", "about",
"it", "its", "this", "that", "these", "those", "i", "you", "he", "she", "we", "they",
"what", "which", "who", "how", "when", "where", "why", "and", "or", "but", "not", "no",
"if", "then", "so", "up", "out", "just",
"a", "an", "the", "is", "are", "was", "were", "be", "been", "being", "have", "has", "had",
"do", "does", "did", "will", "would", "could", "should", "may", "might", "shall", "can",
"to", "of", "in", "for", "on", "with", "at", "by", "from", "as", "into", "about", "it",
"its", "this", "that", "these", "those", "i", "you", "he", "she", "we", "they", "what",
"which", "who", "how", "when", "where", "why", "and", "or", "but", "not", "no", "if",
"then", "so", "up", "out", "just",
];
query
@@ -150,29 +150,94 @@ mod tests {
#[async_trait::async_trait]
impl crate::traits::store::Store for DocStore {
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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> { Ok(()) }
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> { Ok(None) }
async fn list_memory_docs(&self, project_id: ProjectId) -> Result<Vec<MemoryDoc>, EngineError> {
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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> {
Ok(())
}
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> {
Ok(None)
}
async fn list_memory_docs(
&self,
project_id: ProjectId,
) -> Result<Vec<MemoryDoc>, EngineError> {
let docs = self.docs.lock().await;
Ok(docs.iter().filter(|d| d.project_id == project_id).cloned().collect())
Ok(docs
.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,
_: &crate::types::mission::Mission,
) -> Result<(), EngineError> {
Ok(())
}
async fn load_mission(
&self,
_: crate::types::mission::MissionId,
) -> Result<Option<crate::types::mission::Mission>, EngineError> {
Ok(None)
}
async fn list_missions(
&self,
_: ProjectId,
) -> Result<Vec<crate::types::mission::Mission>, EngineError> {
Ok(vec![])
}
async fn update_mission_status(
&self,
_: crate::types::mission::MissionId,
_: crate::types::mission::MissionStatus,
) -> Result<(), EngineError> {
Ok(())
}
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, _: &crate::types::mission::Mission) -> Result<(), EngineError> { Ok(()) }
async fn load_mission(&self, _: crate::types::mission::MissionId) -> Result<Option<crate::types::mission::Mission>, EngineError> { Ok(None) }
async fn list_missions(&self, _: ProjectId) -> Result<Vec<crate::types::mission::Mission>, EngineError> { Ok(vec![]) }
async fn update_mission_status(&self, _: crate::types::mission::MissionId, _: crate::types::mission::MissionStatus) -> Result<(), EngineError> { Ok(()) }
}
#[test]
@@ -229,13 +294,31 @@ mod tests {
async fn retrieve_returns_relevant_docs_by_keyword() {
let project = ProjectId::new();
let store = DocStore::new(vec![
MemoryDoc::new(project, DocType::Lesson, "web_search tool alias", "Use web-search not web_search"),
MemoryDoc::new(project, DocType::Summary, "weather query", "Fetched weather data"),
MemoryDoc::new(project, DocType::Issue, "API timeout", "External API timed out"),
MemoryDoc::new(
project,
DocType::Lesson,
"web_search tool alias",
"Use web-search not web_search",
),
MemoryDoc::new(
project,
DocType::Summary,
"weather query",
"Fetched weather data",
),
MemoryDoc::new(
project,
DocType::Issue,
"API timeout",
"External API timed out",
),
]);
let engine = RetrievalEngine::new(store);
let docs = engine.retrieve_context(project, "web_search error", 5).await.unwrap();
let docs = engine
.retrieve_context(project, "web_search error", 5)
.await
.unwrap();
assert!(!docs.is_empty());
// The lesson about web_search should rank first (keyword + type weight)
assert_eq!(docs[0].doc_type, DocType::Lesson);
@@ -247,16 +330,32 @@ mod tests {
let project_a = ProjectId::new();
let project_b = ProjectId::new();
let store = DocStore::new(vec![
MemoryDoc::new(project_a, DocType::Lesson, "Lesson for project A", "Some lesson"),
MemoryDoc::new(project_b, DocType::Lesson, "Lesson for project B", "Other lesson"),
MemoryDoc::new(
project_a,
DocType::Lesson,
"Lesson for project A",
"Some lesson",
),
MemoryDoc::new(
project_b,
DocType::Lesson,
"Lesson for project B",
"Other lesson",
),
]);
let engine = RetrievalEngine::new(store);
let docs_a = engine.retrieve_context(project_a, "lesson", 5).await.unwrap();
let docs_a = engine
.retrieve_context(project_a, "lesson", 5)
.await
.unwrap();
assert_eq!(docs_a.len(), 1);
assert!(docs_a[0].title.contains("project A"));
let docs_b = engine.retrieve_context(project_b, "lesson", 5).await.unwrap();
let docs_b = engine
.retrieve_context(project_b, "lesson", 5)
.await
.unwrap();
assert_eq!(docs_b.len(), 1);
assert!(docs_b[0].title.contains("project B"));
}
@@ -281,7 +380,10 @@ mod tests {
let store = DocStore::new(vec![]);
let engine = RetrievalEngine::new(store);
let docs = engine.retrieve_context(project, "anything", 5).await.unwrap();
let docs = engine
.retrieve_context(project, "anything", 5)
.await
.unwrap();
assert!(docs.is_empty());
}
@@ -289,8 +391,18 @@ mod tests {
async fn retrieve_spec_ranks_above_summary() {
let project = ProjectId::new();
let store = DocStore::new(vec![
MemoryDoc::new(project, DocType::Summary, "Summary of search", "searched the web"),
MemoryDoc::new(project, DocType::Spec, "Missing search tool", "ALIAS: web_search -> web-search"),
MemoryDoc::new(
project,
DocType::Summary,
"Summary of search",
"searched the web",
),
MemoryDoc::new(
project,
DocType::Spec,
"Missing search tool",
"ALIAS: web_search -> web-search",
),
]);
let engine = RetrievalEngine::new(store);
+342
View File
@@ -64,3 +64,345 @@ impl MemoryStore {
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::traits::store::Store;
use crate::types::capability::{CapabilityLease, LeaseId};
use crate::types::error::EngineError;
use crate::types::event::ThreadEvent;
use crate::types::memory::{DocId, DocType, 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};
use super::MemoryStore;
// ── In-memory Store implementation ───────────────────────
struct InMemoryDocStore {
docs: RwLock<Vec<MemoryDoc>>,
threads: RwLock<Vec<Thread>>,
steps: RwLock<Vec<Step>>,
events: RwLock<Vec<ThreadEvent>>,
projects: RwLock<Vec<Project>>,
leases: RwLock<Vec<CapabilityLease>>,
missions: RwLock<Vec<Mission>>,
}
impl InMemoryDocStore {
fn new() -> Self {
Self {
docs: RwLock::new(Vec::new()),
threads: RwLock::new(Vec::new()),
steps: RwLock::new(Vec::new()),
events: RwLock::new(Vec::new()),
projects: RwLock::new(Vec::new()),
leases: RwLock::new(Vec::new()),
missions: RwLock::new(Vec::new()),
}
}
}
#[async_trait::async_trait]
impl Store for InMemoryDocStore {
// ── Thread operations ────────────────────────────────
async fn save_thread(&self, thread: &Thread) -> Result<(), EngineError> {
let mut threads = self.threads.write().await;
threads.retain(|t| t.id != thread.id);
threads.push(thread.clone());
Ok(())
}
async fn load_thread(&self, id: ThreadId) -> Result<Option<Thread>, EngineError> {
let threads = self.threads.read().await;
Ok(threads.iter().find(|t| t.id == id).cloned())
}
async fn list_threads(&self, project_id: ProjectId) -> Result<Vec<Thread>, EngineError> {
let threads = self.threads.read().await;
Ok(threads
.iter()
.filter(|t| t.project_id == project_id)
.cloned()
.collect())
}
async fn update_thread_state(
&self,
id: ThreadId,
state: ThreadState,
) -> Result<(), EngineError> {
let mut threads = self.threads.write().await;
if let Some(t) = threads.iter_mut().find(|t| t.id == id) {
t.state = state;
}
Ok(())
}
// ── Step operations ──────────────────────────────────
async fn save_step(&self, step: &Step) -> Result<(), EngineError> {
let mut steps = self.steps.write().await;
steps.retain(|s| s.id != step.id);
steps.push(step.clone());
Ok(())
}
async fn load_steps(&self, thread_id: ThreadId) -> Result<Vec<Step>, EngineError> {
let steps = self.steps.read().await;
Ok(steps
.iter()
.filter(|s| s.thread_id == thread_id)
.cloned()
.collect())
}
// ── Event operations ─────────────────────────────────
async fn append_events(&self, events: &[ThreadEvent]) -> Result<(), EngineError> {
let mut stored = self.events.write().await;
stored.extend(events.iter().cloned());
Ok(())
}
async fn load_events(&self, thread_id: ThreadId) -> Result<Vec<ThreadEvent>, EngineError> {
let events = self.events.read().await;
Ok(events
.iter()
.filter(|e| e.thread_id == thread_id)
.cloned()
.collect())
}
// ── Project operations ───────────────────────────────
async fn save_project(&self, project: &Project) -> Result<(), EngineError> {
let mut projects = self.projects.write().await;
projects.retain(|p| p.id != project.id);
projects.push(project.clone());
Ok(())
}
async fn load_project(&self, id: ProjectId) -> Result<Option<Project>, EngineError> {
let projects = self.projects.read().await;
Ok(projects.iter().find(|p| p.id == id).cloned())
}
// ── Memory doc operations ────────────────────────────
async fn save_memory_doc(&self, doc: &MemoryDoc) -> Result<(), EngineError> {
let mut docs = self.docs.write().await;
docs.push(doc.clone());
Ok(())
}
async fn load_memory_doc(&self, id: DocId) -> Result<Option<MemoryDoc>, EngineError> {
let docs = self.docs.read().await;
Ok(docs.iter().find(|d| d.id == id).cloned())
}
async fn list_memory_docs(
&self,
project_id: ProjectId,
) -> Result<Vec<MemoryDoc>, EngineError> {
let docs = self.docs.read().await;
Ok(docs
.iter()
.filter(|d| d.project_id == project_id)
.cloned()
.collect())
}
// ── Capability lease operations ──────────────────────
async fn save_lease(&self, lease: &CapabilityLease) -> Result<(), EngineError> {
let mut leases = self.leases.write().await;
leases.retain(|l| l.id != lease.id);
leases.push(lease.clone());
Ok(())
}
async fn load_active_leases(
&self,
thread_id: ThreadId,
) -> Result<Vec<CapabilityLease>, EngineError> {
let leases = self.leases.read().await;
Ok(leases
.iter()
.filter(|l| l.thread_id == thread_id && !l.revoked)
.cloned()
.collect())
}
async fn revoke_lease(&self, lease_id: LeaseId, _reason: &str) -> Result<(), EngineError> {
let mut leases = self.leases.write().await;
if let Some(l) = leases.iter_mut().find(|l| l.id == lease_id) {
l.revoked = true;
}
Ok(())
}
// ── Mission operations ───────────────────────────────
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> {
let missions = self.missions.read().await;
Ok(missions.iter().find(|m| m.id == id).cloned())
}
async fn list_missions(&self, project_id: ProjectId) -> Result<Vec<Mission>, EngineError> {
let missions = self.missions.read().await;
Ok(missions
.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(())
}
}
fn make_store() -> MemoryStore {
MemoryStore::new(Arc::new(InMemoryDocStore::new()))
}
// ── Tests ────────────────────────────────────────────────
#[tokio::test]
async fn create_doc_and_get() {
let store = make_store();
let project_id = ProjectId::new();
let doc = store
.create_doc(project_id, DocType::Summary, "Test Doc", "Some content")
.await
.unwrap();
assert_eq!(doc.title, "Test Doc");
assert_eq!(doc.content, "Some content");
assert_eq!(doc.doc_type, DocType::Summary);
assert_eq!(doc.project_id, project_id);
assert!(doc.source_thread_id.is_none());
let loaded = store.get_doc(doc.id).await.unwrap();
let loaded = loaded.unwrap();
assert_eq!(loaded.id, doc.id);
assert_eq!(loaded.title, "Test Doc");
assert_eq!(loaded.content, "Some content");
}
#[tokio::test]
async fn create_doc_from_thread_links_source() {
let store = make_store();
let project_id = ProjectId::new();
let thread_id = ThreadId::new();
let doc = store
.create_doc_from_thread(
project_id,
DocType::Lesson,
"Thread Lesson",
"Learned something",
thread_id,
)
.await
.unwrap();
assert_eq!(doc.source_thread_id, Some(thread_id));
assert_eq!(doc.doc_type, DocType::Lesson);
let loaded = store.get_doc(doc.id).await.unwrap().unwrap();
assert_eq!(loaded.source_thread_id, Some(thread_id));
}
#[tokio::test]
async fn list_docs_by_project() {
let store = make_store();
let project_a = ProjectId::new();
let project_b = ProjectId::new();
store
.create_doc(project_a, DocType::Note, "A1", "content a1")
.await
.unwrap();
store
.create_doc(project_a, DocType::Note, "A2", "content a2")
.await
.unwrap();
store
.create_doc(project_b, DocType::Note, "B1", "content b1")
.await
.unwrap();
let docs_a = store.list_docs(project_a, None).await.unwrap();
assert_eq!(docs_a.len(), 2);
assert!(docs_a.iter().all(|d| d.project_id == project_a));
let docs_b = store.list_docs(project_b, None).await.unwrap();
assert_eq!(docs_b.len(), 1);
assert_eq!(docs_b[0].title, "B1");
}
#[tokio::test]
async fn list_docs_filters_by_type() {
let store = make_store();
let project_id = ProjectId::new();
store
.create_doc(project_id, DocType::Summary, "S1", "summary content")
.await
.unwrap();
store
.create_doc(project_id, DocType::Lesson, "L1", "lesson content")
.await
.unwrap();
store
.create_doc(project_id, DocType::Summary, "S2", "another summary")
.await
.unwrap();
let summaries = store
.list_docs(project_id, Some(DocType::Summary))
.await
.unwrap();
assert_eq!(summaries.len(), 2);
assert!(summaries.iter().all(|d| d.doc_type == DocType::Summary));
let lessons = store
.list_docs(project_id, Some(DocType::Lesson))
.await
.unwrap();
assert_eq!(lessons.len(), 1);
assert_eq!(lessons[0].title, "L1");
}
#[tokio::test]
async fn get_nonexistent_returns_none() {
let store = make_store();
let result = store.get_doc(DocId::new()).await.unwrap();
assert!(result.is_none());
}
}
@@ -252,3 +252,400 @@ Rules:
- Include "playbook" only if the thread completed successfully with 2+ tool calls
- Skip docs that duplicate existing knowledge (check with query_memory first)
- Keep content concise — each doc should be a few sentences, not paragraphs"#;
#[cfg(test)]
mod tests {
use std::sync::Arc;
use chrono::Utc;
use crate::capability::registry::CapabilityRegistry;
use crate::traits::effect::{EffectExecutor, ThreadExecutionContext};
use crate::traits::store::Store;
use crate::types::capability::{ActionDef, Capability, CapabilityLease, EffectType, LeaseId};
use crate::types::error::EngineError;
use crate::types::event::ThreadEvent;
use crate::types::memory::{DocId, DocType, MemoryDoc};
use crate::types::mission::{Mission, MissionId, MissionStatus};
use crate::types::project::{Project, ProjectId};
use crate::types::step::{Step, StepId};
use crate::types::thread::{Thread, ThreadId, ThreadState, ThreadType};
use super::{ReflectionExecutor, build_reflection_prompt};
// ── MockStore ──────────────────────────────────────────────
struct MockStore {
docs: tokio::sync::Mutex<Vec<MemoryDoc>>,
}
impl MockStore {
fn new(docs: Vec<MemoryDoc>) -> Arc<Self> {
Arc::new(Self {
docs: tokio::sync::Mutex::new(docs),
})
}
fn empty() -> Arc<Self> {
Self::new(vec![])
}
}
#[async_trait::async_trait]
impl Store for MockStore {
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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> {
Ok(())
}
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> {
Ok(None)
}
async fn list_memory_docs(
&self,
project_id: ProjectId,
) -> Result<Vec<MemoryDoc>, EngineError> {
let docs = self.docs.lock().await;
Ok(docs
.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) -> Result<(), EngineError> {
Ok(())
}
async fn load_mission(&self, _: MissionId) -> Result<Option<Mission>, EngineError> {
Ok(None)
}
async fn list_missions(&self, _: ProjectId) -> Result<Vec<Mission>, EngineError> {
Ok(vec![])
}
async fn update_mission_status(
&self,
_: MissionId,
_: MissionStatus,
) -> Result<(), EngineError> {
Ok(())
}
}
// ── Helpers ────────────────────────────────────────────────
fn make_lease() -> CapabilityLease {
CapabilityLease {
id: LeaseId::new(),
thread_id: ThreadId::new(),
capability_name: "test".into(),
granted_actions: vec![],
granted_at: Utc::now(),
expires_at: None,
max_uses: None,
uses_remaining: None,
revoked: false,
}
}
fn make_ctx() -> ThreadExecutionContext {
ThreadExecutionContext {
thread_id: ThreadId::new(),
thread_type: ThreadType::Reflection,
project_id: ProjectId::new(),
user_id: "test".into(),
step_id: StepId::new(),
}
}
fn make_capability(name: &str, actions: Vec<ActionDef>) -> Capability {
Capability {
name: name.into(),
description: format!("{name} capability"),
actions,
knowledge: vec![],
policies: vec![],
}
}
fn make_action_def(name: &str) -> ActionDef {
ActionDef {
name: name.into(),
description: format!("{name} action"),
parameters_schema: serde_json::json!({"type": "object", "properties": {}}),
effects: vec![EffectType::ReadLocal],
requires_approval: false,
}
}
// ── Tests ──────────────────────────────────────────────────
#[tokio::test]
async fn get_transcript_returns_content() {
let transcript = "Step 1: called web_search\nStep 2: got results\nDone.";
let project_id = ProjectId::new();
let executor = ReflectionExecutor::new(
MockStore::empty(),
Arc::new(CapabilityRegistry::new()),
transcript.to_string(),
project_id,
);
let lease = make_lease();
let ctx = make_ctx();
let result = executor
.execute_action("get_transcript", serde_json::json!({}), &lease, &ctx)
.await
.unwrap();
assert_eq!(result.action_name, "get_transcript");
assert!(!result.is_error);
assert_eq!(result.output["transcript"].as_str().unwrap(), transcript);
}
#[tokio::test]
async fn query_memory_finds_docs() {
let project_id = ProjectId::new();
let docs = vec![
MemoryDoc::new(
project_id,
DocType::Lesson,
"deployment error",
"Fix: restart the service",
),
MemoryDoc::new(
project_id,
DocType::Summary,
"weather check",
"Fetched weather data",
),
];
let store = MockStore::new(docs);
let executor = ReflectionExecutor::new(
store,
Arc::new(CapabilityRegistry::new()),
String::new(),
project_id,
);
let lease = make_lease();
let ctx = make_ctx();
let result = executor
.execute_action(
"query_memory",
serde_json::json!({"query": "deployment error", "max_docs": 5}),
&lease,
&ctx,
)
.await
.unwrap();
assert!(!result.is_error);
let count = result.output["count"].as_u64().unwrap();
assert!(count >= 1, "expected at least 1 doc, got {count}");
let docs_arr = result.output["docs"].as_array().unwrap();
// The deployment error doc should be present
let has_deployment = docs_arr
.iter()
.any(|d| d["title"].as_str().unwrap().contains("deployment"));
assert!(has_deployment, "expected deployment doc in results");
}
#[tokio::test]
async fn check_tool_exists_found() {
let mut registry = CapabilityRegistry::new();
registry.register(make_capability(
"search",
vec![make_action_def("web-search")],
));
let project_id = ProjectId::new();
let executor = ReflectionExecutor::new(
MockStore::empty(),
Arc::new(registry),
String::new(),
project_id,
);
let lease = make_lease();
let ctx = make_ctx();
let result = executor
.execute_action(
"check_tool_exists",
serde_json::json!({"name": "web-search"}),
&lease,
&ctx,
)
.await
.unwrap();
assert!(!result.is_error);
assert_eq!(result.output["exists"], true);
assert!(result.output["similar"].as_array().unwrap().is_empty());
}
#[tokio::test]
async fn check_tool_exists_not_found_suggests_similar() {
let mut registry = CapabilityRegistry::new();
registry.register(make_capability(
"search",
vec![make_action_def("web-search")],
));
let project_id = ProjectId::new();
let executor = ReflectionExecutor::new(
MockStore::empty(),
Arc::new(registry),
String::new(),
project_id,
);
let lease = make_lease();
let ctx = make_ctx();
let result = executor
.execute_action(
"check_tool_exists",
serde_json::json!({"name": "web_search"}),
&lease,
&ctx,
)
.await
.unwrap();
assert!(!result.is_error);
assert_eq!(result.output["exists"], false);
let similar: Vec<String> = result.output["similar"]
.as_array()
.unwrap()
.iter()
.map(|v| v.as_str().unwrap().to_string())
.collect();
assert!(
similar.contains(&"web-search".to_string()),
"expected 'web-search' in similar list, got: {similar:?}"
);
}
#[tokio::test]
async fn list_tools_returns_all() {
let mut registry = CapabilityRegistry::new();
registry.register(make_capability(
"search",
vec![
make_action_def("web-search"),
make_action_def("memory-search"),
],
));
registry.register(make_capability("files", vec![make_action_def("read-file")]));
let project_id = ProjectId::new();
let executor = ReflectionExecutor::new(
MockStore::empty(),
Arc::new(registry),
String::new(),
project_id,
);
let lease = make_lease();
let ctx = make_ctx();
let result = executor
.execute_action("list_tools", serde_json::json!({}), &lease, &ctx)
.await
.unwrap();
assert!(!result.is_error);
assert_eq!(result.output["count"].as_u64().unwrap(), 3);
let tools = result.output["tools"].as_array().unwrap();
let names: Vec<&str> = tools.iter().map(|t| t["name"].as_str().unwrap()).collect();
assert!(names.contains(&"web-search"));
assert!(names.contains(&"memory-search"));
assert!(names.contains(&"read-file"));
}
#[test]
fn build_reflection_prompt_includes_tools() {
let actions = vec![
ActionDef {
name: "get_transcript".into(),
description: "Get the execution transcript".into(),
parameters_schema: serde_json::json!({"type": "object", "properties": {}}),
effects: vec![EffectType::ReadLocal],
requires_approval: false,
},
ActionDef {
name: "query_memory".into(),
description: "Search memory docs".into(),
parameters_schema: serde_json::json!({
"type": "object",
"properties": {
"query": {"type": "string"},
"max_docs": {"type": "integer"}
}
}),
effects: vec![EffectType::ReadLocal],
requires_approval: false,
},
];
let prompt = build_reflection_prompt(&actions, "analyze deployment failure");
assert!(
prompt.contains("get_transcript"),
"prompt should contain get_transcript tool name"
);
assert!(
prompt.contains("query_memory"),
"prompt should contain query_memory tool name"
);
assert!(
prompt.contains("analyze deployment failure"),
"prompt should contain the thread goal"
);
assert!(
prompt.contains("Available tools"),
"prompt should contain the tools section header"
);
}
}
+1 -1
View File
@@ -15,4 +15,4 @@
pub mod executor;
pub mod pipeline;
pub use pipeline::{reflect, reflect_simple, ReflectionResult};
pub use pipeline::{ReflectionResult, reflect, reflect_simple};
@@ -15,7 +15,7 @@ use crate::capability::lease::LeaseManager;
use crate::capability::policy::PolicyEngine;
use crate::capability::registry::CapabilityRegistry;
use crate::executor::ExecutionLoop;
use crate::reflection::executor::{build_reflection_prompt, ReflectionExecutor};
use crate::reflection::executor::{ReflectionExecutor, build_reflection_prompt};
use crate::runtime::messaging::{self, ThreadOutcome};
use crate::traits::llm::LlmBackend;
use crate::traits::store::Store;
@@ -177,7 +177,7 @@ pub async fn reflect_simple(
if thread_failed || had_errors {
let (issue_doc, tokens) =
produce_doc(thread, llm, DocType::Issue, &transcript, ISSUE_PROMPT).await?;
if issue_doc.content.len() > 20 {
if issue_doc.content.chars().count() > 20 {
docs.push(issue_doc);
}
total_tokens.input_tokens += tokens.input_tokens;
@@ -195,7 +195,7 @@ pub async fn reflect_simple(
if has_missing_tools {
let (spec_doc, tokens) =
produce_doc(thread, llm, DocType::Spec, &transcript, SPEC_PROMPT).await?;
if spec_doc.content.len() > 20 {
if spec_doc.content.chars().count() > 20 {
docs.push(spec_doc);
}
total_tokens.input_tokens += tokens.input_tokens;
@@ -213,7 +213,7 @@ pub async fn reflect_simple(
if thread_succeeded && action_count >= 2 {
let (playbook_doc, tokens) =
produce_doc(thread, llm, DocType::Playbook, &transcript, PLAYBOOK_PROMPT).await?;
if playbook_doc.content.len() > 20 {
if playbook_doc.content.chars().count() > 20 {
docs.push(playbook_doc);
}
total_tokens.input_tokens += tokens.input_tokens;
@@ -263,14 +263,16 @@ fn parse_reflection_output(response: &str, source_thread: &Thread) -> Vec<Memory
}
// Fallback: treat the entire response as a summary
if response.len() > 20 {
vec![MemoryDoc::new(
source_thread.project_id,
DocType::Summary,
format!("Summary: {}", source_thread.goal),
response,
)
.with_source_thread(source_thread.id)]
if response.chars().count() > 20 {
vec![
MemoryDoc::new(
source_thread.project_id,
DocType::Summary,
format!("Summary: {}", source_thread.goal),
response,
)
.with_source_thread(source_thread.id),
]
} else {
vec![]
}
@@ -282,7 +284,7 @@ fn parse_doc_entry(value: &serde_json::Value, source_thread: &Thread) -> Option<
let title = value.get("title")?.as_str()?;
let content = value.get("content")?.as_str()?;
if content.len() <= 20 {
if content.chars().count() <= 20 {
return None;
}
@@ -365,7 +367,11 @@ pub(crate) fn build_transcript(thread: &Thread) -> String {
for msg in messages {
let role = format!("{:?}", msg.role);
let content_preview: String = msg.content.chars().take(500).collect();
let truncated = if msg.content.len() > 500 { "..." } else { "" };
let truncated = if msg.content.chars().count() > 500 {
"..."
} else {
""
};
parts.push(format!("[{role}] {content_preview}{truncated}"));
}
@@ -429,8 +435,8 @@ async fn produce_doc(
DocType::Note => format!("Note: {}", thread.goal),
};
let doc = MemoryDoc::new(thread.project_id, doc_type, title, content)
.with_source_thread(thread.id);
let doc =
MemoryDoc::new(thread.project_id, doc_type, title, content).with_source_thread(thread.id);
Ok((doc, output.usage))
}
@@ -503,8 +509,7 @@ mod tests {
#[tokio::test]
async fn reflect_simple_produces_summary() {
let thread = make_completed_thread();
let llm =
MockLlm::with_responses(vec!["Thread accomplished the test task successfully."]);
let llm = MockLlm::with_responses(vec!["Thread accomplished the test task successfully."]);
let result = reflect_simple(&thread, &llm).await.unwrap();
assert_eq!(result.docs.len(), 1);
@@ -597,10 +602,7 @@ mod tests {
]);
let result = reflect_simple(&thread, &llm).await.unwrap();
assert!(result
.docs
.iter()
.any(|d| d.doc_type == DocType::Playbook));
assert!(result.docs.iter().any(|d| d.doc_type == DocType::Playbook));
}
#[tokio::test]
@@ -619,10 +621,7 @@ mod tests {
let llm = MockLlm::with_responses(vec!["Simple summary."]);
let result = reflect_simple(&thread, &llm).await.unwrap();
assert!(!result
.docs
.iter()
.any(|d| d.doc_type == DocType::Playbook));
assert!(!result.docs.iter().any(|d| d.doc_type == DocType::Playbook));
}
// ── parse_reflection_output tests ──────────────────────────
@@ -663,8 +662,7 @@ mod tests {
#[test]
fn parse_skips_short_content() {
let thread = make_completed_thread();
let json =
r#"{"docs": [{"type": "summary", "title": "test", "content": "too short"}]}"#;
let json = r#"{"docs": [{"type": "summary", "title": "test", "content": "too short"}]}"#;
let docs = parse_reflection_output(json, &thread);
assert!(docs.is_empty());
+2 -6
View File
@@ -165,9 +165,7 @@ mod tests {
tracker
.record_success("tool_b", Duration::from_millis(50))
.await;
tracker
.record_failure("tool_b", "not found")
.await;
tracker.record_failure("tool_b", "not found").await;
let m = tracker.get_metrics("tool_b").await.unwrap();
assert_eq!(m.call_count, 2);
@@ -181,9 +179,7 @@ mod tests {
tracker
.record_success("good_tool", Duration::from_millis(10))
.await;
tracker
.record_failure("bad_tool", "always fails")
.await;
tracker.record_failure("bad_tool", "always fails").await;
let unreliable = tracker.unreliable_actions(0.5).await;
assert_eq!(unreliable.len(), 1);
@@ -13,9 +13,7 @@ use tracing::debug;
use crate::runtime::manager::ThreadManager;
use crate::runtime::messaging::ThreadOutcome;
use crate::types::conversation::{
ConversationEntry, ConversationId, ConversationSurface,
};
use crate::types::conversation::{ConversationEntry, ConversationId, ConversationSurface};
use crate::types::error::EngineError;
use crate::types::message::ThreadMessage;
use crate::types::project::ProjectId;
@@ -44,11 +42,7 @@ impl ConversationManager {
}
/// 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) -> ConversationId {
// Check index first
let key = (channel.to_string(), user_id.to_string());
{
@@ -86,11 +80,9 @@ impl ConversationManager {
thread_config: ThreadConfig,
) -> Result<ThreadId, EngineError> {
let mut convs = self.conversations.write().await;
let conv = convs
.get_mut(&conversation_id)
.ok_or(EngineError::Store {
reason: format!("conversation {conversation_id} not found"),
})?;
let conv = convs.get_mut(&conversation_id).ok_or(EngineError::Store {
reason: format!("conversation {conversation_id} not found"),
})?;
// Record the user entry
conv.add_entry(ConversationEntry::user(content));
@@ -248,12 +240,10 @@ fn build_history_from_entries(
history_entries
.iter()
.filter_map(|entry| match &entry.sender {
EntrySender::User => {
Some(crate::types::message::ThreadMessage::user(&entry.content))
}
EntrySender::Agent { .. } => {
Some(crate::types::message::ThreadMessage::assistant(&entry.content))
}
EntrySender::User => Some(crate::types::message::ThreadMessage::user(&entry.content)),
EntrySender::Agent { .. } => Some(crate::types::message::ThreadMessage::assistant(
&entry.content,
)),
EntrySender::System => None, // skip system notifications
})
.collect()
@@ -267,14 +257,14 @@ mod tests {
use crate::capability::registry::CapabilityRegistry;
use crate::traits::effect::EffectExecutor;
use crate::traits::llm::{LlmBackend, LlmCallConfig, LlmOutput};
use crate::types::conversation::EntrySender;
use crate::traits::store::Store;
use crate::types::capability::{ActionDef, CapabilityLease};
use crate::types::conversation::EntrySender;
use crate::types::event::ThreadEvent;
use crate::types::memory::{DocId, MemoryDoc};
use crate::types::project::Project;
use crate::types::step::{ActionResult, LlmResponse, Step, TokenUsage};
use crate::types::thread::ThreadState;
use crate::traits::store::Store;
use std::sync::Mutex;
use std::time::Duration;
@@ -285,59 +275,155 @@ mod tests {
#[async_trait::async_trait]
impl LlmBackend for MockLlm {
async fn complete(
&self, _: &[ThreadMessage], _: &[ActionDef], _: &LlmCallConfig,
&self,
_: &[ThreadMessage],
_: &[ActionDef],
_: &LlmCallConfig,
) -> Result<LlmOutput, EngineError> {
let mut r = self.0.lock().unwrap();
if r.is_empty() {
Ok(LlmOutput { response: LlmResponse::Text("done".into()), usage: TokenUsage::default() })
Ok(LlmOutput {
response: LlmResponse::Text("done".into()),
usage: TokenUsage::default(),
})
} else {
Ok(r.remove(0))
}
}
fn model_name(&self) -> &str { "mock" }
fn model_name(&self) -> &str {
"mock"
}
}
struct MockEffects;
#[async_trait::async_trait]
impl EffectExecutor for MockEffects {
async fn execute_action(&self, _: &str, _: serde_json::Value, _: &CapabilityLease, _: &crate::traits::effect::ThreadExecutionContext) -> Result<ActionResult, EngineError> {
Ok(ActionResult { call_id: String::new(), action_name: String::new(), output: serde_json::json!({}), is_error: false, duration: Duration::from_millis(1) })
async fn execute_action(
&self,
_: &str,
_: serde_json::Value,
_: &CapabilityLease,
_: &crate::traits::effect::ThreadExecutionContext,
) -> Result<ActionResult, EngineError> {
Ok(ActionResult {
call_id: String::new(),
action_name: String::new(),
output: serde_json::json!({}),
is_error: false,
duration: Duration::from_millis(1),
})
}
async fn available_actions(
&self,
_: &[CapabilityLease],
) -> Result<Vec<ActionDef>, EngineError> {
Ok(vec![])
}
async fn available_actions(&self, _: &[CapabilityLease]) -> Result<Vec<ActionDef>, EngineError> { Ok(vec![]) }
}
struct MockStore;
#[async_trait::async_trait]
impl Store for MockStore {
async fn save_thread(&self, _: &crate::types::thread::Thread) -> Result<(), EngineError> { Ok(()) }
async fn load_thread(&self, _: ThreadId) -> Result<Option<crate::types::thread::Thread>, EngineError> { Ok(None) }
async fn list_threads(&self, _: ProjectId) -> Result<Vec<crate::types::thread::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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> { Ok(()) }
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> { Ok(None) }
async fn list_memory_docs(&self, _: ProjectId) -> Result<Vec<MemoryDoc>, EngineError> { Ok(vec![]) }
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, _: crate::types::capability::LeaseId, _: &str) -> Result<(), EngineError> { Ok(()) }
async fn save_mission(&self, _: &crate::types::mission::Mission) -> Result<(), EngineError> { Ok(()) }
async fn load_mission(&self, _: crate::types::mission::MissionId) -> Result<Option<crate::types::mission::Mission>, EngineError> { Ok(None) }
async fn list_missions(&self, _: ProjectId) -> Result<Vec<crate::types::mission::Mission>, EngineError> { Ok(vec![]) }
async fn update_mission_status(&self, _: crate::types::mission::MissionId, _: crate::types::mission::MissionStatus) -> Result<(), EngineError> { Ok(()) }
async fn save_thread(&self, _: &crate::types::thread::Thread) -> Result<(), EngineError> {
Ok(())
}
async fn load_thread(
&self,
_: ThreadId,
) -> Result<Option<crate::types::thread::Thread>, EngineError> {
Ok(None)
}
async fn list_threads(
&self,
_: ProjectId,
) -> Result<Vec<crate::types::thread::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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> {
Ok(())
}
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> {
Ok(None)
}
async fn list_memory_docs(&self, _: ProjectId) -> Result<Vec<MemoryDoc>, EngineError> {
Ok(vec![])
}
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,
_: crate::types::capability::LeaseId,
_: &str,
) -> Result<(), EngineError> {
Ok(())
}
async fn save_mission(
&self,
_: &crate::types::mission::Mission,
) -> Result<(), EngineError> {
Ok(())
}
async fn load_mission(
&self,
_: crate::types::mission::MissionId,
) -> Result<Option<crate::types::mission::Mission>, EngineError> {
Ok(None)
}
async fn list_missions(
&self,
_: ProjectId,
) -> Result<Vec<crate::types::mission::Mission>, EngineError> {
Ok(vec![])
}
async fn update_mission_status(
&self,
_: crate::types::mission::MissionId,
_: crate::types::mission::MissionStatus,
) -> Result<(), EngineError> {
Ok(())
}
}
fn make_conv_manager() -> (Arc<ThreadManager>, ConversationManager) {
let tm = Arc::new(ThreadManager::new(
Arc::new(MockLlm(Mutex::new(vec![
LlmOutput { response: LlmResponse::Text("Hello!".into()), usage: TokenUsage::default() },
]))),
Arc::new(MockLlm(Mutex::new(vec![LlmOutput {
response: LlmResponse::Text("Hello!".into()),
usage: TokenUsage::default(),
}]))),
Arc::new(MockEffects),
Arc::new(MockStore),
Arc::new(CapabilityRegistry::new()),
+170 -45
View File
@@ -66,7 +66,9 @@ impl ThreadManager {
}
/// Subscribe to thread events for live status updates.
pub fn subscribe_events(&self) -> tokio::sync::broadcast::Receiver<crate::types::event::ThreadEvent> {
pub fn subscribe_events(
&self,
) -> tokio::sync::broadcast::Receiver<crate::types::event::ThreadEvent> {
self.event_tx.subscribe()
}
@@ -155,6 +157,7 @@ impl ThreadManager {
let retrieval = crate::memory::RetrievalEngine::new(store_for_retrieval);
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);
@@ -169,7 +172,8 @@ impl ThreadManager {
debug!(thread_id = %thread_id, "thread execution finished");
// Helper to emit events on both the thread and broadcast channel
let emit = |thread: &mut crate::types::thread::Thread, kind: crate::types::event::EventKind| {
let emit = |thread: &mut crate::types::thread::Thread,
kind: crate::types::event::EventKind| {
let event = crate::types::event::ThreadEvent::new(thread.id, kind);
let _ = event_tx.send(event.clone());
thread.events.push(event);
@@ -193,9 +197,19 @@ impl ThreadManager {
) {
tracing::warn!(thread_id = %thread_id, "failed to transition to Reflecting: {e}");
} else {
emit(&mut exec.thread, crate::types::event::EventKind::ReflectionStarted);
emit(
&mut exec.thread,
crate::types::event::EventKind::ReflectionStarted,
);
match crate::reflection::reflect(&exec.thread, &llm_for_reflection, &store_for_task, &caps_for_reflection).await {
match crate::reflection::reflect(
&exec.thread,
&llm_for_reflection,
&store_for_task,
&caps_for_reflection,
)
.await
{
Ok(reflection) => {
let doc_types: Vec<String> = reflection
.docs
@@ -203,31 +217,48 @@ impl ThreadManager {
.map(|d| format!("{:?}", d.doc_type))
.collect();
emit(&mut exec.thread, crate::types::event::EventKind::ReflectionComplete {
docs_produced: reflection.docs.len(),
doc_types,
tokens_used: reflection.tokens_used.total(),
});
emit(
&mut exec.thread,
crate::types::event::EventKind::ReflectionComplete {
docs_produced: reflection.docs.len(),
doc_types,
tokens_used: reflection.tokens_used.total(),
},
);
// Attach reflection results to the trace
crate::executor::trace::attach_reflection(&mut trace, &reflection);
for doc in &reflection.docs {
let _ = store_for_task.save_memory_doc(doc).await;
if let Err(e) = store_for_task.save_memory_doc(doc).await {
tracing::warn!(
thread_id = %thread_id,
doc_title = %doc.title,
"failed to save reflection doc: {e}"
);
}
}
}
Err(e) => {
emit(&mut exec.thread, crate::types::event::EventKind::ReflectionFailed {
error: e.to_string(),
});
emit(
&mut exec.thread,
crate::types::event::EventKind::ReflectionFailed {
error: e.to_string(),
},
);
}
}
// Transition: Reflecting → Done
let _ = exec.thread.transition_to(
if let Err(e) = exec.thread.transition_to(
crate::types::thread::ThreadState::Done,
Some("reflection finished".into()),
);
) {
tracing::warn!(
thread_id = %thread_id,
"failed to transition to Done after reflection: {e}"
);
}
}
}
@@ -238,7 +269,12 @@ impl ThreadManager {
}
// Save final thread state to store
let _ = store_for_task.save_thread(&exec.thread).await;
if let Err(e) = store_for_task.save_thread(&exec.thread).await {
tracing::warn!(
thread_id = %thread_id,
"failed to save final thread state: {e}"
);
}
result
});
@@ -292,10 +328,7 @@ 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> {
pub async fn join_thread(&self, thread_id: ThreadId) -> Result<ThreadOutcome, EngineError> {
let rt = {
let mut running = self.running.write().await;
running.remove(&thread_id)
@@ -345,13 +378,13 @@ impl ThreadManager {
#[cfg(test)]
mod tests {
use super::*;
use crate::traits::llm::{LlmCallConfig, LlmOutput};
use crate::types::capability::{ActionDef, Capability, CapabilityLease, EffectType};
use crate::types::event::ThreadEvent;
use crate::types::memory::{DocId, MemoryDoc};
use crate::types::project::Project;
use crate::types::step::{ActionResult, LlmResponse, Step, TokenUsage};
use crate::types::thread::ThreadState;
use crate::traits::llm::{LlmCallConfig, LlmOutput};
use std::sync::Mutex;
use std::time::Duration;
@@ -428,26 +461,90 @@ mod tests {
#[async_trait::async_trait]
impl Store for MockStore {
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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> { Ok(()) }
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> { Ok(None) }
async fn list_memory_docs(&self, _: ProjectId) -> Result<Vec<MemoryDoc>, EngineError> { Ok(vec![]) }
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, _: crate::types::capability::LeaseId, _: &str) -> Result<(), EngineError> { Ok(()) }
async fn save_mission(&self, _: &crate::types::mission::Mission) -> Result<(), EngineError> { Ok(()) }
async fn load_mission(&self, _: crate::types::mission::MissionId) -> Result<Option<crate::types::mission::Mission>, EngineError> { Ok(None) }
async fn list_missions(&self, _: ProjectId) -> Result<Vec<crate::types::mission::Mission>, EngineError> { Ok(vec![]) }
async fn update_mission_status(&self, _: crate::types::mission::MissionId, _: crate::types::mission::MissionStatus) -> Result<(), EngineError> { Ok(()) }
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 save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> {
Ok(())
}
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> {
Ok(None)
}
async fn list_memory_docs(&self, _: ProjectId) -> Result<Vec<MemoryDoc>, EngineError> {
Ok(vec![])
}
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,
_: crate::types::capability::LeaseId,
_: &str,
) -> Result<(), EngineError> {
Ok(())
}
async fn save_mission(
&self,
_: &crate::types::mission::Mission,
) -> Result<(), EngineError> {
Ok(())
}
async fn load_mission(
&self,
_: crate::types::mission::MissionId,
) -> Result<Option<crate::types::mission::Mission>, EngineError> {
Ok(None)
}
async fn list_missions(
&self,
_: ProjectId,
) -> Result<Vec<crate::types::mission::Mission>, EngineError> {
Ok(vec![])
}
async fn update_mission_status(
&self,
_: crate::types::mission::MissionId,
_: crate::types::mission::MissionStatus,
) -> Result<(), EngineError> {
Ok(())
}
}
fn make_manager(llm: Arc<dyn LlmBackend>) -> ThreadManager {
@@ -484,7 +581,14 @@ mod tests {
let project = ProjectId::new();
let tid = mgr
.spawn_thread("test", ThreadType::Foreground, project, ThreadConfig::default(), None, "user")
.spawn_thread(
"test",
ThreadType::Foreground,
project,
ThreadConfig::default(),
None,
"user",
)
.await
.unwrap();
@@ -515,7 +619,14 @@ mod tests {
let project = ProjectId::new();
let tid = mgr
.spawn_thread("test", ThreadType::Foreground, project, ThreadConfig::default(), None, "user")
.spawn_thread(
"test",
ThreadType::Foreground,
project,
ThreadConfig::default(),
None,
"user",
)
.await
.unwrap();
@@ -536,12 +647,26 @@ mod tests {
let project = ProjectId::new();
let parent = mgr
.spawn_thread("parent", ThreadType::Foreground, project, ThreadConfig::default(), None, "user")
.spawn_thread(
"parent",
ThreadType::Foreground,
project,
ThreadConfig::default(),
None,
"user",
)
.await
.unwrap();
let child = mgr
.spawn_thread("child", ThreadType::Research, project, ThreadConfig::default(), Some(parent), "user")
.spawn_thread(
"child",
ThreadType::Research,
project,
ThreadConfig::default(),
Some(parent),
"user",
)
.await
.unwrap();
+404 -10
View File
@@ -123,10 +123,7 @@ impl MissionManager {
}
/// List all missions in a project.
pub async fn list_missions(
&self,
project_id: ProjectId,
) -> Result<Vec<Mission>, EngineError> {
pub async fn list_missions(&self, project_id: ProjectId) -> Result<Vec<Mission>, EngineError> {
self.store.list_missions(project_id).await
}
@@ -154,17 +151,13 @@ impl MissionManager {
let should_fire = match &mission.cadence {
MissionCadence::Cron { .. } => {
// Fire if next_fire_at has passed
mission
.next_fire_at
.is_some_and(|next| next <= now)
mission.next_fire_at.is_some_and(|next| next <= now)
}
MissionCadence::Manual => false,
MissionCadence::OnEvent { .. } | MissionCadence::OnPush => false,
};
if should_fire
&& let Some(tid) = self.fire_mission(mid, user_id).await?
{
if should_fire && let Some(tid) = self.fire_mission(mid, user_id).await? {
spawned.push(tid);
}
}
@@ -172,3 +165,404 @@ impl MissionManager {
Ok(spawned)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
use std::sync::Mutex;
use std::time::Duration;
use crate::capability::lease::LeaseManager;
use crate::capability::policy::PolicyEngine;
use crate::capability::registry::CapabilityRegistry;
use crate::traits::effect::EffectExecutor;
use crate::traits::llm::{LlmCallConfig, LlmOutput};
use crate::traits::store::Store;
use crate::types::capability::{ActionDef, CapabilityLease};
use crate::types::error::EngineError;
use crate::types::event::ThreadEvent;
use crate::types::memory::{DocId, MemoryDoc};
use crate::types::mission::{Mission, MissionCadence, MissionId, MissionStatus};
use crate::types::project::{Project, ProjectId};
use crate::types::step::{ActionResult, LlmResponse, Step, TokenUsage};
use crate::types::thread::{Thread, ThreadId, ThreadState};
// ── TestStore — in-memory Store that persists missions ───
struct TestStore {
threads: tokio::sync::RwLock<HashMap<ThreadId, Thread>>,
missions: tokio::sync::RwLock<HashMap<MissionId, Mission>>,
}
impl TestStore {
fn new() -> Self {
Self {
threads: tokio::sync::RwLock::new(HashMap::new()),
missions: tokio::sync::RwLock::new(HashMap::new()),
}
}
}
#[async_trait::async_trait]
impl Store for TestStore {
// ── Thread (minimal — save/load needed by ThreadManager) ──
async fn save_thread(&self, thread: &Thread) -> Result<(), EngineError> {
self.threads.write().await.insert(thread.id, thread.clone());
Ok(())
}
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 update_thread_state(
&self,
_: ThreadId,
_: ThreadState,
) -> Result<(), EngineError> {
Ok(())
}
// ── Step (noop) ──
async fn save_step(&self, _: &Step) -> Result<(), EngineError> {
Ok(())
}
async fn load_steps(&self, _: ThreadId) -> Result<Vec<Step>, EngineError> {
Ok(vec![])
}
// ── Event (noop) ──
async fn append_events(&self, _: &[ThreadEvent]) -> Result<(), EngineError> {
Ok(())
}
async fn load_events(&self, _: ThreadId) -> Result<Vec<ThreadEvent>, EngineError> {
Ok(vec![])
}
// ── Project (noop) ──
async fn save_project(&self, _: &Project) -> Result<(), EngineError> {
Ok(())
}
async fn load_project(&self, _: ProjectId) -> Result<Option<Project>, EngineError> {
Ok(None)
}
// ── MemoryDoc (noop) ──
async fn save_memory_doc(&self, _: &MemoryDoc) -> Result<(), EngineError> {
Ok(())
}
async fn load_memory_doc(&self, _: DocId) -> Result<Option<MemoryDoc>, EngineError> {
Ok(None)
}
async fn list_memory_docs(&self, _: ProjectId) -> Result<Vec<MemoryDoc>, EngineError> {
Ok(vec![])
}
// ── Lease (noop) ──
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,
_: crate::types::capability::LeaseId,
_: &str,
) -> Result<(), EngineError> {
Ok(())
}
// ── Mission (fully implemented) ──
async fn save_mission(&self, mission: &Mission) -> Result<(), EngineError> {
self.missions
.write()
.await
.insert(mission.id, mission.clone());
Ok(())
}
async fn load_mission(&self, id: MissionId) -> Result<Option<Mission>, EngineError> {
Ok(self.missions.read().await.get(&id).cloned())
}
async fn list_missions(&self, project_id: ProjectId) -> Result<Vec<Mission>, EngineError> {
Ok(self
.missions
.read()
.await
.values()
.filter(|m| m.project_id == project_id)
.cloned()
.collect())
}
async fn update_mission_status(
&self,
id: MissionId,
status: MissionStatus,
) -> Result<(), EngineError> {
if let Some(mission) = self.missions.write().await.get_mut(&id) {
mission.status = status;
}
Ok(())
}
}
// ── MockLlm — returns canned text responses ─────────────
struct MockLlm {
responses: Mutex<Vec<LlmOutput>>,
}
impl MockLlm {
fn text(msg: &str) -> Arc<Self> {
Arc::new(Self {
responses: Mutex::new(vec![LlmOutput {
response: LlmResponse::Text(msg.into()),
usage: TokenUsage::default(),
}]),
})
}
}
#[async_trait::async_trait]
impl crate::traits::llm::LlmBackend for MockLlm {
async fn complete(
&self,
_: &[crate::types::message::ThreadMessage],
_: &[ActionDef],
_: &LlmCallConfig,
) -> Result<LlmOutput, EngineError> {
let mut r = self.responses.lock().unwrap();
if r.is_empty() {
Ok(LlmOutput {
response: LlmResponse::Text("done".into()),
usage: TokenUsage::default(),
})
} else {
Ok(r.remove(0))
}
}
fn model_name(&self) -> &str {
"mock"
}
}
// ── MockEffects — noop effect executor ───────────────────
struct MockEffects;
#[async_trait::async_trait]
impl EffectExecutor for MockEffects {
async fn execute_action(
&self,
_: &str,
_: serde_json::Value,
_: &CapabilityLease,
_: &crate::traits::effect::ThreadExecutionContext,
) -> Result<ActionResult, EngineError> {
Ok(ActionResult {
call_id: String::new(),
action_name: String::new(),
output: serde_json::json!({}),
is_error: false,
duration: Duration::from_millis(1),
})
}
async fn available_actions(
&self,
_: &[CapabilityLease],
) -> Result<Vec<ActionDef>, EngineError> {
Ok(vec![])
}
}
// ── Helper to build a MissionManager with its dependencies ──
fn make_mission_manager(store: Arc<dyn Store>) -> MissionManager {
let caps = CapabilityRegistry::new();
let thread_manager = Arc::new(ThreadManager::new(
MockLlm::text("done"),
Arc::new(MockEffects),
Arc::clone(&store),
Arc::new(caps),
Arc::new(LeaseManager::new()),
Arc::new(PolicyEngine::new()),
));
MissionManager::new(store, thread_manager)
}
// ── Tests ───────────────────────────────────────────────
#[tokio::test]
async fn create_mission_persists() {
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
.create_mission(
project_id,
"test mission",
"do the thing",
MissionCadence::Manual,
)
.await
.unwrap();
let mission = mgr.get_mission(id).await.unwrap();
assert!(mission.is_some());
let mission = mission.unwrap();
assert_eq!(mission.name, "test mission");
assert_eq!(mission.goal, "do the thing");
assert_eq!(mission.status, MissionStatus::Active);
assert_eq!(mission.project_id, project_id);
}
#[tokio::test]
async fn pause_and_resume() {
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
.create_mission(project_id, "pausable", "goal", MissionCadence::Manual)
.await
.unwrap();
// Pause
mgr.pause_mission(id).await.unwrap();
let mission = mgr.get_mission(id).await.unwrap().unwrap();
assert_eq!(mission.status, MissionStatus::Paused);
// Resume
mgr.resume_mission(id).await.unwrap();
let mission = mgr.get_mission(id).await.unwrap().unwrap();
assert_eq!(mission.status, MissionStatus::Active);
}
#[tokio::test]
async fn complete_removes_from_active() {
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
.create_mission(project_id, "completable", "goal", MissionCadence::Manual)
.await
.unwrap();
mgr.complete_mission(id).await.unwrap();
let mission = mgr.get_mission(id).await.unwrap().unwrap();
assert_eq!(mission.status, MissionStatus::Completed);
assert!(mission.is_terminal());
// Verify removed from active list
let active = mgr.active.read().await;
assert!(!active.contains(&id));
}
#[tokio::test]
async fn fire_mission_spawns_thread() {
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
.create_mission(
project_id,
"fireable",
"build something",
MissionCadence::Manual,
)
.await
.unwrap();
let thread_id = mgr.fire_mission(id, "test-user").await.unwrap();
assert!(
thread_id.is_some(),
"fire_mission should return a thread ID"
);
let tid = thread_id.unwrap();
// Give the spawned thread a moment to finish (MockLlm returns immediately)
tokio::time::sleep(Duration::from_millis(50)).await;
// Verify the thread was recorded in mission history
let mission = mgr.get_mission(id).await.unwrap().unwrap();
assert!(
mission.thread_history.contains(&tid),
"thread should be recorded in mission history"
);
}
#[tokio::test]
async fn fire_terminal_mission_returns_none() {
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
.create_mission(project_id, "terminal", "goal", MissionCadence::Manual)
.await
.unwrap();
// Complete the mission so it becomes terminal
mgr.complete_mission(id).await.unwrap();
let result = mgr.fire_mission(id, "test-user").await.unwrap();
assert!(
result.is_none(),
"firing a terminal mission should return None"
);
}
#[tokio::test]
async fn tick_fires_due_missions() {
let store = Arc::new(TestStore::new());
let mgr = make_mission_manager(Arc::clone(&store) as Arc<dyn Store>);
let project_id = ProjectId::new();
// Create a cron mission with next_fire_at in the past
let id = mgr
.create_mission(
project_id,
"cron mission",
"periodic goal",
MissionCadence::Cron {
expression: "* * * * *".into(),
},
)
.await
.unwrap();
// Set next_fire_at to the past so tick() will fire it
{
let mut missions = store.missions.write().await;
if let Some(mission) = missions.get_mut(&id) {
mission.next_fire_at = Some(chrono::Utc::now() - chrono::Duration::seconds(60));
}
}
let spawned = mgr.tick("test-user").await.unwrap();
assert_eq!(spawned.len(), 1, "tick should fire exactly one due mission");
// Give the spawned thread a moment to finish
tokio::time::sleep(Duration::from_millis(50)).await;
// Verify the thread was recorded
let mission = mgr.get_mission(id).await.unwrap().unwrap();
assert!(
mission.thread_history.contains(&spawned[0]),
"spawned thread should be recorded in mission history"
);
}
}
@@ -210,10 +210,7 @@ mod tests {
// Thread starts
let tid = ThreadId::new();
conv.track_thread(tid);
conv.add_entry(ConversationEntry::system_for_thread(
tid,
"Thread started",
));
conv.add_entry(ConversationEntry::system_for_thread(tid, "Thread started"));
assert_eq!(conv.active_threads.len(), 1);
// Agent responds
@@ -222,7 +219,10 @@ mod tests {
// Thread completes
conv.untrack_thread(tid);
conv.add_entry(ConversationEntry::system_for_thread(tid, "Thread completed"));
conv.add_entry(ConversationEntry::system_for_thread(
tid,
"Thread completed",
));
assert!(conv.active_threads.is_empty());
assert_eq!(conv.entries.len(), 4);
}
+1 -4
View File
@@ -75,10 +75,7 @@ impl ThreadMessage {
}
/// Create an assistant message with action calls.
pub fn assistant_with_actions(
content: Option<String>,
calls: Vec<ActionCall>,
) -> Self {
pub fn assistant_with_actions(content: Option<String>, calls: Vec<ActionCall>) -> Self {
Self {
role: MessageRole::Assistant,
content: content.unwrap_or_default(),
@@ -26,4 +26,3 @@ pub enum Provenance {
/// Retrieved from project memory.
MemoryRetrieval { doc_id: DocId },
}
@@ -132,7 +132,6 @@ pub struct ThreadConfig {
pub max_tool_intent_nudges: u32,
// ── Budget controls (Phase 4, from RLM cross-reference) ──
/// Maximum cumulative input+output tokens before termination.
pub max_tokens_total: Option<u64>,
/// Maximum consecutive steps with errors before termination.