mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-11 14:38:38 +00:00
ec274bdedb
* fix(memory): enforce backend read failure policy * fix(memory): harden failure policy handling * fix(memory): narrow strict read handling * fix(memory): keep timeout handling off saturated executor * fix(memory): preserve legacy fail-closed timeouts
919 lines
27 KiB
Python
919 lines
27 KiB
Python
"""Tests for the official-package OpenViking memory backend."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import gc
|
|
import json
|
|
import socket
|
|
import threading
|
|
import weakref
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from langchain_core.documents import Document
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
|
|
from deerflow.agents.memory.backends.openviking.config import OpenVikingConfig
|
|
from deerflow.agents.memory.backends.openviking.openviking_manager import (
|
|
OpenVikingMemoryManager,
|
|
_canonical_peer_id,
|
|
_session_id,
|
|
)
|
|
from deerflow.agents.memory.manager import (
|
|
MemoryManagerError,
|
|
MemoryReadError,
|
|
_scan_backends,
|
|
reset_memory_manager,
|
|
)
|
|
from deerflow.agents.middlewares.dynamic_context_middleware import DynamicContextMiddleware
|
|
|
|
|
|
class _CommitPolicy:
|
|
def __init__(self, *, mode: str, pending_token_threshold: int = 8_000):
|
|
self.mode = mode
|
|
self.pending_token_threshold = pending_token_threshold
|
|
|
|
|
|
class _PartialWriteError(RuntimeError):
|
|
def __init__(
|
|
self,
|
|
consumed: int,
|
|
*,
|
|
commit_pending: bool = False,
|
|
) -> None:
|
|
super().__init__("partial")
|
|
self.input_messages_consumed = consumed
|
|
self.commit_pending = commit_pending
|
|
|
|
|
|
class _Client:
|
|
supports_request_actor_peer = True
|
|
|
|
def __init__(self, **kwargs: Any):
|
|
self.kwargs = kwargs
|
|
self.closed = False
|
|
self.healthy = True
|
|
|
|
def health(self) -> bool:
|
|
return self.healthy
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class _Recorder:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
commit_policy: _CommitPolicy,
|
|
url: str,
|
|
api_key: str,
|
|
timeout: float,
|
|
extra_headers: dict[str, str],
|
|
) -> None:
|
|
self.commit_policy = commit_policy
|
|
self.connection = {
|
|
"url": url,
|
|
"api_key": api_key,
|
|
"timeout": timeout,
|
|
"extra_headers": extra_headers,
|
|
}
|
|
self._client = _Client(**self.connection)
|
|
self.calls: list[tuple[str, list[Any], str | None, str | None]] = []
|
|
self.flushes: list[tuple[str, str | None]] = []
|
|
self.failures: list[BaseException] = []
|
|
self.closed = False
|
|
|
|
@property
|
|
def client(self) -> _Client:
|
|
return self._client
|
|
|
|
def record(
|
|
self,
|
|
session_id: str,
|
|
messages: list[Any],
|
|
peer_id: str | None = None,
|
|
) -> object:
|
|
self.calls.append((session_id, list(messages), peer_id, _ACTOR_PEER.get()))
|
|
if self.failures:
|
|
raise self.failures.pop(0)
|
|
return object()
|
|
|
|
def flush(self, session_id: str) -> None:
|
|
self.flushes.append((session_id, _ACTOR_PEER.get()))
|
|
if self.failures:
|
|
raise self.failures.pop(0)
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
self._client.close()
|
|
|
|
|
|
class _Retriever:
|
|
def __init__(self, *, client: _Client, **kwargs: Any):
|
|
self.client = client
|
|
self.kwargs = kwargs
|
|
self.limit = kwargs["limit"]
|
|
self.filter = None
|
|
self.session_id = kwargs.get("session_id")
|
|
self.target_uri = kwargs.get("target_uri", "")
|
|
self.calls: list[dict[str, Any]] = []
|
|
|
|
def __copy__(self) -> _Retriever:
|
|
copied = type(self)(client=self.client, **self.kwargs)
|
|
copied.calls = self.calls
|
|
copied.limit = self.limit
|
|
copied.filter = copy.deepcopy(self.filter)
|
|
copied.session_id = self.session_id
|
|
copied.target_uri = copy.deepcopy(self.target_uri)
|
|
return copied
|
|
|
|
def invoke(self, query: str) -> list[Document]:
|
|
self.calls.append(
|
|
{
|
|
"query": query,
|
|
"actor_peer": _ACTOR_PEER.get(),
|
|
"limit": self.limit,
|
|
"filter": copy.deepcopy(self.filter),
|
|
"session_id": self.session_id,
|
|
"target_uri": copy.deepcopy(self.target_uri),
|
|
"search_mode": getattr(self, "search_mode", "find"),
|
|
}
|
|
)
|
|
return [
|
|
Document(
|
|
page_content="Prefers concise answers.",
|
|
metadata={
|
|
"openviking_uri": ("viking://user/memories/preferences/style.md"),
|
|
"openviking_category": "preferences",
|
|
"openviking_score": 0.91,
|
|
},
|
|
)
|
|
]
|
|
|
|
|
|
_ACTOR_PEER: ContextVar[str | None] = ContextVar(
|
|
"test_actor_peer",
|
|
default=None,
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def _use_actor_peer(peer_id: str | None):
|
|
token = _ACTOR_PEER.set(peer_id)
|
|
try:
|
|
yield
|
|
finally:
|
|
_ACTOR_PEER.reset(token)
|
|
|
|
|
|
@pytest.fixture
|
|
def official_integration(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
import deerflow.agents.memory.backends.openviking.openviking_manager as module
|
|
|
|
monkeypatch.setattr(
|
|
module,
|
|
"_load_official_integration",
|
|
lambda: {
|
|
"OpenVikingCommitPolicy": _CommitPolicy,
|
|
"OpenVikingPartialWriteError": _PartialWriteError,
|
|
"OpenVikingRetriever": _Retriever,
|
|
"OpenVikingSessionRecorder": _Recorder,
|
|
"use_actor_peer": _use_actor_peer,
|
|
},
|
|
)
|
|
|
|
|
|
def _backend_config(tmp_path: Path, **overrides: Any) -> dict[str, Any]:
|
|
config: dict[str, Any] = {
|
|
"base_url": "http://openviking:1933",
|
|
"storage_path": str(tmp_path),
|
|
"owner_user_id": "alice",
|
|
"startup_policy": "warn",
|
|
"retrieval": {
|
|
"top_k": 4,
|
|
"max_injection_chars": 1_000,
|
|
"injection_query": "profile preferences and prior decisions",
|
|
},
|
|
}
|
|
config.update(overrides)
|
|
return config
|
|
|
|
|
|
def _manager(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
**overrides: Any,
|
|
) -> OpenVikingMemoryManager:
|
|
monkeypatch.setenv("OPENVIKING_API_KEY", "user-key")
|
|
return OpenVikingMemoryManager.from_config(_backend_config(tmp_path, **overrides))
|
|
|
|
|
|
@pytest.fixture
|
|
def unreachable_openviking_url() -> Iterator[str]:
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as reserved_socket:
|
|
reserved_socket.bind(("127.0.0.1", 0))
|
|
host, port = reserved_socket.getsockname()
|
|
yield f"http://{host}:{port}"
|
|
|
|
|
|
def test_config_uses_single_user_key_and_rejects_legacy_trusted_fields(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.delenv("OPENVIKING_API_KEY", raising=False)
|
|
with pytest.raises(ValueError, match="USER API key"):
|
|
OpenVikingConfig.from_backend_config(_backend_config(tmp_path))
|
|
|
|
monkeypatch.setenv("OPENVIKING_API_KEY", "secret")
|
|
config = OpenVikingConfig.from_backend_config(_backend_config(tmp_path))
|
|
|
|
assert config.owner_user_id == "alice"
|
|
assert config.content_mode == "overview"
|
|
assert config.injection_query == "profile preferences and prior decisions"
|
|
assert "secret" not in repr(config)
|
|
|
|
with pytest.raises(ValueError, match="trusted mode is no longer supported"):
|
|
OpenVikingConfig.from_backend_config(_backend_config(tmp_path, auth_mode="trusted"))
|
|
with pytest.raises(ValueError, match="Unknown OpenViking"):
|
|
OpenVikingConfig.from_backend_config(_backend_config(tmp_path, typo=True))
|
|
with pytest.raises(ValueError, match="reserved prefix"):
|
|
OpenVikingConfig.from_backend_config(_backend_config(tmp_path, default_peer_id="df-agent-default"))
|
|
with pytest.raises(ValueError, match="custom HTTP client fields"):
|
|
OpenVikingConfig.from_backend_config(_backend_config(tmp_path, max_connections=10))
|
|
|
|
|
|
def test_backend_is_discovered_by_registered_name() -> None:
|
|
reset_memory_manager()
|
|
assert _scan_backends()["openviking"] is OpenVikingMemoryManager
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("read_policy", "expected"),
|
|
[
|
|
pytest.param("fail_open", False, id="fail_open"),
|
|
pytest.param("raise", True, id="raise"),
|
|
],
|
|
)
|
|
def test_read_failure_capability_matches_policy(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
read_policy: str,
|
|
expected: bool,
|
|
) -> None:
|
|
manager = _manager(
|
|
tmp_path,
|
|
monkeypatch,
|
|
failure_policy={"read": read_policy},
|
|
)
|
|
|
|
assert manager.read_failures_are_fatal is expected
|
|
|
|
|
|
def test_official_loader_uses_standalone_package() -> None:
|
|
from deerflow.agents.memory.backends.openviking.openviking_manager import (
|
|
_load_official_integration,
|
|
)
|
|
|
|
integration = _load_official_integration()
|
|
|
|
assert integration["OpenVikingSessionRecorder"].__module__.startswith("langchain_openviking")
|
|
assert integration["OpenVikingRetriever"].__module__.startswith("langchain_openviking")
|
|
|
|
|
|
def test_manager_uses_official_recorder_retriever_and_commit_always(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
|
|
assert manager._recorder.commit_policy.mode == "always"
|
|
assert manager._retriever.client is manager._recorder.client
|
|
assert manager._retriever.kwargs["content_mode"] == "overview"
|
|
assert manager._recorder.connection == {
|
|
"url": "http://openviking:1933",
|
|
"api_key": "user-key",
|
|
"timeout": 30.0,
|
|
"extra_headers": {},
|
|
}
|
|
|
|
|
|
def test_context_preserves_existing_fixed_query_behavior(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
|
|
context = manager.get_context(
|
|
"alice",
|
|
agent_name="research",
|
|
thread_id="thread-1",
|
|
)
|
|
|
|
assert context == "- [preferences] Prefers concise answers."
|
|
assert manager._retriever.calls == [
|
|
{
|
|
"query": "profile preferences and prior decisions",
|
|
"actor_peer": "research",
|
|
"limit": 4,
|
|
"filter": None,
|
|
"session_id": _session_id("alice", "research", "thread-1"),
|
|
"target_uri": [
|
|
"viking://user/memories",
|
|
"viking://user/peers/research/memories",
|
|
],
|
|
"search_mode": "search",
|
|
}
|
|
]
|
|
|
|
|
|
def test_context_without_thread_uses_existing_find_path(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
|
|
context = manager.get_context(
|
|
"alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert context == "- [preferences] Prefers concise answers."
|
|
assert manager._retriever.calls == [
|
|
{
|
|
"query": "profile preferences and prior decisions",
|
|
"actor_peer": "research",
|
|
"limit": 4,
|
|
"filter": None,
|
|
"session_id": None,
|
|
"target_uri": [
|
|
"viking://user/memories",
|
|
"viking://user/peers/research/memories",
|
|
],
|
|
"search_mode": "find",
|
|
}
|
|
]
|
|
|
|
|
|
def test_unreachable_context_read_raise_aborts_dynamic_context_injection(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
unreachable_openviking_url: str,
|
|
) -> None:
|
|
manager = _manager(
|
|
tmp_path,
|
|
monkeypatch,
|
|
base_url=unreachable_openviking_url,
|
|
timeout_seconds=0.1,
|
|
failure_policy={"read": "raise"},
|
|
)
|
|
monkeypatch.setattr(
|
|
"deerflow.agents.memory.get_memory_manager",
|
|
lambda: manager,
|
|
)
|
|
monkeypatch.setattr(
|
|
"deerflow.agents.middlewares.dynamic_context_middleware.resolve_runtime_user_id",
|
|
lambda runtime: "alice",
|
|
)
|
|
middleware = DynamicContextMiddleware()
|
|
state = {"messages": [HumanMessage("answer this", id="message-1")]}
|
|
|
|
with pytest.raises(MemoryReadError) as exc_info:
|
|
middleware.before_agent(state, None)
|
|
|
|
assert exc_info.value.__cause__ is not None
|
|
|
|
|
|
def test_strict_scope_mismatch_uses_required_read_error(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(
|
|
tmp_path,
|
|
monkeypatch,
|
|
failure_policy={"read": "raise"},
|
|
)
|
|
monkeypatch.setattr(
|
|
"deerflow.agents.memory.get_memory_manager",
|
|
lambda: manager,
|
|
)
|
|
monkeypatch.setattr(
|
|
"deerflow.agents.middlewares.dynamic_context_middleware.resolve_runtime_user_id",
|
|
lambda runtime: "bob",
|
|
)
|
|
middleware = DynamicContextMiddleware()
|
|
state = {"messages": [HumanMessage("answer this", id="message-1")]}
|
|
|
|
with pytest.raises(MemoryReadError, match="owner_user_id 'alice'") as exc_info:
|
|
middleware.before_agent(state, None)
|
|
with pytest.raises(MemoryReadError, match="owner_user_id 'alice'") as search_error:
|
|
manager.search("preferences", user_id="bob", agent_name="research")
|
|
|
|
for error in (exc_info.value, search_error.value):
|
|
assert isinstance(error.__cause__, MemoryManagerError)
|
|
assert "owner_user_id 'alice'" in str(error.__cause__)
|
|
assert manager._retriever.calls == []
|
|
|
|
|
|
def test_unreachable_context_read_fail_open_returns_no_injected_context(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
unreachable_openviking_url: str,
|
|
) -> None:
|
|
manager = _manager(
|
|
tmp_path,
|
|
monkeypatch,
|
|
base_url=unreachable_openviking_url,
|
|
timeout_seconds=0.1,
|
|
failure_policy={"read": "fail_open"},
|
|
)
|
|
monkeypatch.setattr(
|
|
"deerflow.agents.memory.get_memory_manager",
|
|
lambda: manager,
|
|
)
|
|
monkeypatch.setattr(
|
|
"deerflow.agents.middlewares.dynamic_context_middleware.resolve_runtime_user_id",
|
|
lambda runtime: "alice",
|
|
)
|
|
middleware = DynamicContextMiddleware()
|
|
state = {"messages": [HumanMessage("answer this", id="message-1")]}
|
|
|
|
assert manager.get_context("alice", agent_name="research") == ""
|
|
update = middleware.before_agent(state, None)
|
|
assert update is not None
|
|
assert all(not str(message.id).endswith("__memory") for message in update["messages"])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("read_policy", ["raise", "fail_open"])
|
|
async def test_unreachable_async_context_read_honors_policy(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
unreachable_openviking_url: str,
|
|
read_policy: str,
|
|
) -> None:
|
|
manager = _manager(
|
|
tmp_path,
|
|
monkeypatch,
|
|
base_url=unreachable_openviking_url,
|
|
timeout_seconds=0.1,
|
|
failure_policy={"read": read_policy},
|
|
)
|
|
|
|
if read_policy == "raise":
|
|
with pytest.raises(MemoryReadError) as exc_info:
|
|
await manager.aget_context("alice", agent_name="research")
|
|
assert exc_info.value.__cause__ is not None
|
|
else:
|
|
assert await manager.aget_context("alice", agent_name="research") == ""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unreachable_context_read_raise_aborts_async_dynamic_context_injection(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
unreachable_openviking_url: str,
|
|
) -> None:
|
|
manager = _manager(
|
|
tmp_path,
|
|
monkeypatch,
|
|
base_url=unreachable_openviking_url,
|
|
timeout_seconds=0.1,
|
|
failure_policy={"read": "raise"},
|
|
)
|
|
monkeypatch.setattr(
|
|
"deerflow.agents.memory.get_memory_manager",
|
|
lambda: manager,
|
|
)
|
|
monkeypatch.setattr(
|
|
"deerflow.agents.middlewares.dynamic_context_middleware.resolve_runtime_user_id",
|
|
lambda runtime: "alice",
|
|
)
|
|
middleware = DynamicContextMiddleware()
|
|
state = {"messages": [HumanMessage("answer this", id="message-1")]}
|
|
|
|
with pytest.raises(MemoryReadError) as exc_info:
|
|
await middleware.abefore_agent(state, None)
|
|
|
|
assert exc_info.value.__cause__ is not None
|
|
|
|
|
|
def test_manager_refuses_to_share_single_user_key_across_deerflow_users(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
|
|
with pytest.raises(MemoryManagerError, match="owner_user_id 'alice'"):
|
|
manager.get_context("bob", agent_name="research")
|
|
with pytest.raises(MemoryManagerError, match="owner_user_id 'alice'"):
|
|
manager.add(
|
|
"thread-1",
|
|
[HumanMessage("private", id="h1")],
|
|
user_id="bob",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert manager._retriever.calls == []
|
|
assert manager._recorder.calls == []
|
|
|
|
|
|
def test_manager_records_only_unseen_suffix(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
messages = [
|
|
HumanMessage("Remember Vim.", id="h1"),
|
|
AIMessage("I will remember.", id="a1"),
|
|
]
|
|
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
messages.append(HumanMessage("Also concise.", id="h2"))
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert len(manager._recorder.calls) == 2
|
|
assert manager._recorder.calls[0][1] == messages[:2]
|
|
assert manager._recorder.calls[1][1] == messages[2:]
|
|
assert all(call[2:] == ("research", "research") for call in manager._recorder.calls)
|
|
|
|
|
|
def test_partial_write_progress_is_not_resubmitted(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
messages = [
|
|
HumanMessage("one", id="h1"),
|
|
AIMessage("two", id="a1"),
|
|
HumanMessage("three", id="h2"),
|
|
]
|
|
manager._recorder.failures.append(_PartialWriteError(2))
|
|
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert manager._recorder.calls[0][1] == messages
|
|
assert manager._recorder.calls[1][1] == messages[2:]
|
|
|
|
|
|
def test_pending_commit_is_retried_without_resubmitting_messages(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
messages = [HumanMessage("one", id="h1"), AIMessage("two", id="a1")]
|
|
manager._recorder.failures.append(_PartialWriteError(len(messages), commit_pending=True))
|
|
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
session_id = _session_id("alice", "research", "thread-1")
|
|
assert manager._recorder.flushes == [(session_id, "research")]
|
|
assert len(manager._recorder.calls) == 1
|
|
|
|
|
|
def test_search_maps_documents_without_custom_http_models(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
|
|
results = manager.search(
|
|
"answer style",
|
|
top_k=3,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
category="preferences",
|
|
)
|
|
|
|
assert results == [
|
|
{
|
|
"id": "viking://user/memories/preferences/style.md",
|
|
"content": "Prefers concise answers.",
|
|
"category": "preferences",
|
|
"confidence": 0.91,
|
|
"source": "viking://user/memories/preferences/style.md",
|
|
"score": 0.91,
|
|
}
|
|
]
|
|
assert manager._retriever.calls[0]["filter"] == {
|
|
"op": "must",
|
|
"field": "category",
|
|
"conds": ["preferences"],
|
|
}
|
|
assert manager._retriever.calls[0]["search_mode"] == "find"
|
|
|
|
|
|
def test_capture_keeps_tool_history_but_drops_hidden_injected_context(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
messages = [
|
|
HumanMessage("question", id="h1"),
|
|
AIMessage(
|
|
"",
|
|
id="a-tool",
|
|
tool_calls=[
|
|
{
|
|
"name": "search",
|
|
"args": {"query": "OpenViking"},
|
|
"id": "call-1",
|
|
"type": "tool_call",
|
|
}
|
|
],
|
|
),
|
|
HumanMessage(
|
|
"injected memory",
|
|
id="hidden",
|
|
additional_kwargs={"hide_from_ui": True},
|
|
),
|
|
AIMessage("answer", id="a1"),
|
|
]
|
|
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert manager._recorder.calls[0][1] == [
|
|
messages[0],
|
|
messages[1],
|
|
messages[3],
|
|
]
|
|
|
|
|
|
def test_compacted_history_rebases_without_replaying_known_messages(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(
|
|
tmp_path,
|
|
monkeypatch,
|
|
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("new", id="h20"),
|
|
]
|
|
|
|
manager.add(
|
|
"thread-1",
|
|
compacted,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert manager._recorder.calls[1][1] == [compacted[-1]]
|
|
|
|
|
|
def test_failed_write_does_not_advance_capture_cursor(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
manager._recorder.failures.append(RuntimeError("unavailable"))
|
|
messages = [HumanMessage("retry me", id="h1")]
|
|
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
manager.add(
|
|
"thread-1",
|
|
messages,
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert [call[1] for call in manager._recorder.calls] == [
|
|
messages,
|
|
messages,
|
|
]
|
|
|
|
|
|
def test_corrupt_cursor_fails_closed_instead_of_replaying_history(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
session_id = _session_id("alice", "research", "thread-1")
|
|
cursor = manager._state_path(session_id)
|
|
cursor.parent.mkdir(parents=True)
|
|
cursor.write_text("not-json", encoding="utf-8")
|
|
|
|
with pytest.raises(MemoryManagerError, match="refusing unsafe replay"):
|
|
manager.add(
|
|
"thread-1",
|
|
[HumanMessage("private", id="h1")],
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert manager._recorder.calls == []
|
|
|
|
|
|
def test_peer_mapping_is_stable_and_namespaces_are_disjoint() -> None:
|
|
assert _canonical_peer_id(None, "deerflow") == "deerflow"
|
|
assert _canonical_peer_id("Research", "deerflow") == "research"
|
|
assert _canonical_peer_id("deerflow", "deerflow").startswith("df-agent-")
|
|
assert _canonical_peer_id("-research", "deerflow").startswith("df-agent-")
|
|
assert _canonical_peer_id("df-agent-custom", "deerflow").startswith("df-agent-")
|
|
assert (
|
|
len(
|
|
{
|
|
_canonical_peer_id(None, "deerflow"),
|
|
_canonical_peer_id("deerflow", "deerflow"),
|
|
_canonical_peer_id("-research", "deerflow"),
|
|
_canonical_peer_id("df-agent-custom", "deerflow"),
|
|
}
|
|
)
|
|
== 4
|
|
)
|
|
|
|
|
|
def test_session_locks_do_not_accumulate(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
lock = manager._session_lock("session-1")
|
|
lock_ref = weakref.ref(lock)
|
|
|
|
assert len(manager._session_locks) == 1
|
|
del lock
|
|
gc.collect()
|
|
|
|
assert lock_ref() is None
|
|
assert len(manager._session_locks) == 0
|
|
|
|
|
|
def test_shutdown_uses_existing_manager_lifecycle_contract(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
|
|
assert manager.shutdown_flush(1.0) is True
|
|
assert manager._recorder.closed is True
|
|
assert manager._recorder.client.closed is True
|
|
assert manager.shutdown_flush(1.0) is True
|
|
|
|
|
|
def test_shutdown_timeout_closes_resources_after_active_write_finishes(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
started = threading.Event()
|
|
release = threading.Event()
|
|
original_record = manager._recorder.record
|
|
|
|
def blocking_record(*args: Any, **kwargs: Any) -> object:
|
|
started.set()
|
|
assert release.wait(2.0)
|
|
return original_record(*args, **kwargs)
|
|
|
|
manager._recorder.record = blocking_record
|
|
writer = threading.Thread(
|
|
target=manager.add,
|
|
args=("thread-1", [HumanMessage("hello", id="h1")]),
|
|
kwargs={"user_id": "alice", "agent_name": "research"},
|
|
)
|
|
writer.start()
|
|
assert started.wait(1.0)
|
|
|
|
assert manager.shutdown_flush(0.0) is False
|
|
assert manager._recorder.closed is False
|
|
|
|
release.set()
|
|
writer.join(2.0)
|
|
|
|
assert not writer.is_alive()
|
|
assert manager._recorder.closed is True
|
|
assert manager._recorder.client.closed is True
|
|
assert manager.shutdown_flush(1.0) is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_operations_run_off_the_event_loop(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
event_loop_thread = threading.get_ident()
|
|
worker_threads: list[int] = []
|
|
original_add = OpenVikingMemoryManager.add
|
|
|
|
def recording_add(
|
|
self: OpenVikingMemoryManager,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
worker_threads.append(threading.get_ident())
|
|
original_add(self, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(OpenVikingMemoryManager, "add", recording_add)
|
|
|
|
await manager.aadd(
|
|
"thread-1",
|
|
[HumanMessage("hello", id="h1")],
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
assert worker_threads
|
|
assert worker_threads[0] != event_loop_thread
|
|
|
|
|
|
def test_capture_cursor_contains_no_message_content(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
official_integration: None,
|
|
) -> None:
|
|
manager = _manager(tmp_path, monkeypatch)
|
|
|
|
manager.add(
|
|
"thread-1",
|
|
[HumanMessage("very private text", id="h1")],
|
|
user_id="alice",
|
|
agent_name="research",
|
|
)
|
|
|
|
cursor = next((tmp_path / "openviking" / "sessions").glob("*.json"))
|
|
serialized = cursor.read_text(encoding="utf-8")
|
|
state = json.loads(serialized)
|
|
|
|
assert "very private text" not in serialized
|
|
assert state["submitted_prefix_count"] == 1
|