futurehouse-client 0.3.18.dev186__tar.gz → 0.3.18.dev195__tar.gz

{futurehouse_client-0.3.18.dev186 → futurehouse_client-0.3.18.dev195}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.3.18.dev186
+Version: 0.3.18.dev195
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 Classifier: Operating System :: OS Independent

{futurehouse_client-0.3.18.dev186 → futurehouse_client-0.3.18.dev195}/futurehouse_client/clients/rest_client.py

@@ -15,7 +15,7 @@ import uuid
 from collections.abc import Collection
 from pathlib import Path
 from types import ModuleType
-from typing import Any, ClassVar, assert_never, cast
+from typing import Any, ClassVar, cast
 from uuid import UUID
 
 import cloudpickle
@@ -45,7 +45,6 @@ from tqdm.asyncio import tqdm
 
 from futurehouse_client.clients import JobNames
 from futurehouse_client.models.app import (
-    APIKeyPayload,
     AuthType,
     JobDeploymentConfig,
     PQATaskResponse,
@@ -55,11 +54,7 @@ from futurehouse_client.models.app import (
     TaskResponseVerbose,
 )
 from futurehouse_client.models.rest import ExecutionStatus
-from futurehouse_client.utils.auth import (
-    AUTH_ERRORS_TO_RETRY_ON,
-    AuthError,
-    refresh_token_on_auth_error,
-)
+from futurehouse_client.utils.auth import RefreshingJWT
 from futurehouse_client.utils.general import gather_with_concurrency
 from futurehouse_client.utils.module_utils import (
     OrganizationSelector,
@@ -128,8 +123,6 @@ retry_if_connection_error = retry_if_exception_type((
     FileUploadError,
 ))
 
-# 5 minute default for JWTs
-JWT_TOKEN_CACHE_EXPIRY: int = 300  # seconds
 DEFAULT_AGENT_TIMEOUT: int = 2400  # seconds
 
 
@@ -163,69 +156,85 @@ class RestClient:
         self.api_key = api_key
         self._clients: dict[str, Client | AsyncClient] = {}
         self.headers = headers or {}
-        self.auth_jwt = self._run_auth(jwt=jwt)
+        self.jwt = jwt
         self.organizations: list[str] = self._filter_orgs(organization)
 
     @property
     def client(self) -> Client:
-        """Lazily initialized and cached HTTP client with authentication."""
-        return cast(Client, self.get_client("application/json", with_auth=True))
+        """Authenticated HTTP client for regular API calls."""
+        return cast(Client, self.get_client("application/json", authenticated=True))
 
     @property
     def async_client(self) -> AsyncClient:
-        """Lazily initialized and cached HTTP client with authentication."""
+        """Authenticated async HTTP client for regular API calls."""
         return cast(
             AsyncClient,
-            self.get_client("application/json", with_auth=True, with_async=True),
+            self.get_client("application/json", authenticated=True, async_client=True),
         )
 
     @property
-    def auth_client(self) -> Client:
-        """Lazily initialized and cached HTTP client without authentication."""
-        return cast(Client, self.get_client("application/json", with_auth=False))
+    def unauthenticated_client(self) -> Client:
+        """Unauthenticated HTTP client for auth operations to avoid recursion."""
+        return cast(Client, self.get_client("application/json", authenticated=False))
 
     @property
     def multipart_client(self) -> Client:
-        """Lazily initialized and cached HTTP client for multipart uploads."""
-        return cast(Client, self.get_client(None, with_auth=True))
+        """Authenticated HTTP client for multipart uploads."""
+        return cast(Client, self.get_client(None, authenticated=True))
 
     def get_client(
         self,
         content_type: str | None = "application/json",
-        with_auth: bool = True,
-        with_async: bool = False,
+        authenticated: bool = True,
+        async_client: bool = False,
     ) -> Client | AsyncClient:
         """Return a cached HTTP client or create one if needed.
 
         Args:
             content_type: The desired content type header. Use None for multipart uploads.
-            with_auth: Whether the client should include an Authorization header.
-            with_async: Whether to use an async client.
+            authenticated: Whether the client should include authentication.
+            async_client: Whether to use an async client.
 
         Returns:
             An HTTP client configured with the appropriate headers.
         """
-        # Create a composite key based on content type and auth flag.
-        key = f"{content_type or 'multipart'}_{with_auth}_{with_async}"
+        # Create a composite key based on content type and auth flag
+        key = f"{content_type or 'multipart'}_{authenticated}_{async_client}"
+
         if key not in self._clients:
             headers = copy.deepcopy(self.headers)
-            if with_auth:
-                headers["Authorization"] = f"Bearer {self.auth_jwt}"
+            auth = None
+
+            if authenticated:
+                auth = RefreshingJWT(
+                    # authenticated=False will always return a synchronous client
+                    auth_client=cast(
+                        Client, self.get_client("application/json", authenticated=False)
+                    ),
+                    auth_type=self.auth_type,
+                    api_key=self.api_key,
+                    jwt=self.jwt,
+                )
+
             if content_type:
                 headers["Content-Type"] = content_type
+
             self._clients[key] = (
                 AsyncClient(
                     base_url=self.base_url,
                     headers=headers,
                     timeout=self.REQUEST_TIMEOUT,
+                    auth=auth,
                 )
-                if with_async
+                if async_client
                 else Client(
                     base_url=self.base_url,
                     headers=headers,
                     timeout=self.REQUEST_TIMEOUT,
+                    auth=auth,
                 )
             )
+
         return self._clients[key]
 
     def close(self):
@@ -255,32 +264,6 @@ class RestClient:
             raise ValueError(f"Organization '{organization}' not found.")
         return filtered_orgs
 
-    def _run_auth(self, jwt: str | None = None) -> str:
-        auth_payload: APIKeyPayload | None
-        if self.auth_type == AuthType.API_KEY:
-            auth_payload = APIKeyPayload(api_key=self.api_key)
-        elif self.auth_type == AuthType.JWT:
-            auth_payload = None
-        else:
-            assert_never(self.auth_type)
-        try:
-            # Use the unauthenticated client for login
-            if auth_payload:
-                response = self.auth_client.post(
-                    "/auth/login", json=auth_payload.model_dump()
-                )
-                response.raise_for_status()
-                token_data = response.json()
-            elif jwt:
-                token_data = {"access_token": jwt, "expires_in": JWT_TOKEN_CACHE_EXPIRY}
-            else:
-                raise ValueError("JWT token required for JWT authentication.")
-
-            return token_data["access_token"]
-        except Exception as e:
-            raise RestClientError(f"Error authenticating: {e!s}") from e
-
-    @refresh_token_on_auth_error()
     def _check_job(self, name: str, organization: str) -> dict[str, Any]:
         try:
             response = self.client.get(
@@ -288,19 +271,9 @@ class RestClient:
             )
             response.raise_for_status()
             return response.json()
-        except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
-            raise
         except Exception as e:
             raise JobFetchError(f"Error checking job: {e!s}") from e
 
-    @refresh_token_on_auth_error()
     def _fetch_my_orgs(self) -> list[str]:
         response = self.client.get(f"/v0.1/organizations?filter={True}")
         response.raise_for_status()
@@ -358,7 +331,6 @@ class RestClient:
         if not files:
             raise TaskFetchError(f"No files found in {path}")
 
-    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -400,19 +372,9 @@ class RestClient:
             ):
                 return PQATaskResponse(**data)
             return TaskResponse(**data)
-        except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
-            raise
         except Exception as e:
             raise TaskFetchError(f"Error getting task: {e!s}") from e
 
-    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -456,19 +418,9 @@ class RestClient:
             ):
                 return PQATaskResponse(**data)
             return TaskResponse(**data)
