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

fix(tui): preserve usage telemetry fields

Browse Source
This commit is contained in:
chengyongru
2026-08-25 01:04:25 +08:00
committed by chengyongru
parent 9895c23cb5
commit 8bb3828487
7 changed files with 105 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/channels/websocket/tests/test_websocket_channel.py
+38 -8
View File
@@ -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)
nanobot/providers/base.py
+7 -1
View File
@@ -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
nanobot/webui/session_context.py
+1 -1
View File
@@ -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,
tests/providers/test_usage_contract.py
+8 -1
View File
@@ -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,
}
tests/webui/test_session_context.py
+19 -2
View File
@@ -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,
}
tui/src/protocol.test.ts
+26 -1
View File
@@ -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
tui/src/protocol.ts
+6
View File
@@ -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",