feat(openviking): add full recall prefetch policy

Salvage of PR #48927 by @ehz0ah, which consolidates OpenViking recall
work from #41706 (@huangxun375-stack), #33260, #49975, and #32444.

Replaces stale background post-turn prefetch warming with synchronous
current-query recall. The old queue_prefetch warmed the PREVIOUS user
message while turn-start recall consumed the CURRENT one, so injected
context was always about the wrong topic.

Changes:
- prefetch() now does session-aware /api/v1/search/search with the
  current query, falls back to /api/v1/search/find on failure
- Contract-safe payloads: limit, score_threshold, context_type,
  session_id — no top_k, no search-body mode, no target_uri
- L2 content reads for items with level=2 or empty abstracts, capped
  at full_read_limit (default 2)
- Local ranking (score + query-token overlap + leaf boost), dedup,
  score threshold, and injected-char budget
- queue_prefetch() is now a no-op (background warming removed)
- Additive batched viking_read: uris param accepts up to 3 URIs
- Per-request timeout support on _VikingClient.get/post/delete
- Removes stale _prefetch_result/_prefetch_thread/_prefetch_generation
  state and _invalidate_prefetch_state()
- Strengthened system_prompt_block guidance

Salvage follow-up fixes:
- Expose all 8 recall config knobs in get_config_schema() (PR #48927
  had removed them; #41706 correctly exposed them). Env vars remain
  as internal mechanism but are now visible in setup wizard.
- Lower default timeout 8s→4s, request_timeout 6s→3s, full_read_limit
  3→2 to reduce per-turn blocking latency.

Co-authored-by: Hao Zhe <haozhe4547@gmail.com>
Co-authored-by: Eurekaxun <eurekaxun@163.com>
This commit is contained in:
kshitijk4poor 2026-06-24 18:53:49 +05:30
parent 89540d592b
commit ab9134bf16
3 changed files with 1443 additions and 200 deletions

View file

@ -1,7 +1,10 @@
"""Tests for plugins/memory/openviking/__init__.py — URI normalization and payload handling."""
import json
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Any, cast
from urllib.parse import parse_qs, urlparse
import plugins.memory.openviking as openviking_plugin
from plugins.memory.openviking import OpenVikingMemoryProvider
@ -54,6 +57,74 @@ class RecordingVikingClient:
return {"result": {"memories": [], "resources": []}}
def _recall_context_key(value):
if isinstance(value, list):
return tuple(value)
return value
class FakeRecallClient:
calls = []
responses = {}
def __init__(self, *args, **kwargs):
pass
def post(self, path, payload=None, **kwargs):
payload = payload or {}
self.__class__.calls.append(("post", path, dict(payload)))
context_type = _recall_context_key(payload.get("context_type"))
key = (path, context_type, payload.get("query"), payload.get("session_id"))
if key not in self.__class__.responses:
key = (path, context_type, payload.get("query"))
if key not in self.__class__.responses:
key = (path, context_type)
response = self.__class__.responses[key]
if isinstance(response, Exception):
raise response
return response
def get(self, path, params=None, **kwargs):
params = params or {}
self.__class__.calls.append(("get", path, dict(params)))
response = self.__class__.responses[(path, params.get("uri"))]
if isinstance(response, Exception):
raise response
return response
def make_prefetch_provider(monkeypatch, responses, **env):
monkeypatch.setattr(openviking_plugin, "_VikingClient", FakeRecallClient)
FakeRecallClient.calls = []
FakeRecallClient.responses = responses
for key in (
"OPENVIKING_RECALL_LIMIT",
"OPENVIKING_RECALL_SCORE_THRESHOLD",
"OPENVIKING_RECALL_MAX_INJECTED_CHARS",
"OPENVIKING_RECALL_TIMEOUT_SECONDS",
"OPENVIKING_RECALL_REQUEST_TIMEOUT_SECONDS",
"OPENVIKING_RECALL_FULL_READ_LIMIT",
"OPENVIKING_RECALL_PREFER_ABSTRACT",
"OPENVIKING_RECALL_RESOURCES",
):
monkeypatch.delenv(key, raising=False)
for key, value in env.items():
monkeypatch.setenv(key, str(value))
provider = OpenVikingMemoryProvider()
provider._client = object()
provider._endpoint = "http://openviking.test"
provider._account = "default"
provider._user = "default"
provider._agent = "hermes"
provider._session_id = "session-test"
return provider
def wait_prefetch(provider, query="What should we recall?", session_id="session-test"):
return provider.prefetch(query, session_id=session_id)
class TestOpenVikingSummaryUriNormalization:
def test_normalize_summary_uri_maps_pseudo_files_to_parent_directory(self):
assert OpenVikingMemoryProvider._normalize_summary_uri("viking://user/hermes/.overview.md") == "viking://user/hermes"
@ -61,7 +132,6 @@ class TestOpenVikingSummaryUriNormalization:
assert OpenVikingMemoryProvider._normalize_summary_uri("viking://") == "viking://"
assert OpenVikingMemoryProvider._normalize_summary_uri("viking://user/hermes/memories/profile.md") == "viking://user/hermes/memories/profile.md"
class TestOpenVikingSkillQuerySafety:
def test_derive_returns_empty_string_for_non_string_input(self):
assert openviking_plugin._derive_openviking_user_text(None) == ""
@ -124,7 +194,7 @@ class TestOpenVikingSkillQuerySafety:
assert skill_commands._BUNDLE_USER_INSTRUCTION in bundle
assert skill_commands._BUNDLE_FIRST_SKILL_BLOCK in bundle
def test_queue_prefetch_searches_only_slash_skill_user_instruction(self, monkeypatch):
def test_prefetch_searches_only_slash_skill_user_instruction(self, monkeypatch):
RecordingVikingClient.calls = []
monkeypatch.setattr(openviking_plugin, "_VikingClient", RecordingVikingClient)
provider = OpenVikingMemoryProvider()
@ -143,18 +213,21 @@ class TestOpenVikingSkillQuerySafety:
"make a skill for release triage"
)
provider.queue_prefetch(skill_message)
assert provider._prefetch_thread is not None
provider._prefetch_thread.join(timeout=5.0)
provider.prefetch(skill_message)
assert RecordingVikingClient.calls == [
(
"/api/v1/search/find",
{"query": "make a skill for release triage", "limit": 5},
)
{
"query": "make a skill for release triage",
"limit": 24,
"score_threshold": 0,
"context_type": "memory",
},
),
]
def test_queue_prefetch_searches_only_skill_bundle_user_instruction(self, monkeypatch):
def test_prefetch_searches_only_skill_bundle_user_instruction(self, monkeypatch):
RecordingVikingClient.calls = []
monkeypatch.setattr(openviking_plugin, "_VikingClient", RecordingVikingClient)
provider = OpenVikingMemoryProvider()
@ -174,18 +247,21 @@ class TestOpenVikingSkillQuerySafety:
"Large bundled skill body that must not be searched or embedded."
)
provider.queue_prefetch(skill_message)
assert provider._prefetch_thread is not None
provider._prefetch_thread.join(timeout=5.0)
provider.prefetch(skill_message)
assert RecordingVikingClient.calls == [
(
"/api/v1/search/find",
{"query": "fix the failing retrieval test", "limit": 5},
)
{
"query": "fix the failing retrieval test",
"limit": 24,
"score_threshold": 0,
"context_type": "memory",
},
),
]
def test_queue_prefetch_skips_slash_skill_without_user_instruction(self, monkeypatch):
def test_prefetch_skips_slash_skill_without_user_instruction(self, monkeypatch):
RecordingVikingClient.calls = []
monkeypatch.setattr(openviking_plugin, "_VikingClient", RecordingVikingClient)
provider = OpenVikingMemoryProvider()
@ -197,9 +273,8 @@ class TestOpenVikingSkillQuerySafety:
"Large skill body that must not be searched or embedded."
)
provider.queue_prefetch(skill_message)
assert provider.prefetch(skill_message) == ""
assert provider._prefetch_thread is None
assert RecordingVikingClient.calls == []
def test_sync_turn_stores_only_slash_skill_user_instruction(self, monkeypatch):
@ -265,6 +340,33 @@ class TestOpenVikingSkillQuerySafety:
assert RecordingVikingClient.calls == []
class TestOpenVikingConfigSchema:
def test_recall_policy_options_are_exposed_in_setup_schema(self):
provider = OpenVikingMemoryProvider()
schema = provider.get_config_schema()
env_vars = {entry.get("env_var") for entry in schema}
assert "OPENVIKING_RECALL_LIMIT" in env_vars
assert "OPENVIKING_RECALL_SCORE_THRESHOLD" in env_vars
assert "OPENVIKING_RECALL_MAX_INJECTED_CHARS" in env_vars
assert "OPENVIKING_RECALL_TIMEOUT_SECONDS" in env_vars
assert "OPENVIKING_RECALL_REQUEST_TIMEOUT_SECONDS" in env_vars
assert "OPENVIKING_RECALL_FULL_READ_LIMIT" in env_vars
assert "OPENVIKING_RECALL_PREFER_ABSTRACT" in env_vars
assert "OPENVIKING_RECALL_RESOURCES" in env_vars
assert provider._recall_config() == {
"limit": 6,
"score_threshold": 0.15,
"max_injected_chars": 4000,
"timeout_seconds": 4.0,
"request_timeout_seconds": 3.0,
"full_read_limit": 2,
"prefer_abstract": False,
"resources": False,
}
class TestOpenVikingTurnConversion:
def test_extract_current_turn_anchors_on_latest_matching_user_and_assistant(self):
messages = [
@ -659,6 +761,78 @@ class TestOpenVikingRead:
{"uri": "viking://user/hermes/memories/profile.md"},
)]
def test_read_accepts_uri_batch_and_caps_batch_full_content(self):
provider = OpenVikingMemoryProvider()
uris = [
"viking://user/hermes/memories/a.md",
"viking://user/hermes/memories/b.md",
"viking://user/hermes/memories/c.md",
"viking://user/hermes/memories/d.md",
]
provider._client = FakeVikingClient(
{
(
"/api/v1/content/read",
(("uri", uris[0]),),
): {"result": {"content": "a" * 3000}},
(
"/api/v1/content/read",
(("uri", uris[1]),),
): {"result": {"content": "b content"}},
(
"/api/v1/content/read",
(("uri", uris[2]),),
): {"result": {"content": "c content"}},
}
)
result = json.loads(provider._tool_read({"uris": uris, "level": "full"}))
assert result["requested"] == 4
assert result["returned"] == 3
assert result["truncated"] is True
assert [entry["uri"] for entry in result["results"]] == uris[:3]
assert result["results"][0]["content"].endswith(
"[... truncated, use a more specific URI or full level]"
)
assert len(result["results"][0]["content"]) < 2700
assert provider._client.calls == [
("/api/v1/content/read", {"uri": uris[0]}),
("/api/v1/content/read", {"uri": uris[1]}),
("/api/v1/content/read", {"uri": uris[2]}),
]
def test_read_deduplicates_uri_batch_and_keeps_errors_per_uri(self):
provider = OpenVikingMemoryProvider()
ok_uri = "viking://user/hermes/memories/ok.md"
bad_uri = "viking://user/hermes/memories/bad.md"
provider._client = FakeVikingClient(
{
(
"/api/v1/content/read",
(("uri", ok_uri),),
): {"result": {"content": "ok content"}},
(
"/api/v1/content/read",
(("uri", bad_uri),),
): RuntimeError("read failed"),
}
)
result = json.loads(
provider._tool_read({"uris": [ok_uri, ok_uri, bad_uri], "level": "full"})
)
assert result["requested"] == 2
assert result["returned"] == 2
assert result["truncated"] is False
assert result["results"][0]["content"] == "ok content"
assert result["results"][1] == {
"uri": bad_uri,
"level": "full",
"error": "read failed",
}
def test_overview_file_uri_routes_straight_to_content_read_via_stat_probe(self):
"""Pre-check via fs/stat: file URIs skip the directory-only endpoint entirely."""
provider = OpenVikingMemoryProvider()
@ -789,6 +963,364 @@ class TestOpenVikingRead:
]
class TestOpenVikingAutoRecallPrefetch:
def test_prefetch_e2e_sends_limit_and_reads_l2_content(self, monkeypatch):
records = {"searches": [], "reads": [], "headers": []}
class Handler(BaseHTTPRequestHandler):
def _send_json(self, payload):
body = json.dumps(payload).encode("utf-8")
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, *args):
pass
def do_GET(self):
parsed = urlparse(self.path)
if parsed.path == "/health":
self._send_json({"healthy": True})
return
if parsed.path == "/api/v1/content/read":
query = parse_qs(parsed.query)
uri = query.get("uri", [""])[0]
records["reads"].append(uri)
self._send_json({"result": {"content": "E2E full L2 memory content."}})
return
self.send_error(404)
def do_POST(self):
length = int(self.headers.get("Content-Length", "0") or "0")
payload = json.loads(self.rfile.read(length).decode("utf-8") or "{}")
records["headers"].append(dict(self.headers))
if self.path == "/api/v1/search/search":
records["searches"].append(payload)
if payload.get("context_type") == "memory":
self._send_json({
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/e2e-full.md",
"score": 0.9,
"level": 2,
"category": "events",
"abstract": "E2E abstract should not be injected.",
}
],
"resources": [],
}
})
else:
self._send_json({"result": {"memories": [], "resources": []}})
return
self.send_error(404)
server = HTTPServer(("127.0.0.1", 0), Handler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
endpoint = f"http://127.0.0.1:{server.server_port}"
for key in (
"OPENVIKING_RECALL_LIMIT",
"OPENVIKING_RECALL_SCORE_THRESHOLD",
"OPENVIKING_RECALL_MAX_INJECTED_CHARS",
"OPENVIKING_RECALL_PREFER_ABSTRACT",
"OPENVIKING_RECALL_RESOURCES",
"OPENVIKING_API_KEY",
):
monkeypatch.delenv(key, raising=False)
monkeypatch.setenv("OPENVIKING_ENDPOINT", endpoint)
monkeypatch.setenv("OPENVIKING_ACCOUNT", "acct")
monkeypatch.setenv("OPENVIKING_USER", "user")
monkeypatch.setenv("OPENVIKING_AGENT", "hermes")
provider = OpenVikingMemoryProvider()
try:
provider.initialize("e2e-session")
block = provider.prefetch("What should we recall?", session_id="e2e-session")
finally:
provider.shutdown()
server.shutdown()
server.server_close()
thread.join(timeout=3.0)
assert block.startswith("## OpenViking Context\n")
assert "E2E full L2 memory content." in block
assert "E2E abstract should not be injected." not in block
assert records["reads"] == ["viking://user/peers/hermes/memories/e2e-full.md"]
assert len(records["searches"]) == 1
assert records["searches"][0]["context_type"] == "memory"
assert records["searches"][0]["session_id"] == "e2e-session"
assert "target_uri" not in records["searches"][0]
assert all(payload["limit"] == 24 for payload in records["searches"])
assert all("top_k" not in payload for payload in records["searches"])
assert all("mode" not in payload for payload in records["searches"])
assert all(payload["score_threshold"] == 0 for payload in records["searches"])
normalized_headers = [
{key.lower(): value for key, value in headers.items()}
for headers in records["headers"]
]
assert all(headers.get("x-openviking-actor-peer") == "hermes" for headers in normalized_headers)
assert all(headers.get("x-openviking-account") == "acct" for headers in normalized_headers)
assert all(headers.get("x-openviking-user") == "user" for headers in normalized_headers)
def test_prefetch_searches_current_query_when_no_background_result(self, monkeypatch):
responses = {
(
"/api/v1/search/search",
"memory",
"Who is Caroline?",
"session-test",
): {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/caroline.md",
"score": 0.9,
"level": 1,
"category": "profile",
"abstract": "Caroline is a transgender woman.",
}
]
}
},
}
provider = make_prefetch_provider(monkeypatch, responses)
block = provider.prefetch("Who is Caroline?", session_id="session-test")
assert "Caroline is a transgender woman." in block
def test_prefetch_does_not_consume_other_session_query_result(self, monkeypatch):
responses = {
(
"/api/v1/search/search",
"memory",
"Who is Caroline?",
"session-a",
): {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/caroline.md",
"score": 0.9,
"level": 1,
"category": "profile",
"abstract": "Caroline context should stay scoped.",
}
]
}
},
(
"/api/v1/search/search",
"memory",
"When did Melanie run a charity race?",
"session-b",
): {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/melanie-race.md",
"score": 0.9,
"level": 1,
"category": "events",
"abstract": "Melanie ran the charity race on May 20.",
}
]
}
},
}
provider = make_prefetch_provider(monkeypatch, responses)
first_block = provider.prefetch("Who is Caroline?", session_id="session-a")
block = provider.prefetch(
"When did Melanie run a charity race?",
session_id="session-b",
)
assert "Caroline context should stay scoped." in first_block
assert "Melanie ran the charity race on May 20." in block
assert "Caroline context should stay scoped." not in block
def test_prefetch_filters_low_score_items_with_local_threshold(self, monkeypatch):
responses = {
("/api/v1/search/search", "memory", "What should we recall?", "session-test"): {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/keep.md",
"score": 0.22,
"level": 1,
"category": "preferences",
"abstract": "Keep this relevant memory.",
},
{
"uri": "viking://user/peers/hermes/memories/drop.md",
"score": 0.12,
"level": 1,
"category": "preferences",
"abstract": "Drop this weak memory.",
},
]
}
},
}
provider = make_prefetch_provider(monkeypatch, responses)
block = wait_prefetch(provider)
assert block.startswith("## OpenViking Context\n")
assert "Keep this relevant memory." in block
assert "Drop this weak memory." not in block
search_payloads = [call[2] for call in FakeRecallClient.calls if call[:2] == ("post", "/api/v1/search/search")]
assert len(search_payloads) == 1
assert search_payloads[0]["context_type"] == "memory"
assert "target_uri" not in search_payloads[0]
assert all(payload["limit"] == 24 for payload in search_payloads)
assert all("top_k" not in payload for payload in search_payloads)
assert all("mode" not in payload for payload in search_payloads)
assert all(payload["score_threshold"] == 0 for payload in search_payloads)
def test_prefetch_skips_complete_entries_that_do_not_fit_budget(self, monkeypatch):
long_memory = "X" * 120
responses = {
("/api/v1/search/search", "memory", "What should we recall?", "session-test"): {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/too-large.md",
"score": 0.9,
"level": 1,
"category": "memory",
"abstract": long_memory,
},
{
"uri": "viking://user/peers/hermes/memories/small.md",
"score": 0.8,
"level": 1,
"category": "memory",
"abstract": "Small memory fits.",
},
]
}
},
}
provider = make_prefetch_provider(
monkeypatch,
responses,
OPENVIKING_RECALL_MAX_INJECTED_CHARS="90",
)
block = wait_prefetch(provider)
assert "Small memory fits." in block
assert long_memory not in block
assert "XXX" not in block
def test_prefetch_reads_full_l2_content_by_default(self, monkeypatch):
responses = {
("/api/v1/search/search", "memory", "What should we recall?", "session-test"): {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/full.md",
"score": 0.9,
"level": 2,
"category": "events",
"abstract": "Abstract only.",
}
]
}
},
("/api/v1/content/read", "viking://user/peers/hermes/memories/full.md"): {
"result": {"content": "Full L2 memory content."}
},
}
provider = make_prefetch_provider(monkeypatch, responses)
block = wait_prefetch(provider)
assert "Full L2 memory content." in block
assert "Abstract only." not in block
assert (
"get",
"/api/v1/content/read",
{"uri": "viking://user/peers/hermes/memories/full.md"},
) in FakeRecallClient.calls
def test_prefetch_prefer_abstract_does_not_read_l2_content(self, monkeypatch):
responses = {
("/api/v1/search/search", "memory", "What should we recall?", "session-test"): {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/full.md",
"score": 0.9,
"level": 2,
"category": "events",
"abstract": "Use the abstract.",
}
]
}
},
}
provider = make_prefetch_provider(
monkeypatch,
responses,
OPENVIKING_RECALL_PREFER_ABSTRACT="true",
)
block = wait_prefetch(provider)
assert "Use the abstract." in block
assert not any(call[:2] == ("get", "/api/v1/content/read") for call in FakeRecallClient.calls)
def test_prefetch_honors_configured_limit_candidate_limit_and_resources(self, monkeypatch):
responses = {
("/api/v1/search/search", ("memory", "resource"), "What should we recall?", "session-test"): {
"result": {
"memories": [],
"resources": [
{
"uri": "viking://resources/doc.md",
"score": 0.9,
"level": 1,
"category": "resource",
"abstract": "Resource recall enabled.",
}
]
}
},
}
provider = make_prefetch_provider(
monkeypatch,
responses,
OPENVIKING_RECALL_LIMIT="2",
OPENVIKING_RECALL_RESOURCES="true",
)
block = wait_prefetch(provider)
assert "Resource recall enabled." in block
search_payloads = [call[2] for call in FakeRecallClient.calls if call[:2] == ("post", "/api/v1/search/search")]
assert len(search_payloads) == 1
assert search_payloads[0]["context_type"] == ["memory", "resource"]
assert "target_uri" not in search_payloads[0]
assert all(payload["limit"] == 20 for payload in search_payloads)
assert all("top_k" not in payload for payload in search_payloads)
assert all("mode" not in payload for payload in search_payloads)
def test_queue_prefetch_is_noop_for_openviking_recall(self, monkeypatch):
provider = make_prefetch_provider(monkeypatch, {})
provider.queue_prefetch("What should we recall?", session_id="session-test")
assert FakeRecallClient.calls == []
class TestOpenVikingBrowse:
def test_list_browse_unwraps_and_normalizes_entry_shapes(self):
provider = OpenVikingMemoryProvider()

