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;