Minimize amount of changes and maximize the number of affected features by moving the replies_count patch to the serializer directly

app/controllers/api/v1/statuses_controller.rb

@@ -14,6 +14,16 @@ class Api::V1::StatusesController < Api::BaseController
   override_rate_limit_headers :create, family: :statuses
   override_rate_limit_headers :update, family: :statuses
 
+  # This API was originally unlimited, pagination cannot be introduced without
+  # breaking backwards-compatibility. Arbitrarily high number to cover most
+  # conversations as quasi-unlimited, it would be too much work to render more
+  # than this anyway
+  CONTEXT_LIMIT = 4_096
+
+  # This remains expensive and we don't want to show everything to logged-out users
+  ANCESTORS_LIMIT         = 40
+  DESCENDANTS_LIMIT       = 60
+  DESCENDANTS_DEPTH_LIMIT = 20
 
   def index
     @statuses = preload_collection(@statuses, Status)
@@ -29,10 +39,25 @@ class Api::V1::StatusesController < Api::BaseController
   def context
     cache_if_unauthenticated!
 
-    @status_tree = StatusTree.new(status: @status, account: current_account)
+    ancestors_limit         = CONTEXT_LIMIT
-    @status_node = @status_tree.find_node(@status.id)
+    descendants_limit       = CONTEXT_LIMIT
+    descendants_depth_limit = nil
 
-    render json: @status_node, serializer: REST::StatusTreeNodeSerializer, relationships: StatusRelationshipsPresenter.new(@status_tree.flatten, current_user&.account_id)
+    if current_account.nil?
+      ancestors_limit         = ANCESTORS_LIMIT
+      descendants_limit       = DESCENDANTS_LIMIT
+      descendants_depth_limit = DESCENDANTS_DEPTH_LIMIT
+    end
+
+    ancestors_results   = @status.in_reply_to_id.nil? ? [] : @status.ancestors(ancestors_limit, current_account)
+    descendants_results = @status.descendants(descendants_limit, current_account, descendants_depth_limit)
+    loaded_ancestors    = preload_collection(ancestors_results, Status)
+    loaded_descendants  = preload_collection(descendants_results, Status)
+
+    @context = Context.new(ancestors: loaded_ancestors, descendants: loaded_descendants)
+    statuses = [@status] + @context.ancestors + @context.descendants
+
+    render json: @context, serializer: REST::ContextSerializer, relationships: StatusRelationshipsPresenter.new(statuses, current_user&.account_id)
   end
 
   def create

app/lib/activitypub/tag_manager.rb

@@ -37,8 +37,6 @@ class ActivityPub::TagManager
     return target.uri if target.respond_to?(:local?) && !target.local?
 
     case target.object_type
-    when :status
-      return activity_account_status_url(target.account, target) if target.reblog?
     when :person
       target.instance_actor? ? instance_actor_url : account_url(target)
     when :note, :comment, :activity

app/models/context.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class Context < ActiveModelSerializers::Model
+  attributes :ancestors, :descendants
+end

app/models/status_tree.rb

@@ -1,26 +1,17 @@
 class StatusTree < ActiveModelSerializers::Model
   include PreloadingConcern
 
-  # This API was originally unlimited, pagination cannot be introduced without
-  # breaking backwards-compatibility. Arbitrarily high number to cover most
-  # conversations as quasi-unlimited, it would be too much work to render more
-  # than this anyway
   MAX_COUNT = 4_096
 
-  # This remains expensive and we don't want to show everything to logged-out users
-  ANCESTORS_MAX_COUNT   = 40
-  DESCENDANTS_MAX_COUNT = 60
-  DESCENDANTS_MAX_DEPTH = 20  
-
   attributes :status, :account, :tree
 
   class Node < ActiveModelSerializers::Model
     attributes :status, :tree
 
-    delegate_missing_to :status
-
     delegate :id, to: :status
 
+    delegate_missing_to :status
+
     def object_type = :status
 
     def ancestors
@@ -35,16 +26,12 @@ class StatusTree < ActiveModelSerializers::Model
       tree.children_for(id)
     end
 
-    def replies_count
-      children.size
-    end
-
     def ==(other)
       other.class.in?([Node, Status]) && id == other.id
     end
 
     def inspect
-      "#<StatusTree::Node id: #{id}, parent_id: #{in_reply_to_id || 'nil'}>"
+      "#<StatusTree::Node id: #{id}, in_reply_to_id: #{in_reply_to_id || 'nil'}>"
     end
   end
 
@@ -77,6 +64,10 @@ class StatusTree < ActiveModelSerializers::Model
     "#<StatusTree #{tree.inspect}>"
   end
 
+  def status_node
+    find_node(status.id)
+  end
+
   def find_node(id, subtree = tree)
     subtree.each do |node, children|
       return node if node.id == id
