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

fix(agent): extend sustained goal iteration budget

Browse Source
This commit is contained in:
Xubin Ren
2026-06-01 04:00:15 +08:00
parent cba9ff1f57
commit be2e0172d1
11 changed files with 783 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+50 -11
View File
@@ -45,6 +45,7 @@ from nanobot.session.goal_state import (
sustained_goal_active,
)
from nanobot.session.manager import Session, SessionManager
from nanobot.session import turn_continuation
from nanobot.session.webui_turns import (
WebuiTurnCoordinator,
build_bus_progress_callback,
@@ -112,6 +113,7 @@ class TurnContext:
save_skip: int = 0
outbound: OutboundMessage | None = None
suppress_response: bool = False
on_progress: Callable[..., Awaitable[None]] | None = None
on_stream: Callable[[str], Awaitable[None]] | None = None
@@ -121,6 +123,7 @@ class TurnContext:
pending_queue: asyncio.Queue | None = None
pending_summary: str | None = None
turn_wall_started_at: float = field(default_factory=time.time)
visible_run_started_at: float | None = None
turn_latency_ms: int | None = None
trace: list[StateTraceEntry] = field(default_factory=list)
@@ -565,6 +568,8 @@ class AgentLoop:
Returns True if the message was persisted.
"""
if not turn_continuation.should_persist_user_message(msg.metadata):
return False
media_paths = [p for p in (msg.media or []) if isinstance(p, str) and p]
has_text = isinstance(msg.content, str) and msg.content.strip()
if has_text or media_paths:
@@ -771,6 +776,7 @@ class AgentLoop:
+ "\n\nPlease continue working toward the objective using your tools, "
"or call complete_goal if the work is truly finished."
) if _goal_lines else SUSTAINED_GOAL_CONTINUE_PROMPT
session_metadata = session.metadata if session is not None else None
try:
result = await self.runner.run(AgentRunSpec(
initial_messages=initial_messages,
@@ -796,7 +802,8 @@ class AgentLoop:
llm_timeout_s=runner_wall_llm_timeout_s(
self.sessions,
session.key if session is not None else session_key,
metadata=(session.metadata if session is not None else None),
metadata=session_metadata,
message_metadata=metadata,
),
goal_active_predicate=lambda: sustained_goal_active(session.metadata) if session is not None else False,
goal_continue_message=_goal_continue,
@@ -808,9 +815,15 @@ class AgentLoop:
self._last_usage = result.usage
if result.stop_reason == "max_iterations":
logger.warning("Max iterations ({}) reached", self.max_iterations)
should_stream = turn_continuation.should_stream_budget_response(
stop_reason=result.stop_reason,
pending_queue_available=pending_queue is not None and session is not None,
session_metadata=session_metadata,
message_metadata=metadata,
)
# Push final content through stream so streaming channels (e.g. Feishu)
# update the card instead of leaving it empty.
if on_stream and on_stream_end:
if on_stream and on_stream_end and should_stream:
await on_stream(result.final_content or "")
await on_stream_end(resuming=False)
elif result.stop_reason == "error":
@@ -953,7 +966,8 @@ class AgentLoop:
channel=msg.channel, chat_id=msg.chat_id,
content="", metadata=msg.metadata or {},
))
if msg.channel == "websocket":
continuing = turn_continuation.internal_continuation_pending(msg.metadata)
if msg.channel == "websocket" and not continuing:
turn_lat = self._pending_turn_latency_ms.pop(session_key, None)
await self._webui_turns.handle_turn_end(
msg,
@@ -1017,9 +1031,10 @@ class AgentLoop:
"Re-published {} leftover message(s) to bus for session {}",
leftover, session_key,
)
await self._webui_turns.publish_run_status(msg, "idle")
self._pending_turn_latency_ms.pop(session_key, None)
self._webui_turns.discard(session_key)
if not turn_continuation.internal_continuation_pending(msg.metadata):
await self._webui_turns.publish_run_status(msg, "idle")
self._pending_turn_latency_ms.pop(session_key, None)
self._webui_turns.discard(session_key)
finally:
if pending is None:
await self._webui_turns.publish_run_status(msg, "idle")
@@ -1167,12 +1182,17 @@ class AgentLoop:
)
key = session_key or msg.session_key
t0 = time.time()
ctx = TurnContext(
msg=msg,
session=None,
session_key=key,
state=TurnState.RESTORE,
turn_id=f"{key}:{time.time_ns()}",
turn_wall_started_at=t0,
visible_run_started_at=turn_continuation.internal_continuation_run_started_at(
msg.metadata,
),
on_progress=on_progress,
on_stream=on_stream,
on_stream_end=on_stream_end,
@@ -1378,7 +1398,13 @@ class AgentLoop:
return "ok"
async def _state_run(self, ctx: TurnContext) -> str:
await self._webui_turns.publish_run_status(ctx.msg, "running")
if ctx.visible_run_started_at is None:
ctx.visible_run_started_at = time.time()
await self._webui_turns.publish_run_status(
ctx.msg,
"running",
started_at=ctx.visible_run_started_at,
)
result = await self._run_agent_loop(
ctx.initial_messages,
on_progress=ctx.on_progress,
@@ -1399,15 +1425,25 @@ class AgentLoop:
ctx.all_messages = all_msgs
ctx.stop_reason = stop_reason
ctx.had_injections = had_injections
await turn_continuation.maybe_continue_turn(ctx)
return "ok"
async def _state_save(self, ctx: TurnContext) -> str:
if ctx.final_content is None or not ctx.final_content.strip():
turn_continuation.prepare_save_boundary(ctx)
if (
(ctx.final_content is None or not ctx.final_content.strip())
and not ctx.suppress_response
):
ctx.final_content = EMPTY_FINAL_RESPONSE_MESSAGE
ctx.save_skip = 1 + len(ctx.history) + (1 if ctx.user_persisted_early else 0)
ctx.turn_latency_ms = max(0, int((time.time() - ctx.turn_wall_started_at) * 1000))
latency_started_at = (
ctx.visible_run_started_at
if turn_continuation.internal_continuation_inbound(ctx.msg.metadata)
and ctx.visible_run_started_at is not None
else ctx.turn_wall_started_at
)
ctx.turn_latency_ms = max(0, int((time.time() - latency_started_at) * 1000))
self._save_turn(
ctx.session, ctx.all_messages, ctx.save_skip,
turn_latency_ms=ctx.turn_latency_ms,
@@ -1427,6 +1463,9 @@ class AgentLoop:
return "ok"
async def _state_respond(self, ctx: TurnContext) -> str:
if ctx.suppress_response:
ctx.outbound = None
return "ok"
ctx.outbound = self._assemble_outbound(
ctx.msg,
ctx.final_content,
nanobot/session/goal_state.py
+19 -4
View File
@@ -43,6 +43,19 @@ def sustained_goal_active(metadata: Mapping[str, Any] | None) -> bool:
return isinstance(goal, dict) and goal.get("status") == "active"
def sustained_goal_turn(
metadata: Mapping[str, Any] | None,
*,
message_metadata: Mapping[str, Any] | None = None,
) -> bool:
"""True when this turn should use sustained-goal runtime limits."""
if sustained_goal_active(metadata):
return True
if not message_metadata:
return False
return str(message_metadata.get("original_command") or "").strip() == "/goal"
def parse_goal_state(blob: Any) -> dict[str, Any] | None:
if blob is None:
return None
@@ -98,14 +111,16 @@ def runner_wall_llm_timeout_s(
session_key: str | None,
*,
metadata: Mapping[str, Any] | None = None,
message_metadata: Mapping[str, Any] | None = None,
) -> float | None:
"""Wall-clock cap for :class:`~nanobot.agent.runner.AgentRunner` when streaming an LLM.
Returns ``0.0`` to disable ``asyncio.wait_for`` around the request when a sustained goal is
active; ``None`` means use ``NANOBOT_LLM_TIMEOUT_S``. Pass in-memory ``metadata`` when the
caller already holds :attr:`~nanobot.session.manager.Session.metadata` for this turn.
Returns ``0.0`` to disable ``asyncio.wait_for`` around the request when this is a
sustained-goal turn; ``None`` means use ``NANOBOT_LLM_TIMEOUT_S``. Pass in-memory
``metadata`` when the caller already holds :attr:`~nanobot.session.manager.Session.metadata`
for this turn.
"""
meta: Mapping[str, Any] | None = metadata
if meta is None and session_key:
meta = sessions.get_or_create(session_key).metadata
return 0.0 if sustained_goal_active(meta) else None
return 0.0 if sustained_goal_turn(meta, message_metadata=message_metadata) else None
nanobot/session/turn_continuation.py
+240
View File
@@ -0,0 +1,240 @@
"""Internal turn continuation helpers.
This module keeps budget-boundary continuation policy out of ``AgentLoop``.
The loop calls a small set of helpers; those helpers decide whether an internal
continuation is allowed and, when it is, queue the next turn directly.
"""
from __future__ import annotations
import dataclasses
from typing import Any, Mapping, MutableMapping
from loguru import logger
from nanobot.session.goal_state import (
goal_state_runtime_lines,
sustained_goal_active,
sustained_goal_turn,
)
INTERNAL_CONTINUATION_META = "_internal_continuation"
INTERNAL_CONTINUATION_KIND_META = "_internal_continuation_kind"
INTERNAL_CONTINUATION_PENDING_META = "_internal_continuation_pending"
INTERNAL_CONTINUATION_RUN_STARTED_AT_META = "_internal_continuation_run_started_at"
_GOAL_CONTINUATION_KIND = "sustained_goal"
_GOAL_CONTINUATION_SENDER = "system:continuation"
_GOAL_CONTINUATION_ROUNDS_KEY = "_sustained_goal_continuation_rounds"
_MAX_GOAL_CONTINUATION_ROUNDS = 12
_STRIPPED_INBOUND_META_KEYS = {
"_stream_id",
"_stream_delta",
"_stream_end",
"_resuming",
INTERNAL_CONTINUATION_PENDING_META,
}
def internal_continuation_inbound(metadata: Mapping[str, Any] | None) -> bool:
"""True for an inbound message created by an internal continuation policy."""
return bool(metadata and metadata.get(INTERNAL_CONTINUATION_META) is True)
def internal_continuation_pending(metadata: Mapping[str, Any] | None) -> bool:
"""True when the current turn scheduled an invisible continuation slice."""
return bool(metadata and metadata.get(INTERNAL_CONTINUATION_PENDING_META) is True)
def internal_continuation_run_started_at(metadata: Mapping[str, Any] | None) -> float | None:
"""Return the user-visible run start propagated across continuation slices."""
if not metadata:
return None
value = metadata.get(INTERNAL_CONTINUATION_RUN_STARTED_AT_META)
if not isinstance(value, int | float):
return None
started_at = float(value)
return started_at if started_at > 0 else None
def should_persist_user_message(metadata: Mapping[str, Any] | None) -> bool:
"""Return whether this inbound message should be persisted as user input."""
return not internal_continuation_inbound(metadata)
def should_stream_budget_response(
*,
stop_reason: str,
pending_queue_available: bool,
session_metadata: Mapping[str, Any] | None,
message_metadata: Mapping[str, Any] | None = None,
) -> bool:
"""Return whether the budget-boundary response should be sent to the user."""
return not _continuation_available(
stop_reason=stop_reason,
pending_queue_available=pending_queue_available,
session_metadata=session_metadata,
message_metadata=message_metadata,
)
async def maybe_continue_turn(ctx: Any) -> bool:
"""Queue an internal continuation for *ctx* when policy allows it."""
if ctx.session is None or ctx.pending_queue is None:
return False
if not _continuation_available(
stop_reason=ctx.stop_reason,
pending_queue_available=True,
session_metadata=ctx.session.metadata,
message_metadata=ctx.msg.metadata,
):
return False
metadata = _internal_continuation_metadata(
ctx.msg.metadata,
run_started_at=getattr(ctx, "visible_run_started_at", None),
)
content = _goal_continuation_prompt(ctx.session.metadata)
messages = _strip_terminal_assistant(ctx.all_messages, ctx.final_content)
_increment_goal_continuation_round(ctx.session.metadata)
logger.info("Turn budget reached; scheduling internal continuation")
ctx.msg.metadata[INTERNAL_CONTINUATION_PENDING_META] = True
ctx.final_content = ""
ctx.all_messages = messages
ctx.suppress_response = True
await ctx.pending_queue.put(
dataclasses.replace(
ctx.msg,
sender_id=_GOAL_CONTINUATION_SENDER,
content=content,
media=[],
metadata=metadata,
session_key_override=ctx.session_key,
)
)
return True
def prepare_save_boundary(ctx: Any) -> None:
"""Prepare continuation bookkeeping and the history append boundary."""
if ctx.session is not None:
clear_internal_continuation_state(ctx.session.metadata)
ctx.save_skip = _save_skip_for_turn(
message_metadata=ctx.msg.metadata,
initial_message_count=len(ctx.initial_messages),
history_count=len(ctx.history),
user_persisted_early=ctx.user_persisted_early,
)
def _continuation_available(
*,
stop_reason: str,
pending_queue_available: bool,
session_metadata: Mapping[str, Any] | None,
message_metadata: Mapping[str, Any] | None = None,
) -> bool:
if stop_reason != "max_iterations" or not pending_queue_available:
return False
return _goal_continuation_available(
session_metadata,
message_metadata=message_metadata,
)
def clear_internal_continuation_state(metadata: MutableMapping[str, Any]) -> None:
"""Reset policy bookkeeping once its owning runtime mode is inactive."""
if not sustained_goal_active(metadata):
metadata.pop(_GOAL_CONTINUATION_ROUNDS_KEY, None)
def _save_skip_for_turn(
*,
message_metadata: Mapping[str, Any] | None,
initial_message_count: int,
history_count: int,
user_persisted_early: bool,
) -> int:
"""Return the persisted-message append boundary for this turn."""
if internal_continuation_inbound(message_metadata):
return initial_message_count
return 1 + history_count + (1 if user_persisted_early else 0)
def _goal_continuation_available(
session_metadata: Mapping[str, Any] | None,
*,
message_metadata: Mapping[str, Any] | None = None,
max_rounds: int = _MAX_GOAL_CONTINUATION_ROUNDS,
) -> bool:
if not sustained_goal_turn(session_metadata, message_metadata=message_metadata):
return False
if not sustained_goal_active(session_metadata):
return False
try:
rounds = int((session_metadata or {}).get(_GOAL_CONTINUATION_ROUNDS_KEY) or 0)
except (TypeError, ValueError):
rounds = 0
return rounds < max(0, max_rounds)
def _increment_goal_continuation_round(session_metadata: MutableMapping[str, Any]) -> None:
try:
rounds = int(session_metadata.get(_GOAL_CONTINUATION_ROUNDS_KEY) or 0)
except (TypeError, ValueError):
rounds = 0
session_metadata[_GOAL_CONTINUATION_ROUNDS_KEY] = rounds + 1
def _internal_continuation_metadata(
message_metadata: Mapping[str, Any] | None,
*,
run_started_at: float | None = None,
) -> dict[str, Any]:
metadata = dict(message_metadata or {})
metadata[INTERNAL_CONTINUATION_META] = True
metadata[INTERNAL_CONTINUATION_KIND_META] = _GOAL_CONTINUATION_KIND
if run_started_at is not None:
metadata[INTERNAL_CONTINUATION_RUN_STARTED_AT_META] = float(run_started_at)
for key in _STRIPPED_INBOUND_META_KEYS:
metadata.pop(key, None)
return metadata
def _goal_continuation_prompt(metadata: Mapping[str, Any] | None) -> str:
lines = goal_state_runtime_lines(metadata)
if lines:
goal = "\n".join(lines)
return (
"Continue the active sustained goal after the previous turn reached "
"its tool-call budget.\n\n"
f"{goal}\n\n"
"Continue from the saved context. Do not mention the continuation "
"boundary to the user. Use tools as needed, and call complete_goal "
"when the objective is truly finished."
)
return (
"Continue the active sustained goal after the previous turn reached "
"its tool-call budget. Continue from the saved context. Do not mention "
"the continuation boundary to the user. Use tools as needed, and call "
"complete_goal when the objective is truly finished."
)
def _strip_terminal_assistant(
messages: list[dict[str, Any]],
final_content: str | None,
) -> list[dict[str, Any]]:
"""Drop the synthetic max-iteration assistant message before saving history."""
if not messages:
return messages
last = messages[-1]
if last.get("role") != "assistant":
return messages
if final_content is None or last.get("content") != final_content:
return messages
if last.get("tool_calls"):
return messages
return messages[:-1]
nanobot/session/webui_turns.py
+19 -4
View File
@@ -178,7 +178,13 @@ def websocket_turn_wall_started_at(chat_id: str) -> float | None:
return _WEBSOCKET_TURN_WALL_STARTED_AT.get(chat_id)
async def publish_turn_run_status(bus: MessageBus, msg: InboundMessage, status: str) -> None:
async def publish_turn_run_status(
bus: MessageBus,
msg: InboundMessage,
status: str,
*,
started_at: float | None = None,
) -> None:
"""Notify WebSocket clients while a user turn is executing (timing strip)."""
if msg.channel != "websocket":
return
@@ -189,7 +195,10 @@ async def publish_turn_run_status(bus: MessageBus, msg: InboundMessage, status:
"goal_status": status,
}
if status == "running":
t0 = time.time()
if isinstance(started_at, int | float) and started_at > 0:
t0 = float(started_at)
else:
t0 = time.time()
meta["started_at"] = t0
_WEBSOCKET_TURN_WALL_STARTED_AT[cid] = t0
else:
@@ -300,8 +309,14 @@ class WebuiTurnCoordinator:
def discard(self, session_key: str) -> None:
self._title_contexts.pop(session_key, None)
async def publish_run_status(self, msg: InboundMessage, status: str) -> None:
await publish_turn_run_status(self.bus, msg, status)
async def publish_run_status(
self,
msg: InboundMessage,
status: str,
*,
started_at: float | None = None,
) -> None:
await publish_turn_run_status(self.bus, msg, status, started_at=started_at)
async def handle_turn_end(
self,