Fix reaction response resolution

crates/matrix-sdk-ui/src/timeline/inner.rs

@@ -583,8 +583,6 @@ impl<P: RoomDataProvider> TimelineInner<P> {
         let user_id = self.room_data_provider.own_user_id();
         let annotation_key: AnnotationKey = annotation.into();
 
-        state.in_flight_reaction.remove(&annotation_key);
-
         let reaction_state = state
             .reaction_state
             .get(&AnnotationKey::from(annotation))
@@ -593,14 +591,22 @@ impl<P: RoomDataProvider> TimelineInner<P> {
         let follow_up_action = match (result, reaction_state) {
             (ReactionToggleResult::AddSuccess { event_id, .. }, ReactionState::Redacting(_)) => {
                 // A reaction was added successfully but we've been requested to undo it
+                state
+                    .in_flight_reaction
+                    .insert(annotation_key, ReactionState::Redacting(Some(event_id.to_owned())));
                 ReactionAction::RedactRemote(event_id.to_owned())
             }
             (ReactionToggleResult::RedactSuccess { .. }, ReactionState::Sending(txn_id)) => {
                 // A reaction was was redacted successfully but we've been requested to undo it
-                ReactionAction::SendRemote(txn_id.to_owned())
+                let txn_id = txn_id.to_owned();
+                state
+                    .in_flight_reaction
+                    .insert(annotation_key, ReactionState::Sending(txn_id.clone()));
+                ReactionAction::SendRemote(txn_id)
             }
             _ => {
                 // We're done, so also update the timeline
+                state.in_flight_reaction.remove(&annotation_key);
                 state.reaction_state.remove(&annotation_key);
                 update_timeline_reaction(&mut state, user_id, annotation, result)?;
 

crates/matrix-sdk-ui/src/timeline/tests/reactions.rs

@@ -148,6 +148,57 @@ async fn redact_reaction_from_non_existent_event() {
     assert_no_more_updates(&mut stream).await;
 }
 
+#[async_test]
+async fn toggle_during_request_resolves_new_action() {
+    let timeline = TestTimeline::new();
+    let mut stream = timeline.subscribe().await;
+    let (msg_id, msg_pos) = send_first_message(&timeline, &mut stream).await;
+    let reaction = create_reaction(&msg_id);
+
+    // Add a reaction
+    let action = timeline.toggle_reaction_local(&reaction).await.unwrap();
+    let txn_id = assert_matches!(action, ReactionAction::SendRemote(txn_id) => txn_id);
+    assert_reaction_is_added(&mut stream, &msg_id, msg_pos).await;
+
+    // Toggle before response is received
+    let action = timeline.toggle_reaction_local(&reaction).await.unwrap();
+    assert_matches!(action, ReactionAction::None);
+    assert_reactions_are_removed(&mut stream, &msg_id, msg_pos).await;
+
+    // Receive response and resolve to a redaction action
+    let event_id = EventId::new(server_name!("example.org")); // non existent event
+    let action = timeline
+        .handle_reaction_response(&reaction, &ReactionToggleResult::AddSuccess { txn_id, event_id })
+        .await
+        .unwrap();
+    assert_matches!(action, ReactionAction::RedactRemote(event_id) => event_id);
+    assert_no_more_updates(&mut stream).await;
+
+    // Toggle before response is received
+    let action = timeline.toggle_reaction_local(&reaction).await.unwrap();
+    assert_matches!(action, ReactionAction::None);
+    assert_reaction_is_added(&mut stream, &msg_id, msg_pos).await;
+
+    // Receive response and resolve to a send action
+    let action = timeline
+        .handle_reaction_response(&reaction, &ReactionToggleResult::RedactSuccess)
+        .await
+        .unwrap();
+    let txn_id = assert_matches!(action, ReactionAction::SendRemote(txn_id) => txn_id);
+    assert_no_more_updates(&mut stream).await;
+
+    // Receive response and resolve to no new action
+    let event_id = EventId::new(server_name!("example.org")); // non existent event
+    let action = timeline
+        .handle_reaction_response(&reaction, &ReactionToggleResult::AddSuccess { txn_id, event_id })
+        .await
+        .unwrap();
+    assert_matches!(action, ReactionAction::None);
+    assert_reaction_is_added(&mut stream, &msg_id, msg_pos).await;
+
+    assert_no_more_updates(&mut stream).await;
+}
+
 fn create_reaction(related_message_id: &EventId) -> Annotation {
     let reaction_key = REACTION_KEY.to_owned();
     let msg_id = related_message_id.to_owned();
@@ -189,6 +240,17 @@ async fn assert_reaction_is_updated(
     assert_eq!(reaction_event_id, expected_event_id.map(|it| it.to_owned()).as_ref());
 }
 
+async fn assert_reaction_is_added(
+    stream: &mut (impl Stream<Item = VectorDiff<Arc<TimelineItem>>> + Unpin),
+    related_to: &EventId,
+    message_position: usize,
+) {
+    let own_user_id = &ALICE;
+    let event = assert_event_is_updated(stream, related_to, message_position).await;
+    let reactions = event.reactions().get(&REACTION_KEY.to_owned()).unwrap();
+    assert!(reactions.by_sender(own_user_id).next().is_some());
+}
+
 async fn assert_reactions_are_removed(
     stream: &mut (impl Stream<Item = VectorDiff<Arc<TimelineItem>>> + Unpin),
     related_to: &EventId,