Coverage for src / agent_contracts / runtime / session.py: 100%

1"""Session storage abstraction.""" 

2from __future__ import annotations 

3 

4from typing import Protocol, runtime_checkable, Any 

5import time 

6 

7 

8@runtime_checkable 

9class SessionStore(Protocol): 

10 """Protocol for session persistence. 

11  

12 Implement this protocol to provide session storage 

13 (e.g., PostgreSQL, Redis, DynamoDB). 

14  

15 Example: 

16 >>> class RedisSessionStore: 

17 ... def __init__(self, redis_client): 

18 ... self.redis = redis_client 

19 ...  

20 ... async def load(self, session_id: str) -> dict | None: 

21 ... data = await self.redis.get(f"session:{session_id}") 

22 ... return json.loads(data) if data else None 

23 ...  

24 ... async def save(self, session_id: str, data: dict, ttl: int = 3600): 

25 ... await self.redis.setex( 

26 ... f"session:{session_id}", 

27 ... ttl, 

28 ... json.dumps(data), 

29 ... ) 

30 """ 

31 

32 async def load(self, session_id: str) -> dict[str, Any] | None: 

33 """Load session data. 

34  

35 Args: 

36 session_id: The session identifier 

37  

38 Returns: 

39 Session data dict, or None if not found 

40 """ 

41 ... 

42 

43 async def save( 

44 self, 

45 session_id: str, 

46 data: dict[str, Any], 

47 ttl_seconds: int = 3600, 

48 ) -> None: 

49 """Save session data. 

50  

51 Args: 

52 session_id: The session identifier 

53 data: Session data to persist 

54 ttl_seconds: Time-to-live in seconds (default 1 hour) 

55 """ 

56 ... 

57 

58 async def delete(self, session_id: str) -> None: 

59 """Delete session data. 

60  

61 Args: 

62 session_id: The session identifier 

63 """ 

64 ... 

65 

66 

67class InMemorySessionStore: 

68 """In-memory session store for development/testing. 

69  

70 NOT suitable for production use - data is lost on restart. 

71  

72 Example: 

73 >>> store = InMemorySessionStore() 

74 >>> await store.save("session1", {"user": "test"}) 

75 >>> data = await store.load("session1") 

76 >>> print(data) # {"user": "test"} 

77 """ 

78 

79 def __init__(self) -> None: 

80 """Initialize the store.""" 

81 self._store: dict[str, tuple[dict[str, Any], float]] = {} 

82 

83 async def load(self, session_id: str) -> dict[str, Any] | None: 

84 """Load session data. 

85  

86 Returns None if session doesn't exist or has expired. 

87 """ 

88 entry = self._store.get(session_id) 

89 if entry is None: 

90 return None 

91 

92 data, expires_at = entry 

93 if time.time() > expires_at: 

94 # Expired - clean up and return None 

95 del self._store[session_id] 

96 return None 

97 

98 return data 

99 

100 async def save( 

101 self, 

102 session_id: str, 

103 data: dict[str, Any], 

104 ttl_seconds: int = 3600, 

105 ) -> None: 

106 """Save session data with TTL.""" 

107 expires_at = time.time() + ttl_seconds 

108 self._store[session_id] = (data, expires_at) 

109 

110 async def delete(self, session_id: str) -> None: 

111 """Delete session data.""" 

112 self._store.pop(session_id, None) 

113 

114 def clear(self) -> None: 

115 """Clear all sessions (useful for testing).""" 

116 self._store.clear() 

117 

118 def __len__(self) -> int: 

119 """Return number of stored sessions.""" 

120 return len(self._store)