Coverage for src/local_deep_research/web/dependencies/threadpool.py: 91%
124 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Helpers for running synchronous DB-touching code from async route
3handlers without leaking thread-local state.
5``DatabaseMiddleware`` cleans up the event-loop thread's DB session at
6the end of every request, but it cannot reach into worker threads
7spawned by ``asyncio.to_thread``. ThreadPoolExecutor workers are
8reused across tasks, so a session opened by one user's task would stay
9attached to that worker — the next task that lands on it would inherit
10the previous user's session (memory bloat at best, cross-user state
11mixing at worst).
13``WorkerCleanupAPIRoute`` supplies that owner-thread boundary for plain
14``def`` endpoints. ``WorkerCleanupStreamingResponse`` supplies it for each
15worker that advances a synchronous response body. ``run_db_sync``,
16``run_research_sync`` and ``run_model_discovery_sync`` do the same for
17explicit DB work dispatched from asynchronous endpoints. Long research calls
18and model discovery each use a separate executor so they cannot fill the
19default executor used by authenticated requests.
20"""
22from __future__ import annotations
24import asyncio
25import atexit
26import functools
27import threading
28from collections.abc import AsyncIterable, Iterable, Iterator
29from concurrent.futures import ThreadPoolExecutor
30from contextvars import copy_context
31from typing import TYPE_CHECKING, Any, Callable, TypeVar
33from fastapi.encoders import jsonable_encoder
34from fastapi.routing import APIRoute, APIRouter
35from loguru import logger
36from starlette.responses import Response, StreamingResponse
38if TYPE_CHECKING:
39 from fastapi.dependencies.models import Dependant
41T = TypeVar("T")
43_WORKER_CLEANUP_MARKER = "__ldr_worker_thread_cleanup__"
45# Long API v1 research calls must not occupy asyncio's default executor:
46# every authenticated request uses that executor to open/check its user DB.
47# Keep this pool separate from the short DB dispatches in run_db_sync.
48_API_RESEARCH_WORKERS = 8
49_api_research_executor: ThreadPoolExecutor | None = None
50_api_research_executor_lock = threading.Lock()
53def _get_api_research_executor() -> ThreadPoolExecutor:
54 global _api_research_executor
55 with _api_research_executor_lock:
56 if _api_research_executor is None:
57 _api_research_executor = ThreadPoolExecutor(
58 max_workers=_API_RESEARCH_WORKERS,
59 thread_name_prefix="ldr-api-research",
60 )
61 return _api_research_executor
64def _shutdown_api_research_executor() -> None:
65 global _api_research_executor
66 with _api_research_executor_lock:
67 executor = _api_research_executor
68 _api_research_executor = None
69 if executor is not None:
70 executor.shutdown(wait=True, cancel_futures=True)
73atexit.register(_shutdown_api_research_executor)
75# Model discovery calls each configured provider in turn, and each blocking
76# list-models call can run until its own timeout. A client disconnect does
77# not cancel the request, so page reloads against a slow provider must not
78# fill the default executor or the research workers. A full pool delays only
79# other model discovery.
80_MODEL_DISCOVERY_WORKERS = 8
81_model_discovery_executor: ThreadPoolExecutor | None = None
82_model_discovery_executor_lock = threading.Lock()
85def _get_model_discovery_executor() -> ThreadPoolExecutor:
86 global _model_discovery_executor
87 with _model_discovery_executor_lock:
88 if _model_discovery_executor is None:
89 _model_discovery_executor = ThreadPoolExecutor(
90 max_workers=_MODEL_DISCOVERY_WORKERS,
91 thread_name_prefix="ldr-model-discovery",
92 )
93 return _model_discovery_executor
96def _shutdown_model_discovery_executor() -> None:
97 global _model_discovery_executor
98 with _model_discovery_executor_lock:
99 executor = _model_discovery_executor
100 _model_discovery_executor = None
101 if executor is not None:
102 executor.shutdown(wait=True, cancel_futures=True)
105atexit.register(_shutdown_model_discovery_executor)
108def wrap_sync_route_with_cleanup(dependant: Dependant) -> None:
109 """Wrap a synchronous FastAPI endpoint in an owner-thread cleanup.
111 Mutates ``dependant.call`` in place; does not return the wrapped
112 callable. Takes the ``Dependant`` (rather than the bare callable) so the
113 sync/async/generator classification below can use FastAPI's own
114 ``is_coroutine_callable`` / ``is_gen_callable`` / ``is_async_gen_callable``
115 (``fastapi/dependencies/models.py``) instead of ``inspect.iscoroutinefunction``
116 et al. Those ``inspect`` functions do NOT follow a ``functools.wraps``
117 ``__wrapped__`` chain; FastAPI's own predicates do (via
118 ``inspect.unwrap``), and FastAPI dispatches based on its own predicates
119 -- so a hypothetical ``functools.wraps``-decorated sync passthrough over
120 an ``async def`` endpoint would be misclassified as sync by ``inspect.*``,
121 wrapped here, and then get *awaited* by FastAPI's dispatcher (a 500).
122 Unreachable today (no such endpoint exists in this app, and slowapi's
123 rate-limit decorator preserves sync/async-ness), but free to close and
124 it removes the divergence from FastAPI's own model.
126 These three properties are ``cached_property``s on ``Dependant``,
127 memoized from whatever ``dependant.call`` is at first access. FastAPI's
128 own ``APIRoute.__init__`` (via ``get_route_handler`` ->
129 ``get_request_handler``) already reads ``is_coroutine_callable`` while
130 building the route, before this function ever runs -- so by the time
131 this runs (either from ``WorkerCleanupAPIRoute.__init__`` for direct
132 routes, or from the post-``include_router`` sweep in
133 ``prepare_router_for_worker_cleanup``), the value is already pinned to
134 the real, original endpoint and is unaffected by this function
135 subsequently replacing ``dependant.call`` with the cleanup wrapper.
137 FastAPI executes plain ``def`` endpoints on an AnyIO worker, while ASGI
138 middleware teardown runs on the event-loop thread. Database sessions are
139 thread-local, so middleware cannot release the endpoint worker's sessions.
140 This wrapper's ``finally`` runs on that worker and restores Flask's former
141 request-teardown ownership model.
143 Deliberately NOT a ``yield``-dependency and NOT a ``BackgroundTask``.
144 Both would move cleanup to a *second*, separate
145 ``anyio.to_thread.run_sync`` dispatch, and anyio gives no thread
146 affinity between separate dispatches — the documented hazard in
147 ``get_db_session_dep`` (``web/dependencies/auth.py``) for exactly that
148 split. Cleanup has to happen inside the SAME worker dispatch that ran
149 the endpoint, which is why this wraps ``dependant.call`` itself rather
150 than framing the boundary with FastAPI's own request-lifecycle hooks.
152 That constraint also decides *where inside this one dispatch* cleanup
153 must sit. FastAPI does not serialize a plain (non-``Response``) return
154 value inside ``dependant.call``: ``routing.py``'s route handler calls
155 ``serialize_response``/``jsonable_encoder`` on the raw return value
156 *after* ``dependant.call`` has already returned, back on the event-loop
157 thread. Releasing the session before returning — the natural place for
158 a ``finally``/``with`` cleanup — would therefore free this worker's DB
159 session (and SQLCipher passphrase) before that encoding step runs. Any
160 lazily loaded ORM attribute touched only during encoding would then hit
161 a closed session (``DetachedInstanceError``) or silently serialize as
162 incomplete data instead of raising.
163 ``JSONResponse``, ``HTMLResponse``, ``TemplateResponse`` and
164 ``RedirectResponse`` are unaffected: they render their body
165 synchronously in ``__init__``, i.e. before the endpoint returns them
166 here, so the bytes are already materialized. ``FileResponse`` is the
167 exception, NOT covered by that claim: its ``__init__`` only sets
168 headers (starlette ``responses.py``); the file is opened and streamed
169 later, in its async ``__call__``, i.e. AFTER this wrapper's cleanup.
170 ``favicon`` and ``serve_static`` (``web/fastapi_app.py``) are
171 ``async def``, so this sync-only wrapper never touches them. The
172 legacy ``redirect_static`` shim (``web/routers/research.py``) is a
173 plain ``def`` that returns one and so does pass through here; that is
174 safe only because its body is a static file opened by path, needing no
175 DB session or other resource this cleanup releases. Do not extend the
176 "renders in ``__init__``" claim to ``FileResponse``: any sync route
177 returning one must likewise not depend on request-scoped resources.
178 Everything else (a plain dict, list, Pydantic model, ...) gets encoded
179 here, on this worker, before cleanup runs. This is NOT, contrary to an
180 earlier version of this docstring, simply moving forward a
181 ``jsonable_encoder``-only fallback that every route in this app takes.
182 ``APIRoute.__init__`` derives ``response_model`` from a route's return
183 *type annotation* whenever one isn't given explicitly
184 (``fastapi/routing.py``, ~840-905), and a bare ``-> Any`` annotation
185 (21 routes in ``news_flask_api.py``, 14 of them synchronous and so
186 routed through this wrapper) is truthy, so those routes DO get a
187 ``response_field`` -- "no route declares a ``response_model``" is
188 false. With a ``response_field`` present, FastAPI's own
189 ``serialize_response`` (``routing.py``) takes the
190 ``field.validate()`` + ``field.serialize_json()`` branch instead of the
191 bare ``jsonable_encoder`` branch -- but that still runs, on the
192 event-loop thread, on WHATEVER this wrapper already handed back as this
193 worker's return value. Because that value already went through
194 ``jsonable_encoder`` here, ``field.validate``/``serialize_json`` sees
195 already-plain JSON primitives, not the original rich Python objects, so
196 it has nothing left to reformat -- ``jsonable_encoder`` running first
197 effectively wins the two encoders' formatting differences rather than
198 being replaced by the second one. Measured deltas from running
199 ``jsonable_encoder`` first: an aware ``datetime`` serializes as
200 ``...+00:00`` (``datetime.isoformat()``, used by ``jsonable_encoder``)
201 instead of pydantic's compact ``...Z``; ``Decimal("1.50")`` becomes the
202 float ``1.5`` instead of a decimal-preserving pydantic encoding;
203 ``timedelta`` becomes a float of total seconds instead of pydantic's
204 ISO-8601 duration string. No route in this app hits this today -- every
205 datetime returned from one of these routes is already
206 ``.isoformat()``'d before it reaches this wrapper -- but nothing pins
207 that; see ``tests/web/dependencies/test_threadpool.py`` for a
208 regression test on one such route's serialized shape.
209 """
210 fn = dependant.call
211 if fn is None or getattr(fn, _WORKER_CLEANUP_MARKER, False):
212 return
214 # FastAPI streams generator endpoints after the endpoint call returns.
215 # An ordinary ``finally`` would therefore clean up before iteration even
216 # starts (and sync generator iterations are not worker-affine). There are
217 # no generator endpoints in the production app; leave any future one
218 # untouched here rather than apply an incorrect lifetime model.
219 #
220 # Skipping them silently is only safe because a dedicated assertion
221 # watches for them: ``generator_routes == []`` in
222 # ``tests/web/dependencies/test_threadpool.py``'s
223 # ``test_production_app_wraps_every_sync_api_route``. That test's OTHER
224 # census (``missing == []``) cannot cover this -- it classifies routes
225 # by the very same ``is_gen_callable``/``is_async_gen_callable`` flags
226 # read here, so anything this branch returns early on is excluded from
227 # it by construction. Nor can
228 # ``test_production_routers_do_not_use_bare_streaming_response``, which
229 # AST-scans for a ``StreamingResponse`` construction a generator
230 # endpoint never writes. If the ``generator_routes`` assertion is
231 # dropped, a new ``def ...: yield`` route streams via a bare
232 # ``StreamingResponse``/``iterate_in_threadpool`` with no per-chunk
233 # cleanup -- the #6095 leak class -- and nothing says so.
234 if (
235 dependant.is_coroutine_callable
236 or dependant.is_gen_callable
237 or dependant.is_async_gen_callable
238 ):
239 return
241 @functools.wraps(fn)
242 def _wrapped(*args: Any, **kwargs: Any) -> Any:
243 from ...database.thread_local_session import thread_cleanup
245 with thread_cleanup():
246 result = fn(*args, **kwargs)
247 if isinstance(result, Response):
248 return result
249 # Force JSON encoding now, on the worker thread that still owns
250 # the DB session this endpoint used, instead of leaving it for
251 # FastAPI to do afterward on the event-loop thread with the
252 # session already released.
253 return jsonable_encoder(result)
255 setattr(_wrapped, _WORKER_CLEANUP_MARKER, True)
256 dependant.call = _wrapped
259def iterate_sync_with_cleanup(iterable: Iterable[T]) -> Iterator[T]:
260 """Iterate a synchronous response body with owner-worker cleanup.
262 Starlette advances a synchronous ``StreamingResponse`` body through one
263 ``anyio.to_thread.run_sync(next, ...)`` call per chunk. Those calls happen
264 *after* the endpoint has returned, and they are not guaranteed to use the
265 endpoint's worker (or even the same worker for consecutive chunks).
266 Therefore the endpoint wrapper cannot release sessions opened while the
267 body is being produced.
269 Wrapping each individual ``next()`` call gives every worker that advances
270 the iterator its own deterministic cleanup boundary. Cleanup happens
271 before the chunk is handed back to the event loop, and when iteration
272 finishes or raises INSIDE ``next()``.
274 Once started (at least one ``next()`` has been pulled), closing or
275 finalizing this wrapper throws ``GeneratorExit`` at ``yield item``,
276 after the per-chunk cleanup block has exited. The handler below closes
277 the wrapped iterator inside another cleanup boundary, covering its
278 synchronously executed ``finally``. Work that this ``finally``
279 schedules elsewhere is outside that boundary. Closing this wrapper
280 before its first ``next()`` has no such boundary to throw into --
281 Python marks an unstarted generator closed without ever running the
282 handler below.
284 That close handler is NOT worker-affine, despite running the same
285 ``thread_cleanup()`` the per-chunk block does. This wrapper is a
286 generator, and asyncio's async-generator/garbage-collection finalizer
287 schedules the closing ``aclose()``/``close()`` on the EVENT LOOP
288 thread, so ``thread_cleanup()`` there sweeps the loop thread rather
289 than whichever worker last advanced the iterator. That is benign here:
290 ``DatabaseMiddleware`` already sweeps the loop thread once per
291 request, so the loop thread never accumulates sessions, and the
292 per-chunk boundary above -- which does run on the advancing worker --
293 is what actually covers the workers.
295 This describes what happens WHEN the wrapper is closed, not when a
296 client disconnect causes it to close. The pinned Starlette 1.3.1
297 ``StreamingResponse`` does not explicitly ``aclose()`` its body
298 iterator on a failed send; closure can depend on response/iterator
299 finalization. The direct-close regression test therefore does not
300 establish deterministic cleanup at disconnect. The synchronous route
301 generators close their existing service objects in ``finally``;
302 this wrapper protects that work once close/finalization occurs.
303 """
304 iterator: Iterator[T] | None = None
305 try:
306 while True:
307 try:
308 from ...database.thread_local_session import thread_cleanup
310 with thread_cleanup():
311 if iterator is None:
312 iterator = iter(iterable)
313 item = next(iterator)
314 except StopIteration:
315 return
316 yield item
317 except GeneratorExit:
318 # ``iterable`` is typed broadly (``Iterable[T]``); a plain list/tuple
319 # iterator has no ``close()``. Only generators (the only thing the
320 # five production callers ever pass) do, so guard the call rather
321 # than assuming one exists. Those callers are ``library.py:909``
322 # and ``:1342``, ``rag.py:1361`` and ``:3508``, and
323 # ``research.py:2326``; ``test_threadpool.py``'s
324 # ``test_production_routers_do_not_use_bare_streaming_response``
325 # pins the count at ``>= 5``.
326 close = getattr(iterator, "close", None)
327 if close is not None:
328 from ...database.thread_local_session import thread_cleanup
330 with thread_cleanup():
331 close()
332 raise
335class WorkerCleanupStreamingResponse(StreamingResponse):
336 """Streaming response that cleans workers advancing synchronous bodies.
338 Async iterables already run on the event-loop task and are covered by the
339 request middleware. Only synchronous iterables are handed to AnyIO
340 workers, so only those need the per-iteration wrapper.
341 """
343 def __init__(self, content: Any, *args: Any, **kwargs: Any) -> None:
344 if not isinstance(content, AsyncIterable):
345 content = iterate_sync_with_cleanup(content)
346 super().__init__(content, *args, **kwargs)
349class WorkerCleanupAPIRoute(APIRoute):
350 """APIRoute that cleans thread-local state after synchronous endpoints."""
352 def __init__(self, path: str, endpoint: Callable[..., Any], **kwargs: Any):
353 # Keep ``route.endpoint`` as the exact registered handler. SlowAPI,
354 # route-contract checks, and other FastAPI integrations use that
355 # public attribute for metadata and identity. Request dispatch goes
356 # through ``route.dependant.call``, so wrapping that call target gives
357 # us the worker-thread boundary without changing route identity.
358 super().__init__(path, endpoint, **kwargs)
359 if (
360 self.dependant.call is None
361 ): # pragma: no cover - APIRoute rejects this first
362 raise RuntimeError("FastAPI route has no dispatch callable")
363 wrap_sync_route_with_cleanup(self.dependant)
366def prepare_router_for_worker_cleanup(router: APIRouter) -> None:
367 """Wrap sync dispatch targets on an already-populated router.
369 ``FastAPI.include_router`` copies each source route and rebuilds its
370 dependency graph, so this must run on the destination application's
371 router *after* inclusion. Mutating ``dependant.call`` is intentional:
372 FastAPI dispatches through it while ``route.endpoint`` remains the exact
373 handler registered by the module.
374 """
375 for route in router.routes:
376 if isinstance(route, APIRoute):
377 wrap_sync_route_with_cleanup(route.dependant)
380async def _run_db_sync(
381 fn: Callable[..., T],
382 args: tuple[Any, ...],
383 kwargs: dict[str, Any],
384 executor: ThreadPoolExecutor | None,
385) -> T:
386 def _wrapped() -> T:
387 try:
388 return fn(*args, **kwargs)
389 finally:
390 # Avoid importing at module load — cleanup_current_thread
391 # imports the SQLCipher engine bootstrap, which is too
392 # heavy to run at every web import.
393 try:
394 from ...database.thread_local_session import (
395 cleanup_current_thread,
396 )
398 cleanup_current_thread()
399 except Exception:
400 # Never render a traceback here — this frame holds
401 # credentials; see #6223. ``_wrapped`` closes over
402 # ``args``, and one production caller passes a plaintext
403 # DB password through it
404 # (``socketio_asgi.py``: ``run_db_sync(
405 # db_manager.open_user_database, username, password)``).
406 # ``exc_info=True`` is inert under loguru's ``logger.debug``
407 # and stays that way deliberately. Under
408 # ``LDR_LOGURU_DIAGNOSE`` loguru renders the value of every
409 # identifier on each traced frame's displayed source line,
410 # and the stderr sink's default ``backtrace=True`` extends
411 # the trace upward through the executor frames. Today no
412 # displayed line names ``args`` (this frame's is the
413 # ``cleanup_current_thread()`` call), but that is an accident
414 # of which line raises, not a property anyone maintains.
415 # Both the default and dedicated executors can change their
416 # call frames; a traceback from a closure holding a password
417 # is one refactor away from rendering it, and the
418 # redaction-invariant test cannot catch this call site.
419 # So: never ``logger.exception()`` here.
420 logger.debug(
421 "run_db_sync: cleanup_current_thread failed",
422 exc_info=True,
423 )
425 # ``cleanup_current_thread`` only releases the thread-local DB
426 # session. The worker thread is POOLED, so any other ambient
427 # thread-local state set by ``fn`` outlives this call and is
428 # visible to the next user's task on the same thread. Clear the
429 # rest too, matching what ``thread_local_session.thread_cleanup``
430 # does for dedicated worker threads.
431 #
432 # This runs inside ``_wrapped``, i.e. ON the worker thread, which
433 # is the only place it can work: ``DatabaseMiddleware``'s own
434 # cleanup runs in an ``async def`` on the event-loop thread and so
435 # cannot reach the thread-locals of a pooled worker at all.
436 try:
437 from ...config.thread_settings import clear_settings_context
439 clear_settings_context()
440 except Exception:
441 # Never render a traceback here — this frame holds
442 # credentials; see #6223 and the handler above.
443 logger.debug(
444 "run_db_sync: clear_settings_context failed",
445 exc_info=True,
446 )
448 # Same reasoning for the egress audit context, which was the one
449 # piece of per-thread state this cleanup did not cover. Verified
450 # leaking: a context armed inside a run_db_sync task persisted on
451 # the pooled worker and was inherited by every later task on it.
452 #
453 # Not currently exploitable — the only arming site reachable from
454 # here (``analyze_topic``) clears in its own ``finally`` — and the
455 # leak direction is fail-closed, since the hook only gates
456 # PRIVATE_ONLY/STRICT, so a stale context over-blocks a later
457 # request rather than permitting egress. Closed by construction
458 # anyway, so a future arming site cannot quietly make it matter.
459 try:
460 from ...security.egress.audit_hook import clear_active_context
462 clear_active_context()
463 except Exception:
464 # Never render a traceback here — this frame holds
465 # credentials; see #6223 and the first handler above.
466 logger.debug(
467 "run_db_sync: clear_active_context failed",
468 exc_info=True,
469 )
471 if executor is None:
472 return await asyncio.to_thread(_wrapped)
473 # run_in_executor does not copy contextvars as asyncio.to_thread does.
474 # The research pipeline needs its request owner and egress scope there.
475 context = copy_context()
476 return await asyncio.get_running_loop().run_in_executor(
477 executor, context.run, _wrapped
478 )
481async def run_db_sync(fn: Callable[..., T], /, *args: Any, **kwargs: Any) -> T:
482 """Run short sync DB work on the default executor with worker cleanup.
484 Use this in place of ``await asyncio.to_thread(fn, ...)`` when ``fn``
485 opens a per-user DB session. Long research calls use
486 :func:`run_research_sync` and model discovery uses
487 :func:`run_model_discovery_sync`, so they cannot starve authenticated
488 requests.
489 """
490 return await _run_db_sync(fn, args, kwargs, None)
493async def run_research_sync(
494 fn: Callable[..., T], /, *args: Any, **kwargs: Any
495) -> T:
496 """Run long DB-touching research on its own eight-worker executor.
498 Preserves the caller's contextvars and cleans worker thread-local state.
499 """
500 return await _run_db_sync(fn, args, kwargs, _get_api_research_executor())
503async def run_model_discovery_sync(
504 fn: Callable[..., T], /, *args: Any, **kwargs: Any
505) -> T:
506 """Run model discovery on its own eight-worker executor.
508 Discovery can wait out one provider timeout after another, so it stays
509 off the default executor and the research workers. Preserves the
510 caller's contextvars and cleans worker thread-local state.
511 """
512 return await _run_db_sync(fn, args, kwargs, _get_model_discovery_executor())
515__all__ = [
516 "WorkerCleanupAPIRoute",
517 "WorkerCleanupStreamingResponse",
518 "iterate_sync_with_cleanup",
519 "prepare_router_for_worker_cleanup",
520 "run_db_sync",
521 "run_model_discovery_sync",
522 "run_research_sync",
523 "wrap_sync_route_with_cleanup",
524]