diff --git a/src/ai/agents/_middleware.py b/src/ai/agents/_middleware.py deleted file mode 100644 index 4335c6d3..00000000 --- a/src/ai/agents/_middleware.py +++ /dev/null @@ -1,378 +0,0 @@ -"""Middleware: composable wrappers around all execution surfaces. - -Middleware is run-scoped — pass it to :meth:`Agent.run`:: - - agent.run(model, messages, middleware=[LoggingMiddleware()]) - -Middleware wraps agent runs, model calls, generate calls, tool calls, and -hook calls. Subclass :class:`Middleware` and override the methods you care -about — unimplemented methods pass through to the next middleware (or the -real implementation). - -Ordering: first in the list = outermost. ``[A(), B()]`` means A wraps B -wraps the real call. A sees the call first and the result last. -""" - -from __future__ import annotations - -import contextlib -import contextvars -import dataclasses -from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence -from typing import TYPE_CHECKING, Any - -from ..types import messages as messages_ - -# Compat shim: ``StreamResultLike`` was removed from ``ai.types.proto`` when -# the model layer was reworked. Middleware is dead code under the new -# ``Executor``-based ``api.py`` and is kept around only so the agents -# rewrite can land separately; ``Any`` is enough to keep the existing -# annotations type-checking. -type StreamResultLike = Any - -# --------------------------------------------------------------------------- -# Call context objects — frozen dataclasses with isolated mutable fields. -# -# Mutable container fields (``list``, ``dict``) are shallow-copied at -# construction via ``__post_init__`` so that middleware sees its own copy -# and cannot accidentally mutate the caller's data. To modify fields, -# use ``dataclasses.replace(call, messages=new_msgs)`` before passing -# to ``next``. -# --------------------------------------------------------------------------- - -if TYPE_CHECKING: - import pydantic - - from ..models.core.model import Model - from ..models.core.params import InferenceRequestParams - from ..types import events as events_ - from ..types.tools import Tool - from .agent import Context - from .hooks import HookRegistry - - -@dataclasses.dataclass(frozen=True) -class ModelContext: - """Context for a model streaming call.""" - - model: Model - messages: list[messages_.Message] - tools: Sequence[Tool] | None - output_type: type[pydantic.BaseModel] | None - kwargs: dict[str, Any] - - def __post_init__(self) -> None: - object.__setattr__(self, "messages", list(self.messages)) - if self.tools is not None: - object.__setattr__(self, "tools", list(self.tools)) - object.__setattr__(self, "kwargs", dict(self.kwargs)) - - -@dataclasses.dataclass(frozen=True) -class GenerateContext: - """Context for a non-streaming model call.""" - - model: Model - messages: list[messages_.Message] - tools: Sequence[Tool] | None - output_type: type[pydantic.BaseModel] | None - params: InferenceRequestParams | None - - def __post_init__(self) -> None: - object.__setattr__(self, "messages", list(self.messages)) - if self.tools is not None: - object.__setattr__(self, "tools", list(self.tools)) - - -@dataclasses.dataclass(frozen=True) -class ToolContext: - """Context for a tool execution.""" - - tool_call_id: str - tool_name: str - kwargs: dict[str, Any] - - def __post_init__(self) -> None: - object.__setattr__(self, "kwargs", dict(self.kwargs)) - - -@dataclasses.dataclass(frozen=True) -class HookContext: - """Context for a hook suspension point.""" - - label: str - payload: type[pydantic.BaseModel] - metadata: dict[str, Any] - tool_call_id: str | None = None - registry: HookRegistry | None = None - - def __post_init__(self) -> None: - object.__setattr__(self, "metadata", dict(self.metadata)) - - -# --------------------------------------------------------------------------- -# Middleware base class — override the methods you care about. -# --------------------------------------------------------------------------- - -# Event/message aliases for brevity in signatures. ``_Event`` is intentionally -# typed as ``Any`` so the agent-run chain accepts the wider ``AgentEvent`` -# union (which includes ``ToolCallResult``/``HookEvent``) without a circular -# import from ``ai.agents``. -_Event = Any -_Message = messages_.Message - -# Agent run next-function type: call -> async generator of events. -_AgentRunNext = Callable[["Context"], AsyncGenerator[_Event]] - - -class _Middleware: - """Base middleware class. Override the methods you need. - - Default implementations call ``next(call)`` — a transparent pass-through. - """ - - async def wrap_agent_run( - self, - call: Context, - next: _AgentRunNext, - ) -> AsyncGenerator[_Event]: - """Wrap an agent run. - - ``next(call)`` returns an async generator of ``Event`` objects. - Override to add tracing, durability checkpoints, or other - run-scoped behavior:: - - async def wrap_agent_run(self, call, next): - span = start_span("agent.run") - async for event in next(call): - yield event - span.end() - """ - # aclosing: propagate an early close from the consumer down the - # chain, so inner generators unwind in the task that runs them - async with contextlib.aclosing(next(call)) as events: - async for event in events: - yield event - - async def wrap_model( - self, - call: ModelContext, - next: Callable[[ModelContext], Awaitable[StreamResultLike]], - ) -> StreamResultLike: - """Wrap a model streaming call. - - ``next(call)`` returns a :class:`~ai.types.StreamResultLike` that - is async-iterable over ``Event`` objects. You can do work - before, iterate / transform the stream, or do cleanup after. - - To transform the stream, use - :meth:`~ai.models.StreamResult.from_generator`:: - - async def wrap_model(self, call, next): - stream = await next(call) - async def _add_suffix(): - async for event in stream: - yield event - from ai.models import StreamResult - return StreamResult.from_generator(_add_suffix()) - """ - return await next(call) - - async def wrap_generate( - self, - call: GenerateContext, - next: Callable[[GenerateContext], Awaitable[_Message]], - ) -> _Message: - """Wrap a non-streaming model call.""" - return await next(call) - - async def wrap_tool( - self, - call: ToolContext, - next: Callable[[ToolContext], Awaitable[events_.ToolCallResult]], - ) -> events_.ToolCallResult: - """Wrap a tool execution. - - ``next(call)`` returns a :class:`ToolCallResult`. - """ - return await next(call) - - async def wrap_hook( - self, - call: HookContext, - next: Callable[[HookContext], Awaitable[pydantic.BaseModel]], - ) -> pydantic.BaseModel: - """Wrap a hook suspension point. - - ``next(call)`` blocks until the hook is resolved and returns the - validated payload instance. - """ - return await next(call) - - -# --------------------------------------------------------------------------- -# Run-scoped middleware via ContextVar -# --------------------------------------------------------------------------- - -_active: contextvars.ContextVar[list[_Middleware]] = contextvars.ContextVar( - "middleware", -) - -_EMPTY: list[_Middleware] = [] - - -def get() -> list[_Middleware]: - """Return the middleware stack for the current run (empty if none).""" - return _active.get(_EMPTY) - - -Token = contextvars.Token[list[_Middleware]] - - -def activate(mw: list[_Middleware]) -> Token: - """Set the middleware stack and return a reset token.""" - return _active.set(mw) - - -def deactivate(token: Token) -> None: - """Restore the previous middleware stack.""" - _active.reset(token) - - -# --------------------------------------------------------------------------- -# Chain builders — compose the middleware stack for each surface. -# -# Each builder takes the *real* implementation as a callable and returns -# a callable with the same signature that routes through middleware. -# -# When no middleware is active, the real implementation is returned -# directly — zero overhead. -# --------------------------------------------------------------------------- - - -def _build_model_chain( - real: Callable[[ModelContext], Awaitable[StreamResultLike]], -) -> Callable[[ModelContext], Awaitable[StreamResultLike]]: - mw = get() - if not mw: - return real - - chain = real - for m in reversed(mw): - - def _make( - m: _Middleware, - nxt: Callable[[ModelContext], Awaitable[StreamResultLike]], - ) -> Callable[[ModelContext], Awaitable[StreamResultLike]]: - async def _wrapped(call: ModelContext) -> StreamResultLike: - return await m.wrap_model(call, nxt) - - return _wrapped - - chain = _make(m, chain) - return chain - - -def _build_generate_chain( - real: Callable[[GenerateContext], Awaitable[_Message]], -) -> Callable[[GenerateContext], Awaitable[_Message]]: - mw = get() - if not mw: - return real - - chain = real - for m in reversed(mw): - - def _make( - m: _Middleware, - nxt: Callable[[GenerateContext], Awaitable[_Message]], - ) -> Callable[[GenerateContext], Awaitable[_Message]]: - async def _wrapped(call: GenerateContext) -> _Message: - return await m.wrap_generate(call, nxt) - - return _wrapped - - chain = _make(m, chain) - return chain - - -def _build_tool_chain( - real: Callable[[ToolContext], Awaitable[events_.ToolCallResult]], -) -> Callable[[ToolContext], Awaitable[events_.ToolCallResult]]: - mw = get() - if not mw: - return real - - chain = real - for m in reversed(mw): - - def _make( - m: _Middleware, - nxt: Callable[[ToolContext], Awaitable[events_.ToolCallResult]], - ) -> Callable[[ToolContext], Awaitable[events_.ToolCallResult]]: - async def _wrapped(call: ToolContext) -> events_.ToolCallResult: - return await m.wrap_tool(call, nxt) - - return _wrapped - - chain = _make(m, chain) - return chain - - -def _build_hook_chain( - real: Callable[[HookContext], Awaitable[pydantic.BaseModel]], -) -> Callable[[HookContext], Awaitable[pydantic.BaseModel]]: - mw = get() - if not mw: - return real - - chain = real - for m in reversed(mw): - - def _make( - m: _Middleware, - nxt: Callable[[HookContext], Awaitable[pydantic.BaseModel]], - ) -> Callable[[HookContext], Awaitable[pydantic.BaseModel]]: - async def _wrapped(call: HookContext) -> pydantic.BaseModel: - return await m.wrap_hook(call, nxt) - - return _wrapped - - chain = _make(m, chain) - return chain - - -def _build_agent_run_chain( - real: _AgentRunNext, -) -> _AgentRunNext: - mw = get() - if not mw: - return real - - chain = real - for m in reversed(mw): - - def _make(m: _Middleware, nxt: _AgentRunNext) -> _AgentRunNext: - async def _wrapped(call: Context) -> AsyncGenerator[_Event]: - gen = m.wrap_agent_run(call, nxt) - async with contextlib.aclosing(gen) as events: - async for event in events: - yield event - - return _wrapped - - chain = _make(m, chain) - return chain - - -__all__ = [ - "GenerateContext", - "HookContext", - "ModelContext", - "StreamResultLike", - "ToolContext", - "_Middleware", - "activate", - "deactivate", - "get", -] diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 161afff5..56852db7 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -42,7 +42,6 @@ from ..types import builders from ..types import events as events_ from ..types.messages import MessageBundle -from . import _middleware as middleware_ from . import hooks as hooks_ from . import runtime @@ -757,36 +756,29 @@ async def _execute( message=msg, results=msg.tool_results ) - # Best-effort parse so middleware sees usable kwargs when possible. - # If parsing fails, middleware still gets the raw tool_call_id / - # tool_name and can replace kwargs before _real() executes. + # If the model's args are bad, overrides can still supply them; + # otherwise the parse error becomes the tool result. + parse_error: Exception | None = None try: - base_kwargs = self.kwargs - except Exception: - base_kwargs = {} + kwargs = self.kwargs + except Exception as exc: + kwargs = {} + parse_error = exc if overrides: # Overrides come from user code, not the model — validate # eagerly so programming errors surface immediately. - base_kwargs = _validate_kwargs( - self._tool, {**base_kwargs, **overrides} - ) - - call = middleware_.ToolContext( - tool_call_id=self._part.tool_call_id, - tool_name=self._part.tool_name, - kwargs=base_kwargs, - ) + kwargs = _validate_kwargs(self._tool, {**kwargs, **overrides}) + parse_error = None tool = self._tool - - async def _real( - call: middleware_.ToolContext, - ) -> events_.ToolCallResult: + data.args = kwargs + async with telemetry.span(data) as sp: result: Any model_input: Any = types.messages.MODEL_INPUT_UNSET try: - kwargs = _validate_kwargs(tool, call.kwargs) + if parse_error is not None: + raise parse_error # The returned value decides how the tool runs, so a # plain `def` returning a coroutine works too. returned = tool.fn(**kwargs) @@ -804,8 +796,8 @@ async def _real( ) agg = await _aggregate_from( returned, - tool_call_id=call.tool_call_id, - tool_name=call.tool_name, + tool_call_id=self.id, + tool_name=self.name, aggregator=tool.aggregator, ) result = agg.snapshot() @@ -821,30 +813,26 @@ async def _real( if tool.to_model_input is not None: model_input = tool.to_model_input(result) except (Exception, asyncio.CancelledError) as exc: - return _error_tool_result( - exc, - tool_call_id=call.tool_call_id, - tool_name=call.tool_name, + res = _error_tool_result( + exc, tool_call_id=self.id, tool_name=self.name ) - part = types.messages.ToolResultPart( - tool_call_id=call.tool_call_id, - tool_name=call.tool_name, - result=result, - result_kind=types.messages.ToolResultPart.kind_for(result), - model_input=model_input, - ) - return tool_result(part) - - data.args = base_kwargs - chain = middleware_._build_tool_chain(_real) - async with telemetry.span(data) as sp: - res = await chain(call) - if res.results: - sp.data.result = res.results[0].result - sp.data.is_error = any(p.is_error for p in res.results) + else: + res = tool_result( + types.messages.ToolResultPart( + tool_call_id=self.id, + tool_name=self.name, + result=result, + result_kind=types.messages.ToolResultPart.kind_for( + result + ), + model_input=model_input, + ) + ) + sp.data.result = res.results[0].result + sp.data.is_error = res.results[0].is_error # A tool exception is caught and converted to an error - # result before it reaches this block, so it never hits the - # span's own except path — thread it through explicitly. + # result above, so it never hits the span's own except + # path — thread it through explicitly. if res.exception is not None: sp.error = telemetry.SpanError.from_exception(res.exception) return res @@ -1481,7 +1469,6 @@ def run( *, params: models.InferenceRequestParams | None = None, hook_registry: hooks_.HookRegistry | None = None, - _middleware: list[middleware_._Middleware] | None = None, ) -> AbstractAsyncContextManager[AgentStream[str]]: ... @overload def run[T: pydantic.BaseModel]( @@ -1492,7 +1479,6 @@ def run[T: pydantic.BaseModel]( output_type: type[T], params: models.InferenceRequestParams | None = None, hook_registry: hooks_.HookRegistry | None = None, - _middleware: list[middleware_._Middleware] | None = None, ) -> AbstractAsyncContextManager[AgentStream[T]]: ... def run( self, @@ -1502,7 +1488,6 @@ def run( output_type: type[pydantic.BaseModel] | None = None, params: models.InferenceRequestParams | None = None, hook_registry: hooks_.HookRegistry | None = None, - _middleware: list[middleware_._Middleware] | None = None, ) -> AbstractAsyncContextManager[AgentStream[Any]]: """Run the agent loop, yielding events to the consumer. @@ -1537,7 +1522,6 @@ def run( output_type=output_type, params=params, hook_registry=hook_registry, - _middleware=_middleware, ) @contextlib.asynccontextmanager @@ -1549,7 +1533,6 @@ async def _run( output_type: type[pydantic.BaseModel] | None, params: models.InferenceRequestParams | None, hook_registry: hooks_.HookRegistry | None, - _middleware: list[middleware_._Middleware] | None, ) -> AsyncIterator[AgentStream[Any]]: context = Context( model=model, @@ -1571,9 +1554,9 @@ async def _run( # shares the enclosing run's hooks); otherwise start fresh. registry = hooks_._hook_registry.get(None) or hooks_.HookRegistry() - async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]: + async def _events() -> AsyncGenerator[events_.AgentEvent]: tracker = events_.RunStateTracker() - source = self.loop(call) + source = self.loop(context) async with contextlib.aclosing(runtime.run(source)) as events: async for event in events: # Feed the tracker before the replay filter: replayed @@ -1590,13 +1573,6 @@ async def _real(call: Context) -> AsyncGenerator[events_.AgentEvent]: yield transition async def _stream() -> AsyncGenerator[events_.AgentEvent]: - # Activate middleware for this run (and everything it calls). - # When middleware is None (default), inherit the parent's - # middleware from the context var — this lets nested agents - # share middleware. When middleware is explicitly provided, - # *extend* the parent stack so that outer cross-cutting - # concerns (tracing, durability) are preserved. Pass - # ``_middleware=[]`` to clear the stack entirely. async with telemetry.span( telemetry.RunSpanData( agent=type(self).__name__, @@ -1614,14 +1590,9 @@ async def _stream() -> AsyncGenerator[events_.AgentEvent]: ) as sp: initial_count = len(context.messages) agent_token = _current_agent.set(self) - mw_token: middleware_.Token | None = None - if _middleware is not None: - parent = middleware_.get() - mw_token = middleware_.activate(parent + _middleware) try: - chain = middleware_._build_agent_run_chain(_real) async with contextlib.aclosing( - util.decouple(chain(context), buffer=self.LOOP_BUFFER) + util.decouple(_events(), buffer=self.LOOP_BUFFER) ) as events: async for event in events: # a blocked run can only resume via a hook @@ -1636,8 +1607,6 @@ async def _stream() -> AsyncGenerator[events_.AgentEvent]: sp.data.blocked = False yield event finally: - if mw_token is not None: - middleware_.deactivate(mw_token) _current_agent.reset(agent_token) # Record whatever got produced, even on error or # early close. diff --git a/src/ai/agents/hooks.py b/src/ai/agents/hooks.py index 16ff15ee..049ac0ae 100644 --- a/src/ai/agents/hooks.py +++ b/src/ai/agents/hooks.py @@ -23,14 +23,13 @@ import asyncio import contextlib import contextvars -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any import pydantic from .. import experimental_telemetry as telemetry from .. import types from ..types import messages as messages_ -from . import _middleware as middleware_ from . import runtime as runtime_ if TYPE_CHECKING: @@ -175,32 +174,16 @@ async def hook[T: pydantic.BaseModel]( Defaults to the current registry. """ - call = middleware_.HookContext( - label=_label(hook), - payload=payload, - metadata=metadata or {}, - tool_call_id=tool_call_id, - registry=registry, - ) - - chain = middleware_._build_hook_chain(_hook_impl) - result = await chain(call) - return cast("T", result) - - -async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: - """Core hook logic — the innermost ``next`` in the middleware chain.""" rt = runtime_.get_runtime() - registry = _registry(call.registry) - label = call.label - payload = call.payload - hook_metadata = call.metadata + registry = _registry(registry) + label = _label(hook) + hook_metadata = dict(metadata or {}) data = telemetry.HookSpanData( label=label, hook_type=payload.__name__, metadata=hook_metadata, - tool_call_id=call.tool_call_id, + tool_call_id=tool_call_id, ) # Pre-registered resolution (serverless re-entry). @@ -226,7 +209,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: hook_type=payload.__name__, status="pending", metadata=hook_metadata, - tool_call_id=call.tool_call_id, + tool_call_id=tool_call_id, ) await rt.put_hook(hook_part) @@ -262,7 +245,7 @@ async def _hook_impl(call: middleware_.HookContext) -> pydantic.BaseModel: status="resolved", metadata=hook_metadata, resolution=resolution, - tool_call_id=call.tool_call_id, + tool_call_id=tool_call_id, ) ) diff --git a/tests/agents/test_tools.py b/tests/agents/test_tools.py index b381d648..6fe85539 100644 --- a/tests/agents/test_tools.py +++ b/tests/agents/test_tools.py @@ -415,6 +415,25 @@ async def double(x: int) -> int: result = await tc() assert result.results[0].is_error + assert "JSONDecodeError" in str(result.results[0].result) + + +async def test_tool_call_overrides_repair_malformed_args() -> None: + @ai.tool + async def double(x: int) -> int: + """Double a number.""" + return x * 2 + + part = ai.messages.ToolCallPart( + tool_call_id="tc-1", + tool_name="double", + tool_args='{"x": ', + ) + tc = ai.agents.BoundToolCall(part=part, tool=double) + + result = await tc(x=99) + + assert result.results[0].result == 198 # -- Helpers --------------------------------------------------------------- diff --git a/tests/test_middleware.py b/tests/test_middleware.py deleted file mode 100644 index d06ec08c..00000000 --- a/tests/test_middleware.py +++ /dev/null @@ -1,282 +0,0 @@ -"""Middleware: run-scoped registration, chain ordering, per-surface wrapping.""" - -from __future__ import annotations - -import dataclasses -from collections.abc import AsyncGenerator -from typing import Any - -import pydantic -import pytest - -import ai -from ai.agents import _middleware as middleware -from ai.providers import history_utils -from ai.types import events as agent_events_ - -from .conftest import ( - MOCK_MODEL, - collect_messages, - mock_llm, - text_msg, - tool_call_msg, -) - -# ── Helpers ────────────────────────────────────────────────────── - - -class Confirmation(pydantic.BaseModel): - approved: bool - reason: str = "" - - -# ── wrap_tool ─────────────────────────────────────────────────── - - -async def test_wrap_tool_is_called() -> None: - """Middleware.wrap_tool is invoked for every tool execution.""" - tool_calls: list[middleware.ToolContext] = [] - - class Spy(middleware._Middleware): - async def wrap_tool( - self, call: middleware.ToolContext, next: Any - ) -> Any: - tool_calls.append(call) - return await next(call) - - @ai.tool - async def double(x: int) -> int: - """Double a number.""" - return x * 2 - - my_agent = ai.Agent(tools=[double]) - call1 = [tool_call_msg(tc_id="tc-1", name="double", args='{"x": 7}')] - call2 = [text_msg("14")] - mock_llm([call1, call2]) - - async with my_agent.run( - MOCK_MODEL, [ai.user_message("Double 7")], _middleware=[Spy()] - ) as stream: - async for _m in stream: - pass - - assert len(tool_calls) == 1 - assert tool_calls[0].tool_name == "double" - assert tool_calls[0].tool_call_id == "tc-1" - assert tool_calls[0].kwargs == {"x": 7} - - -# ── wrap_hook ─────────────────────────────────────────────────── - - -async def test_wrap_hook_is_called() -> None: - """Middleware.wrap_hook is invoked for every ai.hook() call.""" - hook_calls: list[middleware.HookContext] = [] - - class Spy(middleware._Middleware): - async def wrap_hook( - self, call: middleware.HookContext, next: Any - ) -> Any: - hook_calls.append(call) - return await next(call) - - class MyAgent(ai.Agent): - async def loop( - self, context: ai.Context - ) -> AsyncGenerator[ai.events.Event]: - async with ai.models.stream(context=context) as stream: - async for event in stream: - yield event - await ai.hook("test_hook", payload=Confirmation) - - my_agent = MyAgent() - - mock_llm([[text_msg("OK")]]) - - async with my_agent.run( - MOCK_MODEL, [ai.user_message("go")], _middleware=[Spy()] - ) as stream: - async for event in stream: - if not isinstance(event, agent_events_.HookEvent): - continue - if event.hook.status == "pending": - ai.resolve_hook("test_hook", {"approved": True, "reason": "ok"}) - - assert len(hook_calls) == 1 - assert hook_calls[0].label == "test_hook" - assert hook_calls[0].payload is Confirmation - - -# ── wrap_agent_run ────────────────────────────────────────────── - - -async def test_wrap_agent_run_ordering() -> None: - """Agent run middleware chain runs in the correct onion order.""" - order: list[str] = [] - - class Outer(middleware._Middleware): - async def wrap_agent_run( - self, call: ai.Context, next: Any - ) -> AsyncGenerator[ai.events.Event]: - order.append("outer-before") - async for event in next(call): - yield event - order.append("outer-after") - - class Inner(middleware._Middleware): - async def wrap_agent_run( - self, call: ai.Context, next: Any - ) -> AsyncGenerator[ai.events.Event]: - order.append("inner-before") - async for event in next(call): - yield event - order.append("inner-after") - - my_agent = ai.Agent() - mock_llm([[text_msg("Hi")]]) - - async with my_agent.run( - MOCK_MODEL, [ai.user_message("Hi")], _middleware=[Outer(), Inner()] - ) as stream: - async for _m in stream: - pass - - assert order == [ - "outer-before", - "inner-before", - "inner-after", - "outer-after", - ] - - -async def test_wrap_tool_context_fields_flow_to_result() -> None: - """ToolContext.tool_name is used in the result message.""" - - class Rewriter(middleware._Middleware): - async def wrap_tool( - self, call: middleware.ToolContext, next: Any - ) -> Any: - # Rewrite the tool_name via dataclasses.replace. - modified = dataclasses.replace(call, tool_name="rewritten-name") - return await next(modified) - - @ai.tool - async def echo(x: int) -> int: - """Echo a number.""" - return x - - my_agent = ai.Agent(tools=[echo]) - call1 = [tool_call_msg(tc_id="tc-1", name="echo", args='{"x": 42}')] - call2 = [text_msg("done")] - mock_llm([call1, call2]) - - async with my_agent.run( - MOCK_MODEL, [ai.user_message("go")], _middleware=[Rewriter()] - ) as stream: - msgs = await collect_messages(stream) - tool_result_msgs = [m for m in msgs if m.role == "tool" and m.tool_results] - - assert len(tool_result_msgs) >= 1 - # The result message should use the rewritten name, not the original. - assert tool_result_msgs[0].tool_results[0].tool_name == "rewritten-name" - - -async def test_wrap_tool_rewriting_tool_call_id_breaks_history() -> None: - """tool_call_id is a correlation key and must stay stable.""" - - class Rewriter(middleware._Middleware): - async def wrap_tool( - self, call: middleware.ToolContext, next: Any - ) -> Any: - modified = dataclasses.replace(call, tool_call_id="rewritten-id") - return await next(modified) - - @ai.tool - async def echo(x: int) -> int: - """Echo a number.""" - return x - - my_agent = ai.Agent(tools=[echo]) - call1 = [tool_call_msg(tc_id="original-id", name="echo", args='{"x": 42}')] - call2 = [text_msg("done")] - mock_llm([call1, call2]) - - with pytest.raises(ExceptionGroup) as exc_info: - async with my_agent.run( - MOCK_MODEL, [ai.user_message("go")], _middleware=[Rewriter()] - ) as stream: - async for _m in stream: - pass - - assert exc_info.group_contains( - history_utils.IntegrityError, match="orphaned-tool-result" - ) - - -# ── Context snapshot isolation ────────────────────────────────── - - -async def test_model_context_messages_are_isolated() -> None: - """Mutating call.messages in middleware does not affect the caller.""" - original_messages = [ai.user_message("Hello")] - - class Mutator(middleware._Middleware): - async def wrap_model( - self, call: middleware.ModelContext, next: Any - ) -> Any: - # Try to mutate the context's messages list in place. - call.messages.append(ai.system_message("injected")) - return await next(call) - - my_agent = ai.Agent() - mock_llm([[text_msg("Hi")]]) - - async with my_agent.run( - MOCK_MODEL, original_messages, _middleware=[Mutator()] - ) as stream: - async for _m in stream: - pass - - # The original list should be unmodified. - assert len(original_messages) == 1 - assert original_messages[0].role == "user" - - -# ── Middleware can repair bad tool args ────────────────────────── - - -async def test_middleware_can_fix_bad_tool_kwargs() -> None: - """A middleware that rewrites call.kwargs can fix malformed tool args.""" - - class ArgFixer(middleware._Middleware): - async def wrap_tool( - self, call: middleware.ToolContext, next: Any - ) -> Any: - # If kwargs are empty (parse failed), supply valid ones. - if not call.kwargs: - fixed = dataclasses.replace(call, kwargs={"x": 99}) - return await next(fixed) - return await next(call) - - @ai.tool - async def double(x: int) -> int: - """Double a number.""" - return x * 2 - - my_agent = ai.Agent(tools=[double]) - # Send completely invalid JSON args — parse will fail, kwargs will be {}. - call1 = [tool_call_msg(tc_id="tc-bad", name="double", args="not json")] - call2 = [text_msg("done")] - mock_llm([call1, call2]) - - async with my_agent.run( - MOCK_MODEL, [ai.user_message("go")], _middleware=[ArgFixer()] - ) as stream: - msgs = await collect_messages(stream) - tool_result_msgs = [m for m in msgs if m.role == "tool" and m.tool_results] - - assert len(tool_result_msgs) >= 1 - # The fixer middleware supplied x=99, so double should return 198. - tr = tool_result_msgs[0].tool_results[0].result - assert tr == 198 - assert tool_result_msgs[0].tool_results[0].is_error is False