kernel-lore-mcp 0.3.2__tar.gz → 0.3.3__tar.gz

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/Cargo.lock

@@ -2638,7 +2638,7 @@ dependencies = [
 
 [[package]]
 name = "kernel_lore_mcp"
-version = "0.3.2"
+version = "0.3.3"
 dependencies = [
  "aho-corasick",
  "anyhow",

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "kernel_lore_mcp"
-version = "0.3.2"
+version = "0.3.3"
 edition = "2024"
 rust-version = "1.85"
 authors = ["Michael Bommarito <michael@bommaritollc.com>"]

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel-lore-mcp
-Version: 0.3.2
+Version: 0.3.3
 Classifier: Development Status :: 2 - Pre-Alpha
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernel-lore-mcp"
-version = "0.3.2"
+version = "0.3.3"
 description = "MCP server exposing fast, structured search over lore.kernel.org to LLM developer tools."
 readme = "README.md"
 license = { text = "MIT" }

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/bin/sync.rs

@@ -114,6 +114,11 @@ struct PhaseHeartbeat {
     handle: Option<std::thread::JoinHandle<()>>,
 }
 
+struct StageHeartbeat {
+    stop: std::sync::Arc<AtomicBool>,
+    handle: Option<std::thread::JoinHandle<()>>,
+}
+
 #[derive(Debug, Clone, Serialize)]
 struct SyncStatusRecord {
     version: String,
@@ -195,6 +200,23 @@ impl SyncStateReporter {
         })
     }
 
+    fn refresh_stage(
+        &self,
+        stage: &str,
+        note: Option<String>,
+        elapsed_secs: f64,
+        memory: MemorySnapshot,
+    ) -> Result<()> {
+        self.update(|status| {
+            status.stage = stage.to_string();
+            status.note = note;
+            status.elapsed_secs = elapsed_secs;
+            status.process_rss_mb = memory.process_rss_bytes.map(bytes_to_mib);
+            status.memory_available_mb = memory.available_bytes.map(bytes_to_mib);
+            status.memory_total_mb = memory.total_bytes.map(bytes_to_mib);
+        })
+    }
+
     fn set_phase_progress(
         &self,
         stage: &str,
@@ -331,6 +353,46 @@ impl Drop for PhaseHeartbeat {
     }
 }
 
+impl Drop for StageHeartbeat {
+    fn drop(&mut self) {
+        self.stop.store(true, Ordering::Relaxed);
+        if let Some(handle) = self.handle.take() {
+            handle.thread().unpark();
+            let _ = handle.join();
+        }
+    }
+}
+
+fn stage_heartbeat(
+    sync_state: SyncStateReporter,
+    stage: &'static str,
+    note: Option<String>,
+    run_started: Instant,
+) -> StageHeartbeat {
+    let stop = std::sync::Arc::new(AtomicBool::new(false));
+    let stop_for_thread = std::sync::Arc::clone(&stop);
+    let stage_name = stage.to_string();
+    let handle = std::thread::spawn(move || {
+        loop {
+            std::thread::park_timeout(PROGRESS_LOG_INTERVAL);
+            if stop_for_thread.load(Ordering::Relaxed) {
+                break;
+            }
+            let elapsed_secs = run_started.elapsed().as_secs_f64();
+            let memory = current_memory_snapshot();
+            if let Err(e) =
+                sync_state.refresh_stage(&stage_name, note.clone(), elapsed_secs, memory)
+            {
+                tracing::warn!(stage = stage_name, error = %e, "failed to write sync status");
+            }
+        }
+    });
+    StageHeartbeat {
+        stop,
+        handle: Some(handle),
+    }
+}
+
 fn main() -> Result<()> {
     tracing_subscriber::fmt()
         .with_writer(std::io::stderr)
@@ -688,16 +750,15 @@ fn main() -> Result<()> {
     if let Some(ref bm25_mutex) = bm25 {
         let memory = current_memory_snapshot();
         let config = bm25_config.expect("bm25 config must exist when writer exists");
+        let note = Some(format!(
+            "threads={} memory_budget_mb={}",
+            config.num_threads,
+            config.memory_budget_mb()
+        ));
         sync_state
-            .set_stage(
-                "bm25_commit",
-                Some(format!(
-                    "threads={} memory_budget_mb={}",
-                    config.num_threads,
-                    config.memory_budget_mb()
-                )),
-            )
+            .set_stage("bm25_commit", note.clone())
             .context("update sync status to bm25_commit")?;
+        let stage_guard = stage_heartbeat(sync_state.clone(), "bm25_commit", note, start);
         tracing::info!(
             bm25_writer_threads = config.num_threads,
             bm25_memory_budget_mb = config.memory_budget_mb(),
@@ -718,6 +779,7 @@ fn main() -> Result<()> {
             memory_total_mb = ?memory.total_bytes.map(bytes_to_mib),
             "bm25 commit done"
         );
+        drop(stage_guard);
     }
 
     // tid rebuild: reasonable to do every sync tick; touches only the
@@ -725,6 +787,7 @@ fn main() -> Result<()> {
     sync_state
         .set_stage("tid_rebuild", None)
         .context("update sync status to tid_rebuild")?;
+    let tid_stage_guard = stage_heartbeat(sync_state.clone(), "tid_rebuild", None, start);
     tracing::info!("tid rebuild starting");
     let tid_result = _core::rebuild_tid(&data_dir).context("rebuild tid side-table")?;
     tracing::info!(
@@ -732,9 +795,11 @@ fn main() -> Result<()> {
         path = %tid_result.0.display(),
         "tid rebuild done"
     );
+    drop(tid_stage_guard);
     sync_state
         .set_stage("path_vocab_rebuild", None)
         .context("update sync status to path_vocab_rebuild")?;
+    let path_stage_guard = stage_heartbeat(sync_state.clone(), "path_vocab_rebuild", None, start);
     tracing::info!("path vocab rebuild starting");
     if let Some(path_count) = rebuild_path_vocab_after_sync(&data_dir)? {
         tracing::info!(
@@ -743,6 +808,7 @@ fn main() -> Result<()> {
             "path vocab rebuild done"
         );
     }
+    drop(path_stage_guard);
 
     // Tally.
     let mut total_ingested: u64 = 0;

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/error.rs

@@ -18,6 +18,9 @@ pub enum Error {
     #[error("query exceeded {limit_ms} ms wall-clock limit")]
     QueryTimeout { limit_ms: u64 },
 
+    #[error("query cancelled by caller")]
+    QueryCancelled,
+
     #[error("invalid cursor: {0}")]
     InvalidCursor(String),
 

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/kernel_lore_mcp/__init__.py

@@ -15,7 +15,7 @@ from typing import Any
 
 __all__ = ["__version__", "native_version"]
 
-__version__ = "0.3.2"
+__version__ = "0.3.3"
 
 
 def __getattr__(name: str) -> Any:

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/kernel_lore_mcp/_core.pyi

@@ -29,6 +29,14 @@ class TidRebuildResult(TypedDict):
     path: str
     rows: int
 
+class _RequestCancelGuard:
+    def close(self) -> None: ...
+
+class _RequestCancelToken:
+    def __init__(self) -> None: ...
+    def cancel(self) -> None: ...
+    def install(self) -> _RequestCancelGuard: ...
+
 def rebuild_tid(data_dir: str | PathLike[str]) -> TidRebuildResult: ...
 def rebuild_bm25(data_dir: str | PathLike[str]) -> int: ...
 def backfill_subject_normalized(data_dir: str | PathLike[str]) -> int: ...

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/kernel_lore_mcp/cost_class.py

@@ -41,7 +41,7 @@ import structlog
 
 from kernel_lore_mcp.config import get_settings
 from kernel_lore_mcp.errors import LoreError
-from kernel_lore_mcp.health import suggest_retry_after_seconds
+from kernel_lore_mcp.health import should_shed_tool_call, suggest_retry_after_seconds
 from kernel_lore_mcp.logging_ import profiling_thresholds
 
 CostClass = Literal["cheap", "moderate", "expensive"]
@@ -147,6 +147,30 @@ def cost_limited(
         # Pair with a zero-timeout acquire for the "reject fast"
         # shape — acquire() would block if locked.
         base_started = current_tool_request_started() or time.monotonic()
+        settings = get_settings()
+        should_shed, shed_reason = should_shed_tool_call(
+            data_dir=settings.data_dir,
+            cost_class=cost,
+        )
+        if should_shed:
+            queue_wait = max(0.0, time.monotonic() - base_started)
+            record_tool_queue_wait(
+                tool_name,
+                cost,
+                queue_wait,
+                "rate_limited",
+            )
+            log.warning(
+                "tool admission shed",
+                tool=tool_name,
+                cost_class=cost,
+                mode=settings.mode,
+                queue_wait_ms=round(queue_wait * 1000, 3),
+                inflight=current_inflight(cost),
+                limit=_LIMITS[cost],
+                reason=shed_reason,
+            )
+            raise _rate_limited(cost, tool_name)
         try:
             await asyncio.wait_for(sem.acquire(), timeout=0.001)
         except TimeoutError as err:
@@ -157,7 +181,6 @@ def cost_limited(
                 queue_wait,
                 "rate_limited",
             )
-            settings = get_settings()
             log.warning(
                 "tool admission rejected",
                 tool=tool_name,

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/kernel_lore_mcp/health.py

@@ -11,6 +11,14 @@ from pathlib import Path
 from typing import Any
 
 _ACTIVE_SYNC_STALE_AFTER_SECONDS = 30
+_SYNC_SHED_STAGES = {
+    "ingest",
+    "bm25_commit",
+    "tid_rebuild",
+    "path_vocab_rebuild",
+    "generation_bump",
+    "save_manifest",
+}
 
 
 def writer_lock_present(data_dir: Path) -> bool:
@@ -62,9 +70,7 @@ def read_sync_state(data_dir: Path) -> dict[str, Any] | None:
     active = bool(payload.get("active"))
     if heartbeat_age_seconds is not None:
         active = (
-            active
-            and writer_active
-            and heartbeat_age_seconds <= _ACTIVE_SYNC_STALE_AFTER_SECONDS
+            active and writer_active and heartbeat_age_seconds <= _ACTIVE_SYNC_STALE_AFTER_SECONDS
         )
     else:
         active = active and writer_active
@@ -122,6 +128,26 @@ def suggest_retry_after_seconds(
     return retry_after, None
 
 
+def should_shed_tool_call(*, data_dir: Path, cost_class: str) -> tuple[bool, str | None]:
+    """Whether the server should fast-reject this tool class.
+
+    Cheap indexed lookups stay available as long as the process is up.
+    Moderate and expensive tools shed aggressively while a writer-heavy
+    sync stage is active so a serving box does not spend its remaining
+    headroom on work we already expect to be slow or likely to timeout.
+    """
+    if cost_class == "cheap":
+        return False, None
+
+    sync = read_sync_state(data_dir)
+    if sync and sync.get("active"):
+        stage = str(sync.get("stage") or "").strip()
+        if stage in _SYNC_SHED_STAGES:
+            return True, f"sync stage `{stage}` is reserving capacity"
+
+    return False, None
+
+
 def _coerce_int(value: object) -> int | None:
     if isinstance(value, bool):
         return None

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/kernel_lore_mcp/routes/status.py

@@ -66,10 +66,9 @@ def _per_list_shards(data_dir: Path) -> dict[str, list[dict[str, str]]]:
     return out
 
 
-def _build_status(settings: Settings) -> dict[str, Any]:
+def _build_status(settings: Settings, *, include_per_list: bool) -> dict[str, Any]:
     data_dir = settings.data_dir
     generation, last_ingest = _read_generation(data_dir)
-    per_list = _per_list_shards(data_dir)
     writer_lock = writer_lock_present(data_dir)
     sync = read_sync_state(data_dir)
 
@@ -85,7 +84,7 @@ def _build_status(settings: Settings) -> dict[str, Any]:
         # that pairs with the Prometheus gauge.
         freshness_ok = last_ingest_age_seconds < 3 * interval
 
-    return {
+    body = {
         "service": "kernel-lore-mcp",
         "version": _native_version(),
         "generation": generation,
@@ -96,10 +95,13 @@ def _build_status(settings: Settings) -> dict[str, Any]:
         "writer_lock_present": writer_lock,
         "sync_active": bool(sync and sync.get("active")),
         "sync": sync,
-        "per_list": per_list,
+        "per_list_omitted": not include_per_list,
         "capabilities": capabilities(data_dir),
         "blind_spots_ref": "blind-spots://coverage",
     }
+    if include_per_list:
+        body["per_list"] = _per_list_shards(data_dir)
+    return body
 
 
 def capabilities(data_dir: Path) -> dict[str, bool]:
@@ -158,7 +160,11 @@ def _native_version() -> str:
         return "unknown"
 
 
-def get_status(settings: Settings | None = None) -> dict[str, Any]:
+def get_status(
+    settings: Settings | None = None,
+    *,
+    include_per_list: bool = False,
+) -> dict[str, Any]:
     """Cached status builder. Used by both the MCP route and tests.
 
     Cache is keyed by `data_dir` so multiple in-process servers (or
@@ -169,7 +175,7 @@ def get_status(settings: Settings | None = None) -> dict[str, Any]:
         from kernel_lore_mcp.config import get_settings
 
         settings = get_settings()
-    cache_key = str(settings.data_dir)
+    cache_key = f"{settings.data_dir}|per_list={int(include_per_list)}"
     ttl = settings.freshness_cache_ttl_seconds
     now = time.monotonic()
     live_writer = writer_lock_present(settings.data_dir)
@@ -178,7 +184,7 @@ def get_status(settings: Settings | None = None) -> dict[str, Any]:
         cached_at, body = _cache[cache_key]
         if now - cached_at < effective_ttl:
             return body
-    body = _build_status(settings)
+    body = _build_status(settings, include_per_list=include_per_list)
     _cache[cache_key] = (now, body)
     return body
 
@@ -190,4 +196,6 @@ def clear_cache() -> None:
 
 async def status_endpoint(request: Request) -> JSONResponse:
     """Starlette/FastMCP custom-route handler."""
-    return JSONResponse(get_status())
+    raw = (request.query_params.get("per_list") or "").strip().lower()
+    include_per_list = raw in {"1", "true", "yes", "on"}
+    return JSONResponse(get_status(include_per_list=include_per_list))

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/kernel_lore_mcp/timeout.py

@@ -24,6 +24,7 @@ from kernel_lore_mcp.config import get_settings
 from kernel_lore_mcp.errors import LoreError
 from kernel_lore_mcp.health import suggest_retry_after_seconds
 from kernel_lore_mcp.logging_ import profiling_thresholds
+from kernel_lore_mcp import _core
 
 log = structlog.get_logger(__name__)
 
@@ -45,9 +46,18 @@ async def run_with_timeout[T](
     ms = timeout_ms or settings.query_wall_clock_ms
     tool_name = getattr(fn, "__name__", str(fn))
     started = time.monotonic()
+    token = _core._RequestCancelToken()
+
+    def _invoke_with_cancel() -> T:
+        guard = token.install()
+        try:
+            return fn(*args)
+        finally:
+            guard.close()
+
     try:
         result = await asyncio.wait_for(
-            asyncio.to_thread(fn, *args),
+            asyncio.to_thread(_invoke_with_cancel),
             timeout=ms / 1000.0,
         )
         elapsed = time.monotonic() - started
@@ -85,10 +95,24 @@ async def run_with_timeout[T](
             echoed_input=echoed_input or {},
             retry_after_seconds=retry_after_seconds,
         ) from None
+    except asyncio.CancelledError:
+        elapsed = time.monotonic() - started
+        record_tool_runtime(tool_name, elapsed, "cancelled")
+        log.warning(
+            "tool runtime cancelled",
+            operation=tool_name,
+            mode=settings.mode,
+            runtime_ms=round(elapsed * 1000, 3),
+            timeout_ms=ms,
+        )
+        raise
     except LoreError as exc:
         elapsed = time.monotonic() - started
         record_tool_runtime(tool_name, elapsed, exc.code)
-        if exc.code != "query_timeout" and elapsed >= profiling_thresholds(settings.mode).tool_seconds:
+        if (
+            exc.code != "query_timeout"
+            and elapsed >= profiling_thresholds(settings.mode).tool_seconds
+        ):
             log.info(
                 "tool runtime completed",
                 operation=tool_name,
@@ -109,6 +133,8 @@ async def run_with_timeout[T](
             timeout_ms=ms,
         )
         raise
+    finally:
+        token.cancel()
 
 
 __all__ = ["run_with_timeout"]

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/lib.rs

@@ -113,6 +113,8 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(crate::python::py_embedding_meta, m)?)?;
     m.add_class::<crate::python::PyReader>()?;
     m.add_class::<crate::python::PyEmbeddingBuilder>()?;
+    m.add_class::<crate::python::PyRequestCancelToken>()?;
+    m.add_class::<crate::python::PyRequestCancelGuard>()?;
     Ok(())
 }
 

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/python.rs

@@ -24,7 +24,53 @@ use crate::ingest;
 use crate::reader::{DiffMode, EqField, MessageRow, Reader as CoreReader, RegexField};
 use crate::router;
 use crate::tid;
-use crate::timeout::DeadlineGuard;
+use crate::timeout::{CancelGuard, CancelToken, DeadlineGuard};
+
+#[pyclass(
+    module = "kernel_lore_mcp._core",
+    name = "_RequestCancelToken",
+    skip_from_py_object
+)]
+#[derive(Clone)]
+pub struct PyRequestCancelToken {
+    inner: CancelToken,
+}
+
+#[pymethods]
+impl PyRequestCancelToken {
+    #[new]
+    fn new() -> Self {
+        Self {
+            inner: CancelToken::new(),
+        }
+    }
+
+    fn cancel(&self) {
+        self.inner.cancel();
+    }
+
+    fn install(&self) -> PyRequestCancelGuard {
+        PyRequestCancelGuard {
+            guard: Some(CancelGuard::install(self.inner.clone())),
+        }
+    }
+}
+
+#[pyclass(
+    module = "kernel_lore_mcp._core",
+    name = "_RequestCancelGuard",
+    skip_from_py_object
+)]
+pub struct PyRequestCancelGuard {
+    guard: Option<CancelGuard>,
+}
+
+#[pymethods]
+impl PyRequestCancelGuard {
+    fn close(&mut self) {
+        self.guard.take();
+    }
+}
 
 /// Install a thread-local deadline for the duration of one reader
 /// query, matching the router-layer wall-clock cap. Cheap — ~one

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/reader.rs

@@ -2571,6 +2571,7 @@ fn parallel_confirm_substring(
 
     const CHUNK: usize = 256;
     let deadline = crate::timeout::current_deadline();
+    let cancel = crate::timeout::current_cancel_token();
     let mut confirmed: Vec<MessageRow> = Vec::with_capacity(limit);
     for chunk in date_sorted_hits.chunks(CHUNK) {
         crate::timeout::check_request_deadline()?;
@@ -2584,6 +2585,7 @@ fn parallel_confirm_substring(
             .par_iter()
             .map(|row| -> Result<Option<MessageRow>> {
                 let _g = deadline.map(crate::timeout::DeadlineGuard::install);
+                let _c = cancel.clone().map(crate::timeout::CancelGuard::install);
                 crate::timeout::check_request_deadline()?;
                 let store = reader.store_for(&row.list)?;
                 let body = store.read_at(row.body_segment_id, row.body_offset)?;
@@ -2653,9 +2655,11 @@ fn collect_trigram_candidates(
         Mutex::new(std::collections::HashSet::new());
     let overflowed = AtomicBool::new(false);
     let deadline = crate::timeout::current_deadline();
+    let cancel = crate::timeout::current_cancel_token();
 
     all_segs.par_iter().try_for_each(|seg_dir| -> Result<()> {
         let _g = deadline.map(crate::timeout::DeadlineGuard::install);
+        let _c = cancel.clone().map(crate::timeout::CancelGuard::install);
         crate::timeout::check_request_deadline()?;
         if overflowed.load(Ordering::Relaxed) {
             return Ok(());
@@ -2710,6 +2714,7 @@ fn parallel_confirm_fuzzy(
 
     const CHUNK: usize = 256;
     let deadline = crate::timeout::current_deadline();
+    let cancel = crate::timeout::current_cancel_token();
     let mut confirmed: Vec<MessageRow> = Vec::with_capacity(limit);
     for chunk in date_sorted_hits.chunks(CHUNK) {
         crate::timeout::check_request_deadline()?;
@@ -2720,6 +2725,7 @@ fn parallel_confirm_fuzzy(
             .par_iter()
             .map(|row| -> Result<Option<MessageRow>> {
                 let _g = deadline.map(crate::timeout::DeadlineGuard::install);
+                let _c = cancel.clone().map(crate::timeout::CancelGuard::install);
                 crate::timeout::check_request_deadline()?;
                 let store = reader.store_for(&row.list)?;
                 let body = store.read_at(row.body_segment_id, row.body_offset)?;

{kernel_lore_mcp-0.3.2 → kernel_lore_mcp-0.3.3}/src/router.rs

@@ -398,20 +398,24 @@ pub fn dispatch(
     // DeadlineGuard::install so scan() checks within the tier honor
     // the same budget.
     let deadline = crate::timeout::current_deadline();
+    let cancel = crate::timeout::current_cancel_token();
     let mut meta_out: TierOut = Ok((None, None));
     let mut trigram_out: TierOut = Ok((None, None));
     let mut bm25_out: TierOut = Ok((None, None));
     rayon::scope(|s| {
         s.spawn(|_| {
             let _g = deadline.map(crate::timeout::DeadlineGuard::install);
+            let _c = cancel.clone().map(crate::timeout::CancelGuard::install);
             meta_out = run_metadata();
         });
         s.spawn(|_| {
             let _g = deadline.map(crate::timeout::DeadlineGuard::install);
+            let _c = cancel.clone().map(crate::timeout::CancelGuard::install);
             trigram_out = run_trigram();
         });
         s.spawn(|_| {
             let _g = deadline.map(crate::timeout::DeadlineGuard::install);
+            let _c = cancel.clone().map(crate::timeout::CancelGuard::install);
             bm25_out = run_bm25();
         });
     });