mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 08:13:11 +03:00
fix(runtime): preserve interrupted turns on gateway exit
This commit is contained in:
+8
-11
@@ -1371,12 +1371,12 @@ class AgentLoop:
|
|||||||
await self.aclose()
|
await self.aclose()
|
||||||
|
|
||||||
def preserve_inflight_turns_on_shutdown(self) -> None:
|
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
|
Normal cancellation intentionally materializes partial output so a
|
||||||
stopped gateway leaves a readable conversation. A managed restart is
|
user-stopped turn leaves a readable conversation. Gateway lifecycle
|
||||||
different: RecoveryCoordinator needs the checkpoint intact to safely
|
shutdown is different: RecoveryCoordinator needs the checkpoint intact
|
||||||
offer the unfinished turn for explicit continuation after restart.
|
to safely offer the unfinished turn for explicit continuation later.
|
||||||
"""
|
"""
|
||||||
self._preserve_inflight_turns_on_shutdown = True
|
self._preserve_inflight_turns_on_shutdown = True
|
||||||
|
|
||||||
@@ -1443,13 +1443,10 @@ class AgentLoop:
|
|||||||
session_key,
|
session_key,
|
||||||
exc_info=True,
|
exc_info=True,
|
||||||
)
|
)
|
||||||
# Preserve partial context from the interrupted turn so
|
# An explicit turn stop materializes partial context so
|
||||||
# the user does not lose tool results and assistant
|
# the next prompt can see completed tool results. Gateway
|
||||||
# messages accumulated before /stop. The checkpoint was
|
# shutdown keeps the durable checkpoint untouched instead,
|
||||||
# already persisted to session metadata by
|
# allowing RecoveryCoordinator to offer Continue safely.
|
||||||
# _emit_checkpoint during tool execution; materializing
|
|
||||||
# it into session history now makes it visible in the
|
|
||||||
# next conversation turn.
|
|
||||||
if (
|
if (
|
||||||
session_key in self._discarding_sessions
|
session_key in self._discarding_sessions
|
||||||
or self._preserve_inflight_turns_on_shutdown
|
or self._preserve_inflight_turns_on_shutdown
|
||||||
|
|||||||
@@ -963,8 +963,10 @@ def _run_gateway(
|
|||||||
with suppress(asyncio.CancelledError):
|
with suppress(asyncio.CancelledError):
|
||||||
await shutdown_task
|
await shutdown_task
|
||||||
cron.stop()
|
cron.stop()
|
||||||
if gateway_runtime.preserves_inflight_turns_on_exit():
|
# A gateway exit interrupts ownership of active turns; it is
|
||||||
agent.preserve_inflight_turns_on_shutdown()
|
# 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()
|
agent.stop()
|
||||||
# Cancel runtime tasks first, then deterministically close
|
# Cancel runtime tasks first, then deterministically close
|
||||||
# exec/MCP resources while the event loop is still alive.
|
# exec/MCP resources while the event loop is still alive.
|
||||||
|
|||||||
@@ -201,52 +201,6 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
|
|||||||
"""Serialize long lifecycle transitions without blocking child cleanup."""
|
"""Serialize long lifecycle transitions without blocking child cleanup."""
|
||||||
return FileLock(f"{self.paths.state_path}.transition.lock")
|
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:
|
def start_background(self, options: ProcessStartOptions) -> RuntimeResult:
|
||||||
"""Start the gateway detached from the current terminal."""
|
"""Start the gateway detached from the current terminal."""
|
||||||
lease = GatewayClientLease(self, kind="gateway-background")
|
lease = GatewayClientLease(self, kind="gateway-background")
|
||||||
@@ -399,15 +353,7 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
|
|||||||
"gateway_foreground_restart_required",
|
"gateway_foreground_restart_required",
|
||||||
status,
|
status,
|
||||||
)
|
)
|
||||||
assert status.pid is not None
|
stop_result = self._stop(timeout_s=timeout_s)
|
||||||
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()
|
|
||||||
if not stop_result.ok:
|
if not stop_result.ok:
|
||||||
return self._result(stop_result)
|
return self._result(stop_result)
|
||||||
with self._lifecycle_lock():
|
with self._lifecycle_lock():
|
||||||
|
|||||||
@@ -677,13 +677,18 @@ class RecoveryCoordinator:
|
|||||||
# contains an activity row without a turn_end. Treat it as an
|
# contains an activity row without a turn_end. Treat it as an
|
||||||
# interrupted turn instead of letting the UI resurrect it as a
|
# interrupted turn instead of letting the UI resurrect it as a
|
||||||
# forever-running spinner.
|
# forever-running spinner.
|
||||||
|
can_continue = self._has_saved_continuation_context(session)
|
||||||
waiting = self._set_state(
|
waiting = self._set_state(
|
||||||
session,
|
session,
|
||||||
status="awaiting_user",
|
status="awaiting_user",
|
||||||
recovery_id=uuid4().hex,
|
recovery_id=uuid4().hex,
|
||||||
attempts=0,
|
attempts=0,
|
||||||
reason="interrupted_without_checkpoint",
|
reason=(
|
||||||
can_continue=False,
|
"interrupted_with_saved_context"
|
||||||
|
if can_continue
|
||||||
|
else "interrupted_without_checkpoint"
|
||||||
|
),
|
||||||
|
can_continue=can_continue,
|
||||||
)
|
)
|
||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
await self._publish(chat_id, waiting)
|
await self._publish(chat_id, waiting)
|
||||||
@@ -892,6 +897,29 @@ class RecoveryCoordinator:
|
|||||||
# corrupt or unavailable; the normal checkpoint path still applies.
|
# corrupt or unavailable; the normal checkpoint path still applies.
|
||||||
return False
|
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
|
@staticmethod
|
||||||
def _websocket_route(session: Session) -> tuple[str, str] | None:
|
def _websocket_route(session: Session) -> tuple[str, str] | None:
|
||||||
return RecoveryCoordinator._websocket_route_for(session.key, session.metadata)
|
return RecoveryCoordinator._websocket_route_for(session.key, session.metadata)
|
||||||
|
|||||||
@@ -164,8 +164,8 @@ async def test_dispatch_cancellation_restores_checkpoint():
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_dispatch_cancellation_keeps_checkpoint_for_managed_restart(tmp_path: Path) -> None:
|
async def test_dispatch_cancellation_keeps_checkpoint_for_gateway_shutdown(tmp_path: Path) -> None:
|
||||||
"""A restart preserves the checkpoint; an explicit stop still restores it."""
|
"""Gateway shutdown preserves the checkpoint; an explicit stop restores it."""
|
||||||
loop = _make_loop(tmp_path)
|
loop = _make_loop(tmp_path)
|
||||||
loop.preserve_inflight_turns_on_shutdown()
|
loop.preserve_inflight_turns_on_shutdown()
|
||||||
loop._restore_runtime_checkpoint = MagicMock() # type: ignore[method-assign]
|
loop._restore_runtime_checkpoint = MagicMock() # type: ignore[method-assign]
|
||||||
|
|||||||
@@ -102,6 +102,9 @@ class _GatewayAgentContractStub:
|
|||||||
) -> OutboundMessage | None:
|
) -> OutboundMessage | None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
def preserve_inflight_turns_on_shutdown(self) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
class _EmptyGatewaySessionManager:
|
class _EmptyGatewaySessionManager:
|
||||||
"""Minimal session-manager contract for gateway assembly tests."""
|
"""Minimal session-manager contract for gateway assembly tests."""
|
||||||
@@ -3712,6 +3715,9 @@ def test_gateway_agent_task_owns_initial_mcp_provider_close(
|
|||||||
async def aclose(self) -> None:
|
async def aclose(self) -> None:
|
||||||
seen["agent_closed"] = True
|
seen["agent_closed"] = True
|
||||||
|
|
||||||
|
def preserve_inflight_turns_on_shutdown(self) -> None:
|
||||||
|
seen["inflight_turns_preserved"] = True
|
||||||
|
|
||||||
def stop(self) -> None:
|
def stop(self) -> None:
|
||||||
seen["agent_stopped"] = True
|
seen["agent_stopped"] = True
|
||||||
|
|
||||||
@@ -3793,6 +3799,7 @@ def test_gateway_agent_task_owns_initial_mcp_provider_close(
|
|||||||
|
|
||||||
assert result.exit_code == 0
|
assert result.exit_code == 0
|
||||||
assert seen["agent_stopped"] is True
|
assert seen["agent_stopped"] is True
|
||||||
|
assert seen["inflight_turns_preserved"] is True
|
||||||
assert seen["agent_closed"] is True
|
assert seen["agent_closed"] is True
|
||||||
assert seen["agent_task_cleaned_up"] is True
|
assert seen["agent_task_cleaned_up"] is True
|
||||||
assert seen["channels_stopped"] is True
|
assert seen["channels_stopped"] is True
|
||||||
|
|||||||
@@ -455,8 +455,7 @@ def test_restart_does_not_detach_a_foreground_gateway(tmp_path, monkeypatch):
|
|||||||
assert result.message == "gateway_foreground_restart_required"
|
assert result.message == "gateway_foreground_restart_required"
|
||||||
|
|
||||||
|
|
||||||
def test_restart_marks_the_exiting_gateway_for_turn_recovery(tmp_path, monkeypatch):
|
def test_restart_stops_then_starts_the_background_gateway(tmp_path, monkeypatch):
|
||||||
"""A managed restart must not look like an explicit stop to the child."""
|
|
||||||
runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Linux")
|
runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Linux")
|
||||||
status = GatewayStatus(
|
status = GatewayStatus(
|
||||||
running=True,
|
running=True,
|
||||||
@@ -469,9 +468,6 @@ def test_restart_marks_the_exiting_gateway_for_turn_recovery(tmp_path, monkeypat
|
|||||||
|
|
||||||
def stop(*, timeout_s: int):
|
def stop(*, timeout_s: int):
|
||||||
assert timeout_s == 20
|
assert timeout_s == 20
|
||||||
assert json.loads(runtime._restart_intent_path.read_text(encoding="utf-8")) == {
|
|
||||||
"pid": 12345
|
|
||||||
}
|
|
||||||
return SimpleNamespace(ok=True, message="gateway_stopped", status=status)
|
return SimpleNamespace(ok=True, message="gateway_stopped", status=status)
|
||||||
|
|
||||||
monkeypatch.setattr(runtime, "_stop", stop)
|
monkeypatch.setattr(runtime, "_stop", stop)
|
||||||
@@ -484,17 +480,6 @@ def test_restart_marks_the_exiting_gateway_for_turn_recovery(tmp_path, monkeypat
|
|||||||
result = runtime.restart(GatewayStartOptions(port=18790))
|
result = runtime.restart(GatewayStartOptions(port=18790))
|
||||||
|
|
||||||
assert result.ok is True
|
assert result.ok is True
|
||||||
assert not runtime._restart_intent_path.exists()
|
|
||||||
|
|
||||||
|
|
||||||
def test_restart_intent_only_applies_to_the_recorded_gateway_process(tmp_path):
|
|
||||||
runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Linux")
|
|
||||||
|
|
||||||
runtime._write_restart_intent(os.getpid())
|
|
||||||
|
|
||||||
assert runtime.preserves_inflight_turns_on_exit() is True
|
|
||||||
runtime._write_restart_intent(os.getpid() + 1)
|
|
||||||
assert runtime.preserves_inflight_turns_on_exit() is False
|
|
||||||
|
|
||||||
|
|
||||||
def test_last_interactive_client_stops_an_on_demand_gateway(tmp_path, monkeypatch):
|
def test_last_interactive_client_stops_an_on_demand_gateway(tmp_path, monkeypatch):
|
||||||
|
|||||||
@@ -75,6 +75,47 @@ async def test_stale_incomplete_transcript_waits_for_confirmation(tmp_path: Path
|
|||||||
assert event.status == "awaiting_user"
|
assert event.status == "awaiting_user"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_materialized_interruption_can_continue_from_saved_context(
|
||||||
|
tmp_path: Path,
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
"""Older shutdowns may have cleared the checkpoint after saving partial history."""
|
||||||
|
sessions = SessionManager(tmp_path)
|
||||||
|
session = sessions.get_or_create("websocket:chat")
|
||||||
|
session.messages.extend(
|
||||||
|
[
|
||||||
|
{"role": "user", "content": "research this"},
|
||||||
|
{"role": "assistant", "content": "I will check."},
|
||||||
|
{"role": "tool", "tool_call_id": "search-1", "content": "saved result"},
|
||||||
|
]
|
||||||
|
)
|
||||||
|
_persist(sessions, session)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"nanobot.webui.transcript.has_unfinished_transcript_tail",
|
||||||
|
lambda _key: True,
|
||||||
|
)
|
||||||
|
|
||||||
|
coordinator, bus, restarted = _coordinator(tmp_path)
|
||||||
|
await coordinator.scan()
|
||||||
|
|
||||||
|
state = restarted.get_or_create("websocket:chat").metadata[RECOVERY_METADATA_KEY]
|
||||||
|
assert state["status"] == "awaiting_user"
|
||||||
|
assert state["reason"] == "interrupted_with_saved_context"
|
||||||
|
assert "can_continue" not in state
|
||||||
|
event = bus.outbound.get_nowait().event
|
||||||
|
assert isinstance(event, RecoveryStateEvent)
|
||||||
|
assert event.can_continue is None
|
||||||
|
|
||||||
|
await coordinator.handle_action(
|
||||||
|
"continue",
|
||||||
|
{"chat_id": "chat", "recovery_id": state["recovery_id"]},
|
||||||
|
)
|
||||||
|
|
||||||
|
continuation = bus.inbound.get_nowait()
|
||||||
|
assert continuation.session_key_override == "websocket:chat"
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_transcript_only_interruption_is_discovered_without_materializing_completed_history(
|
async def test_transcript_only_interruption_is_discovered_without_materializing_completed_history(
|
||||||
tmp_path: Path,
|
tmp_path: Path,
|
||||||
|
|||||||
Reference in New Issue
Block a user