mirror of
https://github.com/openai/codex.git
synced 2026-04-27 08:05:51 +00:00
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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user