diff --git a/runtime/v2/module.go b/runtime/v2/module.go
index 042f4cadf7..4030f2050a 100644
--- a/runtime/v2/module.go
+++ b/runtime/v2/module.go
@@ -161,11 +161,10 @@ func SetupAppBuilder(inputs AppInputs) {
 	app.moduleManager.RegisterInterfaces(inputs.InterfaceRegistrar)
 	app.moduleManager.RegisterLegacyAminoCodec(inputs.LegacyAmino)
 
-	// TODO: this is a bit of a hack, but it's the only way to get the store keys into the app
-	// registerStoreKey could instead set this on StoreOptions directly
 	if inputs.StoreOptions != nil {
 		inputs.AppBuilder.storeOptions = inputs.StoreOptions
 		inputs.AppBuilder.storeOptions.StoreKeys = inputs.AppBuilder.app.storeKeys
+		inputs.AppBuilder.storeOptions.StoreKeys = append(inputs.AppBuilder.storeOptions.StoreKeys, "stf")
 	}
 }
 
diff --git a/server/v2/stf/branch/writer_map.go b/server/v2/stf/branch/writer_map.go
index 8e7edf0e41..ebe02f4e7b 100644
--- a/server/v2/stf/branch/writer_map.go
+++ b/server/v2/stf/branch/writer_map.go
@@ -77,6 +77,8 @@ func (b WriterMap) GetStateChanges() ([]store.StateChanges, error) {
 	return sc, nil
 }
 
+// recurseStateChanges will recursively collect state changes from the tree of
+// WriterMap's and write them to the `changes` map.
 func (b WriterMap) recurseStateChanges(changes map[string][]store.KVPair) error {
 	// depth first
 	if wr, ok := b.state.(WriterMap); ok {