chore: introduce CodexServer as a clearinghouse for all conversations

codex-rs/mcp-server/src/codex_tool_runner.rs

@@ -5,12 +5,13 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use codex_core::Codex;
-use codex_core::codex_wrapper::CodexConversation;
-use codex_core::codex_wrapper::init_codex;
+use codex_core::CodexConversation;
+use codex_core::ConversationManager;
+use codex_core::NewConversation;
 use codex_core::config::Config as CodexConfig;
 use codex_core::protocol::AgentMessageEvent;
 use codex_core::protocol::ApplyPatchApprovalRequestEvent;
+use codex_core::protocol::Event;
 use codex_core::protocol::EventMsg;
 use codex_core::protocol::ExecApprovalRequestEvent;
 use codex_core::protocol::InputItem;
@@ -41,15 +42,14 @@ pub async fn run_codex_tool_session(
     initial_prompt: String,
     config: CodexConfig,
     outgoing: Arc<OutgoingMessageSender>,
-    session_map: Arc<Mutex<HashMap<Uuid, Arc<Codex>>>>,
+    codex_server: Arc<ConversationManager>,
     running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, Uuid>>>,
 ) {
-    let CodexConversation {
-        codex,
+    let NewConversation {
+        conversation_id,
+        conversation,
         session_configured,
-        session_id,
-        ..
-    } = match init_codex(config).await {
+    } = match codex_server.new_conversation(config).await {
         Ok(res) => res,
         Err(e) => {
             let result = CallToolResult {
@@ -65,16 +65,15 @@ pub async fn run_codex_tool_session(
             return;
         }
     };
-    let codex = Arc::new(codex);
-
-    // update the session map so we can retrieve the session in a reply, and then drop it, since
-    // we no longer need it for this function
-    session_map.lock().await.insert(session_id, codex.clone());
-    drop(session_map);
 
+    let session_configured_event = Event {
+        // Use a fake id value for now.
+        id: "".to_string(),
+        msg: EventMsg::SessionConfigured(session_configured.clone()),
+    };
     outgoing
         .send_event_as_notification(
-            &session_configured,
+            &session_configured_event,
             Some(OutgoingNotificationMeta::new(Some(id.clone()))),
         )
         .await;
@@ -89,7 +88,7 @@ pub async fn run_codex_tool_session(
     running_requests_id_to_codex_uuid
         .lock()
         .await
-        .insert(id.clone(), session_id);
+        .insert(id.clone(), conversation_id);
     let submission = Submission {
         id: sub_id.clone(),
         op: Op::UserInput {
@@ -99,18 +98,24 @@ pub async fn run_codex_tool_session(
         },
     };
 
-    if let Err(e) = codex.submit_with_id(submission).await {
+    if let Err(e) = conversation.submit_with_id(submission).await {
         tracing::error!("Failed to submit initial prompt: {e}");
         // unregister the id so we don't keep it in the map
         running_requests_id_to_codex_uuid.lock().await.remove(&id);
         return;
     }
 
-    run_codex_tool_session_inner(codex, outgoing, id, running_requests_id_to_codex_uuid).await;
+    run_codex_tool_session_inner(
+        conversation,
+        outgoing,
+        id,
+        running_requests_id_to_codex_uuid,
+    )
+    .await;
 }
 
 pub async fn run_codex_tool_session_reply(
-    codex: Arc<Codex>,
+    conversation: Arc<CodexConversation>,
     outgoing: Arc<OutgoingMessageSender>,
     request_id: RequestId,
     prompt: String,
@@ -121,7 +126,7 @@ pub async fn run_codex_tool_session_reply(
         .lock()
         .await
         .insert(request_id.clone(), session_id);
-    if let Err(e) = codex
+    if let Err(e) = conversation
         .submit(Op::UserInput {
             items: vec![InputItem::Text { text: prompt }],
         })
@@ -137,7 +142,7 @@ pub async fn run_codex_tool_session_reply(
     }
 
     run_codex_tool_session_inner(
-        codex,
+        conversation,
         outgoing,
         request_id,
         running_requests_id_to_codex_uuid,
@@ -146,7 +151,7 @@ pub async fn run_codex_tool_session_reply(
 }
 
 async fn run_codex_tool_session_inner(
-    codex: Arc<Codex>,
+    codex: Arc<CodexConversation>,
     outgoing: Arc<OutgoingMessageSender>,
     request_id: RequestId,
     running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, Uuid>>>,

codex-rs/mcp-server/src/conversation_loop.rs

@@ -4,7 +4,7 @@ use crate::exec_approval::handle_exec_approval_request;
 use crate::outgoing_message::OutgoingMessageSender;
 use crate::outgoing_message::OutgoingNotificationMeta;
 use crate::patch_approval::handle_patch_approval_request;
-use codex_core::Codex;
+use codex_core::CodexConversation;
 use codex_core::protocol::AgentMessageEvent;
 use codex_core::protocol::ApplyPatchApprovalRequestEvent;
 use codex_core::protocol::EventMsg;
@@ -13,7 +13,7 @@ use mcp_types::RequestId;
 use tracing::error;
 
 pub async fn run_conversation_loop(
-    codex: Arc<Codex>,
+    codex: Arc<CodexConversation>,
     outgoing: Arc<OutgoingMessageSender>,
     request_id: RequestId,
 ) {

codex-rs/mcp-server/src/exec_approval.rs

@@ -1,7 +1,7 @@
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use codex_core::Codex;
+use codex_core::CodexConversation;
 use codex_core::protocol::Op;
 use codex_core::protocol::ReviewDecision;
 use mcp_types::ElicitRequest;
@@ -51,7 +51,7 @@ pub(crate) async fn handle_exec_approval_request(
     command: Vec<String>,
     cwd: PathBuf,
     outgoing: Arc<crate::outgoing_message::OutgoingMessageSender>,
-    codex: Arc<Codex>,
+    codex: Arc<CodexConversation>,
     request_id: RequestId,
     tool_call_id: String,
     event_id: String,
@@ -116,7 +116,7 @@ pub(crate) async fn handle_exec_approval_request(
 async fn on_exec_approval_response(
     event_id: String,
     receiver: tokio::sync::oneshot::Receiver<mcp_types::Result>,
-    codex: Arc<Codex>,
+    codex: Arc<CodexConversation>,
 ) {
     let response = receiver.await;
     let value = match response {

codex-rs/mcp-server/src/message_processor.rs

@@ -14,7 +14,7 @@ use crate::outgoing_message::OutgoingMessageSender;
 use crate::tool_handlers::create_conversation::handle_create_conversation;
 use crate::tool_handlers::send_message::handle_send_message;
 
-use codex_core::Codex;
+use codex_core::ConversationManager;
 use codex_core::config::Config as CodexConfig;
 use codex_core::protocol::Submission;
 use mcp_types::CallToolRequest;
@@ -42,7 +42,7 @@ pub(crate) struct MessageProcessor {
     outgoing: Arc<OutgoingMessageSender>,
     initialized: bool,
     codex_linux_sandbox_exe: Option<PathBuf>,
-    session_map: Arc<Mutex<HashMap<Uuid, Arc<Codex>>>>,
+    codex_server: Arc<ConversationManager>,
     running_requests_id_to_codex_uuid: Arc<Mutex<HashMap<RequestId, Uuid>>>,
     running_session_ids: Arc<Mutex<HashSet<Uuid>>>,
 }
@@ -58,14 +58,14 @@ impl MessageProcessor {
             outgoing: Arc::new(outgoing),
             initialized: false,
             codex_linux_sandbox_exe,
-            session_map: Arc::new(Mutex::new(HashMap::new())),
+            codex_server: Arc::new(ConversationManager::default()),
             running_requests_id_to_codex_uuid: Arc::new(Mutex::new(HashMap::new())),
             running_session_ids: Arc::new(Mutex::new(HashSet::new())),
         }
     }
 
-    pub(crate) fn session_map(&self) -> Arc<Mutex<HashMap<Uuid, Arc<Codex>>>> {
-        self.session_map.clone()
+    pub(crate) fn codex_server(&self) -> Arc<ConversationManager> {
+        self.codex_server.clone()
     }
 
     pub(crate) fn outgoing(&self) -> Arc<OutgoingMessageSender> {
@@ -431,9 +431,9 @@ impl MessageProcessor {
             }
         };
 
-        // Clone outgoing and session map to move into async task.
+        // Clone outgoing and server to move into async task.
         let outgoing = self.outgoing.clone();
-        let session_map = self.session_map.clone();
+        let codex_server = self.codex_server.clone();
         let running_requests_id_to_codex_uuid = self.running_requests_id_to_codex_uuid.clone();
 
         // Spawn an async task to handle the Codex session so that we do not
@@ -445,7 +445,7 @@ impl MessageProcessor {
                 initial_prompt,
                 config,
                 outgoing,
-                session_map,
+                codex_server,
                 running_requests_id_to_codex_uuid,
             )
             .await;
@@ -516,33 +516,27 @@ impl MessageProcessor {
             }
         };
 
-        // load codex from session map
-        let session_map_mutex = Arc::clone(&self.session_map);
-
-        // Clone outgoing and session map to move into async task.
+        // 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 codex = {
-            let session_map = session_map_mutex.lock().await;
-            match session_map.get(&session_id).cloned() {
-                Some(c) => c,
-                None => {
-                    tracing::warn!("Session not found for session_id: {session_id}");
-                    let result = CallToolResult {
-                        content: vec![ContentBlock::TextContent(TextContent {
-                            r#type: "text".to_owned(),
-                            text: format!("Session not found for session_id: {session_id}"),
-                            annotations: None,
-                        })],
-                        is_error: Some(true),
-                        structured_content: None,
-                    };
-                    outgoing
-                        .send_response(request_id, serde_json::to_value(result).unwrap_or_default())
-                        .await;
-                    return;
-                }
+        let codex = match self.codex_server.get_conversation(session_id).await {
+            Ok(c) => c,
+            Err(_) => {
+                tracing::warn!("Session not found for session_id: {session_id}");
+                let result = CallToolResult {
+                    content: vec![ContentBlock::TextContent(TextContent {
+                        r#type: "text".to_owned(),
+                        text: format!("Session not found for session_id: {session_id}"),
+                        annotations: None,
+                    })],
+                    is_error: Some(true),
+                    structured_content: None,
+                };
+                outgoing
+                    .send_response(request_id, serde_json::to_value(result).unwrap_or_default())
+                    .await;
+                return;
             }
         };
 
@@ -609,15 +603,12 @@ impl MessageProcessor {
         };
         tracing::info!("session_id: {session_id}");
 
-        // Obtain the Codex Arc while holding the session_map lock, then release.
-        let codex_arc = {
-            let sessions_guard = self.session_map.lock().await;
-            match sessions_guard.get(&session_id) {
-                Some(codex) => Arc::clone(codex),
-                None => {
-                    tracing::warn!("Session not found for session_id: {session_id}");
-                    return;
-                }
+        // Obtain the Codex conversation from the server.
+        let codex_arc = match self.codex_server.get_conversation(session_id).await {
+            Ok(c) => c,
+            Err(_) => {
+                tracing::warn!("Session not found for session_id: {session_id}");
+                return;
             }
         };
 

codex-rs/mcp-server/src/patch_approval.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use codex_core::Codex;
+use codex_core::CodexConversation;
 use codex_core::protocol::FileChange;
 use codex_core::protocol::Op;
 use codex_core::protocol::ReviewDecision;
@@ -47,7 +47,7 @@ pub(crate) async fn handle_patch_approval_request(
     grant_root: Option<PathBuf>,
     changes: HashMap<PathBuf, FileChange>,
     outgoing: Arc<OutgoingMessageSender>,
-    codex: Arc<Codex>,
+    codex: Arc<CodexConversation>,
     request_id: RequestId,
     tool_call_id: String,
     event_id: String,
@@ -111,7 +111,7 @@ pub(crate) async fn handle_patch_approval_request(
 pub(crate) async fn on_patch_approval_response(
     event_id: String,
     receiver: tokio::sync::oneshot::Receiver<mcp_types::Result>,
-    codex: Arc<Codex>,
+    codex: Arc<CodexConversation>,
 ) {
     let response = receiver.await;
     let value = match response {

codex-rs/mcp-server/src/tool_handlers/create_conversation.rs

@@ -1,16 +1,9 @@
-use std::collections::HashMap;
 use std::path::PathBuf;
-use std::sync::Arc;
 
-use codex_core::Codex;
-use codex_core::codex_wrapper::init_codex;
+use codex_core::NewConversation;
 use codex_core::config::Config as CodexConfig;
 use codex_core::config::ConfigOverrides;
-use codex_core::protocol::EventMsg;
-use codex_core::protocol::SessionConfiguredEvent;
 use mcp_types::RequestId;
-use tokio::sync::Mutex;
-use uuid::Uuid;
 
 use crate::conversation_loop::run_conversation_loop;
 use crate::json_to_toml::json_to_toml;
@@ -81,8 +74,12 @@ pub(crate) async fn handle_create_conversation(
         }
     };
 
-    // Initialize Codex session
-    let codex_conversation = match init_codex(cfg).await {
+    // Initialize Codex session via server API
+    let NewConversation {
+        conversation_id: session_id,
+        conversation,
+        session_configured,
+    } = match message_processor.codex_server().new_conversation(cfg).await {
         Ok(conv) => conv,
         Err(e) => {
             message_processor
@@ -100,41 +97,13 @@ pub(crate) async fn handle_create_conversation(
         }
     };
 
-    // Expect SessionConfigured; if not, return error.
-    let EventMsg::SessionConfigured(SessionConfiguredEvent { model, .. }) =
-        &codex_conversation.session_configured.msg
-    else {
-        message_processor
-            .send_response_with_optional_error(
-                id,
-                Some(ToolCallResponseResult::ConversationCreate(
-                    ConversationCreateResult::Error {
-                        message: "Expected SessionConfigured event".to_string(),
-                    },
-                )),
-                Some(true),
-            )
-            .await;
-        return;
-    };
+    let effective_model = session_configured.model.clone();
 
-    let effective_model = model.clone();
-
-    let session_id = codex_conversation.session_id;
-    let codex_arc = Arc::new(codex_conversation.codex);
-
-    // Store session for future calls
-    insert_session(
-        session_id,
-        codex_arc.clone(),
-        message_processor.session_map(),
-    )
-    .await;
     // Run the conversation loop in the background so this request can return immediately.
     let outgoing = message_processor.outgoing();
     let spawn_id = id.clone();
     tokio::spawn(async move {
-        run_conversation_loop(codex_arc.clone(), outgoing, spawn_id).await;
+        run_conversation_loop(conversation.clone(), outgoing, spawn_id).await;
     });
 
     // Reply with the new conversation id and effective model
@@ -151,12 +120,3 @@ pub(crate) async fn handle_create_conversation(
         )
         .await;
 }
-
-async fn insert_session(
-    session_id: Uuid,
-    codex: Arc<Codex>,
-    session_map: Arc<Mutex<HashMap<Uuid, Arc<Codex>>>>,
-) {
-    let mut guard = session_map.lock().await;
-    guard.insert(session_id, codex);
-}

codex-rs/mcp-server/src/tool_handlers/send_message.rs

@@ -1,12 +1,6 @@
-use std::collections::HashMap;
-use std::sync::Arc;
-
-use codex_core::Codex;
 use codex_core::protocol::Op;
 use codex_core::protocol::Submission;
 use mcp_types::RequestId;
-use tokio::sync::Mutex;
-use uuid::Uuid;
 
 use crate::mcp_protocol::ConversationSendMessageArgs;
 use crate::mcp_protocol::ConversationSendMessageResult;
@@ -41,7 +35,11 @@ pub(crate) async fn handle_send_message(
     }
 
     let session_id = conversation_id.0;
-    let Some(codex) = get_session(session_id, message_processor.session_map()).await else {
+    let Ok(codex) = message_processor
+        .codex_server()
+        .get_conversation(session_id)
+        .await
+    else {
         message_processor
             .send_response_with_optional_error(
                 id,
@@ -114,11 +112,3 @@ pub(crate) async fn handle_send_message(
         )
         .await;
 }
-
-pub(crate) async fn get_session(
-    session_id: Uuid,
-    session_map: Arc<Mutex<HashMap<Uuid, Arc<Codex>>>>,
-) -> Option<Arc<Codex>> {
-    let guard = session_map.lock().await;
-    guard.get(&session_id).cloned()
-}