Fix arcade chat Auth Polling (#96)

# PR Description
When a tool call required authorization, `arcade chat` would hit rate
limits extremely quickly when waiting to authorize with a given link.
This PR introduces using long polling when sending a GET request to
Arcade API `/auth/status`. The `/auth/status` endpoint supports the
`wait` query parameter, where if present, will not respond until either
auth status becomes `completed` or `timeout` is reached. If the
`timeout` is reached, `arcade chat` catches the 408 response and tries
again. For `arcade chat` we set the `wait` query param to 60 seconds.

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)
+        try:
-        auth_response = client.auth.status(auth_response)
+            auth_response = client.auth.status(auth_response, wait=60)
+        except APITimeoutError:
+            continue
 
 
 def get_tool_authorization(

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,
         }

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)
 

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"""
 

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."""