[Purview] Mark responses as responses and fix epoch bug for python long overflow (#4225)

This commit is contained in:
Rishabh Chawla
2026-02-25 10:40:36 -08:00
committed by GitHub
Unverified
parent 6138487888
commit 2e26bb9387
5 changed files with 57 additions and 22 deletions
@@ -98,7 +98,7 @@ internal sealed class PurviewWrapper : IDisposable
try
{
(bool shouldBlockResponse, _) = await this._scopedProcessor.ProcessMessagesAsync(response.Messages, options?.ConversationId, Activity.UploadText, this._purviewSettings, resolvedUserId, cancellationToken).ConfigureAwait(false);
(bool shouldBlockResponse, _) = await this._scopedProcessor.ProcessMessagesAsync(response.Messages, options?.ConversationId, Activity.DownloadText, this._purviewSettings, resolvedUserId, cancellationToken).ConfigureAwait(false);
if (shouldBlockResponse)
{
if (this._logger.IsEnabled(LogLevel.Information))
@@ -186,7 +186,7 @@ internal sealed class PurviewWrapper : IDisposable
sessionIdResponse = sessionId;
}
}
(bool shouldBlockResponse, _) = await this._scopedProcessor.ProcessMessagesAsync(response.Messages, sessionIdResponse, Activity.UploadText, this._purviewSettings, resolvedUserId, cancellationToken).ConfigureAwait(false);
(bool shouldBlockResponse, _) = await this._scopedProcessor.ProcessMessagesAsync(response.Messages, sessionIdResponse, Activity.DownloadText, this._purviewSettings, resolvedUserId, cancellationToken).ConfigureAwait(false);
if (shouldBlockResponse)
{
@@ -51,7 +51,7 @@ public sealed class PurviewWrapperTests : IDisposable
this._mockProcessor.Setup(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<string>(),
It.IsAny<Activity>(),
Activity.UploadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()))
@@ -88,15 +88,24 @@ public sealed class PurviewWrapperTests : IDisposable
It.IsAny<CancellationToken>()))
.ReturnsAsync(innerResponse);
this._mockProcessor.SetupSequence(x => x.ProcessMessagesAsync(
// Prompt check uses UploadText, response check uses DownloadText
this._mockProcessor.Setup(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<string>(),
It.IsAny<Activity>(),
Activity.UploadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync((false, "user-123")) // Prompt allowed
.ReturnsAsync((true, "user-123")); // Response blocked
.ReturnsAsync((false, "user-123")); // Prompt allowed
this._mockProcessor.Setup(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<string>(),
Activity.DownloadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync((true, "user-123")); // Response blocked
// Act
var result = await this._wrapper.ProcessChatContentAsync(messages, null, mockChatClient.Object, CancellationToken.None);
@@ -237,14 +246,21 @@ public sealed class PurviewWrapperTests : IDisposable
// Act
await this._wrapper.ProcessChatContentAsync(messages, options, mockChatClient.Object, CancellationToken.None);
// Assert
// Assert - verify prompt uses UploadText and response uses DownloadText
this._mockProcessor.Verify(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
"conversation-123",
It.IsAny<Activity>(),
Activity.UploadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()), Times.Exactly(2));
It.IsAny<CancellationToken>()), Times.Once);
this._mockProcessor.Verify(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
"conversation-123",
Activity.DownloadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()), Times.Once);
}
#endregion
@@ -264,7 +280,7 @@ public sealed class PurviewWrapperTests : IDisposable
this._mockProcessor.Setup(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<string>(),
It.IsAny<Activity>(),
Activity.UploadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()))
@@ -306,15 +322,24 @@ public sealed class PurviewWrapperTests : IDisposable
ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(innerResponse);
this._mockProcessor.SetupSequence(x => x.ProcessMessagesAsync(
// Prompt check uses UploadText, response check uses DownloadText
this._mockProcessor.Setup(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<string>(),
It.IsAny<Activity>(),
Activity.UploadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync((false, "user-123")) // Prompt allowed
.ReturnsAsync((true, "user-123")); // Response blocked
.ReturnsAsync((false, "user-123")); // Prompt allowed
this._mockProcessor.Setup(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<string>(),
Activity.DownloadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync((true, "user-123")); // Response blocked
// Act
var result = await this._wrapper.ProcessAgentContentAsync(messages, null, null, mockAgent.Object, CancellationToken.None);
@@ -472,10 +497,17 @@ public sealed class PurviewWrapperTests : IDisposable
this._mockProcessor.Verify(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
"conversation-from-props",
It.IsAny<Activity>(),
Activity.UploadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()), Times.Exactly(2));
It.IsAny<CancellationToken>()), Times.Once);
this._mockProcessor.Verify(x => x.ProcessMessagesAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
"conversation-from-props",
Activity.DownloadText,
It.IsAny<PurviewSettings>(),
It.IsAny<string>(),
It.IsAny<CancellationToken>()), Times.Once);
}
[Fact]
@@ -106,7 +106,7 @@ class PurviewPolicyMiddleware(AgentMiddleware):
if context.result and not context.stream:
should_block_response, _ = await self._processor.process_messages(
context.result.messages, # type: ignore[union-attr]
Activity.UPLOAD_TEXT,
Activity.DOWNLOAD_TEXT,
session_id=session_id,
user_id=resolved_user_id,
)
@@ -210,7 +210,7 @@ class PurviewChatPolicyMiddleware(ChatMiddleware):
messages = getattr(result_obj, "messages", None)
if messages:
should_block_response, _ = await self._processor.process_messages(
messages, Activity.UPLOAD_TEXT, session_id=session_id_response, user_id=resolved_user_id
messages, Activity.DOWNLOAD_TEXT, session_id=session_id_response, user_id=resolved_user_id
)
if should_block_response:
from agent_framework import ChatResponse, Message
@@ -158,7 +158,8 @@ class ScopedContentProcessor:
name=f"Agent Framework Message {message_id}",
is_truncated=False,
correlation_id=correlation_id,
sequence_number=time.time_ns(),
# This would be c# ticks equivalent and needs to fit inside c# long
sequence_number=time.time_ns() // 100 + 621355968000000000,
)
activity_meta = ActivityMetadata(activity=activity)
@@ -148,8 +148,10 @@ class TestPurviewPolicyMiddleware:
await middleware.process(context, mock_next)
assert mock_process.call_count == 2
for call in mock_process.call_args_list:
assert call[0][1] == Activity.UPLOAD_TEXT
# First call (pre-check) should be UPLOAD_TEXT for user prompt
assert mock_process.call_args_list[0][0][1] == Activity.UPLOAD_TEXT
# 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