Merge 2436e94386 into sapling-pr-archive-bolinfest

codex-rs/app-server/src/codex_message_processor.rs

@@ -270,6 +270,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;
@@ -2802,6 +2803,10 @@ impl CodexMessageProcessor {
             .await;
     }
 
+    pub(crate) fn subscribe_running_assistant_turn_count(&self) -> watch::Receiver<usize> {
+        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,

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<usize>,
+}
+
+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::<OutboundControlEvent>(CHANNEL_CAPACITY);
 
+    enum TransportRuntime {
+        Stdio,
+        WebSocket {
+            accept_handle: JoinHandle<()>,
+            shutdown_token: CancellationToken,
+        },
+    }
+
     let mut stdio_handles = Vec::<JoinHandle<()>>::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::<ConnectionId, ConnectionState>::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 {

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<usize> {
+        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);

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<Mutex<ThreadWatchState>>,
     outgoing: Option<Arc<OutgoingMessageSender>>,
+    running_turn_count_tx: watch::Sender<usize>,
 }
 
 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<OutgoingMessageSender>) -> 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<usize> {
+        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<ThreadStatusChangedNotification>,
     {
-        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);

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<TransportEvent>,
+    shutdown_token: CancellationToken,
 ) -> IoResult<JoinHandle<()>> {
     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}");
+                        }
+                    }
                 }
             }
         }

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<MaybeTlsStream<tokio::net::TcpStream>>;
+pub(super) type WsClient = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
 
 #[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<Child> {
+pub(super) async fn spawn_websocket_server(
+    codex_home: &Path,
+    bind_addr: SocketAddr,
+) -> Result<Child> {
     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<SocketAddr> {
+pub(super) fn reserve_local_addr() -> Result<SocketAddr> {
     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<WsClient> {
+pub(super) async fn connect_websocket(bind_addr: SocketAddr) -> Result<WsClient> {
     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<WsClient> {
     }
 }
 
-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<JSONRPCResponse> {
+pub(super) async fn read_response_for_id(
+    stream: &mut WsClient,
+    id: i64,
+) -> Result<JSONRPCResponse> {
     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,

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<GracefulCtrlCFixture> {
+    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<std::process::ExitStatus> {
+    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(()),
+        }
+    }
+}

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;

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()) {

codex-rs/core/src/tools/js_repl/mod.rs

@@ -269,7 +269,7 @@ pub struct JsReplManager {
     node_path: Option<PathBuf>,
     node_module_dirs: Vec<PathBuf>,
     tmp_dir: tempfile::TempDir,
-    kernel: Mutex<Option<KernelState>>,
+    kernel: Arc<Mutex<Option<KernelState>>>,
     exec_lock: Arc<tokio::sync::Semaphore>,
     exec_tool_calls: Arc<Mutex<HashMap<String, ExecToolCalls>>>,
 }
@@ -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<Mutex<Child>>,
+        manager_kernel: Arc<Mutex<Option<KernelState>>>,
         recent_stderr: Arc<Mutex<VecDeque<String>>>,
         pending_execs: Arc<Mutex<HashMap<String, tokio::sync::oneshot::Sender<ExecResultMessage>>>>,
         exec_contexts: Arc<Mutex<HashMap<String, ExecContext>>>,
@@ -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::<Vec<_>>();
         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(())
     }
 

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;
@@ -174,11 +175,12 @@ impl CoreShellActionProvider {
 
     async fn prompt(
         &self,
-        command: &[String],
+        program: &Path,
+        argv: &[String],
         workdir: &Path,
         stopwatch: &Stopwatch,
     ) -> anyhow::Result<ReviewDecision> {
-        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();
@@ -202,49 +204,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<SkillMetadata> {
+        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<EscalateAction> {
-        let command = std::iter::once(file.to_string_lossy().to_string())
-            .chain(argv.iter().cloned())
-            .collect::<Vec<_>>();
-        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()),
             },
@@ -258,7 +251,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 => {
@@ -292,7 +285,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<EscalateAction> {
+        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
     }
 }
 
@@ -403,6 +461,13 @@ fn map_exec_result(
     Ok(output)
 }
 
+/// Some functions need the full command as a Vec<String>.
+fn join_program_and_argv(program: &Path, argv: &[String]) -> Vec<String> {
+    std::iter::once(program.to_string_lossy().to_string())
+        .chain(argv.iter().cloned())
+        .collect::<Vec<_>>()
+}
+
 #[cfg(test)]
 mod tests {
     use super::ParsedShellCommand;