diff --git a/src/constants/users.ts b/src/constants/users.ts index f636468e4..6f956fa34 100644 --- a/src/constants/users.ts +++ b/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 diff --git a/src/utils/CopilotAPI.test.ts b/src/utils/CopilotAPI.test.ts new file mode 100644 index 000000000..ac4f6c236 --- /dev/null +++ b/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: (fn: (...args: unknown[]) => Promise, 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', + }) + }) +}) diff --git a/src/utils/CopilotAPI.ts b/src/utils/CopilotAPI.ts index 20b721548..0c6bc2964 100644 --- a/src/utils/CopilotAPI.ts +++ b/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 { 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 => { + 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 {