feat(discovery): pursue bounded follow-up research

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<DiscoveryResult[]>
 }
 
+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> | 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<DiscoveryLoopResult> {
+  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<string, DiscoveryTask>()
+  const pending: DiscoveryTask[] = []
+  const duplicateTaskIds = new Set<string>()
+  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<string, DiscoveryTask>,
+  destination: DiscoveryTask[],
+  duplicateTaskIds: Set<string>,
+): 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<string>()
+  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
+}

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'/,
+    )
+  })
+})