fix(cloud): prewarm selected sandbox image (#3511)

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), {