Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-session/src/lexigram/ai/session/context/session_context.py: 47%

1"""Per-request session context using contextvars for task-scoped isolation.""" 

2 

3from __future__ import annotations 

4 

5import contextvars 

6 

7from lexigram.ai.session.exceptions import SessionError 

8from lexigram.contracts.ai.session import SessionManagerProtocol, SessionState 

9 

10_SESSION_STATE: contextvars.ContextVar[SessionState | None] = contextvars.ContextVar( 

11 "lexigram_session_state", default=None 

12) 

13 

14 

15class SessionContext: 

16 """Per-request session context backed by a ``ContextVar``. 

17 

18 Implements ``SessionContextProtocol``. Each asyncio task (request) has 

19 its own isolated slot in ``_SESSION_STATE``; no cross-request leakage 

20 occurs even under high concurrency. 

21 

22 The typical lifecycle inside a request middleware is:: 

23 

24 token = ctx.bind(state) # request start 

25 try: 

26 ... # handle request 

27 finally: 

28 ctx.unbind(token) # request end 

29 

30 Args: 

31 manager: Session lifecycle manager used by ``get_or_create``. 

32 """ 

33 

34 def __init__(self, manager: SessionManagerProtocol) -> None: 

35 self._manager = manager 

36 

37 @property 

38 def session_id(self) -> str: 

39 """Current session ID from the task-scoped context variable. 

40 

41 Returns: 

42 The bound session ID. 

43 

44 Raises: 

45 SessionError: If no session has been bound for this task. 

46 """ 

47 state = _SESSION_STATE.get() 

48 if state is None: 

49 raise SessionError("No session is currently bound to this context") 

50 return state.session_id 

51 

52 @property 

53 def state(self) -> SessionState: 

54 """Current session state from the task-scoped context variable. 

55 

56 Returns: 

57 The bound ``SessionState``. 

58 

59 Raises: 

60 SessionError: If no session has been bound for this task. 

61 """ 

62 state = _SESSION_STATE.get() 

63 if state is None: 

64 raise SessionError("No session is currently bound to this context") 

65 return state 

66 

67 async def get_or_create(self, user_id: str) -> SessionState: 

68 """Return the bound session, or create a new one if unbound. 

69 

70 Args: 

71 user_id: Owner of the session (only used when creating). 

72 

73 Returns: 

74 An existing or freshly created ``SessionState``. 

75 """ 

76 state = _SESSION_STATE.get() 

77 if state is not None: 

78 return state 

79 new_state = await self._manager.create(user_id=user_id) 

80 _SESSION_STATE.set(new_state) 

81 return new_state 

82 

83 def bind(self, state: SessionState) -> contextvars.Token[SessionState | None]: 

84 """Bind a session state to the current asyncio task. 

85 

86 Should be called at the start of a request by middleware. 

87 

88 Args: 

89 state: The ``SessionState`` to bind. 

90 

91 Returns: 

92 A ``Token`` that must be passed to ``unbind`` to restore the 

93 previous context value. 

94 """ 

95 return _SESSION_STATE.set(state) 

96 

97 def unbind(self, token: contextvars.Token[SessionState | None]) -> None: 

98 """Restore the context variable to its pre-bind value. 

99 

100 Should be called at the end of a request (in a ``finally`` block) 

101 by middleware. 

102 

103 Args: 

104 token: The token returned by the corresponding ``bind`` call. 

105 """ 

106 _SESSION_STATE.reset(token) 

107 

108 

109__all__ = ["SessionContext"]