atagen

macltok

public
Star0
Watch0
722b4c3d0a09ca2aa7b242170395338933a65d52
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/mentions/service.rsRust1593 lines61.8 KB
1
use std::{
2
    collections::{BTreeMap, BTreeSet},
3
    sync::{
4
        Arc,
5
        atomic::{AtomicBool, Ordering},
6
    },
7
    time::Duration,
8
};
9
10
use futures_util::{StreamExt, stream::FuturesUnordered};
11
use matrix_sdk::{
12
    Client, Room, RoomState,
13
    event_cache::RedecryptorReport,
14
    sliding_sync::{ResponseCommitBatch, ResponseCommitHook},
15
};
16
use ruma::{
17
    OwnedRoomId, OwnedUserId, RoomId,
18
    events::{
19
        AnySyncTimelineEvent, ignored_user_list::IgnoredUserListEventContent,
20
        room::encrypted::OriginalSyncRoomEncryptedEvent,
21
    },
22
    serde::Raw,
23
};
24
use serde_json::Value;
25
use tokio::sync::{mpsc, oneshot, watch};
26
27
use super::state::{Candidate, Checkpoint, Data, SavedEvent, SyncState, Window, id, invalid, now};
28
use crate::{
29
    CommandOutput, EngineConfig, Error, ErrorKind, Event, MentionCaptureState, MentionEntry,
30
    MentionGapReason, MentionKey, MentionReason, MentionsCommand, MentionsPage, MentionsUpdate,
31
    TimelineContent, TimelineItem,
32
    engine::{MediaService, MentionsStore, context, executor::EngineBoxFuture},
33
    normalise::{MediaRegistration, MediaScope, normalise_observed},
34
    transport::EventSender,
35
};
36
37
#[cfg(all(test, feature = "sqlite"))]
38
#[path = "service_tests.rs"]
39
mod tests;
40
41
const HOOK_TIMEOUT: Duration = Duration::from_secs(10);
42
const FLUSH_RETRY_MS: u64 = 60_000;
43
44
#[derive(Clone)]
45
pub(crate) struct Handle {
46
    requests: mpsc::Sender<Request>,
47
    timed_out: Arc<AtomicBool>,
48
}
49
50
enum Request {
51
    Batch(ResponseCommitBatch, oneshot::Sender<()>),
52
    Command(
53
        MentionsCommand,
54
        oneshot::Sender<Result<CommandOutput, Error>>,
55
    ),
56
    Stop(oneshot::Sender<()>),
57
}
58
59
pub(crate) struct Service {
60
    client: Client,
61
    account: OwnedUserId,
62
    media: Arc<MediaService>,
63
    /// as many media as the inbox keeps events
64
    media_scope: MediaScope,
65
    store: Option<MentionsStore>,
66
    data: Data,
67
    sync: SyncState,
68
    dirty: bool,
69
    flush_after: u64,
70
    ignored: BTreeSet<OwnedUserId>,
71
    leaving: BTreeSet<OwnedRoomId>,
72
    timed_out: Arc<AtomicBool>,
73
    config: EngineConfig,
74
    requests: mpsc::Receiver<Request>,
75
    updates: watch::Sender<MentionsPage>,
76
    recovery_budget: usize,
77
    latest_scan: bool,
78
    recovery_retry: u64,
79
}
80
81
enum Work {
82
    Capture(
83
        String,
84
        u64,
85
        Result<(crate::EventContext, Vec<MediaRegistration>, Vec<Value>), Error>,
86
    ),
87
    Candidate(String, Result<Value, Error>),
88
    History(
89
        bool,
90
        Option<String>,
91
        Result<ruma::api::client::push::get_notifications::v3::Response, Error>,
92
    ),
93
}
94
95
impl Handle {
96
    /// never fails, a mentions failure becomes a coverage gap rather than a stalled sync
97
    pub(crate) fn hook(&self) -> ResponseCommitHook {
98
        let handle = self.clone();
99
        ResponseCommitHook::new(move |batch| {
100
            let handle = handle.clone();
101
            async move {
102
                let committed = async {
103
                    let (reply, response) = oneshot::channel();
104
                    handle
105
                        .requests
106
                        .send(Request::Batch(batch, reply))
107
                        .await
108
                        .ok()?;
109
                    response.await.ok()
110
                };
111
                match matrix_sdk::timeout::timeout(committed, HOOK_TIMEOUT).await {
112
                    Ok(Some(())) => {}
113
                    Ok(None) => tracing::warn!("mentions service stopped, sync carries on"),
114
                    Err(_) => {
115
                        handle.timed_out.store(true, Ordering::Release);
116
                        tracing::warn!("mentions commit timed out, sync carries on");
117
                    }
118
                }
119
                Ok(())
120
            }
121
        })
122
    }
123
124
    pub(crate) async fn execute(&self, command: MentionsCommand) -> Result<CommandOutput, Error> {
125
        let (reply, response) = oneshot::channel();
126
        self.requests
127
            .send(Request::Command(command, reply))
128
            .await
129
            .map_err(|_| Error::interrupted())?;
130
        response.await.map_err(|_| Error::interrupted())?
131
    }
132
133
    pub(crate) async fn stop(&self) {
134
        let (reply, response) = oneshot::channel();
135
        if self.requests.send(Request::Stop(reply)).await.is_ok() {
136
            let _ = response.await;
137
        }
138
    }
139
}
140
141
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
142
async fn open_store(
143
    store: Option<Result<MentionsStore, Error>>,
144
) -> (
145
    Option<MentionsStore>,
146
    Data,
147
    SyncState,
148
    Vec<MentionGapReason>,
149
) {
150
    let mut store = match store {
151
        None => return (None, Data::default(), SyncState::default(), Vec::new()),
152
        Some(Err(error)) => {
153
            tracing::error!(%error, "mentions store unavailable, the inbox runs in memory");
154
            return (
155
                None,
156
                Data::default(),
157
                SyncState::default(),
158
                vec![MentionGapReason::StoreFailed],
159
            );
160
        }
161
        Some(Ok(store)) => store,
162
    };
163
    let loaded = (|| {
164
        let data: Data = store.read("inbox")?.unwrap_or_default();
165
        let sync: SyncState = store.read("sync")?.unwrap_or_default();
166
        if sync.inbox_revision != data.revision {
167
            return Err(Error::new(
168
                ErrorKind::Store,
169
                "mentions checkpoint does not match its inbox",
170
            ));
171
        }
172
        Ok((data, sync))
173
    })();
174
    let (data, sync) = match loaded {
175
        Ok(loaded) => loaded,
176
        Err(error) => {
177
            tracing::warn!(%error, "resetting an unreadable mentions inbox");
178
            if let Err(error) = store.reset().await {
179
                tracing::error!(%error, "mentions store unavailable, the inbox runs in memory");
180
                return (
181
                    None,
182
                    Data::default(),
183
                    SyncState::default(),
184
                    vec![MentionGapReason::StoreFailed],
185
                );
186
            }
187
            (Data::default(), SyncState::default())
188
        }
189
    };
190
    let gaps = if store.was_reset() {
191
        vec![MentionGapReason::StoreReset]
192
    } else {
193
        Vec::new()
194
    };
195
    (Some(store), data, sync, gaps)
196
}
197
198
#[cfg(not(any(feature = "sqlite", feature = "indexeddb")))]
199
async fn open_store(
200
    store: Option<Result<MentionsStore, Error>>,
201
) -> (
202
    Option<MentionsStore>,
203
    Data,
204
    SyncState,
205
    Vec<MentionGapReason>,
206
) {
207
    let gaps = match store {
208
        None => Vec::new(),
209
        Some(Err(_)) => vec![MentionGapReason::StoreFailed],
210
        Some(Ok(store)) => match store {},
211
    };
212
    (None, Data::default(), SyncState::default(), gaps)
213
}
214
215
impl Service {
216
    pub(crate) async fn new(
217
        client: Client,
218
        media: Arc<MediaService>,
219
        store: Option<Result<MentionsStore, Error>>,
220
        config: EngineConfig,
221
    ) -> Result<(Handle, Self, watch::Receiver<MentionsPage>), Error> {
222
        let account = client
223
            .user_id()
224
            .ok_or_else(|| {
225
                Error::new(
226
                    ErrorKind::Unauthenticated,
227
                    "mentions require an authenticated account",
228
                )
229
            })?
230
            .to_owned();
231
        let (store, mut data, sync, gaps) = open_store(store).await;
232
        for reason in &gaps {
233
            data.gap(None, *reason, None);
234
        }
235
        let page = data.page(&account, None, 50)?;
236
        let media_scope = MediaScope::bounded(config.max_mention_context_events);
237
        media
238
            .register(
239
                &media_scope,
240
                page.entries
241
                    .iter()
242
                    .filter_map(|entry| data.events.get(&id(&entry.key)))
243
                    .flat_map(|event| event.media.clone())
244
                    .collect(),
245
            )
246
            .await;
247
        let (updates, receiver) = watch::channel(page);
248
        let (sender, requests) = mpsc::channel(config.command_capacity.max(1));
249
        let timed_out = Arc::new(AtomicBool::new(false));
250
        let mut service = Self {
251
            client,
252
            account,
253
            media,
254
            media_scope,
255
            store,
256
            data,
257
            sync,
258
            dirty: !gaps.is_empty(),
259
            flush_after: 0,
260
            ignored: BTreeSet::new(),
261
            leaving: BTreeSet::new(),
262
            timed_out: timed_out.clone(),
263
            config,
264
            requests,
265
            updates,
266
            recovery_budget: 20,
267
            latest_scan: true,
268
            recovery_retry: 0,
269
        };
270
        service.ignored = service.ignored_users().await;
271
        Ok((
272
            Handle {
273
                requests: sender,
274
                timed_out,
275
            },
276
            service,
277
            receiver,
278
        ))
279
    }
280
281
    async fn ignored_users(&self) -> BTreeSet<OwnedUserId> {
282
        self.client
283
            .account()
284
            .account_data::<IgnoredUserListEventContent>()
285
            .await
286
            .ok()
287
            .flatten()
288
            .and_then(|raw| raw.deserialize().ok())
289
            .map(|content| content.ignored_users.into_keys().collect())
290
            .unwrap_or_default()
291
    }
292
293
    async fn publish(&mut self) {
294
        let Ok(page) = self.data.page(&self.account, None, 50) else {
295
            return;
296
        };
297
        self.media
298
            .register(
299
                &self.media_scope,
300
                page.entries
301
                    .iter()
302
                    .filter_map(|entry| self.data.events.get(&id(&entry.key)))
303
                    .flat_map(|event| event.media.clone())
304
                    .collect(),
305
            )
306
            .await;
307
        self.updates.send_replace(page);
308
    }
309
310
    /// writes the sync record, and the inbox with it when `dirty`, in one commit so a
311
    /// checkpoint never lands without the inbox changes it covers
312
    async fn flush(&mut self) -> Result<(), Error> {
313
        if self.dirty {
314
            self.advance()?;
315
        }
316
        #[cfg(any(feature = "sqlite", feature = "indexeddb"))]
317
        if let Some(store) = &mut self.store {
318
            use super::state::json_len;
319
            use super::store::Write;
320
321
            let encode = |error| {
322
                Error::with_source(
323
                    ErrorKind::Store,
324
                    "encoding the mentions inbox failed",
325
                    error,
326
                )
327
            };
328
            let sync = serde_json::to_value(&self.sync).map_err(encode)?;
329
            let mut writes = Vec::new();
330
            if self.dirty {
331
                let available = store
332
                    .available(&["inbox", "sync"])
333
                    .saturating_sub(store.sealed_bound("sync", json_len(&sync)));
334
                let mut size = store.sealed_bound("inbox", json_len(&self.data));
335
                while size > available {
336
                    let Some(freed) = self.data.evict_one() else {
337
                        break;
338
                    };
339
                    size = size.saturating_sub(freed * 3);
340
                }
341
                writes.push(Write::Set(
342
                    "inbox".to_owned(),
343
                    serde_json::to_value(&self.data).map_err(encode)?,
344
                ));
345
            }
346
            writes.push(Write::Set("sync".to_owned(), sync));
347
            store.commit(store.revision(), writes).await?;
348
        }
349
        if self.dirty {
350
            self.dirty = false;
351
            self.publish().await;
352
        }
353
        Ok(())
354
    }
355
356
    fn advance(&mut self) -> Result<(), Error> {
357
        self.data.revision = self
358
            .data
359
            .revision
360
            .checked_add(1)
361
            .ok_or_else(|| invalid("mentions revision exhausted"))?;
362
        self.data.bound(
363
            self.config.max_mentions_entries,
364
            self.config.max_mention_context_events,
365
            self.config.max_mention_candidates,
366
        );
367
        self.sync.inbox_revision = self.data.revision;
368
        Ok(())
369
    }
370
371
    async fn apply(&mut self, data: Data) -> Result<(), Error> {
372
        let dirty = std::mem::replace(&mut self.dirty, true);
373
        let previous = std::mem::replace(&mut self.data, data);
374
        let result = self.flush().await;
375
        if result.is_err() {
376
            self.data = previous;
377
            self.dirty = dirty;
378
        }
379
        result
380
    }
381
382
    async fn store_failed(&mut self, error: Error) {
383
        tracing::error!(%error, "saving the mentions inbox failed");
384
        self.data.gap(None, MentionGapReason::StoreFailed, None);
385
        self.dirty = true;
386
        self.flush_after = now().saturating_add(FLUSH_RETRY_MS);
387
        self.publish().await;
388
    }
389
390
    async fn command(&mut self, command: MentionsCommand) -> Result<CommandOutput, Error> {
391
        match command {
392
            MentionsCommand::Page { cursor, limit } => {
393
                let page = self.data.page(&self.account, cursor.as_deref(), limit)?;
394
                self.media
395
                    .register(
396
                        &self.media_scope,
397
                        page.entries
398
                            .iter()
399
                            .filter_map(|entry| self.data.events.get(&id(&entry.key)))
400
                            .flat_map(|event| event.media.clone())
401
                            .collect(),
402
                    )
403
                    .await;
404
                Ok(CommandOutput::MentionsPage(page))
405
            }
406
            MentionsCommand::Context(key) => {
407
                let (context, media) = self.data.context(&key)?;
408
                self.media.register(&self.media_scope, media).await;
409
                Ok(CommandOutput::MentionContext(context))
410
            }
411
            MentionsCommand::RetryContext(key) => {
412
                let room = self
413
                    .client
414
                    .get_room(&key.room_id)
415
                    .filter(|room| room.state() == RoomState::Joined)
416
                    .ok_or_else(|| {
417
                        Error::new(
418
                            ErrorKind::Forbidden,
419
                            "saved context cannot be refetched after leaving the room",
420
                        )
421
                    })?;
422
                drop(room);
423
                let mut next = self.data.clone();
424
                let entry = next
425
                    .entries
426
                    .get_mut(&id(&key))
427
                    .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "mention was not found"))?;
428
                entry.capture = MentionCaptureState::Pending;
429
                entry.revision = self.data.revision + 1;
430
                next.windows.entry(id(&key)).or_insert(Window {
431
                    before: Vec::new(),
432
                    target: Some(id(&key)),
433
                    after: Vec::new(),
434
                    start: None,
435
                    end: None,
436
                });
437
                self.apply(next).await?;
438
                let (context, media) = self.data.context(&key)?;
439
                self.media.register(&self.media_scope, media).await;
440
                Ok(CommandOutput::MentionContext(context))
441
            }
442
            MentionsCommand::SetRead { keys, read } => {
443
                let mut next = self.data.clone();
444
                for key in keys {
445
                    if let Some(entry) = next.entries.get_mut(&id(&key)) {
446
                        entry.read = read;
447
                        entry.revision = self.data.revision + 1;
448
                    }
449
                }
450
                self.apply(next).await?;
451
                Ok(CommandOutput::Done)
452
            }
453
            MentionsCommand::ReadThrough(boundary) => {
454
                if boundary.account != self.account || boundary.revision > self.data.revision {
455
                    return Err(invalid("invalid mentions read boundary"));
456
                }
457
                let mut next = self.data.clone();
458
                for entry in next
459
                    .entries
460
                    .values_mut()
461
                    .filter(|entry| entry.discovery_revision <= boundary.revision)
462
                {
463
                    entry.read = true;
464
                    entry.revision = self.data.revision + 1;
465
                }
466
                self.apply(next).await?;
467
                Ok(CommandOutput::Done)
468
            }
469
            MentionsCommand::Recover => {
470
                self.recovery_budget = 20;
471
                self.latest_scan = true;
472
                self.recovery_retry = 0;
473
                let mut next = self.data.clone();
474
                next.coverage.recovering = true;
475
                next.coverage.historical_supported = None;
476
                self.apply(next).await?;
477
                Ok(CommandOutput::Done)
478
            }
479
        }
