fix(openviking): cool down failed refreshes and publish conn identity atomically

Follow-up hardening on the salvaged _ensure_client() (#21130 fix):

- Failed-config cooldown: after a refresh attempt fails for a given
  resolved config, skip re-probing for 30s. Previously every provider
  access against a down endpoint paid a 3s health probe under
  _client_refresh_lock and emitted a warning (2+ per turn, some on
  user-facing threads: prefetch, tool calls, session end). Retries
  still happen after the cooldown or immediately when config changes,
  and the log message now says so instead of the false 'disabled until
  config changes'.
- Atomic connection snapshot: _conn_snapshot (5-tuple, single
  assignment) is published only after a health check passes.
  _new_client() and on_memory_write's writer read it as one load, so
  background writers can no longer observe a torn mix of old/new
  identity fields mid-refresh or target an endpoint that never passed
  health. Field writes in _ensure_client_locked keep tracking the
  attempted config for the unchanged-config dedupe.
- _env_refresh_enabled moves to the top of initialize(): an exception
  mid-initialize (swallowed by MemoryManager) can no longer leave the
  provider silently stuck in never-refresh mode.
- _search_prefetch_context reuses _new_client() and degrades to ''
  on construction failure instead of propagating.

Mutation-checked: neutering the cooldown or publishing the snapshot on
failed health makes the new regression tests fail.
This commit is contained in:
kshitijk4poor 2026-07-22 15:05:36 +05:30 committed by kshitij
parent 1cfe23c6e4
commit 8f0da78f84
2 changed files with 184 additions and 15 deletions

View file

@ -126,6 +126,11 @@ _GENERATED_MEMORY_SUMMARY_FILENAMES = {
}
_LOCAL_OPENVIKING_HOSTS = {"localhost", "127.0.0.1", "::1"}
_LOCAL_OPENVIKING_AUTOSTART_TIMEOUT = 60.0
# After a refresh attempt fails for a given (unchanged) config, skip re-probing
# for this long. Keeps "unavailable endpoints reconnect on a later access"
# true while preventing every provider access from paying a 3s health probe
# (and emitting a warning) under _client_refresh_lock while a server is down.
_FAILED_CONFIG_RETRY_COOLDOWN_SECONDS = 30.0
_OPENVIKING_SERVER_LOG_RELATIVE_PATH = Path("logs") / "openviking-server.log"
_OPENVIKING_RESPONDED_FAILURE_PREFIX = "OpenViking server responded"
_PENDING_SESSIONS_RELATIVE_DIR = Path("openviking") / "pending_sessions"
@ -1895,6 +1900,16 @@ class OpenVikingMemoryProvider(MemoryProvider):
# Connection settings and _client are one published state. Serialize
# refreshes so callers never observe a new config with the old client.
self._client_refresh_lock = threading.Lock()
# Last connection identity that passed a health check, published as a
# single tuple assignment (atomic in CPython) so lock-free background
# writers (_new_client, on_memory_write) never see a torn mix of old
# and new fields, and never target an endpoint that failed health.
self._conn_snapshot: Optional[tuple] = None
# (settings tuple, monotonic timestamp) of the last refresh attempt
# that failed. While the resolved config still matches and the retry
# cooldown hasn't elapsed, _ensure_client_locked() returns None without
# re-probing — keeping provider accesses cheap while a server is down.
self._failed_refresh: Optional[tuple] = None
self._runtime_start_lock = threading.Lock()
self._runtime_start_thread: Optional[threading.Thread] = None
self._runtime_start_pending = False
@ -2177,6 +2192,10 @@ class OpenVikingMemoryProvider(MemoryProvider):
)
else:
self._client = client
self._conn_snapshot = (
endpoint, self._api_key, self._account, self._user, self._agent,
)
self._failed_refresh = None
status_message = (
f"Local OpenViking server at {endpoint} is reachable; "
"OpenViking memory is active for later turns."
@ -2273,6 +2292,11 @@ class OpenVikingMemoryProvider(MemoryProvider):
self._account = settings["account"]
self._user = settings["user"]
self._agent = settings["agent"]
# Baseline established — subsequent accesses may refresh from env
# (#21130). Set here (not at the end of initialize) so an exception in
# the connection attempt below — swallowed by MemoryManager's guard —
# can't leave the provider silently stuck in never-refresh mode.
self._env_refresh_enabled = True
self._session_id = session_id
self._turn_count = 0
hermes_home = str(kwargs.get("hermes_home") or "").strip()
@ -2318,15 +2342,15 @@ class OpenVikingMemoryProvider(MemoryProvider):
self._client = None
if self._client:
self._conn_snapshot = (
self._endpoint, self._api_key, self._account, self._user, self._agent,
)
self._recover_pending_sessions()
# Register as the last active provider for atexit safety net
global _last_active_provider
_last_active_provider = self
# Baseline established — subsequent accesses may refresh from env (#21130).
self._env_refresh_enabled = True
def _ensure_client(self) -> Optional["_VikingClient"]:
"""Return the active client, rebuilding it if the resolved config changed.
@ -2360,6 +2384,7 @@ class OpenVikingMemoryProvider(MemoryProvider):
account = settings["account"]
user = settings["user"]
agent = settings["agent"]
settings_key = (endpoint, api_key, account, user, agent)
config_unchanged = (
endpoint == getattr(self, "_endpoint", None)
@ -2377,6 +2402,18 @@ class OpenVikingMemoryProvider(MemoryProvider):
or (self._runtime_start_thread and self._runtime_start_thread.is_alive())
):
return self._client
# The last attempt at this exact config failed. Don't pay a
# network probe (3s timeout, under the refresh lock) on every
# access while the server stays down — retry after a cooldown or
# as soon as the resolved config changes.
failed = self._failed_refresh
if failed is not None:
failed_key, failed_at = failed
if (
failed_key == settings_key
and time.monotonic() - failed_at < _FAILED_CONFIG_RETRY_COOLDOWN_SECONDS
):
return None
self._endpoint = endpoint
self._api_key = api_key
@ -2396,9 +2433,16 @@ class OpenVikingMemoryProvider(MemoryProvider):
health_state, health_message = _classify_runtime_openviking_health(client, endpoint)
if health_state == "healthy":
self._client = client
self._conn_snapshot = settings_key
self._failed_refresh = None
return self._client
self._failed_refresh = (settings_key, time.monotonic())
if health_state == "responded":
logger.warning("%s OpenViking memory disabled until config changes.", health_message)
logger.warning(
"%s OpenViking memory disabled; will retry on a later access "
"(after cooldown) or when the config changes.",
health_message,
)
else: # unreachable
self._handle_runtime_openviking_unreachable()
self._client = None
@ -2601,6 +2645,18 @@ class OpenVikingMemoryProvider(MemoryProvider):
t.join(timeout=slice_left)
def _new_client(self) -> _VikingClient:
# Read the connection identity as ONE tuple load: these builders run on
# background writer threads without _client_refresh_lock, and reading
# the five fields individually could observe a torn mix of old and new
# values mid-refresh (new endpoint + old api_key). The snapshot is only
# published after a successful health check; fall back to the field
# reads for legacy/hand-wired paths where no snapshot exists yet.
snapshot = self._conn_snapshot
if snapshot is not None:
endpoint, api_key, account, user, agent = snapshot
return _VikingClient(
endpoint, api_key, account=account, user=user, agent=agent,
)
return _VikingClient(
self._endpoint,
self._api_key,
@ -3032,13 +3088,13 @@ class OpenVikingMemoryProvider(MemoryProvider):
if self._env_refresh_enabled:
client = self._ensure_client()
elif self._client is not None:
client = _VikingClient(
self._endpoint,
self._api_key,
account=self._account,
user=self._user,
agent=self._agent,
)
# Legacy/hand-wired path: no env baseline yet. Build from the
# cached identity, degrading to "" like the rest of prefetch.
try:
client = self._new_client()
except Exception as e:
logger.debug("OpenViking prefetch client build failed: %s", e)
return ""
if client is None:
return ""
@ -4165,10 +4221,7 @@ class OpenVikingMemoryProvider(MemoryProvider):
def _write():
try:
client = _VikingClient(
self._endpoint, self._api_key,
account=self._account, user=self._user, agent=self._agent,
)
client = self._new_client()
client.post("/api/v1/content/write", {
"uri": uri,
"content": content,

View file

@ -2,6 +2,7 @@
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Any, cast
from urllib.parse import parse_qs, urlparse
@ -1886,3 +1887,118 @@ class TestEnsureClientReloadsEnv:
out = provider.handle_tool_call("viking_search", {"query": "x"})
assert "not connected" in out.lower()
class TestEnsureClientFailureHardening:
"""Follow-up hardening on top of #21130: failed-config cooldown and
torn-identity-safe connection snapshots."""
def test_failed_config_probes_once_then_cools_down(self, monkeypatch):
"""A down endpoint must not pay a health probe on every access."""
probes = []
class _StubClient:
def __init__(self, endpoint, api_key="", account="", user="", agent=""):
self.endpoint = endpoint
def health(self):
probes.append(self.endpoint)
return False
monkeypatch.setattr(openviking_plugin, "_VikingClient", _StubClient)
monkeypatch.setenv("OPENVIKING_ENDPOINT", "https://down.example")
monkeypatch.setenv("OPENVIKING_API_KEY", "sk-x")
provider = OpenVikingMemoryProvider()
provider._env_refresh_enabled = True
assert provider._ensure_client() is None
assert provider._ensure_client() is None
assert provider._ensure_client() is None
# Only the first access probes; the rest hit the cooldown.
assert len(probes) == 1
def test_failed_config_retries_after_cooldown(self, monkeypatch):
probes = []
class _StubClient:
def __init__(self, endpoint, api_key="", account="", user="", agent=""):
self.endpoint = endpoint
def health(self):
probes.append(self.endpoint)
return len(probes) > 1 # down on first probe, up on retry
monkeypatch.setattr(openviking_plugin, "_VikingClient", _StubClient)
monkeypatch.setenv("OPENVIKING_ENDPOINT", "https://flaky.example")
monkeypatch.setenv("OPENVIKING_API_KEY", "sk-x")
provider = OpenVikingMemoryProvider()
provider._env_refresh_enabled = True
assert provider._ensure_client() is None
# Simulate the cooldown elapsing.
failed_key, _failed_at = provider._failed_refresh
provider._failed_refresh = (
failed_key,
time.monotonic() - openviking_plugin._FAILED_CONFIG_RETRY_COOLDOWN_SECONDS - 1,
)
assert provider._ensure_client() is not None
assert len(probes) == 2
assert provider._failed_refresh is None
def test_failed_config_retries_immediately_on_config_change(self, monkeypatch):
probes = []
class _StubClient:
def __init__(self, endpoint, api_key="", account="", user="", agent=""):
self.endpoint = endpoint
def health(self):
probes.append(self.endpoint)
return self.endpoint == "https://up.example"
monkeypatch.setattr(openviking_plugin, "_VikingClient", _StubClient)
monkeypatch.setenv("OPENVIKING_ENDPOINT", "https://down.example")
monkeypatch.setenv("OPENVIKING_API_KEY", "sk-x")
provider = OpenVikingMemoryProvider()
provider._env_refresh_enabled = True
assert provider._ensure_client() is None
# /reload lands a new endpoint: cooldown must not block the new config.
monkeypatch.setenv("OPENVIKING_ENDPOINT", "https://up.example")
client = provider._ensure_client()
assert client is not None
assert client.endpoint == "https://up.example"
def test_conn_snapshot_published_only_on_healthy(self, monkeypatch):
class _StubClient:
def __init__(self, endpoint, api_key="", account="", user="", agent=""):
self.endpoint = endpoint
def health(self):
return self.endpoint == "https://up.example"
monkeypatch.setattr(openviking_plugin, "_VikingClient", _StubClient)
monkeypatch.setenv("OPENVIKING_ENDPOINT", "https://up.example")
monkeypatch.setenv("OPENVIKING_API_KEY", "sk-good")
monkeypatch.delenv("OPENVIKING_ACCOUNT", raising=False)
monkeypatch.delenv("OPENVIKING_USER", raising=False)
monkeypatch.delenv("OPENVIKING_AGENT", raising=False)
provider = OpenVikingMemoryProvider()
provider._env_refresh_enabled = True
assert provider._ensure_client() is not None
healthy_snapshot = provider._conn_snapshot
assert healthy_snapshot is not None
assert healthy_snapshot[0] == "https://up.example"
assert healthy_snapshot[1] == "sk-good"
# A failed refresh must NOT advance the snapshot: background writers
# (_new_client) keep targeting the last identity that passed health.
monkeypatch.setenv("OPENVIKING_ENDPOINT", "https://down.example")
assert provider._ensure_client() is None
assert provider._conn_snapshot == healthy_snapshot
built = provider._new_client()
assert built.endpoint == "https://up.example"