Initial middleware implementation

python/packages/main/agent_framework/__init__.py

@@ -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

python/packages/main/agent_framework/_agents.py

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import inspect
 import sys
 from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack
@@ -12,6 +13,17 @@ from ._clients import BaseChatClient, ChatClientProtocol
 from ._logging import get_logger
 from ._mcp import MCPTool
 from ._memory import AggregateContextProvider, Context, ContextProvider
+from ._middleware import (
+    AgentInvocationContext,
+    AgentMiddleware,
+    AgentMiddlewareCallable,
+    AgentMiddlewarePipeline,
+    FunctionInvocationContext,
+    FunctionMiddleware,
+    FunctionMiddlewareCallable,
+    FunctionMiddlewarePipeline,
+    MiddlewareType,
+)
 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 +150,16 @@ 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.
+       middlewares: 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
+    middlewares: MiddlewareType | list[MiddlewareType] | None = None
+    _agent_middleware_pipeline: AgentMiddlewarePipeline = PrivateAttr(default_factory=AgentMiddlewarePipeline)
+    _function_middleware_pipeline: FunctionMiddlewarePipeline = PrivateAttr(default_factory=FunctionMiddlewarePipeline)
 
     async def _notify_thread_of_new_messages(
         self, thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage]
@@ -170,6 +186,223 @@ class BaseAgent(AFBaseModel):
         await deserialize_thread_state(thread, serialized_thread, **kwargs)
         return thread
 
