fix(pi): preserve babysit partial state

Bill Leoutsakos · Bill Leoutsakos · commit dbf510f36d01 · 2026-07-25T14:24:46.000-07:00
diff --git a/apps/sim/executor/handlers/pi/babysit-backend.test.ts b/apps/sim/executor/handlers/pi/babysit-backend.test.ts
@@ -236,6 +236,7 @@ describe('runBabysitPiWithOptions', () => {
     mockReplyAndResolve.mockResolvedValue({
       repliesPosted: 1,
       threadsResolved: 1,
+      resolvedThreadIds: ['thread-1'],
       replyFailures: [],
       resolveFailures: [],
       headMoved: false,
@@ -646,6 +647,43 @@ describe('runBabysitPiWithOptions', () => {
     expect(mockReplyAndResolve.mock.calls[0][4]).toBe(OLD_SHA)
   })
 
+  it('preserves known clean flags when the pin moves after successful writes', async () => {
+    const noChecksGreen = {
+      ...greenChecks,
+      checks: [],
+      contextRequirements: new Map<string, boolean>(),
+    }
+    mockFetchSnapshot
+      .mockResolvedValueOnce(snapshot)
+      .mockResolvedValueOnce(snapshot)
+      .mockResolvedValueOnce({ ...snapshot, headSha: NEW_SHA })
+      .mockResolvedValueOnce({ ...snapshot, headSha: SECOND_SHA })
+    mockFetchThreads.mockResolvedValue({
+      actionable: [trustedThread],
+      skipped: [],
+      totalUnresolved: 1,
+      latestReview: null,
+    })
+    mockFetchChecks.mockResolvedValue(noChecksGreen)
+    const { runner } = makeRunner({})
+    mockWithPiSandbox.mockImplementation(async (callback) => callback(runner))
+
+    const result = await runBabysitPiWithOptions(
+      params({ reviewMentions: ['@review-bot'] }),
+      { onEvent: vi.fn() },
+      { convergenceWaitMs: 0, roundWaitMs: 0 }
+    )
+
+    expect(result).toMatchObject({
+      stopReason: 'head_moved',
+      rounds: 1,
+      commitsPushed: 1,
+      threadsResolved: 1,
+      threadsClean: true,
+      checksGreen: true,
+    })
+  })
+
   it('reports a confirmed push when the convergence read fails transiently', async () => {
     mockFetchSnapshot
       .mockResolvedValueOnce(snapshot)
diff --git a/apps/sim/executor/handlers/pi/babysit-backend.ts b/apps/sim/executor/handlers/pi/babysit-backend.ts
@@ -598,6 +598,7 @@ export async function runBabysitPiWithOptions(
 
   let latestThreads: BabysitThreadsState | undefined
   let latestChecks: BabysitCheckState | undefined
+  let lastKnownChecksGreen = false
   let githubWriteOccurred = false
   try {
     if (signal.aborted) throw new Error('Pi run aborted')
@@ -617,6 +618,7 @@ export async function runBabysitPiWithOptions(
     }
     latestThreads = await fetchBabysitThreads(params, signal)
     latestChecks = await fetchBabysitCheckState(params, pinnedHeadSha, undefined, signal)
+    lastKnownChecksGreen = latestChecks.checksGreen
     const initialRequirements = latestChecks.contextRequirements
     if (latestChecks.startupFailure) {
       const threadsClean =
@@ -726,6 +728,7 @@ export async function runBabysitPiWithOptions(
             initialRequirements,
             signal
           )
+          lastKnownChecksGreen = latestChecks.checksGreen
           if (reviewRequest && !reviewRequest.landed) {
             reviewRequest.landed = await babysitReviewLandedSince(
               params,
@@ -853,6 +856,7 @@ export async function runBabysitPiWithOptions(
           githubWriteOccurred = true
           progress.changedFiles = finalized.changedFiles
           progress.diff = finalized.diff
+          lastKnownChecksGreen = ![...initialRequirements.values()].some((required) => required)
           let convergence: Awaited<ReturnType<typeof waitForHeadConvergence>>
           try {
             convergence = await waitForHeadConvergence(
@@ -938,6 +942,16 @@ export async function runBabysitPiWithOptions(
           laggingHeadSha
         )
         progress.threadsResolved += writeResult.threadsResolved
+        const resolvedThreadIds = new Set(writeResult.resolvedThreadIds ?? [])
+        if (resolvedThreadIds.size > 0) {
+          latestThreads = {
+            ...latestThreads,
+            actionable: latestThreads.actionable.filter(
+              (thread) => !resolvedThreadIds.has(thread.id)
+            ),
+            totalUnresolved: Math.max(0, latestThreads.totalUnresolved - resolvedThreadIds.size),
+          }
+        }
         githubWriteOccurred ||= writeResult.repliesPosted > 0 || writeResult.threadsResolved > 0
         if (writeResult.replyFailures.length) {
           progress.notes.push(`${writeResult.replyFailures.length} thread replies failed.`)
@@ -951,7 +965,15 @@ export async function runBabysitPiWithOptions(
           )
         }
         if (writeResult.stopReason) {
-          return resultFor(totals, writeResult.stopReason, progress, threadsClean, false)
+          const knownThreadsClean =
+            latestThreads.actionable.length === 0 && latestThreads.skipped.length === 0
+          return resultFor(
+            totals,
+            writeResult.stopReason,
+            progress,
+            knownThreadsClean,
+            lastKnownChecksGreen
+          )
         }
         if (writeResult.phaseError) {
           progress.notes.push(
@@ -964,15 +986,27 @@ export async function runBabysitPiWithOptions(
             totals,
             finalized.commitPushed ? 'pushed_awaiting_confirmation' : 'agent_failure',
             progress,
-            threadsClean,
-            false
+            latestThreads.actionable.length === 0 && latestThreads.skipped.length === 0,
+            lastKnownChecksGreen
           )
         }
         if (writeResult.headMoved) {
-          return resultFor(totals, 'head_moved', progress, threadsClean, false)
+          return resultFor(
+            totals,
+            'head_moved',
+            progress,
+            latestThreads.actionable.length === 0 && latestThreads.skipped.length === 0,
+            lastKnownChecksGreen
+          )
         }
         if (writeResult.awaitingConfirmation) {
-          return resultFor(totals, 'pushed_awaiting_confirmation', progress, threadsClean, false)
+          return resultFor(
+            totals,
+            'pushed_awaiting_confirmation',
+            progress,
+            latestThreads.actionable.length === 0 && latestThreads.skipped.length === 0,
+            lastKnownChecksGreen
+          )
         }
 
         if (finalized.commitPushed && params.reviewMentions.length > 0) {
@@ -1016,6 +1050,7 @@ export async function runBabysitPiWithOptions(
           initialRequirements,
           signal
         )
+        lastKnownChecksGreen = latestChecks.checksGreen
         if (reviewRequest) {
           reviewRequest.landed = await babysitReviewLandedSince(
             params,
@@ -1089,7 +1124,7 @@ export async function runBabysitPiWithOptions(
         githubError.reason as BabysitStopReason,
         progress,
         threadsClean,
-        latestChecks?.checksGreen ?? false
+        lastKnownChecksGreen
       )
     }
     if (githubWriteOccurred) {
@@ -1103,8 +1138,10 @@ export async function runBabysitPiWithOptions(
         totals,
         progress.commitsPushed > 0 ? 'pushed_awaiting_confirmation' : 'agent_failure',
         progress,
-        false,
-        false
+        !!latestThreads &&
+          latestThreads.actionable.length === 0 &&
+          latestThreads.skipped.length === 0,
+        lastKnownChecksGreen
       )
     }
     throw createScrubbedPiError(error, secrets, 'Babysit failed')
diff --git a/apps/sim/executor/handlers/pi/babysit-github.test.ts b/apps/sim/executor/handlers/pi/babysit-github.test.ts
@@ -340,6 +340,47 @@ describe('Babysit GitHub orchestration', () => {
     )
   })
 
+  it('reports awaiting confirmation when a third SHA appears after a lagging push', async () => {
+    const thirdSha = 'd'.repeat(40)
+    mockExecuteTool.mockImplementation(async (toolId: string) => {
+      if (toolId === 'github_reply_review_thread') {
+        return { success: true, output: { id: 'reply' } }
+      }
+      if (toolId === 'github_pr_v2') {
+        return {
+          success: true,
+          output: snapshot({
+            head: { sha: thirdSha, ref: 'feature', repo_full_name: 'octo/demo' },
+          }),
+        }
+      }
+      throw new Error(`Unexpected tool ${toolId}`)
+    })
+
+    const result = await replyAndResolveBabysitThreads(
+      params,
+      { headSha: NEXT_SHA, headRef: 'feature', baseRef: 'main' },
+      [
+        {
+          threadId: 'one',
+          classification: 'fixed',
+          reply: 'Fixed.',
+          resolvable: true,
+        },
+      ],
+      undefined,
+      HEAD_SHA
+    )
+
+    expect(result).toMatchObject({
+      repliesPosted: 1,
+      threadsResolved: 0,
+      headMoved: false,
+      awaitingConfirmation: true,
+    })
+    expect(result.stopReason).toBeUndefined()
+  })
+
   it('retains successful reply counts when between-phase revalidation fails', async () => {
     mockExecuteTool.mockImplementation(async (toolId: string) => {
       if (toolId === 'github_reply_review_thread') {
diff --git a/apps/sim/executor/handlers/pi/babysit-github.ts b/apps/sim/executor/handlers/pi/babysit-github.ts
@@ -103,6 +103,7 @@ export interface BabysitCheckState {
 export interface BabysitReplyResolveResult {
   repliesPosted: number
   threadsResolved: number
+  resolvedThreadIds: string[]
   replyFailures: string[]
   resolveFailures: string[]
   headMoved: boolean
@@ -575,6 +576,7 @@ export async function replyAndResolveBabysitThreads(
       return {
         repliesPosted: replySuccesses.length,
         threadsResolved: 0,
+        resolvedThreadIds: [],
         replyFailures,
         resolveFailures: [],
         headMoved: false,
@@ -584,15 +586,21 @@ export async function replyAndResolveBabysitThreads(
     assertBabysitPinned(pin, current)
   } catch (error) {
     if (signal?.aborted) throw error
+    const laggingHeadMoved =
+      error instanceof BabysitGitHubError && error.reason === 'head_moved' && !!laggingHeadSha
     return {
       repliesPosted: replySuccesses.length,
       threadsResolved: 0,
+      resolvedThreadIds: [],
       replyFailures,
       resolveFailures: [],
-      headMoved: error instanceof BabysitGitHubError && error.reason === 'head_moved',
-      awaitingConfirmation: false,
+      headMoved:
+        error instanceof BabysitGitHubError && error.reason === 'head_moved' && !laggingHeadMoved,
+      awaitingConfirmation: laggingHeadMoved,
       ...(error instanceof BabysitGitHubError
-        ? { stopReason: error.reason }
+        ? laggingHeadMoved
+          ? {}
+          : { stopReason: error.reason }
         : {
             phaseError: scrubPiSecrets(
               getErrorMessage(error, 'Failed to revalidate the pull request after replying'),
@@ -602,7 +610,7 @@ export async function replyAndResolveBabysitThreads(
     }
   }
 
-  let threadsResolved = 0
+  const resolvedThreadIds: string[] = []
   const resolveFailures: string[] = []
   for (const decision of replySuccesses.filter((item) => item.resolvable)) {
     try {
@@ -611,7 +619,7 @@ export async function replyAndResolveBabysitThreads(
         { threadId: decision.threadId, apiKey: params.githubToken },
         { signal }
       )
-      if (result.success) threadsResolved += 1
+      if (result.success) resolvedThreadIds.push(decision.threadId)
       else resolveFailures.push(decision.threadId)
     } catch (error) {
       if (signal?.aborted) throw error
@@ -620,7 +628,8 @@ export async function replyAndResolveBabysitThreads(
   }
   return {
     repliesPosted: replySuccesses.length,
-    threadsResolved,
+    threadsResolved: resolvedThreadIds.length,
+    resolvedThreadIds,
     replyFailures,
     resolveFailures,
     headMoved: false,
