diff --git a/agent_core/components/agent_bus/bus.py b/agent_core/components/agent_bus/bus.py index 6768f9b..fc13a9a 100644 --- a/agent_core/components/agent_bus/bus.py +++ b/agent_core/components/agent_bus/bus.py @@ -903,25 +903,45 @@ async def cancel_for_parent(self, parent_task_id: str) -> int: running tasks (asyncio cancel) and queued tasks (purge + status flip) are handled — neither survives the round boundary. - Returns the number of jobs cancelled. + Queued jobs are cleared first, in a separate pass. Cancelling a + running job runs its ``finally``, which drains that session's queue — + so a single interleaved pass would dispatch the very tasks this call + is tearing down, let them reach the LLM, and cancel them again a few + iterations later. + + Returns the number of jobs this call cancelled. """ cancelled = 0 - for jid, entry in list(self._jobs.items()): - if entry.parent_task_id != parent_task_id: - continue - if entry.status in ("completed", "failed", "aborted"): - continue + mine = [ + (jid, entry) for jid, entry in self._jobs.items() + if entry.parent_task_id == parent_task_id + and entry.status not in ("completed", "failed", "aborted") + ] + + for jid, entry in mine: if entry.status == "queued" or entry.task is None: self._purge_queued_job(jid) self._mark_job_aborted(jid, entry) cancelled += 1 + + for jid, entry in mine: + if entry.status in ("completed", "failed", "aborted"): continue - entry.task.cancel() + task = entry.task + if task is None: + continue + task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): - await entry.task + await task + # Counted either way. A job that had actually started absorbs its + # own cancellation — ``_run_and_finalize`` reaches the handler that + # flips the status itself — and only one that never got that far + # still needs the stand-in below. Both were cancelled by this call, + # so gating the count on which of the two happened reported 0 for a + # round that really did tear down every running sub-agent. if entry.status not in ("completed", "failed", "aborted"): self._mark_job_aborted(jid, entry) - cancelled += 1 + cancelled += 1 if cancelled: logger.info( "Cancelled %d orphaned sub-agent jobs for %s", @@ -1630,10 +1650,34 @@ async def _run_and_finalize() -> SubAgentResult: # reconciliation read stays a cursor consume rather than a log scan. self._note_result_recipient(session.task_id, pending.spawn_context) - await _emit_session_task_submitted( - session, job_id, task_prompt, - event_sink=self._event_sink(), - ) + try: + await _emit_session_task_submitted( + session, job_id, task_prompt, + event_sink=self._event_sink(), + ) + + # Dispatch is fire-and-return: a host tool such as ``assign_task`` + # hands back a job id and its invocation scope closes. Nothing + # above is guaranteed to suspend — the event helper returns + # without awaiting when no sink is installed, and a sink whose + # ``append`` never awaits does not yield either — so explicitly + # give the child task one turn before returning its id. + await asyncio.sleep(0) + except BaseException: + # Ownership has already transferred to ``task``. If dispatch is + # cancelled or event publication fails now, the caller never gets + # the job id, so leaving the child alive would create an + # unreachable sub-agent. It would also let the outer failure path + # release a SpawnGuard reservation while the child still runs. + task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + if entry.status not in ("completed", "failed", "aborted"): + _close_session_boundary_aborted(session) + self._mark_job_aborted(job_id, entry) + if session.current_job_id == job_id: + session.current_job_id = None + raise async def _eager_trim_and_offload(self, session: SubAgentSession) -> None: """Compress completed boundaries immediately after a task finishes. diff --git a/tests/test_agent_bus_session_dispatch_handoff.py b/tests/test_agent_bus_session_dispatch_handoff.py new file mode 100644 index 0000000..c1b9c4c --- /dev/null +++ b/tests/test_agent_bus_session_dispatch_handoff.py @@ -0,0 +1,177 @@ +"""A dispatched session task must be running by the time dispatch returns. + +``submit_task_to_session`` is fire-and-return: a host tool hands back a job id +and its invocation scope closes. Nothing on the dispatch path is guaranteed to +suspend — ``_emit_session_task_submitted`` returns without awaiting when no +sink is installed, and a sink whose ``append`` never awaits does not yield +either — so establishing the task takes an explicit yield. + +Without it the job is still ``"submitted"`` when the caller moves on, and a +host that reports job state right after dispatching (a task board, a status +line) describes a worker that has not begun as though it were merely queued. +The window is also the one ``_mark_job_aborted`` exists to paper over: it +releases the guard reservation and substitutes a stand-in result precisely +because a task can be cancelled before ``_run_and_finalize`` reaches its +``finally``. Establishing the task first is the narrower fix — every +cancellation then lands inside the job wrapper on its own. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock + +import pytest + +from agent_core.components.agent_bus import AgentBus +from agent_core.components.agent_bus.spawn_guard import SpawnGuard +from agent_core.models.agent_definition import AgentDefinition +from agent_core.models.task_budget import TaskBudget +from agent_core.runtime.registries import services as registry +from agent_core.runtime.registries.agents import AgentRegistry +from agent_core.runtime.resources.manager import ResourceManager + +TASK_ID = "t-handoff" + + +@pytest.fixture(autouse=True) +def _clear_registry(): + registry.clear() + yield + registry.clear() + + +def _register_researcher() -> None: + reg = AgentRegistry() + reg.register(AgentDefinition( + role_id="researcher", display_name="researcher", system_prompt="p", + allowed_tools=[], color="#000", icon="agent", + )) + registry.register(AgentRegistry, reg) + resource_mgr = MagicMock(spec=ResourceManager) + resource_mgr.get_llm.return_value = MagicMock() + resource_mgr.get_tools_for_role.return_value = [] + registry.register(ResourceManager, resource_mgr) + + +async def _session(bus: AgentBus) -> str: + return await bus.create_session( + task_id=TASK_ID, name="lit", role_id="researcher", + system_prompt="You are a researcher.", + llm_override=MagicMock(), tools_override=[], + ) + + +def _parked_loop(entered: asyncio.Event): + async def loop(**kwargs): + entered.set() + await asyncio.sleep(3600) + return loop + + +@pytest.mark.asyncio +async def test_dispatch_leaves_the_job_running_not_merely_submitted(monkeypatch): + """``entry.status`` flips to ``"running"`` inside ``_run_and_finalize``. + + Reading ``"submitted"`` here means the coroutine had not started — the + condition the trailing yield exists to rule out. + """ + entered = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _parked_loop(entered), + ) + _register_researcher() + bus = AgentBus() + session_id = await _session(bus) + + job_id = await bus.submit_task_to_session(session_id, "q") + + assert bus.get_job_status(job_id) == "running" + assert entered.is_set() + + +@pytest.mark.asyncio +async def test_cancelled_dispatch_does_not_leave_an_unreachable_running_job( + monkeypatch, +): + """Cancellation after task creation must tear the child back down. + + The explicit dispatch yield is a cancellation point after the child has + been registered but before its id is returned. If the submitting task is + cancelled there, the child must not keep running without a caller-visible + handle or consume a guard slot whose reservation was already released. + """ + submit_holder: dict[str, asyncio.Task[str]] = {} + + async def loop(**kwargs): + submit_holder["task"].cancel() + await asyncio.sleep(3600) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", loop, + ) + _register_researcher() + bus = AgentBus() + guard = SpawnGuard(TaskBudget( + max_parallel=2, max_depth=4, max_tokens=30_000, + )) + bus.set_spawn_guard(guard) + session_id = await _session(bus) + + submit_task = asyncio.create_task( + bus.submit_task_to_session( + session_id, "q", estimated_tokens=10_000, + ), + ) + submit_holder["task"] = submit_task + + with pytest.raises(asyncio.CancelledError): + await submit_task + + entry = next(iter(bus._jobs.values())) + assert entry.status == "aborted" + assert entry.task is not None and entry.task.done() + assert bus.get_session(session_id).current_job_id is None + assert guard.active_count == 0 + assert guard.tokens_reserved == 0 + + +@pytest.mark.asyncio +async def test_cancelling_a_round_never_lets_a_queued_task_reach_the_llm( + monkeypatch, +): + """Tearing down a round must not first start what it is tearing down. + + Cancelling a running job runs its ``finally``, which drains the session + queue. Interleaved with the cancellation pass that dispatches the next + queued task — one LLM round-trip in, then cancelled again — so the queue + has to be cleared first. + """ + gate = asyncio.Event() + reached_llm: list[str] = [] + + async def loop(**kwargs): + reached_llm.append(str(kwargs.get("user_message"))) + await gate.wait() + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", loop, + ) + _register_researcher() + bus = AgentBus() + session_id = await _session(bus) + + jobs = [ + await bus.submit_task_to_session(session_id, prompt) + for prompt in ("T1", "T2", "T3") + ] + + cancelled = await bus.cancel_for_parent(TASK_ID) + + assert reached_llm == ["T1"] + # Every job this call tore down is counted, including the running one + # that absorbed its own cancellation. + assert cancelled == 3 + assert [bus.get_job_status(job) for job in jobs] == ["aborted"] * 3 + assert not bus.get_session(session_id).pending_tasks