.NET: Add timeout for wait-for-first-completion

dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProvider.cs

@@ -27,7 +27,7 @@ namespace Microsoft.Agents.AI;
 /// This provider exposes the following tools to the agent:
 /// <list type="bullet">
 /// <item><description><c>background_agents_start_task</c> — Start a background task on a named agent with text input. Returns the task ID.</description></item>
-/// <item><description><c>background_agents_wait_for_first_completion</c> — Block until the first of the specified tasks completes. Returns the completed task's ID.</description></item>
+/// <item><description><c>background_agents_wait_for_first_completion</c> — Wait until the first specified task completes or the configured timeout expires. A timeout leaves the tasks running so the tool can be called again.</description></item>
 /// <item><description><c>background_agents_get_task_results</c> — Retrieve the text output of a completed background task.</description></item>
 /// <item><description><c>background_agents_get_all_tasks</c> — List all background tasks with their IDs, statuses, descriptions, and agent names.</description></item>
 /// <item><description><c>background_agents_continue_task</c> — Send follow-up input to a completed background task's session to resume work.</description></item>
@@ -78,6 +78,7 @@ You have access to background agents that can perform work on your behalf.
     private readonly ProviderSessionState<BackgroundAgentState> _sessionState;
     private readonly ProviderSessionState<BackgroundAgentRuntimeState> _runtimeSessionState;
     private readonly string _instructions;
+    private readonly TimeSpan _waitTimeout;
     private IReadOnlyList<string>? _stateKeys;
 
     /// <summary>
@@ -92,11 +93,30 @@ You have access to background agents that can perform work on your behalf.
     /// <param name="options">Optional settings controlling the provider behavior.</param>
     /// <exception cref="ArgumentNullException"><paramref name="agents"/> is <see langword="null"/>.</exception>
     /// <exception cref="ArgumentException">An agent has a null or empty name, or agent names are not unique.</exception>
+    /// <exception cref="ArgumentOutOfRangeException">
+    /// <see cref="BackgroundAgentsProviderOptions.WaitTimeout"/> is not positive or exceeds the maximum supported delay.
+    /// </exception>
     public BackgroundAgentsProvider(IEnumerable<AIAgent> agents, BackgroundAgentsProviderOptions? options = null)
     {
         _ = Throw.IfNull(agents);
 
         this._agents = ValidateAndBuildAgentDictionary(agents);
+        this._waitTimeout = options?.WaitTimeout ?? BackgroundAgentsProviderOptions.DefaultWaitTimeout;
+        if (this._waitTimeout <= TimeSpan.Zero)
+        {
+            throw new ArgumentOutOfRangeException(
+                nameof(options),
+                this._waitTimeout,
+                $"{nameof(BackgroundAgentsProviderOptions.WaitTimeout)} must be positive.");
+        }
+
+        if (this._waitTimeout > BackgroundAgentsProviderOptions.MaximumWaitTimeout)
+        {
+            throw new ArgumentOutOfRangeException(
+                nameof(options),
+                this._waitTimeout,
+                $"{nameof(BackgroundAgentsProviderOptions.WaitTimeout)} must not exceed {BackgroundAgentsProviderOptions.MaximumWaitTimeout.TotalMilliseconds} milliseconds.");
+        }
 
         string baseInstructions = options?.Instructions ?? DefaultInstructions;
         string agentListText = options?.AgentListBuilder is not null
@@ -672,8 +692,20 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS
                         return "Error: None of the specified task IDs correspond to running tasks.";
                     }
 
-                    // Wait for the first one to complete.
-                    Task completedTask = await Task.WhenAny(waitableTasks.Select(t => t.Task)).ConfigureAwait(false);
+                    // Wait for the first task to complete, but return control without stopping the tasks if the timeout elapses.
+                    Task<Task<AgentResponse>> firstCompletionTask = Task.WhenAny(waitableTasks.Select(t => t.Task));
+                    using var timeoutCts = new CancellationTokenSource();
+                    Task timeoutTask = Task.Delay(this._waitTimeout, timeoutCts.Token);
+                    Task winner = await Task.WhenAny(firstCompletionTask, timeoutTask).ConfigureAwait(false);
+                    timeoutCts.Cancel();
+
+                    if (winner == timeoutTask && !firstCompletionTask.IsCompleted)
+                    {
+                        return FormattableString.Invariant(
+                            $"No background task completed within {this._waitTimeout.TotalSeconds:g} seconds. The tasks are still running; call this tool again if you wish to continue waiting.");
+                    }
+
+                    Task<AgentResponse> completedTask = await firstCompletionTask.ConfigureAwait(false);
 
                     // Find which ID completed.
                     var completedEntry = waitableTasks.First(t => t.Task == completedTask);
