mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-08 05:18:49 +03:00
Code Issues Packages Projects Releases Wiki Activity

feat(webui): track turns and usage in gateway

Browse Source
This commit is contained in:
Xubin Ren
2026-06-06 00:19:31 +08:00
parent 8c9c915df6
commit 8e7af82338
23 changed files with 1492 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/hook.py
+1
View File
@@ -26,6 +26,7 @@ class AgentHookContext:
final_content: str | None = None
stop_reason: str | None = None
error: str | None = None
session_key: str | None = None
@dataclass(slots=True)
nanobot/agent/loop.py
+36 -2
View File
@@ -658,6 +658,31 @@ class AgentLoop:
budget = self.context_window_tokens - max(1, reserved_output) - 1024
return budget if budget > 0 else max(128, self.context_window_tokens // 2)
@staticmethod
def _hook_includes_ephemeral(hook: AgentHook) -> bool:
try:
return hook.include_ephemeral() is True
except Exception:
return False
@staticmethod
def _usage_source_for_turn(
*,
channel: str,
session_key: str | None,
ephemeral: bool,
) -> str:
key = session_key or ""
if key.startswith("dream:") or (ephemeral and key.startswith("dream")):
return "dream"
if key == "heartbeat" or key.startswith("cron:"):
return "cron"
if channel == "api" or key.startswith("api:"):
return "api"
if channel == "system":
return "system"
return "user"
async def _run_agent_loop(
self,
initial_messages: list[dict],
@@ -701,8 +726,12 @@ class AgentLoop:
on_iteration=lambda iteration: setattr(self, "_current_iteration", iteration),
)
hook: AgentHook = loop_hook
if not ephemeral and self._extra_hooks:
hook = CompositeHook([loop_hook] + self._extra_hooks)
extra_hooks = [
h for h in self._extra_hooks
if not ephemeral or self._hook_includes_ephemeral(h)
]
if extra_hooks:
hook = CompositeHook([loop_hook] + extra_hooks)
async def _checkpoint(payload: dict[str, Any]) -> None:
if session is None:
@@ -816,6 +845,11 @@ class AgentLoop:
),
goal_active_predicate=lambda: sustained_goal_active(session.metadata) if session is not None else False,
goal_continue_message=_goal_continue,
usage_source=self._usage_source_for_turn(
channel=channel,
session_key=active_session_key,
ephemeral=ephemeral,
),
))
finally:
reset_workspace_scope(workspace_token)
nanobot/agent/runner.py
+67 -7
View File
@@ -363,14 +363,15 @@ class AgentRunner:
messages_for_model = self._backfill_missing_tool_results(messages_for_model)
except Exception:
messages_for_model = messages
context = AgentHookContext(iteration=iteration, messages=messages)
context = AgentHookContext(
iteration=iteration,
messages=messages,
session_key=spec.session_key,
)
await hook.before_iteration(context)
response = await self._request_model(spec, messages_for_model, hook, context)
raw_usage = self._usage_dict(response.usage)
context.response = response
context.usage = dict(raw_usage)
context.tool_calls = list(response.tool_calls)
self._accumulate_usage(usage, raw_usage)
reasoning_text, cleaned_content = extract_reasoning(
response.reasoning_content,
@@ -378,6 +379,9 @@ class AgentRunner:
response.content,
)
response.content = cleaned_content
raw_usage = self._usage_or_estimate(spec, messages_for_model, response)
context.usage = dict(raw_usage)
self._accumulate_usage(usage, raw_usage)
if reasoning_text and not context.streamed_reasoning:
await hook.emit_reasoning(reasoning_text)
await hook.emit_reasoning_end()
@@ -504,8 +508,9 @@ class AgentRunner:
)
if hook.wants_streaming():
await hook.on_stream_end(context, resuming=False)
retry_messages = self._finalization_retry_messages(messages_for_model)
response = await self._request_finalization_retry(spec, messages_for_model)
retry_usage = self._usage_dict(response.usage)
retry_usage = self._usage_or_estimate(spec, retry_messages, response)
self._accumulate_usage(usage, retry_usage)
raw_usage = self._merge_usage(raw_usage, retry_usage)
context.response = response
@@ -821,11 +826,60 @@ class AgentRunner:
spec: AgentRunSpec,
messages: list[dict[str, Any]],
):
retry_messages = list(messages)
retry_messages.append(build_finalization_retry_message())
retry_messages = self._finalization_retry_messages(messages)
kwargs = self._build_request_kwargs(spec, retry_messages, tools=None)
return await self.provider.chat_with_retry(**kwargs)
@staticmethod
def _finalization_retry_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
retry_messages = list(messages)
retry_messages.append(build_finalization_retry_message())
return retry_messages
def _usage_or_estimate(
self,
spec: AgentRunSpec,
messages: list[dict[str, Any]],
response: LLMResponse,
) -> dict[str, int]:
usage = self._usage_dict(response.usage)
total = self._usage_total(usage)
if total > 0:
usage["total_tokens"] = total
usage.setdefault("provider_tokens", total)
return usage
if response.finish_reason == "error":
return {}
return self._estimate_response_usage(spec, messages, response)
def _estimate_response_usage(
self,
spec: AgentRunSpec,
messages: list[dict[str, Any]],
response: LLMResponse,
) -> dict[str, int]:
try:
tools = spec.tools.get_definitions()
except Exception:
tools = None
prompt_tokens, _ = estimate_prompt_tokens_chain(self.provider, spec.model, messages, tools)
assistant_message = build_assistant_message(
response.content or "",
tool_calls=[tc.to_openai_tool_call() for tc in response.tool_calls],
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
)
completion_tokens = estimate_message_tokens(assistant_message)
total_tokens = max(0, prompt_tokens) + max(0, completion_tokens)
if total_tokens <= 0:
return {}
return {
"prompt_tokens": max(0, prompt_tokens),
"completion_tokens": max(0, completion_tokens),
"total_tokens": total_tokens,
"estimated_tokens": total_tokens,
}
@staticmethod
def _usage_dict(usage: dict[str, Any] | None) -> dict[str, int]:
if not usage:
@@ -838,6 +892,12 @@ class AgentRunner:
continue
return result
@staticmethod
def _usage_total(usage: dict[str, int]) -> int:
return max(0, usage.get("total_tokens", 0) or (
usage.get("prompt_tokens", 0) + usage.get("completion_tokens", 0)
))
@staticmethod
def _accumulate_usage(target: dict[str, int], addition: dict[str, int]) -> None:
for key, value in addition.items():
nanobot/channels/websocket.py
+60 -8
View File
@@ -35,9 +35,6 @@ from nanobot.utils.media_decode import (
)
from nanobot.webui.cli_apps_api import normalize_cli_app_mentions
from nanobot.webui.gateway_services import GatewayServices
from nanobot.webui.http_utils import (
is_localhost as _is_localhost,
)
from nanobot.webui.http_utils import (
normalize_config_path as _normalize_config_path,
)
@@ -240,6 +237,8 @@ _VIDEO_MIME_ALLOWED: frozenset[str] = frozenset({
_UPLOAD_MIME_ALLOWED: frozenset[str] = _IMAGE_MIME_ALLOWED | _VIDEO_MIME_ALLOWED
_DATA_URL_MIME_RE = re.compile(r"^data:([^;]+);base64,", re.DOTALL)
_WEBUI_TURN_ID_RE = re.compile(r"^[A-Za-z0-9._:-]{1,128}$")
_WEBUI_TURN_META_KEY = "webui_turn_id"
def _extract_data_url_mime(url: str) -> str | None:
@@ -263,6 +262,14 @@ def _is_websocket_upgrade(request: WsRequest) -> bool:
return True
def _normalize_webui_turn_id(value: Any) -> str:
if isinstance(value, str):
candidate = value.strip()
if _WEBUI_TURN_ID_RE.fullmatch(candidate):
return candidate
return str(uuid.uuid4())
class WebSocketChannel(BaseChannel):
"""Run a local WebSocket server; forward text/JSON messages to the message bus."""
@@ -296,9 +303,13 @@ class WebSocketChannel(BaseChannel):
self._workspaces = gateway.workspaces
self._stream_text_buffers: dict[tuple[str, str], list[str]] = {}
self._webui_turn_sequences: dict[tuple[str, str], int] = {}
# -- Subscription bookkeeping -------------------------------------------
def _workspace_controls_available(self, connection: Any) -> bool:
return self._http_router.workspace_controls_available(connection)
def _attach(self, connection: Any, chat_id: str) -> None:
"""Idempotently subscribe *connection* to *chat_id*."""
self._subs.setdefault(chat_id, set()).add(connection)
@@ -651,7 +662,7 @@ class WebSocketChannel(BaseChannel):
connection,
lambda: self._workspaces.scope_for_new_chat(
envelope,
controls_available=_is_localhost(connection),
controls_available=self._workspace_controls_available(connection),
),
)
if scope is None:
@@ -688,7 +699,7 @@ class WebSocketChannel(BaseChannel):
envelope,
chat_id=cid,
chat_running=websocket_turn_wall_started_at(cid) is not None,
controls_available=_is_localhost(connection),
controls_available=self._workspace_controls_available(connection),
),
chat_id=cid,
)
@@ -740,7 +751,7 @@ class WebSocketChannel(BaseChannel):
envelope,
chat_id=cid,
chat_running=websocket_turn_wall_started_at(cid) is not None,
controls_available=_is_localhost(connection),
controls_available=self._workspace_controls_available(connection),
),
chat_id=cid,
)
@@ -751,6 +762,7 @@ class WebSocketChannel(BaseChannel):
self._attach(connection, cid)
await self._hydrate_after_subscribe(cid)
metadata: dict[str, Any] = {"remote": getattr(connection, "remote_address", None)}
metadata[_WEBUI_TURN_META_KEY] = _normalize_webui_turn_id(envelope.get("turn_id"))
if envelope.get("webui") is True:
metadata["webui"] = True
cli_apps = normalize_cli_app_mentions(envelope.get("cli_apps"))
@@ -768,10 +780,11 @@ class WebSocketChannel(BaseChannel):
"enabled": True,
"aspect_ratio": aspect_ratio if isinstance(aspect_ratio, str) else None,
}
if envelope.get("webui") is True and self.is_allowed(client_id):
if metadata.get("webui") is True and self.is_allowed(client_id):
self._try_append_webui_user_transcript(
cid,
content,
metadata=metadata,
media_paths=media_paths,
cli_apps=cli_apps,
mcp_presets=mcp_presets,
@@ -849,6 +862,7 @@ class WebSocketChannel(BaseChannel):
chat_id: str,
content: str,
*,
metadata: dict[str, Any] | None,
media_paths: list[str],
cli_apps: list[dict[str, Any]],
mcp_presets: list[dict[str, Any]],
@@ -864,8 +878,30 @@ class WebSocketChannel(BaseChannel):
)
if payload is None:
return
self._annotate_webui_turn(payload, chat_id, metadata, "user")
self._try_append_webui_transcript(chat_id, payload)
def _next_webui_turn_seq(self, chat_id: str, turn_id: str) -> int:
key = (chat_id, turn_id)
seq = self._webui_turn_sequences.get(key, 0) + 1
self._webui_turn_sequences[key] = seq
return seq
def _annotate_webui_turn(
self,
payload: dict[str, Any],
chat_id: str,
metadata: dict[str, Any] | None,
phase: str,
) -> None:
meta = metadata or {}
turn_id = meta.get(_WEBUI_TURN_META_KEY)
if not isinstance(turn_id, str) or not turn_id:
return
payload["turn_id"] = turn_id
payload["turn_phase"] = phase
payload["turn_seq"] = self._next_webui_turn_seq(chat_id, turn_id)
async def send(self, msg: OutboundMessage) -> None:
if msg.metadata.get("_runtime_model_updated"):
await self.send_runtime_model_updated(
@@ -909,7 +945,12 @@ class WebSocketChannel(BaseChannel):
lat_i = int(lat) if isinstance(lat, (int, float)) else None
gs = msg.metadata.get("goal_state")
gs_blob = gs if isinstance(gs, dict) else None
await self.send_turn_end(msg.chat_id, latency_ms=lat_i, goal_state=gs_blob)
await self.send_turn_end(
msg.chat_id,
latency_ms=lat_i,
goal_state=gs_blob,
metadata=msg.metadata,
)
return
if msg.metadata.get("_session_updated"):
scope = msg.metadata.get("_session_update_scope")
@@ -959,6 +1000,8 @@ class WebSocketChannel(BaseChannel):
payload["kind"] = "tool_hint"
elif msg.metadata.get("_progress"):
payload["kind"] = "progress"
phase = "activity" if payload.get("kind") in ("tool_hint", "progress") else "answer"
self._annotate_webui_turn(payload, msg.chat_id, msg.metadata, phase)
transcript_payload = dict(payload)
transcript_payload["text"] = text
self._try_append_webui_transcript(msg.chat_id, transcript_payload)
@@ -989,6 +1032,7 @@ class WebSocketChannel(BaseChannel):
stream_id = meta.get("_stream_id")
if stream_id is not None:
body["stream_id"] = stream_id
self._annotate_webui_turn(body, chat_id, meta, "reasoning")
self._try_append_webui_transcript(chat_id, body)
raw = json.dumps(body, ensure_ascii=False)
for connection in conns:
@@ -1011,6 +1055,7 @@ class WebSocketChannel(BaseChannel):
stream_id = meta.get("_stream_id")
if stream_id is not None:
body["stream_id"] = stream_id
self._annotate_webui_turn(body, chat_id, meta, "reasoning")
self._try_append_webui_transcript(chat_id, body)
raw = json.dumps(body, ensure_ascii=False)
for connection in conns:
@@ -1030,6 +1075,7 @@ class WebSocketChannel(BaseChannel):
"chat_id": chat_id,
"edits": edits,
}
self._annotate_webui_turn(payload, chat_id, metadata, "activity")
self._try_append_webui_transcript(chat_id, payload)
raw = json.dumps(payload, ensure_ascii=False)
for connection in conns:
@@ -1064,6 +1110,7 @@ class WebSocketChannel(BaseChannel):
self._stream_text_buffers.setdefault(stream_key, []).append(delta)
if meta.get("_stream_id") is not None:
body["stream_id"] = meta["_stream_id"]
self._annotate_webui_turn(body, chat_id, meta, "answer")
self._try_append_webui_transcript(chat_id, body)
raw = json.dumps(body, ensure_ascii=False)
for connection in conns:
@@ -1075,6 +1122,7 @@ class WebSocketChannel(BaseChannel):
latency_ms: int | None = None,
*,
goal_state: dict[str, Any] | None = None,
metadata: dict[str, Any] | None = None,
) -> None:
"""Signal that the agent has fully finished processing the current turn."""
conns = list(self._subs.get(chat_id, ()))
@@ -1085,10 +1133,14 @@ class WebSocketChannel(BaseChannel):
body["latency_ms"] = int(latency_ms)
if goal_state is not None:
body["goal_state"] = goal_state
self._annotate_webui_turn(body, chat_id, metadata, "complete")
self._try_append_webui_transcript(chat_id, body)
raw = json.dumps(body, ensure_ascii=False)
for connection in conns:
await self._safe_send_to(connection, raw, label=" turn_end ")
turn_id = body.get("turn_id")
if isinstance(turn_id, str):
self._webui_turn_sequences.pop((chat_id, turn_id), None)
async def send_goal_state(self, chat_id: str, blob: dict[str, Any]) -> None:
"""Push persisted goal-state snapshot for *chat_id* (multi-chat isolation)."""
nanobot/cli/commands.py
+69 -5
View File
@@ -738,6 +738,59 @@ def gateway(
_run_gateway(cfg, port=port)
DESKTOP_BOOTSTRAP_PROVIDER = "openai_codex"
DESKTOP_BOOTSTRAP_MODEL = "openai-codex/gpt-5.1-codex"
def _desktop_provider_error_is_recoverable(error: ValueError) -> bool:
message = str(error)
return "No API key configured" in message or "requires api_key and api_base" in message
def _desktop_provider_needs_bootstrap(config: Config) -> bool:
from nanobot.providers.factory import make_provider
try:
make_provider(config)
return False
except ValueError as e:
if not _desktop_provider_error_is_recoverable(e):
raise
return True
def _reset_desktop_config_to_unconfigured(config: Config) -> bool:
defaults = config.agents.defaults
changed = False
if defaults.model_preset is not None:
defaults.model_preset = None
changed = True
if defaults.provider:
defaults.provider = ""
changed = True
if defaults.model:
defaults.model = ""
changed = True
return changed
def _is_persisted_desktop_bootstrap(config: Config) -> bool:
defaults = config.agents.defaults
return (
defaults.model_preset is None
and defaults.provider == DESKTOP_BOOTSTRAP_PROVIDER
and defaults.model == DESKTOP_BOOTSTRAP_MODEL
and not config.model_presets
)
def _apply_desktop_runtime_bootstrap(config: Config) -> None:
defaults = config.agents.defaults
config.agents.defaults.model_preset = None
defaults.provider = DESKTOP_BOOTSTRAP_PROVIDER
defaults.model = DESKTOP_BOOTSTRAP_MODEL
def _load_or_create_desktop_config(config: str | None, workspace: str | None) -> Config:
"""Load the desktop-owned config, creating it on first launch."""
from nanobot.config.loader import (
@@ -751,7 +804,7 @@ def _load_or_create_desktop_config(config: str | None, workspace: str | None) ->
config_path = Path(config).expanduser().resolve() if config else get_config_path()
set_config_path(config_path)
created = False
changed = False
if config_path.exists():
try:
loaded = resolve_config_env_vars(load_config(config_path))
@@ -760,16 +813,25 @@ def _load_or_create_desktop_config(config: str | None, workspace: str | None) ->
raise typer.Exit(1)
else:
loaded = NanobotConfig()
created = True
changed = True
if workspace:
workspace_path = Path(workspace).expanduser()
loaded.agents.defaults.workspace = str(workspace_path)
created = True
changed = True
if created:
if _is_persisted_desktop_bootstrap(loaded):
changed = _reset_desktop_config_to_unconfigured(loaded) or changed
elif _desktop_provider_needs_bootstrap(loaded):
changed = _reset_desktop_config_to_unconfigured(loaded) or changed
if changed:
save_config(loaded, config_path)
return loaded
runtime_config = loaded.model_copy(deep=True)
if _desktop_provider_needs_bootstrap(runtime_config):
_apply_desktop_runtime_bootstrap(runtime_config)
return runtime_config
def _configure_desktop_gateway(
@@ -889,6 +951,7 @@ def _run_gateway(
from nanobot.providers.image_generation import image_gen_provider_configs
from nanobot.session.manager import SessionManager
from nanobot.session.webui_turns import WebuiTurnCoordinator
from nanobot.webui.token_usage import TokenUsageHook
port = port if port is not None else config.gateway.port
@@ -923,6 +986,7 @@ def _run_gateway(
provider_snapshot_loader=load_provider_snapshot,
runtime_events=runtime_events,
provider_signature=provider_snapshot.signature,
hooks=[TokenUsageHook(timezone_name=config.agents.defaults.timezone)],
)
WebuiTurnCoordinator(
bus=bus,
nanobot/providers/openai_codex_provider.py
+4 -3
View File
@@ -71,7 +71,7 @@ class OpenAICodexProvider(LLMProvider):
try:
try:
content, tool_calls, finish_reason, reasoning_content = await _request_codex(
content, tool_calls, finish_reason, usage, reasoning_content = await _request_codex(
DEFAULT_CODEX_URL, headers, body, verify=True,
on_content_delta=on_content_delta,
on_thinking_delta=on_thinking_delta,
@@ -81,7 +81,7 @@ class OpenAICodexProvider(LLMProvider):
if "CERTIFICATE_VERIFY_FAILED" not in str(e):
raise
logger.warning("SSL verification failed for Codex API; retrying with verify=False")
content, tool_calls, finish_reason, reasoning_content = await _request_codex(
content, tool_calls, finish_reason, usage, reasoning_content = await _request_codex(
DEFAULT_CODEX_URL, headers, body, verify=False,
on_content_delta=on_content_delta,
on_thinking_delta=on_thinking_delta,
@@ -91,6 +91,7 @@ class OpenAICodexProvider(LLMProvider):
content=content,
tool_calls=tool_calls,
finish_reason=finish_reason,
usage=usage,
reasoning_content=reasoning_content,
)
except Exception as e:
@@ -197,7 +198,7 @@ async def _request_codex(
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
) -> tuple[str, list[ToolCallRequest], str, str | None]:
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
idle_timeout_s = int(os.environ.get("NANOBOT_STREAM_IDLE_TIMEOUT_S", "90"))
async with httpx.AsyncClient(timeout=idle_timeout_s, verify=verify) as client:
async with client.stream("POST", url, headers=headers, json=body) as response:
nanobot/providers/openai_responses/parsing.py
+26 -15
View File
@@ -25,12 +25,31 @@ def map_finish_reason(status: str | None) -> str:
return FINISH_REASON_MAP.get(status or "completed", "stop")
def _usage_from_response_obj(response: Any) -> dict[str, int]:
usage_raw = response.get("usage") if isinstance(response, dict) else getattr(response, "usage", None)
if not usage_raw:
return {}
if not isinstance(usage_raw, dict):
dump = getattr(usage_raw, "model_dump", None)
usage_raw = dump() if callable(dump) else vars(usage_raw)
prompt_tokens = int(usage_raw.get("input_tokens") or usage_raw.get("prompt_tokens") or 0)
completion_tokens = int(
usage_raw.get("output_tokens") or usage_raw.get("completion_tokens") or 0
)
total_tokens = int(usage_raw.get("total_tokens") or prompt_tokens + completion_tokens)
return {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": total_tokens,
}
async def iter_sse(response: httpx.Response) -> AsyncGenerator[dict[str, Any], None]:
"""Yield parsed JSON events from a Responses API SSE stream."""
buffer: list[str] = []
def _flush() -> dict[str, Any] | None:
data_lines = [l[5:].strip() for l in buffer if l.startswith("data:")]
data_lines = [line[5:].strip() for line in buffer if line.startswith("data:")]
buffer.clear()
if not data_lines:
return None
@@ -65,7 +84,7 @@ async def consume_sse(
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
) -> tuple[str, list[ToolCallRequest], str]:
"""Consume a Responses API SSE stream into ``(content, tool_calls, finish_reason)``."""
content, tool_calls, finish_reason, _ = await consume_sse_with_reasoning(
content, tool_calls, finish_reason, _, _ = await consume_sse_with_reasoning(
response,
on_content_delta=on_content_delta,
on_tool_call_delta=on_tool_call_delta,
@@ -78,13 +97,14 @@ async def consume_sse_with_reasoning(
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
on_reasoning_delta: Callable[[str], Awaitable[None]] | None = None,
) -> tuple[str, list[ToolCallRequest], str, str | None]:
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
"""Consume a Responses API SSE stream, including visible reasoning summaries."""
content = ""
tool_calls: list[ToolCallRequest] = []
tool_call_buffers: dict[str, dict[str, Any]] = {}
tool_call_args_emitted: set[str] = set()
finish_reason = "stop"
usage: dict[str, int] = {}
reasoning_content: str | None = None
streamed_reasoning = False
@@ -198,6 +218,7 @@ async def consume_sse_with_reasoning(
response_obj = event.get("response") or {}
status = response_obj.get("status")
finish_reason = map_finish_reason(status)
usage = _usage_from_response_obj(response_obj) or usage
if not reasoning_content:
summary = _extract_reasoning_summary_from_output(response_obj.get("output") or [])
if summary:
@@ -208,7 +229,7 @@ async def consume_sse_with_reasoning(
detail = event.get("error") or event.get("message") or event
raise RuntimeError(f"Response failed: {str(detail)[:500]}")
return content, tool_calls, finish_reason, reasoning_content
return content, tool_calls, finish_reason, usage, reasoning_content
def _extract_reasoning_summary_from_output(output: Any) -> str | None:
@@ -280,17 +301,7 @@ def parse_response_output(response: Any) -> LLMResponse:
arguments=args if isinstance(args, dict) else {},
))
usage_raw = response.get("usage") or {}
if not isinstance(usage_raw, dict):
dump = getattr(usage_raw, "model_dump", None)
usage_raw = dump() if callable(dump) else vars(usage_raw)
usage = {}
if usage_raw:
usage = {
"prompt_tokens": int(usage_raw.get("input_tokens") or 0),
"completion_tokens": int(usage_raw.get("output_tokens") or 0),
"total_tokens": int(usage_raw.get("total_tokens") or 0),
}
usage = _usage_from_response_obj(response)
status = response.get("status")
finish_reason = map_finish_reason(status)
nanobot/webui/file_preview.py
+135
View File
@@ -0,0 +1,135 @@
"""Workspace-scoped source preview payloads for the WebUI."""
from __future__ import annotations
import re
from pathlib import Path
from typing import Any
from urllib.parse import unquote, urlparse
from nanobot.security.workspace_access import WorkspaceScope
from nanobot.security.workspace_policy import WorkspaceBoundaryError, resolve_allowed_path
MAX_FILE_PREVIEW_BYTES = 384 * 1024
class WebUIFilePreviewError(ValueError):
"""Raised when a file cannot be previewed through the WebUI."""
def __init__(self, status: int, message: str) -> None:
super().__init__(message)
self.status = status
self.message = message
def file_preview_payload(
raw_path: str | None,
*,
scope: WorkspaceScope,
max_bytes: int = MAX_FILE_PREVIEW_BYTES,
) -> dict[str, Any]:
"""Return a text preview for a file inside the session workspace."""
path = _clean_preview_path(raw_path)
if not path:
raise WebUIFilePreviewError(400, "missing path")
if len(path) > 4096:
raise WebUIFilePreviewError(400, "path is too long")
try:
resolved = resolve_allowed_path(
path,
workspace=scope.project_path,
allowed_root=scope.project_path,
strict=True,
)
except FileNotFoundError as e:
raise WebUIFilePreviewError(404, "file not found") from e
except WorkspaceBoundaryError as e:
raise WebUIFilePreviewError(403, "file is outside the current workspace") from e
except OSError as e:
raise WebUIFilePreviewError(400, "invalid path") from e
if not resolved.is_file():
raise WebUIFilePreviewError(404, "file not found")
try:
with open(resolved, "rb") as f:
raw = f.read(max_bytes + 1)
except OSError as e:
raise WebUIFilePreviewError(500, "failed to read file") from e
if b"\0" in raw[:4096]:
raise WebUIFilePreviewError(415, "binary files cannot be previewed")
truncated = len(raw) > max_bytes
preview_bytes = raw[:max_bytes]
try:
content = preview_bytes.decode("utf-8")
except UnicodeDecodeError:
content = preview_bytes.decode("utf-8", errors="replace")
display_path = _display_path(resolved, scope.project_path)
return {
"path": str(resolved),
"display_path": display_path,
"project_path": str(scope.project_path),
"language": _language_for_path(resolved),
"content": content,
"size": resolved.stat().st_size,
"truncated": truncated,
}
def _clean_preview_path(raw_path: str | None) -> str:
if raw_path is None:
return ""
value = raw_path.strip()
if not value:
return ""
if value.startswith("file://"):
parsed = urlparse(value)
value = unquote(parsed.path)
else:
value = unquote(value)
value = value.split("?", 1)[0].split("#", 1)[0].strip()
if not re.match(r"^[A-Za-z]:[\\/]", value):
value = re.sub(r":\d+(?::\d+)?$", "", value)
return value
def _display_path(path: Path, root: Path) -> str:
try:
return path.relative_to(root).as_posix()
except ValueError:
return path.as_posix()
def _language_for_path(path: Path) -> str:
name = path.name.lower()
ext = path.suffix.lower().lstrip(".")
if name == "dockerfile":
return "dockerfile"
return {
"cjs": "javascript",
"css": "css",
"cts": "typescript",
"html": "html",
"js": "javascript",
"json": "json",
"jsonl": "json",
"jsx": "jsx",
"md": "markdown",
"mdx": "markdown",
"mjs": "javascript",
"mts": "typescript",
"py": "python",
"pyi": "python",
"scss": "scss",
"sh": "bash",
"toml": "toml",
"ts": "typescript",
"tsx": "tsx",
"yaml": "yaml",
"yml": "yaml",
}.get(ext, ext or "text")
nanobot/webui/settings_api.py
+8
View File
@@ -23,6 +23,7 @@ from nanobot.providers.image_generation import (
)
from nanobot.providers.registry import PROVIDERS, find_by_name
from nanobot.security.workspace_access import workspace_sandbox_status
from nanobot.webui.token_usage import token_usage_payload
from nanobot.webui.workspaces import (
read_webui_default_access_mode,
write_webui_default_access_mode,
@@ -747,6 +748,7 @@ def settings_payload(
},
"unified_session": defaults.unified_session,
},
"usage": token_usage_payload(timezone_name=defaults.timezone),
"advanced": {
"restrict_to_workspace": config.tools.restrict_to_workspace,
"workspace_sandbox": sandbox_status.as_dict(),
@@ -771,6 +773,12 @@ def settings_payload(
)
def settings_usage_payload() -> dict[str, Any]:
"""Return the lightweight token usage slice for Overview refreshes."""
config = load_config()
return token_usage_payload(timezone_name=config.agents.defaults.timezone)
def update_agent_settings(query: QueryParams) -> dict[str, Any]:
config = load_config()
defaults = config.agents.defaults
nanobot/webui/settings_routes.py
+8
View File
@@ -27,6 +27,7 @@ from nanobot.webui.settings_api import (
logout_oauth_provider,
provider_models_payload,
settings_payload,
settings_usage_payload,
update_agent_settings,
update_image_generation_settings,
update_model_configuration,
@@ -79,6 +80,8 @@ class WebUISettingsRouter:
async def dispatch(self, request: WsRequest, path: str) -> Response | None:
if path == "/api/settings":
return self._handle_settings(request)
if path == "/api/settings/usage":
return self._handle_settings_usage(request)
if path == "/api/settings/update":
return self._handle_settings_update(request)
if path == "/api/settings/model-configurations/create":
@@ -184,6 +187,11 @@ class WebUISettingsRouter:
)
)
def _handle_settings_usage(self, request: WsRequest) -> Response:
if not self._authorized(request):
return self._unauthorized()
return self._json_response(settings_usage_payload())
def _handle_settings_update(self, request: WsRequest) -> Response:
if not self._authorized(request):
return self._unauthorized()
nanobot/webui/token_usage.py
+331
View File
@@ -0,0 +1,331 @@
"""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
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
from loguru import logger
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.config.paths import get_webui_dir
TOKEN_USAGE_SCHEMA_VERSION = 1
_MAX_STATE_FILE_BYTES = 512 * 1024
_MAX_DAYS_RETAINED = 400
_USAGE_KEYS = (
"prompt_tokens",
"completion_tokens",
"cached_tokens",
"total_tokens",
"provider_tokens",
"estimated_tokens",
)
_REQUEST_KEYS = ("requests", "provider_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 _normalize_usage(raw: dict[str, Any] | None) -> dict[str, int]:
if not isinstance(raw, dict):
return {}
usage = {key: _clean_int(raw.get(key)) for key in _USAGE_KEYS}
fallback_total = usage["prompt_tokens"] + usage["completion_tokens"]
if usage["total_tokens"] <= 0:
usage["total_tokens"] = fallback_total
if usage["estimated_tokens"] <= 0 and usage["provider_tokens"] <= 0:
usage["provider_tokens"] = usage["total_tokens"]
elif usage["estimated_tokens"] > 0 and usage["provider_tokens"] <= 0:
usage["estimated_tokens"] = min(usage["estimated_tokens"], usage["total_tokens"])
elif usage["provider_tokens"] > 0 and usage["estimated_tokens"] <= 0:
usage["provider_tokens"] = min(usage["provider_tokens"], usage["total_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["prompt_tokens"] + cleaned["completion_tokens"]
if cleaned["provider_tokens"] <= 0 and cleaned["estimated_tokens"] <= 0:
cleaned["provider_tokens"] = cleaned["total_tokens"]
requests = {key: _clean_int(row.get(key)) for key in _REQUEST_KEYS}
if (
requests["requests"] > 0
and requests["provider_requests"] <= 0
and requests["estimated_requests"] <= 0
):
if cleaned["estimated_tokens"] > 0 and cleaned["provider_tokens"] <= 0:
requests["estimated_requests"] = requests["requests"]
else:
requests["provider_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 in raw.items():
if not isinstance(row, dict):
continue
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
days_raw = raw.get("days")
if not isinstance(days_raw, dict):
return state
days: dict[str, dict[str, Any]] = {}
for date, row in sorted(days_raw.items())[-_MAX_DAYS_RETAINED:]:
if not isinstance(date, str) or len(date) != 10 or not isinstance(row, dict):
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 write_token_usage_state(raw: dict[str, Any]) -> dict[str, Any]:
state = normalize_token_usage_state(raw)
state["updated_at"] = _utc_now_iso()
encoded = json.dumps(
state,
ensure_ascii=False,
indent=2,
sort_keys=True,
).encode("utf-8")
if len(encoded) > _MAX_STATE_FILE_BYTES:
raise ValueError("token usage state is too large")
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.write(b"\n")
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: dict[str, Any] | 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()
day = _local_day(now, timezone_name=timezone_name)
row = dict(state["days"].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("provider_tokens", 0) <= 0:
row["estimated_requests"] = _clean_int(row.get("estimated_requests")) + 1
else:
row["provider_requests"] = _clean_int(row.get("provider_requests")) + 1
source_key = _clean_source(source)
sources = dict(row.get("sources") or {})
source_row = 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("provider_tokens", 0) <= 0:
source_row["estimated_requests"] = _clean_int(source_row.get("estimated_requests")) + 1
else:
source_row["provider_requests"] = _clean_int(source_row.get("provider_requests")) + 1
sources[source_key] = source_row
row["sources"] = sources
state["days"][day] = row
if len(state["days"]) > _MAX_DAYS_RETAINED:
kept = dict(sorted(state["days"].items())[-_MAX_DAYS_RETAINED:])
state["days"] = kept
return write_token_usage_state(state)
def token_usage_payload(
*,
days: int = 371,
timezone_name: str | None = None,
now: datetime | None = None,
) -> dict[str, Any]:
state = read_token_usage_state()
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(state["days"].items())
if start.isoformat() <= date <= today.isoformat()
]
last_30_start = today - timedelta(days=29)
last_30 = [
row
for date, row in state["days"].items()
if last_30_start.isoformat() <= date <= today.isoformat()
]
last_365_start = today - timedelta(days=364)
last_365 = [
row
for date, row in state["days"].items()
if last_365_start.isoformat() <= date <= today.isoformat()
]
active_dates = {
datetime.fromisoformat(date).date()
for date, row in state["days"].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(state["days"].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
def include_ephemeral(self) -> bool:
return True
async def after_iteration(self, context: AgentHookContext) -> None:
try:
record_token_usage(
context.usage,
source=context.usage_source,
timezone_name=self._timezone_name,
)
except Exception:
logger.exception("failed to record token usage")
nanobot/webui/transcript.py
+153 -10
View File
@@ -572,6 +572,26 @@ def replay_transcript_to_ui_messages(
active_activity_segment_id = segment_id
return segment_id
def _turn_fields(rec: dict[str, Any], fallback_phase: str | None = None) -> dict[str, Any]:
fields: dict[str, Any] = {}
turn_id = rec.get("turn_id")
if isinstance(turn_id, str) and turn_id:
fields["turnId"] = turn_id
phase = rec.get("turn_phase")
if isinstance(phase, str) and phase:
fields["turnPhase"] = phase
elif fallback_phase:
fields["turnPhase"] = fallback_phase
seq = rec.get("turn_seq")
if isinstance(seq, (int, float)):
fields["turnSeq"] = int(seq)
return fields
def _same_turn(message: dict[str, Any], turn_fields: dict[str, Any]) -> bool:
turn_id = turn_fields.get("turnId")
message_turn_id = message.get("turnId")
return not turn_id or not message_turn_id or turn_id == message_turn_id
def _ensure_activity_segment() -> str:
return active_activity_segment_id or _new_activity_segment()
@@ -586,7 +606,13 @@ def replay_transcript_to_ui_messages(
active_activity_segment_id = None
active_file_edit_segment_id = None
def attach_reasoning_chunk(prev: list[dict[str, Any]], chunk: str, idx: int) -> None:
def attach_reasoning_chunk(
prev: list[dict[str, Any]],
chunk: str,
idx: int,
turn_fields: dict[str, Any] | None = None,
) -> None:
turn_fields = turn_fields or {}
for i in range(len(prev) - 1, -1, -1):
candidate = prev[i]
if candidate.get("role") == "user":
@@ -595,6 +621,8 @@ def replay_transcript_to_ui_messages(
break
if candidate.get("role") != "assistant":
continue
if not _same_turn(candidate, turn_fields):
break
content = str(candidate.get("content") or "")
has_answer = len(content) > 0
if (
@@ -608,6 +636,7 @@ def replay_transcript_to_ui_messages(
"reasoning": (str(candidate.get("reasoning") or "")) + chunk,
"reasoningStreaming": True,
"activitySegmentId": candidate.get("activitySegmentId") or _ensure_activity_segment(),
**turn_fields,
}
return
if not has_answer and candidate.get("isStreaming"):
@@ -616,6 +645,7 @@ def replay_transcript_to_ui_messages(
"reasoning": chunk,
"reasoningStreaming": True,
"activitySegmentId": candidate.get("activitySegmentId") or _ensure_activity_segment(),
**turn_fields,
}
return
break
@@ -629,11 +659,16 @@ def replay_transcript_to_ui_messages(
"reasoning": chunk,
"reasoningStreaming": True,
"activitySegmentId": segment,
**turn_fields,
"createdAt": _ts_base + idx,
},
)
def find_active_placeholder(prev: list[dict[str, Any]]) -> str | None:
def find_active_placeholder(
prev: list[dict[str, Any]],
turn_fields: dict[str, Any] | None = None,
) -> str | None:
turn_fields = turn_fields or {}
last = prev[-1] if prev else None
if not last:
return None
@@ -643,6 +678,8 @@ def replay_transcript_to_ui_messages(
return None
if not last.get("isStreaming"):
return None
if not _same_turn(last, turn_fields):
return None
return str(last.get("id"))
def demote_interrupted_assistant(segment: str) -> None:
@@ -721,7 +758,7 @@ def replay_transcript_to_ui_messages(
def absorb_complete(extra: dict[str, Any], idx: int) -> None:
nonlocal active_activity_segment_id, active_file_edit_segment_id
last = messages[-1] if messages else None
if last and is_reasoning_only_placeholder(last):
if last and is_reasoning_only_placeholder(last) and _same_turn(last, extra):
messages[-1] = {
**last,
**extra,
@@ -768,8 +805,13 @@ def replay_transcript_to_ui_messages(
return i
return None
def upsert_file_edits(edits: list[dict[str, Any]], idx: int) -> None:
def upsert_file_edits(
edits: list[dict[str, Any]],
idx: int,
turn_fields: dict[str, Any] | None = None,
) -> None:
nonlocal active_file_edit_segment_id
turn_fields = turn_fields or {}
if not edits:
return
segment = active_file_edit_segment_id
@@ -796,6 +838,7 @@ def replay_transcript_to_ui_messages(
"traces": [],
"fileEdits": [],
"activitySegmentId": segment,
**turn_fields,
"createdAt": _ts_base + idx,
},
)
@@ -827,6 +870,7 @@ def replay_transcript_to_ui_messages(
**last,
"fileEdits": existing,
"activitySegmentId": last.get("activitySegmentId") or segment,
**turn_fields,
}
for idx, rec in enumerate(lines):
@@ -847,6 +891,7 @@ def replay_transcript_to_ui_messages(
"id": _new_id("u", idx),
"role": "user",
"content": text_s,
**_turn_fields(rec, "user"),
"createdAt": _ts_base + idx,
}
if media_att:
@@ -867,7 +912,11 @@ def replay_transcript_to_ui_messages(
if ev == "file_edit":
raw_edits = rec.get("edits")
if isinstance(raw_edits, list):
upsert_file_edits([e for e in raw_edits if isinstance(e, dict)], idx)
upsert_file_edits(
[e for e in raw_edits if isinstance(e, dict)],
idx,
_turn_fields(rec, "activity"),
)
continue
if ev == "delta":
@@ -877,7 +926,8 @@ def replay_transcript_to_ui_messages(
if not isinstance(chunk, str):
continue
close_activity_for_answer()
adopted = find_active_placeholder(messages) if buffer_message_id is None else None
turn_fields = _turn_fields(rec, "answer")
adopted = find_active_placeholder(messages, turn_fields) if buffer_message_id is None else None
if buffer_message_id is None:
if adopted:
buffer_message_id = adopted
@@ -889,6 +939,7 @@ def replay_transcript_to_ui_messages(
"role": "assistant",
"content": "",
"isStreaming": True,
**_turn_fields(rec, "answer"),
"createdAt": _ts_base + idx,
},
)
@@ -896,7 +947,12 @@ def replay_transcript_to_ui_messages(
combined = "".join(buffer_parts)
for i, m in enumerate(messages):
if m.get("id") == buffer_message_id:
messages[i] = {**m, "content": combined, "isStreaming": True}
messages[i] = {
**m,
"content": combined,
"isStreaming": True,
**_turn_fields(rec, "answer"),
}
break
continue
@@ -915,13 +971,19 @@ def replay_transcript_to_ui_messages(
"role": "assistant",
"content": final_text,
"isStreaming": True,
**_turn_fields(rec, "answer"),
"createdAt": _ts_base + idx,
},
)
else:
for i, m in enumerate(messages):
if m.get("id") == buffer_message_id:
messages[i] = {**m, "content": final_text, "isStreaming": True}
messages[i] = {
**m,
"content": final_text,
"isStreaming": True,
**_turn_fields(rec, "answer"),
}
break
buffer_message_id = None
buffer_parts = []
@@ -934,7 +996,7 @@ def replay_transcript_to_ui_messages(
if not isinstance(chunk, str) or not chunk:
continue
close_file_edit_phase_before_activity()
attach_reasoning_chunk(messages, chunk, idx)
attach_reasoning_chunk(messages, chunk, idx, _turn_fields(rec, "reasoning"))
continue
if ev == "reasoning_end":
@@ -956,7 +1018,7 @@ def replay_transcript_to_ui_messages(
if not isinstance(line, str) or not line:
continue
close_file_edit_phase_before_activity()
attach_reasoning_chunk(messages, line, idx)
attach_reasoning_chunk(messages, line, idx, _turn_fields(rec, "reasoning"))
close_reasoning(messages)
continue
if kind in ("tool_hint", "progress"):
@@ -998,6 +1060,7 @@ def replay_transcript_to_ui_messages(
if visible_structured_events
else last.get("toolEvents"),
"activitySegmentId": last.get("activitySegmentId") or segment,
**_turn_fields(rec, "activity"),
}
messages[-1] = merged
else:
@@ -1010,6 +1073,7 @@ def replay_transcript_to_ui_messages(
"traces": trace_lines,
**({"toolEvents": visible_structured_events} if visible_structured_events else {}),
"activitySegmentId": segment,
**_turn_fields(rec, "activity"),
"createdAt": _ts_base + idx,
},
)
@@ -1033,6 +1097,7 @@ def replay_transcript_to_ui_messages(
lat = rec.get("latency_ms")
if isinstance(lat, (int, float)) and lat >= 0:
extra["latencyMs"] = int(lat)
extra.update(_turn_fields(rec, "answer"))
absorb_complete(extra, idx)
if media:
suppress_until_turn_end = True
@@ -1066,6 +1131,84 @@ def replay_transcript_to_ui_messages(
return messages
def _session_content_text(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for item in content:
if isinstance(item, Mapping):
if item.get("type") == "text" and isinstance(item.get("text"), str):
parts.append(str(item["text"]))
elif isinstance(item.get("content"), str):
parts.append(str(item["content"]))
return "\n".join(parts)
return ""
def _session_user_to_transcript_event(message: Mapping[str, Any]) -> dict[str, Any]:
event: dict[str, Any] = {
"event": "user",
"text": _session_content_text(message.get("content")),
}
media = message.get("media")
if isinstance(media, list) and media:
event["media_paths"] = [str(path) for path in media if path]
cli_apps = message.get("cli_apps")
if isinstance(cli_apps, list) and cli_apps:
event["cli_apps"] = [dict(app) for app in cli_apps if isinstance(app, dict)]
mcp_presets = message.get("mcp_presets")
if isinstance(mcp_presets, list) and mcp_presets:
event["mcp_presets"] = [dict(preset) for preset in mcp_presets if isinstance(preset, dict)]
return event
def _restore_session_user_events(
lines: list[dict[str, Any]],
session_messages: list[dict[str, Any]] | None,
) -> list[dict[str, Any]]:
"""Interleave missing user events for legacy transcripts that only persisted replies."""
if not session_messages:
return lines
session_user_count = sum(1 for m in session_messages if m.get("role") == "user")
transcript_user_count = sum(1 for line in lines if line.get("event") == "user")
if session_user_count == 0 or transcript_user_count >= session_user_count:
return lines
non_user_lines = [line for line in lines if line.get("event") != "user"]
line_index = 0
def pop_assistant_turn() -> list[dict[str, Any]]:
nonlocal line_index
turn: list[dict[str, Any]] = []
while line_index < len(non_user_lines):
current = non_user_lines[line_index]
line_index += 1
turn.append(current)
ev = current.get("event")
if ev == "turn_end":
break
if ev in {"message", "stream_end"} and (
line_index >= len(non_user_lines)
or non_user_lines[line_index].get("event") != "turn_end"
):
break
return turn
restored: list[dict[str, Any]] = []
for session_message in session_messages:
role = session_message.get("role")
if role == "user":
restored.append(_session_user_to_transcript_event(session_message))
continue
if role == "assistant":
restored.extend(pop_assistant_turn())
if line_index < len(non_user_lines):
restored.extend(non_user_lines[line_index:])
return restored
def build_webui_thread_response(
session_key: str,
*,
nanobot/webui/ws_http.py
+29 -1
View File
@@ -22,6 +22,7 @@ from websockets.http11 import Response
from nanobot.command.builtin import builtin_command_palette
from nanobot.utils.subagent_channel_display import scrub_subagent_messages_for_channel
from nanobot.webui.file_preview import WebUIFilePreviewError, file_preview_payload
from nanobot.webui.gateway_tokens import GatewayTokenStore, token_response_payload
from nanobot.webui.http_utils import (
case_insensitive_header as _case_insensitive_header,
@@ -162,6 +163,9 @@ class GatewayHTTPHandler:
runtime_capabilities=self._capabilities,
)
def workspace_controls_available(self, connection: Any) -> bool:
return self._runtime_surface == "native" or _is_localhost(connection)
# -- Token management ---------------------------------------------------
def check_api_token(self, request: WsRequest) -> bool:
@@ -291,6 +295,10 @@ class GatewayHTTPHandler:
if m:
return self._handle_webui_thread_get(request, m.group(1))
m = re.match(r"^/api/sessions/([^/]+)/file-preview$", got)
if m:
return self._handle_file_preview(request, m.group(1))
m = re.match(r"^/api/sessions/([^/]+)/delete$", got)
if m:
return self._handle_session_delete(request, m.group(1))
@@ -369,6 +377,24 @@ class GatewayHTTPHandler:
data["workspace_scope"] = scope.payload()
return _http_json_response(data)
def _handle_file_preview(self, request: WsRequest, key: str) -> Response:
if not self.check_api_token(request):
return _http_error(401, "Unauthorized")
decoded_key = _decode_api_key(key)
if decoded_key is None:
return _http_error(400, "invalid session key")
if not _is_websocket_channel_session_key(decoded_key):
return _http_error(404, "session not found")
path = _query_first(_parse_query(request.path), "path")
try:
payload = file_preview_payload(
path,
scope=self.workspaces.scope_for_session_key(decoded_key),
)
except WebUIFilePreviewError as e:
return _http_error(e.status, e.message)
return _http_json_response(payload)
def _handle_session_delete(self, request: WsRequest, key: str) -> Response:
if not self.check_api_token(request):
return _http_error(401, "Unauthorized")
@@ -426,7 +452,9 @@ class GatewayHTTPHandler:
if not self.check_api_token(request):
return _http_error(401, "Unauthorized")
return _http_json_response(
self.workspaces.payload(controls_available=_is_localhost(connection))
self.workspaces.payload(
controls_available=self.workspace_controls_available(connection)
)
)
def _handle_webui_sidebar_state(self, request: WsRequest) -> Response:
tests/agent/test_dream.py
+11
View File
@@ -356,6 +356,17 @@ class TestEphemeralHooks:
await loop.process_direct("test", session_key="cli:normal")
spy.before_iteration.assert_called()
async def test_extra_hooks_can_opt_into_ephemeral(self, tmp_path, _make_loop_with_spy):
"""Usage telemetry can opt into Dream without enabling all hooks."""
loop, spy = _make_loop_with_spy
spy.include_ephemeral.return_value = True
await loop.process_direct(
"test", session_key="dream:hook-test", ephemeral=True,
)
spy.before_iteration.assert_called()
class TestDreamCommitMessage:
async def test_commit_includes_response_summary(self, tmp_path):
tests/agent/test_loop_save_turn.py
+12 -12
View File
@@ -592,16 +592,16 @@ async def test_internal_continuation_queues_turn_without_fake_user_history(
"paused",
[],
[*initial_messages, {"role": "assistant", "content": "paused"}],
"max_iterations",
False,
)
"max_iterations",
False,
)
return (
"done",
[],
[*initial_messages, {"role": "assistant", "content": "done"}],
"completed",
False,
)
"completed",
False,
)
loop._run_agent_loop = fake_run_agent_loop # type: ignore[method-assign]
pending: asyncio.Queue[InboundMessage] = asyncio.Queue()
@@ -665,9 +665,9 @@ async def test_internal_continuation_preserves_streaming_route_metadata(
"paused",
[],
[*initial_messages, {"role": "assistant", "content": "paused"}],
"max_iterations",
False,
)
"max_iterations",
False,
)
assert on_stream is not None
assert on_stream_end is not None
await on_stream("done")
@@ -744,9 +744,9 @@ async def test_websocket_internal_continuation_keeps_single_visible_run(
"paused",
[],
[*initial_messages, {"role": "assistant", "content": "paused"}],
"max_iterations",
False,
)
"max_iterations",
False,
)
return (
"done",
[],
tests/agent/test_runner_hooks.py
+48 -1
View File
@@ -170,6 +170,48 @@ async def test_runner_passes_cached_tokens_to_hook_context():
assert len(captured_usage) == 1
assert captured_usage[0]["cached_tokens"] == 150
assert captured_usage[0]["provider_tokens"] == 220
@pytest.mark.asyncio
async def test_runner_estimates_usage_when_provider_omits_usage(monkeypatch):
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.runner import AgentRunner, AgentRunSpec
provider = MagicMock(spec=LLMProvider)
captured_usage: list[dict] = []
class UsageHook(AgentHook):
async def after_iteration(self, context: AgentHookContext) -> None:
captured_usage.append(dict(context.usage))
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = [{"type": "function", "function": {"name": "lookup"}}]
monkeypatch.setattr(
"nanobot.agent.runner.estimate_prompt_tokens_chain",
lambda provider, model, messages, tools: (123, "test"),
)
monkeypatch.setattr("nanobot.agent.runner.estimate_message_tokens", lambda message: 7)
runner = AgentRunner(provider)
result = await runner.run(AgentRunSpec(
initial_messages=[{"role": "user", "content": "hi"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=UsageHook(),
))
assert result.usage["prompt_tokens"] == 123
assert result.usage["completion_tokens"] == 7
assert result.usage["total_tokens"] == 130
assert result.usage["estimated_tokens"] == 130
assert captured_usage[0]["estimated_tokens"] == 130
@pytest.mark.asyncio
@@ -232,7 +274,12 @@ async def test_runner_calls_run_level_hooks_on_success():
"done",
"completed",
None,
{"prompt_tokens": 3, "completion_tokens": 2},
{
"prompt_tokens": 3,
"completion_tokens": 2,
"total_tokens": 5,
"provider_tokens": 5,
},
["user", "assistant"],
),
("on_finally", "completed", None),
tests/channels/test_websocket_channel.py
+133 -2
View File
@@ -27,7 +27,7 @@ from nanobot.channels.websocket import (
from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import Config, ModelPresetConfig
from nanobot.session import webui_turns as wth
from nanobot.session.manager import SessionManager
from nanobot.session.manager import Session, SessionManager
from nanobot.webui.gateway_services import GatewayServices, build_gateway_services
from nanobot.webui.http_utils import (
issue_route_secret_matches as _issue_route_secret_matches,
@@ -104,6 +104,7 @@ def bus() -> MagicMock:
@pytest.fixture(autouse=True)
def isolate_webui_workspace_state(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path)
monkeypatch.setattr(
"nanobot.webui.workspaces.get_webui_dir",
lambda: tmp_path / "webui",
@@ -277,6 +278,8 @@ async def test_token_issue_route_requires_secret_when_static_token_configured(bu
@pytest.mark.asyncio
async def test_webui_message_envelope_marks_inbound_metadata(bus: MagicMock) -> None:
from nanobot.webui.transcript import read_transcript_lines
channel = _ch(bus)
conn = MagicMock()
conn.remote_address = ("127.0.0.1", 50123)
@@ -284,14 +287,30 @@ async def test_webui_message_envelope_marks_inbound_metadata(bus: MagicMock) ->
await channel._dispatch_envelope(
conn,
"webui-client",
{"type": "message", "chat_id": "chat-1", "content": "hello", "webui": True},
{
"type": "message",
"chat_id": "chat-1",
"content": "hello",
"webui": True,
"turn_id": "turn-1",
},
)
msg = bus.publish_inbound.await_args.args[0]
assert msg.channel == "websocket"
assert msg.chat_id == "chat-1"
assert msg.metadata["webui"] is True
assert msg.metadata["webui_turn_id"] == "turn-1"
assert msg.metadata["_wants_stream"] is True
lines = read_transcript_lines("websocket:chat-1")
assert lines == [{
"event": "user",
"chat_id": "chat-1",
"text": "hello",
"turn_id": "turn-1",
"turn_phase": "user",
"turn_seq": 1,
}]
@pytest.mark.asyncio
@@ -664,6 +683,58 @@ async def test_webui_scope_rejects_non_loopback_custom_scope(bus: MagicMock, tmp
assert sessions.read_session_file("websocket:chat-remote") is None
@pytest.mark.asyncio
async def test_native_webui_scope_allows_custom_scope_without_loopback(
bus: MagicMock,
tmp_path,
) -> None:
default_workspace = tmp_path / "default"
project = tmp_path / "project"
default_workspace.mkdir()
project.mkdir()
sessions = SessionManager(tmp_path / "sessions")
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "host": "127.0.0.1"},
bus,
gateway=_basic_handler(
bus,
session_manager=sessions,
workspace_path=default_workspace,
runtime_surface="native",
),
)
conn = AsyncMock()
conn.remote_address = None
await channel._dispatch_envelope(
conn,
"native-client",
{
"type": "set_workspace_scope",
"chat_id": "chat-native",
"workspace_scope": {
"project_path": str(project),
"access_mode": "full",
},
},
)
payload = json.loads(conn.send.await_args.args[0])
assert payload["event"] == "session_updated"
assert payload["chat_id"] == "chat-native"
assert payload["workspace_scope"]["project_path"] == str(project.resolve())
assert payload["workspace_scope"]["project_name"] == "project"
assert payload["workspace_scope"]["access_mode"] == "full"
assert payload["workspace_scope"]["restrict_to_workspace"] is False
assert payload["workspace_scope"]["sandbox_status"]["restrict_to_workspace"] is False
assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve())
saved = sessions.read_session_file("websocket:chat-native")
assert saved["metadata"]["workspace_scope"] == {
"project_path": str(project.resolve()),
"access_mode": "full",
}
@pytest.mark.asyncio
async def test_send_delivers_json_message_with_media_and_reply() -> None:
bus = MagicMock()
@@ -799,6 +870,7 @@ async def test_send_progress_includes_structured_tool_events() -> None:
metadata={
"_progress": True,
"_tool_hint": True,
"webui_turn_id": "turn-1",
"_tool_events": [
{
"version": 1,
@@ -818,6 +890,9 @@ async def test_send_progress_includes_structured_tool_events() -> None:
payload = json.loads(mock_ws.send.await_args.args[0])
assert payload["event"] == "message"
assert payload["kind"] == "tool_hint"
assert payload["turn_id"] == "turn-1"
assert payload["turn_phase"] == "activity"
assert payload["turn_seq"] == 1
assert payload["tool_events"] == [
{
"version": 1,
@@ -2494,6 +2569,62 @@ def test_handle_webui_thread_get_returns_json(tmp_path, monkeypatch) -> None:
assert body["messages"][0]["content"] == "hi"
def test_handle_file_preview_returns_workspace_file(tmp_path) -> None:
from urllib.parse import quote
from websockets.datastructures import Headers
from websockets.http11 import Request
workspace = tmp_path / "workspace"
source = workspace / "nanobot" / "agent" / "hook.py"
source.parent.mkdir(parents=True)
source.write_text("print('hello')\n", encoding="utf-8")
gateway = _basic_handler(MagicMock(), workspace_path=workspace)
gateway.tokens.api_tokens["tok"] = time.monotonic() + 300.0
key = "websocket:file-preview"
enc = quote(key, safe="")
path = quote("nanobot/agent/hook.py:12", safe="")
req = Request(
f"/api/sessions/{enc}/file-preview?path={path}",
Headers([("Authorization", "Bearer tok")]),
)
resp = gateway.http._handle_file_preview(req, enc)
assert resp.status_code == 200
body = json.loads(resp.body.decode())
assert body["display_path"] == "nanobot/agent/hook.py"
assert body["language"] == "python"
assert body["content"].splitlines() == ["print('hello')"]
assert body["truncated"] is False
def test_handle_file_preview_rejects_paths_outside_workspace(tmp_path) -> None:
from urllib.parse import quote
from websockets.datastructures import Headers
from websockets.http11 import Request
workspace = tmp_path / "workspace"
workspace.mkdir()
outside = tmp_path / "secret.py"
outside.write_text("secret = True\n", encoding="utf-8")
gateway = _basic_handler(MagicMock(), workspace_path=workspace)
gateway.tokens.api_tokens["tok"] = time.monotonic() + 300.0
key = "websocket:file-preview"
enc = quote(key, safe="")
req = Request(
f"/api/sessions/{enc}/file-preview?path={quote(str(outside), safe='')}",
Headers([("Authorization", "Bearer tok")]),
)
resp = gateway.http._handle_file_preview(req, enc)
assert resp.status_code == 403
def test_handle_webui_thread_get_backfills_legacy_missing_user_rows(
tmp_path,
monkeypatch,
tests/cli/test_commands.py
+61 -3
View File
@@ -10,10 +10,9 @@ from typer.testing import CliRunner
from nanobot.bus.events import OutboundMessage
from nanobot.cli.commands import app
from nanobot.providers.factory import make_provider
from nanobot.config.schema import Config
from nanobot.cron.types import CronJob, CronPayload
from nanobot.providers.factory import ProviderSnapshot
from nanobot.providers.factory import ProviderSnapshot, make_provider
from nanobot.providers.openai_codex_provider import _strip_model_prefix
from nanobot.providers.registry import find_by_name
@@ -542,8 +541,8 @@ def test_openai_compat_provider_passes_model_through():
def test_make_provider_uses_github_copilot_backend():
from nanobot.providers.factory import make_provider
from nanobot.config.schema import Config
from nanobot.providers.factory import make_provider
config = Config.model_validate(
{
@@ -1599,6 +1598,65 @@ def test_configure_desktop_gateway_forces_local_websocket_only() -> None:
assert extras["websocket"]["websocket_requires_token"] is True
def test_load_or_create_desktop_config_bootstraps_without_api_key(tmp_path: Path) -> None:
from nanobot.cli.commands import _load_or_create_desktop_config
config_path = tmp_path / "config.json"
loaded = _load_or_create_desktop_config(
str(config_path),
str(tmp_path / "workspace"),
)
assert loaded.agents.defaults.provider == "openai_codex"
assert loaded.agents.defaults.model == "openai-codex/gpt-5.1-codex"
assert loaded.agents.defaults.model_preset is None
assert config_path.exists()
saved = json.loads(config_path.read_text(encoding="utf-8"))
assert saved["agents"]["defaults"]["provider"] == ""
assert saved["agents"]["defaults"]["model"] == ""
assert make_provider(loaded).get_default_model() == "openai-codex/gpt-5.1-codex"
def test_load_or_create_desktop_config_repairs_existing_unconfigured_default(
tmp_path: Path,
) -> None:
from nanobot.cli.commands import _load_or_create_desktop_config
from nanobot.config.loader import save_config
config_path = tmp_path / "config.json"
save_config(Config(), config_path)
loaded = _load_or_create_desktop_config(str(config_path), None)
saved = json.loads(config_path.read_text(encoding="utf-8"))
assert loaded.agents.defaults.provider == "openai_codex"
assert loaded.agents.defaults.model == "openai-codex/gpt-5.1-codex"
assert saved["agents"]["defaults"]["provider"] == ""
assert saved["agents"]["defaults"]["model"] == ""
assert make_provider(loaded).get_default_model() == "openai-codex/gpt-5.1-codex"
def test_load_or_create_desktop_config_unwinds_persisted_bootstrap(
tmp_path: Path,
) -> None:
from nanobot.cli.commands import _load_or_create_desktop_config
from nanobot.config.loader import save_config
config_path = tmp_path / "config.json"
config = Config()
config.agents.defaults.provider = "openai_codex"
config.agents.defaults.model = "openai-codex/gpt-5.1-codex"
save_config(config, config_path)
loaded = _load_or_create_desktop_config(str(config_path), None)
saved = json.loads(config_path.read_text(encoding="utf-8"))
assert loaded.agents.defaults.provider == "openai_codex"
assert loaded.agents.defaults.model == "openai-codex/gpt-5.1-codex"
assert saved["agents"]["defaults"]["provider"] == ""
assert saved["agents"]["defaults"]["model"] == ""
def test_gateway_health_endpoint_binds_and_serves_expected_responses(
monkeypatch, tmp_path: Path
) -> None:
tests/providers/test_openai_codex_provider.py
+5 -4
View File
@@ -11,8 +11,8 @@ from loguru import logger
import nanobot.providers.base as provider_base
from nanobot.providers.openai_codex_provider import (
OpenAICodexProvider,
_codex_error_response,
_build_reasoning_options,
_codex_error_response,
_CodexHTTPError,
_friendly_error,
_request_codex,
@@ -134,7 +134,7 @@ async def test_codex_prompt_cache_key_uses_stable_conversation_prefix(monkeypatc
):
_ = on_thinking_delta, on_tool_call_delta
bodies.append(body)
return "ok", [], "stop", None
return "ok", [], "stop", {}, None
monkeypatch.setattr("nanobot.providers.openai_codex_provider._request_codex", fake_request)
@@ -259,7 +259,7 @@ async def test_codex_retry_uses_structured_timeout_metadata(monkeypatch) -> None
calls += 1
if calls == 1:
raise httpx.ReadTimeout("")
return "ok", [], "stop", None
return "ok", [], "stop", {}, None
async def fake_sleep(delay: float) -> None:
delays.append(delay)
@@ -429,7 +429,7 @@ async def test_codex_stream_surfaces_reasoning_summary(monkeypatch) -> None:
await on_content_delta("answer")
if on_thinking_delta:
await on_thinking_delta("summary")
return "answer", [], "stop", "summary"
return "answer", [], "stop", {"prompt_tokens": 10, "completion_tokens": 5}, "summary"
monkeypatch.setattr("nanobot.providers.openai_codex_provider._request_codex", fake_request)
@@ -447,6 +447,7 @@ async def test_codex_stream_surfaces_reasoning_summary(monkeypatch) -> None:
assert content_deltas == ["answer"]
assert thinking_deltas == ["summary"]
assert response.content == "answer"
assert response.usage == {"prompt_tokens": 10, "completion_tokens": 5}
assert response.reasoning_content == "summary"
tests/providers/test_openai_responses.py
+21 -5
View File
@@ -19,7 +19,6 @@ from nanobot.providers.openai_responses.parsing import (
parse_response_output,
)
# ======================================================================
# converters - split_tool_call_id
# ======================================================================
@@ -478,7 +477,7 @@ class TestConsumeSse:
async def on_reasoning(delta: str) -> None:
deltas.append(delta)
content, tool_calls, finish_reason, reasoning = await consume_sse_with_reasoning(
content, tool_calls, finish_reason, usage, reasoning = await consume_sse_with_reasoning(
response,
on_reasoning_delta=on_reasoning,
)
@@ -486,6 +485,7 @@ class TestConsumeSse:
assert content == "answer"
assert tool_calls == []
assert finish_reason == "stop"
assert usage == {}
assert reasoning == "thinking briefly"
assert deltas == ["thinking ", "briefly"]
@@ -506,7 +506,7 @@ class TestConsumeSse:
},
])
_, _, _, reasoning = await consume_sse_with_reasoning(response)
_, _, _, _, reasoning = await consume_sse_with_reasoning(response)
assert reasoning == "cached summary"
@@ -527,7 +527,7 @@ class TestConsumeSse:
async def on_reasoning(delta: str) -> None:
deltas.append(delta)
_, _, _, reasoning = await consume_sse_with_reasoning(
_, _, _, _, reasoning = await consume_sse_with_reasoning(
response,
on_reasoning_delta=on_reasoning,
)
@@ -545,10 +545,26 @@ class TestConsumeSse:
{"type": "response.completed", "response": {"status": "completed"}},
])
_, _, _, reasoning = await consume_sse_with_reasoning(response)
_, _, _, _, reasoning = await consume_sse_with_reasoning(response)
assert reasoning == "part summary"
@pytest.mark.asyncio
async def test_raw_sse_usage_extracted(self):
response = _SseResponse([
{
"type": "response.completed",
"response": {
"status": "completed",
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
},
},
])
_, _, _, usage, _ = await consume_sse_with_reasoning(response)
assert usage == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
@pytest.mark.asyncio
async def test_tool_call_done_arguments_callback(self):
response = _SseResponse([
tests/utils/test_webui_transcript.py
+118 -2
View File
@@ -43,6 +43,124 @@ def test_replay_delta_and_turn_end(tmp_path, monkeypatch) -> None:
assert msgs[1]["latencyMs"] == 42
def test_replay_preserves_turn_metadata(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path)
key = "websocket:t-turn"
for ev in (
{
"event": "user",
"chat_id": "t-turn",
"text": "q",
"turn_id": "turn-1",
"turn_phase": "user",
"turn_seq": 1,
},
{
"event": "reasoning_delta",
"chat_id": "t-turn",
"text": "think",
"turn_id": "turn-1",
"turn_phase": "reasoning",
"turn_seq": 2,
},
{
"event": "delta",
"chat_id": "t-turn",
"text": "a",
"turn_id": "turn-1",
"turn_phase": "answer",
"turn_seq": 3,
},
{
"event": "turn_end",
"chat_id": "t-turn",
"latency_ms": 12,
"turn_id": "turn-1",
"turn_phase": "complete",
"turn_seq": 4,
},
):
append_transcript_object(key, ev)
msgs = replay_transcript_to_ui_messages(read_transcript_lines(key))
assert msgs[0]["turnId"] == "turn-1"
assert msgs[0]["turnPhase"] == "user"
assert msgs[0]["turnSeq"] == 1
assert msgs[1]["turnId"] == "turn-1"
assert msgs[1]["turnPhase"] == "answer"
assert msgs[1]["turnSeq"] == 3
def test_build_response_restores_session_users_for_legacy_transcript(
tmp_path,
monkeypatch,
) -> None:
monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path)
key = "websocket:legacy-users"
append_transcript_object(
key,
{"event": "message", "chat_id": "legacy-users", "text": "assistant one"},
)
append_transcript_object(
key,
{"event": "message", "chat_id": "legacy-users", "text": "assistant two"},
)
out = build_webui_thread_response(
key,
session_messages=[
{"role": "user", "content": "prompt one", "timestamp": "2026-06-02T10:00:00"},
{"role": "assistant", "content": "session one"},
{"role": "user", "content": "prompt two", "timestamp": "2026-06-02T10:01:00"},
{"role": "assistant", "content": "session two"},
],
)
assert out is not None
assert [(m["role"], m["content"]) for m in out["messages"]] == [
("user", "prompt one"),
("assistant", "assistant one"),
("user", "prompt two"),
("assistant", "assistant two"),
]
def test_build_response_restores_session_users_without_duplicating_new_transcript_users(
tmp_path,
monkeypatch,
) -> None:
monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path)
key = "websocket:mixed-users"
append_transcript_object(
key,
{"event": "message", "chat_id": "mixed-users", "text": "old assistant"},
)
append_transcript_object(key, {"event": "user", "chat_id": "mixed-users", "text": "new prompt"})
append_transcript_object(
key,
{"event": "message", "chat_id": "mixed-users", "text": "new assistant"},
)
out = build_webui_thread_response(
key,
session_messages=[
{"role": "user", "content": "old prompt"},
{"role": "assistant", "content": "old session assistant"},
{"role": "user", "content": "new prompt"},
{"role": "assistant", "content": "new session assistant"},
],
)
assert out is not None
assert [(m["role"], m["content"]) for m in out["messages"]] == [
("user", "old prompt"),
("assistant", "old assistant"),
("user", "new prompt"),
("assistant", "new assistant"),
]
def test_replay_augments_assistant_text() -> None:
msgs = replay_transcript_to_ui_messages(
[
@@ -675,8 +793,6 @@ def test_replay_keeps_new_file_edit_after_reasoning_in_order(tmp_path, monkeypat
def test_build_response_schema(monkeypatch, tmp_path) -> None:
from nanobot.webui.transcript import build_webui_thread_response
monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path)
key = "websocket:t3"
append_transcript_object(key, {"event": "user", "chat_id": "t3", "text": "x"})
tests/webui/test_settings_api.py
+47
View File
@@ -14,6 +14,7 @@ from nanobot.webui.settings_api import (
create_model_configuration,
provider_models_payload,
settings_payload,
settings_usage_payload,
update_agent_settings,
update_model_configuration,
update_network_safety_settings,
@@ -242,6 +243,52 @@ def test_settings_payload_includes_network_safety_fields(
assert payload["advanced"]["ssrf_whitelist_count"] == 1
def test_settings_payload_includes_token_usage_summary(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
config_path = tmp_path / "config.json"
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({"prompt_tokens": 10, "completion_tokens": 5})
payload = settings_payload()
assert payload["usage"]["total_tokens_30d"] == 15
assert payload["usage"]["total_tokens"] == 15
assert payload["usage"]["peak_day_tokens"] == 15
assert payload["usage"]["current_streak_days"] == 1
assert payload["usage"]["longest_streak_days"] == 1
assert payload["usage"]["active_days_30d"] == 1
assert payload["usage"]["requests_30d"] == 1
def test_settings_usage_payload_returns_lightweight_token_usage(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
config_path = tmp_path / "config.json"
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({"prompt_tokens": 20, "completion_tokens": 2})
payload = settings_usage_payload()
assert payload["total_tokens"] == 22
assert payload["requests_30d"] == 1
assert "agent" not in payload
def test_update_network_safety_settings_writes_local_service_flag(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
tests/webui/test_token_usage.py
+109
View File
@@ -0,0 +1,109 @@
from __future__ import annotations
from datetime import datetime, timezone
from nanobot.webui.token_usage import (
record_token_usage,
token_usage_payload,
)
def test_record_token_usage_aggregates_by_local_day(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage(
{"prompt_tokens": 100, "completion_tokens": 40, "cached_tokens": 20},
timezone_name="Asia/Shanghai",
now=datetime(2026, 6, 2, 18, 0, tzinfo=timezone.utc),
)
record_token_usage(
{"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
timezone_name="Asia/Shanghai",
now=datetime(2026, 6, 2, 19, 0, tzinfo=timezone.utc),
)
payload = token_usage_payload(
timezone_name="Asia/Shanghai",
now=datetime(2026, 6, 3, 12, 0, tzinfo=timezone.utc),
)
assert payload["total_tokens_30d"] == 155
assert payload["active_days_30d"] == 1
assert payload["requests_30d"] == 2
assert payload["days"] == [
{
"date": "2026-06-03",
"prompt_tokens": 110,
"completion_tokens": 45,
"cached_tokens": 20,
"total_tokens": 155,
"provider_tokens": 155,
"estimated_tokens": 0,
"requests": 2,
"provider_requests": 2,
"estimated_requests": 0,
"sources": {
"user": {
"prompt_tokens": 110,
"completion_tokens": 45,
"cached_tokens": 20,
"total_tokens": 155,
"provider_tokens": 155,
"estimated_tokens": 0,
"requests": 2,
"provider_requests": 2,
"estimated_requests": 0,
}
},
}
]
def test_record_token_usage_skips_empty_usage(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage({"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
assert payload["days"] == []
assert payload["total_tokens_30d"] == 0
def test_record_token_usage_keeps_estimated_split(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage(
{"prompt_tokens": 100, "completion_tokens": 25, "estimated_tokens": 125},
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
)
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
assert payload["days"][0]["total_tokens"] == 125
assert payload["days"][0]["provider_tokens"] == 0
assert payload["days"][0]["estimated_tokens"] == 125
assert payload["days"][0]["estimated_requests"] == 1
def test_record_token_usage_keeps_source_breakdown(tmp_path, monkeypatch) -> None:
monkeypatch.setattr("nanobot.webui.token_usage.get_webui_dir", lambda: tmp_path / "webui")
record_token_usage(
{"prompt_tokens": 100, "completion_tokens": 25},
source="user",
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
)
record_token_usage(
{"prompt_tokens": 20, "completion_tokens": 5},
source="dream",
now=datetime(2026, 6, 3, tzinfo=timezone.utc),
)
payload = token_usage_payload(now=datetime(2026, 6, 3, tzinfo=timezone.utc))
row = payload["days"][0]
assert row["total_tokens"] == 150
assert row["sources"]["user"]["total_tokens"] == 125
assert row["sources"]["user"]["requests"] == 1
assert row["sources"]["dream"]["total_tokens"] == 25
assert row["sources"]["dream"]["requests"] == 1