480
    }
481
482
    fn observable(&self, room_id: &RoomId) -> Option<Room> {
483
        self.client
484
            .get_room(room_id)
485
            .filter(|room| room.state() == RoomState::Joined || self.leaving.contains(room_id))
486
    }
487
488
    async fn batch(&mut self, batch: ResponseCommitBatch) {
489
        let checkpoint = self.sync.checkpoints.get(&batch.connection_id).cloned();
490
        if checkpoint
491
            .as_ref()
492
            .is_some_and(|checkpoint| checkpoint.pos == batch.next_pos)
493
        {
494
            return;
495
        }
496
        self.ignored = self.ignored_users().await;
497
        self.leaving = batch.response.rooms.left.keys().cloned().collect();
498
        let mut next = self.data.clone();
499
        let mut sync = self.sync.clone();
500
        if self.timed_out.swap(false, Ordering::AcqRel) {
501
            next.gap(None, MentionGapReason::HookTimedOut, None);
502
        }
503
        // a crash after the commit but before the sdk saved `pos` replays from the
504
        // previous position, and that response is a superset of the one committed
505
        let continuous = match &checkpoint {
506
            None => batch.previous_pos.is_none(),
507
            Some(checkpoint) => {
508
                batch.previous_pos.as_ref() == Some(&checkpoint.pos)
509
                    || batch.previous_pos == checkpoint.previous
510
            }
511
        };
512
        if !continuous {
513
            next.gap(None, MentionGapReason::SyncDiscontinuity, None);
514
        }
515
        for (room_id, raw) in &batch.raw_rooms {
516
            if raw.limited && sync.observed_rooms.contains(room_id) {
517
                next.gap(
518
                    Some(room_id.clone()),
519
                    MentionGapReason::LimitedTimeline,
520
                    raw.prev_batch.clone(),
521
                );
522
            }
523
        }
524
        let room_ids: BTreeSet<_> = batch
525
            .raw_rooms
526
            .keys()
527
            .chain(batch.response.rooms.joined.keys())
528
            .chain(batch.response.rooms.left.keys())
529
            .cloned()
530
            .collect();
531
        for room_id in &room_ids {
532
            if self.observable(room_id).is_none() {
533
                continue;
534
            }
535
            sync.observed_rooms.insert(room_id.clone());
536
            let mut observed = BTreeMap::new();
537
            if let Some(raw) = batch.raw_rooms.get(room_id) {
538
                for event in &raw.timeline {
539
                    let Ok(value) = serde_json::from_str::<Value>(event.json().get()) else {
540
                        continue;
541
                    };
542
                    if let Some(event_id) = value["event_id"].as_str() {
543
                        observed.insert(event_id.to_owned(), value);
544
                    }
545
                }
546
            }
547
            let processed = batch
548
                .response
549
                .rooms
550
                .joined
551
                .get(room_id)
552
                .map(|update| &update.timeline.events)
553
                .into_iter()
554
                .chain(
555
                    batch
556
                        .response
557
                        .rooms
558
                        .left
559
                        .get(room_id)
560
                        .map(|update| &update.timeline.events),
561
                )
562
                .flatten();
563
            for event in processed {
564
                let Ok(value) = serde_json::from_str::<Value>(event.raw().json().get()) else {
565
                    continue;
566
                };
567
                if let Some(event_id) = value["event_id"].as_str() {
568
                    observed.insert(event_id.to_owned(), value);
569
                }
570
            }
571
            let mut observed: Vec<_> = observed.into_values().collect();
572
            observed.sort_by_key(|value| value["origin_server_ts"].as_u64().unwrap_or(0));
573
            for value in observed {
574
                self.observe(&mut next, room_id, value).await;
575
            }
576
        }
577
        for room_id in batch.response.rooms.left.keys() {
578
            next.gap(Some(room_id.clone()), MentionGapReason::RoomLeft, None);
579
            next.candidates
580
                .retain(|_, candidate| &candidate.room_id != room_id);
581
            let revision = next.revision + 1;
582
            for entry in next.entries.values_mut().filter(|entry| {
583
                &entry.key.room_id == room_id && entry.capture != MentionCaptureState::Unavailable
584
            }) {
585
                entry.capture = MentionCaptureState::Unavailable;
586
                entry.revision = revision;
587
            }
588
        }
589
        let reevaluate: Vec<_> = next
590
            .entries
591
            .values()
592
            .filter(|entry| room_ids.contains(&entry.key.room_id))
593
            .filter_map(|entry| {
594
                next.events
595
                    .get(&id(&entry.key))
596
                    .and_then(|event| event.raw.clone())
597
                    .map(|raw| (entry.key.room_id.clone(), raw))
598
            })
599
            .collect();
600
        for (room_id, raw) in reevaluate {
601
            self.observe(&mut next, &room_id, raw).await;
602
        }
603
        self.filter_ignored(&mut next);
604
        self.leaving.clear();
605
        sync.checkpoints.insert(
606
            batch.connection_id,
607
            Checkpoint {
608
                pos: batch.next_pos,
609
                previous: batch.previous_pos,
610
            },
611
        );
612
        let changed = next != self.data;
613
        if changed {
614
            self.data = next;
615
            self.dirty = true;
616
        }
617
        self.sync = sync;
618
        if now() < self.flush_after {
619
            if changed && self.advance().is_ok() {
620
                self.publish().await;
621
            }
622
        } else if let Err(error) = self.flush().await {
623
            self.store_failed(error).await;
624
        }
625
    }
