diff --git a/codex-rs/network-proxy/src/credentialed_routes.rs b/codex-rs/network-proxy/src/credentialed_routes.rs index 008d4d8f79..2b4af8ae51 100644 --- a/codex-rs/network-proxy/src/credentialed_routes.rs +++ b/codex-rs/network-proxy/src/credentialed_routes.rs @@ -132,11 +132,6 @@ impl CredentialedRoutesReloader { } } - pub async fn replace_base_state(&self, base_state: ConfigState) -> Result { - *self.base_state.write().await = base_state; - self.reload_now().await - } - pub async fn current_routes(&self) -> CredentialedRoutesConfig { self.credentialed_routes.read().await.clone() } @@ -193,6 +188,19 @@ impl ConfigReloader for CredentialedRoutesReloader { self.build_state(&credentialed_routes).await }) } + + fn replace_base_state(&self, base_state: ConfigState) -> ConfigReloaderFuture<'_, ConfigState> { + Box::pin(async move { + let credentialed_routes = self.credentialed_routes.read().await.clone(); + let state = build_config_state_with_credentialed_routes( + base_state.config.clone(), + base_state.constraints.clone(), + &credentialed_routes, + )?; + *self.base_state.write().await = base_state; + Ok(state) + }) + } } fn build_config_state_with_credentialed_routes( diff --git a/codex-rs/network-proxy/src/credentialed_routes_tests.rs b/codex-rs/network-proxy/src/credentialed_routes_tests.rs index c3104951eb..bfd928380e 100644 --- a/codex-rs/network-proxy/src/credentialed_routes_tests.rs +++ b/codex-rs/network-proxy/src/credentialed_routes_tests.rs @@ -157,3 +157,39 @@ async fn failed_route_state_build_keeps_the_last_working_routes() { CredentialedRoutesConfig::default() ); } + +#[tokio::test] +async fn failed_base_state_replacement_keeps_the_last_working_base() { + let base_state = build_config_state( + NetworkProxyConfig::default(), + NetworkProxyConstraints::default(), + ) + .unwrap(); + let credentialed_routes = + proxy_config(vec![route("connector_123", "https://api.example.com/v1")]); + let reloader = CredentialedRoutesReloader::new( + base_state, + credentialed_routes, + Arc::new(|| async { Ok(CredentialedRoutesConfig::default()) }), + ); + let mut constrained_config = NetworkProxyConfig::default(); + constrained_config + .network + .set_allowed_domains(vec!["existing.example.com".to_string()]); + let constrained_state = build_config_state( + constrained_config, + NetworkProxyConstraints { + allowed_domains: Some(vec!["existing.example.com".to_string()]), + ..NetworkProxyConstraints::default() + }, + ) + .unwrap(); + + assert!( + reloader + .replace_base_state(constrained_state) + .await + .is_err() + ); + assert!(reloader.reload_now().await.is_ok()); +} diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index c3685a4310..47f08c5dc7 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -655,27 +655,8 @@ impl NetworkProxy { } pub async fn replace_config_state(&self, new_state: ConfigState) -> Result<()> { - let current_cfg = self.state.current_cfg().await?; - anyhow::ensure!( - new_state.config.network.enabled == current_cfg.network.enabled, - "cannot update network.enabled on a running proxy" - ); - anyhow::ensure!( - new_state.config.network.proxy_url == current_cfg.network.proxy_url, - "cannot update network.proxy_url on a running proxy" - ); - anyhow::ensure!( - new_state.config.network.socks_url == current_cfg.network.socks_url, - "cannot update network.socks_url on a running proxy" - ); - anyhow::ensure!( - new_state.config.network.enable_socks5 == current_cfg.network.enable_socks5, - "cannot update network.enable_socks5 on a running proxy" - ); - anyhow::ensure!( - new_state.config.network.enable_socks5_udp == current_cfg.network.enable_socks5_udp, - "cannot update network.enable_socks5_udp on a running proxy" - ); + self.ensure_runtime_config_compatible(&new_state.config) + .await?; let settings = NetworkProxyRuntimeSettings::from_config(&new_state.config)?; self.state.replace_config_state(new_state).await?; @@ -687,6 +668,41 @@ impl NetworkProxy { Ok(()) } + pub async fn replace_base_config_state(&self, base_state: ConfigState) -> Result<()> { + self.ensure_runtime_config_compatible(&base_state.config) + .await?; + let new_state = self.state.replace_base_state(base_state).await?; + self.replace_config_state(new_state).await + } + + async fn ensure_runtime_config_compatible( + &self, + new_cfg: &config::NetworkProxyConfig, + ) -> Result<()> { + let current_cfg = self.state.current_cfg().await?; + anyhow::ensure!( + new_cfg.network.enabled == current_cfg.network.enabled, + "cannot update network.enabled on a running proxy" + ); + anyhow::ensure!( + new_cfg.network.proxy_url == current_cfg.network.proxy_url, + "cannot update network.proxy_url on a running proxy" + ); + anyhow::ensure!( + new_cfg.network.socks_url == current_cfg.network.socks_url, + "cannot update network.socks_url on a running proxy" + ); + anyhow::ensure!( + new_cfg.network.enable_socks5 == current_cfg.network.enable_socks5, + "cannot update network.enable_socks5 on a running proxy" + ); + anyhow::ensure!( + new_cfg.network.enable_socks5_udp == current_cfg.network.enable_socks5_udp, + "cannot update network.enable_socks5_udp on a running proxy" + ); + Ok(()) + } + fn runtime_settings(&self) -> NetworkProxyRuntimeSettings { self.runtime_settings .read() diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index e721d0de70..5e82e00137 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -177,6 +177,11 @@ pub trait ConfigReloader: Send + Sync { /// Force a reload, regardless of whether a change was detected. fn reload_now(&self) -> ConfigReloaderFuture<'_, ConfigState>; + + /// Rebuild state after the caller replaces the proxy's base configuration. + fn replace_base_state(&self, base_state: ConfigState) -> ConfigReloaderFuture<'_, ConfigState> { + Box::pin(async move { Ok(base_state) }) + } } pub type ConfigReloaderFuture<'a, T> = Pin> + Send + 'a>>; @@ -352,6 +357,10 @@ impl NetworkProxyState { Ok(()) } + pub(crate) async fn replace_base_state(&self, base_state: ConfigState) -> Result { + self.reloader.replace_base_state(base_state).await + } + pub async fn host_blocked(&self, host: &str, port: u16) -> Result { self.reload_if_needed().await?; let host = match Host::parse(host) {