diff --git a/docs/en/agent-integrations/07-langchain-langgraph.md b/docs/en/agent-integrations/07-langchain-langgraph.md index 52c83d5273..b8763e62da 100644 --- a/docs/en/agent-integrations/07-langchain-langgraph.md +++ b/docs/en/agent-integrations/07-langchain-langgraph.md @@ -264,6 +264,8 @@ Use `close()` after synchronous work or `await store.aclose()` after async work. also supports `with OpenVikingStore(...) as store` and `async with` for automatic cleanup. Cleanup releases only internally created clients; injected clients remain caller-owned. Closing an unused Store creates no client, and a closed Store cannot be reused. +Cleanup waits for active backend requests before releasing the owned client. Use `aclose()` +in async code so this wait does not block the event loop. Following LangGraph's `BaseStore` contract, passing `None` as the value deletes the item and is equivalent to `delete()`. Deleting a key that does not exist is a no-op. diff --git a/docs/zh/agent-integrations/07-langchain-langgraph.md b/docs/zh/agent-integrations/07-langchain-langgraph.md index 9de3f185ee..456020a574 100644 --- a/docs/zh/agent-integrations/07-langchain-langgraph.md +++ b/docs/zh/agent-integrations/07-langchain-langgraph.md @@ -252,6 +252,8 @@ store.close() `with OpenVikingStore(...) as store` 或 `async with` 自动清理。 清理仅释放 Store 内部创建的客户端，注入的客户端仍由调用方管理。 关闭未使用的 Store 不会创建客户端，关闭后不可复用。 +清理会等待正在执行的后端请求结束，再释放内部创建的客户端。 +异步代码应使用 `aclose()`，避免等待过程阻塞事件循环。 按 LangGraph `BaseStore` 约定，传入 `None` 作为值表示删除该条目，等价于 `delete()`。 删除不存在的键不会产生任何影响。 diff --git a/integrations/frameworks/langchain/README.md b/integrations/frameworks/langchain/README.md index 4e9044103c..08349e9457 100644 --- a/integrations/frameworks/langchain/README.md +++ b/integrations/frameworks/langchain/README.md @@ -36,6 +36,7 @@ from openviking_sdk import SyncHTTPClient client = SyncHTTPClient( url="http://127.0.0.1:1933", api_key="your-user-api-key", + auth_mode="api_key", ) client.initialize() retriever = OpenVikingRetriever( @@ -49,6 +50,12 @@ finally: client.close() ``` +For API-key-only integrations, pass `auth_mode="api_key"` alongside `api_key`. +The authentication mode is forwarded to owned synchronous and asynchronous SDK +clients, including clients rebuilt during safe-read recovery. Leaving it unset +preserves the SDK's existing configuration defaults. The same option is available +on the recorder, retriever, message history, store, tools, and context wrappers. + The package also provides `OpenVikingSessionRecorder`, `OpenVikingContextMiddleware`, `OpenVikingStore`, `OpenVikingChatMessageHistory`, and `create_openviking_tools`. diff --git a/integrations/frameworks/langchain/pyproject.toml b/integrations/frameworks/langchain/pyproject.toml index 99d90d1bc6..b7b49094b5 100644 --- a/integrations/frameworks/langchain/pyproject.toml +++ b/integrations/frameworks/langchain/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] dependencies = [ "langchain-core>=1.0.0,<2.0.0", - "openviking-sdk>=0.1.1", + "openviking-sdk>=0.1.7", "pydantic>=2.0.0", ] diff --git a/integrations/frameworks/langchain/src/langchain_openviking/client.py b/integrations/frameworks/langchain/src/langchain_openviking/client.py index 212dae6b23..f87dc32afd 100644 --- a/integrations/frameworks/langchain/src/langchain_openviking/client.py +++ b/integrations/frameworks/langchain/src/langchain_openviking/client.py @@ -76,6 +76,7 @@ class OpenVikingConnection: extra_headers: dict[str, str] | None = None auto_initialize: bool = True async_client: Any = None + auth_mode: str | None = None def __deepcopy__(self, memo: dict[int, Any]) -> OpenVikingConnection: """Copy configuration while preserving caller-owned client identities.""" @@ -85,6 +86,7 @@ def __deepcopy__(self, memo: dict[int, Any]) -> OpenVikingConnection: async_client=self.async_client, url=self.url, api_key=self.api_key, + auth_mode=self.auth_mode, account=self.account, user=self.user, user_id=self.user_id, @@ -561,6 +563,7 @@ def _create_client_from_connection(connection: OpenVikingConnection) -> Any: actor_peer_id=connection.actor_peer_id, timeout=connection.timeout, extra_headers=connection.extra_headers, + **({"auth_mode": connection.auth_mode} if connection.auth_mode is not None else {}), ) if connection.auto_initialize and hasattr(client, "initialize"): @@ -581,6 +584,7 @@ async def _create_async_client_from_connection(connection: OpenVikingConnection) actor_peer_id=connection.actor_peer_id, timeout=connection.timeout, extra_headers=connection.extra_headers, + **({"auth_mode": connection.auth_mode} if connection.auth_mode is not None else {}), ) try: diff --git a/integrations/frameworks/langchain/src/langchain_openviking/context.py b/integrations/frameworks/langchain/src/langchain_openviking/context.py index 8421d3f729..df929d3aa9 100644 --- a/integrations/frameworks/langchain/src/langchain_openviking/context.py +++ b/integrations/frameworks/langchain/src/langchain_openviking/context.py @@ -189,6 +189,7 @@ def __init__( retriever: OpenVikingRetriever | None = None, url: str | None = None, api_key: str | None = None, + auth_mode: str | None = None, account: str | None = None, user: str | None = None, user_id: str | None = None, @@ -210,6 +211,7 @@ def __init__( async_client=async_client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, @@ -223,6 +225,7 @@ def __init__( async_client=async_client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, @@ -639,6 +642,7 @@ def with_openviking_context( async_client: Any = None, url: str | None = None, api_key: str | None = None, + auth_mode: str | None = None, account: str | None = None, user: str | None = None, user_id: str | None = None, @@ -680,6 +684,7 @@ def with_openviking_context( async_client=async_client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, @@ -700,6 +705,7 @@ def with_openviking_context( async_client=async_client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, diff --git a/integrations/frameworks/langchain/src/langchain_openviking/history.py b/integrations/frameworks/langchain/src/langchain_openviking/history.py index 0ef508bef6..51b1d37fc0 100644 --- a/integrations/frameworks/langchain/src/langchain_openviking/history.py +++ b/integrations/frameworks/langchain/src/langchain_openviking/history.py @@ -52,6 +52,7 @@ def __init__( async_client: Any = None, url: str | None = None, api_key: str | None = None, + auth_mode: str | None = None, account: str | None = None, user: str | None = None, user_id: str | None = None, @@ -84,6 +85,7 @@ def __init__( async_client=async_client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, diff --git a/integrations/frameworks/langchain/src/langchain_openviking/middleware.py b/integrations/frameworks/langchain/src/langchain_openviking/middleware.py index 99bbb6f53f..6a3d5be5d4 100644 --- a/integrations/frameworks/langchain/src/langchain_openviking/middleware.py +++ b/integrations/frameworks/langchain/src/langchain_openviking/middleware.py @@ -110,6 +110,7 @@ def __init__( retriever: OpenVikingRetriever | None = None, url: str | None = None, api_key: str | None = None, + auth_mode: str | None = None, account: str | None = None, user: str | None = None, user_id: str | None = None, @@ -141,6 +142,7 @@ def __init__( async_client=async_client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, @@ -153,6 +155,7 @@ def __init__( async_client=async_client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, @@ -168,6 +171,7 @@ def __init__( retriever=self.retriever, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, diff --git a/integrations/frameworks/langchain/src/langchain_openviking/recording.py b/integrations/frameworks/langchain/src/langchain_openviking/recording.py index 4701a1606a..4166a54ff9 100644 --- a/integrations/frameworks/langchain/src/langchain_openviking/recording.py +++ b/integrations/frameworks/langchain/src/langchain_openviking/recording.py @@ -198,6 +198,7 @@ def __init__( async_client: Any = None, url: str | None = None, api_key: str | None = None, + auth_mode: str | None = None, account: str | None = None, user: str | None = None, user_id: str | None = None, @@ -215,6 +216,7 @@ def __init__( async_client=async_client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, diff --git a/integrations/frameworks/langchain/src/langchain_openviking/retrievers.py b/integrations/frameworks/langchain/src/langchain_openviking/retrievers.py index 3a09276be6..bd50bdbd59 100644 --- a/integrations/frameworks/langchain/src/langchain_openviking/retrievers.py +++ b/integrations/frameworks/langchain/src/langchain_openviking/retrievers.py @@ -70,6 +70,7 @@ class OpenVikingRetriever(BaseRetriever): async_client: Any = None url: str | None = None api_key: str | None = None + auth_mode: str | None = None account: str | None = None user: str | None = None user_id: str | None = None @@ -127,6 +128,7 @@ def _get_client(self) -> Any: client=self.client, url=self.url, api_key=self.api_key, + auth_mode=self.auth_mode, account=self.account, user=self.user, user_id=self.user_id, @@ -156,6 +158,7 @@ async def get_async_client(self) -> Any: async_client=self.async_client, url=self.url, api_key=self.api_key, + auth_mode=self.auth_mode, account=self.account, user=self.user, user_id=self.user_id, diff --git a/integrations/frameworks/langchain/src/langchain_openviking/store.py b/integrations/frameworks/langchain/src/langchain_openviking/store.py index e185e942b1..cfeb4d301a 100644 --- a/integrations/frameworks/langchain/src/langchain_openviking/store.py +++ b/integrations/frameworks/langchain/src/langchain_openviking/store.py @@ -7,6 +7,7 @@ import asyncio import json import logging +import threading from collections.abc import Mapping from datetime import datetime, timezone from types import TracebackType @@ -81,6 +82,7 @@ def __init__( client: Any = None, url: str | None = None, api_key: str | None = None, + auth_mode: str | None = None, account: str | None = None, user: str | None = None, user_id: str | None = None, @@ -98,6 +100,7 @@ def __init__( client=client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, @@ -109,18 +112,22 @@ def __init__( self.wait: bool = wait self.search_fetch_limit: int = search_fetch_limit self._client_cache: Any = None + self._client_lock = threading.Lock() + self._call_condition = threading.Condition() + self._active_calls = 0 self._owns_client = client is None self._closed = False def close(self) -> None: """Close the Store and release its internally created client, if any.""" - if self._closed: - return - self._closed = True - client = self._client_cache - self._client_cache = None - if client is not None and self._owns_client: - client.close() + with self._call_condition: + self._closed = True + self._call_condition.wait_for(lambda: self._active_calls == 0) + with self._client_lock: + client = self._client_cache + self._client_cache = None + if client is not None and self._owns_client: + client.close() async def aclose(self) -> None: """Close the Store without blocking the event loop.""" @@ -316,8 +323,7 @@ def _semantic_search( limit: int, offset: int, ) -> list[Any]: - result = call_openviking( - self._client(), + result = self._call( "find", query=query, target_uri=self._index_prefix_uri(namespace_prefix), @@ -357,7 +363,7 @@ def _all_data_uris(self, namespace_prefix: tuple[str, ...]) -> list[str]: seen: set[str] = set() for pattern in ("*.json", "**/*.json"): try: - result = call_openviking(self._client(), "glob", pattern=pattern, uri=base_uri) + result = self._call("glob", pattern=pattern, uri=base_uri) except Exception: continue for uri in _extract_uris(result): @@ -367,17 +373,34 @@ def _all_data_uris(self, namespace_prefix: tuple[str, ...]) -> list[str]: return uris def _client(self) -> Any: - self._raise_if_closed() - if self._client_cache is None: - self._client_cache = ensure_client(self._connection) - return self._client_cache + with self._client_lock: + self._raise_if_closed() + if self._client_cache is None: + client = ensure_client(self._connection) + if self._closed: + if self._owns_client: + client.close() + self._raise_if_closed() + self._client_cache = client + return self._client_cache + + def _call(self, method_name: str, /, **kwargs: Any) -> Any: + with self._call_condition: + self._raise_if_closed() + self._active_calls += 1 + try: + return call_openviking(self._client(), method_name, **kwargs) + finally: + with self._call_condition: + self._active_calls -= 1 + self._call_condition.notify_all() def _raise_if_closed(self) -> None: if self._closed: raise RuntimeError("OpenVikingStore is closed") def _read_record(self, uri: str) -> dict[str, Any]: - content = call_openviking(self._client(), "read", uri=uri) + content = self._call("read", uri=uri) if isinstance(content, bytes): content = content.decode("utf-8") return json.loads(str(content)) @@ -416,8 +439,7 @@ def _write(self, uri: str, content: str) -> None: raise create_exc def _write_create(self, uri: str, content: str) -> None: - call_openviking( - self._client(), + self._call( "write", uri=uri, content=content, @@ -426,8 +448,7 @@ def _write_create(self, uri: str, content: str) -> None: ) def _write_replace(self, uri: str, content: str) -> None: - call_openviking( - self._client(), + self._call( "write", uri=uri, content=content, @@ -437,7 +458,7 @@ def _write_replace(self, uri: str, content: str) -> None: def _remove(self, uri: str) -> None: try: - call_openviking(self._client(), "rm", uri=uri, recursive=False) + self._call("rm", uri=uri, recursive=False) except Exception as exc: if not is_not_found_error(exc): raise diff --git a/integrations/frameworks/langchain/src/langchain_openviking/tools.py b/integrations/frameworks/langchain/src/langchain_openviking/tools.py index ac1c904d89..de2af94ddc 100644 --- a/integrations/frameworks/langchain/src/langchain_openviking/tools.py +++ b/integrations/frameworks/langchain/src/langchain_openviking/tools.py @@ -50,6 +50,7 @@ def create_openviking_tools( client: Any = None, url: str | None = None, api_key: str | None = None, + auth_mode: str | None = None, account: str | None = None, user: str | None = None, user_id: str | None = None, @@ -79,6 +80,7 @@ def get_client() -> Any: client=client, url=url, api_key=api_key, + auth_mode=auth_mode, account=account, user=user, user_id=user_id, diff --git a/integrations/frameworks/langchain/tests/integration/test_recording.py b/integrations/frameworks/langchain/tests/integration/test_recording.py index 983e8125c5..7e947ab6f9 100644 --- a/integrations/frameworks/langchain/tests/integration/test_recording.py +++ b/integrations/frameworks/langchain/tests/integration/test_recording.py @@ -1,5 +1,11 @@ from __future__ import annotations +import copy +import json +import socket +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path from typing import Any import pytest @@ -16,7 +22,9 @@ from langchain_openviking import ( InMemoryOpenVikingClient, OpenVikingChatMessageHistory, + OpenVikingCommitPolicy, OpenVikingContextMiddleware, + OpenVikingSessionRecorder, with_openviking_context, ) @@ -136,3 +144,204 @@ def test_real_create_agent_records_through_middleware_recorder(): "Remember this agent turn.", "Stored by a real LangChain agent.", ] + + +@pytest.fixture +def recording_http_server(monkeypatch): + for name in ("NO_PROXY", "no_proxy"): + monkeypatch.setenv(name, "127.0.0.1,localhost,::1") + requests: list[dict[str, Any]] = [] + failed_read = False + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + requests.append( + { + "path": self.path, + "headers": {key.lower(): value for key, value in self.headers.items()}, + } + ) + encoded = json.dumps({"status": "success", "result": {}}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def do_POST(self): + nonlocal failed_read + payload = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + requests.append( + { + "path": self.path, + "headers": {key.lower(): value for key, value in self.headers.items()}, + "payload": payload, + } + ) + if self.path == "/api/v1/search/find" and not failed_read: + failed_read = True + self.connection.shutdown(socket.SHUT_RDWR) + self.connection.close() + return + result = ( + {"memories": []} if self.path == "/api/v1/search/find" else {"pending_tokens": 1} + ) + encoded = json.dumps({"status": "success", "result": result}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *_args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}", requests + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + assert not thread.is_alive() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["cli", "env"]) +@pytest.mark.parametrize("ambient_mode", ["oidc", "ldap"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_explicit_authentication_survives_copy_recovery_and_owned_close( + tmp_path: Path, monkeypatch, recording_http_server, source, ambient_mode, asynchronous +): + import openviking_sdk + + for name in ( + "OPENVIKING_AUTH_MODE", + "OPENVIKING_USERNAME", + "OPENVIKING_PASSWORD", + "OPENVIKING_OIDC_TOKEN", + "OPENVIKING_ACCOUNT", + "OPENVIKING_USER", + "OPENVIKING_ACTOR_PEER_ID", + ): + monkeypatch.delenv(name, raising=False) + profile = { + "auth_mode": ambient_mode, + "ldap_username": "foreign-user", + "ldap_password": "foreign-password", + "oidc_token": "foreign-token", + } + config = tmp_path / "ovcli.conf" + config.write_text(json.dumps(profile if source == "cli" else {}), encoding="utf-8") + monkeypatch.setenv("OPENVIKING_CLI_CONFIG_FILE", str(config)) + if source == "env": + for name, value in { + "OPENVIKING_AUTH_MODE": ambient_mode, + "OPENVIKING_USERNAME": "foreign-user", + "OPENVIKING_PASSWORD": "foreign-password", + "OPENVIKING_OIDC_TOKEN": "foreign-token", + }.items(): + monkeypatch.setenv(name, value) + + closed: list[Any] = [] + client_class = openviking_sdk.AsyncHTTPClient if asynchronous else openviking_sdk.SyncHTTPClient + original_close = client_class.close + if asynchronous: + + async def close(client): + await original_close(client) + closed.append(client) + else: + + def close(client): + original_close(client) + closed.append(client) + + monkeypatch.setattr(client_class, "close", close) + + url, requests = recording_http_server + original = OpenVikingSessionRecorder( + url=url, + api_key="owner-key", + auth_mode="api_key", + extra_headers={}, + timeout=1, + commit_policy=OpenVikingCommitPolicy(mode="always"), + ) + recorder = copy.deepcopy(original) + try: + if asynchronous: + result = await recorder.arecord( + "auth-session", [HumanMessage(content="Remember this.")] + ) + client = await recorder.get_async_client() + assert await client.find(query="read") == {"memories": []} + else: + result = recorder.record("auth-session", [HumanMessage(content="Remember this.")]) + assert recorder.client.find(query="read") == {"memories": []} + assert result.messages_written == 1 + assert [request["path"] for request in requests] == [ + "/api/v1/sessions/auth-session/messages/batch", + "/api/v1/sessions/auth-session/commit", + "/api/v1/search/find", + "/api/v1/search/find", + ] + assert all(request["headers"]["x-api-key"] == "owner-key" for request in requests) + assert all("authorization" not in request["headers"] for request in requests) + assert len(closed) == 1 + finally: + if asynchronous: + await recorder.aclose() + await original.aclose() + else: + recorder.close() + original.close() + assert len(closed) == 2 + assert closed[0] is not closed[1] + + +@pytest.mark.asyncio +async def test_middleware_applies_explicit_authentication_to_all_owned_components( + tmp_path: Path, monkeypatch, recording_http_server +): + config = tmp_path / "ovcli.conf" + config.write_text("{}", encoding="utf-8") + monkeypatch.setenv("OPENVIKING_CLI_CONFIG_FILE", str(config)) + url, requests = recording_http_server + middleware = OpenVikingContextMiddleware( + url=url, + api_key="owner-key", + auth_mode="api_key", + session_id_resolver=lambda _state, _runtime: "middleware-auth", + include_session_context=True, + commit_policy=OpenVikingCommitPolicy(mode="always"), + ) + # The clients are lazy: later ambient changes must not replace the explicit + # authentication choice when each owned component opens its transport. + monkeypatch.setenv("OPENVIKING_AUTH_MODE", "ldap") + monkeypatch.setenv("OPENVIKING_USERNAME", "foreign-user") + monkeypatch.setenv("OPENVIKING_PASSWORD", "foreign-password") + agent = create_agent( + model=FakeListChatModel(responses=["Stored with explicit authentication."]), + tools=[], + middleware=[middleware], + ) + try: + result = await agent.ainvoke( + {"messages": [HumanMessage(content="Remember this turn.")]}, + config={"configurable": {"thread_id": "middleware-auth"}}, + ) + assert result["messages"][-1].content == "Stored with explicit authentication." + assert any( + request["path"].startswith("/api/v1/sessions/middleware-auth/context") + for request in requests + ) + assert any(request["path"] == "/api/v1/search/search" for request in requests) + assert any(request["path"].endswith("/messages/batch") for request in requests) + assert any(request["path"].endswith("/commit") for request in requests) + assert all(request["headers"]["x-api-key"] == "owner-key" for request in requests) + assert all("authorization" not in request["headers"] for request in requests) + finally: + await middleware.aclose() diff --git a/integrations/frameworks/langchain/tests/test_store.py b/integrations/frameworks/langchain/tests/test_store.py index 024d8a8e43..578c89af86 100644 --- a/integrations/frameworks/langchain/tests/test_store.py +++ b/integrations/frameworks/langchain/tests/test_store.py @@ -3,6 +3,7 @@ import asyncio import json +import threading from datetime import datetime, timezone from unittest.mock import Mock @@ -290,3 +291,105 @@ async def use_store(): operation() with pytest.raises(RuntimeError, match="OpenVikingStore is closed"): await store.aget(("memory",), "note") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pause_point", ["initialize", "safe-read-retry"]) +async def test_store_async_context_drains_cancelled_client_operations( + monkeypatch, tmp_path, pause_point +): + loop = asyncio.get_running_loop() + operation_paused = asyncio.Event() + cleanup_started = asyncio.Event() + read_finished = asyncio.Event() + resume_operation = threading.Event() + context_exited = threading.Event() + instances = [] + reads_after_exit = [] + sdk_client = openviking_sdk.SyncHTTPClient + config = tmp_path / "ovcli.conf" + config.write_text("{}", encoding="utf-8") + monkeypatch.setenv("OPENVIKING_CLI_CONFIG_FILE", str(config)) + + class PausedClient(sdk_client): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.close_count = 0 + self.transport = None + self.index = len(instances) + instances.append(self) + + def initialize(self): + super().initialize() + self.transport = self._async_client._http + if pause_point == "initialize": + loop.call_soon_threadsafe(operation_paused.set) + assert resume_operation.wait(timeout=5), "initialization was not released" + + def read(self, *args, **kwargs): + if pause_point == "safe-read-retry" and self.index == 0: + loop.call_soon_threadsafe(operation_paused.set) + assert resume_operation.wait(timeout=5), "read was not released" + raise ConnectionError("synthetic retryable read failure") + try: + if context_exited.is_set(): + reads_after_exit.append(args) + return json.dumps( + { + "namespace": ["memory"], + "key": "note", + "value": {"text": "stored"}, + "created_at": "2026-01-01T00:00:00+00:00", + "updated_at": "2026-01-01T00:00:00+00:00", + } + ) + finally: + loop.call_soon_threadsafe(read_finished.set) + + def close(self): + self.close_count += 1 + super().close() + + monkeypatch.setattr(openviking_sdk, "SyncHTTPClient", PausedClient) + store = OpenVikingStore( + url="http://127.0.0.1:1933", api_key="synthetic-owner-key", auth_mode="api_key" + ) + + async def use_store(): + async with store: + request = asyncio.create_task(store.aget(("memory",), "note")) + await asyncio.wait_for(operation_paused.wait(), timeout=5) + request.cancel() + with pytest.raises(asyncio.CancelledError): + await request + cleanup_started.set() + context_exited.set() + + context = asyncio.create_task(use_store()) + try: + await asyncio.wait_for(cleanup_started.wait(), timeout=5) + # The cancelled awaiter leaves the SDK operation running in its worker. + completed, _ = await asyncio.wait({context}, timeout=0.1) + resume_operation.set() + await asyncio.wait_for(asyncio.shield(context), timeout=5) + if completed: + await asyncio.wait_for(read_finished.wait(), timeout=5) + await store.aclose() + store.close() + + assert len(instances) == (1 if pause_point == "initialize" else 2) + assert all(client.transport is not None for client in instances) + assert ( + tuple(client.close_count for client in instances), + tuple(client.transport.is_closed for client in instances), + bool(completed), + ) == ((1,) * len(instances), (True,) * len(instances), False) + assert reads_after_exit == [] + with pytest.raises(RuntimeError, match="OpenVikingStore is closed"): + await store.aget(("memory",), "note") + finally: + resume_operation.set() + await asyncio.gather(context, return_exceptions=True) + for client in instances: + if client.transport is not None and not client.transport.is_closed: + await asyncio.to_thread(sdk_client.close, client)