Add --gateway support for Jupyter Kernel Gateway

Cargo.toml

@@ -37,9 +37,9 @@ dirs = "5.0"
 dialoguer = "0.11"
 chrono = { version = "0.4", features = ["serde"] }
 
-# Remote mode dependencies
+# Remote mode dependencies (server + kernel-gateway WS)
 reqwest = { version = "0.11", features = ["json", "native-tls-vendored"] }
-tokio-tungstenite = "0.21"
+tokio-tungstenite = { version = "0.21", features = ["native-tls-vendored"] }
 futures-util = "0.3"
 url = "2.5"
 

src/commands/create_notebook.rs

@@ -106,7 +106,8 @@ pub fn execute(args: CreateArgs) -> Result<()> {
                 .block_on(client.save_notebook(&server_path, &notebook))
                 .context("Failed to create notebook on server")?;
         }
-        crate::execution::types::ExecutionMode::Local => {
+        crate::execution::types::ExecutionMode::Local
+        | crate::execution::types::ExecutionMode::RemoteKernel { .. } => {
             notebook::write_notebook_atomic(&path, &notebook)
                 .context("Failed to write notebook")?;
         }

src/commands/execute_notebook.rs

@@ -43,13 +43,29 @@ pub struct ExecuteNotebookArgs {
     pub end: Option<i32>,
 
     /// Remote server URL (enables remote mode)
-    #[arg(long)]
+    #[arg(long, conflicts_with = "gateway")]
     pub server: Option<String>,
 
     /// Authentication token for remote server
     #[arg(long)]
     pub token: Option<String>,
 
+    /// Kernel gateway URL (enables remote-kernel mode, e.g. http://host:8888)
+    #[arg(long, conflicts_with = "server")]
+    pub gateway: Option<String>,
+
+    /// Authentication token for kernel gateway
+    #[arg(long, requires = "gateway")]
+    pub gateway_token: Option<String>,
+
+    /// Authorization scheme for the gateway token (e.g. "token", "Bearer")
+    #[arg(long, requires = "gateway", default_value = "token")]
+    pub gateway_auth_scheme: String,
+
+    /// Kernel ID on the gateway (auto-discovered if omitted)
+    #[arg(long, requires = "gateway")]
+    pub kernel_id: Option<String>,
+
     /// Output in JSON format instead of text
     #[arg(long)]
     pub json: bool,
@@ -103,9 +119,23 @@ async fn execute_async(args: ExecuteNotebookArgs) -> Result<()> {
 
     let file_path = common::normalize_notebook_path(&args.file);
 
-    // Determine execution mode before reading — remote mode reads from server
-    let mode =
-        crate::commands::common::resolve_execution_mode(args.server.clone(), args.token.clone())?;
+    // Determine execution mode before reading — remote (server) mode reads
+    // from the server, remote-kernel mode reads locally but executes on a
+    // remote kernel gateway, and local mode does both locally.
+    let mode = if let Some(gateway_url) = args.gateway.clone() {
+        let token = args
+            .gateway_token
+            .clone()
+            .context("Must specify --gateway-token when using --gateway")?;
+        ExecutionMode::RemoteKernel {
+            gateway_url,
+            token,
+            kernel_id: args.kernel_id.clone(),
+            auth_scheme: args.gateway_auth_scheme.clone(),
+        }
+    } else {
+        common::resolve_execution_mode(args.server.clone(), args.token.clone())?
+    };
 
     // For remote mode, compute the server-relative path once and reuse it
     // for both the Contents API read and the session identifier.
@@ -129,7 +159,9 @@ async fn execute_async(args: ExecuteNotebookArgs) -> Result<()> {
             )
             .await?
         }
-        ExecutionMode::Local => read_notebook(&file_path).context("Failed to read notebook")?,
+        ExecutionMode::Local | ExecutionMode::RemoteKernel { .. } => {
+            read_notebook(&file_path).context("Failed to read notebook")?
+        }
     };
 
     // Determine cell range
@@ -163,17 +195,18 @@ async fn execute_async(args: ExecuteNotebookArgs) -> Result<()> {
         );
     }
 
+    // Get kernel from notebook metadata if not specified
     let notebook_kernel = notebook
         .metadata
         .kernelspec
         .as_ref()
         .map(|ks| ks.name.as_str());
 
     // For remote mode, reuse the pre-computed server-relative path.
