diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs b/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs index efe376718b..85a4fa54c3 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/BackgroundJobRunner.cs @@ -12,7 +12,7 @@ namespace Microsoft.Agents.AI.Purview; /// /// Service that runs jobs in background threads. /// -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; } } + + /// + /// Shutdown the job runners. + /// + public async Task ShutdownAsync() + { + await this._channelHandler.StopAndWaitForCompletionAsync().ConfigureAwait(false); + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/IBackgroundJobRunner.cs b/dotnet/src/Microsoft.Agents.AI.Purview/IBackgroundJobRunner.cs new file mode 100644 index 0000000000..e9c3d0d54e --- /dev/null +++ b/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; + +/// +/// An interface for a class that manages background jobs. +/// +internal interface IBackgroundJobRunner +{ + /// + /// Shutdown the background jobs. + /// + Task ShutdownAsync(); +} diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewExtensions.cs index 4095345d99..2458db94a3 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewExtensions.cs @@ -41,7 +41,7 @@ private static PurviewWrapper CreateWrapper(TokenCredential tokenCredential, Pur services.AddSingleton(); services.AddSingleton(Channel.CreateBounded(purviewSettings.PendingBackgroundJobLimit)); services.AddSingleton(); - services.AddSingleton(); + services.AddSingleton(); ServiceProvider serviceProvider = services.BuildServiceProvider(); return serviceProvider.GetRequiredService(); diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs index 835b0146a8..14cddbe5c4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/PurviewWrapper.cs @@ -18,7 +18,7 @@ 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; /// /// Creates a new instance. @@ -26,13 +26,13 @@ internal sealed class PurviewWrapper : IDisposable /// The scoped processor used to orchestrate the calls to Purview. /// The settings for Purview integration. /// The logger used for logging. - /// The channel handler used to queue background jobs and add job runners. - public PurviewWrapper(IScopedContentProcessor scopedProcessor, PurviewSettings purviewSettings, ILogger logger, IChannelHandler channelHandler) + /// The runner used to manage background jobs. + 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 messages) @@ -203,7 +203,7 @@ public async Task ProcessAgentContentAsync(IEnumerable _mockProcessor; - private readonly IChannelHandler _channelHandler; + private readonly IBackgroundJobRunner _backgroundJobRunner; private readonly PurviewSettings _settings; private readonly PurviewWrapper _wrapper; public PurviewWrapperTests() { this._mockProcessor = new Mock(); - this._channelHandler = Mock.Of(); this._settings = new PurviewSettings("TestApp") { TenantId = "tenant-123", @@ -33,7 +32,8 @@ public PurviewWrapperTests() 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(); + 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 { @@ -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 {