fix: preserve preview search associations

steipete · omarshahine · steipete · commit 1c8bf04bf9ef · 2026-07-18T15:37:50.000-07:00
Co-authored-by: Omar Shahine &lt;10343873+omarshahine@users.noreply.github.com&gt;
diff --git a/Sources/IMsgCore/MessageStore+MessageConstruction.swift b/Sources/IMsgCore/MessageStore+MessageConstruction.swift
@@ -128,7 +128,7 @@ extension MessageStore {
       .mapValues { Set($0.map(\.chatID)) }
     let selection = MessageRowSelection(store: self, includeChatID: true)
     var previews: [Message] = []
-    var seenPreviewRowIDs = Set<Int64>()
+    var seenPreviewChatIDsByRowID: [Int64: Set<Int64>] = [:]
     var parentCache: ReplyParentCache = [:]
     var pollOptionCache = PollOptionTextCache()
     let maximumDateWindowsPerQuery = 200
@@ -185,7 +185,9 @@ extension MessageStore {
         else {
           continue
         }
-        guard seenPreviewRowIDs.insert(decoded.rowID).inserted else { continue }
+        guard
+          seenPreviewChatIDsByRowID[decoded.rowID, default: []].insert(preview.chatID).inserted
+        else { continue }
         previews.append(preview)
       }
     }
diff --git a/Sources/IMsgCore/MessageStore+Search.swift b/Sources/IMsgCore/MessageStore+Search.swift
@@ -78,13 +78,17 @@ extension MessageStore {
         }
         let queriedMessageCount = messages.count
         if let firstTextRowID = messages.filter({ !isURLPreviewBalloon($0) }).map(\.rowID).min() {
-          let existingRowIDs = Set(messages.map(\.rowID))
+          let existingChatIDsByRowID = Dictionary(grouping: messages, by: \.rowID)
+            .mapValues { Set($0.map(\.chatID)) }
           let linkedPreviews = try linkedURLPreviewLookahead(
             afterRowID: firstTextRowID,
             candidates: messages,
             db: db
           )
-          messages.append(contentsOf: linkedPreviews.filter { !existingRowIDs.contains($0.rowID) })
+          messages.append(
+            contentsOf: linkedPreviews.filter {
+              existingChatIDsByRowID[$0.rowID]?.contains($0.chatID) != true
+            })
         }
         var usedFallbackReplacement = false
         let coalesced = try coalesceURLPreviewMessages(
diff --git a/Tests/IMsgCoreTests/MessageStoreURLPreviewRoutingTests.swift b/Tests/IMsgCoreTests/MessageStoreURLPreviewRoutingTests.swift
@@ -326,3 +326,37 @@ func searchLookaheadKeepsCandidateChatAssociation() throws {
   #expect(messages.first?.chatID == 1)
   #expect(messages.first?.urlPreview == nil)
 }
+
+@Test
+func searchLookaheadCoalescesEveryMultiChatAssociation() throws {
+  let db = try makeURLPreviewTestDB()
+  let now = Date()
+  try db.run("INSERT INTO handle(ROWID, id) VALUES (1, '+123')")
+  try insertURLPreviewTestMessage(
+    db,
+    rowID: 1,
+    chatID: 1,
+    text: "Check this out",
+    guid: "text-guid",
+    date: now
+  )
+  try db.run("INSERT INTO chat_message_join(chat_id, message_id) VALUES (2, 1)")
+  try insertURLPreviewTestMessage(
+    db,
+    rowID: 2,
+    chatID: 1,
+    text: "https://example.com",
+    guid: "preview-guid",
+    replyToGUID: "text-guid",
+    balloonBundleID: MessageStore.urlPreviewBalloonBundleID,
+    date: now.addingTimeInterval(1)
+  )
+  try db.run("INSERT INTO chat_message_join(chat_id, message_id) VALUES (2, 2)")
+
+  let store = try MessageStore(connection: db, path: ":memory:")
+  let messages = try store.searchMessages(query: "Check this", match: "contains", limit: 10)
+
+  #expect(messages.count == 2)
+  #expect(Set(messages.map(\.chatID)) == Set([1, 2]))
+  #expect(messages.allSatisfy { $0.urlPreview?.rowID == 2 })
+}

Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,7 @@ extension MessageStore {`
`128`	`128`	`.mapValues { Set($0.map(\.chatID)) }`
`129`	`129`	`let selection = MessageRowSelection(store: self, includeChatID: true)`
`130`	`130`	`var previews: [Message] = []`
`131`		`- var seenPreviewRowIDs = Set<Int64>()`
	`131`	`+ var seenPreviewChatIDsByRowID: [Int64: Set<Int64>] = [:]`
`132`	`132`	`var parentCache: ReplyParentCache = [:]`
`133`	`133`	`var pollOptionCache = PollOptionTextCache()`
`134`	`134`	`let maximumDateWindowsPerQuery = 200`
`@@ -185,7 +185,9 @@ extension MessageStore {`
`185`	`185`	`else {`
`186`	`186`	`continue`
`187`	`187`	`}`
`188`		`- guard seenPreviewRowIDs.insert(decoded.rowID).inserted else { continue }`
	`188`	`+ guard`
	`189`	`+ seenPreviewChatIDsByRowID[decoded.rowID, default: []].insert(preview.chatID).inserted`
	`190`	`+ else { continue }`
`189`	`191`	`previews.append(preview)`
`190`	`192`	`}`
`191`	`193`	`}`