Python: fix: bound background_agents_wait_for_first_completion with a timeout

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

@@ -24,6 +24,7 @@
 from .._types import AgentResponse, Message
 
 DEFAULT_BACKGROUND_AGENTS_SOURCE_ID = "background_agents"
+_DEFAULT_WAIT_TIMEOUT: Any = object()
 
 DEFAULT_BACKGROUND_AGENTS_INSTRUCTIONS = """\
 ## Background Agents
@@ -249,7 +250,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, bounded by the provider's wait timeout (or the ``timeout_seconds`` argument).
     - ``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.
@@ -271,6 +273,7 @@ def __init__(
         *,
         source_id: str = DEFAULT_BACKGROUND_AGENTS_SOURCE_ID,
         instructions: str | None = None,
+        wait_timeout_seconds: float | None = 300.0,
     ) -> None:
         """Initialize the background agents provider.
 
@@ -286,13 +289,21 @@ 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: Maximum number of seconds
+                ``background_agents_wait_for_first_completion`` blocks before returning control to
+                the model with the current task statuses. ``None`` waits indefinitely. Defaults to
+                300 seconds so a child that never completes cannot suspend the parent's run forever.
 
         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 negative.
         """
         super().__init__(source_id)
 
         self._agents = _validate_and_build_agent_dict(agents)
+        if wait_timeout_seconds is not None and wait_timeout_seconds < 0:
+            raise ValueError("wait_timeout_seconds must be non-negative.")
+        self._wait_timeout_seconds = wait_timeout_seconds
 
         # Build instructions with agent listing.
         base_instructions = instructions if instructions is not None else DEFAULT_BACKGROUND_AGENTS_INSTRUCTIONS
@@ -363,8 +374,20 @@ 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],
+            timeout_seconds: float | None = _DEFAULT_WAIT_TIMEOUT,  # type: ignore[assignment]
+        ) -> str:
+            """Block until the first of the specified background tasks completes, up to a timeout.
+
+            Returns the completed task's ID, or the current task statuses if the timeout elapses
+            first (in which case call this tool again or check task results to proceed).
+
+            Args:
+                task_ids: IDs of background tasks to wait on.
+                timeout_seconds: Maximum time to wait in seconds. When omitted, the provider's
+                    ``wait_timeout_seconds`` is used. Pass ``None`` to wait indefinitely.
+            """
             if not task_ids:
                 return "Error: No task IDs provided."
 
@@ -387,12 +410,30 @@ 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.
+            # Wait for the first one to complete, bounded so a child that never completes
+            # cannot suspend the calling agent's run indefinitely.
+            effective_timeout = (
+                self._wait_timeout_seconds if timeout_seconds is _DEFAULT_WAIT_TIMEOUT else timeout_seconds
+            )
+            if effective_timeout is not None and effective_timeout < 0:
+                return "Error: timeout_seconds must be non-negative."
             done, _ = await asyncio.wait(
                 [t for _, t in waitable],
                 return_when=asyncio.FIRST_COMPLETED,
+                timeout=effective_timeout,
             )
 
+            if not done:
+                # asyncio.wait with timeout=None blocks indefinitely, so this branch is only
+                # reachable when a numeric timeout was in effect.
+                # Refresh state so a task whose runtime disappeared is surfaced as LOST,
+                # then hand control back to the model with an honest status report.
+                tasks = _refresh_task_state(session, provider_state, runtime, source_id=source_id)
+                status_lines = [f"- Task {t.id} [{t.status.value}]" for t in tasks if t.id in task_ids]
+                status_text = "\n".join(status_lines) if status_lines else "No matching tasks found."
+                timeout_label = f"{effective_timeout:g}" if effective_timeout is not None else "unlimited"
+                return f"No task completed within {timeout_label} seconds. Current task statuses:\n{status_text}"
+
             # Find which ID completed.
             completed_id: int | None = None
             for tid, task in waitable:

python/packages/core/tests/core/test_harness_background_agents.py

@@ -65,6 +65,23 @@ async def run(
         return AgentResponse(messages=[Message(role="assistant", contents=[self._response_text])])
 
 
+class _HangingAgent:
+    """Agent stub whose run never completes, simulating a stuck child task."""
+
+    name = "Hanger"
+    description = None
+
+    def create_session(self, *, session_id: str | None = None) -> AgentSession:
+        return AgentSession(session_id=session_id)
+
+    async def run(
+        self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any
+    ) -> AgentResponse[Any]:
+        del messages, stream, session, kwargs
+        await asyncio.Event().wait()
+        raise AssertionError("unreachable")
+
+
 def _make_provider(*agents: _FakeAgent) -> BackgroundAgentsProvider:
     """Create a provider with given agents."""
     return BackgroundAgentsProvider(agents)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
@@ -120,6 +137,15 @@ def test_constructor_valid_agents() -> None:
     assert provider.source_id == "background_agents"
 
 
+def test_constructor_rejects_negative_wait_timeout() -> None:
+    """Should fail fast for invalid provider wait timeout configuration."""
+    with pytest.raises(ValueError, match="wait_timeout_seconds must be non-negative"):
+        BackgroundAgentsProvider(
+            [_FakeAgent("Worker")],  # type: ignore[list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            wait_timeout_seconds=-1,
+        )
+
+
 def test_constructor_custom_source_id() -> None:
     """Should accept custom source_id."""
     provider = BackgroundAgentsProvider([_FakeAgent("Agent1")], source_id="custom_bg")  # type: ignore[list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
@@ -292,6 +318,134 @@ async def test_wait_no_running_tasks() -> None:
     assert "Error" in result or "not running" in result.lower()
 
 
+async def test_wait_for_first_completion_timeout() -> None:
+    """Should return current statuses instead of hanging when no task completes within the timeout."""
+    provider = _make_provider(_HangingAgent())  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Hanger",
+        input="go",
+        description="never finishes",
+    )
+    try:
+        result = await _invoke_tool(
+            tools["background_agents_wait_for_first_completion"],
+            task_ids=[1],
+            timeout_seconds=0.05,
+        )
+        assert "no task completed within" in result.lower()
+        assert "running" in result.lower()
+    finally:
+        runtime = provider._get_runtime(session)
+        for task in list(runtime.in_flight_tasks.values()):
+            task.cancel()
+        await asyncio.gather(*runtime.in_flight_tasks.values(), return_exceptions=True)
+
+
+async def test_wait_timeout_uses_provider_default() -> None:
+    """Should apply the provider's wait_timeout_seconds when the tool timeout is omitted."""
+    provider = BackgroundAgentsProvider(
+        [_HangingAgent()],  # type: ignore[list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        wait_timeout_seconds=0.05,
+    )
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Hanger",
+        input="go",
+        description="never finishes",
+    )
+    try:
+        result = await _invoke_tool(
+            tools["background_agents_wait_for_first_completion"],
+            task_ids=[1],
+        )
+        assert "no task completed within" in result.lower()
+        assert "running" in result.lower()
+    finally:
+        runtime = provider._get_runtime(session)
+        for task in list(runtime.in_flight_tasks.values()):
+            task.cancel()
+        await asyncio.gather(*runtime.in_flight_tasks.values(), return_exceptions=True)
+
+
+async def test_wait_for_first_completion_with_explicit_timeout() -> None:
+    """Should still return the completed task when it finishes before the timeout elapses."""
+    provider = _make_provider(_FakeAgent("Fast", response_text="fast result", delay=0.01))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Fast",
+        input="go",
+        description="fast task",
+    )
+    result = await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[1],
+        timeout_seconds=5.0,
+    )
+    assert "finished" in result.lower()
+    assert "completed" in result.lower()
+
+
+async def test_wait_for_first_completion_explicit_none_waits_indefinitely() -> None:
+    """Should let explicit None override the provider default with an unbounded wait."""
+    provider = BackgroundAgentsProvider(
+        [_FakeAgent("Slow", response_text="slow result", delay=0.01)],  # type: ignore[list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        wait_timeout_seconds=0,
+    )
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Slow",
+        input="go",
+        description="slow task",
+    )
+    result = await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[1],
+        timeout_seconds=None,
+    )
+    assert "finished" in result.lower()
+    assert "completed" in result.lower()
+
+
+async def test_wait_for_first_completion_rejects_negative_tool_timeout() -> None:
+    """Should return an error message instead of raising for a negative tool timeout."""
+    provider = _make_provider(_HangingAgent())  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Hanger",
+        input="go",
+        description="never finishes",
+    )
+    try:
+        result = await _invoke_tool(
+            tools["background_agents_wait_for_first_completion"],
+            task_ids=[1],
+            timeout_seconds=-1,
+        )
+        assert "error" in result.lower()
+        assert "non-negative" in result.lower()
+    finally:
+        runtime = provider._get_runtime(session)
+        for task in list(runtime.in_flight_tasks.values()):
+            task.cancel()
+        await asyncio.gather(*runtime.in_flight_tasks.values(), return_exceptions=True)
+
+
 # --- Get Task Results Tests ---