feat(usage): persist LLM call latency

The wrappers in docsgpt/usage.py already measured how long each call took --
for the llm_gen_finished / llm_stream_finished log lines -- and then threw the
number away. Nothing in the schema recorded it, so an operator could see what
an instance spent but never how slow it was.

Adds token_usage.duration_ms and token_usage.ttft_ms, and threads the
measurements the wrappers already take into the insert. The stream wrapper now
also stamps the moment of the first yielded chunk.

Both columns are nullable, and deliberately so: ttft_ms is NULL for a
non-streaming call and for a stream that failed before yielding anything, and
every row written before this migration is NULL too. A 0 there would drag a
p50 toward an instant first token that never happened.

docsgpt/alembic/versions/0035_token_usage_latency.py

@@ -0,0 +1,42 @@
+"""0035 token usage latency — how long the call took, and time to first token.
+
+Nothing in the schema recorded how long an LLM call took. The wrappers in
+``docsgpt/usage.py`` already measured it for the ``llm_gen_finished`` /
+``llm_stream_finished`` log lines and then threw the number away, so an
+operator could see what an instance spent but never how slow it was.
+
+Two columns, both nullable because both are genuinely unknown for some rows:
+
+* ``duration_ms`` — wall-clock for the whole call.
+* ``ttft_ms`` — time to the first streamed chunk. NULL on non-streaming calls
+  and on streams that failed before yielding anything, which is the honest
+  reading: "no first token", not "instant".
+
+NULL also means "written before this migration", so historical rows never
+drag a p50 towards zero.
+
+Idempotent both ways.
+
+Revision ID: 0035_token_usage_latency
+Revises: 0034_auth_events_actor_target
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+
+revision: str = "0035_token_usage_latency"
+down_revision: Union[str, None] = "0034_auth_events_actor_target"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.execute("ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS duration_ms INTEGER;")
+    op.execute("ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS ttft_ms INTEGER;")
+
+
+def downgrade() -> None:
+    op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS ttft_ms;")
+    op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS duration_ms;")

docsgpt/storage/db/models.py

@@ -250,6 +250,11 @@ token_usage_table = Table(
     # Added in ``0033_quotas``. USD cost of the call at write time; 0 for
     # unpriced and bring-your-own models.
     Column("cost", Numeric(12, 8), nullable=False, server_default="0"),
+    # Added in ``0035_token_usage_latency``. Wall-clock for the call, and time
+    # to the first streamed chunk. ``ttft_ms`` is NULL for non-streaming calls
+    # and for streams that failed before yielding -- "no first token", not 0.
+    Column("duration_ms", Integer),
+    Column("ttft_ms", Integer),
 )
 
 user_logs_table = Table(

docsgpt/storage/db/repositories/token_usage.py

@@ -38,6 +38,8 @@ class TokenUsageRepository:
         cached_tokens: Optional[int] = None,
         cache_write_tokens: Optional[int] = None,
         cost: float = 0.0,
+        duration_ms: Optional[int] = None,
+        ttft_ms: Optional[int] = None,
     ) -> None:
         # Attribution guard: the ``token_usage_attribution_chk`` CHECK
         # constraint requires at least one of ``user_id`` / ``api_key``
@@ -64,6 +66,7 @@ class TokenUsageRepository:
                     user_id, api_key, agent_id,
                     prompt_tokens, generated_tokens,
                     cached_tokens, cache_write_tokens, cost,
+                    duration_ms, ttft_ms,
                     source, request_id, model_id, timestamp
                 )
                 VALUES (
@@ -71,6 +74,7 @@ class TokenUsageRepository:
                     CAST(:agent_id AS uuid),
                     :prompt_tokens, :generated_tokens,
                     :cached_tokens, :cache_write_tokens, :cost,
+                    :duration_ms, :ttft_ms,
                     :source, :request_id, :model_id, COALESCE(:timestamp, now())
                 )
                 """
@@ -84,6 +88,8 @@ class TokenUsageRepository:
                 "cached_tokens": cached_tokens,
                 "cache_write_tokens": cache_write_tokens,
                 "cost": cost,
+                "duration_ms": duration_ms,
+                "ttft_ms": ttft_ms,
                 "source": source,
                 "request_id": request_id,
                 "model_id": model_id,

docsgpt/usage.py

@@ -93,7 +93,7 @@ def _count_prompt_tokens(messages, tools=None, usage_attachments=None, **kwargs)
     return prompt_tokens
 
 
-def _persist_call_usage(llm, call_usage):
+def _persist_call_usage(llm, call_usage, *, duration_ms=None, ttft_ms=None):
     """Write one ``token_usage`` row per LLM call. Always-on; no flag.
 
     Source defaults to ``agent_stream`` and can be overridden per
@@ -101,6 +101,13 @@ def _persist_call_usage(llm, call_usage):
     title / compression / rag_condense / fallback). A ``_request_id``
     stamped on the LLM lets ``count_in_range`` deduplicate the multiple
     rows produced by a single multi-tool agent run.
+
+    Args:
+        llm: The LLM instance the call ran on.
+        call_usage: The call's token counts.
+        duration_ms: Wall-clock for the call, measured by the wrapper.
+        ttft_ms: Time to the first streamed chunk; None for a non-streaming
+            call and for a stream that failed before yielding anything.
     """
     if call_usage["prompt_tokens"] == 0 and call_usage["generated_tokens"] == 0:
         return
