Working input behavior for Jupyter Notebooks

cli/js/40_jupyter.js

@@ -337,7 +337,14 @@ async function formatInner(obj, raw) {
 internals.jupyter = { formatInner };
 
 function enableJupyter() {
-  const { op_jupyter_broadcast } = core.ops;
+  const { op_jupyter_broadcast, op_jupyter_input } = core.ops;
+
+  async function input(
+    prompt,
+    password,
+  ) {
+    return await op_jupyter_input(prompt, password, []);
+  }
 
   async function broadcast(
     msgType,
@@ -412,6 +419,24 @@ function enableJupyter() {
     return;
   }
 
+  // Override confirm and prompt here
+  async function confirm(message = "Confirm") {
+    const answer = await input(`${message} [y/N] `, false)
+    return answer === "Y" || answer === "y";
+  }
+
+  async function prompt(message = "Prompt", defaultValue = "", { password = false } = {}) {
+    const answer = await input(`${message} [default=${defaultValue}] `, password)
+
+    if(answer === "") {
+      return defaultValue;
+    }
+
+    return answer;
+  }
+
+  globalThis.confirm = confirm;
+  globalThis.prompt = prompt;
   globalThis.Deno.jupyter = {
     broadcast,
     display,

cli/ops/jupyter.rs

@@ -10,13 +10,15 @@ use crate::tools::jupyter::server::StdioMsg;
 use deno_core::error::AnyError;
 use deno_core::op2;
 use deno_core::serde_json;
+use deno_core::serde_json::json;
 use deno_core::OpState;
 use tokio::sync::mpsc;
 use tokio::sync::Mutex;
 
 deno_core::extension!(deno_jupyter,
   ops = [
     op_jupyter_broadcast,
+    op_jupyter_input,
   ],
   options = {
     sender: mpsc::UnboundedSender<StdioMsg>,
@@ -30,6 +32,63 @@ deno_core::extension!(deno_jupyter,
   },
 );
 
+/*
+     * https://jupyter-client.readthedocs.io/en/latest/messaging.html#messages-on-the-stdin-router-dealer-channel
+     * The stdin socket of the client is required to have the
+     *  same zmq IDENTITY as the client’s shell socket.
+     *  Because of this, the input_request must be sent with the same IDENTITY
+     *  routing prefix as the execute_reply in order for the frontend to receive the message.
+     * """
+     *
+     * What python does is overrides the prompt and confirm type functions with their own version
+ */
+#[op2(async)]
+#[string]
+pub async fn op_jupyter_input(
+  state: Rc<RefCell<OpState>>,
+  #[string] prompt: String,
+  #[serde] is_password: serde_json::Value,
+  #[serde] buffers: Vec<deno_core::JsBuffer>,
+) -> Result<Option<String>, AnyError> {
+  let (_iopub_socket, last_execution_request, stdin_socket) = {
+    let s = state.borrow();
+
+    (
+      s.borrow::<Arc<Mutex<Connection<zeromq::PubSocket>>>>()
+        .clone(),
+      s.borrow::<Rc<RefCell<Option<JupyterMessage>>>>().clone(),
+      s.borrow::<Arc<Mutex<Connection<zeromq::RouterSocket>>>>()
+        .clone(),
+    )
+  };
+
+  log::info!("Handling input request");
+  let mut stdin = stdin_socket.lock().await;
+
+  let maybe_last_request = last_execution_request.borrow().clone();
+  if let Some(last_request) = maybe_last_request {
+    if !last_request.allow_stdin() {
+      return Ok(None);
+    }
+
+    last_request
+      .new_message_with_identities("input_request")
+      .with_content(json!({
+        "prompt": prompt,
+        "password": is_password,
+      }))
+      .with_buffers(buffers.into_iter().map(|b| b.to_vec().into()).collect())
+      .send(&mut *stdin)
+      .await?;
+
+    let response = JupyterMessage::read(&mut *stdin).await?;
+
+    return Ok(Some(response.value().to_string()));
+  }
+
+  Ok(None)
+}
+
 #[op2(async)]
 pub async fn op_jupyter_broadcast(
   state: Rc<RefCell<OpState>>,

cli/tools/jupyter/server.rs

@@ -38,7 +38,7 @@ pub struct JupyterServer {
   // This is Arc<Mutex<>>, so we don't hold RefCell borrows across await
   // points.
   iopub_socket: Arc<Mutex<Connection<zeromq::PubSocket>>>,
-  stdin_socket: Arc<Mutex<Connection<zeromq::RouterSocket>>>,
+  _stdin_socket: Arc<Mutex<Connection<zeromq::RouterSocket>>>,
   repl_session: repl::ReplSession,
 }
 
@@ -68,6 +68,7 @@ impl JupyterServer {
       let mut op_state = op_state_rc.borrow_mut();
       op_state.put(iopub_socket.clone());
       op_state.put(last_execution_request.clone());
+      op_state.put(stdin_socket.clone());
     }
 
     let cancel_handle = CancelHandle::new_rc();
@@ -76,7 +77,7 @@ impl JupyterServer {
     let mut server = Self {
       execution_count: 0,
       iopub_socket: iopub_socket.clone(),
-      stdin_socket: stdin_socket.clone(),
+      _stdin_socket: stdin_socket.clone(),
       last_execution_request: last_execution_request.clone(),
       repl_session,
     };
@@ -191,44 +192,6 @@ impl JupyterServer {
     }
   }
 
-  async fn handle_input(
-    &self,
-    msg: &JupyterMessage,
-    prompt: &str,
-    password: bool,
-  ) -> Option<String> {
-
-    if !msg.allow_stdin() {
-        return None;
-    }
-
-    /*
-     * TODO
-     * - Get it working in frontend
-     * - Maybe this is the reason it's failing with empty string?
-     * """
-     * https://jupyter-client.readthedocs.io/en/latest/messaging.html#messages-on-the-stdin-router-dealer-channel
-     * The stdin socket of the client is required to have the
-     *  same zmq IDENTITY as the client’s shell socket.
-     *  Because of this, the input_request must be sent with the same IDENTITY
-     *  routing prefix as the execute_reply in order for the frontend to receive the message.
-     * """
-     *
-     * What python does is overrides the prompt and confirm type functions with their own version
-     */
-    log::info!("Handling input request");
-    let mut stdin = self.stdin_socket.lock().await;
-
-    let stdin_request = msg
-        .new_message_with_identities("input_request")
-        .with_content(json!({ "prompt": prompt, "password": password }));
-    log::info!("Sending input request {:?}", stdin_request);
-    stdin_request.send(&mut *stdin).await.ok()?;
-
-    let input_response = JupyterMessage::read(&mut *stdin).await.ok()?;
-    Some(input_response.value().to_string())
-  }
-
   async fn handle_shell(
     &mut self,
     mut connection: Connection<zeromq::RouterSocket>,

runtime/js/99_main.js

@@ -591,6 +591,7 @@ const NOT_IMPORTED_OPS = [
 
   // Related to `Deno.jupyter` API
   "op_jupyter_broadcast",
+  "op_jupyter_input",
 
   // Related to `Deno.test()` API
   "op_test_event_step_result_failed",