Python: Add timeout for wait-for-first-completion

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,
     )

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,

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

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(

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
 

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

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 ---