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

crates/goose/src/agents/agent.rs

@@ -825,11 +825,9 @@ impl Agent {
                     }
                 }
                 Err(e) => {
-                    yield AgentEvent::Message(
-                        Message::assistant().with_text(
-                            format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
-                        )
-                    );
+                    yield AgentEvent::Message(Message::assistant().with_text(
+                        format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
+                    ));
                 }
             }
         }))

crates/goose/src/agents/mcp_client.rs

@@ -1,4 +1,5 @@
 use crate::agents::types::SharedProvider;
+use crate::session_context::SESSION_ID_HEADER;
 use rmcp::model::{Content, ErrorCode, JsonObject};
 /// MCP client implementation for Goose
 use rmcp::{
@@ -334,7 +335,7 @@ impl McpClientTrait for McpClient {
                 ClientRequest::ListResourcesRequest(ListResourcesRequest {
                     params: Some(PaginatedRequestParam { cursor }),
                     method: Default::default(),
-                    extensions: Default::default(),
+                    extensions: inject_session_into_extensions(Default::default()),
                 }),
                 cancel_token,
             )
@@ -358,7 +359,7 @@ impl McpClientTrait for McpClient {
                         uri: uri.to_string(),
                     },
                     method: Default::default(),
-                    extensions: Default::default(),
+                    extensions: inject_session_into_extensions(Default::default()),
                 }),
                 cancel_token,
             )
@@ -380,7 +381,7 @@ impl McpClientTrait for McpClient {
                 ClientRequest::ListToolsRequest(ListToolsRequest {
                     params: Some(PaginatedRequestParam { cursor }),
                     method: Default::default(),
-                    extensions: Default::default(),
+                    extensions: inject_session_into_extensions(Default::default()),
                 }),
                 cancel_token,
             )
@@ -406,7 +407,7 @@ impl McpClientTrait for McpClient {
                         arguments,
                     },
                     method: Default::default(),
-                    extensions: Default::default(),
+                    extensions: inject_session_into_extensions(Default::default()),
                 }),
                 cancel_token,
             )
@@ -428,7 +429,7 @@ impl McpClientTrait for McpClient {
                 ClientRequest::ListPromptsRequest(ListPromptsRequest {
                     params: Some(PaginatedRequestParam { cursor }),
                     method: Default::default(),
-                    extensions: Default::default(),
+                    extensions: inject_session_into_extensions(Default::default()),
                 }),
                 cancel_token,
             )
@@ -458,7 +459,7 @@ impl McpClientTrait for McpClient {
                         arguments,
                     },
                     method: Default::default(),
-                    extensions: Default::default(),
+                    extensions: inject_session_into_extensions(Default::default()),
                 }),
                 cancel_token,
             )
@@ -476,3 +477,118 @@ impl McpClientTrait for McpClient {
         rx
     }
 }
