fix(cloud): prewarm selected sandbox image

Include cloud environment and custom image selections in warm requests and lease matching. Reuse only matching configured warm runs, preserve cold fallback, and ignore stale warm responses after the selection changes.

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(
@@ -687,6 +689,83 @@ describe("TaskCreationSaga", () => {
     });
   });
 
+  it.each([
+    {
+      selection: "sandbox environment",
+      input: { sandboxEnvironmentId: "environment-123" },
+      expectedRunOptions: { sandboxEnvironmentId: "environment-123" },
+    },
+    {
+      selection: "custom image",
+      input: { customImageId: "image-123" },
+      expectedRunOptions: { customImageId: "image-123" },
+    },
+  ])(
+    "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);
+      const createTaskRunMock = vi.fn().mockResolvedValue(createRun());
+      const startTaskRunMock = vi.fn().mockResolvedValue(startedTask);
+      const saga = makeSaga({
+        createTask: createTaskMock,
+        createTaskRun: createTaskRunMock,
+        startTaskRun: startTaskRunMock,
+      });
+
+      const result = await saga.run({
+        content: "Ship the fix",
+        repository: "posthog/posthog",
+        workspaceMode: "cloud",
+        branch: "main",
+        ...input,
+      });
+
+      expect(result.success).toBe(true);
+      expect(createTaskMock.mock.calls[0][0].branch).toBeUndefined();
+      expect(createTaskRunMock).toHaveBeenCalledWith(
+        "task-123",
+        expect.objectContaining(expectedRunOptions),
+      );
+    },
+  );
+
+  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

@@ -692,13 +692,22 @@ export class TaskCreationSaga extends Saga<
           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 };
@@ -788,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

@@ -18,6 +18,7 @@ vi.mock("../../../shell/logger", () => ({
 }));
 
 import { useWarmTask } from "./useWarmTask";
+import { takeWarmTaskLease } from "./warmTaskLease";
 
 interface Props {
   workspaceMode: WorkspaceMode;
@@ -28,6 +29,8 @@ interface Props {
   runtimeAdapter?: string | null;
   model?: string | null;
   reasoningEffort?: string | null;
+  sandboxEnvironmentId?: string | null;
+  customImageId?: string | null;
 }
 
 const cloudTyping: Props = {
@@ -201,6 +204,103 @@ 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 only the latest image when selection changes during the debounce", async () => {
+    const { rerender } = renderHook((props: Props) => useWarmTask(props), {
+      initialProps: { ...cloudTyping, customImageId: "image-123" },
+    });
+
+    rerender({ ...cloudTyping, customImageId: "image-456" });
+    await flushDebounce();
+
+    expect(mockClient.warmTask).toHaveBeenCalledOnce();
+    expect(mockClient.warmTask).toHaveBeenCalledWith({
+      repository: "acme/repo",
+      github_integration: 42,
+      branch: "main",
+      ...NULL_RUNTIME,
+      custom_image_id: "image-456",
+    });
+  });
+
+  it("keeps the latest image lease when warm responses complete out of order", async () => {
+    type WarmResponse = { task_id: string; run_id: string };
+    let resolveFirstWarm!: (value: WarmResponse) => void;
+    let resolveSecondWarm!: (value: WarmResponse) => void;
+    const firstWarm = new Promise<WarmResponse>((resolve) => {
+      resolveFirstWarm = resolve;
+    });
+    const secondWarm = new Promise<WarmResponse>((resolve) => {
+      resolveSecondWarm = resolve;
+    });
+    mockClient.warmTask
+      .mockReturnValueOnce(firstWarm)
+      .mockReturnValueOnce(secondWarm);
+
+    const { rerender } = renderHook((props: Props) => useWarmTask(props), {
+      initialProps: { ...cloudTyping, customImageId: "image-123" },
+    });
+    await flushDebounce();
+
+    rerender({ ...cloudTyping, customImageId: "image-456" });
+    await flushDebounce();
+
+    await act(async () => {
+      resolveSecondWarm({ task_id: "task-2", run_id: "run-2" });
+      await secondWarm;
+    });
+    await act(async () => {
+      resolveFirstWarm({ task_id: "task-1", run_id: "run-1" });
+      await firstWarm;
+    });
+
+    expect(
+      takeWarmTaskLease({
+        repository: "acme/repo",
+        branch: "main",
+        runtimeAdapter: null,
+        model: null,
+        reasoningEffort: null,
+        sandboxEnvironmentId: null,
+        customImageId: "image-456",
+      }),
+    ).toEqual({ taskId: "task-2", runId: "run-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), {