Fix Copilot company pagination

src/constants/users.ts

@@ -1,3 +1,4 @@
 export const MAX_FETCH_ASSIGNEE_COUNT = 15_000
 
 export const MAX_LIMIT_CLIENT_COUNT = 5_000 //used to specify max fetching limit while querying for clients in copilot API.
+export const MAX_LIMIT_COMPANY_COUNT = 5_000

src/utils/CopilotAPI.test.ts

@@ -0,0 +1,74 @@
+const mockListCompanies = jest.fn()
+const mockCopilotApi = jest.fn((_args: unknown) => ({
+  listCompanies: mockListCompanies,
+}))
+
+jest.mock('@/config', () => ({
+  copilotAPIKey: 'test-api-key',
+  APP_ID: 'test-app-id',
+  assemblyApiDomain: 'https://api.example.test',
+}))
+
+jest.mock('copilot-node-sdk', () => ({
+  copilotApi: (args: unknown) => mockCopilotApi(args),
+}))
+
+jest.mock('@/app/api/core/utils/withRetry', () => ({
+  withRetry: <T>(fn: (...args: unknown[]) => Promise<T>, args: unknown[]) => fn(...args),
+}))
+
+import { CopilotAPI } from './CopilotAPI'
+
+const buildCompany = (id: string) => ({
+  id,
+  name: `Company ${id}`,
+  iconImageUrl: null,
+  createdAt: '2026-01-01T00:00:00.000Z',
+})
+
+const buildCompanies = ({ count, prefix }: { count: number; prefix: string }) =>
+  Array.from({ length: count }, (_, index) => buildCompany(`${prefix}-${index}`))
+
+describe('CopilotAPI getCompanies', () => {
+  beforeEach(() => {
+    jest.clearAllMocks()
+  })
+
+  it('delegates directly when the requested limit fits in one Copilot page', async () => {
+    const response = { data: [buildCompany('company-1')] }
+    mockListCompanies.mockResolvedValueOnce(response)
+
+    const copilot = new CopilotAPI('token')
+    const result = await copilot.getCompanies({ limit: 100, isPlaceholder: false })
+
+    expect(result).toEqual(response)
+    expect(mockListCompanies).toHaveBeenCalledTimes(1)
+    expect(mockListCompanies).toHaveBeenCalledWith({ limit: 100, isPlaceholder: false })
+  })
+
+  it('fetches companies in Copilot-sized pages when requesting more than one page', async () => {
+    const firstPage = buildCompanies({ count: 5_000, prefix: 'first' })
+    const secondPage = buildCompanies({ count: 1_000, prefix: 'second' })
+
+    mockListCompanies
+      .mockResolvedValueOnce({ data: firstPage, nextToken: 'next-page' })
+      .mockResolvedValueOnce({ data: secondPage, nextToken: 'unused-page' })
+
+    const copilot = new CopilotAPI('token')
+    const result = await copilot.getCompanies({ limit: 6_000, isPlaceholder: false })
+
+    expect(result.data).toHaveLength(6_000)
+    expect(result.data).toEqual([...firstPage, ...secondPage])
+    expect(mockListCompanies).toHaveBeenCalledTimes(2)
+    expect(mockListCompanies).toHaveBeenNthCalledWith(1, {
+      limit: 5_000,
+      isPlaceholder: false,
+      nextToken: undefined,
+    })
+    expect(mockListCompanies).toHaveBeenNthCalledWith(2, {
+      limit: 1_000,
+      isPlaceholder: false,
+      nextToken: 'next-page',
+    })
+  })
+})

src/utils/CopilotAPI.ts

@@ -2,7 +2,7 @@ import APIError from '@/app/api/core/exceptions/api'
 import httpStatus from 'http-status'
 import { withRetry } from '@/app/api/core/utils/withRetry'
 import { copilotAPIKey as apiKey, APP_ID, assemblyApiDomain } from '@/config'
-import { MAX_LIMIT_CLIENT_COUNT } from '@/constants/users'
+import { MAX_LIMIT_CLIENT_COUNT, MAX_LIMIT_COMPANY_COUNT } from '@/constants/users'
 import {
   AssemblyMetadata,
   ClientRequest,
@@ -44,6 +44,8 @@ import { copilotApi } from 'copilot-node-sdk'
 import { cache } from 'react'
 import { z } from 'zod'
 
+const CopilotPaginationResponseSchema = z.object({ nextToken: z.string().optional() })
+
 export class CopilotAPI {
   copilot: SDK
 
@@ -204,7 +206,38 @@ export class CopilotAPI {
 
   async _getCompanies(args: CopilotListArgs & { isPlaceholder?: boolean } = {}): Promise<CompaniesResponse> {
     console.info('CopilotAPI#_getCompanies', this.token)
-    return CompaniesResponseSchema.parse(await this.copilot.listCompanies(args))
+    const maxLimit = MAX_LIMIT_COMPANY_COUNT
+    const requestedLimit = args.limit || maxLimit
+
+    if (requestedLimit <= maxLimit) {
+      return CompaniesResponseSchema.parse(await this.copilot.listCompanies(args))
+    }
+
+    const fetchCompanies = async ({
+      nextToken,
+      companies,
+    }: {
+      nextToken?: string
+      companies: CompanyResponse[]
+    }): Promise<CompanyResponse[]> => {
+      const remaining = requestedLimit - companies.length
+      if (remaining <= 0) return companies
+
+      const response = await this.copilot.listCompanies({
+        ...args,
+        limit: Math.min(maxLimit, remaining),
+        nextToken,
+      })
+      const parsedResponse = CompaniesResponseSchema.parse(response)
+      const updatedCompanies = [...companies, ...(parsedResponse.data ?? [])]
+      const responseNextToken = CopilotPaginationResponseSchema.parse(response).nextToken
+
+      if (!responseNextToken || updatedCompanies.length >= requestedLimit) return updatedCompanies
+
+      return fetchCompanies({ nextToken: responseNextToken, companies: updatedCompanies })
+    }
+
+    return CompaniesResponseSchema.parse({ data: await fetchCompanies({ companies: [] }) })
   }
 
   async _getCompanyClients(companyId: string): Promise<ClientResponse[]> {