mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
committed by
GitHub
Unverified
parent
3ea33a0616
commit
7561a6aaf0
@@ -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 => {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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(_)
|
||||
|
||||
Reference in New Issue
Block a user