"""Observing one run's model and tool seams.
The capability built here is what an observed run installs for itself: it holds
the run's correlation state, so nothing is stored on the agent and two runs on
one built agent cannot read each other's identity.
Every seam is a *wrapper* hook. A wrapper owns both ends of the operation it
observes -- START, duration and exactly one terminal -- so pairing is a local
variable rather than a table keyed on an upstream id, and there is no window in
which a crash leaves a half-written pairing entry behind. That choice is forced
for tools: ``wrap_tool_execute`` is the only hook that sees every terminal
state, because a ``ModelRetry`` surfaces inside it as ``ToolRetryError`` and
never reaches ``on_tool_execute_error``.
What the seams deliberately do *not* do:
* They never refuse, delay or alter an operation. An observation failure is
contained (see :func:`_emit_safely`); an *observed* failure is re-raised
unchanged, including the original exception instance.
* They never merge two executions into one logical action. Upstream gives each
retry of a tool call a fresh ``tool_call_id``, so there is no stable identity
to merge on, and inferring one from tool name, arguments or timing would
silently fuse two concurrent calls of the same tool.
* They never read an application-defined conversion. Content is projected from
an explicit list of attributes and anything else is skipped, not stringified.
Limits of what these seams can see, measured against the supported floor:
* A capability that refuses a tool from ``before_tool_execute`` runs *outside*
the ``wrap_tool_execute`` frame, so such a tool produces no execution action
at all -- not a START left without a terminal. The failure is still visible on
the run.
* A tool that exceeds its own upstream ``timeout`` arrives as a
``ToolRetryError`` whose *message* mentions a timeout. The type says nothing,
so the attempt is reported as an ordinary error; TIMEOUT is reserved for an
expiry that is visible as a type.
* Losing the worker outside any seam leaves the attempt with no terminal action.
On the supported Pydantic AI floor, a ``BaseException`` escaping a sync tool
cancels the enclosing run task: the tool stays partial while the run records
that observed cancellation.
"""
from __future__ import annotations
from asyncio import CancelledError
from contextvars import ContextVar
from time import monotonic
from typing import Any, Awaitable, Callable, Final, NamedTuple
from uuid import UUID, uuid4
from edri.integrations.pydantic_ai._optional import (
CallDeferred,
DeferredToolRequests,
Hooks,
)
from edri.integrations.pydantic_ai.normalization import REQUEST_USAGE_FIELDS
from edri.integrations.pydantic_ai.observability import (
AgentActionNamespace,
AgentActionOperation,
AgentActionOutcome,
AgentActionPhase,
_ActionEmitter,
_Capture,
_EffectivePolicy,
_RunState,
_emit_safely,
)
__all__ = ["audit_hooks"]
class _RunObservation(NamedTuple):
"""What an observer inside a run needs to publish an action of its own."""
state: _RunState
emitter: _ActionEmitter
policy: _EffectivePolicy
#: The observed run this code is executing inside, or ``None``.
#:
#: A seam EDRI does not register itself -- an MCP ``process_tool_call``, which
#: the application installs on its own toolset -- has no other way to reach the
#: run's emitter and correlation state. It is a context variable rather than
#: anything global for the same reason as :data:`_CURRENT_TOOL`: several runs,
#: and several agents, are in flight at once.
_CURRENT_RUN: ContextVar[_RunObservation | None] = ContextVar(
"edri_pydantic_ai_current_run", default=None
)
def _current_run() -> _RunObservation | None:
"""The observed run enclosing this call, if there is one.
``None`` means EDRI is not observing here -- observability is off, or the
call is outside any run, as with a shared MCP connection's own traffic.
Attribution is then unknown, and an unknown attribution is not published.
"""
return _CURRENT_RUN.get()
#: The tool execution enclosing whatever runs next in this context, as
#: ``(root_action_id, action_id)``. A sync tool body runs on an AnyIO worker
#: thread, but that thread receives a *copy of the context*, so an EDRI dispatch
#: made from inside a tool can find the tool that made it -- which a thread
#: local or an attribute on the agent could not, since several tools of several
#: runs are in flight at once.
#:
#: The root id travels with it so a reader can tell whether the enclosing tool
#: belongs to the same run. An agent run nested inside another agent's tool
#: would otherwise adopt the outer agent's tool as the parent of its own
#: dispatches, which are children of the inner run.
_CURRENT_TOOL: ContextVar[tuple[UUID, UUID] | None] = ContextVar(
"edri_pydantic_ai_current_tool", default=None
)
def _current_tool_action(root_action_id: UUID) -> UUID | None:
"""The tool action enclosing this call within the run owning ``root_action_id``."""
enclosing = _CURRENT_TOOL.get()
if enclosing is None or enclosing[0] != root_action_id:
return None
return enclosing[1]
def _classify_failure(
error: BaseException, *, owns_deadline: bool = False
) -> tuple[AgentActionPhase, AgentActionOutcome] | None:
"""How a failed operation terminated, or ``None`` if it cannot be said.
Checked in this order because the types overlap: ``CallDeferred`` is an
ordinary ``Exception`` subclass, so testing for a generic error first would
publish a deferral as a failure, and ``CancelledError`` is a
``BaseException`` that a generic test would miss entirely. Only an
explicitly typed signal is classified -- a timeout is recognised because
``TimeoutError`` *is* the timeout, never because an exception's text
mentions one.
``owns_deadline`` says whether the observed operation is the one holding
the deadline that could expire (an EDRI dispatch does; a run, a model
request and a tool do not, on this floor). A ``TimeoutError`` crossing an
operation that owns no deadline is that operation failing because something
below it timed out, and the operation below already published the TIMEOUT.
``None`` means EDRI will not describe this terminal at all: a
``KeyboardInterrupt`` or a ``SystemExit`` unwinding the process is not an
outcome of the operation. The attempt is left without a terminal action,
which the inspector represents as exactly that.
"""
if isinstance(error, CallDeferred):
return AgentActionPhase.END, AgentActionOutcome.DEFERRED
if isinstance(error, CancelledError):
return AgentActionPhase.END, AgentActionOutcome.CANCELLED
if owns_deadline and isinstance(error, TimeoutError):
return AgentActionPhase.END, AgentActionOutcome.TIMEOUT
if isinstance(error, Exception):
return AgentActionPhase.ERROR, AgentActionOutcome.ERROR
return None
#: Outcomes that describe control flow rather than a failure. An action carrying
#: one of them publishes no error description: there is no failure to describe,
#: and the exception exists only because that is how the library unwinds.
_CONTROL_OUTCOMES: Final = frozenset(
{AgentActionOutcome.DEFERRED, AgentActionOutcome.CANCELLED}
)
def _emit_failure(
emitter: _ActionEmitter,
error: BaseException,
*,
started: float,
shared: dict[str, Any],
extra: dict[str, Any] | None = None,
owns_deadline: bool = False,
) -> None:
"""Publish the terminal action of a failed operation, if it has one.
Shared by the asynchronous seams and by the synchronous EDRI dispatch
wrapper, so the classification and the ordering it depends on are written
once. The caller re-raises; nothing here alters control flow.
"""
terminal = _classify_failure(error, owns_deadline=owns_deadline)
if terminal is None:
return
phase, outcome = terminal
_emit_safely(
emitter,
phase=phase,
outcome=outcome,
duration=monotonic() - started,
error=None if outcome in _CONTROL_OUTCOMES else error,
**(extra or {}),
**shared,
)
async def _observed(
emitter: _ActionEmitter,
handler: Callable[[], Awaitable[Any]],
*,
shared: dict[str, Any],
request: _Capture | None = None,
terminal_context: Callable[[], dict[str, Any]] | None = None,
terminal: Callable[[Any], dict[str, Any]] | None = None,
) -> Any:
"""Run one operation between a START and exactly one terminal action.
``BaseException`` is caught rather than ``Exception`` so that cancellation
-- which is not an ``Exception`` -- can be observed before the operation
unwinds. Everything caught here is re-raised unchanged, including the
original instance: the observation reports control flow, it never converts
or absorbs it. :func:`_classify_failure` decides what, if anything, the
terminal action says.
"""
_emit_safely(emitter, phase=AgentActionPhase.START, request=request, **shared)
started = monotonic()
try:
result = await handler()
except BaseException as error:
_emit_failure(
emitter,
error,
started=started,
shared=shared,
extra=_extra(terminal_context),
)
raise
extra = {**_extra(terminal_context), **_extra(terminal, result)}
_emit_safely(
emitter,
phase=AgentActionPhase.END,
outcome=extra.pop("outcome", AgentActionOutcome.SUCCESS),
duration=monotonic() - started,
**extra,
**shared,
)
return result
def _extra(source: Callable[..., dict[str, Any]] | None, *args: Any) -> dict[str, Any]:
"""Terminal fields from one describing callable, or nothing.
These callables read upstream objects -- a result's ``output``, a response's
``usage`` -- through properties that can raise. Losing the fields is
acceptable; losing the terminal action, or changing what the observed
operation returns, is not.
"""
if source is None:
return {}
try:
return source(*args)
except Exception:
return {}
def _text(value: Any) -> str | None:
"""``value`` if it is exactly a string. No conversion, ever."""
return value if type(value) is str else None
def _attempt(context: Any) -> int:
"""The one-based attempt number upstream reports for this execution.
``RunContext.retry`` counts retries from zero. It is read defensively
because an attempt number is metadata: a surprising value must not cost the
action, and :class:`AgentAction` rejects anything below one.
"""
retry = getattr(context, "retry", 0)
if type(retry) is not int or retry < 0:
return 1
return retry + 1
def _model_name(request_context: Any) -> str | None:
model = getattr(request_context, "model", None)
return _text(getattr(model, "model_name", None))
def _request_payload(request_context: Any) -> list[dict[str, Any]]:
"""The prompt parts of a model request, projected onto built-ins.
An explicit projection rather than a generic walk: the messages are upstream
objects whose parts hold files, binary content, thinking traces and provider
payloads. Only the kinds named here are published; anything else is skipped
rather than rendered.
A tool call and the return that answered it are both published, so the
request shows what the model was actually given and not only what it asked
for. They carry ``tool_call_id`` because a model can ask for one tool twice
in a single turn, which leaves the name and the position ambiguous -- it is
also what pairs a part here with the tool action that produced it.
Prompt text and tool arguments are published only when they already *are* a
string or a mapping. A tool return is not filtered that way: its value is
whatever the tool returned, an integer or a list as easily as text, and
:mod:`~edri.integrations.pydantic_ai.normalization` is what decides how much
of it survives.
"""
payload: list[dict[str, Any]] = []
for message in getattr(request_context, "messages", ()) or ():
for part in getattr(message, "parts", ()) or ():
kind = _text(getattr(part, "part_kind", None))
if kind in ("system-prompt", "user-prompt"):
content = _text(getattr(part, "content", None))
if content is not None:
payload.append({"kind": kind, "content": content})
elif kind == "tool-call":
arguments = getattr(part, "args", None)
if type(arguments) in (str, dict):
payload.append(
{
"kind": kind,
"tool_name": _text(getattr(part, "tool_name", None)),
"tool_call_id": _text(
getattr(part, "tool_call_id", None)
),
"args": arguments,
}
)
elif kind == "tool-return":
payload.append(
{
"kind": kind,
"tool_name": _text(getattr(part, "tool_name", None)),
"tool_call_id": _text(getattr(part, "tool_call_id", None)),
"content": getattr(part, "content", None),
}
)
return payload
def _response_payload(response: Any) -> dict[str, Any]:
"""The completion a model produced: its text and the tool calls it asked for."""
calls = []
for call in getattr(response, "tool_calls", ()) or ():
arguments = getattr(call, "args", None)
calls.append(
{
"tool_name": _text(getattr(call, "tool_name", None)),
"args": arguments if type(arguments) in (str, dict) else None,
}
)
return {"text": _text(getattr(response, "text", None)), "tool_calls": calls}
[docs]
def audit_hooks(
state: _RunState, emitter: _ActionEmitter, policy: _EffectivePolicy
) -> Hooks:
"""The audit capability for one run.
Registered per run through ``run_sync(capabilities=[...])``, so it is not
installed on the agent and does not persist beyond the invocation it
observes.
"""
hooks = Hooks()
def _run_shared() -> dict[str, Any]:
return {
"action_id": state.root_action_id,
# The run action is its own root: one key buckets every action of
# this invocation, including this one.
"root_action_id": state.root_action_id,
"namespace": AgentActionNamespace.AGENT,
"operation": AgentActionOperation.RUN,
"run_id": state.run_id,
"conversation_id": state.conversation_id,
}
def _child_shared(namespace: str, operation: str) -> dict[str, Any]:
return {
"action_id": uuid4(),
"namespace": namespace,
"operation": operation,
"parent_action_id": state.root_action_id,
"root_action_id": state.root_action_id,
"run_id": state.run_id,
"conversation_id": state.conversation_id,
}
@hooks.on.run
async def _observe_run(ctx: Any, *, handler: Any) -> Any:
"""START and the run's single terminal action.
``wrap_run`` is the outermost seam upstream offers and can enter before
``before_run``, so START means "EDRI observed entry to this seam" and
nothing stronger.
"""
# Written once, here, before anything can read it: this is the first
# point at which upstream run identity exists.
state.run_id = ctx.run_id
state.conversation_id = ctx.conversation_id
# Published for the whole run, so a seam EDRI cannot register -- an
# application's MCP ``process_tool_call`` -- can find this run. The
# token is reset on the way out: the run's identity must not outlive it
# in a context another run could inherit.
run_token = _CURRENT_RUN.set(_RunObservation(state, emitter, policy))
try:
return await _run_observed(ctx, handler)
finally:
_CURRENT_RUN.reset(run_token)
async def _run_observed(ctx: Any, handler: Any) -> Any:
"""The observation itself, reached only after identity is bound above."""
def _terminal(result: Any) -> dict[str, Any]:
if isinstance(getattr(result, "output", None), DeferredToolRequests):
# The run finished this invocation without finishing its work:
# tools were handed back to the caller to resolve. A later
# resume is only the same logical run if upstream says so.
return {"outcome": AgentActionOutcome.DEFERRED}
return {}
return await _observed(
emitter,
lambda: handler(),
shared=_run_shared(),
# Run totals, accumulated across the whole run by the context. A
# per-request delta is a different measurement, published on the
# model actions and never to be summed with this one.
terminal_context=lambda: {"usage": getattr(ctx, "usage", None)},
terminal=_terminal,
)
@hooks.on.model_request
async def _observe_model_request(
ctx: Any, *, request_context: Any, handler: Any
) -> Any:
shared = _child_shared(
AgentActionNamespace.MODEL, AgentActionOperation.REQUEST
)
shared["name"] = _model_name(request_context)
shared["metadata"] = {
"streaming": bool(getattr(request_context, "streaming", False)),
"model_id": _text(getattr(request_context, "model_id", None)),
}
def _terminal(response: Any) -> dict[str, Any]:
return {
# What actually served this request, which is not necessarily
# what the agent declared at setup.
"provider": _text(getattr(response, "provider_name", None)),
"usage": getattr(response, "usage", None),
"usage_counters": REQUEST_USAGE_FIELDS,
"response_data": _Capture(
_response_payload(response), policy.include_model_responses
),
}
return await _observed(
emitter,
lambda: handler(request_context),
shared=shared,
request=_Capture(
_request_payload(request_context), policy.include_model_requests
),
terminal=_terminal,
)
def _tool_shared(
ctx: Any, call: Any, tool_def: Any, operation: str
) -> dict[str, Any]:
shared = _child_shared(AgentActionNamespace.TOOL, operation)
shared["name"] = _text(getattr(tool_def, "name", None))
shared["attempt"] = _attempt(ctx)
shared["metadata"] = {
# Fresh on every retry, so it identifies this execution -- and,
# within one attempt, links validation to execution.
"tool_call_id": _text(getattr(call, "tool_call_id", None)),
"tool_kind": _text(getattr(tool_def, "kind", None)),
}
return shared
@hooks.on.tool_validate
async def _observe_tool_validate(
ctx: Any, *, call: Any, tool_def: Any, args: Any, handler: Any
) -> Any:
return await _observed(
emitter,
lambda: handler(args),
shared=_tool_shared(ctx, call, tool_def, AgentActionOperation.VALIDATE),
# The arguments as the model produced them, before validation: what
# a rejected call was actually asked to do.
request=_Capture(args, policy.include_tool_arguments),
)
@hooks.on.tool_execute
async def _observe_tool_execute(
ctx: Any, *, call: Any, tool_def: Any, args: Any, handler: Any
) -> Any:
shared = _tool_shared(ctx, call, tool_def, AgentActionOperation.EXECUTE)
token = _CURRENT_TOOL.set((state.root_action_id, shared["action_id"]))
try:
return await _observed(
emitter,
lambda: handler(args),
shared=shared,
request=_Capture(args, policy.include_tool_arguments),
terminal=lambda result: {
"response_data": _Capture(result, policy.include_tool_results)
},
)
finally:
_CURRENT_TOOL.reset(token)
return hooks