fix(kb): harden ingestion reconciliation

Author: rschlaefli

docs/async-and-workers.md

@@ -49,7 +49,7 @@ Bare `http.createServer`, two routes: `GET /healthz` and `POST /AddResponse`. No
 `packages/hatchet/src/index.ts:prepareHatchetTasks` registers two local workflows:
 
 - `ingest-kb-resource` accepts the selected resource, version, and attempt identifiers, prepares the exact source bytes, then calls `packages/hatchet/src/kbIngestion.ts:dispatchKBIngestion`. Dispatch awaits `POST /v1/resources`, stores the returned operation identifier, and reuses the same version, digest, source URL, and idempotency key when an attempt is retried.
-- `monitor-kb-ingestions` runs every minute and calls `packages/hatchet/src/kbIngestion.ts:monitorActiveKBIngestions`. It polls `GET /v1/operations/{operation_id}` with bounded concurrency and applies only responses matching the local operation, resource version, and content digest.
+- `monitor-kb-ingestions` runs every minute and calls `packages/hatchet/src/kbIngestion.ts:monitorActiveKBIngestions`. It rotates through at most 32 active operations per tick, polls `GET /v1/operations/{operation_id}` eight at a time, and applies only responses matching the local operation, resource version, and content digest. A succeeded operation becomes `READY` only when its observed digest and actively serving version/digest match the local resource.
 
 Operation events also return through the raw-body `/api/webhooks/kb-ingestion` route registered by `apps/backend-docker/src/app.ts:prepareApp`. `packages/graphql/src/services/knowledgeWebhooks.ts:handleKBIngestionWebhook` accepts the strict canonical event body and the four `X-Ingestion-*` headers, verifies an HMAC-SHA256 signature within the five-minute replay window against the current or previous webhook secret, then applies the same operation/version/digest correlation guards as polling. A resource becomes `READY` only when the platform reports that the expected version and digest are actively serving.
 

packages/hatchet/src/kbIngestion.ts

@@ -10,6 +10,7 @@ import {
 } from './kbIngestionApi.js'
 
 const KB_INGESTION_POLL_CONCURRENCY = 8
