From 416d2a015727252855e0b8e961fee9e9519f5fb2 Mon Sep 17 00:00:00 2001
From: kshitij <kshitijkapoor0611@gmail.com>
Date: Fri, 7 Aug 2026 13:55:10 +0530
Subject: [PATCH] fix(gateway): re-signal interrupts when work is still live at
 settle-window exit
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 gateway/run.py                                | 20 +++++++
 .../test_api_server_active_work_drain.py      | 53 +++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/gateway/run.py b/gateway/run.py
index b2f4f2d95b819..52ae139c5c24e 100644
--- a/gateway/run.py
+++ b/gateway/run.py
@@ -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),
diff --git a/tests/gateway/test_api_server_active_work_drain.py b/tests/gateway/test_api_server_active_work_drain.py
index 9edc39f72e0a8..28740e4d30e39 100644
--- a/tests/gateway/test_api_server_active_work_drain.py
+++ b/tests/gateway/test_api_server_active_work_drain.py
@@ -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,
+        ]
+