Fix #107: atomic pretool latch + atomic writes in SeenStore

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-

better_memory/services/context_seen.py

@@ -8,18 +8,30 @@
 degrades to "nothing seen".
 
 File format: ``context_seen_<session_id>.json`` ->
-``{"turn": int, "seen": {"<kind>:<id>": last_injected_turn},
-"pretool_fired": bool}``. ``pretool_fired`` latches PreToolUse to one real
-firing per session (see :meth:`SeenStore.pretool_fired`).
+``{"turn": int, "seen": {"<kind>:<id>": 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_<session_id>.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
+import tempfile
 from datetime import datetime
 from pathlib import Path
 
-_FILE_RE = re.compile(r"^context_seen_.+\.json$")
+# 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_.-]")
 
 
@@ -32,6 +44,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 +54,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 +62,33 @@ 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")
+            # 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
 
     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 +107,62 @@ def filter_unseen(
         return out
 
     def mark_seen(self, ids: list[tuple[str, str]]) -> None:
-        turn = int(self._data.get("turn") or 0)
+        # Re-read + merge so a concurrent writer's mark_seen entries are
+        # 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 = 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:
-            self._data["seen"][_key(kind, id_)] = turn
+            merged_seen[_key(kind, id_)] = turn
+        self._data = {"turn": 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():

tests/services/test_context_seen.py

@@ -76,3 +76,106 @@ 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_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
+        # 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}
+        # 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
+        # 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
+        # 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,
+        ) == []