-        except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
-            raise
         except Exception as e:
             raise TaskFetchError(f"Error getting task: {e!s}") from e
 
-    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -492,20 +444,10 @@ class RestClient:
             response.raise_for_status()
             trajectory_id = response.json()["trajectory_id"]
             self.trajectory_id = trajectory_id
-        except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
-            raise
         except Exception as e:
             raise TaskFetchError(f"Error creating task: {e!s}") from e
         return trajectory_id
 
-    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -529,15 +471,6 @@ class RestClient:
             response.raise_for_status()
             trajectory_id = response.json()["trajectory_id"]
             self.trajectory_id = trajectory_id
-        except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
-            raise
         except Exception as e:
             raise TaskFetchError(f"Error creating task: {e!s}") from e
         return trajectory_id
@@ -683,7 +616,6 @@ class RestClient:
             for task_id in trajectory_ids
         ]
 
-    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -695,19 +627,11 @@ class RestClient:
             build_id = build_id or self.build_id
             response = self.client.get(f"/v0.1/builds/{build_id}")
             response.raise_for_status()
-        except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
-            raise
+        except Exception as e:
+            raise JobFetchError(f"Error getting build status: {e!s}") from e
         return response.json()
 
     # TODO: Refactor later so we don't have to ignore PLR0915
