diff --git a/Cargo.lock b/Cargo.lock index 03f0c588b..eca2967d3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -720,7 +720,7 @@ checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" [[package]] name = "codegraph" -version = "2.1.6" +version = "2.1.7" dependencies = [ "anyhow", "camino", @@ -743,7 +743,7 @@ dependencies = [ [[package]] name = "codegraph-api" -version = "2.1.6" +version = "2.1.7" dependencies = [ "anyhow", "camino", @@ -760,7 +760,7 @@ dependencies = [ [[package]] name = "codegraph-bench" -version = "2.1.6" +version = "2.1.7" dependencies = [ "anyhow", "camino", @@ -778,7 +778,7 @@ dependencies = [ [[package]] name = "codegraph-binary" -version = "2.1.6" +version = "2.1.7" dependencies = [ "camino", "codegraph-core", @@ -795,7 +795,7 @@ dependencies = [ [[package]] name = "codegraph-context" -version = "2.1.6" +version = "2.1.7" dependencies = [ "codegraph-core", "codegraph-graph", @@ -807,7 +807,7 @@ dependencies = [ [[package]] name = "codegraph-core" -version = "2.1.6" +version = "2.1.7" dependencies = [ "async-graphql", "camino", @@ -818,7 +818,7 @@ dependencies = [ [[package]] name = "codegraph-docs" -version = "2.1.6" +version = "2.1.7" dependencies = [ "anyhow", "codegraph-core", @@ -835,7 +835,7 @@ dependencies = [ [[package]] name = "codegraph-extract" -version = "2.1.6" +version = "2.1.7" dependencies = [ "camino", "codegraph-binary", @@ -874,7 +874,7 @@ dependencies = [ [[package]] name = "codegraph-graph" -version = "2.1.6" +version = "2.1.7" dependencies = [ "async-trait", "bincode", @@ -904,7 +904,7 @@ dependencies = [ [[package]] name = "codegraph-graphql" -version = "2.1.6" +version = "2.1.7" dependencies = [ "anyhow", "async-graphql", @@ -927,7 +927,7 @@ dependencies = [ [[package]] name = "codegraph-installer" -version = "2.1.6" +version = "2.1.7" dependencies = [ "anyhow", "camino", @@ -943,7 +943,7 @@ dependencies = [ [[package]] name = "codegraph-mcp" -version = "2.1.6" +version = "2.1.7" dependencies = [ "anyhow", "axum", @@ -966,7 +966,7 @@ dependencies = [ [[package]] name = "codegraph-sboxes" -version = "2.1.6" +version = "2.1.7" dependencies = [ "camino", "codegraph-core", diff --git a/Cargo.toml b/Cargo.toml index d2af0bb00..a3f168eaf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,7 @@ members = [ ] [workspace.package] -version = "2.1.6" +version = "2.1.7" edition = "2021" rust-version = "1.80" license = "MIT" diff --git a/crates/codegraph-docs/src/graph.rs b/crates/codegraph-docs/src/graph.rs index 75d646afb..86560c796 100644 --- a/crates/codegraph-docs/src/graph.rs +++ b/crates/codegraph-docs/src/graph.rs @@ -36,7 +36,9 @@ const DEFAULT_SHARDING: usize = 64; pub struct DocumentGraph { storage: Arc>, docs: HashMap, - nodes: HashMap, + /// Cache node đọc theo nhu cầu (`node()`) — KHÔNG load sẵn toàn bộ khi + /// open. Mutex (không tokio) vì chỉ giữ trong RAM, không span await. + nodes: std::sync::Mutex>, intern: Interner, path_trie: Search, type_trie: Search, @@ -61,7 +63,7 @@ impl DocumentGraph { Self { storage: storage.clone(), docs: HashMap::new(), - nodes: HashMap::new(), + nodes: std::sync::Mutex::new(HashMap::new()), intern: Interner::new(), path_trie: Search::new(sharding, storage.clone()), type_trie: Search::new(sharding, storage.clone()), @@ -74,14 +76,39 @@ impl DocumentGraph { } } - /// Open an existing graph from persistent storage and rebuild the tries. + /// Open an existing graph from persistent storage — **lazy**: chỉ load + /// doc list + doc metadata (số lượng file, nhỏ) và resume id counters. + /// KHÔNG materialize toàn bộ node metadata — 186k nodes ở repo document + /// lớn làm open chờ hàng chục giây. Node được đọc **theo nhu cầu** từng + /// cái (`node()` — hydrate/collect_path), có cache LRU ở storage layer + /// (`CachedStorage`) và cache in-memory trong `self.nodes`. pub async fn open(storage: Arc>, config: DocConfig) -> Result { let mut graph = Self::new(storage, config); - graph.rebuild().await?; + // Load docs list + metadata (theo doc, không theo node). + let doc_ids = graph.load_doc_list().await?; + for id in doc_ids { + let bytes = { + let guard = graph.storage.read().await; + guard.get_node_meta((DOC_META_BASE + id) as usize).await? + }; + if let Some(bytes) = bytes + && let Ok(doc) = serde_json::from_slice::(&bytes) + { + graph.docs.insert(doc.id, doc); + } + } // Resume id counters từ trạng thái đã persist — reset về `doc_base` - // sẽ đè lên id cũ khi ingest tiếp. + // sẽ đè lên id cũ khi ingest tiếp. next_node_id lấy từ max node id + // trong chain (1 lần đọc chain id, không đọc từng meta). let max_doc = graph.docs.keys().copied().max().unwrap_or(0); - let max_node = graph.nodes.keys().copied().max().unwrap_or(0); + let max_node = { + let guard = graph.storage.read().await; + guard + .get_chain(DOC_NODE_LIST_RECORD as usize) + .await? + .map(|c| c.iter().copied().max().unwrap_or(0)) + .unwrap_or(0) + }; graph.next_doc_id = graph.next_doc_id.max(max_doc + 1); graph.next_node_id = graph.next_node_id.max(max_node + 1); Ok(graph) @@ -104,62 +131,6 @@ impl DocumentGraph { self.upsert_document(doc).await } - /// Rebuild all materialized tries from persisted node/doc metadata. - pub async fn rebuild(&mut self) -> Result<()> { - // Load node list. - let node_ids = { - let guard = self.storage.read().await; - if let Some(chain) = guard.get_chain(DOC_NODE_LIST_RECORD as usize).await? { - chain.to_vec() - } else { - Vec::new() - } - }; - // Load docs list. - let doc_ids = { - let guard = self.storage.read().await; - if let Some(chain) = guard.get_chain(DOC_LIST_RECORD as usize).await? { - chain.to_vec() - } else { - Vec::new() - } - }; - // Load nodes. - for id in &node_ids { - let bytes = { - let guard = self.storage.read().await; - guard.get_node_meta(*id as usize).await? - }; - if let Some(bytes) = bytes - && let Ok(node) = serde_json::from_slice::(&bytes) - { - self.nodes.insert(node.id, node); - } - } - // Load docs. - for id in &doc_ids { - let meta_id = DOC_META_BASE + id; - let bytes = { - let guard = self.storage.read().await; - guard.get_node_meta(meta_id as usize).await? - }; - if let Some(bytes) = bytes - && let Ok(doc) = serde_json::from_slice::(&bytes) - { - self.docs.insert(doc.id, doc); - } - } - // Rebuild tries (in-memory từ node metadata). KHÔNG dùng `Search::clear` - // — nó xoá toàn bộ `clear_node_meta`/`clear_chains` của storage, xoá cả - // node/doc JSON vừa đọc lên (tries của docs start rỗng từ `new()` nên - // không cần clear persistent state). - let nodes: Vec = self.nodes.values().cloned().collect(); - for node in nodes { - self.insert_node_into_tries(&node).await?; - } - Ok(()) - } - /// Ingest a document, replacing any previous version with the same id. pub async fn upsert_document(&mut self, mut doc: Document) -> Result { if let Some(old) = self.docs.get(&doc.id) { @@ -199,9 +170,13 @@ impl DocumentGraph { for node in &doc.nodes { self.insert_node_into_tries(node).await?; } - // Materialize nodes vào map in-memory (hydrate/stats đọc từ đây). - for node in &doc.nodes { - self.nodes.insert(node.id, node.clone()); + // Materialize nodes vào cache in-memory (hydrate/collect_path đọc từ + // đây trước, thiếu thì mới xuống storage). + { + let mut cache = self.nodes.lock().unwrap(); + for node in &doc.nodes { + cache.insert(node.id, node.clone()); + } } self.docs.insert(doc_id, doc.clone()); Ok(doc_id) @@ -221,15 +196,41 @@ impl DocumentGraph { Ok(()) } + /// Đọc một node theo nhu cầu: cache in-memory trước, thiếu thì xuống + /// storage (`CachedStorage` LRU ở giữa). Trả `None` nếu id không tồn tại. + async fn node(&self, node_id: u64) -> Option { + if let Some(n) = self.nodes.lock().unwrap().get(&node_id) { + return Some(n.clone()); + } + let bytes = { + let guard = self.storage.read().await; + guard.get_node_meta(node_id as usize).await.ok().flatten()? + }; + if bytes.is_empty() { + return None; // meta đã bị clear (node removed). + } + let node = serde_json::from_slice::(&bytes).ok()?; + self.nodes.lock().unwrap().insert(node_id, node.clone()); + Some(node) + } + /// Return the document owning `node_id`, if any. - pub fn doc_of(&self, node_id: u64) -> Option<&Document> { - self.nodes.get(&node_id).and_then(|n| self.docs.get(&n.doc)) + pub async fn doc_of(&self, node_id: u64) -> Option<&Document> { + let node = self.node(node_id).await?; + self.docs.get(&node.doc) } /// Hydrate a node into a small payload suitable for LLM reasoning. - pub fn hydrate(&self, node_id: u64) -> Option { - let node = self.nodes.get(&node_id)?; - let path = self.collect_path(node_id); + /// Đọc node + tổ tiên (cho path) + con theo nhu cầu từ storage. + pub async fn hydrate(&self, node_id: u64) -> Option { + let node = self.node(node_id).await?; + let path = self.collect_path(node_id).await; + let mut children = Vec::new(); + for c in &node.children { + if let Some(payload) = Box::pin(self.hydrate(*c)).await { + children.push(payload); + } + } Some(NodePayload { id: node.id, path, @@ -237,11 +238,7 @@ impl DocumentGraph { value: node.value.clone(), key: node.key.clone(), doc: node.doc, - children: node - .children - .iter() - .filter_map(|c| self.hydrate(*c)) - .collect(), + children, }) } @@ -294,11 +291,20 @@ impl DocumentGraph { // ── Stats ───────────────────────────────────────────────────────── - pub fn stats(&self) -> DocStats { - DocStats { + pub async fn stats(&self) -> Result { + // nodes đếm từ chain node id (1 lần đọc chain, không đọc từng meta). + let nodes = { + let guard = self.storage.read().await; + guard + .get_chain(DOC_NODE_LIST_RECORD as usize) + .await? + .map(|c| c.len()) + .unwrap_or(0) + }; + Ok(DocStats { docs: self.docs.len(), - nodes: self.nodes.len(), - } + nodes, + }) } // ── Internal helpers ────────────────────────────────────── @@ -386,6 +392,7 @@ impl DocumentGraph { } async fn remove_document_nodes(&self, doc: &Document) -> Result<()> { + let removed: Vec = doc.nodes.iter().map(|n| n.id).collect(); for node in &doc.nodes { self.storage .write() @@ -393,12 +400,35 @@ impl DocumentGraph { .set_node_meta(node.id as usize, &[]) .await?; } + // Bỏ node id khỏi chain sentinel — stats đếm từ chain nên id cũ + // (doc bị replace/remove) phải ra khỏi danh sách. + let mut list = { + let guard = self.storage.read().await; + guard + .get_chain(DOC_NODE_LIST_RECORD as usize) + .await? + .map(|c| c.to_vec()) + .unwrap_or_default() + }; + list.retain(|id| !removed.contains(id)); + self.storage + .write() + .await + .set_chain(DOC_NODE_LIST_RECORD as usize, &list) + .await?; + // Cache in-memory cũng bỏ theo. + { + let mut cache = self.nodes.lock().unwrap(); + for id in removed { + cache.remove(&id); + } + } Ok(()) } - fn collect_path(&self, mut node_id: u64) -> Vec { + async fn collect_path(&self, mut node_id: u64) -> Vec { let mut path = Vec::new(); - while let Some(node) = self.nodes.get(&node_id) { + while let Some(node) = self.node(node_id).await { if let Some(key) = &node.key { path.push(key.clone()); } @@ -479,7 +509,8 @@ impl DocumentGraph { fn path_tokens(&mut self, node: &Node) -> Vec { let mut tokens = vec![DocToken::root()]; let mut cur = node.id; - while let Some(n) = self.nodes.get(&cur) { + let cache = self.nodes.lock().unwrap(); + while let Some(n) = cache.get(&cur) { if let Some(key) = &n.key { let key_id = self.intern.intern(key.clone()); tokens.push(DocToken::field(key_id)); @@ -524,12 +555,12 @@ mod tests { use super::*; use codegraph_graph::InMemoryStorage; - #[test] - fn new_graph() { + #[tokio::test] + async fn new_graph() { let storage = Arc::new(TokioRwLock::new(InMemoryStorage::default())); let config = DocConfig::default(); let graph = DocumentGraph::new(storage, config); - assert_eq!(graph.stats().docs, 0); + assert_eq!(graph.stats().await.unwrap_or_default().docs, 0); } /// `ingest_file` hai file khác nhau → doc id khác nhau, node không đè nhau; @@ -550,10 +581,10 @@ mod tests { let d1 = graph.ingest_file(p1.to_str().unwrap(), None).await.unwrap(); let d2 = graph.ingest_file(p2.to_str().unwrap(), None).await.unwrap(); assert_ne!(d1, d2); - assert_eq!(graph.stats().docs, 2); + assert_eq!(graph.stats().await.unwrap_or_default().docs, 2); // a.yaml: root+service+name+replicas = 4; b.toml: root+service+name = 3. // Nếu remap local-id sai thì 2 doc đè node nhau → tổng < 7. - assert_eq!(graph.stats().nodes, 7); + assert_eq!(graph.stats().await.unwrap_or_default().nodes, 7); // Re-ingest cùng path → id giữ nguyên. assert_eq!( @@ -565,9 +596,9 @@ mod tests { let mut reopened = DocumentGraph::open(storage, DocConfig::default()) .await .unwrap(); - assert_eq!(reopened.stats().docs, 2); + assert_eq!(reopened.stats().await.unwrap_or_default().docs, 2); // Node list được persist — mở lại phải khôi phục đủ node. - assert_eq!(reopened.stats().nodes, 7); + assert_eq!(reopened.stats().await.unwrap_or_default().nodes, 7); let d3 = reopened .ingest_file(p3.to_str().unwrap(), None) .await diff --git a/crates/codegraph-docs/src/parsers/mod.rs b/crates/codegraph-docs/src/parsers/mod.rs index ce9207c52..a9a96675e 100644 --- a/crates/codegraph-docs/src/parsers/mod.rs +++ b/crates/codegraph-docs/src/parsers/mod.rs @@ -934,7 +934,7 @@ http { let mut graph = DocumentGraph::new(storage, DocConfig::default()); let _doc_id = graph.ingest_file(p.to_str().unwrap(), None).await.unwrap(); // root + events + worker_connections = 3. - assert_eq!(graph.stats().nodes, 3); - assert_eq!(graph.stats().docs, 1); + assert_eq!(graph.stats().await.unwrap_or_default().nodes, 3); + assert_eq!(graph.stats().await.unwrap_or_default().docs, 1); } } diff --git a/crates/codegraph-graph/src/lib.rs b/crates/codegraph-graph/src/lib.rs index d4d04d075..543c5b7d9 100644 --- a/crates/codegraph-graph/src/lib.rs +++ b/crates/codegraph-graph/src/lib.rs @@ -117,6 +117,15 @@ fn backend_unavailable(name: &str) -> Error { )) } +/// Capacity mỗi LRU cache trong `CachedStorage` cho keyspace datasets +/// (docs/binary). Node meta ~vài trăm bytes/entry → 4096 entry ≈ vài MB. +/// Build không bật backend nào → các nhánh dùng nó bị cfg out. +#[cfg_attr( + not(any(feature = "sqlite", feature = "lmdb", feature = "redis")), + allow(dead_code) +)] +const DEFAULT_CACHE_CAPACITY: usize = 4096; + /// Map `search::Error` → `Error::Search`. fn serr_search(e: crate::search::Error) -> Error { Error::Search(e.to_string()) @@ -144,7 +153,12 @@ pub async fn open_keyspace_storage(dsn: &str, keyspace: &str) -> Result Result Result> = Arc::new(TokioRwLock::new(codegraph_graph::InMemoryStorage::default())); let mut graph = DocumentGraph::new(storage, DocConfig::default()); - let doc_id = graph.stats().docs as u64 + 1; + let doc_id = 1; // doc in-memory per-mutation — id không quan trọng let doc = parser .parse(&path, &source, doc_id) .map_err(|e| async_graphql::Error::new(e.to_string()))?; @@ -238,7 +238,7 @@ impl Mutation { .map_err(|e| async_graphql::Error::new(e.to_string()))?; let mut results = Vec::new(); for id in &ids { - if let Some(payload) = state.doc_graph.read().await.hydrate(*id) { + if let Some(payload) = state.doc_graph.read().await.hydrate(*id).await { results.push(json!({ "id": payload.id, "path": payload.path, "kind": format!("{:?}", payload.kind) })); } } @@ -249,7 +249,13 @@ impl Mutation { /// Get document stats. async fn doc_stats(&self, ctx: &Context<'_>) -> GqlResult { let state = ctx.data::>()?; - let stats = state.doc_graph.read().await.stats(); + let stats = state + .doc_graph + .read() + .await + .stats() + .await + .map_err(|e| async_graphql::Error::new(e.to_string()))?; Ok(format!("documents: {}\nnodes: {}", stats.docs, stats.nodes)) } } diff --git a/crates/codegraph-graphql/src/query.rs b/crates/codegraph-graphql/src/query.rs index 94d82e074..393543d39 100644 --- a/crates/codegraph-graphql/src/query.rs +++ b/crates/codegraph-graphql/src/query.rs @@ -346,7 +346,13 @@ impl Query { /// List all documents in the document graph. async fn doc_list(&self, ctx: &Context<'_>) -> GqlResult> { let state = ctx.data::>()?; - let stats = state.doc_graph.read().await.stats(); + let stats = state + .doc_graph + .read() + .await + .stats() + .await + .map_err(|e| async_graphql::Error::new(e.to_string()))?; Ok(vec![DocStatsView { docs: stats.docs, nodes: stats.nodes, @@ -372,7 +378,7 @@ impl Query { .map_err(|e| async_graphql::Error::new(e.to_string()))?; let mut results = Vec::new(); for id in &ids { - if let Some(payload) = state.doc_graph.read().await.hydrate(*id) { + if let Some(payload) = state.doc_graph.read().await.hydrate(*id).await { results.push(DocNodePayload { id: payload.id, path: payload.path, diff --git a/crates/codegraph-mcp/src/docgraph.rs b/crates/codegraph-mcp/src/docgraph.rs new file mode 100644 index 000000000..aeaae373f --- /dev/null +++ b/crates/codegraph-mcp/src/docgraph.rs @@ -0,0 +1,153 @@ +//! SharedDocGraph — document graph dùng chung cho MCP server, mở **lazily**. +//! +//! Mở `DocumentGraph` từ storage (`DocumentGraph::open` → rebuild toàn bộ +//! tries từ node/doc JSON) tốn thời gian tuyến tính với số node — với repo +//! document lớn có thể vượt startup timeout của MCP client. Nên server chỉ +//! giữ root lúc khởi động; lần doc tool đầu tiên mới trigger open + rebuild +//! (dưới rebuild_lock — N call đồng thời chỉ 1 lần open), các call sau dùng +//! handle đã cache. + +use camino::Utf8PathBuf; +use codegraph_docs::DocumentGraph; +use std::sync::{Arc, RwLock}; +use tokio::sync::{Mutex, RwLock as TokioRwLock}; + +/// Doc graph dùng chung: sẵn sàng (in-memory seed / đã open) hoặc lazy theo root. +pub struct SharedDocGraph { + state: RwLock, + /// Serialize open+rebuild — N doc call đồng thời chỉ 1 lần open. + rebuild_lock: Mutex<()>, +} + +enum SharedDocGraphState { + /// Handle sẵn sàng — trả ngay, không chờ. + Ready(Arc>), + /// Chưa open — lần `graph()` đầu mở storage + rebuild từ root này. + Lazy(Utf8PathBuf), +} + +impl SharedDocGraph { + /// Bọc handle đã sẵn sàng (in-memory seed của `CodegraphServer::new`). + pub fn ready(graph: Arc>) -> Self { + Self { + state: RwLock::new(SharedDocGraphState::Ready(graph)), + rebuild_lock: Mutex::new(()), + } + } + + /// Lazy theo workspace root — chưa open gì. Lỗi open khi `graph()` được + /// gọi → fallback in-memory (doc tools vẫn dùng được per-session), giữ + /// nguyên hành vi của đường eager cũ. + pub fn lazy(root: Utf8PathBuf) -> Self { + Self { + state: RwLock::new(SharedDocGraphState::Lazy(root)), + rebuild_lock: Mutex::new(()), + } + } + + /// Handle dùng được: fast path trả handle cached; lazy thì open+rebuild + /// đúng một lần dưới rebuild_lock rồi cache. + pub async fn graph(&self) -> Arc> { + if let SharedDocGraphState::Ready(g) = &*self.state.read().unwrap() { + return g.clone(); + } + let _guard = self.rebuild_lock.lock().await; + if let SharedDocGraphState::Ready(g) = &*self.state.read().unwrap() { + return g.clone(); + } + let root = match &*self.state.read().unwrap() { + SharedDocGraphState::Lazy(root) => root.clone(), + SharedDocGraphState::Ready(g) => return g.clone(), + }; + let graph = match codegraph_extract::open_doc_graph(&root).await { + Ok(g) => Arc::new(TokioRwLock::new(g)), + Err(e) => { + tracing::warn!("doc graph open failed ({e}) — fallback in-memory"); + Arc::new(TokioRwLock::new(DocumentGraph::new( + Arc::new(TokioRwLock::new(codegraph_graph::InMemoryStorage::default())), + codegraph_docs::DocConfig::default(), + ))) + } + }; + *self.state.write().unwrap() = SharedDocGraphState::Ready(graph.clone()); + graph + } + + /// Doc graph đã open chưa (`false` trên lazy instance chưa `graph()` nào). + pub fn is_ready(&self) -> bool { + matches!(&*self.state.read().unwrap(), SharedDocGraphState::Ready(_)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use codegraph_docs::DocConfig; + + /// Lazy instance chưa open gì — `is_ready` false, init không chạm storage. + #[test] + fn lazy_starts_not_ready() { + let shared = SharedDocGraph::lazy("/nonexistent-root-xyz".into()); + assert!(!shared.is_ready()); + } + + /// Ready instance trả đúng handle, `graph()` không đổi instance. + #[tokio::test] + async fn ready_returns_cached_handle() { + let storage: Arc> = + Arc::new(TokioRwLock::new(codegraph_graph::InMemoryStorage::default())); + let graph = Arc::new(TokioRwLock::new(DocumentGraph::new( + storage, + DocConfig::default(), + ))); + let shared = SharedDocGraph::ready(graph.clone()); + assert!(shared.is_ready()); + let got = shared.graph().await; + assert!(Arc::ptr_eq(&graph, &got)); + } + + /// N call `graph()` đồng thời trên lazy instance chỉ open 1 lần — call sau + /// dùng chung handle đã cache (root không tồn tại → fallback in-memory). + #[tokio::test] + async fn lazy_concurrent_calls_share_single_open() { + let shared = Arc::new(SharedDocGraph::lazy("/nonexistent-root-xyz".into())); + let (a, b) = { + let (s1, s2) = (shared.clone(), shared.clone()); + tokio::join!( + async move { s1.graph().await }, + async move { s2.graph().await } + ) + }; + assert!(Arc::ptr_eq(&a, &b)); + assert!(shared.is_ready()); + } + + /// Lazy instance trên root có docs.sqlite đã seed: lần `graph()` đầu + /// rebuild từ storage và thấy đúng dữ liệu (đường của MCP sau khi + /// `with_root_and_format` không chạm storage, doc tool đầu mới open). + #[tokio::test] + async fn lazy_graph_rebuilds_from_persisted_docs() { + let dir = tempfile::tempdir().unwrap(); + let root = camino::Utf8PathBuf::from_path_buf(dir.path().to_path_buf()).unwrap(); + + // "CLI process": ingest 1 doc JSON vào dataset mặc định của root. + let doc_file = dir.path().join("sample.json"); + std::fs::write(&doc_file, r#"{"name": "app", "replicas": 2}"#).unwrap(); + { + let mut graph = codegraph_extract::open_doc_graph(&root).await.unwrap(); + graph + .ingest_file(doc_file.to_string_lossy().as_ref(), None) + .await + .unwrap(); + } + + // "Server process": lazy open thấy lại doc đã persist. + let shared = SharedDocGraph::lazy(root); + assert!(!shared.is_ready()); + let graph = shared.graph().await; + let stats = graph.read().await.stats().await.unwrap(); + assert_eq!(stats.docs, 1); + assert!(stats.nodes > 0); + assert!(shared.is_ready()); + } +} diff --git a/crates/codegraph-mcp/src/lib.rs b/crates/codegraph-mcp/src/lib.rs index 3cc57766b..21cf12067 100644 --- a/crates/codegraph-mcp/src/lib.rs +++ b/crates/codegraph-mcp/src/lib.rs @@ -10,6 +10,7 @@ //! và [`http`] (Streamable HTTP — rmcp cấp một `CodegraphServer` riêng per //! `mcp-session-id`, mỗi phiên bind root riêng). +mod docgraph; #[cfg(feature = "http")] pub mod http; mod session; @@ -17,6 +18,7 @@ pub mod stdio; mod tools; mod usage; +pub use docgraph::SharedDocGraph; #[cfg(feature = "http")] pub use http::serve_http; pub use session::{DetailLevel, InitOutcome, OutputStyle, Session}; @@ -45,7 +47,7 @@ pub const SERVER_VERSION: &str = env!("CARGO_PKG_VERSION"); /// Server MCP. Transport-agnostic: stdio (1 process = 1 session) mount trực /// tiếp, http (tương lai) sẽ xoay vòng session store riêng. pub struct CodegraphServer { - session: Session, + session: Arc, usage: Arc>, /// Session store cho search resumable — sống qua nhiều tool call để resume /// id (trả về khi timeout) có thể retry được. @@ -54,7 +56,8 @@ pub struct CodegraphServer { /// tool trả lỗi rõ ràng. Tương ứng flag `--mermaid` ở CLI. mermaid: bool, /// Document graph for structured document operations (HCL, YAML, JSON, TOML). - doc_graph: Arc>, + /// Lazy: chỉ giữ root lúc startup, open+rebuild trễ tới doc tool đầu tiên. + doc_graph: Arc, } impl CodegraphServer { @@ -73,11 +76,11 @@ impl CodegraphServer { DocConfig::default(), ))); Self { - session: Session::new_with_format(format), + session: Arc::new(Session::new_with_format(format)), usage: Arc::new(Mutex::new(usage::UsageStats::default())), search_sessions: Arc::new(SearchSessionStore::new()), mermaid, - doc_graph, + doc_graph: Arc::new(SharedDocGraph::ready(doc_graph)), } } @@ -94,21 +97,13 @@ impl CodegraphServer { format: OutputStyle, mermaid: bool, ) -> anyhow::Result { - // Document graph mở từ `[docgraph]`/`[storage]` config của root - // (dataset riêng, persist qua các phiên). Lỗi config/backend → fallback - // in-memory thay vì chặn cả server (doc tools vẫn dùng được per-session). - let doc_graph = match codegraph_extract::open_doc_graph(&root).await { - Ok(g) => Arc::new(TokioRwLock::new(g)), - Err(e) => { - tracing::warn!("doc graph open failed ({e}) — fallback in-memory"); - Arc::new(TokioRwLock::new(DocumentGraph::new( - Arc::new(TokioRwLock::new(InMemoryStorage::default())), - DocConfig::default(), - ))) - } - }; + // Document graph mở LAZY theo `[docgraph]`/`[storage]` config của root + // (dataset riêng, persist qua các phiên): startup chỉ giữ root, open + + // rebuild (tuyến tính với số node — có thể lâu trên repo document lớn) + // trễ tới doc tool đầu tiên. Xem `SharedDocGraph`. + let doc_graph = Arc::new(SharedDocGraph::lazy(root.clone())); Ok(Self { - session: Session::with_root_and_format(root, format).await?, + session: Arc::new(Session::with_root_and_format(root, format).await?), usage: Arc::new(Mutex::new(usage::UsageStats::default())), search_sessions: Arc::new(SearchSessionStore::new()), mermaid, @@ -121,6 +116,24 @@ impl CodegraphServer { self.mermaid } + /// Prewarm symbol index ngầm: `initialize` của client không chờ index, + /// nhưng tool call đầu tiên sẽ block cho tới khi `SharedGraphIndex` build + /// xong snapshot (repo lớn → cả phút). Spawn task build ngay sau khi + /// serve bắt đầu — call đầu không còn chờ (hoặc chỉ chờ task này xong). + /// Chỉ có ý nghĩa khi session đã pre-seed root (`with_root_and_format`); + /// session trống → `ensure_ready` refuse, bỏ qua im lặng. + pub fn prewarm_symbol_index(&self) { + let session = Arc::clone(&self.session); + tokio::spawn(async move { + match session.ensure_ready().await { + Ok(index) => { + let _ = index.ensure_fresh().await; + } + Err(e) => tracing::debug!("symbol index prewarm skipped: {e}"), + } + }); + } + /// Dispatch một tool call đã verify tên. Trả [`ToolOutput::Text`] cho thành /// công, [`ToolOutput::Error`] cho lỗi tool (client thấy `is_error`), /// [`Err`] cho lỗi protocol (unknown tool đã bị chặn trước ở `call_tool`). @@ -237,7 +250,8 @@ impl CodegraphServer { let detail = self.session.detail().await; let format = self.session.format().await; - // Document tools — don't require session ready. + // Document tools — lazy doc graph (SharedDocGraph), không cần session + // ready. Open giờ rẻ: `DocumentGraph::open` không materialize nodes. if name.starts_with("codegraph_doc_") { let doc_graph = self.doc_graph.clone(); return match name { diff --git a/crates/codegraph-mcp/src/tools.rs b/crates/codegraph-mcp/src/tools.rs index c2c5fcd8e..193e9968f 100644 --- a/crates/codegraph-mcp/src/tools.rs +++ b/crates/codegraph-mcp/src/tools.rs @@ -3,12 +3,11 @@ use camino::Utf8Path; use codegraph_api::{GraphApi, Pagination}; use codegraph_context::{ContextRequest, Format}; use codegraph_core::{Error, Result, Symbol, SymbolKind, SymbolMatch}; -use codegraph_docs::{tokenize::DocToken, DocumentGraph}; +use codegraph_docs::tokenize::DocToken; use rmcp::model::Tool; use serde::Serialize; use serde_json::{json, Value}; use std::sync::Arc; -use tokio::sync::RwLock as TokioRwLock; /// Định nghĩa một MCP tool — single source of truth cho `tools/list`. struct ToolDef { @@ -1070,11 +1069,13 @@ pub(crate) fn omit_defaults(v: &mut Value) { // ── Document tool dispatch ── pub async fn dispatch_doc_ingest( - doc_graph: Arc>, + doc_graph: Arc, path: &str, format: Option, ) -> Result { let inserted = doc_graph + .graph() + .await .write() .await .ingest_file(path, format.as_deref()) @@ -1084,12 +1085,14 @@ pub async fn dispatch_doc_ingest( } pub async fn dispatch_doc_search( - doc_graph: Arc>, + doc_graph: Arc, _pattern: &str, depth: usize, ) -> Result { let tokens = vec![DocToken::root()]; let ids = doc_graph + .graph() + .await .read() .await .search_path(&tokens, Some(depth)) @@ -1098,9 +1101,10 @@ pub async fn dispatch_doc_search( if ids.is_empty() { return Ok("no nodes matched".to_string()); } + let graph = doc_graph.graph().await; let mut results = Vec::new(); for id in &ids { - if let Some(payload) = doc_graph.read().await.hydrate(*id) { + if let Some(payload) = graph.read().await.hydrate(*id).await { results.push(json!({ "id": payload.id, "path": payload.path, "kind": format!("{:?}", payload.kind), "value": payload.value })); } } @@ -1108,10 +1112,10 @@ pub async fn dispatch_doc_search( } pub async fn dispatch_doc_hydrate( - doc_graph: Arc>, + doc_graph: Arc, node_id: u64, ) -> Result { - let payload = doc_graph.read().await.hydrate(node_id); + let payload = doc_graph.graph().await.read().await.hydrate(node_id).await; match payload { Some(p) => { let json = serde_json::to_string_pretty(&p).map_err(|e| Error::Other(e.to_string()))?; @@ -1121,13 +1125,27 @@ pub async fn dispatch_doc_hydrate( } } -pub async fn dispatch_doc_list(doc_graph: Arc>) -> Result { - let stats = doc_graph.read().await.stats(); +pub async fn dispatch_doc_list(doc_graph: Arc) -> Result { + let stats = doc_graph + .graph() + .await + .read() + .await + .stats() + .await + .map_err(|e| Error::Other(e.to_string()))?; Ok(format!("documents: {}, nodes: {}", stats.docs, stats.nodes)) } -pub async fn dispatch_doc_stats(doc_graph: Arc>) -> Result { - let stats = doc_graph.read().await.stats(); +pub async fn dispatch_doc_stats(doc_graph: Arc) -> Result { + let stats = doc_graph + .graph() + .await + .read() + .await + .stats() + .await + .map_err(|e| Error::Other(e.to_string()))?; Ok(format!("documents: {}\nnodes: {}", stats.docs, stats.nodes)) } diff --git a/crates/codegraph/src/main.rs b/crates/codegraph/src/main.rs index f4a8946da..3a461e1a8 100644 --- a/crates/codegraph/src/main.rs +++ b/crates/codegraph/src/main.rs @@ -328,13 +328,25 @@ async fn ingest_configured_docs(root: &Utf8Path) -> Result<()> { return Ok(()); } let mut graph = open_doc_graph(root).await?; + let bar = indicatif::ProgressBar::new(files.len() as u64); + bar.set_style( + indicatif::ProgressStyle::default_bar() + .template("[{elapsed_precise}] [{wide_bar}] {pos}/{len} ({percent}%) {msg}") + .expect("valid progress bar template") + .progress_chars("#>-"), + ); let mut ingested = 0usize; for (path, format) in &files { + bar.set_message(path.to_string()); match graph.ingest_file(path.as_str(), format.as_deref()).await { Ok(_) => ingested += 1, - Err(e) => eprintln!("doc ingest failed for {path}: {e}"), + Err(e) => { + bar.suspend(|| eprintln!("doc ingest failed for {path}: {e}")); + } } + bar.inc(1); } + bar.finish_and_clear(); eprintln!("ingested {ingested}/{} documents", files.len()); Ok(()) } @@ -703,6 +715,11 @@ async fn cmd_serve( } else { CodegraphServer::new_with_format(format, mermaid) }; + if use_root { + // Build symbol index ngầm sau khi server nhận request — call tool đầu + // không block cả phút trên repo lớn (initialize không bao giờ chờ). + server.prewarm_symbol_index(); + } codegraph_mcp::serve_stdio(server).await } @@ -725,13 +742,13 @@ async fn cmd_doc(root: &Utf8Path, cmd: DocCmd) -> Result<()> { println!("no nodes matched"); } else { for id in &ids { - if let Some(payload) = graph.hydrate(*id) { + if let Some(payload) = graph.hydrate(*id).await { println!("{}: {:?}", id, payload); } } } } - DocCmd::Hydrate { node_id } => match graph.hydrate(node_id) { + DocCmd::Hydrate { node_id } => match graph.hydrate(node_id).await { Some(payload) => { let json = serde_json::to_string_pretty(&payload)?; println!("{json}"); @@ -739,11 +756,11 @@ async fn cmd_doc(root: &Utf8Path, cmd: DocCmd) -> Result<()> { None => println!("node {node_id} not found"), }, DocCmd::List => { - let stats = graph.stats(); + let stats = graph.stats().await?; println!("documents: {}, nodes: {}", stats.docs, stats.nodes); } DocCmd::Stats => { - let stats = graph.stats(); + let stats = graph.stats().await?; println!("documents: {}", stats.docs); println!("nodes: {}", stats.nodes); } diff --git a/packaging/aur/codegraph-rs-bin/PKGBUILD b/packaging/aur/codegraph-rs-bin/PKGBUILD index 8818dc784..657cf7b45 100644 --- a/packaging/aur/codegraph-rs-bin/PKGBUILD +++ b/packaging/aur/codegraph-rs-bin/PKGBUILD @@ -1,6 +1,6 @@ # Maintainer: Hung Pham pkgname=codegraph-rs-bin -pkgver=2.1.6 +pkgver=2.1.7 pkgrel=1 pkgdesc="Local-first code intelligence: tree-sitter knowledge graph + MCP server (prebuilt binary)" arch=('x86_64' 'aarch64') diff --git a/packaging/choco/codegraph.nuspec b/packaging/choco/codegraph.nuspec index 83db65d4c..a9821eb6d 100644 --- a/packaging/choco/codegraph.nuspec +++ b/packaging/choco/codegraph.nuspec @@ -2,7 +2,7 @@ codegraph - 2.1.6 + 2.1.7 codegraph Hung Pham https://github.com/hungpham10/codegraph-rs diff --git a/packaging/winget/codegraph.yaml b/packaging/winget/codegraph.yaml index 4213d0936..ecfd8f4e8 100644 --- a/packaging/winget/codegraph.yaml +++ b/packaging/winget/codegraph.yaml @@ -6,7 +6,7 @@ # release time (or automate it in the release pipeline before submitting to # microsoft/winget-pkgs). PackageIdentifier: hungpham10.codegraph -PackageVersion: 2.1.6 +PackageVersion: 2.1.7 PackageName: codegraph Publisher: Hung Pham PublisherUrl: https://github.com/hungpham10/codegraph-rs @@ -18,7 +18,7 @@ PackageUrl: https://github.com/hungpham10/codegraph-rs InstallerType: zip Installers: - Architecture: x64 - InstallerUrl: https://github.com/hungpham10/codegraph-rs/releases/download/v2.1.6/codegraph-x86_64-pc-windows-msvc.zip + InstallerUrl: https://github.com/hungpham10/codegraph-rs/releases/download/v2.1.7/codegraph-x86_64-pc-windows-msvc.zip InstallerSha256: 0000000000000000000000000000000000000000000000000000000000000000 InstallerType: zip ManifestType: singleton diff --git a/scripts/install.ps1 b/scripts/install.ps1 index 3b22da196..75cc24b8b 100644 --- a/scripts/install.ps1 +++ b/scripts/install.ps1 @@ -5,11 +5,11 @@ # # Usage (pin a version / download the script first): # irm https://raw.githubusercontent.com/hungpham10/codegraph-rs/main/scripts/install.ps1 -OutFile install.ps1 -# .\install.ps1 -Version 2.1.6 +# .\install.ps1 -Version 2.1.7 [CmdletBinding()] param( - # Pin a specific version, e.g. "2.1.6". Empty = latest release. + # Pin a specific version, e.g. "2.1.7". Empty = latest release. [string]$Version )