openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Honor network policy in remote control and recover remote execution (#47410)

Browse Source 
## What changed

- Route remote-control enrollment, pairing, client management, and WebSocket connections through the authenticated HTTP client factory. Keep policy denials distinct from authentication failures and wait for policy changes before reconnecting.
- Preserve policy denials through exec-server connection failures so later connection attempts and capability discovery can recover. Keep direct executors running through temporary policy outages while rejecting permanent policy failures.
- Build cached HTTP transports on blocking workers, serialize construction, and retain successful builds after callers time out so request deadlines cover queued construction.

## Testing

Add regression coverage for policy recovery during remote startup, direct registration, accepted WebSocket reconnection, and capability discovery; queued transport construction timeouts; and remote-control WebSocket fallback with invalid custom CA files.

GitOrigin-RevId: c81db18e83de957175bb213d165281ac701089dd
This commit is contained in:
acrognale-oai
2026-09-23 00:48:53 +00:00
committed by copyberry
parent 8f8ace78cb
commit db3cb33eca
27 changed files with 833 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/Cargo.toml
+2 -2
View File
@@ -7,8 +7,8 @@ license.workspace = true
[lib]
doctest = false
# Compiled through #[path] in src/client.rs; cargo-shear 1.11.2 does not normalize
# the ../tests path and incorrectly reports this linked unit-test module as unlinked.
# cargo-shear 1.11.2 misresolves this #[path] module in src/client.rs.
# The module is compiled and exercised by the exec-server unit test target.
[package.metadata.cargo-shear]
ignored-paths = ["tests/unit/client_provisioning_tests.rs"]
codex-rs/exec-server/src/capability_discovery_cache.rs
+1 -1
View File
@@ -187,7 +187,7 @@ impl ExecutorCapabilityDiscoveryCache {
let response = match environment.discover_capability_roots(params).await {
Ok(response) => response,
Err(error) => {
let retryable = crate::client::is_retryable_recovery_error(&error);
let retryable = crate::client::can_retry_connection_attempt(&error);
let error = error.to_string();
return selected_roots
.into_iter()
codex-rs/exec-server/src/client.rs
+15 -11
View File
@@ -148,9 +148,14 @@ mod provisioning_tests;
mod recovery;
#[path = "client_refresh.rs"]
mod refresh;
pub(crate) use connection_failure::can_retry_connection_attempt;
#[cfg(test)]
pub(crate) use recovery::is_environment_offline_error;
pub(crate) use recovery::is_retryable_recovery_error;
#[path = "client/connection_failure.rs"]
mod connection_failure;
use connection_failure::ConnectionFailure;
pub(crate) use recovery::is_retryable_registry_error;
pub(crate) use recovery::registry_recovery_retry_delay;
use refresh::ConnectionAttempt;
@@ -300,7 +305,7 @@ struct ConnectionState {
enum ConnectionStatus {
Connected(Arc<RpcClient>),
Recovering,
Failed(String),
Failed(ConnectionFailure),
}
impl ConnectionState {
@@ -321,7 +326,8 @@ impl ConnectionState {
let _ = self
.environment_connection_state_tx
.send_if_modified(|current| {
if *current == state {
// A terminal failure must wake callers waiting on recovery.
if *current == state && !matches!(self.status, ConnectionStatus::Failed(_)) {
false
} else {
*current = state;
@@ -539,7 +545,7 @@ impl LazyRemoteExecServerClient {
|| self.startup.result.get().is_some_and(|result| {
result
.as_ref()
.is_err_and(|error| recovery::is_retryable_recovery_error(error))
.is_err_and(|error| can_retry_connection_attempt(error))
})) {
Box::pin(self.reconnect()).await
} else {
@@ -666,6 +672,8 @@ pub enum ExecServerError {
WebSocketConfiguration(String),
#[error("timed out waiting for exec-server initialize handshake after {timeout:?}")]
InitializeTimedOut { timeout: Duration },
#[error(transparent)]
ApplicationNetworkPolicy(#[from] codex_http_client::NetworkPolicyDenied),
#[error("exec-server transport closed")]
Closed,
#[error("{0}")]
@@ -744,9 +752,7 @@ impl ExecServerClient {
ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => Err(
ExecServerError::Disconnected("exec-server environment is recovering".to_string()),
),
ConnectionStatus::Failed(message) => {
Err(ExecServerError::Disconnected(message.clone()))
}
ConnectionStatus::Failed(message) => Err(message.clone().into()),
}
}
@@ -1178,9 +1184,7 @@ impl ExecServerClient {
Some(Ok(()))
}
ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => None,
ConnectionStatus::Failed(message) => {
Some(Err(ExecServerError::Disconnected(message.clone())))
}
ConnectionStatus::Failed(message) => Some(Err(message.clone().into())),
}
}
@@ -1680,8 +1684,8 @@ impl Inner {
// Do not register a process session that can never receive environment
// notifications. Without this check, remote MCP startup could create a
// dead session and wait for process output that will never arrive.
if let Some(message) = self.failure_message() {
return Err(ExecServerError::Disconnected(message));
if let Some(message) = self.connection_failure() {
return Err(message.into());
}
let sessions = self.sessions.load();
if sessions.contains_key(process_id) {
codex-rs/exec-server/src/client/accepted.rs
+2 -2
View File
@@ -196,7 +196,7 @@ impl Inner {
.unwrap_or_else(std::sync::PoisonError::into_inner);
let current_rpc_client = match &connection.status {
ConnectionStatus::Failed(message) => {
return Err(ExecServerError::Disconnected(message.clone()));
return Err(message.clone().into());
}
ConnectionStatus::Connected(rpc_client) => Some(Arc::clone(rpc_client)),
ConnectionStatus::Recovering => None,
@@ -221,7 +221,7 @@ impl Inner {
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let ConnectionStatus::Failed(message) = &connection_state.status {
return Err(ExecServerError::Disconnected(message.clone()));
return Err(message.clone().into());
}
replacement_submission.submit(connection)
}
codex-rs/exec-server/src/client/connection_failure.rs
+66
View File
@@ -0,0 +1,66 @@
//! Preserve policy denials through terminal connection state and allow later policy recovery.
use codex_http_client::NetworkPolicyDenied;
use codex_http_client::RouteAwareRequestError;
use super::ExecServerError;
#[derive(Clone, thiserror::Error, Debug)]
pub(super) enum ConnectionFailure {
#[error("{0}")]
Disconnected(String),
#[error(transparent)]
Policy(NetworkPolicyDenied),
}
impl From<ConnectionFailure> for ExecServerError {
fn from(failure: ConnectionFailure) -> Self {
match failure {
ConnectionFailure::Disconnected(message) => Self::Disconnected(message),
ConnectionFailure::Policy(denial) => Self::ApplicationNetworkPolicy(denial),
}
}
}
impl ExecServerError {
pub(crate) fn application_network_policy_denial(&self) -> Option<NetworkPolicyDenied> {
match self {
Self::ConnectionAttempt(error) => error.application_network_policy_denial(),
Self::ApplicationNetworkPolicy(denial)
| Self::EnvironmentRegistryRequest(RouteAwareRequestError::Policy(denial)) => {
Some(*denial)
}
Self::WebSocketConnect { source, .. } => {
codex_websocket_client::network_policy_denial(source)
}
Self::Spawn(_)
| Self::WebSocketConnectTimeout { .. }
| Self::WebSocketConfiguration(_)
| Self::InitializeTimedOut { .. }
| Self::Closed
| Self::Disconnected(_)
| Self::ProvisioningFailed(_)
| Self::Json(_)
| Self::HttpRequest(_)
| Self::Protocol(_)
| Self::ProvisioningModeConflict { .. }
| Self::Server { .. }
| Self::EnvironmentRegistryHttp { .. }
| Self::EnvironmentRegistryConfig(_)
| Self::EnvironmentRegistryAuth(_)
| Self::EnvironmentRegistryRequest(
RouteAwareRequestError::Request(_)
| RouteAwareRequestError::Route(_)
| RouteAwareRequestError::Build(_)
| RouteAwareRequestError::UnsupportedRedirectScheme(_)
| RouteAwareRequestError::TooManyRedirects
| RouteAwareRequestError::Timeout,
) => None,
}
}
}
pub(crate) fn can_retry_connection_attempt(error: &ExecServerError) -> bool {
error.application_network_policy_denial().is_some()
|| super::recovery::is_retryable_recovery_error(error)
}
codex-rs/exec-server/src/client_recovery.rs
+25 -13
View File
@@ -253,8 +253,8 @@ impl Inner {
pub(super) async fn rpc_client(self: &Arc<Self>) -> Result<Arc<RpcClient>, ExecServerError> {
let mut connection_changed = self.connection_changed.subscribe();
loop {
if let Some(message) = self.failure_message() {
return Err(ExecServerError::Disconnected(message));
if let Some(message) = self.connection_failure() {
return Err(message.into());
}
let rpc_client = {
@@ -310,10 +310,10 @@ impl Inner {
}
pub(super) fn is_failed(&self) -> bool {
self.failure_message().is_some()
self.connection_failure().is_some()
}
pub(super) fn failure_message(&self) -> Option<String> {
pub(super) fn connection_failure(&self) -> Option<super::ConnectionFailure> {
let connection = self
.connection
.lock()
@@ -372,11 +372,13 @@ impl Inner {
let message = format!(
"{disconnect_message}; failed to resume exec-server session: recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}"
);
self.fail(message).await;
self.fail(super::ConnectionFailure::Disconnected(message))
.await;
return;
}
if self.reconnect_strategy.is_none() {
self.fail(disconnect_message).await;
self.fail(super::ConnectionFailure::Disconnected(disconnect_message))
.await;
return;
}
@@ -384,7 +386,8 @@ impl Inner {
let message = format!(
"{disconnect_message}; failed to resume exec-server session: missing session id"
);
self.fail(message).await;
self.fail(super::ConnectionFailure::Disconnected(message))
.await;
return;
};
let uses_registry_backoff = matches!(
@@ -399,10 +402,15 @@ impl Inner {
return;
}
}
Ok(Err(error)) if !is_retryable_recovery_error(&error) => {
break error.to_string();
Ok(Err(error)) => {
if let Some(denial) = error.application_network_policy_denial() {
self.fail(super::ConnectionFailure::Policy(denial)).await;
return;
}
if !is_retryable_recovery_error(&error) {
break error.to_string();
}
}
Ok(Err(_)) => {}
Err(_) => {
break format!("recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}");
}
@@ -425,7 +433,8 @@ impl Inner {
let message =
format!("{disconnect_message}; failed to resume exec-server session: {last_error}");
self.fail(message).await;
self.fail(super::ConnectionFailure::Disconnected(message))
.await;
}
async fn wait_for_process_starts(&self) {
@@ -570,7 +579,7 @@ impl Inner {
Ok(())
}
async fn fail(self: &Arc<Self>, message: String) {
async fn fail(self: &Arc<Self>, message: super::ConnectionFailure) {
let (message, newly_failed) = {
let mut connection = self
.connection
@@ -586,7 +595,7 @@ impl Inner {
};
if newly_failed {
self.notify_connection_changed();
fail_all_in_flight_work(self, message.clone()).await;
fail_all_in_flight_work(self, message.to_string()).await;
}
}
}
@@ -849,6 +858,9 @@ impl ExecServerClient {
}
pub(crate) fn is_retryable_recovery_error(error: &ExecServerError) -> bool {
if error.application_network_policy_denial().is_some() {
return false;
}
if let ExecServerError::ConnectionAttempt(error) = error {
return is_retryable_recovery_error(error.as_ref());
}
codex-rs/exec-server/src/client_refresh.rs
+3 -1
View File
@@ -269,7 +269,9 @@ impl Inner {
ConnectionStatus::Recovering | ConnectionStatus::Failed(_) => None,
};
self.retired.cancel();
connection.set_status(ConnectionStatus::Failed(message.clone()));
connection.set_status(ConnectionStatus::Failed(
super::ConnectionFailure::Disconnected(message.clone()),
));
rpc_client
};
self.connection_changed.send_replace(());
codex-rs/exec-server/src/client_refresh_tests.rs
+55
View File
@@ -881,3 +881,58 @@ async fn retirement_rejects_pending_process_start_before_stream_cleanup() -> Res
}
Ok(())
}
#[tokio::test]
async fn application_policy_registry_denial_does_not_poison_later_noise_startup() -> Result<()> {
let _clock = freeze_clock();
let executor = Executor::start(Validator::default()).await?;
let registry = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("POST"))
.and(wiremock::matchers::path(
"/cloud/environment/environment/connect",
))
.respond_with(
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
"environment_id": "environment",
"url": executor.target.url,
"security_profile": "noise_hybrid_ik_v1",
"executor_registration_id": executor.target.registration,
"executor_public_key": executor.target.identity.public_key(),
"harness_key_authorization": "authorization",
})),
)
.expect(1)
.mount(&registry)
.await;
let controller = codex_http_client::NetworkPolicyController::default();
let policy = controller.policy();
let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
.with_network_policy(policy.clone());
let provider = crate::remote::NoiseRendezvousEnvironmentConfig::new(
registry.uri(),
"environment".into(),
"registry-token".into(),
/*chatgpt_account_id*/ None,
)?
.into_connect_provider(factory.clone())?;
let client = LazyRemoteExecServerClient::new(
ExecServerTransportParams::NoiseRendezvous {
provider,
identity: NoiseChannelIdentity::generate()?,
},
factory,
);
let error = client.get().await.err().unwrap();
assert_eq!(
error.application_network_policy_denial(),
Some(codex_http_client::NetworkPolicyDenied::Unavailable)
);
assert!(!crate::client::is_retryable_recovery_error(&error));
assert!(registry.received_requests().await.unwrap().is_empty());
controller.publish(
policy.revision(),
codex_http_client::DestinationPolicy::Unrestricted,
);
client.get().await?;
Ok(())
}
codex-rs/exec-server/src/environment.rs
+18 -9
View File
@@ -1021,19 +1021,28 @@ impl Environment {
tracing::warn!(%error, "replaying capability discovery after executor recovery");
let recovered =
tokio::time::timeout(std::time::Duration::from_secs(8), async {
while self.readiness_result().is_none_or(|result| result.is_err()) {
loop {
match self.readiness_result() {
Some(Ok(())) => return Some(Ok(())),
Some(Err(error))
if !crate::client::is_retryable_recovery_error(
&error,
) =>
{
return Some(Err(error));
}
Some(Err(_)) | None => {}
}
if connection_state.changed().await.is_err() {
return false;
return None;
}
}
true
})
.await
.unwrap_or(false);
if recovered {
discover().await
} else {
Err(error)
.await;
match recovered {
Ok(Some(Ok(()))) => discover().await,
Ok(Some(Err(error))) => Err(error),
Ok(None) | Err(_) => Err(error),
}
}
response => response,
codex-rs/exec-server/src/remote/direct.rs
+54 -3
View File
@@ -2,6 +2,8 @@ use std::time::Instant;
use codex_api::AuthError;
use codex_api::AuthProvider;
use codex_http_client::NetworkPolicy;
use codex_http_client::NetworkPolicyDenied;
use codex_http_client::Request;
use codex_websocket_client::WebSocketConnection;
use codex_websocket_client::WebSocketConnector;
@@ -43,6 +45,26 @@ struct DirectRegistrationResponse {
}
impl EnvironmentRegistryClient {
async fn register_direct_environment_when_available(
&self,
environment_id: &str,
policy: &NetworkPolicy,
) -> Result<DirectRegistrationResponse, ExecServerError> {
let destination = super::endpoint_url(
&self.base_url,
&format!("/cloud/environment/{environment_id}/direct/register"),
);
let mut retry_attempt = 0;
loop {
match self.register_direct_environment(environment_id).await {
Err(error) if is_policy_recoverable(&error, policy, &destination) => {}
result => return result,
}
sleep(registry_recovery_retry_delay(environment_id, retry_attempt)).await;
retry_attempt = retry_attempt.saturating_add(1);
}
}
#[tracing::instrument(
name = "codex.exec_server.remote.register",
skip_all,
@@ -122,7 +144,10 @@ pub(super) async fn run_direct_environment(
require_tls_or_loopback(&config.base_url, "https")?;
let mut retry_attempt = 0;
let mut registration = client
.register_direct_environment(&config.environment_id)
.register_direct_environment_when_available(
&config.environment_id,
config.http_client_factory.network_policy(),
)
.await?;
loop {
@@ -155,7 +180,12 @@ pub(super) async fn run_direct_environment(
config.telemetry.remote_reconnect("disconnected");
}
Err(error)
if is_retryable_recovery_error(&error)
if (is_retryable_recovery_error(&error)
|| is_policy_recoverable(
&error,
config.http_client_factory.network_policy(),
&registration.url,
))
&& !matches!(
&error,
ExecServerError::WebSocketConnect {
@@ -180,7 +210,10 @@ pub(super) async fn run_direct_environment(
} if response.status() == StatusCode::CONFLICT
) {
registration = client
.register_direct_environment(&config.environment_id)
.register_direct_environment_when_available(
&config.environment_id,
config.http_client_factory.network_policy(),
)
.await?;
}
warn!("direct exec-server connection failed; retrying");
@@ -198,6 +231,24 @@ pub(super) async fn run_direct_environment(
}
}
fn is_policy_recoverable(
error: &ExecServerError,
policy: &NetworkPolicy,
destination: &str,
) -> bool {
match error.application_network_policy_denial() {
Some(NetworkPolicyDenied::Unavailable) => true,
Some(NetworkPolicyDenied::Revoked) => url::Url::parse(destination).is_ok_and(|url| {
matches!(
policy.acquire(&url),
Ok(_) | Err(NetworkPolicyDenied::Unavailable)
)
}),
Some(NetworkPolicyDenied::Destination | NetworkPolicyDenied::UnsupportedTransport)
| None => false,
}
}
async fn connect_direct(
url: &str,
auth_provider: &dyn AuthProvider,
codex-rs/exec-server/src/remote/direct_tests.rs
+25 -4
View File
@@ -6,7 +6,9 @@ use std::time::Instant;
use anyhow::Result;
use codex_api::AuthProvider;
use codex_http_client::DestinationPolicy;
use codex_http_client::HttpClientFactory;
use codex_http_client::NetworkPolicyController;
use codex_http_client::OutboundProxyPolicy;
use http::HeaderMap;
use http::HeaderValue;
@@ -282,8 +284,14 @@ async fn direct_registration_uses_proxy_policy_without_logging_secrets() -> Resu
}
#[tokio::test]
async fn direct_registration_failure_stops_initial_and_conflict_attempts() -> Result<()> {
for successful_registrations in [0, 1] {
async fn direct_registration_handles_policy_outages_and_permanent_failures() -> Result<()> {
for (successful_registrations, policy_unavailable) in [(0, false), (1, false), (0, true)] {
let controller = NetworkPolicyController::default();
let policy = controller.policy();
controller.publish(policy.revision(), DestinationPolicy::Unrestricted);
if policy_unavailable && successful_registrations == 0 {
controller.unavailable(policy.revision());
}
let listener = TcpListener::bind("127.0.0.1:0").await?;
let websocket_url = format!("ws://{}/connect", listener.local_addr()?);
let registry = MockServer::start().await;
@@ -294,7 +302,7 @@ async fn direct_registration_failure_stops_initial_and_conflict_attempts() -> Re
))
.respond_with(move |_: &wiremock::Request| {
let attempt = registration_count.fetch_add(1, Ordering::Relaxed);
if attempt < successful_registrations {
if policy_unavailable || attempt < successful_registrations {
ResponseTemplate::new(200).set_body_json(serde_json::json!({
"environment_id": "environment-requested",
"transport": DIRECT_TRANSPORT,
@@ -313,7 +321,8 @@ async fn direct_registration_failure_stops_initial_and_conflict_attempts() -> Re
"environment-requested".to_string(),
RemoteEnvironmentTransport::Direct,
Arc::new(StaticAuthProvider),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
.with_network_policy(policy.clone().for_current_account()),
)?;
let runtime_paths = ExecServerRuntimePaths::new(
std::env::current_exe()?,
@@ -329,6 +338,18 @@ async fn direct_registration_failure_stops_initial_and_conflict_attempts() -> Re
.await?;
socket.shutdown().await?;
}
if policy_unavailable {
timeout(Duration::from_secs(1), listener.accept())
.await
.expect_err("unavailable policy opened a connection");
controller.publish(policy.revision(), DestinationPolicy::Unrestricted);
let (socket, _) = timeout(Duration::from_secs(5), listener.accept()).await??;
let _websocket = tokio_tungstenite::accept_async(socket).await?;
registry.verify().await;
task.abort();
let _ = task.await;
continue;
}
let error = timeout(Duration::from_secs(5), task)
.await??
.expect_err("registration failure should stop the runner");
codex-rs/exec-server/tests/accepted_websocket.rs
+31 -2
View File
@@ -40,7 +40,10 @@ use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCNotification;
use codex_exec_server_protocol::JSONRPCRequest;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_http_client::DestinationPolicy;
use codex_http_client::HttpClientFactory;
use codex_http_client::NetworkPolicy;
use codex_http_client::NetworkPolicyController;
use codex_http_client::OutboundProxyPolicy;
use codex_utils_path_uri::PathUri;
use common::exec_server::DisconnectableWebSocketProxy;
@@ -132,8 +135,12 @@ async fn accepted_websocket_environment_info_uses_initialization_metadata() -> R
Ok(())
}
#[test_case::test_case(false; "transport_disconnect")]
#[test_case::test_case(true; "policy_unavailable")]
#[tokio::test]
async fn accepted_websocket_interoperates_and_recovers_with_real_direct_executor() -> Result<()> {
async fn accepted_websocket_interoperates_and_recovers_with_real_direct_executor(
policy_unavailable: bool,
) -> Result<()> {
let (websocket_url, mut accepted_sockets, server_task) = start_acceptor().await?;
let proxy = DisconnectableWebSocketProxy::new(&websocket_url).await?;
let registry = MockServer::start().await;
@@ -150,7 +157,15 @@ async fn accepted_websocket_interoperates_and_recovers_with_real_direct_executor
.mount(&registry)
.await;
let http_client_factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault);
let controller = NetworkPolicyController::default();
let policy = controller.policy();
controller.publish(policy.revision(), DestinationPolicy::Unrestricted);
let http_client_factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
.with_network_policy(if policy_unavailable {
policy.clone().for_current_account()
} else {
NetworkPolicy::unmanaged()
});
let config = RemoteEnvironmentConfig::new_with_transport(
registry.uri(),
"environment-1".to_string(),
@@ -289,6 +304,19 @@ async fn accepted_websocket_interoperates_and_recovers_with_real_direct_executor
EnvironmentConnectionState::Connected
);
proxy.pause_and_disconnect().await?;
if policy_unavailable {
controller.unavailable(policy.revision());
assert!(
timeout(Duration::from_secs(/*secs*/ 1), accepted_sockets.recv())
.await
.is_err(),
"unavailable policy reconnected"
);
assert!(
!executor_task.is_finished(),
"policy outage stopped the runner"
);
}
timeout(
TEST_TIMEOUT,
connection_state.wait_for(|state| *state == EnvironmentConnectionState::Disconnected),
@@ -319,6 +347,7 @@ async fn accepted_websocket_interoperates_and_recovers_with_real_direct_executor
.is_err(),
"process reads should wait while recovery is in progress"
);
controller.publish(policy.revision(), DestinationPolicy::Unrestricted);
proxy.resume()?;
let replacement = timeout(TEST_TIMEOUT, accepted_sockets.recv())
.await?
codex-rs/exec-server/tests/environment.rs
+49
View File
@@ -453,3 +453,52 @@ async fn capability_discovery_retries_after_executor_reconnects() -> anyhow::Res
assert!(!cache.take_recovered_discovery());
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn capability_discovery_retries_after_application_network_policy_recovers()
-> anyhow::Result<()> {
let server = exec_server().await?;
let controller = codex_http_client::NetworkPolicyController::default();
let policy = controller.policy();
let manager = Arc::new(EnvironmentManager::without_environments(
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
.with_network_policy(policy.clone()),
));
manager.upsert_environment(
"recovering".to_string(),
server.websocket_url().to_string(),
/*connect_timeout*/ None,
)?;
let cache = ExecutorCapabilityDiscoveryCache::new(manager);
let skill_root = tempfile::tempdir()?;
let selected_roots = vec![SelectedCapabilityRoot {
id: "recovering-skill".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: "recovering".to_string(),
path: PathUri::from_host_native_path(skill_root.path())?,
},
}];
let failed = cache.snapshot(&selected_roots, &HashMap::new()).await;
let error = failed.roots()[0].result.as_ref().unwrap_err();
assert!(error.contains("application network policy"), "{error}");
assert!(!cache.take_recovered_discovery());
controller.publish(
policy.revision(),
codex_http_client::DestinationPolicy::Unrestricted,
);
let recovered = timeout(
Duration::from_secs(/*secs*/ 10),
cache.snapshot(&selected_roots, &HashMap::new()),
)
.await?;
let discovery = recovered.roots()[0]
.result
.as_ref()
.map_err(|error| anyhow::anyhow!("{error}"))?;
assert_eq!(discovery.id, "recovering-skill");
assert!(cache.take_recovered_discovery());
assert!(!cache.take_recovered_discovery());
Ok(())
}