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>
113 lines
4.4 KiB
Python
113 lines
4.4 KiB
Python
"""MCP toolsets, rebuilt for every conversation turn.
|
|
|
|
The rule this enforces is the one thing the identity design depends on: **an MCP
|
|
client is built per turn, with tokens exchanged for that turn**, never held open
|
|
with a fixed token.
|
|
|
|
That is not a limitation of any framework — it is what keeps the blast radius of
|
|
a leaked token to one turn, and what makes `azp`/`sub` in the audit trail mean
|
|
something. PydanticAI takes `headers` at toolset construction, which is exactly
|
|
the shape this needs; the same constraint on other runtimes (langchain-mcp-adapters
|
|
issues #172/#194/#235, CrewAI's MCPServerAdapter) has the same answer.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from pydantic_ai.mcp import MCPToolset, StreamableHttpTransport
|
|
from pydantic_ai.toolsets import (
|
|
AbstractToolset,
|
|
ApprovalRequiredToolset,
|
|
FilteredToolset,
|
|
PrefixedToolset,
|
|
)
|
|
|
|
from .approval import is_write, needs_approval
|
|
from .config import Settings
|
|
from .identity import TokenExchanger
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class McpTarget:
|
|
"""One MCP server this agent is allowed to reach."""
|
|
|
|
name: str
|
|
url: str
|
|
audience_scope: str
|
|
read_only: bool = False
|
|
"""Refuse this server's write tools inside the agent as well as at the gateway.
|
|
|
|
Belt and braces on purpose. The gateway already refuses what the end user's
|
|
roles do not allow, but an admin reaching this agent would otherwise be
|
|
offered `create-order` — a tool this agent has no business calling whoever
|
|
is asking.
|
|
"""
|
|
|
|
|
|
def targets(settings: Settings) -> list[McpTarget]:
|
|
"""The catalog domain read/write, plus sales read-only.
|
|
|
|
sales is here to answer "does this product actually sell?" before a pricing
|
|
or discontinuation decision. Read only — this agent never touches orders,
|
|
and the scope it requests reflects that.
|
|
"""
|
|
return [
|
|
McpTarget("catalog", settings.mcp_catalog_url, settings.mcp_catalog_scope),
|
|
McpTarget("sales", settings.mcp_sales_url, settings.mcp_sales_scope, read_only=True),
|
|
]
|
|
|
|
|
|
async def build_toolsets(
|
|
settings: Settings,
|
|
exchanger: TokenExchanger,
|
|
subject_token: str,
|
|
) -> list[MCPToolset]:
|
|
"""One toolset per MCP server, each carrying its own audience-scoped token.
|
|
|
|
Several narrow tokens rather than one broad one, deliberately: a token
|
|
minted for `mcp-catalog` is refused by the sales gateway, so a prompt
|
|
injection that talks the agent into reaching the wrong domain is refused at
|
|
the endpoint rather than succeeding quietly.
|
|
"""
|
|
def client_factory(**kwargs: Any) -> httpx.AsyncClient:
|
|
# fastmcp builds this client itself and would otherwise take httpx's
|
|
# 5s default connect timeout. See Settings.http_timeout for why that is
|
|
# not survivable against `.local` names, and what the failure looks like
|
|
# when it bites: a bare "Failed to initialize server session".
|
|
kwargs["verify"] = settings.verify
|
|
kwargs["timeout"] = httpx.Timeout(settings.http_timeout, connect=settings.http_timeout)
|
|
return httpx.AsyncClient(**kwargs)
|
|
|
|
toolsets: list[AbstractToolset[None]] = []
|
|
for target in targets(settings):
|
|
token = await exchanger.exchange(subject_token, target.audience_scope)
|
|
toolset: AbstractToolset[None] = MCPToolset(
|
|
StreamableHttpTransport(
|
|
url=target.url,
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
# No `verify=` here: the factory below sets it, and passing
|
|
# both makes fastmcp ignore this one with a warning.
|
|
httpx_client_factory=client_factory,
|
|
),
|
|
id=target.name,
|
|
init_timeout=settings.http_timeout,
|
|
)
|
|
|
|
if target.read_only:
|
|
toolset = FilteredToolset(toolset, lambda ctx, tool: not is_write(tool))
|
|
|
|
# Writes stop and wait for a person. The predicate reads the MCP
|
|
# `readOnlyHint` the server publishes, so the server decides what counts
|
|
# as a write and this agent keeps no list of tool names in step with it.
|
|
toolset = ApprovalRequiredToolset(toolset, needs_approval)
|
|
|
|
# Two servers today, more later. A collision between two domains'
|
|
# namespaces would otherwise be silent — one tool shadowing another of
|
|
# the same name and a different meaning.
|
|
toolsets.append(PrefixedToolset(toolset, target.name))
|
|
|
|
return toolsets
|