Python: [.NET][Python][Purview] Set inline header for requests made with cold scope cache

dotnet/src/Microsoft.Agents.AI.Purview/Models/Jobs/ScopeRetrievalJob.cs

@@ -8,10 +8,6 @@ namespace Microsoft.Agents.AI.Purview.Models.Jobs;
 /// <summary>
 /// Class representing a job that refreshes the protection scopes cache in the background.
 /// </summary>
-/// <remarks>
-/// Used by the parallel protection scopes retrieval path to warm the cache without blocking the
-/// foreground ProcessContent call.
-/// </remarks>
 internal sealed class ScopeRetrievalJob : BackgroundJobBase
 {
     /// <summary>

dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs

@@ -195,21 +195,22 @@ private async Task<ProcessContentResponse> ProcessContentWithProtectionScopesAsy
 
         ProtectionScopesResponse? cacheResponse = await this._cacheProvider.GetAsync<ProtectionScopesCacheKey, ProtectionScopesResponse>(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);
     }
 
     /// <summary>

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<ProtectionScopesRequest>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(new ProtectionScopesResponse { Scopes = [] });
+            .ReturnsAsync(CreateApplicableProtectionScopesResponse());
         this._mockPurviewClient.Setup(x => x.ProcessContentAsync(
             It.IsAny<ProcessContentRequest>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new ProcessContentResponse { PolicyActions = [] });
@@ -668,7 +668,7 @@ public async Task ProcessMessagesAsync_UsesTokenUserIdBeforeAuthorName_Async()
 
         this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync(
             It.IsAny<ProtectionScopesRequest>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(new ProtectionScopesResponse { Scopes = [] });
+            .ReturnsAsync(CreateApplicableProtectionScopesResponse());
         this._mockPurviewClient.Setup(x => x.ProcessContentAsync(
             It.IsAny<ProcessContentRequest>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(new ProcessContentResponse { PolicyActions = [] });
@@ -703,6 +703,9 @@ public async Task ProcessMessagesAsync_UsesProvidedUserId_WhenTokenUserIdIsEmpty
             It.IsAny<ProtectionScopesCacheKey>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync((ProtectionScopesResponse?)null);
 
+        this._mockPurviewClient.Setup(x => x.GetProtectionScopesAsync(
+            It.IsAny<ProtectionScopesRequest>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(CreateApplicableProtectionScopesResponse());
         this._mockPurviewClient.Setup(x => x.ProcessContentAsync(
             It.IsAny<ProcessContentRequest>(), It.IsAny<CancellationToken>()))
             .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<ChatMessage>
@@ -767,19 +770,26 @@ public async Task ProcessMessagesAsync_CacheMiss_QueuesScopeRetrievalJobAndCalls
             .ReturnsAsync((ProtectionScopesResponse?)null);
 
         this._mockPurviewClient.Setup(x => x.ProcessContentAsync(
-            It.IsAny<ProcessContentRequest>(), It.IsAny<CancellationToken>()))
+            It.Is<ProcessContentRequest>(request =>
+                request.ScopeIdentifier == null &&
+                request.ProcessInline),
+            It.IsAny<CancellationToken>()))
             .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<ProcessContentRequest>(), It.IsAny<CancellationToken>()), Times.Once);
+            It.Is<ProcessContentRequest>(request =>
+                request.ScopeIdentifier == null &&
+                request.ProcessInline),
+            It.IsAny<CancellationToken>()), Times.Once);
         this._mockPurviewClient.Verify(x => x.GetProtectionScopesAsync(
             It.IsAny<ProtectionScopesRequest>(), It.IsAny<CancellationToken>()), Times.Never);
-        this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny<ScopeRetrievalJob>()), Times.Once);
+        this._mockChannelHandler.Verify(x => x.QueueJob(
+            It.Is<ScopeRetrievalJob>(job => job.ProcessContentRequest.ProcessInline)), Times.Once);
     }
 
     [Fact]
@@ -799,29 +809,30 @@ public async Task ProcessMessagesAsync_CacheMiss_WithProcessContentBlockAction_R
             It.IsAny<ProtectionScopesCacheKey>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync((ProtectionScopesResponse?)null);
 
-        var pcResponse = new ProcessContentResponse
-        {
-            PolicyActions =
-            [
-                new() { Action = DlpAction.BlockAccess }
-            ]
-        };
-
         this._mockPurviewClient.Setup(x => x.ProcessContentAsync(
-            It.IsAny<ProcessContentRequest>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(pcResponse);
+            It.Is<ProcessContentRequest>(request => request.ProcessInline),
+            It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new ProcessContentResponse
+            {
+                PolicyActions =
+                [
+                    new() { Action = DlpAction.BlockAccess }
+                ]
+            });
 
         // Act
         var result = 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<ScopeRetrievalJob>()), Times.Once);
+        this._mockPurviewClient.Verify(x => x.ProcessContentAsync(
+            It.Is<ProcessContentRequest>(request => request.ProcessInline),
+            It.IsAny<CancellationToken>()), Times.Once);
     }
 
     [Fact]