View file

@ -1,6 +1,7 @@
import json
import os
import stat
import time
import zipfile
from types import SimpleNamespace
from unittest.mock import MagicMock
@ -1208,6 +1209,7 @@ def test_tool_search_sends_limit_not_legacy_top_k():
payload = provider._client.post.call_args.args[1]
assert payload["limit"] == 7
assert "top_k" not in payload
assert "mode" not in payload
def test_tool_search_uses_find_for_normal_search():
@ -1222,6 +1224,7 @@ def test_tool_search_uses_find_for_normal_search():
provider._client.post.assert_called_once_with("/api/v1/search/find", {
"query": "simple lookup",
})
assert "mode" not in provider._client.post.call_args.args[1]
def test_tool_search_uses_session_search_for_deep_search():
@ -1238,6 +1241,7 @@ def test_tool_search_uses_session_search_for_deep_search():
"query": "connect facts",
"session_id": "session-123",
})
assert "mode" not in provider._client.post.call_args.args[1]
def test_tool_add_resource_uploads_existing_local_file(tmp_path):
@ -1590,6 +1594,36 @@ def test_viking_client_delete_uses_identity_headers(monkeypatch):
assert captured["kwargs"]["headers"]["X-OpenViking-Actor-Peer"] == "hermes"
def test_viking_client_post_allows_per_request_timeout(monkeypatch):
client = _VikingClient(
"https://example.com",
api_key="test-key",
account="acct",
user="alice",
agent="hermes",
)
captured = {}
def capture_post(url, **kwargs):
captured["url"] = url
captured["kwargs"] = kwargs
return SimpleNamespace(
status_code=200,
text="",
json=lambda: {"status": "ok", "result": {}},
raise_for_status=lambda: None,
)
monkeypatch.setattr(client._httpx, "post", capture_post)
assert client.post("/api/v1/search/find", {"query": "anything"}, timeout=1.25) == {
"status": "ok",
"result": {},
}
assert captured["url"] == "https://example.com/api/v1/search/find"
assert captured["kwargs"]["timeout"] == 1.25
def test_viking_client_upload_temp_file_uses_multipart_identity_headers(tmp_path, monkeypatch):
sample = tmp_path / "sample.md"
sample.write_text("# Local resource\n", encoding="utf-8")
@ -2140,10 +2174,8 @@ def test_on_session_switch_commits_pending_tokens_without_turn_count():
assert provider._turn_count == 0
def test_on_session_switch_rewound_same_session_only_invalidates_prefetch():
def test_on_session_switch_rewound_same_session_skips_commit_and_rotation():
provider = _make_provider_with_session("same-sid", turn_count=3)
provider._prefetch_generation = 9
provider._prefetch_result = "stale recall"
provider.on_session_switch("same-sid", rewound=True)
@ -2151,17 +2183,6 @@ def test_on_session_switch_rewound_same_session_only_invalidates_prefetch():
provider._client.post.assert_not_called()
assert provider._session_id == "same-sid"
assert provider._turn_count == 3
assert provider._prefetch_generation == 10
assert provider._prefetch_result == ""
def test_on_session_switch_clears_stale_prefetch_result():
provider = _make_provider_with_session("old-sid", turn_count=1)
provider._prefetch_result = "stale recall from old session"
provider.on_session_switch("new-sid")
assert provider._prefetch_result == ""
def test_on_session_switch_waits_for_inflight_sync_thread():
@ -2856,17 +2877,7 @@ def test_on_memory_write_ignores_non_add_actions(action, content, monkeypatch):
assert spawned == []
# ---------------------------------------------------------------------------
# Prefetch staleness: a prefetch worker that finishes AFTER a session switch
# must drop its result instead of repopulating the new session with stale
# recall from the old generation. Bump the generation directly (rather than
# calling on_session_switch, whose own join blocks on the test worker) so
# the test isolates the generation-gating behavior.
# ---------------------------------------------------------------------------
def test_queue_prefetch_drops_result_when_generation_changed_mid_flight():
import threading
def _make_prefetch_provider() -> OpenVikingMemoryProvider:
provider = OpenVikingMemoryProvider()
provider._client = MagicMock()
provider._endpoint = "http://test"
@ -2874,76 +2885,355 @@ def test_queue_prefetch_drops_result_when_generation_changed_mid_flight():
provider._account = "acct"
provider._user = "usr"
provider._agent = "hermes"
provider._session_id = "old-sid"
return provider
started = threading.Event()
release = threading.Event()
def test_queue_prefetch_is_noop_for_openviking_recall(monkeypatch):
provider = _make_prefetch_provider()
constructed_clients = []
class StubClient:
def __init__(self, *a, **kw):
constructed_clients.append((a, kw))
monkeypatch.setattr(openviking_module, "_VikingClient", StubClient)
provider.queue_prefetch("anything", session_id="sid-123")
assert constructed_clients == []
def test_prefetch_sends_contract_safe_memory_context_payload(monkeypatch):
provider = _make_prefetch_provider()
captured_calls = []
class StubClient:
def __init__(self, *a, **kw):
pass
def post(self, path, payload=None, **kwargs):
started.set()
release.wait(timeout=2.0)
captured_calls.append((path, payload))
return {"result": {"memories": [], "resources": []}}
monkeypatch.setattr(openviking_module, "_VikingClient", StubClient)
provider.prefetch("anything")
assert captured_calls == [
(
"/api/v1/search/find",
{
"query": "anything",
"limit": 24,
"score_threshold": 0,
"context_type": "memory",
},
)
]
payload = captured_calls[0][1]
assert "top_k" not in payload
assert "mode" not in payload
assert "target_uri" not in payload
def test_prefetch_uses_session_search_when_session_id_available(monkeypatch):
provider = _make_prefetch_provider()
captured_calls = []
class StubClient:
def __init__(self, *a, **kw):
pass
def post(self, path, payload=None, **kwargs):
captured_calls.append((path, payload))
return {
"result": {
"memories": [
{"uri": "viking://memories/old", "score": 0.9,
"abstract": "stale from old session"},
{
"uri": "viking://user/peers/hermes/memories/events/mem_1.md",
"score": 0.9,
"abstract": "session-aware memory",
},
],
"resources": [],
"skills": [],
}
}
import plugins.memory.openviking as _mod
real_client_cls = _mod._VikingClient
_mod._VikingClient = StubClient
try:
provider.queue_prefetch("anything")
assert started.wait(timeout=2.0), "prefetch worker never entered post()"
# Simulate a session switch by bumping the generation directly.
# The worker captured the pre-bump generation when it was spawned.
provider._prefetch_generation += 1
release.set()
if provider._prefetch_thread:
provider._prefetch_thread.join(timeout=2.0)
finally:
_mod._VikingClient = real_client_cls
monkeypatch.setattr(openviking_module, "_VikingClient", StubClient)
# The stale result from the pre-bump generation must NOT have been written
# into the new generation's prefetch slot.
assert provider._prefetch_result == ""
result = provider.prefetch("anything", session_id="sid-123")
assert captured_calls == [
(
"/api/v1/search/search",
{
"query": "anything",
"limit": 24,
"score_threshold": 0,
"context_type": "memory",
"session_id": "sid-123",
},
)
]
payload = captured_calls[0][1]
assert "top_k" not in payload
assert "mode" not in payload
assert "target_uri" not in payload
assert "session-aware memory" in result
def test_queue_prefetch_sends_limit_not_legacy_top_k():
provider = OpenVikingMemoryProvider()
provider._client = MagicMock()
provider._endpoint = "http://test"
provider._api_key = ""
provider._account = "acct"
provider._user = "usr"
provider._agent = "hermes"
def test_prefetch_falls_back_to_find_when_session_search_fails(monkeypatch):
provider = _make_prefetch_provider()
captured_payloads = []
captured_calls = []
class StubClient:
def __init__(self, *a, **kw):
pass
def post(self, path, payload=None, **kwargs):
captured_payloads.append(payload)
return {"result": {"memories": [], "resources": []}}
captured_calls.append((path, payload))
if path == "/api/v1/search/search":
raise RuntimeError("session unavailable")
return {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/events/mem_2.md",
"score": 0.8,
"abstract": "non-session fallback",
},
],
"resources": [],
"skills": [],
}
}
import plugins.memory.openviking as _mod
real_client_cls = _mod._VikingClient
_mod._VikingClient = StubClient
try:
provider.queue_prefetch("anything")
if provider._prefetch_thread:
provider._prefetch_thread.join(timeout=2.0)
finally:
_mod._VikingClient = real_client_cls
monkeypatch.setattr(openviking_module, "_VikingClient", StubClient)
assert captured_payloads == [{"query": "anything", "limit": 5}]
assert "top_k" not in captured_payloads[0]
result = provider.prefetch("anything", session_id="sid-123")
assert captured_calls == [
(
"/api/v1/search/search",
{
"query": "anything",
"limit": 24,
"score_threshold": 0,
"context_type": "memory",
"session_id": "sid-123",
},
),
(
"/api/v1/search/find",
{
"query": "anything",
"limit": 24,
"score_threshold": 0,
"context_type": "memory",
},
),
]
for _path, payload in captured_calls:
assert "top_k" not in payload
assert "mode" not in payload
assert "target_uri" not in payload
assert "non-session fallback" in result
def test_prefetch_budget_exhaustion_skips_find_fallback_log(caplog):
class StubClient:
def post(self, path, payload=None, **kwargs):
raise AssertionError("local budget exhaustion should not issue HTTP calls")
with caplog.at_level("DEBUG", logger=openviking_module.__name__):
with pytest.raises(TimeoutError):
OpenVikingMemoryProvider._post_prefetch_search(
StubClient(),
"anything",
"sid-123",
limit=24,
context_type="memory",
deadline=time.monotonic() - 1.0,
request_timeout=4.0,
)
assert "falling back to search/find" not in caplog.text
def test_prefetch_reads_l2_content_and_ignores_skills_by_default(monkeypatch):
provider = _make_prefetch_provider()
captured_reads = []
class StubClient:
def __init__(self, *a, **kw):
pass
def post(self, path, payload=None, **kwargs):
return {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/events/mem_3.md",
"score": 0.9,
"level": 2,
"category": "events",
"abstract": "short abstract",
},
],
"resources": [],
"skills": [
{
"uri": "viking://user/skills/release-triage",
"score": 0.7,
"abstract": "skill context",
},
],
}
}
def get(self, path, params=None, **kwargs):
captured_reads.append((path, params or {}))
return {"result": {"content": "full memory content\nwith useful context"}}
monkeypatch.setattr(openviking_module, "_VikingClient", StubClient)
context = provider.prefetch("anything")
assert captured_reads == [
(
"/api/v1/content/read",
{"uri": "viking://user/peers/hermes/memories/events/mem_3.md"},
)
]
assert "full memory content" in context
assert "short abstract" not in context
assert "skill context" not in context
def test_prefetch_reads_empty_abstract_content_within_budget(monkeypatch):
provider = _make_prefetch_provider()
captured_reads = []
class StubClient:
def __init__(self, *a, **kw):
pass
def post(self, path, payload=None, **kwargs):
return {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/one.md",
"score": 0.9,
"abstract": "",
},
],
"resources": [],
"skills": [],
}
}
def get(self, path, params=None, **kwargs):
captured_reads.append((path, params or {}))
uri = (params or {}).get("uri", "")
return {"result": f"content for {uri}"}
monkeypatch.setattr(openviking_module, "_VikingClient", StubClient)
context = provider.prefetch("anything")
assert [params["uri"] for _path, params in captured_reads] == [
"viking://user/peers/hermes/memories/one.md",
]
assert (
"content for viking://user/peers/hermes/memories/one.md"
in context
)
def test_prefetch_caps_full_content_reads(monkeypatch):
provider = _make_prefetch_provider()
captured_reads = []
class StubClient:
def __init__(self, *a, **kw):
pass
def post(self, path, payload=None, **kwargs):
return {
"result": {
"memories": [
{
"uri": f"viking://user/peers/hermes/memories/events/mem_{idx}.md",
"score": 0.9 - (idx * 0.01),
"level": 2,
"category": "events",
"abstract": f"short abstract {idx}",
}
for idx in range(6)
],
"resources": [],
"skills": [],
}
}
def get(self, path, params=None, **kwargs):
captured_reads.append((path, params or {}))
uri = (params or {}).get("uri", "")
return {"result": {"content": f"full content for {uri}"}}
monkeypatch.setattr(openviking_module, "_VikingClient", StubClient)
context = provider.prefetch("anything")
assert len(captured_reads) == 2
assert "full content for viking://user/peers/hermes/memories/events/mem_0.md" in context
assert "full content for viking://user/peers/hermes/memories/events/mem_1.md" in context
assert "short abstract 2" in context
def test_prefetch_uses_bounded_http_timeouts(monkeypatch):
provider = _make_prefetch_provider()
captured_post_kwargs = []
captured_get_kwargs = []
class StubClient:
def __init__(self, *a, **kw):
pass
def post(self, path, payload=None, **kwargs):
captured_post_kwargs.append(kwargs)
return {
"result": {
"memories": [
{
"uri": "viking://user/peers/hermes/memories/events/mem_timeout.md",
"score": 0.9,
"level": 2,
"category": "events",
"abstract": "short abstract",
},
],
"resources": [],
"skills": [],
}
}
def get(self, path, params=None, **kwargs):
captured_get_kwargs.append(kwargs)
return {"result": {"content": "full memory content"}}
monkeypatch.setattr(openviking_module, "_VikingClient", StubClient)
provider.prefetch("anything", session_id="sid-123")
assert 0 < captured_post_kwargs[0]["timeout"] < openviking_module._TIMEOUT
assert 0 < captured_get_kwargs[0]["timeout"] < openviking_module._TIMEOUT