Refactor unified-exec session creation

Co-authored-by: Codex <noreply@openai.com>

codex-rs/core/src/unified_exec/backend.rs

@@ -0,0 +1,77 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+
+use crate::sandboxing::ExecRequest;
+use crate::unified_exec::SpawnLifecycleHandle;
+use crate::unified_exec::UnifiedExecError;
+use crate::unified_exec::UnifiedExecProcess;
+
+pub(crate) type UnifiedExecSessionFactoryHandle = Arc<dyn UnifiedExecSessionFactory>;
+
+#[async_trait]
+pub(crate) trait UnifiedExecSessionFactory: std::fmt::Debug + Send + Sync {
+    async fn open_session(
+        &self,
+        env: &ExecRequest,
+        tty: bool,
+        spawn_lifecycle: SpawnLifecycleHandle,
+    ) -> Result<UnifiedExecProcess, UnifiedExecError>;
+}
+
+#[derive(Debug, Default)]
+pub(crate) struct LocalUnifiedExecSessionFactory;
+
+pub(crate) fn local_unified_exec_session_factory() -> UnifiedExecSessionFactoryHandle {
+    Arc::new(LocalUnifiedExecSessionFactory)
+}
+
+#[async_trait]
+impl UnifiedExecSessionFactory for LocalUnifiedExecSessionFactory {
+    async fn open_session(
+        &self,
+        env: &ExecRequest,
+        tty: bool,
+        spawn_lifecycle: SpawnLifecycleHandle,
+    ) -> Result<UnifiedExecProcess, UnifiedExecError> {
+        open_local_session(env, tty, spawn_lifecycle).await
+    }
+}
+
+async fn open_local_session(
+    env: &ExecRequest,
+    tty: bool,
+    mut spawn_lifecycle: SpawnLifecycleHandle,
+) -> Result<UnifiedExecProcess, UnifiedExecError> {
+    let (program, args) = env
+        .command
+        .split_first()
+        .ok_or(UnifiedExecError::MissingCommandLine)?;
+    let inherited_fds = spawn_lifecycle.inherited_fds();
+
+    let spawn_result = if tty {
+        codex_utils_pty::pty::spawn_process_with_inherited_fds(
+            program,
+            args,
+            env.cwd.as_path(),
+            &env.env,
+            &env.arg0,
+            codex_utils_pty::TerminalSize::default(),
+            &inherited_fds,
+        )
+        .await
+    } else {
+        codex_utils_pty::pipe::spawn_process_no_stdin_with_inherited_fds(
+            program,
+            args,
+            env.cwd.as_path(),
+            &env.env,
+            &env.arg0,
+            &inherited_fds,
+        )
+        .await
+    };
+    let spawned = spawn_result.map_err(|err| UnifiedExecError::create_process(err.to_string()))?;
+    spawn_lifecycle.after_spawn();
+    UnifiedExecProcess::from_spawned(spawned, env.sandbox, spawn_lifecycle).await
+}

codex-rs/core/src/unified_exec/mod.rs

@@ -38,6 +38,7 @@ use crate::codex::TurnContext;
 use crate::sandboxing::SandboxPermissions;
 
 mod async_watcher;
+mod backend;
 mod errors;
 mod head_tail_buffer;
 mod process;
@@ -47,6 +48,8 @@ pub(crate) fn set_deterministic_process_ids_for_tests(enabled: bool) {
     process_manager::set_deterministic_process_ids_for_tests(enabled);
 }
 
+pub(crate) use backend::UnifiedExecSessionFactoryHandle;
+pub(crate) use backend::local_unified_exec_session_factory;
 pub(crate) use errors::UnifiedExecError;
 pub(crate) use process::NoopSpawnLifecycle;
 #[cfg(unix)]
@@ -123,14 +126,26 @@ impl ProcessStore {
 pub(crate) struct UnifiedExecProcessManager {
     process_store: Mutex<ProcessStore>,
     max_write_stdin_yield_time_ms: u64,
+    session_factory: UnifiedExecSessionFactoryHandle,
 }
 
 impl UnifiedExecProcessManager {
     pub(crate) fn new(max_write_stdin_yield_time_ms: u64) -> Self {
+        Self::with_session_factory(
+            max_write_stdin_yield_time_ms,
+            local_unified_exec_session_factory(),
+        )
+    }
+
+    pub(crate) fn with_session_factory(
+        max_write_stdin_yield_time_ms: u64,
+        session_factory: UnifiedExecSessionFactoryHandle,
+    ) -> Self {
         Self {
             process_store: Mutex::new(ProcessStore::default()),
             max_write_stdin_yield_time_ms: max_write_stdin_yield_time_ms
                 .max(MIN_EMPTY_YIELD_TIME_MS),
+            session_factory,
         }
     }
 }

codex-rs/core/src/unified_exec/process_manager.rs

@@ -541,40 +541,11 @@ impl UnifiedExecProcessManager {
         &self,
         env: &ExecRequest,
         tty: bool,
-        mut spawn_lifecycle: SpawnLifecycleHandle,
+        spawn_lifecycle: SpawnLifecycleHandle,
     ) -> Result<UnifiedExecProcess, UnifiedExecError> {
-        let (program, args) = env
-            .command
-            .split_first()
-            .ok_or(UnifiedExecError::MissingCommandLine)?;
-        let inherited_fds = spawn_lifecycle.inherited_fds();
-
-        let spawn_result = if tty {
-            codex_utils_pty::pty::spawn_process_with_inherited_fds(
-                program,
-                args,
-                env.cwd.as_path(),
-                &env.env,
-                &env.arg0,
-                codex_utils_pty::TerminalSize::default(),
-                &inherited_fds,
-            )
+        self.session_factory
+            .open_session(env, tty, spawn_lifecycle)
             .await
-        } else {
-            codex_utils_pty::pipe::spawn_process_no_stdin_with_inherited_fds(
-                program,
-                args,
-                env.cwd.as_path(),
-                &env.env,
-                &env.arg0,
-                &inherited_fds,
-            )
-            .await
-        };
-        let spawned =
-            spawn_result.map_err(|err| UnifiedExecError::create_process(err.to_string()))?;
-        spawn_lifecycle.after_spawn();
-        UnifiedExecProcess::from_spawned(spawned, env.sandbox, spawn_lifecycle).await
     }
 
     pub(super) async fn open_session_with_sandbox(