feat: seed and periodically persist model usage stats to Postgres

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

backend/app/shared/bootstrap.py

@@ -2,9 +2,12 @@
 
 from __future__ import annotations
 
+import asyncio
 from functools import lru_cache
 from typing import Callable
 
+from loguru import logger
+
 from app.application.agent import AgentConversationService, AgentSessionService
 from app.application.agent.agentic_service import AgenticConversationService
 from app.application.documents import DocumentCommandService, DocumentQueryService
@@ -36,6 +39,7 @@ from app.infrastructure.storage.minio_binary_store import MinioDocumentBinarySto
 from app.infrastructure.storage.postgres_document_processing_store import PostgresDocumentProcessingStore
 from app.infrastructure.storage.postgres_document_repository import PostgresDocumentRepository
 from app.infrastructure.storage.postgres_parse_artifact_store import PostgresParseArtifactStore
+from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
 from app.infrastructure.vectorstore.bm25_retriever import BM25Retriever
 from app.infrastructure.vectorstore.cross_encoder_reranker import OpenAICompatibleReranker
 from app.infrastructure.vectorstore.dense_retriever import DenseRetriever
@@ -43,6 +47,7 @@ from app.infrastructure.vectorstore.milvus_vector_index import MilvusVectorIndex
 from app.services.llm.llm_factory import LLMFactory
 from app.domain.compliance.ports import ComplianceRepository
 from app.infrastructure.compliance.repository import PostgresComplianceRepository
+from app.shared.model_usage_tracker import get_model_usage_tracker
 # Keep shared wiring centralized so dependency construction remains consistent.
 
 
@@ -162,6 +167,14 @@ def get_parse_artifact_store():
     return None
 
 
+@lru_cache
+def get_model_usage_store():
+    """Return the Postgres model-usage store, or None when postgres backend is not enabled."""
+    if settings.document_repository_backend == "postgres":
+        return PostgresModelUsageStore()
+    return None
+
+
 @lru_cache
 def get_document_processing_store():
     """Return document processing store for the active repository backend."""
@@ -412,8 +425,67 @@ def get_user_store():
 def preload_runtime_dependencies() -> None:
     """Warm dependencies that are safe and useful to preload during startup."""
     LLMFactory.preload_clients(["qwen", "deepseek"])
+    _start_model_usage_persistence()
 
 
 def cleanup_runtime_dependencies() -> None:
     """Release runtime dependencies that expose explicit cleanup hooks."""
     LLMFactory.cleanup()
+    _stop_model_usage_persistence()
+
+
+_model_usage_flush_task: "asyncio.Task | None" = None
+
+
+def _start_model_usage_persistence() -> None:
+    """Seed ModelUsageTracker from Postgres and start its periodic flush loop.
+
+    No-op when document_repository_backend != "postgres" — ModelUsageTracker
+    then keeps behaving exactly as it always has: purely in-memory, reset on
+    every restart. Never raises: persistence must not block app startup.
+    """
+    global _model_usage_flush_task
+    try:
+        store = get_model_usage_store()
+    except Exception as exc:  # noqa: BLE001 - persistence must never block startup
+        logger.warning("Failed to initialize model usage persistence: {}", exc)
+        return
+    if store is None:
+        return
+
+    tracker = get_model_usage_tracker()
+    try:
+        tracker.seed(store.load_all())
+    except Exception as exc:  # noqa: BLE001 - a bad load must not block startup
+        logger.warning("Failed to load persisted model usage stats: {}", exc)
+
+    async def _flush_loop() -> None:
+        """Snapshot the tracker into Postgres every 60 seconds until cancelled."""
+        while True:
+            await asyncio.sleep(60)
+            try:
+                store.flush(tracker.snapshot())
+            except Exception as exc:  # noqa: BLE001 - one bad cycle must not kill the loop
+                logger.warning("Failed to flush model usage stats: {}", exc)
+
+    _model_usage_flush_task = asyncio.create_task(_flush_loop())
+
+
+def _stop_model_usage_persistence() -> None:
+    """Cancel the periodic flush task and perform one best-effort final flush."""
+    global _model_usage_flush_task
+    if _model_usage_flush_task is not None:
+        _model_usage_flush_task.cancel()
+        _model_usage_flush_task = None
+
+    try:
+        store = get_model_usage_store()
+    except Exception as exc:  # noqa: BLE001 - shutdown must not crash on this
+        logger.warning("Failed to access model usage store during shutdown: {}", exc)
+        return
+    if store is None:
+        return
+    try:
+        store.flush(get_model_usage_tracker().snapshot())
+    except Exception as exc:  # noqa: BLE001 - shutdown must not crash on a flush failure
+        logger.warning("Failed final model usage flush: {}", exc)

backend/tests/observability/test_model_usage_bootstrap.py

