Repository: sparckles/Robyn (language: python)

Issue:
Websocket close method feature
### Description
Currently, Robyn framework lacks the capability to programmatically close a WebSocket connection from within the handler method. This feature is critical for implementing server-side logic that demands the closure of a WebSocket connection under specific conditions, such as the expiration of a user's authorization period.

Almost all major frameworks got this functionality

### Use Case
A practical example of this feature would be in scenarios where a user's temporary access for WebSocket communication expires. For instance, if a user is only authorized to maintain a WebSocket connection for a certain period, the server should have the ability to terminate the connection from within the WebSocket view/handler once this period elapses. This is crucial for maintaining security and efficient resource management.

### Suggested Implementation
The implementation could involve enhancing the existing WebSocket handler class/methods, allowing them to invoke a close connection method. This method should be able to send a closing handshake message to the client and terminate the connection gracefully.

Example pseudo-code:
```python

@websocket.on("message")
  def message(ws, msg) :
      if not is_valid(msg):
         ws.close()  <----- implement this 

```


Discussion / hints:
As you suggested earlier, Firstly, the WebSocket handler class would need to be enhanced to include a close method. This method would be responsible for sending a closing handshake message to the client and terminating the connection.

Here's a simplified example:

class WebSocketHandler:
    def __init__(self, ws):
        self.ws = ws

    def close(self, code=1000, reason=""):
        # Send a closing handshake message to the client
        self.ws.send_close(code, reason)
        # Close the WebSocket connection
        self.ws.close()

Then, in WebSocket view, we could use this close method to terminate the connection under specific conditions:
As you suggested
Example pseudo-code:

@websocket.on("message")
  def message(ws, msg) :
      if not is_valid(msg):
         ws.close()  <----- implement this 
         
Can you please provide reference to the code to contribute to it?

Fix patch:
diff --git a/docs_src/src/pages/documentation/api_reference/websockets.mdx b/docs_src/src/pages/documentation/api_reference/websockets.mdx
index c111d154a..d048bbe40 100644
--- a/docs_src/src/pages/documentation/api_reference/websockets.mdx
+++ b/docs_src/src/pages/documentation/api_reference/websockets.mdx
@@ -2,7 +2,7 @@ export const description =
   'On this page, we’ll dive into the different conversation endpoints you can use to manage conversations programmatically.'
 
 
-## WebSockets {{ tag: 'GET', label: '/hello_world' }}
+## WebSockets {{ tag: 'WebSockets', label: 'WebSockets' }}
 
 <Row>
 <Col>
@@ -214,6 +214,43 @@ To handle real-time communication, Batman learned how to work with WebSockets. H
   </Col>
 </Row>
 
+
+
+<Row>
+  <Col>
+    To programmatically close a WebSocket connection from the server side, Batman learned to use the `close()` method:
+    The `close()` method does the following:
+1. Sends a close message to the client.
+2. Removes the client from the WebSocket registry.
+3. Closes the WebSocket connection.
+
+This method is useful for scenarios where you need to programmatically end a WebSocket connection based on certain conditions or events on the server side.
+
+  </Col>
+  <Col>
+    <CodeGroup title="Request" tag="GET" label="/hello_world">
+
+      ```python {{ title: 'untyped' }}
+      @websocket.on("message")
+      def message(ws, msg):
+          if msg == "disconnect":
+              ws.close()
+              return "Closing connection"
+          return "Message received"
+      ```
+
+      ```python {{title: 'typed'}}
+      @websocket.on("message")
+      def message(ws: WebSocketConnector, msg: str) -> str:
+          if msg == "disconnect":
+              ws.close()
+              return "Closing connection"
+          return "Message received"
+      ```
+    </CodeGroup>
+  </Col>
+</Row>
+
 ---
 
 ## What's next?
diff --git a/robyn/robyn.pyi b/robyn/robyn.pyi
index 595f22536..f354f99bd 100644
--- a/robyn/robyn.pyi
+++ b/robyn/robyn.pyi
@@ -236,6 +236,10 @@ class Headers:
         pass
 
     def is_empty(self) -> bool:
+        """
+        Returns:
+            True if the headers are empty, False otherwise
+        """
         pass
 
 @dataclass
@@ -295,15 +299,20 @@ class Response:
 
     def set_cookie(self, key: str, value: str) -> None:
         """
+        Sets the cookie in the response.
 
-        The function to set a cookie (onto headers).
-
-        @param key: the cookie name
-        @param value: the actual cookie value
+        Args:
+            key (str): The key of the cookie
+            value (str): The value of the cookie
         """
         pass
 
 class Server:
+    """
+    The Server object used to create a Robyn server.
+
+    This object is used to create a Robyn server and add routes, middlewares, etc.
+    """
     def __init__(self) -> None:
         pass
     def add_directory(
@@ -369,10 +378,41 @@ class WebSocketConnector:
     query_params: QueryParams
 
     async def async_broadcast(self, message: str) -> None:
+        """
+        Broadcasts a message to all clients.
+
+        Args:
+            message (str): The message to broadcast
+        """
         pass
     async def async_send_to(self, sender_id: str, message: str) -> None:
+        """
+        Sends a message to a specific client.
+
+        Args:
+            sender_id (str): The id of the sender
+            message (str): The message to send
+        """
         pass
     def sync_broadcast(self, message: str) -> None:
+        """
+        Broadcasts a message to all clients.
+
+        Args:
+            message (str): The message to broadcast
+        """
         pass
     def sync_send_to(self, sender_id: str, message: str) -> None:
+        """
+        Sends a message to a specific client.
+
+        Args:
+            sender_id (str): The id of the sender
+            message (str): The message to send
+        """
+        pass
+    def close(self) -> None:
+        """
+        Closes the connection.
+        """
         pass
diff --git a/src/websockets/mod.rs b/src/websockets/mod.rs
index 8a3394b12..51f51a47c 100644
--- a/src/websockets/mod.rs
+++ b/src/websockets/mod.rs
@@ -3,7 +3,7 @@ pub mod registry;
 use crate::executors::web_socket_executors::execute_ws_function;
 use crate::types::function_info::FunctionInfo;
 use crate::types::multimap::QueryParams;
-use registry::{SendMessageToAll, SendText};
+use registry::{Close, SendMessageToAll, SendText};
 
 use actix::prelude::*;
 use actix::{Actor, AsyncContext, StreamHandler};
@@ -60,7 +60,12 @@ impl Handler<SendText> for WebSocketConnector {
 
     fn handle(&mut self, msg: SendText, ctx: &mut Self::Context) {
         if self.id == msg.recipient_id {
+            let message = msg.message.clone();
             ctx.text(msg.message);
+            if message == "Connection closed" {
+                // Close the WebSocket connection
+                ctx.stop();
+            }
         }
     }
 }
@@ -171,6 +176,10 @@ impl WebSocketConnector {
         Ok(awaitable.into_py(py))
     }
 
+    pub fn close(&self) {
+        self.registry_addr.do_send(Close { id: self.id });
+    }
+
     #[getter]
     pub fn get_id(&self) -> String {
         self.id.to_string()
diff --git a/src/websockets/registry.rs b/src/websockets/registry.rs
index 8f9ddbe9d..702a50ac1 100644
--- a/src/websockets/registry.rs
+++ b/src/websockets/registry.rs
@@ -70,6 +70,8 @@ impl Handler<SendText> for WebSocketRegistry {
 
         if let Some(client_addr) = self.clients.get(&recipient_id) {
             client_addr.do_send(msg);
+        } else {
+            log::warn!("No client found for id: {}", recipient_id);
         }
     }
 }
@@ -96,3 +98,26 @@ impl Handler<SendMessageToAll> for WebSocketRegistry {
         }
     }
 }
+
+pub struct Close {
+    pub id: Uuid,
+}
+
+impl Message for Close {
+    type Result = ();
+}
+
+impl Handler<Close> for WebSocketRegistry {
+    type Result = ();
+
+    fn handle(&mut self, msg: Close, _ctx: &mut Self::Context) {
+        if let Some(client) = self.clients.remove(&msg.id) {
+            // Send a close message to the client before removing it
+            client.do_send(SendText {
+                recipient_id: msg.id,
+                message: "Connection closed".to_string(),
+                sender_id: msg.id,
+            });
+        }
+    }
+}


Test patch:
diff --git a/integration_tests/base_routes.py b/integration_tests/base_routes.py
index 38082f186..b3f90fe14 100644
--- a/integration_tests/base_routes.py
+++ b/integration_tests/base_routes.py
@@ -73,7 +73,10 @@ async def message(ws: WebSocketConnector, msg: str, global_dependencies) -> str:
         await ws.async_broadcast(ws.query_params.get("one"))
         ws.sync_send_to(websocket_id, ws.query_params.get("two"))
         resp = "*chika* *chika* Slim Shady."
-    websocket_state[websocket_id] = (state + 1) % 3
+    elif state == 3:
+        ws.close()
+
+    websocket_state[websocket_id] = (state + 1) % 4
     return resp
 
 
diff --git a/integration_tests/random_number.rs b/integration_tests/random_number.rs
new file mode 100644
index 000000000..f88662fd9
--- /dev/null
+++ b/integration_tests/random_number.rs
@@ -0,0 +1,18 @@
+// rustimport:pyo3
+
+use pyo3::prelude::*;
+
+#[pyfunction]
+fn say_hello() {
+    println!("Hello from random_number, implemented in Rust!")
+}
+
+// Uncomment the below to implement custom pyo3 binding code. Otherwise, 
+// rustimport will generate it for you for all functions annotated with
+// #[pyfunction] and all structs annotated with #[pyclass].
+//
+//#[pymodule]
+//fn random_number(_py: Python, m: &PyModule) -> PyResult<()> {
+//    m.add_function(wrap_pyfunction!(say_hello, m)?)?;
+//    Ok(())
+//}
diff --git a/integration_tests/test_web_sockets.py b/integration_tests/test_web_sockets.py
index 722590dd1..6b16564e2 100644
--- a/integration_tests/test_web_sockets.py
+++ b/integration_tests/test_web_sockets.py
@@ -20,10 +20,14 @@ def test_web_socket_raw_benchmark(session):
     assert ws.recv() == "Whooo??"
 
     ws.send("My name is?")
-    ws.recv() == "hi"
-    ws.recv() == "hello"
+    assert ws.recv() == "hi"
+    assert ws.recv() == "hello"
     assert ws.recv() == "*chika* *chika* Slim Shady."
 
+    # this will close the connection
+    ws.send("test")
+    assert ws.recv() == "Connection closed"
+
 
 def test_web_socket_json(session):
     """