fix(gateway): re-signal interrupts when work is still live at settle-window exit

Review follow-up for the salvaged #79881/#63963 stack: the shutdown
interrupt fires exactly once, but work can materialize AFTER that one
shot on BOTH sibling paths:

- a /v1/runs task admitted before the drain populates
  _active_run_agents only once _create_agent returns
  (queued-before-agent window);
- a _running_agents entry claimed as _AGENT_PENDING_SENTINEL is
  promoted to the real agent by track_agent() on its own schedule,
  after the one-shot walk skipped the sentinel.

Either way the settle loop waited on work nothing signaled, and the
turn went straight to the post-interrupt tool-subprocess kill — the
exact amputation the fix exists to avoid, in a rarer window.

If any work is still live when the settle loop exits, re-invoke
_interrupt_running_agents (which already skips sentinels and folds in
the API-server helper) so late-materializing agents on either path get
the cooperative interrupt. Regression test drives the real stop() path
with an accelerated loop clock and asserts exactly two interrupt
signals.
This commit is contained in:
kshitij 2026-08-07 13:55:10 +05:30 committed by kshitij
parent d9ddfb23d5
commit 416d2a0157
2 changed files with 73 additions and 0 deletions

View File

@ -12950,6 +12950,26 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
self._update_runtime_status("draining")
await asyncio.sleep(0.1)
# The interrupt above fires exactly once, but work can
# materialize AFTER that one shot: a /v1/runs task admitted
# before the drain populates _active_run_agents only once
# _create_agent returns, and a _running_agents entry claimed
# as _AGENT_PENDING_SENTINEL is promoted to a real agent by
# track_agent() on its own schedule. Either way the settle
# loop waited on work nothing signaled. If any is still live
# at settle-loop exit, re-signal so a late-materializing
# agent gets a cooperative interrupt instead of going
# straight to the tool-subprocess kill.
if self._running_agents or self._active_api_run_count():
self._interrupt_running_agents(
_INTERRUPT_REASON_GATEWAY_RESTART
if self._restart_requested
else _INTERRUPT_REASON_GATEWAY_SHUTDOWN
)
logger.debug(
"Re-signaled interrupt for work still live at settle-window exit"
)
# Kill lingering tool subprocesses NOW, before we spend more
# budget on adapter disconnect / session DB close. Under
# systemd (TimeoutStopSec bounded by drain_timeout+headroom),

View File

@ -552,4 +552,57 @@ class TestShutdownSettleWindow:
"still unwinding"
)
@pytest.mark.asyncio
async def test_api_work_still_live_at_settle_exit_is_reinterrupted(
self, monkeypatch
):
"""A /v1/runs agent can materialize AFTER the one-shot interrupt.
The task is counted via ``_active_run_tasks`` from admission, but
``_active_run_agents[run_id]`` is populated only once ``_create_agent``
returns an agent landing in that window missed the single interrupt
and previously went straight to the tool-subprocess kill. The settle
loop must re-signal when API work is still live at exit.
"""
import tools.browser_tool as _bt
import tools.process_registry as _pr
import tools.terminal_tool as _tt
runner, adapter = make_restart_runner()
runner._restart_drain_timeout = 0.01
adapter.disconnect = _make_async_noop()
api = _SettlingApiAdapter(polls_to_settle=10_000) # never settles
runner.adapters = {Platform.TELEGRAM: adapter, Platform.API_SERVER: api}
monkeypatch.setattr(_pr.process_registry, "kill_all", lambda task_id=None: 0)
monkeypatch.setattr(_tt, "cleanup_all_environments", lambda: None)
monkeypatch.setattr(_bt, "cleanup_all_browsers", lambda: None)
# Accelerate the loop clock: each time() call advances 1s of virtual
# time, so the 5s settle deadline expires after a handful of polls
# instead of 5 real seconds. Relative deadline math is preserved.
loop = asyncio.get_running_loop()
_real_time = type(loop).time
_skew = [0.0]
def _fast_time(self):
_skew[0] += 1.0
return _real_time(self) + _skew[0]
monkeypatch.setattr(type(loop), "time", _fast_time)
try:
with patch("gateway.status.remove_pid_file"), \
patch("gateway.status.write_runtime_status"), \
patch("cron.scheduler.mark_job_run"):
await runner.stop()
finally:
monkeypatch.undo()
# One shot from _interrupt_running_agents + one re-signal at settle
# exit because API work was still live.
assert api.interrupt_reasons == [
_INTERRUPT_REASON_GATEWAY_SHUTDOWN,
_INTERRUPT_REASON_GATEWAY_SHUTDOWN,
]