2 files changed, 37 insertions, 5 deletions

diff --git a/app/javascript/mastodon/actions/timelines.js b/app/javascript/mastodon/actions/timelines.js
index 5c0cd93c7..cdaafd89c 100644
--- a/app/javascript/mastodon/actions/timelines.js
+++ b/app/javascript/mastodon/actions/timelines.js
@@ -17,6 +17,8 @@ export const TIMELINE_SCROLL_TOP = 'TIMELINE_SCROLL_TOP';
 export const TIMELINE_CONNECT    = 'TIMELINE_CONNECT';
 export const TIMELINE_DISCONNECT = 'TIMELINE_DISCONNECT';
 
+export const TIMELINE_CONTEXT_UPDATE = 'CONTEXT_UPDATE';
+
 export function refreshTimelineSuccess(timeline, statuses, skipLoading, next) {
   return {
     type: TIMELINE_REFRESH_SUCCESS,
@@ -30,6 +32,16 @@ export function refreshTimelineSuccess(timeline, statuses, skipLoading, next) {
 export function updateTimeline(timeline, status) {
   return (dispatch, getState) => {
     const references = status.reblog ? getState().get('statuses').filter((item, itemId) => (itemId === status.reblog.id || item.get('reblog') === status.reblog.id)).map((_, itemId) => itemId) : [];
+    const parents = [];
+
+    if (status.in_reply_to_id) {
+      let parent = getState().getIn(['statuses', status.in_reply_to_id]);
+
+      while (parent.get('in_reply_to_id')) {
+        parents.push(parent.get('id'));
+        parent = getState().getIn(['statuses', parent.get('in_reply_to_id')]);
+      }
+    }
 
     dispatch({
       type: TIMELINE_UPDATE,
@@ -37,6 +49,14 @@ export function updateTimeline(timeline, status) {
       status,
       references,
     });
+
+    if (parents.length > 0) {
+      dispatch({
+        type: TIMELINE_CONTEXT_UPDATE,
+        status,
+        references: parents,
+      });
+    }
   };
 };
 
diff --git a/app/javascript/mastodon/reducers/contexts.js b/app/javascript/mastodon/reducers/contexts.js
index 9bfc09aa7..d8924e908 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,24 @@ 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 => 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;
   }