valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: introduce ConversationManager as a clearinghouse for all conversations (#2240)

Browse Source 
This PR does two things because after I got deep into the first one I
started pulling on the thread to the second:

- Makes `ConversationManager` the place where all in-memory
conversations are created and stored. Previously, `MessageProcessor` in
the `codex-mcp-server` crate was doing this via its `session_map`, but
this is something that should be done in `codex-core`.
- It unwinds the `ctrl_c: tokio::sync::Notify` that was threaded
throughout our code. I think this made sense at one time, but now that
we handle Ctrl-C within the TUI and have a proper `Op::Interrupt` event,
I don't think this was quite right, so I removed it. For `codex exec`
and `codex proto`, we now use `tokio::signal::ctrl_c()` directly, but we
no longer make `Notify` a field of `Codex` or `CodexConversation`.

Changes of note:

- Adds the files `conversation_manager.rs` and `codex_conversation.rs`
to `codex-core`.
- `Codex` and `CodexSpawnOk` are no longer exported from `codex-core`:
other crates must use `CodexConversation` instead (which is created via
`ConversationManager`).
- `core/src/codex_wrapper.rs` has been deleted in favor of
`ConversationManager`.
- `ConversationManager::new_conversation()` returns `NewConversation`,
which is in line with the `new_conversation` tool we want to add to the
MCP server. Note `NewConversation` includes `SessionConfiguredEvent`, so
we eliminate checks in cases like `codex-rs/core/tests/client.rs` to
verify `SessionConfiguredEvent` is the first event because that is now
internal to `ConversationManager`.
- Quite a bit of code was deleted from
`codex-rs/mcp-server/src/message_processor.rs` since it no longer has to
manage multiple conversations itself: it goes through
`ConversationManager` instead.
- `core/tests/live_agent.rs` has been deleted because I had to update a
bunch of tests and all the tests in here were ignored, and I don't think
anyone ever ran them, so this was just technical debt, at this point.
- Removed `notify_on_sigint()` from `util.rs` (and in a follow-up, I
hope to refactor the blandly-named `util.rs` into more descriptive
files).
- In general, I started replacing local variables named `codex` as
`conversation`, where appropriate, though admittedly I didn't do it
through all the integration tests because that would have added a lot of
noise to this PR.




---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/openai/codex/pull/2240).
* #2264
* #2263
* __->__ #2240
This commit is contained in:
Michael Bolin
2025-08-13 13:38:18 -07:00
committed by GitHub
parent 30ee24521b
commit 08ed618f72
32 changed files with 406 additions and 654 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
codex-rs/core/src/codex.rs
View File

