fix(gateway): mark only still-running sessions resume_pending on drain timeout (#12332)

Follow-up to #12301.

The drain-timeout branch of _stop_impl() was iterating the drain-start
snapshot (active_agents) when marking sessions resume_pending. That
snapshot can include sessions that finished gracefully during the drain
window — marking them would give their next turn a stray
'your previous turn was interrupted by a gateway restart' system note
even though the prior turn actually completed cleanly.

Iterate self._running_agents at timeout time instead, mirroring
_interrupt_running_agents() exactly:
- only sessions still blocking the shutdown get marked
- pending sentinels (AIAgent construction not yet complete) are skipped

Changes:
- gateway/run.py: swap active_agents.keys() for filtered
  self._running_agents.items() iteration in the drain-timeout mark loop.
- tests/gateway/test_restart_resume_pending.py: two regression tests —
  finisher-during-drain not marked, pending sentinel not marked.
This commit is contained in:
Teknium
2026-04-18 17:40:34 -07:00
committed by GitHub
parent cb4addacab
commit c49a58a6d0
2 changed files with 92 additions and 1 deletions

View File

@@ -2383,10 +2383,23 @@ class GatewayRunner:
# existing ``.restart_failure_counts`` stuck-loop counter
# (incremented below, threshold 3), which sets
# ``suspended=True`` and overrides resume_pending.
#
# Iterate self._running_agents (current) rather than the
# drain-start ``active_agents`` snapshot — the snapshot
# may include sessions that finished gracefully during
# the drain window, and marking those falsely would give
# them a stray restart-interruption system note on their
# next turn even though their previous turn completed
# cleanly. Skip pending sentinels for the same reason
# _interrupt_running_agents() does: their agent hasn't
# started yet, there's nothing to interrupt, and the
# session shouldn't carry a misleading resume flag.
_resume_reason = (
"restart_timeout" if self._restart_requested else "shutdown_timeout"
)
for _sk in list(active_agents.keys()):
for _sk, _agent in list(self._running_agents.items()):
if _agent is _AGENT_PENDING_SENTINEL:
continue
try:
self.session_store.mark_resume_pending(_sk, _resume_reason)
except Exception as _e: