diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 34e7932053..18a4d04ad1 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -711,6 +711,7 @@ dependencies = [ "mime_guess", "openssl-sys", "os_info", + "portable-pty", "predicates", "pretty_assertions", "rand 0.9.2", @@ -1443,6 +1444,12 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "downcast-rs" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75b325c5dbd37f80359721ad39aca5a29fb04c89279657cffdda8736d0c0b9d2" + [[package]] name = "dupe" version = "0.9.1" @@ -1688,6 +1695,17 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "filedescriptor" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e40758ed24c9b2eeb76c35fb0aebc66c626084edd827e07e1552279814c6682d" +dependencies = [ + "libc", + "thiserror 1.0.69", + "winapi", +] + [[package]] name = "fixedbitset" version = "0.4.2" @@ -3345,6 +3363,27 @@ dependencies = [ "portable-atomic", ] +[[package]] +name = "portable-pty" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4a596a2b3d2752d94f51fac2d4a96737b8705dddd311a32b9af47211f08671e" +dependencies = [ + "anyhow", + "bitflags 1.3.2", + "downcast-rs", + "filedescriptor", + "lazy_static", + "libc", + "log", + "nix", + "serial2", + "shared_library", + "shell-words", + "winapi", + "winreg", +] + [[package]] name = "potential_utf" version = "0.1.2" @@ -4272,6 +4311,17 @@ dependencies = [ "syn 2.0.104", ] +[[package]] +name = "serial2" +version = "0.2.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26e1e5956803a69ddd72ce2de337b577898801528749565def03515f82bad5bb" +dependencies = [ + "cfg-if", + "libc", + "winapi", +] + [[package]] name = "sha1" version = "0.10.6" @@ -4303,6 +4353,22 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "shared_library" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a9e7e0f2bfae24d8a5b5a66c5b257a83c7412304311512a0c054cd5e619da11" +dependencies = [ + "lazy_static", + "libc", +] + +[[package]] +name = "shell-words" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24188a676b6ae68c3b2cb3a01be17fbf7240ce009799bb56d5b1409051e78fde" + [[package]] name = "shlex" version = "1.3.0" @@ -6025,6 +6091,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "winreg" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80d0f4e272c85def139476380b12f9ac60926689dd2e01d4923222f40580869d" +dependencies = [ + "winapi", +] + [[package]] name = "winsafe" version = "0.0.19" diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 56815ba03c..2f2fa7cbad 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -28,6 +28,7 @@ libc = "0.2.175" mcp-types = { path = "../mcp-types" } mime_guess = "2.0" os_info = "3.12.0" +portable-pty = "0.9.0" rand = "0.9" regex-lite = "0.1.6" reqwest = { version = "0.12", features = ["json", "stream"] } diff --git a/codex-rs/core/src/chat_completions.rs b/codex-rs/core/src/chat_completions.rs index f3ed34c6dc..dc8aa9d7c3 100644 --- a/codex-rs/core/src/chat_completions.rs +++ b/codex-rs/core/src/chat_completions.rs @@ -34,11 +34,12 @@ pub(crate) async fn stream_chat_completions( model_family: &ModelFamily, client: &reqwest::Client, provider: &ModelProviderInfo, + use_streamable_shell_tool: bool, ) -> Result { // Build messages array let mut messages = Vec::::new(); - let full_instructions = prompt.get_full_instructions(model_family); + let full_instructions = prompt.get_full_instructions(model_family, use_streamable_shell_tool); messages.push(json!({"role": "system", "content": full_instructions})); let input = prompt.get_formatted_input(); diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 5534e11f36..27bcb3df41 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -101,6 +101,7 @@ impl ModelClient { &self.config.model_family, &self.client, &self.provider, + self.config.use_experimental_streamable_shell_tool, ) .await?; @@ -146,7 +147,10 @@ impl ModelClient { let store = prompt.store && auth_mode != Some(AuthMode::ChatGPT); - let full_instructions = prompt.get_full_instructions(&self.config.model_family); + let full_instructions = prompt.get_full_instructions( + &self.config.model_family, + self.config.use_experimental_streamable_shell_tool, + ); let tools_json = create_tools_json_for_responses_api(&prompt.tools)?; let reasoning = create_reasoning_param_for_request( &self.config.model_family, diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index 15b8ea89c0..bb342b5e70 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -41,14 +41,30 @@ pub struct Prompt { } impl Prompt { - pub(crate) fn get_full_instructions(&self, model: &ModelFamily) -> Cow<'_, str> { + pub(crate) fn get_full_instructions( + &self, + model: &ModelFamily, + use_streamable_shell_tool: bool, + ) -> Cow<'_, str> { let base = self - .base_instructions_override - .as_deref() - .unwrap_or(BASE_INSTRUCTIONS); - let mut sections: Vec<&str> = vec![base]; + .base_instructions_override.clone() + .unwrap_or_else(|| + if use_streamable_shell_tool { + BASE_INSTRUCTIONS.replace("SHELL_COMMAND_INSTRUCTIONS_HERE", "Do NOT use `shell`. Use only `exec_command` and `write_stdin`.") + } else { + BASE_INSTRUCTIONS.replace("SHELL_COMMAND_INSTRUCTIONS_HERE", r#"When using the shell, you must adhere to the following guidelines: + +- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.) +- Read files in chunks with a max chunk size of 250 lines. Do not use python scripts to attempt to output larger chunks of a file. Command line output will be truncated after 10 kilobytes or 256 lines of output, regardless of the command used. +"#) + } + ); + + // Do NOT use `shell`. Use only `exec_command` and `write_stdin`. + + let mut sections: Vec = vec![base]; if model.needs_special_apply_patch_instructions { - sections.push(APPLY_PATCH_TOOL_INSTRUCTIONS); + sections.push(APPLY_PATCH_TOOL_INSTRUCTIONS.to_owned()); } Cow::Owned(sections.join("\n")) } @@ -148,7 +164,8 @@ mod tests { }; let expected = format!("{BASE_INSTRUCTIONS}\n{APPLY_PATCH_TOOL_INSTRUCTIONS}"); let model_family = find_family_for_model("gpt-4.1").expect("known model slug"); - let full = prompt.get_full_instructions(&model_family); + let full = + prompt.get_full_instructions(&model_family, /*use_streamable_shell_tool*/ false); assert_eq!(full, expected); } } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 7d616f96e0..01acd5c6e5 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -52,6 +52,11 @@ use crate::exec::SandboxType; use crate::exec::StdoutStream; use crate::exec::StreamOutput; use crate::exec::process_exec_tool_call; +use crate::exec_command::EXEC_COMMAND_TOOL_NAME; +use crate::exec_command::ExecCommandParams; +use crate::exec_command::SESSION_MANAGER; +use crate::exec_command::WRITE_STDIN_TOOL_NAME; +use crate::exec_command::WriteStdinParams; use crate::exec_env::create_env; use crate::mcp_connection_manager::McpConnectionManager; use crate::mcp_tool_call::handle_mcp_tool_call; @@ -481,6 +486,7 @@ impl Session { sandbox_policy.clone(), config.include_plan_tool, config.include_apply_patch_tool, + config.use_experimental_streamable_shell_tool, ), user_instructions, base_instructions, @@ -1061,6 +1067,7 @@ async fn submission_loop( new_sandbox_policy.clone(), config.include_plan_tool, config.include_apply_patch_tool, + config.use_experimental_streamable_shell_tool, ); let new_turn_context = TurnContext { @@ -1139,6 +1146,7 @@ async fn submission_loop( sandbox_policy.clone(), config.include_plan_tool, config.include_apply_patch_tool, + config.use_experimental_streamable_shell_tool, ), user_instructions: turn_context.user_instructions.clone(), base_instructions: turn_context.base_instructions.clone(), @@ -1999,6 +2007,41 @@ async fn handle_function_call( .await } "update_plan" => handle_update_plan(sess, arguments, sub_id, call_id).await, + EXEC_COMMAND_TOOL_NAME => { + // TODO(mbolin): Sandbox check. + let exec_params = match serde_json::from_str::(&arguments) { + Ok(params) => params, + Err(e) => { + return ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: format!("failed to parse function arguments: {e}"), + success: Some(false), + }, + }; + } + }; + SESSION_MANAGER + .handle_exec_command_request(call_id, exec_params) + .await + } + WRITE_STDIN_TOOL_NAME => { + let write_stdin_params = match serde_json::from_str::(&arguments) { + Ok(params) => params, + Err(e) => { + return ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: format!("failed to parse function arguments: {e}"), + success: Some(false), + }, + }; + } + }; + SESSION_MANAGER + .handle_write_stdin_request(call_id, write_stdin_params) + .await + } _ => { match sess.mcp_connection_manager.parse_tool_name(&name) { Some((server, tool_name)) => { diff --git a/codex-rs/core/src/config.rs b/codex-rs/core/src/config.rs index 5b6a1ed265..241f80eb7b 100644 --- a/codex-rs/core/src/config.rs +++ b/codex-rs/core/src/config.rs @@ -169,6 +169,8 @@ pub struct Config { /// If set to `true`, the API key will be signed with the `originator` header. pub preferred_auth_method: AuthMode, + + pub use_experimental_streamable_shell_tool: bool, } impl Config { @@ -454,6 +456,8 @@ pub struct ConfigToml { /// Experimental path to a file whose contents replace the built-in BASE_INSTRUCTIONS. pub experimental_instructions_file: Option, + pub experimental_use_exec_command_tool: Option, + /// The value for the `originator` header included with Responses API requests. pub responses_originator_header_internal_override: Option, @@ -729,6 +733,9 @@ impl Config { include_apply_patch_tool: include_apply_patch_tool_val, responses_originator_header, preferred_auth_method: cfg.preferred_auth_method.unwrap_or(AuthMode::ChatGPT), + use_experimental_streamable_shell_tool: cfg + .experimental_use_exec_command_tool + .unwrap_or(false), }; Ok(config) } @@ -1094,6 +1101,7 @@ disable_response_storage = true include_apply_patch_tool: false, responses_originator_header: "codex_cli_rs".to_string(), preferred_auth_method: AuthMode::ChatGPT, + use_experimental_streamable_shell_tool: false, }, o3_profile_config ); @@ -1147,6 +1155,7 @@ disable_response_storage = true include_apply_patch_tool: false, responses_originator_header: "codex_cli_rs".to_string(), preferred_auth_method: AuthMode::ChatGPT, + use_experimental_streamable_shell_tool: false, }; assert_eq!(expected_gpt3_profile_config, gpt3_profile_config); @@ -1215,6 +1224,7 @@ disable_response_storage = true include_apply_patch_tool: false, responses_originator_header: "codex_cli_rs".to_string(), preferred_auth_method: AuthMode::ChatGPT, + use_experimental_streamable_shell_tool: false, }; assert_eq!(expected_zdr_profile_config, zdr_profile_config); diff --git a/codex-rs/core/src/exec_command/exec_command_params.rs b/codex-rs/core/src/exec_command/exec_command_params.rs new file mode 100644 index 0000000000..11a3fd4596 --- /dev/null +++ b/codex-rs/core/src/exec_command/exec_command_params.rs @@ -0,0 +1,57 @@ +use serde::Deserialize; +use serde::Serialize; + +use crate::exec_command::session_id::SessionId; + +#[derive(Debug, Clone, Deserialize)] +pub struct ExecCommandParams { + pub(crate) cmd: String, + + #[serde(default = "default_yield_time")] + pub(crate) yield_time_ms: u64, + + #[serde(default = "max_output_tokens")] + pub(crate) max_output_tokens: u64, + + #[serde(default = "default_shell")] + pub(crate) shell: String, + + #[serde(default = "default_login")] + pub(crate) login: bool, +} + +fn default_yield_time() -> u64 { + 10_000 +} + +fn max_output_tokens() -> u64 { + 10_000 +} + +fn default_login() -> bool { + true +} + +fn default_shell() -> String { + "/bin/bash".to_string() +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct WriteStdinParams { + pub(crate) session_id: SessionId, + pub(crate) chars: String, + + #[serde(default = "write_stdin_default_yield_time_ms")] + pub(crate) yield_time_ms: u64, + + #[serde(default = "write_stdin_default_max_output_tokens")] + pub(crate) max_output_tokens: u64, +} + +fn write_stdin_default_yield_time_ms() -> u64 { + 250 +} + +fn write_stdin_default_max_output_tokens() -> u64 { + 10_000 +} diff --git a/codex-rs/core/src/exec_command/exec_command_session.rs b/codex-rs/core/src/exec_command/exec_command_session.rs new file mode 100644 index 0000000000..131ab84d96 --- /dev/null +++ b/codex-rs/core/src/exec_command/exec_command_session.rs @@ -0,0 +1,54 @@ +use std::sync::Arc; + +use tokio::sync::Mutex; +use tokio::sync::mpsc; + +use crate::exec_command::session_id::SessionId; + +#[allow(dead_code)] +#[derive(Debug)] +pub(crate) struct ExecCommandSession { + pub(crate) id: SessionId, + /// Queue for writing bytes to the process stdin (PTY master write side). + writer_tx: mpsc::Sender>, + /// Stream of output chunks read from the PTY. Wrapped in Mutex so callers can + /// `await` receiving without needing `&mut self`. + output_rx: Arc>>>, +} + +#[allow(dead_code)] +impl ExecCommandSession { + pub(crate) fn new( + id: SessionId, + writer_tx: mpsc::Sender>, + output_rx: mpsc::Receiver>, + ) -> Self { + Self { + id, + writer_tx, + output_rx: Arc::new(Mutex::new(output_rx)), + } + } + + /// Enqueue bytes to be written to the process stdin (PTY master). + pub(crate) async fn write_stdin(&self, bytes: impl AsRef<[u8]>) -> anyhow::Result<()> { + self.writer_tx + .send(bytes.as_ref().to_vec()) + .await + .map_err(|e| anyhow::anyhow!("failed to send to writer: {e}")) + } + + /// Receive the next chunk of output from the process. Returns `None` when the + /// output stream is closed (process exited or reader finished). + pub(crate) async fn recv_output_chunk(&self) -> Option> { + self.output_rx.lock().await.recv().await + } + + pub(crate) fn writer_sender(&self) -> mpsc::Sender> { + self.writer_tx.clone() + } + + pub(crate) fn output_receiver(&self) -> Arc>>> { + self.output_rx.clone() + } +} diff --git a/codex-rs/core/src/exec_command/mod.rs b/codex-rs/core/src/exec_command/mod.rs new file mode 100644 index 0000000000..d48748f633 --- /dev/null +++ b/codex-rs/core/src/exec_command/mod.rs @@ -0,0 +1,13 @@ +mod exec_command_params; +mod exec_command_session; +mod responses_api; +mod session_id; +mod session_manager; + +pub use exec_command_params::ExecCommandParams; +pub use exec_command_params::WriteStdinParams; +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; +pub use responses_api::create_write_stdin_tool_for_responses_api; +pub use session_manager::SESSION_MANAGER; diff --git a/codex-rs/core/src/exec_command/responses_api.rs b/codex-rs/core/src/exec_command/responses_api.rs new file mode 100644 index 0000000000..c07684c326 --- /dev/null +++ b/codex-rs/core/src/exec_command/responses_api.rs @@ -0,0 +1,97 @@ +use std::collections::BTreeMap; + +use crate::openai_tools::JsonSchema; +use crate::openai_tools::ResponsesApiTool; + +pub const EXEC_COMMAND_TOOL_NAME: &str = "exec_command"; +pub const WRITE_STDIN_TOOL_NAME: &str = "write_stdin"; + +pub fn create_exec_command_tool_for_responses_api() -> ResponsesApiTool { + let mut properties = BTreeMap::::new(); + properties.insert( + "cmd".to_string(), + JsonSchema::String { + description: Some("The shell command to execute.".to_string()), + }, + ); + properties.insert( + "yield_time_ms".to_string(), + JsonSchema::Number { + description: Some("The maximum time in milliseconds to wait for output.".to_string()), + }, + ); + properties.insert( + "max_output_tokens".to_string(), + JsonSchema::Number { + description: Some("The maximum number of tokens to output.".to_string()), + }, + ); + properties.insert( + "shell".to_string(), + JsonSchema::String { + description: Some("The shell to use. Defaults to \"/bin/bash\".".to_string()), + }, + ); + properties.insert( + "login".to_string(), + JsonSchema::Boolean { + description: Some( + "Whether to run the command as a login shell. Defaults to true.".to_string(), + ), + }, + ); + + ResponsesApiTool { + name: EXEC_COMMAND_TOOL_NAME.to_owned(), + description: r#"Execute shell commands on the local machine with streaming output."# + .to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["cmd".to_string()]), + additional_properties: Some(false), + }, + } +} + +pub fn create_write_stdin_tool_for_responses_api() -> ResponsesApiTool { + let mut properties = BTreeMap::::new(); + properties.insert( + "session_id".to_string(), + JsonSchema::String { + description: Some("The ID of the exec_command session.".to_string()), + }, + ); + properties.insert( + "chars".to_string(), + JsonSchema::String { + description: Some("The characters to write to stdin.".to_string()), + }, + ); + properties.insert( + "yield_time_ms".to_string(), + JsonSchema::Number { + description: Some( + "The maximum time in milliseconds to wait for output after writing.".to_string(), + ), + }, + ); + properties.insert( + "max_output_tokens".to_string(), + JsonSchema::Number { + description: Some("The maximum number of tokens to output.".to_string()), + }, + ); + + ResponsesApiTool { + name: WRITE_STDIN_TOOL_NAME.to_owned(), + description: r#"Write characters to the stdin of an existing exec_command session."# + .to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["session_id".to_string(), "chars".to_string()]), + additional_properties: Some(false), + }, + } +} diff --git a/codex-rs/core/src/exec_command/session_id.rs b/codex-rs/core/src/exec_command/session_id.rs new file mode 100644 index 0000000000..c97c5d5440 --- /dev/null +++ b/codex-rs/core/src/exec_command/session_id.rs @@ -0,0 +1,5 @@ +use serde::Deserialize; +use serde::Serialize; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub(crate) struct SessionId(pub u32); diff --git a/codex-rs/core/src/exec_command/session_manager.rs b/codex-rs/core/src/exec_command/session_manager.rs new file mode 100644 index 0000000000..60f0d284ad --- /dev/null +++ b/codex-rs/core/src/exec_command/session_manager.rs @@ -0,0 +1,302 @@ +use std::collections::HashMap; +use std::io::Read; +use std::sync::Arc; +use std::sync::LazyLock; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicU32; + +use portable_pty::CommandBuilder; +use portable_pty::PtySize; +use portable_pty::native_pty_system; +use serde_json::json; +use tokio::sync::Mutex; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::time::Duration; +use tokio::time::Instant; +use tokio::time::timeout; + +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::models::FunctionCallOutputPayload; +use crate::models::ResponseInputItem; + +pub static SESSION_MANAGER: LazyLock = LazyLock::new(SessionManager::default); + +#[derive(Debug, Default)] +pub struct SessionManager { + next_session_id: AtomicU32, + sessions: Mutex>, +} + +impl SessionManager { + /// Processes the request and is required to send a response via `outgoing`. + pub async fn handle_exec_command_request( + &self, + call_id: String, + params: ExecCommandParams, + ) -> ResponseInputItem { + // Allocate a session id. + let session_id = SessionId( + self.next_session_id + .fetch_add(1, std::sync::atomic::Ordering::SeqCst), + ); + + let result = create_exec_command_session(session_id, params.clone()).await; + + match result { + Ok((session, mut exit_rx)) => { + // Insert into session map. + let output_receiver = session.output_receiver(); + self.sessions.lock().await.insert(session_id, session); + + // Collect output until either timeout expires or process exits. + // Cap by assuming 4 bytes per token (TODO: use a real tokenizer). + let cap_bytes_u64 = params.max_output_tokens.saturating_mul(4); + let cap_bytes: usize = cap_bytes_u64.min(usize::MAX as u64) as usize; + let cap_hint = cap_bytes.clamp(1024, 8192); + let mut collected: Vec = Vec::with_capacity(cap_hint); + + let deadline = Instant::now() + Duration::from_millis(params.yield_time_ms); + let mut exit_code: Option = None; + + loop { + if Instant::now() >= deadline { + break; + } + let remaining = deadline.saturating_duration_since(Instant::now()); + tokio::select! { + biased; + exit = &mut exit_rx => { + exit_code = exit.ok(); + // Small grace period to pull remaining buffered output + let grace_deadline = Instant::now() + Duration::from_millis(25); + while Instant::now() < grace_deadline { + let recv_next = async { + let mut rx = output_receiver.lock().await; + rx.recv().await + }; + if let Ok(Some(chunk)) = timeout(Duration::from_millis(1), recv_next).await { + let available = cap_bytes.saturating_sub(collected.len()); + if available == 0 { break; } + let take = available.min(chunk.len()); + collected.extend_from_slice(&chunk[..take]); + } else { + break; + } + } + break; + } + chunk = timeout(remaining, async { + let mut rx = output_receiver.lock().await; + rx.recv().await + }) => { + match chunk { + Ok(Some(chunk)) => { + let available = cap_bytes.saturating_sub(collected.len()); + if available == 0 { /* keep draining, but don't store */ } + else { + let take = available.min(chunk.len()); + collected.extend_from_slice(&chunk[..take]); + } + } + Ok(None) => { break; } + Err(_) => { break; } + } + } + } + } + + let text = String::from_utf8_lossy(&collected).to_string(); + let mut structured = json!({ "sessionId": session_id }); + if let Some(code) = exit_code { + structured["exitCode"] = json!(code); + } + + ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: text, + success: Some(true), + }, + } + } + Err(err) => ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: format!("failed to start exec session: {err}"), + success: Some(false), + }, + }, + } + } + + /// Write characters to a session's stdin and collect combined output for up to `yield_time_ms`. + pub async fn handle_write_stdin_request( + &self, + call_id: String, + params: WriteStdinParams, + ) -> ResponseInputItem { + let WriteStdinParams { + session_id, + chars, + yield_time_ms, + max_output_tokens, + } = params; + + // Grab handles without holding the sessions lock across await points. + let (writer_tx, output_rx) = { + let sessions = self.sessions.lock().await; + match sessions.get(&session_id) { + Some(session) => (session.writer_sender(), session.output_receiver()), + None => { + return ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: format!("unknown session id {}", session_id.0), + success: Some(false), + }, + }; + } + } + }; + + // Write stdin if provided. + if !chars.is_empty() && writer_tx.send(chars.into_bytes()).await.is_err() { + return ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: "failed to write to stdin".to_string(), + success: Some(false), + }, + }; + } + + // Collect output up to yield_time_ms, truncating to max_output_tokens bytes. + let mut collected: Vec = Vec::with_capacity(4096); + let deadline = Instant::now() + Duration::from_millis(yield_time_ms); + loop { + let now = Instant::now(); + if now >= deadline { + break; + } + let remaining = deadline - now; + match timeout(remaining, output_rx.lock().await.recv()).await { + Ok(Some(chunk)) => { + // Respect token/byte limit; keep draining but drop once full. + let available = + max_output_tokens.saturating_sub(collected.len() as u64) as usize; + if available > 0 { + let take = available.min(chunk.len()); + collected.extend_from_slice(&chunk[..take]); + } + // Continue loop to drain further within time. + } + Ok(None) => break, // channel closed + Err(_) => break, // timeout + } + } + + // Return text output as a CallToolResult + let text = String::from_utf8_lossy(&collected).to_string(); + ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: text, + success: Some(true), + }, + } + } +} + +/// Spawn PTY and child process per spawn_exec_command_session logic. +async fn create_exec_command_session( + session_id: SessionId, + params: ExecCommandParams, +) -> anyhow::Result<(ExecCommandSession, oneshot::Receiver)> { + let ExecCommandParams { + cmd, + yield_time_ms: _, + max_output_tokens: _, + shell, + login, + } = params; + + // Use the native pty implementation for the system + let pty_system = native_pty_system(); + + // Create a new pty + let pair = pty_system.openpty(PtySize { + rows: 24, + cols: 80, + pixel_width: 0, + pixel_height: 0, + })?; + + // Spawn a shell into the pty + let mut command_builder = CommandBuilder::new(shell); + let shell_mode_opt = if login { "-lc" } else { "-c" }; + command_builder.arg(shell_mode_opt); + command_builder.arg(cmd); + + let mut child = pair.slave.spawn_command(command_builder)?; + + // Channel to forward write requests to the PTY writer. + let (writer_tx, mut writer_rx) = mpsc::channel::>(128); + // Channel for streaming PTY output to readers. + let (output_tx, output_rx) = mpsc::channel::>(256); + + // Reader task: drain PTY and forward chunks to output channel. + let mut reader = pair.master.try_clone_reader()?; + let output_tx_clone = output_tx.clone(); + tokio::task::spawn_blocking(move || { + let mut buf = [0u8; 8192]; + loop { + match reader.read(&mut buf) { + Ok(0) => break, // EOF + Ok(n) => { + // Forward; block if receiver is slow to avoid dropping output. + let _ = output_tx_clone.blocking_send(buf[..n].to_vec()); + } + Err(_) => break, + } + } + }); + + // Writer task: apply stdin writes to the PTY writer. + let writer = pair.master.take_writer()?; + let writer = Arc::new(StdMutex::new(writer)); + tokio::spawn({ + let writer = writer.clone(); + async move { + while let Some(bytes) = writer_rx.recv().await { + let writer = writer.clone(); + // Perform blocking write on a blocking thread. + 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; + } + } + }); + + // Keep the child alive until it exits, then signal exit code. + let (exit_tx, exit_rx) = oneshot::channel::(); + tokio::task::spawn_blocking(move || { + let code = match child.wait() { + Ok(status) => status.exit_code() as i32, + Err(_) => -1, + }; + let _ = exit_tx.send(code); + }); + + // Create and store the session with channels. + let session = ExecCommandSession::new(session_id, writer_tx, output_rx); + Ok((session, exit_rx)) +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 40ccbf6769..ef066c1a0a 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -20,6 +20,7 @@ mod conversation_history; mod environment_context; pub mod error; pub mod exec; +mod exec_command; pub mod exec_env; mod flags; pub mod git_info; diff --git a/codex-rs/core/src/openai_tools.rs b/codex-rs/core/src/openai_tools.rs index dd5eb12516..b5d82efbb5 100644 --- a/codex-rs/core/src/openai_tools.rs +++ b/codex-rs/core/src/openai_tools.rs @@ -37,6 +37,7 @@ pub enum ConfigShellToolType { DefaultShell, ShellWithRequest { sandbox_policy: SandboxPolicy }, LocalShell, + StreamableShell, } #[derive(Debug, Clone)] @@ -53,13 +54,16 @@ impl ToolsConfig { sandbox_policy: SandboxPolicy, include_plan_tool: bool, include_apply_patch_tool: bool, + use_streamable_shell_tool: bool, ) -> Self { - let mut shell_type = if model_family.uses_local_shell_tool { + let mut shell_type = if use_streamable_shell_tool { + ConfigShellToolType::StreamableShell + } else if model_family.uses_local_shell_tool { ConfigShellToolType::LocalShell } else { ConfigShellToolType::DefaultShell }; - if matches!(approval_policy, AskForApproval::OnRequest) { + if matches!(approval_policy, AskForApproval::OnRequest) && !use_streamable_shell_tool { shell_type = ConfigShellToolType::ShellWithRequest { sandbox_policy: sandbox_policy.clone(), } @@ -533,6 +537,14 @@ pub(crate) fn get_openai_tools( 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.plan_tool { @@ -597,6 +609,7 @@ mod tests { SandboxPolicy::ReadOnly, true, model_family.uses_apply_patch_tool, + /*use_experimental_streamable_shell_tool*/ false, ); let tools = get_openai_tools(&config, Some(HashMap::new())); @@ -612,6 +625,7 @@ mod tests { SandboxPolicy::ReadOnly, true, model_family.uses_apply_patch_tool, + /*use_experimental_streamable_shell_tool*/ false, ); let tools = get_openai_tools(&config, Some(HashMap::new())); @@ -627,6 +641,7 @@ mod tests { SandboxPolicy::ReadOnly, false, model_family.uses_apply_patch_tool, + /*use_experimental_streamable_shell_tool*/ false, ); let tools = get_openai_tools( &config, @@ -721,6 +736,7 @@ mod tests { SandboxPolicy::ReadOnly, false, model_family.uses_apply_patch_tool, + /*use_experimental_streamable_shell_tool*/ false, ); let tools = get_openai_tools( @@ -777,6 +793,7 @@ mod tests { SandboxPolicy::ReadOnly, false, model_family.uses_apply_patch_tool, + /*use_experimental_streamable_shell_tool*/ false, ); let tools = get_openai_tools( @@ -828,6 +845,7 @@ mod tests { SandboxPolicy::ReadOnly, false, model_family.uses_apply_patch_tool, + /*use_experimental_streamable_shell_tool*/ false, ); let tools = get_openai_tools( @@ -882,6 +900,7 @@ mod tests { SandboxPolicy::ReadOnly, false, model_family.uses_apply_patch_tool, + /*use_experimental_streamable_shell_tool*/ false, ); let tools = get_openai_tools(