diff --git a/python/packages/core/agent_framework/_harness/_agent.py b/python/packages/core/agent_framework/_harness/_agent.py index 8ad199069f..cd75f420bb 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/_agent.pyi b/python/packages/core/agent_framework/_harness/_agent.pyi index c9f4f31800..fe6562317a 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 3ec3464d2c..c476a31969 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 @@ -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: 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 199fb5660a..172a62cc53 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,282 @@ 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, 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() + with suppress(asyncio.CancelledError): + await asyncio.wait_for(hanging, timeout=1.0) + + +@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. + + 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] + wait_timeout_seconds=bad_timeout, + ) + + +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( + [_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 ---