support MCP elicitations (#6947)

No support for request schema yet, but we'll at least show the message
and allow accept/decline.

<img width="823" height="551" alt="Screenshot 2025-11-21 at 2 44 05 PM"
src="https://github.com/user-attachments/assets/6fbb892d-ca12-4765-921e-9ac4b217534d"
/>
This commit is contained in:
Jeremy Rose
2025-11-21 14:44:53 -08:00
committed by GitHub
Unverified
parent 3ea33a0616
commit 7561a6aaf0
16 changed files with 504 additions and 145 deletions
+51
View File
@@ -34,6 +34,7 @@ use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::TaskStartedEvent;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnContextItem;
use codex_rmcp_client::ElicitationResponse;
use futures::future::BoxFuture;
use futures::prelude::*;
use futures::stream::FuturesOrdered;
@@ -44,6 +45,7 @@ use mcp_types::ListResourcesRequestParams;
use mcp_types::ListResourcesResult;
use mcp_types::ReadResourceRequestParams;
use mcp_types::ReadResourceResult;
use mcp_types::RequestId;
use serde_json;
use serde_json::Value;
use tokio::sync::Mutex;
@@ -940,6 +942,19 @@ impl Session {
}
}
pub async fn resolve_elicitation(
&self,
server_name: String,
id: RequestId,
response: ElicitationResponse,
) -> anyhow::Result<()> {
self.services
.mcp_connection_manager
.read()
.await
.resolve_elicitation(server_name, id, response)
}
/// Records input items: always append to conversation history and
/// persist these response items to rollout.
pub(crate) async fn record_conversation_items(
@@ -1413,6 +1428,13 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
)
.await;
}
Op::ResolveElicitation {
server_name,
request_id,
decision,
} => {
handlers::resolve_elicitation(&sess, server_name, request_id, decision).await;
}
Op::Shutdown => {
if handlers::shutdown(&sess, sub.id.clone()).await {
break;
@@ -1452,6 +1474,9 @@ mod handlers {
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::user_input::UserInput;
use codex_rmcp_client::ElicitationAction;
use codex_rmcp_client::ElicitationResponse;
use mcp_types::RequestId;
use std::sync::Arc;
use tracing::info;
use tracing::warn;
@@ -1535,6 +1560,32 @@ mod handlers {
*previous_context = Some(turn_context);
}
pub async fn resolve_elicitation(
sess: &Arc<Session>,
server_name: String,
request_id: RequestId,
decision: codex_protocol::approvals::ElicitationAction,
) {
let action = match decision {
codex_protocol::approvals::ElicitationAction::Accept => ElicitationAction::Accept,
codex_protocol::approvals::ElicitationAction::Decline => ElicitationAction::Decline,
codex_protocol::approvals::ElicitationAction::Cancel => ElicitationAction::Cancel,
};
let response = ElicitationResponse {
action,
content: None,
};
if let Err(err) = sess
.resolve_elicitation(server_name, request_id, response)
.await
{
warn!(
error = %err,
"failed to resolve elicitation request in session"
);
}
}
pub async fn exec_approval(sess: &Arc<Session>, id: String, decision: ReviewDecision) {
match decision {
ReviewDecision::Abort => {
+161 -115
View File
@@ -11,6 +11,7 @@ use std::collections::HashSet;
use std::env;
use std::ffi::OsString;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use crate::mcp::auth::McpAuthStatusEntry;
@@ -20,14 +21,17 @@ use anyhow::anyhow;
use async_channel::Sender;
use codex_async_utils::CancelErr;
use codex_async_utils::OrCancelExt;
use codex_protocol::approvals::ElicitationRequestEvent;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::McpStartupCompleteEvent;
use codex_protocol::protocol::McpStartupFailure;
use codex_protocol::protocol::McpStartupStatus;
use codex_protocol::protocol::McpStartupUpdateEvent;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::OAuthCredentialsStoreMode;
use codex_rmcp_client::RmcpClient;
use codex_rmcp_client::SendElicitation;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::future::Shared;
@@ -39,6 +43,7 @@ use mcp_types::ListResourcesRequestParams;
use mcp_types::ListResourcesResult;
use mcp_types::ReadResourceRequestParams;
use mcp_types::ReadResourceResult;
use mcp_types::RequestId;
use mcp_types::Resource;
use mcp_types::ResourceTemplate;
use mcp_types::Tool;
@@ -46,6 +51,7 @@ use mcp_types::Tool;
use serde_json::json;
use sha1::Digest;
use sha1::Sha1;
use tokio::sync::oneshot;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use tracing::warn;
@@ -110,6 +116,58 @@ pub(crate) struct ToolInfo {
pub(crate) tool: Tool,
}
type ResponderMap = HashMap<(String, RequestId), oneshot::Sender<ElicitationResponse>>;
#[derive(Clone, Default)]
struct ElicitationRequestManager {
requests: Arc<Mutex<ResponderMap>>,
}
impl ElicitationRequestManager {
fn resolve(
&self,
server_name: String,
id: RequestId,
response: ElicitationResponse,
) -> Result<()> {
self.requests
.lock()
.map_err(|e| anyhow!("failed to lock elicitation requests: {e:?}"))?
.remove(&(server_name, id))
.ok_or_else(|| anyhow!("elicitation request not found"))?
.send(response)
.map_err(|e| anyhow!("failed to send elicitation response: {e:?}"))
}
fn make_sender(&self, server_name: String, tx_event: Sender<Event>) -> SendElicitation {
let elicitation_requests = self.requests.clone();
Box::new(move |id, elicitation| {
let elicitation_requests = elicitation_requests.clone();
let tx_event = tx_event.clone();
let server_name = server_name.clone();
async move {
let (tx, rx) = oneshot::channel();
if let Ok(mut lock) = elicitation_requests.lock() {
lock.insert((server_name.clone(), id.clone()), tx);
}
let _ = tx_event
.send(Event {
id: "mcp_elicitation_request".to_string(),
msg: EventMsg::ElicitationRequest(ElicitationRequestEvent {
server_name,
id,
message: elicitation.message,
}),
})
.await;
rx.await
.context("elicitation request channel closed unexpectedly")
}
.boxed()
})
}
}
#[derive(Clone)]
struct ManagedClient {
client: Arc<RmcpClient>,
@@ -129,19 +187,33 @@ impl AsyncManagedClient {
config: McpServerConfig,
store_mode: OAuthCredentialsStoreMode,
cancel_token: CancellationToken,
tx_event: Sender<Event>,
elicitation_requests: ElicitationRequestManager,
) -> Self {
let tool_filter = ToolFilter::from_config(&config);
let fut = start_server_task(
server_name,
config.transport,
store_mode,
config
.startup_timeout_sec
.unwrap_or(DEFAULT_STARTUP_TIMEOUT),
config.tool_timeout_sec.unwrap_or(DEFAULT_TOOL_TIMEOUT),
tool_filter,
cancel_token,
);
let fut = async move {
if let Err(error) = validate_mcp_server_name(&server_name) {
return Err(error.into());
}
let client =
Arc::new(make_rmcp_client(&server_name, config.transport, store_mode).await?);
match start_server_task(
server_name,
client,
config.startup_timeout_sec.or(Some(DEFAULT_STARTUP_TIMEOUT)),
config.tool_timeout_sec.unwrap_or(DEFAULT_TOOL_TIMEOUT),
tool_filter,
tx_event,
elicitation_requests,
)
.or_cancel(&cancel_token)
.await
{
Ok(result) => result,
Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled),
}
};
Self {
client: fut.boxed().shared(),
}
@@ -156,6 +228,7 @@ impl AsyncManagedClient {
#[derive(Default)]
pub(crate) struct McpConnectionManager {
clients: HashMap<String, AsyncManagedClient>,
elicitation_requests: ElicitationRequestManager,
}
impl McpConnectionManager {
@@ -172,6 +245,7 @@ impl McpConnectionManager {
}
let mut clients = HashMap::new();
let mut join_set = JoinSet::new();
let elicitation_requests = ElicitationRequestManager::default();
for (server_name, cfg) in mcp_servers.into_iter().filter(|(_, cfg)| cfg.enabled) {
let cancel_token = cancel_token.child_token();
let _ = emit_update(
@@ -182,8 +256,14 @@ impl McpConnectionManager {
},
)
.await;
let async_managed_client =
AsyncManagedClient::new(server_name.clone(), cfg, store_mode, cancel_token.clone());
let async_managed_client = AsyncManagedClient::new(
server_name.clone(),
cfg,
store_mode,
cancel_token.clone(),
tx_event.clone(),
elicitation_requests.clone(),
);
clients.insert(server_name.clone(), async_managed_client.clone());
let tx_event = tx_event.clone();
let auth_entry = auth_entries.get(&server_name).cloned();
@@ -217,6 +297,7 @@ impl McpConnectionManager {
});
}
self.clients = clients;
self.elicitation_requests = elicitation_requests.clone();
tokio::spawn(async move {
let outcomes = join_set.join_all().await;
let mut summary = McpStartupCompleteEvent::default();
@@ -250,6 +331,15 @@ impl McpConnectionManager {
.context("failed to get client")
}
pub fn resolve_elicitation(
&self,
server_name: String,
id: RequestId,
response: ElicitationResponse,
) -> Result<()> {
self.elicitation_requests.resolve(server_name, id, response)
}
/// Returns a single map that contains all tools. Each key is the
/// fully-qualified name for the tool.
pub async fn list_all_tools(&self) -> HashMap<String, ToolInfo> {
@@ -580,43 +670,12 @@ impl From<anyhow::Error> for StartupOutcomeError {
async fn start_server_task(
server_name: String,
transport: McpServerTransportConfig,
store_mode: OAuthCredentialsStoreMode,
startup_timeout: Duration, // TODO: cancel_token should handle this.
tool_timeout: Duration,
tool_filter: ToolFilter,
cancel_token: CancellationToken,
) -> Result<ManagedClient, StartupOutcomeError> {
if cancel_token.is_cancelled() {
return Err(StartupOutcomeError::Cancelled);
}
if let Err(error) = validate_mcp_server_name(&server_name) {
return Err(error.into());
}
match start_server_work(
server_name,
transport,
store_mode,
startup_timeout,
tool_timeout,
tool_filter,
)
.or_cancel(&cancel_token)
.await
{
Ok(result) => result,
Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled),
}
}
async fn start_server_work(
server_name: String,
transport: McpServerTransportConfig,
store_mode: OAuthCredentialsStoreMode,
startup_timeout: Duration,
client: Arc<RmcpClient>,
startup_timeout: Option<Duration>, // TODO: cancel_token should handle this.
tool_timeout: Duration,
tool_filter: ToolFilter,
tx_event: Sender<Event>,
elicitation_requests: ElicitationRequestManager,
) -> Result<ManagedClient, StartupOutcomeError> {
let params = mcp_types::InitializeRequestParams {
capabilities: ClientCapabilities {
@@ -639,73 +698,16 @@ async fn start_server_work(
protocol_version: mcp_types::MCP_SCHEMA_VERSION.to_owned(),
};
let client_result = match transport {
McpServerTransportConfig::Stdio {
command,
args,
env,
env_vars,
cwd,
} => {
let command_os: OsString = command.into();
let args_os: Vec<OsString> = args.into_iter().map(Into::into).collect();
match RmcpClient::new_stdio_client(command_os, args_os, env, &env_vars, cwd).await {
Ok(client) => {
let client = Arc::new(client);
client
.initialize(params.clone(), Some(startup_timeout))
.await
.map(|_| client)
}
Err(err) => Err(err.into()),
}
}
McpServerTransportConfig::StreamableHttp {
url,
http_headers,
env_http_headers,
bearer_token_env_var,
} => {
let resolved_bearer_token =
match resolve_bearer_token(&server_name, bearer_token_env_var.as_deref()) {
Ok(token) => token,
Err(error) => return Err(error.into()),
};
match RmcpClient::new_streamable_http_client(
&server_name,
&url,
resolved_bearer_token,
http_headers,
env_http_headers,
store_mode,
)
.await
{
Ok(client) => {
let client = Arc::new(client);
client
.initialize(params.clone(), Some(startup_timeout))
.await
.map(|_| client)
}
Err(err) => Err(err),
}
}
};
let send_elicitation = elicitation_requests.make_sender(server_name.clone(), tx_event);
let client = match client_result {
Ok(client) => client,
Err(error) => {
return Err(error.into());
}
};
client
.initialize(params, startup_timeout, send_elicitation)
.await
.map_err(StartupOutcomeError::from)?;
let tools = match list_tools_for_client(&server_name, &client, startup_timeout).await {
Ok(tools) => tools,
Err(error) => {
return Err(error.into());
}
};
let tools = list_tools_for_client(&server_name, &client, startup_timeout)
.await
.map_err(StartupOutcomeError::from)?;
let managed = ManagedClient {
client: Arc::clone(&client),
@@ -717,12 +719,56 @@ async fn start_server_work(
Ok(managed)
}
async fn make_rmcp_client(
server_name: &str,
transport: McpServerTransportConfig,
store_mode: OAuthCredentialsStoreMode,
) -> Result<RmcpClient, StartupOutcomeError> {
match transport {
McpServerTransportConfig::Stdio {
command,
args,
env,
env_vars,
cwd,
} => {
let command_os: OsString = command.into();
let args_os: Vec<OsString> = args.into_iter().map(Into::into).collect();
RmcpClient::new_stdio_client(command_os, args_os, env, &env_vars, cwd)
.await
.map_err(|err| StartupOutcomeError::from(anyhow!(err)))
}
McpServerTransportConfig::StreamableHttp {
url,
http_headers,
env_http_headers,
bearer_token_env_var,
} => {
let resolved_bearer_token =
match resolve_bearer_token(server_name, bearer_token_env_var.as_deref()) {
Ok(token) => token,
Err(error) => return Err(error.into()),
};
RmcpClient::new_streamable_http_client(
server_name,
&url,
resolved_bearer_token,
http_headers,
env_http_headers,
store_mode,
)
.await
.map_err(StartupOutcomeError::from)
}
}
}
async fn list_tools_for_client(
server_name: &str,
client: &Arc<RmcpClient>,
timeout: Duration,
timeout: Option<Duration>,
) -> Result<Vec<ToolInfo>> {
let resp = client.list_tools(None, Some(timeout)).await?;
let resp = client.list_tools(None, timeout).await?;
Ok(resp
.tools
.into_iter()
+1
View File
@@ -64,6 +64,7 @@ pub(crate) fn should_persist_event_msg(ev: &EventMsg) -> bool {
| EventMsg::ExecCommandOutputDelta(_)
| EventMsg::ExecCommandEnd(_)
| EventMsg::ExecApprovalRequest(_)
| EventMsg::ElicitationRequest(_)
| EventMsg::ApplyPatchApprovalRequest(_)
| EventMsg::BackgroundEvent(_)
| EventMsg::StreamError(_)