diff --git a/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/TxSubmission/Direct.hs b/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/TxSubmission/Direct.hs
index 3a1c94d7b10..50e7212cfe4 100644
--- a/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/TxSubmission/Direct.hs
+++ b/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/TxSubmission/Direct.hs
@@ -53,8 +53,8 @@ directPipelined (TxSubmissionServerPipelined mserver)
     directSender q (SendMsgRequestTxsPipelined txids serverNext)
                    ClientStIdle{recvMsgRequestTxs} = do
       server' <- serverNext
-      SendMsgReplyTxs txs onTxsSent client' <- recvMsgRequestTxs txids
-      _ <- onTxsSent txs
+      SendMsgReplyTxs txs mClient' <- recvMsgRequestTxs txids
+      client' <- mClient'
       directSender (enqueue (CollectTxs txids txs) q) server' client'
 
     directSender q (CollectPipelined (Just server') _) client =
diff --git a/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/TxSubmission/Examples.hs b/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/TxSubmission/Examples.hs
index 5b471e8dbbd..16759ba2799 100644
--- a/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/TxSubmission/Examples.hs
+++ b/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/TxSubmission/Examples.hs
@@ -141,7 +141,7 @@ txSubmissionClient tracer txId txSize maxUnacked =
           traceWith tracer (EventRecvMsgRequestTxs unackedSeq unackedMap
                                                    remainingTxs txids)
           case [ txid | txid <- txids, txid `Map.notMember` unackedMap ] of
-            [] -> pure (SendMsgReplyTxs txs (const $ pure ()) client')
+            [] -> pure (SendMsgReplyTxs txs (pure client'))
               where
                 txs         = map (unackedMap Map.!) txids
                 client'     = client unackedSeq unackedMap' remainingTxs
diff --git a/ouroboros-network/src/Ouroboros/Network/Protocol/TxSubmission/Client.hs b/ouroboros-network/src/Ouroboros/Network/Protocol/TxSubmission/Client.hs
index 6a38b0351b9..dc8a8e194e4 100644
--- a/ouroboros-network/src/Ouroboros/Network/Protocol/TxSubmission/Client.hs
+++ b/ouroboros-network/src/Ouroboros/Network/Protocol/TxSubmission/Client.hs
@@ -80,8 +80,7 @@ data ClientStTxIds blocking txid tx m a where
 
 data ClientStTxs txid tx m a where
   SendMsgReplyTxs   :: [tx]
-                    -> ([tx] -> m ())
-                    -> ClientStIdle txid tx m a
+                    -> m (ClientStIdle txid tx m a)
                     -> ClientStTxs  txid tx m a
 
 
@@ -111,7 +110,7 @@ txSubmissionClientPeer (TxSubmissionClient client) =
                              (Done TokDone result)
 
         MsgRequestTxs txids -> Effect $ do
-          SendMsgReplyTxs txs onTxsSent k <- recvMsgRequestTxs txids
+          SendMsgReplyTxs txs k <- recvMsgRequestTxs txids
           return $ Yield (ClientAgency TokTxs)
                          (MsgReplyTxs txs)
-                         (Effect $ onTxsSent txs >> return (go k))
+                         (Effect $ k >>= return . go)
diff --git a/ouroboros-network/src/Ouroboros/Network/TxSubmission/Outbound.hs b/ouroboros-network/src/Ouroboros/Network/TxSubmission/Outbound.hs
index 180b5efca5c..2f720c1f8ce 100644
--- a/ouroboros-network/src/Ouroboros/Network/TxSubmission/Outbound.hs
+++ b/ouroboros-network/src/Ouroboros/Network/TxSubmission/Outbound.hs
@@ -190,6 +190,6 @@ txSubmissionOutbound tracer maxUnacked TxSubmissionMempoolReader{..} =
               client'      = client unackedSeq unackedMap' lastIdx
 
               -- Trace the transactions to be sent in the response.
-              traceTxsSent = traceWith tracer . TraceTxSubmissionOutboundSendMsgReplyTxs
+              traceTxsSent = traceWith tracer (TraceTxSubmissionOutboundSendMsgReplyTxs txs)
 
-          return $ SendMsgReplyTxs txs traceTxsSent client'
+          return $ SendMsgReplyTxs txs (traceTxsSent >> pure client')