@@ -89,13 +80,13 @@ class StatusTree < ActiveModelSerializers::Model
   def ancestors_for(id)
     ancestors = []
     node = find_node(id)
-    parent_id = node.in_reply_to_id
+    in_reply_to_id = node.in_reply_to_id
 
-    while parent_id
+    while in_reply_to_id
-      parent_node = find_node(parent_id)
+      parent_node = find_node(in_reply_to_id)
       break unless parent_node
       ancestors << parent_node
-      parent_id = parent_node.in_reply_to_id
+      in_reply_to_id = parent_node.in_reply_to_id
     end
 
     ancestors.reverse
@@ -116,24 +107,24 @@ class StatusTree < ActiveModelSerializers::Model
 
   private
 
-  def build_tree_from(nodes, parent_id = nil)
+  def build_tree_from(nodes, in_reply_to_id = nil)
     grouped_nodes = nodes.group_by(&:in_reply_to_id)
 
-    (grouped_nodes[parent_id] || []).each_with_object({}) do |node, tree|
+    (grouped_nodes[in_reply_to_id] || []).each_with_object({}) do |node, tree|
       tree[node] = build_tree_from(nodes - [node], node.id)
     end
   end
 
   def descendants_max_depth
-    account.nil? ? DESCENDANTS_MAX_DEPTH : nil
+    nil
   end
 
   def descendants_max_count
-    account.nil? ? DESCENDANTS_MAX_COUNT : MAX_COUNT
+    MAX_COUNT
   end
 
   def ancestors_max_count
-    account.nil? ? ANCESTORS_MAX_COUNT : MAX_COUNT
+    MAX_COUNT
   end
 
   def collect_descendants(subtree)

app/serializers/rest/context_serializer.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-class REST::StatusTreeNodeSerializer < ActiveModel::Serializer
+class REST::ContextSerializer < ActiveModel::Serializer
   has_many :ancestors,   serializer: REST::StatusSerializer
   has_many :descendants, serializer: REST::StatusSerializer
 end

app/serializers/rest/status_serializer.rb

@@ -91,6 +91,10 @@ class REST::StatusSerializer < ActiveModel::Serializer
     object.untrusted_favourites_count || relationships&.attributes_map&.dig(object.id, :favourites_count) || object.favourites_count
   end
 
+  def replies_count
+    StatusTree.new(status: object, account: current_user&.account).status_node.children.size
+  end
+
   def favourited
     if relationships
       relationships.favourites_map[object.id] || false

spec/serializers/rest/status_serializer_spec.rb

@@ -52,4 +52,77 @@ RSpec.describe REST::StatusSerializer do
       end
     end
   end
+
+  describe '#replies_count' do
+    let(:author) { alice }
+    let(:replier) { bob }
+    let!(:status) { Fabricate(:status, account: author, visibility: :public) }
+
+    context 'when being presented to the account that posted the status' do
+      let(:current_user) { Fabricate(:user, account: author) }
+
+      before do
+        Fabricate(:follow, account: replier, target_account: author)
+        Fabricate(:follow, account: author, target_account: replier)
+      end
+
+      context 'when the status has follower-only replies' do
+        let(:reply) { Fabricate(:status, in_reply_to_id: status.id, account: replier, visibility: :private) }
+
+        before do
+          reply
+        end
+
+        it 'counts 1 reply' do
+          expect(subject['replies_count']).to eq(1)
+        end
+
+        context 'when one of the replies has subsequent replies' do
+          before do
+            Fabricate(:status, in_reply_to_id: reply.id, account: author, visibility: :private)
+          end
+
+          it 'does not count that reply' do
+            expect(subject['replies_count']).to eq 1
+          end
+        end
+      end
+    end
+
+    context 'when being presented to a different account' do
+      let(:current_user) { Fabricate(:user) }
+
+      context 'when the status has follower-only replies from an unfollowed account' do
+        before do
+          Fabricate(:status, in_reply_to_id: status.id, account: replier, visibility: :direct)
+        end
+
+        it 'counts 0 replies' do
+          expect(subject['replies_count']).to be 0
+        end
+      end
+
+      context 'when the replies are public' do
+        before do
+          Fabricate(:status, in_reply_to_id: status.id, account: replier, visibility: :public)
+        end
+
+        it 'counts 1 reply' do
+          expect(subject['replies_count']).to eq 1
+        end
+      end
+
+      context 'when there is one public reply and one private' do
+        before do
+          %i[direct public].each do |visibility|
+            Fabricate(:status, in_reply_to_id: status.id, account: replier, visibility: visibility)
+          end
+        end
+
+        it 'counts 1 reply' do
+          expect(subject['replies_count']).to eq 1
+        end
+      end
+    end
+  end
 end