Coverage for src/local_deep_research/utilities/thread_context.py: 97%

1""" 

2Utility functions for handling research context propagation. 

3 

4This module provides helpers for propagating research context across thread 

5and asyncio boundaries. Built on ``contextvars.ContextVar`` so the context 

6is correctly inherited by frameworks that copy the context to worker 

7threads (e.g. langchain's ``ContextThreadPoolExecutor`` used by LangGraph 

8for parallel tool execution). For stdlib ``ThreadPoolExecutor`` — which 

9does not copy context — use ``preserve_research_context`` below. 

10""" 

11 

12import functools 

13from contextlib import contextmanager 

14from contextvars import ContextVar 

15from typing import Any, Callable, Dict, Generator, Optional 

16 

17from loguru import logger 

18 

19_search_context_var: ContextVar[Optional[Dict[str, Any]]] = ContextVar( 

20 "ldr_search_context", default=None 

21) 

22 

23 

24def set_search_context(context: Dict[str, Any]) -> None: 

25 """ 

26 Sets the research context for the current execution context. 

27 

28 Args: 

29 context: The context to set. 

30 

31 """ 

32 if _search_context_var.get() is not None: 

33 logger.debug( 

34 "Context already set for this thread. It will be overwritten." 

35 ) 

36 _search_context_var.set(context.copy()) 

37 

38 

39def clear_search_context() -> None: 

40 """ 

41 Clears the research context for the current execution context. 

42 

43 Should be called in a finally block after set_search_context() to prevent 

44 context from leaking to subsequent tasks when threads are reused in a pool. 

45 """ 

46 _search_context_var.set(None) 

47 

48 

49def get_search_context() -> Dict[str, Any] | None: 

50 """ 

51 Gets the current research context. 

52 

53 Returns: 

54 The context dictionary, or None if no context is set. 

55 

56 """ 

57 context = _search_context_var.get() 

58 if context is not None: 

59 context = context.copy() 

60 return context 

61 

62 

63@contextmanager 

64def search_context(context: Dict[str, Any]) -> Generator[None, None, None]: 

65 """Context manager that sets and clears search context automatically. 

66 

67 Ensures cleanup even if an exception occurs, preventing context leaks 

68 when threads are reused in a pool. 

69 

70 Example: 

71 with search_context({"research_id": "123"}): 

72 results = engine.run(query) 

73 """ 

74 set_search_context(context) 

75 try: 

76 yield 

77 finally: 

78 clear_search_context() 

79 

80 

81def preserve_research_context(func: Callable) -> Callable: 

82 """ 

83 Decorator that preserves research context across thread boundaries. 

84 

85 Use this decorator on functions that will be executed in ThreadPoolExecutor 

86 to ensure the research context (including research_id) is properly propagated. 

87 

88 When metrics are disabled (e.g., in programmatic mode), this decorator 

89 safely does nothing to avoid database dependencies. 

90 

91 Example: 

92 @preserve_research_context 

93 def search_task(query): 

94 return search_engine.run(query) 

95 """ 

96 # Try to capture current context, but don't fail if it's not set. There 

97 # are legitimate cases where it might not be set, such as for 

98 # programmatic access. 

99 context = get_search_context() 

100 

101 # Capture the submitter thread's egress audit-hook context too. Unlike the 

102 # search context (a ContextVar that langchain's executors copy), the audit 

103 # context lives in a threading.local that stdlib ThreadPoolExecutor workers 

104 # do NOT inherit — so without re-arming it here the PEP-578 socket backstop 

105 # would be inactive on every pool worker. Capture once on the submitter 

106 # thread; re-arm + clear per task below. 

107 try: 

108 from ..security.egress.audit_hook import ( 

109 get_active_context as _get_egress_ctx, 

110 ) 

111 

112 egress_ctx = _get_egress_ctx() 

113 except Exception: 

114 egress_ctx = None 

115 

116 @functools.wraps(func) 

117 def wrapper(*args, **kwargs): 

118 if context is not None: 

119 set_search_context(context) 

120 if egress_ctx is not None: 

121 from ..security.egress.audit_hook import ( 

122 set_active_context as _set_egress_ctx, 

123 ) 

124 

125 _set_egress_ctx(egress_ctx) 

126 

127 try: 

128 return func(*args, **kwargs) 

129 finally: 

130 if egress_ctx is not None: 

131 from ..security.egress.audit_hook import ( 

132 clear_active_context as _clear_egress_ctx, 

133 ) 

134 

135 _clear_egress_ctx() 

136 if context is not None: 

137 clear_search_context() 

138 # Clean up thread-local DB engines created by metrics recording 

139 try: 

140 from ..database.thread_local_session import ( 

141 cleanup_current_thread, 

142 ) 

143 

144 cleanup_current_thread() 

145 except Exception: 

146 logger.debug( 

147 "preserve_research_context: error during cleanup_current_thread", 

148 exc_info=True, 

149 ) 

150 

151 return wrapper