[fleet] Roll forward PIT id in agent paging (#252536)

## Summary
- Roll forward refreshed `pit_id` in Fleet agent paging helpers/callers
and close the PIT using the most recent id.

## Test plan
- `yarn test:jest
x-pack/platform/plugins/shared/fleet/server/services/agents/crud.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/fleet/server/services/agents/action_runner.ts

@@ -213,7 +213,7 @@ export abstract class ActionRunner {
 
   async processAgentsInBatches(): Promise<{ actionId: string }> {
     const start = Date.now();
-    const pitId = this.retryParams.pitId;
+    let pitId = this.retryParams.pitId;
 
     const perPage = this.actionParams.batchSize ?? SO_SEARCH_LIMIT;
 
@@ -239,6 +239,10 @@ export abstract class ActionRunner {
     };
 
     const res = await getAgents();
+    if (res.pit) {
+      pitId = res.pit;
+      this.retryParams.pitId = pitId;
+    }
 
     let currentAgents = res.agents;
     if (currentAgents.length === 0) {
@@ -255,6 +259,10 @@ export abstract class ActionRunner {
       const lastAgent = currentAgents[currentAgents.length - 1];
       this.retryParams.searchAfter = lastAgent.sort!;
       const nextPage = await getAgents();
+      if (nextPage.pit) {
+        pitId = nextPage.pit;
+        this.retryParams.pitId = pitId;
+      }
       currentAgents = nextPage.agents;
       if (currentAgents.length === 0) {
         appContextService

x-pack/platform/plugins/shared/fleet/server/services/agents/crud.test.ts

@@ -81,9 +81,11 @@ describe('Agents CRUD test', () => {
     ids: string[],
     total: number,
     status: AgentStatus,
-    generateSource: (id: string) => Partial<Agent> = () => ({})
+    generateSource: (id: string) => Partial<Agent> = () => ({}),
+    pitId?: string
   ) {
     return {
+      ...(pitId ? { pit_id: pitId } : {}),
       hits: {
         total,
         hits: ids.map((id: string) => ({
@@ -175,6 +177,38 @@ describe('Agents CRUD test', () => {
   });
 
   describe('getAgentsByKuery', () => {
+    it('should roll forward PIT id from search responses', async () => {
+      searchMock.mockResolvedValueOnce(getEsResponse(['1'], 1, 'online', () => ({}), 'pit-2'));
+
+      const firstRes = await getAgentsByKuery(esClientMock, soClientMock, {
+        showAgentless: true,
+        showInactive: false,
+        pitId: 'pit-1',
+      });
+
+      expect(searchMock).toHaveBeenCalledWith(
+        expect.objectContaining({
+          pit: expect.objectContaining({ id: 'pit-1' }),
+        })
+      );
+      expect(firstRes.pit).toBe('pit-2');
+
+      searchMock.mockResolvedValueOnce(getEsResponse(['2'], 1, 'online', () => ({}), 'pit-3'));
+
+      const secondRes = await getAgentsByKuery(esClientMock, soClientMock, {
+        showAgentless: true,
+        showInactive: false,
+        pitId: firstRes.pit,
+      });
+
+      expect(searchMock).toHaveBeenLastCalledWith(
+        expect.objectContaining({
+          pit: expect.objectContaining({ id: 'pit-2' }),
+        })
+      );
+      expect(secondRes.pit).toBe('pit-3');
+    });
+
     it('should return upgradeable on first page', async () => {
       searchMock
         .mockImplementationOnce(() =>

x-pack/platform/plugins/shared/fleet/server/services/agents/crud.ts

@@ -224,6 +224,7 @@ export async function getAgentsByKuery(
   total: number;
   page: number;
   perPage: number;
+  pit?: string;
   statusSummary?: Record<AgentStatus, number>;
   aggregations?: Record<string, estypes.AggregationsAggregate>;
 }> {
@@ -294,7 +295,9 @@ export async function getAgentsByKuery(
     uninstalled: 0,
   };
 
-  const pitIdToUse = pitId || (openPit ? await openPointInTime(esClient, pitKeepAlive) : undefined);
+  const initialPitId =
+    pitId || (openPit ? await openPointInTime(esClient, pitKeepAlive) : undefined);
+  let currentPitId = initialPitId;
 
   const queryAgents = async (
     queryOptions: { from: number; size: number } | { searchAfter: SortResults; size: number }
@@ -334,10 +337,10 @@ export async function getAgentsByKuery(
       fields: Object.keys(runtimeFields),
       sort,
       query: kueryNode ? toElasticsearchQuery(kueryNode) : undefined,
-      ...(pitIdToUse
+      ...(currentPitId
         ? {
             pit: {
-              id: pitIdToUse,
+              id: currentPitId,
               keep_alive: pitKeepAlive,
             },
           }
@@ -361,6 +364,8 @@ export async function getAgentsByKuery(
     throw err;
   }
 
+  currentPitId = res.pit_id ?? currentPitId;
+
   let agents = res.hits.hits.map(searchHitToAgent);
   let total = res.hits.total as number;
   // filtering for a range on the version string will not work,
@@ -372,6 +377,7 @@ export async function getAgentsByKuery(
     // if there are more than SO_SEARCH_LIMIT agents, the logic falls back to same as before
     if (total < SO_SEARCH_LIMIT) {
       const response = await queryAgents({ from: 0, size: SO_SEARCH_LIMIT });
+      currentPitId = response.pit_id ?? currentPitId;
       agents = response.hits.hits
         .map(searchHitToAgent)
         .filter((agent) => isAgentUpgradeAvailable(agent, latestAgentVersion));
@@ -404,7 +410,7 @@ export async function getAgentsByKuery(
     total,
     ...(searchAfter ? { page: 0 } : { page }),
     perPage,
-    ...(pitIdToUse ? { pit: pitIdToUse } : {}),
+    ...(initialPitId ? { pit: currentPitId } : {}),
     ...(aggregations ? { aggregations: res.aggregations } : {}),
     ...(getStatusSummary ? { statusSummary } : {}),
   };

x-pack/platform/plugins/shared/fleet/server/services/spaces/agent_policy.ts

@@ -168,19 +168,22 @@ export async function updateAgentPolicySpaces({
 
   // Update agent actions
   if (agentIndexExists) {
-    const pitId = await openPointInTime(esClient);
+    let pitId = await openPointInTime(esClient);
 
     try {
       let hasMore = true;
       let searchAfter: SortResults | undefined;
       while (hasMore) {
-        const { agents } = await getAgentsByKuery(esClient, newSpaceSoClient, {
+        const { agents, pit } = await getAgentsByKuery(esClient, newSpaceSoClient, {
           kuery: `policy_id:"${agentPolicyId}"`,
           showInactive: true,
           perPage: UPDATE_AGENT_BATCH_SIZE,
           pitId,
           searchAfter,
         });
+        if (pit) {
+          pitId = pit;
+        }
 
         if (agents.length === 0) {
           hasMore = false;