From d87cfda0accd05c065739efd3ce1dbf3aec2b544 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 9 Jul 2026 13:18:54 +0100 Subject: [PATCH] Invalidate thread skills when replacing environments --- .../environment_processor.rs | 16 ++++++--- .../suite/v2/selected_capability_stack.rs | 34 ++++++++++++++++++- codex-rs/exec-server/src/environment.rs | 13 ++++--- 3 files changed, 52 insertions(+), 11 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/environment_processor.rs b/codex-rs/app-server/src/request_processors/environment_processor.rs index 483f389820..31fff8b6dc 100644 --- a/codex-rs/app-server/src/request_processors/environment_processor.rs +++ b/codex-rs/app-server/src/request_processors/environment_processor.rs @@ -25,7 +25,11 @@ impl EnvironmentRequestProcessor { }; for (environment_id, environment) in processor.environment_manager.registered_environments() { - processor.notify_selected_threads_on_readiness_changes(environment_id, environment); + processor.notify_selected_threads_on_readiness_changes( + environment_id, + environment, + /*notify_initially*/ false, + ); } processor } @@ -35,7 +39,8 @@ impl EnvironmentRequestProcessor { params: EnvironmentAddParams, ) -> Result, JSONRPCErrorError> { let environment_id = params.environment_id; - self.environment_manager + let replaced = self + .environment_manager .upsert_environment( environment_id.clone(), params.exec_server_url, @@ -46,7 +51,7 @@ impl EnvironmentRequestProcessor { .environment_manager .get_environment(&environment_id) .ok_or_else(|| internal_error("upserted environment is unavailable"))?; - self.notify_selected_threads_on_readiness_changes(environment_id, environment); + self.notify_selected_threads_on_readiness_changes(environment_id, environment, replaced); Ok(Some(EnvironmentAddResponse {}.into())) } @@ -54,16 +59,17 @@ impl EnvironmentRequestProcessor { &self, environment_id: String, environment: Arc, + notify_initially: bool, ) { let Some(mut readiness_changed) = environment.observe_readiness() else { return; }; - let startup_finished = environment.startup_finished(); + let notify_initially = notify_initially || environment.startup_finished(); let thread_manager = Arc::downgrade(&self.thread_manager); let outgoing = Arc::downgrade(&self.outgoing); let thread_state_manager = self.thread_state_manager.clone(); tokio::spawn(async move { - if !startup_finished && readiness_changed.changed().await.is_err() { + if !notify_initially && readiness_changed.changed().await.is_err() { return; } loop { diff --git a/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs b/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs index 91718fd1a6..4a63d1f3f4 100644 --- a/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs +++ b/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs @@ -285,7 +285,7 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() &mut app_server, &thread_id, &format!("Use ${SKILL_NAME} after reattaching the selected executor"), - fixture.environment_cwd, + fixture.environment_cwd.clone(), ) .await?; let resumed_mcp_pid = wait_for_pid_file(&fixture.pid_file).await?; @@ -308,6 +308,38 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() assert!(output.contains("ECHOING: hello from the selected executor")); assert!(output.contains(EXECUTOR_ENV_VALUE)); + let replacement_thread_id = start_thread( + &mut app_server, + fixture.selected_root.clone(), + fixture.environment_cwd.clone(), + ) + .await?; + assert_eq!( + list_thread_skills(&mut app_server, &replacement_thread_id).await?, + (vec![fixture.thread_skill.clone()], Vec::new()) + ); + let stalled_listener = std::net::TcpListener::bind("127.0.0.1:0")?; + add_environment( + &mut app_server, + &format!("ws://{}", stalled_listener.local_addr()?), + ) + .await?; + timeout( + Duration::from_secs(2), + wait_for_thread_skills_changed(&mut app_server, &replacement_thread_id), + ) + .await + .context("environment replacement did not invalidate thread skills")??; + assert_eq!( + timeout( + Duration::from_secs(2), + list_thread_skills(&mut app_server, &replacement_thread_id), + ) + .await + .context("skills/list waited for the replacement environment")??, + (Vec::new(), Vec::new()) + ); + exec_server.kill().await?; apps_server_handle.abort(); let _ = apps_server_handle.await; diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 20cd475b97..9735866894 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -298,13 +298,14 @@ impl EnvironmentManager { /// Adds or replaces a named remote environment without changing the /// manager's default environment selection. Uses the default WebSocket - /// connection timeout when none is provided. + /// connection timeout when none is provided. Returns whether an existing + /// environment was replaced. pub fn upsert_environment( &self, environment_id: String, exec_server_url: String, connect_timeout: Option, - ) -> Result<(), ExecServerError> { + ) -> Result { if environment_id.is_empty() { return Err(ExecServerError::Protocol( "environment id cannot be empty".to_string(), @@ -329,11 +330,13 @@ impl EnvironmentManager { self.local_runtime_paths.clone(), )); environment.start_connecting(); - self.environments + let replaced = self + .environments .write() .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(environment_id, environment); - Ok(()) + .insert(environment_id, environment) + .is_some(); + Ok(replaced) } /// Adds or replaces a named remote environment that connects through an