626
627
    fn filter_ignored(&self, next: &mut Data) -> bool {
628
        let removed: Vec<_> = next
629
            .entries
630
            .iter()
631
            .filter(|(_, entry)| {
632
                matches!(&entry.event, TimelineItem::Event(event) if self.ignored.contains(&event.sender.user_id))
633
            })
634
            .map(|(key, _)| key.clone())
635
            .collect();
636
        for key in &removed {
637
            next.entries.remove(key);
638
            next.windows.remove(key);
639
        }
640
        next.collect_events();
641
        !removed.is_empty()
642
    }
643
644
    async fn observe(&self, data: &mut Data, room_id: &OwnedRoomId, mut value: Value) {
645
        let Some(event_id) = value["event_id"]
646
            .as_str()
647
            .and_then(|value| value.parse().ok())
648
        else {
649
            return;
650
        };
651
        let key = MentionKey {
652
            room_id: room_id.clone(),
653
            event_id,
654
        };
655
        let identity = id(&key);
656
        let Some(timestamp) = value["origin_server_ts"].as_u64() else {
657
            return;
658
        };
659
        data.coverage.earliest_timestamp_ms = Some(
660
            data.coverage
661
                .earliest_timestamp_ms
662
                .map_or(timestamp, |old| old.min(timestamp)),
663
        );
664
        if value["type"] == "m.room.redaction" {
665
            if let Some(target) = value["redacts"]
666
                .as_str()
667
                .or_else(|| value["content"]["redacts"].as_str())
668
                .and_then(|value| value.parse().ok())
669
            {
670
                let target = id(&MentionKey {
671
                    room_id: room_id.clone(),
672
                    event_id: target,
673
                });
674
                data.redacted.insert(target.clone(), timestamp);
675
                data.entries.remove(&target);
676
                data.windows.remove(&target);
677
                data.candidates.remove(&target);
678
                data.edits.remove(&target);
679
                data.candidates.retain(|_, candidate| {
680
                    candidate.room_id != *room_id
681
                        || candidate.raw["content"]["m.relates_to"]["event_id"].as_str()
682
                            != target.split_once('\n').map(|(_, event)| event)
683
                });
684
                if let Some(saved) = data.events.get_mut(&target) {
685
                    redact(saved);
686
                }
687
                data.touch_contexts(&target);
688
            }
689
            return;
690
        }
691
        if data.redacted.contains_key(&identity)
692
            || value["unsigned"].get("redacted_because").is_some()
693
        {
694
            data.redacted.insert(identity.clone(), timestamp);
695
            data.entries.remove(&identity);
696
            data.windows.remove(&identity);
697
            data.candidates.remove(&identity);
698
            if let Some(saved) = data.events.get_mut(&identity) {
699
                redact(saved);
700
            }
701
            data.touch_contexts(&identity);
702
            return;
703
        }
704
        if value["content"]["m.relates_to"]["rel_type"] == "m.replace" {
705
            if let Some(target) = value["content"]["m.relates_to"]["event_id"]
706
                .as_str()
707
                .and_then(|value| value.parse().ok())
708
            {
709
                let target = id(&MentionKey {
710
                    room_id: room_id.clone(),
711
                    event_id: target,
712
                });
713
                if data.redacted.contains_key(&target) {
714
                    return;
715
                }
716
                let newer = data.edits.get(&target).is_none_or(|old| {
717
                    (old["origin_server_ts"].as_u64(), old["event_id"].as_str())
718
                        < (Some(timestamp), value["event_id"].as_str())
719
                });
720
                if newer {
721
                    if let Some(original) =
722
                        data.events.get(&target).and_then(|saved| saved.raw.clone())
723
                    {
724
                        if original["sender"] != value["sender"] {
725
                            return;
726
                        }
727
                        data.edits.insert(target.clone(), value.clone());
728
                        value = edited(original, &value);
729
                        return Box::pin(self.observe(data, room_id, value)).await;
730
                    }
731
                    data.candidates
732
                        .entry(identity)
733
                        .and_modify(|candidate| candidate.raw = value.clone())
734
                        .or_insert_with(|| Candidate::new(room_id.clone(), value, 0));
735
                }
736
            }
737
            return;
738
        }
739
        if let Some(edit) = value["unsigned"]["m.relations"].get("m.replace").cloned() {
740
            if edit["sender"] == value["sender"]
741
                && edit["content"]["m.new_content"].is_object()
742
                && edit["content"]["m.relates_to"]["event_id"]
743
                    .as_str()
744
                    .is_none_or(|target| target == key.event_id.as_str())
745
                && edit["origin_server_ts"].as_u64().is_some()
746
                && data.edits.get(&identity).is_none_or(|old| {
747
                    (old["origin_server_ts"].as_u64(), old["event_id"].as_str())
748
                        < (edit["origin_server_ts"].as_u64(), edit["event_id"].as_str())
749
                })
750
            {
751
                data.edits.insert(identity.clone(), edit);
752
            }
753
        }
754
        if let Some(edit) = data.edits.get(&identity) {
755
            if value["sender"] == edit["sender"] {
756
                value = edited(value, edit);
757
            }
758
        }
759
        if let Some(edit) = value["unsigned"].get("macltok_edit").cloned() {
760
            data.edits.insert(identity.clone(), edit);
761
        }
762
        if value["type"] == "m.room.encrypted" {
763
            data.candidates
764
                .entry(identity)
765
                .or_insert_with(|| Candidate::new(room_id.clone(), value, 0));
766
            return;
767
        }
768
        let Some((item, media)) = normalise_observed(&value, Some(&self.account), &identity) else {
769
            return;
770
        };
771
        self.supplement(data, &key, &item, &media, &value);
772
        let (reasons, suppressed, unresolved) = self.classify(data, room_id, &value).await;
773
        if unresolved {
774
            data.candidates
775
                .entry(identity.clone())
776
                .or_insert_with(|| Candidate::new(room_id.clone(), value.clone(), now() + 30_000));
777
        } else {
778
            data.candidates.remove(&identity);
779
        }
780
        if reasons.is_empty() {
781
            data.entries.remove(&identity);
782
            data.windows.remove(&identity);
783
            if data.events.contains_key(&identity) || data.edits.contains_key(&identity) {
784
                data.events.insert(
785
                    identity,
786
                    SavedEvent {
787
                        item,
788
                        media,
789
                        raw: Some(value),
790
                    },
791
                );
792
            }
793
            return;
794
        }
795
        let previous = data.entries.get(&identity);
796
        let changed = previous.is_none_or(|entry| {
797
            entry.event != item
798
                || entry.reasons != reasons
799
                || entry.notification_suppressed != suppressed
800
        });
801
        let entry = MentionEntry {
802
            key,
803
            timestamp_ms: timestamp,
804
            discovered_at_ms: previous.map_or_else(now, |entry| entry.discovered_at_ms),
805
            discovery_revision: previous
806
                .map_or(data.revision + 1, |entry| entry.discovery_revision),
807
            revision: if changed {
808
                data.revision + 1
809
            } else {
810
                previous.unwrap().revision
811
            },
812
            event: item.clone(),
813
            reasons,
814
            read: previous.is_some_and(|entry| entry.read),
815
            capture: previous.map_or(MentionCaptureState::Pending, |entry| {
816
                if changed {
817
                    MentionCaptureState::Pending
818
                } else {
819
                    entry.capture.clone()
820
                }
821
            }),
822
            room_name: self.client.get_room(room_id).and_then(|room| room.name()),
823
            notification_suppressed: suppressed,
824
        };
825
        data.entries.insert(identity.clone(), entry);
826
        data.windows.entry(identity.clone()).or_insert(Window {
827
            before: Vec::new(),
828
            target: Some(identity.clone()),
829
            after: Vec::new(),
830
            start: None,
831
            end: None,
832
        });
833
        data.events.insert(
834
            identity,
835
            SavedEvent {
836
                item,
837
                media,
838
                raw: Some(value),
839
            },
840
        );
841
    }
