Wait for the in-process initialization handshake to finish (#60871)

## Why

An immediate request after `start` could compete with the pending `initialized`
notification for channel capacity and fail with `WouldBlock`.

## What changed

Wait for an acknowledgment that `initialized` has been processed before returning
from `start`. Shut down the runtime if sending the notification or receiving its
acknowledgment fails. Ordinary notifications remain fire-and-forget.

## Testing

Require the first request after startup to succeed without retries with channel
capacities of zero (clamped to one) and one.

GitOrigin-RevId: 119a9a4375109f6b85aee0a9a4f8398a69269567
This commit is contained in:
goodman-oai
2026-11-11 22:30:35 +00:00
committed by copyberry
parent 392ced7d3a
commit 3a45bbba3c
+54 -29
View File
@@ -192,9 +192,9 @@ pub enum InProcessServerEvent {
/// Internal message sent from [`InProcessClientHandle`] methods to the runtime task.
///
/// Requests carry a oneshot sender for the response; notifications and server-request
/// replies are fire-and-forget from the caller's perspective (transport errors are
/// caught by `try_send` on the outer channel).
/// Requests carry a oneshot sender for the response. Startup notifications can
/// acknowledge processing; ordinary notifications and server-request replies are
/// fire-and-forget (transport errors are caught by `try_send` on the outer channel).
enum InProcessClientMessage {
Request {
request: Box<ClientRequest>,
@@ -203,6 +203,7 @@ enum InProcessClientMessage {
},
Notification {
notification: ClientNotification,
processed_tx: Option<oneshot::Sender<()>>,
},
ServerRequestResponse {
request_id: RequestId,
@@ -219,7 +220,10 @@ enum InProcessClientMessage {
enum ProcessorCommand {
Request(Box<ClientRequest>, tokio_util::sync::CancellationToken),
Notification(ClientNotification),
Notification {
notification: ClientNotification,
processed_tx: Option<oneshot::Sender<()>>,
},
}
#[derive(Clone)]
@@ -246,7 +250,10 @@ impl InProcessClientSender {
}
pub fn notify(&self, notification: ClientNotification) -> IoResult<()> {
self.try_send_client_message(InProcessClientMessage::Notification { notification })
self.try_send_client_message(InProcessClientMessage::Notification {
notification,
processed_tx: None,
})
}
pub fn respond_to_server_request(&self, request_id: RequestId, result: Result) -> IoResult<()> {
@@ -379,9 +386,10 @@ impl InProcessClientHandle {
/// Starts an in-process app-server runtime and performs initialize handshake.
///
/// This function sends `initialize` followed by `initialized` before returning
/// the handle, so callers receive a ready-to-use runtime. If initialize fails,
/// the runtime is shut down and an `InvalidData` error is returned.
/// This function waits for `initialize` and `initialized` to be processed before
/// returning the handle, so callers receive a ready-to-use runtime. Rejected
/// initialization shuts down the runtime and returns an `InvalidData` error.
/// Transport failures while processing `initialized` also shut down the runtime.
pub async fn start(mut args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
if let Ok(Some(err)) = check_execpolicy_for_warnings(&args.config.config_layer_stack).await {
let (path, range) = crate::exec_policy_warning_location(&err);
@@ -409,7 +417,26 @@ pub async fn start(mut args: InProcessStartArgs) -> IoResult<InProcessClientHand
format!("in-process initialize failed: {}", error.message),
));
}
client.notify(ClientNotification::Initialized)?;
let (processed_tx, processed_rx) = oneshot::channel();
let initialized_result = async {
client
.client
.try_send_client_message(InProcessClientMessage::Notification {
notification: ClientNotification::Initialized,
processed_tx: Some(processed_tx),
})?;
processed_rx.await.map_err(|err| {
IoError::new(
ErrorKind::BrokenPipe,
format!("in-process initialized acknowledgment channel closed: {err}"),
)
})
}
.await;
if let Err(error) = initialized_result {
let _ = client.shutdown().await;
return Err(error);
}
Ok(client)
}
@@ -585,8 +612,11 @@ async fn start_uninitialized(mut args: InProcessStartArgs) -> IoResult<InProcess
processor.send_initialize_notifications().await;
}
}
Some(ProcessorCommand::Notification(notification)) => {
Some(ProcessorCommand::Notification { notification, processed_tx }) => {
processor.process_client_notification(notification).await;
if let Some(processed_tx) = processed_tx {
let _ = processed_tx.send(());
}
}
None => {
break;
@@ -678,8 +708,8 @@ async fn start_uninitialized(mut args: InProcessStartArgs) -> IoResult<InProcess
}
}
}
Some(InProcessClientMessage::Notification { notification }) => {
match processor_tx.try_send(ProcessorCommand::Notification(notification)) {
Some(InProcessClientMessage::Notification { notification, processed_tx }) => {
match processor_tx.try_send(ProcessorCommand::Notification { notification, processed_tx }) {
Ok(()) => {}
Err(mpsc::error::TrySendError::Full(_)) => {
warn!("dropping in-process client notification (queue full)");
@@ -1047,29 +1077,24 @@ mod tests {
}
}
#[test_case::test_case(0; "zero_capacity_is_clamped")]
#[test_case::test_case(1; "one_slot_capacity")]
#[tokio::test]
async fn in_process_start_clamps_zero_channel_capacity() {
async fn in_process_start_accepts_immediate_request(channel_capacity: usize) {
let client = start_test_client_with_options(
SessionSource::Cli,
/*channel_capacity*/ 0,
channel_capacity,
/*capabilities*/ None,
)
.await;
let response = loop {
match client
.request(ClientRequest::ConfigRequirementsRead {
request_id: RequestId::Integer(4),
params: None,
})
.await
{
Ok(response) => break response.expect("request should succeed"),
Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
tokio::task::yield_now().await;
}
Err(err) => panic!("request transport should work: {err}"),
}
};
let response = client
.request(ClientRequest::ConfigRequirementsRead {
request_id: RequestId::Integer(4),
params: None,
})
.await
.expect("the first request should not compete with the initialize handshake")
.expect("request should succeed");
let _parsed: ConfigRequirementsReadResponse =
serde_json::from_value(response).expect("response should match v2 schema");
client