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

refactor(providers): define typed usage contract

Browse Source
This commit is contained in:
chengyongru
2026-08-25 01:04:25 +08:00
committed by chengyongru
parent 89c94d8744
commit 9895c23cb5
84 changed files with 1643 additions and 726 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/agent/test_attachment_references.py
+1 -1
View File
@@ -128,7 +128,7 @@ async def test_pending_document_attachment_keeps_body_out_of_prompt(
nonlocal call_count
call_count += 1
captured_messages.append([dict(message) for message in messages])
return LLMResponse(content=f"answer-{call_count}", tool_calls=[], usage={})
return LLMResponse(content=f"answer-{call_count}", tool_calls=[], usage=None)
loop = _make_loop(workspace)
loop.provider.chat_with_retry = chat_with_retry
tests/agent/test_dream.py
+4 -4
View File
@@ -412,7 +412,7 @@ class TestEphemeralDirect:
provider.supports_tools = True
provider.generation = MagicMock(max_tokens=4096)
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], finish_reason="stop", usage={})
return_value=LLMResponse(content="done", tool_calls=[], finish_reason="stop", usage=None)
)
with (
@@ -556,9 +556,9 @@ class TestEphemeralDirect:
"new_text": "replacement",
},
)],
usage={},
usage=None,
),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage={}),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage=None),
])
resp = await loop.process_direct(
@@ -646,7 +646,7 @@ class TestEphemeralHooks:
provider.generation = MagicMock(max_tokens=4096)
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(
content="done", finish_reason="stop", tool_calls=[], usage={},
content="done", finish_reason="stop", tool_calls=[], usage=None,
)
)
tests/agent/test_evaluator.py
+1 -1
View File
@@ -12,7 +12,7 @@ from nanobot.utils.evaluator import (
class DummyProvider(LLMProvider):
def __init__(self, responses: list[LLMResponse]):
super().__init__()
super().__init__(provider_name="dummy")
self._responses = list(responses)
async def chat(self, *args, **kwargs) -> LLMResponse:
tests/agent/test_history_replay.py
+2 -2
View File
@@ -69,7 +69,7 @@ def test_explicit_message_limit_still_starts_at_user_turn() -> None:
async def test_process_message_replays_with_token_budget_only(tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=32_768)
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[], usage={})
return_value=LLMResponse(content="ok", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
@@ -91,7 +91,7 @@ async def test_process_message_replays_with_token_budget_only(tmp_path: Path) ->
async def test_token_budget_keeps_current_user_as_replay_boundary(tmp_path: Path) -> None:
loop = _make_loop(tmp_path, context_window_tokens=8_000)
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[], usage={})
return_value=LLMResponse(content="ok", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
tests/agent/test_hook_composite.py
+4 -4
View File
@@ -453,7 +453,7 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path):
loop = _make_loop(tmp_path, hooks=[TrackingHook()])
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], usage={})
return_value=LLMResponse(content="done", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
@@ -494,7 +494,7 @@ async def test_agent_loop_turn_hook_factories_receive_context(tmp_path):
loop = _make_loop(tmp_path, hook_factories=[factory("registered")])
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], usage={})
return_value=LLMResponse(content="done", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
@@ -541,7 +541,7 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path):
loop = _make_loop(tmp_path, hooks=[BadHook()])
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="still works", tool_calls=[], usage={})
return_value=LLMResponse(content="still works", tool_calls=[], usage=None)
)
loop.tools.get_definitions = MagicMock(return_value=[])
@@ -562,7 +562,7 @@ async def test_agent_loop_extra_hooks_do_not_swallow_loop_hook_errors(tmp_path):
loop.provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="c1", name="list_dir", arguments={"path": "."})],
usage={},
usage=None,
))
loop.tools.get_definitions = MagicMock(return_value=[])
loop.tools.execute = AsyncMock(return_value="ok")
tests/agent/test_loop_progress.py
+2 -2
View File
@@ -393,9 +393,9 @@ class TestToolEventProgress:
},
)
],
usage={},
usage=None,
)
return LLMResponse(content="Done", tool_calls=[], usage={})
return LLMResponse(content="Done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
tests/agent/test_loop_runner_integration.py
+22 -22
View File
@@ -48,7 +48,7 @@ async def test_ephemeral_runner_enters_and_restores_turn_scopes(tmp_path):
async def chat_with_retry(**_kwargs):
assert goal_mutation_allowed() is True
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
loop.provider.chat_with_retry = AsyncMock(side_effect=chat_with_retry)
loop.tools.get_definitions = MagicMock(return_value=[])
@@ -83,7 +83,7 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
},
)
],
usage={},
usage=None,
),
LLMResponse(
content="closing goal",
@@ -94,7 +94,7 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
arguments={"action": "complete", "recap": "Implemented and tested."},
)
],
usage={},
usage=None,
),
LLMResponse(
content="trying to start another goal",
@@ -105,9 +105,9 @@ async def test_goal_command_can_implement_plan_from_prior_discussion(tmp_path):
arguments={"objective": "Start an unrelated follow-up."},
)
],
usage={},
usage=None,
),
LLMResponse(content="done", tool_calls=[], usage={}),
LLMResponse(content="done", tool_calls=[], usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -160,8 +160,8 @@ async def test_runtime_context_is_persisted_as_next_turn_prompt_prefix(tmp_path)
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="first answer", usage={}),
LLMResponse(content="second answer", usage={}),
LLMResponse(content="first answer", usage=None),
LLMResponse(content="second answer", usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -216,7 +216,7 @@ async def test_webui_quote_reaches_model_without_leaking_into_public_history(tmp
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="answer", usage={}))
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="answer", usage=None))
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
session = loop.sessions.get_or_create("websocket:chat")
@@ -258,9 +258,9 @@ async def test_runtime_context_provider_runs_once_across_tool_iterations(tmp_pat
name="read_file",
arguments={"path": "note.txt"},
)],
usage={},
usage=None,
),
LLMResponse(content="done", usage={}),
LLMResponse(content="done", usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -303,9 +303,9 @@ async def test_non_goal_direct_turn_cannot_reuse_prior_goal_command(tmp_path):
arguments={"objective": "Unauthorized persistent objective."},
)
],
usage={},
usage=None,
),
LLMResponse(content="handled as a one-time task", tool_calls=[], usage={}),
LLMResponse(content="handled as a one-time task", tool_calls=[], usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None)
@@ -383,7 +383,7 @@ async def test_loop_stream_filter_handles_think_only_prefix_without_crashing(tmp
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("<think>hidden")
await on_content_delta("</think>Hello")
return LLMResponse(content="<think>hidden</think>Hello", tool_calls=[], usage={})
return LLMResponse(content="<think>hidden</think>Hello", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -413,7 +413,7 @@ async def test_loop_stream_filter_hides_partial_trailing_think_prefix(tmp_path):
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("Hello <thin")
await on_content_delta("k>hidden</think>World")
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage={})
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -436,7 +436,7 @@ async def test_loop_stream_filter_hides_complete_trailing_think_tag(tmp_path):
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("Hello <think>")
await on_content_delta("hidden</think>World")
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage={})
return LLMResponse(content="Hello <think>hidden</think>World", tool_calls=[], usage=None)
loop.provider.chat_stream_with_retry = chat_stream_with_retry
@@ -459,8 +459,8 @@ async def test_loop_retries_think_only_final_response(tmp_path):
async def chat_with_retry(**kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
return LLMResponse(content="<think>hidden</think>", tool_calls=[], usage={})
return LLMResponse(content="Recovered answer", tool_calls=[], usage={})
return LLMResponse(content="<think>hidden</think>", tool_calls=[], usage=None)
return LLMResponse(content="Recovered answer", tool_calls=[], usage=None)
loop.provider.chat_with_retry = chat_with_retry
@@ -485,7 +485,7 @@ async def test_streamed_flag_not_set_on_llm_error(tmp_path):
provider.get_default_model.return_value = "test-model"
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
error_resp = LLMResponse(
content="503 service unavailable", finish_reason="error", tool_calls=[], usage={},
content="503 service unavailable", finish_reason="error", tool_calls=[], usage=None,
)
loop.provider.chat_with_retry = AsyncMock(return_value=error_resp)
loop.provider.chat_stream_with_retry = AsyncMock(return_value=error_resp)
@@ -523,14 +523,14 @@ async def test_ssrf_soft_block_can_finalize_after_streamed_tool_call(tmp_path):
name="exec",
arguments={"command": "curl http://169.254.169.254/latest/meta-data/"},
)],
usage={},
usage=None,
)
responses = iter([
tool_call_resp,
LLMResponse(
content="I cannot access private URLs. Please share the local file.",
tool_calls=[],
usage={},
usage=None,
),
])
@@ -569,8 +569,8 @@ async def test_next_turn_after_llm_error_keeps_turn_boundary(tmp_path):
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage={}),
LLMResponse(content="Recovered answer", tool_calls=[], usage={}),
LLMResponse(content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage=None),
LLMResponse(content="Recovered answer", tool_calls=[], usage=None),
])
loop = AgentLoop(bus=MessageBus(), provider=provider, workspace=tmp_path, model="test-model")
tests/agent/test_loop_save_turn.py
+5 -6
View File
@@ -20,7 +20,7 @@ from nanobot.bus.outbound_events import (
)
from nanobot.bus.queue import MessageBus
from nanobot.cron.session_turns import CRON_HISTORY_META, CRON_TRIGGER_META
from nanobot.providers.base import LLMProvider, LLMResponse, ProviderConversationState
from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage, ProviderConversationState
from nanobot.providers.factory import ProviderSnapshot
from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
@@ -1883,7 +1883,7 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
async def fake_run_agent_loop(initial_messages, **_kwargs):
loop._last_usage = {"prompt_tokens": 64, "completion_tokens": 9}
loop._last_usage = LLMUsage.reported(input_tokens=64, output_tokens=9)
return (
"done",
[],
@@ -1898,10 +1898,9 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
)
loop.sessions.invalidate("cli:usage")
assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == {
"prompt_tokens": 64,
"completion_tokens": 9,
}
assert loop.sessions.get_or_create("cli:usage").metadata["_last_usage"] == (
LLMUsage.reported(input_tokens=64, output_tokens=9).to_dict()
)
@pytest.mark.asyncio
tests/agent/test_loop_session_policy.py
+1 -1
View File
@@ -30,7 +30,7 @@ def _loop(tmp_path, responses: list[str], **kwargs) -> AgentLoop:
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings()
provider.chat_with_retry = AsyncMock(
side_effect=[LLMResponse(content=response, usage={}) for response in responses]
side_effect=[LLMResponse(content=response, usage=None) for response in responses]
)
return AgentLoop(
bus=MessageBus(),
tests/agent/test_runner_core.py
+190 -31
View File
@@ -14,6 +14,7 @@ from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import (
LLMProvider,
LLMResponse,
LLMUsage,
ProviderCallContext,
ProviderConversationState,
ToolCallRequest,
@@ -22,6 +23,163 @@ from nanobot.providers.base import (
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
def _make_usage_spec(provider, tools):
return make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "hello"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
def test_usage_or_estimate_replaces_reported_zero_for_content(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
monkeypatch.setattr(
"nanobot.agent.runner.estimate_prompt_tokens_chain",
lambda provider, model, messages, definitions: (12, "test"),
)
monkeypatch.setattr("nanobot.agent.runner.estimate_message_tokens", lambda message: 7)
response = LLMResponse(
content="answer",
usage=LLMUsage.reported(input_tokens=0, output_tokens=0),
generation_ms=25,
ttft_ms=5,
)
usage = AgentRunner()._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage == LLMUsage.estimated(input_tokens=12, output_tokens=7).with_timing(
generation_ms=25,
ttft_ms=5,
)
assert usage.source == "estimated"
assert usage.total_tokens == 19
def test_usage_or_estimate_counts_tool_call_output_for_reported_zero(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
captured_message: dict = {}
monkeypatch.setattr(
"nanobot.agent.runner.estimate_prompt_tokens_chain",
lambda provider, model, messages, definitions: (13, "test"),
)
def estimate_output(message):
captured_message.update(message)
return 9
monkeypatch.setattr("nanobot.agent.runner.estimate_message_tokens", estimate_output)
response = LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="lookup",
arguments={"query": "nanobot"},
)
],
finish_reason="tool_calls",
usage=LLMUsage.reported(input_tokens=0, output_tokens=0),
)
usage = AgentRunner()._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage == LLMUsage.estimated(input_tokens=13, output_tokens=9)
assert usage.total_tokens == 22
assert captured_message["tool_calls"][0]["function"]["name"] == "lookup"
@pytest.mark.parametrize(
"provider_usage",
[None, LLMUsage.reported(input_tokens=0, output_tokens=0)],
)
def test_usage_or_estimate_counts_error_without_estimating_tokens(
monkeypatch,
provider_usage: LLMUsage | None,
) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
estimate = MagicMock()
runner = AgentRunner()
monkeypatch.setattr(runner, "_estimate_response_usage", estimate)
response = LLMResponse(
content="upstream failed",
finish_reason="error",
usage=provider_usage,
)
usage = runner._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage is not None
assert usage.total_tokens == 0
assert usage.request_count == 1
assert usage.context_tokens is None
aggregate = LLMUsage.reported(input_tokens=12, output_tokens=3) + usage
assert aggregate.context_tokens == 12
assert aggregate.request_count == 2
estimate.assert_not_called()
def test_usage_or_estimate_trusts_positive_reported_total(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
estimate = MagicMock()
runner = AgentRunner()
monkeypatch.setattr(runner, "_estimate_response_usage", estimate)
response = LLMResponse(
content="answer",
usage=LLMUsage.reported(
input_tokens=15,
output_tokens=18,
total_tokens=175,
),
generation_ms=30,
ttft_ms=6,
)
usage = runner._usage_or_estimate(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
)
assert usage is not None
assert usage.source == "reported"
assert usage.input_tokens == 15
assert usage.output_tokens == 18
assert usage.total_tokens == 175
assert usage.reported_tokens == 175
assert usage.generation_ms == 30
assert usage.ttft_ms == 6
estimate.assert_not_called()
@pytest.mark.asyncio
async def test_runner_preserves_reasoning_fields_and_tool_results():
from nanobot.agent.runner import AgentRunner
@@ -38,10 +196,10 @@ async def test_runner_preserves_reasoning_fields_and_tool_results():
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
reasoning_content="hidden reasoning",
thinking_blocks=[{"type": "thinking", "thinking": "step"}],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -441,7 +599,7 @@ async def test_runner_uses_no_tools_finalization_after_max_iterations():
return LLMResponse(
content="Read the directory twice. More investigation remains.",
tool_calls=[],
usage={"prompt_tokens": 10, "completion_tokens": 7},
usage=LLMUsage.reported(input_tokens=10, output_tokens=7),
)
provider.chat_with_retry = chat_with_retry
@@ -713,10 +871,10 @@ async def test_runner_replaces_empty_tool_result_with_marker():
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="noop", arguments={})],
usage={},
usage=None,
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -751,12 +909,12 @@ async def test_runner_retries_empty_final_response_with_summary_prompt():
return LLMResponse(
content=None,
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 1},
usage=LLMUsage.reported(input_tokens=5, output_tokens=1),
)
return LLMResponse(
content="final answer",
tool_calls=[],
usage={"prompt_tokens": 3, "completion_tokens": 7},
usage=LLMUsage.reported(input_tokens=3, output_tokens=7),
)
provider.chat_with_retry = chat_with_retry
@@ -778,8 +936,9 @@ async def test_runner_retries_empty_final_response_with_summary_prompt():
assert calls[0]["tools"] is not None
assert calls[1]["tools"] is not None
assert calls[2]["tools"] is None
assert result.usage["prompt_tokens"] == 13
assert result.usage["completion_tokens"] == 9
assert result.usage is not None
assert result.usage.input_tokens == 13
assert result.usage.output_tokens == 9
@pytest.mark.asyncio
@@ -851,7 +1010,7 @@ async def test_runner_uses_specific_message_after_empty_finalization_retry():
provider = MagicMock(spec=LLMProvider)
async def chat_with_retry(*, messages, **kwargs):
return LLMResponse(content=None, tool_calls=[], usage={})
return LLMResponse(content=None, tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -891,14 +1050,14 @@ async def test_empty_finalization_retry_discards_candidate_provider_state():
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = True
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content=None, tool_calls=[], usage={}),
LLMResponse(content=None, tool_calls=[], usage={}),
LLMResponse(content=None, tool_calls=[], usage=None),
LLMResponse(content=None, tool_calls=[], usage=None),
LLMResponse(
content="finalized without tools",
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})],
finish_reason="stop",
provider_state=candidate,
usage={},
usage=None,
),
])
tools = MagicMock()
@@ -1037,20 +1196,20 @@ async def test_runner_empty_response_does_not_break_tool_chain():
return LLMResponse(
content=None,
tool_calls=[ToolCallRequest(id="tc1", name="read_file", arguments={"path": "a.txt"})],
usage={"prompt_tokens": 10, "completion_tokens": 5},
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
)
if call_count == 2:
return LLMResponse(content=None, tool_calls=[], usage={"prompt_tokens": 10, "completion_tokens": 1})
return LLMResponse(content=None, tool_calls=[], usage=LLMUsage.reported(input_tokens=10, output_tokens=1))
if call_count == 3:
return LLMResponse(
content=None,
tool_calls=[ToolCallRequest(id="tc2", name="read_file", arguments={"path": "b.txt"})],
usage={"prompt_tokens": 10, "completion_tokens": 5},
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
)
return LLMResponse(
content="Here are the results.",
tool_calls=[],
usage={"prompt_tokens": 10, "completion_tokens": 10},
usage=LLMUsage.reported(input_tokens=10, output_tokens=10),
)
provider.chat_with_retry = chat_with_retry
@@ -1079,9 +1238,8 @@ async def test_runner_empty_response_does_not_break_tool_chain():
@pytest.mark.asyncio
async def test_runner_accumulates_usage_and_preserves_cached_tokens():
"""Runner should accumulate prompt/completion tokens across iterations
and preserve cached_tokens from provider responses."""
async def test_runner_accumulates_usage_and_preserves_cache_reads():
"""Runner accumulates usage across iterations, including cache reads."""
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
@@ -1093,12 +1251,12 @@ async def test_runner_accumulates_usage_and_preserves_cached_tokens():
return LLMResponse(
content="thinking",
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})],
usage={"prompt_tokens": 100, "completion_tokens": 10, "cached_tokens": 80},
usage=LLMUsage.reported(input_tokens=100, output_tokens=10, cache_read_tokens=80),
)
return LLMResponse(
content="done",
tool_calls=[],
usage={"prompt_tokens": 200, "completion_tokens": 20, "cached_tokens": 150},
usage=LLMUsage.reported(input_tokens=200, output_tokens=20, cache_read_tokens=150),
)
provider.chat_with_retry = chat_with_retry
@@ -1116,11 +1274,12 @@ async def test_runner_accumulates_usage_and_preserves_cached_tokens():
))
# Usage should be accumulated across iterations
assert result.usage["prompt_tokens"] == 300 # 100 + 200
assert result.usage["completion_tokens"] == 30 # 10 + 20
assert result.usage["cached_tokens"] == 230 # 80 + 150
assert result.usage["context_tokens"] == 200
assert result.usage["request_count"] == 2
assert result.usage is not None
assert result.usage.input_tokens == 300 # 100 + 200
assert result.usage.output_tokens == 30 # 10 + 20
assert result.usage.cache_read_tokens == 230 # 80 + 150
assert result.usage.context_tokens == 200
assert result.usage.request_count == 2
@pytest.mark.asyncio
@@ -1137,7 +1296,7 @@ async def test_runner_binds_on_retry_wait_to_retry_callback_not_progress():
async def chat_with_retry(**kwargs):
captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry
@@ -1179,7 +1338,7 @@ async def test_runner_passes_temperature_to_provider():
async def chat_with_retry(**kwargs):
captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry
@@ -1208,7 +1367,7 @@ async def test_runner_passes_max_tokens_to_provider():
async def chat_with_retry(**kwargs):
captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry
@@ -1237,7 +1396,7 @@ async def test_runner_passes_reasoning_effort_to_provider():
async def chat_with_retry(**kwargs):
captured.update(kwargs)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = chat_with_retry
tests/agent/test_runner_errors.py
+8 -8
View File
@@ -90,7 +90,7 @@ async def test_llm_error_not_appended_to_session_messages():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage={},
content="429 rate limit exceeded", finish_reason="error", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -158,7 +158,7 @@ async def test_runner_ignores_tool_calls_when_finish_reason_blocks_execution(
content="Request blocked by provider policy.",
finish_reason=finish_reason,
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={"command": "echo nope"})],
usage={},
usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -189,7 +189,7 @@ async def test_runner_tool_error_sets_final_content():
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})],
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
@@ -224,9 +224,9 @@ async def test_runner_preserves_successful_exec_output_that_starts_with_error():
tool_calls=[
ToolCallRequest(id="call_1", name="exec", arguments={"command": "report"})
],
usage={},
usage=None,
)
return LLMResponse(content="done", usage={})
return LLMResponse(content="done", usage=None)
provider.chat_with_retry = chat_with_retry
output = "Error: generated report successfully\n\nExit code: 0"
@@ -266,7 +266,7 @@ async def test_runner_tool_error_preserves_tool_results_in_messages():
ToolCallRequest(id="tc1", name="read_file", arguments={"path": "a"}),
ToolCallRequest(id="tc2", name="exec", arguments={"cmd": "bad"}),
],
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
@@ -333,9 +333,9 @@ async def test_length_finish_with_blank_content_routes_to_length_recovery():
content="",
finish_reason="length",
tool_calls=[ToolCallRequest(id="call_1", name="exec", arguments={})],
usage={},
usage=None,
),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage={}),
LLMResponse(content="done", finish_reason="stop", tool_calls=[], usage=None),
])
tools = MagicMock()
tools.get_definitions.return_value = []
tests/agent/test_runner_fallback.py
+36 -1
View File
@@ -78,7 +78,7 @@ class _FakeProvider(LLMProvider):
*,
responses: list[LLMResponse] | None = None,
):
super().__init__()
super().__init__(provider_name=name)
self.name = name
self._response = response or _make_response()
self._responses = iter(responses) if responses is not None else None
@@ -260,6 +260,41 @@ def test_provider_snapshot_uses_smallest_fallback_context_window() -> None:
assert snapshot.provider._primary_context_window_tokens == 128000
def test_factory_injects_configured_identity_into_primary_and_fallback_leaves() -> None:
from nanobot.config.schema import Config
from nanobot.providers.factory import build_provider_snapshot
config = Config.model_validate({
"agents": {
"defaults": {
"modelPreset": "primary",
"fallbackModels": ["backup"],
}
},
"modelPresets": {
"primary": {"model": "primary-model", "provider": "primary_edge"},
"backup": {"model": "backup-model", "provider": "backup_edge"},
},
"providers": {
"primary_edge": {
"apiKey": "primary-key",
"apiBase": "https://primary.example/v1",
},
"backup_edge": {
"apiKey": "backup-key",
"apiBase": "https://backup.example/v1",
},
},
})
snapshot = build_provider_snapshot(config)
assert isinstance(snapshot.provider, FallbackProvider)
assert snapshot.provider._primary.provider_name == "primary_edge"
fallback = snapshot.provider._provider_factory(snapshot.provider._fallback_presets[0])
assert fallback.provider_name == "backup_edge"
def test_inline_fallback_reasoning_effort_does_not_inherit_primary() -> None:
from nanobot.config.schema import Config
from nanobot.providers.factory import provider_signature
tests/agent/test_runner_goal_continue.py
+8 -8
View File
@@ -25,7 +25,7 @@ async def test_runner_exits_normally_without_predicate():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="all done", tool_calls=[], usage={},
content="all done", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -50,7 +50,7 @@ async def test_runner_exits_normally_with_inactive_goal():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="all done", tool_calls=[], usage={},
content="all done", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -82,7 +82,7 @@ async def test_runner_forces_continue_when_goal_active():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -112,7 +112,7 @@ async def test_runner_respects_max_iterations_even_with_active_goal():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -137,7 +137,7 @@ async def test_runner_goal_continue_not_limited_by_injection_cycle_cap():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -165,7 +165,7 @@ async def test_runner_does_not_force_continue_on_error():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content=None, tool_calls=[], usage={},
content=None, tool_calls=[], usage=None,
finish_reason="error",
))
tools = MagicMock()
@@ -191,7 +191,7 @@ async def test_runner_uses_custom_goal_continue_message():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
@@ -220,7 +220,7 @@ async def test_runner_resolves_goal_continue_message_lazily():
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="still working", tool_calls=[], usage={},
content="still working", tool_calls=[], usage=None,
))
tools = MagicMock()
tools.get_definitions.return_value = []
tests/agent/test_runner_governance.py
+3 -3
View File
@@ -273,7 +273,7 @@ async def test_runner_drops_orphan_tool_results_before_model_request():
async def chat_with_retry(*, messages, **kwargs):
captured_messages[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -312,7 +312,7 @@ async def test_backfill_repairs_model_context_without_shifting_save_turn_boundar
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
response = LLMResponse(content="new answer", tool_calls=[], usage={})
response = LLMResponse(content="new answer", tool_calls=[], usage=None)
provider.chat_with_retry = AsyncMock(return_value=response)
provider.chat_stream_with_retry = AsyncMock(return_value=response)
@@ -397,7 +397,7 @@ async def test_runner_backfill_only_mutates_model_context_not_returned_messages(
async def chat_with_retry(*, messages, **kwargs):
captured_messages[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tests/agent/test_runner_hooks.py
+27 -34
View File
@@ -9,7 +9,7 @@ import pytest
from agent.runner_helpers import make_run_spec
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage, ToolCallRequest
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
@@ -30,7 +30,7 @@ async def test_runner_calls_hooks_in_order():
content="thinking",
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -110,7 +110,7 @@ async def test_runner_streaming_hook_receives_deltas_and_end_signal():
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("he")
await on_content_delta("llo")
return LLMResponse(content="hello", tool_calls=[], usage={})
return LLMResponse(content="hello", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -155,7 +155,7 @@ async def test_runner_measures_stream_generation_without_time_to_first_token():
await on_content_delta("llo")
return LLMResponse(
content="hello",
usage={"prompt_tokens": 100, "completion_tokens": 12},
usage=LLMUsage.reported(input_tokens=100, output_tokens=12),
)
provider.chat_stream_with_retry = chat_stream_with_retry
@@ -181,10 +181,11 @@ async def test_runner_measures_stream_generation_without_time_to_first_token():
hook=StreamingHook(),
))
assert result.usage["generation_ms"] == 600
assert result.usage["measured_completion_tokens"] == 12
assert result.usage["ttft_ms"] == 200
assert result.usage["timed_requests"] == 1
assert result.usage is not None
assert result.usage.generation_ms == 600
assert result.usage.measured_output_tokens == 12
assert result.usage.ttft_ms == 200
assert result.usage.timed_requests == 1
@pytest.mark.asyncio
@@ -240,23 +241,24 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont
@pytest.mark.asyncio
async def test_runner_passes_cached_tokens_to_hook_context():
"""Hook context.usage should contain cached_tokens."""
async def test_runner_passes_cache_read_tokens_to_hook_context():
"""Hook context usage preserves a reported cache-read count."""
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
captured_usage: list[dict] = []
captured_usage: list[LLMUsage] = []
class UsageHook(AgentHook):
async def after_iteration(self, context: AgentHookContext) -> None:
captured_usage.append(dict(context.usage))
assert context.usage is not None
captured_usage.append(context.usage)
async def chat_with_retry(**kwargs):
return LLMResponse(
content="done",
tool_calls=[],
usage={"prompt_tokens": 200, "completion_tokens": 20, "cached_tokens": 150},
usage=LLMUsage.reported(input_tokens=200, output_tokens=20, cache_read_tokens=150),
)
provider.chat_with_retry = chat_with_retry
@@ -274,8 +276,8 @@ async def test_runner_passes_cached_tokens_to_hook_context():
))
assert len(captured_usage) == 1
assert captured_usage[0]["cached_tokens"] == 150
assert captured_usage[0]["provider_tokens"] == 220
assert captured_usage[0].cache_read_tokens == 150
assert captured_usage[0].reported_tokens == 220
@pytest.mark.asyncio
@@ -284,14 +286,15 @@ async def test_runner_estimates_usage_when_provider_omits_usage(monkeypatch):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
captured_usage: list[dict] = []
captured_usage: list[LLMUsage] = []
class UsageHook(AgentHook):
async def after_iteration(self, context: AgentHookContext) -> None:
captured_usage.append(dict(context.usage))
assert context.usage is not None
captured_usage.append(context.usage)
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -312,11 +315,8 @@ async def test_runner_estimates_usage_when_provider_omits_usage(monkeypatch):
hook=UsageHook(),
))
assert result.usage["prompt_tokens"] == 123
assert result.usage["completion_tokens"] == 7
assert result.usage["total_tokens"] == 130
assert result.usage["estimated_tokens"] == 130
assert captured_usage[0]["estimated_tokens"] == 130
assert result.usage == LLMUsage.estimated(input_tokens=123, output_tokens=7)
assert captured_usage[0].estimated_tokens == 130
@pytest.mark.asyncio
@@ -332,7 +332,7 @@ async def test_runner_calls_run_level_hooks_on_success():
return LLMResponse(
content="done",
tool_calls=[],
usage={"prompt_tokens": 3, "completion_tokens": 2},
usage=LLMUsage.reported(input_tokens=3, output_tokens=2),
)
provider.chat_with_retry = chat_with_retry
@@ -350,7 +350,7 @@ async def test_runner_calls_run_level_hooks_on_success():
context.final_content,
context.stop_reason,
context.error,
dict(context.usage),
context.usage,
[msg["role"] for msg in context.messages],
))
@@ -379,14 +379,7 @@ async def test_runner_calls_run_level_hooks_on_success():
"done",
"completed",
None,
{
"prompt_tokens": 3,
"completion_tokens": 2,
"total_tokens": 5,
"provider_tokens": 5,
"request_count": 1,
"context_tokens": 3,
},
LLMUsage.reported(input_tokens=3, output_tokens=2),
["user", "assistant"],
),
("on_finally", "completed", None),
@@ -410,7 +403,7 @@ async def test_runner_run_level_context_is_detached_snapshot():
content="thinking",
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tests/agent/test_runner_injections.py
+28 -28
View File
@@ -263,9 +263,9 @@ async def test_checkpoint1_injects_after_tool_execution():
return LLMResponse(
content="using tool",
tool_calls=[ToolCallRequest(id="c1", name="read_file", arguments={"path": "x"})],
usage={},
usage=None,
)
return LLMResponse(content="final answer", tool_calls=[], usage={})
return LLMResponse(content="final answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -323,8 +323,8 @@ async def test_checkpoint2_injects_after_final_response_with_resuming_stream():
async def chat_stream_with_retry(*, messages, on_content_delta=None, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
@@ -411,8 +411,8 @@ async def test_checkpoint2_preserves_final_response_in_history_before_followup()
call_count["n"] += 1
captured_messages.append([dict(message) for message in messages])
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -474,8 +474,8 @@ async def test_loop_injected_followup_preserves_image_media(tmp_path):
call_count["n"] += 1
captured_messages.append(list(messages))
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
@@ -528,8 +528,8 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path):
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="first answer", tool_calls=[], usage={}),
LLMResponse(content="second answer", tool_calls=[], usage={}),
LLMResponse(content="first answer", tool_calls=[], usage=None),
LLMResponse(content="second answer", tool_calls=[], usage=None),
])
loop = AgentLoop(
bus=MessageBus(),
@@ -653,8 +653,8 @@ async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_p
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
@@ -714,8 +714,8 @@ async def test_runner_merges_multiple_injected_user_messages_without_losing_medi
call_count["n"] += 1
captured_messages.append([dict(message) for message in messages])
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
return LLMResponse(content="first answer", tool_calls=[], usage=None)
return LLMResponse(content="second answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -841,7 +841,7 @@ async def test_injection_cycles_capped_at_max():
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage={})
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -879,7 +879,7 @@ async def test_no_injections_flag_is_false_by_default():
provider = MagicMock()
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -903,7 +903,7 @@ async def test_pending_queue_cleanup_on_dispatch(tmp_path):
loop = _make_loop(tmp_path)
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
loop.provider.chat_with_retry = chat_with_retry
@@ -1329,7 +1329,7 @@ async def test_pending_queue_preserves_overflow_for_next_injection_cycle(tmp_pat
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
captured_messages.append([dict(message) for message in messages])
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage={})
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
@@ -1502,16 +1502,16 @@ async def test_drain_injections_on_fatal_tool_error():
return LLMResponse(
content="stale prefix ",
finish_reason="length",
usage={},
usage=None,
)
if call_count["n"] == 2:
return LLMResponse(
content="",
tool_calls=[ToolCallRequest(id="c1", name="exec", arguments={"cmd": "bad"})],
usage={},
usage=None,
)
# Third call: respond normally to the injected follow-up.
return LLMResponse(content="reply to follow-up", tool_calls=[], usage={})
return LLMResponse(content="reply to follow-up", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -1563,10 +1563,10 @@ async def test_drain_injections_on_llm_error():
content=None,
tool_calls=[],
finish_reason="error",
usage={},
usage=None,
)
# Second call: respond normally to the injected follow-up
return LLMResponse(content="recovered answer", tool_calls=[], usage={})
return LLMResponse(content="recovered answer", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -1614,9 +1614,9 @@ async def test_drain_injections_on_empty_final_response():
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
if call_count["n"] <= _MAX_EMPTY_RETRIES + 1:
return LLMResponse(content="", tool_calls=[], usage={})
return LLMResponse(content="", tool_calls=[], usage=None)
# After retries exhausted + injection drain, respond normally
return LLMResponse(content="answer after empty", tool_calls=[], usage={})
return LLMResponse(content="answer after empty", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -1671,7 +1671,7 @@ async def test_drain_injections_on_max_iterations():
return LLMResponse(
content="",
tool_calls=[ToolCallRequest(id=f"c{call_count['n']}", name="read_file", arguments={"path": "x"})],
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
@@ -1723,7 +1723,7 @@ async def test_drain_injections_set_flag_when_followup_arrives_after_last_iterat
return LLMResponse(
content="",
tool_calls=[ToolCallRequest(id=f"c{call_count['n']}", name="read_file", arguments={"path": "x"})],
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
@@ -1786,7 +1786,7 @@ async def test_injection_cycle_cap_on_error_path():
content=None,
tool_calls=[],
finish_reason="error",
usage={},
usage=None,
)
provider.chat_with_retry = chat_with_retry
tests/agent/test_runner_persistence.py
+7 -7
View File
@@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
from agent.runner_helpers import make_run_spec
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMResponse, ToolCallRequest
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
@@ -25,10 +25,10 @@ async def test_runner_persists_large_tool_results_for_follow_up_calls(tmp_path):
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_big", name="list_dir", arguments={"path": "."})],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -138,10 +138,10 @@ async def test_read_file_result_is_not_offloaded(tmp_path):
return LLMResponse(
content="reading",
tool_calls=[ToolCallRequest(id="call_rf", name="read_file", arguments={"path": "big.txt"})],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -183,10 +183,10 @@ async def test_runner_keeps_going_when_tool_result_persistence_fails():
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tests/agent/test_runner_progress_deltas.py
+10 -10
View File
@@ -23,7 +23,7 @@ async def test_runner_can_disable_provider_progress_delta_streaming():
provider = MagicMock()
provider.supports_progress_deltas = True
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="done", tool_calls=[], usage={})
return_value=LLMResponse(content="done", tool_calls=[], usage=None)
)
provider.chat_stream_with_retry = AsyncMock()
tools = MagicMock()
@@ -59,7 +59,7 @@ async def test_runner_streams_provider_progress_deltas_by_default():
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("he")
await on_content_delta("llo")
return LLMResponse(content="hello", tool_calls=[], usage={})
return LLMResponse(content="hello", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -113,7 +113,7 @@ async def test_runner_routes_hosted_tool_events_to_structured_progress():
"result": {"name": "x_semantic_search"},
})
await on_content_delta("done")
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -264,9 +264,9 @@ async def test_runner_emits_write_file_diff_from_tool_execution_snapshots(tmp_pa
arguments={"path": "big.txt", "content": "line\n" * 24},
)
],
usage={},
usage=None,
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -338,9 +338,9 @@ async def test_runner_emits_edit_file_diff_from_tool_execution_snapshots(tmp_pat
},
)
],
usage={},
usage=None,
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -404,9 +404,9 @@ async def test_runner_marks_file_edit_activity_failed_when_tool_errors(tmp_path)
arguments={"path": "aborted.txt"},
)
],
usage={},
usage=None,
)
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
@@ -469,7 +469,7 @@ async def test_runner_marks_file_edit_activity_failed_when_cancelled(tmp_path):
arguments={"path": "cancelled.txt", "content": "new\n"},
)
],
usage={},
usage=None,
)
provider.chat_stream_with_retry = chat_stream_with_retry
tests/agent/test_runner_reasoning.py
+12 -12
View File
@@ -16,7 +16,7 @@ import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMResponse, ToolCallRequest
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
@@ -53,10 +53,10 @@ async def test_runner_preserves_reasoning_fields_in_assistant_history():
tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={"path": "."})],
reasoning_content="hidden reasoning",
thinking_blocks=[{"type": "thinking", "thinking": "step"}],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
@@ -99,7 +99,7 @@ async def test_runner_emits_anthropic_thinking_blocks():
{"type": "thinking", "thinking": "After careful consideration.", "signature": "sig2"},
],
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_with_retry = chat_with_retry
@@ -135,7 +135,7 @@ async def test_runner_emits_inline_think_content_as_reasoning():
return LLMResponse(
content="<think>Let me think about this...\nThe answer is 42.</think>The answer is 42.",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_with_retry = chat_with_retry
@@ -171,7 +171,7 @@ async def test_runner_prefers_reasoning_content_over_inline_think():
content="<think>inline thinking</think>The answer.",
reasoning_content="dedicated reasoning field",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_with_retry = chat_with_retry
@@ -211,7 +211,7 @@ async def test_runner_emits_reasoning_content_even_when_answer_was_streamed():
content="The answer.",
reasoning_content="step-by-step deduction",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_stream_with_retry = chat_stream_with_retry
@@ -257,7 +257,7 @@ async def test_runner_does_not_double_emit_when_inline_think_already_streamed():
return LLMResponse(
content="<think>working...</think>The answer.",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_stream_with_retry = chat_stream_with_retry
@@ -299,7 +299,7 @@ async def test_runner_closes_reasoning_stream_after_one_shot_response():
content="answer",
reasoning_content="hidden thought",
tool_calls=[],
usage={"prompt_tokens": 5, "completion_tokens": 3},
usage=LLMUsage.reported(input_tokens=5, output_tokens=3),
)
provider.chat_with_retry = chat_with_retry
@@ -350,7 +350,7 @@ async def test_runner_streams_native_thinking_deltas_without_post_hoc_dup():
content="done",
tool_calls=[],
thinking_blocks=[{"type": "thinking", "thinking": "part1part2"}],
usage={"prompt_tokens": 1, "completion_tokens": 2},
usage=LLMUsage.reported(input_tokens=1, output_tokens=2),
)
provider.chat_stream_with_retry = chat_stream_with_retry
@@ -387,7 +387,7 @@ async def test_runner_strips_thinking_tags_from_native_thinking_deltas():
await on_thinking_delta("</thinking>")
if on_content_delta:
await on_content_delta("done")
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
@@ -425,7 +425,7 @@ async def test_runner_ignores_empty_thinking_marker_before_final_reasoning():
content="done",
reasoning_content="Preparing final response",
tool_calls=[],
usage={},
usage=None,
)
provider.chat_stream_with_retry = chat_stream_with_retry
tests/agent/test_runner_tool_execution.py
+8 -8
View File
@@ -106,7 +106,7 @@ async def _run_optional_tool_response(response: LLMResponse):
calls["n"] += 1
if calls["n"] == 1:
return response
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = ToolRegistry()
@@ -272,10 +272,10 @@ async def test_runner_rejects_near_miss_tool_name_without_executing():
)
],
finish_reason="tool_calls",
usage={},
usage=None,
)
captured_second_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = ToolRegistry()
@@ -403,7 +403,7 @@ async def test_runner_treats_legacy_entry_point_error_prefix_as_tool_error(tmp_p
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="legacy_plugin", arguments={})],
usage={},
usage=None,
))
result = await AgentRunner().run(make_run_spec(provider,
@@ -430,9 +430,9 @@ async def test_runner_preserves_structured_plugin_success_that_starts_with_error
tool_calls=[
ToolCallRequest(id="call_1", name="structured_success_plugin", arguments={})
],
usage={},
usage=None,
),
LLMResponse(content="done", tool_calls=[], usage={}),
LLMResponse(content="done", tool_calls=[], usage=None),
])
result = await AgentRunner().run(make_run_spec(provider,
@@ -466,10 +466,10 @@ async def test_runner_blocks_repeated_external_fetches():
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id=f"call_{call_count['n']}", name="web_fetch", arguments={"url": "https://example.com"})],
usage={},
usage=None,
)
captured_final_call[:] = messages
return LLMResponse(content="done", tool_calls=[], usage={})
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tests/agent/test_session_inputs.py
+1 -1
View File
@@ -18,7 +18,7 @@ def _loop(tmp_path: Path) -> AgentLoop:
provider.get_default_model.return_value = "test-model"
provider.generation = SimpleNamespace(max_tokens=4096)
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="Reviewed", tool_calls=[], usage={})
return_value=LLMResponse(content="Reviewed", tool_calls=[], usage=None)
)
return AgentLoop(
bus=MessageBus(),
tests/agent/test_session_model_runtime.py
+1 -1
View File
@@ -18,7 +18,7 @@ from nanobot.utils.llm_runtime import LLMRuntime
class RecordingProvider(LLMProvider):
def __init__(self, name: str) -> None:
super().__init__()
super().__init__(provider_name=name)
self.name = name
self.generation = GenerationSettings(max_tokens=256, temperature=0.1)
self.calls: list[str | None] = []
tests/agent/test_subagent_lifecycle.py
+5 -5
View File
@@ -16,7 +16,7 @@ from nanobot.agent.subagent import (
)
from nanobot.agent.tools.context import current_request_context
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import GenerationSettings, LLMProvider
from nanobot.providers.base import GenerationSettings, LLMProvider, LLMUsage
from nanobot.utils.llm_runtime import LLMRuntime
# ---------------------------------------------------------------------------
@@ -49,7 +49,7 @@ def _make_hook_context(**overrides) -> AgentHookContext:
tool_calls=[],
tool_events=[],
messages=[],
usage={},
usage=None,
error=None,
stop_reason="completed",
final_content="ok",
@@ -97,7 +97,7 @@ class TestSubagentStatus:
assert s.phase == "initializing"
assert s.iteration == 0
assert s.tool_events == []
assert s.usage == {}
assert s.usage is None
assert s.stop_reason is None
assert s.error is None
@@ -677,12 +677,12 @@ class TestSubagentHook:
ctx = _make_hook_context(
iteration=3,
tool_events=[{"name": "read_file", "status": "ok", "detail": ""}],
usage={"prompt_tokens": 100},
usage=LLMUsage.reported(input_tokens=100, output_tokens=0),
)
await hook.after_iteration(ctx)
assert status.iteration == 3
assert len(status.tool_events) == 1
assert status.usage == {"prompt_tokens": 100}
assert status.usage == LLMUsage.reported(input_tokens=100, output_tokens=0)
@pytest.mark.asyncio
async def test_after_iteration_no_status_noop(self):
tests/agent/tools/test_self_tool.py
+10 -9
View File
@@ -15,6 +15,7 @@ from nanobot.agent.tools.self import MyTool
from nanobot.agent.tools.shell import ExecToolConfig
from nanobot.agent.tools.web import WebSearchConfig, WebToolsConfig
from nanobot.config.schema import ModelPresetConfig
from nanobot.providers.base import LLMUsage
# ---------------------------------------------------------------------------
# Helpers
@@ -31,7 +32,7 @@ def _make_mock_loop(**overrides):
loop._start_time = 1000.0
loop.exec_config = ExecToolConfig()
loop.channels_config = MagicMock()
loop._last_usage = {"prompt_tokens": 100, "completion_tokens": 50}
loop._last_usage = LLMUsage.reported(input_tokens=100, output_tokens=50)
loop.last_usage = loop._last_usage
loop._current_iteration = 0
loop.current_iteration = loop._current_iteration
@@ -163,9 +164,9 @@ class TestInspectPathNavigation:
@pytest.mark.asyncio
async def test_inspect_dict_key_via_dotpath(self):
loop = _make_mock_loop()
loop._last_usage = {"prompt_tokens": 100, "completion_tokens": 50}
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.prompt_tokens")
result = await tool.execute(action="check", key="_last_usage.input_tokens")
assert "100" in result
@pytest.mark.asyncio
@@ -624,7 +625,7 @@ class TestSubagentStatusFormatting:
{"name": "grep", "status": "ok", "detail": "searched ERROR"},
{"name": "exec", "status": "error", "detail": "timeout"},
],
usage={"prompt_tokens": 4500, "completion_tokens": 1200},
usage=LLMUsage.reported(input_tokens=4500, output_tokens=1200),
)
result = MyTool._format_value(status)
assert "abc12345" in result
@@ -698,14 +699,14 @@ class TestSubagentHookStatus:
iteration=5,
messages=[],
tool_events=[{"name": "read_file", "status": "ok", "detail": "ok"}],
usage={"prompt_tokens": 100, "completion_tokens": 50},
usage=LLMUsage.reported(input_tokens=100, output_tokens=50),
)
await hook.after_iteration(context)
assert status.iteration == 5
assert len(status.tool_events) == 1
assert status.tool_events[0]["name"] == "read_file"
assert status.usage == {"prompt_tokens": 100, "completion_tokens": 50}
assert status.usage == LLMUsage.reported(input_tokens=100, output_tokens=50)
@pytest.mark.asyncio
async def test_after_iteration_with_error(self):
@@ -821,7 +822,7 @@ class TestInspectTaskStatuses:
phase="awaiting_tools",
iteration=2,
tool_events=[{"name": "read_file", "status": "ok", "detail": "ok"}],
usage={"prompt_tokens": 500, "completion_tokens": 100},
usage=LLMUsage.reported(input_tokens=500, output_tokens=100),
),
}
tool = _make_tool(loop=loop)
@@ -1127,12 +1128,12 @@ class TestLastUsageInSummary:
tool = _make_tool()
result = await tool.execute(action="check")
assert "_last_usage" in result
assert "prompt_tokens" 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 = {}
loop._last_usage = None
loop.last_usage = loop._last_usage
tool = _make_tool(loop=loop)
result = await tool.execute(action="check")
tests/agent/tools/test_subagent_tools.py
+4 -4
View File
@@ -456,7 +456,7 @@ async def test_agent_loop_syncs_updated_max_iterations_before_run(tmp_path):
error=None,
tool_events=[],
messages=[],
usage={},
usage=None,
had_injections=False,
tools_used=[],
)
@@ -501,7 +501,7 @@ async def test_drain_pending_blocks_while_subagents_running(tmp_path):
error=None,
tool_events=[],
messages=[],
usage={},
usage=None,
had_injections=False,
tools_used=[],
provider_state=None,
@@ -587,7 +587,7 @@ async def test_drain_pending_no_block_when_no_subagents(tmp_path):
error=None,
tool_events=[],
messages=[],
usage={},
usage=None,
had_injections=False,
tools_used=[],
provider_state=None,
@@ -637,7 +637,7 @@ async def test_drain_pending_timeout(tmp_path):
error=None,
tool_events=[],
messages=[],
usage={},
usage=None,
had_injections=False,
tools_used=[],
provider_state=None,
tests/bus/test_runtime_events.py
+4 -3
View File
@@ -12,6 +12,7 @@ from nanobot.bus.runtime_events import (
TurnRunStatusChanged,
TurnRuntimeAdmitted,
)
from nanobot.providers.base import LLMUsage
@pytest.mark.asyncio
@@ -99,7 +100,7 @@ async def test_runtime_event_publisher_consumes_turn_metadata_on_complete() -> N
bus.subscribe(seen.append)
publisher.record_turn_runtime("cli:direct", "runtime")
publisher.record_turn_latency("cli:direct", 123)
publisher.record_turn_usage("cli:direct", {"prompt_tokens": 40, "completion_tokens": 2})
publisher.record_turn_usage("cli:direct", LLMUsage.reported(input_tokens=40, output_tokens=2))
await publisher.turn_completed(
channel="cli",
@@ -120,11 +121,11 @@ async def test_runtime_event_publisher_consumes_turn_metadata_on_complete() -> N
assert first.context.metadata == {"source": "test"}
assert first.latency_ms == 123
assert first.runtime == "runtime"
assert first.usage == {"prompt_tokens": 40, "completion_tokens": 2}
assert first.usage == LLMUsage.reported(input_tokens=40, output_tokens=2)
assert isinstance(second, TurnCompleted)
assert second.latency_ms is None
assert second.runtime is None
assert second.usage == {}
assert second.usage is None
@pytest.mark.asyncio
tests/cli/test_restart_command.py
+7 -10
View File
@@ -12,7 +12,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from nanobot.bus.events import InboundMessage
from nanobot.providers.base import LLMResponse
from nanobot.providers.base import LLMResponse, LLMUsage
def _make_loop():
@@ -238,7 +238,7 @@ class TestRestartCommand:
session.get_history.return_value = [{"role": "user"}] * 3
loop.sessions.get_or_create.return_value = session
loop._start_time = time.time() - 125
loop._last_usage = {"prompt_tokens": 0, "completion_tokens": 0}
loop._last_usage = LLMUsage.reported(input_tokens=0, output_tokens=0)
loop.consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(20500, "tiktoken")
)
@@ -305,18 +305,15 @@ class TestRestartCommand:
lambda _message: 7,
)
loop.provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="first", usage={"prompt_tokens": 9, "completion_tokens": 4}),
LLMResponse(content="second", usage={}),
LLMResponse(content="first", usage=LLMUsage.reported(input_tokens=9, output_tokens=4)),
LLMResponse(content="second", usage=None),
])
await loop._run_agent_loop([], runtime=loop.llm_runtime())
assert loop._last_usage["prompt_tokens"] == 9
assert loop._last_usage["completion_tokens"] == 4
assert loop._last_usage == LLMUsage.reported(input_tokens=9, output_tokens=4)
await loop._run_agent_loop([], runtime=loop.llm_runtime())
assert loop._last_usage["prompt_tokens"] == 123
assert loop._last_usage["completion_tokens"] == 7
assert loop._last_usage["estimated_tokens"] == 130
assert loop._last_usage == LLMUsage.estimated(input_tokens=123, output_tokens=7)
@pytest.mark.asyncio
async def test_status_falls_back_to_last_usage_when_context_estimate_missing(self):
@@ -324,7 +321,7 @@ class TestRestartCommand:
session = MagicMock()
session.get_history.return_value = [{"role": "user"}]
loop.sessions.get_or_create.return_value = session
loop._last_usage = {"prompt_tokens": 1200, "completion_tokens": 34}
loop._last_usage = LLMUsage.reported(input_tokens=1200, output_tokens=34)
loop.consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(0, "none")
)
tests/providers/test_azure_openai_provider.py
+2 -1
View File
@@ -380,7 +380,8 @@ async def test_chat_success():
assert isinstance(result, LLMResponse)
assert result.content == "Hello!"
assert result.finish_reason == "stop"
assert result.usage["prompt_tokens"] == 10
assert result.usage is not None
assert result.usage.input_tokens == 10
@pytest.mark.asyncio
tests/providers/test_bedrock_provider.py
+51 -3
View File
@@ -229,8 +229,11 @@ def test_parse_response_maps_text_tools_reasoning_usage_and_stop_reason() -> Non
assert result.content == "hello"
assert result.finish_reason == "tool_calls"
assert result.usage["prompt_tokens"] == 10
assert result.usage["cached_tokens"] == 2
assert result.usage is not None
assert result.usage.input_tokens == 12
assert result.usage.output_tokens == 5
assert result.usage.cache_read_tokens == 2
assert result.usage.cache_write_tokens is None
assert result.reasoning_content == "think"
assert result.thinking_blocks == [{"type": "thinking", "thinking": "think", "signature": "sig"}]
assert result.tool_calls[0].id == "t1"
@@ -276,7 +279,10 @@ async def test_chat_stream_aggregates_text_tool_use_and_usage() -> None:
assert deltas == ["he", "llo"]
assert result.content == "hello"
assert result.finish_reason == "tool_calls"
assert result.usage == {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7}
assert result.usage is not None
assert result.usage.input_tokens == 3
assert result.usage.output_tokens == 4
assert result.usage.total_tokens == 7
assert result.tool_calls[0].name == "search"
assert result.tool_calls[0].arguments == {"q": "x"}
@@ -285,6 +291,48 @@ async def _append_delta(deltas: list[str], text: str) -> None:
deltas.append(text)
@pytest.mark.parametrize(
("wire_usage", "expected_read", "expected_write", "expected_input"),
[
({"inputTokens": 5, "outputTokens": 1}, None, None, 5),
(
{
"inputTokens": 5,
"outputTokens": 1,
"cacheReadInputTokens": 0,
"cacheWriteInputTokens": 0,
},
0,
0,
5,
),
(
{
"inputTokens": 5,
"outputTokens": 1,
"cacheReadInputTokens": 7,
"cacheWriteInputTokens": 3,
},
7,
3,
15,
),
],
)
def test_bedrock_usage_preserves_cache_reporting_and_logical_input(
wire_usage: dict[str, int],
expected_read: int | None,
expected_write: int | None,
expected_input: int,
) -> None:
usage = BedrockProvider._usage(wire_usage)
assert usage is not None
assert usage.cache_read_tokens == expected_read
assert usage.cache_write_tokens == expected_write
assert usage.input_tokens == expected_input
@pytest.mark.asyncio
async def test_chat_error_maps_retry_metadata() -> None:
provider = BedrockProvider(region="us-east-1", client=FakeClient(error=FakeBedrockError()))
tests/providers/test_cached_tokens.py
+88 -22
View File
@@ -14,8 +14,9 @@ class FakeUsage:
class FakePromptDetails:
"""Mimics prompt_tokens_details sub-object."""
def __init__(self, cached_tokens=0):
def __init__(self, cached_tokens=0, cache_write_tokens=None):
self.cached_tokens = cached_tokens
self.cache_write_tokens = cache_write_tokens
class _FakeSpec:
@@ -62,8 +63,9 @@ def test_extract_usage_openai_cached_tokens_dict():
}
}
result = p._parse(response)
assert result.usage["cached_tokens"] == 1200
assert result.usage["prompt_tokens"] == 2000
assert result.usage is not None
assert result.usage.cache_read_tokens == 1200
assert result.usage.input_tokens == 2000
def test_extract_usage_deepseek_cached_tokens_dict():
@@ -80,11 +82,12 @@ def test_extract_usage_deepseek_cached_tokens_dict():
}
}
result = p._parse(response)
assert result.usage["cached_tokens"] == 1200
assert result.usage is not None
assert result.usage.cache_read_tokens == 1200
def test_extract_usage_no_cached_tokens_dict():
"""Response without any cache fields -> no cached_tokens key."""
"""Response without any cache fields preserves an unreported cache count."""
p = _provider()
response = {
"choices": [_DICT_CHOICE],
@@ -95,11 +98,13 @@ def test_extract_usage_no_cached_tokens_dict():
}
}
result = p._parse(response)
assert "cached_tokens" not in result.usage
assert result.usage is not None
assert result.usage.cache_read_tokens is None
assert result.usage.cache_write_tokens is None
def test_extract_usage_openai_cached_zero_dict():
"""cached_tokens=0 should NOT be included (same as existing fields)."""
"""cached_tokens=0 remains distinct from an unreported cache count."""
p = _provider()
response = {
"choices": [_DICT_CHOICE],
@@ -107,11 +112,42 @@ def test_extract_usage_openai_cached_zero_dict():
"prompt_tokens": 2000,
"completion_tokens": 300,
"total_tokens": 2300,
"prompt_tokens_details": {"cached_tokens": 0},
"prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
}
}
result = p._parse(response)
assert "cached_tokens" not in result.usage
assert result.usage is not None
assert result.usage.cache_read_tokens == 0
assert result.usage.cache_write_tokens == 0
def test_extract_usage_preserves_reported_total_and_cache_write_dict():
response = {
"choices": [_DICT_CHOICE],
"usage": {
"prompt_tokens": 15,
"completion_tokens": 18,
"total_tokens": 175,
"prompt_tokens_details": {
"cached_tokens": 0,
"cache_write_tokens": 7,
},
},
}
result = _provider()._parse(response)
assert result.usage is not None
assert result.usage.total_tokens == 175
assert result.usage.reported_tokens == 175
assert result.usage.cache_read_tokens == 0
assert result.usage.cache_write_tokens == 7
def test_extract_usage_missing_is_none():
result = _provider()._parse({"choices": [_DICT_CHOICE]})
assert result.usage is None
# --- object-based response (OpenAI SDK Pydantic model) ---
@@ -127,7 +163,29 @@ def test_extract_usage_openai_cached_tokens_obj():
)
response = FakeUsage(choices=[_FakeChoice()], usage=usage_obj)
result = p._parse(response)
assert result.usage["cached_tokens"] == 1200
assert result.usage is not None
assert result.usage.cache_read_tokens == 1200
def test_extract_usage_preserves_reported_total_and_cache_write_obj():
usage_obj = FakeUsage(
prompt_tokens=15,
completion_tokens=18,
total_tokens=175,
prompt_tokens_details=FakePromptDetails(
cached_tokens=0,
cache_write_tokens=7,
),
)
response = FakeUsage(choices=[_FakeChoice()], usage=usage_obj)
result = _provider()._parse(response)
assert result.usage is not None
assert result.usage.total_tokens == 175
assert result.usage.reported_tokens == 175
assert result.usage.cache_read_tokens == 0
assert result.usage.cache_write_tokens == 7
def test_extract_usage_deepseek_cached_tokens_obj():
@@ -141,7 +199,8 @@ def test_extract_usage_deepseek_cached_tokens_obj():
)
response = FakeUsage(choices=[_FakeChoice()], usage=usage_obj)
result = p._parse(response)
assert result.usage["cached_tokens"] == 1200
assert result.usage is not None
assert result.usage.cache_read_tokens == 1200
def test_extract_usage_stepfun_top_level_cached_tokens_dict():
@@ -157,7 +216,8 @@ def test_extract_usage_stepfun_top_level_cached_tokens_dict():
}
}
result = p._parse(response)
assert result.usage["cached_tokens"] == 512
assert result.usage is not None
assert result.usage.cache_read_tokens == 512
def test_extract_usage_stepfun_top_level_cached_tokens_obj():
@@ -171,7 +231,8 @@ def test_extract_usage_stepfun_top_level_cached_tokens_obj():
)
response = FakeUsage(choices=[_FakeChoice()], usage=usage_obj)
result = p._parse(response)
assert result.usage["cached_tokens"] == 512
assert result.usage is not None
assert result.usage.cache_read_tokens == 512
def test_extract_usage_priority_nested_over_top_level_dict():
@@ -188,11 +249,12 @@ def test_extract_usage_priority_nested_over_top_level_dict():
}
}
result = p._parse(response)
assert result.usage["cached_tokens"] == 100
assert result.usage is not None
assert result.usage.cache_read_tokens == 100
def test_anthropic_maps_cache_fields_to_cached_tokens():
"""Anthropic's cache_read_input_tokens should map to cached_tokens."""
def test_anthropic_adds_native_cache_fields_to_logical_input():
"""Anthropic excludes cache reads/writes from its native input_tokens."""
from nanobot.providers.anthropic_provider import AnthropicProvider
usage_obj = FakeUsage(
@@ -210,14 +272,15 @@ def test_anthropic_maps_cache_fields_to_cached_tokens():
usage=usage_obj,
)
result = AnthropicProvider._parse_response(response)
assert result.usage["cached_tokens"] == 1200
assert result.usage["prompt_tokens"] == 2300
assert result.usage["total_tokens"] == 2500
assert result.usage["cache_creation_input_tokens"] == 300
assert result.usage is not None
assert result.usage.cache_read_tokens == 1200
assert result.usage.cache_write_tokens == 300
assert result.usage.input_tokens == 2300
assert result.usage.total_tokens == 2500
def test_anthropic_no_cache_fields():
"""Anthropic response without cache fields should not have cached_tokens."""
"""Anthropic response without cache fields preserves unreported counts."""
from nanobot.providers.anthropic_provider import AnthropicProvider
usage_obj = FakeUsage(input_tokens=800, output_tokens=200)
@@ -230,4 +293,7 @@ def test_anthropic_no_cache_fields():
usage=usage_obj,
)
result = AnthropicProvider._parse_response(response)
assert "cached_tokens" not in result.usage
assert result.usage is not None
assert result.usage.input_tokens == 800
assert result.usage.cache_read_tokens is None
assert result.usage.cache_write_tokens is None
tests/providers/test_custom_provider.py
+2 -1
View File
@@ -46,7 +46,8 @@ def test_custom_provider_parse_accepts_dict_response() -> None:
assert result.finish_reason == "stop"
assert result.content == "hello from dict"
assert result.usage["total_tokens"] == 3
assert result.usage is not None
assert result.usage.total_tokens == 3
def test_custom_provider_parse_normalizes_text_tool_call() -> None:
tests/providers/test_openai_codex_provider.py
+8 -13
View File
@@ -732,11 +732,7 @@ async def test_codex_compacts_state_at_ninety_percent_before_next_request(
"content": [{"type": "output_text", "text": "old answer"}],
},
],
usage={
"prompt_tokens": 90,
"completion_tokens": 5,
"total_tokens": 95,
},
usage=provider_base.LLMUsage.reported(input_tokens=90, output_tokens=5),
)
bodies: list[dict[str, Any]] = []
@@ -772,11 +768,10 @@ async def test_codex_compacts_state_at_ninety_percent_before_next_request(
model="gpt-5.6-sol",
input_items=body["input"],
output_items=[compact_item],
usage={
"prompt_tokens": 95,
"completion_tokens": 2,
"total_tokens": 97,
},
usage=provider_base.LLMUsage.reported(
input_tokens=95,
output_tokens=2,
),
),
)
return provider_base.LLMResponse(content="done")
@@ -830,7 +825,7 @@ async def test_codex_disables_unsupported_native_compaction_and_continues(
model="gpt-5.6-sol",
input_items=[{"type": "message", "role": "user", "content": "old"}],
output_items=[{"type": "reasoning", "encrypted_content": "opaque"}],
usage={"prompt_tokens": 90, "completion_tokens": 5, "total_tokens": 95},
usage=provider_base.LLMUsage.reported(input_tokens=90, output_tokens=5),
)
bodies: list[dict[str, Any]] = []
@@ -914,7 +909,7 @@ async def test_codex_stream_surfaces_reasoning_summary(monkeypatch) -> None:
return provider_base.LLMResponse(
content="answer",
finish_reason="stop",
usage={"prompt_tokens": 10, "completion_tokens": 5},
usage=provider_base.LLMUsage.reported(input_tokens=10, output_tokens=5),
reasoning_content="summary",
)
@@ -934,7 +929,7 @@ async def test_codex_stream_surfaces_reasoning_summary(monkeypatch) -> None:
assert content_deltas == ["answer"]
assert thinking_deltas == ["summary"]
assert response.content == "answer"
assert response.usage == {"prompt_tokens": 10, "completion_tokens": 5}
assert response.usage == provider_base.LLMUsage.reported(input_tokens=10, output_tokens=5)
assert response.reasoning_content == "summary"
tests/providers/test_openai_responses.py
+101 -30
View File
@@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch
import pytest
from loguru import logger
from nanobot.providers.base import LLMUsage
from nanobot.providers.openai_responses.converters import (
convert_messages,
convert_tools,
@@ -484,7 +485,7 @@ class TestParseResponseOutput:
result = parse_response_output(resp)
assert result.content == "Hello!"
assert result.finish_reason == "stop"
assert result.usage == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
assert result.usage == LLMUsage.reported(input_tokens=10, output_tokens=5)
assert result.tool_calls == []
def test_refusal_response_surfaces_text_without_advancing_state(self):
@@ -652,7 +653,8 @@ class TestParseResponseOutput:
}
result = parse_response_output(mock)
assert result.content == "sdk"
assert result.usage["prompt_tokens"] == 1
assert result.usage is not None
assert result.usage.input_tokens == 1
def test_usage_maps_responses_api_keys(self):
"""Responses API uses input_tokens/output_tokens, not prompt_tokens/completion_tokens."""
@@ -662,9 +664,20 @@ class TestParseResponseOutput:
"usage": {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150},
}
result = parse_response_output(resp)
assert result.usage["prompt_tokens"] == 100
assert result.usage["completion_tokens"] == 50
assert result.usage["total_tokens"] == 150
assert result.usage == LLMUsage.reported(input_tokens=100, output_tokens=50)
def test_non_stream_preserves_provider_reported_total(self):
result = parse_response_output({
"output": [],
"status": "completed",
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 999},
})
assert result.usage == LLMUsage.reported(
input_tokens=10,
output_tokens=5,
total_tokens=999,
)
def test_preserves_every_output_item_as_opaque_state(self):
input_items = [{"role": "user", "content": "inspect the repo"}]
@@ -713,18 +726,18 @@ class TestResponsesConversationState:
{"type": "compaction", "encrypted_content": "compact"},
{"type": "message", "role": "assistant", "content": "new"},
],
usage={
"prompt_tokens": 90,
"completion_tokens": 10,
"total_tokens": 100,
},
usage=LLMUsage.reported(
input_tokens=90,
output_tokens=10,
total_tokens=175,
),
)
assert responses_state_items(state) == [
{"type": "compaction", "encrypted_content": "compact"},
{"type": "message", "role": "assistant", "content": "new"},
]
assert responses_state_context_tokens(state) == 100
assert responses_state_context_tokens(state) == 175
def test_existing_compaction_keeps_canonical_retained_prefix(self):
canonical_input = [
@@ -1090,7 +1103,7 @@ class TestConsumeSse:
assert content == "answer"
assert tool_calls == []
assert finish_reason == "stop"
assert usage == {}
assert usage is None
assert reasoning == "thinking briefly\nChecking result"
assert deltas == ["thinking ", "briefly", "\nChecking result"]
@@ -1224,7 +1237,7 @@ class TestConsumeSse:
assert content == "partial"
assert finish_reason == expected_finish_reason
assert usage == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
assert usage == LLMUsage.reported(input_tokens=10, output_tokens=5)
assert capture.completed is True
assert capture.response == terminal_response
assert capture.output_items == output
@@ -1296,7 +1309,10 @@ class TestConsumeSse:
"status": "completed",
"usage": {
"input_tokens": 10,
"input_tokens_details": {"cached_tokens": 8},
"input_tokens_details": {
"cached_tokens": 8,
"cache_write_tokens": 0,
},
"output_tokens": 5,
"total_tokens": 15,
},
@@ -1306,12 +1322,68 @@ class TestConsumeSse:
_, _, _, usage, _ = await consume_sse_with_reasoning(response)
assert usage == {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
"cached_tokens": 8,
assert usage == LLMUsage.reported(
input_tokens=10,
output_tokens=5,
cache_read_tokens=8,
cache_write_tokens=0,
)
@pytest.mark.asyncio
async def test_stream_and_non_stream_share_usage_normalization(self):
terminal = {
"status": "completed",
"output": [],
"usage": {
"input_tokens": 15,
"input_tokens_details": {
"cached_tokens": 0,
"cache_write_tokens": 7,
},
"output_tokens": 18,
"total_tokens": 175,
},
}
non_stream = parse_response_output(terminal).usage
sse = _SseResponse([
{"type": "response.completed", "response": terminal},
])
_, _, _, streamed, _ = await consume_sse_with_reasoning(sse)
sdk_response = SimpleNamespace(**terminal)
sdk_response.usage = SimpleNamespace(
input_tokens=15,
input_tokens_details=SimpleNamespace(
cached_tokens=0,
cache_write_tokens=7,
),
output_tokens=18,
total_tokens=175,
)
async def sdk_stream():
yield SimpleNamespace(type="response.completed", response=sdk_response)
_, _, _, sdk_streamed, _ = await consume_sdk_stream(sdk_stream())
expected = LLMUsage.reported(
input_tokens=15,
output_tokens=18,
total_tokens=175,
cache_read_tokens=0,
cache_write_tokens=7,
)
assert non_stream == streamed == sdk_streamed == expected
def test_missing_usage_is_not_explicit_zero_usage(self):
missing = parse_response_output({"status": "completed", "output": []})
explicit_zero = parse_response_output({
"status": "completed",
"output": [],
"usage": {"input_tokens": 0, "output_tokens": 0},
})
assert missing.usage is None
assert explicit_zero.usage == LLMUsage.reported(input_tokens=0, output_tokens=0)
@pytest.mark.asyncio
async def test_tool_call_done_arguments_callback(self):
@@ -1778,25 +1850,24 @@ class TestConsumeSdkStream:
@pytest.mark.asyncio
async def test_usage_extracted(self):
usage_obj = MagicMock(
usage_obj = SimpleNamespace(
input_tokens=10,
input_tokens_details=MagicMock(cached_tokens=8),
input_tokens_details=SimpleNamespace(cached_tokens=8),
output_tokens=5,
total_tokens=15,
)
resp_obj = MagicMock(status="completed", usage=usage_obj, output=[])
ev = MagicMock(type="response.completed", response=resp_obj)
resp_obj = SimpleNamespace(status="completed", usage=usage_obj, output=[])
ev = SimpleNamespace(type="response.completed", response=resp_obj)
async def stream():
yield ev
_, _, _, usage, _ = await consume_sdk_stream(stream())
assert usage == {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
"cached_tokens": 8,
}
assert usage == LLMUsage.reported(
input_tokens=10,
output_tokens=5,
cache_read_tokens=8,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
@@ -1851,7 +1922,7 @@ class TestConsumeSdkStream:
assert content == "partial"
assert finish_reason == expected_finish_reason
assert usage == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
assert usage == LLMUsage.reported(input_tokens=10, output_tokens=5)
assert capture.completed is True
assert capture.response == terminal_response
assert capture.output_items == output
tests/providers/test_provider_retry.py
+1 -1
View File
@@ -15,7 +15,7 @@ from nanobot.providers.base import (
class ScriptedProvider(LLMProvider):
def __init__(self, responses):
super().__init__()
super().__init__(provider_name="scripted")
self._responses = list(responses)
self.calls = 0
self.last_kwargs: dict = {}
tests/providers/test_providers_init.py
+1
View File
@@ -32,6 +32,7 @@ def test_importing_providers_package_is_lazy(monkeypatch) -> None:
assert providers.__all__ == [
"LLMProvider",
"LLMResponse",
"LLMUsage",
"AnthropicProvider",
"OpenAICompatProvider",
"OpenAICodexProvider",
tests/providers/test_usage_contract.py
+122
View File
@@ -0,0 +1,122 @@
import pytest
from nanobot.providers.base import LLMUsage
def test_reported_usage_derives_total_and_preserves_unreported_cache() -> None:
usage = LLMUsage.reported(input_tokens=12, output_tokens=3)
assert usage.total_tokens == 15
assert usage.cache_read_tokens is None
assert usage.cache_write_tokens is None
assert usage.source == "reported"
def test_reported_usage_preserves_explicit_total_across_contract_operations() -> None:
usage = LLMUsage.reported(input_tokens=15, output_tokens=18, total_tokens=175)
assert usage.total_tokens == 175
assert usage.reported_tokens == 175
assert usage.estimated_tokens == 0
assert LLMUsage.from_dict(usage.to_dict()) == usage
assert usage.with_timing(generation_ms=25, ttft_ms=5).total_tokens == 175
combined = usage + LLMUsage.estimated(input_tokens=2, output_tokens=1)
assert combined.total_tokens == 178
assert combined.reported_tokens == 175
assert combined.estimated_tokens == 3
def test_reported_usage_normalizes_missing_or_underreported_total() -> None:
missing = LLMUsage.reported(input_tokens=15, output_tokens=18)
underreported = LLMUsage.reported(
input_tokens=15,
output_tokens=18,
total_tokens=12,
)
assert missing.total_tokens == 33
assert underreported.total_tokens == 33
assert underreported.reported_tokens == 33
def test_reported_usage_preserves_explicit_zero_cache() -> None:
usage = LLMUsage.reported(
input_tokens=12,
output_tokens=3,
cache_read_tokens=0,
cache_write_tokens=0,
)
assert usage.cache_read_tokens == 0
assert usage.cache_write_tokens == 0
def test_usage_rejects_inconsistent_token_partitions_and_cache_totals() -> None:
with pytest.raises(ValueError, match="must equal"):
LLMUsage(input_tokens=10, output_tokens=2, total_tokens=12, reported_tokens=11)
with pytest.raises(ValueError, match="at least"):
LLMUsage(input_tokens=10, output_tokens=2, total_tokens=11, reported_tokens=11)
with pytest.raises(ValueError, match="cache token counts"):
LLMUsage.reported(input_tokens=10, output_tokens=2, cache_read_tokens=11)
def test_usage_serialization_is_strict_and_rejects_legacy_or_tampered_data() -> None:
usage = LLMUsage.estimated(input_tokens=10, output_tokens=2)
serialized = usage.to_dict()
assert LLMUsage.from_dict(serialized) == usage
assert LLMUsage.from_dict({"prompt_tokens": 10, "completion_tokens": 2}) is None
assert LLMUsage.from_dict({**serialized, "total_tokens": 99}) is None
assert LLMUsage.from_dict({**serialized, "source": "reported"}) is None
assert LLMUsage.from_dict({**serialized, "legacy_alias": 12}) is None
def test_usage_aggregation_keeps_reported_estimated_split_and_unknown_cache() -> None:
reported = LLMUsage.reported(
input_tokens=10,
output_tokens=2,
total_tokens=20,
cache_read_tokens=4,
)
estimated = LLMUsage.estimated(input_tokens=5, output_tokens=1)
combined = reported + estimated
assert combined.input_tokens == 15
assert combined.output_tokens == 3
assert combined.total_tokens == 26
assert combined.reported_tokens == 20
assert combined.estimated_tokens == 6
assert combined.source == "mixed"
assert combined.cache_read_tokens is None
assert combined.context_tokens == 5
assert combined.request_count == 2
def test_usage_projects_compact_turn_observability_shape() -> None:
usage = LLMUsage.reported(
input_tokens=12,
output_tokens=3,
total_tokens=20,
cache_read_tokens=4,
) + LLMUsage.estimated(input_tokens=18, output_tokens=2)
assert usage.to_turn_dict() == {
"prompt_tokens": 30,
"completion_tokens": 5,
"total_tokens": 40,
"context_tokens": 18,
"request_count": 2,
"estimated_tokens": 20,
}
def test_empty_request_counts_without_replacing_last_context() -> None:
usage = LLMUsage.reported(input_tokens=12, output_tokens=3) + LLMUsage.empty_request()
assert usage.total_tokens == 15
assert usage.context_tokens == 12
assert usage.request_count == 2
tests/providers/test_xai_grok_provider.py
+2 -1
View File
@@ -10,6 +10,7 @@ import httpx
import pytest
from nanobot.config.schema import Config
from nanobot.providers.base import LLMUsage
from nanobot.providers.factory import make_provider
from nanobot.providers.registry import find_by_name
from nanobot.providers.xai_grok_provider import (
@@ -454,7 +455,7 @@ async def test_raw_response_request_streams_text_usage_and_inline_citations(monk
assert result[0] == "Live result [[1]](https://x.com/example/status/1)"
assert result[2] == "stop"
assert result[3] == {"prompt_tokens": 8, "completion_tokens": 4, "total_tokens": 12}
assert result[3] == LLMUsage.reported(input_tokens=8, output_tokens=4)
assert deltas == ["Live result ", "[[1]](https://x.com/example/status/1)"]
assert captured["json"]["tools"] == [{"type": "x_search"}]
tests/test_api_attachment.py
+1 -1
View File
@@ -33,7 +33,7 @@ def _make_mock_agent(response_text: str = "mock response") -> MagicMock:
agent = MagicMock()
agent.process_direct = AsyncMock(return_value=response_text)
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
return agent
tests/test_api_stream.py
+8 -8
View File
@@ -77,7 +77,7 @@ def _make_streaming_agent(tokens: list[str]) -> MagicMock:
return " ".join(tokens)
agent.process_direct = fake_process_direct
agent._last_usage = {}
agent._last_usage = None
return agent
@@ -136,7 +136,7 @@ async def test_stream_false_returns_json(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = AsyncMock(return_value="normal reply")
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
@@ -159,7 +159,7 @@ async def test_stream_default_is_false(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = AsyncMock(return_value="default reply")
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
@@ -215,7 +215,7 @@ async def test_stream_passes_on_stream_callbacks(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = fake_process_direct
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
@@ -248,7 +248,7 @@ async def test_stream_segment_end_does_not_close_sse(aiohttp_client) -> None:
agent.process_direct = fake_process_direct
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
@@ -287,7 +287,7 @@ async def test_stream_uses_final_response_when_no_deltas(aiohttp_client) -> None
agent.process_direct = fake_process_direct
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
@@ -329,7 +329,7 @@ async def test_stream_with_session_id(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = fake_process_direct
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
@@ -358,7 +358,7 @@ async def test_streaming_backend_failure_does_not_emit_success_terminator(aiohtt
agent.process_direct = boom
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
tests/test_build_status.py
+10 -9
View File
@@ -1,5 +1,6 @@
"""Tests for build_status_content cache hit rate display."""
from nanobot.providers.base import LLMUsage
from nanobot.utils.helpers import build_status_content
@@ -8,7 +9,7 @@ def test_status_shows_cache_hit_rate():
version="0.1.0",
model="glm-4-plus",
start_time=1000000.0,
last_usage={"prompt_tokens": 2000, "completion_tokens": 300, "cached_tokens": 1200},
last_usage=LLMUsage.reported(input_tokens=2000, output_tokens=300, cache_read_tokens=1200),
context_window_tokens=128000,
session_msg_count=10,
context_tokens_estimate=5000,
@@ -19,12 +20,12 @@ def test_status_shows_cache_hit_rate():
def test_status_no_cache_info():
"""Without cached_tokens, display should not show cache percentage."""
"""Without a reported cache-read count, omit the cache percentage."""
content = build_status_content(
version="0.1.0",
model="glm-4-plus",
start_time=1000000.0,
last_usage={"prompt_tokens": 2000, "completion_tokens": 300},
last_usage=LLMUsage.reported(input_tokens=2000, output_tokens=300),
context_window_tokens=128000,
session_msg_count=10,
context_tokens_estimate=5000,
@@ -34,13 +35,13 @@ def test_status_no_cache_info():
assert "Tasks: 0 active" in content
def test_status_zero_cached_tokens():
"""cached_tokens=0 should not show cache percentage."""
def test_status_zero_cache_read_tokens():
"""An explicit zero cache-read count should not show cache percentage."""
content = build_status_content(
version="0.1.0",
model="glm-4-plus",
start_time=1000000.0,
last_usage={"prompt_tokens": 2000, "completion_tokens": 300, "cached_tokens": 0},
last_usage=LLMUsage.reported(input_tokens=2000, output_tokens=300, cache_read_tokens=0),
context_window_tokens=128000,
session_msg_count=10,
context_tokens_estimate=5000,
@@ -53,7 +54,7 @@ def test_status_100_percent_cached():
version="0.1.0",
model="glm-4-plus",
start_time=1000000.0,
last_usage={"prompt_tokens": 1000, "completion_tokens": 100, "cached_tokens": 1000},
last_usage=LLMUsage.reported(input_tokens=1000, output_tokens=100, cache_read_tokens=1000),
context_window_tokens=128000,
session_msg_count=5,
context_tokens_estimate=3000,
@@ -67,7 +68,7 @@ def test_status_context_pct_uses_budget_not_total():
version="0.1.0",
model="test",
start_time=1000000.0,
last_usage={"prompt_tokens": 2000, "completion_tokens": 300},
last_usage=LLMUsage.reported(input_tokens=2000, output_tokens=300),
context_window_tokens=128000,
session_msg_count=10,
context_tokens_estimate=120000,
@@ -83,7 +84,7 @@ def test_status_context_pct_capped_at_999():
version="0.1.0",
model="test",
start_time=1000000.0,
last_usage={"prompt_tokens": 2000, "completion_tokens": 300},
last_usage=LLMUsage.reported(input_tokens=2000, output_tokens=300),
context_window_tokens=10000,
session_msg_count=10,
context_tokens_estimate=100000,
tests/test_nanobot_facade.py
+10 -10
View File
@@ -30,6 +30,7 @@ from nanobot.nanobot import (
StreamEvent,
StreamEventType,
)
from nanobot.providers.base import LLMUsage
from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
RuntimeContextBlock,
@@ -601,7 +602,7 @@ async def test_run_no_iterations_leaves_defaults_empty(tmp_path):
result = await bot.run("hi")
assert result.tools_used == []
assert result.messages == []
assert result.usage == {}
assert result.usage is None
assert result.stop_reason is None
assert result.error is None
@@ -622,7 +623,7 @@ async def test_run_populates_observability_fields(tmp_path):
],
final_content="done",
tools_used=["read_file"],
usage={"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
usage=LLMUsage.reported(input_tokens=10, output_tokens=2),
stop_reason="completed",
error=None,
tool_events=[{"tool": "read_file", "status": "ok"}],
@@ -641,7 +642,7 @@ async def test_run_populates_observability_fields(tmp_path):
assert result.content == "done"
assert result.tools_used == ["read_file"]
assert result.usage == {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}
assert result.usage == LLMUsage.reported(input_tokens=10, output_tokens=2)
assert result.stop_reason == "completed"
assert result.error is None
assert result.metadata == {"latency_ms": 42}
@@ -658,7 +659,7 @@ async def test_run_ephemeral_still_captures_runner_observability(tmp_path):
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
content="done",
tool_calls=[],
usage={"total_tokens": 3},
usage=LLMUsage.reported(input_tokens=3, output_tokens=0),
))
bot = Nanobot(AgentLoop(
bus=MessageBus(),
@@ -670,8 +671,7 @@ async def test_run_ephemeral_still_captures_runner_observability(tmp_path):
result = await bot.run("hi", ephemeral=True)
assert result.content == "done"
assert result.usage["total_tokens"] == 3
assert result.usage["provider_tokens"] == 3
assert result.usage == LLMUsage.reported(input_tokens=3, output_tokens=0)
@pytest.mark.asyncio
@@ -1053,7 +1053,7 @@ async def test_run_streamed_wait_returns_full_result_without_consuming_events(tm
],
final_content="done",
tools_used=["read_file"],
usage={"total_tokens": 9},
usage=LLMUsage.reported(input_tokens=9, output_tokens=0),
stop_reason="completed",
)
for hook in hooks:
@@ -1073,7 +1073,7 @@ async def test_run_streamed_wait_returns_full_result_without_consuming_events(tm
assert result.content == "done"
assert result.tools_used == ["read_file"]
assert result.usage == {"total_tokens": 9}
assert result.usage == LLMUsage.reported(input_tokens=9, output_tokens=0)
assert result.stop_reason == "completed"
assert result.metadata == {"latency_ms": 5}
@@ -1397,13 +1397,13 @@ async def test_sdk_capture_prefers_run_level_snapshot():
await hook.after_run(AgentRunHookContext(
messages=final_messages,
tools_used=["read_file"],
usage={"total_tokens": 3},
usage=LLMUsage.reported(input_tokens=3, output_tokens=0),
stop_reason="completed",
))
assert hook.tools_used == ["read_file"]
assert hook.messages == final_messages
assert hook.usage == {"total_tokens": 3}
assert hook.usage == LLMUsage.reported(input_tokens=3, output_tokens=0)
assert hook.stop_reason == "completed"
tests/test_openai_api.py
+11 -10
View File
@@ -17,6 +17,7 @@ from nanobot.api.server import (
create_app,
handle_chat_completions,
)
from nanobot.providers.base import LLMUsage
try:
from aiohttp.test_utils import TestClient, TestServer
@@ -35,7 +36,7 @@ def _make_mock_agent(response_text: str = "mock response") -> MagicMock:
agent = MagicMock()
agent.process_direct = AsyncMock(return_value=response_text)
agent.aclose = AsyncMock()
agent._last_usage = {"prompt_tokens": 100, "completion_tokens": 50}
agent._last_usage = LLMUsage.reported(input_tokens=100, output_tokens=50)
return agent
@@ -87,19 +88,19 @@ def test_chat_completion_response() -> None:
def test_chat_completion_response_with_usage() -> None:
usage = {"prompt_tokens": 150, "completion_tokens": 42}
usage = LLMUsage.reported(input_tokens=150, output_tokens=42)
result = _chat_completion_response("hello world", "test-model", usage)
assert result["usage"]["prompt_tokens"] == 150
assert result["usage"]["completion_tokens"] == 42
assert result["usage"]["total_tokens"] == 192
def test_chat_completion_response_preserves_provider_total_usage() -> None:
usage = {"total_tokens": 77}
def test_chat_completion_response_preserves_explicit_total_usage() -> None:
usage = LLMUsage.reported(input_tokens=70, output_tokens=7, total_tokens=175)
result = _chat_completion_response("hello world", "test-model", usage)
assert result["usage"]["prompt_tokens"] == 0
assert result["usage"]["completion_tokens"] == 0
assert result["usage"]["total_tokens"] == 77
assert result["usage"]["prompt_tokens"] == 70
assert result["usage"]["completion_tokens"] == 7
assert result["usage"]["total_tokens"] == 175
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
@@ -328,7 +329,7 @@ async def test_followup_requests_share_same_session_key(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = fake_process
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
@@ -367,7 +368,7 @@ async def test_fixed_session_requests_are_serialized(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = slow_process
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
@@ -484,7 +485,7 @@ async def test_empty_response_falls_back_without_retry(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = always_empty
agent.aclose = AsyncMock()
agent._last_usage = {}
agent._last_usage = None
app = create_app(agent, model_name="m", api_key=API_KEY)
client = await aiohttp_client(app)
tests/utils/test_searchusage.py
+2 -1
View File
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from nanobot.providers.base import LLMUsage
from nanobot.utils.helpers import build_status_content
from nanobot.utils.searchusage import (
SearchUsageInfo,
@@ -273,7 +274,7 @@ class TestBuildStatusContentWithSearchUsage:
version="0.1.0",
model="claude-opus-4-5",
start_time=1_000_000.0,
last_usage={"prompt_tokens": 1000, "completion_tokens": 200},
last_usage=LLMUsage.reported(input_tokens=1000, output_tokens=200),
context_window_tokens=65536,
session_msg_count=5,
context_tokens_estimate=3000,
tests/webui/test_session_context.py
+4 -10
View File
@@ -1,3 +1,4 @@
from nanobot.providers.base import LLMUsage
from nanobot.session import Session
from nanobot.utils.helpers import estimate_message_tokens
from nanobot.webui.session_context import session_context_payload
@@ -58,19 +59,12 @@ def test_session_context_tolerates_untrusted_summary_metadata() -> None:
def test_session_context_sanitizes_usage_metadata() -> None:
usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175)
session = Session(
key="websocket:context",
metadata={
"_last_usage": {
"prompt_tokens": 120,
"completion_tokens": 8,
"negative": -1,
"boolean": True,
"text": "invalid",
}
},
metadata={"_last_usage": usage.to_dict()},
)
payload = session_context_payload(session)
assert payload["last_usage"] == {"prompt_tokens": 120, "completion_tokens": 8}
assert payload["last_usage"] == usage.to_dict()
tests/webui/test_settings_api.py
+3 -2
View File
@@ -9,6 +9,7 @@ import pytest
from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import Config, InlineFallbackConfig, ModelPresetConfig
from nanobot.providers.base import LLMUsage
from nanobot.providers.registry import find_by_name
from nanobot.session.manager import SessionManager
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
@@ -1467,7 +1468,7 @@ def test_settings_payload_includes_token_usage_summary(
from nanobot.webui.token_usage import record_token_usage
record_token_usage(
{"prompt_tokens": 10, "completion_tokens": 5},
LLMUsage.reported(input_tokens=10, output_tokens=5),
timezone_name=config.agents.defaults.timezone,
)
@@ -1495,7 +1496,7 @@ def test_settings_usage_payload_returns_lightweight_token_usage(
from nanobot.webui.token_usage import record_token_usage
record_token_usage(
{"prompt_tokens": 20, "completion_tokens": 2},
LLMUsage.reported(input_tokens=20, output_tokens=2),
timezone_name=config.agents.defaults.timezone,
)
tests/webui/test_token_usage.py
+140 -24
View File
@@ -1,17 +1,20 @@
from __future__ import annotations
import json
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
import pytest
from nanobot.agent.hook import AgentHookContext
from nanobot.providers.base import LLMUsage
from nanobot.webui.token_usage import (
TokenUsageHook,
read_token_usage_state,
record_response_token_usage,
record_token_usage,
token_usage_payload,
write_token_usage_state,
)
@@ -19,7 +22,7 @@ def _write_state(tmp_path, days: dict) -> None:
state_dir = tmp_path / "webui"
state_dir.mkdir(parents=True, exist_ok=True)
(state_dir / "token-usage.json").write_text(
json.dumps({"days": days}), encoding="utf-8"
json.dumps({"schema_version": 2, "days": days}), encoding="utf-8"
)
@@ -58,7 +61,7 @@ def test_record_scrubs_malformed_day_keys(tmp_path, monkeypatch) -> None:
})
record_token_usage(
{"prompt_tokens": 1, "completion_tokens": 1},
LLMUsage.reported(input_tokens=1, output_tokens=1),
timezone_name="UTC",
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
)
@@ -73,12 +76,16 @@ def test_record_token_usage_aggregates_by_local_day(tmp_path, monkeypatch) -> No
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage(
{"prompt_tokens": 100, "completion_tokens": 40, "cached_tokens": 20},
LLMUsage.reported(
input_tokens=100,
output_tokens=40,
cache_read_tokens=20,
),
timezone_name="Asia/Shanghai",
now=datetime(2026, 6, 2, 18, 0, tzinfo=timezone.utc),
)
record_token_usage(
{"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
LLMUsage.reported(input_tokens=10, output_tokens=5),
timezone_name="Asia/Shanghai",
now=datetime(2026, 6, 2, 19, 0, tzinfo=timezone.utc),
)
@@ -94,25 +101,31 @@ def test_record_token_usage_aggregates_by_local_day(tmp_path, monkeypatch) -> No
assert payload["days"] == [
{
"date": "2026-06-03",
"prompt_tokens": 110,
"completion_tokens": 45,
"cached_tokens": 20,
"input_tokens": 110,
"output_tokens": 45,
"cache_read_tokens": 20,
"cache_write_tokens": 0,
"cache_read_observed_input_tokens": 100,
"cache_write_observed_input_tokens": 0,
"total_tokens": 155,
"provider_tokens": 155,
"reported_tokens": 155,
"estimated_tokens": 0,
"requests": 2,
"provider_requests": 2,
"reported_requests": 2,
"estimated_requests": 0,
"sources": {
"user": {
"prompt_tokens": 110,
"completion_tokens": 45,
"cached_tokens": 20,
"input_tokens": 110,
"output_tokens": 45,
"cache_read_tokens": 20,
"cache_write_tokens": 0,
"cache_read_observed_input_tokens": 100,
"cache_write_observed_input_tokens": 0,
"total_tokens": 155,
"provider_tokens": 155,
"reported_tokens": 155,
"estimated_tokens": 0,
"requests": 2,
"provider_requests": 2,
"reported_requests": 2,
"estimated_requests": 0,
}
},
@@ -120,10 +133,113 @@ def test_record_token_usage_aggregates_by_local_day(tmp_path, monkeypatch) -> No
]
def test_cache_observation_denominators_distinguish_missing_from_zero(
tmp_path,
monkeypatch,
) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
now = datetime(2026, 6, 3, tzinfo=timezone.utc)
record_token_usage(
LLMUsage.reported(input_tokens=100, output_tokens=10),
source="user",
now=now,
)
record_token_usage(
LLMUsage.reported(
input_tokens=40,
output_tokens=5,
cache_read_tokens=0,
cache_write_tokens=0,
),
source="dream",
now=now,
)
row = token_usage_payload(now=now)["days"][0]
assert row["cache_read_tokens"] == 0
assert row["cache_write_tokens"] == 0
assert row["cache_read_observed_input_tokens"] == 40
assert row["cache_write_observed_input_tokens"] == 40
assert row["sources"]["user"]["cache_read_observed_input_tokens"] == 0
assert row["sources"]["user"]["cache_write_observed_input_tokens"] == 0
assert row["sources"]["dream"]["cache_read_observed_input_tokens"] == 40
assert row["sources"]["dream"]["cache_write_observed_input_tokens"] == 40
def _retention_state(sources: tuple[str, ...], *, day_count: int = 400) -> dict:
start = datetime(2025, 1, 1, tzinfo=timezone.utc)
source_usage = {
"input_tokens": 100,
"output_tokens": 10,
"total_tokens": 110,
"reported_tokens": 110,
"requests": 1,
"reported_requests": 1,
}
days = {}
for offset in range(day_count):
day = (start + timedelta(days=offset)).date().isoformat()
days[day] = {
"input_tokens": 100 * len(sources),
"output_tokens": 10 * len(sources),
"total_tokens": 110 * len(sources),
"reported_tokens": 110 * len(sources),
"requests": len(sources),
"reported_requests": len(sources),
"sources": {source: dict(source_usage) for source in sources},
}
return {"schema_version": 2, "days": days}
def test_write_compact_state_keeps_400_days_with_two_sources(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
written = write_token_usage_state(_retention_state(("user", "api")))
persisted = (tmp_path / "webui" / "token-usage.json").read_bytes()
assert len(written["days"]) == 400
assert len(persisted) <= 512 * 1024
assert persisted.endswith(b"\n")
assert json.loads(persisted) == written
def test_write_prunes_only_oldest_days_to_fit_byte_budget(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
sources = ("user", "api", "cron", "dream", "system")
raw = _retention_state(sources)
all_dates = list(raw["days"])
written = write_token_usage_state(raw)
retained_dates = list(written["days"])
persisted = (tmp_path / "webui" / "token-usage.json").read_bytes()
assert 1 <= len(retained_dates) < len(all_dates)
assert retained_dates == all_dates[-len(retained_dates) :]
assert retained_dates[-1] == all_dates[-1]
assert all(set(row["sources"]) == set(sources) for row in written["days"].values())
assert len(persisted) <= 512 * 1024
assert read_token_usage_state() == written
def test_write_raises_when_latest_day_alone_exceeds_byte_budget(
tmp_path,
monkeypatch,
) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
monkeypatch.setattr("nanobot.webui.token_usage._MAX_STATE_FILE_BYTES", 256)
with pytest.raises(ValueError, match="latest token usage day exceeds"):
write_token_usage_state(_retention_state(("user", "api"), day_count=1))
assert not (tmp_path / "webui" / "token-usage.json").exists()
def test_record_token_usage_skips_empty_usage(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage({"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
record_token_usage(LLMUsage.reported(input_tokens=0, output_tokens=0))
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
assert payload["days"] == []
@@ -134,14 +250,14 @@ def test_record_token_usage_keeps_estimated_split(tmp_path, monkeypatch) -> None
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage(
{"prompt_tokens": 100, "completion_tokens": 25, "estimated_tokens": 125},
LLMUsage.estimated(input_tokens=100, output_tokens=25),
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
)
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
assert payload["days"][0]["total_tokens"] == 125
assert payload["days"][0]["provider_tokens"] == 0
assert payload["days"][0]["reported_tokens"] == 0
assert payload["days"][0]["estimated_tokens"] == 125
assert payload["days"][0]["estimated_requests"] == 1
@@ -150,12 +266,12 @@ def test_record_token_usage_keeps_source_breakdown(tmp_path, monkeypatch) -> Non
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage(
{"prompt_tokens": 100, "completion_tokens": 25},
LLMUsage.reported(input_tokens=100, output_tokens=25, total_tokens=175),
source="user",
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
)
record_token_usage(
{"prompt_tokens": 20, "completion_tokens": 5},
LLMUsage.reported(input_tokens=20, output_tokens=5),
source="dream",
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
)
@@ -163,8 +279,8 @@ def test_record_token_usage_keeps_source_breakdown(tmp_path, monkeypatch) -> Non
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
row = payload["days"][0]
assert row["total_tokens"] == 150
assert row["sources"]["user"]["total_tokens"] == 125
assert row["total_tokens"] == 200
assert row["sources"]["user"]["total_tokens"] == 175
assert row["sources"]["user"]["requests"] == 1
assert row["sources"]["dream"]["total_tokens"] == 25
assert row["sources"]["dream"]["requests"] == 1
@@ -175,7 +291,7 @@ def test_record_response_token_usage_uses_response_usage(tmp_path, monkeypatch)
monkeypatch.setattr("nanobot.webui.token_usage._local_day", lambda *_, **__: "2026-06-03")
record_response_token_usage(
SimpleNamespace(usage={"prompt_tokens": 20, "completion_tokens": 5}),
SimpleNamespace(usage=LLMUsage.reported(input_tokens=20, output_tokens=5)),
source="dream",
)
@@ -194,7 +310,7 @@ async def test_token_usage_hook_classifies_source_from_session_key(tmp_path, mon
iteration=0,
messages=[],
session_key="cron:drink-water",
usage={"prompt_tokens": 10, "completion_tokens": 5},
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
)
)