From 0121ee52c9d64643b3e4b54b2fa0afbc99aeddab Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:12:23 +0000 Subject: [PATCH 1/7] Purview: Add streaming buffering and improve docs --- .../PurviewExtensions.cs | 14 + .../src/Microsoft.Agents.AI.Purview/README.md | 54 ++- python/packages/purview/README.md | 51 ++- .../agent_framework_purview/_middleware.py | 323 +++++++++++++----- .../tests/purview/test_chat_middleware.py | 85 ++++- .../purview/tests/purview/test_middleware.py | 61 +++- 6 files changed, 488 insertions(+), 100 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewExtensions.cs index 2458db94a3a..15748a9cc27 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewExtensions.cs @@ -56,6 +56,12 @@ private static PurviewWrapper CreateWrapper(TokenCredential tokenCredential, Pur /// The logger to use for logging. /// The distributed cache to use for caching Purview responses. An in memory cache will be used if this is null. /// The updated + /// + /// This evaluates the run boundary: the messages passed to the agent, and the agent's final response. + /// Content produced while the run executes, such as context provider output or a function call the + /// model returns, is only visible in the final response, after any function has already been invoked. + /// Use the overload to evaluate every model round trip. + /// public static AIAgentBuilder WithPurview(this AIAgentBuilder builder, TokenCredential tokenCredential, PurviewSettings purviewSettings, ILogger? logger = null, IDistributedCache? cache = null) { PurviewWrapper purviewWrapper = CreateWrapper(tokenCredential, purviewSettings, logger, cache); @@ -71,6 +77,14 @@ public static AIAgentBuilder WithPurview(this AIAgentBuilder builder, TokenCrede /// The logger to use for logging. /// The distributed cache to use for caching Purview responses. An in memory cache will be used if this is null. /// The updated + /// + /// This evaluates every model round trip, including a function call the model returns before that + /// function is invoked. That holds only while this client is composed below + /// . Because applies the first + /// Use outermost, call WithPurview after UseFunctionInvocation, or omit + /// UseFunctionInvocation and let the agent add it above this client. Calling WithPurview + /// before UseFunctionInvocation places this client outermost and is not corrected automatically. + /// public static ChatClientBuilder WithPurview(this ChatClientBuilder builder, TokenCredential tokenCredential, PurviewSettings purviewSettings, ILogger? logger = null, IDistributedCache? cache = null) { PurviewWrapper purviewWrapper = CreateWrapper(tokenCredential, purviewSettings, logger, cache); diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/README.md b/dotnet/src/Microsoft.Agents.AI.Purview/README.md index e92d1380e13..6d7b5f3f74a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/README.md +++ b/dotnet/src/Microsoft.Agents.AI.Purview/README.md @@ -191,6 +191,22 @@ var settings = new PurviewSettings("My Sample App") ### Selecting Agent vs Chat Middleware +Both middlewares apply the same policy logic, but they are handed different content, so they do not evaluate the same thing. **Prefer the chat middleware for data loss prevention**; use the agent middleware when a single check at the run boundary is what you want. + +| | Agent middleware | Chat middleware | +| --- | --- | --- | +| Caller's input messages | evaluated | evaluated | +| Context provider output (for example retrieval results or memory) | not evaluated | evaluated | +| Conversation history replayed into the request | not evaluated | evaluated | +| A model's tool call, before the tool executes | not evaluated | evaluated | +| Tool results | evaluated at the end of the run | evaluated on the next request | +| Final response | evaluated | evaluated | +| How often it evaluates | once per run | once per model request | + +The agent middleware receives the messages passed to the agent. Content added while the run executes — context provider output, replayed history, and the function calls and results produced by `FunctionInvokingChatClient` — is assembled downstream of it, so it is visible only in the final response, after any tool has already run. + +The chat middleware receives the fully prepared request on every model round trip, provided it is composed **below** `FunctionInvokingChatClient` (see [Composition order](#composition-order)). A function call returned by the model is then evaluated before that function executes, and its result is evaluated on the following round trip. + Use the agent middleware when you already have / want the full agent pipeline: ``` csharp @@ -217,7 +233,36 @@ IChatClient client = new OpenAIClient( .Build(); ``` -The policy logic is identical; the only difference is the hook point in the pipeline. +Both middlewares can be attached at the same time. The chat middleware then evaluates each model round trip and the agent middleware evaluates the run boundary. + +### Composition order + +`WithPurview` on an `IChatClient` must end up **below** `FunctionInvokingChatClient`, otherwise the middleware sees only the run boundary and function calls execute before it evaluates them. + +``` csharp +// Correct: the agent adds function invocation above the Purview-wrapped client. +IChatClient client = chatClient + .AsBuilder() + .WithPurview(credential, settings) + .Build(); +AIAgent agent = client.AsAIAgent("You are a helpful assistant."); + +// Correct: the first Use is outermost, so function invocation stays above Purview. +IChatClient client = chatClient + .AsBuilder() + .UseFunctionInvocation() + .WithPurview(credential, settings) + .Build(); + +// Not recommended: Purview is outermost, so it evaluates only the completed run. +IChatClient client = chatClient + .AsBuilder() + .WithPurview(credential, settings) + .UseFunctionInvocation() + .Build(); +``` + +The last form is not corrected automatically. When a chat client is turned into an agent, the framework only adds a `FunctionInvokingChatClient` if the pipeline does not already expose one, and the Purview client delegates that lookup inward — so an existing instance is found and the ordering is preserved as written. --- @@ -309,3 +354,10 @@ Every content item on a message is submitted for evaluation, not just its text: as Purview binary content, and `FunctionCallContent`, `FunctionResultContent` and other structured content are serialized to text. Only `UsageContent` is skipped, because it carries token counts rather than user data. + +### References are not dereferenced + +Purview classifies the content it is handed; a reference to content is not the content. A `UriContent`, +a hosted file reference, or a link nested inside a `FunctionResultContent` is submitted as the reference +itself, and the bytes it points at are never fetched or evaluated. A host that needs those bytes +evaluated must resolve them and pass the resolved content through the middleware. diff --git a/python/packages/purview/README.md b/python/packages/purview/README.md index 3b596a75c13..7ab39974c3f 100644 --- a/python/packages/purview/README.md +++ b/python/packages/purview/README.md @@ -215,6 +215,31 @@ settings = PurviewSettings( ### Selecting Agent vs Chat Middleware +Both middlewares apply the same policy logic, but they are handed different content, so +they do not evaluate the same thing. **Prefer the chat middleware for data loss +prevention**; use the agent middleware when a single check at the run boundary is what +you want. + +| | Agent middleware | Chat middleware | +|---|---|---| +| Caller's input messages | evaluated | evaluated | +| Context provider output (for example retrieval results or memory) | not evaluated | evaluated | +| Conversation history replayed into the request | not evaluated | evaluated | +| A model's tool call, before the tool executes | not evaluated | evaluated | +| Tool results | evaluated at the end of the run | evaluated on the next request | +| Final response | evaluated | evaluated | +| How often it evaluates | once per run | once per model request | + +The agent middleware receives the messages passed to `Agent.run()`. Content added while +the run executes — context provider output, replayed history, and the tool calls and +results produced by the function calling loop — is assembled downstream of it, so it is +visible only in the final response, after any tool has already run. + +The chat middleware sits below the function calling loop and receives the fully prepared +request on every model round trip. A tool call returned by the model is therefore +evaluated before that tool executes, and its result is evaluated on the following round +trip. + Use the agent middleware when you already have / want the full agent pipeline: ```python @@ -256,14 +281,15 @@ client = OpenAIChatCompletionClient( agent = Agent(client=client, instructions="You are helpful.") ``` -The policy logic is identical; the difference is only the hook point in the pipeline. +Both middlewares can be attached at the same time. The chat middleware then evaluates +each model round trip and the agent middleware evaluates the run boundary. --- ## Middleware Lifecycle 1. **Before agent execution** (`prompt phase`): all `context.messages` are evaluated. - - If no valid user_id is found, processing is skipped (no policy evaluation) + - A valid user_id is required; if none can be resolved, the request fails rather than proceeding unevaluated - Protection scopes are retrieved (with caching) - Applicable scopes are checked to determine execution mode - In inline mode: content is evaluated immediately @@ -272,7 +298,7 @@ The policy logic is identical; the difference is only the hook point in the pipe 3. **After successful agent execution** (`response phase`): the produced messages are evaluated using the same user_id from the prompt phase. 4. **If blocked**: result messages are replaced with a blocking notice. -The user identifier is discovered from `Message.additional_properties['user_id']` during the prompt phase and reused for the response phase, ensuring both evaluations map consistently to the same user. If no user_id is present, policy evaluation is skipped entirely. +The user identifier is discovered from `Message.additional_properties['user_id']` during the prompt phase and reused for the response phase, ensuring both evaluations map consistently to the same user. See [Fail-closed behaviour](#fail-closed-behaviour) for what happens when no user identifier can be resolved. You can customize the blocking messages using the `blocked_prompt_message` and `blocked_response_message` fields in `PurviewSettings`. For more advanced scenarios, you can wrap the middleware or post-process `context.result` in later middleware. @@ -350,12 +376,29 @@ sent as Purview binary content, and function calls, function results and other s serialized to text. Only `usage` content is skipped, because it carries token counts rather than user data. +### References are not dereferenced + +Purview classifies the content it is handed; a reference to content is not the content. A URI, a hosted +file reference, or a link nested inside a tool result is submitted as the reference itself, and the bytes +it points at are never fetched or evaluated. A host that needs those bytes evaluated must resolve them +and pass the resolved content through the middleware. + +### Streaming responses + +A streamed response is buffered in full and evaluated before any update is released, so it receives the +same evaluation as a non-streaming response. Content is therefore not delivered incrementally while +either middleware is attached: the first update is released only once the whole response has been +evaluated. + +This uses an experimental core API, so attaching the middleware to a streaming call emits an +`ExperimentalWarning`. + --- ## Notes - **User Identification**: When the configured credential resolves to a user token, that token's `user_id` is used for per-user policy scoping. For app-token credentials, provide a `user_id` per request (e.g. in `Message(..., additional_properties={"user_id": ""})`). If no user_id can be provided or inferred, the request fails rather than proceeding unevaluated — see [Security Considerations](#security-considerations). - **Blocking Messages**: Can be customized via `blocked_prompt_message` and `blocked_response_message` in `PurviewSettings`. By default, they are "Prompt blocked by policy" and "Response blocked by policy" respectively. -- **Streaming Responses**: Post-response policy evaluation presently applies only to non-streaming chat responses. +- **Streaming Responses**: Streamed responses are buffered and evaluated in full before any update is released, so content is not delivered incrementally while the middleware is attached — see [Streaming responses](#streaming-responses). - **Error Handling**: Use `ignore_exceptions` and `ignore_payment_required` settings for graceful degradation. When enabled, errors are logged but don't fail the request. - **Caching**: Protection scopes responses and 402 errors are cached by default with a 4-hour TTL. Cache is automatically invalidated when protection scope state changes. - **Cold-cache parallelization**: On a `ProtectionScopes` cache miss, scopes are refreshed in the background while `ProcessContent` runs in the foreground. diff --git a/python/packages/purview/agent_framework_purview/_middleware.py b/python/packages/purview/agent_framework_purview/_middleware.py index 5cdd2649b15..7069f03d26b 100644 --- a/python/packages/purview/agent_framework_purview/_middleware.py +++ b/python/packages/purview/agent_framework_purview/_middleware.py @@ -1,10 +1,22 @@ # Copyright (c) Microsoft. All rights reserved. import logging -from collections.abc import Awaitable, Callable -from typing import Union - -from agent_framework import AgentContext, AgentMiddleware, ChatContext, ChatMiddleware, MiddlewareTermination +from collections.abc import Awaitable, Callable, Sequence +from typing import Any, Union, cast + +from agent_framework import ( + AgentContext, + AgentMiddleware, + AgentResponse, + AgentResponseUpdate, + ChatContext, + ChatMiddleware, + ChatResponse, + ChatResponseUpdate, + Message, + MiddlewareTermination, + ResponseStream, +) from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -21,11 +33,91 @@ logger = logging.getLogger("agent_framework.purview") +def _agent_updates_from_response(response: AgentResponse[Any]) -> list[AgentResponseUpdate]: + """Rebuild stream updates from an assembled agent response.""" + updates = [ + AgentResponseUpdate( + contents=list(message.contents), + role=message.role, + author_name=message.author_name, + message_id=message.message_id, + response_id=response.response_id, + ) + for message in response.messages + ] + if not updates: + updates = [AgentResponseUpdate(role="assistant", response_id=response.response_id)] + return updates + + +def _chat_updates_from_response(response: ChatResponse[Any]) -> list[ChatResponseUpdate]: + """Rebuild stream updates from an assembled chat response.""" + updates = [ + ChatResponseUpdate( + contents=list(message.contents), + role=cast(Any, message.role), + author_name=message.author_name, + message_id=message.message_id, + response_id=response.response_id, + model=response.model, + ) + for message in response.messages + ] + if not updates: + updates = [ChatResponseUpdate(role="assistant", response_id=response.response_id)] + updates[-1].finish_reason = response.finish_reason + return updates + + +async def _should_block_response( + processor: ScopedContentProcessor, + settings: PurviewSettings, + result: Any, + session_id: str | None, + user_id: str | None, +) -> bool: + """Evaluate the messages of an assembled response, honouring the configured error tolerances. + + Returns ``True`` when policy requires the response to be replaced. Errors are + re-raised unless the corresponding ``ignore_*`` setting is enabled. + """ + try: + messages: Sequence[Message] = result.messages + if not messages: + return False + should_block, _ = await processor.process_messages( + messages, + Activity.DOWNLOAD_TEXT, + session_id=session_id, + user_id=user_id, + ) + return should_block + except PurviewPaymentRequiredError as ex: + logger.error(f"Purview payment required error in policy post-check: {ex}") + if not settings.get("ignore_payment_required", False): + raise + except Exception as ex: + logger.error(f"Error in Purview policy post-check: {ex}") + if not settings.get("ignore_exceptions", False): + raise + return False + + class PurviewPolicyMiddleware(AgentMiddleware): """Agent middleware that enforces Purview policies on prompt and response. Accepts a TokenCredential, AsyncTokenCredential, or callable token provider. + Evaluation covers the caller's input messages on the prompt side and the agent's + final response on the response side. Streamed runs are buffered in full and + evaluated before any update is released, so a streamed run receives the same + evaluation as a non-streaming one; content is therefore not delivered + incrementally while this middleware is attached. + + This middleware evaluates the run boundary. It does not see content added inside + the run, such as context-provider output or a model's tool call before the tool + executes. Use :class:`PurviewChatPolicyMiddleware` where those must be evaluated. + Usage: .. code-block:: python @@ -82,8 +174,6 @@ async def process( context.messages, Activity.UPLOAD_TEXT, session_id=session_id ) if should_block_prompt: - from agent_framework import AgentResponse, Message - msg = self._settings.get("blocked_prompt_message", None) or "Prompt blocked by policy" context.result = AgentResponse( @@ -108,43 +198,74 @@ async def process( await call_next() - try: - # Post (response) check only if we have a normal AgentResponse - # Use the same user_id from the request for the response evaluation - session_id_response = self._get_agent_session_id(context) - if session_id_response is None: - session_id_response = session_id - if context.result and not context.stream: - should_block_response, _ = await self._processor.process_messages( - context.result.messages, # type: ignore[union-attr] - Activity.DOWNLOAD_TEXT, - session_id=session_id_response, - user_id=resolved_user_id, - ) - if should_block_response: - from agent_framework import AgentResponse, Message - - msg = self._settings.get("blocked_response_message", None) or "Response blocked by policy" - - context.result = AgentResponse( - messages=[ - Message( - role="system", - contents=[msg], - ) - ] - ) - else: - # Streaming responses are not supported for post-checks - logger.debug("Streaming responses are not supported for Purview policy post-checks") - except PurviewPaymentRequiredError as ex: - logger.error(f"Purview payment required error in policy post-check: {ex}") - if not self._settings.get("ignore_payment_required", False): - raise - except Exception as ex: - logger.error(f"Error in Purview policy post-check: {ex}") - if not self._settings.get("ignore_exceptions", False): - raise + # Post (response) check. The user id resolved during the prompt check is reused + # so the response is evaluated against the same identity as the request. + session_id_response = self._get_agent_session_id(context) + if session_id_response is None: + session_id_response = session_id + + if isinstance(context.result, ResponseStream): + context.result = self._gate_stream( + cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", context.result), + session_id_response, + resolved_user_id, + ) + elif context.result is not None: + should_block_response = await _should_block_response( + self._processor, + self._settings, + context.result, + session_id_response, + resolved_user_id, + ) + if should_block_response: + context.result = self._blocked_response() + + def _blocked_response(self) -> "AgentResponse[Any]": + """Build the replacement response returned when policy blocks the content.""" + msg = self._settings.get("blocked_response_message", None) or "Response blocked by policy" + return AgentResponse(messages=[Message(role="system", contents=[msg])]) + + def _gate_stream( + self, + inner: "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", + session_id: str | None, + user_id: str | None, + ) -> "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]": + """Buffer a streamed run and evaluate its complete content before anything is released. + + The stream is drained in full, the assembled response is evaluated exactly as a + non-streaming response is, and only then are updates released. A blocked response + is replaced, and the released updates are re-derived from the replacement so the + streamed content cannot diverge from the evaluated content. + """ + + async def _consume() -> tuple[Sequence[AgentResponseUpdate], "AgentResponse[Any]"]: + final = await inner.get_final_response() + return list(inner.updates), final + + async def _gate( + _updates: list[AgentResponseUpdate], final: "AgentResponse[Any]" + ) -> tuple["AgentResponse[Any]", bool]: + should_block_response = await _should_block_response( + self._processor, + self._settings, + final, + session_id, + user_id, + ) + if should_block_response: + return self._blocked_response(), True + return final, False + + return cast( + "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", + cast(Any, ResponseStream).buffered_and_gated( + consume=_consume, + gate=_gate, + rederive=_agent_updates_from_response, + ), + ) class PurviewChatPolicyMiddleware(ChatMiddleware): @@ -155,9 +276,13 @@ class PurviewChatPolicyMiddleware(ChatMiddleware): Behavior: * Pre-chat: evaluates outgoing (user + context) messages as an upload activity and can terminate execution if blocked. - * Post-chat: evaluates the received response messages (streaming is not presently supported) - and can replace them with a blocked message. Uses the same user_id from the request - to ensure consistent user identity throughout the evaluation. + * Post-chat: evaluates the received response messages and can replace them with a + blocked message. Uses the same user_id from the request to ensure consistent user + identity throughout the evaluation. + * Streaming: the response is buffered in full and evaluated before any update is + released, so a streamed response receives the same evaluation as a non-streaming + one. Content is therefore not delivered incrementally while this middleware is + attached. Usage: @@ -193,8 +318,6 @@ async def process( context.messages, Activity.UPLOAD_TEXT, session_id=session_id ) if should_block_prompt: - from agent_framework import ChatResponse, Message - blocked_message = Message( role="system", contents=[self._settings.get("blocked_prompt_message", None) or "Prompt blocked by policy"], @@ -214,36 +337,74 @@ async def process( await call_next() - try: - # Post (response) evaluation only if non-streaming and we have messages result shape - # Use the same user_id from the request for the response evaluation - session_id_response = context.options.get("conversation_id") if context.options else None - if session_id_response is None: - session_id_response = session_id - if context.result and not context.stream: - result_obj = context.result - messages = getattr(result_obj, "messages", None) - if messages: - should_block_response, _ = await self._processor.process_messages( - messages, Activity.DOWNLOAD_TEXT, session_id=session_id_response, user_id=resolved_user_id - ) - if should_block_response: - from agent_framework import ChatResponse, Message - - blocked_message = Message( - role="system", - contents=[ - self._settings.get("blocked_response_message", None) or "Response blocked by policy" - ], - ) - context.result = ChatResponse(messages=[blocked_message]) - else: - logger.debug("Streaming responses are not supported for Purview policy post-checks") - except PurviewPaymentRequiredError as ex: - logger.error(f"Purview payment required error in policy post-check: {ex}") - if not self._settings.get("ignore_payment_required", False): - raise - except Exception as ex: - logger.error(f"Error in Purview policy post-check: {ex}") - if not self._settings.get("ignore_exceptions", False): - raise + # Post (response) evaluation. The user id resolved during the prompt check is + # reused so the response is evaluated against the same identity as the request. + session_id_response = context.options.get("conversation_id") if context.options else None + if session_id_response is None: + session_id_response = session_id + + if isinstance(context.result, ResponseStream): + context.result = self._gate_stream( + cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", context.result), + session_id_response, + resolved_user_id, + ) + elif context.result is not None: + should_block_response = await _should_block_response( + self._processor, + self._settings, + context.result, + session_id_response, + resolved_user_id, + ) + if should_block_response: + context.result = self._blocked_response() + + def _blocked_response(self) -> "ChatResponse[Any]": + """Build the replacement response returned when policy blocks the content.""" + blocked_message = Message( + role="system", + contents=[self._settings.get("blocked_response_message", None) or "Response blocked by policy"], + ) + return ChatResponse(messages=[blocked_message]) + + def _gate_stream( + self, + inner: "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", + session_id: str | None, + user_id: str | None, + ) -> "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]": + """Buffer a streamed response and evaluate its complete content before anything is released. + + The stream is drained in full, the assembled response is evaluated exactly as a + non-streaming response is, and only then are updates released. A blocked response + is replaced, and the released updates are re-derived from the replacement so the + streamed content cannot diverge from the evaluated content. + """ + + async def _consume() -> tuple[Sequence[ChatResponseUpdate], "ChatResponse[Any]"]: + final = await inner.get_final_response() + return list(inner.updates), final + + async def _gate( + _updates: list[ChatResponseUpdate], final: "ChatResponse[Any]" + ) -> tuple["ChatResponse[Any]", bool]: + should_block_response = await _should_block_response( + self._processor, + self._settings, + final, + session_id, + user_id, + ) + if should_block_response: + return self._blocked_response(), True + return final, False + + return cast( + "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", + cast(Any, ResponseStream).buffered_and_gated( + consume=_consume, + gate=_gate, + rederive=_chat_updates_from_response, + ), + ) diff --git a/python/packages/purview/tests/purview/test_chat_middleware.py b/python/packages/purview/tests/purview/test_chat_middleware.py index 3c6056d3443..7777b92705d 100644 --- a/python/packages/purview/tests/purview/test_chat_middleware.py +++ b/python/packages/purview/tests/purview/test_chat_middleware.py @@ -1,12 +1,21 @@ # Copyright (c) Microsoft. All rights reserved. """Tests for Purview chat middleware.""" +from collections.abc import AsyncIterator from dataclasses import dataclass from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest -from agent_framework import ChatContext, ChatResponse, Message, MiddlewareTermination +from agent_framework import ( + ChatContext, + ChatResponse, + ChatResponseUpdate, + Content, + Message, + MiddlewareTermination, + ResponseStream, +) from azure.core.credentials import AccessToken from agent_framework_purview import PurviewChatPolicyMiddleware, PurviewSettings @@ -100,7 +109,8 @@ async def mock_next() -> None: assert first_msg.role in ("system", "system") assert "blocked" in first_msg.text.lower() - async def test_streaming_skips_post_check(self, middleware: PurviewChatPolicyMiddleware) -> None: + async def test_streaming_response_is_evaluated_and_blocked(self, middleware: PurviewChatPolicyMiddleware) -> None: + """Streamed content is evaluated in full and replaced when policy blocks it.""" client = DummyChatClient() chat_options = MagicMock() chat_options.model = "test-model" @@ -110,13 +120,59 @@ async def test_streaming_skips_post_check(self, middleware: PurviewChatPolicyMid options=chat_options, stream=True, ) - with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc: + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text="confidential")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ) as mock_proc: + + async def mock_next() -> None: + streaming_context.result = cast(Any, ResponseStream(updates(), finalizer=ChatResponse.from_updates)) + + await middleware.process(streaming_context, mock_next) + released = [update async for update in cast(Any, streaming_context.result)] + + assert mock_proc.call_count == 2 + assert mock_proc.call_args_list[1][0][1] == Activity.DOWNLOAD_TEXT + released_text = "".join(update.text for update in released) + assert "confidential" not in released_text + assert "blocked" in released_text.lower() + + async def test_streaming_response_passes_when_allowed(self, middleware: PurviewChatPolicyMiddleware) -> None: + """Allowed streamed content is released unchanged, reusing the prompt-phase identity.""" + client = DummyChatClient() + chat_options = MagicMock() + chat_options.model = "test-model" + streaming_context = ChatContext( + client=cast(Any, client), + messages=[Message(role="user", contents=["Hello"])], + options=chat_options, + stream=True, + ) + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text="all ")]) + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text="clear")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ) as mock_proc: async def mock_next() -> None: - streaming_context.result = MagicMock() + streaming_context.result = cast(Any, ResponseStream(updates(), finalizer=ChatResponse.from_updates)) await middleware.process(streaming_context, mock_next) - assert mock_proc.call_count == 1 + released = [update async for update in cast(Any, streaming_context.result)] + + assert mock_proc.call_count == 2 + assert mock_proc.call_args_list[1].kwargs["user_id"] == "user-123" + assert "".join(update.text for update in released) == "all clear" async def test_chat_middleware_handles_post_check_exception( self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext @@ -263,19 +319,32 @@ async def mock_next() -> None: # Next should have been called assert context.result is not None - async def test_chat_middleware_handles_result_without_messages_attribute( + async def test_chat_middleware_result_without_messages_attribute_is_not_silently_allowed( self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext ) -> None: - """Test middleware handles result that doesn't have messages attribute.""" + """A result shape that cannot be evaluated surfaces an error rather than passing unchecked.""" with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")): async def mock_next() -> None: # Set result to something without messages attribute chat_context.result = cast(Any, "Some string result") + with pytest.raises(AttributeError): + await middleware.process(chat_context, mock_next) + + async def test_chat_middleware_result_without_messages_attribute_tolerated_when_ignoring_exceptions( + self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext + ) -> None: + """With ignore_exceptions enabled, an unevaluatable result is logged and left unchanged.""" + middleware._settings["ignore_exceptions"] = True + + with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")): + + async def mock_next() -> None: + chat_context.result = cast(Any, "Some string result") + await middleware.process(chat_context, mock_next) - # Should not crash, result should be unchanged assert chat_context.result == "Some string result" async def test_chat_middleware_with_ignore_exceptions(self, mock_credential: AsyncMock) -> None: diff --git a/python/packages/purview/tests/purview/test_middleware.py b/python/packages/purview/tests/purview/test_middleware.py index 5b624367696..e93a9246ba2 100644 --- a/python/packages/purview/tests/purview/test_middleware.py +++ b/python/packages/purview/tests/purview/test_middleware.py @@ -2,11 +2,21 @@ """Tests for Purview middleware.""" +from collections.abc import AsyncIterator from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest -from agent_framework import AgentContext, AgentResponse, AgentSession, Message, MiddlewareTermination +from agent_framework import ( + AgentContext, + AgentResponse, + AgentResponseUpdate, + AgentSession, + Content, + Message, + MiddlewareTermination, + ResponseStream, +) from azure.core.credentials import AccessToken from agent_framework_purview import PurviewPolicyMiddleware, PurviewSettings @@ -156,21 +166,60 @@ async def mock_next() -> None: # Second call (post-check) should be DOWNLOAD_TEXT for agent response assert mock_process.call_args_list[1][0][1] == Activity.DOWNLOAD_TEXT - async def test_middleware_streaming_skips_post_check( + async def test_middleware_streaming_response_is_evaluated_and_blocked( self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock ) -> None: - """Test that streaming results skip post-check evaluation.""" + """Streamed content is evaluated in full and replaced when policy blocks it.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) context.stream = True - with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc: + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="confidential")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ) as mock_proc: + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) + + await middleware.process(context, mock_next) + released = [update async for update in cast(Any, context.result)] + + assert mock_proc.call_count == 2 + assert mock_proc.call_args_list[1][0][1] == Activity.DOWNLOAD_TEXT + released_text = "".join(update.text for update in released) + assert "confidential" not in released_text + assert "blocked" in released_text.lower() + + async def test_middleware_streaming_response_passes_when_allowed( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """Allowed streamed content is released unchanged, reusing the prompt-phase identity.""" + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="all ")]) + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="clear")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ) as mock_proc: async def mock_next() -> None: - context.result = AgentResponse(messages=[Message(role="assistant", contents=["streaming"])]) + context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) await middleware.process(context, mock_next) + released = [update async for update in cast(Any, context.result)] - assert mock_proc.call_count == 1 + assert mock_proc.call_count == 2 + assert mock_proc.call_args_list[1].kwargs["user_id"] == "user-123" + assert "".join(update.text for update in released) == "all clear" async def test_middleware_payment_required_in_pre_check_raises_by_default( self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock From 161a56b375030e6c49355cf67c15188a519ed3b4 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Wed, 23 Sep 2026 18:07:50 +0000 Subject: [PATCH 2/7] Address PR comments --- .../src/Microsoft.Agents.AI.Purview/README.md | 15 +- .../core/agent_framework/_agent_hooks.py | 62 +++----- .../packages/core/agent_framework/_types.py | 114 ++++++++++++++ python/packages/core/tests/core/test_types.py | 84 ++++++++++ python/packages/core/tests/test_types.py | 1 + python/packages/purview/README.md | 15 +- .../agent_framework_purview/_middleware.py | 86 +++++------ .../tests/purview/test_chat_middleware.py | 118 ++++++++++++++ .../purview/tests/purview/test_middleware.py | 145 ++++++++++++++++++ 9 files changed, 540 insertions(+), 100 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/README.md b/dotnet/src/Microsoft.Agents.AI.Purview/README.md index 6d7b5f3f74a..59c0dea9a33 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/README.md +++ b/dotnet/src/Microsoft.Agents.AI.Purview/README.md @@ -355,9 +355,14 @@ as Purview binary content, and `FunctionCallContent`, `FunctionResultContent` an content are serialized to text. Only `UsageContent` is skipped, because it carries token counts rather than user data. -### References are not dereferenced +### Remote references are not dereferenced -Purview classifies the content it is handed; a reference to content is not the content. A `UriContent`, -a hosted file reference, or a link nested inside a `FunctionResultContent` is submitted as the reference -itself, and the bytes it points at are never fetched or evaluated. A host that needs those bytes -evaluated must resolve them and pass the resolved content through the middleware. +Purview classifies the content it is handed; a reference to content is not the content. + +Content that carries its own bytes is evaluated as bytes: `DataContent` is submitted as Purview binary +content. + +A *remote* reference is not. A `UriContent`, a hosted file reference, or a link nested inside a +`FunctionResultContent` is submitted as the reference itself, and the bytes it points at are never +fetched or evaluated. A host that needs those bytes evaluated must resolve them and pass the resolved +content through the middleware. diff --git a/python/packages/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py index bf2f2aaaa14..4814ce39eb1 100644 --- a/python/packages/core/agent_framework/_agent_hooks.py +++ b/python/packages/core/agent_framework/_agent_hooks.py @@ -826,42 +826,6 @@ def write_back(response: AgentResponse[Any], before_content: Any, after: Any) -> # region Enforcement helpers -def _chat_updates_from_response(response: ChatResponse[Any]) -> list[ChatResponseUpdate]: - """Re-derive stream updates from a (transformed) assembled chat response.""" - updates = [ - ChatResponseUpdate( - contents=list(message.contents), - role=cast(Any, message.role), - author_name=message.author_name, - message_id=message.message_id, - response_id=response.response_id, - model=response.model, - ) - for message in response.messages - ] - if not updates: - updates = [ChatResponseUpdate(role="assistant", response_id=response.response_id)] - updates[-1].finish_reason = response.finish_reason - return updates - - -def _agent_updates_from_response(response: AgentResponse[Any]) -> list[AgentResponseUpdate]: - """Re-derive stream updates from a (transformed) assembled agent response.""" - updates = [ - AgentResponseUpdate( - contents=list(message.contents), - role=message.role, - author_name=message.author_name, - message_id=message.message_id, - response_id=response.response_id, - ) - for message in response.messages - ] - if not updates: - updates = [AgentResponseUpdate(role="assistant", response_id=response.response_id)] - return updates - - def _normalized_tools(tools: Any, *, point: str) -> list[Any]: """Normalize a tools value for projection; empty (with a warning) when it cannot be. @@ -1290,6 +1254,9 @@ async def _consume() -> tuple[Sequence[AgentResponseUpdate], AgentResponse[Any]] raise finally: _RUN_STATE.reset(run_token) + # Cancellation or failure part-way through the drain still has to release + # the inner stream, which the cleanup hook below would never reach. + await inner.close() return list(inner.updates), final async def _gate( @@ -1318,12 +1285,15 @@ async def _gate( await self._emit_shutdown(state, "error") raise - return cast( + gated = cast( "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", cast(Any, ResponseStream).buffered_and_gated( - consume=_consume, gate=_gate, rederive=_agent_updates_from_response + consume=_consume, gate=_gate, rederive=AgentResponse.to_updates ), ) + # Closing the gated stream before it is ever pulled never reaches ``_consume``, so + # the inner stream is released through a cleanup hook rather than from inside it. + return gated.with_cleanup_hook(inner.close) class _AgentHooksChatMiddleware(_AgentHooksMiddlewareBase, ChatMiddleware): @@ -1422,8 +1392,13 @@ def _gated_chat_stream( """ async def _consume() -> tuple[Sequence[ChatResponseUpdate], ChatResponse[Any]]: - with gate_handle: - response = await inner.get_final_response() + try: + with gate_handle: + response = await inner.get_final_response() + finally: + # Cancellation or failure part-way through the drain still has to release + # the inner stream, which the cleanup hook below would never reach. + await inner.close() return list(inner.updates), response async def _gate(_updates: list[ChatResponseUpdate], final: ChatResponse[Any]) -> tuple[ChatResponse[Any], bool]: @@ -1444,12 +1419,15 @@ async def _gate(_updates: list[ChatResponseUpdate], final: ChatResponse[Any]) -> await gate_handle.flush() return final, changed - return cast( + gated = cast( "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", cast(Any, ResponseStream).buffered_and_gated( - consume=_consume, gate=_gate, rederive=_chat_updates_from_response + consume=_consume, gate=_gate, rederive=ChatResponse.to_updates ), ) + # Closing the gated stream before it is ever pulled never reaches ``_consume``, so + # the inner stream is released through a cleanup hook rather than from inside it. + return gated.with_cleanup_hook(inner.close) class _AgentHooksFunctionMiddleware(_AgentHooksMiddlewareBase, FunctionMiddleware): diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 17dba475e3e..62b08be46ca 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2126,6 +2126,26 @@ def _process_update(response: ChatResponse | AgentResponse, update: ChatResponse response.continuation_token = update.continuation_token +def _apply_response_tail_to_update( + response: ChatResponse[Any] | AgentResponse[Any], + update: ChatResponseUpdate | AgentResponseUpdate, +) -> None: + """Carry the response-level fields that belong on the final update of a stream. + + ``finish_reason``, ``continuation_token``, ``additional_properties`` and usage are properties + of the response as a whole rather than of any single message, so they ride on the last update. + This mirrors how :func:`_process_update` reads them back. + """ + update.finish_reason = response.finish_reason + update.continuation_token = response.continuation_token + if response.additional_properties: + merged = dict(update.additional_properties) if update.additional_properties else {} + merged.update(response.additional_properties) + update.additional_properties = merged + if response.usage_details is not None: + update.contents.append(Content.from_usage(response.usage_details)) + + def _merge_function_call_content(message: Message, content: Content) -> None: """Merge a streamed function_call chunk into the in-progress call it belongs to. @@ -2651,6 +2671,54 @@ async def from_update_generator( _finalize_response(msg) return msg + def to_updates(self) -> list[ChatResponseUpdate]: + """Split this response into the stream updates it would have been assembled from. + + This is the inverse of :meth:`from_updates`: each message becomes one update, and the + response-level fields are carried on those updates so that re-assembling them reproduces + this response. Usage is carried as usage content on the final update, because + :meth:`from_updates` reads usage from an update's contents rather than from a field. + + This is useful where a complete response has to be emitted as a stream, for example when + middleware buffers a stream, decides on the complete content, and then releases it. + + Returns: + The updates representing this response; always at least one. + + Example: + .. code-block:: python + + from agent_framework import ChatResponse + + response = ChatResponse.from_updates(updates) + assert response.text == ChatResponse.from_updates(response.to_updates()).text + """ + updates = [ + ChatResponseUpdate( + contents=list(message.contents), + role=cast(Any, message.role), + author_name=message.author_name, + message_id=message.message_id, + response_id=self.response_id, + conversation_id=self.conversation_id, + model=self.model, + created_at=self.created_at, + ) + for message in self.messages + ] + if not updates: + updates = [ + ChatResponseUpdate( + role="assistant", + response_id=self.response_id, + conversation_id=self.conversation_id, + model=self.model, + created_at=self.created_at, + ) + ] + _apply_response_tail_to_update(self, updates[-1]) + return updates + @property def text(self) -> str: """Returns the concatenated text of all messages in the response.""" @@ -3063,6 +3131,52 @@ async def from_update_generator( _finalize_response(msg) return msg + def to_updates(self) -> list[AgentResponseUpdate]: + """Split this response into the stream updates it would have been assembled from. + + This is the inverse of :meth:`from_updates`: each message becomes one update, and the + response-level fields are carried on those updates so that re-assembling them reproduces + this response. Usage is carried as usage content on the final update, because + :meth:`from_updates` reads usage from an update's contents rather than from a field. + + This is useful where a complete response has to be emitted as a stream, for example when + middleware buffers a stream, decides on the complete content, and then releases it. + + Returns: + The updates representing this response; always at least one. + + Example: + .. code-block:: python + + from agent_framework import AgentResponse + + response = AgentResponse.from_updates(updates) + assert response.text == AgentResponse.from_updates(response.to_updates()).text + """ + updates = [ + AgentResponseUpdate( + contents=list(message.contents), + role=message.role, + author_name=message.author_name, + message_id=message.message_id, + response_id=self.response_id, + agent_id=self.agent_id, + created_at=self.created_at, + ) + for message in self.messages + ] + if not updates: + updates = [ + AgentResponseUpdate( + role="assistant", + response_id=self.response_id, + agent_id=self.agent_id, + created_at=self.created_at, + ) + ] + _apply_response_tail_to_update(self, updates[-1]) + return updates + def __str__(self) -> str: return self.text diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index a6797b30449..7583197a370 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -1496,6 +1496,49 @@ def test_chat_response_from_streaming_updates_parses_final_assistant_message() - assert response.value.response == "Hello" +def test_chat_response_to_updates_carries_response_level_fields() -> None: + response = ChatResponse( + messages=[ + Message(role="assistant", contents=[Content.from_text("first")], message_id="m1"), + Message(role="tool", contents=[Content.from_text("second")], message_id="m2"), + ], + response_id="resp-1", + conversation_id="conv-1", + model="model-1", + created_at="2024-01-01T00:00:00Z", + finish_reason="stop", + continuation_token="token-1", + usage_details=UsageDetails(input_token_count=3, output_token_count=5), + additional_properties={"custom": "value"}, + ) + + round_tripped = ChatResponse.from_updates(response.to_updates()) + + assert round_tripped.response_id == "resp-1" + assert round_tripped.conversation_id == "conv-1" + assert round_tripped.model == "model-1" + assert round_tripped.created_at == response.created_at + assert round_tripped.finish_reason == "stop" + assert round_tripped.continuation_token == "token-1" + assert round_tripped.additional_properties["custom"] == "value" + assert round_tripped.usage_details is not None + assert round_tripped.usage_details["input_token_count"] == 3 + assert round_tripped.usage_details["output_token_count"] == 5 + assert [message.role for message in round_tripped.messages] == ["assistant", "tool"] + assert round_tripped.text == response.text + + +def test_chat_response_to_updates_without_messages_still_carries_fields() -> None: + response = ChatResponse(messages=[], response_id="resp-1", model="model-1", finish_reason="stop") + + updates = response.to_updates() + + assert len(updates) == 1 + assert updates[0].response_id == "resp-1" + assert updates[0].model == "model-1" + assert updates[0].finish_reason == "stop" + + # region ToolMode @@ -1815,6 +1858,47 @@ def test_agent_run_response_from_updates_uses_last_non_none_agent_id() -> None: assert response.agent_id == "source-agent" +def test_agent_run_response_to_updates_carries_response_level_fields() -> None: + response = AgentResponse( + messages=[ + Message(role="assistant", contents=[Content.from_text("first")], message_id="m1"), + Message(role="tool", contents=[Content.from_text("second")], message_id="m2"), + ], + response_id="resp-1", + agent_id="agent-1", + created_at="2024-01-01T00:00:00Z", + finish_reason="stop", + continuation_token="token-1", + usage_details=UsageDetails(input_token_count=3, output_token_count=5), + additional_properties={"custom": "value"}, + ) + + round_tripped = AgentResponse.from_updates(response.to_updates()) + + assert round_tripped.response_id == "resp-1" + assert round_tripped.agent_id == "agent-1" + assert round_tripped.created_at == response.created_at + assert round_tripped.finish_reason == "stop" + assert round_tripped.continuation_token == "token-1" + assert round_tripped.additional_properties["custom"] == "value" + assert round_tripped.usage_details is not None + assert round_tripped.usage_details["input_token_count"] == 3 + assert round_tripped.usage_details["output_token_count"] == 5 + assert [message.role for message in round_tripped.messages] == ["assistant", "tool"] + assert round_tripped.text == response.text + + +def test_agent_run_response_to_updates_without_messages_still_carries_fields() -> None: + response = AgentResponse(messages=[], response_id="resp-1", agent_id="agent-1", finish_reason="stop") + + updates = response.to_updates() + + assert len(updates) == 1 + assert updates[0].response_id == "resp-1" + assert updates[0].agent_id == "agent-1" + assert updates[0].finish_reason == "stop" + + def test_agent_run_response_str_method(chat_message: Message) -> None: response = AgentResponse(messages=chat_message) assert str(response) == "Hello" diff --git a/python/packages/core/tests/test_types.py b/python/packages/core/tests/test_types.py index 0622fb48e6d..71ab420ce1a 100644 --- a/python/packages/core/tests/test_types.py +++ b/python/packages/core/tests/test_types.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. + class TestPrependInstructionsEmpty: def test_empty_string_instructions_add_no_message(self) -> None: """An unset "" instruction must not inject a contentless system message.""" diff --git a/python/packages/purview/README.md b/python/packages/purview/README.md index 7ab39974c3f..aa3ea9d86e4 100644 --- a/python/packages/purview/README.md +++ b/python/packages/purview/README.md @@ -376,12 +376,17 @@ sent as Purview binary content, and function calls, function results and other s serialized to text. Only `usage` content is skipped, because it carries token counts rather than user data. -### References are not dereferenced +### Remote references are not dereferenced -Purview classifies the content it is handed; a reference to content is not the content. A URI, a hosted -file reference, or a link nested inside a tool result is submitted as the reference itself, and the bytes -it points at are never fetched or evaluated. A host that needs those bytes evaluated must resolve them -and pass the resolved content through the middleware. +Purview classifies the content it is handed; a reference to content is not the content. + +Content that carries its own bytes is evaluated as bytes: a `data:` URI is decoded and the decoded bytes +are submitted as Purview binary content. + +A *remote* reference is not. A `uri` content pointing at a remote location, a hosted file reference, or a +link nested inside a tool result is submitted as the reference itself, and the bytes it points at are +never fetched or evaluated. A host that needs those bytes evaluated must resolve them and pass the +resolved content through the middleware. ### Streaming responses diff --git a/python/packages/purview/agent_framework_purview/_middleware.py b/python/packages/purview/agent_framework_purview/_middleware.py index 7069f03d26b..dff188d4eab 100644 --- a/python/packages/purview/agent_framework_purview/_middleware.py +++ b/python/packages/purview/agent_framework_purview/_middleware.py @@ -33,42 +33,6 @@ logger = logging.getLogger("agent_framework.purview") -def _agent_updates_from_response(response: AgentResponse[Any]) -> list[AgentResponseUpdate]: - """Rebuild stream updates from an assembled agent response.""" - updates = [ - AgentResponseUpdate( - contents=list(message.contents), - role=message.role, - author_name=message.author_name, - message_id=message.message_id, - response_id=response.response_id, - ) - for message in response.messages - ] - if not updates: - updates = [AgentResponseUpdate(role="assistant", response_id=response.response_id)] - return updates - - -def _chat_updates_from_response(response: ChatResponse[Any]) -> list[ChatResponseUpdate]: - """Rebuild stream updates from an assembled chat response.""" - updates = [ - ChatResponseUpdate( - contents=list(message.contents), - role=cast(Any, message.role), - author_name=message.author_name, - message_id=message.message_id, - response_id=response.response_id, - model=response.model, - ) - for message in response.messages - ] - if not updates: - updates = [ChatResponseUpdate(role="assistant", response_id=response.response_id)] - updates[-1].finish_reason = response.finish_reason - return updates - - async def _should_block_response( processor: ScopedContentProcessor, settings: PurviewSettings, @@ -236,12 +200,19 @@ def _gate_stream( The stream is drained in full, the assembled response is evaluated exactly as a non-streaming response is, and only then are updates released. A blocked response - is replaced, and the released updates are re-derived from the replacement so the - streamed content cannot diverge from the evaluated content. + is replaced. The released updates are always re-derived from the response that was + evaluated, so what the caller receives cannot diverge from what policy saw even when + the inner stream's finalizer returns something other than the assembly of its own + updates. """ async def _consume() -> tuple[Sequence[AgentResponseUpdate], "AgentResponse[Any]"]: - final = await inner.get_final_response() + try: + final = await inner.get_final_response() + finally: + # Cancellation or failure part-way through the drain still has to release + # the inner stream, which the cleanup hook below would never reach. + await inner.close() return list(inner.updates), final async def _gate( @@ -256,16 +227,22 @@ async def _gate( ) if should_block_response: return self._blocked_response(), True - return final, False + # Reported as transformed even when the content is allowed through unchanged, so + # the released updates are re-derived from the evaluated response rather than + # replayed from the buffer. + return final, True - return cast( + gated = cast( "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", cast(Any, ResponseStream).buffered_and_gated( consume=_consume, gate=_gate, - rederive=_agent_updates_from_response, + rederive=AgentResponse.to_updates, ), ) + # Closing the gated stream before it is ever pulled never reaches ``_consume``, so + # the inner stream is released through a cleanup hook rather than from inside it. + return gated.with_cleanup_hook(inner.close) class PurviewChatPolicyMiddleware(ChatMiddleware): @@ -378,12 +355,19 @@ def _gate_stream( The stream is drained in full, the assembled response is evaluated exactly as a non-streaming response is, and only then are updates released. A blocked response - is replaced, and the released updates are re-derived from the replacement so the - streamed content cannot diverge from the evaluated content. + is replaced. The released updates are always re-derived from the response that was + evaluated, so what the caller receives cannot diverge from what policy saw even when + the inner stream's finalizer returns something other than the assembly of its own + updates. """ async def _consume() -> tuple[Sequence[ChatResponseUpdate], "ChatResponse[Any]"]: - final = await inner.get_final_response() + try: + final = await inner.get_final_response() + finally: + # Cancellation or failure part-way through the drain still has to release + # the inner stream, which the cleanup hook below would never reach. + await inner.close() return list(inner.updates), final async def _gate( @@ -398,13 +382,19 @@ async def _gate( ) if should_block_response: return self._blocked_response(), True - return final, False + # Reported as transformed even when the content is allowed through unchanged, so + # the released updates are re-derived from the evaluated response rather than + # replayed from the buffer. + return final, True - return cast( + gated = cast( "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", cast(Any, ResponseStream).buffered_and_gated( consume=_consume, gate=_gate, - rederive=_chat_updates_from_response, + rederive=ChatResponse.to_updates, ), ) + # Closing the gated stream before it is ever pulled never reaches ``_consume``, so + # the inner stream is released through a cleanup hook rather than from inside it. + return gated.with_cleanup_hook(inner.close) diff --git a/python/packages/purview/tests/purview/test_chat_middleware.py b/python/packages/purview/tests/purview/test_chat_middleware.py index 7777b92705d..7e30b008018 100644 --- a/python/packages/purview/tests/purview/test_chat_middleware.py +++ b/python/packages/purview/tests/purview/test_chat_middleware.py @@ -174,6 +174,124 @@ async def mock_next() -> None: assert mock_proc.call_args_list[1].kwargs["user_id"] == "user-123" assert "".join(update.text for update in released) == "all clear" + async def test_streaming_releases_the_evaluated_content_not_the_buffered_updates( + self, middleware: PurviewChatPolicyMiddleware + ) -> None: + """The released updates come from the response that was evaluated, not from the buffer.""" + client = DummyChatClient() + chat_options = MagicMock() + chat_options.model = "test-model" + streaming_context = ChatContext( + client=cast(Any, client), + messages=[Message(role="user", contents=["Hello"])], + options=chat_options, + stream=True, + ) + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text="confidential")]) + + def diverging_finalizer(_updates: Any) -> ChatResponse[Any]: + return ChatResponse(messages=[Message(role="assistant", contents=[Content.from_text(text="benign")])]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + + async def mock_next() -> None: + streaming_context.result = cast(Any, ResponseStream(updates(), finalizer=diverging_finalizer)) + + await middleware.process(streaming_context, mock_next) + released = [update async for update in cast(Any, streaming_context.result)] + + released_text = "".join(update.text for update in released) + assert released_text == "benign" + assert "confidential" not in released_text + + async def test_streaming_preserves_response_level_metadata(self, middleware: PurviewChatPolicyMiddleware) -> None: + """Metadata carried by the response survives the buffered stream.""" + client = DummyChatClient() + chat_options = MagicMock() + chat_options.model = "test-model" + streaming_context = ChatContext( + client=cast(Any, client), + messages=[Message(role="user", contents=["Hello"])], + options=chat_options, + stream=True, + ) + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text="all clear")]) + + def finalizer(_updates: Any) -> ChatResponse[Any]: + return ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text(text="all clear")])], + response_id="resp-1", + conversation_id="conv-1", + model="model-1", + continuation_token="token-1", + additional_properties={"custom": "value"}, + ) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + + async def mock_next() -> None: + streaming_context.result = cast(Any, ResponseStream(updates(), finalizer=finalizer)) + + await middleware.process(streaming_context, mock_next) + released = [update async for update in cast(Any, streaming_context.result)] + + assert released[-1].response_id == "resp-1" + assert released[-1].conversation_id == "conv-1" + assert released[-1].model == "model-1" + assert released[-1].continuation_token == "token-1" + assert released[-1].additional_properties is not None + assert released[-1].additional_properties["custom"] == "value" + + async def test_streaming_closes_the_inner_stream_when_never_pulled( + self, middleware: PurviewChatPolicyMiddleware + ) -> None: + """Abandoning the gated stream before the first pull still releases the inner stream.""" + client = DummyChatClient() + chat_options = MagicMock() + chat_options.model = "test-model" + streaming_context = ChatContext( + client=cast(Any, client), + messages=[Message(role="user", contents=["Hello"])], + options=chat_options, + stream=True, + ) + closed = False + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text="all clear")]) + + def _mark_closed() -> None: + nonlocal closed + closed = True + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + + async def mock_next() -> None: + inner = ResponseStream(updates(), finalizer=ChatResponse.from_updates) + inner.with_cleanup_hook(_mark_closed) + streaming_context.result = cast(Any, inner) + + await middleware.process(streaming_context, mock_next) + await cast(Any, streaming_context.result).close() + + assert closed is True + async def test_chat_middleware_handles_post_check_exception( self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext ) -> None: diff --git a/python/packages/purview/tests/purview/test_middleware.py b/python/packages/purview/tests/purview/test_middleware.py index e93a9246ba2..d8a964124f2 100644 --- a/python/packages/purview/tests/purview/test_middleware.py +++ b/python/packages/purview/tests/purview/test_middleware.py @@ -2,6 +2,7 @@ """Tests for Purview middleware.""" +import asyncio from collections.abc import AsyncIterator from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -221,6 +222,150 @@ async def mock_next() -> None: assert mock_proc.call_args_list[1].kwargs["user_id"] == "user-123" assert "".join(update.text for update in released) == "all clear" + async def test_middleware_streaming_releases_the_evaluated_content_not_the_buffered_updates( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """The released updates come from the response that was evaluated, not from the buffer. + + A finalizer is free to return something other than the assembly of its own updates. What + the caller receives has to be what policy actually saw. + """ + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="confidential")]) + + def diverging_finalizer(_updates: Any) -> AgentResponse[Any]: + return AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text(text="benign")])]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=diverging_finalizer)) + + await middleware.process(context, mock_next) + released = [update async for update in cast(Any, context.result)] + + released_text = "".join(update.text for update in released) + assert released_text == "benign" + assert "confidential" not in released_text + + async def test_middleware_streaming_preserves_response_level_metadata( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """Metadata carried by the response survives the buffered stream.""" + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="all clear")]) + + def finalizer(_updates: Any) -> AgentResponse[Any]: + return AgentResponse( + messages=[Message(role="assistant", contents=[Content.from_text(text="all clear")])], + response_id="resp-1", + agent_id="agent-1", + continuation_token="token-1", + additional_properties={"custom": "value"}, + ) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=finalizer)) + + await middleware.process(context, mock_next) + released = [update async for update in cast(Any, context.result)] + + assert released[-1].response_id == "resp-1" + assert released[-1].agent_id == "agent-1" + assert released[-1].continuation_token == "token-1" + assert released[-1].additional_properties is not None + assert released[-1].additional_properties["custom"] == "value" + + async def test_middleware_streaming_closes_the_inner_stream_when_never_pulled( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """Abandoning the gated stream before the first pull still releases the inner stream.""" + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + closed = False + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="all clear")]) + + def _mark_closed() -> None: + nonlocal closed + closed = True + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + + async def mock_next() -> None: + inner = ResponseStream(updates(), finalizer=AgentResponse.from_updates) + inner.with_cleanup_hook(_mark_closed) + context.result = cast(Any, inner) + + await middleware.process(context, mock_next) + await cast(Any, context.result).close() + + assert closed is True + + async def test_middleware_streaming_closes_the_inner_stream_when_the_drain_is_cancelled( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """Cancelling part-way through the drain still releases the inner stream.""" + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + closed = False + started = asyncio.Event() + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="partial")]) + started.set() + await asyncio.Event().wait() + + def _mark_closed() -> None: + nonlocal closed + closed = True + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + + async def mock_next() -> None: + inner = ResponseStream(updates(), finalizer=AgentResponse.from_updates) + inner.with_cleanup_hook(_mark_closed) + context.result = cast(Any, inner) + + await middleware.process(context, mock_next) + + async def drain() -> None: + async for _ in cast(Any, context.result): + pass + + task = asyncio.create_task(drain()) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert closed is True + async def test_middleware_payment_required_in_pre_check_raises_by_default( self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock ) -> None: From 609a1ad00c66b1b5fa4db16b82f48210fce5e1c9 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Wed, 23 Sep 2026 18:24:28 +0000 Subject: [PATCH 3/7] Fix failing type checks --- python/packages/core/tests/core/test_types.py | 8 ++++---- .../purview/tests/purview/test_chat_middleware.py | 4 ++-- python/packages/purview/tests/purview/test_middleware.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 7583197a370..b18fba169b1 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -1507,7 +1507,7 @@ def test_chat_response_to_updates_carries_response_level_fields() -> None: model="model-1", created_at="2024-01-01T00:00:00Z", finish_reason="stop", - continuation_token="token-1", + continuation_token=cast(Any, {"token": "token-1"}), usage_details=UsageDetails(input_token_count=3, output_token_count=5), additional_properties={"custom": "value"}, ) @@ -1519,7 +1519,7 @@ def test_chat_response_to_updates_carries_response_level_fields() -> None: assert round_tripped.model == "model-1" assert round_tripped.created_at == response.created_at assert round_tripped.finish_reason == "stop" - assert round_tripped.continuation_token == "token-1" + assert round_tripped.continuation_token == {"token": "token-1"} assert round_tripped.additional_properties["custom"] == "value" assert round_tripped.usage_details is not None assert round_tripped.usage_details["input_token_count"] == 3 @@ -1868,7 +1868,7 @@ def test_agent_run_response_to_updates_carries_response_level_fields() -> None: agent_id="agent-1", created_at="2024-01-01T00:00:00Z", finish_reason="stop", - continuation_token="token-1", + continuation_token=cast(Any, {"token": "token-1"}), usage_details=UsageDetails(input_token_count=3, output_token_count=5), additional_properties={"custom": "value"}, ) @@ -1879,7 +1879,7 @@ def test_agent_run_response_to_updates_carries_response_level_fields() -> None: assert round_tripped.agent_id == "agent-1" assert round_tripped.created_at == response.created_at assert round_tripped.finish_reason == "stop" - assert round_tripped.continuation_token == "token-1" + assert round_tripped.continuation_token == {"token": "token-1"} assert round_tripped.additional_properties["custom"] == "value" assert round_tripped.usage_details is not None assert round_tripped.usage_details["input_token_count"] == 3 diff --git a/python/packages/purview/tests/purview/test_chat_middleware.py b/python/packages/purview/tests/purview/test_chat_middleware.py index 7e30b008018..cac98fded81 100644 --- a/python/packages/purview/tests/purview/test_chat_middleware.py +++ b/python/packages/purview/tests/purview/test_chat_middleware.py @@ -231,7 +231,7 @@ def finalizer(_updates: Any) -> ChatResponse[Any]: response_id="resp-1", conversation_id="conv-1", model="model-1", - continuation_token="token-1", + continuation_token=cast(Any, {"token": "token-1"}), additional_properties={"custom": "value"}, ) @@ -250,7 +250,7 @@ async def mock_next() -> None: assert released[-1].response_id == "resp-1" assert released[-1].conversation_id == "conv-1" assert released[-1].model == "model-1" - assert released[-1].continuation_token == "token-1" + assert released[-1].continuation_token == {"token": "token-1"} assert released[-1].additional_properties is not None assert released[-1].additional_properties["custom"] == "value" diff --git a/python/packages/purview/tests/purview/test_middleware.py b/python/packages/purview/tests/purview/test_middleware.py index d8a964124f2..932df96bdce 100644 --- a/python/packages/purview/tests/purview/test_middleware.py +++ b/python/packages/purview/tests/purview/test_middleware.py @@ -270,7 +270,7 @@ def finalizer(_updates: Any) -> AgentResponse[Any]: messages=[Message(role="assistant", contents=[Content.from_text(text="all clear")])], response_id="resp-1", agent_id="agent-1", - continuation_token="token-1", + continuation_token=cast(Any, {"token": "token-1"}), additional_properties={"custom": "value"}, ) @@ -288,7 +288,7 @@ async def mock_next() -> None: assert released[-1].response_id == "resp-1" assert released[-1].agent_id == "agent-1" - assert released[-1].continuation_token == "token-1" + assert released[-1].continuation_token == {"token": "token-1"} assert released[-1].additional_properties is not None assert released[-1].additional_properties["custom"] == "value" From ca85853396eaf9c060d3ae3714621ce2ba93213c Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 1 Oct 2026 13:11:50 +0000 Subject: [PATCH 4/7] Address PR comments. --- .../src/Microsoft.Agents.AI.Purview/README.md | 28 ++ .../PurviewDurableHistoryTests.cs | 133 ++++++++ python/packages/purview/README.md | 104 ++++++ .../agent_framework_purview/_middleware.py | 142 ++++---- .../purview/tests/purview/conftest.py | 123 +++++++ .../tests/purview/test_chat_middleware.py | 116 ++++++- .../purview/tests/purview/test_middleware.py | 306 +++++++++++++++++- 7 files changed, 864 insertions(+), 88 deletions(-) create mode 100644 dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewDurableHistoryTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/README.md b/dotnet/src/Microsoft.Agents.AI.Purview/README.md index 59c0dea9a33..5f7ee3d9b2f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/README.md +++ b/dotnet/src/Microsoft.Agents.AI.Purview/README.md @@ -275,6 +275,34 @@ The last form is not corrected automatically. When a chat client is turned into The user id from the prompt message(s) is reused for the response evaluation so both evaluations map consistently to the same user. +### Blocked content and conversation history + +Where Purview sits in the pipeline decides whether a blocked response can still be written +to conversation history. + +**Chat-client level** (`WithPurview` on a `ChatClientBuilder`, or `PurviewChatClient`) — +the response is evaluated and replaced inside `GetResponseAsync`, before the agent stores +the turn, so the replacement is what becomes durable and the model's own content never +reaches the history provider. This holds as long as Purview is composed below any chat +client that stores history itself, which is the same ordering rule described above for +function invocation. + +**Agent level** (`WithPurview` on an `AIAgentBuilder`, or `PurviewAgent`) — the agent +stores the turn as part of the run, and Purview wraps the whole run, so the response is +evaluated after the turn has already been written. A blocked response is replaced for the +caller, but the original content can still be read back from history on a later turn. Use +the chat-client level composition where history must not retain blocked content. + +**Service-managed history** — both of the above describe history the framework stores. +Some chat clients instead keep the conversation on the service and return an id to +continue from, for example a Responses-style API called with `store` enabled. There the +service records the prompt and the response as part of the model call, before the response +comes back to be evaluated, and the agent's session advances to that turn so the next run +resumes from a conversation that still contains the blocked content. Purview still replaces +what the caller receives, but the service's copy is outside the framework's reach and +cannot be withdrawn. Where blocked content must not be retained, turn service-side storage +off so the framework owns the conversation, and use a `ChatHistoryProvider`. + There are several optimizations to speed up Purview calls. Protection scope lookups (the first step in evaluation) are cached to minimize network calls. When a lookup is not cached, the middleware will refresh it in a background worker so the foreground ProcessContent request does not have to wait. If the policies allow content to be processed offline, the middleware will add the process content request to a channel and run it in a background worker. Similarly, the middleware will run a background request if no scopes apply and the interaction only has to be logged in Audit. Payment Required responses from background scope lookups are cached at the tenant level so subsequent requests for the tenant short-circuit. diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewDurableHistoryTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewDurableHistoryTests.cs new file mode 100644 index 00000000000..1b6fd59e4a1 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewDurableHistoryTests.cs @@ -0,0 +1,133 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Purview.Models.Common; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging.Abstractions; +using Moq; + +namespace Microsoft.Agents.AI.Purview.UnitTests; + +/// +/// Pins what becomes durable history when Purview is composed at the chat-client level. +/// The agent persists the response its chat client returned, and Purview has already +/// replaced a blocked response by then, so the replacement is what is stored and the +/// model's own content never reaches the history provider. +/// +public sealed class PurviewDurableHistoryTests +{ + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ChatLevelBlockedResponseDoesNotBecomeDurableHistoryAsync(bool streaming) + { + // Arrange + var provider = new RecordingHistoryProvider(); + var agent = CreateAgent(promptBlocked: false, responseBlocked: true, provider); + var session = await agent.CreateSessionAsync(); + + // Act + var text = await RunAsync(agent, session, streaming); + + // Assert + Assert.Contains("Response blocked by policy", text); + Assert.DoesNotContain(provider.Stored, message => message.Text.Contains("Sensitive response")); + Assert.Contains(provider.Stored, message => message.Text.Contains("Response blocked by policy")); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ChatLevelAllowedResponseStillBecomesDurableHistoryAsync(bool streaming) + { + // Arrange + var provider = new RecordingHistoryProvider(); + var agent = CreateAgent(promptBlocked: false, responseBlocked: false, provider); + var session = await agent.CreateSessionAsync(); + + // Act + var text = await RunAsync(agent, session, streaming); + + // Assert + Assert.Contains("Sensitive response", text); + Assert.Contains(provider.Stored, message => message.Text.Contains("Sensitive response")); + } + + private static async Task RunAsync(AIAgent agent, AgentSession session, bool streaming) + { + var message = new ChatMessage(ChatRole.User, "Test message"); + if (!streaming) + { + return (await agent.RunAsync(message, session)).Text; + } + + var text = string.Empty; + await foreach (var update in agent.RunStreamingAsync(message, session)) + { + text += update.Text; + } + + return text; + } + + private static ChatClientAgent CreateAgent(bool promptBlocked, bool responseBlocked, ChatHistoryProvider provider) + { + var settings = new PurviewSettings("TestApp") + { + TenantId = "tenant-123", + PurviewAppLocation = new PurviewAppLocation(PurviewLocationType.Application, "app-123"), + BlockedPromptMessage = "Prompt blocked by policy", + BlockedResponseMessage = "Response blocked by policy" + }; + + var processor = new Mock(); + processor.Setup(x => x.ProcessMessagesAsync( + It.IsAny>(), + It.IsAny(), + Activity.UploadText, + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync((promptBlocked, "user-123")); + processor.Setup(x => x.ProcessMessagesAsync( + It.IsAny>(), + It.IsAny(), + Activity.DownloadText, + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync((responseBlocked, "user-123")); + + var wrapper = new PurviewWrapper(processor.Object, settings, NullLogger.Instance, Mock.Of()); + + var innerClient = new Mock(); + innerClient.Setup(x => x.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(() => new ChatResponse(new ChatMessage(ChatRole.Assistant, "Sensitive response"))); + + return new ChatClientAgent( + new PurviewChatClient(innerClient.Object, wrapper), + new ChatClientAgentOptions { ChatHistoryProvider = provider }); + } + + /// A chat history provider that records exactly what becomes durable. + private sealed class RecordingHistoryProvider : ChatHistoryProvider + { + public List Stored { get; } = []; + + protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) => + new([.. this.Stored]); + + protected override ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + this.Stored.AddRange(context.RequestMessages); + this.Stored.AddRange(context.ResponseMessages ?? []); + return default; + } + } +} diff --git a/python/packages/purview/README.md b/python/packages/purview/README.md index aa3ea9d86e4..d25e84bf695 100644 --- a/python/packages/purview/README.md +++ b/python/packages/purview/README.md @@ -80,6 +80,8 @@ asyncio.run(main()) If a policy violation is detected on the prompt, the middleware terminates the run and substitutes a system message: `"Prompt blocked by policy"`. If on the response, the result becomes `"Response blocked by policy"`. +If you attach other middleware as well, put the Purview middleware last — see [Middleware order matters](#middleware-order-matters). + --- ## Configuration @@ -228,6 +230,7 @@ you want. | A model's tool call, before the tool executes | not evaluated | evaluated | | Tool results | evaluated at the end of the run | evaluated on the next request | | Final response | evaluated | evaluated | +| A blocked response | replaced for the caller, but [kept in history](#blocked-content-and-conversation-history) | [replaced before the turn is stored](#blocked-content-and-conversation-history), unless history is written per service call | | How often it evaluates | once per run | once per model request | The agent middleware receives the messages passed to `Agent.run()`. Content added while @@ -284,6 +287,9 @@ agent = Agent(client=client, instructions="You are helpful.") Both middlewares can be attached at the same time. The chat middleware then evaluates each model round trip and the agent middleware evaluates the run boundary. +In either list, attach the Purview middleware last — see +[Middleware order matters](#middleware-order-matters). + --- ## Middleware Lifecycle @@ -302,6 +308,94 @@ The user identifier is discovered from `Message.additional_properties['user_id'] You can customize the blocking messages using the `blocked_prompt_message` and `blocked_response_message` fields in `PurviewSettings`. For more advanced scenarios, you can wrap the middleware or post-process `context.result` in later middleware. +### Middleware order matters + +**Attach the Purview middleware last in the middleware list.** + +```python +agent = Agent( + ..., + middleware=[logging_middleware, caching_middleware, purview_middleware], + # ^ last +) +``` + +Middleware are free to replace the response, and those replacements are applied in +middleware order. Whichever middleware comes last gets the final say over what the caller +receives. So any middleware attached *after* Purview can return a response that Purview +never evaluated, and no error or warning is raised when this happens — the response is +simply delivered unevaluated. + +Attaching Purview last also means it evaluates the response *after* other middleware have +finished changing it, so the content that was evaluated is the content the caller +receives. + +The same applies to `PurviewChatPolicyMiddleware` in a chat client's middleware list. + +### Blocked content and conversation history + +Which middleware you use decides whether a blocked response can still be read back from +conversation history. + +**Chat middleware** (`PurviewChatPolicyMiddleware`, in a chat client's middleware list) — +the response is evaluated and replaced during the model call, before an agent run stores +the turn, so the replacement is what becomes durable and the model's own content does not +reach the history provider. + +**Agent middleware** (`PurviewPolicyMiddleware`, in an agent's middleware list) — an agent +run writes its messages to history as part of the run, and this middleware wraps the whole +run, so the turn has already been written by the time the response is evaluated. A blocked +response is replaced for the caller, but the original content can still be read back from +history on a later turn. Use the chat middleware where history must not retain blocked +content. + +Both cover the history providers the agent run writes through. Neither covers code that +writes somewhere else of its own accord, such as a tool that saves its own output. + +#### Not recommended: history written per service call + +The chat middleware description above assumes history is written once per run, which is +the default. An agent configured with `require_per_service_call_history_persistence=True` +instead writes history as each model call returns. + +That writer is added by the agent itself and placed *below* anything in your chat +middleware list, closer to the model. It therefore stores the model's own response before +the chat middleware evaluates it, and a blocked response can still be read back from +history on a later turn. No ordering of your own middleware can change this, because the +writer is always placed below them. + +Avoid per-service-call history persistence where history must not retain blocked content. + +#### Exception: history kept by the service + +Both of the above cover history the framework stores. Some chat clients instead keep the +conversation on the service and return an id to continue from — for example a +Responses-style API run with `store=True`, which is the default for some clients. + +In that mode the framework does not attach a local history provider of its own, and the +service records the prompt and the response as part of the model call, before the response +comes back to be evaluated. The service's copy is outside the framework's reach. The +session's continuation id also advances to that turn, so the next turn resumes from a +conversation that still contains the blocked content. + +Purview still replaces what the caller receives. It cannot remove what the service kept. + +If blocked content must not be retained, set `store` to `False` so the framework owns the +conversation, and attach a history provider: + +```python +client = ChatClient(...) +client.chat_middleware = [PurviewChatPolicyMiddleware(credential, settings)] + +agent = Agent( + client=client, + context_providers=[InMemoryHistoryProvider()], + default_options={"store": False}, +) +``` + +`store` can also be set per run, with `agent.run(..., options={"store": False})`. + --- ## Exceptions @@ -376,6 +470,10 @@ sent as Purview binary content, and function calls, function results and other s serialized to text. Only `usage` content is skipped, because it carries token counts rather than user data. +Evaluation covers the response as it stands when the middleware runs. Because middleware can replace +the response afterwards, what is evaluated depends on where the middleware is attached — see +[Middleware order matters](#middleware-order-matters). + ### Remote references are not dereferenced Purview classifies the content it is handed; a reference to content is not the content. @@ -398,12 +496,18 @@ evaluated. This uses an experimental core API, so attaching the middleware to a streaming call emits an `ExperimentalWarning`. +A streamed response is evaluated once it has been assembled, at the same point a +non-streaming response is, so [Middleware order matters](#middleware-order-matters) +applies to streaming in exactly the same way. + --- ## Notes - **User Identification**: When the configured credential resolves to a user token, that token's `user_id` is used for per-user policy scoping. For app-token credentials, provide a `user_id` per request (e.g. in `Message(..., additional_properties={"user_id": ""})`). If no user_id can be provided or inferred, the request fails rather than proceeding unevaluated — see [Security Considerations](#security-considerations). - **Blocking Messages**: Can be customized via `blocked_prompt_message` and `blocked_response_message` in `PurviewSettings`. By default, they are "Prompt blocked by policy" and "Response blocked by policy" respectively. +- **Middleware Order**: Attach the middleware last in the middleware list. Anything attached after it can return a response it never evaluated — see [Middleware order matters](#middleware-order-matters). - **Streaming Responses**: Streamed responses are buffered and evaluated in full before any update is released, so content is not delivered incrementally while the middleware is attached — see [Streaming responses](#streaming-responses). +- **Blocked Content and History**: The chat middleware replaces a blocked response before the turn is stored; the agent middleware replaces it only for the caller, and history keeps the original. Per-service-call history persistence is not recommended with either — see [Blocked content and conversation history](#blocked-content-and-conversation-history). - **Error Handling**: Use `ignore_exceptions` and `ignore_payment_required` settings for graceful degradation. When enabled, errors are logged but don't fail the request. - **Caching**: Protection scopes responses and 402 errors are cached by default with a 4-hour TTL. Cache is automatically invalidated when protection scope state changes. - **Cold-cache parallelization**: On a `ProtectionScopes` cache miss, scopes are refreshed in the background while `ProcessContent` runs in the foreground. diff --git a/python/packages/purview/agent_framework_purview/_middleware.py b/python/packages/purview/agent_framework_purview/_middleware.py index dff188d4eab..5dcde5d2436 100644 --- a/python/packages/purview/agent_framework_purview/_middleware.py +++ b/python/packages/purview/agent_framework_purview/_middleware.py @@ -169,12 +169,15 @@ async def process( session_id_response = session_id if isinstance(context.result, ResponseStream): - context.result = self._gate_stream( + context.result = self._evaluate_stream( + context, cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", context.result), session_id_response, resolved_user_id, ) - elif context.result is not None: + return + + if context.result is not None: should_block_response = await _should_block_response( self._processor, self._settings, @@ -190,34 +193,32 @@ def _blocked_response(self) -> "AgentResponse[Any]": msg = self._settings.get("blocked_response_message", None) or "Response blocked by policy" return AgentResponse(messages=[Message(role="system", contents=[msg])]) - def _gate_stream( + def _evaluate_stream( self, + context: AgentContext, inner: "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", session_id: str | None, user_id: str | None, ) -> "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]": """Buffer a streamed run and evaluate its complete content before anything is released. - The stream is drained in full, the assembled response is evaluated exactly as a - non-streaming response is, and only then are updates released. A blocked response - is replaced. The released updates are always re-derived from the response that was - evaluated, so what the caller receives cannot diverge from what policy saw even when - the inner stream's finalizer returns something other than the assembly of its own - updates. + The processing stages are registered on the context rather than applied by wrapping + the stream, so several middleware can guard the same run and each one's work is + applied. The pipeline buffers the stream, applies the registered transforms to the + assembled response, and releases updates re-derived from the result. + + Response transforms run in middleware order, so a middleware placed after this one + can replace the response after it has been evaluated. Attach this middleware last + for its evaluation to cover what the caller actually receives; see the package + README on middleware order. + + This covers what the caller receives. The run writes its messages to conversation + history as part of the run itself, below this middleware, so that write has already + happened by the time the response is evaluated here; see the package README on + blocked content and conversation history. """ - async def _consume() -> tuple[Sequence[AgentResponseUpdate], "AgentResponse[Any]"]: - try: - final = await inner.get_final_response() - finally: - # Cancellation or failure part-way through the drain still has to release - # the inner stream, which the cleanup hook below would never reach. - await inner.close() - return list(inner.updates), final - - async def _gate( - _updates: list[AgentResponseUpdate], final: "AgentResponse[Any]" - ) -> tuple["AgentResponse[Any]", bool]: + async def _transform(final: "AgentResponse[Any]") -> "AgentResponse[Any]": should_block_response = await _should_block_response( self._processor, self._settings, @@ -226,23 +227,17 @@ async def _gate( user_id, ) if should_block_response: - return self._blocked_response(), True - # Reported as transformed even when the content is allowed through unchanged, so - # the released updates are re-derived from the evaluated response rather than - # replayed from the buffer. - return final, True - - gated = cast( - "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", - cast(Any, ResponseStream).buffered_and_gated( - consume=_consume, - gate=_gate, - rederive=AgentResponse.to_updates, - ), - ) - # Closing the gated stream before it is ever pulled never reaches ``_consume``, so - # the inner stream is released through a cleanup hook rather than from inside it. - return gated.with_cleanup_hook(inner.close) + return self._blocked_response() + # Returning the evaluated response, rather than None, is what makes the + # released updates be re-derived from it. A finalizer may produce a response + # that is not the assembly of its own updates, so releasing the buffered + # updates instead would release content this evaluation never saw. + return final + + context.stream_result_transforms.append(_transform) + context.stream_result_to_updates = AgentResponse.to_updates + context.stream_buffer_updates = True + return inner class PurviewChatPolicyMiddleware(ChatMiddleware): @@ -321,12 +316,15 @@ async def process( session_id_response = session_id if isinstance(context.result, ResponseStream): - context.result = self._gate_stream( + context.result = self._evaluate_stream( + context, cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", context.result), session_id_response, resolved_user_id, ) - elif context.result is not None: + return + + if context.result is not None: should_block_response = await _should_block_response( self._processor, self._settings, @@ -345,34 +343,32 @@ def _blocked_response(self) -> "ChatResponse[Any]": ) return ChatResponse(messages=[blocked_message]) - def _gate_stream( + def _evaluate_stream( self, + context: ChatContext, inner: "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", session_id: str | None, user_id: str | None, ) -> "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]": """Buffer a streamed response and evaluate its complete content before anything is released. - The stream is drained in full, the assembled response is evaluated exactly as a - non-streaming response is, and only then are updates released. A blocked response - is replaced. The released updates are always re-derived from the response that was - evaluated, so what the caller receives cannot diverge from what policy saw even when - the inner stream's finalizer returns something other than the assembly of its own - updates. + The processing stages are registered on the context rather than applied by wrapping + the stream, so several middleware can guard the same call and each one's work is + applied. The pipeline buffers the stream, applies the registered transforms to the + assembled response, and releases updates re-derived from the result. + + Response transforms run in middleware order, so a middleware placed after this one + can replace the response after it has been evaluated. Attach this middleware last + for its evaluation to cover what the caller actually receives; see the package + README on middleware order. + + This covers what the caller receives, and the evaluated response is what an agent + run above this middleware stores. Per-service-call history writes happen below it + and are not covered; see the package README on blocked content and conversation + history. """ - async def _consume() -> tuple[Sequence[ChatResponseUpdate], "ChatResponse[Any]"]: - try: - final = await inner.get_final_response() - finally: - # Cancellation or failure part-way through the drain still has to release - # the inner stream, which the cleanup hook below would never reach. - await inner.close() - return list(inner.updates), final - - async def _gate( - _updates: list[ChatResponseUpdate], final: "ChatResponse[Any]" - ) -> tuple["ChatResponse[Any]", bool]: + async def _transform(final: "ChatResponse[Any]") -> "ChatResponse[Any]": should_block_response = await _should_block_response( self._processor, self._settings, @@ -381,20 +377,14 @@ async def _gate( user_id, ) if should_block_response: - return self._blocked_response(), True - # Reported as transformed even when the content is allowed through unchanged, so - # the released updates are re-derived from the evaluated response rather than - # replayed from the buffer. - return final, True - - gated = cast( - "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", - cast(Any, ResponseStream).buffered_and_gated( - consume=_consume, - gate=_gate, - rederive=ChatResponse.to_updates, - ), - ) - # Closing the gated stream before it is ever pulled never reaches ``_consume``, so - # the inner stream is released through a cleanup hook rather than from inside it. - return gated.with_cleanup_hook(inner.close) + return self._blocked_response() + # Returning the evaluated response, rather than None, is what makes the + # released updates be re-derived from it. A finalizer may produce a response + # that is not the assembly of its own updates, so releasing the buffered + # updates instead would release content this evaluation never saw. + return final + + context.stream_result_transforms.append(_transform) + context.stream_result_to_updates = ChatResponse.to_updates + context.stream_buffer_updates = True + return inner diff --git a/python/packages/purview/tests/purview/conftest.py b/python/packages/purview/tests/purview/conftest.py index 0089053b7c9..d5924bb46b8 100644 --- a/python/packages/purview/tests/purview/conftest.py +++ b/python/packages/purview/tests/purview/conftest.py @@ -1,7 +1,22 @@ # Copyright (c) Microsoft. All rights reserved. """Shared pytest fixtures for Purview tests.""" +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from typing import Any + import pytest +from agent_framework import ( + Agent, + AgentSession, + BaseChatClient, + ChatMiddlewareLayer, + ChatResponse, + ChatResponseUpdate, + Content, + HistoryProvider, + Message, + ResponseStream, +) from agent_framework_purview._models import ( Activity, @@ -18,6 +33,70 @@ ) +class StubChatClient(ChatMiddlewareLayer[Any], BaseChatClient[Any]): + """Chat client that answers with one fixed assistant message, streamed or not.""" + + def __init__(self, text: str = "model reply", **kwargs: Any) -> None: + super().__init__(middleware=[], **kwargs) + self.text = text + + def _inner_get_response( # type: ignore[override] + self, + *, + messages: Sequence[Message], + stream: bool, + options: Mapping[str, Any], + **kwargs: Any, + ) -> "Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]": + if stream: + + async def _updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text(text=self.text)]) + + return ResponseStream(_updates(), finalizer=ChatResponse.from_updates) + + async def _response() -> "ChatResponse[Any]": + return ChatResponse(messages=[Message(role="assistant", contents=[self.text])]) + + return _response() + + +@pytest.fixture +def stub_chat_client() -> type[StubChatClient]: + """Factory for a chat client that answers with a fixed assistant message.""" + return StubChatClient + + +@pytest.fixture +def stored_texts() -> Callable[[AgentSession, HistoryProvider], list[str]]: + """Read back the message texts a history provider made durable for a session.""" + + def _read(session: AgentSession, provider: HistoryProvider) -> list[str]: + stored = session.state.get(provider.source_id, {}).get("messages", []) + texts: list[str] = [] + for message in stored: + text = message.get("text") if isinstance(message, dict) else getattr(message, "text", None) + texts.append(text if text is not None else str(message)) + return texts + + return _read + + +@pytest.fixture +def run_agent() -> Callable[..., Awaitable[str]]: + """Run an agent through either path and return the text the caller received.""" + + async def _run(agent: Agent, prompt: str, session: AgentSession, *, stream: bool) -> str: + if stream: + text = "" + async for update in agent.run(prompt, session=session, stream=True): + text += update.text + return text + return (await agent.run(prompt, session=session)).text + + return _run + + @pytest.fixture def content_to_process_factory(): """Factory fixture to create ContentToProcess objects with test data.""" @@ -66,3 +145,47 @@ def _create_request( ) return _create_request + + +async def run_agent_middleware( + middleware: Any, + context: Any, + set_result: Callable[[], Awaitable[None]], +) -> Any: + """Run one agent middleware through the real pipeline and return the finished result. + + Middleware declares the stream processing it needs on the context; the pipeline is what + applies that declaration to the stream once the chain unwinds. Calling ``process`` + directly would therefore leave a streamed result unguarded, so tests go through the + pipeline to exercise the same wiring production uses. + """ + from agent_framework._middleware import AgentMiddlewarePipeline # pyright: ignore[reportPrivateUsage] + + items = tuple(middleware) if isinstance(middleware, (list, tuple)) else (middleware,) + + async def final_handler(ctx: Any) -> Any: + await set_result() + return ctx.result + + return await AgentMiddlewarePipeline(*items).execute(context, final_handler) + + +async def run_chat_middleware( + middleware: Any, + context: Any, + set_result: Callable[[], Awaitable[None]], +) -> Any: + """Run one chat middleware through the real pipeline and return the finished result. + + See :func:`run_agent_middleware` for why the pipeline is used instead of calling + ``process`` directly. + """ + from agent_framework._middleware import ChatMiddlewarePipeline # pyright: ignore[reportPrivateUsage] + + items = tuple(middleware) if isinstance(middleware, (list, tuple)) else (middleware,) + + async def final_handler(ctx: Any) -> Any: + await set_result() + return ctx.result + + return await ChatMiddlewarePipeline(*items).execute(context, final_handler) diff --git a/python/packages/purview/tests/purview/test_chat_middleware.py b/python/packages/purview/tests/purview/test_chat_middleware.py index cac98fded81..6d229724cbf 100644 --- a/python/packages/purview/tests/purview/test_chat_middleware.py +++ b/python/packages/purview/tests/purview/test_chat_middleware.py @@ -8,15 +8,19 @@ import pytest from agent_framework import ( + Agent, + AgentSession, ChatContext, ChatResponse, ChatResponseUpdate, Content, + InMemoryHistoryProvider, Message, MiddlewareTermination, ResponseStream, ) from azure.core.credentials import AccessToken +from conftest import run_chat_middleware # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from agent_framework_purview import PurviewChatPolicyMiddleware, PurviewSettings from agent_framework_purview._models import Activity @@ -133,7 +137,7 @@ async def updates() -> AsyncIterator[ChatResponseUpdate]: async def mock_next() -> None: streaming_context.result = cast(Any, ResponseStream(updates(), finalizer=ChatResponse.from_updates)) - await middleware.process(streaming_context, mock_next) + await run_chat_middleware(middleware, streaming_context, mock_next) released = [update async for update in cast(Any, streaming_context.result)] assert mock_proc.call_count == 2 @@ -167,7 +171,7 @@ async def updates() -> AsyncIterator[ChatResponseUpdate]: async def mock_next() -> None: streaming_context.result = cast(Any, ResponseStream(updates(), finalizer=ChatResponse.from_updates)) - await middleware.process(streaming_context, mock_next) + await run_chat_middleware(middleware, streaming_context, mock_next) released = [update async for update in cast(Any, streaming_context.result)] assert mock_proc.call_count == 2 @@ -203,7 +207,7 @@ def diverging_finalizer(_updates: Any) -> ChatResponse[Any]: async def mock_next() -> None: streaming_context.result = cast(Any, ResponseStream(updates(), finalizer=diverging_finalizer)) - await middleware.process(streaming_context, mock_next) + await run_chat_middleware(middleware, streaming_context, mock_next) released = [update async for update in cast(Any, streaming_context.result)] released_text = "".join(update.text for update in released) @@ -244,7 +248,7 @@ def finalizer(_updates: Any) -> ChatResponse[Any]: async def mock_next() -> None: streaming_context.result = cast(Any, ResponseStream(updates(), finalizer=finalizer)) - await middleware.process(streaming_context, mock_next) + await run_chat_middleware(middleware, streaming_context, mock_next) released = [update async for update in cast(Any, streaming_context.result)] assert released[-1].response_id == "resp-1" @@ -287,7 +291,7 @@ async def mock_next() -> None: inner.with_cleanup_hook(_mark_closed) streaming_context.result = cast(Any, inner) - await middleware.process(streaming_context, mock_next) + await run_chat_middleware(middleware, streaming_context, mock_next) await cast(Any, streaming_context.result).close() assert closed is True @@ -610,3 +614,105 @@ async def mock_next() -> None: # Check post-check call includes session_id post_check_call = mock_proc.call_args_list[1] assert post_check_call[1]["session_id"] == "conv-999" + + +class TestPurviewChatPolicyMiddlewareDurableHistory: + """How this middleware's response check relates to what is written to history. + + An agent run that stores history once per run does so above this middleware, so the + evaluated response is what gets stored. Per-service-call history writes happen below + it, as each model call returns, so those are not covered. + """ + + @pytest.fixture + def middleware(self) -> PurviewChatPolicyMiddleware: + credential = AsyncMock() + credential.get_token = AsyncMock(return_value=AccessToken("fake-token", 9999999999)) + return PurviewChatPolicyMiddleware(credential, PurviewSettings(app_name="Test App", tenant_id="test-tenant")) + + @staticmethod + def _build(middleware: PurviewChatPolicyMiddleware, client_factory: Any, per_service_call: bool) -> Any: + client = client_factory() + client.chat_middleware = [middleware] + provider = InMemoryHistoryProvider() + agent = Agent( + client=client, + context_providers=[provider], + require_per_service_call_history_persistence=per_service_call, + ) + return agent, provider + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + async def test_blocked_response_is_replaced_before_the_turn_is_stored( + self, + middleware: PurviewChatPolicyMiddleware, + stub_chat_client: Any, + stored_texts: Any, + run_agent: Any, + stream: bool, + ) -> None: + """With once-per-run history, the replacement is what becomes durable.""" + agent, provider = self._build(middleware, stub_chat_client, per_service_call=False) + session = AgentSession() + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ): + text = await run_agent(agent, "hello there", session, stream=stream) + + assert "Response blocked by policy" in text + assert stored_texts(session, provider) == ["hello there", "Response blocked by policy"] + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + async def test_blocked_response_stays_in_history_when_written_per_service_call( + self, + middleware: PurviewChatPolicyMiddleware, + stub_chat_client: Any, + stored_texts: Any, + run_agent: Any, + stream: bool, + ) -> None: + """Per-service-call writes happen below this middleware, so they keep the original. + + Pinned so the documented behaviour cannot drift unnoticed: per-service-call history + persistence is not recommended alongside this middleware. + """ + agent, provider = self._build(middleware, stub_chat_client, per_service_call=True) + session = AgentSession() + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ): + text = await run_agent(agent, "hello there", session, stream=stream) + + assert "Response blocked by policy" in text + assert stored_texts(session, provider) == ["hello there", "model reply"] + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + @pytest.mark.parametrize("per_service_call", [False, True], ids=["per_run", "per_service_call"]) + async def test_allowed_response_still_becomes_history( + self, + middleware: PurviewChatPolicyMiddleware, + stub_chat_client: Any, + stored_texts: Any, + run_agent: Any, + stream: bool, + per_service_call: bool, + ) -> None: + """Evaluating the response must not disturb an allowed call's history write.""" + agent, provider = self._build(middleware, stub_chat_client, per_service_call) + session = AgentSession() + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + text = await run_agent(agent, "hello there", session, stream=stream) + + assert text == "model reply" + assert stored_texts(session, provider) == ["hello there", "model reply"] diff --git a/python/packages/purview/tests/purview/test_middleware.py b/python/packages/purview/tests/purview/test_middleware.py index 932df96bdce..9b0d8264d8c 100644 --- a/python/packages/purview/tests/purview/test_middleware.py +++ b/python/packages/purview/tests/purview/test_middleware.py @@ -3,22 +3,26 @@ """Tests for Purview middleware.""" import asyncio -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from agent_framework import ( + Agent, AgentContext, + AgentMiddleware, AgentResponse, AgentResponseUpdate, AgentSession, Content, + InMemoryHistoryProvider, Message, MiddlewareTermination, ResponseStream, ) from azure.core.credentials import AccessToken +from conftest import run_agent_middleware # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from agent_framework_purview import PurviewPolicyMiddleware, PurviewSettings from agent_framework_purview._models import Activity @@ -186,7 +190,7 @@ async def updates() -> AsyncIterator[AgentResponseUpdate]: async def mock_next() -> None: context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) - await middleware.process(context, mock_next) + await run_agent_middleware(middleware, context, mock_next) released = [update async for update in cast(Any, context.result)] assert mock_proc.call_count == 2 @@ -215,7 +219,7 @@ async def updates() -> AsyncIterator[AgentResponseUpdate]: async def mock_next() -> None: context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) - await middleware.process(context, mock_next) + await run_agent_middleware(middleware, context, mock_next) released = [update async for update in cast(Any, context.result)] assert mock_proc.call_count == 2 @@ -248,7 +252,7 @@ def diverging_finalizer(_updates: Any) -> AgentResponse[Any]: async def mock_next() -> None: context.result = cast(Any, ResponseStream(updates(), finalizer=diverging_finalizer)) - await middleware.process(context, mock_next) + await run_agent_middleware(middleware, context, mock_next) released = [update async for update in cast(Any, context.result)] released_text = "".join(update.text for update in released) @@ -283,7 +287,7 @@ def finalizer(_updates: Any) -> AgentResponse[Any]: async def mock_next() -> None: context.result = cast(Any, ResponseStream(updates(), finalizer=finalizer)) - await middleware.process(context, mock_next) + await run_agent_middleware(middleware, context, mock_next) released = [update async for update in cast(Any, context.result)] assert released[-1].response_id == "resp-1" @@ -318,7 +322,7 @@ async def mock_next() -> None: inner.with_cleanup_hook(_mark_closed) context.result = cast(Any, inner) - await middleware.process(context, mock_next) + await run_agent_middleware(middleware, context, mock_next) await cast(Any, context.result).close() assert closed is True @@ -352,7 +356,7 @@ async def mock_next() -> None: inner.with_cleanup_hook(_mark_closed) context.result = cast(Any, inner) - await middleware.process(context, mock_next) + await run_agent_middleware(middleware, context, mock_next) async def drain() -> None: async for _ in cast(Any, context.result): @@ -625,3 +629,291 @@ async def mock_next() -> None: # Check post-check call includes session_id post_check_call = mock_proc.call_args_list[1] assert post_check_call[1]["session_id"] == "thread-999" + + +class TestPurviewPolicyMiddlewareDurableHistory: + """How this middleware's response check relates to what the run writes to history. + + An agent run writes its messages to conversation history as part of the run itself, + which happens below this middleware. The response check therefore runs after that + write, and governs what the caller receives rather than what is stored. + """ + + @pytest.fixture + def middleware(self) -> PurviewPolicyMiddleware: + credential = AsyncMock() + credential.get_token = AsyncMock(return_value=AccessToken("fake-token", 9999999999)) + return PurviewPolicyMiddleware(credential, PurviewSettings(app_name="Test App", tenant_id="test-tenant")) + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + async def test_blocked_response_is_replaced_for_the_caller_but_stays_in_history( + self, + middleware: PurviewPolicyMiddleware, + stub_chat_client: Any, + stored_texts: Any, + run_agent: Any, + stream: bool, + ) -> None: + """The caller receives the replacement, while history keeps what the run wrote. + + Pinned so the documented behaviour cannot drift unnoticed: callers who need + history to exclude blocked content use the chat middleware instead. + """ + provider = InMemoryHistoryProvider() + session = AgentSession() + agent = Agent(client=stub_chat_client(), context_providers=[provider], middleware=[middleware]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ): + text = await run_agent(agent, "hello there", session, stream=stream) + + assert "Response blocked by policy" in text + assert stored_texts(session, provider) == ["hello there", "model reply"] + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + async def test_allowed_response_still_becomes_history( + self, + middleware: PurviewPolicyMiddleware, + stub_chat_client: Any, + stored_texts: Any, + run_agent: Any, + stream: bool, + ) -> None: + """Evaluating the response must not disturb an allowed run's history write.""" + provider = InMemoryHistoryProvider() + session = AgentSession() + agent = Agent(client=stub_chat_client(), context_providers=[provider], middleware=[middleware]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ): + text = await run_agent(agent, "hello there", session, stream=stream) + + assert text == "model reply" + assert stored_texts(session, provider) == ["hello there", "model reply"] + + +class _SuffixMiddleware(AgentMiddleware): + """A second middleware that also inspects and rewrites a streamed response. + + It registers the same kind of whole-response processing that Purview does, so a run + guarded by both exercises the case where two middleware each need the assembled + response before anything is released to the caller. + """ + + def __init__(self, suffix: str) -> None: + self.suffix = suffix + self.seen_text: str | None = None + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + if not isinstance(context.result, ResponseStream): + return + + async def _transform(final: "AgentResponse[Any]") -> "AgentResponse[Any]": + self.seen_text = final.text + return AgentResponse(messages=[Message(role="assistant", contents=[final.text + self.suffix])]) + + context.stream_result_transforms.append(_transform) + context.stream_result_to_updates = AgentResponse.to_updates + context.stream_buffer_updates = True + + +class _ResponseSubstitutingMiddleware(AgentMiddleware): + """A middleware that supplies its own response instead of passing on what it received. + + Caching and fallback middleware behave this way: they hand back a stored answer and + discard whatever the run produced. Registered on the ordinary rewrite stage, which is + where such middleware belongs. + """ + + def __init__(self) -> None: + self.ran = False + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + if not isinstance(context.result, ResponseStream): + return + + def _rewrite(_: "AgentResponse[Any]") -> "AgentResponse[Any]": + self.ran = True + return AgentResponse(messages=[Message(role="assistant", contents=["confidential"])]) + + context.stream_result_transforms.append(_rewrite) + context.stream_buffer_updates = True + context.stream_result_to_updates = AgentResponse.to_updates + + +class TestPurviewPolicyMiddlewareSharedStream: + """Purview guarding a streamed run alongside another middleware that rewrites responses. + + Response rewrites run in middleware order, so these use the arrangement the package + README recommends: Purview attached last, after any middleware that rewrites responses. + """ + + @pytest.fixture + def mock_credential(self) -> AsyncMock: + credential = AsyncMock() + credential.get_token = AsyncMock(return_value=AccessToken("fake-token", 9999999999)) + return credential + + @pytest.fixture + def middleware(self, mock_credential: AsyncMock) -> PurviewPolicyMiddleware: + return PurviewPolicyMiddleware(mock_credential, PurviewSettings(app_name="Test App", tenant_id="test-tenant")) + + @pytest.fixture + def mock_agent(self) -> MagicMock: + agent = MagicMock() + agent.name = "test-agent" + return agent + + async def test_both_middleware_change_the_same_streamed_run_when_allowed( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """Neither middleware's work is lost when both act on one streamed run.""" + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + other = _SuffixMiddleware(" [checked]") + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="all clear")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ) as mock_proc: + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) + + await run_agent_middleware([other, middleware], context, mock_next) + released = [update async for update in cast(Any, context.result)] + + assert other.seen_text == "all clear" + assert mock_proc.call_count == 2 + assert "".join(update.text for update in released) == "all clear [checked]" + + async def test_the_evaluated_response_is_the_one_the_caller_receives( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """Attached last, the evaluation sees the response after other middleware rewrote it.""" + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + other = _SuffixMiddleware(" [checked]") + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="all clear")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (False, "user-123")], + ) as mock_proc: + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) + + await run_agent_middleware([other, middleware], context, mock_next) + released = [update async for update in cast(Any, context.result)] + + evaluated = mock_proc.call_args_list[1][0][0] + evaluated_text = "".join(part.text for message in evaluated for part in message.contents if part.text) + assert "[checked]" in evaluated_text + assert "".join(update.text for update in released) == evaluated_text + + async def test_purview_still_blocks_when_another_middleware_shares_the_stream( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """A second middleware on the same run cannot let blocked content reach the caller.""" + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + other = _SuffixMiddleware(" [checked]") + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="confidential")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ): + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) + + await run_agent_middleware([other, middleware], context, mock_next) + released = [update async for update in cast(Any, context.result)] + + released_text = "".join(update.text for update in released) + assert "confidential" not in released_text + assert "blocked" in released_text.lower() + + async def test_policy_replacement_survives_a_middleware_that_substitutes_the_response( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """A middleware that discards the response it is given cannot reinstate blocked content. + + Cache and fallback middleware return a stored answer rather than the one the run + produced. Attached last, the evaluation sees that substituted answer. + """ + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + other = _ResponseSubstitutingMiddleware() + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="something else")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ): + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) + + await run_agent_middleware([other, middleware], context, mock_next) + released = [update async for update in cast(Any, context.result)] + + released_text = "".join(update.text for update in released) + assert other.ran is True + assert "confidential" not in released_text + assert "blocked" in released_text.lower() + + async def test_a_middleware_attached_after_purview_can_replace_the_evaluated_response( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + ) -> None: + """Pins the documented limitation: attaching Purview before a rewriting middleware. + + Response transforms run in middleware order, so a middleware attached after this + one gets the last word and its response is never evaluated. The package README + tells users to attach Purview last for this reason. This test records what + actually happens when they do not, so the documentation cannot drift away from + the behaviour. + """ + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + other = _ResponseSubstitutingMiddleware() + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="something else")]) + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ): + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates)) + + await run_agent_middleware([middleware, other], context, mock_next) + released = [update async for update in cast(Any, context.result)] + + assert "".join(update.text for update in released) == "confidential" From e486bf924fa744904fd6640ff1f58aecdbf7a606 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 1 Oct 2026 13:30:53 +0000 Subject: [PATCH 5/7] Fix formatting --- .../PurviewDurableHistoryTests.cs | 1 - 1 file changed, 1 deletion(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewDurableHistoryTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewDurableHistoryTests.cs index 1b6fd59e4a1..b7acbdba9b3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewDurableHistoryTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewDurableHistoryTests.cs @@ -1,7 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. using System.Collections.Generic; -using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Purview.Models.Common; From fe42feca47f4ca53749d550fcee4e103c2cbf343 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:31:37 +0000 Subject: [PATCH 6/7] Fix broken tests. --- .../purview/tests/purview/conftest.py | 54 +++++++++++-------- .../tests/purview/test_chat_middleware.py | 17 +++--- .../purview/tests/purview/test_middleware.py | 23 ++++---- 3 files changed, 54 insertions(+), 40 deletions(-) diff --git a/python/packages/purview/tests/purview/conftest.py b/python/packages/purview/tests/purview/conftest.py index d5924bb46b8..3aec2b33490 100644 --- a/python/packages/purview/tests/purview/conftest.py +++ b/python/packages/purview/tests/purview/conftest.py @@ -147,11 +147,8 @@ def _create_request( return _create_request -async def run_agent_middleware( - middleware: Any, - context: Any, - set_result: Callable[[], Awaitable[None]], -) -> Any: +@pytest.fixture +def run_agent_middleware() -> Callable[..., Awaitable[Any]]: """Run one agent middleware through the real pipeline and return the finished result. Middleware declares the stream processing it needs on the context; the pipeline is what @@ -159,33 +156,46 @@ async def run_agent_middleware( directly would therefore leave a streamed result unguarded, so tests go through the pipeline to exercise the same wiring production uses. """ - from agent_framework._middleware import AgentMiddlewarePipeline # pyright: ignore[reportPrivateUsage] - items = tuple(middleware) if isinstance(middleware, (list, tuple)) else (middleware,) + async def _run( + middleware: Any, + context: Any, + set_result: Callable[[], Awaitable[None]], + ) -> Any: + from agent_framework._middleware import AgentMiddlewarePipeline # pyright: ignore[reportPrivateUsage] + + items = tuple(middleware) if isinstance(middleware, (list, tuple)) else (middleware,) - async def final_handler(ctx: Any) -> Any: - await set_result() - return ctx.result + async def final_handler(ctx: Any) -> Any: + await set_result() + return ctx.result - return await AgentMiddlewarePipeline(*items).execute(context, final_handler) + return await AgentMiddlewarePipeline(*items).execute(context, final_handler) + + return _run -async def run_chat_middleware( - middleware: Any, - context: Any, - set_result: Callable[[], Awaitable[None]], -) -> Any: +@pytest.fixture +def run_chat_middleware() -> Callable[..., Awaitable[Any]]: """Run one chat middleware through the real pipeline and return the finished result. See :func:`run_agent_middleware` for why the pipeline is used instead of calling ``process`` directly. """ - from agent_framework._middleware import ChatMiddlewarePipeline # pyright: ignore[reportPrivateUsage] - items = tuple(middleware) if isinstance(middleware, (list, tuple)) else (middleware,) + async def _run( + middleware: Any, + context: Any, + set_result: Callable[[], Awaitable[None]], + ) -> Any: + from agent_framework._middleware import ChatMiddlewarePipeline # pyright: ignore[reportPrivateUsage] + + items = tuple(middleware) if isinstance(middleware, (list, tuple)) else (middleware,) - async def final_handler(ctx: Any) -> Any: - await set_result() - return ctx.result + async def final_handler(ctx: Any) -> Any: + await set_result() + return ctx.result - return await ChatMiddlewarePipeline(*items).execute(context, final_handler) + return await ChatMiddlewarePipeline(*items).execute(context, final_handler) + + return _run diff --git a/python/packages/purview/tests/purview/test_chat_middleware.py b/python/packages/purview/tests/purview/test_chat_middleware.py index 6d229724cbf..c771d9bc183 100644 --- a/python/packages/purview/tests/purview/test_chat_middleware.py +++ b/python/packages/purview/tests/purview/test_chat_middleware.py @@ -20,7 +20,6 @@ ResponseStream, ) from azure.core.credentials import AccessToken -from conftest import run_chat_middleware # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from agent_framework_purview import PurviewChatPolicyMiddleware, PurviewSettings from agent_framework_purview._models import Activity @@ -113,7 +112,9 @@ async def mock_next() -> None: assert first_msg.role in ("system", "system") assert "blocked" in first_msg.text.lower() - async def test_streaming_response_is_evaluated_and_blocked(self, middleware: PurviewChatPolicyMiddleware) -> None: + async def test_streaming_response_is_evaluated_and_blocked( + self, middleware: PurviewChatPolicyMiddleware, run_chat_middleware: Any + ) -> None: """Streamed content is evaluated in full and replaced when policy blocks it.""" client = DummyChatClient() chat_options = MagicMock() @@ -146,7 +147,9 @@ async def mock_next() -> None: assert "confidential" not in released_text assert "blocked" in released_text.lower() - async def test_streaming_response_passes_when_allowed(self, middleware: PurviewChatPolicyMiddleware) -> None: + async def test_streaming_response_passes_when_allowed( + self, middleware: PurviewChatPolicyMiddleware, run_chat_middleware: Any + ) -> None: """Allowed streamed content is released unchanged, reusing the prompt-phase identity.""" client = DummyChatClient() chat_options = MagicMock() @@ -179,7 +182,7 @@ async def mock_next() -> None: assert "".join(update.text for update in released) == "all clear" async def test_streaming_releases_the_evaluated_content_not_the_buffered_updates( - self, middleware: PurviewChatPolicyMiddleware + self, middleware: PurviewChatPolicyMiddleware, run_chat_middleware: Any ) -> None: """The released updates come from the response that was evaluated, not from the buffer.""" client = DummyChatClient() @@ -214,7 +217,9 @@ async def mock_next() -> None: assert released_text == "benign" assert "confidential" not in released_text - async def test_streaming_preserves_response_level_metadata(self, middleware: PurviewChatPolicyMiddleware) -> None: + async def test_streaming_preserves_response_level_metadata( + self, middleware: PurviewChatPolicyMiddleware, run_chat_middleware: Any + ) -> None: """Metadata carried by the response survives the buffered stream.""" client = DummyChatClient() chat_options = MagicMock() @@ -259,7 +264,7 @@ async def mock_next() -> None: assert released[-1].additional_properties["custom"] == "value" async def test_streaming_closes_the_inner_stream_when_never_pulled( - self, middleware: PurviewChatPolicyMiddleware + self, middleware: PurviewChatPolicyMiddleware, run_chat_middleware: Any ) -> None: """Abandoning the gated stream before the first pull still releases the inner stream.""" client = DummyChatClient() diff --git a/python/packages/purview/tests/purview/test_middleware.py b/python/packages/purview/tests/purview/test_middleware.py index 9b0d8264d8c..b4c9ae90117 100644 --- a/python/packages/purview/tests/purview/test_middleware.py +++ b/python/packages/purview/tests/purview/test_middleware.py @@ -22,7 +22,6 @@ ResponseStream, ) from azure.core.credentials import AccessToken -from conftest import run_agent_middleware # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from agent_framework_purview import PurviewPolicyMiddleware, PurviewSettings from agent_framework_purview._models import Activity @@ -172,7 +171,7 @@ async def mock_next() -> None: assert mock_process.call_args_list[1][0][1] == Activity.DOWNLOAD_TEXT async def test_middleware_streaming_response_is_evaluated_and_blocked( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """Streamed content is evaluated in full and replaced when policy blocks it.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) @@ -200,7 +199,7 @@ async def mock_next() -> None: assert "blocked" in released_text.lower() async def test_middleware_streaming_response_passes_when_allowed( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """Allowed streamed content is released unchanged, reusing the prompt-phase identity.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) @@ -227,7 +226,7 @@ async def mock_next() -> None: assert "".join(update.text for update in released) == "all clear" async def test_middleware_streaming_releases_the_evaluated_content_not_the_buffered_updates( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """The released updates come from the response that was evaluated, not from the buffer. @@ -260,7 +259,7 @@ async def mock_next() -> None: assert "confidential" not in released_text async def test_middleware_streaming_preserves_response_level_metadata( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """Metadata carried by the response survives the buffered stream.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) @@ -297,7 +296,7 @@ async def mock_next() -> None: assert released[-1].additional_properties["custom"] == "value" async def test_middleware_streaming_closes_the_inner_stream_when_never_pulled( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """Abandoning the gated stream before the first pull still releases the inner stream.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) @@ -328,7 +327,7 @@ async def mock_next() -> None: assert closed is True async def test_middleware_streaming_closes_the_inner_stream_when_the_drain_is_cancelled( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """Cancelling part-way through the drain still releases the inner stream.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) @@ -773,7 +772,7 @@ def mock_agent(self) -> MagicMock: return agent async def test_both_middleware_change_the_same_streamed_run_when_allowed( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """Neither middleware's work is lost when both act on one streamed run.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) @@ -800,7 +799,7 @@ async def mock_next() -> None: assert "".join(update.text for update in released) == "all clear [checked]" async def test_the_evaluated_response_is_the_one_the_caller_receives( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """Attached last, the evaluation sees the response after other middleware rewrote it.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) @@ -828,7 +827,7 @@ async def mock_next() -> None: assert "".join(update.text for update in released) == evaluated_text async def test_purview_still_blocks_when_another_middleware_shares_the_stream( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """A second middleware on the same run cannot let blocked content reach the caller.""" context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) @@ -855,7 +854,7 @@ async def mock_next() -> None: assert "blocked" in released_text.lower() async def test_policy_replacement_survives_a_middleware_that_substitutes_the_response( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """A middleware that discards the response it is given cannot reinstate blocked content. @@ -887,7 +886,7 @@ async def mock_next() -> None: assert "blocked" in released_text.lower() async def test_a_middleware_attached_after_purview_can_replace_the_evaluated_response( - self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any ) -> None: """Pins the documented limitation: attaching Purview before a rewriting middleware. From 26103d1b9bcc99b1b85a71ef55c7a6ab545f7655 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 2 Oct 2026 12:05:03 +0000 Subject: [PATCH 7/7] Address PR comments --- .../PurviewWrapper.cs | 53 ++++++- .../PurviewWrapperTests.cs | 134 ++++++++++++++++++ .../agent_framework_purview/_middleware.py | 56 ++++++-- .../tests/purview/test_chat_middleware.py | 75 ++++++++++ .../purview/tests/purview/test_middleware.py | 112 +++++++++++++++ 5 files changed, 418 insertions(+), 12 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs index f5c7477bd56..fa59faf61fc 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs @@ -35,6 +35,55 @@ public PurviewWrapper(IScopedContentProcessor scopedProcessor, PurviewSettings p this._backgroundJobRunner = backgroundJobRunner; } + /// + /// Builds the replacement response returned when policy blocks a response that was produced. + /// + /// The response that was evaluated and blocked. + /// The notice to return in place of the blocked content. + /// A response carrying the original control metadata with only the messages replaced. + /// + /// The control fields are carried over so the caller can still identify the call and resume it; + /// only the messages are replaced. is deliberately + /// not carried, because it holds the content that was blocked. + /// + private static ChatResponse BuildBlockedResponse(ChatResponse evaluated, string blockedMessage) => + new(new ChatMessage(ChatRole.System, blockedMessage)) + { + ResponseId = evaluated.ResponseId, + ConversationId = evaluated.ConversationId, + ModelId = evaluated.ModelId, + CreatedAt = evaluated.CreatedAt, + FinishReason = evaluated.FinishReason, + Usage = evaluated.Usage, + ContinuationToken = evaluated.ContinuationToken, + AdditionalProperties = evaluated.AdditionalProperties, + }; + + /// + /// Builds the replacement response returned when policy blocks a response that was produced. + /// + /// The response that was evaluated and blocked. + /// The notice to return in place of the blocked content. + /// A response carrying the original control metadata with only the messages replaced. + /// + /// The control fields are carried over so the caller can still identify the run and resume it; + /// only the messages are replaced. is deliberately + /// not carried, because it holds the content that was blocked. + /// + private static AgentResponse BuildBlockedResponse(AgentResponse evaluated, string blockedMessage) => +#pragma warning disable MEAI001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed. + new(new ChatMessage(ChatRole.System, blockedMessage)) + { + ResponseId = evaluated.ResponseId, + AgentId = evaluated.AgentId, + CreatedAt = evaluated.CreatedAt, + FinishReason = evaluated.FinishReason, + Usage = evaluated.Usage, + ContinuationToken = evaluated.ContinuationToken, + AdditionalProperties = evaluated.AdditionalProperties, + }; +#pragma warning restore MEAI001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed. + private static string GetSessionIdFromAgentSession(AgentSession? session, IEnumerable messages) { if (session is ChatClientAgentSession chatClientAgentSession && @@ -107,7 +156,7 @@ public async Task ProcessChatContentAsync(IEnumerable this._logger.LogInformation("Response blocked by policy. Sending message: {Message}", this._purviewSettings.BlockedResponseMessage); } - return new ChatResponse(new ChatMessage(ChatRole.System, this._purviewSettings.BlockedResponseMessage)); + return BuildBlockedResponse(response, this._purviewSettings.BlockedResponseMessage); } } catch (Exception ex) @@ -197,7 +246,7 @@ public async Task ProcessAgentContentAsync(IEnumerable + { + new(ChatRole.User, "Test message") + }; + var mockChatClient = new Mock(); + var createdAt = DateTimeOffset.UtcNow; + var innerResponse = new ChatResponse(new ChatMessage(ChatRole.Assistant, "Sensitive response")) + { + ResponseId = "resp-1", + ConversationId = "conv-1", + ModelId = "model-1", + CreatedAt = createdAt, + FinishReason = ChatFinishReason.Stop, + Usage = new UsageDetails { InputTokenCount = 3, OutputTokenCount = 5 }, + AdditionalProperties = new AdditionalPropertiesDictionary { ["custom"] = "value" }, + RawRepresentation = new { ProviderPayload = "Sensitive response" } + }; + + mockChatClient.Setup(x => x.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(innerResponse); + + this._mockProcessor.Setup(x => x.ProcessMessagesAsync( + It.IsAny>(), + It.IsAny(), + Activity.UploadText, + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync((false, "user-123")); + + this._mockProcessor.Setup(x => x.ProcessMessagesAsync( + It.IsAny>(), + It.IsAny(), + Activity.DownloadText, + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync((true, "user-123")); + + // Act + var result = await this._wrapper.ProcessChatContentAsync(messages, null, mockChatClient.Object, CancellationToken.None); + + // Assert - the call is still identifiable and resumable + Assert.Equal("resp-1", result.ResponseId); + Assert.Equal("conv-1", result.ConversationId); + Assert.Equal("model-1", result.ModelId); + Assert.Equal(createdAt, result.CreatedAt); + Assert.Equal(ChatFinishReason.Stop, result.FinishReason); + Assert.Equal(3, result.Usage?.InputTokenCount); + Assert.Equal("value", result.AdditionalProperties?["custom"]); + + // Assert - only the messages were replaced, and the blocked payload is not handed back + Assert.Single(result.Messages); + Assert.Equal(ChatRole.System, result.Messages[0].Role); + Assert.Equal("Response blocked by policy", result.Messages[0].Text); + Assert.Null(result.RawRepresentation); + } + + [Fact] + public async Task ProcessAgentContentAsync_WithBlockedResponse_PreservesControlMetadataAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var mockAgent = new Mock(); + var createdAt = DateTimeOffset.UtcNow; + var innerResponse = new AgentResponse(new ChatMessage(ChatRole.Assistant, "Sensitive response")) + { + ResponseId = "resp-1", + AgentId = "agent-1", + CreatedAt = createdAt, + FinishReason = ChatFinishReason.Stop, + Usage = new UsageDetails { InputTokenCount = 3, OutputTokenCount = 5 }, + AdditionalProperties = new AdditionalPropertiesDictionary { ["custom"] = "value" }, + RawRepresentation = new { ProviderPayload = "Sensitive response" } + }; + + mockAgent.Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(innerResponse); + + this._mockProcessor.Setup(x => x.ProcessMessagesAsync( + It.IsAny>(), + It.IsAny(), + Activity.UploadText, + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync((false, "user-123")); + + this._mockProcessor.Setup(x => x.ProcessMessagesAsync( + It.IsAny>(), + It.IsAny(), + Activity.DownloadText, + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync((true, "user-123")); + + // Act + var result = await this._wrapper.ProcessAgentContentAsync(messages, null, null, mockAgent.Object, CancellationToken.None); + + // Assert - the run is still identifiable and resumable + Assert.Equal("resp-1", result.ResponseId); + Assert.Equal("agent-1", result.AgentId); + Assert.Equal(createdAt, result.CreatedAt); + Assert.Equal(ChatFinishReason.Stop, result.FinishReason); + Assert.Equal(3, result.Usage?.InputTokenCount); + Assert.Equal("value", result.AdditionalProperties?["custom"]); + + // Assert - only the messages were replaced, and the blocked payload is not handed back + Assert.Single(result.Messages); + Assert.Equal(ChatRole.System, result.Messages[0].Role); + Assert.Equal("Response blocked by policy", result.Messages[0].Text); + Assert.Null(result.RawRepresentation); + } + + #endregion + public void Dispose() { this._wrapper.Dispose(); diff --git a/python/packages/purview/agent_framework_purview/_middleware.py b/python/packages/purview/agent_framework_purview/_middleware.py index 5dcde5d2436..2d3faf08cf1 100644 --- a/python/packages/purview/agent_framework_purview/_middleware.py +++ b/python/packages/purview/agent_framework_purview/_middleware.py @@ -186,12 +186,30 @@ async def process( resolved_user_id, ) if should_block_response: - context.result = self._blocked_response() + context.result = self._blocked_response(context.result) - def _blocked_response(self) -> "AgentResponse[Any]": - """Build the replacement response returned when policy blocks the content.""" + def _blocked_response(self, evaluated: "AgentResponse[Any] | None" = None) -> "AgentResponse[Any]": + """Build the replacement response returned when policy blocks the content. + + When a response was actually produced, its control fields are carried over and only + the messages are replaced, so the caller can still identify the operation and resume + it. The structured value and the raw provider payload are deliberately not carried, + because both hold the content that was blocked. + """ msg = self._settings.get("blocked_response_message", None) or "Response blocked by policy" - return AgentResponse(messages=[Message(role="system", contents=[msg])]) + blocked = Message(role="system", contents=[msg]) + if evaluated is None: + return AgentResponse(messages=[blocked]) + return AgentResponse( + messages=[blocked], + response_id=evaluated.response_id, + agent_id=evaluated.agent_id, + created_at=evaluated.created_at, + finish_reason=cast("Any", evaluated.finish_reason), + usage_details=evaluated.usage_details, + continuation_token=evaluated.continuation_token, + additional_properties=dict(evaluated.additional_properties) if evaluated.additional_properties else None, + ) def _evaluate_stream( self, @@ -227,7 +245,7 @@ async def _transform(final: "AgentResponse[Any]") -> "AgentResponse[Any]": user_id, ) if should_block_response: - return self._blocked_response() + return self._blocked_response(final) # Returning the evaluated response, rather than None, is what makes the # released updates be re-derived from it. A finalizer may produce a response # that is not the assembly of its own updates, so releasing the buffered @@ -333,15 +351,33 @@ async def process( resolved_user_id, ) if should_block_response: - context.result = self._blocked_response() + context.result = self._blocked_response(context.result) - def _blocked_response(self) -> "ChatResponse[Any]": - """Build the replacement response returned when policy blocks the content.""" + def _blocked_response(self, evaluated: "ChatResponse[Any] | None" = None) -> "ChatResponse[Any]": + """Build the replacement response returned when policy blocks the content. + + When a response was actually produced, its control fields are carried over and only + the messages are replaced, so the caller can still identify the call and resume it. + The structured value and the raw provider payload are deliberately not carried, + because both hold the content that was blocked. + """ blocked_message = Message( role="system", contents=[self._settings.get("blocked_response_message", None) or "Response blocked by policy"], ) - return ChatResponse(messages=[blocked_message]) + if evaluated is None: + return ChatResponse(messages=[blocked_message]) + return ChatResponse( + messages=[blocked_message], + response_id=evaluated.response_id, + conversation_id=evaluated.conversation_id, + model=evaluated.model, + created_at=evaluated.created_at, + finish_reason=cast("Any", evaluated.finish_reason), + usage_details=evaluated.usage_details, + continuation_token=evaluated.continuation_token, + additional_properties=dict(evaluated.additional_properties) if evaluated.additional_properties else None, + ) def _evaluate_stream( self, @@ -377,7 +413,7 @@ async def _transform(final: "ChatResponse[Any]") -> "ChatResponse[Any]": user_id, ) if should_block_response: - return self._blocked_response() + return self._blocked_response(final) # Returning the evaluated response, rather than None, is what makes the # released updates be re-derived from it. A finalizer may produce a response # that is not the assembly of its own updates, so releasing the buffered diff --git a/python/packages/purview/tests/purview/test_chat_middleware.py b/python/packages/purview/tests/purview/test_chat_middleware.py index c771d9bc183..a4cb40a9d47 100644 --- a/python/packages/purview/tests/purview/test_chat_middleware.py +++ b/python/packages/purview/tests/purview/test_chat_middleware.py @@ -721,3 +721,78 @@ async def test_allowed_response_still_becomes_history( assert text == "model reply" assert stored_texts(session, provider) == ["hello there", "model reply"] + + +class TestPurviewChatPolicyMiddlewareBlockedResponseEnvelope: + """What a blocked response keeps from the response it replaced. + + Blocking removes the content, not the call. The caller still needs to be able to + identify the call and resume it, so the control fields are carried over and only the + messages are replaced. + """ + + @pytest.fixture + def middleware(self) -> PurviewChatPolicyMiddleware: + credential = AsyncMock() + credential.get_token = AsyncMock(return_value=AccessToken("fake-token", 9999999999)) + return PurviewChatPolicyMiddleware(credential, PurviewSettings(app_name="Test App", tenant_id="test-tenant")) + + @pytest.fixture + def evaluated(self) -> ChatResponse[Any]: + return ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_text(text="confidential")])], + response_id="resp-1", + conversation_id="conv-1", + model="model-1", + created_at="2024-01-01T00:00:00Z", + finish_reason="stop", + continuation_token=cast(Any, {"token": "resume-me"}), + additional_properties={"custom": "value"}, + raw_representation={"provider_payload": "confidential"}, + ) + + def test_control_fields_are_carried_over( + self, middleware: PurviewChatPolicyMiddleware, evaluated: ChatResponse[Any] + ) -> None: + blocked = middleware._blocked_response(evaluated) + + assert blocked.response_id == "resp-1" + assert blocked.conversation_id == "conv-1" + assert blocked.model == "model-1" + assert blocked.created_at == evaluated.created_at + assert blocked.finish_reason == "stop" + assert blocked.continuation_token == {"token": "resume-me"} + assert blocked.additional_properties["custom"] == "value" + + def test_only_the_messages_are_replaced( + self, middleware: PurviewChatPolicyMiddleware, evaluated: ChatResponse[Any] + ) -> None: + blocked = middleware._blocked_response(evaluated) + + assert len(blocked.messages) == 1 + assert blocked.messages[0].role == "system" + assert "confidential" not in blocked.text + + def test_the_provider_payload_is_not_carried_over( + self, middleware: PurviewChatPolicyMiddleware, evaluated: ChatResponse[Any] + ) -> None: + """It holds the content that was blocked, so carrying it would hand that content back.""" + blocked = middleware._blocked_response(evaluated) + + assert blocked.raw_representation is None + + def test_properties_are_copied_rather_than_shared( + self, middleware: PurviewChatPolicyMiddleware, evaluated: ChatResponse[Any] + ) -> None: + blocked = middleware._blocked_response(evaluated) + blocked.additional_properties["custom"] = "tampered" + + assert evaluated.additional_properties["custom"] == "value" + + def test_a_blocked_prompt_has_no_response_to_carry_from(self, middleware: PurviewChatPolicyMiddleware) -> None: + """Nothing was produced, so there is no envelope to preserve.""" + blocked = middleware._blocked_response() + + assert blocked.response_id is None + assert blocked.continuation_token is None + assert blocked.messages[0].role == "system" diff --git a/python/packages/purview/tests/purview/test_middleware.py b/python/packages/purview/tests/purview/test_middleware.py index b4c9ae90117..fe2cedf518e 100644 --- a/python/packages/purview/tests/purview/test_middleware.py +++ b/python/packages/purview/tests/purview/test_middleware.py @@ -916,3 +916,115 @@ async def mock_next() -> None: released = [update async for update in cast(Any, context.result)] assert "".join(update.text for update in released) == "confidential" + + +class TestPurviewPolicyMiddlewareBlockedResponseEnvelope: + """What a blocked response keeps from the response it replaced. + + Blocking removes the content, not the call. The caller still needs to be able to + identify the run and resume it, so the control fields are carried over and only the + messages are replaced. + """ + + @pytest.fixture + def middleware(self) -> PurviewPolicyMiddleware: + credential = AsyncMock() + credential.get_token = AsyncMock(return_value=AccessToken("fake-token", 9999999999)) + return PurviewPolicyMiddleware(credential, PurviewSettings(app_name="Test App", tenant_id="test-tenant")) + + @pytest.fixture + def mock_agent(self) -> MagicMock: + agent = MagicMock() + agent.name = "test-agent" + return agent + + @pytest.fixture + def evaluated(self) -> AgentResponse[Any]: + return AgentResponse( + messages=[Message(role="assistant", contents=[Content.from_text(text="confidential")])], + response_id="resp-1", + agent_id="agent-1", + created_at="2024-01-01T00:00:00Z", + finish_reason="stop", + continuation_token=cast(Any, {"token": "resume-me"}), + additional_properties={"custom": "value"}, + raw_representation={"provider_payload": "confidential"}, + ) + + def test_control_fields_are_carried_over( + self, middleware: PurviewPolicyMiddleware, evaluated: AgentResponse[Any] + ) -> None: + blocked = middleware._blocked_response(evaluated) + + assert blocked.response_id == "resp-1" + assert blocked.agent_id == "agent-1" + assert blocked.created_at == evaluated.created_at + assert blocked.finish_reason == "stop" + assert blocked.continuation_token == {"token": "resume-me"} + assert blocked.additional_properties["custom"] == "value" + + def test_only_the_messages_are_replaced( + self, middleware: PurviewPolicyMiddleware, evaluated: AgentResponse[Any] + ) -> None: + blocked = middleware._blocked_response(evaluated) + + assert len(blocked.messages) == 1 + assert blocked.messages[0].role == "system" + assert "confidential" not in blocked.text + + def test_the_provider_payload_is_not_carried_over( + self, middleware: PurviewPolicyMiddleware, evaluated: AgentResponse[Any] + ) -> None: + """It holds the content that was blocked, so carrying it would hand that content back.""" + blocked = middleware._blocked_response(evaluated) + + assert blocked.raw_representation is None + + def test_properties_are_copied_rather_than_shared( + self, middleware: PurviewPolicyMiddleware, evaluated: AgentResponse[Any] + ) -> None: + blocked = middleware._blocked_response(evaluated) + blocked.additional_properties["custom"] = "tampered" + + assert evaluated.additional_properties["custom"] == "value" + + def test_a_blocked_prompt_has_no_response_to_carry_from(self, middleware: PurviewPolicyMiddleware) -> None: + """Nothing was produced, so there is no envelope to preserve.""" + blocked = middleware._blocked_response() + + assert blocked.response_id is None + assert blocked.continuation_token is None + assert blocked.messages[0].role == "system" + + async def test_a_blocked_streamed_run_keeps_the_envelope( + self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any + ) -> None: + context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["Hello"])]) + context.stream = True + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="confidential")]) + + def finalizer(collected: Any) -> AgentResponse[Any]: + final = AgentResponse.from_updates(collected) + final.continuation_token = cast(Any, {"token": "resume-me"}) + final.response_id = "resp-1" + return final + + with patch.object( + middleware._processor, + "process_messages", + side_effect=[(False, "user-123"), (True, "user-123")], + ): + + async def mock_next() -> None: + context.result = cast(Any, ResponseStream(updates(), finalizer=finalizer)) + + await run_agent_middleware([middleware], context, mock_next) + stream = cast(Any, context.result) + released = [update async for update in stream] + final = await stream.get_final_response() + + assert "confidential" not in "".join(update.text for update in released) + assert final.continuation_token == {"token": "resume-me"} + assert final.response_id == "resp-1"