[EDR Workflows][Fleet] Accurate endpoint count across multiple agent policies (#193705)

This PR updates the method for counting endpoint statuses. Previously,
we fetched agent status using a single agent policy ID. With this
change, we now pass an array of policy IDs, allowing us to include the
returned stats for endpoints that share the same integration policy
assigned to multiple agent policies.

![Screenshot 2024-09-23 at 13 53
57](https://github.com/user-attachments/assets/570027b7-79d7-4c9a-aa64-c0ecfe76cb7f)
![Screenshot 2024-09-23 at 13 53
24](https://github.com/user-attachments/assets/17d62c24-9d46-4133-a817-ea5849930435)
![Screenshot 2024-09-23 at 13 53
45](https://github.com/user-attachments/assets/c9fb5ed7-e4a0-4faa-a24d-253def10f163)

x-pack/plugins/fleet/server/routes/agent/handlers.ts

@@ -328,7 +328,8 @@ export const getAgentStatusForAgentPolicyHandler: FleetRequestHandler<
       soClient,
       request.query.policyId,
       request.query.kuery,
-      coreContext.savedObjects.client.getCurrentNamespace()
+      coreContext.savedObjects.client.getCurrentNamespace(),
+      request.query.policyIds
     );
 
     const body: GetAgentStatusResponse = { results };

x-pack/plugins/fleet/server/services/agents/status.test.ts

@@ -7,6 +7,8 @@
 
 import { errors as EsErrors } from '@elastic/elasticsearch';
 
+import { AGENTS_INDEX } from '../../../common';
+
 import { createAppContextStartContractMock } from '../../mocks';
 
 import { appContextService } from '../app_context';
@@ -168,4 +170,67 @@ describe('getAgentStatusForAgentPolicy', () => {
 
     expect(esClient.search).toHaveBeenCalledTimes(2);
   });
+
+  it('calls esClient.search with correct parameters when agentPolicyIds are provided', async () => {
+    const esClient = {
+      search: jest.fn().mockResolvedValue({
+        aggregations: {
+          status: {
+            buckets: [
+              { key: 'online', doc_count: 2 },
+              { key: 'error', doc_count: 1 },
+            ],
+          },
+        },
+      }),
+    };
+
+    const soClient = {
+      find: jest.fn().mockResolvedValue({
+        saved_objects: [
+          { id: 'agentPolicyId1', attributes: { name: 'Policy 1' } },
+          { id: 'agentPolicyId2', attributes: { name: 'Policy 2' } },
+        ],
+      }),
+    };
+
+    const agentPolicyIds = ['agentPolicyId1', 'agentPolicyId2'];
+    const filterKuery = 'filterKuery';
+    const spaceId = 'spaceId';
+
+    await getAgentStatusForAgentPolicy(
+      esClient as any,
+      soClient as any,
+      undefined,
+      filterKuery,
+      spaceId,
+      agentPolicyIds
+    );
+
+    expect(esClient.search).toHaveBeenCalledWith(
+      expect.objectContaining({
+        index: AGENTS_INDEX,
+        size: 0,
+        query: expect.objectContaining({
+          bool: expect.objectContaining({
+            must: expect.arrayContaining([
+              expect.objectContaining({
+                terms: {
+                  policy_id: agentPolicyIds,
+                },
+              }),
+            ]),
+          }),
+        }),
+        aggregations: expect.objectContaining({
+          status: expect.objectContaining({
+            terms: expect.objectContaining({
+              field: 'status',
+              size: expect.any(Number),
+            }),
+          }),
+        }),
+      })
+    );
+  });
 });

x-pack/plugins/fleet/server/services/agents/status.ts

@@ -40,12 +40,23 @@ export async function getAgentStatusById(
   return (await getAgentById(esClient, soClient, agentId)).status!;
 }
 
+/**
+ * getAgentStatusForAgentPolicy
+ * @param esClient
+ * @param soClient
+ * @param agentPolicyId @deprecated use agentPolicyIds instead since the move to multi-policy
+ * @param filterKuery
+ * @param spaceId
+ * @param agentPolicyIds
+ */
+
 export async function getAgentStatusForAgentPolicy(
   esClient: ElasticsearchClient,
   soClient: SavedObjectsClientContract,
   agentPolicyId?: string,
   filterKuery?: string,
-  spaceId?: string
+  spaceId?: string,
+  agentPolicyIds?: string[]
 ) {
   const logger = appContextService.getLogger();
   const runtimeFields = await buildAgentStatusRuntimeField(soClient);
@@ -71,8 +82,14 @@ export async function getAgentStatusForAgentPolicy(
     );
     clauses.push(kueryAsElasticsearchQuery);
   }
-
-  if (agentPolicyId) {
+  // If agentPolicyIds is provided, we filter by those, otherwise we filter by depreciated agentPolicyId
+  if (agentPolicyIds) {
+    clauses.push({
+      terms: {
+        policy_id: agentPolicyIds,
+      },
+    });
+  } else if (agentPolicyId) {
     clauses.push({
       term: {
         policy_id: agentPolicyId,

x-pack/plugins/fleet/server/types/rest_spec/agent.ts

@@ -241,6 +241,7 @@ export const PostBulkUpdateAgentTagsRequestSchema = {
 export const GetAgentStatusRequestSchema = {
   query: schema.object({
     policyId: schema.maybe(schema.string()),
+    policyIds: schema.maybe(schema.arrayOf(schema.string())),
     kuery: schema.maybe(
       schema.string({
         validate: (value: string) => {

x-pack/plugins/security_solution/public/management/pages/policy/store/policy_details/middleware/policy_settings_middleware.ts

@@ -97,7 +97,7 @@ export const policySettingsMiddlewareRunner: MiddlewareRunner = async (
     // Agent summary is secondary data, so its ok for it to come after the details
     // page is populated with the main content
     if (policyItem.policy_id) {
-      const { results } = await sendGetFleetAgentStatusForPolicy(http, policyItem.policy_id);
+      const { results } = await sendGetFleetAgentStatusForPolicy(http, policyItem.policy_ids);
       dispatch({
         type: 'serverReturnedPolicyDetailsAgentSummaryData',
         payload: {

x-pack/plugins/security_solution/public/management/services/policies/ingest.ts

@@ -86,20 +86,20 @@ export const sendPutPackagePolicy = (
  * Get a status summary for all Agents that are currently assigned to a given agent policy
  *
  * @param http
- * @param policyId
+ * @param policyIds
  * @param options
  */
 export const sendGetFleetAgentStatusForPolicy = (
   http: HttpStart,
   /** the Agent (fleet) policy id */
-  policyId: string,
+  policyIds: string[],
   options: Exclude<HttpFetchOptions, 'query'> = {}
 ): Promise<GetAgentStatusResponse> => {
   return http.get(INGEST_API_FLEET_AGENT_STATUS, {
     ...options,
     version: API_VERSIONS.public.v1,
     query: {
-      policyId,
+      policyIds,
     },
   });
 };