diff --git a/src/anthropic/lib/credentials/_types.py b/src/anthropic/lib/credentials/_types.py index 2dfe0593c..e87068834 100644 --- a/src/anthropic/lib/credentials/_types.py +++ b/src/anthropic/lib/credentials/_types.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Dict, Callable, Optional, Protocol +from typing import Any, Dict, Union, Callable, Optional, Protocol, Coroutine from dataclasses import field, dataclass from typing_extensions import override, runtime_checkable @@ -80,7 +80,7 @@ def for_base_url(self, base_url: str) -> AccessTokenProvider: # Innermost layer: returns the raw external JWT string (used as the # ``identity_token_provider`` argument to :class:`WorkloadIdentityCredentials`). -IdentityTokenProvider = Callable[[], str] +IdentityTokenProvider = Union[Callable[[], str], Callable[[], Coroutine[Any, Any, str]]] @dataclass(frozen=True) diff --git a/src/anthropic/lib/credentials/_workload.py b/src/anthropic/lib/credentials/_workload.py index b00ab21c8..ffe384dd3 100644 --- a/src/anthropic/lib/credentials/_workload.py +++ b/src/anthropic/lib/credentials/_workload.py @@ -2,6 +2,8 @@ import copy import time +import asyncio +import inspect import logging from types import TracebackType from typing import Any, Dict, Type, Union, NoReturn, Optional @@ -135,6 +137,22 @@ def __str__(self) -> str: return base +def _resolve_identity_token(provider: IdentityTokenProvider) -> str: + raw = provider() + if inspect.iscoroutine(raw): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop and loop.is_running(): + import concurrent.futures + + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return str(pool.submit(asyncio.run, raw).result()) + return str(asyncio.run(raw)) + return str(raw) + + class WorkloadIdentityCredentials: """Exchanges an external OIDC JWT for an Anthropic access token via the RFC 7523 ``jwt-bearer`` grant. @@ -251,7 +269,7 @@ def __call__(self, *, force_refresh: bool = False) -> AccessToken: # file (e.g. a k8s projected SA token) may have rotated. force_refresh # is a no-op: this provider has no cache to bypass. del force_refresh - jwt = SecretStr(self._identity_token_provider()) + jwt = SecretStr(_resolve_identity_token(self._identity_token_provider)) assertion_bytes = len(jwt.get_secret_value().encode("utf-8")) if assertion_bytes > _MAX_ASSERTION_BYTES: diff --git a/tests/lib/test_credentials.py b/tests/lib/test_credentials.py index cf88bab56..a126e1b51 100644 --- a/tests/lib/test_credentials.py +++ b/tests/lib/test_credentials.py @@ -1001,6 +1001,31 @@ def test_exchange(self, respx_mock: MockRouter) -> None: assert "workspace_id" not in body assert "scope" not in body + @pytest.mark.respx() + def test_exchange_with_async_identity_token_provider(self, respx_mock: MockRouter) -> None: + respx_mock.post(TOKEN_URL).mock( + return_value=httpx2.Response( + 200, + json={"access_token": "sk-ant-oat01-async", "token_type": "Bearer", "expires_in": 600}, + ) + ) + + async def async_provider() -> str: + return "ext.jwt.async.value" + + creds = WorkloadIdentityCredentials( + identity_token_provider=async_provider, + federation_rule_id="fdrl_01abc", + organization_id="00000000-0000-0000-0000-000000000000", + ) + + token = creds() + assert token.token == "sk-ant-oat01-async" + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert len(calls) == 1 + body = json.loads(calls[0].request.content) + assert body["assertion"] == "ext.jwt.async.value" + @pytest.mark.respx() def test_service_account_included(self, respx_mock: MockRouter) -> None: respx_mock.post(TOKEN_URL).mock(return_value=httpx2.Response(200, json={"access_token": "t", "expires_in": 60}))