mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
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:
parent
cb1d782fb3
commit
dc6f0d7421
13
CHANGELOG.md
13
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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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) ──
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
53
backend/tests/blocking_io/test_deermem_watermark_sidecar.py
Normal file
53
backend/tests/blocking_io/test_deermem_watermark_sidecar.py
Normal 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)
|
||||
@ -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"]
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user