diff --git a/.github/workflows/rust-ci.yml b/.github/workflows/rust-ci.yml index bae33e30d9..280939c611 100644 --- a/.github/workflows/rust-ci.yml +++ b/.github/workflows/rust-ci.yml @@ -62,6 +62,8 @@ jobs: components: rustfmt - name: cargo fmt run: cargo fmt -- --config imports_granularity=Item --check + - name: Verify codegen for mcp-types + run: ./mcp-types/check_lib_rs.py cargo_shear: name: cargo shear diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index b706afb2f0..69ef7bc449 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -54,6 +54,7 @@ tracing = { version = "0.1.41", features = ["log"] } tree-sitter = "0.25.9" tree-sitter-bash = "0.25.0" uuid = { version = "1", features = ["serde", "v4"] } +which = "6" wildmatch = "2.4.0" @@ -69,9 +70,6 @@ openssl-sys = { version = "*", features = ["vendored"] } [target.aarch64-unknown-linux-musl.dependencies] openssl-sys = { version = "*", features = ["vendored"] } -[target.'cfg(target_os = "windows")'.dependencies] -which = "6" - [dev-dependencies] assert_cmd = "2" core_test_support = { path = "tests/common" } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index f3e9800c10..46d155be74 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -19,13 +19,14 @@ use codex_apply_patch::ApplyPatchAction; use codex_apply_patch::MaybeApplyPatchVerified; use codex_apply_patch::maybe_parse_apply_patch_verified; use codex_protocol::mcp_protocol::ConversationId; -use codex_protocol::protocol::ConversationHistoryResponseEvent; +use codex_protocol::protocol::ConversationPathResponseEvent; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::TaskStartedEvent; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use futures::prelude::*; use mcp_types::CallToolResult; +use serde::Deserialize; use serde::Serialize; use serde_json; use tokio::sync::oneshot; @@ -112,6 +113,7 @@ use crate::safety::assess_command_safety; use crate::safety::assess_safety_for_untrusted_command; use crate::shell; use crate::turn_diff_tracker::TurnDiffTracker; +use crate::unified_exec::UnifiedExecSessionManager; use crate::user_instructions::UserInstructions; use crate::user_notification::UserNotification; use crate::util::backoff; @@ -280,6 +282,7 @@ pub(crate) struct Session { /// Manager for external MCP servers/tools. mcp_connection_manager: McpConnectionManager, session_manager: ExecSessionManager, + unified_exec_manager: UnifiedExecSessionManager, /// External notifier command (will be passed as args to exec()). When /// `None` this feature is disabled. @@ -471,6 +474,7 @@ impl Session { include_web_search_request: config.tools_web_search_request, use_streamable_shell_tool: config.use_experimental_streamable_shell_tool, include_view_image_tool: config.include_view_image_tool, + experimental_unified_exec_tool: config.use_experimental_unified_exec_tool, }), user_instructions, base_instructions, @@ -484,6 +488,7 @@ impl Session { tx_event: tx_event.clone(), mcp_connection_manager, session_manager: ExecSessionManager::default(), + unified_exec_manager: UnifiedExecSessionManager::default(), notify, state: Mutex::new(state), rollout: Mutex::new(Some(rollout_recorder)), @@ -1149,6 +1154,7 @@ async fn submission_loop( include_web_search_request: config.tools_web_search_request, use_streamable_shell_tool: config.use_experimental_streamable_shell_tool, include_view_image_tool: config.include_view_image_tool, + experimental_unified_exec_tool: config.use_experimental_unified_exec_tool, }); let new_turn_context = TurnContext { @@ -1251,6 +1257,8 @@ async fn submission_loop( use_streamable_shell_tool: config .use_experimental_streamable_shell_tool, include_view_image_tool: config.include_view_image_tool, + experimental_unified_exec_tool: config + .use_experimental_unified_exec_tool, }), user_instructions: turn_context.user_instructions.clone(), base_instructions: turn_context.base_instructions.clone(), @@ -1397,14 +1405,29 @@ async fn submission_loop( sess.send_event(event).await; break; } - Op::GetHistory => { + Op::GetPath => { let sub_id = sub.id.clone(); - + // Flush rollout writes before returning the path so readers observe a consistent file. + let (path, rec_opt) = { + let guard = sess.rollout.lock_unchecked(); + match guard.as_ref() { + Some(rec) => (rec.get_rollout_path(), Some(rec.clone())), + None => { + error!("rollout recorder not found"); + continue; + } + } + }; + if let Some(rec) = rec_opt + && let Err(e) = rec.flush().await + { + warn!("failed to flush rollout recorder before GetHistory: {e}"); + } let event = Event { id: sub_id.clone(), - msg: EventMsg::ConversationHistory(ConversationHistoryResponseEvent { + msg: EventMsg::ConversationPath(ConversationPathResponseEvent { conversation_id: sess.conversation_id, - entries: sess.state.lock_unchecked().history.contents(), + path, }), }; sess.send_event(event).await; @@ -2082,6 +2105,72 @@ async fn handle_response_item( Ok(output) } +async fn handle_unified_exec_tool_call( + sess: &Session, + call_id: String, + session_id: Option, + arguments: Vec, + timeout_ms: Option, +) -> ResponseInputItem { + let parsed_session_id = if let Some(session_id) = session_id { + match session_id.parse::() { + Ok(parsed) => Some(parsed), + Err(output) => { + return ResponseInputItem::FunctionCallOutput { + call_id: call_id.to_string(), + output: FunctionCallOutputPayload { + content: format!("invalid session_id: {session_id} due to error {output}"), + success: Some(false), + }, + }; + } + } + } else { + None + }; + + let request = crate::unified_exec::UnifiedExecRequest { + session_id: parsed_session_id, + input_chunks: &arguments, + timeout_ms, + }; + + let result = sess.unified_exec_manager.handle_request(request).await; + + let output_payload = match result { + Ok(value) => { + #[derive(Serialize)] + struct SerializedUnifiedExecResult<'a> { + session_id: Option, + output: &'a str, + } + + match serde_json::to_string(&SerializedUnifiedExecResult { + session_id: value.session_id.map(|id| id.to_string()), + output: &value.output, + }) { + Ok(serialized) => FunctionCallOutputPayload { + content: serialized, + success: Some(true), + }, + Err(err) => FunctionCallOutputPayload { + content: format!("failed to serialize unified exec output: {err}"), + success: Some(false), + }, + } + } + Err(err) => FunctionCallOutputPayload { + content: format!("unified exec failed: {err}"), + success: Some(false), + }, + }; + + ResponseInputItem::FunctionCallOutput { + call_id, + output: output_payload, + } +} + async fn handle_function_call( sess: &Session, turn_context: &TurnContext, @@ -2109,6 +2198,38 @@ async fn handle_function_call( ) .await } + "unified_exec" => { + #[derive(Deserialize)] + struct UnifiedExecArgs { + input: Vec, + #[serde(default)] + session_id: Option, + #[serde(default)] + timeout_ms: Option, + } + + let args = match serde_json::from_str::(&arguments) { + Ok(args) => args, + Err(err) => { + return ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: format!("failed to parse function arguments: {err}"), + success: Some(false), + }, + }; + } + }; + + handle_unified_exec_tool_call( + sess, + call_id, + args.session_id, + args.input, + args.timeout_ms, + ) + .await + } "view_image" => { #[derive(serde::Deserialize)] struct SeeImageArgs { diff --git a/codex-rs/core/src/config.rs b/codex-rs/core/src/config.rs index 8ae9931f22..316f7276c0 100644 --- a/codex-rs/core/src/config.rs +++ b/codex-rs/core/src/config.rs @@ -172,6 +172,9 @@ pub struct Config { pub use_experimental_streamable_shell_tool: bool, + /// If set to `true`, used only the experimental unified exec tool. + pub use_experimental_unified_exec_tool: bool, + /// Include the `view_image` tool that lets the agent attach a local image path to context. pub include_view_image_tool: bool, @@ -487,6 +490,7 @@ pub struct ConfigToml { pub experimental_instructions_file: Option, pub experimental_use_exec_command_tool: Option, + pub experimental_use_unified_exec_tool: Option, pub projects: Option>, @@ -837,6 +841,9 @@ impl Config { use_experimental_streamable_shell_tool: cfg .experimental_use_exec_command_tool .unwrap_or(false), + use_experimental_unified_exec_tool: cfg + .experimental_use_unified_exec_tool + .unwrap_or(true), include_view_image_tool, active_profile: active_profile_name, disable_paste_burst: cfg.disable_paste_burst.unwrap_or(false), @@ -1212,6 +1219,7 @@ model_verbosity = "high" tools_web_search_request: false, preferred_auth_method: AuthMode::ChatGPT, use_experimental_streamable_shell_tool: false, + use_experimental_unified_exec_tool: true, include_view_image_tool: true, active_profile: Some("o3".to_string()), disable_paste_burst: false, @@ -1269,6 +1277,7 @@ model_verbosity = "high" tools_web_search_request: false, preferred_auth_method: AuthMode::ChatGPT, use_experimental_streamable_shell_tool: false, + use_experimental_unified_exec_tool: true, include_view_image_tool: true, active_profile: Some("gpt3".to_string()), disable_paste_burst: false, @@ -1341,6 +1350,7 @@ model_verbosity = "high" tools_web_search_request: false, preferred_auth_method: AuthMode::ChatGPT, use_experimental_streamable_shell_tool: false, + use_experimental_unified_exec_tool: true, include_view_image_tool: true, active_profile: Some("zdr".to_string()), disable_paste_burst: false, @@ -1399,6 +1409,7 @@ model_verbosity = "high" tools_web_search_request: false, preferred_auth_method: AuthMode::ChatGPT, use_experimental_streamable_shell_tool: false, + use_experimental_unified_exec_tool: true, include_view_image_tool: true, active_profile: Some("gpt5".to_string()), disable_paste_burst: false, diff --git a/codex-rs/core/src/conversation_manager.rs b/codex-rs/core/src/conversation_manager.rs index c0695f16f3..a5f66d4553 100644 --- a/codex-rs/core/src/conversation_manager.rs +++ b/codex-rs/core/src/conversation_manager.rs @@ -150,13 +150,13 @@ impl ConversationManager { /// caller's `config`). The new conversation will have a fresh id. pub async fn fork_conversation( &self, - conversation_history: Vec, num_messages_to_drop: usize, config: Config, + path: PathBuf, ) -> CodexResult { // Compute the prefix up to the cut point. - let history = - truncate_after_dropping_last_messages(conversation_history, num_messages_to_drop); + let history = RolloutRecorder::get_rollout_history(&path).await?; + let history = truncate_after_dropping_last_messages(history, num_messages_to_drop); // Spawn a new conversation with the computed initial history. let auth_manager = self.auth_manager.clone(); @@ -171,36 +171,36 @@ impl ConversationManager { /// Return a prefix of `items` obtained by dropping the last `n` user messages /// and all items that follow them. -fn truncate_after_dropping_last_messages(items: Vec, n: usize) -> InitialHistory { +fn truncate_after_dropping_last_messages(history: InitialHistory, n: usize) -> InitialHistory { if n == 0 { - let rolled: Vec = items.into_iter().map(RolloutItem::ResponseItem).collect(); - return InitialHistory::Forked(rolled); + return InitialHistory::Forked(history.get_rollout_items()); } - // Walk backwards counting only `user` Message items, find cut index. - let mut count = 0usize; - let mut cut_index = 0usize; - for (idx, item) in items.iter().enumerate().rev() { - if let ResponseItem::Message { role, .. } = item + // Work directly on rollout items, and cut the vector at the nth-from-last user message input. + let items: Vec = history.get_rollout_items(); + + // Find indices of user message inputs in rollout order. + let mut user_positions: Vec = Vec::new(); + for (idx, item) in items.iter().enumerate() { + if let RolloutItem::ResponseItem(ResponseItem::Message { role, .. }) = item && role == "user" { - count += 1; - if count == n { - // Cut everything from this user message to the end. - cut_index = idx; - break; - } + user_positions.push(idx); } } - if cut_index == 0 { - // No prefix remains after dropping; start a new conversation. + + // If fewer than n user messages exist, treat as empty. + if user_positions.len() < n { + return InitialHistory::New; + } + + // Cut strictly before the nth-from-last user message (do not keep the nth itself). + let cut_idx = user_positions[user_positions.len() - n]; + let rolled: Vec = items.into_iter().take(cut_idx).collect(); + + if rolled.is_empty() { InitialHistory::New } else { - let rolled: Vec = items - .into_iter() - .take(cut_index) - .map(RolloutItem::ResponseItem) - .collect(); InitialHistory::Forked(rolled) } } @@ -256,7 +256,13 @@ mod tests { assistant_msg("a4"), ]; - let truncated = truncate_after_dropping_last_messages(items.clone(), 1); + // Wrap as InitialHistory::Forked with response items only. + let initial: Vec = items + .iter() + .cloned() + .map(RolloutItem::ResponseItem) + .collect(); + let truncated = truncate_after_dropping_last_messages(InitialHistory::Forked(initial), 1); let got_items = truncated.get_rollout_items(); let expected_items = vec![ RolloutItem::ResponseItem(items[0].clone()), @@ -268,7 +274,12 @@ mod tests { serde_json::to_value(&expected_items).unwrap() ); - let truncated2 = truncate_after_dropping_last_messages(items, 2); + let initial2: Vec = items + .iter() + .cloned() + .map(RolloutItem::ResponseItem) + .collect(); + let truncated2 = truncate_after_dropping_last_messages(InitialHistory::Forked(initial2), 2); assert!(matches!(truncated2, InitialHistory::New)); } } diff --git a/codex-rs/core/src/exec_command/exec_command_session.rs b/codex-rs/core/src/exec_command/exec_command_session.rs index 7503150c9a..b524506eaf 100644 --- a/codex-rs/core/src/exec_command/exec_command_session.rs +++ b/codex-rs/core/src/exec_command/exec_command_session.rs @@ -24,6 +24,9 @@ pub(crate) struct ExecCommandSession { /// JoinHandle for the child wait task. wait_handle: StdMutex>>, + + /// Tracks whether the underlying process has exited. + exit_status: std::sync::Arc, } impl ExecCommandSession { @@ -34,6 +37,7 @@ impl ExecCommandSession { reader_handle: JoinHandle<()>, writer_handle: JoinHandle<()>, wait_handle: JoinHandle<()>, + exit_status: std::sync::Arc, ) -> Self { Self { writer_tx, @@ -42,6 +46,7 @@ impl ExecCommandSession { reader_handle: StdMutex::new(Some(reader_handle)), writer_handle: StdMutex::new(Some(writer_handle)), wait_handle: StdMutex::new(Some(wait_handle)), + exit_status, } } @@ -52,6 +57,10 @@ impl ExecCommandSession { pub(crate) fn output_receiver(&self) -> broadcast::Receiver> { self.output_tx.subscribe() } + + pub(crate) fn has_exited(&self) -> bool { + self.exit_status.load(std::sync::atomic::Ordering::SeqCst) + } } impl Drop for ExecCommandSession { diff --git a/codex-rs/core/src/exec_command/mod.rs b/codex-rs/core/src/exec_command/mod.rs index 9cdaa4d31d..103b76ad67 100644 --- a/codex-rs/core/src/exec_command/mod.rs +++ b/codex-rs/core/src/exec_command/mod.rs @@ -6,6 +6,7 @@ mod session_manager; pub use exec_command_params::ExecCommandParams; pub use exec_command_params::WriteStdinParams; +pub(crate) use exec_command_session::ExecCommandSession; pub use responses_api::EXEC_COMMAND_TOOL_NAME; pub use responses_api::WRITE_STDIN_TOOL_NAME; pub use responses_api::create_exec_command_tool_for_responses_api; diff --git a/codex-rs/core/src/exec_command/session_manager.rs b/codex-rs/core/src/exec_command/session_manager.rs index c547409cd1..9578610cc7 100644 --- a/codex-rs/core/src/exec_command/session_manager.rs +++ b/codex-rs/core/src/exec_command/session_manager.rs @@ -3,6 +3,7 @@ use std::io::ErrorKind; use std::io::Read; use std::sync::Arc; use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU32; use portable_pty::CommandBuilder; @@ -19,6 +20,7 @@ use crate::exec_command::exec_command_params::ExecCommandParams; use crate::exec_command::exec_command_params::WriteStdinParams; use crate::exec_command::exec_command_session::ExecCommandSession; use crate::exec_command::session_id::SessionId; +use crate::truncate::truncate_middle; use codex_protocol::models::FunctionCallOutputPayload; #[derive(Debug, Default)] @@ -327,11 +329,14 @@ async fn create_exec_command_session( // Keep the child alive until it exits, then signal exit code. let (exit_tx, exit_rx) = oneshot::channel::(); + let exit_status = Arc::new(AtomicBool::new(false)); + let wait_exit_status = exit_status.clone(); let wait_handle = tokio::task::spawn_blocking(move || { let code = match child.wait() { Ok(status) => status.exit_code() as i32, Err(_) => -1, }; + wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst); let _ = exit_tx.send(code); }); @@ -343,116 +348,11 @@ async fn create_exec_command_session( reader_handle, writer_handle, wait_handle, + exit_status, ); Ok((session, exit_rx)) } -/// Truncate the middle of a UTF-8 string to at most `max_bytes` bytes, -/// preserving the beginning and the end. Returns the possibly truncated -/// string and `Some(original_token_count)` (estimated at 4 bytes/token) -/// if truncation occurred; otherwise returns the original string and `None`. -fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option) { - // No truncation needed - if s.len() <= max_bytes { - return (s.to_string(), None); - } - let est_tokens = (s.len() as u64).div_ceil(4); - if max_bytes == 0 { - // Cannot keep any content; still return a full marker (never truncated). - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); - } - - // Helper to truncate a string to a given byte length on a char boundary. - fn truncate_on_boundary(input: &str, max_len: usize) -> &str { - if input.len() <= max_len { - return input; - } - let mut end = max_len; - while end > 0 && !input.is_char_boundary(end) { - end -= 1; - } - &input[..end] - } - - // Given a left/right budget, prefer newline boundaries; otherwise fall back - // to UTF-8 char boundaries. - fn pick_prefix_end(s: &str, left_budget: usize) -> usize { - if let Some(head) = s.get(..left_budget) - && let Some(i) = head.rfind('\n') - { - return i + 1; // keep the newline so suffix starts on a fresh line - } - truncate_on_boundary(s, left_budget).len() - } - - fn pick_suffix_start(s: &str, right_budget: usize) -> usize { - let start_tail = s.len().saturating_sub(right_budget); - if let Some(tail) = s.get(start_tail..) - && let Some(i) = tail.find('\n') - { - return start_tail + i + 1; // start after newline - } - // Fall back to a char boundary at or after start_tail. - let mut idx = start_tail.min(s.len()); - while idx < s.len() && !s.is_char_boundary(idx) { - idx += 1; - } - idx - } - - // Refine marker length and budgets until stable. Marker is never truncated. - let mut guess_tokens = est_tokens; // worst-case: everything truncated - for _ in 0..4 { - let marker = format!("…{guess_tokens} tokens truncated…"); - let marker_len = marker.len(); - let keep_budget = max_bytes.saturating_sub(marker_len); - if keep_budget == 0 { - // No room for any content within the cap; return a full, untruncated marker - // that reflects the entire truncated content. - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); - } - - let left_budget = keep_budget / 2; - let right_budget = keep_budget - left_budget; - let prefix_end = pick_prefix_end(s, left_budget); - let mut suffix_start = pick_suffix_start(s, right_budget); - if suffix_start < prefix_end { - suffix_start = prefix_end; - } - let kept_content_bytes = prefix_end + (s.len() - suffix_start); - let truncated_content_bytes = s.len().saturating_sub(kept_content_bytes); - let new_tokens = (truncated_content_bytes as u64).div_ceil(4); - if new_tokens == guess_tokens { - let mut out = String::with_capacity(marker_len + kept_content_bytes + 1); - out.push_str(&s[..prefix_end]); - out.push_str(&marker); - // Place marker on its own line for symmetry when we keep line boundaries. - out.push('\n'); - out.push_str(&s[suffix_start..]); - return (out, Some(est_tokens)); - } - guess_tokens = new_tokens; - } - - // Fallback: use last guess to build output. - let marker = format!("…{guess_tokens} tokens truncated…"); - let marker_len = marker.len(); - let keep_budget = max_bytes.saturating_sub(marker_len); - if keep_budget == 0 { - return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); - } - let left_budget = keep_budget / 2; - let right_budget = keep_budget - left_budget; - let prefix_end = pick_prefix_end(s, left_budget); - let suffix_start = pick_suffix_start(s, right_budget); - let mut out = String::with_capacity(marker_len + prefix_end + (s.len() - suffix_start) + 1); - out.push_str(&s[..prefix_end]); - out.push_str(&marker); - out.push('\n'); - out.push_str(&s[suffix_start..]); - (out, Some(est_tokens)) -} - #[cfg(test)] mod tests { use super::*; @@ -616,50 +516,4 @@ Output: abc"#; assert_eq!(expected, text); } - - #[test] - fn truncate_middle_no_newlines_fallback() { - // A long string with no newlines that exceeds the cap. - let s = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"; - let max_bytes = 16; // force truncation - let (out, original) = truncate_middle(s, max_bytes); - // For very small caps, we return the full, untruncated marker, - // even if it exceeds the cap. - assert_eq!(out, "…16 tokens truncated…"); - // Original string length is 62 bytes => ceil(62/4) = 16 tokens. - assert_eq!(original, Some(16)); - } - - #[test] - fn truncate_middle_prefers_newline_boundaries() { - // Build a multi-line string of 20 numbered lines (each "NNN\n"). - let mut s = String::new(); - for i in 1..=20 { - s.push_str(&format!("{i:03}\n")); - } - // Total length: 20 lines * 4 bytes per line = 80 bytes. - assert_eq!(s.len(), 80); - - // Choose a cap that forces truncation while leaving room for - // a few lines on each side after accounting for the marker. - let max_bytes = 64; - // Expect exact output: first 4 lines, marker, last 4 lines, and correct token estimate (80/4 = 20). - assert_eq!( - truncate_middle(&s, max_bytes), - ( - r#"001 -002 -003 -004 -…12 tokens truncated… -017 -018 -019 -020 -"# - .to_string(), - Some(20) - ) - ); - } } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 77b4037e25..5223b6d327 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -35,6 +35,8 @@ mod mcp_tool_call; mod message_history; mod model_provider_info; pub mod parse_command; +mod truncate; +mod unified_exec; mod user_instructions; pub use model_provider_info::BUILT_IN_OSS_MODEL_PROVIDER_ID; pub use model_provider_info::ModelProviderInfo; diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index 9fc7275583..eecda78887 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -17,7 +17,7 @@ use anyhow::Result; use anyhow::anyhow; use codex_mcp_client::McpClient; use mcp_types::ClientCapabilities; -use mcp_types::McpClientInfo; +use mcp_types::Implementation; use mcp_types::Tool; use serde_json::json; @@ -159,10 +159,14 @@ impl McpConnectionManager { // indicates this should be an empty object. elicitation: Some(json!({})), }, - client_info: McpClientInfo { + client_info: Implementation { name: "codex-mcp-client".to_owned(), version: env!("CARGO_PKG_VERSION").to_owned(), title: Some("Codex".into()), + // This field is used by Codex when it is an MCP + // server: it should not be used when Codex is + // an MCP client. + user_agent: None, }, protocol_version: mcp_types::MCP_SCHEMA_VERSION.to_owned(), }; diff --git a/codex-rs/core/src/openai_tools.rs b/codex-rs/core/src/openai_tools.rs index 9521e4ee93..ec56fdfd2c 100644 --- a/codex-rs/core/src/openai_tools.rs +++ b/codex-rs/core/src/openai_tools.rs @@ -70,6 +70,7 @@ pub(crate) struct ToolsConfig { pub apply_patch_tool_type: Option, pub web_search_request: bool, pub include_view_image_tool: bool, + pub experimental_unified_exec_tool: bool, } pub(crate) struct ToolsConfigParams<'a> { @@ -81,6 +82,7 @@ pub(crate) struct ToolsConfigParams<'a> { pub(crate) include_web_search_request: bool, pub(crate) use_streamable_shell_tool: bool, pub(crate) include_view_image_tool: bool, + pub(crate) experimental_unified_exec_tool: bool, } impl ToolsConfig { @@ -94,6 +96,7 @@ impl ToolsConfig { include_web_search_request, use_streamable_shell_tool, include_view_image_tool, + experimental_unified_exec_tool, } = params; let mut shell_type = if *use_streamable_shell_tool { ConfigShellToolType::StreamableShell @@ -126,6 +129,7 @@ impl ToolsConfig { apply_patch_tool_type, web_search_request: *include_web_search_request, include_view_image_tool: *include_view_image_tool, + experimental_unified_exec_tool: *experimental_unified_exec_tool, } } } @@ -200,6 +204,53 @@ fn create_shell_tool() -> OpenAiTool { }) } +fn create_unified_exec_tool() -> OpenAiTool { + let mut properties = BTreeMap::new(); + properties.insert( + "input".to_string(), + JsonSchema::Array { + items: Box::new(JsonSchema::String { description: None }), + description: Some( + "When no session_id is provided, treat the array as the command and arguments \ + to launch. When session_id is set, concatenate the strings (in order) and write \ + them to the session's stdin." + .to_string(), + ), + }, + ); + properties.insert( + "session_id".to_string(), + JsonSchema::String { + description: Some( + "Identifier for an existing interactive session. If omitted, a new command \ + is spawned." + .to_string(), + ), + }, + ); + properties.insert( + "timeout_ms".to_string(), + JsonSchema::Number { + description: Some( + "Maximum time in milliseconds to wait for output after writing the input." + .to_string(), + ), + }, + ); + + OpenAiTool::Function(ResponsesApiTool { + name: "unified_exec".to_string(), + description: + "Runs a command in a PTY. Provide a session_id to reuse an existing interactive session.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["input".to_string()]), + additional_properties: Some(false), + }, + }) +} + fn create_shell_tool_for_sandbox(sandbox_policy: &SandboxPolicy) -> OpenAiTool { let mut properties = BTreeMap::new(); properties.insert( @@ -531,23 +582,27 @@ pub(crate) fn get_openai_tools( ) -> Vec { let mut tools: Vec = Vec::new(); - match &config.shell_type { - ConfigShellToolType::DefaultShell => { - tools.push(create_shell_tool()); - } - ConfigShellToolType::ShellWithRequest { sandbox_policy } => { - tools.push(create_shell_tool_for_sandbox(sandbox_policy)); - } - ConfigShellToolType::LocalShell => { - tools.push(OpenAiTool::LocalShell {}); - } - ConfigShellToolType::StreamableShell => { - tools.push(OpenAiTool::Function( - crate::exec_command::create_exec_command_tool_for_responses_api(), - )); - tools.push(OpenAiTool::Function( - crate::exec_command::create_write_stdin_tool_for_responses_api(), - )); + if config.experimental_unified_exec_tool { + tools.push(create_unified_exec_tool()); + } else { + match &config.shell_type { + ConfigShellToolType::DefaultShell => { + tools.push(create_shell_tool()); + } + ConfigShellToolType::ShellWithRequest { sandbox_policy } => { + tools.push(create_shell_tool_for_sandbox(sandbox_policy)); + } + ConfigShellToolType::LocalShell => { + tools.push(OpenAiTool::LocalShell {}); + } + ConfigShellToolType::StreamableShell => { + tools.push(OpenAiTool::Function( + crate::exec_command::create_exec_command_tool_for_responses_api(), + )); + tools.push(OpenAiTool::Function( + crate::exec_command::create_write_stdin_tool_for_responses_api(), + )); + } } } @@ -574,10 +629,8 @@ pub(crate) fn get_openai_tools( if config.include_view_image_tool { tools.push(create_view_image_tool()); } - if let Some(mcp_tools) = mcp_tools { // Ensure deterministic ordering to maximize prompt cache hits. - // HashMap iteration order is non-deterministic, so sort by fully-qualified tool name. let mut entries: Vec<(String, mcp_types::Tool)> = mcp_tools.into_iter().collect(); entries.sort_by(|a, b| a.0.cmp(&b.0)); @@ -639,12 +692,13 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, include_view_image_tool: true, + experimental_unified_exec_tool: true, }); let tools = get_openai_tools(&config, Some(HashMap::new())); assert_eq_tool_names( &tools, - &["local_shell", "update_plan", "web_search", "view_image"], + &["unified_exec", "update_plan", "web_search", "view_image"], ); } @@ -660,12 +714,13 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, include_view_image_tool: true, + experimental_unified_exec_tool: true, }); let tools = get_openai_tools(&config, Some(HashMap::new())); assert_eq_tool_names( &tools, - &["shell", "update_plan", "web_search", "view_image"], + &["unified_exec", "update_plan", "web_search", "view_image"], ); } @@ -681,6 +736,7 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, include_view_image_tool: true, + experimental_unified_exec_tool: true, }); let tools = get_openai_tools( &config, @@ -723,7 +779,7 @@ mod tests { assert_eq_tool_names( &tools, &[ - "shell", + "unified_exec", "web_search", "view_image", "test_server/do_something_cool", @@ -786,6 +842,7 @@ mod tests { include_web_search_request: false, use_streamable_shell_tool: false, include_view_image_tool: true, + experimental_unified_exec_tool: true, }); // Intentionally construct a map with keys that would sort alphabetically. @@ -838,11 +895,11 @@ mod tests { ]); let tools = get_openai_tools(&config, Some(tools_map)); - // Expect shell first, followed by MCP tools sorted by fully-qualified name. + // Expect unified_exec first, followed by MCP tools sorted by fully-qualified name. assert_eq_tool_names( &tools, &[ - "shell", + "unified_exec", "view_image", "test_server/cool", "test_server/do", @@ -863,6 +920,7 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, include_view_image_tool: true, + experimental_unified_exec_tool: true, }); let tools = get_openai_tools( @@ -890,7 +948,7 @@ mod tests { assert_eq_tool_names( &tools, - &["shell", "web_search", "view_image", "dash/search"], + &["unified_exec", "web_search", "view_image", "dash/search"], ); assert_eq!( @@ -925,6 +983,7 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, include_view_image_tool: true, + experimental_unified_exec_tool: true, }); let tools = get_openai_tools( @@ -950,7 +1009,7 @@ mod tests { assert_eq_tool_names( &tools, - &["shell", "web_search", "view_image", "dash/paginate"], + &["unified_exec", "web_search", "view_image", "dash/paginate"], ); assert_eq!( tools[3], @@ -982,6 +1041,7 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, include_view_image_tool: true, + experimental_unified_exec_tool: true, }); let tools = get_openai_tools( @@ -1005,7 +1065,10 @@ mod tests { )])), ); - assert_eq_tool_names(&tools, &["shell", "web_search", "view_image", "dash/tags"]); + assert_eq_tool_names( + &tools, + &["unified_exec", "web_search", "view_image", "dash/tags"], + ); assert_eq!( tools[3], OpenAiTool::Function(ResponsesApiTool { @@ -1039,6 +1102,7 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, include_view_image_tool: true, + experimental_unified_exec_tool: true, }); let tools = get_openai_tools( @@ -1062,7 +1126,10 @@ mod tests { )])), ); - assert_eq_tool_names(&tools, &["shell", "web_search", "view_image", "dash/value"]); + assert_eq_tool_names( + &tools, + &["unified_exec", "web_search", "view_image", "dash/value"], + ); assert_eq!( tools[3], OpenAiTool::Function(ResponsesApiTool { diff --git a/codex-rs/core/src/rollout/policy.rs b/codex-rs/core/src/rollout/policy.rs index 7c58260479..c0d6d3af3a 100644 --- a/codex-rs/core/src/rollout/policy.rs +++ b/codex-rs/core/src/rollout/policy.rs @@ -65,6 +65,6 @@ pub(crate) fn should_persist_event_msg(ev: &EventMsg) -> bool { | EventMsg::PlanUpdate(_) | EventMsg::TurnAborted(_) | EventMsg::ShutdownComplete - | EventMsg::ConversationHistory(_) => false, + | EventMsg::ConversationPath(_) => false, } } diff --git a/codex-rs/core/src/rollout/recorder.rs b/codex-rs/core/src/rollout/recorder.rs index f5609f7206..63495bd8d0 100644 --- a/codex-rs/core/src/rollout/recorder.rs +++ b/codex-rs/core/src/rollout/recorder.rs @@ -77,7 +77,13 @@ pub enum RolloutRecorderParams { enum RolloutCmd { AddItems(Vec), - Shutdown { ack: oneshot::Sender<()> }, + /// Ensure all prior writes are processed; respond when flushed. + Flush { + ack: oneshot::Sender<()>, + }, + Shutdown { + ack: oneshot::Sender<()>, + }, } impl RolloutRecorderParams { @@ -185,6 +191,17 @@ impl RolloutRecorder { .map_err(|e| IoError::other(format!("failed to queue rollout items: {e}"))) } + /// Flush all queued writes and wait until they are committed by the writer task. + pub async fn flush(&self) -> std::io::Result<()> { + let (tx, rx) = oneshot::channel(); + self.tx + .send(RolloutCmd::Flush { ack: tx }) + .await + .map_err(|e| IoError::other(format!("failed to queue rollout flush: {e}")))?; + rx.await + .map_err(|e| IoError::other(format!("failed waiting for rollout flush: {e}"))) + } + pub(crate) async fn get_rollout_history(path: &Path) -> std::io::Result { info!("Resuming rollout from {path:?}"); tracing::error!("Resuming rollout from {path:?}"); @@ -211,11 +228,11 @@ impl RolloutRecorder { match serde_json::from_value::(v.clone()) { Ok(rollout_line) => match rollout_line.item { RolloutItem::SessionMeta(session_meta_line) => { - tracing::error!( - "Parsed conversation ID from rollout file: {:?}", - session_meta_line.meta.id - ); - conversation_id = Some(session_meta_line.meta.id); + // Use the FIRST SessionMeta encountered in the file as the canonical + // conversation id and main session information. Keep all items intact. + if conversation_id.is_none() { + conversation_id = Some(session_meta_line.meta.id); + } items.push(RolloutItem::SessionMeta(session_meta_line)); } RolloutItem::ResponseItem(item) => { @@ -251,6 +268,10 @@ impl RolloutRecorder { })) } + pub(crate) fn get_rollout_path(&self) -> PathBuf { + self.rollout_path.clone() + } + pub async fn shutdown(&self) -> std::io::Result<()> { let (tx_done, rx_done) = oneshot::channel(); match self.tx.send(RolloutCmd::Shutdown { ack: tx_done }).await { @@ -351,6 +372,14 @@ async fn rollout_writer( } } } + RolloutCmd::Flush { ack } => { + // Ensure underlying file is flushed and then ack. + if let Err(e) = writer.file.flush().await { + let _ = ack.send(()); + return Err(e); + } + let _ = ack.send(()); + } RolloutCmd::Shutdown { ack } => { let _ = ack.send(()); } diff --git a/codex-rs/core/src/truncate.rs b/codex-rs/core/src/truncate.rs new file mode 100644 index 0000000000..ab0158720b --- /dev/null +++ b/codex-rs/core/src/truncate.rs @@ -0,0 +1,180 @@ +//! Utilities for truncating large chunks of output while preserving a prefix +//! and suffix on UTF-8 boundaries. + +/// Truncate the middle of a UTF-8 string to at most `max_bytes` bytes, +/// preserving the beginning and the end. Returns the possibly truncated +/// string and `Some(original_token_count)` (estimated at 4 bytes/token) +/// if truncation occurred; otherwise returns the original string and `None`. +pub(crate) fn truncate_middle(s: &str, max_bytes: usize) -> (String, Option) { + if s.len() <= max_bytes { + return (s.to_string(), None); + } + + let est_tokens = (s.len() as u64).div_ceil(4); + if max_bytes == 0 { + return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); + } + + fn truncate_on_boundary(input: &str, max_len: usize) -> &str { + if input.len() <= max_len { + return input; + } + let mut end = max_len; + while end > 0 && !input.is_char_boundary(end) { + end -= 1; + } + &input[..end] + } + + fn pick_prefix_end(s: &str, left_budget: usize) -> usize { + if let Some(head) = s.get(..left_budget) + && let Some(i) = head.rfind('\n') + { + return i + 1; + } + truncate_on_boundary(s, left_budget).len() + } + + fn pick_suffix_start(s: &str, right_budget: usize) -> usize { + let start_tail = s.len().saturating_sub(right_budget); + if let Some(tail) = s.get(start_tail..) + && let Some(i) = tail.find('\n') + { + return start_tail + i + 1; + } + + let mut idx = start_tail.min(s.len()); + while idx < s.len() && !s.is_char_boundary(idx) { + idx += 1; + } + idx + } + + let mut guess_tokens = est_tokens; + for _ in 0..4 { + let marker = format!("…{guess_tokens} tokens truncated…"); + let marker_len = marker.len(); + let keep_budget = max_bytes.saturating_sub(marker_len); + if keep_budget == 0 { + return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); + } + + let left_budget = keep_budget / 2; + let right_budget = keep_budget - left_budget; + let prefix_end = pick_prefix_end(s, left_budget); + let mut suffix_start = pick_suffix_start(s, right_budget); + if suffix_start < prefix_end { + suffix_start = prefix_end; + } + + let kept_content_bytes = prefix_end + (s.len() - suffix_start); + let truncated_content_bytes = s.len().saturating_sub(kept_content_bytes); + let new_tokens = (truncated_content_bytes as u64).div_ceil(4); + + if new_tokens == guess_tokens { + let mut out = String::with_capacity(marker_len + kept_content_bytes + 1); + out.push_str(&s[..prefix_end]); + out.push_str(&marker); + out.push('\n'); + out.push_str(&s[suffix_start..]); + return (out, Some(est_tokens)); + } + + guess_tokens = new_tokens; + } + + let marker = format!("…{guess_tokens} tokens truncated…"); + let marker_len = marker.len(); + let keep_budget = max_bytes.saturating_sub(marker_len); + if keep_budget == 0 { + return (format!("…{est_tokens} tokens truncated…"), Some(est_tokens)); + } + + let left_budget = keep_budget / 2; + let right_budget = keep_budget - left_budget; + let prefix_end = pick_prefix_end(s, left_budget); + let suffix_start = pick_suffix_start(s, right_budget); + + let mut out = String::with_capacity(marker_len + prefix_end + (s.len() - suffix_start) + 1); + out.push_str(&s[..prefix_end]); + out.push_str(&marker); + out.push('\n'); + out.push_str(&s[suffix_start..]); + (out, Some(est_tokens)) +} + +#[cfg(test)] +mod tests { + use super::truncate_middle; + + #[test] + fn truncate_middle_no_newlines_fallback() { + let s = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ*"; + let max_bytes = 32; + let (out, original) = truncate_middle(s, max_bytes); + assert!(out.starts_with("abc")); + assert!(out.contains("tokens truncated")); + assert!(out.ends_with("XYZ*")); + assert_eq!(original, Some((s.len() as u64).div_ceil(4))); + } + + #[test] + fn truncate_middle_prefers_newline_boundaries() { + let mut s = String::new(); + for i in 1..=20 { + s.push_str(&format!("{i:03}\n")); + } + assert_eq!(s.len(), 80); + + let max_bytes = 64; + let (out, tokens) = truncate_middle(&s, max_bytes); + assert!(out.starts_with("001\n002\n003\n004\n")); + assert!(out.contains("tokens truncated")); + assert!(out.ends_with("017\n018\n019\n020\n")); + assert_eq!(tokens, Some(20)); + } + + #[test] + fn truncate_middle_handles_utf8_content() { + let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with ascii text\n"; + let max_bytes = 32; + let (out, tokens) = truncate_middle(s, max_bytes); + + assert!(out.contains("tokens truncated")); + assert!(!out.contains('\u{fffd}')); + assert_eq!(tokens, Some((s.len() as u64).div_ceil(4))); + } + + #[test] + fn truncate_middle_prefers_newline_boundaries_2() { + // Build a multi-line string of 20 numbered lines (each "NNN\n"). + let mut s = String::new(); + for i in 1..=20 { + s.push_str(&format!("{i:03}\n")); + } + // Total length: 20 lines * 4 bytes per line = 80 bytes. + assert_eq!(s.len(), 80); + + // Choose a cap that forces truncation while leaving room for + // a few lines on each side after accounting for the marker. + let max_bytes = 64; + // Expect exact output: first 4 lines, marker, last 4 lines, and correct token estimate (80/4 = 20). + assert_eq!( + truncate_middle(&s, max_bytes), + ( + r#"001 +002 +003 +004 +…12 tokens truncated… +017 +018 +019 +020 +"# + .to_string(), + Some(20) + ) + ); + } +} diff --git a/codex-rs/core/src/unified_exec/errors.rs b/codex-rs/core/src/unified_exec/errors.rs new file mode 100644 index 0000000000..6bf5bf7ec5 --- /dev/null +++ b/codex-rs/core/src/unified_exec/errors.rs @@ -0,0 +1,22 @@ +use thiserror::Error; + +#[derive(Debug, Error)] +pub(crate) enum UnifiedExecError { + #[error("Failed to create unified exec session: {pty_error}")] + CreateSession { + #[source] + pty_error: anyhow::Error, + }, + #[error("Unknown session id {session_id}")] + UnknownSessionId { session_id: i32 }, + #[error("failed to write to stdin")] + WriteToStdin, + #[error("missing command line for unified exec request")] + MissingCommandLine, +} + +impl UnifiedExecError { + pub(crate) fn create_session(error: anyhow::Error) -> Self { + Self::CreateSession { pty_error: error } + } +} diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs new file mode 100644 index 0000000000..25ab3d4dc0 --- /dev/null +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -0,0 +1,653 @@ +use portable_pty::CommandBuilder; +use portable_pty::PtySize; +use portable_pty::native_pty_system; +use std::collections::HashMap; +use std::collections::VecDeque; +use std::io::ErrorKind; +use std::io::Read; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicI32; +use std::sync::atomic::Ordering; +use tokio::sync::Mutex; +use tokio::sync::Notify; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tokio::time::Duration; +use tokio::time::Instant; + +use crate::exec_command::ExecCommandSession; +use crate::truncate::truncate_middle; + +mod errors; + +pub(crate) use errors::UnifiedExecError; + +const DEFAULT_TIMEOUT_MS: u64 = 1_000; +const MAX_TIMEOUT_MS: u64 = 60_000; +const UNIFIED_EXEC_OUTPUT_MAX_BYTES: usize = 128 * 1024; // 128 KiB + +#[derive(Debug)] +pub(crate) struct UnifiedExecRequest<'a> { + pub session_id: Option, + pub input_chunks: &'a [String], + pub timeout_ms: Option, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct UnifiedExecResult { + pub session_id: Option, + pub output: String, +} + +#[derive(Debug, Default)] +pub(crate) struct UnifiedExecSessionManager { + next_session_id: AtomicI32, + sessions: Mutex>, +} + +#[derive(Debug)] +struct ManagedUnifiedExecSession { + session: ExecCommandSession, + output_buffer: OutputBuffer, + /// Notifies waiters whenever new output has been appended to + /// `output_buffer`, allowing clients to poll for fresh data. + output_notify: Arc, + output_task: JoinHandle<()>, +} + +#[derive(Debug, Default)] +struct OutputBufferState { + chunks: VecDeque>, + total_bytes: usize, +} + +impl OutputBufferState { + fn push_chunk(&mut self, chunk: Vec) { + self.total_bytes = self.total_bytes.saturating_add(chunk.len()); + self.chunks.push_back(chunk); + + let mut excess = self + .total_bytes + .saturating_sub(UNIFIED_EXEC_OUTPUT_MAX_BYTES); + + while excess > 0 { + match self.chunks.front_mut() { + Some(front) if excess >= front.len() => { + excess -= front.len(); + self.total_bytes = self.total_bytes.saturating_sub(front.len()); + self.chunks.pop_front(); + } + Some(front) => { + front.drain(..excess); + self.total_bytes = self.total_bytes.saturating_sub(excess); + break; + } + None => break, + } + } + } + + fn drain(&mut self) -> Vec> { + let drained: Vec> = self.chunks.drain(..).collect(); + self.total_bytes = 0; + drained + } +} + +type OutputBuffer = Arc>; +type OutputHandles = (OutputBuffer, Arc); + +impl ManagedUnifiedExecSession { + fn new(session: ExecCommandSession) -> Self { + let output_buffer = Arc::new(Mutex::new(OutputBufferState::default())); + let output_notify = Arc::new(Notify::new()); + let mut receiver = session.output_receiver(); + let buffer_clone = Arc::clone(&output_buffer); + let notify_clone = Arc::clone(&output_notify); + let output_task = tokio::spawn(async move { + while let Ok(chunk) = receiver.recv().await { + let mut guard = buffer_clone.lock().await; + guard.push_chunk(chunk); + drop(guard); + notify_clone.notify_waiters(); + } + }); + + Self { + session, + output_buffer, + output_notify, + output_task, + } + } + + fn writer_sender(&self) -> mpsc::Sender> { + self.session.writer_sender() + } + + fn output_handles(&self) -> OutputHandles { + ( + Arc::clone(&self.output_buffer), + Arc::clone(&self.output_notify), + ) + } + + fn has_exited(&self) -> bool { + self.session.has_exited() + } +} + +impl Drop for ManagedUnifiedExecSession { + fn drop(&mut self) { + self.output_task.abort(); + } +} + +impl UnifiedExecSessionManager { + pub async fn handle_request( + &self, + request: UnifiedExecRequest<'_>, + ) -> Result { + let (timeout_ms, timeout_warning) = match request.timeout_ms { + Some(requested) if requested > MAX_TIMEOUT_MS => ( + MAX_TIMEOUT_MS, + Some(format!( + "Warning: requested timeout {requested}ms exceeds maximum of {MAX_TIMEOUT_MS}ms; clamping to {MAX_TIMEOUT_MS}ms.\n" + )), + ), + Some(requested) => (requested, None), + None => (DEFAULT_TIMEOUT_MS, None), + }; + + let mut new_session: Option = None; + let session_id; + let writer_tx; + let output_buffer; + let output_notify; + + if let Some(existing_id) = request.session_id { + let mut sessions = self.sessions.lock().await; + match sessions.get(&existing_id) { + Some(session) => { + if session.has_exited() { + sessions.remove(&existing_id); + return Err(UnifiedExecError::UnknownSessionId { + session_id: existing_id, + }); + } + let (buffer, notify) = session.output_handles(); + session_id = existing_id; + writer_tx = session.writer_sender(); + output_buffer = buffer; + output_notify = notify; + } + None => { + return Err(UnifiedExecError::UnknownSessionId { + session_id: existing_id, + }); + } + } + drop(sessions); + } else { + let command = request.input_chunks.to_vec(); + let new_id = self.next_session_id.fetch_add(1, Ordering::SeqCst); + let session = create_unified_exec_session(&command).await?; + let managed_session = ManagedUnifiedExecSession::new(session); + let (buffer, notify) = managed_session.output_handles(); + writer_tx = managed_session.writer_sender(); + output_buffer = buffer; + output_notify = notify; + session_id = new_id; + new_session = Some(managed_session); + }; + + if request.session_id.is_some() { + let joined_input = request.input_chunks.join(" "); + if !joined_input.is_empty() && writer_tx.send(joined_input.into_bytes()).await.is_err() + { + return Err(UnifiedExecError::WriteToStdin); + } + } + + let mut collected: Vec = Vec::with_capacity(4096); + let start = Instant::now(); + let deadline = start + Duration::from_millis(timeout_ms); + + loop { + let drained_chunks; + let mut wait_for_output = None; + { + let mut guard = output_buffer.lock().await; + drained_chunks = guard.drain(); + if drained_chunks.is_empty() { + wait_for_output = Some(output_notify.notified()); + } + } + + if drained_chunks.is_empty() { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining == Duration::ZERO { + break; + } + + let notified = wait_for_output.unwrap_or_else(|| output_notify.notified()); + tokio::pin!(notified); + tokio::select! { + _ = &mut notified => {} + _ = tokio::time::sleep(remaining) => break, + } + continue; + } + + for chunk in drained_chunks { + collected.extend_from_slice(&chunk); + } + + if Instant::now() >= deadline { + break; + } + } + + let (output, _maybe_tokens) = truncate_middle( + &String::from_utf8_lossy(&collected), + UNIFIED_EXEC_OUTPUT_MAX_BYTES, + ); + let output = if let Some(warning) = timeout_warning { + format!("{warning}{output}") + } else { + output + }; + + let should_store_session = if let Some(session) = new_session.as_ref() { + !session.has_exited() + } else if request.session_id.is_some() { + let mut sessions = self.sessions.lock().await; + if let Some(existing) = sessions.get(&session_id) { + if existing.has_exited() { + sessions.remove(&session_id); + false + } else { + true + } + } else { + false + } + } else { + true + }; + + if should_store_session { + if let Some(session) = new_session { + self.sessions.lock().await.insert(session_id, session); + } + Ok(UnifiedExecResult { + session_id: Some(session_id), + output, + }) + } else { + Ok(UnifiedExecResult { + session_id: None, + output, + }) + } + } +} + +async fn create_unified_exec_session( + command: &[String], +) -> Result { + if command.is_empty() { + return Err(UnifiedExecError::MissingCommandLine); + } + + let pty_system = native_pty_system(); + + let pair = pty_system + .openpty(PtySize { + rows: 24, + cols: 80, + pixel_width: 0, + pixel_height: 0, + }) + .map_err(UnifiedExecError::create_session)?; + + // Safe thanks to the check at the top of the function. + let mut command_builder = CommandBuilder::new(command[0].clone()); + for arg in &command[1..] { + command_builder.arg(arg); + } + + let mut child = pair + .slave + .spawn_command(command_builder) + .map_err(UnifiedExecError::create_session)?; + let killer = child.clone_killer(); + + let (writer_tx, mut writer_rx) = mpsc::channel::>(128); + let (output_tx, _) = tokio::sync::broadcast::channel::>(256); + + let mut reader = pair + .master + .try_clone_reader() + .map_err(UnifiedExecError::create_session)?; + let output_tx_clone = output_tx.clone(); + let reader_handle = tokio::task::spawn_blocking(move || { + let mut buf = [0u8; 8192]; + loop { + match reader.read(&mut buf) { + Ok(0) => break, + Ok(n) => { + let _ = output_tx_clone.send(buf[..n].to_vec()); + } + Err(ref e) if e.kind() == ErrorKind::Interrupted => continue, + Err(ref e) if e.kind() == ErrorKind::WouldBlock => { + std::thread::sleep(Duration::from_millis(5)); + continue; + } + Err(_) => break, + } + } + }); + + let writer = pair + .master + .take_writer() + .map_err(UnifiedExecError::create_session)?; + let writer = Arc::new(StdMutex::new(writer)); + let writer_handle = tokio::spawn({ + let writer = writer.clone(); + async move { + while let Some(bytes) = writer_rx.recv().await { + let writer = writer.clone(); + let _ = tokio::task::spawn_blocking(move || { + if let Ok(mut guard) = writer.lock() { + use std::io::Write; + let _ = guard.write_all(&bytes); + let _ = guard.flush(); + } + }) + .await; + } + } + }); + + let exit_status = Arc::new(AtomicBool::new(false)); + let wait_exit_status = Arc::clone(&exit_status); + let wait_handle = tokio::task::spawn_blocking(move || { + let _ = child.wait(); + wait_exit_status.store(true, Ordering::SeqCst); + }); + + Ok(ExecCommandSession::new( + writer_tx, + output_tx, + killer, + reader_handle, + writer_handle, + wait_handle, + exit_status, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn push_chunk_trims_only_excess_bytes() { + let mut buffer = OutputBufferState::default(); + buffer.push_chunk(vec![b'a'; UNIFIED_EXEC_OUTPUT_MAX_BYTES]); + buffer.push_chunk(vec![b'b']); + buffer.push_chunk(vec![b'c']); + + assert_eq!(buffer.total_bytes, UNIFIED_EXEC_OUTPUT_MAX_BYTES); + assert_eq!(buffer.chunks.len(), 3); + assert_eq!( + buffer.chunks.front().unwrap().len(), + UNIFIED_EXEC_OUTPUT_MAX_BYTES - 2 + ); + assert_eq!(buffer.chunks.pop_back().unwrap(), vec![b'c']); + assert_eq!(buffer.chunks.pop_back().unwrap(), vec![b'b']); + } + + #[cfg(unix)] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn unified_exec_persists_across_requests_jif() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + + let open_shell = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["bash".to_string(), "-i".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + let session_id = open_shell.session_id.expect("expected session_id"); + + manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &[ + "export".to_string(), + "CODEX_INTERACTIVE_SHELL_VAR=codex\n".to_string(), + ], + timeout_ms: Some(2_500), + }) + .await?; + + let out_2 = manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &["echo $CODEX_INTERACTIVE_SHELL_VAR\n".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + assert!(out_2.output.contains("codex")); + + Ok(()) + } + + #[cfg(unix)] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn multi_unified_exec_sessions() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + + let shell_a = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["/bin/bash".to_string(), "-i".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + let session_a = shell_a.session_id.expect("expected session id"); + + manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_a), + input_chunks: &["export CODEX_INTERACTIVE_SHELL_VAR=codex\n".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + + let out_2 = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &[ + "echo".to_string(), + "$CODEX_INTERACTIVE_SHELL_VAR\n".to_string(), + ], + timeout_ms: Some(1_500), + }) + .await?; + assert!(!out_2.output.contains("codex")); + + let out_3 = manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_a), + input_chunks: &["echo $CODEX_INTERACTIVE_SHELL_VAR\n".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + assert!(out_3.output.contains("codex")); + + Ok(()) + } + + #[cfg(unix)] + #[tokio::test] + async fn unified_exec_timeouts() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + + let open_shell = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["bash".to_string(), "-i".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + let session_id = open_shell.session_id.expect("expected session id"); + + manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &[ + "export".to_string(), + "CODEX_INTERACTIVE_SHELL_VAR=codex\n".to_string(), + ], + timeout_ms: Some(1_500), + }) + .await?; + + let out_2 = manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &["sleep 5 && echo $CODEX_INTERACTIVE_SHELL_VAR\n".to_string()], + timeout_ms: Some(10), + }) + .await?; + assert!(!out_2.output.contains("codex")); + + tokio::time::sleep(Duration::from_secs(7)).await; + + let empty = Vec::new(); + let out_3 = manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &empty, + timeout_ms: Some(100), + }) + .await?; + + assert!(out_3.output.contains("codex")); + + Ok(()) + } + + #[cfg(unix)] + #[tokio::test] + async fn requests_with_large_timeout_are_capped() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + + let result = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["echo".to_string(), "codex".to_string()], + timeout_ms: Some(120_000), + }) + .await?; + + assert!(result.output.starts_with( + "Warning: requested timeout 120000ms exceeds maximum of 60000ms; clamping to 60000ms.\n" + )); + assert!(result.output.contains("codex")); + + Ok(()) + } + + #[cfg(unix)] + #[tokio::test] + async fn completed_commands_do_not_persist_sessions() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + let result = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["/bin/echo".to_string(), "codex".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + + assert!(result.session_id.is_none()); + assert!(result.output.contains("codex")); + + assert!(manager.sessions.lock().await.is_empty()); + + Ok(()) + } + + #[cfg(unix)] + #[tokio::test] + async fn correct_path_resolution() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + let result = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["echo".to_string(), "codex".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + + assert!(result.session_id.is_none()); + assert!(result.output.contains("codex")); + + assert!(manager.sessions.lock().await.is_empty()); + + Ok(()) + } + + #[cfg(unix)] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn reusing_completed_session_returns_unknown_session() -> Result<(), UnifiedExecError> { + let manager = UnifiedExecSessionManager::default(); + + let open_shell = manager + .handle_request(UnifiedExecRequest { + session_id: None, + input_chunks: &["/bin/bash".to_string(), "-i".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + let session_id = open_shell.session_id.expect("expected session id"); + + manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &["exit\n".to_string()], + timeout_ms: Some(1_500), + }) + .await?; + + tokio::time::sleep(Duration::from_millis(200)).await; + + let err = manager + .handle_request(UnifiedExecRequest { + session_id: Some(session_id), + input_chunks: &[], + timeout_ms: Some(100), + }) + .await + .expect_err("expected unknown session error"); + + match err { + UnifiedExecError::UnknownSessionId { session_id: err_id } => { + assert_eq!(err_id, session_id); + } + other => panic!("expected UnknownSessionId, got {other:?}"), + } + + assert!(!manager.sessions.lock().await.contains_key(&session_id)); + + Ok(()) + } +} diff --git a/codex-rs/core/tests/suite/fork_conversation.rs b/codex-rs/core/tests/suite/fork_conversation.rs index 92e43dbfeb..08e3f29d28 100644 --- a/codex-rs/core/tests/suite/fork_conversation.rs +++ b/codex-rs/core/tests/suite/fork_conversation.rs @@ -1,12 +1,16 @@ use codex_core::CodexAuth; +use codex_core::ContentItem; use codex_core::ConversationManager; use codex_core::ModelProviderInfo; use codex_core::NewConversation; +use codex_core::ResponseItem; use codex_core::built_in_model_providers; -use codex_core::protocol::ConversationHistoryResponseEvent; +use codex_core::protocol::ConversationPathResponseEvent; use codex_core::protocol::EventMsg; use codex_core::protocol::InputItem; use codex_core::protocol::Op; +use codex_core::protocol::RolloutItem; +use codex_core::protocol::RolloutLine; use core_test_support::load_default_config_for_test; use core_test_support::wait_for_event; use tempfile::TempDir; @@ -71,84 +75,121 @@ async fn fork_conversation_twice_drops_to_first_message() { let _ = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; } - // Request history from the base conversation. - codex.submit(Op::GetHistory).await.unwrap(); + // Request history from the base conversation to obtain rollout path. + codex.submit(Op::GetPath).await.unwrap(); let base_history = - wait_for_event(&codex, |ev| matches!(ev, EventMsg::ConversationHistory(_))).await; - - // Capture entries from the base history and compute expected prefixes after each fork. - let entries_after_three = match &base_history { - EventMsg::ConversationHistory(ConversationHistoryResponseEvent { entries, .. }) => { - entries.clone() - } + wait_for_event(&codex, |ev| matches!(ev, EventMsg::ConversationPath(_))).await; + let base_path = match &base_history { + EventMsg::ConversationPath(ConversationPathResponseEvent { path, .. }) => path.clone(), _ => panic!("expected ConversationHistory event"), }; - // History layout for this test: - // [0] user instructions, - // [1] environment context, - // [2] "first" user message, - // [3] "second" user message, - // [4] "third" user message. - // Fork 1: drops the last user message and everything after. - let expected_after_first = vec![ - entries_after_three[0].clone(), - entries_after_three[1].clone(), - entries_after_three[2].clone(), - entries_after_three[3].clone(), - ]; + // GetHistory flushes before returning the path; no wait needed. - // Fork 2: drops the last user message and everything after. - // [0] user instructions, - // [1] environment context, - // [2] "first" user message, - let expected_after_second = vec![ - entries_after_three[0].clone(), - entries_after_three[1].clone(), - entries_after_three[2].clone(), - ]; + // Helper: read rollout items (excluding SessionMeta) from a JSONL path. + let read_items = |p: &std::path::Path| -> Vec { + let text = std::fs::read_to_string(p).expect("read rollout file"); + let mut items: Vec = Vec::new(); + for line in text.lines() { + if line.trim().is_empty() { + continue; + } + let v: serde_json::Value = serde_json::from_str(line).expect("jsonl line"); + let rl: RolloutLine = serde_json::from_value(v).expect("rollout line"); + match rl.item { + RolloutItem::SessionMeta(_) => {} + other => items.push(other), + } + } + items + }; - // Fork once with n=1 → drops the last user message and everything after. + // Compute expected prefixes after each fork by truncating base rollout at nth-from-last user input. + let base_items = read_items(&base_path); + let find_user_input_positions = |items: &[RolloutItem]| -> Vec { + let mut pos = Vec::new(); + for (i, it) in items.iter().enumerate() { + if let RolloutItem::ResponseItem(ResponseItem::Message { role, content, .. }) = it + && role == "user" + { + // Consider any user message as an input boundary; recorder stores both EventMsg and ResponseItem. + // We specifically look for input items, which are represented as ContentItem::InputText. + if content + .iter() + .any(|c| matches!(c, ContentItem::InputText { .. })) + { + pos.push(i); + } + } + } + pos + }; + let user_inputs = find_user_input_positions(&base_items); + + // After dropping last user input (n=1), cut strictly before that input if present, else empty. + let cut1 = user_inputs + .get(user_inputs.len().saturating_sub(1)) + .copied() + .unwrap_or(0); + let expected_after_first: Vec = base_items[..cut1].to_vec(); + + // After dropping again (n=1 on fork1), compute expected relative to fork1's rollout. + + // Fork once with n=1 → drops the last user input and everything after. let NewConversation { conversation: codex_fork1, .. } = conversation_manager - .fork_conversation(entries_after_three.clone(), 1, config_for_fork.clone()) + .fork_conversation(1, config_for_fork.clone(), base_path.clone()) .await .expect("fork 1"); - codex_fork1.submit(Op::GetHistory).await.unwrap(); + codex_fork1.submit(Op::GetPath).await.unwrap(); let fork1_history = wait_for_event(&codex_fork1, |ev| { - matches!(ev, EventMsg::ConversationHistory(_)) + matches!(ev, EventMsg::ConversationPath(_)) }) .await; - let entries_after_first_fork = match &fork1_history { - EventMsg::ConversationHistory(ConversationHistoryResponseEvent { entries, .. }) => { - assert!(matches!( - fork1_history, - EventMsg::ConversationHistory(ConversationHistoryResponseEvent { ref entries, .. }) if *entries == expected_after_first - )); - entries.clone() - } + let fork1_path = match &fork1_history { + EventMsg::ConversationPath(ConversationPathResponseEvent { path, .. }) => path.clone(), _ => panic!("expected ConversationHistory event after first fork"), }; + // GetHistory on fork1 flushed; the file is ready. + let fork1_items = read_items(&fork1_path); + pretty_assertions::assert_eq!( + serde_json::to_value(&fork1_items).unwrap(), + serde_json::to_value(&expected_after_first).unwrap() + ); + // Fork again with n=1 → drops the (new) last user message, leaving only the first. let NewConversation { conversation: codex_fork2, .. } = conversation_manager - .fork_conversation(entries_after_first_fork.clone(), 1, config_for_fork.clone()) + .fork_conversation(1, config_for_fork.clone(), fork1_path.clone()) .await .expect("fork 2"); - codex_fork2.submit(Op::GetHistory).await.unwrap(); + codex_fork2.submit(Op::GetPath).await.unwrap(); let fork2_history = wait_for_event(&codex_fork2, |ev| { - matches!(ev, EventMsg::ConversationHistory(_)) + matches!(ev, EventMsg::ConversationPath(_)) }) .await; - assert!(matches!( - fork2_history, - EventMsg::ConversationHistory(ConversationHistoryResponseEvent { ref entries, .. }) if *entries == expected_after_second - )); + let fork2_path = match &fork2_history { + EventMsg::ConversationPath(ConversationPathResponseEvent { path, .. }) => path.clone(), + _ => panic!("expected ConversationHistory event after second fork"), + }; + // GetHistory on fork2 flushed; the file is ready. + let fork1_items = read_items(&fork1_path); + let fork1_user_inputs = find_user_input_positions(&fork1_items); + let cut_last_on_fork1 = fork1_user_inputs + .get(fork1_user_inputs.len().saturating_sub(1)) + .copied() + .unwrap_or(0); + let expected_after_second: Vec = fork1_items[..cut_last_on_fork1].to_vec(); + let fork2_items = read_items(&fork2_path); + pretty_assertions::assert_eq!( + serde_json::to_value(&fork2_items).unwrap(), + serde_json::to_value(&expected_after_second).unwrap() + ); } diff --git a/codex-rs/core/tests/suite/prompt_caching.rs b/codex-rs/core/tests/suite/prompt_caching.rs index 5ac3da6003..e77df01f20 100644 --- a/codex-rs/core/tests/suite/prompt_caching.rs +++ b/codex-rs/core/tests/suite/prompt_caching.rs @@ -191,7 +191,8 @@ async fn prompt_tools_are_consistent_across_requests() { let expected_instructions: &str = include_str!("../../prompt.md"); // our internal implementation is responsible for keeping tools in sync // with the OpenAI schema, so we just verify the tool presence here - let expected_tools_names: &[&str] = &["shell", "update_plan", "apply_patch", "view_image"]; + let expected_tools_names: &[&str] = + &["unified_exec", "update_plan", "apply_patch", "view_image"]; let body0 = requests[0].body_json::().unwrap(); assert_eq!( body0["instructions"], diff --git a/codex-rs/exec/src/event_processor_with_human_output.rs b/codex-rs/exec/src/event_processor_with_human_output.rs index ae4708976f..3c879397bd 100644 --- a/codex-rs/exec/src/event_processor_with_human_output.rs +++ b/codex-rs/exec/src/event_processor_with_human_output.rs @@ -559,7 +559,7 @@ impl EventProcessor for EventProcessorWithHumanOutput { } }, EventMsg::ShutdownComplete => return CodexStatus::Shutdown, - EventMsg::ConversationHistory(_) => {} + EventMsg::ConversationPath(_) => {} EventMsg::UserMessage(_) => {} } CodexStatus::Running diff --git a/codex-rs/mcp-client/src/main.rs b/codex-rs/mcp-client/src/main.rs index a9f4e335d1..d25bca4ba3 100644 --- a/codex-rs/mcp-client/src/main.rs +++ b/codex-rs/mcp-client/src/main.rs @@ -17,10 +17,10 @@ use anyhow::Context; use anyhow::Result; use codex_mcp_client::McpClient; use mcp_types::ClientCapabilities; +use mcp_types::Implementation; use mcp_types::InitializeRequestParams; use mcp_types::ListToolsRequestParams; use mcp_types::MCP_SCHEMA_VERSION; -use mcp_types::McpClientInfo; use tracing_subscriber::EnvFilter; #[tokio::main] @@ -60,10 +60,13 @@ async fn main() -> Result<()> { sampling: None, elicitation: None, }, - client_info: McpClientInfo { + client_info: Implementation { name: "codex-mcp-client".to_owned(), version: env!("CARGO_PKG_VERSION").to_owned(), title: Some("Codex".to_string()), + // This field is used by Codex when it is an MCP server: it should + // not be used when Codex is an MCP client. + user_agent: None, }, protocol_version: MCP_SCHEMA_VERSION.to_owned(), }; diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index feeb7f018a..d3351faa88 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -277,7 +277,7 @@ async fn run_codex_tool_session_inner( | EventMsg::GetHistoryEntryResponse(_) | EventMsg::PlanUpdate(_) | EventMsg::TurnAborted(_) - | EventMsg::ConversationHistory(_) + | EventMsg::ConversationPath(_) | EventMsg::UserMessage(_) | EventMsg::ShutdownComplete => { // For now, we do not do anything extra for these diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 6e63dee5f7..83e3d1cb75 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -234,7 +234,7 @@ impl MessageProcessor { }, instructions: None, protocol_version: params.protocol_version.clone(), - server_info: mcp_types::McpServerInfo { + server_info: mcp_types::Implementation { name: "codex-mcp-server".to_string(), version: env!("CARGO_PKG_VERSION").to_string(), title: Some("Codex".to_string()), diff --git a/codex-rs/mcp-server/tests/common/mcp_process.rs b/codex-rs/mcp-server/tests/common/mcp_process.rs index 3f4f51d9b5..ec7953ded6 100644 --- a/codex-rs/mcp-server/tests/common/mcp_process.rs +++ b/codex-rs/mcp-server/tests/common/mcp_process.rs @@ -26,13 +26,13 @@ use codex_protocol::mcp_protocol::SendUserTurnParams; use mcp_types::CallToolRequestParams; use mcp_types::ClientCapabilities; +use mcp_types::Implementation; use mcp_types::InitializeRequestParams; use mcp_types::JSONRPC_VERSION; use mcp_types::JSONRPCMessage; use mcp_types::JSONRPCNotification; use mcp_types::JSONRPCRequest; use mcp_types::JSONRPCResponse; -use mcp_types::McpClientInfo; use mcp_types::ModelContextProtocolNotification; use mcp_types::ModelContextProtocolRequest; use mcp_types::RequestId; @@ -134,10 +134,11 @@ impl McpProcess { roots: None, sampling: None, }, - client_info: McpClientInfo { + client_info: Implementation { name: "elicitation test".into(), title: Some("Elicitation Test".into()), version: "0.0.0".into(), + user_agent: None, }, protocol_version: mcp_types::MCP_SCHEMA_VERSION.into(), }; diff --git a/codex-rs/mcp-types/check_lib_rs.py b/codex-rs/mcp-types/check_lib_rs.py new file mode 100755 index 0000000000..37b623a260 --- /dev/null +++ b/codex-rs/mcp-types/check_lib_rs.py @@ -0,0 +1,21 @@ +#!/usr/bin/env python3 + +import subprocess +import sys +from pathlib import Path + + +def main() -> int: + crate_dir = Path(__file__).resolve().parent + generator = crate_dir / "generate_mcp_types.py" + + result = subprocess.run( + [sys.executable, str(generator), "--check"], + cwd=crate_dir, + check=False, + ) + return result.returncode + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/codex-rs/mcp-types/generate_mcp_types.py b/codex-rs/mcp-types/generate_mcp_types.py index d71a6e59a2..81d152498d 100755 --- a/codex-rs/mcp-types/generate_mcp_types.py +++ b/codex-rs/mcp-types/generate_mcp_types.py @@ -5,15 +5,22 @@ import argparse import json import subprocess import sys +import tempfile from dataclasses import ( dataclass, ) +from difflib import unified_diff from pathlib import Path +from shutil import copy2 # Helper first so it is defined when other functions call it. from typing import Any, Literal + +def eprint(*args: Any, **kwargs: Any) -> None: + print(*args, file=sys.stderr, **kwargs) + SCHEMA_VERSION = "2025-06-18" JSONRPC_VERSION = "2.0" @@ -43,16 +50,31 @@ def main() -> int: default_schema_file = ( Path(__file__).resolve().parent / "schema" / SCHEMA_VERSION / "schema.json" ) + default_lib_rs = Path(__file__).resolve().parent / "src/lib.rs" parser.add_argument( "schema_file", nargs="?", default=default_schema_file, help="schema.json file to process", ) + parser.add_argument( + "--check", + action="store_true", + help="Regenerate lib.rs in a sandbox and ensure the checked-in file matches", + ) args = parser.parse_args() - schema_file = args.schema_file + schema_file = Path(args.schema_file) + crate_dir = Path(__file__).resolve().parent - lib_rs = Path(__file__).resolve().parent / "src/lib.rs" + if args.check: + return run_check(schema_file, crate_dir, default_lib_rs) + + generate_lib_rs(schema_file, default_lib_rs, fmt=True) + return 0 + + +def generate_lib_rs(schema_file: Path, lib_rs: Path, fmt: bool) -> None: + lib_rs.parent.mkdir(parents=True, exist_ok=True) global DEFINITIONS # Allow helper functions to access the schema. @@ -185,13 +207,72 @@ fn default_jsonrpc() -> String {{ JSONRPC_VERSION.to_owned() }} for chunk in out: f.write(chunk) - subprocess.check_call( - ["cargo", "fmt", "--", "--config", "imports_granularity=Item"], - cwd=lib_rs.parent.parent, - stderr=subprocess.DEVNULL, - ) + if fmt: + subprocess.check_call( + ["cargo", "fmt", "--", "--config", "imports_granularity=Item"], + cwd=lib_rs.parent.parent, + stderr=subprocess.DEVNULL, + ) - return 0 + +def run_check(schema_file: Path, crate_dir: Path, checked_in_lib: Path) -> int: + config_path = crate_dir.parent / "rustfmt.toml" + eprint(f"Running --check with schema {schema_file}") + + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + eprint(f"Created temporary workspace at {tmp_path}") + manifest_path = tmp_path / "Cargo.toml" + eprint(f"Copying Cargo.toml into {manifest_path}") + copy2(crate_dir / "Cargo.toml", manifest_path) + manifest_text = manifest_path.read_text(encoding="utf-8") + manifest_text = manifest_text.replace( + "version = { workspace = true }", 'version = "0.0.0"', + ) + manifest_text = manifest_text.replace("\n[lints]\nworkspace = true\n", "\n") + manifest_path.write_text(manifest_text, encoding="utf-8") + src_dir = tmp_path / "src" + src_dir.mkdir(parents=True, exist_ok=True) + eprint(f"Generating lib.rs into {src_dir}") + generated_lib = src_dir / "lib.rs" + + generate_lib_rs(schema_file, generated_lib, fmt=False) + + eprint("Formatting generated lib.rs with rustfmt") + subprocess.check_call( + [ + "rustfmt", + "--config-path", + str(config_path), + str(generated_lib), + ], + cwd=tmp_path, + stderr=subprocess.DEVNULL, + ) + + eprint("Running cargo build in temporary workspace") + subprocess.check_call(["cargo", "build", "--quiet"], cwd=tmp_path) + + eprint("Comparing generated lib.rs with checked-in version") + checked_in_contents = checked_in_lib.read_text(encoding="utf-8") + generated_contents = generated_lib.read_text(encoding="utf-8") + + if checked_in_contents == generated_contents: + eprint("lib.rs matches checked-in version") + return 0 + + diff = unified_diff( + checked_in_contents.splitlines(keepends=True), + generated_contents.splitlines(keepends=True), + fromfile=str(checked_in_lib), + tofile=str(generated_lib), + ) + diff_text = "".join(diff) + eprint("Generated lib.rs does not match the checked-in version. Diff:") + if diff_text: + eprint(diff_text, end="") + eprint("Re-run generate_mcp_types.py without --check to update src/lib.rs.") + return 1 def add_definition(name: str, definition: dict[str, Any], out: list[str]) -> None: @@ -265,8 +346,11 @@ class StructField: name: str type_name: str serde: str | None = None + comment: str | None = None def append(self, out: list[str], supports_const: bool) -> None: + if self.comment: + out.append(f" // {self.comment}\n") if self.serde: out.append(f" {self.serde}\n") if self.viz == "const": @@ -312,6 +396,18 @@ def define_struct( else: fields.append(StructField("pub", rs_prop.name, prop_type, rs_prop.serde)) + # Special-case: add Codex-specific user_agent to Implementation + if name == "Implementation": + fields.append( + StructField( + "pub", + "user_agent", + "Option", + '#[serde(default, skip_serializing_if = "Option::is_none")]', + "This is an extra field that the Codex MCP server sends as part of InitializeResult.", + ) + ) + if implements_request_trait(name): add_trait_impl(name, "ModelContextProtocolRequest", fields, out) elif implements_notification_trait(name): diff --git a/codex-rs/mcp-types/src/lib.rs b/codex-rs/mcp-types/src/lib.rs index 75778cb27c..37d790ab6a 100644 --- a/codex-rs/mcp-types/src/lib.rs +++ b/codex-rs/mcp-types/src/lib.rs @@ -482,20 +482,12 @@ pub struct ImageContent { /// Describes the name and version of an MCP implementation, with an optional title for UI representation. #[derive(Debug, Clone, PartialEq, Deserialize, Serialize, TS)] -pub struct McpClientInfo { - pub name: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub title: Option, - pub version: String, -} - -/// Describes the name and version of an MCP implementation, with an optional title for UI representation. -#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, TS)] -pub struct McpServerInfo { +pub struct Implementation { pub name: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub title: Option, pub version: String, + // This is an extra field that the Codex MCP server sends as part of InitializeResult. #[serde(default, skip_serializing_if = "Option::is_none")] pub user_agent: Option, } @@ -513,7 +505,7 @@ impl ModelContextProtocolRequest for InitializeRequest { pub struct InitializeRequestParams { pub capabilities: ClientCapabilities, #[serde(rename = "clientInfo")] - pub client_info: McpClientInfo, + pub client_info: Implementation, #[serde(rename = "protocolVersion")] pub protocol_version: String, } @@ -527,7 +519,7 @@ pub struct InitializeResult { #[serde(rename = "protocolVersion")] pub protocol_version: String, #[serde(rename = "serverInfo")] - pub server_info: McpServerInfo, + pub server_info: Implementation, } impl From for serde_json::Value { diff --git a/codex-rs/mcp-types/tests/suite/initialize.rs b/codex-rs/mcp-types/tests/suite/initialize.rs index 2bd3c789c9..73ff120274 100644 --- a/codex-rs/mcp-types/tests/suite/initialize.rs +++ b/codex-rs/mcp-types/tests/suite/initialize.rs @@ -1,10 +1,10 @@ use mcp_types::ClientCapabilities; use mcp_types::ClientRequest; +use mcp_types::Implementation; use mcp_types::InitializeRequestParams; use mcp_types::JSONRPC_VERSION; use mcp_types::JSONRPCMessage; use mcp_types::JSONRPCRequest; -use mcp_types::McpClientInfo; use mcp_types::RequestId; use serde_json::json; @@ -58,10 +58,11 @@ fn deserialize_initialize_request() { sampling: None, elicitation: None, }, - client_info: McpClientInfo { + client_info: Implementation { name: "acme-client".into(), title: Some("Acme".to_string()), version: "1.2.3".into(), + user_agent: None, }, protocol_version: "2025-06-18".into(), } diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 23a86a594e..13a81c9de6 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -115,7 +115,6 @@ pub enum ResponseItem { status: Option, action: WebSearchAction, }, - #[serde(other)] Other, } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 7e358f12e2..be31228560 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -149,7 +149,7 @@ pub enum Op { /// Request the full in-memory conversation transcript for the current session. /// Reply is delivered via `EventMsg::ConversationHistory`. - GetHistory, + GetPath, /// Request the list of MCP tools available across all configured servers. /// Reply is delivered via `EventMsg::McpListToolsResponse`. @@ -499,7 +499,7 @@ pub enum EventMsg { /// Notification that the agent is shutting down. ShutdownComplete, - ConversationHistory(ConversationHistoryResponseEvent), + ConversationPath(ConversationPathResponseEvent), } // Individual event payload types matching each `EventMsg` variant. @@ -801,9 +801,9 @@ pub struct WebSearchEndEvent { /// Response payload for `Op::GetHistory` containing the current session's /// in-memory transcript. #[derive(Debug, Clone, Deserialize, Serialize, TS)] -pub struct ConversationHistoryResponseEvent { +pub struct ConversationPathResponseEvent { pub conversation_id: ConversationId, - pub entries: Vec, + pub path: PathBuf, } #[derive(Debug, Clone, Deserialize, Serialize, TS)] diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index 039c9403da..4445c9391a 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -1,9 +1,11 @@ +use std::path::PathBuf; + use crate::app::App; use crate::backtrack_helpers; use crate::pager_overlay::Overlay; use crate::tui; use crate::tui::TuiEvent; -use codex_core::protocol::ConversationHistoryResponseEvent; +use codex_core::protocol::ConversationPathResponseEvent; use codex_protocol::mcp_protocol::ConversationId; use color_eyre::eyre::Result; use crossterm::event::KeyCode; @@ -98,7 +100,7 @@ impl App { ) { self.backtrack.pending = Some((base_id, drop_last_messages, prefill)); self.app_event_tx.send(crate::app_event::AppEvent::CodexOp( - codex_core::protocol::Op::GetHistory, + codex_core::protocol::Op::GetPath, )); } @@ -265,7 +267,7 @@ impl App { pub(crate) async fn on_conversation_history_for_backtrack( &mut self, tui: &mut tui::Tui, - ev: ConversationHistoryResponseEvent, + ev: ConversationPathResponseEvent, ) -> Result<()> { if let Some((base_id, _, _)) = self.backtrack.pending.as_ref() && ev.conversation_id == *base_id @@ -281,14 +283,14 @@ impl App { async fn fork_and_switch_to_new_conversation( &mut self, tui: &mut tui::Tui, - ev: ConversationHistoryResponseEvent, + ev: ConversationPathResponseEvent, drop_count: usize, prefill: String, ) { let cfg = self.chat_widget.config_ref().clone(); // Perform the fork via a thin wrapper for clarity/testability. let result = self - .perform_fork(ev.entries.clone(), drop_count, cfg.clone()) + .perform_fork(ev.path.clone(), drop_count, cfg.clone()) .await; match result { Ok(new_conv) => { @@ -301,13 +303,11 @@ impl App { /// Thin wrapper around ConversationManager::fork_conversation. async fn perform_fork( &self, - entries: Vec, + path: PathBuf, drop_count: usize, cfg: codex_core::config::Config, ) -> codex_core::error::Result { - self.server - .fork_conversation(entries, drop_count, cfg) - .await + self.server.fork_conversation(drop_count, cfg, path).await } /// Install a forked conversation into the ChatWidget and update UI to reflect selection. diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 1896f0a46e..cbd6f3d8f0 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -1,4 +1,4 @@ -use codex_core::protocol::ConversationHistoryResponseEvent; +use codex_core::protocol::ConversationPathResponseEvent; use codex_core::protocol::Event; use codex_file_search::FileMatch; @@ -58,5 +58,5 @@ pub(crate) enum AppEvent { UpdateSandboxPolicy(SandboxPolicy), /// Forwarded conversation history snapshot from the current conversation. - ConversationHistory(ConversationHistoryResponseEvent), + ConversationHistory(ConversationPathResponseEvent), } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index dca821f255..1ccfddea80 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -1083,7 +1083,7 @@ impl ChatWidget { self.on_user_message_event(ev); } } - EventMsg::ConversationHistory(ev) => { + EventMsg::ConversationPath(ev) => { self.app_event_tx .send(crate::app_event::AppEvent::ConversationHistory(ev)); }