diff --git a/nanobot/channels/websocket/tests/test_websocket_channel.py b/nanobot/channels/websocket/tests/test_websocket_channel.py index 96f08068d..d0125a74d 100644 --- a/nanobot/channels/websocket/tests/test_websocket_channel.py +++ b/nanobot/channels/websocket/tests/test_websocket_channel.py @@ -3319,6 +3319,11 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None: channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus)) mock_ws = AsyncMock() channel._attach(mock_ws, "chat-1") + usage = LLMUsage.reported( + input_tokens=80, + output_tokens=20, + cache_read_tokens=40, + ).with_timing(generation_ms=500, ttft_ms=125) await channel.send(OutboundMessage( channel="websocket", @@ -3326,7 +3331,7 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None: content="", event=TurnEndEvent( latency_ms=1500, - usage=LLMUsage.reported(input_tokens=80, output_tokens=20, cache_read_tokens=40), + usage=usage, context_window_tokens=128_000, ), )) @@ -3336,11 +3341,19 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None: "event": "turn_end", "chat_id": "chat-1", "latency_ms": 1500, - "usage": LLMUsage.reported( - input_tokens=80, - output_tokens=20, - cache_read_tokens=40, - ).to_turn_dict(), + "usage": { + "prompt_tokens": 80, + "completion_tokens": 20, + "total_tokens": 100, + "context_tokens": 80, + "cached_tokens": 40, + "request_count": 1, + "estimated_tokens": 0, + "generation_ms": 500, + "measured_completion_tokens": 20, + "ttft_ms": 125, + "timed_requests": 1, + }, "context_window_tokens": 128_000, }, {"event": "session_updated", "chat_id": "chat-1", "scope": "thread"}, @@ -5307,7 +5320,12 @@ async def test_handle_session_context_get_reads_detached_session() -> None: from nanobot.session import Session - usage = LLMUsage.reported(input_tokens=12, output_tokens=3, total_tokens=175) + usage = LLMUsage.reported( + input_tokens=12, + output_tokens=3, + total_tokens=175, + cache_read_tokens=6, + ).with_timing(generation_ms=300, ttft_ms=45) session = Session( key="websocket:context-route", messages=[{"role": "user", "content": "hello"}], @@ -5328,7 +5346,19 @@ async def test_handle_session_context_get_reads_detached_session() -> None: assert response.status_code == 200 body = json.loads(response.body.decode()) assert body["replay_messages"] == 1 - assert body["last_usage"] == usage.to_dict() + assert body["last_usage"] == { + "prompt_tokens": 12, + "completion_tokens": 3, + "total_tokens": 175, + "context_tokens": 12, + "cached_tokens": 6, + "request_count": 1, + "estimated_tokens": 0, + "generation_ms": 300, + "measured_completion_tokens": 3, + "ttft_ms": 45, + "timed_requests": 1, + } manager.read_session_snapshot.assert_called_once_with(session.key) diff --git a/nanobot/providers/base.py b/nanobot/providers/base.py index 749d9963d..82c12e88e 100644 --- a/nanobot/providers/base.py +++ b/nanobot/providers/base.py @@ -451,7 +451,7 @@ class LLMUsage: } def to_turn_dict(self) -> dict[str, int]: - """Project canonical usage into the WebUI's compact per-turn shape.""" + """Project canonical usage into the compact WebUI/TUI per-turn shape.""" result: dict[str, int] = { "prompt_tokens": self.input_tokens, "completion_tokens": self.output_tokens, @@ -465,6 +465,12 @@ class LLMUsage: result["cached_tokens"] = self.cache_read_tokens if self.cache_write_tokens is not None: result["cache_write_tokens"] = self.cache_write_tokens + if self.generation_ms > 0 and self.measured_output_tokens > 0: + result["generation_ms"] = self.generation_ms + result["measured_completion_tokens"] = self.measured_output_tokens + if self.timed_requests > 0: + result["ttft_ms"] = self.ttft_ms + result["timed_requests"] = self.timed_requests return result @classmethod diff --git a/nanobot/webui/session_context.py b/nanobot/webui/session_context.py index 5ec3b18ab..d409c7fcb 100644 --- a/nanobot/webui/session_context.py +++ b/nanobot/webui/session_context.py @@ -38,7 +38,7 @@ def session_context_payload(session: Session) -> dict[str, Any]: estimate_message_tokens({"role": "system", "content": summary}) if summary else 0 ) stored_usage = LLMUsage.from_dict(session.metadata.get("_last_usage")) - last_usage = stored_usage.to_dict() if stored_usage is not None else None + last_usage = stored_usage.to_turn_dict() if stored_usage is not None else None return { "schema_version": 1, diff --git a/tests/providers/test_usage_contract.py b/tests/providers/test_usage_contract.py index 9b9cbb733..477d89fa6 100644 --- a/tests/providers/test_usage_contract.py +++ b/tests/providers/test_usage_contract.py @@ -102,7 +102,10 @@ def test_usage_projects_compact_turn_observability_shape() -> None: output_tokens=3, total_tokens=20, cache_read_tokens=4, - ) + LLMUsage.estimated(input_tokens=18, output_tokens=2) + ).with_timing(generation_ms=250, ttft_ms=50) + LLMUsage.estimated( + input_tokens=18, + output_tokens=2, + ).with_timing(generation_ms=100, ttft_ms=None) assert usage.to_turn_dict() == { "prompt_tokens": 30, @@ -111,6 +114,10 @@ def test_usage_projects_compact_turn_observability_shape() -> None: "context_tokens": 18, "request_count": 2, "estimated_tokens": 20, + "generation_ms": 350, + "measured_completion_tokens": 5, + "ttft_ms": 50, + "timed_requests": 1, } diff --git a/tests/webui/test_session_context.py b/tests/webui/test_session_context.py index b125b1b5f..68bc5fe9c 100644 --- a/tests/webui/test_session_context.py +++ b/tests/webui/test_session_context.py @@ -59,7 +59,12 @@ def test_session_context_tolerates_untrusted_summary_metadata() -> None: def test_session_context_sanitizes_usage_metadata() -> None: - usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175) + usage = LLMUsage.reported( + input_tokens=120, + output_tokens=8, + total_tokens=175, + cache_read_tokens=48, + ).with_timing(generation_ms=400, ttft_ms=75) session = Session( key="websocket:context", metadata={"_last_usage": usage.to_dict()}, @@ -67,4 +72,16 @@ def test_session_context_sanitizes_usage_metadata() -> None: payload = session_context_payload(session) - assert payload["last_usage"] == usage.to_dict() + assert payload["last_usage"] == { + "prompt_tokens": 120, + "completion_tokens": 8, + "total_tokens": 175, + "context_tokens": 120, + "cached_tokens": 48, + "request_count": 1, + "estimated_tokens": 0, + "generation_ms": 400, + "measured_completion_tokens": 8, + "ttft_ms": 75, + "timed_requests": 1, + } diff --git a/tui/src/protocol.test.ts b/tui/src/protocol.test.ts index 5b7c2bd47..34bbd9472 100644 --- a/tui/src/protocol.test.ts +++ b/tui/src/protocol.test.ts @@ -726,6 +726,19 @@ describe("gateway protocol", () => { estimated_session_tokens: 2176, archived_summary: "Older work was compacted.", archived_summary_at: "2026-08-13T10:00:00Z", + last_usage: { + prompt_tokens: 1200, + completion_tokens: 80, + total_tokens: 1280, + context_tokens: 1200, + cached_tokens: 900, + request_count: 1, + estimated_tokens: 0, + generation_ms: 4000, + measured_completion_tokens: 80, + ttft_ms: 250, + timed_requests: 1, + }, })))) as unknown as typeof fetch try { @@ -738,7 +751,19 @@ describe("gateway protocol", () => { estimatedSessionTokens: 2176, archivedSummary: "Older work was compacted.", archivedSummaryAt: "2026-08-13T10:00:00Z", - lastUsage: null, + lastUsage: { + prompt_tokens: 1200, + completion_tokens: 80, + total_tokens: 1280, + context_tokens: 1200, + cached_tokens: 900, + request_count: 1, + estimated_tokens: 0, + generation_ms: 4000, + measured_completion_tokens: 80, + ttft_ms: 250, + timed_requests: 1, + }, }) } finally { globalThis.fetch = original diff --git a/tui/src/protocol.ts b/tui/src/protocol.ts index 885b5e96e..692ee2b3c 100644 --- a/tui/src/protocol.ts +++ b/tui/src/protocol.ts @@ -222,7 +222,10 @@ export interface TokenUsage { prompt_tokens?: number completion_tokens?: number cached_tokens?: number + cache_write_tokens?: number total_tokens?: number + context_tokens?: number + request_count?: number provider_tokens?: number estimated_tokens?: number cost_usd?: number @@ -372,7 +375,10 @@ function isTokenUsage(value: unknown): value is TokenUsage { "prompt_tokens", "completion_tokens", "cached_tokens", + "cache_write_tokens", "total_tokens", + "context_tokens", + "request_count", "provider_tokens", "estimated_tokens", "cost_usd",