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/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
+/// 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/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewWrapperTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewWrapperTests.cs
index 08910060a28..6406afd0ce2 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewWrapperTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewWrapperTests.cs
@@ -605,6 +605,140 @@ public async Task ProcessAgentContentAsync_PassesResolvedUserId_ToResponseProces
#endregion
+ #region Blocked response envelope tests
+
+ [Fact]
+ public async Task ProcessChatContentAsync_WithBlockedResponse_PreservesControlMetadataAsync()
+ {
+ // Arrange
+ var messages = new List
+ {
+ 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/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py
index e2d419fa9e4..830ddf2faae 100644
--- a/python/packages/core/agent_framework/_agent_hooks.py
+++ b/python/packages/core/agent_framework/_agent_hooks.py
@@ -825,42 +825,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.
@@ -1313,7 +1277,7 @@ async def _release_error(exc: BaseException) -> None:
await self._emit_shutdown(state, "cancelled" if isinstance(exc, asyncio.CancelledError) else "error")
context._stream_terminal_result_transforms.append(_transform) # pyright: ignore[reportPrivateUsage]
- context._stream_terminal_result_to_updates = _agent_updates_from_response # pyright: ignore[reportPrivateUsage]
+ context._stream_terminal_result_to_updates = AgentResponse.to_updates # pyright: ignore[reportPrivateUsage]
context._stream_terminal_result_is_authoritative = True # pyright: ignore[reportPrivateUsage]
context._stream_release_hooks.append(_release) # pyright: ignore[reportPrivateUsage]
context._stream_release_error_hooks.append(_release_error) # pyright: ignore[reportPrivateUsage]
@@ -1443,7 +1407,7 @@ def _release_error(_: BaseException) -> None:
gate_handle.drop()
context._stream_terminal_result_transforms.append(_transform) # pyright: ignore[reportPrivateUsage]
- context._stream_terminal_result_to_updates = _chat_updates_from_response # pyright: ignore[reportPrivateUsage]
+ context._stream_terminal_result_to_updates = ChatResponse.to_updates # pyright: ignore[reportPrivateUsage]
context._stream_terminal_result_is_authoritative = True # pyright: ignore[reportPrivateUsage]
context._stream_release_hooks.append(_release) # pyright: ignore[reportPrivateUsage]
context._stream_release_error_hooks.append(_release_error) # pyright: ignore[reportPrivateUsage]
diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py
index e148e1d2e4e..f655161c230 100644
--- a/python/packages/core/agent_framework/_types.py
+++ b/python/packages/core/agent_framework/_types.py
@@ -2239,6 +2239,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.
@@ -2764,6 +2784,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."""
@@ -3176,6 +3244,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 e607e1b9e34..b5840094aae 100644
--- a/python/packages/core/tests/core/test_types.py
+++ b/python/packages/core/tests/core/test_types.py
@@ -1570,6 +1570,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=cast(Any, {"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": "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
@@ -1889,6 +1932,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=cast(Any, {"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": "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/purview/README.md b/python/packages/purview/README.md
index 3b596a75c13..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
@@ -215,6 +217,32 @@ 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 |
+| 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
+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 +284,18 @@ 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.
+
+In either list, attach the Purview middleware last — see
+[Middleware order matters](#middleware-order-matters).
---
## 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,10 +304,98 @@ 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.
+### 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
@@ -350,12 +470,44 @@ 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.
+
+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
+
+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`.
+
+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.
-- **Streaming Responses**: Post-response policy evaluation presently applies only to non-streaming chat responses.
+- **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 5cdd2649b15..2d3faf08cf1 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,55 @@
logger = logging.getLogger("agent_framework.purview")
+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 +138,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 +162,100 @@ 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._evaluate_stream(
+ context,
+ cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", context.result),
+ session_id_response,
+ resolved_user_id,
+ )
+ return
+
+ if 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(context.result)
+
+ 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"
+ 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,
+ 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 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 _transform(final: "AgentResponse[Any]") -> "AgentResponse[Any]":
+ should_block_response = await _should_block_response(
+ self._processor,
+ self._settings,
+ final,
+ session_id,
+ user_id,
+ )
+ if should_block_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
+ # 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):
@@ -155,9 +266,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 +308,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 +327,100 @@ 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._evaluate_stream(
+ context,
+ cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", context.result),
+ session_id_response,
+ resolved_user_id,
+ )
+ return
+
+ if 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(context.result)
+
+ 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"],
+ )
+ 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,
+ 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 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 _transform(final: "ChatResponse[Any]") -> "ChatResponse[Any]":
+ should_block_response = await _should_block_response(
+ self._processor,
+ self._settings,
+ final,
+ session_id,
+ user_id,
+ )
+ if should_block_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
+ # 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..3aec2b33490 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,57 @@ def _create_request(
)
return _create_request
+
+
+@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
+ 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.
+ """
+
+ 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
+
+ return await AgentMiddlewarePipeline(*items).execute(context, final_handler)
+
+ return _run
+
+
+@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.
+ """
+
+ 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
+
+ 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 3c6056d3443..a4cb40a9d47 100644
--- a/python/packages/purview/tests/purview/test_chat_middleware.py
+++ b/python/packages/purview/tests/purview/test_chat_middleware.py
@@ -1,12 +1,24 @@
# 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 (
+ Agent,
+ AgentSession,
+ ChatContext,
+ ChatResponse,
+ ChatResponseUpdate,
+ Content,
+ InMemoryHistoryProvider,
+ Message,
+ MiddlewareTermination,
+ ResponseStream,
+)
from azure.core.credentials import AccessToken
from agent_framework_purview import PurviewChatPolicyMiddleware, PurviewSettings
@@ -100,7 +112,10 @@ 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, run_chat_middleware: Any
+ ) -> 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 +125,181 @@ 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 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
+ 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, run_chat_middleware: Any
+ ) -> 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 = cast(Any, ResponseStream(updates(), finalizer=ChatResponse.from_updates))
+
+ 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
+ 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, run_chat_middleware: Any
+ ) -> 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 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)
+ assert released_text == "benign"
+ assert "confidential" not in released_text
+
+ 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()
+ 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=cast(Any, {"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 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"
+ assert released[-1].conversation_id == "conv-1"
+ assert released[-1].model == "model-1"
+ assert released[-1].continuation_token == {"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, run_chat_middleware: Any
+ ) -> 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:
- streaming_context.result = MagicMock()
+ inner = ResponseStream(updates(), finalizer=ChatResponse.from_updates)
+ inner.with_cleanup_hook(_mark_closed)
+ streaming_context.result = cast(Any, inner)
+
+ await run_chat_middleware(middleware, streaming_context, mock_next)
+ await cast(Any, streaming_context.result).close()
- await middleware.process(streaming_context, mock_next)
- assert mock_proc.call_count == 1
+ assert closed is True
async def test_chat_middleware_handles_post_check_exception(
self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext
@@ -263,19 +446,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:
@@ -423,3 +619,180 @@ 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"]
+
+
+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 5b624367696..fe2cedf518e 100644
--- a/python/packages/purview/tests/purview/test_middleware.py
+++ b/python/packages/purview/tests/purview/test_middleware.py
@@ -2,11 +2,25 @@
"""Tests for Purview middleware."""
+import asyncio
+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 AgentContext, AgentResponse, AgentSession, Message, MiddlewareTermination
+from agent_framework import (
+ Agent,
+ AgentContext,
+ AgentMiddleware,
+ AgentResponse,
+ AgentResponseUpdate,
+ AgentSession,
+ Content,
+ InMemoryHistoryProvider,
+ Message,
+ MiddlewareTermination,
+ ResponseStream,
+)
from azure.core.credentials import AccessToken
from agent_framework_purview import PurviewPolicyMiddleware, PurviewSettings
@@ -156,21 +170,204 @@ 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(
- self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
+ async def test_middleware_streaming_response_is_evaluated_and_blocked(
+ self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock, run_agent_middleware: Any
) -> 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 = AgentResponse(messages=[Message(role="assistant", contents=["streaming"])])
+ 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
+ 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, 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"])])
+ 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 = cast(Any, ResponseStream(updates(), finalizer=AgentResponse.from_updates))
+
+ await run_agent_middleware(middleware, 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].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, run_agent_middleware: Any
+ ) -> 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 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)
+ 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, 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"])])
+ 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=cast(Any, {"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 run_agent_middleware(middleware, 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": "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, 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"])])
+ 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 run_agent_middleware(middleware, 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, 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"])])
+ 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()
- assert mock_proc.call_count == 1
+ 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 run_agent_middleware(middleware, 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
@@ -431,3 +628,403 @@ 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, 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"])])
+ 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, 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"])])
+ 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, 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"])])
+ 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, run_agent_middleware: Any
+ ) -> 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, run_agent_middleware: Any
+ ) -> 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"
+
+
+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"