From 4eb322f4b0a206e4c7bc2e8eb4aa12d2e956988b Mon Sep 17 00:00:00 2001
From: Yxxx <45665172+marsjane@users.noreply.github.com>
Date: Sat, 8 Aug 2026 22:59:07 +0800
Subject: [PATCH] repl: Retry iopub connection on failure (#62358)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## 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>
---
 crates/repl/src/kernels/ssh_kernel.rs | 42 ++++++++++++++++++++++-----
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/crates/repl/src/kernels/ssh_kernel.rs b/crates/repl/src/kernels/ssh_kernel.rs
index 797b111a14..0317270b2a 100644
--- a/crates/repl/src/kernels/ssh_kernel.rs
+++ b/crates/repl/src/kernels/ssh_kernel.rs
@@ -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(