fix(cloud-agent-next): avoid OOM by downloading snapshot in sandbox (#908)

## Summary

- Move snapshot download from worker memory to sandbox via `curl` — the
full JSON never materializes in the 128MB worker
- Replace worker-side diff extraction (`extractDiffsFromMessages`) and
application (`applySessionDiff`) with a sandbox-side Node script that
reads, deduplicates, and applies diffs directly on disk
- Refactor `restoreSessionSnapshot` to accept a file path (curl writes
directly) instead of a string payload
- Remove `resolve` and `zod` imports that were only used by deleted
methods

Depends on PR #884 which adds the streaming `/api/session/:id/export`
endpoint.

## Test plan

- [x] All 601 existing tests pass (updated 4 cold-start tests for new
exec-based flow)
- [x] Typecheck passes
- [ ] Integration test: cold-start resume with real session to verify
curl download + kilo import + diff application

Author: eshurakov

cloud-agent-next/Dockerfile

@@ -12,7 +12,7 @@ RUN mkdir -p -m 755 /etc/apt/keyrings \
 	&& mkdir -p -m 755 /etc/apt/sources.list.d \
 	&& echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | tee /etc/apt/sources.list.d/github-cli.list > /dev/null \
 	&& apt update \
-	&& apt install gh -y
+	&& apt install gh jq -y
 
 # Install GitLab CLI (glab) - download official .deb from GitLab releases
 RUN GLAB_VERSION="1.80.4" \
@@ -39,9 +39,10 @@ RUN npm install -g pnpm @kilocode/cli@${KILOCODE_CLI_VERSION}
 COPY wrapper /tmp/wrapper-build/wrapper
 COPY src/shared /tmp/wrapper-build/src/shared
 
-# Build the wrapper bundle
+# Build the wrapper bundle and restore-session script
 RUN cd /tmp/wrapper-build/wrapper && \
     bun build src/main.ts --outfile=/usr/local/bin/kilocode-wrapper.js --target=bun --minify && \
+    bun build src/restore-session.ts --outfile=/usr/local/bin/kilo-restore-session.js --target=bun --minify && \
     rm -rf /tmp/wrapper-build
 
 # DO NOT override USER, WORKDIR, or ENTRYPOINT from the base image

cloud-agent-next/Dockerfile.dev

@@ -12,7 +12,7 @@ ARG KILOCODE_CLI_VERSION="next"
 # This builds kilo-cli from source and copies the linux-x64 binary here.
 
 # Install ripgrep and Generate locales to suppress setlocale warnings
-RUN apt-get update && apt-get install -y --no-install-recommends ripgrep locales && \
+RUN apt-get update && apt-get install -y --no-install-recommends ripgrep jq locales && \
 	   sed -i -e 's/# en_US.UTF-8 UTF-8/en_US.UTF-8 UTF-8/' /etc/locale.gen && \
 	   locale-gen && \
 	   apt-get clean && \
@@ -34,14 +34,15 @@ RUN npm install -g pnpm @kilocode/cli@${KILOCODE_CLI_VERSION}
 #COPY kilo /usr/local/bin/kilo
 #RUN chmod +x /usr/local/bin/kilo
 
-# === Build wrapper bundle inside container ===
-# This ensures the wrapper is built with the same Bun version that will run it,
+# === Build wrapper bundle and restore-session script inside container ===
+# This ensures both scripts are built with the same Bun version that will run them,
 # avoiding any compatibility issues between build and runtime environments.
 COPY wrapper /tmp/wrapper-build/wrapper
 COPY src/shared /tmp/wrapper-build/src/shared
 
 RUN cd /tmp/wrapper-build/wrapper && \
     bun build src/main.ts --outfile=/usr/local/bin/kilocode-wrapper.js --target=bun --minify && \
+    bun build src/restore-session.ts --outfile=/usr/local/bin/kilo-restore-session.js --target=bun --minify && \
     rm -rf /tmp/wrapper-build
 
 # DO NOT override USER, WORKDIR, or ENTRYPOINT from the base image

cloud-agent-next/src/session-service.test.ts

@@ -1,5 +1,3 @@
-import { resolve } from 'node:path';
-import { z } from 'zod';
 import type {
   ExecutionSession,
   SandboxInstance,
@@ -933,185 +931,6 @@ export class SessionService {
     };
   }
 
-  /**
-   * Write a snapshot payload to a temp file, run `kilo import`, then clean up.
-   *
-   * @param snapshotPayload - Pre-fetched JSON string of the session snapshot.
-   */
-  private async restoreSessionSnapshot(
-    session: ExecutionSession,
-    sessionId: string,
-    userId: string,
-    snapshotPayload: string
-  ): Promise<void> {
-    const tmpPath = `/tmp/kilo-session-export-${sessionId}.json`;
-    let wroteSnapshot = false;
-    try {
-      await session.writeFile(tmpPath, snapshotPayload);
-      wroteSnapshot = true;
-
-      const importResult = await session.exec(`kilo import "${tmpPath}"`);
-      if (importResult.exitCode !== 0) {
-        logger
-          .withFields({
-            sessionId,
-            userId,
-            exitCode: importResult.exitCode,
-            stderr: importResult.stderr,
-            stdout: importResult.stdout,
-          })
-          .error('Session snapshot import failed');
-        throw new Error(`Session snapshot import failed with exit code ${importResult.exitCode}`);
-      }
-    } catch (error) {
-      logger
-        .withFields({
-          sessionId,
-          userId,
-          error: error instanceof Error ? error.message : String(error),
-        })
-        .error('Session snapshot restore failed');
-      throw error instanceof Error ? error : new Error(String(error));
-    } finally {
-      if (wroteSnapshot) {
-        try {
-          await session.deleteFile(tmpPath);
-        } catch (error) {
-          logger
-            .withFields({
-              sessionId,
-              userId,
-              error: error instanceof Error ? error.message : String(error),
-            })
-            .debug('Failed to delete session snapshot temp file');
-        }
-      }
-    }
-  }
-
-  /**
-   * Apply file-level changes from a pre-parsed diff array on top of the freshly
-   * cloned repo.  Called during cold-start resume after `restoreSessionSnapshot`.
-   *
-   * Each diff entry contains the full `after` content, so we simply write (or
-   * delete) files to recreate the workspace state from the previous session.
-   *
-   * @param diffs - Pre-parsed `FileDiff[]` array (or `null` when no diff exists).
-   */
-  private async applySessionDiff(
-    session: ExecutionSession,
-    sessionId: string,
-    userId: string,
-    workspacePath: string,
-    diffs: Array<{
-      file: string;
-      after: string;
-      status?: 'added' | 'deleted' | 'modified';
-    }> | null
-  ): Promise<void> {
-    if (!Array.isArray(diffs) || diffs.length === 0) {
-      return;
-    }
-
-    let applied = 0;
-    let skipped = 0;
-
-    for (const diff of diffs) {
-      if (!diff.file) {
-        skipped++;
-        continue;
-      }
-
-      // Skip binary files — they have empty after content
-      if (diff.status !== 'deleted' && !diff.after) {
-        skipped++;
-        continue;
-      }
-
-      const filePath = resolve(workspacePath, diff.file);
-      if (!filePath.startsWith(workspacePath + '/')) {
-        logger
-          .withFields({ sessionId, userId, file: diff.file })
-          .warn('Skipping diff entry with path outside workspace');
-        skipped++;
-        continue;
-      }
-
-      try {
-        if (diff.status === 'deleted') {
-          await session.deleteFile(filePath);
-          applied++;
-        } else {
-          // Ensure parent directory exists.
-          // Use single-quoted path to prevent shell metacharacter injection.
-          const lastSlash = filePath.lastIndexOf('/');
-          if (lastSlash > 0) {
-            const parentDir = filePath.substring(0, lastSlash);
-            const escaped = parentDir.replaceAll("'", "'\\''");
-            await session.exec(`mkdir -p '${escaped}'`);
-          }
-          await session.writeFile(filePath, diff.after);
-          applied++;
-        }
-      } catch (error) {
-        logger
-          .withFields({
-            sessionId,
-            userId,
-            file: diff.file,
-            status: diff.status,
-            error: error instanceof Error ? error.message : String(error),
-          })
-          .warn('Failed to apply file diff (non-fatal, continuing)');
-        skipped++;
-      }
-    }
-
-    logger
-      .withFields({ sessionId, userId, applied, skipped, total: diffs.length })
-      .info('Applied session diff');
-  }
-
-  /**
-   * Extract last-write-wins file diffs from streamed snapshot messages.
-   * The snapshot format is `{ info, messages: [{ info: { summary: { diffs } }, parts }] }`.
-   * Returns null when no diffs exist.
-   */
-  private static extractDiffsFromMessages(
-    payload: string
-  ): Array<{ file: string; after: string; status?: 'added' | 'deleted' | 'modified' }> | null {
-    const fileDiffSchema = z.object({
-      file: z.string(),
-      after: z.string().default(''),
-      status: z.enum(['added', 'deleted', 'modified']).default('modified'),
-    });
-
-    try {
-      const parsed = JSON.parse(payload) as {
-        messages?: Array<{ info?: { summary?: { diffs?: unknown[] } } }>;
-      };
-      if (!parsed.messages) return null;
-
-      const byFile = new Map<string, z.infer<typeof fileDiffSchema>>();
-
-      for (const msg of parsed.messages) {
-        const diffs = msg.info?.summary?.diffs;
-        if (!Array.isArray(diffs)) continue;
-
-        for (const d of diffs) {
-          const result = fileDiffSchema.safeParse(d);
-          if (!result.success) continue;
-          byFile.set(result.data.file, result.data);
-        }
-      }
-
-      if (byFile.size === 0) return null;
-      return [...byFile.values()];
-    } catch {
-      return null;
-    }
-  }
-
   /**
    * Initialize a cloud-agent session by resuming an existing kilo session.
    *
@@ -1475,38 +1294,62 @@ export class SessionService {
       await writeAuthFile(sandbox, context.sessionHome, kilocodeToken);
       await writeGlobalRules(sandbox, context.sessionHome, sessionId);
 
-      // Fetch snapshot from session-ingest DO and buffer it for sandbox writeFile (string-only API).
-      const internalSecret = await env.INTERNAL_API_SECRET_PROD.get();
-      const response = await env.SESSION_INGEST.fetch(
-        new Request(`https://session-ingest/internal/session/${metadata.kiloSessionId}/export`, {
-          headers: {
-            'X-Internal-Secret': internalSecret,
-            'X-Kilo-User-Id': userId,
-          },
-        })
+      // Single restore script handles download, import, and diff application inside
+      // the sandbox — the snapshot never enters worker memory.
+      logger.info('Starting cold-start session restore');
+
+      const escapedId = metadata.kiloSessionId.replaceAll("'", "'\\''");
+      const escapedWorkspace = context.workspacePath.replaceAll("'", "'\\''");
+      const restoreResult = await session.exec(
+        `bun /usr/local/bin/kilo-restore-session.js '${escapedId}' '${escapedWorkspace}'`
       );
 
-      if (response.status === 404) {
+      if (restoreResult.exitCode !== 0) {
+        logger
+          .withFields({
+            sessionId,
+            userId,
+            exitCode: restoreResult.exitCode,
+            stderr: restoreResult.stderr,
+            stdout: restoreResult.stdout,
+          })
+          .error('Cold-start session restore failed');
+
+        // Parse stdout JSON for structured error info
+        let code: number | undefined;
+        try {
+          const parsed = JSON.parse(restoreResult.stdout?.trim() ?? '{}') as Record<
+            string,
+            unknown
+          >;
+          if (typeof parsed.code === 'number') {
+            code = parsed.code;
+          }
+        } catch {
+          // non-JSON stdout, ignore
+        }
+
+        if (code === 404) {
+          throw new SessionSnapshotRestoreError(
+            'Session snapshot restore failed: session not found',
+            404
+          );
+        }
         throw new SessionSnapshotRestoreError(
-          `Session snapshot restore failed: session not found`,
-          404
+          `Cold-start session restore failed: exit ${restoreResult.exitCode}`,
+          code
         );
       }
-      if (!response.ok) {
-        throw new Error(`Session export failed: ${response.status}`);
-      }
 
-      const snapshotPayload = await response.text();
-
-      // Extract diffs client-side from the streamed snapshot messages.
-      const diffs = SessionService.extractDiffsFromMessages(snapshotPayload);
-
-      await this.restoreSessionSnapshot(session, sessionId, userId, snapshotPayload);
-
-      // Apply file-level changes from the previous session on top of the fresh clone.
-      // This runs after kilo import (conversation restore) since the CLI doesn't need
-      // the file state during import — it only restores its internal session DB.
-      await this.applySessionDiff(session, sessionId, userId, context.workspacePath, diffs);
+      // Log structured summary from restore script
+      try {
+        const summary = JSON.parse(restoreResult.stdout?.trim() ?? '{}') as Record<string, unknown>;
+        logger
+          .withFields({ sessionId, userId, ...summary })
+          .info('Cold-start session restore completed');
+      } catch {
+        // non-JSON stdout, non-fatal
+      }
 
       // Re-run setup commands (fresh clone, need to reinstall)
       if (metadata.setupCommands && metadata.setupCommands.length > 0) {

cloud-agent-next/src/session-service.ts

@@ -7,4 +7,12 @@ await Bun.build({
   sourcemap: 'external',
 });
 
-console.log('Build complete: dist/wrapper.js');
+await Bun.build({
+  entrypoints: ['./src/restore-session.ts'],
+  outdir: './dist',
+  naming: 'restore-session.js',
+  target: 'bun',
+  minify: true,
+});
+
+console.log('Build complete: dist/wrapper.js, dist/restore-session.js');