Fix process callbacks, PID capture, and getLogs race

For background execution, the exit code file was written before output
labelers finished writing to the log. Added a marker file that signals
when labelers complete, ensuring all output is captured before getLogs
reads the log file.

Author: ghostwriternr

.changeset/fix-process-callbacks.md

@@ -0,0 +1,5 @@
+---
+'@cloudflare/sandbox': patch
+---
+
+Fix process callbacks, PID capture, and getLogs race condition for fast commands

packages/sandbox-container/src/services/process-service.ts

@@ -150,7 +150,9 @@ export class ProcessService {
         command,
         async (event) => {
           // Route events to process record listeners
-          if (event.type === 'stdout' && event.data) {
+          if (event.type === 'start' && event.pid !== undefined) {
+            await this.store.update(processRecord.id, { pid: event.pid });
+          } else if (event.type === 'stdout' && event.data) {
             processRecord.stdout += event.data;
             processRecord.outputListeners.forEach((listener) => {
               listener('stdout', event.data!);

packages/sandbox-container/src/services/process-store.ts

@@ -61,16 +61,16 @@ export class ProcessStore {
       return;
     }
 
-    const updated = { ...existing, ...data };
-    this.processes.set(id, updated);
+    // Mutate in place to preserve reference held by ProcessService for event routing
+    Object.assign(existing, data);
 
     // Persist terminal states to disk and free memory
     const isTerminal = ['completed', 'failed', 'killed', 'error'].includes(
-      updated.status
+      existing.status
     );
     if (isTerminal) {
       try {
-        await this.writeProcessFile(id, updated);
+        await this.writeProcessFile(id, existing);
       } catch (error) {
         // Write failed, still delete to prevent memory leak
         // Explicit tradeoff: container stability > process history

packages/sandbox-container/src/services/session-manager.ts

@@ -220,18 +220,76 @@ export class SessionManager {
 
       const session = sessionResult.data;
 
-      // Get async generator
       const generator = session.execStream(command, { commandId, cwd, env });
 
-      // CRITICAL: Await first event to ensure command is tracked before returning
-      // This prevents race condition where killCommand() is called before trackCommand()
+      /*
+       * EVENT DRAINING STRATEGY
+       *
+       * We process the first 2 events synchronously before returning. Why 2?
+       *
+       * Event sequence for commands:
+       *   Fast command (echo 123):  start → stdout → complete
+       *   No-output command (true): start → complete
+       *   Long-running (sleep 10):  start → ... (events trickle in over time)
+       *
+       * - Event 1 is always 'start' (with PID) - must await to track the command
+       * - Event 2 is typically 'stdout' (first output) or 'complete' (no output)
+       *
+       * By processing 2 events, getLogs() returns correct data for fast commands
+       * like "echo 123" without blocking long-running commands unnecessarily.
+       * This is a heuristic - commands with multiple stdout chunks may have
+       * partial output if getLogs() is called immediately after startProcess().
+       */
+
       const firstResult = await generator.next();
 
-      if (!firstResult.done) {
-        await onEvent(firstResult.value);
+      if (firstResult.done) {
+        return {
+          success: true,
+          data: { continueStreaming: Promise.resolve() }
+        };
+      }
+
+      await onEvent(firstResult.value);
+
+      if (
+        firstResult.value.type === 'complete' ||
+        firstResult.value.type === 'error'
+      ) {
+        for await (const event of generator) {
+          await onEvent(event);
+        }
+        return {
+          success: true,
+          data: { continueStreaming: Promise.resolve() }
+        };
+      }
+
+      const secondResult = await generator.next();
+
+      if (secondResult.done) {
+        return {
+          success: true,
+          data: { continueStreaming: Promise.resolve() }
+        };
+      }
+
+      await onEvent(secondResult.value);
+
+      if (
+        secondResult.value.type === 'complete' ||
+        secondResult.value.type === 'error'
+      ) {
+        for await (const event of generator) {
+          await onEvent(event);
+        }
+        return {
+          success: true,
+          data: { continueStreaming: Promise.resolve() }
+        };
       }
 
-      // Create background task for remaining events
+      // Process still running after 2 events - continue in background
       const continueStreaming = (async () => {
         try {
           for await (const event of generator) {

packages/sandbox-container/src/session.ts

@@ -310,6 +310,10 @@ export class Session {
     const logFile = join(this.sessionDir!, `${commandId}.log`);
     const exitCodeFile = join(this.sessionDir!, `${commandId}.exit`);
     const pidFile = join(this.sessionDir!, `${commandId}.pid`);
+    const labelersDoneFile = join(
+      this.sessionDir!,
+      `${commandId}.labelers.done`
+    );
 
     this.logger.info('Streaming command execution started', {
       sessionId: this.id,
@@ -340,10 +344,14 @@ export class Session {
         throw new Error('Shell stdin is not available');
       }
 
+      // Wait for PID file to be created (bash script writes it after starting command)
+      const pid = await this.waitForPidFile(pidFile);
+
       yield {
         type: 'start',
         timestamp: new Date().toISOString(),
-        command
+        command,
+        pid
       };
 
       // Hybrid approach: poll log file until exit code is written
@@ -400,14 +408,41 @@ export class Session {
         await Bun.sleep(CONFIG.STREAM_CHUNK_DELAY_MS);
       }
 
-      // Read final chunks
+      /*
+       * Wait for labelers done marker file.
+       * The exit code file is written by the command subshell, but labelers
+       * run in parallel background processes. The background monitor creates
+       * the labelers done file after waiting for labelers to finish.
+       */
+      const maxWaitMs = 5000;
+      const startWait = Date.now();
+      let labelersDone = false;
+      while (Date.now() - startWait < maxWaitMs) {
+        const doneFile = Bun.file(labelersDoneFile);
+        if (await doneFile.exists()) {
+          labelersDone = true;
+          break;
+        }
+        await Bun.sleep(CONFIG.STREAM_CHUNK_DELAY_MS);
+      }
+
+      if (!labelersDone) {
+        this.logger.warn('Output capture timeout - logs may be incomplete', {
+          commandId,
+          sessionId: this.id,
+          timeoutMs: maxWaitMs
+        });
+      }
+
+      // Read final chunks from log file after labelers are done
       const file = Bun.file(logFile);
       if (await file.exists()) {
-        const content = await file.text();
-        const newContent = content.slice(position);
+        const logContent = await file.text();
+        const finalContent = logContent.slice(position);
 
-        if (newContent) {
-          const lines = newContent.split('\n');
+        // Process final chunks
+        if (finalContent) {
+          const lines = finalContent.split('\n');
           for (const line of lines) {
             if (!line) continue;
 
@@ -428,6 +463,13 @@ export class Session {
         }
       }
 
+      // Clean up labelers done file
+      try {
+        await rm(labelersDoneFile, { force: true });
+      } catch {
+        // Ignore cleanup errors
+      }
+
       // Parse exit code (already read during polling loop)
       const exitCode = parseInt(exitCodeContent, 10);
       if (Number.isNaN(exitCode)) {
@@ -629,6 +671,7 @@ export class Session {
     const stdoutPipe = join(this.sessionDir!, `${cmdId}.stdout.pipe`);
     const stderrPipe = join(this.sessionDir!, `${cmdId}.stderr.pipe`);
     const pidFile = join(this.sessionDir!, `${cmdId}.pid`);
+    const labelersDoneFile = join(this.sessionDir!, `${cmdId}.labelers.done`);
 
     // Escape paths for safe shell usage
     const safeStdoutPipe = this.escapeShellPath(stdoutPipe);
@@ -637,6 +680,7 @@ export class Session {
     const safeExitCodeFile = this.escapeShellPath(exitCodeFile);
     const safeSessionDir = this.escapeShellPath(this.sessionDir!);
     const safePidFile = this.escapeShellPath(pidFile);
+    const safeLabelersDoneFile = this.escapeShellPath(labelersDoneFile);
 
     const indentLines = (input: string, spaces: number) => {
       const prefix = ' '.repeat(spaces);
@@ -722,6 +766,7 @@ export class Session {
         script += `    (\n`;
         script += `      wait "$r1" "$r2" 2>/dev/null\n`;
         script += `      rm -f "$sp" "$ep"\n`;
+        script += `      touch ${safeLabelersDoneFile}\n`;
         script += `    ) &\n`;
         script += `    # Restore directory immediately\n`;
         script += `    cd "$PREV_DIR"\n`;
@@ -745,6 +790,7 @@ export class Session {
         script += `  (\n`;
         script += `    wait "$r1" "$r2" 2>/dev/null\n`;
         script += `    rm -f "$sp" "$ep"\n`;
+        script += `    touch ${safeLabelersDoneFile}\n`;
         script += `  ) &\n`;
       }
     } else {
@@ -1003,6 +1049,35 @@ export class Session {
     }
   }
 
+  /**
+   * Wait for PID file to be created and return the PID
+   * Returns undefined if file doesn't appear within timeout
+   */
+  private async waitForPidFile(
+    pidFile: string,
+    timeoutMs: number = 1000
+  ): Promise<number | undefined> {
+    const startTime = Date.now();
+
+    while (Date.now() - startTime < timeoutMs) {
+      try {
+        const file = Bun.file(pidFile);
+        if (await file.exists()) {
+          const content = await file.text();
+          const pid = parseInt(content.trim(), 10);
+          if (!Number.isNaN(pid)) {
+            return pid;
+          }
+        }
+      } catch {
+        // Ignore errors, keep polling
+      }
+      await Bun.sleep(10); // Poll every 10ms
+    }
+
+    return undefined;
+  }
+
   /**
    * Escape shell path for safe usage in bash scripts
    */

packages/sandbox/src/sandbox.ts

@@ -1289,6 +1289,16 @@ export class Sandbox<Env = unknown> extends Container<Env> implements ISandbox {
         options.onStart(processObj);
       }
 
+      // Start background streaming if output/exit callbacks are provided
+      if (options?.onOutput || options?.onExit) {
+        // Fire and forget - don't await, let it run in background
+        this.startProcessCallbackStream(response.processId, options).catch(
+          () => {
+            // Error already handled in startProcessCallbackStream
+          }
+        );
+      }
+
       return processObj;
     } catch (error) {
       if (options?.onError && error instanceof Error) {
@@ -1299,6 +1309,56 @@ export class Sandbox<Env = unknown> extends Container<Env> implements ISandbox {
     }
   }
 
+  /**
+   * Start background streaming for process callbacks
+   * Opens SSE stream to container and routes events to callbacks
+   */
+  private async startProcessCallbackStream(
+    processId: string,
+    options: ProcessOptions
+  ): Promise<void> {
+    try {
+      const stream = await this.client.processes.streamProcessLogs(processId);
+
+      for await (const event of parseSSEStream<{
+        type: string;
+        data?: string;
+        exitCode?: number;
+        processId?: string;
+      }>(stream)) {
+        switch (event.type) {
+          case 'stdout':
+            if (event.data && options.onOutput) {
+              options.onOutput('stdout', event.data);
+            }
+            break;
+          case 'stderr':
+            if (event.data && options.onOutput) {
+              options.onOutput('stderr', event.data);
+            }
+            break;
+          case 'exit':
+          case 'complete':
+            if (options.onExit) {
+              options.onExit(event.exitCode ?? null);
+            }
+            return; // Stream complete
+        }
+      }
+    } catch (error) {
+      // Call onError if streaming fails
+      if (options.onError && error instanceof Error) {
+        options.onError(error);
+      }
+      // Don't rethrow - background streaming failure shouldn't crash the caller
+      this.logger.error(
+        'Background process streaming failed',
+        error instanceof Error ? error : new Error(String(error)),
+        { processId }
+      );
+    }
+  }
+
   async listProcesses(sessionId?: string): Promise<Process[]> {
     const session = sessionId ?? (await this.ensureDefaultSession());
     const response = await this.client.processes.listProcesses();

packages/shared/src/types.ts

@@ -209,6 +209,7 @@ export interface ExecEvent {
   result?: ExecResult;
   error?: string;
   sessionId?: string;
+  pid?: number; // Present on 'start' event
 }
 
 export interface LogEvent {