Coverage for src/local_deep_research/utilities/threading_utils.py: 100%

1import threading 

2import uuid 

3from typing import Any, Callable, Hashable, Tuple 

4 

5from cachetools import cached, keys 

6 

7g_thread_local_store = threading.local() 

8 

9 

10def thread_specific_cache(*args: Any, **kwargs: Any) -> Callable: 

11 """ 

12 A version of `cached()` that is local to a single thread. In other words, 

13 cache entries will only be valid in the thread where they were created. 

14 

15 WHY THIS IS STILL THREAD-KEYED UNDER FastAPI, AND NOT A `ContextVar`. 

16 

17 This survived the Flask -> FastAPI migration on purpose; it is not a 

18 leftover. The two Flask helpers that used to live in this module 

19 (`thread_with_app_context`, `thread_context`) were deleted in that 

20 migration because they wrapped `flask.g` / `AppContext`, and identity and 

21 research-context propagation genuinely did move to `ContextVar` — see 

22 `utilities/request_context.py` and `utilities/thread_context.py`. This 

23 decorator is a different case, because what it caches is not context but a 

24 THREAD-AFFINE RESOURCE. 

25 

26 Its only production caller is `db_utils._get_cached_user_session`, which 

27 caches a SQLAlchemy `Session`. The post-auth session accessor on 

28 `db_manager` builds a fresh `sessionmaker(bind=engine)()` on every call and 

29 does nothing to scope it 

30 per thread, so this key is the only thing keeping one cached `Session` from 

31 being handed to two anyio threadpool workers running concurrently for the 

32 same user. A `Session` is documented as not thread-safe: concurrent 

33 identity-map mutation and autoflush on one instance is the failure this 

34 prevents. 

35 

36 A `ContextVar`-keyed cache would NOT be equivalent, and would be unsafe 

37 here. A single request context can straddle two worker threads: Starlette 

38 drives sync-generator dependencies through `contextmanager_in_threadpool`, 

39 which dispatches `__enter__` and `__exit__` as two separate 

40 `anyio.to_thread.run_sync` calls, and anyio selects a worker via 

41 `idle_workers.pop()` with no task affinity. (This hazard is written up at 

42 `web/dependencies/auth.py::get_db_session_dep`, and handled deliberately by 

43 the streaming generators in `web/routers/library.py`.) Keying on context 

44 would therefore return ONE `Session` to a context spread across two 

45 threads — reintroducing exactly the cross-thread sharing that keying on 

46 the thread prevents. `ContextVar` is the right primitive for propagating 

47 identity; `threading.local()` is the right primitive for a resource whose 

48 safety is defined per OS thread. Do not swap one for the other without 

49 also giving `Session` acquisition a different lifetime model. 

50 

51 WHAT THE THREAD KEY DOES *NOT* DO. It provides no cross-user isolation. 

52 That comes from `username` being part of `keys.hashkey(*args_, **kwargs_)` 

53 below, i.e. from applying this decorator BELOW username resolution 

54 (`_get_cached_user_session(username, _namespace)`) rather than above it. 

55 See the note on the lock further down, and 

56 `tests/security/test_cross_user_isolation_invariants.py`, which pins both 

57 properties separately. Conflating the two is how the key-completeness bug 

58 described below got misdiagnosed for weeks. 

59 

60 Args: 

61 *args: Will be forwarded to `cached()`. 

62 **kwargs: Will be forwarded to `cached()`. 

63 

64 Returns: 

65 The wrapped function. 

66 

67 """ 

68 

69 def _key_func(*args_: Any, **kwargs_: Any) -> Tuple[Hashable, ...]: 

70 base_hash = keys.hashkey(*args_, **kwargs_) 

71 

72 if hasattr(g_thread_local_store, "thread_id"): 

73 # We already gave this thread a unique ID. Use that. 

74 thread_id = g_thread_local_store.thread_id 

75 else: 

76 # Give this thread a new unique ID. 

77 thread_id = uuid.uuid4().hex 

78 g_thread_local_store.thread_id = thread_id 

79 

80 return (thread_id,) + base_hash 

81 

82 # cachetools' `cached()` is NOT thread-safe unless a lock is supplied. 

83 # The cache object is shared by every calling thread, and concurrent 

84 # __setitem__/eviction can desync an LRUCache's internal ordering from 

85 # its data dict, after which `popitem()` raises KeyError. With the 

86 # per-thread key below the key space is (thread x args), so a small 

87 # maxsize under many worker threads means near-constant eviction — 

88 # precisely the condition that provokes it. Measured: without this lock, 

89 # ~432k KeyErrors across 480k concurrent calls; with it, zero. 

90 # 

91 # NOTE ON WHAT THIS LOCK DOES *NOT* FIX. An earlier version of this 

92 # comment claimed the unlocked cache could return "the entry stored 

93 # under a DIFFERENT key", and credited the lock with fixing the 

94 # cross-user settings leak this branch chased for weeks. That mechanism 

95 # is not possible: `cachetools.Cache.__getitem__` is `return 

96 # self.__data[key]`, a plain dict lookup on a tuple-of-str key, so a hit 

97 # always returns that key's own value. The failure mode here is loud 

98 # (KeyError), never a silently wrong value. 

99 # 

100 # The cross-user leak was a KEY-COMPLETENESS bug, fixed separately in 

101 # `utilities/db_utils.py` by moving the cache below username resolution 

102 # (`_get_cached_user_session(username, _namespace)`). Before that the key 

103 # was effectively `(thread_uuid,)` for web callers, so every user served 

104 # by one reused worker thread shared a single entry. Measured on the two 

105 # shapes: key-above-resolution leaks ~472k wrong-user sessions per 480k 

106 # calls with or without this lock; key-below-resolution leaks zero. 

107 # 

108 # Both changes are worth keeping — this one prevents crashes, that one 

109 # prevents the leak. They are not interchangeable, and this is the wrong 

110 # place to look for the leak fix. 

111 kwargs.setdefault("lock", threading.RLock()) 

112 return cached(*args, **kwargs, key=_key_func)