langchain 1.0.5__py3-none-any.whl → 1.2.4__py3-none-any.whl

langchain/agents/middleware/types.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field, replace
 from inspect import iscoroutinefunction
 from typing import (
@@ -19,19 +19,22 @@ from typing import (
 if TYPE_CHECKING:
     from collections.abc import Awaitable
 
+    from langgraph.types import Command
+
 # Needed as top level import for Pydantic schema generation on AgentState
+import warnings
 from typing import TypeAlias
 
-from langchain_core.messages import (  # noqa: TC002
+from langchain_core.messages import (
     AIMessage,
     AnyMessage,
     BaseMessage,
+    SystemMessage,
     ToolMessage,
 )
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt.tool_node import ToolCallRequest, ToolCallWrapper
-from langgraph.types import Command  # noqa: TC002
 from langgraph.typing import ContextT
 from typing_extensions import NotRequired, Required, TypedDict, TypeVar, Unpack
 
@@ -69,59 +72,194 @@ ResponseT = TypeVar("ResponseT")
 
 
 class _ModelRequestOverrides(TypedDict, total=False):
-    """Possible overrides for ModelRequest.override() method."""
+    """Possible overrides for `ModelRequest.override()` method."""
 
     model: BaseChatModel
-    system_prompt: str | None
+    system_message: SystemMessage | None
     messages: list[AnyMessage]
     tool_choice: Any | None
-    tools: list[BaseTool | dict]
-    response_format: ResponseFormat | None
+    tools: list[BaseTool | dict[str, Any]]
+    response_format: ResponseFormat[Any] | None
     model_settings: dict[str, Any]
+    state: AgentState[Any]
 
 
-@dataclass
+@dataclass(init=False)
 class ModelRequest:
     """Model request information for the agent."""
 
     model: BaseChatModel
-    system_prompt: str | None
-    messages: list[AnyMessage]  # excluding system prompt
+    messages: list[AnyMessage]  # excluding system message
+    system_message: SystemMessage | None
     tool_choice: Any | None
-    tools: list[BaseTool | dict]
-    response_format: ResponseFormat | None
-    state: AgentState
+    tools: list[BaseTool | dict[str, Any]]
+    response_format: ResponseFormat[Any] | None
+    state: AgentState[Any]
     runtime: Runtime[ContextT]  # type: ignore[valid-type]
     model_settings: dict[str, Any] = field(default_factory=dict)
 
+    def __init__(
+        self,
+        *,
+        model: BaseChatModel,
+        messages: list[AnyMessage],
+        system_message: SystemMessage | None = None,
+        system_prompt: str | None = None,
+        tool_choice: Any | None = None,
+        tools: list[BaseTool | dict[str, Any]] | None = None,
+        response_format: ResponseFormat[Any] | None = None,
+        state: AgentState[Any] | None = None,
+        runtime: Runtime[ContextT] | None = None,
+        model_settings: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize ModelRequest with backward compatibility for system_prompt.
+
+        Args:
+            model: The chat model to use.
+            messages: List of messages (excluding system prompt).
+            tool_choice: Tool choice configuration.
+            tools: List of available tools.
+            response_format: Response format specification.
+            state: Agent state.
+            runtime: Runtime context.
+            model_settings: Additional model settings.
+            system_message: System message instance (preferred).
+            system_prompt: System prompt string (deprecated, converted to SystemMessage).
+
+        Raises:
+            ValueError: If both `system_prompt` and `system_message` are provided.
+        """
+        # Handle system_prompt/system_message conversion and validation
+        if system_prompt is not None and system_message is not None:
+            msg = "Cannot specify both system_prompt and system_message"
+            raise ValueError(msg)
+
+        if system_prompt is not None:
+            system_message = SystemMessage(content=system_prompt)
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=DeprecationWarning)
+            self.model = model
+            self.messages = messages
+            self.system_message = system_message
+            self.tool_choice = tool_choice
+            self.tools = tools if tools is not None else []
+            self.response_format = response_format
+            self.state = state if state is not None else {"messages": []}
+            self.runtime = runtime  # type: ignore[assignment]
+            self.model_settings = model_settings if model_settings is not None else {}
+
+    @property
+    def system_prompt(self) -> str | None:
+        """Get system prompt text from system_message.
+
+        Returns:
+            The content of the system message if present, otherwise `None`.
+        """
+        if self.system_message is None:
+            return None
+        return self.system_message.text
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Set an attribute with a deprecation warning.
+
+        Direct attribute assignment on `ModelRequest` is deprecated. Use the
+        `override()` method instead to create a new request with modified attributes.
+
+        Args:
+            name: Attribute name.
+            value: Attribute value.
+        """
+        # Special handling for system_prompt - convert to system_message
+        if name == "system_prompt":
+            warnings.warn(
+                "Direct attribute assignment to ModelRequest.system_prompt is deprecated. "
+                "Use request.override(system_message=SystemMessage(...)) instead to create "
+                "a new request with the modified system message.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if value is None:
+                object.__setattr__(self, "system_message", None)
+            else:
+                object.__setattr__(self, "system_message", SystemMessage(content=value))
+            return
+
+        warnings.warn(
+            f"Direct attribute assignment to ModelRequest.{name} is deprecated. "
+            f"Use request.override({name}=...) instead to create a new request "
+            f"with the modified attribute.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        object.__setattr__(self, name, value)
+
     def override(self, **overrides: Unpack[_ModelRequestOverrides]) -> ModelRequest:
         """Replace the request with a new request with the given overrides.
 
         Returns a new `ModelRequest` instance with the specified attributes replaced.
+
         This follows an immutable pattern, leaving the original request unchanged.
 
         Args:
-            **overrides: Keyword arguments for attributes to override. Supported keys:
-                - model: BaseChatModel instance
-                - system_prompt: Optional system prompt string
-                - messages: List of messages
-                - tool_choice: Tool choice configuration
-                - tools: List of available tools
-                - response_format: Response format specification
-                - model_settings: Additional model settings
+            **overrides: Keyword arguments for attributes to override.
+
+                Supported keys:
+
+                - `model`: `BaseChatModel` instance
+                - `system_prompt`: deprecated, use `system_message` instead
+                - `system_message`: `SystemMessage` instance
+                - `messages`: `list` of messages
+                - `tool_choice`: Tool choice configuration
+                - `tools`: `list` of available tools
+                - `response_format`: Response format specification
+                - `model_settings`: Additional model settings
+                - `state`: Agent state dictionary
 
         Returns:
-            New ModelRequest instance with specified overrides applied.
+            New `ModelRequest` instance with specified overrides applied.
 
         Examples:
-            ```python
-            # Create a new request with different model
-            new_request = request.override(model=different_model)
+            !!! example "Create a new request with different model"
 
-            # Override multiple attributes
-            new_request = request.override(system_prompt="New instructions", tool_choice="auto")
-            ```
+                ```python
+                new_request = request.override(model=different_model)
+                ```
+
+            !!! example "Override system message (preferred)"
+
+                ```python
+                from langchain_core.messages import SystemMessage
+
+                new_request = request.override(
+                    system_message=SystemMessage(content="New instructions")
+                )
+                ```
+
+            !!! example "Override multiple attributes"
+
+                ```python
+                new_request = request.override(
+                    model=ChatOpenAI(model="gpt-4o"),
+                    system_message=SystemMessage(content="New instructions"),
+                )
+                ```
+
+        Raises:
+            ValueError: If both `system_prompt` and `system_message` are provided.
         """
+        # Handle system_prompt/system_message conversion
+        if "system_prompt" in overrides and "system_message" in overrides:
+            msg = "Cannot specify both system_prompt and system_message"
+            raise ValueError(msg)
+
+        if "system_prompt" in overrides:
+            system_prompt = cast("str | None", overrides.pop("system_prompt"))  # type: ignore[typeddict-item]
+            if system_prompt is None:
+                overrides["system_message"] = None
+            else:
+                overrides["system_message"] = SystemMessage(content=system_prompt)
+
         return replace(self, **overrides)
 
 
@@ -129,24 +267,25 @@ class ModelRequest:
 class ModelResponse:
     """Response from model execution including messages and optional structured output.
 
-    The result will usually contain a single AIMessage, but may include
-    an additional ToolMessage if the model used a tool for structured output.
+    The result will usually contain a single `AIMessage`, but may include an additional
+    `ToolMessage` if the model used a tool for structured output.
     """
 
     result: list[BaseMessage]
     """List of messages from model execution."""
 
     structured_response: Any = None
-    """Parsed structured output if response_format was specified, None otherwise."""
+    """Parsed structured output if `response_format` was specified, `None` otherwise."""
 
 
 # Type alias for middleware return type - allows returning either full response or just AIMessage
-ModelCallResult: TypeAlias = "ModelResponse | AIMessage"
-"""Type alias for model call handler return value.
+ModelCallResult: TypeAlias = ModelResponse | AIMessage
+"""`TypeAlias` for model call handler return value.
 
 Middleware can return either:
-- ModelResponse: Full response with messages and optional structured output
-- AIMessage: Simplified return for simple use cases
+
+- `ModelResponse`: Full response with messages and optional structured output
+- `AIMessage`: Simplified return for simple use cases
 """
 
 
@@ -182,7 +321,7 @@ class AgentState(TypedDict, Generic[ResponseT]):
 class _InputAgentState(TypedDict):  # noqa: PYI049
     """Input state schema for the agent."""
 
-    messages: Required[Annotated[list[AnyMessage | dict], add_messages]]
+    messages: Required[Annotated[list[AnyMessage | dict[str, Any]], add_messages]]
 
 
 class _OutputAgentState(TypedDict, Generic[ResponseT]):  # noqa: PYI049
@@ -192,9 +331,13 @@ class _OutputAgentState(TypedDict, Generic[ResponseT]):  # noqa: PYI049
     structured_response: NotRequired[ResponseT]
 
 
-StateT = TypeVar("StateT", bound=AgentState, default=AgentState)
-StateT_co = TypeVar("StateT_co", bound=AgentState, default=AgentState, covariant=True)
-StateT_contra = TypeVar("StateT_contra", bound=AgentState, contravariant=True)
+StateT = TypeVar("StateT", bound=AgentState[Any], default=AgentState[Any])
+StateT_co = TypeVar("StateT_co", bound=AgentState[Any], default=AgentState[Any], covariant=True)
+StateT_contra = TypeVar("StateT_contra", bound=AgentState[Any], contravariant=True)
+
+
+class _DefaultAgentState(AgentState[Any]):
+    """AgentMiddleware default state."""
 
 
 class AgentMiddleware(Generic[StateT, ContextT]):
@@ -204,10 +347,10 @@ class AgentMiddleware(Generic[StateT, ContextT]):
     between steps in the main agent loop.
     """
 
-    state_schema: type[StateT] = cast("type[StateT]", AgentState)
+    state_schema: type[StateT] = cast("type[StateT]", _DefaultAgentState)
     """The schema for state passed to the middleware nodes."""
 
-    tools: list[BaseTool]
+    tools: Sequence[BaseTool]
     """Additional tools registered by the middleware."""
 
     @property
@@ -219,28 +362,76 @@ class AgentMiddleware(Generic[StateT, ContextT]):
         return self.__class__.__name__
 
     def before_agent(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
-        """Logic to run before the agent execution starts."""
+        """Logic to run before the agent execution starts.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Agent state updates to apply before agent execution.
+        """
 
     async def abefore_agent(
         self, state: StateT, runtime: Runtime[ContextT]
     ) -> dict[str, Any] | None:
-        """Async logic to run before the agent execution starts."""
+        """Async logic to run before the agent execution starts.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Agent state updates to apply before agent execution.
+        """
 
     def before_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
-        """Logic to run before the model is called."""
+        """Logic to run before the model is called.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Agent state updates to apply before model call.
+        """
 
     async def abefore_model(
         self, state: StateT, runtime: Runtime[ContextT]
     ) -> dict[str, Any] | None:
-        """Async logic to run before the model is called."""
+        """Async logic to run before the model is called.
+
+        Args:
+            state: The agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Agent state updates to apply before model call.
+        """
 
     def after_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
-        """Logic to run after the model is called."""
+        """Logic to run after the model is called.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Agent state updates to apply after model call.
+        """
 
     async def aafter_model(
         self, state: StateT, runtime: Runtime[ContextT]
     ) -> dict[str, Any] | None:
-        """Async logic to run after the model is called."""
+        """Async logic to run after the model is called.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Agent state updates to apply after model call.
+        """
 
     def wrap_model_call(
         self,
@@ -249,6 +440,8 @@ class AgentMiddleware(Generic[StateT, ContextT]):
     ) -> ModelCallResult:
         """Intercept and control model execution via handler callback.
 
+        Async version is `awrap_model_call`
+
         The handler callback executes the model request and returns a `ModelResponse`.
         Middleware can call the handler multiple times for retry logic, skip calling
         it to short-circuit, or modify the request/response. Multiple middleware
@@ -257,61 +450,71 @@ class AgentMiddleware(Generic[StateT, ContextT]):
         Args:
             request: Model request to execute (includes state and runtime).
             handler: Callback that executes the model request and returns
-                `ModelResponse`. Call this to execute the model. Can be called multiple
-                times for retry logic. Can skip calling it to short-circuit.
+                `ModelResponse`.
+
+                Call this to execute the model.
+
+                Can be called multiple times for retry logic.
+
+                Can skip calling it to short-circuit.
 
         Returns:
-            `ModelCallResult`
+            The model call result.
 
         Examples:
-            Retry on error:
-            ```python
-            def wrap_model_call(self, request, handler):
-                for attempt in range(3):
+            !!! example "Retry on error"
+
+                ```python
+                def wrap_model_call(self, request, handler):
+                    for attempt in range(3):
+                        try:
+                            return handler(request)
+                        except Exception:
+                            if attempt == 2:
+                                raise
+                ```
+
+            !!! example "Rewrite response"
+
+                ```python
+                def wrap_model_call(self, request, handler):
+                    response = handler(request)
+                    ai_msg = response.result[0]
+                    return ModelResponse(
+                        result=[AIMessage(content=f"[{ai_msg.content}]")],
+                        structured_response=response.structured_response,
+                    )
+                ```
+
+            !!! example "Error to fallback"
+
+                ```python
+                def wrap_model_call(self, request, handler):
                     try:
                         return handler(request)
                     except Exception:
-                        if attempt == 2:
-                            raise
-            ```
-
-            Rewrite response:
-            ```python
-            def wrap_model_call(self, request, handler):
-                response = handler(request)
-                ai_msg = response.result[0]
-                return ModelResponse(
-                    result=[AIMessage(content=f"[{ai_msg.content}]")],
-                    structured_response=response.structured_response,
-                )
-            ```
-
-            Error to fallback:
-            ```python
-            def wrap_model_call(self, request, handler):
-                try:
-                    return handler(request)
-                except Exception:
-                    return ModelResponse(result=[AIMessage(content="Service unavailable")])
-            ```
-
-            Cache/short-circuit:
-            ```python
-            def wrap_model_call(self, request, handler):
-                if cached := get_cache(request):
-                    return cached  # Short-circuit with cached result
-                response = handler(request)
-                save_cache(request, response)
-                return response
-            ```
-
-            Simple AIMessage return (converted automatically):
-            ```python
-            def wrap_model_call(self, request, handler):
-                response = handler(request)
-                # Can return AIMessage directly for simple cases
-                return AIMessage(content="Simplified response")
-            ```
+                        return ModelResponse(result=[AIMessage(content="Service unavailable")])
+                ```
+
+            !!! example "Cache/short-circuit"
+
+                ```python
+                def wrap_model_call(self, request, handler):
+                    if cached := get_cache(request):
+                        return cached  # Short-circuit with cached result
+                    response = handler(request)
+                    save_cache(request, response)
+                    return response
+                ```
+
+            !!! example "Simple `AIMessage` return (converted automatically)"
+
+                ```python
+                def wrap_model_call(self, request, handler):
+                    response = handler(request)
+                    # Can return AIMessage directly for simple cases
+                    return AIMessage(content="Simplified response")
+                ```
         """
         msg = (
             "Synchronous implementation of wrap_model_call is not available. "
@@ -333,6 +536,7 @@ class AgentMiddleware(Generic[StateT, ContextT]):
         """Intercept and control async model execution via handler callback.
 
         The handler callback executes the model request and returns a `ModelResponse`.
+
         Middleware can call the handler multiple times for retry logic, skip calling
         it to short-circuit, or modify the request/response. Multiple middleware
         compose with first in list as outermost layer.
@@ -340,23 +544,29 @@ class AgentMiddleware(Generic[StateT, ContextT]):
         Args:
             request: Model request to execute (includes state and runtime).
             handler: Async callback that executes the model request and returns
-                `ModelResponse`. Call this to execute the model. Can be called multiple
-                times for retry logic. Can skip calling it to short-circuit.
+                `ModelResponse`.
+
+                Call this to execute the model.
+
+                Can be called multiple times for retry logic.
+
+                Can skip calling it to short-circuit.
 
         Returns:
-            ModelCallResult
+            The model call result.
 
         Examples:
-            Retry on error:
-            ```python
-            async def awrap_model_call(self, request, handler):
-                for attempt in range(3):
-                    try:
-                        return await handler(request)
-                    except Exception:
-                        if attempt == 2:
-                            raise
-            ```
+            !!! example "Retry on error"
+
+                ```python
+                async def awrap_model_call(self, request, handler):
+                    for attempt in range(3):
+                        try:
+                            return await handler(request)
+                        except Exception:
+                            if attempt == 2:
+                                raise
+                ```
         """
         msg = (
             "Asynchronous implementation of awrap_model_call is not available. "
@@ -371,70 +581,98 @@ class AgentMiddleware(Generic[StateT, ContextT]):
         raise NotImplementedError(msg)
 
     def after_agent(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
-        """Logic to run after the agent execution completes."""
+        """Logic to run after the agent execution completes.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Agent state updates to apply after agent execution.
+        """
 
     async def aafter_agent(
         self, state: StateT, runtime: Runtime[ContextT]
     ) -> dict[str, Any] | None:
-        """Async logic to run after the agent execution completes."""
+        """Async logic to run after the agent execution completes.
+
+        Args:
+            state: The current agent state.
+            runtime: The runtime context.
+
+        Returns:
+            Agent state updates to apply after agent execution.
+        """
 
     def wrap_tool_call(
         self,
         request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], ToolMessage | Command],
-    ) -> ToolMessage | Command:
+        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
+    ) -> ToolMessage | Command[Any]:
         """Intercept tool execution for retries, monitoring, or modification.
 
+        Async version is `awrap_tool_call`
+
         Multiple middleware compose automatically (first defined = outermost).
+
         Exceptions propagate unless `handle_tool_errors` is configured on `ToolNode`.
 
         Args:
             request: Tool call request with call `dict`, `BaseTool`, state, and runtime.
+
                 Access state via `request.state` and runtime via `request.runtime`.
-            handler: Callable to execute the tool (can be called multiple times).
+            handler: `Callable` to execute the tool (can be called multiple times).
 
         Returns:
             `ToolMessage` or `Command` (the final result).
 
-        The handler callable can be invoked multiple times for retry logic.
+        The handler `Callable` can be invoked multiple times for retry logic.
+
         Each call to handler is independent and stateless.
 
         Examples:
-            Modify request before execution:
-
-            ```python
-            def wrap_tool_call(self, request, handler):
-                request.tool_call["args"]["value"] *= 2
-                return handler(request)
-            ```
+            !!! example "Modify request before execution"
+
+                ```python
+                def wrap_tool_call(self, request, handler):
+                    modified_call = {
+                        **request.tool_call,
+                        "args": {
+                            **request.tool_call["args"],
+                            "value": request.tool_call["args"]["value"] * 2,
+                        },
+                    }
+                    request = request.override(tool_call=modified_call)
+                    return handler(request)
+                ```
+
+            !!! example "Retry on error (call handler multiple times)"
+
+                ```python
+                def wrap_tool_call(self, request, handler):
+                    for attempt in range(3):
+                        try:
+                            result = handler(request)
+                            if is_valid(result):
+                                return result
+                        except Exception:
+                            if attempt == 2:
+                                raise
+                    return result
+                ```
 
-            Retry on error (call handler multiple times):
+            !!! example "Conditional retry based on response"
 
-            ```python
-            def wrap_tool_call(self, request, handler):
-                for attempt in range(3):
-                    try:
+                ```python
+                def wrap_tool_call(self, request, handler):
+                    for attempt in range(3):
                         result = handler(request)
-                        if is_valid(result):
+                        if isinstance(result, ToolMessage) and result.status != "error":
                             return result
-                    except Exception:
-                        if attempt == 2:
-                            raise
-                return result
-            ```
-
-            Conditional retry based on response:
-
-            ```python
-            def wrap_tool_call(self, request, handler):
-                for attempt in range(3):
-                    result = handler(request)
-                    if isinstance(result, ToolMessage) and result.status != "error":
+                        if attempt < 2:
+                            continue
                         return result
-                    if attempt < 2:
-                        continue
-                    return result
-            ```
+                ```
         """
         msg = (
             "Synchronous implementation of wrap_tool_call is not available. "
@@ -451,8 +689,8 @@ class AgentMiddleware(Generic[StateT, ContextT]):
     async def awrap_tool_call(
         self,
         request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
-    ) -> ToolMessage | Command:
+        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],
+    ) -> ToolMessage | Command[Any]:
         """Intercept and control async tool execution via handler callback.
 
         The handler callback executes the tool call and returns a `ToolMessage` or
@@ -462,40 +700,48 @@ class AgentMiddleware(Generic[StateT, ContextT]):
 
         Args:
             request: Tool call request with call `dict`, `BaseTool`, state, and runtime.
+
                 Access state via `request.state` and runtime via `request.runtime`.
             handler: Async callable to execute the tool and returns `ToolMessage` or
-                `Command`. Call this to execute the tool. Can be called multiple times
-                for retry logic. Can skip calling it to short-circuit.
+                `Command`.
+
+                Call this to execute the tool.
+
+                Can be called multiple times for retry logic.
+
+                Can skip calling it to short-circuit.
 
         Returns:
             `ToolMessage` or `Command` (the final result).
 
-        The handler callable can be invoked multiple times for retry logic.
+        The handler `Callable` can be invoked multiple times for retry logic.
+
         Each call to handler is independent and stateless.
 
         Examples:
-            Async retry on error:
-            ```python
-            async def awrap_tool_call(self, request, handler):
-                for attempt in range(3):
-                    try:
-                        result = await handler(request)
-                        if is_valid(result):
-                            return result
-                    except Exception:
-                        if attempt == 2:
-                            raise
-                return result
-            ```
-
-            ```python
-            async def awrap_tool_call(self, request, handler):
-                if cached := await get_cache_async(request):
-                    return ToolMessage(content=cached, tool_call_id=request.tool_call["id"])
-                result = await handler(request)
-                await save_cache_async(request, result)
-                return result
-            ```
+            !!! example "Async retry on error"
+
+                ```python
+                async def awrap_tool_call(self, request, handler):
+                    for attempt in range(3):
+                        try:
+                            result = await handler(request)
+                            if is_valid(result):
+                                return result
+                        except Exception:
+                            if attempt == 2:
+                                raise
+                    return result
+                ```
+
+                ```python
+                async def awrap_tool_call(self, request, handler):
+                    if cached := await get_cache_async(request):
+                        return ToolMessage(content=cached, tool_call_id=request.tool_call["id"])
+                    result = await handler(request)
+                    await save_cache_async(request, result)
+                    return result
+                ```
         """
         msg = (
             "Asynchronous implementation of awrap_tool_call is not available. "
@@ -515,16 +761,18 @@ class _CallableWithStateAndRuntime(Protocol[StateT_contra, ContextT]):
 
     def __call__(
         self, state: StateT_contra, runtime: Runtime[ContextT]
-    ) -> dict[str, Any] | Command | None | Awaitable[dict[str, Any] | Command | None]:
+    ) -> dict[str, Any] | Command[Any] | None | Awaitable[dict[str, Any] | Command[Any] | None]:
         """Perform some logic with the state and runtime."""
         ...
 
 
-class _CallableReturningPromptString(Protocol[StateT_contra, ContextT]):  # type: ignore[misc]
-    """Callable that returns a prompt string given `ModelRequest` (contains state and runtime)."""
+class _CallableReturningSystemMessage(Protocol[StateT_contra, ContextT]):  # type: ignore[misc]
+    """Callable that returns a prompt string or SystemMessage given `ModelRequest`."""
 
-    def __call__(self, request: ModelRequest) -> str | Awaitable[str]:
-        """Generate a system prompt string based on the request."""
+    def __call__(
+        self, request: ModelRequest
+    ) -> str | SystemMessage | Awaitable[str | SystemMessage]:
+        """Generate a system prompt string or SystemMessage based on the request."""
         ...
 
 
@@ -554,8 +802,8 @@ class _CallableReturningToolResponse(Protocol):
     def __call__(
         self,
         request: ToolCallRequest,
-        handler: Callable[[ToolCallRequest], ToolMessage | Command],
-    ) -> ToolMessage | Command:
+        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
+    ) -> ToolMessage | Command[Any]:
         """Intercept tool execution via handler callback."""
         ...
 
@@ -574,26 +822,32 @@ def hook_config(
     can jump to, which establishes conditional edges in the agent graph.
 
     Args:
-        can_jump_to: Optional list of valid jump destinations. Can be:
-            - "tools": Jump to the tools node
-            - "model": Jump back to the model node
-            - "end": Jump to the end of the graph
+        can_jump_to: Optional list of valid jump destinations.
+
+            Can be:
+
+            - `'tools'`: Jump to the tools node
+            - `'model'`: Jump back to the model node
+            - `'end'`: Jump to the end of the graph
 
     Returns:
         Decorator function that marks the method with configuration metadata.
 
     Examples:
-        Using decorator on a class method:
-        ```python
-        class MyMiddleware(AgentMiddleware):
-            @hook_config(can_jump_to=["end", "model"])
-            def before_model(self, state: AgentState) -> dict[str, Any] | None:
-                if some_condition(state):
-                    return {"jump_to": "end"}
-                return None
-        ```
+        !!! example "Using decorator on a class method"
+
+            ```python
+            class MyMiddleware(AgentMiddleware):
+                @hook_config(can_jump_to=["end", "model"])
+                def before_model(self, state: AgentState) -> dict[str, Any] | None:
+                    if some_condition(state):
+                        return {"jump_to": "end"}
+                    return None
+            ```
+
+        Alternative: Use the `can_jump_to` parameter in `before_model`/`after_model`
+        decorators:
 
-        Alternative: Use the `can_jump_to` parameter in `before_model`/`after_model` decorators:
         ```python
         @before_model(can_jump_to=["end"])
         def conditional_middleware(state: AgentState) -> dict[str, Any] | None:
@@ -644,48 +898,76 @@ def before_model(
     """Decorator used to dynamically create a middleware with the `before_model` hook.
 
     Args:
-        func: The function to be decorated. Must accept:
-            `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
-        state_schema: Optional custom state schema type. If not provided, uses the default
-            `AgentState` schema.
+        func: The function to be decorated.
+
+            Must accept: `state: StateT, runtime: Runtime[ContextT]` - State and runtime
+                context
+        state_schema: Optional custom state schema type.
+
+            If not provided, uses the default `AgentState` schema.
         tools: Optional list of additional tools to register with this middleware.
         can_jump_to: Optional list of valid jump destinations for conditional edges.
-            Valid values are: `"tools"`, `"model"`, `"end"`
-        name: Optional name for the generated middleware class. If not provided,
-            uses the decorated function's name.
+
+            Valid values are: `'tools'`, `'model'`, `'end'`
+        name: Optional name for the generated middleware class.
+
+            If not provided, uses the decorated function's name.
 
     Returns:
         Either an `AgentMiddleware` instance (if func is provided directly) or a
-        decorator function that can be applied to a function it is wrapping.
+            decorator function that can be applied to a function it is wrapping.
 
     The decorated function should return:
-        - `dict[str, Any]` - State updates to merge into the agent state
-        - `Command` - A command to control flow (e.g., jump to different node)
-        - `None` - No state updates or flow control
+
+    - `dict[str, Any]` - State updates to merge into the agent state
+    - `Command` - A command to control flow (e.g., jump to different node)
+    - `None` - No state updates or flow control
 
     Examples:
-        Basic usage:
-        ```python
-        @before_model
-        def log_before_model(state: AgentState, runtime: Runtime) -> None:
-            print(f"About to call model with {len(state['messages'])} messages")
-        ```
+        !!! example "Basic usage"
 
-        With conditional jumping:
-        ```python
-        @before_model(can_jump_to=["end"])
-        def conditional_before_model(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
-            if some_condition(state):
-                return {"jump_to": "end"}
-            return None
-        ```
+            ```python
+            @before_model
+            def log_before_model(state: AgentState, runtime: Runtime) -> None:
+                print(f"About to call model with {len(state['messages'])} messages")
+            ```
 
-        With custom state schema:
-        ```python
-        @before_model(state_schema=MyCustomState)
-        def custom_before_model(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
-            return {"custom_field": "updated_value"}
-        ```
+        !!! example "With conditional jumping"
+
+            ```python
+            @before_model(can_jump_to=["end"])
+            def conditional_before_model(
+                state: AgentState, runtime: Runtime
+            ) -> dict[str, Any] | None:
+                if some_condition(state):
+                    return {"jump_to": "end"}
+                return None
+            ```
+
+        !!! example "With custom state schema"
+
+            ```python
+            @before_model(state_schema=MyCustomState)
+            def custom_before_model(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
+                return {"custom_field": "updated_value"}
+            ```
+
+        !!! example "Streaming custom events before model call"
+
+            Use `runtime.stream_writer` to emit custom events before each model invocation.
+            Events are received when streaming with `stream_mode="custom"`.
+
+            ```python
+            @before_model
+            async def notify_model_call(state: AgentState, runtime: Runtime) -> None:
+                '''Notify user before model is called.'''
+                runtime.stream_writer(
+                    {
+                        "type": "status",
+                        "message": "Thinking...",
+                    }
+                )
+            ```
     """
 
     def decorator(
@@ -700,10 +982,10 @@ def before_model(
         if is_async:
 
             async def async_wrapped(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+                _self: AgentMiddleware[StateT, ContextT],
                 state: StateT,
                 runtime: Runtime[ContextT],
-            ) -> dict[str, Any] | Command | None:
+            ) -> dict[str, Any] | Command[Any] | None:
                 return await func(state, runtime)  # type: ignore[misc]
 
             # Preserve can_jump_to metadata on the wrapped function
@@ -725,10 +1007,10 @@ def before_model(
             )()
 
         def wrapped(
-            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            _self: AgentMiddleware[StateT, ContextT],
             state: StateT,
             runtime: Runtime[ContextT],
-        ) -> dict[str, Any] | Command | None:
+        ) -> dict[str, Any] | Command[Any] | None:
             return func(state, runtime)  # type: ignore[return-value]
 
         # Preserve can_jump_to metadata on the wrapped function
@@ -786,39 +1068,66 @@ def after_model(
     """Decorator used to dynamically create a middleware with the `after_model` hook.
 
     Args:
-        func: The function to be decorated. Must accept:
-            `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
-        state_schema: Optional custom state schema type. If not provided, uses the
-            default `AgentState` schema.
+        func: The function to be decorated.
+
+            Must accept: `state: StateT, runtime: Runtime[ContextT]` - State and runtime
+            context
+        state_schema: Optional custom state schema type.
+
+            If not provided, uses the default `AgentState` schema.
         tools: Optional list of additional tools to register with this middleware.
         can_jump_to: Optional list of valid jump destinations for conditional edges.
-            Valid values are: `"tools"`, `"model"`, `"end"`
-        name: Optional name for the generated middleware class. If not provided,
-            uses the decorated function's name.
+
+            Valid values are: `'tools'`, `'model'`, `'end'`
+        name: Optional name for the generated middleware class.
+
+            If not provided, uses the decorated function's name.
 
     Returns:
         Either an `AgentMiddleware` instance (if func is provided) or a decorator
-        function that can be applied to a function.
+            function that can be applied to a function.
 
     The decorated function should return:
-        - `dict[str, Any]` - State updates to merge into the agent state
-        - `Command` - A command to control flow (e.g., jump to different node)
-        - `None` - No state updates or flow control
+
+    - `dict[str, Any]` - State updates to merge into the agent state
+    - `Command` - A command to control flow (e.g., jump to different node)
+    - `None` - No state updates or flow control
 
     Examples:
-        Basic usage for logging model responses:
-        ```python
-        @after_model
-        def log_latest_message(state: AgentState, runtime: Runtime) -> None:
-            print(state["messages"][-1].content)
-        ```
+        !!! example "Basic usage for logging model responses"
 
-        With custom state schema:
-        ```python
-        @after_model(state_schema=MyCustomState, name="MyAfterModelMiddleware")
-        def custom_after_model(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
-            return {"custom_field": "updated_after_model"}
-        ```
+            ```python
+            @after_model
+            def log_latest_message(state: AgentState, runtime: Runtime) -> None:
+                print(state["messages"][-1].content)
+            ```
+
+        !!! example "With custom state schema"
+
+            ```python
+            @after_model(state_schema=MyCustomState, name="MyAfterModelMiddleware")
+            def custom_after_model(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
+                return {"custom_field": "updated_after_model"}
+            ```
+
+        !!! example "Streaming custom events after model call"
+
+            Use `runtime.stream_writer` to emit custom events after model responds.
+            Events are received when streaming with `stream_mode="custom"`.
+
+            ```python
+            @after_model
+            async def notify_model_response(state: AgentState, runtime: Runtime) -> None:
+                '''Notify user after model has responded.'''
+                last_message = state["messages"][-1]
+                has_tool_calls = hasattr(last_message, "tool_calls") and last_message.tool_calls
+                runtime.stream_writer(
+                    {
+                        "type": "status",
+                        "message": "Using tools..." if has_tool_calls else "Response ready!",
+                    }
+                )
+            ```
     """
 
     def decorator(
@@ -833,10 +1142,10 @@ def after_model(
         if is_async:
 
             async def async_wrapped(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+                _self: AgentMiddleware[StateT, ContextT],
                 state: StateT,
                 runtime: Runtime[ContextT],
-            ) -> dict[str, Any] | Command | None:
+            ) -> dict[str, Any] | Command[Any] | None:
                 return await func(state, runtime)  # type: ignore[misc]
 
             # Preserve can_jump_to metadata on the wrapped function
@@ -856,10 +1165,10 @@ def after_model(
             )()
 
         def wrapped(
-            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            _self: AgentMiddleware[StateT, ContextT],
             state: StateT,
             runtime: Runtime[ContextT],
-        ) -> dict[str, Any] | Command | None:
+        ) -> dict[str, Any] | Command[Any] | None:
             return func(state, runtime)  # type: ignore[return-value]
 
         # Preserve can_jump_to metadata on the wrapped function
@@ -917,48 +1226,99 @@ def before_agent(
     """Decorator used to dynamically create a middleware with the `before_agent` hook.
 
     Args:
-        func: The function to be decorated. Must accept:
-            `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
-        state_schema: Optional custom state schema type. If not provided, uses the
-            default `AgentState` schema.
+        func: The function to be decorated.
+
+            Must accept: `state: StateT, runtime: Runtime[ContextT]` - State and runtime
+            context
+        state_schema: Optional custom state schema type.
+
+            If not provided, uses the default `AgentState` schema.
         tools: Optional list of additional tools to register with this middleware.
         can_jump_to: Optional list of valid jump destinations for conditional edges.
-            Valid values are: `"tools"`, `"model"`, `"end"`
-        name: Optional name for the generated middleware class. If not provided,
-            uses the decorated function's name.
+
+            Valid values are: `'tools'`, `'model'`, `'end'`
+        name: Optional name for the generated middleware class.
+
+            If not provided, uses the decorated function's name.
 
     Returns:
         Either an `AgentMiddleware` instance (if func is provided directly) or a
-        decorator function that can be applied to a function it is wrapping.
+            decorator function that can be applied to a function it is wrapping.
 
     The decorated function should return:
-        - `dict[str, Any]` - State updates to merge into the agent state
-        - `Command` - A command to control flow (e.g., jump to different node)
-        - `None` - No state updates or flow control
+
+    - `dict[str, Any]` - State updates to merge into the agent state
+    - `Command` - A command to control flow (e.g., jump to different node)
+    - `None` - No state updates or flow control
 
     Examples:
-        Basic usage:
-        ```python
-        @before_agent
-        def log_before_agent(state: AgentState, runtime: Runtime) -> None:
-            print(f"Starting agent with {len(state['messages'])} messages")
-        ```
+        !!! example "Basic usage"
 
-        With conditional jumping:
-        ```python
-        @before_agent(can_jump_to=["end"])
-        def conditional_before_agent(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
-            if some_condition(state):
-                return {"jump_to": "end"}
-            return None
-        ```
+            ```python
+            @before_agent
+            def log_before_agent(state: AgentState, runtime: Runtime) -> None:
+                print(f"Starting agent with {len(state['messages'])} messages")
+            ```
 
-        With custom state schema:
-        ```python
-        @before_agent(state_schema=MyCustomState)
-        def custom_before_agent(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
-            return {"custom_field": "initialized_value"}
-        ```
+        !!! example "With conditional jumping"
+
+            ```python
+            @before_agent(can_jump_to=["end"])
+            def conditional_before_agent(
+                state: AgentState, runtime: Runtime
+            ) -> dict[str, Any] | None:
+                if some_condition(state):
+                    return {"jump_to": "end"}
+                return None
+            ```
+
+        !!! example "With custom state schema"
+
+            ```python
+            @before_agent(state_schema=MyCustomState)
+            def custom_before_agent(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
+                return {"custom_field": "initialized_value"}
+            ```
+
+        !!! example "Streaming custom events"
+
+            Use `runtime.stream_writer` to emit custom events during agent execution.
+            Events are received when streaming with `stream_mode="custom"`.
+
+            ```python
+            from langchain.agents import create_agent
+            from langchain.agents.middleware import before_agent, AgentState
+            from langchain.messages import HumanMessage
+            from langgraph.runtime import Runtime
+
+
+            @before_agent
+            async def notify_start(state: AgentState, runtime: Runtime) -> None:
+                '''Notify user that agent is starting.'''
+                runtime.stream_writer(
+                    {
+                        "type": "status",
+                        "message": "Initializing agent session...",
+                    }
+                )
+                # Perform prerequisite tasks here
+                runtime.stream_writer({"type": "status", "message": "Agent ready!"})
+
+
+            agent = create_agent(
+                model="openai:gpt-5.2",
+                tools=[...],
+                middleware=[notify_start],
+            )
+
+            # Consume with stream_mode="custom" to receive events
+            async for mode, event in agent.astream(
+                {"messages": [HumanMessage("Hello")]},
+                stream_mode=["updates", "custom"],
+            ):
+                if mode == "custom":
+                    print(f"Status: {event}")
+            ```
     """
 
     def decorator(
@@ -973,10 +1333,10 @@ def before_agent(
         if is_async:
 
             async def async_wrapped(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+                _self: AgentMiddleware[StateT, ContextT],
                 state: StateT,
                 runtime: Runtime[ContextT],
-            ) -> dict[str, Any] | Command | None:
+            ) -> dict[str, Any] | Command[Any] | None:
                 return await func(state, runtime)  # type: ignore[misc]
 
             # Preserve can_jump_to metadata on the wrapped function
@@ -998,10 +1358,10 @@ def before_agent(
             )()
 
         def wrapped(
-            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            _self: AgentMiddleware[StateT, ContextT],
             state: StateT,
             runtime: Runtime[ContextT],
-        ) -> dict[str, Any] | Command | None:
+        ) -> dict[str, Any] | Command[Any] | None:
             return func(state, runtime)  # type: ignore[return-value]
 
         # Preserve can_jump_to metadata on the wrapped function
@@ -1058,40 +1418,68 @@ def after_agent(
 ):
     """Decorator used to dynamically create a middleware with the `after_agent` hook.
 
+    Async version is `aafter_agent`.
+
     Args:
-        func: The function to be decorated. Must accept:
-            `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
-        state_schema: Optional custom state schema type. If not provided, uses the
-            default `AgentState` schema.
+        func: The function to be decorated.
+
+            Must accept: `state: StateT, runtime: Runtime[ContextT]` - State and runtime
+            context
+        state_schema: Optional custom state schema type.
+
+            If not provided, uses the default `AgentState` schema.
         tools: Optional list of additional tools to register with this middleware.
         can_jump_to: Optional list of valid jump destinations for conditional edges.
-            Valid values are: `"tools"`, `"model"`, `"end"`
-        name: Optional name for the generated middleware class. If not provided,
-            uses the decorated function's name.
+
+            Valid values are: `'tools'`, `'model'`, `'end'`
+        name: Optional name for the generated middleware class.
+
+            If not provided, uses the decorated function's name.
 
     Returns:
         Either an `AgentMiddleware` instance (if func is provided) or a decorator
-        function that can be applied to a function.
+            function that can be applied to a function.
 
     The decorated function should return:
-        - `dict[str, Any]` - State updates to merge into the agent state
-        - `Command` - A command to control flow (e.g., jump to different node)
-        - `None` - No state updates or flow control
+
+    - `dict[str, Any]` - State updates to merge into the agent state
+    - `Command` - A command to control flow (e.g., jump to different node)
+    - `None` - No state updates or flow control
 
     Examples:
-        Basic usage for logging agent completion:
-        ```python
-        @after_agent
-        def log_completion(state: AgentState, runtime: Runtime) -> None:
-            print(f"Agent completed with {len(state['messages'])} messages")
-        ```
+        !!! example "Basic usage for logging agent completion"
 
-        With custom state schema:
-        ```python
-        @after_agent(state_schema=MyCustomState, name="MyAfterAgentMiddleware")
-        def custom_after_agent(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
-            return {"custom_field": "finalized_value"}
-        ```
+            ```python
+            @after_agent
+            def log_completion(state: AgentState, runtime: Runtime) -> None:
+                print(f"Agent completed with {len(state['messages'])} messages")
+            ```
+
+        !!! example "With custom state schema"
+
+            ```python
+            @after_agent(state_schema=MyCustomState, name="MyAfterAgentMiddleware")
+            def custom_after_agent(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
+                return {"custom_field": "finalized_value"}
+            ```
+
+        !!! example "Streaming custom events on completion"
+
+            Use `runtime.stream_writer` to emit custom events when agent completes.
+            Events are received when streaming with `stream_mode="custom"`.
+
+            ```python
+            @after_agent
+            async def notify_completion(state: AgentState, runtime: Runtime) -> None:
+                '''Notify user that agent has completed.'''
+                runtime.stream_writer(
+                    {
+                        "type": "status",
+                        "message": "Agent execution complete!",
+                        "total_messages": len(state["messages"]),
+                    }
+                )
+            ```
     """
 
     def decorator(
@@ -1106,10 +1494,10 @@ def after_agent(
         if is_async:
 
             async def async_wrapped(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+                _self: AgentMiddleware[StateT, ContextT],
                 state: StateT,
                 runtime: Runtime[ContextT],
-            ) -> dict[str, Any] | Command | None:
+            ) -> dict[str, Any] | Command[Any] | None:
                 return await func(state, runtime)  # type: ignore[misc]
 
             # Preserve can_jump_to metadata on the wrapped function
@@ -1129,10 +1517,10 @@ def after_agent(
             )()
 
         def wrapped(
-            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            _self: AgentMiddleware[StateT, ContextT],
             state: StateT,
             runtime: Runtime[ContextT],
-        ) -> dict[str, Any] | Command | None:
+        ) -> dict[str, Any] | Command[Any] | None:
             return func(state, runtime)  # type: ignore[return-value]
 
         # Preserve can_jump_to metadata on the wrapped function
@@ -1159,7 +1547,7 @@ def after_agent(
 
 @overload
 def dynamic_prompt(
-    func: _CallableReturningPromptString[StateT, ContextT],
+    func: _CallableReturningSystemMessage[StateT, ContextT],
 ) -> AgentMiddleware[StateT, ContextT]: ...
 
 
@@ -1167,16 +1555,16 @@ def dynamic_prompt(
 def dynamic_prompt(
     func: None = None,
 ) -> Callable[
-    [_CallableReturningPromptString[StateT, ContextT]],
+    [_CallableReturningSystemMessage[StateT, ContextT]],
     AgentMiddleware[StateT, ContextT],
 ]: ...
 
 
 def dynamic_prompt(
-    func: _CallableReturningPromptString[StateT, ContextT] | None = None,
+    func: _CallableReturningSystemMessage[StateT, ContextT] | None = None,
 ) -> (
     Callable[
-        [_CallableReturningPromptString[StateT, ContextT]],
+        [_CallableReturningSystemMessage[StateT, ContextT]],
         AgentMiddleware[StateT, ContextT],
     ]
     | AgentMiddleware[StateT, ContextT]
@@ -1188,18 +1576,22 @@ def dynamic_prompt(
     a string that will be set as the system prompt for the model request.
 
     Args:
-        func: The function to be decorated. Must accept:
-            `request: ModelRequest` - Model request (contains state and runtime)
+        func: The function to be decorated.
+
+            Must accept: `request: ModelRequest` - Model request (contains state and
+            runtime)
 
     Returns:
-        Either an AgentMiddleware instance (if func is provided) or a decorator function
-        that can be applied to a function.
+        Either an `AgentMiddleware` instance (if func is provided) or a decorator
+            function that can be applied to a function.
 
     The decorated function should return:
-        - `str` - The system prompt to use for the model request
+        - `str` – The system prompt string to use for the model request
+        - `SystemMessage` – A complete system message to use for the model request
 
     Examples:
         Basic usage with dynamic content:
+
         ```python
         @dynamic_prompt
         def my_prompt(request: ModelRequest) -> str:
@@ -1208,6 +1600,7 @@ def dynamic_prompt(
         ```
 
         Using state to customize the prompt:
+
         ```python
         @dynamic_prompt
         def context_aware_prompt(request: ModelRequest) -> str:
@@ -1218,25 +1611,29 @@ def dynamic_prompt(
         ```
 
         Using with agent:
+
         ```python
         agent = create_agent(model, middleware=[my_prompt])
         ```
     """
 
     def decorator(
-        func: _CallableReturningPromptString[StateT, ContextT],
+        func: _CallableReturningSystemMessage[StateT, ContextT],
     ) -> AgentMiddleware[StateT, ContextT]:
         is_async = iscoroutinefunction(func)
 
         if is_async:
 
             async def async_wrapped(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+                _self: AgentMiddleware[StateT, ContextT],
                 request: ModelRequest,
                 handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
             ) -> ModelCallResult:
                 prompt = await func(request)  # type: ignore[misc]
-                request.system_prompt = prompt
+                if isinstance(prompt, SystemMessage):
+                    request = request.override(system_message=prompt)
+                else:
+                    request = request.override(system_message=SystemMessage(content=prompt))
                 return await handler(request)
 
             middleware_name = cast("str", getattr(func, "__name__", "DynamicPromptMiddleware"))
@@ -1252,22 +1649,28 @@ def dynamic_prompt(
             )()
 
         def wrapped(
-            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            _self: AgentMiddleware[StateT, ContextT],
             request: ModelRequest,
             handler: Callable[[ModelRequest], ModelResponse],
         ) -> ModelCallResult:
-            prompt = cast("str", func(request))
-            request.system_prompt = prompt
+            prompt = cast("Callable[[ModelRequest], SystemMessage | str]", func)(request)
+            if isinstance(prompt, SystemMessage):
+                request = request.override(system_message=prompt)
+            else:
+                request = request.override(system_message=SystemMessage(content=prompt))
             return handler(request)
 
         async def async_wrapped_from_sync(
-            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            _self: AgentMiddleware[StateT, ContextT],
             request: ModelRequest,
             handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
         ) -> ModelCallResult:
             # Delegate to sync function
-            prompt = cast("str", func(request))
-            request.system_prompt = prompt
+            prompt = cast("Callable[[ModelRequest], SystemMessage | str]", func)(request)
+            if isinstance(prompt, SystemMessage):
+                request = request.override(system_message=prompt)
+            else:
+                request = request.override(system_message=SystemMessage(content=prompt))
             return await handler(request)
 
         middleware_name = cast("str", getattr(func, "__name__", "DynamicPromptMiddleware"))
@@ -1322,68 +1725,77 @@ def wrap_model_call(
 ):
     """Create middleware with `wrap_model_call` hook from a function.
 
-    Converts a function with handler callback into middleware that can intercept
-    model calls, implement retry logic, handle errors, and rewrite responses.
+    Converts a function with handler callback into middleware that can intercept model
+    calls, implement retry logic, handle errors, and rewrite responses.
 
     Args:
         func: Function accepting (request, handler) that calls handler(request)
             to execute the model and returns `ModelResponse` or `AIMessage`.
+
             Request contains state and runtime.
-        state_schema: Custom state schema. Defaults to `AgentState`.
+        state_schema: Custom state schema.
+
+            Defaults to `AgentState`.
         tools: Additional tools to register with this middleware.
-        name: Middleware class name. Defaults to function name.
+        name: Middleware class name.
+
+            Defaults to function name.
 
     Returns:
         `AgentMiddleware` instance if func provided, otherwise a decorator.
 
     Examples:
-        Basic retry logic:
-        ```python
-        @wrap_model_call
-        def retry_on_error(request, handler):
-            max_retries = 3
-            for attempt in range(max_retries):
+        !!! example "Basic retry logic"
+
+            ```python
+            @wrap_model_call
+            def retry_on_error(request, handler):
+                max_retries = 3
+                for attempt in range(max_retries):
+                    try:
+                        return handler(request)
+                    except Exception:
+                        if attempt == max_retries - 1:
+                            raise
+            ```
+
+        !!! example "Model fallback"
+
+            ```python
+            @wrap_model_call
+            def fallback_model(request, handler):
+                # Try primary model
                 try:
                     return handler(request)
                 except Exception:
-                    if attempt == max_retries - 1:
-                        raise
-        ```
+                    pass
 
-        Model fallback:
-        ```python
-        @wrap_model_call
-        def fallback_model(request, handler):
-            # Try primary model
-            try:
+                # Try fallback model
+                request = request.override(model=fallback_model_instance)
                 return handler(request)
-            except Exception:
-                pass
+            ```
 
-            # Try fallback model
-            request.model = fallback_model_instance
-            return handler(request)
-        ```
+        !!! example "Rewrite response content (full `ModelResponse`)"
 
-        Rewrite response content (full ModelResponse):
-        ```python
-        @wrap_model_call
-        def uppercase_responses(request, handler):
-            response = handler(request)
-            ai_msg = response.result[0]
-            return ModelResponse(
-                result=[AIMessage(content=ai_msg.content.upper())],
-                structured_response=response.structured_response,
-            )
-        ```
+            ```python
+            @wrap_model_call
+            def uppercase_responses(request, handler):
+                response = handler(request)
+                ai_msg = response.result[0]
+                return ModelResponse(
+                    result=[AIMessage(content=ai_msg.content.upper())],
+                    structured_response=response.structured_response,
+                )
+            ```
 
-        Simple AIMessage return (converted automatically):
-        ```python
-        @wrap_model_call
-        def simple_response(request, handler):
-            # AIMessage is automatically converted to ModelResponse
-            return AIMessage(content="Simple response")
-        ```
+        !!! example "Simple `AIMessage` return (converted automatically)"
+
+            ```python
+            @wrap_model_call
+            def simple_response(request, handler):
+                # AIMessage is automatically converted to ModelResponse
+                return AIMessage(content="Simple response")
+            ```
     """
 
     def decorator(
@@ -1394,7 +1806,7 @@ def wrap_model_call(
         if is_async:
 
             async def async_wrapped(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+                _self: AgentMiddleware[StateT, ContextT],
                 request: ModelRequest,
                 handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
             ) -> ModelCallResult:
@@ -1415,7 +1827,7 @@ def wrap_model_call(
             )()
 
         def wrapped(
-            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            _self: AgentMiddleware[StateT, ContextT],
             request: ModelRequest,
             handler: Callable[[ModelRequest], ModelResponse],
         ) -> ModelCallResult:
@@ -1470,63 +1882,80 @@ def wrap_tool_call(
 ):
     """Create middleware with `wrap_tool_call` hook from a function.
 
+    Async version is `awrap_tool_call`.
+
     Converts a function with handler callback into middleware that can intercept
     tool calls, implement retry logic, monitor execution, and modify responses.
 
     Args:
         func: Function accepting (request, handler) that calls
             handler(request) to execute the tool and returns final `ToolMessage` or
-            `Command`. Can be sync or async.
+            `Command`.
+
+            Can be sync or async.
         tools: Additional tools to register with this middleware.
-        name: Middleware class name. Defaults to function name.
+        name: Middleware class name.
+
+            Defaults to function name.
 
     Returns:
         `AgentMiddleware` instance if func provided, otherwise a decorator.
 
     Examples:
-        Retry logic:
-        ```python
-        @wrap_tool_call
-        def retry_on_error(request, handler):
-            max_retries = 3
-            for attempt in range(max_retries):
-                try:
-                    return handler(request)
-                except Exception:
-                    if attempt == max_retries - 1:
-                        raise
-        ```
+        !!! example "Retry logic"
 
-        Async retry logic:
-        ```python
-        @wrap_tool_call
-        async def async_retry(request, handler):
-            for attempt in range(3):
-                try:
-                    return await handler(request)
-                except Exception:
-                    if attempt == 2:
-                        raise
-        ```
+            ```python
+            @wrap_tool_call
+            def retry_on_error(request, handler):
+                max_retries = 3
+                for attempt in range(max_retries):
+                    try:
+                        return handler(request)
+                    except Exception:
+                        if attempt == max_retries - 1:
+                            raise
+            ```
 
-        Modify request:
-        ```python
-        @wrap_tool_call
-        def modify_args(request, handler):
-            request.tool_call["args"]["value"] *= 2
-            return handler(request)
-        ```
+        !!! example "Async retry logic"
 
-        Short-circuit with cached result:
-        ```python
-        @wrap_tool_call
-        def with_cache(request, handler):
-            if cached := get_cache(request):
-                return ToolMessage(content=cached, tool_call_id=request.tool_call["id"])
-            result = handler(request)
-            save_cache(request, result)
-            return result
-        ```
+            ```python
+            @wrap_tool_call
+            async def async_retry(request, handler):
+                for attempt in range(3):
+                    try:
+                        return await handler(request)
+                    except Exception:
+                        if attempt == 2:
+                            raise
+            ```
+
+        !!! example "Modify request"
+
+            ```python
+            @wrap_tool_call
+            def modify_args(request, handler):
+                modified_call = {
+                    **request.tool_call,
+                    "args": {
+                        **request.tool_call["args"],
+                        "value": request.tool_call["args"]["value"] * 2,
+                    },
+                }
+                request = request.override(tool_call=modified_call)
+                return handler(request)
+            ```
+
+        !!! example "Short-circuit with cached result"
+
+            ```python
+            @wrap_tool_call
+            def with_cache(request, handler):
+                if cached := get_cache(request):
+                    return ToolMessage(content=cached, tool_call_id=request.tool_call["id"])
+                result = handler(request)
+                save_cache(request, result)
+                return result
+            ```
     """
 
     def decorator(
@@ -1537,10 +1966,10 @@ def wrap_tool_call(
         if is_async:
 
             async def async_wrapped(
-                self: AgentMiddleware,  # noqa: ARG001
+                _self: AgentMiddleware,
                 request: ToolCallRequest,
-                handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
-            ) -> ToolMessage | Command:
+                handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],
+            ) -> ToolMessage | Command[Any]:
                 return await func(request, handler)  # type: ignore[arg-type,misc]
 
             middleware_name = name or cast(
@@ -1558,10 +1987,10 @@ def wrap_tool_call(
             )()
 
         def wrapped(
-            self: AgentMiddleware,  # noqa: ARG001
+            _self: AgentMiddleware,
             request: ToolCallRequest,
-            handler: Callable[[ToolCallRequest], ToolMessage | Command],
-        ) -> ToolMessage | Command:
+            handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
+        ) -> ToolMessage | Command[Any]:
             return func(request, handler)
 
         middleware_name = name or cast("str", getattr(func, "__name__", "WrapToolCallMiddleware"))