Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,21 @@ with no setup. For the CLI/MCP server it needs a local model directory
- Distill one from any sentence-transformer (Apache-2.0 Jina-Code by default) in
~30 s on CPU: `python scripts/distill_static_model.py`.

#### `CODEGRAPH_SKIP_MEMORY_CHECK` — force the embedding model past the RAM gate

Before loading the ONNX model, the server checks available memory and, if under
~1.5 GB, skips the model to avoid an OOM-kill (running graph-only instead).
Set `CODEGRAPH_SKIP_MEMORY_CHECK=1` (also accepts `true`/`yes`) to bypass that
check and always load the model.

Use it if embeddings are disabled even though the machine has plenty of free
RAM.
A reading of `0 MB available` is treated as a detection failure and the model
loads anyway (macOS parks reclaimable memory in inactive/speculative pages that
some memory readers do not count as free), so this override is mainly for other
cases where the reported figure is low but wrong.
It works in both MCP and one-shot `--run-tool` modes.

#### `--profile` — narrow the MCP tool surface

The full 32-tool surface is convenient but inflates the agent's prompt-context cost. A profile exposes only the slice you need (also settable via the `CODEGRAPH_TOOL_PROFILE` env var):
Expand Down
50 changes: 41 additions & 9 deletions crates/codegraph-server/src/ai_query/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,15 @@ fn available_memory_mb() -> u64 {
sys.available_memory() / (1024 * 1024)
}

/// Whether an available-memory reading counts as real pressure for the embed
/// loop. A 0 MB reading is a detection failure, not genuine pressure — some
/// macOS `sysinfo` versions report 0 available because reclaimable memory is
/// not counted as free (issue #13) — so it must not ratchet the batch size
/// down or force checkpoints.
fn embed_memory_pressured(avail_mb: u64) -> bool {
avail_mb > 0 && avail_mb < EMBED_LOW_MEM_MB
}

