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