diff --git a/CHANGELOG.md b/CHANGELOG.md
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## Version 0.4.1 (2017-06-28)
+
+- Fixed runtime type error that could occur when shrinking state machine commands (#91, @jystic)
+
 ## Version 0.4 (2017-06-28)
 
 - Abstract state machine testing, check out the [process registry example](https://github.com/hedgehogqa/haskell-hedgehog/blob/master/hedgehog-example/test/Test/Example/Registry.hs) to see how it works (#89, @jystic)
diff --git a/README.md b/README.md
--- a/README.md
+++ b/README.md
@@ -76,7 +76,7 @@
 ```
 
  [hackage]: http://hackage.haskell.org/package/hedgehog
- [hackage-shield]: https://img.shields.io/badge/hackage-v0.4-blue.svg
+ [hackage-shield]: https://img.shields.io/badge/hackage-v0.4.1-blue.svg
 
  [travis]: https://travis-ci.org/hedgehogqa/haskell-hedgehog
  [travis-shield]: https://travis-ci.org/hedgehogqa/haskell-hedgehog.svg?branch=master
diff --git a/hedgehog.cabal b/hedgehog.cabal
--- a/hedgehog.cabal
+++ b/hedgehog.cabal
@@ -1,7 +1,7 @@
 name:
   hedgehog
 version:
-  0.4
+  0.4.1
 license:
   BSD3
 author:
diff --git a/src/Hedgehog/Internal/State.hs b/src/Hedgehog/Internal/State.hs
--- a/src/Hedgehog/Internal/State.hs
+++ b/src/Hedgehog/Internal/State.hs
@@ -55,8 +55,6 @@
 import           Data.Map (Map)
 import qualified Data.Map as Map
 import qualified Data.Maybe as Maybe
-import           Data.Set (Set)
-import qualified Data.Set as Set
 import           Data.Typeable (Typeable, TypeRep, Proxy(..), typeRep)
 
 import           Hedgehog.Internal.Gen (Gen)
@@ -338,31 +336,52 @@
       showString " :<- " .
       showsPrec 11 input
 
+-- | Extract the variable name and the type from a symbolic value.
+--
+takeSymbolic :: forall a. Symbolic a -> (Var, TypeRep)
+takeSymbolic (Symbolic var) =
+  (var, typeRep (Proxy :: Proxy a))
+
+-- | Insert a symbolic variable in to a map of variables to types.
+--
+insertSymbolic :: Symbolic a -> Map Var TypeRep -> Map Var TypeRep
+insertSymbolic s =
+  let
+    (var, typ) =
+      takeSymbolic s
+  in
+    Map.insert var typ
+
 -- | Collects all the symbolic values in a data structure and produces a set of
 --   all the variables they refer to.
 --
-takeVariables :: HTraversable t => t Symbolic -> Set Var
+takeVariables :: forall t. HTraversable t => t Symbolic -> Map Var TypeRep
 takeVariables xs =
   let
-    go x@(Symbolic var) = do
-      modify (Set.insert var)
+    go x = do
+      modify (insertSymbolic x)
       pure x
   in
-    flip execState Set.empty $ htraverse go xs
+    flip execState Map.empty $ htraverse go xs
 
 -- | Checks that the symbolic values in the data structure refer only to the
---   variables in the provided set.
+--   variables in the provided set, and that they are of the correct type.
 --
-variablesOK :: HTraversable t => t Symbolic -> Set Var -> Bool
+variablesOK :: HTraversable t => t Symbolic -> Map Var TypeRep -> Bool
 variablesOK xs allowed =
-  Set.null (takeVariables xs `Set.difference` allowed)
+  let
+    vars =
+      takeVariables xs
+  in
+    Map.null (vars `Map.difference` allowed) &&
+    and (Map.intersectionWith (==) vars allowed)
 
 -- | Drops invalid actions from the sequence.
 --
 dropInvalid :: (forall v. state v) -> [Action m state] -> [Action m state]
 dropInvalid initial =
   let
-    loop step@(Action input output@(Symbolic var) _execute require update _ensure) = do
+    loop step@(Action input output _execute require update _ensure) = do
       ((state0, vars0), steps0) <- get
 
       when (require state0 input && variablesOK input vars0) $
@@ -371,14 +390,14 @@
             update state0 input output
 
           vars =
-            Set.insert var vars0
+            insertSymbolic output vars0
 
           steps =
             steps0 ++ [step]
         in
           put ((state, vars), steps)
   in
-    snd . flip execState ((initial, Set.empty), []) . traverse_ loop
+    snd . flip execState ((initial, Map.empty), []) . traverse_ loop
 
 -- | Generates a single action from a set of possible commands.
 --