feat: Prefer references to fully downloaded messages for chat assignment (#5645)

src/chat.rs

@@ -4620,7 +4620,7 @@ impl Context {
                     .0
             }
             SyncId::Msgids(msgids) => {
-                let msg = message::get_latest_by_rfc724_mids(self, msgids)
+                let msg = message::get_by_rfc724_mids(self, msgids)
                     .await?
                     .with_context(|| format!("No message found for Message-IDs {msgids:?}"))?;
                 ChatId::lookup_by_message(&msg)

src/message.rs

@@ -1921,21 +1921,30 @@ pub(crate) async fn rfc724_mid_exists_and(
     Ok(res)
 }
 
-/// Given a list of Message-IDs, returns the latest message found in the database.
+/// Given a list of Message-IDs, returns the most relevant message found in the database.
 ///
+/// Relevance here is `(download_state == Done, index)`, where `index` is an index of Message-ID in
+/// `mids`. This means Message-IDs should be ordered from the least late to the latest one (like in
+/// the References header).
 /// Only messages that are not in the trash chat are considered.
-pub(crate) async fn get_latest_by_rfc724_mids(
+pub(crate) async fn get_by_rfc724_mids(
     context: &Context,
     mids: &[String],
 ) -> Result<Option<Message>> {
+    let mut latest = None;
     for id in mids.iter().rev() {
-        if let Some((msg_id, _)) = rfc724_mid_exists(context, id).await? {
-            if let Some(msg) = Message::load_from_db_optional(context, msg_id).await? {
-                return Ok(Some(msg));
-            }
+        let Some((msg_id, _)) = rfc724_mid_exists(context, id).await? else {
+            continue;
+        };
+        let Some(msg) = Message::load_from_db_optional(context, msg_id).await? else {
+            continue;
+        };
+        if msg.download_state == DownloadState::Done {
+            return Ok(Some(msg));
         }
+        latest.get_or_insert(msg);
     }
-    Ok(None)
+    Ok(latest)
 }
 
 /// How a message is primarily displayed.

src/receive_imf.rs

@@ -489,7 +489,9 @@ pub(crate) async fn receive_imf_inner(
             can_info_msg = false;
             Some(Message::load_from_db(context, insert_msg_id).await?)
         } else if let Some(field) = mime_parser.get_header(HeaderDef::InReplyTo) {
-            if let Some(instance) = get_rfc724_mid_in_list(context, field).await? {
+            if let Some(instance) =
+                message::get_by_rfc724_mids(context, &parse_message_ids(field)).await?
+            {
                 can_info_msg = instance.download_state() == DownloadState::Done;
                 Some(instance)
             } else {
@@ -707,9 +709,13 @@ async fn add_parts(
         better_msg = Some(stock_str::msg_location_enabled_by(context, from_id).await);
     }
 
-    let parent = get_parent_message(context, mime_parser)
-        .await?
-        .filter(|p| Some(p.id) != replace_msg_id);
+    let parent = get_parent_message(
+        context,
+        mime_parser.get_header(HeaderDef::References),
+        mime_parser.get_header(HeaderDef::InReplyTo),
+    )
+    .await?
+    .filter(|p| Some(p.id) != replace_msg_id);
 
     let is_dc_message = if mime_parser.has_chat_version() {
         MessengerMessage::Yes
@@ -2792,53 +2798,35 @@ async fn get_previous_message(
     Ok(None)
 }
 
-/// Given a list of Message-IDs, returns the latest message found in the database.
-///
-/// Only messages that are not in the trash chat are considered.
-async fn get_rfc724_mid_in_list(context: &Context, mid_list: &str) -> Result<Option<Message>> {
-    message::get_latest_by_rfc724_mids(context, &parse_message_ids(mid_list)).await
-}
-
 /// Returns the last message referenced from References: header found in the database.
 ///
 /// If none found, tries In-Reply-To: as a fallback for classic MUAs that don't set the
 /// References: header.
 async fn get_parent_message(
     context: &Context,
-    mime_parser: &MimeMessage,
+    references: Option<&str>,
+    in_reply_to: Option<&str>,
 ) -> Result<Option<Message>> {
-    if let Some(field) = mime_parser.get_header(HeaderDef::References) {
-        if let Some(msg) = get_rfc724_mid_in_list(context, field).await? {
-            return Ok(Some(msg));
-        }
+    let mut mids = Vec::new();
+    if let Some(field) = in_reply_to {
+        mids = parse_message_ids(field);
     }
-
-    if let Some(field) = mime_parser.get_header(HeaderDef::InReplyTo) {
-        if let Some(msg) = get_rfc724_mid_in_list(context, field).await? {
-            return Ok(Some(msg));
-        }
+    if let Some(field) = references {
+        mids.append(&mut parse_message_ids(field));
     }
-
-    Ok(None)
+    message::get_by_rfc724_mids(context, &mids).await
 }
 
 pub(crate) async fn get_prefetch_parent_message(
     context: &Context,
     headers: &[mailparse::MailHeader<'_>],
 ) -> Result<Option<Message>> {
-    if let Some(field) = headers.get_header_value(HeaderDef::References) {
-        if let Some(msg) = get_rfc724_mid_in_list(context, &field).await? {
-            return Ok(Some(msg));
-        }
-    }
-
-    if let Some(field) = headers.get_header_value(HeaderDef::InReplyTo) {
-        if let Some(msg) = get_rfc724_mid_in_list(context, &field).await? {
-            return Ok(Some(msg));
-        }
-    }
-
-    Ok(None)
+    get_parent_message(
+        context,
+        headers.get_header_value(HeaderDef::References).as_deref(),
+        headers.get_header_value(HeaderDef::InReplyTo).as_deref(),
+    )
+    .await
 }
 
 /// Looks up contact IDs from the database given the list of recipients.

src/receive_imf/tests.rs

@@ -2887,6 +2887,18 @@ async fn test_incoming_contact_request() -> Result<()> {
     }
 }
 
+async fn get_parent_message(
+    context: &Context,
+    mime_parser: &MimeMessage,
+) -> Result<Option<Message>> {
+    super::get_parent_message(
+        context,
+        mime_parser.get_header(HeaderDef::References),
+        mime_parser.get_header(HeaderDef::InReplyTo),
+    )
+    .await
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn test_get_parent_message() -> Result<()> {
     let t = TestContext::new_alice().await;
@@ -4624,6 +4636,50 @@ async fn test_references() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_prefer_references_to_downloaded_msgs() -> Result<()> {
+    let mut tcm = TestContextManager::new();
+    let alice = &tcm.alice().await;
+    let bob = &tcm.bob().await;
+    bob.set_config(Config::DownloadLimit, Some("1")).await?;
+    let fiona = &tcm.fiona().await;
+    let alice_bob_id = tcm.send_recv(bob, alice, "hi").await.from_id;
+    let alice_fiona_id = tcm.send_recv(fiona, alice, "hi").await.from_id;
+    let alice_chat_id = create_group_chat(alice, ProtectionStatus::Unprotected, "Group").await?;
+    add_contact_to_chat(alice, alice_chat_id, alice_bob_id).await?;
+    // W/o fiona the test doesn't work -- the last message is assigned to the 1:1 chat due to
+    // `is_probably_private_reply()`.
+    add_contact_to_chat(alice, alice_chat_id, alice_fiona_id).await?;
+    let sent = alice.send_text(alice_chat_id, "Hi").await;
+    let received = bob.recv_msg(&sent).await;
+    assert_eq!(received.download_state, DownloadState::Done);
+    let bob_chat_id = received.chat_id;
+
+    let file_bytes = include_bytes!("../../test-data/image/screenshot.gif");
+    let mut msg = Message::new(Viewtype::File);
+    msg.set_file_from_bytes(alice, "file", file_bytes, None)
+        .await?;
+    let mut sent = alice.send_msg(alice_chat_id, &mut msg).await;
+    sent.payload = sent
+        .payload
+        .replace("References:", "X-Microsoft-Original-References:")
+        .replace("In-Reply-To:", "X-Microsoft-Original-In-Reply-To:");
+    let received = bob.recv_msg(&sent).await;
+    assert_eq!(received.download_state, DownloadState::Available);
+    assert_ne!(received.chat_id, bob_chat_id);
+    assert_eq!(received.chat_id, bob.get_chat(alice).await.id);
+
+    let mut msg = Message::new(Viewtype::File);
+    msg.set_file_from_bytes(alice, "file", file_bytes, None)
+        .await?;
+    let sent = alice.send_msg(alice_chat_id, &mut msg).await;
+    let received = bob.recv_msg(&sent).await;
+    assert_eq!(received.download_state, DownloadState::Available);
+    assert_eq!(received.chat_id, bob_chat_id);
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn test_list_from() -> Result<()> {
     let t = &TestContext::new_alice().await;