From b0afee85b3a94f829c4743fce2bbf75c494e6312 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Tue, 14 Jul 2026 09:13:12 +0530 Subject: [PATCH] fix(memory): size HNSW index_batch resize off the id high-water mark (#2139) ## Description `HNSWVectorIndex.index_batch()` used the live memory map size to decide whether to resize hnswlib before adding new labels. hnswlib does not reclaim capacity slots when labels are removed with `mark_deleted`, so after delete/evict churn the live count can be much lower than the assigned-id high-water mark. That lets a batch add skip resizing and then fail in `add_items` with `number of elements exceeds the specified limit`. ## Fix - Size the batch resize check from `self._next_hnsw_id`, which has already been incremented for the new batch labels. - Match the single-item `index()` path's high-water-mark capacity behavior. - Add a regression test that deletes most entries from a small index and then batch-adds enough new memories to require a resize. - Merge current `main` to refresh mergeability and stale lint results. ## Testing ```text uvx ruff@0.15.17 check headroom/memory/adapters/hnsw.py tests/test_memory/test_hnsw_batch_capacity.py headroom/memory/factory.py All checks passed! uvx ruff@0.15.17 format --check headroom/memory/adapters/hnsw.py tests/test_memory/test_hnsw_batch_capacity.py headroom/memory/factory.py 3 files already formatted git diff --check headroomlabs/main...HEAD # no output uv run --extra dev python -m pytest tests/test_memory/test_hnsw_batch_capacity.py -q 1 passed, 18 warnings ``` ## Review Readiness - [x] Ready for review - [x] Regression test added - [x] CHANGELOG updated Co-authored-by: JerrettDavis Co-authored-by: Tejas Chopra --- CHANGELOG.md | 1 + headroom/memory/adapters/hnsw.py | 12 +++- tests/test_memory/test_hnsw_batch_capacity.py | 61 +++++++++++++++++++ 3 files changed, 72 insertions(+), 2 deletions(-) create mode 100644 tests/test_memory/test_hnsw_batch_capacity.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 372b70ebf..675b0f670 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -112,6 +112,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **wrap/doctor:** make the Claude Remote Control gate warning accurate and stop it firing for users who never had the feature ([#1779](https://github.com/headroomlabs-ai/headroom/issues/1779)). Claude Code 2.1.196 added a client-side check that **deterministically** disables first-party Remote Control (`/remote-control` / `/rc`) whenever `ANTHROPIC_BASE_URL` points at a non-`api.anthropic.com` host — which Headroom always does. The old notice hedged ("may hide the Remote Control menu"); it now states the disable as fact, names the `/rc` command, and detects the installed Claude Code version so the wording is exact (`2.1.196` when known, `2.1.196+` when not). The gate is upstream and RC's control-plane talks to `claude.ai` (not the API host), so Headroom cannot restore it — the warning tells you to run Claude without Headroom for RC sessions. The warning is suppressed for auth modes that never had Remote Control (API-key/PAYG via `ANTHROPIC_API_KEY`/`ANTHROPIC_AUTH_TOKEN`, and Bedrock/Vertex/Foundry cloud IAM) and on Claude Code builds older than 2.1.196 where RC is unaffected by a custom base URL. Both the `headroom wrap claude` launch banner and `headroom doctor` co-report the sibling base-URL gates Headroom *does* restore — on-demand tool loading (#746, automatic) and the 1M context window (#1158, via `--1m`) — and the wrap-side co-report is session-accurate: it says "already restored via --1m" when the flag is in effect and reports tool deferral OFF (not falsely "kept on") when the user chose `--tool-search false`/`ENABLE_TOOL_SEARCH=false`; the `ENABLE_TOOL_SEARCH=...` banner line got the same accuracy fix. `is_custom_anthropic_base_url` now recognizes scheme-less values (`myproxy.local:8080`, `127.0.0.1:8787`) as custom hosts and degrades gracefully on malformed URLs instead of crashing `doctor`. `doctor` resolves the Claude Code version lazily, so runs with no custom base URL never pay the `claude --version` subprocess. No request bytes are touched (cache-safe); this is UX/notice-only. * **install:** default the docker image to `ghcr.io/headroomlabs-ai/headroom:latest` instead of the dead `ghcr.io/chopratejas/headroom:latest`. After the repo moved to the `headroomlabs-ai` org, GHCR did not redirect the old package, so `headroom install` / `headroom init` and the install scripts pulled a frozen `0.27.0` image while current releases publish to the new path ([#1867](https://github.com/headroomlabs-ai/headroom/issues/1867)). * **transforms/content-router:** stop a profile-derived `read_protection_window` kwarg from weakening an explicit `--protect-tool-results` guarantee. `ContentRouter.apply()` computes `read_protection_window` from `protect_recent_reads_fraction`, where `0.0` (the sentinel `--protect-tool-results` sets) means "protect all excluded-tool output regardless of conversation depth" per #1374's documented contract — but the method then unconditionally overwrote that window with a `read_protection_window` kwarg whenever one was present. `proxy_pipeline_kwargs()` supplies that kwarg on every request from the active `AgentSavingsProfile.protect_recent` (the default `coding` profile sets `protect_recent=2`), so in practice only the last 2 messages ever kept read-protection and older excluded-tool output silently fell through to lossy compression. The runtime kwarg may now only narrow the window when `protect_recent_reads_fraction > 0`; it can no longer shrink the "protect everything" guarantee set by `--protect-tool-results`. +* **memory:** size the HNSW `index_batch` resize off the assigned-id high-water mark, not the live entry count. hnswlib never reclaims a slot on `mark_deleted` (used by remove/evict), so its usable capacity is bounded by the number of ids ever assigned (`_next_hnsw_id`). `index_batch` computed `required_capacity = len(self._memory_to_hnsw) + len(new_memories)` — the *live* count — which drops below `_next_hnsw_id` after deletion/eviction churn, so the resize was skipped and `add_items` raised `RuntimeError: number of elements exceeds the specified limit`, crashing the save path on the HNSW backend. It now resizes off `_next_hnsw_id`, matching the single-item `index()` guard. * **memory:** apply the `turn_id` scope filter even when `agent_id` is absent. In `SQLiteMemoryStore._build_query_conditions` the `turn_id` condition was nested inside the `agent_id` block, so a query filtered by `user_id` + `session_id` + `turn_id` (no `agent_id`) dropped the `turn_id` predicate entirely and returned every memory in the session instead of the single turn — an over-broad result that leaks sibling-turn memories into recall (and makes `count()` wrong for that scope). `agent_id` and `turn_id` are now applied independently. * **transforms:** stop the lossless `diff` fold from silently dropping lines out of non-diff content. `ContentRouter._lossless_first` tries every `compact_lossless` fold on all content, but the `diff` kind (`diff_strip_index`) is the only one with no exact-inverse check — it removes any line shaped like `index ..`. Applied to arbitrary text/log/search payloads that happen to contain such a line, that line was deleted with no CCR marker, so it was unrecoverable — a violation of the lossless no-loss contract the method's own docstring promises. The `diff` fold now runs only when the strategy is `DIFF` or the content is diff-shaped (`_looks_like_diff`); genuine diffs still have their `index` bookkeeping folded. * **memory:** include the Ollama server URL in the embedder cache key so a second backend can't get an embedder bound to the wrong server. `_create_embedder` cached by `(backend, model)` only, but the Ollama embedder is constructed with `base_url=config.ollama_base_url`. Two configs in the same process that shared a backend and model but pointed at different Ollama servers (e.g. a per-project storage router) collided on one cache slot, so the second silently reused the first's embedder and embedded against the wrong host. The cache key now also includes `ollama_base_url`. diff --git a/headroom/memory/adapters/hnsw.py b/headroom/memory/adapters/hnsw.py index c678325ac..9d7f54d47 100644 --- a/headroom/memory/adapters/hnsw.py +++ b/headroom/memory/adapters/hnsw.py @@ -458,8 +458,16 @@ class HNSWVectorIndex: new_memories.append((memory, embedding, hnsw_id)) self._next_hnsw_id += 1 - # Resize if needed - required_capacity = len(self._memory_to_hnsw) + len(new_memories) + # Resize if needed. hnswlib never frees a slot on mark_deleted + # (remove/evict), so its capacity is bounded by the high-water mark + # of assigned ids (_next_hnsw_id, already incremented for the new + # memories above), NOT the live entry count. After deletions or + # evictions the live count is well below _next_hnsw_id, so keying the + # resize off `len(self._memory_to_hnsw)` under-provisions and the + # add_items below raises "number of elements exceeds the specified + # limit". This mirrors the single-item index() guard, which resizes + # off _next_hnsw_id. + required_capacity = self._next_hnsw_id if required_capacity > self._max_elements: new_max = max(self._max_elements * 2, required_capacity + 1000) self._resize_index(new_max) diff --git a/tests/test_memory/test_hnsw_batch_capacity.py b/tests/test_memory/test_hnsw_batch_capacity.py new file mode 100644 index 000000000..ab611e9d1 --- /dev/null +++ b/tests/test_memory/test_hnsw_batch_capacity.py @@ -0,0 +1,61 @@ +"""HNSW index_batch must resize based on the assigned-id high-water mark, not +the live entry count, so batch adds after eviction/deletion churn don't overflow +hnswlib's max_elements.""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +import numpy as np +import pytest + +from headroom.memory.models import Memory + +try: + from headroom.memory.adapters.hnsw import _check_hnswlib_available + + HNSW_AVAILABLE = _check_hnswlib_available() +except ImportError: + HNSW_AVAILABLE = False + + +@pytest.fixture +def temp_hnsw_path(): + with tempfile.NamedTemporaryFile(suffix=".hnsw", delete=False) as f: + yield Path(f.name) + + +def _mem(i: int, dim: int = 8) -> Memory: + rng = np.random.default_rng(i) + return Memory( + content=f"m{i}", + user_id="u", + embedding=rng.standard_normal(dim).astype(np.float32), + ) + + +@pytest.mark.skipif(not HNSW_AVAILABLE, reason="hnswlib not installed") +@pytest.mark.asyncio +async def test_index_batch_after_deletion_churn_does_not_overflow(temp_hnsw_path): + from headroom.memory.adapters.hnsw import HNSWVectorIndex + + # Small ceiling so we hit it quickly. mark_deleted (remove) never frees a + # slot, so the assigned-id counter climbs toward max_elements while the live + # count stays low. + index = HNSWVectorIndex(dimension=8, max_elements=8, save_path=temp_hnsw_path) + + singles = [_mem(i) for i in range(6)] + for m in singles: + await index.index(m) # assigned ids 0..5; next id high-water = 6 + + # Delete 5 of them (mark_deleted; the 5 hnswlib slots are NOT reclaimed). + for m in singles[:5]: + await index.remove(m.id) + + # A batch of 3 now needs slots 6,7,8 -> hnswlib must hold 9 labels. The old + # check used the live count (1) + 3 = 4 <= 8 and skipped the resize, so + # add_items raised "number of elements exceeds the specified limit". + added = await index.index_batch([_mem(100), _mem(101), _mem(102)]) + + assert added == 3