fix(cloud-agent-next): decide worktree deletion exclusivity from a local reference index

services/cloud-agent-next/src/persistence/SandboxControl.ts

@@ -3,7 +3,19 @@ import {
   cloudAgentWorktreeIdSchema,
   WORKTREE_RUNTIME_HISTORY_UNAVAILABLE,
 } from '@kilocode/session-ingest-contracts';
-import { resolveSandboxExclusivity } from '../sandbox-control/worktree-ownership.js';
+import {
+  RECONCILIATION_LIMITS,
+  reconcileSandboxReferences,
+} from '../sandbox-control/worktree-ownership.js';
+import {
+  addSessionReference,
+  hasForeignReference,
+  markReferencesReconciled,
+  removeSessionReference,
+  removeWorktreeReferences,
+  worktreeIdFromDirectory,
+  type SessionReferenceState,
+} from '../sandbox-control/session-references.js';
 import { getWorktreeWorkspacePath } from '../workspace.js';
 import {
   cleanWorktreeRuntime,
@@ -157,6 +169,8 @@ import {
   saveDeadlines,
   savePhysicalRecord,
   saveRouteTable,
+  loadSessionReferences,
+  saveSessionReferences,
   saveTransitionLog,
   loadSessionCredentialGrants,
   saveSessionCredentialGrants,
@@ -781,7 +795,7 @@ export class SandboxControl extends DurableObject<Env> {
     ) {
       return null;
     }
-    const worktreeId = route?.worktreeId ?? this.worktreeIdFromDirectory(input.directory);
+    const worktreeId = route?.worktreeId ?? worktreeIdFromDirectory(input.directory);
     if (
       this.runtimeDeleted ||
       this.exclusiveDeletionWorktreeId ||
@@ -1240,7 +1254,7 @@ export class SandboxControl extends DurableObject<Env> {
       )
         throw new Error('Session interaction scope is stale');
       this.assertWorktreeAdmission(route.worktreeId);
-      this.assertWorktreeAdmission(this.worktreeIdFromDirectory(session.directory));
+      this.assertWorktreeAdmission(worktreeIdFromDirectory(session.directory));
     });
   }
 
@@ -2167,10 +2181,20 @@ export class SandboxControl extends DurableObject<Env> {
       ) {
         throw new Error('Sandbox credential containment mismatch');
       }