842
843
    fn supplement(
844
        &self,
845
        data: &mut Data,
846
        key: &MentionKey,
847
        item: &TimelineItem,
848
        media: &[MediaRegistration],
849
        raw: &Value,
850
    ) {
851
        let TimelineItem::Event(event) = item else {
852
            return;
853
        };
854
        let identity = id(key);
855
        if data
856
            .events
857
            .get(&identity)
858
            .is_some_and(|saved| saved.item != *item)
859
        {
860
            data.touch_contexts(&identity);
861
        }
862
        let mut retained = data.events.contains_key(&identity);
863
        let saved_events = &data.events;
864
        for (target, window) in &mut data.windows {
865
            let Some(entry) = data.entries.get_mut(target) else {
866
                continue;
867
            };
868
            if entry.key.room_id != key.room_id
869
                || (entry.timestamp_ms, target.as_str()) >= (event.timestamp_ms, identity.as_str())
870
                || window.after.contains(&identity)
871
            {
872
                continue;
873
            }
874
            let previous = window.after.clone();
875
            window.after.push(identity.clone());
876
            window.after.sort_by_key(|key| {
877
                if key == &identity {
878
                    (event.timestamp_ms, key.clone())
879
                } else {
880
                    (
881
                        saved_events
882
                            .get(key)
883
                            .and_then(|saved| match &saved.item {
884
                                TimelineItem::Event(event) => Some(event.timestamp_ms),
885
                                _ => None,
886
                            })
887
                            .unwrap_or(0),
888
                        key.clone(),
889
                    )
890
                }
891
            });
892
            window.after.truncate(10);
893
            if window.after == previous {
894
                continue;
895
            }
896
            window.end = None;
897
            retained |= window.after.contains(&identity);
898
            if entry.capture == MentionCaptureState::Partial
899
                && window.after.len() == 10
900
                && (window.before.len() == 10 || window.start.is_none())
901
            {
902
                entry.capture = MentionCaptureState::Ready;
903
            }
904
            entry.revision = data.revision + 1;
905
        }
906
        if retained {
907
            data.events.insert(
908
                identity,
909
                SavedEvent {
910
                    item: item.clone(),
911
                    media: media.to_vec(),
912
                    raw: Some(raw.clone()),
913
                },
914
            );
915
        }
916
    }
