fix(runtime): preserve interrupted turns on gateway exit

This commit is contained in:
Xubin Ren
2026-08-24 00:58:04 +08:00
parent c7e2a474a0
commit 2cdfba38b2
8 changed files with 94 additions and 88 deletions
+8 -11
View File
@@ -1371,12 +1371,12 @@ class AgentLoop:
await self.aclose()
def preserve_inflight_turns_on_shutdown(self) -> None:
"""Keep durable checkpoints when the owning gateway is restarting.
"""Keep durable checkpoints when the owning gateway exits.
Normal cancellation intentionally materializes partial output so a
stopped gateway leaves a readable conversation. A managed restart is
different: RecoveryCoordinator needs the checkpoint intact to safely
offer the unfinished turn for explicit continuation after restart.
user-stopped turn leaves a readable conversation. Gateway lifecycle
shutdown is different: RecoveryCoordinator needs the checkpoint intact
to safely offer the unfinished turn for explicit continuation later.
"""
self._preserve_inflight_turns_on_shutdown = True
@@ -1443,13 +1443,10 @@ class AgentLoop:
session_key,
exc_info=True,
)
# Preserve partial context from the interrupted turn so
# the user does not lose tool results and assistant
# messages accumulated before /stop. The checkpoint was
# already persisted to session metadata by
# _emit_checkpoint during tool execution; materializing
# it into session history now makes it visible in the
# next conversation turn.
# An explicit turn stop materializes partial context so
# the next prompt can see completed tool results. Gateway
# shutdown keeps the durable checkpoint untouched instead,
# allowing RecoveryCoordinator to offer Continue safely.
if (
session_key in self._discarding_sessions
or self._preserve_inflight_turns_on_shutdown
+4 -2
View File
@@ -963,8 +963,10 @@ def _run_gateway(
with suppress(asyncio.CancelledError):
await shutdown_task
cron.stop()
if gateway_runtime.preserves_inflight_turns_on_exit():
agent.preserve_inflight_turns_on_shutdown()
# A gateway exit interrupts ownership of active turns; it is
# not the same as the user stopping a turn. Keep checkpoints
# so the next gateway can offer an explicit Continue action.
agent.preserve_inflight_turns_on_shutdown()
agent.stop()
# Cancel runtime tasks first, then deterministically close
# exec/MCP resources while the event loop is still alive.
+1 -55
View File
@@ -201,52 +201,6 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
"""Serialize long lifecycle transitions without blocking child cleanup."""
return FileLock(f"{self.paths.state_path}.transition.lock")
@property
def _restart_intent_path(self) -> Path:
"""Return the short-lived marker used to distinguish restart from stop.
A gateway receives the same operating-system termination request for a
graceful ``restart`` and an explicit ``stop``. The marker lets the
exiting process preserve its durable turn checkpoint only for the
former. It is intentionally local to one gateway instance.
"""
return self.paths.state_path.with_name(f"{self.paths.state_path.name}.restart")
def preserves_inflight_turns_on_exit(self) -> bool:
"""Whether this gateway was asked to exit as part of a managed restart."""
try:
raw_intent: object = json.loads(
self._restart_intent_path.read_text(encoding="utf-8")
)
except (json.JSONDecodeError, OSError):
return False
if not isinstance(raw_intent, dict):
return False
intent = cast(dict[str, object], raw_intent)
return intent.get("pid") == os.getpid()
def _write_restart_intent(self, pid: int) -> None:
self.paths.run_dir.mkdir(parents=True, exist_ok=True)
target = self._restart_intent_path
with tempfile.NamedTemporaryFile(
"w",
encoding="utf-8",
dir=target.parent,
prefix=f".{target.name}.",
delete=False,
) as handle:
json.dump({"pid": pid}, handle)
handle.flush()
os.fsync(handle.fileno())
temporary = Path(handle.name)
os.replace(temporary, target)
def _clear_restart_intent(self) -> None:
try:
self._restart_intent_path.unlink()
except FileNotFoundError:
pass
def start_background(self, options: ProcessStartOptions) -> RuntimeResult:
"""Start the gateway detached from the current terminal."""
lease = GatewayClientLease(self, kind="gateway-background")
@@ -399,15 +353,7 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
"gateway_foreground_restart_required",
status,
)
assert status.pid is not None
self._write_restart_intent(status.pid)
try:
stop_result = self._stop(timeout_s=timeout_s)
finally:
# The old process reads the marker while handling shutdown.
# Never let a stale marker turn a later explicit stop into a
# recoverable restart.
self._clear_restart_intent()
stop_result = self._stop(timeout_s=timeout_s)
if not stop_result.ok:
return self._result(stop_result)
with self._lifecycle_lock():
+30 -2
View File
@@ -677,13 +677,18 @@ class RecoveryCoordinator:
# contains an activity row without a turn_end. Treat it as an
# interrupted turn instead of letting the UI resurrect it as a
# forever-running spinner.
can_continue = self._has_saved_continuation_context(session)
waiting = self._set_state(
session,
status="awaiting_user",
recovery_id=uuid4().hex,
attempts=0,
reason="interrupted_without_checkpoint",
can_continue=False,
reason=(
"interrupted_with_saved_context"
if can_continue
else "interrupted_without_checkpoint"
),
can_continue=can_continue,
)
self.sessions.save(session)
await self._publish(chat_id, waiting)
@@ -892,6 +897,29 @@ class RecoveryCoordinator:
# corrupt or unavailable; the normal checkpoint path still applies.
return False
@staticmethod
def _has_saved_continuation_context(session: Session) -> bool:
"""Whether an interrupted turn left model-visible context to continue from."""
last_user = next(
(
index
for index in range(len(session.messages) - 1, -1, -1)
if session.messages[index].get("role") == "user"
),
None,
)
if last_user is None:
return False
tail = session.messages[last_user + 1 :]
return bool(tail) and (
tail[-1].get("role") == "tool"
or any(message.get("_recovery_interrupted") is True for message in tail)
or any(
message.get("role") == "assistant" and bool(message.get("tool_calls"))
for message in tail
)
)
@staticmethod
def _websocket_route(session: Session) -> tuple[str, str] | None:
return RecoveryCoordinator._websocket_route_for(session.key, session.metadata)