From f607fb644f784b0cf6821d7d57b2a6b7bdf509e7 Mon Sep 17 00:00:00 2001 From: CD Cabrera Date: Tue, 8 Sep 2026 15:00:28 -0400 Subject: [PATCH 1/2] feat(collections): pf-4402 start delay, initial hydrate --- .../__snapshots__/collections.test.ts.snap | 6 +- .../__snapshots__/server.task.test.ts.snap | 23 +++ src/__tests__/collections.test.ts | 93 +++++++++++- src/__tests__/server.task.test.ts | 19 +++ src/collections.ts | 139 +++++++++++++++++- src/server.collections.ts | 14 +- src/server.task.ts | 31 +++- 7 files changed, 311 insertions(+), 14 deletions(-) diff --git a/src/__tests__/__snapshots__/collections.test.ts.snap b/src/__tests__/__snapshots__/collections.test.ts.snap index c673bf61..c2b713fc 100644 --- a/src/__tests__/__snapshots__/collections.test.ts.snap +++ b/src/__tests__/__snapshots__/collections.test.ts.snap @@ -1,12 +1,14 @@ // Jest Snapshot v1, https://jestjs.io/docs/snapshot-testing -exports[`registerCollections should call onSettle with all results (fulfilled and rejected) 1`] = ` +exports[`registerCollections should call onSettle with all results, fulfilled and rejected 1`] = ` { "fulfilled": [ { "records": [ { "id": "1", + "sourceId": "mock", + "sourceType": "mock", }, ], }, @@ -26,6 +28,8 @@ exports[`registerCollections should call onSettle with all results (fulfilled an "records": [ { "id": "1", + "sourceId": "mock", + "sourceType": "mock", }, ], }, diff --git a/src/__tests__/__snapshots__/server.task.test.ts.snap b/src/__tests__/__snapshots__/server.task.test.ts.snap index a339fd21..d37299e3 100644 --- a/src/__tests__/__snapshots__/server.task.test.ts.snap +++ b/src/__tests__/__snapshots__/server.task.test.ts.snap @@ -60,6 +60,29 @@ exports[`deferTask should cancel a task 1`] = ` ] `; +exports[`deferTask should delay the first execution when delayStartMs is provided 1`] = ` +[ + [ + { + "type": "start", + "value": [Function], + }, + ], + [ + { + "type": "run:delay", + "value": [Function], + }, + ], + [ + { + "type": "run", + "value": [Function], + }, + ], +] +`; + exports[`deferTask should enforce a timeout 1`] = ` [ [ diff --git a/src/__tests__/collections.test.ts b/src/__tests__/collections.test.ts index 6639b9cb..726cd4b3 100644 --- a/src/__tests__/collections.test.ts +++ b/src/__tests__/collections.test.ts @@ -216,9 +216,91 @@ describe('registerCollections', () => { await expect(registerCollections(collections)).resolves.not.toThrow(); }); + it('should immediately hydrate serverRecordsRegistry when config.initial is provided', async () => { + const initialRecords = [{ id: 'init-1', sourceId: 'local', sourceType: 'api' }] as any; + let resolveHandler: (res: any) => void; + const asyncPromise = new Promise(resolve => { + resolveHandler = resolve as (res: any) => void; + }); + const handler = jest.fn().mockImplementation(() => asyncPromise); + + const collections: any[] = [ + ['dual-phase-coll', handler, { initial: { records: initialRecords } }] + ]; + + const registrationPromise = registerCollections(collections); + + // Immediate check: serverRecordsRegistry has initial records before handler finishes + expect(getServerRecordsRegistry({ collectionName: 'dual-phase-coll' })).toEqual({ records: initialRecords }); + + resolveHandler!({ records: [{ id: 'live-1', sourceId: 'live', sourceType: 'api' }] }); + await registrationPromise; + }); + + it('should retain previous viable records when retainLastViable is true and update returns empty records', async () => { + const initialRecords = [{ id: 'init-1', sourceId: 'local', sourceType: 'api' }] as any; + const handler = jest.fn().mockResolvedValue({ records: [] }); + + const collections: any[] = [ + ['retained-coll', handler, { initial: { records: initialRecords }, retainLastViable: true }] + ]; + + await registerCollections(collections); + + // Retains initialRecords because update returned empty records + expect(getServerRecordsRegistry({ collectionName: 'retained-coll' })).toEqual({ records: initialRecords }); + }); + + it('should retain previous viable records when retainLastViable is true and update throws an error', async () => { + const initialRecords = [{ id: 'init-1', sourceId: 'local', sourceType: 'api' }] as any; + const handler = jest.fn().mockRejectedValue(new Error('Network failure')); + + const collections: any[] = [ + ['error-retained-coll', handler, { initial: { records: initialRecords }, retainLastViable: true }] + ]; + + await registerCollections(collections); + + // Retains initialRecords because update threw an error + expect(getServerRecordsRegistry({ collectionName: 'error-retained-coll' })).toEqual({ records: initialRecords }); + }); + + it('should support custom predicate function for retainLastViable', async () => { + const initialRecords = [ + { id: 'init-1', sourceId: 'mock', sourceType: 'mock' }, + { id: 'init-2', sourceId: 'mock', sourceType: 'mock' }, + { id: 'init-3', sourceId: 'mock', sourceType: 'mock' } + ]; + // Crawl returned only 1 record (loss of > 50% data) + const handler = jest.fn().mockResolvedValue({ records: [{ id: 'init-1', sourceId: 'mock', sourceType: 'mock' }] }); + const customPredicate = jest.fn().mockImplementation(({ previous, current }) => { + const prevCount = previous?.records?.length || 0; + const newCount = current?.records?.length || 0; + + return newCount < prevCount * 0.5; + }); + + const collections: any[] = [ + ['custom-predicate-coll', handler, { + initial: { records: initialRecords }, + retainLastViable: customPredicate + }] + ]; + + await registerCollections(collections); + + expect(customPredicate).toHaveBeenCalledWith(expect.objectContaining({ + name: 'custom-predicate-coll', + previous: { records: initialRecords }, + current: { records: [{ id: 'init-1', sourceId: 'mock', sourceType: 'mock' }] }, + isSuccess: true + })); + expect(getServerRecordsRegistry({ collectionName: 'custom-predicate-coll' })).toEqual({ records: initialRecords }); + }); + it('should call onRequired when all required collections are settled', async () => { const onRequired = jest.fn(); - const handler = jest.fn().mockResolvedValue({ records: [{ id: '1' }] }); + const handler = jest.fn().mockResolvedValue({ records: [{ id: '1', sourceId: 'mock', sourceType: 'mock' }] }); const collections: any[] = [ ['req', handler, { isRequired: true }] ]; @@ -226,11 +308,11 @@ describe('registerCollections', () => { await registerCollections(collections, { onRequired }); expect(onRequired).toHaveBeenCalledWith([ - expect.objectContaining({ name: 'req', response: { records: [{ id: '1' }] } }) + expect.objectContaining({ name: 'req', response: { records: [{ id: '1', sourceId: 'mock', sourceType: 'mock' }] } }) ]); }); - it('should call onSettle with all results (fulfilled and rejected)', async () => { + it('should call onSettle with all results, fulfilled and rejected', async () => { let settlePromiseResolve: (value: any) => void; const settlePromise = new Promise(resolve => { settlePromiseResolve = resolve; @@ -238,7 +320,7 @@ describe('registerCollections', () => { const onSettle = jest.fn(results => settlePromiseResolve(results)); - const handler1 = jest.fn().mockResolvedValue({ records: [{ id: '1' }] }); + const handler1 = jest.fn().mockResolvedValue({ records: [{ id: '1', sourceId: 'mock', sourceType: 'mock' }] }); const handler2 = jest.fn().mockRejectedValue(new Error('Fail')); const collections: any[] = [ @@ -250,7 +332,8 @@ describe('registerCollections', () => { const results: any = await settlePromise; expect(results).toMatchSnapshot(); - expect(results.fulfilled).toContainEqual({ records: [{ id: '1' }] }); + + expect(results.fulfilled).toContainEqual({ records: [{ id: '1', sourceId: 'mock', sourceType: 'mock' }] }); expect(results.rejected).toContainEqual(expect.objectContaining({ name: 'c2' })); }); diff --git a/src/__tests__/server.task.test.ts b/src/__tests__/server.task.test.ts index b2a75849..c1fd5aba 100644 --- a/src/__tests__/server.task.test.ts +++ b/src/__tests__/server.task.test.ts @@ -158,6 +158,25 @@ describe('deferTask', () => { expect(mockDebug.mock.calls).toMatchSnapshot(); }); + + it('should delay the first execution when delayStartMs is provided', async () => { + const mockDebug = jest.fn(); + const mockFunc = jest.fn().mockReturnValue('delayed'); + const handle = deferTask(mockFunc, { debug: mockDebug, delayStartMs: 1000, intervalMs: 100, repeat: 1 })(); + const promise = handle.start(); + + // Task delayed + await jest.advanceTimersByTimeAsync(500); + expect(mockFunc).not.toHaveBeenCalled(); + + // Delay finishes and task executes + await jest.advanceTimersByTimeAsync(500); + const result = await promise; + + expect(result).toBe('delayed'); + expect(mockFunc).toHaveBeenCalledTimes(1); + expect(mockDebug.mock.calls).toMatchSnapshot(); + }); }); describe('delay', () => { diff --git a/src/collections.ts b/src/collections.ts index f22e5856..ee2dd6c0 100644 --- a/src/collections.ts +++ b/src/collections.ts @@ -1,3 +1,4 @@ +import { isPlainObject } from './server.helpers'; import { formatUnknownError, log } from './logger'; import { type GlobalOptions } from './options'; @@ -49,11 +50,17 @@ interface McpCollectionResult { * 1. `handler` `{Function}`: callback function accepting an optional argument * 2. `_config` `{Object}`: Application level record source configuration. Unavailable to * record collection plugins. + * - `_config.initial`: Optional initial collection records or loader function executed + * immediately at server startup prior to background scheduled runs or worker execution. + * Hydrates the server records registry at $t=0$. * - `_config.runParallel`: Optional internal import specifier (`#specifier`) to run the * collection handler in a worker thread via the heavy pool. The referenced * module must export `collectionCallback`. Applied in {@link composeCollections}. * - `_config.runSchedule`: Optional object to dynamically decide if the record source * should run in a scheduled interval using {@link DeferTaskOptions} + * - `_config.retainLastViable`: Optional boolean or custom function to retain previously + * viable collection records in the registry if an update fails, drops to zero records, + * or triggers custom retention conditions. * - `_config.isRequired`: Optional boolean used to control server startup when * collections are required for operation. * - `_config._isInternal`: Optional boolean. Applied internally. Attempting to manually @@ -63,13 +70,16 @@ type McpCollection = [ name: string, handler: (arg?: unknown) => McpCollectionResult | Promise, _config?: { + initial?: McpCollectionResult | (() => McpCollectionResult | Promise); runParallel?: `#${string}`; runSchedule?: { continueOnError?: boolean; cancelMs?: number; + delayStartMs?: number; intervalMs?: number; repeat?: number }; + retainLastViable?: RetainLastViableOption; // priority?: number; isRequired?: boolean; // group?: string; @@ -77,6 +87,35 @@ type McpCollection = [ } ]; +/** + * Context provided to a {@link RetainLastViableCollection} evaluation. + * + * @property name - Collection name being evaluated. + * @property {McpCollectionResult|undefined} [previous] - Previous "viable" collection stored in the registry. + * @property {McpCollectionResult|undefined} [current] - Updated collection response created by the latest run. + * @property [error] - Error, exception, thrown during collection updates. + * @property isSuccess - Did the collection callback resolve without throwing an error? + */ +type RetainLastViableContext = { + name: string; + previous?: McpCollectionResult | undefined; + current?: McpCollectionResult | undefined; + error?: unknown | undefined; + isSuccess: boolean; +}; + +/** + * Custom function to determine whether the previous collection should be kept. Returning `true` + * keeps the previous collection; `false` lets it get updated. Useful when a remote collection + * fails. + */ +type RetainLastViableCollection = (context: RetainLastViableContext) => boolean | Promise; + +/** + * Config option for `retainLastViable`. Supports `boolean` shorthand or a custom function. + */ +type RetainLastViableOption = boolean | RetainLastViableCollection; + /** * A function that creates a collection registered with the MCP server. */ @@ -281,6 +320,43 @@ const onUpdateServerRecordsRegistry = ( }; }; +/** + * Default "last viable" check, see {@link RetainLastViableCollection}. + * Retains the previous response if: + * 1. If the previous data existed and records had length (`previous.records.length > 0`), AND + * 2. The new update threw an error (`!isSuccess`) OR returned zero records (`current.records.length === 0`). + * + * @param context - Retention context. + */ +const defaultRetainCollection: RetainLastViableCollection = context => { + const { previous, current, isSuccess } = context || {} as RetainLastViableContext; + const prevCount = Array.isArray(previous?.records) ? previous.records.length : 0; + const newCount = Array.isArray(current?.records) ? current.records.length : 0; + + return prevCount > 0 && (!isSuccess || newCount === 0); +}; + +/** + * Is this a collection record? + * + * @param value - Value to check. + */ +const isMcpCollectionRecord = (value: unknown): value is McpCollectionRecord => + isPlainObject(value) && + typeof (value as McpCollectionRecord).id === 'string' && (value as McpCollectionRecord).id.length > 0 && + typeof (value as McpCollectionRecord).sourceId === 'string' && (value as McpCollectionRecord).sourceId.length > 0 && + typeof (value as McpCollectionRecord).sourceType === 'string' && (value as McpCollectionRecord).sourceType.length > 0; + +/** + * Is this a collection result? + * + * @param value - Value to check. + */ +const isMcpCollectionResult = (value: unknown): value is McpCollectionResult => + isPlainObject(value) && + Array.isArray((value as McpCollectionResult).records) && + (value as McpCollectionResult).records.every(isMcpCollectionRecord); + /** * Registers a set of collections asynchronously. * @@ -309,22 +385,73 @@ const registerCollections = async ( ): Promise => { log.debug(`Reviewing registration for ${collections.length} collections.`); + // Step 1: Immediate hydration for collections with `_config.initial` + for (const [name, , config] of collections) { + if (config?.initial) { + try { + const initialResult = typeof config.initial === 'function' + ? await config.initial() + : config.initial; + + if (isMcpCollectionResult(initialResult)) { + await setServerRecordsRegistry({ name, response: initialResult, error: undefined }); + } else { + throw new Error(`Invalid collection response "${name}"`); + } + } catch (err) { + log.warn(`Failed to hydrate initial data for collection "${name}": ${formatUnknownError(err)}`); + } + } + } + + // Step 2: Main collection execution (handles scheduled/worker/background callbacks) // Wrapper for each loader; handle incremental updates - const registrationPromises = collections.map(async ([name, callback]) => { + const registrationPromises = collections.map(async ([name, callback, config]) => { let error: unknown | undefined; let response: McpCollectionResult | undefined; let isSuccess = false; try { response = await callback(); - isSuccess = true; + + if (isMcpCollectionResult(response)) { + isSuccess = true; + } else { + throw new Error(`Invalid collection response "${name}"`); + } } catch (err) { error = err; log.error(`Error loading collection ${name}: ${formatUnknownError(err)}`); } + const previous = getServerRecordsRegistry({ collectionName: name }) as McpCollectionResult | undefined; + let shouldRetain = false; + + if (config?.retainLastViable) { + try { + const context: RetainLastViableContext = { + name, + previous, + current: response, + error, + isSuccess + }; + + shouldRetain = await Promise.resolve( + typeof config.retainLastViable === 'function' + ? (config.retainLastViable as RetainLastViableCollection)(context) + : defaultRetainCollection(context) + ); + } catch (err) { + log.warn(`Error evaluating "retainLastViable" collection "${name}": ${formatUnknownError(err)}`); + } + } + try { - if (response) { + if (shouldRetain) { + log.warn(`Collection "${name}" update triggered retention policy; keeping previous viable response (${previous?.records?.length || 0} records).`); + response = previous; + } else if (response) { await setServerRecordsRegistry({ name, response, error }); } } catch (err) { @@ -400,11 +527,17 @@ const registerCollections = async ( }; export { + defaultRetainCollection, getServerRecordsRegistry, + isMcpCollectionRecord, + isMcpCollectionResult, onUpdateServerRecordsRegistry, registerCollections, setServerRecordsRegistry, type OnUpdateServerRecordsRegistryOptions, + type RetainLastViableContext, + type RetainLastViableOption, + type RetainLastViableCollection, type McpCollection, type McpCollectionCreator, type McpCollectionRecord, diff --git a/src/server.collections.ts b/src/server.collections.ts index cddf0370..c6ed2d08 100644 --- a/src/server.collections.ts +++ b/src/server.collections.ts @@ -35,7 +35,7 @@ options: GlobalOptions = getOptions()): McpCollectionCreator => () => { }); }; - return config ? [name, handler, config] : [name, handler]; + return config ? [name, handler, { ...config }] : [name, handler]; }; /** @@ -55,7 +55,10 @@ options: GlobalOptions = getOptions()): McpCollectionCreator => () => { const [name, callback, config] = creator(options); const deferOptions = { ...(typeof runSchedule?.cancelMs === 'number' ? { cancelMs: runSchedule.cancelMs } : {}), - ...(typeof runSchedule?.intervalMs === 'number' ? { intervalMs: runSchedule.intervalMs } : {}) + ...(typeof runSchedule?.intervalMs === 'number' ? { intervalMs: runSchedule.intervalMs } : {}), + ...(typeof runSchedule?.delayStartMs === 'number' ? { delayStartMs: runSchedule.delayStartMs } : {}), + ...(typeof runSchedule?.continueOnError === 'boolean' ? { continueOnError: runSchedule.continueOnError } : {}), + ...(typeof runSchedule?.repeat === 'number' ? { repeat: runSchedule.repeat } : {}) }; const handler = async (args?: unknown): Promise => { @@ -71,7 +74,7 @@ options: GlobalOptions = getOptions()): McpCollectionCreator => () => { return response || { records: [] }; }; - return config ? [name, handler, config] : [name, handler]; + return config ? [name, handler, { ...config }] : [name, handler]; }; /** @@ -118,7 +121,10 @@ const composeCollections = async ( updatedCreator = makeParallelProxyCreator({ creator, moduleSpecifier: runHostValue, exportName: 'collectionCallback' }); } - if (typeof runScheduleConfig?.cancelMs === 'number' || typeof runScheduleConfig?.intervalMs === 'number') { + if (typeof runScheduleConfig?.cancelMs === 'number' || + typeof runScheduleConfig?.intervalMs === 'number' || + typeof runScheduleConfig?.delayStartMs === 'number' || + typeof runScheduleConfig?.repeat === 'number') { // Layer scheduling so the defer-task guardrails apply to the entire execution, including any worker-pool proxy. updatedCreator = makeScheduledProxyCreator({ creator: updatedCreator, runSchedule: runScheduleConfig }); } diff --git a/src/server.task.ts b/src/server.task.ts index 68feb2b7..745f7197 100644 --- a/src/server.task.ts +++ b/src/server.task.ts @@ -42,6 +42,7 @@ interface DeferTaskHandle { * Defaults to `false`. `stop()` and `cancelMs` still terminate the loop. * @property {DeferTaskDebugHandler} [debug] - Debug callback for lifecycle events. * See {@link deferTask}. + * @property [delayStartMs] - Delay before the **FIRST** run in ms. (default `0`) * @property [intervalMs] - Max time for both per-execution timeout AND * the randomized base delay between repetitions. The per-execution timeout is * derived as `intervalMs * 1.5` so a run is never killed by the same value used @@ -54,6 +55,7 @@ interface DeferTaskOptions { cancelMs?: number; continueOnError?: boolean; debug?: DeferTaskDebugHandler; + delayStartMs?: number; intervalMs?: number; repeat?: number | undefined; errorMessage?: string; @@ -165,6 +167,7 @@ const deferTask = ( cancelMs, continueOnError = false, debug = () => {}, + delayStartMs, repeat, intervalMs, errorMessage = 'Task timed out' @@ -172,7 +175,8 @@ const deferTask = ( const validRepeat = typeof repeat === 'number' && repeat > 0 ? repeat : 1; const updatedRepeat = Number.isFinite(validRepeat) ? validRepeat : undefined; - const updatedIntervalMs = intervalMs ?? 1000; + const updatedIntervalMs = typeof intervalMs === 'number' && Number.isFinite(intervalMs) ? intervalMs : 1000; + const updatedDelayStartMs = typeof delayStartMs === 'number' && Number.isFinite(delayStartMs) ? delayStartMs : 0; const runTimeoutMs = updatedIntervalMs * 1.5; let updatedCancelMs = cancelMs; @@ -202,6 +206,30 @@ const deferTask = ( return undefined; } + // Initial startup delay + if (state.count === 0 && updatedDelayStartMs > 0) { + debug({ + type: 'run:delay', + value: () => ({ ...state }) + }); + + const randomizedStartMs = updatedDelayStartMs * (0.9 + Math.random() * 0.2); + + try { + await delay({ ms: randomizedStartMs, signal: state.controller?.signal }); + } catch (error) { + if (error instanceof DOMException && error.name === 'AbortError') { + return undefined; + } + + return Promise.reject(error); + } + + if (!state.isRunning) { + return undefined; + } + } + const startFunc = timeoutFunction(() => { if (state.isRunning) { state.count += 1; @@ -249,6 +277,7 @@ const deferTask = ( return Promise.reject(error); }); + // Subsequent interval delay if (state.isRunning && (updatedRepeat === undefined || state.count < updatedRepeat)) { const randomizedMs = updatedIntervalMs * (0.9 + Math.random() * 0.2); From 972fb593d1fe559f3f924e6447884bd20e505937 Mon Sep 17 00:00:00 2001 From: CD Cabrera Date: Tue, 8 Sep 2026 17:16:58 -0400 Subject: [PATCH 2/2] fix: review update --- src/__tests__/collections.test.ts | 30 +++++++++++++++++++----------- src/collections.ts | 5 +++-- 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/src/__tests__/collections.test.ts b/src/__tests__/collections.test.ts index 726cd4b3..f1b597a7 100644 --- a/src/__tests__/collections.test.ts +++ b/src/__tests__/collections.test.ts @@ -225,13 +225,13 @@ describe('registerCollections', () => { const handler = jest.fn().mockImplementation(() => asyncPromise); const collections: any[] = [ - ['dual-phase-coll', handler, { initial: { records: initialRecords } }] + ['dual-phase-collection', handler, { initial: { records: initialRecords } }] ]; const registrationPromise = registerCollections(collections); // Immediate check: serverRecordsRegistry has initial records before handler finishes - expect(getServerRecordsRegistry({ collectionName: 'dual-phase-coll' })).toEqual({ records: initialRecords }); + expect(getServerRecordsRegistry({ collectionName: 'dual-phase-collection' })).toEqual({ records: initialRecords }); resolveHandler!({ records: [{ id: 'live-1', sourceId: 'live', sourceType: 'api' }] }); await registrationPromise; @@ -242,13 +242,13 @@ describe('registerCollections', () => { const handler = jest.fn().mockResolvedValue({ records: [] }); const collections: any[] = [ - ['retained-coll', handler, { initial: { records: initialRecords }, retainLastViable: true }] + ['retained-collection', handler, { initial: { records: initialRecords }, retainLastViable: true }] ]; await registerCollections(collections); // Retains initialRecords because update returned empty records - expect(getServerRecordsRegistry({ collectionName: 'retained-coll' })).toEqual({ records: initialRecords }); + expect(getServerRecordsRegistry({ collectionName: 'retained-collection' })).toEqual({ records: initialRecords }); }); it('should retain previous viable records when retainLastViable is true and update throws an error', async () => { @@ -256,22 +256,22 @@ describe('registerCollections', () => { const handler = jest.fn().mockRejectedValue(new Error('Network failure')); const collections: any[] = [ - ['error-retained-coll', handler, { initial: { records: initialRecords }, retainLastViable: true }] + ['error-retained-collection', handler, { initial: { records: initialRecords }, retainLastViable: true }] ]; await registerCollections(collections); // Retains initialRecords because update threw an error - expect(getServerRecordsRegistry({ collectionName: 'error-retained-coll' })).toEqual({ records: initialRecords }); + expect(getServerRecordsRegistry({ collectionName: 'error-retained-collection' })).toEqual({ records: initialRecords }); }); - it('should support custom predicate function for retainLastViable', async () => { + it('should support a custom function for retainLastViable', async () => { const initialRecords = [ { id: 'init-1', sourceId: 'mock', sourceType: 'mock' }, { id: 'init-2', sourceId: 'mock', sourceType: 'mock' }, { id: 'init-3', sourceId: 'mock', sourceType: 'mock' } ]; - // Crawl returned only 1 record (loss of > 50% data) + // Crawl returned only 1 record const handler = jest.fn().mockResolvedValue({ records: [{ id: 'init-1', sourceId: 'mock', sourceType: 'mock' }] }); const customPredicate = jest.fn().mockImplementation(({ previous, current }) => { const prevCount = previous?.records?.length || 0; @@ -281,7 +281,7 @@ describe('registerCollections', () => { }); const collections: any[] = [ - ['custom-predicate-coll', handler, { + ['custom-func-collection', handler, { initial: { records: initialRecords }, retainLastViable: customPredicate }] @@ -290,12 +290,20 @@ describe('registerCollections', () => { await registerCollections(collections); expect(customPredicate).toHaveBeenCalledWith(expect.objectContaining({ - name: 'custom-predicate-coll', + name: 'custom-func-collection', previous: { records: initialRecords }, current: { records: [{ id: 'init-1', sourceId: 'mock', sourceType: 'mock' }] }, isSuccess: true })); - expect(getServerRecordsRegistry({ collectionName: 'custom-predicate-coll' })).toEqual({ records: initialRecords }); + expect(getServerRecordsRegistry({ collectionName: 'custom-func-collection' })).toEqual({ records: initialRecords }); + }); + + it('should not write invalid collections to the registry', async () => { + const handler = jest.fn().mockResolvedValue({ records: [{ id: 'invalid-record' }] }); + + await registerCollections([['invalid-collection', handler]]); + + expect(getServerRecordsRegistry({ collectionName: 'invalid-collection' })).toBeUndefined(); }); it('should call onRequired when all required collections are settled', async () => { diff --git a/src/collections.ts b/src/collections.ts index ee2dd6c0..4c5346d3 100644 --- a/src/collections.ts +++ b/src/collections.ts @@ -412,10 +412,11 @@ const registerCollections = async ( let isSuccess = false; try { - response = await callback(); + const initialResponse = await callback(); - if (isMcpCollectionResult(response)) { + if (isMcpCollectionResult(initialResponse)) { isSuccess = true; + response = initialResponse; } else { throw new Error(`Invalid collection response "${name}"`); }