mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-09-04 02:01:48 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(agent): make checkpoint recovery ownership explicit

Browse Source
This commit is contained in:
chengyongru
2026-08-26 16:33:59 +08:00
committed by chengyongru
parent 9f5a56f1ec
commit 56aa7296f3
4 changed files with 47 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/agent/test_loop_save_turn.py
+23 -25
View File
@@ -37,7 +37,14 @@ from nanobot.session.keys import (
UNIFIED_SESSION_KEY,
)
from nanobot.session.manager import Session
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY, PENDING_FOLLOWUPS_KEY
from nanobot.session.recovery import (
PENDING_FOLLOWUP_ID_KEY,
PENDING_FOLLOWUPS_KEY,
PROVIDER_STATE_CHECKPOINT_VERSION,
PROVIDER_STATE_CHECKPOINT_VERSION_KEY,
RUNTIME_CHECKPOINT_KEY,
restore_runtime_checkpoint,
)
from nanobot.session.turn_continuation import (
INTERNAL_CONTINUATION_META,
INTERNAL_CONTINUATION_RUN_STARTED_AT_META,
@@ -682,12 +689,11 @@ def test_save_turn_stamps_latency_on_last_assistant() -> None:
def test_restore_runtime_checkpoint_rehydrates_completed_and_pending_tools() -> None:
loop = _mk_loop()
session = Session(
key="test:checkpoint",
provider_state=_provider_state(),
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"assistant_message": {
"role": "assistant",
"content": "working",
@@ -723,10 +729,10 @@ def test_restore_runtime_checkpoint_rehydrates_completed_and_pending_tools() ->
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is None
assert session.metadata.get(RUNTIME_CHECKPOINT_KEY) is None
assert session.messages[0]["role"] == "assistant"
assert session.messages[1]["tool_call_id"] == "call_done"
assert session.messages[2]["tool_call_id"] == "call_pending"
@@ -735,17 +741,14 @@ def test_restore_runtime_checkpoint_rehydrates_completed_and_pending_tools() ->
def test_restore_final_response_checkpoint_preserves_matching_provider_state() -> None:
loop = _mk_loop()
state = _provider_state()
session = Session(
key="test:final-checkpoint",
provider_state=state,
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"phase": "final_response",
AgentLoop._PROVIDER_STATE_CHECKPOINT_VERSION_KEY: (
AgentLoop._PROVIDER_STATE_CHECKPOINT_VERSION
),
PROVIDER_STATE_CHECKPOINT_VERSION_KEY: PROVIDER_STATE_CHECKPOINT_VERSION,
"assistant_message": {
"role": "assistant",
"content": "finished",
@@ -756,21 +759,20 @@ def test_restore_final_response_checkpoint_preserves_matching_provider_state() -
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.messages[-1]["content"] == "finished"
assert session.provider_state is state
assert session.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is None
assert session.metadata.get(RUNTIME_CHECKPOINT_KEY) is None
def test_restore_legacy_final_checkpoint_discards_unproven_provider_state() -> None:
loop = _mk_loop()
session = Session(
key="test:legacy-final-checkpoint",
provider_state=_provider_state(),
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"phase": "final_response",
"assistant_message": {
"role": "assistant",
@@ -782,7 +784,7 @@ def test_restore_legacy_final_checkpoint_discards_unproven_provider_state() -> N
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.messages[-1]["content"] == "finished"
@@ -790,7 +792,6 @@ def test_restore_legacy_final_checkpoint_discards_unproven_provider_state() -> N
def test_restore_completed_tools_checkpoint_preserves_matching_provider_state() -> None:
loop = _mk_loop()
tool_result = {
"role": "tool",
"tool_call_id": "call_done",
@@ -802,11 +803,9 @@ def test_restore_completed_tools_checkpoint_preserves_matching_provider_state()
key="test:completed-tools-checkpoint",
provider_state=state,
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"phase": "tools_completed",
AgentLoop._PROVIDER_STATE_CHECKPOINT_VERSION_KEY: (
AgentLoop._PROVIDER_STATE_CHECKPOINT_VERSION
),
PROVIDER_STATE_CHECKPOINT_VERSION_KEY: PROVIDER_STATE_CHECKPOINT_VERSION,
"assistant_message": {
"role": "assistant",
"content": None,
@@ -824,7 +823,7 @@ def test_restore_completed_tools_checkpoint_preserves_matching_provider_state()
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.messages[-1]["content"] == "compacted result"
@@ -832,7 +831,6 @@ def test_restore_completed_tools_checkpoint_preserves_matching_provider_state()
def test_restore_runtime_checkpoint_dedupes_overlapping_tail() -> None:
loop = _mk_loop()
session = Session(
key="test:checkpoint-overlap",
messages=[
@@ -860,7 +858,7 @@ def test_restore_runtime_checkpoint_dedupes_overlapping_tail() -> None:
},
],
metadata={
AgentLoop._RUNTIME_CHECKPOINT_KEY: {
RUNTIME_CHECKPOINT_KEY: {
"assistant_message": {
"role": "assistant",
"content": "working",
@@ -896,10 +894,10 @@ def test_restore_runtime_checkpoint_dedupes_overlapping_tail() -> None:
},
)
restored = loop._restore_runtime_checkpoint(session)
restored = restore_runtime_checkpoint(session)
assert restored is True
assert session.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is None
assert session.metadata.get(RUNTIME_CHECKPOINT_KEY) is None
assert len(session.messages) == 3
assert session.messages[0]["role"] == "assistant"
assert session.messages[1]["tool_call_id"] == "call_done"
tests/agent/test_stop_preserves_context.py
+19 -51
View File
@@ -18,6 +18,7 @@ import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.session.recovery import RUNTIME_CHECKPOINT_KEY
def _make_provider():
@@ -41,51 +42,6 @@ def _make_loop(tmp_path: Path) -> AgentLoop:
return AgentLoop(bus=bus, provider=provider, workspace=tmp_path)
class TestStopPreservesContext:
"""Verify that /stop restores partial context via checkpoint."""
def test_restore_checkpoint_method_exists(self, tmp_path):
"""AgentLoop should have _restore_runtime_checkpoint."""
loop = _make_loop(tmp_path)
assert hasattr(loop, "_restore_runtime_checkpoint")
def test_checkpoint_key_constant(self, tmp_path):
"""The runtime checkpoint key should be defined."""
loop = _make_loop(tmp_path)
assert loop._RUNTIME_CHECKPOINT_KEY == "runtime_checkpoint"
def test_cancel_dispatch_restores_checkpoint(self, tmp_path):
"""When a task is cancelled, the checkpoint should be restored."""
loop = _make_loop(tmp_path)
session = MagicMock()
session.metadata = {
"runtime_checkpoint": {
"phase": "awaiting_tools",
"iteration": 0,
"assistant_message": {
"role": "assistant",
"content": "Let me search for that.",
"tool_calls": [{"id": "tc_1", "type": "function",
"function": {"name": "web_search", "arguments": "{}"}}],
},
"completed_tool_results": [
{"role": "tool", "tool_call_id": "tc_1",
"content": "Search results: ..."},
],
"pending_tool_calls": [],
}
}
session.messages = [
{"role": "user", "content": "Search for something"},
]
loop.sessions.get_or_create.return_value = session
restored = loop._restore_runtime_checkpoint(session)
assert restored is True
assert len(session.messages) > 1
assert "runtime_checkpoint" not in session.metadata
@pytest.mark.asyncio
async def test_dispatch_cancellation_restores_checkpoint():
"""Regression for #2966: /stop interrupting _dispatch must materialize the
@@ -93,9 +49,8 @@ async def test_dispatch_cancellation_restores_checkpoint():
unwinds, so the next turn can see the partial work.
This exercises the real _dispatch path (locks, pending queues, the
CancelledError handler) rather than poking _restore_runtime_checkpoint in
isolation, so a future refactor that drops the cancel-time restore is
caught by CI instead of silently regressing.
CancelledError handler), so a future refactor that drops the cancel-time
restore is caught by CI instead of silently regressing.
"""
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
@@ -112,7 +67,7 @@ async def test_dispatch_cancellation_restores_checkpoint():
mock_subagent_manager.return_value.cancel_by_session = AsyncMock(return_value=0)
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
checkpoint_key = loop._RUNTIME_CHECKPOINT_KEY
checkpoint_key = RUNTIME_CHECKPOINT_KEY
session = SimpleNamespace(
key="test:c1",
metadata={
@@ -168,7 +123,19 @@ async def test_dispatch_cancellation_keeps_checkpoint_for_gateway_shutdown(tmp_p
"""Gateway shutdown preserves the checkpoint; an explicit stop restores it."""
loop = _make_loop(tmp_path)
loop.preserve_inflight_turns_on_shutdown()
loop._restore_runtime_checkpoint = MagicMock() # type: ignore[method-assign]
checkpoint_key = RUNTIME_CHECKPOINT_KEY
checkpoint = {
"phase": "final_response",
"assistant_message": {"role": "assistant", "content": "finished"},
"completed_tool_results": [],
"pending_tool_calls": [],
}
session = SimpleNamespace(
metadata={checkpoint_key: checkpoint},
messages=[],
provider_state=None,
)
loop.sessions.get_or_create.return_value = session
async def _cancel(*_args: object, **_kwargs: object) -> None:
raise asyncio.CancelledError()
@@ -182,4 +149,5 @@ async def test_dispatch_cancellation_keeps_checkpoint_for_gateway_shutdown(tmp_p
InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="work")
)
loop._restore_runtime_checkpoint.assert_not_called()
assert session.metadata[checkpoint_key] == checkpoint
assert session.messages == []