From 4e3c1d7ac343a20be26d681b44464432b12c6ce9 Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Thu, 16 Jul 2026 16:15:08 +0100 Subject: [PATCH] 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. --- .../api-client/src/posthog-client.test.ts | 33 ++++++ packages/api-client/src/posthog-client.ts | 8 ++ .../core/src/task-detail/taskCreationHost.ts | 2 + .../src/task-detail/taskCreationSaga.test.ts | 79 ++++++++++++++ .../core/src/task-detail/taskCreationSaga.ts | 19 +++- .../task-detail/components/TaskInput.tsx | 2 + .../task-detail/hooks/useWarmTask.test.tsx | 100 ++++++++++++++++++ .../features/task-detail/hooks/useWarmTask.ts | 26 ++++- .../task-detail/hooks/warmTaskLease.ts | 4 + .../task-detail/taskCreationHostImpl.ts | 2 + 10 files changed, 272 insertions(+), 3 deletions(-) diff --git a/packages/api-client/src/posthog-client.test.ts b/packages/api-client/src/posthog-client.test.ts index 8c29ded301..3e6509ac08 100644 --- a/packages/api-client/src/posthog-client.test.ts +++ b/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, diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index 7543fb83d4..f2967294ee 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/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 } + : {}), }), }, }); diff --git a/packages/core/src/task-detail/taskCreationHost.ts b/packages/core/src/task-detail/taskCreationHost.ts index cbf55daa39..a419ccf60a 100644 --- a/packages/core/src/task-detail/taskCreationHost.ts +++ b/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, diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index f3d626053a..b53f5f990f 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/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", diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index b3e0c5218a..d74a341d35 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/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, diff --git a/packages/ui/src/features/task-detail/components/TaskInput.tsx b/packages/ui/src/features/task-detail/components/TaskInput.tsx index cf4b0766c4..a9f7a5881e 100644 --- a/packages/ui/src/features/task-detail/components/TaskInput.tsx +++ b/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 = diff --git a/packages/ui/src/features/task-detail/hooks/useWarmTask.test.tsx b/packages/ui/src/features/task-detail/hooks/useWarmTask.test.tsx index 2d53cf82c3..0f9042f85d 100644 --- a/packages/ui/src/features/task-detail/hooks/useWarmTask.test.tsx +++ b/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((resolve) => { + resolveFirstWarm = resolve; + }); + const secondWarm = new Promise((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), { diff --git a/packages/ui/src/features/task-detail/hooks/useWarmTask.ts b/packages/ui/src/features/task-detail/hooks/useWarmTask.ts index 73a06dd16c..1315104a10 100644 --- a/packages/ui/src/features/task-detail/hooks/useWarmTask.ts +++ b/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,18 +34,23 @@ export function useWarmTask({ runtimeAdapter, model, reasoningEffort, + sandboxEnvironmentId, + customImageId, }: UseWarmTaskOptions): void { const enabled = useFeatureFlag(TASKS_PREWARM_SANDBOX_FLAG); const client = useOptionalAuthenticatedClient(); const debounceRef = useRef | null>(null); const lastWarmedKeyRef = useRef(null); + const latestKeyRef = useRef(null); const isCloud = workspaceMode === "cloud"; const normalizedBranch = branch ?? null; 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,8 +66,11 @@ export function useWarmTask({ runtimeAdapter: normalizedRuntimeAdapter, model: normalizedModel, reasoningEffort: normalizedReasoningEffort, + sandboxEnvironmentId: normalizedSandboxEnvironmentId, + customImageId: normalizedCustomImageId, })}` : null; + latestKeyRef.current = key; useEffect(() => { const clearDebounce = (): void => { @@ -84,6 +94,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,9 +107,13 @@ 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) { + if (warm && latestKeyRef.current === key) { rememberWarmTaskLease( buildWarmTaskLeaseKey({ repository, @@ -105,13 +121,17 @@ export function useWarmTask({ runtimeAdapter: warmRuntimeAdapter, model: warmModel, reasoningEffort: warmReasoningEffort, + sandboxEnvironmentId: warmSandboxEnvironmentId, + customImageId: warmCustomImageId, }), { taskId: warm.task_id, runId: warm.run_id }, ); } }) .catch((error) => { - lastWarmedKeyRef.current = null; + if (latestKeyRef.current === key) { + lastWarmedKeyRef.current = null; + } log.warn("Failed to warm task", { error }); }); }, WARM_DEBOUNCE_MS); @@ -127,5 +147,7 @@ export function useWarmTask({ normalizedRuntimeAdapter, normalizedModel, normalizedReasoningEffort, + normalizedSandboxEnvironmentId, + normalizedCustomImageId, ]); } diff --git a/packages/ui/src/features/task-detail/hooks/warmTaskLease.ts b/packages/ui/src/features/task-detail/hooks/warmTaskLease.ts index ae5628c5fc..86074fc47d 100644 --- a/packages/ui/src/features/task-detail/hooks/warmTaskLease.ts +++ b/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(":"); } diff --git a/packages/ui/src/features/task-detail/taskCreationHostImpl.ts b/packages/ui/src/features/task-detail/taskCreationHostImpl.ts index af8c4078af..e4d99882e7 100644 --- a/packages/ui/src/features/task-detail/taskCreationHostImpl.ts +++ b/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); }