From 8fb4d6fffe384f4f2b3617ac54ea7451e284b791 Mon Sep 17 00:00:00 2001
From: Eric Eastwood <erice@element.io>
Date: Wed, 29 Sep 2021 22:03:13 -0500
Subject: [PATCH] Force /context to return  state for the given historical
 event

---
 synapse/handlers/message.py |  2 ++
 synapse/handlers/room.py    | 19 +++++++++++++++----
 synapse/rest/client/room.py |  1 +
 synapse/storage/state.py    |  9 ++++++++-
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index 79799d774f..f3dad1ab5c 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -611,6 +611,7 @@ class EventCreationHandler:
                 auth_event_ids=auth_event_ids,
                 depth=depth,
             )
+            #logger.info("auth_event_ids before=%s", auth_event_ids)
             auth_events = await self.store.get_events_as_list(auth_event_ids)
             # Create a StateMap[str]
             auth_event_state_map = {
@@ -622,6 +623,7 @@ class EventCreationHandler:
                 current_state_ids=auth_event_state_map,
                 for_verification=False,
             )
+            #logger.info("auth_event_ids after=%s", auth_event_ids)
 
         event, context = await self.create_new_client_event(
             builder=builder,
diff --git a/synapse/handlers/room.py b/synapse/handlers/room.py
index 8fede5e935..ef556cd058 100644
--- a/synapse/handlers/room.py
+++ b/synapse/handlers/room.py
@@ -1166,11 +1166,22 @@ class RoomContextHandler:
         # first? Shouldn't we be consistent with /sync?
         # https://github.com/matrix-org/matrix-doc/issues/687
 
-        state = await self.state_store.get_state_for_events(
-            [last_event_id], state_filter=state_filter
-        )
+        logger.info("get_event_context event_id=%s =====================", event_id)
+        
+        event_id_to_get_state_from = last_event_id
 
-        state_events = list(state[last_event_id].values())
+        # For historical events, we want to get the state at the specified event.
+        # TODO: maybe we can change how we're getting events_before and events_after
+        # here so it still works correctly without this hack
+        if event.internal_metadata.is_historical():
+            event_id_to_get_state_from = event_id
+
+        state = await self.state_store.get_state_for_events(
+            [event_id_to_get_state_from], state_filter=state_filter
+        )
+        logger.info("get_event_context event_id=%s state=%s", event_id, state)
+
+        state_events = list(state[event_id_to_get_state_from].values())
         if event_filter:
             state_events = event_filter.filter(state_events)
 
diff --git a/synapse/rest/client/room.py b/synapse/rest/client/room.py
index bf46dc60f2..a73b2f9582 100644
--- a/synapse/rest/client/room.py
+++ b/synapse/rest/client/room.py
@@ -695,6 +695,7 @@ class RoomEventContextServlet(RestServlet):
         results = await self.room_context_handler.get_event_context(
             requester, room_id, event_id, limit, event_filter
         )
+        logger.info("get /context event_id=%s results=%s", event_id, results)
 
         if not results:
             raise SynapseError(404, "Event not found.", errcode=Codes.NOT_FOUND)
diff --git a/synapse/storage/state.py b/synapse/storage/state.py
index 5e86befde4..90ded4ef81 100644
--- a/synapse/storage/state.py
+++ b/synapse/storage/state.py
@@ -491,7 +491,6 @@ class StateGroupStorage:
             [ev_id for sd in group_to_state.values() for ev_id in sd.values()],
             get_prev_content=False,
         )
-
         event_to_state = {
             event_id: {
                 k: state_event_map[v]
@@ -500,6 +499,14 @@ class StateGroupStorage:
             }
             for event_id, group in event_to_groups.items()
         }
+        logger.info(
+            "event_to_groups event_ids=%s event_to_groups=%s group_to_state=%s state_event_map=%s event_to_state=%s",
+            event_ids,
+            event_to_groups,
+            group_to_state,
+            state_event_map,
+            event_to_state
+        )
 
         return {event: event_to_state[event] for event in event_ids}