| 1 | //! `recall_archive` tool — search prior cycle archives (issue #127). |
| 2 | //! |
| 3 | //! Companion to the checkpoint-restart cycle architecture (#124). When the |
| 4 | //! agent's `<carry_forward>` briefing missed something, this tool scans the |
| 5 | //! on-disk JSONL archives at `~/.deepseek/sessions/<id>/cycles/*.jsonl` and |
| 6 | //! returns the top-N matching messages. |
| 7 | //! |
| 8 | //! ## Scoring |
| 9 | //! |
| 10 | //! v1: a simplified BM25 over tokenized message text. No external embedding |
| 11 | //! model, no cache — every call walks the archives. Acceptable because the |
| 12 | //! per-cycle archive is bounded by the 110K cycle threshold and most sessions |
| 13 | //! cross at most a handful of cycles. v2 (later) can add an |
| 14 | //! `~/.deepseek/embeddings/` cache built on archive write. |
| 15 | |
| 16 | use std::collections::HashMap; |
| 17 | use std::fs::read_dir; |
| 18 | use std::path::PathBuf; |
| 19 | |
| 20 | use async_trait::async_trait; |
| 21 | use serde::Serialize; |
| 22 | use serde_json::{Value, json}; |
| 23 | |
| 24 | use super::spec::{ |
| 25 | ApprovalRequirement, ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, |
| 26 | optional_u64, required_str, |
| 27 | }; |
| 28 | use crate::cycle_manager::open_archive; |
| 29 | use crate::models::{ContentBlock, Message}; |
| 30 | |
| 31 | const DEFAULT_MAX_RESULTS: usize = 3; |
| 32 | const HARD_MAX_RESULTS: usize = 10; |
| 33 | const CONTEXT_WINDOW_CHARS: usize = 240; |
| 34 | |
| 35 | /// BM25 hyper-parameters. Standard defaults from the literature. |
| 36 | const K1: f64 = 1.5; |
| 37 | const B: f64 = 0.75; |
| 38 | |
| 39 | pub struct RecallArchiveTool; |
| 40 | |
| 41 | #[derive(Debug, Clone, Serialize)] |
| 42 | struct RecallHit { |
| 43 | cycle: u32, |
| 44 | /// 0-based message index within the cycle. |
| 45 | message_index: usize, |
| 46 | role: String, |
| 47 | score: f64, |
| 48 | /// Short window around the best match, with `…` markers when truncated. |
| 49 | excerpt: String, |
| 50 | } |
| 51 | |
| 52 | #[async_trait] |
| 53 | impl ToolSpec for RecallArchiveTool { |
| 54 | fn name(&self) -> &'static str { |
| 55 | "recall_archive" |
| 56 | } |
| 57 | |
| 58 | fn description(&self) -> &'static str { |
| 59 | "Search prior context cycles for content not in your briefing. Use sparingly — \ |
| 60 | frequent recalls mean your briefing was too sparse; refine your next briefing." |
| 61 | } |
| 62 | |
| 63 | fn input_schema(&self) -> Value { |
| 64 | json!({ |
| 65 | "type": "object", |
| 66 | "properties": { |
| 67 | "query": { |
| 68 | "type": "string", |
| 69 | "description": "Search query. Tokenized and BM25-scored against archived messages." |
| 70 | }, |
| 71 | "cycle": { |
| 72 | "type": "integer", |
| 73 | "description": "Optional: limit to a specific prior cycle number." |
| 74 | }, |
| 75 | "max_results": { |
| 76 | "type": "integer", |
| 77 | "description": "Maximum hits to return (default 3, hard-capped at 10)." |
| 78 | } |
| 79 | }, |
| 80 | "required": ["query"] |
| 81 | }) |
| 82 | } |
| 83 | |
| 84 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 85 | vec![ToolCapability::ReadOnly] |
| 86 | } |
| 87 | |
| 88 | fn approval_requirement(&self) -> ApprovalRequirement { |
| 89 | ApprovalRequirement::Auto |
| 90 | } |
| 91 | |
| 92 | async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolResult, ToolError> { |
| 93 | let query = required_str(&input, "query")?.trim().to_string(); |
| 94 | if query.is_empty() { |
| 95 | return Err(ToolError::invalid_input("query cannot be empty")); |
| 96 | } |
| 97 | |
| 98 | let max_results = (optional_u64(&input, "max_results", DEFAULT_MAX_RESULTS as u64) |
| 99 | as usize) |
| 100 | .clamp(1, HARD_MAX_RESULTS); |
| 101 | let cycle_filter = input.get("cycle").and_then(Value::as_u64).map(|n| n as u32); |
| 102 | |
| 103 | let session_id = context.state_namespace.as_str(); |
| 104 | let archives = list_archives(session_id).map_err(|err| { |
| 105 | ToolError::execution_failed(format!("Failed to enumerate cycle archives: {err}")) |
| 106 | })?; |
| 107 | |
| 108 | if archives.is_empty() { |
| 109 | return Ok(ToolResult::success(json!({ |
| 110 | "hits": [], |
| 111 | "note": "No prior cycle archives exist. The session has not crossed a cycle boundary yet." |
| 112 | }).to_string())); |
| 113 | } |
| 114 | |
| 115 | let documents = load_messages(&archives, cycle_filter).map_err(|err| { |
| 116 | ToolError::execution_failed(format!("Failed to read cycle archives: {err}")) |
| 117 | })?; |
| 118 | |
| 119 | if documents.is_empty() { |
| 120 | let note = match cycle_filter { |
| 121 | Some(c) => format!("Cycle {c} has no messages in its archive."), |
| 122 | None => "Cycle archives exist but contain no message text.".to_string(), |
| 123 | }; |
| 124 | return Ok(ToolResult::success( |
| 125 | json!({"hits": [], "note": note}).to_string(), |
| 126 | )); |
| 127 | } |
| 128 | |
| 129 | let query_tokens = tokenize(&query); |
| 130 | if query_tokens.is_empty() { |
| 131 | return Err(ToolError::invalid_input( |
| 132 | "query has no scoring tokens after tokenization", |
| 133 | )); |
| 134 | } |
| 135 | |
| 136 | let hits = score_bm25(&documents, &query_tokens, max_results); |
| 137 | |
| 138 | let payload = json!({ |
| 139 | "query": query, |
| 140 | "cycles_searched": archives.len(), |
| 141 | "messages_scanned": documents.len(), |
| 142 | "hits": hits, |
| 143 | }); |
| 144 | |
| 145 | Ok(ToolResult::success(payload.to_string())) |
| 146 | } |
| 147 | } |
| 148 | |
| 149 | /// One archived message + its provenance, ready to score. |
| 150 | struct ArchivedDoc { |
| 151 | cycle: u32, |
| 152 | message_index: usize, |
| 153 | role: String, |
| 154 | text: String, |
| 155 | tokens: Vec<String>, |
| 156 | } |
| 157 | |
| 158 | fn archive_root(session_id: &str) -> Result<PathBuf, std::io::Error> { |
| 159 | let home = dirs::home_dir().ok_or_else(|| { |
| 160 | std::io::Error::new( |
| 161 | std::io::ErrorKind::NotFound, |
| 162 | "Could not resolve home directory for cycle archive root", |
| 163 | ) |
| 164 | })?; |
| 165 | Ok(home |
| 166 | .join(".deepseek") |
| 167 | .join("sessions") |
| 168 | .join(session_id) |
| 169 | .join("cycles")) |
| 170 | } |
| 171 | |
| 172 | /// Enumerate all archive files for a session, sorted by cycle number ascending. |
| 173 | fn list_archives(session_id: &str) -> Result<Vec<(u32, PathBuf)>, std::io::Error> { |
| 174 | let root = archive_root(session_id)?; |
| 175 | if !root.exists() { |
| 176 | return Ok(Vec::new()); |
| 177 | } |
| 178 | let mut archives: Vec<(u32, PathBuf)> = Vec::new(); |
| 179 | for entry in read_dir(&root)? { |
| 180 | let entry = entry?; |
| 181 | let path = entry.path(); |
| 182 | if path.extension().and_then(|s| s.to_str()) != Some("jsonl") { |
| 183 | continue; |
| 184 | } |
| 185 | let stem = match path.file_stem().and_then(|s| s.to_str()) { |
| 186 | Some(s) => s, |
| 187 | None => continue, |
| 188 | }; |
| 189 | let Ok(cycle_n) = stem.parse::<u32>() else { |
| 190 | continue; |
| 191 | }; |
| 192 | archives.push((cycle_n, path)); |
| 193 | } |
| 194 | archives.sort_by_key(|(n, _)| *n); |
| 195 | Ok(archives) |
| 196 | } |
| 197 | |
| 198 | /// Read messages from each archive into a flat scoreable list. |
| 199 | fn load_messages( |
| 200 | archives: &[(u32, PathBuf)], |
| 201 | cycle_filter: Option<u32>, |
| 202 | ) -> Result<Vec<ArchivedDoc>, anyhow::Error> { |
| 203 | let mut docs: Vec<ArchivedDoc> = Vec::new(); |
| 204 | for (cycle_n, path) in archives { |
| 205 | if let Some(filter) = cycle_filter |
| 206 | && *cycle_n != filter |
| 207 | { |
| 208 | continue; |
| 209 | } |
| 210 | let (header, reader) = open_archive(path)?; |
| 211 | for (idx, message_result) in reader.enumerate() { |
| 212 | let message = message_result?; |
| 213 | let text = message_text(&message); |
| 214 | if text.trim().is_empty() { |
| 215 | continue; |
| 216 | } |
| 217 | let tokens = tokenize(&text); |
| 218 | if tokens.is_empty() { |
| 219 | continue; |
| 220 | } |
| 221 | docs.push(ArchivedDoc { |
| 222 | cycle: header.cycle, |
| 223 | message_index: idx, |
| 224 | role: message.role, |
| 225 | text, |
| 226 | tokens, |
| 227 | }); |
| 228 | } |
| 229 | } |
| 230 | Ok(docs) |
| 231 | } |
| 232 | |
| 233 | /// Concatenate all text-bearing content blocks of a message. |
| 234 | fn message_text(message: &Message) -> String { |
| 235 | let mut out = String::new(); |
| 236 | let mut push = |s: &str| { |
| 237 | if !out.is_empty() { |
| 238 | out.push('\n'); |
| 239 | } |
| 240 | out.push_str(s); |
| 241 | }; |
| 242 | for block in &message.content { |
| 243 | match block { |
| 244 | ContentBlock::Text { text, .. } => push(text), |
| 245 | ContentBlock::ToolUse { name, input, .. } => { |
| 246 | push(&format!("[tool_use {name}] {input}")); |
| 247 | } |
| 248 | ContentBlock::ToolResult { content, .. } => { |
| 249 | push(&format!("[tool_result] {content}")); |
| 250 | } |
| 251 | ContentBlock::Thinking { thinking } => { |
| 252 | push(&format!("[thinking] {thinking}")); |
| 253 | } |
| 254 | ContentBlock::ServerToolUse { name, input, .. } => { |
| 255 | push(&format!("[server_tool_use {name}] {input}")); |
| 256 | } |
| 257 | ContentBlock::ToolSearchToolResult { content, .. } => { |
| 258 | push(&format!("[tool_search_result] {content}")); |
| 259 | } |
| 260 | ContentBlock::CodeExecutionToolResult { content, .. } => { |
| 261 | push(&format!("[code_execution_result] {content}")); |
| 262 | } |
| 263 | } |
| 264 | } |
| 265 | out |
| 266 | } |
| 267 | |
| 268 | /// Lower-case, split on non-alphanumerics, drop short tokens. Same recipe as |
| 269 | /// most lightweight BM25 implementations. |
| 270 | fn tokenize(text: &str) -> Vec<String> { |
| 271 | text.to_ascii_lowercase() |
| 272 | .split(|c: char| !c.is_alphanumeric()) |
| 273 | .filter(|s| s.len() >= 2) |
| 274 | .map(str::to_string) |
| 275 | .collect() |
| 276 | } |
| 277 | |
| 278 | /// Score documents against a query using BM25, return the top-N. |
| 279 | fn score_bm25(docs: &[ArchivedDoc], query_tokens: &[String], max_results: usize) -> Vec<RecallHit> { |
| 280 | if docs.is_empty() || query_tokens.is_empty() { |
| 281 | return Vec::new(); |
| 282 | } |
| 283 | |
| 284 | let n = docs.len() as f64; |
| 285 | let avgdl: f64 = docs.iter().map(|d| d.tokens.len() as f64).sum::<f64>() / n.max(1.0); |
| 286 | |
| 287 | // Document frequency per query term. |
| 288 | let mut df: HashMap<&str, u64> = HashMap::new(); |
| 289 | for token in query_tokens { |
| 290 | let mut count = 0u64; |
| 291 | for doc in docs { |
| 292 | if doc.tokens.iter().any(|t| t == token) { |
| 293 | count += 1; |
| 294 | } |
| 295 | } |
| 296 | df.insert(token.as_str(), count); |
| 297 | } |
| 298 | |
| 299 | let mut scored: Vec<(f64, &ArchivedDoc)> = docs |
| 300 | .iter() |
| 301 | .map(|doc| (bm25_doc_score(doc, query_tokens, &df, n, avgdl), doc)) |
| 302 | .filter(|(score, _)| *score > 0.0) |
| 303 | .collect(); |
| 304 | |
| 305 | scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); |
| 306 | scored.truncate(max_results); |
| 307 | |
| 308 | scored |
| 309 | .into_iter() |
| 310 | .map(|(score, doc)| RecallHit { |
| 311 | cycle: doc.cycle, |
| 312 | message_index: doc.message_index, |
| 313 | role: doc.role.clone(), |
| 314 | score: round_score(score), |
| 315 | excerpt: best_window(&doc.text, query_tokens, CONTEXT_WINDOW_CHARS), |
| 316 | }) |
| 317 | .collect() |
| 318 | } |
| 319 | |
| 320 | fn bm25_doc_score( |
| 321 | doc: &ArchivedDoc, |
| 322 | query_tokens: &[String], |
| 323 | df: &HashMap<&str, u64>, |
| 324 | n: f64, |
| 325 | avgdl: f64, |
| 326 | ) -> f64 { |
| 327 | let dl = doc.tokens.len() as f64; |
| 328 | if dl == 0.0 { |
| 329 | return 0.0; |
| 330 | } |
| 331 | let mut score = 0.0; |
| 332 | for token in query_tokens { |
| 333 | let tf = doc.tokens.iter().filter(|t| *t == token).count() as f64; |
| 334 | if tf == 0.0 { |
| 335 | continue; |
| 336 | } |
| 337 | let df_t = df.get(token.as_str()).copied().unwrap_or(0) as f64; |
| 338 | let idf = ((n - df_t + 0.5) / (df_t + 0.5) + 1.0).ln(); |
| 339 | let denom = tf + K1 * (1.0 - B + B * (dl / avgdl.max(1.0))); |
| 340 | score += idf * (tf * (K1 + 1.0)) / denom.max(f64::EPSILON); |
| 341 | } |
| 342 | score |
| 343 | } |
| 344 | |
| 345 | fn round_score(score: f64) -> f64 { |
| 346 | (score * 1000.0).round() / 1000.0 |
| 347 | } |
| 348 | |
| 349 | /// Find the substring of `text` of at most `window_chars` characters that |
| 350 | /// contains the densest cluster of query tokens. Returns it with `…` markers |
| 351 | /// when truncated. Falls back to a head-of-text excerpt when no tokens hit. |
| 352 | fn best_window(text: &str, query_tokens: &[String], window_chars: usize) -> String { |
| 353 | let lower = text.to_ascii_lowercase(); |
| 354 | let mut hit_positions: Vec<usize> = Vec::new(); |
| 355 | for token in query_tokens { |
| 356 | let mut start = 0usize; |
| 357 | while let Some(pos) = lower[start..].find(token.as_str()) { |
| 358 | hit_positions.push(start + pos); |
| 359 | start += pos + token.len(); |
| 360 | } |
| 361 | } |
| 362 | if hit_positions.is_empty() { |
| 363 | return head_excerpt(text, window_chars); |
| 364 | } |
| 365 | hit_positions.sort_unstable(); |
| 366 | |
| 367 | // Greedy: center the window on the first hit, walk forward as long as |
| 368 | // additional hits fit in the window. |
| 369 | let center = hit_positions[0]; |
| 370 | let half = window_chars / 2; |
| 371 | let start = center.saturating_sub(half); |
| 372 | let end = (start + window_chars).min(text.len()); |
| 373 | let start = align_char_boundary(text, start, false); |
| 374 | let end = align_char_boundary(text, end, true); |
| 375 | let prefix = if start > 0 { "…" } else { "" }; |
| 376 | let suffix = if end < text.len() { "…" } else { "" }; |
| 377 | format!("{prefix}{}{suffix}", &text[start..end]) |
| 378 | } |
| 379 | |
| 380 | fn head_excerpt(text: &str, max_chars: usize) -> String { |
| 381 | if text.len() <= max_chars { |
| 382 | return text.to_string(); |
| 383 | } |
| 384 | let cut = align_char_boundary(text, max_chars, true); |
| 385 | format!("{}…", &text[..cut]) |
| 386 | } |
| 387 | |
| 388 | /// Walk left or right until `idx` lands on a UTF-8 char boundary. |
| 389 | fn align_char_boundary(text: &str, mut idx: usize, walk_right: bool) -> usize { |
| 390 | if idx >= text.len() { |
| 391 | return text.len(); |
| 392 | } |
| 393 | while idx > 0 && idx < text.len() && !text.is_char_boundary(idx) { |
| 394 | if walk_right { |
| 395 | idx += 1; |
| 396 | } else { |
| 397 | idx -= 1; |
| 398 | } |
| 399 | } |
| 400 | idx |
| 401 | } |
| 402 | |
| 403 | #[cfg(test)] |
| 404 | mod tests { |
| 405 | use super::*; |
| 406 | use crate::cycle_manager::archive_cycle; |
| 407 | use crate::models::{ContentBlock, Message}; |
| 408 | use chrono::Utc; |
| 409 | use tempfile::TempDir; |
| 410 | |
| 411 | fn user_msg(text: &str) -> Message { |
| 412 | Message { |
| 413 | role: "user".to_string(), |
| 414 | content: vec![ContentBlock::Text { |
| 415 | text: text.to_string(), |
| 416 | cache_control: None, |
| 417 | }], |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | fn asst_msg(text: &str) -> Message { |
| 422 | Message { |
| 423 | role: "assistant".to_string(), |
| 424 | content: vec![ContentBlock::Text { |
| 425 | text: text.to_string(), |
| 426 | cache_control: None, |
| 427 | }], |
| 428 | } |
| 429 | } |
| 430 | |
| 431 | /// Serializes home-redirecting tests since cargo runs tests in parallel |
| 432 | /// by default. Held for the full test (no `.await` while holding it). |
| 433 | static HOME_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); |
| 434 | |
| 435 | /// Guard that points `dirs::home_dir()` at a tempdir for the test's |
| 436 | /// lifetime and restores the original on drop. On Unix this means |
| 437 | /// `HOME`; on Windows it means `USERPROFILE`. We set both so the same |
| 438 | /// guard works portably. Holds `HOME_LOCK` to serialize. |
| 439 | struct HomeGuard { |
| 440 | _tmp: TempDir, |
| 441 | original_home: Option<String>, |
| 442 | original_userprofile: Option<String>, |
| 443 | _lock: std::sync::MutexGuard<'static, ()>, |
| 444 | } |
| 445 | impl HomeGuard { |
| 446 | fn new() -> Self { |
| 447 | let lock = HOME_LOCK.lock().unwrap_or_else(|p| p.into_inner()); |
| 448 | let tmp = TempDir::new().expect("tempdir"); |
| 449 | let original_home = std::env::var("HOME").ok(); |
| 450 | let original_userprofile = std::env::var("USERPROFILE").ok(); |
| 451 | // SAFETY: serialized by HOME_LOCK; only this thread mutates the |
| 452 | // env vars for the duration of the guard. |
| 453 | unsafe { |
| 454 | std::env::set_var("HOME", tmp.path()); |
| 455 | std::env::set_var("USERPROFILE", tmp.path()); |
| 456 | } |
| 457 | Self { |
| 458 | _tmp: tmp, |
| 459 | original_home, |
| 460 | original_userprofile, |
| 461 | _lock: lock, |
| 462 | } |
| 463 | } |
| 464 | } |
| 465 | impl Drop for HomeGuard { |
| 466 | fn drop(&mut self) { |
| 467 | // SAFETY: still holding HOME_LOCK. |
| 468 | unsafe { |
| 469 | match self.original_home.take() { |
| 470 | Some(v) => std::env::set_var("HOME", v), |
| 471 | None => std::env::remove_var("HOME"), |
| 472 | } |
| 473 | match self.original_userprofile.take() { |
| 474 | Some(v) => std::env::set_var("USERPROFILE", v), |
| 475 | None => std::env::remove_var("USERPROFILE"), |
| 476 | } |
| 477 | } |
| 478 | } |
| 479 | } |
| 480 | |
| 481 | fn fresh_session_id() -> String { |
| 482 | format!("test-{}", uuid::Uuid::new_v4()) |
| 483 | } |
| 484 | |
| 485 | fn ctx_for_session(workspace: &std::path::Path, session_id: &str) -> ToolContext { |
| 486 | ToolContext::new(workspace).with_state_namespace(session_id.to_string()) |
| 487 | } |
| 488 | |
| 489 | #[test] |
| 490 | fn tokenize_lowers_splits_drops_short() { |
| 491 | // Filter is `len >= 2`, so "a" and "0" drop; "42" stays. |
| 492 | let toks = tokenize("Hello, World! a 42 OAuth-2.0"); |
| 493 | assert_eq!(toks, vec!["hello", "world", "42", "oauth"]); |
| 494 | } |
| 495 | |
| 496 | #[test] |
| 497 | fn message_text_concatenates_blocks() { |
| 498 | let m = Message { |
| 499 | role: "user".to_string(), |
| 500 | content: vec![ |
| 501 | ContentBlock::Text { |
| 502 | text: "first".to_string(), |
| 503 | cache_control: None, |
| 504 | }, |
| 505 | ContentBlock::Text { |
| 506 | text: "second".to_string(), |
| 507 | cache_control: None, |
| 508 | }, |
| 509 | ], |
| 510 | }; |
| 511 | assert_eq!(message_text(&m), "first\nsecond"); |
| 512 | } |
| 513 | |
| 514 | #[test] |
| 515 | fn list_archives_handles_missing_dir() { |
| 516 | let _home = HomeGuard::new(); |
| 517 | let sid = fresh_session_id(); |
| 518 | let archives = list_archives(&sid).expect("list_archives"); |
| 519 | assert!(archives.is_empty()); |
| 520 | } |
| 521 | |
| 522 | #[test] |
| 523 | fn list_archives_sorts_by_cycle_number() { |
| 524 | let _home = HomeGuard::new(); |
| 525 | let sid = fresh_session_id(); |
| 526 | let now = Utc::now(); |
| 527 | archive_cycle(&sid, 3, &[user_msg("c3")], "deepseek-v4-pro", now).unwrap(); |
| 528 | archive_cycle(&sid, 1, &[user_msg("c1")], "deepseek-v4-pro", now).unwrap(); |
| 529 | archive_cycle(&sid, 2, &[user_msg("c2")], "deepseek-v4-pro", now).unwrap(); |
| 530 | let archives = list_archives(&sid).unwrap(); |
| 531 | let cycles: Vec<u32> = archives.iter().map(|(n, _)| *n).collect(); |
| 532 | assert_eq!(cycles, vec![1, 2, 3]); |
| 533 | } |
| 534 | |
| 535 | #[tokio::test] |
| 536 | async fn execute_returns_empty_when_no_archives() { |
| 537 | let _home = HomeGuard::new(); |
| 538 | let sid = fresh_session_id(); |
| 539 | let workspace = TempDir::new().unwrap(); |
| 540 | let ctx = ctx_for_session(workspace.path(), &sid); |
| 541 | let tool = RecallArchiveTool; |
| 542 | let result = tool |
| 543 | .execute(json!({"query": "anything"}), &ctx) |
| 544 | .await |
| 545 | .unwrap(); |
| 546 | assert!(result.content.contains("No prior cycle archives")); |
| 547 | } |
| 548 | |
| 549 | #[tokio::test] |
| 550 | async fn execute_finds_matching_messages() { |
| 551 | let _home = HomeGuard::new(); |
| 552 | let sid = fresh_session_id(); |
| 553 | let workspace = TempDir::new().unwrap(); |
| 554 | let ctx = ctx_for_session(workspace.path(), &sid); |
| 555 | let now = Utc::now(); |
| 556 | let messages = vec![ |
| 557 | user_msg("How does the cycle restart strategy work?"), |
| 558 | asst_msg("It archives messages to JSONL when crossing the 110K threshold."), |
| 559 | user_msg("What happens if briefing is too short?"), |
| 560 | asst_msg("Use recall_archive to retrieve specific past content from JSONL files."), |
| 561 | ]; |
| 562 | archive_cycle(&sid, 1, &messages, "deepseek-v4-pro", now).unwrap(); |
| 563 | |
| 564 | let tool = RecallArchiveTool; |
| 565 | let result = tool |
| 566 | .execute( |
| 567 | json!({"query": "JSONL archive briefing", "max_results": 3}), |
| 568 | &ctx, |
| 569 | ) |
| 570 | .await |
| 571 | .unwrap(); |
| 572 | assert!( |
| 573 | result.content.contains("\"cycle\":1"), |
| 574 | "got: {}", |
| 575 | result.content |
| 576 | ); |
| 577 | assert!( |
| 578 | result.content.contains("\"hits\""), |
| 579 | "got: {}", |
| 580 | result.content |
| 581 | ); |
| 582 | assert!(result.content.contains("JSONL"), "got: {}", result.content); |
| 583 | } |
| 584 | |
| 585 | #[tokio::test] |
| 586 | async fn execute_filters_by_cycle() { |
| 587 | let _home = HomeGuard::new(); |
| 588 | let sid = fresh_session_id(); |
| 589 | let workspace = TempDir::new().unwrap(); |
| 590 | let ctx = ctx_for_session(workspace.path(), &sid); |
| 591 | let now = Utc::now(); |
| 592 | archive_cycle( |
| 593 | &sid, |
| 594 | 1, |
| 595 | &[user_msg("alpha pattern")], |
| 596 | "deepseek-v4-pro", |
| 597 | now, |
| 598 | ) |
| 599 | .unwrap(); |
| 600 | archive_cycle( |
| 601 | &sid, |
| 602 | 2, |
| 603 | &[user_msg("alpha pattern")], |
| 604 | "deepseek-v4-pro", |
| 605 | now, |
| 606 | ) |
| 607 | .unwrap(); |
| 608 | |
| 609 | let tool = RecallArchiveTool; |
| 610 | let result = tool |
| 611 | .execute( |
| 612 | json!({"query": "alpha", "cycle": 2, "max_results": 5}), |
| 613 | &ctx, |
| 614 | ) |
| 615 | .await |
| 616 | .unwrap(); |
| 617 | assert!( |
| 618 | result.content.contains("\"cycle\":2"), |
| 619 | "got: {}", |
| 620 | result.content |
| 621 | ); |
| 622 | assert!( |
| 623 | !result.content.contains("\"cycle\":1"), |
| 624 | "got: {}", |
| 625 | result.content |
| 626 | ); |
| 627 | } |
| 628 | |
| 629 | #[tokio::test] |
| 630 | async fn execute_caps_max_results_at_hard_max() { |
| 631 | let _home = HomeGuard::new(); |
| 632 | let sid = fresh_session_id(); |
| 633 | let workspace = TempDir::new().unwrap(); |
| 634 | let ctx = ctx_for_session(workspace.path(), &sid); |
| 635 | let now = Utc::now(); |
| 636 | let mut messages: Vec<Message> = Vec::new(); |
| 637 | for i in 0..30 { |
| 638 | messages.push(user_msg(&format!("alpha message number {i}"))); |
| 639 | } |
| 640 | archive_cycle(&sid, 1, &messages, "deepseek-v4-pro", now).unwrap(); |
| 641 | |
| 642 | let tool = RecallArchiveTool; |
| 643 | let result = tool |
| 644 | .execute(json!({"query": "alpha", "max_results": 999}), &ctx) |
| 645 | .await |
| 646 | .unwrap(); |
| 647 | let count = result.content.matches("\"message_index\":").count(); |
| 648 | assert!(count <= HARD_MAX_RESULTS, "got {count} hits"); |
| 649 | } |
| 650 | |
| 651 | #[tokio::test] |
| 652 | async fn execute_rejects_empty_query() { |
| 653 | let _home = HomeGuard::new(); |
| 654 | let sid = fresh_session_id(); |
| 655 | let workspace = TempDir::new().unwrap(); |
| 656 | let ctx = ctx_for_session(workspace.path(), &sid); |
| 657 | let tool = RecallArchiveTool; |
| 658 | let err = tool |
| 659 | .execute(json!({"query": " "}), &ctx) |
| 660 | .await |
| 661 | .unwrap_err(); |
| 662 | assert!(matches!(err, ToolError::InvalidInput { .. })); |
| 663 | } |
| 664 | |
| 665 | #[test] |
| 666 | fn best_window_centers_on_first_hit() { |
| 667 | let text = "lorem ipsum dolor sit amet, the quick brown fox jumps over the lazy dog"; |
| 668 | let win = best_window(text, &["fox".to_string()], 30); |
| 669 | assert!(win.contains("fox"), "got: {win}"); |
| 670 | } |
| 671 | |
| 672 | #[test] |
| 673 | fn best_window_falls_back_to_head_when_no_hits() { |
| 674 | let text = "the quick brown fox jumps"; |
| 675 | let win = best_window(text, &["zzz".to_string()], 10); |
| 676 | assert!(win.starts_with("the quick"), "got: {win}"); |
| 677 | } |
| 678 | |
| 679 | #[test] |
| 680 | fn align_char_boundary_handles_multibyte() { |
| 681 | let text = "héllo world"; |
| 682 | // Index 2 is mid-byte for `é` (UTF-8 encoded as 2 bytes). |
| 683 | let aligned = align_char_boundary(text, 2, true); |
| 684 | assert!(text.is_char_boundary(aligned), "boundary check"); |
| 685 | } |
| 686 | |
| 687 | #[test] |
| 688 | fn bm25_returns_relevant_docs_drops_irrelevant() { |
| 689 | // BM25 length normalization can let very short matching docs outrank |
| 690 | // longer ones with higher term-frequency, so we only assert the |
| 691 | // weak invariant: matching docs are returned, non-matching docs are |
| 692 | // filtered out. |
| 693 | let docs = vec![ |
| 694 | ArchivedDoc { |
| 695 | cycle: 1, |
| 696 | message_index: 0, |
| 697 | role: "user".to_string(), |
| 698 | text: "cat dog cat dog cat".to_string(), |
| 699 | tokens: tokenize("cat dog cat dog cat"), |
| 700 | }, |
| 701 | ArchivedDoc { |
| 702 | cycle: 1, |
| 703 | message_index: 1, |
| 704 | role: "user".to_string(), |
| 705 | text: "fish bird".to_string(), |
| 706 | tokens: tokenize("fish bird"), |
| 707 | }, |
| 708 | ArchivedDoc { |
| 709 | cycle: 1, |
| 710 | message_index: 2, |
| 711 | role: "user".to_string(), |
| 712 | text: "cat sleeps".to_string(), |
| 713 | tokens: tokenize("cat sleeps"), |
| 714 | }, |
| 715 | ]; |
| 716 | let hits = score_bm25(&docs, &["cat".to_string()], 3); |
| 717 | let indices: Vec<usize> = hits.iter().map(|h| h.message_index).collect(); |
| 718 | assert!(indices.contains(&0), "doc 0 (3x cat) should appear"); |
| 719 | assert!(indices.contains(&2), "doc 2 (1x cat) should appear"); |
| 720 | assert!(!indices.contains(&1), "zero-score doc filtered"); |
| 721 | assert!(hits[0].score > 0.0, "top hit has positive score"); |
| 722 | } |
| 723 | } |
| 724 |