@@ -20,7 +20,6 @@ use futures::prelude::*;
use mcp_types::CallToolResult;
use serde::Serialize;
use serde_json;
use tokio::sync::Notify;
use tokio::sync::oneshot;
use tokio::task::AbortHandle;
use tracing::debug;
@@ -124,11 +123,7 @@ pub struct CodexSpawnOk {
impl Codex {
/// Spawn a new [`Codex`] and initialize the session.
pub async fn spawn(
config: Config,
auth: Option<CodexAuth>,
ctrl_c: Arc<Notify>,
) -> CodexResult<CodexSpawnOk> {
pub async fn spawn(config: Config, auth: Option<CodexAuth>) -> CodexResult<CodexSpawnOk> {
// experimental resume path (undocumented)
let resume_path = config.experimental_resume.clone();
info!("resume_path: {resume_path:?}");
@@ -156,9 +151,9 @@ 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,
));
// This task will run until Op::Shutdown is received.
tokio::spawn(submission_loop(session_id, config, auth, rx_sub, tx_event));
let codex = Codex {
next_id: AtomicU64::new(0),
tx_sub,
@@ -210,7 +205,6 @@ impl Codex {
pub(crate) struct Session {
client: ModelClient,
pub(crate) tx_event: Sender<Event>,
ctrl_c: Arc<Notify>,
/// The session's current working directory. All relative paths provided by
/// the model as well as sandbox policies are resolved against this path
@@ -493,7 +487,6 @@ impl Session {
let result = process_exec_tool_call(
exec_args.params,
exec_args.sandbox_type,
exec_args.ctrl_c,
exec_args.sandbox_policy,
exec_args.codex_linux_sandbox_exe,
exec_args.stdout_stream,
@@ -578,7 +571,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();
@@ -709,7 +702,6 @@ async fn submission_loop(
auth: Option<CodexAuth>,
rx_sub: Receiver<Submission>,
tx_event: Sender<Event>,
ctrl_c: Arc<Notify>,
) {
let mut sess: Option<Arc<Session>> = None;
// shorthand - send an event when there is no active session
@@ -724,21 +716,8 @@ async fn submission_loop(
tx_event.send(event).await.ok();
};
loop {
let interrupted = ctrl_c.notified();
let sub = tokio::select! {
res = rx_sub.recv() => match res {
Ok(sub) => sub,
Err(_) => break,
},
_ = interrupted => {
if let Some(sess) = sess.as_ref(){
sess.abort();
}
continue;
},
};
// To break out of this loop, send Op::Shutdown.
while let Ok(sub) = rx_sub.recv().await {
debug!(?sub, "Submission");
match sub.op {
Op::Interrupt => {
@@ -877,7 +856,6 @@ async fn submission_loop(
config.include_plan_tool,
),
tx_event: tx_event.clone(),
ctrl_c: Arc::clone(&ctrl_c),
user_instructions,
base_instructions,
approval_policy,
@@ -1787,7 +1765,6 @@ fn parse_container_exec_arguments(
pub struct ExecInvokeArgs<'a> {
pub params: ExecParams,
pub sandbox_type: SandboxType,
pub ctrl_c: Arc<Notify>,
pub sandbox_policy: &'a SandboxPolicy,
pub codex_linux_sandbox_exe: &'a Option<PathBuf>,
pub stdout_stream: Option<StdoutStream>,
@@ -1972,7 +1949,6 @@ async fn handle_container_exec_with_params(
ExecInvokeArgs {
params: params.clone(),
sandbox_type,
ctrl_c: sess.ctrl_c.clone(),
sandbox_policy: &sess.sandbox_policy,
codex_linux_sandbox_exe: &sess.codex_linux_sandbox_exe,
stdout_stream: Some(StdoutStream {
@@ -2104,7 +2080,6 @@ async fn handle_sandbox_error(
ExecInvokeArgs {
params,
sandbox_type: SandboxType::None,
ctrl_c: sess.ctrl_c.clone(),
sandbox_policy: &sess.sandbox_policy,
codex_linux_sandbox_exe: &sess.codex_linux_sandbox_exe,
stdout_stream: Some(StdoutStream {

30
codex-rs/core/src/codex_conversation.rs Normal file
View File

@@ -0,0 +1,30 @@
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,
}
/// Conduit for the bidirectional stream of messages that compose a conversation
/// in Codex.
impl CodexConversation {
pub(crate) fn new(codex: Codex) -> Self {
Self { codex }
}
pub async fn submit(&self, op: Op) -> CodexResult<String> {
self.codex.submit(op).await
}
/// Use sparingly: this is intended to be removed soon.
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
}
}

59
codex-rs/core/src/codex_wrapper.rs
View File

@@ -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<Notify>,
}
/// 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<CodexConversation> {
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,
})
}

96
codex-rs/core/src/conversation_manager.rs Normal file
View File

@@ -0,0 +1,96 @@
use std::collections::HashMap;
use std::sync::Arc;
use codex_login::CodexAuth;
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 CodexSpawnOk {
codex,
init_id,
session_id: conversation_id,
} = Codex::spawn(config, auth).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));
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))
}
}

7
codex-rs/core/src/error.rs
View File

@@ -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,

16
codex-rs/core/src/exec.rs
View File

@@ -6,7 +6,6 @@ use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process::ExitStatus;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;
@@ -15,7 +14,6 @@ use tokio::io::AsyncRead;
use tokio::io::AsyncReadExt;
use tokio::io::BufReader;
use tokio::process::Child;
use tokio::sync::Notify;
use crate::error::CodexErr;
use crate::error::Result;
@@ -80,7 +78,6 @@ pub struct StdoutStream {
pub async fn process_exec_tool_call(
params: ExecParams,
sandbox_type: SandboxType,
ctrl_c: Arc<Notify>,
sandbox_policy: &SandboxPolicy,
codex_linux_sandbox_exe: &Option<PathBuf>,
stdout_stream: Option<StdoutStream>,
@@ -89,7 +86,7 @@ pub async fn process_exec_tool_call(
let raw_output_result: std::result::Result<RawExecToolCallOutput, CodexErr> = match sandbox_type
{
SandboxType::None => exec(params, sandbox_policy, ctrl_c, stdout_stream.clone()).await,
SandboxType::None => exec(params, sandbox_policy, stdout_stream.clone()).await,
SandboxType::MacosSeatbelt => {
let timeout = params.timeout_duration();
let ExecParams {
@@ -103,7 +100,7 @@ pub async fn process_exec_tool_call(
env,
)
.await?;
consume_truncated_output(child, ctrl_c, timeout, stdout_stream.clone()).await
consume_truncated_output(child, timeout, stdout_stream.clone()).await
}
SandboxType::LinuxSeccomp => {
let timeout = params.timeout_duration();
@@ -124,7 +121,7 @@ pub async fn process_exec_tool_call(
)
.await?;
consume_truncated_output(child, ctrl_c, timeout, stdout_stream).await
consume_truncated_output(child, timeout, stdout_stream).await
}
};
let duration = start.elapsed();
@@ -286,7 +283,6 @@ pub struct ExecToolCallOutput {
async fn exec(
params: ExecParams,
sandbox_policy: &SandboxPolicy,
ctrl_c: Arc<Notify>,
stdout_stream: Option<StdoutStream>,
) -> Result<RawExecToolCallOutput> {
let timeout = params.timeout_duration();
@@ -311,14 +307,13 @@ async fn exec(
env,
)
.await?;
consume_truncated_output(child, ctrl_c, timeout, stdout_stream).await
consume_truncated_output(child, timeout, stdout_stream).await
}
/// Consumes the output of a child process, truncating it so it is suitable for
/// use as the output of a `shell` tool call. Also enforces specified timeout.
pub(crate) async fn consume_truncated_output(
mut child: Child,
ctrl_c: Arc<Notify>,
timeout: Duration,
stdout_stream: Option<StdoutStream>,
) -> Result<RawExecToolCallOutput> {
@@ -352,7 +347,6 @@ pub(crate) async fn consume_truncated_output(
true,
));
let interrupted = ctrl_c.notified();
let exit_status = tokio::select! {
result = tokio::time::timeout(timeout, child.wait()) => {
match result {
@@ -366,7 +360,7 @@ pub(crate) async fn consume_truncated_output(
}
}
}
_ = interrupted => {
_ = tokio::signal::ctrl_c() => {
child.start_kill()?;
synthetic_exit_status(128 + SIGKILL_CODE)
}

8
codex-rs/core/src/lib.rs
View File

@@ -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;

4
codex-rs/core/src/shell.rs
View File

@@ -167,9 +167,6 @@ mod tests {
for (input, expected_cmd, expected_output) in cases {
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Notify;
use crate::exec::ExecParams;
use crate::exec::SandboxType;
@@ -219,7 +216,6 @@ mod tests {
justification: None,
},
SandboxType::None,
Arc::new(Notify::new()),
&SandboxPolicy::DangerFullAccess,
&None,
None,

21
codex-rs/core/src/util.rs
View File

@@ -1,32 +1,11 @@
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use rand::Rng;
use tokio::sync::Notify;
use tracing::debug;
const INITIAL_DELAY_MS: u64 = 200;
const BACKOFF_FACTOR: f64 = 2.0;
/// Make a CancellationToken that is fulfilled when SIGINT occurs.
pub fn notify_on_sigint() -> Arc<Notify> {
let notify = Arc::new(Notify::new());
tokio::spawn({
let notify = Arc::clone(&notify);
async move {
loop {
tokio::signal::ctrl_c().await.ok();
debug!("Keyboard interrupt");
notify.notify_waiters();
}
}
});
notify
}
pub(crate) fn backoff(attempt: u64) -> Duration {
let exp = BACKOFF_FACTOR.powi(attempt.saturating_sub(1) as i32);
let base = (INITIAL_DELAY_MS as f64 * exp) as u64;