Coverage for src/local_deep_research/database/thread_metrics.py: 98%
66 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Thread-safe metrics database access.
4This module provides a way for background threads to write metrics
5to the user's encrypted database by creating thread-local connections
6with the provided password.
7"""
9import threading
10from contextlib import contextmanager
11from typing import Optional
13from loguru import logger
14from sqlalchemy.orm import Session
16from .encrypted_db import db_manager
17from .rekey_lock import RekeyOperationalError
18from ..utilities.request_context import get_current_username # noqa: F401
21class ThreadSafeMetricsWriter:
22 """
23 Thread-safe writer for metrics to encrypted user databases.
24 Creates encrypted connections per thread using provided passwords.
25 """
27 def __init__(self):
28 self._thread_local = threading.local()
30 def set_user_password(self, username: str, password: str):
31 """
32 Store user password for the current thread.
33 This allows the thread to create its own encrypted connection.
35 IMPORTANT: This is safe because:
36 1. Password is already in memory (user is logged in)
37 2. It's only stored thread-locally
38 3. It's explicitly cleared by `clear_passwords()` (invoked from
39 `cleanup_current_thread()`) so pooled worker threads do not
40 retain credentials between tasks.
41 """
43 if not hasattr(self._thread_local, "passwords"):
44 self._thread_local.passwords = {}
45 self._thread_local.passwords[username] = password
47 def clear_passwords(self):
48 """Remove all passwords cached on the current thread's local store.
50 Called by `cleanup_current_thread()` so that pooled worker threads
51 don't retain plaintext credentials between tasks.
52 """
53 if hasattr(self._thread_local, "passwords"):
54 self._thread_local.passwords.clear()
56 @contextmanager
57 def get_session(self, username: str = None) -> Session:
58 """
59 Get a database session for metrics in the current thread.
60 Creates a new encrypted connection if needed.
62 Args:
63 username: The username for database access. If not provided,
64 will attempt to get it from Flask session.
65 """
66 # If username not provided, get it from request context.
67 # ValueError is the right exception type here — this is a database
68 # service called from threadpools, schedulers, and worker threads,
69 # not from an HTTP handler. werkzeug.exceptions.Unauthorized used
70 # to be raised here, but its 401-meaning is meaningless outside an
71 # HTTP context, and the surrounding try/except (which caught
72 # ImportError/RuntimeError) was structurally broken — the import
73 # never raised the exceptions the except clause caught.
74 if username is None:
75 username = get_current_username()
76 if not username: 76 ↛ 80line 76 didn't jump to line 80 because the condition on line 76 was always true
77 raise ValueError("No username in request context")
79 # Get password for this user in this thread
80 if not hasattr(self._thread_local, "passwords"):
81 raise ValueError("No password set for thread metrics access")
83 password = self._thread_local.passwords.get(username)
85 if not password:
86 raise ValueError(
87 f"No password available for user {username} in this thread"
88 )
90 # Create a thread-safe session for this user
91 session = None
92 try:
93 try:
94 session = db_manager.create_thread_safe_session_for_metrics(
95 username, password
96 )
97 except RekeyOperationalError as error:
98 # The database is unavailable for an operational reason (a
99 # rekey sidecar or epoch fault, or pending rekey recovery).
100 # This writer's callers do not know the typed error, so it
101 # becomes the same "no session" failure as below, raised
102 # outside this block so the fault is not chained onto it.
103 # Type name only: ``password`` is a live local in this frame.
104 logger.warning(
105 f"Metrics session unavailable for user {username}: "
106 f"{type(error).__name__}"
107 )
108 if not session:
109 raise ValueError( # noqa: TRY301 — except does session rollback before re-raise
110 f"Failed to create session for user {username}"
111 )
112 yield session
113 session.commit()
114 except Exception as exc:
115 # ``logger.warning`` (no traceback) rather than
116 # ``logger.exception``: ``password`` is a live local in this
117 # frame, so rendering a traceback under ``diagnose=True``
118 # would dump the plaintext SQLCipher master password
119 # (unrecoverable — TRUST.md §5). The exception still
120 # propagates via the ``raise`` below, so no detail is
121 # swallowed for the caller (#4182). Only the exception
122 # CLASS is logged (never the message/traceback): it cannot
123 # carry the password, and without it every failure here
124 # (lock timeout vs poisoned session vs open failure) is
125 # indistinguishable in the logs.
126 logger.warning(
127 f"Session error for {username} (exc_class={type(exc).__name__})"
128 )
129 if session:
130 session.rollback()
131 raise
132 finally:
133 if session: 133 ↛ exitline 133 didn't return from function 'get_session' because the condition on line 133 was always true
134 from ..utilities.resource_utils import safe_close
136 safe_close(session, "thread metrics session")
138 def write_token_metrics(
139 self, username: str, research_id: Optional[int], token_data: dict
140 ):
141 """
142 Write token metrics from any thread.
144 Args:
145 username: The username (for database access)
146 research_id: The research ID
147 token_data: Dictionary with token metrics data
148 """
149 with self.get_session(username) as session:
150 self.add_token_metrics_to_session(session, research_id, token_data)
152 @staticmethod
153 def add_token_metrics_to_session(
154 session: Session, research_id: Optional[int], token_data: dict
155 ) -> None:
156 """Stage token and model usage on a caller-owned session.
158 The caller owns the transaction. This method never commits, rolls
159 back, closes, or provisions another connection.
160 """
161 # Import here to avoid circular imports
162 from .models import ModelUsage, TokenUsage
164 # The caller resolves the provider-reported total; fall back to
165 # the sum only when it did not supply one.
166 total_tokens = token_data.get("total_tokens")
167 if total_tokens is None:
168 total_tokens = token_data.get("prompt_tokens", 0) + token_data.get(
169 "completion_tokens", 0
170 )
172 # Create TokenUsage record
173 token_usage = TokenUsage(
174 research_id=research_id,
175 model_name=token_data.get("model_name"),
176 model_provider=token_data.get("provider"),
177 prompt_tokens=token_data.get("prompt_tokens", 0),
178 completion_tokens=token_data.get("completion_tokens", 0),
179 total_tokens=total_tokens,
180 # Research context
181 research_query=token_data.get("research_query"),
182 research_mode=token_data.get("research_mode"),
183 research_phase=token_data.get("research_phase"),
184 search_iteration=token_data.get("search_iteration"),
185 # Performance metrics
186 response_time_ms=token_data.get("response_time_ms"),
187 success_status=token_data.get("success_status", "success"),
188 error_type=token_data.get("error_type"),
189 # Search engine context
190 search_engines_planned=token_data.get("search_engines_planned"),
191 search_engine_selected=token_data.get("search_engine_selected"),
192 # Call stack tracking
193 calling_file=token_data.get("calling_file"),
194 calling_function=token_data.get("calling_function"),
195 call_stack=token_data.get("call_stack"),
196 # Context overflow detection
197 context_limit=token_data.get("context_limit"),
198 context_truncated=token_data.get("context_truncated", False),
199 tokens_truncated=token_data.get("tokens_truncated"),
200 truncation_ratio=token_data.get("truncation_ratio"),
201 # Raw Ollama metrics
202 ollama_prompt_eval_count=token_data.get("ollama_prompt_eval_count"),
203 ollama_eval_count=token_data.get("ollama_eval_count"),
204 ollama_total_duration=token_data.get("ollama_total_duration"),
205 ollama_load_duration=token_data.get("ollama_load_duration"),
206 ollama_prompt_eval_duration=token_data.get(
207 "ollama_prompt_eval_duration"
208 ),
209 ollama_eval_duration=token_data.get("ollama_eval_duration"),
210 )
211 session.add(token_usage)
213 # Update or create model usage statistics (matches MainThread
214 # path in token_counter.py _save_to_db).
215 model_name = token_data.get("model_name")
216 model_usage = (
217 session.query(ModelUsage)
218 .filter_by(
219 model_name=model_name,
220 model_provider=token_data.get("provider"),
221 )
222 .first()
223 )
225 if model_usage:
226 model_usage.total_tokens += total_tokens
227 model_usage.total_calls += 1
228 else:
229 model_usage = ModelUsage(
230 model_name=model_name,
231 model_provider=token_data.get("provider"),
232 total_tokens=total_tokens,
233 total_calls=1,
234 )
235 session.add(model_usage)
238# Global instance for thread-safe metrics
239metrics_writer = ThreadSafeMetricsWriter()