diff --git a/src/discovery.ts b/src/discovery.ts index f9584cb..ab8ec43 100644 --- a/src/discovery.ts +++ b/src/discovery.ts @@ -1,3 +1,5 @@ +import { isDeepStrictEqual } from 'node:util' + export interface DiscoveryTask { id: string goal: string @@ -29,6 +31,122 @@ export interface KnowledgeDiscoveryDispatcher { ): Promise } +export type DiscoveryLoopStopReason = 'complete' | 'max-rounds' | 'max-tasks' | 'aborted' + +export interface DiscoveryLoopRound { + round: number + tasks: DiscoveryTask[] + results: DiscoveryResult[] + queuedFollowUps: DiscoveryTask[] +} + +export interface DiscoveryLoopResult { + stopReason: DiscoveryLoopStopReason + tasksDispatched: number + results: DiscoveryResult[] + rounds: DiscoveryLoopRound[] + /** Tasks retained, not dropped, when a configured limit stops the loop. */ + pendingTasks: DiscoveryTask[] + /** Structurally identical task identities ignored to prevent cycles. */ + duplicateTaskIds: string[] +} + +export interface RunDiscoveryLoopOptions { + dispatcher: KnowledgeDiscoveryDispatcher + initialTasks: readonly DiscoveryTask[] + /** Maximum follow-up depth including the initial dispatch. Default 3. */ + maxRounds?: number + /** Maximum tasks dispatched across all rounds. Default 24. */ + maxTasks?: number + /** Forwarded to the dispatcher. Default 4. */ + concurrency?: number + signal?: AbortSignal + onRound?: (round: DiscoveryLoopRound) => Promise | void +} + +/** + * Dispatch discovery tasks and recursively pursue worker-proposed follow-ups. + * + * The runner owns only bounded scheduling and accounting. Workers own search, + * source verification, and domain-specific candidate shapes. + */ +export async function runDiscoveryLoop( + options: RunDiscoveryLoopOptions, +): Promise { + const maxRounds = positiveInteger(options.maxRounds ?? 3, 'maxRounds') + const maxTasks = positiveInteger(options.maxTasks ?? 24, 'maxTasks') + const concurrency = positiveInteger(options.concurrency ?? 4, 'concurrency') + const known = new Map() + const pending: DiscoveryTask[] = [] + const duplicateTaskIds = new Set() + enqueueTasks(options.initialTasks, known, pending, duplicateTaskIds) + + const rounds: DiscoveryLoopRound[] = [] + const results: DiscoveryResult[] = [] + let tasksDispatched = 0 + let aborted = false + + while (pending.length > 0 && rounds.length < maxRounds && tasksDispatched < maxTasks) { + if (options.signal?.aborted) { + aborted = true + break + } + const capacity = maxTasks - tasksDispatched + const tasks = pending.splice(0, capacity) + let dispatched: DiscoveryResult[] + try { + dispatched = orderCompleteDispatch( + tasks, + await options.dispatcher.dispatch(tasks, { + concurrency, + signal: options.signal, + }), + ) + } catch (cause) { + if (!options.signal?.aborted) throw cause + pending.unshift(...tasks) + aborted = true + break + } + tasksDispatched += tasks.length + results.push(...dispatched) + + const queuedFollowUps: DiscoveryTask[] = [] + for (const result of dispatched) { + enqueueTasks(result.followUpTasks ?? [], known, queuedFollowUps, duplicateTaskIds) + } + pending.push(...queuedFollowUps) + const round: DiscoveryLoopRound = { + round: rounds.length + 1, + tasks, + results: dispatched, + queuedFollowUps, + } + rounds.push(round) + await options.onRound?.(round) + if (options.signal?.aborted) { + aborted = true + break + } + } + + const stopReason: DiscoveryLoopStopReason = aborted + ? 'aborted' + : pending.length === 0 + ? 'complete' + : tasksDispatched >= maxTasks + ? 'max-tasks' + : 'max-rounds' + return { + stopReason, + tasksDispatched, + results, + rounds, + pendingTasks: pending, + duplicateTaskIds: [...duplicateTaskIds], + } +} + export function createLocalDiscoveryDispatcher( worker: KnowledgeDiscoveryWorker, ): KnowledgeDiscoveryDispatcher { @@ -45,11 +163,64 @@ export function createLocalDiscoveryDispatcher( } } await Promise.all(Array.from({ length: Math.min(concurrency, tasks.length) }, runNext)) - return results.sort( - (a, b) => - tasks.findIndex((task) => task.id === a.taskId) - - tasks.findIndex((task) => task.id === b.taskId), - ) + return orderCompleteDispatch(tasks, results) }, } } + +function enqueueTasks( + tasks: readonly DiscoveryTask[], + known: Map, + destination: DiscoveryTask[], + duplicateTaskIds: Set, +): void { + for (const task of tasks) { + assertTask(task) + const prior = known.get(task.id) + if (!prior) { + known.set(task.id, task) + destination.push(task) + continue + } + if (!isDeepStrictEqual(prior, task)) { + throw new Error(`Discovery task id '${task.id}' refers to conflicting tasks`) + } + duplicateTaskIds.add(task.id) + } +} + +function assertTask(task: DiscoveryTask): void { + if (!task.id.trim()) throw new Error('Discovery task id must not be empty') + if (!task.goal.trim()) throw new Error(`Discovery task '${task.id}' goal must not be empty`) +} + +function orderCompleteDispatch( + tasks: readonly DiscoveryTask[], + results: readonly DiscoveryResult[], +): DiscoveryResult[] { + const order = new Map(tasks.map((task, index) => [task.id, index])) + const received = new Set() + for (const result of results) { + if (!order.has(result.taskId)) { + throw new Error(`Discovery dispatcher returned unknown task '${result.taskId}'`) + } + if (received.has(result.taskId)) { + throw new Error(`Discovery dispatcher returned task '${result.taskId}' more than once`) + } + received.add(result.taskId) + } + const missing = tasks.filter((task) => !received.has(task.id)).map((task) => task.id) + if (missing.length > 0) { + throw new Error(`Discovery dispatcher omitted tasks: ${missing.join(', ')}`) + } + return [...results].sort( + (a, b) => (order.get(a.taskId) as number) - (order.get(b.taskId) as number), + ) +} + +function positiveInteger(value: number, name: string): number { + if (!Number.isInteger(value) || value < 1) { + throw new Error(`Discovery loop ${name} must be a positive integer`) + } + return value +} diff --git a/tests/discovery.test.ts b/tests/discovery.test.ts new file mode 100644 index 0000000..3735876 --- /dev/null +++ b/tests/discovery.test.ts @@ -0,0 +1,148 @@ +import { describe, expect, it } from 'vitest' +import { + createLocalDiscoveryDispatcher, + type DiscoveryTask, + runDiscoveryLoop, +} from '../src/discovery' + +describe('runDiscoveryLoop', () => { + it('pursues follow-up tasks across bounded concurrent rounds', async () => { + const dispatcher = createLocalDiscoveryDispatcher({ + run: async (task) => ({ + taskId: task.id, + summary: `researched ${task.goal}`, + sourceUris: [`https://example.test/${task.id}`], + followUpTasks: + task.id === 'root' + ? [ + { id: 'tools', goal: 'Find tools' }, + { id: 'mcp', goal: 'Find MCP servers' }, + ] + : task.id === 'tools' + ? [{ id: 'licenses', goal: 'Verify licenses' }] + : undefined, + }), + }) + + const result = await runDiscoveryLoop({ + dispatcher, + initialTasks: [{ id: 'root', goal: 'Find open-source options' }], + maxRounds: 3, + maxTasks: 10, + concurrency: 2, + }) + + expect(result.stopReason).toBe('complete') + expect(result.tasksDispatched).toBe(4) + expect(result.rounds.map((round) => round.tasks.map((task) => task.id))).toEqual([ + ['root'], + ['tools', 'mcp'], + ['licenses'], + ]) + expect(result.results.flatMap((entry) => entry.sourceUris ?? [])).toHaveLength(4) + expect(result.pendingTasks).toEqual([]) + }) + + it('returns every undispatched task when a round or task limit stops the search', async () => { + const followUps: DiscoveryTask[] = [ + { id: 'a', goal: 'A' }, + { id: 'b', goal: 'B' }, + ] + const dispatcher = createLocalDiscoveryDispatcher({ + run: (task) => ({ + taskId: task.id, + summary: task.goal, + followUpTasks: task.id === 'root' ? followUps : undefined, + }), + }) + + const roundLimited = await runDiscoveryLoop({ + dispatcher, + initialTasks: [{ id: 'root', goal: 'Root' }], + maxRounds: 1, + }) + expect(roundLimited.stopReason).toBe('max-rounds') + expect(roundLimited.pendingTasks).toEqual(followUps) + + const taskLimited = await runDiscoveryLoop({ + dispatcher, + initialTasks: [{ id: 'root', goal: 'Root' }], + maxTasks: 2, + }) + expect(taskLimited.stopReason).toBe('max-tasks') + expect(taskLimited.results.map((entry) => entry.taskId)).toEqual(['root', 'a']) + expect(taskLimited.pendingTasks).toEqual([{ id: 'b', goal: 'B' }]) + }) + + it('deduplicates exact cycles and rejects conflicting task identities', async () => { + const task = { id: 'same', goal: 'Same' } + const exact = await runDiscoveryLoop({ + dispatcher: createLocalDiscoveryDispatcher({ + run: (current) => ({ + taskId: current.id, + summary: current.goal, + followUpTasks: [task], + }), + }), + initialTasks: [task], + }) + expect(exact.stopReason).toBe('complete') + expect(exact.duplicateTaskIds).toEqual(['same']) + + await expect( + runDiscoveryLoop({ + dispatcher: createLocalDiscoveryDispatcher({ + run: (current) => ({ + taskId: current.id, + summary: current.goal, + followUpTasks: [{ id: current.id, goal: 'Different' }], + }), + }), + initialTasks: [task], + }), + ).rejects.toThrow(/conflicting tasks/) + }) + + it('rejects incomplete dispatcher output instead of losing a task', async () => { + await expect( + runDiscoveryLoop({ + dispatcher: { + dispatch: async () => [{ taskId: 'a', summary: 'only one' }], + }, + initialTasks: [ + { id: 'a', goal: 'A' }, + { id: 'b', goal: 'B' }, + ], + }), + ).rejects.toThrow(/omitted tasks: b/) + }) + + it('returns completed work and pending tasks when cancelled', async () => { + const controller = new AbortController() + const result = await runDiscoveryLoop({ + dispatcher: createLocalDiscoveryDispatcher({ + run: (task) => ({ + taskId: task.id, + summary: task.goal, + followUpTasks: task.id === 'root' ? [{ id: 'next', goal: 'Next' }] : undefined, + }), + }), + initialTasks: [{ id: 'root', goal: 'Root' }], + signal: controller.signal, + onRound: () => controller.abort(), + }) + + expect(result.stopReason).toBe('aborted') + expect(result.results.map((entry) => entry.taskId)).toEqual(['root']) + expect(result.pendingTasks).toEqual([{ id: 'next', goal: 'Next' }]) + }) + + it('validates local dispatcher output when used without the loop', async () => { + const dispatcher = createLocalDiscoveryDispatcher({ + run: () => ({ taskId: 'unknown', summary: 'wrong identity' }), + }) + await expect(dispatcher.dispatch([{ id: 'expected', goal: 'Expected' }])).rejects.toThrow( + /unknown task 'unknown'/, + ) + }) +})