Add connected agents to cluster agents GraphQL response

https://gitlab.com/gitlab-org/gitlab/-/merge_requests/69820

Changelog: added
EE: true

doc/api/graphql/reference/index.md

@@ -5249,6 +5249,29 @@ The edge type for [`ComplianceFramework`](#complianceframework).
 | <a id="complianceframeworkedgecursor"></a>`cursor` | [`String!`](#string) | A cursor for use in pagination. |
 | <a id="complianceframeworkedgenode"></a>`node` | [`ComplianceFramework`](#complianceframework) | The item at the end of the edge. |
 
+#### `ConnectedAgentConnection`
+
+The connection type for [`ConnectedAgent`](#connectedagent).
+
+##### Fields
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="connectedagentconnectionedges"></a>`edges` | [`[ConnectedAgentEdge]`](#connectedagentedge) | A list of edges. |
+| <a id="connectedagentconnectionnodes"></a>`nodes` | [`[ConnectedAgent]`](#connectedagent) | A list of nodes. |
+| <a id="connectedagentconnectionpageinfo"></a>`pageInfo` | [`PageInfo!`](#pageinfo) | Information to aid in pagination. |
+
+#### `ConnectedAgentEdge`
+
+The edge type for [`ConnectedAgent`](#connectedagent).
+
+##### Fields
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="connectedagentedgecursor"></a>`cursor` | [`String!`](#string) | A cursor for use in pagination. |
+| <a id="connectedagentedgenode"></a>`node` | [`ConnectedAgent`](#connectedagent) | The item at the end of the edge. |
+
 #### `ContainerRepositoryConnection`
 
 The connection type for [`ContainerRepository`](#containerrepository).
@@ -8267,6 +8290,7 @@ GitLab CI/CD configuration template.
 
 | Name | Type | Description |
 | ---- | ---- | ----------- |
+| <a id="clusteragentconnections"></a>`connections` | [`ConnectedAgentConnection`](#connectedagentconnection) | Active connections for the cluster agent. (see [Connections](#connections)) |
 | <a id="clusteragentcreatedat"></a>`createdAt` | [`Time`](#time) | Timestamp the cluster agent was created. |
 | <a id="clusteragentcreatedbyuser"></a>`createdByUser` | [`UserCore`](#usercore) | User object, containing information about the person who created the agent. |
 | <a id="clusteragentid"></a>`id` | [`ID!`](#id) | ID of the cluster agent. |
@@ -8428,6 +8452,18 @@ Conan metadata.
 | <a id="conanmetadatarecipepath"></a>`recipePath` | [`String!`](#string) | Recipe path of the Conan package. |
 | <a id="conanmetadataupdatedat"></a>`updatedAt` | [`Time!`](#time) | Date of most recent update. |
 
+### `ConnectedAgent`
+
+Connection details for an Agent.
+
+#### Fields
+
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| <a id="connectedagentconnectedat"></a>`connectedAt` | [`Time`](#time) | When the connection was established. |
+| <a id="connectedagentconnectionid"></a>`connectionId` | [`BigInt`](#bigint) | ID of the connection. |
+| <a id="connectedagentmetadata"></a>`metadata` | [`JSON`](#json) | Information about the Agent. |
+
 ### `ContainerExpirationPolicy`
 
 A tag expiration policy designed to keep only the images that matter most.

ee/app/graphql/resolvers/kas/agent_connections_resolver.rb

+# frozen_string_literal: true
+
+module Resolvers
+  module Kas
+    class AgentConnectionsResolver < BaseResolver
+      type Types::Kas::AgentConnectionType, null: true
+
+      alias_method :agent, :object
+
+      delegate :project, to: :agent
+
+      def resolve
+        return [] unless can_read_connected_agents?
+
+        BatchLoader::GraphQL.for(agent.id).batch(key: project, default_value: []) do |agent_ids, loader|
+          agents = get_connected_agents.group_by(&:agent_id).slice(*agent_ids)
+
+          agents.each do |agent_id, connections|
+            loader.call(agent_id, connections)
+          end
+        end
+      end
+
+      private
+
+      def can_read_connected_agents?
+        project.licensed_feature_available?(:cluster_agents) && current_user.can?(:admin_cluster, project)
+      end
+
+      def get_connected_agents
+        kas_client.get_connected_agents(project: project)
+      rescue GRPC::BadStatus => e
+        raise Gitlab::Graphql::Errors::ResourceNotAvailable, e.class.name
+      end
+
+      def kas_client
+        @kas_client ||= Gitlab::Kas::Client.new
+      end
+    end
+  end
+end

ee/app/graphql/types/clusters/agent_type.rb

@@ -48,6 +48,13 @@ module Types
             null: true,
             description: 'Web path of the cluster agent.'
 
+      field :connections,
+            Types::Kas::AgentConnectionType.connection_type,
+            null: true,
+            description: 'Active connections for the cluster agent',
+            complexity: 5,
+            resolver: ::Resolvers::Kas::AgentConnectionsResolver
+
       def project
         Gitlab::Graphql::Loaders::BatchModelLoader.new(Project, object.project_id).find
       end

ee/app/graphql/types/kas/agent_connection_type.rb

+# frozen_string_literal: true
+
+module Types
+  module Kas
+    # rubocop: disable Graphql/AuthorizeTypes
+    class AgentConnectionType < BaseObject
+      graphql_name 'ConnectedAgent'
+      description 'Connection details for an Agent'
+
+      field :connected_at,
+            Types::TimeType,
+            null: true,
+            description: 'When the connection was established.'
+
+      field :connection_id,
+            GraphQL::Types::BigInt,
+            null: true,
+            description: 'ID of the connection.'
+
+      field :metadata, # rubocop:disable Graphql/JSONType
+            GraphQL::Types::JSON,
+            method: :agent_meta,
+            null: true,
+            description: 'Information about the Agent.'
+
+      def connected_at
+        Time.at(object.connected_at.seconds)
+      end
+    end
+    # rubocop: enable Graphql/AuthorizeTypes
+  end
+end

ee/spec/graphql/resolvers/kas/agent_connections_resolver_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Resolvers::Kas::AgentConnectionsResolver do
+  include GraphqlHelpers
+
+  it { expect(described_class.type).to eq(Types::Kas::AgentConnectionType) }
+  it { expect(described_class.null).to be_truthy }
+
+  describe '#resolve' do
+    let_it_be(:project) { create(:project) }
+    let_it_be(:agent1) { create(:cluster_agent, project: project) }
+    let_it_be(:agent2) { create(:cluster_agent, project: project) }
+
+    let(:user) { create(:user, maintainer_projects: [project]) }
+    let(:ctx) { Hash(current_user: user) }
+    let(:feature_available) { true }
+
+    let(:connection1) { double(agent_id: agent1.id) }
+    let(:connection2) { double(agent_id: agent1.id) }
+    let(:connection3) { double(agent_id: agent2.id) }
+    let(:connected_agents) { [connection1, connection2, connection3] }
+    let(:kas_client) { instance_double(Gitlab::Kas::Client, get_connected_agents: connected_agents) }
+
+    subject do
+      batch_sync do
+        resolve(described_class, obj: agent1, ctx: ctx)
+      end
+    end
+
+    before do
+      stub_licensed_features(cluster_agents: feature_available)
+
+      allow(Gitlab::Kas::Client).to receive(:new).and_return(kas_client)
+    end
+
+    it 'returns active connections for the agent' do
+      expect(subject).to contain_exactly(connection1, connection2)
+    end
+
+    it 'queries KAS once when multiple agents are requested' do
+      expect(kas_client).to receive(:get_connected_agents).once
+
+      response = batch_sync do
+        resolve(described_class, obj: agent1, ctx: ctx)
+        resolve(described_class, obj: agent2, ctx: ctx)
+      end
+
+      expect(response).to contain_exactly(connection3)
+    end
+
+    context 'an error is returned from the KAS client' do
+      before do
+        allow(kas_client).to receive(:get_connected_agents).and_raise(GRPC::DeadlineExceeded)
+      end
+
+      it 'raises a graphql error' do
+        expect { subject }.to raise_error(Gitlab::Graphql::Errors::ResourceNotAvailable, 'GRPC::DeadlineExceeded')
+      end
+    end
+
+    context 'feature is not available' do
+      let(:feature_available) { false }
+
+      it { is_expected.to be_empty }
+    end
+
+    context 'user does not have permission' do
+      let(:user) { create(:user) }
+
+      it { is_expected.to be_empty }
+    end
+  end
+end

ee/spec/graphql/types/clusters/agent_type_spec.rb

@@ -3,7 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe GitlabSchema.types['ClusterAgent'] do
-  let(:fields) { %i[created_at created_by_user id name project updated_at tokens web_path] }
+  let(:fields) { %i[created_at created_by_user id name project updated_at tokens web_path connections] }
 
   it { expect(described_class.graphql_name).to eq('ClusterAgent') }
 

ee/spec/graphql/types/kas/agent_connection_type_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Types::Kas::AgentConnectionType do
+  include GraphqlHelpers
+
+  let(:fields) { %i[connected_at connection_id metadata] }
+
+  it { expect(described_class.graphql_name).to eq('ConnectedAgent') }
+  it { expect(described_class.description).to eq('Connection details for an Agent') }
+  it { expect(described_class).to have_graphql_fields(fields) }
+
+  describe '#connected_at' do
+    let(:connected_at) { double(Google::Protobuf::Timestamp, seconds: 123456, nanos: 654321) }
+    let(:object) { double(Gitlab::Agent::AgentTracker::ConnectedAgentInfo, connected_at: connected_at) }
+
+    it 'converts the seconds value to a timestamp' do
+      expect(resolve_field(:connected_at, object)).to eq(Time.at(connected_at.seconds))
+    end
+  end
+end

ee/spec/requests/api/graphql/project/cluster_agents_spec.rb

@@ -22,6 +22,7 @@ RSpec.describe 'Project.cluster_agents' do
   end
 
   before do
+    allow(Gitlab::Kas::Client).to receive(:new).and_return(double(get_connected_agents: []))
     stub_licensed_features(cluster_agents: true)
   end
 

lib/gitlab/kas/client.rb

@@ -7,6 +7,7 @@ module Gitlab
       JWT_AUDIENCE = 'gitlab-kas'
 
       STUB_CLASSES = {
+        agent_tracker: Gitlab::Agent::AgentTracker::Rpc::AgentTracker::Stub,
         configuration_project: Gitlab::Agent::ConfigurationProject::Rpc::ConfigurationProject::Stub
       }.freeze
 
@@ -17,6 +18,15 @@ module Gitlab
         raise ConfigurationError, 'KAS internal URL is not configured' unless Gitlab::Kas.internal_url.present?
       end
 
+      def get_connected_agents(project:)
+        request = Gitlab::Agent::AgentTracker::Rpc::GetConnectedAgentsRequest.new(project_id: project.id)
+
+        stub_for(:agent_tracker)
+          .get_connected_agents(request, metadata: metadata)
+          .agents
+          .to_a
+      end
+
       def list_agent_config_files(project:)
         request = Gitlab::Agent::ConfigurationProject::Rpc::ListAgentConfigFilesRequest.new(
           repository: repository(project),
@@ -49,7 +59,7 @@ module Gitlab
       end
 
       def kas_endpoint_url
-        Gitlab::Kas.internal_url.sub(%r{^grpc://|^grpcs://}, '')
+        Gitlab::Kas.internal_url.sub(%r{^grpcs?://}, '')
       end
 
       def credentials

spec/lib/gitlab/kas/client_spec.rb

@@ -4,6 +4,7 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Kas::Client do
   let_it_be(:project) { create(:project) }
+  let_it_be(:agent) { create(:cluster_agent, project: project) }
 
   describe '#initialize' do
     context 'kas is not enabled' do
@@ -44,6 +45,32 @@ RSpec.describe Gitlab::Kas::Client do
       expect(token).to receive(:audience=).with(described_class::JWT_AUDIENCE)
     end
 
+    describe '#get_connected_agents' do
+      let(:stub) { instance_double(Gitlab::Agent::AgentTracker::Rpc::AgentTracker::Stub) }
+      let(:request) { instance_double(Gitlab::Agent::AgentTracker::Rpc::GetConnectedAgentsRequest) }
+      let(:response) { double(Gitlab::Agent::AgentTracker::Rpc::GetConnectedAgentsResponse, agents: connected_agents) }
+
+      let(:connected_agents) { [double] }
+
+      subject { described_class.new.get_connected_agents(project: project) }
+
+      before do
+        expect(Gitlab::Agent::AgentTracker::Rpc::AgentTracker::Stub).to receive(:new)
+          .with('example.kas.internal', :this_channel_is_insecure, timeout: described_class::TIMEOUT)
+          .and_return(stub)
+
+        expect(Gitlab::Agent::AgentTracker::Rpc::GetConnectedAgentsRequest).to receive(:new)
+          .with(project_id: project.id)
+          .and_return(request)
+
+        expect(stub).to receive(:get_connected_agents)
+          .with(request, metadata: { 'authorization' => 'bearer test-token' })
+          .and_return(response)
+      end
+
+      it { expect(subject).to eq(connected_agents) }
+    end
+
     describe '#list_agent_config_files' do
       let(:stub) { instance_double(Gitlab::Agent::ConfigurationProject::Rpc::ConfigurationProject::Stub) }