diff --git a/src/harness/middleware/mod.rs b/src/harness/middleware/mod.rs index 5e8b323..1c12eee 100644 --- a/src/harness/middleware/mod.rs +++ b/src/harness/middleware/mod.rs @@ -40,6 +40,7 @@ use crate::error::{Result, TinyAgentsError}; use crate::harness::cache::{CacheLayoutEvent, PromptCacheLayout}; use crate::harness::context::RunContext; use crate::harness::events::AgentEvent; +use crate::harness::message::Message; use crate::harness::model::{ModelDelta, ModelRequest, ModelResponse}; use crate::harness::summarization::{ ConcatSummarizer, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy, @@ -720,6 +721,98 @@ impl Middleware for ContextCom } } +// ── MicrocompactMiddleware ──────────────────────────────────────────────────── + +impl MicrocompactMiddleware { + /// Creates a micro-compaction middleware that keeps the newest `keep_recent` + /// tool-result bodies verbatim and blanks older ones with `placeholder`. + /// Event emission is off by default; enable it with + /// [`MicrocompactMiddleware::with_events`]. + pub fn new(keep_recent: usize, placeholder: impl Into) -> Self { + Self { + label: "microcompact", + keep_recent, + placeholder: placeholder.into(), + emit_events: false, + } + } + + /// Enable or disable emitting an + /// [`AgentEvent::Compressed`][crate::harness::events::AgentEvent::Compressed] + /// event whenever at least one tool body is cleared. Off by default so the + /// middleware can be a silent transcript rewrite. + pub fn with_events(mut self, emit_events: bool) -> Self { + self.emit_events = emit_events; + self + } + + /// The number of most-recent tool-result bodies kept verbatim. + pub fn keep_recent(&self) -> usize { + self.keep_recent + } + + /// The placeholder text swapped in for cleared tool-result bodies. + pub fn placeholder(&self) -> &str { + &self.placeholder + } +} + +#[async_trait] +impl Middleware for MicrocompactMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + // Indices of every tool-result message, oldest → newest. + let tool_idxs: Vec = request + .messages + .iter() + .enumerate() + .filter(|(_, m)| matches!(m, Message::Tool(_))) + .map(|(i, _)| i) + .collect(); + if tool_idxs.len() <= self.keep_recent { + return Ok(()); + } + + let from_tokens = if self.emit_events { + total_message_tokens(&request.messages) + } else { + 0 + }; + + let cut = tool_idxs.len() - self.keep_recent; + let mut cleared = 0usize; + for &i in &tool_idxs[..cut] { + // Skip messages already reduced to the placeholder; otherwise swap the + // body for it (idempotent, preserves the tool_call_id). + if request.messages[i].text() == self.placeholder { + continue; + } + if let Message::Tool(t) = &request.messages[i] { + let id = t.tool_call_id.clone(); + request.messages[i] = Message::tool(id, self.placeholder.clone()); + cleared += 1; + } + } + + if self.emit_events && cleared > 0 { + let to_tokens = total_message_tokens(&request.messages); + ctx.emit(AgentEvent::Compressed { + from_tokens, + to_tokens, + }); + } + Ok(()) + } +} + // ── PromptCacheGuardMiddleware ──────────────────────────────────────────────── impl PromptCacheGuardMiddleware { diff --git a/src/harness/middleware/test.rs b/src/harness/middleware/test.rs index 8b9b10b..e58fbe8 100644 --- a/src/harness/middleware/test.rs +++ b/src/harness/middleware/test.rs @@ -348,6 +348,157 @@ async fn context_compression_none_window_falls_back_to_trigger_tokens() { assert_eq!(mw.records().len(), 1); } +// ── MicrocompactMiddleware ──────────────────────────────────────────────────── +// +// These cases are the byte-for-byte parity contract ported from the OpenHuman +// in-house `MicrocompactMiddleware` this type replaces: older tool-result bodies +// are blanked to the caller's placeholder, the newest `keep_recent` are kept +// verbatim, non-tool messages are never touched, and the pass is idempotent. + +const CLEARED: &str = "[Old tool result content cleared]"; + +#[tokio::test] +async fn microcompact_clears_older_tool_bodies_and_keeps_recent() { + let mw = MicrocompactMiddleware::new(1, CLEARED); + let mut request = ModelRequest { + messages: vec![ + Message::system("sys"), + Message::user("hello"), + Message::tool("t1", "FIRST_BODY"), + Message::assistant("thinking"), + Message::tool("t2", "SECOND_BODY"), + Message::tool("t3", "THIRD_BODY"), + ], + ..Default::default() + }; + + mw.before_model(&mut ctx(), &(), &mut request) + .await + .unwrap(); + + // 3 tool messages, keep_recent=1 → the two oldest cleared, newest kept. + assert_eq!(request.messages[2].text(), CLEARED); + assert_eq!(request.messages[4].text(), CLEARED); + assert_eq!(request.messages[5].text(), "THIRD_BODY"); + // Non-tool messages are never touched. + assert_eq!(request.messages[0].text(), "sys"); + assert_eq!(request.messages[1].text(), "hello"); + assert_eq!(request.messages[3].text(), "thinking"); + // Cleared tool messages keep their tool_call_id. + match &request.messages[2] { + Message::Tool(t) => assert_eq!(t.tool_call_id, "t1"), + other => panic!("expected tool message, got {other:?}"), + } +} + +#[tokio::test] +async fn microcompact_is_a_noop_when_within_keep_recent() { + let mw = MicrocompactMiddleware::new(5, CLEARED); + let mut request = ModelRequest { + messages: vec![Message::tool("t1", "A"), Message::tool("t2", "B")], + ..Default::default() + }; + mw.before_model(&mut ctx(), &(), &mut request) + .await + .unwrap(); + assert_eq!(request.messages[0].text(), "A"); + assert_eq!(request.messages[1].text(), "B"); +} + +#[tokio::test] +async fn microcompact_is_idempotent() { + let mw = MicrocompactMiddleware::new(1, CLEARED); + let mut request = ModelRequest { + messages: vec![Message::tool("t1", "FIRST"), Message::tool("t2", "SECOND")], + ..Default::default() + }; + mw.before_model(&mut ctx(), &(), &mut request) + .await + .unwrap(); + assert_eq!(request.messages[0].text(), CLEARED); + // Second pass leaves the already-cleared body as the placeholder. + mw.before_model(&mut ctx(), &(), &mut request) + .await + .unwrap(); + assert_eq!(request.messages[0].text(), CLEARED); + assert_eq!(request.messages[1].text(), "SECOND"); +} + +#[tokio::test] +async fn microcompact_emits_no_event_by_default() { + // Default construction is silent: bodies are still cleared, but no + // Compressed event is emitted (parity with the OpenHuman in-house version). + let mw = Arc::new(MicrocompactMiddleware::new(1, CLEARED)); + 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: vec![Message::tool("t1", "FIRST"), Message::tool("t2", "SECOND")], + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + + assert_eq!(request.messages[0].text(), CLEARED); + let events: Vec = recorder.events().into_iter().map(|r| r.event).collect(); + assert!( + !events + .iter() + .any(|e| matches!(e, AgentEvent::Compressed { .. })), + "no Compressed event should be emitted when events are off" + ); +} + +#[tokio::test] +async fn microcompact_emits_compressed_event_when_enabled() { + let mw = Arc::new(MicrocompactMiddleware::new(1, CLEARED).with_events(true)); + 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: vec![ + Message::tool("t1", "x".repeat(400)), + Message::tool("t2", "y".repeat(400)), + ], + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + + let compressed: Vec<(u64, u64)> = recorder + .events() + .into_iter() + .filter_map(|r| match r.event { + AgentEvent::Compressed { + from_tokens, + to_tokens, + } => Some((from_tokens, to_tokens)), + _ => None, + }) + .collect(); + assert_eq!( + compressed.len(), + 1, + "one Compressed event when a body cleared" + ); + assert!( + compressed[0].0 > compressed[0].1, + "tokens dropped after clear" + ); +} + #[tokio::test] async fn usage_accounting_accumulates_across_calls() { let mw = Arc::new(UsageAccountingMiddleware::new()); diff --git a/src/harness/middleware/types.rs b/src/harness/middleware/types.rs index 40ca649..f73778c 100644 --- a/src/harness/middleware/types.rs +++ b/src/harness/middleware/types.rs @@ -491,6 +491,44 @@ pub struct ContextCompressionMiddleware { pub(crate) records: Mutex>, } +// ── MicrocompactMiddleware ──────────────────────────────────────────────────── + +/// Middleware that clears the bodies of older **tool-result** messages while +/// keeping the `keep_recent` most recent ones verbatim. +/// +/// In `before_model` it walks `request.messages`, finds the tool-result +/// messages, and — once there are more than `keep_recent` of them — replaces the +/// content of every tool result *except* the newest `keep_recent` with a fixed +/// [`placeholder`](Self::placeholder) string, preserving each message's +/// `tool_call_id`. Non-tool messages (system/user/assistant) are never touched +/// and no chat turn is dropped, so this bounds the cost of a long, tool-heavy +/// thread without the semantic loss of summarization. +/// +/// This is the "micro-compaction" companion to +/// [`ContextCompressionMiddleware`]: the latter summarizes *older chat history* +/// when the transcript nears the context window, whereas this one only ever +/// blanks *stale tool payloads* that the model no longer needs verbatim. The two +/// compose cleanly. +/// +/// The operation is **idempotent**: a body already equal to the placeholder is +/// left as-is, so repeated `before_model` passes converge. When there are at +/// most `keep_recent` tool results the middleware is a complete no-op. +/// +/// The placeholder text is caller-supplied (via +/// [`MicrocompactMiddleware::new`]) so host applications can keep their own +/// model-facing wording stable. Event emission is **opt-in** (default off, see +/// [`MicrocompactMiddleware::with_events`]): when enabled and at least one body +/// is cleared, an +/// [`AgentEvent::Compressed`][crate::harness::events::AgentEvent::Compressed] +/// event carrying the before/after token estimate is emitted; when disabled the +/// middleware mutates the request silently. +pub struct MicrocompactMiddleware { + pub(crate) label: &'static str, + pub(crate) keep_recent: usize, + pub(crate) placeholder: String, + pub(crate) emit_events: bool, +} + // ── PromptCacheGuardMiddleware ──────────────────────────────────────────────── /// Middleware that watches the prompt cache layout for accidental prefix