fix(cloud): prewarm selected sandbox image

Include cloud environment and custom image selections in warm requests and lease matching. Reuse matching configured warm runs while preserving cold fallback when no configured lease is available.

Generated-By: PostHog Code
Task-Id: 6b752fb1-4fa1-45e4-920a-38306e4b7cee

Author: tatoalo

packages/api-client/src/posthog-client.test.ts

@@ -544,6 +544,39 @@ describe("PostHogAPIClient", () => {
       );
     });
 
+    it("forwards the selected sandbox environment and custom image", async () => {
+      const fetch = vi.fn().mockResolvedValue({
+        ok: true,
+        text: async () =>
+          JSON.stringify({ task_id: "task-1", run_id: "run-1" }),
+      });
+      const client = makeClient(fetch);
+
+      await client.warmTask({
+        repository: "PostHog/posthog",
+        github_integration: 42,
+        sandbox_environment_id: "environment-123",
+        custom_image_id: "image-123",
+      });
+
+      expect(fetch).toHaveBeenCalledWith(
+        expect.objectContaining({
+          overrides: {
+            body: JSON.stringify({
+              repository: "PostHog/posthog",
+              github_integration: 42,
+              branch: null,
+              runtime_adapter: null,
+              model: null,
+              reasoning_effort: null,
+              sandbox_environment_id: "environment-123",
+              custom_image_id: "image-123",
+            }),
+          },
+        }),
+      );
+    });
+
     it("sends a null branch when none is provided", async () => {
       const fetch = vi.fn().mockResolvedValue({
         ok: true,

packages/api-client/src/posthog-client.ts

@@ -2817,6 +2817,8 @@ export class PostHogAPIClient {
     runtime_adapter?: string | null;
     model?: string | null;
     reasoning_effort?: string | null;
+    sandbox_environment_id?: string | null;
+    custom_image_id?: string | null;
   }): Promise<{ task_id: string; run_id: string } | null> {
     const teamId = await this.getTeamId();
     const urlPath = `/api/projects/${teamId}/tasks/warm/`;
@@ -2833,6 +2835,12 @@ export class PostHogAPIClient {
           runtime_adapter: options.runtime_adapter ?? null,
           model: options.model ?? null,
           reasoning_effort: options.reasoning_effort ?? null,
+          ...(options.sandbox_environment_id
+            ? { sandbox_environment_id: options.sandbox_environment_id }
+            : {}),
+          ...(options.custom_image_id
+            ? { custom_image_id: options.custom_image_id }
+            : {}),
         }),
       },
     });

packages/core/src/task-detail/taskCreationHost.ts

