feat(usage): add unified provider usage backend

This commit is contained in:
chengyongru
2026-08-25 01:22:25 +08:00
committed by chengyongru
parent 8bb3828487
commit 2ac802b2d5
24 changed files with 1605 additions and 808 deletions
+6
View File
@@ -49,6 +49,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus from nanobot.bus.runtime_events import RuntimeEventBus
from nanobot.command import CommandContext, CommandRouter, register_builtin_commands from nanobot.command import CommandContext, CommandRouter, register_builtin_commands
from nanobot.config.schema import AgentDefaults, ModelPresetConfig 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.base import LLMProvider, LLMUsage, ProviderConversationState
from nanobot.providers.factory import ProviderSnapshot from nanobot.providers.factory import ProviderSnapshot
from nanobot.runtime_context import ( from nanobot.runtime_context import (
@@ -1202,6 +1203,11 @@ class AgentLoop:
message_metadata=metadata, message_metadata=metadata,
), ),
provider_state=provider_state, provider_state=provider_state,
llm_usage_source=source_from_request(
active_session_key,
channel=channel,
metadata=metadata,
),
)) ))
finally: finally:
turn_scope_stack.close() turn_scope_stack.close()
+11 -9
View File
@@ -20,6 +20,7 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, cast
from loguru import logger from loguru import logger
from nanobot.llm_usage.context import llm_usage_source
from nanobot.runtime_context import public_history_messages from nanobot.runtime_context import public_history_messages
from nanobot.session.manager import ( from nanobot.session.manager import (
MIN_COMPACTED_REPLAY_MESSAGES, MIN_COMPACTED_REPLAY_MESSAGES,
@@ -915,15 +916,16 @@ class Consolidator:
if not messages: if not messages:
return None return None
try: try:
response = await runtime.provider.chat_with_retry( with llm_usage_source("dream"):
model=runtime.model, response = await runtime.provider.chat_with_retry(
messages=request_messages, model=runtime.model,
tools=request_tools, messages=request_messages,
tool_choice="none", tools=request_tools,
temperature=runtime.generation.temperature, tool_choice="none",
max_tokens=runtime.generation.max_tokens, temperature=runtime.generation.temperature,
reasoning_effort=runtime.generation.reasoning_effort, max_tokens=runtime.generation.max_tokens,
) reasoning_effort=runtime.generation.reasoning_effort,
)
except Exception: except Exception:
logger.warning("Consolidation provider call failed, raw-dumping to history") logger.warning("Consolidation provider call failed, raw-dumping to history")
self.store.raw_archive(messages, session_key=session_key) self.store.raw_archive(messages, session_key=session_key)
+23 -10
View File
@@ -20,6 +20,12 @@ from nanobot.agent.context_governance import (
) )
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result 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 ( from nanobot.providers.base import (
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
@@ -118,6 +124,7 @@ class AgentRunSpec:
goal_continue_message: GoalContinueMessage | None = None goal_continue_message: GoalContinueMessage | None = None
finalize_on_max_iterations: bool = True finalize_on_max_iterations: bool = True
provider_state: ProviderConversationState | None = None provider_state: ProviderConversationState | None = None
llm_usage_source: LLMUsageSource | None = None
@dataclass(slots=True) @dataclass(slots=True)
@@ -392,6 +399,9 @@ class AgentRunner:
hook = spec.hook or AgentHook() hook = spec.hook or AgentHook()
messages = list(spec.initial_messages) messages = list(spec.initial_messages)
context = AgentRunHookContext(messages=deepcopy(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: try:
await hook.before_run(context) await hook.before_run(context)
@@ -424,17 +434,20 @@ class AgentRunner:
await hook.after_run(context) await hook.after_run(context)
return result return result
finally: finally:
context.messages = deepcopy(messages) try:
if context.exception is None: context.messages = deepcopy(messages)
await hook.on_finally(context) if context.exception is None:
else:
try:
await hook.on_finally(context) await hook.on_finally(context)
except Exception: else:
logger.exception( try:
"AgentHook.on_finally error after {}", await hook.on_finally(context)
context.stop_reason or "run exception", 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( async def _run_core(
self, self,
+9 -1
View File
@@ -8,7 +8,7 @@ import warnings
from collections.abc import Mapping from collections.abc import Mapping
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import Any, Callable, TypedDict from typing import Any, Callable, NotRequired, TypedDict
from loguru import logger from loguru import logger
@@ -28,6 +28,7 @@ from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.events import InboundMessage from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.config.schema import AgentDefaults, ToolsConfig 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.providers.base import LLMProvider, LLMUsage
from nanobot.security.workspace_access import ( from nanobot.security.workspace_access import (
WorkspaceScope, WorkspaceScope,
@@ -43,6 +44,7 @@ class _SubagentOrigin(TypedDict):
channel: str channel: str
chat_id: str chat_id: str
session_key: str | None session_key: str | None
llm_usage_source: NotRequired[LLMUsageSource]
@dataclass(slots=True) @dataclass(slots=True)
@@ -252,6 +254,7 @@ class SubagentManager:
"channel": origin_channel, "channel": origin_channel,
"chat_id": origin_chat_id, "chat_id": origin_chat_id,
"session_key": session_key, "session_key": session_key,
"llm_usage_source": current_llm_usage_source(),
} }
status = SubagentStatus( status = SubagentStatus(
@@ -315,6 +318,7 @@ class SubagentManager:
"channel": origin_channel, "channel": origin_channel,
"chat_id": origin_chat_id, "chat_id": origin_chat_id,
"session_key": session_key, "session_key": session_key,
"llm_usage_source": current_llm_usage_source(),
} }
status = SubagentStatus( status = SubagentStatus(
task_id=task_id, task_id=task_id,
@@ -417,6 +421,10 @@ class SubagentManager:
session_key=sess_key, session_key=sess_key,
workspace=root, workspace=root,
llm_timeout_s=llm_timeout, llm_timeout_s=llm_timeout,
llm_usage_source=origin.get(
"llm_usage_source",
current_llm_usage_source(),
),
)) ))
finally: finally:
if token is not None: if token is not None:
+18 -24
View File
@@ -313,6 +313,8 @@ def _run_gateway(
from nanobot.cron.service import CronJobSkippedError, CronService from nanobot.cron.service import CronJobSkippedError, CronService
from nanobot.cron.session_turns import is_bound_cron_job from nanobot.cron.session_turns import is_bound_cron_job
from nanobot.cron.types import CronJob 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 ( from nanobot.providers.factory import (
ProviderSnapshot, ProviderSnapshot,
build_provider_snapshot, build_provider_snapshot,
@@ -330,7 +332,6 @@ def _run_gateway(
) )
from nanobot.triggers.local_runner import run_local_trigger_queue from nanobot.triggers.local_runner import run_local_trigger_queue
from nanobot.triggers.local_store import LocalTriggerStore 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 port = port if port is not None else config.gateway.port
webui_url = _webui_browser_url(config) webui_url = _webui_browser_url(config)
@@ -361,7 +362,8 @@ def _run_gateway(
runtime_events = RuntimeEventBus() runtime_events = RuntimeEventBus()
fallback_model_observer = build_webui_fallback_model_observer(bus) 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): if isinstance(snapshot.provider, FallbackProvider):
snapshot.provider.set_fallback_model_observer(fallback_model_observer) snapshot.provider.set_fallback_model_observer(fallback_model_observer)
return snapshot return snapshot
@@ -371,20 +373,19 @@ def _run_gateway(
**kwargs: Any, **kwargs: Any,
) -> ProviderSnapshot: ) -> ProviderSnapshot:
try: try:
return _observe_fallback_models(load_provider_snapshot(*args, **kwargs)) return _observe_provider(load_provider_snapshot(*args, **kwargs))
except ValueError as exc: except ValueError as exc:
if unconfigured_provider_error is None: if unconfigured_provider_error is None:
raise 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: if unconfigured_provider_error is not None:
provider_snapshot = build_unconfigured_provider_snapshot( provider_snapshot = _observe_provider(
config, build_unconfigured_provider_snapshot(config, unconfigured_provider_error)
unconfigured_provider_error,
) )
else: else:
try: try:
provider_snapshot = _observe_fallback_models(build_provider_snapshot(config)) provider_snapshot = _observe_provider(build_provider_snapshot(config))
except ValueError as exc: except ValueError as exc:
console.print(f"[red]Error: {exc}[/red]") console.print(f"[red]Error: {exc}[/red]")
raise typer.Exit(1) from exc raise typer.Exit(1) from exc
@@ -443,7 +444,6 @@ def _run_gateway(
runtime_events=runtime_events, runtime_events=runtime_events,
turn_delivery_factory=turn_delivery_factory, turn_delivery_factory=turn_delivery_factory,
provider_signature=provider_snapshot.signature, provider_signature=provider_snapshot.signature,
hooks=[TokenUsageHook(timezone_name=config.agents.defaults.timezone)],
local_trigger_store=trigger_store, local_trigger_store=trigger_store,
hook_factories=[create_file_edit_activity_hook], hook_factories=[create_file_edit_activity_hook],
tool_registry=tools, tool_registry=tools,
@@ -564,13 +564,6 @@ def _run_gateway(
except Exception: except Exception:
logger.exception("Dream cron job failed") logger.exception("Dream cron job failed")
finally: 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) sha = _commit_dream_changes(store)
if sha: if sha:
logger.info("Dream commit: {}", sha) logger.info("Dream commit: {}", sha)
@@ -630,14 +623,15 @@ def _run_gateway(
evaluator_prompt = resolve_evaluator_prompt(config.workspace_path) evaluator_prompt = resolve_evaluator_prompt(config.workspace_path)
# Fail closed: stay silent on evaluator failure instead of notifying. # Fail closed: stay silent on evaluator failure instead of notifying.
should_notify = await evaluate_response( with llm_usage_source("cron"):
response=response, should_notify = await evaluate_response(
task_context=prompt, response=response,
provider=agent.provider, task_context=prompt,
model=agent.model, provider=agent.provider,
evaluator_prompt=evaluator_prompt, model=agent.model,
default_notify=False, evaluator_prompt=evaluator_prompt,
) default_notify=False,
)
if should_notify: if should_notify:
logger.info("Heartbeat: completed, delivering response") logger.info("Heartbeat: completed, delivering response")
-7
View File
@@ -479,13 +479,6 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage:
elapsed = time.monotonic() - t0 elapsed = time.monotonic() - t0
content = f"Dream failed after {elapsed:.1f}s: {e}" content = f"Dream failed after {elapsed:.1f}s: {e}"
finally: 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(): if store.git.is_initialized():
commit_msg = build_dream_commit_message("dream: manual run", diff_body) commit_msg = build_dream_commit_message("dream: manual run", diff_body)
sha = store.git.auto_commit(commit_msg) sha = store.git.auto_commit(commit_msg)
+86
View File
@@ -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",
]
+70
View File
@@ -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)
+38
View File
@@ -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")
+560
View File
@@ -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)
+143 -7
View File
@@ -6,6 +6,7 @@ import asyncio
import json import json
import os import os
import re import re
import time
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
from contextlib import suppress from contextlib import suppress
@@ -13,19 +14,23 @@ from copy import deepcopy
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime, timezone from datetime import datetime, timezone
from email.utils import parsedate_to_datetime from email.utils import parsedate_to_datetime
from typing import Any, Literal, cast from typing import TYPE_CHECKING, Any, Literal, cast
import json_repair import json_repair
from loguru import logger from loguru import logger
from nanobot.utils.helpers import sanitize_surrogates_deep 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" STREAM_IDLE_TIMEOUT_ENV = "NANOBOT_STREAM_IDLE_TIMEOUT_S"
DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0 DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0
MAX_STREAM_IDLE_TIMEOUT_S = 3600.0 MAX_STREAM_IDLE_TIMEOUT_S = 3600.0
RETRY_AFTER_BUFFER = 1 RETRY_AFTER_BUFFER = 1
RetryEventCallback = Callable[[str], Awaitable[None]] RetryEventCallback = Callable[[str], Awaitable[None]]
LLMCallObserver = Callable[["LLMCallRecord"], None]
def resolve_stream_idle_timeout_s( def resolve_stream_idle_timeout_s(
@@ -682,6 +687,95 @@ class LLMProvider(ABC):
self.api_base = api_base self.api_base = api_base
self.provider_name = provider_name self.provider_name = provider_name
self.generation: GenerationSettings = GenerationSettings() 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( def can_resume_conversation_state(
self, self,
@@ -1068,18 +1162,39 @@ class LLMProvider(ABC):
async def _safe_chat(self, **kwargs: Any) -> LLMResponse: async def _safe_chat(self, **kwargs: Any) -> LLMResponse:
"""Call chat() and convert unexpected exceptions to error responses.""" """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: try:
provider_context = kwargs.pop("provider_context", None) provider_context = kwargs.pop("provider_context", None)
if isinstance(provider_context, ProviderCallContext): if isinstance(provider_context, ProviderCallContext):
return await self.chat_with_context( response = await self.chat_with_context(
provider_context=provider_context, provider_context=provider_context,
**kwargs, **kwargs,
) )
return await self.chat(**kwargs) else:
response = await self.chat(**kwargs)
except asyncio.CancelledError: 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 raise
except Exception as exc: 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( async def chat_stream(
self, self,
@@ -1142,18 +1257,39 @@ class LLMProvider(ABC):
async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse: async def _safe_chat_stream(self, **kwargs: Any) -> LLMResponse:
"""Call chat_stream() and convert unexpected exceptions to error responses.""" """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: try:
provider_context = kwargs.pop("provider_context", None) provider_context = kwargs.pop("provider_context", None)
if isinstance(provider_context, ProviderCallContext): if isinstance(provider_context, ProviderCallContext):
return await self.chat_stream_with_context( response = await self.chat_stream_with_context(
provider_context=provider_context, provider_context=provider_context,
**kwargs, **kwargs,
) )
return await self.chat_stream(**kwargs) else:
response = await self.chat_stream(**kwargs)
except asyncio.CancelledError: 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 raise
except Exception as exc: 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( async def chat_stream_with_retry(
self, self,
+7
View File
@@ -13,6 +13,7 @@ from loguru import logger
from nanobot.providers.base import ( from nanobot.providers.base import (
GenerationSettings, GenerationSettings,
LLMCallObserver,
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
ProviderCallContext, ProviderCallContext,
@@ -151,6 +152,11 @@ class FallbackProvider(LLMProvider):
"""Attach a process-level observer without changing request call signatures.""" """Attach a process-level observer without changing request call signatures."""
self._fallback_model_observer = observer 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 @property
def supports_progress_deltas(self) -> bool: def supports_progress_deltas(self) -> bool:
return bool(getattr(self._primary, "supports_progress_deltas", False)) return bool(getattr(self._primary, "supports_progress_deltas", False))
@@ -506,6 +512,7 @@ class FallbackProvider(LLMProvider):
) )
try: try:
fallback_provider = self._provider_factory(fallback) fallback_provider = self._provider_factory(fallback)
fallback_provider.set_llm_call_observer(self._llm_call_observer)
except Exception as exc: except Exception as exc:
logger.warning( logger.warning(
"Failed to create provider for fallback '{}': {}", fallback_model, exc "Failed to create provider for fallback '{}': {}", fallback_model, exc
+20 -18
View File
@@ -37,6 +37,7 @@ from nanobot.bus.runtime_events import (
TurnRuntimeAdmitted, TurnRuntimeAdmitted,
UserInputAccepted, UserInputAccepted,
) )
from nanobot.llm_usage.context import llm_usage_source
from nanobot.providers.base import LLMProvider, LLMUsage from nanobot.providers.base import LLMProvider, LLMUsage
from nanobot.providers.fallback_provider import FallbackModelObserver from nanobot.providers.fallback_provider import FallbackModelObserver
from nanobot.runtime_context import public_history_message 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)}" prompt += f"\nAssistant: {truncate_text(assistant_text, 1_000)}"
try: try:
response = await provider.chat_with_retry( with llm_usage_source("system"):
[ response = await provider.chat_with_retry(
{ [
"role": "system", {
"content": ( "role": "system",
"You write short, neutral chat titles. " "content": (
"Return only the title text." "You write short, neutral chat titles. "
), "Return only the title text."
}, ),
{"role": "user", "content": prompt}, },
], {"role": "user", "content": prompt},
tools=None, ],
model=model, tools=None,
max_tokens=TITLE_GENERATION_MAX_TOKENS, model=model,
temperature=0.2, max_tokens=TITLE_GENERATION_MAX_TOKENS,
reasoning_effort=TITLE_GENERATION_REASONING_EFFORT, temperature=0.2,
retry_mode="standard", reasoning_effort=TITLE_GENERATION_REASONING_EFFORT,
) retry_mode="standard",
)
except Exception: except Exception:
logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True) logger.debug("Failed to generate webui session title for {}", session_key, exc_info=True)
return False return False
+2 -2
View File
@@ -284,9 +284,9 @@ class WebUISettingsRouter:
if not self._authorized(request): if not self._authorized(request):
return self._unauthorized() return self._unauthorized()
if route == ("root", "settings"): if route == ("root", "settings"):
return self._handle_settings() return await asyncio.to_thread(self._handle_settings)
if route == ("root", "usage"): if route == ("root", "usage"):
return self._handle_settings_usage() return await asyncio.to_thread(self._handle_settings_usage)
domain, action = route domain, action = route
domain_request = self._domain_request( domain_request = self._domain_request(
+3 -3
View File
@@ -20,6 +20,7 @@ from nanobot.channels.contracts import (
channel_update_instance_config, channel_update_instance_config,
) )
from nanobot.config.schema import 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.optional_features import OptionalFeatureError, with_channel_runtime_status
from nanobot.security.workspace_access import workspace_sandbox_status from nanobot.security.workspace_access import workspace_sandbox_status
from nanobot.webui.settings_capabilities import network_safety_payload from nanobot.webui.settings_capabilities import network_safety_payload
@@ -31,7 +32,6 @@ from nanobot.webui.settings_contracts import (
query_first, query_first,
query_first_alias, query_first_alias,
) )
from nanobot.webui.token_usage import token_usage_payload
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.webui.settings_services import WebUISettingsServices from nanobot.webui.settings_services import WebUISettingsServices
@@ -121,7 +121,7 @@ def system_settings_payload(
}, },
"unified_session": defaults.unified_session, "unified_session": defaults.unified_session,
}, },
"usage": token_usage_payload(timezone_name=defaults.timezone), "usage": llm_usage_payload(timezone_name=defaults.timezone),
"advanced": { "advanced": {
"restrict_to_workspace": config.tools.restrict_to_workspace, "restrict_to_workspace": config.tools.restrict_to_workspace,
"workspace_sandbox": sandbox_status.as_dict(), "workspace_sandbox": sandbox_status.as_dict(),
@@ -139,7 +139,7 @@ def system_settings_payload(
def settings_usage_payload(config: Config) -> dict[str, Any]: def settings_usage_payload(config: Config) -> dict[str, Any]:
"""Return the lightweight token usage slice for Overview refreshes.""" """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]: def update_agent_system_settings(config: Config, query: QueryParams) -> tuple[bool, bool]:
-392
View File
@@ -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")
+29
View File
@@ -1,5 +1,6 @@
"""Tests for SubagentManager.""" """Tests for SubagentManager."""
import asyncio
from pathlib import Path from pathlib import Path
from unittest.mock import AsyncMock, MagicMock 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.agent.tools.filesystem import FileToolsConfig
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.config.schema import ToolsConfig from nanobot.config.schema import ToolsConfig
from nanobot.llm_usage.context import llm_usage_source
from nanobot.providers.base import GenerationSettings, LLMProvider from nanobot.providers.base import GenerationSettings, LLMProvider
from nanobot.security.workspace_access import build_workspace_scope from nanobot.security.workspace_access import build_workspace_scope
from nanobot.utils.llm_runtime import LLMRuntime 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] spec = sm.runner.run.call_args.args[0]
assert spec.fail_on_tool_error is False 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"
+19
View File
@@ -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"
+275
View File
@@ -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()
+180
View File
@@ -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"),
]
+23 -16
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import builtins import builtins
import json import json
import time
from types import SimpleNamespace from types import SimpleNamespace
import httpx import httpx
@@ -9,6 +10,8 @@ import pytest
from nanobot.config.loader import load_config, save_config from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import Config, InlineFallbackConfig, ModelPresetConfig 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.base import LLMUsage
from nanobot.providers.registry import find_by_name from nanobot.providers.registry import find_by_name
from nanobot.session.manager import SessionManager from nanobot.session.manager import SessionManager
@@ -1463,14 +1466,16 @@ def test_settings_payload_includes_token_usage_summary(
config = Config() config = Config()
save_config(config, config_path) save_config(config, config_path)
monkeypatch.setattr("nanobot.config.loader._current_config_path", 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") get_llm_usage_store().record(LLMCallRecord(
started_at_ms=int(time.time() * 1000),
from nanobot.webui.token_usage import record_token_usage duration_ms=1,
provider="openai",
record_token_usage( model="gpt-5",
LLMUsage.reported(input_tokens=10, output_tokens=5), source="user",
timezone_name=config.agents.defaults.timezone, stream=False,
) finish_reason="stop",
usage=LLMUsage.reported(input_tokens=10, output_tokens=5),
))
payload = settings_payload() payload = settings_payload()
@@ -1491,14 +1496,16 @@ def test_settings_usage_payload_returns_lightweight_token_usage(
config = Config() config = Config()
save_config(config, config_path) save_config(config, config_path)
monkeypatch.setattr("nanobot.config.loader._current_config_path", 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") get_llm_usage_store().record(LLMCallRecord(
started_at_ms=int(time.time() * 1000),
from nanobot.webui.token_usage import record_token_usage duration_ms=1,
provider="openai",
record_token_usage( model="gpt-5",
LLMUsage.reported(input_tokens=20, output_tokens=2), source="user",
timezone_name=config.agents.defaults.timezone, stream=False,
) finish_reason="stop",
usage=LLMUsage.reported(input_tokens=20, output_tokens=2),
))
payload = settings_usage_payload() payload = settings_usage_payload()
+40
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio import asyncio
import json import json
import threading
from collections.abc import Awaitable, Callable, Mapping from collections.abc import Awaitable, Callable, Mapping
from pathlib import Path from pathlib import Path
from types import SimpleNamespace 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 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 @pytest.mark.asyncio
async def test_mcp_reload_callback_is_bounded( async def test_mcp_reload_callback_is_bounded(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
-319
View File
@@ -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
+43
View File
@@ -724,6 +724,13 @@ export interface SettingsPayload {
requests: number; requests: number;
reported_requests?: number; reported_requests?: number;
estimated_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< sources?: Record<
"user" | "api" | "cron" | "dream" | "system" | string, "user" | "api" | "cron" | "dream" | "system" | string,
{ {
@@ -739,6 +746,13 @@ export interface SettingsPayload {
requests: number; requests: number;
reported_requests?: number; reported_requests?: number;
estimated_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; longest_streak_days: number;
active_days_30d: number; active_days_30d: number;
requests_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; updated_at?: string | null;
}; };
advanced: { advanced: {