@@ -149,6 +156,8 @@ def _persist_call_usage(llm, call_usage):
                 ),
                 request_id=getattr(llm, "_request_id", None),
                 model_id=model_id,
+                duration_ms=duration_ms,
+                ttft_ms=ttft_ms,
             )
     except Exception:
         logger.exception("token_usage persist failed")
@@ -258,10 +267,12 @@ def gen_token_usage(func):
             error = exc
             raise
         finally:
+            duration_ms = int((time.monotonic() - started_at) * 1000)
             call_usage = _prefer_provider_usage(self, call_usage)
             self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"]
             self.token_usage["generated_tokens"] += call_usage["generated_tokens"]
-            _persist_call_usage(self, call_usage)
+            # A non-streaming call has no first-token moment; ttft stays NULL.
+            _persist_call_usage(self, call_usage, duration_ms=duration_ms)
             emit = getattr(self, "_emit_gen_finished_log", None)
             if callable(emit):
                 try:
@@ -269,7 +280,7 @@ def gen_token_usage(func):
                         model,
                         prompt_tokens=call_usage["prompt_tokens"],
                         completion_tokens=call_usage["generated_tokens"],
-                        latency_ms=int((time.monotonic() - started_at) * 1000),
+                        latency_ms=duration_ms,
                         cached_tokens=call_usage.get("cached_tokens"),
                         cache_write_tokens=call_usage.get("cache_write_tokens"),
                         error=error,
@@ -293,10 +304,13 @@ def stream_token_usage(func):
         )
         batch = []
         started_at = time.monotonic()
+        first_chunk_at: float | None = None
         error: BaseException | None = None
         try:
             result = func(self, model, messages, stream, tools, **kwargs)
             for r in result:
+                if first_chunk_at is None:
+                    first_chunk_at = time.monotonic()
                 batch.append(r)
                 yield r
         except Exception as exc:
@@ -306,12 +320,22 @@ def stream_token_usage(func):
             error = exc
             raise
         finally:
+            duration_ms = int((time.monotonic() - started_at) * 1000)
+            # NULL, not 0, when the stream failed before yielding: "no first
+            # token" must not read as an instant one in a p50.
+            ttft_ms = (
+                int((first_chunk_at - started_at) * 1000)
+                if first_chunk_at is not None
+                else None
+            )
             for line in batch:
                 call_usage["generated_tokens"] += _count_tokens(line)
             call_usage = _prefer_provider_usage(self, call_usage)
             self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"]
             self.token_usage["generated_tokens"] += call_usage["generated_tokens"]
-            _persist_call_usage(self, call_usage)
+            _persist_call_usage(
+                self, call_usage, duration_ms=duration_ms, ttft_ms=ttft_ms
+            )
             emit = getattr(self, "_emit_stream_finished_log", None)
             if callable(emit):
                 try:
@@ -319,7 +343,7 @@ def stream_token_usage(func):
                         model,
                         prompt_tokens=call_usage["prompt_tokens"],
                         completion_tokens=call_usage["generated_tokens"],
-                        latency_ms=int((time.monotonic() - started_at) * 1000),
+                        latency_ms=duration_ms,
                         cached_tokens=call_usage.get("cached_tokens"),
                         cache_write_tokens=call_usage.get("cache_write_tokens"),
                         error=error,

tests/storage/db/test_migration_0035.py

