Stéphane Tailland
1a13743078
The agent asked for confirmation because the system prompt told it to, and read the answer itself. Both halves are now enforced. Each MCP toolset is wrapped in ApprovalRequiredToolset whose predicate reads the readOnlyHint the server publishes — the gateway decides what counts as a write, and the agent keeps no list of tool names in step with it. output_type=[str, DeferredToolRequests] lets a turn end holding a proposed call instead of making it; the executor parks the A2A task in input-required and renders the arguments the model actually sent, not what it said it would do. The reply is matched against a closed set of words, so "oui mais change le prix" denies the call and the correction comes back for approval on its own merits — verified end to end, along with a draft the user approved that was still refused because msrp sat below buyprice. Also wired: a ProductDraft validated after approval and before the call leaves; UsageLimits on every run; FilteredToolset denying sales' write tools inside the agent as well as at the gateway; PrefixedToolset namespacing the two servers; dynamic instructions naming the caller; Agent.instrument_all() for a span per model call and per tool call; and pydantic_ai.retries on inference only, so a 429 from the gateway is waited out rather than failing mid-reasoning. Inference gets its own httpx client: validate_response raises on any 4xx, and applying that to the shared client would turn a token-exchange refusal into an exception carrying no body — which is where Keycloak says what was wrong. Its transport is built explicitly, because a transport passed to the client makes verify= inert and the lab CA would otherwise go undeclared. Tests grew from 7 to 31, covering what approval decides and what a draft must satisfy. evals/catalogue_evals.py adds a behavioural suite, kept out of pytest because every case spends real tokens; its first case checks that a count comes from the tool rather than from what the model remembers of a public dataset. CLAUDE.md corrected: instrument=True is not an Agent parameter in 2.35.3, pydantic_ai.retries needs tenacity, TemplateStr needs pydantic-handlebars, and approval does not remove the in-memory history — resuming still needs the messages, so replicas: 1 stands until they live in the A2A task store. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
"""Token exchange: the end user's token in, one narrowly-scoped token per MCP
|
|
server out.
|
|
|
|
Two measured facts about Keycloak 26.6 shape this, both of which cost time to
|
|
discover and neither of which is obvious from the RFC:
|
|
|
|
**`actor_token` is silently ignored.** Standard Token Exchange accepts the
|
|
parameter, returns a valid token, and emits no `act` claim — no error, no
|
|
warning. Delegation therefore cannot be built on RFC 8693 `act`. The actor is
|
|
identified by `azp`, which carries this agent's client id.
|
|
|
|
**The audience is requested as a scope.** Passing `audience=mcp-catalog` returns
|
|
"Requested audience not available". Passing `scope=openid mcp-catalog-audience`
|
|
works, and the resulting token carries `aud: [mcp-catalog]`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import time
|
|
from dataclasses import dataclass
|
|
|
|
import httpx2
|
|
|
|
GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
|
|
SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"
|
|
|
|
|
|
class TokenExchangeError(RuntimeError):
|
|
"""Keycloak refused the exchange. Carries its response, which names the cause."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExchangedToken:
|
|
access_token: str
|
|
expires_at: float
|
|
|
|
@property
|
|
def expired(self) -> bool:
|
|
# A conversation turn lasts seconds and a token five minutes, so this
|
|
# margin only matters for a turn that stalls on a slow model.
|
|
return time.monotonic() >= self.expires_at - 30
|
|
|
|
|
|
class TokenExchanger:
|
|
"""Exchanges a user token for audience-scoped tokens.
|
|
|
|
Caches per (subject token, scope) for the lifetime of one turn. Deliberately
|
|
not a long-lived cache keyed by user: a token that outlives the turn it was
|
|
minted for is a token that can be replayed after the user's session ended.
|
|
"""
|
|
|
|
def __init__(self, client: httpx2.AsyncClient, *, token_endpoint: str, client_id: str, client_secret: str) -> None:
|
|
self._client = client
|
|
self._endpoint = token_endpoint
|
|
self._client_id = client_id
|
|
self._client_secret = client_secret
|
|
self._cache: dict[tuple[str, str], ExchangedToken] = {}
|
|
|
|
async def exchange(self, subject_token: str, audience_scope: str) -> str:
|
|
key = (subject_token, audience_scope)
|
|
cached = self._cache.get(key)
|
|
if cached and not cached.expired:
|
|
return cached.access_token
|
|
|
|
response = await self._client.post(
|
|
self._endpoint,
|
|
data={
|
|
"grant_type": GRANT_TYPE,
|
|
"client_id": self._client_id,
|
|
"client_secret": self._client_secret,
|
|
"subject_token": subject_token,
|
|
"subject_token_type": SUBJECT_TOKEN_TYPE,
|
|
# Not `audience=...`. See the module docstring.
|
|
"scope": f"openid {audience_scope}",
|
|
},
|
|
)
|
|
if response.status_code != 200:
|
|
raise TokenExchangeError(
|
|
f"exchange for {audience_scope!r} failed: HTTP {response.status_code} {response.text[:200]}"
|
|
)
|
|
|
|
payload = response.json()
|
|
token = ExchangedToken(
|
|
access_token=payload["access_token"],
|
|
expires_at=time.monotonic() + float(payload.get("expires_in", 60)),
|
|
)
|
|
self._cache[key] = token
|
|
return token.access_token
|
|
|
|
|
|
def caller_of(access_token: str) -> str:
|
|
"""The username inside a token, for the prompt and nothing else.
|
|
|
|
Read without verifying the signature on purpose: this is a label shown to
|
|
the model, never a decision. Every decision that depends on this identity is
|
|
taken by Keycloak or by the gateway, both of which do verify.
|
|
"""
|
|
try:
|
|
payload = access_token.split(".")[1]
|
|
payload += "=" * (-len(payload) % 4)
|
|
claims = json.loads(base64.urlsafe_b64decode(payload))
|
|
except Exception: # noqa: BLE001 — a malformed token is the gateway's problem
|
|
return "an unidentified caller"
|
|
return claims.get("preferred_username") or claims.get("sub") or "an unidentified caller"
|