-    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -887,13 +811,6 @@ class RestClient:
                 build_context = response.json()
                 self.build_id = build_context["build_id"]
             except HTTPStatusError as e:
-                if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                    raise AuthError(
-                        e.response.status_code,
-                        f"Authentication failed: {e}",
-                        request=e.request,
-                        response=e.response,
-                    ) from e
                 error_detail = response.json()
                 error_message = error_detail.get("detail", str(e))
                 raise JobCreationError(
@@ -974,7 +891,6 @@ class RestClient:
         except Exception as e:
             raise FileUploadError(f"Error uploading directory {dir_path}: {e}") from e
 
-    @refresh_token_on_auth_error()
     def _upload_single_file(
         self,
         job_name: str,
@@ -1048,20 +964,10 @@ class RestClient:
                         )
 
             logger.info(f"Successfully uploaded {file_name}")
-        except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
-            raise
         except Exception as e:
             logger.exception(f"Error uploading file {file_path}")
             raise FileUploadError(f"Error uploading file {file_path}: {e}") from e
 
-    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -1098,13 +1004,6 @@ class RestClient:
             response.raise_for_status()
             return response.json()
         except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
             logger.exception(
                 f"Error listing files for job {job_name}, trajectory {trajectory_id}, upload_id {upload_id}: {e.response.text}"
             )
@@ -1117,7 +1016,6 @@ class RestClient:
             )
             raise RestClientError(f"Error listing files: {e!s}") from e
 
-    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -1165,13 +1063,6 @@ class RestClient:
 
             logger.info(f"File {file_path} downloaded to {destination_path}")
         except HTTPStatusError as e:
-            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
-                raise AuthError(
-                    e.response.status_code,
-                    f"Authentication failed: {e}",
-                    request=e.request,
-                    response=e.response,
-                ) from e
             logger.exception(
                 f"Error downloading file {file_path} for job {job_name}, trajectory_id {trajectory_id}: {e.response.text}"
             )

futurehouse_client-0.3.18.dev195/futurehouse_client/utils/auth.py

@@ -0,0 +1,89 @@
+import logging
+from typing import ClassVar, Final
+
+import httpx
+
+from futurehouse_client.models.app import APIKeyPayload, AuthType
+
+logger = logging.getLogger(__name__)
+
+INVALID_REFRESH_TYPE_MSG: Final[str] = (
+    "API key auth is required to refresh auth tokens."
+)
+JWT_TOKEN_CACHE_EXPIRY: int = 300  # seconds
+
+
+def _run_auth(
+    client: httpx.Client,
+    auth_type: AuthType = AuthType.API_KEY,
+    api_key: str | None = None,
+    jwt: str | None = None,
+) -> str:
+    auth_payload: APIKeyPayload | None
+    if auth_type == AuthType.API_KEY:
+        auth_payload = APIKeyPayload(api_key=api_key)
+    elif auth_type == AuthType.JWT:
+        auth_payload = None
+    try:
+        if auth_payload:
+            response = client.post("/auth/login", json=auth_payload.model_dump())
+            response.raise_for_status()
+            token_data = response.json()
+        elif jwt:
+            token_data = {"access_token": jwt, "expires_in": JWT_TOKEN_CACHE_EXPIRY}
+        else:
+            raise ValueError("JWT token required for JWT authentication.")
+
+        return token_data["access_token"]
+    except Exception as e:
+        raise Exception("Failed to authenticate") from e  # noqa: TRY002
+
+
+class RefreshingJWT(httpx.Auth):
+    """Automatically (re-)inject a JWT and transparently retry exactly once when we hit a 401/403."""
+
+    RETRY_STATUSES: ClassVar[set[int]] = {
+        httpx.codes.UNAUTHORIZED,
+        httpx.codes.FORBIDDEN,
+    }
+
+    def __init__(
+        self,
+        auth_client: httpx.Client,
+        auth_type: AuthType = AuthType.API_KEY,
+        api_key: str | None = None,
+        jwt: str | None = None,
+    ):
+        self.auth_type = auth_type
+        self.auth_client = auth_client
+        self.api_key = api_key
+        self._jwt = _run_auth(
+            client=auth_client,
+            jwt=jwt,
+            auth_type=auth_type,
+            api_key=api_key,
+        )
+
+    def refresh_token(self):
+        if self.auth_type == AuthType.JWT:
+            logger.error(INVALID_REFRESH_TYPE_MSG)
+            raise ValueError(INVALID_REFRESH_TYPE_MSG)
+        self._jwt = _run_auth(
+            client=self.auth_client,
+            auth_type=self.auth_type,
+            api_key=self.api_key,
+        )
+
+    def auth_flow(self, request):
+        request.headers["Authorization"] = f"Bearer {self._jwt}"
+        response = yield request
+
+        # If it failed, refresh once and replay the request
+        if response.status_code in self.RETRY_STATUSES:
+            logger.info(
+                "Received %s, refreshing token and retrying …",
+                response.status_code,
+            )
+            self.refresh_token()
+            request.headers["Authorization"] = f"Bearer {self._jwt}"
+            yield request  # second (and final) attempt, again or use a while loop

