diff --git a/src/harness/middleware/library/context.rs b/src/harness/middleware/library/context.rs index ff8c3d0..300073a 100644 --- a/src/harness/middleware/library/context.rs +++ b/src/harness/middleware/library/context.rs @@ -7,8 +7,9 @@ use super::*; use crate::harness::cache::{CacheLayoutEvent, PromptCacheLayout}; use crate::harness::middleware::{ - ContextCompressionMiddleware, DEFAULT_CACHE_GUARD_EVENT_CAP, DEFAULT_COMPRESSION_RECORD_CAP, - MessageTrimMiddleware, MicrocompactMiddleware, PromptCacheGuardMiddleware, + CompressionFailurePolicy, ContextCompressionMiddleware, DEFAULT_CACHE_GUARD_EVENT_CAP, + DEFAULT_COMPRESSION_RECORD_CAP, MessageTrimMiddleware, MicrocompactMiddleware, + PromptCacheGuardMiddleware, }; use crate::harness::summarization::{ ConcatSummarizer, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy, @@ -57,6 +58,10 @@ impl ContextCompressionMiddleware { } /// Creates a compression middleware with a custom [`Summarizer`]. + /// + /// The failure policy defaults to + /// [`CompressionFailurePolicy::FallbackTrim`]; override it with + /// [`with_failure_policy`](Self::with_failure_policy). pub fn with_summarizer(policy: SummarizationPolicy, summarizer: Box) -> Self { Self { label: "context_compression", @@ -64,9 +69,23 @@ impl ContextCompressionMiddleware { summarizer, records: std::sync::Mutex::new(std::collections::VecDeque::new()), max_records: DEFAULT_COMPRESSION_RECORD_CAP, + on_failure: CompressionFailurePolicy::default(), } } + /// Sets the [`CompressionFailurePolicy`] applied when the [`Summarizer`] + /// returns an `Err`. Defaults to + /// [`CompressionFailurePolicy::FallbackTrim`]. + pub fn with_failure_policy(mut self, on_failure: CompressionFailurePolicy) -> Self { + self.on_failure = on_failure; + self + } + + /// Returns the configured [`CompressionFailurePolicy`]. + pub fn failure_policy(&self) -> CompressionFailurePolicy { + self.on_failure + } + /// Sets the maximum number of [`SummaryRecord`]s retained before the /// oldest is evicted. `0` disables recording entirely. pub fn with_max_records(mut self, max_records: usize) -> Self { @@ -123,7 +142,40 @@ impl Middleware for ContextCom } let from_tokens = total_message_tokens(&request.messages); - let record = self.summarizer.summarize(&to_summarize).await?; + let record = match self.summarizer.summarize(&to_summarize).await { + Ok(record) => record, + Err(err) => { + // A summarizer failure hits precisely the longest, most valuable + // transcripts (the ones that reached the compaction threshold). + // Emit a diagnostic and recover per the configured + // `CompressionFailurePolicy` instead of aborting the whole run. + ctx.emit(AgentEvent::MiddlewareFailed { + name: self.label.to_string(), + error: err.to_string(), + }); + match self.on_failure { + // Legacy behaviour: propagate and let the run fail. + CompressionFailurePolicy::Abort => return Err(err), + // Keep the (over-threshold) transcript verbatim and continue. + CompressionFailurePolicy::PassThrough => return Ok(()), + // Deterministic front-drop to the policy's trigger budget, + // preserving system messages (see `TrimStrategy::MaxTokens`). + CompressionFailurePolicy::FallbackTrim => { + let trimmed = trim_messages( + &request.messages, + &TrimStrategy::MaxTokens(self.policy.trigger_budget()), + ); + let to_tokens = total_message_tokens(&trimmed); + request.messages = trimmed; + ctx.emit(AgentEvent::Compressed { + from_tokens, + to_tokens, + }); + return Ok(()); + } + } + } + }; // `plan` returns `to_keep` as `[system prompts..., recent turns...]`. // Insert the summary *after* the leading system prompts, not at index 0: diff --git a/src/harness/middleware/test.rs b/src/harness/middleware/test.rs index 65f61e9..ae952ee 100644 --- a/src/harness/middleware/test.rs +++ b/src/harness/middleware/test.rs @@ -10,7 +10,7 @@ use crate::harness::context::{RunConfig, RunContext}; use crate::harness::events::{AgentEvent, RecordingListener}; use crate::harness::message::{AssistantMessage, ContentBlock, Message, UserMessage}; use crate::harness::model::{ModelRequest, ModelResponse, PromptSegment, SegmentRole}; -use crate::harness::summarization::{SummarizationPolicy, TrimStrategy}; +use crate::harness::summarization::{SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy}; use crate::harness::tool::{ToolCall, ToolResult}; use crate::harness::usage::Usage; @@ -397,6 +397,173 @@ async fn context_compression_compresses_at_or_above_threshold() { assert!(compressed[0].1 > 0); } +/// A [`Summarizer`] that always fails — models the real gap the +/// `ConcatSummarizer` can never exhibit (a model-backed summarizer whose +/// provider call is rejected). +struct FailingSummarizer; + +#[async_trait] +impl Summarizer for FailingSummarizer { + async fn summarize(&self, _messages: &[Message]) -> Result { + Err(TinyAgentsError::Model("summarizer boom".to_string())) + } +} + +/// Build an over-threshold transcript `[system, big-1, big-2, big-3]` under a +/// 100-token window / 0.5 threshold (→ 50-token trigger budget), so the +/// compression middleware is guaranteed to fire. +fn over_threshold_request() -> (SummarizationPolicy, Vec) { + let policy = SummarizationPolicy { + keep_last: 1, + ..SummarizationPolicy::default() + } + .with_context_window(100) + .with_threshold_fraction(0.5); + let big = "a".repeat(200); + let messages = vec![ + Message::system("You are a helpful assistant."), + user(&format!("{big}-1")), + user(&format!("{big}-2")), + user(&format!("{big}-3")), + ]; + (policy, messages) +} + +#[tokio::test] +async fn context_compression_falls_back_to_trim_when_summarizer_errors() { + // Regression for the "summarizer failure aborts the run" gap: under the + // default FallbackTrim policy a failing summarizer must NOT abort. The + // middleware front-drops the transcript to the trigger budget (keeping the + // leading system prompt) and continues, emitting a MiddlewareFailed + // diagnostic plus a Compressed event for the trim. + let (policy, before) = over_threshold_request(); + let mw = Arc::new(ContextCompressionMiddleware::with_summarizer( + policy, + Box::new(FailingSummarizer), + )); + // Default failure policy is the trim fallback. + assert_eq!(mw.failure_policy(), CompressionFailurePolicy::FallbackTrim); + + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + + let mut request = ModelRequest { + messages: before.clone(), + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .expect("summarizer failure must not abort the run under FallbackTrim"); + + // Trimmed from the front — strictly fewer messages — with the system prompt + // preserved (MaxTokens drops system messages last). + assert!(request.messages.len() < before.len()); + assert!(matches!(request.messages[0], Message::System(_))); + // No summary record could be produced (the summarizer failed). + assert!(mw.records().is_empty()); + + let events: Vec = recorder.events().into_iter().map(|r| r.event).collect(); + assert!( + events.iter().any(|e| matches!( + e, + AgentEvent::MiddlewareFailed { name, error } + if name == "context_compression" && error.contains("summarizer boom") + )), + "a MiddlewareFailed diagnostic naming the failure must be emitted: {events:?}" + ); + assert!( + events + .iter() + .any(|e| matches!(e, AgentEvent::Compressed { .. })), + "the fallback trim must emit a Compressed event: {events:?}" + ); +} + +#[tokio::test] +async fn context_compression_abort_policy_propagates_summarizer_error() { + // Opt back into the legacy behaviour: Abort propagates the error so the run + // fails, but still emits the MiddlewareFailed diagnostic first. + let (policy, before) = over_threshold_request(); + let mw = Arc::new( + ContextCompressionMiddleware::with_summarizer(policy, Box::new(FailingSummarizer)) + .with_failure_policy(CompressionFailurePolicy::Abort), + ); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + + let mut request = ModelRequest { + messages: before, + ..Default::default() + }; + let result = stack.run_before_model(&mut c, &(), &mut request).await; + assert!( + result.is_err(), + "Abort must propagate the summarizer error and fail the run" + ); + + let events: Vec = recorder.events().into_iter().map(|r| r.event).collect(); + assert!( + events.iter().any(|e| matches!( + e, + AgentEvent::MiddlewareFailed { name, .. } if name == "context_compression" + )), + "Abort must still emit the MiddlewareFailed diagnostic: {events:?}" + ); +} + +#[tokio::test] +async fn context_compression_pass_through_policy_keeps_transcript_and_continues() { + // PassThrough leaves the (over-threshold) transcript untouched and lets the + // run continue, emitting only the MiddlewareFailed diagnostic. + let (policy, before) = over_threshold_request(); + let mw = Arc::new( + ContextCompressionMiddleware::with_summarizer(policy, Box::new(FailingSummarizer)) + .with_failure_policy(CompressionFailurePolicy::PassThrough), + ); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + + let mut request = ModelRequest { + messages: before.clone(), + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .expect("PassThrough must not abort the run"); + + // Transcript is untouched. + assert_eq!(request.messages, before); + let events: Vec = recorder.events().into_iter().map(|r| r.event).collect(); + assert!( + events.iter().any(|e| matches!( + e, + AgentEvent::MiddlewareFailed { name, .. } if name == "context_compression" + )), + "PassThrough must emit the MiddlewareFailed diagnostic: {events:?}" + ); + // No compression happened, so no Compressed event. + assert!( + !events + .iter() + .any(|e| matches!(e, AgentEvent::Compressed { .. })), + "PassThrough must not emit a Compressed event: {events:?}" + ); +} + #[tokio::test] async fn context_compression_records_are_bounded_by_max_records() { // A long-running loop that compresses repeatedly must not grow the diff --git a/src/harness/middleware/types.rs b/src/harness/middleware/types.rs index 2b1ddb4..92d26de 100644 --- a/src/harness/middleware/types.rs +++ b/src/harness/middleware/types.rs @@ -489,12 +489,50 @@ pub struct MessageTrimMiddleware { /// [`ContextCompressionMiddleware`] retains before evicting the oldest. pub const DEFAULT_COMPRESSION_RECORD_CAP: usize = 1024; +/// How [`ContextCompressionMiddleware`] recovers when its [`Summarizer`] +/// returns an `Err`. +/// +/// The default +/// [`ConcatSummarizer`][crate::harness::summarization::ConcatSummarizer] is +/// infallible, but the [`Summarizer`] trait allows failure (for example a +/// model-backed summarizer whose provider call is rejected). A summarizer +/// failure strikes precisely on the longest, most valuable transcripts — the +/// ones that reached the compaction threshold — so propagating the error and +/// aborting the whole run is the worst outcome. This policy selects the +/// recovery behaviour; the default is [`FallbackTrim`](Self::FallbackTrim). +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum CompressionFailurePolicy { + /// Propagate the summarizer error, aborting the run. This is the legacy + /// behaviour from before the policy existed. + Abort, + + /// Deterministically front-drop the oldest messages (system messages + /// preserved) until the transcript fits the policy's trigger budget, then + /// continue. Loses old context but keeps the run alive. The default. + #[default] + FallbackTrim, + + /// Leave the transcript untouched and continue. The next model call sees the + /// full (over-threshold) transcript; use when the caller would rather risk a + /// provider-side context error than drop history. + PassThrough, +} + +/// Middleware that condenses older transcript history into a summary when the +/// transcript nears the model's context window. See +/// [`ContextCompressionMiddleware::new`] and +/// [`with_summarizer`](ContextCompressionMiddleware::with_summarizer) for +/// construction, and [`CompressionFailurePolicy`] / +/// [`with_failure_policy`](ContextCompressionMiddleware::with_failure_policy) +/// for how a summarizer error is recovered. pub struct ContextCompressionMiddleware { pub(crate) label: &'static str, pub(crate) policy: SummarizationPolicy, pub(crate) summarizer: Box, pub(crate) records: Mutex>, pub(crate) max_records: usize, + /// Recovery behaviour when [`Summarizer::summarize`] returns `Err`. + pub(crate) on_failure: CompressionFailurePolicy, } // ── MicrocompactMiddleware ──────────────────────────────────────────────────── diff --git a/src/harness/model/mod.rs b/src/harness/model/mod.rs index aa24cb1..ea3e6b8 100644 --- a/src/harness/model/mod.rs +++ b/src/harness/model/mod.rs @@ -753,10 +753,10 @@ impl StreamAccumulator { fn push_tool_chunk(&mut self, call_id: &str, content: &str, tool_name: Option<&str>) { if let Some(entry) = self.tool_chunks.iter_mut().find(|(id, ..)| id == call_id) { entry.1.push_str(content); - if entry.2.is_none() { - if let Some(name) = tool_name.filter(|n| !n.is_empty()) { - entry.2 = Some(name.to_string()); - } + if entry.2.is_none() + && let Some(name) = tool_name.filter(|n| !n.is_empty()) + { + entry.2 = Some(name.to_string()); } } else { self.tool_chunks.push(( diff --git a/src/harness/stream/mod.rs b/src/harness/stream/mod.rs index 14411f0..1712142 100644 --- a/src/harness/stream/mod.rs +++ b/src/harness/stream/mod.rs @@ -143,6 +143,5 @@ pub fn stream(chunks: &[StreamChunk], modes: &[StreamMode]) -> Vec .collect() } -#[cfg(test)] #[cfg(test)] mod test;