diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index cb9578be09..a2635389b5 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -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", ] diff --git a/codex-rs/mcp-server/Cargo.toml b/codex-rs/mcp-server/Cargo.toml index 9c4df86294..82653c6bc5 100644 --- a/codex-rs/mcp-server/Cargo.toml +++ b/codex-rs/mcp-server/Cargo.toml @@ -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 } diff --git a/codex-rs/mcp-server/src/active_turn_registry.rs b/codex-rs/mcp-server/src/active_turn_registry.rs new file mode 100644 index 0000000000..74652810b4 --- /dev/null +++ b/codex-rs/mcp-server/src/active_turn_registry.rs @@ -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, +} + +#[derive(Default)] +struct ActiveTurnState { + by_request: HashMap, + 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 { + 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; diff --git a/codex-rs/mcp-server/src/active_turn_registry_tests.rs b/codex-rs/mcp-server/src/active_turn_registry_tests.rs new file mode 100644 index 0000000000..55c7ffe3df --- /dev/null +++ b/codex-rs/mcp-server/src/active_turn_registry_tests.rs @@ -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)); +} diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index 39dc9f9e8c..dd551a580a 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -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, thread_manager: Arc, - running_requests_id_to_codex_uuid: Arc>>, + active_turns: Arc, ) { 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, request_id: RequestId, prompt: String, - running_requests_id_to_codex_uuid: Arc>>, + active_turns: Arc, ) { - 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, outgoing: Arc, request_id: RequestId, - running_requests_id_to_codex_uuid: Arc>>, + active_turns: Arc, ) { 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; } } diff --git a/codex-rs/mcp-server/src/exec_approval.rs b/codex-rs/mcp-server/src/exec_approval.rs index c1427a0578..5296315363 100644 --- a/codex-rs/mcp-server/src/exec_approval.rs +++ b/codex-rs/mcp-server/src/exec_approval.rs @@ -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; } diff --git a/codex-rs/mcp-server/src/extension_event_sink.rs b/codex-rs/mcp-server/src/extension_event_sink.rs new file mode 100644 index 0000000000..7c802cd246 --- /dev/null +++ b/codex-rs/mcp-server/src/extension_event_sink.rs @@ -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, + active_turns: Arc, +) -> Arc { + Arc::new(McpExtensionEventSink { + outgoing, + active_turns, + }) +} + +struct McpExtensionEventSink { + outgoing: Arc, + active_turns: Arc, +} + +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; diff --git a/codex-rs/mcp-server/src/extension_event_sink_tests.rs b/codex-rs/mcp-server/src/extension_event_sink_tests.rs new file mode 100644 index 0000000000..806b573a6d --- /dev/null +++ b/codex-rs/mcp-server/src/extension_event_sink_tests.rs @@ -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, + Arc, + Arc, + mpsc::UnboundedReceiver, +) { + 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| ¶ms["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| ¶ms["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()); +} diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index 079fd07eff..998fe2a33a 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -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; diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 3bc7f14cdd..eb5cec4e81 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -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, - running_requests_id_to_codex_uuid: Arc>>, + active_turns: Arc, } impl MessageProcessor { @@ -66,7 +65,10 @@ impl MessageProcessor { let user_instructions_provider = Arc::new(CodexHomeUserInstructionsProvider::new( config.codex_home.clone(), )); - let mut extensions = ExtensionRegistryBuilder::::new(); + let active_turns = Arc::new(ActiveTurnRegistry::default()); + let mut extensions = ExtensionRegistryBuilder::::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) { @@ -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) { diff --git a/codex-rs/mcp-server/src/outgoing_message.rs b/codex-rs/mcp-server/src/outgoing_message.rs index bb32f2a4b8..8e8a2c683c 100644 --- a/codex-rs/mcp-server/src/outgoing_message.rs +++ b/codex-rs/mcp-server/src/outgoing_message.rs @@ -79,15 +79,14 @@ impl OutgoingMessageSender { } } - pub(crate) async fn send_response(&self, id: RequestId, response: T) { + pub(crate) fn send_response(&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, @@ -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::(); 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 { diff --git a/codex-rs/mcp-server/src/patch_approval.rs b/codex-rs/mcp-server/src/patch_approval.rs index 56eca276b3..9cc1815908 100644 --- a/codex-rs/mcp-server/src/patch_approval.rs +++ b/codex-rs/mcp-server/src/patch_approval.rs @@ -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; } diff --git a/codex-rs/mcp-server/tests/common/mcp_process.rs b/codex-rs/mcp-server/tests/common/mcp_process.rs index 9b4badd4ce..f1ad725d24 100644 --- a/codex-rs/mcp-server/tests/common/mcp_process.rs +++ b/codex-rs/mcp-server/tests/common/mcp_process.rs @@ -311,23 +311,41 @@ impl McpProcess { ) -> anyhow::Result> { 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> { + 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> { + 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) = ¬ification.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); diff --git a/codex-rs/mcp-server/tests/suite/codex_tool.rs b/codex-rs/mcp-server/tests/suite/codex_tool.rs index 75db8b0e6c..0505f30be7 100644 --- a/codex-rs/mcp-server/tests/suite/codex_tool.rs +++ b/codex-rs/mcp-server/tests/suite/codex_tool.rs @@ -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, grant_root: Option,