{futurehouse_client-0.3.18.dev186 → futurehouse_client-0.3.18.dev195}/futurehouse_client.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.3.18.dev186
+Version: 0.3.18.dev195
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 Classifier: Operating System :: OS Independent

{futurehouse_client-0.3.18.dev186 → futurehouse_client-0.3.18.dev195}/tests/test_rest.py

@@ -2,7 +2,6 @@
 import asyncio
 import os
 import time
-from unittest.mock import patch
 
 import pytest
 from futurehouse_client.clients import (
@@ -10,10 +9,9 @@ from futurehouse_client.clients import (
     PQATaskResponse,
     TaskResponseVerbose,
 )
-from futurehouse_client.clients.rest_client import RestClient
+from futurehouse_client.clients.rest_client import RestClient, TaskFetchError
 from futurehouse_client.models.app import Stage, TaskRequest
 from futurehouse_client.models.rest import ExecutionStatus
-from futurehouse_client.utils.auth import AuthError, refresh_token_on_auth_error
 from pytest_subtests import SubTests
 
 ADMIN_API_KEY = os.environ["PLAYWRIGHT_ADMIN_API_KEY"]
@@ -40,7 +38,7 @@ def pub_client():
 
 
 @pytest.fixture
-def task_data():
+def task_req():
     """Create a sample task request."""
     return TaskRequest(
         name=JobNames.from_string("dummy"),
@@ -49,7 +47,7 @@ def task_data():
 
 
 @pytest.fixture
-def pqa_task_data():
+def pqa_task_req():
     return TaskRequest(
         name=JobNames.from_string("crow"),
         query="How many moons does earth have?",
@@ -58,40 +56,40 @@ def pqa_task_data():
 
 @pytest.mark.timeout(300)
 @pytest.mark.flaky(reruns=3)
-def test_futurehouse_dummy_env_crow(admin_client: RestClient, task_data: TaskRequest):
-    admin_client.create_task(task_data)
+def test_futurehouse_dummy_env_crow(admin_client: RestClient, task_req: TaskRequest):
+    admin_client.create_task(task_req)
     while (task_status := admin_client.get_task().status) in {"queued", "in progress"}:
         time.sleep(5)
     assert task_status == "success"
 
 
 def test_insufficient_permissions_request(
-    pub_client: RestClient, task_data: TaskRequest
+    pub_client: RestClient, task_req: TaskRequest
 ):
     # Create a new instance so that cached credentials aren't reused
-    with pytest.raises(AuthError) as exc_info:
-        pub_client.create_task(task_data)
+    with pytest.raises(TaskFetchError) as exc_info:
+        pub_client.create_task(task_req)
 
-    assert "403 Forbidden" in str(exc_info.value)
+    assert "Error creating task" in str(exc_info.value)
 
 
 @pytest.mark.timeout(300)
 @pytest.mark.asyncio
 async def test_job_response(  # noqa: PLR0915
-    subtests: SubTests, admin_client: RestClient, pqa_task_data: TaskRequest
+    subtests: SubTests, admin_client: RestClient, pqa_task_req: TaskRequest
 ):
-    task_id = admin_client.create_task(pqa_task_data)
-    atask_id = await admin_client.acreate_task(pqa_task_data)
+    task_id = admin_client.create_task(pqa_task_req)
+    atask_id = await admin_client.acreate_task(pqa_task_req)
 
     with subtests.test("Test TaskResponse with queued task"):
         task_response = admin_client.get_task(task_id)
         assert task_response.status in {"queued", "in progress"}
-        assert task_response.job_name == pqa_task_data.name
-        assert task_response.query == pqa_task_data.query
+        assert task_response.job_name == pqa_task_req.name
+        assert task_response.query == pqa_task_req.query
         task_response = await admin_client.aget_task(atask_id)
         assert task_response.status in {"queued", "in progress"}
-        assert task_response.job_name == pqa_task_data.name
-        assert task_response.query == pqa_task_data.query
+        assert task_response.job_name == pqa_task_req.name
+        assert task_response.query == pqa_task_req.query
 
     for _ in range(TEST_MAX_POLLS):
         task_response = admin_client.get_task(task_id)
@@ -111,8 +109,8 @@ async def test_job_response(  # noqa: PLR0915
         # assert it has general fields
         assert task_response.status == "success"
         assert task_response.task_id is not None
-        assert pqa_task_data.name in task_response.job_name
-        assert pqa_task_data.query in task_response.query
+        assert pqa_task_req.name in task_response.job_name
+        assert pqa_task_req.query in task_response.query
         # assert it has PQA specific fields
         assert task_response.answer is not None
         # assert it's not verbose
@@ -125,8 +123,8 @@ async def test_job_response(  # noqa: PLR0915
         # assert it has general fields
         assert task_response.status == "success"
         assert task_response.task_id is not None
-        assert pqa_task_data.name in task_response.job_name
-        assert pqa_task_data.query in task_response.query
+        assert pqa_task_req.name in task_response.job_name
+        assert pqa_task_req.query in task_response.query
         # assert it has PQA specific fields
         assert task_response.answer is not None
         # assert it's not verbose
@@ -151,9 +149,9 @@ async def test_job_response(  # noqa: PLR0915
 @pytest.mark.timeout(300)
 @pytest.mark.flaky(reruns=3)
 def test_run_until_done_futurehouse_dummy_env_crow(
-    admin_client: RestClient, task_data: TaskRequest
+    admin_client: RestClient, task_req: TaskRequest
 ):
-    tasks_to_do = [task_data, task_data]
+    tasks_to_do = [task_req, task_req]
 
     results = admin_client.run_tasks_until_done(tasks_to_do)
 
@@ -165,9 +163,9 @@ def test_run_until_done_futurehouse_dummy_env_crow(
 @pytest.mark.flaky(reruns=3)
 @pytest.mark.asyncio
 async def test_arun_until_done_futurehouse_dummy_env_crow(
-    admin_client: RestClient, task_data: TaskRequest
+    admin_client: RestClient, task_req: TaskRequest
 ):
-    tasks_to_do = [task_data, task_data]
+    tasks_to_do = [task_req, task_req]
 
     results = await admin_client.arun_tasks_until_done(tasks_to_do)
 
@@ -179,9 +177,9 @@ async def test_arun_until_done_futurehouse_dummy_env_crow(
 @pytest.mark.flaky(reruns=3)
 @pytest.mark.asyncio
 async def test_timeout_run_until_done_futurehouse_dummy_env_crow(
-    admin_client: RestClient, task_data: TaskRequest
+    admin_client: RestClient, task_req: TaskRequest
 ):
-    tasks_to_do = [task_data, task_data]
+    tasks_to_do = [task_req, task_req]
 
     results = await admin_client.arun_tasks_until_done(
         tasks_to_do, verbose=True, timeout=5, progress_bar=True
@@ -202,59 +200,3 @@ async def test_timeout_run_until_done_futurehouse_dummy_env_crow(
     assert all(not isinstance(task, PQATaskResponse) for task in results), (
         "Should be verbose."
     )
-
-
-def test_auth_refresh_flow(admin_client: RestClient):
-    refresh_calls = 0
-    func_calls = 0
-
-    def mock_run_auth(*args, **kwargs):
-        nonlocal refresh_calls
-        refresh_calls += 1
-        return f"fresh-token-{refresh_calls}"
-
-    @refresh_token_on_auth_error()
-    def test_func(self, *args):
-        nonlocal func_calls
-        func_calls += 1
-
-        if func_calls == 1:
-            raise AuthError(401, "Auth failed", None, None)
-        return "success"
-
-    with patch.object(admin_client, "_run_auth", mock_run_auth):
-        result = test_func(admin_client)
-
-        assert result == "success"
-        assert func_calls == 2, "Function should be called twice"
-        assert refresh_calls == 1, "Auth should be refreshed once"
-        assert admin_client.auth_jwt == "fresh-token-1"
-
-
-@pytest.mark.asyncio
-async def test_async_auth_refresh_flow(admin_client: RestClient):
-    refresh_calls = 0
-    func_calls = 0
-
-    def mock_run_auth(*args, **kwargs):
-        nonlocal refresh_calls
-        refresh_calls += 1
-        return f"fresh-token-{refresh_calls}"
-
-    @refresh_token_on_auth_error()
-    async def test_async_func(self, *args):
-        nonlocal func_calls
-        func_calls += 1
-
-        if func_calls == 1:
-            raise AuthError(401, "Auth failed", None, None)
-        await asyncio.sleep(1)
-        return "success"
-
-    with patch.object(admin_client, "_run_auth", mock_run_auth):
-        result = await test_async_func(admin_client)
-
-        assert result == "success"
-        assert func_calls == 2, "Function should be called twice"
-        assert refresh_calls == 1, "Auth should be refreshed once"
-        assert admin_client.auth_jwt == "fresh-token-1"

futurehouse_client-0.3.18.dev186/futurehouse_client/utils/auth.py

@@ -1,107 +0,0 @@
-import asyncio
-import logging
-from collections.abc import Callable, Coroutine
-from functools import wraps
-from typing import Any, Final, Optional, ParamSpec, TypeVar, overload
-
-import httpx
-from httpx import HTTPStatusError
-
-logger = logging.getLogger(__name__)
-
-T = TypeVar("T")
-P = ParamSpec("P")
-
-AUTH_ERRORS_TO_RETRY_ON: Final[set[int]] = {
-    httpx.codes.UNAUTHORIZED,
-    httpx.codes.FORBIDDEN,
-}
-
-
-class AuthError(Exception):
-    """Raised when authentication fails with 401/403 status."""
-
-    def __init__(self, status_code: int, message: str, request=None, response=None):
-        self.status_code = status_code
-        self.request = request
-        self.response = response
-        super().__init__(message)
-
-
-def is_auth_error(e: Exception) -> bool:
-    if isinstance(e, AuthError):
-        return True
-    if isinstance(e, HTTPStatusError):
-        return e.response.status_code in AUTH_ERRORS_TO_RETRY_ON
-    return False
-
-
-def get_status_code(e: Exception) -> Optional[int]:
-    if isinstance(e, AuthError):
-        return e.status_code
-    if isinstance(e, HTTPStatusError):
-        return e.response.status_code
-    return None
-
-
-@overload
-def refresh_token_on_auth_error(
-    func: Callable[P, Coroutine[Any, Any, T]],
-) -> Callable[P, Coroutine[Any, Any, T]]: ...
-
-
-@overload
-def refresh_token_on_auth_error(
-    func: None = None, *, max_retries: int = ...
-) -> Callable[[Callable[P, T]], Callable[P, T]]: ...
-
-
-def refresh_token_on_auth_error(func=None, max_retries=1):
-    """Decorator that refreshes JWT token on 401/403 auth errors."""
-
-    def decorator(fn):
-        @wraps(fn)
-        def sync_wrapper(self, *args, **kwargs):
-            retries = 0
-            while True:
-                try:
-                    return fn(self, *args, **kwargs)
-                except Exception as e:
-                    if is_auth_error(e) and retries < max_retries:
-                        retries += 1
-                        status = get_status_code(e) or "Unknown"
-                        logger.info(
-                            f"Received auth error {status}, "
-                            f"refreshing token and retrying (attempt {retries}/{max_retries})..."
-                        )
-                        self.auth_jwt = self._run_auth()
-                        self._clients = {}
-                        continue
-                    raise
-
-        @wraps(fn)
-        async def async_wrapper(self, *args, **kwargs):
-            retries = 0
-            while True:
-                try:
-                    return await fn(self, *args, **kwargs)
-                except Exception as e:
-                    if is_auth_error(e) and retries < max_retries:
-                        retries += 1
-                        status = get_status_code(e) or "Unknown"
-                        logger.info(
-                            f"Received auth error {status}, "
-                            f"refreshing token and retrying (attempt {retries}/{max_retries})..."
-                        )
-                        self.auth_jwt = self._run_auth()
-                        self._clients = {}
-                        continue
-                    raise
-
-        if asyncio.iscoroutinefunction(fn):
-            return async_wrapper
-        return sync_wrapper
-
-    if callable(func):
-        return decorator(func)
-    return decorator