mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-04 16:38:49 +00:00
07c2677eed
normalize_token_usage_state only length-checked persisted day keys, so a hand-edited or foreign 10-char key (e.g. "not-a-dat3" or "2026-13-01") in token-usage.json survived reads and atomic rewrites. token_usage_payload then parsed every day key with an unguarded datetime.fromisoformat, so one such key failed every /api/settings and /api/settings/usage request until the file was repaired by hand. Validate day keys in normalize_token_usage_state, the shared boundary that every read, record, and rewrite already funnels through. Malformed keys are dropped like other malformed rows and scrubbed from the file on the next write; valid state is unchanged.
204 lines
7.3 KiB
Python
204 lines
7.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from nanobot.agent.hook import AgentHookContext
|
|
from nanobot.webui.token_usage import (
|
|
TokenUsageHook,
|
|
record_response_token_usage,
|
|
record_token_usage,
|
|
token_usage_payload,
|
|
)
|
|
|
|
|
|
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"
|
|
)
|
|
|
|
|
|
def test_payload_tolerates_malformed_persisted_day_keys(tmp_path, monkeypatch) -> None:
|
|
"""Day keys that are not real dates must not break settings payloads.
|
|
|
|
normalize_token_usage_state only length-checks day keys, so a hand-edited
|
|
10-char key survives reads and atomic rewrites; token_usage_payload then
|
|
parsed it with an unguarded fromisoformat, failing every /api/settings and
|
|
/api/settings/usage request until the file was fixed by hand.
|
|
"""
|
|
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
|
_write_state(tmp_path, {
|
|
"not-a-dat3": {"total_tokens": 7, "requests": 1},
|
|
"2026-13-01": {"total_tokens": 9, "requests": 1},
|
|
"2026-06-02": {"total_tokens": 5, "requests": 1},
|
|
})
|
|
|
|
payload = token_usage_payload(
|
|
timezone_name="UTC",
|
|
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
|
|
)
|
|
|
|
assert payload["total_tokens"] == 5
|
|
assert payload["total_tokens_30d"] == 5
|
|
assert payload["requests_30d"] == 1
|
|
assert payload["active_days_30d"] == 1
|
|
|
|
|
|
def test_record_scrubs_malformed_day_keys(tmp_path, monkeypatch) -> None:
|
|
"""Rewrites drop malformed day keys instead of persisting them forever."""
|
|
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
|
_write_state(tmp_path, {
|
|
"not-a-dat3": {"total_tokens": 7, "requests": 1},
|
|
"2026-06-02": {"total_tokens": 5, "requests": 1},
|
|
})
|
|
|
|
record_token_usage(
|
|
{"prompt_tokens": 1, "completion_tokens": 1},
|
|
timezone_name="UTC",
|
|
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
|
|
)
|
|
|
|
raw = json.loads((tmp_path / "webui" / "token-usage.json").read_text(encoding="utf-8"))
|
|
assert "not-a-dat3" not in raw["days"]
|
|
assert "2026-06-02" in raw["days"]
|
|
assert "2026-06-03" in raw["days"]
|
|
|
|
|
|
def test_record_token_usage_aggregates_by_local_day(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": 40, "cached_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},
|
|
timezone_name="Asia/Shanghai",
|
|
now=datetime(2026, 6, 2, 19, 0, tzinfo=timezone.utc),
|
|
)
|
|
|
|
payload = token_usage_payload(
|
|
timezone_name="Asia/Shanghai",
|
|
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
|
|
)
|
|
|
|
assert payload["total_tokens_30d"] == 155
|
|
assert payload["active_days_30d"] == 1
|
|
assert payload["requests_30d"] == 2
|
|
assert payload["days"] == [
|
|
{
|
|
"date": "2026-06-03",
|
|
"prompt_tokens": 110,
|
|
"completion_tokens": 45,
|
|
"cached_tokens": 20,
|
|
"total_tokens": 155,
|
|
"provider_tokens": 155,
|
|
"estimated_tokens": 0,
|
|
"requests": 2,
|
|
"provider_requests": 2,
|
|
"estimated_requests": 0,
|
|
"sources": {
|
|
"user": {
|
|
"prompt_tokens": 110,
|
|
"completion_tokens": 45,
|
|
"cached_tokens": 20,
|
|
"total_tokens": 155,
|
|
"provider_tokens": 155,
|
|
"estimated_tokens": 0,
|
|
"requests": 2,
|
|
"provider_requests": 2,
|
|
"estimated_requests": 0,
|
|
}
|
|
},
|
|
}
|
|
]
|
|
|
|
|
|
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})
|
|
|
|
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
|
|
assert payload["days"] == []
|
|
assert payload["total_tokens_30d"] == 0
|
|
|
|
|
|
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},
|
|
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]["estimated_tokens"] == 125
|
|
assert payload["days"][0]["estimated_requests"] == 1
|
|
|
|
|
|
def test_record_token_usage_keeps_source_breakdown(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},
|
|
source="user",
|
|
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
|
|
)
|
|
record_token_usage(
|
|
{"prompt_tokens": 20, "completion_tokens": 5},
|
|
source="dream",
|
|
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
|
|
)
|
|
|
|
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["sources"]["user"]["requests"] == 1
|
|
assert row["sources"]["dream"]["total_tokens"] == 25
|
|
assert row["sources"]["dream"]["requests"] == 1
|
|
|
|
|
|
def test_record_response_token_usage_uses_response_usage(tmp_path, monkeypatch) -> None:
|
|
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
|
monkeypatch.setattr("nanobot.webui.token_usage._local_day", lambda *_, **__: "2026-06-03")
|
|
|
|
record_response_token_usage(
|
|
SimpleNamespace(usage={"prompt_tokens": 20, "completion_tokens": 5}),
|
|
source="dream",
|
|
)
|
|
|
|
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
|
|
assert payload["days"][0]["sources"]["dream"]["total_tokens"] == 25
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_token_usage_hook_classifies_source_from_session_key(tmp_path, monkeypatch) -> None:
|
|
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
|
|
monkeypatch.setattr("nanobot.webui.token_usage._local_day", lambda *_, **__: "2026-06-03")
|
|
|
|
hook = TokenUsageHook()
|
|
await hook.after_iteration(
|
|
AgentHookContext(
|
|
iteration=0,
|
|
messages=[],
|
|
session_key="cron:drink-water",
|
|
usage={"prompt_tokens": 10, "completion_tokens": 5},
|
|
)
|
|
)
|
|
|
|
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
|
|
|
|
assert payload["days"][0]["sources"]["cron"]["total_tokens"] == 15
|