From c6184363d8586258f68cfda83717e50e426c8f6d Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Thu, 19 Feb 2026 12:58:07 +0100
Subject: [PATCH] Session server

---
 README.md          | 14 +++++---------
 examples/server.rs | 29 ++++++++---------------------
 src/server.rs      | 42 +++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 54 insertions(+), 31 deletions(-)

diff --git a/README.md b/README.md
index 6ec622d..7a7535a 100644
--- a/README.md
+++ b/README.md
@@ -66,16 +66,12 @@ impl Method for Data {
 
 let server = SessionServer::bind("127.0.0.1:8080").await?;
 
-loop {
-    let session = server.accept().await;
+server
+    .session_loop(async |session, addr| {
+        // This will run on every new client
 
-    session
-        .on::<Data, _>(async |_, req| {
-            // Response is required
-            Ok("Hello from server".to_string())
-        })
-        .await;
-}
+        Ok(())
+    }).await;
 ```
 
 ## Protocol
diff --git a/examples/server.rs b/examples/server.rs
index 0d6fb8e..77583b1 100644
--- a/examples/server.rs
+++ b/examples/server.rs
@@ -1,7 +1,5 @@
 use serde::{Deserialize, Serialize};
-use tokio::net::TcpListener;
-
-use session_rs::{Method, session::Session, ws::WebSocket};
+use session_rs::{Method, server::SessionServer};
 
 #[derive(Debug, Serialize, Deserialize)]
 struct Data;
@@ -15,23 +13,10 @@ impl Method for Data {
 
 #[tokio::main(flavor = "current_thread")]
 async fn main() -> session_rs::Result<()> {
-    let listener = TcpListener::bind("127.0.0.1:8080").await?;
-    println!("Server listening on ws://127.0.0.1:8080");
-
-    loop {
-        let (stream, addr) = listener.accept().await?;
-        println!("New connection: {}", addr);
-
-        tokio::spawn(async move {
-            // Wrap session in Arc so tasks can share it
-            let session = Session::from_ws(
-                WebSocket::handshake(stream)
-                    .await
-                    .expect("Failed to initialize websocket"),
-            );
-
-            session.start_receiver();
+    let server = SessionServer::bind("127.0.0.1:8080").await?;
 
+    server
+        .session_loop(async |session, _| {
             session
                 .on::<Data, _>(async |_, req| {
                     println!("Msg from client: {req}");
@@ -43,6 +28,8 @@ async fn main() -> session_rs::Result<()> {
                     Ok("Hello from server".to_string())
                 })
                 .await;
-        });
-    }
+
+            Ok(())
+        })
+        .await
 }
diff --git a/src/server.rs b/src/server.rs
index 7754b94..4b4db24 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1 +1,41 @@
-pub struct SessionServer();
+use std::{net::SocketAddr, sync::Arc};
+
+use tokio::net::TcpListener;
+
+use crate::{session::Session, ws::WebSocket};
+
+pub struct SessionServer {
+    listener: TcpListener,
+}
+
+impl SessionServer {
+    pub async fn bind(addr: &str) -> crate::Result<Self> {
+        Ok(Self {
+            listener: TcpListener::bind(addr).await?,
+        })
+    }
+
+    pub async fn accept(&self) -> crate::Result<(Session, SocketAddr)> {
+        let (stream, addr) = self.listener.accept().await?;
+
+        let ws = WebSocket::handshake(stream).await?;
+
+        Ok((Session::from_ws(ws), addr))
+    }
+
+    pub async fn session_loop<Fut: Future<Output = crate::Result<()>> + Send + 'static>(
+        &self,
+        on_conn: impl Fn(Session, SocketAddr) -> Fut + 'static,
+    ) -> crate::Result<()> {
+        let conn_handler = Arc::new(on_conn);
+
+        loop {
+            let (session, addr) = self.accept().await?;
+            let conn_handler = conn_handler.clone();
+
+            session.start_receiver();
+
+            tokio::spawn(conn_handler(session, addr));
+        }
+    }
+}