diff --git a/mock_tests/test_auth.py b/mock_tests/test_auth.py index 192f0eb6d..e2ce21322 100644 --- a/mock_tests/test_auth.py +++ b/mock_tests/test_auth.py @@ -1,8 +1,9 @@ import asyncio import json +import threading import time import warnings -from typing import Union +from typing import List, Union import grpc import pytest @@ -84,6 +85,214 @@ def test_client_credentials(weaviate_auth_mock: HTTPServer, start_grpc_server: g weaviate_auth_mock.check_assertions() +@pytest.mark.asyncio +async def test_client_credentials_refresh_async( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server +) -> None: + """Test the refresh_session branch of the async token refresher. + + Client-credentials tokens carry no refresh token, so the refresher must get a whole + new token from the saved credentials. + """ + token_requests = 0 + + def handler(request: Request) -> Response: + nonlocal token_requests + token_requests += 1 + return Response( + json.dumps({"access_token": ACCESS_TOKEN, "expires_in": 1}), + content_type="application/json", + ) + + weaviate_auth_mock.expect_request("/auth").respond_with_handler(handler) + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + + async with weaviate.use_async_with_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthClientCredentials( + client_secret=CLIENT_SECRET, scope=SCOPE + ), + ) as client: + await client.collections.list_all() + first = token_requests + await asyncio.sleep(3) # refresh interval is max(expires_in - 30, 1) -> 1s + assert token_requests > first # a fresh token was fetched with the credentials + + +def _reject_refreshes(weaviate_auth_mock: HTTPServer) -> List[float]: + """Make the IdP reject every refresh (400 invalid_grant); returns the hit timestamps.""" + hits: List[float] = [] + + def handler(request: Request) -> Response: + hits.append(time.monotonic()) + return Response( + json.dumps({"error": "invalid_grant", "error_description": "refresh token expired"}), + status=400, + content_type="application/json", + ) + + weaviate_auth_mock.expect_request("/auth").respond_with_handler(handler) + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + return hits + + +@pytest.mark.asyncio +async def test_token_refresh_survives_failures_async( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, recwarn +) -> None: + """A failing refresh must not kill the refresher: warn, retry, stay alive. + + A 400 invalid_grant makes authlib raise OAuthError, which is not an httpx.HTTPError. + """ + hits = _reject_refreshes(weaviate_auth_mock) + + async with weaviate.use_async_with_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, + refresh_token=REFRESH_TOKEN, + expires_in=1, # force an immediate (and failing) refresh + ), + ) as client: + task = getattr(client._connection, "_ConnectionBase__token_refresh_task") # noqa: B009 + assert task is not None + await asyncio.sleep(3) + assert not task.done() # the refresher survived the failures + await client.collections.list_all() # ... and the client still works + + assert len(hits) >= 2 # it kept retrying + # recwarn's "default" filter shows an identical warning once per location + assert len([w for w in recwarn if str(w.message).startswith("Con001")]) >= 1 + assert task.done() # close() cancelled and awaited it + + +def test_token_refresh_survives_failures( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, recwarn +) -> None: + """Sync twin of the test above. + + The daemon thread used to die silently on anything but an httpx.HTTPError. + """ + hits = _reject_refreshes(weaviate_auth_mock) + + threads_before = set(threading.enumerate()) + with weaviate.connect_to_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, refresh_token=REFRESH_TOKEN, expires_in=1 + ), + ) as client: + refreshers = [ + t for t in set(threading.enumerate()) - threads_before if t.name == "TokenRefresh" + ] + assert len(refreshers) == 1 + time.sleep(3) + assert refreshers[0].is_alive() # survived the failures + client.collections.list_all() + + assert len(hits) >= 2 + # recwarn's "default" filter shows an identical warning once per location + assert len([w for w in recwarn if str(w.message).startswith("Con001")]) >= 1 + refreshers[0].join(timeout=2) + assert not refreshers[0].is_alive() # close() stops the daemon thread promptly + + +@pytest.mark.asyncio +async def test_async_auth_starts_no_threads( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server +) -> None: + """The async client must refresh tokens with an asyncio task, not threads. + + Under WASM/Pyodide threads cannot start at all, so the TokenRefresh daemon thread + and the event-loop sidecar thread would make every async OIDC flow crash connect(). + """ + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + weaviate_auth_mock.expect_request("/auth").respond_with_json( + { + "access_token": ACCESS_TOKEN, + "expires_in": 500, + "refresh_token": REFRESH_TOKEN, + } + ) + + # compare thread OBJECTS, not names: earlier sync tests leave stale TokenRefresh + # daemon threads alive, which would mask a regression in a name-set comparison + threads_before = set(threading.enumerate()) + tasks_before = asyncio.all_tasks() + async with weaviate.use_async_with_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, refresh_token=REFRESH_TOKEN, expires_in=500 + ), + ) as client: + await client.collections.list_all() + new_thread_names = {t.name for t in set(threading.enumerate()) - threads_before} + assert "TokenRefresh" not in new_thread_names + assert "eventLoop" not in new_thread_names + refresh_tasks = [ + t for t in asyncio.all_tasks() - tasks_before if "token_refresh" in repr(t.get_coro()) + ] + assert len(refresh_tasks) == 1 # the refresher runs as an asyncio task instead + # ... and close() must cancel it AND await it: done as soon as close() returns + assert refresh_tasks[0].done() + + +def test_sync_reconnect_leaves_exactly_one_refresher_thread( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server +) -> None: + """close() must end the daemon thread promptly, and a later connect() must not revive it. + + The thread used to re-read the connection's shutdown event on every loop, so after + close()+connect() it picked up the NEW (unset) event and kept refreshing next to the + new thread — two refreshers per client. + """ + weaviate_auth_mock.expect_request("/auth").respond_with_json( + {"access_token": ACCESS_TOKEN, "expires_in": 500, "refresh_token": REFRESH_TOKEN} + ) + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + + def refreshers() -> List[threading.Thread]: + return [t for t in set(threading.enumerate()) - threads_before if t.name == "TokenRefresh"] + + threads_before = set(threading.enumerate()) + client = weaviate.connect_to_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, refresh_token=REFRESH_TOKEN, expires_in=500 + ), + ) + (first,) = refreshers() + client.close() + first.join(timeout=2) + assert not first.is_alive() # not asleep until the next (470s away) wake-up + + client.connect() + client.collections.list_all() + alive = [t for t in refreshers() if t.is_alive()] + assert len(alive) == 1 and alive[0] is not first + client.close() + alive[0].join(timeout=2) + assert not alive[0].is_alive() + + @pytest.mark.parametrize("header_name", ["Authorization", "authorization"]) def test_auth_header_priority( recwarn, weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, header_name: str diff --git a/weaviate/connect/v4.py b/weaviate/connect/v4.py index d70bfc77b..89f69b751 100644 --- a/weaviate/connect/v4.py +++ b/weaviate/connect/v4.py @@ -62,7 +62,6 @@ JSONPayload, _get_proxies, ) -from weaviate.connect.event_loop import _EventLoopSingleton from weaviate.connect.integrations import _IntegrationConfig from weaviate.embedded import EmbeddedV4 from weaviate.exceptions import ( @@ -154,6 +153,8 @@ def __init__( self._connected = False self._skip_init_checks = skip_init_checks self._grpc_config = grpc_config + self._shutdown_background_event: Optional[Event] = None + self.__token_refresh_task: Optional["asyncio.Task[None]"] = None client_type = "sync" if isinstance(self, ConnectionSync) else "async" embedded_suffix = "-embedded" if self.embedded_db is not None else "" @@ -531,79 +532,95 @@ def _create_background_token_refresh(self, _auth: Optional[_Auth] = None) -> Non if "refresh_token" not in self._client.token and _auth is None: return - # make an event loop sidecar thread for running async token refreshing - event_loop = ( - _EventLoopSingleton.get_instance() - if isinstance(self._client, AsyncOAuth2Client) - else None - ) - - expires_in: int = self._client.token.get( - "expires_in", 60 - ) # use 1minute as token lifetime if not supplied - self._shutdown_background_event = Event() - - def refresh_token() -> None: - if isinstance(self._client, AsyncOAuth2Client): - assert event_loop is not None - self._client.token = event_loop.run_until_complete( - self._client.refresh_token, - url=self._client.metadata["token_endpoint"], - ) - elif isinstance(self._client, OAuth2Client): - self._client.token = self._client.refresh_token( - url=self._client.metadata["token_endpoint"] - ) - - def refresh_session() -> None: - assert _auth is not None - if isinstance(self._client, AsyncOAuth2Client): - assert event_loop is not None - new_session = event_loop.run_until_complete( - _auth.aresult, result=_auth.get_auth_session() - ) - self._client.token = event_loop.run_until_complete(new_session.fetch_token) - elif isinstance(self._client, OAuth2Client): - new_session = _auth.result(_auth.get_auth_session()) - self._client.token = new_session.fetch_token() - - def update_refresh_time() -> int: - assert isinstance(self._client, (OAuth2Client, AsyncOAuth2Client)) - return self._client.token.get("expires_in", 60) - 30 + # stop the refresher from an earlier connect(), if any + self._cancel_background_token_refresh() + + # refresh 30s before the token expires (assume 1 minute if the token does not say); + # the loops below always wait at least 1s + refresh_in: int = self._client.token.get("expires_in", 60) - 30 + # the refresher keeps its own event: after close() + connect() the old refresher + # must stop on this one instead of running on with the new one + shutdown = Event() + self._shutdown_background_event = shutdown + + if isinstance(self._client, AsyncOAuth2Client): + # async client: refresh in an asyncio task on the current loop, not in a thread + # (threads cannot start under WASM/Pyodide) + self.__token_refresh_task = asyncio.get_running_loop().create_task( + self.__periodic_token_refresh_async(refresh_in, _auth, shutdown) + ) + return def periodic_refresh_token(refresh_time: int, _auth: Optional[_Auth]) -> None: - while ( - self._shutdown_background_event is not None - and not self._shutdown_background_event.is_set() - ): - # use refresh token when available - time.sleep(max(refresh_time, 1)) + # wait on the event instead of sleeping, so close() can end the thread right away + while not shutdown.wait(timeout=max(refresh_time, 1)): try: - if self._client is None: + # use one client for the whole round: close()/connect() may replace + # self._client while a refresh is running + client = self._client + if not isinstance(client, OAuth2Client): continue - elif ( - isinstance(self._client, (OAuth2Client, AsyncOAuth2Client)) - and "refresh_token" in self._client.token - ): - refresh_token() + if "refresh_token" in client.token: + client.token = client.refresh_token(url=client.metadata["token_endpoint"]) else: - # client credentials usually does not contain a refresh token => get a new token using the - # saved credentials - refresh_session() - refresh_time = update_refresh_time() - except HTTPError as exc: - # retry again after one second, might be an unstable connection + # client credentials usually does not contain a refresh token => get a + # new token using the saved credentials + assert _auth is not None + new_session = _auth.result(_auth.get_auth_session()) + client.token = new_session.fetch_token() + refresh_time = client.token.get("expires_in", 60) - 30 + except Exception as exc: + # retry in one second; any error must keep the refresher alive, not only + # network errors refresh_time = 1 _Warnings.token_refresh_failed(exc) demon = Thread( target=periodic_refresh_token, - args=(expires_in, _auth), + args=(refresh_in, _auth), daemon=True, name="TokenRefresh", ) demon.start() + def _cancel_background_token_refresh(self) -> Optional["asyncio.Task[None]"]: + """Stop the token refresher: set the shutdown event (sync thread) and cancel the async task. + + Returns the cancelled task, if any, so close() can wait for it to finish. + """ + if self._shutdown_background_event is not None: + self._shutdown_background_event.set() + task, self.__token_refresh_task = self.__token_refresh_task, None + if task is not None: + task.cancel() + return task + + async def __periodic_token_refresh_async( + self, refresh_time: int, _auth: Optional[_Auth], shutdown: Event + ) -> None: + """Async version of ``periodic_refresh_token``, run as a task; close() cancels it.""" + while not shutdown.is_set(): + await asyncio.sleep(max(refresh_time, 1)) + try: + client = self._client + if not isinstance(client, AsyncOAuth2Client): + continue + if "refresh_token" in client.token: + client.token = await client.refresh_token(url=client.metadata["token_endpoint"]) + else: + # client credentials usually does not contain a refresh token => get a + # new token using the saved credentials + assert _auth is not None + new_session = await _Auth.aresult(_auth.get_auth_session()) + client.token = await new_session.fetch_token() + refresh_time = client.token.get("expires_in", 60) - 30 + except asyncio.CancelledError: + raise + except Exception as exc: + # retry in one second; any error must keep the refresher alive + refresh_time = 1 + _Warnings.token_refresh_failed(exc) + def __get_latest_headers(self) -> Dict[str, str]: if "authorization" in self._headers: return self._headers @@ -711,9 +728,13 @@ def exc(e: Exception) -> None: def close(self, colour: executor.Colour) -> executor.Result[None]: if self.embedded_db is not None: self.embedded_db.stop() + refresh_task = self._cancel_background_token_refresh() if colour == "async": async def execute() -> None: + if refresh_task is not None: + # wait for the task to finish before the client is closed + await asyncio.gather(refresh_task, return_exceptions=True) if self._client is not None: assert isinstance(self._client, AsyncClient) await self._client.aclose()