openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-08-23 13:09:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enable skills in the MCP server (#36339)

Browse Source 
## What changed

- Install the skills extension for MCP Codex tool sessions so host skills are
  included in the developer instructions.
- Route extension warnings to the matching active MCP turn as `codex/event`
  notifications, preserving request and thread metadata and ordering warnings
  before the final tool response.
- Track reply turns independently and bound forwarded warning messages to 256
  UTF-8 bytes.

## Testing

- Add unit coverage for active-turn routing, request ID collisions, warning
  ordering, filtering, and truncation.
- Add an MCP integration test for host skill instructions and skills context
  budget warnings.

GitOrigin-RevId: 9e5c699665f6889bf1103e78c3d6ca53824a10a0
This commit is contained in:
felixxia-oai
2026-07-31 14:30:28 +00:00
committed by copyberry
parent 448118f544
commit 5548c95d66
14 changed files with 734 additions and 205 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
codex-rs/Cargo.lock generated
View File

@@ -3541,12 +3541,15 @@ dependencies = [
"codex-core",
"codex-exec-server",
"codex-extension-api",
"codex-features",
"codex-git-attribution",
"codex-home",
"codex-image-generation-extension",
"codex-login",
"codex-otel",
"codex-protocol",
"codex-shell-command",
"codex-skills-extension",
"codex-utils-absolute-path",
"codex-utils-cli",
"codex-utils-json-to-toml",
@@ -3563,6 +3566,7 @@ dependencies = [
"tokio",
"tracing",
"tracing-subscriber",
"uuid",
"wiremock",
]

4
codex-rs/mcp-server/Cargo.toml
View File

@@ -25,9 +25,12 @@ codex-home = { workspace = true }
codex-image-generation-extension = { workspace = true }
codex-exec-server = { workspace = true }
codex-extension-api = { workspace = true }
codex-features = { workspace = true }
codex-git-attribution = { workspace = true }
codex-login = { workspace = true }
codex-otel = { workspace = true }
codex-protocol = { workspace = true }
codex-skills-extension = { workspace = true }
codex-utils-cli = { workspace = true }
codex-utils-json-to-toml = { workspace = true }
rmcp = { workspace = true }
@@ -44,6 +47,7 @@ tokio = { workspace = true, features = [
] }
tracing = { workspace = true, features = ["log"] }
tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] }
uuid = { workspace = true, features = ["v7"] }
[dev-dependencies]
app_test_support = { workspace = true }

97
codex-rs/mcp-server/src/active_turn_registry.rs Normal file
View File

@@ -0,0 +1,97 @@
use std::collections::HashMap;
use std::sync::Mutex;
use std::sync::MutexGuard;
use std::sync::PoisonError;
use codex_protocol::ThreadId;
use rmcp::model::RequestId;
/// Tracks active MCP requests by both request ID and Codex turn.
///
/// Warning and response enqueue callbacks run while holding the same lock. This
/// makes route removal and response enqueue atomic with respect to warning
/// enqueue, so a warning is either queued before its tool response or dropped.
#[derive(Default)]
pub(crate) struct ActiveTurnRegistry {
state: Mutex<ActiveTurnState>,
}
#[derive(Default)]
struct ActiveTurnState {
by_request: HashMap<RequestId, ActiveTurn>,
by_thread_and_turn: HashMap<(ThreadId, String), RequestId>,
}
struct ActiveTurn {
thread_id: ThreadId,
turn_id: String,
}
impl ActiveTurnRegistry {
pub(crate) fn register(&self, request_id: RequestId, thread_id: ThreadId, turn_id: String) {
let active_turn = ActiveTurn {
thread_id,
turn_id: turn_id.clone(),
};
let mut state = self.lock();
if let Some(previous_turn) = state.by_request.insert(request_id.clone(), active_turn) {
state
.by_thread_and_turn
.remove(&(previous_turn.thread_id, previous_turn.turn_id));
}
state
.by_thread_and_turn
.insert((thread_id, turn_id), request_id);
}
pub(crate) fn thread_id(&self, request_id: &RequestId) -> Option<ThreadId> {
self.lock()
.by_request
.get(request_id)
.map(|active_turn| active_turn.thread_id)
}
pub(crate) fn with_request_id(
&self,
thread_id: ThreadId,
turn_id: &str,
f: impl FnOnce(&RequestId),
) -> bool {
let state = self.lock();
let Some(request_id) = state
.by_thread_and_turn
.get(&(thread_id, turn_id.to_string()))
else {
return false;
};
f(request_id);
true
}
pub(crate) fn unregister(&self, request_id: &RequestId) {
let mut state = self.lock();
Self::remove(&mut state, request_id);
}
pub(crate) fn finish(&self, request_id: &RequestId, finish: impl FnOnce()) {
let mut state = self.lock();
Self::remove(&mut state, request_id);
finish();
}
fn remove(state: &mut ActiveTurnState, request_id: &RequestId) {
if let Some(active_turn) = state.by_request.remove(request_id) {
state
.by_thread_and_turn
.remove(&(active_turn.thread_id, active_turn.turn_id));
}
}
fn lock(&self) -> MutexGuard<'_, ActiveTurnState> {
self.state.lock().unwrap_or_else(PoisonError::into_inner)
}
}
#[cfg(test)]
#[path = "active_turn_registry_tests.rs"]
mod tests;

70
codex-rs/mcp-server/src/active_turn_registry_tests.rs Normal file
View File

