1 files changed, 23 insertions, 5 deletions

diff --git a/app/javascript/mastodon/reducers/contexts.js b/app/javascript/mastodon/reducers/contexts.js
index 9bfc09aa7..64d584a01 100644
--- a/app/javascript/mastodon/reducers/contexts.js
+++ b/app/javascript/mastodon/reducers/contexts.js
@@ -1,6 +1,6 @@
 import { CONTEXT_FETCH_SUCCESS } from '../actions/statuses';
-import { TIMELINE_DELETE } from '../actions/timelines';
-import { Map as ImmutableMap, List as ImmutableList, fromJS } from 'immutable';
+import { TIMELINE_DELETE, TIMELINE_CONTEXT_UPDATE } from '../actions/timelines';
+import { Map as ImmutableMap, List as ImmutableList } from 'immutable';
 
 const initialState = ImmutableMap({
   ancestors: ImmutableMap(),
@@ -8,8 +8,8 @@ const initialState = ImmutableMap({
 });
 
 const normalizeContext = (state, id, ancestors, descendants) => {
-  const ancestorsIds   = ancestors.map(ancestor => ancestor.get('id'));
-  const descendantsIds = descendants.map(descendant => descendant.get('id'));
+  const ancestorsIds   = ImmutableList(ancestors.map(ancestor => ancestor.id));
+  const descendantsIds = ImmutableList(descendants.map(descendant => descendant.id));
 
   return state.withMutations(map => {
     map.setIn(['ancestors', id], ancestorsIds);
@@ -31,12 +31,30 @@ const deleteFromContexts = (state, id) => {
   return state;
 };
 
+const updateContext = (state, status, references) => {
+  return state.update('descendants', map => {
+    references.forEach(parentId => {
+      map = map.update(parentId, ImmutableList(), list => {
+        if (list.includes(status.id)) {
+          return list;
+        }
+
+        return list.push(status.id);
+      });
+    });
+
+    return map;
+  });
+};
+
 export default function contexts(state = initialState, action) {
   switch(action.type) {
   case CONTEXT_FETCH_SUCCESS:
-    return normalizeContext(state, action.id, fromJS(action.ancestors), fromJS(action.descendants));
+    return normalizeContext(state, action.id, action.ancestors, action.descendants);
   case TIMELINE_DELETE:
     return deleteFromContexts(state, action.id);
+  case TIMELINE_CONTEXT_UPDATE:
+    return updateContext(state, action.status, action.references);
   default:
     return state;
   }