fix(code-review): pin sandbox retries

packages/worker-utils/src/cloud-agent-next-client.ts

@@ -33,6 +33,7 @@ export type CloudAgentPrepareSessionInput = {
   upstreamBranch?: string;
   callbackTarget?: CallbackTarget;
   createdOnPlatform?: string;
+  sandboxRetryOfCloudAgentSessionId?: string;
   gateThreshold?: 'off' | 'all' | 'warning' | 'critical';
   runtimeSkills?: Array<{
     name: string;
@@ -44,6 +45,7 @@ export type CloudAgentPrepareSessionInput = {
 export type CloudAgentPrepareSessionOutput = {
   cloudAgentSessionId: string;
   kiloSessionId: string;
+  sandboxRetryPrepared?: true;
 };
 
 export type CloudAgentInitiateInput = {
@@ -294,7 +296,26 @@ export function createCloudAgentNextFetchClient(baseUrl: string): CloudAgentNext
           `Unexpected prepareSession response shape: ${JSON.stringify(data).slice(0, 500)}`
         );
       }
-      return data as unknown as CloudAgentPrepareSessionOutput;
+      if (data.sandboxRetryPrepared !== undefined && data.sandboxRetryPrepared !== true) {
+        throw new Error(
+          `Unexpected prepareSession response shape: ${JSON.stringify(data).slice(0, 500)}`
+        );
+      }
+      if (
+        input.sandboxRetryOfCloudAgentSessionId !== undefined &&
+        data.sandboxRetryPrepared !== true
+      ) {
+        throw new Error('prepareSession did not acknowledge sandbox retry preparation');
+      }
+
+      const output: CloudAgentPrepareSessionOutput = {
+        cloudAgentSessionId: data.cloudAgentSessionId,
+        kiloSessionId: data.kiloSessionId,
+      };
+      if (data.sandboxRetryPrepared === true) {
+        output.sandboxRetryPrepared = true;
+      }
+      return output;
     },
 
     async initiateFromPreparedSession(headers, input) {

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

@@ -617,10 +617,13 @@ export class CloudAgentSession extends DurableObject<WorkerEnv> {
             return { status: 'inspection-failed', error: 'Session metadata unavailable' };
           return createAgentSandbox(this.env, metadata).stopWrappers(request);
         },
-        recordSharedSandboxFailover: routeKey =>
-          this.sharedSandboxFailoverRecorder
-            ? this.sharedSandboxFailoverRecorder(routeKey)
-            : recordSharedSandboxFailover(this.env.SHARED_SANDBOX_OVERRIDES, routeKey),
+        recordSharedSandboxFailover: async routeKey => {
+          if (this.sharedSandboxFailoverRecorder) {
+            await this.sharedSandboxFailoverRecorder(routeKey);
+            return;
+          }
+          await recordSharedSandboxFailover(this.env.SHARED_SANDBOX_OVERRIDES, routeKey);
+        },
         requestAlarmAtOrBefore: deadline => this.scheduleAlarmAtOrBefore(deadline),
         getSessionIdForLogs: () => this.sessionId,
       });

services/cloud-agent-next/src/router/handlers/session-prepare.ts

@@ -360,18 +360,24 @@ const prepareSessionHandler = internalApiProtectedProcedure
               userId: ctx.userId,
               authToken: ctx.authToken,
               botId: ctx.botId,
+              sandboxRetryOfCloudAgentSessionId: input.sandboxRetryOfCloudAgentSessionId,
             })
           : await registerNewSession(requestWithProfile, {
               env: ctx.env,
               userId: ctx.userId,
               authToken: ctx.authToken,
               botId: ctx.botId,
+              sandboxRetryOfCloudAgentSessionId: input.sandboxRetryOfCloudAgentSessionId,
             });
 
-      return {
+      const output = {
         cloudAgentSessionId: result.cloudAgentSessionId,
         kiloSessionId: result.kiloSessionId,
       };
+      if (result.sandboxRetryPrepared === true) {
+        return { ...output, sandboxRetryPrepared: true };
+      }
+      return output;
     });
   });
 

services/cloud-agent-next/src/router/schemas.ts

@@ -441,6 +441,9 @@ export const PrepareSessionInput = z
     callbackTarget: CallbackTargetSchema.optional().describe(
       'Optional callback target configuration for execution completion notifications'
     ),
+    sandboxRetryOfCloudAgentSessionId: sessionIdSchema
+      .optional()
+      .describe('Failed Cloud Agent session that this fresh retry must not reuse'),
 
     // Organization context
     kilocodeOrganizationId: z
@@ -546,6 +549,7 @@ export const PrepareSessionInput = z
 export const PrepareSessionOutput = z.object({
   cloudAgentSessionId: z.string().describe('The generated cloud-agent session ID'),
   kiloSessionId: z.string().describe('The Kilo CLI session ID'),
+  sandboxRetryPrepared: z.literal(true).optional(),
 });
 
 /**

services/cloud-agent-next/src/session-prepare.test.ts

@@ -115,6 +115,7 @@ vi.mock('./session-service.js', () => ({
 
 import { appRouter } from './router.js';
 import { profileResolutionPolicyForSessionCreateOrigin } from './router/handlers/session-prepare.js';
+import { fetchSessionMetadata } from './session-service.js';
 import type { TRPCContext, SessionId } from './types.js';
 
 function createMockDOStub(
@@ -265,6 +266,7 @@ describe('prepareSession endpoint', () => {
     recordInternalCompensationMock.mockResolvedValue(undefined);
     mergeProfileConfigurationMock.mockResolvedValue({});
     assertKiloModelAvailableMock.mockResolvedValue(undefined);
+    vi.mocked(fetchSessionMetadata).mockResolvedValue(null);
   });
 
   it('rejects request without internal API key header', async () => {
@@ -704,6 +706,125 @@ describe('prepareSession endpoint', () => {
     );
   });
 
+  it('prepares a shared sandbox retry on the failover slot', async () => {
+    const sourceSessionId = 'agent_00000000-0000-4000-8000-000000000001';
+    const routeKey = 'usr-000000000000000000000000000000000000000000000000';
+    const failoverSandboxId = 'usr-b4593afcaf2e9e1dfb1611150b786cfe8aeba3c77352a3df';
+    generateSandboxRoutingTargetMock.mockResolvedValueOnce({
+      kind: 'shared',
+      routeKey,
+    });
+    vi.mocked(fetchSessionMetadata).mockResolvedValueOnce({
+      metadataSchemaVersion: 2,
+      identity: { sessionId: sourceSessionId, userId: 'test-user-123' },
+      auth: {},
+      workspace: {
+        sandboxId: routeKey,
+        sandboxRoute: { kind: 'shared', routeKey },
+      },
+      lifecycle: { version: 1, timestamp: Date.now() },
+    });
+    const overrideStore = {
+      get: vi.fn().mockResolvedValue(null),
+      put: vi.fn().mockResolvedValue(undefined),
+    };
+    const doStub = createMockDOStub();
+    const context = createInternalApiContext({ doStub });
+    Object.assign(context.env, { SHARED_SANDBOX_OVERRIDES: overrideStore });
+    const caller = appRouter.createCaller(context);
+
+    const result = await caller.prepareSession({
+      prompt: 'Retry on a fresh sandbox slot',
+      mode: 'code',
+      model: 'claude-3',
+      githubRepo: 'acme/repo',
+      sandboxRetryOfCloudAgentSessionId: sourceSessionId,
+    });
+
+    expect(result).toEqual({
+      cloudAgentSessionId: 'agent_12345678-1234-1234-1234-123456789abc',
+      kiloSessionId: 'cli-session-abc123',
+      sandboxRetryPrepared: true,
+    });
+    expect(fetchSessionMetadata).toHaveBeenCalledWith(
+      context.env,
+      'test-user-123',
+      sourceSessionId
+    );
+    expect(overrideStore.get).toHaveBeenCalledWith(`shared-sandbox-route:${routeKey}`);
+    expect(overrideStore.put).toHaveBeenCalledWith(
+      `shared-sandbox-route:${routeKey}`,
+      'shared-slot-v1'
+    );
+    expect(doStub.registerSession).toHaveBeenCalledWith(
+      expect.objectContaining({
+        workspace: {
+          sandboxId: failoverSandboxId,
+          shallow: false,
+          sandboxRoute: {
+            kind: 'shared',
+            routeKey,
+            suffix: 'shared-slot-v1',
+          },
+        },
+      })
+    );
+    expect(recordSandboxIdentityMock).toHaveBeenCalledWith(
+      expect.objectContaining({ sandboxId: failoverSandboxId }),
+      expect.any(Object)
+    );
+  });
+
+  it('rejects a shared sandbox retry when the source already uses the failover slot', async () => {
+    const sourceSessionId = 'agent_00000000-0000-4000-8000-000000000001';
+    const routeKey = 'usr-000000000000000000000000000000000000000000000000';
+    const failoverSandboxId = 'usr-b4593afcaf2e9e1dfb1611150b786cfe8aeba3c77352a3df';
+    generateSandboxRoutingTargetMock.mockResolvedValueOnce({
+      kind: 'shared',
+      routeKey,
+    });
+    vi.mocked(fetchSessionMetadata).mockResolvedValueOnce({
+      metadataSchemaVersion: 2,
+      identity: { sessionId: sourceSessionId, userId: 'test-user-123' },
+      auth: {},
+      workspace: {
+        sandboxId: failoverSandboxId,
+        sandboxRoute: { kind: 'shared', routeKey, suffix: 'shared-slot-v1' },
+      },
+      lifecycle: { version: 1, timestamp: Date.now() },
+    });
+    const overrideStore = {
+      get: vi.fn().mockResolvedValue('shared-slot-v1'),
+      put: vi.fn().mockResolvedValue(undefined),
+    };
+    const doStub = createMockDOStub();
+    const context = createInternalApiContext({ doStub });
+    Object.assign(context.env, { SHARED_SANDBOX_OVERRIDES: overrideStore });
+    const caller = appRouter.createCaller(context);
+
+    await expect(
+      caller.prepareSession({
+        prompt: 'Do not retry from a failed failover slot',
+        mode: 'code',
+        model: 'claude-3',
+        githubRepo: 'acme/repo',
+        sandboxRetryOfCloudAgentSessionId: sourceSessionId,
+      })
+    ).rejects.toMatchObject({ code: 'BAD_REQUEST' });
+
+    expect(overrideStore.get).not.toHaveBeenCalled();
+    expect(overrideStore.put).not.toHaveBeenCalled();
+    expect(recordSandboxIdentityMock).not.toHaveBeenCalled();
+    expect(createCliSessionMock).not.toHaveBeenCalled();
+    expect(doStub.registerSession).not.toHaveBeenCalled();
+    expect(recordSessionFailureMock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        failure: { stage: 'sandbox_identity', code: 'sandbox_id_derivation_failed' },
+      }),
+      expect.any(Object)
+    );
+  });
+
   it('creates auto-initiated devcontainer sessions with grouped DIND sandbox intent', async () => {
     generateSandboxRoutingTargetMock.mockResolvedValueOnce({
       kind: 'isolated',