-    // For local mode, use absolute path for working directory determination.
+    // For local and remote-kernel modes, use absolute path for working directory determination.
     let notebook_identifier = match &mode {
-        ExecutionMode::Remote { .. } => server_path.unwrap(),
-        ExecutionMode::Local => {
+        ExecutionMode::Remote { .. } => server_path.clone().expect("set for Remote mode"),
+        ExecutionMode::Local | ExecutionMode::RemoteKernel { .. } => {
             let abs =
                 std::fs::canonicalize(&file_path).context("Failed to resolve notebook path")?;
             abs.to_str()
@@ -203,8 +236,10 @@ async fn execute_async(args: ExecuteNotebookArgs) -> Result<()> {
     let mut execution_results: HashMap<usize, crate::execution::types::ExecutionResult> =
         HashMap::new();
 
-    let is_streaming = matches!(mode, ExecutionMode::Remote { .. })
-        && matches!(format, OutputFormat::Text | OutputFormat::Markdown);
+    let is_streaming = matches!(
+        mode,
+        ExecutionMode::Remote { .. } | ExecutionMode::RemoteKernel { .. }
+    ) && matches!(format, OutputFormat::Text | OutputFormat::Markdown);
 
     // For streaming remote text mode, print notebook header before execution
     if is_streaming {
@@ -308,7 +343,7 @@ async fn execute_async(args: ExecuteNotebookArgs) -> Result<()> {
 
     // Persist changes based on mode
     match mode {
-        ExecutionMode::Local => {
+        ExecutionMode::Local | ExecutionMode::RemoteKernel { .. } => {
             // Write notebook to file
             write_notebook_atomic(&file_path, &notebook).context("Failed to write notebook")?;
         }

src/commands/read.rs

@@ -63,7 +63,10 @@ pub fn execute(args: ReadArgs) -> Result<()> {
                 &server_path,
             ))?
         }
-        crate::execution::types::ExecutionMode::Local => notebook::read_notebook(&file_path)?,
+        crate::execution::types::ExecutionMode::Local
+        | crate::execution::types::ExecutionMode::RemoteKernel { .. } => {
+            notebook::read_notebook(&file_path)?
+        }
     };
 
     // Determine format: markdown (default) or JSON

src/commands/search.rs

@@ -123,7 +123,8 @@ pub fn execute(args: SearchArgs) -> Result<()> {
                 &server_path,
             ))?
         }
-        crate::execution::types::ExecutionMode::Local => {
+        crate::execution::types::ExecutionMode::Local
+        | crate::execution::types::ExecutionMode::RemoteKernel { .. } => {
             // Phase 1: Text pre-filter - quick scan of raw file
             let file_content = fs::read_to_string(&file_path)?;
             if !re.is_match(&file_content) {
@@ -201,7 +202,10 @@ fn execute_with_errors(args: &SearchArgs) -> Result<()> {
                 &server_path,
             ))?
         }
-        crate::execution::types::ExecutionMode::Local => notebook::read_notebook(&file_path)?,
+        crate::execution::types::ExecutionMode::Local
+        | crate::execution::types::ExecutionMode::RemoteKernel { .. } => {
+            notebook::read_notebook(&file_path)?
+        }
     };
     let mut results = Vec::new();
 

src/execution/mod.rs

@@ -6,6 +6,7 @@
 
 pub mod local;
 pub mod remote;
+pub mod remote_kernel;
 pub mod types;
 
 use anyhow::Result;
@@ -48,5 +49,17 @@ pub fn create_backend(config: ExecutionConfig) -> Result<Box<dyn ExecutionBacken
         ExecutionMode::Remote { server_url, token } => Ok(Box::new(remote::RemoteExecutor::new(
             config, server_url, token,
         )?)),
+        ExecutionMode::RemoteKernel {
+            gateway_url,
+            token,
+            kernel_id,
+            auth_scheme,
+        } => Ok(Box::new(remote_kernel::RemoteKernelExecutor::new(
+            config,
+            gateway_url,
+            token,
+            kernel_id,
+            auth_scheme,
+        )?)),
     }
 }

src/execution/remote/websocket.rs

@@ -32,6 +32,38 @@ impl KernelWebSocket {
         Ok(Self { write, read })
     }
 
+    /// Connect to a kernel via WebSocket with an `Authorization` header.
+    ///
+    /// Used by the kernel-gateway path, which authenticates the WS upgrade
+    /// with e.g. `Authorization: token <gateway-token>`. `auth_value` is the
+    /// full header value (e.g. `"token notebooks"`).
+    pub async fn connect_with_auth(ws_url: &str, auth_value: &str) -> Result<Self> {
+        use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+
+        let mut req = ws_url
+            .into_client_request()
+            .context("Failed to build kernel WebSocket request")?;
+        req.headers_mut().insert(
+            "Authorization",
+            auth_value
+                .parse()
+                .context("Invalid Authorization header value")?,
+        );
+        req.headers_mut().insert(
+            "Sec-WebSocket-Protocol",
+            "v1.kernel.websocket.jupyter.org"
+                .parse()
+                .context("Invalid Sec-WebSocket-Protocol value")?,
+        );
+
+        let (ws_stream, _) = connect_async(req)
+            .await
+            .context("Failed to connect to kernel WebSocket")?;
+
+        let (write, read) = ws_stream.split();
+        Ok(Self { write, read })
+    }
+
     /// Parse Jupyter's binary message format
     fn parse_binary_message(data: &[u8]) -> Option<JupyterMessage> {
         // Read number of buffers (first 8 bytes, little-endian)