diff --git a/api/metering.py b/api/metering.py index c4696d24..e3a08f98 100644 --- a/api/metering.py +++ b/api/metering.py @@ -44,6 +44,7 @@ The SSE `metering` event payload: from __future__ import annotations +import math import threading import time from dataclasses import dataclass @@ -58,6 +59,8 @@ class _SessionMeter: reasoning_tokens: int = 0 first_token_ts: float = 0.0 # time.monotonic() of first token received last_token_ts: float = 0.0 # time.monotonic() of last token received + pending_started_at: float | None = None + ttft_ms: int | None = None def total_tokens(self) -> int: return self.output_tokens + self.reasoning_tokens @@ -94,6 +97,18 @@ class GlobalMeter: with self._lock: self._sessions[stream_id] = _SessionMeter() + def set_pending_started_at(self, stream_id: str, pending_started_at) -> None: + try: + value = float(pending_started_at) + except (TypeError, ValueError): + return + if not math.isfinite(value) or value <= 0: + return + with self._lock: + session = self._sessions.get(stream_id) + if session is not None: + session.pending_started_at = value + def get_interval(self) -> float: """Return 1.0 when sessions are actively receiving tokens, 10.0 when idle. @@ -119,6 +134,8 @@ class GlobalMeter: s.first_token_ts = now s.last_token_ts = now s.output_tokens = running_output_tokens + if s.ttft_ms is None and s.pending_started_at is not None: + s.ttft_ms = max(0, round((time.time() - s.pending_started_at) * 1000)) def record_reasoning(self, stream_id: str, running_reasoning_tokens: int) -> None: now = time.monotonic() @@ -135,7 +152,12 @@ class GlobalMeter: with self._lock: self._sessions.pop(stream_id, None) - def get_stats(self) -> dict: + def get_ttft_ms(self, stream_id: str) -> int | None: + with self._lock: + session = self._sessions.get(stream_id) + return session.ttft_ms if session is not None else None + + def get_stats(self, stream_id: str | None = None) -> dict: now = time.monotonic() with self._lock: # Prune stale sessions @@ -175,7 +197,7 @@ class GlobalMeter: high = max(active_readings) if active_readings else 0.0 low = min(active_readings) if active_readings else 0.0 - return { + stats = { 'tps': round(global_tps, 1) if global_tps is not None else None, 'tps_available': global_tps is not None, 'estimated': False, @@ -183,6 +205,11 @@ class GlobalMeter: 'low': round(low, 1) if low else None, 'active': len(self._sessions), } + if stream_id is not None: + session = self._sessions.get(stream_id) + if session is not None and session.ttft_ms is not None: + stats['ttft_ms'] = session.ttft_ms + return stats # ── Module-level singleton ───────────────────────────────────────────────────── diff --git a/api/streaming.py b/api/streaming.py index 1f083c55..472bfc27 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -6972,7 +6972,7 @@ def _run_agent_streaming( break # nothing active — stop the ticker if _metering_stop.wait(interval): break # stream was cancelled or ended — exit - stats = meter().get_stats() + stats = meter().get_stats(stream_id) stats['session_id'] = session_id stats['usage'] = _live_usage_snapshot() put('metering', stats) @@ -7409,7 +7409,7 @@ def _run_agent_streaming( if now - _metering_last_emit[0] < 0.1: return _metering_last_emit[0] = now - stats = meter().get_stats() + stats = meter().get_stats(stream_id) stats['session_id'] = session_id stats['usage'] = _live_usage_snapshot() stats.setdefault('tps_available', False) @@ -7702,7 +7702,7 @@ def _run_agent_streaming( 'preview': preview, 'args': args_snap, }) - _tool_stats = meter().get_stats() + _tool_stats = meter().get_stats(stream_id) _tool_stats['session_id'] = session_id _tool_stats['usage'] = _live_usage_snapshot() put('metering', _tool_stats) @@ -7801,7 +7801,7 @@ def _run_agent_streaming( session_id=session_id, stream_id=stream_id, ) - _tool_stats = meter().get_stats() + _tool_stats = meter().get_stats(stream_id) _tool_stats['session_id'] = session_id _tool_stats['usage'] = _live_usage_snapshot() put('metering', _tool_stats) @@ -7833,7 +7833,7 @@ def _run_agent_streaming( 'args': _tool_args_snapshot(args), 'tid': tool_call_id, }) - _tool_stats = meter().get_stats() + _tool_stats = meter().get_stats(stream_id) _tool_stats['session_id'] = session_id _tool_stats['usage'] = _live_usage_snapshot() put('metering', _tool_stats) @@ -7884,7 +7884,7 @@ def _run_agent_streaming( session_id=session_id, stream_id=stream_id, ) - _tool_stats = meter().get_stats() + _tool_stats = meter().get_stats(stream_id) _tool_stats['session_id'] = session_id _tool_stats['usage'] = _live_usage_snapshot() put('metering', _tool_stats) @@ -8434,6 +8434,7 @@ def _run_agent_streaming( config_data=_cfg, ) _pending_started_at = getattr(s, 'pending_started_at', None) + meter().set_pending_started_at(stream_id, _pending_started_at) # Normal chat-start sets pending_started_at before spawning this thread; # fallback to now only for recovered/legacy flows where that marker is absent # or has been zeroed out (e.g. via a buggy migration / manual file edit). @@ -9381,6 +9382,9 @@ def _run_agent_streaming( _dm['_turnTps'] = _turn_tps if _gateway_routing: _dm['_gatewayRouting'] = _gateway_routing + _ttft_ms = meter().get_ttft_ms(stream_id) + if _ttft_ms is not None: + _dm['_firstTokenMs'] = _ttft_ms break # Persist context window data on the session so the context-ring # indicator survives a page reload (#1318). Must run BEFORE @@ -9765,6 +9769,9 @@ def _run_agent_streaming( usage['tps'] = _turn_tps if _gateway_routing: usage['gateway_routing'] = _gateway_routing + _ttft_ms = meter().get_ttft_ms(stream_id) + if _ttft_ms is not None: + usage['ttft_ms'] = _ttft_ms # Include context window data from the agent's compressor for the UI indicator. # The session-level persistence happens above (before s.save()) so the values # survive a page reload; this block only populates the live SSE usage payload. @@ -9986,7 +9993,7 @@ def _run_agent_streaming( _done_payload['terminal_reason'] = 'max_iterations' put('done', _done_payload) # Emit one last metering packet for the live message-header TPS label. - meter_stats = meter().get_stats() + meter_stats = meter().get_stats(stream_id) meter_stats['session_id'] = session_id meter_stats.setdefault('tps_available', False) meter_stats.setdefault('estimated', False) diff --git a/tests/test_issue6006_ttft_instrumentation.py b/tests/test_issue6006_ttft_instrumentation.py new file mode 100644 index 00000000..39de1194 --- /dev/null +++ b/tests/test_issue6006_ttft_instrumentation.py @@ -0,0 +1,104 @@ +import api.metering as metering + + +class _Clock: + def __init__(self): + self.wall = 1000.0 + self.mono = 10.0 + + def time(self): + return self.wall + + def monotonic(self): + return self.mono + + +def _meter(monkeypatch): + clock = _Clock() + monkeypatch.setattr(metering.time, 'time', clock.time) + monkeypatch.setattr(metering.time, 'monotonic', clock.monotonic) + return metering.GlobalMeter(), clock + + +def test_first_output_records_stream_ttft_ms(monkeypatch): + meter, clock = _meter(monkeypatch) + meter.begin_session('one') + meter.set_pending_started_at('one', 1000.0) + + clock.wall += 0.250 + meter.record_token('one', 1) + + assert meter.get_stats('one')['ttft_ms'] == 250 + + +def test_reasoning_before_output_does_not_start_ttft(monkeypatch): + meter, clock = _meter(monkeypatch) + meter.begin_session('one') + meter.set_pending_started_at('one', 1000.0) + + clock.wall += 0.100 + meter.record_reasoning('one', 1) + assert 'ttft_ms' not in meter.get_stats('one') + + clock.wall += 0.150 + meter.record_token('one', 1) + assert meter.get_stats('one')['ttft_ms'] == 250 + + clock.wall += 1.0 + meter.record_token('one', 2) + assert meter.get_stats('one')['ttft_ms'] == 250 + + +def test_missing_start_and_zero_token_omit_ttft(monkeypatch): + meter, clock = _meter(monkeypatch) + meter.begin_session('missing') + meter.set_pending_started_at('missing', None) + meter.record_token('missing', 1) + assert 'ttft_ms' not in meter.get_stats('missing') + + meter.begin_session('invalid') + meter.set_pending_started_at('invalid', 0) + meter.record_token('invalid', 1) + assert 'ttft_ms' not in meter.get_stats('invalid') + + meter.begin_session('zero') + meter.set_pending_started_at('zero', clock.wall) + assert 'ttft_ms' not in meter.get_stats('zero') + meter.end_session('zero', 0) + assert meter.get_stats('zero')['active'] == 2 + + +def test_concurrent_stream_ttft_isolation(monkeypatch): + meter, clock = _meter(monkeypatch) + meter.begin_session('one') + meter.begin_session('two') + meter.set_pending_started_at('one', 1000.0) + meter.set_pending_started_at('two', 1000.0) + + clock.wall += 0.125 + meter.record_token('one', 1) + clock.wall += 0.375 + meter.record_token('two', 1) + + assert meter.get_stats('one')['ttft_ms'] == 125 + assert meter.get_stats('two')['ttft_ms'] == 500 + + +def test_streaming_propagates_ttft_to_journal_done_and_message(monkeypatch): + meter, clock = _meter(monkeypatch) + stream_id = 'propagation' + meter.begin_session(stream_id) + meter.set_pending_started_at(stream_id, 1000.0) + clock.wall += 0.250 + meter.record_token(stream_id, 1) + + metering_event = meter.get_stats(stream_id) + journal = [metering_event] + usage = {'ttft_ms': metering_event['ttft_ms']} + assistant = {'role': 'assistant', '_firstTokenMs': meter.get_ttft_ms(stream_id)} + final_session = {'messages': [assistant]} + + assert journal[0]['ttft_ms'] == 250 + assert usage['ttft_ms'] == 250 + assert assistant['_firstTokenMs'] == 250 + assert final_session['messages'][-1]['_firstTokenMs'] == 250