/// Split an identifier into camelCase/snake_case words (deduped, lowercased),
/// reusing the BM25 tokenizer: `getUserById` -> "get user by id".
fn split_identifier_words(name: &str) -> String {
Expand Down Expand Up @@ -392,13 +401,13 @@ impl QueryEngine {
pos = end;
chunks_done += 1;

if chunks_done % 10 == 0 && pos < total {
if chunks_done.is_multiple_of(10) && pos < total {
tracing::info!("[QueryEngine] Embedded {}/{} symbols", pos, total);
}

// RAM backpressure: shed batch size under memory pressure.
let pressured = chunks_done % EMBED_MEM_CHECK_CHUNKS == 0
&& available_memory_mb() < EMBED_LOW_MEM_MB;
let pressured = chunks_done.is_multiple_of(EMBED_MEM_CHECK_CHUNKS)
&& embed_memory_pressured(available_memory_mb());
if pressured && chunk_size > 4 {
chunk_size = (chunk_size / 2).max(4);
tracing::warn!(
Expand Down Expand Up @@ -546,8 +555,8 @@ impl QueryEngine {
pos = end;
chunks_done += 1;

let pressured = chunks_done % EMBED_MEM_CHECK_CHUNKS == 0
&& available_memory_mb() < EMBED_LOW_MEM_MB;
let pressured = chunks_done.is_multiple_of(EMBED_MEM_CHECK_CHUNKS)
&& embed_memory_pressured(available_memory_mb());
if pressured && chunk_size > 4 {
chunk_size = (chunk_size / 2).max(4);
tracing::warn!(
Expand Down Expand Up @@ -2632,6 +2641,24 @@ mod tests {
(engine, graph)
}

#[test]
fn embed_memory_pressured_low_reading_is_pressure() {
assert!(embed_memory_pressured(1));
assert!(embed_memory_pressured(EMBED_LOW_MEM_MB - 1));
}

#[test]
fn embed_memory_pressured_zero_reading_is_detection_failure_not_pressure() {
assert!(!embed_memory_pressured(0));
}

#[test]
fn embed_memory_pressured_at_or_above_floor_is_not_pressure() {
assert!(!embed_memory_pressured(EMBED_LOW_MEM_MB));
assert!(!embed_memory_pressured(EMBED_LOW_MEM_MB + 1));
assert!(!embed_memory_pressured(64 * 1024));
}

#[tokio::test]
async fn test_engine_creation() {
let (engine, _) = create_test_engine().await;
Expand Down Expand Up @@ -3964,7 +3991,10 @@ mod tests {

#[test]
fn split_identifier_words_handles_camel_and_snake() {
assert_eq!(split_identifier_words("authenticate_user"), "authenticate user");
assert_eq!(
split_identifier_words("authenticate_user"),
"authenticate user"
);
assert_eq!(split_identifier_words("getUserById"), "get user by id");
// The existing tokenizer keeps acronym+word runs joined (HTML|Parser is
// NOT split) and drops 1-char tokens — known limitations worth revisiting
Expand All @@ -3986,11 +4016,13 @@ mod tests {
// Enabled: split name words are front-loaded for the static embedder.
let with_split =
QueryEngine::build_embed_text(&node, 0, "getUserById", false, true, &graph);
assert!(with_split.starts_with("get user by id"), "got: {with_split}");
assert!(
with_split.starts_with("get user by id"),
"got: {with_split}"
);

// Disabled (default): original transformer-path text is unchanged.
let without =
QueryEngine::build_embed_text(&node, 0, "getUserById", false, false, &graph);
let without = QueryEngine::build_embed_text(&node, 0, "getUserById", false, false, &graph);
assert!(without.starts_with("getUserById"), "got: {without}");
assert!(!without.contains("get user by id"));
}
Expand Down
18 changes: 16 additions & 2 deletions crates/codegraph-server/src/mcp/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -219,10 +219,24 @@ mod imp {
let shared_engine = {
let mut sys = sysinfo::System::new();
sys.refresh_memory();
if sys.available_memory() < 1_500_000_000 {
tracing::warn!("Engine: <1.5 GB free — running graph-only (no shared model)");
let avail = sys.available_memory();
let gate = crate::memory::evaluate_memory_gate(
avail,
crate::memory::MODEL_MIN_FREE_BYTES,
crate::memory::memory_check_bypassed(),
);
if let crate::memory::MemoryGate::Skip(avail) = gate {
tracing::warn!(
"Engine: only {} MB available — skipping shared embedding model to avoid OOM; running graph-only. Set CODEGRAPH_SKIP_MEMORY_CHECK=1 to override.",
avail / 1_000_000
);
None
} else {
if gate == crate::memory::MemoryGate::ProceedDetectionFailed {
tracing::warn!(
"Engine: available memory read as 0 MB — treating as a detection failure (common on macOS, where reclaimable memory is not counted as free) and proceeding with the shared model load. Set CODEGRAPH_SKIP_MEMORY_CHECK=1 to always bypass this check."
);
}
match VectorEngine::from_backend(model_cache_dir(), &cfg.embedding_model) {
Ok(e) => {
tracing::info!("Engine: shared embedding model loaded");
Expand Down
137 changes: 125 additions & 12 deletions crates/codegraph-server/src/memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,62 @@ fn project_data_dir(workspace_path: &Path) -> Result<PathBuf, MemoryError> {
.join(slug))
}

/// Minimum available memory to load the ONNX embedding model + runtime
/// without risking an OOM-kill (a native crash Rust can't catch).
pub(crate) const MODEL_MIN_FREE_BYTES: u64 = 1_500_000_000; // ~1.5 GB

/// Decision of the pre-model-load RAM gate. Separated from the live sysinfo
/// reading so the policy is unit-testable.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum MemoryGate {
/// Enough available memory (or the check was bypassed) — load the model.
Load,
/// Genuinely constrained — skip the model and run graph-only. Carries the
/// detected byte count for the log/error message.
Skip(u64),
/// The available-memory reading was 0, which is a detection failure rather
/// than a real state: a running process always holds memory, and on macOS
/// reclaimable memory is parked in inactive/speculative/purgeable pages
/// that some `sysinfo` versions report as 0 available (issue #13). Proceed
/// with the load instead of hard-disabling on a number we don't trust.
ProceedDetectionFailed,
}

/// Decide whether to load the embedding model from the available-memory
/// reading, the minimum threshold, and whether the user forced a bypass.
///
/// Pure: no I/O, no env reads — the caller supplies the inputs. This is the
/// gate policy; keeping it separate lets tests cover the 0-MB detection-failure
/// path that can't be reproduced by driving `sysinfo` live.
pub(crate) fn evaluate_memory_gate(
available_bytes: u64,
min_free_bytes: u64,
bypass: bool,
) -> MemoryGate {
if bypass {
return MemoryGate::Load;
}
if available_bytes == 0 {
return MemoryGate::ProceedDetectionFailed;
}
if available_bytes < min_free_bytes {
return MemoryGate::Skip(available_bytes);
}
MemoryGate::Load
}

/// Whether the user forced the RAM gate off via `CODEGRAPH_SKIP_MEMORY_CHECK`
/// (`1`/`true`/`yes`, case-insensitive). The escape hatch works in both MCP and
/// one-shot `--run-tool` modes since it is read at model-load time.
pub(crate) fn memory_check_bypassed() -> bool {
std::env::var("CODEGRAPH_SKIP_MEMORY_CHECK")
.map(|v| {
let v = v.trim();
v == "1" || v.eq_ignore_ascii_case("true") || v.eq_ignore_ascii_case("yes")
})
.unwrap_or(false)
}

/// Memory manager for the LSP server
///
/// Opens the database on-demand for each operation and closes it immediately after.
Expand All @@ -192,7 +248,10 @@ pub struct MemoryManager {
impl MemoryManager {
/// Create a new MemoryManager
pub fn new(extension_path: Option<PathBuf>) -> Self {
Self::with_model(extension_path, codegraph_memory::EmbeddingBackend::default())
Self::with_model(
extension_path,
codegraph_memory::EmbeddingBackend::default(),
)
}

/// Create a new MemoryManager with a specific embedding backend
Expand Down Expand Up @@ -284,21 +343,30 @@ impl MemoryManager {
let mut sys = sysinfo::System::new();
sys.refresh_memory();
let avail = sys.available_memory();
const MIN_FREE_BYTES: u64 = 1_500_000_000; // ~1.5 GB for model + ort runtime
tracing::info!(
"[MemoryManager::initialize] available memory: {} MB",
avail / 1_000_000
);
if avail < MIN_FREE_BYTES {
crate::crash_phase::mark("onnx_skipped_lowmem");
tracing::warn!(
"[MemoryManager::initialize] only {} MB free — skipping embedding model to avoid OOM; semantic search disabled (graph-only)",
avail / 1_000_000
);
return Err(MemoryError::Other(format!(
"insufficient memory ({} MB free) to load embedding model; running graph-only",
avail / 1_000_000
)));
match evaluate_memory_gate(avail, MODEL_MIN_FREE_BYTES, memory_check_bypassed()) {
MemoryGate::Skip(avail) => {
crate::crash_phase::mark("onnx_skipped_lowmem");
tracing::warn!(
"[MemoryManager::initialize] only {} MB available — skipping embedding model to avoid OOM; semantic search disabled (graph-only). Set CODEGRAPH_SKIP_MEMORY_CHECK=1 to override.",
avail / 1_000_000
);
return Err(MemoryError::Other(format!(
"insufficient memory ({} MB available) to load embedding model; running graph-only (set CODEGRAPH_SKIP_MEMORY_CHECK=1 to override)",
avail / 1_000_000
)));
}
MemoryGate::ProceedDetectionFailed => {
// 0 MB is a detection failure (see MemoryGate) — don't
// disable embeddings on a Mac that actually has RAM.
tracing::warn!(
"[MemoryManager::initialize] available memory read as 0 MB — treating as a detection failure (common on macOS, where reclaimable memory is not counted as free) and proceeding with the embedding model load. Set CODEGRAPH_SKIP_MEMORY_CHECK=1 to always bypass this check."
);
}
MemoryGate::Load => {}
}
}

Expand Down Expand Up @@ -612,6 +680,51 @@ pub use codegraph_memory::{
mod tests {
use super::*;

const MIN: u64 = MODEL_MIN_FREE_BYTES;

#[test]
fn gate_loads_when_ample_memory() {
assert_eq!(
evaluate_memory_gate(8_000_000_000, MIN, false),
MemoryGate::Load
);
}

#[test]
fn gate_loads_exactly_at_threshold() {
// `< min` is the skip condition, so exactly `min` must load.
assert_eq!(evaluate_memory_gate(MIN, MIN, false), MemoryGate::Load);
}

#[test]
fn gate_skips_when_genuinely_low() {
assert_eq!(
evaluate_memory_gate(500_000_000, MIN, false),
MemoryGate::Skip(500_000_000)
);
}

#[test]
fn gate_proceeds_on_zero_reading_as_detection_failure() {
// The issue-13 case: sysinfo reports 0 available on a healthy Mac.
// 0 is never a real state, so proceed rather than disable embeddings.
assert_eq!(
evaluate_memory_gate(0, MIN, false),
MemoryGate::ProceedDetectionFailed
);
}

#[test]
fn gate_bypass_forces_load_regardless_of_reading() {
// CODEGRAPH_SKIP_MEMORY_CHECK=1 overrides even a genuine low reading.
assert_eq!(evaluate_memory_gate(0, MIN, true), MemoryGate::Load);
assert_eq!(evaluate_memory_gate(1, MIN, true), MemoryGate::Load);
assert_eq!(
evaluate_memory_gate(500_000_000, MIN, true),
MemoryGate::Load
);
}

#[test]
fn generation_zero_maps_to_historical_path() {
let dir = Path::new("/home/u/.codegraph");
Expand Down
5 changes: 5 additions & 0 deletions mcp-package/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,11 @@ Pass flags after `--`:
| `--graph-only` | off | Skip embeddings — graph + structural tools only. No ONNX model load, 10-50× faster indexing. For CI / one-shot graph queries. |
| `--run-tool <name>` | — | One-shot: index, run a single tool, print result, exit. No MCP handshake. Pair with `--tool-args '<json>'`. |

Before loading the ONNX embedding model, the server checks available memory and runs graph-only if under ~1.5 GB.
If embeddings are disabled even though the machine has plenty of free RAM, set `CODEGRAPH_SKIP_MEMORY_CHECK=1` (also accepts `true`/`yes`) to bypass the check.
A reading of `0 MB available` is treated as a detection failure and the model loads anyway (common on macOS).
Works in both MCP and one-shot `--run-tool` modes.

### Agent rules (recommended)

Pre-configured rule files that teach your AI agent to use CodeGraph tools before falling back to grep / multi-file reads:
Expand Down
Loading