917
918
    async fn classify(
919
        &self,
920
        data: &Data,
921
        room_id: &OwnedRoomId,
922
        value: &Value,
923
    ) -> (Vec<MentionReason>, bool, bool) {
924
        let Some(sender) = value["sender"]
925
            .as_str()
926
            .and_then(|sender| sender.parse::<OwnedUserId>().ok())
927
        else {
928
            return (vec![], false, false);
929
        };
930
        if sender == self.account
931
            || self.ignored.contains(&sender)
932
            || value["type"] != "m.room.message"
933
        {
934
            return (vec![], false, false);
935
        }
936
        let Some(room) = self.observable(room_id) else {
937
            return (vec![], false, false);
938
        };
939
        let raw: Raw<AnySyncTimelineEvent> =
940
            Raw::from_json_string(value.to_string()).expect("serialized event");
941
        let actions = room.event_push_actions(&raw).await.ok().flatten();
942
        let notify = actions.as_ref().is_some_and(|actions| {
943
            actions
944
                .iter()
945
                .any(|action| matches!(action, ruma::push::Action::Notify))
946
        });
947
        let highlight = actions.as_ref().is_some_and(|actions| actions.iter().any(|action| matches!(action, ruma::push::Action::SetTweak(ruma::push::Tweak::Highlight(value)) if value != &ruma::push::HighlightTweakValue::No)));
948
        let direct = value["content"]["m.mentions"]["user_ids"]
949
            .as_array()
950
            .is_some_and(|users| {
951
                users
952
                    .iter()
953
                    .any(|user| user.as_str() == Some(self.account.as_str()))
954
            });
955
        let mut reasons = Vec::new();
956
        if direct {
957
            reasons.push(MentionReason::Direct);
958
        }
959
        let mut unresolved = false;
960
        if let Some(reply) = value["content"]["m.relates_to"]["m.in_reply_to"]["event_id"]
961
            .as_str()
962
            .and_then(|value| value.parse().ok())
963
        {
964
            let target = data.events.get(&id(&MentionKey {
965
                room_id: room_id.clone(),
966
                event_id: reply,
967
            }));
968
            let reply_to_self = value["unsigned"]["macltok_reply_sender"].as_str() == Some(self.account.as_str()) || target.is_some_and(|target| matches!(&target.item, TimelineItem::Event(event) if event.sender.user_id == self.account));
969
            if reply_to_self && (direct || (notify && highlight)) {
970
                reasons.push(MentionReason::Reply);
971
            } else if target.is_none()
972
                && value["unsigned"]["macltok_reply_sender"].as_str().is_none()
973
                && (direct || (notify && highlight))
974
            {
975
                unresolved = true;
976
            }
977
        }
978
        if value["content"]["m.mentions"]["room"] == true {
979
            match room.power_levels().await {
980
                Ok(levels) if levels.user_can_trigger_room_notification(&sender) => {
981
                    if notify && highlight {
982
                        reasons.push(MentionReason::Room);
983
                    } else if actions.is_none() {
984
                        unresolved = true;
985
                    }
986
                }
987
                Err(_) => unresolved = true,
988
                _ => {}
989
            }
990
        }
991
        if reasons.is_empty() && notify && highlight {
992
            reasons.push(MentionReason::Keyword);
993
        }
994
        (reasons, !notify, unresolved)
995
    }
