Cancel abandoned cloud thread list requests (#61342)

## Why

An in-flight `ThreadList` request can block the cloud worker from processing subsequent requests even after its caller drops it.

## What changed

Skip queued `ThreadList` requests whose response receiver is closed, and cancel in-flight list handling when the receiver closes. Thread listing has no conversation state to commit, so the worker can safely move on to the next command.

## Testing

Add a regression test that holds a list RPC open, drops the caller's request, and verifies that a subsequent `GetAccount` request completes before the list RPC is released.

GitOrigin-RevId: 7c74e3305d58b85224d9882ea7e24b1ae791ec52
This commit is contained in:
Benjamin Carlsson
2026-11-13 18:45:02 +00:00
committed by copyberry
parent 7844815bd7
commit 61c106ed66
2 changed files with 93 additions and 6 deletions
@@ -1,14 +1,16 @@
#![cfg(test)]
//! Ensures only a successfully attached thread can receive new turns.
//! Exercises native worker progress and reply attachment.
use super::*;
use futures::stream;
use pretty_assertions::assert_eq;
use std::sync::Mutex;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::body_bytes;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
@@ -223,6 +225,83 @@ async fn failed_archive_loaded_list_tracks_attachment_not_cached_metadata() -> R
Ok(())
}
#[tokio::test]
async fn dropping_an_in_flight_list_unblocks_the_native_worker() -> Result<()> {
let server = MockServer::start().await;
let home = tempfile::tempdir()?;
let config = crate::legacy_core::config::ConfigBuilder::default()
.codex_home(home.path().to_owned())
.fallback_cwd(Some(home.path().to_owned()))
.without_project_context()
.build()
.await?;
let native = cloud::Client::new(
&config.http_client_factory(),
&server.uri(),
cloud::Credentials::new("synthetic-bearer", "synthetic-account")?,
)?;
let account = protocol::GetAccountResponse {
user_id: None,
account: None,
requires_openai_auth: true,
workspace_routing: None,
};
let (client, commands, output) = CloudClient::channel(/*capacity*/ 32);
let client = AppServerClient::Cloud(client);
let mut requests = requests::Requests {
client: Arc::new(native),
account: account.clone(),
loaded: None,
warned_partial_history: None,
};
tokio::spawn(async move {
run_worker(commands, &mut requests, &mut Replies::default(), &output).await;
});
let (list_arrived, mut arrival) = mpsc::unbounded_channel();
let (release_list, held) = std::sync::mpsc::channel::<()>();
let held = Mutex::new(held);
Mock::given(method("POST"))
.and(path(
"/oaiproto.codex_app_server_backend.ThreadService/ListThreads",
))
.and(header("authorization", "Bearer synthetic-bearer"))
.respond_with(move |_: &wiremock::Request| {
let _ = list_arrived.send(());
// Wiremock serves on a dedicated thread. Hold the RPC until the sender drops.
let _ = held.lock().unwrap().recv();
ResponseTemplate::new(503)
})
.expect(1)
.mount(&server)
.await;
let mut list = Box::pin(client.request(protocol::ClientRequest::ThreadList {
request_id: protocol::RequestId::Integer(1),
params: serde_json::from_value(serde_json::json!({"limit": 10}))?,
}));
tokio::select! {
seen = tokio::time::timeout(crate::SHUTDOWN_TIMEOUT, arrival.recv()) => {
assert_eq!(seen?, Some(()));
}
_ = &mut list => panic!("list finished before the fixture held it"),
}
drop(list);
let account_response = tokio::time::timeout(
crate::SHUTDOWN_TIMEOUT,
client.request_typed::<protocol::GetAccountResponse>(protocol::ClientRequest::GetAccount {
request_id: protocol::RequestId::Integer(2),
params: protocol::GetAccountParams {
refresh_token: false,
},
}),
)
.await;
drop(release_list);
assert_eq!(account_response??, account);
client.shutdown().await?;
Ok(())
}
#[tokio::test]
async fn failed_replacement_keeps_old_thread_until_new_attachment_succeeds() -> Result<()> {
let (output, mut notifications) = mpsc::unbounded_channel();
@@ -277,7 +277,7 @@ async fn run_worker(
output: &mpsc::UnboundedSender<AppServerEvent>,
) {
loop {
let command = tokio::select! {
let mut command = tokio::select! {
_ = output.closed() => return,
() = replies.forward_next(output) => continue,
command = commands.recv() => match command { Some(command) => command, None => return },
@@ -335,10 +335,18 @@ async fn run_worker(
)),
};
let result = match request {
Ok((request, mode, project)) => tokio::select! {
_ = output.closed() => return,
result = handle_request(request, mode, project, requests, replies, output) => result,
},
Ok((request, mode, project)) => {
// Inventory is a read with no conversation state to commit.
let is_thread_list = matches!(&request, protocol::ClientRequest::ThreadList { .. });
if is_thread_list && command.response.is_closed() {
continue;
}
tokio::select! {
_ = output.closed() => return,
_ = command.response.closed(), if is_thread_list => continue,
result = handle_request(request, mode, project, requests, replies, output) => result,
}
}
Err(error) => Err(error),
};
if result.is_err()