[EDR Workflows][Fleet] Accurate endpoint count across multiple agent policies

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,
     },
   });
 };