From ff8c051407ca21341e98c3b60fa8985a3664dcad Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 9 Feb 2026 14:52:33 -0800 Subject: [PATCH] feat: retain NetworkProxy, when appropriate --- .../app-server/src/codex_message_processor.rs | 22 ++- codex-rs/core/src/codex.rs | 23 +++ codex-rs/core/src/config/mod.rs | 170 +++++++++++++++++- codex-rs/core/src/state/service.rs | 5 + codex-rs/core/src/tasks/user_shell.rs | 2 +- codex-rs/core/src/tools/handlers/shell.rs | 6 +- .../core/src/unified_exec/process_manager.rs | 2 +- codex-rs/network-proxy/src/config.rs | 4 +- 8 files changed, 223 insertions(+), 11 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 0c24ed30c0..b0a8a37778 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1639,13 +1639,30 @@ impl CodexMessageProcessor { let timeout_ms = params .timeout_ms .and_then(|timeout_ms| u64::try_from(timeout_ms).ok()); + let started_network_proxy = match self.config.network.as_ref() { + Some(spec) => match spec.start_proxy().await { + Ok(started) => Some(started), + Err(err) => { + let error = JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: format!("failed to start managed network proxy: {err}"), + data: None, + }; + self.outgoing.send_error(request, error).await; + return; + } + }, + None => None, + }; let windows_sandbox_level = WindowsSandboxLevel::from_config(&self.config); let exec_params = ExecParams { command: params.command, cwd, expiration: timeout_ms.into(), env, - network: self.config.network.clone(), + network: started_network_proxy + .as_ref() + .map(|(proxy, _handle)| proxy.clone()), sandbox_permissions: SandboxPermissions::UseDefault, windows_sandbox_level, justification: None, @@ -1673,9 +1690,12 @@ impl CodexMessageProcessor { let outgoing = self.outgoing.clone(); let request_for_task = request; let sandbox_cwd = self.config.cwd.clone(); + let started_network_proxy_for_task = started_network_proxy; let use_linux_sandbox_bwrap = self.config.features.enabled(Feature::UseLinuxSandboxBwrap); tokio::spawn(async move { + let _network_proxy_handle = + started_network_proxy_for_task.map(|(_proxy, handle)| handle); match codex_core::exec::process_exec_tool_call( exec_params, &effective_policy, diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 53b8aa8d5e..56b6eb44bc 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -44,6 +44,7 @@ use crate::turn_metadata::resolve_turn_metadata_header_with_timeout; use crate::util::error_or_panic; use async_channel::Receiver; use async_channel::Sender; +use codex_network_proxy::NetworkProxy; use codex_protocol::ThreadId; use codex_protocol::approvals::ExecPolicyAmendment; use codex_protocol::config_types::ModeKind; @@ -533,6 +534,7 @@ pub(crate) struct TurnContext { pub(crate) personality: Option, pub(crate) approval_policy: AskForApproval, pub(crate) sandbox_policy: SandboxPolicy, + pub(crate) network: Option, pub(crate) windows_sandbox_level: WindowsSandboxLevel, pub(crate) shell_environment_policy: ShellEnvironmentPolicy, pub(crate) tools_config: ToolsConfig, @@ -797,6 +799,7 @@ impl Session { session_configuration: &SessionConfiguration, per_turn_config: Config, model_info: ModelInfo, + network: Option, sub_id: String, ) -> TurnContext { let reasoning_effort = session_configuration.collaboration_mode.reasoning_effort(); @@ -836,6 +839,7 @@ impl Session { personality: session_configuration.personality, approval_policy: session_configuration.approval_policy.value(), sandbox_policy: session_configuration.sandbox_policy.get().clone(), + network, windows_sandbox_level: session_configuration.windows_sandbox_level, shell_environment_policy: per_turn_config.shell_environment_policy.clone(), tools_config, @@ -1058,6 +1062,15 @@ impl Session { }; session_configuration.thread_name = thread_name.clone(); let mut state = SessionState::new(session_configuration.clone()); + let (network_proxy, network_proxy_handle) = match config.network.as_ref() { + Some(spec) => { + let (proxy, handle) = spec.start_proxy().await.map_err(|err| { + anyhow::anyhow!("failed to start managed network proxy: {err}") + })?; + (Some(proxy), Some(handle)) + } + None => (None, None), + }; let services = SessionServices { mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())), @@ -1080,6 +1093,8 @@ impl Session { skills_manager, file_watcher, agent_control, + network_proxy, + network_proxy_handle, state_db: state_db_ctx.clone(), model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), @@ -1553,6 +1568,7 @@ impl Session { &session_configuration, per_turn_config, model_info, + self.services.network_proxy.clone(), sub_id, ); @@ -3696,6 +3712,7 @@ async fn spawn_review_thread( personality: parent_turn_context.personality, approval_policy: parent_turn_context.approval_policy, sandbox_policy: parent_turn_context.sandbox_policy.clone(), + network: parent_turn_context.network.clone(), windows_sandbox_level: parent_turn_context.windows_sandbox_level, shell_environment_policy: parent_turn_context.shell_environment_policy.clone(), cwd: parent_turn_context.cwd.clone(), @@ -6182,6 +6199,8 @@ mod tests { skills_manager, file_watcher, agent_control, + network_proxy: None, + network_proxy_handle: None, state_db: None, model_client: ModelClient::new( Some(auth_manager.clone()), @@ -6205,6 +6224,7 @@ mod tests { &session_configuration, per_turn_config, model_info, + None, "turn_id".to_string(), ); @@ -6315,6 +6335,8 @@ mod tests { skills_manager, file_watcher, agent_control, + network_proxy: None, + network_proxy_handle: None, state_db: None, model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), @@ -6338,6 +6360,7 @@ mod tests { &session_configuration, per_turn_config, model_info, + None, "turn_id".to_string(), )); diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 34a1772266..dcd3c1cc5d 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -26,6 +26,7 @@ use crate::config_loader::ConstrainedWithSource; use crate::config_loader::LoaderOverrides; use crate::config_loader::McpServerIdentity; use crate::config_loader::McpServerRequirement; +use crate::config_loader::NetworkConstraints; use crate::config_loader::ResidencyRequirement; use crate::config_loader::Sourced; use crate::config_loader::load_config_layers_state; @@ -45,9 +46,18 @@ use crate::project_doc::LOCAL_PROJECT_DOC_FILENAME; use crate::protocol::AskForApproval; use crate::protocol::SandboxPolicy; use crate::windows_sandbox::WindowsSandboxLevelExt; +use async_trait::async_trait; use codex_app_server_protocol::Tools; use codex_app_server_protocol::UserSavedConfig; +use codex_network_proxy::ConfigReloader; +use codex_network_proxy::ConfigState; use codex_network_proxy::NetworkProxy; +use codex_network_proxy::NetworkProxyConfig; +use codex_network_proxy::NetworkProxyConstraints; +use codex_network_proxy::NetworkProxyHandle; +use codex_network_proxy::NetworkProxyState; +use codex_network_proxy::build_config_state; +use codex_network_proxy::validate_policy_against_constraints; use codex_protocol::config_types::AltScreenMode; use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModeKind; @@ -71,6 +81,7 @@ use std::collections::HashMap; use std::io::ErrorKind; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; #[cfg(test)] use tempfile::tempdir; @@ -101,6 +112,143 @@ pub(crate) const DEFAULT_AGENT_MAX_THREADS: Option = Some(6); pub const CONFIG_TOML_FILE: &str = "config.toml"; +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NetworkProxySpec { + config: NetworkProxyConfig, + constraints: NetworkProxyConstraints, + cfg_path: PathBuf, +} + +#[derive(Clone)] +struct StaticNetworkProxyReloader { + state: ConfigState, +} + +impl StaticNetworkProxyReloader { + fn new(state: ConfigState) -> Self { + Self { state } + } +} + +#[async_trait] +impl ConfigReloader for StaticNetworkProxyReloader { + async fn maybe_reload(&self) -> anyhow::Result> { + Ok(None) + } + + async fn reload_now(&self) -> anyhow::Result { + Ok(self.state.clone()) + } +} + +impl NetworkProxySpec { + fn from_constraints( + merged_toml: TomlValue, + requirements: NetworkConstraints, + cfg_path: PathBuf, + ) -> std::io::Result { + let config: NetworkProxyConfig = merged_toml.try_into().map_err(|err| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("failed to deserialize network proxy config: {err}"), + ) + })?; + let (config, constraints) = Self::apply_requirements(config, &requirements); + validate_policy_against_constraints(&config, &constraints).map_err(|err| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("network proxy constraints are invalid: {err}"), + ) + })?; + Ok(Self { + config, + constraints, + cfg_path, + }) + } + + pub async fn start_proxy(&self) -> std::io::Result<(NetworkProxy, NetworkProxyHandle)> { + let state = build_config_state( + self.config.clone(), + self.constraints.clone(), + self.cfg_path.clone(), + ) + .map_err(|err| { + std::io::Error::other(format!("failed to build network proxy state: {err}")) + })?; + let reloader = Arc::new(StaticNetworkProxyReloader::new(state.clone())); + let state = NetworkProxyState::with_reloader(state, reloader); + let proxy = NetworkProxy::builder() + .state(Arc::new(state)) + .build() + .await + .map_err(|err| { + std::io::Error::other(format!("failed to build network proxy: {err}")) + })?; + let handle = proxy + .run() + .await + .map_err(|err| std::io::Error::other(format!("failed to run network proxy: {err}")))?; + Ok((proxy, handle)) + } + + fn apply_requirements( + mut config: NetworkProxyConfig, + requirements: &NetworkConstraints, + ) -> (NetworkProxyConfig, NetworkProxyConstraints) { + let mut constraints = NetworkProxyConstraints::default(); + + if let Some(enabled) = requirements.enabled { + config.network.enabled = enabled; + constraints.enabled = Some(enabled); + } + if let Some(http_port) = requirements.http_port { + config.network.proxy_url = format!("http://127.0.0.1:{http_port}"); + } + if let Some(socks_port) = requirements.socks_port { + config.network.socks_url = format!("http://127.0.0.1:{socks_port}"); + } + if let Some(allow_upstream_proxy) = requirements.allow_upstream_proxy { + config.network.allow_upstream_proxy = allow_upstream_proxy; + constraints.allow_upstream_proxy = Some(allow_upstream_proxy); + } + if let Some(dangerously_allow_non_loopback_proxy) = + requirements.dangerously_allow_non_loopback_proxy + { + config.network.dangerously_allow_non_loopback_proxy = + dangerously_allow_non_loopback_proxy; + constraints.dangerously_allow_non_loopback_proxy = + Some(dangerously_allow_non_loopback_proxy); + } + if let Some(dangerously_allow_non_loopback_admin) = + requirements.dangerously_allow_non_loopback_admin + { + config.network.dangerously_allow_non_loopback_admin = + dangerously_allow_non_loopback_admin; + constraints.dangerously_allow_non_loopback_admin = + Some(dangerously_allow_non_loopback_admin); + } + if let Some(allowed_domains) = requirements.allowed_domains.clone() { + config.network.allowed_domains = allowed_domains.clone(); + constraints.allowed_domains = Some(allowed_domains); + } + if let Some(denied_domains) = requirements.denied_domains.clone() { + config.network.denied_domains = denied_domains.clone(); + constraints.denied_domains = Some(denied_domains); + } + if let Some(allow_unix_sockets) = requirements.allow_unix_sockets.clone() { + config.network.allow_unix_sockets = allow_unix_sockets.clone(); + constraints.allow_unix_sockets = Some(allow_unix_sockets); + } + if let Some(allow_local_binding) = requirements.allow_local_binding { + config.network.allow_local_binding = allow_local_binding; + constraints.allow_local_binding = Some(allow_local_binding); + } + + (config, constraints) + } +} + #[cfg(test)] pub(crate) fn test_config() -> Config { let codex_home = tempdir().expect("create temp dir"); @@ -154,7 +302,7 @@ pub struct Config { pub enforce_residency: Constrained>, /// Effective network configuration applied to all spawned processes. - pub network: Option, + pub network: Option, /// True if the user passed in an override or set a value in config.toml /// for either of approval_policy or sandbox_mode. @@ -1650,6 +1798,7 @@ impl Config { // Ensure that every field of ConfigRequirements is applied to the final // Config. + let merged_toml = config_layer_stack.effective_config(); let ConfigRequirements { approval_policy: mut constrained_approval_policy, sandbox_policy: mut constrained_sandbox_policy, @@ -1657,7 +1806,7 @@ impl Config { mcp_servers, exec_policy: _, enforce_residency, - network: _network_requirements, + network: network_requirements, } = requirements; apply_requirement_constrained_value( @@ -1682,6 +1831,21 @@ impl Config { let mcp_servers = constrain_mcp_servers(cfg.mcp_servers.clone(), mcp_servers.as_ref()) .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, format!("{e}")))?; + let network = match network_requirements { + Some(Sourced { value, source }) => { + let cfg_path = codex_home.join(CONFIG_TOML_FILE); + let network = NetworkProxySpec::from_constraints(merged_toml, value, cfg_path) + .map_err(|err| { + std::io::Error::new( + err.kind(), + format!("failed to build managed network proxy from {source}: {err}"), + ) + })?; + Some(network) + } + None => None, + }; + let config = Self { model, review_model, @@ -1694,7 +1858,7 @@ impl Config { approval_policy: constrained_approval_policy.value, sandbox_policy: constrained_sandbox_policy.value, enforce_residency: enforce_residency.value, - network: None, + network, did_user_set_custom_approval_policy_or_sandbox_mode, forced_auto_mode_downgraded_on_windows, shell_environment_policy, diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 0438119d42..3c2ebbe81e 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -14,6 +14,8 @@ use crate::skills::SkillsManager; use crate::state_db::StateDbHandle; use crate::tools::sandboxing::ApprovalStore; use crate::unified_exec::UnifiedExecProcessManager; +use codex_network_proxy::NetworkProxy; +use codex_network_proxy::NetworkProxyHandle; use codex_otel::OtelManager; use tokio::sync::Mutex; use tokio::sync::RwLock; @@ -38,6 +40,9 @@ pub(crate) struct SessionServices { pub(crate) skills_manager: Arc, pub(crate) file_watcher: Arc, pub(crate) agent_control: AgentControl, + pub(crate) network_proxy: Option, + #[allow(dead_code)] + pub(crate) network_proxy_handle: Option, pub(crate) state_db: Option, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index c034e26f47..d33d1fe0bf 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -168,7 +168,7 @@ pub(crate) async fn execute_user_shell_command( exec_env, &sandbox_policy, stdout_stream, - turn_context.config.network.clone(), + turn_context.network.clone(), ) .or_cancel(&cancellation_token) .await; diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 18c594d997..b9e2a97d6d 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -53,7 +53,7 @@ impl ShellHandler { cwd: turn_context.resolve_path(params.workdir.clone()), expiration: params.timeout_ms.into(), env: create_env(&turn_context.shell_environment_policy, Some(thread_id)), - network: turn_context.config.network.clone(), + network: turn_context.network.clone(), sandbox_permissions: params.sandbox_permissions.unwrap_or_default(), windows_sandbox_level: turn_context.windows_sandbox_level, justification: params.justification.clone(), @@ -82,7 +82,7 @@ impl ShellCommandHandler { cwd: turn_context.resolve_path(params.workdir.clone()), expiration: params.timeout_ms.into(), env: create_env(&turn_context.shell_environment_policy, Some(thread_id)), - network: turn_context.config.network.clone(), + network: turn_context.network.clone(), sandbox_permissions: params.sandbox_permissions.unwrap_or_default(), windows_sandbox_level: turn_context.windows_sandbox_level, justification: params.justification.clone(), @@ -444,7 +444,7 @@ mod tests { assert_eq!(exec_params.command, expected_command); assert_eq!(exec_params.cwd, expected_cwd); assert_eq!(exec_params.env, expected_env); - assert_eq!(exec_params.network, turn_context.config.network); + assert_eq!(exec_params.network, turn_context.network); assert_eq!(exec_params.expiration.timeout_ms(), timeout_ms); assert_eq!(exec_params.sandbox_permissions, sandbox_permissions); assert_eq!(exec_params.justification, justification); diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index cee73fae25..0ed05a0a08 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -520,7 +520,7 @@ impl UnifiedExecProcessManager { command: request.command.clone(), cwd, env, - network: context.turn.config.network.clone(), + network: context.turn.network.clone(), tty: request.tty, sandbox_permissions: request.sandbox_permissions, justification: request.justification.clone(), diff --git a/codex-rs/network-proxy/src/config.rs b/codex-rs/network-proxy/src/config.rs index e4ef202a45..8ce914b717 100644 --- a/codex-rs/network-proxy/src/config.rs +++ b/codex-rs/network-proxy/src/config.rs @@ -8,13 +8,13 @@ use std::net::SocketAddr; use tracing::warn; use url::Url; -#[derive(Debug, Clone, Serialize, Deserialize, Default)] +#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)] pub struct NetworkProxyConfig { #[serde(default)] pub network: NetworkProxySettings, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct NetworkProxySettings { #[serde(default)] pub enabled: bool,