From ab42fdb58be95763eb762a6d169bde4eb55e26a5 Mon Sep 17 00:00:00 2001
From: Marc Scholten <marc@digitallyinduced.com>
Date: Sun, 19 Jan 2025 10:58:55 -0800
Subject: [PATCH] allow passing custom initialState to WSApp

---
 IHP/ControllerSupport.hs | 10 +++++-----
 IHP/RouterSupport.hs     |  5 +++--
 IHP/WebSocket.hs         |  6 +++---
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/IHP/ControllerSupport.hs b/IHP/ControllerSupport.hs
index 8108c4c57..09b3a3756 100644
--- a/IHP/ControllerSupport.hs
+++ b/IHP/ControllerSupport.hs
@@ -139,8 +139,8 @@ prepareRLSIfNeeded modelContext = do
         Nothing -> pure modelContext
 
 {-# INLINE startWebSocketApp #-}
-startWebSocketApp :: forall webSocketApp application. (?applicationContext :: ApplicationContext, ?context :: RequestContext, InitControllerContext application, ?application :: application, Typeable application, WebSockets.WSApp webSocketApp) => IO ResponseReceived -> Network.Wai.Application
-startWebSocketApp onHTTP request respond = do
+startWebSocketApp :: forall webSocketApp application. (?applicationContext :: ApplicationContext, ?context :: RequestContext, InitControllerContext application, ?application :: application, Typeable application, WebSockets.WSApp webSocketApp) => webSocketApp -> IO ResponseReceived -> Network.Wai.Application
+startWebSocketApp initialState onHTTP request respond = do
     let ?modelContext = ?applicationContext.modelContext
     requestContext <- createRequestContext ?applicationContext request respond
     let ?requestContext = requestContext
@@ -156,7 +156,7 @@ startWebSocketApp onHTTP request respond = do
             try (initContext @application) >>= \case
                 Left (exception :: SomeException) -> putStrLn $ "Unexpected exception in initContext, " <> tshow exception
                 Right context -> do
-                    WebSockets.startWSApp @webSocketApp connection
+                    WebSockets.startWSApp initialState connection
 
     let connectionOptions = WebSockets.connectionOptions @webSocketApp
 
@@ -166,8 +166,8 @@ startWebSocketApp onHTTP request respond = do
             Just response -> respond response
             Nothing -> onHTTP
 {-# INLINE startWebSocketAppAndFailOnHTTP #-}
-startWebSocketAppAndFailOnHTTP :: forall webSocketApp application. (?applicationContext :: ApplicationContext, ?context :: RequestContext, InitControllerContext application, ?application :: application, Typeable application, WebSockets.WSApp webSocketApp) => Network.Wai.Application
-startWebSocketAppAndFailOnHTTP = startWebSocketApp @webSocketApp @application (respond $ responseLBS HTTP.status400 [(hContentType, "text/plain")] "This endpoint is only available via a WebSocket")
+startWebSocketAppAndFailOnHTTP :: forall webSocketApp application. (?applicationContext :: ApplicationContext, ?context :: RequestContext, InitControllerContext application, ?application :: application, Typeable application, WebSockets.WSApp webSocketApp) => webSocketApp -> Network.Wai.Application
+startWebSocketAppAndFailOnHTTP initialState = startWebSocketApp @webSocketApp @application initialState (respond $ responseLBS HTTP.status400 [(hContentType, "text/plain")] "This endpoint is only available via a WebSocket")
     where
         respond = ?context.respond
 
diff --git a/IHP/RouterSupport.hs b/IHP/RouterSupport.hs
index 863d8a7a6..69f8e632d 100644
--- a/IHP/RouterSupport.hs
+++ b/IHP/RouterSupport.hs
@@ -784,7 +784,7 @@ webSocketAppWithCustomPath :: forall webSocketApp application.
 webSocketAppWithCustomPath path = do
         Attoparsec.char '/'
         string path
-        pure (startWebSocketAppAndFailOnHTTP @webSocketApp)
+        pure (startWebSocketAppAndFailOnHTTP (WS.initialState @webSocketApp))
 {-# INLINABLE webSocketAppWithCustomPath #-}
 
 webSocketAppWithCustomPathAndHTTPFallback :: forall webSocketApp application.
@@ -800,7 +800,8 @@ webSocketAppWithCustomPathAndHTTPFallback :: forall webSocketApp application.
 webSocketAppWithCustomPathAndHTTPFallback path = do
         Attoparsec.char '/'
         string path
-        pure (startWebSocketApp @webSocketApp (runActionWithNewContext (WS.initialState @webSocketApp)))
+        let action = WS.initialState @webSocketApp
+        pure (startWebSocketApp action (runActionWithNewContext action))
 {-# INLINABLE webSocketAppWithCustomPathAndHTTPFallback #-}
 
 
diff --git a/IHP/WebSocket.hs b/IHP/WebSocket.hs
index 37789a6d8..6075dca16 100644
--- a/IHP/WebSocket.hs
+++ b/IHP/WebSocket.hs
@@ -60,9 +60,9 @@ class WSApp state where
     connectionOptions :: WebSocket.ConnectionOptions
     connectionOptions = WebSocket.defaultConnectionOptions
 
-startWSApp :: forall state. (WSApp state, ?applicationContext :: ApplicationContext, ?requestContext :: RequestContext, ?context :: ControllerContext, ?modelContext :: ModelContext) => Websocket.Connection -> IO ()
-startWSApp connection = do
-    state <- newIORef (initialState @state)
+startWSApp :: forall state. (WSApp state, ?applicationContext :: ApplicationContext, ?requestContext :: RequestContext, ?context :: ControllerContext, ?modelContext :: ModelContext) => state -> Websocket.Connection -> IO ()
+startWSApp initialState connection = do
+    state <- newIORef initialState
     let ?state = state
 
     result <- Exception.try ((withPingPong defaultPingPongOptions connection (\connection -> let ?connection = connection in run @state)) `Exception.finally` (let ?connection = connection in onClose @state))