Reconnect gRPC code-mode sessions after host restarts (#38257)

## What changed

- Reopen a cached code-mode session when its gRPC host stops, while
  serializing concurrent reconnection attempts and coordinating shutdown.
- Scope cell IDs to the new host generation so callbacks remain consistent
  and stale `wait` or `terminate` requests are rejected.
- Accept both `unix://` and `unix:` endpoints for gRPC hosts on Unix systems.

## Testing

- Cover host restart recovery, concurrent execution after reconnection,
  generation-aware callbacks and cell operations, stale cell rejection, and
  Unix socket execution.

GitOrigin-RevId: 548e168fdcef7f7d54bd32262e614886bf7bdd32
This commit is contained in:
Channing Conger
2026-08-12 22:00:20 +00:00
committed by copyberry
parent 020f6c963e
commit bde723ae7d
6 changed files with 841 additions and 5 deletions
+13 -5
View File
@@ -39,12 +39,16 @@ mod callbacks;
mod completion;
mod conversion;
mod deadline;
mod generation;
mod operations;
mod reconnect;
mod state;
mod transport;
type GrpcClient = CodeModeHostClient<GrpcTransport>;
const SHUTDOWN_ERROR: &str = "code mode session is shutting down";
/// Creates code-mode sessions over an HTTP/2 gRPC connection.
#[derive(Clone)]
pub struct GrpcCodeModeSessionProvider {
@@ -52,7 +56,7 @@ pub struct GrpcCodeModeSessionProvider {
}
impl GrpcCodeModeSessionProvider {
/// Connects lazily to an `http://` or `https://` gRPC endpoint.
/// Connects lazily to an `http://`, `https://`, or `unix://` gRPC endpoint.
pub fn new(endpoint: impl Into<String>) -> Self {
Self::with_http_client_factory(
endpoint,
@@ -166,9 +170,13 @@ impl CodeModeSessionProvider for GrpcCodeModeSessionProvider {
limits: CodeModeSessionCellExecutionLimits,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(async move {
self.open_binding(delegate, limits)
.await
.map(|session| session as _)
let session = Arc::new(reconnect::ReconnectableSession::new(
self.clone(),
delegate,
limits,
));
session.initialize().await?;
Ok(session as _)
})
}
}
@@ -246,7 +254,7 @@ impl SessionInner {
.unwrap_or_else(PoisonError::into_inner)
.require_open()?;
if self.shutdown_requested.load(Ordering::Acquire) {
return Err("code mode session is shutting down".to_string());
return Err(SHUTDOWN_ERROR.to_string());
}
Ok(())
}