Fix mypy type errors and update purview middleware for MiddlewareTermination

- Fix type annotations in _middleware.py for context.result assignments
- Fix type annotations in _agents.py for chat_client.get_response calls
- Update ResponseStream result_hooks type to allow Awaitable[TFinal | None]
- Convert Sequence to list in FunctionInvocationLayer init
- Add type: ignore comments for cooperative inheritance patterns
- Update AgentRunContext.result type to use ResponseStream instead of AsyncIterable
- Update purview middleware to use raise MiddlewareTermination instead of terminate flag
- Ensure MiddlewareTermination is not caught by generic Exception handlers
- Update purview tests to expect MiddlewareTermination

Author: eavanvalkenburg

python/packages/core/agent_framework/_agents.py

@@ -874,7 +874,7 @@ async def _run_non_streaming() -> AgentResponse[Any]:
                     options=options,
                     kwargs=kwargs,
                 )
-                response = await self.chat_client.get_response(
+                response = await self.chat_client.get_response(  # type: ignore[call-overload]
                     messages=ctx["thread_messages"],
                     stream=False,
                     options=ctx["chat_options"],
@@ -944,8 +944,8 @@ async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]:
                 options=options,
                 kwargs=kwargs,
             )
-            ctx = ctx_holder["ctx"]
-            return self.chat_client.get_response(
+            ctx: _RunContext = ctx_holder["ctx"]  # type: ignore[assignment]  # Safe: we just assigned it
+            return self.chat_client.get_response(  # type: ignore[call-overload, no-any-return]
                 messages=ctx["thread_messages"],
                 stream=True,
                 options=ctx["chat_options"],

python/packages/core/agent_framework/_middleware.py

@@ -125,7 +125,7 @@ class AgentRunContext:
         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 AgentResponse.
-                For streaming: should be AsyncIterable[AgentResponseUpdate].
+                For streaming: should be ResponseStream[AgentResponseUpdate, AgentResponse].
         kwargs: Additional keyword arguments passed to the agent run method.
 
     Examples:
@@ -160,7 +160,7 @@ def __init__(
         options: Mapping[str, Any] | None = None,
         stream: bool = False,
         metadata: Mapping[str, Any] | None = None,
-        result: AgentResponse | AsyncIterable[AgentResponseUpdate] | None = None,
+        result: AgentResponse | ResponseStream[AgentResponseUpdate, AgentResponse] | None = None,
         kwargs: Mapping[str, Any] | None = None,
         stream_transform_hooks: Sequence[
             Callable[[AgentResponseUpdate], AgentResponseUpdate | Awaitable[AgentResponseUpdate]]
@@ -767,7 +767,7 @@ async def execute(
             The agent response after processing through all middleware.
         """
         if not self._middleware:
-            context.result = final_handler(context)
+            context.result = final_handler(context)  # type: ignore[assignment]
             if isinstance(context.result, Awaitable):
                 context.result = await context.result
             return context.result
@@ -776,7 +776,7 @@ def create_next_handler(index: int) -> Callable[[AgentRunContext], Awaitable[Non
             if index >= len(self._middleware):
 
                 async def final_wrapper(c: AgentRunContext) -> None:
-                    c.result = final_handler(c)
+                    c.result = final_handler(c)  # type: ignore[assignment]
                     if inspect.isawaitable(c.result):
                         c.result = await c.result
 
@@ -904,7 +904,7 @@ async def execute(
         final_handler: Callable[
             [ChatContext], Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]
         ],
-    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+    ) -> ChatResponse | ResponseStream[ChatResponseUpdate, ChatResponse] | None:
         """Execute the chat middleware pipeline.
 
         Args:
@@ -915,7 +915,7 @@ async def execute(
             The chat response after processing through all middleware.
         """
         if not self._middleware:
-            context.result = final_handler(context)
+            context.result = final_handler(context)  # type: ignore[assignment]
             if isinstance(context.result, Awaitable):
                 context.result = await context.result
             if context.stream and not isinstance(context.result, ResponseStream):
@@ -926,7 +926,7 @@ def create_next_handler(index: int) -> Callable[[ChatContext], Awaitable[None]]:
             if index >= len(self._middleware):
 
                 async def final_wrapper(c: ChatContext) -> None:
-                    c.result = final_handler(c)
+                    c.result = final_handler(c)  # type: ignore[assignment]
                     if inspect.isawaitable(c.result):
                         c.result = await c.result
 
@@ -1027,15 +1027,15 @@ def get_response(
             *middleware["chat"],
         )
         if not pipeline.has_middlewares:
-            return super_get_response(
+            return super_get_response(  # type: ignore[no-any-return]
                 messages=messages,
                 stream=stream,
                 options=options,
                 **kwargs,
             )
 
         context = ChatContext(
-            chat_client=self,
+            chat_client=self,  # type: ignore[arg-type]
             messages=prepare_messages(messages),
             options=options,
             stream=stream,
@@ -1063,13 +1063,13 @@ async def _execute_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]:
             return ResponseStream.from_awaitable(_execute_stream())
 
         # For non-streaming, return the coroutine directly
-        return _execute()
+        return _execute()  # type: ignore[return-value]
 
     def _middleware_handler(
         self, context: ChatContext
     ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
         """Internal middleware handler to adapt to pipeline."""
-        return super().get_response(
+        return super().get_response(  # type: ignore[misc, no-any-return]
             messages=context.messages,
             stream=context.stream,
             options=context.options or {},
@@ -1089,7 +1089,7 @@ def __init__(
         middleware_list = categorize_middleware(middleware)
         self.agent_middleware = middleware_list["agent"]
         # Pass middleware to super so BaseAgent can store it for dynamic rebuild
-        super().__init__(*args, middleware=middleware, **kwargs)
+        super().__init__(*args, middleware=middleware, **kwargs)  # type: ignore[call-arg]
         if chat_client := getattr(self, "chat_client", None):
             client_chat_middleware = getattr(chat_client, "chat_middleware", [])
             client_chat_middleware.extend(middleware_list["chat"])
@@ -1157,10 +1157,10 @@ def run(
 
         # Execute with middleware if available
         if not pipeline.has_middlewares:
-            return super().run(messages, stream=stream, thread=thread, options=options, **combined_kwargs)
+            return super().run(messages, stream=stream, thread=thread, options=options, **combined_kwargs)  # type: ignore[misc, no-any-return]
 
         context = AgentRunContext(
-            agent=self,
+            agent=self,  # type: ignore[arg-type]
             messages=prepare_messages(messages),
             thread=thread,
             options=options,
@@ -1189,12 +1189,12 @@ async def _execute_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse
             return ResponseStream.from_awaitable(_execute_stream())
 
         # For non-streaming, return the coroutine directly
-        return _execute()
+        return _execute()  # type: ignore[return-value]
 
     def _middleware_handler(
         self, context: AgentRunContext
     ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
-        return super().run(
+        return super().run(  # type: ignore[misc, no-any-return]
             context.messages,
             stream=context.stream,
             thread=context.thread,

python/packages/core/agent_framework/_tools.py

@@ -2047,7 +2047,9 @@ def __init__(
         function_invocation_configuration: FunctionInvocationConfiguration | None = None,
         **kwargs: Any,
     ) -> None:
-        self.function_middleware: list[FunctionMiddlewareTypes] = function_middleware or []
+        self.function_middleware: list[FunctionMiddlewareTypes] = (
+            list(function_middleware) if function_middleware else []
+        )
         self.function_invocation_configuration = normalize_function_invocation_configuration(
             function_invocation_configuration
         )

python/packages/core/agent_framework/_types.py

@@ -2464,7 +2464,7 @@ def __init__(
         finalizer: Callable[[Sequence[TUpdate]], TFinal | Awaitable[TFinal]] | None = None,
         transform_hooks: list[Callable[[TUpdate], TUpdate | Awaitable[TUpdate] | None]] | None = None,
         cleanup_hooks: list[Callable[[], Awaitable[None] | None]] | None = None,
-        result_hooks: list[Callable[[TFinal], TFinal | Awaitable[TFinal] | None]] | None = None,
+        result_hooks: list[Callable[[TFinal], TFinal | Awaitable[TFinal | None] | None]] | None = None,
     ) -> None:
         """A Async Iterable stream of updates.
 
@@ -2489,7 +2489,7 @@ def __init__(
         self._transform_hooks: list[Callable[[TUpdate], TUpdate | Awaitable[TUpdate] | None]] = (
             transform_hooks if transform_hooks is not None else []
         )
-        self._result_hooks: list[Callable[[TFinal], TFinal | Awaitable[TFinal] | None]] = (
+        self._result_hooks: list[Callable[[TFinal], TFinal | Awaitable[TFinal | None] | None]] = (
             result_hooks if result_hooks is not None else []
         )
         self._cleanup_hooks: list[Callable[[], Awaitable[None] | None]] = (
@@ -2748,7 +2748,7 @@ def with_transform_hook(
 
     def with_result_hook(
         self,
-        hook: Callable[[TFinal], TFinal | Awaitable[TFinal] | None],
+        hook: Callable[[TFinal], TFinal | Awaitable[TFinal | None] | None],
     ) -> ResponseStream[TUpdate, TFinal]:
         """Register a result hook executed after finalization."""
         self._result_hooks.append(hook)

python/packages/purview/agent_framework_purview/_middleware.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Awaitable, Callable
 
-from agent_framework import AgentMiddleware, AgentRunContext, ChatContext, ChatMiddleware
+from agent_framework import AgentMiddleware, AgentRunContext, ChatContext, ChatMiddleware, MiddlewareTermination
 from agent_framework._logging import get_logger
 from azure.core.credentials import TokenCredential
 from azure.core.credentials_async import AsyncTokenCredential
@@ -62,8 +62,9 @@ async def process(
                 context.result = AgentResponse(
                     messages=[ChatMessage(role=Role.SYSTEM, text=self._settings.blocked_prompt_message)]
                 )
-                context.terminate = True
-                return
+                raise MiddlewareTermination
+        except MiddlewareTermination:
+            raise
         except PurviewPaymentRequiredError as ex:
             logger.error(f"Purview payment required error in policy pre-check: {ex}")
             if not self._settings.ignore_payment_required:
@@ -151,8 +152,9 @@ async def process(
 
                 blocked_message = ChatMessage(role="system", text=self._settings.blocked_prompt_message)
                 context.result = ChatResponse(messages=[blocked_message])
-                context.terminate = True
-                return
+                raise MiddlewareTermination
+        except MiddlewareTermination:
+            raise
         except PurviewPaymentRequiredError as ex:
             logger.error(f"Purview payment required error in policy pre-check: {ex}")
             if not self._settings.ignore_payment_required:

python/packages/purview/tests/test_chat_middleware.py

@@ -5,7 +5,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from agent_framework import ChatContext, ChatMessage, Role
+from agent_framework import ChatContext, ChatMessage, MiddlewareTermination, Role
 from azure.core.credentials import AccessToken
 
 from agent_framework_purview import PurviewChatPolicyMiddleware, PurviewSettings
@@ -71,8 +71,8 @@ async def test_blocks_prompt(self, middleware: PurviewChatPolicyMiddleware, chat
             async def mock_next(ctx: ChatContext) -> None:  # should not run
                 raise AssertionError("next should not be called when prompt blocked")
 
-            await middleware.process(chat_context, mock_next)
-            assert chat_context.terminate
+            with pytest.raises(MiddlewareTermination):
+                await middleware.process(chat_context, mock_next)
             assert chat_context.result
             assert hasattr(chat_context.result, "messages")
             msg = chat_context.result.messages[0]

python/packages/purview/tests/test_middleware.py

@@ -5,7 +5,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from agent_framework import AgentResponse, AgentRunContext, ChatMessage, Role
+from agent_framework import AgentResponse, AgentRunContext, ChatMessage, MiddlewareTermination, Role
 from azure.core.credentials import AccessToken
 
 from agent_framework_purview import PurviewPolicyMiddleware, PurviewSettings
@@ -79,11 +79,11 @@ async def mock_next(ctx: AgentRunContext) -> None:
                 nonlocal next_called
                 next_called = True
 
-            await middleware.process(context, mock_next)
+            with pytest.raises(MiddlewareTermination):
+                await middleware.process(context, mock_next)
 
             assert not next_called
             assert context.result is not None
-            assert context.terminate
             assert len(context.result.messages) == 1
             assert context.result.messages[0].role == Role.SYSTEM
             assert "blocked by policy" in context.result.messages[0].text.lower()