Python: Agent and Function middleware (#770)

* Initial middleware implementation * Small fixes * Small updates * Small updates in samples * Moved middleware functionality to decorator * Removed obsolete file * Renamed AgentInvocationContext to AzureRunContext * Added unit tests * Small settings update for test discovery in VS Code * Added unit tests * Reverted changes in environment settings * Added context result override * Renaming and updates to logic * Added more samples * Updated DEV_SETUP.md * Addressed PR feedback * Addressed PR feedback * Removed unused parameter * Small fix * Small fix in telemetry logic * Revert "Small fix in telemetry logic" This reverts commit 6f82660d2d. * Small fix --------- Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2025-09-18 16:30:05 -07:00
parent 538be4c149
commit 99860a5d07
16 changed files with 3071 additions and 40 deletions
@@ -13,6 +13,7 @@ from ._clients import *  # noqa: F403
 from ._logging import *  # noqa: F403
 from ._mcp import *  # noqa: F403
 from ._memory import *  # noqa: F403
+from ._middleware import *  # noqa: F403
 from ._threads import *  # noqa: F403
 from ._tools import *  # noqa: F403
 from ._types import *  # noqa: F403
@@ -12,6 +12,7 @@ from ._clients import BaseChatClient, ChatClientProtocol
 from ._logging import get_logger
 from ._mcp import MCPTool
 from ._memory import AggregateContextProvider, Context, ContextProvider
+from ._middleware import Middleware, use_agent_middleware
 from ._pydantic import AFBaseModel
 from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages
 from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, ToolProtocol
@@ -138,12 +139,14 @@ class BaseAgent(AFBaseModel):
       description: The description of the agent.
       display_name: The display name of the agent, which is either the name or id.
       context_providers: The collection of multiple context providers to include during agent invocation.
+       middleware: List of middleware to intercept agent and function invocations.
    """

    id: str = Field(default_factory=lambda: str(uuid4()))
    name: str | None = None
    description: str | None = None
    context_providers: AggregateContextProvider | None = None
+    middleware: Middleware | list[Middleware] | None = None

    async def _notify_thread_of_new_messages(
        self, thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage]
@@ -189,6 +192,7 @@ class BaseAgent(AFBaseModel):
 # region ChatAgent


+@use_agent_middleware
@use_agent_telemetry
 class ChatAgent(BaseAgent):
    """A Chat Client Agent."""
@@ -231,6 +235,7 @@ class ChatAgent(BaseAgent):
        additional_properties: dict[str, Any] | None = None,
        chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
+        middleware: Middleware | list[Middleware] | None = None,
        **kwargs: Any,
    ) -> None:
        """Create a ChatAgent.
@@ -266,6 +271,7 @@ class ChatAgent(BaseAgent):
            chat_message_store_factory: factory function to create an instance of ChatMessageStore. If not provided,
                the default in-memory store will be used.
            context_providers: The collection of multiple context providers to include during agent invocation.
+            middleware: List of middleware to intercept agent and function invocations.
            kwargs: any additional keyword arguments.
                Unused, can be used by subclasses of this Agent.
        """
@@ -287,6 +293,7 @@ class ChatAgent(BaseAgent):
            "chat_client": chat_client,
            "chat_message_store_factory": chat_message_store_factory,
            "context_providers": aggregate_context_providers,
+            "middleware": middleware,
            "chat_options": ChatOptions(
                ai_model_id=model,
                frequency_penalty=frequency_penalty,
@@ -10,6 +10,7 @@ from pydantic import BaseModel, Field
 from ._logging import get_logger
 from ._mcp import MCPTool
 from ._memory import AggregateContextProvider, ContextProvider
+from ._middleware import Middleware
 from ._pydantic import AFBaseModel
 from ._threads import ChatMessageStore
 from ._tools import ToolProtocol
@@ -344,7 +345,14 @@ class BaseChatClient(AFBaseModel, ABC):
            )
        prepped_messages = self.prepare_messages(messages)
        self._prepare_tool_choice(chat_options=chat_options)
-        return await self._inner_get_response(messages=prepped_messages, chat_options=chat_options, **kwargs)
+
+        # Remove middleware pipeline from kwargs as it's only used by function invocation wrappers
+        if "_function_middleware_pipeline" in kwargs:
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k != "_function_middleware_pipeline"}
+        else:
+            filtered_kwargs = kwargs
+
+        return await self._inner_get_response(messages=prepped_messages, chat_options=chat_options, **filtered_kwargs)

    async def get_streaming_response(
        self,
@@ -424,8 +432,15 @@ class BaseChatClient(AFBaseModel, ABC):
            )
        prepped_messages = self.prepare_messages(messages)
        self._prepare_tool_choice(chat_options=chat_options)
+
+        # Remove middleware pipeline from kwargs as it's only used by function invocation wrappers
+        if "_function_middleware_pipeline" in kwargs:
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k != "_function_middleware_pipeline"}
+        else:
+            filtered_kwargs = kwargs
+
        async for update in self._inner_get_streaming_response(
-            messages=prepped_messages, chat_options=chat_options, **kwargs
+            messages=prepped_messages, chat_options=chat_options, **filtered_kwargs
        ):
            yield update

@@ -465,6 +480,7 @@ class BaseChatClient(AFBaseModel, ABC):
        | None = None,
        chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
+        middleware: Middleware | list[Middleware] | None = None,
        **kwargs: Any,
    ) -> "ChatAgent":
        """Create an agent with the given name and instructions.
@@ -476,6 +492,7 @@ class BaseChatClient(AFBaseModel, ABC):
            chat_message_store_factory: Factory function to create an instance of ChatMessageStore. If not provided,
                the default in-memory store will be used.
            context_providers: Context providers to include during agent invocation.
+            middleware: List of middleware to intercept agent and function invocations.
            **kwargs: Additional keyword arguments to pass to the agent.
                See ChatAgent for all the available options.

@@ -491,6 +508,7 @@ class BaseChatClient(AFBaseModel, ABC):
            tools=tools,
            chat_message_store_factory=chat_message_store_factory,
            context_providers=context_providers,
+            middleware=middleware,
            **kwargs,
        )

