Fix Babysit wait-only budgeting

Author: Bill Leoutsakos

apps/sim/executor/handlers/pi/babysit-backend.test.ts

@@ -162,6 +162,7 @@ function commandResult(stdout = '', stderr = '', exitCode = 0) {
 }
 
 function makeRunner(options: {
+  cloneResult?: ReturnType<typeof commandResult>
   prepareStdout?: string | string[]
   pushResult?: ReturnType<typeof commandResult>
   roundFile?: string
@@ -186,7 +187,7 @@ function makeRunner(options: {
       ) => {
         runCalls.push({ command, envs: runOptions.envs, timeoutMs: runOptions.timeoutMs })
         if (command.includes('git clone')) {
-          return commandResult('__GIT_CONFIG_DIGEST__=digest-1\n')
+          return options.cloneResult ?? commandResult('__GIT_CONFIG_DIGEST__=digest-1\n')
         }
         if (command.includes('pi -p --mode json')) {
           runOptions.onStdout?.('{"type":"agent_end"}\n')
@@ -312,6 +313,58 @@ describe('runBabysitPiWithOptions', () => {
     expect(mockReviewLanded).toHaveBeenCalledTimes(1)
   })
 
+  it('preserves known clean flags when the Babysit clone fails after the review request', async () => {
+    mockFetchSnapshot.mockResolvedValue(snapshot)
+    mockFetchThreads.mockResolvedValue({
+      actionable: [],
+      skipped: [],
+      totalUnresolved: 0,
+      latestReview: null,
+    })
+    mockFetchChecks.mockResolvedValue(greenChecks)
+    const { runner } = makeRunner({
+      cloneResult: commandResult('', 'clone failed', 1),
+    })
+    mockWithPiSandbox.mockImplementation(async (callback) => callback(runner))
+
+    const result = await runBabysitPiWithOptions(params(), { onEvent: vi.fn() })
+
+    expect(result).toMatchObject({
+      stopReason: 'agent_failure',
+      threadsClean: true,
+      checksGreen: true,
+      rounds: 0,
+      commitsPushed: 0,
+    })
+  })
+
+  it('uses remaining time for wait-only polling without reserving an agent round', async () => {
+    mockFetchSnapshot.mockResolvedValue(snapshot)
+    mockFetchThreads.mockResolvedValue({
+      actionable: [],
+      skipped: [],
+      totalUnresolved: 0,
+      latestReview: null,
+    })
+    mockFetchChecks.mockResolvedValue(greenChecks)
+    const { runner } = makeRunner({})
+    mockWithPiSandbox.mockImplementation(async (callback) => callback(runner))
+
+    const result = await runBabysitPiWithOptions(
+      params({ executionBudgetMs: 2 * 60 * 1000 }),
+      { onEvent: vi.fn() },
+      { roundWaitMs: 30_000 }
+    )
+
+    expect(mockReviewLanded).toHaveBeenCalledTimes(1)
+    expect(result).toMatchObject({
+      stopReason: 'clean',
+      threadsClean: true,
+      checksGreen: true,
+      rounds: 0,
+    })
+  })
+
   it('refuses excess failing checks before fetching discarded diagnostics', async () => {
     const failures = Array.from({ length: 21 }, (_, index) => ({
       ...failingCheck,

apps/sim/executor/handlers/pi/babysit-backend.ts

@@ -686,12 +686,24 @@ export async function runBabysitPiWithOptions(
         progress.notes.push(
           `Clone failed: ${scrubGitSecrets(clone.stderr || clone.stdout, params.githubToken)}`
         )
-        return resultFor(totals, 'agent_failure', progress, false, false)
+        return resultFor(
+          totals,
+          'agent_failure',
+          progress,
+          threadsAreClean(latestThreads),
+          lastKnownChecksGreen
+        )
       }
       const gitConfigDigest = extractMarkerValues(clone.stdout, GIT_CONFIG_DIGEST_MARKER)[0]
       if (!gitConfigDigest) {
         progress.notes.push('Clone did not report the repository Git-config digest.')
-        return resultFor(totals, 'agent_failure', progress, false, false)
+        return resultFor(
+          totals,
+          'agent_failure',
+          progress,
+          threadsAreClean(latestThreads),
+          lastKnownChecksGreen
+        )
       }
       if (params.search) {
         await runner.writeFile(PI_SEARCH_EXTENSION_PATH, PI_SEARCH_EXTENSION_SOURCE)
@@ -725,10 +737,7 @@ export async function runBabysitPiWithOptions(
           latestThreads!.actionable.length > 0 || latestChecks!.blockingFailing.length > 0
         if (!needsAgent) {
           const remaining = lifetime - (Date.now() - startedAt)
-          if (
-            remaining <=
-            options.roundWaitMs + MIN_ROUND_BUDGET_MS + ROUND_FINALIZATION_RESERVE_MS
-          ) {
+          if (remaining <= options.roundWaitMs) {
             const reason = outstandingReason(
               'budget_exhausted',
               latestThreads!,