1use std::collections::BTreeMap;
2
3use crate::memory::artifacts::{
4 build_compacted_tool_content, has_recovery_envelope, is_compacted_tool_content,
5 ToolResultArtifactConfig,
6};
7use crate::memory::token_utils::count_messages_tokens;
8use crate::tools::ToolResultRetention;
9use crate::types::{Message, MessageRole, ToolArtifactRef};
10use crate::workspace::{persist_text_artifact, read_validated_text_artifact};
11
12use super::MemoryManager;
13
14#[derive(Debug, Clone)]
15pub(crate) struct MicrocompactionPlan {
16 candidates: Vec<MicrocompactionCandidate>,
17 current_usage: u64,
18 target_usage: u64,
19 pub(crate) candidate_count: usize,
20 pub(crate) estimated_reclaimable_tokens: u64,
21}
22
23#[derive(Debug, Clone)]
24struct MicrocompactionCandidate {
25 message_index: usize,
26 tool_name: String,
27 existing_artifact: Option<ToolArtifactRef>,
28}
29
30#[derive(Debug, Clone, Default)]
31pub(crate) struct MicrocompactionApplication {
32 pub(crate) messages: Vec<Message>,
33 pub(crate) archived_count: usize,
34 pub(crate) reclaimed_tokens: u64,
35 pub(crate) artifact_failure_count: usize,
36}
37
38impl MicrocompactionPlan {
39 pub(crate) fn empty() -> Self {
40 Self {
41 candidates: Vec::new(),
42 current_usage: 0,
43 target_usage: 0,
44 candidate_count: 0,
45 estimated_reclaimable_tokens: 0,
46 }
47 }
48
49 pub(crate) fn has_candidates(&self) -> bool {
50 !self.candidates.is_empty()
51 }
52}
53
54impl MemoryManager {
55 pub(crate) fn validate_model_recovery_surface(
56 &mut self,
57 messages: &[Message],
58 tool_schemas: &[serde_json::Value],
59 ) -> Result<(), String> {
60 let recovery_tool_available = tool_schemas.iter().any(|schema| {
61 schema["function"]["name"].as_str() == Some(crate::constants::READ_FILE_TOOL_NAME)
62 });
63 self.recovery_tool_available = recovery_tool_available;
64 if !recovery_tool_available
65 && messages
66 .iter()
67 .any(|message| is_compacted_tool_content(&message.content))
68 {
69 return Err(
70 "microcompaction_recovery_unavailable: compacted tool results require \
71 a model-visible read_file tool"
72 .to_string(),
73 );
74 }
75 Ok(())
76 }
77
78 pub(crate) fn plan_cycle_microcompaction(
79 &self,
80 messages: &[Message],
81 current_cycle: u32,
82 current_usage: u64,
83 ) -> MicrocompactionPlan {
84 let cleaned = self.remove_previous_summary(messages);
85 let sanitized = crate::memory::message_sanitizer::filter_empty_assistant_messages(&cleaned);
86 self.plan_microcompaction(&sanitized, current_cycle, current_usage)
87 }
88
89 pub(crate) fn plan_microcompaction(
90 &self,
91 messages: &[Message],
92 current_cycle: u32,
93 current_usage: u64,
94 ) -> MicrocompactionPlan {
95 if !self.recovery_tool_available
96 || messages.is_empty()
97 || current_usage <= self.microcompact_trigger_threshold()
98 {
99 return MicrocompactionPlan::empty();
100 }
101
102 let policy = self.config.microcompaction_policy;
103 let tool_call_names = build_tool_call_name_map(messages);
104 let inferred_cycles = infer_message_cycles(messages);
105 let max_inferred_cycle = inferred_cycles.last().copied().unwrap_or_default();
106 let effective_current_cycle = current_cycle.max(max_inferred_cycle.saturating_add(1));
107 let protected_cycle = effective_current_cycle.saturating_sub(policy.keep_recent_cycles);
108 let target = self.microcompact_target_threshold();
109 let marker_config = self.marker_config();
110 let mut candidates = Vec::new();
111 let mut estimated_reclaimable_tokens = 0u64;
112
113 for (message_index, (message, inferred_cycle)) in
114 messages.iter().zip(inferred_cycles).enumerate()
115 {
116 let Some(tool_name) = eligible_tool_name(
117 message,
118 inferred_cycle,
119 protected_cycle,
120 policy.min_result_chars as usize,
121 &tool_call_names,
122 &self.tool_result_retentions,
123 ) else {
124 continue;
125 };
126 let artifact_path = message
127 .artifact_ref
128 .as_ref()
129 .map(|artifact| artifact.path.as_str())
130 .unwrap_or(".vv-agent/artifacts/task/call-00000000000000000000000000000000.txt");
131 let excerpt_source = message
132 .metadata
133 .get("_vv_agent_microcompact_excerpt")
134 .and_then(serde_json::Value::as_str)
135 .unwrap_or(&message.content);
136 let estimated_marker = build_compacted_tool_content(
137 excerpt_source,
138 artifact_path,
139 tool_name,
140 &marker_config,
141 );
142 let before = count_messages_tokens(std::slice::from_ref(message), &self.config.model);
143 let after = count_messages_tokens(
144 &[Message::tool(
145 estimated_marker,
146 message.tool_call_id.clone().unwrap_or_default(),
147 )],
148 &self.config.model,
149 );
150 let reclaimable = before.saturating_sub(after);
151 if reclaimable == 0 {
152 continue;
153 }
154 estimated_reclaimable_tokens = estimated_reclaimable_tokens.saturating_add(reclaimable);
155 candidates.push(MicrocompactionCandidate {
156 message_index,
157 tool_name: tool_name.to_string(),
158 existing_artifact: message.artifact_ref.clone(),
159 });
160 }
161
162 MicrocompactionPlan {
163 candidate_count: candidates.len(),
164 candidates,
165 current_usage,
166 target_usage: target,
167 estimated_reclaimable_tokens,
168 }
169 }
170
171 pub(crate) fn apply_microcompaction(
172 &self,
173 messages: &[Message],
174 plan: &MicrocompactionPlan,
175 ) -> MicrocompactionApplication {
176 if !plan.has_candidates() {
177 return MicrocompactionApplication {
178 messages: messages.to_vec(),
179 ..MicrocompactionApplication::default()
180 };
181 }
182 let Some(backend) = self.workspace_backend.as_ref() else {
183 return MicrocompactionApplication {
184 messages: messages.to_vec(),
185 artifact_failure_count: plan.candidate_count,
186 ..MicrocompactionApplication::default()
187 };
188 };
189
190 let marker_config = self.marker_config();
191 let mut updated = messages.to_vec();
192 let mut archived_count = 0usize;
193 let mut reclaimed_tokens = 0u64;
194 let mut artifact_failure_count = 0usize;
195 let mut projected_usage = plan.current_usage;
196 for candidate in &plan.candidates {
197 if projected_usage <= plan.target_usage {
198 break;
199 }
200 let message = &messages[candidate.message_index];
201 let Some((artifact, complete_content)) =
202 archive_tool_message(self, backend, message, candidate)
203 else {
204 artifact_failure_count += 1;
205 continue;
206 };
207 let marker = build_compacted_tool_content(
208 &complete_content,
209 &artifact.path,
210 &candidate.tool_name,
211 &marker_config,
212 );
213 let before = count_messages_tokens(std::slice::from_ref(message), &self.config.model);
214 let mut replacement = message.clone();
215 replacement.content = marker;
216 replacement.artifact_ref = Some(artifact);
217 replacement
218 .metadata
219 .remove("_vv_agent_microcompact_excerpt");
220 let after =
221 count_messages_tokens(std::slice::from_ref(&replacement), &self.config.model);
222 let actual_reclaimed = before.saturating_sub(after);
223 if actual_reclaimed == 0 {
224 continue;
225 }
226 reclaimed_tokens = reclaimed_tokens.saturating_add(actual_reclaimed);
227 updated[candidate.message_index] = replacement;
228 archived_count += 1;
229 projected_usage = projected_usage.saturating_sub(actual_reclaimed);
230 }
231 MicrocompactionApplication {
232 messages: updated,
233 archived_count,
234 reclaimed_tokens,
235 artifact_failure_count,
236 }
237 }
238
239 pub fn microcompact_trigger_threshold(&self) -> u64 {
240 (self.autocompact_threshold() as f64 * self.config.microcompaction_policy.trigger_ratio)
241 .floor() as u64
242 }
243
244 pub fn microcompact_target_threshold(&self) -> u64 {
245 (self.autocompact_threshold() as f64 * self.config.microcompaction_policy.target_ratio)
246 .floor() as u64
247 }
248
249 pub fn should_preemptive_microcompact(&self, message_length: u64) -> bool {
250 let threshold = self.microcompact_trigger_threshold();
251 threshold > 0 && message_length > threshold
252 }
253
254 pub fn microcompact_messages(
255 &self,
256 messages: &[Message],
257 cycle_index: Option<u32>,
258 ) -> (Vec<Message>, usize) {
259 let Some(cycle_index) = cycle_index else {
260 return (messages.to_vec(), 0);
261 };
262 let current_usage = count_messages_tokens(messages, &self.config.model);
263 let plan = self.plan_microcompaction(messages, cycle_index, current_usage);
264 let application = self.apply_microcompaction(messages, &plan);
265 (application.messages, application.archived_count)
266 }
267
268 fn marker_config(&self) -> ToolResultArtifactConfig {
269 ToolResultArtifactConfig {
270 artifact_namespace: self.artifact_namespace.clone(),
271 compact_threshold: self.config.tool_result_compact_threshold,
272 keep_last: self.config.tool_result_keep_last,
273 excerpt_head: self.config.tool_result_excerpt_head,
274 excerpt_tail: self.config.tool_result_excerpt_tail,
275 }
276 }
277}
278
279fn archive_tool_message(
280 manager: &MemoryManager,
281 backend: &std::sync::Arc<dyn crate::workspace::WorkspaceBackend>,
282 message: &Message,
283 candidate: &MicrocompactionCandidate,
284) -> Option<(ToolArtifactRef, String)> {
285 if let Some(artifact) = candidate.existing_artifact.as_ref() {
286 let complete_content = read_validated_text_artifact(backend.as_ref(), artifact).ok()?;
287 return Some((artifact.clone(), complete_content));
288 }
289 if has_recovery_envelope(&message.content) {
290 return None;
291 }
292 let artifact = persist_text_artifact(
293 backend.clone(),
294 &manager.artifact_namespace,
295 message.tool_call_id.as_deref().unwrap_or("tool-result"),
296 &message.content,
297 )
298 .ok()?;
299 Some((artifact, message.content.clone()))
300}
301
302fn build_tool_call_name_map(messages: &[Message]) -> BTreeMap<String, String> {
303 let mut tool_call_names = BTreeMap::new();
304 for message in messages {
305 if message.role != MessageRole::Assistant {
306 continue;
307 }
308 for tool_call in &message.tool_calls {
309 tool_call_names.insert(tool_call.id.clone(), tool_call.name.clone());
310 }
311 }
312 tool_call_names
313}
314
315fn infer_message_cycles(messages: &[Message]) -> Vec<u32> {
316 let mut current_cycle = 0;
317 let mut inferred = Vec::with_capacity(messages.len());
318 for message in messages {
319 if message.role == MessageRole::Assistant {
320 current_cycle += 1;
321 }
322 inferred.push(current_cycle);
323 }
324 inferred
325}
326
327fn eligible_tool_name<'a>(
328 message: &'a Message,
329 inferred_cycle: u32,
330 protected_cycle: u32,
331 min_result_chars: usize,
332 tool_call_names: &'a BTreeMap<String, String>,
333 retentions: &BTreeMap<String, ToolResultRetention>,
334) -> Option<&'a str> {
335 if message.role != MessageRole::Tool
336 || inferred_cycle >= protected_cycle
337 || message.content.chars().count() <= min_result_chars
338 || is_compacted_tool_content(&message.content)
339 || (message.artifact_ref.is_none() && has_recovery_envelope(&message.content))
340 {
341 return None;
342 }
343 let tool_name = tool_call_names.get(message.tool_call_id.as_deref()?)?;
344 (retentions.get(tool_name).copied().unwrap_or_default() == ToolResultRetention::Archive)
345 .then_some(tool_name.as_str())
346}
347
348#[cfg(test)]
349mod tests {
350 use std::sync::Arc;
351
352 use crate::memory::{MemoryManager, MemoryManagerConfig};
353 use crate::types::Message;
354 use crate::workspace::MemoryWorkspaceBackend;
355
356 use super::{MicrocompactionCandidate, MicrocompactionPlan};
357
358 #[test]
359 fn application_skips_zero_gain_replacement_and_continues_to_later_candidate() {
360 let manager = MemoryManager::new(MemoryManagerConfig::default())
361 .with_workspace_backend(Arc::new(MemoryWorkspaceBackend::default()));
362 let messages = vec![
363 Message::tool("small result", "first"),
364 Message::tool("large result ".repeat(1_000), "second"),
365 ];
366 let plan = MicrocompactionPlan {
367 candidates: vec![
368 MicrocompactionCandidate {
369 message_index: 0,
370 tool_name: "first_tool".to_string(),
371 existing_artifact: None,
372 },
373 MicrocompactionCandidate {
374 message_index: 1,
375 tool_name: "second_tool".to_string(),
376 existing_artifact: None,
377 },
378 ],
379 current_usage: 100,
380 target_usage: 99,
381 candidate_count: 2,
382 estimated_reclaimable_tokens: 200,
383 };
384
385 let applied = manager.apply_microcompaction(&messages, &plan);
386
387 assert_eq!(applied.archived_count, 1);
388 assert_eq!(applied.artifact_failure_count, 0);
389 assert_eq!(applied.messages[0], messages[0]);
390 assert!(applied.messages[1]
391 .content
392 .starts_with(crate::memory::TOOL_RESULT_COMPACT_MARKER));
393 }
394}