fix: MCP leaks in app-server (#17223)

The disconnect path now reuses the same teardown flow as explicit
unsubscribe, and the thread-state bookkeeping consistently reports only
threads that lost their last subscriber

https://github.com/openai/codex/issues/16895
This commit is contained in:
jif-oai
2026-04-10 15:31:26 +01:00
committed by GitHub
parent 8035cb03f1
commit 8d58899297
3 changed files with 190 additions and 84 deletions

View File

@@ -2,6 +2,7 @@ use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use app_test_support::create_mock_responses_server_sequence_unchecked;
use app_test_support::to_response;
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use codex_app_server_protocol::ClientInfo;
@@ -12,6 +13,10 @@ use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadLoadedListParams;
use codex_app_server_protocol::ThreadLoadedListResponse;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
use futures::SinkExt;
use futures::StreamExt;
use hmac::Hmac;
@@ -332,6 +337,37 @@ async fn websocket_transport_allows_unauthenticated_non_loopback_startup_by_defa
Ok(())
}
#[tokio::test]
async fn websocket_disconnect_unloads_last_subscribed_thread() -> Result<()> {
let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await;
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), &server.uri(), "never")?;
let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?;
let mut ws1 = connect_websocket(bind_addr).await?;
send_initialize_request(&mut ws1, /*id*/ 1, "ws_thread_owner").await?;
read_response_for_id(&mut ws1, /*id*/ 1).await?;
let thread_id = start_thread(&mut ws1, /*id*/ 2).await?;
assert_loaded_threads(&mut ws1, /*id*/ 3, &[thread_id.as_str()]).await?;
ws1.close(None).await.context("failed to close websocket")?;
drop(ws1);
let mut ws2 = connect_websocket(bind_addr).await?;
send_initialize_request(&mut ws2, /*id*/ 4, "ws_reconnect_client").await?;
read_response_for_id(&mut ws2, /*id*/ 4).await?;
wait_for_loaded_threads(&mut ws2, /*first_id*/ 5, &[]).await?;
process
.kill()
.await
.context("failed to stop websocket app-server process")?;
Ok(())
}
pub(super) async fn spawn_websocket_server(codex_home: &Path) -> Result<(Child, SocketAddr)> {
spawn_websocket_server_with_args(codex_home, "ws://127.0.0.1:0", &[]).await
}
@@ -564,6 +600,78 @@ pub(super) async fn send_initialize_request(
.await
}
async fn start_thread(stream: &mut WsClient, id: i64) -> Result<String> {
send_request(
stream,
"thread/start",
id,
Some(serde_json::to_value(ThreadStartParams {
model: Some("mock-model".to_string()),
..Default::default()
})?),
)
.await?;
let response = read_response_for_id(stream, id).await?;
let ThreadStartResponse { thread, .. } = to_response::<ThreadStartResponse>(response)?;
Ok(thread.id)
}
async fn assert_loaded_threads(stream: &mut WsClient, id: i64, expected: &[&str]) -> Result<()> {
let response = request_loaded_threads(stream, id).await?;
let mut actual = response.data;
actual.sort();
let mut expected = expected
.iter()
.map(|thread_id| (*thread_id).to_string())
.collect::<Vec<_>>();
expected.sort();
assert_eq!(actual, expected);
assert_eq!(response.next_cursor, None);
Ok(())
}
async fn wait_for_loaded_threads(
stream: &mut WsClient,
first_id: i64,
expected: &[&str],
) -> Result<()> {
let mut next_id = first_id;
let expected = expected
.iter()
.map(|thread_id| (*thread_id).to_string())
.collect::<Vec<_>>();
timeout(DEFAULT_READ_TIMEOUT, async {
loop {
let response = request_loaded_threads(stream, next_id).await?;
next_id += 1;
let mut actual = response.data;
actual.sort();
if actual == expected {
return Ok::<(), anyhow::Error>(());
}
sleep(Duration::from_millis(50)).await;
}
})
.await
.context("timed out waiting for loaded thread list")??;
Ok(())
}
async fn request_loaded_threads(
stream: &mut WsClient,
id: i64,
) -> Result<ThreadLoadedListResponse> {
send_request(
stream,
"thread/loaded/list",
id,
Some(serde_json::to_value(ThreadLoadedListParams::default())?),
)
.await?;
let response = read_response_for_id(stream, id).await?;
to_response::<ThreadLoadedListResponse>(response)
}
async fn send_config_read_request(stream: &mut WsClient, id: i64) -> Result<()> {
send_request(
stream,