996
997
    pub(crate) async fn run(mut self) {
998
        let reports_client = self.client.clone();
999
        let reports = reports_client
1000
            .event_cache()
1001
            .subscribe_to_decryption_reports();
1002
        futures_util::pin_mut!(reports);
1003
        let mut ignores = self.client.subscribe_to_ignore_user_list_changes();
1004
        let mut work: FuturesUnordered<EngineBoxFuture<'static, Work>> = FuturesUnordered::new();
1005
        let mut captures = BTreeSet::new();
1006
        let mut retries = BTreeSet::new();
1007
        let mut history = false;
1008
        let mut next_scan = now() + 60_000;
1009
        let mut next_expiry = 0;
1010
        let mut tick = Box::pin(matrix_sdk::sleep::sleep(Duration::from_secs(1)));
1011
        loop {
1012
            self.schedule(&mut work, &mut captures, &mut retries, &mut history);
1013
            tokio::select! {
1014
                biased;
1015
                request = self.requests.recv() => match request {
1016
                    Some(Request::Stop(reply)) => {
1017
                        if self.dirty && let Err(error) = self.flush().await {
1018
                            tracing::error!(%error, "saving the mentions inbox at shutdown failed");
1019
                        }
1020
                        self.media.release(&self.media_scope).await;
1021
                        let _ = reply.send(());
1022
                        return;
1023
                    }
1024
                    Some(Request::Batch(batch, reply)) => { self.batch(batch).await; let _ = reply.send(()); }
1025
                    Some(Request::Command(command, reply)) => { let result = self.command(command).await; let _ = reply.send(result); }
1026
                    None => return,
1027
                },
1028
                Some(result) = work.next(), if !work.is_empty() => {
1029
                    let mut next = std::mem::take(&mut self.data);
1030
                    // a candidate that failed again only moved its retry time, which needs no write
1031
                    let changed = self.complete(&mut next, result, &mut captures, &mut retries, &mut history).await;
1032
                    self.data = next;
1033
                    self.dirty |= changed;
1034
                }
1035
                Some(report) = reports.next() => {
1036
                    let room = match report { Ok(RedecryptorReport::ResolvedUtds { room_id, .. }) => Some(room_id), _ => None };
1037
                    for candidate in self.data.candidates.values_mut().filter(|candidate| candidate.encrypted() && room.as_ref().is_none_or(|room| *room == candidate.room_id)) { candidate.retry_at_ms = 0; }
1038
                }
1039
                Some(_) = ignores.next() => {
1040
                    self.ignored = self.ignored_users().await;
1041
                    let mut next = std::mem::take(&mut self.data);
1042
                    self.dirty |= self.filter_ignored(&mut next);
1043
                    self.data = next;
1044
                }
1045
                () = &mut tick => {
1046
                    tick = Box::pin(matrix_sdk::sleep::sleep(Duration::from_secs(1)));
1047
                    if now() >= next_scan && !history { self.recovery_budget = 20; self.latest_scan = true; next_scan = now() + 60_000; }
1048
                    if now() >= next_expiry { self.dirty |= self.data.expire_candidates(now()); next_expiry = now() + 60_000; }
1049
                    if self.dirty && now() >= self.flush_after && let Err(error) = self.flush().await { self.store_failed(error).await; }
1050
                }
1051
            }
1052
        }
1053
    }
1054
1055
    async fn complete(
1056
        &mut self,
1057
        next: &mut Data,
1058
        result: Work,
1059
        captures: &mut BTreeSet<String>,
1060
        retries: &mut BTreeSet<String>,
1061
        history: &mut bool,
1062
    ) -> bool {
1063
        match result {
1064
            Work::Capture(key, revision, result) => {
1065
                captures.remove(&key);
1066
                let observed = result
1067
                    .as_ref()
1068
                    .ok()
1069
                    .map(|(_, _, observed)| observed.clone())
1070
                    .unwrap_or_default();
1071
                let room_id = next
1072
                    .entries
1073
                    .get(&key)
1074
                    .filter(|entry| entry.revision == revision)
1075
                    .map(|entry| entry.key.room_id.clone())
1076
                    .filter(|room_id| {
1077
                        self.client
1078
                            .get_room(room_id)
1079
                            .is_some_and(|room| room.state() == RoomState::Joined)
1080
                    });
1081
                self.captured(next, key, revision, result);
1082
                if let Some(room_id) = room_id {
1083
                    for value in observed {
1084
                        let identity = value["event_id"]
1085
                            .as_str()
1086
                            .map(|event_id| format!("{room_id}\n{event_id}"));
1087
                        if identity.as_ref().is_some_and(|key| {
1088
                            next.events
1089
                                .get(key)
1090
                                .is_some_and(|event| event.raw.is_some())
1091
                        }) {
1092
                            continue;
1093
                        }
1094
                        self.observe(next, &room_id, value).await;
1095
                    }
1096
                }
1097
            }
1098
            Work::Candidate(key, result) => {
1099
                retries.remove(&key);
1100
                if let Some(mut candidate) = next.candidates.remove(&key) {
1101
                    match result {
1102
                        Ok(value) => {
1103
                            self.observe(next, &candidate.room_id, value).await;
1104
                            if let Some(unresolved) = next.candidates.get_mut(&key) {
1105
                                unresolved.attempts = candidate.attempts.saturating_add(1);
1106
                                unresolved.first_seen_ms = candidate.first_seen_ms;
1107
                                unresolved.retry_at_ms =
1108
                                    now() + (1000u64 << unresolved.attempts.min(8));
1109
                            }
1110
                        }
1111
                        Err(error) => {
1112
                            if !candidate.encrypted()
1113
                                && (matches!(
1114
                                    error.kind,
1115
                                    ErrorKind::EventNotFound | ErrorKind::Forbidden
1116
                                ) || error
1117
                                    .server
1118
                                    .as_ref()
1119
                                    .is_some_and(|server| matches!(server.status, 403 | 404)))
1120
                            {
1121
                                next.gap(
1122
                                    Some(candidate.room_id),
1123
                                    MentionGapReason::ContextUnavailable,
1124
                                    None,
1125
                                );
1126
                            } else {
1127
                                candidate.attempts = candidate.attempts.saturating_add(1);
1128
                                candidate.retry_at_ms =
1129
                                    retry_at(&error, 1000u64 << candidate.attempts.min(8));
1130
                                next.candidates.insert(key, candidate);
1131
                                return false;
1132
                            }
1133
                        }
1134
                    }
1135
                }
1136
            }
1137
            Work::History(latest, from, result) => {
1138
                *history = false;
1139
                self.recovery_budget = self.recovery_budget.saturating_sub(1);
1140
                match result {
1141
                    Ok(page) => {
1142
                        next.coverage.historical_supported = Some(true);
1143
                        let mut reached_anchor = false;
1144
                        for notification in page.notifications {
1145
                            if let Some(event_id) = notification
1146
                                .event
1147
                                .get_field::<String>("event_id")
1148
                                .ok()
1149
                                .flatten()
1150
                            {
1151
                                let identity = format!("{}\n{}", notification.room_id, event_id);
1152
                                if latest && from.is_none() && next.latest_head.is_none() {
1153
                                    next.latest_head = Some(identity.clone());
1154
                                }
1155
                                if latest && next.latest_anchor.as_ref() == Some(&identity) {
1156
                                    reached_anchor = true;
1157
                                    break;
1158
                                }
1159
                            }
1160
                            if self
1161
                                .client
1162
                                .get_room(&notification.room_id)
1163
                                .is_some_and(|room| room.state() == RoomState::Joined)
1164
                                && let Ok(value) =
1165
                                    serde_json::from_str(notification.event.json().get())
1166
                            {
1167
                                self.observe(next, &notification.room_id, value).await;
1168
                            }
1169
                        }
1170
                        let token = page
1171
                            .next_token
1172
                            .filter(|token| Some(token) != from.as_ref() && !reached_anchor);
1173
                        if latest {
1174
                            next.latest_cursor = token.clone();
1175
                            if token.is_none() {
1176
                                self.latest_scan = false;
1177
                                next.latest_anchor =
1178
                                    next.latest_head.take().or(next.latest_anchor.take());
1179
                            }
1180
                        } else {
1181
                            next.history_cursor = token.clone();
1182
                            next.history_exhausted = token.is_none();
1183
                        }
1184
                    }
1185
                    Err(error) => {
1186
                        if error
1187
                            .server
1188
                            .as_ref()
1189
                            .is_some_and(|server| matches!(server.status, 404 | 405 | 501))
1190
                        {
1191
                            next.coverage.historical_supported = Some(false);
1192
                            self.recovery_budget = 0;
1193
                        }
1194
                        if error
1195
                            .server
1196
                            .as_ref()
1197
                            .is_some_and(|server| server.status == 400)
1198
                        {
1199
                            if latest {
1200
                                next.latest_cursor = None;
1201
                                next.latest_head = None;
1202
                            } else {
1203
                                next.history_cursor = None;
1204
                            }
1205
                        }
1206
                        next.gap(None, MentionGapReason::NotificationHistoryUnavailable, from);
1207
                        self.recovery_retry = retry_at(&error, 60_000);
1208
                    }
1209
                }
1210
                next.coverage.recovering = self.recovery_budget > 0
1211
                    && (self.latest_scan || !next.history_exhausted)
1212
                    && next.coverage.historical_supported != Some(false);
1213
            }
1214
        }
1215
        true
1216
    }
