diff --git a/crates/ironclaw_engine/src/executor/compaction.rs b/crates/ironclaw_engine/src/executor/compaction.rs new file mode 100644 index 00000000..1985bc06 --- /dev/null +++ b/crates/ironclaw_engine/src/executor/compaction.rs @@ -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, + compaction_count: u32, +) -> Result { + // 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, + /// 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)); + } +} diff --git a/crates/ironclaw_engine/src/executor/loop_engine.rs b/crates/ironclaw_engine/src/executor/loop_engine.rs index b8e766dd..bdb420b4 100644 --- a/crates/ironclaw_engine/src/executor/loop_engine.rs +++ b/crates/ironclaw_engine/src/executor/loop_engine.rs @@ -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 diff --git a/crates/ironclaw_engine/src/executor/mod.rs b/crates/ironclaw_engine/src/executor/mod.rs index 43b7e3c7..6f9592b6 100644 --- a/crates/ironclaw_engine/src/executor/mod.rs +++ b/crates/ironclaw_engine/src/executor/mod.rs @@ -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; diff --git a/crates/ironclaw_engine/src/lib.rs b/crates/ironclaw_engine/src/lib.rs index 8fee3666..9281574d 100644 --- a/crates/ironclaw_engine/src/lib.rs +++ b/crates/ironclaw_engine/src/lib.rs @@ -65,3 +65,7 @@ pub use executor::ExecutionLoop; pub use memory::MemoryStore; pub use memory::RetrievalEngine; + +// ── Re-exports: reflection ──────────────────────────────────── + +pub use reflection::ReflectionResult; diff --git a/crates/ironclaw_engine/src/reflection/mod.rs b/crates/ironclaw_engine/src/reflection/mod.rs index dbfa6ab8..f125b951 100644 --- a/crates/ironclaw_engine/src/reflection/mod.rs +++ b/crates/ironclaw_engine/src/reflection/mod.rs @@ -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}; diff --git a/crates/ironclaw_engine/src/reflection/pipeline.rs b/crates/ironclaw_engine/src/reflection/pipeline.rs new file mode 100644 index 00000000..0850fc0b --- /dev/null +++ b/crates/ironclaw_engine/src/reflection/pipeline.rs @@ -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, + /// 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, +) -> Result { + 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 = 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, + 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)) +} diff --git a/crates/ironclaw_engine/src/traits/llm.rs b/crates/ironclaw_engine/src/traits/llm.rs index 04654224..43a3bce2 100644 --- a/crates/ironclaw_engine/src/traits/llm.rs +++ b/crates/ironclaw_engine/src/traits/llm.rs @@ -21,6 +21,9 @@ pub struct LlmCallConfig { pub temperature: Option, /// 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, } diff --git a/crates/ironclaw_engine/src/types/error.rs b/crates/ironclaw_engine/src/types/error.rs index 24030eb4..94752ec0 100644 --- a/crates/ironclaw_engine/src/types/error.rs +++ b/crates/ironclaw_engine/src/types/error.rs @@ -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; diff --git a/crates/ironclaw_engine/src/types/thread.rs b/crates/ironclaw_engine/src/types/thread.rs index 8e5cf481..6c5b2c37 100644 --- a/crates/ironclaw_engine/src/types/thread.rs +++ b/crates/ironclaw_engine/src/types/thread.rs @@ -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, + /// Maximum consecutive steps with errors before termination. + /// Resets to 0 on any successful step (matching official RLM behavior). + pub max_consecutive_errors: Option, + /// 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, } } }