@@ -0,0 +1,632 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterable, Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar
+
+from ._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+    from ._agents import AgentProtocol
+    from ._tools import AIFunction
+
+TAgent = TypeVar("TAgent", bound="AgentProtocol")
+
+__all__ = [
+    "AgentMiddleware",
+    "AgentRunContext",
+    "FunctionInvocationContext",
+    "FunctionMiddleware",
+    "Middleware",
+    "use_agent_middleware",
+]
+
+
+@dataclass
+class AgentRunContext:
+    """Context object for agent middleware invocations.
+
+    Attributes:
+        agent: The agent being invoked.
+        messages: The messages being sent to the agent.
+        is_streaming: Whether this is a streaming invocation.
+        metadata: Metadata dictionary for sharing data between agent middleware.
+        result: Agent execution result. Can be observed after calling next()
+                to see the actual execution result or can be set to override the execution result.
+                For non-streaming: should be AgentRunResponse
+                For streaming: should be AsyncIterable[AgentRunResponseUpdate]
+    """
+
+    agent: "AgentProtocol"
+    messages: list[ChatMessage]
+    is_streaming: bool = False
+    metadata: dict[str, Any] = field(default_factory=lambda: {})
+    result: AgentRunResponse | AsyncIterable[AgentRunResponseUpdate] | None = None
+
+
+@dataclass
+class FunctionInvocationContext:
+    """Context object for function middleware invocations.
+
+    Attributes:
+        function: The function being invoked.
+        arguments: The validated arguments for the function.
+        metadata: Metadata dictionary for sharing data between function middleware.
+        result: Function execution result. Can be observed after calling next()
+                to see the actual execution result or can be set to override the execution result.
+    """
+
+    function: "AIFunction[Any, Any]"
+    arguments: "BaseModel"
+    metadata: dict[str, Any] = field(default_factory=lambda: {})
+    result: Any = None
+
+
+class AgentMiddleware(ABC):
+    """Abstract base class for agent middleware that can intercept agent invocations."""
+
+    @abstractmethod
+    async def process(
+        self,
+        context: AgentRunContext,
+        next: Callable[[AgentRunContext], Awaitable[None]],
+    ) -> None:
+        """Process an agent invocation.
+
+        Args:
+            context: Agent invocation context containing agent, messages, and metadata.
+                    Use context.is_streaming to determine if this is a streaming call.
+                    Middleware can set context.result to override execution, or observe
+                    the actual execution result after calling next().
+                    For non-streaming: AgentRunResponse
+                    For streaming: AsyncIterable[AgentRunResponseUpdate]
+            next: Function to call the next middleware or final agent execution.
+                  Does not return anything - all data flows through the context.
+
+        Note:
+            Middleware should not return anything. All data manipulation should happen
+            within the context object. Set context.result to override execution,
+            or observe context.result after calling next() for actual results.
+        """
+        ...
+
+
+class FunctionMiddleware(ABC):
+    """Abstract base class for function middleware that can intercept function invocations."""
+
+    @abstractmethod
+    async def process(
+        self,
+        context: FunctionInvocationContext,
+        next: Callable[[FunctionInvocationContext], Awaitable[None]],
+    ) -> None:
+        """Process a function invocation.
+
+        Args:
+            context: Function invocation context containing function, arguments, and metadata.
+                    Middleware can set context.result to override execution, or observe
+                    the actual execution result after calling next().
+            next: Function to call the next middleware or final function execution.
+                  Does not return anything - all data flows through the context.
+
+        Note:
+            Middleware should not return anything. All data manipulation should happen
+            within the context object. Set context.result to override execution,
+            or observe context.result after calling next() for actual results.
+        """
+        ...
+
+
+# Pure function type definitions for convenience
+AgentMiddlewareCallable = Callable[[AgentRunContext, Callable[[AgentRunContext], Awaitable[None]]], Awaitable[None]]
+
+FunctionMiddlewareCallable = Callable[
+    [FunctionInvocationContext, Callable[[FunctionInvocationContext], Awaitable[None]]], Awaitable[None]
+]
+
+# Type alias for all middleware types
+Middleware: TypeAlias = AgentMiddleware | AgentMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable
+
+
+class AgentMiddlewareWrapper(AgentMiddleware):
+    """Wrapper to convert pure functions into AgentMiddleware protocol objects."""
+
+    def __init__(self, func: AgentMiddlewareCallable):
+        self.func = func
+
+    async def process(
+        self,
+        context: AgentRunContext,
+        next: Callable[[AgentRunContext], Awaitable[None]],
+    ) -> None:
+        await self.func(context, next)
+
+
+class FunctionMiddlewareWrapper(FunctionMiddleware):
+    """Wrapper to convert pure functions into FunctionMiddleware protocol objects."""
+
+    def __init__(self, func: FunctionMiddlewareCallable):
+        self.func = func
+
+    async def process(
+        self,
+        context: FunctionInvocationContext,
+        next: Callable[[FunctionInvocationContext], Awaitable[None]],
+    ) -> None:
+        await self.func(context, next)
+
+
+class BaseMiddlewarePipeline(ABC):
+    """Base class for middleware pipeline execution."""
+
+    def __init__(self) -> None:
+        """Initialize the base middleware pipeline."""
+        self._middlewares: list[Any] = []
+
+    @abstractmethod
+    def _register_middleware(self, middleware: Any) -> None:
+        """Register a middleware item. Must be implemented by subclasses."""
+        ...
+
+    @property
+    def has_middlewares(self) -> bool:
+        """Check if there are any middlewares registered."""
+        return bool(self._middlewares)
+
+    def _create_handler_chain(
+        self,
+        final_handler: Callable[[Any], Awaitable[Any]],
+        result_container: dict[str, Any],
+        result_key: str = "result",
+    ) -> Callable[[Any], Awaitable[None]]:
+        """Create a chain of middleware handlers.
+
+        Args:
+            final_handler: The final handler to execute
+            result_container: Container to store the result
+            result_key: Key to use in the result container
+
+        Returns:
+            The first handler in the chain
+        """
+
+        def create_next_handler(index: int) -> Callable[[Any], Awaitable[None]]:
+            if index >= len(self._middlewares):
+
+                async def final_wrapper(c: Any) -> None:
+                    # Execute actual handler and populate context for observability
+                    result = await final_handler(c)
+                    result_container[result_key] = result
+                    c.result = result
+
+                return final_wrapper
+
+            middleware = self._middlewares[index]
+            next_handler = create_next_handler(index + 1)
+
+            async def current_handler(c: Any) -> None:
+                await middleware.process(c, next_handler)
+
+            return current_handler
+
+        return create_next_handler(0)
+
+
+class AgentMiddlewarePipeline(BaseMiddlewarePipeline):
+    """Executes agent middleware in a chain."""
+
+    def __init__(self, middlewares: list[AgentMiddleware | AgentMiddlewareCallable] | None = None):
+        """Initialize the agent middleware pipeline.
+
+        Args:
+            middlewares: List of agent middleware to include in the pipeline.
+        """
+        super().__init__()
+        self._middlewares: list[AgentMiddleware] = []
+
+        if middlewares:
+            for middleware in middlewares:
+                self._register_middleware(middleware)
+
+    def _register_middleware(self, middleware: AgentMiddleware | AgentMiddlewareCallable) -> None:
+        """Register an agent middleware item."""
+        if isinstance(middleware, AgentMiddleware):
+            self._middlewares.append(middleware)
+        elif callable(middleware):
+            self._middlewares.append(AgentMiddlewareWrapper(middleware))
+
+    async def execute(
+        self,
+        agent: "AgentProtocol",
+        messages: list[ChatMessage],
+        context: AgentRunContext,
+        final_handler: Callable[[AgentRunContext], Awaitable[AgentRunResponse]],
+    ) -> AgentRunResponse | None:
+        """Execute the agent middleware pipeline for non-streaming.
+
+        Args:
+            agent: The agent being invoked.
+            messages: The messages to send to the agent.
+            context: The agent invocation context.
+            final_handler: The final handler that performs the actual agent execution.
+
+        Returns:
+            The agent response after processing through all middleware.
+        """
+        # Update context with agent and messages
+        context.agent = agent
+        context.messages = messages
+        context.is_streaming = False
+
+        if not self._middlewares:
+            return await final_handler(context)
+
+        # Store the final result
+        result_container: dict[str, AgentRunResponse | None] = {"response": None}
+
+        def create_next_handler(index: int) -> Callable[[AgentRunContext], Awaitable[None]]:
+            if index >= len(self._middlewares):
+
+                async def final_wrapper(c: AgentRunContext) -> None:
+                    # Execute actual handler and populate context for observability
+                    result = await final_handler(c)
+                    result_container["result"] = result
+                    c.result = result
+
+                return final_wrapper
+
+            middleware = self._middlewares[index]
+            next_handler = create_next_handler(index + 1)
+
+            async def current_handler(c: AgentRunContext) -> None:
+                await middleware.process(c, next_handler)
+                # After middleware execution, check if response was overridden
+                if c.result is not None and isinstance(c.result, AgentRunResponse):
+                    result_container["result"] = c.result
+
+            return current_handler
+
+        first_handler = create_next_handler(0)
+        await first_handler(context)
+
+        # Return the result from result container or overridden result
+        if context.result is not None and isinstance(context.result, AgentRunResponse):
+            return context.result
+
+        # If no result was set (next() not called), return empty AgentRunResponse
+        response = result_container.get("result")
+        if response is None:
+            return AgentRunResponse()
+        return response
+
+    async def execute_stream(
+        self,
+        agent: "AgentProtocol",
+        messages: list[ChatMessage],
+        context: AgentRunContext,
+        final_handler: Callable[[AgentRunContext], AsyncIterable[AgentRunResponseUpdate]],
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Execute the agent middleware pipeline for streaming.
+
+        Args:
+            agent: The agent being invoked.
+            messages: The messages to send to the agent.
+            context: The agent invocation context.
+            final_handler: The final handler that performs the actual agent streaming execution.
+
+        Yields:
+            Agent response updates after processing through all middleware.
+        """
+        # Update context with agent and messages
+        context.agent = agent
+        context.messages = messages
+        context.is_streaming = True
+
+        if not self._middlewares:
+            async for update in final_handler(context):
+                yield update
+            return
+
+        # Store the final result
+        result_container: dict[str, AsyncIterable[AgentRunResponseUpdate] | None] = {"result_stream": None}
+
+        def create_next_handler(index: int) -> Callable[[AgentRunContext], Awaitable[None]]:
+            if index >= len(self._middlewares):
+
+                async def final_wrapper(c: AgentRunContext) -> None:  # noqa: RUF029
+                    # Execute actual handler and populate context for observability
+                    result = final_handler(c)
+                    result_container["result_stream"] = result
+                    c.result = result
+
+                return final_wrapper
+
+            middleware = self._middlewares[index]
+            next_handler = create_next_handler(index + 1)
+
+            async def current_handler(c: AgentRunContext) -> None:
+                await middleware.process(c, next_handler)
+
+            return current_handler
+
+        first_handler = create_next_handler(0)
+        await first_handler(context)
+
+        # Yield from the result stream in result container or overridden result
+        if context.result is not None and hasattr(context.result, "__aiter__"):
+            async for update in context.result:  # type: ignore
+                yield update
+            return
+
+        result_stream = result_container["result_stream"]
+        if result_stream is None:
+            # If no result stream was set (next() not called), yield nothing
+            return
+
+        async for update in result_stream:
+            yield update
+
+
+class FunctionMiddlewarePipeline(BaseMiddlewarePipeline):
+    """Executes function middleware in a chain."""
+
+    def __init__(self, middlewares: list[FunctionMiddleware | FunctionMiddlewareCallable] | None = None):
+        """Initialize the function middleware pipeline.
+
+        Args:
+            middlewares: List of function middleware to include in the pipeline.
+        """
+        super().__init__()
+        self._middlewares: list[FunctionMiddleware] = []
+
+        if middlewares:
+            for middleware in middlewares:
+                self._register_middleware(middleware)
+
+    def _register_middleware(self, middleware: FunctionMiddleware | FunctionMiddlewareCallable) -> None:
+        """Register a function middleware item."""
+        # Check if it's a class instance inheriting from FunctionMiddleware
+        if isinstance(middleware, FunctionMiddleware):
+            self._middlewares.append(middleware)
+        elif callable(middleware):
+            self._middlewares.append(FunctionMiddlewareWrapper(middleware))
+
+    async def execute(
+        self,
+        function: Any,
+        arguments: "BaseModel",
+        context: FunctionInvocationContext,
+        final_handler: Callable[[FunctionInvocationContext], Awaitable[Any]],
+    ) -> Any:
+        """Execute the function middleware pipeline.
+
+        Args:
+            function: The function being invoked.
+            arguments: The validated arguments for the function.
+            context: The function invocation context.
+            final_handler: The final handler that performs the actual function execution.
+
+        Returns:
+            The function result after processing through all middleware.
+        """
+        # Update context with function and arguments
+        context.function = function
+        context.arguments = arguments
+
+        if not self._middlewares:
+            return await final_handler(context)
+
+        # Store the final result
+        result_container: dict[str, Any] = {"result": None}
+
+        # Custom final handler that handles pre-existing results
+        async def function_final_handler(c: FunctionInvocationContext) -> Any:
+            # If result was set before calling next(), skip execution
+            if c.result is not None:
+                return c.result
+            # Execute actual handler and populate context for observability
+            return await final_handler(c)
+
+        first_handler = self._create_handler_chain(function_final_handler, result_container, "result")
+        await first_handler(context)
+
+        # Return the result from result container or overridden result
+        if context.result is not None:
+            return context.result
+        return result_container["result"]
+
+
+# Decorator for adding middleware support to agent classes
+def use_agent_middleware(agent_class: type[TAgent]) -> type[TAgent]:
+    """Class decorator that adds middleware support to an agent class.
+
+    This decorator adds middleware functionality to any agent class.
+    It wraps the run() and run_stream() methods to provide middleware execution.
+
+    Args:
+        agent_class: The agent class to add middleware support to.
+
+    Returns:
+        The modified agent class with middleware support.
+    """
+    import inspect
+
+    # Store original methods
+    original_run = agent_class.run  # type: ignore[attr-defined]
+    original_run_stream = agent_class.run_stream  # type: ignore[attr-defined]
+
+    def _initialize_middleware_pipelines(self: Any, middlewares: Middleware | list[Middleware] | None) -> None:
+        """Initialize agent and function middleware pipelines from the provided middleware list."""
+        if not middlewares:
+            return
+
+        middleware_list: list[Middleware] = middlewares if isinstance(middlewares, list) else [middlewares]  # type: ignore
+
+        # Separate agent and function middleware using isinstance checks
+        agent_middlewares: list[AgentMiddleware | AgentMiddlewareCallable] = []
+        function_middlewares: list[FunctionMiddleware | FunctionMiddlewareCallable] = []
+
+        for middleware in middleware_list:
+            if isinstance(middleware, AgentMiddleware):
+                agent_middlewares.append(middleware)
+            elif isinstance(middleware, FunctionMiddleware):
+                function_middlewares.append(middleware)
+            elif callable(middleware):  # type: ignore[arg-type]
+                # Check function signature to determine type
+                try:
+                    sig = inspect.signature(middleware)
+                    params = list(sig.parameters.values())
+                    if len(params) >= 1:
+                        first_param = params[0]
+                        # Check if first parameter is AgentRunContext or FunctionInvocationContext
+                        if (
+                            hasattr(first_param.annotation, "__name__")
+                            and first_param.annotation.__name__ == "AgentRunContext"
+                        ):
+                            agent_middlewares.append(middleware)  # type: ignore
+                        elif (
+                            hasattr(first_param.annotation, "__name__")
+                            and first_param.annotation.__name__ == "FunctionInvocationContext"
+                        ):
+                            function_middlewares.append(middleware)  # type: ignore
+                        else:
+                            # Default to agent middleware if uncertain
+                            agent_middlewares.append(middleware)  # type: ignore
+                    else:
+                        agent_middlewares.append(middleware)  # type: ignore
+                except Exception:
+                    # If signature inspection fails, assume it's an agent middleware
+                    agent_middlewares.append(middleware)  # type: ignore
+            else:
+                # Fallback
+                agent_middlewares.append(middleware)  # type: ignore
+
+        self._agent_middleware_pipeline = AgentMiddlewarePipeline(agent_middlewares)
+        self._function_middleware_pipeline = FunctionMiddlewarePipeline(function_middlewares)
+
+    async def middleware_enabled_run(
+        self: Any,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: Any = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        """Middleware-enabled run method."""
+        # Initialize middleware pipelines if not already done
+        if (
+            hasattr(self, "middleware")
+            and self.middleware
+            and not (
+                hasattr(self, "_agent_middleware_pipeline")
+                and hasattr(self, "_function_middleware_pipeline")
+                and (
+                    self._agent_middleware_pipeline.has_middlewares
+                    or self._function_middleware_pipeline.has_middlewares
+                )
+            )
+        ):
+            _initialize_middleware_pipelines(self, self.middleware)
+
+        # Ensure pipelines exist even if empty
+        if not hasattr(self, "_agent_middleware_pipeline"):
+            self._agent_middleware_pipeline = AgentMiddlewarePipeline()
+        if not hasattr(self, "_function_middleware_pipeline"):
+            self._function_middleware_pipeline = FunctionMiddlewarePipeline()
+
+        # Add function middleware pipeline to kwargs if available
+        if self._function_middleware_pipeline.has_middlewares:
+            kwargs["_function_middleware_pipeline"] = self._function_middleware_pipeline
+
+        normalized_messages = self._normalize_messages(messages)
+
+        # Execute with middleware if available
+        if self._agent_middleware_pipeline.has_middlewares:
+            context = AgentRunContext(
+                agent=self,  # type: ignore[arg-type]
+                messages=normalized_messages,
+                is_streaming=False,
+            )
+
+            async def _execute_handler(ctx: AgentRunContext) -> AgentRunResponse:
+                return await original_run(self, ctx.messages, thread=thread, **kwargs)  # type: ignore
+
+            result = await self._agent_middleware_pipeline.execute(
+                self,  # type: ignore[arg-type]
+                normalized_messages,
+                context,
+                _execute_handler,
+            )
+
+            return result if result else AgentRunResponse()
+
+        # No middleware, execute directly
+        return await original_run(self, normalized_messages, thread=thread, **kwargs)  # type: ignore[return-value]
+
+    def middleware_enabled_run_stream(
+        self: Any,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: Any = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Middleware-enabled run_stream method."""
+        # Initialize middleware pipelines if not already done
+        if (
+            hasattr(self, "middleware")
+            and self.middleware
+            and not (
+                hasattr(self, "_agent_middleware_pipeline")
+                and hasattr(self, "_function_middleware_pipeline")
+                and (
+                    self._agent_middleware_pipeline.has_middlewares
+                    or self._function_middleware_pipeline.has_middlewares
+                )
+            )
+        ):
+            _initialize_middleware_pipelines(self, self.middleware)
+
+        # Ensure pipelines exist even if empty
+        if not hasattr(self, "_agent_middleware_pipeline"):
+            self._agent_middleware_pipeline = AgentMiddlewarePipeline()
+        if not hasattr(self, "_function_middleware_pipeline"):
+            self._function_middleware_pipeline = FunctionMiddlewarePipeline()
+
+        # Add function middleware pipeline to kwargs if available
+        if self._function_middleware_pipeline.has_middlewares:
+            kwargs["_function_middleware_pipeline"] = self._function_middleware_pipeline
+
+        normalized_messages = self._normalize_messages(messages)
+
+        # Execute with middleware if available
+        if self._agent_middleware_pipeline.has_middlewares:
+            context = AgentRunContext(
+                agent=self,  # type: ignore[arg-type]
+                messages=normalized_messages,
+                is_streaming=True,
+            )
+
+            async def _execute_stream_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
+                async for update in original_run_stream(self, ctx.messages, thread=thread, **kwargs):  # type: ignore[misc]
+                    yield update
+
+            async def _stream_generator() -> AsyncIterable[AgentRunResponseUpdate]:
+                async for update in self._agent_middleware_pipeline.execute_stream(
+                    self,  # type: ignore[arg-type]
+                    normalized_messages,
+                    context,
+                    _execute_stream_handler,
+                ):
+                    yield update
+
+            return _stream_generator()
+
+        # No middleware, execute directly
+        return original_run_stream(self, normalized_messages, thread=thread, **kwargs)  # type: ignore
+
+    agent_class.run = middleware_enabled_run  # type: ignore
+    agent_class.run_stream = middleware_enabled_run_stream  # type: ignore
+
+    return agent_class
@@ -576,6 +576,7 @@ async def _auto_invoke_function(
    tool_map: dict[str, AIFunction[BaseModel, Any]],
    sequence_index: int | None = None,
    request_index: int | None = None,
+    middleware_pipeline: Any = None,  # Optional MiddlewarePipeline
 ) -> "Contents":
    """Invoke a function call requested by the agent, applying filters that are defined in the agent."""
    from ._types import FunctionResultContent
@@ -590,14 +591,43 @@ async def _auto_invoke_function(
    merged_args: dict[str, Any] = (custom_args or {}) | parsed_args
    args = tool.input_model.model_validate(merged_args)
    exception = None
-    try:
-        function_result = await tool.invoke(
+
+    # Execute through middleware pipeline if available
+    if middleware_pipeline and hasattr(middleware_pipeline, "has_middlewares") and middleware_pipeline.has_middlewares:
+        from ._middleware import FunctionInvocationContext
+
+        middleware_context = FunctionInvocationContext(
+            function=tool,
            arguments=args,
-            tool_call_id=function_call_content.call_id,
-        )  # type: ignore[arg-type]
-    except Exception as ex:
-        exception = ex
-        function_result = None
+        )
+
+        async def final_function_handler(context_obj: Any) -> Any:
+            return await tool.invoke(
+                arguments=context_obj.arguments,
+                tool_call_id=function_call_content.call_id,
+            )
+
+        try:
+            function_result = await middleware_pipeline.execute(
+                function=tool,
+                arguments=args,
+                context=middleware_context,
+                final_handler=final_function_handler,
+            )
+        except Exception as ex:
+            exception = ex
+            function_result = None
+    else:
+        # No middleware - execute directly
+        try:
+            function_result = await tool.invoke(
+                arguments=args,
+                tool_call_id=function_call_content.call_id,
+            )  # type: ignore[arg-type]
+        except Exception as ex:
+            exception = ex
+            function_result = None
+
    return FunctionResultContent(
        call_id=function_call_content.call_id,
        exception=exception,
@@ -631,6 +661,7 @@ async def execute_function_calls(
    | Callable[..., Any] \
    | MutableMapping[str, Any] \
    | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]",
+    middleware_pipeline: Any = None,  # Optional MiddlewarePipeline to avoid circular imports
 ) -> list["Contents"]:
    tool_map = _get_tool_map(tools)
    # Run all function calls concurrently
@@ -641,6 +672,7 @@ async def execute_function_calls(
            tool_map=tool_map,
            sequence_index=seq_idx,
            request_index=attempt_idx,
+            middleware_pipeline=middleware_pipeline,
        )
        for seq_idx, function_call in enumerate(function_calls)
    ])
