diff --git a/consensus/state.go b/consensus/state.go
index 9dead0cfb..4b541ccb1 100644
--- a/consensus/state.go
+++ b/consensus/state.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"reflect"
 	"sync"
 	"time"
@@ -348,6 +349,23 @@ func (cs *ConsensusState) OnStart() error {
 		return err
 	}
 
+	// If the latest block was applied in the tmsp handshake,
+	// we may not have written the current height to the wal,
+	// so check here and write it if not found.
+	// TODO: remove this and run the handhsake/replay
+	// through the consensus state with a mock app
+	gr, found, err := cs.wal.group.Search("#HEIGHT: ", makeHeightSearchFunc(cs.Height))
+	if (err == io.EOF || !found) && cs.Step == RoundStepNewHeight {
+		log.Warn("Height not found in wal. Writing new height", "height", cs.Height)
+		rs := cs.RoundStateEvent()
+		cs.wal.Save(rs)
+	} else if err != nil {
+		return err
+	}
+	if gr != nil {
+		gr.Close()
+	}
+
 	// we need the timeoutRoutine for replay so
 	//  we don't block on the tick chan.
 	// NOTE: we will get a build up of garbage go routines
@@ -362,14 +380,6 @@ func (cs *ConsensusState) OnStart() error {
 		// let's go for it anyways, maybe we're fine
 	}
 
-	// If the latest block was applied in the tmsp handshake,
-	// we may not have written the current height to the wal,
-	// so write it here in case
-	if cs.Step == RoundStepNewHeight {
-		log.Warn("wal.writeHeight", "height", cs.Height)
-		cs.wal.writeHeight(cs.Height)
-	}
-
 	// now start the receiveRoutine
 	go cs.receiveRoutine(0)
 
diff --git a/state/execution.go b/state/execution.go
index 47ca01498..f8e2c1d82 100644
--- a/state/execution.go
+++ b/state/execution.go
@@ -418,6 +418,9 @@ func (h *Handshaker) ReplayBlocks(appHash []byte, appBlockHeight int, appConnCon
 		// replay all blocks starting with appBlockHeight+1
 		var eventCache types.Fireable // nil
 
+		// TODO: use stateBlockHeight instead and let the consensus state
+		// do the replay
+
 		var appHash []byte
 		for i := appBlockHeight + 1; i <= storeBlockHeight; i++ {
 			h.nBlocks += 1
@@ -443,6 +446,5 @@ func (h *Handshaker) ReplayBlocks(appHash []byte, appBlockHeight int, appConnCon
 		}
 		return nil
 	}
-
 	return nil
 }