fix(memory): persist the DeerMem extraction watermark across restarts and replicas (#7205)

* fix(memory): persist the DeerMem extraction watermark across restarts and replicas

The updater's per-thread extraction watermark lived only in an in-process
LRU, so a restart or a second Gateway replica on the same storage root fed
already-extracted turns to the memory LLM again.

FileMemoryStorage now keeps it in users/{user}/agents/{agent}/.metadata/
watermarks.json, keyed by thread id, written atomically under the per-user
lock without touching memory.json and capped by watermark_max_keys. Each
update feeds after the furthest of the LRU entry and the persisted one, and
writes through after an extraction or a pre-screen skip; bypass_watermark is
unchanged and sidecar IO failures fall back to the LRU. Memory clears remove
the file with the other .metadata sidecars. Content-based message identities
are stored as SHA-256 digests so the sidecar holds no conversation text.

* docs(changelog): link the memory watermark entry to #7204 and #7205

* fix(memory): never replace an unreadable watermark sidecar

set_extraction_watermark rewrites the whole watermarks.json, and the shared
sidecar reader turns an OSError into an empty mapping, so a failed read could
replace every other thread's watermark with the current one (or delete the
file). The watermark sidecar now has its own reader: a missing file reads as
empty, malformed content reads as empty (nothing to keep; the next write
repairs it), and any other OSError propagates before anything is written. The
updater already logs that and keeps the extraction successful.

* fix(memory): keep the persisted watermark when the last message has no identity

_message_identity returns None for a last message with neither an id nor
text. The updater wrote that None through, which deleted the thread's entry
from watermarks.json, losing a watermark a peer may have advanced and forcing
a full feed on every replica. The write-through now skips None: the persisted
entry still marks an extracted prefix and the furthest-of-both read keeps
using it.
This commit is contained in:
Nan Gao 2026-10-21 07:26:44 +08:00 committed by GitHub
parent cb1d782fb3
commit dc6f0d7421
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 558 additions and 38 deletions

View File

@ -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

View File

@ -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

View File

@ -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.<provider>.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.

View File

@ -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)

View File

@ -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) ──

View File

@ -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):

View File

@ -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)

View File

@ -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

View File

@ -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)

View File

@ -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"]