diff --git a/plugins/memory/openviking/__init__.py b/plugins/memory/openviking/__init__.py index db9a4c1501..c1643f2ee8 100644 --- a/plugins/memory/openviking/__init__.py +++ b/plugins/memory/openviking/__init__.py @@ -110,14 +110,10 @@ class _VikingClient: def _url(self, path: str) -> str: return f"{self._endpoint}{path}" - def _auth_headers(self) -> dict: - h = { - "X-OpenViking-Account": self._account, - "X-OpenViking-User": self._user, - } - if self._api_key: - h["X-API-Key"] = self._api_key - return h + def _multipart_headers(self) -> dict: + headers = self._headers() + headers.pop("Content-Type", None) + return headers def _parse_response(self, resp) -> dict: try: @@ -167,7 +163,7 @@ class _VikingClient: resp = self._httpx.post( self._url("/api/v1/resources/temp_upload"), files={"file": (file_path.name, f, mime_type)}, - headers=self._auth_headers(), + headers=self._multipart_headers(), timeout=_TIMEOUT, ) data = self._parse_response(resp) diff --git a/tests/plugins/memory/test_openviking_provider.py b/tests/plugins/memory/test_openviking_provider.py index 467c8bd30e..d5b115600f 100644 --- a/tests/plugins/memory/test_openviking_provider.py +++ b/tests/plugins/memory/test_openviking_provider.py @@ -93,6 +93,62 @@ def test_tool_add_resource_uploads_existing_local_file(tmp_path): assert result["root_uri"] == "viking://resources/sample" +def test_tool_add_resource_uploads_existing_local_directory_and_cleans_zip(tmp_path): + docs = tmp_path / "docs" + docs.mkdir() + (docs / "guide.md").write_text("# Guide\n", encoding="utf-8") + nested = docs / "nested" + nested.mkdir() + (nested / "api.md").write_text("# API\n", encoding="utf-8") + provider = OpenVikingMemoryProvider() + provider._client = MagicMock() + uploaded_paths = [] + provider._client.upload_temp_file.side_effect = ( + lambda path: uploaded_paths.append(path) or "upload_docs.zip" + ) + provider._client.post.return_value = { + "status": "ok", + "result": {"root_uri": "viking://resources/docs"}, + } + + result = json.loads(provider._tool_add_resource({ + "url": str(docs), + "reason": "directory test", + "wait": True, + })) + + assert uploaded_paths + assert uploaded_paths[0].suffix == ".zip" + assert not uploaded_paths[0].exists() + provider._client.post.assert_called_once_with("/api/v1/resources", { + "reason": "directory test", + "wait": True, + "source_name": "docs", + "temp_file_id": "upload_docs.zip", + }) + assert result["status"] == "added" + assert result["root_uri"] == "viking://resources/docs" + + +def test_tool_add_resource_cleans_local_directory_zip_when_add_fails(tmp_path): + docs = tmp_path / "docs" + docs.mkdir() + (docs / "guide.md").write_text("# Guide\n", encoding="utf-8") + provider = OpenVikingMemoryProvider() + provider._client = MagicMock() + uploaded_paths = [] + provider._client.upload_temp_file.side_effect = ( + lambda path: uploaded_paths.append(path) or "upload_docs.zip" + ) + provider._client.post.side_effect = RuntimeError("add failed") + + with pytest.raises(RuntimeError, match="add failed"): + provider._tool_add_resource({"url": str(docs)}) + + assert uploaded_paths + assert not uploaded_paths[0].exists() + + def test_tool_add_resource_sends_remote_url_as_path(): provider = OpenVikingMemoryProvider() provider._client = MagicMock() @@ -109,6 +165,41 @@ def test_tool_add_resource_sends_remote_url_as_path(): }) +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") + client = _VikingClient( + "https://example.com", + api_key="test-key", + account="test-account", + user="test-user", + agent="test-agent", + ) + captured_kwargs = {} + + def capture_httpx_post(url, **kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + status_code=200, + text="", + json=lambda: {"status": "ok", "result": {"temp_file_id": "upload_sample.md"}}, + raise_for_status=lambda: None, + ) + + monkeypatch.setattr(client._httpx, "post", capture_httpx_post) + + assert client.upload_temp_file(sample) == "upload_sample.md" + + assert "files" in captured_kwargs + assert "json" not in captured_kwargs + headers = captured_kwargs["headers"] + assert headers["X-OpenViking-Account"] == "test-account" + assert headers["X-OpenViking-User"] == "test-user" + assert headers["X-OpenViking-Agent"] == "test-agent" + assert headers["X-API-Key"] == "test-key" + assert "Content-Type" not in headers + + def test_viking_client_raises_structured_server_error(): client = _VikingClient.__new__(_VikingClient) response = SimpleNamespace(