Python: bound background_agents_wait_for_first_completion with a timeout

python/packages/core/agent_framework/_harness/_agent.py

@@ -24,7 +24,7 @@
 from .._skills import SkillsProvider
 from .._telemetry import FeatureIndex, mark_feature_used
 from .._types import ChatOptions
-from ._background_agents import BackgroundAgentsProvider
+from ._background_agents import DEFAULT_BACKGROUND_AGENTS_WAIT_TIMEOUT_SECONDS, BackgroundAgentsProvider
 from ._file_access import AgentFileStore, FileAccessProvider, FileSystemAgentFileStore
 from ._file_memory import FileMemoryProvider
 from ._loop import DEFAULT_MAX_ITERATIONS, AgentLoopMiddleware
@@ -160,6 +160,7 @@ def _assemble_context_providers(
     skills_paths: str | Path | Sequence[str | Path] | None,
     background_agents: Sequence[SupportsAgentRun] | None,
     background_agents_instructions: str | None,
+    background_agents_wait_timeout_seconds: float | None,
     shell_context_provider: ContextProvider | None,
     extra_context_providers: Sequence[ContextProvider] | None,
 ) -> list[ContextProvider]:
@@ -205,7 +206,13 @@ def _assemble_context_providers(
 
     # Background agents are opt-in: only added when agents are provided.
     if background_agents:
-        providers.append(BackgroundAgentsProvider(background_agents, instructions=background_agents_instructions))
+        providers.append(
+            BackgroundAgentsProvider(
+                background_agents,
+                instructions=background_agents_instructions,
+                wait_timeout_seconds=background_agents_wait_timeout_seconds,
+            )
+        )
 
     # Shell environment provider is opt-in: only added when a shell tool was wired.
     if shell_context_provider is not None:
@@ -329,6 +336,7 @@ def create_harness_agent(
     skills_paths: str | Path | Sequence[str | Path] | None = None,
     background_agents: Sequence[SupportsAgentRun] | None = None,
     background_agents_instructions: str | None = None,
+    background_agents_wait_timeout_seconds: float | None = DEFAULT_BACKGROUND_AGENTS_WAIT_TIMEOUT_SECONDS,
     shell_executor: ShellExecutor | None = None,
     shell_environment_provider_options: ShellEnvironmentProviderOptions | None = None,
     disable_web_search: bool = False,
@@ -478,6 +486,10 @@ def create_harness_agent(
         background_agents_instructions: Optional instruction override for the
             ``BackgroundAgentsProvider``. May include ``{background_agents}`` placeholder
             which will be replaced with the agent listing.
+        background_agents_wait_timeout_seconds: Default upper bound, in seconds, applied when the
+            agent waits for a background task to complete. Set to ``None`` to wait without a bound.
+            Bounding the wait keeps a background agent that never completes from suspending this
+            agent's run indefinitely.
         shell_executor: Optional shell tool that enables shell command execution. When
             provided, the shell tool and a ``ShellEnvironmentProvider`` are automatically
             added (provided the client supports shell tools; otherwise a warning is logged
@@ -601,6 +613,7 @@ def create_harness_agent(
         skills_paths=skills_paths,
         background_agents=background_agents,
         background_agents_instructions=background_agents_instructions,
+        background_agents_wait_timeout_seconds=background_agents_wait_timeout_seconds,
         shell_context_provider=shell_provider,
         extra_context_providers=context_providers,
     )

python/packages/core/agent_framework/_harness/_agent.pyi

@@ -77,6 +77,7 @@ def create_harness_agent(
     skills_paths: str | Path | Sequence[str | Path] | None = None,
     background_agents: Sequence[SupportsAgentRun] | None = None,
     background_agents_instructions: str | None = None,
+    background_agents_wait_timeout_seconds: float | None = ...,
     shell_executor: _ShellExecutorLike | None = None,
     shell_environment_provider_options: _ShellEnvironmentProviderOptionsLike | None = None,
     disable_web_search: bool = False,

python/packages/core/agent_framework/_harness/_background_agents.py

@@ -11,6 +11,7 @@
 
 import asyncio
 import logging
+import math
 from collections.abc import Awaitable, MutableMapping, Sequence
 from dataclasses import dataclass, field
 from enum import Enum
@@ -28,6 +29,16 @@
 
 DEFAULT_BACKGROUND_AGENTS_SOURCE_ID = "background_agents"
 
+#: Default upper bound, in seconds, for ``background_agents_wait_for_first_completion``.
+#: Chosen to be generous enough that healthy long-running child agents are never cut short,
+#: while still guaranteeing the parent's function-calling loop regains control.
+DEFAULT_BACKGROUND_AGENTS_WAIT_TIMEOUT_SECONDS = 300.0
+
+#: Upper bound, in seconds, on a single internal wait slice. The wait tool waits in slices of at
+#: most this length so that ``_refresh_task_state`` runs between slices and can promote a task
+#: whose runtime reference has disappeared to ``LOST`` well before the overall timeout elapses.
+_WAIT_SLICE_SECONDS = 5.0
+
 DEFAULT_BACKGROUND_AGENTS_INSTRUCTIONS = """\
 ## Background Agents
 
@@ -36,6 +47,9 @@
 - Use the `background_agents_*` tools to start tasks on background agents and check their results.
 - Creating a background task does not block, and background tasks run concurrently.
 - Important: Always wait for outstanding tasks to finish before you finish processing.
+- `background_agents_wait_for_first_completion` is bounded by a timeout. If it reports that it timed \
+out, the tasks it lists as still running have not finished: wait again, or use \
+background_agents_get_all_tasks to check their status. Do not treat a timeout as completion.
 - Important: After retrieving results from a completed task, clear it with \
 background_agents_clear_completed_task to free memory, unless you plan to continue it with \
 background_agents_continue_task.
@@ -152,6 +166,38 @@ def _log_abandoned_background_task(task: asyncio.Task[Any]) -> None:
         logger.debug("Abandoned background task raised: %s", exception)
 
 
+def _validate_wait_timeout(timeout_seconds: float | None) -> float | None:
+    """Validate a wait timeout, returning it unchanged when acceptable.
+
+    Args:
+        timeout_seconds: Timeout in seconds, or ``None`` to wait without a bound.
+
+    Returns:
+        The validated timeout.
+
+    Raises:
+        ValueError: If the timeout is not a positive number.
+    """
+    if timeout_seconds is None:
+        return None
+    # bool is an int subclass, and a timeout of True/False is always a caller mistake.
+    if isinstance(timeout_seconds, bool) or not isinstance(timeout_seconds, (int, float)):
+        raise ValueError(f"Background agent wait timeout must be a number or None; got {timeout_seconds!r}.")
+    # Convert before range checks: a very large int raises OverflowError here, which would escape as
+    # something other than the documented ValueError.
+    try:
+        timeout = float(timeout_seconds)
+    except OverflowError as exc:
+        raise ValueError(f"Background agent wait timeout is too large; got {timeout_seconds!r}.") from exc
+    # Rejects NaN and infinity together: an infinite deadline would restore the unbounded wait that
+    # this timeout exists to prevent. Callers who want that must pass None explicitly.
+    if not math.isfinite(timeout):
+        raise ValueError(f"Background agent wait timeout must be a finite number or None; got {timeout_seconds!r}.")
+    if timeout <= 0:
+        raise ValueError(f"Background agent wait timeout must be greater than 0; got {timeout_seconds!r}.")
+    return timeout
+
+
 def _validate_and_build_agent_dict(agents: Sequence[SupportsAgentRun]) -> dict[str, SupportsAgentRun]:
     """Validate agents and build a case-insensitive lookup dict.
 
@@ -259,6 +305,58 @@ def _refresh_task_state(
     return tasks
 
 
+async def _wait_first_completed(
+    session: AgentSession,
+    state: dict[str, Any],
+    runtime: _RuntimeState,
+    waitable: list[tuple[int, asyncio.Task[AgentResponse[Any]]]],
+    *,
+    timeout: float | None,
+    source_id: str,
+) -> tuple[set[asyncio.Task[AgentResponse[Any]]], bool]:
+    """Wait for the first of ``waitable`` to finish, bounded by ``timeout``.
+
+    The wait is performed in slices of at most ``_WAIT_SLICE_SECONDS``. Between slices
+    ``_refresh_task_state`` runs so that a task whose runtime reference has disappeared is promoted
+    to ``LOST``, ending the wait early rather than stalling for the whole timeout.
+
+    Returns:
+        A tuple of the tasks that completed and whether the deadline expired. The task set is empty
+        both when the deadline expired and when a requested task reached a terminal state without
+        its asyncio task finishing; the boolean distinguishes the two.
+    """
+    loop = asyncio.get_running_loop()
+    deadline = None if timeout is None else loop.time() + timeout
+    pending_tasks = [task for _, task in waitable]
+    waited_ids = set(task_id for task_id, _ in waitable)
+
+    while True:
+        slice_timeout = _WAIT_SLICE_SECONDS
+        if deadline is not None:
+            remaining = deadline - loop.time()
+            if remaining <= 0:
+                return set(), True
+            slice_timeout = min(slice_timeout, remaining)
+
+        done, _ = await asyncio.wait(
+            pending_tasks,
+            timeout=slice_timeout,
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+        if done:
+            return done, False
+
+        # Nothing finished in this slice. Refresh so a vanished runtime reference becomes LOST.
+        # Stop as soon as any requested task reaches a terminal state: the caller asked for the
+        # first result, so one task going LOST must not be masked by another still running.
+        tasks = _refresh_task_state(session, state, runtime, source_id=source_id)
+        watched = [t for t in tasks if t.id in waited_ids]
+        if any(t.status != BackgroundTaskStatus.RUNNING for t in watched) or not any(
+            t.status == BackgroundTaskStatus.RUNNING for t in watched
+        ):
+            return set(), False
+
+
 # ---------------------------------------------------------------------------
 # Provider class
 # ---------------------------------------------------------------------------
@@ -275,7 +373,8 @@ class BackgroundAgentsProvider(ContextProvider):
     This provider exposes the following tools to the agent:
 
     - ``background_agents_start_task`` — Start a background task on a named agent with text input.
-    - ``background_agents_wait_for_first_completion`` — Block until the first of the specified tasks completes.
+    - ``background_agents_wait_for_first_completion`` — Block until the first of the specified tasks
+      completes, or until ``wait_timeout_seconds`` elapses.
     - ``background_agents_get_task_results`` — Retrieve the text output of a completed background task.
     - ``background_agents_get_all_tasks`` — List all background tasks with their IDs, statuses, and descriptions.
     - ``background_agents_continue_task`` — Send follow-up input to a completed task's session to resume work.
@@ -297,6 +396,7 @@ def __init__(
         *,
         source_id: str = DEFAULT_BACKGROUND_AGENTS_SOURCE_ID,
         instructions: str | None = None,
+        wait_timeout_seconds: float | None = DEFAULT_BACKGROUND_AGENTS_WAIT_TIMEOUT_SECONDS,
     ) -> None:
         """Initialize the background agents provider.
 
@@ -312,13 +412,20 @@ def __init__(
             source_id: Unique source ID for serializable task state in session.
             instructions: Optional instruction override. May include ``{background_agents}``
                 placeholder which will be replaced with the agent listing.
+            wait_timeout_seconds: Default upper bound, in seconds, applied to
+                ``background_agents_wait_for_first_completion`` when the model does not supply its
+                own ``timeout_seconds``. Set to ``None`` to wait without a bound. Bounding the wait
+                keeps a child agent that never completes from suspending the parent's
+                function-calling loop indefinitely.
 
         Raises:
-            ValueError: If agents is empty, an agent has no name, or names are not unique.
+            ValueError: If agents is empty, an agent has no name, names are not unique, or
+                ``wait_timeout_seconds`` is not a positive number or ``None``.
         """
         super().__init__(source_id)
 
         self._agents = _validate_and_build_agent_dict(agents)
+        self._wait_timeout_seconds = _validate_wait_timeout(wait_timeout_seconds)
 
         # Build instructions with agent listing.
         base_instructions = instructions if instructions is not None else DEFAULT_BACKGROUND_AGENTS_INSTRUCTIONS
@@ -501,14 +608,32 @@ def background_agents_start_task(agent_name: str, input: str, description: str)
         background_agents_start_task._invoke_sync_on_event_loop = True  # pyright: ignore[reportPrivateUsage]
 
         @tool(name="background_agents_wait_for_first_completion", approval_mode="never_require")
-        async def background_agents_wait_for_first_completion(task_ids: list[int]) -> str:
-            """Block until the first of the specified background tasks completes. Returns the completed task's ID."""
+        async def background_agents_wait_for_first_completion(
+            task_ids: list[int],
+            # The provider default is bound as this parameter's default so an explicitly passed
+            # None means "wait without a bound" rather than being indistinguishable from omission.
+            timeout_seconds: float | None = self._wait_timeout_seconds,
+        ) -> str:
+            """Block until the first of the specified background tasks completes, or the timeout elapses.
+
+            Returns the completed task's ID, or the current status of each task if the wait timed out.
+            Pass timeout_seconds to override the provider's default wait timeout, or null to wait
+            without a bound.
+            """
             if runtime.closed:
                 return "Error: Session is being released; cannot wait for background tasks."
 
             if not task_ids:
                 return "Error: No task IDs provided."
 
+            # A model-supplied timeout is reported back as an error string rather than raised: this
+            # tool exists to keep the function-calling loop responsive, so a bad argument must not
+            # fail the tool invocation.
+            try:
+                timeout = _validate_wait_timeout(timeout_seconds)
+            except ValueError as exc:
+                return f"Error: {exc}"
+
             # Collect in-flight tasks matching the requested IDs.
             waitable: list[tuple[int, asyncio.Task[AgentResponse[Any]]]] = []
             for tid in task_ids:
@@ -528,12 +653,35 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st
                     )
                 return "Error: None of the specified task IDs correspond to running tasks."
 
-            # Wait for the first one to complete.
-            done, _ = await asyncio.wait(
-                [t for _, t in waitable],
-                return_when=asyncio.FIRST_COMPLETED,
+            # Wait for the first one to complete, bounded by the effective timeout. The wait is
+            # sliced so _refresh_task_state runs between slices: a task whose runtime reference has
+            # disappeared is promoted to LOST there, which ends the wait early instead of stalling
+            # for the full timeout.
+            done, timed_out = await _wait_first_completed(
+                session,
+                provider_state,
+                runtime,
+                waitable,
+                timeout=timeout,
+                source_id=source_id,
             )
 
+            if not done:
+                tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
+                statuses = ", ".join(f"task {t.id}: {t.status.value}" for t in tasks if t.id in task_ids)
+                if timed_out:
+                    return (
+                        f"Timed out after {timeout} seconds waiting for tasks {task_ids} to complete. "
+                        f"Current status: {statuses or 'unknown'}. "
+                        "The tasks may still be running; wait again or check their status."
+                    )
+                # The wait ended because a requested task reached a terminal state (for example
+                # LOST) without its asyncio task producing a result.
+                return (
+                    f"Stopped waiting for tasks {task_ids}: no running task remains to wait for. "
+                    f"Current status: {statuses or 'unknown'}."
+                )
+
             # Find which ID completed.
             completed_id: int | None = None
             for tid, task in waitable: