atagen

macltok

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