-- | Low-level WebSocket client for Substrate
--
-- The core primitive is:
--
-- @
-- substrateRpc :: SubstrateConnection -> Value -> Stream (Of PlexusStreamItem) IO ()
-- @
--
-- This establishes a subscription and yields stream items until completion.
module Substrate.Client
( -- * Connection
SubstrateConnection
, connect
, disconnect
-- * Core RPC primitive
, substrateRpc
-- * Configuration
, SubstrateConfig(..)
, defaultConfig
) where
import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.Async (Async, async, cancel, wait)
import Control.Concurrent.STM
import Control.Exception (SomeException, catch)
import Control.Monad (forever, void)
import Data.Aeson
import qualified Data.Aeson.KeyMap as KM
import Data.IORef
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import qualified Data.Text as T
import Network.WebSockets (Connection)
import qualified Network.WebSockets as WS
import Streaming
import qualified Streaming.Prelude as Str
import Plexus.Types
-- | Substrate connection configuration
data SubstrateConfig = SubstrateConfig
{ substrateHost :: String
, substratePort :: Int
, substratePath :: String
}
deriving stock (Show, Eq)
-- | Default configuration for local development
defaultConfig :: SubstrateConfig
defaultConfig = SubstrateConfig
{ substrateHost = "127.0.0.1"
, substratePort = 4444
, substratePath = "/"
}
-- | Pending request with its queue, waiting for subscription ID
data PendingRequest = PendingRequest
{ prQueue :: TQueue Value -- ^ Queue to receive notifications
, prResponse :: TMVar RpcResponse -- ^ Response channel
}
-- | Active connection to substrate
data SubstrateConnection = SubstrateConnection
{ scConnection :: Connection
, scNextId :: IORef Int
, scSubscriptions :: TVar (Map SubscriptionId (TQueue Value))
, scPendingReqs :: TVar (Map RequestId PendingRequest)
, scReaderThread :: Async ()
}
-- | Connect to substrate
connect :: SubstrateConfig -> IO SubstrateConnection
connect SubstrateConfig{..} = do
-- Initialize state
nextId <- newIORef 1
subs <- newTVarIO Map.empty
pendingReqs <- newTVarIO Map.empty
connVar <- newEmptyTMVarIO
readerVar <- newEmptyTMVarIO
-- Run WebSocket client in a background thread
void $ forkIO $
WS.runClient substrateHost substratePort substratePath $ \conn -> do
-- Start reader in yet another thread (so we can signal back)
reader <- async $ readerLoop conn subs pendingReqs
-- Signal connection is ready
atomically $ putTMVar connVar conn
atomically $ putTMVar readerVar reader
-- Keep alive until reader exits (ignore cancellation)
void (wait reader) `catch` \(_ :: SomeException) -> pure ()
-- Wait for connection (with small delay to let the fork start)
threadDelay 100000 -- 100ms
conn <- atomically $ takeTMVar connVar
reader <- atomically $ takeTMVar readerVar
pure SubstrateConnection
{ scConnection = conn
, scNextId = nextId
, scSubscriptions = subs
, scPendingReqs = pendingReqs
, scReaderThread = reader
}
-- | Disconnect from substrate
disconnect :: SubstrateConnection -> IO ()
disconnect SubstrateConnection{..} = do
cancel scReaderThread
WS.sendClose scConnection ("bye" :: Text) `catch` \(_ :: SomeException) -> pure ()
-- | Reader loop - dispatches incoming messages to the right handler
readerLoop
:: Connection
-> TVar (Map SubscriptionId (TQueue Value))
-> TVar (Map RequestId PendingRequest)
-> IO ()
readerLoop conn subs pendingReqs = forever $ do
msg <- WS.receiveData conn
case eitherDecode msg of
Left err -> putStrLn $ "Failed to decode message: " <> err
Right val -> dispatch val
where
dispatch :: Value -> IO ()
dispatch val = case val of
Object o
-- Check if it's a subscription notification (has "params.subscription" but no "id")
| Just (Object params) <- KM.lookup "params" o
, Just _ <- KM.lookup "subscription" params
, Nothing <- KM.lookup "id" o
-> handleNotification val
-- Otherwise it's a response (has "id")
| Just _ <- KM.lookup "id" o
-> handleResponse val
_ -> putStrLn $ "Unknown message format: " <> show val
handleNotification :: Value -> IO ()
handleNotification val =
case fromJSON val of
Success (SubscriptionNotification _ _ params) -> do
let subId = subParamsSubscription params
result = subParamsResult params
mQueue <- Map.lookup subId <$> readTVarIO subs
case mQueue of
Just queue -> atomically $ writeTQueue queue result
Nothing -> putStrLn $ "Unknown subscription: " <> show subId
Error err -> putStrLn $ "Failed to parse notification: " <> err
handleResponse :: Value -> IO ()
handleResponse val =
case fromJSON val of
Success resp -> do
let rid = rpcRespId resp
mPending <- Map.lookup rid <$> readTVarIO pendingReqs
case mPending of
Just (PendingRequest queue respVar) -> do
-- For successful responses, register the subscription queue BEFORE
-- signaling completion. This prevents race with notifications.
case resp of
RpcSuccess _ result ->
case fromJSON result of
Success subId -> atomically $ do
modifyTVar' subs $ Map.insert subId queue
putTMVar respVar resp
Error _ ->
-- Can't parse subId, just signal response
atomically $ putTMVar respVar resp
RpcError{} ->
atomically $ putTMVar respVar resp
Nothing -> putStrLn $ "Unknown request id: " <> show rid
Error err -> putStrLn $ "Failed to parse response: " <> err
-- | The core RPC primitive
--
-- @substrateRpc conn method params@ sends a subscription request and returns a stream
-- of 'PlexusStreamItem' values. The stream completes when a 'StreamDone' or
-- 'StreamError' item is received.
--
-- Example:
--
-- @
-- import qualified Streaming.Prelude as S
--
-- main = do
-- conn <- connect defaultConfig
-- S.print $ substrateRpc conn "bash_execute" (toJSON ["echo hello"])
-- @
substrateRpc
:: SubstrateConnection
-> Text -- ^ Method name (e.g., "bash_execute")
-> Value -- ^ Parameters
-> Stream (Of PlexusStreamItem) IO ()
substrateRpc conn method params = do
-- Get next request ID
rid <- liftIO $ atomicModifyIORef' (scNextId conn) $ \n -> (n + 1, RequestId n)
-- Create queue for this subscription and response channel
queue <- liftIO newTQueueIO
respVar <- liftIO newEmptyTMVarIO
-- Register pending request (includes queue so reader can register subscription)
liftIO $ atomically $ modifyTVar' (scPendingReqs conn) $
Map.insert rid (PendingRequest queue respVar)
-- Send subscription request
let req = mkSubscribeRequest rid method params
liftIO $ WS.sendTextData (scConnection conn) (encode req)
-- Wait for subscription confirmation
resp <- liftIO $ atomically $ takeTMVar respVar
-- Clean up pending request
liftIO $ atomically $ modifyTVar' (scPendingReqs conn) $ Map.delete rid
case resp of
RpcError _ err -> do
-- Emit error as StreamError so caller can handle it properly
Str.yield $ StreamError
{ itemPlexusHash = "" -- No hash available for subscription errors
, itemProvenance = Provenance ["substrate"]
, itemError = T.pack $ "Subscription error: " <> show err
, itemRecoverable = False
}
RpcSuccess _ result -> do
-- The reader loop already registered the queue in scSubscriptions
-- when it processed the response, so we just need the subId for cleanup
case fromJSON result of
Error err -> do
liftIO $ putStrLn $ "Failed to parse subscription id: " <> err
pure ()
Success subId -> do
-- Stream items until done, cleanup when finished
streamItems queue subId <* liftIO (cleanup subId)
where
streamItems :: TQueue Value -> SubscriptionId -> Stream (Of PlexusStreamItem) IO ()
streamItems queue subId = do
val <- liftIO $ atomically $ readTQueue queue
case fromJSON val of
Error err -> do
liftIO $ putStrLn $ "Failed to parse stream item: " <> err
streamItems queue subId
Success item -> do
Str.yield item
case item of
StreamDone{} -> pure () -- End of stream
StreamError{} -> pure () -- Error terminates stream
_ -> streamItems queue subId
cleanup :: SubscriptionId -> IO ()
cleanup subId = atomically $ modifyTVar' (scSubscriptions conn) $
Map.delete subId