use anyhow::Result;
use std::sync::Arc;
use std::time::Instant;
use memrec_common::{
JsonRpcRequest, JsonRpcResponse, JsonRpcError,
protocol::{RequestAction, RequestParams, ResponseResult,
MemoryResult, MemoryListResult, SuccessResult,
StatsResult, SearchHit, SemanticSearchResult,
ProjectInfoResult, GetProjectInfoParams, VersionResult},
Memory, MemoryType,
};
use uuid::Uuid;
use crate::storage::{MemoryStorage, VectorStorage, SearchFilter};
use crate::embedding::FastEmbedGenerator;
use crate::project::{detect_project_id, find_project_root};
pub struct Router {
storage: Arc<dyn MemoryStorage>,
vector_store: Arc<dyn VectorStorage>,
embedder: Arc<FastEmbedGenerator>,
}
impl Router {
pub fn new(
storage: Arc<dyn MemoryStorage>,
vector_store: Arc<dyn VectorStorage>,
embedder: Arc<FastEmbedGenerator>,
) -> Self {
Self { storage, vector_store, embedder }
}
pub fn new_simple(storage: Arc<dyn MemoryStorage>) -> Self {
let embedder = Arc::new(FastEmbedGenerator::new().unwrap_or_default());
let vector_store = Arc::new(crate::storage::VectorStore::new(embedder.dimension()));
Self { storage, vector_store, embedder }
}
pub async fn route(&self, request: JsonRpcRequest) -> JsonRpcResponse {
match request.method {
RequestAction::Add => self.handle_add(request.params, request.id).await,
RequestAction::Get => self.handle_get(request.params, request.id).await,
RequestAction::List => self.handle_list(request.params, request.id).await,
RequestAction::Delete => self.handle_delete(request.params, request.id).await,
RequestAction::Stats => self.handle_stats(request.id).await,
RequestAction::SearchMemory => self.handle_search_memory(request.params, request.id).await,
RequestAction::GetProjectInfo => self.handle_project_info(request.params, request.id).await,
RequestAction::GetVersion => self.handle_version(request.id).await,
_ => JsonRpcResponse::error(
JsonRpcError {
code: -32601,
message: format!("Method not found: {}", request.method),
data: None,
},
request.id
)
}
}
async fn handle_add(&self, params: Option<RequestParams>, id: u64) -> JsonRpcResponse {
match params {
Some(RequestParams::Add(p)) => {
let project_id = if p.is_global {
Some(Uuid::nil())
} else {
p.project_id.or_else(|| detect_project_id(p.working_dir.as_deref()).ok())
};
let mut memory = Memory::new(p.content.clone(), p.memory_type)
.with_tags(p.tags);
if let Some(pid) = project_id {
memory = memory.with_project(pid);
}
match self.storage.save(&memory).await {
Ok(_) => {
let embedding = self.embedder.embed(&p.content).ok();
if let Some(embed) = embedding {
let payload = crate::storage::VectorPayload {
project_id: memory.project_id,
memory_type: memory.memory_type.to_string(),
tags: memory.tags.clone(),
content_preview: p.content.chars().take(200).collect(),
importance: memory.importance,
chunk_group_id: memory.chunk_group_id,
chunk_index: memory.chunk_index,
chunk_total: memory.chunk_total,
};
self.vector_store.add(&memory.id, &embed, payload).await.ok();
}
JsonRpcResponse::success(
ResponseResult::Memory(MemoryResult { memory }),
id
)
}
Err(e) => JsonRpcResponse::error(
JsonRpcError { code: -32000, message: e.to_string(), data: None },
id
)
}
}
_ => JsonRpcResponse::error(
JsonRpcError { code: -32602, message: "Invalid params for Add".to_string(), data: None },
id
)
}
}
async fn handle_get(&self, params: Option<RequestParams>, id: u64) -> JsonRpcResponse {
match params {
Some(RequestParams::Get(p)) => {
if p.merge {
self.handle_get_with_merge(&p.id, id).await
} else {
match self.storage.get(&p.id).await {
Ok(Some(memory)) => {
self.storage.update(&memory).await.ok();
JsonRpcResponse::success(
ResponseResult::Memory(MemoryResult { memory }),
id
)
}
Ok(None) => JsonRpcResponse::error(
JsonRpcError { code: -32001, message: "Memory not found".to_string(), data: None },
id
),
Err(e) => JsonRpcResponse::error(
JsonRpcError { code: -32000, message: e.to_string(), data: None },
id
)
}
}
}
_ => JsonRpcResponse::error(
JsonRpcError { code: -32602, message: "Invalid params for Get".to_string(), data: None },
id
)
}
}
async fn handle_get_with_merge(&self, id: &Uuid, rpc_id: u64) -> JsonRpcResponse {
match self.storage.get(id).await {
Ok(Some(memory)) => {
if !memory.is_chunked() {
self.storage.update(&memory).await.ok();
return JsonRpcResponse::success(
ResponseResult::Memory(MemoryResult { memory }),
rpc_id
);
}
let group_id = memory.chunk_group_id.unwrap();
match self.storage.get_chunks_by_group(&group_id).await {
Ok(chunks) => {
if chunks.is_empty() {
return JsonRpcResponse::error(
JsonRpcError { code: -32005, message: "No chunks found".to_string(), data: None },
rpc_id
);
}
let mut sorted_chunks = chunks;
sorted_chunks.sort_by_key(|c| c.chunk_index.unwrap_or(0));
let merged_content = sorted_chunks.iter()
.map(|c| c.content.as_str())
.collect::<Vec<_>>()
.join("\n");
let first_chunk = sorted_chunks.first().unwrap();
let mut merged_memory = Memory::new(merged_content, first_chunk.memory_type)
.with_tags(first_chunk.tags.clone());
merged_memory.id = group_id;
merged_memory.project_id = first_chunk.project_id;
JsonRpcResponse::success(
ResponseResult::Memory(MemoryResult { memory: merged_memory }),
rpc_id
)
}
Err(e) => JsonRpcResponse::error(
JsonRpcError { code: -32000, message: e.to_string(), data: None },
rpc_id
)
}
}
Ok(None) => JsonRpcResponse::error(
JsonRpcError { code: -32001, message: "Memory not found".to_string(), data: None },
rpc_id
),
Err(e) => JsonRpcResponse::error(
JsonRpcError { code: -32000, message: e.to_string(), data: None },
rpc_id
)
}
}
async fn handle_list(&self, params: Option<RequestParams>, id: u64) -> JsonRpcResponse {
let (limit, project_only, global_only, project_id) = match params {
Some(RequestParams::List(p)) => (p.limit, p.project_only, p.global_only, p.project_id),
_ => (20, false, false, None),
};
let memories = self.storage.list(limit * 5).await.unwrap_or_default();
let filtered: Vec<Memory> = memories.into_iter()
.filter(|m| {
if m.is_deleted {
return false;
}
if project_only {
if let Some(pid) = project_id {
m.project_id == Some(pid)
} else {
m.project_id.is_some() && !m.project_id.unwrap().is_nil()
}
} else if global_only {
m.project_id.is_none() || m.project_id.unwrap().is_nil()
} else {
true
}
})
.take(limit)
.collect();
let total = filtered.len();
JsonRpcResponse::success(
ResponseResult::MemoryList(MemoryListResult { memories: filtered, total }),
id
)
}
async fn handle_delete(&self, params: Option<RequestParams>, id: u64) -> JsonRpcResponse {
match params {
Some(RequestParams::Delete(p)) => {
match self.storage.delete(&p.id).await {
Ok(deleted) => {
let message = if deleted { "Memory hard deleted" } else { "Memory soft deleted" };
JsonRpcResponse::success(
ResponseResult::Success(SuccessResult { message: message.to_string() }),
id
)
}
Err(e) => JsonRpcResponse::error(
JsonRpcError { code: -32000, message: e.to_string(), data: None },
id
)
}
}
_ => JsonRpcResponse::error(
JsonRpcError { code: -32602, message: "Invalid params for Delete".to_string(), data: None },
id
)
}
}
async fn handle_stats(&self, id: u64) -> JsonRpcResponse {
match (self.storage.count().await, self.storage.count_deleted().await) {
(Ok(total), Ok(deleted)) => {
JsonRpcResponse::success(
ResponseResult::Stats(StatsResult {
total_memories: total + deleted,
active_memories: total,
deleted_memories: deleted,
storage_usage: 0.0,
avg_importance: 0.0,
}),
id
)
}
(Err(e), _) | (_, Err(e)) => JsonRpcResponse::error(
JsonRpcError { code: -32000, message: e.to_string(), data: None },
id
)
}
}
async fn handle_search_memory(&self, params: Option<RequestParams>, id: u64) -> JsonRpcResponse {
match params {
Some(RequestParams::SearchMemory(p)) => {
let start = Instant::now();
let embedding = match self.embedder.embed(&p.query) {
Ok(e) => e,
Err(e) => return JsonRpcResponse::error(
JsonRpcError { code: -32002, message: format!("Embedding error: {}", e), data: None },
id
),
};
let embed_time = start.elapsed().as_millis() as u64;
let project_id = if p.cross_project {
None
} else if p.global_only {
Some(Uuid::nil())
} else if p.project_only {
p.project_id.or_else(|| detect_project_id(p.working_dir.as_deref()).ok())
} else {
p.project_id.or_else(|| detect_project_id(p.working_dir.as_deref()).ok())
};
let include_global = !p.project_only && !p.cross_project;
let filter = SearchFilter {
project_id,
include_global,
memory_type: p.memory_type.map(|t| t.to_string()),
min_score: p.min_score,
};
let search_start = Instant::now();
let hits = match self.vector_store.search(&embedding, filter, p.top_k).await {
Ok(h) => h,
Err(e) => return JsonRpcResponse::error(
JsonRpcError { code: -32003, message: format!("Search error: {}", e), data: None },
id
),
};
let search_time = search_start.elapsed().as_millis() as u64;
let mut results: Vec<SearchHit> = vec![];
for h in hits {
let memory = self.storage.get(&h.memory_id).await.ok().flatten();
let memory_type = memory.as_ref().map(|m| m.memory_type).unwrap_or(MemoryType::Conversation);
let created_at = memory.map(|m| m.created_at).unwrap_or_default();
results.push(SearchHit {
memory_id: h.memory_id,
score: h.score,
memory_type,
content_preview: h.payload.content_preview.clone(),
project_id: h.payload.project_id,
tags: h.payload.tags.clone(),
is_chunked: h.payload.chunk_group_id.is_some(),
chunk_group_id: h.payload.chunk_group_id,
chunk_index: h.payload.chunk_index,
chunk_total: h.payload.chunk_total,
created_at,
});
}
let total = results.len();
JsonRpcResponse::success(
ResponseResult::SemanticSearchResult(SemanticSearchResult {
results,
total,
query_embedding_time_ms: embed_time,
search_time_ms: search_time,
}),
id
)
}
_ => JsonRpcResponse::error(
JsonRpcError { code: -32602, message: "Invalid params for SearchMemory".to_string(), data: None },
id
)
}
}
async fn handle_project_info(&self, params: Option<RequestParams>, id: u64) -> JsonRpcResponse {
let params: GetProjectInfoParams = match params {
Some(RequestParams::GetProjectInfo(p)) => p,
_ => GetProjectInfoParams::default(),
};
match detect_project_id(params.working_dir.as_deref()) {
Ok(project_id) => {
let project_root = find_project_root(params.working_dir.as_deref())
.unwrap_or_default()
.to_string_lossy()
.to_string();
let mr_pid_path = std::path::Path::new(&project_root).join(".mr_pid");
let mr_pid_exists = mr_pid_path.exists();
let memory_count = self.storage.list_by_project(&project_id).await
.unwrap_or_default()
.len();
JsonRpcResponse::success(
ResponseResult::ProjectInfo(ProjectInfoResult {
project_id,
project_name: None,
project_root,
memory_count,
mr_pid_exists,
}),
id
)
}
Err(e) => JsonRpcResponse::error(
JsonRpcError { code: -32004, message: e.to_string(), data: None },
id
)
}
}
async fn handle_version(&self, id: u64) -> JsonRpcResponse {
JsonRpcResponse::success(
ResponseResult::Version(VersionResult {
version: env!("CARGO_PKG_VERSION").to_string(),
}),
id
)
}
pub fn parse_request(&self, raw: &str) -> Result<JsonRpcRequest> {
serde_json::from_str(raw)
.map_err(|e| anyhow::anyhow!("Failed to parse request: {}", e))
}
pub fn serialize_response(&self, response: &JsonRpcResponse) -> Result<String> {
serde_json::to_string(response)
.map_err(|e| anyhow::anyhow!("Failed to serialize response: {}", e))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::{MemoryStore, VectorStore};
use crate::storage::rocksdb::RocksDBStore;
use tempfile::tempdir;
use memrec_common::protocol::{SearchMemoryParams, GetProjectInfoParams};
#[tokio::test]
async fn test_router_search_memory() {
let dir = tempdir().unwrap();
let rocksdb = RocksDBStore::open(dir.path()).unwrap();
let storage = Arc::new(MemoryStore::new(rocksdb));
let embedder = Arc::new(FastEmbedGenerator::new().unwrap());
let vector_store = Arc::new(VectorStore::new(embedder.dimension()));
let router = Router::new(storage, vector_store, embedder);
let request = JsonRpcRequest::new(
RequestAction::SearchMemory,
Some(RequestParams::SearchMemory(SearchMemoryParams {
query: "test query".to_string(),
project_id: None,
include_global: true,
project_only: false,
global_only: false,
cross_project: false,
memory_type: None,
top_k: 10,
min_score: 0.0,
working_dir: None,
})),
1,
);
let response = router.route(request).await;
assert!(response.result.is_some());
}
#[tokio::test]
async fn test_router_project_info() {
let dir = tempdir().unwrap();
let rocksdb = RocksDBStore::open(dir.path()).unwrap();
let storage = Arc::new(MemoryStore::new(rocksdb));
let router = Router::new_simple(storage);
let request = JsonRpcRequest::new(
RequestAction::GetProjectInfo,
Some(RequestParams::GetProjectInfo(GetProjectInfoParams::default())),
1,
);
let response = router.route(request).await;
assert!(response.result.is_some() || response.error.is_some());
}
}