Per-Request Auth Token Refresh

docs/api/exceptions.rst

@@ -129,7 +129,7 @@ OAuth Error Handling
             client_id=client_id,
             client_secret=client_secret
         )
-        token = await credential_helper.get_token()
+        token = credential_helper.get_token()
     except OAuthError as e:
         logger.error(f"OAuth authentication failed: {e}")
         # Handle OAuth failure - check credentials

docs/authentication.rst

@@ -76,14 +76,6 @@ Set these environment variables for OAuth authentication:
             audience=os.getenv("OAUTH_AUDIENCE", "crisp-athena-live"),
         )
 
-        # Test token acquisition
-        try:
-            token = await credential_helper.get_token()
-            print(f"Successfully acquired token (length: {len(token)})")
-        except Exception as e:
-            print(f"Failed to acquire OAuth token: {e}")
-            return
-
         # Create authenticated channel
         channel = await create_channel_with_credentials(
             host=os.getenv("ATHENA_HOST"),
@@ -261,7 +253,7 @@ Handle OAuth-specific errors gracefully:
     from resolver_athena_client.client.exceptions import AuthenticationError
 
     try:
-        token = await credential_helper.get_token()
+        token = credential_helper.get_token()
     except AuthenticationError as e:
         logger.error(f"OAuth authentication failed: {e}")
         # Handle authentication failure
@@ -356,7 +348,7 @@ Test your authentication setup:
                 client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
             )
 
-            token = await credential_helper.get_token()
+            token = credential_helper.get_token()
             print(f"✓ Authentication successful (token length: {len(token)})")
             return True
 

examples/classify_single_example.py

@@ -213,15 +213,6 @@ async def main() -> int:
         audience=audience,
     )
 
-    # Test token acquisition
-    try:
-        logger.info("Acquiring OAuth token...")
-        token = await credential_helper.get_token()
-        logger.info("Successfully acquired token (length: %d)", len(token))
-    except Exception:
-        logger.exception("Failed to acquire OAuth token")
-        return 1
-
     # Configure client options
     options = AthenaOptions(
         host=host,

examples/example.py

@@ -163,15 +163,6 @@ async def main() -> int:
         audience=audience,
     )
 
-    # Test token acquisition
-    try:
-        logger.info("Acquiring OAuth token...")
-        token = await credential_helper.get_token()
-        logger.info("Successfully acquired token (length: %d)", len(token))
-    except Exception:
-        logger.exception("Failed to acquire OAuth token")
-        return 1
-
     # Get available deployment
     channel = await create_channel_with_credentials(host, credential_helper)
     async with DeploymentSelector(channel) as deployment_selector:

src/resolver_athena_client/client/__init__.py

@@ -5,6 +5,7 @@
 
 from resolver_athena_client.client.channel import (
     CredentialHelper,
+    TokenData,
     create_channel_with_credentials,
 )
 from resolver_athena_client.client.exceptions import (
@@ -26,6 +27,7 @@
     "CredentialError",
     "CredentialHelper",
     "OAuthError",
+    "TokenData",
     "TokenExpiredError",
     "create_channel_with_credentials",
     "get_output_error_summary",

src/resolver_athena_client/client/channel.py

@@ -1,9 +1,9 @@
 """Channel creation utilities for the Athena client."""
 
-import asyncio
 import json
+import threading
 import time
-from typing import override
+from typing import NamedTuple, override
 
 import grpc
 import httpx
@@ -16,37 +16,22 @@
 )
 
 