+    def _initialize_middleware_pipelines(self, middlewares: MiddlewareType | list[MiddlewareType] | None) -> None:
+        """Initialize agent and function middleware pipelines from the provided middlewares.
+
+        This method classifies middleware by type and creates appropriate pipelines.
+
+        Args:
+            middlewares: List of middleware to classify and initialize into pipelines.
+        """
+        # Separate middlewares by type
+        agent_middlewares: list[AgentMiddleware | AgentMiddlewareCallable] = []
+        function_middlewares: list[FunctionMiddleware | FunctionMiddlewareCallable] = []
+
+        if middlewares:
+            middlewares_list = middlewares if isinstance(middlewares, list) else [middlewares]
+            for middleware in middlewares_list:
+                # Classify middleware by type checking
+                if isinstance(middleware, AgentMiddleware):
+                    agent_middlewares.append(middleware)
+                elif isinstance(middleware, FunctionMiddleware):
+                    function_middlewares.append(middleware)
+                elif callable(middleware) and inspect.iscoroutinefunction(middleware):
+                    # It's a function, classify by signature inspection
+                    sig = inspect.signature(middleware)
+                    params = list(sig.parameters.values())
+
+                    # Look at the first parameter's type annotation
+                    if len(params) >= 1:
+                        first_param_annotation = params[0].annotation
+                        if first_param_annotation != inspect.Parameter.empty:
+                            # Check if it's an AgentInvocationContext or FunctionInvocationContext
+                            if first_param_annotation is AgentInvocationContext:
+                                agent_middlewares.append(middleware)  # type: ignore
+                            elif first_param_annotation is FunctionInvocationContext:
+                                function_middlewares.append(middleware)  # type: ignore
+
+        self._agent_middleware_pipeline = AgentMiddlewarePipeline(agent_middlewares)
+        self._function_middleware_pipeline = FunctionMiddlewarePipeline(function_middlewares)
+
+    async def _run_impl(
+        self,
+        messages: list[ChatMessage],
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        """Core agent execution logic that subclasses must implement.
+
+        This method contains the actual agent logic and is called by the public run method
+        after middleware initialization and message normalization.
+
+        Args:
+            messages: Normalized list of ChatMessage objects.
+            thread: The conversation thread associated with the messages.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            An agent response.
+
+        Raises:
+            NotImplementedError: This method must be implemented by subclasses.
+        """
+        raise NotImplementedError("Subclasses must implement _run_impl method")
+
+    def _run_stream_impl(
+        self,
+        messages: list[ChatMessage],
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Core agent streaming execution logic that subclasses must implement.
+
+        This method contains the actual agent streaming logic and is called by the public
+        run_stream method after middleware initialization and message normalization.
+
+        Args:
+            messages: Normalized list of ChatMessage objects.
+            thread: The conversation thread associated with the messages.
+            kwargs: Additional keyword arguments.
+
+        Yields:
+            Agent response updates.
+
+        Raises:
+            NotImplementedError: This method must be implemented by subclasses.
+        """
+        raise NotImplementedError("Subclasses must implement _run_stream_impl method")
+
+    async def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        """Get a response from the agent with automatic middleware execution.
+
+        This method handles middleware initialization, message normalization, and execution
+        automatically. Custom agents should implement _run_impl instead of overriding this method.
+
+        Args:
+            messages: The message(s) to send to the agent.
+            thread: The conversation thread associated with the message(s).
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            An agent response item.
+        """
+        # Initialize middleware pipelines if not already done
+        if self.middlewares and not (
+            self._agent_middleware_pipeline.has_middlewares or self._function_middleware_pipeline.has_middlewares
+        ):
+            self._initialize_middleware_pipelines(self.middlewares)
+
+        # Normalize messages to ChatMessage list
+        if messages is None:
+            normalized_messages: list[ChatMessage] = []
+        elif isinstance(messages, str):
+            normalized_messages = [ChatMessage(role=Role.USER, text=messages)]
+        elif isinstance(messages, ChatMessage):
+            normalized_messages = [messages]
+        elif isinstance(messages, list):
+            normalized_messages = []
+            for msg in messages:
+                if isinstance(msg, str):
+                    normalized_messages.append(ChatMessage(role=Role.USER, text=msg))
+                elif isinstance(msg, ChatMessage):
+                    normalized_messages.append(msg)
+
+        # Execute with middleware if available
+        if self._agent_middleware_pipeline.has_middlewares:
+            from ._middleware import AgentInvocationContext
+
+            context = AgentInvocationContext(
+                agent=self,  # type: ignore[arg-type]
+                messages=normalized_messages,
+                is_streaming=False,
+            )
+
+            async def _execute_handler(ctx: AgentInvocationContext) -> AgentRunResponse:
+                return await self._run_impl(ctx.messages, thread=thread, **kwargs)
+
+            return await self._agent_middleware_pipeline.execute(
+                self,  # type: ignore[arg-type]
+                normalized_messages,
+                context,
+                _execute_handler,
+            )
+
+        # No middleware, execute directly
+        return await self._run_impl(normalized_messages, thread=thread, **kwargs)
+
+    def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Run the agent as a stream with automatic middleware execution.
+
+        This method handles middleware initialization, message normalization, and execution
+        automatically. Custom agents should implement _run_stream_impl instead of overriding this method.
+
+        Args:
+            messages: The message(s) to send to the agent.
+            thread: The conversation thread associated with the message(s).
+            kwargs: Additional keyword arguments.
+
+        Yields:
+            Agent response items.
+        """
+        # Initialize middleware pipelines if not already done
+        if self.middlewares and not (
+            self._agent_middleware_pipeline.has_middlewares or self._function_middleware_pipeline.has_middlewares
+        ):
+            self._initialize_middleware_pipelines(self.middlewares)
+
+        # Normalize messages to ChatMessage list
+        if messages is None:
+            normalized_messages: list[ChatMessage] = []
+        elif isinstance(messages, str):
+            normalized_messages = [ChatMessage(role=Role.USER, text=messages)]
+        elif isinstance(messages, ChatMessage):
+            normalized_messages = [messages]
+        elif isinstance(messages, list):
+            normalized_messages = []
+            for msg in messages:
+                if isinstance(msg, str):
+                    normalized_messages.append(ChatMessage(role=Role.USER, text=msg))
+                elif isinstance(msg, ChatMessage):
+                    normalized_messages.append(msg)
+
+        # Execute with middleware if available
+        if self._agent_middleware_pipeline.has_middlewares:
+            from ._middleware import AgentInvocationContext
+
+            context = AgentInvocationContext(
+                agent=self,  # type: ignore[arg-type]
+                messages=normalized_messages,
+                is_streaming=True,
+            )
+
+            async def _execute_stream_handler(ctx: AgentInvocationContext) -> AsyncIterable[AgentRunResponseUpdate]:
+                async for update in self._run_stream_impl(ctx.messages, thread=thread, **kwargs):
+                    yield update
+
+            return self._agent_middleware_pipeline.execute_stream(
+                self,  # type: ignore[arg-type]
+                normalized_messages,
+                context,
+                _execute_stream_handler,
+            )
+
+        # No middleware, execute directly
+        return self._run_stream_impl(normalized_messages, thread=thread, **kwargs)
+
 
 # region ChatAgent
 
@@ -216,6 +449,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,
+        middlewares: MiddlewareType | list[MiddlewareType] | None = None,
         **kwargs: Any,
     ) -> None:
         """Create a ChatAgent.
@@ -251,6 +485,8 @@ 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.
+            middlewares: List of middleware to intercept agent and function invocations.
+                Can include class instances implementing middleware protocols or pure functions.
             kwargs: any additional keyword arguments.
                 Unused, can be used by subclasses of this Agent.
         """
@@ -272,6 +508,7 @@ class ChatAgent(BaseAgent):
             "chat_client": chat_client,
             "chat_message_store_factory": chat_message_store_factory,
             "context_providers": aggregate_context_providers,
+            "middlewares": middlewares,
             "chat_options": ChatOptions(
                 ai_model_id=model,
                 frequency_penalty=frequency_penalty,
@@ -304,6 +541,9 @@ class ChatAgent(BaseAgent):
         self._update_agent_name()
         self._local_mcp_tools = local_mcp_tools  # type: ignore[assignment]
 
+        # Initialize middleware pipelines
+        self._initialize_middleware_pipelines(middlewares)
+
     async def __aenter__(self) -> "Self":
         """Async context manager entry.
 
@@ -338,9 +578,9 @@ class ChatAgent(BaseAgent):
         if hasattr(self.chat_client, "_update_agent_name") and callable(self.chat_client._update_agent_name):  # type: ignore[reportAttributeAccessIssue, attr-defined]
             self.chat_client._update_agent_name(self.name)  # type: ignore[reportAttributeAccessIssue, attr-defined]
 
-    async def run(
+    async def _run_impl(
         self,
-        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        messages: list[ChatMessage],
         *,
         thread: AgentThread | None = None,
         frequency_penalty: float | None = None,
@@ -367,16 +607,10 @@ class ChatAgent(BaseAgent):
         additional_properties: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> AgentRunResponse:
-        """Run the agent with the given messages and options.
-
-        Remarks:
-            Since you won't always call the agent.run directly, but it get's called
-            through orchestration, it is advised to set your default values for
-            all the chat client parameters in the agent constructor.
-            If both parameters are used, the ones passed to the run methods take precedence.
+        """Core ChatAgent execution logic called by the base run method after middleware processing.
 
         Args:
-            messages: The messages to process.
+            messages: Normalized list of ChatMessage objects.
             thread: The thread to use for the agent.
             frequency_penalty: the frequency penalty to use.
             logit_bias: the logit bias to use.
@@ -396,7 +630,60 @@ class ChatAgent(BaseAgent):
             additional_properties: additional properties to include in the request.
             kwargs: Additional keyword arguments for the agent.
                 will only be passed to functions that are called.
+
+        Returns:
+            An agent response.
         """
+        return await self._run_internal(
+            messages=messages,
+            thread=thread,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            max_tokens=max_tokens,
+            metadata=metadata,
+            model=model,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            store=store,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,  # type: ignore[arg-type]
+            top_p=top_p,
+            user=user,
+            additional_properties=additional_properties,
+            **kwargs,
+        )
+
+    async def _run_internal(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | dict[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | dict[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        """Internal run method that performs the actual agent execution without middleware."""
         input_messages = self._normalize_messages(messages)
         context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
         thread, thread_messages = await self._prepare_thread_and_messages(
@@ -467,7 +754,83 @@ class ChatAgent(BaseAgent):
             additional_properties=response.additional_properties,
         )
 
-    async def run_stream(
+    def _run_stream_impl(
+        self,
+        messages: list[ChatMessage],
+        *,
+        thread: AgentThread | None = None,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Core ChatAgent streaming execution logic called by the base run_stream method after middleware processing.
+
+        Args:
+            messages: Normalized list of ChatMessage objects.
+            thread: The thread to use for the agent.
+            frequency_penalty: the frequency penalty to use.
+            logit_bias: the logit bias to use.
+            max_tokens: The maximum number of tokens to generate.
+            metadata: additional metadata to include in the request.
+            model: The model to use for the agent.
+            presence_penalty: the presence penalty to use.
+            response_format: the format of the response.
+            seed: the random seed to use.
+            stop: the stop sequence(s) for the request.
+            store: whether to store the response.
+            temperature: the sampling temperature to use.
+            tool_choice: the tool choice for the request.
+            tools: the tools to use for the request.
+            top_p: the nucleus sampling probability to use.
+            user: the user to associate with the request.
+            additional_properties: additional properties to include in the request.
+            kwargs: any additional keyword arguments.
+                will only be passed to functions that are called.
+
+        Yields:
+            Agent response updates.
+        """
+        return self._run_stream_internal(
+            messages=messages,
+            thread=thread,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            max_tokens=max_tokens,
+            metadata=metadata,
+            model=model,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            store=store,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=tools,  # type: ignore[arg-type]
+            top_p=top_p,
+            user=user,
+            additional_properties=additional_properties,
+            **kwargs,
+        )
+
+    async def _run_stream_internal(
         self,
         messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
         *,
@@ -494,37 +857,7 @@ class ChatAgent(BaseAgent):
         additional_properties: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> AsyncIterable[AgentRunResponseUpdate]:
-        """Stream the agent with the given messages and options.
-
-        Remarks:
-            Since you won't always call the agent.run_stream directly, but it get's called
-            through orchestration, it is advised to set your default values for
-            all the chat client parameters in the agent constructor.
-            If both parameters are used, the ones passed to the run methods take precedence.
-
-        Args:
-            messages: The messages to process.
-            thread: The thread to use for the agent.
-            frequency_penalty: the frequency penalty to use.
-            logit_bias: the logit bias to use.
-            max_tokens: The maximum number of tokens to generate.
-            metadata: additional metadata to include in the request.
-            model: The model to use for the agent.
-            presence_penalty: the presence penalty to use.
-            response_format: the format of the response.
-            seed: the random seed to use.
-            stop: the stop sequence(s) for the request.
-            store: whether to store the response.
-            temperature: the sampling temperature to use.
-            tool_choice: the tool choice for the request.
-            tools: the tools to use for the request.
-            top_p: the nucleus sampling probability to use.
-            user: the user to associate with the request.
-            additional_properties: additional properties to include in the request.
-            kwargs: any additional keyword arguments.
-                will only be passed to functions that are called.
-
-        """
+        """Internal run_stream method that performs the actual agent execution without middleware."""
         input_messages = self._normalize_messages(messages)
         context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
         thread, thread_messages = await self._prepare_thread_and_messages(

python/packages/main/agent_framework/_clients.py

@@ -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 MiddlewareType
 from ._pydantic import AFBaseModel
 from ._threads import ChatMessageStore
 from ._tools import ToolProtocol
@@ -465,6 +466,7 @@ class BaseChatClient(AFBaseModel, ABC):
         | None = None,
         chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
         context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
+        middlewares: MiddlewareType | list[MiddlewareType] | None = None,
         **kwargs: Any,
     ) -> "ChatAgent":
         """Create an agent with the given name and instructions.
@@ -476,6 +478,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.
+            middlewares: 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 +494,7 @@ class BaseChatClient(AFBaseModel, ABC):
             tools=tools,
             chat_message_store_factory=chat_message_store_factory,
             context_providers=context_providers,
+            middlewares=middlewares,
             **kwargs,
         )
 

python/packages/main/agent_framework/_middleware.py

@@ -0,0 +1,414 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterable, Awaitable, Callable
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+    from ._agents import AgentProtocol
+    from ._tools import AIFunction
+    from ._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage
+
+__all__ = [
+    "AgentInvocationContext",
+    "AgentMiddleware",
+    "FunctionInvocationContext",
+    "FunctionMiddleware",
+    "MiddlewareType",
+]
+
+
+class AgentInvocationContext:
+    """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.
+        request_id: Unique identifier for the current request.
+        metadata: Metadata dictionary for sharing data between agent middleware.
+    """
+
+    def __init__(
+        self,
+        agent: "AgentProtocol",
+        messages: list["ChatMessage"],
+        is_streaming: bool = False,
+        request_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize agent invocation context.
+
+        Args:
+            agent: The agent being invoked.
+            messages: The messages being sent to the agent.
+            is_streaming: Whether this is a streaming invocation.
+            request_id: Unique identifier for the request. Auto-generated if None.
+            metadata: Metadata dictionary.
+        """
+        self.agent = agent
+        self.messages = messages
+        self.is_streaming = is_streaming
+        self.request_id = request_id or str(uuid4())
+        self.metadata = metadata or {}
+
+
+class FunctionInvocationContext:
+    """Context object for function middleware invocations.
+
+    Attributes:
+        function: The function being invoked.
+        arguments: The validated arguments for the function.
+        request_id: Unique identifier for the current request.
+        metadata: Metadata dictionary for sharing data between function middleware.
+    """
+
+    def __init__(
+        self,
+        function: "AIFunction[Any, Any]",
+        arguments: "BaseModel",
+        request_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize function invocation context.
+
+        Args:
+            function: The function being invoked.
+            arguments: The validated arguments for the function.
+            request_id: Unique identifier for the request. Auto-generated if None.
+            metadata: Metadata dictionary.
+        """
+        self.function = function
+        self.arguments = arguments
+        self.request_id = request_id or str(uuid4())
+        self.metadata = metadata or {}
+
+
+class AgentMiddleware(ABC):
+    """Abstract base class for agent middleware that can intercept agent invocations."""
+
+    @abstractmethod
+    async def process(
+        self,
+        context: AgentInvocationContext,
+        next: Callable[[AgentInvocationContext], 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.should_skip=True and provide context.response
+                    or context.response_stream to override the agent execution.
+            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.should_skip=True and provide
+            context.response or context.response_stream to override execution.
+        """
+        ...
+
+
+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.should_skip=True and provide context.result
+                    to override the function execution.
+            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.should_skip=True and provide
+            context.result to override execution.
+        """
+        ...
+
+
+# Pure function type definitions for convenience
+AgentMiddlewareCallable = Callable[
+    [AgentInvocationContext, Callable[[AgentInvocationContext], Awaitable[None]]], Awaitable[None]
+]
+
+FunctionMiddlewareCallable = Callable[
+    [FunctionInvocationContext, Callable[[FunctionInvocationContext], Awaitable[None]]], Awaitable[None]
+]
+
+# Type alias for all middleware types
+MiddlewareType = AgentMiddleware | AgentMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable
+
+
+class AgentMiddlewareWrapper:
+    """Wrapper to convert pure functions into AgentMiddleware protocol objects."""
+
+    def __init__(self, func: AgentMiddlewareCallable):
+        self.func = func
+
+    async def process(
+        self,
+        context: AgentInvocationContext,
+        next: Callable[[AgentInvocationContext], Awaitable[None]],
+    ) -> None:
+        await self.func(context, next)
+
+
+class FunctionMiddlewareWrapper:
+    """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 AgentMiddlewarePipeline:
+    """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.
+        """
+        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 callable(middleware):
+            # Check if it's already a protocol implementation
+            if callable(middleware) and not hasattr(middleware, "func"):
+                # It's a class instance implementing the protocol
+                self._middlewares.append(middleware)  # type: ignore
+            else:
+                # It's a pure function, wrap it
+                self._middlewares.append(AgentMiddlewareWrapper(middleware))  # type: ignore
+        else:
+            self._middlewares.append(middleware)  # type: ignore
+
+    async def execute(
+        self,
+        agent: "AgentProtocol",
+        messages: list["ChatMessage"],
+        context: AgentInvocationContext,
+        final_handler: Callable[[AgentInvocationContext], Awaitable["AgentRunResponse"]],
+    ) -> "AgentRunResponse":
+        """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[[AgentInvocationContext], Awaitable[None]]:
+            if index >= len(self._middlewares):
+
+                async def final_wrapper(c: AgentInvocationContext) -> None:
+                    result_container["response"] = await final_handler(c)
+
+                return final_wrapper
+
+            middleware = self._middlewares[index]
+            next_handler = create_next_handler(index + 1)
+
+            async def current_handler(c: AgentInvocationContext) -> None:
+                await middleware.process(c, next_handler)
+
+            return current_handler
+
+        first_handler = create_next_handler(0)
+        await first_handler(context)
+
+        # Return the response from result container
+        response = result_container["response"]
+        if response is None:
+            raise RuntimeError("No response set after middleware execution")
+        return response
+
+    async def execute_stream(
+        self,
+        agent: "AgentProtocol",
+        messages: list["ChatMessage"],
+        context: AgentInvocationContext,
+        final_handler: Callable[[AgentInvocationContext], 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] = {"response_stream": None}
+
+        def create_next_handler(index: int) -> Callable[[AgentInvocationContext], Awaitable[None]]:
+            if index >= len(self._middlewares):
+
+                async def final_wrapper(c: AgentInvocationContext) -> None:  # noqa: RUF029
+                    result_container["response_stream"] = final_handler(c)
+
+                return final_wrapper
+
+            middleware = self._middlewares[index]
+            next_handler = create_next_handler(index + 1)
+
+            async def current_handler(c: AgentInvocationContext) -> None:
+                await middleware.process(c, next_handler)
+
+            return current_handler
+
+        first_handler = create_next_handler(0)
+        await first_handler(context)
+
+        # Yield from the response stream in result container
+        response_stream = result_container["response_stream"]
+        if response_stream is None:
+            raise RuntimeError("No response stream set after middleware execution")
+
+        async for update in response_stream:
+            yield update
+
+    @property
+    def has_middlewares(self) -> bool:
+        """Check if there are any middlewares registered."""
+        return bool(self._middlewares)
+
+
+class FunctionMiddlewarePipeline:
+    """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.
+        """
+        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."""
+        if callable(middleware):
+            # Check if it's already a protocol implementation
+            if callable(middleware) and not hasattr(middleware, "func"):
+                # It's a class instance implementing the protocol
+                self._middlewares.append(middleware)  # type: ignore
+            else:
+                # It's a pure function, wrap it
+                self._middlewares.append(FunctionMiddlewareWrapper(middleware))  # type: ignore
+        else:
+            self._middlewares.append(middleware)  # type: ignore
+
+    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}
+
+        def create_next_handler(index: int) -> Callable[[FunctionInvocationContext], Awaitable[None]]:
+            if index >= len(self._middlewares):
+
+                async def final_wrapper(c: FunctionInvocationContext) -> None:
+                    result_container["result"] = await final_handler(c)
+
+                return final_wrapper
+
+            middleware = self._middlewares[index]
+            next_handler = create_next_handler(index + 1)
+
+            async def current_handler(c: FunctionInvocationContext) -> None:
+                await middleware.process(c, next_handler)
+
+            return current_handler
+
+        first_handler = create_next_handler(0)
+        await first_handler(context)
+
+        # Return the result from result container
+        result = result_container["result"]
+        if result is None:
+            raise RuntimeError("No result set after middleware execution")
+        return result
+
+    @property
+    def has_middlewares(self) -> bool:
+        """Check if there are any middlewares registered."""
+        return bool(self._middlewares)

python/packages/main/agent_framework/_tools.py

@@ -571,6 +571,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
@@ -585,14 +586,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,
@@ -626,6 +656,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
@@ -636,6 +667,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)
     ])
@@ -701,11 +733,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]
@@ -810,11 +845,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")

python/packages/workflow/agent_framework_workflow/_agent.py

@@ -107,7 +107,7 @@ class WorkflowAgent(BaseAgent):
         thread = thread or self.get_new_thread()
         response_id = str(uuid.uuid4())
 
-        async for update in self._run_stream_impl(input_messages, response_id):
+        async for update in self._run_stream_impl(input_messages, thread=thread, response_id=response_id):
             response_updates.append(update)
 
         # Convert updates to final response.
@@ -141,7 +141,7 @@ class WorkflowAgent(BaseAgent):
         response_updates: list[AgentRunResponseUpdate] = []
         response_id = str(uuid.uuid4())
 
-        async for update in self._run_stream_impl(input_messages, response_id):
+        async for update in self._run_stream_impl(input_messages, thread=thread, response_id=response_id):
             response_updates.append(update)
             yield update
 
@@ -154,18 +154,25 @@ class WorkflowAgent(BaseAgent):
 
     async def _run_stream_impl(
         self,
-        input_messages: list[ChatMessage],
-        response_id: str,
+        messages: list[ChatMessage],
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
     ) -> AsyncIterable[AgentRunResponseUpdate]:
         """Internal implementation of streaming execution.
 
         Args:
-            input_messages: Normalized input messages to process.
-            response_id: The unique response ID for this workflow execution.
+            messages: Normalized input messages to process.
+            thread: The conversation thread (unused in workflow agent).
+            kwargs: Additional keyword arguments, including response_id.
 
         Yields:
             AgentRunResponseUpdate objects representing the workflow execution progress.
         """
+        # Extract response_id from kwargs
+        response_id = kwargs.get("response_id")
+        if not response_id:
+            raise ValueError("response_id must be provided in kwargs")
         # Determine the event stream based on whether we have function responses
         if bool(self.pending_requests):
             # This is a continuation - use send_responses_streaming to send function responses back
@@ -174,7 +181,7 @@ class WorkflowAgent(BaseAgent):
             # Extract function responses from input messages, and ensure that
             # only function responses are present in messages if there is any
             # pending request.
-            function_responses = self._extract_function_responses(input_messages)
+            function_responses = self._extract_function_responses(messages)
 
             # Pop pending requests if fulfilled.
             for request_id in list(self.pending_requests.keys()):
@@ -188,7 +195,7 @@ class WorkflowAgent(BaseAgent):
         else:
             # Execute workflow with streaming (initial run or no function responses)
             # Pass the new input messages directly to the workflow
-            event_stream = self.workflow.run_stream(input_messages)
+            event_stream = self.workflow.run_stream(messages)
 
         # Process events from the stream
         async for event in event_stream:

python/samples/getting_started/middleware/class_based_middleware.py

@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import time
+from collections.abc import Awaitable, Callable
+from random import randint
+from typing import Annotated
+
+from agent_framework import AgentInvocationContext, AgentMiddleware, FunctionInvocationContext, FunctionMiddleware
+from agent_framework.foundry import FoundryChatClient
+from azure.identity.aio import AzureCliCredential
+from pydantic import Field
+
+
+def get_weather(
+    location: Annotated[str, Field(description="The location to get the weather for.")],
+) -> str:
+    """Get the weather for a given location."""
+    conditions = ["sunny", "cloudy", "rainy", "stormy"]
+    return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
+
+
+class SecurityAgentMiddleware(AgentMiddleware):
+    """Agent middleware that checks for security violations."""
+
+    async def process(
+        self,
+        context: AgentInvocationContext,
+        next: Callable[[AgentInvocationContext], Awaitable[None]],
+    ) -> None:
+        # Check for potential security violations in the query
+        # Look at the last user message
+        last_message = context.messages[-1] if context.messages else None
+        if last_message and last_message.text:
+            query = last_message.text
+            if "password" in query.lower() or "secret" in query.lower():
+                print("[SecurityAgentMiddleware] Security Warning: Detected sensitive information, blocking request.")
+                # Simply don't call next() to prevent execution
+                return
+
+        print("[SecurityAgentMiddleware] Security check passed.")
+        await next(context)
+
+
+class LoggingFunctionMiddleware(FunctionMiddleware):
+    """Function middleware that logs function calls."""
+
+    async def process(
+        self,
+        context: FunctionInvocationContext,
+        next: Callable[[FunctionInvocationContext], Awaitable[None]],
+    ) -> None:
+        function_name = context.function.name
+        print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.")
+
+        start_time = time.time()
+
+        await next(context)
+
+        end_time = time.time()
+        duration = end_time - start_time
+
+        print(f"[LoggingFunctionMiddleware] Function {function_name} completed in {duration:.3f}s.")
+
+
+async def main() -> None:
+    """Example demonstrating class-based middleware."""
+    print("=== Class-based Middleware Example ===")
+
+    # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
+    # authentication option.
+    async with (
+        AzureCliCredential() as credential,
+        FoundryChatClient(async_credential=credential).create_agent(
+            name="WeatherAgent",
+            instructions="You are a helpful weather assistant.",
+            tools=get_weather,
+            middlewares=[SecurityAgentMiddleware(), LoggingFunctionMiddleware()],
+        ) as agent,
+    ):
+        # Test with normal query
+        print("\n--- Normal Query ---")
+        query = "What's the weather like in Seattle?"
+        print(f"User: {query}")
+        result = await agent.run(query)
+        print(f"Agent: {result}\n")
+
+        # Test with security-related query
+        print("--- Security Test ---")
+        query = "What's the password for the weather service?"
+        print(f"User: {query}")
+        result = await agent.run(query)
+        print(f"Agent: {result}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/samples/getting_started/middleware/function_based_middleware.py

@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import time
+from collections.abc import Callable
+from random import randint
+from typing import Annotated, Any
+
+from agent_framework import (
+    AgentInvocationContext,
+    FunctionInvocationContext,
+)
+from agent_framework.foundry import FoundryChatClient
+from azure.identity.aio import AzureCliCredential
+from pydantic import Field
+
+
+def get_weather(
+    location: Annotated[str, Field(description="The location to get the weather for.")],
+) -> str:
+    """Get the weather for a given location."""
+    conditions = ["sunny", "cloudy", "rainy", "stormy"]
+    return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
+
+
+async def security_agent_middleware(
+    context: AgentInvocationContext,
+    next: Callable[[AgentInvocationContext], Any],
+) -> None:
+    """Agent middleware that checks for security violations."""
+    # Check for potential security violations in the query
+    # For this example, we'll check the last user message
+    last_message = context.messages[-1] if context.messages else None
+    if last_message and last_message.text:
+        query = last_message.text
+        if "password" in query.lower() or "secret" in query.lower():
+            print("[SecurityAgentMiddleware] Security Warning: Detected potential sensitive information.")
+            # Simply don't call next() to prevent execution
+            return
+
+    print("[SecurityAgentMiddleware] Security check passed.")
+    await next(context)
+
+
+async def logging_function_middleware(
+    context: FunctionInvocationContext,
+    next: Callable[[FunctionInvocationContext], Any],
+) -> None:
+    """Function middleware that logs function calls."""
+    function_name = context.function.name
+    print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.")
+
+    start_time = time.time()
+
+    await next(context)
+
+    end_time = time.time()
+    duration = end_time - start_time
+
+    print(f"[LoggingFunctionMiddleware] Function {function_name} completed in {duration:.3f}s.")
+
+
+async def main() -> None:
+    """Example demonstrating function-based middleware."""
+    print("=== Function-based Middleware Example ===")
+
+    # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
+    # authentication option.
+    async with (
+        AzureCliCredential() as credential,
+        FoundryChatClient(async_credential=credential).create_agent(
+            name="WeatherAgent",
+            instructions="You are a helpful weather assistant.",
+            tools=get_weather,
+            middlewares=[security_agent_middleware, logging_function_middleware],
+        ) as agent,
+    ):
+        # Test with normal query
+        print("\n--- Normal Query ---")
+        query = "What's the weather like in Tokyo?"
+        print(f"User: {query}")
+        result = await agent.run(query)
+        print(f"Agent: {result}\n")
+
+        # Test with security violation
+        print("--- Security Test ---")
+        query = "What's the secret weather password?"
+        print(f"User: {query}")
+        result = await agent.run(query)
+        print(f"Agent: {result}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())