Fix state consistency in MonadHead instacesHEADmaster

1 files changed, 4 insertions, 2 deletions

diff --git a/main/Test.hs b/main/Test.hs
index 1755031..9038083 100644
--- a/main/Test.hs
+++ b/main/Test.hs
@@ -282,8 +282,10 @@ instance MonadStorage CommandM where
 
 instance MonadHead LocalState CommandM where
     updateLocalHead f = do
-        Just h <- gets tsHead
-        (Just h', x) <- maybe (fail "failed to reload head") (flip updateHead f) =<< reloadHead h
+        h <- maybe (fail "failed to reload head") return =<< reloadHead =<< getOrLoadHead
+        ( Just h', x ) <- updateHead' h $ \h' -> do
+            modify $ \s -> s { tsHead = Just h' }
+            f (headStoredObject h')
         modify $ \s -> s { tsHead = Just h' }
         return x