diff --git a/Chatto/Source/ChatController/BaseChatViewController+Changes.swift b/Chatto/Source/ChatController/BaseChatViewController+Changes.swift
index a55da6f..2650661 100644
--- a/Chatto/Source/ChatController/BaseChatViewController+Changes.swift
+++ b/Chatto/Source/ChatController/BaseChatViewController+Changes.swift
@@ -26,12 +26,8 @@ import Foundation
 
 extension BaseChatViewController: ChatDataSourceDelegateProtocol {
 
-    public enum UpdateContext {
-        case Normal
-        case FirstLoad
-        case Pagination
-        case Reload
-        case MessageCountReduction
+    public func chatDataSourceDidUpdate(chatDataSource: ChatDataSourceProtocol, context: UpdateContext) {
+        self.enqueueModelUpdate(context: context)
     }
 
     public func chatDataSourceDidUpdate(chatDataSource: ChatDataSourceProtocol) {
diff --git a/Chatto/Source/ChatController/ChatDataSourceProtocol.swift b/Chatto/Source/ChatController/ChatDataSourceProtocol.swift
index a833287..36143d0 100644
--- a/Chatto/Source/ChatController/ChatDataSourceProtocol.swift
+++ b/Chatto/Source/ChatController/ChatDataSourceProtocol.swift
@@ -24,8 +24,17 @@
 
 import Foundation
 
+public enum UpdateContext {
+    case Normal
+    case FirstLoad
+    case Pagination
+    case Reload
+    case MessageCountReduction
+}
+
 public protocol ChatDataSourceDelegateProtocol: class {
     func chatDataSourceDidUpdate(chatDataSource: ChatDataSourceProtocol)
+    func chatDataSourceDidUpdate(chatDataSource: ChatDataSourceProtocol, context: UpdateContext)
 }
 
 public protocol ChatDataSourceProtocol: class {