-class TokenMetadataPlugin(grpc.AuthMetadataPlugin):
-    """Plugin that adds authorization token to gRPC metadata."""
+class TokenData(NamedTuple):
+    """Immutable snapshot of token state.
 
-    def __init__(self, token: str) -> None:
-        """Initialize the plugin with the auth token.
-
-        Args:
-        ----
-            token: The authorization token to add to requests
-
-        """
-        self._token: str = token
-
-    @override
-    def __call__(
-        self,
-        _: grpc.AuthMetadataContext,
-        callback: grpc.AuthMetadataPluginCallback,
-    ) -> None:
-        """Pass authentication metadata to the provided callback.
-
-        This method will be invoked asynchronously in a separate thread.
+    Storing token, expiry, and scheme together as a single object
+    ensures that validity checks and token reads are always consistent,
+    eliminating TOCTOU races between ``get_token`` and
+    ``invalidate_token``.
+    """
 
-        Args:
-        ----
-            callback: An AuthMetadataPluginCallback to be invoked either
-            synchronously or asynchronously.
+    access_token: str
+    expires_at: float
+    scheme: str
 
-        """
-        metadata = (("authorization", f"Token {self._token}"),)
-        callback(metadata, None)
+    def is_valid(self) -> bool:
+        """Check if this token is still valid (with a 30-second buffer)."""
+        return time.time() < (self.expires_at - 30)
 
 
 class CredentialHelper:
@@ -80,56 +65,49 @@ def __init__(
         self._client_secret: str = client_secret
         self._auth_url: str = auth_url
         self._audience: str = audience
-        self._token: str | None = None
-        self._token_expires_at: float | None = None
-        self._lock: asyncio.Lock = asyncio.Lock()
+        self._token_data: TokenData | None = None
+        self._lock: threading.Lock = threading.Lock()
 
-    async def get_token(self) -> str:
-        """Get a valid authentication token.
+    def get_token(self) -> TokenData:
+        """Get valid token data, refreshing if necessary.
 
-        This method will return a cached token if it's still valid,
-        or fetch a new token if needed.
+        Uses double-checked locking: the happy path (token is valid)
+        avoids acquiring the lock entirely.  The lock is only taken
+        when the token needs to be refreshed.
 
         Returns
         -------
-            A valid authentication token
+            A valid ``TokenData`` containing access token, expiry, and scheme
 
         Raises
         ------
             OAuthError: If token acquisition fails
-            TokenExpiredError: If token has expired and refresh fails
+            RuntimeError: If token is unexpectedly None after refresh
 
         """
-        async with self._lock:
-            if self._is_token_valid():
-                if self._token is None:
-                    msg = "Token should be valid but is None"
-                    raise RuntimeError(msg)
-                return self._token
-
-            await self._refresh_token()
-            if self._token is None:
-                msg = "Token refresh failed"
-                raise RuntimeError(msg)
-            return self._token
+        token_data = self._token_data
+        if token_data is not None and token_data.is_valid():
+            return token_data
 
-    def _is_token_valid(self) -> bool:
-        """Check if the current token is valid and not expired.
+        with self._lock:
+            token_data = self._token_data
+            if token_data is not None and token_data.is_valid():
+                return token_data
 
-        Returns
-        -------
-            True if token is valid, False otherwise
+            self._refresh_token()
 
-        """
-        if not self._token or not self._token_expires_at:
-            return False
-
-        # Add 30 second buffer before expiration
-        return time.time() < (self._token_expires_at - 30)
+            token_data = self._token_data
+            if token_data is None:
+                msg = "Token is unexpectedly None after refresh"
+                raise RuntimeError(msg)
+            return token_data
 
-    async def _refresh_token(self) -> None:
+    def _refresh_token(self) -> None:
         """Refresh the authentication token by making an OAuth request.
 
+        This is a synchronous call (suitable for the gRPC metadata-plugin
+        thread) and must be called while ``self._lock`` is held.
+
         Raises
         ------
             OAuthError: If the OAuth request fails
@@ -145,21 +123,24 @@ async def _refresh_token(self) -> None:
         headers = {"content-type": "application/json"}
 
         try:
-            async with httpx.AsyncClient() as client:
-                response = await client.post(
+            with httpx.Client() as client:
+                response = client.post(
                     self._auth_url,
                     json=payload,
                     headers=headers,
                     timeout=30.0,
                 )
                 _ = response.raise_for_status()
 
-                token_data = response.json()
-                self._token = token_data["access_token"]
-                expires_in = token_data.get(
-                    "expires_in", 3600
-                )  # Default 1 hour
-                self._token_expires_at = time.time() + expires_in
+            raw = response.json()
+            access_token: str = raw["access_token"]
+            expires_in: int = raw.get("expires_in", 3600)  # Default 1 hour
+            scheme: str = raw.get("token_type", "Bearer").capitalize()
+            self._token_data = TokenData(
+                access_token=access_token,
+                expires_at=time.time() + expires_in,
+                scheme=scheme,
+            )
 
         except httpx.HTTPStatusError as e:
             error_detail = ""
@@ -190,11 +171,58 @@ async def _refresh_token(self) -> None:
             msg = f"Unexpected error during OAuth: {e}"
             raise OAuthError(msg) from e
 
-    async def invalidate_token(self) -> None:
+    def invalidate_token(self) -> None:
         """Invalidate the current token to force a refresh on next use."""
-        async with self._lock:
-            self._token = None
-            self._token_expires_at = None
+        with self._lock:
+            self._token_data = None
+
+
+class _AutoRefreshTokenAuthMetadataPlugin(grpc.AuthMetadataPlugin):
+    """gRPC auth plugin that fetches a fresh token for every RPC.
+
+    The plugin delegates to ``CredentialHelper.get_token()`` which
+    handles caching, expiry checks, and thread-safe refresh internally.
+    This callback is invoked by gRPC on a *separate* thread, so the
+    underlying ``CredentialHelper`` must use ``threading.Lock`` (not
+    ``asyncio.Lock``).
+    """
+
+    def __init__(self, credential_helper: CredentialHelper) -> None:
+        """Initialize with a credential helper.
+
+        Args:
+        ----
+            credential_helper: The helper that manages token lifecycle
+
+        """
+        self._credential_helper: CredentialHelper = credential_helper
+
+    @override
+    def __call__(
+        self,
+        _: grpc.AuthMetadataContext,
+        callback: grpc.AuthMetadataPluginCallback,
+    ) -> None:
+        """Supply authorization metadata for an RPC.
+
+        Called by the gRPC runtime on a background thread before each
+        RPC.  On success the token is forwarded as a Bearer token; on
+        failure the error is passed to the callback so gRPC can surface
+        it as an RPC error.
+
+        Args:
+        ----
+            callback: gRPC callback to receive metadata or an error
+
+        """
+        try:
+            token_data = self._credential_helper.get_token()
+            scheme = token_data.scheme
+            token = token_data.access_token
+            metadata = (("authorization", f"{scheme} {token}"),)
+            callback(metadata, None)
+        except Exception as err:  # noqa: BLE001
+            callback(None, err)  # pyright: ignore[reportArgumentType]
 
 
 async def create_channel_with_credentials(
@@ -215,19 +243,23 @@ async def create_channel_with_credentials(
     Raises:
     ------
         InvalidHostError: If host is empty
-        OAuthError: If OAuth authentication fails
+
+    Note:
+    ----
+        OAuth errors are no longer raised at channel-creation time.
+        Instead, they surface as RPC errors when the per-request auth
+        metadata plugin attempts to acquire a token.
 
     """
     if not host:
         raise InvalidHostError(InvalidHostError.default_message)
 
-    # Get a valid token from the credential helper
-    token = await credential_helper.get_token()
-
-    # Create credentials with token authentication
+    # Create credentials with per-RPC token refresh
     credentials = grpc.composite_channel_credentials(
         grpc.ssl_channel_credentials(),
-        grpc.access_token_call_credentials(token),
+        grpc.metadata_call_credentials(
+            _AutoRefreshTokenAuthMetadataPlugin(credential_helper)
+        ),
     )
 
     # Configure gRPC options for persistent connections