agent: process SMP messages concurrently between different connections

2026-08-29 01:18:40 +00:00 · 2026-05-23 10:09:01 +00:00
parent f0b7a4be73
commit 04685840fc
8 changed files with 309 additions and 139 deletions
@@ -145,6 +145,7 @@ where

 import Control.Applicative ((<|>))
 import Control.Concurrent.STM (retry)
+import Data.IORef
 import Control.Logger.Simple
 import Control.Monad
 import Control.Monad.Except
@@ -270,19 +271,25 @@ getSMPAgentClient_ clientId cfg initServers@InitialAgentServers {smp, xftp, netC
      liftIO $ checkServers "SMP" smp >> checkServers "XFTP" xftp
      currentTs <- liftIO getCurrentTime
      notices <- liftIO $ withTransaction store (`getClientNotices` presetServers) `catchAll_` pure []
-      c@AgentClient {acThread} <- liftIO . newAgentClient clientId initServers currentTs notices =<< ask
+      env <- ask
+      cRef <- liftIO $ newIORef (error "agent client not initialized")
+      let processMsg t = do
+            c <- readIORef cRef
+            agentOperationBracket c AORcvNetwork waitUntilActive (processSMPTransmissions c t) `runReaderT` env
+              `catchOwn` \e -> atomically $ writeTBQueue (subQ c) ("", "", AEvt SAEConn $ ERR $ CRITICAL True $ "subscriber error: " <> show e)
+      c@AgentClient {acThread} <- liftIO $ newAgentClient clientId initServers currentTs notices processMsg env
+      liftIO $ writeIORef cRef c
      t <- runAgentThreads c `forkFinally` const (liftIO $ disconnectAgentClient c)
      atomically . writeTVar acThread . Just =<< mkWeakThreadId t
      pure c
    checkServers protocol srvs =
      forM_ (M.assocs srvs) $ \(userId, srvs') -> checkUserServers ("getSMPAgentClient " <> protocol <> " " <> tshow userId) srvs'
    runAgentThreads c
-      | backgroundMode = run c "subscriber" $ subscriber c
+      | backgroundMode = forever $ liftIO $ threadDelay maxBound
      | otherwise = do
          restoreServersStats c
          raceAny_
-            [ run c "subscriber" $ subscriber c,
-              run c "runNtfSupervisor" $ runNtfSupervisor c,
+            [ run c "runNtfSupervisor" $ runNtfSupervisor c,
              run c "cleanupManager" $ cleanupManager c,
              run c "logServersStats" $ logServersStats c
            ]
@@ -2982,14 +2989,6 @@ getNextSMPServer :: AgentClient -> UserId -> [SMPServer] -> AM SMPServerWithAuth
 getNextSMPServer c userId = getNextServer c userId storageSrvs
 {-# INLINE getNextSMPServer #-}

-subscriber :: AgentClient -> AM' ()
-subscriber c@AgentClient {msgQ, subQ} = run $ forever $ do
-  t <- atomically $ readTBQueue msgQ
-  agentOperationBracket c AORcvNetwork waitUntilActive $
-    processSMPTransmissions c t
-  where
-    run a = a `catchOwn` \e -> notify $ CRITICAL True $ "Agent subscriber stopped: " <> show e
-    notify err = atomically $ writeTBQueue subQ ("", "", AEvt SAEConn $ ERR err)

 cleanupManager :: AgentClient -> AM' ()
 cleanupManager c@AgentClient {subQ} = do
@@ -338,7 +338,7 @@ data AgentClient = AgentClient
  { acThread :: TVar (Maybe (Weak ThreadId)),
    active :: TVar Bool,
    subQ :: TBQueue ATransmission,
-    msgQ :: TBQueue (ServerTransmissionBatch SMPVersion ErrorType BrokerMsg),
+    processServerMsg :: ServerTransmissionBatch SMPVersion ErrorType BrokerMsg -> IO (),
    smpServers :: TMap UserId (UserServers 'PSMP),
    smpClients :: TMap SMPTransportSession SMPClientVar,
    useClientServices :: TMap UserId Bool,
@@ -505,15 +505,14 @@ data UserNetworkType = UNNone | UNCellular | UNWifi | UNEthernet | UNOther
  deriving (Eq, Show)

 -- | Creates an SMP agent client instance that receives commands and sends responses via 'TBQueue's.
-newAgentClient :: Int -> InitialAgentServers -> UTCTime -> Map (Maybe SMPServer) (Maybe SystemSeconds) -> Env -> IO AgentClient
-newAgentClient clientId InitialAgentServers {smp, ntf, xftp, netCfg, useServices, presetDomains, presetServers} currentTs notices agentEnv = do
+newAgentClient :: Int -> InitialAgentServers -> UTCTime -> Map (Maybe SMPServer) (Maybe SystemSeconds) -> (ServerTransmissionBatch SMPVersion ErrorType BrokerMsg -> IO ()) -> Env -> IO AgentClient
+newAgentClient clientId InitialAgentServers {smp, ntf, xftp, netCfg, useServices, presetDomains, presetServers} currentTs notices processServerMsg agentEnv = do
  let cfg = config agentEnv
      qSize = tbqSize cfg
  proxySessTs <- newTVarIO =<< getCurrentTime
  acThread <- newTVarIO Nothing
  active <- newTVarIO True
  subQ <- newTBQueueIO qSize
-  msgQ <- newTBQueueIO qSize
  smpServers <- newTVarIO $ M.map mkUserServers smp
  smpClients <- TM.emptyIO
  useClientServices <- newTVarIO useServices
@@ -553,7 +552,7 @@ newAgentClient clientId InitialAgentServers {smp, ntf, xftp, netCfg, useServices
      { acThread,
        active,
        subQ,
-        msgQ,
+        processServerMsg,
        smpServers,
        smpClients,
        useClientServices,
@@ -733,7 +732,7 @@ getSMPProxyClient c@AgentClient {active, smpClients, smpProxiedRelays, workerSeq
        Nothing -> Left $ BROKER (B.unpack $ strEncode srv) TIMEOUT

 smpConnectClient :: AgentClient -> NetworkRequestMode -> SMPTransportSession -> TMap SMPServer ProxiedRelayVar -> SMPClientVar -> AM SMPConnectedClient
-smpConnectClient c@AgentClient {smpClients, msgQ, proxySessTs, presetDomains} nm tSess@(userId, srv, _) prs v =
+smpConnectClient c@AgentClient {smpClients, proxySessTs, presetDomains} nm tSess@(userId, srv, _) prs v =
  newProtocolClient c tSess smpClients connectClient v
    `catchAllErrors` \e -> lift (resubscribeSMPSession c tSess) >> throwE e
  where
@@ -746,7 +745,7 @@ smpConnectClient c@AgentClient {smpClients, msgQ, proxySessTs, presetDomains} nm
      env <- ask
      smp <- liftError (protocolClientError SMP $ B.unpack $ strEncode srv) $ do
        ts <- readTVarIO proxySessTs
-        ExceptT $ getProtocolClient g nm tSess cfg' presetDomains (Just msgQ) ts $ smpClientDisconnected c tSess env v' prs
+        ExceptT $ getProtocolClient g nm tSess cfg' presetDomains (Just $ processServerMsg c) ts $ smpClientDisconnected c tSess env v' prs
      atomically $ SS.setSessionId tSess (sessionId $ thParams smp) $ currentSubs c
      updateClientService service smp
      pure SMPConnectedClient {connectedClient = smp, proxiedRelays = prs}
@@ -2835,8 +2834,8 @@ data ClientInfo
  deriving (Show)

 getAgentQueuesInfo :: AgentClient -> IO AgentQueuesInfo
-getAgentQueuesInfo AgentClient {msgQ, subQ, smpClients} = do
-  msgQInfo <- atomically $ getTBQueueInfo msgQ
+getAgentQueuesInfo AgentClient {subQ, smpClients} = do
+  let msgQInfo = TBQueueInfo {qLength = 0, qFull = False}
  subQInfo <- atomically $ getTBQueueInfo subQ
  smpClientsMap <- readTVarIO smpClients
  let smpClientsMap' = M.mapKeys (decodeLatin1 . strEncode) smpClientsMap
@@ -128,6 +128,7 @@ where

 import Control.Applicative ((<|>))
 import Control.Concurrent (ThreadId, forkFinally, forkIO, killThread, mkWeakThreadId)
+import Control.Concurrent.MVar
 import Control.Concurrent.Async
 import Control.Concurrent.STM
 import Control.Exception (Exception, Handler (..), IOException, SomeAsyncException, SomeException)
@@ -199,7 +200,8 @@ data PClient v err msg = PClient
    sentCommands :: TMap CorrId (Request err msg),
    sndQ :: TBQueue (Maybe (Request err msg), ByteString),
    rcvQ :: TBQueue (NonEmpty (Transmission (Either err msg))),
-    msgQ :: Maybe (TBQueue (ServerTransmissionBatch v err msg))
+    processServerMsg :: Maybe (ServerTransmissionBatch v err msg -> IO ()),
+    processLock :: MVar ()
  }

 smpClientStub :: TVar ChaChaDRG -> ByteString -> VersionSMP -> Maybe (THandleAuth 'TClient) -> IO SMPClient
@@ -213,6 +215,7 @@ smpClientStub g sessionId thVersion thAuth = do
  timeoutErrorCount <- newTVarIO 0
  sndQ <- newTBQueueIO 100
  rcvQ <- newTBQueueIO 100
+  processLock <- newMVar ()
  let NetworkConfig {tcpConnectTimeout, tcpTimeout} = defaultNetworkConfig
  return
    ProtocolClient
@@ -244,7 +247,8 @@ smpClientStub g sessionId thVersion thAuth = do
              sentCommands,
              sndQ,
              rcvQ,
-              msgQ = Nothing
+              processServerMsg = Nothing,
+              processLock
            }
      }

@@ -562,10 +566,10 @@ type SMPTransportSession = TransportSession BrokerMsg
 -- | Connects to 'ProtocolServer' using passed client configuration
 -- and queue for messages and notifications.
 --
-- A single queue can be used for multiple 'SMPClient' instances,
+-- A single callback can be used for multiple 'SMPClient' instances,
 -- as 'SMPServerTransmission' includes server information.
-getProtocolClient :: forall v err msg. Protocol v err msg => TVar ChaChaDRG -> NetworkRequestMode -> TransportSession msg -> ProtocolClientConfig v -> [HostName] -> Maybe (TBQueue (ServerTransmissionBatch v err msg)) -> UTCTime -> (ProtocolClient v err msg -> IO ()) -> IO (Either (ProtocolClientError err) (ProtocolClient v err msg))
-getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize, networkConfig, clientALPN, serviceCredentials, serverVRange, agreeSecret, proxyServer, useSNI} presetDomains msgQ proxySessTs disconnected = do
+getProtocolClient :: forall v err msg. Protocol v err msg => TVar ChaChaDRG -> NetworkRequestMode -> TransportSession msg -> ProtocolClientConfig v -> [HostName] -> Maybe (ServerTransmissionBatch v err msg -> IO ()) -> UTCTime -> (ProtocolClient v err msg -> IO ()) -> IO (Either (ProtocolClientError err) (ProtocolClient v err msg))
+getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize, networkConfig, clientALPN, serviceCredentials, serverVRange, agreeSecret, proxyServer, useSNI} presetDomains processServerMsg proxySessTs disconnected = do
  case chooseTransportHost networkConfig (host srv) of
    Right useHost ->
      (getCurrentTime >>= mkProtocolClient useHost >>= runClient useTransport useHost)
@@ -583,6 +587,7 @@ getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qS
      sentCommands <- TM.emptyIO
      sndQ <- newTBQueueIO qSize
      rcvQ <- newTBQueueIO qSize
+      processLock <- newMVar ()
      return
        PClient
          { connected,
@@ -597,7 +602,8 @@ getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qS
            sentCommands,
            sndQ,
            rcvQ,
-            msgQ
+            processServerMsg,
+            processLock
          }

    runClient :: (ServiceName, ATransport 'TClient) -> TransportHost -> PClient v err msg -> IO (Either (ProtocolClientError err) (ProtocolClient v err msg))
@@ -686,8 +692,10 @@ getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qS
    processMsgs :: ProtocolClient v err msg -> NonEmpty (Transmission (Either err msg)) -> IO ()
    processMsgs c ts = do
      ts' <- catMaybes <$> mapM (processMsg c) (L.toList ts)
-      forM_ msgQ $ \q ->
-        mapM_ (atomically . writeTBQueue q . serverTransmission c) (L.nonEmpty ts')
+      forM_ processServerMsg $ \process ->
+        forM_ (L.nonEmpty ts') $ \ts'' ->
+          withMVar (processLock $ client_ c) $ \_ ->
+            process $ serverTransmission c ts''

    processMsg :: ProtocolClient v err msg -> Transmission (Either err msg) -> IO (Maybe (EntityId, ServerTransmission err msg))
    processMsg ProtocolClient {client_ = PClient {sentCommands}} (corrId, entId, respOrErr)
@@ -714,7 +722,7 @@ getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qS
            Just e -> Left $ PCEProtocolError e
            _ -> Right r
        sendMsg :: ServerTransmission err msg -> IO (Maybe (EntityId, ServerTransmission err msg))
-        sendMsg t = case msgQ of
+        sendMsg t = case processServerMsg of
          Just _ -> pure $ Just (entId, t)
          Nothing ->
            Nothing <$ case clientResp of
@@ -872,7 +880,10 @@ processSUBResponse_ c rId = \case
  r' -> pure . Left $ unexpectedResponse r'

 writeSMPMessage :: SMPClient -> RecipientId -> BrokerMsg -> IO ()
-writeSMPMessage c rId msg = atomically $ mapM_ (`writeTBQueue` serverTransmission c [(rId, STEvent (Right msg))]) (msgQ $ client_ c)
+writeSMPMessage c rId msg =
+  forM_ (processServerMsg $ client_ c) $ \process ->
+    withMVar (processLock $ client_ c) $ \_ ->
+      process $ serverTransmission c [(rId, STEvent (Right msg))]

 serverTransmission :: ProtocolClient v err msg -> NonEmpty (RecipientId, ServerTransmission err msg) -> ServerTransmissionBatch v err msg
 serverTransmission ProtocolClient {thParams, client_ = PClient {transportSession}} ts = (transportSession, thParams, ts)
@@ -138,7 +138,7 @@ data SMPClientAgent p = SMPClientAgent
    dbService :: Maybe DBService,
    active :: TVar Bool,
    startedAt :: UTCTime,
-    msgQ :: TBQueue (ServerTransmissionBatch SMPVersion ErrorType BrokerMsg),
+    processMsg :: ServerTransmissionBatch SMPVersion ErrorType BrokerMsg -> IO (),
    agentQ :: TBQueue SMPClientAgentEvent,
    randomDrg :: TVar ChaChaDRG,
    smpClients :: TMap SMPServer SMPClientVar,
@@ -158,11 +158,10 @@ data SMPClientAgent p = SMPClientAgent

 type OwnServer = Bool

-newSMPClientAgent :: SParty p -> SMPClientAgentConfig -> Maybe DBService -> TVar ChaChaDRG -> IO (SMPClientAgent p)
-newSMPClientAgent agentParty agentCfg@SMPClientAgentConfig {msgQSize, agentQSize} dbService randomDrg = do
+newSMPClientAgent :: SParty p -> SMPClientAgentConfig -> (ServerTransmissionBatch SMPVersion ErrorType BrokerMsg -> IO ()) -> Maybe DBService -> TVar ChaChaDRG -> IO (SMPClientAgent p)
+newSMPClientAgent agentParty agentCfg@SMPClientAgentConfig {agentQSize} processMsg dbService randomDrg = do
  active <- newTVarIO True
  startedAt <- getCurrentTime
-  msgQ <- newTBQueueIO msgQSize
  agentQ <- newTBQueueIO agentQSize
  smpClients <- TM.emptyIO
  smpSessions <- TM.emptyIO
@@ -179,7 +178,7 @@ newSMPClientAgent agentParty agentCfg@SMPClientAgentConfig {msgQSize, agentQSize
        dbService,
        active,
        startedAt,
-        msgQ,
+        processMsg,
        agentQ,
        randomDrg,
        smpClients,
@@ -257,7 +256,7 @@ isOwnServer SMPClientAgent {agentCfg} ProtocolServer {host} =

 -- | Run an SMP client for SMPClientVar
 connectClient :: SMPClientAgent p -> SMPServer -> SMPClientVar -> IO (Either SMPClientError SMPClient)
-connectClient ca@SMPClientAgent {agentCfg, dbService, smpClients, smpSessions, msgQ, randomDrg, startedAt} srv v = case dbService of
+connectClient ca@SMPClientAgent {agentCfg, dbService, smpClients, smpSessions, processMsg, randomDrg, startedAt} srv v = case dbService of
  Just dbs -> runExceptT $ do
    creds <- ExceptT $ getCredentials dbs srv
    smp <- ExceptT $ getClient cfg {serviceCredentials = Just creds}
@@ -267,7 +266,7 @@ connectClient ca@SMPClientAgent {agentCfg, dbService, smpClients, smpSessions, m
  Nothing -> getClient cfg
  where
    cfg = smpCfg agentCfg
-    getClient cfg' = getProtocolClient randomDrg NRMBackground (1, srv, Nothing) cfg' [] (Just msgQ) startedAt clientDisconnected
+    getClient cfg' = getProtocolClient randomDrg NRMBackground (1, srv, Nothing) cfg' [] (Just processMsg) startedAt clientDisconnected

    clientDisconnected :: SMPClient -> IO ()
    clientDisconnected smp = do
@@ -54,7 +54,7 @@ import GHC.IORef (atomicSwapIORef)
 import GHC.Stats (getRTSStats)
 import Network.Socket (ServiceName, Socket, socketToHandle)
 import Numeric.Natural (Natural)
-import Simplex.Messaging.Client (ProtocolClientError (..), SMPClientError, ServerTransmission (..))
+import Simplex.Messaging.Client (ProtocolClientError (..), SMPClientError, ServerTransmission (..), ServerTransmissionBatch)
 import Simplex.Messaging.Client.Agent
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Encoding.String
@@ -68,7 +68,7 @@ import Simplex.Messaging.Notifications.Server.Store (NtfSTMStore, TokenNtfMessag
 import Simplex.Messaging.Notifications.Server.Store.Postgres
 import Simplex.Messaging.Notifications.Server.Store.Types
 import Simplex.Messaging.Notifications.Transport
-import Simplex.Messaging.Protocol (EntityId (..), ErrorType (..), NotifierId, Party (..), ProtocolServer (host), SMPServer, ServiceSub (..), SignedTransmission, Transmission, pattern NoEntity, pattern SMPServer, encodeTransmission, tGetServer, tPut)
+import Simplex.Messaging.Protocol (BrokerMsg, EntityId (..), ErrorType (..), NotifierId, Party (..), ProtocolServer (host), SMPServer, ServiceSub (..), SignedTransmission, Transmission, pattern NoEntity, pattern SMPServer, encodeTransmission, tGetServer, tPut)
 import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.Server
 import Simplex.Messaging.Server.Control (CPClientRole (..))
@@ -77,7 +77,7 @@ import Simplex.Messaging.Server.Stats (PeriodStats (..), PeriodStatCounts (..),
 import Simplex.Messaging.Session
 import Simplex.Messaging.SystemTime
 import Simplex.Messaging.TMap (TMap)
-import Simplex.Messaging.Transport (ASrvTransport, ATransport (..), THandle (..), THandleAuth (..), THandleParams (..), TProxy, Transport (..), TransportPeer (..), defaultSupportedParams)
+import Simplex.Messaging.Transport (ASrvTransport, ATransport (..), SMPVersion, THandle (..), THandleAuth (..), THandleParams (..), TProxy, Transport (..), TransportPeer (..), defaultSupportedParams)
 import Simplex.Messaging.Transport.Buffer (trimCR)
 import Simplex.Messaging.Transport.Server (AddHTTP, runTransportServer, runLocalTCPServer)
 import Simplex.Messaging.Util
@@ -101,7 +101,11 @@ runNtfServer cfg = do
  runNtfServerBlocking started cfg

 runNtfServerBlocking :: TMVar Bool -> NtfServerConfig -> IO ()
-runNtfServerBlocking started cfg = runReaderT (ntfServer cfg started) =<< newNtfServerEnv cfg
+runNtfServerBlocking started cfg = runReaderT (ntfServer cfg started) =<< newNtfServerEnv cfg processMsg
+  where
+    processMsg envRef t = do
+      env <- readIORef envRef
+      receiveSMPMessage env t

 type M a = ReaderT NtfEnv IO a

@@ -525,97 +529,91 @@ subscribeNtfs NtfSubscriber {smpSubscribers, subscriberSeq, smpAgent = ca} st sm
      void $ updateSubStatus st srvId' nId NSPending
      subscribeQueuesNtfs ca smpServer' [sub]

+receiveSMPMessage :: NtfEnv -> ServerTransmissionBatch SMPVersion ErrorType BrokerMsg -> IO ()
+receiveSMPMessage env ((_, srv@(SMPServer (h :| _) _ _), _), THandleParams {sessionId}, ts) =
+  (`runReaderT` env) $ do
+    st <- asks store
+    ps <- asks pushServer
+    stats <- asks serverStats
+    let ca = smpAgent $ subscriber env
+    forM_ ts $ \(ntfId, t) -> case t of
+      STUnexpectedError e -> logError $ "SMP client unexpected error: " <> tshow e -- uncorrelated response, should not happen
+      STResponse {} -> pure () -- it was already reported as timeout error
+      STEvent msgOrErr -> do
+        let smpQueue = SMPQueueNtf srv ntfId
+        case msgOrErr of
+          Right (SMP.NMSG nmsgNonce encNMsgMeta) -> do
+            ntfTs <- liftIO getSystemTime
+            liftIO $ updatePeriodStats (activeSubs stats) ntfId
+            let newNtf = PNMessageData {smpQueue, ntfTs, nmsgNonce, encNMsgMeta}
+                srvHost = safeDecodeUtf8 $ strEncode h
+                isOwn = isOwnServer ca srv
+            liftIO (addTokenLastNtf st newNtf) >>= \case
+              Right (tkn, lastNtfs) -> do
+                pushNotification ps (Just srvHost) isOwn tkn $ PNMessage lastNtfs
+                liftIO $ incNtfStat_ stats ntfReceived
+                when isOwn $ liftIO $ incServerStat srvHost (ntfReceivedOwn stats)
+              Left AUTH -> liftIO $ do
+                incNtfStat_ stats ntfReceivedAuth
+                when isOwn $ incServerStat srvHost (ntfReceivedAuthOwn stats)
+              Left _ -> pure ()
+          Right SMP.END ->
+            whenM (atomically $ activeClientSession' ca sessionId srv) $
+              void $ liftIO $ updateSrvSubStatus st smpQueue NSEnd
+          Right SMP.DELD ->
+            void $ liftIO $ updateSrvSubStatus st smpQueue NSDeleted
+          Right (SMP.ERR e) -> logError $ "SMP server error: " <> tshow e
+          Right _ -> logError "SMP server unexpected response"
+          Left e -> logError $ "SMP client error: " <> tshow e
+
 ntfSubscriber :: NtfSubscriber -> M ()
-ntfSubscriber NtfSubscriber {smpAgent = ca@SMPClientAgent {msgQ, agentQ}} =
-  race_ receiveSMP receiveAgent
+ntfSubscriber NtfSubscriber {smpAgent = ca@SMPClientAgent {agentQ}} = do
+  st <- asks store
+  batchSize <- asks $ subsBatchSize . config
+  liftIO $ forever $
+    atomically (readTBQueue agentQ) >>= \case
+      CAConnected srv serviceId -> do
+        let asService = if isJust serviceId then "as service " else ""
+        logInfo $ "SMP server reconnected " <> asService <> showServer' srv
+      CADisconnected srv nIds -> do
+        updated <- batchUpdateSrvSubStatus st srv Nothing nIds NSInactive
+        logSubStatus srv "disconnected" (L.length nIds) updated
+      CASubscribed srv serviceId nIds -> do
+        updated <- batchUpdateSrvSubStatus st srv serviceId nIds NSActive
+        let asService = if isJust serviceId then " as service" else ""
+        logSubStatus srv ("subscribed" <> asService) (L.length nIds) updated
+      CASubError srv errs -> do
+        forM_ (L.nonEmpty $ mapMaybe (\(nId, err) -> (nId,) <$> queueSubErrorStatus err) $ L.toList errs) $ \subStatuses -> do
+          updated <- batchUpdateSrvSubErrors st srv subStatuses
+          logSubErrors srv subStatuses updated
+          -- TODO [certs rcv] resubscribe queues with statuses NSErr and NSService
+      CAServiceDisconnected srv serviceSub ->
+        logNote $ "SMP server service disconnected " <> showService srv serviceSub
+      CAServiceSubscribed srv serviceSub@(ServiceSub _ n idsHash) (ServiceSub _ n' idsHash')
+        | n /= n' -> logWarn $ msg <> ", confirmed subs: " <> tshow n'
+        | idsHash /= idsHash' -> logWarn $ msg <> ", different IDs hash"
+        | otherwise -> logNote msg
+        where
+          msg = "SMP server service subscribed " <> showService srv serviceSub
+      CAServiceSubError srv serviceSub e ->
+        -- Errors that require re-subscribing queues directly are reported as CAServiceUnavailable.
+        -- See smpSubscribeService in Simplex.Messaging.Client.Agent
+        logError $ "SMP server service subscription error " <> showService srv serviceSub <> ": " <> tshow e
+      CAServiceUnavailable srv serviceSub -> do
+        logError $ "SMP server service unavailable: " <> showService srv serviceSub
+        removeServiceAndAssociations st srv >>= \case
+          Right (srvId, updated) -> do
+            logSubStatus srv "removed service association" updated updated
+            void $ subscribeSrvSubs ca st batchSize (srv, srvId, Nothing)
+          Left e -> logError $ "SMP server update and resubscription error " <> tshow e
  where
-    receiveSMP = do
-      st <- asks store
-      ps <- asks pushServer
-      stats <- asks serverStats
-      forever $ do
-        ((_, srv@(SMPServer (h :| _) _ _), _), THandleParams {sessionId}, ts) <- atomically $ readTBQueue msgQ
-        forM_ ts $ \(ntfId, t) -> case t of
-          STUnexpectedError e -> logError $ "SMP client unexpected error: " <> tshow e -- uncorrelated response, should not happen
-          STResponse {} -> pure () -- it was already reported as timeout error
-          STEvent msgOrErr -> do
-            let smpQueue = SMPQueueNtf srv ntfId
-            case msgOrErr of
-              Right (SMP.NMSG nmsgNonce encNMsgMeta) -> do
-                ntfTs <- liftIO getSystemTime
-                liftIO $ updatePeriodStats (activeSubs stats) ntfId
-                let newNtf = PNMessageData {smpQueue, ntfTs, nmsgNonce, encNMsgMeta}
-                    srvHost = safeDecodeUtf8 $ strEncode h
-                    isOwn = isOwnServer ca srv
-                liftIO (addTokenLastNtf st newNtf) >>= \case
-                  Right (tkn, lastNtfs) -> do
-                    pushNotification ps (Just srvHost) isOwn tkn $ PNMessage lastNtfs
-                    liftIO $ incNtfStat_ stats ntfReceived
-                    when isOwn $ liftIO $ incServerStat srvHost (ntfReceivedOwn stats)
-                  Left AUTH -> liftIO $ do
-                    incNtfStat_ stats ntfReceivedAuth
-                    when isOwn $ incServerStat srvHost (ntfReceivedAuthOwn stats)
-                  Left _ -> pure ()
-              Right SMP.END ->
-                whenM (atomically $ activeClientSession' ca sessionId srv) $
-                  void $ liftIO $ updateSrvSubStatus st smpQueue NSEnd
-              Right SMP.DELD ->
-                void $ liftIO $ updateSrvSubStatus st smpQueue NSDeleted
-              Right (SMP.ERR e) -> logError $ "SMP server error: " <> tshow e
-              Right _ -> logError "SMP server unexpected response"
-              Left e -> logError $ "SMP client error: " <> tshow e
-
-    receiveAgent = do
-      st <- asks store
-      batchSize <- asks $ subsBatchSize . config
-      liftIO $ forever $
-        atomically (readTBQueue agentQ) >>= \case
-          CAConnected srv serviceId -> do
-            let asService = if isJust serviceId then "as service " else ""
-            logInfo $ "SMP server reconnected " <> asService <> showServer' srv
-          CADisconnected srv nIds -> do
-            updated <- batchUpdateSrvSubStatus st srv Nothing nIds NSInactive
-            logSubStatus srv "disconnected" (L.length nIds) updated
-          CASubscribed srv serviceId nIds -> do
-            updated <- batchUpdateSrvSubStatus st srv serviceId nIds NSActive
-            let asService = if isJust serviceId then " as service" else ""
-            logSubStatus srv ("subscribed" <> asService) (L.length nIds) updated
-          CASubError srv errs -> do
-            forM_ (L.nonEmpty $ mapMaybe (\(nId, err) -> (nId,) <$> queueSubErrorStatus err) $ L.toList errs) $ \subStatuses -> do
-              updated <- batchUpdateSrvSubErrors st srv subStatuses
-              logSubErrors srv subStatuses updated
-              -- TODO [certs rcv] resubscribe queues with statuses NSErr and NSService
-          CAServiceDisconnected srv serviceSub ->
-            logNote $ "SMP server service disconnected " <> showService srv serviceSub
-          CAServiceSubscribed srv serviceSub@(ServiceSub _ n idsHash) (ServiceSub _ n' idsHash')
-            | n /= n' -> logWarn $ msg <> ", confirmed subs: " <> tshow n'
-            | idsHash /= idsHash' -> logWarn $ msg <> ", different IDs hash"
-            | otherwise -> logNote msg
-            where
-              msg = "SMP server service subscribed " <> showService srv serviceSub
-          CAServiceSubError srv serviceSub e ->
-            -- Errors that require re-subscribing queues directly are reported as CAServiceUnavailable.
-            -- See smpSubscribeService in Simplex.Messaging.Client.Agent
-            logError $ "SMP server service subscription error " <> showService srv serviceSub <> ": " <> tshow e
-          CAServiceUnavailable srv serviceSub -> do
-            logError $ "SMP server service unavailable: " <> showService srv serviceSub
-            removeServiceAndAssociations st srv >>= \case
-              Right (srvId, updated) -> do
-                logSubStatus srv "removed service association" updated updated
-                void $ subscribeSrvSubs ca st batchSize (srv, srvId, Nothing)
-              Left e -> logError $ "SMP server update and resubscription error " <> tshow e
-      where
-        showService srv (ServiceSub serviceId n _) = showServer' srv  <> ", service ID " <> decodeLatin1 (strEncode serviceId) <> ", " <> tshow n <> " subs"
-
+    showService srv (ServiceSub serviceId n _) = showServer' srv <> ", service ID " <> decodeLatin1 (strEncode serviceId) <> ", " <> tshow n <> " subs"
    logSubErrors :: SMPServer -> NonEmpty (SMP.NotifierId, NtfSubStatus) -> Int -> IO ()
-    logSubErrors srv subs updated = forM_ (L.group $ L.sort $ L.map snd subs) $ \ss -> do
+    logSubErrors srv subs updated = forM_ (L.group $ L.sort $ L.map snd subs) $ \ss ->
      logError $ "SMP server subscription errors " <> showServer' srv <> ": " <> tshow (L.head ss) <> " (" <> tshow (length ss) <> " errors, " <> tshow updated <> " subs updated)"
-
    queueSubErrorStatus :: SMPClientError -> Maybe NtfSubStatus
    queueSubErrorStatus = \case
      PCEProtocolError AUTH -> Just NSAuth
-      -- TODO [certs rcv] we could allow making individual subscriptions within service session to handle SERVICE error.
-      -- This would require full stack changes in SMP server, SMP client and SMP service agent.
      PCEProtocolError SERVICE -> Just NSService
      PCEProtocolError e -> updateErr "SMP error " e
      PCEResponseError e -> updateErr "ResponseError " e
@@ -623,12 +621,11 @@ ntfSubscriber NtfSubscriber {smpAgent = ca@SMPClientAgent {msgQ, agentQ}} =
      PCETransportError e -> updateErr "TransportError " e
      PCECryptoError e -> updateErr "CryptoError " e
      PCEIncompatibleHost -> Just $ NSErr "IncompatibleHost"
-      PCEServiceUnavailable -> Just NSService -- this error should not happen on individual subscriptions
+      PCEServiceUnavailable -> Just NSService
      PCEResponseTimeout -> Nothing
      PCENetworkError _ -> Nothing
      PCEIOError _ -> Nothing
      where
-        -- Note on moving to PostgreSQL: the idea of logging errors without e is removed here
        updateErr :: Show e => ByteString -> e -> Maybe NtfSubStatus
        updateErr errType e = Just $ NSErr $ errType <> bshow e

@@ -28,6 +28,7 @@ module Simplex.Messaging.Notifications.Server.Env
  ) where

 import Control.Concurrent (ThreadId)
+import Data.IORef
 import qualified Control.Exception as E
 import Control.Logger.Simple
 import Control.Monad
@@ -45,7 +46,7 @@ import qualified Data.X509.Validation as XV
 import Network.Socket
 import qualified Network.TLS as TLS
 import Numeric.Natural
-import Simplex.Messaging.Client (ProtocolClientError (..), SMPClientError)
+import Simplex.Messaging.Client (ProtocolClientError (..), SMPClientError, ServerTransmissionBatch)
 import Simplex.Messaging.Client.Agent
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Notifications.Protocol
@@ -54,14 +55,14 @@ import Simplex.Messaging.Notifications.Server.Stats
 import Simplex.Messaging.Notifications.Server.Store.Postgres
 import Simplex.Messaging.Notifications.Server.Store.Types
 import Simplex.Messaging.Notifications.Transport (NTFVersion, VersionRangeNTF)
-import Simplex.Messaging.Protocol (BasicAuth, CorrId, Party (..), SMPServer, SParty (..), ServiceId, Transmission)
+import Simplex.Messaging.Protocol (BasicAuth, BrokerMsg, CorrId, ErrorType, Party (..), SMPServer, SParty (..), ServiceId, Transmission)
 import Simplex.Messaging.Server.Env.STM (StartOptions (..))
 import Simplex.Messaging.Server.Expiration
 import Simplex.Messaging.Server.QueueStore.Postgres.Config (PostgresStoreCfg (..))
 import Simplex.Messaging.Session
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Transport (ASrvTransport, SMPServiceRole (..), ServiceCredentials (..), THandleParams, TransportPeer (..))
+import Simplex.Messaging.Transport (ASrvTransport, SMPServiceRole (..), SMPVersion, ServiceCredentials (..), THandleParams, TransportPeer (..))
 import Simplex.Messaging.Transport.Credentials (genCredentials, tlsCredentials)
 import Simplex.Messaging.Transport.Server (AddHTTP, ServerCredentials, TransportServerConfig, loadFingerprint, loadServerCredential)
 import Simplex.Messaging.Util (liftEitherWith, tshow)
@@ -119,17 +120,20 @@ data NtfEnv = NtfEnv
    serverStats :: NtfServerStats
  }

-newNtfServerEnv :: NtfServerConfig -> IO NtfEnv
-newNtfServerEnv config@NtfServerConfig {pushQSize, smpAgentCfg, apnsConfig, dbStoreConfig, ntfCredentials, useServiceCreds} = do
+newNtfServerEnv :: NtfServerConfig -> (IORef NtfEnv -> ServerTransmissionBatch SMPVersion ErrorType BrokerMsg -> IO ()) -> IO NtfEnv
+newNtfServerEnv config@NtfServerConfig {pushQSize, smpAgentCfg, apnsConfig, dbStoreConfig, ntfCredentials, useServiceCreds} mkProcessMsg = do
  random <- C.newRandom
  store <- newNtfDbStore dbStoreConfig
  tlsServerCreds <- loadServerCredential ntfCredentials
  XV.Fingerprint fp <- loadFingerprint ntfCredentials
  let dbService = if useServiceCreds then Just $ mkDbService random store else Nothing
-  subscriber <- newNtfSubscriber smpAgentCfg dbService random
+  envRef <- newIORef $ error "NtfEnv not initialized"
+  subscriber <- newNtfSubscriber smpAgentCfg (mkProcessMsg envRef) dbService random
  pushServer <- newNtfPushServer pushQSize apnsConfig
  serverStats <- newNtfServerStats =<< getCurrentTime
-  pure NtfEnv {config, subscriber, pushServer, store, random, tlsServerCreds, serverIdentity = C.KeyHash fp, serverStats}
+  let env = NtfEnv {config, subscriber, pushServer, store, random, tlsServerCreds, serverIdentity = C.KeyHash fp, serverStats}
+  writeIORef envRef env
+  pure env
  where
    mkDbService g st = DBService {getCredentials, updateServiceId}
      where
@@ -158,11 +162,11 @@ data NtfSubscriber = NtfSubscriber

 type SMPSubscriberVar = SessionVar SMPSubscriber

-newNtfSubscriber :: SMPClientAgentConfig -> Maybe DBService -> TVar ChaChaDRG -> IO NtfSubscriber
-newNtfSubscriber smpAgentCfg dbService random = do
+newNtfSubscriber :: SMPClientAgentConfig -> (ServerTransmissionBatch SMPVersion ErrorType BrokerMsg -> IO ()) -> Maybe DBService -> TVar ChaChaDRG -> IO NtfSubscriber
+newNtfSubscriber smpAgentCfg processMsg dbService random = do
  smpSubscribers <- TM.emptyIO
  subscriberSeq <- newTVarIO 0
-  smpAgent <- newSMPClientAgent SNotifierService smpAgentCfg dbService random
+  smpAgent <- newSMPClientAgent SNotifierService smpAgentCfg processMsg dbService random
  pure NtfSubscriber {smpSubscribers, subscriberSeq, smpAgent}

 data SMPSubscriber = SMPSubscriber
@@ -706,7 +706,7 @@ mkJournalStoreConfig queueStoreCfg storePath msgQueueQuota maxJournalMsgCount ma

 newSMPProxyAgent :: SMPClientAgentConfig -> TVar ChaChaDRG -> IO ProxyAgent
 newSMPProxyAgent smpAgentCfg random = do
-  smpAgent <- newSMPClientAgent SSender smpAgentCfg Nothing random
+  smpAgent <- newSMPClientAgent SSender smpAgentCfg (\_ -> pure ()) Nothing random
  pure ProxyAgent {smpAgent}

 readWriteQueueStore :: forall q. StoreQueueClass q => Bool -> (RecipientId -> QueueRec -> IO q) -> FilePath -> STMQueueStore q -> IO (StoreLog 'WriteMode)