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

* refactor(agent): make run usage explicit

* fix(api): capture usage per run

nanobot/agent/loop.py

@@ -28,7 +28,12 @@ from nanobot.agent.cron_turns import CronTurnCoordinator
 from nanobot.agent.hook import AgentHook, AgentTurnHookFactory
 from nanobot.agent.memory import Consolidator
 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.tools.context import RequestContext, bind_request_context, reset_request_context
 from nanobot.agent.tools.exec_session import ExecSessionManager
@@ -204,11 +209,6 @@ class AgentLoop:
     def tool_names(self) -> list[str]:
         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
     def provider(self) -> LLMProvider:
         """Provider selected for future turn admissions."""
@@ -379,7 +379,6 @@ class AgentLoop:
             default_restrict_to_workspace=restrict_to_workspace,
         )
         self._start_time = time.time()
-        self._last_usage: LLMUsage | None = None
         self._extra_hooks: list[AgentHook] = hooks or []
         self._hook_factories: list[AgentTurnHookFactory] = hook_factories or []
 
@@ -973,7 +972,7 @@ class AgentLoop:
         tools: ToolRegistry | None = None,
         request_context: RequestContext | None = None,
         provider_state: ProviderConversationState | None = None,
-    ) -> tuple[str | None, list[str], list[dict[str, Any]], str, bool]:
+    ) -> AgentRunResult:
         """Run the agent iteration loop.
 
         *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
         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()
 
@@ -1227,7 +1226,6 @@ class AgentLoop:
             reset_workspace_scope(workspace_token)
             reset_request_context(request_token)
             reset_file_states(file_state_token)
-        self._last_usage = result.usage
         if session is not None and not ephemeral:
             session.provider_state = result.provider_state
         if result.stop_reason == "max_iterations":
@@ -1250,7 +1248,7 @@ class AgentLoop:
                 await on_stream_end(resuming=False)
         elif result.stop_reason == "error":
             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:
         """Scan idle sessions no more often than the configured interval."""
@@ -2045,12 +2043,11 @@ class AgentLoop:
             request_context=ctx.request_context,
             provider_state=ctx.provider_state,
         )
-        final_content, _, all_msgs, stop_reason, had_injections = result
-        ctx.final_content = final_content
-        ctx.all_messages = all_msgs
-        ctx.stop_reason = stop_reason
-        ctx.had_injections = had_injections
-        ctx.usage = self._last_usage
+        ctx.final_content = result.final_content
+        ctx.all_messages = result.messages
+        ctx.stop_reason = result.stop_reason
+        ctx.had_injections = result.had_injections
+        ctx.usage = result.usage
         ctx.delivery.record_usage(ctx.usage)
         if ctx.kind is TurnKind.USER:
             await turn_continuation.maybe_continue_turn(ctx)

nanobot/agent/tools/runtime_control.py

@@ -12,7 +12,6 @@ if TYPE_CHECKING:
     from nanobot.agent.tools.shell import ExecToolConfig
     from nanobot.agent.tools.web import WebToolsConfig
     from nanobot.config.schema import ModelPresetConfig
-    from nanobot.providers.base import LLMUsage
     from nanobot.utils.llm_runtime import LLMRuntime
 
 
@@ -35,7 +34,6 @@ RUNTIME_SNAPSHOT_KEYS = frozenset({
     "web_config",
     "exec_config",
     "subagents",
-    "_last_usage",
 })
 
 RUNTIME_COMMAND_KEYS = frozenset({
@@ -66,7 +64,6 @@ class RuntimeSnapshot:
     web_config: dict[str, object]
     exec_config: dict[str, object]
     subagent_statuses: dict[str, dict[str, object]]
-    last_usage: Mapping[str, JsonScalar]
     scratchpad: dict[str, JsonValue]
 
     def as_mapping(self) -> Mapping[str, object]:
@@ -86,7 +83,6 @@ class RuntimeSnapshot:
             "web_config": self.web_config,
             "exec_config": self.exec_config,
             "subagents": {"_task_statuses": self.subagent_statuses},
-            "_last_usage": self.last_usage,
         }
         assert values.keys() == RUNTIME_SNAPSHOT_KEYS
         return values
@@ -151,9 +147,6 @@ class _RuntimeControlTarget(Protocol):
     @property
     def tool_names(self) -> list[str]: ...
 
-    @property
-    def last_usage(self) -> LLMUsage | None: ...
-
     def set_runtime_model(self, model: str) -> 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),
             exec_config=_snapshot_exec_config(target.exec_config),
             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),
         )
 

nanobot/agent/tools/self.py

@@ -90,7 +90,6 @@ class MyTool(Tool):
         "tool_names",
         "current_iteration",
         "_current_iteration",  # updated by runner only
-        "_last_usage",
         "exec_config",  # inspect allowed (e.g. check sandbox), modify blocked
         "web_config",  # inspect allowed (e.g. check enable), modify blocked
         "model_presets",  # config-derived catalog; changes require config reload
@@ -150,7 +149,7 @@ class MyTool(Tool):
             "Actions: check, set.\n"
             "- check (no key): full config overview — start here.\n"
             "- 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. "
             "Scratchpad keys persist across turns but not restarts.\n"
             "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"
             "\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"
             "- 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"
@@ -448,8 +447,6 @@ class MyTool(Tool):
             "subagents",
         ):
             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:
             parts.append(self._format_value(snapshot.scratchpad, "scratchpad"))
         return "\n".join(parts)

nanobot/api/server.py

@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable, cast
 from aiohttp import web
 from loguru import logger
 
+from nanobot.agent.hook import AgentHook, AgentRunHookContext
 from nanobot.config.paths import get_media_dir
 from nanobot.providers.base import LLMUsage
 from nanobot.utils.helpers import safe_filename
@@ -53,6 +54,17 @@ _PREPARE_AGENT_KEY = web.AppKey[Callable[[], Awaitable[None]] | None]("prepare_a
 _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(
     app: Any,
     key: web.AppKey[Any],
@@ -399,6 +411,7 @@ async def handle_chat_completions(request: web.Request) -> web.Response | web.St
         return resp
 
     # -- non-streaming path (original logic) --
+    usage_capture = _UsageCaptureHook()
     try:
         async with session_lock:
             try:
@@ -410,6 +423,7 @@ async def handle_chat_completions(request: web.Request) -> web.Response | web.St
                         session_key=session_key,
                         channel="api",
                         chat_id=API_CHAT_ID,
+                        hooks=[usage_capture],
                     )
                 response_text = _response_text(response)
                 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 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)
     )
 
 

nanobot/command/builtin.py

@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Literal, cast
 from nanobot import __version__
 from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
 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.restart import set_restart_notice_to_env
 from nanobot.utils.workspace_prompts import initialize_workspace_prompt
@@ -265,8 +266,8 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
             session,
             runtime=runtime,
         )
+    last_usage = LLMUsage.from_dict(session.metadata.get("_last_usage"))
     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
 
     # 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,
         content=build_status_content(
             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,
             session_msg_count=len(session.get_history(max_messages=0)),
             context_tokens_estimate=ctx_est,

nanobot/skills/my/SKILL.md

@@ -1,6 +1,6 @@
 ---
 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

nanobot/skills/my/references/examples.md

@@ -15,8 +15,6 @@ Concrete scenarios showing when and how to use the my tool effectively.
 ```
 → my(action="check", key="max_iterations")
   → 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."
 ```
 
@@ -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="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."
-```