@@ -698,7 +730,7 @@ private AITool[] CreateTools(BackgroundAgentState state, BackgroundAgentRuntimeS
                 new AIFunctionFactoryOptions
                 {
                     Name = "background_agents_wait_for_first_completion",
-                    Description = "Block until the first of the specified background tasks completes. Provide one or more task IDs. Returns a status message containing the ID of the task that completed first.",
+                    Description = "Wait until the first of the specified background tasks completes or the configured timeout expires. Provide one or more task IDs. Returns a status message containing the ID of the task that completed first. On timeout, the tasks remain running and this tool can be called again to continue waiting.",
                     SerializerOptions = serializerOptions,
                 }),
 

dotnet/src/Microsoft.Agents.AI/Harness/BackgroundAgents/BackgroundAgentsProviderOptions.cs

@@ -13,6 +13,9 @@ namespace Microsoft.Agents.AI;
 [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
 public sealed class BackgroundAgentsProviderOptions
 {
+    internal static readonly TimeSpan DefaultWaitTimeout = TimeSpan.FromMinutes(5);
+    internal static readonly TimeSpan MaximumWaitTimeout = TimeSpan.FromMilliseconds(uint.MaxValue - 1L);
+
     /// <summary>
     /// Gets or sets custom instructions provided to the agent for using the background agent tools.
     /// </summary>
@@ -36,4 +39,14 @@ public sealed class BackgroundAgentsProviderOptions
     /// a formatted string describing the available background agents.
     /// </value>
     public Func<IReadOnlyDictionary<string, AIAgent>, string>? AgentListBuilder { get; set; }
+
+    /// <summary>
+    /// Gets or sets the maximum amount of time the wait tool blocks for a background task to complete.
+    /// </summary>
+    /// <value>
+    /// The default is five minutes. The value must be greater than <see cref="TimeSpan.Zero"/> and must not
+    /// exceed 4,294,967,294 milliseconds, the maximum delay supported by the targeted .NET runtimes.
+    /// When the timeout elapses, the tool returns control to the agent and leaves the background tasks running.
+    /// </value>
+    public TimeSpan WaitTimeout { get; set; } = DefaultWaitTimeout;
 }

dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/BackgroundAgents/BackgroundAgentsProviderTests.cs

@@ -96,6 +96,76 @@ public void Constructor_ValidAgents_Succeeds()
         Assert.NotNull(provider);
     }
 
+    /// <summary>
+    /// Verify that the default wait timeout is five minutes.
+    /// </summary>
+    [Fact]
+    public void Options_DefaultWaitTimeout_IsFiveMinutes()
+    {
+        // Arrange & Act
+        var options = new BackgroundAgentsProviderOptions();
+
+        // Assert
+        Assert.Equal(TimeSpan.FromMinutes(5), options.WaitTimeout);
+    }
+
+    /// <summary>
+    /// Verify that the constructor rejects non-positive wait timeouts.
+    /// </summary>
+    /// <param name="seconds">The timeout value in seconds.</param>
+    [Theory]
+    [InlineData(0)]
+    [InlineData(-1)]
+    public void Constructor_NonPositiveWaitTimeout_Throws(int seconds)
+    {
+        // Arrange
+        var agent = CreateMockAgent("Research", "Research agent");
+        var options = new BackgroundAgentsProviderOptions
+        {
+            WaitTimeout = TimeSpan.FromSeconds(seconds),
+        };
+
+        // Act & Assert
+        Assert.Throws<ArgumentOutOfRangeException>(() => new BackgroundAgentsProvider(new[] { agent }, options));
+    }
+
+    /// <summary>
+    /// Verify that the constructor rejects wait timeouts above the maximum supported delay.
+    /// </summary>
+    [Fact]
+    public void Constructor_ExcessiveWaitTimeout_Throws()
+    {
+        // Arrange
+        var agent = CreateMockAgent("Research", "Research agent");
+        var options = new BackgroundAgentsProviderOptions
+        {
+            WaitTimeout = TimeSpan.MaxValue,
+        };
+
+        // Act & Assert
+        Assert.Throws<ArgumentOutOfRangeException>(() => new BackgroundAgentsProvider(new[] { agent }, options));
+    }
+
+    /// <summary>
+    /// Verify that the constructor accepts the maximum supported wait timeout.
+    /// </summary>
+    [Fact]
+    public void Constructor_MaximumWaitTimeout_Succeeds()
+    {
+        // Arrange
+        var agent = CreateMockAgent("Research", "Research agent");
+        var options = new BackgroundAgentsProviderOptions
+        {
+            WaitTimeout = BackgroundAgentsProviderOptions.MaximumWaitTimeout,
+        };
+
+        // Act
+        var provider = new BackgroundAgentsProvider(new[] { agent }, options);
+
+        // Assert
+        Assert.NotNull(provider);
+    }
+
     #endregion
 
     #region ProvideAIContextAsync Tests