@@ -0,0 +1,70 @@
use codex_protocol::ThreadId;
use pretty_assertions::assert_eq;
use rmcp::model::RequestId;
use super::ActiveTurnRegistry;
#[test]
fn colliding_request_id_strings_use_independent_turn_routes() {
let active_turns = ActiveTurnRegistry::default();
let thread_id = ThreadId::new();
let numeric_request_id = RequestId::Number(7);
let string_request_id = RequestId::String("7".into());
let numeric_turn_id = "numeric-turn".to_string();
let string_turn_id = "string-turn".to_string();
active_turns.register(
numeric_request_id.clone(),
thread_id,
numeric_turn_id.clone(),
);
active_turns.register(string_request_id.clone(), thread_id, string_turn_id.clone());
let mut routed_request_ids = Vec::new();
assert!(
active_turns.with_request_id(thread_id, &numeric_turn_id, |request_id| {
routed_request_ids.push(request_id.clone());
})
);
assert!(
active_turns.with_request_id(thread_id, &string_turn_id, |request_id| {
routed_request_ids.push(request_id.clone());
})
);
assert_eq!(
routed_request_ids,
vec![numeric_request_id.clone(), string_request_id.clone()]
);
active_turns.unregister(&numeric_request_id);
assert!(!active_turns.with_request_id(thread_id, &numeric_turn_id, |_| {}));
let mut remaining_request_id = None;
assert!(
active_turns.with_request_id(thread_id, &string_turn_id, |request_id| {
remaining_request_id = Some(request_id.clone());
})
);
assert_eq!(remaining_request_id, Some(string_request_id));
}
#[test]
fn reused_request_id_replaces_the_previous_turn_route() {
let active_turns = ActiveTurnRegistry::default();
let thread_id = ThreadId::new();
let request_id = RequestId::Number(7);
let first_turn_id = "first-turn".to_string();
let second_turn_id = "second-turn".to_string();
active_turns.register(request_id.clone(), thread_id, first_turn_id.clone());
active_turns.register(request_id.clone(), thread_id, second_turn_id.clone());
assert!(!active_turns.with_request_id(thread_id, &first_turn_id, |_| {}));
let mut routed_request_id = None;
assert!(
active_turns.with_request_id(thread_id, &second_turn_id, |request_id| {
routed_request_id = Some(request_id.clone());
})
);
assert_eq!(routed_request_id, Some(request_id));
}

140
codex-rs/mcp-server/src/codex_tool_runner.rs
View File

