diff --git a/ChatBot.Tests/Models/ChatSessionCompressionTests.cs b/ChatBot.Tests/Models/ChatSessionCompressionTests.cs
index 6b9ecf6..4ef094c 100644
--- a/ChatBot.Tests/Models/ChatSessionCompressionTests.cs
+++ b/ChatBot.Tests/Models/ChatSessionCompressionTests.cs
@@ -1,8 +1,10 @@
 using ChatBot.Models;
+using ChatBot.Models.Dto;
 using ChatBot.Services.Interfaces;
 using FluentAssertions;
 using Moq;
 using OllamaSharp.Models.Chat;
+using System.Collections.Concurrent;
 
 namespace ChatBot.Tests.Models;
 
@@ -17,15 +19,17 @@ public class ChatSessionCompressionTests
         session.SetCompressionService(compressionServiceMock.Object);
 
         // Setup compression service to return compressed messages
+        var compressedMessages = new List<ChatMessage>
+        {
+            new ChatMessage { Role = ChatRole.System.ToString(), Content = "System prompt" },
+            new ChatMessage { Role = ChatRole.User.ToString(), Content = "Compressed user message" }
+        };
         compressionServiceMock
-            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
-            .ReturnsAsync(
-                new List<ChatMessage>
-                {
-                    new() { Role = ChatRole.System, Content = "System prompt" },
-                    new() { Role = ChatRole.User, Content = "Compressed user message" },
-                }
-            );
+            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(compressedMessages);
+        compressionServiceMock
+            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
+            .Returns(true);
 
         // Add messages to session
         for (int i = 0; i < 10; i++)
@@ -57,9 +61,13 @@ public class ChatSessionCompressionTests
         session.SetCompressionService(compressionServiceMock.Object);
 
         // Setup compression service to throw an exception
+        var exception = new Exception("Compression failed");
         compressionServiceMock
-            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
-            .ThrowsAsync(new Exception("Compression failed"));
+            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
+            .ThrowsAsync(exception);
+        compressionServiceMock
+            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
+            .Returns(true);
 
         // Add messages to session
         for (int i = 0; i < 5; i++)
@@ -86,6 +94,11 @@ public class ChatSessionCompressionTests
         var session = new ChatSession();
         var compressionServiceMock = new Mock<IHistoryCompressionService>();
         session.SetCompressionService(compressionServiceMock.Object);
+        
+        // Setup compression service to return false for ShouldCompress when count is below threshold
+        compressionServiceMock
+            .Setup(x => x.ShouldCompress(It.Is<int>(c => c < 5), It.Is<int>(t => t == 5)))
+            .Returns(false);
 
         // Add messages to session (below threshold)
         session.AddMessage(new ChatMessage { Role = ChatRole.User, Content = "Message 1" });
@@ -100,7 +113,7 @@ public class ChatSessionCompressionTests
 
         // Assert - Should not call compression service
         compressionServiceMock.Verify(
-            x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()),
+            x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()),
             Times.Never
         );
 
@@ -116,17 +129,21 @@ public class ChatSessionCompressionTests
         var compressionServiceMock = new Mock<IHistoryCompressionService>();
         session.SetCompressionService(compressionServiceMock.Object);
 
-        // Setup compression service to delay to simulate processing time
+        // Setup compression service to simulate processing time
+        var delayedResult = new List<ChatMessage>
+        {
+            new ChatMessage { Role = ChatRole.System.ToString(), Content = "Compressed" }
+        };
         compressionServiceMock
-            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
-            .ReturnsAsync(
-                (List<ChatMessage> messages, int target) =>
-                    Task.Delay(50)
-                        .ContinueWith(_ => new List<ChatMessage>
-                        {
-                            new() { Role = ChatRole.System, Content = "Compressed" },
-                        })
-            );
+            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
+            .Returns(async (List<ChatMessage> messages, int target, CancellationToken ct) =>
+            {
+                await Task.Delay(50);
+                return delayedResult;
+            });
+        compressionServiceMock
+            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
+            .Returns(true);
 
         var tasks = new List<Task>();
         int messageCount = 5;
@@ -176,29 +193,28 @@ public class ChatSessionCompressionTests
 
         // Setup compression service to preserve system message
         compressionServiceMock
-            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
-            .ReturnsAsync(
-                (List<ChatMessage> messages, int target) =>
+            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
+            .Returns((List<ChatMessage> messages, int target, CancellationToken ct) =>
+            {
+                var systemMessage = messages.FirstOrDefault(m => m.Role == ChatRole.System.ToString());
+                var compressed = new List<ChatMessage>();
+
+                if (systemMessage != null)
                 {
-                    var systemMessage = messages.FirstOrDefault(m => m.Role == ChatRole.System);
-                    var compressed = new List<ChatMessage>();
-
-                    if (systemMessage != null)
-                    {
-                        compressed.Add(systemMessage);
-                    }
-
-                    compressed.Add(
-                        new ChatMessage
-                        {
-                            Role = ChatRole.User,
-                            Content = "Compressed user messages",
-                        }
-                    );
-
-                    return compressed;
+                    compressed.Add(systemMessage);
                 }
-            );
+
+                compressed.Add(new ChatMessage
+                {
+                    Role = ChatRole.User.ToString(),
+                    Content = "Compressed user messages"
+                });
+
+                return Task.FromResult(compressed);
+            });
+        compressionServiceMock
+            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
+            .Returns(true);
 
         // Add system message and some user messages
         session.AddMessage(new ChatMessage { Role = ChatRole.System, Content = "System prompt" });