From 42e34151964bdd738d243655ba369fd5b16611f6 Mon Sep 17 00:00:00 2001 From: xodn348 Date: Mon, 11 May 2026 05:39:45 +0000 Subject: [PATCH] fix(client): raise auth exception synchronously in Client.init() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, _start_auth_task() fired authentication off in a daemon thread without joining, so any InvalidApiKeyException from fetch_auth_token was silently swallowed and Client.init() always set _initialized=True regardless of whether the key was valid. - v3.fetch_auth_token now raises InvalidApiKeyException instead of returning None when the server returns no token or the request fails outright. - _fetch_auth_async propagates exceptions rather than catching them. - _start_auth_task in synchronous contexts (no running event loop) joins the auth thread and re-raises, surfacing failures immediately to the caller. - Update test mocks to use AsyncMock for the async fetch_auth_token method, and fix the wrong patch target in test_session.py (api → client module). Closes #1336 --- agentops/client/api/versions/v3.py | 58 ++++++++------- agentops/client/client.py | 70 ++++++++++--------- tests/integration/test_auth_flow.py | 11 +-- tests/integration/test_session_concurrency.py | 6 +- tests/unit/test_session.py | 14 ++-- 5 files changed, 86 insertions(+), 73 deletions(-) diff --git a/agentops/client/api/versions/v3.py b/agentops/client/api/versions/v3.py index 0c8a81937..ee1d95a5a 100644 --- a/agentops/client/api/versions/v3.py +++ b/agentops/client/api/versions/v3.py @@ -7,6 +7,7 @@ from agentops.client.api.base import BaseApiClient from agentops.client.api.types import AuthTokenResponse from agentops.client.http.http_client import HttpClient +from agentops.exceptions import InvalidApiKeyException from agentops.logging import logger from termcolor import colored @@ -32,38 +33,35 @@ async def fetch_auth_token(self, api_key: str) -> AuthTokenResponse: api_key: The API key to authenticate with Returns: - AuthTokenResponse containing token and project information, or None if failed - """ - try: - path = "/v3/auth/token" - data = {"api_key": api_key} - headers = self.prepare_headers() - - # Build full URL - url = self._get_full_url(path) - - # Make async request - response_data = await HttpClient.async_request( - method="POST", url=url, data=data, headers=headers, timeout=30 - ) + AuthTokenResponse containing token and project information - token = response_data.get("token") - if not token: - logger.warning("Authentication failed: Perhaps an invalid API key?") - return None - - # Check project premium status - if response_data.get("project_prem_status") != "pro": - logger.info( - colored( - "\x1b[34mYou're on the agentops free plan 🤔\x1b[0m", - "blue", - ) + Raises: + InvalidApiKeyException: If the API key is invalid or authentication fails + """ + path = "/v3/auth/token" + data = {"api_key": api_key} + headers = self.prepare_headers() + + # Build full URL + url = self._get_full_url(path) + + # Make async request + response_data = await HttpClient.async_request( + method="POST", url=url, data=data, headers=headers, timeout=30 + ) + + if response_data is None or not response_data.get("token"): + raise InvalidApiKeyException(api_key, self.endpoint) + + # Check project premium status + if response_data.get("project_prem_status") != "pro": + logger.info( + colored( + "\x1b[34mYou're on the agentops free plan 🤔\x1b[0m", + "blue", ) + ) - return response_data - - except Exception: - return None + return response_data # Add V3-specific API methods here diff --git a/agentops/client/client.py b/agentops/client/client.py index 87d6afd3d..59c9504e1 100644 --- a/agentops/client/client.py +++ b/agentops/client/client.py @@ -94,55 +94,61 @@ def _set_auth_data(self, token: str, project_id: str): HttpClient.set_project_id(project_id) async def _fetch_auth_async(self, api_key: str) -> Optional[dict]: - """Asynchronously fetch authentication token.""" - try: - response = await self.api.v3.fetch_auth_token(api_key) - if response: - self._set_auth_data(response["token"], response["project_id"]) + """Asynchronously fetch authentication token. + + Raises: + InvalidApiKeyException: propagated from v3.fetch_auth_token on failure + """ + response = await self.api.v3.fetch_auth_token(api_key) + self._set_auth_data(response["token"], response["project_id"]) - # Update V4 client with token - self.api.v4.set_auth_token(response["token"]) + # Update V4 client with token + self.api.v4.set_auth_token(response["token"]) - # Update tracer config with real project ID - tracing_config = {"project_id": response["project_id"]} - tracer.update_config(tracing_config) + # Update tracer config with real project ID + tracing_config = {"project_id": response["project_id"]} + tracer.update_config(tracing_config) - logger.debug("Successfully fetched authentication token asynchronously") - return response - else: - logger.debug("Authentication failed - will continue without authentication") - return None - except Exception: - return None + logger.debug("Successfully fetched authentication token asynchronously") + return response def _start_auth_task(self, api_key: str): - """Start the async authentication task.""" + """Start the async authentication task. + + In synchronous contexts (no running event loop), blocks until auth completes + and re-raises any authentication exception so init() can surface it to the caller. + In async contexts, schedules a background task; the caller must handle the result. + """ if self._auth_task and not self._auth_task.done(): return # Task already running try: loop = asyncio.get_event_loop() if loop.is_running(): - # Use existing event loop + # Async context: schedule as a background task; exceptions won't propagate here self._auth_task = loop.create_task(self._fetch_auth_async(api_key)) - else: - # Create new event loop in background thread - def run_async_auth(): - asyncio.run(self._fetch_auth_async(api_key)) + return + except RuntimeError: + pass - import threading + # Synchronous context: run auth in a thread so asyncio.run() works, then join + # and re-raise so the caller (init()) surfaces the error immediately. + auth_exc: list = [] - auth_thread = threading.Thread(target=run_async_auth, daemon=True) - auth_thread.start() - except RuntimeError: - # Create new event loop in background thread - def run_async_auth(): + def run_async_auth(): + try: asyncio.run(self._fetch_auth_async(api_key)) + except Exception as exc: + auth_exc.append(exc) - import threading + auth_thread = threading.Thread(target=run_async_auth, daemon=True) + auth_thread.start() + auth_thread.join(timeout=35) # slightly longer than the HTTP timeout (30s) - auth_thread = threading.Thread(target=run_async_auth, daemon=True) - auth_thread.start() + if auth_thread.is_alive(): + raise TimeoutError("AgentOps authentication timed out after 35 seconds") + if auth_exc: + raise auth_exc[0] def init(self, **kwargs: Any) -> None: # Return type updated to None # Recreate the Config object to parse environment variables at the time of initialization diff --git a/tests/integration/test_auth_flow.py b/tests/integration/test_auth_flow.py index da8ccda80..91280fb58 100644 --- a/tests/integration/test_auth_flow.py +++ b/tests/integration/test_auth_flow.py @@ -1,5 +1,5 @@ import pytest -from unittest.mock import patch, MagicMock +from unittest.mock import patch, MagicMock, AsyncMock from agentops.client import Client from agentops.exceptions import InvalidApiKeyException, ApiServerException @@ -20,7 +20,9 @@ def test_auth_flow(mock_api_key): with patch("agentops.client.client.ApiClient") as mock_api_client: # Create mock API instance mock_api = MagicMock() - mock_api.v3.fetch_auth_token.return_value = {"token": "mock_token", "project_id": "mock_project_id"} + mock_api.v3.fetch_auth_token = AsyncMock( + return_value={"token": "mock_token", "project_id": "mock_project_id"} + ) mock_api_client.return_value = mock_api # Initialize the client @@ -38,15 +40,16 @@ def test_auth_flow(mock_api_key): @pytest.mark.vcr() def test_auth_flow_invalid_key(): - """Test authentication flow with invalid API key.""" + """Test authentication flow with invalid API key raises an exception.""" with patch("agentops.client.client.ApiClient") as mock_api_client: # Create mock API instance that raises an error mock_api = MagicMock() - mock_api.v3.fetch_auth_token.side_effect = ApiServerException("Invalid API key") + mock_api.v3.fetch_auth_token = AsyncMock(side_effect=ApiServerException("Invalid API key")) mock_api_client.return_value = mock_api client = Client() with pytest.raises((InvalidApiKeyException, ApiServerException)) as exc_info: client.init(api_key="invalid-key") + assert not client.initialized assert "Invalid API key" in str(exc_info.value) diff --git a/tests/integration/test_session_concurrency.py b/tests/integration/test_session_concurrency.py index 928dd972a..17b96a7da 100644 --- a/tests/integration/test_session_concurrency.py +++ b/tests/integration/test_session_concurrency.py @@ -1,6 +1,6 @@ import pytest import concurrent.futures -from unittest.mock import patch, MagicMock +from unittest.mock import patch, MagicMock, AsyncMock from fastapi import FastAPI from fastapi.testclient import TestClient import agentops @@ -37,7 +37,9 @@ def setup_agentops(mock_api_key): with patch("agentops.client.client.ApiClient") as mock_api_client: # Create mock API instance mock_api = MagicMock() - mock_api.v3.fetch_auth_token.return_value = {"token": "mock_token", "project_id": "mock_project_id"} + mock_api.v3.fetch_auth_token = AsyncMock( + return_value={"token": "mock_token", "project_id": "mock_project_id"} + ) mock_api_client.return_value = mock_api # Mock global tracer to avoid actual initialization diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index d334dae5c..34460c1a8 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -1,5 +1,5 @@ import pytest -from unittest.mock import patch, MagicMock +from unittest.mock import patch, MagicMock, AsyncMock # Tests for the new session management functionality # These tests call the actual public API but mock the underlying implementation @@ -25,10 +25,12 @@ def mock_tracing_core(): @pytest.fixture(scope="function") def mock_api_client(): """Mock the API client to avoid actual API calls""" - with patch("agentops.client.api.ApiClient") as mock_api: + with patch("agentops.client.client.ApiClient") as mock_api: # Configure the v3.fetch_auth_token method to return a valid response mock_v3 = MagicMock() - mock_v3.fetch_auth_token.return_value = {"token": "mock-jwt-token", "project_id": "mock-project-id"} + mock_v3.fetch_auth_token = AsyncMock( + return_value={"token": "mock-jwt-token", "project_id": "mock-project-id"} + ) mock_api.return_value.v3 = mock_v3 yield mock_api @@ -423,9 +425,11 @@ def test_session_management_integration(): mock_tracer.initialized = True # Mock API client - with patch("agentops.client.api.ApiClient") as mock_api: + with patch("agentops.client.client.ApiClient") as mock_api: mock_v3 = MagicMock() - mock_v3.fetch_auth_token.return_value = {"token": "mock-jwt-token", "project_id": "mock-project-id"} + mock_v3.fetch_auth_token = AsyncMock( + return_value={"token": "mock-jwt-token", "project_id": "mock-project-id"} + ) mock_api.return_value.v3 = mock_v3 # Initialize AgentOps