diff --git a/tests/test_api_client.py b/tests/test_api_client.py index 6b77711f0..971dbce2e 100644 --- a/tests/test_api_client.py +++ b/tests/test_api_client.py @@ -24,7 +24,6 @@ from validmind.utils import md_to_html from validmind.vm_models.figure import Figure - loop = asyncio.new_event_loop() @@ -281,7 +280,9 @@ def test_log_text_generates_text_and_logs_metadata( 200, json={ "content_id": "dataset_summary_text", - "text": md_to_html("## Generated Summary\nGenerated content.", mathml=True), + "text": md_to_html( + "## Generated Summary\nGenerated content.", mathml=True + ), }, ) @@ -390,6 +391,50 @@ def test_log_text_rejects_invalid_context(self): context={"content_ids": ["valid", ""]}, ) + # -- request-path refresh wiring (ZD-682) ----------------------------- + # These assert the request wrappers actually invoke the refresh hook and + # that the 401 retry fires end-to-end, complementing the unit tests that + # exercise _ensure_fresh_oidc_token in isolation. + + @patch("validmind.api_client._ensure_fresh_oidc_token") + @patch("aiohttp.ClientSession.get") + def test_get_refreshes_before_request_and_retries_on_401( + self, mock_get, mock_ensure + ): + mock_ensure.return_value = True # forced refresh reports success -> retry + mock_get.side_effect = [ + MockAsyncResponse(401, text="unauthorized"), + MockAsyncResponse(200, json={"ok": True}), + ] + result = self.run_async(api_client._get, "endpoint") + self.assertEqual(result, {"ok": True}) + self.assertEqual(mock_get.call_count, 2) + self.assertEqual(mock_ensure.call_count, 2) + mock_ensure.assert_any_call(force=True) + + @patch("validmind.api_client._ensure_fresh_oidc_token") + @patch("aiohttp.ClientSession.post") + def test_post_refreshes_before_request(self, mock_post, mock_ensure): + mock_post.return_value = MockAsyncResponse(200, json={"ok": True}) + result = self.run_async(api_client._post, "endpoint", data={"a": "b"}) + self.assertEqual(result, {"ok": True}) + mock_ensure.assert_called_once_with() + + @patch("validmind.api_client._ensure_fresh_oidc_token") + @patch("requests.get") + def test_ping_refreshes_before_request_and_retries_on_401( + self, mock_get, mock_ensure + ): + mock_ensure.return_value = True + mock_get.side_effect = [ + MockResponse(401, text="unauthorized"), + MockResponse(200, json={"model": {"name": "n", "cuid": "c"}}), + ] + api_client._ping() + self.assertEqual(mock_get.call_count, 2) + self.assertEqual(mock_ensure.call_count, 2) + mock_ensure.assert_any_call(force=True) + class TestAPIClientOIDC(unittest.TestCase): """OIDC device-flow authentication via vm.init().""" @@ -656,6 +701,195 @@ def test_api_url_alias_sets_host(self, mock_obtain, mock_ping): ) self.assertEqual(api_client.get_api_host(), "http://localhost/from-api-url/") + # -- request-path token refresh (ZD-682) ------------------------------ + + def _init_oidc( + self, + *, + access_token="tok", + expires_at, + refresh_token="refresh-token", + issuer="https://issuer.example.com/", + ): + """Put the client in OIDC mode with a chosen token expiry.""" + with ( + patch("validmind.api_client._ping"), + patch("validmind.api_client._obtain_oidc_tokens") as mock_obtain, + ): + mock_obtain.return_value = { + "issuer": issuer, + "client_id": "cid", + "access_token": access_token, + "expires_at": expires_at, + "refresh_token": refresh_token, + "id_token": None, + } + api_client.init( + model="model-cuid", + api_host="http://localhost/track/", + api_key="", + api_secret="", + issuer=issuer, + client_id="cid", + document="documentation", + ) + + def test_ensure_fresh_oidc_token_noop_when_token_valid(self): + self._init_oidc(expires_at="2099-01-01T00:00:00+00:00") + with patch("validmind.oidc_device.try_refresh_cached_tokens") as mock_refresh: + self.assertTrue(api_client._ensure_fresh_oidc_token()) + mock_refresh.assert_not_called() + self.assertEqual(api_client._access_token, "tok") + + def test_ensure_fresh_oidc_token_refreshes_when_expired(self): + self._init_oidc(access_token="old-tok", expires_at="2000-01-01T00:00:00+00:00") + entry = { + "issuer": "https://issuer.example.com/", + "client_id": "cid", + "access_token": "old-tok", + "expires_at": "2000-01-01T00:00:00+00:00", + "refresh_token": "refresh-token", + } + new_tokens = { + "issuer": "https://issuer.example.com/", + "client_id": "cid", + "access_token": "new-tok", + "expires_at": "2099-01-01T00:00:00+00:00", + "refresh_token": "refresh-token-2", + } + with ( + patch("validmind.credentials_store.get_cached_entry", return_value=entry), + patch("validmind.credentials_store.upsert_cached_entry") as mock_upsert, + patch( + "validmind.oidc_device.try_refresh_cached_tokens", + return_value=new_tokens, + ) as mock_refresh, + ): + self.assertTrue(api_client._ensure_fresh_oidc_token()) + mock_refresh.assert_called_once() + mock_upsert.assert_called_once() + self.assertEqual(api_client._access_token, "new-tok") + + def test_ensure_fresh_oidc_token_force_refreshes_valid_token(self): + # Backstop for the reactive-401 path: force refresh even when unexpired. + self._init_oidc(access_token="old-tok", expires_at="2099-01-01T00:00:00+00:00") + entry = { + "issuer": "https://issuer.example.com/", + "client_id": "cid", + "access_token": "old-tok", + "expires_at": "2099-01-01T00:00:00+00:00", + "refresh_token": "refresh-token", + } + new_tokens = { + "issuer": "https://issuer.example.com/", + "client_id": "cid", + "access_token": "new-tok", + "expires_at": "2099-01-01T00:00:00+00:00", + "refresh_token": "refresh-token-2", + } + with ( + patch("validmind.credentials_store.get_cached_entry", return_value=entry), + patch("validmind.credentials_store.upsert_cached_entry"), + patch( + "validmind.oidc_device.try_refresh_cached_tokens", + return_value=new_tokens, + ) as mock_refresh, + ): + self.assertTrue(api_client._ensure_fresh_oidc_token(force=True)) + mock_refresh.assert_called_once() + self.assertEqual(api_client._access_token, "new-tok") + + def test_ensure_fresh_oidc_token_returns_false_without_refresh_token(self): + self._init_oidc(expires_at="2000-01-01T00:00:00+00:00", refresh_token=None) + entry = { + "issuer": "https://issuer.example.com/", + "client_id": "cid", + "access_token": "tok", + "expires_at": "2000-01-01T00:00:00+00:00", + "refresh_token": None, + } + with ( + patch("validmind.credentials_store.get_cached_entry", return_value=entry), + patch("validmind.oidc_device.try_refresh_cached_tokens") as mock_refresh, + ): + self.assertFalse(api_client._ensure_fresh_oidc_token()) + mock_refresh.assert_not_called() + + def test_ensure_fresh_oidc_token_noop_in_api_key_mode(self): + with patch("validmind.api_client._ping"): + api_client.init( + api_key=os.environ["VM_API_KEY"], + api_secret=os.environ["VM_API_SECRET"], + api_host=os.environ["VM_API_HOST"], + model=os.environ["VM_API_MODEL"], + document="documentation", + ) + self.assertFalse(api_client._ensure_fresh_oidc_token()) + + def test_raise_for_api_error_gives_clear_message_on_oidc_401(self): + self._init_oidc(expires_at="2099-01-01T00:00:00+00:00") + with self.assertRaises(ValidMindAuthError) as ctx: + api_client._raise_for_api_error(401, "unauthorized") + self.assertIn("vm.init()", str(ctx.exception)) + + def test_locked_caller_adopts_concurrently_refreshed_token(self): + # Models the "losing" caller in a refresh race deterministically: it saw + # a stale token (so the cheap hot-path check is bypassed and it enters + # the locked section), but by the time it holds the lock another caller + # has already refreshed and persisted a fresh token. The in-lock re-check + # must make it adopt that token rather than issue a second refresh — this + # is what collapses a stampede to a single token-endpoint call. (Mutual + # exclusion itself is threading.Lock's job and isn't re-tested here.) + self._init_oidc(access_token="old-tok", expires_at="2000-01-01T00:00:00+00:00") + # In-memory expiry is stale, so the cheap check is bypassed and control + # reaches the locked section, exactly as a contended caller would. + self.assertTrue(api_client._oidc_token_is_stale()) + + fresh_entry = { + "issuer": "https://issuer.example.com/", + "client_id": "cid", + "access_token": "new-tok", + "expires_at": "2099-01-01T00:00:00+00:00", + "refresh_token": "refresh-token-2", + } + with ( + patch( + "validmind.credentials_store.get_cached_entry", return_value=fresh_entry + ), + patch("validmind.oidc_device.try_refresh_cached_tokens") as mock_refresh, + ): + self.assertTrue(api_client._ensure_fresh_oidc_token()) + mock_refresh.assert_not_called() + # Adopted the token another caller already refreshed, no second fetch. + self.assertEqual(api_client._access_token, "new-tok") + + def test_ensure_fresh_oidc_token_drops_cache_on_refresh_failure(self): + # On a failed refresh (e.g. revoked refresh token / invalid_grant), the + # cached entry is deleted so it isn't re-attempted on every request — + # matching init()'s _obtain_oidc_tokens. + self._init_oidc(expires_at="2000-01-01T00:00:00+00:00") + expired = { + "issuer": "https://issuer.example.com/", + "client_id": "cid", + "access_token": "old-tok", + "expires_at": "2000-01-01T00:00:00+00:00", + "refresh_token": "refresh-token", + } + with ( + patch("validmind.credentials_store.get_cached_entry", return_value=expired), + patch("validmind.credentials_store.delete_cached_entry") as mock_delete, + patch( + "validmind.oidc_device.try_refresh_cached_tokens", + side_effect=ValidMindAuthError("invalid_grant"), + ), + ): + self.assertFalse(api_client._ensure_fresh_oidc_token()) + mock_delete.assert_called_once() + # In-memory token is cleared too, so the next request fails fast at header + # build with the re-auth message rather than sending a doomed request. + self.assertIsNone(api_client._access_token) + self.assertIsNone(api_client._oidc_expires_at) + if __name__ == "__main__": unittest.main() diff --git a/validmind/api_client.py b/validmind/api_client.py index c5c4603c0..b3de39431 100644 --- a/validmind/api_client.py +++ b/validmind/api_client.py @@ -11,6 +11,7 @@ import atexit import json import os +import threading from io import BytesIO from typing import Any, Dict, List, Optional, Tuple, Union from urllib.parse import urlencode, urljoin @@ -42,6 +43,10 @@ _auth_mode = "api_key" _access_token: Optional[str] = None _oidc_login_context: Optional[Dict[str, str]] = None +# Expiry (ISO-8601) of the in-memory OIDC access token, for cheap request-path +# staleness checks; guarded together with the token by _oidc_refresh_lock. +_oidc_expires_at: Optional[str] = None +_oidc_refresh_lock = threading.Lock() __api_session: Optional[aiohttp.ClientSession] = None @@ -147,12 +152,21 @@ def _get_url( async def _get( endpoint: str, params: Optional[Dict[str, str]] = None ) -> Dict[str, Any]: + _ensure_fresh_oidc_token() url = _get_url(endpoint, params) session = _get_session() async with session.get(url) as r: + # A 401 can still slip through if the token was revoked ahead of its + # stated expiry; force a refresh and retry once (safe — GET has no body). + if r.status == 401 and _ensure_fresh_oidc_token(force=True): + session = _get_session() + async with session.get(url) as retry: + if retry.status != 200: + _raise_for_api_error(retry.status, await retry.text()) + return await retry.json() if r.status != 200: - raise_api_error(await r.text()) + _raise_for_api_error(r.status, await r.text()) return await r.json() @@ -163,6 +177,7 @@ async def _post( data: Optional[Union[dict, FormData]] = None, files: Optional[Dict[str, Tuple[str, BytesIO, str]]] = None, ) -> Dict[str, Any]: + _ensure_fresh_oidc_token() url = _get_url(endpoint, params) session = _get_session() @@ -186,20 +201,29 @@ async def _post( _data = data async with session.post(url, data=_data) as r: + # Not retried on 401: the request body / upload stream may already be + # consumed. Proactive refresh above covers the expiry case; a 401 here + # means a genuinely rejected token, surfaced as a clear auth error. if r.status != 200: - raise_api_error(await r.text()) + _raise_for_api_error(r.status, await r.text()) return await r.json() def _ping() -> Dict[str, Any]: """Validates that we can connect to the ValidMind API (does not use the async session).""" + _ensure_fresh_oidc_token() r = requests.get( url=_get_url("ping"), headers=_get_api_headers(), ) + if r.status_code == 401 and _ensure_fresh_oidc_token(force=True): + r = requests.get( + url=_get_url("ping"), + headers=_get_api_headers(), + ) if r.status_code != 200: - raise_api_error(r.text) + _raise_for_api_error(r.status_code, r.text) client_info = r.json() @@ -224,6 +248,36 @@ def _ping() -> Dict[str, Any]: ) +def _refresh_oidc_from_cache( + issuer: str, + client_id: str, + refresh_token: str, + scope: Optional[str], + audience: Optional[str] = None, +) -> Optional[Dict[str, Any]]: + """Exchange a cached refresh token for new tokens, persisting the result. + + Shared by the ``init()`` path (``_obtain_oidc_tokens``) and the request path + (``_ensure_fresh_oidc_token``) so their refresh behaviour can't drift. On + success the refreshed entry is cached and returned. On a failed refresh (e.g. + revoked token / ``invalid_grant``) the stale entry is deleted and None is + returned, so a dead refresh token isn't re-attempted; the caller decides how + to recover (``init()`` falls back to device flow, the request path stops). + """ + from .credentials_store import delete_cached_entry, upsert_cached_entry + from .oidc_device import try_refresh_cached_tokens + + try: + new_tokens = try_refresh_cached_tokens( + issuer, client_id, refresh_token, scope, audience=audience + ) + except ValidMindAuthError: + delete_cached_entry(issuer, client_id, audience=audience) + return None + upsert_cached_entry(issuer, client_id, new_tokens, audience=audience) + return new_tokens + + def _obtain_oidc_tokens( issuer: str, client_id: str, @@ -232,14 +286,13 @@ def _obtain_oidc_tokens( ) -> Dict[str, Any]: """Return a credentials entry dict with access_token, expires_at, refresh_token, etc.""" from .credentials_store import ( - delete_cached_entry, get_cached_entry, is_expired, normalize_client_id, normalize_issuer, upsert_cached_entry, ) - from .oidc_device import run_device_flow, try_refresh_cached_tokens + from .oidc_device import run_device_flow norm_issuer = normalize_issuer(issuer) norm_client_id = normalize_client_id(client_id) @@ -247,20 +300,11 @@ def _obtain_oidc_tokens( if cached and not is_expired(cached): return cached if cached and cached.get("refresh_token"): - try: - new_tokens = try_refresh_cached_tokens( - norm_issuer, - norm_client_id, - cached["refresh_token"], - scope, - audience=audience, - ) - upsert_cached_entry( - norm_issuer, norm_client_id, new_tokens, audience=audience - ) + new_tokens = _refresh_oidc_from_cache( + norm_issuer, norm_client_id, cached["refresh_token"], scope, audience + ) + if new_tokens is not None: return new_tokens - except ValidMindAuthError: - delete_cached_entry(norm_issuer, norm_client_id, audience=audience) tokens = run_device_flow(norm_issuer, norm_client_id, scope, audience=audience) upsert_cached_entry(norm_issuer, norm_client_id, tokens, audience=audience) return tokens @@ -276,6 +320,100 @@ def _select_oidc_bearer_token(entry: Dict[str, Any]) -> str: return entry["access_token"] +def _set_oidc_access_token(entry: Dict[str, Any]) -> None: + """Adopt the bearer token from a credentials entry as the active token. + + Records its expiry for cheap request-path staleness checks and drops the + pooled aiohttp session so the next request rebuilds it with the new token. + """ + global _access_token, _oidc_expires_at + _access_token = _select_oidc_bearer_token(entry) + _oidc_expires_at = entry.get("expires_at") + _invalidate_async_session() + + +def _clear_oidc_access_token() -> None: + """Drop the in-memory OIDC token and pooled session after an unrecoverable refresh. + + Pairs with deleting the cached entry: with no usable token in memory, the next + request fails fast at header build with the clear "run vm.init()" error instead + of sending a doomed request just to get a 401 back. + """ + global _access_token, _oidc_expires_at + _access_token = None + _oidc_expires_at = None + _invalidate_async_session() + + +def _oidc_token_is_stale() -> bool: + """Cheap in-memory expiry check (no file I/O) using the 120s skew in is_expired.""" + from .credentials_store import is_expired + + return is_expired({"expires_at": _oidc_expires_at}) + + +def _ensure_fresh_oidc_token(force: bool = False) -> bool: + """Refresh the OIDC access token on the request path when it has expired. + + The token captured at ``init()`` is otherwise reused unchanged for the life of + the process, so a session that outlives the provider's access-token lifetime + starts failing on every call until ``init()`` is re-run. This uses the cached + refresh token (requires ``offline_access``, the default scope) to obtain a new + access token in place. Returns True when a usable token is in place, False when + no refresh was possible (non-OIDC mode, no cached refresh token, or a failed + refresh — in which case the cached entry is dropped, matching init()). + """ + if _auth_mode != "oidc" or _oidc_login_context is None: + return False + # Hot path: skip file I/O and the lock while the current token is still valid. + if not force and not _oidc_token_is_stale(): + return True + + from .credentials_store import get_cached_entry, is_expired + + ctx = _oidc_login_context + issuer = ctx["issuer"] + client_id = ctx["client_id"] + scope = ctx.get("scope") + audience = ctx.get("audience") or None + + with _oidc_refresh_lock: + entry = get_cached_entry(issuer, client_id, audience=audience) + if entry is None: + return False + # Another caller may have refreshed the token while we waited on the lock. + if not force and not is_expired(entry): + _set_oidc_access_token(entry) + return True + refresh_token = entry.get("refresh_token") + if not refresh_token: + return False + # Shared with init(): on failure the cached entry is dropped so a + # bad/revoked refresh token isn't re-attempted on every request. Recovery + # is re-running vm.init() (surfaced by the clear 401 auth error); unlike + # init() the request path can't fall back to interactive device flow. + new_tokens = _refresh_oidc_from_cache( + issuer, client_id, refresh_token, scope, audience + ) + if new_tokens is None: + # Hard failure: the cache entry was dropped; clear the in-memory token + # too so the next request fails fast with the clear re-auth message. + _clear_oidc_access_token() + return False + _set_oidc_access_token(new_tokens) + return True + + +def _raise_for_api_error(status: int, text: str) -> None: + """Raise a clear, actionable auth error for an OIDC 401, else defer to raise_api_error.""" + if status == 401 and _auth_mode == "oidc": + raise ValidMindAuthError( + "ValidMind rejected the OAuth access token (HTTP 401); it may have " + "expired or been revoked. Run vm.init() again to re-authenticate." + ) + raise_api_error(text) + + def init( api_key: Optional[str] = None, api_secret: Optional[str] = None, @@ -330,7 +468,7 @@ def init( ValidMindAuthError: If OIDC configuration conflicts or login fails. """ global _api_key, _api_secret, _api_host, _model_cuid, _monitoring, _document - global _auth_mode, _access_token, _oidc_login_context + global _auth_mode, _access_token, _oidc_login_context, _oidc_expires_at if api_key == "...": # special case to detect when running a notebook placeholder (...) @@ -402,20 +540,21 @@ def init( entry = _obtain_oidc_tokens( oidc_issuer, oidc_client_id, scope_val, audience=oidc_audience_opt ) - _access_token = _select_oidc_bearer_token(entry) _oidc_login_context = { "issuer": entry["issuer"], "client_id": entry["client_id"], "scope": scope_val, "audience": entry.get("audience") or oidc_audience_val, } - _invalidate_async_session() + # Sets _access_token / _oidc_expires_at and invalidates the async session. + _set_oidc_access_token(entry) else: if env_key is None or env_secret is None: raise MissingAPICredentialsError() _auth_mode = "api_key" _access_token = None _oidc_login_context = None + _oidc_expires_at = None _api_key = env_key _api_secret = env_secret _api_host = resolved_host @@ -656,6 +795,7 @@ def _validate_log_text_context( def generate_qualitative_text(text_generation_data: Dict[str, Any]) -> Dict[str, Any]: """Generate qualitative text using the ValidMind AI API.""" + _ensure_fresh_oidc_token() r = requests.post( url=_get_url("ai/generate/qualitative_text_generation"), headers=_get_api_headers(), @@ -663,7 +803,7 @@ def generate_qualitative_text(text_generation_data: Dict[str, Any]) -> Dict[str, ) if r.status_code != 200: - raise_api_error(r.text) + _raise_for_api_error(r.status_code, r.text) return r.json() @@ -901,6 +1041,7 @@ def log_metric( def generate_test_result_description(test_result_data: Dict[str, Any]) -> str: + _ensure_fresh_oidc_token() r = requests.post( url=_get_url("ai/generate/test_result_description"), headers=_get_api_headers(), @@ -908,6 +1049,6 @@ def generate_test_result_description(test_result_data: Dict[str, Any]) -> str: ) if r.status_code != 200: - raise_api_error(r.text) + _raise_for_api_error(r.status_code, r.text) return r.json()