Coverage for src/local_deep_research/web/dependencies/threadpool.py: 100%
81 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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, and ``run_db_sync`` does
16the same for explicit ``asyncio.to_thread`` DB work inside asynchronous
17endpoints.
18"""
20from __future__ import annotations
22import asyncio
23import functools
24from collections.abc import AsyncIterable, Iterable, Iterator
25from typing import TYPE_CHECKING, Any, Callable, TypeVar
27from fastapi.encoders import jsonable_encoder
28from fastapi.routing import APIRoute, APIRouter
29from loguru import logger
30from starlette.responses import Response, StreamingResponse
32if TYPE_CHECKING:
33 from fastapi.dependencies.models import Dependant
35T = TypeVar("T")
37_WORKER_CLEANUP_MARKER = "__ldr_worker_thread_cleanup__"
40def wrap_sync_route_with_cleanup(dependant: "Dependant") -> None:
41 """Wrap a synchronous FastAPI endpoint in an owner-thread cleanup.
43 Mutates ``dependant.call`` in place; does not return the wrapped
44 callable. Takes the ``Dependant`` (rather than the bare callable) so the
45 sync/async/generator classification below can use FastAPI's own
46 ``is_coroutine_callable`` / ``is_gen_callable`` / ``is_async_gen_callable``
47 (``fastapi/dependencies/models.py``) instead of ``inspect.iscoroutinefunction``
48 et al. Those ``inspect`` functions do NOT follow a ``functools.wraps``
49 ``__wrapped__`` chain; FastAPI's own predicates do (via
50 ``inspect.unwrap``), and FastAPI dispatches based on its own predicates
51 -- so a hypothetical ``functools.wraps``-decorated sync passthrough over
52 an ``async def`` endpoint would be misclassified as sync by ``inspect.*``,
53 wrapped here, and then get *awaited* by FastAPI's dispatcher (a 500).
54 Unreachable today (no such endpoint exists in this app, and slowapi's
55 rate-limit decorator preserves sync/async-ness), but free to close and
56 it removes the divergence from FastAPI's own model.
58 These three properties are ``cached_property``s on ``Dependant``,
59 memoized from whatever ``dependant.call`` is at first access. FastAPI's
60 own ``APIRoute.__init__`` (via ``get_route_handler`` ->
61 ``get_request_handler``) already reads ``is_coroutine_callable`` while
62 building the route, before this function ever runs -- so by the time
63 this runs (either from ``WorkerCleanupAPIRoute.__init__`` for direct
64 routes, or from the post-``include_router`` sweep in
65 ``prepare_router_for_worker_cleanup``), the value is already pinned to
66 the real, original endpoint and is unaffected by this function
67 subsequently replacing ``dependant.call`` with the cleanup wrapper.
69 FastAPI executes plain ``def`` endpoints on an AnyIO worker, while ASGI
70 middleware teardown runs on the event-loop thread. Database sessions are
71 thread-local, so middleware cannot release the endpoint worker's sessions.
72 This wrapper's ``finally`` runs on that worker and restores Flask's former
73 request-teardown ownership model.
75 Deliberately NOT a ``yield``-dependency and NOT a ``BackgroundTask``.
76 Both would move cleanup to a *second*, separate
77 ``anyio.to_thread.run_sync`` dispatch, and anyio gives no thread
78 affinity between separate dispatches — the documented hazard in
79 ``get_db_session_dep`` (``web/dependencies/auth.py``) for exactly that
80 split. Cleanup has to happen inside the SAME worker dispatch that ran
81 the endpoint, which is why this wraps ``dependant.call`` itself rather
82 than framing the boundary with FastAPI's own request-lifecycle hooks.
84 That constraint also decides *where inside this one dispatch* cleanup
85 must sit. FastAPI does not serialize a plain (non-``Response``) return
86 value inside ``dependant.call``: ``routing.py``'s route handler calls
87 ``serialize_response``/``jsonable_encoder`` on the raw return value
88 *after* ``dependant.call`` has already returned, back on the event-loop
89 thread. Releasing the session before returning — the natural place for
90 a ``finally``/``with`` cleanup — would therefore free this worker's DB
91 session (and SQLCipher passphrase) before that encoding step runs. Any
92 lazily loaded ORM attribute touched only during encoding would then hit
93 a closed session (``DetachedInstanceError``) or silently serialize as
94 incomplete data instead of raising.
95 ``JSONResponse``, ``HTMLResponse``, ``TemplateResponse`` and
96 ``RedirectResponse`` are unaffected: they render their body
97 synchronously in ``__init__``, i.e. before the endpoint returns them
98 here, so the bytes are already materialized. ``FileResponse`` is the
99 exception, NOT covered by that claim: its ``__init__`` only sets
100 headers (starlette ``responses.py``); the file is opened and streamed
101 later, in its async ``__call__``. That distinction is moot for this
102 wrapper specifically -- every route in this app that returns a
103 ``FileResponse`` (``favicon``, ``serve_static`` in
104 ``web/fastapi_app.py``) is ``async def``, so this sync-only wrapper
105 never touches them -- but do not extend the "renders in ``__init__``"
106 claim to ``FileResponse`` if a future sync route ever returns one.
107 Everything else (a plain dict, list, Pydantic model, ...) gets encoded
108 here, on this worker, before cleanup runs. This is NOT, contrary to an
109 earlier version of this docstring, simply moving forward a
110 ``jsonable_encoder``-only fallback that every route in this app takes.
111 ``APIRoute.__init__`` derives ``response_model`` from a route's return
112 *type annotation* whenever one isn't given explicitly
113 (``fastapi/routing.py``, ~840-905), and a bare ``-> Any`` annotation
114 (21 routes in ``news_flask_api.py``, 14 of them synchronous and so
115 routed through this wrapper) is truthy, so those routes DO get a
116 ``response_field`` -- "no route declares a ``response_model``" is
117 false. With a ``response_field`` present, FastAPI's own
118 ``serialize_response`` (``routing.py``) takes the
119 ``field.validate()`` + ``field.serialize_json()`` branch instead of the
120 bare ``jsonable_encoder`` branch -- but that still runs, on the
121 event-loop thread, on WHATEVER this wrapper already handed back as this
122 worker's return value. Because that value already went through
123 ``jsonable_encoder`` here, ``field.validate``/``serialize_json`` sees
124 already-plain JSON primitives, not the original rich Python objects, so
125 it has nothing left to reformat -- ``jsonable_encoder`` running first
126 effectively wins the two encoders' formatting differences rather than
127 being replaced by the second one. Measured deltas from running
128 ``jsonable_encoder`` first: an aware ``datetime`` serializes as
129 ``...+00:00`` (``datetime.isoformat()``, used by ``jsonable_encoder``)
130 instead of pydantic's compact ``...Z``; ``Decimal("1.50")`` becomes the
131 float ``1.5`` instead of a decimal-preserving pydantic encoding;
132 ``timedelta`` becomes a float of total seconds instead of pydantic's
133 ISO-8601 duration string. No route in this app hits this today -- every
134 datetime returned from one of these routes is already
135 ``.isoformat()``'d before it reaches this wrapper -- but nothing pins
136 that; see ``tests/web/dependencies/test_threadpool.py`` for a
137 regression test on one such route's serialized shape.
138 """
139 fn = dependant.call
140 if fn is None or getattr(fn, _WORKER_CLEANUP_MARKER, False):
141 return
143 # FastAPI streams generator endpoints after the endpoint call returns.
144 # An ordinary ``finally`` would therefore clean up before iteration even
145 # starts (and sync generator iterations are not worker-affine). There are
146 # no generator endpoints in the production app; leave any future one
147 # untouched here rather than apply an incorrect lifetime model.
148 #
149 # Skipping them silently is only safe because a dedicated assertion
150 # watches for them: ``generator_routes == []`` in
151 # ``tests/web/dependencies/test_threadpool.py``'s
152 # ``test_production_app_wraps_every_sync_api_route``. That test's OTHER
153 # census (``missing == []``) cannot cover this -- it classifies routes
154 # by the very same ``is_gen_callable``/``is_async_gen_callable`` flags
155 # read here, so anything this branch returns early on is excluded from
156 # it by construction. Nor can
157 # ``test_production_routers_do_not_use_bare_streaming_response``, which
158 # AST-scans for a ``StreamingResponse`` construction a generator
159 # endpoint never writes. If the ``generator_routes`` assertion is
160 # dropped, a new ``def ...: yield`` route streams via a bare
161 # ``StreamingResponse``/``iterate_in_threadpool`` with no per-chunk
162 # cleanup -- the #6095 leak class -- and nothing says so.
163 if (
164 dependant.is_coroutine_callable
165 or dependant.is_gen_callable
166 or dependant.is_async_gen_callable
167 ):
168 return
170 @functools.wraps(fn)
171 def _wrapped(*args: Any, **kwargs: Any) -> Any:
172 from ...database.thread_local_session import thread_cleanup
174 with thread_cleanup():
175 result = fn(*args, **kwargs)
176 if isinstance(result, Response):
177 return result
178 # Force JSON encoding now, on the worker thread that still owns
179 # the DB session this endpoint used, instead of leaving it for
180 # FastAPI to do afterward on the event-loop thread with the
181 # session already released.
182 return jsonable_encoder(result)
184 setattr(_wrapped, _WORKER_CLEANUP_MARKER, True)
185 dependant.call = _wrapped
188def iterate_sync_with_cleanup(iterable: Iterable[T]) -> Iterator[T]:
189 """Iterate a synchronous response body with owner-worker cleanup.
191 Starlette advances a synchronous ``StreamingResponse`` body through one
192 ``anyio.to_thread.run_sync(next, ...)`` call per chunk. Those calls happen
193 *after* the endpoint has returned, and they are not guaranteed to use the
194 endpoint's worker (or even the same worker for consecutive chunks).
195 Therefore the endpoint wrapper cannot release sessions opened while the
196 body is being produced.
198 Wrapping each individual ``next()`` call gives every worker that advances
199 the iterator its own deterministic cleanup boundary. Cleanup happens
200 before the chunk is handed back to the event loop, and when iteration
201 finishes or raises INSIDE ``next()``.
203 Once started (at least one ``next()`` has been pulled), closing or
204 finalizing this wrapper throws ``GeneratorExit`` at ``yield item``,
205 after the per-chunk cleanup block has exited. The handler below closes
206 the wrapped iterator inside another cleanup boundary, covering its
207 synchronously executed ``finally``. Work that this ``finally``
208 schedules elsewhere is outside that boundary. Closing this wrapper
209 before its first ``next()`` has no such boundary to throw into --
210 Python marks an unstarted generator closed without ever running the
211 handler below.
213 That close handler is NOT worker-affine, despite running the same
214 ``thread_cleanup()`` the per-chunk block does. This wrapper is a
215 generator, and asyncio's async-generator/garbage-collection finalizer
216 schedules the closing ``aclose()``/``close()`` on the EVENT LOOP
217 thread, so ``thread_cleanup()`` there sweeps the loop thread rather
218 than whichever worker last advanced the iterator. That is benign here:
219 ``DatabaseMiddleware`` already sweeps the loop thread once per
220 request, so the loop thread never accumulates sessions, and the
221 per-chunk boundary above -- which does run on the advancing worker --
222 is what actually covers the workers.
224 This describes what happens WHEN the wrapper is closed, not when a
225 client disconnect causes it to close. The pinned Starlette 1.3.1
226 ``StreamingResponse`` does not explicitly ``aclose()`` its body
227 iterator on a failed send; closure can depend on response/iterator
228 finalization. The direct-close regression test therefore does not
229 establish deterministic cleanup at disconnect. The synchronous route
230 generators close their existing service objects in ``finally``;
231 this wrapper protects that work once close/finalization occurs.
232 """
233 iterator: Iterator[T] | None = None
234 try:
235 while True:
236 try:
237 from ...database.thread_local_session import thread_cleanup
239 with thread_cleanup():
240 if iterator is None:
241 iterator = iter(iterable)
242 item = next(iterator)
243 except StopIteration:
244 return
245 yield item
246 except GeneratorExit:
247 # ``iterable`` is typed broadly (``Iterable[T]``); a plain list/tuple
248 # iterator has no ``close()``. Only generators (the only thing the
249 # five production callers ever pass) do, so guard the call rather
250 # than assuming one exists. Those callers are ``library.py:909``
251 # and ``:1342``, ``rag.py:1361`` and ``:3508``, and
252 # ``research.py:2326``; ``test_threadpool.py``'s
253 # ``test_production_routers_do_not_use_bare_streaming_response``
254 # pins the count at ``>= 5``.
255 close = getattr(iterator, "close", None)
256 if close is not None:
257 from ...database.thread_local_session import thread_cleanup
259 with thread_cleanup():
260 close()
261 raise
264class WorkerCleanupStreamingResponse(StreamingResponse):
265 """Streaming response that cleans workers advancing synchronous bodies.
267 Async iterables already run on the event-loop task and are covered by the
268 request middleware. Only synchronous iterables are handed to AnyIO
269 workers, so only those need the per-iteration wrapper.
270 """
272 def __init__(self, content: Any, *args: Any, **kwargs: Any) -> None:
273 if not isinstance(content, AsyncIterable):
274 content = iterate_sync_with_cleanup(content)
275 super().__init__(content, *args, **kwargs)
278class WorkerCleanupAPIRoute(APIRoute):
279 """APIRoute that cleans thread-local state after synchronous endpoints."""
281 def __init__(self, path: str, endpoint: Callable[..., Any], **kwargs: Any):
282 # Keep ``route.endpoint`` as the exact registered handler. SlowAPI,
283 # route-contract checks, and other FastAPI integrations use that
284 # public attribute for metadata and identity. Request dispatch goes
285 # through ``route.dependant.call``, so wrapping that call target gives
286 # us the worker-thread boundary without changing route identity.
287 super().__init__(path, endpoint, **kwargs)
288 if (
289 self.dependant.call is None
290 ): # pragma: no cover - APIRoute rejects this first
291 raise RuntimeError("FastAPI route has no dispatch callable")
292 wrap_sync_route_with_cleanup(self.dependant)
295def prepare_router_for_worker_cleanup(router: APIRouter) -> None:
296 """Wrap sync dispatch targets on an already-populated router.
298 ``FastAPI.include_router`` copies each source route and rebuilds its
299 dependency graph, so this must run on the destination application's
300 router *after* inclusion. Mutating ``dependant.call`` is intentional:
301 FastAPI dispatches through it while ``route.endpoint`` remains the exact
302 handler registered by the module.
303 """
304 for route in router.routes:
305 if isinstance(route, APIRoute):
306 wrap_sync_route_with_cleanup(route.dependant)
309async def run_db_sync(fn: Callable[..., T], /, *args: Any, **kwargs: Any) -> T:
310 """Run a sync DB-touching function on the asyncio default
311 threadpool, then clean up the worker's thread-local DB session.
313 Use this in place of ``await asyncio.to_thread(fn, ...)`` when
314 ``fn`` opens a ``get_user_db_session(...)`` block. If ``fn`` is
315 purely CPU-bound and never opens a DB session, prefer
316 ``asyncio.to_thread`` directly — the cleanup call is cheap but
317 not free.
318 """
320 def _wrapped() -> T:
321 try:
322 return fn(*args, **kwargs)
323 finally:
324 # Avoid importing at module load — cleanup_current_thread
325 # imports the SQLCipher engine bootstrap, which is too
326 # heavy to run at every web import.
327 try:
328 from ...database.thread_local_session import (
329 cleanup_current_thread,
330 )
332 cleanup_current_thread()
333 except Exception:
334 # Never render a traceback here — this frame holds
335 # credentials; see #6223. ``_wrapped`` closes over
336 # ``args``, and one production caller passes a plaintext
337 # DB password through it
338 # (``socketio_asgi.py``: ``run_db_sync(
339 # db_manager.open_user_database, username, password)``).
340 # ``exc_info=True`` is inert under loguru's ``logger.debug``
341 # and stays that way deliberately. Under
342 # ``LDR_LOGURU_DIAGNOSE`` loguru renders the value of every
343 # identifier on each traced frame's displayed source line,
344 # and the stderr sink's default ``backtrace=True`` extends
345 # the trace upward through the executor frames. Today no
346 # displayed line names ``args`` (this frame's is the
347 # ``cleanup_current_thread()`` call, and
348 # ``asyncio.to_thread`` hands the executor a zero-argument
349 # ``partial``, so ``_WorkItem.args`` is empty), but that is
350 # an accident of which line raises, not a property anyone
351 # maintains: any traceback formatted from inside a closure
352 # holding a password is one refactor away from rendering
353 # it, and the redaction-invariant test cannot catch this
354 # call site. So: never ``logger.exception()`` here.
355 logger.debug(
356 "run_db_sync: cleanup_current_thread failed",
357 exc_info=True,
358 )
360 # ``cleanup_current_thread`` only releases the thread-local DB
361 # session. The worker thread is POOLED, so any other ambient
362 # thread-local state set by ``fn`` outlives this call and is
363 # visible to the next user's task on the same thread. Clear the
364 # rest too, matching what ``thread_local_session.thread_cleanup``
365 # does for dedicated worker threads.
366 #
367 # This runs inside ``_wrapped``, i.e. ON the worker thread, which
368 # is the only place it can work: ``DatabaseMiddleware``'s own
369 # cleanup runs in an ``async def`` on the event-loop thread and so
370 # cannot reach the thread-locals of a pooled worker at all.
371 try:
372 from ...config.thread_settings import clear_settings_context
374 clear_settings_context()
375 except Exception:
376 # Never render a traceback here — this frame holds
377 # credentials; see #6223 and the handler above.
378 logger.debug(
379 "run_db_sync: clear_settings_context failed",
380 exc_info=True,
381 )
383 # Same reasoning for the egress audit context, which was the one
384 # piece of per-thread state this cleanup did not cover. Verified
385 # leaking: a context armed inside a run_db_sync task persisted on
386 # the pooled worker and was inherited by every later task on it.
387 #
388 # Not currently exploitable — the only arming site reachable from
389 # here (``analyze_topic``) clears in its own ``finally`` — and the
390 # leak direction is fail-closed, since the hook only gates
391 # PRIVATE_ONLY/STRICT, so a stale context over-blocks a later
392 # request rather than permitting egress. Closed by construction
393 # anyway, so a future arming site cannot quietly make it matter.
394 try:
395 from ...security.egress.audit_hook import clear_active_context
397 clear_active_context()
398 except Exception:
399 # Never render a traceback here — this frame holds
400 # credentials; see #6223 and the first handler above.
401 logger.debug(
402 "run_db_sync: clear_active_context failed",
403 exc_info=True,
404 )
406 return await asyncio.to_thread(_wrapped)
409__all__ = [
410 "WorkerCleanupAPIRoute",
411 "WorkerCleanupStreamingResponse",
412 "iterate_sync_with_cleanup",
413 "prepare_router_for_worker_cleanup",
414 "run_db_sync",
415 "wrap_sync_route_with_cleanup",
416]