Route MCP elicitations through tool streams (#9943)

Signed-off-by: Douwe M Osinga <douwe@sidewalklabs.com> Co-authored-by: Douwe M Osinga <douwe@sidewalklabs.com>
2026-06-23 21:31:25 -04:00
parent 35711b8691
commit 6782d1f506
11 changed files with 632 additions and 129 deletions
@@ -1,9 +1,9 @@
 use anyhow::Result;
 use serde_json::Value;
-use std::collections::{HashMap, VecDeque};
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::sync::{Mutex, OwnedMutexGuard, RwLock};
+use tokio::sync::{mpsc, Mutex, OwnedMutexGuard, RwLock};
 use tokio::time::timeout;
 use tracing::warn;
 use uuid::Uuid;
@@ -46,14 +46,14 @@ impl PendingResponseClaim {

 pub(crate) struct ActionRequiredManager {
    pending: Arc<RwLock<HashMap<String, Arc<Mutex<PendingRequest>>>>>,
-    queued_requests: Mutex<HashMap<String, VecDeque<Message>>>,
+    action_required_senders: Mutex<HashMap<(String, String), mpsc::Sender<Message>>>,
 }

 impl ActionRequiredManager {
    fn new() -> Self {
        Self {
            pending: Arc::new(RwLock::new(HashMap::new())),
-            queued_requests: Mutex::new(HashMap::new()),
+            action_required_senders: Mutex::new(HashMap::new()),
        }
    }

@@ -66,6 +66,7 @@ impl ActionRequiredManager {
    pub(crate) async fn request_and_wait(
        &self,
        session_id: String,
+        tool_call_request_id: String,
        message: String,
        schema: Value,
        timeout_duration: Duration,
@@ -87,12 +88,28 @@ impl ActionRequiredManager {
            MessageContent::action_required_elicitation(id.clone(), message, schema),
        );

-        self.queued_requests
+        let sender = self
+            .action_required_senders
            .lock()
            .await
-            .entry(session_id)
-            .or_default()
-            .push_back(action_required_message);
+            .get(&(session_id.clone(), tool_call_request_id.clone()))
+            .cloned();
+
+        let Some(sender) = sender else {
+            self.pending.write().await.remove(&id);
+            return Err(anyhow::anyhow!(
+                "Tool call request not found for elicitation: {}",
+                tool_call_request_id
+            ));
+        };
+
+        if sender.send(action_required_message).await.is_err() {
+            self.pending.write().await.remove(&id);
+            return Err(anyhow::anyhow!(
+                "Tool call action-required stream closed: {}",
+                tool_call_request_id
+            ));
+        }

        let result = self
            .wait_for_response(&id, pending_request, rx, timeout_duration)
@@ -179,13 +196,39 @@ impl ActionRequiredManager {
        }
    }

-    pub(crate) async fn drain_requests_for_session(&self, session_id: &str) -> Vec<Message> {
-        self.queued_requests
+    pub(crate) async fn register_action_required_stream(
+        &self,
+        session_id: String,
+        tool_call_request_id: String,
+    ) -> mpsc::Receiver<Message> {
+        let (tx, rx) = mpsc::channel(8);
+        self.action_required_senders
            .lock()
            .await
-            .remove(session_id)
-            .map(|queue| queue.into_iter().collect())
-            .unwrap_or_default()
+            .insert((session_id, tool_call_request_id), tx);
+        rx
+    }
+
+    pub(crate) async fn has_action_required_stream(
+        &self,
+        session_id: &str,
+        tool_call_request_id: &str,
+    ) -> bool {
+        self.action_required_senders
+            .lock()
+            .await
+            .contains_key(&(session_id.to_string(), tool_call_request_id.to_string()))
+    }
+
+    pub(crate) async fn unregister_action_required_stream(
+        &self,
+        session_id: &str,
+        tool_call_request_id: &str,
+    ) {
+        self.action_required_senders
+            .lock()
+            .await
+            .remove(&(session_id.to_string(), tool_call_request_id.to_string()));
    }
 }

@@ -205,32 +248,26 @@ mod tests {
        }
    }

-    async fn wait_for_elicitation_messages(
-        manager: &ActionRequiredManager,
-        session_id: &str,
-    ) -> Vec<Message> {
-        tokio::time::timeout(Duration::from_secs(1), async {
-            loop {
-                let messages = manager.drain_requests_for_session(session_id).await;
-                if !messages.is_empty() {
-                    return messages;
-                }
-                tokio::task::yield_now().await;
-            }
-        })
-        .await
-        .unwrap_or_else(|_| panic!("timed out waiting for elicitation message for {session_id}"))
+    async fn recv_elicitation_message(rx: &mut mpsc::Receiver<Message>) -> Message {
+        tokio::time::timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .expect("timed out waiting for elicitation message")
+            .expect("action-required stream closed")
    }

    #[tokio::test]
    async fn wrong_session_does_not_consume_pending_response() {
        let manager = Arc::new(ActionRequiredManager::new());
+        let mut action_required_rx = manager
+            .register_action_required_stream("session-a".to_string(), "tool-call-a".to_string())
+            .await;
        let waiter = {
            let manager = manager.clone();
            tokio::spawn(async move {
                manager
                    .request_and_wait(
                        "session-a".to_string(),
+                        "tool-call-a".to_string(),
                        "Need input".to_string(),
                        json!({ "type": "object" }),
                        Duration::from_secs(5),
@@ -239,9 +276,8 @@ mod tests {
            })
        };

-        let messages = wait_for_elicitation_messages(&manager, "session-a").await;
-        assert_eq!(messages.len(), 1);
-        let request_id = elicitation_id(&messages[0]);
+        let message = recv_elicitation_message(&mut action_required_rx).await;
+        let request_id = elicitation_id(&message);

        let err = match manager.claim_response("session-b", &request_id).await {
            Ok(_) => panic!("wrong session should not claim pending response"),
@@ -264,14 +300,21 @@ mod tests {
    }

    #[tokio::test]
-    async fn drains_only_requested_session() {
+    async fn streams_only_requested_tool_call() {
        let manager = Arc::new(ActionRequiredManager::new());
+        let mut stream_a = manager
+            .register_action_required_stream("session-a".to_string(), "tool-call-a".to_string())
+            .await;
+        let mut stream_b = manager
+            .register_action_required_stream("session-b".to_string(), "tool-call-b".to_string())
+            .await;
        let waiter_a = {
            let manager = manager.clone();
            tokio::spawn(async move {
                manager
                    .request_and_wait(
                        "session-a".to_string(),
+                        "tool-call-a".to_string(),
                        "Need input A".to_string(),
                        json!({ "type": "object" }),
                        Duration::from_secs(5),
@@ -285,6 +328,7 @@ mod tests {
                manager
                    .request_and_wait(
                        "session-b".to_string(),
+                        "tool-call-b".to_string(),
                        "Need input B".to_string(),
                        json!({ "type": "object" }),
                        Duration::from_secs(5),
@@ -293,16 +337,78 @@ mod tests {
            })
        };

-        let session_a_messages = wait_for_elicitation_messages(&manager, "session-a").await;
-        assert_eq!(session_a_messages.len(), 1);
-        let request_id_a = elicitation_id(&session_a_messages[0]);
+        let message_a = recv_elicitation_message(&mut stream_a).await;
+        let request_id_a = elicitation_id(&message_a);
+        assert!(stream_a.try_recv().is_err());

-        let empty_messages = manager.drain_requests_for_session("session-a").await;
-        assert!(empty_messages.is_empty());
+        let message_b = recv_elicitation_message(&mut stream_b).await;
+        let request_id_b = elicitation_id(&message_b);

-        let session_b_messages = wait_for_elicitation_messages(&manager, "session-b").await;
-        assert_eq!(session_b_messages.len(), 1);
-        let request_id_b = elicitation_id(&session_b_messages[0]);
+        manager
+            .claim_response("session-a", &request_id_a)
+            .await
+            .unwrap()
+            .submit(ElicitationOutcome::Accept(json!({ "answer": "a" })))
+            .unwrap();
+        manager
+            .claim_response("session-b", &request_id_b)
+            .await
+            .unwrap()
+            .submit(ElicitationOutcome::Accept(json!({ "answer": "b" })))
+            .unwrap();
+
+        assert_eq!(
+            waiter_a.await.unwrap().unwrap(),
+            ElicitationOutcome::Accept(json!({ "answer": "a" }))
+        );
+        assert_eq!(
+            waiter_b.await.unwrap().unwrap(),
+            ElicitationOutcome::Accept(json!({ "answer": "b" }))
+        );
+    }
+
+    #[tokio::test]
+    async fn streams_are_namespaced_by_session() {
+        let manager = Arc::new(ActionRequiredManager::new());
+        let mut stream_a = manager
+            .register_action_required_stream("session-a".to_string(), "tool-call-a".to_string())
+            .await;
+        let mut stream_b = manager
+            .register_action_required_stream("session-b".to_string(), "tool-call-a".to_string())
+            .await;
+        let waiter_a = {
+            let manager = manager.clone();
+            tokio::spawn(async move {
+                manager
+                    .request_and_wait(
+                        "session-a".to_string(),
+                        "tool-call-a".to_string(),
+                        "Need input A".to_string(),
+                        json!({ "type": "object" }),
+                        Duration::from_secs(5),
+                    )
+                    .await
+            })
+        };
+        let waiter_b = {
+            let manager = manager.clone();
+            tokio::spawn(async move {
+                manager
+                    .request_and_wait(
+                        "session-b".to_string(),
+                        "tool-call-a".to_string(),
+                        "Need input B".to_string(),
+                        json!({ "type": "object" }),
+                        Duration::from_secs(5),
+                    )
+                    .await
+            })
+        };
+
+        let message_a = recv_elicitation_message(&mut stream_a).await;
+        let request_id_a = elicitation_id(&message_a);
+        let message_b = recv_elicitation_message(&mut stream_b).await;
+        let request_id_b = elicitation_id(&message_b);

        manager
            .claim_response("session-a", &request_id_a)
@@ -330,12 +436,16 @@ mod tests {
    #[tokio::test]
    async fn claimed_response_can_complete_after_timeout_deadline() {
        let manager = Arc::new(ActionRequiredManager::new());
+        let mut action_required_rx = manager
+            .register_action_required_stream("session-a".to_string(), "tool-call-a".to_string())
+            .await;
        let waiter = {
            let manager = manager.clone();
            tokio::spawn(async move {
                manager
                    .request_and_wait(
                        "session-a".to_string(),
+                        "tool-call-a".to_string(),
                        "Need input".to_string(),
                        json!({ "type": "object" }),
                        Duration::from_millis(25),
@@ -344,9 +454,8 @@ mod tests {
            })
        };

-        let messages = wait_for_elicitation_messages(&manager, "session-a").await;
-        assert_eq!(messages.len(), 1);
-        let request_id = elicitation_id(&messages[0]);
+        let message = recv_elicitation_message(&mut action_required_rx).await;
+        let request_id = elicitation_id(&message);

        let claim = manager
            .claim_response("session-a", &request_id)
@@ -368,12 +477,19 @@ mod tests {
    #[tokio::test]
    async fn request_and_wait_returns_decline_and_cancel_actions() {
        let manager = Arc::new(ActionRequiredManager::new());
+        let mut decline_rx = manager
+            .register_action_required_stream("session-a".to_string(), "tool-call-a".to_string())
+            .await;
+        let mut cancel_rx = manager
+            .register_action_required_stream("session-b".to_string(), "tool-call-b".to_string())
+            .await;
        let decline_waiter = {
            let manager = manager.clone();
            tokio::spawn(async move {
                manager
                    .request_and_wait(
                        "session-a".to_string(),
+                        "tool-call-a".to_string(),
                        "Need input A".to_string(),
                        json!({ "type": "object" }),
                        Duration::from_secs(5),
@@ -387,6 +503,7 @@ mod tests {
                manager
                    .request_and_wait(
                        "session-b".to_string(),
+                        "tool-call-b".to_string(),
                        "Need input B".to_string(),
                        json!({ "type": "object" }),
                        Duration::from_secs(5),
@@ -395,10 +512,10 @@ mod tests {
            })
        };

-        let decline_messages = wait_for_elicitation_messages(&manager, "session-a").await;
-        let decline_request_id = elicitation_id(&decline_messages[0]);
-        let cancel_messages = wait_for_elicitation_messages(&manager, "session-b").await;
-        let cancel_request_id = elicitation_id(&cancel_messages[0]);
+        let decline_message = recv_elicitation_message(&mut decline_rx).await;
+        let decline_request_id = elicitation_id(&decline_message);
+        let cancel_message = recv_elicitation_message(&mut cancel_rx).await;
+        let cancel_request_id = elicitation_id(&cancel_message);

        manager
            .claim_response("session-a", &decline_request_id)
@@ -422,4 +539,48 @@ mod tests {
            ElicitationOutcome::Cancel
        );
    }
+
+    #[tokio::test]
+    async fn missing_tool_call_stream_errors() {
+        let manager = Arc::new(ActionRequiredManager::new());
+
+        let result = manager
+            .request_and_wait(
+                "session-a".to_string(),
+                "missing-tool-call".to_string(),
+                "Need input".to_string(),
+                json!({ "type": "object" }),
+                Duration::from_secs(5),
+            )
+            .await;
+
+        let err = result.expect_err("request should fail without a registered stream");
+        assert!(err
+            .to_string()
+            .contains("Tool call request not found for elicitation"));
+    }
+
+    #[tokio::test]
+    async fn closed_tool_call_stream_errors() {
+        let manager = Arc::new(ActionRequiredManager::new());
+        let rx = manager
+            .register_action_required_stream("session-a".to_string(), "tool-call-a".to_string())
+            .await;
+        drop(rx);
+
+        let result = manager
+            .request_and_wait(
+                "session-a".to_string(),
+                "tool-call-a".to_string(),
+                "Need input".to_string(),
+                json!({ "type": "object" }),
+                Duration::from_secs(5),
+            )
+            .await;
+
+        let err = result.expect_err("request should fail when stream is closed");
+        assert!(err
+            .to_string()
+            .contains("Tool call action-required stream closed"));
+    }
 }
@@ -16,7 +16,7 @@ use super::mcp_client::GooseMcpHostInfo;
 use super::platform_tools;
 use super::tool_confirmation_router::ToolConfirmationRouter;
 use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
-use crate::action_required_manager::{ActionRequiredManager, ElicitationOutcome};
+use crate::action_required_manager::ElicitationOutcome;
 use crate::agents::extension::{ExtensionConfig, ExtensionResult, ToolInfo};
 use crate::agents::extension_manager::{
    get_parameter_names, ExtensionManager, ExtensionManagerCapabilities,
@@ -274,6 +274,7 @@ impl Default for Agent {
 }

 pub enum ToolStreamItem<T> {
+    ActionRequired(Message),
    Message(ServerNotification),
    Result(T),
 }
@@ -285,17 +286,22 @@ pub type ToolStream =
 // final result of the tool call. MCP notifications are not request-scoped, but
 // this lets us capture all notifications emitted during the tool call for
 // simpler consumption
-pub fn tool_stream<S, F>(rx: S, done: F) -> ToolStream
+pub fn tool_stream<S, A, F>(rx: S, action_required_rx: A, done: F) -> ToolStream
 where
    S: Stream<Item = ServerNotification> + Send + Unpin + 'static,
+    A: Stream<Item = Message> + Send + Unpin + 'static,
    F: Future<Output = ToolResult<CallToolResult>> + Send + 'static,
 {
    Box::pin(async_stream::stream! {
        tokio::pin!(done);
        let mut rx = rx;
+        let mut action_required_rx = action_required_rx;

        loop {
            tokio::select! {
+                Some(msg) = action_required_rx.next() => {
+                    yield ToolStreamItem::ActionRequired(msg);
+                }
                Some(msg) = rx.next() => {
                    yield ToolStreamItem::Message(msg);
                }
@@ -572,6 +578,7 @@ impl Agent {

        ToolCallResult {
            notification_stream: result.notification_stream,
+            action_required_stream: result.action_required_stream,
            result: Box::new(fut.boxed()),
        }
    }
@@ -645,24 +652,6 @@ impl Agent {
            | RetryResult::SuccessChecksPassed => Ok(false),
        }
    }
-    async fn drain_elicitation_messages(&self, session_id: &str) -> Vec<Message> {
-        let mut messages = Vec::new();
-        let manager = self.config.session_manager.clone();
-        for mut elicitation_message in ActionRequiredManager::global()
-            .drain_requests_for_session(session_id)
-            .await
-        {
-            if elicitation_message.id.is_none() {
-                elicitation_message = elicitation_message.with_generated_id();
-            }
-            if let Err(e) = manager.add_message(session_id, &elicitation_message).await {
-                warn!("Failed to save elicitation message to session: {}", e);
-            }
-            messages.push(elicitation_message);
-        }
-        messages
-    }
-
    async fn load_project_instructions(&self, session: &Session) -> Option<String> {
        let project_id = session.project_id.as_deref()?;
        let entry = crate::sources::read_project(project_id).ok()?;
@@ -783,11 +772,16 @@ impl Agent {
                            result
                                .notification_stream
                                .unwrap_or_else(|| Box::new(stream::empty())),
+                            result
+                                .action_required_stream
+                                .unwrap_or_else(|| Box::new(stream::empty())),
                            result.result,
                        ),
-                        Err(e) => {
-                            tool_stream(Box::new(stream::empty()), futures::future::ready(Err(e)))
-                        }
+                        Err(e) => tool_stream(
+                            Box::new(stream::empty()),
+                            Box::new(stream::empty()),
+                            futures::future::ready(Err(e)),
+                        ),
                    },
                ));
            }
@@ -2152,10 +2146,6 @@ impl Agent {
                                            break;
                                        }

-                                        for msg in self.drain_elicitation_messages(&session_config.id).await {
-                                            yield AgentEvent::Message(msg);
-                                        }
-
                                        tokio::select! {
                                            biased;

@@ -2163,6 +2153,15 @@ impl Agent {
                                                match tool_item {
                                                    Some((request_id, item)) => {
                                                        match item {
+                                                            ToolStreamItem::ActionRequired(mut msg) => {
+                                                                if msg.id.is_none() {
+                                                                    msg = msg.with_generated_id();
+                                                                }
+                                                                if let Err(e) = session_manager.add_message(&session_config.id, &msg).await {
+                                                                    warn!("Failed to save elicitation message to session: {}", e);
+                                                                }
+                                                                yield AgentEvent::Message(msg);
+                                                            }
                                                            ToolStreamItem::Result(output) => {
                                                                if let Ok(ref call_result) = output {
                                                                    if let Some(ref meta) = call_result.meta {
@@ -2200,17 +2199,10 @@ impl Agent {
                                                }
                                            }

-                                            _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
-                                                // Continue loop to drain elicitation messages
-                                            }
+                                            _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {}
                                        }
                                    }

-                                    // check for remaining elicitation messages after all tools complete
-                                    for msg in self.drain_elicitation_messages(&session_config.id).await {
-                                        yield AgentEvent::Message(msg);
-                                    }
-
                                    if all_install_successful && !enable_extension_request_ids.is_empty() {
                                        if let Err(e) = self.save_extension_state(&session_config).await {
                                            warn!("Failed to save extension state after runtime changes: {}", e);
@@ -2,6 +2,7 @@ use anyhow::Result;
 use axum::http::{HeaderMap, HeaderName, HeaderValue};
 use chrono::{DateTime, Utc};
 use futures::stream::{FuturesUnordered, StreamExt};
+use futures::Stream;
 use futures::{future, FutureExt};
 use once_cell::sync::Lazy;
 use rmcp::service::{ClientInitializeError, ServiceError};
@@ -13,9 +14,11 @@ use rmcp::transport::{
 };
 use std::collections::HashMap;
 use std::path::PathBuf;
+use std::pin::Pin;
 use std::process::Stdio;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
+use std::task::{Context, Poll};
 use std::time::Duration;
 use tempfile::{tempdir, TempDir};
 use tokio::io::AsyncReadExt;
@@ -32,6 +35,7 @@ use super::extension::{
 };
 use super::tool_execution::{ToolCallContext, ToolCallResult};
 use super::types::SharedProvider;
+use crate::action_required_manager::ActionRequiredManager;
 use crate::agents::extension::{Envs, ProcessExit};
 use crate::agents::extension_malware_check;
 use crate::agents::mcp_client::{
@@ -54,6 +58,49 @@ use serde_json::Value;

 type McpClientBox = Arc<dyn McpClientTrait>;

+struct ActionRequiredStream {
+    inner: ReceiverStream<crate::conversation::message::Message>,
+    session_id: String,
+    tool_call_request_id: String,
+}
+
+impl ActionRequiredStream {
+    fn new(
+        receiver: tokio::sync::mpsc::Receiver<crate::conversation::message::Message>,
+        session_id: String,
+        tool_call_request_id: String,
+    ) -> Self {
+        Self {
+            inner: ReceiverStream::new(receiver),
+            session_id,
+            tool_call_request_id,
+        }
+    }
+}
+
+impl Stream for ActionRequiredStream {
+    type Item = crate::conversation::message::Message;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        Pin::new(&mut self.inner).poll_next(cx)
+    }
+}
+
+impl Drop for ActionRequiredStream {
+    fn drop(&mut self) {
+        let session_id = self.session_id.clone();
+        let tool_call_request_id = self.tool_call_request_id.clone();
+        let Ok(handle) = tokio::runtime::Handle::try_current() else {
+            return;
+        };
+        handle.spawn(async move {
+            ActionRequiredManager::global()
+                .unregister_action_required_stream(&session_id, &tool_call_request_id)
+                .await;
+        });
+    }
+}
+
 static RE_ENV_BRACES: Lazy<regex::Regex> =
    Lazy::new(|| regex::Regex::new(r"\$\{\s*([A-Za-z_][A-Za-z0-9_]*)\s*\}").expect("valid regex"));

@@ -1733,11 +1780,33 @@ impl ExtensionManager {
        let client = resolved.client.clone();
        let hydration_client = client.clone();
        let notifications_receiver = client.subscribe().await;
+        let session_id = ctx.session_id.clone();
+        let action_required_tool_call_request_id = ctx.tool_call_request_id.clone();
+        let action_required_receiver =
+            if let Some(tool_call_request_id) = action_required_tool_call_request_id.clone() {
+                if ActionRequiredManager::global()
+                    .has_action_required_stream(&session_id, &tool_call_request_id)
+                    .await
+                {
+                    None
+                } else {
+                    let registered_tool_call_request_id = tool_call_request_id.clone();
+                    let receiver = ActionRequiredManager::global()
+                        .register_action_required_stream(session_id.clone(), tool_call_request_id)
+                        .await;
+                    Some((
+                        receiver,
+                        session_id.clone(),
+                        registered_tool_call_request_id,
+                    ))
+                }
+            } else {
+                None
+            };
        let actual_tool_name = resolved.actual_tool_name.clone();
        let resolved_tool = resolved;
        let should_hydrate_mcp_app = self.host_supports_mcp_apps();
        let read_cancellation_token = cancellation_token.clone();
-        let session_id = ctx.session_id.clone();
        let owned_ctx = ToolCallContext::new(
            ctx.session_id.clone(),
            ctx.working_dir.clone(),
@@ -1751,7 +1820,7 @@ impl ExtensionManager {
                owned_ctx.session_id,
                owned_ctx.working_dir,
            );
-            let mut result = client
+            let call_result = client
                .call_tool(&owned_ctx, &actual_tool_name, arguments, cancellation_token)
                .await
                .map_err(|e| match e {
@@ -1759,7 +1828,9 @@ impl ExtensionManager {
                    _ => {
                        ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), e.maybe_to_value())
                    }
-                })?;
+                });
+
+            let mut result = call_result?;

            remove_untrusted_mcp_app_meta(&mut result);

@@ -1782,6 +1853,15 @@ impl ExtensionManager {
        Ok(ToolCallResult {
            result: Box::new(fut.boxed()),
            notification_stream: Some(Box::new(ReceiverStream::new(notifications_receiver))),
+            action_required_stream: action_required_receiver.map(
+                |(rx, session_id, tool_call_request_id)| {
+                    Box::new(ActionRequiredStream::new(
+                        rx,
+                        session_id,
+                        tool_call_request_id,
+                    )) as _
+                },
+            ),
        })
    }

@@ -1,7 +1,7 @@
 use crate::action_required_manager::{ActionRequiredManager, ElicitationOutcome};
 use crate::agents::tool_execution::ToolCallContext;
 use crate::agents::types::SharedProvider;
-use crate::session_context::{SESSION_ID_HEADER, WORKING_DIR_HEADER};
+use crate::session_context::{SESSION_ID_HEADER, TOOL_CALL_REQUEST_ID_HEADER, WORKING_DIR_HEADER};
 use rmcp::model::{
    CreateElicitationRequestParams, CreateElicitationResult, ElicitationAction, ErrorCode,
    ExtensionCapabilities, Extensions, JsonObject, ListRootsResult, LoggingMessageNotification,
@@ -26,7 +26,9 @@ use rmcp::{
    ClientHandler, ErrorData, Peer, RoleClient, ServiceError, ServiceExt,
 };
 use serde_json::Value;
-use std::{path::PathBuf, sync::Arc, time::Duration};
+use std::{
+    collections::HashMap, path::PathBuf, sync::Arc, sync::Mutex as StdMutex, time::Duration,
+};
 use tokio::sync::{
    mpsc::{self, Sender},
    Mutex,
@@ -148,10 +150,34 @@ pub trait McpClientTrait: Send + Sync {
    }
 }

+struct ActiveToolCallGuard {
+    active_tool_calls: Arc<StdMutex<HashMap<String, Vec<String>>>>,
+    session_id: String,
+    tool_call_request_id: String,
+}
+
+impl Drop for ActiveToolCallGuard {
+    fn drop(&mut self) {
+        let mut active_tool_calls = self
+            .active_tool_calls
+            .lock()
+            .expect("active_tool_calls mutex poisoned");
+        if let Some(calls) = active_tool_calls.get_mut(&self.session_id) {
+            if let Some(pos) = calls.iter().position(|id| id == &self.tool_call_request_id) {
+                calls.remove(pos);
+            }
+            if calls.is_empty() {
+                active_tool_calls.remove(&self.session_id);
+            }
+        }
+    }
+}
+
 pub struct GooseClient {
    notification_handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>,
    provider: SharedProvider,
    session_id: Mutex<Option<String>>,
+    active_tool_calls: Arc<StdMutex<HashMap<String, Vec<String>>>>,
    client_name: String,
    capabilities: GooseMcpClientCapabilities,
    working_dir: Arc<tokio::sync::RwLock<PathBuf>>,
@@ -169,6 +195,7 @@ impl GooseClient {
            notification_handlers: handlers,
            provider,
            session_id: Mutex::new(None),
+            active_tool_calls: Arc::new(StdMutex::new(HashMap::new())),
            client_name,
            capabilities,
            working_dir: Arc::new(tokio::sync::RwLock::new(working_dir)),
@@ -207,6 +234,62 @@ impl GooseClient {
            .map(|value| value.to_string())
    }

+    fn tool_call_request_id_from_extensions(extensions: &Extensions) -> Option<String> {
+        let meta = extensions.get::<Meta>()?;
+        meta.0
+            .iter()
+            .find(|(key, _)| key.eq_ignore_ascii_case(TOOL_CALL_REQUEST_ID_HEADER))
+            .and_then(|(_, value)| value.as_str())
+            .map(|value| value.to_string())
+    }
+
+    fn register_active_tool_call(
+        &self,
+        session_id: &str,
+        tool_call_request_id: &str,
+    ) -> ActiveToolCallGuard {
+        self.active_tool_calls
+            .lock()
+            .expect("active_tool_calls mutex poisoned")
+            .entry(session_id.to_string())
+            .or_default()
+            .push(tool_call_request_id.to_string());
+        ActiveToolCallGuard {
+            active_tool_calls: self.active_tool_calls.clone(),
+            session_id: session_id.to_string(),
+            tool_call_request_id: tool_call_request_id.to_string(),
+        }
+    }
+
+    fn resolve_tool_call_request_id(
+        &self,
+        session_id: &str,
+        extensions: &Extensions,
+    ) -> Result<String, ErrorData> {
+        if let Some(tool_call_request_id) = Self::tool_call_request_id_from_extensions(extensions) {
+            return Ok(tool_call_request_id);
+        }
+
+        let active_tool_calls = self
+            .active_tool_calls
+            .lock()
+            .expect("active_tool_calls mutex poisoned");
+        match active_tool_calls.get(session_id).map(Vec::as_slice) {
+            Some([tool_call_request_id]) => Ok(tool_call_request_id.clone()),
+            Some(calls) if calls.len() > 1 => Err(ErrorData::new(
+                ErrorCode::INTERNAL_ERROR,
+                "Cannot correlate elicitation request: multiple tool calls are active and the \
+                 server did not echo the tool call request id",
+                None,
+            )),
+            _ => Err(ErrorData::new(
+                ErrorCode::INTERNAL_ERROR,
+                "Could not resolve tool call request id for elicitation request",
+                None,
+            )),
+        }
+    }
+
    fn resolved_extensions(&self) -> ExtensionCapabilities {
        if let Some(host_info) = &self.capabilities.host_info {
            if host_info.explicit_extensions {
@@ -396,6 +479,8 @@ impl ClientHandler for GooseClient {
                    None,
                )
            })?;
+        let tool_call_request_id =
+            self.resolve_tool_call_request_id(&session_id, &context.extensions)?;

        let (message, schema_value) = match &request {
            CreateElicitationRequestParams::FormElicitationParams {
@@ -418,7 +503,13 @@ impl ClientHandler for GooseClient {
        };

        ActionRequiredManager::global()
-            .request_and_wait(session_id, message, schema_value, Duration::from_secs(300))
+            .request_and_wait(
+                session_id,
+                tool_call_request_id,
+                message,
+                schema_value,
+                Duration::from_secs(300),
+            )
            .await
            .map(|response| match response {
                ElicitationOutcome::Accept(user_data) => {
@@ -548,19 +639,34 @@ impl McpClient {
        &self,
        session_id: &str,
        working_dir: Option<&str>,
+        tool_call_request_id: Option<&str>,
        request: ClientRequest,
        cancel_token: CancellationToken,
    ) -> Result<ServerResult, Error> {
-        let request = inject_session_context_into_request(request, Some(session_id), working_dir);
+        let request = inject_session_context_into_request(
+            request,
+            Some(session_id),
+            working_dir,
+            tool_call_request_id,
+        );
+        let active_tool_call = tool_call_request_id.filter(|id| !id.is_empty());
        // The inner mutex is held only for the send; the actual response wait
-        // happens outside the lock so concurrent calls can overlap.
-        let handle = {
+        // happens outside the lock so concurrent calls can overlap. The guard
+        // unregisters the active tool call on drop, covering cancellation and
+        // dropped reply streams as well as normal completion.
+        let (handle, _active_tool_call_guard) = {
            let client = self.client.lock().await;
            client.service().set_session_id(session_id).await;
-            client
+            let guard = active_tool_call.map(|tool_call_request_id| {
+                client
+                    .service()
+                    .register_active_tool_call(session_id, tool_call_request_id)
+            });
+            let handle = client
                .send_cancellable_request(request, PeerRequestOptions::no_options())
-                .await
-        }?;
+                .await?;
+            (handle, guard)
+        };

        await_response(handle, self.timeout, &cancel_token).await
    }
@@ -616,6 +722,7 @@ impl McpClientTrait for McpClient {
            .send_request_with_context(
                session_id,
                None,
+                None,
                ClientRequest::ListResourcesRequest(RequestOptionalParam::with_param(
                    PaginatedRequestParams::default().with_cursor(cursor),
                )),
@@ -639,6 +746,7 @@ impl McpClientTrait for McpClient {
            .send_request_with_context(
                session_id,
                None,
+                None,
                ClientRequest::ReadResourceRequest(Request::new(ReadResourceRequestParams::new(
                    uri.to_string(),
                ))),
@@ -662,6 +770,7 @@ impl McpClientTrait for McpClient {
            .send_request_with_context(
                session_id,
                None,
+                None,
                ClientRequest::ListToolsRequest(RequestOptionalParam::with_param(
                    PaginatedRequestParams::default().with_cursor(cursor),
                )),
@@ -692,6 +801,7 @@ impl McpClientTrait for McpClient {
            .send_request_with_context(
                &ctx.session_id,
                ctx.working_dir_str(),
+                ctx.tool_call_request_id.as_deref(),
                request,
                cancel_token,
            )
@@ -713,6 +823,7 @@ impl McpClientTrait for McpClient {
            .send_request_with_context(
                session_id,
                None,
+                None,
                ClientRequest::ListPromptsRequest(RequestOptionalParam::with_param(
                    PaginatedRequestParams::default().with_cursor(cursor),
                )),
@@ -745,6 +856,7 @@ impl McpClientTrait for McpClient {
            .send_request_with_context(
                session_id,
                None,
+                None,
                ClientRequest::GetPromptRequest(Request::new(params)),
                cancel_token,
            )
@@ -773,9 +885,11 @@ fn inject_session_context_into_extensions(
    mut extensions: Extensions,
    session_id: Option<&str>,
    working_dir: Option<&str>,
+    tool_call_request_id: Option<&str>,
 ) -> Extensions {
    let session_id = session_id.filter(|id| !id.is_empty());
    let working_dir = working_dir.filter(|dir| !dir.is_empty());
+    let tool_call_request_id = tool_call_request_id.filter(|id| !id.is_empty());
    let mut meta_map = extensions
        .get::<Meta>()
        .map(|meta| meta.0.clone())
@@ -783,7 +897,9 @@ fn inject_session_context_into_extensions(

    // JsonObject is case-sensitive, so we use retain for case-insensitive removal
    meta_map.retain(|k, _| {
-        !k.eq_ignore_ascii_case(SESSION_ID_HEADER) && !k.eq_ignore_ascii_case(WORKING_DIR_HEADER)
+        !k.eq_ignore_ascii_case(SESSION_ID_HEADER)
+            && !k.eq_ignore_ascii_case(WORKING_DIR_HEADER)
+            && !k.eq_ignore_ascii_case(TOOL_CALL_REQUEST_ID_HEADER)
    });

    if let Some(session_id) = session_id {
@@ -800,6 +916,13 @@ fn inject_session_context_into_extensions(
        );
    }

+    if let Some(tool_call_request_id) = tool_call_request_id {
+        meta_map.insert(
+            TOOL_CALL_REQUEST_ID_HEADER.to_string(),
+            Value::String(tool_call_request_id.to_string()),
+        );
+    }
+
    extensions.insert(Meta(meta_map));
    extensions
 }
@@ -808,36 +931,61 @@ fn inject_session_context_into_request(
    request: ClientRequest,
    session_id: Option<&str>,
    working_dir: Option<&str>,
+    tool_call_request_id: Option<&str>,
 ) -> ClientRequest {
    match request {
        ClientRequest::ListResourcesRequest(mut req) => {
-            req.extensions =
-                inject_session_context_into_extensions(req.extensions, session_id, working_dir);
+            req.extensions = inject_session_context_into_extensions(
+                req.extensions,
+                session_id,
+                working_dir,
+                None,
+            );
            ClientRequest::ListResourcesRequest(req)
        }
        ClientRequest::ReadResourceRequest(mut req) => {
-            req.extensions =
-                inject_session_context_into_extensions(req.extensions, session_id, working_dir);
+            req.extensions = inject_session_context_into_extensions(
+                req.extensions,
+                session_id,
+                working_dir,
+                None,
+            );
            ClientRequest::ReadResourceRequest(req)
        }
        ClientRequest::ListToolsRequest(mut req) => {
-            req.extensions =
-                inject_session_context_into_extensions(req.extensions, session_id, working_dir);
+            req.extensions = inject_session_context_into_extensions(
+                req.extensions,
+                session_id,
+                working_dir,
+                None,
+            );
            ClientRequest::ListToolsRequest(req)
        }
        ClientRequest::CallToolRequest(mut req) => {
-            req.extensions =
-                inject_session_context_into_extensions(req.extensions, session_id, working_dir);
+            req.extensions = inject_session_context_into_extensions(
+                req.extensions,
+                session_id,
+                working_dir,
+                tool_call_request_id,
+            );
            ClientRequest::CallToolRequest(req)
        }
        ClientRequest::ListPromptsRequest(mut req) => {
-            req.extensions =
-                inject_session_context_into_extensions(req.extensions, session_id, working_dir);
+            req.extensions = inject_session_context_into_extensions(
+                req.extensions,
+                session_id,
+                working_dir,
+                None,
+            );
            ClientRequest::ListPromptsRequest(req)
        }
        ClientRequest::GetPromptRequest(mut req) => {
-            req.extensions =
-                inject_session_context_into_extensions(req.extensions, session_id, working_dir);
+            req.extensions = inject_session_context_into_extensions(
+                req.extensions,
+                session_id,
+                working_dir,
+                None,
+            );
            ClientRequest::GetPromptRequest(req)
        }
        other => other,
@@ -953,7 +1101,7 @@ mod tests {
            }

            let extensions =
-                inject_session_context_into_extensions(Extensions::new(), ext_session, None);
+                inject_session_context_into_extensions(Extensions::new(), ext_session, None, None);

            let resolved = client.resolve_session_id(&extensions).await;

@@ -962,6 +1110,83 @@ mod tests {
        });
    }

+    #[test]
+    fn test_resolve_tool_call_request_id_from_extensions() {
+        let client = new_client(GoosePlatform::GooseCli);
+        let _guard = client.register_active_tool_call("session-a", "active-tool-call");
+        let extensions = inject_session_context_into_extensions(
+            Extensions::new(),
+            Some("session-a"),
+            None,
+            Some("extension-tool-call"),
+        );
+
+        let resolved = client
+            .resolve_tool_call_request_id("session-a", &extensions)
+            .unwrap();
+
+        assert_eq!(resolved, "extension-tool-call");
+    }
+
+    #[test]
+    fn test_resolve_tool_call_request_id_from_active_call() {
+        let client = new_client(GoosePlatform::GooseCli);
+        let _guard = client.register_active_tool_call("session-a", "active-tool-call");
+
+        let resolved = client
+            .resolve_tool_call_request_id("session-a", &Extensions::new())
+            .unwrap();
+
+        assert_eq!(resolved, "active-tool-call");
+    }
+
+    #[test]
+    fn test_resolve_tool_call_request_id_errors_when_calls_overlap() {
+        let client = new_client(GoosePlatform::GooseCli);
+        let _guard_a = client.register_active_tool_call("session-a", "active-tool-call-a");
+        let _guard_b = client.register_active_tool_call("session-a", "active-tool-call-b");
+
+        let error = client
+            .resolve_tool_call_request_id("session-a", &Extensions::new())
+            .expect_err("ambiguous elicitation should not resolve to an arbitrary call");
+
+        assert_eq!(error.code, ErrorCode::INTERNAL_ERROR);
+    }
+
+    #[test]
+    fn test_resolve_tool_call_request_id_prefers_echoed_id_while_calls_overlap() {
+        let client = new_client(GoosePlatform::GooseCli);
+        let _guard_a = client.register_active_tool_call("session-a", "active-tool-call-a");
+        let _guard_b = client.register_active_tool_call("session-a", "active-tool-call-b");
+        let extensions = inject_session_context_into_extensions(
+            Extensions::new(),
+            Some("session-a"),
+            None,
+            Some("active-tool-call-a"),
+        );
+
+        let resolved = client
+            .resolve_tool_call_request_id("session-a", &extensions)
+            .unwrap();
+
+        assert_eq!(resolved, "active-tool-call-a");
+    }
+
+    #[test]
+    fn test_dropping_guard_unregisters_active_tool_call() {
+        let client = new_client(GoosePlatform::GooseCli);
+        let guard_a = client.register_active_tool_call("session-a", "active-tool-call-a");
+        let _guard_b = client.register_active_tool_call("session-a", "active-tool-call-b");
+
+        drop(guard_a);
+
+        let resolved = client
+            .resolve_tool_call_request_id("session-a", &Extensions::new())
+            .unwrap();
+
+        assert_eq!(resolved, "active-tool-call-b");
+    }
+
    #[test_case(list_resources_request; "list_resources")]
    #[test_case(read_resource_request; "read_resource")]
    #[test_case(list_tools_request; "list_tools")]
@@ -980,7 +1205,7 @@ mod tests {
        );

        let request = request_builder(extensions);
-        let request = inject_session_context_into_request(request, Some(session_id), None);
+        let request = inject_session_context_into_request(request, Some(session_id), None, None);
        let extensions = request_extensions(&request).expect("request should have extensions");
        let meta = extensions
            .get::<Meta>()
@@ -994,13 +1219,20 @@ mod tests {
            meta.0.get("other-key"),
            Some(&Value::String("preserve-me".to_string()))
        );
+        if matches!(request, ClientRequest::CallToolRequest(_)) {
+            assert!(!meta.0.contains_key(TOOL_CALL_REQUEST_ID_HEADER));
+        }
    }

    #[test]
    fn test_session_id_in_mcp_meta() {
        let session_id = "test-session-789";
-        let extensions =
-            inject_session_context_into_extensions(Default::default(), Some(session_id), None);
+        let extensions = inject_session_context_into_extensions(
+            Default::default(),
+            Some(session_id),
+            None,
+            None,
+        );
        let mcp_meta = extensions.get::<Meta>().unwrap();

        assert_eq!(
@@ -1052,12 +1284,43 @@ mod tests {
            .unwrap(),
        );

-        let extensions = inject_session_context_into_extensions(extensions, session_id, None);
+        let extensions = inject_session_context_into_extensions(extensions, session_id, None, None);
        let mcp_meta = extensions.get::<Meta>().unwrap();

        assert_eq!(&mcp_meta.0, expected_meta.as_object().unwrap());
    }

+    #[test]
+    fn test_tool_call_request_id_injected_only_for_call_tool() {
+        let session_id = "test-session-id";
+        let tool_call_request_id = "tool-request-1";
+
+        let call_request = inject_session_context_into_request(
+            call_tool_request(Extensions::new()),
+            Some(session_id),
+            None,
+            Some(tool_call_request_id),
+        );
+        let call_meta = request_extensions(&call_request)
+            .and_then(|extensions| extensions.get::<Meta>())
+            .expect("call request should have meta");
+        assert_eq!(
+            call_meta.0.get(TOOL_CALL_REQUEST_ID_HEADER),
+            Some(&Value::String(tool_call_request_id.to_string()))
+        );
+
+        let tools_request = inject_session_context_into_request(
+            list_tools_request(Extensions::new()),
+            Some(session_id),
+            None,
+            Some(tool_call_request_id),
+        );
+        let tools_meta = request_extensions(&tools_request)
+            .and_then(|extensions| extensions.get::<Meta>())
+            .expect("list tools request should have meta");
+        assert!(!tools_meta.0.contains_key(TOOL_CALL_REQUEST_ID_HEADER));
+    }
+
    #[test]
    fn test_client_info_advertises_mcp_apps_ui_extension() {
        let client = new_client(GoosePlatform::GooseDesktop);
@@ -143,7 +143,7 @@ impl CodeExecutionClient {
    /// Build a PctxRegistry with all tool callbacks registered
    fn build_callback_registry(
        &self,
-        session_id: &str,
+        ctx: &ToolCallContext,
        code_mode: &CodeMode,
    ) -> Result<PctxRegistry, String> {
        let manager = self
@@ -163,7 +163,7 @@ impl CodeExecutionClient {
                    .unwrap_or_default(),
                &cfg.name
            );
-            let callback = create_tool_callback(session_id.to_string(), full_name, manager.clone());
+            let callback = create_tool_callback(ctx.clone(), full_name, manager.clone());
            registry
                .add_callback(&cfg.id(), callback)
                .map_err(|e| format!("Failed to register callback: {e}"))?;
@@ -236,7 +236,7 @@ impl CodeExecutionClient {
    /// Handle the execute typescript tool call
    async fn handle_execute_typescript(
        &self,
-        session_id: &str,
+        ctx: &ToolCallContext,
        arguments: Option<JsonObject>,
    ) -> Result<Vec<Content>, String> {
        let args: ExecuteWithToolGraph = arguments
@@ -245,8 +245,9 @@ impl CodeExecutionClient {
            .map_err(|e| format!("Failed to parse arguments: {e}"))?
            .ok_or("Missing arguments for execute_typescript")?;

+        let session_id = &ctx.session_id;
        let code_mode = self.get_code_mode(session_id).await?;
-        let registry = self.build_callback_registry(session_id, &code_mode)?;
+        let registry = self.build_callback_registry(ctx, &code_mode)?;
        let code = args.input.code.clone();
        let disclosure = self.disclosure;

@@ -273,12 +274,12 @@ impl CodeExecutionClient {
 }

 fn create_tool_callback(
-    session_id: String,
+    ctx: ToolCallContext,
    full_name: String,
    manager: Arc<crate::agents::ExtensionManager>,
 ) -> CallbackFn {
    Arc::new(move |args: Option<Value>| {
-        let session_id = session_id.clone();
+        let ctx = ctx.clone();
        let full_name = full_name.clone();
        let manager = manager.clone();
        Box::pin(async move {
@@ -289,7 +290,6 @@ fn create_tool_callback(
                }
                params
            };
-            let ctx = crate::agents::ToolCallContext::new(session_id, None, None);
            match manager
                .dispatch_tool_call(&ctx, tool_call, CancellationToken::new())
                .await
@@ -457,7 +457,7 @@ impl McpClientTrait for CodeExecutionClient {
                    .await
            }
            "execute_bash" => self.handle_execute_bash(session_id, arguments).await,
-            "execute_typescript" => self.handle_execute_typescript(session_id, arguments).await,
+            "execute_typescript" => self.handle_execute_typescript(ctx, arguments).await,
            _ => Err(format!("Unknown tool: {name}")),
        };

@@ -9,11 +9,13 @@ use tokio_util::sync::CancellationToken;
 use std::path::PathBuf;

 use crate::config::permission::PermissionLevel;
+use crate::conversation::message::Message;
 use crate::mcp_utils::ToolResult;
 use crate::permission::Permission;
 use rmcp::model::{Content, ServerNotification};

 /// Context passed through the tool call dispatch chain.
+#[derive(Clone)]
 pub struct ToolCallContext {
    pub session_id: String,
    pub working_dir: Option<PathBuf>,
@@ -43,6 +45,7 @@ impl ToolCallContext {
 pub struct ToolCallResult {
    pub result: Box<dyn Future<Output = ToolResult<rmcp::model::CallToolResult>> + Send + Unpin>,
    pub notification_stream: Option<Box<dyn Stream<Item = ServerNotification> + Send + Unpin>>,
+    pub action_required_stream: Option<Box<dyn Stream<Item = Message> + Send + Unpin>>,
 }

 impl From<ToolResult<rmcp::model::CallToolResult>> for ToolCallResult {
@@ -50,13 +53,14 @@ impl From<ToolResult<rmcp::model::CallToolResult>> for ToolCallResult {
        Self {
            result: Box::new(futures::future::ready(result)),
            notification_stream: None,
+            action_required_stream: None,
        }
    }
 }

 use super::agent::{tool_stream, ToolStream};
 use crate::agents::Agent;
-use crate::conversation::message::{Message, ToolRequest};
+use crate::conversation::message::ToolRequest;
 use crate::session::Session;
 use crate::tool_inspection::get_security_finding_id_from_results;

@@ -133,9 +137,11 @@ impl Agent {
                    tool_futures.push((req_id, match tool_result {
                        Ok(result) => tool_stream(
                            result.notification_stream.unwrap_or_else(|| Box::new(stream::empty())),
+                            result.action_required_stream.unwrap_or_else(|| Box::new(stream::empty())),
                            result.result,
                        ),
                        Err(e) => tool_stream(
+                            Box::new(stream::empty()),
                            Box::new(stream::empty()),
                            futures::future::ready(Err(e)),
                        ),
@@ -1,6 +1,7 @@
 use tokio::task_local;

 pub const SESSION_ID_HEADER: &str = "agent-session-id";
+pub const TOOL_CALL_REQUEST_ID_HEADER: &str = "agent-tool-call-request-id";
 pub const WORKING_DIR_HEADER: &str = "agent-working-dir";

 task_local! {
@@ -9,6 +9,6 @@ STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
 STDERR: time=2025-12-11T17:58:47.642-05:00 level=INFO msg="session initialized"
 STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":0}}}
 STDOUT: {"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"get_file_contents","description":"Get file contents from GitHub","inputSchema":{"type":"object","properties":{"owner":{"type":"string"},"repo":{"type":"string"},"path":{"type":"string"},"sha":{"type":"string"}},"required":["owner","repo","path"]}}]}}
-STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":1},"name":"get_file_contents","arguments":{"owner":"block","path":"README.md","repo":"goose","sha":"ab62b863c1666232a67048b6c4e10007a2a5b83c"}}}
+STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","agent-tool-call-request-id":"test-id","progressToken":1},"name":"get_file_contents","arguments":{"owner":"block","path":"README.md","repo":"goose","sha":"ab62b863c1666232a67048b6c4e10007a2a5b83c"}}}
 STDOUT: {"jsonrpc":"2.0","id":2,"result":{"content":[{"type":"text","text":"successfully downloaded text file (SHA: de9bdde7f260549bf3a083651842f30ab29cf4e9)"},{"type":"resource","resource":{"uri":"repo://block/goose/sha/ab62b863c1666232a67048b6c4e10007a2a5b83c/contents/README.md","mimeType":"text/plain; charset=utf-8","text":"\u003cdiv align=\"center\"\u003e\n\n# goose\n\n_a local, extensible, open source AI agent that automates engineering tasks_\n\n\u003cp align=\"center\"\u003e\n  \u003ca href=\"https://opensource.org/licenses/Apache-2.0\"\u003e\n    \u003cimg src=\"https://img.shields.io/badge/License-Apache_2.0-blue.svg\"\u003e\n  \u003c/a\u003e\n  \u003ca href=\"https://discord.gg/7GaTvbDwga\"\u003e\n    \u003cimg src=\"https://img.shields.io/discord/1287729918100246654?logo=discord\u0026logoColor=white\u0026label=Join+Us\u0026color=blueviolet\" alt=\"Discord\"\u003e\n  \u003c/a\u003e\n  \u003ca href=\"https://github.com/block/goose/actions/workflows/ci.yml\"\u003e\n     \u003cimg src=\"https://img.shields.io/github/actions/workflow/status/block/goose/ci.yml?branch=main\" alt=\"CI\"\u003e\n  \u003c/a\u003e\n\u003c/p\u003e\n\u003c/div\u003e\n\ngoose is your on-machine AI agent, capable of automating complex development tasks from start to finish. More than just code suggestions, goose can build entire projects from scratch, write and execute code, debug failures, orchestrate workflows, and interact with external APIs - _autonomously_.\n\nWhether you're prototyping an idea, refining existing code, or managing intricate engineering pipelines, goose adapts to your workflow and executes tasks with precision.\n\nDesigned for maximum flexibility, goose works with any LLM and supports multi-model configuration to optimize performance and cost, seamlessly integrates with MCP servers, and is available as both a desktop app as well as CLI - making it the ultimate AI assistant for developers who want to move faster and focus on innovation.\n\n[![Watch the video](https://github.com/user-attachments/assets/ddc71240-3928-41b5-8210-626dfb28af7a)](https://youtu.be/D-DpDunrbpo)\n\n# Quick Links\n- [Quickstart](https://goose-docs.ai/docs/quickstart)\n- [Installation](https://goose-docs.ai/docs/getting-started/installation)\n- [Tutorials](https://goose-docs.ai/docs/category/tutorials)\n- [Documentation](https://goose-docs.ai/docs/category/getting-started)\n\n\n# a little goose humor 🦢\n\n\u003e Why did the developer choose goose as their AI agent?\n\u003e \n\u003e Because it always helps them \"migrate\" their code to production! 🚀\n\n# goose around with us\n- [Discord](https://discord.gg/block-opensource)\n- [YouTube](https://www.youtube.com/@goose-oss)\n- [LinkedIn](https://www.linkedin.com/company/goose-oss)\n- [Twitter/X](https://x.com/goose_oss)\n- [Bluesky](https://bsky.app/profile/opensource.block.xyz)\n- [Nostr](https://njump.me/opensource@block.xyz)\n"}}]}}
 STDERR: time=2025-12-11T17:58:48.133-05:00 level=INFO msg="server session disconnected" session_id=""
@@ -6,20 +6,20 @@ STDOUT: {"method":"notifications/tools/list_changed","jsonrpc":"2.0"}
 STDOUT: {"method":"notifications/tools/list_changed","jsonrpc":"2.0"}
 STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":0}}}
 STDOUT: {"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"echo","description":"Echo a message","inputSchema":{"type":"object","properties":{"message":{"type":"string"}},"required":["message"]}},{"name":"get-sum","description":"Get the sum of two numbers","inputSchema":{"type":"object","properties":{"a":{"type":"number"},"b":{"type":"number"}},"required":["a","b"]}},{"name":"trigger-long-running-operation","description":"Trigger a long-running operation","inputSchema":{"type":"object","properties":{"duration":{"type":"number"},"steps":{"type":"number"}},"required":["duration","steps"]}},{"name":"get-structured-content","description":"Get structured content","inputSchema":{"type":"object","properties":{"location":{"type":"string"}},"required":["location"]}},{"name":"trigger-sampling-request","description":"Trigger a sampling request","inputSchema":{"type":"object","properties":{"prompt":{"type":"string"},"maxTokens":{"type":"number"}},"required":["prompt","maxTokens"]}}]}}
-STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":1},"name":"echo","arguments":{"message":"Hello, world!"}}}
+STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","agent-tool-call-request-id":"test-id","progressToken":1},"name":"echo","arguments":{"message":"Hello, world!"}}}
 STDOUT: {"result":{"content":[{"type":"text","text":"Echo: Hello, world!"}]},"jsonrpc":"2.0","id":2}
-STDIN: {"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":2},"name":"get-sum","arguments":{"a":1,"b":2}}}
+STDIN: {"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","agent-tool-call-request-id":"test-id","progressToken":2},"name":"get-sum","arguments":{"a":1,"b":2}}}
 STDOUT: {"result":{"content":[{"type":"text","text":"The sum of 1 and 2 is 3."}]},"jsonrpc":"2.0","id":3}
-STDIN: {"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":3},"name":"trigger-long-running-operation","arguments":{"duration":1,"steps":5}}}
+STDIN: {"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","agent-tool-call-request-id":"test-id","progressToken":3},"name":"trigger-long-running-operation","arguments":{"duration":1,"steps":5}}}
 STDOUT: {"method":"notifications/progress","params":{"progress":1,"total":5,"progressToken":3},"jsonrpc":"2.0"}
 STDOUT: {"method":"notifications/progress","params":{"progress":2,"total":5,"progressToken":3},"jsonrpc":"2.0"}
 STDOUT: {"method":"notifications/progress","params":{"progress":3,"total":5,"progressToken":3},"jsonrpc":"2.0"}
 STDOUT: {"method":"notifications/progress","params":{"progress":4,"total":5,"progressToken":3},"jsonrpc":"2.0"}
 STDOUT: {"method":"notifications/progress","params":{"progress":5,"total":5,"progressToken":3},"jsonrpc":"2.0"}
 STDOUT: {"result":{"content":[{"type":"text","text":"Long running operation completed. Duration: 1 seconds, Steps: 5."}]},"jsonrpc":"2.0","id":4}
-STDIN: {"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":4},"name":"get-structured-content","arguments":{"location":"New York"}}}
+STDIN: {"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","agent-tool-call-request-id":"test-id","progressToken":4},"name":"get-structured-content","arguments":{"location":"New York"}}}
 STDOUT: {"result":{"content":[{"type":"text","text":"{\"temperature\":33,\"conditions\":\"Cloudy\",\"humidity\":82}"}],"structuredContent":{"temperature":33,"conditions":"Cloudy","humidity":82}},"jsonrpc":"2.0","id":5}
-STDIN: {"jsonrpc":"2.0","id":6,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":5},"name":"trigger-sampling-request","arguments":{"maxTokens":100,"prompt":"Please provide a quote from The Great Gatsby"}}}
+STDIN: {"jsonrpc":"2.0","id":6,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","agent-tool-call-request-id":"test-id","progressToken":5},"name":"trigger-sampling-request","arguments":{"maxTokens":100,"prompt":"Please provide a quote from The Great Gatsby"}}}
 STDOUT: {"method":"sampling/createMessage","params":{"messages":[{"role":"user","content":{"type":"text","text":"Resource trigger-sampling-request context: Please provide a quote from The Great Gatsby"}}],"systemPrompt":"You are a helpful test server.","maxTokens":100,"temperature":0.7},"jsonrpc":"2.0","id":0}
 STDIN: {"jsonrpc":"2.0","id":0,"result":{"model":"mock","stopReason":"endTurn","role":"assistant","content":{"type":"text","text":"\"So we beat on, boats against the current, borne back ceaselessly into the past.\" — F. Scott Fitzgerald, The Great Gatsby (1925)"}}}
 STDOUT: {"result":{"content":[{"type":"text","text":"LLM sampling result: \n{\n  \"model\": \"mock\",\n  \"stopReason\": \"endTurn\",\n  \"role\": \"assistant\",\n  \"content\": {\n    \"type\": \"text\",\n    \"text\": \"\\\"So we beat on, boats against the current, borne back ceaselessly into the past.\\\" — F. Scott Fitzgerald, The Great Gatsby (1925)\"\n  }\n}"}]},"jsonrpc":"2.0","id":6}
@@ -27,5 +27,5 @@ STDOUT: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabi
 STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
 STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":0}}}
 STDOUT: {"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"divide","description":"Divide two numbers","inputSchema":{"type":"object","properties":{"dividend":{"type":"number"},"divisor":{"type":"number"}},"required":["dividend","divisor"]}}]}}
-STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":1},"name":"divide","arguments":{"dividend":10,"divisor":2}}}
+STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","agent-tool-call-request-id":"test-id","progressToken":1},"name":"divide","arguments":{"dividend":10,"divisor":2}}}
 STDOUT: {"jsonrpc":"2.0","id":2,"result":{"content":[{"type":"text","text":"5.0"}],"structuredContent":{"result":5.0},"isError":false}}
@@ -3,5 +3,5 @@ STDOUT: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabi
 STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
 STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":0}}}
 STDOUT: {"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"fetch","description":"Fetch a URL","inputSchema":{"type":"object","properties":{"url":{"type":"string"}},"required":["url"]}}]}}
-STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","progressToken":1},"name":"fetch","arguments":{"url":"https://example.com"}}}
+STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"agent-session-id":"test-session-id","agent-tool-call-request-id":"test-id","progressToken":1},"name":"fetch","arguments":{"url":"https://example.com"}}}
 STDOUT: {"jsonrpc":"2.0","id":2,"result":{"content":[{"type":"text","text":"Contents of https://example.com/:\nThis domain is for use in documentation examples without needing permission. Avoid use in operations.\n\n[Learn more](https://iana.org/domains/example)"}],"isError":false}}