Zane
82 lines
2.2 KiB
Rust
82 lines
2.2 KiB
Rust
use tokio::task_local;
|
|
|
|
pub const SESSION_ID_HEADER: &str = "agent-session-id";
|
|
pub const WORKING_DIR_HEADER: &str = "agent-working-dir";
|
|
|
|
task_local! {
|
|
pub static SESSION_ID: Option<String>;
|
|
}
|
|
|
|
pub async fn with_session_id<F>(session_id: Option<String>, f: F) -> F::Output
|
|
where
|
|
F: std::future::Future,
|
|
{
|
|
if let Some(id) = session_id {
|
|
SESSION_ID.scope(Some(id), f).await
|
|
} else {
|
|
f.await
|
|
}
|
|
}
|
|
|
|
pub fn current_session_id() -> Option<String> {
|
|
SESSION_ID.try_with(|id| id.clone()).ok().flatten()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn test_session_id_available_when_set() {
|
|
with_session_id(Some("test-session-123".to_string()), async {
|
|
assert_eq!(current_session_id(), Some("test-session-123".to_string()));
|
|
})
|
|
.await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_session_id_none_when_not_set() {
|
|
let id = current_session_id();
|
|
assert_eq!(id, None);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_session_id_none_when_explicitly_none() {
|
|
with_session_id(None, async {
|
|
assert_eq!(current_session_id(), None);
|
|
})
|
|
.await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_session_id_scoped_correctly() {
|
|
assert_eq!(current_session_id(), None);
|
|
|
|
with_session_id(Some("outer-session".to_string()), async {
|
|
assert_eq!(current_session_id(), Some("outer-session".to_string()));
|
|
|
|
with_session_id(Some("inner-session".to_string()), async {
|
|
assert_eq!(current_session_id(), Some("inner-session".to_string()));
|
|
})
|
|
.await;
|
|
|
|
assert_eq!(current_session_id(), Some("outer-session".to_string()));
|
|
})
|
|
.await;
|
|
|
|
assert_eq!(current_session_id(), None);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_session_id_across_await_points() {
|
|
with_session_id(Some("persistent-session".to_string()), async {
|
|
assert_eq!(current_session_id(), Some("persistent-session".to_string()));
|
|
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
|
|
|
|
assert_eq!(current_session_id(), Some("persistent-session".to_string()));
|
|
})
|
|
.await;
|
|
}
|
|
}
|