use std::collections::HashMap;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::{tool, tool_router};
use super::super::types::{GetObservationsParams, TimelineParams};
use super::errors::{self, McpToolError, McpToolResult};
use super::MemoryServer;
use crate::retrieval::search;
use crate::{db, memory};
#[tool_router(router = tool_router_context, vis = "pub(super)")]
impl MemoryServer {
#[tool(
description = "Get chronological observations around a specific point. Useful for understanding what happened before/after a change. Provide anchor ID or search query to find the center point."
)]
pub(super) fn timeline(
&self,
Parameters(params): Parameters<TimelineParams>,
) -> McpToolResult<String> {
const TOOL: &str = "timeline";
let start = std::time::Instant::now();
crate::log::info(
"mcp",
&format!(
"timeline called anchor={:?} query={:?} project={:?} before={} after={}",
params.anchor,
params.query,
params.project,
params.depth_before.unwrap_or(5),
params.depth_after.unwrap_or(5),
),
);
self.with_conn(TOOL, |conn| {
let anchor_id = if let Some(id) = params.anchor {
let anchor = db::get_observations_by_ids(conn, &[id], params.project.as_deref())
.map_err(|e| {
crate::log::warn("mcp", &format!("timeline anchor lookup failed: {}", e));
McpToolError::db_query(TOOL, e)
})?;
if anchor.is_empty() {
return Err(McpToolError::not_found(
TOOL,
format!("No observation found for anchor id {id}"),
));
}
id
} else if let Some(query) = ¶ms.query {
let results = search::search(
conn,
Some(query),
params.project.as_deref(),
None,
1,
0,
true,
)
.map_err(|e| {
crate::log::warn("mcp", &format!("timeline search failed: {}", e));
McpToolError::db_query(TOOL, e)
})?;
results
.first()
.map(|observation| observation.id)
.ok_or_else(|| {
McpToolError::not_found(TOOL, format!("No results for query '{query}'"))
})?
} else {
return Err(McpToolError::invalid_request(
TOOL,
"anchor or query required",
));
};
let results = db::get_timeline_around(
conn,
anchor_id,
params.depth_before.unwrap_or(5),
params.depth_after.unwrap_or(5),
params.project.as_deref(),
)
.map_err(|e| {
crate::log::warn("mcp", &format!("timeline failed: {}", e));
McpToolError::db_query(TOOL, e)
})?;
crate::log::info(
"mcp",
&format!(
"timeline done anchor={} count={} {}ms",
anchor_id,
results.len(),
start.elapsed().as_millis()
),
);
errors::to_json_pretty(TOOL, &results)
})
}
#[tool(
description = "Fetch complete details by IDs. Use after search(): pass selected IDs and the exact source from search.next_step.source or each result.source. Supports source='memory' for curated memories and source='observation' for legacy observations."
)]
pub(super) fn get_observations(
&self,
Parameters(params): Parameters<GetObservationsParams>,
) -> McpToolResult<String> {
const TOOL: &str = "get_observations";
let start = std::time::Instant::now();
let source = params.source.as_deref().unwrap_or("memory");
crate::log::info(
"mcp",
&format!(
"get_observations called ids={:?} project={:?} source={}",
params.ids, params.project, source
),
);
self.with_conn(TOOL, |conn| {
let results = match source {
"observation" => {
let observations_result =
db::get_observations_by_ids(conn, ¶ms.ids, params.project.as_deref());
let observations = observations_result.map_err(|e| {
crate::log::warn("mcp", &format!("get_observations failed: {}", e));
McpToolError::db_query(TOOL, e)
})?;
ensure_requested_ids_found(
TOOL,
source,
¶ms.ids,
observations.iter().map(|observation| observation.id),
)?;
let accessed_ids: Vec<i64> = observations
.iter()
.map(|observation| observation.id)
.collect();
if !accessed_ids.is_empty() {
if let Err(err) = db::update_last_accessed(conn, &accessed_ids) {
crate::log::warn(
"mcp",
&format!("update_last_accessed failed: {}", err),
);
}
}
observation_details_with_compressed_sources(conn, &observations).map_err(
|e| {
crate::log::warn(
"mcp",
&format!("load compressed observation sources failed: {}", e),
);
McpToolError::db_query(TOOL, e)
},
)?
}
"memory" => {
let memories_result =
memory::get_memories_by_ids(conn, ¶ms.ids, params.project.as_deref());
let memories = memories_result.map_err(|e| {
crate::log::warn("mcp", &format!("get_memories failed: {}", e));
McpToolError::db_query(TOOL, e)
})?;
ensure_requested_ids_found(
TOOL,
source,
¶ms.ids,
memories.iter().map(|memory| memory.id),
)?;
memory_details_with_topic_traces(conn, &memories, params.project.as_deref())
.map_err(|e| {
crate::log::warn("mcp", &format!("load topic traces failed: {}", e));
McpToolError::db_query(TOOL, e)
})?
}
other => {
return Err(McpToolError::unsupported_source(
TOOL,
format!("unsupported source '{other}'; expected 'memory' or 'observation'"),
));
}
};
crate::log::info(
"mcp",
&format!(
"get_observations done source={} count={} {}ms",
source,
params.ids.len(),
start.elapsed().as_millis()
),
);
errors::to_json_pretty(TOOL, &results)
})
}
}
fn observation_details_with_compressed_sources(
conn: &rusqlite::Connection,
observations: &[db::Observation],
) -> anyhow::Result<serde_json::Value> {
let mut value = serde_json::to_value(observations)?;
let Some(items) = value.as_array_mut() else {
return Ok(value);
};
let observation_ids: Vec<i64> = observations
.iter()
.map(|observation| observation.id)
.collect();
let sources_by_observation = db::load_compressed_observation_sources(conn, &observation_ids)?;
for (item, observation) in items.iter_mut().zip(observations) {
let Some(sources) = sources_by_observation.get(&observation.id) else {
continue;
};
if !sources.is_empty() {
item["compressed_sources"] = serde_json::to_value(sources)?;
}
}
Ok(value)
}
fn memory_details_with_topic_traces(
conn: &rusqlite::Connection,
memories: &[memory::Memory],
requested_project: Option<&str>,
) -> anyhow::Result<serde_json::Value> {
const TRACE_LIMIT: i64 = 12;
let mut value = serde_json::to_value(memories)?;
let Some(items) = value.as_array_mut() else {
return Ok(value);
};
let mut trace_cache = HashMap::new();
for (item, memory) in items.iter_mut().zip(memories) {
let Some(topic_key) = memory.topic_key.as_deref() else {
continue;
};
let trace_project = match requested_project {
Some(project) if project == memory.project => project,
Some(_) => continue,
None => memory.project.as_str(),
};
let cache_key = (trace_project.to_string(), topic_key.to_string());
if !trace_cache.contains_key(&cache_key) {
let trace = db::load_trace_by_topic_key(conn, trace_project, topic_key, TRACE_LIMIT)?;
trace_cache.insert(cache_key.clone(), trace);
}
let trace = trace_cache
.get(&cache_key)
.expect("trace cache should contain loaded key");
if !trace.is_empty() {
item["topic_trace"] = serde_json::to_value(trace)?;
}
}
Ok(value)
}
fn ensure_requested_ids_found(
tool: &'static str,
source: &str,
requested_ids: &[i64],
found_ids: impl Iterator<Item = i64>,
) -> McpToolResult<()> {
if requested_ids.is_empty() {
return Ok(());
}
let found: std::collections::HashSet<i64> = found_ids.collect();
let missing: Vec<i64> = requested_ids
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>()
.into_iter()
.filter(|id| !found.contains(id))
.collect();
if missing.is_empty() {
return Ok(());
}
Err(McpToolError::not_found(
tool,
format!("{source} id(s) not found: {missing:?}"),
))
}
#[cfg(test)]
mod tests {
use anyhow::Result;
use rusqlite::Connection;
use super::*;
fn memory_row(project: &str, scope: &str) -> crate::memory::Memory {
crate::memory::Memory {
id: 1,
session_id: Some("session-1".to_string()),
project: project.to_string(),
topic_key: Some("global-contract".to_string()),
title: "Global contract".to_string(),
text: "Global memory body".to_string(),
memory_type: "decision".to_string(),
files: None,
created_at_epoch: 10,
updated_at_epoch: 10,
status: "active".to_string(),
branch: None,
scope: scope.to_string(),
}
}
fn conn_with_trace() -> Result<Connection> {
let conn = Connection::open_in_memory()?;
crate::migrate::run_migrations(&conn)?;
conn.execute_batch("PRAGMA foreign_keys=OFF;")?;
crate::db::insert_topic_segment(
&conn,
&crate::db::TopicSegmentInput {
host_id: 1,
project_id: 1,
session_row_id: 1,
project: "/repo",
topic_key: "global-contract",
title: "Repo-only trace",
summary: "Trace belongs to /repo.",
status: "resolved",
segment_index: 0,
covered_from_event_id: 10,
covered_to_event_id: 12,
evidence_event_ids: "[10,12]",
files: None,
confidence: 0.8,
},
)?;
Ok(conn)
}
#[test]
fn topic_trace_is_omitted_for_requested_project_mismatch() -> Result<()> {
let conn = conn_with_trace()?;
let memories = vec![memory_row("/repo", "global")];
let value = memory_details_with_topic_traces(&conn, &memories, Some("/other"))?;
assert_eq!(value[0]["scope"], "global");
assert!(value[0]["topic_trace"].is_null());
Ok(())
}
#[test]
fn topic_trace_is_attached_for_matching_requested_project() -> Result<()> {
let conn = conn_with_trace()?;
let memories = vec![memory_row("/repo", "project")];
let value = memory_details_with_topic_traces(&conn, &memories, Some("/repo"))?;
assert_eq!(value[0]["topic_trace"][0]["title"], "Repo-only trace");
Ok(())
}
}