mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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 <mxjerrett@gmail.com> Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
This commit is contained in:
parent
5cece7bf58
commit
b0afee85b3
3 changed files with 72 additions and 2 deletions
|
|
@ -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 <hex>..<hex>`. 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`.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
61
tests/test_memory/test_hnsw_batch_capacity.py
Normal file
61
tests/test_memory/test_hnsw_batch_capacity.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Add a link
Reference in a new issue