From 2ac802b2d56ae70171baf962ae48c48e6df57f15 Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Sat, 22 Aug 2026 02:00:28 +0800 Subject: [PATCH] feat(usage): add unified provider usage backend --- nanobot/agent/loop.py | 6 + nanobot/agent/memory.py | 20 +- nanobot/agent/runner.py | 33 +- nanobot/agent/subagent.py | 10 +- nanobot/cli/gateway_runtime.py | 42 +- nanobot/command/builtin.py | 7 - nanobot/llm_usage/__init__.py | 86 ++++ nanobot/llm_usage/context.py | 70 +++ nanobot/llm_usage/models.py | 38 ++ nanobot/llm_usage/store.py | 560 +++++++++++++++++++++ nanobot/providers/base.py | 150 +++++- nanobot/providers/fallback_provider.py | 7 + nanobot/session/webui_turns.py | 38 +- nanobot/webui/settings_routes.py | 4 +- nanobot/webui/settings_system.py | 6 +- nanobot/webui/token_usage.py | 392 --------------- tests/agent/test_subagent.py | 29 ++ tests/llm_usage/test_llm_usage_context.py | 19 + tests/llm_usage/test_llm_usage_store.py | 275 ++++++++++ tests/providers/test_llm_usage_observer.py | 180 +++++++ tests/webui/test_settings_api.py | 39 +- tests/webui/test_settings_routes.py | 40 ++ tests/webui/test_token_usage.py | 319 ------------ webui/src/lib/types.ts | 43 ++ 24 files changed, 1605 insertions(+), 808 deletions(-) create mode 100644 nanobot/llm_usage/__init__.py create mode 100644 nanobot/llm_usage/context.py create mode 100644 nanobot/llm_usage/models.py create mode 100644 nanobot/llm_usage/store.py delete mode 100644 nanobot/webui/token_usage.py create mode 100644 tests/llm_usage/test_llm_usage_context.py create mode 100644 tests/llm_usage/test_llm_usage_store.py create mode 100644 tests/providers/test_llm_usage_observer.py delete mode 100644 tests/webui/test_token_usage.py diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 2f34664f2..458345ed4 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -49,6 +49,7 @@ from nanobot.bus.queue import MessageBus from nanobot.bus.runtime_events import RuntimeEventBus from nanobot.command import CommandContext, CommandRouter, register_builtin_commands from nanobot.config.schema import AgentDefaults, ModelPresetConfig +from nanobot.llm_usage.context import source_from_request from nanobot.providers.base import LLMProvider, LLMUsage, ProviderConversationState from nanobot.providers.factory import ProviderSnapshot from nanobot.runtime_context import ( @@ -1202,6 +1203,11 @@ class AgentLoop: message_metadata=metadata, ), provider_state=provider_state, + llm_usage_source=source_from_request( + active_session_key, + channel=channel, + metadata=metadata, + ), )) finally: turn_scope_stack.close() diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index ba964a208..98e89c03d 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -20,6 +20,7 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, cast from loguru import logger +from nanobot.llm_usage.context import llm_usage_source from nanobot.runtime_context import public_history_messages from nanobot.session.manager import ( MIN_COMPACTED_REPLAY_MESSAGES, @@ -915,15 +916,16 @@ class Consolidator: if not messages: return None try: - response = await runtime.provider.chat_with_retry( - model=runtime.model, - messages=request_messages, - tools=request_tools, - tool_choice="none", - temperature=runtime.generation.temperature, - max_tokens=runtime.generation.max_tokens, - reasoning_effort=runtime.generation.reasoning_effort, - ) + with llm_usage_source("dream"): + response = await runtime.provider.chat_with_retry( + model=runtime.model, + messages=request_messages, + tools=request_tools, + tool_choice="none", + temperature=runtime.generation.temperature, + max_tokens=runtime.generation.max_tokens, + reasoning_effort=runtime.generation.reasoning_effort, + ) except Exception: logger.warning("Consolidation provider call failed, raw-dumping to history") self.store.raw_archive(messages, session_key=session_key) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 14131d2d5..feffd7cc7 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -20,6 +20,12 @@ from nanobot.agent.context_governance import ( ) from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result +from nanobot.llm_usage.context import ( + LLMUsageSource, + bind_llm_usage_source, + reset_llm_usage_source, + source_from_session_key, +) from nanobot.providers.base import ( LLMProvider, LLMResponse, @@ -118,6 +124,7 @@ class AgentRunSpec: goal_continue_message: GoalContinueMessage | None = None finalize_on_max_iterations: bool = True provider_state: ProviderConversationState | None = None + llm_usage_source: LLMUsageSource | None = None @dataclass(slots=True) @@ -392,6 +399,9 @@ class AgentRunner: hook = spec.hook or AgentHook() messages = list(spec.initial_messages) context = AgentRunHookContext(messages=deepcopy(messages)) + llm_usage_source_token = bind_llm_usage_source( + spec.llm_usage_source or source_from_session_key(spec.session_key) + ) try: await hook.before_run(context) @@ -424,17 +434,20 @@ class AgentRunner: await hook.after_run(context) return result finally: - context.messages = deepcopy(messages) - if context.exception is None: - await hook.on_finally(context) - else: - try: + try: + context.messages = deepcopy(messages) + if context.exception is None: await hook.on_finally(context) - except Exception: - logger.exception( - "AgentHook.on_finally error after {}", - context.stop_reason or "run exception", - ) + else: + try: + await hook.on_finally(context) + except Exception: + logger.exception( + "AgentHook.on_finally error after {}", + context.stop_reason or "run exception", + ) + finally: + reset_llm_usage_source(llm_usage_source_token) async def _run_core( self, diff --git a/nanobot/agent/subagent.py b/nanobot/agent/subagent.py index e88e43d38..f5716a18f 100644 --- a/nanobot/agent/subagent.py +++ b/nanobot/agent/subagent.py @@ -8,7 +8,7 @@ import warnings from collections.abc import Mapping from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Callable, TypedDict +from typing import Any, Callable, NotRequired, TypedDict from loguru import logger @@ -28,6 +28,7 @@ from nanobot.agent.tools.registry import ToolRegistry from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus from nanobot.config.schema import AgentDefaults, ToolsConfig +from nanobot.llm_usage.context import LLMUsageSource, current_llm_usage_source from nanobot.providers.base import LLMProvider, LLMUsage from nanobot.security.workspace_access import ( WorkspaceScope, @@ -43,6 +44,7 @@ class _SubagentOrigin(TypedDict): channel: str chat_id: str session_key: str | None + llm_usage_source: NotRequired[LLMUsageSource] @dataclass(slots=True) @@ -252,6 +254,7 @@ class SubagentManager: "channel": origin_channel, "chat_id": origin_chat_id, "session_key": session_key, + "llm_usage_source": current_llm_usage_source(), } status = SubagentStatus( @@ -315,6 +318,7 @@ class SubagentManager: "channel": origin_channel, "chat_id": origin_chat_id, "session_key": session_key, + "llm_usage_source": current_llm_usage_source(), } status = SubagentStatus( task_id=task_id, @@ -417,6 +421,10 @@ class SubagentManager: session_key=sess_key, workspace=root, llm_timeout_s=llm_timeout, + llm_usage_source=origin.get( + "llm_usage_source", + current_llm_usage_source(), + ), )) finally: if token is not None: diff --git a/nanobot/cli/gateway_runtime.py b/nanobot/cli/gateway_runtime.py index d7082cd27..8d5d0f3a0 100644 --- a/nanobot/cli/gateway_runtime.py +++ b/nanobot/cli/gateway_runtime.py @@ -313,6 +313,8 @@ def _run_gateway( from nanobot.cron.service import CronJobSkippedError, CronService from nanobot.cron.session_turns import is_bound_cron_job from nanobot.cron.types import CronJob + from nanobot.llm_usage import record_llm_call + from nanobot.llm_usage.context import llm_usage_source from nanobot.providers.factory import ( ProviderSnapshot, build_provider_snapshot, @@ -330,7 +332,6 @@ def _run_gateway( ) from nanobot.triggers.local_runner import run_local_trigger_queue from nanobot.triggers.local_store import LocalTriggerStore - from nanobot.webui.token_usage import TokenUsageHook port = port if port is not None else config.gateway.port webui_url = _webui_browser_url(config) @@ -361,7 +362,8 @@ def _run_gateway( runtime_events = RuntimeEventBus() fallback_model_observer = build_webui_fallback_model_observer(bus) - def _observe_fallback_models(snapshot: ProviderSnapshot) -> ProviderSnapshot: + def _observe_provider(snapshot: ProviderSnapshot) -> ProviderSnapshot: + snapshot.provider.set_llm_call_observer(record_llm_call) if isinstance(snapshot.provider, FallbackProvider): snapshot.provider.set_fallback_model_observer(fallback_model_observer) return snapshot @@ -371,20 +373,19 @@ def _run_gateway( **kwargs: Any, ) -> ProviderSnapshot: try: - return _observe_fallback_models(load_provider_snapshot(*args, **kwargs)) + return _observe_provider(load_provider_snapshot(*args, **kwargs)) except ValueError as exc: if unconfigured_provider_error is None: raise - return build_unconfigured_provider_snapshot(config, str(exc)) + return _observe_provider(build_unconfigured_provider_snapshot(config, str(exc))) if unconfigured_provider_error is not None: - provider_snapshot = build_unconfigured_provider_snapshot( - config, - unconfigured_provider_error, + provider_snapshot = _observe_provider( + build_unconfigured_provider_snapshot(config, unconfigured_provider_error) ) else: try: - provider_snapshot = _observe_fallback_models(build_provider_snapshot(config)) + provider_snapshot = _observe_provider(build_provider_snapshot(config)) except ValueError as exc: console.print(f"[red]Error: {exc}[/red]") raise typer.Exit(1) from exc @@ -443,7 +444,6 @@ def _run_gateway( runtime_events=runtime_events, turn_delivery_factory=turn_delivery_factory, provider_signature=provider_snapshot.signature, - hooks=[TokenUsageHook(timezone_name=config.agents.defaults.timezone)], local_trigger_store=trigger_store, hook_factories=[create_file_edit_activity_hook], tool_registry=tools, @@ -564,13 +564,6 @@ def _run_gateway( except Exception: logger.exception("Dream cron job failed") finally: - from nanobot.webui.token_usage import record_response_token_usage - - record_response_token_usage( - resp, - source="dream", - timezone_name=config.agents.defaults.timezone, - ) sha = _commit_dream_changes(store) if sha: logger.info("Dream commit: {}", sha) @@ -630,14 +623,15 @@ def _run_gateway( evaluator_prompt = resolve_evaluator_prompt(config.workspace_path) # Fail closed: stay silent on evaluator failure instead of notifying. - should_notify = await evaluate_response( - response=response, - task_context=prompt, - provider=agent.provider, - model=agent.model, - evaluator_prompt=evaluator_prompt, - default_notify=False, - ) + with llm_usage_source("cron"): + should_notify = await evaluate_response( + response=response, + task_context=prompt, + provider=agent.provider, + model=agent.model, + evaluator_prompt=evaluator_prompt, + default_notify=False, + ) if should_notify: logger.info("Heartbeat: completed, delivering response") diff --git a/nanobot/command/builtin.py b/nanobot/command/builtin.py index 0d3e1c2cb..ac8f8421a 100644 --- a/nanobot/command/builtin.py +++ b/nanobot/command/builtin.py @@ -479,13 +479,6 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage: elapsed = time.monotonic() - t0 content = f"Dream failed after {elapsed:.1f}s: {e}" finally: - from nanobot.webui.token_usage import record_response_token_usage - - record_response_token_usage( - resp, - source="dream", - timezone_name=getattr(loop.context, "timezone", None), - ) if store.git.is_initialized(): commit_msg = build_dream_commit_message("dream: manual run", diff_body) sha = store.git.auto_commit(commit_msg) diff --git a/nanobot/llm_usage/__init__.py b/nanobot/llm_usage/__init__.py new file mode 100644 index 000000000..95d51a9c5 --- /dev/null +++ b/nanobot/llm_usage/__init__.py @@ -0,0 +1,86 @@ +"""Unified, content-free LLM usage backend.""" + +from __future__ import annotations + +import threading +from pathlib import Path +from typing import Any + +from loguru import logger + +from nanobot.config.paths import get_data_dir +from nanobot.llm_usage.models import LLMCallRecord +from nanobot.llm_usage.store import LLMUsageStore + +_STORES_LOCK = threading.Lock() +_STORES: dict[Path, LLMUsageStore] = {} + + +def empty_usage_payload() -> dict[str, Any]: + return { + "days": [], + "total_tokens": 0, + "total_tokens_30d": 0, + "total_tokens_365d": 0, + "reported_tokens_30d": 0, + "estimated_tokens_30d": 0, + "cache_read_tokens_30d": 0, + "cache_read_observed_input_tokens_30d": 0, + "cache_read_rate_30d": None, + "peak_day_tokens": 0, + "current_streak_days": 0, + "longest_streak_days": 0, + "active_days_30d": 0, + "requests_30d": 0, + "failed_requests_30d": 0, + "providers_30d": [], + "updated_at": None, + } + + +def llm_usage_store_path() -> Path: + return get_data_dir() / "llm_usage.sqlite3" + + +def get_llm_usage_store(path: Path | None = None) -> LLMUsageStore: + resolved = (path or llm_usage_store_path()).resolve(strict=False) + with _STORES_LOCK: + store = _STORES.get(resolved) + if store is None: + store = LLMUsageStore(resolved) + _STORES[resolved] = store + return store + + +def record_llm_call(call: LLMCallRecord) -> None: + """Default fail-open callback attached to gateway provider snapshots.""" + try: + get_llm_usage_store().record(call) + except Exception: + logger.exception("failed to record LLM usage") + + +def llm_usage_payload( + *, + days: int = 371, + timezone_name: str | None = None, +) -> dict[str, Any]: + try: + return get_llm_usage_store().usage_payload( + days=days, + timezone_name=timezone_name, + ) + except Exception: + logger.exception("failed to query LLM usage") + return empty_usage_payload() + + +__all__ = [ + "LLMCallRecord", + "LLMUsageStore", + "empty_usage_payload", + "get_llm_usage_store", + "record_llm_call", + "llm_usage_store_path", + "llm_usage_payload", +] diff --git a/nanobot/llm_usage/context.py b/nanobot/llm_usage/context.py new file mode 100644 index 000000000..c6055e2b6 --- /dev/null +++ b/nanobot/llm_usage/context.py @@ -0,0 +1,70 @@ +"""Request-local metadata for LLM usage records.""" + +from __future__ import annotations + +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar, Token +from typing import Literal + +LLMUsageSource = Literal["user", "api", "cron", "dream", "system"] + +_CURRENT_SOURCE: ContextVar[LLMUsageSource] = ContextVar( + "nanobot_llm_usage_source", + default="system", +) + + +def source_from_session_key(session_key: str | None) -> LLMUsageSource: + """Classify a private session key without persisting that key.""" + key = session_key or "" + if key.startswith("dream:"): + return "dream" + if key == "heartbeat" or key.startswith("cron:"): + return "cron" + if key.startswith("api:"): + return "api" + if key.startswith("system:"): + return "system" + return "user" + + +def source_from_request( + session_key: str | None, + *, + channel: str | None, + metadata: Mapping[str, object] | None, +) -> LLMUsageSource: + """Classify a turn from trusted ingress metadata without retaining identifiers.""" + values = metadata or {} + if isinstance(values.get("_cron_trigger"), Mapping): + return "cron" + if isinstance(values.get("_local_trigger"), Mapping): + return "cron" + if channel == "api": + return "api" + if channel == "system": + return "system" + return source_from_session_key(session_key) + + +def current_llm_usage_source() -> LLMUsageSource: + return _CURRENT_SOURCE.get() + + +def bind_llm_usage_source(source: LLMUsageSource) -> Token[LLMUsageSource]: + return _CURRENT_SOURCE.set(source) + + +def reset_llm_usage_source(token: Token[LLMUsageSource]) -> None: + _CURRENT_SOURCE.reset(token) + + +@contextmanager +def llm_usage_source(source: LLMUsageSource) -> Generator[None]: + """Bind a coarse usage source for nested provider calls.""" + token = bind_llm_usage_source(source) + try: + yield + finally: + reset_llm_usage_source(token) diff --git a/nanobot/llm_usage/models.py b/nanobot/llm_usage/models.py new file mode 100644 index 000000000..7ceb6f5c0 --- /dev/null +++ b/nanobot/llm_usage/models.py @@ -0,0 +1,38 @@ +"""Content-free records emitted for physical LLM provider calls.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from nanobot.llm_usage.context import LLMUsageSource +from nanobot.providers.base import LLMUsage + + +@dataclass(frozen=True, slots=True) +class LLMCallRecord: + """The small, chart-oriented result of one provider call attempt. + + Request messages, response text, reasoning, and tool payloads deliberately do + not belong to this contract. Sessions already own that content. + """ + + started_at_ms: int + duration_ms: int + provider: str + model: str + source: LLMUsageSource + stream: bool + finish_reason: str + usage: LLMUsage | None = None + error_status_code: int | None = None + error_kind: str | None = None + + def __post_init__(self) -> None: + if self.started_at_ms < 0 or self.duration_ms < 0: + raise ValueError("LLM usage timestamps must be non-negative") + if not self.provider.strip() or not self.model.strip(): + raise ValueError("LLM usage provider and model must be non-empty") + if self.source not in {"user", "api", "cron", "dream", "system"}: + raise ValueError("invalid LLM usage source") + if not self.finish_reason.strip(): + raise ValueError("LLM usage finish_reason must be non-empty") diff --git a/nanobot/llm_usage/store.py b/nanobot/llm_usage/store.py new file mode 100644 index 000000000..224cabb61 --- /dev/null +++ b/nanobot/llm_usage/store.py @@ -0,0 +1,560 @@ +"""SQLite persistence and chart queries for LLM usage records.""" + +from __future__ import annotations + +import os +import sqlite3 +import threading +import time +from collections.abc import Iterable +from copy import deepcopy +from datetime import date, datetime, timedelta, timezone +from pathlib import Path +from typing import Any, cast +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError + +from nanobot.llm_usage.models import LLMCallRecord + +SCHEMA_VERSION = 1 +MAX_DAYS_RETAINED = 400 +MAX_CALLS_RETAINED = 100_000 + +_ERROR_KINDS = frozenset({ + "authentication", + "cancelled", + "configuration", + "connection", + "content_filter", + "context_length", + "empty", + "http", + "invalid_request", + "overloaded", + "permission", + "rate_limit", + "refusal", + "server_error", + "timeout", +}) +_FINISH_REASONS = frozenset({ + "cancelled", + "content_filter", + "error", + "function_call", + "length", + "refusal", + "stop", + "tool_calls", +}) + +_USAGE_COLUMNS = ( + "input_tokens", + "output_tokens", + "cache_read_tokens", + "cache_write_tokens", + "cache_read_observed_input_tokens", + "cache_write_observed_input_tokens", + "total_tokens", + "reported_tokens", + "estimated_tokens", + "generation_ms", + "measured_output_tokens", + "ttft_ms", + "timed_requests", +) +_REQUEST_COLUMNS = ( + "requests", + "successful_requests", + "failed_requests", + "reported_requests", + "estimated_requests", +) +_AGGREGATE_SQL = """ + COALESCE(SUM(input_tokens), 0) AS input_tokens, + COALESCE(SUM(output_tokens), 0) AS output_tokens, + COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens, + COALESCE(SUM(cache_write_tokens), 0) AS cache_write_tokens, + COALESCE(SUM( + CASE WHEN cache_read_tokens IS NOT NULL THEN input_tokens ELSE 0 END + ), 0) AS cache_read_observed_input_tokens, + COALESCE(SUM( + CASE WHEN cache_write_tokens IS NOT NULL THEN input_tokens ELSE 0 END + ), 0) AS cache_write_observed_input_tokens, + COALESCE(SUM(total_tokens), 0) AS total_tokens, + COALESCE(SUM(reported_tokens), 0) AS reported_tokens, + COALESCE(SUM(estimated_tokens), 0) AS estimated_tokens, + COALESCE(SUM(generation_ms), 0) AS generation_ms, + COALESCE(SUM(measured_output_tokens), 0) AS measured_output_tokens, + COALESCE(SUM(ttft_ms), 0) AS ttft_ms, + COALESCE(SUM(timed_requests), 0) AS timed_requests, + COUNT(*) AS requests, + COALESCE(SUM(CASE WHEN finish_reason IN ('error', 'cancelled') THEN 0 ELSE 1 END), 0) + AS successful_requests, + COALESCE(SUM(CASE WHEN finish_reason IN ('error', 'cancelled') THEN 1 ELSE 0 END), 0) + AS failed_requests, + COALESCE(SUM( + CASE WHEN total_tokens IS NOT NULL AND NOT ( + estimated_tokens > 0 AND reported_tokens = 0 + ) THEN 1 ELSE 0 END + ), 0) AS reported_requests, + COALESCE(SUM( + CASE WHEN estimated_tokens > 0 AND reported_tokens = 0 THEN 1 ELSE 0 END + ), 0) AS estimated_requests, + COALESCE(SUM(duration_ms), 0) AS duration_ms +""" + + +def _zone(timezone_name: str | None) -> timezone | ZoneInfo: + if not timezone_name: + return timezone.utc + try: + return ZoneInfo(timezone_name) + except ZoneInfoNotFoundError: + return timezone.utc + + +def _clean_error_kind(value: str | None) -> str | None: + if value is None: + return None + cleaned = value.strip().lower() + if not cleaned: + return None + return cleaned if cleaned in _ERROR_KINDS else "other" + + +def _clean_finish_reason(value: str) -> str: + cleaned = value.strip().lower() + return cleaned if cleaned in _FINISH_REASONS else "other" + + +def _clean_status_code(value: int | None) -> int | None: + if value is None: + return None + try: + status = int(value) + except (TypeError, ValueError): + return None + return status if 100 <= status <= 599 else None + + +def _as_int_row(row: sqlite3.Row) -> dict[str, int]: + return { + key: max(0, int(row[key] or 0)) + for key in (*_USAGE_COLUMNS, *_REQUEST_COLUMNS, "duration_ms") + } + + +def _empty_totals() -> dict[str, int]: + return {key: 0 for key in (*_USAGE_COLUMNS, *_REQUEST_COLUMNS, "duration_ms")} + + +def _sum_rows(rows: Iterable[dict[str, Any]]) -> dict[str, int]: + totals = _empty_totals() + for row in rows: + for key in totals: + totals[key] += max(0, int(row.get(key) or 0)) + return totals + + +class LLMUsageStore: + """A small synchronous WAL database shared by gateway threads/processes.""" + + def __init__(self, path: Path) -> None: + self.path = path + self._lock = threading.RLock() + self._connection: sqlite3.Connection | None = None + self._connection_pid: int | None = None + self._last_prune_utc_day: int | None = None + self._writes_since_size_prune = 0 + self._write_version = 0 + self._cached_payload_key: tuple[int, str, str, int, int] | None = None + self._cached_payload: dict[str, Any] | None = None + + def _connect(self) -> sqlite3.Connection: + pid = os.getpid() + if self._connection is not None and self._connection_pid == pid: + return self._connection + if self._connection is not None: + self._connection.close() + self._cached_payload_key = None + self._cached_payload = None + self.path.parent.mkdir(parents=True, exist_ok=True) + connection = sqlite3.connect( + self.path, + timeout=0.25, + isolation_level=None, + check_same_thread=False, + ) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA busy_timeout = 250") + connection.execute("PRAGMA journal_mode = WAL") + connection.execute("PRAGMA synchronous = NORMAL") + connection.execute("PRAGMA temp_store = MEMORY") + connection.create_function("llm_usage_local_day", 2, self._local_day, deterministic=True) + connection.executescript( + """ + CREATE TABLE IF NOT EXISTS llm_calls ( + id INTEGER PRIMARY KEY, + started_at_ms INTEGER NOT NULL, + duration_ms INTEGER NOT NULL, + provider TEXT NOT NULL, + model TEXT NOT NULL, + source TEXT NOT NULL, + stream INTEGER NOT NULL, + finish_reason TEXT NOT NULL, + input_tokens INTEGER, + output_tokens INTEGER, + total_tokens INTEGER, + cache_read_tokens INTEGER, + cache_write_tokens INTEGER, + reported_tokens INTEGER, + estimated_tokens INTEGER, + generation_ms INTEGER, + measured_output_tokens INTEGER, + ttft_ms INTEGER, + timed_requests INTEGER, + error_status_code INTEGER, + error_kind TEXT + ); + CREATE INDEX IF NOT EXISTS llm_calls_started_at_idx + ON llm_calls(started_at_ms); + CREATE INDEX IF NOT EXISTS llm_calls_provider_model_time_idx + ON llm_calls(provider, model, started_at_ms); + """ + ) + connection.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") + self._connection = connection + self._connection_pid = pid + return connection + + def _read_connection(self) -> sqlite3.Connection: + connection = sqlite3.connect( + self.path, + timeout=0.25, + isolation_level=None, + ) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA busy_timeout = 250") + connection.execute("PRAGMA query_only = ON") + connection.execute("PRAGMA temp_store = MEMORY") + connection.create_function("llm_usage_local_day", 2, self._local_day, deterministic=True) + return connection + + @staticmethod + def _local_day(started_at_ms: object, timezone_name: object) -> str | None: + if not isinstance(started_at_ms, int) or not isinstance(timezone_name, str): + return None + dt = datetime.fromtimestamp(started_at_ms / 1000, timezone.utc) + return dt.astimezone(_zone(timezone_name)).date().isoformat() + + def close(self) -> None: + with self._lock: + if self._connection is not None: + self._connection.close() + self._connection = None + self._connection_pid = None + self._cached_payload_key = None + self._cached_payload = None + + def record(self, call: LLMCallRecord) -> None: + usage = call.usage + usage_data = usage.to_dict() if usage is not None else {} + values: tuple[object, ...] = ( + call.started_at_ms, + call.duration_ms, + call.provider[:120], + call.model[:240], + call.source, + int(call.stream), + _clean_finish_reason(call.finish_reason), + *( + usage_data.get(key) + for key in ( + "input_tokens", + "output_tokens", + "total_tokens", + "cache_read_tokens", + "cache_write_tokens", + "reported_tokens", + "estimated_tokens", + "generation_ms", + "measured_output_tokens", + "ttft_ms", + "timed_requests", + ) + ), + _clean_status_code(call.error_status_code), + _clean_error_kind(call.error_kind), + ) + with self._lock: + connection = self._connect() + connection.execute( + """ + INSERT INTO llm_calls ( + started_at_ms, duration_ms, provider, model, source, stream, + finish_reason, input_tokens, output_tokens, total_tokens, + cache_read_tokens, cache_write_tokens, reported_tokens, + estimated_tokens, generation_ms, measured_output_tokens, + ttft_ms, timed_requests, error_status_code, error_kind + ) VALUES ( + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? + ) + """, + values, + ) + self._write_version += 1 + self._cached_payload_key = None + self._cached_payload = None + self._prune_if_due(connection) + + def _prune_if_due(self, connection: sqlite3.Connection) -> None: + utc_day = int(time.time() // 86_400) + self._writes_since_size_prune += 1 + prune_age = self._last_prune_utc_day != utc_day + prune_size = self._writes_since_size_prune >= 1_024 + if not prune_age and not prune_size: + return + if prune_age: + cutoff_ms = int( + (datetime.now(timezone.utc) - timedelta(days=MAX_DAYS_RETAINED)).timestamp() + * 1000 + ) + connection.execute("DELETE FROM llm_calls WHERE started_at_ms < ?", (cutoff_ms,)) + connection.execute( + """ + DELETE FROM llm_calls + WHERE id <= COALESCE(( + SELECT id FROM llm_calls ORDER BY id DESC LIMIT 1 OFFSET ? + ), -1) + """, + (MAX_CALLS_RETAINED,), + ) + self._last_prune_utc_day = utc_day + self._writes_since_size_prune = 0 + + def count(self) -> int: + with self._lock: + row = self._connect().execute("SELECT COUNT(*) AS count FROM llm_calls").fetchone() + return int(row["count"] if row is not None else 0) + + def _aggregate( + self, + *, + connection: sqlite3.Connection, + start_ms: int | None, + end_ms: int, + group_by: tuple[str, ...] = (), + limit: int | None = None, + ) -> list[sqlite3.Row]: + selected = f"{', '.join(group_by)}, " if group_by else "" + where = "started_at_ms < ?" + params: list[object] = [end_ms] + if start_ms is not None: + where = "started_at_ms >= ? AND started_at_ms < ?" + params = [start_ms, end_ms] + query = f"SELECT {selected}{_AGGREGATE_SQL} FROM llm_calls WHERE {where}" + if group_by: + query += f" GROUP BY {', '.join(group_by)} ORDER BY total_tokens DESC" + if limit is not None: + query += " LIMIT ?" + params.append(limit) + return list(connection.execute(query, params).fetchall()) + + def _daily_rows( + self, + *, + connection: sqlite3.Connection, + start_ms: int, + end_ms: int, + timezone_name: str, + ) -> list[dict[str, Any]]: + query = f""" + SELECT llm_usage_local_day(started_at_ms, ?) AS date, source, + {_AGGREGATE_SQL} + FROM llm_calls + WHERE started_at_ms >= ? AND started_at_ms < ? + GROUP BY date, source + ORDER BY date, source + """ + rows = connection.execute( + query, + (timezone_name, start_ms, end_ms), + ).fetchall() + by_date: dict[str, dict[str, Any]] = {} + for row in rows: + day = cast(str | None, row["date"]) + if day is None: + continue + values = _as_int_row(row) + aggregate = by_date.setdefault( + day, + {"date": day, **_empty_totals(), "sources": {}}, + ) + for key, value in values.items(): + aggregate[key] += value + aggregate["sources"][str(row["source"])] = values + return list(by_date.values()) + + @staticmethod + def _midnight_ms(value: date, zone: timezone | ZoneInfo) -> int: + return int(datetime.combine(value, datetime.min.time(), tzinfo=zone).timestamp() * 1000) + + def usage_payload( + self, + *, + days: int = 371, + timezone_name: str | None = None, + now: datetime | None = None, + ) -> dict[str, Any]: + zone = _zone(timezone_name) + current = now or datetime.now(timezone.utc) + if current.tzinfo is None: + current = current.replace(tzinfo=timezone.utc) + today = current.astimezone(zone).date() + safe_days = max(1, days) + zone_name = getattr(zone, "key", "UTC") + + with self._lock: + data_version_row = self._connect().execute("PRAGMA data_version").fetchone() + data_version = int(data_version_row[0]) if data_version_row is not None else 0 + write_version = self._write_version + cache_key = ( + safe_days, + zone_name, + today.isoformat(), + write_version, + data_version, + ) + if self._cached_payload_key == cache_key and self._cached_payload is not None: + return deepcopy(self._cached_payload) + + connection = self._read_connection() + try: + connection.execute("BEGIN") + end_ms = self._midnight_ms(today + timedelta(days=1), zone) + retained_start = today - timedelta(days=MAX_DAYS_RETAINED - 1) + retained_start_ms = self._midnight_ms(retained_start, zone) + daily = self._daily_rows( + connection=connection, + start_ms=retained_start_ms, + end_ms=end_ms, + timezone_name=zone_name, + ) + + requested_start = today - timedelta(days=safe_days - 1) + visible_days = [row for row in daily if row["date"] >= requested_start.isoformat()] + last_30_start_ms = self._midnight_ms(today - timedelta(days=29), zone) + + last_30_date = (today - timedelta(days=29)).isoformat() + last_365_date = (today - timedelta(days=364)).isoformat() + all_totals = _sum_rows(daily) + totals_30 = _sum_rows(row for row in daily if row["date"] >= last_30_date) + totals_365 = _sum_rows(row for row in daily if row["date"] >= last_365_date) + + provider_rows = self._aggregate( + connection=connection, + start_ms=last_30_start_ms, + end_ms=end_ms, + group_by=("provider", "model"), + limit=50, + ) + providers_30d = [ + { + "provider": str(row["provider"]), + "model": str(row["model"]), + **_as_int_row(row), + } + for row in provider_rows + ] + + active_dates = { + date.fromisoformat(row["date"]) for row in daily if row["total_tokens"] > 0 + } + current_streak = 0 + cursor = today + while cursor in active_dates: + current_streak += 1 + cursor -= timedelta(days=1) + longest_streak = 0 + running_streak = 0 + previous: date | None = None + for cursor in sorted(active_dates): + running_streak = running_streak + 1 if previous == cursor - timedelta(days=1) else 1 + longest_streak = max(longest_streak, running_streak) + previous = cursor + + latest = ( + connection + .execute("SELECT MAX(started_at_ms) AS updated_at_ms FROM llm_calls") + .fetchone() + ) + updated_at_ms = int(latest["updated_at_ms"] or 0) if latest is not None else 0 + denominator = totals_30["cache_read_observed_input_tokens"] + payload = { + "days": visible_days, + "total_tokens": all_totals["total_tokens"], + "total_tokens_30d": totals_30["total_tokens"], + "total_tokens_365d": totals_365["total_tokens"], + "reported_tokens_30d": totals_30["reported_tokens"], + "estimated_tokens_30d": totals_30["estimated_tokens"], + "cache_read_tokens_30d": totals_30["cache_read_tokens"], + "cache_read_observed_input_tokens_30d": denominator, + "cache_read_rate_30d": ( + totals_30["cache_read_tokens"] / denominator if denominator else None + ), + "peak_day_tokens": max( + (int(row["total_tokens"]) for row in daily), + default=0, + ), + "current_streak_days": current_streak, + "longest_streak_days": longest_streak, + "active_days_30d": sum( + 1 + for row in daily + if row["date"] >= last_30_date and row["total_tokens"] > 0 + ), + "requests_30d": totals_30["requests"], + "failed_requests_30d": totals_30["failed_requests"], + "providers_30d": providers_30d, + "updated_at": ( + datetime.fromtimestamp(updated_at_ms / 1000, timezone.utc) + .isoformat() + .replace("+00:00", "Z") + if updated_at_ms + else None + ), + } + finally: + connection.close() + + with self._lock: + latest_data_version_row = self._connect().execute("PRAGMA data_version").fetchone() + latest_data_version = ( + int(latest_data_version_row[0]) + if latest_data_version_row is not None + else 0 + ) + if self._write_version == write_version and latest_data_version == data_version: + self._cached_payload_key = cache_key + self._cached_payload = payload + return deepcopy(payload) + + def recent_calls(self, *, limit: int = 100) -> list[dict[str, Any]]: + """Return bounded metadata rows for diagnostics; never returns content.""" + safe_limit = min(max(1, limit), 1_000) + with self._lock: + rows = ( + self._connect() + .execute( + """ + SELECT * FROM llm_calls ORDER BY started_at_ms DESC, id DESC LIMIT ? + """, + (safe_limit,), + ) + .fetchall() + ) + return [dict(row) for row in rows] + + def record_many(self, calls: Iterable[LLMCallRecord]) -> None: + for call in calls: + self.record(call) diff --git a/nanobot/providers/base.py b/nanobot/providers/base.py index 82c12e88e..a794c1adb 100644 --- a/nanobot/providers/base.py +++ b/nanobot/providers/base.py @@ -6,6 +6,7 @@ import asyncio import json import os import re +import time from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable from contextlib import suppress @@ -13,19 +14,23 @@ from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime, timezone from email.utils import parsedate_to_datetime -from typing import Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, cast import json_repair from loguru import logger from nanobot.utils.helpers import sanitize_surrogates_deep +if TYPE_CHECKING: + from nanobot.llm_usage.models import LLMCallRecord + STREAM_IDLE_TIMEOUT_ENV = "NANOBOT_STREAM_IDLE_TIMEOUT_S" DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0 MAX_STREAM_IDLE_TIMEOUT_S = 3600.0 RETRY_AFTER_BUFFER = 1 RetryEventCallback = Callable[[str], Awaitable[None]] +LLMCallObserver = Callable[["LLMCallRecord"], None] def resolve_stream_idle_timeout_s( @@ -682,6 +687,95 @@ class LLMProvider(ABC): self.api_base = api_base self.provider_name = provider_name self.generation: GenerationSettings = GenerationSettings() + self._llm_call_observer: LLMCallObserver | None = None + + def set_llm_call_observer(self, observer: LLMCallObserver | None) -> None: + """Attach a fail-open observer for each physical retry-managed call.""" + self._llm_call_observer = observer + + def _usage_for_call( + self, + response: LLMResponse, + kwargs: dict[str, Any], + ) -> LLMUsage | None: + usage = response.usage + if usage is None or usage.total_tokens == 0: + if response.finish_reason in {"error", "cancelled"}: + return None + messages = kwargs.get("messages") + if not isinstance(messages, list): + return usage + tools_value = kwargs.get("tools") + tools = cast(list[dict[str, Any]], tools_value) if isinstance(tools_value, list) else None + model_value = kwargs.get("model") + model = model_value if isinstance(model_value, str) else self.get_default_model() + try: + from nanobot.utils.helpers import ( + build_assistant_message, + estimate_message_tokens, + estimate_prompt_tokens_chain, + ) + + input_tokens, _ = estimate_prompt_tokens_chain( + self, + model, + cast(list[dict[str, Any]], messages), + tools, + ) + assistant_message = build_assistant_message( + response.content or "", + tool_calls=[call.to_openai_tool_call() for call in response.tool_calls], + reasoning_content=response.reasoning_content, + thinking_blocks=response.thinking_blocks, + ) + usage = LLMUsage.estimated( + input_tokens=max(0, input_tokens), + output_tokens=max(0, estimate_message_tokens(assistant_message)), + ) + except Exception: + logger.exception("failed to estimate usage for {}", self.provider_name) + return usage + return usage.with_timing( + generation_ms=response.generation_ms, + ttft_ms=response.ttft_ms, + ) + + def _observe_llm_call( + self, + response: LLMResponse, + kwargs: dict[str, Any], + *, + started_at_ms: int, + started_at_ns: int, + stream: bool, + ) -> LLMResponse: + observer = self._llm_call_observer + if observer is None: + return response + usage = self._usage_for_call(response, kwargs) + if usage is not None: + response.usage = usage + model_value = kwargs.get("model") + model = model_value if isinstance(model_value, str) and model_value else self.get_default_model() + try: + from nanobot.llm_usage.context import current_llm_usage_source + from nanobot.llm_usage.models import LLMCallRecord + + observer(LLMCallRecord( + started_at_ms=started_at_ms, + duration_ms=max(0, (time.monotonic_ns() - started_at_ns) // 1_000_000), + provider=self.provider_name, + model=model, + source=current_llm_usage_source(), + stream=stream, + finish_reason=response.finish_reason, + usage=usage, + error_status_code=response.error_status_code, + error_kind=response.error_kind, + )) + except Exception: + logger.exception("LLM call observer failed for {}", self.provider_name) + return response def can_resume_conversation_state( self, @@ -1068,18 +1162,39 @@ class LLMProvider(ABC): async def _safe_chat(self, **kwargs: Any) -> LLMResponse: """Call chat() and convert unexpected exceptions to error responses.""" + started_at_ms = time.time_ns() // 1_000_000 + started_at_ns = time.monotonic_ns() try: provider_context = kwargs.pop("provider_context", None) if isinstance(provider_context, ProviderCallContext): - return await self.chat_with_context( + response = await self.chat_with_context( provider_context=provider_context, **kwargs, ) - return await self.chat(**kwargs) + else: + response = await self.chat(**kwargs) except asyncio.CancelledError: + self._observe_llm_call( + LLMResponse( + content=None, + finish_reason="cancelled", + error_kind="cancelled", + ), + kwargs, + started_at_ms=started_at_ms, + started_at_ns=started_at_ns, + stream=False, + ) raise except Exception as exc: - return LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error") + response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error") + return self._observe_llm_call( + response, + kwargs, + started_at_ms=started_at_ms, + started_at_ns=started_at_ns, + stream=False, + ) async def chat_stream( self, @@ -1142,18 +1257,39 @@ class LLMProvider(ABC): async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse: """Call chat_stream() and convert unexpected exceptions to error responses.""" + started_at_ms = time.time_ns() // 1_000_000 + started_at_ns = time.monotonic_ns() try: provider_context = kwargs.pop("provider_context", None) if isinstance(provider_context, ProviderCallContext): - return await self.chat_stream_with_context( + response = await self.chat_stream_with_context( provider_context=provider_context, **kwargs, ) - return await self.chat_stream(**kwargs) + else: + response = await self.chat_stream(**kwargs) except asyncio.CancelledError: + self._observe_llm_call( + LLMResponse( + content=None, + finish_reason="cancelled", + error_kind="cancelled", + ), + kwargs, + started_at_ms=started_at_ms, + started_at_ns=started_at_ns, + stream=True, + ) raise except Exception as exc: - return LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error") + response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error") + return self._observe_llm_call( + response, + kwargs, + started_at_ms=started_at_ms, + started_at_ns=started_at_ns, + stream=True, + ) async def chat_stream_with_retry( self, diff --git a/nanobot/providers/fallback_provider.py b/nanobot/providers/fallback_provider.py index 0920a724f..24a55227a 100644 --- a/nanobot/providers/fallback_provider.py +++ b/nanobot/providers/fallback_provider.py @@ -13,6 +13,7 @@ from loguru import logger from nanobot.providers.base import ( GenerationSettings, + LLMCallObserver, LLMProvider, LLMResponse, ProviderCallContext, @@ -151,6 +152,11 @@ class FallbackProvider(LLMProvider): """Attach a process-level observer without changing request call signatures.""" self._fallback_model_observer = observer + def set_llm_call_observer(self, observer: LLMCallObserver | None) -> None: + """Attach usage recording to the primary and future fallback leaves.""" + super().set_llm_call_observer(observer) + self._primary.set_llm_call_observer(observer) + @property def supports_progress_deltas(self) -> bool: return bool(getattr(self._primary, "supports_progress_deltas", False)) @@ -506,6 +512,7 @@ class FallbackProvider(LLMProvider): ) try: fallback_provider = self._provider_factory(fallback) + fallback_provider.set_llm_call_observer(self._llm_call_observer) except Exception as exc: logger.warning( "Failed to create provider for fallback '{}': {}", fallback_model, exc diff --git a/nanobot/session/webui_turns.py b/nanobot/session/webui_turns.py index a4fc5d0b7..a403b69b0 100644 --- a/nanobot/session/webui_turns.py +++ b/nanobot/session/webui_turns.py @@ -37,6 +37,7 @@ from nanobot.bus.runtime_events import ( TurnRuntimeAdmitted, UserInputAccepted, ) +from nanobot.llm_usage.context import llm_usage_source from nanobot.providers.base import LLMProvider, LLMUsage from nanobot.providers.fallback_provider import FallbackModelObserver from nanobot.runtime_context import public_history_message @@ -208,24 +209,25 @@ async def maybe_generate_webui_title( prompt += f"\nAssistant: {truncate_text(assistant_text, 1_000)}" try: - response = await provider.chat_with_retry( - [ - { - "role": "system", - "content": ( - "You write short, neutral chat titles. " - "Return only the title text." - ), - }, - {"role": "user", "content": prompt}, - ], - tools=None, - model=model, - max_tokens=TITLE_GENERATION_MAX_TOKENS, - temperature=0.2, - reasoning_effort=TITLE_GENERATION_REASONING_EFFORT, - retry_mode="standard", - ) + with llm_usage_source("system"): + response = await provider.chat_with_retry( + [ + { + "role": "system", + "content": ( + "You write short, neutral chat titles. " + "Return only the title text." + ), + }, + {"role": "user", "content": prompt}, + ], + tools=None, + model=model, + max_tokens=TITLE_GENERATION_MAX_TOKENS, + temperature=0.2, + reasoning_effort=TITLE_GENERATION_REASONING_EFFORT, + retry_mode="standard", + ) except Exception: logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True) return False diff --git a/nanobot/webui/settings_routes.py b/nanobot/webui/settings_routes.py index 25f79ccdf..cbb5b4df0 100644 --- a/nanobot/webui/settings_routes.py +++ b/nanobot/webui/settings_routes.py @@ -284,9 +284,9 @@ class WebUISettingsRouter: if not self._authorized(request): return self._unauthorized() if route == ("root", "settings"): - return self._handle_settings() + return await asyncio.to_thread(self._handle_settings) if route == ("root", "usage"): - return self._handle_settings_usage() + return await asyncio.to_thread(self._handle_settings_usage) domain, action = route domain_request = self._domain_request( diff --git a/nanobot/webui/settings_system.py b/nanobot/webui/settings_system.py index f57394742..fa4a5f7b2 100644 --- a/nanobot/webui/settings_system.py +++ b/nanobot/webui/settings_system.py @@ -20,6 +20,7 @@ from nanobot.channels.contracts import ( channel_update_instance_config, ) from nanobot.config.schema import Config +from nanobot.llm_usage import llm_usage_payload from nanobot.optional_features import OptionalFeatureError, with_channel_runtime_status from nanobot.security.workspace_access import workspace_sandbox_status from nanobot.webui.settings_capabilities import network_safety_payload @@ -31,7 +32,6 @@ from nanobot.webui.settings_contracts import ( query_first, query_first_alias, ) -from nanobot.webui.token_usage import token_usage_payload if TYPE_CHECKING: from nanobot.webui.settings_services import WebUISettingsServices @@ -121,7 +121,7 @@ def system_settings_payload( }, "unified_session": defaults.unified_session, }, - "usage": token_usage_payload(timezone_name=defaults.timezone), + "usage": llm_usage_payload(timezone_name=defaults.timezone), "advanced": { "restrict_to_workspace": config.tools.restrict_to_workspace, "workspace_sandbox": sandbox_status.as_dict(), @@ -139,7 +139,7 @@ def system_settings_payload( def settings_usage_payload(config: Config) -> dict[str, Any]: """Return the lightweight token usage slice for Overview refreshes.""" - return token_usage_payload(timezone_name=config.agents.defaults.timezone) + return llm_usage_payload(timezone_name=config.agents.defaults.timezone) def update_agent_system_settings(config: Config, query: QueryParams) -> tuple[bool, bool]: diff --git a/nanobot/webui/token_usage.py b/nanobot/webui/token_usage.py deleted file mode 100644 index 59264e620..000000000 --- a/nanobot/webui/token_usage.py +++ /dev/null @@ -1,392 +0,0 @@ -"""Workspace-scoped token usage telemetry for WebUI overview surfaces.""" - -from __future__ import annotations - -import json -import os -import threading -import time -from datetime import datetime, timedelta, timezone -from pathlib import Path -from typing import Any, Mapping, cast -from zoneinfo import ZoneInfo, ZoneInfoNotFoundError - -from loguru import logger - -from nanobot.agent.hook import AgentHook, AgentHookContext -from nanobot.config.paths import get_webui_dir -from nanobot.providers.base import LLMUsage - -TOKEN_USAGE_SCHEMA_VERSION = 2 -_MAX_STATE_FILE_BYTES = 512 * 1024 -_MAX_DAYS_RETAINED = 400 -_USAGE_KEYS = ( - "input_tokens", - "output_tokens", - "cache_read_tokens", - "cache_write_tokens", - "cache_read_observed_input_tokens", - "cache_write_observed_input_tokens", - "total_tokens", - "reported_tokens", - "estimated_tokens", -) -_REQUEST_KEYS = ("requests", "reported_requests", "estimated_requests") -_SOURCE_KEYS = ("user", "api", "cron", "dream", "system") -_WRITE_LOCK = threading.Lock() - - -def token_usage_state_path() -> Path: - return get_webui_dir() / "token-usage.json" - - -def default_token_usage_state() -> dict[str, Any]: - return { - "schema_version": TOKEN_USAGE_SCHEMA_VERSION, - "days": {}, - "updated_at": None, - } - - -def _utc_now_iso() -> str: - return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) - - -def _zone(timezone_name: str | None) -> timezone | ZoneInfo: - if not timezone_name: - return timezone.utc - try: - return ZoneInfo(timezone_name) - except ZoneInfoNotFoundError: - return timezone.utc - - -def _local_day(now: datetime | None = None, *, timezone_name: str | None = None) -> str: - dt = now or datetime.now(timezone.utc) - if dt.tzinfo is None: - dt = dt.replace(tzinfo=timezone.utc) - return dt.astimezone(_zone(timezone_name)).date().isoformat() - - -def _clean_int(value: Any) -> int: - try: - return max(0, int(value or 0)) - except (TypeError, ValueError): - return 0 - - -def _clean_source(value: str | None) -> str: - return value if value in _SOURCE_KEYS else "system" - - -def _source_from_session_key(session_key: str | None) -> str: - key = session_key or "" - if key.startswith("dream:"): - return "dream" - if key == "heartbeat" or key.startswith("cron:"): - return "cron" - if key.startswith("api:"): - return "api" - if key.startswith("system:"): - return "system" - return "user" - - -def _normalize_usage(raw: LLMUsage | None) -> dict[str, int]: - if raw is None: - return {} - usage = { - "input_tokens": raw.input_tokens, - "output_tokens": raw.output_tokens, - "cache_read_tokens": raw.cache_read_tokens or 0, - "cache_write_tokens": raw.cache_write_tokens or 0, - "cache_read_observed_input_tokens": ( - raw.input_tokens if raw.cache_read_tokens is not None else 0 - ), - "cache_write_observed_input_tokens": ( - raw.input_tokens if raw.cache_write_tokens is not None else 0 - ), - "total_tokens": raw.total_tokens, - "reported_tokens": raw.reported_tokens, - "estimated_tokens": raw.estimated_tokens, - } - return usage if usage["total_tokens"] > 0 else {} - - -def _normalize_usage_row(row: dict[str, Any]) -> dict[str, int]: - cleaned = {key: _clean_int(row.get(key)) for key in _USAGE_KEYS} - if cleaned["total_tokens"] <= 0: - cleaned["total_tokens"] = cleaned["input_tokens"] + cleaned["output_tokens"] - if cleaned["reported_tokens"] <= 0 and cleaned["estimated_tokens"] <= 0: - cleaned["reported_tokens"] = cleaned["total_tokens"] - requests = {key: _clean_int(row.get(key)) for key in _REQUEST_KEYS} - if ( - requests["requests"] > 0 - and requests["reported_requests"] <= 0 - and requests["estimated_requests"] <= 0 - ): - if cleaned["estimated_tokens"] > 0 and cleaned["reported_tokens"] <= 0: - requests["estimated_requests"] = requests["requests"] - else: - requests["reported_requests"] = requests["requests"] - return {**cleaned, **requests} - - -def _normalize_sources(raw: Any, fallback: dict[str, int]) -> dict[str, dict[str, int]]: - sources: dict[str, dict[str, int]] = {} - if isinstance(raw, dict): - for source, row_value in cast(dict[Any, Any], raw).items(): - if not isinstance(row_value, dict): - continue - row = cast(dict[str, Any], row_value) - normalized = _normalize_usage_row(row) - if normalized["total_tokens"] <= 0 and normalized["requests"] <= 0: - continue - source_key = _clean_source(str(source)) - current = sources.get(source_key) - if current is None: - sources[source_key] = normalized - else: - for key in (*_USAGE_KEYS, *_REQUEST_KEYS): - current[key] = _clean_int(current.get(key)) + normalized[key] - if not sources and (fallback["total_tokens"] > 0 or fallback["requests"] > 0): - sources["user"] = {key: fallback[key] for key in (*_USAGE_KEYS, *_REQUEST_KEYS)} - return sources - - -def normalize_token_usage_state(raw: Any) -> dict[str, Any]: - state = default_token_usage_state() - if not isinstance(raw, dict): - return state - raw = cast(dict[str, Any], raw) - if raw.get("schema_version") != TOKEN_USAGE_SCHEMA_VERSION: - return state - days_raw = raw.get("days") - if not isinstance(days_raw, dict): - return state - - days: dict[str, dict[str, Any]] = {} - for date, row_value in sorted(cast(dict[Any, Any], days_raw).items())[-_MAX_DAYS_RETAINED:]: - if not isinstance(date, str) or len(date) != 10 or not isinstance(row_value, dict): - continue - row = cast(dict[str, Any], row_value) - try: - datetime.fromisoformat(date) - except ValueError: - # A hand-edited or foreign day key that is not a real date would - # otherwise reach token_usage_payload's date parsing and fail every - # settings request; drop it like any other malformed row. - continue - normalized = _normalize_usage_row(row) - if normalized["total_tokens"] <= 0 and normalized["requests"] <= 0: - continue - days[date] = { - "date": date, - **normalized, - "sources": _normalize_sources(row.get("sources"), normalized), - } - - state["days"] = days - updated_at = raw.get("updated_at") - state["updated_at"] = updated_at if isinstance(updated_at, str) else None - return state - - -def read_token_usage_state() -> dict[str, Any]: - path = token_usage_state_path() - if not path.is_file(): - return default_token_usage_state() - try: - if path.stat().st_size > _MAX_STATE_FILE_BYTES: - logger.warning("token usage state too large, ignoring: {}", path) - return default_token_usage_state() - with open(path, encoding="utf-8") as f: - raw = json.load(f) - except (OSError, json.JSONDecodeError) as e: - logger.warning("read token usage state failed {}: {}", path, e) - return default_token_usage_state() - return normalize_token_usage_state(raw) - - -def _encode_token_usage_state(state: dict[str, Any]) -> bytes: - """Encode the persisted state compactly, including its trailing newline.""" - payload = json.dumps( - state, - ensure_ascii=False, - separators=(",", ":"), - sort_keys=True, - ) - return f"{payload}\n".encode("utf-8") - - -def write_token_usage_state(raw: dict[str, Any]) -> dict[str, Any]: - # Day-count retention is applied by normalization first. The byte budget - # then trims only the oldest remaining days, preserving a contiguous suffix. - state = normalize_token_usage_state(raw) - state["updated_at"] = _utc_now_iso() - days = cast(dict[str, dict[str, Any]], state["days"]) - encoded = _encode_token_usage_state(state) - while len(encoded) > _MAX_STATE_FILE_BYTES and len(days) > 1: - del days[min(days)] - encoded = _encode_token_usage_state(state) - if len(encoded) > _MAX_STATE_FILE_BYTES: - raise ValueError("latest token usage day exceeds the state byte limit") - - path = token_usage_state_path() - path.parent.mkdir(parents=True, exist_ok=True) - tmp = path.with_suffix(".json.tmp") - with open(tmp, "wb") as f: - f.write(encoded) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, path) - try: - dir_fd = os.open(path.parent, os.O_RDONLY) - except OSError: - return state - try: - os.fsync(dir_fd) - finally: - os.close(dir_fd) - return state - - -def record_token_usage( - usage: LLMUsage | None, - *, - source: str = "user", - timezone_name: str | None = None, - now: datetime | None = None, -) -> dict[str, Any]: - normalized = _normalize_usage(usage) - if not normalized: - return read_token_usage_state() - - with _WRITE_LOCK: - state = read_token_usage_state() - days_by_date = cast(dict[str, dict[str, Any]], state["days"]) - day = _local_day(now, timezone_name=timezone_name) - row: dict[str, Any] = dict(days_by_date.get(day) or {"date": day, "requests": 0}) - for key in _USAGE_KEYS: - row[key] = _clean_int(row.get(key)) + normalized.get(key, 0) - row["requests"] = _clean_int(row.get("requests")) + 1 - if normalized.get("estimated_tokens", 0) > 0 and normalized.get("reported_tokens", 0) <= 0: - row["estimated_requests"] = _clean_int(row.get("estimated_requests")) + 1 - else: - row["reported_requests"] = _clean_int(row.get("reported_requests")) + 1 - - source_key = _clean_source(source) - sources: dict[str, dict[str, Any]] = dict( - cast(Mapping[str, dict[str, Any]], row.get("sources") or {}) - ) - source_row: dict[str, Any] = dict(sources.get(source_key) or {"requests": 0}) - for key in _USAGE_KEYS: - source_row[key] = _clean_int(source_row.get(key)) + normalized.get(key, 0) - source_row["requests"] = _clean_int(source_row.get("requests")) + 1 - if normalized.get("estimated_tokens", 0) > 0 and normalized.get("reported_tokens", 0) <= 0: - source_row["estimated_requests"] = _clean_int(source_row.get("estimated_requests")) + 1 - else: - source_row["reported_requests"] = _clean_int(source_row.get("reported_requests")) + 1 - sources[source_key] = source_row - row["sources"] = sources - - days_by_date[day] = row - if len(days_by_date) > _MAX_DAYS_RETAINED: - state["days"] = dict(sorted(days_by_date.items())[-_MAX_DAYS_RETAINED:]) - return write_token_usage_state(state) - - -def record_response_token_usage( - response: Any, - *, - source: str, - timezone_name: str | None = None, -) -> None: - try: - record_token_usage( - getattr(response, "usage", None), - source=source, - timezone_name=timezone_name, - ) - except Exception: - logger.exception("failed to record {} token usage", source) - - -def token_usage_payload( - *, - days: int = 371, - timezone_name: str | None = None, - now: datetime | None = None, -) -> dict[str, Any]: - state = read_token_usage_state() - days_by_date = cast(dict[str, dict[str, Any]], state["days"]) - today = datetime.fromisoformat(_local_day(now, timezone_name=timezone_name)).date() - start = today - timedelta(days=max(1, days) - 1) - day_rows = [ - row - for date, row in sorted(days_by_date.items()) - if start.isoformat() <= date <= today.isoformat() - ] - last_30_start = today - timedelta(days=29) - last_30 = [ - row - for date, row in days_by_date.items() - if last_30_start.isoformat() <= date <= today.isoformat() - ] - last_365_start = today - timedelta(days=364) - last_365 = [ - row - for date, row in days_by_date.items() - if last_365_start.isoformat() <= date <= today.isoformat() - ] - active_dates = { - datetime.fromisoformat(date).date() - for date, row in days_by_date.items() - if _clean_int(row.get("total_tokens")) > 0 - } - current_streak = 0 - cursor = today - while cursor in active_dates: - current_streak += 1 - cursor -= timedelta(days=1) - - longest_streak = 0 - running_streak = 0 - for cursor in sorted(active_dates): - if cursor - timedelta(days=1) in active_dates: - running_streak += 1 - else: - running_streak = 1 - longest_streak = max(longest_streak, running_streak) - - all_rows = list(days_by_date.values()) - return { - "days": day_rows, - "total_tokens": sum(_clean_int(row.get("total_tokens")) for row in all_rows), - "total_tokens_30d": sum(_clean_int(row.get("total_tokens")) for row in last_30), - "total_tokens_365d": sum(_clean_int(row.get("total_tokens")) for row in last_365), - "peak_day_tokens": max([_clean_int(row.get("total_tokens")) for row in all_rows] or [0]), - "current_streak_days": current_streak, - "longest_streak_days": longest_streak, - "active_days_30d": sum(1 for row in last_30 if _clean_int(row.get("total_tokens")) > 0), - "requests_30d": sum(_clean_int(row.get("requests")) for row in last_30), - "updated_at": state.get("updated_at"), - } - - -class TokenUsageHook(AgentHook): - """Persist provider-reported token usage without coupling it to chat messages.""" - - def __init__(self, *, timezone_name: str | None = None) -> None: - super().__init__() - self._timezone_name = timezone_name - - async def after_iteration(self, context: AgentHookContext) -> None: - try: - record_token_usage( - context.usage, - source=_source_from_session_key(context.session_key), - timezone_name=self._timezone_name, - ) - except Exception: - logger.exception("failed to record token usage") diff --git a/tests/agent/test_subagent.py b/tests/agent/test_subagent.py index 724a31f78..a38b17c05 100644 --- a/tests/agent/test_subagent.py +++ b/tests/agent/test_subagent.py @@ -1,5 +1,6 @@ """Tests for SubagentManager.""" +import asyncio from pathlib import Path from unittest.mock import AsyncMock, MagicMock @@ -10,6 +11,7 @@ from nanobot.agent.subagent import SubagentManager, SubagentStatus from nanobot.agent.tools.filesystem import FileToolsConfig from nanobot.bus.queue import MessageBus from nanobot.config.schema import ToolsConfig +from nanobot.llm_usage.context import llm_usage_source from nanobot.providers.base import GenerationSettings, LLMProvider from nanobot.security.workspace_access import build_workspace_scope from nanobot.utils.llm_runtime import LLMRuntime @@ -198,3 +200,30 @@ async def test_subagent_forwards_fail_on_tool_error_to_runner(tmp_path): spec = sm.runner.run.call_args.args[0] assert spec.fail_on_tool_error is False + + +@pytest.mark.asyncio +async def test_spawned_subagent_inherits_llm_usage_source(tmp_path): + provider = MagicMock(spec=LLMProvider) + provider.get_default_model.return_value = "test" + sm = SubagentManager( + workspace=tmp_path, + bus=MessageBus(), + max_tool_result_chars=16_000, + ) + sm.runner.run = AsyncMock( + return_value=AgentRunResult(final_content="ok", messages=[], stop_reason="completed") + ) + sm._announce_result = AsyncMock() + + with llm_usage_source("cron"): + await sm.spawn( + "automation task", + session_key="websocket:bound-automation", + runtime=_runtime(provider), + ) + tasks = list(sm._running_tasks.values()) + await asyncio.gather(*tasks) + + spec = sm.runner.run.call_args.args[0] + assert spec.llm_usage_source == "cron" diff --git a/tests/llm_usage/test_llm_usage_context.py b/tests/llm_usage/test_llm_usage_context.py new file mode 100644 index 000000000..49728e693 --- /dev/null +++ b/tests/llm_usage/test_llm_usage_context.py @@ -0,0 +1,19 @@ +from nanobot.llm_usage.context import source_from_request + + +def test_automation_metadata_overrides_user_session_source() -> None: + assert source_from_request( + "websocket:ordinary-session", + channel="websocket", + metadata={"_cron_trigger": {"job_id": "job"}}, + ) == "cron" + assert source_from_request( + "websocket:ordinary-session", + channel="websocket", + metadata={"_local_trigger": {"trigger_id": "trigger"}}, + ) == "cron" + + +def test_api_and_system_channels_have_explicit_sources() -> None: + assert source_from_request("shared-session", channel="api", metadata={}) == "api" + assert source_from_request("shared-session", channel="system", metadata={}) == "system" diff --git a/tests/llm_usage/test_llm_usage_store.py b/tests/llm_usage/test_llm_usage_store.py new file mode 100644 index 000000000..a6b0d3b69 --- /dev/null +++ b/tests/llm_usage/test_llm_usage_store.py @@ -0,0 +1,275 @@ +"""Tests for the SQLite LLM usage store.""" + +from __future__ import annotations + +import sqlite3 +import threading +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import pytest + +from nanobot.llm_usage.context import LLMUsageSource +from nanobot.llm_usage.models import LLMCallRecord +from nanobot.llm_usage.store import SCHEMA_VERSION, LLMUsageStore +from nanobot.providers.base import LLMUsage + + +def _timestamp(value: str) -> int: + return int(datetime.fromisoformat(value).timestamp() * 1000) + + +def _call( + started_at: str, + *, + provider: str = "openai", + model: str = "gpt-5", + source: LLMUsageSource = "user", + usage: LLMUsage | None = None, + finish_reason: str = "stop", + error_kind: str | None = None, +) -> LLMCallRecord: + return LLMCallRecord( + started_at_ms=_timestamp(started_at), + duration_ms=250, + provider=provider, + model=model, + source=source, + stream=True, + finish_reason=finish_reason, + usage=usage, + error_status_code=429 if finish_reason == "error" else None, + error_kind=error_kind or ("rate_limit" if finish_reason == "error" else None), + ) + + +def test_store_keeps_only_content_free_call_metadata(tmp_path: Path) -> None: + path = tmp_path / "llm_usage.sqlite3" + store = LLMUsageStore(path) + store.record( + _call( + "2026-06-03T00:00:00+00:00", + usage=LLMUsage.reported(input_tokens=100, output_tokens=20), + ) + ) + + row = store.recent_calls(limit=1)[0] + assert row["provider"] == "openai" + assert row["model"] == "gpt-5" + assert row["total_tokens"] == 120 + assert not { + "messages", + "prompt", + "content", + "response", + "tool_calls", + "error_type", + "error_code", + } & set(row) + + with sqlite3.connect(path) as connection: + version = connection.execute("PRAGMA user_version").fetchone()[0] + mode = connection.execute("PRAGMA journal_mode").fetchone()[0] + assert version == SCHEMA_VERSION + assert str(mode).lower() == "wal" + + +def test_usage_payload_aggregates_cache_coverage_sources_and_failures(tmp_path: Path) -> None: + store = LLMUsageStore(tmp_path / "llm_usage.sqlite3") + store.record_many( + [ + _call( + "2026-06-02T23:30:00+00:00", + usage=LLMUsage.reported( + input_tokens=100, + output_tokens=20, + cache_read_tokens=40, + cache_write_tokens=10, + ), + ), + _call( + "2026-06-03T01:00:00+00:00", + source="api", + usage=LLMUsage.reported(input_tokens=50, output_tokens=5), + ), + _call( + "2026-06-03T02:00:00+00:00", + provider="anthropic", + model="claude-sonnet-4", + source="dream", + usage=LLMUsage.estimated(input_tokens=30, output_tokens=10), + ), + _call( + "2026-06-03T03:00:00+00:00", + provider="anthropic", + model="claude-sonnet-4", + source="system", + finish_reason="error", + ), + ] + ) + + payload = store.usage_payload( + timezone_name="Asia/Shanghai", + now=datetime(2026, 6, 3, 8, tzinfo=timezone.utc), + ) + + assert payload["total_tokens_30d"] == 215 + assert payload["reported_tokens_30d"] == 175 + assert payload["estimated_tokens_30d"] == 40 + assert payload["requests_30d"] == 4 + assert payload["failed_requests_30d"] == 1 + assert payload["cache_read_tokens_30d"] == 40 + assert payload["cache_read_observed_input_tokens_30d"] == 100 + assert payload["cache_read_rate_30d"] == 0.4 + + day = payload["days"][0] + assert day["date"] == "2026-06-03" + assert day["requests"] == 4 + assert day["reported_requests"] == 2 + assert day["estimated_requests"] == 1 + assert day["sources"]["api"]["cache_read_observed_input_tokens"] == 0 + assert day["sources"]["user"]["cache_read_observed_input_tokens"] == 100 + assert {(row["provider"], row["model"]) for row in payload["providers_30d"]} == { + ("openai", "gpt-5"), + ("anthropic", "claude-sonnet-4"), + } + + +def test_usage_payload_preserves_zero_cache_observation(tmp_path: Path) -> None: + store = LLMUsageStore(tmp_path / "llm_usage.sqlite3") + store.record( + _call( + "2026-06-03T00:00:00+00:00", + usage=LLMUsage.reported( + input_tokens=80, + output_tokens=5, + cache_read_tokens=0, + ), + ) + ) + + payload = store.usage_payload( + now=datetime(2026, 6, 3, 12, tzinfo=timezone.utc), + ) + + assert payload["cache_read_tokens_30d"] == 0 + assert payload["cache_read_observed_input_tokens_30d"] == 80 + assert payload["cache_read_rate_30d"] == 0.0 + + +def test_recent_calls_is_bounded(tmp_path: Path) -> None: + store = LLMUsageStore(tmp_path / "llm_usage.sqlite3") + call = _call( + "2026-06-03T00:00:00+00:00", + usage=LLMUsage.reported(input_tokens=1, output_tokens=1), + ) + store.record_many(call for _ in range(1_005)) + + assert len(store.recent_calls(limit=10_000)) == 1_000 + + +def test_cancelled_calls_are_failures_and_error_kind_is_coarse(tmp_path: Path) -> None: + store = LLMUsageStore(tmp_path / "llm_usage.sqlite3") + call = _call( + "2026-06-03T00:00:00+00:00", + finish_reason="cancelled", + error_kind="provider payload: secret text", + ) + store.record(call) + + payload = store.usage_payload( + now=datetime(2026, 6, 3, 12, tzinfo=timezone.utc), + ) + + assert payload["failed_requests_30d"] == 1 + assert store.recent_calls(limit=1)[0]["error_kind"] == "other" + + +def test_usage_payload_cache_is_isolated_and_invalidated_on_write(tmp_path: Path) -> None: + store = LLMUsageStore(tmp_path / "llm_usage.sqlite3") + first_call = _call( + "2026-06-03T00:00:00+00:00", + usage=LLMUsage.reported(input_tokens=10, output_tokens=2), + ) + store.record(first_call) + kwargs = {"now": datetime(2026, 6, 3, 12, tzinfo=timezone.utc)} + + first = store.usage_payload(**kwargs) + first["days"].clear() + cached = store.usage_payload(**kwargs) + assert cached["total_tokens"] == 12 + assert cached["days"] + + store.record(first_call) + refreshed = store.usage_payload(**kwargs) + assert refreshed["total_tokens"] == 24 + + +def test_usage_payload_cache_is_invalidated_when_connection_pid_changes( + tmp_path: Path, +) -> None: + path = tmp_path / "llm_usage.sqlite3" + store = LLMUsageStore(path) + other_store = LLMUsageStore(path) + first_call = _call( + "2026-06-03T00:00:00+00:00", + usage=LLMUsage.reported(input_tokens=1, output_tokens=0), + ) + kwargs = {"now": datetime(2026, 6, 3, 12, tzinfo=timezone.utc)} + store.record(first_call) + assert store.usage_payload(**kwargs)["total_tokens"] == 1 + + other_store.record(_call( + "2026-06-03T00:01:00+00:00", + usage=LLMUsage.reported(input_tokens=2, output_tokens=0), + )) + store._connection_pid = -1 + + assert store.usage_payload(**kwargs)["total_tokens"] == 3 + other_store.close() + + +def test_usage_query_does_not_hold_writer_lock( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = LLMUsageStore(tmp_path / "llm_usage.sqlite3") + call = _call( + "2026-06-03T00:00:00+00:00", + usage=LLMUsage.reported(input_tokens=10, output_tokens=2), + ) + store.record(call) + query_started = threading.Event() + release_query = threading.Event() + record_finished = threading.Event() + original_daily_rows = store._daily_rows + + def slow_daily_rows(**kwargs: Any): + query_started.set() + assert release_query.wait(timeout=2) + return original_daily_rows(**kwargs) + + def record_call() -> None: + store.record(call) + record_finished.set() + + monkeypatch.setattr(store, "_daily_rows", slow_daily_rows) + query_thread = threading.Thread(target=lambda: store.usage_payload( + now=datetime(2026, 6, 3, 12, tzinfo=timezone.utc), + )) + record_thread = threading.Thread(target=record_call) + + query_thread.start() + assert query_started.wait(timeout=2) + record_thread.start() + try: + assert record_finished.wait(timeout=0.5) + finally: + release_query.set() + query_thread.join(timeout=2) + record_thread.join(timeout=2) + + assert not query_thread.is_alive() + assert not record_thread.is_alive() diff --git a/tests/providers/test_llm_usage_observer.py b/tests/providers/test_llm_usage_observer.py new file mode 100644 index 000000000..6c22ca5ac --- /dev/null +++ b/tests/providers/test_llm_usage_observer.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Iterator +from types import SimpleNamespace + +import pytest + +from nanobot.llm_usage.context import llm_usage_source +from nanobot.llm_usage.models import LLMCallRecord +from nanobot.providers.base import LLMProvider, LLMResponse, LLMUsage +from nanobot.providers.fallback_provider import FallbackProvider + + +class _SequenceProvider(LLMProvider): + _CHAT_RETRY_DELAYS = (0,) + + def __init__(self, responses: Iterator[LLMResponse]) -> None: + super().__init__(provider_name="test-provider") + self._responses = responses + + async def chat(self, **_kwargs: object) -> LLMResponse: + return next(self._responses) + + def get_default_model(self) -> str: + return "test-model" + + +class _NoRetryProvider(_SequenceProvider): + _CHAT_RETRY_DELAYS = () + + +class _BlockingProvider(LLMProvider): + async def chat(self, **_kwargs: object) -> LLMResponse: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async def chat_stream(self, **_kwargs: object) -> LLMResponse: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + def get_default_model(self) -> str: + return "test-model" + + +@pytest.mark.asyncio +async def test_observer_receives_every_retry_attempt() -> None: + provider = _SequenceProvider( + iter( + [ + LLMResponse( + content="temporary failure", + finish_reason="error", + error_kind="timeout", + ), + LLMResponse( + content="ok", + usage=LLMUsage.reported( + input_tokens=100, + output_tokens=20, + cache_read_tokens=60, + ), + ), + ] + ) + ) + events: list[LLMCallRecord] = [] + provider.set_llm_call_observer(events.append) + + with llm_usage_source("api"): + response = await provider.chat_with_retry( + messages=[{"role": "user", "content": "hello"}], + model="selected-model", + ) + + assert response.finish_reason == "stop" + assert len(events) == 2 + assert [event.finish_reason for event in events] == ["error", "stop"] + assert all(event.provider == "test-provider" for event in events) + assert all(event.model == "selected-model" for event in events) + assert all(event.source == "api" for event in events) + assert events[1].usage is not None + assert events[1].usage.cache_read_tokens == 60 + + +@pytest.mark.asyncio +async def test_observer_estimates_missing_success_usage_without_storing_content() -> None: + provider = _SequenceProvider(iter([LLMResponse(content="hello")])) + events: list[LLMCallRecord] = [] + provider.set_llm_call_observer(events.append) + + response = await provider.chat_with_retry( + messages=[{"role": "user", "content": "hello"}], + ) + + assert response.usage is not None + assert response.usage.source == "estimated" + assert events[0].usage == response.usage + assert "content" not in LLMCallRecord.__dataclass_fields__ + + +@pytest.mark.asyncio +async def test_observer_failure_never_breaks_provider_call() -> None: + provider = _SequenceProvider(iter([LLMResponse(content="ok")])) + + def _fail(_event: LLMCallRecord) -> None: + raise RuntimeError("disk unavailable") + + provider.set_llm_call_observer(_fail) + response = await provider.chat_with_retry( + messages=[{"role": "user", "content": "hello"}], + ) + + assert response.content == "ok" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_observer_records_cancelled_provider_attempt(stream: bool) -> None: + provider = _BlockingProvider(provider_name="blocking-provider") + events: list[LLMCallRecord] = [] + provider.set_llm_call_observer(events.append) + call = provider.chat_stream_with_retry if stream else provider.chat_with_retry + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for( + call(messages=[{"role": "user", "content": "hello"}]), + timeout=0.01, + ) + + assert len(events) == 1 + assert events[0].finish_reason == "cancelled" + assert events[0].error_kind == "cancelled" + assert events[0].usage is None + + +@pytest.mark.asyncio +async def test_fallback_provider_propagates_observer_to_every_leaf() -> None: + primary = _NoRetryProvider( + iter( + [ + LLMResponse( + content="primary unavailable", + finish_reason="error", + error_kind="timeout", + ) + ] + ) + ) + fallback = _SequenceProvider( + iter( + [ + LLMResponse( + content="fallback ok", + usage=LLMUsage.reported(input_tokens=12, output_tokens=3), + ) + ] + ) + ) + preset = SimpleNamespace( + model="fallback-model", + max_tokens=256, + temperature=0.2, + reasoning_effort=None, + context_window_tokens=4_096, + ) + provider = FallbackProvider(primary, [preset], lambda _preset: fallback) + events: list[LLMCallRecord] = [] + provider.set_llm_call_observer(events.append) + + response = await provider.chat_with_retry( + messages=[{"role": "user", "content": "hello"}], + model="primary-model", + ) + + assert response.content == "fallback ok" + assert [(event.model, event.finish_reason) for event in events] == [ + ("primary-model", "error"), + ("fallback-model", "stop"), + ] diff --git a/tests/webui/test_settings_api.py b/tests/webui/test_settings_api.py index caab2b864..de40dea33 100644 --- a/tests/webui/test_settings_api.py +++ b/tests/webui/test_settings_api.py @@ -2,6 +2,7 @@ from __future__ import annotations import builtins import json +import time from types import SimpleNamespace import httpx @@ -9,6 +10,8 @@ import pytest from nanobot.config.loader import load_config, save_config from nanobot.config.schema import Config, InlineFallbackConfig, ModelPresetConfig +from nanobot.llm_usage import get_llm_usage_store +from nanobot.llm_usage.models import LLMCallRecord from nanobot.providers.base import LLMUsage from nanobot.providers.registry import find_by_name from nanobot.session.manager import SessionManager @@ -1463,14 +1466,16 @@ def test_settings_payload_includes_token_usage_summary( config = Config() save_config(config, config_path) monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path) - monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui") - - from nanobot.webui.token_usage import record_token_usage - - record_token_usage( - LLMUsage.reported(input_tokens=10, output_tokens=5), - timezone_name=config.agents.defaults.timezone, - ) + get_llm_usage_store().record(LLMCallRecord( + started_at_ms=int(time.time() * 1000), + duration_ms=1, + provider="openai", + model="gpt-5", + source="user", + stream=False, + finish_reason="stop", + usage=LLMUsage.reported(input_tokens=10, output_tokens=5), + )) payload = settings_payload() @@ -1491,14 +1496,16 @@ def test_settings_usage_payload_returns_lightweight_token_usage( config = Config() save_config(config, config_path) monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path) - monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui") - - from nanobot.webui.token_usage import record_token_usage - - record_token_usage( - LLMUsage.reported(input_tokens=20, output_tokens=2), - timezone_name=config.agents.defaults.timezone, - ) + get_llm_usage_store().record(LLMCallRecord( + started_at_ms=int(time.time() * 1000), + duration_ms=1, + provider="openai", + model="gpt-5", + source="user", + stream=False, + finish_reason="stop", + usage=LLMUsage.reported(input_tokens=20, output_tokens=2), + )) payload = settings_usage_payload() diff --git a/tests/webui/test_settings_routes.py b/tests/webui/test_settings_routes.py index 84d8f3ed5..3c4aafafb 100644 --- a/tests/webui/test_settings_routes.py +++ b/tests/webui/test_settings_routes.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import threading from collections.abc import Awaitable, Callable, Mapping from pathlib import Path from types import SimpleNamespace @@ -98,6 +99,45 @@ async def test_mcp_list_serializes_local_runtime_failure_snapshot(tmp_path) -> N assert snapshot_calls == 1 +@pytest.mark.asyncio +async def test_usage_query_runs_off_the_event_loop(monkeypatch) -> None: + calling_thread = threading.get_ident() + worker_threads: list[int] = [] + + def usage_payload(**_kwargs): + worker_threads.append(threading.get_ident()) + return {"days": []} + + monkeypatch.setattr("nanobot.webui.settings_routes.settings_usage_payload", usage_payload) + request = SimpleNamespace(path="/api/settings/usage", headers=Headers()) + + response = await _router().dispatch(None, request, request.path) + + assert response is not None + assert response.status_code == 200 + assert worker_threads and worker_threads[0] != calling_thread + + +@pytest.mark.asyncio +async def test_full_settings_query_runs_off_the_event_loop(monkeypatch) -> None: + calling_thread = threading.get_ident() + worker_threads: list[int] = [] + router = _router() + + def settings_response(): + worker_threads.append(threading.get_ident()) + return http_json_response({"ok": True}) + + monkeypatch.setattr(router, "_handle_settings", settings_response) + request = SimpleNamespace(path="/api/settings", headers=Headers()) + + response = await router.dispatch(None, request, request.path) + + assert response is not None + assert response.status_code == 200 + assert worker_threads and worker_threads[0] != calling_thread + + @pytest.mark.asyncio async def test_mcp_reload_callback_is_bounded( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/webui/test_token_usage.py b/tests/webui/test_token_usage.py deleted file mode 100644 index 07f2cc202..000000000 --- a/tests/webui/test_token_usage.py +++ /dev/null @@ -1,319 +0,0 @@ -from __future__ import annotations - -import json -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, -) - - -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({"schema_version": 2, "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( - LLMUsage.reported(input_tokens=1, output_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( - 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( - LLMUsage.reported(input_tokens=10, output_tokens=5), - 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", - "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, - "reported_tokens": 155, - "estimated_tokens": 0, - "requests": 2, - "reported_requests": 2, - "estimated_requests": 0, - "sources": { - "user": { - "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, - "reported_tokens": 155, - "estimated_tokens": 0, - "requests": 2, - "reported_requests": 2, - "estimated_requests": 0, - } - }, - } - ] - - -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(LLMUsage.reported(input_tokens=0, output_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( - 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]["reported_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( - LLMUsage.reported(input_tokens=100, output_tokens=25, total_tokens=175), - source="user", - now=datetime(2026, 6, 3, tzinfo=timezone.utc), - ) - record_token_usage( - LLMUsage.reported(input_tokens=20, output_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"] == 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 - - -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=LLMUsage.reported(input_tokens=20, output_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=LLMUsage.reported(input_tokens=10, output_tokens=5), - ) - ) - - payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc)) - - assert payload["days"][0]["sources"]["cron"]["total_tokens"] == 15 diff --git a/webui/src/lib/types.ts b/webui/src/lib/types.ts index 9e110f6e8..c69c85367 100644 --- a/webui/src/lib/types.ts +++ b/webui/src/lib/types.ts @@ -724,6 +724,13 @@ export interface SettingsPayload { requests: number; reported_requests?: number; estimated_requests?: number; + successful_requests?: number; + failed_requests?: number; + generation_ms?: number; + measured_output_tokens?: number; + ttft_ms?: number; + timed_requests?: number; + duration_ms?: number; sources?: Record< "user" | "api" | "cron" | "dream" | "system" | string, { @@ -739,6 +746,13 @@ export interface SettingsPayload { requests: number; reported_requests?: number; estimated_requests?: number; + successful_requests?: number; + failed_requests?: number; + generation_ms?: number; + measured_output_tokens?: number; + ttft_ms?: number; + timed_requests?: number; + duration_ms?: number; } >; }>; @@ -750,6 +764,35 @@ export interface SettingsPayload { longest_streak_days: number; active_days_30d: number; requests_30d: number; + failed_requests_30d?: number; + reported_tokens_30d?: number; + estimated_tokens_30d?: number; + cache_read_tokens_30d?: number; + cache_read_observed_input_tokens_30d?: number; + cache_read_rate_30d?: number | null; + providers_30d?: Array<{ + provider: string; + model: string; + input_tokens: number; + output_tokens: number; + cache_read_tokens: number; + cache_write_tokens: number; + cache_read_observed_input_tokens: number; + cache_write_observed_input_tokens: number; + total_tokens: number; + reported_tokens: number; + estimated_tokens: number; + requests: number; + successful_requests: number; + failed_requests: number; + reported_requests: number; + estimated_requests: number; + generation_ms: number; + measured_output_tokens: number; + ttft_ms: number; + timed_requests: number; + duration_ms: number; + }>; updated_at?: string | null; }; advanced: {