From 0eb992219b31ed8f51fa16b49827cde58e10c103 Mon Sep 17 00:00:00 2001
From: Kyle Kelley <rgbkrk@gmail.com>
Date: Thu, 12 Dec 2024 09:26:16 -0800
Subject: [PATCH] Set User Agent for Jupyter websocket connections (#21910)

Some VPN configurations require that websockets present a user agent.
This adds it in directly for the repl usage. I wish there was a way to
reuse the user agent from the `cx.http_client`, but I'm not seeing a
simple way to do that for the moment since it's not on the `HttpClient`
trait.

No release notes since this feature hasn't been announced/exposed.

Release Notes:

- N/A
---
 Cargo.lock                                |  1 +
 crates/repl/Cargo.toml                    |  1 +
 crates/repl/src/kernels/remote_kernels.rs | 35 +++++++++++++++++++++--
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index aa1798732b..bae0781067 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -10363,6 +10363,7 @@ dependencies = [
  "alacritty_terminal",
  "anyhow",
  "async-dispatcher",
+ "async-tungstenite 0.28.1",
  "base64 0.22.1",
  "client",
  "collections",
diff --git a/crates/repl/Cargo.toml b/crates/repl/Cargo.toml
index 293a58e762..a49958c49e 100644
--- a/crates/repl/Cargo.toml
+++ b/crates/repl/Cargo.toml
@@ -16,6 +16,7 @@ doctest = false
 alacritty_terminal.workspace = true
 anyhow.workspace = true
 async-dispatcher.workspace = true
+async-tungstenite = { workspace = true, features = ["async-std", "async-tls"] }
 base64.workspace = true
 client.workspace = true
 collections.workspace = true
diff --git a/crates/repl/src/kernels/remote_kernels.rs b/crates/repl/src/kernels/remote_kernels.rs
index e1b41276fa..5a3a22ce87 100644
--- a/crates/repl/src/kernels/remote_kernels.rs
+++ b/crates/repl/src/kernels/remote_kernels.rs
@@ -3,6 +3,11 @@ use gpui::{Task, View, WindowContext};
 use http_client::{AsyncBody, HttpClient, Request};
 use jupyter_protocol::{ExecutionState, JupyterKernelspec, JupyterMessage, KernelInfoReply};
 
+use async_tungstenite::{
+    async_std::connect_async,
+    tungstenite::{client::IntoClientRequest, http::HeaderValue},
+};
+
 use futures::StreamExt;
 use smol::io::AsyncReadExt as _;
 
@@ -11,8 +16,8 @@ use crate::Session;
 use super::RunningKernel;
 use anyhow::Result;
 use jupyter_websocket_client::{
-    JupyterWebSocketReader, JupyterWebSocketWriter, KernelLaunchRequest, KernelSpecsResponse,
-    RemoteServer,
+    JupyterWebSocket, JupyterWebSocketReader, JupyterWebSocketWriter, KernelLaunchRequest,
+    KernelSpecsResponse, RemoteServer,
 };
 use std::{fmt::Debug, sync::Arc};
 
@@ -151,7 +156,31 @@ impl RemoteRunningKernel {
             )
             .await?;
 
-            let (kernel_socket, _response) = remote_server.connect_to_kernel(&kernel_id).await?;
+            let ws_url = format!(
+                "{}/api/kernels/{}/channels?token={}",
+                remote_server.base_url.replace("http", "ws"),
+                kernel_id,
+                remote_server.token
+            );
+
+            let mut req: Request<()> = ws_url.into_client_request()?;
+            let headers = req.headers_mut();
+
+            headers.insert(
+                "User-Agent",
+                HeaderValue::from_str(&format!(
+                    "Zed/{} ({}; {})",
+                    "repl",
+                    std::env::consts::OS,
+                    std::env::consts::ARCH
+                ))?,
+            );
+
+            let response = connect_async(req).await;
+
+            let (ws_stream, _response) = response?;
+
+            let kernel_socket = JupyterWebSocket { inner: ws_stream };
 
             let (mut w, mut r): (JupyterWebSocketWriter, JupyterWebSocketReader) =
                 kernel_socket.split();