mirror of
https://github.com/openai/codex.git
synced 2026-04-29 08:56:38 +00:00
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 <noreply@openai.com>
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user