From b3eb81b92807e66581d1be8486f99a4d18fcb5a5 Mon Sep 17 00:00:00 2001
From: Bilal Rafique <mejrafique@gmail.com>
Date: Mon, 8 Sep 2025 17:27:36 +0500
Subject: [PATCH] Add wait after transaction broadcast when processing Tx after
 broadcast

---
 pkg/lumera/modules/tx/impl.go      | 41 ++++++++++++++++++++++++++++++
 pkg/lumera/modules/tx/interface.go |  3 +++
 pkg/lumera/modules/tx/tx_mock.go   | 15 +++++++++++
 pkg/testutil/lumera.go             |  6 +++++
 4 files changed, 65 insertions(+)

diff --git a/pkg/lumera/modules/tx/impl.go b/pkg/lumera/modules/tx/impl.go
index 37481751..d342601b 100644
--- a/pkg/lumera/modules/tx/impl.go
+++ b/pkg/lumera/modules/tx/impl.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math"
 	"strconv"
+	"time"
 
 	"github.com/LumeraProtocol/supernode/v2/pkg/logtrace"
 	lumeracodec "github.com/LumeraProtocol/supernode/v2/pkg/lumera/codec"
@@ -184,6 +185,20 @@ func (m *module) BroadcastTransaction(ctx context.Context, txBytes []byte) (*sdk
 	return resp, nil
 }
 
+// GetTransaction queries a transaction by its hash
+func (m *module) GetTransaction(ctx context.Context, txHash string) (*sdktx.GetTxResponse, error) {
+	req := &sdktx.GetTxRequest{
+		Hash: txHash,
+	}
+
+	resp, err := m.client.GetTx(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get transaction: %w", err)
+	}
+
+	return resp, nil
+}
+
 // CalculateFee calculates the transaction fee based on gas usage and config
 func (m *module) CalculateFee(gasAmount uint64, config *TxConfig) string {
 	// Determine gas price (numeric) and denom. Accept both plain number (e.g., "0.025")
@@ -272,6 +287,32 @@ func (m *module) ProcessTransaction(ctx context.Context, msgs []types.Msg, accou
 		return result, fmt.Errorf("failed to broadcast transaction: %w", err)
 	}
 
+	if result != nil && result.TxResponse != nil && result.TxResponse.Code == 0 && len(result.TxResponse.Events) == 0 {
+		logtrace.Info(ctx, "Transaction broadcast successful, waiting for inclusion to get events...", nil)
+
+		// Retry 5 times with 1 second intervals
+		var txResp *sdktx.GetTxResponse
+		for i := 0; i < 5; i++ {
+			time.Sleep(1 * time.Second)
+
+			txResp, err = m.GetTransaction(ctx, result.TxResponse.TxHash)
+			if err == nil && txResp != nil && txResp.TxResponse != nil {
+				// Successfully got the transaction with events
+				logtrace.Info(ctx, fmt.Sprintf("Retrieved transaction with %d events", len(txResp.TxResponse.Events)), nil)
+				result.TxResponse = txResp.TxResponse
+				break
+			}
+
+			if err != nil {
+				logtrace.Warn(ctx, fmt.Sprintf("Attempt %d: failed to query transaction: %v", i+1, err), nil)
+			}
+		}
+	}
+
+	if len(result.TxResponse.Events) == 0 {
+		logtrace.Error(ctx, "Failed to retrieve transaction events after 5 attempts", nil)
+	}
+
 	return result, nil
 }
 
diff --git a/pkg/lumera/modules/tx/interface.go b/pkg/lumera/modules/tx/interface.go
index 23d184ea..eedaf7ab 100644
--- a/pkg/lumera/modules/tx/interface.go
+++ b/pkg/lumera/modules/tx/interface.go
@@ -35,6 +35,9 @@ type Module interface {
 	// BroadcastTransaction broadcasts a signed transaction and returns the result
 	BroadcastTransaction(ctx context.Context, txBytes []byte) (*sdktx.BroadcastTxResponse, error)
 
+	// GetTransaction queries a transaction by its hash
+	GetTransaction(ctx context.Context, txHash string) (*sdktx.GetTxResponse, error)
+
 	// CalculateFee calculates the transaction fee based on gas usage and config
 	CalculateFee(gasAmount uint64, config *TxConfig) string
 
diff --git a/pkg/lumera/modules/tx/tx_mock.go b/pkg/lumera/modules/tx/tx_mock.go
index c0eac536..2628151b 100644
--- a/pkg/lumera/modules/tx/tx_mock.go
+++ b/pkg/lumera/modules/tx/tx_mock.go
@@ -58,6 +58,21 @@ func (mr *MockModuleMockRecorder) BroadcastTransaction(ctx, txBytes any) *gomock
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BroadcastTransaction", reflect.TypeOf((*MockModule)(nil).BroadcastTransaction), ctx, txBytes)
 }
 
+// GetTransaction mocks base method.
+func (m *MockModule) GetTransaction(ctx context.Context, txHash string) (*tx.GetTxResponse, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetTransaction", ctx, txHash)
+	ret0, _ := ret[0].(*tx.GetTxResponse)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetTransaction indicates an expected call of GetTransaction.
+func (mr *MockModuleMockRecorder) GetTransaction(ctx, txHash any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTransaction", reflect.TypeOf((*MockModule)(nil).GetTransaction), ctx, txHash)
+}
+
 // BuildAndSignTransaction mocks base method.
 func (m *MockModule) BuildAndSignTransaction(ctx context.Context, msgs []types.Msg, accountInfo *types0.BaseAccount, gasLimit uint64, fee string, config *TxConfig) ([]byte, error) {
 	m.ctrl.T.Helper()
diff --git a/pkg/testutil/lumera.go b/pkg/testutil/lumera.go
index 7f5b778c..3f556a97 100644
--- a/pkg/testutil/lumera.go
+++ b/pkg/testutil/lumera.go
@@ -210,6 +210,12 @@ func (m *MockTxModule) ProcessTransaction(ctx context.Context, msgs []sdktypes.M
 	return &sdktx.BroadcastTxResponse{}, nil
 }
 
+// GetTransaction queries a transaction by its hash
+func (m *MockTxModule) GetTransaction(ctx context.Context, txHash string) (*sdktx.GetTxResponse, error) {
+	// Mock implementation returns empty transaction response
+	return &sdktx.GetTxResponse{}, nil
+}
+
 // MockNodeModule implements the node.Module interface for testing
 type MockNodeModule struct{}