Skip to main content

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}