From 56f6383b131517576c9966f07e828f7cd9b9a99e Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 15:04:07 +0800 Subject: [PATCH 1/2] fix(agent-bus): establish a dispatched session task, and cancel a round cleanly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``_dispatch_session_task`` returned without a guaranteed suspension point: ``_emit_session_task_submitted`` returns before awaiting when no sink is installed, and a sink whose ``append`` never awaits does not yield either. So ``submit_task_to_session`` — fire-and-return by contract — could hand a job id back while its coroutine had not started. A host reporting job state right after dispatch then described a worker that had not begun as merely ``"submitted"``, and every cancellation in that window relied on ``_mark_job_aborted`` standing in for a ``finally`` the coroutine never reached. One trailing yield establishes the task instead. That exposed two defects in ``cancel_for_parent`` which its own test had been hiding — the test only passed because none of its three jobs had ever started: - the running job's cancellation runs its ``finally``, which drains the session queue, so the single interleaved pass dispatched the very tasks the call was tearing down. Each one reached the LLM before being cancelled again a few iterations later. Queued jobs are now cleared in a first pass, so the drain finds nothing to start. - the return count only counted jobs that still needed the stand-in. A job that had actually started absorbs its own cancellation, so a round that really did tear down every running sub-agent reported 0. Both outcomes are now counted. Co-Authored-By: Claude Opus 5 (1M context) --- agent_core/components/agent_bus/bus.py | 53 +++++-- ...test_agent_bus_session_dispatch_handoff.py | 129 ++++++++++++++++++ 2 files changed, 173 insertions(+), 9 deletions(-) create mode 100644 tests/test_agent_bus_session_dispatch_handoff.py diff --git a/agent_core/components/agent_bus/bus.py b/agent_core/components/agent_bus/bus.py index 6768f9b..92bbdbf 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", @@ -1635,6 +1655,21 @@ async def _run_and_finalize() -> SubAgentResult: 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 — ``_emit_session_task_submitted`` returns + # without awaiting when no sink is installed, and a sink whose + # ``append`` never awaits does not yield either — so without this the + # task can still be unstarted when the caller moves on. Until + # ``_run_and_finalize`` begins, neither the ``finally`` that returns + # the guard reservation nor the handler that turns a cancellation into + # a real ``SubAgentResult`` exists yet, so a cancel landing in that + # window kills the job outright and the caller gets the synthetic + # "cancelled before publishing a report" instead of the partial result + # the sub-agent actually had. One yield establishes the task; every + # cancel after it lands inside the job wrapper. + await asyncio.sleep(0) + 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..b746db2 --- /dev/null +++ b/tests/test_agent_bus_session_dispatch_handoff.py @@ -0,0 +1,129 @@ +"""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.models.agent_definition import AgentDefinition +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_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 From 9015e1c5c8ea1c50f0f083788dbe8707c365c31a Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 15:20:43 +0800 Subject: [PATCH 2/2] fix(agent-bus): tear down cancelled session dispatch --- agent_core/components/agent_bus/bus.py | 45 ++++++++++------- ...test_agent_bus_session_dispatch_handoff.py | 48 +++++++++++++++++++ 2 files changed, 75 insertions(+), 18 deletions(-) diff --git a/agent_core/components/agent_bus/bus.py b/agent_core/components/agent_bus/bus.py index 92bbdbf..fc13a9a 100644 --- a/agent_core/components/agent_bus/bus.py +++ b/agent_core/components/agent_bus/bus.py @@ -1650,25 +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 — ``_emit_session_task_submitted`` returns - # without awaiting when no sink is installed, and a sink whose - # ``append`` never awaits does not yield either — so without this the - # task can still be unstarted when the caller moves on. Until - # ``_run_and_finalize`` begins, neither the ``finally`` that returns - # the guard reservation nor the handler that turns a cancellation into - # a real ``SubAgentResult`` exists yet, so a cancel landing in that - # window kills the job outright and the caller gets the synthetic - # "cancelled before publishing a report" instead of the partial result - # the sub-agent actually had. One yield establishes the task; every - # cancel after it lands inside the job wrapper. - await asyncio.sleep(0) + # 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 index b746db2..c1b9c4c 100644 --- a/tests/test_agent_bus_session_dispatch_handoff.py +++ b/tests/test_agent_bus_session_dispatch_handoff.py @@ -24,7 +24,9 @@ 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 @@ -89,6 +91,52 @@ async def test_dispatch_leaves_the_job_running_not_merely_submitted(monkeypatch) 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,