From 8f0da78f84e78f4038773deb994a2138de89673e Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:05:36 +0530 Subject: [PATCH] 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. --- plugins/memory/openviking/__init__.py | 83 ++++++++++++--- tests/openviking_plugin/test_openviking.py | 116 +++++++++++++++++++++ 2 files changed, 184 insertions(+), 15 deletions(-) diff --git a/plugins/memory/openviking/__init__.py b/plugins/memory/openviking/__init__.py index 5e753a2f26d1..e59f9467ac23 100644 --- a/plugins/memory/openviking/__init__.py +++ b/plugins/memory/openviking/__init__.py @@ -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, diff --git a/tests/openviking_plugin/test_openviking.py b/tests/openviking_plugin/test_openviking.py index 5a7c4255c52b..c6edbc010545 100644 --- a/tests/openviking_plugin/test_openviking.py +++ b/tests/openviking_plugin/test_openviking.py @@ -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"