diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a6ce915d5..171b088c85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1265,6 +1265,18 @@ This release closes that milestone with **621 merged pull requests**. ### Fixed +- **memory:** DeerMem no longer sends conversation turns it already extracted + to the memory LLM again when a thread's next turn runs on another Gateway + instance or after a restart. The extraction watermark (the last message of + each thread already extracted) lived only in a per-process cache; it is now + also written to `users/{user}/agents/{agent}/.metadata/watermarks.json` under + the shared `storage_path` (bounded by `watermark_max_keys`, without changing + the memory revision), and each update starts after the furthest of the + process copy and the persisted one. Clearing memory deletes the file with + the agent's other `.metadata` sidecars. Messages without an id are identified + by a SHA-256 digest of their text, so the sidecar holds no conversation text. + (issue #7204, [#7205]) + - **skills:** `GET /api/skills/custom/{skill_name}/history` now answers an invalid skill name (not hyphen-case, or longer than 64 characters) with `400`, as delete and rollback already do, instead of `500` with a logged traceback on @@ -11920,3 +11932,4 @@ with **180 merged pull requests** since the first 2.0 milestone tag. [#7104]: https://github.com/bytedance/deer-flow/pull/7104 [#7123]: https://github.com/bytedance/deer-flow/pull/7123 [#7155]: https://github.com/bytedance/deer-flow/pull/7155 +[#7205]: https://github.com/bytedance/deer-flow/pull/7205 diff --git a/CHANGELOG_zh.md b/CHANGELOG_zh.md index c4a63943cd..02d697cbe4 100644 --- a/CHANGELOG_zh.md +++ b/CHANGELOG_zh.md @@ -1015,6 +1015,15 @@ ### 修复 +- **记忆：** 线程的下一轮落到另一个 Gateway 实例上或进程重启后，DeerMem 不再把已经 + 抽取过的对话轮次再次交给记忆 LLM。抽取水位（每个线程已抽取到的最后一条消息）此前只 + 存在进程内缓存中；现在还会写入共享 `storage_path` 下的 + `users/{user}/agents/{agent}/.metadata/watermarks.json`（条目数受 `watermark_max_keys` + 约束，不改变记忆 revision），每次更新从进程内副本与持久化副本中更靠后的那一个之后开始。 + 清空记忆时，该文件与该 agent 的其他 `.metadata` sidecar 一并删除。没有 id 的消息以其 + 文本的 SHA-256 摘要识别，sidecar 中不保存对话文本。 + (issue #7204, [#7205]) + - **技能：** `GET /api/skills/custom/{skill_name}/history` 收到非法技能名（不是 hyphen-case，或超过 64 个字符）时，现在与删除、回滚一样返回 `400`，不再返回 `500` 并在每次请求时记录一次带堆栈的错误日志。只有名称校验映射为 `400`：损坏的历史文件 @@ -9471,3 +9480,4 @@ DeerFlow 2.0 是围绕"超级智能体"框架的彻底重写，核心包含子 [#7104]: https://github.com/bytedance/deer-flow/pull/7104 [#7123]: https://github.com/bytedance/deer-flow/pull/7123 [#7155]: https://github.com/bytedance/deer-flow/pull/7155 +[#7205]: https://github.com/bytedance/deer-flow/pull/7205 diff --git a/backend/docs/CONFIGURATION.md b/backend/docs/CONFIGURATION.md index 6076dd8063..40804e4abd 100644 --- a/backend/docs/CONFIGURATION.md +++ b/backend/docs/CONFIGURATION.md @@ -764,7 +764,7 @@ Notes: - The declaration also drives the `agent_storage.backend: file` divergence warning, the inbound webhook dedupe warning, the `auth.local.throttle_storage: memory` warning, and the WeChat QR-login guard, which otherwise only look at the worker count. - Login lockouts are shared through the database: `auth.local.throttle_storage` (startup-only) defaults to `auto`, which keeps the per-IP failed-login counters for `POST /api/v1/auth/login/local` in the `login_throttle` table whenever `database.backend` is `sqlite` or `postgres`, so every instance enforces one `max_login_attempts` limit per client IP and a lockout on one instance holds on all of them. `memory` keeps the historical per-process counter (N instances give an attacker N × `max_login_attempts` guesses; a declared multi-instance deployment logs a warning), and `db` forces the table (it falls back to memory with a warning when the database backend is `memory`, and refuses to start when the configured database's engine is unavailable; `auto` falls back with a warning in that case). `max_login_attempts` and `lockout_seconds` stay live-read. - IM chat-to-thread bindings are shared through the database too: with `database.backend: sqlite` or `postgres` the `ChannelManager` keeps the `channel_name:chat_id[:topic_id]` → thread mapping of unbound IM conversations in the `channel_thread_bindings` table (migration `0038_channel_thread_bindings`), so a conversation created on one instance continues on the same thread when its next message lands on another. On the first start after upgrading, an existing `{base_dir}/channels/store.json` is imported once — only into an empty table, with `INSERT … ON CONFLICT DO NOTHING` so concurrently starting instances cannot duplicate a binding — and renamed `store.json.migrated`; a populated table leaves the file untouched. `memory` keeps the per-process JSON file. -- With the default DeerMem backend, the declaration also warns when the derived SQLite retrieval index sits inside the shared `storage_path` (the default `{storage_path}/.retrieval`). Set `memory.backend_config.retrieval_index_path` to an instance-local directory (relative values resolve against `storage_path`): SQLite WAL is unsupported on network filesystems, and the index is rebuilt from the Markdown facts at startup and re-synced per user scope after a peer writes. +- With the default DeerMem backend, the declaration also warns when the derived SQLite retrieval index sits inside the shared `storage_path` (the default `{storage_path}/.retrieval`). Set `memory.backend_config.retrieval_index_path` to an instance-local directory (relative values resolve against `storage_path`): SQLite WAL is unsupported on network filesystems, and the index is rebuilt from the Markdown facts at startup and re-synced per user scope after a peer writes. DeerMem's extraction watermark lives in the shared `storage_path` too (`users/{user}/agents/{agent}/.metadata/watermarks.json`), so a thread whose next turn lands on another instance, or on a restarted one, does not send already-extracted turns to the memory LLM again. - A declared multi-instance deployment that enables a feature storing credentials (today `channel_connections.enabled: true`) also needs one shared `DEER_FLOW_CREDENTIALS_KEY`; see [Credentials encryption key](#credentials-encryption-key-deer_flow_credentials_key). - IM channels: the gate refuses `channels.leader_lease.enabled: false` while a provider other than GitHub is enabled (under `channels.*`, or connectable from the web UI through `channel_connections..enabled`), because every instance would otherwise open the platform connection with the same bot credentials. With `auto` (the default) or `true` the instances share each provider's connection through the `channel_leases` table; see [IM channel leader lease](#im-channel-leader-lease-channelsleader_lease). - Restart-required: the gate runs once at startup. Restart all Gateway instances together after changing it. diff --git a/backend/packages/harness/deerflow/agents/memory/AGENTS.md b/backend/packages/harness/deerflow/agents/memory/AGENTS.md index 2d3395717f..c1e95543e5 100644 --- a/backend/packages/harness/deerflow/agents/memory/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/memory/AGENTS.md @@ -298,10 +298,12 @@ It does not require a separate signal-to-fact match. Search increments access heat only for facts it returns. Prompt injection and `get_context()` do not increment access heat. -Usage and audit sidecars live below the agent `.metadata/` directory. +Usage, audit, and per-thread extraction watermark sidecars live below the +agent `.metadata/` directory. They must not change canonical Markdown timestamps or revisions. Write audits only after canonical persistence succeeds. User delete and clear operations must remove matching sidecar data. +Updates feed after the furthest of the LRU and persisted watermark (best-effort IO). Staleness review reuses the regular updater call. It can keep, remove, or extend eligible aged facts. @@ -382,8 +384,8 @@ Keep these cross-component constraints in sync: - The pod grace period must include retrieval wait, flush time, and shutdown margin. - `retrieval_adapter` selects FTS5, a custom factory, or the empty fallback. - Eviction weights must total `1.0`. -- `watermark_max_keys: 0` makes the conversation watermark cache unbounded. -- A dropped watermark can re-extract one batch on the next turn. +- `watermark_max_keys` caps the watermark LRU and each `watermarks.json`; `0` is unbounded. +- A watermark missing from both can re-extract one batch on the next turn. #### Write-side near-duplicate fact gate (opt-in) diff --git a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/config.py b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/config.py index 2801816979..e9caa9eb4f 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/config.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/config.py @@ -320,16 +320,17 @@ class DeerMemConfig(BaseModel): "Set programmatically by the host factory (not from YAML)." ), ) - # ── Watermark cache (in-memory, bounded LRU) ───────────────────────── + # ── Watermark (in-memory LRU + persisted per-agent sidecar) ────────── watermark_max_keys: int = Field( default=4096, ge=0, description=( "Soft cap on the in-memory conversation-watermark cache (one entry " - "per distinct thread/user/agent). The cache is a bounded LRU: when " - "over capacity the least-recently-used entry is dropped, and a " - "dropped key re-extracts one batch on that thread's next turn (the " - "same as a restart). 0 = unbounded." + "per distinct thread/user/agent) and on each agent's persisted " + "watermarks.json sidecar (one entry per thread). Over capacity the " + "cache drops its least-recently-used entry and the sidecar its " + "least-recently-advanced thread; a thread missing from both " + "re-extracts one batch on its next turn. 0 = unbounded." ), ) # ── Message processing (externalized patterns / prompts) ── diff --git a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/paths.py b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/paths.py index 9e612fccdf..886707fed8 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/paths.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/paths.py @@ -149,6 +149,11 @@ def agent_eviction_audit_path(memory_path: Path, agent_name: str) -> Path: return agent_metadata_directory(memory_path, agent_name) / "eviction-audit.json" +def agent_watermark_path(memory_path: Path, agent_name: str) -> Path: + """Return the per-thread extraction watermark sidecar path for one agent.""" + return agent_metadata_directory(memory_path, agent_name) / "watermarks.json" + + def fact_file_path(memory_path: Path, fact_id: str, *, agent_name: str) -> Path: """Return the sharded Markdown path for one agent-owned fact.""" if not fact_id or not re.fullmatch(r"[A-Za-z0-9_-]+", fact_id): diff --git a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/storage.py b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/storage.py index 8dee4bfc57..aae7b06d7b 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/storage.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/storage.py @@ -39,6 +39,7 @@ from .paths import ( agent_facts_directory, agent_metadata_directory, agent_usage_path, + agent_watermark_path, fact_file_path, memory_file_path, safe_user_id, @@ -456,6 +457,11 @@ def _parse_listed_fact(path: Path) -> dict[str, Any] | None: return None +def _is_extraction_watermark(value: Any) -> bool: + """A persisted watermark is a non-empty list of strings (a message identity).""" + return isinstance(value, list) and bool(value) and all(isinstance(part, str) for part in value) + + def _fsync_parent_directory(directory: Path) -> None: """Make a completed rename durable on POSIX filesystems.""" if os.name == "nt": @@ -604,6 +610,26 @@ class MemoryStorage(abc.ABC): ) -> None: """Remove sidecar data for selected facts, or the whole scope.""" + def get_extraction_watermark( + self, + thread_id: str, + *, + agent_name: str, + user_id: str | None = None, + ) -> list[str] | None: + """Return one thread's persisted extraction watermark; None = none kept.""" + return None + + def set_extraction_watermark( + self, + thread_id: str, + watermark: list[str] | None, + *, + agent_name: str, + user_id: str | None = None, + ) -> None: + """Persist one thread's extraction watermark (None drops it); optional.""" + def close(self) -> None: """Release optional storage resources.""" @@ -652,6 +678,30 @@ class FileMemoryStorage(MemoryStorage): return expected_type() return value + @staticmethod + def _read_watermark_sidecar(path: Path) -> dict[str, Any]: + """Read the watermark sidecar, raising when the file cannot be read. + + Unlike ``_read_json_sidecar`` an ``OSError`` other than a missing file + propagates: the writer rewrites the whole file, and treating an + unreadable one as empty would replace every other thread's watermark. + Content that is not a JSON object holds no usable watermark, so it reads + as empty and the next write repairs it. + """ + try: + raw = path.read_bytes() + except FileNotFoundError: + return {} + try: + value = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + logger.warning("Ignoring malformed DeerMem watermark sidecar %s", path) + return {} + if not isinstance(value, dict): + logger.warning("Ignoring DeerMem watermark sidecar with invalid root type: %s", path) + return {} + return value + def get_fact_usage( self, *, @@ -866,6 +916,68 @@ class FileMemoryStorage(MemoryStorage): except OSError: pass + def get_extraction_watermark( + self, + thread_id: str, + *, + agent_name: str, + user_id: str | None = None, + ) -> list[str] | None: + path = self._get_memory_file_path(agent_name, user_id=user_id) + key = self._cache_key(agent_name, user_id=user_id) + with ( + self._scope_lock(key), + _process_file_lock( + path.parent / ".memory.lock", + float(getattr(self._config, "file_lock_timeout_seconds", 10)), + ), + ): + raw = self._read_watermark_sidecar(agent_watermark_path(path, agent_name)) + value = raw.get(thread_id) + return list(value) if _is_extraction_watermark(value) else None + + def set_extraction_watermark( + self, + thread_id: str, + watermark: list[str] | None, + *, + agent_name: str, + user_id: str | None = None, + ) -> None: + """Write through one thread's watermark without touching the manifest. + + The sidecar is keyed by thread id and bounded by + ``watermark_max_keys`` (0 = unbounded). The advanced thread is + re-inserted last, so the cap drops the least-recently-advanced threads; + a dropped thread re-extracts one batch on its next turn. A sidecar that + cannot be read raises before anything is written, so the other threads' + watermarks survive. + """ + path = self._get_memory_file_path(agent_name, user_id=user_id) + sidecar_path = agent_watermark_path(path, agent_name) + key = self._cache_key(agent_name, user_id=user_id) + with ( + self._scope_lock(key), + _process_file_lock( + path.parent / ".memory.lock", + float(getattr(self._config, "file_lock_timeout_seconds", 10)), + ), + ): + raw = self._read_watermark_sidecar(sidecar_path) + watermarks = {other: value for other, value in raw.items() if other != thread_id and _is_extraction_watermark(value)} + if watermark is not None: + watermarks[thread_id] = list(watermark) + cap = self._config.watermark_max_keys + if cap > 0 and len(watermarks) > cap: + watermarks = dict(list(watermarks.items())[-cap:]) + if watermarks: + _atomic_write( + sidecar_path, + json.dumps(watermarks, ensure_ascii=False, indent=2).encode("utf-8"), + ) + else: + sidecar_path.unlink(missing_ok=True) + @staticmethod def _cache_key(agent_name: str | None = None, *, user_id: str | None = None) -> tuple[str | None, str | None]: return (user_id, agent_name) diff --git a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py index 40aa8d2516..c813648ce5 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py @@ -4,6 +4,7 @@ import asyncio import atexit import concurrent.futures import copy +import hashlib import html import json import logging @@ -849,10 +850,12 @@ def _message_identity(msg: Any) -> tuple[str, ...] | None: valid when summarization removes the conversation front (an index watermark would point at the wrong message after a front removal, silently skipping un-extracted turns). Prefers the langgraph message ``id`` (unique, robust to - duplicate content); falls back to ``(type, content)`` when no id is set - (e.g. plain ``HumanMessage(content=...)`` in tests). Returns ``None`` for a - message with neither id nor extractable text -- the caller then feeds the - full list, which is safe over-extraction and never loss. + duplicate content); falls back to ``(type, sha256(content))`` when no id is + set (e.g. plain ``HumanMessage(content=...)`` in tests) -- a digest because + the identity is also persisted to the watermark sidecar, which must not hold + conversation text. Returns ``None`` for a message with neither id nor + extractable text -- the caller then feeds the full list, which is safe + over-extraction and never loss. """ mid = getattr(msg, "id", None) if isinstance(mid, str) and mid: @@ -861,7 +864,7 @@ def _message_identity(msg: Any) -> tuple[str, ...] | None: if not text: return None msg_type = getattr(msg, "type", "") or "" - return ("content", msg_type, text) + return ("content", msg_type, hashlib.sha256(text.encode("utf-8")).hexdigest()) class MemoryUpdater: @@ -887,10 +890,12 @@ class MemoryUpdater: self._prompts_dir = prompts_dir self._callbacks = callbacks # Watermark: last-extracted message identity per (thread_id, user_id, - # agent_name), held in memory so a restart re-extracts one batch. The - # cache is a bounded LRU (config.watermark_max_keys) so a long-lived - # gateway handling many threads cannot grow it without limit; a dropped - # key re-extracts one batch on that thread's next turn. + # agent_name). This process caches it in a bounded LRU + # (config.watermark_max_keys) so a long-lived gateway handling many + # threads cannot grow it without limit, and writes it through to the + # storage's per-agent sidecar so a restart or a peer replica sharing the + # storage root continues from it. A key missing from both re-extracts + # one batch on that thread's next turn. self._watermarks: OrderedDict[tuple[str | None, str | None, str | None], tuple[str, ...] | None] = OrderedDict() # ── Data access + fact CRUD (formerly module-level functions; use self._storage) ── @@ -1722,17 +1727,71 @@ class MemoryUpdater: value: tuple[str, ...] | None, ) -> None: """Store ``value`` for ``key``, evicting the least-recently-used entry - when the bounded LRU cache exceeds ``config.watermark_max_keys``. + when the bounded LRU cache exceeds ``config.watermark_max_keys``, and + write a real identity through to the storage sidecar. A dropped key is safe: the next turn for that thread finds no watermark - and re-extracts one batch (the documented restart behavior). ``0`` = - unbounded (no eviction). + and re-extracts one batch. ``0`` = unbounded (no eviction). """ self._watermarks[key] = value self._watermarks.move_to_end(key) cap = self._config.watermark_max_keys if cap > 0 and len(self._watermarks) > cap: self._watermarks.popitem(last=False) + self._store_persisted_watermark(key, value) + + @staticmethod + def _persisted_watermark_scope(key: tuple[str | None, str | None, str | None]) -> tuple[str, str | None, str] | None: + """Return ``(thread_id, user_id, agent_name)`` when ``key`` can be + persisted; the sidecar is per agent and keyed by thread id.""" + thread_id, user_id, agent_name = key + if not isinstance(thread_id, str) or not thread_id or agent_name is None: + return None + return thread_id, user_id, agent_name + + def _load_persisted_watermark(self, key: tuple[str | None, str | None, str | None]) -> tuple[str, ...] | None: + """Read ``key``'s watermark from the storage sidecar. + + ``None`` when the storage keeps none or the read fails: the sidecar is + best-effort, so the in-process cache then decides alone. + """ + scope = self._persisted_watermark_scope(key) + getter = getattr(self._storage, "get_extraction_watermark", None) + if scope is None or getter is None: + return None + thread_id, user_id, agent_name = scope + try: + value = getter(thread_id, agent_name=agent_name, user_id=user_id) + except Exception: + logger.warning("Failed to read the persisted memory watermark (thread=%s)", thread_id, exc_info=True) + return None + if isinstance(value, list | tuple) and value and all(isinstance(part, str) for part in value): + return tuple(value) + return None + + def _store_persisted_watermark( + self, + key: tuple[str | None, str | None, str | None], + value: tuple[str, ...] | None, + ) -> None: + """Write ``key``'s watermark to the storage sidecar, best-effort. + + A ``None`` identity (the last message has neither id nor text) is not + written: the persisted entry, possibly advanced by a peer, still marks + an extracted prefix, and the furthest-of-both read keeps using it. + A failure is logged and leaves the in-process cache authoritative for + this process; the extraction it follows has already been persisted, so + it must not turn into a failed update. + """ + scope = self._persisted_watermark_scope(key) + setter = getattr(self._storage, "set_extraction_watermark", None) + if value is None or scope is None or setter is None: + return + thread_id, user_id, agent_name = scope + try: + setter(thread_id, list(value), agent_name=agent_name, user_id=user_id) + except Exception: + logger.warning("Failed to persist the memory watermark (thread=%s)", thread_id, exc_info=True) def _feed_after_watermark( self, @@ -1742,19 +1801,25 @@ class MemoryUpdater: """Return the slice of ``messages`` not yet extracted. The watermark stores the identity of the last-extracted message (see - :func:`_message_identity`). If that message is still present, everything - *after* it is fed; if it is absent (front removed by summarization, or - the first-ever extraction for this key) the full list is fed. Re-feeding - is safe over-extraction -- it never skips a turn, which is the only - failure direction that would lose facts. + :func:`_message_identity`). Two copies are consulted: this process's LRU + entry and the storage sidecar (written by this process, an earlier one, + or a peer replica sharing the storage root). Each marks a prefix that + was already extracted, so the feed starts after whichever is found + furthest along ``messages``; a stale copy therefore costs nothing. If + neither is present (front removed by summarization, or the first-ever + extraction for this key) the full list is fed. Re-feeding is safe + over-extraction -- it never skips a turn, which is the only failure + direction that would lose facts. """ - last_id = self._watermark_get(watermark_key) - if last_id is None: + candidates = {self._watermark_get(watermark_key), self._load_persisted_watermark(watermark_key)} - {None} + if not candidates: return messages - for i, msg in enumerate(messages): - if _message_identity(msg) == last_id: - return messages[i + 1 :] - return messages + identities = [_message_identity(msg) for msg in messages] + start = 0 + for last_id in candidates: + if last_id in identities: + start = max(start, identities.index(last_id) + 1) + return messages[start:] def _do_update_memory_sync_impl( self, @@ -1779,10 +1844,11 @@ class MemoryUpdater: Watermark: the middleware passes the full conversation each turn, so without skipping already-extracted turns every update re-feeds old messages. The watermark stores the identity of the last-extracted - message (content/id based, in-memory only) so it stays correct when - summarization removes the conversation front; a restart loses it and - re-extracts one batch. ``bypass_watermark`` is set by the emergency - (summarization) flush path: the subset it carries is a one-shot + message (content/id based, cached in-process and persisted per agent + by the storage) so it stays correct when summarization removes the + conversation front and survives a restart or a move to a peer replica + (see :meth:`_feed_after_watermark`). ``bypass_watermark`` is set by the + emergency (summarization) flush path: the subset it carries is a one-shot "extract before removal" snapshot, so it is fed in full and does not read or advance the conversation watermark (advancing it from the subset's own length would regress the watermark and skip un-extracted diff --git a/backend/tests/blocking_io/test_deermem_watermark_sidecar.py b/backend/tests/blocking_io/test_deermem_watermark_sidecar.py new file mode 100644 index 0000000000..d00a954505 --- /dev/null +++ b/backend/tests/blocking_io/test_deermem_watermark_sidecar.py @@ -0,0 +1,53 @@ +"""Regression anchor: the DeerMem watermark sidecar stays off the event loop.""" + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from typing import Any + +import pytest +from langchain_core.messages import AIMessage, HumanMessage + +from deerflow.agents.memory.backends.deermem.deermem.config import DeerMemConfig +from deerflow.agents.memory.backends.deermem.deermem.core.storage import FileMemoryStorage +from deerflow.agents.memory.backends.deermem.deermem.core.updater import MemoryUpdater + +pytestmark = pytest.mark.asyncio + + +class _Response: + content = json.dumps({"user": {}, "history": {}, "newFacts": [], "factsToRemove": []}) + usage_metadata: dict[str, int] | None = None + + +class _LLM: + def __init__(self) -> None: + self.calls = 0 + + def invoke(self, prompt: Any, config: Any = None) -> _Response: + del prompt, config + self.calls += 1 + return _Response() + + +def _updater(root: Path, llm: _LLM) -> MemoryUpdater: + # Config validation stats storage_path; the host builds DeerMem at startup. + config = DeerMemConfig(storage_path=str(root), retrieval_adapter="") + return MemoryUpdater(config, FileMemoryStorage(config), llm) + + +async def test_async_update_reads_and_writes_the_watermark_sidecar_off_the_loop(tmp_path: Path) -> None: + llm = _LLM() + messages = [HumanMessage(content="I prefer green tea.", id="h1"), AIMessage(content="Noted.", id="a1")] + + first = await asyncio.to_thread(_updater, tmp_path, llm) + assert await first.aupdate_memory(messages, thread_id="t1", agent_name="a", user_id="u") + # A fresh updater (a restarted process) reads the sidecar and skips the batch. + restarted = await asyncio.to_thread(_updater, tmp_path, llm) + assert await restarted.aupdate_memory(messages, thread_id="t1", agent_name="a", user_id="u") + + assert llm.calls == 1 + sidecar = tmp_path / "users" / "u" / "agents" / "a" / ".metadata" / "watermarks.json" + assert await asyncio.to_thread(sidecar.is_file) diff --git a/backend/tests/test_updater_watermark.py b/backend/tests/test_updater_watermark.py index 989394ea42..a1425fdc39 100644 --- a/backend/tests/test_updater_watermark.py +++ b/backend/tests/test_updater_watermark.py @@ -1,14 +1,22 @@ -"""Tests for the in-memory watermark (skip already-extracted messages).""" +"""Tests for the extraction watermark (skip already-extracted messages). + +The updater keeps a bounded in-process LRU and, with file storage, persists the +same watermark to ``agents/{agent}/.metadata/watermarks.json`` so a restarted +process or a peer replica on the same storage root does not re-extract turns. +""" from __future__ import annotations +import errno import json +from pathlib import Path from typing import Any +import pytest from langchain_core.messages import AIMessage, HumanMessage from deerflow.agents.memory.backends.deermem.deermem.config import DeerMemConfig -from deerflow.agents.memory.backends.deermem.deermem.core.storage import MemoryStorage +from deerflow.agents.memory.backends.deermem.deermem.core.storage import FileMemoryStorage, MemoryStorage from deerflow.agents.memory.backends.deermem.deermem.core.updater import MemoryUpdater, _message_identity @@ -181,3 +189,253 @@ def test_watermark_cache_is_bounded_lru() -> None: # t1's next turn finds no watermark -> re-extracts one batch (not skipped). updater.update_memory(msgs, thread_id="t1", agent_name="a", user_id="u") assert llm.invoke_count == 4 # t1, t2, t3, then t1 re-extract + + +# ── Persisted watermark (``.metadata/watermarks.json``) ────────────────── + + +class _RecordingLLM(_FakeLLM): + """``_FakeLLM`` that also keeps each prompt, to check what was fed.""" + + def __init__(self) -> None: + super().__init__() + self.prompts: list[str] = [] + + def invoke(self, prompt: Any, config: Any = None) -> Any: + self.prompts.append(str(prompt)) + return super().invoke(prompt, config) + + +def _file_updater(root: Path, llm: _FakeLLM | None = None, **overrides: Any) -> MemoryUpdater: + """Build an updater over real file storage, the way a fresh process or a + peer replica sharing ``storage_path`` would.""" + config = _config(storage_path=str(root), retrieval_adapter="", **overrides) + return MemoryUpdater(config, FileMemoryStorage(config), llm or _RecordingLLM()) + + +def _sidecar(root: Path, *, user_id: str = "u", agent_name: str = "a") -> Path: + return root / "users" / user_id / "agents" / agent_name / ".metadata" / "watermarks.json" + + +def test_watermark_survives_a_restart(tmp_path: Path) -> None: + messages = _msgs("first") + first_llm = _RecordingLLM() + _file_updater(tmp_path, first_llm).update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + assert first_llm.invoke_count == 1 + + # A new process starts with an empty LRU but reads the persisted watermark. + llm = _RecordingLLM() + restarted = _file_updater(tmp_path, llm) + assert restarted.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") is True + assert llm.invoke_count == 0 + + messages += _msgs("second") + restarted.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + assert llm.invoke_count == 1 + assert "reply-second" in llm.prompts[0] + assert "reply-first" not in llm.prompts[0] + + +def test_peer_updaters_continue_from_each_others_watermark(tmp_path: Path) -> None: + """Two replicas on one storage root: whichever runs the next update starts + from the furthest watermark either recorded, even when its own cached entry + for the thread is stale.""" + llm_a, llm_b = _RecordingLLM(), _RecordingLLM() + replica_a = _file_updater(tmp_path, llm_a) + replica_b = _file_updater(tmp_path, llm_b) + messages = _msgs("first") + replica_a.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + + messages += _msgs("second") + replica_b.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + assert llm_b.invoke_count == 1 + assert "reply-first" not in llm_b.prompts[0] + + # Replica A's cache still points at the first turn; the sidecar is further. + messages += _msgs("third") + replica_a.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + assert llm_a.invoke_count == 2 + assert "reply-third" in llm_a.prompts[1] + assert "reply-second" not in llm_a.prompts[1] + + +def test_clearing_memory_removes_the_persisted_watermark(tmp_path: Path) -> None: + """Scoped and user-wide clears delete the watermark sidecar together with + the other ``.metadata`` sidecars of the cleared agent.""" + messages = _msgs("first") + updater = _file_updater(tmp_path) + + updater.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + assert _sidecar(tmp_path).is_file() + updater.clear_memory_data("a", user_id="u") + assert not _sidecar(tmp_path).exists() + + updater.update_memory(messages, thread_id="t2", agent_name="a", user_id="u") + assert _sidecar(tmp_path).is_file() + updater.clear_all_memory_data(user_id="u") + assert not _sidecar(tmp_path).parent.exists() + + # A fresh process finds no watermark for the cleared scope. + llm = _RecordingLLM() + _file_updater(tmp_path, llm).update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + assert llm.invoke_count == 1 + + +def test_persisted_watermarks_are_capped_by_watermark_max_keys(tmp_path: Path) -> None: + messages = _msgs("first") + updater = _file_updater(tmp_path, watermark_max_keys=2) + for thread_id in ("t1", "t2", "t3"): + updater.update_memory(messages, thread_id=thread_id, agent_name="a", user_id="u") + assert sorted(json.loads(_sidecar(tmp_path).read_text(encoding="utf-8"))) == ["t2", "t3"] + + llm = _RecordingLLM() + restarted = _file_updater(tmp_path, llm, watermark_max_keys=2) + restarted.update_memory(messages, thread_id="t3", agent_name="a", user_id="u") + assert llm.invoke_count == 0 + # The least-recently-advanced thread was dropped: it re-extracts one batch. + restarted.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + assert llm.invoke_count == 1 + + +def test_persisted_watermark_is_scoped_to_its_user_and_agent(tmp_path: Path) -> None: + messages = _msgs("first") + _file_updater(tmp_path).update_memory(messages, thread_id="t1", agent_name="a", user_id="u1") + + llm = _RecordingLLM() + other = _file_updater(tmp_path, llm) + other.update_memory(messages, thread_id="t1", agent_name="a", user_id="u2") + other.update_memory(messages, thread_id="t1", agent_name="b", user_id="u1") + assert llm.invoke_count == 2 + + other.update_memory(messages, thread_id="t1", agent_name="a", user_id="u1") + assert llm.invoke_count == 2 + + +def test_emergency_flush_neither_reads_nor_advances_the_persisted_watermark(tmp_path: Path) -> None: + msgs = _msgs("a", "b") + _file_updater(tmp_path).update_memory(msgs[:2], thread_id="t1", agent_name="a", user_id="u") + before = _sidecar(tmp_path).read_bytes() + + llm = _RecordingLLM() + restarted = _file_updater(tmp_path, llm) + restarted.update_memory(msgs[:2], thread_id="t1", agent_name="a", user_id="u", bypass_watermark=True) + assert llm.invoke_count == 1 + assert _sidecar(tmp_path).read_bytes() == before + + restarted.update_memory(msgs, thread_id="t1", agent_name="a", user_id="u") + assert llm.invoke_count == 2 + assert "reply-a" not in llm.prompts[1] + + +def test_an_unidentifiable_last_message_keeps_the_persisted_watermark(tmp_path: Path) -> None: + """A last message with neither id nor text has no identity. Writing that + through would delete a persisted watermark a peer may have advanced and + force a full feed on every replica.""" + messages = _msgs("first") + _file_updater(tmp_path).update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + before = _sidecar(tmp_path).read_bytes() + + llm = _RecordingLLM() + updater = _file_updater(tmp_path, llm) + updater.update_memory(messages + _msgs("second") + [AIMessage(content="")], thread_id="t1", agent_name="a", user_id="u") + assert llm.invoke_count == 1 + assert _sidecar(tmp_path).read_bytes() == before + + +def test_persisted_watermark_stores_no_message_text(tmp_path: Path) -> None: + """Messages without an id fall back to a content identity, which is + persisted as a digest so the sidecar never holds conversation text.""" + messages = [HumanMessage(content="My passport number is X1234567"), AIMessage(content="Noted, I will remember it.")] + _file_updater(tmp_path).update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + + raw = _sidecar(tmp_path).read_text(encoding="utf-8") + assert "X1234567" not in raw + assert "remember" not in raw + llm = _RecordingLLM() + _file_updater(tmp_path, llm).update_memory(messages, thread_id="t1", agent_name="a", user_id="u") + assert llm.invoke_count == 0 + + +def test_watermark_sidecar_failures_fall_back_to_the_process_cache(tmp_path: Path) -> None: + class _BrokenSidecarStorage(FileMemoryStorage): + def get_extraction_watermark(self, *args: Any, **kwargs: Any) -> Any: + raise OSError("sidecar unavailable") + + def set_extraction_watermark(self, *args: Any, **kwargs: Any) -> None: + raise OSError("sidecar unavailable") + + config = _config(storage_path=str(tmp_path), retrieval_adapter="") + llm = _RecordingLLM() + updater = MemoryUpdater(config, _BrokenSidecarStorage(config), llm) + messages = _msgs("first") + + assert updater.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") is True + assert updater.update_memory(messages, thread_id="t1", agent_name="a", user_id="u") is True + assert llm.invoke_count == 1 + + +def test_file_storage_round_trips_extraction_watermarks_without_a_manifest_write(tmp_path: Path) -> None: + storage = FileMemoryStorage(_config(storage_path=str(tmp_path), retrieval_adapter="")) + assert storage.get_extraction_watermark("t1", agent_name="a", user_id="u") is None + + storage.set_extraction_watermark("t1", ["id", "m1"], agent_name="a", user_id="u") + storage.set_extraction_watermark("t2", ["id", "m2"], agent_name="a", user_id="u") + assert storage.get_extraction_watermark("t1", agent_name="a", user_id="u") == ["id", "m1"] + # A watermark is not memory content: no manifest (and no revision) is written. + assert not (tmp_path / "users" / "u" / "memory.json").exists() + + storage.set_extraction_watermark("t1", None, agent_name="a", user_id="u") + storage.set_extraction_watermark("t2", None, agent_name="a", user_id="u") + assert storage.get_extraction_watermark("t1", agent_name="a", user_id="u") is None + assert not _sidecar(tmp_path).exists() + + +def test_malformed_watermark_sidecar_entries_are_ignored(tmp_path: Path) -> None: + storage = FileMemoryStorage(_config(storage_path=str(tmp_path), retrieval_adapter="")) + sidecar = _sidecar(tmp_path) + sidecar.parent.mkdir(parents=True) + sidecar.write_text("not json", encoding="utf-8") + assert storage.get_extraction_watermark("t1", agent_name="a", user_id="u") is None + + sidecar.write_text(json.dumps({"t1": "id", "t2": ["id", 2], "t3": []}), encoding="utf-8") + for thread_id in ("t1", "t2", "t3"): + assert storage.get_extraction_watermark(thread_id, agent_name="a", user_id="u") is None + + storage.set_extraction_watermark("t4", ["id", "m4"], agent_name="a", user_id="u") + assert json.loads(sidecar.read_text(encoding="utf-8")) == {"t4": ["id", "m4"]} + + +def test_unreadable_watermark_sidecar_is_never_replaced(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A read error is not an empty sidecar: every write rewrites the whole + file, so treating it as empty would drop the other threads' watermarks.""" + config = _config(storage_path=str(tmp_path), retrieval_adapter="") + storage = FileMemoryStorage(config) + storage.set_extraction_watermark("t1", ["id", "m1"], agent_name="a", user_id="u") + storage.set_extraction_watermark("t2", ["id", "m2"], agent_name="a", user_id="u") + sidecar = _sidecar(tmp_path) + before = sidecar.read_bytes() + + def _fail_on_sidecar(read: Any) -> Any: + def wrapper(self: Path, *args: Any, **kwargs: Any) -> Any: + if self.name == "watermarks.json": + raise OSError(errno.EIO, "Input/output error", str(self)) + return read(self, *args, **kwargs) + + return wrapper + + monkeypatch.setattr(Path, "read_bytes", _fail_on_sidecar(Path.read_bytes)) + monkeypatch.setattr(Path, "read_text", _fail_on_sidecar(Path.read_text)) + + with pytest.raises(OSError): + storage.set_extraction_watermark("t3", ["id", "m3"], agent_name="a", user_id="u") + with pytest.raises(OSError): + storage.set_extraction_watermark("t1", None, agent_name="a", user_id="u") + # Through the updater the extraction still succeeds; only the write-through is skipped. + llm = _RecordingLLM() + assert MemoryUpdater(config, storage, llm).update_memory(_msgs("first"), thread_id="t3", agent_name="a", user_id="u") is True + assert llm.invoke_count == 1 + + monkeypatch.undo() + assert sidecar.read_bytes() == before + assert storage.get_extraction_watermark("t1", agent_name="a", user_id="u") == ["id", "m1"] + assert storage.get_extraction_watermark("t2", agent_name="a", user_id="u") == ["id", "m2"]