@@ -2,9 +2,9 @@
//! Tokio task. Separated from `message_processor.rs` to keep that file small
//! and to make future feature-growth easier to manage.
use std::collections::HashMap;
use std::sync::Arc;
use crate::active_turn_registry::ActiveTurnRegistry;
use crate::exec_approval::handle_exec_approval_request;
use crate::outgoing_message::OutgoingMessageSender;
use crate::outgoing_message::OutgoingNotificationMeta;
@@ -28,7 +28,7 @@ use rmcp::model::CallToolResult;
use rmcp::model::ContentBlock;
use rmcp::model::RequestId;
use serde_json::json;
use tokio::sync::Mutex;
use uuid::Uuid;
/// To adhere to MCP `tools/call` response format, include the Codex
/// `threadId` in the `structured_content` field of the response.
@@ -61,7 +61,7 @@ pub async fn run_codex_tool_session(
config: CodexConfig,
outgoing: Arc<OutgoingMessageSender>,
thread_manager: Arc<ThreadManager>,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ThreadId>>>,
active_turns: Arc<ActiveTurnRegistry>,
) {
let NewThread {
thread_id,
@@ -76,7 +76,7 @@ pub async fn run_codex_tool_session(
let result = CallToolResult::error(vec![ContentBlock::text(format!(
"Failed to start Codex session: {e}"
))]);
outgoing.send_response(id.clone(), result).await;
outgoing.send_response(id.clone(), result);
return;
}
};
@@ -86,26 +86,20 @@ pub async fn run_codex_tool_session(
id: "".to_string(),
msg: EventMsg::SessionConfigured(session_configured.clone()),
};
outgoing
.send_event_as_notification(
&session_configured_event,
Some(OutgoingNotificationMeta {
request_id: Some(id.clone()),
thread_id: Some(thread_id),
}),
)
.await;
outgoing.send_event_as_notification(
&session_configured_event,
Some(OutgoingNotificationMeta {
request_id: Some(id.clone()),
thread_id: Some(thread_id),
}),
);
// Use the original MCP request ID as the `sub_id` for the Codex submission so that
// any events emitted for this tool-call can be correlated with the
// originating `tools/call` request.
let sub_id = id.to_string();
running_requests_id_to_codex_uuid
.lock()
.await
.insert(id.clone(), thread_id);
// Preserve the legacy event ID for initial `codex` calls. Each call starts
// a new thread, so the thread and turn pair remains unique.
let turn_id = id.to_string();
active_turns.register(id.clone(), thread_id, turn_id.clone());
let submission = Submission {
id: sub_id.clone(),
id: turn_id,
op: Op::UserInput {
items: vec![UserInput::Text {
text: initial_prompt.clone(),
@@ -129,20 +123,11 @@ pub async fn run_codex_tool_session(
format!("Failed to submit initial prompt: {e}"),
Some(true),
);
outgoing.send_response(id.clone(), result).await;
// unregister the id so we don't keep it in the map
running_requests_id_to_codex_uuid.lock().await.remove(&id);
active_turns.finish(&id, || outgoing.send_response(id.clone(), result));
return;
}
run_codex_tool_session_inner(
thread_id,
thread,
outgoing,
id,
running_requests_id_to_codex_uuid,
)
.await;
run_codex_tool_session_inner(thread_id, thread, outgoing, id, active_turns).await;
}
pub async fn run_codex_tool_session_reply(
@@ -151,23 +136,29 @@ pub async fn run_codex_tool_session_reply(
outgoing: Arc<OutgoingMessageSender>,
request_id: RequestId,
prompt: String,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ThreadId>>>,
active_turns: Arc<ActiveTurnRegistry>,
) {
running_requests_id_to_codex_uuid
.lock()
.await
.insert(request_id.clone(), thread_id);
// Replies share a thread, so use Core's UUIDv7 submission ID convention
// instead of a reusable MCP request ID.
let turn_id = Uuid::now_v7().to_string();
active_turns.register(request_id.clone(), thread_id, turn_id.clone());
if let Err(e) = thread
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: prompt,
// MCP tool prompts are plain text with no UI element ranges.
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
.submit_with_id(Submission {
id: turn_id,
op: Op::UserInput {
items: vec![UserInput::Text {
text: prompt,
// MCP tool prompts are plain text with no UI element ranges.
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
},
client_user_message_id: None,
trace: None,
parent_turn_id: None,
})
.await
{
@@ -177,23 +168,13 @@ pub async fn run_codex_tool_session_reply(
format!("Failed to submit user input: {e}"),
Some(true),
);
outgoing.send_response(request_id.clone(), result).await;
// unregister the id so we don't keep it in the map
running_requests_id_to_codex_uuid
.lock()
.await
.remove(&request_id);
active_turns.finish(&request_id, || {
outgoing.send_response(request_id.clone(), result);
});
return;
}
run_codex_tool_session_inner(
thread_id,
thread,
outgoing,
request_id,
running_requests_id_to_codex_uuid,
)
.await;
run_codex_tool_session_inner(thread_id, thread, outgoing, request_id, active_turns).await;
}
async fn run_codex_tool_session_inner(
@@ -201,7 +182,7 @@ async fn run_codex_tool_session_inner(
thread: Arc<CodexThread>,
outgoing: Arc<OutgoingMessageSender>,
request_id: RequestId,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ThreadId>>>,
active_turns: Arc<ActiveTurnRegistry>,
) {
let request_id_str = request_id.to_string();
@@ -210,15 +191,13 @@ async fn run_codex_tool_session_inner(
loop {
match thread.next_event().await {
Ok(event) => {
outgoing
.send_event_as_notification(
&event,
Some(OutgoingNotificationMeta {
request_id: Some(request_id.clone()),
thread_id: Some(thread_id),
}),
)
.await;
outgoing.send_event_as_notification(
&event,
Some(OutgoingNotificationMeta {
request_id: Some(request_id.clone()),
thread_id: Some(thread_id),
}),
);
match event.msg {
EventMsg::ExecApprovalRequest(ev) => {
@@ -267,7 +246,9 @@ async fn run_codex_tool_session_inner(
err_event.message,
Some(true),
);
outgoing.send_response(request_id.clone(), result).await;
active_turns.finish(&request_id, || {
outgoing.send_response(request_id.clone(), result);
});
break;
}
EventMsg::Warning(_)
@@ -317,12 +298,9 @@ async fn run_codex_tool_session_inner(
let result = create_call_tool_result_with_thread_id(
thread_id, text, /*is_error*/ None,
);
outgoing.send_response(request_id.clone(), result).await;
// unregister the id so we don't keep it in the map
running_requests_id_to_codex_uuid
.lock()
.await
.remove(&request_id);
active_turns.finish(&request_id, || {
outgoing.send_response(request_id.clone(), result);
});
break;
}
EventMsg::SessionConfigured(_) => {
@@ -415,7 +393,9 @@ async fn run_codex_tool_session_inner(
format!("Codex runtime error: {e}"),
Some(true),
);
outgoing.send_response(request_id.clone(), result).await;
active_turns.finish(&request_id, || {
outgoing.send_response(request_id.clone(), result);
});
break;
}
}

4
codex-rs/mcp-server/src/exec_approval.rs
View File

@@ -86,9 +86,7 @@ pub(crate) async fn handle_exec_approval_request(
let message = format!("Failed to serialize ExecApprovalElicitRequestParams: {err}");
error!("{message}");
outgoing
.send_error(request_id.clone(), ErrorData::invalid_params(message, None))
.await;
outgoing.send_error(request_id.clone(), ErrorData::invalid_params(message, None));
return;
}

89
codex-rs/mcp-server/src/extension_event_sink.rs Normal file
View File

@@ -0,0 +1,89 @@
use std::sync::Arc;
use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionWarning;
use codex_protocol::ThreadId;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::WarningEvent;
use crate::active_turn_registry::ActiveTurnRegistry;
use crate::outgoing_message::OutgoingMessageSender;
use crate::outgoing_message::OutgoingNotificationMeta;
const MAX_EXTENSION_WARNING_BYTES: usize = 256;
pub(crate) fn extension_event_sink(
outgoing: Arc<OutgoingMessageSender>,
active_turns: Arc<ActiveTurnRegistry>,
) -> Arc<dyn ExtensionEventSink> {
Arc::new(McpExtensionEventSink {
outgoing,
active_turns,
})
}
struct McpExtensionEventSink {
outgoing: Arc<OutgoingMessageSender>,
active_turns: Arc<ActiveTurnRegistry>,
}
impl ExtensionEventSink for McpExtensionEventSink {
fn emit(&self, event: Event) {
tracing::debug!(
event_id = %event.id,
msg = ?event.msg,
"dropping unsupported extension event"
);
}
fn emit_warning(&self, warning: ExtensionWarning) {
let ExtensionWarning {
thread_id,
turn_id,
message,
} = warning;
let Ok(thread_id) = ThreadId::from_string(&thread_id) else {
tracing::warn!(%thread_id, "dropping extension warning with invalid thread id");
return;
};
let Some(turn_id) = turn_id else {
tracing::debug!(%thread_id, "dropping extension warning without a turn id");
return;
};
let mut message = message;
if message.len() > MAX_EXTENSION_WARNING_BYTES {
let mut truncate_at = MAX_EXTENSION_WARNING_BYTES;
while !message.is_char_boundary(truncate_at) {
truncate_at -= 1;
}
message.truncate(truncate_at);
}
let found_active_turn =
self.active_turns
.with_request_id(thread_id, &turn_id, |request_id| {
let event = Event {
id: turn_id.clone(),
msg: EventMsg::Warning(WarningEvent { message }),
};
self.outgoing.send_event_as_notification(
&event,
Some(OutgoingNotificationMeta {
request_id: Some(request_id.clone()),
thread_id: Some(thread_id),
}),
);
});
if !found_active_turn {
tracing::debug!(
%thread_id,
%turn_id,
"dropping extension warning without a matching active turn"
);
}
}
}
#[cfg(test)]
#[path = "extension_event_sink_tests.rs"]
mod tests;

190
codex-rs/mcp-server/src/extension_event_sink_tests.rs Normal file
View File

@@ -0,0 +1,190 @@
use std::sync::Arc;
use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionWarning;
use codex_protocol::ThreadId;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::WarningEvent;
use pretty_assertions::assert_eq;
use rmcp::model::RequestId;
use serde_json::json;
use tokio::sync::mpsc;
use super::extension_event_sink;
use crate::active_turn_registry::ActiveTurnRegistry;
use crate::outgoing_message::OutgoingMessage;
use crate::outgoing_message::OutgoingMessageSender;
use crate::outgoing_message::OutgoingNotification;
fn test_sink() -> (
Arc<dyn ExtensionEventSink>,
Arc<OutgoingMessageSender>,
Arc<ActiveTurnRegistry>,
mpsc::UnboundedReceiver<OutgoingMessage>,
) {
let (outgoing_tx, outgoing_rx) = mpsc::unbounded_channel();
let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx));
let active_turns = Arc::new(ActiveTurnRegistry::default());
(
extension_event_sink(Arc::clone(&outgoing), Arc::clone(&active_turns)),
outgoing,
active_turns,
outgoing_rx,
)
}
#[tokio::test(flavor = "current_thread")]
async fn warning_is_enqueued_before_following_response() {
let (sink, outgoing, active_turns, mut outgoing_rx) = test_sink();
let request_id = RequestId::Number(7);
let thread_id = ThreadId::new();
let turn_id = "turn-7".to_string();
active_turns.register(request_id.clone(), thread_id, turn_id.clone());
sink.emit_warning(ExtensionWarning {
thread_id: thread_id.to_string(),
turn_id: Some(turn_id),
message: "warning".to_string(),
});
active_turns.finish(&request_id, || {
outgoing.send_response(request_id.clone(), json!({}));
});
let Some(OutgoingMessage::Notification(notification)) = outgoing_rx.recv().await else {
panic!("warning notification should be enqueued before the response");
};
assert_eq!(notification.method, "codex/event");
assert_eq!(
notification.params.as_ref().map(|params| &params["msg"]),
Some(&json!({
"message": "warning",
"type": "warning",
}))
);
}
#[tokio::test(flavor = "current_thread")]
async fn colliding_request_id_strings_route_warnings_independently() {
let (sink, _outgoing, active_turns, mut outgoing_rx) = test_sink();
let thread_id = ThreadId::new();
let numeric_turn_id = "019c0000-0000-7000-8000-000000000001".to_string();
let string_turn_id = "019c0000-0000-7000-8000-000000000002".to_string();
active_turns.register(RequestId::Number(7), thread_id, numeric_turn_id.clone());
active_turns.register(
RequestId::String("7".into()),
thread_id,
string_turn_id.clone(),
);
sink.emit_warning(ExtensionWarning {
thread_id: thread_id.to_string(),
turn_id: Some(numeric_turn_id.clone()),
message: "numeric warning".to_string(),
});
sink.emit_warning(ExtensionWarning {
thread_id: thread_id.to_string(),
turn_id: Some(string_turn_id.clone()),
message: "string warning".to_string(),
});
let Some(OutgoingMessage::Notification(numeric_warning)) = outgoing_rx.recv().await else {
panic!("expected warning for numeric request id");
};
assert_eq!(
numeric_warning,
OutgoingNotification {
method: "codex/event".to_string(),
params: Some(json!({
"_meta": {
"requestId": 7,
"threadId": thread_id,
},
"id": numeric_turn_id,
"msg": {
"message": "numeric warning",
"type": "warning",
},
})),
}
);
let Some(OutgoingMessage::Notification(string_warning)) = outgoing_rx.recv().await else {
panic!("expected warning for string request id");
};
assert_eq!(
string_warning,
OutgoingNotification {
method: "codex/event".to_string(),
params: Some(json!({
"_meta": {
"requestId": "7",
"threadId": thread_id,
},
"id": string_turn_id,
"msg": {
"message": "string warning",
"type": "warning",
},
})),
}
);
}
#[tokio::test(flavor = "current_thread")]
async fn warning_requires_a_matching_active_turn() {
let (sink, _outgoing, active_turns, mut outgoing_rx) = test_sink();
let request_id = RequestId::Number(7);
let thread_id = ThreadId::new();
active_turns.register(request_id, thread_id, "turn-7".to_string());
sink.emit_warning(ExtensionWarning {
thread_id: thread_id.to_string(),
turn_id: Some("different-turn".to_string()),
message: "warning".to_string(),
});
tokio::task::yield_now().await;
assert!(outgoing_rx.try_recv().is_err());
}
#[tokio::test(flavor = "current_thread")]
async fn warning_is_truncated_to_256_utf8_bytes() {
let (sink, _outgoing, active_turns, mut outgoing_rx) = test_sink();
let request_id = RequestId::Number(7);
let thread_id = ThreadId::new();
let turn_id = "turn-7".to_string();
active_turns.register(request_id.clone(), thread_id, turn_id.clone());
sink.emit_warning(ExtensionWarning {
thread_id: thread_id.to_string(),
turn_id: Some(turn_id),
message: format!("{}é", "a".repeat(255)),
});
let Some(OutgoingMessage::Notification(notification)) = outgoing_rx.recv().await else {
panic!("expected warning notification");
};
assert_eq!(
notification.params.as_ref().map(|params| &params["msg"]),
Some(&json!({
"message": "a".repeat(255),
"type": "warning",
}))
);
}
#[tokio::test(flavor = "current_thread")]
async fn generic_extension_events_are_dropped() {
let (sink, _outgoing, _running_requests, mut outgoing_rx) = test_sink();
sink.emit(Event {
id: "turn".to_string(),
msg: EventMsg::Warning(WarningEvent {
message: "warning".to_string(),
}),
});
tokio::task::yield_now().await;
assert!(outgoing_rx.try_recv().is_err());
}

2
codex-rs/mcp-server/src/lib.rs
View File

@@ -29,9 +29,11 @@ use tracing::info;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::prelude::*;
mod active_turn_registry;
mod codex_tool_config;
mod codex_tool_runner;
mod exec_approval;
mod extension_event_sink;
pub(crate) mod message_processor;
mod outgoing_message;
mod patch_approval;

165
codex-rs/mcp-server/src/message_processor.rs
View File

@@ -1,4 +1,3 @@
use std::collections::HashMap;
use std::sync::Arc;
use codex_arg0::Arg0DispatchPaths;
@@ -11,7 +10,6 @@ use codex_home::CodexHomeUserInstructionsProvider;
use codex_login::AuthManager;
use codex_login::default_client::USER_AGENT_SUFFIX;
use codex_login::default_client::get_codex_user_agent;
use codex_protocol::ThreadId;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::Submission;
use rmcp::model::CallToolRequestParams;
@@ -29,13 +27,14 @@ use rmcp::model::JsonRpcResponse;
use rmcp::model::RequestId;
use rmcp::model::ServerCapabilities;
use serde_json::json;
use tokio::sync::Mutex;
use tokio::task;
use crate::active_turn_registry::ActiveTurnRegistry;
use crate::codex_tool_config::CodexToolCallParam;
use crate::codex_tool_config::CodexToolCallReplyParam;
use crate::codex_tool_config::create_tool_for_codex_tool_call_param;
use crate::codex_tool_config::create_tool_for_codex_tool_call_reply_param;
use crate::extension_event_sink::extension_event_sink;
use crate::outgoing_message::OutgoingMessageSender;
pub(crate) struct MessageProcessor {
@@ -43,7 +42,7 @@ pub(crate) struct MessageProcessor {
initialized: bool,
arg0_paths: Arg0DispatchPaths,
thread_manager: Arc<ThreadManager>,
running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, ThreadId>>>,
active_turns: Arc<ActiveTurnRegistry>,
}
impl MessageProcessor {
@@ -66,7 +65,10 @@ impl MessageProcessor {
let user_instructions_provider = Arc::new(CodexHomeUserInstructionsProvider::new(
config.codex_home.clone(),
));
let mut extensions = ExtensionRegistryBuilder::<Config>::new();
let active_turns = Arc::new(ActiveTurnRegistry::default());
let mut extensions = ExtensionRegistryBuilder::<Config>::with_event_sink(
extension_event_sink(Arc::clone(&outgoing), Arc::clone(&active_turns)),
);
codex_git_attribution::install(
&mut extensions,
auth_manager.clone(),
@@ -78,6 +80,21 @@ impl MessageProcessor {
auth_manager.clone(),
|config: &Config| Some(config.codex_home.clone()),
);
let skill_providers = codex_skills_extension::SkillProviders::new()
.with_host_provider(Arc::new(codex_skills_extension::HostSkillProvider::new()));
codex_skills_extension::install_with_providers_and_metrics(
&mut extensions,
skill_providers,
codex_otel::global(),
|config: &Config| codex_skills_extension::SkillsExtensionConfig {
include_instructions: config.include_skill_instructions,
bundled_skills_enabled: config.bundled_skills_enabled(),
orchestrator_skills_enabled: config.orchestrator_skills_enabled,
shadow_selection_enabled: config
.features
.enabled(codex_features::Feature::SkillSearch),
},
);
let thread_manager = Arc::new(ThreadManager::new(
config.as_ref(),
Arc::clone(&auth_manager),
@@ -99,7 +116,7 @@ impl MessageProcessor {
initialized: false,
arg0_paths,
thread_manager,
running_requests_id_to_codex_uuid: Arc::new(Mutex::new(HashMap::new())),
active_turns,
}
}
@@ -156,17 +173,15 @@ impl MessageProcessor {
.await;
}
ClientRequest::CustomRequest(custom) => {
let method = custom.method.clone();
self.outgoing
.send_error(
request_id,
ErrorData::new(
ErrorCode::METHOD_NOT_FOUND,
format!("method not found: {method}"),
Some(json!({ "method": method })),
),
)
.await;
let method = custom.method;
self.outgoing.send_error(
request_id,
ErrorData::new(
ErrorCode::METHOD_NOT_FOUND,
format!("method not found: {method}"),
Some(json!({ "method": method })),
),
);
}
request => {
self.handle_unsupported_request(request_id, request.method())
@@ -219,12 +234,10 @@ impl MessageProcessor {
tracing::info!("initialize -> params: {:?}", params);
if self.initialized {
self.outgoing
.send_error(
id,
ErrorData::invalid_request("initialize called more than once", None),
)
.await;
self.outgoing.send_error(
id,
ErrorData::invalid_request("initialize called more than once", None),
);
return;
}
@@ -243,15 +256,13 @@ impl MessageProcessor {
let mut server_info_value = match serde_json::to_value(&server_info) {
Ok(value) => value,
Err(err) => {
self.outgoing
.send_error(
id,
ErrorData::internal_error(
format!("failed to serialize server info: {err}"),
None,
),
)
.await;
self.outgoing.send_error(
id,
ErrorData::internal_error(
format!("failed to serialize server info: {err}"),
None,
),
);
return;
}
};
@@ -264,20 +275,18 @@ impl MessageProcessor {
.enable_tool_list_changed()
.build();
let result = InitializeResult::new(capabilities)
.with_protocol_version(params.protocol_version.clone())
.with_protocol_version(params.protocol_version)
.with_server_info(server_info);
let mut result_value = match serde_json::to_value(result) {
Ok(value) => value,
Err(err) => {
self.outgoing
.send_error(
id,
ErrorData::internal_error(
format!("failed to serialize initialize response: {err}"),
None,
),
)
.await;
self.outgoing.send_error(
id,
ErrorData::internal_error(
format!("failed to serialize initialize response: {err}"),
None,
),
);
return;
}
};
@@ -287,12 +296,12 @@ impl MessageProcessor {
}
self.initialized = true;
self.outgoing.send_response(id, result_value).await;
self.outgoing.send_response(id, result_value);
}
async fn handle_ping(&self, id: RequestId) {
tracing::info!("ping");
self.outgoing.send_response(id, json!({})).await;
self.outgoing.send_response(id, json!({}));
}
fn handle_list_resources(&self, params: Option<rmcp::model::PaginatedRequestParams>) {
@@ -334,7 +343,7 @@ impl MessageProcessor {
create_tool_for_codex_tool_call_reply_param(),
]);
self.outgoing.send_response(id, result).await;
self.outgoing.send_response(id, result);
}
async fn handle_call_tool(&self, id: RequestId, params: CallToolRequestParams) {
@@ -353,7 +362,7 @@ impl MessageProcessor {
let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(format!(
"Unknown tool '{name}'"
))]);
self.outgoing.send_response(id, result).await;
self.outgoing.send_response(id, result);
}
}
}
@@ -372,7 +381,7 @@ impl MessageProcessor {
let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
format!("Failed to load Codex configuration from overrides: {e}"),
)]);
self.outgoing.send_response(id, result).await;
self.outgoing.send_response(id, result);
return;
}
},
@@ -380,7 +389,7 @@ impl MessageProcessor {
let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
format!("Failed to parse configuration for Codex tool: {e}"),
)]);
self.outgoing.send_response(id, result).await;
self.outgoing.send_response(id, result);
return;
}
},
@@ -388,7 +397,7 @@ impl MessageProcessor {
let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
"Missing arguments for codex tool-call; the `prompt` field is required.",
)]);
self.outgoing.send_response(id, result).await;
self.outgoing.send_response(id, result);
return;
}
};
@@ -396,7 +405,7 @@ impl MessageProcessor {
// Clone outgoing and server to move into async task.
let outgoing = self.outgoing.clone();
let thread_manager = self.thread_manager.clone();
let running_requests_id_to_codex_uuid = self.running_requests_id_to_codex_uuid.clone();
let active_turns = Arc::clone(&self.active_turns);
// Spawn an async task to handle the Codex session so that we do not
// block the synchronous message-processing loop.
@@ -408,7 +417,7 @@ impl MessageProcessor {
config,
outgoing,
thread_manager,
running_requests_id_to_codex_uuid,
active_turns,
)
.await;
});
@@ -431,7 +440,7 @@ impl MessageProcessor {
let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
format!("Failed to parse configuration for Codex tool: {e}"),
)]);
self.outgoing.send_response(request_id, result).await;
self.outgoing.send_response(request_id, result);
return;
}
},
@@ -442,7 +451,7 @@ impl MessageProcessor {
let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
"Missing arguments for codex-reply tool-call; the `thread_id` and `prompt` fields are required.",
)]);
self.outgoing.send_response(request_id, result).await;
self.outgoing.send_response(request_id, result);
return;
}
};
@@ -454,14 +463,14 @@ impl MessageProcessor {
let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(format!(
"Failed to parse thread_id: {e}"
))]);
self.outgoing.send_response(request_id, result).await;
self.outgoing.send_response(request_id, result);
return;
}
};
// Clone outgoing to move into async task.
let outgoing = self.outgoing.clone();
let running_requests_id_to_codex_uuid = self.running_requests_id_to_codex_uuid.clone();
let active_turns = Arc::clone(&self.active_turns);
let codex = match self.thread_manager.get_thread(thread_id).await {
Ok(c) => c,
@@ -472,7 +481,7 @@ impl MessageProcessor {
format!("Session not found for thread_id: {thread_id}"),
Some(true),
);
outgoing.send_response(request_id, result).await;
outgoing.send_response(request_id, result);
return;
}
};
@@ -481,7 +490,7 @@ impl MessageProcessor {
let prompt = codex_tool_call_reply_param.prompt.clone();
tokio::spawn({
let outgoing = outgoing.clone();
let running_requests_id_to_codex_uuid = running_requests_id_to_codex_uuid.clone();
let active_turns = Arc::clone(&active_turns);
async move {
crate::codex_tool_runner::run_codex_tool_session_reply(
@@ -490,7 +499,7 @@ impl MessageProcessor {
outgoing,
request_id,
prompt,
running_requests_id_to_codex_uuid,
active_turns,
)
.await;
}
@@ -507,16 +516,14 @@ impl MessageProcessor {
}
async fn handle_unsupported_request(&self, id: RequestId, method: &str) {
self.outgoing
.send_error(
id,
ErrorData::new(
ErrorCode::METHOD_NOT_FOUND,
format!("method not found: {method}"),
Some(json!({ "method": method })),
),
)
.await;
self.outgoing.send_error(
id,
ErrorData::new(
ErrorCode::METHOD_NOT_FOUND,
format!("method not found: {method}"),
Some(json!({ "method": method })),
),
);
}
// ---------------------------------------------------------------------
@@ -531,15 +538,12 @@ impl MessageProcessor {
// Create a stable string form early for logging and submission id.
let request_id_string = request_id.to_string();
// Obtain the thread id while holding the first lock, then release.
let thread_id = {
let map_guard = self.running_requests_id_to_codex_uuid.lock().await;
match map_guard.get(&request_id) {
Some(id) => *id,
None => {
tracing::warn!("Session not found for request_id: {request_id_string}");
return;
}
// Resolve the thread for the active MCP request.
let thread_id = match self.active_turns.thread_id(&request_id) {
Some(id) => id,
None => {
tracing::warn!("Session not found for request_id: {request_id_string}");
return;
}
};
tracing::info!("thread_id: {thread_id}");
@@ -567,11 +571,8 @@ impl MessageProcessor {
tracing::error!("Failed to submit interrupt to Codex: {e}");
return;
}
// unregister the id so we don't keep it in the map
self.running_requests_id_to_codex_uuid
.lock()
.await
.remove(&request_id);
// Stop routing extension events to the cancelled turn.
self.active_turns.unregister(&request_id);
}
fn handle_progress_notification(&self, params: rmcp::model::ProgressNotificationParam) {

38
codex-rs/mcp-server/src/outgoing_message.rs
View File

@@ -79,15 +79,14 @@ impl OutgoingMessageSender {
}
}
pub(crate) async fn send_response<T: Serialize>(&self, id: RequestId, response: T) {
pub(crate) fn send_response<T: Serialize>(&self, id: RequestId, response: T) {
let mut result = match serde_json::to_value(response) {
Ok(result) => result,
Err(err) => {
self.send_error(
id,
ErrorData::internal_error(format!("failed to serialize response: {err}"), None),
)
.await;
);
return;
}
};
@@ -109,7 +108,7 @@ impl OutgoingMessageSender {
/// This is used with the MCP server, but not the more general JSON-RPC app
/// server. Prefer [`OutgoingMessageSender::send_server_notification`] where
/// possible.
pub(crate) async fn send_event_as_notification(
pub(crate) fn send_event_as_notification(
&self,
event: &Event,
meta: Option<OutgoingNotificationMeta>,
@@ -129,17 +128,16 @@ impl OutgoingMessageSender {
self.send_notification(OutgoingNotification {
method: "codex/event".to_string(),
params: Some(params.clone()),
})
.await;
params: Some(params),
});
}
pub(crate) async fn send_notification(&self, notification: OutgoingNotification) {
pub(crate) fn send_notification(&self, notification: OutgoingNotification) {
let outgoing_message = OutgoingMessage::Notification(notification);
let _ = self.sender.send(outgoing_message);
}
pub(crate) async fn send_error(&self, id: RequestId, error: ErrorData) {
pub(crate) fn send_error(&self, id: RequestId, error: ErrorData) {
let outgoing_message = OutgoingMessage::Error(OutgoingError { id, error });
let _ = self.sender.send(outgoing_message);
}
@@ -301,12 +299,10 @@ mod tests {
let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel::<OutgoingMessage>();
let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx);
outgoing_message_sender
.send_response(
RequestId::Number(1),
rmcp::model::CallToolResult::success(Vec::new()),
)
.await;
outgoing_message_sender.send_response(
RequestId::Number(1),
rmcp::model::CallToolResult::success(Vec::new()),
);
let Some(OutgoingMessage::Response(response)) = outgoing_rx.recv().await else {
panic!("expected a tool-call response");
@@ -346,9 +342,7 @@ mod tests {
}),
};
outgoing_message_sender
.send_event_as_notification(&event, /*meta*/ None)
.await;
outgoing_message_sender.send_event_as_notification(&event, /*meta*/ None);
let result = outgoing_rx.recv().await.unwrap();
let OutgoingMessage::Notification(OutgoingNotification { method, params }) = result else {
@@ -399,9 +393,7 @@ mod tests {
thread_id: None,
};
outgoing_message_sender
.send_event_as_notification(&event, Some(meta))
.await;
outgoing_message_sender.send_event_as_notification(&event, Some(meta));
let result = outgoing_rx.recv().await.unwrap();
let OutgoingMessage::Notification(OutgoingNotification { method, params }) = result else {
@@ -467,9 +459,7 @@ mod tests {
thread_id: Some(thread_id),
};
outgoing_message_sender
.send_event_as_notification(&event, Some(meta))
.await;
outgoing_message_sender.send_event_as_notification(&event, Some(meta));
let result = outgoing_rx.recv().await.unwrap();
let OutgoingMessage::Notification(OutgoingNotification { method, params }) = result else {

4
codex-rs/mcp-server/src/patch_approval.rs
View File

@@ -78,9 +78,7 @@ pub(crate) async fn handle_patch_approval_request(
let message = format!("Failed to serialize PatchApprovalElicitRequestParams: {err}");
error!("{message}");
outgoing
.send_error(request_id.clone(), ErrorData::invalid_params(message, None))
.await;
outgoing.send_error(request_id.clone(), ErrorData::invalid_params(message, None));
return;
}

44
codex-rs/mcp-server/tests/common/mcp_process.rs
View File

@@ -311,23 +311,41 @@ impl McpProcess {
) -> anyhow::Result<JsonRpcNotification<CustomNotification>> {
eprintln!("in read_stream_until_legacy_task_complete_notification()");
self.read_stream_until_codex_event("task_complete").await
}
pub async fn read_stream_until_codex_event(
&mut self,
event_type: &str,
) -> anyhow::Result<JsonRpcNotification<CustomNotification>> {
self.read_stream_until_codex_event_matching(event_type, |_| true)
.await
}
pub async fn read_stream_until_codex_event_matching(
&mut self,
event_type: &str,
predicate: impl Fn(&serde_json::Value) -> bool,
) -> anyhow::Result<JsonRpcNotification<CustomNotification>> {
eprintln!("in read_stream_until_codex_event({event_type})");
loop {
let message = self.read_jsonrpc_message().await?;
match message {
JsonRpcMessage::Notification(notification) => {
let is_match = if notification.notification.method == "codex/event" {
if let Some(params) = &notification.notification.params {
params
.get("msg")
.and_then(|m| m.get("type"))
.and_then(|t| t.as_str())
== Some("task_complete")
} else {
false
}
} else {
false
};
let is_match = notification.notification.method == "codex/event"
&& notification
.notification
.params
.as_ref()
.is_some_and(|params| {
params
.get("msg")
.and_then(|m| m.get("type"))
.and_then(|t| t.as_str())
== Some(event_type)
&& predicate(params)
});
if is_match {
return Ok(notification);

88
codex-rs/mcp-server/tests/suite/codex_tool.rs
View File

@@ -378,6 +378,12 @@ async fn codex_tool_passes_base_instructions() -> anyhow::Result<()> {
// Run `codex mcp` with a specific config.toml.
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), &server.uri())?;
let skill_dir = codex_home.path().join("skills").join("demo");
std::fs::create_dir_all(&skill_dir)?;
std::fs::write(
skill_dir.join("SKILL.md"),
"---\nname: demo\ndescription: Demo skill.\n---\n# Demo\n\nUse this skill.\n",
)?;
write_chatgpt_auth(
codex_home.path(),
ChatGptAuthFixture::new("chatgpt-token").account_id("workspace-123"),
@@ -471,6 +477,11 @@ async fn codex_tool_passes_base_instructions() -> anyhow::Result<()> {
1
);
assert_eq!(developer_text.matches("Generated with Codex.").count(), 1);
assert_eq!(
developer_text.matches("- demo: Demo skill.").count(),
1,
"host skill catalog should be included exactly once"
);
assert!(
developer_contents
.iter()
@@ -492,6 +503,83 @@ async fn codex_tool_passes_base_instructions() -> anyhow::Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_codex_tool_forwards_skills_extension_warnings() {
skip_if_no_network!();
codex_tool_forwards_skills_extension_warnings()
.await
.expect("codex tool should forward skills extension warnings");
}
async fn codex_tool_forwards_skills_extension_warnings() -> anyhow::Result<()> {
let server =
create_mock_responses_server(vec![create_final_assistant_message_sse_response("Enjoy!")?])
.await;
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), &server.uri())?;
let skills_dir = codex_home.path().join("skills");
for index in 0..200 {
let name = format!("skill-{index:03}");
let skill_dir = skills_dir.join(&name);
std::fs::create_dir_all(&skill_dir)?;
let description = format!("Skill {index}: {}", "x".repeat(200));
std::fs::write(
skill_dir.join("SKILL.md"),
format!(
"---\nname: {name}\ndescription: {description}\n---\n# {name}\n\nUse this skill.\n"
),
)?;
}
let mut mcp_process = McpProcess::new(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp_process.initialize()).await??;
let codex_request_id = mcp_process
.send_codex_tool_call(CodexToolCallParam {
prompt: "How are you?".to_string(),
..Default::default()
})
.await?;
let warning = timeout(
DEFAULT_READ_TIMEOUT,
mcp_process.read_stream_until_codex_event_matching("warning", |params| {
params["msg"]["message"]
.as_str()
.is_some_and(|message| message.contains("skills context budget"))
}),
)
.await??;
let warning_json = serde_json::to_value(&warning)?;
let params = warning
.notification
.params
.ok_or_else(|| anyhow::anyhow!("warning notification should include params"))?;
assert_eq!(
warning_json["params"]["_meta"]["requestId"],
codex_request_id
);
assert_eq!(warning_json["params"]["id"], codex_request_id.to_string());
assert!(
warning_json["params"]["_meta"]["threadId"]
.as_str()
.is_some_and(|thread_id| !thread_id.is_empty())
);
assert_eq!(params["msg"]["type"], "warning");
assert!(
params["msg"]["message"]
.as_str()
.is_some_and(|message| message.contains("skills context budget"))
);
timeout(
DEFAULT_READ_TIMEOUT,
mcp_process.read_stream_until_response_message(RequestId::Number(codex_request_id)),
)
.await??;
Ok(())
}
fn create_expected_patch_approval_elicitation_request_params(
changes: HashMap<PathBuf, FileChange>,
grant_root: Option<PathBuf>,