Merge pull request #575 from posit-dev/upkeep/shell-types

Streamline handling of Shell replies

.vscode/launch.json

@@ -51,6 +51,24 @@
             },
             "args": [],
             "cwd": "${workspaceFolder}"
+        },
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Integration tests: Amalthea",
+            "cargo": {
+                "args": [
+                    "test",
+                    "--no-run",
+                    "--package=amalthea"
+                ],
+                "filter": {
+                    "name": "client",
+                    "kind": "test"
+                }
+            },
+            "args": [],
+            "cwd": "${workspaceFolder}"
         }
     ]
 }

crates/amalthea/src/error.rs

@@ -8,6 +8,7 @@
 use std::fmt;
 use std::sync::mpsc::SendError;
 
+use crate::wire::exception::Exception;
 use crate::wire::jupyter_message::Message;
 
 /// Type representing all errors that can occur inside the Amalthea implementation.
@@ -45,6 +46,9 @@ pub enum Error {
     InvalidInputRequest(String),
     InvalidConsoleInput(String),
     Anyhow(anyhow::Error),
+    ShellErrorReply(Exception),
+    /// Execute errors also include the execution count
+    ShellErrorExecuteReply(Exception, u32),
 }
 
 impl std::error::Error for Error {}
@@ -201,6 +205,15 @@ impl fmt::Display for Error {
             Error::Anyhow(err) => {
                 write!(f, "{err:?}")
             },
+            Error::ShellErrorReply(error) => {
+                write!(f, "Got an error reply on Shell: {error:?}")
+            },
+            Error::ShellErrorExecuteReply(error, count) => {
+                write!(
+                    f,
+                    "Got an execute error reply on Shell for request {count}: {error:?}"
+                )
+            },
         }
     }
 }

crates/amalthea/src/fixtures/dummy_frontend.rs

@@ -335,6 +335,14 @@ impl DummyFrontend {
         self.recv_iopub_stream(expect, Stream::Stderr)
     }
 
+    pub fn recv_iopub_comm_close(&self) -> String {
+        let msg = self.recv_iopub();
+
+        assert_matches!(msg, Message::CommClose(data) => {
+            data.content.comm_id
+        })
+    }
+
     /// Receive from IOPub Stream
     ///
     /// Stdout and Stderr Stream messages are buffered, so to reliably test against them

crates/amalthea/src/kernel.rs

@@ -61,7 +61,7 @@ pub fn connect(
     name: &str,
     connection_file: ConnectionFile,
     registration_file: Option<RegistrationFile>,
-    shell_handler: Arc<Mutex<dyn ShellHandler>>,
+    shell_handler: Box<dyn ShellHandler>,
     control_handler: Arc<Mutex<dyn ControlHandler>>,
     lsp_handler: Option<Arc<Mutex<dyn ServerHandler>>>,
     dap_handler: Option<Arc<Mutex<dyn ServerHandler>>>,
@@ -104,18 +104,15 @@ pub fn connect(
     )?;
     let shell_port = port_finalize(&shell_socket, connection_file.shell_port)?;
 
-    let shell_clone = shell_handler.clone();
     let iopub_tx_clone = iopub_tx.clone();
-    let lsp_handler_clone = lsp_handler.clone();
-    let dap_handler_clone = dap_handler.clone();
     spawn!(format!("{name}-shell"), move || {
         shell_thread(
             shell_socket,
             iopub_tx_clone,
             comm_manager_tx,
-            shell_clone,
-            lsp_handler_clone,
-            dap_handler_clone,
+            shell_handler,
+            lsp_handler,
+            dap_handler,
         )
     });
 
@@ -330,7 +327,7 @@ fn shell_thread(
     socket: Socket,
     iopub_tx: Sender<IOPubMessage>,
     comm_manager_tx: Sender<CommManagerEvent>,
-    shell_handler: Arc<Mutex<dyn ShellHandler>>,
+    shell_handler: Box<dyn ShellHandler>,
     lsp_handler: Option<Arc<Mutex<dyn ServerHandler>>>,
     dap_handler: Option<Arc<Mutex<dyn ServerHandler>>>,
 ) -> Result<(), Error> {

crates/amalthea/src/language/shell_handler.rs

@@ -11,9 +11,7 @@ use crate::comm::comm_channel::Comm;
 use crate::socket::comm::CommSocket;
 use crate::wire::complete_reply::CompleteReply;
 use crate::wire::complete_request::CompleteRequest;
-use crate::wire::exception::Exception;
 use crate::wire::execute_reply::ExecuteReply;
-use crate::wire::execute_reply_exception::ExecuteReplyException;
 use crate::wire::execute_request::ExecuteRequest;
 use crate::wire::inspect_reply::InspectReply;
 use crate::wire::inspect_request::InspectRequest;
@@ -31,7 +29,7 @@ pub trait ShellHandler: Send {
     async fn handle_info_request(
         &mut self,
         req: &KernelInfoRequest,
-    ) -> Result<KernelInfoReply, Exception>;
+    ) -> crate::Result<KernelInfoReply>;
 
     /// Handles a request to test a fragment of code to see whether it is a
     /// complete expression.
@@ -40,7 +38,7 @@ pub trait ShellHandler: Send {
     async fn handle_is_complete_request(
         &self,
         req: &IsCompleteRequest,
-    ) -> Result<IsCompleteReply, Exception>;
+    ) -> crate::Result<IsCompleteReply>;
 
     /// Handles a request to execute code.
     ///
@@ -52,21 +50,17 @@ pub trait ShellHandler: Send {
         &mut self,
         originator: Originator,
         req: &ExecuteRequest,
-    ) -> Result<ExecuteReply, ExecuteReplyException>;
+    ) -> crate::Result<ExecuteReply>;
 
     /// Handles a request to provide completions for the given code fragment.
     ///
     /// Docs: https://jupyter-client.readthedocs.io/en/stable/messaging.html#completion
-    async fn handle_complete_request(
-        &self,
-        req: &CompleteRequest,
-    ) -> Result<CompleteReply, Exception>;
+    async fn handle_complete_request(&self, req: &CompleteRequest) -> crate::Result<CompleteReply>;
 
     /// Handles a request to inspect a fragment of code.
     ///
     /// Docs: https://jupyter-client.readthedocs.io/en/stable/messaging.html#introspection
-    async fn handle_inspect_request(&self, req: &InspectRequest)
-        -> Result<InspectReply, Exception>;
+    async fn handle_inspect_request(&self, req: &InspectRequest) -> crate::Result<InspectReply>;
 
     /// Handles a request to open a comm.
     ///
@@ -76,5 +70,5 @@ pub trait ShellHandler: Send {
     ///
     /// * `target` - The target name of the comm, such as `positron.variables`
     /// * `comm` - The comm channel to use to communicate with the frontend
-    async fn handle_comm_open(&self, target: Comm, comm: CommSocket) -> Result<bool, Exception>;
+    async fn handle_comm_open(&self, target: Comm, comm: CommSocket) -> crate::Result<bool>;
 }