diff --git a/plugins/memory/openviking/__init__.py b/plugins/memory/openviking/__init__.py index f46d71321e6aa..0ada506dce465 100644 --- a/plugins/memory/openviking/__init__.py +++ b/plugins/memory/openviking/__init__.py @@ -89,6 +89,12 @@ def __init__(self, endpoint: str, api_key: str = "", if self._httpx is None: raise ImportError("httpx is required for OpenViking: pip install httpx") + # Bypass system proxy for internal endpoints to avoid corporate proxy + # returning 403 for RFC-1918 addresses (e.g. 10.37.43.162). + # trust_env=False ignores http_proxy/https_proxy env vars. + httpx_module = self._httpx + self._httpx = httpx_module.Client(trust_env=False) + def _headers(self) -> dict: h = { "Content-Type": "application/json", @@ -207,8 +213,8 @@ def health(self) -> bool: "name": "viking_remember", "description": ( "Explicitly store a fact or memory in the OpenViking knowledge base. " - "Use for important information the agent should remember long-term. " - "The system automatically categorizes and indexes the memory." + "The memory is indexed immediately and searchable across all sessions. " + "Use for important information the agent should remember long-term." ), "parameters": { "type": "object", @@ -334,7 +340,38 @@ def system_prompt_block(self) -> str: ) def prefetch(self, query: str, *, session_id: str = "") -> str: - """Return prefetched results from the background thread.""" + """Search relevant memories for the current query and return results. + + Uses the current question (not the previous turn's) for recall. + Falls back to the cached background prefetch result if the synchronous + search times out. + """ + if not self._client or not query: + return "" + + # Try synchronous search with current query first + try: + client = _VikingClient(self._endpoint, self._api_key) + resp = client.post("/api/v1/search/find", { + "query": query, + "top_k": 5, + }) + result = resp.get("result", {}) + parts = [] + for ctx_type in ("memories", "resources"): + items = result.get(ctx_type, []) + for item in items[:3]: + uri = item.get("uri", "") + abstract = item.get("abstract", "") + score = item.get("score", 0) + if abstract: + parts.append(f"- [{score:.2f}] {abstract} ({uri})") + if parts: + return f"## OpenViking Context\n" + "\n".join(parts) + except Exception: + pass + + # Fallback: try cached background prefetch result if self._prefetch_thread and self._prefetch_thread.is_alive(): self._prefetch_thread.join(timeout=3.0) with self._prefetch_lock: @@ -599,9 +636,15 @@ def _tool_remember(self, args: dict) -> str: ], }) + # Commit immediately so the memory becomes searchable across sessions. + try: + self._client.post(f"/api/v1/sessions/{self._session_id}/commit") + except Exception as e: + logger.debug("OpenViking commit after remember failed: %s", e) + return json.dumps({ "status": "stored", - "message": "Memory recorded. Will be extracted and indexed on session commit.", + "message": "Memory recorded and indexed for cross-session recall.", }) def _tool_add_resource(self, args: dict) -> str: diff --git a/run_agent.py b/run_agent.py index 333dda3927fda..fdeb86227ec49 100644 --- a/run_agent.py +++ b/run_agent.py @@ -8058,7 +8058,7 @@ def run_conversation( api_start_time = time.time() retry_count = 0 - max_retries = 3 + max_retries = 10 primary_recovery_attempted = False max_compression_attempts = 3 codex_auth_retry_attempted=False diff --git a/tests/plugins/memory/test_openviking_provider.py b/tests/plugins/memory/test_openviking_provider.py new file mode 100644 index 0000000000000..8e69ce445777d --- /dev/null +++ b/tests/plugins/memory/test_openviking_provider.py @@ -0,0 +1,206 @@ +"""Tests for the OpenVikingMemoryProvider.""" + +import pytest +from unittest.mock import patch, MagicMock + + +class FakeHttpx: + """Fake httpx module for testing without network. + + Supports both module-level calls (httpx.get/httpx.post) and + Client-instance calls (httpx.Client().get/httpx.Client().post) + since _VikingClient now uses httpx.Client(trust_env=False). + """ + + def __init__(self, health_ok=True): + self.health_ok = health_ok + self.calls = [] + + def __call__(self, **kwargs): + """Return self when called as httpx.Client(trust_env=False).""" + return self + + def Client(self, **kwargs): + """Return self when _VikingClient calls httpx.Client(trust_env=False).""" + return self + + def get(self, url, **kwargs): + self.calls.append(("get", url, kwargs)) + m = MagicMock() + m.status_code = 200 if self.health_ok else 500 + m.raise_for_status = MagicMock( + side_effect=Exception("server error") if not self.health_ok else None + ) + m.json = MagicMock(return_value={"result": []}) + return m + + def post(self, url, json=None, **kwargs): + self.calls.append(("post", url, json, kwargs)) + m = MagicMock() + m.raise_for_status = MagicMock() + m.json = MagicMock(return_value={"result": { + "memories": [ + {"uri": "viking://user/memories/test-1", "abstract": "Test memory 1", "score": 0.95}, + {"uri": "viking://user/memories/test-2", "abstract": "Test memory 2", "score": 0.80}, + ], + "resources": [], + }}) + return m + + +@pytest.fixture +def fake_httpx(): + return FakeHttpx() + + +@pytest.fixture +def provider(fake_httpx, monkeypatch, tmp_path): + monkeypatch.setenv("OPENVIKING_ENDPOINT", "http://localhost:1933") + monkeypatch.setenv("OPENVIKING_API_KEY", "test-key") + monkeypatch.setattr("plugins.memory.openviking._get_httpx", lambda: fake_httpx) + from plugins.memory.openviking import OpenVikingMemoryProvider + p = OpenVikingMemoryProvider() + p.initialize("session-test-1", hermes_home=str(tmp_path), platform="cli") + return p + + +class TestPrefetchSync: + """Test that prefetch() uses the current query for synchronous recall.""" + + def test_prefetch_returns_sync_results(self, provider, fake_httpx): + """prefetch() should search with the current query and return results.""" + result = provider.prefetch("my test query") + + # Should have called the search endpoint + post_calls = [(c[0], c[1], c[2]) for c in fake_httpx.calls if c[0] == "post"] + assert len(post_calls) == 1 + assert "/api/v1/search/find" in post_calls[0][1] + assert post_calls[0][2]["query"] == "my test query" + + # Should return formatted results + assert "## OpenViking Context" in result + assert "Test memory 1" in result + assert "viking://user/memories/test-1" in result + + def test_prefetch_uses_current_query_not_cached(self, provider, fake_httpx): + """prefetch() should use the current query, not a stale cached one.""" + # Queue a background prefetch for a different query + provider.queue_prefetch("stale background query") + + # Now call prefetch with a new current query + fake_httpx.calls.clear() + result = provider.prefetch("current fresh query") + + # The synchronous search should use "current fresh query", not "stale background query" + post_calls = [(c[0], c[1], c[2]) for c in fake_httpx.calls if c[0] == "post"] + assert len(post_calls) == 1 + assert post_calls[0][2]["query"] == "current fresh query" + assert "## OpenViking Context" in result + + def test_prefetch_empty_query_returns_nothing(self, provider, fake_httpx): + """prefetch() with empty query should return empty string.""" + result = provider.prefetch("") + assert result == "" + + def test_prefetch_empty_client_returns_nothing(self, monkeypatch, tmp_path): + """prefetch() should return empty when client is not initialized.""" + monkeypatch.setenv("OPENVIKING_ENDPOINT", "") + from plugins.memory.openviking import OpenVikingMemoryProvider + p = OpenVikingMemoryProvider() + p.initialize("session-test-2", hermes_home=str(tmp_path), platform="cli") + result = p.prefetch("some query") + assert result == "" + + def test_prefetch_falls_back_to_cache_on_sync_failure(self, provider, fake_httpx): + """If sync search fails, prefetch() should fall back to cached background results.""" + # Simulate sync failure by raising on httpx.post + original_post = fake_httpx.post + def failing_post(*args, **kwargs): + m = MagicMock() + m.raise_for_status = MagicMock(side_effect=Exception("network error")) + return m + fake_httpx.post = failing_post + + # Pre-populate background cache + provider.queue_prefetch("background query") + import time + time.sleep(0.1) + + # prefetch should return empty (no cache from current session) + result = provider.prefetch("current query") + # The sync fails, fallback finds nothing cached, returns empty + assert result == "" + + +class TestQueuePrefetch: + """Test that queue_prefetch() fires background searches.""" + + def test_queue_prefetch_starts_thread(self, provider): + """queue_prefetch() should start a background thread.""" + provider.queue_prefetch("background search query") + assert provider._prefetch_thread is not None + assert provider._prefetch_thread.daemon is True + + +class TestSyncTurn: + """Test that sync_turn() records conversation turns.""" + + def test_sync_turn_increments_turn_count(self, provider, fake_httpx): + """sync_turn() should increment the turn counter.""" + assert provider._turn_count == 0 + provider.sync_turn("user message", "assistant reply") + assert provider._turn_count == 1 + provider.sync_turn("next user message", "next assistant reply") + assert provider._turn_count == 2 + + def test_on_session_end_commits(self, provider, fake_httpx): + """on_session_end() should commit the session.""" + provider.sync_turn("user", "assistant") + provider.on_session_end([]) + + # Should have called commit + commit_calls = [(c[0], c[1]) for c in fake_httpx.calls if "/commit" in str(c[1])] + assert len(commit_calls) >= 1 + + +class TestVikingRemember: + """Test that viking_remember() commits immediately for cross-session recall.""" + + def test_remember_commits_immediately(self, provider, fake_httpx): + """viking_remember() should commit immediately so the memory is searchable across sessions.""" + result = provider.handle_tool_call("viking_remember", {"content": "My favorite color is blue"}) + + # Should have added the message + post_calls = [(c[0], c[1]) for c in fake_httpx.calls if c[0] == "post"] + session_msg_calls = [c for c in post_calls if "/messages" in str(c[1])] + assert len(session_msg_calls) == 1 + + # Should have committed immediately + commit_calls = [c for c in post_calls if "/commit" in str(c[1])] + assert len(commit_calls) == 1 + + assert "stored" in result + + def test_remember_with_category_commits(self, provider, fake_httpx): + """viking_remember() with category should also commit immediately.""" + provider.handle_tool_call("viking_remember", { + "content": "Likes coffee", + "category": "preference", + }) + + post_calls = [(c[0], c[1]) for c in fake_httpx.calls if c[0] == "post"] + commit_calls = [c for c in post_calls if "/commit" in str(c[1])] + assert len(commit_calls) == 1 + + +class TestVikingSearch: + """Test that viking_search() returns formatted results.""" + + def test_search_returns_formatted_results(self, provider, fake_httpx): + """viking_search() should return formatted search results.""" + result = provider.handle_tool_call("viking_search", {"query": "favorite color"}) + + post_calls = [(c[0], c[1]) for c in fake_httpx.calls if c[0] == "post"] + assert any("/search/find" in str(c[1]) for c in post_calls) + assert "results" in result +