Merging 496e683 into trunk-temp/pr-3481/94f6a581-6a94-45ce-adcc-124c6fda40e4

Author: trunk-io[bot]

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

@@ -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);
   }