diff --git a/codex-rs/cli/src/proto.rs b/codex-rs/cli/src/proto.rs index 6c1de7eaa9..5082a791b0 100644 --- a/codex-rs/cli/src/proto.rs +++ b/codex-rs/cli/src/proto.rs @@ -3,11 +3,11 @@ use std::sync::Arc; use clap::Parser; use codex_common::CliConfigOverrides; -use codex_core::Codex; -use codex_core::CodexSpawnOk; use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_core::protocol::Submission; +use codex_core::server::CodexConversation; +use codex_core::server::CodexServer; use codex_core::util::notify_on_sigint; use codex_login::CodexAuth; use tokio::io::AsyncBufReadExt; @@ -36,10 +36,12 @@ pub async fn run_main(opts: ProtoCli) -> anyhow::Result<()> { .map_err(anyhow::Error::msg)?; let config = Config::load_with_cli_overrides(overrides_vec, ConfigOverrides::default())?; - let auth = CodexAuth::from_codex_home(&config.codex_home)?; + let _auth = CodexAuth::from_codex_home(&config.codex_home)?; let ctrl_c = notify_on_sigint(); - let CodexSpawnOk { codex, .. } = Codex::spawn(config, auth, ctrl_c.clone()).await?; - let codex = Arc::new(codex); + // Use server API to start a conversation + let server = CodexServer::default(); + let new_conv = server.new_conversation(config).await?; + let codex: Arc = server.get_conversation(new_conv.conversation_id).await?; // Task that reads JSON lines from stdin and forwards to Submission Queue let sq_fut = { diff --git a/codex-rs/core/src/codex_wrapper.rs b/codex-rs/core/src/codex_wrapper.rs deleted file mode 100644 index dc10ec8d84..0000000000 --- a/codex-rs/core/src/codex_wrapper.rs +++ /dev/null @@ -1,59 +0,0 @@ -use std::sync::Arc; - -use crate::Codex; -use crate::CodexSpawnOk; -use crate::config::Config; -use crate::protocol::Event; -use crate::protocol::EventMsg; -use crate::util::notify_on_sigint; -use codex_login::CodexAuth; -use tokio::sync::Notify; -use uuid::Uuid; - -/// Represents an active Codex conversation, including the first event -/// (which is [`EventMsg::SessionConfigured`]). -pub struct CodexConversation { - pub codex: Codex, - pub session_id: Uuid, - pub session_configured: Event, - pub ctrl_c: Arc, -} - -/// Spawn a new [`Codex`] and initialize the session. -/// -/// Returns the wrapped [`Codex`] **and** the `SessionInitialized` event that -/// is received as a response to the initial `ConfigureSession` submission so -/// that callers can surface the information to the UI. -pub async fn init_codex(config: Config) -> anyhow::Result { - let ctrl_c = notify_on_sigint(); - let auth = CodexAuth::from_codex_home(&config.codex_home)?; - let CodexSpawnOk { - codex, - init_id, - session_id, - } = Codex::spawn(config, auth, ctrl_c.clone()).await?; - - // The first event must be `SessionInitialized`. Validate and forward it to - // the caller so that they can display it in the conversation history. - let event = codex.next_event().await?; - if event.id != init_id - || !matches!( - &event, - Event { - id: _id, - msg: EventMsg::SessionConfigured(_), - } - ) - { - return Err(anyhow::anyhow!( - "expected SessionInitialized but got {event:?}" - )); - } - - Ok(CodexConversation { - codex, - session_id, - session_configured: event, - ctrl_c, - }) -} diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index 2931d30636..da7c528207 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -3,6 +3,7 @@ use serde_json; use std::io; use thiserror::Error; use tokio::task::JoinError; +use uuid::Uuid; pub type Result = std::result::Result; @@ -44,6 +45,9 @@ pub enum CodexErr { #[error("stream disconnected before completion: {0}")] Stream(String), + #[error("no conversation with id: {0}")] + ConversationNotFound(Uuid), + /// Returned by run_command_stream when the spawned child process timed out (10s). #[error("timeout waiting for child process to exit")] Timeout, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index b36689f057..5724643be9 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -13,7 +13,6 @@ mod client_common; pub mod codex; pub use codex::Codex; pub use codex::CodexSpawnOk; -pub mod codex_wrapper; pub mod config; pub mod config_profile; pub mod config_types; @@ -44,6 +43,7 @@ pub mod protocol; mod rollout; pub(crate) mod safety; pub mod seatbelt; +pub mod server; pub mod shell; pub mod spawn; pub mod turn_diff_tracker; diff --git a/codex-rs/core/src/server.rs b/codex-rs/core/src/server.rs new file mode 100644 index 0000000000..dea35c8afe --- /dev/null +++ b/codex-rs/core/src/server.rs @@ -0,0 +1,117 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use codex_login::CodexAuth; +use tokio::sync::Notify; +use tokio::sync::RwLock; +use tokio::sync::futures::Notified; +use uuid::Uuid; + +use crate::Codex; +use crate::CodexSpawnOk; +use crate::config::Config; +use crate::error::CodexErr; +use crate::error::Result as CodexResult; +use crate::protocol::Event; +use crate::protocol::EventMsg; +use crate::protocol::Op; +use crate::protocol::SessionConfiguredEvent; +use crate::protocol::Submission; + +/// Codex "server" that manages multiple conversations. +pub struct CodexServer { + conversations: Arc>>>, +} + +/// Represents an active Codex conversation, including the first event +/// (which is [`EventMsg::SessionConfigured`]). +pub struct CodexConversation { + codex: Codex, + cancellation_token: Arc, +} + +impl CodexConversation { + pub async fn submit(&self, op: Op) -> CodexResult { + self.codex.submit(op).await + } + + pub async fn submit_with_id(&self, sub: Submission) -> CodexResult<()> { + self.codex.submit_with_id(sub).await + } + + pub async fn next_event(&self) -> CodexResult { + self.codex.next_event().await + } + + pub fn on_cancel(&self) -> Notified { + self.cancellation_token.notified() + } +} + +pub struct NewConversation { + pub conversation_id: Uuid, + pub session_configured: SessionConfiguredEvent, +} + +impl Default for CodexServer { + fn default() -> Self { + Self { + conversations: Arc::new(RwLock::new(HashMap::new())), + } + } +} + +impl CodexServer { + pub async fn new_conversation(&self, config: Config) -> anyhow::Result { + // TODO(mbolin): Determine whether this should be wired up to SIGINT. + let cancellation_token = Arc::new(Notify::new()); + let auth = CodexAuth::from_codex_home(&config.codex_home)?; + + let CodexSpawnOk { + codex, + init_id, + session_id: conversation_id, + } = Codex::spawn(config, auth, cancellation_token.clone()).await?; + + // The first event must be `SessionInitialized`. Validate and forward it + // to the caller so that they can display it in the conversation + // history. + let event = codex.next_event().await?; + let session_configured = match event { + Event { + id, + msg: EventMsg::SessionConfigured(session_configured), + } if id == init_id => session_configured, + _ => { + return Err(anyhow::anyhow!( + "expected SessionInitialized but got {event:?}" + )); + } + }; + + let codex_conversation = Arc::new(CodexConversation { + codex, + cancellation_token, + }); + self.conversations + .write() + .await + .insert(conversation_id, codex_conversation); + + Ok(NewConversation { + conversation_id, + session_configured, + }) + } + + pub async fn get_conversation( + &self, + conversation_id: Uuid, + ) -> CodexResult> { + let conversations = self.conversations.read().await; + conversations + .get(&conversation_id) + .cloned() + .ok_or_else(|| CodexErr::ConversationNotFound(conversation_id)) + } +} diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 6ed57898b2..89385a85a8 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -6,12 +6,9 @@ mod event_processor_with_json_output; use std::io::IsTerminal; use std::io::Read; use std::path::PathBuf; -use std::sync::Arc; pub use cli::Cli; use codex_core::BUILT_IN_OSS_MODEL_PROVIDER_ID; -use codex_core::codex_wrapper::CodexConversation; -use codex_core::codex_wrapper::{self}; use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_core::config_types::SandboxMode; @@ -21,6 +18,8 @@ use codex_core::protocol::EventMsg; use codex_core::protocol::InputItem; use codex_core::protocol::Op; use codex_core::protocol::TaskCompleteEvent; +use codex_core::server::CodexServer; +use codex_core::server::NewConversation; use codex_core::util::is_inside_git_repo; use codex_ollama::DEFAULT_OSS_MODEL; use event_processor_with_human_output::EventProcessorWithHumanOutput; @@ -185,25 +184,24 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option) -> any std::process::exit(1); } - let CodexConversation { - codex: codex_wrapper, + let server = CodexServer::default(); + let NewConversation { + conversation_id, session_configured, - ctrl_c, - .. - } = codex_wrapper::init_codex(config).await?; - let codex = Arc::new(codex_wrapper); + } = server.new_conversation(config).await?; info!("Codex initialized with event: {session_configured:?}"); + let conversation = server.get_conversation(conversation_id).await?; let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); { - let codex = codex.clone(); + let conversation = conversation.clone(); tokio::spawn(async move { loop { - let interrupted = ctrl_c.notified(); + let interrupted = conversation.on_cancel(); tokio::select! { _ = interrupted => { // Forward an interrupt to the codex so it can abort any in‑flight task. - let _ = codex + let _ = conversation .submit( Op::Interrupt, ) @@ -213,7 +211,7 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option) -> any // will emit a `TurnInterrupted` (Error) event which is drained later. break; } - res = codex.next_event() => match res { + res = conversation.next_event() => match res { Ok(event) => { debug!("Received event: {event:?}"); @@ -243,9 +241,9 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option) -> any .into_iter() .map(|path| InputItem::LocalImage { path }) .collect(); - let initial_images_event_id = codex.submit(Op::UserInput { items }).await?; + let initial_images_event_id = conversation.submit(Op::UserInput { items }).await?; info!("Sent images with event ID: {initial_images_event_id}"); - while let Ok(event) = codex.next_event().await { + while let Ok(event) = conversation.next_event().await { if event.id == initial_images_event_id && matches!( event.msg, @@ -261,7 +259,7 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option) -> any // Send the prompt. let items: Vec = vec![InputItem::Text { text: prompt }]; - let initial_prompt_task_id = codex.submit(Op::UserInput { items }).await?; + let initial_prompt_task_id = conversation.submit(Op::UserInput { items }).await?; info!("Sent prompt with event ID: {initial_prompt_task_id}"); // Run the loop until the task is complete. @@ -270,7 +268,7 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option) -> any match shutdown { CodexStatus::Running => continue, CodexStatus::InitiateShutdown => { - codex.submit(Op::Shutdown).await?; + conversation.submit(Op::Shutdown).await?; } CodexStatus::Shutdown => { break; diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index b91c4a7609..2551191a4b 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -5,18 +5,18 @@ 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::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; use codex_core::protocol::Op; use codex_core::protocol::Submission; use codex_core::protocol::TaskCompleteEvent; +use codex_core::server::CodexConversation; +use codex_core::server::CodexServer; use mcp_types::CallToolResult; use mcp_types::ContentBlock; use mcp_types::RequestId; @@ -41,15 +41,10 @@ pub async fn run_codex_tool_session( initial_prompt: String, config: CodexConfig, outgoing: Arc, - session_map: Arc>>>, + codex_server: Arc, running_requests_id_to_codex_uuid: Arc>>, ) { - let CodexConversation { - codex, - session_configured, - session_id, - .. - } = match init_codex(config).await { + let new_conv = match codex_server.new_conversation(config).await { Ok(res) => res, Err(e) => { let result = CallToolResult { @@ -65,16 +60,32 @@ 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_id = new_conv.conversation_id; + let session_configured = new_conv.session_configured; + let codex: Arc = match codex_server.get_conversation(session_id).await { + Ok(c) => c, + Err(e) => { + let result = CallToolResult { + content: vec![ContentBlock::TextContent(TextContent { + r#type: "text".to_string(), + text: format!("Failed to get conversation handle: {e}"), + annotations: None, + })], + is_error: Some(true), + structured_content: None, + }; + outgoing.send_response(id.clone(), result.into()).await; + return; + } + }; + let session_configured_event = Event { + id: "init".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; @@ -110,7 +121,7 @@ pub async fn run_codex_tool_session( } pub async fn run_codex_tool_session_reply( - codex: Arc, + codex: Arc, outgoing: Arc, request_id: RequestId, prompt: String, @@ -146,7 +157,7 @@ pub async fn run_codex_tool_session_reply( } async fn run_codex_tool_session_inner( - codex: Arc, + codex: Arc, outgoing: Arc, request_id: RequestId, running_requests_id_to_codex_uuid: Arc>>, diff --git a/codex-rs/mcp-server/src/conversation_loop.rs b/codex-rs/mcp-server/src/conversation_loop.rs index 80c34760c5..1995bd4be1 100644 --- a/codex-rs/mcp-server/src/conversation_loop.rs +++ b/codex-rs/mcp-server/src/conversation_loop.rs @@ -4,16 +4,16 @@ 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::protocol::AgentMessageEvent; use codex_core::protocol::ApplyPatchApprovalRequestEvent; use codex_core::protocol::EventMsg; use codex_core::protocol::ExecApprovalRequestEvent; +use codex_core::server::CodexConversation; use mcp_types::RequestId; use tracing::error; pub async fn run_conversation_loop( - codex: Arc, + codex: Arc, outgoing: Arc, request_id: RequestId, ) { diff --git a/codex-rs/mcp-server/src/exec_approval.rs b/codex-rs/mcp-server/src/exec_approval.rs index 54abfb35bd..5fb5db3d61 100644 --- a/codex-rs/mcp-server/src/exec_approval.rs +++ b/codex-rs/mcp-server/src/exec_approval.rs @@ -1,9 +1,9 @@ use std::path::PathBuf; use std::sync::Arc; -use codex_core::Codex; use codex_core::protocol::Op; use codex_core::protocol::ReviewDecision; +use codex_core::server::CodexConversation; use mcp_types::ElicitRequest; use mcp_types::ElicitRequestParamsRequestedSchema; use mcp_types::JSONRPCErrorError; @@ -51,7 +51,7 @@ pub(crate) async fn handle_exec_approval_request( command: Vec, cwd: PathBuf, outgoing: Arc, - codex: Arc, + codex: Arc, 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, - codex: Arc, + codex: Arc, ) { let response = receiver.await; let value = match response { diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 2f99cda723..a64438a832 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -14,9 +14,9 @@ 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::config::Config as CodexConfig; use codex_core::protocol::Submission; +use codex_core::server::CodexServer; use mcp_types::CallToolRequest; use mcp_types::CallToolRequestParams; use mcp_types::CallToolResult; @@ -42,7 +42,7 @@ pub(crate) struct MessageProcessor { outgoing: Arc, initialized: bool, codex_linux_sandbox_exe: Option, - session_map: Arc>>>, + codex_server: Arc, running_requests_id_to_codex_uuid: Arc>>, running_session_ids: Arc>>, } @@ -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(CodexServer::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>>> { - self.session_map.clone() + pub(crate) fn codex_server(&self) -> Arc { + self.codex_server.clone() } pub(crate) fn outgoing(&self) -> Arc { @@ -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; } }; diff --git a/codex-rs/mcp-server/src/patch_approval.rs b/codex-rs/mcp-server/src/patch_approval.rs index db99ee5f27..67e13625b7 100644 --- a/codex-rs/mcp-server/src/patch_approval.rs +++ b/codex-rs/mcp-server/src/patch_approval.rs @@ -2,10 +2,10 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; -use codex_core::Codex; use codex_core::protocol::FileChange; use codex_core::protocol::Op; use codex_core::protocol::ReviewDecision; +use codex_core::server::CodexConversation; use mcp_types::ElicitRequest; use mcp_types::ElicitRequestParamsRequestedSchema; use mcp_types::JSONRPCErrorError; @@ -47,7 +47,7 @@ pub(crate) async fn handle_patch_approval_request( grant_root: Option, changes: HashMap, outgoing: Arc, - codex: Arc, + codex: Arc, 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, - codex: Arc, + codex: Arc, ) { let response = receiver.await; let value = match response { diff --git a/codex-rs/mcp-server/src/tool_handlers/create_conversation.rs b/codex-rs/mcp-server/src/tool_handlers/create_conversation.rs index 559bf72905..d904ffe70f 100644 --- a/codex-rs/mcp-server/src/tool_handlers/create_conversation.rs +++ b/codex-rs/mcp-server/src/tool_handlers/create_conversation.rs @@ -1,16 +1,11 @@ -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::config::Config as CodexConfig; use codex_core::config::ConfigOverrides; -use codex_core::protocol::EventMsg; -use codex_core::protocol::SessionConfiguredEvent; +use codex_core::server::CodexConversation; +use codex_core::server::NewConversation; 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 +76,11 @@ 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, + session_configured, + } = match message_processor.codex_server().new_conversation(cfg).await { Ok(conv) => conv, Err(e) => { message_processor @@ -100,36 +98,33 @@ 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(); + + // Obtain the Codex conversation handle from the server + let codex_arc: Arc = match message_processor + .codex_server() + .get_conversation(session_id) + .await + { + Ok(conv) => conv, + Err(e) => { + message_processor + .send_response_with_optional_error( + id, + Some(ToolCallResponseResult::ConversationCreate( + ConversationCreateResult::Error { + message: format!( + "Failed to get conversation handle for {session_id}: {e}" + ), + }, + )), + Some(true), + ) + .await; + return; + } }; - 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(); @@ -152,11 +147,4 @@ pub(crate) async fn handle_create_conversation( .await; } -async fn insert_session( - session_id: Uuid, - codex: Arc, - session_map: Arc>>>, -) { - let mut guard = session_map.lock().await; - guard.insert(session_id, codex); -} +// No longer need to insert into a local session map; CodexServer manages sessions. diff --git a/codex-rs/mcp-server/src/tool_handlers/send_message.rs b/codex-rs/mcp-server/src/tool_handlers/send_message.rs index 894176bef6..ae0b4a5b29 100644 --- a/codex-rs/mcp-server/src/tool_handlers/send_message.rs +++ b/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,4 @@ pub(crate) async fn handle_send_message( ) .await; } - -pub(crate) async fn get_session( - session_id: Uuid, - session_map: Arc>>>, -) -> Option> { - let guard = session_map.lock().await; - guard.get(&session_id).cloned() -} +// Session lookup is delegated to CodexServer now. diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 3d312ffce0..3892db68e7 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -2,8 +2,6 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; -use codex_core::codex_wrapper::CodexConversation; -use codex_core::codex_wrapper::init_codex; use codex_core::config::Config; use codex_core::parse_command::ParsedCommand; use codex_core::protocol::AgentMessageDeltaEvent; @@ -53,6 +51,8 @@ use crate::history_cell::HistoryCell; use crate::history_cell::PatchEventType; use crate::live_wrap::RowBuilder; use crate::user_approval_widget::ApprovalRequest; +use codex_core::server::CodexConversation; +use codex_core::server::CodexServer; use codex_file_search::FileMatch; use ratatui::style::Stylize; @@ -172,23 +172,31 @@ impl ChatWidget<'_> { // Create the Codex asynchronously so the UI loads as quickly as possible. let config_for_agent_loop = config.clone(); tokio::spawn(async move { - let CodexConversation { - codex, - session_configured, - .. - } = match init_codex(config_for_agent_loop).await { - Ok(vals) => vals, + let server = CodexServer::default(); + let new_conv = match server.new_conversation(config_for_agent_loop).await { + Ok(v) => v, Err(e) => { // TODO: surface this error to the user. tracing::error!("failed to initialize codex: {e}"); return; } }; + let session_configured = new_conv.session_configured.clone(); + let codex: Arc = + match server.get_conversation(new_conv.conversation_id).await { + Ok(c) => c, + Err(e) => { + tracing::error!("failed to get conversation handle: {e}"); + return; + } + }; - // Forward the captured `SessionInitialized` event that was consumed - // inside `init_codex()` so it can be rendered in the UI. - app_event_tx_clone.send(AppEvent::CodexEvent(session_configured.clone())); - let codex = Arc::new(codex); + // Forward the captured `SessionConfigured` event so it can be rendered in the UI. + let ev = codex_core::protocol::Event { + id: "init".to_string(), + msg: codex_core::protocol::EventMsg::SessionConfigured(session_configured.clone()), + }; + app_event_tx_clone.send(AppEvent::CodexEvent(ev)); let codex_clone = codex.clone(); tokio::spawn(async move { while let Some(op) = codex_op_rx.recv().await {