improvement(sandbox): log per-phase timing for each sandbox execution (#9587)

* improvement(sandbox): log per-phase timing for each sandbox execution

* test(sandbox): time the unchanged-CLI warm path on a controlled clock

Author: waleedlatif1

apps/sim/lib/execution/remote-sandbox/index.ts

@@ -43,6 +43,10 @@ import {
   SANDBOX_OUTPUT_DIR_MAX_DEPTH,
   SANDBOX_OUTPUT_DIR_SENTINEL,
 } from '@/lib/execution/remote-sandbox/paths'
+import {
+  createSandboxPhaseTimer,
+  type SandboxPhaseTimer,
+} from '@/lib/execution/remote-sandbox/phase-timing'
 import { resolvePiSandboxLifetimeMs } from '@/lib/execution/remote-sandbox/pi-lifetime'
 import { resolveProvider } from '@/lib/execution/remote-sandbox/provider'
 import {
@@ -116,13 +120,16 @@ async function leaseSandbox(
   selected: ResolvedSandbox | null,
   signal: AbortSignal,
   meterUsage: boolean | undefined,
-  session: SandboxSessionRequest | undefined
+  session: SandboxSessionRequest | undefined,
+  phases: SandboxPhaseTimer
 ): Promise<SandboxLease> {
   const provider = resolveProvider()
   const sessionCapable = Boolean(session && !meterUsage && provider.findSessionSandbox)
 
   if (session && sessionCapable) {
+    const endAccess = phases.start('access')
     await prepareSandboxSessionAccess(session.key, signal)
+    endAccess()
     return withSandboxSessionLock(session.key, signal, async (leaseSignal) => {
       const { created, status } = await ensureSessionSandbox({
         provider,
@@ -132,6 +139,7 @@ async function leaseSandbox(
         session,
         signal: leaseSignal,
         bootstrapTimeoutMs: remainingSandboxBudgetMs(signal),
+        phases,
       })
       return {
         created,
@@ -214,6 +222,16 @@ function raceSandboxAbort<T>(operation: Promise<T>, signal: AbortSignal): Promis
   })
 }
 
+/** Attaches the phase breakdown to whichever result the execution returns, after its release. */
+async function withPhaseTimings(
+  execute: (phases: SandboxPhaseTimer) => Promise<SandboxExecutionResult>
+): Promise<SandboxExecutionResult> {
+  const phases = createSandboxPhaseTimer()
+  const result = await execute(phases)
+  result.phaseMs = phases.timings
+  return result
+}
+
 async function withSandboxExecutionBudget<T>(
   timeoutMs: number,
   parentSignal: AbortSignal | undefined,
@@ -864,22 +882,26 @@ async function acceptSessionOutputHistory(
 async function executeInSandboxWithinBudget(
   // The budget wrapper always injects the signal; the required-signal type states that
   // invariant instead of a cast hiding it.
-  req: SandboxExecutionRequest & { signal: AbortSignal }
+  req: SandboxExecutionRequest & { signal: AbortSignal },
+  phases: SandboxPhaseTimer
 ): Promise<SandboxExecutionResult> {
   const { code, language, signal } = req
   const kind = req.sandboxKind ?? 'code'
   throwIfAborted(signal)
 
   // Resolved before the sandbox is created so a selection that cannot be honored
   // fails without spending a provider create.
+  const endResolve = phases.start('resolve')
   const selected = await resolveWorkspaceSandbox({
     kind,
     language,
     workspaceId: req.workspaceId,
     sandboxId: req.sandboxId,
   })
+  endResolve()
   throwIfAborted(signal)
 
+  const endLease = phases.start('lease')
   const lease = await leaseSandbox(
     kind,
     {
@@ -890,8 +912,10 @@ async function executeInSandboxWithinBudget(
     selected,
     signal,
     req.meterUsage,
-    req.session
+    req.session,
+    phases
   )
+  endLease()
   const created = lease.created
   const sandbox = created.sandbox
   const sandboxId = sandbox.sandboxId
@@ -906,6 +930,7 @@ async function executeInSandboxWithinBudget(
     // the finally below. Dependencies land before the inputs so user code and its
     // mounts always see a complete environment.
     //
+    const endPrepare = phases.start('prepare')
     if (lease.session && req.session)
       await recordSessionFileInput(
         req.session.key,
@@ -926,6 +951,7 @@ async function executeInSandboxWithinBudget(
       signal,
       Boolean(lease.session)
     )
+    endPrepare()
     const executionEnvironment = {
       ...selected?.envs,
       ...req.session?.envs,
@@ -944,13 +970,15 @@ async function executeInSandboxWithinBudget(
 
     let execution: SandboxCodeResult
     try {
+      const endRun = phases.start('run')
       execution = await sandbox.runCode(code, {
         timeoutMs: codeBudgetMs(req, signal),
         javascriptPreload: buildJavaScriptRuntimeBindingsSource(req.runtimeBindings ?? []),
         maxOutputBytes: MAX_SANDBOX_PROCESS_OUTPUT_BYTES,
         signal,
         ...(hasExecutionEnvironment ? { envs: executionEnvironment } : {}),
       })
+      endRun()
     } catch (error) {
       throwIfAborted(signal)
       throw error
@@ -1009,11 +1037,13 @@ async function executeInSandboxWithinBudget(
       ...sessionField,
     }
     try {
+      const endCollect = phases.start('collect')
       const { exportedFiles, exportedFileContent, collectedFiles } = await collectExportedFiles(
         sandbox,
         req,
         { signal }
       )
+      endCollect()
       throwIfAborted(signal)
       billableResult.exportedFileContent = exportedFileContent
       billableResult.exportedFiles = exportedFiles
@@ -1044,6 +1074,7 @@ async function executeInSandboxWithinBudget(
     if (cost && billableOutputError) {
       attachTrustedSandboxOutputCost(billableOutputError, cost)
     }
+    const endRelease = phases.start('release')
     try {
       await privateInputFiles?.cleanup()
       await lease.release()
@@ -1052,6 +1083,7 @@ async function executeInSandboxWithinBudget(
         providerId: created.providerId,
         sandboxId,
       })
+      endRelease()
     }
   }
 }
@@ -1061,34 +1093,41 @@ export function executeInSandbox(req: SandboxExecutionRequest): Promise<SandboxE
     req.timeoutMs,
     req.signal,
     Boolean(req.session && !req.meterUsage),
-    (signal) => executeInSandboxWithinBudget({ ...req, signal })
+    (signal) =>
+      withPhaseTimings((phases) => executeInSandboxWithinBudget({ ...req, signal }, phases))
   )
 }
 
 async function executeShellInSandboxWithinBudget(
-  req: SandboxShellExecutionRequest & { signal: AbortSignal }
+  req: SandboxShellExecutionRequest & { signal: AbortSignal },
+  phases: SandboxPhaseTimer
 ): Promise<SandboxExecutionResult> {
   const { code, envs, signal } = req
   const kind = req.sandboxKind ?? 'shell'
   throwIfAborted(signal)
 
   // No language is passed: a shell execution runs commands rather than a language
   // runtime, so whichever language the sandbox carries is the one it installs.
+  const endResolve = phases.start('resolve')
   const selected = await resolveWorkspaceSandbox({
     kind,
     workspaceId: req.workspaceId,
     sandboxId: req.sandboxId,
   })
+  endResolve()
   throwIfAborted(signal)
 
+  const endLease = phases.start('lease')
   const lease = await leaseSandbox(
     kind,
     { imageRef: selected?.imageRef, lifetimeMs: remainingSandboxBudgetMs(signal) },
     selected,
     signal,
     req.meterUsage,
-    req.session
+    req.session,
+    phases
   )
+  endLease()
   const created = lease.created
   const sandbox = created.sandbox
   const sandboxId = sandbox.sandboxId
@@ -1102,6 +1141,7 @@ async function executeShellInSandboxWithinBudget(
     // Inside the try so a failed install or mount still releases the sandbox via
     // the finally below. The install shares the caller's budget rather than adding
     // to it — see the note in `executeInSandbox`.
+    const endPrepare = phases.start('prepare')
     if (lease.session && req.session)
       await recordSessionFileInput(
         req.session.key,
@@ -1123,9 +1163,11 @@ async function executeShellInSandboxWithinBudget(
       signal,
       Boolean(lease.session)
     )
+    endPrepare()
 
     let result: SandboxCommandResult
     try {
+      const endRun = phases.start('run')
       result = await sandbox.runCommand(code, {
         envs: {
           ...selected?.envs,
@@ -1141,6 +1183,7 @@ async function executeShellInSandboxWithinBudget(
         rootUser: !lease.session,
         atMostOnce: true,
       })
+      endRun()
     } catch (error) {
       throwIfAborted(signal)
       throw error
@@ -1185,11 +1228,13 @@ async function executeShellInSandboxWithinBudget(
       ...sessionField,
     }
     try {
+      const endCollect = phases.start('collect')
       const { exportedFiles, exportedFileContent, collectedFiles } = await collectExportedFiles(
         sandbox,
         req,
         { signal }
       )
+      endCollect()
       throwIfAborted(signal)
       billableResult.exportedFileContent = exportedFileContent
       billableResult.exportedFiles = exportedFiles
@@ -1220,6 +1265,7 @@ async function executeShellInSandboxWithinBudget(
     if (cost && billableOutputError) {
       attachTrustedSandboxOutputCost(billableOutputError, cost)
     }
+    const endRelease = phases.start('release')
     try {
       await privateInputFiles?.cleanup()
       await lease.release()
@@ -1228,6 +1274,7 @@ async function executeShellInSandboxWithinBudget(
         providerId: created.providerId,
         sandboxId,
       })
+      endRelease()
     }
   }
 }
@@ -1239,7 +1286,8 @@ export function executeShellInSandbox(
     req.timeoutMs,
     req.signal,
     Boolean(req.session && !req.meterUsage),
-    (signal) => executeShellInSandboxWithinBudget({ ...req, signal })
+    (signal) =>
+      withPhaseTimings((phases) => executeShellInSandboxWithinBudget({ ...req, signal }, phases))
   )
 }
 

apps/sim/lib/execution/remote-sandbox/phase-timing.ts

@@ -0,0 +1,39 @@
+/**
+ * Phases of one sandbox execution, in the order they run. `lease` spans the whole acquisition
+ * (session lock included); for a session, `lookup` and `cli` are the parts of it spent finding or
+ * creating the machine and verifying its deployment tooling.
+ */
+type SandboxPhase =
+  | 'resolve'
+  | 'access'
+  | 'lease'
+  | 'lookup'
+  | 'cli'
+  | 'prepare'
+  | 'run'
+  | 'collect'
+  | 'release'
+
+/** Wall-clock milliseconds per phase; a phase that did not finish is absent. */
+export type SandboxPhaseTimings = Partial<Record<SandboxPhase, number>>
+
+/** Accumulates per-phase wall time for one execution's log line. */
+export interface SandboxPhaseTimer {
+  readonly timings: SandboxPhaseTimings
+  /** Starts timing a phase; the returned function records its duration. */
+  start(phase: SandboxPhase): () => void
+}
+
+/** Creates a timer whose callbacks record rounded milliseconds; unfinished phases stay absent. */
+export function createSandboxPhaseTimer(): SandboxPhaseTimer {
+  const timings: SandboxPhaseTimings = {}
+  return {
+    timings,
+    start(phase) {
+      const startedAt = performance.now()
+      return () => {
+        timings[phase] = Math.round(performance.now() - startedAt)
+      }
+    },
+  }
+}

apps/sim/lib/execution/remote-sandbox/session-sandbox.test.ts

@@ -821,6 +821,46 @@ describe('session sandbox lease', () => {
     expect(writes).toHaveBeenCalledTimes(firstWrites)
   })
 
+  it('attributes a warm shell call’s wall time to the phase that spent it', async () => {
+    const { handle } = fakeSandbox('phase-timing')
+    mockFindSessionSandbox.mockResolvedValue(handle)
+    const request = {
+      code: 'true',
+      envs: {},
+      timeoutMs: 10_000,
+      session: { key: 'phase-timing', cli: CLI },
+    }
+    await executeShellInSandbox(request)
+
+    let nowMs = 0
+    const clock = vi.spyOn(performance, 'now').mockImplementation(() => nowMs)
+    const writes = vi.fn(handle.writeFile.bind(handle))
+    handle.writeFile = writes
+    const runCommand = handle.runCommand.bind(handle)
+    handle.runCommand = async (command, options) => {
+      nowMs += options.envs?.SIM_CLI_PATH ? 30 : 60
+      return runCommand(command, options)
+    }
+    mockFindSessionSandbox.mockImplementation(async () => {
+      nowMs += 40
+      return handle
+    })
+
+    try {
+      const result = await executeShellInSandbox(request)
+
+      const phases = result.phaseMs ?? {}
+      expect(result.sandboxSession).toBe('reused')
+      expect(writes).not.toHaveBeenCalled()
+      expect(phases.lookup).toBe(40)
+      expect(phases.cli).toBe(30)
+      expect(phases.run).toBe(60)
+      expect(phases.lease).toBeGreaterThanOrEqual(70)
+    } finally {
+      clock.mockRestore()
+    }
+  })
+
   it('pins shell CLI resolution ahead of caller environment paths', async () => {
     const { handle, calls } = fakeSandbox('shell-cli-path')
     mockFindSessionSandbox.mockResolvedValue(handle)

apps/sim/lib/execution/remote-sandbox/session.ts

@@ -3,6 +3,7 @@ import { getErrorMessage } from '@sim/utils/errors'
 import { createSandboxPricing } from '@/lib/billing/sandbox-pricing'
 import { mapSettledWithConcurrency } from '@/lib/core/utils/concurrency'
 import { type CreatedSandbox, createSelectedSandbox } from '@/lib/execution/remote-sandbox/create'
+import type { SandboxPhaseTimer } from '@/lib/execution/remote-sandbox/phase-timing'
 import { resolveProvider } from '@/lib/execution/remote-sandbox/provider'
 import type { ResolvedSandbox } from '@/lib/execution/remote-sandbox/resolve'
 import { ensureSessionCli } from '@/lib/execution/remote-sandbox/session-cli'
@@ -56,12 +57,14 @@ export async function ensureSessionSandbox(args: {
   session: SandboxSessionRequest
   signal: AbortSignal
   bootstrapTimeoutMs: number
+  phases?: SandboxPhaseTimer
 }): Promise<SessionAcquisition> {
   const { provider, kind, options, selected, session, signal } = args
   signal.throwIfAborted()
   if (!provider.findSessionSandbox) throw new Error('This deployment has no persistent workbench')
   const lifetimeMs = SESSION_SANDBOX_IDLE_MS + (options.lifetimeMs ?? 0)
   const requestedAtMs = Date.now()
+  const endLookup = args.phases?.start('lookup')
   const existing = await provider.findSessionSandbox(session.key, {
     ...(options.language ? { language: options.language } : {}),
     lifetimeMs,
@@ -83,6 +86,7 @@ export async function ensureSessionSandbox(args: {
         true,
         provider
       )
+  endLookup?.()
   // Before the abort check: a stopped caller still leaves this machine for the next call to
   // reuse, and a reused machine without history can never be certified.
   if (!existing)
@@ -96,7 +100,9 @@ export async function ensureSessionSandbox(args: {
     await created.sandbox.extendLifetime?.(lifetimeMs)
   signal.throwIfAborted()
   if (session.cli) {
+    const endCli = args.phases?.start('cli')
     await ensureSessionCli(created.sandbox, session.cli, signal, args.bootstrapTimeoutMs)
+    endCli?.()
   }
   signal.throwIfAborted()
   return { created, status: existing ? 'reused' : 'created' }