Improve the Notebook Editor

ui/src/ide/editor/notebook/NbButtons.tsx

@@ -25,7 +25,10 @@ function NbButtons(props){
             ))}
         </select>
         <div className='ms-auto'>{props.kernelName}</div>
-        <div className="kStatus">{props.kernelStatus}</div>
+        <div className="kStatus">
+          <span className={`kernelStatus ks-${props.kernelStatus}`}></span>
+          <button className="reconnectButton" onClick={props.startWebSocket}><img src='./images/editor/reconnect-icon.svg' title='Reconnect Kernel'></img></button>
+        </div>
       </div>
     )
   }

ui/src/ide/editor/notebook/NotebookEditor.scss

@@ -125,3 +125,46 @@
     border-radius: 2px;
 }
 
+.kernelStatus{
+    display: inline-block;
+    height: 22px;
+    width: 22px;
+    border-radius: 11px;
+    border: 1px solid #1f1f1f;
+    -webkit-box-shadow: 3px 3px 10px 3px #dddddd;
+    -moz-box-shadow: 3px 3px 10px 3px #dddddd;
+    box-shadow: 3px 3px 10px 3px #dddddd;
+    margin-right: 20px;
+}
+
+
+.ks-idle{
+    background-color: rgb(19, 149, 19);
+}
+
+.ks-busy{
+    background-color: rgb(251, 54, 54);
+}
+
+.ks-connected{
+    background-color:  rgb(19, 149, 19);
+}
+
+.reconnectButton{
+    border: none;
+    width: 22px;
+    height: 22px;
+    -webkit-box-shadow: 3px 3px 10px 3px #dddddd;
+    -moz-box-shadow: 3px 3px 10px 3px #dddddd;
+    box-shadow: 3px 3px 10px 3px #dddddd;
+    border-radius: 10px;
+    padding: 0;
+    padding-inline: 0;
+    padding-block: 0;
+}
+
+.reconnectButton img{
+    width: 22px;
+    height: 22px;
+    margin-top: -10px;
+}

ui/src/ide/editor/notebook/NotebookEditor.tsx

@@ -489,6 +489,7 @@ export default function NotebookEditor(props) {
           kernelName={kernelName}
           kernelStatus={kernelStatus}
           changeCellType={changeCellType}
+          startWebSocket={startWebSocket}
         />
         {debugMode && (
           <div>

websocket/channels.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -38,6 +39,7 @@ type KernelWebSocketConnection struct {
 	IOPubWindowByteQueue []interface{}
 	KernelInfoChannel    zmq4.Socket
 	Subprotocol          string
+	mu                   sync.Mutex
 }
 
 func (kwsConn *KernelWebSocketConnection) getAllowedMessageTypes() []string {
@@ -76,11 +78,11 @@ func (kwsConn *KernelWebSocketConnection) writeMessage() { //msg interface{}, bi
 			if socket.Socket == nil {
 				continue
 			}
-			log.Debug().Msgf("Polling .....")
 			switch s := socket.Socket; s {
 
 			case iopub_channel:
 				msg, _ := s.Recv(0)
+				log.Info().Msgf("Received from IoPub socket: %s\n", msg)
 				if IsJSON([]byte(msg)) {
 					var msgData map[string]interface{}
 					err := json.Unmarshal([]byte(msg), &msgData)
@@ -93,6 +95,17 @@ func (kwsConn *KernelWebSocketConnection) writeMessage() { //msg interface{}, bi
 						jsonResponse[key] = value
 					}
 
+					if _, ok := jsonResponse["execution_state"]; ok {
+						jsonResponse["channel"] = "iopub"
+						jsonData, err := json.Marshal(jsonResponse)
+						if err != nil {
+							log.Info().Msgf("Error marshaling JSON: %s", err)
+							continue
+						}
+						kwsConn.Send <- []byte(jsonData)
+						jsonResponse = make(map[string]interface{})
+					}
+
 				} else {
 					if len(jsonResponse) > 0 {
 						jsonResponse["channel"] = "iopub"

websocket/kernel_websocket_handler.go

@@ -110,10 +110,12 @@ func (kwsConn *KernelWebSocketConnection) writeMessages() {
 			kwsConn.Conn.WriteMessage(websocket.CloseMessage, []byte{})
 			return
 		}
+		kwsConn.mu.Lock()
 		if err := kwsConn.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
 			log.Info().Msgf("Error writing message: %s", err)
 			return
 		}
+		kwsConn.mu.Unlock()
 	}
 }