use super::{error, MemoryBinding, Result};
use crate::SharedMessage;
use agent_core::context_archive::{
self, ArchiveDescriptor, ArchivePart, ArchiveRedactor, ArchiveRef, ArchiveStore,
ArchivedMessage, Forgotten, Unscanned,
};
use serde_json::Value;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
fn redactor() -> Arc<ArchiveRedactor> {
Arc::new(|text| {
let mut value = serde_json::from_str(text).unwrap_or_else(|_| Value::String(text.into()));
crate::runtime::trace::export::redact_value(&mut value);
match value {
Value::String(s) => s,
other => other.to_string(),
}
})
}
fn valid_id(id: &str) -> bool {
id.len() == 32 && id.bytes().all(|b| b.is_ascii_hexdigit())
}
fn reader(
base: &std::path::Path,
scope_key: &str,
abort: Arc<AtomicBool>,
) -> std::io::Result<Option<ArchiveStore>> {
let store = ArchiveStore::open_existing(base, scope_key, "context-windows-v1")?;
Ok(store.map(|store| store.with_abort(abort)))
}
fn no_archive() -> std::io::Error {
std::io::Error::new(std::io::ErrorKind::NotFound, context_archive::NOT_FOUND)
}
async fn archive_worker<T: Send + 'static>(
work: impl FnOnce(Arc<AtomicBool>) -> std::io::Result<T> + Send + 'static,
) -> Result<T> {
struct AbortOnDrop(Arc<AtomicBool>);
impl Drop for AbortOnDrop {
fn drop(&mut self) {
self.0.store(true, Ordering::Release);
}
}
let abort = AbortOnDrop(Arc::new(AtomicBool::new(false)));
let flag = abort.0.clone();
tokio::task::spawn_blocking(move || {
if flag.load(Ordering::Acquire) {
return Err(std::io::Error::new(
std::io::ErrorKind::Interrupted,
"archive operation canceled before it started",
));
}
work(flag)
})
.await
.map_err(|_| error("history worker failed"))?
.map_err(|e| error(e.to_string()))
}
impl MemoryBinding {
pub async fn history_seal(
&self,
logical_id: &str,
messages: &[SharedMessage],
note: &str,
) -> Result<(ArchiveRef, String)> {
let scope = self.scope()?.clone();
let base = self.base().to_owned();
let logical = logical_id.to_owned();
let source = messages.to_vec();
let note = note.to_owned();
archive_worker(move |abort| {
ArchiveStore::new(&base, scope.key(), &logical)?
.with_redactor(redactor())
.with_abort(abort)
.seal_with_note(&source, ¬e)
})
.await
}
pub async fn history_seal_window(
&self,
logical_id: &str,
messages: &[SharedMessage],
note: &str,
) -> Result<(Vec<ArchivePart>, String)> {
let scope = self.scope()?.clone();
let base = self.base().to_owned();
let logical = logical_id.to_owned();
let source = messages.to_vec();
let note = note.to_owned();
archive_worker(move |abort| {
let window = ArchiveStore::new(&base, scope.key(), &logical)?
.with_redactor(redactor())
.with_abort(abort)
.seal_window(&source, ¬e)?;
Ok((window.parts, window.note))
})
.await
}
pub async fn history_search(
&self,
query: &str,
limit: usize,
) -> Result<Vec<ArchiveDescriptor>> {
if query.len() > 512 || limit > 32 {
return Err(error("history search exceeds bounds"));
}
let scope = self.scope()?.clone();
let base = self.base().to_owned();
let query = query.to_owned();
archive_worker(move |abort| match reader(&base, scope.key(), abort)? {
Some(store) => store.search(&query, limit),
None => Ok(Vec::new()),
})
.await
}
pub async fn history_search_scan(
&self,
query: &str,
limit: usize,
) -> Result<(Vec<ArchiveDescriptor>, Option<Unscanned>)> {
if query.len() > 512 || limit > 32 {
return Err(error("history search exceeds bounds"));
}
let scope = self.scope()?.clone();
let base = self.base().to_owned();
let query = query.to_owned();
archive_worker(move |abort| match reader(&base, scope.key(), abort)? {
Some(store) => store.search_scan(&query, limit),
None => Ok((Vec::new(), None)),
})
.await
}
pub async fn history_fetch(
&self,
id: &str,
start: usize,
limit: usize,
) -> Result<Vec<ArchivedMessage>> {
if !valid_id(id) || limit > 128 {
return Err(error("invalid history range"));
}
let scope = self.scope()?.clone();
let base = self.base().to_owned();
let id = id.to_owned();
archive_worker(move |abort| {
reader(&base, scope.key(), abort)?
.ok_or_else(no_archive)?
.fetch_with_provenance(&id, start, limit)
})
.await
}
pub async fn history_note(&self, id: &str) -> Result<String> {
if !valid_id(id) {
return Err(error("invalid history id"));
}
let scope = self.scope()?.clone();
let base = self.base().to_owned();
let id = id.to_owned();
archive_worker(move |abort| {
reader(&base, scope.key(), abort)?
.ok_or_else(no_archive)?
.fetch_note(&id)
})
.await
}
pub async fn history_forget(&self, id: &str) -> Result<Forgotten> {
if !valid_id(id) {
return Err(error("invalid history id"));
}
let scope = self.scope()?.clone();
let base = self.base().to_owned();
let id = id.to_owned();
archive_worker(move |abort| {
ArchiveStore::new(&base, scope.key(), "context-windows-v1")?
.with_abort(abort)
.forget(&id)
})
.await
}
}
#[cfg(test)]
mod budget_tests {
use super::*;
use agent_core::memory::store::ProjectScope;
use serde_json::json;
use std::path::Path;
fn binding(root: &Path) -> MemoryBinding {
MemoryBinding::new(root.into(), Ok(ProjectScope::for_root(root).unwrap()))
}
async fn roundtrip(binding: &MemoryBinding) {
let messages = vec![
Arc::new(json!({"role":"system","content":"x".repeat(9 * 1024 * 1024)})),
Arc::new(
json!({"role":"assistant","metadata":"x".repeat(9 * 1024 * 1024),"content":[
{"type":"text","text":"exact source αβ\n"},
{"type":"image","source":{"data":"x".repeat(1024 * 1024)}}
]}),
),
Arc::new(json!({"role":"user","content":[{
"type":"tool_result","tool_use_id":"r1","content":"exact tool log\n"
}]})),
];
let original = messages.clone();
let (receipt, note) = binding
.history_seal("budget-session", &messages, "bounded note")
.await
.unwrap();
assert_eq!(receipt.source_message_count, 3);
assert_eq!(receipt.message_count, 2);
assert_eq!(note, "bounded note");
let fetched = binding.history_fetch(&receipt.id, 0, 8).await.unwrap();
assert_eq!(fetched[0].source_index, 1);
assert_eq!(fetched[0].block_indices, vec![0]);
assert_eq!(
fetched[0].message["content"][0]["text"],
"exact source αβ\n"
);
assert_eq!(fetched[1].message, messages[2]);
assert_eq!(messages, original);
let (retry, retry_note) = binding
.history_seal("budget-session", &messages, "different note")
.await
.unwrap();
assert_eq!(retry.id, receipt.id);
assert_eq!(retry_note, "bounded note");
let before = binding.history_search("", 8).await.unwrap();
let too_big = vec![Arc::new(
json!({"role":"assistant", "content":"x".repeat(context_archive::MAX_INPUT_BYTES + 1)}),
)];
let err = binding
.history_seal("budget-session", &too_big, "")
.await
.unwrap_err();
assert!(err.to_string().contains("input byte budget exceeded"));
assert_eq!(binding.history_search("", 8).await.unwrap(), before);
binding.history_forget(&receipt.id).await.unwrap();
assert!(binding
.history_seal("budget-session", &messages, "bounded note")
.await
.is_err());
}
#[tokio::test]
async fn legacy_history_eligible_budget_roundtrip_and_atomic_refusal() {
let temp = tempfile::tempdir().unwrap();
roundtrip(&binding(temp.path())).await;
}
#[tokio::test]
async fn legacy_history_seals_past_128_archives_and_reports_a_complete_search() {
let temp = tempfile::tempdir().unwrap();
let binding = binding(temp.path());
let mut ids = Vec::new();
for index in 0..129 {
let messages = vec![Arc::new(
json!({"role":"user","content":format!("window {index}.")}),
)];
let (reference, _) = binding
.history_seal("synthetic-session", &messages, "")
.await
.unwrap();
ids.push(reference.id);
}
for id in &ids {
assert_eq!(binding.history_fetch(id, 0, 1).await.unwrap().len(), 1);
}
ids.sort();
ids.dedup();
assert_eq!(ids.len(), 129);
let (rows, stopped_after) = binding.history_search_scan("window 0.", 32).await.unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(stopped_after, None);
assert_eq!(
binding.history_fetch(&rows[0].id, 0, 1).await.unwrap()[0].message["content"],
"window 0."
);
let (rows, stopped_after) = binding.history_search_scan("", 32).await.unwrap();
assert_eq!((rows.len(), stopped_after), (32, None));
assert_eq!(binding.history_search("", 32).await.unwrap(), rows);
assert!(binding.history_search_scan("", 33).await.is_err());
assert!(binding.history_search_scan(&"q".repeat(513), 1).await.is_err());
}
#[cfg(unix)]
#[tokio::test]
async fn legacy_history_reads_create_nothing_and_change_nothing() {
use std::os::unix::fs::MetadataExt;
let temp = tempfile::tempdir().unwrap();
let root = temp.path().to_path_buf();
let binding = binding(&root);
let missing = "a".repeat(32);
let not_there = context_archive::NOT_FOUND;
assert!(binding.history_search("", 8).await.unwrap().is_empty());
let (rows, stopped_after) = binding.history_search_scan("", 8).await.unwrap();
assert!(rows.is_empty() && stopped_after.is_none());
let error = binding.history_fetch(&missing, 0, 1).await.unwrap_err();
assert!(error.to_string().contains(not_there), "{error}");
let error = binding.history_note(&missing).await.unwrap_err();
assert!(error.to_string().contains(not_there), "{error}");
assert_eq!(std::fs::read_dir(&root).unwrap().count(), 0);
let message = vec![Arc::new(json!({"role":"user","content":"kept words"}))];
let (reference, _) = binding
.history_seal("synthetic-session", &message, "kept note")
.await
.unwrap();
let stamps = || {
let mut found = Vec::new();
let mut pending = vec![root.clone()];
while let Some(path) = pending.pop() {
let meta = std::fs::symlink_metadata(&path).unwrap();
if meta.is_dir() {
let entries = std::fs::read_dir(&path).unwrap();
pending.extend(entries.map(|entry| entry.unwrap().path()));
}
let times = (meta.ctime(), meta.ctime_nsec(), meta.mtime(), meta.mtime_nsec());
found.push((path, meta.mode(), times));
}
found.sort();
found
};
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let before = stamps();
assert_eq!(before.len(), 6, "base, context-archives, project, lock, shelf, archive");
let fetched = binding.history_fetch(&reference.id, 0, 1).await.unwrap();
assert_eq!(fetched[0].message["content"], "kept words");
assert_eq!(binding.history_note(&reference.id).await.unwrap(), "kept note");
assert_eq!(binding.history_search("kept", 8).await.unwrap().len(), 1);
let (rows, _) = binding.history_search_scan("", 8).await.unwrap();
assert_eq!(rows.len(), 1);
let error = binding.history_fetch(&missing, 0, 1).await.unwrap_err();
assert!(error.to_string().contains(not_there), "{error}");
assert_eq!(stamps(), before);
}
fn lines(count: usize) -> Vec<SharedMessage> {
(0..count)
.map(|index| Arc::new(json!({"role":"user","content":format!("line {index}.")})))
.collect()
}
#[tokio::test]
async fn legacy_history_seal_window_splits_what_history_seal_refuses() {
let temp = tempfile::tempdir().unwrap();
let binding = binding(temp.path());
let small = lines(3);
let (parts, note) = binding
.history_seal_window("window-session", &small, "the note")
.await
.unwrap();
let (reference, stored) = binding
.history_seal("window-session", &small, "another note")
.await
.unwrap();
let whole = ArchivePart {
archive: reference,
first_source_index: 0,
reduced_messages: 0,
};
assert_eq!(parts, [whole]);
assert_eq!((note.as_str(), stored.as_str()), ("the note", "the note"));
assert_eq!(binding.history_search("", 8).await.unwrap().len(), 1);
let large = lines(context_archive::MAX_MESSAGES + 1);
let refused = binding
.history_seal("window-session", &large, "")
.await
.unwrap_err();
assert!(refused.to_string().contains("budget exceeded"), "{refused}");
assert_eq!(binding.history_search("", 8).await.unwrap().len(), 1);
let (parts, note) = binding
.history_seal_window("window-session", &large, "the note")
.await
.unwrap();
assert_eq!(parts.len(), 2, "{parts:?}");
assert_eq!(note, "the note");
let mut covered = 0;
for part in &parts {
assert_eq!(part.first_source_index, covered);
assert_eq!(part.reduced_messages, 0);
let first = binding.history_fetch(&part.archive.id, 0, 1).await.unwrap();
assert_eq!(first[0].source_index, 0);
assert_eq!(first[0].message, large[covered]);
covered += part.archive.source_message_count;
}
assert_eq!(covered, large.len());
assert_eq!(binding.history_note(&parts[0].archive.id).await.unwrap(), "");
assert_eq!(binding.history_note(&parts[1].archive.id).await.unwrap(), "the note");
let again = binding
.history_seal_window("window-session", &large, "a later note")
.await
.unwrap();
assert_eq!(again, (parts, "the note".to_owned()));
assert_eq!(binding.history_search("", 8).await.unwrap().len(), 3);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_legacy_history_operations_queue_instead_of_reporting_busy() {
let temp = tempfile::tempdir().unwrap();
let binding = binding(temp.path());
let tasks: Vec<_> = (0..8)
.map(|worker| {
let binding = binding.clone();
tokio::spawn(async move {
for round in 0..6 {
let text = format!("worker {worker} round {round}");
let messages = vec![Arc::new(json!({"role":"user","content":text}))];
let logical = format!("synthetic-session-{worker}");
let (reference, note) =
binding.history_seal(&logical, &messages, "note").await?;
assert_eq!(note, "note");
assert_eq!(binding.history_fetch(&reference.id, 0, 1).await?.len(), 1);
binding.history_search("worker", 4).await?;
}
Ok::<_, crate::RuntimeError>(())
})
})
.collect();
for task in tasks {
task.await.unwrap().unwrap();
}
assert_eq!(binding.history_search("", 32).await.unwrap().len(), 32);
}
#[tokio::test]
async fn a_dropped_caller_aborts_its_archive_operation_before_it_starts() {
let (entered_tx, entered_rx) = std::sync::mpsc::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
let (seen_tx, seen_rx) = std::sync::mpsc::channel();
let call = archive_worker(move |abort| {
entered_tx.send(()).unwrap();
release_rx.recv().unwrap();
seen_tx.send(abort.load(Ordering::Acquire)).unwrap();
Ok(())
});
tokio::select! {
biased;
_ = tokio::task::spawn_blocking(move || entered_rx.recv().unwrap()) => {}
_ = call => panic!("worker cannot finish before release"),
}
release_tx.send(()).unwrap();
assert!(tokio::task::spawn_blocking(move || seen_rx.recv().unwrap())
.await
.unwrap());
let completed = archive_worker(|abort| Ok(abort.load(Ordering::Acquire)))
.await
.unwrap();
assert!(!completed, "an awaited operation is never aborted");
}
}