fix(cloud): keep latest warm image lease

Ignore stale warm responses after the sandbox selection changes so an older image cannot replace the lease for the latest environment or custom image.

Generated-By: PostHog Code
Task-Id: 6b752fb1-4fa1-45e4-920a-38306e4b7cee

Author: tatoalo

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;
@@ -238,6 +239,68 @@ describe("useWarmTask", () => {
     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), {

packages/ui/src/features/task-detail/hooks/useWarmTask.ts

@@ -42,6 +42,7 @@ export function useWarmTask({
 
   const debounceRef = useRef<ReturnType<typeof setTimeout> | null>(null);
   const lastWarmedKeyRef = useRef<string | null>(null);
+  const latestKeyRef = useRef<string | null>(null);
 
   const isCloud = workspaceMode === "cloud";
   const normalizedBranch = branch ?? null;
@@ -69,6 +70,7 @@ export function useWarmTask({
           customImageId: normalizedCustomImageId,
         })}`
       : null;
+  latestKeyRef.current = key;
 
   useEffect(() => {
     const clearDebounce = (): void => {
@@ -111,7 +113,7 @@ export function useWarmTask({
           ...(warmCustomImageId ? { custom_image_id: warmCustomImageId } : {}),
         })
         .then((warm) => {
-          if (warm) {
+          if (warm && latestKeyRef.current === key) {
             rememberWarmTaskLease(
               buildWarmTaskLeaseKey({
                 repository,
@@ -127,7 +129,9 @@ export function useWarmTask({
           }
         })
         .catch((error) => {
-          lastWarmedKeyRef.current = null;
+          if (latestKeyRef.current === key) {
+            lastWarmedKeyRef.current = null;
+          }
           log.warn("Failed to warm task", { error });
         });
     }, WARM_DEBOUNCE_MS);