+const KB_INGESTION_POLL_LIMIT = 32
 
 export type KBIngestionLogger = {
   info?: (
@@ -40,6 +41,7 @@ export type MonitorKBIngestionsDependencies = {
   prisma: KBIngestionPrisma
   client?: KBIngestionApiClient
   env?: NodeJS.ProcessEnv
+  now?: () => Date
   logger?: KBIngestionLogger
 }
 
@@ -373,6 +375,20 @@ async function reconcileResource({
       return
     }
 
+    if (
+      operation.status === 'succeeded' &&
+      (operation.observedSha256 !== contentSha256 ||
+        operation.serving.activeResourceVersion !== resource.resourceVersion ||
+        operation.serving.activeSha256 !== contentSha256)
+    ) {
+      await logErrorBestEffort(
+        logger,
+        'KB ingestion serving correlation failed',
+        identifiers
+      )
+      return
+    }
+
     const transition = mapOperationStatus(operation.status)
     const sourceStatuses =
       operation.status === 'accepted'
@@ -408,15 +424,29 @@ export async function monitorActiveKBIngestions(
   dependencies: MonitorKBIngestionsDependencies
 ): Promise<void> {
   const env = dependencies.env ?? process.env
-  const resources = await dependencies.prisma.kBResource.findMany({
-    where: {
-      status: {
-        in: [KBResourceStatus.QUEUED, KBResourceStatus.PROCESSING],
-      },
-      ingestionAttemptId: { not: null },
-      contentSha256: { not: null },
-      externalOperationId: { not: null },
+  const activeWhere = {
+    status: {
+      in: [KBResourceStatus.QUEUED, KBResourceStatus.PROCESSING],
     },
+    ingestionAttemptId: { not: null },
+    contentSha256: { not: null },
+    externalOperationId: { not: null },
+  }
+  const activeCount = await dependencies.prisma.kBResource.count({
+    where: activeWhere,
+  })
+  if (activeCount === 0) {
+    return
+  }
+
+  const pollWindow = Math.floor(
+    (dependencies.now?.() ?? new Date()).getTime() / 60_000
+  )
+  const skip =
+    ((pollWindow % activeCount) * KB_INGESTION_POLL_LIMIT) % activeCount
+  const query = {
+    where: activeWhere,
+    orderBy: { id: 'asc' as const },
     select: {
       id: true,
       kbId: true,
@@ -425,9 +455,19 @@ export async function monitorActiveKBIngestions(
       contentSha256: true,
       externalOperationId: true,
     },
+  }
+  const resources = await dependencies.prisma.kBResource.findMany({
+    ...query,
+    skip,
+    take: KB_INGESTION_POLL_LIMIT,
   })
-  if (resources.length === 0) {
-    return
+  if (resources.length < KB_INGESTION_POLL_LIMIT && skip > 0) {
+    resources.push(
+      ...(await dependencies.prisma.kBResource.findMany({
+        ...query,
+        take: Math.min(KB_INGESTION_POLL_LIMIT - resources.length, skip),
+      }))
+    )
   }
   const client = dependencies.client ?? createKBIngestionApiClient({ env })
 

packages/hatchet/src/kbIngestionApi.ts

@@ -197,10 +197,12 @@ function parseOperationStatus(value: unknown): KBOperationStatusResponse {
       'updated_at',
     ]) ||
     !isBoundedString(operation.operation_id, 255) ||
+    typeof operation.status !== 'string' ||
     !['accepted', 'running', 'succeeded', 'failed', 'superseded'].includes(
-      String(operation.status)
+      operation.status
     ) ||
-    !['create', 'update', 'delete'].includes(String(operation.operation)) ||
+    typeof operation.operation !== 'string' ||
+    !['create', 'update', 'delete'].includes(operation.operation) ||
     !isBoundedString(operation.project_id, 255) ||
     !isBoundedString(operation.producer, 255) ||
     !isBoundedString(operation.external_resource_id, 512) ||

packages/hatchet/test/kbIngestion.test.ts

@@ -261,7 +261,12 @@ describe('KB ingestion reconciliation', () => {
   function monitorPrisma(resources: Record<string, unknown>[]) {
     return {
       kBResource: {
-        findMany: vi.fn().mockResolvedValue(resources),
+        count: vi.fn().mockResolvedValue(resources.length),
+        findMany: vi
+          .fn()
+          .mockImplementation(async ({ skip = 0, take = resources.length }) =>
+            resources.slice(skip, skip + take)
+          ),
         updateMany: vi.fn().mockResolvedValue({ count: 1 }),
       },
     }
@@ -327,9 +332,16 @@ describe('KB ingestion reconciliation', () => {
     await monitorActiveKBIngestions({
       prisma: prisma as never,
       client: client({
-        getOperation: vi
-          .fn()
-          .mockResolvedValue(operation({ status: 'succeeded' })),
+        getOperation: vi.fn().mockResolvedValue(
+          operation({
+            status: 'succeeded',
+            observedSha256: CONTENT_SHA256,
+            serving: {
+              activeResourceVersion: 3,
+              activeSha256: CONTENT_SHA256,
+            },
+          })
+        ),
       }),
     })
 
@@ -341,6 +353,59 @@ describe('KB ingestion reconciliation', () => {
     expect(update.data.ingestedAt).toBeInstanceOf(Date)
   })
 
+  it.each([
+    {
+      observedSha256: null,
+      serving: {
+        activeResourceVersion: 3,
+        activeSha256: CONTENT_SHA256,
+      },
+    },
+    {
+      observedSha256: CONTENT_SHA256,
+      serving: {
+        activeResourceVersion: 2,
+        activeSha256: CONTENT_SHA256,
+      },
+    },
+    {
+      observedSha256: CONTENT_SHA256,
+      serving: {
+        activeResourceVersion: 3,
+        activeSha256: 'a'.repeat(64),
+      },
+    },
+  ])(
+    'refuses success until the exact version and digest are actively serving',
+    async (servingIdentity) => {
+      const prisma = monitorPrisma([activeResource])
+      const logger = { error: vi.fn() }
+
+      await monitorActiveKBIngestions({
+        prisma: prisma as never,
+        client: client({
+          getOperation: vi.fn().mockResolvedValue(
+            operation({
+              status: 'succeeded',
+              ...servingIdentity,
+            })
+          ),
+        }),
+        logger,
+      })
+
+      expect(prisma.kBResource.updateMany).not.toHaveBeenCalled()
+      expect(logger.error).toHaveBeenCalledWith(
+        'KB ingestion serving correlation failed',
+        {
+          resourceId: RESOURCE_ID,
+          kbId: KB_ID,
+          ingestionAttemptId: ATTEMPT_ID,
+        }
+      )
+    }
+  )
+
   it('refuses a status response that does not match every correlation field', async () => {
     const prisma = monitorPrisma([activeResource])
     const logger = { error: vi.fn() }
@@ -397,4 +462,44 @@ describe('KB ingestion reconciliation', () => {
     expect(getOperation).toHaveBeenCalledTimes(17)
     expect(peak).toBe(8)
   })
+
+  it('rotates a bounded 32-resource reconciliation window each minute', async () => {
+    const resources = Array.from({ length: 49 }, (_, index) => ({
+      ...activeResource,
+      id: `00000000-0000-4000-8000-${String(index).padStart(12, '0')}`,
+      externalOperationId: `${OPERATION_ID}_${index}`,
+    }))
+    const prisma = {
+      kBResource: {
+        count: vi.fn().mockResolvedValue(resources.length),
+        findMany: vi.fn().mockImplementation(async ({ skip = 0, take }) => {
+          return resources.slice(skip, skip + take)
+        }),
+        updateMany: vi.fn().mockResolvedValue({ count: 1 }),
+      },
+    }
+    const getOperation = vi.fn().mockImplementation(async (operationId) => {
+      const index = Number(operationId.split('_').at(-1))
+      return operation({
+        operationId,
+        externalResourceId: resources[index]!.id,
+      })
+    })
+
+    await monitorActiveKBIngestions({
+      prisma: prisma as never,
+      client: client({ getOperation }),
+      now: () => new Date(60_000),
+    })
+
+    expect(prisma.kBResource.findMany).toHaveBeenNthCalledWith(
+      1,
+      expect.objectContaining({ skip: 32, take: 32 })
+    )
+    expect(prisma.kBResource.findMany).toHaveBeenNthCalledWith(
+      2,
+      expect.objectContaining({ take: 15 })
+    )
+    expect(getOperation).toHaveBeenCalledTimes(32)
+  })
 })

packages/hatchet/test/kbIngestionApi.test.ts

@@ -194,6 +194,26 @@ describe('canonical ingestion API client', () => {
       )
     ).rejects.toThrow('Ingestion API returned an invalid response')
   })
+
+  it.each([{ status: ['succeeded'] }, { operation: ['update'] }])(
+    'rejects non-string operation enums',
+    async (override) => {
+      const fetchRequest = vi.fn().mockResolvedValue({
+        ok: true,
+        status: 200,
+        json: vi.fn().mockResolvedValue({
+          ...operationResponse,
+          ...override,
+        }),
+      })
+
+      await expect(
+        createKBIngestionApiClient({ env, fetchRequest }).getOperation(
+          operationResponse.operation_id
+        )
+      ).rejects.toThrow('Ingestion API returned an invalid response')
+    }
+  )
 })
 
 describe('ingestion source preparation', () => {