Stop VF hydrator from ignoring original-tweet verdicts on retweets

home-mixer/candidate_hydrators/vf_candidate_hydrator.rs

@@ -111,7 +111,7 @@ impl Hydrator<ScoredPostsQuery, PostCandidate> for VFCandidateHydrator {
 
         let mut hydrated_candidates = Vec::with_capacity(candidates.len());
         for candidate in candidates {
-            let primary_result = all_results.get(&candidate.tweet_id);
+            let primary_result = all_results.get(&primary_vf_id(candidate));
             let (visibility_action, visibility_reason) = visibility_fields(primary_result);
 
             let drop_ancillary = should_drop_ancillary(candidate, &all_results);
@@ -137,6 +137,10 @@ impl Hydrator<ScoredPostsQuery, PostCandidate> for VFCandidateHydrator {
     }
 }
 
+pub(crate) fn primary_vf_id(candidate: &PostCandidate) -> u64 {
+    candidate.retweeted_tweet_id.unwrap_or(candidate.tweet_id)
+}
+
 pub(crate) fn should_drop_ancillary(
     candidate: &PostCandidate,
     vf_results: &HashMap<u64, Result<TweetVisibility>>,
@@ -318,3 +322,151 @@ mod tests {
         assert!(!should_drop_ancillary(&candidate, &results));
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use xai_safety_label_store::types::SafetyLabelMap;
+    use xai_visibility_filtering::models::{
+        Action, DropReason, SafetyResult, SafetyResultReason,
+    };
+
+    struct MapClient {
+        results: HashMap<u64, FilteredReason>,
+    }
+
+    fn vis(reason: FilteredReason) -> TweetVisibility {
+        TweetVisibility {
+            reason: Some(reason),
+            safety_labels: Ok(SafetyLabelMap::default()),
+        }
+    }
+
+    fn interstitial() -> FilteredReason {
+        FilteredReason::SafetyResult(SafetyResult {
+            reason: Some(SafetyResultReason::NsfwHighPrecision),
+            action: Action::Interstitial,
+        })
+    }
+
+    fn allow() -> FilteredReason {
+        FilteredReason::SafetyResult(SafetyResult {
+            reason: None,
+            action: Action::Allow,
+        })
+    }
+
+    fn drop_reason() -> FilteredReason {
+        FilteredReason::SafetyResult(SafetyResult {
+            reason: Some(SafetyResultReason::NsfwHighPrecision),
+            action: Action::Drop(DropReason {}),
+        })
+    }
+
+    #[async_trait]
+    impl VfClient for MapClient {
+        async fn get_result(
+            &self,
+            post_ids: Vec<u64>,
+            _safety_level: SafetyLevel,
+            _for_user_id: u64,
+            _context: Option<TwitterContextViewer>,
+        ) -> HashMap<u64, Result<TweetVisibility>> {
+            post_ids
+                .into_iter()
+                .filter_map(|id| {
+                    self.results
+                        .get(&id)
+                        .cloned()
+                        .map(|reason| (id, Ok(vis(reason))))
+                })
+                .collect()
+        }
+    }
+
+    async fn hydrate(
+        results: HashMap<u64, FilteredReason>,
+        candidates: &[PostCandidate],
+    ) -> Vec<std::result::Result<PostCandidate, String>> {
+        let client = Arc::new(MapClient { results });
+        VFCandidateHydrator::new(client.clone(), client)
+            .await
+            .hydrate(&ScoredPostsQuery::default(), candidates)
+            .await
+    }
+
+    #[tokio::test]
+    async fn native_post_still_uses_its_own_id() {
+        let results = hydrate(
+            HashMap::from([(20, interstitial())]),
+            &[PostCandidate {
+                tweet_id: 20,
+                in_network: Some(true),
+                ..Default::default()
+            }],
+        )
+        .await;
+        let hydrated = results[0].as_ref().unwrap();
+        assert!(matches!(
+            hydrated.visibility_reason,
+            Some(FilteredReason::SafetyResult(ref s)) if s.action == Action::Interstitial
+        ));
+        assert_eq!(hydrated.drop_ancillary_posts, Some(false));
+    }
+
+    #[tokio::test]
+    async fn retweet_uses_original_interstitial() {
+        let results = hydrate(
+            HashMap::from([(10, allow()), (20, interstitial())]),
+            &[PostCandidate {
+                tweet_id: 10,
+                retweeted_tweet_id: Some(20),
+                in_network: Some(true),
+                ..Default::default()
+            }],
+        )
+        .await;
+        let hydrated = results[0].as_ref().unwrap();
+        assert!(matches!(
+            hydrated.visibility_reason,
+            Some(FilteredReason::SafetyResult(ref s)) if s.action == Action::Interstitial
+        ));
+        assert_eq!(hydrated.drop_ancillary_posts, Some(false));
+    }
+
+    #[tokio::test]
+    async fn wrapper_interstitial_is_not_the_primary() {
+        let results = hydrate(
+            HashMap::from([(10, interstitial())]),
+            &[PostCandidate {
+                tweet_id: 10,
+                retweeted_tweet_id: Some(20),
+                in_network: Some(true),
+                ..Default::default()
+            }],
+        )
+        .await;
+        let hydrated = results[0].as_ref().unwrap();
+        assert_eq!(hydrated.visibility_reason, None);
+    }
+
+    #[tokio::test]
+    async fn original_drop_still_ancillary_drops() {
+        let results = hydrate(
+            HashMap::from([(10, allow()), (20, drop_reason())]),
+            &[PostCandidate {
+                tweet_id: 10,
+                retweeted_tweet_id: Some(20),
+                in_network: Some(true),
+                ..Default::default()
+            }],
+        )
+        .await;
+        let hydrated = results[0].as_ref().unwrap();
+        assert!(matches!(
+            hydrated.visibility_reason,
+            Some(FilteredReason::SafetyResult(ref s)) if matches!(s.action, Action::Drop(_))
+        ));
+        assert_eq!(hydrated.drop_ancillary_posts, Some(true));
+    }
+}

home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs

@@ -1,4 +1,6 @@
-use crate::candidate_hydrators::vf_candidate_hydrator::{should_drop_ancillary, visibility_fields};
+use crate::candidate_hydrators::vf_candidate_hydrator::{
+    primary_vf_id, should_drop_ancillary, visibility_fields,
+};
 use crate::models::candidate::PostCandidate;
 use crate::models::query::ScoredPostsQuery;
 use crate::params::EnableXaiVfClient;
@@ -66,7 +68,7 @@ impl Hydrator<ScoredPostsQuery, PostCandidate> for VFFollowingCandidateHydrator
 
         let mut hydrated_candidates = Vec::with_capacity(candidates.len());
         for candidate in candidates {
-            let primary_result = all_results.get(&candidate.tweet_id);
+            let primary_result = all_results.get(&primary_vf_id(candidate));
             let (visibility_action, visibility_reason) = visibility_fields(primary_result);
 
             let drop_ancillary = should_drop_ancillary(candidate, &all_results);
@@ -91,3 +93,101 @@ impl Hydrator<ScoredPostsQuery, PostCandidate> for VFFollowingCandidateHydrator
         candidate.drop_ancillary_posts = hydrated.drop_ancillary_posts;
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use xai_safety_label_store::types::SafetyLabelMap;
+    use xai_twittercontext_proto::TwitterContextViewer;
+    use xai_visibility_filtering::models::{Action, SafetyResult, SafetyResultReason};
+    use xai_visibility_filtering::vf_client::{SafetyLevel, TweetVisibility};
+
+    struct MapClient {
+        results: HashMap<u64, FilteredReason>,
+    }
+
+    fn vis(reason: FilteredReason) -> TweetVisibility {
+        TweetVisibility {
+            reason: Some(reason),
+            safety_labels: Ok(SafetyLabelMap::default()),
+        }
+    }
+
+    fn interstitial() -> FilteredReason {
+        FilteredReason::SafetyResult(SafetyResult {
+            reason: Some(SafetyResultReason::NsfwHighPrecision),
+            action: Action::Interstitial,
+        })
+    }
+
+    fn allow() -> FilteredReason {
+        FilteredReason::SafetyResult(SafetyResult {
+            reason: None,
+            action: Action::Allow,
+        })
+    }
+
+    #[async_trait]
+    impl VfClient for MapClient {
+        async fn get_result(
+            &self,
+            post_ids: Vec<u64>,
+            _safety_level: SafetyLevel,
+            _for_user_id: u64,
+            _context: Option<TwitterContextViewer>,
+        ) -> HashMap<u64, Result<TweetVisibility>> {
+            post_ids
+                .into_iter()
+                .filter_map(|id| {
+                    self.results
+                        .get(&id)
+                        .cloned()
+                        .map(|reason| (id, Ok(vis(reason))))
+                })
+                .collect()
+        }
+    }
+
+    fn hydrator(results: HashMap<u64, FilteredReason>) -> VFFollowingCandidateHydrator {
+        let client = Arc::new(MapClient { results });
+        VFFollowingCandidateHydrator::new(client.clone(), client)
+    }
+
+    #[tokio::test]
+    async fn native_post_still_uses_its_own_id() {
+        let results = hydrator(HashMap::from([(20, interstitial())]))
+            .hydrate(
+                &ScoredPostsQuery::default(),
+                &[PostCandidate {
+                    tweet_id: 20,
+                    ..Default::default()
+                }],
+            )
+            .await;
+        let hydrated = results[0].as_ref().unwrap();
+        assert!(matches!(
+            hydrated.visibility_reason,
+            Some(FilteredReason::SafetyResult(ref s)) if s.action == Action::Interstitial
+        ));
+    }
+
+    #[tokio::test]
+    async fn retweet_uses_original_interstitial() {
+        let results = hydrator(HashMap::from([(10, allow()), (20, interstitial())]))
+            .hydrate(
+                &ScoredPostsQuery::default(),
+                &[PostCandidate {
+                    tweet_id: 10,
+                    retweeted_tweet_id: Some(20),
+                    ..Default::default()
+                }],
+            )
+            .await;
+        let hydrated = results[0].as_ref().unwrap();
+        assert!(matches!(
+            hydrated.visibility_reason,
+            Some(FilteredReason::SafetyResult(ref s)) if s.action == Action::Interstitial
+        ));
+        assert_eq!(hydrated.drop_ancillary_posts, Some(false));
+    }
+}