diff --git a/src/core/audit.py b/src/core/audit.py index 179d207..ebbccb8 100644 --- a/src/core/audit.py +++ b/src/core/audit.py @@ -2,15 +2,22 @@ Audit module — handles startup reconciliation of offline edits and deletes. """ +import logging from core.config import settings from core.discord_fetcher import fetch_channel_messages, fetch_message_by_id from core.state import get_last_ingested_message_id, ensure_channel_in_state, update_last_ingested_message_id -import logging logger = logging.getLogger(__name__) +def _message_id_to_int(message_id: str) -> int | None: + try: + return int(str(message_id)) + except (TypeError, ValueError): + return None + + async def run_startup_audit( channel_ids: list[str], window: int = 1000, @@ -127,15 +134,17 @@ async def run_startup_audit( limit=window, before=last_message_id, ) - + discord_index: dict[str, str] = {} + recent_lookup: dict[str, dict] = {} cursor_id_str = str(cursor_msg["id"]) discord_index[cursor_id_str] = cursor_content - + for msg in recent: docs = _docs_from_discord_messages([msg], channel_id) msg_id_str = str(msg["id"]) discord_index[msg_id_str] = docs[0].page_content if docs else "" + recent_lookup[msg_id_str] = msg if not discord_index: logger.info( "[audit] Channel %s: no messages before cursor to audit", @@ -143,6 +152,14 @@ async def run_startup_audit( ) summary[channel_id] = {"deleted": 0, "updated": 0, "skipped": "no_messages_before_cursor"} continue + + cursor_id_int = _message_id_to_int(cursor_id_str) + audited_ids = [ + parsed_id + for parsed_id in (_message_id_to_int(msg_id) for msg_id in discord_index) + if parsed_id is not None + ] + audit_lower_bound = min(audited_ids) if audited_ids else cursor_id_int try: _bootstrap_collection() except Exception as e: @@ -191,6 +208,22 @@ async def run_startup_audit( if not msg_id: continue msg_id_str = str(msg_id) + msg_id_int = _message_id_to_int(msg_id_str) + if msg_id_int is None: + logger.debug("[audit] Skipping non-numeric message_id=%s", msg_id_str) + continue + if ( + audit_lower_bound is not None + and cursor_id_int is not None + and not audit_lower_bound <= msg_id_int <= cursor_id_int + ): + logger.debug( + "[audit] Skipping message_id=%s outside audited window [%s, %s]", + msg_id_str, + audit_lower_bound, + cursor_id_int, + ) + continue logger.debug( "[audit] Checking msg_id=%s (type=%s) - in discord_index=%s", msg_id_str, type(msg_id), msg_id_str in discord_index, @@ -218,9 +251,7 @@ async def run_startup_audit( if msg_id_str == cursor_id_str: msg_dict = cursor_msg else: - msg_dict = next( - (m for m in recent if str(m["id"]) == msg_id_str), None - ) + msg_dict = recent_lookup.get(msg_id_str) if msg_dict: update_message_in_store(msg_dict, channel_id) updated += 1 @@ -248,4 +279,4 @@ async def run_startup_audit( channel_id, deleted, updated, ) - return summary \ No newline at end of file + return summary diff --git a/tests/unit/test_audit.py b/tests/unit/test_audit.py new file mode 100644 index 0000000..39bf281 --- /dev/null +++ b/tests/unit/test_audit.py @@ -0,0 +1,144 @@ +import asyncio +from types import SimpleNamespace + +from core import audit +from core import ingestion + + +def _message(message_id: str, content: str) -> dict: + return { + "id": message_id, + "content": content, + "author": {"username": "alice"}, + "timestamp": f"2026-03-10T08:0{message_id[-1]}:00+00:00", + "embeds": [], + } + + +def _point(message_id: str, channel_id: str, content: str): + return SimpleNamespace( + payload={ + "page_content": f"[alice] {content}", + "metadata": { + "channel_id": channel_id, + "message_id": message_id, + "source": "discord", + }, + } + ) + + +def test_startup_audit_does_not_delete_messages_older_than_audited_window(monkeypatch): + channel_id = "chan-1" + deleted_ids: list[str] = [] + + cursor_msg = _message("1005", "cursor") + recent_messages = [ + _message("1004", "recent-1"), + _message("1003", "recent-2"), + ] + qdrant_points = [ + _point("1002", channel_id, "older-outside-window"), + _point("1003", channel_id, "recent-2"), + _point("1004", channel_id, "recent-1"), + _point("1005", channel_id, "cursor"), + ] + + class FakeClient: + def scroll(self, **kwargs): + filter_obj = kwargs.get("scroll_filter") + conditions = getattr(filter_obj, "must", []) if filter_obj else [] + message_ids = [ + getattr(getattr(cond, "match", None), "value", None) + for cond in conditions + if getattr(cond, "key", None) == "metadata.message_id" + ] + if message_ids == ["1005"]: + return [_point("1005", channel_id, "cursor")], None + return qdrant_points, None + + def count(self, **_kwargs): + return SimpleNamespace(count=len(qdrant_points)) + + async def fake_fetch_message_by_id(**_kwargs): + return cursor_msg + + async def fake_fetch_channel_messages(**_kwargs): + return recent_messages + + monkeypatch.setattr(audit, "get_last_ingested_message_id", lambda _channel_id: "1005") + monkeypatch.setattr(audit, "fetch_message_by_id", fake_fetch_message_by_id) + monkeypatch.setattr(audit, "fetch_channel_messages", fake_fetch_channel_messages) + monkeypatch.setattr(audit, "ensure_channel_in_state", lambda _channel_id: None) + monkeypatch.setattr(ingestion, "_bootstrap_collection", lambda: None) + monkeypatch.setattr(ingestion, "get_qdrant_client", lambda: FakeClient()) + monkeypatch.setattr( + ingestion, + "delete_message_from_store", + lambda _channel_id, message_id: deleted_ids.append(message_id) or 1, + ) + monkeypatch.setattr(ingestion, "update_message_in_store", lambda *_args, **_kwargs: 0) + + summary = asyncio.run(audit.run_startup_audit([channel_id], window=2)) + + assert deleted_ids == [] + assert summary[channel_id]["deleted"] == 0 + assert summary[channel_id]["updated"] == 0 + + +def test_startup_audit_deletes_missing_message_within_audited_window(monkeypatch): + channel_id = "chan-1" + deleted_ids: list[str] = [] + + cursor_msg = _message("1005", "cursor") + recent_messages = [ + _message("1003", "recent-2"), + _message("1002", "recent-3"), + ] + qdrant_points = [ + _point("1002", channel_id, "recent-3"), + _point("1003", channel_id, "recent-2"), + _point("1004", channel_id, "deleted-offline"), + _point("1005", channel_id, "cursor"), + ] + + class FakeClient: + def scroll(self, **kwargs): + filter_obj = kwargs.get("scroll_filter") + conditions = getattr(filter_obj, "must", []) if filter_obj else [] + message_ids = [ + getattr(getattr(cond, "match", None), "value", None) + for cond in conditions + if getattr(cond, "key", None) == "metadata.message_id" + ] + if message_ids == ["1005"]: + return [_point("1005", channel_id, "cursor")], None + return qdrant_points, None + + def count(self, **_kwargs): + return SimpleNamespace(count=len(qdrant_points)) + + async def fake_fetch_message_by_id(**_kwargs): + return cursor_msg + + async def fake_fetch_channel_messages(**_kwargs): + return recent_messages + + monkeypatch.setattr(audit, "get_last_ingested_message_id", lambda _channel_id: "1005") + monkeypatch.setattr(audit, "fetch_message_by_id", fake_fetch_message_by_id) + monkeypatch.setattr(audit, "fetch_channel_messages", fake_fetch_channel_messages) + monkeypatch.setattr(audit, "ensure_channel_in_state", lambda _channel_id: None) + monkeypatch.setattr(ingestion, "_bootstrap_collection", lambda: None) + monkeypatch.setattr(ingestion, "get_qdrant_client", lambda: FakeClient()) + monkeypatch.setattr( + ingestion, + "delete_message_from_store", + lambda _channel_id, message_id: deleted_ids.append(message_id) or 1, + ) + monkeypatch.setattr(ingestion, "update_message_in_store", lambda *_args, **_kwargs: 0) + + summary = asyncio.run(audit.run_startup_audit([channel_id], window=2)) + + assert deleted_ids == ["1004"] + assert summary[channel_id]["deleted"] == 1 + assert summary[channel_id]["updated"] == 0