From 5163850025943b3e97dbf9c2d75235013b4f9116 Mon Sep 17 00:00:00 2001 From: Max Johnson <162359438+maxj-oai@users.noreply.github.com> Date: Tue, 24 Feb 2026 16:27:59 -0800 Subject: [PATCH 1/4] codex-rs/app-server: graceful websocket restart on Ctrl-C (#12517) ## Summary - add graceful websocket app-server restart on Ctrl-C by draining until no assistant turns are running - stop the websocket acceptor and disconnect existing connections once the drain condition is met - add a websocket integration test that verifies Ctrl-C waits for an in-flight turn before exit ## Verification - `cargo check -p codex-app-server --quiet` - `cargo test -p codex-app-server --test all suite::v2::connection_handling_websocket` - I (maxj) tested remote and local Codex.app --------- Co-authored-by: Codex --- .../app-server/src/codex_message_processor.rs | 5 + codex-rs/app-server/src/lib.rs | 150 ++++++++++- codex-rs/app-server/src/message_processor.rs | 6 + codex-rs/app-server/src/thread_status.rs | 58 ++++- codex-rs/app-server/src/transport.rs | 49 ++-- .../suite/v2/connection_handling_websocket.rs | 28 ++- .../v2/connection_handling_websocket_unix.rs | 237 ++++++++++++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 2 + 8 files changed, 493 insertions(+), 42 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 2d18ff7a51..b2f9f528bd 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -269,6 +269,7 @@ use std::time::SystemTime; use tokio::sync::Mutex; use tokio::sync::broadcast; use tokio::sync::oneshot; +use tokio::sync::watch; use toml::Value as TomlValue; use tracing::error; use tracing::info; @@ -2796,6 +2797,10 @@ impl CodexMessageProcessor { .await; } + pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver { + self.thread_watch_manager.subscribe_running_turn_count() + } + /// Best-effort: ensure initialized connections are subscribed to this thread. pub(crate) async fn try_attach_thread_listener( &mut self, diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 1bdbf97a13..b44661fd01 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -98,6 +98,72 @@ enum OutboundControlEvent { }, /// Remove state for a closed/disconnected connection. Closed { connection_id: ConnectionId }, + /// Disconnect all connection-oriented clients during graceful restart. + DisconnectAll, +} + +#[derive(Default)] +struct ShutdownState { + requested: bool, + forced: bool, + last_logged_running_turn_count: Option, +} + +enum ShutdownAction { + Noop, + Finish, +} + +impl ShutdownState { + fn requested(&self) -> bool { + self.requested + } + + fn forced(&self) -> bool { + self.forced + } + + fn on_ctrl_c(&mut self, connection_count: usize, running_turn_count: usize) { + if self.requested { + self.forced = true; + return; + } + + self.requested = true; + self.last_logged_running_turn_count = None; + info!( + "received Ctrl-C; entering graceful restart drain (connections={}, runningAssistantTurns={}, requests still accepted until no assistant turns are running)", + connection_count, running_turn_count, + ); + } + + fn update(&mut self, running_turn_count: usize, connection_count: usize) -> ShutdownAction { + if !self.requested { + return ShutdownAction::Noop; + } + + if self.forced || running_turn_count == 0 { + if self.forced { + info!( + "received second Ctrl-C; forcing restart with {running_turn_count} running assistant turn(s) and {connection_count} connection(s)" + ); + } else { + info!( + "Ctrl-C restart: no assistant turns running; stopping acceptor and disconnecting {connection_count} connection(s)" + ); + } + return ShutdownAction::Finish; + } + + if self.last_logged_running_turn_count != Some(running_turn_count) { + info!( + "Ctrl-C restart: waiting for {running_turn_count} running assistant turn(s) to finish" + ); + self.last_logged_running_turn_count = Some(running_turn_count); + } + + ShutdownAction::Noop + } } fn config_warning_from_error( @@ -253,19 +319,37 @@ pub async fn run_main_with_transport( let (outbound_control_tx, mut outbound_control_rx) = mpsc::channel::(CHANNEL_CAPACITY); + enum TransportRuntime { + Stdio, + WebSocket { + accept_handle: JoinHandle<()>, + shutdown_token: CancellationToken, + }, + } + let mut stdio_handles = Vec::>::new(); - let mut websocket_accept_handle = None; - match transport { + let transport_runtime = match transport { AppServerTransport::Stdio => { start_stdio_connection(transport_event_tx.clone(), &mut stdio_handles).await?; + TransportRuntime::Stdio } AppServerTransport::WebSocket { bind_address } => { - websocket_accept_handle = - Some(start_websocket_acceptor(bind_address, transport_event_tx.clone()).await?); + let shutdown_token = CancellationToken::new(); + let accept_handle = start_websocket_acceptor( + bind_address, + transport_event_tx.clone(), + shutdown_token.clone(), + ) + .await?; + TransportRuntime::WebSocket { + accept_handle, + shutdown_token, + } } - } - let single_client_mode = matches!(transport, AppServerTransport::Stdio); + }; + let single_client_mode = matches!(&transport_runtime, TransportRuntime::Stdio); let shutdown_when_no_connections = single_client_mode; + let graceful_ctrl_c_restart_enabled = !single_client_mode; // Parse CLI overrides once and derive the base Config eagerly so later // components do not need to work with raw TOML values. @@ -434,6 +518,16 @@ pub async fn run_main_with_transport( OutboundControlEvent::Closed { connection_id } => { outbound_connections.remove(&connection_id); } + OutboundControlEvent::DisconnectAll => { + info!( + "disconnecting {} outbound websocket connection(s) for graceful restart", + outbound_connections.len() + ); + for connection_state in outbound_connections.values() { + connection_state.request_disconnect(); + } + outbound_connections.clear(); + } } } envelope = outgoing_rx.recv() => { @@ -464,11 +558,46 @@ pub async fn run_main_with_transport( config_warnings, }); let mut thread_created_rx = processor.thread_created_receiver(); + let mut running_turn_count_rx = processor.subscribe_running_assistant_turn_count(); let mut connections = HashMap::::new(); + let websocket_accept_shutdown = match &transport_runtime { + TransportRuntime::WebSocket { shutdown_token, .. } => Some(shutdown_token.clone()), + TransportRuntime::Stdio => None, + }; async move { let mut listen_for_threads = true; + let mut shutdown_state = ShutdownState::default(); loop { + let running_turn_count = { + let running_turn_count = running_turn_count_rx.borrow(); + *running_turn_count + }; + if matches!( + shutdown_state.update(running_turn_count, connections.len()), + ShutdownAction::Finish + ) { + if let Some(shutdown_token) = &websocket_accept_shutdown { + shutdown_token.cancel(); + } + let _ = outbound_control_tx + .send(OutboundControlEvent::DisconnectAll) + .await; + break; + } + tokio::select! { + ctrl_c_result = tokio::signal::ctrl_c(), if graceful_ctrl_c_restart_enabled && !shutdown_state.forced() => { + if let Err(err) = ctrl_c_result { + warn!("failed to listen for Ctrl-C during graceful restart drain: {err}"); + } + let running_turn_count = *running_turn_count_rx.borrow(); + shutdown_state.on_ctrl_c(connections.len(), running_turn_count); + } + changed = running_turn_count_rx.changed(), if graceful_ctrl_c_restart_enabled && shutdown_state.requested() => { + if changed.is_err() { + warn!("running-turn watcher closed during graceful restart drain"); + } + } event = transport_event_rx.recv() => { let Some(event) = event else { break; @@ -619,8 +748,13 @@ pub async fn run_main_with_transport( let _ = processor_handle.await; let _ = outbound_handle.await; - if let Some(handle) = websocket_accept_handle { - handle.abort(); + if let TransportRuntime::WebSocket { + accept_handle, + shutdown_token, + } = transport_runtime + { + shutdown_token.cancel(); + let _ = accept_handle.await; } for handle in stdio_handles { diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 18d6a121e7..a4bb6a7389 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -51,6 +51,7 @@ use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; use futures::FutureExt; use tokio::sync::broadcast; +use tokio::sync::watch; use tokio::time::Duration; use tokio::time::timeout; use toml::Value as TomlValue; @@ -438,6 +439,11 @@ impl MessageProcessor { .await; } + pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver { + self.codex_message_processor + .subscribe_running_assistant_turn_count() + } + /// Handle a standalone JSON-RPC response originating from the peer. pub(crate) async fn process_response(&mut self, response: JSONRPCResponse) { tracing::info!("<- response: {:?}", response); diff --git a/codex-rs/app-server/src/thread_status.rs b/codex-rs/app-server/src/thread_status.rs index 5693631111..2eb1d8d036 100644 --- a/codex-rs/app-server/src/thread_status.rs +++ b/codex-rs/app-server/src/thread_status.rs @@ -15,11 +15,13 @@ use std::sync::Arc; use tokio::sync::Mutex; #[cfg(test)] use tokio::sync::mpsc; +use tokio::sync::watch; #[derive(Clone)] pub(crate) struct ThreadWatchManager { state: Arc>, outgoing: Option>, + running_turn_count_tx: watch::Sender, } pub(crate) struct ThreadWatchActiveGuard { @@ -71,16 +73,20 @@ impl Default for ThreadWatchManager { impl ThreadWatchManager { pub(crate) fn new() -> Self { + let (running_turn_count_tx, _running_turn_count_rx) = watch::channel(0); Self { state: Arc::new(Mutex::new(ThreadWatchState::default())), outgoing: None, + running_turn_count_tx, } } pub(crate) fn new_with_outgoing(outgoing: Arc) -> Self { + let (running_turn_count_tx, _running_turn_count_rx) = watch::channel(0); Self { state: Arc::new(Mutex::new(ThreadWatchState::default())), outgoing: Some(outgoing), + running_turn_count_tx, } } @@ -113,6 +119,21 @@ impl ThreadWatchManager { .collect() } + #[cfg(test)] + pub(crate) async fn running_turn_count(&self) -> usize { + self.state + .lock() + .await + .runtime_by_thread_id + .values() + .filter(|runtime| runtime.running) + .count() + } + + pub(crate) fn subscribe_running_turn_count(&self) -> watch::Receiver { + self.running_turn_count_tx.subscribe() + } + pub(crate) async fn note_turn_started(&self, thread_id: &str) { self.update_runtime_for_thread(thread_id, |runtime| { runtime.is_loaded = true; @@ -193,10 +214,17 @@ impl ThreadWatchManager { where F: FnOnce(&mut ThreadWatchState) -> Option, { - let notification = { + let (notification, running_turn_count) = { let mut state = self.state.lock().await; - mutate(&mut state) + let notification = mutate(&mut state); + let running_turn_count = state + .runtime_by_thread_id + .values() + .filter(|runtime| runtime.running) + .count(); + (notification, running_turn_count) }; + let _ = self.running_turn_count_tx.send(running_turn_count); if let Some(notification) = notification && let Some(outgoing) = &self.outgoing @@ -588,6 +616,32 @@ mod tests { ); } + #[tokio::test] + async fn has_running_turns_tracks_runtime_running_flag_only() { + let manager = ThreadWatchManager::new(); + manager + .upsert_thread(test_thread( + INTERACTIVE_THREAD_ID, + codex_app_server_protocol::SessionSource::Cli, + )) + .await; + + assert_eq!(manager.running_turn_count().await, 0); + + let _permission_guard = manager + .note_permission_requested(INTERACTIVE_THREAD_ID) + .await; + assert_eq!(manager.running_turn_count().await, 0); + + manager.note_turn_started(INTERACTIVE_THREAD_ID).await; + assert_eq!(manager.running_turn_count().await, 1); + + manager + .note_turn_completed(INTERACTIVE_THREAD_ID, false) + .await; + assert_eq!(manager.running_turn_count().await, 0); + } + #[tokio::test] async fn status_change_emits_notification() { let (outgoing_tx, mut outgoing_rx) = mpsc::channel(8); diff --git a/codex-rs/app-server/src/transport.rs b/codex-rs/app-server/src/transport.rs index b074a69578..027522153e 100644 --- a/codex-rs/app-server/src/transport.rs +++ b/codex-rs/app-server/src/transport.rs @@ -68,12 +68,6 @@ fn print_websocket_startup_banner(addr: SocketAddr) { } } -#[allow(clippy::print_stderr)] -fn print_websocket_connection(peer_addr: SocketAddr) { - let connected_label = colorize("websocket client connected from", Style::new().dimmed()); - eprintln!("{connected_label} {peer_addr}"); -} - #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum AppServerTransport { Stdio, @@ -193,7 +187,7 @@ impl OutboundConnectionState { self.disconnect_sender.is_some() } - fn request_disconnect(&self) { + pub(crate) fn request_disconnect(&self) { if let Some(disconnect_sender) = &self.disconnect_sender { disconnect_sender.cancel(); } @@ -271,6 +265,7 @@ pub(crate) async fn start_stdio_connection( pub(crate) async fn start_websocket_acceptor( bind_address: SocketAddr, transport_event_tx: mpsc::Sender, + shutdown_token: CancellationToken, ) -> IoResult> { let listener = TcpListener::bind(bind_address).await?; let local_addr = listener.local_addr()?; @@ -280,23 +275,31 @@ pub(crate) async fn start_websocket_acceptor( let connection_counter = Arc::new(AtomicU64::new(1)); Ok(tokio::spawn(async move { loop { - match listener.accept().await { - Ok((stream, peer_addr)) => { - print_websocket_connection(peer_addr); - let connection_id = - ConnectionId(connection_counter.fetch_add(1, Ordering::Relaxed)); - let transport_event_tx_for_connection = transport_event_tx.clone(); - tokio::spawn(async move { - run_websocket_connection( - connection_id, - stream, - transport_event_tx_for_connection, - ) - .await; - }); + tokio::select! { + _ = shutdown_token.cancelled() => { + info!("websocket acceptor shutting down"); + break; } - Err(err) => { - error!("failed to accept websocket connection: {err}"); + accept_result = listener.accept() => { + match accept_result { + Ok((stream, peer_addr)) => { + info!(%peer_addr, "websocket client connected"); + let connection_id = + ConnectionId(connection_counter.fetch_add(1, Ordering::Relaxed)); + let transport_event_tx_for_connection = transport_event_tx.clone(); + tokio::spawn(async move { + run_websocket_connection( + connection_id, + stream, + transport_event_tx_for_connection, + ) + .await; + }); + } + Err(err) => { + error!("failed to accept websocket connection: {err}"); + } + } } } } diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index ddd4326fc9..0bbbc70b55 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs @@ -28,9 +28,9 @@ use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message as WebSocketMessage; -const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(5); +pub(super) const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(5); -type WsClient = WebSocketStream>; +pub(super) type WsClient = WebSocketStream>; #[tokio::test] async fn websocket_transport_routes_per_connection_handshake_and_responses() -> Result<()> { @@ -78,7 +78,10 @@ async fn websocket_transport_routes_per_connection_handshake_and_responses() -> Ok(()) } -async fn spawn_websocket_server(codex_home: &Path, bind_addr: SocketAddr) -> Result { +pub(super) async fn spawn_websocket_server( + codex_home: &Path, + bind_addr: SocketAddr, +) -> Result { let program = codex_utils_cargo_bin::cargo_bin("codex-app-server") .context("should find app-server binary")?; let mut cmd = Command::new(program); @@ -106,14 +109,14 @@ async fn spawn_websocket_server(codex_home: &Path, bind_addr: SocketAddr) -> Res Ok(process) } -fn reserve_local_addr() -> Result { +pub(super) fn reserve_local_addr() -> Result { let listener = std::net::TcpListener::bind("127.0.0.1:0")?; let addr = listener.local_addr()?; drop(listener); Ok(addr) } -async fn connect_websocket(bind_addr: SocketAddr) -> Result { +pub(super) async fn connect_websocket(bind_addr: SocketAddr) -> Result { let url = format!("ws://{bind_addr}"); let deadline = Instant::now() + Duration::from_secs(10); loop { @@ -129,7 +132,11 @@ async fn connect_websocket(bind_addr: SocketAddr) -> Result { } } -async fn send_initialize_request(stream: &mut WsClient, id: i64, client_name: &str) -> Result<()> { +pub(super) async fn send_initialize_request( + stream: &mut WsClient, + id: i64, + client_name: &str, +) -> Result<()> { let params = InitializeParams { client_info: ClientInfo { name: client_name.to_string(), @@ -157,7 +164,7 @@ async fn send_config_read_request(stream: &mut WsClient, id: i64) -> Result<()> .await } -async fn send_request( +pub(super) async fn send_request( stream: &mut WsClient, method: &str, id: i64, @@ -179,7 +186,10 @@ async fn send_jsonrpc(stream: &mut WsClient, message: JSONRPCMessage) -> Result< .context("failed to send websocket frame") } -async fn read_response_for_id(stream: &mut WsClient, id: i64) -> Result { +pub(super) async fn read_response_for_id( + stream: &mut WsClient, + id: i64, +) -> Result { let target_id = RequestId::Integer(id); loop { let message = read_jsonrpc_message(stream).await?; @@ -235,7 +245,7 @@ async fn assert_no_message(stream: &mut WsClient, wait_for: Duration) -> Result< } } -fn create_config_toml( +pub(super) fn create_config_toml( codex_home: &Path, server_uri: &str, approval_policy: &str, diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs new file mode 100644 index 0000000000..c95f78f017 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs @@ -0,0 +1,237 @@ +use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT; +use super::connection_handling_websocket::WsClient; +use super::connection_handling_websocket::connect_websocket; +use super::connection_handling_websocket::create_config_toml; +use super::connection_handling_websocket::read_response_for_id; +use super::connection_handling_websocket::reserve_local_addr; +use super::connection_handling_websocket::send_initialize_request; +use super::connection_handling_websocket::send_request; +use super::connection_handling_websocket::spawn_websocket_server; +use anyhow::Context; +use anyhow::Result; +use anyhow::bail; +use app_test_support::create_final_assistant_message_sse_response; +use app_test_support::to_response; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::UserInput as V2UserInput; +use core_test_support::responses; +use futures::SinkExt; +use futures::StreamExt; +use std::process::Command as StdCommand; +use tempfile::TempDir; +use tokio::process::Child; +use tokio::time::Duration; +use tokio::time::Instant; +use tokio::time::sleep; +use tokio::time::timeout; +use tokio_tungstenite::tungstenite::Message as WebSocketMessage; +use wiremock::Mock; +use wiremock::matchers::method; +use wiremock::matchers::path_regex; + +#[tokio::test] +async fn websocket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Result<()> { + let GracefulCtrlCFixture { + _codex_home, + _server, + mut process, + mut ws, + } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; + + send_sigint(&process)?; + assert_process_does_not_exit_within(&mut process, Duration::from_millis(300)).await?; + + let status = wait_for_process_exit_within( + &mut process, + Duration::from_secs(10), + "timed out waiting for graceful Ctrl-C restart shutdown", + ) + .await?; + assert!(status.success(), "expected graceful exit, got {status}"); + + expect_websocket_disconnect(&mut ws).await?; + + Ok(()) +} + +#[tokio::test] +async fn websocket_transport_second_ctrl_c_forces_exit_while_turn_running() -> Result<()> { + let GracefulCtrlCFixture { + _codex_home, + _server, + mut process, + mut ws, + } = start_ctrl_c_restart_fixture(Duration::from_secs(3)).await?; + + send_sigint(&process)?; + assert_process_does_not_exit_within(&mut process, Duration::from_millis(300)).await?; + + send_sigint(&process)?; + let status = wait_for_process_exit_within( + &mut process, + Duration::from_secs(2), + "timed out waiting for forced Ctrl-C restart shutdown", + ) + .await?; + assert!(status.success(), "expected graceful exit, got {status}"); + + expect_websocket_disconnect(&mut ws).await?; + + Ok(()) +} + +struct GracefulCtrlCFixture { + _codex_home: TempDir, + _server: wiremock::MockServer, + process: Child, + ws: WsClient, +} + +async fn start_ctrl_c_restart_fixture(turn_delay: Duration) -> Result { + let server = responses::start_mock_server().await; + let delayed_turn_response = create_final_assistant_message_sse_response("Done")?; + Mock::given(method("POST")) + .and(path_regex(".*/responses$")) + .respond_with(responses::sse_response(delayed_turn_response).set_delay(turn_delay)) + .up_to_n_times(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + + let bind_addr = reserve_local_addr()?; + let process = spawn_websocket_server(codex_home.path(), bind_addr).await?; + let mut ws = connect_websocket(bind_addr).await?; + + send_initialize_request(&mut ws, 1, "ws_graceful_shutdown").await?; + let init_response = read_response_for_id(&mut ws, 1).await?; + assert_eq!(init_response.id, RequestId::Integer(1)); + + send_thread_start_request(&mut ws, 2).await?; + let thread_start_response = read_response_for_id(&mut ws, 2).await?; + let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; + + send_turn_start_request(&mut ws, 3, &thread.id).await?; + let turn_start_response = read_response_for_id(&mut ws, 3).await?; + assert_eq!(turn_start_response.id, RequestId::Integer(3)); + + wait_for_responses_post(&server, Duration::from_secs(5)).await?; + + Ok(GracefulCtrlCFixture { + _codex_home: codex_home, + _server: server, + process, + ws, + }) +} + +async fn send_thread_start_request(stream: &mut WsClient, id: i64) -> Result<()> { + send_request( + stream, + "thread/start", + id, + Some(serde_json::to_value(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + })?), + ) + .await +} + +async fn send_turn_start_request(stream: &mut WsClient, id: i64, thread_id: &str) -> Result<()> { + send_request( + stream, + "turn/start", + id, + Some(serde_json::to_value(TurnStartParams { + thread_id: thread_id.to_string(), + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + })?), + ) + .await +} + +async fn wait_for_responses_post(server: &wiremock::MockServer, wait_for: Duration) -> Result<()> { + let deadline = Instant::now() + wait_for; + loop { + let requests = server + .received_requests() + .await + .context("failed to read mock server requests")?; + if requests + .iter() + .any(|request| request.method == "POST" && request.url.path().ends_with("/responses")) + { + return Ok(()); + } + if Instant::now() >= deadline { + bail!("timed out waiting for /responses request"); + } + sleep(Duration::from_millis(10)).await; + } +} + +fn send_sigint(process: &Child) -> Result<()> { + let pid = process + .id() + .context("websocket app-server process has no pid")?; + let status = StdCommand::new("kill") + .arg("-INT") + .arg(pid.to_string()) + .status() + .context("failed to invoke kill -INT")?; + if !status.success() { + bail!("kill -INT exited with {status}"); + } + Ok(()) +} + +async fn assert_process_does_not_exit_within(process: &mut Child, window: Duration) -> Result<()> { + match timeout(window, process.wait()).await { + Err(_) => Ok(()), + Ok(Ok(status)) => bail!("process exited too early during graceful drain: {status}"), + Ok(Err(err)) => Err(err).context("failed waiting for process"), + } +} + +async fn wait_for_process_exit_within( + process: &mut Child, + window: Duration, + timeout_context: &'static str, +) -> Result { + timeout(window, process.wait()) + .await + .context(timeout_context)? + .context("failed waiting for websocket app-server process exit") +} + +async fn expect_websocket_disconnect(stream: &mut WsClient) -> Result<()> { + loop { + let frame = timeout(DEFAULT_READ_TIMEOUT, stream.next()) + .await + .context("timed out waiting for websocket disconnect")?; + match frame { + None => return Ok(()), + Some(Ok(WebSocketMessage::Close(_))) => return Ok(()), + Some(Ok(WebSocketMessage::Ping(payload))) => { + stream + .send(WebSocketMessage::Pong(payload)) + .await + .context("failed to reply to ping while waiting for disconnect")?; + } + Some(Ok(WebSocketMessage::Pong(_))) => {} + Some(Ok(WebSocketMessage::Frame(_))) => {} + Some(Ok(WebSocketMessage::Text(_))) => {} + Some(Ok(WebSocketMessage::Binary(_))) => {} + Some(Err(_)) => return Ok(()), + } + } +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 0e428536ca..67e7ae3ea6 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -5,6 +5,8 @@ mod collaboration_mode_list; mod compaction; mod config_rpc; mod connection_handling_websocket; +#[cfg(unix)] +mod connection_handling_websocket_unix; mod dynamic_tools; mod experimental_api; mod experimental_feature_list; From 63c2ac96cdc2e2b61bb23813a5cc0276eada0af5 Mon Sep 17 00:00:00 2001 From: Curtis 'Fjord' Hawthorne Date: Tue, 24 Feb 2026 17:12:02 -0800 Subject: [PATCH 2/4] fix(js_repl): surface uncaught kernel errors and reset cleanly (#12636) ## Summary Improve `js_repl` behavior when the Node kernel hits a process-level failure (for example, an uncaught exception or unhandled Promise rejection). Instead of only surfacing a generic `js_repl kernel exited unexpectedly` after stdout EOF, `js_repl` now returns a clearer exec error for the active request, then resets the kernel cleanly. ## Why Some sandbox-denied operations can trigger Node errors that become process-level failures (for example, an unhandled EventEmitter `'error'` event). In that case: - the kernel process exits, - the host sees stdout EOF, - the user gets a generic kernel-exit error, - and the next request can briefly race with stale kernel state. This change improves that failure mode without monkeypatching Node APIs. ## Changes ### Kernel-side (`js_repl` Node process) - Add process-level handlers for: - `uncaughtException` - `unhandledRejection` - When one of these fires: - best-effort emit a normal `exec_result` error for the active exec - include actionable guidance to catch/handle async errors (including Promise rejections and EventEmitter `'error'` events) - exit intentionally so the host can reset/restart the kernel ### Host-side (`JsReplManager`) - Clear dead kernel state as soon as the stdout reader observes unexpected kernel exit/EOF. - This lets the next `js_repl` exec start a fresh kernel instead of hitting a stale broken-pipe path. ### Tests - Add regression coverage for: - uncaught async exception -> exec error + kernel recovery on next exec - Update forced-kernel-exit test to validate recovery behavior (next exec restarts cleanly) ## Impact - Better user-facing error for kernel crashes caused by uncaught/unhandled async failures. - Cleaner recovery behavior after kernel exit. ## Validation - `cargo test -p codex-core --lib tools::js_repl::tests::js_repl_uncaught_exception_returns_exec_error_and_recovers -- --exact` - `cargo test -p codex-core --lib tools::js_repl::tests::js_repl_forced_kernel_exit_recovers_on_next_exec -- --exact` - `just fmt` --- codex-rs/core/src/tools/js_repl/kernel.js | 63 ++++++++++ codex-rs/core/src/tools/js_repl/mod.rs | 141 ++++++++++++++++++++-- 2 files changed, 194 insertions(+), 10 deletions(-) diff --git a/codex-rs/core/src/tools/js_repl/kernel.js b/codex-rs/core/src/tools/js_repl/kernel.js index 283a9f4516..f5459fde79 100644 --- a/codex-rs/core/src/tools/js_repl/kernel.js +++ b/codex-rs/core/src/tools/js_repl/kernel.js @@ -84,6 +84,8 @@ let previousModule = null; /** @type {Binding[]} */ let previousBindings = []; let cellCounter = 0; +let activeExecId = null; +let fatalExitScheduled = false; const builtinModuleSet = new Set([ ...builtinModules, @@ -395,6 +397,54 @@ function send(message) { process.stdout.write("\n"); } +function formatErrorMessage(error) { + if (error && typeof error === "object" && "message" in error) { + return error.message ? String(error.message) : String(error); + } + return String(error); +} + +function sendFatalExecResultSync(kind, error) { + if (!activeExecId) { + return; + } + const payload = { + type: "exec_result", + id: activeExecId, + ok: false, + output: "", + error: `js_repl kernel ${kind}: ${formatErrorMessage(error)}; kernel reset. Catch or handle async errors (including Promise rejections and EventEmitter 'error' events) to avoid kernel termination.`, + }; + try { + fs.writeSync(process.stdout.fd, `${JSON.stringify(payload)}\n`); + } catch { + // Best effort only; the host will still surface stdout EOF diagnostics. + } +} + +function scheduleFatalExit(kind, error) { + if (fatalExitScheduled) { + process.exitCode = 1; + return; + } + fatalExitScheduled = true; + sendFatalExecResultSync(kind, error); + + try { + fs.writeSync( + process.stderr.fd, + `js_repl kernel ${kind}: ${formatErrorMessage(error)}\n`, + ); + } catch { + // ignore + } + + // The host will observe stdout EOF, reset kernel state, and restart on demand. + setImmediate(() => { + process.exit(1); + }); +} + function formatLog(args) { return args .map((arg) => @@ -431,6 +481,7 @@ function withCapturedConsole(ctx, fn) { } async function handleExec(message) { + activeExecId = message.id; const tool = (toolName, args) => { if (typeof toolName !== "string" || !toolName) { return Promise.reject(new Error("codex.tool expects a tool name string")); @@ -527,6 +578,10 @@ async function handleExec(message) { output: "", error: error && error.message ? error.message : String(error), }); + } finally { + if (activeExecId === message.id) { + activeExecId = null; + } } } @@ -540,6 +595,14 @@ function handleToolResult(message) { let queue = Promise.resolve(); +process.on("uncaughtException", (error) => { + scheduleFatalExit("uncaught exception", error); +}); + +process.on("unhandledRejection", (reason) => { + scheduleFatalExit("unhandled rejection", reason); +}); + const input = createInterface({ input: process.stdin, crlfDelay: Infinity }); input.on("line", (line) => { if (!line.trim()) { diff --git a/codex-rs/core/src/tools/js_repl/mod.rs b/codex-rs/core/src/tools/js_repl/mod.rs index 6e6bb62281..a5bd81049f 100644 --- a/codex-rs/core/src/tools/js_repl/mod.rs +++ b/codex-rs/core/src/tools/js_repl/mod.rs @@ -269,7 +269,7 @@ pub struct JsReplManager { node_path: Option, node_module_dirs: Vec, tmp_dir: tempfile::TempDir, - kernel: Mutex>, + kernel: Arc>>, exec_lock: Arc, exec_tool_calls: Arc>>, } @@ -287,7 +287,7 @@ impl JsReplManager { node_path, node_module_dirs, tmp_dir, - kernel: Mutex::new(None), + kernel: Arc::new(Mutex::new(None)), exec_lock: Arc::new(tokio::sync::Semaphore::new(1)), exec_tool_calls: Arc::new(Mutex::new(HashMap::new())), }); @@ -669,6 +669,7 @@ impl JsReplManager { tokio::spawn(Self::read_stdout( stdout, Arc::clone(&child), + Arc::clone(&self.kernel), Arc::clone(&recent_stderr), Arc::clone(&pending_execs), Arc::clone(&exec_contexts), @@ -811,6 +812,7 @@ impl JsReplManager { async fn read_stdout( stdout: tokio::process::ChildStdout, child: Arc>, + manager_kernel: Arc>>, recent_stderr: Arc>>, pending_execs: Arc>>>, exec_contexts: Arc>>, @@ -964,6 +966,16 @@ impl JsReplManager { .clone() .unwrap_or_else(|| "js_repl kernel exited unexpectedly".to_string()); + { + let mut kernel = manager_kernel.lock().await; + let should_clear = kernel + .as_ref() + .is_some_and(|state| Arc::ptr_eq(&state.child, &child)); + if should_clear { + kernel.take(); + } + } + let mut pending = pending_execs.lock().await; let pending_exec_ids = pending.keys().cloned().collect::>(); for (_id, tx) in pending.drain() { @@ -1743,7 +1755,7 @@ mod tests { } #[tokio::test] - async fn js_repl_kernel_failure_includes_model_diagnostics() -> anyhow::Result<()> { + async fn js_repl_forced_kernel_exit_recovers_on_next_exec() -> anyhow::Result<()> { if !can_run_js_repl_runtime_tests().await { return Ok(()); } @@ -1772,9 +1784,24 @@ mod tests { Arc::clone(&state.child) }; JsReplManager::kill_kernel_child(&child, "test_crash").await; - tokio::time::sleep(Duration::from_millis(50)).await; + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let cleared = { + let guard = manager.kernel.lock().await; + guard + .as_ref() + .is_none_or(|state| !Arc::ptr_eq(&state.child, &child)) + }; + if cleared { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("host should clear dead kernel state promptly"); - let err = manager + let result = manager .execute( session, turn, @@ -1784,13 +1811,107 @@ mod tests { timeout_ms: Some(10_000), }, ) - .await - .expect_err("expected kernel failure after forced kill"); + .await?; + assert!(result.output.contains("after-kill")); + Ok(()) + } + + #[tokio::test] + async fn js_repl_uncaught_exception_returns_exec_error_and_recovers() -> anyhow::Result<()> { + if !can_run_js_repl_runtime_tests().await { + return Ok(()); + } + + let (session, turn) = crate::codex::make_session_and_context().await; + let session = Arc::new(session); + let turn = Arc::new(turn); + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::default())); + let manager = turn.js_repl.manager().await?; + + manager + .execute( + Arc::clone(&session), + Arc::clone(&turn), + Arc::clone(&tracker), + JsReplArgs { + code: "console.log('warmup');".to_string(), + timeout_ms: Some(10_000), + }, + ) + .await?; + + let child = { + let guard = manager.kernel.lock().await; + let state = guard.as_ref().expect("kernel should exist after warmup"); + Arc::clone(&state.child) + }; + + let err = tokio::time::timeout( + Duration::from_secs(3), + manager.execute( + Arc::clone(&session), + Arc::clone(&turn), + Arc::clone(&tracker), + JsReplArgs { + code: "setTimeout(() => { throw new Error('boom'); }, 0);\nawait new Promise(() => {});".to_string(), + timeout_ms: Some(10_000), + }, + ), + ) + .await + .expect("uncaught exception should fail promptly") + .expect_err("expected uncaught exception to fail the exec"); let message = err.to_string(); - assert!(message.contains("js_repl diagnostics:")); - assert!(message.contains("\"reason\":\"write_failed\"")); - assert!(message.contains("\"kernel_status\":\"exited(")); + assert!(message.contains("js_repl kernel uncaught exception: boom")); + assert!(message.contains("kernel reset.")); + assert!(message.contains("Catch or handle async errors")); + assert!(!message.contains("js_repl kernel exited unexpectedly")); + + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let exited = { + let mut child = child.lock().await; + child.try_wait()?.is_some() + }; + if exited { + return Ok::<(), anyhow::Error>(()); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("uncaught exception should terminate the previous kernel process")?; + + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let cleared = { + let guard = manager.kernel.lock().await; + guard + .as_ref() + .is_none_or(|state| !Arc::ptr_eq(&state.child, &child)) + }; + if cleared { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("host should clear dead kernel state promptly"); + + let next = manager + .execute( + session, + turn, + tracker, + JsReplArgs { + code: "console.log('after reset');".to_string(), + timeout_ms: Some(10_000), + }, + ) + .await?; + assert!(next.output.contains("after reset")); Ok(()) } From 448fb6ac22ba580802db3a734ba303b370d6daa1 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 24 Feb 2026 17:15:54 -0800 Subject: [PATCH 3/4] fix: clarify the value of SkillMetadata.path (#12729) Rename `SkillMetadata.path` to `SkillMetadata.path_to_skills_md` for clarity. Would ideally change the type to `AbsolutePathBuf`, but that can be done later. --- .../app-server/src/codex_message_processor.rs | 4 +- codex-rs/core/src/codex.rs | 4 +- codex-rs/core/src/mcp/skill_dependencies.rs | 2 +- codex-rs/core/src/mentions.rs | 2 +- codex-rs/core/src/skills/injection.rs | 32 +++++----- codex-rs/core/src/skills/invocation_utils.rs | 10 ++-- codex-rs/core/src/skills/loader.rs | 58 +++++++++---------- codex-rs/core/src/skills/model.rs | 5 +- codex-rs/core/src/skills/render.rs | 2 +- codex-rs/tui/src/bottom_pane/chat_composer.rs | 2 +- codex-rs/tui/src/bottom_pane/mod.rs | 2 +- codex-rs/tui/src/chatwidget.rs | 12 ++-- codex-rs/tui/src/chatwidget/skills.rs | 14 ++--- codex-rs/tui/src/chatwidget/tests.rs | 4 +- 14 files changed, 80 insertions(+), 73 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index b2f9f528bd..276b4ff4ad 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -6526,7 +6526,7 @@ fn skills_to_info( skills .iter() .map(|skill| { - let enabled = !disabled_paths.contains(&skill.path); + let enabled = !disabled_paths.contains(&skill.path_to_skills_md); codex_app_server_protocol::SkillMetadata { name: skill.name.clone(), description: skill.description.clone(), @@ -6557,7 +6557,7 @@ fn skills_to_info( .collect(), } }), - path: skill.path.clone(), + path: skill.path_to_skills_md.clone(), scope: skill.scope.into(), enabled, } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 397318c8e8..d3d789c4ed 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -4734,9 +4734,9 @@ fn skills_to_info( .collect(), } }), - path: skill.path.clone(), + path: skill.path_to_skills_md.clone(), scope: skill.scope, - enabled: !disabled_paths.contains(&skill.path), + enabled: !disabled_paths.contains(&skill.path_to_skills_md), }) .collect() } diff --git a/codex-rs/core/src/mcp/skill_dependencies.rs b/codex-rs/core/src/mcp/skill_dependencies.rs index f2138cc293..7ea040aede 100644 --- a/codex-rs/core/src/mcp/skill_dependencies.rs +++ b/codex-rs/core/src/mcp/skill_dependencies.rs @@ -434,7 +434,7 @@ mod tests { dependencies: Some(SkillDependencies { tools }), policy: None, permissions: None, - path: PathBuf::from("skill"), + path_to_skills_md: PathBuf::from("skill"), scope: SkillScope::User, } } diff --git a/codex-rs/core/src/mentions.rs b/codex-rs/core/src/mentions.rs index edbcf01610..008456d61a 100644 --- a/codex-rs/core/src/mentions.rs +++ b/codex-rs/core/src/mentions.rs @@ -55,7 +55,7 @@ pub(crate) fn build_skill_name_counts( let mut exact_counts: HashMap = HashMap::new(); let mut lower_counts: HashMap = HashMap::new(); for skill in skills { - if disabled_paths.contains(&skill.path) { + if disabled_paths.contains(&skill.path_to_skills_md) { continue; } *exact_counts.entry(skill.name.clone()).or_insert(0) += 1; diff --git a/codex-rs/core/src/skills/injection.rs b/codex-rs/core/src/skills/injection.rs index 475ea6c127..22767d6dc1 100644 --- a/codex-rs/core/src/skills/injection.rs +++ b/codex-rs/core/src/skills/injection.rs @@ -37,18 +37,18 @@ pub(crate) async fn build_skill_injections( let mut invocations = Vec::new(); for skill in mentioned_skills { - match fs::read_to_string(&skill.path).await { + match fs::read_to_string(&skill.path_to_skills_md).await { Ok(contents) => { emit_skill_injected_metric(otel, skill, "ok"); invocations.push(SkillInvocation { skill_name: skill.name.clone(), skill_scope: skill.scope, - skill_path: skill.path.clone(), + skill_path: skill.path_to_skills_md.clone(), invocation_type: InvocationType::Explicit, }); result.items.push(ResponseItem::from(SkillInstructions { name: skill.name.clone(), - path: skill.path.to_string_lossy().into_owned(), + path: skill.path_to_skills_md.to_string_lossy().into_owned(), contents, })); } @@ -57,7 +57,7 @@ pub(crate) async fn build_skill_injections( let message = format!( "Failed to load skill {name} at {path}: {err:#}", name = skill.name, - path = skill.path.display() + path = skill.path_to_skills_md.display() ); result.warnings.push(message); } @@ -121,9 +121,9 @@ pub(crate) fn collect_explicit_skill_mentions( if let Some(skill) = selection_context .skills .iter() - .find(|skill| skill.path.as_path() == path.as_path()) + .find(|skill| skill.path_to_skills_md.as_path() == path.as_path()) { - seen_paths.insert(skill.path.clone()); + seen_paths.insert(skill.path_to_skills_md.clone()); seen_names.insert(skill.name.clone()); selected.push(skill.clone()); } @@ -304,23 +304,27 @@ fn select_skills_from_mentions( .collect(); for skill in selection_context.skills { - if selection_context.disabled_paths.contains(&skill.path) - || seen_paths.contains(&skill.path) + if selection_context + .disabled_paths + .contains(&skill.path_to_skills_md) + || seen_paths.contains(&skill.path_to_skills_md) { continue; } - let path_str = skill.path.to_string_lossy(); + let path_str = skill.path_to_skills_md.to_string_lossy(); if mention_skill_paths.contains(path_str.as_ref()) { - seen_paths.insert(skill.path.clone()); + seen_paths.insert(skill.path_to_skills_md.clone()); seen_names.insert(skill.name.clone()); selected.push(skill.clone()); } } for skill in selection_context.skills { - if selection_context.disabled_paths.contains(&skill.path) - || seen_paths.contains(&skill.path) + if selection_context + .disabled_paths + .contains(&skill.path_to_skills_md) + || seen_paths.contains(&skill.path_to_skills_md) { continue; } @@ -347,7 +351,7 @@ fn select_skills_from_mentions( } if seen_names.insert(skill.name.clone()) { - seen_paths.insert(skill.path.clone()); + seen_paths.insert(skill.path_to_skills_md.clone()); selected.push(skill.clone()); } } @@ -479,7 +483,7 @@ mod tests { dependencies: None, policy: None, permissions: None, - path: PathBuf::from(path), + path_to_skills_md: PathBuf::from(path), scope: codex_protocol::protocol::SkillScope::User, } } diff --git a/codex-rs/core/src/skills/invocation_utils.rs b/codex-rs/core/src/skills/invocation_utils.rs index fd847091de..65f077f96a 100644 --- a/codex-rs/core/src/skills/invocation_utils.rs +++ b/codex-rs/core/src/skills/invocation_utils.rs @@ -32,10 +32,10 @@ pub(crate) fn build_implicit_skill_path_indexes( let mut by_scripts_dir = HashMap::new(); let mut by_skill_doc_path = HashMap::new(); for skill in skills { - let skill_doc_path = normalize_path(skill.path.as_path()); + let skill_doc_path = normalize_path(skill.path_to_skills_md.as_path()); by_skill_doc_path.insert(skill_doc_path, skill.clone()); - if let Some(skill_dir) = skill.path.parent() { + if let Some(skill_dir) = skill.path_to_skills_md.parent() { let scripts_dir = normalize_path(&skill_dir.join("scripts")); by_scripts_dir.insert(scripts_dir, skill); } @@ -118,7 +118,7 @@ pub(crate) async fn ensure_skill_approval_for_command( let cache_key = SkillApprovalCacheKey { skill_name: skill.name.clone(), - skill_path: skill.path.clone(), + skill_path: skill.path_to_skills_md.clone(), skill_scope: skill.scope, }; let already_approved = { @@ -162,7 +162,7 @@ pub(crate) async fn maybe_emit_implicit_skill_invocation( let invocation = SkillInvocation { skill_name: candidate.name, skill_scope: candidate.scope, - skill_path: candidate.path, + skill_path: candidate.path_to_skills_md, invocation_type: InvocationType::Implicit, }; let skill_scope = match invocation.skill_scope { @@ -337,7 +337,7 @@ mod tests { dependencies: None, policy: None, permissions: None, - path: skill_doc_path, + path_to_skills_md: skill_doc_path, scope: codex_protocol::protocol::SkillScope::User, } } diff --git a/codex-rs/core/src/skills/loader.rs b/codex-rs/core/src/skills/loader.rs index 6508d4fd10..960830649b 100644 --- a/codex-rs/core/src/skills/loader.rs +++ b/codex-rs/core/src/skills/loader.rs @@ -167,7 +167,7 @@ where let mut seen: HashSet = HashSet::new(); outcome .skills - .retain(|skill| seen.insert(skill.path.clone())); + .retain(|skill| seen.insert(skill.path_to_skills_md.clone())); fn scope_rank(scope: SkillScope) -> u8 { // Higher-priority scopes first (matches root scan order for dedupe). @@ -183,7 +183,7 @@ where scope_rank(a.scope) .cmp(&scope_rank(b.scope)) .then_with(|| a.name.cmp(&b.name)) - .then_with(|| a.path.cmp(&b.path)) + .then_with(|| a.path_to_skills_md.cmp(&b.path_to_skills_md)) }); outcome @@ -528,7 +528,7 @@ fn parse_skill_file(path: &Path, scope: SkillScope) -> Result, // This is an experimental field. pub permissions: Option, - pub path: PathBuf, + /// Path to the SKILLS.md file that declares this skill. + pub path_to_skills_md: PathBuf, pub scope: SkillScope, } @@ -76,7 +77,7 @@ pub struct SkillLoadOutcome { impl SkillLoadOutcome { pub fn is_skill_enabled(&self, skill: &SkillMetadata) -> bool { - !self.disabled_paths.contains(&skill.path) + !self.disabled_paths.contains(&skill.path_to_skills_md) } pub fn is_skill_allowed_for_implicit_invocation(&self, skill: &SkillMetadata) -> bool { diff --git a/codex-rs/core/src/skills/render.rs b/codex-rs/core/src/skills/render.rs index 9627778dce..e6243f1693 100644 --- a/codex-rs/core/src/skills/render.rs +++ b/codex-rs/core/src/skills/render.rs @@ -11,7 +11,7 @@ pub fn render_skills_section(skills: &[SkillMetadata]) -> Option { lines.push("### Available skills".to_string()); for skill in skills { - let path_str = skill.path.to_string_lossy().replace('\\', "/"); + let path_str = skill.path_to_skills_md.to_string_lossy().replace('\\', "/"); let name = skill.name.as_str(); let description = skill.description.as_str(); lines.push(format!("- {name}: {description} (file: {path_str})")); diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index 8d9e4b250d..7892119966 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -3559,7 +3559,7 @@ impl ChatComposer { description, insert_text: format!("${skill_name}"), search_terms, - path: Some(skill.path.to_string_lossy().into_owned()), + path: Some(skill.path_to_skills_md.to_string_lossy().into_owned()), category_tag: (skill.scope == codex_protocol::protocol::SkillScope::Repo) .then(|| "[Repo]".to_string()), }); diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index 6801292360..5781e055d2 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -1495,7 +1495,7 @@ mod tests { dependencies: None, policy: None, permissions: None, - path: PathBuf::from("test-skill"), + path_to_skills_md: PathBuf::from("test-skill"), scope: SkillScope::User, }]), }); diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index ac8da4d55d..15a6a58d62 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -4066,12 +4066,14 @@ impl ChatWidget { .strip_prefix("skill://") .unwrap_or(binding.path.as_str()); let path = Path::new(path); - if let Some(skill) = skills.iter().find(|skill| skill.path.as_path() == path) - && selected_skill_paths.insert(skill.path.clone()) + if let Some(skill) = skills + .iter() + .find(|skill| skill.path_to_skills_md.as_path() == path) + && selected_skill_paths.insert(skill.path_to_skills_md.clone()) { items.push(UserInput::Skill { name: skill.name.clone(), - path: skill.path.clone(), + path: skill.path_to_skills_md.clone(), }); } } @@ -4079,13 +4081,13 @@ impl ChatWidget { let skill_mentions = find_skill_mentions_with_tool_mentions(&mentions, skills); for skill in skill_mentions { if bound_names.contains(skill.name.as_str()) - || !selected_skill_paths.insert(skill.path.clone()) + || !selected_skill_paths.insert(skill.path_to_skills_md.clone()) { continue; } items.push(UserInput::Skill { name: skill.name.clone(), - path: skill.path.clone(), + path: skill.path_to_skills_md.clone(), }); } } diff --git a/codex-rs/tui/src/chatwidget/skills.rs b/codex-rs/tui/src/chatwidget/skills.rs index 1d71cc61ec..72af6b89b6 100644 --- a/codex-rs/tui/src/chatwidget/skills.rs +++ b/codex-rs/tui/src/chatwidget/skills.rs @@ -78,7 +78,7 @@ impl ChatWidget { let display_name = skill_display_name(&core_skill).to_string(); let description = skill_description(&core_skill).to_string(); let name = core_skill.name.clone(); - let path = core_skill.path; + let path = core_skill.path_to_skills_md; SkillsToggleItem { name: display_name, skill_name: name, @@ -191,7 +191,7 @@ fn protocol_skill_to_core(skill: &ProtocolSkillMetadata) -> SkillMetadata { }), policy: None, permissions: None, - path: skill.path.clone(), + path_to_skills_md: skill.path.clone(), scope: skill.scope, } } @@ -229,23 +229,23 @@ pub(crate) fn find_skill_mentions_with_tool_mentions( let mut matches: Vec = Vec::new(); for skill in skills { - if seen_paths.contains(&skill.path) { + if seen_paths.contains(&skill.path_to_skills_md) { continue; } - let path_str = skill.path.to_string_lossy(); + let path_str = skill.path_to_skills_md.to_string_lossy(); if mention_skill_paths.contains(path_str.as_ref()) { - seen_paths.insert(skill.path.clone()); + seen_paths.insert(skill.path_to_skills_md.clone()); seen_names.insert(skill.name.clone()); matches.push(skill.clone()); } } for skill in skills { - if seen_paths.contains(&skill.path) { + if seen_paths.contains(&skill.path_to_skills_md) { continue; } if mentions.names.contains(&skill.name) && seen_names.insert(skill.name.clone()) { - seen_paths.insert(skill.path.clone()); + seen_paths.insert(skill.path_to_skills_md.clone()); matches.push(skill.clone()); } } diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index 5fa9854138..1e938e2d4b 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -931,7 +931,7 @@ async fn submission_prefers_selected_duplicate_skill_path() { dependencies: None, policy: None, permissions: None, - path: repo_skill_path, + path_to_skills_md: repo_skill_path, scope: SkillScope::Repo, }, SkillMetadata { @@ -942,7 +942,7 @@ async fn submission_prefers_selected_duplicate_skill_path() { dependencies: None, policy: None, permissions: None, - path: user_skill_path.clone(), + path_to_skills_md: user_skill_path.clone(), scope: SkillScope::User, }, ])); From 2436e94386e166234397278b310309210f49333c Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 24 Feb 2026 17:16:55 -0800 Subject: [PATCH 4/4] feat: codex-shell-escalation should check skills for excalation policy --- .../tools/runtimes/shell/unix_escalation.rs | 149 +++++++++++++----- 1 file changed, 107 insertions(+), 42 deletions(-) diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 811b1fdfe6..5f02de7d98 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -8,6 +8,7 @@ use crate::exec::is_likely_sandbox_denied; use crate::features::Feature; use crate::sandboxing::SandboxPermissions; use crate::shell::ShellType; +use crate::skills::SkillMetadata; use crate::tools::runtimes::build_command_spec; use crate::tools::sandboxing::SandboxAttempt; use crate::tools::sandboxing::ToolCtx; @@ -164,11 +165,12 @@ impl CoreShellActionProvider { async fn prompt( &self, - command: &[String], + program: &Path, + argv: &[String], workdir: &Path, stopwatch: &Stopwatch, ) -> anyhow::Result { - let command = command.to_vec(); + let command = join_program_and_argv(program, argv); let workdir = workdir.to_path_buf(); let session = self.session.clone(); let turn = self.turn.clone(); @@ -192,49 +194,40 @@ impl CoreShellActionProvider { }) .await) } -} -#[async_trait::async_trait] -impl EscalationPolicy for CoreShellActionProvider { - async fn determine_action( + /// Because we should be intercepting execve(2) calls, `program` should be + /// an absolute path. The idea is that we check to see whether it matches + /// any skills. + async fn find_skill(&self, program: &Path) -> Option { + let force_reload = false; + let skills_outcome = self + .session + .services + .skills_manager + .skills_for_cwd(&self.turn.cwd, force_reload) + .await; + + for skill in skills_outcome.skills { + let Some(skill_root) = skill.path_to_skills_md.parent() else { + continue; + }; + if program.starts_with(skill_root.join("scripts")) { + return Some(skill); + } + } + + None + } + + async fn process_decision( &self, - file: &Path, + decision: Decision, + needs_escalation: bool, + program: &Path, argv: &[String], workdir: &Path, ) -> anyhow::Result { - let command = std::iter::once(file.to_string_lossy().to_string()) - .chain(argv.iter().cloned()) - .collect::>(); - let (commands, used_complex_parsing) = - if let Some(commands) = parse_shell_lc_plain_commands(&command) { - (commands, false) - } else if let Some(single_command) = parse_shell_lc_single_command_prefix(&command) { - (vec![single_command], true) - } else { - (vec![command.clone()], false) - }; - - let fallback = |cmd: &[String]| { - crate::exec_policy::render_decision_for_unmatched_command( - self.approval_policy, - &self.sandbox_policy, - cmd, - self.sandbox_permissions, - used_complex_parsing, - ) - }; - let evaluation = { - let policy = self.policy.read().await; - policy.check_multiple(commands.iter(), &fallback) - }; - // When true, means the Evaluation was due to *.rules, not the - // fallback function. - let decision_driven_by_policy = - Self::decision_driven_by_policy(&evaluation.matched_rules, evaluation.decision); - let needs_escalation = - self.sandbox_permissions.requires_escalated_permissions() || decision_driven_by_policy; - - Ok(match evaluation.decision { + let action = match decision { Decision::Forbidden => EscalateAction::Deny { reason: Some("Execution forbidden by policy".to_string()), }, @@ -248,7 +241,7 @@ impl EscalationPolicy for CoreShellActionProvider { reason: Some("Execution forbidden by policy".to_string()), } } else { - match self.prompt(&command, workdir, &self.stopwatch).await? { + match self.prompt(program, argv, workdir, &self.stopwatch).await? { ReviewDecision::Approved | ReviewDecision::ApprovedExecpolicyAmendment { .. } | ReviewDecision::ApprovedForSession => { @@ -282,7 +275,72 @@ impl EscalationPolicy for CoreShellActionProvider { } } Decision::Allow => EscalateAction::Escalate, - }) + }; + Ok(action) + } +} + +#[async_trait::async_trait] +impl EscalationPolicy for CoreShellActionProvider { + async fn determine_action( + &self, + program: &Path, + argv: &[String], + workdir: &Path, + ) -> anyhow::Result { + if let Some(skill) = self.find_skill(program).await { + // TODO(mbolin): Check to see whether the program has been approved + // for this session already, in which case we can return + // `EscalateAction::Escalate`, though we will also have to pass + // along the sandbox constraints, once those are supported. + tracing::info!("Matches program {program:?} to skill {skill:?}"); + + // For now, we always prompt for scripts that look like they belong + // to skills, which means we ignore exec policy rules for those + // scripts. + return self + .process_decision(Decision::Prompt, true, program, argv, workdir) + .await; + } + + let command = join_program_and_argv(program, argv); + let (commands, used_complex_parsing) = + if let Some(commands) = parse_shell_lc_plain_commands(&command) { + (commands, false) + } else if let Some(single_command) = parse_shell_lc_single_command_prefix(&command) { + (vec![single_command], true) + } else { + (vec![command.clone()], false) + }; + + let fallback = |cmd: &[String]| { + crate::exec_policy::render_decision_for_unmatched_command( + self.approval_policy, + &self.sandbox_policy, + cmd, + self.sandbox_permissions, + used_complex_parsing, + ) + }; + let evaluation = { + let policy = self.policy.read().await; + policy.check_multiple(commands.iter(), &fallback) + }; + // When true, means the Evaluation was due to *.rules, not the + // fallback function. + let decision_driven_by_policy = + Self::decision_driven_by_policy(&evaluation.matched_rules, evaluation.decision); + let needs_escalation = + self.sandbox_permissions.requires_escalated_permissions() || decision_driven_by_policy; + + self.process_decision( + evaluation.decision, + needs_escalation, + program, + argv, + workdir, + ) + .await } } @@ -421,6 +479,13 @@ fn map_exec_result( Ok(output) } +/// Some functions need the full command as a Vec. +fn join_program_and_argv(program: &Path, argv: &[String]) -> Vec { + std::iter::once(program.to_string_lossy().to_string()) + .chain(argv.iter().cloned()) + .collect::>() +} + #[cfg(test)] mod tests { use super::ParsedShellCommand;