From 963bdee8d03a44a89f1ba924c44237ded1ec551b Mon Sep 17 00:00:00 2001 From: manojmeruva Date: Thu, 27 Aug 2026 14:40:44 +0530 Subject: [PATCH 1/2] Python: bound background_agents_wait_for_first_completion with a timeout background_agents_wait_for_first_completion called asyncio.wait() with no timeout, so a child agent that never completed suspended the parent's function-calling loop indefinitely. Because the tool holds direct asyncio.Task references, _refresh_task_state never ran while the wait was parked, so a task whose runtime reference had disappeared was never promoted to LOST, and the model could not poll task status to recover. Add a provider-level wait_timeout_seconds default of 300 seconds and an optional per-call timeout_seconds override; either may be None to preserve the previous unbounded behavior. The wait now runs in bounded slices, refreshing task state between them so a LOST task ends the wait early rather than stalling for the full timeout. On timeout the tool returns current task statuses to the model instead of raising. Fixes #7454 --- .../core/agent_framework/_harness/_agent.py | 17 +- .../_harness/_background_agents.py | 142 ++++++++++++- .../core/test_harness_background_agents.py | 190 ++++++++++++++++++ 3 files changed, 339 insertions(+), 10 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_agent.py b/python/packages/core/agent_framework/_harness/_agent.py index 8ad199069f4..cd75f420bb1 100644 --- a/python/packages/core/agent_framework/_harness/_agent.py +++ b/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, ) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 3ec3464d2cd..7c66f1e86f5 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -28,6 +28,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 +46,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 +165,30 @@ 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}.") + if timeout_seconds != timeout_seconds: # NaN never compares greater than 0. + raise ValueError("Background agent wait timeout must not be NaN.") + if timeout_seconds <= 0: + raise ValueError(f"Background agent wait timeout must be greater than 0; got {timeout_seconds!r}.") + return float(timeout_seconds) + + def _validate_and_build_agent_dict(agents: Sequence[SupportsAgentRun]) -> dict[str, SupportsAgentRun]: """Validate agents and build a case-insensitive lookup dict. @@ -259,6 +296,53 @@ 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, +) -> set[asyncio.Task[AgentResponse[Any]]]: + """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: + The set of tasks that completed, which is empty when the timeout elapsed or every awaited + task stopped being tracked. + """ + loop = asyncio.get_running_loop() + deadline = None if timeout is None else loop.time() + timeout + pending_tasks = [task for _, task in waitable] + waited_ids = [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() + slice_timeout = min(slice_timeout, remaining) + + done, _ = await asyncio.wait( + pending_tasks, + timeout=slice_timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + if done: + return done + + # Nothing finished in this slice. Refresh so a vanished runtime reference becomes LOST, + # and stop waiting if none of the requested tasks is still considered running. + tasks = _refresh_task_state(session, state, runtime, source_id=source_id) + if not any(t.id in waited_ids and t.status == BackgroundTaskStatus.RUNNING for t in tasks): + return set() + + # --------------------------------------------------------------------------- # Provider class # --------------------------------------------------------------------------- @@ -275,7 +359,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 +382,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 +398,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 +594,31 @@ 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, 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. + """ 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. + if timeout_seconds is None: + timeout = self._wait_timeout_seconds + else: + 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 +638,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. - 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 = 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) + 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." + ) + # 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 199fb5660ad..f34ed8aa828 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -293,6 +293,196 @@ async def test_wait_no_running_tasks() -> None: assert "Error" in result or "not running" in result.lower() +# --- Wait Timeout Tests --- + + +async def test_wait_times_out_instead_of_blocking_forever() -> None: + """A child that never completes must not suspend the caller indefinitely.""" + provider = BackgroundAgentsProvider( + [_FakeAgent("Hanger")], # 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) + + runtime = provider._get_runtime(session) + unblock = asyncio.Event() + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Hanger", + input="go", + description="hangs forever", + ) + # Replace the real task with one that never finishes on its own. + runtime.in_flight_tasks[1].cancel() + hanging = asyncio.create_task(unblock.wait()) + runtime.in_flight_tasks[1] = hanging # type: ignore[assignment] # pyrefly: ignore[bad-assignment] # ty: ignore[invalid-assignment] + + result = await asyncio.wait_for( + _invoke_tool(tools["background_agents_wait_for_first_completion"], task_ids=[1]), + timeout=5.0, + ) + + assert "timed out" in result.lower() + assert "running" in result.lower() + + unblock.set() + with suppress(asyncio.CancelledError): + await asyncio.wait_for(hanging, timeout=1.0) + + +async def test_wait_timeout_seconds_argument_overrides_provider_default() -> None: + """A model-supplied timeout_seconds should bound the wait even when the default is unbounded.""" + provider = BackgroundAgentsProvider( + [_FakeAgent("Hanger")], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=None, + ) + session = _make_session() + tools = await _get_tools(provider, session) + + runtime = provider._get_runtime(session) + unblock = asyncio.Event() + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Hanger", + input="go", + description="hangs forever", + ) + runtime.in_flight_tasks[1].cancel() + hanging = asyncio.create_task(unblock.wait()) + runtime.in_flight_tasks[1] = hanging # type: ignore[assignment] # pyrefly: ignore[bad-assignment] # ty: ignore[invalid-assignment] + + result = await asyncio.wait_for( + _invoke_tool( + tools["background_agents_wait_for_first_completion"], + task_ids=[1], + timeout_seconds=0.05, + ), + timeout=5.0, + ) + + assert "timed out" in result.lower() + + unblock.set() + with suppress(asyncio.CancelledError): + await asyncio.wait_for(hanging, timeout=1.0) + + +async def test_wait_returns_normally_when_task_completes_before_timeout() -> None: + """A generous timeout must not disturb the normal completion path.""" + provider = BackgroundAgentsProvider( + [_FakeAgent("Fast", response_text="fast result", delay=0.01)], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=30.0, + ) + 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], + ) + + assert "finished" in result.lower() + assert "completed" in result.lower() + assert "timed out" not in result.lower() + + +async def test_wait_returns_early_when_task_becomes_lost() -> None: + """A task whose runtime reference disappears should end the wait before the timeout.""" + provider = BackgroundAgentsProvider( + [_FakeAgent("Worker")], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=30.0, + ) + session = _make_session() + tools = await _get_tools(provider, session) + + runtime = provider._get_runtime(session) + unblock = asyncio.Event() + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Worker", + input="go", + description="loses its reference", + ) + runtime.in_flight_tasks[1].cancel() + hanging = asyncio.create_task(unblock.wait()) + runtime.in_flight_tasks[1] = hanging # type: ignore[assignment] # pyrefly: ignore[bad-assignment] # ty: ignore[invalid-assignment] + + async def _drop_reference() -> None: + # Simulate the runtime losing track of the task while the wait is parked. + await asyncio.sleep(0.05) + runtime.in_flight_tasks.pop(1, None) + + dropper = asyncio.create_task(_drop_reference()) + start = asyncio.get_running_loop().time() + + result = await asyncio.wait_for( + _invoke_tool(tools["background_agents_wait_for_first_completion"], task_ids=[1]), + timeout=30.0, + ) + elapsed = asyncio.get_running_loop().time() - start + + # Returned well before the 30s timeout because the task was promoted to LOST. + assert elapsed < 20.0 + assert "lost" in result.lower() + + await dropper + unblock.set() + with suppress(asyncio.CancelledError): + await asyncio.wait_for(hanging, timeout=1.0) + + +@pytest.mark.parametrize("bad_timeout", [0, -1, -0.5, float("nan"), True]) +def test_constructor_rejects_invalid_wait_timeout(bad_timeout: Any) -> None: + """Provider-level timeout misconfiguration should fail fast with ValueError.""" + with pytest.raises(ValueError, match="wait timeout"): + BackgroundAgentsProvider( + [_FakeAgent("Worker")], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=bad_timeout, + ) + + +def test_constructor_accepts_none_wait_timeout() -> None: + """``None`` should be accepted and preserve unbounded waiting.""" + provider = BackgroundAgentsProvider( + [_FakeAgent("Worker")], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=None, + ) + assert provider._wait_timeout_seconds is None + + +@pytest.mark.parametrize("bad_timeout", [0, -1, -0.5]) +async def test_wait_returns_error_for_invalid_timeout_argument(bad_timeout: float) -> None: + """A bad model-supplied timeout must return an error string, not raise.""" + provider = _make_provider(_FakeAgent("Worker", delay=0.01)) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Worker", + input="go", + description="task", + ) + result = await _invoke_tool( + tools["background_agents_wait_for_first_completion"], + task_ids=[1], + timeout_seconds=bad_timeout, + ) + + assert "Error" in result + assert "wait timeout" in result + + # --- Get Task Results Tests --- From 9b5e35c6d5ff8489ec324c33a31a34d312d745a1 Mon Sep 17 00:00:00 2001 From: manojmeruva Date: Thu, 27 Aug 2026 17:45:02 +0530 Subject: [PATCH 2/2] Address Copilot review feedback on background-agent wait timeout - Reject non-finite timeouts. Positive infinity previously passed validation and produced an infinite deadline, restoring the unbounded wait this change exists to prevent. Convert to float first so a very large int raises the documented ValueError instead of OverflowError. - End the wait as soon as any requested task reaches a terminal state. When waiting on several IDs, one task becoming LOST was previously masked by another still running, parking the caller until the full timeout. - Distinguish a terminal-state wakeup from deadline expiry. An early return no longer reports "Timed out after N seconds" (or "after None seconds" in unbounded mode) when the deadline had not elapsed. - Bind the provider default as the tool parameter's default so an explicit timeout_seconds=None waits without a bound instead of being indistinguishable from omission. - Declare background_agents_wait_timeout_seconds in _agent.pyi so type checkers and editors accept the new keyword. --- .../core/agent_framework/_harness/_agent.pyi | 1 + .../_harness/_background_agents.py | 76 +++++++++------ .../core/test_harness_background_agents.py | 92 ++++++++++++++++++- 3 files changed, 139 insertions(+), 30 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_agent.pyi b/python/packages/core/agent_framework/_harness/_agent.pyi index c9f4f318007..fe6562317a6 100644 --- a/python/packages/core/agent_framework/_harness/_agent.pyi +++ b/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, diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 7c66f1e86f5..c476a319697 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/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 @@ -182,11 +183,19 @@ def _validate_wait_timeout(timeout_seconds: float | None) -> float | 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}.") - if timeout_seconds != timeout_seconds: # NaN never compares greater than 0. - raise ValueError("Background agent wait timeout must not be NaN.") - if timeout_seconds <= 0: + # 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 float(timeout_seconds) + return timeout def _validate_and_build_agent_dict(agents: Sequence[SupportsAgentRun]) -> dict[str, SupportsAgentRun]: @@ -304,7 +313,7 @@ async def _wait_first_completed( *, timeout: float | None, source_id: str, -) -> set[asyncio.Task[AgentResponse[Any]]]: +) -> 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 @@ -312,20 +321,21 @@ async def _wait_first_completed( to ``LOST``, ending the wait early rather than stalling for the whole timeout. Returns: - The set of tasks that completed, which is empty when the timeout elapsed or every awaited - task stopped being tracked. + 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 = [task_id for task_id, _ 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() + return set(), True slice_timeout = min(slice_timeout, remaining) done, _ = await asyncio.wait( @@ -334,13 +344,17 @@ async def _wait_first_completed( return_when=asyncio.FIRST_COMPLETED, ) if done: - return done + return done, False - # Nothing finished in this slice. Refresh so a vanished runtime reference becomes LOST, - # and stop waiting if none of the requested tasks is still considered running. + # 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) - if not any(t.id in waited_ids and t.status == BackgroundTaskStatus.RUNNING for t in tasks): - return set() + 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 # --------------------------------------------------------------------------- @@ -595,12 +609,16 @@ def background_agents_start_task(agent_name: str, input: str, description: str) @tool(name="background_agents_wait_for_first_completion", approval_mode="never_require") async def background_agents_wait_for_first_completion( - task_ids: list[int], timeout_seconds: float | None = None + 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. + 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." @@ -611,13 +629,10 @@ async def background_agents_wait_for_first_completion( # 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. - if timeout_seconds is None: - timeout = self._wait_timeout_seconds - else: - try: - timeout = _validate_wait_timeout(timeout_seconds) - except ValueError as exc: - return f"Error: {exc}" + 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]]]] = [] @@ -642,7 +657,7 @@ async def background_agents_wait_for_first_completion( # 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 = await _wait_first_completed( + done, timed_out = await _wait_first_completed( session, provider_state, runtime, @@ -654,10 +669,17 @@ async def background_agents_wait_for_first_completion( 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"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." + f"Stopped waiting for tasks {task_ids}: no running task remains to wait for. " + f"Current status: {statuses or 'unknown'}." ) # Find which ID completed. 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 f34ed8aa828..172a62cc534 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -431,9 +431,11 @@ async def _drop_reference() -> None: ) elapsed = asyncio.get_running_loop().time() - start - # Returned well before the 30s timeout because the task was promoted to LOST. + # Returned well before the 30s timeout because the task was promoted to LOST, and the message + # must not claim the deadline expired. assert elapsed < 20.0 assert "lost" in result.lower() + assert "timed out" not in result.lower() await dropper unblock.set() @@ -441,9 +443,16 @@ async def _drop_reference() -> None: await asyncio.wait_for(hanging, timeout=1.0) -@pytest.mark.parametrize("bad_timeout", [0, -1, -0.5, float("nan"), True]) +@pytest.mark.parametrize( + "bad_timeout", + [0, -1, -0.5, float("nan"), True, float("inf"), float("-inf"), 10**400, "30"], +) def test_constructor_rejects_invalid_wait_timeout(bad_timeout: Any) -> None: - """Provider-level timeout misconfiguration should fail fast with ValueError.""" + """Provider-level timeout misconfiguration should fail fast with ValueError. + + Infinity is rejected because an infinite deadline would restore the unbounded wait this + timeout exists to prevent; callers who want that must pass ``None`` explicitly. + """ with pytest.raises(ValueError, match="wait timeout"): BackgroundAgentsProvider( [_FakeAgent("Worker")], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] @@ -451,6 +460,83 @@ def test_constructor_rejects_invalid_wait_timeout(bad_timeout: Any) -> None: ) +async def test_wait_explicit_none_timeout_argument_waits_unbounded() -> None: + """An explicit timeout_seconds=None must be unbounded, not fall back to the provider default.""" + provider = BackgroundAgentsProvider( + [_FakeAgent("Slow", response_text="slow result", delay=0.2)], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=0.01, + ) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Slow", + input="go", + description="slower than the provider default", + ) + # The provider default of 0.01s would time out; an explicit None must wait for the result. + result = await _invoke_tool( + tools["background_agents_wait_for_first_completion"], + task_ids=[1], + timeout_seconds=None, + ) + + assert "completed" in result.lower() + assert "timed out" not in result.lower() + + +async def test_wait_returns_when_one_of_several_tasks_is_lost() -> None: + """One requested task going LOST should end the wait even while another is still running.""" + provider = BackgroundAgentsProvider( + [_FakeAgent("Worker")], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=30.0, + ) + session = _make_session() + tools = await _get_tools(provider, session) + + runtime = provider._get_runtime(session) + unblock = asyncio.Event() + + for index in (1, 2): + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Worker", + input="go", + description=f"task {index}", + ) + runtime.in_flight_tasks[index].cancel() + waiter = asyncio.create_task(unblock.wait()) + runtime.in_flight_tasks[index] = waiter # type: ignore[assignment] # pyrefly: ignore[bad-assignment] # ty: ignore[invalid-assignment] + + hanging = [runtime.in_flight_tasks[1], runtime.in_flight_tasks[2]] + + async def _drop_one_reference() -> None: + # Task 1 loses its reference while task 2 keeps running. + await asyncio.sleep(0.05) + runtime.in_flight_tasks.pop(1, None) + + dropper = asyncio.create_task(_drop_one_reference()) + start = asyncio.get_running_loop().time() + + result = await asyncio.wait_for( + _invoke_tool(tools["background_agents_wait_for_first_completion"], task_ids=[1, 2]), + timeout=30.0, + ) + elapsed = asyncio.get_running_loop().time() - start + + # Returned on the lost task rather than parking until the 30s deadline. + assert elapsed < 20.0 + assert "lost" in result.lower() + assert "timed out" not in result.lower() + + await dropper + unblock.set() + for task in hanging: + with suppress(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1.0) + + def test_constructor_accepts_none_wait_timeout() -> None: """``None`` should be accepted and preserve unbounded waiting.""" provider = BackgroundAgentsProvider(