From c2d1893992386a8f8b86f7fd61ba30fd1567e53d Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Tue, 22 Sep 2026 10:40:30 +0100 Subject: [PATCH] fix(admin): keyset export paging, rollup-free spend, consistent totals Review pass over the preceding commits. - The activity export paged by OFFSET. The journals are append-only and the feed is newest-first, so rows written mid-export shift the window down and repeat rows already emitted. Pages by keyset on (created_at, feed, id) now. - top_token_users and the per-user breakdown counted run-level rollup rows. A scheduled run already has a row per LLM call, so its spend was billed twice -- invisible while the column was tokens, obvious once it was dollars. Both now exclude them, matching every other spend query. - total_cost was summed from the per-model split, which drops rows with no model_id and so undercounted the figure printed above the series. Summed from the series instead. - The CSV export encoded its detail cell without the fallback the NDJSON branch had, so a non-JSON-native value would have failed the stream. - The activity view reset pagination in an effect, which fetched the stale page against the new filters before fetching again. Reset in the setters. - The audit taxonomy is no longer re-exported from the helper module; the one caller that wanted it imports from where it lives. --- docsgpt/api/admin/activity.py | 10 ++-- docsgpt/api/admin/routes.py | 4 +- docsgpt/api/audit.py | 7 +-- docsgpt/api/user/sources/upload.py | 12 ++-- docsgpt/storage/db/repositories/activity.py | 58 ++++++++++++++----- .../storage/db/repositories/admin_stats.py | 27 ++++++++- frontend/src/admin/Activity.tsx | 26 ++++++--- tests/api/test_audit_helper.py | 3 +- .../handlers/test_round_usage_persistence.py | 13 +++-- 9 files changed, 113 insertions(+), 47 deletions(-) diff --git a/docsgpt/api/admin/activity.py b/docsgpt/api/admin/activity.py index 832b2686..f2f7bd1e 100644 --- a/docsgpt/api/admin/activity.py +++ b/docsgpt/api/admin/activity.py @@ -43,7 +43,7 @@ _EXPORT_CHUNK = 1_000 _MAX_SEARCH_LENGTH = 200 -def _csv_list(name: str, allowed: tuple[str, ...]) -> Optional[list[str]]: +def _facet_list(name: str, allowed: tuple[str, ...]) -> Optional[list[str]]: """Parse a repeated/comma-separated query arg, dropping unknown values. Args: @@ -96,8 +96,8 @@ def _search_arg() -> Optional[str]: def _filters() -> dict: """The filter set shared by the feed and the export.""" return { - "feeds": _csv_list("feed", _FEEDS), - "categories": _csv_list("category", ACTIVITY_CATEGORIES), + "feeds": _facet_list("feed", _FEEDS), + "categories": _facet_list("category", ACTIVITY_CATEGORIES), "events": _event_list(), "actor_id": request.args.get("actor_id") or None, "user_id": request.args.get("user_id") or None, @@ -216,7 +216,9 @@ def _csv_rows(filters: dict, limit: int) -> Iterator[str]: for row in _export_rows(filters, limit): serialized = _serialize(row) # ``detail`` is a JSON object; a CSV cell holds its compact encoding. - serialized["detail"] = json.dumps(serialized.get("detail") or {}) + serialized["detail"] = json.dumps( + serialized.get("detail") or {}, default=str + ) writer.writerow(serialized) yield _drain(buffer) diff --git a/docsgpt/api/admin/routes.py b/docsgpt/api/admin/routes.py index ff51149d..318c6b76 100644 --- a/docsgpt/api/admin/routes.py +++ b/docsgpt/api/admin/routes.py @@ -359,7 +359,9 @@ class AdminUsageResource(Resource): "group_by": group_by, "series": series, "total_tokens": int(total), - "total_cost": round(sum(row["cost"] for row in by_model), 4), + # Summed from the series, not from ``by_model``: the latter + # drops rows with no model_id, so it would undercount. + "total_cost": round(sum(row["cost"] for row in series), 4), "by_model": by_model, "latency": latency, "top_users": top_users, diff --git a/docsgpt/api/audit.py b/docsgpt/api/audit.py index 0ca37827..26c6bd4a 100644 --- a/docsgpt/api/audit.py +++ b/docsgpt/api/audit.py @@ -14,6 +14,9 @@ Two properties matter and are enforced here rather than at every call site: 2. **Request context is optional.** Celery tasks (ingestion finishing, a scheduled run) have no Flask request; the row simply records without an IP or user agent instead of raising. + +The event → category taxonomy the admin activity feed filters on lives in +``docsgpt/audit_events.py``, which the storage layer needs too. """ from __future__ import annotations @@ -23,14 +26,10 @@ from typing import Any, Optional from flask import has_request_context, request -from docsgpt.audit_events import ACTIVITY_CATEGORIES, category_for from docsgpt.storage.db.repositories.auth_events import AuthEventsRepository logger = logging.getLogger(__name__) -# Re-exported so route modules have one import for recording and classifying. -__all__ = ["ACTIVITY_CATEGORIES", "category_for", "record_event"] - def record_event( conn, diff --git a/docsgpt/api/user/sources/upload.py b/docsgpt/api/user/sources/upload.py index b052b300..31f2d700 100644 --- a/docsgpt/api/user/sources/upload.py +++ b/docsgpt/api/user/sources/upload.py @@ -397,10 +397,6 @@ class UploadFile(Resource): return make_response(jsonify({"success": False}), 400) # Predetermined id matches the dedup-claim row; loser GET sees same. response_task_id = predetermined_task_id or task.id - # ``source_uuid`` was minted above and passed to the worker as - # ``source_id``; the worker uses it verbatim for every SSE event, - # so the frontend can correlate inbound ``source.ingest.*`` to - # this upload regardless of whether an idempotency key was set. # Audited here, not in the worker: this is the user action. The # ingest may still fail, which the source's own status records. try: @@ -416,9 +412,15 @@ class UploadFile(Resource): ) except Exception as err: current_app.logger.warning( - "Could not audit upload for source %s: %s", source_uuid, err, + "Could not audit upload for source %s: %s", + source_uuid, + err, exc_info=True, ) + # ``source_uuid`` was minted above and passed to the worker as + # ``source_id``; the worker uses it verbatim for every SSE event, + # so the frontend can correlate inbound ``source.ingest.*`` to + # this upload regardless of whether an idempotency key was set. response_payload: dict = { "success": True, "task_id": response_task_id, diff --git a/docsgpt/storage/db/repositories/activity.py b/docsgpt/storage/db/repositories/activity.py index be49b1e4..691f2c7e 100644 --- a/docsgpt/storage/db/repositories/activity.py +++ b/docsgpt/storage/db/repositories/activity.py @@ -46,6 +46,10 @@ ACTIVITY_COLUMNS = ( # the category filter cannot match it. _FEED_CATEGORIES = {"auth": None, "device": "device", "guardrail": "safety"} +# ``feed`` and ``id`` break ties so a page boundary is deterministic and the +# keyset cursor below has something total to compare against. +_ORDER_BY = "ORDER BY created_at DESC, feed DESC, id DESC" + _AUTH_BRANCH = f""" SELECT 'auth' AS feed, id::text AS id, @@ -158,14 +162,14 @@ class ActivityRepository: def _query( self, *, - feeds: Optional[Sequence[str]], - categories: Optional[Sequence[str]], - events: Optional[Sequence[str]], - actor_id: Optional[str], - user_id: Optional[str], - since: Optional[datetime], - until: Optional[datetime], - search: Optional[str], + feeds: Optional[Sequence[str]] = None, + categories: Optional[Sequence[str]] = None, + events: Optional[Sequence[str]] = None, + actor_id: Optional[str] = None, + user_id: Optional[str] = None, + since: Optional[datetime] = None, + until: Optional[datetime] = None, + search: Optional[str] = None, ) -> tuple[str, dict]: selected = self._feeds_for(feeds, categories) if not selected: @@ -247,11 +251,33 @@ class ActivityRepository: return [] params.update({"limit": int(limit), "offset": int(offset)}) result = self._conn.execute( - text(f"{sql} ORDER BY created_at DESC, feed, id DESC LIMIT :limit OFFSET :offset"), + text(f"{sql} {_ORDER_BY} LIMIT :limit OFFSET :offset"), params, ) return [dict(row._mapping) for row in result.fetchall()] + def _page_after( + self, cursor: Optional[tuple], limit: int, filters: dict + ) -> list[dict]: + """One keyset page strictly older than ``cursor``. + + OFFSET paging would drift here: the feed is ordered newest-first and + the journals are append-only, so rows inserted mid-export push the + window down and repeat rows that were already emitted. + """ + sql, params = self._query(**filters) + if not sql: + return [] + if cursor is not None: + sql += " AND " if "WHERE" in sql else " WHERE " + sql += "(created_at, feed, id) < (:cur_at, :cur_feed, :cur_id)" + params.update( + {"cur_at": cursor[0], "cur_feed": cursor[1], "cur_id": cursor[2]} + ) + params["limit"] = int(limit) + result = self._conn.execute(text(f"{sql} {_ORDER_BY} LIMIT :limit"), params) + return [dict(row._mapping) for row in result.fetchall()] + def count( self, *, @@ -291,8 +317,9 @@ class ActivityRepository: ) -> Iterator[dict]: """Yield every matching row, newest first, for the export endpoint. - Pages through the feed rather than materializing it, so an export of a - large instance streams instead of building the whole result in memory. + Pages by keyset, not offset: the journals are append-only and the feed + is newest-first, so rows written mid-export would otherwise shift the + window and repeat rows already emitted. Args: chunk_size: Rows fetched per round trip. @@ -304,17 +331,18 @@ class ActivityRepository: Row dicts carrying :data:`ACTIVITY_COLUMNS`. """ emitted = 0 - offset = 0 + cursor: Optional[tuple] = None while emitted < max_rows: - page = self.list( - limit=min(chunk_size, max_rows - emitted), offset=offset, **filters + page = self._page_after( + cursor, min(chunk_size, max_rows - emitted), filters ) if not page: return for row in page: yield row + last = page[-1] + cursor = (last["created_at"], last["feed"], last["id"]) emitted += len(page) - offset += len(page) def event_names(self) -> list[dict]: """Distinct ``(event, category)`` pairs across all three journals. diff --git a/docsgpt/storage/db/repositories/admin_stats.py b/docsgpt/storage/db/repositories/admin_stats.py index 23fdd3d9..bd4241f9 100644 --- a/docsgpt/storage/db/repositories/admin_stats.py +++ b/docsgpt/storage/db/repositories/admin_stats.py @@ -14,6 +14,7 @@ from typing import Optional from sqlalchemy import Connection, text from docsgpt.storage.db.base_repository import row_to_dict +from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository def _round_ms(value) -> Optional[int]: @@ -76,6 +77,10 @@ class AdminStatsRepository: Ordered by tokens, not cost: a cheap model can dominate token volume while a costly one dominates the bill, and the operator wants to see both columns rather than have the ranking pick for them. + + Run-level rollup rows are excluded, matching every other spend query: + a scheduled run already has a row per LLM call, so counting the rollup + too would bill it twice. """ result = self._conn.execute( text( @@ -85,12 +90,17 @@ class AdminStatsRepository: COALESCE(SUM(cost), 0) AS cost FROM token_usage WHERE timestamp >= :since AND user_id IS NOT NULL + AND source <> ALL(:rollup_sources) GROUP BY user_id ORDER BY tokens DESC LIMIT :limit """ ), - {"since": since, "limit": int(limit)}, + { + "since": since, + "limit": int(limit), + "rollup_sources": list(TokenUsageRepository.ROLLUP_SOURCES), + }, ) return [ {"user_id": r[0], "tokens": int(r[1]), "cost": float(r[2])} @@ -140,6 +150,9 @@ class AdminStatsRepository: cannot answer either question an operator actually has -- what is this person costing, and what is driving it. + Run-level rollup rows are excluded, so a scheduled run is counted once + (by its per-call rows), not twice. + Args: user_id: The billable user (auth ``sub``). since: Inclusive window start. @@ -148,6 +161,7 @@ class AdminStatsRepository: Returns: ``{"totals": {...}, "by_model": [...], "by_source": [...]}``. """ + rollups = list(TokenUsageRepository.ROLLUP_SOURCES) totals = self._conn.execute( text( """ @@ -156,9 +170,10 @@ class AdminStatsRepository: count(*) AS calls FROM token_usage WHERE user_id = :u AND timestamp >= :since + AND source <> ALL(:rollup_sources) """ ), - {"u": user_id, "since": since}, + {"u": user_id, "since": since, "rollup_sources": rollups}, ).one() def _breakdown(key_expr: str) -> list[dict]: @@ -170,12 +185,18 @@ class AdminStatsRepository: COALESCE(SUM(cost), 0) AS cost FROM token_usage WHERE user_id = :u AND timestamp >= :since + AND source <> ALL(:rollup_sources) GROUP BY key ORDER BY tokens DESC, key LIMIT :limit """ ), - {"u": user_id, "since": since, "limit": int(limit)}, + { + "u": user_id, + "since": since, + "limit": int(limit), + "rollup_sources": rollups, + }, ) return [ {"key": r.key, "tokens": int(r.tokens), "cost": float(r.cost)} diff --git a/frontend/src/admin/Activity.tsx b/frontend/src/admin/Activity.tsx index 0cb4ce64..65c0bdb7 100644 --- a/frontend/src/admin/Activity.tsx +++ b/frontend/src/admin/Activity.tsx @@ -137,6 +137,18 @@ export default function Activity() { const [searchDraft, setSearchDraft] = useState(''); const [search, setSearch] = useState(''); + /** + * Narrowing the feed invalidates whatever page you were on, so every filter + * setter resets it. Resetting in an effect instead would fetch the old page + * against the new filters first, then fetch again. + */ + const withPageReset = + (set: (value: T) => void) => + (value: T) => { + setPage(1); + set(value); + }; + const filters: ActivityFilters = useMemo( () => ({ category: categories.length ? categories : undefined, @@ -180,11 +192,6 @@ export default function Activity() { load(); }, [load]); - // Any filter change invalidates the current page number. - useEffect(() => { - setPage(1); - }, [filters]); - const exportAs = async (format: 'csv' | 'ndjson') => { setExporting(true); try { @@ -225,7 +232,7 @@ export default function Activity() { const hasFilters = categories.length > 0 || events.length > 0 || search !== ''; - const applySearch = () => setSearch(searchDraft.trim()); + const applySearch = () => withPageReset(setSearch)(searchDraft.trim()); return (
@@ -241,14 +248,14 @@ export default function Activity() { setRangeDays(range.days)} + onClick={() => withPageReset(setRangeDays)(range.days)} > {range.label} @@ -271,6 +278,7 @@ export default function Activity() { variant="ghost" size="sm" onClick={() => { + setPage(1); setCategories([]); setEvents([]); setSearchDraft(''); diff --git a/tests/api/test_audit_helper.py b/tests/api/test_audit_helper.py index f017a7ac..d51a584c 100644 --- a/tests/api/test_audit_helper.py +++ b/tests/api/test_audit_helper.py @@ -6,7 +6,8 @@ from unittest.mock import MagicMock, Mock, patch import pytest -from docsgpt.api.audit import ACTIVITY_CATEGORIES, category_for, record_event +from docsgpt.api.audit import record_event +from docsgpt.audit_events import ACTIVITY_CATEGORIES, category_for @pytest.fixture diff --git a/tests/llm/handlers/test_round_usage_persistence.py b/tests/llm/handlers/test_round_usage_persistence.py index b7f1636f..e28409a9 100644 --- a/tests/llm/handlers/test_round_usage_persistence.py +++ b/tests/llm/handlers/test_round_usage_persistence.py @@ -167,11 +167,11 @@ class TestPerRoundUsageRows: return messages, None rows = [] - with patch.object( - usage_mod, - "_persist_call_usage", - side_effect=lambda llm_, cu: rows.append(dict(cu)), - ): + + def _record(llm_, cu, *, duration_ms=None, ttft_ms=None): + rows.append({**cu, "duration_ms": duration_ms, "ttft_ms": ttft_ms}) + + with patch.object(usage_mod, "_persist_call_usage", side_effect=_record): with patch.object(handler, "handle_tool_calls", fake_tool_calls): first = llm.gen_stream( model="m", messages=[{"role": "user", "content": "hi"}] @@ -185,6 +185,9 @@ class TestPerRoundUsageRows: counts = lambda row: {k: row[k] for k in ("prompt_tokens", "generated_tokens")} # noqa: E731 assert counts(rows[0]) == {"prompt_tokens": 100, "generated_tokens": 10} assert counts(rows[1]) == {"prompt_tokens": 200, "generated_tokens": 20} + # Each round is timed independently, and each streamed a chunk. + assert all(row["duration_ms"] is not None for row in rows) + assert all(row["ttft_ms"] is not None for row in rows) class TestPreferProviderUsageClaim: