From db847220801d4065e1b13f5979907db84be5cbed Mon Sep 17 00:00:00 2001
From: Gabriel Peal <gpeal@users.noreply.github.com>
Date: Wed, 23 Jul 2025 15:40:00 -0700
Subject: [PATCH] Fix flaky test (#1664)

Co-authored-by: aibrahim-oai <aibrahim@openai.com>
---
 codex-rs/chatgpt/src/apply_command.rs       | 24 +++++++++++++++------
 codex-rs/chatgpt/tests/apply_command_e2e.rs | 14 ++----------
 codex-rs/cli/src/main.rs                    |  2 +-
 3 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/codex-rs/chatgpt/src/apply_command.rs b/codex-rs/chatgpt/src/apply_command.rs
index 4209d958..52ab205a 100644
--- a/codex-rs/chatgpt/src/apply_command.rs
+++ b/codex-rs/chatgpt/src/apply_command.rs
@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use clap::Parser;
 use codex_common::CliConfigOverrides;
 use codex_core::config::Config;
@@ -17,7 +19,10 @@ pub struct ApplyCommand {
     #[clap(flatten)]
     pub config_overrides: CliConfigOverrides,
 }
-pub async fn run_apply_command(apply_cli: ApplyCommand) -> anyhow::Result<()> {
+pub async fn run_apply_command(
+    apply_cli: ApplyCommand,
+    cwd: Option<PathBuf>,
+) -> anyhow::Result<()> {
     let config = Config::load_with_cli_overrides(
         apply_cli
             .config_overrides
@@ -29,10 +34,13 @@ pub async fn run_apply_command(apply_cli: ApplyCommand) -> anyhow::Result<()> {
     init_chatgpt_token_from_auth(&config.codex_home).await?;
 
     let task_response = get_task(&config, apply_cli.task_id).await?;
-    apply_diff_from_task(task_response).await
+    apply_diff_from_task(task_response, cwd).await
 }
 
-pub async fn apply_diff_from_task(task_response: GetTaskResponse) -> anyhow::Result<()> {
+pub async fn apply_diff_from_task(
+    task_response: GetTaskResponse,
+    cwd: Option<PathBuf>,
+) -> anyhow::Result<()> {
     let diff_turn = match task_response.current_diff_task_turn {
         Some(turn) => turn,
         None => anyhow::bail!("No diff turn found"),
@@ -42,13 +50,17 @@ pub async fn apply_diff_from_task(task_response: GetTaskResponse) -> anyhow::Res
         _ => None,
     });
     match output_diff {
-        Some(output_diff) => apply_diff(&output_diff.diff).await,
+        Some(output_diff) => apply_diff(&output_diff.diff, cwd).await,
         None => anyhow::bail!("No PR output item found"),
     }
 }
 
-async fn apply_diff(diff: &str) -> anyhow::Result<()> {
-    let toplevel_output = tokio::process::Command::new("git")
+async fn apply_diff(diff: &str, cwd: Option<PathBuf>) -> anyhow::Result<()> {
+    let mut cmd = tokio::process::Command::new("git");
+    if let Some(cwd) = cwd {
+        cmd.current_dir(cwd);
+    }
+    let toplevel_output = cmd
         .args(vec!["rev-parse", "--show-toplevel"])
         .output()
         .await?;
diff --git a/codex-rs/chatgpt/tests/apply_command_e2e.rs b/codex-rs/chatgpt/tests/apply_command_e2e.rs
index e395e4f1..45c33bed 100644
--- a/codex-rs/chatgpt/tests/apply_command_e2e.rs
+++ b/codex-rs/chatgpt/tests/apply_command_e2e.rs
@@ -78,17 +78,7 @@ async fn test_apply_command_creates_fibonacci_file() {
         .await
         .expect("Failed to load fixture");
 
-    let original_dir = std::env::current_dir().expect("Failed to get current dir");
-    std::env::set_current_dir(repo_path).expect("Failed to change directory");
-    struct DirGuard(std::path::PathBuf);
-    impl Drop for DirGuard {
-        fn drop(&mut self) {
-            let _ = std::env::set_current_dir(&self.0);
-        }
-    }
-    let _guard = DirGuard(original_dir);
-
-    apply_diff_from_task(task_response)
+    apply_diff_from_task(task_response, Some(repo_path.to_path_buf()))
         .await
         .expect("Failed to apply diff from task");
 
@@ -173,7 +163,7 @@ console.log(fib(10));
         .await
         .expect("Failed to load fixture");
 
-    let apply_result = apply_diff_from_task(task_response).await;
+    let apply_result = apply_diff_from_task(task_response, Some(repo_path.to_path_buf())).await;
 
     assert!(
         apply_result.is_err(),
diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs
index 7e23782d..e397b0ca 100644
--- a/codex-rs/cli/src/main.rs
+++ b/codex-rs/cli/src/main.rs
@@ -145,7 +145,7 @@ async fn cli_main(codex_linux_sandbox_exe: Option<PathBuf>) -> anyhow::Result<()
         },
         Some(Subcommand::Apply(mut apply_cli)) => {
             prepend_config_flags(&mut apply_cli.config_overrides, cli.config_overrides);
-            run_apply_command(apply_cli).await?;
+            run_apply_command(apply_cli, None).await?;
         }
     }