atagen

macltok

public
Star0
Watch0
2c13c8db61741baf208b9c98b78c4a9e97d1fe61
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git

Blame: src/engine/context.rs

89bc9f8d246datagen2026-10-24 13:40 UTC1
use std::sync::Arc;
89bc9f8d246datagen2026-10-24 13:40 UTC2
cbb715e11963atagen2026-10-24 23:24 UTC3
use matrix_sdk::{
cbb715e11963atagen2026-10-24 23:24 UTC4
    Room, deserialized_responses::TimelineEvent as SdkTimelineEvent, room::MessagesOptions,
cbb715e11963atagen2026-10-24 23:24 UTC5
};
cbb715e11963atagen2026-10-24 23:24 UTC6
use ruma::{OwnedEventId, OwnedUserId, RoomId, UInt};
89bc9f8d246datagen2026-10-24 13:40 UTC7
89bc9f8d246datagen2026-10-24 13:40 UTC8
use crate::{
89bc9f8d246datagen2026-10-24 13:40 UTC9
    Error, ErrorKind, EventContext, EventContextCommand, EventContextDirection, EventContextPage,
89bc9f8d246datagen2026-10-24 13:40 UTC10
    TimelineItem,
cbb715e11963atagen2026-10-24 23:24 UTC11
    engine::MediaService,
cbb715e11963atagen2026-10-24 23:24 UTC12
    normalise::{MediaRegistration, normalise_observed},
89bc9f8d246datagen2026-10-24 13:40 UTC13
};
89bc9f8d246datagen2026-10-24 13:40 UTC14
89bc9f8d246datagen2026-10-24 13:40 UTC15
const MAX_CONTEXT_LIMIT: u16 = 50;
89bc9f8d246datagen2026-10-24 13:40 UTC16
89bc9f8d246datagen2026-10-24 13:40 UTC17
pub(crate) async fn execute(
89bc9f8d246datagen2026-10-24 13:40 UTC18
    room: Room,
cbb715e11963atagen2026-10-24 23:24 UTC19
    current_user_id: Option<OwnedUserId>,
cbb715e11963atagen2026-10-24 23:24 UTC20
    media: Arc<MediaService>,
89bc9f8d246datagen2026-10-24 13:40 UTC21
    command: EventContextCommand,
89bc9f8d246datagen2026-10-24 13:40 UTC22
) -> Result<crate::CommandOutput, Error> {
89bc9f8d246datagen2026-10-24 13:40 UTC23
    if let EventContextCommand::Before { token, limit, .. }
89bc9f8d246datagen2026-10-24 13:40 UTC24
    | EventContextCommand::After { token, limit, .. } = &command
89bc9f8d246datagen2026-10-24 13:40 UTC25
    {
89bc9f8d246datagen2026-10-24 13:40 UTC26
        validate_limit(*limit)?;
89bc9f8d246datagen2026-10-24 13:40 UTC27
        if token.is_empty() {
89bc9f8d246datagen2026-10-24 13:40 UTC28
            return Err(Error::new(
89bc9f8d246datagen2026-10-24 13:40 UTC29
                ErrorKind::Invalid,
89bc9f8d246datagen2026-10-24 13:40 UTC30
                "context pagination token is empty",
89bc9f8d246datagen2026-10-24 13:40 UTC31
            ));
89bc9f8d246datagen2026-10-24 13:40 UTC32
        }
89bc9f8d246datagen2026-10-24 13:40 UTC33
        return execute_page(
89bc9f8d246datagen2026-10-24 13:40 UTC34
            room,
89bc9f8d246datagen2026-10-24 13:40 UTC35
            current_user_id,
cbb715e11963atagen2026-10-24 23:24 UTC36
            media,
89bc9f8d246datagen2026-10-24 13:40 UTC37
            token,
89bc9f8d246datagen2026-10-24 13:40 UTC38
            *limit,
89bc9f8d246datagen2026-10-24 13:40 UTC39
            matches!(command, EventContextCommand::Before { .. }),
89bc9f8d246datagen2026-10-24 13:40 UTC40
        )
89bc9f8d246datagen2026-10-24 13:40 UTC41
        .await;
89bc9f8d246datagen2026-10-24 13:40 UTC42
    }
89bc9f8d246datagen2026-10-24 13:40 UTC43
cbb715e11963atagen2026-10-24 23:24 UTC44
    let (event_id, before_limit, after_limit) = match &command {
89bc9f8d246datagen2026-10-24 13:40 UTC45
        EventContextCommand::Fetch {
89bc9f8d246datagen2026-10-24 13:40 UTC46
            event_id,
89bc9f8d246datagen2026-10-24 13:40 UTC47
            before_limit,
89bc9f8d246datagen2026-10-24 13:40 UTC48
            after_limit,
89bc9f8d246datagen2026-10-24 13:40 UTC49
            ..
89bc9f8d246datagen2026-10-24 13:40 UTC50
        } => {
89bc9f8d246datagen2026-10-24 13:40 UTC51
            validate_limit(*before_limit)?;
89bc9f8d246datagen2026-10-24 13:40 UTC52
            validate_limit(*after_limit)?;
cbb715e11963atagen2026-10-24 23:24 UTC53
            (event_id.clone(), *before_limit, *after_limit)
89bc9f8d246datagen2026-10-24 13:40 UTC54
        }
89bc9f8d246datagen2026-10-24 13:40 UTC55
        EventContextCommand::Before { .. } | EventContextCommand::After { .. } => unreachable!(),
89bc9f8d246datagen2026-10-24 13:40 UTC56
    };
cbb715e11963atagen2026-10-24 23:24 UTC57
    let (context, registrations, _) =
cbb715e11963atagen2026-10-24 23:24 UTC58
        capture(room, current_user_id, event_id, before_limit, after_limit).await?;
cbb715e11963atagen2026-10-24 23:24 UTC59
    media.register(registrations).await;
cbb715e11963atagen2026-10-24 23:24 UTC60
    Ok(crate::CommandOutput::EventContext(context))
cbb715e11963atagen2026-10-24 23:24 UTC61
}
89bc9f8d246datagen2026-10-24 13:40 UTC62
cbb715e11963atagen2026-10-24 23:24 UTC63
pub(crate) async fn capture(
cbb715e11963atagen2026-10-24 23:24 UTC64
    room: Room,
cbb715e11963atagen2026-10-24 23:24 UTC65
    current_user_id: Option<OwnedUserId>,
cbb715e11963atagen2026-10-24 23:24 UTC66
    event_id: OwnedEventId,
cbb715e11963atagen2026-10-24 23:24 UTC67
    before_limit: u16,
cbb715e11963atagen2026-10-24 23:24 UTC68
    after_limit: u16,
cbb715e11963atagen2026-10-24 23:24 UTC69
) -> Result<(EventContext, Vec<MediaRegistration>, Vec<serde_json::Value>), Error> {
cbb715e11963atagen2026-10-24 23:24 UTC70
    validate_limit(before_limit)?;
cbb715e11963atagen2026-10-24 23:24 UTC71
    validate_limit(after_limit)?;
cbb715e11963atagen2026-10-24 23:24 UTC72
    let raw = room
cbb715e11963atagen2026-10-24 23:24 UTC73
        .event_with_context(&event_id, false, UInt::from(0_u16), None)
89bc9f8d246datagen2026-10-24 13:40 UTC74
        .await
89bc9f8d246datagen2026-10-24 13:40 UTC75
        .map_err(|source| {
cbb715e11963atagen2026-10-24 23:24 UTC76
            let error = Error::from_sdk("fetching event context failed", source);
cbb715e11963atagen2026-10-24 23:24 UTC77
            if error
cbb715e11963atagen2026-10-24 23:24 UTC78
                .server
cbb715e11963atagen2026-10-24 23:24 UTC79
                .as_ref()
cbb715e11963atagen2026-10-24 23:24 UTC80
                .is_some_and(|server| server.status == 404)
cbb715e11963atagen2026-10-24 23:24 UTC81
            {
cbb715e11963atagen2026-10-24 23:24 UTC82
                Error::new(ErrorKind::EventNotFound, "event was not found")
cbb715e11963atagen2026-10-24 23:24 UTC83
            } else {
cbb715e11963atagen2026-10-24 23:24 UTC84
                error
cbb715e11963atagen2026-10-24 23:24 UTC85
            }
89bc9f8d246datagen2026-10-24 13:40 UTC86
        })?;
cbb715e11963atagen2026-10-24 23:24 UTC87
    if raw.event.is_none() {
cbb715e11963atagen2026-10-24 23:24 UTC88
        return Err(Error::new(ErrorKind::EventNotFound, "event was not found"));
89bc9f8d246datagen2026-10-24 13:40 UTC89
    }
cbb715e11963atagen2026-10-24 23:24 UTC90
    let before = fetch_side(
cbb715e11963atagen2026-10-24 23:24 UTC91
        &room,
cbb715e11963atagen2026-10-24 23:24 UTC92
        raw.prev_batch_token.as_deref(),
cbb715e11963atagen2026-10-24 23:24 UTC93
        before_limit,
cbb715e11963atagen2026-10-24 23:24 UTC94
        true,
cbb715e11963atagen2026-10-24 23:24 UTC95
        current_user_id.as_deref(),
cbb715e11963atagen2026-10-24 23:24 UTC96
    );
cbb715e11963atagen2026-10-24 23:24 UTC97
    let after = fetch_side(
cbb715e11963atagen2026-10-24 23:24 UTC98
        &room,
cbb715e11963atagen2026-10-24 23:24 UTC99
        raw.next_batch_token.as_deref(),
cbb715e11963atagen2026-10-24 23:24 UTC100
        after_limit,
cbb715e11963atagen2026-10-24 23:24 UTC101
        false,
cbb715e11963atagen2026-10-24 23:24 UTC102
        current_user_id.as_deref(),
cbb715e11963atagen2026-10-24 23:24 UTC103
    );
cbb715e11963atagen2026-10-24 23:24 UTC104
    let (events_before, mut before_media, start_token, mut observed) = before.await?;
cbb715e11963atagen2026-10-24 23:24 UTC105
    let (events_after, after_media, end_token, following) = after.await?;
cbb715e11963atagen2026-10-24 23:24 UTC106
    observed.extend(following);
cbb715e11963atagen2026-10-24 23:24 UTC107
    if let Some(value) = raw
cbb715e11963atagen2026-10-24 23:24 UTC108
        .event
cbb715e11963atagen2026-10-24 23:24 UTC109
        .as_ref()
cbb715e11963atagen2026-10-24 23:24 UTC110
        .and_then(|event| serde_json::from_str(event.raw().json().get()).ok())
cbb715e11963atagen2026-10-24 23:24 UTC111
    {
cbb715e11963atagen2026-10-24 23:24 UTC112
        observed.push(value);
cbb715e11963atagen2026-10-24 23:24 UTC113
    }
cbb715e11963atagen2026-10-24 23:24 UTC114
    let event = raw.event.as_ref().and_then(|event| {
cbb715e11963atagen2026-10-24 23:24 UTC115
        normalise_sdk_event(event, current_user_id.as_deref(), room.room_id()).map(
cbb715e11963atagen2026-10-24 23:24 UTC116
            |(item, media)| {
cbb715e11963atagen2026-10-24 23:24 UTC117
                before_media.extend(media);
cbb715e11963atagen2026-10-24 23:24 UTC118
                item
cbb715e11963atagen2026-10-24 23:24 UTC119
            },
cbb715e11963atagen2026-10-24 23:24 UTC120
        )
cbb715e11963atagen2026-10-24 23:24 UTC121
    });
cbb715e11963atagen2026-10-24 23:24 UTC122
    before_media.extend(after_media);
cbb715e11963atagen2026-10-24 23:24 UTC123
    Ok((
cbb715e11963atagen2026-10-24 23:24 UTC124
        EventContext {
cbb715e11963atagen2026-10-24 23:24 UTC125
            events_before,
cbb715e11963atagen2026-10-24 23:24 UTC126
            event,
cbb715e11963atagen2026-10-24 23:24 UTC127
            events_after,
cbb715e11963atagen2026-10-24 23:24 UTC128
            start_token,
cbb715e11963atagen2026-10-24 23:24 UTC129
            end_token,
cbb715e11963atagen2026-10-24 23:24 UTC130
        },
cbb715e11963atagen2026-10-24 23:24 UTC131
        before_media,
cbb715e11963atagen2026-10-24 23:24 UTC132
        observed,
cbb715e11963atagen2026-10-24 23:24 UTC133
    ))
cbb715e11963atagen2026-10-24 23:24 UTC134
}
89bc9f8d246datagen2026-10-24 13:40 UTC135
cbb715e11963atagen2026-10-24 23:24 UTC136
async fn fetch_side(
cbb715e11963atagen2026-10-24 23:24 UTC137
    room: &Room,
cbb715e11963atagen2026-10-24 23:24 UTC138
    token: Option<&str>,
cbb715e11963atagen2026-10-24 23:24 UTC139
    limit: u16,
cbb715e11963atagen2026-10-24 23:24 UTC140
    backwards: bool,
cbb715e11963atagen2026-10-24 23:24 UTC141
    current_user_id: Option<&ruma::UserId>,
cbb715e11963atagen2026-10-24 23:24 UTC142
) -> Result<
cbb715e11963atagen2026-10-24 23:24 UTC143
    (
cbb715e11963atagen2026-10-24 23:24 UTC144
        Vec<TimelineItem>,
cbb715e11963atagen2026-10-24 23:24 UTC145
        Vec<MediaRegistration>,
cbb715e11963atagen2026-10-24 23:24 UTC146
        Option<String>,
cbb715e11963atagen2026-10-24 23:24 UTC147
        Vec<serde_json::Value>,
cbb715e11963atagen2026-10-24 23:24 UTC148
    ),
