From 718446068b8d93d3084113abedc5ece117b98433 Mon Sep 17 00:00:00 2001 From: Rishabh4275 Date: Tue, 24 Feb 2026 12:53:54 -0800 Subject: [PATCH] [Purview] Mark responses as responses and fix epoch bug for python long overflow --- .../PurviewWrapper.cs | 4 +- .../PurviewWrapperTests.cs | 62 ++++++++++++++----- .../agent_framework_purview/_middleware.py | 4 +- .../agent_framework_purview/_processor.py | 3 +- .../purview/tests/purview/test_middleware.py | 6 +- 5 files changed, 57 insertions(+), 22 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs index 9b6cdc2ffd..27882375d7 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs @@ -98,7 +98,7 @@ public async Task ProcessChatContentAsync(IEnumerable 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 @@ public async Task ProcessAgentContentAsync(IEnumerable x.ProcessMessagesAsync( It.IsAny>(), It.IsAny(), - It.IsAny(), + Activity.UploadText, It.IsAny(), It.IsAny(), It.IsAny())) @@ -88,15 +88,24 @@ public async Task ProcessChatContentAsync_WithAllowedPromptAndBlockedResponse_Re It.IsAny())) .ReturnsAsync(innerResponse); - this._mockProcessor.SetupSequence(x => x.ProcessMessagesAsync( + // Prompt check uses UploadText, response check uses DownloadText + this._mockProcessor.Setup(x => x.ProcessMessagesAsync( It.IsAny>(), It.IsAny(), - It.IsAny(), + Activity.UploadText, + It.IsAny(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync((false, "user-123")); // Prompt allowed + + this._mockProcessor.Setup(x => x.ProcessMessagesAsync( + It.IsAny>(), + It.IsAny(), + Activity.DownloadText, It.IsAny(), It.IsAny(), It.IsAny())) - .ReturnsAsync((false, "user-123")) // Prompt allowed - .ReturnsAsync((true, "user-123")); // Response blocked + .ReturnsAsync((true, "user-123")); // Response blocked // Act var result = await this._wrapper.ProcessChatContentAsync(messages, null, mockChatClient.Object, CancellationToken.None); @@ -237,14 +246,21 @@ public async Task ProcessChatContentAsync_UsesConversationIdFromOptions_Async() // 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>(), "conversation-123", - It.IsAny(), + Activity.UploadText, + It.IsAny(), + It.IsAny(), + It.IsAny()), Times.Once); + this._mockProcessor.Verify(x => x.ProcessMessagesAsync( + It.IsAny>(), + "conversation-123", + Activity.DownloadText, It.IsAny(), It.IsAny(), - It.IsAny()), Times.Exactly(2)); + It.IsAny()), Times.Once); } #endregion @@ -264,7 +280,7 @@ public async Task ProcessAgentContentAsync_WithBlockedPrompt_ReturnsBlockedMessa this._mockProcessor.Setup(x => x.ProcessMessagesAsync( It.IsAny>(), It.IsAny(), - It.IsAny(), + Activity.UploadText, It.IsAny(), It.IsAny(), It.IsAny())) @@ -306,15 +322,24 @@ public async Task ProcessAgentContentAsync_WithAllowedPromptAndBlockedResponse_R ItExpr.IsAny()) .ReturnsAsync(innerResponse); - this._mockProcessor.SetupSequence(x => x.ProcessMessagesAsync( + // Prompt check uses UploadText, response check uses DownloadText + this._mockProcessor.Setup(x => x.ProcessMessagesAsync( It.IsAny>(), It.IsAny(), - It.IsAny(), + Activity.UploadText, It.IsAny(), It.IsAny(), It.IsAny())) - .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>(), + It.IsAny(), + Activity.DownloadText, + It.IsAny(), + It.IsAny(), + It.IsAny())) + .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 async Task ProcessAgentContentAsync_ExtractsThreadIdFromMessageAdditional this._mockProcessor.Verify(x => x.ProcessMessagesAsync( It.IsAny>(), "conversation-from-props", - It.IsAny(), + Activity.UploadText, + It.IsAny(), + It.IsAny(), + It.IsAny()), Times.Once); + this._mockProcessor.Verify(x => x.ProcessMessagesAsync( + It.IsAny>(), + "conversation-from-props", + Activity.DownloadText, It.IsAny(), It.IsAny(), - It.IsAny()), Times.Exactly(2)); + It.IsAny()), Times.Once); } [Fact] diff --git a/python/packages/purview/agent_framework_purview/_middleware.py b/python/packages/purview/agent_framework_purview/_middleware.py index 2296122135..55619d0a39 100644 --- a/python/packages/purview/agent_framework_purview/_middleware.py +++ b/python/packages/purview/agent_framework_purview/_middleware.py @@ -106,7 +106,7 @@ async def process( 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 @@ async def process( 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 diff --git a/python/packages/purview/agent_framework_purview/_processor.py b/python/packages/purview/agent_framework_purview/_processor.py index a7fc030cbf..e911fae7a5 100644 --- a/python/packages/purview/agent_framework_purview/_processor.py +++ b/python/packages/purview/agent_framework_purview/_processor.py @@ -158,7 +158,8 @@ async def _map_messages( 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) diff --git a/python/packages/purview/tests/purview/test_middleware.py b/python/packages/purview/tests/purview/test_middleware.py index 3a34c48344..b49b5d46a0 100644 --- a/python/packages/purview/tests/purview/test_middleware.py +++ b/python/packages/purview/tests/purview/test_middleware.py @@ -148,8 +148,10 @@ async def mock_next() -> None: 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