diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 6cd70a34037..631a8bd7705 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -229,7 +229,9 @@ def _finalize_task( else: task_info.status = BackgroundTaskStatus.COMPLETED task_info.result_text = completed_task.result().text - runtime.in_flight_tasks.pop(task_info.id, None) + # A stale waiter may observe a task ID already reused by a continuation. + if runtime.in_flight_tasks.get(task_info.id) is completed_task: + runtime.in_flight_tasks.pop(task_info.id, None) def _refresh_task_state( @@ -559,18 +561,21 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st # Find which ID completed. completed_id: int | None = None + completed_task: asyncio.Task[AgentResponse[Any]] | None = None for tid, task in waitable: if task in done: completed_id = tid + completed_task = task break # Finalize the completed task. tasks = _get_tasks(provider_state) task_info = next((t for t in tasks if t.id == completed_id), None) - if task_info is not None and completed_id is not None: - completed_task = runtime.in_flight_tasks.get(completed_id) - if completed_task is not None: - _finalize_task(task_info, completed_task, runtime) + if task_info is not None and completed_task is not None: + # Report the waited-on execution without overwriting a newer one. + is_current_task = runtime.in_flight_tasks.get(task_info.id) is completed_task + _finalize_task(task_info, completed_task, runtime) + if is_current_task: _save_tasks(provider_state, tasks) _save_provider_state(session, provider_state, source_id=source_id) 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 88e776a0c81..7e14c074d6f 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -275,6 +275,128 @@ async def test_get_all_tasks_shows_tasks() -> None: # --- Wait for Completion Tests --- +@pytest.mark.parametrize("continuation_state", ["running", "completed", "finalized", "failed"]) +@pytest.mark.parametrize("first_outcome", ["completed", "failed", "cancelled"]) +async def test_wait_does_not_finalize_a_continued_task( + monkeypatch: pytest.MonkeyPatch, continuation_state: str, first_outcome: str +) -> None: + """A waiter for an earlier execution must not consume a continuation's task.""" + first_release = asyncio.Event() + second_started = asyncio.Event() + second_release = asyncio.Event() + wait_started = asyncio.Event() + wait_completed = asyncio.Event() + release_waiter = asyncio.Event() + + class ContinuingAgent(_FakeAgent): + async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse[Any]: + if messages == "first": + await first_release.wait() + if first_outcome == "failed": + raise RuntimeError("first failed") + if first_outcome == "cancelled": + raise asyncio.CancelledError + return AgentResponse(messages=[Message(role="assistant", contents=["first result"])]) + second_started.set() + await second_release.wait() + if continuation_state == "failed": + raise RuntimeError("continuation failed") + return AgentResponse(messages=[Message(role="assistant", contents=["second result"])]) + + real_wait = asyncio.wait + + async def paused_wait(*args: Any, **kwargs: Any) -> Any: + wait_started.set() + result = await real_wait(*args, **kwargs) + wait_completed.set() + await release_waiter.wait() + return result + + monkeypatch.setattr(asyncio, "wait", paused_wait) + provider = _make_provider(ContinuingAgent("Worker")) + session = _make_session() + tools = await _get_tools(provider, session) + await _invoke_tool( + tools["background_agents_start_task"], agent_name="Worker", input="first", description="first execution" + ) + waiter = asyncio.create_task(_invoke_tool(tools["background_agents_wait_for_first_completion"], task_ids=[1])) + try: + await wait_started.wait() + first_release.set() + await wait_completed.wait() + first_result = { + "completed": "first result", + "failed": "Task failed: first failed", + "cancelled": "Task failed: Task was canceled.", + }[first_outcome] + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == first_result + assert "continued" in await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="second") + await second_started.wait() + continuation = provider._runtime[session.session_id].in_flight_tasks[1] + if continuation_state != "running": + second_release.set() + await asyncio.gather(continuation, return_exceptions=True) + if continuation_state == "finalized": + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "second result" + expected_state = session.to_dict() + release_waiter.set() + expected_status = "completed" if first_outcome == "completed" else "failed" + assert await waiter == f"Task 1 finished with status: {expected_status}." + assert session.to_dict() == expected_state + if continuation_state != "finalized": + assert provider._runtime[session.session_id].in_flight_tasks[1] is continuation + second_release.set() + await asyncio.gather(continuation, return_exceptions=True) + expected_result = "Task failed: continuation failed" if continuation_state == "failed" else "second result" + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == expected_result + finally: + first_release.set() + second_release.set() + release_waiter.set() + await asyncio.gather(waiter, return_exceptions=True) + await provider.release_session(session) + + +async def test_concurrent_waiters_report_the_same_completed_execution(monkeypatch: pytest.MonkeyPatch) -> None: + """Both waiters must report completion even when one has already finalized it.""" + agent_release = asyncio.Event() + both_waiting = asyncio.Event() + wait_count = 0 + real_wait = asyncio.wait + + class WaitingAgent(_FakeAgent): + async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse[Any]: + await agent_release.wait() + return AgentResponse(messages=[Message(role="assistant", contents=["shared result"])]) + + async def counted_wait(*args: Any, **kwargs: Any) -> Any: + nonlocal wait_count + wait_count += 1 + if wait_count == 2: + both_waiting.set() + return await real_wait(*args, **kwargs) + + monkeypatch.setattr(asyncio, "wait", counted_wait) + provider = _make_provider(WaitingAgent("Worker")) + session = _make_session() + tools = await _get_tools(provider, session) + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="start", description="shared") + waiters = [ + asyncio.create_task(_invoke_tool(tools["background_agents_wait_for_first_completion"], task_ids=[1])) + for _ in range(2) + ] + try: + await both_waiting.wait() + agent_release.set() + assert await asyncio.gather(*waiters) == ["Task 1 finished with status: completed."] * 2 + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "shared result" + assert not provider._runtime[session.session_id].in_flight_tasks + finally: + agent_release.set() + await asyncio.gather(*waiters, return_exceptions=True) + await provider.release_session(session) + + async def test_wait_for_first_completion() -> None: """Should wait and return when a task completes.""" provider = _make_provider(_FakeAgent("Fast", response_text="fast result", delay=0.01))