Skip to content

Commit 98d4d2c

Browse files
committed
Add CompressionFailurePolicy so a summarizer error no longer aborts the run
ContextCompressionMiddleware propagated a Summarizer error through before_model, failing the whole run exactly on the longest transcripts that reached the compaction threshold. Add a CompressionFailurePolicy (Abort | FallbackTrim | PassThrough) on the middleware, defaulting to a deterministic front-drop trim to the policy's trigger budget (system messages preserved) and emitting a MiddlewareFailed diagnostic on fallback. Add tests covering all three policies with a failing summarizer.
1 parent b1f0d82 commit 98d4d2c

3 files changed

Lines changed: 261 additions & 4 deletions

File tree

‎src/harness/middleware/library/context.rs‎

Lines changed: 55 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,9 @@
77
use super::*;
88
use crate::harness::cache::{CacheLayoutEvent, PromptCacheLayout};
99
use crate::harness::middleware::{
10-
ContextCompressionMiddleware, DEFAULT_CACHE_GUARD_EVENT_CAP, DEFAULT_COMPRESSION_RECORD_CAP,
11-
MessageTrimMiddleware, MicrocompactMiddleware, PromptCacheGuardMiddleware,
10+
CompressionFailurePolicy, ContextCompressionMiddleware, DEFAULT_CACHE_GUARD_EVENT_CAP,
11+
DEFAULT_COMPRESSION_RECORD_CAP, MessageTrimMiddleware, MicrocompactMiddleware,
12+
PromptCacheGuardMiddleware,
1213
};
1314
use crate::harness::summarization::{
1415
ConcatSummarizer, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy,
@@ -57,16 +58,34 @@ impl ContextCompressionMiddleware {
5758
}
5859

5960
/// Creates a compression middleware with a custom [`Summarizer`].
61+
///
62+
/// The failure policy defaults to
63+
/// [`CompressionFailurePolicy::FallbackTrim`]; override it with
64+
/// [`with_failure_policy`](Self::with_failure_policy).
6065
pub fn with_summarizer(policy: SummarizationPolicy, summarizer: Box<dyn Summarizer>) -> Self {
6166
Self {
6267
label: "context_compression",
6368
policy,
6469
summarizer,
6570
records: std::sync::Mutex::new(std::collections::VecDeque::new()),
6671
max_records: DEFAULT_COMPRESSION_RECORD_CAP,
72+
on_failure: CompressionFailurePolicy::default(),
6773
}
6874
}
6975

76+
/// Sets the [`CompressionFailurePolicy`] applied when the [`Summarizer`]
77+
/// returns an `Err`. Defaults to
78+
/// [`CompressionFailurePolicy::FallbackTrim`].
79+
pub fn with_failure_policy(mut self, on_failure: CompressionFailurePolicy) -> Self {
80+
self.on_failure = on_failure;
81+
self
82+
}
83+
84+
/// Returns the configured [`CompressionFailurePolicy`].
85+
pub fn failure_policy(&self) -> CompressionFailurePolicy {
86+
self.on_failure
87+
}
88+
7089
/// Sets the maximum number of [`SummaryRecord`]s retained before the
7190
/// oldest is evicted. `0` disables recording entirely.
7291
pub fn with_max_records(mut self, max_records: usize) -> Self {
@@ -123,7 +142,40 @@ impl<State: Send + Sync, Ctx: Send + Sync> Middleware<State, Ctx> for ContextCom
123142
}
124143

125144
let from_tokens = total_message_tokens(&request.messages);
126-
let record = self.summarizer.summarize(&to_summarize).await?;
145+
let record = match self.summarizer.summarize(&to_summarize).await {
146+
Ok(record) => record,
147+
Err(err) => {
148+
// A summarizer failure hits precisely the longest, most valuable
149+
// transcripts (the ones that reached the compaction threshold).
150+
// Emit a diagnostic and recover per the configured
151+
// `CompressionFailurePolicy` instead of aborting the whole run.
152+
ctx.emit(AgentEvent::MiddlewareFailed {
153+
name: self.label.to_string(),
154+
error: err.to_string(),
155+
});
156+
match self.on_failure {
157+
// Legacy behaviour: propagate and let the run fail.
158+
CompressionFailurePolicy::Abort => return Err(err),
159+
// Keep the (over-threshold) transcript verbatim and continue.
160+
CompressionFailurePolicy::PassThrough => return Ok(()),
161+
// Deterministic front-drop to the policy's trigger budget,
162+
// preserving system messages (see `TrimStrategy::MaxTokens`).
163+
CompressionFailurePolicy::FallbackTrim => {
164+
let trimmed = trim_messages(
165+
&request.messages,
166+
&TrimStrategy::MaxTokens(self.policy.trigger_budget()),
167+
);
168+
let to_tokens = total_message_tokens(&trimmed);
169+
request.messages = trimmed;
170+
ctx.emit(AgentEvent::Compressed {
171+
from_tokens,
172+
to_tokens,
173+
});
174+
return Ok(());
175+
}
176+
}
177+
}
178+
};
127179

128180
// `plan` returns `to_keep` as `[system prompts..., recent turns...]`.
129181
// Insert the summary *after* the leading system prompts, not at index 0:

‎src/harness/middleware/test.rs‎

Lines changed: 168 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use crate::harness::context::{RunConfig, RunContext};
1010
use crate::harness::events::{AgentEvent, RecordingListener};
1111
use crate::harness::message::{AssistantMessage, ContentBlock, Message, UserMessage};
1212
use crate::harness::model::{ModelRequest, ModelResponse, PromptSegment, SegmentRole};
13-
use crate::harness::summarization::{SummarizationPolicy, TrimStrategy};
13+
use crate::harness::summarization::{SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy};
1414
use crate::harness::tool::{ToolCall, ToolResult};
1515
use crate::harness::usage::Usage;
1616

@@ -397,6 +397,173 @@ async fn context_compression_compresses_at_or_above_threshold() {
397397
assert!(compressed[0].1 > 0);
398398
}
399399

400+
/// A [`Summarizer`] that always fails — models the real gap the
401+
/// `ConcatSummarizer` can never exhibit (a model-backed summarizer whose
402+
/// provider call is rejected).
403+
struct FailingSummarizer;
404+
405+
#[async_trait]
406+
impl Summarizer for FailingSummarizer {
407+
async fn summarize(&self, _messages: &[Message]) -> Result<SummaryRecord> {
408+
Err(TinyAgentsError::Model("summarizer boom".to_string()))
409+
}
410+
}
411+
412+
/// Build an over-threshold transcript `[system, big-1, big-2, big-3]` under a
413+
/// 100-token window / 0.5 threshold (→ 50-token trigger budget), so the
414+
/// compression middleware is guaranteed to fire.
415+
fn over_threshold_request() -> (SummarizationPolicy, Vec<Message>) {
416+
let policy = SummarizationPolicy {
417+
keep_last: 1,
418+
..SummarizationPolicy::default()
419+
}
420+
.with_context_window(100)
421+
.with_threshold_fraction(0.5);
422+
let big = "a".repeat(200);
423+
let messages = vec![
424+
Message::system("You are a helpful assistant."),
425+
user(&format!("{big}-1")),
426+
user(&format!("{big}-2")),
427+
user(&format!("{big}-3")),
428+
];
429+
(policy, messages)
430+
}
431+
432+
#[tokio::test]
433+
async fn context_compression_falls_back_to_trim_when_summarizer_errors() {
434+
// Regression for the "summarizer failure aborts the run" gap: under the
435+
// default FallbackTrim policy a failing summarizer must NOT abort. The
436+
// middleware front-drops the transcript to the trigger budget (keeping the
437+
// leading system prompt) and continues, emitting a MiddlewareFailed
438+
// diagnostic plus a Compressed event for the trim.
439+
let (policy, before) = over_threshold_request();
440+
let mw = Arc::new(ContextCompressionMiddleware::with_summarizer(
441+
policy,
442+
Box::new(FailingSummarizer),
443+
));
444+
// Default failure policy is the trim fallback.
445+
assert_eq!(mw.failure_policy(), CompressionFailurePolicy::FallbackTrim);
446+
447+
let mut stack: MiddlewareStack<()> = MiddlewareStack::new();
448+
stack.push(mw.clone());
449+
450+
let recorder = Arc::new(RecordingListener::new());
451+
let mut c = ctx();
452+
c.events.subscribe(recorder.clone());
453+
454+
let mut request = ModelRequest {
455+
messages: before.clone(),
456+
..Default::default()
457+
};
458+
stack
459+
.run_before_model(&mut c, &(), &mut request)
460+
.await
461+
.expect("summarizer failure must not abort the run under FallbackTrim");
462+
463+
// Trimmed from the front — strictly fewer messages — with the system prompt
464+
// preserved (MaxTokens drops system messages last).
465+
assert!(request.messages.len() < before.len());
466+
assert!(matches!(request.messages[0], Message::System(_)));
467+
// No summary record could be produced (the summarizer failed).
468+
assert!(mw.records().is_empty());
469+
470+
let events: Vec<AgentEvent> = recorder.events().into_iter().map(|r| r.event).collect();
471+
assert!(
472+
events.iter().any(|e| matches!(
473+
e,
474+
AgentEvent::MiddlewareFailed { name, error }
475+
if name == "context_compression" && error.contains("summarizer boom")
476+
)),
477+
"a MiddlewareFailed diagnostic naming the failure must be emitted: {events:?}"
478+
);
479+
assert!(
480+
events
481+
.iter()
482+
.any(|e| matches!(e, AgentEvent::Compressed { .. })),
483+
"the fallback trim must emit a Compressed event: {events:?}"
484+
);
485+
}
486+
487+
#[tokio::test]
488+
async fn context_compression_abort_policy_propagates_summarizer_error() {
489+
// Opt back into the legacy behaviour: Abort propagates the error so the run
490+
// fails, but still emits the MiddlewareFailed diagnostic first.
491+
let (policy, before) = over_threshold_request();
492+
let mw = Arc::new(
493+
ContextCompressionMiddleware::with_summarizer(policy, Box::new(FailingSummarizer))
494+
.with_failure_policy(CompressionFailurePolicy::Abort),
495+
);
496+
let mut stack: MiddlewareStack<()> = MiddlewareStack::new();
497+
stack.push(mw.clone());
498+
499+
let recorder = Arc::new(RecordingListener::new());
500+
let mut c = ctx();
501+
c.events.subscribe(recorder.clone());
502+
503+
let mut request = ModelRequest {
504+
messages: before,
505+
..Default::default()
506+
};
507+
let result = stack.run_before_model(&mut c, &(), &mut request).await;
508+
assert!(
509+
result.is_err(),
510+
"Abort must propagate the summarizer error and fail the run"
511+
);
512+
513+
let events: Vec<AgentEvent> = recorder.events().into_iter().map(|r| r.event).collect();
514+
assert!(
515+
events.iter().any(|e| matches!(
516+
e,
517+
AgentEvent::MiddlewareFailed { name, .. } if name == "context_compression"
518+
)),
519+
"Abort must still emit the MiddlewareFailed diagnostic: {events:?}"
520+
);
521+
}
522+
523+
#[tokio::test]
524+
async fn context_compression_pass_through_policy_keeps_transcript_and_continues() {
525+
// PassThrough leaves the (over-threshold) transcript untouched and lets the
526+
// run continue, emitting only the MiddlewareFailed diagnostic.
527+
let (policy, before) = over_threshold_request();
528+
let mw = Arc::new(
529+
ContextCompressionMiddleware::with_summarizer(policy, Box::new(FailingSummarizer))
530+
.with_failure_policy(CompressionFailurePolicy::PassThrough),
531+
);
532+
let mut stack: MiddlewareStack<()> = MiddlewareStack::new();
533+
stack.push(mw.clone());
534+
535+
let recorder = Arc::new(RecordingListener::new());
536+
let mut c = ctx();
537+
c.events.subscribe(recorder.clone());
538+
539+
let mut request = ModelRequest {
540+
messages: before.clone(),
541+
..Default::default()
542+
};
543+
stack
544+
.run_before_model(&mut c, &(), &mut request)
545+
.await
546+
.expect("PassThrough must not abort the run");
547+
548+
// Transcript is untouched.
549+
assert_eq!(request.messages, before);
550+
let events: Vec<AgentEvent> = recorder.events().into_iter().map(|r| r.event).collect();
551+
assert!(
552+
events.iter().any(|e| matches!(
553+
e,
554+
AgentEvent::MiddlewareFailed { name, .. } if name == "context_compression"
555+
)),
556+
"PassThrough must emit the MiddlewareFailed diagnostic: {events:?}"
557+
);
558+
// No compression happened, so no Compressed event.
559+
assert!(
560+
!events
561+
.iter()
562+
.any(|e| matches!(e, AgentEvent::Compressed { .. })),
563+
"PassThrough must not emit a Compressed event: {events:?}"
564+
);
565+
}
566+
400567
#[tokio::test]
401568
async fn context_compression_records_are_bounded_by_max_records() {
402569
// A long-running loop that compresses repeatedly must not grow the

‎src/harness/middleware/types.rs‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -489,12 +489,50 @@ pub struct MessageTrimMiddleware {
489489
/// [`ContextCompressionMiddleware`] retains before evicting the oldest.
490490
pub const DEFAULT_COMPRESSION_RECORD_CAP: usize = 1024;
491491

492+
/// How [`ContextCompressionMiddleware`] recovers when its [`Summarizer`]
493+
/// returns an `Err`.
494+
///
495+
/// The default
496+
/// [`ConcatSummarizer`][crate::harness::summarization::ConcatSummarizer] is
497+
/// infallible, but the [`Summarizer`] trait allows failure (for example a
498+
/// model-backed summarizer whose provider call is rejected). A summarizer
499+
/// failure strikes precisely on the longest, most valuable transcripts — the
500+
/// ones that reached the compaction threshold — so propagating the error and
501+
/// aborting the whole run is the worst outcome. This policy selects the
502+
/// recovery behaviour; the default is [`FallbackTrim`](Self::FallbackTrim).
503+
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
504+
pub enum CompressionFailurePolicy {
505+
/// Propagate the summarizer error, aborting the run. This is the legacy
506+
/// behaviour from before the policy existed.
507+
Abort,
508+
509+
/// Deterministically front-drop the oldest messages (system messages
510+
/// preserved) until the transcript fits the policy's trigger budget, then
511+
/// continue. Loses old context but keeps the run alive. The default.
512+
#[default]
513+
FallbackTrim,
514+
515+
/// Leave the transcript untouched and continue. The next model call sees the
516+
/// full (over-threshold) transcript; use when the caller would rather risk a
517+
/// provider-side context error than drop history.
518+
PassThrough,
519+
}
520+
521+
/// Middleware that condenses older transcript history into a summary when the
522+
/// transcript nears the model's context window. See
523+
/// [`ContextCompressionMiddleware::new`] and
524+
/// [`with_summarizer`](ContextCompressionMiddleware::with_summarizer) for
525+
/// construction, and [`CompressionFailurePolicy`] /
526+
/// [`with_failure_policy`](ContextCompressionMiddleware::with_failure_policy)
527+
/// for how a summarizer error is recovered.
492528
pub struct ContextCompressionMiddleware {
493529
pub(crate) label: &'static str,
494530
pub(crate) policy: SummarizationPolicy,
495531
pub(crate) summarizer: Box<dyn Summarizer>,
496532
pub(crate) records: Mutex<VecDeque<SummaryRecord>>,
497533
pub(crate) max_records: usize,
534+
/// Recovery behaviour when [`Summarizer::summarize`] returns `Err`.
535+
pub(crate) on_failure: CompressionFailurePolicy,
498536
}
499537

500538
// ── MicrocompactMiddleware ────────────────────────────────────────────────────

0 commit comments

Comments
 (0)