diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index c1c70b64a3..f51b79531d 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -249,7 +249,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 +272,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,6 +288,10 @@ 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. @@ -293,6 +299,7 @@ def __init__( super().__init__(source_id) self._agents = _validate_and_build_agent_dict(agents) + 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 +370,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 = None, + ) -> 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. + """ if not task_ids: return "Error: No task IDs provided." @@ -387,12 +406,28 @@ 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 None 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: diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index 98e6fa2a67..bcf49553f1 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/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] @@ -292,6 +309,139 @@ 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_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) + + +async def test_wait_for_first_completion_rejects_negative_provider_timeout() -> None: + """Should return an error message instead of raising for a negative provider timeout.""" + provider = BackgroundAgentsProvider( + [_HangingAgent()], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=-1, + ) + 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 "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 ---