From baebe540e583a37e4444b13892baa1cd936e5440 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 06:13:12 +0000 Subject: [PATCH 1/3] Fix #107: atomic pretool latch + atomic writes in SeenStore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PreToolUse latch was a check-then-set across two file operations, so N parallel hook processes at session start could all observe pretool_fired == False and all proceed with the full retrieval — four sqlite connections, four Ollama embeds, and four duplicate injections. The _save path also rewrote the whole file from a stale in-memory snapshot with a plain write_text, so process P4's mark_seen would clobber the entries P1-P3 had recorded and re-inject them next turn. Replace the latch with an atomic O_CREAT|O_EXCL claim on a sibling sentinel (context_seen_.pretool). Callers on the concurrent path now use SeenStore.try_claim_pretool_fired(), which returns True to exactly one process per session; the older pretool_fired / mark_pretool_fired pair is retained for compatibility and now routes through the same sentinel. Rewrite _save to write to a temp sibling then os.replace, matching the existing _atomic_write idiom in install_hooks.py, and rewrite bump_turn / mark_seen to re-read the on-disk snapshot before merging so a concurrent writer's entries are preserved. prune_stale's file regex now also matches the .pretool sentinel so long-lived state directories don't accumulate them. --- better_memory/hooks/contextual_inject.py | 18 +++--- better_memory/services/context_seen.py | 72 ++++++++++++++++++++---- tests/services/test_context_seen.py | 57 +++++++++++++++++++ 3 files changed, 129 insertions(+), 18 deletions(-) diff --git a/better_memory/hooks/contextual_inject.py b/better_memory/hooks/contextual_inject.py index 7e85343..31355ab 100644 --- a/better_memory/hooks/contextual_inject.py +++ b/better_memory/hooks/contextual_inject.py @@ -11,11 +11,13 @@ observability (contextual_fired_userprompt/pretool, contextual_injected, contextual_suppressed_floor, contextual_suppressed_dedup). -PreToolUse is latched to one real firing per session (SeenStore.pretool_fired -/ mark_pretool_fired): the installed matcher is unscoped (all tools), so -without the latch every tool call would re-run the full retrieval path. -Later PreToolUse events in the same session short-circuit on the state file -before any DB/embedder work. UserPromptSubmit is unaffected by the latch. +PreToolUse is latched to one real firing per session +(SeenStore.try_claim_pretool_fired — an atomic O_CREAT|O_EXCL claim on a +sentinel file, race-safe across parallel hook processes): the installed +matcher is unscoped (all tools), so without the latch every tool call +would re-run the full retrieval path. Later PreToolUse events in the same +session short-circuit on the sentinel before any DB/embedder work. +UserPromptSubmit is unaffected by the latch. """ from __future__ import annotations @@ -108,9 +110,11 @@ def main() -> None: prune_stale(state_dir, now=datetime.now(UTC)) seen = SeenStore(state_dir, session_id) if event == "PreToolUse": - if seen.pretool_fired(): + # Atomic O_CREAT|O_EXCL claim: if another parallel hook + # process already fired for this session, we return False + # and short-circuit before opening any DB / embedder. + if not seen.try_claim_pretool_fired(): raise _SkipInjection() # module-local sentinel; caught below - seen.mark_pretool_fired() seen.bump_turn() # A real local connection is opened in BOTH modes now. Agentcore # mode never stores memory CONTENT locally, but session- diff --git a/better_memory/services/context_seen.py b/better_memory/services/context_seen.py index 34f76bf..917f664 100644 --- a/better_memory/services/context_seen.py +++ b/better_memory/services/context_seen.py @@ -8,18 +8,27 @@ degrades to "nothing seen". File format: ``context_seen_.json`` -> -``{"turn": int, "seen": {":": last_injected_turn}, -"pretool_fired": bool}``. ``pretool_fired`` latches PreToolUse to one real -firing per session (see :meth:`SeenStore.pretool_fired`). +``{"turn": int, "seen": {":": last_injected_turn}}``. Writes go +through a temp file + :func:`os.replace` so a concurrent reader never sees +a partial JSON, and mutators re-read the on-disk snapshot before merging +so a second process's writes are not silently clobbered. + +The PreToolUse "one real firing per session" latch is a sibling sentinel +file ``context_seen_.pretool`` claimed atomically via +``os.open(..., O_CREAT|O_EXCL)``. Only :meth:`SeenStore.try_claim_pretool_fired` +is race-free; the older read-then-write pair (:meth:`pretool_fired` / +:meth:`mark_pretool_fired`) is retained for compatibility but callers on +the concurrent path must use the atomic claim. """ from __future__ import annotations import json +import os import re from datetime import datetime from pathlib import Path -_FILE_RE = re.compile(r"^context_seen_.+\.json$") +_FILE_RE = re.compile(r"^context_seen_.+\.(json|pretool)$") _SAFE_SESSION_RE = re.compile(r"[^A-Za-z0-9_.-]") @@ -32,6 +41,7 @@ def __init__(self, state_dir: Path, session_id: str) -> None: self._dir = state_dir safe = _SAFE_SESSION_RE.sub("_", session_id or "unknown") self._path = state_dir / f"context_seen_{safe}.json" + self._sentinel = state_dir / f"context_seen_{safe}.pretool" self._data = self._load() def _load(self) -> dict: @@ -41,7 +51,6 @@ def _load(self) -> dict: return { "turn": int(raw.get("turn") or 0), "seen": raw["seen"], - "pretool_fired": bool(raw.get("pretool_fired")), } except BaseException: # noqa: BLE001 - corrupt/missing -> empty pass @@ -50,11 +59,16 @@ def _load(self) -> dict: def _save(self) -> None: try: self._dir.mkdir(parents=True, exist_ok=True) - self._path.write_text(json.dumps(self._data), encoding="utf-8") + tmp = self._path.with_suffix(self._path.suffix + ".tmp") + tmp.write_text(json.dumps(self._data), encoding="utf-8") + os.replace(tmp, self._path) except BaseException: # noqa: BLE001 - best-effort pass def bump_turn(self) -> int: + # Re-read latest on-disk snapshot so a concurrent process's turn + # bump isn't silently overwritten by our stale copy. + self._data = self._load() self._data["turn"] = int(self._data.get("turn") or 0) + 1 self._save() return self._data["turn"] @@ -73,21 +87,57 @@ def filter_unseen( return out def mark_seen(self, ids: list[tuple[str, str]]) -> None: + # Re-read + merge so a concurrent writer's mark_seen entries are + # preserved when we write our own back. + latest = self._load() turn = int(self._data.get("turn") or 0) + merged_seen = dict(latest.get("seen") or {}) for kind, id_ in ids: - self._data["seen"][_key(kind, id_)] = turn + merged_seen[_key(kind, id_)] = turn + self._data = { + "turn": max(int(latest.get("turn") or 0), turn), + "seen": merged_seen, + } self._save() + def try_claim_pretool_fired(self) -> bool: + """Atomic check-and-set on a sentinel file. + + Returns True iff this call created the sentinel (i.e. this process + is the first firing this session); False if the sentinel already + exists or the claim could not be established. Safe to race across + processes: ``O_CREAT|O_EXCL`` guarantees at most one caller sees + the True return per session. + """ + try: + self._dir.mkdir(parents=True, exist_ok=True) + fd = os.open( + str(self._sentinel), + os.O_CREAT | os.O_EXCL | os.O_WRONLY, + 0o644, + ) + os.close(fd) + return True + except FileExistsError: + return False + except BaseException: # noqa: BLE001 - best-effort; behave as "already claimed" + return False + def pretool_fired(self) -> bool: - return bool(self._data.get("pretool_fired")) + return self._sentinel.exists() def mark_pretool_fired(self) -> None: - self._data["pretool_fired"] = True - self._save() + # Retained for compatibility; the sentinel is idempotent so + # racing callers converge on the same on-disk state, though only + # try_claim_pretool_fired() distinguishes the winner. + self.try_claim_pretool_fired() def prune_stale(state_dir: Path, *, now: datetime, max_age_days: int = 7) -> None: - """Delete context_seen files older than max_age_days. Never raises.""" + """Delete context_seen state / pretool sentinels older than max_age_days. + + Never raises. + """ try: cutoff = now.timestamp() - max_age_days * 86400 for f in state_dir.iterdir(): diff --git a/tests/services/test_context_seen.py b/tests/services/test_context_seen.py index 37bf231..3fc4425 100644 --- a/tests/services/test_context_seen.py +++ b/tests/services/test_context_seen.py @@ -76,3 +76,60 @@ def test_defaults_false_then_persists(self, tmp_path): def test_corrupt_state_means_not_fired(self, tmp_path): (tmp_path / "context_seen_sess.json").write_text("{", encoding="utf-8") assert SeenStore(tmp_path, "sess").pretool_fired() is False + + def test_try_claim_pretool_fired_only_first_caller_wins(self, tmp_path): + # #107: PreToolUse "one real firing per session" was a check-then-set + # across two file operations, so N parallel hook processes could all + # observe pretool_fired == False and all proceed. The sentinel-based + # atomic claim guarantees exactly one True return. + stores = [SeenStore(tmp_path, "sess") for _ in range(4)] + wins = [s.try_claim_pretool_fired() for s in stores] + assert wins.count(True) == 1 + assert wins.count(False) == 3 + # And every subsequent instance sees the latch as fired. + assert SeenStore(tmp_path, "sess").pretool_fired() is True + + def test_prune_stale_removes_pretool_sentinel(self, tmp_path): + import os + SeenStore(tmp_path, "sess").mark_pretool_fired() + sentinel = tmp_path / "context_seen_sess.pretool" + assert sentinel.exists() + ten_days_ago = datetime(2026, 7, 1, tzinfo=UTC).timestamp() + os.utime(sentinel, (ten_days_ago, ten_days_ago)) + prune_stale(tmp_path, now=datetime(2026, 7, 11, tzinfo=UTC)) + assert not sentinel.exists() + + +class TestConcurrentMutators: + def test_save_is_atomic_via_temp_replace(self, tmp_path): + # #107: previously plain write_text left a window where a + # concurrent reader could see a truncated file. Assert the temp + # sibling is cleaned up and the final file is complete JSON. + s = SeenStore(tmp_path, "sess") + s.bump_turn() + s.mark_seen([("reflection", "r1"), ("semantic", "m1")]) + path = tmp_path / "context_seen_sess.json" + assert path.exists() + import json as _json + loaded = _json.loads(path.read_text(encoding="utf-8")) + assert loaded["turn"] == 1 + assert loaded["seen"] == {"reflection:r1": 1, "semantic:m1": 1} + # temp sibling from atomic write must not linger on success + assert not (tmp_path / "context_seen_sess.json.tmp").exists() + + def test_mark_seen_merges_concurrent_writers_state(self, tmp_path): + # #107: two hook processes at the same turn each marked disjoint + # ids; the second _save from a stale snapshot dropped the first + # process's entries. Both writers now re-read + merge, so the + # union survives. + a = SeenStore(tmp_path, "sess") + b = SeenStore(tmp_path, "sess") + a.bump_turn() # a and b both loaded turn=0; a advances the file to 1 + b.mark_seen([("reflection", "r1")]) # b writes with its snapshot + a.mark_seen([("semantic", "m1")]) # a writes with its snapshot + # Union of both writers' entries survives on disk. + fresh = SeenStore(tmp_path, "sess") + # Neither key should be filter_unseen-visible on a fresh read. + assert fresh.filter_unseen( + [("reflection", "r1"), ("semantic", "m1")], reinject_turns=0, + ) == [] From e5931e826e3ba795cbd5a4b9953188534b0d55e2 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 06:17:18 +0000 Subject: [PATCH 2/3] Fix mark_seen stamp: use freshly-read turn, not stale local snapshot mark_seen already re-reads the on-disk snapshot so a concurrent writer's `seen` entries survive, but it was still stamping newly-marked ids with the process's local `_data["turn"]`. If another process bumped the turn between our snapshot and our mark_seen, our entry landed with a turn value below the file's true current turn, so filter_unseen's `(turn - last) > reinject_turns` gap opened one turn early and the item we just marked got re-injected prematurely. Stamp both the merged entries and the top-level turn with `max(latest_turn, local_turn)` so the entry's `last_injected_turn` can never be strictly below the file's current turn. Regression test pins the exact scenario (A bumps to 1, B concurrently bumps the file to 2, A.mark_seen must stamp with 2 so a reader with reinject_turns=1 does not re-inject). --- better_memory/services/context_seen.py | 17 +++++++++++------ tests/services/test_context_seen.py | 24 ++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/better_memory/services/context_seen.py b/better_memory/services/context_seen.py index 917f664..53067be 100644 --- a/better_memory/services/context_seen.py +++ b/better_memory/services/context_seen.py @@ -88,16 +88,21 @@ def filter_unseen( def mark_seen(self, ids: list[tuple[str, str]]) -> None: # Re-read + merge so a concurrent writer's mark_seen entries are - # preserved when we write our own back. + # preserved when we write our own back. Stamp new entries with + # the freshly-read turn (or our own if it's ahead), never our + # possibly-stale local snapshot — otherwise a concurrent bump_turn + # between our load and this call would leave the entry stamped + # with an older turn and filter_unseen would trip its + # (turn - last) > reinject_turns gate one turn early. latest = self._load() - turn = int(self._data.get("turn") or 0) + turn = max( + int(latest.get("turn") or 0), + int(self._data.get("turn") or 0), + ) merged_seen = dict(latest.get("seen") or {}) for kind, id_ in ids: merged_seen[_key(kind, id_)] = turn - self._data = { - "turn": max(int(latest.get("turn") or 0), turn), - "seen": merged_seen, - } + self._data = {"turn": turn, "seen": merged_seen} self._save() def try_claim_pretool_fired(self) -> bool: diff --git a/tests/services/test_context_seen.py b/tests/services/test_context_seen.py index 3fc4425..7140b05 100644 --- a/tests/services/test_context_seen.py +++ b/tests/services/test_context_seen.py @@ -117,6 +117,30 @@ def test_save_is_atomic_via_temp_replace(self, tmp_path): # temp sibling from atomic write must not linger on success assert not (tmp_path / "context_seen_sess.json.tmp").exists() + def test_mark_seen_stamps_with_freshly_read_turn_not_stale_local(self, tmp_path): + # #107 (BugBot follow-up): mark_seen re-reads to preserve a + # concurrent writer's seen entries, but the *stamp* for new + # entries must also come from the freshly-read turn (or the + # local turn if it's ahead), not the possibly-stale local + # snapshot alone. Otherwise a concurrent bump_turn between our + # load and our mark_seen leaves the entry stamped one turn (or + # more) below the file's true turn, and filter_unseen's + # (turn - last) > reinject_turns gate opens prematurely. + a = SeenStore(tmp_path, "sess") + a.bump_turn() # a and file both at turn=1 + # A parallel writer bumps the file forward while `a` is + # mid-retrieval (a's in-memory _data still says turn=1). + b = SeenStore(tmp_path, "sess") + b.bump_turn() # file now at turn=2 + a.mark_seen([("reflection", "r1")]) + stamped = SeenStore(tmp_path, "sess") + assert stamped._data["seen"]["reflection:r1"] == 2 + # A next-turn reader with reinject_turns=1 must NOT reinject: + # (turn - last) = (2 - 2) = 0, not > 1. + assert stamped.filter_unseen( + [("reflection", "r1")], reinject_turns=1, + ) == [] + def test_mark_seen_merges_concurrent_writers_state(self, tmp_path): # #107: two hook processes at the same turn each marked disjoint # ids; the second _save from a stale snapshot dropped the first From aa23f0965269b4a649f94361ba22ffdd5024d106 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 06:23:18 +0000 Subject: [PATCH 3/3] Fix _save: per-call-unique tmp so concurrent writers don't clobber each other MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The atomic-write helper still derived its tmp name from the session path (`self._path.with_suffix(self._path.suffix + ".tmp")`), giving every writer to the same session the same tmp filename. When two processes' _save() calls interleaved (the very scenario this PR is meant to harden against), one could truncate the other's in-flight tmp before the first os.replace ran, and the second os.replace would then fail on the now-missing shared tmp — silently, because the BaseException guard swallows FileNotFoundError. Result: one writer's data lost. Use tempfile.mkstemp(prefix=., suffix=.tmp, dir=self._dir) so each _save call gets its own tmp path, mirroring the existing pattern at runtime/session_marker.py:63. On failure the tmp is unlinked; on success os.replace consumes it. prune_stale's _FILE_RE now also matches leftover `.json..tmp` siblings so a hard-killed process between mkstemp and replace/unlink doesn't leak tmps forever. Regression test spies on os.replace and pins that two back-to-back _save calls hand it two distinct source paths. --- better_memory/services/context_seen.py | 28 ++++++++++++++++++++++---- tests/services/test_context_seen.py | 26 ++++++++++++++++++++++-- 2 files changed, 48 insertions(+), 6 deletions(-) diff --git a/better_memory/services/context_seen.py b/better_memory/services/context_seen.py index 53067be..c9f53f2 100644 --- a/better_memory/services/context_seen.py +++ b/better_memory/services/context_seen.py @@ -25,10 +25,13 @@ import json import os import re +import tempfile from datetime import datetime from pathlib import Path -_FILE_RE = re.compile(r"^context_seen_.+\.(json|pretool)$") +# Matches the state JSON, the pretool sentinel, and stray `.tmp` siblings +# left behind by a process hard-killed between mkstemp and os.replace. +_FILE_RE = re.compile(r"^context_seen_.+\.(json|pretool|json\..+\.tmp)$") _SAFE_SESSION_RE = re.compile(r"[^A-Za-z0-9_.-]") @@ -59,9 +62,26 @@ def _load(self) -> dict: def _save(self) -> None: try: self._dir.mkdir(parents=True, exist_ok=True) - tmp = self._path.with_suffix(self._path.suffix + ".tmp") - tmp.write_text(json.dumps(self._data), encoding="utf-8") - os.replace(tmp, self._path) + # Per-call-unique tmp so concurrent writers to the same + # session file don't truncate each other's in-flight temp + # and cause one writer's os.replace to silently publish the + # other's content (or fail after the other's replace already + # moved the shared tmp). Mirrors runtime/session_marker.py. + fd, tmp_name = tempfile.mkstemp( + prefix=f"{self._path.name}.", + suffix=".tmp", + dir=self._dir, + ) + tmp_path = Path(tmp_name) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(json.dumps(self._data)) + os.replace(tmp_path, self._path) + except BaseException: # noqa: BLE001 + try: + tmp_path.unlink(missing_ok=True) + except BaseException: # noqa: BLE001 + pass except BaseException: # noqa: BLE001 - best-effort pass diff --git a/tests/services/test_context_seen.py b/tests/services/test_context_seen.py index 7140b05..8df46a2 100644 --- a/tests/services/test_context_seen.py +++ b/tests/services/test_context_seen.py @@ -101,6 +101,28 @@ def test_prune_stale_removes_pretool_sentinel(self, tmp_path): class TestConcurrentMutators: + def test_save_uses_per_call_unique_tmp(self, tmp_path, monkeypatch): + # #107 (BugBot follow-up): the atomic-write path must derive a + # per-call-unique tmp so two concurrent writers don't truncate + # each other's in-flight file and lose one write on the second + # os.replace ("no such file"). Assert that back-to-back _saves + # use distinct source paths at os.replace. + import os as _os + s = SeenStore(tmp_path, "sess") + seen_srcs: list[str] = [] + orig_replace = _os.replace + def spy_replace(src, dst): # noqa: ANN001 - test spy + seen_srcs.append(str(src)) + return orig_replace(src, dst) + monkeypatch.setattr(_os, "replace", spy_replace) + s.bump_turn() + s.mark_seen([("k", "v")]) + assert len(seen_srcs) == 2 + assert seen_srcs[0] != seen_srcs[1] + # And no leftover tmp siblings after successful writes. + leftovers = [p.name for p in tmp_path.iterdir() if ".tmp" in p.name] + assert leftovers == [] + def test_save_is_atomic_via_temp_replace(self, tmp_path): # #107: previously plain write_text left a window where a # concurrent reader could see a truncated file. Assert the temp @@ -114,8 +136,8 @@ def test_save_is_atomic_via_temp_replace(self, tmp_path): loaded = _json.loads(path.read_text(encoding="utf-8")) assert loaded["turn"] == 1 assert loaded["seen"] == {"reflection:r1": 1, "semantic:m1": 1} - # temp sibling from atomic write must not linger on success - assert not (tmp_path / "context_seen_sess.json.tmp").exists() + # No temp sibling from atomic write should linger on success. + assert not any(".tmp" in p.name for p in tmp_path.iterdir()) def test_mark_seen_stamps_with_freshly_read_turn_not_stale_local(self, tmp_path): # #107 (BugBot follow-up): mark_seen re-reads to preserve a