From 611254d759a061db8af4e010598186ca61563b96 Mon Sep 17 00:00:00 2001 From: lbruton Date: Tue, 27 Oct 2026 13:51:03 -0500 Subject: [PATCH 1/2] fix(DEVS-145): resolve three post-merge Codex findings on PR #78 Codex posted three review threads on #78 seconds before it merged, so the /pr-resolve pass never saw them. - Shared embed cooldown in both directions (P1). EmbeddingBudget.before_batch checked only its own last batch, so a transcript batch could start right after a memory embed. It now also honors EmbedThrottle.remaining(). A memory embed queued on mlx-embed behind a transcript batch could also start at once; _embed_after_cooldown now calls EmbedThrottle.wait_blocking() on mlx-embed and stamps the throttle there. - No vector search under another model (P2). vector_search re-reads the stored memory model and raises MemoryModelMismatch instead of querying Milvus; same-size vectors from another model would otherwise match. MemoryService already answers from FTS when the vector channel fails. - CLI audit project inventory (P2). memory_cli audit passes memory_service.project_inventory() (memory.project, else tag, from each known root's .claude/project.json), so unknown projects are reported. The report carries inventory_checked, and the CLI says when the check could not run instead of printing a clean result. Tests: 9 new (3 throttle, 2 index, 1 service, 3 CLI), written first and failing before the fix. Full suite 725 passed, 0 skipped with the test DSNs set; ruff clean; Vale clean on the changed docs. --- AGENTS.md | 4 +- CLAUDE.md | 4 +- README.md | 2 +- embedding_control.py | 38 +++++++++++++----- memory_cli.py | 5 +++ memory_index.py | 29 ++++++++++++-- memory_service.py | 23 ++++++++--- tests/test_embedding_control.py | 68 +++++++++++++++++++++++++++++++++ tests/test_memory_cli.py | 53 +++++++++++++++++++++++++ tests/test_memory_index.py | 31 +++++++++++++++ tests/test_memory_service.py | 12 ++++++ 11 files changed, 244 insertions(+), 25 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f72f55e..7185e57 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -45,8 +45,8 @@ Semantic search over agent session transcripts. Independent project, originally - **EmbeddingGemma padding mask (SESF-57)** — mlx-embeddings 0.0.5 hands the 0/1 tokenizer mask to SDPA as an *additive* float mask, so padding leaked into shorter docs in a batch. `rag_engine._embed_gemma` runs the gemma forward itself with a boolean key-padding mask (`_key_padding_mask`). Batched now matches single to cos ≥ 0.99999, and a batch of one (every query) is unchanged. `tests/test_embed_padding.py` pins this, so don't revert to calling `model(input_ids, attention_mask=...)`. The 512-token sliding window is still not applied on local layers (SESF-62, measure first). Stored vectors were rewritten in place by `reembed.py` (`cleanup.py reembed {plan|run|drive|status|verify}`; chunked under the MLX crash limit; old vectors backed up to `~/.sessionflow/audit/reembed-.f16`). - **Memory store is Postgres (DEVS-145)** — memories commit to Postgres first; the Milvus `memories` collection and `~/.sessionflow/memories_fts.db` are derived indexes. The DSN comes from `SESSIONFLOW_MEMORY_DSN` only: SessionFlow calls no secrets tool and reads no DSN file, so a launchd install sets it in the LaunchAgent environment. Memory starts after transcript startup on its own thread; a missing DSN or a down database leaves `/health` `memory.ready` false with the reason and transcript search untouched. Blocking store calls go through `memory_store.run_db` (the `memory-db` executor), never the event loop or `mlx-embed`. - **Never write a memory index directly (DEVS-145)** — every write goes through `MemoryService` to Postgres, which resets that memory's rows in `memory_index_acks` in the same transaction. The sync loop (`SESSIONFLOW_MEMORY_SYNC_INTERVAL`, plus a full reconcile at start and every 6 hours) applies what each consumer (`milvus`, `fts:`) lacks. A Milvus write re-reads the row under `pg_advisory_xact_lock` and upserts only if the embedded hash still matches; the FTS file keeps a per-record watermark so an older write cannot revive a deleted record. `change_seq` orders work but does not follow commit order, so pending work is decided by acknowledgements, not by a cursor. -- **Memory embeds share the MLX thread (DEVS-145)** — `memory_index.embed_for_index` waits on `embedding_control.EmbedThrottle` (200 ms floor) and runs on `mlx-embed`; it never uses the transcript budget, which drops batches. `memory_cli` never embeds: `rebuild` and `reindex` only reset acknowledgements, and `restore-check` asks the running server's loopback `/memory/restore-check` route to index into a temporary `memories_restore` collection. -- **Memory embedding generation (DEVS-145)** — `memory_index_meta` names the memory model and generation. `cleanup.py reembed` starts a new generation on a fresh run and resets the Milvus acknowledgements; the server re-embeds from Postgres text. A server whose local model differs refuses memory vector writes and reports it in `/health` (`memory.index.milvus_blocked`). +- **Memory embeds share the MLX thread (DEVS-145)** — `memory_index.embed_for_index` waits on `embedding_control.EmbedThrottle` (200 ms floor) and runs on `mlx-embed`, where it waits again and stamps the throttle; it never uses the transcript budget, which drops batches. `EmbeddingBudget.before_batch` honors the same throttle, so neither consumer starts inside the other's cooldown. `memory_cli` never embeds: `rebuild` and `reindex` only reset acknowledgements, and `restore-check` asks the running server's loopback `/memory/restore-check` route to index into a temporary `memories_restore` collection. +- **Memory embedding generation (DEVS-145)** — `memory_index_meta` names the memory model and generation. `cleanup.py reembed` starts a new generation on a fresh run and resets the Milvus acknowledgements; the server re-embeds from Postgres text. A server whose local model differs refuses memory vector writes and vector searches (`MemoryModelMismatch`; search answers from FTS) and reports it in `/health` (`memory.index.milvus_blocked`). - **Memory secret check (DEVS-145)** — `add_memory` rejects text with any `secret_redaction.scan_spans` hit and names the rule. Paths and code identifiers can trip it; the fix is a narrow pattern in the operator allowlist (`SESSIONFLOW_REDACT_ALLOWLIST`), never a looser rule. - **Transcript cleanup leaves memory alone (DEVS-145)** — age expiry, `cleanup_sessions`, `sanitize --drop`, `reset` and `migrate-schema` act on `sessions` and `fts.db` only; `reset` and `migrate-schema` say so. - **Memory tests (DEVS-145)** — Postgres tests need `SESSIONFLOW_TEST_PG_DSN` (one fresh schema per test through the `memory_pg` fixture) and skip with that reason when it is unset; restore-check tests also need `SESSIONFLOW_TEST_PG_RESTORE_DSN` and `pg_dump`/`pg_restore` 17. An autouse fixture clears `SESSIONFLOW_MEMORY_DSN`. `stub_rag_engine` evicts project modules first imported under the stub, so a later test imports them against the real `rag_engine`. pytest prints fixture DSNs in setup errors, so mask `password=` in any shared output. diff --git a/CLAUDE.md b/CLAUDE.md index a10c750..f0f29ee 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -46,8 +46,8 @@ Semantic search over Claude Code session transcripts. Independent project, origi - **EmbeddingGemma padding mask (SESF-57)** — mlx-embeddings 0.0.5 hands the 0/1 tokenizer mask to SDPA as an *additive* float mask, so padding leaked into shorter docs in a batch. `rag_engine._embed_gemma` runs the gemma forward itself with a boolean key-padding mask (`_key_padding_mask`). Batched now matches single to cos ≥ 0.99999, and a batch of one (every query) is unchanged. `tests/test_embed_padding.py` pins this, so don't revert to calling `model(input_ids, attention_mask=...)`. The 512-token sliding window is still not applied on local layers (SESF-62, measure first). Stored vectors were rewritten in place by `reembed.py` (`cleanup.py reembed {plan|run|drive|status|verify}`; chunked under the MLX crash limit; old vectors backed up to `~/.sessionflow/audit/reembed-.f16`). - **Memory store is Postgres (DEVS-145)** — memories commit to Postgres first; the Milvus `memories` collection and `~/.sessionflow/memories_fts.db` are derived indexes. The DSN comes from `SESSIONFLOW_MEMORY_DSN` only: SessionFlow calls no secrets tool and reads no DSN file, so a launchd install sets it in the LaunchAgent environment. Memory starts after transcript startup on its own thread; a missing DSN or a down database leaves `/health` `memory.ready` false with the reason and transcript search untouched. Blocking store calls go through `memory_store.run_db` (the `memory-db` executor), never the event loop or `mlx-embed`. - **Never write a memory index directly (DEVS-145)** — every write goes through `MemoryService` to Postgres, which resets that memory's rows in `memory_index_acks` in the same transaction. The sync loop (`SESSIONFLOW_MEMORY_SYNC_INTERVAL`, plus a full reconcile at start and every 6 hours) applies what each consumer (`milvus`, `fts:`) lacks. A Milvus write re-reads the row under `pg_advisory_xact_lock` and upserts only if the embedded hash still matches; the FTS file keeps a per-record watermark so an older write cannot revive a deleted record. `change_seq` orders work but does not follow commit order, so pending work is decided by acknowledgements, not by a cursor. -- **Memory embeds share the MLX thread (DEVS-145)** — `memory_index.embed_for_index` waits on `embedding_control.EmbedThrottle` (200 ms floor) and runs on `mlx-embed`; it never uses the transcript budget, which drops batches. `memory_cli` never embeds: `rebuild` and `reindex` only reset acknowledgements, and `restore-check` asks the running server's loopback `/memory/restore-check` route to index into a temporary `memories_restore` collection. -- **Memory embedding generation (DEVS-145)** — `memory_index_meta` names the memory model and generation. `cleanup.py reembed` starts a new generation on a fresh run and resets the Milvus acknowledgements; the server re-embeds from Postgres text. A server whose local model differs refuses memory vector writes and reports it in `/health` (`memory.index.milvus_blocked`). +- **Memory embeds share the MLX thread (DEVS-145)** — `memory_index.embed_for_index` waits on `embedding_control.EmbedThrottle` (200 ms floor) and runs on `mlx-embed`, where it waits again and stamps the throttle; it never uses the transcript budget, which drops batches. `EmbeddingBudget.before_batch` honors the same throttle, so neither consumer starts inside the other's cooldown. `memory_cli` never embeds: `rebuild` and `reindex` only reset acknowledgements, and `restore-check` asks the running server's loopback `/memory/restore-check` route to index into a temporary `memories_restore` collection. +- **Memory embedding generation (DEVS-145)** — `memory_index_meta` names the memory model and generation. `cleanup.py reembed` starts a new generation on a fresh run and resets the Milvus acknowledgements; the server re-embeds from Postgres text. A server whose local model differs refuses memory vector writes and vector searches (`MemoryModelMismatch`; search answers from FTS) and reports it in `/health` (`memory.index.milvus_blocked`). - **Memory secret check (DEVS-145)** — `add_memory` rejects text with any `secret_redaction.scan_spans` hit and names the rule. Paths and code identifiers can trip it; the fix is a narrow pattern in the operator allowlist (`SESSIONFLOW_REDACT_ALLOWLIST`), never a looser rule. - **Transcript cleanup leaves memory alone (DEVS-145)** — age expiry, `cleanup_sessions`, `sanitize --drop`, `reset` and `migrate-schema` act on `sessions` and `fts.db` only; `reset` and `migrate-schema` say so. - **Memory tests (DEVS-145)** — Postgres tests need `SESSIONFLOW_TEST_PG_DSN` (one fresh schema per test through the `memory_pg` fixture) and skip with that reason when it is unset; restore-check tests also need `SESSIONFLOW_TEST_PG_RESTORE_DSN` and `pg_dump`/`pg_restore` 17. An autouse fixture clears `SESSIONFLOW_MEMORY_DSN`. `stub_rag_engine` evicts project modules first imported under the stub, so a later test imports them against the real `rag_engine`. pytest prints fixture DSNs in setup errors, so mask `password=` in any shared output. diff --git a/README.md b/README.md index d82bace..6aef582 100644 --- a/README.md +++ b/README.md @@ -112,7 +112,7 @@ SessionFlow also stores memories: notes an agent writes on purpose, such as sess `setup.sh` copies `hooks/memory-session-start.py` into the Claude Code and Codex hook folders. At session start it shows the three newest session summaries and recent turns, and falls back to a local cache when the server cannot answer. Setup registers the hook for Codex and prints the Claude Code `~/.claude/settings.json` SessionStart entry for you to add by hand. ```bash -python memory_cli.py audit # Likely duplicates, aged open items, counts +python memory_cli.py audit # Likely duplicates, aged open items, counts, unknown projects python memory_cli.py rebuild --yes # Re-index every memory from Postgres python memory_cli.py import export.json --report report.json python memory_cli.py restore-check memory.dump --target-dsn "$EMPTY_DSN" diff --git a/embedding_control.py b/embedding_control.py index e3ff14e..2448735 100644 --- a/embedding_control.py +++ b/embedding_control.py @@ -117,8 +117,8 @@ class EmbedThrottle: """Cooldown between embed batches that waits instead of denying work. Memory index embeds use it in place of the transcript budget, which drops - batches. Transcript batches stamp it too, so a memory embed never starts - inside a transcript batch's cooldown. The cooldown never goes below 200 ms. + batches. Both consumers stamp it and honor it, so neither starts inside the + other's cooldown. The cooldown never goes below 200 ms. """ def __init__( @@ -126,11 +126,13 @@ def __init__( cooldown_ms: int = MIN_COOLDOWN_MS, clock: Callable[[], float] = time.monotonic, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + blocking_sleep: Callable[[float], None] = time.sleep, ) -> None: - """Set the cooldown; ``clock`` and ``sleep`` are injectable for tests.""" + """Set the cooldown; ``clock`` and both sleeps are injectable for tests.""" self.cooldown_ms = max(int(cooldown_ms), MIN_COOLDOWN_MS) self._clock = clock self._sleep = sleep + self._blocking_sleep = blocking_sleep self._last_done = 0.0 # An asyncio.Lock binds to one event loop, so each loop gets its own. self._locks: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock]" = ( @@ -155,14 +157,28 @@ async def batch(self) -> AsyncIterator[None]: finally: self.done() + def remaining(self) -> float: + """Return the seconds left in the cooldown after the last finished batch.""" + if not self._last_done: + return 0.0 + return max(0.0, self.cooldown_ms / 1000.0 - (self._clock() - self._last_done)) + async def wait(self) -> None: """Sleep until the cooldown after the last finished batch has passed.""" - if not self._last_done: - return - remaining = self.cooldown_ms / 1000.0 - (self._clock() - self._last_done) + remaining = self.remaining() if remaining > 0: await self._sleep(remaining) + def wait_blocking(self) -> None: + """Block the calling thread until the cooldown has passed. + + Runs on the ``mlx-embed`` thread just before an embed, so work queued + behind another consumer's batch still starts a full cooldown after it. + """ + remaining = self.remaining() + if remaining > 0: + self._blocking_sleep(remaining) + def done(self) -> None: """Record that a batch finished now.""" self._last_done = self._clock() @@ -227,13 +243,15 @@ def before_batch(self, batch_size: int, estimated_chars: int = 0) -> EmbeddingDe f"Batch size {batch_size} exceeds configured limit {self.batch_size}", ) - elapsed = time.monotonic() - self.last_batch_finished_at - cooldown = self.cooldown_ms / 1000.0 - if self.last_batch_finished_at and elapsed < cooldown: + remaining = get_embed_throttle().remaining() + if self.last_batch_finished_at: + elapsed = time.monotonic() - self.last_batch_finished_at + remaining = max(remaining, self.cooldown_ms / 1000.0 - elapsed) + if remaining > 0: return EmbeddingDecision( False, "Embedding cooldown active", - retry_after_seconds=max(0.0, cooldown - elapsed), + retry_after_seconds=remaining, ) return EmbeddingDecision(True) diff --git a/memory_cli.py b/memory_cli.py index 0731bc4..1d2986e 100644 --- a/memory_cli.py +++ b/memory_cli.py @@ -115,6 +115,7 @@ async def no_query_embedding(text: str) -> list: return memory_service.MemoryService( store=_open_store(), index=_DeferredIndex(), embed_query=no_query_embedding, allowlist=rag_engine.load_allowlist(allowlist_path), + project_inventory=memory_service.project_inventory(), ) @@ -159,6 +160,10 @@ def cmd_audit(args: argparse.Namespace) -> int: if index is None or index.unavailable: _err("Similar-memory clusters are unavailable: Milvus could not be read. " "Stale items and counts are still reported.") + if report.get("inventory_checked") is False: + _err("Unknown-project check skipped: no known project root has a readable " + ".claude/project.json. Open one session in each project so its " + "SessionStart hook registers the root, then run the audit again.") if args.format == "json": print(_json(report)) return 0 diff --git a/memory_index.py b/memory_index.py index 5e80d95..821189e 100644 --- a/memory_index.py +++ b/memory_index.py @@ -38,6 +38,10 @@ "type", "status", "issue_ids", "source", "category", "created_at") +class MemoryModelMismatch(RuntimeError): + """This machine's embedding model is not the memory model, so its vectors cannot be compared.""" + + def _iso(value: Any) -> str: if isinstance(value, datetime): return value.astimezone(timezone.utc).isoformat() @@ -65,8 +69,18 @@ async def embed_for_index(texts: List[str]) -> List[List[float]]: async with batch(): return await _embed_on_mlx_thread(texts) await throttle.wait() + return await _embed_on_mlx_thread(texts) + + +def _embed_after_cooldown(texts: List[str]) -> List[List[float]]: + # Runs on mlx-embed. A transcript batch can finish between the async wait + # and this call, and the next one checks the stamp on this thread. + throttle = embedding_control.get_embed_throttle() + wait_blocking = getattr(throttle, "wait_blocking", None) + if wait_blocking is not None: + wait_blocking() try: - return await _embed_on_mlx_thread(texts) + return rag_engine.embed_texts(texts) finally: throttle.done() @@ -75,8 +89,8 @@ async def _embed_on_mlx_thread(texts: List[str]) -> List[List[float]]: loop = asyncio.get_running_loop() executor = rag_engine._embed_executor if executor is not None: - return await loop.run_in_executor(executor, rag_engine.embed_texts, texts) - return await asyncio.to_thread(rag_engine.embed_texts, texts) + return await loop.run_in_executor(executor, _embed_after_cooldown, texts) + return await asyncio.to_thread(_embed_after_cooldown, texts) class MemoryFTS: @@ -463,7 +477,14 @@ def _vector_search_sync(self, vector: List[float], filters: Dict[str, Any], async def vector_search(self, vector: List[float], *, filters: Dict[str, Any], limit: int) -> List[Dict[str, Any]]: - """Return Milvus hits (``memory_id``, ``version``, ``generation``, ``score``), best first.""" + """Return Milvus hits (``memory_id``, ``version``, ``generation``, ``score``), best first. + + Raises ``MemoryModelMismatch`` when the stored memory model differs from + this machine's: same-size vectors from another model would still match. + """ + await self._refresh_model_block() + if self.milvus_blocked: + raise MemoryModelMismatch(self.milvus_blocked) return await run_db(self._vector_search_sync, vector, filters, limit) async def fts_search(self, query: str, *, filters: Dict[str, Any], diff --git a/memory_service.py b/memory_service.py index 8271165..4c436ea 100644 --- a/memory_service.py +++ b/memory_service.py @@ -127,8 +127,8 @@ def _memory_day(row: Dict[str, Any]): return row["created_at"].astimezone(memory_store.CHICAGO).date() -def resolve_project_roots(project: str) -> List[str]: - """Return the known project roots whose ``.claude/project.json`` names ``project``. +def _known_project_slugs() -> List[tuple[str, str]]: + """Return ``(root, project)`` for each registered root with a readable ``.claude/project.json``. The roots are the ones the server has registered from transcripts, so no folder layout is assumed. The project is ``memory.project``, else ``tag``. @@ -138,16 +138,26 @@ def resolve_project_roots(project: str) -> List[str]: import file_watcher - roots = [] + pairs = [] for root in file_watcher.known_project_roots(): try: data = json.loads((Path(root) / ".claude" / "project.json").read_text(encoding="utf-8")) except (OSError, ValueError): continue slug = (data.get("memory") or {}).get("project") or str(data.get("tag") or "").lower() - if slug == project: - roots.append(root) - return roots + if slug: + pairs.append((root, slug)) + return pairs + + +def resolve_project_roots(project: str) -> List[str]: + """Return the known project roots whose ``.claude/project.json`` names ``project``.""" + return [root for root, slug in _known_project_slugs() if slug == project] + + +def project_inventory() -> set[str]: + """Return every project named by a known root's ``.claude/project.json``.""" + return {slug for _, slug in _known_project_slugs()} def _parse_time(name: str, value: Any) -> Optional[datetime]: @@ -453,6 +463,7 @@ def find(x: str) -> str: unknown = sorted(p for p in projects if self.project_inventory and p not in self.project_inventory) return {"projects": projects, "unknown_projects": unknown, + "inventory_checked": bool(self.project_inventory), "similar_threshold": self.similar_threshold, "audit_age_days": self.audit_age_days} diff --git a/tests/test_embedding_control.py b/tests/test_embedding_control.py index b356fd0..733c047 100644 --- a/tests/test_embedding_control.py +++ b/tests/test_embedding_control.py @@ -233,3 +233,71 @@ async def both(): assert len(starts) == 2 assert starts[1] - starts[0] >= 0.2 - 1e-9 + + +def test_throttle_reports_the_cooldown_left(): + from embedding_control import EmbedThrottle + + clock = [10.0] + throttle = EmbedThrottle(cooldown_ms=200, clock=lambda: clock[0]) + + assert throttle.remaining() == 0.0 + throttle.done() + clock[0] += 0.05 + assert throttle.remaining() == _approx(0.15) + clock[0] += 1.0 + assert throttle.remaining() == 0.0 + + +def test_transcript_batch_waits_out_a_memory_embeds_cooldown(monkeypatch): + import asyncio + + import embedding_control + import memory_index + import rag_engine + + throttle = embedding_control.EmbedThrottle(cooldown_ms=200) + monkeypatch.setattr(embedding_control, "get_embed_throttle", lambda: throttle) + monkeypatch.setattr(rag_engine, "_embed_executor", None) + monkeypatch.setattr(rag_engine, "embed_texts", lambda texts: [[0.0] for _ in texts]) + asyncio.run(memory_index.embed_for_index(["memory"])) + + budget = embedding_control.EmbeddingBudget(batch_size=4, cooldown_ms=200, max_turns_per_run=10) + decision = budget.before_batch(batch_size=4) + + assert decision.allowed is False + assert 0 < decision.retry_after_seconds <= 0.2 + + +def test_memory_embed_queued_behind_a_transcript_batch_waits_on_the_mlx_thread(monkeypatch): + import asyncio + + import embedding_control + import memory_index + import rag_engine + + clock = [100.0] + slept = [] + + def fake_blocking_sleep(seconds): + slept.append(seconds) + clock[0] += seconds + + throttle = embedding_control.EmbedThrottle( + cooldown_ms=200, clock=lambda: clock[0], blocking_sleep=fake_blocking_sleep) + monkeypatch.setattr(embedding_control, "get_embed_throttle", lambda: throttle) + monkeypatch.setattr(rag_engine, "_embed_executor", None) + starts = [] + + def fake_embed(texts): + starts.append(clock[0]) + return [[0.0] for _ in texts] + + monkeypatch.setattr(rag_engine, "embed_texts", fake_embed) + throttle.done() # a transcript batch finished after the memory embed was queued + + asyncio.run(memory_index._embed_on_mlx_thread(["memory"])) + + assert slept == [_approx(0.2)] + assert starts == [_approx(100.2)] + assert throttle.remaining() == _approx(0.2) diff --git a/tests/test_memory_cli.py b/tests/test_memory_cli.py index 77c09d3..ed0015a 100644 --- a/tests/test_memory_cli.py +++ b/tests/test_memory_cli.py @@ -115,6 +115,59 @@ async def audit(self): assert '"unknown_projects"' in capsys.readouterr().out +def _project_root(tmp_path, name, config): + import json + + root = tmp_path / name + (root / ".claude").mkdir(parents=True) + (root / ".claude" / "project.json").write_text(json.dumps(config), encoding="utf-8") + return str(root) + + +def test_project_inventory_reads_each_known_root(monkeypatch, tmp_path): + import file_watcher + import memory_service + + roots = [ + _project_root(tmp_path, "a", {"memory": {"project": "sessionflow"}, "tag": "ignored"}), + _project_root(tmp_path, "b", {"tag": "DevOps"}), + str(tmp_path / "no-config"), + ] + monkeypatch.setattr(file_watcher, "known_project_roots", lambda: roots) + + assert memory_service.project_inventory() == {"sessionflow", "devops"} + + +def test_cli_audit_reports_projects_missing_from_the_inventory( + cli_store, no_embedding, monkeypatch, tmp_path, capsys): + import file_watcher + + root = _project_root(tmp_path, "sf", {"memory": {"project": "sessionflow"}}) + monkeypatch.setattr(file_watcher, "known_project_roots", lambda: [root]) + monkeypatch.setattr(memory_cli, "_open_neighbour_index", lambda: None) + cli_store.add(_rec("known project")) + cli_store.add(_rec("typo project", project="sesionflow")) + + assert memory_cli.main(["audit"]) == 0 + + assert "Unknown projects: sesionflow" in capsys.readouterr().out + + +def test_cli_audit_says_when_no_inventory_is_available( + cli_store, no_embedding, monkeypatch, capsys): + import file_watcher + + monkeypatch.setattr(file_watcher, "known_project_roots", lambda: []) + monkeypatch.setattr(memory_cli, "_open_neighbour_index", lambda: None) + cli_store.add(_rec("any project", project="sesionflow")) + + assert memory_cli.main(["audit", "--format", "json"]) == 0 + + captured = capsys.readouterr() + assert '"inventory_checked": false' in captured.out + assert "Unknown-project check skipped" in captured.err + + class _StoredVectorMilvus: """A Milvus stand-in holding one stored vector and answering the threshold search.""" diff --git a/tests/test_memory_index.py b/tests/test_memory_index.py index beca3ea..40a2c2a 100644 --- a/tests/test_memory_index.py +++ b/tests/test_memory_index.py @@ -29,6 +29,7 @@ class FakeMilvus: def __init__(self): self.collections: dict[str, dict[int, dict]] = {} self.fail_next_upsert = False + self.searches: list[dict] = [] def has_collection(self, name): return name in self.collections @@ -71,6 +72,10 @@ def delete(self, collection_name, ids=None, filter=None): def query(self, collection_name, filter="", output_fields=None, limit=None, **kwargs): return [dict(r) for r in self.collections.get(collection_name, {}).values()] + def search(self, collection_name, data, **kwargs): + self.searches.append({"collection_name": collection_name, **kwargs}) + return [[]] + def rows(self): return {r["memory_id"]: r for r in self.collections.get("memories", {}).values()} @@ -347,6 +352,32 @@ def test_running_index_stops_writing_after_the_stored_model_changes(store, milvu assert second.memory_id in milvus.rows() +def test_mismatched_local_model_runs_no_vector_search(store, milvus, tmp_path): + _index(store, milvus, tmp_path, host="minitron", model_name="embeddinggemma") + other = _index(store, milvus, tmp_path, host="blue", model_name="modernbert") + + with pytest.raises(memory_index.MemoryModelMismatch, match="modernbert"): + asyncio.run(other.vector_search(_vec("query"), filters={"project": "sessionflow"}, + limit=5)) + + assert milvus.searches == [] + + +def test_vector_search_checks_the_stored_model_before_searching(store, milvus, tmp_path): + idx = _index(store, milvus, tmp_path) + assert asyncio.run(idx.vector_search(_vec("query"), filters={"project": "sessionflow"}, + limit=5)) == [] + assert len(milvus.searches) == 1 + + store.bump_generation(model_name="modernbert", vector_dim=768) + + with pytest.raises(memory_index.MemoryModelMismatch): + asyncio.run(idx.vector_search(_vec("query"), filters={"project": "sessionflow"}, + limit=5)) + assert len(milvus.searches) == 1 + assert "modernbert" in idx.health()["milvus_blocked"] + + def test_milvus_write_refuses_a_stale_model_between_ticks(store, milvus, tmp_path): idx = _index(store, milvus, tmp_path) written = store.add(_rec("raced the swap")) diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index 39e914b..d04e6be 100644 --- a/tests/test_memory_service.py +++ b/tests/test_memory_service.py @@ -319,6 +319,18 @@ def test_milvus_error_falls_back_to_keyword_search(svc, index): assert [r["memory_id"] for r in out["results"]] == [written["memory_id"]] +def test_model_mismatch_falls_back_to_keyword_search(svc, index): + import memory_index + + written = _add(svc, "keyword under another model") + index.vector_error = memory_index.MemoryModelMismatch("local model modernbert differs") + index.vector_hits = [{"memory_id": written["memory_id"], "version": 1, "generation": 1}] + index.fts_hits = [{"memory_id": written["memory_id"], "version": 1, "generation": 1}] + out = asyncio.run(svc.search("keyword", {"project": "sessionflow"}, 5)) + assert out["channels"] == ["fts"] + assert [r["memory_id"] for r in out["results"]] == [written["memory_id"]] + + def test_canonical_filters_are_rechecked(svc, index): decision = _add(svc, "a decision", type="decision", agent_id="codex", extra={"source": "wrap", "category": "infra"}) From 7d7e8d0a8f87b896756a1838c7ea1e7c04f02d6d Mon Sep 17 00:00:00 2001 From: lbruton Date: Tue, 27 Oct 2026 13:56:09 -0500 Subject: [PATCH 2/2] fix(DEVS-145): stamp a memory embed's cooldown once, on mlx-embed Codex review on #80: _embed_after_cooldown stamps the throttle on mlx-embed, and leaving EmbedThrottle.batch() stamped it again on the event loop. A transcript batch already sleeping out the first stamp could then be denied on add_turns' single retry, and the watcher advances the transcript offset anyway, so those turns were skipped. EmbedThrottle.batch takes stamp=False; embed_for_index uses it, so the only stamp is the one on mlx-embed. Tests: batch(stamp=False) leaves no stamp, and one memory embed stamps exactly once, off the event loop thread. Full suite 727 passed, 0 skipped; ruff clean. --- embedding_control.py | 8 +++--- memory_index.py | 4 ++- tests/test_embedding_control.py | 44 +++++++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/embedding_control.py b/embedding_control.py index 2448735..dca66a4 100644 --- a/embedding_control.py +++ b/embedding_control.py @@ -140,11 +140,12 @@ def __init__( ) @asynccontextmanager - async def batch(self) -> AsyncIterator[None]: + async def batch(self, stamp: bool = True) -> AsyncIterator[None]: """Hold the throttle for one batch: wait out the cooldown, run, then stamp it. Concurrent batches on the same event loop run one at a time, so each - starts a full cooldown after the previous one finished. + starts a full cooldown after the previous one finished. Pass + ``stamp=False`` when the batch stamps the throttle itself. """ loop = asyncio.get_running_loop() lock = self._locks.get(loop) @@ -155,7 +156,8 @@ async def batch(self) -> AsyncIterator[None]: try: yield finally: - self.done() + if stamp: + self.done() def remaining(self) -> float: """Return the seconds left in the cooldown after the last finished batch.""" diff --git a/memory_index.py b/memory_index.py index 821189e..f5efb1f 100644 --- a/memory_index.py +++ b/memory_index.py @@ -66,7 +66,9 @@ async def embed_for_index(texts: List[str]) -> List[List[float]]: throttle = embedding_control.get_embed_throttle() batch = getattr(throttle, "batch", None) if batch is not None: - async with batch(): + # _embed_after_cooldown stamps on mlx-embed; a second stamp here would + # push back a transcript batch already waiting out the first one. + async with batch(stamp=False): return await _embed_on_mlx_thread(texts) await throttle.wait() return await _embed_on_mlx_thread(texts) diff --git a/tests/test_embedding_control.py b/tests/test_embedding_control.py index 733c047..a4b1450 100644 --- a/tests/test_embedding_control.py +++ b/tests/test_embedding_control.py @@ -301,3 +301,47 @@ def fake_embed(texts): assert slept == [_approx(0.2)] assert starts == [_approx(100.2)] assert throttle.remaining() == _approx(0.2) + + +def test_throttle_batch_can_leave_the_stamp_to_the_caller(): + import asyncio + + from embedding_control import EmbedThrottle + + clock = [5.0] + throttle = EmbedThrottle(cooldown_ms=200, clock=lambda: clock[0]) + + async def run(): + async with throttle.batch(stamp=False): + pass + + asyncio.run(run()) + assert throttle.remaining() == 0.0 + + +def test_memory_embed_stamps_the_throttle_once_on_the_mlx_thread(monkeypatch): + # A second stamp from the event loop would move the deadline under a + # transcript batch that is already sleeping out the first one. + import asyncio + import threading + + import embedding_control + import memory_index + import rag_engine + + stamps = [] + + class _Counting(embedding_control.EmbedThrottle): + def done(self): + stamps.append(threading.current_thread().name) + super().done() + + throttle = _Counting(cooldown_ms=200) + monkeypatch.setattr(embedding_control, "get_embed_throttle", lambda: throttle) + monkeypatch.setattr(rag_engine, "_embed_executor", None) + monkeypatch.setattr(rag_engine, "embed_texts", lambda texts: [[0.0] for _ in texts]) + + asyncio.run(memory_index.embed_for_index(["memory"])) + + assert len(stamps) == 1 + assert stamps[0] != threading.main_thread().name