@@ -300,6 +370,80 @@ public async Task WaitForFirstCompletion_EmptyList_ReturnsErrorAsync()
         Assert.Contains("Error", GetStringResult(result));
     }
 
+    /// <summary>
+    /// Verify that WaitForFirstCompletion returns after the configured timeout without stopping the task.
+    /// </summary>
+    [Fact]
+    public async Task WaitForFirstCompletion_TimeoutLeavesTaskRunningAsync()
+    {
+        // Arrange
+        var tcs = new TaskCompletionSource<AgentResponse>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var agent = CreateMockAgentWithRunResult("Research", tcs.Task);
+        var provider = new BackgroundAgentsProvider(
+            new[] { agent },
+            new BackgroundAgentsProviderOptions { WaitTimeout = TimeSpan.FromMilliseconds(50) });
+        var (tools, session) = await CreateToolsForSessionAsync(provider);
+        AIFunction startBackgroundTask = GetTool(tools, "background_agents_start_task");
+        AIFunction waitForFirst = GetTool(tools, "background_agents_wait_for_first_completion");
+        AIFunction getResults = GetTool(tools, "background_agents_get_task_results");
+
+        await startBackgroundTask.InvokeAsync(new AIFunctionArguments
+        {
+            ["agentName"] = "Research",
+            ["input"] = "Task 1",
+            ["description"] = "First task",
+        });
+
+        // Act
+        object? result = await waitForFirst.InvokeAsync(new AIFunctionArguments
+        {
+            ["taskIds"] = new List<int> { 1 },
+        });
+
+        // Assert
+        Assert.Equal(
+            "No background task completed within 0.05 seconds. The tasks are still running; call this tool again if you wish to continue waiting.",
+            GetStringResult(result));
+
+        BackgroundAgentRuntimeState runtimeState = GetRuntimeState(provider, session);
+        Assert.False(runtimeState.InFlightTasks[1].IsCompleted);
+
+        object? taskResult = await getResults.InvokeAsync(new AIFunctionArguments
+        {
+            ["taskId"] = 1,
+        });
+        Assert.Contains("still running", GetStringResult(taskResult));
+
+        tcs.SetResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "done")));
+        await runtimeState.InFlightTasks[1];
+        await waitForFirst.InvokeAsync(new AIFunctionArguments
+        {
+            ["taskIds"] = new List<int> { 1 },
+        });
+    }
+
+    /// <summary>
+    /// Verify that the wait timeout is controlled by the provider rather than exposed to the model.
+    /// </summary>
+    [Fact]
+    public async Task WaitForFirstCompletion_OnlyTaskIdsAreModelSettableAsync()
+    {
+        // Arrange
+        var agent = CreateMockAgent("Research", "Research agent");
+        var (tools, _) = await CreateToolsWithProviderAsync(agent);
+        AIFunction waitForFirst = GetTool(tools, "background_agents_wait_for_first_completion");
+
+        // Act
+        JsonElement properties = waitForFirst.JsonSchema.GetProperty("properties");
+
+        // Assert
+        JsonProperty property = Assert.Single(properties.EnumerateObject());
+        Assert.Equal("taskIds", property.Name);
+        Assert.Contains("configured timeout expires", waitForFirst.Description);
+        Assert.Contains("tasks remain running", waitForFirst.Description);
+        Assert.Contains("called again", waitForFirst.Description);
+    }
+
     #endregion
 
     #region GetBackgroundTaskResults Tests