Skip to main content

swink_agent_memory/
compaction.rs

1//! Summarization-aware context compaction.
2//!
3//! Wraps the core sliding-window strategy with the ability to inject a
4//! pre-computed summary of dropped messages. The summary is generated
5//! asynchronously (outside the agent loop) and stored for the next
6//! synchronous compaction pass.
7
8use std::sync::{Arc, Mutex};
9
10#[allow(deprecated)]
11use swink_agent::sliding_window;
12use swink_agent::{AgentMessage, AssistantMessage, ContentBlock, LlmMessage};
13
14/// Result of a compaction operation (diagnostic type for future use).
15#[non_exhaustive]
16#[derive(Debug)]
17pub struct CompactionResult {
18    /// The compacted messages.
19    pub messages: Vec<AgentMessage>,
20    /// Number of messages removed during compaction.
21    pub removed_count: usize,
22    /// The summary that was injected, if any.
23    pub summary: Option<String>,
24}
25
26impl CompactionResult {
27    /// Creates a new compaction result from its constituent parts.
28    #[must_use]
29    pub fn new(messages: Vec<AgentMessage>, removed_count: usize, summary: Option<String>) -> Self {
30        Self {
31            messages,
32            removed_count,
33            summary,
34        }
35    }
36}
37
38/// Summarization-aware context compactor.
39///
40/// Combines the core [`sliding_window`] strategy with an optional summary
41/// that replaces dropped messages. The summary is stored externally via
42/// [`set_summary`](Self::set_summary) and injected after the anchor
43/// messages during the next compaction pass.
44///
45/// # Usage
46///
47/// ```rust,ignore
48/// let compactor = SummarizingCompactor::new(100_000, 50_000, 2);
49/// let agent = Agent::new()
50///     .with_transform_context(compactor.compaction_fn());
51///
52/// // After each turn, if messages were dropped:
53/// compactor.set_summary("Summary of earlier conversation...");
54/// ```
55pub struct SummarizingCompactor {
56    /// Pre-computed summary to inject during compaction.
57    summary: Arc<Mutex<Option<String>>>,
58    /// Normal token budget.
59    normal_budget: usize,
60    /// Overflow token budget.
61    overflow_budget: usize,
62    /// Number of anchor messages to preserve.
63    anchor: usize,
64}
65
66impl SummarizingCompactor {
67    /// Create a new compactor with the given budget parameters.
68    pub fn new(normal_budget: usize, overflow_budget: usize, anchor: usize) -> Self {
69        Self {
70            summary: Arc::new(Mutex::new(None)),
71            normal_budget,
72            overflow_budget,
73            anchor,
74        }
75    }
76
77    /// Returns a closure compatible with `Agent::with_transform_context()`.
78    ///
79    /// Behaves like [`sliding_window`] but injects a stored summary after
80    /// the anchor messages when compaction occurs. If no summary is stored,
81    /// behaves identically to `sliding_window`.
82    ///
83    /// The summary is consumed after injection — it will not be re-injected
84    /// on subsequent compaction passes.
85    pub fn compaction_fn(&self) -> impl Fn(&mut Vec<AgentMessage>, bool) + Send + Sync {
86        let summary = Arc::clone(&self.summary);
87        #[allow(deprecated)]
88        let base = sliding_window(self.normal_budget, self.overflow_budget, self.anchor);
89        let anchor = self.anchor;
90
91        move |messages: &mut Vec<AgentMessage>, overflow: bool| {
92            let len_before = messages.len();
93
94            // Run the base sliding window compaction.
95            base(messages, overflow);
96
97            let len_after = messages.len();
98
99            // If messages were dropped and we have a stored summary, inject it.
100            if len_after < len_before {
101                let mut guard = summary
102                    .lock()
103                    .unwrap_or_else(std::sync::PoisonError::into_inner);
104                if let Some(text) = guard.take() {
105                    let summary_msg = AgentMessage::Llm(LlmMessage::Assistant(
106                        AssistantMessage::new(
107                            vec![ContentBlock::Text {
108                                text: format!("[Context summary of earlier conversation]\n{text}"),
109                            }],
110                            String::new(),
111                            String::new(),
112                        )
113                        .with_timestamp(0),
114                    ));
115
116                    // Insert after anchor messages.
117                    let insert_pos = anchor.min(messages.len());
118                    messages.insert(insert_pos, summary_msg);
119                }
120            }
121        }
122    }
123
124    /// Store a summary to be injected on the next compaction pass.
125    ///
126    /// This should be called after each turn where context was compacted,
127    /// passing a summary of the messages that were dropped. In the future
128    /// this will be generated via an LLM call; for now callers provide the
129    /// text directly.
130    pub fn set_summary(&self, text: impl Into<String>) {
131        let mut guard = self
132            .summary
133            .lock()
134            .unwrap_or_else(std::sync::PoisonError::into_inner);
135        *guard = Some(text.into());
136    }
137
138    /// Clear the stored summary.
139    pub fn clear_summary(&self) {
140        let mut guard = self
141            .summary
142            .lock()
143            .unwrap_or_else(std::sync::PoisonError::into_inner);
144        *guard = None;
145    }
146
147    /// Returns true if a summary is currently stored.
148    pub fn has_summary(&self) -> bool {
149        let guard = self
150            .summary
151            .lock()
152            .unwrap_or_else(std::sync::PoisonError::into_inner);
153        guard.is_some()
154    }
155}
156
157#[cfg(test)]
158#[path = "compaction_tests.rs"]
159mod tests;