Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(runtime): add user-controlled turn recovery

Browse Source
This commit is contained in:
Xubin Ren
2026-08-24 00:58:04 +08:00
parent ffa58aa5ef
commit 12029f8812
60 changed files with 4027 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+103 -132
View File
@@ -79,6 +79,15 @@ from nanobot.session.model_selection import (
SESSION_MODEL_PRESET_METADATA_KEY,
model_preset_from_metadata,
)
from nanobot.session.recovery import (
PENDING_FOLLOWUP_ID_KEY,
RECOVERY_INBOUND_METADATA_KEY,
RecoveryAdmission,
acknowledge_pending_followups,
record_pending_followup,
restore_pending_interruption,
restore_runtime_checkpoint,
)
from nanobot.session.summary import SessionSummary
from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator
from nanobot.utils.cancellation import task_is_cancelling
@@ -291,12 +300,14 @@ class AgentLoop:
restart_mode: str = "auto",
local_trigger_store: LocalTriggerStore | None = None,
idle_compact_check_interval_seconds: int = 0,
recovery_admission: RecoveryAdmission | None = None,
):
from nanobot.config.schema import ToolsConfig
_tc = tools_config or ToolsConfig()
defaults = AgentDefaults()
self.bus = bus
self._recovery_admission = recovery_admission
if turn_delivery_factory is not None:
if turn_delivery_factory.bus is not bus:
raise ValueError("turn delivery factory must use the agent message bus")
@@ -409,6 +420,7 @@ class AgentLoop:
# When a session has an active task, new messages for that session
# are routed here instead of creating a new task.
self._pending_queues: dict[str, asyncio.Queue[InboundMessage]] = {}
self._preserve_inflight_turns_on_shutdown = False
self._deferred_automation_turns: dict[str, list[InboundMessage]] = {}
self._cron_turns = CronTurnCoordinator(
publish_inbound=self.bus.publish_inbound,
@@ -726,6 +738,9 @@ class AgentLoop:
extra[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta
session.add_message("user", text, **extra)
self._mark_pending_user_turn(session)
followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id:
acknowledge_pending_followups(session, [followup_id])
self.sessions.save(session)
return True
return False
@@ -1061,6 +1076,9 @@ class AgentLoop:
row["subagent_task_id"] = task_id
row[HIDDEN_HISTORY_META] = subagent_marker
row["injected_event"] = "subagent_result"
followup_id = metadata.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id:
row[PENDING_FOLLOWUP_ID_KEY] = followup_id
return row
items: list[dict[str, Any]] = []
@@ -1285,6 +1303,17 @@ class AgentLoop:
break
if deferred:
continue
# A newer WebUI message must supersede an explicit recovery
# before it is injected into that recovery's pending queue.
# Without this admission point, a recovered turn could finish
# first and only then observe the user's newer request.
if (
effective_key in self._pending_queues
and msg.channel == "websocket"
and self._recovery_admission is not None
and not await self._recovery_admission.admit(msg)
):
continue
# If this session already has an active pending queue (i.e. a task
# is processing this session), route the message there for mid-turn
# injection instead of creating a competing task.
@@ -1303,6 +1332,17 @@ class AgentLoop:
msg,
session_key_override=effective_key,
)
session = self.sessions.get_or_create(effective_key)
followup_id = record_pending_followup(session, pending_msg)
if followup_id is not None:
pending_msg = dataclasses.replace(
pending_msg,
metadata={
**pending_msg.metadata,
PENDING_FOLLOWUP_ID_KEY: followup_id,
},
)
self.sessions.save(session)
try:
self._pending_queues[effective_key].put_nowait(pending_msg)
except asyncio.QueueFull:
@@ -1310,6 +1350,7 @@ class AgentLoop:
"Pending queue full for session {}, falling back to queued task",
effective_key,
)
msg = pending_msg
else:
logger.info(
"Routed follow-up message to pending queue for session {}",
@@ -1319,17 +1360,45 @@ class AgentLoop:
# Compute the effective session key before dispatching
# This ensures /stop command can find tasks correctly when unified session is enabled
task = asyncio.create_task(self._dispatch(msg))
active_tasks = self._active_tasks.setdefault(effective_key, set())
active_tasks: set[asyncio.Task[Any]] = self._active_tasks.setdefault(
effective_key,
set(),
)
active_tasks.add(task)
task.add_done_callback(active_tasks.discard)
finally:
await self.aclose()
def preserve_inflight_turns_on_shutdown(self) -> None:
"""Keep durable checkpoints when the owning gateway is restarting.
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.
"""
self._preserve_inflight_turns_on_shutdown = True
async def _dispatch(self, msg: InboundMessage) -> None:
"""Process a message: per-session serial, cross-session concurrent."""
session_key = self._effective_session_key(msg)
if session_key != msg.session_key:
msg = dataclasses.replace(msg, session_key_override=session_key)
recovery_task_registered = False
recovery_admission = self._recovery_admission
current_task: asyncio.Task[Any] | None = None
if recovery_admission is not None:
recovery_id = msg.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
if isinstance(recovery_id, str) and recovery_id:
current_task = asyncio.current_task()
if current_task is not None:
recovery_admission.register_recovery_task(session_key, current_task)
recovery_task_registered = True
if not await recovery_admission.admit(msg):
logger.info("Skipped stale recovery for session {}", session_key)
if recovery_task_registered and current_task is not None:
recovery_admission.unregister_recovery_task(session_key, current_task)
return
lock = self._get_session_lock(session_key)
gate = self._concurrency_gate or nullcontext()
@@ -1380,7 +1449,10 @@ class AgentLoop:
# _emit_checkpoint during tool execution; materializing
# it into session history now makes it visible in the
# next conversation turn.
if session_key in self._discarding_sessions:
if (
session_key in self._discarding_sessions
or self._preserve_inflight_turns_on_shutdown
):
raise
try:
key = self._effective_session_key(msg)
@@ -1437,6 +1509,12 @@ class AgentLoop:
await delivery.idle()
await self._publish_next_deferred_automation_turn(session_key)
finally:
if (
recovery_task_registered
and current_task is not None
and recovery_admission is not None
):
recovery_admission.unregister_recovery_task(session_key, current_task)
if pending is None:
await delivery.idle()
await self._publish_next_deferred_automation_turn(session_key)
@@ -1738,7 +1816,10 @@ class AgentLoop:
if self._restore_runtime_checkpoint(session):
self.sessions.save(session)
if self._restore_pending_user_turn(session):
if (
RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
and restore_pending_interruption(session)
):
self.sessions.save(session)
async def _compact_session(self, ctx: TurnContext) -> None:
@@ -2093,8 +2174,21 @@ class AgentLoop:
if m.get("role") == "tool" and m.get("tool_call_id")
}
last_assistant_idx: int | None = None
saved_followup_ids: set[str] = set()
for m in messages[skip:]:
entry = dict(m)
followup_id_value = cast(object, entry.pop(PENDING_FOLLOWUP_ID_KEY, None))
followup_ids = (
[followup_id_value]
if isinstance(followup_id_value, str)
else [
followup_id
for followup_id in cast(list[object], followup_id_value)
if isinstance(followup_id, str)
]
if isinstance(followup_id_value, list)
else []
)
internal_meta = cast(object, entry.pop("_meta", None))
runtime_context_meta = (
cast(dict[str, Any], internal_meta).get(
@@ -2147,6 +2241,8 @@ class AgentLoop:
entry[RUNTIME_CONTEXT_HISTORY_META] = runtime_context_meta
entry.setdefault("timestamp", datetime.now().isoformat())
session.messages.append(entry)
if role == "user":
saved_followup_ids.update(followup_id for followup_id in followup_ids if followup_id)
if role == "assistant":
last_assistant_idx = len(session.messages) - 1
declared_tool_call_ids.update(
@@ -2161,6 +2257,8 @@ class AgentLoop:
)
if turn_latency_ms is not None and last_assistant_idx is not None:
session.messages[last_assistant_idx]["latency_ms"] = int(turn_latency_ms)
if saved_followup_ids:
acknowledge_pending_followups(session, saved_followup_ids)
session.updated_at = datetime.now()
def _persist_subagent_followup(self, session: Session, msg: InboundMessage) -> bool:
@@ -2195,7 +2293,7 @@ class AgentLoop:
def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None:
"""Persist the latest in-flight turn state into session metadata."""
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
self.sessions.save(session)
self.sessions.save_runtime_checkpoint(session)
def _mark_pending_user_turn(self, session: Session) -> None:
session.metadata[self._PENDING_USER_TURN_KEY] = True
@@ -2207,136 +2305,9 @@ class AgentLoop:
if self._RUNTIME_CHECKPOINT_KEY in session.metadata:
session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None)
@staticmethod
def _checkpoint_message_key(message: dict[str, Any]) -> tuple[Any, ...]:
return (
message.get("role"),
message.get("content"),
message.get("tool_call_id"),
message.get("name"),
message.get("tool_calls"),
message.get("reasoning_content"),
message.get("thinking_blocks"),
)
def _restore_runtime_checkpoint(self, session: Session) -> bool:
"""Materialize an unfinished turn into session history before a new request."""
from datetime import datetime
checkpoint = cast(
object,
session.metadata.get(self._RUNTIME_CHECKPOINT_KEY),
)
if not isinstance(checkpoint, dict):
return False
checkpoint_data = cast(dict[str, Any], checkpoint)
assistant_message = cast(object, checkpoint_data.get("assistant_message"))
completed_tool_results = cast(
Iterable[object],
checkpoint_data.get("completed_tool_results") or [],
)
pending_tool_calls = cast(
Iterable[object],
checkpoint_data.get("pending_tool_calls") or [],
)
restored_messages: list[dict[str, Any]] = []
if isinstance(assistant_message, dict):
restored = dict(cast(dict[str, Any], assistant_message))
restored.setdefault("timestamp", datetime.now().isoformat())
restored_messages.append(restored)
for message in completed_tool_results:
if isinstance(message, dict):
restored = dict(cast(dict[str, Any], message))
restored.setdefault("timestamp", datetime.now().isoformat())
restored_messages.append(restored)
for tool_call in pending_tool_calls:
if not isinstance(tool_call, dict):
continue
tool_call_data = cast(dict[str, Any], tool_call)
tool_id = tool_call_data.get("id")
function_data = cast(
dict[str, Any],
tool_call_data.get("function") or {},
)
name = function_data.get("name") or "tool"
restored_messages.append(
{
"role": "tool",
"tool_call_id": tool_id,
"name": name,
"content": "Error: Task interrupted before this tool finished.",
"timestamp": datetime.now().isoformat(),
}
)
overlap = 0
max_overlap = min(len(session.messages), len(restored_messages))
for size in range(max_overlap, 0, -1):
existing = session.messages[-size:]
restored = restored_messages[:size]
if all(
self._checkpoint_message_key(left) == self._checkpoint_message_key(right)
for left, right in zip(existing, restored)
):
overlap = size
break
appended_messages = restored_messages[overlap:]
session.messages.extend(appended_messages)
assistant_message_data = (
cast(dict[str, Any], assistant_message)
if isinstance(assistant_message, dict)
else None
)
provider_state_is_synchronized = (
checkpoint_data.get(self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY)
== self._PROVIDER_STATE_CHECKPOINT_VERSION
)
phase = checkpoint_data.get("phase")
exact_final_response = (
phase == "final_response"
and assistant_message_data is not None
and assistant_message_data.get("role") == "assistant"
and not bool(checkpoint_data.get("completed_tool_results"))
and not bool(checkpoint_data.get("pending_tool_calls"))
)
exact_completed_tools = (
phase == "tools_completed"
and assistant_message_data is not None
and assistant_message_data.get("role") == "assistant"
and not bool(checkpoint_data.get("pending_tool_calls"))
)
if not (
provider_state_is_synchronized
and (exact_final_response or exact_completed_tools)
):
session.provider_state = None
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
return True
def _restore_pending_user_turn(self, session: Session) -> bool:
"""Close a turn that only persisted the user message before crashing."""
from datetime import datetime
if not session.metadata.get(self._PENDING_USER_TURN_KEY):
return False
if session.messages and session.messages[-1].get("role") == "user":
session.messages.append(
{
"role": "assistant",
"content": "Error: Task interrupted before a response was generated.",
"timestamp": datetime.now().isoformat(),
}
)
session.provider_state = None
session.updated_at = datetime.now()
self._clear_pending_user_turn(session)
return True
return restore_runtime_checkpoint(session)
async def process_direct(
self,
nanobot/agent/runner.py
+18
View File
@@ -37,6 +37,7 @@ from nanobot.runtime_context import (
reattach_runtime_context,
)
from nanobot.session.history_visibility import is_hidden_history_message
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY
from nanobot.utils.helpers import (
IncrementalThinkExtractor,
build_assistant_message,
@@ -234,6 +235,23 @@ class AgentRunner:
merged.get("content"),
injection.get("content"),
)
followup_id = injection.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id:
existing = cast(object, merged.get(PENDING_FOLLOWUP_ID_KEY))
followup_ids = (
[existing]
if isinstance(existing, str)
else [
item
for item in cast(list[object], existing)
if isinstance(item, str)
]
if isinstance(existing, list)
else []
)
if followup_id not in followup_ids:
followup_ids.append(followup_id)
merged[PENDING_FOLLOWUP_ID_KEY] = followup_ids
messages[-1] = merged
continue
messages.append(injection)