fix(cloud-agent-next): key repository snapshots by setup commands

services/cloud-agent-next/docs/control-plane.md

@@ -379,12 +379,14 @@ These keep their current owners and evidence; the rewrite ports them, it does no
 Scope: `workspace_*` routes on `cloudflare-containers` whose directory is the constant isolated
 path `/workspace/app`. A snapshot holds one repository at one path, so any other route has no key.
 
-- **Key.** When a route is first prepared, the Sandbox DO hashes the owner and the repository URL
-  with HMAC-SHA256 under a Worker secret. Env is not part of the key: setup re-runs on every start.
-  It stores the digest on the route (`repo_key`) and keeps it across attempts. No key is computed, and
-  no snapshot used, when the owner is not enrolled (`CONTAINER_REPO_SNAPSHOT_IDS`), the route
-  has no repository, the provider cannot capture, or the secret is missing. Scope is per user;
-  per org is a change to that one field.
+- **Key.** When a route is first prepared, the Sandbox DO hashes the owner, the repository URL and
+  the setup commands with HMAC-SHA256 under a Worker secret. Env is not part of the key: setup
+  re-runs on every start and rewrites what it derives from the env. The setup commands join the
+  hash so a snapshot is reused only by a route with the same setup; a different setup for one
+  repository gets its own key. It stores the digest on the route (`repo_key`) and keeps it across
+  attempts. No key is computed, and no snapshot used, when the owner is not enrolled
+  (`CONTAINER_REPO_SNAPSHOT_IDS`), the route has no repository, the provider cannot capture, or the
+  secret is missing. Scope is per user; per org is a change to that one field.
 - **Start.** `launch` receives the key of the routes waiting for the allocation: exactly one key,
   else none. `SandboxContainers` appends the image, and starts from the
   repository snapshot in the `REPO_SNAPSHOTS` KV index when it has one for that digest, else from
@@ -403,8 +405,9 @@ path `/workspace/app`. A snapshot holds one repository at one path, so any other
    the entry. A completed image fallback records `image`, so the next start tries the retained
    snapshot again unless its failure budget is exhausted.
 - **Capture.** `session.prepare` carries `capture: true` when the route has a key and the provider
-  can capture. A route with setup commands has no key: capture runs after setup, so the snapshot
-  would publish the whole container under a key that names only the user and the repository. The
+  can capture. The key names the setup commands as well as the user and the repository, because
+  capture runs after setup: a snapshot built with one setup is never adopted by a route with
+  another. The
   wrapper decides whether to capture (a fresh clone, or an adopted snapshot that is
   due for a refresh) and then sends `workspace.capture` after setup, with origin bare. The Sandbox DO
   calls the provider off its serial queue, bounded at 5 min 5 s, and answers `workspace.captured`
@@ -539,8 +542,8 @@ setup the wrapper writes the stamp (before the snapshot, so a restored container
 allocation's stamp), sets `origin` to the bare URL, clears the reflogs and `FETCH_HEAD`, sends
 `workspace.capture` and waits for `workspace.captured`. It then restores the authenticated URL,
 whatever the outcome, and continues to the Kilo runtime. A snapshot therefore holds no agent edits,
-no Kilo home and no git credential. Process env never reaches the disk. A route with setup commands
-has no key, so it restores and captures nothing: setup output is never published to a snapshot.
+no Kilo home and no git credential. Process env never reaches the disk. Capture runs after setup, so
+the key names the setup commands too: a snapshot is reused only by a route with the same setup.
 
 ### Prompts and turn outcome
 

services/cloud-agent-next/src/control-plane/sandbox/repo-key.test.ts

@@ -54,15 +54,6 @@ describe('repoSnapshotEligible', () => {
       })
     ).toBe(false);
   });
-
-  it('excludes a route with setup commands and treats empty or absent the same', () => {
-    const enrolled = gate({ enrolledIds: '*' });
-    expect(repoSnapshotEligible(enrolled, { ...route, setupCommands: ['pnpm install'] })).toBe(
-      false
-    );
-    expect(repoSnapshotEligible(enrolled, { ...route, setupCommands: [] })).toBe(true);
-    expect(repoSnapshotEligible(enrolled, route)).toBe(true);
-  });
 });
 
 describe('repoSnapshotEligibility', () => {
@@ -113,4 +104,36 @@ describe('computeRepoKey', () => {
     expect(otherUser).not.toBe(base);
     expect(otherRepo).not.toBe(base);
   });
+
+  it('folds the setup commands into the key and keeps the no-setup key unchanged', async () => {
+    const noSetup = await computeRepoKey({ secret: 'secret', userId: USER, ...route });
+    const empty = await computeRepoKey({
+      secret: 'secret',
+      userId: USER,
+      ...route,
+      setupCommands: [],
+    });
+    const install = await computeRepoKey({
+      secret: 'secret',
+      userId: USER,
+      ...route,
+      setupCommands: ['pnpm install'],
+    });
+    const other = await computeRepoKey({
+      secret: 'secret',
+      userId: USER,
+      ...route,
+      setupCommands: ['pnpm build'],
+    });
+    const installAgain = await computeRepoKey({
+      secret: 'secret',
+      userId: USER,
+      ...route,
+      setupCommands: ['pnpm install'],
+    });
+    expect(empty).toBe(noSetup);
+    expect(install).not.toBe(noSetup);
+    expect(other).not.toBe(install);
+    expect(installAgain).toBe(install);
+  });
 });

