diff --git a/src/Infrastructure/BotSharp.Abstraction/Repositories/IBotSharpRepository.cs b/src/Infrastructure/BotSharp.Abstraction/Repositories/IBotSharpRepository.cs
index 1834dfd5..b8ca82a7 100644
--- a/src/Infrastructure/BotSharp.Abstraction/Repositories/IBotSharpRepository.cs
+++ b/src/Infrastructure/BotSharp.Abstraction/Repositories/IBotSharpRepository.cs
@@ -33,7 +33,7 @@ public interface IBotSharpRepository
     List<DialogElement> GetConversationDialogs(string conversationId);
     void UpdateConversationDialogElements(string conversationId, List<DialogContentUpdateModel> updateElements);
     void AppendConversationDialogs(string conversationId, List<DialogElement> dialogs);
-    List<StateKeyValue> GetConversationStates(string conversationId);
+    ConversationState GetConversationStates(string conversationId);
     void UpdateConversationStates(string conversationId, List<StateKeyValue> states);
     void UpdateConversationStatus(string conversationId, string status);
     Conversation GetConversation(string conversationId);
diff --git a/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationStateService.cs b/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationStateService.cs
index a6c4642e..0c106612 100644
--- a/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationStateService.cs
+++ b/src/Infrastructure/BotSharp.Core/Conversations/Services/ConversationStateService.cs
@@ -48,7 +48,7 @@ public class ConversationStateService : IConversationStateService, IDisposable
 
         if (_states.TryGetValue(name, out var values))
         {
-            preValue = values.LastOrDefault()?.Data ?? string.Empty;
+            preValue = values?.LastOrDefault()?.Data ?? string.Empty;
         }
 
         if (!_states.ContainsKey(name) || preValue != currentValue)
@@ -82,15 +82,14 @@ public class ConversationStateService : IConversationStateService, IDisposable
     {
         _conversationId = conversationId;
 
-        var savedStates = _db.GetConversationStates(_conversationId).ToList();
-        _states = new ConversationState(savedStates);
+        _states = _db.GetConversationStates(_conversationId);
         var curStates = new Dictionary<string, string>();
 
-        if (!savedStates.IsNullOrEmpty())
+        if (!_states.IsNullOrEmpty())
         {
-            foreach (var state in savedStates)
+            foreach (var state in _states)
             {
-                var value = state.Values.LastOrDefault()?.Data ?? string.Empty;
+                var value = state.Value?.LastOrDefault()?.Data ?? string.Empty;
                 curStates[state.Key] = value;
                 _logger.LogInformation($"[STATE] {state.Key} : {value}");
             }
@@ -134,7 +133,7 @@ public class ConversationStateService : IConversationStateService, IDisposable
         var curStates = new Dictionary<string, string>();
         foreach (var state in _states)
         {
-            curStates[state.Key] = state.Value.LastOrDefault()?.Data ?? string.Empty;
+            curStates[state.Key] = state.Value?.LastOrDefault()?.Data ?? string.Empty;
         }
         return curStates;
     }
diff --git a/src/Infrastructure/BotSharp.Core/Repository/BotSharpDbContext.cs b/src/Infrastructure/BotSharp.Core/Repository/BotSharpDbContext.cs
index ec8ac5f6..c3e5ad5e 100644
--- a/src/Infrastructure/BotSharp.Core/Repository/BotSharpDbContext.cs
+++ b/src/Infrastructure/BotSharp.Core/Repository/BotSharpDbContext.cs
@@ -152,7 +152,7 @@ public class BotSharpDbContext : Database, IBotSharpRepository
         throw new NotImplementedException();
     }
 
-    public List<StateKeyValue> GetConversationStates(string conversationId)
+    public ConversationState GetConversationStates(string conversationId)
     {
         throw new NotImplementedException();
     }
diff --git a/src/Infrastructure/BotSharp.Core/Repository/FileRepository.cs b/src/Infrastructure/BotSharp.Core/Repository/FileRepository.cs
index 030ac91a..fb20db1f 100644
--- a/src/Infrastructure/BotSharp.Core/Repository/FileRepository.cs
+++ b/src/Infrastructure/BotSharp.Core/Repository/FileRepository.cs
@@ -753,7 +753,7 @@ public class FileRepository : IBotSharpRepository
         }
     }
 
-    public List<StateKeyValue> GetConversationStates(string conversationId)
+    public ConversationState GetConversationStates(string conversationId)
     {
         var states = new List<StateKeyValue>();
         var convDir = FindConversationDirectory(conversationId);
@@ -763,7 +763,7 @@ public class FileRepository : IBotSharpRepository
             states = CollectConversationStates(stateFile);
         }
 
-        return states;
+        return new ConversationState(states);
     }
 
     public void UpdateConversationStates(string conversationId, List<StateKeyValue> states)
@@ -821,7 +821,7 @@ public class FileRepository : IBotSharpRepository
             var curStates = new Dictionary<string, string>();
             states.ForEach(x =>
             {
-                curStates[x.Key] = x.Values.LastOrDefault()?.Data ?? string.Empty;
+                curStates[x.Key] = x.Values?.LastOrDefault()?.Data ?? string.Empty;
             });
             record.States = curStates;
         }
diff --git a/src/Plugins/BotSharp.Plugin.MongoStorage/Repository/MongoRepository.cs b/src/Plugins/BotSharp.Plugin.MongoStorage/Repository/MongoRepository.cs
index 7c57ec43..5baf65c2 100644
--- a/src/Plugins/BotSharp.Plugin.MongoStorage/Repository/MongoRepository.cs
+++ b/src/Plugins/BotSharp.Plugin.MongoStorage/Repository/MongoRepository.cs
@@ -746,9 +746,9 @@ public class MongoRepository : IBotSharpRepository
         _dc.Conversations.UpdateOne(filterConv, updateConv);
     }
 
-    public List<StateKeyValue> GetConversationStates(string conversationId)
+    public ConversationState GetConversationStates(string conversationId)
     {
-        var states = new List<StateKeyValue>();
+        var states = new ConversationState();
         if (string.IsNullOrEmpty(conversationId)) return states;
 
         var filter = Builders<ConversationStateDocument>.Filter.Eq(x => x.ConversationId, conversationId);
@@ -756,7 +756,7 @@ public class MongoRepository : IBotSharpRepository
         if (foundStates == null || foundStates.States.IsNullOrEmpty()) return states;
 
         var savedStates = foundStates.States.Select(x => StateMongoElement.ToDomainElement(x)).ToList();
-        return savedStates;
+        return new ConversationState(savedStates);
     }
 
     public void UpdateConversationStates(string conversationId, List<StateKeyValue> states)
@@ -806,7 +806,7 @@ public class MongoRepository : IBotSharpRepository
         var curStates = new Dictionary<string, string>();
         states.States.ForEach(x =>
         {
-            curStates[x.Key] = x.Values.LastOrDefault()?.Data ?? string.Empty;
+            curStates[x.Key] = x.Values?.LastOrDefault()?.Data ?? string.Empty;
         });
 
         return new Conversation