diff --git a/src/BE/Controllers/Chats/Chats/ChatController.cs b/src/BE/Controllers/Chats/Chats/ChatController.cs
index 9ea43738..a268dd10 100644
--- a/src/BE/Controllers/Chats/Chats/ChatController.cs
+++ b/src/BE/Controllers/Chats/Chats/ChatController.cs
@@ -67,17 +67,20 @@ public async Task<IActionResult> StartConversationStreamed(
         }
 
         // ensure chat.ChatSpan contains all span ids that in request, otherwise return error
-        if (req.Spans.Any(x => !chat.ChatSpans.Any(y => y.SpanId == x.Id)))
+        if (req.MessageId == null)
         {
-            return BadRequest("Invalid span id");
+            if (req.Spans.Any(x => !chat.ChatSpans.Any(y => y.SpanId == x.Id)))
+            {
+                return BadRequest("Invalid span id");
+            }
         }
 
         // get span id -> model id mapping but request only contains span id, so we need to get model id from chat.ChatSpan
         Dictionary<byte, short> spanModelMapping = req.Spans.ToDictionary(x => x.Id, x => chat.ChatSpans.First(y => y.SpanId == x.Id).ModelId);
         Dictionary<short, UserModel> userModels = await userModelManager.GetUserModels(currentUser.Id, [.. spanModelMapping.Values], cancellationToken);
 
-        // ensure all model ids are valid
-        if (userModels.Count != spanModelMapping.Count)
+        // ensure spanModelMapping contains all userModels
+        if (spanModelMapping.Values.Any(x => !userModels.ContainsKey(x)))
         {
             return BadRequest("Invalid span model");
         }
@@ -310,6 +313,7 @@ private static async Task<ChatSpanResponse> ProcessChatSpan(
             [
                 MessageContent.FromText(icc.FullResponse.TextSegment),
             ],
+            SpanId = span.Id,
             CreatedAt = DateTime.UtcNow,
             ParentId = req.MessageId,
         };
diff --git a/src/BE/Controllers/Chats/UserChats/UserChatsController.cs b/src/BE/Controllers/Chats/UserChats/UserChatsController.cs
index 576a7f9e..d47c374a 100644
--- a/src/BE/Controllers/Chats/UserChats/UserChatsController.cs
+++ b/src/BE/Controllers/Chats/UserChats/UserChatsController.cs
@@ -99,18 +99,19 @@ public async Task<ActionResult<ChatsResponse>> CreateChat([FromBody] CreateChatR
         };
 
         Chat? lastChat = await db.Chats
-            .Include(x => x.ChatSpans)
+            .Include(x => x.ChatSpans.OrderBy(x => x.SpanId))
             .Where(x => x.UserId == currentUser.Id && !x.IsDeleted && x.ChatSpans.Any())
             .OrderByDescending(x => x.CreatedAt)
             .FirstOrDefaultAsync(cancellationToken);
         if (lastChat != null && lastChat.ChatSpans.All(cs => validModels.ContainsKey(cs.ModelId)))
         {
-            chat.ChatSpans = lastChat.ChatSpans.Select(cs => new ChatSpan()
+            chat.ChatSpans = lastChat.ChatSpans.Select((cs, i) => new ChatSpan()
             {
                 ModelId = cs.ModelId,
                 Model = validModels[cs.ModelId].Model,
                 EnableSearch = cs.EnableSearch,
                 Temperature = cs.Temperature,
+                SpanId = (byte)i,
             }).ToList();
         }
         else
diff --git a/src/BE/DB/ChatSpan.cs b/src/BE/DB/ChatSpan.cs
index 3d1bc8bd..301a0693 100644
--- a/src/BE/DB/ChatSpan.cs
+++ b/src/BE/DB/ChatSpan.cs
@@ -27,9 +27,6 @@ public partial class ChatSpan
     [InverseProperty("ChatSpans")]
     public virtual Chat Chat { get; set; } = null!;
 
-    [InverseProperty("ChatSpan")]
-    public virtual ICollection<Message> Messages { get; set; } = new List<Message>();
-
     [ForeignKey("ModelId")]
     [InverseProperty("ChatSpans")]
     public virtual Model Model { get; set; } = null!;
diff --git a/src/BE/DB/ChatsDB.cs b/src/BE/DB/ChatsDB.cs
index 76b377e2..1c4545ba 100644
--- a/src/BE/DB/ChatsDB.cs
+++ b/src/BE/DB/ChatsDB.cs
@@ -209,8 +209,6 @@ protected override void OnModelCreating(ModelBuilder modelBuilder)
             entity.HasOne(d => d.Parent).WithMany(p => p.InverseParent).HasConstraintName("FK_Message_ParentMessage");
 
             entity.HasOne(d => d.Usage).WithOne(p => p.Message).HasConstraintName("FK_Message_UserModelUsage");
-
-            entity.HasOne(d => d.ChatSpan).WithMany(p => p.Messages).HasConstraintName("FK_Message_ChatSpan");
         });
 
         modelBuilder.Entity<MessageContent>(entity =>
diff --git a/src/BE/DB/Message.cs b/src/BE/DB/Message.cs
index 05adca35..e3402fe2 100644
--- a/src/BE/DB/Message.cs
+++ b/src/BE/DB/Message.cs
@@ -34,10 +34,6 @@ public partial class Message
     [InverseProperty("Messages")]
     public virtual ChatRole ChatRole { get; set; } = null!;
 
-    [ForeignKey("ChatId, SpanId")]
-    [InverseProperty("Messages")]
-    public virtual ChatSpan? ChatSpan { get; set; }
-
     [InverseProperty("Parent")]
     public virtual ICollection<Message> InverseParent { get; set; } = new List<Message>();
 
diff --git a/src/scripts/db-migration/2024/20241217-chat-span.sql b/src/scripts/db-migration/2024/20241217-chat-span.sql
index d6b85943..3a4effd3 100644
--- a/src/scripts/db-migration/2024/20241217-chat-span.sql
+++ b/src/scripts/db-migration/2024/20241217-chat-span.sql
@@ -865,4 +865,5 @@ ALTER TABLE dbo.MessageContent SET (LOCK_ESCALATION = TABLE)
 GO
 COMMIT
 
-update [Message] set [SpanId] = null where ChatRoleId = 2;
\ No newline at end of file
+update [Message] set [SpanId] = null where ChatRoleId = 2;
+ALTER TABLE dbo.Message DROP CONSTRAINT FK_Message_ChatSpan;
\ No newline at end of file