feat: retain NetworkProxy, when appropriate (#11207)

As of this PR, `SessionServices` retains a
`Option<StartedNetworkProxy>`, if appropriate.

Now the `network` field on `Config` is `Option<NetworkProxySpec>`
instead of `Option<NetworkProxy>`.

Over in `Session::new()`, we invoke `NetworkProxySpec::start_proxy()` to
create the `StartedNetworkProxy`, which is a new struct that retains the
`NetworkProxy` as well as the `NetworkProxyHandle`. (Note that `Drop` is
implemented for `NetworkProxyHandle` to ensure the proxies are shutdown
when it is dropped.)

The `NetworkProxy` from the `StartedNetworkProxy` is threaded through to
the appropriate places.


---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/openai/codex/pull/11207).
* #11285
* __->__ #11207

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;
@@ -113,6 +114,7 @@ use crate::config::Config;
 use crate::config::Constrained;
 use crate::config::ConstraintResult;
 use crate::config::GhostSnapshotConfig;
+use crate::config::StartedNetworkProxy;
 use crate::config::resolve_web_search_mode_for_turn;
 use crate::config::types::McpServerConfig;
 use crate::config::types::ShellEnvironmentPolicy;
@@ -172,6 +174,7 @@ use crate::protocol::RequestUserInputEvent;
 use crate::protocol::ReviewDecision;
 use crate::protocol::SandboxPolicy;
 use crate::protocol::SessionConfiguredEvent;
+use crate::protocol::SessionNetworkProxyRuntime;
 use crate::protocol::SkillDependencies as ProtocolSkillDependencies;
 use crate::protocol::SkillErrorInfo;
 use crate::protocol::SkillInterface as ProtocolSkillInterface;
@@ -539,6 +542,7 @@ pub(crate) struct TurnContext {
     pub(crate) personality: Option<Personality>,
     pub(crate) approval_policy: AskForApproval,
     pub(crate) sandbox_policy: SandboxPolicy,
+    pub(crate) network: Option<NetworkProxy>,
     pub(crate) windows_sandbox_level: WindowsSandboxLevel,
     pub(crate) shell_environment_policy: ShellEnvironmentPolicy,
     pub(crate) tools_config: ToolsConfig,
@@ -803,6 +807,7 @@ impl Session {
         session_configuration: &SessionConfiguration,
         per_turn_config: Config,
         model_info: ModelInfo,
+        network: Option<NetworkProxy>,
         sub_id: String,
     ) -> TurnContext {
         let reasoning_effort = session_configuration.collaboration_mode.reasoning_effort();
@@ -842,6 +847,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,
@@ -1064,6 +1070,21 @@ impl Session {
             };
         session_configuration.thread_name = thread_name.clone();
         let mut state = SessionState::new(session_configuration.clone());
+        let network_proxy =
+            match config.network.as_ref() {
+                Some(spec) => Some(spec.start_proxy().await.map_err(|err| {
+                    anyhow::anyhow!("failed to start managed network proxy: {err}")
+                })?),
+                None => None,
+            };
+        let session_network_proxy = network_proxy.as_ref().map(|started| {
+            let proxy = started.proxy();
+            SessionNetworkProxyRuntime {
+                http_addr: proxy.http_addr().to_string(),
+                socks_addr: proxy.socks_addr().to_string(),
+                admin_addr: proxy.admin_addr().to_string(),
+            }
+        });
 
         let services = SessionServices {
             mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())),
@@ -1086,6 +1107,7 @@ impl Session {
             skills_manager,
             file_watcher,
             agent_control,
+            network_proxy,
             state_db: state_db_ctx.clone(),
             model_client: ModelClient::new(
                 Some(Arc::clone(&auth_manager)),
@@ -1144,6 +1166,7 @@ impl Session {
                 history_log_id,
                 history_entry_count,
                 initial_messages,
+                network_proxy: session_network_proxy,
                 rollout_path,
             }),
         })
@@ -1559,6 +1582,10 @@ impl Session {
             &session_configuration,
             per_turn_config,
             model_info,
+            self.services
+                .network_proxy
+                .as_ref()
+                .map(StartedNetworkProxy::proxy),
             sub_id,
         );
 
