use crate::cli::mcp::protocol::JsonRpcError;
use crate::edit::EditChange;
use lru::LruCache;
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use tokio::sync::Mutex;
pub const EDIT_CACHE_MAX_ENTRY_BYTES: usize = 256 * 1024;
pub const EDIT_CACHE_TOTAL_CAP_BYTES: usize = 8 * 1024 * 1024;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EditCacheEntry {
pub file_path: PathBuf,
pub preview_token: String,
pub original_text: String,
pub modified_text: String,
pub changes: Vec<EditChange>,
pub timestamp: chrono::DateTime<chrono::Utc>,
}
impl EditCacheEntry {
pub fn estimated_size(&self) -> usize {
let path_bytes = self.file_path.as_os_str().len();
#[cfg(target_os = "windows")]
let path_bytes = path_bytes * 2;
path_bytes
+ self.preview_token.len()
+ self.original_text.len()
+ self.modified_text.len()
+ self
.changes
.iter()
.map(|c| c.estimated_size())
.sum::<usize>()
+ 64 }
}
#[derive(Debug)]
pub struct EntryTooLargeError {
pub size: usize,
pub limit: usize,
}
impl std::fmt::Display for EntryTooLargeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"edit-preview entry ({} bytes) exceeds per-entry limit ({} bytes)",
self.size, self.limit
)
}
}
impl std::error::Error for EntryTooLargeError {}
pub struct EditCache {
entries: Mutex<(LruCache<PathBuf, EditCacheEntry>, usize)>,
}
impl Default for EditCache {
fn default() -> Self {
Self {
entries: Mutex::new((
LruCache::new(NonZeroUsize::new(10_000).unwrap()),
0,
)),
}
}
}
impl EditCache {
pub fn new() -> Self {
Self::default()
}
async fn get_abs_path_and_cache_file(
&self,
project_storage: &Path,
file_path: &Path,
) -> (PathBuf, PathBuf) {
let abs_path = if file_path.is_absolute() {
file_path.to_path_buf()
} else {
tokio::fs::canonicalize(file_path)
.await
.unwrap_or_else(|_| file_path.to_path_buf())
};
let hash = blake3::hash(abs_path.to_string_lossy().as_bytes()).to_hex();
let cache_file = project_storage
.join("edit_cache")
.join(format!("{}.json", hash));
(abs_path, cache_file)
}
pub async fn set(
&self,
project_storage: &Path,
entry: EditCacheEntry,
) -> Result<Result<(), EntryTooLargeError>, JsonRpcError> {
let entry_size = entry.estimated_size();
if entry_size > EDIT_CACHE_MAX_ENTRY_BYTES {
return Ok(Err(EntryTooLargeError {
size: entry_size,
limit: EDIT_CACHE_MAX_ENTRY_BYTES,
}));
}
let (abs_path, cache_file) = self
.get_abs_path_and_cache_file(project_storage, &entry.file_path)
.await;
let cache_dir = cache_file.parent().unwrap();
if tokio::fs::metadata(cache_dir).await.is_err() {
tokio::fs::create_dir_all(cache_dir).await.map_err(|e| {
JsonRpcError::internal_error(format!(
"Failed to create edit cache directory: {}",
e
))
})?;
}
let json = serde_json::to_string_pretty(&entry).map_err(|e| {
JsonRpcError::internal_error(format!("Failed to serialize edit cache: {}", e))
})?;
tokio::fs::write(&cache_file, json).await.map_err(|e| {
JsonRpcError::internal_error(format!("Failed to write edit cache to disk: {}", e))
})?;
{
let mut guard = self.entries.lock().await;
let (ref mut entries, ref mut total_bytes) = *guard;
if let Some(existing) = entries.pop(&abs_path) {
*total_bytes = total_bytes.saturating_sub(existing.estimated_size());
}
while *total_bytes + entry_size > EDIT_CACHE_TOTAL_CAP_BYTES {
if let Some((_, removed)) = entries.pop_lru() {
*total_bytes = total_bytes.saturating_sub(removed.estimated_size());
} else {
break;
}
}
entries.put(abs_path, entry);
*total_bytes = total_bytes.saturating_add(entry_size);
}
Ok(Ok(()))
}
pub async fn get(&self, project_storage: &Path, file_path: &Path) -> Option<EditCacheEntry> {
let (abs_path, cache_file) = self
.get_abs_path_and_cache_file(project_storage, file_path)
.await;
{
let mut guard = self.entries.lock().await;
if let Some(entry) = guard.0.get_mut(&abs_path) {
return Some(entry.clone());
}
}
if let Ok(json) = tokio::fs::read_to_string(&cache_file).await {
if let Ok(entry) = serde_json::from_str::<EditCacheEntry>(&json) {
let entry_size = entry.estimated_size();
if entry_size <= EDIT_CACHE_MAX_ENTRY_BYTES {
let mut guard = self.entries.lock().await;
let (ref mut entries, ref mut total_bytes) = *guard;
if let Some(existing) = entries.pop(&abs_path) {
*total_bytes = total_bytes.saturating_sub(existing.estimated_size());
}
while *total_bytes + entry_size > EDIT_CACHE_TOTAL_CAP_BYTES {
if let Some((_, removed)) = entries.pop_lru() {
*total_bytes = total_bytes.saturating_sub(removed.estimated_size());
} else {
break;
}
}
*total_bytes = total_bytes.saturating_add(entry_size);
entries.put(abs_path, entry.clone());
}
return Some(entry);
}
}
None
}
pub async fn clear(&self, project_storage: &Path, file_path: &Path) {
let (abs_path, cache_file) = self
.get_abs_path_and_cache_file(project_storage, file_path)
.await;
{
let mut guard = self.entries.lock().await;
if let Some(removed) = guard.0.pop(&abs_path) {
guard.1 = guard.1.saturating_sub(removed.estimated_size());
}
}
let _ = tokio::fs::remove_file(cache_file).await;
}
pub async fn hot_cache_bytes(&self) -> usize {
self.entries.lock().await.1
}
}
pub static GLOBAL_EDIT_CACHE: Lazy<EditCache> = Lazy::new(EditCache::new);
#[cfg(test)]
mod tests {
use super::*;
use crate::edit::EditChange;
use std::path::PathBuf;
fn make_entry(original_len: usize, modified_len: usize) -> EditCacheEntry {
EditCacheEntry {
file_path: PathBuf::from("/test/file.rs"),
preview_token: "test-token".to_string(),
original_text: "x".repeat(original_len),
modified_text: "y".repeat(modified_len),
changes: vec![EditChange::ReplaceText {
start: 0,
end: original_len.min(1),
new_text: "y".to_string(),
}],
timestamp: chrono::Utc::now(),
}
}
#[tokio::test]
async fn test_edit_cache_rejects_oversized_entry() {
let temp_dir = tempfile::tempdir().unwrap();
let cache = EditCache::new();
let oversized = make_entry(EDIT_CACHE_MAX_ENTRY_BYTES + 1000, 10);
let result = cache
.set(temp_dir.path(), oversized)
.await
.expect("no IO error");
assert!(
result.is_err(),
"oversized entry should be rejected, but was accepted"
);
let err = result.unwrap_err();
assert_eq!(err.limit, EDIT_CACHE_MAX_ENTRY_BYTES);
assert!(err.size > EDIT_CACHE_MAX_ENTRY_BYTES);
}
#[tokio::test]
async fn test_edit_cache_total_residency_capped() {
let temp_dir = tempfile::tempdir().unwrap();
let cache = EditCache::new();
let entry_size = 100_000; let entries_to_fill = (EDIT_CACHE_TOTAL_CAP_BYTES / entry_size) + 3;
for i in 0..entries_to_fill {
let mut entry = make_entry(entry_size / 2, entry_size / 2);
entry.file_path = PathBuf::from(format!("/test/file_{}.rs", i));
let result = cache
.set(temp_dir.path(), entry)
.await
.expect("no IO error");
assert!(result.is_ok(), "entry {} should be accepted", i);
}
assert!(
cache.hot_cache_bytes().await <= EDIT_CACHE_TOTAL_CAP_BYTES + 10_000,
"hot cache bytes ({}) should not exceed cap ({})",
cache.hot_cache_bytes().await,
EDIT_CACHE_TOTAL_CAP_BYTES
);
}
#[tokio::test]
async fn test_edit_cache_backfill_replaces_existing_hot_entry_bytes() {
let temp_dir = tempfile::tempdir().unwrap();
let cache = EditCache::new();
let mut initial = make_entry(32, 32);
initial.file_path = PathBuf::from("/test/backfill.rs");
cache
.set(temp_dir.path(), initial.clone())
.await
.unwrap()
.unwrap();
{
let mut guard = cache.entries.lock().await;
let mut disk_entry = make_entry(160, 160);
disk_entry.file_path = initial.file_path.clone();
let disk_entry_size = disk_entry.estimated_size();
guard.1 -= initial.estimated_size();
guard.0.put(initial.file_path.clone(), disk_entry.clone());
guard.1 += disk_entry_size;
}
let backed_fill = cache
.get(temp_dir.path(), &initial.file_path)
.await
.expect("entry should be recovered from cold storage");
assert_eq!(backed_fill.file_path, initial.file_path);
let expected = backed_fill.estimated_size();
assert_eq!(
cache.hot_cache_bytes().await,
expected,
"backfill should replace the existing hot entry bytes exactly"
);
}
#[test]
fn test_edit_cache_entry_estimated_size() {
let entry = make_entry(100, 200);
let size = entry.estimated_size();
assert!(size > 300, "estimated size should account for text content");
}
}