From 0940731bd19f879ce4d890e0e68b87808718ea56 Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Thu, 19 Feb 2026 04:05:22 +0100
Subject: [PATCH] Request receiver

---
 examples/client.rs | 21 +++++----------------
 src/lib.rs         |  8 ++++++++
 src/session.rs     | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 57 insertions(+), 18 deletions(-)

diff --git a/examples/client.rs b/examples/client.rs
index 4239b5c..3d7564d 100644
--- a/examples/client.rs
+++ b/examples/client.rs
@@ -1,8 +1,8 @@
 use serde::{Deserialize, Serialize};
-use session_rs::{Method, session::Session, ws::Frame};
+use session_rs::{Method, session::Session};
 
 #[derive(Debug, Serialize, Deserialize)]
-struct Data {}
+struct Data;
 
 impl Method for Data {
     const NAME: &'static str = "data";
@@ -14,23 +14,12 @@ impl Method for Data {
 async fn main() -> session_rs::Result<()> {
     let session = Session::connect("127.0.0.1:8080", "/").await?;
 
-    tokio::spawn({
-        let session = session.clone();
-        async move {
-            loop {
-                match session.ws.read().await {
-                    Ok(Frame::Text(text)) => {
-                        println!("Server says: {}", text);
-                    }
-                    Ok(_) => {}
-                    Err(_) => break,
-                }
-            }
-        }
-    });
+    session.start_receiver();
 
     session.request::<Data>(()).await?;
 
+    session.on::<Data>(|i, d| println!("Ok {i} {d:?}")).await;
+
     tokio::time::sleep(tokio::time::Duration::from_millis(1000)).await;
 
     session.close().await?;
diff --git a/src/lib.rs b/src/lib.rs
index 4761d38..9913bc4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -12,6 +12,14 @@ pub trait Method {
     type Response: Serialize + for<'de> Deserialize<'de>;
 }
 
+pub struct GenericMethod;
+
+impl Method for GenericMethod {
+    const NAME: &'static str = "generic_do_not_use";
+    type Request = serde_json::Value;
+    type Response = serde_json::Value;
+}
+
 #[derive(Debug)]
 pub enum Error {
     WebSocket(ws::Error),
diff --git a/src/session.rs b/src/session.rs
index 07980e5..8c739b0 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -1,9 +1,9 @@
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
 
 use serde::{Deserialize, Serialize};
 use tokio::sync::Mutex;
 
-use crate::{Method, ws::WebSocket};
+use crate::{GenericMethod, Method, ws::WebSocket};
 
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "lowercase", tag = "type")]
@@ -27,6 +27,7 @@ pub enum Message<M: Method> {
 pub struct Session {
     pub ws: WebSocket,
     id: Arc<Mutex<u32>>,
+    methods: Arc<Mutex<HashMap<String, Box<dyn Fn(u32, serde_json::Value) + Send + Sync>>>>,
 }
 
 impl Session {
@@ -34,6 +35,7 @@ impl Session {
         Self {
             ws: self.ws.clone(),
             id: self.id.clone(),
+            methods: self.methods.clone(),
         }
     }
 }
@@ -43,6 +45,7 @@ impl Session {
         Self {
             ws,
             id: Arc::new(Mutex::new(0)),
+            methods: Arc::new(Mutex::new(HashMap::new())),
         }
     }
 
@@ -51,6 +54,45 @@ impl Session {
     }
 }
 
+impl Session {
+    pub fn start_receiver(&self) {
+        let s = self.clone();
+        tokio::spawn(async move {
+            loop {
+                match s.ws.read().await {
+                    Ok(crate::ws::Frame::Text(text)) => {
+                        let Ok(msg) = serde_json::from_str::<Message<GenericMethod>>(&text) else {
+                            continue;
+                        };
+
+                        match msg {
+                            Message::Request { id, method, data } => {
+                                if let Some(m) = s.methods.lock().await.get(&method) {
+                                    (m)(id, data)
+                                }
+                            }
+                            _ => {}
+                        }
+                    }
+                    Ok(_) => {}
+                    Err(_) => break,
+                }
+            }
+        });
+    }
+
+    pub async fn on<M: Method>(&self, handler: impl Fn(u32, M::Request) + Send + Sync + 'static) {
+        self.methods.lock().await.insert(
+            M::NAME.to_string(),
+            Box::new(move |id, value| {
+                if let Ok(req) = serde_json::from_value(value) {
+                    (handler)(id, req)
+                }
+            }),
+        );
+    }
+}
+
 impl Session {
     pub async fn send<M: Method>(&self, data: &Message<M>) -> crate::Result<()> {
         self.ws