From 6b48e7a626296038fae4aed3e620ba92c17cf13a Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Sat, 7 Feb 2026 01:07:40 -0800 Subject: [PATCH] feat(core): own network proxy singleton in thread manager --- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/network_proxy.rs | 72 +++++++++++++++ codex-rs/core/src/thread_manager.rs | 131 ++++++++++++++++++++++++++-- 3 files changed, 198 insertions(+), 6 deletions(-) create mode 100644 codex-rs/core/src/network_proxy.rs diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index f1534cf0f7..95449f51a7 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -49,6 +49,7 @@ mod mcp_tool_call; mod mentions; mod message_history; mod model_provider_info; +mod network_proxy; pub mod parse_command; pub mod path_utils; pub mod personality_migration; diff --git a/codex-rs/core/src/network_proxy.rs b/codex-rs/core/src/network_proxy.rs new file mode 100644 index 0000000000..78214fb145 --- /dev/null +++ b/codex-rs/core/src/network_proxy.rs @@ -0,0 +1,72 @@ +use anyhow::Context; +use anyhow::Result; +use std::path::Path; +use std::process::Stdio; +use tokio::process::Child; +use tokio::process::Command; +use tokio::time::Duration; +use tokio::time::timeout; + +pub(crate) struct ManagedNetworkProxy { + child: Child, +} + +impl ManagedNetworkProxy { + pub(crate) async fn maybe_start(codex_home: &Path, enabled: bool) -> Result> { + if !enabled { + return Ok(None); + } + + let mut child = Command::new(network_proxy_binary()) + .env("CODEX_HOME", codex_home) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .context("failed to spawn codex-network-proxy process")?; + + // If the proxy process exits immediately, treat startup as failed so callers can + // log once and proceed without assuming policy enforcement is active. + if let Ok(status_result) = timeout(Duration::from_millis(250), child.wait()).await { + let status = status_result.context("failed to wait for codex-network-proxy")?; + anyhow::bail!("codex-network-proxy exited early with status {status}"); + } + + Ok(Some(Self { child })) + } + + pub(crate) async fn shutdown(&mut self) -> Result<()> { + if self + .child + .try_wait() + .context("failed to check codex-network-proxy state")? + .is_some() + { + return Ok(()); + } + + self.child + .start_kill() + .context("failed to signal codex-network-proxy for shutdown")?; + let _ = self.child.wait().await; + Ok(()) + } +} + +fn network_proxy_binary() -> String { + std::env::var("CODEX_NETWORK_PROXY_BIN").unwrap_or_else(|_| "codex-network-proxy".to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn maybe_start_returns_none_when_disabled() { + let codex_home = tempfile::tempdir().expect("create codex home"); + let proxy = ManagedNetworkProxy::maybe_start(codex_home.path(), false) + .await + .expect("startup should succeed"); + assert!(proxy.is_none()); + } +} diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index e702d9a004..7fca9784e4 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -14,6 +14,7 @@ use crate::error::Result as CodexResult; use crate::file_watcher::FileWatcher; use crate::file_watcher::FileWatcherEvent; use crate::models_manager::manager::ModelsManager; +use crate::network_proxy::ManagedNetworkProxy; use crate::protocol::Event; use crate::protocol::EventMsg; use crate::protocol::SessionConfiguredEvent; @@ -36,12 +37,62 @@ use tempfile::TempDir; use tokio::runtime::Handle; #[cfg(any(test, feature = "test-support"))] use tokio::runtime::RuntimeFlavor; +use tokio::sync::Mutex; use tokio::sync::RwLock; use tokio::sync::broadcast; use tracing::warn; const THREAD_CREATED_CHANNEL_CAPACITY: usize = 1024; +/// Process-wide proxy runtime owned by a single ThreadManagerState. +/// +/// This keeps proxy lifecycle in one place: +/// - one startup gate (prevents duplicate binds), +/// - one proxy handle. +struct NetworkProxyRuntime { + codex_home: PathBuf, + proxy: Mutex>, + start_gate: Mutex<()>, +} + +impl NetworkProxyRuntime { + fn new(codex_home: PathBuf) -> Self { + Self { + codex_home, + proxy: Mutex::new(None), + start_gate: Mutex::new(()), + } + } + + async fn ensure_started(&self, enabled: bool) -> CodexResult<()> { + if !enabled { + return Ok(()); + } + + // Serialize startup attempts so concurrent thread starts cannot race and bind twice. + let _gate = self.start_gate.lock().await; + if self.proxy.lock().await.is_some() { + return Ok(()); + } + + let proxy = ManagedNetworkProxy::maybe_start(self.codex_home.as_path(), enabled) + .await + .map_err(|err| CodexErr::Fatal(format!("failed to start network proxy: {err:#}")))?; + *self.proxy.lock().await = proxy; + Ok(()) + } + + async fn shutdown(&self) -> CodexResult<()> { + if let Some(mut proxy) = self.proxy.lock().await.take() { + proxy + .shutdown() + .await + .map_err(|err| CodexErr::Fatal(format!("failed to stop network proxy: {err:#}")))?; + } + Ok(()) + } +} + fn build_file_watcher(codex_home: PathBuf, skills_manager: Arc) -> Arc { #[cfg(any(test, feature = "test-support"))] if let Ok(handle) = Handle::try_current() @@ -109,6 +160,7 @@ pub(crate) struct ThreadManagerState { models_manager: Arc, skills_manager: Arc, file_watcher: Arc, + network_proxy_runtime: NetworkProxyRuntime, session_source: SessionSource, #[cfg(any(test, feature = "test-support"))] #[allow(dead_code)] @@ -125,6 +177,7 @@ impl ThreadManager { let (thread_created_tx, _) = broadcast::channel(THREAD_CREATED_CHANNEL_CAPACITY); let skills_manager = Arc::new(SkillsManager::new(codex_home.clone())); let file_watcher = build_file_watcher(codex_home.clone(), Arc::clone(&skills_manager)); + let network_proxy_runtime = NetworkProxyRuntime::new(codex_home.clone()); Self { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), @@ -132,6 +185,7 @@ impl ThreadManager { models_manager: Arc::new(ModelsManager::new(codex_home, auth_manager.clone())), skills_manager, file_watcher, + network_proxy_runtime, auth_manager, session_source, #[cfg(any(test, feature = "test-support"))] @@ -165,6 +219,7 @@ impl ThreadManager { let (thread_created_tx, _) = broadcast::channel(THREAD_CREATED_CHANNEL_CAPACITY); let skills_manager = Arc::new(SkillsManager::new(codex_home.clone())); let file_watcher = build_file_watcher(codex_home.clone(), Arc::clone(&skills_manager)); + let network_proxy_runtime = NetworkProxyRuntime::new(codex_home.clone()); Self { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), @@ -176,6 +231,7 @@ impl ThreadManager { )), skills_manager, file_watcher, + network_proxy_runtime, auth_manager, session_source: SessionSource::Exec, #[cfg(any(test, feature = "test-support"))] @@ -301,16 +357,12 @@ impl ThreadManager { /// as `Arc`, it is possible that other references to it exist elsewhere. /// Returns the thread if the thread was found and removed. pub async fn remove_thread(&self, thread_id: &ThreadId) -> Option> { - self.state.threads.write().await.remove(thread_id) + self.state.remove_thread(thread_id).await } /// Closes all threads open in this ThreadManager pub async fn remove_and_close_all_threads(&self) -> CodexResult<()> { - for thread in self.state.threads.read().await.values() { - thread.submit(Op::Shutdown).await?; - } - self.state.threads.write().await.clear(); - Ok(()) + self.state.remove_and_close_all_threads().await } /// Fork an existing thread by taking messages up to the given position (not including @@ -434,6 +486,11 @@ impl ThreadManagerState { session_source: SessionSource, dynamic_tools: Vec, ) -> CodexResult { + // Requirements currently control whether proxy wiring is active for this process. + let proxy_enabled = proxy_enabled_from_requirements(&config); + if let Err(err) = self.ensure_network_proxy_started(proxy_enabled).await { + warn!("failed to start network proxy singleton: {err:#}"); + } self.file_watcher.register_config(&config); let CodexSpawnOk { codex, thread_id, .. @@ -485,6 +542,46 @@ impl ThreadManagerState { pub(crate) fn notify_thread_created(&self, thread_id: ThreadId) { let _ = self.thread_created_tx.send(thread_id); } + + pub(crate) async fn ensure_network_proxy_started(&self, enabled: bool) -> CodexResult<()> { + self.network_proxy_runtime.ensure_started(enabled).await + } + + pub(crate) async fn shutdown_network_proxy(&self) -> CodexResult<()> { + self.network_proxy_runtime.shutdown().await + } + + pub(crate) async fn remove_and_close_all_threads(&self) -> CodexResult<()> { + // Clone thread handles first so we don't hold the map lock while awaiting shutdown. + let threads = self + .threads + .read() + .await + .values() + .cloned() + .collect::>(); + for thread in threads { + thread.submit(Op::Shutdown).await?; + } + self.threads.write().await.clear(); + + if let Err(err) = self.shutdown_network_proxy().await { + warn!("failed to stop network proxy singleton: {err:#}"); + } + + Ok(()) + } +} + +fn proxy_enabled_from_requirements(config: &Config) -> bool { + // Presence of any requirements network block enables proxy wiring; explicit enabled=false + // is treated as an opt-out to keep control with requirements authors. + config + .config_layer_stack + .requirements_toml() + .network + .as_ref() + .is_some_and(|network| network.enabled.unwrap_or(true)) } /// Return a prefix of `items` obtained by cutting strictly before the nth user message @@ -614,4 +711,26 @@ mod tests { serde_json::to_value(&expected).unwrap() ); } + + #[tokio::test] + async fn network_proxy_runtime_noops_when_network_disabled() { + let codex_home = tempfile::tempdir().expect("create codex_home tempdir"); + let runtime = NetworkProxyRuntime::new(codex_home.path().to_path_buf()); + runtime + .ensure_started(false) + .await + .expect("disabled network should skip startup"); + assert!(runtime.proxy.lock().await.is_none()); + } + + #[tokio::test] + async fn network_proxy_runtime_shutdown_noops_without_proxy() { + let codex_home = tempfile::tempdir().expect("create codex_home tempdir"); + let runtime = NetworkProxyRuntime::new(codex_home.path().to_path_buf()); + runtime + .shutdown() + .await + .expect("shutdown without active proxy should succeed"); + assert!(runtime.proxy.lock().await.is_none()); + } }