fix(mobile): prewarm the selected sandbox image (port #3511) (#3596)

Author: Gilbert09

apps/mobile/src/features/tasks/api.test.ts

@@ -63,6 +63,30 @@ describe("runTaskInCloud", () => {
     expect(init.body).toBeUndefined();
   });
 
+  it("forwards the selected sandbox environment and custom image", async () => {
+    await runTaskInCloud("task-1", {
+      sandboxEnvironmentId: "environment-123",
+      customImageId: "image-123",
+    });
+
+    expect(bodyOf(mockFetch.mock.calls[0])).toMatchObject({
+      sandbox_environment_id: "environment-123",
+      custom_image_id: "image-123",
+    });
+  });
+
+  it("omits the sandbox environment and custom image when unset", async () => {
+    await runTaskInCloud("task-1", {
+      model: "claude-opus-4-8",
+      sandboxEnvironmentId: null,
+      customImageId: null,
+    });
+
+    const body = bodyOf(mockFetch.mock.calls[0]);
+    expect(body).not.toHaveProperty("sandbox_environment_id");
+    expect(body).not.toHaveProperty("custom_image_id");
+  });
+
   it("sends rtk_enabled=false when the run opts out", async () => {
     await runTaskInCloud("task-1", { rtkEnabled: false });
 

apps/mobile/src/features/tasks/api.ts

@@ -314,6 +314,8 @@ export async function warmTask(options: {
   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 baseUrl = getBaseUrl();
   const projectId = getProjectId();
@@ -329,6 +331,12 @@ export async function warmTask(options: {
         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 }
+          : {}),
       }),
     },
   );
@@ -431,6 +439,10 @@ export interface RunTaskInCloudOptions {
   model?: string;
   /** Reasoning effort: "low" | "medium" | "high" (model-dependent). */
   reasoningEffort?: string;
+  /** Sandbox environment / custom base image to run on. Sent so a reused warm
+   *  sandbox matches the selection instead of a mismatched default. */
+  sandboxEnvironmentId?: string | null;
+  customImageId?: string | null;
   /** Permission mode: "default" | "acceptEdits" | "plan" | "auto". */
   initialPermissionMode?: string;
   /** Source that triggered this run. */
@@ -463,6 +475,8 @@ export async function runTaskInCloud(
       options.runtimeAdapter !== undefined ||
       options.model !== undefined ||
       options.reasoningEffort !== undefined ||
+      options.sandboxEnvironmentId !== undefined ||
+      options.customImageId !== undefined ||
       options.initialPermissionMode !== undefined ||
       options.runSource !== undefined ||
       options.signalReportId !== undefined ||
@@ -488,6 +502,12 @@ export async function runTaskInCloud(
         payload.reasoning_effort = options.reasoningEffort;
       }
     }
+    if (options?.sandboxEnvironmentId) {
+      payload.sandbox_environment_id = options.sandboxEnvironmentId;
+    }
+    if (options?.customImageId) {
+      payload.custom_image_id = options.customImageId;
+    }
     if (options?.initialPermissionMode) {
       payload.initial_permission_mode = options.initialPermissionMode;
     }

apps/mobile/src/features/tasks/api.warm.test.ts

@@ -90,6 +90,58 @@ describe("warmTask", () => {
     );
   });
 
+  it("forwards the selected sandbox environment and custom image", async () => {
+    mockFetch.mockResolvedValueOnce(
+      new Response(JSON.stringify({ task_id: "task-1", run_id: "run-1" }), {
+        status: 200,
+      }),
+    );
+
+    await warmTask({
+      repository: "posthog/posthog",
+      github_integration: 7,
+      branch: "main",
+      sandbox_environment_id: "environment-123",
+      custom_image_id: "image-123",
+    });
+
+    expect(mockFetch).toHaveBeenCalledWith(
+      "https://app.posthog.test/api/projects/42/tasks/warm/",
+      expect.objectContaining({
+        body: JSON.stringify({
+          repository: "posthog/posthog",
+          github_integration: 7,
+          branch: "main",
+          runtime_adapter: null,
+          model: null,
+          reasoning_effort: null,
+          sandbox_environment_id: "environment-123",
+          custom_image_id: "image-123",
+        }),
+      }),
+    );
+  });
+
+  it("omits the sandbox environment and custom image when unset", async () => {
+    mockFetch.mockResolvedValueOnce(
+      new Response(JSON.stringify({ task_id: "task-1", run_id: "run-1" }), {
+        status: 200,
+      }),
+    );
+
+    await warmTask({
+      repository: "posthog/posthog",
+      github_integration: 7,
+      sandbox_environment_id: null,
+      custom_image_id: null,
+    });
+
+    const [, init] = mockFetch.mock.calls[0] as [string, RequestInit];
+    const body = JSON.parse(init.body as string);
+    expect(body).not.toHaveProperty("sandbox_environment_id");
+    expect(body).not.toHaveProperty("custom_image_id");
+  });
+
   it("serializes a missing branch as null", async () => {
     mockFetch.mockResolvedValueOnce(
       new Response(JSON.stringify({ task_id: "task-1", run_id: "run-1" }), {

apps/mobile/src/features/tasks/hooks/useWarmTask.test.tsx

@@ -32,6 +32,8 @@ interface Props {
   runtimeAdapter?: string | null;
   model?: string | null;
   reasoningEffort?: string | null;
+  sandboxEnvironmentId?: string | null;
+  customImageId?: string | null;
 }
 
 const composing: Props = {
@@ -198,6 +200,42 @@ describe("useWarmTask", () => {
     },
   );
 
+  it("forwards the sandbox environment and custom image", async () => {
+    render({
+      ...composing,
+      sandboxEnvironmentId: "environment-123",
+      customImageId: "image-123",
+    });
+    await flushDebounce();
+
+    expect(mockWarmTask).toHaveBeenCalledWith({
+      repository: "acme/repo",
+      github_integration: 42,
+      branch: "main",
+      ...NULL_RUNTIME,
+      sandbox_environment_id: "environment-123",
+      custom_image_id: "image-123",
+    });
+  });
+
+  it("re-warms when the custom image changes", async () => {
+    const { rerender } = render({ ...composing, customImageId: "image-123" });
+    await flushDebounce();
+    expect(mockWarmTask).toHaveBeenCalledOnce();
+
+    rerender({ ...composing, customImageId: "image-456" });
+    await flushDebounce();
+
+    expect(mockWarmTask).toHaveBeenCalledTimes(2);
+    expect(mockWarmTask).toHaveBeenLastCalledWith({
+      repository: "acme/repo",
+      github_integration: 42,
+      branch: "main",
+      ...NULL_RUNTIME,
+      custom_image_id: "image-456",
+    });
+  });
+
   it("warms again for a new selection after a failed warm", async () => {
     mockWarmTask.mockRejectedValueOnce(new Error("boom"));
     const { rerender } = render(composing);

apps/mobile/src/features/tasks/hooks/useWarmTask.ts

@@ -16,6 +16,8 @@ interface UseWarmTaskOptions {
   runtimeAdapter?: string | null;
   model?: string | null;
   reasoningEffort?: string | null;
+  sandboxEnvironmentId?: string | null;
+  customImageId?: string | null;
 }
 
 export function useWarmTask({
@@ -26,6 +28,8 @@ export function useWarmTask({
   runtimeAdapter,
   model,
   reasoningEffort,
+  sandboxEnvironmentId,
+  customImageId,
 }: UseWarmTaskOptions): void {
   const enabled = useFeatureFlag(TASKS_PREWARM_SANDBOX_FLAG);
 
@@ -36,14 +40,16 @@ 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 &&
     !!repository &&
     githubIntegrationId != null &&
     !composerIsEmpty;
   const key =
     repository && githubIntegrationId != null
-      ? `${githubIntegrationId}:${repository}:${normalizedBranch ?? ""}:${normalizedRuntimeAdapter ?? ""}:${normalizedModel ?? ""}:${normalizedReasoningEffort ?? ""}`
+      ? `${githubIntegrationId}:${repository}:${normalizedBranch ?? ""}:${normalizedRuntimeAdapter ?? ""}:${normalizedModel ?? ""}:${normalizedReasoningEffort ?? ""}:${normalizedSandboxEnvironmentId ?? ""}:${normalizedCustomImageId ?? ""}`
       : null;
 
   useEffect(() => {
@@ -68,6 +74,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;
@@ -78,6 +86,10 @@ export function useWarmTask({
         runtime_adapter: warmRuntimeAdapter,
         model: warmModel,
         reasoning_effort: warmReasoningEffort,
+        ...(warmSandboxEnvironmentId
+          ? { sandbox_environment_id: warmSandboxEnvironmentId }
+          : {}),
+        ...(warmCustomImageId ? { custom_image_id: warmCustomImageId } : {}),
       }).catch((error) => {
         lastWarmedKeyRef.current = null;
         log.warn("Failed to warm task", error);
@@ -94,5 +106,7 @@ export function useWarmTask({
     normalizedRuntimeAdapter,
     normalizedModel,
     normalizedReasoningEffort,
+    normalizedSandboxEnvironmentId,
+    normalizedCustomImageId,
   ]);
 }