@@ -9,6 +9,7 @@ import (
 	"math/big"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/ethereum/go-ethereum/beacon/engine"
@@ -40,9 +41,14 @@ type EngineClient struct {
 	genesisHash   common.Hash       // Hash of the genesis block
 	initialHeight uint64
 	feeRecipient  common.Address // Address to receive transaction fees
+
+	lastHeadBlockHash      common.Hash // Store last non-finalized HeadBlockHash
+	lastSafeBlockHash      common.Hash // Store last non-finalized SafeBlockHash
+	lastFinalizedBlockHash common.Hash // Store last finalized block hash
+	mu                     sync.Mutex  // Mutex to protect concurrent access to mutable state
 }
 
-// NewPureEngineExecutionClient creates a new instance of EngineAPIExecutionClient
+// NewEngineExecutionClient creates a new instance of EngineClient
 func NewEngineExecutionClient(
 	ethURL,
 	engineURL string,
@@ -86,6 +92,8 @@ func NewEngineExecutionClient(
 
 // InitChain initializes the blockchain with the given genesis parameters
 func (c *EngineClient) InitChain(ctx context.Context, genesisTime time.Time, initialHeight uint64, chainID string) ([]byte, uint64, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	if initialHeight != 1 {
 		return nil, 0, fmt.Errorf("initialHeight must be 1, got %d", initialHeight)
 	}
@@ -116,6 +124,8 @@ func (c *EngineClient) InitChain(ctx context.Context, genesisTime time.Time, ini
 
 // GetTxs retrieves transactions from the current execution payload
 func (c *EngineClient) GetTxs(ctx context.Context) ([][]byte, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	var result struct {
 		Pending map[string]map[string]*types.Transaction `json:"pending"`
 		Queued  map[string]map[string]*types.Transaction `json:"queued"`
@@ -153,6 +163,8 @@ func (c *EngineClient) GetTxs(ctx context.Context) ([][]byte, error) {
 
 // ExecuteTxs executes the given transactions at the specified block height and timestamp
 func (c *EngineClient) ExecuteTxs(ctx context.Context, txs [][]byte, blockHeight uint64, timestamp time.Time, prevStateRoot []byte) (updatedStateRoot []byte, maxBytes uint64, err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	// convert rollkit tx to eth tx
 	ethTxs := make([]*types.Transaction, len(txs))
 	for i, tx := range txs {
@@ -264,12 +276,26 @@ func (c *EngineClient) ExecuteTxs(ctx context.Context, txs [][]byte, blockHeight
 }
 
 func (c *EngineClient) setFinal(ctx context.Context, blockHash common.Hash, isFinal bool) error {
-	args := engine.ForkchoiceStateV1{
-		HeadBlockHash: blockHash,
-		SafeBlockHash: blockHash,
-	}
+	var args engine.ForkchoiceStateV1
+
 	if isFinal {
-		args.FinalizedBlockHash = blockHash
+		// Use the last stored values for head and safe block hashes
+		args = engine.ForkchoiceStateV1{
+			HeadBlockHash:      c.lastHeadBlockHash,
+			SafeBlockHash:      c.lastSafeBlockHash,
+			FinalizedBlockHash: blockHash,
+		}
+		// Store the current blockHash as the latest finalized block hash
+		c.lastFinalizedBlockHash = blockHash
+	} else {
+		// Store the current blockHash as the latest head and safe block hashes
+		c.lastHeadBlockHash = blockHash
+		c.lastSafeBlockHash = blockHash
+		args = engine.ForkchoiceStateV1{
+			HeadBlockHash:      blockHash,
+			SafeBlockHash:      blockHash,
+			FinalizedBlockHash: c.lastFinalizedBlockHash,
+		}
 	}
 
 	var forkchoiceResult engine.ForkChoiceResponse
@@ -290,6 +316,8 @@ func (c *EngineClient) setFinal(ctx context.Context, blockHash common.Hash, isFi
 
 // SetFinal marks the block at the given height as finalized
 func (c *EngineClient) SetFinal(ctx context.Context, blockHeight uint64) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	blockHash, _, _, _, err := c.getBlockInfo(ctx, blockHeight)
 	if err != nil {
 		return fmt.Errorf("failed to get block info: %w", err)
 
@@ -199,7 +199,6 @@ func TestEngineExecution(t *testing.T) {
 			require.GreaterOrEqual(tt, lastTxs, 0, "Number of transactions should be non-negative")
 
 			prevStateRoot = newStateRoot
-			fmt.Println("all good blockheight", blockHeight)
 		}
 	})
 }
Original file line number	Diff line number	Diff line change
`@@ -199,7 +199,6 @@ func TestEngineExecution(t *testing.T) {`
`199`	`199`	`require.GreaterOrEqual(tt, lastTxs, 0, "Number of transactions should be non-negative")`
`200`	`200`
`201`	`201`	`prevStateRoot = newStateRoot`
`202`		`- fmt.Println("all good blockheight", blockHeight)`
`203`	`202`	`}`
`204`	`203`	`})`
`205`	`204`	`}`