cbb715e11963atagen2026-10-24 23:24 UTC149
    Error,
cbb715e11963atagen2026-10-24 23:24 UTC150
> {
cbb715e11963atagen2026-10-24 23:24 UTC151
    if limit == 0 {
cbb715e11963atagen2026-10-24 23:24 UTC152
        return Ok((Vec::new(), Vec::new(), token.map(str::to_owned), Vec::new()));
cbb715e11963atagen2026-10-24 23:24 UTC153
    }
cbb715e11963atagen2026-10-24 23:24 UTC154
    let Some(token) = token else {
cbb715e11963atagen2026-10-24 23:24 UTC155
        return Ok((Vec::new(), Vec::new(), None, Vec::new()));
cbb715e11963atagen2026-10-24 23:24 UTC156
    };
cbb715e11963atagen2026-10-24 23:24 UTC157
    let mut options = if backwards {
cbb715e11963atagen2026-10-24 23:24 UTC158
        MessagesOptions::backward()
cbb715e11963atagen2026-10-24 23:24 UTC159
    } else {
cbb715e11963atagen2026-10-24 23:24 UTC160
        MessagesOptions::forward()
cbb715e11963atagen2026-10-24 23:24 UTC161
    };
cbb715e11963atagen2026-10-24 23:24 UTC162
    options = options.from(token);
cbb715e11963atagen2026-10-24 23:24 UTC163
    options.limit = UInt::from(limit);
cbb715e11963atagen2026-10-24 23:24 UTC164
    let messages = room
cbb715e11963atagen2026-10-24 23:24 UTC165
        .messages(options)
cbb715e11963atagen2026-10-24 23:24 UTC166
        .await
cbb715e11963atagen2026-10-24 23:24 UTC167
        .map_err(|source| Error::from_sdk("fetching event context events failed", source))?;
cbb715e11963atagen2026-10-24 23:24 UTC168
    let mut events = messages.chunk;
cbb715e11963atagen2026-10-24 23:24 UTC169
    events.truncate(usize::from(limit));
cbb715e11963atagen2026-10-24 23:24 UTC170
    if backwards {
cbb715e11963atagen2026-10-24 23:24 UTC171
        events.reverse();
cbb715e11963atagen2026-10-24 23:24 UTC172
    }
cbb715e11963atagen2026-10-24 23:24 UTC173
    let next_token = messages.end;
cbb715e11963atagen2026-10-24 23:24 UTC174
    let observed = events
cbb715e11963atagen2026-10-24 23:24 UTC175
        .iter()
cbb715e11963atagen2026-10-24 23:24 UTC176
        .filter_map(|event| serde_json::from_str(event.raw().json().get()).ok())
cbb715e11963atagen2026-10-24 23:24 UTC177
        .collect();
cbb715e11963atagen2026-10-24 23:24 UTC178
    let (items, media) = normalise_events(events, current_user_id, room.room_id());
cbb715e11963atagen2026-10-24 23:24 UTC179
    Ok((items, media, next_token, observed))
89bc9f8d246datagen2026-10-24 13:40 UTC180
}
89bc9f8d246datagen2026-10-24 13:40 UTC181
89bc9f8d246datagen2026-10-24 13:40 UTC182
async fn execute_page(
89bc9f8d246datagen2026-10-24 13:40 UTC183
    room: Room,
cbb715e11963atagen2026-10-24 23:24 UTC184
    current_user_id: Option<OwnedUserId>,
cbb715e11963atagen2026-10-24 23:24 UTC185
    media: Arc<MediaService>,
89bc9f8d246datagen2026-10-24 13:40 UTC186
    token: &str,
89bc9f8d246datagen2026-10-24 13:40 UTC187
    limit: u16,
89bc9f8d246datagen2026-10-24 13:40 UTC188
    backwards: bool,
89bc9f8d246datagen2026-10-24 13:40 UTC189
) -> Result<crate::CommandOutput, Error> {
89bc9f8d246datagen2026-10-24 13:40 UTC190
    let mut options = if backwards {
89bc9f8d246datagen2026-10-24 13:40 UTC191
        MessagesOptions::backward()
89bc9f8d246datagen2026-10-24 13:40 UTC192
    } else {
89bc9f8d246datagen2026-10-24 13:40 UTC193
        MessagesOptions::forward()
89bc9f8d246datagen2026-10-24 13:40 UTC194
    };
89bc9f8d246datagen2026-10-24 13:40 UTC195
    options = options.from(token);
cbb715e11963atagen2026-10-24 23:24 UTC196
    options.limit = UInt::from(limit);
89bc9f8d246datagen2026-10-24 13:40 UTC197
    let messages = room
89bc9f8d246datagen2026-10-24 13:40 UTC198
        .messages(options)
89bc9f8d246datagen2026-10-24 13:40 UTC199
        .await
89bc9f8d246datagen2026-10-24 13:40 UTC200
        .map_err(|source| Error::from_sdk("fetching event context page failed", source))?;
cbb715e11963atagen2026-10-24 23:24 UTC201
    let next_token = messages.end.clone();
cbb715e11963atagen2026-10-24 23:24 UTC202
    let mut events = messages.chunk;
cbb715e11963atagen2026-10-24 23:24 UTC203
    events.truncate(usize::from(limit));
89bc9f8d246datagen2026-10-24 13:40 UTC204
    if backwards {
89bc9f8d246datagen2026-10-24 13:40 UTC205
        events.reverse();
89bc9f8d246datagen2026-10-24 13:40 UTC206
    }
cbb715e11963atagen2026-10-24 23:24 UTC207
    let (events, registrations) =
cbb715e11963atagen2026-10-24 23:24 UTC208
        normalise_events(events, current_user_id.as_deref(), room.room_id());
cbb715e11963atagen2026-10-24 23:24 UTC209
    media.register(registrations).await;
89bc9f8d246datagen2026-10-24 13:40 UTC210
    let end_reached = next_token.is_none();
89bc9f8d246datagen2026-10-24 13:40 UTC211
    let direction = if backwards {
89bc9f8d246datagen2026-10-24 13:40 UTC212
        EventContextDirection::Before
89bc9f8d246datagen2026-10-24 13:40 UTC213
    } else {
89bc9f8d246datagen2026-10-24 13:40 UTC214
        EventContextDirection::After
89bc9f8d246datagen2026-10-24 13:40 UTC215
    };
89bc9f8d246datagen2026-10-24 13:40 UTC216
    Ok(crate::CommandOutput::EventContextPage(EventContextPage {
89bc9f8d246datagen2026-10-24 13:40 UTC217
        direction,
89bc9f8d246datagen2026-10-24 13:40 UTC218
        events,
89bc9f8d246datagen2026-10-24 13:40 UTC219
        next_token,
89bc9f8d246datagen2026-10-24 13:40 UTC220
        end_reached,
89bc9f8d246datagen2026-10-24 13:40 UTC221
    }))
89bc9f8d246datagen2026-10-24 13:40 UTC222
}
89bc9f8d246datagen2026-10-24 13:40 UTC223
cbb715e11963atagen2026-10-24 23:24 UTC224
fn normalise_events(
cbb715e11963atagen2026-10-24 23:24 UTC225
    events: Vec<SdkTimelineEvent>,
cbb715e11963atagen2026-10-24 23:24 UTC226
    current_user_id: Option<&ruma::UserId>,
cbb715e11963atagen2026-10-24 23:24 UTC227
    room_id: &RoomId,
cbb715e11963atagen2026-10-24 23:24 UTC228
) -> (Vec<TimelineItem>, Vec<MediaRegistration>) {
cbb715e11963atagen2026-10-24 23:24 UTC229
    let mut items = Vec::with_capacity(events.len());
cbb715e11963atagen2026-10-24 23:24 UTC230
    let mut media = Vec::new();
cbb715e11963atagen2026-10-24 23:24 UTC231
    for event in events {
cbb715e11963atagen2026-10-24 23:24 UTC232
        if let Some((item, registrations)) = normalise_sdk_event(&event, current_user_id, room_id) {
cbb715e11963atagen2026-10-24 23:24 UTC233
            items.push(item);
cbb715e11963atagen2026-10-24 23:24 UTC234
            media.extend(registrations);
cbb715e11963atagen2026-10-24 23:24 UTC235
        }
89bc9f8d246datagen2026-10-24 13:40 UTC236
    }
cbb715e11963atagen2026-10-24 23:24 UTC237
    (items, media)
89bc9f8d246datagen2026-10-24 13:40 UTC238
}
89bc9f8d246datagen2026-10-24 13:40 UTC239
cbb715e11963atagen2026-10-24 23:24 UTC240
fn normalise_sdk_event(
cbb715e11963atagen2026-10-24 23:24 UTC241
    event: &SdkTimelineEvent,
89bc9f8d246datagen2026-10-24 13:40 UTC242
    current_user_id: Option<&ruma::UserId>,
cbb715e11963atagen2026-10-24 23:24 UTC243
    room_id: &RoomId,
cbb715e11963atagen2026-10-24 23:24 UTC244
) -> Option<(TimelineItem, Vec<MediaRegistration>)> {
cbb715e11963atagen2026-10-24 23:24 UTC245
    let value: serde_json::Value = serde_json::from_str(event.raw().json().get()).ok()?;
cbb715e11963atagen2026-10-24 23:24 UTC246
    let event_id = value.get("event_id")?.as_str()?;
cbb715e11963atagen2026-10-24 23:24 UTC247
    let identity = format!("{room_id}\n{event_id}");
cbb715e11963atagen2026-10-24 23:24 UTC248
    normalise_observed(&value, current_user_id, &identity)
89bc9f8d246datagen2026-10-24 13:40 UTC249
}
89bc9f8d246datagen2026-10-24 13:40 UTC250
cbb715e11963atagen2026-10-24 23:24 UTC251
fn validate_limit(limit: u16) -> Result<(), Error> {
cbb715e11963atagen2026-10-24 23:24 UTC252
    if limit > MAX_CONTEXT_LIMIT {
cbb715e11963atagen2026-10-24 23:24 UTC253
        Err(Error::new(ErrorKind::Invalid, "context limit exceeds 50"))
cbb715e11963atagen2026-10-24 23:24 UTC254
    } else {
cbb715e11963atagen2026-10-24 23:24 UTC255
        Ok(())
cbb715e11963atagen2026-10-24 23:24 UTC256
    }
89bc9f8d246datagen2026-10-24 13:40 UTC257
}