+
+/// Replaces session ID, case-insensitively, in Extensions._meta.
+fn inject_session_into_extensions(
+    mut extensions: rmcp::model::Extensions,
+) -> rmcp::model::Extensions {
+    use rmcp::model::Meta;
+
+    if let Some(session_id) = crate::session_context::current_session_id() {
+        let mut meta_map = extensions
+            .get::<Meta>()
+            .map(|meta| meta.0.clone())
+            .unwrap_or_default();
+
+        // JsonObject is case-sensitive, so we use retain for case-insensitive removal
+        meta_map.retain(|k, _| !k.eq_ignore_ascii_case(SESSION_ID_HEADER));
+
+        meta_map.insert(SESSION_ID_HEADER.to_string(), Value::String(session_id));
+
+        extensions.insert(Meta(meta_map));
+    }
+
+    extensions
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rmcp::model::Meta;
+
+    #[tokio::test]
+    async fn test_session_id_in_mcp_meta() {
+        use serde_json::json;
+
+        let session_id = "test-session-789";
+        crate::session_context::with_session_id(Some(session_id.to_string()), async {
+            let extensions = inject_session_into_extensions(Default::default());
+            let meta = extensions.get::<Meta>().unwrap();
+
+            assert_eq!(
+                &meta.0,
+                json!({
+                    SESSION_ID_HEADER: session_id
+                })
+                .as_object()
+                .unwrap()
+            );
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn test_no_session_id_in_mcp_when_absent() {
+        let extensions = inject_session_into_extensions(Default::default());
+        let meta = extensions.get::<Meta>();
+
+        assert!(meta.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_all_mcp_operations_include_session() {
+        use serde_json::json;
+
+        let session_id = "consistent-session-id";
+        crate::session_context::with_session_id(Some(session_id.to_string()), async {
+            let ext1 = inject_session_into_extensions(Default::default());
+            let ext2 = inject_session_into_extensions(Default::default());
+            let ext3 = inject_session_into_extensions(Default::default());
+
+            for ext in [&ext1, &ext2, &ext3] {
+                assert_eq!(
+                    &ext.get::<Meta>().unwrap().0,
+                    json!({
+                        SESSION_ID_HEADER: session_id
+                    })
+                    .as_object()
+                    .unwrap()
+                );
+            }
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn test_session_id_case_insensitive_replacement() {
+        use rmcp::model::{Extensions, Meta};
+        use serde_json::{from_value, json};
+
+        let session_id = "new-session-id";
+        crate::session_context::with_session_id(Some(session_id.to_string()), async {
+            let mut extensions = Extensions::new();
+            extensions.insert(
+                from_value::<Meta>(json!({
+                    "GOOSE-SESSION-ID": "old-session-1",
+                    "Goose-Session-Id": "old-session-2",
+                    "other-key": "preserve-me"
+                }))
+                .unwrap(),
+            );
+
+            let extensions = inject_session_into_extensions(extensions);
+            let meta = extensions.get::<Meta>().unwrap();
+
+            assert_eq!(
+                &meta.0,
+                json!({
+                    SESSION_ID_HEADER: session_id,
+                    "other-key": "preserve-me"
+                })
+                .as_object()
+                .unwrap()
+            );
+        })
+        .await;
+    }
+}

crates/goose/src/agents/subagent_handler.rs

@@ -164,10 +164,15 @@ fn get_agent_messages(
         } else {
             None
         };
-        let mut stream = agent
-            .reply(conversation.clone(), session_config, None)
-            .await
-            .map_err(|e| anyhow!("Failed to get reply from agent: {}", e))?;
+
+        let session_id = session_config.as_ref().map(|s| s.id.clone());
+        let mut stream = crate::session_context::with_session_id(session_id, async {
+            agent
+                .reply(conversation.clone(), session_config, None)
+                .await
+        })
+        .await
+        .map_err(|e| anyhow!("Failed to get reply from agent: {}", e))?;
         while let Some(message_result) = stream.next().await {
             match message_result {
                 Ok(AgentEvent::Message(msg)) => conversation.push(msg),

crates/goose/src/lib.rs

@@ -17,6 +17,7 @@ pub mod scheduler_factory;
 pub mod scheduler_trait;
 pub mod security;
 pub mod session;
+pub mod session_context;
 pub mod token_counter;
 pub mod tool_inspection;
 pub mod tool_monitor;

crates/goose/src/providers/api_client.rs

@@ -1,3 +1,4 @@
+use crate::session_context::SESSION_ID_HEADER;
 use anyhow::Result;
 use async_trait::async_trait;
 use reqwest::{
@@ -369,6 +370,10 @@ impl<'a> ApiRequestBuilder<'a> {
         let mut request = request_builder(url, &self.client.client);
         request = request.headers(self.headers.clone());
 
+        if let Some(session_id) = crate::session_context::current_session_id() {
+            request = request.header(SESSION_ID_HEADER, session_id);
+        }
+
         request = match &self.client.auth {
             AuthMethod::BearerToken(token) => {
                 request.header("Authorization", format!("Bearer {}", token))
@@ -398,3 +403,55 @@ impl fmt::Debug for ApiClient {
             .finish_non_exhaustive()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_session_id_header_injection() {
+        let client = ApiClient::new(
+            "http://localhost:8080".to_string(),
+            AuthMethod::BearerToken("test-token".to_string()),
+        )
+        .unwrap();
+
+        // Execute request within session context
+        crate::session_context::with_session_id(Some("test-session-456".to_string()), async {
+            let builder = client.request("/test");
+            let request = builder
+                .send_request(|url, client| client.get(url))
+                .await
+                .unwrap();
+
+            let headers = request.build().unwrap().headers().clone();
+
+            assert!(headers.contains_key(SESSION_ID_HEADER));
+            assert_eq!(
+                headers.get(SESSION_ID_HEADER).unwrap().to_str().unwrap(),
+                "test-session-456"
+            );
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn test_no_session_id_header_when_absent() {
+        let client = ApiClient::new(
+            "http://localhost:8080".to_string(),
+            AuthMethod::BearerToken("test-token".to_string()),
+        )
+        .unwrap();
+
+        // Build a request without session context
+        let builder = client.request("/test");
+        let request = builder
+            .send_request(|url, client| client.get(url))
+            .await
+            .unwrap();
+
+        let headers = request.build().unwrap().headers().clone();
+
+        assert!(!headers.contains_key(SESSION_ID_HEADER));
+    }
+}

crates/goose/src/providers/mod.rs

@@ -1,5 +1,5 @@
 pub mod anthropic;
-mod api_client;
+pub mod api_client;
 pub mod azure;
 pub mod azureauth;
 pub mod base;

crates/goose/src/providers/openai.rs

@@ -112,6 +112,20 @@ impl OpenAiProvider {
         })
     }
 
+    #[doc(hidden)]
+    pub fn new(api_client: ApiClient, model: ModelConfig) -> Self {
+        Self {
+            api_client,
+            base_path: "v1/chat/completions".to_string(),
+            organization: None,
+            project: None,
+            model,
+            custom_headers: None,
+            supports_streaming: true,
+            name: Self::metadata().name,
+        }
+    }
+
     pub fn from_custom_config(
         model: ModelConfig,
         config: DeclarativeProviderConfig,

crates/goose/src/scheduler.rs

@@ -1216,9 +1216,13 @@ async fn run_scheduled_job_internal(
         retry_config: None,
     };
 
-    match agent
-        .reply(conversation.clone(), Some(session_config.clone()), None)
-        .await
+    let session_id = Some(session_config.id.clone());
+    match crate::session_context::with_session_id(session_id, async {
+        agent
+            .reply(conversation.clone(), Some(session_config.clone()), None)
+            .await
+    })
+    .await
     {
         Ok(mut stream) => {
             use futures::StreamExt;

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;
+    }
+}