diff --git a/AGENTS.md b/AGENTS.md index 7185e57..90c3bc5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -44,7 +44,7 @@ Semantic search over agent session transcripts. Independent project, originally - **Secret redaction guard (SESF-41)** — ingestion-time redaction hooked once in `add_turns` (covers the embedding, Milvus `document`, and FTS `content` sinks plus `add_turns_async`, and runs before `_extract_issue_ids`). Pure engine in `secret_redaction.py`: `redact(text, *, mode, allowlist) -> (redacted_text, hits)`. Config: `SESSIONFLOW_REDACT` (on/off, default **on**); `SESSIONFLOW_REDACT_MODE` (`enforce`|`report`, default **report** when unset — detect + count, store raw so operators can size the false-positive rate before enforcing); `SESSIONFLOW_REDACT_ALLOWLIST` (path to an operator regex allowlist, one pattern per line). Per-rule detection counts surface via `get_stats` under the `redaction` key. Tier-3 entropy is a length-gated Shannon scanner (detect-secrets' entropy plugins are unused — `scan_line` ignores their limit); GitHub/GitLab use custom regexes (detect-secrets reports a truncated prefix). - **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. +- **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, this host's FTS first, and skips a Milvus record another server already acknowledged. Each tick and reconcile is bounded by `SESSIONFLOW_MEMORY_TICK_TIMEOUT` (600 s): a timeout logs every asyncio task's await chain, cancels the work and lands in `/health` `memory.sync_error`, and `memory.index.stale` turns true when `last_tick` falls behind. A failed reconcile retries after 300 s, and a pass is skipped while three of the four memory-db workers are busy (`memory_store.db_calls_in_flight`), because a cancelled `run_db` await leaves its thread running. `kill -USR1 ` logs the same task dump on demand; thread profilers such as py-spy cannot see suspended coroutines. 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`, 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. diff --git a/CLAUDE.md b/CLAUDE.md index f0f29ee..467b6ad 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -45,7 +45,7 @@ Semantic search over Claude Code session transcripts. Independent project, origi - **Retroactive sanitizer (SESF-42)** — removes secrets *already* indexed (the cleanup half of SESF-35; SESF-41 is the forward guard). Orchestrator in `sanitize.py` (`scan` dry-run / `apply` redact|drop), driven by `cleanup.py sanitize` (CLI) and the `sanitize_index` MCP tool. **Dry-run is the default and writes nothing**; `--apply`/`apply` **refuse without `--yes`/`confirm`**. Reuses the SESF-41 detector via `secret_redaction.scan_spans` (per-occurrence, value-free audit spans — snippets mask every detected span in-window, including advisory Tier-3 whose forcing keyword sits outside the ±24 window). New Milvus primitives `rag_engine.upsert_document` (upsert-by-PK + **metadata-preserving** FTS delete-then-insert; `new_vector=None` keeps the stored vector on FTS-only re-converge) and `delete_by_doc_id` (`DeleteResult(deleted, fts_ok)`). **FTS rewrite/delete failure keeps a turn unfinished/retryable** — FTS healing only hydrates *missing* doc_ids, so a stale row left in place would never be re-redacted. Per-run audit JSONL at `~/.sessionflow/audit/redaction-.jsonl` + checkpoint `~/.sessionflow/sanitize_state.json` (both `0700` dir / `0600` file, value-free). Sanitization is irreversible — **redaction ≠ rotation; rotate any indexed key**. - **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. +- **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, this host's FTS first, and skips a Milvus record another server already acknowledged. Each tick and reconcile is bounded by `SESSIONFLOW_MEMORY_TICK_TIMEOUT` (600 s): a timeout logs every asyncio task's await chain, cancels the work and lands in `/health` `memory.sync_error`, and `memory.index.stale` turns true when `last_tick` falls behind. A failed reconcile retries after 300 s, and a pass is skipped while three of the four memory-db workers are busy (`memory_store.db_calls_in_flight`), because a cancelled `run_db` await leaves its thread running. `kill -USR1 ` logs the same task dump on demand; thread profilers such as py-spy cannot see suspended coroutines. 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`, 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. diff --git a/README.md b/README.md index 6aef582..61ece9e 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,7 @@ SessionFlow also stores memories: notes an agent writes on purpose, such as sess |----------|---------|-------------| | `SESSIONFLOW_MEMORY_DSN` | unset | Postgres connection string. SessionFlow reads it from the environment only. | | `SESSIONFLOW_MEMORY_SYNC_INTERVAL` | `5` | Seconds between index sync passes. | +| `SESSIONFLOW_MEMORY_TICK_TIMEOUT` | `600` | Seconds one sync pass may run. A pass that runs longer is cancelled, its waiting tasks are logged, and `/health` reports it under `memory.sync_error`. `memory.index.stale` is true when the last pass is older than this plus three intervals. | | `SESSIONFLOW_MEMORY_SIMILAR_THRESHOLD` | `0.88` | Similarity at which `add_memory` lists a memory as similar. | | `SESSIONFLOW_MEMORY_AUDIT_AGE_DAYS` | `21` | Age at which the audit lists an open item or blocker. | | `SESSIONFLOW_HOST_NAME` | short host name | Names this machine's keyword index in the acknowledgements. | diff --git a/http_server.py b/http_server.py index 14c7e84..6d4b98c 100644 --- a/http_server.py +++ b/http_server.py @@ -14,6 +14,7 @@ import concurrent.futures import contextlib from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone import json import logging import os @@ -490,6 +491,15 @@ async def _fts_heal_worker() -> None: # The full index-versus-Postgres reconcile runs on the first tick and then at # most this often. MEMORY_RECONCILE_INTERVAL = 6 * 3600 +# After a failed or timed-out reconcile, passes only tick until this many +# seconds have passed, so a hanging reconcile cannot slow every pass. +MEMORY_RECONCILE_RETRY_INTERVAL = 300 +# A tick or reconcile still running after this many seconds is cancelled and +# reported. A full batch of 256 Milvus embeds fits well inside it. +MEMORY_TICK_TIMEOUT = float(os.getenv("SESSIONFLOW_MEMORY_TICK_TIMEOUT", "600")) +# /health calls the loop stale once its last tick is older than the tick +# timeout plus this many sync intervals. +MEMORY_STALE_INTERVALS = 3 @dataclass @@ -499,14 +509,87 @@ class MemorySyncState: Attributes: index: The ``MemoryIndex`` (or any object with async ``tick`` and ``reconcile``). clock: Monotonic clock, injectable for tests. - last_reconcile: Clock value of the last full reconcile, or None. + last_reconcile: Clock value of the last successful reconcile, or None. + last_reconcile_attempt: Clock value of the last reconcile attempt, or None. last_error: Text of the last tick or reconcile failure, or None. + tick_timeout: Seconds a tick or reconcile may run before it is cancelled. """ index: object clock: object = field(default=time.monotonic) last_reconcile: float | None = None + last_reconcile_attempt: float | None = None last_error: str | None = None + tick_timeout: float = MEMORY_TICK_TIMEOUT + + +def _await_chain(task: asyncio.Task) -> list[str]: + """Return the coroutine frames a task is suspended in, outermost first, ending at its future.""" + lines = [] + obj = task.get_coro() + while obj is not None: + frame = (getattr(obj, "cr_frame", None) or getattr(obj, "gi_frame", None) + or getattr(obj, "ag_frame", None)) + if frame is None: + lines.append(f" awaiting {type(obj).__name__}: {repr(obj)[:160]}") + break + code = frame.f_code + lines.append(f" {code.co_qualname} ({os.path.basename(code.co_filename)}:{frame.f_lineno})") + obj = (getattr(obj, "cr_await", None) or getattr(obj, "gi_yieldfrom", None) + or getattr(obj, "ag_await", None)) + return lines + + +def format_async_tasks() -> str: + """Return every pending asyncio task on the running loop with the chain it is awaiting. + + Thread profilers cannot see suspended coroutines; this can. + """ + out = [] + for task in sorted(asyncio.all_tasks(), key=lambda t: t.get_name()): + out.append(f"Task {task.get_name()}:") + out.extend(_await_chain(task)) + return "\n".join(out) + + +def log_async_tasks(reason: str) -> None: + """Log every pending asyncio task and what it awaits, at WARNING.""" + logger.warning("Asyncio tasks (%s):\n%s", reason, format_async_tasks()) + + +def install_task_dump_handler() -> bool: + """Log the asyncio tasks on SIGUSR1 (`kill -USR1 `); returns False where unsupported.""" + sig = getattr(signal, "SIGUSR1", None) + if sig is None: + return False + try: + asyncio.get_running_loop().add_signal_handler(sig, log_async_tasks, "SIGUSR1") + except (NotImplementedError, RuntimeError, ValueError): + return False + return True + + +async def _run_bounded(awaitable, timeout: float, what: str): + """Await ``awaitable`` for at most ``timeout`` seconds. + + On timeout, log where every task is waiting before cancelling the work, then + raise TimeoutError. Cancelling the caller cancels the work too. + """ + task = asyncio.ensure_future(awaitable) + try: + done, _ = await asyncio.wait({task}, timeout=timeout) + except asyncio.CancelledError: + task.cancel() + raise + if done: + return task.result() + message = f"memory sync {what} timed out after {timeout:g} s" + log_async_tasks(message) + task.cancel() + await asyncio.wait({task}, timeout=5) + if task.done() and not task.cancelled(): + task.exception() + raise TimeoutError(message) def _scrub(exc: BaseException) -> str: @@ -515,21 +598,47 @@ def _scrub(exc: BaseException) -> str: return redacted +def _store_workers_saturated() -> str | None: + """Return a pause message while all but one memory-db worker are busy, else None.""" + busy = memory_store.db_calls_in_flight() + if busy < memory_store.DB_WORKERS - 1: + return None + return (f"memory sync paused: {busy} of {memory_store.DB_WORKERS} store workers " + "are busy, likely with calls left running by a timeout") + + async def _memory_sync_run_once(state: MemorySyncState, *, first_tick: bool) -> None: """Run one memory sync tick: the reconcile when due, then the pending-acknowledgement pass. - The tick runs even when the reconcile fails, and a failed reconcile stays - due for the next pass. A failure is recorded on ``state`` and never ends - the loop. + The tick runs even when the reconcile fails. A failed reconcile stays due + and runs again once ``MEMORY_RECONCILE_RETRY_INTERVAL`` has passed. Each + step is bounded by ``state.tick_timeout``. Worker capacity is checked + before the pass and again before the tick, and the work stops while all + but one memory-db worker are busy, so calls left running by a timeout + cannot take the last worker from search. A failure, timeout or pause is + recorded on ``state`` and never ends the loop. """ + paused = _store_workers_saturated() + if paused: + if state.last_error != paused: + logger.warning("%s", paused) + state.last_error = paused + return now = state.clock() errors: list[str] = [] - if first_tick or state.last_reconcile is None or ( + retry_wait = ( + not first_tick + and state.last_reconcile_attempt is not None + and state.last_reconcile_attempt != state.last_reconcile + and now - state.last_reconcile_attempt < MEMORY_RECONCILE_RETRY_INTERVAL + ) + if not retry_wait and (first_tick or state.last_reconcile is None or ( now - state.last_reconcile >= MEMORY_RECONCILE_INTERVAL - ): + )): + reconciled = False try: - reset = await state.index.reconcile() - state.last_reconcile = now + reset = await _run_bounded(state.index.reconcile(), state.tick_timeout, "reconcile") + reconciled = True if reset: logger.info("Memory reconcile reset %d acknowledgements", reset) except asyncio.CancelledError: @@ -537,8 +646,20 @@ async def _memory_sync_run_once(state: MemorySyncState, *, first_tick: bool) -> except Exception as exc: errors.append(_scrub(exc)) logger.warning("Memory reconcile failed: %s", errors[-1]) + finally: + # Stamped when the attempt ends: a timeout can take longer than the retry delay. + # On success both stamps are equal, which is how retry_wait tells success from failure. + state.last_reconcile_attempt = state.clock() + if reconciled: + state.last_reconcile = state.last_reconcile_attempt + paused = _store_workers_saturated() + if paused: + logger.warning("%s", paused) + errors.append(paused) + state.last_error = "; ".join(errors) + return try: - await state.index.tick() + await _run_bounded(state.index.tick(), state.tick_timeout, "tick") except asyncio.CancelledError: raise except Exception as exc: @@ -630,6 +751,16 @@ def jsonable(value): return value +def _sync_staleness(last_tick: str | None, sync: MemorySyncState | None) -> dict: + """Return the age of the last finished tick and whether the loop counts as stale.""" + if not last_tick: + return {"last_tick_age_s": None, "stale": False} + age = (datetime.now(timezone.utc) - datetime.fromisoformat(last_tick)).total_seconds() + timeout = sync.tick_timeout if sync is not None else MEMORY_TICK_TIMEOUT + limit = timeout + MEMORY_STALE_INTERVALS * MEMORY_SYNC_INTERVAL + return {"last_tick_age_s": round(age, 1), "stale": age > limit} + + def _memory_health_payload() -> dict: """Return memory readiness, the start-up error, acknowledgement lag and dump status.""" payload = {"ready": _memory_status["ready"], "error": _memory_status["error"], @@ -639,6 +770,7 @@ def _memory_health_payload() -> dict: sync = _memory_runtime.get("sync") if index is not None: payload["index"] = index.health() + payload["index"].update(_sync_staleness(payload["index"].get("last_tick"), sync)) if sync is not None: payload["sync_error"] = sync.last_error if store is not None and _memory_status["ready"]: @@ -1235,6 +1367,8 @@ async def lifespan(app: Starlette): heartbeat = HeartbeatThread(HEARTBEAT_FILE) heartbeat.start() print(f"[HTTP] Heartbeat thread started ({HEARTBEAT_FILE})", file=sys.stderr) + if install_task_dump_handler(): + print(f"[HTTP] kill -USR1 {os.getpid()} logs the asyncio task stacks", file=sys.stderr) # Own executor, not the loop default: shutdown must not wait on a model # load that is still running when the server is stopped mid-startup. diff --git a/memory_index.py b/memory_index.py index f5efb1f..4b6d8dd 100644 --- a/memory_index.py +++ b/memory_index.py @@ -371,6 +371,10 @@ async def _apply_milvus(self, memory_id: str) -> bool: row = await run_db(self.store.canonical, memory_id) if row is None: return False + # Another server may have indexed this version since the batch was read. + if await run_db(self.store.is_acked, MILVUS_CONSUMER, memory_id, + version=row["version"], generation=row["generation"]): + return True vector = None if row["deleted_at"] is None: vector = (await self.embed([row["text"]]))[0] @@ -545,19 +549,23 @@ def lagging(self) -> bool: # --- Loop, reconcile and rebuild ----------------------------------------- async def tick(self) -> Dict[str, int]: - """Apply every record whose acknowledgement differs, up to ``TICK_BATCH`` per consumer.""" + """Apply every record whose acknowledgement differs, up to ``TICK_BATCH`` per consumer. + + This machine's FTS file goes first: it needs no embedding, so a slow or + stuck Milvus pass cannot hold it back. + """ await self._refresh_model_block() done = {MILVUS_CONSUMER: 0, self.fts_consumer: 0} seen = 0 + batch = await run_db(self.store.pending, self.fts_consumer, TICK_BATCH) + seen += len(batch) + for row in batch: + done[self.fts_consumer] += int(await self._apply_fts(row["memory_id"])) if not self.milvus_blocked: batch = await run_db(self.store.pending, MILVUS_CONSUMER, TICK_BATCH) seen += len(batch) for row in batch: done[MILVUS_CONSUMER] += int(await self._apply_milvus(row["memory_id"])) - batch = await run_db(self.store.pending, self.fts_consumer, TICK_BATCH) - seen += len(batch) - for row in batch: - done[self.fts_consumer] += int(await self._apply_fts(row["memory_id"])) # Lagging until a tick finds nothing left to apply. self._lagging = seen > 0 self.last_tick = datetime.now(timezone.utc) diff --git a/memory_store.py b/memory_store.py index e59d82f..8f997be 100644 --- a/memory_store.py +++ b/memory_store.py @@ -16,6 +16,7 @@ import hashlib import json import os +import threading import uuid from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager @@ -210,18 +211,43 @@ def fetch_dsn() -> str: _db_executor: Optional[ThreadPoolExecutor] = None +DB_WORKERS = 4 +_in_flight = 0 +_in_flight_lock = threading.Lock() + + def db_executor() -> ThreadPoolExecutor: """Return the executor that runs every blocking store call in the server.""" global _db_executor if _db_executor is None: - _db_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="memory-db") + _db_executor = ThreadPoolExecutor(max_workers=DB_WORKERS, thread_name_prefix="memory-db") return _db_executor +def db_calls_in_flight() -> int: + """Return how many store calls are running on memory-db threads now. + + A cancelled ``run_db`` await does not stop its thread, so this counts calls + that outlived a timeout as well. + """ + return _in_flight + + +def _counted(fn: Callable[..., Any], args: tuple, kwargs: dict) -> Any: + global _in_flight + with _in_flight_lock: + _in_flight += 1 + try: + return fn(*args, **kwargs) + finally: + with _in_flight_lock: + _in_flight -= 1 + + async def run_db(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: """Run a blocking store call on the memory-db executor.""" loop = asyncio.get_running_loop() - return await loop.run_in_executor(db_executor(), lambda: fn(*args, **kwargs)) + return await loop.run_in_executor(db_executor(), _counted, fn, args, kwargs) # --- Nightly dump status ---------------------------------------------- @@ -761,6 +787,16 @@ def ack(self, consumer: str, memory_id: str, *, version: int, generation: int, ).fetchone() return row is not None + def is_acked(self, consumer: str, memory_id: str, *, version: int, generation: int) -> bool: + """Return True when ``consumer`` already holds ``version`` at ``generation`` or later.""" + with self._pool.connection() as conn: + row = conn.execute( + "SELECT 1 FROM memory_index_acks WHERE consumer = %s AND memory_id = %s " + "AND (version, generation) >= (%s, %s)", + (consumer, uuid.UUID(str(memory_id)), version, generation), + ).fetchone() + return row is not None + def reset_acks(self, memory_ids: Optional[Sequence[str]] = None, consumer: Optional[str] = None) -> int: """Delete acknowledgements so the records become pending again; returns the count. diff --git a/tests/test_fts_heal_worker.py b/tests/test_fts_heal_worker.py index 16a6b3f..631e6ab 100644 --- a/tests/test_fts_heal_worker.py +++ b/tests/test_fts_heal_worker.py @@ -262,15 +262,20 @@ async def reconcile(self): raise RuntimeError("milvus scan failed") index = _ReconcileBroken() - state = http_server.MemorySyncState(index=index, clock=lambda: 0.0) + clock = [0.0] + state = http_server.MemorySyncState(index=index, clock=lambda: clock[0]) asyncio.run(http_server._memory_sync_run_once(state, first_tick=True)) assert index.calls == ["reconcile", "tick"] assert state.last_reconcile is None assert "milvus scan failed" in state.last_error + # Inside the retry delay the pass only ticks; after it, the reconcile runs again. asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) - assert index.calls == ["reconcile", "tick", "reconcile", "tick"] + assert index.calls == ["reconcile", "tick", "tick"] + clock[0] += http_server.MEMORY_RECONCILE_RETRY_INTERVAL + asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) + assert index.calls == ["reconcile", "tick", "tick", "reconcile", "tick"] def test_memory_sync_cancellation_propagates_from_reconcile(): @@ -281,3 +286,247 @@ async def reconcile(self): state = http_server.MemorySyncState(index=_Cancelled(), clock=lambda: 0.0) with pytest.raises(asyncio.CancelledError): asyncio.run(http_server._memory_sync_run_once(state, first_tick=True)) + + +# --------------------------------------------------------------------------- +# A tick or reconcile that never returns is bounded, reported in /health, and +# leaves a log of where it was waiting. The loop goes on. +# --------------------------------------------------------------------------- + + +async def _waits_forever_inner(): + await asyncio.Event().wait() + + +class _StuckTick(_FakeMemoryIndex): + async def tick(self): + self.calls.append("tick") + await _waits_forever_inner() + + +def test_memory_sync_tick_timeout_is_recorded_and_the_loop_goes_on(): + index = _StuckTick() + state = http_server.MemorySyncState(index=index, clock=lambda: 0.0, tick_timeout=0.05) + + asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) + + assert "timed out" in state.last_error + assert "tick" in state.last_error + + +def test_memory_sync_clears_the_timeout_error_after_a_good_tick(): + index = _StuckTick() + state = http_server.MemorySyncState(index=index, clock=lambda: 0.0, tick_timeout=0.05) + asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) + state.index = _FakeMemoryIndex() + + asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) + + assert state.last_error is None + + +def test_memory_sync_retry_delay_starts_when_a_slow_reconcile_ends(): + clock = [0.0] + + class _SlowBroken(_FakeMemoryIndex): + async def reconcile(self): + self.calls.append("reconcile") + clock[0] += http_server.MEMORY_RECONCILE_RETRY_INTERVAL + 1 + raise RuntimeError("reconcile gave up") + + index = _SlowBroken() + state = http_server.MemorySyncState(index=index, clock=lambda: clock[0]) + + asyncio.run(http_server._memory_sync_run_once(state, first_tick=True)) + asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) + + assert index.calls == ["reconcile", "tick", "tick"] + + +def test_memory_sync_success_marks_attempt_and_reconcile_together(): + clock = [100.0] + + class _Slow(_FakeMemoryIndex): + async def reconcile(self): + self.calls.append("reconcile") + clock[0] += 50 + return 0 + + state = http_server.MemorySyncState(index=_Slow(), clock=lambda: clock[0]) + asyncio.run(http_server._memory_sync_run_once(state, first_tick=True)) + + assert state.last_reconcile == state.last_reconcile_attempt == 150.0 + + +def test_memory_sync_pauses_while_store_workers_are_still_busy(monkeypatch): + # A timed-out store call keeps its memory-db thread; the loop must not keep + # adding work until search has no worker left. + memory_store = importlib.import_module("memory_store") + index = _FakeMemoryIndex() + state = http_server.MemorySyncState(index=index, clock=lambda: 0.0) + monkeypatch.setattr(memory_store, "db_calls_in_flight", + lambda: memory_store.DB_WORKERS - 1) + + asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) + + assert index.calls == [] + assert "paused" in state.last_error + + +def test_memory_sync_rechecks_store_workers_before_the_tick(monkeypatch): + memory_store = importlib.import_module("memory_store") + busy = [memory_store.DB_WORKERS - 2] + + class _LeavesAWorkerBusy(_FakeMemoryIndex): + async def reconcile(self): + self.calls.append("reconcile") + busy[0] += 1 # a timed-out store call that keeps its thread + raise RuntimeError("reconcile timed out") + + index = _LeavesAWorkerBusy() + state = http_server.MemorySyncState(index=index, clock=lambda: 0.0) + monkeypatch.setattr(memory_store, "db_calls_in_flight", lambda: busy[0]) + + asyncio.run(http_server._memory_sync_run_once(state, first_tick=True)) + + assert index.calls == ["reconcile"] + assert "reconcile timed out" in state.last_error + assert "paused" in state.last_error + + +def test_memory_sync_runs_while_one_store_call_is_in_flight(monkeypatch): + memory_store = importlib.import_module("memory_store") + index = _FakeMemoryIndex() + state = http_server.MemorySyncState(index=index, clock=lambda: 0.0, last_reconcile=0.0) + monkeypatch.setattr(memory_store, "db_calls_in_flight", lambda: 1) + + asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) + + assert index.calls == ["tick"] + assert state.last_error is None + + +def test_run_db_counts_calls_running_on_the_store_threads(): + import threading + + memory_store = importlib.import_module("memory_store") + release = threading.Event() + started = threading.Event() + + def blocking(): + started.set() + release.wait(5) + return "done" + + async def run(): + task = asyncio.ensure_future(memory_store.run_db(blocking)) + await asyncio.get_running_loop().run_in_executor(None, started.wait, 5) + during = memory_store.db_calls_in_flight() + release.set() + result = await task + return during, result + + during, result = asyncio.run(run()) + assert during == 1 and result == "done" + assert memory_store.db_calls_in_flight() == 0 + + +def test_memory_sync_timeout_logs_where_the_tick_was_waiting(caplog): + import logging + + state = http_server.MemorySyncState(index=_StuckTick(), clock=lambda: 0.0, tick_timeout=0.05) + + with caplog.at_level(logging.WARNING, logger="sessionflow"): + asyncio.run(http_server._memory_sync_run_once(state, first_tick=False)) + + text = caplog.text + assert "_waits_forever_inner" in text, text + assert "Event.wait" in text or "wait" in text + + +def test_memory_sync_reconcile_timeout_still_ticks_and_stays_due(): + class _StuckReconcile(_FakeMemoryIndex): + async def reconcile(self): + self.calls.append("reconcile") + await _waits_forever_inner() + + index = _StuckReconcile() + state = http_server.MemorySyncState(index=index, clock=lambda: 0.0, tick_timeout=0.05) + + asyncio.run(http_server._memory_sync_run_once(state, first_tick=True)) + + assert index.calls == ["reconcile", "tick"] + assert state.last_reconcile is None + assert "reconcile timed out" in state.last_error + + +def test_format_async_tasks_follows_the_await_chain(): + async def middle(): + await _waits_forever_inner() + + async def run(): + task = asyncio.create_task(middle(), name="probe-task") + await asyncio.sleep(0) + text = http_server.format_async_tasks() + task.cancel() + return text + + text = asyncio.run(run()) + assert "probe-task" in text + assert "middle" in text and "_waits_forever_inner" in text + + +def test_sigusr1_logs_the_async_tasks(caplog): + import logging + import signal + + async def run(): + assert http_server.install_task_dump_handler() + try: + os.kill(os.getpid(), signal.SIGUSR1) + await asyncio.sleep(0.05) + finally: + asyncio.get_running_loop().remove_signal_handler(signal.SIGUSR1) + + with caplog.at_level(logging.WARNING, logger="sessionflow"): + asyncio.run(run()) + assert "SIGUSR1" in caplog.text + + +class _HealthIndex: + def __init__(self, last_tick): + self.last_tick = last_tick + + def health(self): + return {"host": "minitron", "consumers": ["milvus", "fts:minitron"], + "milvus_blocked": None, + "last_tick": self.last_tick.isoformat() if self.last_tick else None} + + +@pytest.mark.parametrize("age_s, stale", [(5.0, False), (10_000.0, True)]) +def test_memory_health_flags_a_stale_sync_loop(monkeypatch, age_s, stale): + from datetime import datetime, timedelta, timezone + + index = _HealthIndex(datetime.now(timezone.utc) - timedelta(seconds=age_s)) + state = http_server.MemorySyncState(index=index, tick_timeout=60.0) + monkeypatch.setitem(http_server._memory_runtime, "index", index) + monkeypatch.setitem(http_server._memory_runtime, "sync", state) + monkeypatch.setitem(http_server._memory_status, "ready", False) + + payload = http_server._memory_health_payload() + + assert payload["index"]["stale"] is stale + assert payload["index"]["last_tick_age_s"] == pytest.approx(age_s, abs=5.0) + + +def test_memory_health_is_not_stale_before_the_first_tick(monkeypatch): + index = _HealthIndex(None) + monkeypatch.setitem(http_server._memory_runtime, "index", index) + monkeypatch.setitem(http_server._memory_runtime, "sync", + http_server.MemorySyncState(index=index, tick_timeout=60.0)) + monkeypatch.setitem(http_server._memory_status, "ready", False) + + payload = http_server._memory_health_payload() + + assert payload["index"]["stale"] is False + assert payload["index"]["last_tick_age_s"] is None diff --git a/tests/test_memory_index.py b/tests/test_memory_index.py index 40a2c2a..366528a 100644 --- a/tests/test_memory_index.py +++ b/tests/test_memory_index.py @@ -352,6 +352,48 @@ def test_running_index_stops_writing_after_the_stored_model_changes(store, milvu assert second.memory_id in milvus.rows() +def test_fts_is_applied_even_while_the_milvus_embed_hangs(store, milvus, tmp_path): + async def hangs(texts): + await asyncio.Event().wait() + + idx = _index(store, milvus, tmp_path, embed=hangs) + written = store.add(_rec("fts first")) + + async def tick_until_timeout(): + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(idx.tick(), timeout=2.0) + + asyncio.run(tick_until_timeout()) + + assert store.pending("fts:minitron") == [] + assert [r["memory_id"] for r in store.pending("milvus")] == [written.memory_id] + + +def test_milvus_apply_skips_a_record_another_server_already_indexed(store, milvus, tmp_path): + other = _index(store, milvus, tmp_path, host="blue") + written = store.add(_rec("indexed by blue")) + assert asyncio.run(other.apply(written.memory_id)) == "ok" + + async def forbidden(texts): + raise AssertionError("an indexed record must not be embedded again") + + idx = _index(store, milvus, tmp_path, embed=forbidden) + + assert asyncio.run(idx._apply_milvus(written.memory_id)) is True + + +def test_milvus_apply_still_writes_a_newer_version(store, milvus, tmp_path): + other = _index(store, milvus, tmp_path, host="blue") + written = store.add(_rec("v1 by blue")) + asyncio.run(other.apply(written.memory_id)) + store.update(written.memory_id, {"text": "v2 here", + "content_hash": hashlib.sha256(b"v2 here").hexdigest()}) + idx = _index(store, milvus, tmp_path) + + assert asyncio.run(idx._apply_milvus(written.memory_id)) is True + assert milvus.rows()[written.memory_id]["version"] == 2 + + 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")