fix(knowledge): check the page deadline after the lease heartbeat

Author: waleedlatif1

apps/sim/lib/knowledge/connectors/member-observations.test.ts

@@ -529,6 +529,27 @@ describe('writeProjectionPages', () => {
     }
   })
 
+  /** The lease heartbeat runs before each page and can itself use up what is left of the budget. */
+  it('opens no transaction when the heartbeat before a page outlasts the deadline', async () => {
+    let clock = Date.now()
+    const now = vi.spyOn(Date, 'now').mockImplementation(() => clock)
+    try {
+      const transaction = vi.fn()
+      const write = vi.fn()
+      const beforePage = vi.fn(async () => {
+        clock += 2_000
+      })
+
+      await expect(
+        writeProjectionPages(['a'], transaction, write, { beforePage, deadlineAt: clock + 1_000 })
+      ).resolves.toEqual({ written: 0, finished: false })
+      expect(beforePage).toHaveBeenCalledOnce()
+      expect(transaction).not.toHaveBeenCalled()
+    } finally {
+      now.mockRestore()
+    }
+  })
+
   it('writes every page when no deadline is given', async () => {
     queueTableRows(schemaMock.document, [
       { id: 'a', chunkCount: 200 },

apps/sim/lib/knowledge/connectors/member-observations.ts

@@ -425,9 +425,9 @@ export async function writeProjectionPages(
   let pending = [...documentIds]
   let written = 0
   while (pending.length > 0) {
+    await options.beforePage?.()
     if (options.deadlineAt !== undefined && Date.now() >= options.deadlineAt)
       return { written, finished: false }
-    await options.beforePage?.()
     const { rows, rest } = await transaction(async (tx) => {
       const { page, rest } = await lockProjectionPage(tx, pending)
       return { rows: page.length > 0 ? await write(tx, page) : 0, rest }
@@ -471,8 +471,8 @@ export async function rewriteConnectorDocumentAcls(input: {
   let rewritten = 0
   let after: ConnectorDocumentCursor | undefined
   for (;;) {
-    if (expired()) return { rewritten, finished: false }
     await input.beforePage?.()
+    if (expired()) return { rewritten, finished: false }
     const window = await db
       .select({
         id: document.id,