@@ -0,0 +1,111 @@
+"""Unit tests for the model-usage persistence wiring in app.shared.bootstrap.
+
+get_model_usage_store()'s settings-gating is tested the same way
+tests/test_reranker_bootstrap.py tests get_reranker() — by patching
+"app.shared.bootstrap.settings" wholesale, matching this codebase's
+established convention for testing @lru_cache settings-gated factories.
+The remaining tests isolate _start_model_usage_persistence() /
+_stop_model_usage_persistence() from get_model_usage_store() entirely (via
+monkeypatch on the module-level function), so no real database or event loop
+is needed anywhere in this file — asyncio.create_task itself is also mocked.
+"""
+
+from __future__ import annotations
+
+import sys
+from unittest.mock import MagicMock, patch
+
+# Patch psycopg2 before importing anything that transitively imports it, in
+# case this file is collected before test_model_usage_persistence.py.
+mock_psycopg2 = MagicMock()
+mock_psycopg2.extras = MagicMock()
+sys.modules.setdefault("psycopg2", mock_psycopg2)
+sys.modules.setdefault("psycopg2.extras", mock_psycopg2.extras)
+sys.modules.setdefault("psycopg2.pool", MagicMock())
+
+from app.shared import bootstrap
+from app.shared.model_usage_tracker import ModelUsageEntry, ModelUsageTracker
+
+
+def test_get_model_usage_store_returns_none_when_not_postgres_backend():
+    """get_model_usage_store() must be None unless document_repository_backend == 'postgres'."""
+    bootstrap.get_model_usage_store.cache_clear()
+
+    with patch("app.shared.bootstrap.settings") as mock_settings:
+        mock_settings.document_repository_backend = "json"
+        result = bootstrap.get_model_usage_store()
+
+    bootstrap.get_model_usage_store.cache_clear()
+    assert result is None
+
+
+def test_get_model_usage_store_returns_instance_when_postgres_backend():
+    """get_model_usage_store() must return a PostgresModelUsageStore when enabled.
+
+    ThreadedConnectionPool is mocked so no real connection is attempted; the
+    postgres_host/port/user/password/db values PostgresModelUsageStore reads
+    come from app.config.settings.settings directly (not from the
+    app.shared.bootstrap.settings reference mocked below), so they don't need
+    to be set here — only document_repository_backend gates this factory.
+    """
+    bootstrap.get_model_usage_store.cache_clear()
+
+    with patch("psycopg2.pool.ThreadedConnectionPool"), \
+         patch(
+             "app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema"
+         ), \
+         patch("app.shared.bootstrap.settings") as mock_settings:
+        mock_settings.document_repository_backend = "postgres"
+        result = bootstrap.get_model_usage_store()
+
+    bootstrap.get_model_usage_store.cache_clear()
+    from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
+    assert isinstance(result, PostgresModelUsageStore)
+
+
+def test_start_model_usage_persistence_seeds_tracker_and_starts_flush_loop(monkeypatch):
+    """When a store is available, startup must seed the tracker and schedule the flush task."""
+    fake_store = MagicMock()
+    fake_store.load_all.return_value = {
+        "deepseek:deepseek-v4-flash": ModelUsageEntry(
+            provider="deepseek", model="deepseek-v4-flash", total_tokens=99,
+        ),
+    }
+    tracker = ModelUsageTracker()
+    monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: fake_store)
+    monkeypatch.setattr(bootstrap, "get_model_usage_tracker", lambda: tracker)
+
+    with patch("asyncio.create_task") as mock_create_task:
+        bootstrap._start_model_usage_persistence()
+        # Close the coroutine object passed to the mock so pytest doesn't warn
+        # about "coroutine was never awaited" — it was never meant to run here.
+        mock_create_task.call_args[0][0].close()
+
+    assert tracker.get("deepseek", "deepseek-v4-flash").total_tokens == 99
+    mock_create_task.assert_called_once()
+
+    bootstrap._stop_model_usage_persistence()  # reset the module-level task handle
+
+
+def test_start_model_usage_persistence_is_a_no_op_without_a_store(monkeypatch):
+    """No store configured (json backend) — startup must not touch asyncio or the tracker."""
+    monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: None)
+
+    with patch("asyncio.create_task") as mock_create_task:
+        bootstrap._start_model_usage_persistence()
+
+    mock_create_task.assert_not_called()
+
+
+def test_stop_model_usage_persistence_cancels_task_and_flushes(monkeypatch):
+    """Shutdown must cancel the running flush task and perform one final flush."""
+    fake_store = MagicMock()
+    monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: fake_store)
+    fake_task = MagicMock()
+    bootstrap._model_usage_flush_task = fake_task
+
+    bootstrap._stop_model_usage_persistence()
+
+    fake_task.cancel.assert_called_once()
+    fake_store.flush.assert_called_once()
+    assert bootstrap._model_usage_flush_task is None