[security-defend-workflows] Roll forward PIT id in osquery paging (#252535)

## Summary
- Thread the latest PIT id through osquery agent group paging and close
with the final refreshed id.

## Test plan
- `yarn test:jest
x-pack/platform/plugins/shared/osquery/server/lib/parse_agent_groups.test.ts`

Refs: ralph issue-236; #252458

Made with [Cursor](https://cursor.com)

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: gsoldevila

x-pack/platform/plugins/shared/osquery/server/lib/parse_agent_groups.test.ts

@@ -19,14 +19,15 @@ function createPaginatedMockResponse(totalAgents: number, chunkSize = 9000) {
   return jest.fn(
     ({
       searchAfter,
+      pitId,
     }: {
       searchAfter?: SortResults;
       kuery?: string;
       showInactive?: boolean;
       pitId?: string;
       perPage?: number;
       page?: number;
-    }): Promise<{ agents: Agent[]; total: number }> => {
+    }): Promise<{ agents: Agent[]; total: number; pit?: string }> => {
       const start = searchAfter ? (searchAfter[0] as number) + 1 : 0;
       const end = Math.min(start + chunkSize, agentIds.length);
       const chunk = agentIds.slice(start, end);
@@ -37,6 +38,7 @@ function createPaginatedMockResponse(totalAgents: number, chunkSize = 9000) {
           sort: [start + index],
         })) as Agent[],
         total: agentIds.length,
+        ...(pitId ? { pit: `${pitId}-next` } : {}),
       });
     }
   );
@@ -137,31 +139,39 @@ describe('aggregateResults', () => {
         results: generateResults(),
         total: 18001,
         searchAfter: ['firstSort'],
+        pitId: 'refreshedPit-1',
       })
       .mockResolvedValueOnce({
         results: generateResults(2),
         total: 18001,
         searchAfter: ['secondSort'],
+        pitId: 'refreshedPit-2',
       })
       .mockResolvedValueOnce({
         results: ['result_18001'],
         total: 18001,
         searchAfter: ['thirdSort'],
+        pitId: 'refreshedPit-3',
       });
 
     const result = await aggregateResults(generatorMock, mockElasticsearchClient, mockContext);
     expect(generatorMock).toHaveBeenCalledWith(1, expect.any(Number));
     expect(generatorMock).toHaveBeenCalledWith(1, expect.any(Number), undefined, 'mockedPitId');
-    expect(generatorMock).toHaveBeenCalledWith(2, expect.any(Number), ['firstSort'], 'mockedPitId');
+    expect(generatorMock).toHaveBeenCalledWith(
+      2,
+      expect.any(Number),
+      ['firstSort'],
+      'refreshedPit-1'
+    );
     expect(generatorMock).toHaveBeenCalledWith(
       3,
       expect.any(Number),
       ['secondSort'],
-      'mockedPitId'
+      'refreshedPit-2'
     );
     expect(mockOpenPointInTime).toHaveBeenCalledTimes(1);
     expect(mockClosePointInTime).toHaveBeenCalledTimes(1);
-    expect(mockClosePointInTime).toHaveBeenCalledWith({ id: 'mockedPitId' });
+    expect(mockClosePointInTime).toHaveBeenCalledWith({ id: 'refreshedPit-3' });
     expect(result.length).toEqual(18001);
   });
 });
@@ -249,11 +259,11 @@ describe('parseAgentSelection', () => {
         index: '.fleet-agents',
         keep_alive: '10m',
       });
-      expect(mockClosePointInTime).toHaveBeenCalledWith({ id: 'mockedPitId' });
+      expect(mockClosePointInTime).toHaveBeenCalledWith({ id: 'mockedPitId-next-next' });
 
       const thirdCall = mockAgentService.listAgents.mock.calls[2][0];
       expect(thirdCall.searchAfter).toBeDefined();
-      expect(thirdCall.pitId).toBe('mockedPitId');
+      expect(thirdCall.pitId).toBe('mockedPitId-next');
     });
 
     it('should continue even if PIT close fails', async () => {

x-pack/platform/plugins/shared/osquery/server/lib/parse_agent_groups.ts

@@ -28,7 +28,7 @@ export const aggregateResults = async (
     perPage: number,
     searchAfter?: SortResults,
     pitId?: string
-  ) => Promise<{ results: string[]; total: number; searchAfter?: SortResults }>,
+  ) => Promise<{ results: string[]; total: number; searchAfter?: SortResults; pitId?: string }>,
   esClient: ElasticsearchClient,
   context: OsqueryAppContext
 ) => {
@@ -39,30 +39,37 @@ export const aggregateResults = async (
     // One page only, no need for PIT
     results = initialResults;
   } else {
-    const { id: pitId } = await esClient.openPointInTime({
-      index: AGENTS_INDEX,
-      keep_alive: '10m',
-    });
+    let pitId = (
+      await esClient.openPointInTime({
+        index: AGENTS_INDEX,
+        keep_alive: '10m',
+      })
+    ).id;
     let currentSort: SortResults | undefined;
     // Refetch first page with PIT
-    const { results: pitInitialResults, searchAfter } = await generator(
+    const {
+      results: pitInitialResults,
+      searchAfter,
+      pitId: returnedPitId,
+    } = await generator(
       1,
       PER_PAGE,
       currentSort, // No searchAfter for first page, its built based on first page results
       pitId
     );
     results = pitInitialResults;
     currentSort = searchAfter;
+    pitId = returnedPitId ?? pitId;
     let currPage = 2;
     while (currPage <= totalPages) {
-      const { results: additionalResults, searchAfter: additionalSearchAfter } = await generator(
-        currPage++,
-        PER_PAGE,
-        currentSort,
-        pitId
-      );
+      const {
+        results: additionalResults,
+        searchAfter: additionalSearchAfter,
+        pitId: additionalPitId,
+      } = await generator(currPage++, PER_PAGE, currentSort, pitId);
       results.push(...additionalResults);
       currentSort = additionalSearchAfter;
+      pitId = additionalPitId ?? pitId;
     }
 
     try {
@@ -171,6 +178,7 @@ export const parseAgentSelection = async (
               res.agents.length > 0 && res.agents[res.agents.length - 1].sort
                 ? res.agents[res.agents.length - 1].sort
                 : undefined,
+            pitId: res.pit,
           };
         },
         esClient,
@@ -208,6 +216,7 @@ export const parseAgentSelection = async (
                 res.agents.length > 0 && res.agents[res.agents.length - 1].sort
                   ? res.agents[res.agents.length - 1].sort
                   : undefined,
+              pitId: res.pit,
             };
           },
           esClient,