refactor: enhance state retrieval logic in StateService by consolidating state fetching and improving error handling for better clarity

ggazzo · ggazzo · commit 55e2213e495d · 2025-09-13T19:43:15.000-03:00
diff --git a/packages/federation-sdk/src/services/state.service.ts b/packages/federation-sdk/src/services/state.service.ts
@@ -250,36 +250,44 @@ export class StateService {
 		// Retrieves a snapshot of a room's state at a given event, in the form of event IDs. This performs the same function as calling /state/{roomId}, however this returns just the event IDs rather than the full events.
 		const { stateId } = event;
 
-		const { delta: lastStateDelta, prevStateIds = [] } =
-			(await this.stateRepository.getStateById(stateId)) ?? {};
-
-		this.logger.debug({ delta: lastStateDelta, prevStateIds }, 'last state');
+		const storedState = await this.stateRepository.getStateById(stateId);
 
-		if (!lastStateDelta) {
+		if (!storedState) {
 			this.logger.error(eventId, 'last state delta not found');
 			throw new Error(`State at event ${eventId} not found`);
 		}
 
+		this.logger.debug(
+			{ delta: storedState?.delta, prevStateIds: storedState.prevStateIds },
+			'last state',
+		);
+
 		const state = new Map<StateMapKey, PersistentEventBase>();
 
-		if (prevStateIds.length === 0) {
+		// TODO: not sure which case is this
+		if (storedState.prevStateIds.length === 0) {
 			const previous = await this.eventRepository.findById(
-				lastStateDelta.eventId,
+				storedState.delta.eventId,
 			);
 			if (!previous) {
-				throw new Error(`Event ${lastStateDelta.eventId} not found`);
+				throw new Error(`Event ${eventId} not found`);
 			}
 
 			state.set(
-				lastStateDelta.identifier,
+				storedState.delta.identifier,
 				PersistentEventFactory.createFromRawEvent(previous.event, roomVersion),
 			);
 
 			return state;
 		}
 
+		const eventsToFetch = [
+			includeEvent && eventId,
+			...storedState.prevStateIds,
+		].filter(Boolean) as string[];
+
 		const stateMappings = await this.stateRepository
-			.getStateMappingsByStateIdsOrdered(prevStateIds)
+			.getStateMappingsByStateIdsOrdered(eventsToFetch)
 			.toArray();
 
 		for await (const { delta } of stateMappings) {
@@ -295,26 +303,6 @@ export class StateService {
 			);
 		}
 
-		if (!includeEvent) {
-			return state;
-		}
-
-		this.logger.debug({ eventId }, 'including event in state');
-
-		// update the last state
-		const { identifier: lastStateKey, eventId: lastStateEventId } =
-			lastStateDelta;
-
-		const lastEvent = await this.eventRepository.findById(lastStateEventId);
-		if (!lastEvent) {
-			throw new Error(`Event ${lastStateEventId} not found`);
-		}
-
-		state.set(
-			lastStateKey,
-			PersistentEventFactory.createFromRawEvent(lastEvent.event, roomVersion),
-		);
-
 		return state;
 	}
 