@@ -125,6 +125,8 @@ export interface ITaskCreationHost {
     runtimeAdapter?: string | null;
     model?: string | null;
     reasoningEffort?: string | null;
+    sandboxEnvironmentId?: string | null;
+    customImageId?: string | null;
   }): { taskId: string; runId: string } | null;
   uploadRunAttachments(
     client: TaskCreationApiClient,

packages/core/src/task-detail/taskCreationSaga.test.ts

@@ -569,6 +569,8 @@ describe("TaskCreationSaga", () => {
       runtimeAdapter: null,
       model: null,
       reasoningEffort: null,
+      sandboxEnvironmentId: null,
+      customImageId: null,
     });
     // The bundle must land on the warm run before createTask triggers activation.
     expect(mockHost.uploadRunAttachments).toHaveBeenCalledWith(
@@ -699,8 +701,9 @@ describe("TaskCreationSaga", () => {
       expectedRunOptions: { customImageId: "image-123" },
     },
   ])(
-    "starts a cold run for a selected $selection",
+    "falls back to a cold run without a matching warm $selection lease",
     async ({ input, expectedRunOptions }) => {
+      mockHost.takeWarmTaskLease.mockReturnValue(null);
       const createdTask = createTask();
       const startedTask = createTask({ latest_run: createRun() });
       const createTaskMock = vi.fn().mockResolvedValue(createdTask);
@@ -729,6 +732,40 @@ describe("TaskCreationSaga", () => {
     },
   );
 
+  it("reuses a warm run built from the selected custom image", async () => {
+    mockHost.takeWarmTaskLease.mockReturnValue({
+      taskId: "warm-task",
+      runId: "warm-run",
+    });
+    const warmActivatedTask = createTask({
+      id: "warm-task",
+      latest_run: createRun({ id: "warm-run", task: "warm-task" }),
+    });
+    const createTaskMock = vi.fn().mockResolvedValue(warmActivatedTask);
+    const createTaskRunMock = vi.fn();
+    const saga = makeSaga({
+      createTask: createTaskMock,
+      createTaskRun: createTaskRunMock,
+    });
+
+    const result = await saga.run({
+      content: "Ship the fix",
+      repository: "posthog/posthog",
+      workspaceMode: "cloud",
+      branch: "main",
+      customImageId: "image-123",
+    });
+
+    expect(result.success).toBe(true);
+    expect(createTaskMock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        branch: "main",
+        custom_image_id: "image-123",
+      }),
+    );
+    expect(createTaskRunMock).not.toHaveBeenCalled();
+  });
+
   it("uses the selected user GitHub integration for cloud task creation", async () => {
     const createdTask = createTask({
       github_user_integration: "user-integration-123",

packages/core/src/task-detail/taskCreationSaga.ts

@@ -685,24 +685,29 @@ export class TaskCreationSaga extends Saga<
       augmented,
     };
 
-    if (input.sandboxEnvironmentId || input.customImageId) {
-      return { ...base, suppressWarmReuse: true };
-    }
-
     const lease = input.repository
       ? this.deps.host.takeWarmTaskLease({
           repository: input.repository,
           branch: input.branch ?? null,
           runtimeAdapter: input.adapter ?? null,
           model: input.model ?? null,
           reasoningEffort: input.reasoningLevel ?? null,
+          sandboxEnvironmentId: input.sandboxEnvironmentId ?? null,
+          customImageId: input.customImageId ?? null,
         })
       : null;
 
+    const requiresConfiguredWarm = Boolean(
+      input.sandboxEnvironmentId || input.customImageId,
+    );
+
     const needsAttachments =
       transport.filePaths.length > 0 || transport.skillBundles.length > 0;
     if (!needsAttachments) {
-      return base;
+      return {
+        ...base,
+        suppressWarmReuse: requiresConfiguredWarm && !lease,
+      };
     }
     if (!lease) {
       return { ...base, suppressWarmReuse: true };
@@ -792,6 +797,14 @@ export class TaskCreationSaga extends Saga<
             input.workspaceMode === "cloud"
               ? (input.reasoningLevel ?? null)
               : undefined,
+          sandbox_environment_id:
+            input.workspaceMode === "cloud" && !warmPayload?.suppressWarmReuse
+              ? input.sandboxEnvironmentId
+              : undefined,
+          custom_image_id:
+            input.workspaceMode === "cloud" && !warmPayload?.suppressWarmReuse
+              ? input.customImageId
+              : undefined,
           signal_report: input.signalReportId ?? undefined,
           channel: input.channelId ?? undefined,
           pending_user_message: warmPayload?.pendingUserMessage,

packages/ui/src/features/task-detail/components/TaskInput.tsx

@@ -730,6 +730,8 @@ export function TaskInput({
     runtimeAdapter: adapter ?? null,
     model: effectiveModel,
     reasoningEffort: effectiveReasoningLevel,
+    sandboxEnvironmentId: workspaceMode === "cloud" ? selectedCloudEnvId : null,
+    customImageId: workspaceMode === "cloud" ? selectedCustomImageId : null,
   });
 
   const branchForTaskCreation =

packages/ui/src/features/task-detail/hooks/useWarmTask.test.tsx

@@ -28,6 +28,8 @@ interface Props {
   runtimeAdapter?: string | null;
   model?: string | null;
   reasoningEffort?: string | null;
+  sandboxEnvironmentId?: string | null;
+  customImageId?: string | null;
 }
 
 const cloudTyping: Props = {
@@ -201,6 +203,41 @@ describe("useWarmTask", () => {
     expect(mockClient.warmTask).toHaveBeenCalledTimes(2);
   });
 
+  it("forwards sandbox configuration and re-warms when the image changes", async () => {
+    const { rerender } = renderHook((props: Props) => useWarmTask(props), {
+      initialProps: {
+        ...cloudTyping,
+        sandboxEnvironmentId: "environment-123",
+        customImageId: "image-123",
+      },
+    });
+    await flushDebounce();
+    expect(mockClient.warmTask).toHaveBeenLastCalledWith({
+      repository: "acme/repo",
+      github_integration: 42,
+      branch: "main",
+      ...NULL_RUNTIME,
+      sandbox_environment_id: "environment-123",
+      custom_image_id: "image-123",
+    });
+
+    rerender({
+      ...cloudTyping,
+      sandboxEnvironmentId: "environment-123",
+      customImageId: "image-456",
+    });
+    await flushDebounce();
+    expect(mockClient.warmTask).toHaveBeenLastCalledWith({
+      repository: "acme/repo",
+      github_integration: 42,
+      branch: "main",
+      ...NULL_RUNTIME,
+      sandbox_environment_id: "environment-123",
+      custom_image_id: "image-456",
+    });
+    expect(mockClient.warmTask).toHaveBeenCalledTimes(2);
+  });
+
   it("warms again for a new selection after a failed warm", async () => {
     mockClient.warmTask.mockRejectedValueOnce(new Error("boom"));
     const { rerender } = renderHook((props: Props) => useWarmTask(props), {

packages/ui/src/features/task-detail/hooks/useWarmTask.ts

@@ -21,6 +21,8 @@ interface UseWarmTaskOptions {
   runtimeAdapter?: string | null;
   model?: string | null;
   reasoningEffort?: string | null;
+  sandboxEnvironmentId?: string | null;
+  customImageId?: string | null;
 }
 
 export function useWarmTask({
@@ -32,6 +34,8 @@ export function useWarmTask({
   runtimeAdapter,
   model,
   reasoningEffort,
+  sandboxEnvironmentId,
+  customImageId,
 }: UseWarmTaskOptions): void {
   const enabled = useFeatureFlag(TASKS_PREWARM_SANDBOX_FLAG);
   const client = useOptionalAuthenticatedClient();
@@ -44,6 +48,8 @@ export function useWarmTask({
   const normalizedRuntimeAdapter = runtimeAdapter ?? null;
   const normalizedModel = model ?? null;
   const normalizedReasoningEffort = reasoningEffort ?? null;
+  const normalizedSandboxEnvironmentId = sandboxEnvironmentId ?? null;
+  const normalizedCustomImageId = customImageId ?? null;
   const eligible =
     enabled &&
     isCloud &&
@@ -59,6 +65,8 @@ export function useWarmTask({
           runtimeAdapter: normalizedRuntimeAdapter,
           model: normalizedModel,
           reasoningEffort: normalizedReasoningEffort,
+          sandboxEnvironmentId: normalizedSandboxEnvironmentId,
+          customImageId: normalizedCustomImageId,
         })}`
       : null;
 
@@ -84,6 +92,8 @@ export function useWarmTask({
     const warmRuntimeAdapter = normalizedRuntimeAdapter;
     const warmModel = normalizedModel;
     const warmReasoningEffort = normalizedReasoningEffort;
+    const warmSandboxEnvironmentId = normalizedSandboxEnvironmentId;
+    const warmCustomImageId = normalizedCustomImageId;
     debounceRef.current = setTimeout(() => {
       debounceRef.current = null;
       lastWarmedKeyRef.current = key;
@@ -95,6 +105,10 @@ export function useWarmTask({
           runtime_adapter: warmRuntimeAdapter,
           model: warmModel,
           reasoning_effort: warmReasoningEffort,
+          ...(warmSandboxEnvironmentId
+            ? { sandbox_environment_id: warmSandboxEnvironmentId }
+            : {}),
+          ...(warmCustomImageId ? { custom_image_id: warmCustomImageId } : {}),
         })
         .then((warm) => {
           if (warm) {
@@ -105,6 +119,8 @@ export function useWarmTask({
                 runtimeAdapter: warmRuntimeAdapter,
                 model: warmModel,
                 reasoningEffort: warmReasoningEffort,
+                sandboxEnvironmentId: warmSandboxEnvironmentId,
+                customImageId: warmCustomImageId,
               }),
               { taskId: warm.task_id, runId: warm.run_id },
             );
@@ -127,5 +143,7 @@ export function useWarmTask({
     normalizedRuntimeAdapter,
     normalizedModel,
     normalizedReasoningEffort,
+    normalizedSandboxEnvironmentId,
+    normalizedCustomImageId,
   ]);
 }

packages/ui/src/features/task-detail/hooks/warmTaskLease.ts

@@ -9,6 +9,8 @@ export interface WarmTaskLeaseKeyParts {
   runtimeAdapter?: string | null;
   model?: string | null;
   reasoningEffort?: string | null;
+  sandboxEnvironmentId?: string | null;
+  customImageId?: string | null;
 }
 
 export function buildWarmTaskLeaseKey(parts: WarmTaskLeaseKeyParts): string {
@@ -18,6 +20,8 @@ export function buildWarmTaskLeaseKey(parts: WarmTaskLeaseKeyParts): string {
     parts.runtimeAdapter ?? "",
     parts.model ?? "",
     parts.reasoningEffort ?? "",
+    parts.sandboxEnvironmentId ?? "",
+    parts.customImageId ?? "",
   ].join(":");
 }
 

packages/ui/src/features/task-detail/taskCreationHostImpl.ts

@@ -162,6 +162,8 @@ export class TrpcTaskCreationHost implements ITaskCreationHost {
     runtimeAdapter?: string | null;
     model?: string | null;
     reasoningEffort?: string | null;
+    sandboxEnvironmentId?: string | null;
+    customImageId?: string | null;
   }): { taskId: string; runId: string } | null {
     return takeWarmTaskLease(args);
   }