diff --git a/Cargo.toml b/Cargo.toml index 56206109a..6a2121536 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,7 +24,9 @@ authors = ["Cleboost "] [workspace.dependencies] # core serde = { version = "1", features = ["derive"] } -serde_json = "1" +# preserve_order: json!-built objects emit keys in written (documented) order +# thay vì alphabet — derived structs luôn serialize theo declaration order. +serde_json = { version = "1", features = ["preserve_order"] } thiserror = "2" async-trait = "0.1" anyhow = "1" diff --git a/crates/codegraph-api/src/lib.rs b/crates/codegraph-api/src/lib.rs index 35ffc9d77..cf1387be5 100644 --- a/crates/codegraph-api/src/lib.rs +++ b/crates/codegraph-api/src/lib.rs @@ -9,17 +9,149 @@ use codegraph_core::{ CallSiteResult, ClassInfo, DependenciesReport, Error, FileInfo, FlowResult, FunctionScope, MemberInfo, ResolveResult, Result, SearchFlowResult, Symbol, SymbolKind, SymbolMatch, }; -use codegraph_graph::{GraphIndex, SharedGraphIndex}; -use std::sync::Arc; +use codegraph_graph::{GraphIndex, SearchCursor, SharedGraphIndex}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; pub struct GraphApi { shared_index: Arc, + /// Session store cho search resumable (resume id → cursor). + sessions: Arc, +} + +// ==================== Search session store ==================== + +/// Cursor session lưu **phía server** — LLM chỉ cầm một id ngắn (hex) và echo +/// lại khi retry. Id vô nghĩa ngoài tiến trình này: index version đổi (re-ingest) +/// hoặc server restart → session stale, báo LLM retry không có `resume`. +struct StoredResume { + created: Instant, + /// Version index lúc tạo — đổi (re-ingest) → cursor mất giá trị. + index_version: u64, + cursor: SearchCursor, +} + +/// Store in-process cho resume id → cursor. Không persist; purge theo TTL khi +/// `put` (đủ cho use-case retry trong vài phút). +pub struct SearchSessionStore { + inner: Mutex>, + ttl: Duration, + max_sessions: usize, + next_id: AtomicU64, +} + +impl Default for SearchSessionStore { + fn default() -> Self { + Self::new() + } +} + +impl SearchSessionStore { + pub fn new() -> Self { + Self { + inner: Mutex::new(HashMap::new()), + ttl: Duration::from_secs(600), + max_sessions: 512, + next_id: AtomicU64::new(0), + } + } + + /// Lưu cursor, trả id hex ngắn. Trước khi thêm: purge session quá TTL, chặn + /// số session tối đa (evict session già nhất). + pub fn put(&self, cursor: SearchCursor, index_version: u64) -> String { + let mut map = self.inner.lock().unwrap(); + let now = Instant::now(); + map.retain(|_, s| now.duration_since(s.created) < self.ttl); + while map.len() >= self.max_sessions { + let oldest = map + .iter() + .min_by_key(|(_, s)| s.created) + .map(|(k, _)| k.clone()); + if let Some(k) = oldest { + map.remove(&k); + } else { + break; + } + } + let id = Self::gen_id(&self.next_id); + map.insert( + id.clone(), + StoredResume { + created: now, + index_version, + cursor, + }, + ); + id + } + + /// Đọc cursor theo id — `None` nếu không có / quá TTL. + pub fn get(&self, id: &str) -> Option<(u64, SearchCursor)> { + let map = self.inner.lock().unwrap(); + map.get(id).map(|s| (s.index_version, s.cursor.clone())) + } + + /// Xoá session (khi search hoàn tất, không còn page nào). + pub fn remove(&self, id: &str) { + self.inner.lock().unwrap().remove(id); + } + + /// Id hex 16 ký tự: epoch-nanos + counter tiến trình — đủ unique trong + /// tiến trình, không cần crate random. + fn gen_id(counter: &AtomicU64) -> String { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + let n = counter.fetch_add(1, Ordering::Relaxed); + let v = (nanos as u64) ^ (n.wrapping_mul(0x9E37_79B9_7F4A_7C15)); + format!("{:016x}", v) + } +} + +/// Kết quả search resumable từ `GraphApi` — tầng MCP dựng message từ đây. +#[derive(Debug)] +pub struct ResumeSearchOutcome { + pub page: Vec, + pub total: usize, + pub timed_out: bool, + /// Số đơn vị đã xử lý lúc ngắt (names khi đang phase A, symbols khi phase B). + pub progress: usize, + /// Resume id để retry: `Some` khi timed_out HOẶC còn page sau. `None` = + /// xong và hết page (session đã xoá). + pub resume: Option, + /// Version index mà search chạy trên — đổi giữa các lần retry → resume + /// không còn giá trị. + pub index_version: u64, +} + +/// Phân trang cho search symbol: `limit` chặn số symbol mỗi trang (`0` = +/// không giới hạn), `offset` bỏ qua `offset` symbol đầu. +#[derive(Debug, Clone, Copy)] +pub struct Pagination { + pub limit: u32, + pub offset: u32, } impl GraphApi { pub fn new_with_index(index: Arc) -> Self { Self { shared_index: index, + sessions: Arc::new(SearchSessionStore::new()), + } + } + + /// Dùng chung session store (resume id) — server MCP giữ store ở vòng đời + /// server để resume id sống qua nhiều tool call. + pub fn new_with_sessions( + index: Arc, + sessions: Arc, + ) -> Self { + Self { + shared_index: index, + sessions, } } @@ -36,6 +168,27 @@ impl GraphApi { .await } + /// Resumable + deadline-aware của [`Self::search`] — nền cho + /// `codegraph_search`. `timeout_ms = 0` = không giới hạn thời gian. + /// `resume` = id trả về từ lần timeout trước (phải cùng query). + pub async fn search_resumable( + &self, + query: &str, + limit: u32, + resume: Option, + timeout_ms: u64, + ) -> Result { + self.search_symbol_paged_resumable( + query, + None, + SymbolMatch::Contains, + Pagination { limit, offset: 0 }, + resume, + timeout_ms, + ) + .await + } + /// Search symbol nâng cao — kind filter + match mode + phân trang. /// Trả về (page, total). pub async fn search_symbol_paged( @@ -52,6 +205,89 @@ impl GraphApi { .await } + /// Resumable + deadline-aware của [`Self::search_symbol_paged`] — nền cho + /// `codegraph_search_symbol`. `timeout_ms = 0` = không giới hạn. + /// + /// `resume` được validate (index version + query/mode/kind phải khớp) — + /// sai → lỗi báo LLM retry không có `resume`. + pub async fn search_symbol_paged_resumable( + &self, + query: &str, + kind: Option, + mode: SymbolMatch, + pagination: Pagination, + resume: Option, + timeout_ms: u64, + ) -> Result { + let idx = self.index().await; + let version = idx.version(); + let q = query.to_lowercase(); + + // ── Validate resume id (nếu có) ── + let cursor = match &resume { + Some(id) => { + let (stored_version, stored) = self.sessions.get(id).ok_or_else(|| { + Error::Invalid("resume id expired or unknown — retry without resume".into()) + })?; + if stored_version != version { + return Err(Error::Invalid( + "index was re-built since this resume was created — retry without resume" + .into(), + )); + } + if stored.query != q || stored.mode != mode || stored.kind != kind { + return Err(Error::Invalid( + "resume id was created for a different query — retry without resume".into(), + )); + } + Some(stored) + } + None => None, + }; + + // ── Deadline ── + let deadline = if timeout_ms == 0 { + None + } else { + Some(Instant::now() + Duration::from_millis(timeout_ms)) + }; + + let out = idx + .search_symbol_paged_resumable( + query, + kind, + mode, + codegraph_graph::Pagination { + limit: pagination.limit as usize, + offset: pagination.offset as usize, + }, + cursor, + deadline, + ) + .await?; + + // ── Quản lý session: lưu khi còn tiếp tục (timeout / còn page), xoá + // khi xong hẳn. ── + let resume_id = match &out.cursor { + Some(c) => Some(self.sessions.put(c.clone(), version)), + None => { + if let Some(id) = &resume { + self.sessions.remove(id); + } + None + } + }; + + Ok(ResumeSearchOutcome { + page: out.page, + total: out.total, + timed_out: out.timed_out, + progress: out.progress, + resume: resume_id, + index_version: version, + }) + } + /// Methods của class (compact projection). pub async fn class_methods(&self, id: u64) -> Vec { self.index().await.list_methods_of_class(id) diff --git a/crates/codegraph-api/tests/api.rs b/crates/codegraph-api/tests/api.rs index 8cdd85ef0..7f791203a 100644 --- a/crates/codegraph-api/tests/api.rs +++ b/crates/codegraph-api/tests/api.rs @@ -1,5 +1,7 @@ -use codegraph_api::GraphApi; -use codegraph_core::{CallRecord, EffectType, ScopeLevel, Symbol, SymbolKind, SYMBOL_BASE}; +use codegraph_api::{GraphApi, Pagination}; +use codegraph_core::{ + CallRecord, EffectType, ScopeLevel, Symbol, SymbolKind, SymbolMatch, SYMBOL_BASE, +}; use codegraph_graph::{GraphIndex, ParseResult, SharedGraphIndex}; use std::collections::HashMap; use std::sync::Arc; @@ -167,7 +169,177 @@ async fn files_stats_and_context() { include_source: false, limit: 5, format: codegraph_context::Format::Markdown, + strip_prefix: None, }; let md = api.context_markdown(&req).await.unwrap(); assert!(md.contains("caller")); } + +/// Seed N symbol tên "order_*" (mỗi tên 1 symbol) — đủ lớn để search chậm hơn +/// 1ms (debug build) và có tổng > limit (test phân trang). +async fn seed_many(db: &str, count: usize) { + let mut idx = GraphIndex::open(db).await.unwrap(); + let mut results = Vec::new(); + for (id, i) in (SYMBOL_BASE..).zip(0..count) { + let name = format!("order_{i:05}"); + results.push(ParseResult { + path: "src/a.ts".into(), + language: "typescript".into(), + bytes: 10, + lines: 4, + symbols: vec![sym(id, &name)], + chains: HashMap::new(), + calls: vec![], + }); + } + idx.ingest(&results).await.unwrap(); +} + +/// Resume roundtrip: timeout → lấy resume id → retry cùng args + resume → kết +/// quả đầy đủ, không lặp/không mất. Resume id sai → lỗi bảo retry không resume. +#[tokio::test] +async fn search_resumable_timeout_retry_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("db.sqlite"); + let db_str = format!("sqlite://{}", db_path.to_string_lossy()); + seed_many(&db_str, 6000).await; + let api = api(&db_str).await; + + // Call 1: timeout_ms=1 — trên seed 6000 symbol debug build chắc chắn trễ + // hơn 1ms. Nếu máy quá nhanh (không timeout) test vẫn đúng — chỉ bỏ qua + // nhánh retry. total = 5000 vì name engine chặn cứng MAX_RESULTS tên distinct. + let capped = 5000; + let first = api.search_resumable("order", 20, None, 1).await.unwrap(); + let resume_id = if first.timed_out { + assert!(first.resume.is_some(), "timeout must carry a resume id"); + first.resume.unwrap() + } else { + // Hoàn tất ngay — verify kết quả rồi dừng (không cần retry). + let ids: std::collections::HashSet = first.page.iter().map(|s| s.id).collect(); + assert_eq!(ids.len(), first.page.len(), "no duplicate results"); + assert_eq!(first.total, capped); + assert_eq!(first.page.len(), 20); + return; + }; + + // Retry: cùng args + resume, không giới hạn thời gian → hoàn tất. + let out = api + .search_resumable("order", 20, Some(resume_id.clone()), 0) + .await + .unwrap(); + assert!(!out.timed_out); + assert_eq!(out.total, capped, "total must match the full scan (capped)"); + let ids: std::collections::HashSet = out.page.iter().map(|s| s.id).collect(); + assert_eq!( + ids.len(), + out.page.len(), + "no duplicate results after resume" + ); + assert_eq!(out.page.len(), 20); + + // Resume id không tồn tại → lỗi (LLM nên retry không resume). + assert!( + api.search_resumable("order", 20, Some("deadbeef00000000".into()), 0) + .await + .is_err(), + "unknown resume id must be rejected" + ); + // Resume id không khớp query → lỗi. + assert!( + api.search_resumable("totally_different", 20, Some(resume_id), 0) + .await + .is_err(), + "resume id for a different query must be rejected" + ); +} + +/// Phân trang qua resume (Paged cursor): call 1 limit=10 (timeout_ms=0) hoàn +/// tất + còn page sau → resume id; call 2 cùng resume + offset=10 → page rời, +/// tổng nhất quán. +#[tokio::test] +async fn search_symbol_paged_resume_paging() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("db.sqlite"); + let db_str = format!("sqlite://{}", db_path.to_string_lossy()); + seed_many(&db_str, 1500).await; + let api = api(&db_str).await; + + let first = api + .search_symbol_paged_resumable( + "order", + None, + SymbolMatch::Contains, + Pagination { + limit: 10, + offset: 0, + }, + None, + 0, + ) + .await + .unwrap(); + assert!(!first.timed_out); + assert_eq!(first.total, 1500); + assert_eq!(first.page.len(), 10); + assert!( + first.resume.is_some(), + "more pages remain -> response must carry a resume id" + ); + let resume_id = first.resume.unwrap(); + + // Trang 2 qua resume (Paged cursor — không quét lại). + let second = api + .search_symbol_paged_resumable( + "order", + None, + SymbolMatch::Contains, + Pagination { + limit: 10, + offset: 10, + }, + Some(resume_id.clone()), + 0, + ) + .await + .unwrap(); + assert!(!second.timed_out); + assert_eq!(second.total, 1500, "total must be stable across pages"); + let page1: std::collections::HashSet = first.page.iter().map(|s| s.id).collect(); + let page2: std::collections::HashSet = second.page.iter().map(|s| s.id).collect(); + assert!(page1.is_disjoint(&page2), "pages must be disjoint"); + assert_eq!(second.page.len(), 10); + + // Page cuối rời. + let last = api + .search_symbol_paged_resumable( + "order", + None, + SymbolMatch::Contains, + Pagination { + limit: 10, + offset: 1490, + }, + Some(resume_id.clone()), + 0, + ) + .await + .unwrap(); + assert_eq!(last.page.len(), 10); + assert!(last.resume.is_none(), "last page: no more resume"); + + // Resume id này thuộc query "order" — dùng cho query khác → lỗi. + assert!(api + .search_symbol_paged_resumable( + "zzz", + None, + SymbolMatch::Contains, + Pagination { + limit: 10, + offset: 0 + }, + Some(resume_id), + 0 + ) + .await + .is_err()); +} diff --git a/crates/codegraph-context/src/lib.rs b/crates/codegraph-context/src/lib.rs index 1fdebb868..29c92040c 100644 --- a/crates/codegraph-context/src/lib.rs +++ b/crates/codegraph-context/src/lib.rs @@ -26,6 +26,10 @@ pub struct ContextRequest { pub include_source: bool, pub limit: u32, pub format: Format, + /// Workspace root — strip khỏi `file` trong markdown (path tương đối tiết + /// kiệm token cho LLM). `None` = giữ absolute (CLI/HTTP không biết root). + #[serde(default)] + pub strip_prefix: Option, } impl Default for ContextRequest { @@ -36,6 +40,7 @@ impl Default for ContextRequest { include_source: false, limit: 5, format: Format::Markdown, + strip_prefix: None, } } } @@ -59,7 +64,7 @@ pub async fn build(index: &Arc, req: &ContextRequest) -> Resul let response = build_response(index, req).await?; match req.format { Format::Json => Ok(serde_json::to_string_pretty(&response).unwrap_or_default()), - Format::Markdown => Ok(render_markdown(&response)), + Format::Markdown => Ok(render_markdown(&response, req.strip_prefix.as_deref())), } } @@ -113,7 +118,19 @@ pub async fn build_response( }) } -fn render_markdown(resp: &ContextResponse) -> String { +/// Strip `root/` prefix khỏi path (boundary-aware) — `None` giữ nguyên. +fn rel_path<'a>(p: &'a str, strip: Option<&str>) -> &'a str { + if let Some(root) = strip { + if let Some(rest) = p.strip_prefix(root) { + if let Some(rest) = rest.strip_prefix('/') { + return rest; + } + } + } + p +} + +fn render_markdown(resp: &ContextResponse, strip: Option<&str>) -> String { let mut out = String::new(); let _ = writeln!(out, "# Context: `{}`", resp.query); if resp.hits.is_empty() { @@ -126,7 +143,7 @@ fn render_markdown(resp: &ContextResponse) -> String { "\n## `{}` — {} — `{}:{}`", h.symbol.name, h.symbol.kind.as_str(), - h.symbol.file, + rel_path(&h.symbol.file, strip), h.symbol.line ); if let Some(sig) = &h.symbol.signature { @@ -138,13 +155,25 @@ fn render_markdown(resp: &ContextResponse) -> String { if !h.callers.is_empty() { let _ = writeln!(out, "\n**Callers** ({}):", h.callers.len()); for c in &h.callers { - let _ = writeln!(out, "- `{}` — `{}:{}`", c.name, c.file, c.line); + let _ = writeln!( + out, + "- `{}` — `{}:{}`", + c.name, + rel_path(&c.file, strip), + c.line + ); } } if !h.callees.is_empty() { let _ = writeln!(out, "\n**Callees** ({}):", h.callees.len()); for c in &h.callees { - let _ = writeln!(out, "- `{}` — `{}:{}`", c.name, c.file, c.line); + let _ = writeln!( + out, + "- `{}` — `{}:{}`", + c.name, + rel_path(&c.file, strip), + c.line + ); } } } diff --git a/crates/codegraph-core/src/semgraph.rs b/crates/codegraph-core/src/semgraph.rs index 47ecd105a..68cff5154 100644 --- a/crates/codegraph-core/src/semgraph.rs +++ b/crates/codegraph-core/src/semgraph.rs @@ -174,6 +174,17 @@ pub enum ScopeLevel { Parameter, } +impl ScopeLevel { + pub fn as_str(self) -> &'static str { + match self { + Self::Global => "global", + Self::ObjectField => "object_field", + Self::Local => "local", + Self::Parameter => "parameter", + } + } +} + /// Phân loại tác động bên ngoài của một call (để impact/report). #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)] #[serde(rename_all = "snake_case")] diff --git a/crates/codegraph-graph/src/lib.rs b/crates/codegraph-graph/src/lib.rs index 0457c3182..3c64ee90a 100644 --- a/crates/codegraph-graph/src/lib.rs +++ b/crates/codegraph-graph/src/lib.rs @@ -35,6 +35,7 @@ //! var-type alias, gom SaveCallRecords) → files → rebuild engines → bump version. pub use crate::search::Search; +use crate::search::SearchResume; #[cfg(feature = "lmdb")] pub use crate::storage::lmdb::LmdbStorage; #[cfg(feature = "sqlite")] @@ -49,6 +50,7 @@ use codegraph_core::{ use serde::{Deserialize, Serialize}; use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; +use std::time::Instant; use tokio::sync::RwLock; #[cfg(feature = "bloom-search")] @@ -136,6 +138,9 @@ pub struct GraphIndex { names: Search, /// `record - 1` → tên (song song với thứ tự insert name engine). name_records: Vec, + /// Các tên distinct (lowercase) **đã sort** — dùng cho scan Prefix/Suffix/ + /// Exact không cần sort lại (resume chỉ lưu vị trí, không lưu mảng). + sorted_name_keys: Vec, /// symbol id → Symbol (registry — nguồn chân lý in-memory). symbols: HashMap, /// tên (lowercase) → symbol ids (mở rộng trùng tên khi search/resolve). @@ -156,6 +161,71 @@ pub struct GraphIndex { version: u64, } +// ── Search resumable (deadline-aware, checkpointable) ── + +/// Phase của một search resumable. Cursor nằm server-side (session store ở +/// tầng API) — LLM chỉ nhìn thấy một session id; state không serialize ra +/// ngoài mà chỉ được tái lập từ (query, mode, kind, phase). +#[derive(Debug, Clone)] +pub enum SearchCursorPhase { + /// Mode `Contains`: đang giữa lúc chạy name engine (resumable). `names` + /// engine trả records không theo thứ tự tên — phase A xong phải sort + /// trước khi chuyển sang `Expand`. + Engine(SearchResume), + /// Mode `Prefix`/`Suffix`/`Exact`: đang quét `sorted_name_keys` từ + /// `name_pos` (mảng đã sort sẵn — resume chỉ cần lưu vị trí). + ScanNames { name_pos: usize }, + /// Phase A xong: stream ids theo từng tên đã sort — filter `kind`, gom + /// vào `collected`. `name_index[name]` luôn tăng dần theo id nên + /// collected đã sort theo (name, id) — không cần sort cuối. + Expand { + names: Vec, + name_idx: usize, + id_idx: usize, + collected: Vec, + }, + /// Search đã hoàn tất: `collected` + `total` giữ để phân trang tiếp mà + /// không quét lại. Không chứa query — `SearchCursor.query` lo phần đó. + Paged { collected: Vec, total: usize }, +} + +/// Server-side cursor cho search resumable — validate theo (query, mode, +/// kind); `limit`/`offset` KHÔNG nằm trong cursor (page có thể đổi giữa +/// chừng khi retry). +#[derive(Debug, Clone)] +pub struct SearchCursor { + /// Query lowercase (khớp với cursor — thay đổi → resume không hợp lệ). + pub query: String, + pub mode: SymbolMatch, + pub kind: Option, + pub phase: SearchCursorPhase, +} + +/// Kết quả của [`GraphIndex::search_symbol_paged_resumable`]. +#[derive(Debug)] +pub struct PagedSearchOutcome { + /// Trang kết quả (chỉ đầy khi search hoàn tất; rỗng khi timed out). + pub page: Vec, + /// Tổng số khớp — chỉ chính xác khi search hoàn tất. + pub total: usize, + /// Deadline hết hạn giữa chừng → `cursor` là phase dở, phải retry. + pub timed_out: bool, + /// Tiến độ đo được lúc ngắt (names khớp khi đang phase A, symbols gom + /// được khi đang phase B) — dùng cho message báo LLM. + pub progress: usize, + /// Cursor tiếp tục: `Some(phase dở)` khi timed_out; `Some(Paged)` khi + /// hoàn tất nhưng còn page sau; `None` khi xong và hết page. + pub cursor: Option, +} + +/// Phân trang cho search symbol: `limit` chặn số symbol mỗi trang +/// (`0` = không giới hạn), `offset` bỏ qua `offset` symbol đầu. +#[derive(Debug, Clone, Copy)] +pub struct Pagination { + pub limit: usize, + pub offset: usize, +} + impl GraphIndex { /// Index in-memory (test/dev, không persist). pub fn in_memory() -> Self { @@ -333,6 +403,7 @@ impl GraphIndex { names: Search::new(CHAIN_SHARDING, name_storage), storage, name_records: Vec::new(), + sorted_name_keys: Vec::new(), symbols: HashMap::new(), name_index: HashMap::new(), scope_index: HashMap::new(), @@ -531,6 +602,7 @@ impl GraphIndex { p.phase("rebuild name-search engine", distinct.len()); } let mut record = 0usize; + self.sorted_name_keys = distinct.iter().map(|s| s.to_string()).collect(); for name in distinct { record += 1; let metas: Vec> = vec![None; name.len()]; @@ -583,6 +655,7 @@ impl GraphIndex { self.edges.clear(); self.files.clear(); self.name_records.clear(); + self.sorted_name_keys.clear(); self.next_id = SYMBOL_BASE; // ── Phase 1: register + remap ── @@ -1606,74 +1679,258 @@ impl GraphIndex { limit: usize, offset: usize, ) -> Result<(Vec, usize)> { + let out = self + .search_symbol_paged_resumable( + query, + kind, + mode, + Pagination { limit, offset }, + None, + None, + ) + .await?; + Ok((out.page, out.total)) + } + /// Phiên bản resumable + deadline-aware của [`search_symbol_paged`]: ngắt + /// giữa chừng khi `deadline` hết hạn, trả `PagedSearchOutcome { timed_out: + /// true, cursor: Some(phase dở) }` — caller gọi lại với `resume = + /// Some(cursor)` để tiếp tục từ đúng vị trí (không lặp phần đã duyệt). + /// + /// - Phase A (Contains: engine resumable; khác: scan `sorted_name_keys`) + /// sinh danh sách tên khớp **đã sort** — chỉ hoàn tất sau khi quét hết + /// (không thể sort từng phần vì tên mới có thể chèn vào giữa). + /// - Phase B (`Expand`) stream ids theo tên đã sort — `name_index[name]` + /// tăng dần theo id nên collected đã sort (name, id) — **không cần sort + /// cuối** (diệt O(n log n) của bản cũ). total chỉ chính xác khi quét xong. + /// - Hoàn tất + còn page sau → `cursor = Some(Paged)` để phân trang tiếp + /// không cần quét lại. + /// + /// `resume` phải khớp (query, mode, kind) — sai → `InvalidArgument`. + pub async fn search_symbol_paged_resumable( + &self, + query: &str, + kind: Option, + mode: SymbolMatch, + pagination: Pagination, + resume: Option, + deadline: Option, + ) -> Result { let q = query.to_lowercase(); - let mut seen = HashSet::new(); - let mut ids: Vec = Vec::new(); - match mode { - // Substring qua name engine (radix — nhanh hơn duyệt toàn bộ tên). - SymbolMatch::Contains => { - let hits = match self.names.search(q.as_bytes(), None).await { - Ok(h) => h, - Err(_) => return Ok((Vec::new(), 0)), - }; - for (record, _) in hits { - if record == 0 { - continue; - } - let Some(name) = self.name_records.get(record - 1) else { - continue; - }; - let Some(name_ids) = self.name_index.get(name) else { - continue; + + // Validate resume: query/mode/kind phải khớp (limit/offset không — page + // có thể đổi giữa chừng khi retry). + if let Some(c) = &resume + && (c.query != q || c.mode != mode || c.kind != kind) + { + return Err(Error::Invalid( + "resume cursor does not match this query".into(), + )); + } + + // ── Khôi phục / khởi tạo phase ── + let (mut phase, mut timed_out) = match resume.map(|c| c.phase) { + Some(p) => (p, false), + None => ( + match mode { + SymbolMatch::Contains => SearchCursorPhase::Engine(SearchResume::default()), + _ => SearchCursorPhase::ScanNames { name_pos: 0 }, + }, + false, + ), + }; + + // ── Phase A: sinh danh sách tên khớp (sort) ── + match &mut phase { + SearchCursorPhase::Engine(sr) => { + let page = self + .names + .search_resumable(q.as_bytes(), None, Some(sr.clone()), deadline) + .await?; + if page.timed_out { + phase = SearchCursorPhase::Engine(page.resume.unwrap_or_default()); + timed_out = true; + } else { + // record → tên, sort → Expand. + let mut names: Vec = page + .record_ids + .iter() + .filter_map(|&r| { + if r == 0 { + return None; + } + self.name_records.get(r - 1).cloned() + }) + .collect(); + names.sort(); + phase = SearchCursorPhase::Expand { + names, + name_idx: 0, + id_idx: 0, + collected: Vec::new(), }; - for &id in name_ids { - if seen.insert(id) { - ids.push(id); - } - } } } - // Prefix/suffix/exact duyệt name_index (bộ nhỏ hơn symbol registry). - SymbolMatch::Prefix | SymbolMatch::Suffix | SymbolMatch::Exact => { - for (name, name_ids) in &self.name_index { - let matched = match mode { + SearchCursorPhase::ScanNames { name_pos } => { + let mut matched: Vec = Vec::new(); + let mut pos = *name_pos; + loop { + if let Some(dl) = deadline + && Instant::now() >= dl + { + phase = SearchCursorPhase::ScanNames { name_pos: pos }; + timed_out = true; + break; + } + if pos >= self.sorted_name_keys.len() { + break; + } + let name = &self.sorted_name_keys[pos]; + let ok = match mode { SymbolMatch::Prefix => name.starts_with(&q), SymbolMatch::Suffix => name.ends_with(&q), SymbolMatch::Exact => name == &q, _ => false, }; - if !matched { - continue; + if ok { + matched.push(name.clone()); } - for &id in name_ids { - if seen.insert(id) { - ids.push(id); - } + pos += 1; + } + if !timed_out { + phase = SearchCursorPhase::Expand { + names: matched, + name_idx: 0, + id_idx: 0, + collected: Vec::new(), + }; + } + } + // Phase A xong rồi (timed out ở phase B trước) — không làm gì. + _ => {} + } + + // ── Phase B: stream ids theo tên đã sort → collected ── + if !timed_out + && let SearchCursorPhase::Expand { + names, + name_idx, + id_idx, + collected, + } = &mut phase + { + loop { + if let Some(dl) = deadline + && Instant::now() >= dl + { + timed_out = true; + break; + } + if *name_idx >= names.len() { + break; + } + let name = &names[*name_idx]; + let Some(name_ids) = self.name_index.get(name) else { + *name_idx += 1; + *id_idx = 0; + continue; + }; + if *id_idx >= name_ids.len() { + *name_idx += 1; + *id_idx = 0; + continue; + } + let id = name_ids[*id_idx]; + *id_idx += 1; + if let Some(k) = kind { + if self.symbols.get(&id).is_some_and(|s| s.kind == k) { + collected.push(id); } + } else { + collected.push(id); } } } - let mut all: Vec = ids - .into_iter() - .filter(|&id| match kind { - Some(k) => self.symbols.get(&id).is_some_and(|s| s.kind == k), - None => true, - }) - .collect(); - all.sort_by(|&a, &b| { - let na = self.symbols.get(&a).map(|s| s.name.as_str()).unwrap_or(""); - let nb = self.symbols.get(&b).map(|s| s.name.as_str()).unwrap_or(""); - na.cmp(nb).then(a.cmp(&b)) - }); - let total = all.len(); - let limit = if limit == 0 { usize::MAX } else { limit }; - let page = all - .into_iter() - .skip(offset) - .take(limit) - .filter_map(|id| self.symbols.get(&id).cloned()) + + // ── Trang kết quả + cursor ── + if timed_out { + let progress = match &phase { + SearchCursorPhase::Engine(sr) => sr.record_ids.len(), + SearchCursorPhase::ScanNames { name_pos } => *name_pos, + SearchCursorPhase::Expand { collected, .. } => collected.len(), + SearchCursorPhase::Paged { .. } => 0, + }; + return Ok(PagedSearchOutcome { + page: Vec::new(), + total: 0, + timed_out: true, + progress, + cursor: Some(SearchCursor { + query: q, + mode, + kind, + phase, + }), + }); + } + + // Hoàn tất: lấy collected + total từ Expand, hoặc dùng thẳng từ Paged. + let (collected, total) = match &phase { + SearchCursorPhase::Expand { collected, .. } => { + let total = collected.len(); + (collected.clone(), total) + } + SearchCursorPhase::Paged { collected, total } => (collected.clone(), *total), + // Không thể tới đây khi chưa hoàn tất phase A. + _ => (Vec::new(), 0), + }; + + let cap = if pagination.limit == 0 { + usize::MAX + } else { + pagination.limit + }; + let page: Vec = collected + .iter() + .skip(pagination.offset) + .take(cap) + .filter_map(|id| self.symbols.get(id).cloned()) .collect(); - Ok((page, total)) + let page_end = pagination.offset + page.len(); + let more = page_end < total; + + Ok(PagedSearchOutcome { + page, + total, + timed_out: false, + progress: total, + cursor: more.then_some(SearchCursor { + query: q, + mode, + kind, + phase: SearchCursorPhase::Paged { collected, total }, + }), + }) + } + + /// Resumable + deadline-aware của `search_symbol_filtered` (mode Contains, + /// không lọc kind) — nền cho `codegraph_search`. `limit` chặn số symbol + /// trả về; kết quả sort theo (name, id). + pub async fn search_symbol_resumable( + &self, + query: &str, + limit: usize, + resume: Option, + deadline: Option, + ) -> Result { + self.search_symbol_paged_resumable( + query, + None, + SymbolMatch::Contains, + Pagination { limit, offset: 0 }, + resume, + deadline, + ) + .await } /// Số liệu tổng hợp. @@ -2226,8 +2483,8 @@ mod tests { .unwrap(); assert_eq!(total, 1); assert_eq!(hits[0].name, "validate"); - // contains + pagination. Sort byte-wise (case-sensitive): uppercase - // "Order*" đứng trước "getOrders". + // contains + pagination. Sort theo tên lowercase (nhất quán với search + // case-insensitive): "getorders" đứng trước "order*". let (page0, total) = idx .search_symbol_paged("order", None, SymbolMatch::Contains, 2, 0) .await @@ -2237,15 +2494,216 @@ mod tests { "OrderService, OrderController, OrderRepository + getOrders" ); assert_eq!(page0.len(), 2); - assert_eq!(page0[0].name, "OrderController"); - assert_eq!(page0[1].name, "OrderRepository"); + assert_eq!(page0[0].name, "getOrders"); + assert_eq!(page0[1].name, "OrderController"); let (page1, _) = idx .search_symbol_paged("order", None, SymbolMatch::Contains, 2, 2) .await .unwrap(); assert_eq!(page1.len(), 2); - assert_eq!(page1[0].name, "OrderService"); - assert_eq!(page1[1].name, "getOrders"); + assert_eq!(page1[0].name, "OrderRepository"); + assert_eq!(page1[1].name, "OrderService"); + } + + /// Resumable search == direct search, mọi mode + kind filter + offset + tên + /// trùng. Dùng deadline ĐÃ HẾT HẠN ở call đầu — monotonic clock nên check + /// `now >= deadline` luôn true → chắc chắn timed_out ngay, sinh checkpoint + /// hợp lệ; call sau resume không deadline → hoàn tất. Kết quả phải khớp + /// `search_symbol_paged` (không lặp, không mất, cùng thứ tự). + #[tokio::test] + async fn search_paged_resumable_matches_direct() { + let mut idx = GraphIndex::in_memory(); + let mut results = Vec::new(); + let mut id = SYMBOL_BASE; + // 600 tên "order_*" (Function). + for i in 0..600 { + let name = format!("order_{i:04}"); + results.push(result( + "a.ts", + vec![sym("a.ts", &name, id)], + HashMap::new(), + vec![], + )); + id += 1; + } + // 300 symbol Class TRÙNG TÊN "OrderService" — dedup theo (name, id), + // mỗi id là 1 kết quả riêng (test kind filter + duplicate names). + for _ in 0..300 { + let mut s = sym("a.ts", "OrderService", id); + s.kind = SymbolKind::Class; + results.push(result("a.ts", vec![s], HashMap::new(), vec![])); + id += 1; + } + // 100 tên "x_order_*" — khớp contains, không khớp prefix/exact/suffix "service". + for i in 0..100 { + let name = format!("x_order_{i:03}"); + results.push(result( + "a.ts", + vec![sym("a.ts", &name, id)], + HashMap::new(), + vec![], + )); + id += 1; + } + idx.ingest(&results).await.unwrap(); + + // Driver: call đầu deadline hết hạn → timed_out (sinh checkpoint); call + // sau resume không deadline → hoàn tất. So với direct (không deadline). + async fn chained( + idx: &GraphIndex, + q: &str, + kind: Option, + mode: SymbolMatch, + limit: usize, + offset: usize, + ) -> PagedSearchOutcome { + let first = idx + .search_symbol_paged_resumable( + q, + kind, + mode, + Pagination { limit, offset }, + None, + Some(Instant::now()), + ) + .await + .unwrap(); + assert!(first.timed_out, "expired deadline must time out"); + assert!(first.cursor.is_some(), "timeout must carry a cursor"); + let out = idx + .search_symbol_paged_resumable( + q, + kind, + mode, + Pagination { limit, offset }, + first.cursor, + None, + ) + .await + .unwrap(); + assert!(!out.timed_out, "resume without deadline must complete"); + out + } + + let cases: Vec<(&str, Option, SymbolMatch, usize, usize)> = vec![ + ("order", None, SymbolMatch::Contains, 20, 0), + ( + "order", + Some(SymbolKind::Function), + SymbolMatch::Contains, + 20, + 0, + ), + ( + "order", + Some(SymbolKind::Class), + SymbolMatch::Contains, + 30, + 7, + ), + ("order", None, SymbolMatch::Prefix, 10, 5), + ("service", None, SymbolMatch::Suffix, 20, 0), + ( + "orderservice", + Some(SymbolKind::Class), + SymbolMatch::Exact, + 20, + 0, + ), + ]; + for (q, kind, mode, limit, offset) in cases { + let (direct_page, direct_total) = idx + .search_symbol_paged(q, kind, mode, limit, offset) + .await + .unwrap(); + let chained = chained(&idx, q, kind, mode, limit, offset).await; + assert_eq!( + chained.total, direct_total, + "total differs for {q} {mode:?} {kind:?}" + ); + let got: Vec<(u64, String)> = chained + .page + .iter() + .map(|s| (s.id, s.name.clone())) + .collect(); + let want: Vec<(u64, String)> = + direct_page.iter().map(|s| (s.id, s.name.clone())).collect(); + assert_eq!(got, want, "page differs for {q} {mode:?} {kind:?}"); + } + } + + /// Multi-step resume chain: seed đủ lớn, deadline ngắn thật → mỗi call làm + /// ≥1 bước rồi timed out (có thể nhiều lần), chain tới khi hoàn tất. Kết + /// quả cuối phải khớp direct. Vòng lặp có cận phòng hờ (entry-expiry dưới + /// tải nặng có thể làm 1 call không tiến) — fail hẳn thay vì treo. + #[tokio::test] + async fn search_paged_resumable_multistep_chain() { + let mut idx = GraphIndex::in_memory(); + let mut results = Vec::new(); + for (id, i) in (SYMBOL_BASE..).zip(0..4000) { + let name = format!("order_{i:04}"); + results.push(result( + "a.ts", + vec![sym("a.ts", &name, id)], + HashMap::new(), + vec![], + )); + } + idx.ingest(&results).await.unwrap(); + + let (direct_page, direct_total) = idx + .search_symbol_paged("order", None, SymbolMatch::Contains, 10, 0) + .await + .unwrap(); + assert_eq!(direct_total, 4000); + + // Call đầu deadline hết hạn → chắc chắn timed_out (tạo checkpoint). + let mut cursor = idx + .search_symbol_paged_resumable( + "order", + None, + SymbolMatch::Contains, + Pagination { + limit: 10, + offset: 0, + }, + None, + Some(Instant::now()), + ) + .await + .unwrap() + .cursor; + assert!(cursor.is_some()); + + // Resume với deadline thật ngắn — lặp tới khi hoàn tất (mỗi lần tiến ≥1 bước). + let mut completed = None; + for _ in 0..2000 { + let out = idx + .search_symbol_paged_resumable( + "order", + None, + SymbolMatch::Contains, + Pagination { + limit: 10, + offset: 0, + }, + cursor, + Some(Instant::now() + std::time::Duration::from_millis(10)), + ) + .await + .unwrap(); + if out.timed_out { + cursor = out.cursor; + continue; + } + completed = Some(out); + break; + } + let out = completed.expect("resume chain must terminate"); + assert_eq!(out.total, direct_total); + let got: Vec<(u64, String)> = out.page.iter().map(|s| (s.id, s.name.clone())).collect(); + let want: Vec<(u64, String)> = direct_page.iter().map(|s| (s.id, s.name.clone())).collect(); + assert_eq!(got, want); } /// dependencies_report — module prefix từ call names, internal vs external. diff --git a/crates/codegraph-graph/src/radix.rs b/crates/codegraph-graph/src/radix.rs index 8654b8001..8973c1976 100644 --- a/crates/codegraph-graph/src/radix.rs +++ b/crates/codegraph-graph/src/radix.rs @@ -118,6 +118,44 @@ pub type SearchMatcher = Arc OnMatchCallback + S /// shortcuts/cache dựa trên `old_prefix` + `breakpoint` rồi để radix commit. pub type OnSplitCallback = Arc Result<()> + Send + Sync>; +// ==================== Resumable DFS ==================== + +/// Frame trên work-stack của `Radix::search_dfs_resumable`. +/// +/// Chỉ lưu 4 số — `prefix`/`continuations`/`children` được recompute từ +/// `node_id` khi xử lý (matcher deterministic theo `(prefix, pattern, +/// pattern_pos)`), nên checkpoint nhỏ và resume chính xác. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DfsFrame { + pub node_id: usize, + pub pattern_pos: usize, + pub cont_idx: usize, + pub child_idx: usize, +} + +/// Trạng thái duyệt hiện tại của `Radix::search_dfs_resumable` khi bị deadline +/// ngắt giữa chừng. +#[derive(Debug, Clone)] +pub enum DfsState { + /// Đang dò xuống children (chưa tìm thấy match hoàn chỉnh). + Search(Vec), + /// Đang collect toàn bộ records trong subtree của `root` (sau khi matcher + /// báo `found`). Stack = `(node_id, child_idx)` — duyệt pre-order. + Collect { + root: usize, + stack: Vec<(usize, usize)>, + }, +} + +/// Checkpoint của một lần duyệt bị ngắt — resume từ đây. +#[derive(Debug, Clone, Default)] +pub struct DfsCheckpoint { + /// `None` = đã duyệt xong (caller advance sang candidate khác). + pub state: Option, + /// Records đã collect được tính tới lúc ngắt. + pub records: Vec, +} + /// Callback khi chạm tới một node cụ thể, chứa thông tin đầy đủ về node /// đó dưới dạng metadata, có cấu trúc dạng node, metadata và trả về id của /// node, lưu ý vì đây là callback access nên nó có thể bị trùng hoặc gọi lại @@ -612,127 +650,209 @@ impl Radix { /// Trả về record IDs của match đầu tiên theo DFS trong mỗi subtree (khớp /// hành vi `search_index::search_like`). Không kèm meta/key length — đó là /// concern của caller (`Search` lưu chúng trong Storage). + /// + /// Wrapper không deadline cho tests; production path (`Search`) dùng + /// [`Self::search_dfs_resumable`] để cancel giữa chừng. + #[cfg_attr(not(test), allow(dead_code))] pub async fn search_dfs( &self, begin: usize, pattern: &[T], matcher: SearchMatcher, ) -> Result> { - let mut records = Vec::new(); - - if pattern.is_empty() { - return Err(Error::NotFound); - } - - let node_id = if begin == EMPTY { - self.storage - .read() - .await - .get_root(shard_of(pattern[0], self.sharding)) - .await? - } else { - begin - }; - - if node_id == EMPTY { - return Ok(Vec::new()); - } - - self.search_dfs_iter(node_id, pattern, matcher, 0, &mut records) + let (records, _) = self + .search_dfs_resumable(begin, pattern, matcher, None, None) .await?; Ok(records) } - /// DFS dùng `matcher`: đọc prefix của `node_id`, hỏi matcher, rồi quyết - /// định collect subtree / đệ quy xuống children theo `continuations`. + /// Như [`search_dfs`](Self::search_dfs) nhưng **resumable + deadline-aware**: + /// duyệt bằng explicit work-stack (không async recursion) nên ngắt được giữa + /// chừng khi `deadline` hết hạn. Khi ngắt: trả `(records, Some(checkpoint))` — + /// caller gọi lại với `resume = Some(checkpoint)` để tiếp tục chính xác từ vị + /// trí dừng; hoàn tất không timeout: `None` ở vị trí checkpoint. /// - /// `pattern_pos` tại node entry luôn là vị trí pattern bắt đầu dò trên - /// prefix của node này (data_pos = 0). - #[inline] - async fn search_dfs_iter( + /// Semantics giữ nguyên `search_dfs`: node đầu tiên (theo DFS) có pattern + /// khớp hoàn chỉnh trong prefix → collect toàn bộ records của subtree đó rồi + /// dừng (short-circuit); prefix hết mà pattern chưa khớp hết → dò xuống + /// children theo `continuations` matcher trả về. + pub async fn search_dfs_resumable( &self, - node_id: usize, + begin: usize, pattern: &[T], matcher: SearchMatcher, - pattern_pos: usize, - out: &mut Vec, - ) -> Result<()> { - let (prefix_bytes, _record) = { self.storage.read().await.get_node(node_id).await? }; - let prefix = Self::to_vec(&prefix_bytes); - - let result = matcher(&prefix, pattern, pattern_pos); - - // Match hoàn chỉnh → collect toàn bộ records trong subtree. - if result.found { - self.collect_subtree_records(node_id, out).await?; - return Ok(()); + resume: Option, + deadline: Option, + ) -> Result<(Vec, Option)> { + if pattern.is_empty() { + return Err(Error::NotFound); } - // Với mỗi vị trí pattern mà matcher cho phép tiếp tục, đi xuống - // child có element đầu khớp pattern[pp]. Short-circuit ở match đầu - // tiên trong subtree (khớp dfs_search cũ của search_index). - let children = { self.storage.read().await.get_children(node_id).await? }; - for pp in result.continuations { - if pp == 0 || pp >= pattern.len() { - continue; + // Trạng thái: từ checkpoint (resume) hoặc khởi tạo từ `begin`. + let (mut state, mut records) = if let Some(cp) = resume { + (cp.state, cp.records) + } else { + let node_id = if begin == EMPTY { + self.storage + .read() + .await + .get_root(shard_of(pattern[0], self.sharding)) + .await? + } else { + begin + }; + if node_id == EMPTY { + return Ok((Vec::new(), None)); } + ( + Some(DfsState::Search(vec![DfsFrame { + node_id, + pattern_pos: 0, + cont_idx: 0, + child_idx: 0, + }])), + Vec::new(), + ) + }; - let next_elem = pattern[pp]; - for &child in &children { - let (cp_bytes, _) = { self.storage.read().await.get_node(child).await? }; - let cp = Self::to_vec(&cp_bytes); + // Mỗi vòng lặp xử lý đúng 1 bước duyệt; giữa các bước check deadline. + // `state = None` → duyệt xong (Search không match / Collect xong). + while let Some(cur) = state.take() { + if let Some(dl) = deadline + && std::time::Instant::now() >= dl + { + return Ok(( + Vec::new(), + Some(DfsCheckpoint { + state: Some(cur), + records, + }), + )); + } - if !cp.is_empty() && cp[0] == next_elem { - // Prune nhánh: bloom của child không chứa `pattern[pp..]` - // (substring) → subtree chắc chắn không có match tiếp tục, - // bỏ nhánh. Bloom có 0 false negative nên không bao giờ bỏ - // nhánh có match thật. Chỉ prune khi substring đủ ngắn và - // child có bloom (không có → fallback full traversal). - #[cfg(feature = "bloom-search")] - { - let remaining_len = pattern.len() - pp; - if remaining_len <= bloom_cfg::MATCH_CAP { - let bloom_bytes = - { self.storage.read().await.get_node_bloom(child).await? }; - if let Some(bloom_bytes) = bloom_bytes - && let Some(bf) = BloomFilter::deserialize(&bloom_bytes) - && !bf.contains(&Self::from_vec(&pattern[pp..])) - { + state = match cur { + DfsState::Search(mut stack) => { + // Bước tới: pop frame, đọc prefix, hỏi matcher. Found → chuyển + // sang Collect; ngược lại tìm child khớp element tiếp theo. + let mut next: Option = None; + while next.is_none() { + let Some(mut frame) = stack.pop() else { + break; // stack rỗng — không có match trong subtree này. + }; + + let (prefix_bytes, _record) = + { self.storage.read().await.get_node(frame.node_id).await? }; + let prefix = Self::to_vec(&prefix_bytes); + let result = matcher(&prefix, pattern, frame.pattern_pos); + + // Match hoàn chỉnh → collect toàn bộ subtree rồi dừng. + if result.found { + next = Some(DfsState::Collect { + root: frame.node_id, + stack: vec![(frame.node_id, 0)], + }); + break; + } + + let children = { + self.storage + .read() + .await + .get_children(frame.node_id) + .await? + }; + let mut descended = false; + while frame.cont_idx < result.continuations.len() { + let pp = result.continuations[frame.cont_idx]; + if pp == 0 || pp >= pattern.len() { + frame.cont_idx += 1; + frame.child_idx = 0; continue; } + + let next_elem = pattern[pp]; + while frame.child_idx < children.len() { + let child = children[frame.child_idx]; + frame.child_idx += 1; + let (cp_bytes, _) = + { self.storage.read().await.get_node(child).await? }; + let cp = Self::to_vec(&cp_bytes); + if cp.is_empty() || cp[0] != next_elem { + continue; + } + + // Prune nhánh: bloom của child không chứa + // `pattern[pp..]` (substring) → subtree chắc chắn + // không có match tiếp tục, bỏ nhánh. Bloom có 0 + // false negative nên không bao giờ bỏ nhánh có + // match thật. Chỉ prune khi substring đủ ngắn và + // child có bloom (không có → fallback traversal). + #[cfg(feature = "bloom-search")] + { + let remaining_len = pattern.len() - pp; + if remaining_len <= bloom_cfg::MATCH_CAP { + let bloom_bytes = { + self.storage.read().await.get_node_bloom(child).await? + }; + if let Some(bloom_bytes) = bloom_bytes + && let Some(bf) = BloomFilter::deserialize(&bloom_bytes) + && !bf.contains(&Self::from_vec(&pattern[pp..])) + { + continue; + } + } + } + + // Đi xuống child — đẩy frame hiện tại lại (với vị + // trí đã tiến) + frame con mới. + stack.push(frame); + stack.push(DfsFrame { + node_id: child, + pattern_pos: pp, + cont_idx: 0, + child_idx: 0, + }); + descended = true; + break; + } + if descended { + break; + } + frame.cont_idx += 1; + frame.child_idx = 0; + } + if descended { + next = Some(DfsState::Search(stack)); + break; } + // Frame này đã dò hết continuations — pop frame tiếp theo. } - - Box::pin(self.search_dfs_iter(child, pattern, matcher.clone(), pp, out)) - .await?; - if !out.is_empty() { - return Ok(()); + // `None` = stack rỗng không có match → candidate xong. + next + } + DfsState::Collect { root, mut stack } => { + // Collect subtree theo pre-order (record của node trước, sau + // đó mới children — giống bản đệ quy cũ). + if let Some((node_id, child_idx)) = stack.pop() { + let (_prefix_bytes, record) = + { self.storage.read().await.get_node(node_id).await? }; + if record != EMPTY { + records.push(record); + } + let children = { self.storage.read().await.get_children(node_id).await? }; + if child_idx < children.len() { + stack.push((node_id, child_idx + 1)); + stack.push((children[child_idx], 0)); + } + Some(DfsState::Collect { root, stack }) + } else { + None // Collect xong — candidate đã có records, dừng. } } - } + }; } - Ok(()) - } - - /// Collect toàn bộ record IDs trong subtree của `node_id` (DFS). - #[inline] - async fn collect_subtree_records( - &self, - node_id: usize, - records: &mut Vec, - ) -> Result<()> { - let (_prefix_bytes, record) = { self.storage.read().await.get_node(node_id).await? }; - if record != EMPTY { - records.push(record); - } - - let children = { self.storage.read().await.get_children(node_id).await? }; - for &child in &children { - Box::pin(self.collect_subtree_records(child, records)).await?; - } - - Ok(()) + Ok((records, None)) } /// Chẻ `parent` tại `breakpoint`: diff --git a/crates/codegraph-graph/src/search.rs b/crates/codegraph-graph/src/search.rs index 51d1c36da..43d408ce3 100644 --- a/crates/codegraph-graph/src/search.rs +++ b/crates/codegraph-graph/src/search.rs @@ -22,11 +22,13 @@ use std::collections::HashSet; use std::sync::{Arc, Mutex}; +use std::time::Instant; use tokio::sync::RwLock; use crate::radix::{ - self, EMPTY, Element, OnMatchCallback, OnNodeAccessCallback, Radix, SearchMatcher, + self, DfsCheckpoint, EMPTY, Element, OnMatchCallback, OnNodeAccessCallback, Radix, + SearchMatcher, }; use crate::storage::{InMemoryStorage, Storage}; @@ -35,6 +37,37 @@ use crate::storage::{InMemoryStorage, Storage}; /// Giới hạn cứng số kết quả trả về (khớp `codegraph-graph::HARD_LIMIT`). const MAX_RESULTS: usize = 5000; +// ==================== Resumable search ==================== + +/// Trạng thái resume của một lần [`Search::search_resumable`] bị ngắt bởi +/// deadline — caller gọi lại với cùng pattern + `resume` này để tiếp tục từ +/// đúng vị trí dừng (candidates recompute từ storage — deterministic trong một +/// snapshot, nên chỉ cần lưu vị trí + trạng thái DFS). +#[derive(Debug, Clone, Default)] +pub struct SearchResume { + /// Candidate tiếp theo cần xử lý (index vào shortcut candidates). + pub cand_idx: usize, + /// Trạng thái DFS của candidate hiện tại (`None` = giữa các candidate — + /// chưa xử lý candidate nào dở). + pub dfs: Option, + /// Records đã collect (dedup chéo candidates) tính tới lúc ngắt. + pub record_ids: Vec, + /// Vị trí trong phase resolve (filter `depth`) nếu bị ngắt ở đó. + pub resolve_idx: usize, +} + +/// Kết quả của [`Search::search_resumable`]. +#[derive(Debug, Clone, Default)] +pub struct SearchPage { + /// Records khớp (record idx). Khi `timed_out` — records đã collect tới lúc + /// ngắt (cũng nằm trong `resume.record_ids`). + pub record_ids: Vec, + /// `Some` = bị ngắt giữa chừng — caller phải gọi lại với `resume` này. + pub resume: Option, + /// `true` khi `resume` có nghĩa (deadline đã hết hạn giữa chừng). + pub timed_out: bool, +} + // ==================== Error ==================== #[derive(Debug)] @@ -419,6 +452,48 @@ impl Search { pattern: &[T], depth: Option, ) -> Result>)>> { + let page = self.search_resumable(pattern, depth, None, None).await?; + if page.record_ids.is_empty() { + return Err(Error::NotFound); + } + // Resolve meta (API cũ giữ nguyên) — record_ids đã được filter `depth` + // ở search_resumable nên chỉ cần đọc meta. + let storage = self.storage.read().await; + let mut results = Vec::new(); + for &rid in &page.record_ids { + if rid == EMPTY { + continue; + } + let meta = storage.get_meta(rid).await?; + results.push((rid, meta)); + if results.len() >= MAX_RESULTS { + break; + } + } + if results.is_empty() { + Err(Error::NotFound) + } else { + Ok(results) + } + } + + /// Như [`search`](Self::search) nhưng **resumable + deadline-aware** — dùng + /// khi index lớn làm query chạy lâu. Khi `deadline` hết hạn giữa chừng: trả + /// `SearchPage { timed_out: true, resume: Some(...) }` — caller gọi lại với + /// `resume` để tiếp tục từ đúng vị trí dừng (không lặp phần đã duyệt, không + /// mất records đã collect). Hoàn tất: `timed_out: false`, `resume: None`. + /// + /// Khác `search`: trả `record_ids` (`Vec`) không kèm meta — callers + /// hiện tại (name engine, chain engine) không dùng meta; `search` giữ API cũ. + /// + /// `resume: None` = search mới. `deadline: None` = chạy tới cùng. + pub async fn search_resumable( + &self, + pattern: &[T], + depth: Option, + resume: Option, + deadline: Option, + ) -> Result { if pattern.is_empty() { return Err(Error::NotFound); } @@ -426,7 +501,8 @@ impl Search { let first_elem = pattern[0]; let si = radix::shard_of(first_elem, self.sharding); - // Query candidates trực tiếp từ storage. + // Query candidates trực tiếp từ storage (deterministic per snapshot — + // resume chỉ cần cand_idx, không cần lưu candidates). let candidates = self .storage .read() @@ -437,60 +513,104 @@ impl Search { // depth = max hop → max key length (số element) = depth + 1. let max_len = depth.map(|d| d + 1); - // Mỗi candidate: `Radix::search_dfs` chạy matcher KMP trong subtree và - // trả record IDs (logic trie không còn nằm ở đây). Dedup chéo candidates - // — subtree của candidate này có thể chứa subtree của candidate khác. + let (mut cand_idx, mut dfs, mut record_ids, resolve_idx) = match resume { + Some(r) => (r.cand_idx, r.dfs, r.record_ids, r.resolve_idx), + None => (0, None, Vec::new(), 0), + }; + // `seen` = tập record_ids đã collect (dedup chéo candidates — subtree + // của candidate này có thể chứa subtree của candidate khác). + let mut seen: HashSet = record_ids.iter().copied().collect(); let matcher = kmp_matcher(pattern); - let mut seen = HashSet::new(); - let mut record_ids = Vec::new(); - for &node_id in &candidates { - if record_ids.len() >= MAX_RESULTS { - break; + + // ── Candidate loop ── + while cand_idx < candidates.len() { + if let Some(dl) = deadline + && Instant::now() >= dl + { + return Ok(SearchPage { + record_ids: record_ids.clone(), + resume: Some(SearchResume { + cand_idx, + dfs, + record_ids: record_ids.clone(), + resolve_idx: 0, + }), + timed_out: true, + }); } - for rid in self + + let node_id = candidates[cand_idx]; + let (records, ckpt) = self .trie - .search_dfs(node_id, pattern, matcher.clone()) - .await? - { - if seen.insert(rid) { - record_ids.push(rid); + .search_dfs_resumable(node_id, pattern, matcher.clone(), dfs.take(), deadline) + .await?; + match ckpt { + // Timeout giữa candidate — lưu trạng thái DFS, tiếp tục lần sau. + Some(cp) => { + dfs = Some(cp); + } + // Candidate xong — dedup records (records của candidate này) vào + // kết quả chung. + None => { + for rid in records { + if seen.insert(rid) { + record_ids.push(rid); + if record_ids.len() >= MAX_RESULTS { + break; + } + } + } if record_ids.len() >= MAX_RESULTS { break; } + cand_idx += 1; + dfs = None; } } } - if record_ids.is_empty() { - return Err(Error::NotFound); - } - - // Resolve: filter `depth` (key length trong storage) + đọc meta. - let mut results = Vec::new(); - { + // ── Resolve: filter `depth` (key length trong storage) — deadline-aware. + // Bỏ qua nếu không giới hạn depth (name engine — chiếm đa số query). + if let Some(m) = max_len { + let mut out = Vec::new(); + let mut ridx = resolve_idx; let storage = self.storage.read().await; - for &rid in &record_ids { + loop { + if let Some(dl) = deadline + && Instant::now() >= dl + { + return Ok(SearchPage { + record_ids: out.clone(), + resume: Some(SearchResume { + cand_idx: candidates.len(), + dfs: None, + record_ids: out.clone(), + resolve_idx: ridx, + }), + timed_out: true, + }); + } + if ridx >= record_ids.len() { + break; + } + let rid = record_ids[ridx]; + ridx += 1; if rid == EMPTY { continue; } - if let Some(m) = max_len - && storage.get_key_len(rid).await?.unwrap_or(usize::MAX) > m - { + if storage.get_key_len(rid).await?.unwrap_or(usize::MAX) > m { continue; } - let meta = storage.get_meta(rid).await?; - results.push((rid, meta)); - if results.len() >= MAX_RESULTS { - break; - } + out.push(rid); } + record_ids = out; } - if results.is_empty() { - Err(Error::NotFound) - } else { - Ok(results) - } + Ok(SearchPage { + record_ids, + resume: None, + timed_out: false, + }) } /// Tìm tất cả `(full_key, record)` có key bắt đầu bằng `prefix` (prefix match). diff --git a/crates/codegraph-graph/src/shared.rs b/crates/codegraph-graph/src/shared.rs index 0e8d6bc94..6c328f08f 100644 --- a/crates/codegraph-graph/src/shared.rs +++ b/crates/codegraph-graph/src/shared.rs @@ -81,6 +81,12 @@ impl SharedGraphIndex { /// Chỉ gọi khi `dsn.is_some()`. async fn current_version(&self) -> Option { let dsn = self.dsn.as_ref()?; + // `path` chỉ dùng bởi các backend có probe file độc lập (sqlite/lmdb); + // build không bật backend nào → biến thừa, cho phép bỏ qua lint. + #[cfg_attr( + not(any(feature = "sqlite", feature = "lmdb")), + allow(unused_variables) + )] let path = trim_scheme(dsn); match self.scheme() { #[cfg(feature = "sqlite")] diff --git a/crates/codegraph-mcp/src/lib.rs b/crates/codegraph-mcp/src/lib.rs index a9a86ab59..0a3a87540 100644 --- a/crates/codegraph-mcp/src/lib.rs +++ b/crates/codegraph-mcp/src/lib.rs @@ -15,19 +15,18 @@ pub mod stdio; mod tools; mod usage; -pub use session::{InitOutcome, Session}; +pub use session::{DetailLevel, InitOutcome, OutputStyle, Session}; pub use stdio::serve_stdio; -use std::future::Future; use std::sync::{Arc, Mutex}; -use codegraph_api::GraphApi; +use codegraph_api::{GraphApi, SearchSessionStore}; use rmcp::handler::server::ServerHandler; use rmcp::model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, Implementation, - ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerInfo, + CacheScope, CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, + Implementation, ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerInfo, }; -use rmcp::service::{MaybeSendFuture, RequestContext}; +use rmcp::service::RequestContext; use rmcp::{ErrorData as McpError, RoleServer}; use serde_json::{json, Value}; @@ -41,23 +40,42 @@ pub const SERVER_VERSION: &str = env!("CARGO_PKG_VERSION"); pub struct CodegraphServer { session: Session, usage: Arc>, + /// Session store cho search resumable — sống qua nhiều tool call để resume + /// id (trả về khi timeout) có thể retry được. + search_sessions: Arc, } impl CodegraphServer { /// Server với session trống — `codegraph_init` sẽ bind root trong phiên. pub fn new() -> Self { + Self::new_with_format(OutputStyle::default()) + } + + /// `new()` nhưng seed output format từ CLI lúc khởi động + /// (`codegraph serve --mcp --format=...`). + pub fn new_with_format(format: OutputStyle) -> Self { Self { - session: Session::new(), + session: Session::new_with_format(format), usage: Arc::new(Mutex::new(usage::UsageStats::default())), + search_sessions: Arc::new(SearchSessionStore::new()), } } /// Pre-seed root từ `--path` lúc khởi động (tương đương đã `codegraph_init` /// với root đó, không index thêm). Giữ CLI/watcher flow không vỡ. pub async fn with_root(root: camino::Utf8PathBuf) -> anyhow::Result { + Self::with_root_and_format(root, OutputStyle::default()).await + } + + /// `with_root()` nhưng seed output format từ CLI lúc khởi động. + pub async fn with_root_and_format( + root: camino::Utf8PathBuf, + format: OutputStyle, + ) -> anyhow::Result { Ok(Self { - session: Session::with_root(root).await?, + session: Session::with_root_and_format(root, format).await?, usage: Arc::new(Mutex::new(usage::UsageStats::default())), + search_sessions: Arc::new(SearchSessionStore::new()), }) } @@ -75,7 +93,14 @@ impl CodegraphServer { u.reset(); } drop(u); - let text = serde_json::to_string_pretty(&report).map_err(|e| { + let mut v = serde_json::to_value(&report).map_err(|e| { + McpError::internal_error( + "usage report failed", + Some(json!({"reason": e.to_string()})), + ) + })?; + tools::omit_defaults(&mut v); + let text = serde_json::to_string_pretty(&v).map_err(|e| { McpError::internal_error( "usage report failed", Some(json!({"reason": e.to_string()})), @@ -102,14 +127,29 @@ impl CodegraphServer { // Default = KHÔNG index — bind nhanh, không block user. Agent muốn // data thì chủ động gọi codegraph_index {} (hoặc truyền index=true). let do_index = args.get("index").and_then(|v| v.as_bool()).unwrap_or(false); + // Detail level mặc định cho symbol trong list tools (minimal/medium/verbose). + let detail = args + .get("detail") + .and_then(|v| v.as_str()) + .and_then(DetailLevel::parse) + .unwrap_or_default(); + // Output format (minimize/medium) — None giữ nguyên seed từ CLI. + let format = args + .get("format") + .and_then(|v| v.as_str()) + .and_then(OutputStyle::parse); return match self .session - .init(camino::Utf8PathBuf::from(path), do_index) + .init(camino::Utf8PathBuf::from(path), do_index, detail, format) .await { Ok(out) => { - let mut v = - json!({ "root": out.root.as_str(), "initialized": out.dir.as_str() }); + let mut v = json!({ + "root": out.root.as_str(), + "initialized": out.dir.as_str(), + "detail": detail.as_str(), + "format": self.session.format().await.as_str(), + }); if let Some(stats) = &out.indexed { v["indexed"] = session::stats_json(stats); } @@ -147,12 +187,14 @@ impl CodegraphServer { Ok(sgi) => sgi, Err(e) => return Ok(ToolOutput::Error(e.to_string())), }; - let api = GraphApi::new_with_index(sgi.clone()); + let api = GraphApi::new_with_sessions(sgi.clone(), self.search_sessions.clone()); // ensure_ready chỉ Ok khi session có root — đây chỉ là phòng hờ. let Some(root) = self.session.root().await else { return Ok(ToolOutput::Error("session root unavailable".into())); }; + let detail = self.session.detail().await; + let format = self.session.format().await; let dispatch = match name { "codegraph_sandbox" => tools::dispatch_sandbox(&root, sgi.clone(), args.clone()).await, "codegraph_diff" => tools::dispatch_diff(&root, sgi.clone(), args.clone()).await, @@ -162,14 +204,14 @@ impl CodegraphServer { "codegraph_origin_simulate" => { tools::dispatch_origin_simulate(&root, sgi.clone(), args.clone()).await } - _ => tools::dispatch_with_api(&api, name, args).await, + _ => tools::dispatch_with_api(&api, &root, detail, format, name, args).await, }; match dispatch { Ok(text) => { // Ước lượng source bytes mà answer "thay thế" (file refs trong answer). let source_bytes = match serde_json::from_str::(&text) { - Ok(v) => usage::estimate_source_bytes(&api, &v).await, + Ok(v) => usage::estimate_source_bytes(&api, &v, root.as_str()).await, Err(_) => 0, }; Ok(ToolOutput::Text { text, source_bytes }) @@ -194,7 +236,9 @@ enum ToolOutput { impl ToolOutput { fn json(v: &Value) -> Self { - match serde_json::to_string_pretty(v) { + let mut v = v.clone(); + tools::omit_defaults(&mut v); + match serde_json::to_string_pretty(&v) { Ok(text) => ToolOutput::Text { text, source_bytes: 0, @@ -211,54 +255,52 @@ impl ServerHandler for CodegraphServer { .with_instructions(SERVER_INSTRUCTIONS.to_string()) } - fn list_tools( + async fn list_tools( &self, _request: Option, _context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - async move { - Ok(ListToolsResult { - tools: tools::rmcp_tools(), - ..Default::default() - }) - } + ) -> Result { + // Protocol 2026-07-28 (SEP-2549) bắt buộc `ttlMs`/`cacheScope` trên + // result; client strict (vd ZCode) validate theo schema đó → phải set. + // ttl_ms = 0: kết quả coi như stale ngay, không cache phía client. + Ok(ListToolsResult::with_all_items(tools::rmcp_tools()) + .with_ttl_ms(0) + .with_cache_scope(CacheScope::Public)) } - fn call_tool( + async fn call_tool( &self, request: CallToolRequestParams, _context: RequestContext, - ) -> impl Future> + MaybeSendFuture + '_ { - async move { - let name = request.name.as_ref(); - let args = request.arguments.map(Value::Object).unwrap_or(Value::Null); + ) -> Result { + let name = request.name.as_ref(); + let args = request.arguments.map(Value::Object).unwrap_or(Value::Null); - // Tên tool không tồn tại → protocol error (client thấy lỗi JSON-RPC - // method-not-found, không thấy một "tool ảo"). Chặn sớm trước khi - // chạy vào run_tool để không cho nhầm tool lạ chạy nhánh `_`. - if !tools::is_known_tool(name) { - return Err(McpError::method_not_found::< - rmcp::model::CallToolRequestMethod, - >()); - } + // Tên tool không tồn tại → protocol error (client thấy lỗi JSON-RPC + // method-not-found, không thấy một "tool ảo"). Chặn sớm trước khi + // chạy vào run_tool để không cho nhầm tool lạ chạy nhánh `_`. + if !tools::is_known_tool(name) { + return Err(McpError::method_not_found::< + rmcp::model::CallToolRequestMethod, + >()); + } - match self.run_tool(name, args).await { - Ok(ToolOutput::Text { text, source_bytes }) => { - self.usage - .lock() - .unwrap() - .record(name, text.len() as u64, source_bytes, false); - Ok(CallToolResult::success(vec![ContentBlock::text(text)]).into()) - } - Ok(ToolOutput::Error(msg)) => { - self.usage - .lock() - .unwrap() - .record(name, msg.len() as u64, 0, true); - Ok(CallToolResult::error(vec![ContentBlock::text(msg)]).into()) - } - Err(e) => Err(e), + match self.run_tool(name, args).await { + Ok(ToolOutput::Text { text, source_bytes }) => { + self.usage + .lock() + .unwrap() + .record(name, text.len() as u64, source_bytes, false); + Ok(CallToolResult::success(vec![ContentBlock::text(text)]).into()) + } + Ok(ToolOutput::Error(msg)) => { + self.usage + .lock() + .unwrap() + .record(name, msg.len() as u64, 0, true); + Ok(CallToolResult::error(vec![ContentBlock::text(msg)]).into()) } + Err(e) => Err(e), } } } diff --git a/crates/codegraph-mcp/src/server-instructions.md b/crates/codegraph-mcp/src/server-instructions.md index ee383b7cd..783c9a298 100644 --- a/crates/codegraph-mcp/src/server-instructions.md +++ b/crates/codegraph-mcp/src/server-instructions.md @@ -14,7 +14,9 @@ before querying: and **non-blocking: it does NOT index by default** (`index` defaults to `false`). After binding, call `codegraph_index {}` to build/refresh the index (or pass `"index": true` to `codegraph_init` to index immediately). - Re-running with a different `path` re-points the session. + Re-running with a different `path` re-points the session. Optionally set + the default output detail for list tools with + `"detail": "minimal" | "medium" | "verbose"` (see below). - `codegraph_deinit {}` — release the session (the `.codegraph/` and index files stay on disk). An unbound session **refuses every query tool** until `codegraph_init` binds it again. @@ -72,11 +74,148 @@ when multiple symbols share a name. When a name is ambiguous the tool returns anywhere, default), `prefix`, `suffix` (e.g. `match="suffix", query="Service"` finds every `*Service` class), and `exact`. Use `total` + `offset` to page. +## Large indexes: timeout + resume + +On very large indexes a broad search (`codegraph_search` / `codegraph_search_symbol`) +can exceed its time budget. Both tools accept `timeout_ms` (default `2000`; +`0` = no limit). When the budget runs out mid-search the tool **errors** and +does NOT return partial results — the message includes `"resume": ""` and a +progress count: + +``` +codegraph_search_symbol timed out after 2000ms (collected 134 symbols so far). +Retry the same call with the same arguments plus "resume": "" to continue +the search from where it stopped. +``` + +To explore effectively and continuously: **retry the exact same call with the +same arguments plus the `resume` id** — the search continues exactly where it +stopped (nothing is re-scanned, nothing is lost) and eventually returns the +full results. You can keep retrying as many times as needed; each retry that +times out yields a fresh resume id. + +- Resume ids are **short-lived and in-process**: re-indexing the workspace + (version bump) or restarting the server invalidates them. If a resume id is + rejected, retry the search **without** `resume`. +- A resume id is tied to its query/mode/kind — passing it with different + arguments is rejected; retry without `resume`. +- When `codegraph_search_symbol` completes with more pages available, the + response includes a `resume` id in addition to `total`/`has_more` — pass it + on the next call (with a new `offset`) to page further **without re-scanning** + the index. +- `codegraph_search` on success returns a plain array (no `resume` field); if + you need more results, narrow the query or use `codegraph_search_symbol`. + ## Trust the results Codegraph returns AST-derived structural data. Do NOT re-verify with grep — that's slower, less accurate, and wastes context. +## Output detail & token usage + +Symbols in list-tool responses (`codegraph_search`, `codegraph_callers`, +`codegraph_callees`, `codegraph_impact`, `codegraph_search_symbol`, +`codegraph_search_by_annotation`, `codegraph_list_classes`, +`codegraph_list_interfaces`, and the symbol embedded in `codegraph_flow`) are +compacted by default to keep responses token-lean. Under `format=medium` the +`detail` level selects which fields appear; under `format=minimize` (default) +`detail` is ignored — see [Response formats](#response-formats-binance-style-minimal). + +- **Session-wide default** is set at bind time: `codegraph_init {"path": ..., + "detail": "minimal"}` (or re-run `codegraph_init` to change it later). +- **Per-call override** — any list tool accepts a `detail` arg that wins over + the session default for that one call. + +Levels: +- `minimal` — `{id, name, kind, file, line}`. Fewest tokens; best for + scanning long lists. +- `medium` (default) — adds `signature` (the declaration line). Enough for + most reasoning. +- `verbose` — the full `Symbol` (doc comments, annotations, scope, type_ref, + end_line, language). Use only when you actually need those fields; + `codegraph_symbol {"id": ...}` returns the full symbol for a single target. + +`file` paths in responses are **relative to the workspace root** (the `root` +returned by `codegraph_init`). To keep context lean, prefer smaller `limit` +values and `id`-based lookups over re-running broad searches. + +## Response formats (Binance-style minimal) + +Every response is minimal by default. A `format` knob selects between two +styles — set at server startup (`codegraph serve --mcp --format=...`, default +`minimize`), per session (`codegraph_init {"format": ...}`), or per call +(`"format": ...` arg on any tool, which wins over both): + +- **`minimize`** (default) — symbol items are **positional arrays** with a + fixed, documented order (see the schema below). No keys, no per-item JSON + overhead — this is the "remove the key, keep only the value" style. +- **`medium`** — objects keep their keys; fields whose value is the default + (`null`, `false`, `""`, `[]`, `{}`, and numeric `0` for the sentinels + `scope_id` / `type_ref` / `end_line`) are **omitted entirely**. Counts and + totals (`total`, `limit`, `offset`, `symbols`, `files`, ...) always stay, + even when `0`, so summary responses stay readable. + +The omission rule applies to **every object in both formats** — wrapper +metadata such as `resume: null`, `has_more: false`, `truncated: false`, +`deleted: false` disappears when it holds the default value. **Absent means +default.** Arrays never omit positions. + +### Symbol array schema (`format=minimize`) + +Each symbol is a fixed 14-element array. The order is part of the contract — +never reorder or truncate it: + +| # | field | type | absent = | +|---|-------|------|----------| +| 0 | `id` | number | — | +| 1 | `name` | string | — | +| 2 | `kind` | string (`function`, `method`, `class`, …) | — | +| 3 | `scope` | string (`global`, `object_field`, `local`, `parameter`) | — | +| 4 | `scope_id` | number | `0` = global | +| 5 | `type_ref` | number | `0` = none | +| 6 | `type_name` | string \| `null` | `null` = none | +| 7 | `file` | string | relative to workspace root | +| 8 | `line` | number | — | +| 9 | `end_line` | number | `0` = not recorded | +| 10 | `signature` | string \| `null` | `null` = none | +| 11 | `doc` | string \| `null` | `null` = none | +| 12 | `annotations` | array | `[]` = none | +| 13 | `language` | string | — | + +`format=minimize` **ignores** `detail` — the schema is always these 14 fields. +Use `format=medium` (optionally with `detail=verbose`) when you want a lean +projection or a fully self-describing object instead. + +### Example + +`codegraph_search_symbol {"query": "greet"}` (minimize, default): + +```json +{ + "results": [ + [100, "greet", "function", "global", 0, 0, null, "app.py", 1, 2, + "def greet(name: str) -> str:", null, [], "python"] + ], + "total": 1, + "limit": 20, + "offset": 0 +} +``` + +`codegraph_search_symbol {"query": "greet", "format": "medium"}`: + +```json +{ + "results": [ + { "id": 100, "name": "greet", "kind": "function", + "file": "app.py", "line": 1, "signature": "def greet(name: str) -> str:" } + ], + "total": 1, + "limit": 20, + "offset": 0 +} +``` + ## Symbols are numbers Symbols are identified by numeric `id` (global registry, ≥ 100). Call-chain @@ -126,18 +265,18 @@ Arguments: - `diff`: the unified diff text. Supports multi-file diffs, added/removed/ renamed files, and `\ No newline at end of file`. -Response shape: +Response shape (default-valued fields omitted per the omission rule): ```json { "draft": true, "summary": { "files_in_diff": 2, "files_matched": 2, "symbols_affected": 1, - "flows_affected": 1, "new_files": [], "unmatched_files": [] + "flows_affected": 1 }, "files": [{ "path": "src/foo.rs", "matched": true, "matched_path": "/abs/workspace/src/foo.rs", - "added_lines": 3, "removed_lines": 2, "deleted": false, + "added_lines": 3, "removed_lines": 2, "symbols": [{ "symbol": { "id": 141, "name": "foo", "file": "src/foo.rs", "line": 10, "end_line": 25 }, "impact": "modified" }], "flows": [{ "flow": { "id": 141, "name": "foo", "file": "src/foo.rs", "line": 10 }, @@ -160,7 +299,8 @@ Key points: span of the whole affected region. - A file that doesn't match anything in the index lands in `summary.unmatched_files` (never indexed) or `summary.new_files` (added file - with no removed lines). + with no removed lines). Both keys are **omitted when empty** (`[]`), like + `deleted: false` and any other default value. ## Diff simulation — `codegraph_diff_simulate` @@ -176,14 +316,14 @@ Arguments (besides `diff`): - `args`, `mocks`, `branch_policy`, `loop_cap`: same contract as `codegraph_sandbox`. -Response shape: +Response shape (default-valued fields omitted): ```json { "draft": true, "entry": "compute", "base_ref": "HEAD", "affected_functions": ["compute", "cap"], - "before": { "present": true, "return": 50, "sequence": ["if:1", "call:fetch"], "missing_mocks": [] }, - "after": { "present": true, "return": 6, "sequence": ["if:1", "call:fetch", "call:extra"], "missing_mocks": [] }, - "delta": { "sequence_added": ["call:extra"], "sequence_removed": [] } + "before": { "present": true, "return": 50, "sequence": ["if:1", "call:fetch"] }, + "after": { "present": true, "return": 6, "sequence": ["if:1", "call:fetch", "call:extra"] }, + "delta": { "sequence_added": ["call:extra"] } } ``` @@ -194,10 +334,12 @@ evaluated), loops run up to `loop_cap`, and **numeric arithmetic on values is not modeled**. So the reliable signal is `delta.sequence_added/removed` — e.g. an MR that adds/removes a call, a branch, or switches a callee shows up as a sequence delta; an MR that only changes an arithmetic expression does not. -A function that doesn't exist in `base_ref` (new in the MR) reports -`before.present: false`; a callee without a mock reports +A function that doesn't exist in `base_ref` (new in the MR) reports `before` +**without** a `present` field (absent = not present; only `reason` remains). A +callee without a mock reports `link_error: no mock configured for callee(s): …` (compile aborts before -running — supply it in `mocks` and retry). +running — supply it in `mocks` and retry). `missing_mocks` and empty +`sequence_removed` are omitted when empty. ## Origin/ref simulation — `codegraph_origin_simulate` @@ -215,13 +357,13 @@ Arguments: - `args`, `mocks`, `branch_policy`, `loop_cap`: same contract as `codegraph_sandbox`. -Response shape: +Response shape (default-valued fields omitted): ```json { "draft": true, "entry": "compute", "ref": "origin/main", - "origin": { "present": true, "return": 50, "sequence": ["if:1", "call:fetch"], "missing_mocks": [] }, - "working_tree": { "present": true, "return": 6, "sequence": ["if:1", "call:fetch", "call:extra"], "missing_mocks": [] }, - "delta": { "sequence_added": ["call:extra"], "sequence_removed": [] } + "origin": { "present": true, "return": 50, "sequence": ["if:1", "call:fetch"] }, + "working_tree": { "present": true, "return": 6, "sequence": ["if:1", "call:fetch", "call:extra"] }, + "delta": { "sequence_added": ["call:extra"] } } ``` diff --git a/crates/codegraph-mcp/src/session.rs b/crates/codegraph-mcp/src/session.rs index 61c7e574a..8ad681e51 100644 --- a/crates/codegraph-mcp/src/session.rs +++ b/crates/codegraph-mcp/src/session.rs @@ -23,6 +23,70 @@ use serde_json::{json, Value}; use std::sync::Arc; use tokio::sync::RwLock; +/// Mức chi tiết mặc định của Symbol trong response các list tool — set tại +/// `codegraph_init {"detail": ...}`, có thể ghi đè từng call bằng arg `detail`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum DetailLevel { + /// `{id, name, kind, file, line}` — tối ưu token cho reasoning. + Minimal, + /// Mặc định: thêm `signature` (dòng khai báo đầu tiên). + #[default] + Medium, + /// Full `Symbol` (doc, annotations, scope, type_ref, ...) — như cũ. + Verbose, +} + +impl DetailLevel { + /// Parse từ tên arg (`minimal`/`medium`/`verbose`) — `None` nếu lạ. + pub fn parse(s: &str) -> Option { + Some(match s { + "minimal" => Self::Minimal, + "medium" => Self::Medium, + "verbose" => Self::Verbose, + _ => return None, + }) + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Minimal => "minimal", + Self::Medium => "medium", + Self::Verbose => "verbose", + } + } +} + +/// Định dạng response kiểu Binance-style minimal — set tại +/// `codegraph_init {"format": ...}`, ghi đè từng call bằng arg `format`, và có +/// thể seed từ CLI lúc khởi động (`codegraph serve --mcp --format=...`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum OutputStyle { + /// Mặc định — nhỏ gọn nhất: symbol thành mảng vị trí cố định (chỉ value, + /// order được document; value thiếu = sentinel null/0/""/[]). + #[default] + Minimize, + /// Giữ key, lược bỏ field có value mặc định (None/0/""/[]/{}/false). + Medium, +} + +impl OutputStyle { + /// Parse từ tên arg (`minimize`/`medium`) — `None` nếu lạ. + pub fn parse(s: &str) -> Option { + Some(match s { + "minimize" => Self::Minimize, + "medium" => Self::Medium, + _ => return None, + }) + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Minimize => "minimize", + Self::Medium => "medium", + } + } +} + /// Trạng thái session. enum SessionState { /// Chưa có root nào được bind (hoặc đã `codegraph_deinit`). @@ -45,6 +109,8 @@ pub struct InitOutcome { pub struct Session { root: RwLock>, state: RwLock, + detail: RwLock, + format: RwLock, } impl Default for Session { @@ -56,15 +122,27 @@ impl Default for Session { impl Session { /// Session trống — chưa có root nào; `codegraph_init` sẽ bind. pub fn new() -> Self { + Self::new_with_format(OutputStyle::default()) + } + + /// `new()` nhưng seed sẵn output format từ CLI lúc khởi động. + pub fn new_with_format(format: OutputStyle) -> Self { Self { root: RwLock::new(None), state: RwLock::new(SessionState::Empty), + detail: RwLock::new(DetailLevel::default()), + format: RwLock::new(format), } } /// Pre-seed root lúc khởi động (`--path`). Có `.codegraph/` → load storage /// ngay (Ready); chưa init → Empty, chờ `codegraph_init` bind lại. pub async fn with_root(root: Utf8PathBuf) -> Result { + Self::with_root_and_format(root, OutputStyle::default()).await + } + + /// `with_root()` nhưng seed sẵn output format từ CLI lúc khởi động. + pub async fn with_root_and_format(root: Utf8PathBuf, format: OutputStyle) -> Result { let state = if project_dir(&root).exists() { let dsn = ExtractConfig::load(&root).storage_dsn(&root); let shared_index = Arc::new(SharedGraphIndex::open(dsn.clone()).await?); @@ -75,6 +153,8 @@ impl Session { Ok(Self { root: RwLock::new(Some(root)), state: RwLock::new(state), + detail: RwLock::new(DetailLevel::default()), + format: RwLock::new(format), }) } @@ -94,12 +174,20 @@ impl Session { .unwrap_or(false) } - /// `codegraph_init { path, index }`: normalize/validate path, bind root, - /// tạo `.codegraph/` + config, index CHỈ khi `do_index = true` (mặc định - /// không index — bind nhanh, không block user; agent chủ động gọi - /// `codegraph_index {}` khi cần data), rồi load storage theo config vừa - /// tạo → session chuyển sang `Ready`. - pub async fn init(&self, path: Utf8PathBuf, do_index: bool) -> Result { + /// `codegraph_init { path, index, detail, format }`: normalize/validate path, + /// bind root, tạo `.codegraph/` + config, index CHỈ khi `do_index = true` + /// (mặc định không index — bind nhanh, không block user; agent chủ động gọi + /// `codegraph_index {}` khi cần data), rồi load storage theo config vừa tạo + /// → session chuyển sang `Ready`. `detail` là mức chi tiết mặc định cho + /// symbol trong response các list tool (minimal/medium/verbose); `format` là + /// output style (minimize/medium) — `None` giữ nguyên giá trị seed từ CLI. + pub async fn init( + &self, + path: Utf8PathBuf, + do_index: bool, + detail: DetailLevel, + format: Option, + ) -> Result { let root = normalize_root(path)?; let dir = init_project(&root)?; let indexed = if do_index { @@ -115,11 +203,25 @@ impl Session { // Root set trước state — mọi `ensure_ready` đồng thời đọc root mới sẽ // tự swap state theo DSN mới (xem `ensure_ready`). *self.root.write().await = Some(root.clone()); + *self.detail.write().await = detail; + if let Some(f) = format { + *self.format.write().await = f; + } let mut st = self.state.write().await; *st = SessionState::Ready { dsn, shared_index }; Ok(InitOutcome { root, dir, indexed }) } + /// Detail level hiện tại (default mặc định cho symbol trong list tools). + pub async fn detail(&self) -> DetailLevel { + *self.detail.read().await + } + + /// Output format hiện tại (minimize/medium) cho mọi response. + pub async fn format(&self) -> OutputStyle { + *self.format.read().await + } + /// `codegraph_deinit`: nhả session — trả root cũ (nếu có). `.codegraph/` /// và index để nguyên trên đĩa; `codegraph_init` có thể bind lại sau đó. pub async fn deinit(&self) -> Result> { diff --git a/crates/codegraph-mcp/src/stdio.rs b/crates/codegraph-mcp/src/stdio.rs index 705e96775..1f89d9555 100644 --- a/crates/codegraph-mcp/src/stdio.rs +++ b/crates/codegraph-mcp/src/stdio.rs @@ -16,7 +16,8 @@ pub async fn serve_stdio(service: S) -> anyhow::Result<()> where S: rmcp::ServerHandler, { - service.serve(rmcp::transport::io::stdio()) + service + .serve(rmcp::transport::io::stdio()) .await? .waiting() .await?; diff --git a/crates/codegraph-mcp/src/tools.rs b/crates/codegraph-mcp/src/tools.rs index cfea47e31..d0430987d 100644 --- a/crates/codegraph-mcp/src/tools.rs +++ b/crates/codegraph-mcp/src/tools.rs @@ -1,11 +1,13 @@ +use crate::session::{DetailLevel, OutputStyle}; use camino::{Utf8Path, Utf8PathBuf}; -use codegraph_api::GraphApi; +use codegraph_api::{GraphApi, Pagination}; use codegraph_context::{ContextRequest, Format}; use codegraph_core::{is_marker, Error, Result, Symbol, SymbolKind, SymbolMatch}; use codegraph_extract::Orchestrator; use codegraph_graph::{GraphIndex, SharedGraphIndex}; use codegraph_sboxes::{compile_with_mocks, BranchPolicy, SboxConfig}; use rmcp::model::Tool; +use serde::Serialize; use serde_json::{json, Value}; use std::sync::Arc; @@ -38,10 +40,14 @@ fn tool_defs() -> Vec { vec![ tool( "codegraph_search", - "Search symbols by name (substring, case-insensitive).", + "Search symbols by name (substring, case-insensitive). On large indexes this can take a while — pass timeout_ms (default 2000) and, if the call returns a timeout error containing \"resume\": \"\", retry the SAME call with that resume id to continue the search from where it stopped.", json!({ "type": "object", "properties": { "query": { "type": "string" }, - "limit": { "type": "integer", "default": 20 } + "limit": { "type": "integer", "default": 10 }, + "resume": { "type": "string", "description": "Resume id from a previous timeout — retry the same call with this to continue where it stopped." }, + "timeout_ms": { "type": "integer", "default": 2000, "description": "Soft time budget in ms; 0 = no limit. On timeout the tool errors with a resume id." }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Symbol detail for this call (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } }, "required": ["query"] }), ), tool( @@ -49,7 +55,8 @@ fn tool_defs() -> Vec { "Look up a symbol by id or exact name. Duplicate names → ambiguous with the full match list; retry with symbol_id.", json!({ "type": "object", "properties": { "id": { "type": "integer" }, - "name": { "type": "string" } + "name": { "type": "string" }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol as a fixed-order positional array (default), medium = full object with default-valued fields omitted." } } }), ), tool( @@ -57,14 +64,18 @@ fn tool_defs() -> Vec { "Find functions that (transitively) call the given symbol.", json!({ "type": "object", "properties": { "node": { "type": "integer" }, - "depth": { "type": "integer", "default": 1 } + "depth": { "type": "integer", "default": 1 }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Symbol detail for this call (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } }, "required": ["node"] }), ), tool( "codegraph_callees", "Find functions called directly by the given symbol.", json!({ "type": "object", "properties": { - "node": { "type": "integer" } + "node": { "type": "integer" }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Symbol detail for this call (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } }, "required": ["node"] }), ), tool( @@ -72,14 +83,18 @@ fn tool_defs() -> Vec { "Impact radius: who transitively depends on this symbol.", json!({ "type": "object", "properties": { "node": { "type": "integer" }, - "max_depth": { "type": "integer", "default": 3 } + "max_depth": { "type": "integer", "default": 3 }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Symbol detail for this call (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } }, "required": ["node"] }), ), tool( "codegraph_flow", "Call chain of a symbol: markers (LOOP, IF_TRUE, …) + callee names + call sites with line/condition/effect.", json!({ "type": "object", "properties": { - "node": { "type": "integer" } + "node": { "type": "integer" }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Detail for the embedded symbol (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } }, "required": ["node"] }), ), tool( @@ -104,7 +119,7 @@ fn tool_defs() -> Vec { "Functions that call a library call whose name contains the query (includes unresolved external calls).", json!({ "type": "object", "properties": { "query": { "type": "string" }, - "limit": { "type": "integer", "default": 20 } + "limit": { "type": "integer", "default": 10 } }, "required": ["query"] }), ), tool( @@ -120,10 +135,12 @@ fn tool_defs() -> Vec { // ── Admin tools (init / deinit / index) — thao tác trên session slot ── tool( "codegraph_init", - "Bind this MCP session to a workspace root (idempotent): creates .codegraph/ with .gitignore, version, and config.toml. Pass path (absolute workspace root) to select the directory for this session. index defaults to false — binding is quick and non-blocking (it does NOT index); call codegraph_index {} afterwards (or pass index=true here) only when you need a fresh index to query. Re-running with a different path re-points the session.", + "Bind this MCP session to a workspace root (idempotent): creates .codegraph/ with .gitignore, version, and config.toml. Pass path (absolute workspace root) to select the directory for this session. index defaults to false — binding is quick and non-blocking (it does NOT index); call codegraph_index {} afterwards (or pass index=true here) only when you need a fresh index to query. detail sets the default symbol detail for list-tool responses (default medium). Re-running with a different path re-points the session.", json!({ "type": "object", "properties": { "path": { "type": "string", "description": "Absolute path of the workspace root to bind this session to." }, - "index": { "type": "boolean", "default": false } + "index": { "type": "boolean", "default": false }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "default": "medium", "description": "Default symbol detail for list-tool responses: minimal = id/name/kind/file/line (fewest tokens), medium = + signature, verbose = full Symbol (doc, annotations, ...). Per-call detail overrides this." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "default": "minimize", "description": "Output format for every response: minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted. Per-call format overrides this." } }, "required": ["path"] }), ), tool( @@ -139,13 +156,17 @@ fn tool_defs() -> Vec { // ── Enhanced symbol search (semgraph_search_symbol) ── tool( "codegraph_search_symbol", - "Search symbols by name with optional kind filter, match mode, and pagination. match: 'contains' (substring anywhere, default), 'prefix' (name starts with), 'suffix' (name ENDS with — e.g. query=\"Service\" finds every *Service class), 'exact' (exact name, case-insensitive). Use 'total' with 'offset' to fetch further pages until offset >= total.", + "Search symbols by name with optional kind filter, match mode, and pagination. match: 'contains' (substring anywhere, default), 'prefix' (name starts with), 'suffix' (name ENDS with — e.g. query=\"Service\" finds every *Service class), 'exact' (exact name, case-insensitive). Use 'total' with 'offset' to fetch further pages until offset >= total. On large indexes pass timeout_ms (default 2000); if the call returns a timeout error containing \"resume\": \"\", retry the SAME call with that resume id to continue. When more results remain, the response includes a resume id you can pass to page further without re-scanning.", json!({ "type": "object", "properties": { "query": { "type": "string" }, "kind": { "type": "string", "enum": ["function", "method", "class", "interface", "enum", "variable", "constant", "parameter", "field", "module", "file"] }, "match": { "type": "string", "enum": ["contains", "prefix", "suffix", "exact"], "default": "contains" }, "limit": { "type": "integer", "default": 20 }, - "offset": { "type": "integer", "default": 0 } + "offset": { "type": "integer", "default": 0 }, + "resume": { "type": "string", "description": "Resume id from a previous timeout (or from a previous response with more pages) — retry the same call with this to continue where it stopped." }, + "timeout_ms": { "type": "integer", "default": 2000, "description": "Soft time budget in ms; 0 = no limit. On timeout the tool errors with a resume id." }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Symbol detail for this call (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } }, "required": ["query"] }), ), // ── Class queries (semgraph_get_class_methods / get_class / list_classes / list_interfaces) ── @@ -155,7 +176,8 @@ fn tool_defs() -> Vec { json!({ "type": "object", "properties": { "class_name": { "type": "string" }, "id": { "type": "integer" }, - "compact": { "type": "boolean", "default": true } + "compact": { "type": "boolean", "default": true }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } } }), ), tool( @@ -163,7 +185,8 @@ fn tool_defs() -> Vec { "Get class/interface/enum details with fields and methods as separate lists.", json!({ "type": "object", "properties": { "class_name": { "type": "string" }, - "id": { "type": "integer" } + "id": { "type": "integer" }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = embedded class symbol as a fixed-order positional array (default), medium = objects with default-valued fields omitted." } } }), ), tool( @@ -171,7 +194,9 @@ fn tool_defs() -> Vec { "List all class symbols in the index (paginated).", json!({ "type": "object", "properties": { "limit": { "type": "integer", "default": 20 }, - "offset": { "type": "integer", "default": 0 } + "offset": { "type": "integer", "default": 0 }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Symbol detail for this call (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } } }), ), tool( @@ -179,7 +204,9 @@ fn tool_defs() -> Vec { "List all interface symbols in the index (paginated).", json!({ "type": "object", "properties": { "limit": { "type": "integer", "default": 20 }, - "offset": { "type": "integer", "default": 0 } + "offset": { "type": "integer", "default": 0 }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Symbol detail for this call (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } } }), ), tool( @@ -187,7 +214,8 @@ fn tool_defs() -> Vec { "Get a function's parameters and local variables. Disambiguate duplicate function names with 'id' from codegraph_search (pass 'id' alone).", json!({ "type": "object", "properties": { "func_name": { "type": "string" }, - "id": { "type": "integer" } + "id": { "type": "integer" }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = function/params/locals as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } } }), ), // ── Annotation / call / dependency queries ── @@ -197,8 +225,10 @@ fn tool_defs() -> Vec { json!({ "type": "object", "properties": { "annotation": { "type": "string" }, "kind": { "type": "string", "enum": ["function", "method", "class", "interface", "enum", "variable", "constant", "parameter", "field", "module", "file"] }, - "limit": { "type": "integer", "default": 50 }, - "offset": { "type": "integer", "default": 0 } + "limit": { "type": "integer", "default": 20 }, + "offset": { "type": "integer", "default": 0 }, + "detail": { "type": "string", "enum": ["minimal", "medium", "verbose"], "description": "Symbol detail for this call (overrides session default): minimal = id/name/kind/file/line, medium = + signature, verbose = full Symbol." }, + "format": { "type": "string", "enum": ["minimize", "medium"], "description": "Output format for this call (overrides session default): minimize = symbol items as fixed-order positional arrays (default), medium = objects with default-valued fields omitted." } }, "required": ["annotation"] }), ), tool( @@ -206,7 +236,7 @@ fn tool_defs() -> Vec { "Find functions that call a given class/method name inside their bodies (e.g. \"LogManager\" or \"LogManager.getLogger\"). Matches ALL call names captured by the parser — including external library calls that don't resolve to in-repo symbols. Each result includes per-call-site context: line, surrounding condition, whether inside a loop, and the call arguments.", json!({ "type": "object", "properties": { "call_name": { "type": "string" }, - "limit": { "type": "integer", "default": 50 } + "limit": { "type": "integer", "default": 20 } }, "required": ["call_name"] }), ), tool( @@ -271,32 +301,83 @@ fn tool_defs() -> Vec { ] } -pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Result { +pub async fn dispatch_with_api( + api: &GraphApi, + root: &Utf8Path, + session_detail: DetailLevel, + session_format: OutputStyle, + name: &str, + args: Value, +) -> Result { match name { "codegraph_search" => { let q = arg_str(&args, "query")?; - let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as u32; - let hits = api.search(q, limit).await?; - serde_json::to_string_pretty(&hits).map_err(|e| Error::Invalid(e.to_string())) + let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(10) as u32; + let resume = args + .get("resume") + .and_then(|v| v.as_str()) + .map(str::to_string); + let timeout_ms = args + .get("timeout_ms") + .and_then(|v| v.as_u64()) + .unwrap_or(2000); + let out = api.search_resumable(q, limit, resume, timeout_ms).await?; + if out.timed_out { + // Không trả kết quả nửa chừng — báo lỗi kèm resume id để LLM retry + // cùng args + resume → search tiếp tục đúng vị trí dừng. + return Err(Error::Other(format!( + "codegraph_search timed out after {}ms (collected {} symbols so far). \ + Retry the same call with the same arguments plus \"resume\": \"{}\" \ + to continue the search from where it stopped.", + timeout_ms, + out.progress, + out.resume.as_deref().unwrap_or("") + ))); + } + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let out: Vec = out + .page + .iter() + .map(|s| symbol_json(root.as_str(), s, detail, format)) + .collect(); + emit_value(root.as_str(), Value::Array(out)) } "codegraph_symbol" => { + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); if let Some(id) = args.get("id").and_then(|v| v.as_u64()) { let s = api.symbol_by_id(id).await; - return serde_json::to_string_pretty(&s).map_err(|e| Error::Invalid(e.to_string())); + return match s { + Some(s) => emit_value( + root.as_str(), + symbol_json(root.as_str(), &s, detail, format), + ), + None => emit_value(root.as_str(), Value::Null), + }; } if let Some(name) = args.get("name").and_then(|v| v.as_str()) { let r = api.resolve(name, 0).await?; if r.ambiguous { // Trùng tên — trả matches để LLM retry với symbol_id. + let matches: Vec = r + .matches + .iter() + .map(|s| symbol_json(root.as_str(), s, detail, format)) + .collect(); return Ok(format!( "ambiguous ({} matches):\n{}", - r.matches.len(), - serde_json::to_string_pretty(&r.matches) - .map_err(|e| Error::Invalid(e.to_string()))? + matches.len(), + emit_value(root.as_str(), Value::Array(matches))? )); } - return serde_json::to_string_pretty(&r.symbol) - .map_err(|e| Error::Invalid(e.to_string())); + return match r.symbol { + Some(s) => emit_value( + root.as_str(), + symbol_json(root.as_str(), &s, detail, format), + ), + None => emit_value(root.as_str(), Value::Null), + }; } Err(Error::Invalid("provide id or name".into())) } @@ -304,28 +385,56 @@ pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Resul let id = arg_u64(&args, "node")?; let depth = args.get("depth").and_then(|v| v.as_u64()).unwrap_or(1) as u32; let hits = api.callers(id, depth).await?; - serde_json::to_string_pretty(&hits).map_err(|e| Error::Invalid(e.to_string())) + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let out: Vec = hits + .iter() + .map(|s| symbol_json(root.as_str(), s, detail, format)) + .collect(); + emit_value(root.as_str(), Value::Array(out)) } "codegraph_callees" => { let id = arg_u64(&args, "node")?; let hits = api.callees(id).await?; - serde_json::to_string_pretty(&hits).map_err(|e| Error::Invalid(e.to_string())) + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let out: Vec = hits + .iter() + .map(|s| symbol_json(root.as_str(), s, detail, format)) + .collect(); + emit_value(root.as_str(), Value::Array(out)) } "codegraph_impact" => { let id = arg_u64(&args, "node")?; let depth = args.get("max_depth").and_then(|v| v.as_u64()).unwrap_or(3) as u32; - let report = api.impact(id, depth).await?; - serde_json::to_string_pretty(&report).map_err(|e| Error::Invalid(e.to_string())) + let hits = api.impact(id, depth).await?; + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let out: Vec = hits + .iter() + .map(|s| symbol_json(root.as_str(), s, detail, format)) + .collect(); + emit_value(root.as_str(), Value::Array(out)) } "codegraph_flow" => { let id = arg_u64(&args, "node")?; let flow = api.flow(id).await?; - serde_json::to_string_pretty(&flow).map_err(|e| Error::Invalid(e.to_string())) + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + emit_value( + root.as_str(), + json!({ + "symbol": symbol_json(root.as_str(), &flow.symbol, detail, format), + "chain": flow.chain, + "chain_desc": flow.chain_desc, + "calls": flow.calls, + }), + ) } "codegraph_search_flow" => { let pattern = arg_str(&args, "pattern")?; let hits = api.search_flow_pattern(pattern).await?; - serde_json::to_string_pretty(&hits).map_err(|e| Error::Invalid(e.to_string())) + emit(root.as_str(), &hits) } "codegraph_context" => { let req = ContextRequest { @@ -337,23 +446,37 @@ pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Resul .unwrap_or(false), limit: args.get("limit").and_then(|v| v.as_u64()).unwrap_or(5) as u32, format: Format::Markdown, + strip_prefix: Some(root.as_str().to_string()), }; Ok(api.context_markdown(&req).await?) } "codegraph_references" => { let q = arg_str(&args, "query")?; - let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as u32; + let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(10) as u32; let report = api.references(q, limit).await?; - serde_json::to_string_pretty(&report).map_err(|e| Error::Invalid(e.to_string())) + emit(root.as_str(), &report) } "codegraph_files" => { let prefix = args.get("path").and_then(|v| v.as_str()).unwrap_or(""); - let files = api.files(prefix).await; - serde_json::to_string_pretty(&files).map_err(|e| Error::Invalid(e.to_string())) + // Index lưu path absolute; output relativize theo root. Filter khớp + // CẢ prefix absolute (path gốc) lẫn prefix tương đối (path hiển thị). + let files = api.files("").await; + let files: Vec<_> = if prefix.is_empty() { + files + } else { + files + .into_iter() + .filter(|f| { + f.path.starts_with(prefix) + || strip_root_prefix(&f.path, root.as_str()).starts_with(prefix) + }) + .collect() + }; + emit(root.as_str(), &files) } "codegraph_status" => { let stats = api.stats().await; - serde_json::to_string_pretty(&stats).map_err(|e| Error::Invalid(e.to_string())) + emit(root.as_str(), &stats) } "codegraph_search_symbol" => { let q = arg_str(&args, "query")?; @@ -368,17 +491,52 @@ pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Resul .unwrap_or(SymbolMatch::Contains); let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as u32; let offset = args.get("offset").and_then(|v| v.as_u64()).unwrap_or(0) as u32; - let (results, total) = api - .search_symbol_paged(q, kind, mode, limit, offset) + let resume = args + .get("resume") + .and_then(|v| v.as_str()) + .map(str::to_string); + let timeout_ms = args + .get("timeout_ms") + .and_then(|v| v.as_u64()) + .unwrap_or(2000); + let out = api + .search_symbol_paged_resumable( + q, + kind, + mode, + Pagination { limit, offset }, + resume, + timeout_ms, + ) .await?; - serde_json::to_string_pretty(&json!({ - "results": results, - "total": total, - "limit": limit, - "offset": offset, - "has_more": offset as usize + results.len() < total, - })) - .map_err(|e| Error::Invalid(e.to_string())) + if out.timed_out { + return Err(Error::Other(format!( + "codegraph_search_symbol timed out after {}ms (collected {} symbols so far). \ + Retry the same call with the same arguments plus \"resume\": \"{}\" \ + to continue the search from where it stopped.", + timeout_ms, + out.progress, + out.resume.as_deref().unwrap_or("") + ))); + } + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let results: Vec = out + .page + .into_iter() + .map(|s| symbol_json(root.as_str(), &s, detail, format)) + .collect(); + emit_value( + root.as_str(), + json!({ + "results": results, + "total": out.total, + "limit": limit, + "offset": offset, + "has_more": offset as usize + results.len() < out.total, + "resume": out.resume, + }), + ) } "codegraph_class_methods" => { let target = resolve_target( @@ -390,7 +548,7 @@ pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Resul ) .await?; match target { - Target::Ambiguous(v) => Ok(json_str(v)), + Target::Ambiguous(v) => emit_value(root.as_str(), v), Target::Symbol(sym) => { if !matches!( sym.kind, @@ -419,13 +577,15 @@ pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Resul .map(|m| serde_json::to_value(&m).unwrap_or(Value::Null)) .collect() }; - serde_json::to_string_pretty(&json!({ - "class_name": sym.name, - "methods": methods, - "compact": compact, - "total": methods.len(), - })) - .map_err(|e| Error::Invalid(e.to_string())) + emit_value( + root.as_str(), + json!({ + "class_name": sym.name, + "methods": methods, + "compact": compact, + "total": methods.len(), + }), + ) } } } @@ -439,10 +599,20 @@ pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Resul ) .await?; match target { - Target::Ambiguous(v) => Ok(json_str(v)), + Target::Ambiguous(v) => emit_value(root.as_str(), v), Target::Symbol(sym) => match api.class_info(sym.id).await { - Some(info) => serde_json::to_string_pretty(&info) - .map_err(|e| Error::Invalid(e.to_string())), + Some(info) => { + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + emit_value( + root.as_str(), + json!({ + "class": symbol_json(root.as_str(), &info.class, detail, format), + "fields": info.fields, + "methods": info.methods, + }), + ) + } None => Err(Error::Invalid(format!( "symbol {:?} (id {}) is not a class/interface/enum", sym.name, sym.id @@ -454,42 +624,80 @@ pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Resul let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as u32; let offset = args.get("offset").and_then(|v| v.as_u64()).unwrap_or(0) as u32; let (results, total) = api.list_by_kind(SymbolKind::Class, limit, offset).await; - serde_json::to_string_pretty(&json!({ - "kind": "class", - "results": results, - "total": total, - "limit": limit, - "offset": offset, - })) - .map_err(|e| Error::Invalid(e.to_string())) + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let results: Vec = results + .into_iter() + .map(|s| symbol_json(root.as_str(), &s, detail, format)) + .collect(); + emit_value( + root.as_str(), + json!({ + "kind": "class", + "results": results, + "total": total, + "limit": limit, + "offset": offset, + }), + ) } "codegraph_list_interfaces" => { let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as u32; let offset = args.get("offset").and_then(|v| v.as_u64()).unwrap_or(0) as u32; let (results, total) = api.list_by_kind(SymbolKind::Interface, limit, offset).await; - serde_json::to_string_pretty(&json!({ - "kind": "interface", - "results": results, - "total": total, - "limit": limit, - "offset": offset, - })) - .map_err(|e| Error::Invalid(e.to_string())) + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let results: Vec = results + .into_iter() + .map(|s| symbol_json(root.as_str(), &s, detail, format)) + .collect(); + emit_value( + root.as_str(), + json!({ + "kind": "interface", + "results": results, + "total": total, + "limit": limit, + "offset": offset, + }), + ) } "codegraph_function_scope" => { let target = resolve_target(api, &args, "id", "func_name", &[]).await?; match target { - Target::Ambiguous(v) => Ok(json_str(v)), + Target::Ambiguous(v) => emit_value(root.as_str(), v), Target::Symbol(sym) => match api.function_scope(sym.id).await { - Some(scope) => serde_json::to_string_pretty(&scope) - .map_err(|e| Error::Invalid(e.to_string())), - None => Ok(serde_json::to_string_pretty(&json!({ - "function": sym.name, - "parameters": [], - "locals": [], - "total": 0, - })) - .map_err(|e| Error::Invalid(e.to_string()))?), + Some(scope) => { + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let parameters: Vec = scope + .parameters + .iter() + .map(|s| symbol_json(root.as_str(), s, detail, format)) + .collect(); + let locals: Vec = scope + .locals + .iter() + .map(|s| symbol_json(root.as_str(), s, detail, format)) + .collect(); + emit_value( + root.as_str(), + json!({ + "function": symbol_json(root.as_str(), &scope.function, detail, format), + "parameters": parameters, + "locals": locals, + }), + ) + } + None => emit_value( + root.as_str(), + json!({ + "function": sym.name, + "parameters": [], + "locals": [], + "total": 0, + }), + ), }, } } @@ -499,35 +707,45 @@ pub async fn dispatch_with_api(api: &GraphApi, name: &str, args: Value) -> Resul .get("kind") .and_then(|v| v.as_str()) .and_then(SymbolKind::parse); - let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(50) as u32; + let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as u32; let offset = args.get("offset").and_then(|v| v.as_u64()).unwrap_or(0) as u32; let (results, total, truncated) = api .search_by_annotation(annotation, kind, offset, limit) .await; - serde_json::to_string_pretty(&json!({ - "annotation": annotation, - "kind": kind.map(|k| k.as_str()), - "results": results, - "total": total, - "offset": offset, - "truncated": truncated, - })) - .map_err(|e| Error::Invalid(e.to_string())) + let detail = detail_from_args(&args, session_detail); + let format = format_from_args(&args, session_format); + let results: Vec = results + .into_iter() + .map(|s| symbol_json(root.as_str(), &s, detail, format)) + .collect(); + emit_value( + root.as_str(), + json!({ + "annotation": annotation, + "kind": kind.map(|k| k.as_str()), + "results": results, + "total": total, + "offset": offset, + "truncated": truncated, + }), + ) } "codegraph_search_by_call" => { let call_name = arg_str(&args, "call_name")?; - let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(50) as u32; + let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as u32; let hits = api.references(call_name, limit).await?; - serde_json::to_string_pretty(&json!({ - "call_name": call_name, - "results": hits, - "total": hits.len(), - })) - .map_err(|e| Error::Invalid(e.to_string())) + emit_value( + root.as_str(), + json!({ + "call_name": call_name, + "results": hits, + "total": hits.len(), + }), + ) } "codegraph_dependencies" => { let report = api.dependencies().await; - serde_json::to_string_pretty(&report).map_err(|e| Error::Invalid(e.to_string())) + emit(root.as_str(), &report) } _ => Err(Error::Invalid(format!("unknown tool: {name}"))), } @@ -609,10 +827,6 @@ async fn resolve_target( } } -fn json_str(v: Value) -> String { - serde_json::to_string_pretty(&v).unwrap_or_else(|_| v.to_string()) -} - fn arg_str<'a>(v: &'a Value, k: &str) -> Result<&'a str> { v.get(k) .and_then(|x| x.as_str()) @@ -624,6 +838,165 @@ fn arg_u64(v: &Value, k: &str) -> Result { .ok_or_else(|| Error::Invalid(format!("missing int arg: {k}"))) } +// ── Symbol detail + path relativization ── +// List tools trả symbol theo `DetailLevel` của session (`codegraph_init +// {"detail": ...}`), ghi đè từng call bằng arg `detail`. Mọi response đi qua +// `emit_value`/`emit` để `file`/`path` relativize theo workspace root — LLM +// không cần thấy tiền tố absolute lặp lại trên từng dòng. + +/// Detail level cho một tool: arg `detail` ghi đè session default. +fn detail_from_args(args: &Value, session: DetailLevel) -> DetailLevel { + args.get("detail") + .and_then(|v| v.as_str()) + .and_then(DetailLevel::parse) + .unwrap_or(session) +} + +/// Output style cho một tool: arg `format` ghi đè session default. +fn format_from_args(args: &Value, session: OutputStyle) -> OutputStyle { + args.get("format") + .and_then(|v| v.as_str()) + .and_then(OutputStyle::parse) + .unwrap_or(session) +} + +/// Symbol JSON theo `detail` + `style`. `Minimize` (mặc định) → mảng vị trí cố +/// định (order được document trong server-instructions.md; file đã relativize +/// theo root — relativize_paths chỉ chạm object key, không chạm phần tử mảng); +/// `Medium` → object giữ key (field default bị lược sau trong `omit_defaults`). +fn symbol_json(root: &str, s: &Symbol, detail: DetailLevel, style: OutputStyle) -> Value { + match style { + OutputStyle::Minimize => json!([ + s.id, + s.name, + s.kind.as_str(), + s.scope.as_str(), + s.scope_id, + s.type_ref, + s.type_name, + strip_root_prefix(&s.file, root), + s.line, + s.end_line, + s.signature, + s.doc, + s.annotations, + s.language, + ]), + OutputStyle::Medium => match detail { + DetailLevel::Minimal => json!({ + "id": s.id, + "name": s.name, + "kind": s.kind.as_str(), + "file": s.file, + "line": s.line, + }), + DetailLevel::Medium => json!({ + "id": s.id, + "name": s.name, + "kind": s.kind.as_str(), + "file": s.file, + "line": s.line, + "signature": s.signature, + }), + DetailLevel::Verbose => serde_json::to_value(s).unwrap_or(Value::Null), + }, + } +} + +/// Strip `root/` prefix khỏi một path — chỉ khi root là tiền tố theo boundary +/// (`root` + `/`), tránh cắt nhầm `/root2/...`. Giữ nguyên nếu không khớp. +pub(crate) fn strip_root_prefix<'a>(path: &'a str, root: &str) -> &'a str { + if let Some(rest) = path.strip_prefix(root) { + if let Some(rest) = rest.strip_prefix('/') { + return rest; + } + } + path +} + +/// Keys mang đường dẫn file trong response — relativize theo workspace root. +const PATH_KEYS: [&str; 3] = ["file", "path", "matched_path"]; + +/// Strip `root/` prefix khỏi mọi đường dẫn file trong cây JSON (in-place). +fn relativize_paths(v: &mut Value, root: &str) { + match v { + Value::Object(map) => { + for (k, val) in map.iter_mut() { + if PATH_KEYS.contains(&k.as_str()) { + if let Some(s) = val.as_str() { + *val = Value::String(strip_root_prefix(s, root).to_string()); + } + } + relativize_paths(val, root); + } + } + Value::Array(arr) => { + for item in arr.iter_mut() { + relativize_paths(item, root); + } + } + _ => {} + } +} + +/// Serialize payload JSON kèm relativize path theo root — mọi response tool +/// đi qua đây để `file`/`path` trả về tương đối so với workspace root. +fn emit_value(root: &str, v: Value) -> Result { + let mut v = v; + relativize_paths(&mut v, root); + omit_defaults(&mut v); + serde_json::to_string_pretty(&v).map_err(|e| Error::Invalid(e.to_string())) +} + +/// `emit_value` cho bất kỳ type serializable nào (chuyển qua `to_value`). +fn emit(root: &str, v: &T) -> Result { + let value = serde_json::to_value(v).map_err(|e| Error::Invalid(e.to_string()))?; + emit_value(root, value) +} + +/// Keys có `0` = "absent" (sentinel) — value 0 bị lược như default. Các số khác +/// (counts/totals như `total`, `symbols`, `lines`, ...) giữ nguyên 0 vì ý nghĩa. +const ZERO_SENTINEL_KEYS: [&str; 3] = ["scope_id", "type_ref", "end_line"]; + +/// Value có phải "default" cần lược không (Binance-style minimal): +/// null / false / "" / [] / {} — và số 0 cho sentinel keys. +fn is_default_value(key: &str, v: &Value) -> bool { + match v { + Value::Null => true, + Value::Bool(b) => !*b, + Value::String(s) => s.is_empty(), + Value::Array(a) => a.is_empty(), + Value::Object(m) => m.is_empty(), + Value::Number(n) => ZERO_SENTINEL_KEYS.contains(&key) && n.as_f64() == Some(0.0), + } +} + +/// Lược bỏ key có value mặc định trong mọi OBJECT (in-place). ARRAY không bao +/// giờ bị xóa phần tử — schema mảng vị trí cố định (style `minimize`) phải giữ +/// nguyên độ dài; chỉ object con bên trong được xử lý tiếp. +/// +/// Giữ thứ tự key (preserve_order): `mem::take` + rebuild — `Map::remove` là +/// swap-remove (đảo thứ tự), `shift_remove` không có sẵn trên mọi bản serde_json. +pub(crate) fn omit_defaults(v: &mut Value) { + match v { + Value::Object(map) => { + let old = std::mem::take(map); + for (k, mut child) in old { + omit_defaults(&mut child); + if !is_default_value(&k, &child) { + map.insert(k, child); + } + } + } + Value::Array(arr) => { + for item in arr.iter_mut() { + omit_defaults(item); + } + } + _ => {} + } +} + // ── Sandbox tool (codegraph_sandbox) ── // Cần workspace root (config.toml `[sandbox]` + mock dirs) và snapshot index, // nên dispatch riêng qua `SharedGraphIndex` — không qua `GraphApi`. @@ -741,18 +1114,20 @@ pub async fn dispatch_sandbox( .iter() .filter_map(|id| idx.symbol_by_id(*id).map(|s| s.name)) .collect(); - serde_json::to_string_pretty(&json!({ - "entry": flow.symbol.name, - "entry_id": entry_id, - "group": group_names, - "args": call_args, - "return": ret, - "mocks": trace.mocks, - "conds": trace.conds, - "missing_mocks": trace.missing, - "sequence": trace.sequence(), - })) - .map_err(|e| Error::Invalid(e.to_string())) + emit_value( + root.as_str(), + json!({ + "entry": flow.symbol.name, + "entry_id": entry_id, + "group": group_names, + "args": call_args, + "return": ret, + "mocks": trace.mocks, + "conds": trace.conds, + "missing_mocks": trace.missing, + "sequence": trace.sequence(), + }), + ) } /// Phân tích unified diff (MR / patch / `git diff`) thành bản DRAFT tác động @@ -769,7 +1144,7 @@ pub async fn dispatch_diff( let idx = shared.ensure_fresh().await; let report = idx.diff_assess(&parsed, Some(root.as_std_path())).await; - serde_json::to_string_pretty(&report).map_err(|e| Error::Invalid(e.to_string())) + emit(root.as_str(), &report) } /// Chạy sandbox trên flow của `entry_name` trong một index cụ thể. Trả JSON @@ -956,7 +1331,7 @@ pub async fn dispatch_diff_simulate( let _ = std::fs::remove_dir_all(&tmp); let payload = result?; - serde_json::to_string_pretty(&payload).map_err(|e| Error::Invalid(e.to_string())) + emit_value(root.as_str(), payload) } /// Ref → simulate: chạy sandbox trên flow entry trên cây git tại `ref` (index @@ -1002,5 +1377,250 @@ pub async fn dispatch_origin_simulate( let _ = std::fs::remove_dir_all(&tmp); let payload = result?; - serde_json::to_string_pretty(&payload).map_err(|e| Error::Invalid(e.to_string())) + emit_value(root.as_str(), payload) +} + +#[cfg(test)] +mod tests { + use super::*; + use codegraph_core::{ScopeLevel, Symbol}; + + fn sample_symbol() -> Symbol { + Symbol { + id: 123, + name: "fetch_user".into(), + kind: SymbolKind::Function, + scope: ScopeLevel::Global, + scope_id: 0, + type_ref: 0, + type_name: None, + file: "/workspace/src/user.rs".into(), + line: 10, + end_line: 22, + signature: Some("fn fetch_user(id: u64) -> User".into()), + doc: Some("/// Lấy user theo id.".into()), + annotations: vec![], + language: "rust".into(), + } + } + + #[test] + fn detail_level_parse_roundtrip() { + assert_eq!(DetailLevel::parse("minimal"), Some(DetailLevel::Minimal)); + assert_eq!(DetailLevel::parse("medium"), Some(DetailLevel::Medium)); + assert_eq!(DetailLevel::parse("verbose"), Some(DetailLevel::Verbose)); + assert_eq!(DetailLevel::parse("bogus"), None); + assert_eq!(DetailLevel::default(), DetailLevel::Medium); + } + + #[test] + fn detail_from_args_overrides_session() { + let args = json!({ "detail": "verbose" }); + assert_eq!( + detail_from_args(&args, DetailLevel::Minimal), + DetailLevel::Verbose + ); + let no_arg = json!({ "query": "x" }); + assert_eq!( + detail_from_args(&no_arg, DetailLevel::Minimal), + DetailLevel::Minimal + ); + } + + #[test] + fn output_style_parse_roundtrip() { + assert_eq!(OutputStyle::parse("minimize"), Some(OutputStyle::Minimize)); + assert_eq!(OutputStyle::parse("medium"), Some(OutputStyle::Medium)); + assert_eq!(OutputStyle::parse("bogus"), None); + assert_eq!(OutputStyle::default(), OutputStyle::Minimize); + assert_eq!(OutputStyle::Minimize.as_str(), "minimize"); + assert_eq!(OutputStyle::Medium.as_str(), "medium"); + } + + #[test] + fn format_from_args_overrides_session() { + let args = json!({ "format": "medium" }); + assert_eq!( + format_from_args(&args, OutputStyle::Minimize), + OutputStyle::Medium + ); + let no_arg = json!({ "query": "x" }); + assert_eq!( + format_from_args(&no_arg, OutputStyle::Medium), + OutputStyle::Medium + ); + } + + #[test] + fn symbol_json_shapes_medium() { + let s = sample_symbol(); + // Style Medium giữ key; lược field default diễn ra sau ở emit_value/omit_defaults. + let minimal = symbol_json("/workspace", &s, DetailLevel::Minimal, OutputStyle::Medium); + assert_eq!(minimal["id"], 123); + assert_eq!(minimal["name"], "fetch_user"); + assert_eq!(minimal["kind"], "function"); + assert_eq!(minimal["file"], "/workspace/src/user.rs"); + assert_eq!(minimal["line"], 10); + assert!(minimal.get("signature").is_none()); + assert!(minimal.get("doc").is_none()); + + let medium = symbol_json("/workspace", &s, DetailLevel::Medium, OutputStyle::Medium); + assert_eq!(medium["signature"], "fn fetch_user(id: u64) -> User"); + assert!(medium.get("doc").is_none()); + + let verbose = symbol_json("/workspace", &s, DetailLevel::Verbose, OutputStyle::Medium); + assert_eq!(verbose["doc"], "/// Lấy user theo id."); + assert_eq!(verbose["end_line"], 22); + assert_eq!(verbose["language"], "rust"); + assert_eq!(verbose["type_name"], Value::Null); + } + + #[test] + fn symbol_json_minimize_array() { + let s = sample_symbol(); + // Mảng vị trí cố định: [id, name, kind, scope, scope_id, type_ref, + // type_name, file, line, end_line, signature, doc, annotations, language]. + let arr = symbol_json( + "/workspace", + &s, + DetailLevel::Verbose, + OutputStyle::Minimize, + ); + let a = arr.as_array().expect("minimize → array"); + assert_eq!(a.len(), 14); + assert_eq!(a[0], json!(123)); + assert_eq!(a[1], json!("fetch_user")); + assert_eq!(a[2], json!("function")); + assert_eq!(a[3], json!("global")); + assert_eq!(a[4], json!(0), "scope_id sentinel — vị trí giữ nguyên"); + assert_eq!(a[5], json!(0), "type_ref sentinel"); + assert_eq!(a[6], Value::Null, "type_name None"); + assert_eq!(a[7], json!("src/user.rs"), "file relativize theo root"); + assert_eq!(a[8], json!(10)); + assert_eq!(a[9], json!(22)); + assert_eq!(a[10], json!("fn fetch_user(id: u64) -> User")); + assert_eq!(a[11], json!("/// Lấy user theo id.")); + assert_eq!(a[12], json!([]), "annotations rỗng — phần tử giữ nguyên"); + assert_eq!(a[13], json!("rust")); + // detail bị bỏ qua ở minimize — mọi level ra cùng schema 14 vị trí. + let lean = symbol_json( + "/workspace", + &s, + DetailLevel::Minimal, + OutputStyle::Minimize, + ); + assert_eq!(lean.as_array().map(Vec::len), Some(14)); + } + + #[test] + fn omit_defaults_strips_defaults_keeps_counts() { + let mut v = json!({ + "results": [{ + "id": 1, "name": "a", "kind": "function", "scope": "global", + "scope_id": 0, "type_ref": 0, "type_name": null, "file": "a.rs", + "line": 3, "end_line": 0, "signature": null, "doc": "", + "annotations": [], "language": "" + }], + "total": 0, + "limit": 20, + "offset": 0, + "has_more": false, + "resume": null, + "kind": null, + "nested": { "a": [], "b": 0, "c": "" } + }); + omit_defaults(&mut v); + let r = &v["results"][0]; + assert_eq!(r.get("scope_id"), None, "0 sentinel lược"); + assert_eq!(r.get("type_ref"), None, "0 sentinel lược"); + assert_eq!(r.get("end_line"), None, "0 sentinel lược"); + assert_eq!(r.get("type_name"), None, "null lược"); + assert_eq!(r.get("signature"), None, "null lược"); + assert_eq!(r.get("doc"), None, "'' lược"); + assert_eq!(r.get("annotations"), None, "[] lược"); + assert_eq!(r.get("language"), None, "'' lược"); + assert_eq!(r["line"], 3, "line không phải sentinel — giữ"); + assert_eq!(r["name"], "a", "name giữ"); + assert_eq!(v.get("has_more"), None, "false lược"); + assert_eq!(v.get("resume"), None, "null lược"); + assert_eq!(v.get("kind"), None, "null lược"); + assert_eq!(v["total"], 0, "count giữ 0"); + assert_eq!(v["offset"], 0, "count giữ 0"); + assert_eq!(v["nested"]["b"], 0, "số không-sentinel giữ"); + assert_eq!(v["nested"].get("a"), None); + assert_eq!(v["nested"].get("c"), None); + } + + #[test] + fn omit_defaults_keeps_array_positions() { + // Schema mảng vị trí cố định — phần tử []/null/0 KHÔNG bị xóa khỏi mảng. + let mut v = json!({ + "results": [[123, "a", "function", "global", 0, 0, null, "a.rs", 1, 0, null, null, [], "rust"]] + }); + omit_defaults(&mut v); + let arr = v["results"][0].as_array().expect("mảng giữ nguyên"); + assert_eq!(arr.len(), 14); + assert_eq!(arr[4], json!(0)); + assert_eq!(arr[12], json!([])); + } + + #[test] + fn strip_root_prefix_is_boundary_aware() { + assert_eq!(strip_root_prefix("/workspace/a.rs", "/workspace"), "a.rs"); + assert_eq!(strip_root_prefix("/workspace/", "/workspace"), ""); + assert_eq!(strip_root_prefix("/workspace", "/workspace"), "/workspace"); + assert_eq!( + strip_root_prefix("/workspace2/a.rs", "/workspace"), + "/workspace2/a.rs" + ); + assert_eq!(strip_root_prefix("a.rs", "/workspace"), "a.rs"); + } + + #[test] + fn relativize_paths_rewrites_path_keys() { + let mut v = json!({ + "file": "/workspace/a.rs", + "path": "/workspace/c/d.rs", + "matched_path": "/workspace/e.rs", + "root": "/workspace", + "name": "/workspace/not-a-path-key", + "nested": [ { "file": "/workspace/x.rs", "label": "/workspace/y.rs" } ], + }); + relativize_paths(&mut v, "/workspace"); + assert_eq!(v["file"], "a.rs"); + assert_eq!(v["path"], "c/d.rs"); + assert_eq!(v["matched_path"], "e.rs"); + assert_eq!(v["root"], "/workspace", "key 'root' không relativize"); + assert_eq!( + v["name"], "/workspace/not-a-path-key", + "key khác không phải path" + ); + assert_eq!(v["nested"][0]["file"], "x.rs"); + assert_eq!(v["nested"][0]["label"], "/workspace/y.rs"); + } + + #[test] + fn emit_value_relativizes_and_roundtrips() { + let payload = json!({ + "hits": [ { "file": "/workspace/src/a.rs", "line": 1, "note": null, "skip": false } ] + }); + let text = emit_value("/workspace", payload).unwrap(); + let parsed: Value = serde_json::from_str(&text).unwrap(); + assert_eq!(parsed["hits"][0]["file"], "src/a.rs"); + assert!(parsed["hits"][0].get("note").is_none(), "null bị lược"); + assert!(parsed["hits"][0].get("skip").is_none(), "false bị lược"); + } + + #[test] + fn list_tools_result_serializes_cache_fields() { + // Protocol 2026-07-28 (SEP-2549) yêu cầu ttlMs/cacheScope trên tools/list; + // thiếu field → client strict (vd ZCode) reject toàn bộ response. + let result = rmcp::model::ListToolsResult::with_all_items(rmcp_tools()) + .with_ttl_ms(0) + .with_cache_scope(rmcp::model::CacheScope::Public); + let v = serde_json::to_value(&result).unwrap(); + assert_eq!(v["ttlMs"], 0); + assert_eq!(v["cacheScope"], "public"); + assert_eq!(v["tools"].as_array().map(Vec::len), Some(tool_defs().len())); + } } diff --git a/crates/codegraph-mcp/src/usage.rs b/crates/codegraph-mcp/src/usage.rs index 065128edf..5acca5e5c 100644 --- a/crates/codegraph-mcp/src/usage.rs +++ b/crates/codegraph-mcp/src/usage.rs @@ -92,21 +92,35 @@ impl UsageStats { /// Ước lượng source bytes mà một answer JSON "thay thế": gom mọi giá trị của /// key `file` (path của symbol trả về), map sang `FileInfo.bytes` trong index. -/// Duyệt toàn bộ cây JSON — an toàn với mọi shape của answer. -pub async fn estimate_source_bytes(api: &codegraph_api::GraphApi, answer_json: &Value) -> u64 { +/// Duyệt toàn bộ cây JSON — an toàn với mọi shape của answer. `root` là +/// workspace root: answer relativize `file` theo root (xem `tools::emit_value`), +/// nên lookup key cũng strip root để khớp. +pub async fn estimate_source_bytes( + api: &codegraph_api::GraphApi, + answer_json: &Value, + root: &str, +) -> u64 { let mut paths = Vec::new(); collect_file_paths(answer_json, &mut paths); if paths.is_empty() { return 0; } - // FileInfo.bytes của từng file (lazy — chỉ build khi cần). + // FileInfo.bytes của từng file (lazy — chỉ build khi cần). Key theo path + // tương đối với root — cùng dạng với `file` trong answer đã relativize. let files = api.files("").await; - let bytes_by_path: std::collections::HashMap<&str, u64> = - files.iter().map(|f| (f.path.as_str(), f.bytes)).collect(); + let bytes_by_path: std::collections::HashMap<&str, u64> = files + .iter() + .map(|f| { + ( + crate::tools::strip_root_prefix(f.path.as_str(), root), + f.bytes, + ) + }) + .collect(); let mut seen = std::collections::HashSet::new(); let mut total = 0u64; for p in paths { - if let Some(b) = bytes_by_path.get(p.as_str()) { + if let Some(b) = bytes_by_path.get(crate::tools::strip_root_prefix(&p, root)) { if seen.insert(p) { total += b; } diff --git a/crates/codegraph/src/main.rs b/crates/codegraph/src/main.rs index aa8a77071..76178f285 100644 --- a/crates/codegraph/src/main.rs +++ b/crates/codegraph/src/main.rs @@ -50,9 +50,32 @@ enum Cmd { Serve { #[arg(long)] mcp: bool, + /// Output format cho mọi response (Binance-style minimal): + /// minimize (mặc định) = symbol thành mảng vị trí cố định; medium = giữ + /// key, lược field có value mặc định. Ghi đè được theo session + /// (codegraph_init {"format": ...}) và từng call (arg "format"). + #[arg(long, value_enum, default_value_t = OutputFormat::Minimize)] + format: OutputFormat, }, } +/// Giá trị `--format` của CLI — map sang `codegraph_mcp::OutputStyle`. +#[derive(Clone, Copy, Debug, Default, clap::ValueEnum)] +enum OutputFormat { + #[default] + Minimize, + Medium, +} + +impl OutputFormat { + fn style(self) -> codegraph_mcp::OutputStyle { + match self { + Self::Minimize => codegraph_mcp::OutputStyle::Minimize, + Self::Medium => codegraph_mcp::OutputStyle::Medium, + } + } +} + #[tokio::main] async fn main() -> Result<()> { tracing_subscriber::fmt() @@ -80,10 +103,17 @@ async fn main() -> Result<()> { match cmd { Cmd::Init { no_index, progress } => cmd_init(&root, !no_index, progress).await, Cmd::Deinit => cmd_deinit(&root), - Cmd::Serve { mcp } => cmd_serve(&root, mcp).await, + Cmd::Serve { mcp, format } => cmd_serve(&root, mcp, format.style()).await, } } +/// Workspace đã init chưa — dấu hiệu là thư mục `.codegraph/` tồn tại (do +/// `codegraph init` tạo). Backend-agnostic: không phụ thuộc db file tồn tại +/// (lmdb dùng thư mục, redis không có file địa phương). +fn is_initialized(root: &Utf8Path) -> bool { + codegraph_extract::project_dir(root).exists() +} + /// Không có subcommand → in help. Banner console cũ bị bỏ: giao diện chính giờ /// là MCP (agent dùng `codegraph_init`/`codegraph_status` qua tools). async fn cmd_default(_root: &Utf8Path) -> Result<()> { @@ -107,13 +137,6 @@ async fn open_index(root: &Utf8Path) -> Result { } } -/// Workspace đã init chưa — dấu hiệu là thư mục `.codegraph/` tồn tại (do -/// `codegraph init` tạo). Backend-agnostic: không phụ thuộc db file tồn tại -/// (lmdb dùng thư mục, redis không có file địa phương). -fn is_initialized(root: &Utf8Path) -> bool { - codegraph_extract::project_dir(root).exists() -} - /// `codegraph init`: tạo `.codegraph/` + config, index ngay nếu `do_index` /// (progress bar khi `show_progress`). không gọi installer nữa. async fn cmd_init(root: &Utf8Path, do_index: bool, show_progress: bool) -> Result<()> { @@ -164,7 +187,7 @@ fn cmd_deinit(root: &Utf8Path) -> Result<()> { } /// `codegraph serve --mcp`: chạy MCP server trên stdio. -async fn cmd_serve(root: &Utf8Path, mcp: bool) -> Result<()> { +async fn cmd_serve(root: &Utf8Path, mcp: bool, format: codegraph_mcp::OutputStyle) -> Result<()> { if !mcp { return Err(anyhow!("only --mcp transport supported")); } @@ -182,9 +205,9 @@ async fn cmd_serve(root: &Utf8Path, mcp: bool) -> Result<()> { watcher::spawn(root.to_path_buf(), dsn.clone()); } let server = if use_root { - CodegraphServer::with_root(root.to_path_buf()).await? + CodegraphServer::with_root_and_format(root.to_path_buf(), format).await? } else { - CodegraphServer::new() + CodegraphServer::new_with_format(format) }; codegraph_mcp::serve_stdio(server).await -} \ No newline at end of file +}