diff --git a/python/packages/core/agent_framework/_harness/_agent.py b/python/packages/core/agent_framework/_harness/_agent.py index 8ad199069f..c3cdea84b2 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: int, 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: int = 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: Maximum seconds the background-agent wait + tool blocks for a task to complete. Must be a positive integer. + Defaults to 300 seconds. On expiry, the tool returns normally and leaves the + background tasks running. Only used when ``background_agents`` is provided. 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 @@ -525,7 +537,8 @@ def create_harness_agent( Raises: ValueError: If max_context_window_tokens is provided and <= 0, or max_output_tokens is provided and <= 0, or max_output_tokens >= - max_context_window_tokens when both are provided. + max_context_window_tokens when both are provided, or + background_agents_wait_timeout_seconds is invalid when background agents are provided. """ if max_context_window_tokens is not None and max_context_window_tokens <= 0: raise ValueError("max_context_window_tokens must be positive.") @@ -601,6 +614,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..c5a1bf0029 100644 --- a/python/packages/core/agent_framework/_harness/_agent.pyi +++ b/python/packages/core/agent_framework/_harness/_agent.pyi @@ -14,6 +14,7 @@ from .._sessions import ContextProvider, HistoryProvider from .._skills import SkillsProvider from .._tools import ToolTypes from .._types import ChatOptions +from ._background_agents import DEFAULT_BACKGROUND_AGENTS_WAIT_TIMEOUT_SECONDS from ._file_access import AgentFileStore from ._loop import DEFAULT_MAX_ITERATIONS, NextMessageCallable, ShouldContinueCallable from ._mode import AgentModeProvider @@ -77,6 +78,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: int = DEFAULT_BACKGROUND_AGENTS_WAIT_TIMEOUT_SECONDS, 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..6cd70a3403 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -27,6 +27,7 @@ logger = logging.getLogger(__name__) DEFAULT_BACKGROUND_AGENTS_SOURCE_ID = "background_agents" +DEFAULT_BACKGROUND_AGENTS_WAIT_TIMEOUT_SECONDS = 300 DEFAULT_BACKGROUND_AGENTS_INSTRUCTIONS = """\ ## Background Agents @@ -275,7 +276,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`` — Wait until the first specified task completes or the + configured timeout expires. A timeout leaves the tasks running so the tool can be called again. - ``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 +299,7 @@ def __init__( *, source_id: str = DEFAULT_BACKGROUND_AGENTS_SOURCE_ID, instructions: str | None = None, + wait_timeout_seconds: int = DEFAULT_BACKGROUND_AGENTS_WAIT_TIMEOUT_SECONDS, ) -> None: """Initialize the background agents provider. @@ -312,13 +315,23 @@ 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 seconds the wait tool blocks for a task to complete. + Must be a positive integer. Defaults to 300 seconds. 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 integer. """ super().__init__(source_id) self._agents = _validate_and_build_agent_dict(agents) + if ( + isinstance(wait_timeout_seconds, bool) + or not isinstance(wait_timeout_seconds, int) + or wait_timeout_seconds <= 0 + ): + raise ValueError("wait_timeout_seconds must be a positive integer.") + 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 @@ -502,7 +515,11 @@ 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]) -> str: - """Block until the first of the specified background tasks completes. Returns the completed task's ID.""" + """Wait until the first of the specified tasks completes or the configured timeout expires. + + Returns the completed task's ID. On timeout, the tasks remain running and this tool + can be called again to continue waiting. + """ if runtime.closed: return "Error: Session is being released; cannot wait for background tasks." @@ -531,8 +548,14 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st # Wait for the first one to complete. done, _ = await asyncio.wait( [t for _, t in waitable], + timeout=self._wait_timeout_seconds, return_when=asyncio.FIRST_COMPLETED, ) + if not done: + return ( + f"No background task completed within {self._wait_timeout_seconds} seconds. " + "The tasks are still running; call this tool again if you wish to continue waiting." + ) # Find which ID completed. completed_id: int | None = None diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py index 62200518ee..7b14765868 100644 --- a/python/packages/core/agent_framework/_harness/_loop.py +++ b/python/packages/core/agent_framework/_harness/_loop.py @@ -34,7 +34,9 @@ from pydantic import BaseModel, Field from typing_extensions import Self -from .._agents import _LOOP_ITERATION_TOKEN_KEY # pyright: ignore[reportPrivateUsage] -- shared loop-turn marker, see _agents.py +from .._agents import ( + _LOOP_ITERATION_TOKEN_KEY, # pyright: ignore[reportPrivateUsage] -- shared loop-turn marker, see _agents.py +) from .._feature_stage import ExperimentalFeature, experimental from .._middleware import AgentContext, AgentMiddleware, MiddlewareTermination from .._sessions import SessionContext @@ -490,8 +492,7 @@ async def _fire_turn_scoped_after_providers( if response is None or run_after is None: return if not any( - getattr(provider, "after_run_once_per_turn", False) - for provider in getattr(agent, "context_providers", []) + getattr(provider, "after_run_once_per_turn", False) for provider in getattr(agent, "context_providers", []) ): return session_context = SessionContext( diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 09bd592e92..353a1a2efb 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -713,11 +713,7 @@ def _process_request_info_event( Note: Text requests use the function-call envelope so callers can reply with a matching function result. """ - if ( - isinstance(event.data, Content) - and event.data.user_input_request - and event.data.type != "text" - ): + if isinstance(event.data, Content) and event.data.user_input_request and event.data.type != "text": # Preserve specialized requests that callers already understand how to present. return event.data diff --git a/python/packages/core/tests/core/test_harness_agent.py b/python/packages/core/tests/core/test_harness_agent.py index b74902116e..04f6488e7c 100644 --- a/python/packages/core/tests/core/test_harness_agent.py +++ b/python/packages/core/tests/core/test_harness_agent.py @@ -841,6 +841,24 @@ def test_create_harness_agent_background_agents_custom_instructions() -> None: assert "Helper" in bg_providers[0]._instructions +def test_create_harness_agent_background_agents_custom_wait_timeout() -> None: + """Custom wait timeout should be passed to BackgroundAgentsProvider.""" + from agent_framework._harness._background_agents import BackgroundAgentsProvider + + bg_agent = _FakeBackgroundAgent("Helper", "A helper agent") + agent = create_harness_agent( + client=_FakeChatClient(), + max_context_window_tokens=128_000, + max_output_tokens=16_384, + disable_web_search=True, + background_agents=[bg_agent], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + background_agents_wait_timeout_seconds=12, + ) + providers = agent.context_providers or [] + bg_provider = next(p for p in providers if isinstance(p, BackgroundAgentsProvider)) + assert bg_provider._wait_timeout_seconds == 12 + + def test_create_harness_agent_empty_background_agents_list() -> None: """An empty background_agents list should NOT add a BackgroundAgentsProvider.""" from agent_framework._harness._background_agents import BackgroundAgentsProvider 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..88e776a0c8 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -127,6 +127,22 @@ def test_constructor_custom_source_id() -> None: assert provider.source_id == "custom_bg" +def test_constructor_uses_default_wait_timeout() -> None: + """Should use a bounded five-minute wait by default.""" + provider = _make_provider(_FakeAgent("Worker")) + assert provider._wait_timeout_seconds == 300 + + +@pytest.mark.parametrize("wait_timeout_seconds", [0, -1, True, 1.5, float("inf"), float("nan")]) +def test_constructor_rejects_invalid_wait_timeout(wait_timeout_seconds: object) -> None: + """Should reject wait timeouts that are not positive integers.""" + with pytest.raises(ValueError, match="positive integer"): + BackgroundAgentsProvider( + [_FakeAgent("Worker")], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=wait_timeout_seconds, # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + ) + + # --- Tool Injection Tests --- @@ -146,6 +162,18 @@ async def test_before_run_injects_six_tools() -> None: assert set(tools.keys()) == expected_names +async def test_wait_timeout_is_not_model_settable() -> None: + """The model-facing wait tool should only accept task IDs.""" + provider = _make_provider(_FakeAgent("Worker")) + tools = await _get_tools(provider, _make_session()) + wait_tool = tools["background_agents_wait_for_first_completion"] + wait_parameters = wait_tool.parameters() + assert set(wait_parameters["properties"]) == {"task_ids"} + assert "configured timeout expires" in wait_tool.description + assert "tasks remain running" in wait_tool.description + assert "called again" in wait_tool.description + + async def test_before_run_injects_instructions() -> None: """before_run should inject instructions mentioning agent names.""" provider = _make_provider(_FakeAgent("ResearchBot", "Does research")) @@ -293,6 +321,55 @@ async def test_wait_no_running_tasks() -> None: assert "Error" in result or "not running" in result.lower() +async def test_wait_timeout_returns_without_stopping_task(monkeypatch: pytest.MonkeyPatch) -> None: + """Should return normally on timeout and leave the child task running.""" + observed_timeout: float | None = None + + async def _return_timeout( + tasks: Any, + *, + timeout: float | None = None, + return_when: Any = asyncio.ALL_COMPLETED, + ) -> tuple[set[Any], set[Any]]: + nonlocal observed_timeout + observed_timeout = timeout + return set(), set(tasks) + + monkeypatch.setattr(asyncio, "wait", _return_timeout) + provider = BackgroundAgentsProvider( + [_FakeAgent("Slow", delay=10.0)], # type: ignore[list-item] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + wait_timeout_seconds=7, + ) + 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", + ) + + try: + result = await _invoke_tool( + tools["background_agents_wait_for_first_completion"], + task_ids=[1], + ) + assert result == ( + "No background task completed within 7 seconds. " + "The tasks are still running; call this tool again if you wish to continue waiting." + ) + assert observed_timeout == 7 + + runtime = provider._get_runtime(session) + assert not runtime.in_flight_tasks[1].done() + + task_result = await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + assert "still running" in task_result.lower() + finally: + await provider.release_session(session) + + # --- Get Task Results Tests ---