diff --git a/nanobot/agent/hook.py b/nanobot/agent/hook.py index ed2c95498..b0a722468 100644 --- a/nanobot/agent/hook.py +++ b/nanobot/agent/hook.py @@ -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) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index f31589cb9..9bbc272c6 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -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) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index a6eda06fa..8746b5c27 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -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(): diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index d0de15787..85f60148c 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -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).""" diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index 62f804ca2..57d41df25 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -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, diff --git a/nanobot/providers/openai_codex_provider.py b/nanobot/providers/openai_codex_provider.py index fc92e8ae8..aeee832e4 100644 --- a/nanobot/providers/openai_codex_provider.py +++ b/nanobot/providers/openai_codex_provider.py @@ -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: diff --git a/nanobot/providers/openai_responses/parsing.py b/nanobot/providers/openai_responses/parsing.py index 846165562..fbfc9813c 100644 --- a/nanobot/providers/openai_responses/parsing.py +++ b/nanobot/providers/openai_responses/parsing.py @@ -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) diff --git a/nanobot/webui/file_preview.py b/nanobot/webui/file_preview.py new file mode 100644 index 000000000..958f3ce85 --- /dev/null +++ b/nanobot/webui/file_preview.py @@ -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") diff --git a/nanobot/webui/settings_api.py b/nanobot/webui/settings_api.py index d4c8849c9..3f3df3957 100644 --- a/nanobot/webui/settings_api.py +++ b/nanobot/webui/settings_api.py @@ -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 diff --git a/nanobot/webui/settings_routes.py b/nanobot/webui/settings_routes.py index 9e0caab57..ff5b7d7df 100644 --- a/nanobot/webui/settings_routes.py +++ b/nanobot/webui/settings_routes.py @@ -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() diff --git a/nanobot/webui/token_usage.py b/nanobot/webui/token_usage.py new file mode 100644 index 000000000..326b29c03 --- /dev/null +++ b/nanobot/webui/token_usage.py @@ -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") diff --git a/nanobot/webui/transcript.py b/nanobot/webui/transcript.py index f0a073618..cfcb2371d 100644 --- a/nanobot/webui/transcript.py +++ b/nanobot/webui/transcript.py @@ -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, *, diff --git a/nanobot/webui/ws_http.py b/nanobot/webui/ws_http.py index 2b60bfca8..4edf7eab2 100644 --- a/nanobot/webui/ws_http.py +++ b/nanobot/webui/ws_http.py @@ -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: diff --git a/tests/agent/test_dream.py b/tests/agent/test_dream.py index 5b9bd32a0..f9c6f1d7b 100644 --- a/tests/agent/test_dream.py +++ b/tests/agent/test_dream.py @@ -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): diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index 874f0b435..295bc4888 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -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", [], diff --git a/tests/agent/test_runner_hooks.py b/tests/agent/test_runner_hooks.py index 7ba107618..7eb8c2496 100644 --- a/tests/agent/test_runner_hooks.py +++ b/tests/agent/test_runner_hooks.py @@ -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), diff --git a/tests/channels/test_websocket_channel.py b/tests/channels/test_websocket_channel.py index a441547bd..6afeec154 100644 --- a/tests/channels/test_websocket_channel.py +++ b/tests/channels/test_websocket_channel.py @@ -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, diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index 353d04ecb..8f13725c6 100644 --- a/tests/cli/test_commands.py +++ b/tests/cli/test_commands.py @@ -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: diff --git a/tests/providers/test_openai_codex_provider.py b/tests/providers/test_openai_codex_provider.py index e1994555c..cbd3ed488 100644 --- a/tests/providers/test_openai_codex_provider.py +++ b/tests/providers/test_openai_codex_provider.py @@ -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" diff --git a/tests/providers/test_openai_responses.py b/tests/providers/test_openai_responses.py index 49ae86493..e9d8545e1 100644 --- a/tests/providers/test_openai_responses.py +++ b/tests/providers/test_openai_responses.py @@ -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([ diff --git a/tests/utils/test_webui_transcript.py b/tests/utils/test_webui_transcript.py index 1c7372902..55d886cd8 100644 --- a/tests/utils/test_webui_transcript.py +++ b/tests/utils/test_webui_transcript.py @@ -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"}) diff --git a/tests/webui/test_settings_api.py b/tests/webui/test_settings_api.py index 80a56ed91..d48dd6bd1 100644 --- a/tests/webui/test_settings_api.py +++ b/tests/webui/test_settings_api.py @@ -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, diff --git a/tests/webui/test_token_usage.py b/tests/webui/test_token_usage.py new file mode 100644 index 000000000..2f07162ba --- /dev/null +++ b/tests/webui/test_token_usage.py @@ -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