diff --git a/server/src/handlers/meta.rs b/server/src/handlers/meta.rs index 6db96fa..a13f920 100644 --- a/server/src/handlers/meta.rs +++ b/server/src/handlers/meta.rs @@ -361,14 +361,28 @@ impl McpTool for OmniSearchHandler { serde_json::to_value(&kg_entities).unwrap_or_default() }); + let mut matched_tasks = std::collections::HashSet::new(); + let mut matched_snippets = std::collections::HashSet::new(); + let mut matched_adrs = std::collections::HashSet::new(); + for (id, typ, _, _, _) in &matches { + match typ.as_str() { + "task" => { + matched_tasks.insert(id.as_str()); + } + "snippet" => { + matched_snippets.insert(id.as_str()); + } + "adr" => { + matched_adrs.insert(id.as_str()); + } + _ => {} + } + } + let tasks_json = state.tasks.read_with(|all_tasks| { let filtered: Vec<_> = all_tasks .iter() - .filter(|t| { - matches - .iter() - .any(|(id, typ, _, _, _)| id == &t.id && typ == "task") - }) + .filter(|t| matched_tasks.contains(t.id.as_str())) .collect(); serde_json::to_value(&filtered).unwrap_or_default() }); @@ -376,11 +390,7 @@ impl McpTool for OmniSearchHandler { let snippets_json = state.snippets.read_with(|all_snippets| { let filtered: Vec<_> = all_snippets .iter() - .filter(|s| { - matches - .iter() - .any(|(id, typ, _, _, _)| id == &s.name && typ == "snippet") - }) + .filter(|s| matched_snippets.contains(s.name.as_str())) .collect(); serde_json::to_value(&filtered).unwrap_or_default() }); @@ -388,11 +398,7 @@ impl McpTool for OmniSearchHandler { let adrs_json = state.adrs.read_with(|all_adrs| { let filtered: Vec<_> = all_adrs .iter() - .filter(|a| { - matches - .iter() - .any(|(id, typ, _, _, _)| id == &a.id && typ == "adr") - }) + .filter(|a| matched_adrs.contains(a.id.as_str())) .collect(); serde_json::to_value(&filtered).unwrap_or_default() }); diff --git a/server/src/main.rs b/server/src/main.rs index 5bf7ca2..289758e 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -115,16 +115,22 @@ async fn run_server(state: Arc) -> Result<(), Box = app_state_clone - .clients - .read() - .unwrap_or_else(|e| e.into_inner()) - .values() - .cloned() - .collect(); - for client_tx in senders { - let _ = client_tx.try_send(msg.clone()); + loop { + match rx.recv().await { + Ok(msg) => { + let senders: Vec<_> = app_state_clone + .clients + .read() + .unwrap_or_else(|e| e.into_inner()) + .values() + .cloned() + .collect(); + for client_tx in senders { + let _ = client_tx.try_send(msg.clone()); + } + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue, } } }); diff --git a/server/src/state.rs b/server/src/state.rs index 28694ff..3084cf5 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -124,26 +124,27 @@ impl MemoryState { let idx = new_idx.clone(); tokio::task::spawn_blocking(move || { - state.graph.read_with(|g| { - for e in g.entities.values() { - idx.add_entity_sync(e); - } - }); - state.tasks.read_with(|tasks| { - for t in tasks { - idx.add_task_sync(t); - } - }); - state.snippets.read_with(|snippets| { - for s in snippets { - idx.add_snippet_sync(s); - } - }); - state.adrs.read_with(|adrs| { - for a in adrs { - idx.add_adr_sync(a); - } - }); + let entities: Vec<_> = state + .graph + .read_with(|g| g.entities.values().cloned().collect()); + for e in entities { + idx.add_entity_sync(&e); + } + + let tasks = state.tasks.read_with(|t| t.clone()); + for t in tasks { + idx.add_task_sync(&t); + } + + let snippets = state.snippets.read_with(|s| s.clone()); + for s in snippets { + idx.add_snippet_sync(&s); + } + + let adrs = state.adrs.read_with(|a| a.clone()); + for a in adrs { + idx.add_adr_sync(&a); + } }) .await .unwrap_or_else(|e| { diff --git a/server/src/store.rs b/server/src/store.rs index e755777..35f06e2 100644 --- a/server/src/store.rs +++ b/server/src/store.rs @@ -28,18 +28,20 @@ impl let key_inner = key_clone.clone(); let json_data = { let lock = cache_clone.read().unwrap_or_else(|e| e.into_inner()); - serde_json::to_vec(&*lock).unwrap_or_default() + serde_json::to_vec(&*lock).ok() }; - let _ = tokio::task::spawn_blocking(move || { - if let Ok(write_txn) = db_inner.begin_write() { - if let Ok(mut table) = write_txn.open_table(STORE_TABLE) { - let _ = table.insert(key_inner.as_str(), json_data.as_slice()); + if let Some(json_data) = json_data { + let _ = tokio::task::spawn_blocking(move || { + if let Ok(write_txn) = db_inner.begin_write() { + if let Ok(mut table) = write_txn.open_table(STORE_TABLE) { + let _ = table.insert(key_inner.as_str(), json_data.as_slice()); + } + let _ = write_txn.commit(); } - let _ = write_txn.commit(); - } - }) - .await; + }) + .await; + } } });