diff --git a/Tapeti.Flow/ConfigExtensions.cs b/Tapeti.Flow/ConfigExtensions.cs
index ffaad41..38ad5f3 100644
--- a/Tapeti.Flow/ConfigExtensions.cs
+++ b/Tapeti.Flow/ConfigExtensions.cs
@@ -2,9 +2,9 @@
 {
     public static class ConfigExtensions
     {
-        public static TapetiConfig WithFlow(this TapetiConfig config)
+        public static TapetiConfig WithFlow(this TapetiConfig config, IFlowRepository<Default.FlowState> flowRepository)
         {
-            config.Use(new FlowMiddleware());
+            config.Use(new FlowMiddleware(flowRepository));
             return config;
         }
     }
diff --git a/Tapeti.Flow/FlowMiddleware.cs b/Tapeti.Flow/FlowMiddleware.cs
index 739d0b8..cba399d 100644
--- a/Tapeti.Flow/FlowMiddleware.cs
+++ b/Tapeti.Flow/FlowMiddleware.cs
@@ -6,12 +6,19 @@ namespace Tapeti.Flow
 {
     public class FlowMiddleware : ITapetiExtension
     {
+        private IFlowRepository<Default.FlowState> flowRepository;
+
+        public FlowMiddleware(IFlowRepository<Default.FlowState> flowRepository)
+        {
+            this.flowRepository = flowRepository;
+        }
+
         public void RegisterDefaults(IDependencyContainer container)
         {
             container.RegisterDefault<IFlowProvider, FlowProvider>();
             container.RegisterDefault<IFlowStarter, FlowStarter>();
             container.RegisterDefault<IFlowHandler, FlowProvider>();
-            container.RegisterDefault<IFlowRepository<FlowState>, NonPersistentFlowRepository<FlowState>>();
+            container.RegisterDefault<IFlowRepository<FlowState>>(() => flowRepository ?? new NonPersistentFlowRepository<Default.FlowState>());
             container.RegisterDefault<IFlowStore, FlowStore>();
         }