diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs index c23553f1855..da7aac682d6 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs @@ -8,10 +8,6 @@ namespace Microsoft.Agents.AI.Purview.Models.Jobs; /// /// Class representing a job that refreshes the protection scopes cache in the background. /// -/// -/// Used by the parallel protection scopes retrieval path to warm the cache without blocking the -/// foreground ProcessContent call. -/// internal sealed class ScopeRetrievalJob : BackgroundJobBase { /// diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs index fab7c28d9ae..f49f1239a89 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs @@ -195,21 +195,22 @@ private async Task ProcessContentWithProtectionScopesAsy ProtectionScopesResponse? cacheResponse = await this._cacheProvider.GetAsync(cacheKey, cancellationToken).ConfigureAwait(false); - if (cacheResponse != null) + if (cacheResponse == null) { - return await this.ProcessWithCachedScopesAsync(pcRequest, cacheResponse, cacheKey, cancellationToken).ConfigureAwait(false); - } + pcRequest.ProcessInline = true; + try + { + this._channelHandler.QueueJob(new ScopeRetrievalJob(psRequest, cacheKey, pcRequest)); + } + catch (PurviewJobException) + { + // QueueJob logs admission failures. Scope refresh is best effort. + } - try - { - this._channelHandler.QueueJob(new ScopeRetrievalJob(psRequest, cacheKey, pcRequest)); - } - catch (PurviewJobException) - { - // QueueJob already logs failures. Scope warmup is best effort; don't block ProcessContent. + return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions: null, cancellationToken).ConfigureAwait(false); } - return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions: null, cancellationToken).ConfigureAwait(false); + return await this.ProcessWithCachedScopesAsync(pcRequest, cacheResponse, cacheKey, cancellationToken).ConfigureAwait(false); } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs index d1b9d535589..b59cc7a3929 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs @@ -627,7 +627,7 @@ public async Task ProcessMessagesAsync_UsesTokenUserIdBeforeMessageAdditionalPro this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( It.IsAny(), It.IsAny())) - .ReturnsAsync(new ProtectionScopesResponse { Scopes = [] }); + .ReturnsAsync(CreateApplicableProtectionScopesResponse()); this._mockPurviewClient.Setup(x => x.ProcessContentAsync( It.IsAny(), It.IsAny())) .ReturnsAsync(new ProcessContentResponse { PolicyActions = [] }); @@ -668,7 +668,7 @@ public async Task ProcessMessagesAsync_UsesTokenUserIdBeforeAuthorName_Async() this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( It.IsAny(), It.IsAny())) - .ReturnsAsync(new ProtectionScopesResponse { Scopes = [] }); + .ReturnsAsync(CreateApplicableProtectionScopesResponse()); this._mockPurviewClient.Setup(x => x.ProcessContentAsync( It.IsAny(), It.IsAny())) .ReturnsAsync(new ProcessContentResponse { PolicyActions = [] }); @@ -703,6 +703,9 @@ public async Task ProcessMessagesAsync_UsesProvidedUserId_WhenTokenUserIdIsEmpty It.IsAny(), It.IsAny())) .ReturnsAsync((ProtectionScopesResponse?)null); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(CreateApplicableProtectionScopesResponse()); this._mockPurviewClient.Setup(x => x.ProcessContentAsync( It.IsAny(), It.IsAny())) .ReturnsAsync(new ProcessContentResponse { PolicyActions = [] }); @@ -750,7 +753,7 @@ public async Task ProcessMessagesAsync_ExtractsUserIdFromMessageAuthorName_WhenV } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_QueuesScopeRetrievalJobAndCallsProcessContentAsync() + public async Task ProcessMessagesAsync_CacheMiss_CallsProcessContentInlineAndQueuesScopeRetrievalAsync() { // Arrange var messages = new List @@ -767,19 +770,26 @@ public async Task ProcessMessagesAsync_CacheMiss_QueuesScopeRetrievalJobAndCalls .ReturnsAsync((ProtectionScopesResponse?)null); this._mockPurviewClient.Setup(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny())) + It.Is(request => + request.ScopeIdentifier == null && + request.ProcessInline), + It.IsAny())) .ReturnsAsync(new ProcessContentResponse()); // Act await this._processor.ProcessMessagesAsync( messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); - // Assert: ProcessContent runs in the foreground; GetProtectionScopes is queued as a background job. + // Assert this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Once); + It.Is(request => + request.ScopeIdentifier == null && + request.ProcessInline), + It.IsAny()), Times.Once); this._mockPurviewClient.Verify(x => x.GetProtectionScopesAsync( It.IsAny(), It.IsAny()), Times.Never); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + this._mockChannelHandler.Verify(x => x.QueueJob( + It.Is(job => job.ProcessContentRequest.ProcessInline)), Times.Once); } [Fact] @@ -799,17 +809,16 @@ public async Task ProcessMessagesAsync_CacheMiss_WithProcessContentBlockAction_R It.IsAny(), It.IsAny())) .ReturnsAsync((ProtectionScopesResponse?)null); - var pcResponse = new ProcessContentResponse - { - PolicyActions = - [ - new() { Action = DlpAction.BlockAccess } - ] - }; - this._mockPurviewClient.Setup(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(pcResponse); + It.Is(request => request.ProcessInline), + It.IsAny())) + .ReturnsAsync(new ProcessContentResponse + { + PolicyActions = + [ + new() { Action = DlpAction.BlockAccess } + ] + }); // Act var result = await this._processor.ProcessMessagesAsync( @@ -817,11 +826,13 @@ public async Task ProcessMessagesAsync_CacheMiss_WithProcessContentBlockAction_R // Assert Assert.True(result.shouldBlock); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.Is(request => request.ProcessInline), + It.IsAny()), Times.Once); } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_StillCallsProcessContentWhenScopeJobCannotQueueAsync() + public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalCannotQueue_CallsProcessContentInlineAsync() { // Arrange var messages = new List @@ -841,17 +852,19 @@ public async Task ProcessMessagesAsync_CacheMiss_StillCallsProcessContentWhenSco .Throws(new PurviewJobException("queue unavailable")); this._mockPurviewClient.Setup(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny())) + It.Is(request => request.ProcessInline), + It.IsAny())) .ReturnsAsync(new ProcessContentResponse()); // Act await this._processor.ProcessMessagesAsync( messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); - // Assert: scope warmup is attempted, and ProcessContent still runs when it can't be queued. + // Assert this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Once); + It.Is(request => request.ProcessInline), + It.IsAny()), Times.Once); } [Fact] @@ -878,11 +891,11 @@ await Assert.ThrowsAsync(() => this._mockPurviewClient.Verify(x => x.ProcessContentAsync( It.IsAny(), It.IsAny()), Times.Never); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); } [Fact] - public async Task BackgroundJobRunner_ScopeRetrievalPaymentRequired_CachesForSubsequentCallsAsync() + public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_CachesAndQueuesContentActivityAsync() { // Arrange Func, Task>? runner = null; @@ -890,40 +903,65 @@ public async Task BackgroundJobRunner_ScopeRetrievalPaymentRequired_CachesForSub Mock purviewClient = new(); Mock cacheProvider = new(); PurviewSettings settings = new("TestApp") { MaxConcurrentJobConsumers = 1 }; - ProtectionScopesRequest request = new("user-123", "tenant-123") - { - Activities = ProtectionScopeActivities.UploadText, - Locations = - [ - new("microsoft.graph.policyLocationApplication", "app-123") - ] - }; + ProtectionScopesRequest request = CreateProtectionScopesRequest(); ProtectionScopesCacheKey cacheKey = new(request); + ScopeRetrievalJob job = new(request, cacheKey, CreateProcessContentRequest()); + ProtectionScopesResponse response = new() { ScopeIdentifier = "scope-123", Scopes = [] }; Channel channel = Channel.CreateUnbounded(); channelHandler.Setup(x => x.AddRunner(It.IsAny, Task>>())) .Callback, Task>>(callback => runner = callback); + purviewClient.Setup(x => x.GetProtectionScopesAsync(It.IsAny(), It.IsAny())) + .ReturnsAsync(response); + _ = new BackgroundJobRunner(channelHandler.Object, purviewClient.Object, cacheProvider.Object, NullLogger.Instance, settings); + + // Act + Assert.NotNull(runner); + await channel.Writer.WriteAsync(job); + channel.Writer.Complete(); + await runner(channel); + + // Assert + cacheProvider.Verify(x => x.SetAsync(cacheKey, response, It.IsAny()), Times.Once); + channelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + } + + [Fact] + public async Task BackgroundJobRunner_ScopeRetrievalApplicableScope_CachesWithoutContentActivityAsync() + { + // Arrange + Func, Task>? runner = null; + Mock channelHandler = new(); + Mock purviewClient = new(); + Mock cacheProvider = new(); + PurviewSettings settings = new("TestApp") { MaxConcurrentJobConsumers = 1 }; + ProtectionScopesRequest request = CreateProtectionScopesRequest(); + ProtectionScopesCacheKey cacheKey = new(request); + ScopeRetrievalJob job = new(request, cacheKey, CreateProcessContentRequest()); + ProtectionScopesResponse response = CreateApplicableProtectionScopesResponse(); + Channel channel = Channel.CreateUnbounded(); + + channelHandler.Setup(x => x.AddRunner(It.IsAny, Task>>())) + .Callback, Task>>(callback => runner = callback); purviewClient.Setup(x => x.GetProtectionScopesAsync(It.IsAny(), It.IsAny())) - .ThrowsAsync(new PurviewPaymentRequiredException("Payment required")); + .ReturnsAsync(response); _ = new BackgroundJobRunner(channelHandler.Object, purviewClient.Object, cacheProvider.Object, NullLogger.Instance, settings); // Act Assert.NotNull(runner); - await channel.Writer.WriteAsync(new ScopeRetrievalJob(request, cacheKey, CreateProcessContentRequest())); + await channel.Writer.WriteAsync(job); channel.Writer.Complete(); await runner(channel); // Assert - cacheProvider.Verify(x => x.SetAsync( - It.Is(key => key.TenantId == "tenant-123"), - It.Is(entry => entry.Message == "Payment required"), - It.IsAny()), Times.Once); + cacheProvider.Verify(x => x.SetAsync(cacheKey, response, It.IsAny()), Times.Once); + channelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); } [Fact] - public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_QueuesContentActivityJobAsync() + public async Task BackgroundJobRunner_ScopeRetrievalPaymentRequired_CachesForSubsequentCallsAsync() { // Arrange Func, Task>? runner = null; @@ -937,9 +975,8 @@ public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_QueuesCon channelHandler.Setup(x => x.AddRunner(It.IsAny, Task>>())) .Callback, Task>>(callback => runner = callback); - purviewClient.Setup(x => x.GetProtectionScopesAsync(It.IsAny(), It.IsAny())) - .ReturnsAsync(new ProtectionScopesResponse { Scopes = [] }); + .ThrowsAsync(new PurviewPaymentRequiredException("Payment required")); _ = new BackgroundJobRunner(channelHandler.Object, purviewClient.Object, cacheProvider.Object, NullLogger.Instance, settings); @@ -950,7 +987,10 @@ public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_QueuesCon await runner(channel); // Assert - channelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + cacheProvider.Verify(x => x.SetAsync( + It.Is(key => key.TenantId == "tenant-123"), + It.Is(entry => entry.Message == "Payment required"), + It.IsAny()), Times.Once); } #endregion @@ -969,6 +1009,26 @@ private static ProtectionScopesRequest CreateProtectionScopesRequest() }; } + private static ProtectionScopesResponse CreateApplicableProtectionScopesResponse(ExecutionMode executionMode = ExecutionMode.EvaluateInline) + { + return new ProtectionScopesResponse + { + ScopeIdentifier = "scope-123", + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("microsoft.graph.policyLocationApplication", "app-123") + ], + ExecutionMode = executionMode + } + ] + }; + } + private static ProcessContentRequest CreateProcessContentRequest() { PurviewTextContent content = new("Test content"); diff --git a/python/packages/purview/agent_framework_purview/_processor.py b/python/packages/purview/agent_framework_purview/_processor.py index 8fda0acf34c..d6c24a078f2 100644 --- a/python/packages/purview/agent_framework_purview/_processor.py +++ b/python/packages/purview/agent_framework_purview/_processor.py @@ -230,6 +230,7 @@ async def _process_with_scopes(self, pc_request: ProcessContentRequest) -> Proce if cached_ps_resp is not None and isinstance(cached_ps_resp, ProtectionScopesResponse): return await self._process_with_cached_scopes(pc_request, cached_ps_resp, cache_key) + pc_request.process_inline = True task = asyncio.create_task(self._refresh_protection_scopes_background(ps_req, cache_key, pc_request)) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) diff --git a/python/packages/purview/tests/purview/test_processor.py b/python/packages/purview/tests/purview/test_processor.py index e147e8666ab..9d7388a416b 100644 --- a/python/packages/purview/tests/purview/test_processor.py +++ b/python/packages/purview/tests/purview/test_processor.py @@ -283,6 +283,7 @@ async def test_process_with_scopes_calls_client_methods( # On cache miss, ProcessContent runs in the foreground and the response is returned. assert response.id == "response-123" mock_client.process_content.assert_called_once() + assert mock_client.process_content.call_args.args[0].process_inline is True # Protection scopes are refreshed in a background task. await asyncio.gather(*list(processor._background_tasks))