use async_trait::async_trait;
use serde_json::{json, Value};
use std::path::PathBuf;
use crate::error::{Error, Result};
use crate::llm::ToolSpec;
use crate::session::SessionReader;
use crate::tools::Tool;
pub struct EpisodicRecall {
workspace: PathBuf,
}
impl EpisodicRecall {
pub fn new(workspace: impl Into<PathBuf>) -> Self {
Self {
workspace: workspace.into(),
}
}
}
#[async_trait]
impl Tool for EpisodicRecall {
fn spec(&self) -> ToolSpec {
ToolSpec {
name: "episodic_recall".into(),
description: "Search past session transcripts for messages matching a query. Reads JSONL session files from the workspace's session directory. Returns matching entries with surrounding context.".into(),
parameters: json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search query (case-insensitive substring match against message content, role, and tool call names/arguments)"
},
"session_id": {
"type": "string",
"description": "Optional session ID to restrict the search to a single session. If omitted, searches all sessions."
},
"limit": {
"type": "integer",
"description": "Maximum number of matching entries to return (default 5)",
"default": 5
},
"context_lines": {
"type": "integer",
"description": "Number of surrounding messages to include before and after each match (default 2)",
"default": 2
}
},
"required": ["query"]
}),
}
}
fn side_effect_class(&self) -> crate::tools::ToolSideEffect {
crate::tools::ToolSideEffect::ReadOnly
}
async fn execute(&self, arguments: Value) -> Result<String> {
let query = arguments["query"]
.as_str()
.ok_or_else(|| Error::BadToolArgs {
name: "episodic_recall".into(),
message: "missing required parameter: query".to_string(),
})?;
let session_id = arguments["session_id"].as_str().map(|s| s.to_string());
let limit = arguments["limit"].as_i64().unwrap_or(5) as usize;
let context_lines = arguments["context_lines"].as_i64().unwrap_or(2) as usize;
let all_sessions = match SessionReader::list_sessions(&self.workspace) {
Ok(sessions) => sessions,
Err(e) => {
return Ok(format!("failed to list sessions: {e}"));
}
};
if all_sessions.is_empty() {
return Ok("no past sessions found".to_string());
}
let sessions: Vec<&PathBuf> = if let Some(ref sid) = session_id {
all_sessions
.iter()
.filter(|dir| {
let dir_name = dir.file_name().and_then(|n| n.to_str()).unwrap_or("");
dir_name == sid.as_str() || dir_name.ends_with(sid.as_str())
})
.collect()
} else {
all_sessions.iter().collect()
};
if sessions.is_empty() {
return Ok(format!(
"no sessions found matching '{}'",
session_id.unwrap_or_default()
));
}
let query_lower = query.to_lowercase();
let mut results: Vec<String> = Vec::new();
for session_dir in sessions {
let entries = match SessionReader::load_transcript(session_dir) {
Ok(e) => e,
Err(_) => continue,
};
if entries.is_empty() {
continue;
}
let mut match_indices: Vec<usize> = Vec::new();
for (i, entry) in entries.iter().enumerate() {
if entry.role.to_lowercase().contains(&query_lower) {
match_indices.push(i);
continue;
}
if entry.content.to_lowercase().contains(&query_lower) {
match_indices.push(i);
continue;
}
for tc in &entry.tool_calls {
if tc.name.to_lowercase().contains(&query_lower)
|| tc
.arguments
.to_string()
.to_lowercase()
.contains(&query_lower)
{
match_indices.push(i);
break;
}
}
}
if match_indices.is_empty() {
continue;
}
match_indices.sort();
match_indices.dedup();
let session_name = session_dir
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("unknown");
results.push(format!("--- Session: {} ---", session_name));
let mut rendered_indices: std::collections::BTreeSet<usize> =
std::collections::BTreeSet::new();
for &idx in &match_indices {
let start = idx.saturating_sub(context_lines);
let end = (idx + context_lines + 1).min(entries.len());
for i in start..end {
rendered_indices.insert(i);
}
}
for &i in &rendered_indices {
let entry = &entries[i];
let marker = if match_indices.contains(&i) {
" >>> "
} else {
" "
};
let role_padded = format!("{:>9}", entry.role);
let content_preview = if entry.content.len() > 200 {
format!("{}...", crate::truncate_str(&entry.content, 197))
} else {
entry.content.clone()
};
let tool_info = if !entry.tool_calls.is_empty() {
let names: Vec<&str> =
entry.tool_calls.iter().map(|tc| tc.name.as_str()).collect();
format!(" [tool_calls: {}]", names.join(", "))
} else {
String::new()
};
results.push(format!(
"{}[{}] {}{} {}",
marker, entry.id, role_padded, tool_info, content_preview
));
}
results.push(String::new()); }
if results.is_empty() {
return Ok(format!("no matches found for '{}'", query));
}
let mut output_lines: Vec<String> = Vec::new();
let mut match_count = 0;
for line in results {
if line.starts_with(" >>> ") {
match_count += 1;
if match_count > limit {
output_lines.push(format!(
"... (truncated, showing {}/{} matches)",
limit, match_count
));
break;
}
}
output_lines.push(line);
}
Ok(output_lines.join("\n"))
}
}
pub fn episodic_recall_summary(workspace: &std::path::Path, limit: usize) -> String {
let sessions = match SessionReader::list_sessions(workspace) {
Ok(s) => s,
Err(_) => return String::new(),
};
if sessions.is_empty() {
return String::new();
}
let mut lines: Vec<String> = Vec::new();
lines.push(format!(
"# Recent Sessions (last {}; use `episodic_recall` to search)",
limit
));
let recent: Vec<&PathBuf> = sessions.iter().rev().take(limit).collect();
for session_dir in recent {
let meta = match SessionReader::load_meta(session_dir) {
Ok(m) => m,
Err(_) => continue,
};
let session_name = session_dir
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("unknown");
let goal_preview = if meta.goal.len() > 80 {
format!("{}...", crate::truncate_str(&meta.goal, 77))
} else {
meta.goal.clone()
};
lines.push(format!(
"- {} ({} msgs, {}): {}",
session_name, meta.message_count, meta.status, goal_preview
));
}
lines.join("\n")
}
#[cfg(test)]
mod tests {
use super::*;
use crate::message::Message;
use crate::session::SessionWriter;
use std::time::Duration;
fn setup_session_with_messages(
ws: &std::path::Path,
goal: &str,
messages: &[(&str, &str)],
) -> String {
let mut writer = SessionWriter::create(ws, goal, "test-model", "test-provider").unwrap();
let session_dir = writer.session_dir().to_path_buf();
for (role, content) in messages {
let msg = match *role {
"user" => Message::user(content.to_string()),
"assistant" => Message::assistant(content.to_string()),
"system" => Message::system(content.to_string()),
_ => Message::user(content.to_string()),
};
writer.append(&msg, None, None).unwrap();
}
writer.finish("completed").unwrap();
session_dir
.file_name()
.unwrap()
.to_str()
.unwrap()
.to_string()
}
#[test]
fn test_a_episodic_recall_finds_messages() {
let tmp = crate::test_util::IsolatedWorkspace::new();
let ws = tmp.path();
setup_session_with_messages(
ws,
"test goal",
&[
("system", "You are a helpful assistant."),
("user", "Hello, how are you?"),
("assistant", "I'm doing great, thanks!"),
("user", "What is Rust?"),
("assistant", "Rust is a systems programming language."),
],
);
let tool = EpisodicRecall::new(ws);
let result = tokio::runtime::Runtime::new()
.unwrap()
.block_on(tool.execute(json!({"query": "Rust"})));
assert!(result.is_ok());
let output = result.unwrap();
assert!(output.contains("Rust"), "output: {output}");
assert!(output.contains("systems programming"), "output: {output}");
}
#[test]
fn test_b_episodic_recall_with_session_filter() {
let tmp = crate::test_util::IsolatedWorkspace::new();
let ws = tmp.path();
let sid1 = setup_session_with_messages(
ws,
"first-session-unique-goal",
&[("user", "Hello from session 1")],
);
std::thread::sleep(Duration::from_secs(1));
let sid2 = setup_session_with_messages(
ws,
"second-session-unique-goal",
&[("user", "Hello from session 2")],
);
assert_ne!(sid1, sid2, "session IDs should differ");
let tool = EpisodicRecall::new(ws);
let result = tokio::runtime::Runtime::new()
.unwrap()
.block_on(tool.execute(json!({
"query": "Hello",
"session_id": &sid2
})));
assert!(result.is_ok());
let output = result.unwrap();
assert!(output.contains("session 2"), "output: {output}");
assert!(!output.contains("session 1"), "output: {output}");
}
#[test]
fn test_c_episodic_recall_empty_when_no_sessions() {
let tmp = crate::test_util::IsolatedWorkspace::new();
let ws = tmp.path();
let tool = EpisodicRecall::new(ws);
let result = tokio::runtime::Runtime::new()
.unwrap()
.block_on(tool.execute(json!({"query": "anything"})));
assert!(result.is_ok());
let output = result.unwrap();
assert_eq!(output, "no past sessions found");
}
#[test]
fn test_d_episodic_recall_summary() {
let tmp = crate::test_util::IsolatedWorkspace::new();
let ws = tmp.path();
let summary = episodic_recall_summary(ws, 5);
assert_eq!(summary, "");
setup_session_with_messages(
ws,
"fix the bug in parser",
&[("user", "Please fix the parser bug")],
);
let summary = episodic_recall_summary(ws, 5);
assert!(!summary.is_empty(), "summary should not be empty");
assert!(
summary.contains("fix the bug in parser"),
"summary: {summary}"
);
assert!(summary.contains("1 msgs"), "summary: {summary}");
}
}