-      const result = await this.mutateRoutes(table => {
+      const result = await this.mutateRoutesAndReferences((table, references) => {
         this.assertWorktreeAdmission(worktreeId);
         const attached = attachRoute(table, input, ownerId);
-        return { value: attached, changed: attached.changed };
+        const added = addSessionReference(references, {
+          sessionId: input.sessionId,
+          kiloSessionId: input.kiloSessionId,
+          directory: input.directory,
+          ...(worktreeId !== undefined ? { worktreeId } : {}),
+        });
+        return {
+          value: attached,
+          routesChanged: attached.changed,
+          referencesChanged: added.changed,
+        };
       });
       if (result.changed) {
         await this.appendLog(
@@ -2288,10 +2312,21 @@ export class SandboxControl extends DurableObject<Env> {
       const result = await this.withCredentialUpdate(() =>
         this.ctx.storage.transaction(async () => {
           const table = await loadRouteTable(this.ctx.storage);
+          const removing = runtimeDetached ? table.get(sessionId) : undefined;
           const detached = runtimeDetached
             ? detachRoute(table, sessionId)
             : { table, existed: false };
           await saveRouteTable(this.ctx.storage, detached.table);
+          if (detached.existed && removing) {
+            const references = await loadSessionReferences(this.ctx.storage);
+            const tombstoned = addSessionReference(references, {
+              sessionId: removing.sessionId,
+              kiloSessionId: removing.kiloSessionId,
+              directory: removing.directory,
+              ...(removing.worktreeId !== undefined ? { worktreeId: removing.worktreeId } : {}),
+            });
+            if (tombstoned.changed) await saveSessionReferences(this.ctx.storage, references);
+          }
           const grants = await loadSessionCredentialGrants(this.ctx.storage);
           if (grants.some(grant => grant.members.some(member => member.sessionId === sessionId))) {
             if (this.providerKind === 'vercel') {
@@ -2326,6 +2361,16 @@ export class SandboxControl extends DurableObject<Env> {
     return { existed };
   }
 
+  async forgetSessionReference(sessionId: string): Promise<void> {
+    await this.ensureOperationalInitialized();
+    await this.ctx.storage.transaction(async () => {
+      const references = await loadSessionReferences(this.ctx.storage);
+      if (removeSessionReference(references, sessionId).changed) {
+        await saveSessionReferences(this.ctx.storage, references);
+      }
+    });
+  }
+
   deleteWorktreeResources(
     raw: SandboxWorktreeCleanupInput
   ): Promise<{ deleted: true; sessionIds: string[] }> {
@@ -2401,7 +2446,7 @@ export class SandboxControl extends DurableObject<Env> {
       await this.revokeWorktreeCredentials(input.worktreeId);
     }
     const deletedIds = new Set(journal.sessionIds);
-    const detached = await this.mutateRoutes(table => {
+    const detached = await this.mutateRoutesAndReferences((table, references) => {
       const sessionIds: string[] = [];
       for (const [sessionId, route] of table) {
         if (
@@ -2412,7 +2457,12 @@ export class SandboxControl extends DurableObject<Env> {
           sessionIds.push(sessionId);
         }
       }
-      return { value: sessionIds, changed: sessionIds.length > 0 };
+      const removed = removeWorktreeReferences(references, input.worktreeId);
+      return {
+        value: sessionIds,
+        routesChanged: sessionIds.length > 0,
+        referencesChanged: removed.changed,
+      };
     });
     await Promise.allSettled(detached.flatMap(id => this.sessionForwarding.get(id) ?? []));
     for (const id of detached) this.sessionForwarding.delete(id);
@@ -2505,41 +2555,62 @@ export class SandboxControl extends DurableObject<Env> {
         )
       )
         return true;
-      const receipts = await loadWorktreeDeletionJournals(this.ctx.storage);
-      const releasedWorktreeIds = [...receipts]
-        .filter(([, receipt]) => receipt.resourcesCleaned)
-        .map(([id]) => id);
-      const released = new Set<string>(releasedWorktreeIds);
-      const requestedIds = new Set(input.sessionIds);
-      const otherRoutes = [...(await loadRouteTable(this.ctx.storage)).values()].some(route => {
-        const worktreeId = route.worktreeId ?? this.worktreeIdFromDirectory(route.directory);
-        if (worktreeId && released.has(worktreeId)) return false;
-        return (
-          route.directory !== directory ||
-          !requestedIds.has(route.kiloSessionId) ||
-          (worktreeId !== undefined && worktreeId !== input.worktreeId)
-        );
-      });
-      const exclusive =
-        !otherRoutes &&
-        (await withTimeout(
-          resolveSandboxExclusivity(this.env, {
+      const storage = this.ctx.storage;
+      const receipts = await loadWorktreeDeletionJournals(storage);
+      const released = new Set(
+        [...receipts].filter(([, receipt]) => receipt.resourcesCleaned).map(([id]) => id)
+      );
+      const target = {
+        worktreeId: input.worktreeId,
+        directory,
+        sessionIds: new Set(input.sessionIds),
+        releasedWorktreeIds: released,
+      };
+      const references = await loadSessionReferences(storage);
+      const foreignEvidence = async () => [
+        ...references.entries,
+        ...(await loadRouteTable(storage)).values(),
+      ];
+      if (hasForeignReference(references, await foreignEvidence(), target)) {
+        await this.releaseWorktreeAdmission(input.worktreeId);
+        return false;
+      }
+      if (!references.reconciled) {
+        const result = await reconcileSandboxReferences(
+          this.env,
+          {
             worktreeId: input.worktreeId,
             kiloUserId: input.kiloUserId,
             organizationId: input.organizationId,
             location: input.location,
-            releasedWorktreeIds,
-          }),
-          DEADLINE_MS.stopAttempt,
-          'Worktree ownership lookup timed out'
-        ));
-      if (exclusive) return true;
+            releasedWorktreeIds: [...released],
+          },
+          RECONCILIATION_LIMITS
+        );
+        if (!result.complete || result.foreign || result.unavailable) {
+          await this.releaseWorktreeAdmission(input.worktreeId);
+          return false;
+        }
+        const confirmed = await storage.transaction(async () => {
+          const current = await loadSessionReferences(storage);
+          if (current.reconciled) return true;
+          const routes = await loadRouteTable(storage);
+          if (hasForeignReference(current, [...current.entries, ...routes.values()], target)) {
+            return false;
+          }
+          await saveSessionReferences(storage, markReferencesReconciled(current));
+          return true;
+        });
+        if (!confirmed) {
+          await this.releaseWorktreeAdmission(input.worktreeId);
+          return false;
+        }
+      }
+      return true;
     } catch (error) {
       await this.releaseWorktreeAdmission(input.worktreeId);
       throw error;
     }
-    await this.releaseWorktreeAdmission(input.worktreeId);
-    return false;
   }
 
   private async releaseWorktreeAdmission(worktreeId: string): Promise<void> {
@@ -2549,17 +2620,12 @@ export class SandboxControl extends DurableObject<Env> {
     if (!this.runtimeDeleted) await this.scheduleAlarm(await loadDeadlines(this.ctx.storage));
   }
 
-  private worktreeIdFromDirectory(directory: string): string | undefined {
-    const parsed = cloudAgentWorktreeIdSchema.safeParse(directory.split('/').at(-1));
-    return parsed.success ? parsed.data : undefined;
-  }
-
   private async assertRequestWorktreeAdmission(
     input: SandboxControlOutboundRequest
   ): Promise<void> {
     const session = input.session;
     if (!session) return;
-    const worktreeId = this.worktreeIdFromDirectory(session.directory);
+    const worktreeId = worktreeIdFromDirectory(session.directory);
     if (input.operation === 'session.sync' && this.exclusiveDeletionWorktreeId) {
       const allowed = await this.ctx.storage.transaction(async () => {
         const exclusiveWorktreeId = this.exclusiveDeletionWorktreeId;
@@ -5136,7 +5202,7 @@ export class SandboxControl extends DurableObject<Env> {
     if (!route || route.ownerId !== input.ownerId) {
       return { allowed: false, reason: 'session_not_attached' };
     }
-    const worktreeId = route.worktreeId ?? this.worktreeIdFromDirectory(route.directory);
+    const worktreeId = route.worktreeId ?? worktreeIdFromDirectory(route.directory);
     if (
       this.runtimeDeleted ||
       this.exclusiveDeletionWorktreeId ||
@@ -5654,13 +5720,18 @@ export class SandboxControl extends DurableObject<Env> {
     return notified.every(Boolean);
   }
 
-  private mutateRoutes<T>(
-    mutation: (table: Map<string, SessionRoute>) => { value: T; changed: boolean }
+  private mutateRoutesAndReferences<T>(
+    mutation: (
+      table: Map<string, SessionRoute>,
+      references: SessionReferenceState
+    ) => { value: T; routesChanged: boolean; referencesChanged: boolean }
   ): Promise<T> {
     return this.ctx.storage.transaction(async () => {
       const table = await loadRouteTable(this.ctx.storage);
-      const updated = mutation(table);
-      if (updated.changed) await saveRouteTable(this.ctx.storage, table);
+      const references = await loadSessionReferences(this.ctx.storage);
+      const updated = mutation(table, references);
+      if (updated.routesChanged) await saveRouteTable(this.ctx.storage, table);
+      if (updated.referencesChanged) await saveSessionReferences(this.ctx.storage, references);
       return updated.value;
     });
   }

services/cloud-agent-next/src/sandbox-control/durable-state.test.ts

@@ -9,13 +9,24 @@ import {
   readSandboxControlState,
   loadRouteTable,
   loadSessionCredentialGrants,
+  loadSessionReferences,
   loadTransitionLog,
   saveDeadlines,
   savePhysicalRecord,
   saveRouteTable,
   saveSessionCredentialGrants,
+  saveSessionReferences,
   saveTransitionLog,
+  SESSION_REFERENCES_KEY,
 } from './durable-state.js';
+import {
+  MAX_REFERENCE_BYTES,
+  MAX_REFERENCE_ENTRIES,
+  addSessionReference,
+  emptySessionReferenceState,
+  markReferencesReconciled,
+  serializedReferenceBytes,
+} from './session-references.js';
 import { createControlPlaneCredential } from './managed-credential.js';
 import {
   claimCreate,
@@ -276,6 +287,17 @@ describe('sandbox control durable state', () => {
       grant.userId
     );
     await saveRouteTable(storage, table);
+    await saveSessionReferences(
+      storage,
+      markReferencesReconciled(
+        addSessionReference(emptySessionReferenceState(), {
+          sessionId: grant.scopeId,
+          kiloSessionId: ROOT_ID,
+          directory: grant.directory,
+          worktreeId: grant.scopeId,
+        }).state
+      )
+    );
     await saveDeadlines(storage, { heartbeatExpiry: 3000 });
     await saveTransitionLog(storage, [{ at: 1001, kind: 'physical', to: 'running' }]);
     await storage.put('owner', grant.userId);
@@ -285,8 +307,75 @@ describe('sandbox control durable state', () => {
     expect(await loadSessionCredentialGrants(storage)).toEqual([]);
     expect(await loadPhysicalRecord(storage)).toStrictEqual(initialPhysicalRecord(false));
     expect(await loadRouteTable(storage)).toEqual(emptyRouteTable());
+    expect(await loadSessionReferences(storage)).toEqual(emptySessionReferenceState());
     expect(await loadDeadlines(storage)).toEqual({});
     expect(await loadTransitionLog(storage)).toEqual([]);
     expect(await storage.get('owner')).toBe(grant.userId);
   });
+
+  it('defaults absent session references to the empty state and round-trips a reconciled index', async () => {
+    const storage = memoryStorage();
+    expect(await loadSessionReferences(storage)).toEqual(emptySessionReferenceState());
+    const state = emptySessionReferenceState();
+    addSessionReference(state, {
+      sessionId: SESSION_ID,
+      kiloSessionId: ROOT_ID,
+      directory: '/workspace/paths/org/project/worktree_11111111-1111-4111-8111-111111111111',
+    });
+    addSessionReference(state, {
+      sessionId: 'workspace_other',
+      kiloSessionId: 'ses_other',
+      directory: '/workspace/paths/org/other',
+      worktreeId: 'worktree_22222222-2222-4222-8222-222222222222',
+    });
+    markReferencesReconciled(state);
+
+    await saveSessionReferences(storage, state);
+
+    expect(await loadSessionReferences(storage)).toEqual(state);
+  });
+
+  it('rejects malformed session references instead of defaulting to the empty state', async () => {
+    const storage = memoryStorage();
+    for (const value of [
+      { reconciled: 'yes', overflowed: false, entries: [] },
+      {
+        reconciled: true,
+        overflowed: false,
+        entries: [{ sessionId: SESSION_ID, kiloSessionId: ROOT_ID }],
+      },
+      {
+        reconciled: true,
+        overflowed: false,
+        entries: [{ sessionId: SESSION_ID, kiloSessionId: ROOT_ID, directory: '', extra: 'field' }],
+      },
+    ]) {
+      await storage.put(SESSION_REFERENCES_KEY, value);
+      await expect(loadSessionReferences(storage)).rejects.toThrow();
+    }
+  });
+
+  it('rejects a persisted session reference index over the entry or byte limit', async () => {
+    const storage = memoryStorage();
+    const entry = (index: number, directory: string) => ({
+      sessionId: `ses_${index}`,
+      kiloSessionId: `kilo_${index}`,
+      directory,
+    });
+    await storage.put(SESSION_REFERENCES_KEY, {
+      reconciled: false,
+      overflowed: false,
+      entries: Array.from({ length: MAX_REFERENCE_ENTRIES + 1 }, (_, index) => entry(index, 'd')),
+    });
+    await expect(loadSessionReferences(storage)).rejects.toThrow();
+
+    const oversized = Array.from({ length: 200 }, (_, index) => entry(index, 'd'.repeat(512)));
+    expect(serializedReferenceBytes(oversized)).toBeGreaterThan(MAX_REFERENCE_BYTES);
+    await storage.put(SESSION_REFERENCES_KEY, {
+      reconciled: false,
+      overflowed: false,
+      entries: oversized,
+    });
+    await expect(loadSessionReferences(storage)).rejects.toThrow();
+  });
 });