fix(pi): keep babysit sandboxes active

Author: Bill Leoutsakos

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

@@ -13,6 +13,7 @@ const {
   mockRequestReview,
   mockReviewLanded,
   mockResolvePiSandboxLifetime,
+  mockSleepUntilAborted,
 } = vi.hoisted(() => ({
   mockWithPiSandbox: vi.fn(),
   mockFetchSnapshot: vi.fn(),
@@ -23,6 +24,7 @@ const {
   mockRequestReview: vi.fn(),
   mockReviewLanded: vi.fn(),
   mockResolvePiSandboxLifetime: vi.fn(),
+  mockSleepUntilAborted: vi.fn(),
 }))
 
 vi.mock('@/lib/execution/remote-sandbox', () => ({
@@ -32,6 +34,9 @@ vi.mock('@/lib/execution/cancellation', () => ({
   isRedisCancellationEnabled: () => false,
   isExecutionCancelled: vi.fn().mockResolvedValue(false),
 }))
+vi.mock('@/lib/data-drains/destinations/utils', () => ({
+  sleepUntilAborted: mockSleepUntilAborted,
+}))
 vi.mock('@/lib/execution/remote-sandbox/pi-lifetime', async (importOriginal) => {
   const original =
     await importOriginal<typeof import('@/lib/execution/remote-sandbox/pi-lifetime')>()
@@ -218,6 +223,7 @@ describe('runBabysitPiWithOptions', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     mockResolvePiSandboxLifetime.mockReturnValue(undefined)
+    mockSleepUntilAborted.mockResolvedValue(undefined)
     mockFetchDiagnostics.mockResolvedValue(new Map([['check:ci', 'failure output']]))
     mockReplyAndResolve.mockResolvedValue({
       repliesPosted: 1,
@@ -507,11 +513,13 @@ describe('runBabysitPiWithOptions', () => {
     const result = await runBabysitPiWithOptions(
       params({ maxRounds: 1 }),
       { onEvent: vi.fn() },
-      { roundWaitMs: 0 }
+      { roundWaitMs: 1 }
     )
 
     expect(result).toMatchObject({ stopReason: 'clean', rounds: 0, checksGreen: true })
     expect(runCalls.some(({ command }) => command.includes('pi -p --mode json'))).toBe(false)
+    expect(mockSleepUntilAborted).toHaveBeenCalledWith(1, expect.any(AbortSignal))
+    expect(runCalls.some(({ command }) => command === 'true')).toBe(true)
   })
 
   it('returns pushed_awaiting_confirmation after replying against a lagging GitHub record', async () => {

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

@@ -89,6 +89,7 @@ const MAX_REVIEW_PROMPT_BYTES = 250_000
 const MAX_CHECK_PROMPT_BYTES = 400_000
 const MIN_ROUND_BUDGET_MS = 5 * 60 * 1000
 const ROUND_FINALIZATION_RESERVE_MS = 2 * FINALIZE_TIMEOUT_MS
+const SANDBOX_KEEPALIVE_INTERVAL_MS = 4 * 60 * 1000
 
 const BABYSIT_GUIDANCE =
   'You are fixing an existing pull request in a long-lived automated sandbox. Make only minimal ' +
@@ -487,6 +488,25 @@ async function waitForHeadConvergence(
   return 'lagging'
 }
 
+async function waitWithSandboxKeepalive(
+  runner: PiSandboxRunner,
+  durationMs: number,
+  signal: AbortSignal
+): Promise<void> {
+  let remainingMs = durationMs
+  while (remainingMs > 0) {
+    const intervalMs = Math.min(remainingMs, SANDBOX_KEEPALIVE_INTERVAL_MS)
+    await sleepUntilAborted(intervalMs, signal)
+    if (signal.aborted) throw new Error('Pi run aborted')
+    const keepalive = await raceAbort(
+      runner.run('true', { timeoutMs: FINALIZE_TIMEOUT_MS }),
+      signal
+    )
+    if (keepalive.exitCode !== 0) throw new Error('Babysit sandbox keepalive failed')
+    remainingMs -= intervalMs
+  }
+}
+
 function outstandingReason(
   fallback: BabysitStopReason,
   threads: BabysitThreadsState,
@@ -664,8 +684,7 @@ export async function runBabysitPiWithOptions(
             )
             return resultFor(totals, reason, progress, threadsClean, latestChecks!.checksGreen)
           }
-          await sleepUntilAborted(options.roundWaitMs, signal)
-          if (signal.aborted) throw new Error('Pi run aborted')
+          await waitWithSandboxKeepalive(runner, options.roundWaitMs, signal)
           snapshot = await fetchBabysitSnapshot(params, signal)
           assertBabysitPinned(
             { headSha: pinnedHeadSha, headRef: pinnedHeadRef, baseRef: pinnedBaseRef },
@@ -957,8 +976,7 @@ export async function runBabysitPiWithOptions(
           remainingBeforeWait >
             options.roundWaitMs + MIN_ROUND_BUDGET_MS + ROUND_FINALIZATION_RESERVE_MS
         ) {
-          await sleepUntilAborted(options.roundWaitMs, signal)
-          if (signal.aborted) throw new Error('Pi run aborted')
+          await waitWithSandboxKeepalive(runner, options.roundWaitMs, signal)
         }
 
         snapshot = await fetchBabysitSnapshot(params, signal)

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

@@ -161,6 +161,16 @@ describe('Babysit GitHub orchestration', () => {
     })
   })
 
+  it('accepts a complete empty rollup when the commit has no checks', async () => {
+    mockExecuteTool.mockResolvedValueOnce(checkPage([], { state: null }))
+
+    const state = await fetchBabysitCheckState(params, HEAD_SHA)
+
+    expect(state.checks).toEqual([])
+    expect(state.checksGreen).toBe(true)
+    expect(state.contextRequirements).toEqual(new Map())
+  })
+
   it('treats EXPECTED and incomplete checks as pending and optional failures as non-blocking', async () => {
     mockExecuteTool.mockResolvedValueOnce(
       checkPage([

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

@@ -433,14 +433,10 @@ export async function fetchBabysitCheckState(
       cursor = endCursor
       if (page === MAX_PAGES - 1) throw new Error('Check listing exceeded its page bound')
     }
-    const canSynthesizeMissingContexts = !!initialRequirements?.size
-    if (expectedTotal === undefined || (expectedTotal === 0 && !canSynthesizeMissingContexts)) {
+    if (expectedTotal === undefined) {
       throw new Error('GitHub returned no check rollup for the pinned commit')
     }
-    if (
-      (rollupState === null || rollupState === undefined) &&
-      !(expectedTotal === 0 && canSynthesizeMissingContexts)
-    ) {
+    if ((rollupState === null || rollupState === undefined) && expectedTotal !== 0) {
       throw new Error('GitHub returned checks without a rollup state')
     }
   } catch (error) {