-    public async Task ProcessMessagesAsync_CacheMiss_StillCallsProcessContentWhenScopeJobCannotQueueAsync()
+    public async Task ProcessMessagesAsync_CacheMiss_WhenScopeRetrievalCannotQueue_CallsProcessContentInlineAsync()
     {
         // Arrange
         var messages = new List<ChatMessage>
@@ -841,17 +852,19 @@ public async Task ProcessMessagesAsync_CacheMiss_StillCallsProcessContentWhenSco
             .Throws(new PurviewJobException("queue unavailable"));
 
         this._mockPurviewClient.Setup(x => x.ProcessContentAsync(
-            It.IsAny<ProcessContentRequest>(), It.IsAny<CancellationToken>()))
+            It.Is<ProcessContentRequest>(request => request.ProcessInline),
+            It.IsAny<CancellationToken>()))
             .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<ScopeRetrievalJob>()), Times.Once);
         this._mockPurviewClient.Verify(x => x.ProcessContentAsync(
-            It.IsAny<ProcessContentRequest>(), It.IsAny<CancellationToken>()), Times.Once);
+            It.Is<ProcessContentRequest>(request => request.ProcessInline),
+            It.IsAny<CancellationToken>()), Times.Once);
     }
 
     [Fact]
@@ -878,52 +891,77 @@ await Assert.ThrowsAsync<PurviewPaymentRequiredException>(() =>
 
         this._mockPurviewClient.Verify(x => x.ProcessContentAsync(
             It.IsAny<ProcessContentRequest>(), It.IsAny<CancellationToken>()), Times.Never);
-        this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny<ScopeRetrievalJob>()), Times.Never);
+        this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny<BackgroundJobBase>()), Times.Never);
     }
 
     [Fact]
-    public async Task BackgroundJobRunner_ScopeRetrievalPaymentRequired_CachesForSubsequentCallsAsync()
+    public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_CachesAndQueuesContentActivityAsync()
     {
         // Arrange
         Func<Channel<BackgroundJobBase>, Task>? runner = null;
         Mock<IChannelHandler> channelHandler = new();
         Mock<IPurviewClient> purviewClient = new();
         Mock<ICacheProvider> 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<BackgroundJobBase> channel = Channel.CreateUnbounded<BackgroundJobBase>();
 
         channelHandler.Setup(x => x.AddRunner(It.IsAny<Func<Channel<BackgroundJobBase>, Task>>()))
             .Callback<Func<Channel<BackgroundJobBase>, Task>>(callback => runner = callback);
+        purviewClient.Setup(x => x.GetProtectionScopesAsync(It.IsAny<ProtectionScopesRequest>(), It.IsAny<CancellationToken>()))
+            .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<CancellationToken>()), Times.Once);
+        channelHandler.Verify(x => x.QueueJob(It.IsAny<ContentActivityJob>()), Times.Once);
+    }
+
+    [Fact]
+    public async Task BackgroundJobRunner_ScopeRetrievalApplicableScope_CachesWithoutContentActivityAsync()
+    {
+        // Arrange
+        Func<Channel<BackgroundJobBase>, Task>? runner = null;
+        Mock<IChannelHandler> channelHandler = new();
+        Mock<IPurviewClient> purviewClient = new();
+        Mock<ICacheProvider> 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<BackgroundJobBase> channel = Channel.CreateUnbounded<BackgroundJobBase>();
+
+        channelHandler.Setup(x => x.AddRunner(It.IsAny<Func<Channel<BackgroundJobBase>, Task>>()))
+            .Callback<Func<Channel<BackgroundJobBase>, Task>>(callback => runner = callback);
         purviewClient.Setup(x => x.GetProtectionScopesAsync(It.IsAny<ProtectionScopesRequest>(), It.IsAny<CancellationToken>()))
-            .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<PaymentRequiredCacheKey>(key => key.TenantId == "tenant-123"),
-            It.Is<PaymentRequiredCacheEntry>(entry => entry.Message == "Payment required"),
-            It.IsAny<CancellationToken>()), Times.Once);
+        cacheProvider.Verify(x => x.SetAsync(cacheKey, response, It.IsAny<CancellationToken>()), Times.Once);
+        channelHandler.Verify(x => x.QueueJob(It.IsAny<ContentActivityJob>()), Times.Never);
     }
 
     [Fact]
-    public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_QueuesContentActivityJobAsync()
+    public async Task BackgroundJobRunner_ScopeRetrievalPaymentRequired_CachesForSubsequentCallsAsync()
     {
         // Arrange
         Func<Channel<BackgroundJobBase>, Task>? runner = null;
@@ -937,9 +975,8 @@ public async Task BackgroundJobRunner_ScopeRetrievalNoApplicableScopes_QueuesCon
 
         channelHandler.Setup(x => x.AddRunner(It.IsAny<Func<Channel<BackgroundJobBase>, Task>>()))
             .Callback<Func<Channel<BackgroundJobBase>, Task>>(callback => runner = callback);
-
         purviewClient.Setup(x => x.GetProtectionScopesAsync(It.IsAny<ProtectionScopesRequest>(), It.IsAny<CancellationToken>()))
-            .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<ContentActivityJob>()), Times.Once);
+        cacheProvider.Verify(x => x.SetAsync(
+            It.Is<PaymentRequiredCacheKey>(key => key.TenantId == "tenant-123"),
+            It.Is<PaymentRequiredCacheEntry>(entry => entry.Message == "Payment required"),
+            It.IsAny<CancellationToken>()), 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");

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)

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))