1217
1218
    fn schedule(
1219
        &mut self,
1220
        work: &mut FuturesUnordered<EngineBoxFuture<'static, Work>>,
1221
        captures: &mut BTreeSet<String>,
1222
        retries: &mut BTreeSet<String>,
1223
        history: &mut bool,
1224
    ) {
1225
        for (key, entry) in self.data.entries.iter().filter(|(_, entry)| matches!(entry.capture, MentionCaptureState::Pending) || matches!(entry.capture, MentionCaptureState::Failed {retry_at_ms, ..} if retry_at_ms <= now())).take(128) {
1226
            if captures.len() >= 2 { break; }
1227
            if captures.contains(key) { continue; }
1228
            let Some(room) = self.client.get_room(&entry.key.room_id).filter(|room| room.state() == RoomState::Joined) else { continue; };
1229
            captures.insert(key.clone());
1230
            let key = key.clone(); let event_id = entry.key.event_id.clone(); let account = self.account.clone(); let revision = entry.revision;
1231
            work.push(Box::pin(async move { Work::Capture(key, revision, context::capture(room, Some(account), event_id, 10, 10).await) }));
1232
        }
1233
        for (key, candidate) in &self.data.candidates {
1234
            if retries.len() >= 2 {
1235
                break;
1236
            }
1237
            if candidate.retry_at_ms > now() || retries.contains(key) {
1238
                continue;
1239
            }
1240
            let Some(room) = self
1241
                .client
1242
                .get_room(&candidate.room_id)
1243
                .filter(|room| room.state() == RoomState::Joined)
1244
            else {
1245
                continue;
1246
            };
1247
            let key = key.clone();
1248
            let candidate = candidate.clone();
1249
            let key_for_work = key.clone();
1250
            retries.insert(key);
1251
            work.push(Box::pin(async move {
1252
                let result = async {
1253
                    if candidate.raw["type"] == "m.room.encrypted" {
1254
                        let raw: Raw<OriginalSyncRoomEncryptedEvent> =
1255
                            Raw::from_json_string(candidate.raw.to_string()).map_err(|error| {
1256
                                Error::with_source(
1257
                                    ErrorKind::Invalid,
1258
                                    "invalid encrypted mention candidate",
1259
                                    error,
1260
                                )
1261
                            })?;
1262
                        let event = room.decrypt_event(&raw, None).await.map_err(|error| {
1263
                            Error::from_sdk("decrypting a mention candidate failed", error)
1264
                        })?;
1265
                        let value: Value =
1266
                            serde_json::from_str(event.raw().json().get()).map_err(|error| {
1267
                                Error::with_source(
1268
                                    ErrorKind::Invalid,
1269
                                    "invalid decrypted mention",
1270
                                    error,
1271
                                )
1272
                            })?;
1273
                        if value["type"] == "m.room.encrypted" {
1274
                            return Err(Error::new(
1275
                                ErrorKind::Encryption,
1276
                                "mention candidate still awaits keys",
1277
                            ));
1278
                        }
1279
                        return Ok(value);
1280
                    }
1281
                    let target = candidate.raw["content"]["m.relates_to"]["event_id"]
1282
                        .as_str()
1283
                        .filter(|_| {
1284
                            candidate.raw["content"]["m.relates_to"]["rel_type"] == "m.replace"
1285
                        });
1286
                    if let Some(target) =
1287
                        target.and_then(|target| target.parse::<ruma::OwnedEventId>().ok())
1288
                    {
1289
                        let event = room.event(&target, None).await.map_err(|error| {
1290
                            Error::from_sdk("fetching an edited mention failed", error)
1291
                        })?;
1292
                        let original: Value = serde_json::from_str(event.raw().json().get())
1293
                            .map_err(|error| {
1294
                                Error::with_source(
1295
                                    ErrorKind::Invalid,
1296
                                    "invalid edited mention",
1297
                                    error,
1298
                                )
1299
                            })?;
1300
                        if original["sender"] == candidate.raw["sender"] {
1301
                            let mut original = edited(original, &candidate.raw);
1302
                            original["unsigned"]["macltok_edit"] = candidate.raw;
1303
                            return Ok(original);
1304
                        }
1305
                        return Err(Error::new(
1306
                            ErrorKind::Forbidden,
1307
                            "an edit belongs to a different sender",
1308
                        ));
1309
                    }
1310
                    let mut value = candidate.raw;
1311
                    if let Some(target) =
1312
                        value["content"]["m.relates_to"]["m.in_reply_to"]["event_id"]
1313
                            .as_str()
1314
                            .and_then(|target| target.parse::<ruma::OwnedEventId>().ok())
1315
                    {
1316
                        let event = room.event(&target, None).await.map_err(|error| {
1317
                            Error::from_sdk("fetching a mention reply target failed", error)
1318
                        })?;
1319
                        let target: Value = serde_json::from_str(event.raw().json().get())
1320
                            .map_err(|error| {
1321
                                Error::with_source(
1322
                                    ErrorKind::Invalid,
1323
                                    "invalid mention reply target",
1324
                                    error,
1325
                                )
1326
                            })?;
1327
                        value["unsigned"]["macltok_reply_sender"] = target["sender"].clone();
1328
                    }
1329
                    Ok(value)
1330
                }
1331
                .await;
1332
                Work::Candidate(key_for_work, result)
1333
            }));
1334
        }
1335
        if !*history
1336
            && self.recovery_budget > 0
1337
            && now() >= self.recovery_retry
1338
            && (self.latest_scan || !self.data.history_exhausted)
1339
            && self.data.coverage.historical_supported != Some(false)
1340
        {
1341
            *history = true;
1342
            let latest = self.latest_scan;
1343
            let from = if latest {
1344
                self.data.latest_cursor.clone()
1345
            } else {
1346
                self.data.history_cursor.clone()
1347
            };
1348
            let client = self.client.clone();
1349
            work.push(Box::pin(async move {
1350
                let mut request = ruma::api::client::push::get_notifications::v3::Request::new();
1351
                request.from = from.clone();
1352
                request.limit = Some(ruma::uint!(50));
1353
                let result = client.send(request).await.map_err(|error| {
1354
                    Error::from_sdk("fetching mention history failed", error.into())
1355
                });
1356
                Work::History(latest, from, result)
1357
            }));
1358
        }
1359
    }