services/cloud-agent-next/src/control-plane/sandbox/repo-key.ts

@@ -10,6 +10,8 @@ export type RepoKeyInput = {
   userId: string;
   repoUrl: string | undefined;
   directory: string;
+  /** Setup commands the snapshot was captured after; part of the snapshot identity. */
+  setupCommands?: string[];
 };
 
 export type RepoSnapshotGate = {
@@ -22,8 +24,7 @@ export type RepoSnapshotGate = {
 export type RepoSnapshotIneligibilityReason =
   | 'not_enrolled'
   | 'no_repo_url'
-  | 'directory_not_isolated'
-  | 'has_setup_commands';
+  | 'directory_not_isolated';
 
 export type RepoSnapshotEligibility =
   | { eligible: true }
@@ -38,20 +39,15 @@ export type RepoSnapshotEligibility =
  * enrolled, and the route must have a repository at the constant path the
  * snapshots are taken at. A snapshot holds one repository at one path, so a
  * per-session path or a shared sandbox could only capture something no later
- * session can reuse. A route with any setup command is ineligible: capture runs
- * after setup, so the snapshot would carry setup output under a key that names
- * only the user and the repository.
+ * session can reuse.
  *
  * Enrollment follows the `SANDBOX_SELECTION_IDS` allowlist convention:
  * one list matched against the user ID and then the org ID.
  */
 export function repoSnapshotEligibility(
   gate: RepoSnapshotGate,
-  route: Pick<RepoKeyInput, 'repoUrl' | 'directory'> & { setupCommands?: string[] }
+  route: Pick<RepoKeyInput, 'repoUrl' | 'directory'>
 ): RepoSnapshotEligibility {
-  if ((route.setupCommands?.length ?? 0) > 0) {
-    return { eligible: false, reason: 'has_setup_commands' };
-  }
   if (!(isOrgInList(gate.enrolledIds, gate.userId) || isOrgInList(gate.enrolledIds, gate.orgId))) {
     return { eligible: false, reason: 'not_enrolled' };
   }
@@ -64,22 +60,28 @@ export function repoSnapshotEligibility(
 
 export function repoSnapshotEligible(
   gate: RepoSnapshotGate,
-  route: Pick<RepoKeyInput, 'repoUrl' | 'directory'> & { setupCommands?: string[] }
+  route: Pick<RepoKeyInput, 'repoUrl' | 'directory'>
 ): boolean {
   return repoSnapshotEligibility(gate, route).eligible;
 }
 
 /**
- * The keyed hash that names a repository snapshot: scope and repository. Env is
- * not part of it on purpose: setup re-runs on every start and rewrites what it
- * derives from the env. Setup-command output is kept out by not computing a key
- * for a route that has setup commands, rather than by narrowing the hash. It is
- * hashed under a Worker secret so the digest cannot be guessed to address
- * another user's snapshot.
+ * The keyed hash that names a repository snapshot: scope, repository and the setup
+ * commands that produced it. Env is not part of it: setup re-runs on every start
+ * and rewrites what it derives from the env. Because capture runs after setup, the
+ * setup commands join the hash so a snapshot is reused only by a route whose setup
+ * matches; a different setup for the same repository gets its own key instead of
+ * adopting output it did not produce. The digest is hashed under a Worker secret so
+ * it cannot be guessed to address another user's snapshot.
  */
 export async function computeRepoKey(input: RepoKeyInput): Promise<string | null> {
   if (input.secret === null || input.repoUrl === undefined) return null;
-  const message = JSON.stringify([REPO_KEY_DOMAIN, input.userId, input.repoUrl]);
+  const setup = await setupCommandsDigest(input.setupCommands);
+  const message = JSON.stringify(
+    setup === null
+      ? [REPO_KEY_DOMAIN, input.userId, input.repoUrl]
+      : [REPO_KEY_DOMAIN, input.userId, input.repoUrl, setup]
+  );
   const encoder = new TextEncoder();
   const key = await crypto.subtle.importKey(
     'raw',
@@ -91,3 +93,16 @@ export async function computeRepoKey(input: RepoKeyInput): Promise<string | null
   const digest = new Uint8Array(await crypto.subtle.sign('HMAC', key, encoder.encode(message)));
   return Array.from(digest, byte => byte.toString(16).padStart(2, '0')).join('');
 }
+
+/**
+ * A digest of the setup commands, or null when there are none. Absent and empty
+ * are the same, so a route without setup keeps the key computed before setup
+ * joined the identity. The commands are hashed rather than embedded so a secret
+ * bearing command never enters the HMAC message.
+ */
+async function setupCommandsDigest(setupCommands: string[] | undefined): Promise<string | null> {
+  if (setupCommands === undefined || setupCommands.length === 0) return null;
+  const bytes = new TextEncoder().encode(JSON.stringify(setupCommands));
+  const digest = new Uint8Array(await crypto.subtle.digest('SHA-256', bytes));
+  return Array.from(digest, byte => byte.toString(16).padStart(2, '0')).join('');
+}

services/cloud-agent-next/src/control-plane/sandbox/sandbox-do.ts

@@ -1819,8 +1819,9 @@ export class SandboxControlV2 extends DurableObject<Env> {
   /**
    * The repository snapshot key for a route, or null when none applies. It is hashed
    * from the input spec: the grant rewrites `spec.env` with per-attempt credential
-   * aliases, so the stored spec would give every session its own key. A route with
-   * setup commands has no key, so it neither restores nor captures a snapshot.
+   * aliases, so the stored spec would give every session its own key. The setup
+   * commands are folded into the key, so a snapshot is reused only by a route with
+   * the same setup; a different setup neither restores nor captures under it.
    *
    * Emits one `snapshot_key` diagnostic per route so a missing snapshot can be
    * traced to the gate that rejected it (or to a missing key secret).
@@ -1839,10 +1840,7 @@ export class SandboxControlV2 extends DurableObject<Env> {
       userId: source.userId,
       orgId: source.orgId,
     };
-    const eligibility = repoSnapshotEligibility(gate, {
-      ...route,
-      setupCommands: spec.setupCommands,
-    });
+    const eligibility = repoSnapshotEligibility(gate, route);
     if (!eligibility.eligible) {
       this.logSnapshotKey(spec.sessionId, false, eligibility.reason);
       return null;
@@ -1852,7 +1850,12 @@ export class SandboxControlV2 extends DurableObject<Env> {
       1_000,
       'Repository snapshot key secret lookup timed out'
     ).catch(() => null);
-    const key = await computeRepoKey({ secret, userId: source.userId, ...route });
+    const key = await computeRepoKey({
+      secret,
+      userId: source.userId,
+      ...route,
+      setupCommands: spec.setupCommands,
+    });
     if (key === null) {
       this.logSnapshotKey(spec.sessionId, false, 'secret_missing');
       return null;

services/cloud-agent-next/test/integration/sandbox-control-v2-repo-snapshots.test.ts

@@ -612,18 +612,36 @@ describe('repository snapshot launch', () => {
     expect(frame.spec.capture).toBeUndefined();
   });
 
-  it('launches without a key and never asks for a capture when setup commands are present', async () => {
-    const provider = createFakeProvider();
+  it('uses a key and asks the wrapper to capture when setup commands are present', async () => {
+    const provider = createFakeProvider({ startSource: 'repository' });
     const stub = await setup('sbx__repo_launch_setup', provider);
     await stub.prepare(prepareInput({ setupCommands: ['pnpm install'] }));
-
-    expect(await repoKeyOf(stub)).toBeNull();
+    const key = await repoKeyOf(stub);
+    expect(key).not.toBeNull();
 
     const { wrapper } = await connect('sbx__repo_launch_setup', provider);
 
-    expect(provider.launchOptions).toEqual([{}]);
+    expect(provider.launchOptions).toEqual([{ repoKey: key }]);
     const frame = await prepareFrameOf(wrapper);
-    expect(frame.spec.capture).toBeUndefined();
+    expect(frame.spec.capture).toBe(true);
+  });
+
+  it('gives different keys to different setup commands for one repository', async () => {
+    const setupA = 'workspace_aaaa1111-1111-1111-1111-111111111111';
+    const setupB = 'workspace_bbbb2222-2222-2222-2222-222222222222';
+    const providerA = createFakeProvider();
+    const stubA = await setup('sbx__repo_setup_a', providerA);
+    await stubA.prepare(prepareInput({ sessionId: setupA, setupCommands: ['pnpm install'] }));
+
+    const providerB = createFakeProvider();
+    const stubB = await setup('sbx__repo_setup_b', providerB);
+    await stubB.prepare(prepareInput({ sessionId: setupB, setupCommands: ['pnpm build'] }));
+
+    const keyA = await repoKeyOf(stubA, setupA);
+    const keyB = await repoKeyOf(stubB, setupB);
+    expect(keyA).not.toBeNull();
+    expect(keyB).not.toBeNull();
+    expect(keyA).not.toBe(keyB);
   });
 
   it('skips a snapshot whose start never connected and starts from the image', async () => {