diff --git a/codex-rs/core/src/features.rs b/codex-rs/core/src/features.rs index 69442815e7..12c440cf10 100644 --- a/codex-rs/core/src/features.rs +++ b/codex-rs/core/src/features.rs @@ -40,6 +40,8 @@ pub enum Feature { // Experimental /// Use the single unified PTY-backed exec tool. UnifiedExec, + /// Use the unified exec wrapper tool with named sessions. + UnifiedExecWrapper, /// Enable experimental RMCP features such as OAuth login. RmcpClient, /// Include the freeform apply_patch tool. @@ -293,6 +295,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Experimental, default_enabled: false, }, + FeatureSpec { + id: Feature::UnifiedExecWrapper, + key: "unified_exec_wrapper", + stage: Stage::Experimental, + default_enabled: false, + }, FeatureSpec { id: Feature::RmcpClient, key: "rmcp_client", diff --git a/codex-rs/core/src/shell.rs b/codex-rs/core/src/shell.rs index 2338f41cd4..c4cc2543ab 100644 --- a/codex-rs/core/src/shell.rs +++ b/codex-rs/core/src/shell.rs @@ -58,6 +58,28 @@ impl Shell { } } } + + pub fn login_command(&self) -> Vec { + self.command_with_login(true) + } + + pub fn non_login_command(&self) -> Vec { + self.command_with_login(false) + } + + fn command_with_login(&self, login_shell: bool) -> Vec { + let shell_path = self.shell_path.to_string_lossy().to_string(); + match self.shell_type { + ShellType::PowerShell | ShellType::Cmd => vec![shell_path], + ShellType::Zsh | ShellType::Bash | ShellType::Sh => { + if login_shell { + vec![shell_path, "-l".to_string()] + } else { + vec![shell_path] + } + } + } + } } #[cfg(unix)] @@ -450,6 +472,17 @@ mod tests { ); } + #[test] + fn shell_command_login_variants() { + let sh_shell = Shell { + shell_type: ShellType::Sh, + shell_path: PathBuf::from("/bin/sh"), + }; + + assert_eq!(sh_shell.login_command(), vec!["/bin/sh", "-l"]); + assert_eq!(sh_shell.non_login_command(), vec!["/bin/sh"]); + } + #[tokio::test] async fn test_current_shell_detects_zsh() { let shell = Command::new("sh") diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index f2500a413b..f0ccbc627b 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -58,6 +58,28 @@ struct WriteStdinArgs { max_output_tokens: Option, } +#[derive(Debug, Deserialize)] +struct NewSessionArgs { + session_name: String, + #[serde(default)] + workdir: Option, + #[serde(default = "default_write_stdin_yield_time_ms")] + yield_time_ms: u64, + #[serde(default)] + max_output_tokens: Option, +} + +#[derive(Debug, Deserialize)] +struct FeedCharsArgs { + session_name: String, + #[serde(default)] + chars: String, + #[serde(default = "default_write_stdin_yield_time_ms")] + yield_time_ms: u64, + #[serde(default)] + max_output_tokens: Option, +} + fn default_exec_yield_time_ms() -> u64 { 10000 } @@ -207,6 +229,87 @@ impl ToolHandler for UnifiedExecHandler { FunctionCallError::RespondToModel(format!("exec_command failed: {err:?}")) })? } + "new_session" => { + let args: NewSessionArgs = serde_json::from_str(&arguments).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to parse new_session arguments: {err:?}" + )) + })?; + + if args.session_name.trim().is_empty() { + return Err(FunctionCallError::RespondToModel( + "session_name must not be empty".to_string(), + )); + } + + let process_id = manager.allocate_process_id().await; + if let Err(err) = manager + .register_session_name(&args.session_name, &process_id) + .await + { + manager.release_process_id(&process_id).await; + return Err(FunctionCallError::RespondToModel(err.to_string())); + } + + let workdir = args + .workdir + .filter(|value| !value.is_empty()) + .map(|dir| context.turn.resolve_path(Some(dir))); + + let response = manager + .exec_command( + ExecCommandRequest { + command: session.user_shell().login_command(), + process_id, + yield_time_ms: args.yield_time_ms, + max_output_tokens: args.max_output_tokens, + workdir, + with_escalated_permissions: None, + justification: None, + }, + &context, + ) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!("exec_command failed: {err:?}")) + })?; + + if response.process_id.is_none() { + manager.clear_session_name(&args.session_name).await; + } + + response + } + "feed_chars" => { + let args: FeedCharsArgs = serde_json::from_str(&arguments).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to parse feed_chars arguments: {err:?}" + )) + })?; + + let Some(process_id) = manager + .process_id_for_session_name(&args.session_name) + .await + else { + return Err(FunctionCallError::RespondToModel(format!( + "Session '{}' not found", + args.session_name + ))); + }; + + manager + .write_stdin(WriteStdinRequest { + call_id: &call_id, + process_id: &process_id, + input: &args.chars, + yield_time_ms: args.yield_time_ms, + max_output_tokens: args.max_output_tokens, + }) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!("write_stdin failed: {err:?}")) + })? + } "write_stdin" => { let args: WriteStdinArgs = serde_json::from_str(&arguments).map_err(|err| { FunctionCallError::RespondToModel(format!( diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 89a71b0edc..e08fa34a4a 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -42,6 +42,8 @@ impl ToolsConfig { let shell_type = if !features.enabled(Feature::ShellTool) { ConfigShellToolType::Disabled + } else if features.enabled(Feature::UnifiedExecWrapper) { + ConfigShellToolType::UnifiedExecWrapper } else if features.enabled(Feature::UnifiedExec) { ConfigShellToolType::UnifiedExec } else { @@ -251,6 +253,98 @@ fn create_write_stdin_tool() -> ToolSpec { }) } +fn create_new_session_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "session_name".to_string(), + JsonSchema::String { + description: Some("Unique name for the session".to_string()), + }, + ); + properties.insert( + "workdir".to_string(), + JsonSchema::String { + description: Some( + "Optional working directory for the session; defaults to the turn cwd.".to_string(), + ), + }, + ); + properties.insert( + "yield_time_ms".to_string(), + JsonSchema::Number { + description: Some( + "How long to wait (in milliseconds) before returning the initial output." + .to_string(), + ), + }, + ); + properties.insert( + "max_output_tokens".to_string(), + JsonSchema::Number { + description: Some( + "Maximum number of tokens to return. Excess output will be truncated.".to_string(), + ), + }, + ); + + ToolSpec::Function(ResponsesApiTool { + name: "new_session".to_string(), + description: "Open a new interactive exec session in a container. Normally used for launching an interactive shell. Multiple sessions may be running at a time.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["session_name".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + +fn create_feed_chars_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "session_name".to_string(), + JsonSchema::String { + description: Some("Session to feed characters to".to_string()), + }, + ); + properties.insert( + "chars".to_string(), + JsonSchema::String { + description: Some("Characters to feed; may be empty".to_string()), + }, + ); + properties.insert( + "yield_time_ms".to_string(), + JsonSchema::Number { + description: Some( + "How long to wait (in milliseconds) for output before flushing STDOUT/STDERR." + .to_string(), + ), + }, + ); + properties.insert( + "max_output_tokens".to_string(), + JsonSchema::Number { + description: Some( + "Maximum number of tokens to return. Excess output will be truncated.".to_string(), + ), + }, + ); + + ToolSpec::Function(ResponsesApiTool { + name: "feed_chars".to_string(), + description: + "Feed characters to a session's STDIN, wait briefly, flush STDOUT/STDERR, and return the results." + .to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["session_name".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + fn create_shell_tool() -> ToolSpec { let mut properties = BTreeMap::new(); properties.insert( @@ -1013,6 +1107,12 @@ pub(crate) fn build_specs( builder.register_handler("exec_command", unified_exec_handler.clone()); builder.register_handler("write_stdin", unified_exec_handler); } + ConfigShellToolType::UnifiedExecWrapper => { + builder.push_spec(create_new_session_tool()); + builder.push_spec(create_feed_chars_tool()); + builder.register_handler("new_session", unified_exec_handler.clone()); + builder.register_handler("feed_chars", unified_exec_handler); + } ConfigShellToolType::Disabled => { // Do nothing. } @@ -1163,6 +1263,7 @@ mod tests { ConfigShellToolType::Default => Some("shell"), ConfigShellToolType::Local => Some("local_shell"), ConfigShellToolType::UnifiedExec => None, + ConfigShellToolType::UnifiedExecWrapper => Some("unified_exec_wrapper"), ConfigShellToolType::Disabled => None, ConfigShellToolType::ShellCommand => Some("shell_command"), } diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 34b62df342..4aa4aea23e 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -109,17 +109,43 @@ pub(crate) struct UnifiedExecSessionManager { pub(crate) struct SessionStore { sessions: HashMap, reserved_sessions_id: HashSet, + session_names: HashMap, } impl SessionStore { fn remove(&mut self, session_id: &str) -> Option { self.reserved_sessions_id.remove(session_id); + self.session_names.retain(|_, id| id != session_id); self.sessions.remove(session_id) } pub(crate) fn clear(&mut self) { self.reserved_sessions_id.clear(); self.sessions.clear(); + self.session_names.clear(); + } + + fn process_id_for_name(&self, session_name: &str) -> Option { + self.session_names.get(session_name).cloned() + } + + fn insert_session_name( + &mut self, + session_name: &str, + process_id: &str, + ) -> Result<(), UnifiedExecError> { + if self.session_names.contains_key(session_name) { + return Err(UnifiedExecError::create_session(format!( + "Session '{session_name}' already in use" + ))); + } + self.session_names + .insert(session_name.to_string(), process_id.to_string()); + Ok(()) + } + + fn clear_session_name(&mut self, session_name: &str) { + self.session_names.remove(session_name); } } @@ -456,4 +482,45 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn session_names_track_process_ids() { + let (session, _turn) = test_session_and_turn(); + let process_id = session + .services + .unified_exec_manager + .allocate_process_id() + .await; + + session + .services + .unified_exec_manager + .register_session_name("default", &process_id) + .await + .expect("session name reserved"); + + pretty_assertions::assert_eq!( + session + .services + .unified_exec_manager + .process_id_for_session_name("default") + .await, + Some(process_id.clone()) + ); + + session + .services + .unified_exec_manager + .release_process_id(&process_id) + .await; + + assert!( + session + .services + .unified_exec_manager + .process_id_for_session_name("default") + .await + .is_none() + ); + } } diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index 88d65ca142..95c7d423d8 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -116,6 +116,25 @@ impl UnifiedExecSessionManager { store.remove(process_id); } + pub(crate) async fn register_session_name( + &self, + session_name: &str, + process_id: &str, + ) -> Result<(), UnifiedExecError> { + let mut store = self.session_store.lock().await; + store.insert_session_name(session_name, process_id) + } + + pub(crate) async fn process_id_for_session_name(&self, session_name: &str) -> Option { + let store = self.session_store.lock().await; + store.process_id_for_name(session_name) + } + + pub(crate) async fn clear_session_name(&self, session_name: &str) { + let mut store = self.session_store.lock().await; + store.clear_session_name(session_name); + } + pub(crate) async fn exec_command( &self, request: ExecCommandRequest, diff --git a/codex-rs/core/tests/suite/tools.rs b/codex-rs/core/tests/suite/tools.rs index 1f743c422f..5865dc78a3 100644 --- a/codex-rs/core/tests/suite/tools.rs +++ b/codex-rs/core/tests/suite/tools.rs @@ -296,6 +296,7 @@ async fn collect_tools(use_unified_exec: bool) -> Result> { } else { config.features.disable(Feature::UnifiedExec); } + config.features.disable(Feature::UnifiedExecWrapper); }); let test = builder.build(&server).await?; diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 0804811a3f..6c81324cfb 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -108,6 +108,7 @@ pub enum ConfigShellToolType { Default, Local, UnifiedExec, + UnifiedExecWrapper, Disabled, ShellCommand, }