From e41043d05dc98901cddb5224a2aa4e19dbf954c1 Mon Sep 17 00:00:00 2001 From: Eoin Doherty Date: Fri, 28 Aug 2026 14:23:58 -0700 Subject: [PATCH 1/4] Fix Purview scope cache miss bypass Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../BackgroundJobRunner.cs | 31 +- .../Models/Jobs/ScopeRetrievalJob.cs | 44 --- .../ScopedContentProcessor.cs | 85 ++++- .../ScopedContentProcessorTests.cs | 322 ++++++++++++------ 4 files changed, 303 insertions(+), 179 deletions(-) delete mode 100644 dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs b/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs index 03f73d8007a..85a4fa54c3c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs @@ -1,14 +1,10 @@ // Copyright (c) Microsoft. All rights reserved. using System; -using System.Collections.Generic; using System.Threading; using System.Threading.Channels; using System.Threading.Tasks; -using Microsoft.Agents.AI.Purview.Models.Common; using Microsoft.Agents.AI.Purview.Models.Jobs; -using Microsoft.Agents.AI.Purview.Models.Requests; -using Microsoft.Agents.AI.Purview.Models.Responses; using Microsoft.Extensions.Logging; namespace Microsoft.Agents.AI.Purview; @@ -20,7 +16,6 @@ internal sealed class BackgroundJobRunner : IBackgroundJobRunner { private readonly IChannelHandler _channelHandler; private readonly IPurviewClient _purviewClient; - private readonly ICacheProvider _cacheProvider; private readonly ILogger _logger; /// @@ -28,14 +23,12 @@ internal sealed class BackgroundJobRunner : IBackgroundJobRunner /// /// The channel handler used to manage job channels. /// The Purview client used to send requests to Purview. - /// The cache provider used to store protection scopes results. /// The logger used to log information about background jobs. /// The settings used to configure Purview client behavior. - public BackgroundJobRunner(IChannelHandler channelHandler, IPurviewClient purviewClient, ICacheProvider cacheProvider, ILogger logger, PurviewSettings purviewSettings) + public BackgroundJobRunner(IChannelHandler channelHandler, IPurviewClient purviewClient, ILogger logger, PurviewSettings purviewSettings) { this._channelHandler = channelHandler; this._purviewClient = purviewClient; - this._cacheProvider = cacheProvider; this._logger = logger; for (int i = 0; i < purviewSettings.MaxConcurrentJobConsumers; i++) @@ -74,28 +67,6 @@ private async Task RunJobAsync(BackgroundJobBase job) break; case ContentActivityJob contentActivityJob: _ = await this._purviewClient.SendContentActivitiesAsync(contentActivityJob.Request, CancellationToken.None).ConfigureAwait(false); - break; - case ScopeRetrievalJob scopeRetrievalJob: - try - { - ProtectionScopesResponse response = await this._purviewClient.GetProtectionScopesAsync(scopeRetrievalJob.Request, CancellationToken.None).ConfigureAwait(false); - await this._cacheProvider.SetAsync(scopeRetrievalJob.CacheKey, response, CancellationToken.None).ConfigureAwait(false); - (bool shouldProcess, List _, ExecutionMode _) = ScopedContentProcessor.CheckApplicableScopes(scopeRetrievalJob.ProcessContentRequest, response); - if (!shouldProcess) - { - ProcessContentRequest pcRequest = scopeRetrievalJob.ProcessContentRequest; - ContentActivitiesRequest caRequest = new(pcRequest.UserId, pcRequest.TenantId, pcRequest.ContentToProcess, pcRequest.CorrelationId); - this._channelHandler.QueueJob(new ContentActivityJob(caRequest)); - } - } - catch (PurviewPaymentRequiredException ex) - { - await this._cacheProvider.SetAsync( - new PaymentRequiredCacheKey(scopeRetrievalJob.Request.TenantId), - new PaymentRequiredCacheEntry(ex.Message), - CancellationToken.None).ConfigureAwait(false); - } - break; } } diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs deleted file mode 100644 index c23553f1855..00000000000 --- a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using Microsoft.Agents.AI.Purview.Models.Common; -using Microsoft.Agents.AI.Purview.Models.Requests; - -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 -{ - /// - /// Initializes a new instance of the class. - /// - /// The protection scopes request to send to Purview. - /// The cache key used to store the response. - /// The original process content request that triggered scope retrieval. - public ScopeRetrievalJob(ProtectionScopesRequest request, ProtectionScopesCacheKey cacheKey, ProcessContentRequest processContentRequest) - { - this.Request = request; - this.CacheKey = cacheKey; - this.ProcessContentRequest = processContentRequest; - } - - /// - /// Gets the protection scopes request. - /// - public ProtectionScopesRequest Request { get; } - - /// - /// Gets the cache key used to store the response. - /// - public ProtectionScopesCacheKey CacheKey { get; } - - /// - /// Gets the original process content request that triggered scope retrieval. - /// - public ProcessContentRequest ProcessContentRequest { get; } -} diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs index fab7c28d9ae..c9806a26c85 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs @@ -1,14 +1,18 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; +using System.Text.Json; +using System.Text.Json.Serialization.Metadata; using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Purview.Models.Common; using Microsoft.Agents.AI.Purview.Models.Jobs; using Microsoft.Agents.AI.Purview.Models.Requests; using Microsoft.Agents.AI.Purview.Models.Responses; +using Microsoft.Agents.AI.Purview.Serialization; using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Purview; @@ -21,6 +25,7 @@ internal sealed class ScopedContentProcessor : IScopedContentProcessor private readonly IPurviewClient _purviewClient; private readonly ICacheProvider _cacheProvider; private readonly IChannelHandler _channelHandler; + private readonly ConcurrentDictionary> _scopeRetrievals = new(); /// /// Create a new instance of . @@ -195,21 +200,89 @@ 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); + cacheResponse = await this.GetAndCacheProtectionScopesAsync(psRequest, cacheKey, cancellationToken).ConfigureAwait(false); + } + + return await this.ProcessWithCachedScopesAsync(pcRequest, cacheResponse, cacheKey, cancellationToken).ConfigureAwait(false); + } + + /// + /// Retrieve and cache protection scopes while sharing concurrent retrievals for the same cache key. + /// + private async Task GetAndCacheProtectionScopesAsync( + ProtectionScopesRequest psRequest, + ProtectionScopesCacheKey cacheKey, + CancellationToken cancellationToken) + { + JsonTypeInfo keyTypeInfo = + (JsonTypeInfo)PurviewSerializationUtils.SerializationSettings.GetTypeInfo(typeof(ProtectionScopesCacheKey)); + string serializedCacheKey = JsonSerializer.Serialize(cacheKey, keyTypeInfo); + TaskCompletionSource candidate = new(TaskCreationOptions.RunContinuationsAsynchronously); + TaskCompletionSource retrieval = this._scopeRetrievals.GetOrAdd(serializedCacheKey, candidate); + if (ReferenceEquals(candidate, retrieval)) + { + _ = this.PopulateProtectionScopesAsync(candidate, psRequest, cacheKey, cancellationToken); } try { - this._channelHandler.QueueJob(new ScopeRetrievalJob(psRequest, cacheKey, pcRequest)); + return await retrieval.Task.ConfigureAwait(false); } - catch (PurviewJobException) + finally { - // QueueJob already logs failures. Scope warmup is best effort; don't block ProcessContent. + _ = ((ICollection>>)this._scopeRetrievals) + .Remove(new KeyValuePair>(serializedCacheKey, retrieval)); } + } - return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions: null, cancellationToken).ConfigureAwait(false); + /// + /// Complete a shared protection scopes retrieval without allowing exceptions to escape a fire-and-forget task. + /// + private async Task PopulateProtectionScopesAsync( + TaskCompletionSource completionSource, + ProtectionScopesRequest psRequest, + ProtectionScopesCacheKey cacheKey, + CancellationToken cancellationToken) + { + try + { + ProtectionScopesResponse response = await this.RetrieveAndCacheProtectionScopesAsync(psRequest, cacheKey, cancellationToken).ConfigureAwait(false); + completionSource.TrySetResult(response); + } + catch (OperationCanceledException ex) + { + completionSource.TrySetCanceled(ex.CancellationToken); + } + catch (Exception ex) + { + completionSource.TrySetException(ex); + } + } + + /// + /// Retrieve protection scopes and cache either the response or a payment-required result. + /// + private async Task RetrieveAndCacheProtectionScopesAsync( + ProtectionScopesRequest psRequest, + ProtectionScopesCacheKey cacheKey, + CancellationToken cancellationToken) + { + try + { + ProtectionScopesResponse response = await this._purviewClient.GetProtectionScopesAsync(psRequest, cancellationToken).ConfigureAwait(false); + await this._cacheProvider.SetAsync(cacheKey, response, cancellationToken).ConfigureAwait(false); + return response; + } + catch (PurviewPaymentRequiredException ex) + { + await this._cacheProvider.SetAsync( + new PaymentRequiredCacheKey(psRequest.TenantId), + new PaymentRequiredCacheEntry(ex.Message), + cancellationToken).ConfigureAwait(false); + throw; + } } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs index d1b9d535589..8f80b98c95d 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs @@ -3,14 +3,12 @@ using System; using System.Collections.Generic; using System.Threading; -using System.Threading.Channels; using System.Threading.Tasks; using Microsoft.Agents.AI.Purview.Models.Common; using Microsoft.Agents.AI.Purview.Models.Jobs; using Microsoft.Agents.AI.Purview.Models.Requests; using Microsoft.Agents.AI.Purview.Models.Responses; using Microsoft.Extensions.AI; -using Microsoft.Extensions.Logging.Abstractions; using Moq; namespace Microsoft.Agents.AI.Purview.UnitTests; @@ -627,7 +625,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 +666,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 +701,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 +751,7 @@ public async Task ProcessMessagesAsync_ExtractsUserIdFromMessageAuthorName_WhenV } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_QueuesScopeRetrievalJobAndCallsProcessContentAsync() + public async Task ProcessMessagesAsync_CacheMiss_RetrievesScopesBeforeInlineProcessingAsync() { // Arrange var messages = new List @@ -766,24 +767,43 @@ public async Task ProcessMessagesAsync_CacheMiss_QueuesScopeRetrievalJobAndCalls It.IsAny(), It.IsAny())) .ReturnsAsync((ProtectionScopesResponse?)null); + ProtectionScopesResponse scopes = CreateApplicableProtectionScopesResponse(); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(scopes); + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(new ProcessContentResponse()); + It.Is(request => + request.ScopeIdentifier == scopes.ScopeIdentifier && + request.ProcessInline), + It.IsAny())) + .ReturnsAsync(new ProcessContentResponse + { + PolicyActions = + [ + new() { Action = DlpAction.BlockAccess } + ] + }); // Act - await this._processor.ProcessMessagesAsync( + var result = 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. - this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Once); + // Assert + Assert.True(result.shouldBlock); this._mockPurviewClient.Verify(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny()), Times.Never); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + It.IsAny(), It.IsAny()), Times.Once); + this._mockCacheProvider.Verify(x => x.SetAsync( + It.IsAny(), scopes, It.IsAny()), Times.Once); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.Is(request => + request.ScopeIdentifier == scopes.ScopeIdentifier && + request.ProcessInline), + It.IsAny()), Times.Once); } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_WithProcessContentBlockAction_ReturnsShouldBlockTrueAsync() + public async Task ProcessMessagesAsync_CacheMiss_WithOfflineScope_QueuesScopedProcessContentAsync() { // Arrange var messages = new List @@ -799,29 +819,26 @@ 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); + ProtectionScopesResponse scopes = CreateApplicableProtectionScopesResponse(ExecutionMode.EvaluateOffline); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(scopes); // Act - var result = await this._processor.ProcessMessagesAsync( + await this._processor.ProcessMessagesAsync( messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); // Assert - Assert.True(result.shouldBlock); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Never); + this._mockChannelHandler.Verify(x => x.QueueJob( + It.Is(job => + job.Request.ScopeIdentifier == scopes.ScopeIdentifier && + !job.Request.ProcessInline)), Times.Once); } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_StillCallsProcessContentWhenScopeJobCannotQueueAsync() + public async Task ProcessMessagesAsync_CacheMiss_WithNoApplicableScope_QueuesContentActivityAsync() { // Arrange var messages = new List @@ -837,134 +854,241 @@ public async Task ProcessMessagesAsync_CacheMiss_StillCallsProcessContentWhenSco It.IsAny(), It.IsAny())) .ReturnsAsync((ProtectionScopesResponse?)null); - this._mockChannelHandler.Setup(x => x.QueueJob(It.IsAny())) - .Throws(new PurviewJobException("queue unavailable")); - - this._mockPurviewClient.Setup(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(new ProcessContentResponse()); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(new ProtectionScopesResponse { ScopeIdentifier = "scope-123", Scopes = [] }); // 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. - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + // Assert this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Once); + It.IsAny(), It.IsAny()), Times.Never); + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); } [Fact] - public async Task ProcessMessagesAsync_WithCachedPaymentRequiredState_ThrowsPaymentRequiredAsync() + public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalFails_DoesNotProcessContentAsync() { // Arrange var messages = new List { - new (ChatRole.User, "Test message") + new(ChatRole.User, "Test message") }; var settings = CreateValidPurviewSettings(); var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync((ProtectionScopesResponse?)null); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .ThrowsAsync(new PurviewRequestException("scope retrieval failed")); - this._mockCacheProvider.Setup(x => x.GetAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(new PaymentRequiredCacheEntry("Payment required")); - - // Act + Assert - await Assert.ThrowsAsync(() => + // Act & Assert + await Assert.ThrowsAsync(() => this._processor.ProcessMessagesAsync( messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); - 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 ProcessMessagesAsync_CacheMiss_WhenScopeCachingFails_DoesNotProcessContentAsync() { // Arrange - Func, Task>? runner = null; - Mock channelHandler = new(); - 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") - ] + var messages = new List + { + new(ChatRole.User, "Test message") }; - ProtectionScopesCacheKey cacheKey = new(request); - Channel channel = Channel.CreateUnbounded(); + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + ProtectionScopesResponse scopes = CreateApplicableProtectionScopesResponse(); + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync((ProtectionScopesResponse?)null); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(scopes); + this._mockCacheProvider.Setup(x => x.SetAsync( + It.IsAny(), scopes, It.IsAny())) + .ThrowsAsync(new PurviewRequestException("scope caching failed")); - channelHandler.Setup(x => x.AddRunner(It.IsAny, Task>>())) - .Callback, Task>>(callback => runner = callback); + // Act & Assert + await Assert.ThrowsAsync(() => + this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Never); + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); + } - purviewClient.Setup(x => x.GetProtectionScopesAsync(It.IsAny(), It.IsAny())) - .ThrowsAsync(new PurviewPaymentRequiredException("Payment required")); + [Fact] + public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalIsCanceled_DoesNotProcessContentAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync((ProtectionScopesResponse?)null); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .ThrowsAsync(new OperationCanceledException()); - _ = new BackgroundJobRunner(channelHandler.Object, purviewClient.Object, cacheProvider.Object, NullLogger.Instance, settings); + // Act & Assert + await Assert.ThrowsAnyAsync(() => + this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Never); + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); + } - // Act - Assert.NotNull(runner); - await channel.Writer.WriteAsync(new ScopeRetrievalJob(request, cacheKey, CreateProcessContentRequest())); - channel.Writer.Complete(); - await runner(channel); + [Fact] + public async Task ProcessMessagesAsync_CacheMiss_WhenPaymentIsRequired_CachesAndRethrowsAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync((ProtectionScopesResponse?)null); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .ThrowsAsync(new PurviewPaymentRequiredException("Payment required")); - // Assert - cacheProvider.Verify(x => x.SetAsync( + // Act & Assert + await Assert.ThrowsAsync(() => + this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); + this._mockCacheProvider.Verify(x => x.SetAsync( It.Is(key => key.TenantId == "tenant-123"), It.Is(entry => entry.Message == "Payment required"), It.IsAny()), Times.Once); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Never); } [Fact] - public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_QueuesContentActivityJobAsync() + public async Task ProcessMessagesAsync_ConcurrentCacheMisses_ShareScopeRetrievalAsync() { // Arrange - Func, Task>? runner = null; - Mock channelHandler = new(); - Mock purviewClient = new(); - Mock cacheProvider = new(); - PurviewSettings settings = new("TestApp") { MaxConcurrentJobConsumers = 1 }; - ProtectionScopesRequest request = CreateProtectionScopesRequest(); - ScopeRetrievalJob job = new(request, new ProtectionScopesCacheKey(request), CreateProcessContentRequest()); - Channel channel = Channel.CreateUnbounded(); + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + TaskCompletionSource scopesSource = new(TaskCreationOptions.RunContinuationsAsynchronously); + TaskCompletionSource retrievalStarted = new(TaskCreationOptions.RunContinuationsAsynchronously); + TaskCompletionSource secondCacheRead = new(TaskCreationOptions.RunContinuationsAsynchronously); + int cacheReadCount = 0; - channelHandler.Setup(x => x.AddRunner(It.IsAny, Task>>())) - .Callback, Task>>(callback => runner = callback); + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(() => + { + if (Interlocked.Increment(ref cacheReadCount) == 2) + { + secondCacheRead.TrySetResult(true); + } - purviewClient.Setup(x => x.GetProtectionScopesAsync(It.IsAny(), It.IsAny())) - .ReturnsAsync(new ProtectionScopesResponse { Scopes = [] }); + return null; + }); + this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny())) + .Returns(() => + { + retrievalStarted.TrySetResult(true); + return scopesSource.Task; + }); + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(new ProcessContentResponse()); - _ = new BackgroundJobRunner(channelHandler.Object, purviewClient.Object, cacheProvider.Object, NullLogger.Instance, settings); + Task<(bool shouldBlock, string? userId)> first = this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + await retrievalStarted.Task; + Task<(bool shouldBlock, string? userId)> second = this._processor.ProcessMessagesAsync( + messages, "session-456", Activity.UploadText, settings, "user-123", CancellationToken.None); + await secondCacheRead.Task; // Act - Assert.NotNull(runner); - await channel.Writer.WriteAsync(job); - channel.Writer.Complete(); - await runner(channel); + scopesSource.SetResult(CreateApplicableProtectionScopesResponse()); + await Task.WhenAll(first, second); // Assert - channelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + this._mockPurviewClient.Verify(x => x.GetProtectionScopesAsync( + It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + public async Task ProcessMessagesAsync_WithCachedPaymentRequiredState_ThrowsPaymentRequiredAsync() + { + // Arrange + var messages = new List + { + new (ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(new PaymentRequiredCacheEntry("Payment required")); + + // Act + Assert + await Assert.ThrowsAsync(() => + this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); + + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Never); + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); } #endregion #region Helper Methods - private static ProtectionScopesRequest CreateProtectionScopesRequest() + private static ProtectionScopesResponse CreateApplicableProtectionScopesResponse(ExecutionMode executionMode = ExecutionMode.EvaluateInline) { - return new ProtectionScopesRequest("user-123", "tenant-123") + return new ProtectionScopesResponse { - Activities = ProtectionScopeActivities.UploadText, - Locations = + ScopeIdentifier = "scope-123", + Scopes = [ - new("microsoft.graph.policyLocationApplication", "app-123") + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("microsoft.graph.policyLocationApplication", "app-123") + ], + ExecutionMode = executionMode + } ] }; } From 6d0d3a93456d6de70ebf143430aff5c70e3d10cf Mon Sep 17 00:00:00 2001 From: Eoin Doherty Date: Mon, 31 Aug 2026 10:15:00 -0700 Subject: [PATCH 2/4] Optimize Purview cold scope processing Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../BackgroundJobRunner.cs | 31 +- .../Models/Jobs/ScopeRetrievalJob.cs | 40 +++ .../ScopedContentProcessor.cs | 86 +---- .../ScopedContentProcessorTests.cs | 324 +++++++----------- 4 files changed, 199 insertions(+), 282 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs b/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs index 85a4fa54c3c..03f73d8007a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs @@ -1,10 +1,14 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Collections.Generic; using System.Threading; using System.Threading.Channels; using System.Threading.Tasks; +using Microsoft.Agents.AI.Purview.Models.Common; using Microsoft.Agents.AI.Purview.Models.Jobs; +using Microsoft.Agents.AI.Purview.Models.Requests; +using Microsoft.Agents.AI.Purview.Models.Responses; using Microsoft.Extensions.Logging; namespace Microsoft.Agents.AI.Purview; @@ -16,6 +20,7 @@ internal sealed class BackgroundJobRunner : IBackgroundJobRunner { private readonly IChannelHandler _channelHandler; private readonly IPurviewClient _purviewClient; + private readonly ICacheProvider _cacheProvider; private readonly ILogger _logger; /// @@ -23,12 +28,14 @@ internal sealed class BackgroundJobRunner : IBackgroundJobRunner /// /// The channel handler used to manage job channels. /// The Purview client used to send requests to Purview. + /// The cache provider used to store protection scopes results. /// The logger used to log information about background jobs. /// The settings used to configure Purview client behavior. - public BackgroundJobRunner(IChannelHandler channelHandler, IPurviewClient purviewClient, ILogger logger, PurviewSettings purviewSettings) + public BackgroundJobRunner(IChannelHandler channelHandler, IPurviewClient purviewClient, ICacheProvider cacheProvider, ILogger logger, PurviewSettings purviewSettings) { this._channelHandler = channelHandler; this._purviewClient = purviewClient; + this._cacheProvider = cacheProvider; this._logger = logger; for (int i = 0; i < purviewSettings.MaxConcurrentJobConsumers; i++) @@ -67,6 +74,28 @@ private async Task RunJobAsync(BackgroundJobBase job) break; case ContentActivityJob contentActivityJob: _ = await this._purviewClient.SendContentActivitiesAsync(contentActivityJob.Request, CancellationToken.None).ConfigureAwait(false); + break; + case ScopeRetrievalJob scopeRetrievalJob: + try + { + ProtectionScopesResponse response = await this._purviewClient.GetProtectionScopesAsync(scopeRetrievalJob.Request, CancellationToken.None).ConfigureAwait(false); + await this._cacheProvider.SetAsync(scopeRetrievalJob.CacheKey, response, CancellationToken.None).ConfigureAwait(false); + (bool shouldProcess, List _, ExecutionMode _) = ScopedContentProcessor.CheckApplicableScopes(scopeRetrievalJob.ProcessContentRequest, response); + if (!shouldProcess) + { + ProcessContentRequest pcRequest = scopeRetrievalJob.ProcessContentRequest; + ContentActivitiesRequest caRequest = new(pcRequest.UserId, pcRequest.TenantId, pcRequest.ContentToProcess, pcRequest.CorrelationId); + this._channelHandler.QueueJob(new ContentActivityJob(caRequest)); + } + } + catch (PurviewPaymentRequiredException ex) + { + await this._cacheProvider.SetAsync( + new PaymentRequiredCacheKey(scopeRetrievalJob.Request.TenantId), + new PaymentRequiredCacheEntry(ex.Message), + CancellationToken.None).ConfigureAwait(false); + } + break; } } diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs new file mode 100644 index 00000000000..da7aac682d6 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs @@ -0,0 +1,40 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Agents.AI.Purview.Models.Common; +using Microsoft.Agents.AI.Purview.Models.Requests; + +namespace Microsoft.Agents.AI.Purview.Models.Jobs; + +/// +/// Class representing a job that refreshes the protection scopes cache in the background. +/// +internal sealed class ScopeRetrievalJob : BackgroundJobBase +{ + /// + /// Initializes a new instance of the class. + /// + /// The protection scopes request to send to Purview. + /// The cache key used to store the response. + /// The original process content request that triggered scope retrieval. + public ScopeRetrievalJob(ProtectionScopesRequest request, ProtectionScopesCacheKey cacheKey, ProcessContentRequest processContentRequest) + { + this.Request = request; + this.CacheKey = cacheKey; + this.ProcessContentRequest = processContentRequest; + } + + /// + /// Gets the protection scopes request. + /// + public ProtectionScopesRequest Request { get; } + + /// + /// Gets the cache key used to store the response. + /// + public ProtectionScopesCacheKey CacheKey { get; } + + /// + /// Gets the original process content request that triggered scope retrieval. + /// + public ProcessContentRequest ProcessContentRequest { get; } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs index c9806a26c85..037751fa46b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs @@ -1,18 +1,14 @@ // Copyright (c) Microsoft. All rights reserved. using System; -using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; -using System.Text.Json; -using System.Text.Json.Serialization.Metadata; using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Purview.Models.Common; using Microsoft.Agents.AI.Purview.Models.Jobs; using Microsoft.Agents.AI.Purview.Models.Requests; using Microsoft.Agents.AI.Purview.Models.Responses; -using Microsoft.Agents.AI.Purview.Serialization; using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Purview; @@ -25,7 +21,6 @@ internal sealed class ScopedContentProcessor : IScopedContentProcessor private readonly IPurviewClient _purviewClient; private readonly ICacheProvider _cacheProvider; private readonly IChannelHandler _channelHandler; - private readonly ConcurrentDictionary> _scopeRetrievals = new(); /// /// Create a new instance of . @@ -202,89 +197,14 @@ private async Task ProcessContentWithProtectionScopesAsy if (cacheResponse == null) { - cacheResponse = await this.GetAndCacheProtectionScopesAsync(psRequest, cacheKey, cancellationToken).ConfigureAwait(false); + pcRequest.ProcessInline = true; + this._channelHandler.QueueJob(new ScopeRetrievalJob(psRequest, cacheKey, pcRequest)); + return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions: null, cancellationToken).ConfigureAwait(false); } return await this.ProcessWithCachedScopesAsync(pcRequest, cacheResponse, cacheKey, cancellationToken).ConfigureAwait(false); } - /// - /// Retrieve and cache protection scopes while sharing concurrent retrievals for the same cache key. - /// - private async Task GetAndCacheProtectionScopesAsync( - ProtectionScopesRequest psRequest, - ProtectionScopesCacheKey cacheKey, - CancellationToken cancellationToken) - { - JsonTypeInfo keyTypeInfo = - (JsonTypeInfo)PurviewSerializationUtils.SerializationSettings.GetTypeInfo(typeof(ProtectionScopesCacheKey)); - string serializedCacheKey = JsonSerializer.Serialize(cacheKey, keyTypeInfo); - TaskCompletionSource candidate = new(TaskCreationOptions.RunContinuationsAsynchronously); - TaskCompletionSource retrieval = this._scopeRetrievals.GetOrAdd(serializedCacheKey, candidate); - if (ReferenceEquals(candidate, retrieval)) - { - _ = this.PopulateProtectionScopesAsync(candidate, psRequest, cacheKey, cancellationToken); - } - - try - { - return await retrieval.Task.ConfigureAwait(false); - } - finally - { - _ = ((ICollection>>)this._scopeRetrievals) - .Remove(new KeyValuePair>(serializedCacheKey, retrieval)); - } - } - - /// - /// Complete a shared protection scopes retrieval without allowing exceptions to escape a fire-and-forget task. - /// - private async Task PopulateProtectionScopesAsync( - TaskCompletionSource completionSource, - ProtectionScopesRequest psRequest, - ProtectionScopesCacheKey cacheKey, - CancellationToken cancellationToken) - { - try - { - ProtectionScopesResponse response = await this.RetrieveAndCacheProtectionScopesAsync(psRequest, cacheKey, cancellationToken).ConfigureAwait(false); - completionSource.TrySetResult(response); - } - catch (OperationCanceledException ex) - { - completionSource.TrySetCanceled(ex.CancellationToken); - } - catch (Exception ex) - { - completionSource.TrySetException(ex); - } - } - - /// - /// Retrieve protection scopes and cache either the response or a payment-required result. - /// - private async Task RetrieveAndCacheProtectionScopesAsync( - ProtectionScopesRequest psRequest, - ProtectionScopesCacheKey cacheKey, - CancellationToken cancellationToken) - { - try - { - ProtectionScopesResponse response = await this._purviewClient.GetProtectionScopesAsync(psRequest, cancellationToken).ConfigureAwait(false); - await this._cacheProvider.SetAsync(cacheKey, response, cancellationToken).ConfigureAwait(false); - return response; - } - catch (PurviewPaymentRequiredException ex) - { - await this._cacheProvider.SetAsync( - new PaymentRequiredCacheKey(psRequest.TenantId), - new PaymentRequiredCacheEntry(ex.Message), - cancellationToken).ConfigureAwait(false); - throw; - } - } - /// /// Apply locally-cached protection scopes to the request and dispatch ProcessContent appropriately. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs index 8f80b98c95d..f560b373460 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs @@ -3,12 +3,14 @@ using System; using System.Collections.Generic; using System.Threading; +using System.Threading.Channels; using System.Threading.Tasks; using Microsoft.Agents.AI.Purview.Models.Common; using Microsoft.Agents.AI.Purview.Models.Jobs; using Microsoft.Agents.AI.Purview.Models.Requests; using Microsoft.Agents.AI.Purview.Models.Responses; using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging.Abstractions; using Moq; namespace Microsoft.Agents.AI.Purview.UnitTests; @@ -751,7 +753,7 @@ public async Task ProcessMessagesAsync_ExtractsUserIdFromMessageAuthorName_WhenV } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_RetrievesScopesBeforeInlineProcessingAsync() + public async Task ProcessMessagesAsync_CacheMiss_CallsProcessContentInlineAndQueuesScopeRetrievalAsync() { // Arrange var messages = new List @@ -767,43 +769,31 @@ public async Task ProcessMessagesAsync_CacheMiss_RetrievesScopesBeforeInlineProc It.IsAny(), It.IsAny())) .ReturnsAsync((ProtectionScopesResponse?)null); - ProtectionScopesResponse scopes = CreateApplicableProtectionScopesResponse(); - this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(scopes); - this._mockPurviewClient.Setup(x => x.ProcessContentAsync( It.Is(request => - request.ScopeIdentifier == scopes.ScopeIdentifier && + request.ScopeIdentifier == null && request.ProcessInline), It.IsAny())) - .ReturnsAsync(new ProcessContentResponse - { - PolicyActions = - [ - new() { Action = DlpAction.BlockAccess } - ] - }); + .ReturnsAsync(new ProcessContentResponse()); // Act - var result = await this._processor.ProcessMessagesAsync( + await this._processor.ProcessMessagesAsync( messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); // Assert - Assert.True(result.shouldBlock); - this._mockPurviewClient.Verify(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny()), Times.Once); - this._mockCacheProvider.Verify(x => x.SetAsync( - It.IsAny(), scopes, It.IsAny()), Times.Once); this._mockPurviewClient.Verify(x => x.ProcessContentAsync( It.Is(request => - request.ScopeIdentifier == scopes.ScopeIdentifier && + 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.Is(job => job.ProcessContentRequest.ProcessInline)), Times.Once); } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_WithOfflineScope_QueuesScopedProcessContentAsync() + public async Task ProcessMessagesAsync_CacheMiss_WithProcessContentBlockAction_ReturnsShouldBlockTrueAsync() { // Arrange var messages = new List @@ -819,26 +809,30 @@ public async Task ProcessMessagesAsync_CacheMiss_WithOfflineScope_QueuesScopedPr It.IsAny(), It.IsAny())) .ReturnsAsync((ProtectionScopesResponse?)null); - ProtectionScopesResponse scopes = CreateApplicableProtectionScopesResponse(ExecutionMode.EvaluateOffline); - this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(scopes); + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + It.Is(request => request.ProcessInline), + It.IsAny())) + .ReturnsAsync(new ProcessContentResponse + { + PolicyActions = + [ + new() { Action = DlpAction.BlockAccess } + ] + }); // Act - await this._processor.ProcessMessagesAsync( + var result = await this._processor.ProcessMessagesAsync( messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); // Assert + Assert.True(result.shouldBlock); this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Never); - this._mockChannelHandler.Verify(x => x.QueueJob( - It.Is(job => - job.Request.ScopeIdentifier == scopes.ScopeIdentifier && - !job.Request.ProcessInline)), Times.Once); + It.Is(request => request.ProcessInline), + It.IsAny()), Times.Once); } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_WithNoApplicableScope_QueuesContentActivityAsync() + public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalCannotQueue_DoesNotCallProcessContentAsync() { // Arrange var messages = new List @@ -854,225 +848,159 @@ public async Task ProcessMessagesAsync_CacheMiss_WithNoApplicableScope_QueuesCon It.IsAny(), It.IsAny())) .ReturnsAsync((ProtectionScopesResponse?)null); - this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(new ProtectionScopesResponse { ScopeIdentifier = "scope-123", Scopes = [] }); - - // Act - await this._processor.ProcessMessagesAsync( - messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); - - // Assert - this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Never); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); - } - - [Fact] - public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalFails_DoesNotProcessContentAsync() - { - // Arrange - var messages = new List - { - new(ChatRole.User, "Test message") - }; - var settings = CreateValidPurviewSettings(); - var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; - this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) - .ReturnsAsync(tokenInfo); - this._mockCacheProvider.Setup(x => x.GetAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync((ProtectionScopesResponse?)null); - this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny())) - .ThrowsAsync(new PurviewRequestException("scope retrieval failed")); + this._mockChannelHandler.Setup(x => x.QueueJob(It.IsAny())) + .Throws(new PurviewJobException("queue unavailable")); // Act & Assert - await Assert.ThrowsAsync(() => + await Assert.ThrowsAsync(() => this._processor.ProcessMessagesAsync( messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); this._mockPurviewClient.Verify(x => x.ProcessContentAsync( It.IsAny(), It.IsAny()), Times.Never); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_WhenScopeCachingFails_DoesNotProcessContentAsync() + public async Task ProcessMessagesAsync_WithCachedPaymentRequiredState_ThrowsPaymentRequiredAsync() { // Arrange var messages = new List { - new(ChatRole.User, "Test message") + new (ChatRole.User, "Test message") }; var settings = CreateValidPurviewSettings(); var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; - ProtectionScopesResponse scopes = CreateApplicableProtectionScopesResponse(); this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) .ReturnsAsync(tokenInfo); - this._mockCacheProvider.Setup(x => x.GetAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync((ProtectionScopesResponse?)null); - this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(scopes); - this._mockCacheProvider.Setup(x => x.SetAsync( - It.IsAny(), scopes, It.IsAny())) - .ThrowsAsync(new PurviewRequestException("scope caching failed")); - - // Act & Assert - await Assert.ThrowsAsync(() => - this._processor.ProcessMessagesAsync( - messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); - this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Never); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); - } - [Fact] - public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalIsCanceled_DoesNotProcessContentAsync() - { - // Arrange - var messages = new List - { - new(ChatRole.User, "Test message") - }; - var settings = CreateValidPurviewSettings(); - var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; - this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) - .ReturnsAsync(tokenInfo); - this._mockCacheProvider.Setup(x => x.GetAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync((ProtectionScopesResponse?)null); - this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny())) - .ThrowsAsync(new OperationCanceledException()); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(new PaymentRequiredCacheEntry("Payment required")); - // Act & Assert - await Assert.ThrowsAnyAsync(() => + // Act + Assert + await Assert.ThrowsAsync(() => this._processor.ProcessMessagesAsync( messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( It.IsAny(), It.IsAny()), Times.Never); this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_WhenPaymentIsRequired_CachesAndRethrowsAsync() + public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_CachesAndQueuesContentActivityAsync() { // Arrange - var messages = new List - { - new(ChatRole.User, "Test message") - }; - var settings = CreateValidPurviewSettings(); - var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; - this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) - .ReturnsAsync(tokenInfo); - this._mockCacheProvider.Setup(x => x.GetAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync((ProtectionScopesResponse?)null); - this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny())) - .ThrowsAsync(new PurviewPaymentRequiredException("Payment required")); + 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 = 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 - await Assert.ThrowsAsync(() => - this._processor.ProcessMessagesAsync( - messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); - this._mockCacheProvider.Verify(x => x.SetAsync( - It.Is(key => key.TenantId == "tenant-123"), - It.Is(entry => entry.Message == "Payment required"), - It.IsAny()), Times.Once); - this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Never); + // 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 ProcessMessagesAsync_ConcurrentCacheMisses_ShareScopeRetrievalAsync() + public async Task BackgroundJobRunner_ScopeRetrievalApplicableScope_CachesWithoutContentActivityAsync() { // Arrange - var messages = new List - { - new(ChatRole.User, "Test message") - }; - var settings = CreateValidPurviewSettings(); - var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; - TaskCompletionSource scopesSource = new(TaskCreationOptions.RunContinuationsAsynchronously); - TaskCompletionSource retrievalStarted = new(TaskCreationOptions.RunContinuationsAsynchronously); - TaskCompletionSource secondCacheRead = new(TaskCreationOptions.RunContinuationsAsynchronously); - int cacheReadCount = 0; - - this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) - .ReturnsAsync(tokenInfo); - this._mockCacheProvider.Setup(x => x.GetAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(() => - { - if (Interlocked.Increment(ref cacheReadCount) == 2) - { - secondCacheRead.TrySetResult(true); - } - - return null; - }); - this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync( - It.IsAny(), It.IsAny())) - .Returns(() => - { - retrievalStarted.TrySetResult(true); - return scopesSource.Task; - }); - this._mockPurviewClient.Setup(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(new ProcessContentResponse()); - - Task<(bool shouldBlock, string? userId)> first = this._processor.ProcessMessagesAsync( - messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); - await retrievalStarted.Task; - Task<(bool shouldBlock, string? userId)> second = this._processor.ProcessMessagesAsync( - messages, "session-456", Activity.UploadText, settings, "user-123", CancellationToken.None); - await secondCacheRead.Task; + 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())) + .ReturnsAsync(response); + + _ = new BackgroundJobRunner(channelHandler.Object, purviewClient.Object, cacheProvider.Object, NullLogger.Instance, settings); // Act - scopesSource.SetResult(CreateApplicableProtectionScopesResponse()); - await Task.WhenAll(first, second); + Assert.NotNull(runner); + await channel.Writer.WriteAsync(job); + channel.Writer.Complete(); + await runner(channel); // Assert - this._mockPurviewClient.Verify(x => x.GetProtectionScopesAsync( - It.IsAny(), 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 ProcessMessagesAsync_WithCachedPaymentRequiredState_ThrowsPaymentRequiredAsync() + public async Task BackgroundJobRunner_ScopeRetrievalPaymentRequired_CachesForSubsequentCallsAsync() { // Arrange - var messages = new List - { - new (ChatRole.User, "Test message") - }; - var settings = CreateValidPurviewSettings(); - var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; - this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) - .ReturnsAsync(tokenInfo); + Func, Task>? runner = null; + Mock channelHandler = new(); + Mock purviewClient = new(); + Mock cacheProvider = new(); + PurviewSettings settings = new("TestApp") { MaxConcurrentJobConsumers = 1 }; + ProtectionScopesRequest request = CreateProtectionScopesRequest(); + ScopeRetrievalJob job = new(request, new ProtectionScopesCacheKey(request), CreateProcessContentRequest()); + 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")); - this._mockCacheProvider.Setup(x => x.GetAsync( - It.IsAny(), It.IsAny())) - .ReturnsAsync(new PaymentRequiredCacheEntry("Payment required")); + _ = new BackgroundJobRunner(channelHandler.Object, purviewClient.Object, cacheProvider.Object, NullLogger.Instance, settings); - // Act + Assert - await Assert.ThrowsAsync(() => - this._processor.ProcessMessagesAsync( - messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); + // Act + Assert.NotNull(runner); + await channel.Writer.WriteAsync(job); + channel.Writer.Complete(); + await runner(channel); - this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Never); - this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Never); + // Assert + cacheProvider.Verify(x => x.SetAsync( + It.Is(key => key.TenantId == "tenant-123"), + It.Is(entry => entry.Message == "Payment required"), + It.IsAny()), Times.Once); } #endregion #region Helper Methods + private static ProtectionScopesRequest CreateProtectionScopesRequest() + { + return new ProtectionScopesRequest("user-123", "tenant-123") + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("microsoft.graph.policyLocationApplication", "app-123") + ] + }; + } + private static ProtectionScopesResponse CreateApplicableProtectionScopesResponse(ExecutionMode executionMode = ExecutionMode.EvaluateInline) { return new ProtectionScopesResponse From 36424bab16a35599460643b0993445fbe38a1e84 Mon Sep 17 00:00:00 2001 From: Eoin Doherty Date: Mon, 31 Aug 2026 15:16:24 -0700 Subject: [PATCH 3/4] Run Python Purview cold requests inline Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/purview/agent_framework_purview/_processor.py | 1 + python/packages/purview/tests/purview/test_processor.py | 1 + 2 files changed, 2 insertions(+) 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)) From 0b74e0d73133e03c4859dc77cbe0ee3f77624858 Mon Sep 17 00:00:00 2001 From: Eoin Doherty Date: Mon, 31 Aug 2026 18:01:54 -0700 Subject: [PATCH 4/4] Keep Purview scope refresh best effort Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../ScopedContentProcessor.cs | 10 +++++++++- .../ScopedContentProcessorTests.cs | 20 +++++++++++++------ 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs index 037751fa46b..f49f1239a89 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs @@ -198,7 +198,15 @@ private async Task ProcessContentWithProtectionScopesAsy if (cacheResponse == null) { pcRequest.ProcessInline = true; - this._channelHandler.QueueJob(new ScopeRetrievalJob(psRequest, cacheKey, pcRequest)); + try + { + this._channelHandler.QueueJob(new ScopeRetrievalJob(psRequest, cacheKey, pcRequest)); + } + catch (PurviewJobException) + { + // QueueJob logs admission failures. Scope refresh is best effort. + } + return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions: null, 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 f560b373460..b59cc7a3929 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs @@ -832,7 +832,7 @@ public async Task ProcessMessagesAsync_CacheMiss_WithProcessContentBlockAction_R } [Fact] - public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalCannotQueue_DoesNotCallProcessContentAsync() + public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalCannotQueue_CallsProcessContentInlineAsync() { // Arrange var messages = new List @@ -851,12 +851,20 @@ public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalCannotQueue_D this._mockChannelHandler.Setup(x => x.QueueJob(It.IsAny())) .Throws(new PurviewJobException("queue unavailable")); - // Act & Assert - await Assert.ThrowsAsync(() => - this._processor.ProcessMessagesAsync( - messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None)); + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + 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 + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); this._mockPurviewClient.Verify(x => x.ProcessContentAsync( - It.IsAny(), It.IsAny()), Times.Never); + It.Is(request => request.ProcessInline), + It.IsAny()), Times.Once); } [Fact]