mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Resume exec-server sessions after disconnect (#28512)
Supersedes #28288 (closed). ## Why A short WebSocket interruption currently ends every client-side process handle, even though exec-server keeps the server session and its processes alive for a short time. This is especially visible for executor-backed stdio MCP servers: a temporary connection loss becomes a permanent `Transport closed` error. The server already has the information needed to resume the session, but the client opens a fresh session instead of using it. This change reconnects below the process and MCP layers. Existing process handles stay valid, missed output is recovered, and the same server-side processes continue running. ## State machine One logical `ExecServerClient` stays alive while its underlying RPC connection changes generations. ```text transport closes +------------------------------------------------+ | v +-------------+ +-------------+ | Connected | | Recovering | +-------------+ +-------------+ ^ | | session resumed, processes caught up | retryable error +------------------------------------------------+ loops until deadline | | deadline or permanent error v +-------------+ | Failed | +-------------+ ``` ### `Connected` - New RPC calls use the current connection. - Process notifications are published in sequence order. - A disconnect only starts recovery if it came from the current connection generation. Late events from older generations cannot replace the active connection. ### `Recovering` - New calls wait instead of choosing a half-connected RPC client. - Existing process handles, wake subscriptions, and event subscriptions stay open. - Streaming HTTP response bodies fail immediately because their byte streams cannot be resumed safely. - Recovery first waits for process starts that were already in flight. A start whose result became ambiguous is cleaned up after reconnection instead of being silently adopted. - The client reconnects with the learned `session_id`. The server may briefly report that the old connection is still attached, so that error is retried until the detach finishes. - The notification consumer starts before the resume handshake completes. This prevents a busy process from filling the notification queue and blocking the initialize response. - Before installing the new connection, the client catches up every recoverable process with `process/read`. ### `Failed` - Recovery stops after 25 seconds or after a permanent error. - Waiting calls are released with one stable disconnect error. - Existing process sessions receive a terminal failure instead of waiting forever. ## Recovering process events Output, exit, and close events share one sequence. During normal operation, the client buffers early events until every lower sequence has been published. After reconnection, the client reads each process starting after its last published sequence: 1. Retained output chunks are inserted by sequence number. 2. Exit and close state are reconstructed in their sequence positions. 3. Events already received as live notifications are ignored as duplicates. 4. Newly contiguous events are published in order. 5. If the server no longer retains enough output to fill a sequence gap, only that process is terminated and failed. The recovered connection remains usable for other processes. The server reports its full next event sequence for unbounded reads, including exit and close events. Closed processes remain readable for the same 30-second window used to retain detached sessions. ## Other details - Detached server sessions are retained for 30 seconds, leaving margin around the client's 25-second recovery deadline. - Session attach and detach update the active notification sender under the same attachment lock, so an old connection cannot clear a newly attached sender. - A dedicated error code distinguishes the temporary "session is still attached" race from permanent initialization errors. - Process starts are identity-checked on both client and server. Cleanup from an older start cannot remove a newer process that reused the same ID. - Mutating requests that were already in flight when the transport closed are not replayed, because the client cannot know whether the server applied them. Requests started after recovery is known wait for the replacement connection. - We assume the server/client version stays in sync (on the before/after this PR) ## User impact Long-running commands and stdio MCP servers can survive a temporary exec-server WebSocket interruption without changing process IDs or losing output produced during the outage.
This commit is contained in:
jif
@@ -1,5 +1,6 @@
|
||||
mod common;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Context;
|
||||
@@ -21,6 +22,7 @@ use tempfile::TempDir;
|
||||
use test_case::test_case;
|
||||
use tokio::sync::watch;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::sleep;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use common::DELAYED_OUTPUT_AFTER_EXIT_PARENT_ARG;
|
||||
@@ -30,7 +32,7 @@ use common::exec_server::exec_server;
|
||||
|
||||
struct ProcessContext {
|
||||
backend: Arc<dyn ExecBackend>,
|
||||
server: Option<ExecServerHarness>,
|
||||
_server: Option<ExecServerHarness>,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
@@ -55,13 +57,13 @@ async fn create_process_context(use_remote: bool) -> Result<ProcessContext> {
|
||||
let environment = Environment::create_for_tests(Some(server.websocket_url().to_string()))?;
|
||||
Ok(ProcessContext {
|
||||
backend: environment.get_exec_backend(),
|
||||
server: Some(server),
|
||||
_server: Some(server),
|
||||
})
|
||||
} else {
|
||||
let environment = Environment::create_for_tests(/*exec_server_url*/ None)?;
|
||||
Ok(ProcessContext {
|
||||
backend: environment.get_exec_backend(),
|
||||
server: None,
|
||||
_server: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -634,88 +636,145 @@ async fn assert_exec_process_preserves_queued_events_before_subscribe(
|
||||
#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")]
|
||||
// Serialize tests that launch a real exec-server process through the full CLI.
|
||||
#[serial_test::serial(remote_exec_server)]
|
||||
async fn remote_exec_process_reports_transport_disconnect() -> Result<()> {
|
||||
let mut context = create_process_context(/*use_remote*/ true).await?;
|
||||
let session = context
|
||||
.backend
|
||||
async fn remote_exec_process_recovers_after_transport_disconnect() -> Result<()> {
|
||||
let server = exec_server().await?;
|
||||
let mut proxy = server.disconnectable_websocket_proxy().await?;
|
||||
let environment = Environment::create_for_tests(Some(proxy.websocket_url().to_string()))?;
|
||||
let backend = environment.get_exec_backend();
|
||||
let temp_dir = TempDir::new()?;
|
||||
let gate_path = temp_dir.path().join("release-output");
|
||||
let emitted_path = temp_dir.path().join("output-emitted");
|
||||
let session = backend
|
||||
.start(ExecParams {
|
||||
process_id: ProcessId::from("proc-disconnect"),
|
||||
process_id: ProcessId::from("proc-recover"),
|
||||
argv: vec![
|
||||
"/bin/sh".to_string(),
|
||||
"-c".to_string(),
|
||||
"sleep 10".to_string(),
|
||||
concat!(
|
||||
"printf 'ready:%s\\n' \"$$\"; ",
|
||||
"while [ ! -f \"$GATE\" ]; do /bin/sleep 0.01; done; ",
|
||||
"printf 'during:%s\\n' \"$$\"; ",
|
||||
": > \"$EMITTED\"; ",
|
||||
"IFS= read -r line; ",
|
||||
"printf 'after:%s:%s\\n' \"$$\" \"$line\"; ",
|
||||
"exit 7",
|
||||
)
|
||||
.to_string(),
|
||||
],
|
||||
cwd: PathUri::from_path(std::env::current_dir()?)?,
|
||||
env_policy: /*env_policy*/ None,
|
||||
env: Default::default(),
|
||||
env: HashMap::from([
|
||||
(
|
||||
"GATE".to_string(),
|
||||
gate_path.to_string_lossy().into_owned(),
|
||||
),
|
||||
(
|
||||
"EMITTED".to_string(),
|
||||
emitted_path.to_string_lossy().into_owned(),
|
||||
),
|
||||
]),
|
||||
tty: false,
|
||||
pipe_stdin: false,
|
||||
pipe_stdin: true,
|
||||
arg0: None,
|
||||
})
|
||||
.await?;
|
||||
|
||||
let process = Arc::clone(&session.process);
|
||||
let mut events = process.subscribe_events();
|
||||
let process_for_pending_read = Arc::clone(&process);
|
||||
let pending_read = tokio::spawn(async move {
|
||||
process_for_pending_read
|
||||
let mut output = Vec::new();
|
||||
let mut last_seq = 0;
|
||||
while !output.ends_with(b"\n") {
|
||||
match timeout(Duration::from_secs(5), events.recv()).await?? {
|
||||
ExecProcessEvent::Output(chunk) => {
|
||||
assert_eq!(chunk.seq, last_seq + 1);
|
||||
last_seq = chunk.seq;
|
||||
output.extend_from_slice(&chunk.chunk.into_inner());
|
||||
}
|
||||
event => anyhow::bail!("expected ready output before disconnect, got {event:?}"),
|
||||
}
|
||||
}
|
||||
let ready = String::from_utf8(output.clone())?;
|
||||
let pid = ready
|
||||
.strip_prefix("ready:")
|
||||
.and_then(|line| line.strip_suffix('\n'))
|
||||
.context("ready output should contain the process id")?
|
||||
.to_string();
|
||||
|
||||
proxy.pause_and_disconnect().await?;
|
||||
tokio::fs::write(&gate_path, b"").await?;
|
||||
timeout(Duration::from_secs(5), async {
|
||||
while tokio::fs::metadata(&emitted_path).await.is_err() {
|
||||
sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.context("process did not emit output while disconnected")?;
|
||||
|
||||
let process_for_read = Arc::clone(&process);
|
||||
let mut pending_read = tokio::spawn(async move {
|
||||
process_for_read
|
||||
.read(
|
||||
/*after_seq*/ None,
|
||||
/*after_seq*/ Some(last_seq),
|
||||
/*max_bytes*/ None,
|
||||
/*wait_ms*/ Some(60_000),
|
||||
/*wait_ms*/ Some(0),
|
||||
)
|
||||
.await
|
||||
});
|
||||
let server = context
|
||||
.server
|
||||
.as_mut()
|
||||
.expect("remote context should include exec-server harness");
|
||||
server.shutdown().await?;
|
||||
|
||||
let event = timeout(Duration::from_secs(2), events.recv()).await??;
|
||||
let ExecProcessEvent::Failed(event_message) = event else {
|
||||
anyhow::bail!("expected process failure event, got {event:?}");
|
||||
};
|
||||
assert!(
|
||||
event_message.starts_with("exec-server transport disconnected"),
|
||||
"unexpected failure event: {event_message}"
|
||||
timeout(Duration::from_millis(200), &mut pending_read)
|
||||
.await
|
||||
.is_err(),
|
||||
"process reads should wait while recovery is in progress"
|
||||
);
|
||||
proxy.resume()?;
|
||||
|
||||
let recovered_read = timeout(Duration::from_secs(5), pending_read)
|
||||
.await
|
||||
.context("timed out waiting for a read after recovery")??;
|
||||
let recovered_read = recovered_read?;
|
||||
assert_eq!(recovered_read.failure, None);
|
||||
let recovered_output = recovered_read
|
||||
.chunks
|
||||
.into_iter()
|
||||
.flat_map(|chunk| chunk.chunk.into_inner())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
String::from_utf8(recovered_output)?,
|
||||
format!("during:{pid}\n")
|
||||
);
|
||||
|
||||
let pending_response = timeout(Duration::from_secs(2), pending_read).await???;
|
||||
let pending_message = pending_response
|
||||
.failure
|
||||
.expect("pending read should surface disconnect as a failure");
|
||||
assert!(
|
||||
pending_message.starts_with("exec-server transport disconnected"),
|
||||
"unexpected pending failure message: {pending_message}"
|
||||
);
|
||||
let write = timeout(Duration::from_secs(5), process.write(b"hello\n".to_vec()))
|
||||
.await
|
||||
.context("timed out waiting for a write after recovery")??;
|
||||
assert_eq!(write.status, WriteStatus::Accepted);
|
||||
|
||||
let mut wake_rx = process.subscribe_wake();
|
||||
let response = read_process_until_change(process, &mut wake_rx, /*after_seq*/ None).await?;
|
||||
let message = response
|
||||
.failure
|
||||
.expect("disconnect should surface as a failure");
|
||||
assert!(
|
||||
message.starts_with("exec-server transport disconnected"),
|
||||
"unexpected failure message: {message}"
|
||||
);
|
||||
assert!(
|
||||
response.closed,
|
||||
"disconnect should close the process session"
|
||||
);
|
||||
|
||||
let write_result = timeout(
|
||||
Duration::from_secs(2),
|
||||
session.process.write(b"hello".to_vec()),
|
||||
)
|
||||
.await
|
||||
.context("timed out waiting for write after disconnect")?;
|
||||
let write_error = write_result.expect_err("write after disconnect should fail");
|
||||
assert!(
|
||||
write_error
|
||||
.to_string()
|
||||
.starts_with("exec-server transport disconnected"),
|
||||
"unexpected write error: {write_error}"
|
||||
let mut saw_exit = false;
|
||||
loop {
|
||||
match timeout(Duration::from_secs(5), events.recv()).await?? {
|
||||
ExecProcessEvent::Output(chunk) => {
|
||||
assert_eq!(chunk.seq, last_seq + 1);
|
||||
last_seq = chunk.seq;
|
||||
output.extend_from_slice(&chunk.chunk.into_inner());
|
||||
}
|
||||
ExecProcessEvent::Exited { seq, exit_code } => {
|
||||
assert_eq!(seq, last_seq + 1);
|
||||
assert_eq!(exit_code, 7);
|
||||
last_seq = seq;
|
||||
saw_exit = true;
|
||||
}
|
||||
ExecProcessEvent::Closed { seq } => {
|
||||
assert!(saw_exit, "closed must be delivered after exit");
|
||||
assert_eq!(seq, last_seq + 1);
|
||||
break;
|
||||
}
|
||||
ExecProcessEvent::Failed(message) => {
|
||||
anyhow::bail!("process recovery failed: {message}");
|
||||
}
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
String::from_utf8(output)?,
|
||||
format!("ready:{pid}\nduring:{pid}\nafter:{pid}:hello\n")
|
||||
);
|
||||
|
||||
Ok(())
|
||||
|
||||
Reference in New Issue
Block a user