adk_runner/intra_compaction.rs
1//! Intra-invocation context compaction trigger logic.
2//!
3//! This module provides [`IntraInvocationCompactor`], which checks whether the
4//! estimated token count of conversation events exceeds a threshold and triggers
5//! summarization using the existing [`BaseEventsSummarizer`] trait. The actual
6//! summarization logic lives in `adk-agent/src/compaction.rs` (`LlmEventSummarizer`);
7//! this module only handles the trigger decision and overlap preservation.
8
9use adk_core::intra_compaction::{IntraCompactionConfig, estimate_tokens};
10use adk_core::{BaseEventsSummarizer, Event, Result};
11use std::sync::Arc;
12use std::sync::atomic::{AtomicBool, Ordering};
13
14/// Intra-invocation compactor that monitors token usage during a single
15/// invocation and triggers summarization when the context exceeds a threshold.
16///
17/// The compactor enforces at-most-once compaction per LLM call cycle via an
18/// [`AtomicBool`] guard. Call [`reset_cycle`](Self::reset_cycle) at the start
19/// of each LLM call to re-arm the guard.
20///
21/// # Example
22///
23/// ```rust,ignore
24/// use adk_runner::IntraInvocationCompactor;
25/// use adk_core::IntraCompactionConfig;
26///
27/// let compactor = IntraInvocationCompactor::new(
28/// IntraCompactionConfig::default(),
29/// summarizer,
30/// );
31///
32/// // Before each LLM call:
33/// compactor.reset_cycle();
34/// if let Some(compacted) = compactor.maybe_compact(&events).await? {
35/// events = compacted;
36/// }
37/// ```
38pub struct IntraInvocationCompactor {
39 config: IntraCompactionConfig,
40 /// Reuses the existing `BaseEventsSummarizer` trait from adk-core.
41 summarizer: Arc<dyn BaseEventsSummarizer>,
42 /// Track whether compaction already ran this LLM call cycle.
43 compacted_this_cycle: AtomicBool,
44}
45
46impl IntraInvocationCompactor {
47 /// Create a new compactor with the given config and summarizer.
48 pub fn new(config: IntraCompactionConfig, summarizer: Arc<dyn BaseEventsSummarizer>) -> Self {
49 Self { config, summarizer, compacted_this_cycle: AtomicBool::new(false) }
50 }
51
52 /// Check if compaction is needed and perform it if so.
53 ///
54 /// Returns `Some(compacted_events)` if compaction was triggered, or `None`
55 /// if no compaction was needed (below threshold or already compacted this cycle).
56 ///
57 /// On summarizer error, logs a warning and returns `None` (uncompacted history
58 /// is used).
59 pub async fn maybe_compact(&self, events: &[Event]) -> Result<Option<Vec<Event>>> {
60 // Guard: at most once per cycle
61 if self
62 .compacted_this_cycle
63 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
64 .is_err()
65 {
66 return Ok(None);
67 }
68
69 let estimated = estimate_tokens(events, self.config.chars_per_token);
70 if estimated <= self.config.token_threshold {
71 // Below threshold — reset the guard so a future call in the same
72 // cycle can still trigger if tokens grow.
73 self.compacted_this_cycle.store(false, Ordering::SeqCst);
74 return Ok(None);
75 }
76
77 // Determine which events to summarize vs. preserve
78 let overlap = self.config.overlap_event_count.min(events.len());
79 let summarize_end = events.len().saturating_sub(overlap);
80
81 if summarize_end == 0 {
82 // All events are in the overlap window — nothing to summarize
83 self.compacted_this_cycle.store(false, Ordering::SeqCst);
84 return Ok(None);
85 }
86
87 let events_to_summarize = &events[..summarize_end];
88 let overlap_events = &events[summarize_end..];
89
90 // Call the summarizer — on error, log and return uncompacted
91 match self.summarizer.summarize_events(events_to_summarize).await {
92 Ok(Some(summary_event)) => {
93 let mut compacted = Vec::with_capacity(1 + overlap);
94 compacted.push(summary_event);
95 compacted.extend_from_slice(overlap_events);
96 Ok(Some(compacted))
97 }
98 Ok(None) => {
99 // Summarizer returned None (e.g., empty input) — no compaction
100 Ok(None)
101 }
102 Err(e) => {
103 tracing::warn!(
104 error = %e,
105 "intra-invocation compaction failed, continuing with uncompacted history"
106 );
107 Ok(None)
108 }
109 }
110 }
111
112 /// Reset the per-cycle guard. Call this at the start of each LLM call.
113 pub fn reset_cycle(&self) {
114 self.compacted_this_cycle.store(false, Ordering::SeqCst);
115 }
116}