1"""Types for middleware and agents."""23from __future__ import annotations45import warnings6from collections.abc import Awaitable, Callable, Sequence7from dataclasses import dataclass, field, replace8from inspect import iscoroutinefunction9from typing import (10    TYPE_CHECKING,11    Annotated,12    Any,13    Generic,14    Literal,15    cast,16    overload,17)1819# Needed as top level import for Pydantic schema generation on AgentState20from langchain_core.messages import (21    AIMessage,22    AnyMessage,23    BaseMessage,24    SystemMessage,25    ToolMessage,26)27from langgraph.channels.ephemeral_value import EphemeralValue28from langgraph.graph.message import add_messages29from langgraph.prebuilt.tool_node import ToolCallRequest, ToolCallWrapper30from langgraph.runtime import Runtime31from langgraph.types import Command32from langgraph.typing import ContextT33from typing_extensions import NotRequired, Required, TypedDict, TypeVar, Unpack3435if TYPE_CHECKING:36    from langchain_core.language_models.chat_models import BaseChatModel37    from langchain_core.tools import BaseTool38    from langgraph.stream._mux import TransformerFactory3940    from langchain.agents.structured_output import ResponseFormat4142__all__ = [43    "AgentMiddleware",44    "AgentState",45    "ContextT",46    "ExtendedModelResponse",47    "InputAgentState",48    "ModelCallResult",49    "ModelRequest",50    "ModelResponse",51    "OmitFromSchema",52    "OutputAgentState",53    "ResponseT",54    "StateT_co",55    "ToolCallRequest",56    "ToolCallWrapper",57    "after_agent",58    "after_model",59    "before_agent",60    "before_model",61    "dynamic_prompt",62    "hook_config",63    "wrap_tool_call",64]6566JumpTo = Literal["tools", "model", "end"]67"""Destination to jump to when a middleware node returns."""6869ResponseT = TypeVar("ResponseT", default=Any)707172class _ModelRequestOverrides(TypedDict, total=False):73    """Possible overrides for `ModelRequest.override()` method."""7475    model: BaseChatModel76    system_message: SystemMessage | None77    messages: list[AnyMessage]78    tool_choice: Any | None79    tools: list[BaseTool | dict[str, Any]]80    response_format: ResponseFormat[Any] | None81    model_settings: dict[str, Any]82    state: AgentState[Any]838485@dataclass(init=False)86class ModelRequest(Generic[ContextT]):87    """Model request information for the agent.8889    Type Parameters:90        ContextT: The type of the runtime context.9192            Defaults to `None` if not specified.93    """9495    model: BaseChatModel96    messages: list[AnyMessage]  # excluding system message97    system_message: SystemMessage | None98    tool_choice: Any | None99    tools: list[BaseTool | dict[str, Any]]100    response_format: ResponseFormat[Any] | None101    state: AgentState[Any]102    runtime: Runtime[ContextT]103    model_settings: dict[str, Any] = field(default_factory=dict)104105    def __init__(106        self,107        *,108        model: BaseChatModel,109        messages: list[AnyMessage],110        system_message: SystemMessage | None = None,111        system_prompt: str | None = None,112        tool_choice: Any | None = None,113        tools: list[BaseTool | dict[str, Any]] | None = None,114        response_format: ResponseFormat[Any] | None = None,115        state: AgentState[Any] | None = None,116        runtime: Runtime[ContextT] | None = None,117        model_settings: dict[str, Any] | None = None,118    ) -> None:119        """Initialize `ModelRequest` with backward compatibility for `system_prompt`.120121        Args:122            model: The chat model to use.123            messages: List of messages (excluding system prompt).124            tool_choice: Tool choice configuration.125            tools: List of available tools.126            response_format: Response format specification.127            state: Agent state.128            runtime: Runtime context.129            model_settings: Additional model settings.130            system_message: System message instance (preferred).131            system_prompt: System prompt string (deprecated, converted to `SystemMessage`).132133        Raises:134            ValueError: If both `system_prompt` and `system_message` are provided.135        """136        # Handle system_prompt/system_message conversion and validation137        if system_prompt is not None and system_message is not None:138            msg = "Cannot specify both system_prompt and system_message"139            raise ValueError(msg)140141        if system_prompt is not None:142            system_message = SystemMessage(content=system_prompt)143144        with warnings.catch_warnings():145            warnings.simplefilter("ignore", category=DeprecationWarning)146            self.model = model147            self.messages = messages148            self.system_message = system_message149            self.tool_choice = tool_choice150            self.tools = tools if tools is not None else []151            self.response_format = response_format152            self.state = state if state is not None else {"messages": []}153            self.runtime = runtime  # type: ignore[assignment]154            self.model_settings = model_settings if model_settings is not None else {}155156    @property157    def system_prompt(self) -> str | None:158        """Get system prompt text from system_message.159160        Returns:161            The content of the system message if present, otherwise `None`.162        """163        if self.system_message is None:164            return None165        return self.system_message.text166167    def __setattr__(self, name: str, value: Any) -> None:168        """Set an attribute with a deprecation warning.169170        Direct attribute assignment on `ModelRequest` is deprecated. Use the171        `override()` method instead to create a new request with modified attributes.172173        Args:174            name: Attribute name.175            value: Attribute value.176        """177        # Special handling for system_prompt - convert to system_message178        if name == "system_prompt":179            warnings.warn(180                "Direct attribute assignment to ModelRequest.system_prompt is deprecated. "181                "Use request.override(system_message=SystemMessage(...)) instead to create "182                "a new request with the modified system message.",183                DeprecationWarning,184                stacklevel=2,185            )186            if value is None:187                object.__setattr__(self, "system_message", None)188            else:189                object.__setattr__(self, "system_message", SystemMessage(content=value))190            return191192        warnings.warn(193            f"Direct attribute assignment to ModelRequest.{name} is deprecated. "194            f"Use request.override({name}=...) instead to create a new request "195            f"with the modified attribute.",196            DeprecationWarning,197            stacklevel=2,198        )199        object.__setattr__(self, name, value)200201    def override(self, **overrides: Unpack[_ModelRequestOverrides]) -> ModelRequest[ContextT]:202        """Replace the request with a new request with the given overrides.203204        Returns a new `ModelRequest` instance with the specified attributes replaced.205206        This follows an immutable pattern, leaving the original request unchanged.207208        Args:209            **overrides: Keyword arguments for attributes to override.210211                Supported keys:212213                - `model`: `BaseChatModel` instance214                - `system_prompt`: deprecated, use `system_message` instead215                - `system_message`: `SystemMessage` instance216                - `messages`: `list` of messages217                - `tool_choice`: Tool choice configuration218                - `tools`: `list` of available tools219                - `response_format`: Response format specification220                - `model_settings`: Additional model settings221                - `state`: Agent state dictionary222223        Returns:224            New `ModelRequest` instance with specified overrides applied.225226        Examples:227            !!! example "Create a new request with different model"228229                ```python230                new_request = request.override(model=different_model)231                ```232233            !!! example "Override system message (preferred)"234235                ```python236                from langchain_core.messages import SystemMessage237238                new_request = request.override(239                    system_message=SystemMessage(content="New instructions")240                )241                ```242243            !!! example "Override multiple attributes"244245                ```python246                new_request = request.override(247                    model=ChatOpenAI(model="gpt-5.5"),248                    system_message=SystemMessage(content="New instructions"),249                )250                ```251252        Raises:253            ValueError: If both `system_prompt` and `system_message` are provided.254        """255        # Handle system_prompt/system_message conversion256        if "system_prompt" in overrides and "system_message" in overrides:257            msg = "Cannot specify both system_prompt and system_message"258            raise ValueError(msg)259260        if "system_prompt" in overrides:261            system_prompt = cast("str | None", overrides.pop("system_prompt"))  # type: ignore[typeddict-item]262            if system_prompt is None:263                overrides["system_message"] = None264            else:265                overrides["system_message"] = SystemMessage(content=system_prompt)266267        return replace(self, **overrides)268269270@dataclass271class ModelResponse(Generic[ResponseT]):272    """Response from model execution including messages and optional structured output.273274    The result will usually contain a single `AIMessage`, but may include an additional275    `ToolMessage` if the model used a tool for structured output.276277    Type Parameters:278        ResponseT: The type of the structured response. Defaults to `Any` if not specified.279    """280281    result: list[BaseMessage]282    """List of messages from model execution."""283284    structured_response: ResponseT | None = None285    """Parsed structured output if `response_format` was specified, `None` otherwise."""286287288@dataclass289class ExtendedModelResponse(Generic[ResponseT]):290    """Model response with an optional 'Command' from 'wrap_model_call' middleware.291292    Use this to return a 'Command' alongside the model response from a293    'wrap_model_call' handler. The command is applied as an additional state294    update after the model node completes, using the graph's reducers (e.g.295    'add_messages' for the 'messages' key).296297    Because each 'Command' is applied through the reducer, messages in the298    command are **added alongside** the model response messages rather than299    replacing them. For non-reducer state fields, later commands overwrite300    earlier ones (outermost middleware wins over inner).301302    Type Parameters:303        ResponseT: The type of the structured response. Defaults to 'Any' if not specified.304    """305306    model_response: ModelResponse[ResponseT]307    """The underlying model response."""308309    command: Command[Any] | None = None310    """Optional command to apply as an additional state update."""311312313ModelCallResult = ModelResponse[ResponseT] | AIMessage | ExtendedModelResponse[ResponseT]314"""Return type for model call handlers.315316Middleware can return either:317318- `ModelResponse`: Full response with messages and optional structured output319- `AIMessage`: Simplified return for simple use cases320- `ExtendedModelResponse`: Response with an optional `Command` for additional state updates321    `goto`, `resume`, and `graph` are not yet supported on these commands.322    A `NotImplementedError` will be raised if you try to use them.323"""324325326@dataclass327class OmitFromSchema:328    """Annotation used to mark state attributes as omitted from input or output schemas."""329330    input: bool = True331    """Whether to omit the attribute from the input schema."""332333    output: bool = True334    """Whether to omit the attribute from the output schema."""335336337OmitFromInput = OmitFromSchema(input=True, output=False)338"""Annotation used to mark state attributes as omitted from input schema."""339340OmitFromOutput = OmitFromSchema(input=False, output=True)341"""Annotation used to mark state attributes as omitted from output schema."""342343PrivateStateAttr = OmitFromSchema(input=True, output=True)344"""Annotation used to mark state attributes as purely internal for a given middleware."""345346347class AgentState(TypedDict, Generic[ResponseT]):348    """State schema for the agent."""349350    messages: Required[Annotated[list[AnyMessage], add_messages]]351    jump_to: NotRequired[Annotated[JumpTo | None, EphemeralValue, PrivateStateAttr]]352    structured_response: NotRequired[Annotated[ResponseT, OmitFromInput]]353354355class InputAgentState(TypedDict):356    """Input state schema for the agent."""357358    messages: Required[Annotated[list[AnyMessage | dict[str, Any]], add_messages]]359360361class OutputAgentState(TypedDict, Generic[ResponseT]):362    """Output state schema for the agent."""363364    messages: Required[Annotated[list[AnyMessage], add_messages]]365    structured_response: NotRequired[ResponseT]366367368# Deprecated aliases kept for backwards compatibility with external consumers that369# imported the previously private names. Remove in a future release.370_InputAgentState = InputAgentState371_OutputAgentState = OutputAgentState372373374StateT = TypeVar("StateT", bound=AgentState[Any], default=AgentState[Any])375StateT_co = TypeVar("StateT_co", bound=AgentState[Any], default=AgentState[Any], covariant=True)376StateT_contra = TypeVar("StateT_contra", bound=AgentState[Any], contravariant=True)377378379class _DefaultAgentState(AgentState[Any]):380    """AgentMiddleware default state."""381382383class AgentMiddleware(Generic[StateT, ContextT, ResponseT]):384    """Base middleware class for an agent.385386    Subclass this and implement any of the defined methods to customize agent behavior387    between steps in the main agent loop.388389    Type Parameters:390        StateT: The type of the agent state. Defaults to `AgentState[Any]`.391        ContextT: The type of the runtime context. Defaults to `None`.392        ResponseT: The type of the structured response. Defaults to `Any`.393    """394395    state_schema: type[StateT] = cast("type[StateT]", _DefaultAgentState)396    """The schema for state passed to the middleware nodes."""397398    tools: Sequence[BaseTool]399    """Additional tools registered by the middleware."""400401    transformers: Sequence[TransformerFactory] = ()402    """Stream transformer factories registered by the middleware.403404    Each entry is a scope-aware factory invoked as `factory(scope)` so every405    invocation receives a fresh instance. Factories are merged with the406    `transformers` argument of [`create_agent`][langchain.agents.create_agent]407    at graph compile time, after the `ToolCallTransformer` and before any408    user-supplied entries.409    """410411    @property412    def name(self) -> str:413        """The name of the middleware instance.414415        Defaults to the class name, but can be overridden for custom naming.416        """417        return self.__class__.__name__418419    def before_agent(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:420        """Logic to run before the agent execution starts.421422        Args:423            state: The current agent state.424            runtime: The runtime context.425426        Returns:427            Agent state updates to apply before agent execution.428        """429430    async def abefore_agent(431        self, state: StateT, runtime: Runtime[ContextT]432    ) -> dict[str, Any] | None:433        """Async logic to run before the agent execution starts.434435        Args:436            state: The current agent state.437            runtime: The runtime context.438439        Returns:440            Agent state updates to apply before agent execution.441        """442443    def before_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:444        """Logic to run before the model is called.445446        Args:447            state: The current agent state.448            runtime: The runtime context.449450        Returns:451            Agent state updates to apply before model call.452        """453454    async def abefore_model(455        self, state: StateT, runtime: Runtime[ContextT]456    ) -> dict[str, Any] | None:457        """Async logic to run before the model is called.458459        Args:460            state: The agent state.461            runtime: The runtime context.462463        Returns:464            Agent state updates to apply before model call.465        """466467    def after_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:468        """Logic to run after the model is called.469470        Args:471            state: The current agent state.472            runtime: The runtime context.473474        Returns:475            Agent state updates to apply after model call.476        """477478    async def aafter_model(479        self, state: StateT, runtime: Runtime[ContextT]480    ) -> dict[str, Any] | None:481        """Async logic to run after the model is called.482483        Args:484            state: The current agent state.485            runtime: The runtime context.486487        Returns:488            Agent state updates to apply after model call.489        """490491    def wrap_model_call(492        self,493        request: ModelRequest[ContextT],494        handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],495    ) -> ModelResponse[ResponseT] | AIMessage | ExtendedModelResponse[ResponseT]:496        """Intercept and control model execution via handler callback.497498        Async version is `awrap_model_call`499500        The handler callback executes the model request and returns a `ModelResponse`.501        Middleware can call the handler multiple times for retry logic, skip calling502        it to short-circuit, or modify the request/response. Multiple middleware503        compose with first in list as outermost layer.504505        Args:506            request: Model request to execute (includes state and runtime).507            handler: Callback that executes the model request and returns508                `ModelResponse`.509510                Call this to execute the model.511512                Can be called multiple times for retry logic.513514                Can skip calling it to short-circuit.515516        Returns:517            The model call result.518519        Examples:520            !!! example "Retry on error"521522                ```python523                def wrap_model_call(self, request, handler):524                    for attempt in range(3):525                        try:526                            return handler(request)527                        except Exception:528                            if attempt == 2:529                                raise530                ```531532            !!! example "Rewrite response"533534                ```python535                def wrap_model_call(self, request, handler):536                    response = handler(request)537                    ai_msg = response.result[0]538                    return ModelResponse(539                        result=[AIMessage(content=f"[{ai_msg.content}]")],540                        structured_response=response.structured_response,541                    )542                ```543544            !!! example "Error to fallback"545546                ```python547                def wrap_model_call(self, request, handler):548                    try:549                        return handler(request)550                    except Exception:551                        return ModelResponse(result=[AIMessage(content="Service unavailable")])552                ```553554            !!! example "Cache/short-circuit"555556                ```python557                def wrap_model_call(self, request, handler):558                    if cached := get_cache(request):559                        return cached  # Short-circuit with cached result560                    response = handler(request)561                    save_cache(request, response)562                    return response563                ```564565            !!! example "Simple `AIMessage` return (converted automatically)"566567                ```python568                def wrap_model_call(self, request, handler):569                    response = handler(request)570                    # Can return AIMessage directly for simple cases571                    return AIMessage(content="Simplified response")572                ```573        """574        msg = (575            "Synchronous implementation of wrap_model_call is not available. "576            "You are likely encountering this error because you defined only the async version "577            "(awrap_model_call) and invoked your agent in a synchronous context "578            "(e.g., using `stream()` or `invoke()`). "579            "To resolve this, either: "580            "(1) subclass AgentMiddleware and implement the synchronous wrap_model_call method, "581            "(2) use the @wrap_model_call decorator on a standalone sync function, or "582            "(3) invoke your agent asynchronously using `astream()` or `ainvoke()`."583        )584        raise NotImplementedError(msg)585586    async def awrap_model_call(587        self,588        request: ModelRequest[ContextT],589        handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],590    ) -> ModelResponse[ResponseT] | AIMessage | ExtendedModelResponse[ResponseT]:591        """Intercept and control async model execution via handler callback.592593        The handler callback executes the model request and returns a `ModelResponse`.594595        Middleware can call the handler multiple times for retry logic, skip calling596        it to short-circuit, or modify the request/response. Multiple middleware597        compose with first in list as outermost layer.598599        Args:600            request: Model request to execute (includes state and runtime).601            handler: Async callback that executes the model request and returns602                `ModelResponse`.603604                Call this to execute the model.605606                Can be called multiple times for retry logic.607608                Can skip calling it to short-circuit.609610        Returns:611            The model call result.612613        Examples:614            !!! example "Retry on error"615616                ```python617                async def awrap_model_call(self, request, handler):618                    for attempt in range(3):619                        try:620                            return await handler(request)621                        except Exception:622                            if attempt == 2:623                                raise624                ```625        """626        msg = (627            "Asynchronous implementation of awrap_model_call is not available. "628            "You are likely encountering this error because you defined only the sync version "629            "(wrap_model_call) and invoked your agent in an asynchronous context "630            "(e.g., using `astream()` or `ainvoke()`). "631            "To resolve this, either: "632            "(1) subclass AgentMiddleware and implement the asynchronous awrap_model_call method, "633            "(2) use the @wrap_model_call decorator on a standalone async function, or "634            "(3) invoke your agent synchronously using `stream()` or `invoke()`."635        )636        raise NotImplementedError(msg)637638    def after_agent(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:639        """Logic to run after the agent execution completes.640641        Args:642            state: The current agent state.643            runtime: The runtime context.644645        Returns:646            Agent state updates to apply after agent execution.647        """648649    async def aafter_agent(650        self, state: StateT, runtime: Runtime[ContextT]651    ) -> dict[str, Any] | None:652        """Async logic to run after the agent execution completes.653654        Args:655            state: The current agent state.656            runtime: The runtime context.657658        Returns:659            Agent state updates to apply after agent execution.660        """661662    def wrap_tool_call(663        self,664        request: ToolCallRequest,665        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],666    ) -> ToolMessage | Command[Any]:667        """Intercept tool execution for retries, monitoring, or modification.668669        Async version is `awrap_tool_call`670671        Multiple middleware compose automatically (first defined = outermost).672673        Exceptions propagate unless `handle_tool_errors` is configured on `ToolNode`.674675        Args:676            request: Tool call request with call `dict`, `BaseTool`, state, and runtime.677678                Access state via `request.state` and runtime via `request.runtime`.679            handler: `Callable` to execute the tool (can be called multiple times).680681        Returns:682            `ToolMessage` or `Command` (the final result).683684        The handler `Callable` can be invoked multiple times for retry logic.685686        Each call to handler is independent and stateless.687688        Examples:689            !!! example "Modify request before execution"690691                ```python692                def wrap_tool_call(self, request, handler):693                    modified_call = {694                        **request.tool_call,695                        "args": {696                            **request.tool_call["args"],697                            "value": request.tool_call["args"]["value"] * 2,698                        },699                    }700                    request = request.override(tool_call=modified_call)701                    return handler(request)702                ```703704            !!! example "Retry on error (call handler multiple times)"705706                ```python707                def wrap_tool_call(self, request, handler):708                    for attempt in range(3):709                        try:710                            result = handler(request)711                            if is_valid(result):712                                return result713                        except Exception:714                            if attempt == 2:715                                raise716                    return result717                ```718719            !!! example "Conditional retry based on response"720721                ```python722                def wrap_tool_call(self, request, handler):723                    for attempt in range(3):724                        result = handler(request)725                        if isinstance(result, ToolMessage) and result.status != "error":726                            return result727                        if attempt < 2:728                            continue729                        return result730                ```731        """732        msg = (733            "Synchronous implementation of wrap_tool_call is not available. "734            "You are likely encountering this error because you defined only the async version "735            "(awrap_tool_call) and invoked your agent in a synchronous context "736            "(e.g., using `stream()` or `invoke()`). "737            "To resolve this, either: "738            "(1) subclass AgentMiddleware and implement the synchronous wrap_tool_call method, "739            "(2) use the @wrap_tool_call decorator on a standalone sync function, or "740            "(3) invoke your agent asynchronously using `astream()` or `ainvoke()`."741        )742        raise NotImplementedError(msg)743744    async def awrap_tool_call(745        self,746        request: ToolCallRequest,747        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],748    ) -> ToolMessage | Command[Any]:749        """Intercept and control async tool execution via handler callback.750751        The handler callback executes the tool call and returns a `ToolMessage` or752        `Command`. Middleware can call the handler multiple times for retry logic, skip753        calling it to short-circuit, or modify the request/response. Multiple middleware754        compose with first in list as outermost layer.755756        Args:757            request: Tool call request with call `dict`, `BaseTool`, state, and runtime.758759                Access state via `request.state` and runtime via `request.runtime`.760            handler: Async callable to execute the tool and returns `ToolMessage` or761                `Command`.762763                Call this to execute the tool.764765                Can be called multiple times for retry logic.766767                Can skip calling it to short-circuit.768769        Returns:770            `ToolMessage` or `Command` (the final result).771772        The handler `Callable` can be invoked multiple times for retry logic.773774        Each call to handler is independent and stateless.775776        Examples:777            !!! example "Async retry on error"778779                ```python780                async def awrap_tool_call(self, request, handler):781                    for attempt in range(3):782                        try:783                            result = await handler(request)784                            if is_valid(result):785                                return result786                        except Exception:787                            if attempt == 2:788                                raise789                    return result790                ```791792                ```python793                async def awrap_tool_call(self, request, handler):794                    if cached := await get_cache_async(request):795                        return ToolMessage(content=cached, tool_call_id=request.tool_call["id"])796                    result = await handler(request)797                    await save_cache_async(request, result)798                    return result799                ```800        """801        msg = (802            "Asynchronous implementation of awrap_tool_call is not available. "803            "You are likely encountering this error because you defined only the sync version "804            "(wrap_tool_call) and invoked your agent in an asynchronous context "805            "(e.g., using `astream()` or `ainvoke()`). "806            "To resolve this, either: "807            "(1) subclass AgentMiddleware and implement the asynchronous awrap_tool_call method, "808            "(2) use the @wrap_tool_call decorator on a standalone async function, or "809            "(3) invoke your agent synchronously using `stream()` or `invoke()`."810        )811        raise NotImplementedError(msg)812813814_SyncCallableWithStateAndRuntime = Callable[815    [StateT_contra, Runtime[ContextT]], dict[str, Any] | Command[Any] | None816]817_AsyncCallableWithStateAndRuntime = Callable[818    [StateT_contra, Runtime[ContextT]], Awaitable[dict[str, Any] | Command[Any] | None]819]820_CallableWithStateAndRuntime = (821    _SyncCallableWithStateAndRuntime[StateT_contra, ContextT]822    | _AsyncCallableWithStateAndRuntime[StateT_contra, ContextT]823)824825_SyncCallableReturningSystemMessage = Callable[[ModelRequest[ContextT]], str | SystemMessage]826_AsyncCallableReturningSystemMessage = Callable[827    [ModelRequest[ContextT]], Awaitable[str | SystemMessage]828]829_CallableReturningSystemMessage = (830    _SyncCallableReturningSystemMessage[ContextT] | _AsyncCallableReturningSystemMessage[ContextT]831)832833# Sync/async signatures for `wrap_model_call` interception; see `@wrap_model_call`.834_SyncCallableReturningModelResponse = Callable[835    [ModelRequest[ContextT], Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]]],836    ModelCallResult,837]838_AsyncCallableReturningModelResponse = Callable[839    [840        ModelRequest[ContextT],841        Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],842    ],843    Awaitable[ModelCallResult],844]845_CallableReturningModelResponse = (846    _SyncCallableReturningModelResponse[ContextT, ResponseT]847    | _AsyncCallableReturningModelResponse[ContextT, ResponseT]848)849850# Sync/async signatures for `wrap_tool_call` interception; see `@wrap_tool_call`.851_SyncCallableReturningToolResponse = Callable[852    [ToolCallRequest, Callable[[ToolCallRequest], ToolMessage | Command[Any]]],853    ToolMessage | Command[Any],854]855_AsyncCallableReturningToolResponse = Callable[856    [ToolCallRequest, Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]]],857    Awaitable[ToolMessage | Command[Any]],858]859_CallableReturningToolResponse = (860    _SyncCallableReturningToolResponse | _AsyncCallableReturningToolResponse861)862863864CallableT = TypeVar("CallableT", bound=Callable[..., Any])865866867def hook_config(868    *,869    can_jump_to: list[JumpTo] | None = None,870) -> Callable[[CallableT], CallableT]:871    """Decorator to configure hook behavior in middleware methods.872873    Use this decorator on `before_model` or `after_model` methods in middleware classes874    to configure their behavior. Currently supports specifying which destinations they875    can jump to, which establishes conditional edges in the agent graph.876877    Args:878        can_jump_to: Optional list of valid jump destinations.879880            Can be:881882            - `'tools'`: Jump to the tools node883            - `'model'`: Jump back to the model node884            - `'end'`: Jump to the end of the graph885886    Returns:887        Decorator function that marks the method with configuration metadata.888889    Examples:890        !!! example "Using decorator on a class method"891892            ```python893            class MyMiddleware(AgentMiddleware):894                @hook_config(can_jump_to=["end", "model"])895                def before_model(self, state: AgentState) -> dict[str, Any] | None:896                    if some_condition(state):897                        return {"jump_to": "end"}898                    return None899            ```900901        Alternative: Use the `can_jump_to` parameter in `before_model`/`after_model`902        decorators:903904        ```python905        @before_model(can_jump_to=["end"])906        def conditional_middleware(state: AgentState) -> dict[str, Any] | None:907            if should_exit(state):908                return {"jump_to": "end"}909            return None910        ```911    """912913    def decorator(func: CallableT) -> CallableT:914        if can_jump_to is not None:915            func.__can_jump_to__ = can_jump_to  # type: ignore[attr-defined]916        return func917918    return decorator919920921@overload922def before_model(923    func: _CallableWithStateAndRuntime[StateT, ContextT],924) -> AgentMiddleware[StateT, ContextT]: ...925926927@overload928def before_model(929    func: None = None,930    *,931    state_schema: type[StateT] | None = None,932    tools: list[BaseTool] | None = None,933    can_jump_to: list[JumpTo] | None = None,934    name: str | None = None,935) -> Callable[936    [_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]937]: ...938939940def before_model(941    func: _CallableWithStateAndRuntime[StateT, ContextT] | None = None,942    *,943    state_schema: type[StateT] | None = None,944    tools: list[BaseTool] | None = None,945    can_jump_to: list[JumpTo] | None = None,946    name: str | None = None,947) -> (948    Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]949    | AgentMiddleware[StateT, ContextT]950):951    """Decorator used to dynamically create a middleware with the `before_model` hook.952953    Args:954        func: The function to be decorated.955956            Must accept: `state: StateT, runtime: Runtime[ContextT]` - State and runtime957                context958        state_schema: Optional custom state schema type.959960            If not provided, uses the default `AgentState` schema.961        tools: Optional list of additional tools to register with this middleware.962        can_jump_to: Optional list of valid jump destinations for conditional edges.963964            Valid values are: `'tools'`, `'model'`, `'end'`965        name: Optional name for the generated middleware class.966967            If not provided, uses the decorated function's name.968969    Returns:970        Either an `AgentMiddleware` instance (if func is provided directly) or a971            decorator function that can be applied to a function it is wrapping.972973    The decorated function should return:974975    - `dict[str, Any]` - State updates to merge into the agent state976    - `Command` - A command to control flow (e.g., jump to different node)977    - `None` - No state updates or flow control978979    Examples:980        !!! example "Basic usage"981982            ```python983            @before_model984            def log_before_model(state: AgentState, runtime: Runtime) -> None:985                print(f"About to call model with {len(state['messages'])} messages")986            ```987988        !!! example "With conditional jumping"989990            ```python991            @before_model(can_jump_to=["end"])992            def conditional_before_model(993                state: AgentState, runtime: Runtime994            ) -> dict[str, Any] | None:995                if some_condition(state):996                    return {"jump_to": "end"}997                return None998            ```9991000        !!! example "With custom state schema"10011002            ```python1003            @before_model(state_schema=MyCustomState)1004            def custom_before_model(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:1005                return {"custom_field": "updated_value"}1006            ```10071008        !!! example "Streaming custom events before model call"10091010            Use `runtime.stream_writer` to emit custom events before each model invocation.1011            Events are received when streaming with `stream_mode="custom"`.10121013            ```python1014            @before_model1015            async def notify_model_call(state: AgentState, runtime: Runtime) -> None:1016                '''Notify user before model is called.'''1017                runtime.stream_writer(1018                    {1019                        "type": "status",1020                        "message": "Thinking...",1021                    }1022                )1023            ```1024    """10251026    def decorator(1027        func: _CallableWithStateAndRuntime[StateT, ContextT],1028    ) -> AgentMiddleware[StateT, ContextT]:1029        is_async = iscoroutinefunction(func)10301031        func_can_jump_to = (1032            can_jump_to if can_jump_to is not None else getattr(func, "__can_jump_to__", [])1033        )10341035        if is_async:10361037            async def async_wrapped(1038                _self: AgentMiddleware[StateT, ContextT],1039                state: StateT,1040                runtime: Runtime[ContextT],1041            ) -> dict[str, Any] | Command[Any] | None:1042                # `iscoroutinefunction` narrows `func` at runtime, but type checkers1043                # cannot narrow this sync-or-async callable union.1044                return await cast("_AsyncCallableWithStateAndRuntime[StateT, ContextT]", func)(1045                    state, runtime1046                )10471048            # Preserve can_jump_to metadata on the wrapped function1049            if func_can_jump_to:1050                async_wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]10511052            middleware_name = name or cast(1053                "str", getattr(func, "__name__", "BeforeModelMiddleware")1054            )10551056            # `type(...)` builds the correct middleware subclass at runtime, but1057            # type checkers cannot infer its generic `AgentMiddleware` parameters.1058            return cast(1059                "AgentMiddleware[StateT, ContextT]",1060                type(1061                    middleware_name,1062                    (AgentMiddleware,),1063                    {1064                        "state_schema": state_schema or AgentState,1065                        "tools": tools or [],1066                        "abefore_model": async_wrapped,1067                    },1068                )(),1069            )10701071        def wrapped(1072            _self: AgentMiddleware[StateT, ContextT],1073            state: StateT,1074            runtime: Runtime[ContextT],1075        ) -> dict[str, Any] | Command[Any] | None:1076            # `iscoroutinefunction` narrows `func` at runtime, but type checkers1077            # cannot narrow this sync-or-async callable union.1078            return cast("_SyncCallableWithStateAndRuntime[StateT, ContextT]", func)(state, runtime)10791080        # Preserve can_jump_to metadata on the wrapped function1081        if func_can_jump_to:1082            wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]10831084        # Use function name as default if no name provided1085        middleware_name = name or cast("str", getattr(func, "__name__", "BeforeModelMiddleware"))10861087        # `type(...)` builds the correct middleware subclass at runtime, but1088        # type checkers cannot infer its generic `AgentMiddleware` parameters.1089        return cast(1090            "AgentMiddleware[StateT, ContextT]",1091            type(1092                middleware_name,1093                (AgentMiddleware,),1094                {1095                    "state_schema": state_schema or AgentState,1096                    "tools": tools or [],1097                    "before_model": wrapped,1098                },1099            )(),1100        )11011102    if func is not None:1103        return decorator(func)1104    return decorator110511061107@overload1108def after_model(1109    func: _CallableWithStateAndRuntime[StateT, ContextT],1110) -> AgentMiddleware[StateT, ContextT]: ...111111121113@overload1114def after_model(1115    func: None = None,1116    *,1117    state_schema: type[StateT] | None = None,1118    tools: list[BaseTool] | None = None,1119    can_jump_to: list[JumpTo] | None = None,1120    name: str | None = None,1121) -> Callable[1122    [_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]1123]: ...112411251126def after_model(1127    func: _CallableWithStateAndRuntime[StateT, ContextT] | None = None,1128    *,1129    state_schema: type[StateT] | None = None,1130    tools: list[BaseTool] | None = None,1131    can_jump_to: list[JumpTo] | None = None,1132    name: str | None = None,1133) -> (1134    Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]1135    | AgentMiddleware[StateT, ContextT]1136):1137    """Decorator used to dynamically create a middleware with the `after_model` hook.11381139    Args:1140        func: The function to be decorated.11411142            Must accept: `state: StateT, runtime: Runtime[ContextT]` - State and runtime1143            context1144        state_schema: Optional custom state schema type.11451146            If not provided, uses the default `AgentState` schema.1147        tools: Optional list of additional tools to register with this middleware.1148        can_jump_to: Optional list of valid jump destinations for conditional edges.11491150            Valid values are: `'tools'`, `'model'`, `'end'`1151        name: Optional name for the generated middleware class.11521153            If not provided, uses the decorated function's name.11541155    Returns:1156        Either an `AgentMiddleware` instance (if func is provided) or a decorator1157            function that can be applied to a function.11581159    The decorated function should return:11601161    - `dict[str, Any]` - State updates to merge into the agent state1162    - `Command` - A command to control flow (e.g., jump to different node)1163    - `None` - No state updates or flow control11641165    Examples:1166        !!! example "Basic usage for logging model responses"11671168            ```python1169            @after_model1170            def log_latest_message(state: AgentState, runtime: Runtime) -> None:1171                print(state["messages"][-1].content)1172            ```11731174        !!! example "With custom state schema"11751176            ```python1177            @after_model(state_schema=MyCustomState, name="MyAfterModelMiddleware")1178            def custom_after_model(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:1179                return {"custom_field": "updated_after_model"}1180            ```11811182        !!! example "Streaming custom events after model call"11831184            Use `runtime.stream_writer` to emit custom events after model responds.1185            Events are received when streaming with `stream_mode="custom"`.11861187            ```python1188            @after_model1189            async def notify_model_response(state: AgentState, runtime: Runtime) -> None:1190                '''Notify user after model has responded.'''1191                last_message = state["messages"][-1]1192                has_tool_calls = hasattr(last_message, "tool_calls") and last_message.tool_calls1193                runtime.stream_writer(1194                    {1195                        "type": "status",1196                        "message": "Using tools..." if has_tool_calls else "Response ready!",1197                    }1198                )1199            ```1200    """12011202    def decorator(1203        func: _CallableWithStateAndRuntime[StateT, ContextT],1204    ) -> AgentMiddleware[StateT, ContextT]:1205        is_async = iscoroutinefunction(func)1206        # Extract can_jump_to from decorator parameter or from function metadata1207        func_can_jump_to = (1208            can_jump_to if can_jump_to is not None else getattr(func, "__can_jump_to__", [])1209        )12101211        if is_async:12121213            async def async_wrapped(1214                _self: AgentMiddleware[StateT, ContextT],1215                state: StateT,1216                runtime: Runtime[ContextT],1217            ) -> dict[str, Any] | Command[Any] | None:1218                # `iscoroutinefunction` narrows `func` at runtime, but type checkers1219                # cannot narrow this sync-or-async callable union.1220                return await cast("_AsyncCallableWithStateAndRuntime[StateT, ContextT]", func)(1221                    state, runtime1222                )12231224            # Preserve can_jump_to metadata on the wrapped function1225            if func_can_jump_to:1226                async_wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]12271228            middleware_name = name or cast("str", getattr(func, "__name__", "AfterModelMiddleware"))12291230            # `type(...)` builds the correct middleware subclass at runtime, but1231            # type checkers cannot infer its generic `AgentMiddleware` parameters.1232            return cast(1233                "AgentMiddleware[StateT, ContextT]",1234                type(1235                    middleware_name,1236                    (AgentMiddleware,),1237                    {1238                        "state_schema": state_schema or AgentState,1239                        "tools": tools or [],1240                        "aafter_model": async_wrapped,1241                    },1242                )(),1243            )12441245        def wrapped(1246            _self: AgentMiddleware[StateT, ContextT],1247            state: StateT,1248            runtime: Runtime[ContextT],1249        ) -> dict[str, Any] | Command[Any] | None:1250            # `iscoroutinefunction` narrows `func` at runtime, but type checkers1251            # cannot narrow this sync-or-async callable union.1252            return cast("_SyncCallableWithStateAndRuntime[StateT, ContextT]", func)(state, runtime)12531254        # Preserve can_jump_to metadata on the wrapped function1255        if func_can_jump_to:1256            wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]12571258        # Use function name as default if no name provided1259        middleware_name = name or cast("str", getattr(func, "__name__", "AfterModelMiddleware"))12601261        # `type(...)` builds the correct middleware subclass at runtime, but1262        # type checkers cannot infer its generic `AgentMiddleware` parameters.1263        return cast(1264            "AgentMiddleware[StateT, ContextT]",1265            type(1266                middleware_name,1267                (AgentMiddleware,),1268                {1269                    "state_schema": state_schema or AgentState,1270                    "tools": tools or [],1271                    "after_model": wrapped,1272                },1273            )(),1274        )12751276    if func is not None:1277        return decorator(func)1278    return decorator127912801281@overload1282def before_agent(1283    func: _CallableWithStateAndRuntime[StateT, ContextT],1284) -> AgentMiddleware[StateT, ContextT]: ...128512861287@overload1288def before_agent(1289    func: None = None,1290    *,1291    state_schema: type[StateT] | None = None,1292    tools: list[BaseTool] | None = None,1293    can_jump_to: list[JumpTo] | None = None,1294    name: str | None = None,1295) -> Callable[1296    [_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]1297]: ...129812991300def before_agent(1301    func: _CallableWithStateAndRuntime[StateT, ContextT] | None = None,1302    *,1303    state_schema: type[StateT] | None = None,1304    tools: list[BaseTool] | None = None,1305    can_jump_to: list[JumpTo] | None = None,1306    name: str | None = None,1307) -> (1308    Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]1309    | AgentMiddleware[StateT, ContextT]1310):1311    """Decorator used to dynamically create a middleware with the `before_agent` hook.13121313    Args:1314        func: The function to be decorated.13151316            Must accept: `state: StateT, runtime: Runtime[ContextT]` - State and runtime1317            context1318        state_schema: Optional custom state schema type.13191320            If not provided, uses the default `AgentState` schema.1321        tools: Optional list of additional tools to register with this middleware.1322        can_jump_to: Optional list of valid jump destinations for conditional edges.13231324            Valid values are: `'tools'`, `'model'`, `'end'`1325        name: Optional name for the generated middleware class.13261327            If not provided, uses the decorated function's name.13281329    Returns:1330        Either an `AgentMiddleware` instance (if func is provided directly) or a1331            decorator function that can be applied to a function it is wrapping.13321333    The decorated function should return:13341335    - `dict[str, Any]` - State updates to merge into the agent state1336    - `Command` - A command to control flow (e.g., jump to different node)1337    - `None` - No state updates or flow control13381339    Examples:1340        !!! example "Basic usage"13411342            ```python1343            @before_agent1344            def log_before_agent(state: AgentState, runtime: Runtime) -> None:1345                print(f"Starting agent with {len(state['messages'])} messages")1346            ```13471348        !!! example "With conditional jumping"13491350            ```python1351            @before_agent(can_jump_to=["end"])1352            def conditional_before_agent(1353                state: AgentState, runtime: Runtime1354            ) -> dict[str, Any] | None:1355                if some_condition(state):1356                    return {"jump_to": "end"}1357                return None1358            ```13591360        !!! example "With custom state schema"13611362            ```python1363            @before_agent(state_schema=MyCustomState)1364            def custom_before_agent(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:1365                return {"custom_field": "initialized_value"}1366            ```13671368        !!! example "Streaming custom events"13691370            Use `runtime.stream_writer` to emit custom events during agent execution.1371            Events are received when streaming with `stream_mode="custom"`.13721373            ```python1374            from langchain.agents import create_agent1375            from langchain.agents.middleware import before_agent, AgentState1376            from langchain.messages import HumanMessage1377            from langgraph.runtime import Runtime137813791380            @before_agent1381            async def notify_start(state: AgentState, runtime: Runtime) -> None:1382                '''Notify user that agent is starting.'''1383                runtime.stream_writer(1384                    {1385                        "type": "status",1386                        "message": "Initializing agent session...",1387                    }1388                )1389                # Perform prerequisite tasks here1390                runtime.stream_writer({"type": "status", "message": "Agent ready!"})139113921393            agent = create_agent(1394                model="openai:gpt-5.2",1395                tools=[...],1396                middleware=[notify_start],1397            )13981399            # Consume with stream_mode="custom" to receive events1400            async for mode, event in agent.astream(1401                {"messages": [HumanMessage("Hello")]},1402                stream_mode=["updates", "custom"],1403            ):1404                if mode == "custom":1405                    print(f"Status: {event}")1406            ```1407    """14081409    def decorator(1410        func: _CallableWithStateAndRuntime[StateT, ContextT],1411    ) -> AgentMiddleware[StateT, ContextT]:1412        is_async = iscoroutinefunction(func)14131414        func_can_jump_to = (1415            can_jump_to if can_jump_to is not None else getattr(func, "__can_jump_to__", [])1416        )14171418        if is_async:14191420            async def async_wrapped(1421                _self: AgentMiddleware[StateT, ContextT],1422                state: StateT,1423                runtime: Runtime[ContextT],1424            ) -> dict[str, Any] | Command[Any] | None:1425                # `iscoroutinefunction` narrows `func` at runtime, but type checkers1426                # cannot narrow this sync-or-async callable union.1427                return await cast("_AsyncCallableWithStateAndRuntime[StateT, ContextT]", func)(1428                    state, runtime1429                )14301431            # Preserve can_jump_to metadata on the wrapped function1432            if func_can_jump_to:1433                async_wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]14341435            middleware_name = name or cast(1436                "str", getattr(func, "__name__", "BeforeAgentMiddleware")1437            )14381439            # `type(...)` builds the correct middleware subclass at runtime, but1440            # type checkers cannot infer its generic `AgentMiddleware` parameters.1441            return cast(1442                "AgentMiddleware[StateT, ContextT]",1443                type(1444                    middleware_name,1445                    (AgentMiddleware,),1446                    {1447                        "state_schema": state_schema or AgentState,1448                        "tools": tools or [],1449                        "abefore_agent": async_wrapped,1450                    },1451                )(),1452            )14531454        def wrapped(1455            _self: AgentMiddleware[StateT, ContextT],1456            state: StateT,1457            runtime: Runtime[ContextT],1458        ) -> dict[str, Any] | Command[Any] | None:1459            # `iscoroutinefunction` narrows `func` at runtime, but type checkers1460            # cannot narrow this sync-or-async callable union.1461            return cast("_SyncCallableWithStateAndRuntime[StateT, ContextT]", func)(state, runtime)14621463        # Preserve can_jump_to metadata on the wrapped function1464        if func_can_jump_to:1465            wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]14661467        # Use function name as default if no name provided1468        middleware_name = name or cast("str", getattr(func, "__name__", "BeforeAgentMiddleware"))14691470        # `type(...)` builds the correct middleware subclass at runtime, but1471        # type checkers cannot infer its generic `AgentMiddleware` parameters.1472        return cast(1473            "AgentMiddleware[StateT, ContextT]",1474            type(1475                middleware_name,1476                (AgentMiddleware,),1477                {1478                    "state_schema": state_schema or AgentState,1479                    "tools": tools or [],1480                    "before_agent": wrapped,1481                },1482            )(),1483        )14841485    if func is not None:1486        return decorator(func)1487    return decorator148814891490@overload1491def after_agent(1492    func: _CallableWithStateAndRuntime[StateT, ContextT],1493) -> AgentMiddleware[StateT, ContextT]: ...149414951496@overload1497def after_agent(1498    func: None = None,1499    *,1500    state_schema: type[StateT] | None = None,1501    tools: list[BaseTool] | None = None,1502    can_jump_to: list[JumpTo] | None = None,1503    name: str | None = None,1504) -> Callable[1505    [_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]1506]: ...150715081509def after_agent(1510    func: _CallableWithStateAndRuntime[StateT, ContextT] | None = None,1511    *,1512    state_schema: type[StateT] | None = None,1513    tools: list[BaseTool] | None = None,1514    can_jump_to: list[JumpTo] | None = None,1515    name: str | None = None,1516) -> (1517    Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]1518    | AgentMiddleware[StateT, ContextT]1519):1520    """Decorator used to dynamically create a middleware with the `after_agent` hook.15211522    Async version is `aafter_agent`.15231524    Args:1525        func: The function to be decorated.15261527            Must accept: `state: StateT, runtime: Runtime[ContextT]` - State and runtime1528            context1529        state_schema: Optional custom state schema type.15301531            If not provided, uses the default `AgentState` schema.1532        tools: Optional list of additional tools to register with this middleware.1533        can_jump_to: Optional list of valid jump destinations for conditional edges.15341535            Valid values are: `'tools'`, `'model'`, `'end'`1536        name: Optional name for the generated middleware class.15371538            If not provided, uses the decorated function's name.15391540    Returns:1541        Either an `AgentMiddleware` instance (if func is provided) or a decorator1542            function that can be applied to a function.15431544    The decorated function should return:15451546    - `dict[str, Any]` - State updates to merge into the agent state1547    - `Command` - A command to control flow (e.g., jump to different node)1548    - `None` - No state updates or flow control15491550    Examples:1551        !!! example "Basic usage for logging agent completion"15521553            ```python1554            @after_agent1555            def log_completion(state: AgentState, runtime: Runtime) -> None:1556                print(f"Agent completed with {len(state['messages'])} messages")1557            ```15581559        !!! example "With custom state schema"15601561            ```python1562            @after_agent(state_schema=MyCustomState, name="MyAfterAgentMiddleware")1563            def custom_after_agent(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:1564                return {"custom_field": "finalized_value"}1565            ```15661567        !!! example "Streaming custom events on completion"15681569            Use `runtime.stream_writer` to emit custom events when agent completes.1570            Events are received when streaming with `stream_mode="custom"`.15711572            ```python1573            @after_agent1574            async def notify_completion(state: AgentState, runtime: Runtime) -> None:1575                '''Notify user that agent has completed.'''1576                runtime.stream_writer(1577                    {1578                        "type": "status",1579                        "message": "Agent execution complete!",1580                        "total_messages": len(state["messages"]),1581                    }1582                )1583            ```1584    """15851586    def decorator(1587        func: _CallableWithStateAndRuntime[StateT, ContextT],1588    ) -> AgentMiddleware[StateT, ContextT]:1589        is_async = iscoroutinefunction(func)1590        # Extract can_jump_to from decorator parameter or from function metadata1591        func_can_jump_to = (1592            can_jump_to if can_jump_to is not None else getattr(func, "__can_jump_to__", [])1593        )15941595        if is_async:15961597            async def async_wrapped(1598                _self: AgentMiddleware[StateT, ContextT],1599                state: StateT,1600                runtime: Runtime[ContextT],1601            ) -> dict[str, Any] | Command[Any] | None:1602                # `iscoroutinefunction` narrows `func` at runtime, but type checkers1603                # cannot narrow this sync-or-async callable union.1604                return await cast("_AsyncCallableWithStateAndRuntime[StateT, ContextT]", func)(1605                    state, runtime1606                )16071608            # Preserve can_jump_to metadata on the wrapped function1609            if func_can_jump_to:1610                async_wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]16111612            middleware_name = name or cast("str", getattr(func, "__name__", "AfterAgentMiddleware"))16131614            # `type(...)` builds the correct middleware subclass at runtime, but1615            # type checkers cannot infer its generic `AgentMiddleware` parameters.1616            return cast(1617                "AgentMiddleware[StateT, ContextT]",1618                type(1619                    middleware_name,1620                    (AgentMiddleware,),1621                    {1622                        "state_schema": state_schema or AgentState,1623                        "tools": tools or [],1624                        "aafter_agent": async_wrapped,1625                    },1626                )(),1627            )16281629        def wrapped(1630            _self: AgentMiddleware[StateT, ContextT],1631            state: StateT,1632            runtime: Runtime[ContextT],1633        ) -> dict[str, Any] | Command[Any] | None:1634            # `iscoroutinefunction` narrows `func` at runtime, but type checkers1635            # cannot narrow this sync-or-async callable union.1636            return cast("_SyncCallableWithStateAndRuntime[StateT, ContextT]", func)(state, runtime)16371638        # Preserve can_jump_to metadata on the wrapped function1639        if func_can_jump_to:1640            wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]16411642        # Use function name as default if no name provided1643        middleware_name = name or cast("str", getattr(func, "__name__", "AfterAgentMiddleware"))16441645        # `type(...)` builds the correct middleware subclass at runtime, but1646        # type checkers cannot infer its generic `AgentMiddleware` parameters.1647        return cast(1648            "AgentMiddleware[StateT, ContextT]",1649            type(1650                middleware_name,1651                (AgentMiddleware,),1652                {1653                    "state_schema": state_schema or AgentState,1654                    "tools": tools or [],1655                    "after_agent": wrapped,1656                },1657            )(),1658        )16591660    if func is not None:1661        return decorator(func)1662    return decorator166316641665@overload1666def dynamic_prompt(1667    func: _CallableReturningSystemMessage[ContextT],1668) -> AgentMiddleware[StateT, ContextT]: ...166916701671@overload1672def dynamic_prompt(1673    func: None = None,1674) -> Callable[1675    [_CallableReturningSystemMessage[ContextT]],1676    AgentMiddleware[StateT, ContextT],1677]: ...167816791680def dynamic_prompt(1681    func: _CallableReturningSystemMessage[ContextT] | None = None,1682) -> (1683    Callable[1684        [_CallableReturningSystemMessage[ContextT]],1685        AgentMiddleware[StateT, ContextT],1686    ]1687    | AgentMiddleware[StateT, ContextT]1688):1689    """Decorator used to dynamically generate system prompts for the model.16901691    This is a convenience decorator that creates middleware using `wrap_model_call`1692    specifically for dynamic prompt generation. The decorated function should return1693    a string that will be set as the system prompt for the model request.16941695    Args:1696        func: The function to be decorated.16971698            Must accept: `request: ModelRequest` - Model request (contains state and1699            runtime)17001701    Returns:1702        Either an `AgentMiddleware` instance (if func is provided) or a decorator1703            function that can be applied to a function.17041705    The decorated function should return:1706        - `str` – The system prompt string to use for the model request1707        - `SystemMessage` – A complete system message to use for the model request17081709    Examples:1710        Basic usage with dynamic content:17111712        ```python1713        @dynamic_prompt1714        def my_prompt(request: ModelRequest) -> str:1715            user_name = request.runtime.context.get("user_name", "User")1716            return f"You are a helpful assistant helping {user_name}."1717        ```17181719        Using state to customize the prompt:17201721        ```python1722        @dynamic_prompt1723        def context_aware_prompt(request: ModelRequest) -> str:1724            msg_count = len(request.state["messages"])1725            if msg_count > 10:1726                return "You are in a long conversation. Be concise."1727            return "You are a helpful assistant."1728        ```17291730        Using with agent:17311732        ```python1733        agent = create_agent(model, middleware=[my_prompt])1734        ```1735    """17361737    def decorator(1738        func: _CallableReturningSystemMessage[ContextT],1739    ) -> AgentMiddleware[StateT, ContextT]:1740        is_async = iscoroutinefunction(func)17411742        if is_async:17431744            async def async_wrapped(1745                _self: AgentMiddleware[StateT, ContextT],1746                request: ModelRequest[ContextT],1747                handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[Any]]],1748            ) -> ModelResponse[Any] | AIMessage:1749                prompt = await cast("_AsyncCallableReturningSystemMessage[ContextT]", func)(request)1750                if isinstance(prompt, SystemMessage):1751                    request = request.override(system_message=prompt)1752                else:1753                    request = request.override(system_message=SystemMessage(content=prompt))1754                return await handler(request)17551756            middleware_name = cast("str", getattr(func, "__name__", "DynamicPromptMiddleware"))17571758            # `type(...)` builds the correct middleware subclass at runtime, but1759            # type checkers cannot infer its generic `AgentMiddleware` parameters.1760            return cast(1761                "AgentMiddleware[StateT, ContextT]",1762                type(1763                    middleware_name,1764                    (AgentMiddleware,),1765                    {1766                        "state_schema": AgentState,1767                        "tools": [],1768                        "awrap_model_call": async_wrapped,1769                    },1770                )(),1771            )17721773        def wrapped(1774            _self: AgentMiddleware[StateT, ContextT],1775            request: ModelRequest[ContextT],1776            handler: Callable[[ModelRequest[ContextT]], ModelResponse[Any]],1777        ) -> ModelResponse[Any] | AIMessage:1778            prompt = cast("_SyncCallableReturningSystemMessage[ContextT]", func)(request)1779            if isinstance(prompt, SystemMessage):1780                request = request.override(system_message=prompt)1781            else:1782                request = request.override(system_message=SystemMessage(content=prompt))1783            return handler(request)17841785        async def async_wrapped_from_sync(1786            _self: AgentMiddleware[StateT, ContextT],1787            request: ModelRequest[ContextT],1788            handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[Any]]],1789        ) -> ModelResponse[Any] | AIMessage:1790            # Delegate to sync function1791            prompt = cast("_SyncCallableReturningSystemMessage[ContextT]", func)(request)1792            if isinstance(prompt, SystemMessage):1793                request = request.override(system_message=prompt)1794            else:1795                request = request.override(system_message=SystemMessage(content=prompt))1796            return await handler(request)17971798        middleware_name = cast("str", getattr(func, "__name__", "DynamicPromptMiddleware"))17991800        # `type(...)` builds the correct middleware subclass at runtime, but1801        # type checkers cannot infer its generic `AgentMiddleware` parameters.1802        return cast(1803            "AgentMiddleware[StateT, ContextT]",1804            type(1805                middleware_name,1806                (AgentMiddleware,),1807                {1808                    "state_schema": AgentState,1809                    "tools": [],1810                    "wrap_model_call": wrapped,1811                    "awrap_model_call": async_wrapped_from_sync,1812                },1813            )(),1814        )18151816    if func is not None:1817        return decorator(func)1818    return decorator181918201821@overload1822def wrap_model_call(1823    func: _CallableReturningModelResponse[ContextT, ResponseT],1824) -> AgentMiddleware[StateT, ContextT]: ...182518261827@overload1828def wrap_model_call(1829    func: None = None,1830    *,1831    state_schema: type[StateT] | None = None,1832    tools: list[BaseTool] | None = None,1833    name: str | None = None,1834) -> Callable[1835    [_CallableReturningModelResponse[ContextT, ResponseT]],1836    AgentMiddleware[StateT, ContextT],1837]: ...183818391840def wrap_model_call(1841    func: _CallableReturningModelResponse[ContextT, ResponseT] | None = None,1842    *,1843    state_schema: type[StateT] | None = None,1844    tools: list[BaseTool] | None = None,1845    name: str | None = None,1846) -> (1847    Callable[1848        [_CallableReturningModelResponse[ContextT, ResponseT]],1849        AgentMiddleware[StateT, ContextT],1850    ]1851    | AgentMiddleware[StateT, ContextT]1852):1853    """Create middleware with `wrap_model_call` hook from a function.18541855    Converts a function with handler callback into middleware that can intercept model1856    calls, implement retry logic, handle errors, and rewrite responses.18571858    Args:1859        func: Function accepting (request, handler) that calls handler(request)1860            to execute the model and returns `ModelResponse` or `AIMessage`.18611862            Request contains state and runtime.1863        state_schema: Custom state schema.18641865            Defaults to `AgentState`.1866        tools: Additional tools to register with this middleware.1867        name: Middleware class name.18681869            Defaults to function name.18701871    Returns:1872        `AgentMiddleware` instance if func provided, otherwise a decorator.18731874    Examples:1875        !!! example "Basic retry logic"18761877            ```python1878            @wrap_model_call1879            def retry_on_error(request, handler):1880                max_retries = 31881                for attempt in range(max_retries):1882                    try:1883                        return handler(request)1884                    except Exception:1885                        if attempt == max_retries - 1:1886                            raise1887            ```18881889        !!! example "Model fallback"18901891            ```python1892            @wrap_model_call1893            def fallback_model(request, handler):1894                # Try primary model1895                try:1896                    return handler(request)1897                except Exception:1898                    pass18991900                # Try fallback model1901                request = request.override(model=fallback_model_instance)1902                return handler(request)1903            ```19041905        !!! example "Rewrite response content (full `ModelResponse`)"19061907            ```python1908            @wrap_model_call1909            def uppercase_responses(request, handler):1910                response = handler(request)1911                ai_msg = response.result[0]1912                return ModelResponse(1913                    result=[AIMessage(content=ai_msg.content.upper())],1914                    structured_response=response.structured_response,1915                )1916            ```19171918        !!! example "Simple `AIMessage` return (converted automatically)"19191920            ```python1921            @wrap_model_call1922            def simple_response(request, handler):1923                # AIMessage is automatically converted to ModelResponse1924                return AIMessage(content="Simple response")1925            ```1926    """19271928    def decorator(1929        func: _CallableReturningModelResponse[ContextT, ResponseT],1930    ) -> AgentMiddleware[StateT, ContextT]:1931        is_async = iscoroutinefunction(func)19321933        if is_async:19341935            async def async_wrapped(1936                _self: AgentMiddleware[StateT, ContextT],1937                request: ModelRequest[ContextT],1938                handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],1939            ) -> ModelCallResult[ResponseT]:1940                return await cast(1941                    "_AsyncCallableReturningModelResponse[ContextT, ResponseT]", func1942                )(request, handler)19431944            middleware_name = name or cast(1945                "str", getattr(func, "__name__", "WrapModelCallMiddleware")1946            )19471948            # `type(...)` builds the correct middleware subclass at runtime, but1949            # type checkers cannot infer its generic `AgentMiddleware` parameters.1950            return cast(1951                "AgentMiddleware[StateT, ContextT]",1952                type(1953                    middleware_name,1954                    (AgentMiddleware,),1955                    {1956                        "state_schema": state_schema or AgentState,1957                        "tools": tools or [],1958                        "awrap_model_call": async_wrapped,1959                    },1960                )(),1961            )19621963        def wrapped(1964            _self: AgentMiddleware[StateT, ContextT],1965            request: ModelRequest[ContextT],1966            handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],1967        ) -> ModelCallResult[ResponseT]:1968            return cast("_SyncCallableReturningModelResponse[ContextT, ResponseT]", func)(1969                request, handler1970            )19711972        middleware_name = name or cast("str", getattr(func, "__name__", "WrapModelCallMiddleware"))19731974        # `type(...)` builds the correct middleware subclass at runtime, but1975        # type checkers cannot infer its generic `AgentMiddleware` parameters.1976        return cast(1977            "AgentMiddleware[StateT, ContextT]",1978            type(1979                middleware_name,1980                (AgentMiddleware,),1981                {1982                    "state_schema": state_schema or AgentState,1983                    "tools": tools or [],1984                    "wrap_model_call": wrapped,1985                },1986            )(),1987        )19881989    if func is not None:1990        return decorator(func)1991    return decorator199219931994@overload1995def wrap_tool_call(1996    func: _CallableReturningToolResponse,1997) -> AgentMiddleware: ...199819992000@overload