Finetune to reduce LLM token

Cargo.toml

@@ -24,7 +24,9 @@ authors = ["Cleboost <clement.balarot@gmail.com>"]
 [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"

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<SharedGraphIndex>,
+    /// Session store cho search resumable (resume id → cursor).
+    sessions: Arc<SearchSessionStore>,
+}
+
+// ==================== 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<HashMap<String, StoredResume>>,
+    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<Symbol>,
+    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<String>,
+    /// 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<SharedGraphIndex>) -> 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<SharedGraphIndex>,
+        sessions: Arc<SearchSessionStore>,
+    ) -> 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<String>,
+        timeout_ms: u64,
+    ) -> Result<ResumeSearchOutcome> {
+        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<SymbolKind>,
+        mode: SymbolMatch,
+        pagination: Pagination,
+        resume: Option<String>,
+        timeout_ms: u64,
+    ) -> Result<ResumeSearchOutcome> {
+        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<MemberInfo> {
         self.index().await.list_methods_of_class(id)