mirror of
https://github.com/openai/codex.git
synced 2026-11-13 21:30:47 +00:00
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:
committed by
copyberry
parent
7844815bd7
commit
61c106ed66
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user