From ba0aa74d9c239fdffaa8ab0f95f58b72e9aab096 Mon Sep 17 00:00:00 2001
From: Prasun Anand <prasunanand.bitsp@gmail.com>
Date: Sun, 22 Dec 2024 11:54:02 +0530
Subject: [PATCH] Improve the Notebook Editor

---
 ui/public/images/editor/reconnect-icon.svg    |  3 ++
 ui/src/ide/editor/notebook/NbButtons.tsx      |  5 ++-
 .../ide/editor/notebook/NotebookEditor.scss   | 43 +++++++++++++++++++
 ui/src/ide/editor/notebook/NotebookEditor.tsx |  1 +
 websocket/channels.go                         | 15 ++++++-
 websocket/kernel_websocket_handler.go         |  2 +
 6 files changed, 67 insertions(+), 2 deletions(-)
 create mode 100644 ui/public/images/editor/reconnect-icon.svg

diff --git a/ui/public/images/editor/reconnect-icon.svg b/ui/public/images/editor/reconnect-icon.svg
new file mode 100644
index 0000000..bfe7921
--- /dev/null
+++ b/ui/public/images/editor/reconnect-icon.svg
@@ -0,0 +1,3 @@
+
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 122.88 122.88">
+<defs><style>.cls-1{fill-rule:evenodd;}</style></defs><title>sync</title><path class="cls-1" d="M64.89,32.65,59.81,58.5l-5.16-7.77C43.54,55.19,37.3,62.54,36.38,73.86c-9.13-16-3.59-30.25,8-38.63L39.09,27.3l25.8,5.35ZM61.44,0A61.46,61.46,0,1,1,18,18,61.21,61.21,0,0,1,61.44,0ZM97.56,25.32a51.08,51.08,0,1,0,15,36.12,51,51,0,0,0-15-36.12ZM56.64,91.8,61.72,66l5.16,7.77C78,69.26,84.23,61.91,85.15,50.59c9.13,16,3.59,30.25-8,38.63l5.26,7.93L56.64,91.8Z"/></svg>
\ No newline at end of file
diff --git a/ui/src/ide/editor/notebook/NbButtons.tsx b/ui/src/ide/editor/notebook/NbButtons.tsx
index b5c7303..38c1051 100644
--- a/ui/src/ide/editor/notebook/NbButtons.tsx
+++ b/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>
     )
   }
diff --git a/ui/src/ide/editor/notebook/NotebookEditor.scss b/ui/src/ide/editor/notebook/NotebookEditor.scss
index f99e91e..5cc6d75 100644
--- a/ui/src/ide/editor/notebook/NotebookEditor.scss
+++ b/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;
+}
\ No newline at end of file
diff --git a/ui/src/ide/editor/notebook/NotebookEditor.tsx b/ui/src/ide/editor/notebook/NotebookEditor.tsx
index 6716423..8935639 100644
--- a/ui/src/ide/editor/notebook/NotebookEditor.tsx
+++ b/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>
diff --git a/websocket/channels.go b/websocket/channels.go
index e746a67..7cfa219 100644
--- a/websocket/channels.go
+++ b/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"
diff --git a/websocket/kernel_websocket_handler.go b/websocket/kernel_websocket_handler.go
index 372271d..2579669 100644
--- a/websocket/kernel_websocket_handler.go
+++ b/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()
 	}
 }