vv_agent/memory/
post_compact_restore.rs

1use std::path::{Path, PathBuf};
2
3use serde_json::Value;
4
5use crate::memory::token_utils::estimate_tokens;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct PostCompactRestoreConfig {
9    pub total_budget_tokens: u64,
10    pub max_tokens_per_file: u64,
11    pub max_files: usize,
12    pub token_model: String,
13}
14
15impl Default for PostCompactRestoreConfig {
16    fn default() -> Self {
17        Self {
18            total_budget_tokens: 30_000,
19            max_tokens_per_file: 5_000,
20            max_files: 10,
21            token_model: String::new(),
22        }
23    }
24}
25
26pub fn restore_key_files(
27    summary_data: &Value,
28    workspace: Option<&Path>,
29    config: &PostCompactRestoreConfig,
30) -> String {
31    let Some(workspace) = workspace else {
32        return String::new();
33    };
34    let Some(raw_files) = summary_data
35        .get("files_examined_or_modified")
36        .and_then(Value::as_array)
37    else {
38        return String::new();
39    };
40    let Ok(workspace_root) = workspace.canonicalize() else {
41        return String::new();
42    };
43
44    let mut indexed_files = raw_files
45        .iter()
46        .enumerate()
47        .filter_map(|(index, item)| item.as_object().map(|file| (index, file)))
48        .collect::<Vec<_>>();
49    indexed_files.sort_by_key(|(index, file)| {
50        let action = file
51            .get("action")
52            .and_then(Value::as_str)
53            .unwrap_or("read")
54            .trim()
55            .to_ascii_lowercase();
56        (action_priority(&action), *index)
57    });
58
59    let mut restored_parts = Vec::new();
60    let mut total_tokens = 0_u64;
61    for (_, file_info) in indexed_files.into_iter().take(config.max_files) {
62        let path_value = file_info
63            .get("path")
64            .and_then(Value::as_str)
65            .unwrap_or_default()
66            .trim();
67        if path_value.is_empty() {
68            continue;
69        }
70        let Some(resolved_path) = resolve_workspace_file(&workspace_root, path_value) else {
71            continue;
72        };
73        if !resolved_path.is_file() {
74            continue;
75        }
76        let Ok(content) = std::fs::read_to_string(&resolved_path) else {
77            continue;
78        };
79
80        let action = file_info
81            .get("action")
82            .and_then(Value::as_str)
83            .unwrap_or("read")
84            .trim()
85            .to_ascii_lowercase();
86        let action = if action.is_empty() { "read" } else { &action };
87        let content = truncate_to_token_budget(
88            &content,
89            config.max_tokens_per_file.max(1),
90            &config.token_model,
91        );
92        let candidate =
93            format!("<file path=\"{path_value}\" action=\"{action}\">\n{content}\n</file>");
94        let candidate_tokens = estimate_tokens(&candidate, &config.token_model);
95        if total_tokens + candidate_tokens > config.total_budget_tokens {
96            break;
97        }
98        restored_parts.push(candidate);
99        total_tokens += candidate_tokens;
100    }
101
102    if restored_parts.is_empty() {
103        return String::new();
104    }
105    format!(
106        "<Post-Compaction File Context>\nThe following files were relevant in the previous conversation context:\n\n{}\n</Post-Compaction File Context>",
107        restored_parts.join("\n\n")
108    )
109}
110
111fn action_priority(action: &str) -> u8 {
112    match action {
113        "modified" => 0,
114        "created" => 1,
115        "deleted" => 2,
116        "read" => 3,
117        _ => 99,
118    }
119}
120
121fn resolve_workspace_file(workspace_root: &Path, relative_path: &str) -> Option<PathBuf> {
122    let candidate = workspace_root.join(relative_path).canonicalize().ok()?;
123    candidate.starts_with(workspace_root).then_some(candidate)
124}
125
126fn truncate_to_token_budget(content: &str, max_tokens: u64, model: &str) -> String {
127    if estimate_tokens(content, model) <= max_tokens {
128        return content.to_string();
129    }
130
131    let notice = "\n... [truncated after compaction restore]";
132    let mut low = 0_usize;
133    let mut high = content.len();
134    let mut best = notice.to_string();
135    while low <= high {
136        let middle = (low + high) / 2;
137        let prefix = floor_char_boundary(content, middle);
138        let candidate = format!("{}{}", content[..prefix].trim_end(), notice);
139        let candidate_tokens = estimate_tokens(&candidate, model);
140        if candidate_tokens <= max_tokens {
141            best = candidate;
142            low = middle + 1;
143        } else if middle == 0 {
144            break;
145        } else {
146            high = middle - 1;
147        }
148    }
149    best
150}
151
152fn floor_char_boundary(content: &str, mut index: usize) -> usize {
153    index = index.min(content.len());
154    while index > 0 && !content.is_char_boundary(index) {
155        index -= 1;
156    }
157    index
158}
vv_agent/memory/post_compact_restore.rs

vv_agent/memory/
post_compact_restore.rs