@@ -706,11 +738,14 @@ def _handle_function_calls_response(
                if not tools and (chat_options := kwargs.get("chat_options")) and isinstance(chat_options, ChatOptions):
                    tools = chat_options.tools
                if function_calls and tools:
+                    # Extract function middleware pipeline from kwargs if available
+                    middleware_pipeline = kwargs.get("_function_middleware_pipeline")
                    function_results = await execute_function_calls(
                        custom_args=kwargs,
                        attempt_idx=attempt_idx,
                        function_calls=function_calls,
                        tools=tools,  # type: ignore
+                        middleware_pipeline=middleware_pipeline,
                    )
                    # add a single ChatMessage to the response with the results
                    result_message = ChatMessage(role="tool", contents=function_results)  # type: ignore[call-overload]
@@ -815,11 +850,14 @@ def _handle_function_calls_streaming_response(
                    tools = chat_options.tools

                if function_calls and tools:
+                    # Extract function middleware pipeline from kwargs if available
+                    middleware_pipeline = kwargs.get("_function_middleware_pipeline")
                    function_results = await execute_function_calls(
                        custom_args=kwargs,
                        attempt_idx=attempt_idx,
                        function_calls=function_calls,
                        tools=tools,  # type: ignore[reportArgumentType]
+                        middleware_pipeline=middleware_pipeline,
                    )
                    function_result_msg = ChatMessage(role="tool", contents=function_results)
                    yield ChatResponseUpdate(contents=function_results, role="tool")
@@ -1,27 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-
-from typing import Generic, Protocol, TypeVar, runtime_checkable
-
-TInput = TypeVar("TInput")
-TResponse = TypeVar("TResponse")
-
-__all__ = ["InputGuardrail", "OutputGuardrail"]
-
-
-@runtime_checkable
-class InputGuardrail(Protocol, Generic[TInput]):
-    """A protocol for input guardrails that can validate and transform input messages."""
-
-    def __call__(self, message: TInput) -> TInput:
-        """Validate and possibly transform the input message."""
-        ...
-
-
-@runtime_checkable
-class OutputGuardrail(Protocol, Generic[TResponse]):
-    """A protocol for output guardrails that can validate and transform output messages."""
-
-    def __call__(self, message: TResponse) -> TResponse:
-        """Validate and possibly transform the output message."""
-        ...
@@ -1026,8 +1026,12 @@ def _capture_messages(

    prepped = prepare_messages(messages)
    for index, message in enumerate(prepped):
+        try:
+            message_data = message.model_dump(exclude_none=True)
+        except Exception:
+            message_data = {"role": message.role.value, "contents": message.contents}
        logger.info(
-            message.model_dump_json(exclude_none=True),
+            message_data,
            extra={
                OtelAttr.EVENT_NAME: OtelAttr.CHOICE if output else ROLE_EVENT_MAP.get(message.role.value),
                OtelAttr.PROVIDER_NAME: provider_name,