refactor(agent): make run usage explicit (#5546)

* refactor(agent): make run usage explicit

* fix(api): capture usage per run
This commit is contained in:
chengyongru
2026-08-26 15:18:53 +08:00
committed by GitHub
parent 0c84725b13
commit 4f6c0aedfa
20 changed files with 184 additions and 233 deletions
+14 -17
View File
@@ -28,7 +28,12 @@ from nanobot.agent.cron_turns import CronTurnCoordinator
from nanobot.agent.hook import AgentHook, AgentTurnHookFactory from nanobot.agent.hook import AgentHook, AgentTurnHookFactory
from nanobot.agent.memory import Consolidator from nanobot.agent.memory import Consolidator
from nanobot.agent.model_runtime import ModelRuntimeResolver from nanobot.agent.model_runtime import ModelRuntimeResolver
from nanobot.agent.runner import _MAX_INJECTIONS_PER_TURN, AgentRunner, AgentRunSpec from nanobot.agent.runner import (
_MAX_INJECTIONS_PER_TURN,
AgentRunner,
AgentRunResult,
AgentRunSpec,
)
from nanobot.agent.subagent import SubagentManager from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context
from nanobot.agent.tools.exec_session import ExecSessionManager from nanobot.agent.tools.exec_session import ExecSessionManager
@@ -204,11 +209,6 @@ class AgentLoop:
def tool_names(self) -> list[str]: def tool_names(self) -> list[str]:
return self.tools.tool_names return self.tools.tool_names
@property
def last_usage(self) -> LLMUsage | None:
"""Latest aggregate usage exposed through the runtime-control snapshot."""
return self._last_usage
@property @property
def provider(self) -> LLMProvider: def provider(self) -> LLMProvider:
"""Provider selected for future turn admissions.""" """Provider selected for future turn admissions."""
@@ -379,7 +379,6 @@ class AgentLoop:
default_restrict_to_workspace=restrict_to_workspace, default_restrict_to_workspace=restrict_to_workspace,
) )
self._start_time = time.time() self._start_time = time.time()
self._last_usage: LLMUsage | None = None
self._extra_hooks: list[AgentHook] = hooks or [] self._extra_hooks: list[AgentHook] = hooks or []
self._hook_factories: list[AgentTurnHookFactory] = hook_factories or [] self._hook_factories: list[AgentTurnHookFactory] = hook_factories or []
@@ -973,7 +972,7 @@ class AgentLoop:
tools: ToolRegistry | None = None, tools: ToolRegistry | None = None,
request_context: RequestContext | None = None, request_context: RequestContext | None = None,
provider_state: ProviderConversationState | None = None, provider_state: ProviderConversationState | None = None,
) -> tuple[str | None, list[str], list[dict[str, Any]], str, bool]: ) -> AgentRunResult:
"""Run the agent iteration loop. """Run the agent iteration loop.
*on_stream*: called with each content delta during streaming. *on_stream*: called with each content delta during streaming.
@@ -981,7 +980,7 @@ class AgentLoop:
``resuming=True`` means the active turn continues. ``merge_next=True`` means ``resuming=True`` means the active turn continues. ``merge_next=True`` means
the next text segment belongs to the same user-visible assistant message. the next text segment belongs to the same user-visible assistant message.
Returns (final_content, tools_used, messages, stop_reason, had_injections). Returns the complete result produced by ``AgentRunner``.
""" """
self._sync_subagent_runtime_limits() self._sync_subagent_runtime_limits()
@@ -1227,7 +1226,6 @@ class AgentLoop:
reset_workspace_scope(workspace_token) reset_workspace_scope(workspace_token)
reset_request_context(request_token) reset_request_context(request_token)
reset_file_states(file_state_token) reset_file_states(file_state_token)
self._last_usage = result.usage
if session is not None and not ephemeral: if session is not None and not ephemeral:
session.provider_state = result.provider_state session.provider_state = result.provider_state
if result.stop_reason == "max_iterations": if result.stop_reason == "max_iterations":
@@ -1250,7 +1248,7 @@ class AgentLoop:
await on_stream_end(resuming=False) await on_stream_end(resuming=False)
elif result.stop_reason == "error": elif result.stop_reason == "error":
logger.error("LLM returned error: {}", (result.final_content or "")[:200]) logger.error("LLM returned error: {}", (result.final_content or "")[:200])
return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections return result
def _check_expired_sessions_if_due(self) -> None: def _check_expired_sessions_if_due(self) -> None:
"""Scan idle sessions no more often than the configured interval.""" """Scan idle sessions no more often than the configured interval."""
@@ -2045,12 +2043,11 @@ class AgentLoop:
request_context=ctx.request_context, request_context=ctx.request_context,
provider_state=ctx.provider_state, provider_state=ctx.provider_state,
) )
final_content, _, all_msgs, stop_reason, had_injections = result ctx.final_content = result.final_content
ctx.final_content = final_content ctx.all_messages = result.messages
ctx.all_messages = all_msgs ctx.stop_reason = result.stop_reason
ctx.stop_reason = stop_reason ctx.had_injections = result.had_injections
ctx.had_injections = had_injections ctx.usage = result.usage
ctx.usage = self._last_usage
ctx.delivery.record_usage(ctx.usage) ctx.delivery.record_usage(ctx.usage)
if ctx.kind is TurnKind.USER: if ctx.kind is TurnKind.USER:
await turn_continuation.maybe_continue_turn(ctx) await turn_continuation.maybe_continue_turn(ctx)
-8
View File
@@ -12,7 +12,6 @@ if TYPE_CHECKING:
from nanobot.agent.tools.shell import ExecToolConfig from nanobot.agent.tools.shell import ExecToolConfig
from nanobot.agent.tools.web import WebToolsConfig from nanobot.agent.tools.web import WebToolsConfig
from nanobot.config.schema import ModelPresetConfig from nanobot.config.schema import ModelPresetConfig
from nanobot.providers.base import LLMUsage
from nanobot.utils.llm_runtime import LLMRuntime from nanobot.utils.llm_runtime import LLMRuntime
@@ -35,7 +34,6 @@ RUNTIME_SNAPSHOT_KEYS = frozenset({
"web_config", "web_config",
"exec_config", "exec_config",
"subagents", "subagents",
"_last_usage",
}) })
RUNTIME_COMMAND_KEYS = frozenset({ RUNTIME_COMMAND_KEYS = frozenset({
@@ -66,7 +64,6 @@ class RuntimeSnapshot:
web_config: dict[str, object] web_config: dict[str, object]
exec_config: dict[str, object] exec_config: dict[str, object]
subagent_statuses: dict[str, dict[str, object]] subagent_statuses: dict[str, dict[str, object]]
last_usage: Mapping[str, JsonScalar]
scratchpad: dict[str, JsonValue] scratchpad: dict[str, JsonValue]
def as_mapping(self) -> Mapping[str, object]: def as_mapping(self) -> Mapping[str, object]:
@@ -86,7 +83,6 @@ class RuntimeSnapshot:
"web_config": self.web_config, "web_config": self.web_config,
"exec_config": self.exec_config, "exec_config": self.exec_config,
"subagents": {"_task_statuses": self.subagent_statuses}, "subagents": {"_task_statuses": self.subagent_statuses},
"_last_usage": self.last_usage,
} }
assert values.keys() == RUNTIME_SNAPSHOT_KEYS assert values.keys() == RUNTIME_SNAPSHOT_KEYS
return values return values
@@ -151,9 +147,6 @@ class _RuntimeControlTarget(Protocol):
@property @property
def tool_names(self) -> list[str]: ... def tool_names(self) -> list[str]: ...
@property
def last_usage(self) -> LLMUsage | None: ...
def set_runtime_model(self, model: str) -> LLMRuntime: ... def set_runtime_model(self, model: str) -> LLMRuntime: ...
def set_runtime_context_window(self, context_window_tokens: int) -> LLMRuntime: ... def set_runtime_context_window(self, context_window_tokens: int) -> LLMRuntime: ...
@@ -191,7 +184,6 @@ class AgentRuntimeControl:
web_config=_snapshot_web_config(target.web_config), web_config=_snapshot_web_config(target.web_config),
exec_config=_snapshot_exec_config(target.exec_config), exec_config=_snapshot_exec_config(target.exec_config),
subagent_statuses=_snapshot_subagent_statuses(target.subagents), subagent_statuses=_snapshot_subagent_statuses(target.subagents),
last_usage=target.last_usage.to_dict() if target.last_usage is not None else {},
scratchpad=_snapshot_json_mapping(self.__scratchpad), scratchpad=_snapshot_json_mapping(self.__scratchpad),
) )
+2 -5
View File
@@ -90,7 +90,6 @@ class MyTool(Tool):
"tool_names", "tool_names",
"current_iteration", "current_iteration",
"_current_iteration", # updated by runner only "_current_iteration", # updated by runner only
"_last_usage",
"exec_config", # inspect allowed (e.g. check sandbox), modify blocked "exec_config", # inspect allowed (e.g. check sandbox), modify blocked
"web_config", # inspect allowed (e.g. check enable), modify blocked "web_config", # inspect allowed (e.g. check enable), modify blocked
"model_presets", # config-derived catalog; changes require config reload "model_presets", # config-derived catalog; changes require config reload
@@ -150,7 +149,7 @@ class MyTool(Tool):
"Actions: check, set.\n" "Actions: check, set.\n"
"- check (no key): full config overview — start here.\n" "- check (no key): full config overview — start here.\n"
"- check (key): drill into a value. Dot-paths allowed " "- check (key): drill into a value. Dot-paths allowed "
"(e.g. '_last_usage.input_tokens', 'web_config.enable').\n" "(e.g. 'web_config.enable').\n"
"- set (key, value): change config or store notes in your scratchpad. " "- set (key, value): change config or store notes in your scratchpad. "
"Scratchpad keys persist across turns but not restarts.\n" "Scratchpad keys persist across turns but not restarts.\n"
"Key values: _current_iteration (current progress), " "Key values: _current_iteration (current progress), "
@@ -162,7 +161,7 @@ class MyTool(Tool):
"Note: web_config and exec_config are readable but read-only.\n" "Note: web_config and exec_config are readable but read-only.\n"
"\n" "\n"
"When to use:\n" "When to use:\n"
"- User asks about your model, settings, or token usage → check that key.\n" "- User asks about your model or settings → check that key.\n"
"- User asks to switch to a named model preset → set model_preset to that preset name.\n" "- User asks to switch to a named model preset → set model_preset to that preset name.\n"
"- A tool fails or behaves unexpectedly → check the related config to diagnose.\n" "- A tool fails or behaves unexpectedly → check the related config to diagnose.\n"
"- User asks you to remember a preference for this session → set to store it in your scratchpad.\n" "- User asks you to remember a preference for this session → set to store it in your scratchpad.\n"
@@ -448,8 +447,6 @@ class MyTool(Tool):
"subagents", "subagents",
): ):
parts.append(self._format_value(values[k], k)) parts.append(self._format_value(values[k], k))
if snapshot.last_usage:
parts.append(self._format_value(snapshot.last_usage, "_last_usage"))
if snapshot.scratchpad: if snapshot.scratchpad:
parts.append(self._format_value(snapshot.scratchpad, "scratchpad")) parts.append(self._format_value(snapshot.scratchpad, "scratchpad"))
return "\n".join(parts) return "\n".join(parts)
+15 -1
View File
@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable, cast
from aiohttp import web from aiohttp import web
from loguru import logger from loguru import logger
from nanobot.agent.hook import AgentHook, AgentRunHookContext
from nanobot.config.paths import get_media_dir from nanobot.config.paths import get_media_dir
from nanobot.providers.base import LLMUsage from nanobot.providers.base import LLMUsage
from nanobot.utils.helpers import safe_filename from nanobot.utils.helpers import safe_filename
@@ -53,6 +54,17 @@ _PREPARE_AGENT_KEY = web.AppKey[Callable[[], Awaitable[None]] | None]("prepare_a
_MISSING = object() _MISSING = object()
class _UsageCaptureHook(AgentHook):
"""Capture the aggregate usage owned by one API run."""
def __init__(self) -> None:
super().__init__()
self.usage: LLMUsage | None = None
async def after_run(self, context: AgentRunHookContext) -> None:
self.usage = context.usage
def _app_value( def _app_value(
app: Any, app: Any,
key: web.AppKey[Any], key: web.AppKey[Any],
@@ -399,6 +411,7 @@ async def handle_chat_completions(request: web.Request) -> web.Response | web.St
return resp return resp
# -- non-streaming path (original logic) -- # -- non-streaming path (original logic) --
usage_capture = _UsageCaptureHook()
try: try:
async with session_lock: async with session_lock:
try: try:
@@ -410,6 +423,7 @@ async def handle_chat_completions(request: web.Request) -> web.Response | web.St
session_key=session_key, session_key=session_key,
channel="api", channel="api",
chat_id=API_CHAT_ID, chat_id=API_CHAT_ID,
hooks=[usage_capture],
) )
response_text = _response_text(response) response_text = _response_text(response)
if not response_text or not response_text.strip(): if not response_text or not response_text.strip():
@@ -426,7 +440,7 @@ async def handle_chat_completions(request: web.Request) -> web.Response | web.St
return _error_json(500, "Internal server error", err_type="server_error") return _error_json(500, "Internal server error", err_type="server_error")
return web.json_response( return web.json_response(
_chat_completion_response(response_text, model_name, getattr(agent_loop, "_last_usage", None)) _chat_completion_response(response_text, model_name, usage_capture.usage)
) )
+3 -2
View File
@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Literal, cast
from nanobot import __version__ from nanobot import __version__
from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text
from nanobot.providers.base import LLMUsage
from nanobot.utils.helpers import build_status_content from nanobot.utils.helpers import build_status_content
from nanobot.utils.restart import set_restart_notice_to_env from nanobot.utils.restart import set_restart_notice_to_env
from nanobot.utils.workspace_prompts import initialize_workspace_prompt from nanobot.utils.workspace_prompts import initialize_workspace_prompt
@@ -265,8 +266,8 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
session, session,
runtime=runtime, runtime=runtime,
) )
last_usage = LLMUsage.from_dict(session.metadata.get("_last_usage"))
if ctx_est <= 0: if ctx_est <= 0:
last_usage = loop._last_usage # pyright: ignore[reportPrivateUsage]
ctx_est = last_usage.input_tokens if last_usage is not None else 0 ctx_est = last_usage.input_tokens if last_usage is not None else 0
# Fetch web search provider usage (best-effort, never blocks the response) # Fetch web search provider usage (best-effort, never blocks the response)
@@ -289,7 +290,7 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
chat_id=ctx.msg.chat_id, chat_id=ctx.msg.chat_id,
content=build_status_content( content=build_status_content(
version=__version__, model=runtime.model, version=__version__, model=runtime.model,
start_time=loop._start_time, last_usage=loop._last_usage, # pyright: ignore[reportPrivateUsage] start_time=loop._start_time, last_usage=last_usage, # pyright: ignore[reportPrivateUsage]
context_window_tokens=runtime.context_window_tokens, context_window_tokens=runtime.context_window_tokens,
session_msg_count=len(session.get_history(max_messages=0)), session_msg_count=len(session.get_history(max_messages=0)),
context_tokens_estimate=ctx_est, context_tokens_estimate=ctx_est,
+1 -1
View File
@@ -1,6 +1,6 @@
--- ---
name: my name: my
description: Inspect and optionally adjust the agent's runtime state. Use to check the current model or preset, context window, iteration progress and limits, token usage, workspace and tool configuration, subagent status, and request routing metadata such as channel, chat ID, and sender ID; diagnose unavailable capabilities; change allowed runtime settings; or store temporary session scratchpad values. description: Inspect and optionally adjust the agent's runtime state. Use to check the current model or preset, context window, iteration progress and limits, workspace and tool configuration, subagent status, and request routing metadata such as channel, chat ID, and sender ID; diagnose unavailable capabilities; change allowed runtime settings; or store temporary session scratchpad values.
--- ---
# Self-Awareness # Self-Awareness
-11
View File
@@ -15,8 +15,6 @@ Concrete scenarios showing when and how to use the my tool effectively.
``` ```
→ my(action="check", key="max_iterations") → my(action="check", key="max_iterations")
→ 40 → 40
→ my(action="check", key="_last_usage")
→ {"input_tokens": 62000, "output_tokens": 3000}
→ "I hit the iteration limit (40). The task was complex. I can ask the user if they want to increase it." → "I hit the iteration limit (40). The task was complex. I can ask the user if they want to increase it."
``` ```
@@ -66,12 +64,3 @@ Concrete scenarios showing when and how to use the my tool effectively.
→ my(action="set", key="test_framework", value="pytest") → my(action="set", key="test_framework", value="pytest")
→ my(action="set", key="has_docker", value=true) → my(action="set", key="has_docker", value=true)
``` ```
## Budget Awareness
### Token-conscious behavior
```
→ my(action="check", key="_last_usage")
→ {"input_tokens": 58000, "output_tokens": 12000}
→ "I've consumed ~70k tokens. I'll keep my remaining responses focused."
```
+3 -3
View File
@@ -145,7 +145,7 @@ async def test_pending_document_attachment_keeps_body_out_of_prompt(
) )
) )
final_content, _, _, _, had_injections = await loop._run_agent_loop( result = await loop._run_agent_loop(
[{"role": "user", "content": "hello"}], [{"role": "user", "content": "hello"}],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
channel="cli", channel="cli",
@@ -153,8 +153,8 @@ async def test_pending_document_attachment_keeps_body_out_of_prompt(
pending_queue=pending_queue, pending_queue=pending_queue,
) )
assert final_content == "answer-2" assert result.final_content == "answer-2"
assert had_injections is True assert result.had_injections is True
injected_user_content = [ injected_user_content = [
message["content"] message["content"]
for message in captured_messages[-1] for message in captured_messages[-1]
+8 -1
View File
@@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest import pytest
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.agent.runner import AgentRunResult
from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.events import InboundMessage from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
@@ -235,7 +236,13 @@ class TestAgentLoopTTLParam:
session = loop.sessions.get_or_create("cli:direct") session = loop.sessions.get_or_create("cli:direct")
session.get_history = MagicMock(return_value=[]) session.get_history = MagicMock(return_value=[])
loop.context.build_messages = MagicMock(return_value=[]) loop.context.build_messages = MagicMock(return_value=[])
loop._run_agent_loop = AsyncMock(return_value=("ok", [], [], "stop", False)) loop._run_agent_loop = AsyncMock(
return_value=AgentRunResult(
final_content="ok",
messages=[],
stop_reason="stop",
)
)
loop._save_turn = MagicMock() loop._save_turn = MagicMock()
msg = InboundMessage( msg = InboundMessage(
+7 -7
View File
@@ -457,12 +457,12 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path):
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
content, tools_used, messages, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[{"role": "user", "content": "hi"}], [{"role": "user", "content": "hi"}],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
) )
assert content == "done" assert result.final_content == "done"
assert "before_run" in events assert "before_run" in events
assert "before_iter:0" in events assert "before_iter:0" in events
assert "after_iter:0" in events assert "after_iter:0" in events
@@ -545,12 +545,12 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path):
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[{"role": "user", "content": "hi"}], [{"role": "user", "content": "hi"}],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
) )
assert content == "still works" assert result.final_content == "still works"
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -590,11 +590,11 @@ async def test_agent_loop_no_hooks_backward_compat(tmp_path):
loop.tools.execute = AsyncMock(return_value="ok") loop.tools.execute = AsyncMock(return_value="ok")
loop.max_iterations = 2 loop.max_iterations = 2
content, tools_used, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime() [], runtime=loop.llm_runtime()
) )
assert content == ( assert result.final_content == (
"I reached the maximum number of tool call iterations (2) " "I reached the maximum number of tool call iterations (2) "
"without completing the task. You can try breaking the task into smaller steps." "without completing the task. You can try breaking the task into smaller steps."
) )
assert tools_used == ["list_dir", "list_dir"] assert result.tools_used == ["list_dir", "list_dir"]
+8 -8
View File
@@ -83,11 +83,11 @@ class TestToolEventProgress:
) -> None: ) -> None:
progress.append((content, tool_hint, tool_events)) progress.append((content, tool_hint, tool_events))
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress [], runtime=loop.llm_runtime(), on_progress=on_progress
) )
assert final_content == "Done" assert result.final_content == "Done"
assert progress == [ assert progress == [
("Visible", False, None), ("Visible", False, None),
( (
@@ -154,11 +154,11 @@ class TestToolEventProgress:
if file_edit_events: if file_edit_events:
file_events.extend(file_edit_events) file_events.extend(file_edit_events)
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress [], runtime=loop.llm_runtime(), on_progress=on_progress
) )
assert final_content == "Done" assert result.final_content == "Done"
assert [event["phase"] for event in file_events] == ["start", "end"] assert [event["phase"] for event in file_events] == ["start", "end"]
assert file_events[0] == { assert file_events[0] == {
"version": 1, "version": 1,
@@ -224,11 +224,11 @@ class TestToolEventProgress:
prepare_file_edit_trackers, prepare_file_edit_trackers,
) )
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress [], runtime=loop.llm_runtime(), on_progress=on_progress
) )
assert final_content == "Done" assert result.final_content == "Done"
assert target.read_text(encoding="utf-8") == "new\n" assert target.read_text(encoding="utf-8") == "new\n"
prepare_file_edit_trackers.assert_not_called() prepare_file_edit_trackers.assert_not_called()
@@ -1028,14 +1028,14 @@ class TestToolEventProgress:
) -> None: ) -> None:
progress.append((content, tool_hint, tool_events)) progress.append((content, tool_hint, tool_events))
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], [],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
on_progress=on_progress, on_progress=on_progress,
on_stream=on_stream, on_stream=on_stream,
) )
assert final_content == "Done" assert result.final_content == "Done"
assert streamed == ["I will", " inspect it."] assert streamed == ["I will", " inspect it."]
assert progress[0][0] == 'custom_tool("foo.txt")' assert progress[0][0] == 'custom_tool("foo.txt")'
assert all(item[0] != "I will inspect it." for item in progress) assert all(item[0] != "I will inspect it." for item in progress)
+13 -13
View File
@@ -338,11 +338,11 @@ async def test_loop_max_iterations_message_stays_stable(tmp_path):
loop.tools.execute = AsyncMock(return_value="ok") loop.tools.execute = AsyncMock(return_value="ok")
loop.max_iterations = 2 loop.max_iterations = 2
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime() [], runtime=loop.llm_runtime()
) )
assert final_content == ( assert result.final_content == (
"I reached the maximum number of tool call iterations (2) " "I reached the maximum number of tool call iterations (2) "
"without completing the task. You can try breaking the task into smaller steps." "without completing the task. You can try breaking the task into smaller steps."
) )
@@ -359,16 +359,16 @@ async def test_loop_goal_turn_uses_standard_iteration_budget(tmp_path):
loop.tools.execute = AsyncMock(return_value="ok") loop.tools.execute = AsyncMock(return_value="ok")
loop.max_iterations = 2 loop.max_iterations = 2
final_content, _, _, stop_reason, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], [],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
metadata={"original_command": "/goal"}, metadata={"original_command": "/goal"},
) )
assert stop_reason == "max_iterations" assert result.stop_reason == "max_iterations"
assert loop.provider.chat_with_retry.await_count == 3 assert loop.provider.chat_with_retry.await_count == 3
assert loop.provider.chat_with_retry.await_args_list[-1].kwargs["tools"] is None assert loop.provider.chat_with_retry.await_args_list[-1].kwargs["tools"] is None
assert final_content == ( assert result.final_content == (
"I reached the maximum number of tool call iterations (2) " "I reached the maximum number of tool call iterations (2) "
"without completing the task. You can try breaking the task into smaller steps." "without completing the task. You can try breaking the task into smaller steps."
) )
@@ -393,14 +393,14 @@ async def test_loop_stream_filter_handles_think_only_prefix_without_crashing(tmp
async def on_stream_end(*, resuming: bool = False) -> None: async def on_stream_end(*, resuming: bool = False) -> None:
endings.append(resuming) endings.append(resuming)
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], [],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
on_stream=on_stream, on_stream=on_stream,
on_stream_end=on_stream_end, on_stream_end=on_stream_end,
) )
assert final_content == "Hello" assert result.final_content == "Hello"
assert deltas == ["Hello"] assert deltas == ["Hello"]
assert endings == [False] assert endings == [False]
@@ -420,11 +420,11 @@ async def test_loop_stream_filter_hides_partial_trailing_think_prefix(tmp_path):
async def on_stream(delta: str) -> None: async def on_stream(delta: str) -> None:
deltas.append(delta) deltas.append(delta)
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_stream=on_stream [], runtime=loop.llm_runtime(), on_stream=on_stream
) )
assert final_content == "Hello World" assert result.final_content == "Hello World"
assert deltas == ["Hello", " World"] assert deltas == ["Hello", " World"]
@@ -443,11 +443,11 @@ async def test_loop_stream_filter_hides_complete_trailing_think_tag(tmp_path):
async def on_stream(delta: str) -> None: async def on_stream(delta: str) -> None:
deltas.append(delta) deltas.append(delta)
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_stream=on_stream [], runtime=loop.llm_runtime(), on_stream=on_stream
) )
assert final_content == "Hello World" assert result.final_content == "Hello World"
assert deltas == ["Hello", " World"] assert deltas == ["Hello", " World"]
@@ -464,11 +464,11 @@ async def test_loop_retries_think_only_final_response(tmp_path):
loop.provider.chat_with_retry = chat_with_retry loop.provider.chat_with_retry = chat_with_retry
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime() [], runtime=loop.llm_runtime()
) )
assert final_content == "Recovered answer" assert result.final_content == "Recovered answer"
assert call_count["n"] == 2 assert call_count["n"] == 2
+54 -74
View File
@@ -9,6 +9,7 @@ from loguru import logger
from nanobot.agent.context import ContextBuilder from nanobot.agent.context import ContextBuilder
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.agent.runner import AgentRunResult
from nanobot.agent.tools.context import RequestContext, request_context from nanobot.agent.tools.context import RequestContext, request_context
from nanobot.bus.events import InboundMessage from nanobot.bus.events import InboundMessage
from nanobot.bus.outbound_events import ( from nanobot.bus.outbound_events import (
@@ -54,6 +55,23 @@ from nanobot.session.webui_turns import (
from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META
def _agent_run_result(
final_content: str,
messages: list[dict],
*,
stop_reason: str = "completed",
had_injections: bool = False,
usage: LLMUsage | None = None,
) -> AgentRunResult:
return AgentRunResult(
final_content=final_content,
messages=messages,
stop_reason=stop_reason,
had_injections=had_injections,
usage=usage,
)
def _mk_loop() -> AgentLoop: def _mk_loop() -> AgentLoop:
loop = AgentLoop.__new__(AgentLoop) loop = AgentLoop.__new__(AgentLoop)
from nanobot.config.schema import AgentDefaults from nanobot.config.schema import AgentDefaults
@@ -1261,16 +1279,14 @@ async def test_process_message_persists_media_only_turn_without_text(tmp_path: P
async def test_process_message_does_not_duplicate_early_persisted_user_message(tmp_path: Path) -> None: async def test_process_message_does_not_duplicate_early_persisted_user_message(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path) loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop._run_agent_loop = AsyncMock(return_value=( loop._run_agent_loop = AsyncMock(return_value=_agent_run_result(
"done", "done",
None,
[ [
{"role": "system", "content": "system"}, {"role": "system", "content": "system"},
{"role": "user", "content": "hello"}, {"role": "user", "content": "hello"},
{"role": "assistant", "content": "done"}, {"role": "assistant", "content": "done"},
], ],
"stop", stop_reason="stop",
False,
)) # type: ignore[method-assign] )) # type: ignore[method-assign]
result = await loop._process_message( result = await loop._process_message(
@@ -1308,19 +1324,14 @@ async def test_internal_continuation_queues_turn_without_fake_user_history(
async def fake_run_agent_loop(initial_messages, *, metadata=None, **_kwargs): async def fake_run_agent_loop(initial_messages, *, metadata=None, **_kwargs):
calls.append({"initial_messages": initial_messages, "metadata": metadata}) calls.append({"initial_messages": initial_messages, "metadata": metadata})
if len(calls) == 1: if len(calls) == 1:
return ( return _agent_run_result(
"paused", "paused",
[],
[*initial_messages, {"role": "assistant", "content": "paused"}], [*initial_messages, {"role": "assistant", "content": "paused"}],
"max_iterations", stop_reason="max_iterations",
False,
) )
return ( return _agent_run_result(
"done", "done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}], [*initial_messages, {"role": "assistant", "content": "done"}],
"completed",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -1382,23 +1393,18 @@ async def test_internal_continuation_preserves_streaming_route_metadata(
nonlocal calls nonlocal calls
calls += 1 calls += 1
if calls == 1: if calls == 1:
return ( return _agent_run_result(
"paused", "paused",
[],
[*initial_messages, {"role": "assistant", "content": "paused"}], [*initial_messages, {"role": "assistant", "content": "paused"}],
"max_iterations", stop_reason="max_iterations",
False,
) )
assert on_stream is not None assert on_stream is not None
assert on_stream_end is not None assert on_stream_end is not None
await on_stream("done") await on_stream("done")
await on_stream_end(resuming=False) await on_stream_end(resuming=False)
return ( return _agent_run_result(
"done", "done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}], [*initial_messages, {"role": "assistant", "content": "done"}],
"completed",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -1460,19 +1466,14 @@ async def test_websocket_internal_continuation_keeps_single_visible_run(
nonlocal calls nonlocal calls
calls += 1 calls += 1
if calls == 1: if calls == 1:
return ( return _agent_run_result(
"paused", "paused",
[],
[*initial_messages, {"role": "assistant", "content": "paused"}], [*initial_messages, {"role": "assistant", "content": "paused"}],
"max_iterations", stop_reason="max_iterations",
False,
) )
return ( return _agent_run_result(
"done", "done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}], [*initial_messages, {"role": "assistant", "content": "done"}],
"completed",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -1521,16 +1522,14 @@ async def test_process_message_keeps_delivery_chat_for_thread_session(tmp_path:
{"role": "user", "content": "runtime + hello"}, {"role": "user", "content": "runtime + hello"},
] ]
) )
loop._run_agent_loop = AsyncMock(return_value=( # type: ignore[method-assign] loop._run_agent_loop = AsyncMock(return_value=_agent_run_result( # type: ignore[method-assign]
"done", "done",
[],
[ [
{"role": "system", "content": "system"}, {"role": "system", "content": "system"},
{"role": "user", "content": "runtime + hello"}, {"role": "user", "content": "runtime + hello"},
{"role": "assistant", "content": "done"}, {"role": "assistant", "content": "done"},
], ],
"stop", stop_reason="stop",
False,
)) ))
result = await loop._process_message( result = await loop._process_message(
@@ -1571,16 +1570,14 @@ async def test_process_message_uses_explicit_session_for_goal_context(
{"role": "user", "content": "runtime + system"}, {"role": "user", "content": "runtime + system"},
] ]
) )
loop._run_agent_loop = AsyncMock(return_value=( # type: ignore[method-assign] loop._run_agent_loop = AsyncMock(return_value=_agent_run_result( # type: ignore[method-assign]
"ok", "ok",
[],
[ [
{"role": "system", "content": "system"}, {"role": "system", "content": "system"},
{"role": "user", "content": "runtime + system"}, {"role": "user", "content": "runtime + system"},
{"role": "assistant", "content": "ok"}, {"role": "assistant", "content": "ok"},
], ],
"stop", stop_reason="stop",
False,
)) ))
result = await loop._process_message( result = await loop._process_message(
@@ -1711,9 +1708,8 @@ async def test_next_turn_after_crash_closes_pending_user_turn_before_new_input(t
]) ])
loop.sessions.save(session) loop.sessions.save(session)
loop._run_agent_loop = AsyncMock(return_value=( loop._run_agent_loop = AsyncMock(return_value=_agent_run_result(
"new answer", "new answer",
None,
[ [
{"role": "system", "content": "system"}, {"role": "system", "content": "system"},
{"role": "user", "content": "old question"}, {"role": "user", "content": "old question"},
@@ -1721,8 +1717,7 @@ async def test_next_turn_after_crash_closes_pending_user_turn_before_new_input(t
{"role": "user", "content": "new question"}, {"role": "user", "content": "new question"},
{"role": "assistant", "content": "new answer"}, {"role": "assistant", "content": "new answer"},
], ],
"stop", stop_reason="stop",
False,
)) # type: ignore[method-assign] )) # type: ignore[method-assign]
result = await loop._process_message( result = await loop._process_message(
@@ -1816,12 +1811,10 @@ async def test_stop_preserves_runtime_checkpoint_for_next_turn(tmp_path: Path) -
assert interrupted.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is not None assert interrupted.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is not None
async def resumed_run_agent_loop(initial_messages, **_kwargs): async def resumed_run_agent_loop(initial_messages, **_kwargs):
return ( return _agent_run_result(
"next answer", "next answer",
None,
[*initial_messages, {"role": "assistant", "content": "next answer"}], [*initial_messages, {"role": "assistant", "content": "next answer"}],
"stop", stop_reason="stop",
False,
) )
loop._run_agent_loop = resumed_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = resumed_run_agent_loop # type: ignore[method-assign]
@@ -1872,12 +1865,10 @@ async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_
seen["initial_messages"] = initial_messages seen["initial_messages"] = initial_messages
seen["runtime"] = kwargs["runtime"] seen["runtime"] = kwargs["runtime"]
seen["request_context"] = kwargs["request_context"] seen["request_context"] = kwargs["request_context"]
return ( return _agent_run_result(
"done", "done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}], [*initial_messages, {"role": "assistant", "content": "done"}],
"stop", stop_reason="stop",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -1944,15 +1935,14 @@ async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_
async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) -> None: async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path) loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
turn_usage = LLMUsage.reported(input_tokens=64, output_tokens=9)
async def fake_run_agent_loop(initial_messages, **_kwargs): async def fake_run_agent_loop(initial_messages, **_kwargs):
loop._last_usage = LLMUsage.reported(input_tokens=64, output_tokens=9) return _agent_run_result(
return (
"done", "done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}], [*initial_messages, {"role": "assistant", "content": "done"}],
"stop", stop_reason="stop",
False, usage=turn_usage,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -1962,7 +1952,7 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
loop.sessions.invalidate("cli:usage") loop.sessions.invalidate("cli:usage")
assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == ( assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == (
LLMUsage.reported(input_tokens=64, output_tokens=9).to_dict() turn_usage.to_dict()
) )
@@ -1974,12 +1964,10 @@ async def test_system_subagent_followup_does_not_log_content(tmp_path: Path) ->
) )
async def fake_run_agent_loop(initial_messages, **_kwargs): async def fake_run_agent_loop(initial_messages, **_kwargs):
return ( return _agent_run_result(
"done", "done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}], [*initial_messages, {"role": "assistant", "content": "done"}],
"stop", stop_reason="stop",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -2032,12 +2020,10 @@ async def test_system_subagent_followup_uses_common_turn_lifecycle(tmp_path: Pat
setattr(loop, name, record) setattr(loop, name, record)
async def fake_run_agent_loop(initial_messages, **_kwargs): async def fake_run_agent_loop(initial_messages, **_kwargs):
return ( return _agent_run_result(
"done", "done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}], [*initial_messages, {"role": "assistant", "content": "done"}],
"stop", stop_reason="stop",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -2077,12 +2063,10 @@ async def test_multiple_subagent_followups_all_persist_as_standalone_history(tmp
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
async def fake_run_agent_loop(initial_messages, **_kwargs): async def fake_run_agent_loop(initial_messages, **_kwargs):
return ( return _agent_run_result(
"ack", "ack",
[],
[*initial_messages, {"role": "assistant", "content": "ack"}], [*initial_messages, {"role": "assistant", "content": "ack"}],
"stop", stop_reason="stop",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -2212,12 +2196,10 @@ async def test_system_subagent_followup_uses_thread_session_and_slack_metadata(t
async def fake_run_agent_loop(initial_messages, **kwargs): async def fake_run_agent_loop(initial_messages, **kwargs):
seen["initial_messages"] = initial_messages seen["initial_messages"] = initial_messages
seen["request_context"] = kwargs["request_context"] seen["request_context"] = kwargs["request_context"]
return ( return _agent_run_result(
"done", "done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}], [*initial_messages, {"role": "assistant", "content": "done"}],
"stop", stop_reason="stop",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
@@ -2269,9 +2251,8 @@ async def test_turn_after_unanswered_user_keeps_tool_call_pairing(tmp_path: Path
async def fake_run_agent_loop(initial_messages, **_kwargs): async def fake_run_agent_loop(initial_messages, **_kwargs):
assert [m["role"] for m in initial_messages] == ["system", "user"] assert [m["role"] for m in initial_messages] == ["system", "user"]
return ( return _agent_run_result(
"done", "done",
[],
[ [
*initial_messages, *initial_messages,
{ {
@@ -2286,8 +2267,7 @@ async def test_turn_after_unanswered_user_keeps_tool_call_pairing(tmp_path: Path
{"role": "tool", "tool_call_id": "call_ls", "name": "exec", "content": "file.txt"}, {"role": "tool", "tool_call_id": "call_ls", "name": "exec", "content": "file.txt"},
{"role": "assistant", "content": "done"}, {"role": "assistant", "content": "done"},
], ],
"stop", stop_reason="stop",
False,
) )
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign] loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
+13 -13
View File
@@ -614,7 +614,7 @@ async def test_loop_injected_followup_preserves_image_media(tmp_path):
media=[str(image_path)], media=[str(image_path)],
)) ))
final_content, _, _, _, had_injections = await loop._run_agent_loop( result = await loop._run_agent_loop(
[{"role": "user", "content": "hello"}], [{"role": "user", "content": "hello"}],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
channel="cli", channel="cli",
@@ -622,8 +622,8 @@ async def test_loop_injected_followup_preserves_image_media(tmp_path):
pending_queue=pending_queue, pending_queue=pending_queue,
) )
assert final_content == "second answer" assert result.final_content == "second answer"
assert had_injections is True assert result.had_injections is True
assert call_count["n"] == 2 assert call_count["n"] == 2
injected_user_messages = [ injected_user_messages = [
message for message in captured_messages[-1] message for message in captured_messages[-1]
@@ -708,7 +708,7 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path):
}, },
)) ))
_, _, all_messages, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[{"role": "user", "content": "initial message from user A"}], [{"role": "user", "content": "initial message from user A"}],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
session=session, session=session,
@@ -741,7 +741,7 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path):
), ),
] ]
injected = [message for message in all_messages if message.get("role") == "user"][-1] injected = [message for message in result.messages if message.get("role") == "user"][-1]
assert "follow-up from the second speaker" in str(injected["content"]) assert "follow-up from the second speaker" in str(injected["content"])
model_messages = provider.chat_with_retry.await_args_list[-1].kwargs["messages"] model_messages = provider.chat_with_retry.await_args_list[-1].kwargs["messages"]
assert "telegram | group-1 | user-b | message-2" in str(model_messages) assert "telegram | group-1 | user-b | message-2" in str(model_messages)
@@ -753,7 +753,7 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path):
"identity", "identity",
] ]
loop._save_turn(session, all_messages, skip=1) loop._save_turn(session, result.messages, skip=1)
persisted = [message for message in session.messages if message.get("role") == "user"][-1] persisted = [message for message in session.messages if message.get("role") == "user"][-1]
assert "telegram | group-1 | user-b | message-2" in str(persisted["content"]) assert "telegram | group-1 | user-b | message-2" in str(persisted["content"])
assert "telegram | group-1 | user-c | message-3" in str(persisted["content"]) assert "telegram | group-1 | user-c | message-3" in str(persisted["content"])
@@ -805,7 +805,7 @@ async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_p
metadata={"injected_event": "subagent_result", "subagent_task_id": "sub-1"}, metadata={"injected_event": "subagent_result", "subagent_task_id": "sub-1"},
)) ))
final_content, _, all_msgs, _, had_injections = await loop._run_agent_loop( result = await loop._run_agent_loop(
[{"role": "user", "content": "hello"}], [{"role": "user", "content": "hello"}],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
channel="cli", channel="cli",
@@ -813,10 +813,10 @@ async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_p
pending_queue=pending_queue, pending_queue=pending_queue,
) )
assert final_content == "second answer" assert result.final_content == "second answer"
assert had_injections is True assert result.had_injections is True
assert call_count["n"] == 2 assert call_count["n"] == 2
injected_users = [message for message in all_msgs if message.get("role") == "user"][-2:] injected_users = [message for message in result.messages if message.get("role") == "user"][-2:]
assert [message["content"] for message in injected_users] == ["visible follow-up", payload] assert [message["content"] for message in injected_users] == ["visible follow-up", payload]
assert injected_users[1][HIDDEN_HISTORY_META] == { assert injected_users[1][HIDDEN_HISTORY_META] == {
"kind": "subagent_result", "kind": "subagent_result",
@@ -1469,7 +1469,7 @@ async def test_pending_queue_preserves_overflow_for_next_injection_cycle(tmp_pat
content=f"follow-up-{idx}", content=f"follow-up-{idx}",
)) ))
final_content, _, _, _, had_injections = await loop._run_agent_loop( result = await loop._run_agent_loop(
[{"role": "user", "content": "hello"}], [{"role": "user", "content": "hello"}],
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
channel="cli", channel="cli",
@@ -1477,8 +1477,8 @@ async def test_pending_queue_preserves_overflow_for_next_injection_cycle(tmp_pat
pending_queue=pending_queue, pending_queue=pending_queue,
) )
assert final_content == "answer-3" assert result.final_content == "answer-3"
assert had_injections is True assert result.had_injections is True
assert call_count["n"] == 3 assert call_count["n"] == 3
flattened_user_content = "\n".join( flattened_user_content = "\n".join(
message["content"] message["content"]
-34
View File
@@ -32,8 +32,6 @@ def _make_mock_loop(**overrides):
loop._start_time = 1000.0 loop._start_time = 1000.0
loop.exec_config = ExecToolConfig() loop.exec_config = ExecToolConfig()
loop.channels_config = MagicMock() loop.channels_config = MagicMock()
loop._last_usage = LLMUsage.reported(input_tokens=100, output_tokens=50)
loop.last_usage = loop._last_usage
loop._current_iteration = 0 loop._current_iteration = 0
loop.current_iteration = loop._current_iteration loop.current_iteration = loop._current_iteration
loop.provider_retry_mode = "standard" loop.provider_retry_mode = "standard"
@@ -112,7 +110,6 @@ class TestInspectSummary:
assert "workspace" in result assert "workspace" in result
assert "provider_retry_mode" in result assert "provider_retry_mode" in result
assert "max_tool_result_chars" in result assert "max_tool_result_chars" in result
assert "_last_usage" in result
assert "_current_iteration" in result assert "_current_iteration" in result
@@ -161,14 +158,6 @@ class TestInspectPathNavigation:
result = await tool.execute(action="check", key="web_config.enable") result = await tool.execute(action="check", key="web_config.enable")
assert "True" in result assert "True" in result
@pytest.mark.asyncio
async def test_inspect_dict_key_via_dotpath(self):
loop = _make_mock_loop()
loop._last_usage = LLMUsage.reported(input_tokens=100, output_tokens=50)
tool = _make_tool(loop=loop)
result = await tool.execute(action="check", key="_last_usage.input_tokens")
assert "100" in result
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_inspect_blocked_in_path(self): async def test_inspect_blocked_in_path(self):
tool = _make_tool() tool = _make_tool()
@@ -1117,29 +1106,6 @@ class TestCurrentIteration:
assert "read-only" in result assert "read-only" in result
# ---------------------------------------------------------------------------
# _last_usage in check summary (Fix #5)
# ---------------------------------------------------------------------------
class TestLastUsageInSummary:
@pytest.mark.asyncio
async def test_last_usage_shown_in_summary(self):
tool = _make_tool()
result = await tool.execute(action="check")
assert "_last_usage" in result
assert "input_tokens" in result
@pytest.mark.asyncio
async def test_last_usage_not_shown_when_empty(self):
loop = _make_mock_loop()
loop._last_usage = None
loop.last_usage = loop._last_usage
tool = _make_tool(loop=loop)
result = await tool.execute(action="check")
assert "_last_usage" not in result
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# request context (audit session tracking) # request context (audit session tracking)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
+11 -7
View File
@@ -236,9 +236,11 @@ class TestRestartCommand:
loop, _bus = _make_loop() loop, _bus = _make_loop()
session = MagicMock() session = MagicMock()
session.get_history.return_value = [{"role": "user"}] * 3 session.get_history.return_value = [{"role": "user"}] * 3
session.metadata = {
"_last_usage": LLMUsage.reported(input_tokens=0, output_tokens=0).to_dict()
}
loop.sessions.get_or_create.return_value = session loop.sessions.get_or_create.return_value = session
loop._start_time = time.time() - 125 loop._start_time = time.time() - 125
loop._last_usage = LLMUsage.reported(input_tokens=0, output_tokens=0)
loop.consolidator.estimate_session_prompt_tokens = MagicMock( loop.consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(20500, "tiktoken") return_value=(20500, "tiktoken")
) )
@@ -309,19 +311,21 @@ class TestRestartCommand:
LLMResponse(content="second", usage=None), LLMResponse(content="second", usage=None),
]) ])
await loop._run_agent_loop([], runtime=loop.llm_runtime()) first = await loop._run_agent_loop([], runtime=loop.llm_runtime())
assert loop._last_usage == LLMUsage.reported(input_tokens=9, output_tokens=4) assert first.usage == LLMUsage.reported(input_tokens=9, output_tokens=4)
await loop._run_agent_loop([], runtime=loop.llm_runtime()) second = await loop._run_agent_loop([], runtime=loop.llm_runtime())
assert loop._last_usage == LLMUsage.estimated(input_tokens=123, output_tokens=7) assert second.usage == LLMUsage.estimated(input_tokens=123, output_tokens=7)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_status_falls_back_to_last_usage_when_context_estimate_missing(self): async def test_status_falls_back_to_session_usage_when_context_estimate_missing(self):
loop, _bus = _make_loop() loop, _bus = _make_loop()
session = MagicMock() session = MagicMock()
session.get_history.return_value = [{"role": "user"}] session.get_history.return_value = [{"role": "user"}]
session.metadata = {
"_last_usage": LLMUsage.reported(input_tokens=1200, output_tokens=34).to_dict()
}
loop.sessions.get_or_create.return_value = session loop.sessions.get_or_create.return_value = session
loop._last_usage = LLMUsage.reported(input_tokens=1200, output_tokens=34)
loop.consolidator.estimate_session_prompt_tokens = MagicMock( loop.consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(0, "none") return_value=(0, "none")
) )
-1
View File
@@ -33,7 +33,6 @@ def _make_mock_agent(response_text: str = "mock response") -> MagicMock:
agent = MagicMock() agent = MagicMock()
agent.process_direct = AsyncMock(return_value=response_text) agent.process_direct = AsyncMock(return_value=response_text)
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
return agent return agent
-8
View File
@@ -77,7 +77,6 @@ def _make_streaming_agent(tokens: list[str]) -> MagicMock:
return " ".join(tokens) return " ".join(tokens)
agent.process_direct = fake_process_direct agent.process_direct = fake_process_direct
agent._last_usage = None
return agent return agent
@@ -136,7 +135,6 @@ async def test_stream_false_returns_json(aiohttp_client) -> None:
agent = MagicMock() agent = MagicMock()
agent.process_direct = AsyncMock(return_value="normal reply") agent.process_direct = AsyncMock(return_value="normal reply")
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
@@ -159,7 +157,6 @@ async def test_stream_default_is_false(aiohttp_client) -> None:
agent = MagicMock() agent = MagicMock()
agent.process_direct = AsyncMock(return_value="default reply") agent.process_direct = AsyncMock(return_value="default reply")
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
@@ -215,7 +212,6 @@ async def test_stream_passes_on_stream_callbacks(aiohttp_client) -> None:
agent = MagicMock() agent = MagicMock()
agent.process_direct = fake_process_direct agent.process_direct = fake_process_direct
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
@@ -248,7 +244,6 @@ async def test_stream_segment_end_does_not_close_sse(aiohttp_client) -> None:
agent.process_direct = fake_process_direct agent.process_direct = fake_process_direct
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
@@ -287,7 +282,6 @@ async def test_stream_uses_final_response_when_no_deltas(aiohttp_client) -> None
agent.process_direct = fake_process_direct agent.process_direct = fake_process_direct
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
@@ -329,7 +323,6 @@ async def test_stream_with_session_id(aiohttp_client) -> None:
agent = MagicMock() agent = MagicMock()
agent.process_direct = fake_process_direct agent.process_direct = fake_process_direct
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
@@ -358,7 +351,6 @@ async def test_streaming_backend_failure_does_not_emit_success_terminator(aiohtt
agent.process_direct = boom agent.process_direct = boom
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
+25 -12
View File
@@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest import pytest
import pytest_asyncio import pytest_asyncio
from nanobot.agent.hook import AgentHook, AgentRunHookContext
from nanobot.api.server import ( from nanobot.api.server import (
API_CHAT_ID, API_CHAT_ID,
API_SESSION_KEY, API_SESSION_KEY,
@@ -36,7 +37,6 @@ def _make_mock_agent(response_text: str = "mock response") -> MagicMock:
agent = MagicMock() agent = MagicMock()
agent.process_direct = AsyncMock(return_value=response_text) agent.process_direct = AsyncMock(return_value=response_text)
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = LLMUsage.reported(input_tokens=100, output_tokens=50)
return agent return agent
@@ -296,7 +296,18 @@ async def test_single_user_message_must_have_user_role() -> None:
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed") @pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_successful_request_uses_fixed_api_session(aiohttp_client, mock_agent) -> None: async def test_successful_request_uses_fixed_api_session_and_run_usage(
aiohttp_client,
mock_agent,
) -> None:
usage = LLMUsage.reported(input_tokens=100, output_tokens=50)
async def process_direct(*, hooks: list[AgentHook], **_kwargs: object) -> str:
for hook in hooks:
await hook.after_run(AgentRunHookContext(messages=[], usage=usage))
return "mock response"
mock_agent.process_direct = AsyncMock(side_effect=process_direct)
app = create_app(mock_agent, model_name="test-model", api_key=API_KEY) app = create_app(mock_agent, model_name="test-model", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
resp = await client.post( resp = await client.post(
@@ -308,13 +319,18 @@ async def test_successful_request_uses_fixed_api_session(aiohttp_client, mock_ag
body = await resp.json() body = await resp.json()
assert body["choices"][0]["message"]["content"] == "mock response" assert body["choices"][0]["message"]["content"] == "mock response"
assert body["model"] == "test-model" assert body["model"] == "test-model"
mock_agent.process_direct.assert_called_once_with( assert body["usage"] == {
content="hello", "prompt_tokens": 100,
media=None, "completion_tokens": 50,
session_key=API_SESSION_KEY, "total_tokens": 150,
channel="api", }
chat_id=API_CHAT_ID, call_kwargs = mock_agent.process_direct.call_args.kwargs
) assert call_kwargs["content"] == "hello"
assert call_kwargs["media"] is None
assert call_kwargs["session_key"] == API_SESSION_KEY
assert call_kwargs["channel"] == "api"
assert call_kwargs["chat_id"] == API_CHAT_ID
assert len(call_kwargs["hooks"]) == 1
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed") @pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
@@ -329,7 +345,6 @@ async def test_followup_requests_share_same_session_key(aiohttp_client) -> None:
agent = MagicMock() agent = MagicMock()
agent.process_direct = fake_process agent.process_direct = fake_process
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
@@ -368,7 +383,6 @@ async def test_fixed_session_requests_are_serialized(aiohttp_client) -> None:
agent = MagicMock() agent = MagicMock()
agent.process_direct = slow_process agent.process_direct = slow_process
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
@@ -485,7 +499,6 @@ async def test_empty_response_falls_back_without_retry(aiohttp_client) -> None:
agent = MagicMock() agent = MagicMock()
agent.process_direct = always_empty agent.process_direct = always_empty
agent.aclose = AsyncMock() agent.aclose = AsyncMock()
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY) app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app) client = await aiohttp_client(app)
+2 -2
View File
@@ -144,11 +144,11 @@ class TestMessageToolSuppressLogic:
async def on_progress(content: str, *, tool_hint: bool = False) -> None: async def on_progress(content: str, *, tool_hint: bool = False) -> None:
progress.append((content, tool_hint)) progress.append((content, tool_hint))
final_content, _, _, _, _ = await loop._run_agent_loop( result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress [], runtime=loop.llm_runtime(), on_progress=on_progress
) )
assert final_content == "Done" assert result.final_content == "Done"
assert progress == [ assert progress == [
("Visible", False), ("Visible", False),
('read foo.txt', True), ('read foo.txt', True),