mirror of
https://github.com/usestrix/strix.git
synced 2026-10-16 04:28:23 +00:00
feat(llm): authenticate to Anthropic via Workload Identity Federation
Exchange an IdP-issued JWT for a short-lived Anthropic access token and refresh it for the duration of the scan. Configured with the same ANTHROPIC_FEDERATION_RULE_ID / ORGANIZATION_ID / SERVICE_ACCOUNT_ID / WORKSPACE_ID / IDENTITY_TOKEN(_FILE) env vars the Anthropic SDKs read, with a GitHub Actions OIDC fallback. The minted token is kept in memory only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
parent
4695290682
commit
5cec577083
8 changed files with 640 additions and 5 deletions
|
|
@ -22,3 +22,47 @@ export LLM_API_KEY="sk-ant-..."
|
|||
1. Go to [console.anthropic.com](https://console.anthropic.com)
|
||||
2. Navigate to API Keys
|
||||
3. Create a new key
|
||||
|
||||
## Workload Identity Federation (no API key)
|
||||
|
||||
Strix can authenticate to Anthropic with [Workload Identity Federation](https://platform.claude.com/docs/en/manage-claude/workload-identity-federation) instead of a static key. It exchanges a JWT from your identity provider for a short-lived access token and re-exchanges it before it expires, so long scans keep running. The token is held in memory only.
|
||||
|
||||
Set the same variables the Anthropic SDKs read, and leave `LLM_API_KEY` unset (an API key always wins):
|
||||
|
||||
```bash
|
||||
export STRIX_LLM="anthropic/claude-sonnet-4-6"
|
||||
export ANTHROPIC_FEDERATION_RULE_ID="fdrl_..."
|
||||
export ANTHROPIC_ORGANIZATION_ID="00000000-0000-0000-0000-000000000000"
|
||||
export ANTHROPIC_SERVICE_ACCOUNT_ID="svac_..."
|
||||
export ANTHROPIC_WORKSPACE_ID="wrkspc_..." # required if the rule covers several workspaces
|
||||
export ANTHROPIC_IDENTITY_TOKEN_FILE="/var/run/secrets/anthropic.com/token"
|
||||
```
|
||||
|
||||
Strix re-reads `ANTHROPIC_IDENTITY_TOKEN_FILE` on every exchange, so the file must hold a fresh JWT each time (Kubernetes projected tokens rotate automatically). `ANTHROPIC_IDENTITY_TOKEN` (a literal JWT) also works for short runs.
|
||||
|
||||
Federation applies only to models that call Anthropic's API directly (`anthropic/...` or a bare `claude-...` name). Claude through Bedrock, Vertex AI, or OpenRouter keeps that provider's authentication.
|
||||
|
||||
### GitHub Actions
|
||||
|
||||
Inside a job with `id-token: write` and no token file configured, Strix fetches a fresh GitHub OIDC token (audience `https://api.anthropic.com`) for each exchange:
|
||||
|
||||
```yaml
|
||||
permissions:
|
||||
id-token: write
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
strix:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
STRIX_LLM: anthropic/claude-sonnet-4-6
|
||||
ANTHROPIC_FEDERATION_RULE_ID: fdrl_...
|
||||
ANTHROPIC_ORGANIZATION_ID: 00000000-0000-0000-0000-000000000000
|
||||
ANTHROPIC_SERVICE_ACCOUNT_ID: svac_...
|
||||
ANTHROPIC_WORKSPACE_ID: wrkspc_...
|
||||
steps:
|
||||
- uses: actions/checkout@v5
|
||||
- run: strix -n -t ./ --scan-mode quick
|
||||
```
|
||||
|
||||
`run.json` records `auth_mode: workload_identity` for federated runs. If an exchange is denied, the response is an opaque `401`; the reason is listed on the Console's **Settings → Workload identity → Authentication history** page.
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from strix.config.loader import (
|
|||
persist_current,
|
||||
)
|
||||
from strix.config.settings import (
|
||||
AnthropicWifSettings,
|
||||
ContextSettings,
|
||||
DedupeSettings,
|
||||
IntegrationSettings,
|
||||
|
|
@ -28,6 +29,7 @@ from strix.config.settings import (
|
|||
|
||||
|
||||
__all__ = [
|
||||
"AnthropicWifSettings",
|
||||
"ContextSettings",
|
||||
"DedupeSettings",
|
||||
"IntegrationSettings",
|
||||
|
|
|
|||
230
strix/config/anthropic_wif.py
Normal file
230
strix/config/anthropic_wif.py
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
"""Anthropic Workload Identity Federation: trade an IdP-issued JWT for a
|
||||
short-lived Anthropic access token, and keep it fresh for the whole scan.
|
||||
|
||||
Configured with the same env vars the Anthropic SDKs read
|
||||
(``ANTHROPIC_FEDERATION_RULE_ID``, ``ANTHROPIC_ORGANIZATION_ID``,
|
||||
``ANTHROPIC_SERVICE_ACCOUNT_ID``, ``ANTHROPIC_WORKSPACE_ID``, and
|
||||
``ANTHROPIC_IDENTITY_TOKEN_FILE`` or ``ANTHROPIC_IDENTITY_TOKEN``), so one
|
||||
environment serves Strix and any SDK workload alongside it. Inside a GitHub
|
||||
Actions job with ``id-token: write`` and no token configured, the JWT is fetched
|
||||
from the runner's OIDC endpoint instead.
|
||||
|
||||
The minted ``sk-ant-oat01-`` token lives in memory only; it is never persisted.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
import requests
|
||||
|
||||
from strix.config.loader import load_settings
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
DEFAULT_BASE_URL = "https://api.anthropic.com"
|
||||
TOKEN_PATH = "/v1/oauth/token" # noqa: S105 # nosec B105 - URL path, not a secret
|
||||
GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
|
||||
GITHUB_AUDIENCE = "https://api.anthropic.com"
|
||||
|
||||
_TIMEOUT_S = 30
|
||||
# The Anthropic SDKs' refresh schedule: try early, and only fail once the cached
|
||||
# token is too close to expiry to send.
|
||||
_ADVISORY_REFRESH_S = 120
|
||||
_MANDATORY_REFRESH_S = 30
|
||||
|
||||
_lock = threading.Lock()
|
||||
_cached: tuple[str, float] | None = None # (access token, expires-at epoch seconds)
|
||||
|
||||
|
||||
class WifAuthError(RuntimeError):
|
||||
"""The federated token exchange (or fetching the JWT for it) failed."""
|
||||
|
||||
|
||||
def is_configured() -> bool:
|
||||
"""Every field the exchange needs is set, and a JWT source exists."""
|
||||
wif = load_settings().anthropic_wif
|
||||
ids = (wif.federation_rule_id, wif.organization_id, wif.service_account_id)
|
||||
return all(ids) and _has_identity_source()
|
||||
|
||||
|
||||
def is_anthropic_model(model_name: str) -> bool:
|
||||
"""Whether LiteLLM sends ``model_name`` to Anthropic's own API.
|
||||
|
||||
Claude on Bedrock or Vertex carries that provider's prefix and keeps that
|
||||
provider's auth.
|
||||
"""
|
||||
name = model_name.strip().lower()
|
||||
for prefix in ("litellm/", "any-llm/"):
|
||||
name = name.removeprefix(prefix)
|
||||
return name.startswith("anthropic/") or ("/" not in name and name.startswith("claude"))
|
||||
|
||||
|
||||
def applies(model_name: str | None, api_key: str | None = None) -> bool:
|
||||
"""Whether ``model_name`` authenticates by Workload Identity Federation.
|
||||
|
||||
An explicit API key wins, as it does in the Anthropic SDKs.
|
||||
"""
|
||||
if api_key or not model_name:
|
||||
return False
|
||||
return is_anthropic_model(model_name) and is_configured()
|
||||
|
||||
|
||||
def _has_identity_source() -> bool:
|
||||
return bool(
|
||||
load_settings().anthropic_wif.identity_token_file
|
||||
or os.environ.get("ANTHROPIC_IDENTITY_TOKEN")
|
||||
or _github_oidc_available()
|
||||
)
|
||||
|
||||
|
||||
def _github_oidc_available() -> bool:
|
||||
return bool(
|
||||
os.environ.get("ACTIONS_ID_TOKEN_REQUEST_URL")
|
||||
and os.environ.get("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
|
||||
)
|
||||
|
||||
|
||||
def read_identity_token() -> str:
|
||||
"""The IdP JWT for one exchange, read fresh every time.
|
||||
|
||||
JWTs carrying a ``jti`` are single-use, and projected token files rotate on
|
||||
disk, so a cached JWT would be rejected on the next exchange.
|
||||
"""
|
||||
token_file = load_settings().anthropic_wif.identity_token_file
|
||||
if token_file:
|
||||
try:
|
||||
token = Path(token_file).expanduser().read_text(encoding="utf-8").strip()
|
||||
except OSError as exc:
|
||||
raise WifAuthError(f"cannot read ANTHROPIC_IDENTITY_TOKEN_FILE: {exc}") from exc
|
||||
if not token:
|
||||
raise WifAuthError(f"ANTHROPIC_IDENTITY_TOKEN_FILE {token_file} is empty")
|
||||
return token
|
||||
literal = (os.environ.get("ANTHROPIC_IDENTITY_TOKEN") or "").strip()
|
||||
if literal:
|
||||
return literal
|
||||
if _github_oidc_available():
|
||||
return _fetch_github_oidc_token()
|
||||
raise WifAuthError(
|
||||
"no identity token: set ANTHROPIC_IDENTITY_TOKEN_FILE or ANTHROPIC_IDENTITY_TOKEN"
|
||||
)
|
||||
|
||||
|
||||
def _fetch_github_oidc_token() -> str:
|
||||
url = os.environ["ACTIONS_ID_TOKEN_REQUEST_URL"]
|
||||
try:
|
||||
resp = requests.get(
|
||||
url,
|
||||
params={"audience": GITHUB_AUDIENCE},
|
||||
headers={"Authorization": f"Bearer {os.environ['ACTIONS_ID_TOKEN_REQUEST_TOKEN']}"},
|
||||
timeout=_TIMEOUT_S,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
raise WifAuthError(f"GitHub OIDC token request failed: {exc}") from exc
|
||||
if resp.status_code != 200:
|
||||
raise WifAuthError(f"GitHub OIDC token request failed: HTTP {resp.status_code}")
|
||||
value = _json(resp).get("value")
|
||||
if not isinstance(value, str) or not value:
|
||||
raise WifAuthError("GitHub OIDC token response has no value")
|
||||
return value
|
||||
|
||||
|
||||
def exchange() -> tuple[str, float]:
|
||||
"""Run one ``jwt-bearer`` exchange; return ``(access_token, expires_at)``."""
|
||||
settings = load_settings()
|
||||
wif = settings.anthropic_wif
|
||||
body: dict[str, str] = {
|
||||
"grant_type": GRANT_TYPE,
|
||||
"assertion": read_identity_token(),
|
||||
"federation_rule_id": wif.federation_rule_id or "",
|
||||
"organization_id": wif.organization_id or "",
|
||||
"service_account_id": wif.service_account_id or "",
|
||||
}
|
||||
if wif.workspace_id:
|
||||
body["workspace_id"] = wif.workspace_id
|
||||
url = token_url(settings.llm.api_base)
|
||||
started = time.time()
|
||||
try:
|
||||
resp = requests.post(url, json=body, timeout=_TIMEOUT_S)
|
||||
except requests.RequestException as exc:
|
||||
raise WifAuthError(f"token exchange request failed: {exc}") from exc
|
||||
payload = _json(resp)
|
||||
if resp.status_code != 200:
|
||||
raise WifAuthError(_describe_error(resp, payload))
|
||||
access = payload.get("access_token")
|
||||
expires_in = payload.get("expires_in")
|
||||
if not isinstance(access, str) or not access or not isinstance(expires_in, int | float):
|
||||
raise WifAuthError("token exchange response is missing access_token or expires_in")
|
||||
logger.info("Anthropic workload identity token minted (expires in %ss)", int(expires_in))
|
||||
return access, started + float(expires_in)
|
||||
|
||||
|
||||
def token_url(api_base: str | None) -> str:
|
||||
"""The exchange endpoint on the same host the Messages calls go to."""
|
||||
base = (api_base or DEFAULT_BASE_URL).rstrip("/")
|
||||
base = base.removesuffix("/v1")
|
||||
return f"{base}{TOKEN_PATH}"
|
||||
|
||||
|
||||
def _json(resp: requests.Response) -> dict[str, Any]:
|
||||
try:
|
||||
data: object = resp.json()
|
||||
except ValueError:
|
||||
return {}
|
||||
return cast("dict[str, Any]", data) if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
def _describe_error(resp: requests.Response, payload: dict[str, Any]) -> str:
|
||||
error = payload.get("error")
|
||||
detail = ""
|
||||
if isinstance(error, dict):
|
||||
fields = cast("dict[str, Any]", error)
|
||||
detail = f"{fields.get('type', '')}: {fields.get('message', '')}".strip(": ")
|
||||
elif isinstance(error, str):
|
||||
detail = f"{error}: {payload.get('error_description', '')}".strip(": ")
|
||||
request_id = resp.headers.get("request-id")
|
||||
message = f"token exchange failed: HTTP {resp.status_code}"
|
||||
if detail:
|
||||
message += f" ({detail})"
|
||||
if request_id:
|
||||
message += f" [request-id {request_id}]"
|
||||
if resp.status_code == 401:
|
||||
message += "; see the authentication history under Settings > Workload identity"
|
||||
return message
|
||||
|
||||
|
||||
def get_valid_token() -> str:
|
||||
"""A usable access token, exchanging a new one when the cached one runs low.
|
||||
|
||||
Inside the advisory window a failed exchange keeps serving the cached token;
|
||||
inside the mandatory window it raises.
|
||||
"""
|
||||
global _cached # noqa: PLW0603
|
||||
with _lock:
|
||||
now = time.time()
|
||||
if _cached is not None and now < _cached[1] - _ADVISORY_REFRESH_S:
|
||||
return _cached[0]
|
||||
try:
|
||||
_cached = exchange()
|
||||
except WifAuthError:
|
||||
if _cached is not None and now < _cached[1] - _MANDATORY_REFRESH_S:
|
||||
logger.warning(
|
||||
"Anthropic workload identity refresh failed; using the cached token",
|
||||
exc_info=True,
|
||||
)
|
||||
return _cached[0]
|
||||
raise
|
||||
return _cached[0]
|
||||
|
||||
|
||||
def reset_cache() -> None:
|
||||
global _cached # noqa: PLW0603
|
||||
with _lock:
|
||||
_cached = None
|
||||
|
|
@ -403,4 +403,11 @@ def subscription_model(model_name: str | None) -> str | None:
|
|||
|
||||
|
||||
def auth_mode(model_name: str | None) -> str:
|
||||
return "subscription" if subscription_model(model_name) else "api_key"
|
||||
if subscription_model(model_name):
|
||||
return "subscription"
|
||||
from strix.config import anthropic_wif
|
||||
from strix.config.loader import load_settings
|
||||
|
||||
if anthropic_wif.applies(model_name, load_settings().llm.api_key):
|
||||
return "workload_identity"
|
||||
return "api_key"
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import functools
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
|
|
@ -39,7 +40,7 @@ from openai.types.responses.response_usage import ResponseUsage
|
|||
from openai.types.shared import Reasoning
|
||||
|
||||
from strix.agents.prompt import CACHE_POINT
|
||||
from strix.config import codex
|
||||
from strix.config import anthropic_wif, codex
|
||||
from strix.config.loader import load_settings
|
||||
from strix.config.tool_call_ids import TurnCallIdRewriter, dedupe_input
|
||||
from strix.config.tool_call_limits import TurnToolCallLimiter
|
||||
|
|
@ -50,6 +51,7 @@ if TYPE_CHECKING:
|
|||
from collections.abc import AsyncIterator
|
||||
|
||||
from agents.agent_output import AgentOutputSchemaBase
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
from agents.handoffs import Handoff
|
||||
from agents.items import ModelResponse, TResponseInputItem, TResponseStreamEvent
|
||||
from agents.models.interface import ModelTracing
|
||||
|
|
@ -549,6 +551,39 @@ class _CredentialedLitellmProvider(ModelProvider):
|
|||
)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _wif_litellm_model_cls() -> type[LitellmModel]:
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
|
||||
class _WifLitellmModel(LitellmModel):
|
||||
"""An Anthropic model authenticated by Workload Identity Federation.
|
||||
|
||||
Every request re-reads the access token, so a scan outlives each
|
||||
short-lived token; the exchange runs off the event loop.
|
||||
"""
|
||||
|
||||
async def _fetch_response(self, *args: Any, **kwargs: Any) -> Any:
|
||||
self.api_key = await asyncio.to_thread(anthropic_wif.get_valid_token)
|
||||
return await super()._fetch_response(*args, **kwargs)
|
||||
|
||||
return _WifLitellmModel
|
||||
|
||||
|
||||
class _WifLitellmProvider(ModelProvider):
|
||||
"""LiteLLM route that federates Anthropic models and leaves the rest alone."""
|
||||
|
||||
def __init__(self, fallback: ModelProvider) -> None:
|
||||
self._fallback = fallback
|
||||
|
||||
def get_model(self, model_name: str | None) -> Model:
|
||||
from agents.models.default_models import get_default_model
|
||||
|
||||
name = model_name or get_default_model()
|
||||
if not anthropic_wif.is_anthropic_model(name):
|
||||
return self._fallback.get_model(model_name)
|
||||
return _wif_litellm_model_cls()(model=name)
|
||||
|
||||
|
||||
class StrixProvider(MultiProvider):
|
||||
"""Route any non-OpenAI prefix through LiteLLM with the prefix preserved,
|
||||
so users type ``deepseek/deepseek-chat`` rather than
|
||||
|
|
@ -580,7 +615,14 @@ class StrixProvider(MultiProvider):
|
|||
def _create_fallback_provider(self, prefix: str) -> ModelProvider:
|
||||
if prefix == "litellm" and (self._override_api_key or self._override_base_url):
|
||||
return _CredentialedLitellmProvider(self._override_api_key, self._override_base_url)
|
||||
return super()._create_fallback_provider(prefix)
|
||||
provider = super()._create_fallback_provider(prefix)
|
||||
if (
|
||||
prefix == "litellm"
|
||||
and not load_settings().llm.api_key
|
||||
and anthropic_wif.is_configured()
|
||||
):
|
||||
return _WifLitellmProvider(provider)
|
||||
return provider
|
||||
|
||||
def _resolve_prefixed_model(
|
||||
self,
|
||||
|
|
@ -699,6 +741,9 @@ def configure_sdk_model_defaults(settings: Settings) -> None:
|
|||
return
|
||||
_configure_litellm_compatibility()
|
||||
_configure_openrouter_attribution(llm.model)
|
||||
if anthropic_wif.applies(llm.model, llm.api_key):
|
||||
# Fail at startup, not on the first turn, when the federation is misconfigured.
|
||||
anthropic_wif.get_valid_token()
|
||||
if llm.api_key:
|
||||
set_default_openai_key(llm.api_key, use_for_tracing=False)
|
||||
_configure_litellm_default("api_key", llm.api_key)
|
||||
|
|
|
|||
|
|
@ -202,11 +202,28 @@ class ViewerSettings(BaseSettings):
|
|||
app_url: str = Field(default="https://app.strix.ai", alias="STRIX_APP_URL")
|
||||
|
||||
|
||||
class AnthropicWifSettings(BaseSettings):
|
||||
"""Anthropic Workload Identity Federation, under the Anthropic SDKs' own env names.
|
||||
|
||||
The literal ``ANTHROPIC_IDENTITY_TOKEN`` JWT is read straight from the
|
||||
environment instead, so it never lands in the persisted config file.
|
||||
"""
|
||||
|
||||
model_config = _BASE_CONFIG
|
||||
|
||||
federation_rule_id: str | None = Field(default=None, alias="ANTHROPIC_FEDERATION_RULE_ID")
|
||||
organization_id: str | None = Field(default=None, alias="ANTHROPIC_ORGANIZATION_ID")
|
||||
service_account_id: str | None = Field(default=None, alias="ANTHROPIC_SERVICE_ACCOUNT_ID")
|
||||
workspace_id: str | None = Field(default=None, alias="ANTHROPIC_WORKSPACE_ID")
|
||||
identity_token_file: str | None = Field(default=None, alias="ANTHROPIC_IDENTITY_TOKEN_FILE")
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
model_config = _BASE_CONFIG
|
||||
|
||||
llm: LlmSettings = Field(default_factory=LlmSettings)
|
||||
dedupe: DedupeSettings = Field(default_factory=DedupeSettings)
|
||||
anthropic_wif: AnthropicWifSettings = Field(default_factory=AnthropicWifSettings)
|
||||
runtime: RuntimeSettings = Field(default_factory=RuntimeSettings)
|
||||
context: ContextSettings = Field(default_factory=ContextSettings)
|
||||
telemetry: TelemetrySettings = Field(default_factory=TelemetrySettings)
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from rich.console import Console
|
|||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
|
||||
from strix.config import IntegrationSettings, codex, load_settings
|
||||
from strix.config import IntegrationSettings, anthropic_wif, codex, load_settings
|
||||
from strix.interface.utils import (
|
||||
check_docker_connection,
|
||||
image_exists,
|
||||
|
|
@ -53,7 +53,30 @@ def validate_environment() -> None:
|
|||
if not settings.llm.model:
|
||||
missing_required_vars.append("STRIX_LLM")
|
||||
|
||||
if not settings.llm.api_key:
|
||||
uses_wif = anthropic_wif.applies(settings.llm.model, settings.llm.api_key)
|
||||
if uses_wif:
|
||||
try:
|
||||
anthropic_wif.get_valid_token()
|
||||
except anthropic_wif.WifAuthError as exc:
|
||||
console.print(
|
||||
f"[red]Anthropic workload identity federation failed:[/] {exc}",
|
||||
highlight=False,
|
||||
)
|
||||
report_error("workload_identity_exchange_failed")
|
||||
sys.exit(1)
|
||||
logger.info("Anthropic workload identity token exchange OK")
|
||||
elif (
|
||||
settings.llm.api_key
|
||||
and anthropic_wif.is_anthropic_model(settings.llm.model or "")
|
||||
and anthropic_wif.is_configured()
|
||||
):
|
||||
console.print(
|
||||
"[yellow]LLM_API_KEY is set, so it is used instead of the configured "
|
||||
"Anthropic workload identity federation.[/] Unset it to federate.",
|
||||
highlight=False,
|
||||
)
|
||||
|
||||
if not settings.llm.api_key and not uses_wif:
|
||||
missing_optional_vars.append("LLM_API_KEY")
|
||||
|
||||
if not settings.llm.api_base:
|
||||
|
|
|
|||
267
tests/test_anthropic_wif.py
Normal file
267
tests/test_anthropic_wif.py
Normal file
|
|
@ -0,0 +1,267 @@
|
|||
"""Tests for Anthropic Workload Identity Federation: exchange, refresh, routing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from agents.extensions.models.litellm_model import LitellmModel
|
||||
|
||||
from strix.config import anthropic_wif, codex, loader
|
||||
from strix.config.models import StrixProvider
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
_WIF_ENV = {
|
||||
"ANTHROPIC_FEDERATION_RULE_ID": "fdrl_rule",
|
||||
"ANTHROPIC_ORGANIZATION_ID": "00000000-0000-0000-0000-000000000000",
|
||||
"ANTHROPIC_SERVICE_ACCOUNT_ID": "svac_sa",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def token_file(tmp_path: Path) -> Path:
|
||||
path = tmp_path / "jwt"
|
||||
path.write_text("jwt-1\n", encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolated(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, token_file: Path) -> None:
|
||||
for key in (
|
||||
"STRIX_LLM",
|
||||
"LLM_API_KEY",
|
||||
"OPENAI_API_KEY",
|
||||
"LLM_API_BASE",
|
||||
"OPENAI_API_BASE",
|
||||
"OPENAI_BASE_URL",
|
||||
"LITELLM_BASE_URL",
|
||||
"OLLAMA_API_BASE",
|
||||
"ANTHROPIC_WORKSPACE_ID",
|
||||
"ANTHROPIC_IDENTITY_TOKEN",
|
||||
"ACTIONS_ID_TOKEN_REQUEST_URL",
|
||||
"ACTIONS_ID_TOKEN_REQUEST_TOKEN",
|
||||
):
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
for key, value in _WIF_ENV.items():
|
||||
monkeypatch.setenv(key, value)
|
||||
monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN_FILE", str(token_file))
|
||||
monkeypatch.setattr(loader, "_override", tmp_path / "missing-config.json")
|
||||
monkeypatch.setattr(loader, "_cached", None)
|
||||
anthropic_wif.reset_cache()
|
||||
|
||||
|
||||
def _response(status: int, body: dict[str, Any], headers: dict[str, str] | None = None) -> Any:
|
||||
resp = mock.MagicMock(spec=requests.Response)
|
||||
resp.status_code = status
|
||||
resp.json.return_value = body
|
||||
resp.headers = headers or {}
|
||||
return resp
|
||||
|
||||
|
||||
def _ok(token: str = "sk-ant-oat01-a", expires_in: int = 600) -> Any: # noqa: S107
|
||||
return _response(200, {"access_token": token, "token_type": "Bearer", "expires_in": expires_in})
|
||||
|
||||
|
||||
def test_configured_only_with_ids_and_a_token_source(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
assert anthropic_wif.is_configured()
|
||||
monkeypatch.delenv("ANTHROPIC_IDENTITY_TOKEN_FILE")
|
||||
monkeypatch.setattr(loader, "_cached", None)
|
||||
assert not anthropic_wif.is_configured()
|
||||
monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "jwt-literal")
|
||||
assert anthropic_wif.is_configured()
|
||||
monkeypatch.delenv("ANTHROPIC_SERVICE_ACCOUNT_ID")
|
||||
monkeypatch.setattr(loader, "_cached", None)
|
||||
assert not anthropic_wif.is_configured()
|
||||
|
||||
|
||||
def test_exchange_sends_jwt_bearer_grant(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("ANTHROPIC_WORKSPACE_ID", "wrkspc_ws")
|
||||
with mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()) as post:
|
||||
token, expires_at = anthropic_wif.exchange()
|
||||
assert token == "sk-ant-oat01-a" # noqa: S105
|
||||
assert expires_at > 0
|
||||
url = post.call_args.args[0]
|
||||
assert url == "https://api.anthropic.com/v1/oauth/token"
|
||||
assert post.call_args.kwargs["json"] == {
|
||||
"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
|
||||
"assertion": "jwt-1",
|
||||
"federation_rule_id": "fdrl_rule",
|
||||
"organization_id": "00000000-0000-0000-0000-000000000000",
|
||||
"service_account_id": "svac_sa",
|
||||
"workspace_id": "wrkspc_ws",
|
||||
}
|
||||
|
||||
|
||||
def test_exchange_omits_unset_workspace() -> None:
|
||||
with mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()) as post:
|
||||
anthropic_wif.exchange()
|
||||
assert "workspace_id" not in post.call_args.kwargs["json"]
|
||||
|
||||
|
||||
def test_token_url_follows_api_base() -> None:
|
||||
assert anthropic_wif.token_url("https://gw.example/v1/") == "https://gw.example/v1/oauth/token"
|
||||
assert anthropic_wif.token_url(None) == "https://api.anthropic.com/v1/oauth/token"
|
||||
|
||||
|
||||
def test_exchange_rereads_the_token_file(token_file: Path) -> None:
|
||||
with mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()) as post:
|
||||
anthropic_wif.exchange()
|
||||
token_file.write_text("jwt-2", encoding="utf-8")
|
||||
anthropic_wif.exchange()
|
||||
assertions = [call.kwargs["json"]["assertion"] for call in post.call_args_list]
|
||||
assert assertions == ["jwt-1", "jwt-2"]
|
||||
|
||||
|
||||
def test_exchange_error_names_status_but_not_the_jwt() -> None:
|
||||
body = {
|
||||
"type": "error",
|
||||
"error": {"type": "authentication_error", "message": "Authentication failed"},
|
||||
}
|
||||
resp = _response(401, body, {"request-id": "req_1"})
|
||||
with (
|
||||
mock.patch.object(anthropic_wif.requests, "post", return_value=resp),
|
||||
pytest.raises(anthropic_wif.WifAuthError) as exc_info,
|
||||
):
|
||||
anthropic_wif.exchange()
|
||||
message = str(exc_info.value)
|
||||
assert "HTTP 401" in message
|
||||
assert "authentication_error" in message
|
||||
assert "req_1" in message
|
||||
assert "jwt-1" not in message
|
||||
|
||||
|
||||
def test_github_actions_oidc_token_used_without_a_file(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("ANTHROPIC_IDENTITY_TOKEN_FILE")
|
||||
monkeypatch.setattr(loader, "_cached", None)
|
||||
monkeypatch.setenv("ACTIONS_ID_TOKEN_REQUEST_URL", "https://gh.example/token?x=1")
|
||||
monkeypatch.setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "gh-req")
|
||||
gh = _response(200, {"value": "gh-jwt"})
|
||||
with mock.patch.object(anthropic_wif.requests, "get", return_value=gh) as get:
|
||||
assert anthropic_wif.read_identity_token() == "gh-jwt"
|
||||
assert get.call_args.kwargs["params"] == {"audience": "https://api.anthropic.com"}
|
||||
assert get.call_args.kwargs["headers"] == {"Authorization": "Bearer gh-req"}
|
||||
|
||||
|
||||
def test_cached_token_reused_until_advisory_window() -> None:
|
||||
with (
|
||||
mock.patch.object(anthropic_wif.time, "time", return_value=1000.0),
|
||||
mock.patch.object(anthropic_wif.requests, "post", return_value=_ok(expires_in=600)) as post,
|
||||
):
|
||||
assert anthropic_wif.get_valid_token() == "sk-ant-oat01-a"
|
||||
assert anthropic_wif.get_valid_token() == "sk-ant-oat01-a"
|
||||
assert post.call_count == 1
|
||||
|
||||
|
||||
def test_refreshes_inside_advisory_window() -> None:
|
||||
with (
|
||||
mock.patch.object(anthropic_wif.requests, "post", return_value=_ok("sk-ant-oat01-a")),
|
||||
mock.patch.object(anthropic_wif.time, "time", return_value=1000.0),
|
||||
):
|
||||
anthropic_wif.get_valid_token()
|
||||
# expires at 1600; 1500 is inside the 120s advisory window.
|
||||
with (
|
||||
mock.patch.object(anthropic_wif.time, "time", return_value=1500.0),
|
||||
mock.patch.object(anthropic_wif.requests, "post", return_value=_ok("sk-ant-oat01-b")),
|
||||
):
|
||||
assert anthropic_wif.get_valid_token() == "sk-ant-oat01-b"
|
||||
|
||||
|
||||
def test_advisory_refresh_failure_keeps_cached_token() -> None:
|
||||
with (
|
||||
mock.patch.object(anthropic_wif.time, "time", return_value=1000.0),
|
||||
mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()),
|
||||
):
|
||||
anthropic_wif.get_valid_token()
|
||||
with (
|
||||
mock.patch.object(anthropic_wif.time, "time", return_value=1500.0),
|
||||
mock.patch.object(
|
||||
anthropic_wif.requests, "post", side_effect=requests.ConnectionError("down")
|
||||
),
|
||||
):
|
||||
assert anthropic_wif.get_valid_token() == "sk-ant-oat01-a"
|
||||
|
||||
|
||||
def test_mandatory_refresh_failure_raises() -> None:
|
||||
with (
|
||||
mock.patch.object(anthropic_wif.time, "time", return_value=1000.0),
|
||||
mock.patch.object(anthropic_wif.requests, "post", return_value=_ok()),
|
||||
):
|
||||
anthropic_wif.get_valid_token()
|
||||
with (
|
||||
mock.patch.object(anthropic_wif.time, "time", return_value=1580.0),
|
||||
mock.patch.object(
|
||||
anthropic_wif.requests, "post", side_effect=requests.ConnectionError("down")
|
||||
),
|
||||
pytest.raises(anthropic_wif.WifAuthError),
|
||||
):
|
||||
anthropic_wif.get_valid_token()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "expected"),
|
||||
[
|
||||
("anthropic/claude-sonnet-5-5", True),
|
||||
("litellm/anthropic/claude-sonnet-5-5", True),
|
||||
("claude-sonnet-5-5", True),
|
||||
("bedrock/anthropic.claude-sonnet-5-5", False),
|
||||
("vertex_ai/claude-sonnet-5-5", False),
|
||||
("openrouter/anthropic/claude-sonnet-5-5", False),
|
||||
("openai/gpt-5", False),
|
||||
],
|
||||
)
|
||||
def test_applies_only_to_anthropic_direct(model: str, *, expected: bool) -> None:
|
||||
assert anthropic_wif.applies(model) is expected
|
||||
|
||||
|
||||
def test_explicit_api_key_wins() -> None:
|
||||
assert not anthropic_wif.applies("anthropic/claude-sonnet-5-5", "sk-ant-api03-x")
|
||||
|
||||
|
||||
def test_auth_mode_reports_workload_identity(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
assert codex.auth_mode("anthropic/claude-sonnet-5-5") == "workload_identity"
|
||||
assert codex.auth_mode("openai/gpt-5") == "api_key"
|
||||
monkeypatch.setenv("LLM_API_KEY", "sk-ant-api03-x")
|
||||
monkeypatch.setattr(loader, "_cached", None)
|
||||
assert codex.auth_mode("anthropic/claude-sonnet-5-5") == "api_key"
|
||||
|
||||
|
||||
def _unwrap(model: object) -> object:
|
||||
while hasattr(model, "_inner"):
|
||||
model = model._inner
|
||||
return model
|
||||
|
||||
|
||||
def test_provider_federates_anthropic_requests() -> None:
|
||||
model = _unwrap(StrixProvider().get_model("anthropic/claude-sonnet-5-5"))
|
||||
sent: dict[str, Any] = {}
|
||||
|
||||
async def fake_fetch(self: Any, *_args: Any, **_kwargs: Any) -> str:
|
||||
sent["api_key"] = self.api_key
|
||||
return "response"
|
||||
|
||||
with (
|
||||
mock.patch.object(LitellmModel, "_fetch_response", fake_fetch),
|
||||
mock.patch.object(anthropic_wif, "get_valid_token", return_value="sk-ant-oat01-z"),
|
||||
):
|
||||
result = asyncio.run(model._fetch_response()) # type: ignore[attr-defined]
|
||||
assert result == "response"
|
||||
assert sent["api_key"] == "sk-ant-oat01-z"
|
||||
|
||||
|
||||
def test_provider_leaves_other_litellm_models_alone() -> None:
|
||||
model = _unwrap(StrixProvider().get_model("deepseek/deepseek-chat"))
|
||||
assert type(model).__name__ == "LitellmModel"
|
||||
|
||||
|
||||
def test_provider_uses_api_key_when_set(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LLM_API_KEY", "sk-ant-api03-x")
|
||||
monkeypatch.setattr(loader, "_cached", None)
|
||||
model = _unwrap(StrixProvider().get_model("anthropic/claude-sonnet-5-5"))
|
||||
assert type(model).__name__ == "LitellmModel"
|
||||
Loading…
Add table
Reference in a new issue