From 491105c7c28ba9433587001977ac236bd991fc8e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 20 Jul 2026 10:39:06 +0000 Subject: [PATCH] fix(scheduler): abort queued wake inject on cancel; purge goal-check on terminal cancel When cancel_session_tasks skipped a queued agent_wake inject at dispatch, the wake stayed staged and poll_due_wakes never retried it. Call _finalize_agent_wake_delivery(confirmed=False) on the dispatch skip path. Terminal cancel also flushed staged goal-check wakes from the pending queue after the user aborted the session. Purge goal-check wakes at cancel start so flush_pending_wakes cannot resurrect them. Co-authored-by: Yuxuan Liu --- src/system/kernel/scheduler.py | 11 ++++ tests/test_kernel_runtime.py | 106 +++++++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+) diff --git a/src/system/kernel/scheduler.py b/src/system/kernel/scheduler.py index e1690b9..24d2fcf 100644 --- a/src/system/kernel/scheduler.py +++ b/src/system/kernel/scheduler.py @@ -608,6 +608,16 @@ async def cancel_session_tasks(self, session_id: str) -> bool: 避免新入队请求在取消窗口内被误判为 cancelled 而 skip。 """ self._cancelled_sessions.add(session_id) + try: + from agent_core.tools.agent_wake import cancel_pending_wakes + + cancel_pending_wakes(session_id, label="goal-check") + except Exception as exc: + logger.debug( + "KernelScheduler: cancel goal-check wakes failed session_id=%s: %s", + session_id, + exc, + ) tasks = self._session_active_tasks.get(session_id, set()) active_tasks = [t for t in list(tasks) if not t.done()] cancelled_any = bool(active_tasks) @@ -665,6 +675,7 @@ async def _dispatch_loop(self) -> None: request.request_id, asyncio.CancelledError("session cancelled while queued"), ) + self._finalize_agent_wake_delivery(request, confirmed=False) self._queue.task_done() self._maybe_clear_cancelled(request.session_id) continue diff --git a/tests/test_kernel_runtime.py b/tests/test_kernel_runtime.py index 8b35b2b..b6409e4 100644 --- a/tests/test_kernel_runtime.py +++ b/tests/test_kernel_runtime.py @@ -1185,6 +1185,112 @@ async def test_cancel_skip_aborts_agent_wake_delivery() -> None: assert len(poll_due_wakes()) == 1 +@pytest.mark.asyncio +async def test_cancel_dispatch_skip_aborts_agent_wake_delivery() -> None: + """dispatch 在队列层跳过已 cancel 的 wake inject 时应 abort,供 poll 重试。""" + from agent_core.tools.agent_wake import ( + clear_all_wakes_for_tests, + deliver_wake_via_inject, + list_wakes, + poll_due_wakes, + register_wake, + ) + from agent_core.tools.bash_job_notify import set_notify_dependencies + + clear_all_wakes_for_tests() + scheduler = _minimal_scheduler_for_cancel_tests() + await scheduler.start() + session_id = "cli:root" + wid = register_wake( + session_id=session_id, + fire_at=time.time() - 1, + message="wake queued then cancel", + wake_id="wake-dispatch-cancel", + ) + wake = poll_due_wakes()[0] + set_notify_dependencies(scheduler=scheduler, core_pool=None) + deliver_wake_via_inject(wake=wake, scheduler=scheduler) + assert list_wakes()[0]["staged"] is True + assert scheduler._queue.qsize() >= 1 + + await scheduler.cancel_session_tasks(session_id) + for _ in range(50): + if list_wakes()[0]["staged"] is False: + break + await asyncio.sleep(0.02) + + assert list_wakes()[0]["staged"] is False + assert len(poll_due_wakes()) == 1 + await scheduler.stop() + + +@pytest.mark.asyncio +async def test_terminal_cancel_does_not_flush_staged_goal_check() -> None: + """终端 cancel 不应在 finally flush 已暂存的 goal-check 唤醒。""" + from agent_core.tools.agent_wake import ( + clear_all_wakes_for_tests, + deliver_wake_via_inject, + list_wakes, + poll_due_wakes, + register_wake, + ) + from agent_core.tools.bash_job_notify import set_notify_dependencies + + clear_all_wakes_for_tests() + scheduler = _minimal_scheduler_for_cancel_tests() + session_id = "cli:root" + hold = asyncio.Event() + release = asyncio.Event() + + async def _slow_run(*_a: Any, **_k: Any) -> Any: + from agent_core.interfaces import AgentRunResult + + hold.set() + await release.wait() + return AgentRunResult(output_text="done") + + scheduler._kernel.run = AsyncMock(side_effect=_slow_run) # type: ignore[attr-defined] + scheduler._core_pool.acquire = AsyncMock( # type: ignore[attr-defined] + return_value=SimpleNamespace( + prepare_turn=AsyncMock(return_value=1), + _finalize_turn=AsyncMock(), + _session_logger=None, + ) + ) + set_notify_dependencies(scheduler=scheduler, core_pool=None) + await scheduler.start() + + user_req = KernelRequest.create(text="block", session_id=session_id) + user_task = asyncio.create_task(scheduler.submit(user_req)) + for _ in range(50): + if hold.is_set(): + break + await asyncio.sleep(0.02) + assert scheduler.session_inflight_request_count(session_id) == 1 + + register_wake( + session_id=session_id, + fire_at=time.time() - 1, + message="goal nudge", + label="goal-check", + wake_id="wake-goal-cancel", + ) + wake = poll_due_wakes()[0] + assert deliver_wake_via_inject(wake=wake, scheduler=scheduler) is True + assert scheduler._queue.qsize() == 0 + + await scheduler.cancel_session_tasks(session_id) + for _ in range(50): + if scheduler.session_inflight_request_count(session_id) == 0: + break + await asyncio.sleep(0.02) + release.set() + + assert list_wakes() == [] + assert scheduler._queue.qsize() == 0 + await scheduler.stop() + + @pytest.mark.asyncio async def test_cancel_session_tasks_clears_mark_when_lock_waiter_aborted() -> None: """stop 时若第二个请求在等锁,cancel 后 session 应恢复可投递。"""