chore: introduce CodexServer as a clearinghouse for all conversations

Author: bolinfest

codex-rs/cli/src/proto.rs

@@ -1,15 +1,12 @@
 use std::io::IsTerminal;
-use std::sync::Arc;
 
 use clap::Parser;
 use codex_common::CliConfigOverrides;
-use codex_core::Codex;
-use codex_core::CodexSpawnOk;
+use codex_core::ConversationManager;
+use codex_core::NewConversation;
 use codex_core::config::Config;
 use codex_core::config::ConfigOverrides;
 use codex_core::protocol::Submission;
-use codex_core::util::notify_on_sigint;
-use codex_login::CodexAuth;
 use tokio::io::AsyncBufReadExt;
 use tokio::io::BufReader;
 use tracing::error;
@@ -36,21 +33,23 @@ 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 ctrl_c = notify_on_sigint();
-    let CodexSpawnOk { codex, .. } = Codex::spawn(config, auth, ctrl_c.clone()).await?;
-    let codex = Arc::new(codex);
+    // Use conversation_manager API to start a conversation
+    let conversation_manager = ConversationManager::default();
+    let NewConversation {
+        conversation_id: _,
+        conversation,
+        session_configured: _,
+    } = conversation_manager.new_conversation(config).await?;
 
     // Task that reads JSON lines from stdin and forwards to Submission Queue
     let sq_fut = {
-        let codex = codex.clone();
-        let ctrl_c = ctrl_c.clone();
+        let conversation = conversation.clone();
         async move {
             let stdin = BufReader::new(tokio::io::stdin());
             let mut lines = stdin.lines();
             loop {
                 let result = tokio::select! {
-                    _ = ctrl_c.notified() => {
+                    _ = conversation.on_abort() => {
                         info!("Interrupted, exiting");
                         break
                     },
@@ -65,7 +64,7 @@ pub async fn run_main(opts: ProtoCli) -> anyhow::Result<()> {
                         }
                         match serde_json::from_str::<Submission>(line) {
                             Ok(sub) => {
-                                if let Err(e) = codex.submit_with_id(sub).await {
+                                if let Err(e) = conversation.submit_with_id(sub).await {
                                     error!("{e:#}");
                                     break;
                                 }
@@ -88,8 +87,8 @@ pub async fn run_main(opts: ProtoCli) -> anyhow::Result<()> {
     let eq_fut = async move {
         loop {
             let event = tokio::select! {
-                _ = ctrl_c.notified() => break,
-                event = codex.next_event() => event,
+                _ = conversation.on_abort() => break,
+                event = conversation.next_event() => event,
             };
             match event {
                 Ok(event) => {

codex-rs/core/src/codex.rs

@@ -126,7 +126,7 @@ impl Codex {
     pub async fn spawn(
         config: Config,
         auth: Option<CodexAuth>,
-        ctrl_c: Arc<Notify>,
+        on_abort: Arc<Notify>,
     ) -> CodexResult<CodexSpawnOk> {
         // experimental resume path (undocumented)
         let resume_path = config.experimental_resume.clone();
@@ -156,7 +156,7 @@ impl Codex {
         // Generate a unique ID for the lifetime of this Codex session.
         let session_id = Uuid::new_v4();
         tokio::spawn(submission_loop(
-            session_id, config, auth, rx_sub, tx_event, ctrl_c,
+            session_id, config, auth, rx_sub, tx_event, on_abort,
         ));
         let codex = Codex {
             next_id: AtomicU64::new(0),
@@ -577,7 +577,7 @@ impl Session {
             .await
     }
 
-    pub fn abort(&self) {
+    fn abort(&self) {
         info!("Aborting existing session");
         let mut state = self.state.lock().unwrap();
         state.pending_approvals.clear();
@@ -708,7 +708,7 @@ async fn submission_loop(
     auth: Option<CodexAuth>,
     rx_sub: Receiver<Submission>,
     tx_event: Sender<Event>,
-    ctrl_c: Arc<Notify>,
+    on_abort: Arc<Notify>,
 ) {
     let mut sess: Option<Arc<Session>> = None;
     // shorthand - send an event when there is no active session
@@ -724,13 +724,13 @@ async fn submission_loop(
     };
 
     loop {
-        let interrupted = ctrl_c.notified();
+        let aborted = on_abort.notified();
         let sub = tokio::select! {
             res = rx_sub.recv() => match res {
                 Ok(sub) => sub,
                 Err(_) => break,
             },
-            _ = interrupted => {
+            _ = aborted => {
                 if let Some(sess) = sess.as_ref(){
                     sess.abort();
                 }
@@ -876,7 +876,7 @@ async fn submission_loop(
                         config.include_plan_tool,
                     ),
                     tx_event: tx_event.clone(),
-                    ctrl_c: Arc::clone(&ctrl_c),
+                    ctrl_c: Arc::clone(&on_abort),
                     user_instructions,
                     base_instructions,
                     approval_policy,

codex-rs/core/src/codex_conversation.rs

@@ -0,0 +1,42 @@
+use std::sync::Arc;
+
+use tokio::sync::Notify;
+use tokio::sync::futures::Notified;
+
+use crate::codex::Codex;
+use crate::error::Result as CodexResult;
+use crate::protocol::Event;
+use crate::protocol::Op;
+use crate::protocol::Submission;
+
+pub struct CodexConversation {
+    codex: Codex,
+    on_abort: Arc<Notify>,
+}
+
+impl CodexConversation {
+    pub(crate) fn new(codex: Codex, on_abort: Arc<Notify>) -> Self {
+        Self { codex, on_abort }
+    }
+
+    pub async fn submit(&self, op: Op) -> CodexResult<String> {
+        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<Event> {
+        self.codex.next_event().await
+    }
+
+    pub fn abort(&self) {
+        self.on_abort.notify_waiters();
+    }
+
+    /// await this to get notified when the user _aborts_ the conversation.
+    pub fn on_abort(&self) -> Notified {
+        self.on_abort.notified()
+    }
+}

codex-rs/core/src/codex_wrapper.rs

@@ -0,0 +1,99 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use codex_login::CodexAuth;
+use tokio::sync::Notify;
+use tokio::sync::RwLock;
+use uuid::Uuid;
+
+use crate::codex::Codex;
+use crate::codex::CodexSpawnOk;
+use crate::codex_conversation::CodexConversation;
+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::SessionConfiguredEvent;
+
+/// Represents a newly created Codex conversation, including the first event
+/// (which is [`EventMsg::SessionConfigured`]).
+pub struct NewConversation {
+    pub conversation_id: Uuid,
+    pub conversation: Arc<CodexConversation>,
+    pub session_configured: SessionConfiguredEvent,
+}
+
+/// [`ConversationManager`] is responsible for creating conversations and
+/// maintaining them in memory.
+pub struct ConversationManager {
+    conversations: Arc<RwLock<HashMap<Uuid, Arc<CodexConversation>>>>,
+}
+
+impl Default for ConversationManager {
+    fn default() -> Self {
+        Self {
+            conversations: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+}
+
+impl ConversationManager {
+    pub async fn new_conversation(&self, config: Config) -> CodexResult<NewConversation> {
+        let auth = CodexAuth::from_codex_home(&config.codex_home)?;
+        self.new_conversation_with_auth(config, auth).await
+    }
+
+    /// Used for integration tests: should not be used by ordinary business
+    /// logic.
+    pub async fn new_conversation_with_auth(
+        &self,
+        config: Config,
+        auth: Option<CodexAuth>,
+    ) -> CodexResult<NewConversation> {
+        let on_abort = Arc::new(Notify::new());
+
+        let CodexSpawnOk {
+            codex,
+            init_id,
+            session_id: conversation_id,
+        } = Codex::spawn(config, auth, on_abort.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(CodexErr::SessionConfiguredNotFirstEvent);
+            }
+        };
+
+        let conversation = Arc::new(CodexConversation::new(codex, on_abort));
+        self.conversations
+            .write()
+            .await
+            .insert(conversation_id, conversation.clone());
+
+        Ok(NewConversation {
+            conversation_id,
+            conversation,
+            session_configured,
+        })
+    }
+
+    pub async fn get_conversation(
+        &self,
+        conversation_id: Uuid,
+    ) -> CodexResult<Arc<CodexConversation>> {
+        let conversations = self.conversations.read().await;
+        conversations
+            .get(&conversation_id)
+            .cloned()
+            .ok_or_else(|| CodexErr::ConversationNotFound(conversation_id))
+    }
+}

codex-rs/core/src/conversation_manager.rs

@@ -3,6 +3,7 @@ use serde_json;
 use std::io;
 use thiserror::Error;
 use tokio::task::JoinError;
+use uuid::Uuid;
 
 pub type Result<T> = std::result::Result<T, CodexErr>;
 
@@ -44,6 +45,12 @@ pub enum CodexErr {
     #[error("stream disconnected before completion: {0}")]
     Stream(String),
 
+    #[error("no conversation with id: {0}")]
+    ConversationNotFound(Uuid),
+
+    #[error("session configured event was not the first event in the stream")]
+    SessionConfiguredNotFirstEvent,
+
     /// Returned by run_command_stream when the spawned child process timed out (10s).
     #[error("timeout waiting for child process to exit")]
     Timeout,

codex-rs/core/src/error.rs

@@ -11,9 +11,8 @@ mod chat_completions;
 mod client;
 mod client_common;
 pub mod codex;
-pub use codex::Codex;
-pub use codex::CodexSpawnOk;
-pub mod codex_wrapper;
+mod codex_conversation;
+pub use codex_conversation::CodexConversation;
 pub mod config;
 pub mod config_profile;
 pub mod config_types;
@@ -34,6 +33,9 @@ pub use model_provider_info::ModelProviderInfo;
 pub use model_provider_info::WireApi;
 pub use model_provider_info::built_in_model_providers;
 pub use model_provider_info::create_oss_provider_with_base_url;
+mod conversation_manager;
+pub use conversation_manager::ConversationManager;
+pub use conversation_manager::NewConversation;
 pub mod model_family;
 mod models;
 mod openai_model_info;