fix(pi): preserve Daytona babysit budget

Author: Bill Leoutsakos

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

@@ -12,6 +12,7 @@ const {
   mockReplyAndResolve,
   mockRequestReview,
   mockReviewLanded,
+  mockResolvePiSandboxLifetime,
 } = vi.hoisted(() => ({
   mockWithPiSandbox: vi.fn(),
   mockFetchSnapshot: vi.fn(),
@@ -21,6 +22,7 @@ const {
   mockReplyAndResolve: vi.fn(),
   mockRequestReview: vi.fn(),
   mockReviewLanded: vi.fn(),
+  mockResolvePiSandboxLifetime: vi.fn(),
 }))
 
 vi.mock('@/lib/execution/remote-sandbox', () => ({
@@ -30,6 +32,14 @@ vi.mock('@/lib/execution/cancellation', () => ({
   isRedisCancellationEnabled: () => false,
   isExecutionCancelled: vi.fn().mockResolvedValue(false),
 }))
+vi.mock('@/lib/execution/remote-sandbox/pi-lifetime', async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import('@/lib/execution/remote-sandbox/pi-lifetime')>()
+  return {
+    ...original,
+    resolvePiSandboxLifetimeMs: mockResolvePiSandboxLifetime,
+  }
+})
 vi.mock('@/executor/handlers/pi/babysit-github', async (importOriginal) => {
   const original = await importOriginal<typeof import('@/executor/handlers/pi/babysit-github')>()
   return {
@@ -44,7 +54,11 @@ vi.mock('@/executor/handlers/pi/babysit-github', async (importOriginal) => {
   }
 })
 
-import { runBabysitPiWithOptions } from '@/executor/handlers/pi/babysit-backend'
+import { getMaxExecutionTimeout } from '@/lib/core/execution-limits'
+import {
+  resolveBabysitExecutionBudgetMs,
+  runBabysitPiWithOptions,
+} from '@/executor/handlers/pi/babysit-backend'
 import { BABYSIT_ROUND_PATH } from '@/executor/handlers/pi/babysit-round'
 import type { PiBabysitRunParams } from '@/executor/handlers/pi/backend'
 import { DIFF_PATH } from '@/executor/handlers/pi/cloud-shared'
@@ -195,6 +209,7 @@ function makeRunner(options: {
 describe('runBabysitPiWithOptions', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    mockResolvePiSandboxLifetime.mockReturnValue(undefined)
     mockFetchDiagnostics.mockResolvedValue(new Map([['check:ci', 'failure output']]))
     mockReplyAndResolve.mockResolvedValue({
       repliesPosted: 1,
@@ -213,6 +228,10 @@ describe('runBabysitPiWithOptions', () => {
     mockReviewLanded.mockResolvedValue(false)
   })
 
+  it('uses the platform execution budget when the provider has no absolute lifetime', () => {
+    expect(resolveBabysitExecutionBudgetMs()).toBe(getMaxExecutionTimeout())
+  })
+
   it('returns clean before sandbox creation when the PR already needs nothing', async () => {
     mockFetchSnapshot.mockResolvedValue(snapshot)
     mockFetchThreads.mockResolvedValue({

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

@@ -7,13 +7,11 @@
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import { truncate } from '@sim/utils/string'
+import { getMaxExecutionTimeout } from '@/lib/core/execution-limits'
 import { sleepUntilAborted } from '@/lib/data-drains/destinations/utils'
 import { isExecutionCancelled, isRedisCancellationEnabled } from '@/lib/execution/cancellation'
 import { type PiSandboxRunner, withPiSandbox } from '@/lib/execution/remote-sandbox'
-import {
-  PI_SANDBOX_MAX_LIFETIME_MS,
-  resolvePiSandboxLifetimeMs,
-} from '@/lib/execution/remote-sandbox/pi-lifetime'
+import { resolvePiSandboxLifetimeMs } from '@/lib/execution/remote-sandbox/pi-lifetime'
 import {
   assertBabysitPinned,
   type BabysitCheckState,
@@ -513,6 +511,11 @@ function outstandingReason(
   return fallback
 }
 
+/** Resolves the host-side run budget without imposing E2B's lifetime on other providers. */
+export function resolveBabysitExecutionBudgetMs(executionBudgetMs?: number): number {
+  return executionBudgetMs ?? resolvePiSandboxLifetimeMs() ?? getMaxExecutionTimeout()
+}
+
 /** Injectable variant used by deterministic multi-round tests. */
 export async function runBabysitPiWithOptions(
   params: PiBabysitRunParams,
@@ -540,8 +543,7 @@ export async function runBabysitPiWithOptions(
   )
   const { signal } = cancellation
   const startedAt = Date.now()
-  const lifetime =
-    params.executionBudgetMs ?? resolvePiSandboxLifetimeMs() ?? PI_SANDBOX_MAX_LIFETIME_MS
+  const lifetime = resolveBabysitExecutionBudgetMs(params.executionBudgetMs)
   if (lifetime < MIN_ROUND_BUDGET_MS) {
     cancellation.cleanup()
     throw new Error(