From 9bb8225079b54f36829b0d7fb94b5c00d9f8ec61 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:24:55 -0700 Subject: [PATCH] fix(memory): harden OpenViking retries and watermarks (#4552) --- README.md | 6 +- backend/AGENTS.md | 22 ++- .../memory/backends/openviking/client.py | 23 ++- .../memory/backends/openviking/config.py | 12 +- .../backends/openviking/openviking_manager.py | 59 +++++++- .../tests/test_openviking_memory_backend.py | 143 +++++++++++++++++- config.example.yaml | 3 + docs/OPENVIKING.md | 15 +- 8 files changed, 259 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index a04992226..46a3b7ce8 100644 --- a/README.md +++ b/README.md @@ -972,8 +972,10 @@ DeerFlow also includes an optional `openviking` memory backend. It connects to an independent OpenViking server over HTTP, submits completed turns through OpenViking Sessions, and recalls remote memories for prompt injection while leaving DeerMem as the default. The initial integration supports -`memory.mode: middleware`. Submitted-message watermarks prevent a failed -Session commit from duplicating already accepted messages on retry; see +`memory.mode: middleware`. Bounded submitted-message watermarks cover long and +compacted histories and prevent a failed Session commit from duplicating +already accepted messages on retry; the shared HTTP client also has explicit +connection limits and jittered retries. See [OpenViking memory backend](docs/OPENVIKING.md) for configuration and Docker startup. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 83720d2c3..4bc61f7dc 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -875,14 +875,20 @@ The cached value is reused for both the blocking (`runs.wait`) and streaming (`_ results into the shared contract. It hashes `(user_id, agent_name)` into a safe OpenViking trusted-user identity for hard scope isolation and keeps bounded message watermarks below - `{storage_path}/openviking/sessions/`. The watermark separately records - submitted and committed message IDs: once batch submission succeeds, a - later update never resubmits those messages or retries an ambiguous commit. - A future batch can commit the still-open Session together with new messages. Session - locks are weakly cached, async entrypoints offload synchronous HTTP and file - IO, and graceful shutdown rejects new work before draining all in-flight - client operations within its timeout. It does not implement DeerMem fact - CRUD/import/export and must not import the OpenViking embedded runtime. + `{storage_path}/openviking/sessions/`. The watermark combines a constant-size + ordered-prefix digest for append-only histories with a bounded recent-ID + fallback for compaction and separately records submitted and committed + progress. Once batch submission succeeds, a later update never resubmits + those messages or retries an ambiguous commit; a future batch can commit the + still-open Session together with new messages. Schema-v2 recent-ID + watermarks migrate without duplicating their anchored history. The shared + HTTP client has explicit total/keep-alive connection limits and jittered + exponential retry delays, and its configuration representation omits the API + key. Session locks are weakly cached, async entrypoints offload synchronous + HTTP and file IO, and graceful shutdown rejects new work before draining all + in-flight client operations within its timeout. It does not implement + DeerMem fact CRUD/import/export and must not import the OpenViking embedded + runtime. - `memory.mode: tool` skips `MemoryMiddleware` and registers `memory_search`, `memory_add`, `memory_update`, and `memory_delete` on the agent. The model decides when to search, add, update, or delete facts; this is opt-in/experimental and should not be described as better than middleware mode without eval evidence. - Both modes share `FileMemoryStorage`, per-user/per-agent isolation, manual CRUD primitives, and the updater backend. Injection is mode-aware: middleware mode injects global `user`/`history` summaries plus the selected agent's facts, while tool mode injects only the global summaries and leaves every agent fact behind `memory_search` to avoid duplicating automatically injected and retrieval-returned context. `memory.injection_enabled: false` suppresses the complete block in either mode. - Middleware mode queue debounces (30s default), batches updates, and commits global summaries plus the selected/default agent's fact delta through a user-level lock, optimistic user-memory revisions, per-fact revisions, and a recoverable target-file journal. Only explicitly marked point operations may rebase a stale shared revision, and only while every addressed fact still satisfies its original absent/revision precondition. Snapshot-derived clear/trim/consolidation operations instead reload the complete document and recompute their intent on a manifest conflict, with a bounded retry. Typed manifest/fact conflict subclasses keep that decision independent of exception text, and same-ID creates and stale same-fact writes fail. Scope-lock objects are weakly cached so inactive users do not grow a process-lifetime map. Cache validation does not scale with the fact-file count: its token combines the shared JSON's `(mtime_ns, size, revision)`, so the persisted revision invalidates stale caches even when a coarse-mtime filesystem reports identical metadata for same-size writes; direct out-of-band Markdown edits require `reload()`. Atomic replacement also syncs the parent directory on POSIX so the rename is durable. DeerMem translates private storage conflict/corruption exceptions to the backend-neutral MemoryManager contract; the Gateway maps them to HTTP 409 and a stable HTTP 500 response respectively. A normal default-manager read automatically migrates legacy facts from the global JSON into `__default__`; it also adopts the earlier implicit `lead-agent` fact bucket only when that directory has no custom-agent `config.yaml`, and rejects unexpected files instead of deleting them. The v1-to-v2 migration is one-way for the running application: operators must stop DeerFlow and snapshot the configured storage root before upgrade. Before any destructive v2 write, every migrated JSON source is durably retained as `{manifest_filename}.v1.bak`; a missing-write or mismatched existing backup aborts without modifying v1 data. Legacy per-agent JSON is deleted only after its non-empty summaries are safely adopted or confirmed identical; summary conflicts keep the source file and fail loudly. diff --git a/backend/packages/harness/deerflow/agents/memory/backends/openviking/client.py b/backend/packages/harness/deerflow/agents/memory/backends/openviking/client.py index 894959f66..46fdcfdbf 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/openviking/client.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/openviking/client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import random import time from typing import Any @@ -48,7 +49,16 @@ class OpenVikingHttpClient: write=config.write_timeout_seconds, pool=config.pool_timeout_seconds, ) - self._client = httpx.Client(base_url=config.base_url, timeout=timeout, transport=transport) + limits = httpx.Limits( + max_connections=config.max_connections, + max_keepalive_connections=config.max_keepalive_connections, + ) + self._client = httpx.Client( + base_url=config.base_url, + timeout=timeout, + limits=limits, + transport=transport, + ) def close(self) -> None: self._client.close() @@ -207,16 +217,16 @@ class OpenVikingHttpClient: response = self._client.request(method, path, headers=self._headers(identity), **kwargs) except httpx.TimeoutException as exc: if attempt + 1 < attempts: - time.sleep(0.05 * (2**attempt)) + time.sleep(_retry_delay(attempt)) continue raise OpenVikingTimeoutError(operation, "OpenViking request timed out") from exc except httpx.TransportError as exc: if attempt + 1 < attempts: - time.sleep(0.05 * (2**attempt)) + time.sleep(_retry_delay(attempt)) continue raise OpenVikingUnavailableError(operation, "OpenViking is unavailable") from exc if response.status_code in {429, 502, 503, 504} and attempt + 1 < attempts: - time.sleep(0.05 * (2**attempt)) + time.sleep(_retry_delay(attempt)) continue if response.status_code >= 400: try: @@ -230,3 +240,8 @@ class OpenVikingHttpClient: raise error_type(operation, message, status_code=response.status_code, code=code) return response raise OpenVikingUnavailableError(operation, "OpenViking request failed") + + +def _retry_delay(attempt: int) -> float: + base_delay = 0.05 * (2**attempt) + return base_delay + random.uniform(0.0, base_delay) diff --git a/backend/packages/harness/deerflow/agents/memory/backends/openviking/config.py b/backend/packages/harness/deerflow/agents/memory/backends/openviking/config.py index 51aa630c4..3225ad6e7 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/openviking/config.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/openviking/config.py @@ -3,7 +3,7 @@ from __future__ import annotations import os -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Literal from urllib.parse import urlparse @@ -20,11 +20,13 @@ class OpenVikingConfig: storage_path: str auth_mode: Literal["trusted", "dev"] account: str - api_key: str | None + api_key: str | None = field(repr=False) connect_timeout_seconds: float read_timeout_seconds: float write_timeout_seconds: float pool_timeout_seconds: float + max_connections: int + max_keepalive_connections: int max_retries: int search_top_k: int score_threshold: float | None @@ -56,6 +58,8 @@ class OpenVikingConfig: read_timeout_seconds=float(cfg.pop("read_timeout_seconds", 10.0)), write_timeout_seconds=float(cfg.pop("write_timeout_seconds", 10.0)), pool_timeout_seconds=float(cfg.pop("pool_timeout_seconds", 2.0)), + max_connections=int(cfg.pop("max_connections", 100)), + max_keepalive_connections=int(cfg.pop("max_keepalive_connections", 20)), max_retries=int(cfg.pop("max_retries", 1)), search_top_k=int(retrieval.pop("top_k", 8)), score_threshold=_optional_float(retrieval.pop("score_threshold", None)), @@ -95,6 +99,10 @@ class OpenVikingConfig: for field_name in ("connect_timeout_seconds", "read_timeout_seconds", "write_timeout_seconds", "pool_timeout_seconds"): if getattr(self, field_name) <= 0: raise ValueError(f"OpenViking {field_name} must be > 0") + if not 1 <= self.max_connections <= 1000: + raise ValueError("OpenViking max_connections must be between 1 and 1000") + if not 0 <= self.max_keepalive_connections <= self.max_connections: + raise ValueError("OpenViking max_keepalive_connections must be between 0 and max_connections") if not 0 <= self.max_retries <= 5: raise ValueError("OpenViking max_retries must be between 0 and 5") if not 1 <= self.search_top_k <= 100: diff --git a/backend/packages/harness/deerflow/agents/memory/backends/openviking/openviking_manager.py b/backend/packages/harness/deerflow/agents/memory/backends/openviking/openviking_manager.py index e1997d80e..0bfeca100 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/openviking/openviking_manager.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/openviking/openviking_manager.py @@ -258,10 +258,22 @@ class OpenVikingMemoryManager(MemoryManager): state = self._load_state(session_id) submitted_ids = list(state.get("submitted_message_ids", state.get("seen_message_ids", []))) committed_ids = list(state.get("committed_message_ids", state.get("seen_message_ids", []))) - submitted = set(submitted_ids) converted = _convert_messages(messages, self._should_keep_hidden_message) - pending = [message for message in converted if message.message_id not in submitted] + prefix_count = _matching_submitted_prefix_count(state, submitted_ids, converted) + if prefix_count is not None: + pending = converted[prefix_count:] + else: + submitted = set(submitted_ids) + pending = [message for message in converted if message.message_id not in submitted] if not pending: + if prefix_count is None and converted: + state = { + **state, + "schema_version": 3, + "submitted_prefix_count": len(converted), + "submitted_prefix_digest": _message_sequence_digest(converted), + } + self._save_state(session_id, state) return try: self._client.ensure_session(identity, session_id) @@ -278,10 +290,14 @@ class OpenVikingMemoryManager(MemoryManager): return submitted_ids.extend(message.message_id for message in pending) state = { - "schema_version": 2, + "schema_version": 3, "session_id": session_id, "submitted_message_ids": submitted_ids[-self._config.max_seen_message_ids :], "committed_message_ids": committed_ids[-self._config.max_seen_message_ids :], + "submitted_prefix_count": len(converted), + "submitted_prefix_digest": _message_sequence_digest(converted), + "committed_prefix_count": state.get("committed_prefix_count"), + "committed_prefix_digest": state.get("committed_prefix_digest"), "last_commit_task_id": state.get("last_commit_task_id"), "last_archive_uri": state.get("last_archive_uri"), } @@ -301,8 +317,10 @@ class OpenVikingMemoryManager(MemoryManager): state = { **state, - "schema_version": 2, + "schema_version": 3, "committed_message_ids": state["submitted_message_ids"], + "committed_prefix_count": state["submitted_prefix_count"], + "committed_prefix_digest": state["submitted_prefix_digest"], "last_commit_task_id": commit.task_id, "last_archive_uri": commit.archive_uri, } @@ -399,6 +417,39 @@ def _session_id(identity: OpenVikingIdentity, thread_id: str) -> str: return f"df_{digest[:48]}" +def _matching_submitted_prefix_count( + state: dict[str, Any], + submitted_ids: list[str], + messages: list[OpenVikingMessage], +) -> int | None: + count = state.get("submitted_prefix_count") + digest = state.get("submitted_prefix_digest") + if isinstance(count, int) and 0 <= count <= len(messages) and isinstance(digest, str): + if _message_sequence_digest(messages[:count]) == digest: + return count + return None + + # Schema v2 only retained a recent suffix of submitted IDs. When that + # suffix still appears intact, it safely anchors the append-only prefix and + # avoids a one-time duplicate submission during migration to schema v3. + if submitted_ids and len(submitted_ids) <= len(messages): + message_ids = [message.message_id for message in messages] + width = len(submitted_ids) + for start in range(len(message_ids) - width, -1, -1): + if message_ids[start : start + width] == submitted_ids: + return start + width + return None + + +def _message_sequence_digest(messages: list[OpenVikingMessage]) -> str: + digest = hashlib.sha256() + for message in messages: + encoded = message.message_id.encode() + digest.update(len(encoded).to_bytes(8, "big")) + digest.update(encoded) + return digest.hexdigest() + + def _convert_messages( messages: list[Any], should_keep_hidden_message: Callable[[Any], bool] | None, diff --git a/backend/tests/test_openviking_memory_backend.py b/backend/tests/test_openviking_memory_backend.py index 329a77f1e..0cf47820e 100644 --- a/backend/tests/test_openviking_memory_backend.py +++ b/backend/tests/test_openviking_memory_backend.py @@ -58,6 +58,38 @@ def test_config_rejects_dev_auth_without_explicit_opt_in(tmp_path: Path) -> None OpenVikingConfig.from_backend_config(_backend_config(tmp_path, auth_mode="dev")) +def test_config_repr_does_not_expose_api_key(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENVIKING_API_KEY", "super-secret-api-key") + + config = OpenVikingConfig.from_backend_config(_backend_config(tmp_path)) + + assert "super-secret-api-key" not in repr(config) + + +def test_config_parses_and_validates_connection_limits(tmp_path: Path) -> None: + config = OpenVikingConfig.from_backend_config( + _backend_config( + tmp_path, + max_connections=48, + max_keepalive_connections=12, + ) + ) + + assert config.max_connections == 48 + assert config.max_keepalive_connections == 12 + + with pytest.raises(ValueError, match="max_connections"): + OpenVikingConfig.from_backend_config(_backend_config(tmp_path, max_connections=0)) + with pytest.raises(ValueError, match="max_keepalive_connections"): + OpenVikingConfig.from_backend_config( + _backend_config( + tmp_path, + max_connections=10, + max_keepalive_connections=11, + ) + ) + + def test_backend_is_discovered_by_registered_name() -> None: reset_memory_manager() assert _scan_backends()["openviking"] is OpenVikingMemoryManager @@ -122,6 +154,55 @@ def test_http_client_maps_authentication_error(tmp_path: Path) -> None: assert exc_info.value.code == "UNAUTHENTICATED" +def test_http_client_configures_connection_limits(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + class _RecordingClient: + def __init__(self, **kwargs: Any) -> None: + captured.update(kwargs) + + monkeypatch.setattr(httpx, "Client", _RecordingClient) + config = OpenVikingConfig.from_backend_config( + _backend_config( + tmp_path, + max_connections=32, + max_keepalive_connections=8, + ) + ) + + OpenVikingHttpClient(config) + + limits = captured["limits"] + assert limits.max_connections == 32 + assert limits.max_keepalive_connections == 8 + + +def test_http_client_adds_jitter_to_retry_delay(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + attempts = 0 + jitter_ranges: list[tuple[float, float]] = [] + sleeps: list[float] = [] + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + return httpx.Response(503) + return httpx.Response(200, json={"status": "ok"}) + + def fake_uniform(lower: float, upper: float) -> float: + jitter_ranges.append((lower, upper)) + return 0.02 + + monkeypatch.setattr("random.uniform", fake_uniform) + monkeypatch.setattr("time.sleep", sleeps.append) + config = OpenVikingConfig.from_backend_config(_backend_config(tmp_path, max_retries=1)) + client = OpenVikingHttpClient(config, transport=httpx.MockTransport(handler)) + + assert client.health() is True + assert jitter_ranges == [(0.0, 0.05)] + assert sleeps == [pytest.approx(0.07)] + + class _FakeClient: def __init__(self) -> None: self.ensured: list[tuple[OpenVikingIdentity, str]] = [] @@ -180,8 +261,8 @@ class _FakeClient: self.closed = True -def _manager(tmp_path: Path) -> tuple[OpenVikingMemoryManager, _FakeClient]: - manager = OpenVikingMemoryManager.from_config(_backend_config(tmp_path)) +def _manager(tmp_path: Path, **overrides: Any) -> tuple[OpenVikingMemoryManager, _FakeClient]: + manager = OpenVikingMemoryManager.from_config(_backend_config(tmp_path, **overrides)) fake = _FakeClient() manager._client = fake # type: ignore[assignment] return manager, fake @@ -250,6 +331,64 @@ def test_manager_does_not_resubmit_messages_after_failed_commit(tmp_path: Path) assert recovered_state["last_commit_task_id"] == "task-1" +def test_manager_does_not_resubmit_history_beyond_recent_id_window(tmp_path: Path) -> None: + manager, client = _manager(tmp_path, max_seen_message_ids=16) + messages = [HumanMessage(f"message {index}", id=f"h{index}") for index in range(20)] + + manager.add("thread-1", messages, user_id="alice", agent_name="research") + manager.add("thread-1", messages, user_id="alice", agent_name="research") + + assert len(client.added) == 1 + + messages.append(HumanMessage("message 20", id="h20")) + manager.add("thread-1", messages, user_id="alice", agent_name="research") + + assert len(client.added) == 2 + assert [message.message_id for message in client.added[1][2]] == ["df_h20"] + watermark = next((tmp_path / "openviking" / "sessions").glob("*.json")) + state = json.loads(watermark.read_text(encoding="utf-8")) + assert state["schema_version"] == 3 + assert state["submitted_prefix_count"] == 21 + assert len(state["submitted_message_ids"]) == 16 + + +def test_manager_rebases_watermark_after_history_compaction(tmp_path: Path) -> None: + manager, client = _manager(tmp_path, max_seen_message_ids=16) + messages = [HumanMessage(f"message {index}", id=f"h{index}") for index in range(20)] + manager.add("thread-1", messages, user_id="alice", agent_name="research") + + compacted = [*messages[-8:], HumanMessage("message 20", id="h20")] + manager.add("thread-1", compacted, user_id="alice", agent_name="research") + manager.add("thread-1", compacted, user_id="alice", agent_name="research") + + assert len(client.added) == 2 + assert [message.message_id for message in client.added[1][2]] == ["df_h20"] + watermark = next((tmp_path / "openviking" / "sessions").glob("*.json")) + state = json.loads(watermark.read_text(encoding="utf-8")) + assert state["submitted_prefix_count"] == 9 + + +def test_manager_migrates_legacy_recent_id_watermark(tmp_path: Path) -> None: + manager, client = _manager(tmp_path, max_seen_message_ids=16) + messages = [HumanMessage(f"message {index}", id=f"h{index}") for index in range(20)] + manager.add("thread-1", messages, user_id="alice", agent_name="research") + + watermark = next((tmp_path / "openviking" / "sessions").glob("*.json")) + legacy_state = json.loads(watermark.read_text(encoding="utf-8")) + legacy_state["schema_version"] = 2 + legacy_state.pop("submitted_prefix_count", None) + legacy_state.pop("submitted_prefix_digest", None) + legacy_state.pop("committed_prefix_count", None) + legacy_state.pop("committed_prefix_digest", None) + watermark.write_text(json.dumps(legacy_state), encoding="utf-8") + + messages.append(HumanMessage("message 20", id="h20")) + manager.add("thread-1", messages, user_id="alice", agent_name="research") + + assert len(client.added) == 2 + assert [message.message_id for message in client.added[1][2]] == ["df_h20"] + + def test_manager_identity_is_stable_and_agent_isolated(tmp_path: Path) -> None: manager, client = _manager(tmp_path) messages = [HumanMessage("hello", id="h1"), AIMessage("hi", id="a1")] diff --git a/config.example.yaml b/config.example.yaml index 81f4bbf66..cfda3c502 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1690,6 +1690,9 @@ memory: # auth_mode: trusted # account: deerflow # api_key_env: OPENVIKING_API_KEY + # max_connections: 100 + # max_keepalive_connections: 20 + # max_seen_message_ids: 512 # recent-ID fallback for compacted histories # startup_policy: fail_fast # failure_policy: # read: fail_open diff --git a/docs/OPENVIKING.md b/docs/OPENVIKING.md index 2412d02f0..6bee05eed 100644 --- a/docs/OPENVIKING.md +++ b/docs/OPENVIKING.md @@ -14,8 +14,9 @@ runtime. - Explicit search through the backend-neutral `MemoryManager.search` API. - Hard isolation by hashing each DeerFlow `(user_id, agent_name)` scope into a separate OpenViking trusted user identity. -- Local message watermarks under DeerFlow's runtime home to avoid resubmitting - the same conversation history in a single-Gateway deployment. +- Local bounded message watermarks under DeerFlow's runtime home. An ordered + prefix digest handles append-only histories of any length, while a recent-ID + window handles history compaction without resubmitting known messages. The current backend does not implement DeerMem fact CRUD, import/export, or the Settings memory document. Keep `mode: middleware`; tool mode is rejected @@ -57,6 +58,9 @@ memory: auth_mode: trusted account: deerflow api_key_env: OPENVIKING_API_KEY + max_connections: 100 + max_keepalive_connections: 20 + max_seen_message_ids: 512 startup_policy: fail_fast failure_policy: read: fail_open @@ -72,6 +76,11 @@ For a locally installed DeerFlow process, use the host, use `http://host.docker.internal:1933` and set `allow_insecure_http: true`. +`max_connections` and `max_keepalive_connections` bound the shared HTTP +connection pool. `max_seen_message_ids` bounds only the recent-ID fallback used +when a conversation is compacted or rewritten; append-only histories are +tracked by a constant-size prefix digest and do not depend on that window. + ## Docker first-time startup Create the standard DeerFlow local files if they no longer exist: @@ -200,6 +209,8 @@ The DeerFlow entrypoint is . watermark before committing the Session. If commit then fails, later updates do not resubmit those messages or retry the ambiguous commit; a future batch can commit the still-open Session together with new messages. +- Retried health, session lookup, and search requests use exponential backoff + with jitter so concurrent Gateway workers do not retry in lockstep. - OpenViking commit is eventually consistent: accepting a commit archives the messages immediately, while summary and memory extraction finish in a background task.