diff --git a/arcade/arcade/cli/utils.py b/arcade/arcade/cli/utils.py index a9fd983f..d2bf72d3 100644 --- a/arcade/arcade/cli/utils.py +++ b/arcade/arcade/cli/utils.py @@ -1,5 +1,4 @@ import importlib.util -import time from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Union @@ -14,6 +13,7 @@ from typer.core import TyperGroup from typer.models import Context from arcade.client.client import Arcade +from arcade.client.errors import APITimeoutError from arcade.client.schema import AuthResponse from arcade.core.catalog import ToolCatalog from arcade.core.config_model import Config @@ -364,8 +364,10 @@ def wait_for_authorization_completion(client: Arcade, tool_authorization: dict | auth_response = AuthResponse.model_validate(tool_authorization) while auth_response.status != "completed": - time.sleep(0.5) - auth_response = client.auth.status(auth_response) + try: + auth_response = client.auth.status(auth_response, wait=60) + except APITimeoutError: + continue def get_tool_authorization( diff --git a/arcade/arcade/client/base.py b/arcade/arcade/client/base.py index c88b1868..f9f3359e 100644 --- a/arcade/arcade/client/base.py +++ b/arcade/arcade/client/base.py @@ -5,6 +5,7 @@ import httpx from httpx import Timeout from arcade.client.errors import ( + APITimeoutError, BadRequestError, InternalServerError, NotFoundError, @@ -81,6 +82,7 @@ class BaseArcadeClient: 401: UnauthorizedError, 403: PermissionDeniedError, 404: NotFoundError, + 408: APITimeoutError, 429: RateLimitError, 500: InternalServerError, } diff --git a/arcade/arcade/client/client.py b/arcade/arcade/client/client.py index ba0b3296..123cb8d7 100644 --- a/arcade/arcade/client/client.py +++ b/arcade/arcade/client/client.py @@ -1,5 +1,6 @@ from typing import Any, TypeVar, Union +from httpx import Timeout from openai import AsyncOpenAI, OpenAI from openai.resources.chat import AsyncChat, Chat @@ -62,7 +63,10 @@ class AuthResource(BaseResource[ClientT]): return AuthResponse(**data) def status( - self, auth_id_or_response: Union[str, AuthResponse], scopes: list[str] | None = None + self, + auth_id_or_response: Union[str, AuthResponse], + scopes: list[str] | None = None, + wait: int | None = None, ) -> AuthResponse: """ Poll for the status of an authorization @@ -80,10 +84,28 @@ class AuthResource(BaseResource[ClientT]): else: auth_id = auth_id_or_response + # Calculate the new timeout based on the wait parameter + new_timeout = self._client._timeout + if wait is not None: + if isinstance(self._client._timeout, Timeout): + new_timeout = Timeout( + connect=self._client._timeout.connect, + read=(self._client._timeout.read or 0) + wait, + write=self._client._timeout.write, + pool=self._client._timeout.pool, + ) + else: + new_timeout = self._client._timeout + wait + data = self._client._execute_request( # type: ignore[attr-defined] "GET", f"{self._resource_path}/status", - params={"authorizationId": auth_id, "scopes": " ".join(scopes) if scopes else None}, + params={ + "authorizationId": auth_id, + "scopes": " ".join(scopes) if scopes else None, + "wait": wait, + }, + timeout=new_timeout, ) return AuthResponse(**data) @@ -216,7 +238,10 @@ class AsyncAuthResource(BaseResource[AsyncArcadeClient]): return AuthResponse(**data) async def status( - self, auth_id_or_response: Union[str, AuthResponse], scopes: list[str] | None = None + self, + auth_id_or_response: Union[str, AuthResponse], + scopes: list[str] | None = None, + wait: int | None = None, ) -> AuthResponse: """ Poll for the status of an authorization asynchronously @@ -234,10 +259,24 @@ class AsyncAuthResource(BaseResource[AsyncArcadeClient]): else: auth_id = auth_id_or_response + # Calculate the new timeout based on the wait parameter + new_timeout = self._client._timeout + if wait is not None: + if isinstance(self._client._timeout, Timeout): + new_timeout = Timeout( + connect=self._client._timeout.connect, + read=(self._client._timeout.read or 0) + wait, + write=self._client._timeout.write, + pool=self._client._timeout.pool, + ) + else: + new_timeout = self._client._timeout + wait + data = await self._client._execute_request( # type: ignore[attr-defined] "GET", f"{self._resource_path}/status", params={"authorizationId": auth_id, "scopes": " ".join(scopes) if scopes else None}, + timeout=new_timeout, ) return AuthResponse(**data) diff --git a/arcade/arcade/client/errors.py b/arcade/arcade/client/errors.py index c9f6c648..c95d8dc1 100644 --- a/arcade/arcade/client/errors.py +++ b/arcade/arcade/client/errors.py @@ -71,6 +71,12 @@ class NotFoundError(APIStatusError): status_code = 404 +class APITimeoutError(APIStatusError): + """408 Request Timeout""" + + status_code = 408 + + class RateLimitError(APIStatusError): """429 Too Many Requests""" diff --git a/arcade/tests/client/test_client.py b/arcade/tests/client/test_client.py index fe5a259c..0b607665 100644 --- a/arcade/tests/client/test_client.py +++ b/arcade/tests/client/test_client.py @@ -5,6 +5,7 @@ from httpx import HTTPStatusError, Response from arcade.client import Arcade, AsyncArcade, AuthProvider from arcade.client.errors import ( + APITimeoutError, BadRequestError, EngineNotHealthyError, InternalServerError, @@ -111,6 +112,7 @@ def mock_async_response(): (401, UnauthorizedError), (403, PermissionDeniedError), (404, NotFoundError), + (408, APITimeoutError), (500, InternalServerError), ], ) @@ -169,6 +171,13 @@ def test_arcade_auth_poll_authorization(test_sync_client, mock_response, monkeyp assert auth_response == AuthResponse(**AUTH_RESPONSE_DATA) +def test_arcade_auth_long_poll_authorization(test_sync_client, mock_response, monkeypatch): + """Test Arcade.auth.poll_authorization method with long polling.""" + monkeypatch.setattr(Arcade, "_execute_request", lambda *args, **kwargs: AUTH_RESPONSE_DATA) + auth_response = test_sync_client.auth.status("auth_123", wait=1) + assert auth_response == AuthResponse(**AUTH_RESPONSE_DATA) + + def test_arcade_tool_run(test_sync_client, mock_response, monkeypatch): """Test Arcade.tools.run method.""" monkeypatch.setattr(Arcade, "_execute_request", lambda *args, **kwargs: TOOL_RESPONSE_DATA) @@ -283,6 +292,20 @@ async def test_async_arcade_auth_poll_authorization( assert auth_response == AuthResponse(**AUTH_RESPONSE_DATA) +@pytest.mark.asyncio +async def test_async_arcade_auth_long_poll_authorization( + test_async_client, mock_async_response, monkeypatch +): + """Test AsyncArcade.auth.poll_authorization method with long polling.""" + + async def mock_execute_request(*args, **kwargs): + return AUTH_RESPONSE_DATA + + monkeypatch.setattr(AsyncArcade, "_execute_request", mock_execute_request) + auth_response = await test_async_client.auth.status("auth_123", wait=1) + assert auth_response == AuthResponse(**AUTH_RESPONSE_DATA) + + @pytest.mark.asyncio async def test_async_arcade_tool_run(test_async_client, mock_async_response, monkeypatch): """Test AsyncArcade.tools.run method."""