1360
1361
    fn captured(
1362
        &self,
1363
        data: &mut Data,
1364
        key: String,
1365
        revision: u64,
1366
        result: Result<(crate::EventContext, Vec<MediaRegistration>, Vec<Value>), Error>,
1367
    ) {
1368
        let Some(entry) = data.entries.get(&key) else {
1369
            return;
1370
        };
1371
        if entry.revision != revision {
1372
            return;
1373
        }
1374
        let room_id = entry.key.room_id.clone();
1375
        if !self
1376
            .client
1377
            .get_room(&room_id)
1378
            .is_some_and(|room| room.state() == RoomState::Joined)
1379
        {
1380
            return;
1381
        }
1382
        match result {
1383
            Ok((context, media, _)) => {
1384
                data.capture_attempts.remove(&key);
1385
                let live_after = data
1386
                    .windows
1387
                    .get(&key)
1388
                    .map(|window| window.after.clone())
1389
                    .unwrap_or_default();
1390
                let mut insert = |item: TimelineItem| -> Option<String> {
1391
                    let TimelineItem::Event(event) = &item else {
1392
                        return None;
1393
                    };
1394
                    let event_id = event.event_id.clone()?;
1395
                    let event_key = id(&MentionKey {
1396
                        room_id: room_id.clone(),
1397
                        event_id,
1398
                    });
1399
                    let registrations = media_for_item(&item, &media);
1400
                    let mut saved = SavedEvent {
1401
                        item,
1402
                        media: registrations,
1403
                        raw: None,
1404
                    };
1405
                    if let Some(existing) = data.events.get(&event_key) {
1406
                        saved = existing.clone();
1407
                    }
1408
                    if data.redacted.contains_key(&event_key) {
1409
                        redact(&mut saved);
1410
                    }
1411
                    data.events.insert(event_key.clone(), saved);
1412
                    Some(event_key)
1413
                };
1414
                let before: Vec<_> = context
1415
                    .events_before
1416
                    .into_iter()
1417
                    .filter_map(&mut insert)
1418
                    .collect();
1419
                let target = insert(context.event);
1420
                let mut after: Vec<_> = context
1421
                    .events_after
1422
                    .into_iter()
1423
                    .filter_map(&mut insert)
1424
                    .collect();
1425
                let extended = live_after.iter().any(|key| !after.contains(key));
1426
                after.extend(live_after);
1427
                after.sort_by_key(|key| {
1428
                    (
1429
                        data.events
1430
                            .get(key)
1431
                            .and_then(|saved| match &saved.item {
1432
                                TimelineItem::Event(event) => Some(event.timestamp_ms),
1433
                                _ => None,
1434
                            })
1435
                            .unwrap_or(0),
1436
                        key.clone(),
1437
                    )
1438
                });
1439
                after.dedup();
1440
                after.truncate(10);
1441
                let entry = data.entries.get_mut(&key).unwrap();
1442
                entry.capture = if target.is_none() {
1443
                    MentionCaptureState::Unavailable
1444
                } else if after.len() < 10 || (before.len() < 10 && context.start_token.is_some()) {
1445
                    MentionCaptureState::Partial
1446
                } else {
1447
                    MentionCaptureState::Ready
1448
                };
1449
                entry.revision = data.revision + 1;
1450
                data.windows.insert(
1451
                    key,
1452
                    Window {
1453
                        before,
1454
                        target,
1455
                        after,
1456
                        start: context.start_token,
1457
                        end: if extended { None } else { context.end_token },
1458
                    },
1459
                );
1460
            }
1461
            Err(error) => {
1462
                let attempts = data.capture_attempts.entry(key.clone()).or_default();
1463
                *attempts = attempts.saturating_add(1);
1464
                let delay = (1000u64 << (*attempts).min(11)).min(3_600_000);
1465
                let entry = data.entries.get_mut(&key).unwrap();
1466
                entry.capture = if matches!(
1467
                    error.kind,
1468
                    ErrorKind::EventNotFound | ErrorKind::Forbidden | ErrorKind::RoomUnknown
1469
                ) || error
1470
                    .server
1471
                    .as_ref()
1472
                    .is_some_and(|server| matches!(server.status, 403 | 404))
1473
                {
1474
                    MentionCaptureState::Unavailable
1475
                } else {
1476
                    MentionCaptureState::Failed {
1477
                        message: error.to_string(),
1478
                        retry_at_ms: retry_at(&error, delay),
1479
                    }
1480
                };
1481
                entry.revision = data.revision + 1;
1482
            }
1483
        }
1484
    }
1485
}
1486
1487
fn edited(mut original: Value, edit: &Value) -> Value {
1488
    original["content"] = edit["content"]["m.new_content"].clone();
1489
    original["unsigned"]["macltok_edited"] = Value::Bool(true);
1490
    original
1491
}
1492
1493
fn retry_at(error: &Error, fallback_ms: u64) -> u64 {
1494
    let mut source = std::error::Error::source(error);
1495
    while let Some(current) = source {
1496
        if let Some(sdk) = current.downcast_ref::<matrix_sdk::Error>() {
1497
            if let Some(ruma::api::error::ErrorKind::LimitExceeded(data)) =
1498
                sdk.client_api_error_kind()
1499
            {
1500
                if let Some(delay) = data.retry_after {
1501
                    let delay = match delay {
1502
                        ruma::api::error::RetryAfter::Delay(delay) => delay,
1503
                        ruma::api::error::RetryAfter::DateTime(time) => time
1504
                            .duration_since(ruma::time::SystemTime::now())
1505
                            .unwrap_or_default(),
1506
                    };
1507
                    return now().saturating_add(
1508
                        u64::try_from(delay.as_millis())
1509
                            .unwrap_or(u64::MAX)
1510
                            .max(fallback_ms),
1511
                    );
1512
                }
1513
            }
1514
        }
1515
        source = current.source();
1516
    }
1517
    now().saturating_add(fallback_ms)
1518
}
1519
1520
fn redact(event: &mut SavedEvent) {
1521
    if let TimelineItem::Event(item) = &mut event.item {
1522
        item.content = TimelineContent::Redacted;
1523
        item.reply_to = None;
1524
        item.reactions.clear();
1525
    }
1526
    event.media.clear();
1527
    event.raw = None;
1528
}
1529
1530
fn media_for_item(
1531
    item: &TimelineItem,
1532
    registrations: &[MediaRegistration],
1533
) -> Vec<MediaRegistration> {
1534
    let TimelineItem::Event(event) = item else {
1535
        return Vec::new();
1536
    };
1537
    let mut ids = Vec::new();
1538
    match &event.content {
1539
        TimelineContent::Message(message) => {
1540
            ids.extend(message.media.iter().map(|media| &media.id));
1541
            ids.extend(
1542
                message
1543
                    .link_previews
1544
                    .iter()
1545
                    .filter_map(|preview| preview.image.as_ref().map(|image| &image.media_id)),
1546
            );
1547
        }
1548
        TimelineContent::Sticker { media, .. } => ids.push(&media.id),
1549
        _ => {}
1550
    }
1551
    registrations
1552
        .iter()
1553
        .filter(|registration| ids.contains(&&registration.id))
1554
        .cloned()
1555
        .collect()
1556
}
1557
1558
pub(crate) async fn watch_updates(mut updates: watch::Receiver<MentionsPage>, events: EventSender) {
1559
    let mut previous: BTreeMap<String, MentionEntry> = BTreeMap::new();
1560
    loop {
1561
        let snapshot = updates.borrow_and_update().clone();
1562
        let current: BTreeMap<_, _> = snapshot
1563
            .entries
1564
            .iter()
1565
            .map(|entry| (id(&entry.key), entry.clone()))
1566
            .collect();
1567
        let upserts = current
1568
            .iter()
1569
            .filter(|(key, entry)| previous.get(*key) != Some(entry))
1570
            .map(|(_, entry)| entry.clone())
1571
            .collect();
1572
        let removed = previous
1573
            .iter()
1574
            .filter(|(key, _)| !current.contains_key(*key))
1575
            .map(|(_, entry)| entry.key.clone())
1576
            .collect();
1577
        previous = current;
1578
        if events
1579
            .send(Event::MentionsChanged(MentionsUpdate {
1580
                snapshot,
1581
                upserts,
1582
                removed,
1583
            }))
1584
            .await
1585
            .is_err()
1586
        {
1587
            return;
1588
        }
1589
        if updates.changed().await.is_err() {
1590
            return;
1591
        }
1592
    }
1593
}