feat(engine): Phase 4 — budget controls, compaction, reflection pipeline

Budget enforcement in ExecutionLoop:
- max_tokens_total: cumulative token limit, checked before each iteration
- max_duration: wall-clock timeout for entire thread
- max_consecutive_errors: consecutive error steps threshold (resets on
  success, matching official RLM behavior)
- All produce ThreadOutcome::Failed with descriptive messages

Context compaction (from RLM paper, 85% threshold):
- estimate_tokens(): char-based estimation (chars/4, matching RLM)
- should_compact(): triggers when tokens >= threshold_pct * context_limit
- compact_messages(): asks LLM to summarize progress, replaces history
  with [system, summary, continuation_note], preserves intermediate results
- Configurable via ThreadConfig: model_context_limit, compaction_threshold

Dual model routing:
- LlmCallConfig gains depth field (0=root, 1+=sub-call)
- Implementations can route to cheaper models for sub-calls
- ExecutionLoop passes thread depth to every LLM call

Reflection pipeline (reflection/pipeline.rs):
- reflect(thread, llm): analyzes completed thread via LLM
- Produces Summary doc (always), Lesson doc (if errors), Issue doc (if failed)
- Builds transcript from thread messages + error events
- Returns ReflectionResult with docs + token usage

ThreadConfig extended with: max_tokens_total, max_consecutive_errors,
model_context_limit, enable_compaction, compaction_threshold, depth, max_depth.

