1use tokio::task_local;
2
3pub const SESSION_ID_HEADER: &str = "aster-session-id";
4
5task_local! {
6 pub static SESSION_ID: Option<String>;
7}
8
9pub async fn with_session_id<F>(session_id: Option<String>, f: F) -> F::Output
10where
11 F: std::future::Future,
12{
13 if let Some(id) = session_id {
14 SESSION_ID.scope(Some(id), f).await
15 } else {
16 f.await
17 }
18}
19
20pub fn current_session_id() -> Option<String> {
21 SESSION_ID.try_with(|id| id.clone()).ok().flatten()
22}
23
24#[cfg(test)]
25mod tests {
26 use super::*;
27
28 #[tokio::test]
29 async fn test_session_id_available_when_set() {
30 with_session_id(Some("test-session-123".to_string()), async {
31 assert_eq!(current_session_id(), Some("test-session-123".to_string()));
32 })
33 .await;
34 }
35
36 #[tokio::test]
37 async fn test_session_id_none_when_not_set() {
38 let id = current_session_id();
39 assert_eq!(id, None);
40 }
41
42 #[tokio::test]
43 async fn test_session_id_none_when_explicitly_none() {
44 with_session_id(None, async {
45 assert_eq!(current_session_id(), None);
46 })
47 .await;
48 }
49
50 #[tokio::test]
51 async fn test_session_id_scoped_correctly() {
52 assert_eq!(current_session_id(), None);
53
54 with_session_id(Some("outer-session".to_string()), async {
55 assert_eq!(current_session_id(), Some("outer-session".to_string()));
56
57 with_session_id(Some("inner-session".to_string()), async {
58 assert_eq!(current_session_id(), Some("inner-session".to_string()));
59 })
60 .await;
61
62 assert_eq!(current_session_id(), Some("outer-session".to_string()));
63 })
64 .await;
65
66 assert_eq!(current_session_id(), None);
67 }
68
69 #[tokio::test]
70 async fn test_session_id_across_await_points() {
71 with_session_id(Some("persistent-session".to_string()), async {
72 assert_eq!(current_session_id(), Some("persistent-session".to_string()));
73
74 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
75
76 assert_eq!(current_session_id(), Some("persistent-session".to_string()));
77 })
78 .await;
79 }
80}