Coverage for src/local_deep_research/database/thread_metrics.py: 98%
59 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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 ..utilities.request_context import get_current_username # noqa: F401
20class ThreadSafeMetricsWriter:
21 """
22 Thread-safe writer for metrics to encrypted user databases.
23 Creates encrypted connections per thread using provided passwords.
24 """
26 def __init__(self):
27 self._thread_local = threading.local()
29 def set_user_password(self, username: str, password: str):
30 """
31 Store user password for the current thread.
32 This allows the thread to create its own encrypted connection.
34 IMPORTANT: This is safe because:
35 1. Password is already in memory (user is logged in)
36 2. It's only stored thread-locally
37 3. It's explicitly cleared by `clear_passwords()` (invoked from
38 `cleanup_current_thread()`) so pooled worker threads do not
39 retain credentials between tasks.
40 """
42 if not hasattr(self._thread_local, "passwords"):
43 self._thread_local.passwords = {}
44 self._thread_local.passwords[username] = password
46 def clear_passwords(self):
47 """Remove all passwords cached on the current thread's local store.
49 Called by `cleanup_current_thread()` so that pooled worker threads
50 don't retain plaintext credentials between tasks.
51 """
52 if hasattr(self._thread_local, "passwords"):
53 self._thread_local.passwords.clear()
55 @contextmanager
56 def get_session(self, username: str = None) -> Session:
57 """
58 Get a database session for metrics in the current thread.
59 Creates a new encrypted connection if needed.
61 Args:
62 username: The username for database access. If not provided,
63 will attempt to get it from Flask session.
64 """
65 # If username not provided, get it from request context.
66 # ValueError is the right exception type here — this is a database
67 # service called from threadpools, schedulers, and worker threads,
68 # not from an HTTP handler. werkzeug.exceptions.Unauthorized used
69 # to be raised here, but its 401-meaning is meaningless outside an
70 # HTTP context, and the surrounding try/except (which caught
71 # ImportError/RuntimeError) was structurally broken — the import
72 # never raised the exceptions the except clause caught.
73 if username is None:
74 username = get_current_username()
75 if not username: 75 ↛ 79line 75 didn't jump to line 79 because the condition on line 75 was always true
76 raise ValueError("No username in request context")
78 # Get password for this user in this thread
79 if not hasattr(self._thread_local, "passwords"):
80 raise ValueError("No password set for thread metrics access")
82 password = self._thread_local.passwords.get(username)
84 if not password:
85 raise ValueError(
86 f"No password available for user {username} in this thread"
87 )
89 # Create a thread-safe session for this user
90 session = None
91 try:
92 session = db_manager.create_thread_safe_session_for_metrics(
93 username, password
94 )
95 if not session:
96 raise ValueError( # noqa: TRY301 — except does session rollback before re-raise
97 f"Failed to create session for user {username}"
98 )
99 yield session
100 session.commit()
101 except Exception:
102 # ``logger.warning`` (no traceback) rather than
103 # ``logger.exception``: ``password`` is a live local in this
104 # frame, so rendering a traceback under ``diagnose=True``
105 # would dump the plaintext SQLCipher master password
106 # (unrecoverable — TRUST.md §5). The exception still
107 # propagates via the ``raise`` below, so no detail is
108 # swallowed for the caller (#4182).
109 logger.warning(f"Session error for {username}")
110 if session:
111 session.rollback()
112 raise
113 finally:
114 if session: 114 ↛ exitline 114 didn't return from function 'get_session' because the condition on line 114 was always true
115 from ..utilities.resource_utils import safe_close
117 safe_close(session, "thread metrics session")
119 def write_token_metrics(
120 self, username: str, research_id: Optional[int], token_data: dict
121 ):
122 """
123 Write token metrics from any thread.
125 Args:
126 username: The username (for database access)
127 research_id: The research ID
128 token_data: Dictionary with token metrics data
129 """
130 with self.get_session(username) as session:
131 # Import here to avoid circular imports
132 from .models import ModelUsage, TokenUsage
134 # The caller resolves the provider-reported total; fall back to
135 # the sum only when it did not supply one.
136 total_tokens = token_data.get("total_tokens")
137 if total_tokens is None:
138 total_tokens = token_data.get(
139 "prompt_tokens", 0
140 ) + token_data.get("completion_tokens", 0)
142 # Create TokenUsage record
143 token_usage = TokenUsage(
144 research_id=research_id,
145 model_name=token_data.get("model_name"),
146 model_provider=token_data.get("provider"),
147 prompt_tokens=token_data.get("prompt_tokens", 0),
148 completion_tokens=token_data.get("completion_tokens", 0),
149 total_tokens=total_tokens,
150 # Research context
151 research_query=token_data.get("research_query"),
152 research_mode=token_data.get("research_mode"),
153 research_phase=token_data.get("research_phase"),
154 search_iteration=token_data.get("search_iteration"),
155 # Performance metrics
156 response_time_ms=token_data.get("response_time_ms"),
157 success_status=token_data.get("success_status", "success"),
158 error_type=token_data.get("error_type"),
159 # Search engine context
160 search_engines_planned=token_data.get("search_engines_planned"),
161 search_engine_selected=token_data.get("search_engine_selected"),
162 # Call stack tracking
163 calling_file=token_data.get("calling_file"),
164 calling_function=token_data.get("calling_function"),
165 call_stack=token_data.get("call_stack"),
166 # Context overflow detection
167 context_limit=token_data.get("context_limit"),
168 context_truncated=token_data.get("context_truncated", False),
169 tokens_truncated=token_data.get("tokens_truncated"),
170 truncation_ratio=token_data.get("truncation_ratio"),
171 # Raw Ollama metrics
172 ollama_prompt_eval_count=token_data.get(
173 "ollama_prompt_eval_count"
174 ),
175 ollama_eval_count=token_data.get("ollama_eval_count"),
176 ollama_total_duration=token_data.get("ollama_total_duration"),
177 ollama_load_duration=token_data.get("ollama_load_duration"),
178 ollama_prompt_eval_duration=token_data.get(
179 "ollama_prompt_eval_duration"
180 ),
181 ollama_eval_duration=token_data.get("ollama_eval_duration"),
182 )
183 session.add(token_usage)
185 # Update or create model usage statistics (matches MainThread
186 # path in token_counter.py _save_to_db).
187 model_name = token_data.get("model_name")
188 model_usage = (
189 session.query(ModelUsage)
190 .filter_by(model_name=model_name)
191 .first()
192 )
194 if model_usage:
195 model_usage.total_tokens += total_tokens
196 model_usage.total_calls += 1
197 else:
198 model_usage = ModelUsage(
199 model_name=model_name,
200 model_provider=token_data.get("provider"),
201 total_tokens=total_tokens,
202 total_calls=1,
203 )
204 session.add(model_usage)
207# Global instance for thread-safe metrics
208metrics_writer = ThreadSafeMetricsWriter()