diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 7769f262a9e..39462a8e45c 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -55,6 +55,11 @@ struct WriteStdinArgs { max_output_tokens: Option, } +#[derive(Debug, Deserialize)] +struct TerminateSessionArgs { + session_id: i32, +} + fn default_exec_yield_time_ms() -> u64 { 10000 } @@ -137,7 +142,7 @@ impl ToolHandler for UnifiedExecHandler { codex_protocol::protocol::AskForApproval::OnRequest ) { - manager.release_process_id(&process_id).await; + manager.terminate_session(&process_id).await; return Err(FunctionCallError::RespondToModel(format!( "approval policy is {policy:?}; reject command — you cannot ask for escalated permissions if the approval policy is {policy:?}", policy = context.turn.approval_policy @@ -161,7 +166,7 @@ impl ToolHandler for UnifiedExecHandler { ) .await? { - manager.release_process_id(&process_id).await; + manager.terminate_session(&process_id).await; return Ok(output); } @@ -208,6 +213,32 @@ impl ToolHandler for UnifiedExecHandler { response } + "terminate_session" => { + let args: TerminateSessionArgs = + serde_json::from_str(&arguments).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to parse terminate_session arguments: {err:?}" + )) + })?; + let success = manager + .terminate_session(&args.session_id.to_string()) + .await; + UnifiedExecResponse { + event_call_id: context.call_id.clone(), + chunk_id: crate::unified_exec::generate_chunk_id(), + wall_time: std::time::Duration::ZERO, + output: if success { + "Session terminated.".to_string() + } else { + "Session not found.".to_string() + }, + raw_output: Vec::new(), + process_id: None, + exit_code: None, + original_token_count: None, + session_command: None, + } + } other => { return Err(FunctionCallError::RespondToModel(format!( "unsupported unified exec function {other}" diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 0a66b414039..b63f758b900 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -266,6 +266,27 @@ fn create_write_stdin_tool() -> ToolSpec { }) } +fn create_terminate_session_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "session_id".to_string(), + JsonSchema::Number { + description: Some("Identifier of the unified exec session to terminate.".to_string()), + }, + ); + + ToolSpec::Function(ResponsesApiTool { + name: "terminate_session".to_string(), + description: "Terminates a running unified exec session.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["session_id".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + fn create_shell_tool() -> ToolSpec { let properties = BTreeMap::from([ ( @@ -1138,8 +1159,10 @@ pub(crate) fn build_specs( ConfigShellToolType::UnifiedExec => { builder.push_spec(create_exec_command_tool()); builder.push_spec(create_write_stdin_tool()); + builder.push_spec(create_terminate_session_tool()); builder.register_handler("exec_command", unified_exec_handler.clone()); - builder.register_handler("write_stdin", unified_exec_handler); + builder.register_handler("write_stdin", unified_exec_handler.clone()); + builder.register_handler("terminate_session", unified_exec_handler); } ConfigShellToolType::Disabled => { // Do nothing. @@ -1393,6 +1416,7 @@ mod tests { for spec in [ create_exec_command_tool(), create_write_stdin_tool(), + create_terminate_session_tool(), create_list_mcp_resources_tool(), create_list_mcp_resource_templates_tool(), create_read_mcp_resource_tool(), @@ -1539,6 +1563,7 @@ mod tests { &[ "exec_command", "write_stdin", + "terminate_session", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1560,6 +1585,7 @@ mod tests { &[ "exec_command", "write_stdin", + "terminate_session", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1645,6 +1671,7 @@ mod tests { &[ "exec_command", "write_stdin", + "terminate_session", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1665,6 +1692,7 @@ mod tests { &[ "exec_command", "write_stdin", + "terminate_session", "list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource", @@ -1689,7 +1717,12 @@ mod tests { let (tools, _) = build_specs(&tools_config, Some(HashMap::new())).build(); // Only check the shell variant and a couple of core tools. - let mut subset = vec!["exec_command", "write_stdin", "update_plan"]; + let mut subset = vec![ + "exec_command", + "write_stdin", + "terminate_session", + "update_plan", + ]; if let Some(shell_tool) = shell_tool_name(&tools_config) { subset.push(shell_tool); } @@ -1711,6 +1744,7 @@ mod tests { assert!(!find_tool(&tools, "exec_command").supports_parallel_tool_calls); assert!(!find_tool(&tools, "write_stdin").supports_parallel_tool_calls); + assert!(!find_tool(&tools, "terminate_session").supports_parallel_tool_calls); assert!(find_tool(&tools, "grep_files").supports_parallel_tool_calls); assert!(find_tool(&tools, "list_dir").supports_parallel_tool_calls); assert!(find_tool(&tools, "read_file").supports_parallel_tool_calls); diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 2e80156114f..d37fcb098e5 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -110,9 +110,14 @@ impl UnifiedExecProcessManager { } } - pub(crate) async fn release_process_id(&self, process_id: &str) { + pub(crate) async fn terminate_session(&self, process_id: &str) -> bool { let mut store = self.process_store.lock().await; - store.remove(process_id); + if let Some(entry) = store.remove(process_id) { + entry.process.terminate(); + true + } else { + store.reserved_process_ids.remove(process_id) + } } pub(crate) async fn exec_command( @@ -138,7 +143,7 @@ impl UnifiedExecProcessManager { let process = match process { Ok(process) => Arc::new(process), Err(err) => { - self.release_process_id(&request.process_id).await; + self.terminate_session(&request.process_id).await; return Err(err); } }; @@ -207,7 +212,7 @@ impl UnifiedExecProcessManager { ) .await; - self.release_process_id(&request.process_id).await; + self.terminate_session(&request.process_id).await; process.check_for_sandbox_denial_with_text(&text).await?; } else { // Long‑lived command: persist the process so write_stdin can reuse