diff --git a/CHANGELOG.md b/CHANGELOG.md
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+# 0.2.1
+
+* Fix missing export of `relayStateKey` and change its type.
+
 # 0.2.0
 
 * Added parsing for RelayState from form data, as sent by e.g. Shibboleth when a `target` query string parameter is passed to the unsolicited SSO endpoint.
diff --git a/src/Network/Wai/SAML2.hs b/src/Network/Wai/SAML2.hs
--- a/src/Network/Wai/SAML2.hs
+++ b/src/Network/Wai/SAML2.hs
@@ -21,6 +21,7 @@
     assertionKey,
     errorKey,
     saml2Vault,
+    relayStateKey,
 
     -- * Re-exports
     module Network.Wai.SAML2.Config,
@@ -31,6 +32,7 @@
 --------------------------------------------------------------------------------
 
 import qualified Data.ByteString as BS
+import Data.Maybe (fromMaybe)
 import qualified Data.Vault.Lazy as V
 
 import Network.Wai 
@@ -169,7 +171,7 @@
 -- | 'relayStateKey' is a vault key for retrieving the relay state
 -- from request vaults if the 'saml2Vault' 'Middleware' is used
 -- and the assertion is valid.
-relayStateKey :: V.Key (Maybe BS.ByteString)
+relayStateKey :: V.Key BS.ByteString
 relayStateKey = unsafePerformIO V.newKey
 
 -- | 'errorKey' is a vault key for retrieving SAML2 errors from request vaults
@@ -189,9 +191,13 @@
                 vault = V.insert errorKey err (vault req)
             } sendResponse 
           callback (Right result) app req sendResponse = do
+            let mRelayState = relayState result 
+            let vlt = vault req
+
             app req{
                 vault = V.insert assertionKey (assertion result)
-                      $ V.insert relayStateKey (relayState result) (vault req)
+                      $ fromMaybe vlt $ mRelayState >>= \rs -> 
+                            pure $ V.insert relayStateKey rs vlt
             } sendResponse
 
 --------------------------------------------------------------------------------
@@ -203,3 +209,5 @@
     -- | The assertion obtained from the response that has been validated.
     assertion :: !Assertion
 } deriving (Eq, Show)
+
+--------------------------------------------------------------------------------
diff --git a/wai-saml2.cabal b/wai-saml2.cabal
--- a/wai-saml2.cabal
+++ b/wai-saml2.cabal
@@ -4,10 +4,10 @@
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: ddb0bf4bfb873f58ec1a0de52ad6c91db0957cdde75805acb20f431229689da2
+-- hash: 0ab786ffa39a1a85044697458c860145df0f561c8b7ba0211f4fc6718c81d5ef
 
 name:           wai-saml2
-version:        0.2.0.0
+version:        0.2.1.0
 synopsis:       SAML2 assertion validation as WAI middleware
 description:    A Haskell library which implements SAML2 assertion validation as WAI middleware
 category:       Security