diff --git a/plugins/memory/openviking/__init__.py b/plugins/memory/openviking/__init__.py index 1202e11a1604..f6296805b5d5 100644 --- a/plugins/memory/openviking/__init__.py +++ b/plugins/memory/openviking/__init__.py @@ -294,12 +294,23 @@ class _VikingClient: @staticmethod def _needs_trusted_identity_retry(exc: Exception) -> bool: + """Detect errors that indicate missing tenant-scoped identity headers. + + Trusted mode can ask for ``X-OpenViking-Account`` / + ``X-OpenViking-User`` using slightly different wording across + OpenViking versions. Match that trusted-mode missing-identity shape + instead of enumerating every exact string, while keeping deliberate + API-key permission denials non-retriable. + """ message = str(exc) - return ( - "Trusted mode requests must include X-OpenViking-Account" in message - or "Trusted mode requests must include X-OpenViking-User" in message - or "Trusted mode requests must include X-OpenViking-Account or explicit account_id" in message - ) + if "Trusted mode requests must include" not in message: + return False + if "X-OpenViking-Account" not in message and "X-OpenViking-User" not in message: + return False + status_code = getattr(exc, "status_code", None) + if status_code is not None and status_code != 400: + return False + return True def _send_with_trusted_identity_retry(self, send, *, multipart: bool = False) -> dict: try: diff --git a/tests/plugins/memory/test_openviking_provider.py b/tests/plugins/memory/test_openviking_provider.py index 24c10f78433f..d70386c74a44 100644 --- a/tests/plugins/memory/test_openviking_provider.py +++ b/tests/plugins/memory/test_openviking_provider.py @@ -1956,12 +1956,16 @@ def test_viking_client_does_not_retry_root_tenant_error_as_trusted_mode(monkeypa captured_headers.append(kwargs.get("headers") or {}) if len(captured_headers) == 1: return SimpleNamespace( - status_code=400, + status_code=403, text="", json=lambda: { "error": { - "code": "INVALID_ARGUMENT", - "message": "ROOT requests to tenant-scoped APIs must include X-OpenViking-Account and X-OpenViking-User headers.", + "code": "PERMISSION_DENIED", + "message": ( + "ROOT API keys cannot access tenant-scoped data APIs in api_key mode. " + "Use a user/admin API key for data access, or trusted mode for upstream " + "identity assertion." + ), }, }, raise_for_status=lambda: None, @@ -1975,7 +1979,7 @@ def test_viking_client_does_not_retry_root_tenant_error_as_trusted_mode(monkeypa monkeypatch.setattr(client._httpx, "post", capture_post) - with pytest.raises(openviking_module._OpenVikingHTTPError, match="ROOT requests"): + with pytest.raises(openviking_module._OpenVikingHTTPError, match="ROOT API keys cannot access"): client.post("/api/v1/search/search", {"query": "status"}) assert len(captured_headers) == 1