use super::*;use crate::shell::ShellType;use crate::shell::default_user_shell;use codex_exec_server::Environment;use codex_tools::UnifiedExecShellMode;use codex_tools::ZshForkConfig;use codex_utils_absolute_path::AbsolutePathBuf;use codex_utils_output_truncation::TruncationPolicy;use pretty_assertions::assert_eq;use std::sync::Arc;use crate::session::tests::make_session_and_context;use crate::tools::context::ExecCommandToolOutput;use crate::tools::context::ToolCallSource;use crate::tools::context::ToolInvocation;use crate::tools::context::ToolPayload;use crate::tools::hook_names::HookToolName;use crate::tools::registry::CoreToolRuntime;use crate::turn_diff_tracker::TurnDiffTracker;use tokio::sync::Mutex;const TEST_TRUNCATION_POLICY: TruncationPolicy = TruncationPolicy::Tokens(10_000);async fn invocation_for_payload( tool_name: &str, call_id: &str, payload: ToolPayload,) -> ToolInvocation { let (session, turn) = make_session_and_context().await; ToolInvocation { session: session.into(), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: call_id.to_string(), tool_name: codex_tools::ToolName::plain(tool_name), source: ToolCallSource::Direct, payload, }}#[test]fn test_get_command_uses_default_shell_when_unspecified() -> anyhow::Result<()> { let json = r#"{"cmd": "echo hello"}"#; let args: ExecCommandArgs = parse_arguments(json)?; assert!(args.shell.is_none()); let resolved = get_command( &args, Arc::new(default_user_shell()), &UnifiedExecShellMode::Direct, /*allow_login_shell*/ true, ) .map_err(anyhow::Error::msg)?; let command = resolved.command; assert_eq!(command.len(), 3); assert_eq!(command[2], "echo hello"); Ok(())}#[test]fn test_get_command_respects_explicit_bash_shell() -> anyhow::Result<()> { let json = r#"{"cmd": "echo hello", "shell": "/bin/bash"}"#; let args: ExecCommandArgs = parse_arguments(json)?; assert_eq!(args.shell.as_deref(), Some("/bin/bash")); let resolved = get_command( &args, Arc::new(default_user_shell()), &UnifiedExecShellMode::Direct, /*allow_login_shell*/ true, ) .map_err(anyhow::Error::msg)?; let command = resolved.command; assert_eq!(command.last(), Some(&"echo hello".to_string())); if command .iter() .any(|arg| arg.eq_ignore_ascii_case("-Command")) { assert!(command.contains(&"-NoProfile".to_string())); } Ok(())}#[test]fn test_get_command_respects_explicit_powershell_shell() -> anyhow::Result<()> { let temp_dir = tempfile::tempdir()?; let powershell_path = temp_dir.path().join(if cfg!(windows) { "powershell.exe" } else { "powershell" }); std::fs::write(&powershell_path, "")?; let json = serde_json::json!({ "cmd": "echo hello", "shell": powershell_path, }) .to_string(); let args: ExecCommandArgs = parse_arguments(&json)?; assert_eq!( args.shell.as_deref(), Some(powershell_path.to_string_lossy().as_ref()) ); let resolved = get_command( &args, Arc::new(default_user_shell()), &UnifiedExecShellMode::Direct, /*allow_login_shell*/ true, ) .map_err(anyhow::Error::msg)?; let command = resolved.command; assert_eq!(command[2], "echo hello"); assert_eq!(resolved.shell_type, ShellType::PowerShell); Ok(())}#[test]fn test_get_command_respects_explicit_cmd_shell() -> anyhow::Result<()> { let json = r#"{"cmd": "echo hello", "shell": "cmd"}"#; let args: ExecCommandArgs = parse_arguments(json)?; assert_eq!(args.shell.as_deref(), Some("cmd")); let resolved = get_command( &args, Arc::new(default_user_shell()), &UnifiedExecShellMode::Direct, /*allow_login_shell*/ true, ) .map_err(anyhow::Error::msg)?; let command = resolved.command; assert_eq!(command[2], "echo hello"); Ok(())}#[test]fn test_get_command_rejects_explicit_login_when_disallowed() -> anyhow::Result<()> { let json = r#"{"cmd": "echo hello", "login": true}"#; let args: ExecCommandArgs = parse_arguments(json)?; let err = get_command( &args, Arc::new(default_user_shell()), &UnifiedExecShellMode::Direct, /*allow_login_shell*/ false, ) .expect_err("explicit login should be rejected"); assert!( err.contains("login shell is disabled by config"), "unexpected error: {err}" ); Ok(())}#[test]fn test_get_command_rejects_explicit_shell_in_zsh_fork_mode() -> anyhow::Result<()> { let json = r#"{"cmd": "echo hello", "shell": "/bin/bash"}"#; let args: ExecCommandArgs = parse_arguments(json)?; let shell_zsh_path = AbsolutePathBuf::from_absolute_path(if cfg!(windows) { r"C:\opt\codex\zsh" } else { "/opt/codex/zsh" })?; let shell_mode = UnifiedExecShellMode::ZshFork(ZshForkConfig { shell_zsh_path, main_execve_wrapper_exe: AbsolutePathBuf::from_absolute_path(if cfg!(windows) { r"C:\opt\codex\codex-execve-wrapper" } else { "/opt/codex/codex-execve-wrapper" })?, }); let err = get_command( &args, Arc::new(default_user_shell()), &shell_mode, /*allow_login_shell*/ true, ) .expect_err("explicit shell should be rejected"); assert!( err.contains("`shell` is not supported for local zsh-fork exec"), "unexpected error: {err}" ); Ok(())}#[tokio::test]async fn shell_mode_for_environment_uses_direct_mode_for_remote_environments() -> anyhow::Result<()>{ let shell_zsh_path = AbsolutePathBuf::from_absolute_path(if cfg!(windows) { r"C:\opt\codex\zsh" } else { "/opt/codex/zsh" })?; let shell_mode = UnifiedExecShellMode::ZshFork(ZshForkConfig { shell_zsh_path, main_execve_wrapper_exe: AbsolutePathBuf::from_absolute_path(if cfg!(windows) { r"C:\opt\codex\codex-execve-wrapper" } else { "/opt/codex/codex-execve-wrapper" })?, }); let local_environment = Environment::default_for_tests(); let remote_environment = Environment::create_for_tests(Some("ws://127.0.0.1:1/remote-exec-server".to_string()))?; assert_eq!( shell_mode_for_environment(&shell_mode, &local_environment), shell_mode ); assert_eq!( shell_mode_for_environment(&shell_mode, &remote_environment), UnifiedExecShellMode::Direct ); Ok(())}#[tokio::test]async fn exec_command_pre_tool_use_payload_uses_raw_command() { let payload = ToolPayload::Function { arguments: serde_json::json!({ "cmd": "printf exec command" }).to_string(), }; let (session, turn) = make_session_and_context().await; let handler = ExecCommandHandler::default(); assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-43".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), source: crate::tools::context::ToolCallSource::Direct, payload, }), Some(crate::tools::registry::PreToolUsePayload { tool_name: HookToolName::bash(), tool_input: serde_json::json!({ "command": "printf exec command" }), }) );}#[tokio::test]async fn exec_command_pre_tool_use_payload_skips_write_stdin() { let payload = ToolPayload::Function { arguments: serde_json::json!({ "chars": "echo hi" }).to_string(), }; let (session, turn) = make_session_and_context().await; let handler = WriteStdinHandler; assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-44".to_string(), tool_name: codex_tools::ToolName::plain("write_stdin"), source: crate::tools::context::ToolCallSource::Direct, payload, }), None );}#[tokio::test]async fn exec_command_post_tool_use_payload_uses_output_for_noninteractive_one_shot_commands() { let payload = ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo three", "tty": false }).to_string(), }; let output = ExecCommandToolOutput { event_call_id: "call-43".to_string(), chunk_id: "chunk-1".to_string(), wall_time: std::time::Duration::from_millis(498), raw_output: b"three".to_vec(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, hook_command: Some("echo three".to_string()), }; let invocation = invocation_for_payload("exec_command", "call-43", payload).await; let handler = ExecCommandHandler::default(); assert_eq!( handler.post_tool_use_payload(&invocation, &output), Some(crate::tools::registry::PostToolUsePayload { tool_name: HookToolName::bash(), tool_use_id: "call-43".to_string(), tool_input: serde_json::json!({ "command": "echo three" }), tool_response: serde_json::json!("three"), }) );}#[tokio::test]async fn exec_command_post_tool_use_payload_uses_output_for_interactive_completion() { let payload = ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo three", "tty": true }).to_string(), }; let output = ExecCommandToolOutput { event_call_id: "call-44".to_string(), chunk_id: "chunk-1".to_string(), wall_time: std::time::Duration::from_millis(498), raw_output: b"three".to_vec(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, hook_command: Some("echo three".to_string()), }; let invocation = invocation_for_payload("exec_command", "call-44", payload).await; let handler = ExecCommandHandler::default(); assert_eq!( handler.post_tool_use_payload(&invocation, &output), Some(crate::tools::registry::PostToolUsePayload { tool_name: HookToolName::bash(), tool_use_id: "call-44".to_string(), tool_input: serde_json::json!({ "command": "echo three" }), tool_response: serde_json::json!("three"), }) );}#[tokio::test]async fn exec_command_post_tool_use_payload_skips_running_sessions() { let payload = ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo three", "tty": false }).to_string(), }; let output = ExecCommandToolOutput { event_call_id: "event-45".to_string(), chunk_id: "chunk-1".to_string(), wall_time: std::time::Duration::from_millis(498), raw_output: b"three".to_vec(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: Some(45), exit_code: None, original_token_count: None, hook_command: Some("echo three".to_string()), }; let invocation = invocation_for_payload("exec_command", "call-45", payload).await; let handler = ExecCommandHandler::default(); assert_eq!(handler.post_tool_use_payload(&invocation, &output), None);}#[tokio::test]async fn write_stdin_post_tool_use_payload_uses_original_exec_call_id_and_command_on_completion() { let payload = ToolPayload::Function { arguments: serde_json::json!({ "session_id": 45, "chars": "", }) .to_string(), }; let output = ExecCommandToolOutput { event_call_id: "exec-call-45".to_string(), chunk_id: "chunk-2".to_string(), wall_time: std::time::Duration::from_millis(498), raw_output: b"finished\n".to_vec(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, hook_command: Some("sleep 1; echo finished".to_string()), }; let invocation = invocation_for_payload("write_stdin", "write-stdin-call", payload).await; let handler = WriteStdinHandler; assert_eq!( handler.post_tool_use_payload(&invocation, &output), Some(crate::tools::registry::PostToolUsePayload { tool_name: HookToolName::bash(), tool_use_id: "exec-call-45".to_string(), tool_input: serde_json::json!({ "command": "sleep 1; echo finished" }), tool_response: serde_json::json!("finished\n"), }) );}#[tokio::test]async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separate() { let payload = ToolPayload::Function { arguments: serde_json::json!({ "session_id": 45, "chars": "" }).to_string(), }; let output_a = ExecCommandToolOutput { event_call_id: "exec-call-a".to_string(), chunk_id: "chunk-a".to_string(), wall_time: std::time::Duration::from_millis(498), raw_output: b"alpha\n".to_vec(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, hook_command: Some("sleep 2; echo alpha".to_string()), }; let output_b = ExecCommandToolOutput { event_call_id: "exec-call-b".to_string(), chunk_id: "chunk-b".to_string(), wall_time: std::time::Duration::from_millis(498), raw_output: b"beta\n".to_vec(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, hook_command: Some("sleep 1; echo beta".to_string()), }; let invocation_b = invocation_for_payload("write_stdin", "write-call-b", payload.clone()).await; let invocation_a = invocation_for_payload("write_stdin", "write-call-a", payload).await; let handler = WriteStdinHandler; let payloads = [ handler.post_tool_use_payload(&invocation_b, &output_b), handler.post_tool_use_payload(&invocation_a, &output_a), ]; assert_eq!( payloads, [ Some(crate::tools::registry::PostToolUsePayload { tool_name: HookToolName::bash(), tool_use_id: "exec-call-b".to_string(), tool_input: serde_json::json!({ "command": "sleep 1; echo beta" }), tool_response: serde_json::json!("beta\n"), }), Some(crate::tools::registry::PostToolUsePayload { tool_name: HookToolName::bash(), tool_use_id: "exec-call-a".to_string(), tool_input: serde_json::json!({ "command": "sleep 2; echo alpha" }), tool_response: serde_json::json!("alpha\n"), }), ] );}