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

1""" 

2Per-request context for the authenticated user. 

3 

4Service-layer code historically read `flask_session.get("username")` to 

5discover who the current user is. Under FastAPI there is no Flask request 

6context, so those reads return None and the caller crashes (or silently 

7drops to an "anonymous" branch). 

8 

9This module exposes a `contextvars`-backed username context that the 

10DatabaseMiddleware sets at the start of each authenticated request and 

11clears at the end. Service code can call `get_current_username()` as a 

12Flask-free fallback. 

13 

14contextvars (not `threading.local`) is the correct primitive here because 

15it isolates per-asyncio-task as well as per-thread, which matches FastAPI's 

16mixed sync/async execution model. 

17""" 

18 

19from __future__ import annotations 

20 

21import contextvars 

22from contextlib import contextmanager 

23from typing import Generator, Optional 

24 

25_username_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( 

26 "ldr_request_username", default=None 

27) 

28_session_id_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( 

29 "ldr_request_session_id", default=None 

30) 

31 

32 

33def set_request_user( 

34 username: Optional[str], session_id: Optional[str] = None 

35) -> tuple: 

36 """Set the request-scoped username and session_id. 

37 

38 Returns reset handles so the caller can restore the previous values 

39 via `reset_request_user(handles)`. 

40 """ 

41 user_handle = _username_var.set(username) 

42 session_handle = _session_id_var.set(session_id) 

43 return (user_handle, session_handle) 

44 

45 

46def reset_request_user(handles: tuple) -> None: 

47 """Restore the previous request-scoped username and session_id.""" 

48 user_handle, session_handle = handles 

49 try: 

50 _username_var.reset(user_handle) 

51 except (ValueError, LookupError): 

52 pass 

53 try: 

54 _session_id_var.reset(session_handle) 

55 except (ValueError, LookupError): 

56 pass 

57 

58 

59def get_current_username() -> Optional[str]: 

60 """Return the username for the current request, or None. 

61 

62 Reads the contextvar populated by ``DatabaseMiddleware`` (and by 

63 background workers / scheduler jobs that explicitly push a context 

64 via ``request_user(...)``). Returns None outside any of those 

65 contexts. 

66 """ 

67 return _username_var.get() 

68 

69 

70def get_current_session_id() -> Optional[str]: 

71 """Return the session_id for the current request, or None. 

72 

73 Same source as ``get_current_username`` — the contextvar populated 

74 by ``DatabaseMiddleware`` / explicit ``request_user(...)``. 

75 """ 

76 return _session_id_var.get() 

77 

78 

79@contextmanager 

80def request_user( 

81 username: Optional[str], session_id: Optional[str] = None 

82) -> Generator[None, None, None]: 

83 """Context manager: set request user inside the block, restore on exit.""" 

84 tokens = set_request_user(username, session_id) 

85 try: 

86 yield 

87 finally: 

88 reset_request_user(tokens)