fix(tui): preserve usage telemetry fields

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
@@ -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)) channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus))
mock_ws = AsyncMock() mock_ws = AsyncMock()
channel._attach(mock_ws, "chat-1") 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( await channel.send(OutboundMessage(
channel="websocket", channel="websocket",
@@ -3326,7 +3331,7 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
content="", content="",
event=TurnEndEvent( event=TurnEndEvent(
latency_ms=1500, latency_ms=1500,
usage=LLMUsage.reported(input_tokens=80, output_tokens=20, cache_read_tokens=40), usage=usage,
context_window_tokens=128_000, context_window_tokens=128_000,
), ),
)) ))
@@ -3336,11 +3341,19 @@ async def test_send_turn_end_includes_latency_ms_when_present() -> None:
"event": "turn_end", "event": "turn_end",
"chat_id": "chat-1", "chat_id": "chat-1",
"latency_ms": 1500, "latency_ms": 1500,
"usage": LLMUsage.reported( "usage": {
input_tokens=80, "prompt_tokens": 80,
output_tokens=20, "completion_tokens": 20,
cache_read_tokens=40, "total_tokens": 100,
).to_turn_dict(), "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, "context_window_tokens": 128_000,
}, },
{"event": "session_updated", "chat_id": "chat-1", "scope": "thread"}, {"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 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( session = Session(
key="websocket:context-route", key="websocket:context-route",
messages=[{"role": "user", "content": "hello"}], 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 assert response.status_code == 200
body = json.loads(response.body.decode()) body = json.loads(response.body.decode())
assert body["replay_messages"] == 1 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) manager.read_session_snapshot.assert_called_once_with(session.key)
+7 -1
View File
@@ -451,7 +451,7 @@ class LLMUsage:
} }
def to_turn_dict(self) -> dict[str, int]: 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] = { result: dict[str, int] = {
"prompt_tokens": self.input_tokens, "prompt_tokens": self.input_tokens,
"completion_tokens": self.output_tokens, "completion_tokens": self.output_tokens,
@@ -465,6 +465,12 @@ class LLMUsage:
result["cached_tokens"] = self.cache_read_tokens result["cached_tokens"] = self.cache_read_tokens
if self.cache_write_tokens is not None: if self.cache_write_tokens is not None:
result["cache_write_tokens"] = self.cache_write_tokens 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 return result
@classmethod @classmethod
+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 estimate_message_tokens({"role": "system", "content": summary}) if summary else 0
) )
stored_usage = LLMUsage.from_dict(session.metadata.get("_last_usage")) 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 { return {
"schema_version": 1, "schema_version": 1,
+8 -1
View File
@@ -102,7 +102,10 @@ def test_usage_projects_compact_turn_observability_shape() -> None:
output_tokens=3, output_tokens=3,
total_tokens=20, total_tokens=20,
cache_read_tokens=4, 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() == { assert usage.to_turn_dict() == {
"prompt_tokens": 30, "prompt_tokens": 30,
@@ -111,6 +114,10 @@ def test_usage_projects_compact_turn_observability_shape() -> None:
"context_tokens": 18, "context_tokens": 18,
"request_count": 2, "request_count": 2,
"estimated_tokens": 20, "estimated_tokens": 20,
"generation_ms": 350,
"measured_completion_tokens": 5,
"ttft_ms": 50,
"timed_requests": 1,
} }
+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: 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( session = Session(
key="websocket:context", key="websocket:context",
metadata={"_last_usage": usage.to_dict()}, metadata={"_last_usage": usage.to_dict()},
@@ -67,4 +72,16 @@ def test_session_context_sanitizes_usage_metadata() -> None:
payload = session_context_payload(session) 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,
}
+26 -1
View File
@@ -726,6 +726,19 @@ describe("gateway protocol", () => {
estimated_session_tokens: 2176, estimated_session_tokens: 2176,
archived_summary: "Older work was compacted.", archived_summary: "Older work was compacted.",
archived_summary_at: "2026-08-13T10:00:00Z", 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 })))) as unknown as typeof fetch
try { try {
@@ -738,7 +751,19 @@ describe("gateway protocol", () => {
estimatedSessionTokens: 2176, estimatedSessionTokens: 2176,
archivedSummary: "Older work was compacted.", archivedSummary: "Older work was compacted.",
archivedSummaryAt: "2026-08-13T10:00:00Z", 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 { } finally {
globalThis.fetch = original globalThis.fetch = original
+6
View File
@@ -222,7 +222,10 @@ export interface TokenUsage {
prompt_tokens?: number prompt_tokens?: number
completion_tokens?: number completion_tokens?: number
cached_tokens?: number cached_tokens?: number
cache_write_tokens?: number
total_tokens?: number total_tokens?: number
context_tokens?: number
request_count?: number
provider_tokens?: number provider_tokens?: number
estimated_tokens?: number estimated_tokens?: number
cost_usd?: number cost_usd?: number
@@ -372,7 +375,10 @@ function isTokenUsage(value: unknown): value is TokenUsage {
"prompt_tokens", "prompt_tokens",
"completion_tokens", "completion_tokens",
"cached_tokens", "cached_tokens",
"cache_write_tokens",
"total_tokens", "total_tokens",
"context_tokens",
"request_count",
"provider_tokens", "provider_tokens",
"estimated_tokens", "estimated_tokens",
"cost_usd", "cost_usd",