.NET Purview Middleware: Improve Background Job Runner Injection

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

@@ -12,7 +12,7 @@ namespace Microsoft.Agents.AI.Purview;
 /// <summary>
 /// Service that runs jobs in background threads.
 /// </summary>
-internal sealed class BackgroundJobRunner
+internal sealed class BackgroundJobRunner : IBackgroundJobRunner
 {
     private readonly IChannelHandler _channelHandler;
     private readonly IPurviewClient _purviewClient;
@@ -70,4 +70,12 @@ private async Task RunJobAsync(BackgroundJobBase job)
                 break;
         }
     }
+
+    /// <summary>
+    /// Shutdown the job runners.
+    /// </summary>
+    public async Task ShutdownAsync()
+    {
+        await this._channelHandler.StopAndWaitForCompletionAsync().ConfigureAwait(false);
+    }
 }

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

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Threading.Tasks;
+
+namespace Microsoft.Agents.AI.Purview;
+
+/// <summary>
+/// An interface for a class that managed background jobs.
+/// </summary>
+internal interface IBackgroundJobRunner
+{
+    /// <summary>
+    /// Shutdown the background jobs.
+    /// </summary>
+    Task ShutdownAsync();
+}

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

@@ -41,7 +41,7 @@ private static PurviewWrapper CreateWrapper(TokenCredential tokenCredential, Pur
         services.AddSingleton<PurviewWrapper>();
         services.AddSingleton(Channel.CreateBounded<BackgroundJobBase>(purviewSettings.PendingBackgroundJobLimit));
         services.AddSingleton<IChannelHandler, ChannelHandler>();
-        services.AddSingleton<BackgroundJobRunner>();
+        services.AddSingleton<IBackgroundJobRunner, BackgroundJobRunner>();
         ServiceProvider serviceProvider = services.BuildServiceProvider();
 
         return serviceProvider.GetRequiredService<PurviewWrapper>();

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

@@ -18,21 +18,21 @@ internal sealed class PurviewWrapper : IDisposable
     private readonly ILogger _logger;
     private readonly IScopedContentProcessor _scopedProcessor;
     private readonly PurviewSettings _purviewSettings;
-    private readonly IChannelHandler _channelHandler;
+    private readonly IBackgroundJobRunner _backgroundJobRunner;
 
     /// <summary>
     /// Creates a new <see cref="PurviewWrapper"/> instance.
     /// </summary>
     /// <param name="scopedProcessor">The scoped processor used to orchestrate the calls to Purview.</param>
     /// <param name="purviewSettings">The settings for Purview integration.</param>
     /// <param name="logger">The logger used for logging.</param>
-    /// <param name="channelHandler">The channel handler used to queue background jobs and add job runners.</param>
-    public PurviewWrapper(IScopedContentProcessor scopedProcessor, PurviewSettings purviewSettings, ILogger logger, IChannelHandler channelHandler)
+    /// <param name="backgroundJobRunner">The runner used manage background jobs.</param>
+    public PurviewWrapper(IScopedContentProcessor scopedProcessor, PurviewSettings purviewSettings, ILogger logger, IBackgroundJobRunner backgroundJobRunner)
     {
         this._scopedProcessor = scopedProcessor;
         this._purviewSettings = purviewSettings;
         this._logger = logger;
-        this._channelHandler = channelHandler;
+        this._backgroundJobRunner = backgroundJobRunner;
     }
 
     private static string GetThreadIdFromAgentThread(AgentThread? thread, IEnumerable<ChatMessage> messages)
@@ -203,7 +203,7 @@ public async Task<AgentResponse> ProcessAgentContentAsync(IEnumerable<ChatMessag
     public void Dispose()
     {
 #pragma warning disable VSTHRD002 // Need to wait for pending jobs to complete.
-        this._channelHandler.StopAndWaitForCompletionAsync().GetAwaiter().GetResult();
+        this._backgroundJobRunner.ShutdownAsync().GetAwaiter().GetResult();
 #pragma warning restore VSTHRD002 // Need to wait for pending jobs to complete.
     }
 }

dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewWrapperTests.cs

@@ -18,22 +18,22 @@ namespace Microsoft.Agents.AI.Purview.UnitTests;
 public sealed class PurviewWrapperTests : IDisposable
 {
     private readonly Mock<IScopedContentProcessor> _mockProcessor;
-    private readonly IChannelHandler _channelHandler;
+    private readonly IBackgroundJobRunner _backgroundJobRunner;
     private readonly PurviewSettings _settings;
     private readonly PurviewWrapper _wrapper;
 
     public PurviewWrapperTests()
     {
         this._mockProcessor = new Mock<IScopedContentProcessor>();
-        this._channelHandler = Mock.Of<IChannelHandler>();
         this._settings = new PurviewSettings("TestApp")
         {
             TenantId = "tenant-123",
             PurviewAppLocation = new PurviewAppLocation(PurviewLocationType.Application, "app-123"),
             BlockedPromptMessage = "Prompt blocked by policy",
             BlockedResponseMessage = "Response blocked by policy"
         };
-        this._wrapper = new PurviewWrapper(this._mockProcessor.Object, this._settings, NullLogger.Instance, this._channelHandler);
+        this._backgroundJobRunner = Mock.Of<IBackgroundJobRunner>();
+        this._wrapper = new PurviewWrapper(this._mockProcessor.Object, this._settings, NullLogger.Instance, this._backgroundJobRunner);
     }
 
     #region ProcessChatContentAsync Tests
@@ -151,7 +151,7 @@ public async Task ProcessChatContentAsync_WithIgnoreExceptions_ContinuesOnPrompt
             IgnoreExceptions = true,
             PurviewAppLocation = new PurviewAppLocation(PurviewLocationType.Application, "app-123")
         };
-        var wrapper = new PurviewWrapper(this._mockProcessor.Object, settingsWithIgnore, NullLogger.Instance, this._channelHandler);
+        var wrapper = new PurviewWrapper(this._mockProcessor.Object, settingsWithIgnore, NullLogger.Instance, this._backgroundJobRunner);
 
         var messages = new List<ChatMessage>
         {
@@ -371,7 +371,7 @@ public async Task ProcessAgentContentAsync_WithIgnoreExceptions_ContinuesOnError
             IgnoreExceptions = true,
             PurviewAppLocation = new PurviewAppLocation(PurviewLocationType.Application, "app-123")
         };
-        var wrapper = new PurviewWrapper(this._mockProcessor.Object, settingsWithIgnore, NullLogger.Instance, this._channelHandler);
+        var wrapper = new PurviewWrapper(this._mockProcessor.Object, settingsWithIgnore, NullLogger.Instance, this._backgroundJobRunner);
 
         var messages = new List<ChatMessage>
         {