78 tests passing, zero clippy warnings.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
2026-03-21 22:55:34 -07:00
co-authored by Claude Opus 4.6
parent ff1107179a
commit 4bc7ffdf0c
9 changed files with 547 additions and 9 deletions
@@ -0,0 +1,178 @@
//! Context compaction and token counting.
//!
//! When message history approaches the model's context limit, compaction
//! asks the LLM to summarize progress and resets the history. This follows
//! the official RLM pattern (compaction at 85% of context limit).
use std::sync::Arc;
use tracing::debug;
use crate::traits::llm::{LlmBackend, LlmCallConfig};
use crate::types::error::EngineError;
use crate::types::message::{MessageRole, ThreadMessage};
use crate::types::step::{LlmResponse, TokenUsage};
/// Characters per token estimate when no tokenizer is available.
/// Conservative estimate (official RLM uses 4).
const CHARS_PER_TOKEN: usize = 4;
/// Estimate token count for a list of messages.
///
/// Uses character length / `CHARS_PER_TOKEN` as a rough estimate.
/// The official RLM uses tiktoken when available; we use this fallback
/// since we don't depend on a Python tokenizer.
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)
})
.sum();
total_chars.div_ceil(CHARS_PER_TOKEN)
}
/// Check if compaction should be triggered.
///
/// Returns `true` when estimated token count exceeds `threshold_pct` of
/// the model's context limit.
pub fn should_compact(
messages: &[ThreadMessage],
model_context_limit: usize,
threshold_pct: f64,
) -> bool {
let tokens = estimate_tokens(messages);
let threshold = (model_context_limit as f64 * threshold_pct) as usize;
tokens >= threshold
}
/// The compaction prompt sent to the LLM.
const COMPACTION_PROMPT: &str = "\
Summarize your progress so far in a concise but complete way. Include:
1. What you have accomplished
2. Key intermediate results and variable values
3. What still needs to be done
4. Any errors encountered and how they were handled
Preserve all information needed to continue the task. Be specific about data values.";
/// Compact the message history by asking the LLM to summarize.
///
/// Returns the new (shorter) message list and the token usage from the
/// summarization call. The original messages are replaced with:
/// `[system_prompt, summary, continuation_note]`
///
/// The full original messages are returned separately so the caller can
/// store them (e.g., in a `history` variable or event log).
pub async fn compact_messages(
messages: &[ThreadMessage],
llm: &Arc<dyn LlmBackend>,
compaction_count: u32,
) -> Result<CompactionResult, EngineError> {
// Build a summarization request from existing messages + prompt
let mut summarize_messages = messages.to_vec();
summarize_messages.push(ThreadMessage::user(COMPACTION_PROMPT.to_string()));
let config = LlmCallConfig {
force_text: true,
..LlmCallConfig::default()
};
let output = llm.complete(&summarize_messages, &[], &config).await?;
let summary_text = match output.response {
LlmResponse::Text(t) => t,
LlmResponse::ActionCalls { content, .. } | LlmResponse::Code { content, .. } => {
content.unwrap_or_else(|| "[compaction produced no summary]".into())
}
};
// Preserve the system prompt (first message if it's a system message)
let system_msg = messages
.iter()
.find(|m| m.role == MessageRole::System)
.cloned();
// Build compacted history
let mut compacted = Vec::new();
if let Some(sys) = system_msg {
compacted.push(sys);
}
compacted.push(ThreadMessage::assistant(summary_text.clone()));
compacted.push(ThreadMessage::system(format!(
"Your conversation has been compacted {n} time(s). \
The summary above captures your progress. Continue working on the task.",
n = compaction_count + 1,
)));
let tokens_before = estimate_tokens(messages);
let tokens_after = estimate_tokens(&compacted);
debug!(
tokens_before,
tokens_after,
compaction_count = compaction_count + 1,
"context compacted"
);
Ok(CompactionResult {
compacted_messages: compacted,
summary: summary_text,
tokens_used: output.usage,
tokens_before,
tokens_after,
})
}
/// Result of a compaction operation.
pub struct CompactionResult {
/// The new (shorter) message list.
pub compacted_messages: Vec<ThreadMessage>,
/// The summary text produced by the LLM.
pub summary: String,
/// Tokens used by the summarization LLM call.
pub tokens_used: TokenUsage,
/// Estimated token count before compaction.
pub tokens_before: usize,
/// Estimated token count after compaction.
pub tokens_after: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn estimate_tokens_empty() {
assert_eq!(estimate_tokens(&[]), 0);
}
#[test]
fn estimate_tokens_basic() {
let msgs = vec![
ThreadMessage::system("Hello world"), // 11 chars + 4 overhead = 15 / 4 = 3.75
ThreadMessage::user("Hi"), // 2 chars + 4 = 6 / 4 = 1.5
];
let tokens = estimate_tokens(&msgs);
// (11+4 + 2+4) / 4 = 21/4 = 5.25 → 6 (ceiling)
assert!(tokens > 0);
assert!(tokens < 100);
}
#[test]
fn should_compact_below_threshold() {
let msgs = vec![ThreadMessage::user("short message")];
assert!(!should_compact(&msgs, 128_000, 0.85));
}
#[test]
fn should_compact_above_threshold() {
// Create a message large enough to trigger compaction at low limit
let big = "x".repeat(1000);
let msgs = vec![ThreadMessage::user(big)];
// 1000 chars / 4 = 250 tokens. Context limit 200, threshold 85% = 170
assert!(should_compact(&msgs, 200, 0.85));
}
}
@@ -64,7 +64,10 @@ impl ExecutionLoop {
let max_iterations = self.thread.config.max_iterations;
let max_nudges = self.thread.config.max_tool_intent_nudges;
let nudge_enabled = self.thread.config.enable_tool_intent_nudge;
let start_time = std::time::Instant::now();
let mut nudge_count: u32 = 0;
let mut consecutive_errors: u32 = 0;
let mut compaction_count: u32 = 0;
for iteration in 0..max_iterations {
// 1. Check signals
@@ -80,24 +83,90 @@ impl ExecutionLoop {
}
}
// 2. Get active leases
// 2. Check budget limits
if let Some(max_tokens) = self.thread.config.max_tokens_total
&& self.thread.total_tokens_used >= max_tokens
{
warn!(
thread_id = %self.thread.id,
used = self.thread.total_tokens_used,
limit = max_tokens,
"token limit exceeded"
);
self.thread.transition_to(
ThreadState::Completed,
Some("token limit exceeded".into()),
)?;
return Ok(ThreadOutcome::Failed {
error: format!(
"Token limit exceeded: {} of {} tokens",
self.thread.total_tokens_used, max_tokens
),
});
}
if let Some(max_dur) = self.thread.config.max_duration {
let elapsed = start_time.elapsed();
if elapsed >= max_dur {
warn!(
thread_id = %self.thread.id,
elapsed = ?elapsed,
limit = ?max_dur,
"thread timeout"
);
self.thread
.transition_to(ThreadState::Completed, Some("timeout".into()))?;
return Ok(ThreadOutcome::Failed {
error: format!("Thread timeout: {elapsed:?} of {max_dur:?}"),
});
}
}
// 3. Check compaction
if self.thread.config.enable_compaction {
let ctx_limit = self.thread.config.model_context_limit;
let threshold = self.thread.config.compaction_threshold;
if crate::executor::compaction::should_compact(
&self.thread.messages,
ctx_limit,
threshold,
) {
debug!(
thread_id = %self.thread.id,
compaction_count,
"triggering context compaction"
);
let result = crate::executor::compaction::compact_messages(
&self.thread.messages,
&self.llm,
compaction_count,
)
.await?;
self.thread.total_tokens_used += result.tokens_used.total();
self.thread.messages = result.compacted_messages;
compaction_count += 1;
}
}
// 4. Get active leases
let active_leases = self.leases.active_for_thread(self.thread.id).await;
// 3. Build context
// 5. Build context
let (messages, actions) =
build_step_context(&self.thread.messages, &active_leases, &self.effects).await?;
// 4. Create step
// 6. Create step
let mut step = Step::new(self.thread.id, iteration + 1);
step.status = StepStatus::LlmCalling;
self.thread.add_event(EventKind::StepStarted {
step_id: step.id,
});
// 5. Call LLM
// 7. Call LLM
let force_text = iteration >= max_iterations.saturating_sub(1);
let config = LlmCallConfig {
force_text,
depth: self.thread.config.depth,
..LlmCallConfig::default()
};
@@ -308,10 +377,37 @@ impl ExecutionLoop {
return Ok(outcome);
}
// If code had errors, the error text is already in the
// metadata message — the LLM can self-correct on next turn.
// Track consecutive errors for budget enforcement
if code_result.had_error {
consecutive_errors += 1;
} else {
consecutive_errors = 0;
}
}
}
// Check consecutive error threshold after each step
if let Some(max_errors) = self.thread.config.max_consecutive_errors
&& consecutive_errors >= max_errors
{
warn!(
thread_id = %self.thread.id,
consecutive_errors,
max_errors,
"consecutive error threshold exceeded"
);
self.thread.transition_to(
ThreadState::Failed,
Some(format!(
"consecutive error threshold: {consecutive_errors} errors"
)),
)?;
return Ok(ThreadOutcome::Failed {
error: format!(
"Consecutive error threshold exceeded: {consecutive_errors} of {max_errors}"
),
});
}
}
// Max iterations reached
@@ -5,6 +5,7 @@
//! - [`context`] — context building for LLM calls
//! - [`intent`] — tool intent nudge detection
pub mod compaction;
pub mod context;
pub mod intent;
pub mod loop_engine;
+4
View File
@@ -65,3 +65,7 @@ pub use executor::ExecutionLoop;
pub use memory::MemoryStore;
pub use memory::RetrievalEngine;
// ── Re-exports: reflection ────────────────────────────────────
pub use reflection::ReflectionResult;
+9 -3
View File
@@ -1,5 +1,11 @@
//! Post-thread reflection pipeline.
//!
//! After a thread completes, the reflection pipeline produces structured
//! knowledge (summaries, lessons, playbooks, issue docs) from the thread's
//! execution trace. Implemented in Phase 4.
//! After a thread completes, [`reflect()`] uses the LLM to produce structured
//! knowledge (MemoryDocs) from the thread's execution trace:
//! - Summary — what the thread accomplished
//! - Lesson — what was learned from errors/workarounds
//! - Issue — unresolved problems for follow-up
pub mod pipeline;
pub use pipeline::{reflect, ReflectionResult};
@@ -0,0 +1,210 @@
//! Reflection pipeline — produces structured knowledge from completed threads.
//!
//! After a thread completes, the reflection pipeline uses the LLM to:
//! 1. Summarize what the thread accomplished
//! 2. Extract lessons from failures and workarounds
//! 3. Detect unresolved issues
//! 4. Identify missing capabilities
//!
//! Each produces a MemoryDoc stored in the thread's project scope.
use std::sync::Arc;
use tracing::debug;
use crate::traits::llm::{LlmBackend, LlmCallConfig};
use crate::types::error::EngineError;
use crate::types::event::EventKind;
use crate::types::memory::{DocType, MemoryDoc};
use crate::types::message::ThreadMessage;
use crate::types::step::{LlmResponse, TokenUsage};
use crate::types::thread::Thread;
/// Result of running the reflection pipeline on a completed thread.
pub struct ReflectionResult {
/// Memory docs produced by reflection.
pub docs: Vec<MemoryDoc>,
/// Total tokens used by reflection LLM calls.
pub tokens_used: TokenUsage,
}
/// Run the reflection pipeline on a completed thread.
///
/// Produces structured knowledge (MemoryDocs) from the thread's messages
/// and events. Uses the LLM for summarization and analysis.
pub async fn reflect(
thread: &Thread,
llm: &Arc<dyn LlmBackend>,
) -> Result<ReflectionResult, EngineError> {
let mut docs = Vec::new();
let mut total_tokens = TokenUsage::default();
// Build a transcript of the thread's work for the LLM to analyze
let transcript = build_transcript(thread);
// 1. Summary doc
let (summary_doc, tokens) =
produce_doc(thread, llm, DocType::Summary, &transcript, SUMMARY_PROMPT).await?;
docs.push(summary_doc);
total_tokens.input_tokens += tokens.input_tokens;
total_tokens.output_tokens += tokens.output_tokens;
// 2. Lessons (only if there were errors)
let had_errors = thread.events.iter().any(|e| {
matches!(
e.kind,
EventKind::ActionFailed { .. } | EventKind::StepFailed { .. }
)
});
if had_errors {
let (lesson_doc, tokens) =
produce_doc(thread, llm, DocType::Lesson, &transcript, LESSON_PROMPT).await?;
docs.push(lesson_doc);
total_tokens.input_tokens += tokens.input_tokens;
total_tokens.output_tokens += tokens.output_tokens;
}
// 3. Issues (if thread failed or had unresolved problems)
let thread_failed = thread.state == crate::types::thread::ThreadState::Failed;
if thread_failed || had_errors {
let (issue_doc, tokens) =
produce_doc(thread, llm, DocType::Issue, &transcript, ISSUE_PROMPT).await?;
// Only add if the LLM produced non-trivial content
if issue_doc.content.len() > 20 {
docs.push(issue_doc);
}
total_tokens.input_tokens += tokens.input_tokens;
total_tokens.output_tokens += tokens.output_tokens;
}
debug!(
thread_id = %thread.id,
docs_produced = docs.len(),
total_tokens = total_tokens.total(),
"reflection complete"
);
Ok(ReflectionResult {
docs,
tokens_used: total_tokens,
})
}
// ── Prompts ─────────────────────────────────────────────────
const SUMMARY_PROMPT: &str = "\
Summarize what this thread accomplished in 2-4 sentences. Include:
- The goal and whether it was achieved
- Key results or outputs
- Tools/actions that were used
Be factual and concise.";
const LESSON_PROMPT: &str = "\
Extract lessons learned from this thread's execution. Focus on:
- Errors encountered and how they were resolved (or not)
- Workarounds that were discovered
- Surprising findings about tool behavior
- Patterns that could be reused in similar tasks
Write each lesson as a single clear sentence. If there are no meaningful lessons, write 'No lessons.'.";
const ISSUE_PROMPT: &str = "\
Identify any unresolved issues from this thread. Focus on:
- Errors that were not resolved
- Tasks that could not be completed
- Missing tools or capabilities that were needed
- Data quality issues encountered
If there are no unresolved issues, write 'No issues.'.";
// ── Helpers ─────────────────────────────────────────────────
/// Build a concise transcript of the thread's work.
fn build_transcript(thread: &Thread) -> String {
let mut parts = Vec::new();
parts.push(format!("Goal: {}", thread.goal));
parts.push(format!("Steps: {}", thread.step_count));
parts.push(format!("Tokens used: {}", thread.total_tokens_used));
parts.push(format!("State: {:?}", thread.state));
// Include messages (truncated for very long threads)
let max_messages = 30;
let messages = if thread.messages.len() > max_messages {
&thread.messages[thread.messages.len() - max_messages..]
} else {
&thread.messages
};
parts.push("\n--- Messages ---".into());
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 { "" };
parts.push(format!("[{role}] {content_preview}{truncated}"));
}
// Include notable events
let error_events: Vec<String> = thread
.events
.iter()
.filter_map(|e| match &e.kind {
EventKind::ActionFailed { action_name, error, .. } => {
Some(format!("Action '{action_name}' failed: {error}"))
}
EventKind::StepFailed { error, .. } => Some(format!("Step failed: {error}")),
_ => None,
})
.collect();
if !error_events.is_empty() {
parts.push("\n--- Errors ---".into());
for err in error_events {
parts.push(err);
}
}
parts.join("\n")
}
/// Produce a single MemoryDoc by asking the LLM to analyze the transcript.
async fn produce_doc(
thread: &Thread,
llm: &Arc<dyn LlmBackend>,
doc_type: DocType,
transcript: &str,
prompt: &str,
) -> Result<(MemoryDoc, TokenUsage), EngineError> {
let messages = vec![
ThreadMessage::system(format!(
"You are analyzing a completed agent thread. Here is the transcript:\n\n{transcript}"
)),
ThreadMessage::user(prompt.to_string()),
];
let config = LlmCallConfig {
force_text: true,
..LlmCallConfig::default()
};
let output = llm.complete(&messages, &[], &config).await?;
let content = match output.response {
LlmResponse::Text(t) => t,
LlmResponse::ActionCalls { content, .. } | LlmResponse::Code { content, .. } => {
content.unwrap_or_default()
}
};
let title = match doc_type {
DocType::Summary => format!("Summary: {}", thread.goal),
DocType::Lesson => format!("Lessons: {}", thread.goal),
DocType::Issue => format!("Issues: {}", thread.goal),
DocType::Playbook => format!("Playbook: {}", thread.goal),
DocType::Spec => format!("Spec: {}", thread.goal),
DocType::Note => format!("Note: {}", thread.goal),
};
let doc = MemoryDoc::new(thread.project_id, doc_type, title, content)
.with_source_thread(thread.id);
Ok((doc, output.usage))
}
+3
View File
@@ -21,6 +21,9 @@ pub struct LlmCallConfig {
pub temperature: Option<f32>,
/// When true, the LLM should not return action calls.
pub force_text: bool,
/// Depth in the recursive call tree (0 = root, 1+ = sub-call).
/// Implementations can use this to route to cheaper models for sub-calls.
pub depth: u32,
/// Opaque metadata forwarded to the LLM provider.
pub metadata: HashMap<String, String>,
}
+12
View File
@@ -43,6 +43,18 @@ pub enum EngineError {
#[error("max iterations reached: {limit}")]
MaxIterations { limit: usize },
#[error("token limit exceeded: {used} of {limit}")]
TokenLimitExceeded { used: u64, limit: u64 },
#[error("consecutive error threshold exceeded: {count} errors (limit: {threshold})")]
ConsecutiveErrors { count: u32, threshold: u32 },
#[error("thread timeout: {elapsed:?} of {limit:?}")]
Timeout {
elapsed: std::time::Duration,
limit: std::time::Duration,
},
}
use crate::types::project::ProjectId;
@@ -130,6 +130,27 @@ pub struct ThreadConfig {
pub enable_tool_intent_nudge: bool,
/// Maximum number of tool intent nudges per thread.
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.
/// Resets to 0 on any successful step (matching official RLM behavior).
pub max_consecutive_errors: Option<u32>,
/// Model context limit in tokens (for compaction threshold calculation).
/// Default: 128,000. Used to trigger compaction at 85% usage.
pub model_context_limit: usize,
/// Whether to enable automatic compaction when context grows large.
pub enable_compaction: bool,
/// Compaction threshold as fraction of model_context_limit (0.0-1.0).
/// Default: 0.85 (matching official RLM).
pub compaction_threshold: f64,
/// Depth of this thread in the recursive call tree.
/// Root threads are depth 0. Sub-calls via rlm_query() increment depth.
pub depth: u32,
/// Maximum recursion depth for rlm_query() sub-calls.
pub max_depth: u32,
}
impl Default for ThreadConfig {
@@ -140,6 +161,13 @@ impl Default for ThreadConfig {
enable_reflection: false,
enable_tool_intent_nudge: true,
max_tool_intent_nudges: 2,
max_tokens_total: None,
max_consecutive_errors: None,
model_context_limit: 128_000,
enable_compaction: false,
compaction_threshold: 0.85,
depth: 0,
max_depth: 1,
}
}
}