repl: Retry iopub connection on failure (#62358)

## Context

This is a rebase of #52302, which was approved by @Veykril in May but
still needed a rebase before it could be merged — it seems the original
author @MostlyKIGuess forgot about it. I rely on this fix daily for REPL
over SSH on an unstable connection and need it rather urgently, so I
cherry-picked the original commit (with authorship preserved) and
resubmitted it as a new PR.

Closes #51834
Supersedes #52302

## How to Review

The code itself was already reviewed and approved in #52302. This PR
only rebases it onto latest `main`, with one simple conflict resolution
against #53014: when all retries are exhausted, the error now uses the
improved message from that PR (suggesting `pip install ipykernel` on the
remote host) instead of the plain one. Comparing this diff against the
original PR's diff should make the review straightforward.

## Self-Review Checklist

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content adheres to Zed's UI standards
([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
and
[icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md)
guidelines)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- repl: Fixed iopub connection failures over SSH on slow or unstable
connections by retrying with exponential backoff.

Co-authored-by: mostlykiguess <bruvistrue93@gmail.com>
This commit is contained in:
Yxxxandmostlykiguess authored and GitHub committed 2026-08-08 14:59:07 +00:00
1 parent 0d952edd5a
commit 4eb322f4b0
1 file changed
+34 -8
+34 -8
View File
@@ -160,8 +160,9 @@ impl SshRunningKernel {
attempt + 1
);
// giving the tunnel a moment to fully establish forwarding
// waiting same as WSL
cx.background_executor()
.timer(std::time::Duration::from_millis(500))
.timer(std::time::Duration::from_secs(2))
.await;
break;
}
@@ -209,13 +210,38 @@ impl SshRunningKernel {
serde_json::from_value(local_connection_info)?;
let session_id = uuid::Uuid::new_v4().to_string();
let output_socket = runtimelib::create_client_iopub_connection(
&connection_info_struct,
"",
&session_id,
)
.await
.context("Failed to create iopub connection. Is `ipykernel` installed in the remote environment? Try running `pip install ipykernel` on the remote host.")?;
let output_socket = {
let max_retries = 5;
let mut retry = 0;
loop {
match runtimelib::create_client_iopub_connection(
&connection_info_struct,
"",
&session_id,
)
.await
{
Ok(socket) => break socket,
Err(err) => {
if retry >= max_retries {
anyhow::bail!(
"Failed to create iopub connection: {}. Is `ipykernel` installed in the remote environment? Try running `pip install ipykernel` on the remote host.",
err
);
}
log::debug!(
"Retrying iopub connection (attempt {}/{})",
retry + 1,
max_retries
);
cx.background_executor()
.timer(std::time::Duration::from_millis(200 << retry))
.await;
retry += 1;
}
}
}
};
let peer_identity = runtimelib::peer_identity_for_session(&session_id)?;
let shell_socket = runtimelib::create_client_shell_connection_with_identity(