@@ -3702,6 +3729,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(),
@@ -6188,6 +6216,7 @@ mod tests {
             skills_manager,
             file_watcher,
             agent_control,
+            network_proxy: None,
             state_db: None,
             model_client: ModelClient::new(
                 Some(auth_manager.clone()),
@@ -6211,6 +6240,7 @@ mod tests {
             &session_configuration,
             per_turn_config,
             model_info,
+            None,
             "turn_id".to_string(),
         );
 
@@ -6321,6 +6351,7 @@ mod tests {
             skills_manager,
             file_watcher,
             agent_control,
+            network_proxy: None,
             state_db: None,
             model_client: ModelClient::new(
                 Some(Arc::clone(&auth_manager)),
@@ -6344,6 +6375,7 @@ mod tests {
             &session_configuration,
             per_turn_config,
             model_info,
+            None,
             "turn_id".to_string(),
         ));
 

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

@@ -47,7 +47,6 @@ use crate::protocol::SandboxPolicy;
 use crate::windows_sandbox::WindowsSandboxLevelExt;
 use codex_app_server_protocol::Tools;
 use codex_app_server_protocol::UserSavedConfig;
-use codex_network_proxy::NetworkProxy;
 use codex_protocol::config_types::AltScreenMode;
 use codex_protocol::config_types::ForcedLoginMethod;
 use codex_protocol::config_types::ModeKind;
@@ -80,6 +79,7 @@ use toml_edit::DocumentMut;
 
 mod constraint;
 pub mod edit;
+mod network_proxy_spec;
 pub mod profile;
 pub mod schema;
 pub mod service;
@@ -88,6 +88,8 @@ pub use constraint::Constrained;
 pub use constraint::ConstraintError;
 pub use constraint::ConstraintResult;
 
+pub use network_proxy_spec::NetworkProxySpec;
+pub use network_proxy_spec::StartedNetworkProxy;
 pub use service::ConfigService;
 pub use service::ConfigServiceError;
 
@@ -154,7 +156,7 @@ pub struct Config {
     pub enforce_residency: Constrained<Option<ResidencyRequirement>>,
 
     /// Effective network configuration applied to all spawned processes.
-    pub network: Option<NetworkProxy>,
+    pub network: Option<NetworkProxySpec>,
 
     /// True if the user passed in an override or set a value in config.toml
     /// for either of approval_policy or sandbox_mode.
@@ -1657,7 +1659,7 @@ impl Config {
             mcp_servers,
             exec_policy: _,
             enforce_residency,
-            network: _network_requirements,
+            network: network_requirements,
         } = requirements;
 
         apply_requirement_constrained_value(
@@ -1682,6 +1684,20 @@ 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 network = NetworkProxySpec::from_constraints(&config_layer_stack, value)
+                    .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 +1710,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,

codex-rs/core/src/config/network_proxy_spec.rs

@@ -0,0 +1,167 @@
+use crate::config;
+use crate::config_loader::NetworkConstraints;
+use async_trait::async_trait;
+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::host_and_port_from_network_addr;
+use codex_network_proxy::validate_policy_against_constraints;
+use std::sync::Arc;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct NetworkProxySpec {
+    config: NetworkProxyConfig,
+    constraints: NetworkProxyConstraints,
+}
+
+pub struct StartedNetworkProxy {
+    proxy: NetworkProxy,
+    _handle: NetworkProxyHandle,
+}
+
+impl StartedNetworkProxy {
+    fn new(proxy: NetworkProxy, handle: NetworkProxyHandle) -> Self {
+        Self {
+            proxy,
+            _handle: handle,
+        }
+    }
+
+    pub fn proxy(&self) -> NetworkProxy {
+        self.proxy.clone()
+    }
+}
+
+#[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<Option<ConfigState>> {
+        Ok(None)
+    }
+
+    async fn reload_now(&self) -> anyhow::Result<ConfigState> {
+        Ok(self.state.clone())
+    }
+
+    fn source_label(&self) -> String {
+        "StaticNetworkProxyReloader".to_string()
+    }
+}
+
+impl NetworkProxySpec {
+    pub fn proxy_host_and_port(&self) -> String {
+        host_and_port_from_network_addr(&self.config.network.proxy_url, 3128)
+    }
+
+    pub(crate) fn from_constraints(
+        _config_layer_stack: &config::ConfigLayerStack,
+        requirements: NetworkConstraints,
+    ) -> std::io::Result<Self> {
+        // TODO(mbolin): Use ConfigLayerStack once we are ready to start
+        // honoring network configuration in config.toml.
+        let config = NetworkProxyConfig::default();
+        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,
+        })
+    }
+
+    pub async fn start_proxy(&self) -> std::io::Result<StartedNetworkProxy> {
+        let state =
+            build_config_state(self.config.clone(), self.constraints.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(StartedNetworkProxy::new(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)
+    }
+}

codex-rs/core/src/state/service.rs

@@ -5,6 +5,7 @@ use crate::RolloutRecorder;
 use crate::agent::AgentControl;
 use crate::analytics_client::AnalyticsEventsClient;
 use crate::client::ModelClient;
+use crate::config::StartedNetworkProxy;
 use crate::exec_policy::ExecPolicyManager;
 use crate::file_watcher::FileWatcher;
 use crate::hooks::Hooks;
@@ -38,6 +39,7 @@ pub(crate) struct SessionServices {
     pub(crate) skills_manager: Arc<SkillsManager>,
     pub(crate) file_watcher: Arc<FileWatcher>,
     pub(crate) agent_control: AgentControl,
+    pub(crate) network_proxy: Option<StartedNetworkProxy>,
     pub(crate) state_db: Option<StateDbHandle>,
     /// Session-scoped model client shared across turns.
     pub(crate) model_client: ModelClient,

codex-rs/core/src/tasks/user_shell.rs

@@ -147,7 +147,7 @@ pub(crate) async fn execute_user_shell_command(
             &turn_context.shell_environment_policy,
             Some(session.conversation_id),
         ),
-        network: turn_context.config.network.clone(),
+        network: turn_context.network.clone(),
         // TODO(zhao-oai): Now that we have ExecExpiration::Cancellation, we
         // should use that instead of an "arbitrarily large" timeout here.
         expiration: USER_SHELL_TIMEOUT_MS.into(),

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);

codex-rs/core/src/tools/handlers/unified_exec.rs

@@ -192,7 +192,7 @@ impl ToolHandler for UnifiedExecHandler {
                             yield_time_ms,
                             max_output_tokens,
                             workdir,
-                            network: context.turn.config.network.clone(),
+                            network: context.turn.network.clone(),
                             tty,
                             sandbox_permissions,
                             justification,