Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 0 additions & 16 deletions codex-rs/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 0 additions & 1 deletion codex-rs/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,6 @@ members = [
"exec-server",
"execpolicy",
"execpolicy-legacy",
"ext/connectors",
"ext/extension-api",
"ext/goal",
"ext/guardian",
Expand Down
95 changes: 95 additions & 0 deletions codex-rs/codex-mcp/src/connection_manager_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,101 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel
);
}

#[tokio::test]
async fn elicitation_route_ids_are_unique_across_managers() {
let first_manager = ElicitationRequestManager::new(
AskForApproval::OnRequest,
PermissionProfile::default(),
/*reviewer*/ None,
);
let second_manager = ElicitationRequestManager::new(
AskForApproval::OnRequest,
PermissionProfile::default(),
/*reviewer*/ None,
);
let (first_tx, first_rx) = async_channel::bounded(1);
let (second_tx, second_rx) = async_channel::bounded(1);
let first_sender = first_manager.make_sender("server".to_string(), first_tx);
let second_sender = second_manager.make_sender("server".to_string(), second_tx);

let first_task = tokio::spawn(first_sender(
NumberOrString::Number(1),
codex_rmcp_client::Elicitation::OpenAiForm {
meta: None,
message: "First?".to_string(),
requested_schema: serde_json::json!({}),
},
));
let second_task = tokio::spawn(second_sender(
NumberOrString::Number(1),
codex_rmcp_client::Elicitation::OpenAiForm {
meta: None,
message: "Second?".to_string(),
requested_schema: serde_json::json!({}),
},
));

let EventMsg::ElicitationRequest(first_request) = first_rx
.recv()
.await
.expect("first elicitation should be emitted")
.msg
else {
panic!("expected first elicitation request");
};
let EventMsg::ElicitationRequest(second_request) = second_rx
.recv()
.await
.expect("second elicitation should be emitted")
.msg
else {
panic!("expected second elicitation request");
};
assert_ne!(first_request.id, second_request.id);

let first_response = ElicitationResponse {
action: ElicitationAction::Accept,
content: Some(serde_json::json!({"manager": "first"})),
meta: None,
};
let second_response = ElicitationResponse {
action: ElicitationAction::Decline,
content: Some(serde_json::json!({"manager": "second"})),
meta: None,
};
let first_id = match first_request.id {
codex_protocol::mcp::RequestId::String(value) => NumberOrString::String(Arc::from(value)),
codex_protocol::mcp::RequestId::Integer(value) => NumberOrString::Number(value),
};
let second_id = match second_request.id {
codex_protocol::mcp::RequestId::String(value) => NumberOrString::String(Arc::from(value)),
codex_protocol::mcp::RequestId::Integer(value) => NumberOrString::Number(value),
};
first_manager
.resolve("server".to_string(), first_id, first_response.clone())
.await
.expect("first manager should resolve its routed request");
second_manager
.resolve("server".to_string(), second_id, second_response.clone())
.await
.expect("second manager should resolve its routed request");

assert_eq!(
first_task
.await
.expect("first elicitation task should join")
.expect("first elicitation should resolve"),
first_response
);
assert_eq!(
second_task
.await
.expect("second elicitation task should join")
.expect("second elicitation should resolve"),
second_response
);
}

#[test]
fn test_normalize_tools_short_non_duplicated_names() {
let tools = vec![
Expand Down
12 changes: 10 additions & 2 deletions codex-rs/codex-mcp/src/elicitation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;

use crate::mcp::McpPermissionPromptAutoApproveContext;
use crate::mcp::mcp_permission_prompt_is_auto_approved;
Expand Down Expand Up @@ -48,13 +50,16 @@ pub trait ElicitationReviewer: Send + Sync {

pub type ElicitationReviewerHandle = Arc<dyn ElicitationReviewer>;

static NEXT_ELICITATION_ROUTE_ID: AtomicU64 = AtomicU64::new(1);

#[derive(Clone)]
pub(crate) struct ElicitationRequestManager {
requests: Arc<Mutex<ResponderMap>>,
pub(crate) approval_policy: Arc<StdMutex<AskForApproval>>,
pub(crate) permission_profile: Arc<StdMutex<PermissionProfile>>,
auto_deny: Arc<StdMutex<bool>>,
reviewer: Option<ElicitationReviewerHandle>,
route_id: u64,
}

impl ElicitationRequestManager {
Expand All @@ -69,6 +74,7 @@ impl ElicitationRequestManager {
permission_profile: Arc::new(StdMutex::new(permission_profile)),
auto_deny: Arc::new(StdMutex::new(false)),
reviewer,
route_id: NEXT_ELICITATION_ROUTE_ID.fetch_add(1, Ordering::Relaxed),
}
}

Expand Down Expand Up @@ -110,6 +116,7 @@ impl ElicitationRequestManager {
let permission_profile = self.permission_profile.clone();
let auto_deny = self.auto_deny.clone();
let reviewer = self.reviewer.clone();
let route_id = self.route_id;
Box::new(move |id, elicitation| {
let elicitation_requests = elicitation_requests.clone();
let tx_event = tx_event.clone();
Expand Down Expand Up @@ -214,17 +221,18 @@ impl ElicitationRequestManager {
},
};
let (tx, rx) = oneshot::channel();
let response_id = RequestId::String(Arc::from(format!("{route_id}:{id}")));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve request ID variant in routed keys

If an MCP server has overlapping elicitations whose JSON-RPC IDs are 1 and "1", this stringified routed ID makes both entries use the same key (<route>:1), so the second insert overwrites the first responder and one request can never be resolved. The previous typed RequestId key distinguished numeric and string IDs, so the route wrapper should encode the variant as well as the value.

Useful? React with 👍 / 👎.

{
let mut lock = elicitation_requests.lock().await;
lock.insert((server_name.clone(), id.clone()), tx);
lock.insert((server_name.clone(), response_id.clone()), tx);
Comment on lines +224 to +227

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Route responses to the originating manager

When a refreshed/candidate manager emits an elicitation before refresh_mcp_servers swaps it in, or a superseded manager still has an unanswered prompt, the client returns this routed ID but Session::resolve_elicitation still calls mcp_connection_manager.load_full().resolve_elicitation(...) on only the current manager. Because the routed ID is stored only in the originating manager's local requests map here, the prefix does not actually route across managers and the response is rejected as elicitation request not found; please keep a session-level registry or dispatch by the route prefix to the owning manager.

Useful? React with 👍 / 👎.

}
let _ = tx_event
.send(Event {
id: "mcp_elicitation_request".to_string(),
msg: EventMsg::ElicitationRequest(ElicitationRequestEvent {
turn_id: None,
server_name,
id: match id.clone() {
id: match response_id {
rmcp::model::NumberOrString::String(value) => {
ProtocolRequestId::String(value.to_string())
}
Expand Down
Loading
Loading