atagen

macltok

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