fix(gateway): auto-resume sessions after drain-timeout restart (#11852) (#12301)

The shutdown banner promised "send any message after restart to resume
where you left off" but the code did the opposite: a drain-timeout
restart skipped the .clean_shutdown marker, which made the next startup
call suspend_recently_active(), which marked the session suspended,
which made get_or_create_session() spawn a fresh session_id with a
'Session automatically reset. Use /resume...' notice — contradicting
the banner.

Introduce a resume_pending state on SessionEntry that is distinct from
suspended. Drain-timeout shutdown flags active sessions resume_pending
instead of letting startup-wide suspension destroy them. The next
message on the same session_key preserves the session_id, reloads the
transcript, and the agent receives a reason-aware restart-resume
system note that subsumes the existing tool-tail auto-continue note
(PR #9934).

Terminal escalation still flows through the existing
.restart_failure_counts stuck-loop counter (PR #7536, threshold 3) —
no parallel counter on SessionEntry. suspended still wins over
resume_pending in get_or_create_session() so genuinely stuck sessions
converge to a clean slate.

Spec: PR #11852 (BrennerSpear). Implementation follows the spec with
the approved correction (reuse .restart_failure_counts rather than
adding a resume_attempts field).

Changes:
- gateway/session.py: SessionEntry.resume_pending/resume_reason/
  last_resume_marked_at + to_dict/from_dict; SessionStore
  .mark_resume_pending()/clear_resume_pending(); get_or_create_session()
  returns existing entry when resume_pending (suspended still wins);
  suspend_recently_active() skips resume_pending entries.
- gateway/run.py: _stop_impl() drain-timeout branch marks active
  sessions resume_pending before _interrupt_running_agents();
  _run_agent() injects reason-aware restart-resume system note that
  subsumes the tool-tail case; successful-turn cleanup also clears
  resume_pending next to _clear_restart_failure_count();
  _notify_active_sessions_of_shutdown() softens the restart banner to
  'I'll try to resume where you left off' (honest about stuck-loop
  escalation).
- tests/gateway/test_restart_resume_pending.py: 29 new tests covering
  SessionEntry roundtrip, mark/clear helpers, get_or_create_session
  precedence (suspended > resume_pending), suspend_recently_active
  skip, drain-timeout mark reason (restart vs shutdown), system-note
  injection decision tree (including tool-tail subsumption), banner
  wording, and stuck-loop escalation override.
This commit is contained in:
Teknium
2026-04-18 17:32:17 -07:00
committed by GitHub
parent ad99e32371
commit cb4addacab
3 changed files with 782 additions and 5 deletions

View File

@@ -1539,7 +1539,7 @@ class GatewayRunner:
action = "restarting" if self._restart_requested else "shutting down"
hint = (
"Your current task will be interrupted. "
"Send any message after restart to resume where it left off."
"Send any message after restart and I'll try to resume where you left off."
if self._restart_requested
else "Your current task will be interrupted."
)
@@ -2373,6 +2373,27 @@ class GatewayRunner:
timeout,
self._running_agent_count(),
)
# Mark forcibly-interrupted sessions as resume_pending BEFORE
# interrupting the agents. This preserves each session's
# session_id + transcript so the next message on the same
# session_key auto-resumes from the existing conversation
# instead of getting routed through suspend_recently_active()
# and converted into a fresh session. Terminal escalation
# for genuinely stuck sessions still flows through the
# existing ``.restart_failure_counts`` stuck-loop counter
# (incremented below, threshold 3), which sets
# ``suspended=True`` and overrides resume_pending.
_resume_reason = (
"restart_timeout" if self._restart_requested else "shutdown_timeout"
)
for _sk in list(active_agents.keys()):
try:
self.session_store.mark_resume_pending(_sk, _resume_reason)
except Exception as _e:
logger.debug(
"mark_resume_pending failed for %s: %s",
_sk[:20], _e,
)
self._interrupt_running_agents(
"Gateway restarting" if self._restart_requested else "Gateway shutting down"
)
@@ -4152,8 +4173,20 @@ class GatewayRunner:
# Successful turn — clear any stuck-loop counter for this session.
# This ensures the counter only accumulates across CONSECUTIVE
# restarts where the session was active (never completed).
#
# Also clear the resume_pending flag (set by drain-timeout
# shutdown) — the turn ran to completion, so recovery
# succeeded and subsequent messages should no longer receive
# the restart-interruption system note.
if session_key:
self._clear_restart_failure_count(session_key)
try:
self.session_store.clear_resume_pending(session_key)
except Exception as _e:
logger.debug(
"clear_resume_pending failed for %s: %s",
session_key[:20], _e,
)
# Surface error details when the agent failed silently (final_response=None)
if not response and agent_result.get("failed"):
@@ -9427,7 +9460,40 @@ class GatewayRunner:
# restart, crash, SIGTERM). Prepend a system note so the model
# finishes processing the pending tool results before addressing
# the user's new message. (#4493)
if agent_history and agent_history[-1].get("role") == "tool":
#
# Session-level resume_pending (set on drain-timeout shutdown)
# escalates the wording — the transcript's last role may be
# anything (tool, assistant with unfinished work, etc.), so we
# give a stronger, reason-aware instruction that subsumes the
# tool-tail case.
_resume_entry = None
if session_key:
try:
_resume_entry = self.session_store._entries.get(session_key)
except Exception:
_resume_entry = None
_is_resume_pending = bool(
_resume_entry is not None and getattr(_resume_entry, "resume_pending", False)
)
if _is_resume_pending:
_reason = getattr(_resume_entry, "resume_reason", None) or "restart_timeout"
_reason_phrase = (
"a gateway restart"
if _reason == "restart_timeout"
else "a gateway shutdown"
if _reason == "shutdown_timeout"
else "a gateway interruption"
)
message = (
f"[System note: Your previous turn in this session was interrupted "
f"by {_reason_phrase}. The conversation history below is intact. "
f"If it contains unfinished tool result(s), process them first and "
f"summarize what was accomplished, then address the user's new "
f"message below.]\n\n"
+ message
)
elif agent_history and agent_history[-1].get("role") == "tool":
message = (
"[System note: Your previous turn was interrupted before you could "
"process the last tool result(s). The conversation history contains "

View File

@@ -377,7 +377,19 @@ class SessionEntry:
# this session (create a new session_id) so the user starts fresh.
# Set by /stop to break stuck-resume loops (#7536).
suspended: bool = False
# When True the session was interrupted by a gateway restart/shutdown
# drain timeout, but recovery is still expected. Unlike ``suspended``,
# ``resume_pending`` preserves the existing session_id on next access —
# the user stays on the same transcript and the agent auto-continues
# from where it left off. Cleared after the next successful turn.
# Escalation to ``suspended`` is handled by the existing
# ``.restart_failure_counts`` stuck-loop counter (#7536), not by a
# parallel counter on this entry.
resume_pending: bool = False
resume_reason: Optional[str] = None # e.g. "restart_timeout"
last_resume_marked_at: Optional[datetime] = None
def to_dict(self) -> Dict[str, Any]:
result = {
"session_key": self.session_key,
@@ -397,6 +409,13 @@ class SessionEntry:
"cost_status": self.cost_status,
"memory_flushed": self.memory_flushed,
"suspended": self.suspended,
"resume_pending": self.resume_pending,
"resume_reason": self.resume_reason,
"last_resume_marked_at": (
self.last_resume_marked_at.isoformat()
if self.last_resume_marked_at
else None
),
}
if self.origin:
result["origin"] = self.origin.to_dict()
@@ -414,7 +433,15 @@ class SessionEntry:
platform = Platform(data["platform"])
except ValueError as e:
logger.debug("Unknown platform value %r: %s", data["platform"], e)
last_resume_marked_at = None
_lrma = data.get("last_resume_marked_at")
if _lrma:
try:
last_resume_marked_at = datetime.fromisoformat(_lrma)
except (TypeError, ValueError):
last_resume_marked_at = None
return cls(
session_key=data["session_key"],
session_id=data["session_id"],
@@ -434,6 +461,9 @@ class SessionEntry:
cost_status=data.get("cost_status", "unknown"),
memory_flushed=data.get("memory_flushed", False),
suspended=data.get("suspended", False),
resume_pending=data.get("resume_pending", False),
resume_reason=data.get("resume_reason"),
last_resume_marked_at=last_resume_marked_at,
)
@@ -710,9 +740,23 @@ class SessionStore:
entry = self._entries[session_key]
# Auto-reset sessions marked as suspended (e.g. after /stop
# broke a stuck loop — #7536).
# broke a stuck loop — #7536). ``suspended`` is the hard
# forced-wipe signal and always wins over ``resume_pending``,
# so repeated interrupted restarts that escalate via the
# existing ``.restart_failure_counts`` stuck-loop counter
# still converge to a clean slate.
if entry.suspended:
reset_reason = "suspended"
elif entry.resume_pending:
# Restart-interrupted session: preserve the session_id
# and return the existing entry so the transcript
# reloads intact. ``resume_pending`` is cleared after
# the NEXT successful turn completes (not here), which
# means a re-interrupted retry keeps trying — the
# stuck-loop counter handles terminal escalation.
entry.updated_at = now
self._save()
return entry
else:
reset_reason = self._should_reset(entry, source)
if not reset_reason:
@@ -802,6 +846,55 @@ class SessionStore:
return True
return False
def mark_resume_pending(
self,
session_key: str,
reason: str = "restart_timeout",
) -> bool:
"""Mark a session as resumable after a restart interruption.
Unlike ``suspend_session()``, this preserves the existing
``session_id`` and the transcript. The next call to
``get_or_create_session()`` for this key returns the same entry
so the user auto-resumes on the same conversation lane.
Returns True if the session existed and was marked.
"""
with self._lock:
self._ensure_loaded_locked()
if session_key in self._entries:
entry = self._entries[session_key]
# Never override an explicit ``suspended`` — that is a hard
# forced-wipe signal (from /stop or stuck-loop escalation).
if entry.suspended:
return False
entry.resume_pending = True
entry.resume_reason = reason
entry.last_resume_marked_at = _now()
self._save()
return True
return False
def clear_resume_pending(self, session_key: str) -> bool:
"""Clear the resume-pending flag after a successful resumed turn.
Called from the gateway after ``run_conversation()`` returns a
final response for a session that had ``resume_pending=True``,
signalling that recovery succeeded.
Returns True if a flag was cleared.
"""
with self._lock:
self._ensure_loaded_locked()
entry = self._entries.get(session_key)
if entry is None or not entry.resume_pending:
return False
entry.resume_pending = False
entry.resume_reason = None
entry.last_resume_marked_at = None
self._save()
return True
def prune_old_entries(self, max_age_days: int) -> int:
"""Drop SessionEntry records older than max_age_days.
@@ -861,6 +954,12 @@ class SessionStore:
(#7536). Only suspends sessions updated within *max_age_seconds*
to avoid resetting long-idle sessions that are harmless to resume.
Returns the number of sessions that were suspended.
Entries flagged ``resume_pending=True`` are skipped — those were
marked intentionally by the drain-timeout path as recoverable.
Terminal escalation for genuinely stuck ``resume_pending`` sessions
is handled by the existing ``.restart_failure_counts`` stuck-loop
counter, which runs after this method on startup.
"""
from datetime import timedelta
@@ -869,6 +968,8 @@ class SessionStore:
with self._lock:
self._ensure_loaded_locked()
for entry in self._entries.values():
if entry.resume_pending:
continue
if not entry.suspended and entry.updated_at >= cutoff:
entry.suspended = True
count += 1