openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-17 12:23:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: StreamableShell with exec_command and write_stdin tools

Browse Source
This commit is contained in:
Michael Bolin
2025-08-21 19:42:22 -07:00
parent 750ca9e21d
commit f7cdaaa246
15 changed files with 804 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
codex-rs/Cargo.lock generated
View File

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

1
codex-rs/core/Cargo.toml
View File

@@ -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"] }

3
codex-rs/core/src/chat_completions.rs
View File

@@ -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<ResponseStream> {
// Build messages array
let mut messages = Vec::<serde_json::Value>::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();

6
codex-rs/core/src/client.rs
View File

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

31
codex-rs/core/src/client_common.rs
View File

@@ -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<String> = 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);
}
}

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

@@ -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,46 @@ 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::<ExecCommandParams>(&arguments) {
Ok(params) => params,
Err(e) => {
return ResponseInputItem::FunctionCallOutput {
call_id,
output: FunctionCallOutputPayload {
content: format!("failed to parse function arguments: {e}"),
success: Some(false),
},
};
}
};
let result = SESSION_MANAGER
.handle_exec_command_request(exec_params)
.await;
let function_call_output: FunctionCallOutputPayload = result.into();
ResponseInputItem::FunctionCallOutput {
call_id,
output: function_call_output,
}
}
WRITE_STDIN_TOOL_NAME => {
let write_stdin_params = match serde_json::from_str::<WriteStdinParams>(&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)) => {

10
codex-rs/core/src/config.rs
View File

@@ -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<PathBuf>,
pub experimental_use_exec_command_tool: Option<bool>,
/// The value for the `originator` header included with Responses API requests.
pub responses_originator_header_internal_override: Option<String>,
@@ -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);

57
codex-rs/core/src/exec_command/exec_command_params.rs Normal file
View File

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

54
codex-rs/core/src/exec_command/exec_command_session.rs Normal file
View File

@@ -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<Vec<u8>>,
/// Stream of output chunks read from the PTY. Wrapped in Mutex so callers can
/// `await` receiving without needing `&mut self`.
output_rx: Arc<Mutex<mpsc::Receiver<Vec<u8>>>>,
}
#[allow(dead_code)]
impl ExecCommandSession {
pub(crate) fn new(
id: SessionId,
writer_tx: mpsc::Sender<Vec<u8>>,
output_rx: mpsc::Receiver<Vec<u8>>,
) -> 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<Vec<u8>> {
self.output_rx.lock().await.recv().await
}
pub(crate) fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> {
self.writer_tx.clone()
}
pub(crate) fn output_receiver(&self) -> Arc<Mutex<mpsc::Receiver<Vec<u8>>>> {
self.output_rx.clone()
}
}

13
codex-rs/core/src/exec_command/mod.rs Normal file
View File

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

97
codex-rs/core/src/exec_command/responses_api.rs Normal file
View File

@@ -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::<String, JsonSchema>::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::<String, JsonSchema>::new();
properties.insert(
"session_id".to_string(),
JsonSchema::Number {
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),
},
}
}

5
codex-rs/core/src/exec_command/session_id.rs Normal file
View File

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

391
codex-rs/core/src/exec_command/session_manager.rs Normal file
View File

@@ -0,0 +1,391 @@
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 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<SessionManager> = LazyLock::new(SessionManager::default);
#[derive(Debug, Default)]
pub struct SessionManager {
next_session_id: AtomicU32,
sessions: Mutex<HashMap<SessionId, ExecCommandSession>>,
}
pub struct ExecCommandOutput {
wall_time: Duration,
exit_status: ExitStatus,
original_token_count: Option<u64>,
output: String,
}
impl ExecCommandOutput {
fn to_text_output(&self) -> String {
let wall_time_secs = self.wall_time.as_secs_f32();
let termination_status = match self.exit_status {
ExitStatus::Exited(code) => format!("Process exited with code {code}"),
ExitStatus::Ongoing(session_id) => {
format!("Process running with session ID {}", session_id.0)
}
};
let truncation_status = match self.original_token_count {
Some(tokens) => {
format!("\nWarning: truncated output (original token count: {tokens})")
}
None => "".to_string(),
};
format!(
r#"Wall time: {wall_time_secs:.3} seconds
{termination_status}{truncation_status}
Output:
{output}"#,
output = self.output
)
}
}
pub enum ExitStatus {
Exited(i32),
Ongoing(SessionId),
}
impl From<Result<ExecCommandOutput, String>> for FunctionCallOutputPayload {
fn from(val: Result<ExecCommandOutput, String>) -> Self {
match val {
Ok(output) => FunctionCallOutputPayload {
content: output.to_text_output(),
success: Some(true),
},
Err(err) => FunctionCallOutputPayload {
content: err,
success: Some(false),
},
}
}
}
impl SessionManager {
/// Processes the request and is required to send a response via `outgoing`.
pub async fn handle_exec_command_request(
&self,
params: ExecCommandParams,
) -> Result<ExecCommandOutput, String> {
// Allocate a session id.
let session_id = SessionId(
self.next_session_id
.fetch_add(1, std::sync::atomic::Ordering::SeqCst),
);
let (session, mut exit_rx) = create_exec_command_session(session_id, params.clone())
.await
.map_err(|err| {
format!(
"failed to create exec command session for session id {}: {err}",
session_id.0
)
})?;
// 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<u8> = Vec::with_capacity(cap_hint);
let start_time = Instant::now();
let deadline = start_time + Duration::from_millis(params.yield_time_ms);
let mut exit_code: Option<i32> = 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 output = String::from_utf8_lossy(&collected).to_string();
let exit_status = if let Some(code) = exit_code {
ExitStatus::Exited(code)
} else {
ExitStatus::Ongoing(session_id)
};
// TODO(mbolin): Ensure max_output_tokens is honored correctly.
let original_token_count = None;
Ok(ExecCommandOutput {
wall_time: Instant::now().duration_since(start_time),
exit_status,
original_token_count,
output,
})
}
/// 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<u8> = 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<i32>)> {
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::<Vec<u8>>(128);
// Channel for streaming PTY output to readers.
let (output_tx, output_rx) = mpsc::channel::<Vec<u8>>(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::<i32>();
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))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::exec_command::session_id::SessionId;
#[test]
fn to_text_output_exited_no_truncation() {
let out = ExecCommandOutput {
wall_time: Duration::from_millis(1234),
exit_status: ExitStatus::Exited(0),
original_token_count: None,
output: "hello".to_string(),
};
let text = out.to_text_output();
let expected = r#"Wall time: 1.234 seconds
Process exited with code 0
Output:
hello"#;
assert_eq!(expected, text);
}
#[test]
fn to_text_output_ongoing_with_truncation() {
let out = ExecCommandOutput {
wall_time: Duration::from_millis(500),
exit_status: ExitStatus::Ongoing(SessionId(42)),
original_token_count: Some(1000),
output: "abc".to_string(),
};
let text = out.to_text_output();
let expected = r#"Wall time: 0.500 seconds
Process running with session ID 42
Warning: truncated output (original token count: 1000)
Output:
abc"#;
assert_eq!(expected, text);
}
}

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

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

23
codex-rs/core/src/openai_tools.rs
View File

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