diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentRuntimeState.cs b/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentRuntimeState.cs index f8e2f3accc3..7aad0db29c9 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentRuntimeState.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentRuntimeState.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using System.Text.Json.Serialization; +using System.Threading; using System.Threading.Tasks; namespace Microsoft.Agents.AI; @@ -17,6 +18,18 @@ namespace Microsoft.Agents.AI; /// internal sealed class BackgroundAgentRuntimeState { + /// + /// Gets an object used to synchronize access to the runtime references held by this instance. + /// + /// + /// Background task registration happens on the agent's tool-invocation path, while + /// may be called concurrently by a host. + /// All mutations of the dictionaries below, and of , must be performed under this lock + /// so that a task can never be registered into an already-released runtime. + /// + [JsonIgnore] + public object SyncRoot { get; } = new(); + /// /// Gets the mapping of task IDs to their in-flight instances. /// @@ -29,4 +42,25 @@ internal sealed class BackgroundAgentRuntimeState /// [JsonIgnore] public Dictionary BackgroundTaskSessions { get; } = []; + + /// + /// Gets the mapping of task IDs to the controlling their run. + /// + /// + /// A source is created when a task is started or continued, and is disposed and removed when the task is + /// finalized, cleared, or when the session is released via . + /// + [JsonIgnore] + public Dictionary TaskCancellations { get; } = []; + + /// + /// Gets or sets a value indicating whether this runtime has been released via + /// . + /// + /// + /// Once released, all in-flight tasks have been cancelled and awaited, and the runtime references have been + /// dropped. Tools that would start new background work refuse to run against a released runtime. + /// + [JsonIgnore] + public bool IsReleased { get; set; } } diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs b/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs index 755aa45769f..aec6a9883ea 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs @@ -35,6 +35,11 @@ namespace Microsoft.Agents.AI; /// /// /// +/// Background tasks are tracked per session and keep running until they complete. When a host is finished with a +/// session it should call to cancel and await any in-flight tasks, so that +/// abandoned work does not continue to invoke models and tools in the background. +/// +/// /// Security considerations: The agents passed to the constructor are delegated /// arbitrary work by the parent agent — the parent sends them text input (which may include content /// derived from the parent's own untrusted context) and receives back whatever text they produce. A @@ -60,7 +65,16 @@ You have access to background agents that can perform work on your behalf. {background_agents} """; + private const string ReleasedRuntimeStartError = + "Error: The background agents runtime for this session has been released. No new background tasks can be started."; + + private const string ReleasedRuntimeContinueError = + "Error: The background agents runtime for this session has been released. Background tasks can no longer be continued."; + + private const string ReleasedTaskCanceledMessage = "Task was canceled because the session was released."; + private readonly Dictionary _agents; + private static readonly TimeSpan s_defaultReleaseTimeout = TimeSpan.FromSeconds(30); private readonly ProviderSessionState _sessionState; private readonly ProviderSessionState _runtimeSessionState; private readonly string _instructions; @@ -148,6 +162,189 @@ public IReadOnlyList GetIncompleteTasks(AgentSession? sessio return incomplete; } + /// + /// Releases all runtime state held for the specified session, cancelling and awaiting any in-flight background tasks. + /// + /// The agent session whose background runtime should be released. + /// + /// to cancel any background tasks that are still running; to require that + /// all background tasks have already completed. + /// + /// + /// The maximum amount of time to wait for cancelled tasks to finish. Defaults to 30 seconds when . + /// Use to wait indefinitely. If the timeout elapses, the remaining tasks are + /// abandoned rather than blocking the caller. + /// + /// The to monitor for cancellation requests while waiting. + /// A task that represents the asynchronous release operation. + /// is . + /// is negative and is not . + /// is and one or more background tasks are still running. + /// + /// + /// Background tasks continue to execute — invoking models and tools — even after a host stops using the session + /// that started them. Hosts should call this method when a conversation ends, or from their own eviction policy, + /// so that abandoned work is stopped instead of running to completion with results nobody will read. + /// + /// + /// This method is idempotent: releasing an already-released session does nothing. Once released, the + /// background_agents_start_task and background_agents_continue_task tools refuse to run for that + /// session, and any tasks that were still running are recorded as + /// so a restored session does not report phantom running work. + /// + /// + public async Task ReleaseSessionAsync( + AgentSession session, + bool cancelRunning = true, + TimeSpan? timeout = null, + CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(session); + + TimeSpan effectiveTimeout = timeout ?? s_defaultReleaseTimeout; + if (effectiveTimeout < TimeSpan.Zero && effectiveTimeout != Timeout.InfiniteTimeSpan) + { + throw new ArgumentOutOfRangeException( + nameof(timeout), + effectiveTimeout, + "The timeout must not be negative, unless it is Timeout.InfiniteTimeSpan."); + } + + BackgroundAgentRuntimeState runtimeState = this._runtimeSessionState.GetOrInitializeState(session); + BackgroundAgentState state = this._sessionState.GetOrInitializeState(session); + + KeyValuePair>[] trackedTasks; + HashSet pendingTaskIds; + + lock (runtimeState.SyncRoot) + { + if (runtimeState.IsReleased) + { + return; + } + + trackedTasks = runtimeState.InFlightTasks.ToArray(); + + // Snapshot which tasks were still pending before anything is cancelled. Tasks that had already + // finished keep their real outcome; only these pending ones are reported as released. + pendingTaskIds = [.. trackedTasks.Where(t => !t.Value.IsCompleted).Select(t => t.Key)]; + + if (!cancelRunning && pendingTaskIds.Count > 0) + { + throw new InvalidOperationException( + $"Cannot release the session because {pendingTaskIds.Count} background task(s) are still running. Pass cancelRunning: true to cancel them."); + } + + runtimeState.IsReleased = true; + + foreach (int taskId in pendingTaskIds) + { + if (runtimeState.TaskCancellations.TryGetValue(taskId, out CancellationTokenSource? cts)) + { + try + { + cts.Cancel(); + } + catch (ObjectDisposedException) + { + // The source was already disposed by a concurrent finalization; nothing to cancel. + } + } + } + } + + try + { + await WaitForTasksAsync(trackedTasks.Select(t => t.Value), effectiveTimeout, cancellationToken).ConfigureAwait(false); + } + finally + { + lock (runtimeState.SyncRoot) + { + // Finalize every tracked task that actually finished, so successful results and real failure + // reasons are preserved rather than being overwritten with a release failure. + foreach (var kvp in trackedTasks) + { + BackgroundTaskInfo? tracked = state.Tasks.FirstOrDefault(t => t.Id == kvp.Key); + if (tracked is null || tracked.Status != BackgroundTaskStatus.Running || !kvp.Value.IsCompleted) + { + continue; + } + + FinalizeTask(tracked, kvp.Value, runtimeState); + + if (kvp.Value.IsCanceled && pendingTaskIds.Contains(kvp.Key)) + { + // Report the actual reason rather than the generic cancellation message. + tracked.ErrorText = ReleasedTaskCanceledMessage; + } + } + + foreach (int taskId in runtimeState.TaskCancellations.Keys.ToArray()) + { + DisposeTaskCancellation(runtimeState, taskId); + } + + runtimeState.InFlightTasks.Clear(); + runtimeState.BackgroundTaskSessions.Clear(); + + // Anything still running was abandoned (for example after the timeout elapsed). + foreach (BackgroundTaskInfo task in state.Tasks) + { + if (task.Status == BackgroundTaskStatus.Running) + { + task.Status = BackgroundTaskStatus.Failed; + task.ErrorText = ReleasedTaskCanceledMessage; + } + } + } + + this._sessionState.SaveState(session, state); + this._runtimeSessionState.SaveState(session, runtimeState); + } + } + + /// + /// Waits for the specified tasks to finish, observing their exceptions and giving up once the timeout elapses. + /// + private static async Task WaitForTasksAsync(IEnumerable tasks, TimeSpan timeout, CancellationToken cancellationToken) + { + // Attach an observer to every task, including those that already completed, so that a fault is always + // observed. Otherwise clearing the last reference to a faulted task can surface an UnobservedTaskException. + Task[] observers = tasks.Select(t => t.ContinueWith( + static antecedent => _ = antecedent.Exception, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default)).ToArray(); + + // Observers never fault, so awaiting them cannot throw. + Task[] pending = observers.Where(o => !o.IsCompleted).ToArray(); + if (pending.Length == 0) + { + return; + } + + Task all = Task.WhenAll(pending); + + if (timeout == Timeout.InfiniteTimeSpan && !cancellationToken.CanBeCanceled) + { + await all.ConfigureAwait(false); + return; + } + + using var delayCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + Task delay = Task.Delay(timeout, delayCts.Token); + + Task winner = await Task.WhenAny(all, delay).ConfigureAwait(false); + delayCts.Cancel(); + + if (winner != all) + { + // The wait was abandoned; the remaining tasks are left to finish on their own. + cancellationToken.ThrowIfCancellationRequested(); + } + } + /// /// Validates the agent collection and builds a case-insensitive name dictionary. /// @@ -204,25 +401,28 @@ private static string BuildDefaultAgentListText(IReadOnlyDictionary? inFlight)) - { - // In-flight reference lost (e.g., after restart/deserialization). - task.Status = BackgroundTaskStatus.Lost; - changed = true; - continue; - } + if (!runtimeState.InFlightTasks.TryGetValue(task.Id, out Task? inFlight)) + { + // In-flight reference lost (e.g., after restart/deserialization). + task.Status = BackgroundTaskStatus.Lost; + changed = true; + continue; + } - if (inFlight.IsCompleted) - { - FinalizeTask(task, inFlight, runtimeState); - changed = true; + if (inFlight.IsCompleted) + { + FinalizeTask(task, inFlight, runtimeState); + changed = true; + } } } @@ -235,6 +435,7 @@ private void TryRefreshTaskState(BackgroundAgentState state, BackgroundAgentRunt /// /// Finalizes a task by extracting results from the completed Task and updating the BackgroundTaskInfo. /// + /// Callers must hold . private static void FinalizeTask(BackgroundTaskInfo taskInfo, Task completedTask, BackgroundAgentRuntimeState runtimeState) { if (completedTask.Status == TaskStatus.RanToCompletion) @@ -256,6 +457,55 @@ private static void FinalizeTask(BackgroundTaskInfo taskInfo, Task + /// Removes and disposes the tracked for the specified task, if any. + /// + /// Callers must hold . + private static void DisposeTaskCancellation(BackgroundAgentRuntimeState runtimeState, int taskId) + { + if (runtimeState.TaskCancellations.TryGetValue(taskId, out CancellationTokenSource? cts)) + { + runtimeState.TaskCancellations.Remove(taskId); + cts.Dispose(); + } + } + + /// + /// Starts a background run for the specified task, tracking both the resulting task and a + /// that allows the run to be cancelled when the session is released. + /// + /// + /// if the run was started and tracked; if the session was + /// released before the run could be registered, in which case nothing is started. + /// + private static bool StartTrackedRun(BackgroundAgentRuntimeState runtimeState, int taskId, AIAgent agent, string input, AgentSession subSession) + { + lock (runtimeState.SyncRoot) + { + // Re-check under the lock: the session may have been released while the caller awaited session creation. + // Starting here would produce a task that is never tracked and therefore never cancelled. + if (runtimeState.IsReleased) + { + return false; + } + + // Replace any cancellation source left over from a previous run of the same task. + DisposeTaskCancellation(runtimeState, taskId); + + var cts = new CancellationTokenSource(); + runtimeState.TaskCancellations[taskId] = cts; + + // Wrap in Task.Run to fork the ExecutionContext. AIAgent.RunAsync is a non-async + // method that synchronously sets the static AsyncLocal CurrentRunContext. Without + // this isolation, the background agent's RunAsync would overwrite the outer (calling) + // agent's CurrentRunContext, corrupting all subsequent tool invocations in the + // same FICC batch. + runtimeState.InFlightTasks[taskId] = Task.Run(() => agent.RunAsync(input, subSession, cancellationToken: cts.Token), cts.Token); + return true; + } } private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeState runtimeState, AgentSession? session) @@ -270,6 +520,11 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS [Description("The request to pass to the background agent.")] string input, [Description("A description of the task used to identify the task later.")] string description) => { + if (runtimeState.IsReleased) + { + return ReleasedRuntimeStartError; + } + if (!this._agents.TryGetValue(agentName, out AIAgent? agent)) { return $"Error: No background agent found with name '{agentName}'. Available agents: {string.Join(", ", this._agents.Keys)}"; @@ -288,13 +543,18 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS // Create a dedicated session for this background task so it can be continued later. AgentSession subSession = await agent.CreateSessionAsync().ConfigureAwait(false); - // Wrap in Task.Run to fork the ExecutionContext. AIAgent.RunAsync is a non-async - // method that synchronously sets the static AsyncLocal CurrentRunContext. Without - // this isolation, the background agent's RunAsync would overwrite the outer (calling) - // agent's CurrentRunContext, corrupting all subsequent tool invocations in the - // same FICC batch. - runtimeState.InFlightTasks[taskId] = Task.Run(() => agent.RunAsync(input, subSession)); - runtimeState.BackgroundTaskSessions[taskId] = subSession; + if (!StartTrackedRun(runtimeState, taskId, agent, input, subSession)) + { + // The session was released while the background session was being created. + state.Tasks.Remove(taskInfo); + this._sessionState.SaveState(session, state); + return ReleasedRuntimeStartError; + } + + lock (runtimeState.SyncRoot) + { + runtimeState.BackgroundTaskSessions[taskId] = subSession; + } this._sessionState.SaveState(session, state); return $"Background task {taskId} started on agent '{agentName}'."; @@ -317,11 +577,14 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS // Collect in-flight tasks matching the requested IDs (including already-completed ones, // since Task.WhenAny returns immediately for completed tasks). var waitableTasks = new List<(int Id, Task Task)>(); - foreach (int id in taskIds) + lock (runtimeState.SyncRoot) { - if (runtimeState.InFlightTasks.TryGetValue(id, out Task? inFlight)) + foreach (int id in taskIds) { - waitableTasks.Add((id, inFlight)); + if (runtimeState.InFlightTasks.TryGetValue(id, out Task? inFlight)) + { + waitableTasks.Add((id, inFlight)); + } } } @@ -351,7 +614,14 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS BackgroundTaskInfo? taskInfo = state.Tasks.FirstOrDefault(t => t.Id == completedEntry.Id); if (taskInfo is not null) { - FinalizeTask(taskInfo, completedEntry.Task, runtimeState); + lock (runtimeState.SyncRoot) + { + if (taskInfo.Status == BackgroundTaskStatus.Running) + { + FinalizeTask(taskInfo, completedEntry.Task, runtimeState); + } + } + this._sessionState.SaveState(session, state); } @@ -420,6 +690,11 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS AIFunctionFactory.Create( (int taskId, string text) => { + if (runtimeState.IsReleased) + { + return ReleasedRuntimeContinueError; + } + this.TryRefreshTaskState(state, runtimeState, session); BackgroundTaskInfo? taskInfo = state.Tasks.FirstOrDefault(t => t.Id == taskId); @@ -443,7 +718,13 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS return $"Error: Agent '{taskInfo.AgentName}' is no longer available."; } - if (!runtimeState.BackgroundTaskSessions.TryGetValue(taskId, out AgentSession? subSession)) + AgentSession? subSession; + lock (runtimeState.SyncRoot) + { + _ = runtimeState.BackgroundTaskSessions.TryGetValue(taskId, out subSession); + } + + if (subSession is null) { return $"Error: Session for task {taskId} is no longer available."; } @@ -454,7 +735,13 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS taskInfo.ErrorText = null; // Wrap in Task.Run to isolate the ExecutionContext (see StartBackgroundTask comment). - runtimeState.InFlightTasks[taskId] = Task.Run(() => agent.RunAsync(text, subSession)); + if (!StartTrackedRun(runtimeState, taskId, agent, text, subSession)) + { + taskInfo.Status = BackgroundTaskStatus.Failed; + taskInfo.ErrorText = ReleasedTaskCanceledMessage; + this._sessionState.SaveState(session, state); + return ReleasedRuntimeContinueError; + } this._sessionState.SaveState(session, state); return $"Task {taskId} continued with new input."; @@ -486,8 +773,12 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS state.Tasks.Remove(taskInfo); // Clean up runtime references. - runtimeState.InFlightTasks.Remove(taskId); - runtimeState.BackgroundTaskSessions.Remove(taskId); + lock (runtimeState.SyncRoot) + { + runtimeState.InFlightTasks.Remove(taskId); + runtimeState.BackgroundTaskSessions.Remove(taskId); + DisposeTaskCancellation(runtimeState, taskId); + } this._sessionState.SaveState(session, state); return $"Task {taskId} cleared."; diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs index f071d5e084f..75cc0cabb21 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs @@ -885,6 +885,422 @@ public async Task CustomAgentListBuilder_UsedForAgentListAsync() #endregion + #region ReleaseSessionAsync Tests + + /// + /// Verify that releasing a session cancels and awaits an in-flight background task. + /// + [Fact] + public async Task ReleaseSessionAsync_CancelsInFlightTaskAsync() + { + // Arrange + var callbackEntered = new TaskCompletionSource(); + var observedCancellation = new TaskCompletionSource(); + var agent = CreateMockAgentWithCancellableCallback("Research", async ct => + { + callbackEntered.SetResult(true); + try + { + await Task.Delay(Timeout.Infinite, ct); + } + catch (OperationCanceledException) + { + observedCancellation.SetResult(true); + throw; + } + + return new AgentResponse(new ChatMessage(ChatRole.Assistant, "never")); + }); + + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + AIFunction startBackgroundTask = GetTool(tools, "background_agents_start_task"); + + await startBackgroundTask.InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + // Wait until the run is actually executing, otherwise cancellation may prevent the + // delegate from ever being scheduled and the cancellation signal would never be set. + Assert.True(await callbackEntered.Task); + + // Act + await provider.ReleaseSessionAsync(session); + + // Assert — the background run observed cancellation and no tasks remain running. + Assert.True(await observedCancellation.Task); + Assert.Empty(provider.GetIncompleteTasks(session)); + } + + /// + /// Verify that releasing a session more than once is a no-op. + /// + [Fact] + public async Task ReleaseSessionAsync_IsIdempotentAsync() + { + // Arrange + var agent = CreateMockAgentWithCancellableCallback("Research", async ct => + { + await Task.Delay(Timeout.Infinite, ct); + return new AgentResponse(new ChatMessage(ChatRole.Assistant, "never")); + }); + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + AIFunction startBackgroundTask = GetTool(tools, "background_agents_start_task"); + + await startBackgroundTask.InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + // Act + await provider.ReleaseSessionAsync(session); + await provider.ReleaseSessionAsync(session); + + // Assert — the second release did not throw and the state remains released. + Assert.Empty(provider.GetIncompleteTasks(session)); + } + + /// + /// Verify that releasing one session does not affect background tasks in another session. + /// + [Fact] + public async Task ReleaseSessionAsync_DoesNotAffectOtherSessionsAsync() + { + // Arrange + var agent = CreateMockAgentWithCancellableCallback("Research", async ct => + { + await Task.Delay(Timeout.Infinite, ct); + return new AgentResponse(new ChatMessage(ChatRole.Assistant, "never")); + }); + var provider = new BackgroundAgentsProvider(new[] { agent }); + + var (toolsA, sessionA) = await CreateToolsForSessionAsync(provider); + var (toolsB, sessionB) = await CreateToolsForSessionAsync(provider); + + await GetTool(toolsA, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task A", + ["description"] = "Session A task", + }); + + await GetTool(toolsB, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task B", + ["description"] = "Session B task", + }); + + // Act + await provider.ReleaseSessionAsync(sessionA); + + // Assert — session B's task is untouched. + Assert.Empty(provider.GetIncompleteTasks(sessionA)); + Assert.Single(provider.GetIncompleteTasks(sessionB)); + } + + /// + /// Verify that releasing with cancelRunning false throws when tasks are still running. + /// + [Fact] + public async Task ReleaseSessionAsync_CancelRunningFalseWithRunningTask_ThrowsAsync() + { + // Arrange + var tcs = new TaskCompletionSource(); + var agent = CreateMockAgentWithRunResult("Research", tcs.Task); + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + + await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + // Act & Assert + await Assert.ThrowsAsync( + () => provider.ReleaseSessionAsync(session, cancelRunning: false)); + + // Assert — the task is left running because the release was rejected. + Assert.Single(provider.GetIncompleteTasks(session)); + + tcs.SetResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "done"))); + } + + /// + /// Verify that releasing with cancelRunning false succeeds when all tasks have completed. + /// + [Fact] + public async Task ReleaseSessionAsync_CancelRunningFalseWithCompletedTask_SucceedsAsync() + { + // Arrange + var agent = CreateMockAgentWithRunResult("Research", Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "done")))); + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + + await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + // Ensure the run has completed before releasing. + await GetTool(tools, "background_agents_wait_for_first_completion").InvokeAsync(new AIFunctionArguments + { + ["taskIds"] = new List { 1 }, + }); + + Assert.Empty(provider.GetIncompleteTasks(session)); + + // Act + await provider.ReleaseSessionAsync(session, cancelRunning: false); + + // Assert — subsequent starts are rejected, confirming the runtime was released. + object? result = await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 2", + ["description"] = "Second task", + }); + + Assert.Contains("released", GetStringResult(result)); + } + + /// + /// Verify that a task still running when the session is released is recorded as failed. + /// + [Fact] + public async Task ReleaseSessionAsync_MarksRunningTasksAsFailedAsync() + { + // Arrange + var agent = CreateMockAgentWithCancellableCallback("Research", async ct => + { + await Task.Delay(Timeout.Infinite, ct); + return new AgentResponse(new ChatMessage(ChatRole.Assistant, "never")); + }); + + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + + await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + // Act + await provider.ReleaseSessionAsync(session); + + // Assert + object? result = await GetTool(tools, "background_agents_get_all_tasks").InvokeAsync(new AIFunctionArguments()); + string text = GetStringResult(result); + Assert.Contains("[Failed]", text); + } + + /// + /// Verify that the start and continue tools return an error after the session is released. + /// + [Fact] + public async Task ReleaseSessionAsync_ToolsReturnErrorAfterReleaseAsync() + { + // Arrange + var agent = CreateMockAgentWithRunResult("Research", Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "done")))); + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + + await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + await provider.ReleaseSessionAsync(session); + + // Act + object? startResult = await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 2", + ["description"] = "Second task", + }); + + object? continueResult = await GetTool(tools, "background_agents_continue_task").InvokeAsync(new AIFunctionArguments + { + ["taskId"] = 1, + ["text"] = "More work", + }); + + // Assert + Assert.Contains("released", GetStringResult(startResult)); + Assert.Contains("released", GetStringResult(continueResult)); + } + + /// + /// Verify that a task ignoring cancellation does not block the release beyond the timeout. + /// + [Fact] + public async Task ReleaseSessionAsync_TimeoutAbandonsUncooperativeTaskAsync() + { + // Arrange — the run never observes its cancellation token. + var tcs = new TaskCompletionSource(); + var agent = CreateMockAgentWithRunResult("Research", tcs.Task); + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + + await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + // Act + await provider.ReleaseSessionAsync(session, timeout: TimeSpan.FromMilliseconds(50)); + + // Assert — release completed despite the task still being pending. + Assert.False(tcs.Task.IsCompleted); + Assert.Empty(provider.GetIncompleteTasks(session)); + + tcs.SetResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "late"))); + } + + /// + /// Verify that releasing a null session throws. + /// + [Fact] + public async Task ReleaseSessionAsync_NullSession_ThrowsAsync() + { + // Arrange + var agent = CreateMockAgent("Research", "Research agent"); + var provider = new BackgroundAgentsProvider(new[] { agent }); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.ReleaseSessionAsync(null!)); + } + + /// + /// Verify that a task which completed but has not yet been refreshed keeps its result instead of + /// being reported as canceled by the release. + /// + [Fact] + public async Task ReleaseSessionAsync_CompletedButUnrefreshedTask_KeepsResultAsync() + { + // Arrange + var runEntered = new TaskCompletionSource(); + var agent = CreateMockAgentWithCancellableCallback("Research", _ => + { + runEntered.TrySetResult(true); + return Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "Result 1"))); + }); + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + + await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + // Wait until the run has actually started, so it cannot be cancelled before it produces a result. + // The persisted state is never refreshed, so the task is still marked Running when the release begins. + Assert.True(await runEntered.Task); + + // Act + await provider.ReleaseSessionAsync(session); + + // Assert — the successful result is preserved rather than overwritten with a release failure. + object? result = await GetTool(tools, "background_agents_get_task_results").InvokeAsync(new AIFunctionArguments + { + ["taskId"] = 1, + }); + + Assert.Equal("Result 1", GetStringResult(result)); + } + + /// + /// Verify that an invalid timeout is rejected before the session is released. + /// + [Fact] + public async Task ReleaseSessionAsync_NegativeTimeout_ThrowsWithoutReleasingAsync() + { + // Arrange + var tcs = new TaskCompletionSource(); + var agent = CreateMockAgentWithRunResult("Research", tcs.Task); + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + + await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }); + + // Act & Assert + await Assert.ThrowsAsync( + () => provider.ReleaseSessionAsync(session, timeout: TimeSpan.FromSeconds(-5))); + + // Assert — the session was not released, so the task is untouched and new tasks can still start. + Assert.Single(provider.GetIncompleteTasks(session)); + + object? startResult = await GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 2", + ["description"] = "Second task", + }); + + Assert.DoesNotContain("released", GetStringResult(startResult)); + + tcs.SetResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "done"))); + } + + /// + /// Verify that a start racing with a release does not register an untracked background task. + /// + [Fact] + public async Task ReleaseSessionAsync_DuringStart_RefusesToRegisterTaskAsync() + { + // Arrange — session creation blocks so the release can happen mid-start. + var sessionCreationGate = new TaskCompletionSource(); + var runStarted = false; + var agent = CreateMockAgentWithGatedSession( + "Research", + sessionCreationGate.Task, + () => + { + runStarted = true; + return Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "done"))); + }); + + var (tools, provider, session) = await CreateToolsWithSessionAsync(agent); + + Task startTask = GetTool(tools, "background_agents_start_task").InvokeAsync(new AIFunctionArguments + { + ["agentName"] = "Research", + ["input"] = "Task 1", + ["description"] = "First task", + }).AsTask(); + + // Act — release while the start is awaiting session creation, then let creation finish. + await provider.ReleaseSessionAsync(session); + sessionCreationGate.SetResult(true); + + object? startResult = await startTask; + + // Assert — the start was refused, no run was launched, and no task was recorded. + Assert.Contains("released", GetStringResult(startResult)); + Assert.False(runStarted); + Assert.Empty(provider.GetIncompleteTasks(session)); + + object? allTasks = await GetTool(tools, "background_agents_get_all_tasks").InvokeAsync(new AIFunctionArguments()); + Assert.Equal("No tasks.", GetStringResult(allTasks)); + } + + #endregion + #region Helper Methods private static AIAgent CreateMockAgent(string? name, string? description) @@ -944,6 +1360,69 @@ private static AIAgent CreateMockAgentWithCallback(string name, Func> callback) + { + var mock = new Mock(); + mock.SetupGet(a => a.Name).Returns(name); + mock.Protected() + .Setup>( + "CreateSessionCoreAsync", + ItExpr.IsAny()) + .Returns(new ValueTask(new ChatClientAgentSession())); + mock.Protected() + .Setup>( + "RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns((IEnumerable _, AgentSession _, AgentRunOptions _, CancellationToken ct) => callback(ct)); + return mock.Object; + } + + private static AIAgent CreateMockAgentWithGatedSession(string name, Task sessionGate, Func> callback) + { + var mock = new Mock(); + mock.SetupGet(a => a.Name).Returns(name); + mock.Protected() + .Setup>( + "CreateSessionCoreAsync", + ItExpr.IsAny()) + .Returns(async () => + { + await sessionGate; + return new ChatClientAgentSession(); + }); + mock.Protected() + .Setup>( + "RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .Returns(callback); + return mock.Object; + } + + private static async Task<(IEnumerable Tools, BackgroundAgentsProvider Provider, AgentSession Session)> CreateToolsWithSessionAsync(AIAgent agent) + { + var provider = new BackgroundAgentsProvider(new[] { agent }); + var (tools, session) = await CreateToolsForSessionAsync(provider); + return (tools, provider, session); + } + + private static async Task<(IEnumerable Tools, AgentSession Session)> CreateToolsForSessionAsync(BackgroundAgentsProvider provider) + { + var mockAgent = new Mock().Object; + var session = new ChatClientAgentSession(); +#pragma warning disable MAAI001 + var context = new AIContextProvider.InvokingContext(mockAgent, session, new AIContext()); +#pragma warning restore MAAI001 + + AIContext result = await provider.InvokingAsync(context); + return (result.Tools!, session); + } + private static AIContextProvider.InvokingContext CreateInvokingContext() { var mockAgent = new Mock().Object;