@@ -0,0 +1,85 @@
+"""Migration round-trip test for 0035_token_usage_latency."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+from sqlalchemy import text
+
+
+pytestmark = pytest.mark.integration
+
+
+def _alembic_ini() -> Path:
+    return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini"
+
+
+def _run_alembic(url: str, *args: str) -> None:
+    subprocess.check_call(
+        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
+        timeout=60,
+        env={**os.environ, "POSTGRES_URI": url},
+    )
+
+
+def _alembic_version(conn) -> str:
+    return conn.execute(text("SELECT version_num FROM alembic_version")).scalar()
+
+
+def _column_exists(conn, table: str, column: str) -> bool:
+    row = conn.execute(
+        text(
+            "SELECT 1 FROM information_schema.columns "
+            "WHERE table_name = :t AND column_name = :c AND table_schema = 'public'"
+        ),
+        {"t": table, "c": column},
+    ).fetchone()
+    return row is not None
+
+
+_0035 = "0035_token_usage_latency"
+_0034 = "0034_auth_events_actor_target"
+
+
+class TestMigration0035RoundTrip:
+    def test_head_has_latency_columns(self, pg_engine):
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) >= _0035
+            assert _column_exists(conn, "token_usage", "duration_ms")
+            assert _column_exists(conn, "token_usage", "ttft_ms")
+
+    def test_downgrade_drops_then_upgrade_restores(self, pg_engine):
+        url = pg_engine.url.render_as_string(hide_password=False)
+        _run_alembic(url, "downgrade", _0034)
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) == _0034
+            assert not _column_exists(conn, "token_usage", "duration_ms")
+            assert not _column_exists(conn, "token_usage", "ttft_ms")
+        _run_alembic(url, "upgrade", "head")
+        with pg_engine.connect() as conn:
+            assert _column_exists(conn, "token_usage", "duration_ms")
+
+    def test_existing_rows_read_null_latency(self, pg_engine):
+        """Rows written before 0035 must read NULL, not 0 — a 0ms call is a lie."""
+        url = pg_engine.url.render_as_string(hide_password=False)
+        _run_alembic(url, "downgrade", _0034)
+        with pg_engine.begin() as conn:
+            conn.execute(
+                text(
+                    "INSERT INTO token_usage (user_id, prompt_tokens, generated_tokens) "
+                    "VALUES ('u-mig35', 10, 1)"
+                )
+            )
+        _run_alembic(url, "upgrade", "head")
+        with pg_engine.connect() as conn:
+            row = conn.execute(
+                text(
+                    "SELECT duration_ms, ttft_ms FROM token_usage "
+                    "WHERE user_id = 'u-mig35'"
+                )
+            ).fetchone()
+        assert tuple(row) == (None, None)

tests/test_usage_latency.py

@@ -0,0 +1,94 @@
+"""Latency is measured by the usage wrappers and persisted with the call.
+
+``duration_ms`` and ``ttft_ms`` were already computed for the finish log lines
+and then discarded; these pin that they now reach ``token_usage``.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+
+from docsgpt.usage import gen_token_usage, stream_token_usage
+
+
+class _LLM:
+    """Minimal stand-in for an LLM instance the wrappers decorate."""
+
+    def __init__(self):
+        self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
+        self.decoded_token = {"sub": "u1"}
+        self.user_api_key = None
+        self.agent_id = None
+
+
+@pytest.fixture
+def persisted():
+    calls: list[dict] = []
+
+    def _capture(llm, call_usage, *, duration_ms=None, ttft_ms=None):
+        calls.append({"duration_ms": duration_ms, "ttft_ms": ttft_ms})
+
+    with patch("docsgpt.usage._persist_call_usage", _capture):
+        yield calls
+
+
+@pytest.mark.unit
+class TestNonStreaming:
+    def test_records_a_duration_and_no_first_token(self, persisted):
+        @gen_token_usage
+        def _gen(self, model, messages, stream, tools, **kwargs):
+            return "hello"
+
+        assert _gen(_LLM(), "m", [], False, None) == "hello"
+        assert persisted[0]["duration_ms"] >= 0
+        # A non-streaming call has no first-token moment.
+        assert persisted[0]["ttft_ms"] is None
+
+    def test_a_failed_call_still_records_its_duration(self, persisted):
+        @gen_token_usage
+        def _gen(self, model, messages, stream, tools, **kwargs):
+            raise RuntimeError("upstream down")
+
+        with pytest.raises(RuntimeError):
+            _gen(_LLM(), "m", [], False, None)
+        assert persisted[0]["duration_ms"] >= 0
+
+
+@pytest.mark.unit
+class TestStreaming:
+    def test_records_time_to_first_chunk(self, persisted):
+        @stream_token_usage
+        def _stream(self, model, messages, stream, tools, **kwargs):
+            yield "a"
+            yield "b"
+
+        assert list(_stream(_LLM(), "m", [], True, None)) == ["a", "b"]
+        row = persisted[0]
+        assert row["ttft_ms"] is not None
+        # First token cannot land after the call finished.
+        assert row["ttft_ms"] <= row["duration_ms"]
+
+    def test_a_stream_that_never_yields_has_no_first_token(self, persisted):
+        """NULL, not 0 — an instant p50 would be a lie."""
+
+        @stream_token_usage
+        def _stream(self, model, messages, stream, tools, **kwargs):
+            raise RuntimeError("refused")
+            yield  # pragma: no cover - unreachable, makes this a generator
+
+        with pytest.raises(RuntimeError):
+            list(_stream(_LLM(), "m", [], True, None))
+        assert persisted[0]["ttft_ms"] is None
+        assert persisted[0]["duration_ms"] >= 0
+
+    def test_a_stream_cut_short_keeps_the_first_token_it_saw(self, persisted):
+        @stream_token_usage
+        def _stream(self, model, messages, stream, tools, **kwargs):
+            yield "a"
+            raise RuntimeError("dropped")
+
+        with pytest.raises(RuntimeError):
+            list(_stream(_LLM(), "m", [], True, None))
+        assert persisted[0]["ttft_ms"] is not None