Propagate session ID in LLM and MCP requests (#5165)

crates/goose/src/session_context.rs

@@ -0,0 +1,80 @@
+use tokio::task_local;
+
+pub const SESSION_ID_HEADER: &str = "goose-session-id";
+
+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;
+    }
+}