openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add exec-server readiness and health probes

Browse Source
This commit is contained in:
Richard Lee
2026-05-12 08:27:01 -07:00
parent c968b85fab
commit e29d90d704
9 changed files with 280 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -2693,6 +2693,7 @@ dependencies = [
"anyhow",
"arc-swap",
"async-trait",
"axum",
"base64 0.22.1",
"bytes",
"codex-app-server-protocol",
codex-rs/exec-server/Cargo.toml
+1
View File
@@ -13,6 +13,7 @@ workspace = true
[dependencies]
arc-swap = { workspace = true }
async-trait = { workspace = true }
axum = { workspace = true, features = ["http1", "tokio", "ws"] }
base64 = { workspace = true }
bytes = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-rs/exec-server/README.md
+7
View File
@@ -24,6 +24,13 @@ The CLI entrypoint supports:
- `ws://IP:PORT` (default)
- `--remote URL --executor-id ID [--name NAME]`
When running with `--listen ws://IP:PORT`, the same listener also serves basic
HTTP probes:
- `GET /readyz` returns `200 OK` while the listener is accepting new websocket
connections, then `503 Service Unavailable` once graceful shutdown begins.
- `GET /healthz` returns `200 OK` while the listener is still able to serve HTTP.
Remote mode registers the local exec-server with the executor registry,
then reconnects to the service-provided rendezvous websocket as the executor.
It requires a bearer token in `CODEX_EXEC_SERVER_REMOTE_BEARER_TOKEN`.
codex-rs/exec-server/src/connection.rs
+105 -55
View File
@@ -1,5 +1,9 @@
use axum::extract::ws::Message as AxumWebSocketMessage;
use axum::extract::ws::WebSocket as AxumWebSocket;
use codex_app_server_protocol::JSONRPCMessage;
use futures::Sink;
use futures::SinkExt;
use futures::Stream;
use futures::StreamExt;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
@@ -123,11 +127,30 @@ impl JsonRpcConnection {
pub(crate) fn from_websocket<S>(stream: WebSocketStream<S>, connection_label: String) -> Self
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
let (websocket_writer, websocket_reader) = stream.split();
Self::from_websocket_parts(websocket_writer, websocket_reader, connection_label)
}
pub(crate) fn from_axum_websocket(stream: AxumWebSocket, connection_label: String) -> Self {
let (websocket_writer, websocket_reader) = stream.split();
Self::from_websocket_parts(websocket_writer, websocket_reader, connection_label)
}
fn from_websocket_parts<W, R, M, E>(
mut websocket_writer: W,
mut websocket_reader: R,
connection_label: String,
) -> Self
where
W: Sink<M, Error = E> + Unpin + Send + 'static,
R: Stream<Item = Result<M, E>> + Unpin + Send + 'static,
M: JsonRpcWebSocketMessage,
E: std::fmt::Display + Send + 'static,
{
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (disconnected_tx, disconnected_rx) = watch::channel(false);
let (mut websocket_writer, mut websocket_reader) = stream.split();
let reader_label = connection_label.clone();
let incoming_tx_for_reader = incoming_tx.clone();
@@ -135,61 +158,36 @@ impl JsonRpcConnection {
let reader_task = tokio::spawn(async move {
loop {
match websocket_reader.next().await {
Some(Ok(Message::Text(text))) => {
match serde_json::from_str::<JSONRPCMessage>(text.as_ref()) {
Ok(message) => {
if incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Message(message))
.await
.is_err()
{
break;
}
}
Err(err) => {
send_malformed_message(
&incoming_tx_for_reader,
Some(format!(
"failed to parse websocket JSON-RPC message from {reader_label}: {err}"
)),
)
.await;
Some(Ok(message)) => match message.parse_jsonrpc_frame() {
Ok(JsonRpcWebSocketFrame::Message(message)) => {
if incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Message(message))
.await
.is_err()
{
break;
}
}
}
Some(Ok(Message::Binary(bytes))) => {
match serde_json::from_slice::<JSONRPCMessage>(bytes.as_ref()) {
Ok(message) => {
if incoming_tx_for_reader
.send(JsonRpcConnectionEvent::Message(message))
.await
.is_err()
{
break;
}
}
Err(err) => {
send_malformed_message(
&incoming_tx_for_reader,
Some(format!(
"failed to parse websocket JSON-RPC message from {reader_label}: {err}"
)),
)
.await;
}
Err(err) => {
send_malformed_message(
&incoming_tx_for_reader,
Some(format!(
"failed to parse websocket JSON-RPC message from {reader_label}: {err}"
)),
)
.await;
}
}
Some(Ok(Message::Close(_))) => {
send_disconnected(
&incoming_tx_for_reader,
&disconnected_tx_for_reader,
/*reason*/ None,
)
.await;
break;
}
Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {}
Some(Ok(_)) => {}
Ok(JsonRpcWebSocketFrame::Close) => {
send_disconnected(
&incoming_tx_for_reader,
&disconnected_tx_for_reader,
/*reason*/ None,
)
.await;
break;
}
Ok(JsonRpcWebSocketFrame::Ignore) => {}
},
Some(Err(err)) => {
send_disconnected(
&incoming_tx_for_reader,
@@ -218,8 +216,7 @@ impl JsonRpcConnection {
while let Some(message) = outgoing_rx.recv().await {
match serialize_jsonrpc_message(&message) {
Ok(encoded) => {
if let Err(err) = websocket_writer.send(Message::Text(encoded.into())).await
{
if let Err(err) = websocket_writer.send(M::from_text(encoded)).await {
send_disconnected(
&incoming_tx,
&disconnected_tx,
@@ -271,6 +268,59 @@ impl JsonRpcConnection {
}
}
enum JsonRpcWebSocketFrame {
Message(JSONRPCMessage),
Close,
Ignore,
}
trait JsonRpcWebSocketMessage: Send + 'static {
fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error>;
fn from_text(text: String) -> Self;
}
impl JsonRpcWebSocketMessage for Message {
fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
match self {
Message::Text(text) => {
serde_json::from_str(text.as_ref()).map(JsonRpcWebSocketFrame::Message)
}
Message::Binary(bytes) => {
serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message)
}
Message::Close(_) => Ok(JsonRpcWebSocketFrame::Close),
Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => {
Ok(JsonRpcWebSocketFrame::Ignore)
}
}
}
fn from_text(text: String) -> Self {
Self::Text(text.into())
}
}
impl JsonRpcWebSocketMessage for AxumWebSocketMessage {
fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
match self {
AxumWebSocketMessage::Text(text) => {
serde_json::from_str(text.as_ref()).map(JsonRpcWebSocketFrame::Message)
}
AxumWebSocketMessage::Binary(bytes) => {
serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message)
}
AxumWebSocketMessage::Close(_) => Ok(JsonRpcWebSocketFrame::Close),
AxumWebSocketMessage::Ping(_) | AxumWebSocketMessage::Pong(_) => {
Ok(JsonRpcWebSocketFrame::Ignore)
}
}
}
fn from_text(text: String) -> Self {
Self::Text(text.into())
}
}
async fn send_disconnected(
incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,
disconnected_tx: &watch::Sender<bool>,
codex-rs/exec-server/src/server/transport.rs
+67 -47
View File
@@ -1,3 +1,11 @@
use axum::Router;
use axum::extract::ConnectInfo;
use axum::extract::State;
use axum::extract::ws::WebSocketUpgrade;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::routing::any;
use axum::routing::get;
use std::io::Result as IoResult;
use std::io::Write as _;
use std::net::SocketAddr;
@@ -5,9 +13,7 @@ use tokio::io;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use tracing::info;
use tracing::warn;
@@ -132,59 +138,73 @@ async fn run_websocket_listener(
std::io::stdout().flush()?;
let shutdown_token = CancellationToken::new();
let connection_tasks = TaskTracker::new();
let shutdown_signal = shutdown_signal();
tokio::pin!(shutdown_signal);
loop {
let accepted = tokio::select! {
accepted = listener.accept() => accepted?,
shutdown_result = &mut shutdown_signal => {
if let Err(err) = shutdown_result {
warn!("failed to listen for exec-server shutdown signal: {err}");
}
info!("received SIGTERM; shutting down codex-exec-server");
break;
}
};
let (stream, peer_addr) = accepted;
let processor = processor.clone();
let connection_shutdown_token = shutdown_token.clone();
connection_tasks.spawn(async move {
let websocket = tokio::select! {
websocket = accept_async(stream) => websocket,
_ = connection_shutdown_token.cancelled() => {
return;
}
};
match websocket {
Ok(websocket) => {
processor
.run_connection(
JsonRpcConnection::from_websocket(
websocket,
format!("exec-server websocket {peer_addr}"),
),
connection_shutdown_token,
)
.await;
}
Err(err) => {
warn!(
"failed to accept exec-server websocket connection from {peer_addr}: {err}"
);
}
}
let router = Router::new()
.route("/", any(websocket_upgrade_handler))
.route("/readyz", get(readiness_handler))
.route("/healthz", get(health_check_handler))
.with_state(ExecServerWebSocketState {
processor: processor.clone(),
shutdown_token: shutdown_token.clone(),
});
}
let graceful_shutdown_token = shutdown_token.clone();
let server = axum::serve(
listener,
router.into_make_service_with_connect_info::<SocketAddr>(),
)
.with_graceful_shutdown(async move {
if let Err(err) = shutdown_signal().await {
warn!("failed to listen for exec-server shutdown signal: {err}");
}
info!("received SIGTERM; shutting down codex-exec-server");
graceful_shutdown_token.cancel();
});
let server_result = server.await;
shutdown_token.cancel();
connection_tasks.close();
connection_tasks.wait().await;
processor.shutdown().await;
info!("codex-exec-server shutdown complete");
server_result?;
Ok(())
}
#[derive(Clone)]
struct ExecServerWebSocketState {
processor: ConnectionProcessor,
shutdown_token: CancellationToken,
}
async fn readiness_handler(State(state): State<ExecServerWebSocketState>) -> StatusCode {
if state.shutdown_token.is_cancelled() {
StatusCode::SERVICE_UNAVAILABLE
} else {
StatusCode::OK
}
}
async fn health_check_handler() -> StatusCode {
StatusCode::OK
}
async fn websocket_upgrade_handler(
websocket: WebSocketUpgrade,
ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
State(state): State<ExecServerWebSocketState>,
) -> impl IntoResponse {
info!(%peer_addr, "exec-server websocket client connected");
websocket.on_upgrade(move |stream| async move {
state
.processor
.run_connection(
JsonRpcConnection::from_axum_websocket(
stream,
format!("exec-server websocket {peer_addr}"),
),
state.shutdown_token,
)
.await;
})
}
async fn shutdown_signal() -> IoResult<()> {
#[cfg(unix)]
{
codex-rs/exec-server/src/server/transport_tests.rs
+31
View File
@@ -1,6 +1,8 @@
use std::net::SocketAddr;
use std::time::Duration;
use axum::extract::State;
use axum::http::StatusCode;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
@@ -15,13 +17,17 @@ use tokio::time::timeout;
use super::DEFAULT_LISTEN_URL;
use super::ExecServerListenTransport;
use super::ExecServerWebSocketState;
use super::health_check_handler;
use super::parse_listen_url;
use super::readiness_handler;
use super::run_stdio_connection_with_io;
use crate::ExecServerRuntimePaths;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeParams;
use crate::protocol::InitializeResponse;
use crate::server::processor::ConnectionProcessor;
#[test]
fn parse_listen_url_accepts_default_websocket_url() {
@@ -125,6 +131,31 @@ fn parse_listen_url_accepts_websocket_url() {
);
}
#[tokio::test]
async fn readiness_handler_reports_shutdown_state() {
let shutdown_token = tokio_util::sync::CancellationToken::new();
let state = ExecServerWebSocketState {
processor: ConnectionProcessor::new(test_runtime_paths()),
shutdown_token: shutdown_token.clone(),
};
assert_eq!(
readiness_handler(State(state.clone())).await,
StatusCode::OK
);
shutdown_token.cancel();
assert_eq!(
readiness_handler(State(state)).await,
StatusCode::SERVICE_UNAVAILABLE
);
}
#[tokio::test]
async fn health_check_handler_stays_ok_during_shutdown() {
assert_eq!(health_check_handler().await, StatusCode::OK);
}
#[test]
fn parse_listen_url_rejects_invalid_websocket_url() {
let err = parse_listen_url("ws://localhost:1234")
codex-rs/exec-server/tests/common/exec_server.rs
+5
View File
@@ -144,6 +144,11 @@ impl ExecServerHarness {
Ok(())
}
pub(crate) async fn send_raw_binary(&mut self, bytes: Vec<u8>) -> anyhow::Result<()> {
self.websocket.send(Message::Binary(bytes.into())).await?;
Ok(())
}
pub(crate) async fn next_event(&mut self) -> anyhow::Result<JSONRPCMessage> {
self.next_event_with_timeout(EVENT_TIMEOUT).await
}
codex-rs/exec-server/tests/health.rs
+27
View File
@@ -0,0 +1,27 @@
#![cfg(unix)]
mod common;
use common::exec_server::exec_server;
use pretty_assertions::assert_eq;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_serves_kubernetes_style_probe_routes() -> anyhow::Result<()> {
let mut server = exec_server().await?;
let http_base_url = server
.websocket_url()
.strip_prefix("ws://")
.expect("websocket URL should use ws://");
let readyz = reqwest::get(format!("http://{http_base_url}/readyz")).await?;
assert_eq!(readyz.status(), reqwest::StatusCode::OK);
let healthz = reqwest::get(format!("http://{http_base_url}/healthz")).await?;
assert_eq!(healthz.status(), reqwest::StatusCode::OK);
let legacy_health = reqwest::get(format!("http://{http_base_url}/health")).await?;
assert_eq!(legacy_health.status(), reqwest::StatusCode::NOT_FOUND);
server.shutdown().await?;
Ok(())
}
codex-rs/exec-server/tests/websocket.rs
+36
View File
@@ -60,3 +60,39 @@ async fn exec_server_reports_malformed_websocket_json_and_keeps_running() -> any
server.shutdown().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_accepts_binary_websocket_json() -> anyhow::Result<()> {
let mut server = exec_server().await?;
let initialize_id = codex_app_server_protocol::RequestId::Integer(1);
let initialize = JSONRPCMessage::Request(codex_app_server_protocol::JSONRPCRequest {
id: initialize_id.clone(),
method: "initialize".to_string(),
params: Some(serde_json::to_value(InitializeParams {
client_name: "exec-server-binary-test".to_string(),
resume_session_id: None,
})?),
trace: None,
});
server
.send_raw_binary(serde_json::to_vec(&initialize)?)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
panic!("expected initialize response for binary input");
};
assert_eq!(id, initialize_id);
let initialize_response: InitializeResponse = serde_json::from_value(result)?;
Uuid::parse_str(&initialize_response.session_id)?;
server.shutdown().await?;
Ok(())
}