feat: make the write path a mechanism instead of a prompt rule

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>
This commit is contained in:
Stéphane Tailland
2026-08-29 09:15:33 +02:00
parent aace433d7e
commit 1a13743078
12 changed files with 815 additions and 49 deletions
+10
View File
@@ -6,12 +6,22 @@ import logging
import uvicorn
from pydantic_ai import Agent
from .a2a_server import build_app
from .config import Settings
def main() -> None:
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
# OpenTelemetry spans for every model call and every tool call. The MCP
# server already records *which* tool ran; this is what records why — the
# reasoning step that chose it, and what it cost.
#
# `instrument_all` rather than a per-agent flag: 2.35.3 has no `instrument=`
# on Agent, and this process runs exactly one agent anyway.
Agent.instrument_all()
settings = Settings() # type: ignore[call-arg] # client_secret comes from the environment
uvicorn.run(build_app(settings), host=settings.host, port=settings.port)
+42 -6
View File
@@ -42,7 +42,7 @@ from a2a.types import (
)
from starlette.applications import Starlette
from .agent import CatalogueAgent
from .agent import CatalogueAgent, build_llm_client
from .config import Settings
logger = logging.getLogger(__name__)
@@ -128,10 +128,22 @@ class CatalogueExecutor(AgentExecutor):
)
return
user_input = context.get_user_input()
try:
answer = await self._agent.run(
context.get_user_input(), subject_token, context.context_id
)
# A conversation already holding a proposed write is answering it,
# not asking something new. The decision is read from the reply
# deterministically — a model must not be the thing that decides
# whether the user said yes.
if self._agent.awaiting(context.context_id) is not None:
approved = _reads_as_approval(user_input)
result = await self._agent.resume(
context.context_id,
subject_token,
approved=approved,
reason="" if approved else f"the user replied {user_input!r}",
)
else:
result = await self._agent.run(user_input, subject_token, context.context_id)
except Exception as exc: # noqa: BLE001 — surfaced to the caller, logged here
logger.exception("turn failed")
await updater.failed(
@@ -139,12 +151,34 @@ class CatalogueExecutor(AgentExecutor):
)
return
await updater.complete(message=updater.new_agent_message(parts=[Part(text=answer)]))
if result.awaiting_approval:
# `input-required` rather than `completed`: the task stays open and
# carries the arguments the model proposed, so the record of what was
# agreed to is protocol state rather than a sentence in a transcript.
await updater.requires_input(
message=updater.new_agent_message(parts=[Part(text=result.text)])
)
return
await updater.complete(message=updater.new_agent_message(parts=[Part(text=result.text)]))
async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
await TaskUpdater(event_queue, context.task_id, context.context_id).cancel()
#: What counts as approval. Deliberately a closed set, matched exactly after
#: stripping: an approval must be unambiguous, and anything else — a question, a
#: correction, silence dressed as agreement — cancels. The prompt shown to the
#: user names the word to send.
APPROVAL_WORDS = frozenset(
{"approve", "approved", "approuve", "approuvé", "oui", "yes", "ok", "confirm", "confirme"}
)
def _reads_as_approval(text: str) -> bool:
return text.strip().strip(".!").casefold() in APPROVAL_WORDS
def _subject_token(context: RequestContext) -> str | None:
headers = {}
call = context.call_context
@@ -160,9 +194,10 @@ def build_app(settings: Settings) -> Starlette:
# httpx2, not httpx: PydanticAI deprecates httpx clients for
# OpenAI-compatible providers and drops them in v3.
http_client = httpx2.AsyncClient(verify=settings.verify, timeout=settings.http_timeout)
llm_client = build_llm_client(settings)
card = build_agent_card(settings)
handler = DefaultRequestHandlerV2(
agent_executor=CatalogueExecutor(CatalogueAgent(settings, http_client)),
agent_executor=CatalogueExecutor(CatalogueAgent(settings, http_client, llm_client)),
task_store=InMemoryTaskStore(),
agent_card=card,
)
@@ -172,6 +207,7 @@ def build_app(settings: Settings) -> Starlette:
async def lifespan(_app: Starlette) -> AsyncIterator[None]:
yield
await http_client.aclose()
await llm_client.aclose()
return Starlette(
routes=create_agent_card_routes(card)
+212 -19
View File
@@ -3,17 +3,43 @@
from __future__ import annotations
from collections import OrderedDict
from dataclasses import dataclass
from typing import Any
import httpx
import httpx2
from pydantic_ai import Agent
from pydantic_ai import (
Agent,
DeferredToolRequests,
DeferredToolResults,
ToolApproved,
ToolDenied,
)
from pydantic_ai.messages import ModelMessage
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.tools import RunContext
from pydantic_ai.retries import AsyncTenacityTransport, RetryConfig, wait_retry_after
from pydantic_ai.usage import UsageLimits
from tenacity import retry_if_exception_type, stop_after_attempt
from .approval import validate_draft
from .config import Settings
from .identity import TokenExchanger
from .identity import TokenExchanger, caller_of
from .mcp_clients import build_toolsets
@dataclass
class CallerDeps:
"""What the run knows about who it is acting for.
Only the username, and only for the prompt: nothing here grants anything.
The identity that matters travels in the exchanged token, and the gateway
is what reads it.
"""
user: str
SYSTEM_PROMPT = """\
You are Catalogue, the product management agent for the Classic Models catalog.
@@ -30,16 +56,64 @@ How to work:
missing. Do not infer availability from what you can see.
Writing to the catalog:
- You may create products. Before doing so, restate exactly what you are about
to create and ask the user to confirm. Never write on an implied instruction.
- Every field you write comes from the user or from a tool. If a required field
is missing, ask for it — do not choose a plausible value, a price least of all.
- You may create products. Every field you write comes from the user or from a
tool. If a required field is missing, ask for it — do not choose a plausible
value, a price least of all.
- You do not need to ask for confirmation yourself. A write is held for the
user's approval before it leaves, and they see exactly what you proposed.
- After a write, report what was created, including the identifier the tool
returned, so the user can check it.
- A write may be refused. That is the gateway applying the user's own role, not
a mistake to work around: report the refusal and stop.
- A write may be refused, by the user or by the gateway applying their role.
That is not a mistake to work around: report the refusal and stop.
"""
#: Appended per turn. Dynamic rather than baked into the constant above so the
#: agent states the identity it is acting under — the thing that decides which
#: tools it can even see.
def caller_instructions(ctx: RunContext[CallerDeps]) -> str:
"""Appended per run, so the agent states the identity it is acting under.
A function rather than `TemplateStr`: that wants `pydantic-handlebars`, and a
template engine earns nothing against a single interpolation.
"""
return (
f"You are acting on behalf of {ctx.deps.user}. The tools you can see are "
"the ones that identity is allowed to call; if a capability seems "
"missing, say so rather than looking for another way to it."
)
def build_llm_client(settings: Settings) -> httpx2.AsyncClient:
"""A client of its own for inference, because it retries and the other must not.
The LLM gateway answers `429 AI token rate limit exceeded` when a
tool-calling loop feeds a large result back on every iteration, and it says
when to come back. `wait_retry_after` honours that header, so a turn survives
the limit instead of failing halfway through its reasoning.
Not the shared client: `validate_response` raises on any 4xx, which would
turn a token exchange refusal into an exception carrying no body — and the
body is where Keycloak says what was wrong.
"""
config = RetryConfig(
retry=retry_if_exception_type(httpx.HTTPStatusError),
wait=wait_retry_after(max_wait=60.0),
stop=stop_after_attempt(4),
reraise=True,
)
return httpx2.AsyncClient(
transport=AsyncTenacityTransport(
config=config,
# Built explicitly: a transport passed to the client makes `verify=`
# on the client inert, so the wrapped transport is the only place the
# lab CA can be declared. Leave it out and every call fails with a
# bare "Connection error".
wrapped=httpx2.AsyncHTTPTransport(verify=settings.verify),
validate_response=lambda r: r.raise_for_status(),
),
timeout=settings.http_timeout,
)
def build_model(
settings: Settings, http_client: httpx2.AsyncClient, access_token: str
@@ -74,11 +148,45 @@ def build_model(
# conversation, not something to persist and re-authorise later.
MAX_CONVERSATIONS = 200
#: A turn that loops has to stop somewhere. Reached, it raises rather than
#: quietly returning a partial answer.
DEFAULT_TOOL_CALL_LIMIT = 20
@dataclass
class PendingApproval:
"""A turn parked on the user's decision.
Holds the conversation as it stood plus the calls the model proposed, so the
resumed run continues the same reasoning rather than starting again from a
summary of it.
"""
requests: DeferredToolRequests
messages: list[ModelMessage]
@dataclass
class TurnResult:
"""What one turn produced: an answer, or a decision to make."""
text: str
pending: DeferredToolRequests | None = None
@property
def awaiting_approval(self) -> bool:
return self.pending is not None
class CatalogueAgent:
"""Holds what is stable across turns; rebuilds what must not be."""
def __init__(self, settings: Settings, http_client: httpx2.AsyncClient) -> None:
def __init__(
self,
settings: Settings,
http_client: httpx2.AsyncClient,
llm_client: httpx2.AsyncClient,
) -> None:
self._settings = settings
self._exchanger = TokenExchanger(
http_client,
@@ -86,32 +194,117 @@ class CatalogueAgent:
client_id=settings.client_id,
client_secret=settings.client_secret,
)
self._http_client = http_client
self._http_client = llm_client
self._history: OrderedDict[str, list[ModelMessage]] = OrderedDict()
self._pending: OrderedDict[str, PendingApproval] = OrderedDict()
async def run(self, prompt: str, subject_token: str, conversation_id: str) -> str:
async def run(self, prompt: str, subject_token: str, conversation_id: str) -> TurnResult:
"""One turn: exchange tokens, build the toolsets, answer, drop them.
The Agent is constructed here rather than once in __init__ because its
toolsets carry this turn's tokens. Construction is cheap; a long-lived
agent holding a long-lived token is not.
History is carried across turns of the same conversation. Without it the
confirm-before-writing rule cannot terminate: the agent restates what it
is about to create, the user says yes, and a historyless agent has no
idea what "yes" refers to.
`output_type` admits `DeferredToolRequests` so a write can end the turn
without being made. That is the whole mechanism: the run stops holding
the arguments the model proposed, and nothing reaches the gateway until
someone approves them.
"""
agent, deps = await self._build(subject_token)
async with agent:
result = await agent.run(
prompt,
deps=deps,
message_history=self._history.get(conversation_id),
usage_limits=UsageLimits(tool_calls_limit=self._settings.max_tool_calls),
)
return self._finish(conversation_id, result)
async def resume(
self, conversation_id: str, subject_token: str, *, approved: bool, reason: str = ""
) -> TurnResult:
"""Continue a turn that was waiting on a decision.
Approval is not enough on its own. The proposed arguments are validated
first, and a draft that fails is denied with the reason — so a write that
the user waved through is still refused if the model built it wrong.
"""
pending = self._pending.pop(conversation_id, None)
if pending is None:
raise LookupError("nothing is awaiting approval in this conversation")
results = DeferredToolResults()
for call in pending.requests.approvals:
if not approved:
results.approvals[call.tool_call_id] = ToolDenied(reason or "refused by the user")
continue
if problem := validate_draft(call.tool_name, call.args_as_dict()):
results.approvals[call.tool_call_id] = ToolDenied(
f"the proposed record is not valid and was not sent: {problem}"
)
continue
results.approvals[call.tool_call_id] = ToolApproved()
agent, deps = await self._build(subject_token)
async with agent:
result = await agent.run(
deps=deps,
message_history=pending.messages,
deferred_tool_results=results,
usage_limits=UsageLimits(tool_calls_limit=self._settings.max_tool_calls),
)
return self._finish(conversation_id, result)
def awaiting(self, conversation_id: str) -> PendingApproval | None:
return self._pending.get(conversation_id)
async def _build(self, subject_token: str) -> tuple[Agent, CallerDeps]:
toolsets = await build_toolsets(self._settings, self._exchanger, subject_token)
llm_token = await self._exchanger.exchange(subject_token, self._settings.llm_scope)
model = build_model(self._settings, self._http_client, llm_token)
agent = Agent(model, instructions=SYSTEM_PROMPT, toolsets=toolsets)
async with agent:
result = await agent.run(prompt, message_history=self._history.get(conversation_id))
agent = Agent(
model,
instructions=[SYSTEM_PROMPT, caller_instructions],
deps_type=CallerDeps,
toolsets=toolsets,
# A turn may end with an answer or with a decision to make. Without
# DeferredToolRequests in the union, a held tool call is an error
# rather than a state.
output_type=[str, DeferredToolRequests],
)
return agent, CallerDeps(user=caller_of(subject_token))
def _finish(self, conversation_id: str, result: Any) -> TurnResult:
self._remember(conversation_id, result.all_messages())
return result.output
output = result.output
if isinstance(output, DeferredToolRequests):
self._pending[conversation_id] = PendingApproval(output, result.all_messages())
self._pending.move_to_end(conversation_id)
while len(self._pending) > MAX_CONVERSATIONS:
self._pending.popitem(last=False)
return TurnResult(describe(output), pending=output)
return TurnResult(str(output))
def _remember(self, conversation_id: str, messages: list[ModelMessage]) -> None:
self._history[conversation_id] = messages
self._history.move_to_end(conversation_id)
while len(self._history) > MAX_CONVERSATIONS:
self._history.popitem(last=False)
def describe(requests: DeferredToolRequests) -> str:
"""What the user is being asked to approve, in their own reading.
Rendered from the arguments the model actually proposed, not from what it
said it would do. Those are the same thing only when the model is behaving,
and this exists for when it is not.
"""
lines = ["This will change data. Nothing has been sent yet.", ""]
for call in requests.approvals:
lines.append(f"**{call.tool_name}**")
body = (call.args_as_dict() or {}).get("body") or call.args_as_dict() or {}
for key, value in sorted(body.items()):
lines.append(f"- {key}: {value}")
lines.append("")
lines.append("Reply `approve` to proceed. Anything else cancels it.")
return "\n".join(lines)
+116
View File
@@ -0,0 +1,116 @@
"""Which tool calls need a human, and what a valid draft looks like.
Two things live here, and they are deliberately separate from the system prompt.
A prompt asks the model to behave; these decide. The model can be talked out of
an instruction — by a confused user, or by product data carrying an injection —
and cannot be talked out of either of the checks below.
Neither of them is authorisation. Who may write is settled by the MCP server
against the end user's own roles, before this agent sees a tool at all. What is
here makes the agent predictable, not privileged.
"""
from __future__ import annotations
from decimal import Decimal, InvalidOperation
from typing import Any
from pydantic import BaseModel, ConfigDict, Field, field_validator
from pydantic_ai.tools import ToolDefinition
# Fallback only. Tool names come from the API's operationId, so the verb is the
# first segment — but the annotations below are the real source, and this is
# what answers when a server declares none.
WRITE_VERBS = frozenset({"create", "update", "patch", "delete", "post", "put"})
def is_write(tool: ToolDefinition) -> bool:
"""True when calling this tool changes something.
Read from the MCP annotations the server publishes — `readOnlyHint` is part
of the MCP specification, and the gateway now emits it for every tool. So the
server decides what counts as a write, and the agent obeys; it does not carry
its own list of tool names to keep in step.
"""
annotations = (tool.metadata or {}).get("annotations") or {}
hint = annotations.get("readOnlyHint")
if isinstance(hint, bool):
return not hint
# No annotation: fall back to the verb, and treat anything unrecognised as a
# write. Asking for approval needlessly is a nuisance; skipping it is not.
name = tool.name.rsplit("_", 1)[-1]
return name.split("-", 1)[0].lower() in WRITE_VERBS or "-" not in name
def needs_approval(ctx: Any, tool: ToolDefinition, args: dict[str, Any]) -> bool:
"""The signature `ApprovalRequiredToolset` expects. Writes stop; reads run."""
return is_write(tool)
class ProductDraft(BaseModel):
"""What `create-product` is allowed to be called with.
The MCP server already validates against the API's own schema — required
fields, string lengths. This adds what a schema cannot say: that a price is
a number above zero, that the margin is not inverted, that nothing arrived
as an empty string because the model had nothing to put there.
It runs after the user approves and before the call leaves, so a draft that
fails here is refused with a reason rather than written and regretted.
"""
model_config = ConfigDict(extra="forbid")
productcode: str = Field(min_length=1, max_length=15)
productname: str = Field(min_length=1, max_length=70)
productline: str = Field(min_length=1)
productscale: str = Field(min_length=1, max_length=10)
productvendor: str = Field(min_length=1, max_length=50)
productdescription: str = Field(min_length=1)
quantityinstock: int = Field(ge=0, le=32767)
buyprice: Decimal = Field(gt=0)
msrp: Decimal = Field(gt=0)
@field_validator("buyprice", "msrp", mode="before")
@classmethod
def _decimal_from_string(cls, v: Any) -> Any:
# The API takes these as strings; a model that sends "" or "n/a" must
# fail here rather than at the backend, where the error is a 400 with no
# useful shape.
if isinstance(v, str):
try:
return Decimal(v)
except InvalidOperation as exc:
raise ValueError(f"not a number: {v!r}") from exc
return v
@field_validator("msrp")
@classmethod
def _sells_above_cost(cls, v: Decimal, info: Any) -> Decimal:
buy = info.data.get("buyprice")
if buy is not None and v <= buy:
raise ValueError(f"msrp {v} is not above buyprice {buy}")
return v
#: Tool name (unprefixed) → the model its arguments must satisfy. A write with no
#: entry here is approved on the user's word alone; that is the current state for
#: every domain but catalog, and it is visible rather than implied.
DRAFTS: dict[str, type[BaseModel]] = {"create-product": ProductDraft}
def validate_draft(tool_name: str, args: dict[str, Any]) -> str | None:
"""Return why this call must not proceed, or None if it may.
Arguments arrive shaped by the gateway: a request body sits under `body`,
query parameters under `query_*`, path parameters under `path_*`. Only the
body is checked — the rest is the server's own schema to enforce.
"""
model = DRAFTS.get(tool_name.rsplit("_", 1)[-1])
if model is None:
return None
try:
model.model_validate(args.get("body") or {})
except Exception as exc: # noqa: BLE001 — the message is the product here
return str(exc)
return None
+4
View File
@@ -49,6 +49,10 @@ class Settings(BaseSettings):
# do OIDC outbound, but it proves nothing about the end user — this agent
# can do better and does.
llm_scope: str = "llm-proxy-audience"
# A turn that loops has to stop somewhere. Twenty tool calls is far
# beyond any legitimate catalogue question and well under the point
# where the token budget would notice.
max_tool_calls: int = 20
# --- transport ----------------------------------------------------------
# The lab uses a private CA (sttlab-local-ca). Point this at the CA bundle
+18
View File
@@ -16,6 +16,8 @@ works, and the resulting token carries `aud: [mcp-catalog]`.
from __future__ import annotations
import base64
import json
import time
from dataclasses import dataclass
@@ -86,3 +88,19 @@ class TokenExchanger:
)
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"
+41 -14
View File
@@ -18,7 +18,14 @@ 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
@@ -30,6 +37,14 @@ class McpTarget:
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]:
@@ -41,7 +56,7 @@ def targets(settings: Settings) -> list[McpTarget]:
"""
return [
McpTarget("catalog", settings.mcp_catalog_url, settings.mcp_catalog_scope),
McpTarget("sales", settings.mcp_sales_url, settings.mcp_sales_scope),
McpTarget("sales", settings.mcp_sales_url, settings.mcp_sales_scope, read_only=True),
]
@@ -66,20 +81,32 @@ async def build_toolsets(
kwargs["timeout"] = httpx.Timeout(settings.http_timeout, connect=settings.http_timeout)
return httpx.AsyncClient(**kwargs)
toolsets: list[MCPToolset] = []
toolsets: list[AbstractToolset[None]] = []
for target in targets(settings):
token = await exchanger.exchange(subject_token, target.audience_scope)
toolsets.append(
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,
)
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