feat(sessions): add stop run control for cloud runs (#3382)

Author: tatoalo

apps/code/src/main/di/container.ts

@@ -419,6 +419,12 @@ container.bind(CLOUD_TASK_AUTH).toDynamicValue((ctx) => ({
     ctx
       .get<AuthService>(MAIN_AUTH_SERVICE)
       .authenticatedFetch(fetch, url, init),
+  getCloudContext: async () => {
+    const auth = ctx.get<AuthService>(MAIN_AUTH_SERVICE);
+    const { apiHost } = await auth.getValidAccessToken();
+    const teamId = auth.getState().currentProjectId;
+    return teamId === null ? null : { apiHost, teamId };
+  },
 }));
 container.bind(MAIN_CLOUD_TASK_SERVICE).toService(CLOUD_TASK_SERVICE);
 container.load(contextMenuCoreModule);

packages/core/src/archive/archiveOrchestration.test.ts

@@ -26,6 +26,7 @@ class Harness {
     restoreCommandCenter: vi.fn(),
     getFocusedWorktreePath: vi.fn().mockReturnValue(null),
     disableFocus: vi.fn().mockResolvedValue(undefined),
+    stopCloudRun: vi.fn().mockResolvedValue(true),
     disconnectFromTask: vi.fn().mockResolvedValue(undefined),
     archive: vi.fn().mockResolvedValue(undefined),
     logError: vi.fn(),
@@ -122,6 +123,43 @@ describe("archiveTask", () => {
 
     expect(harness.deps.clearTerminalStates).not.toHaveBeenCalled();
   });
+
+  it("stops a running cloud task before archiving it", async () => {
+    await archiveTask(TASK_ID, harness.deps);
+
+    expect(harness.deps.stopCloudRun).toHaveBeenCalledWith(TASK_ID);
+    expect(
+      vi.mocked(harness.deps.stopCloudRun).mock.invocationCallOrder[0],
+    ).toBeLessThan(
+      vi.mocked(harness.deps.archive).mock.invocationCallOrder[0] ?? Infinity,
+    );
+  });
+
+  it("does not archive when a running cloud task cannot be stopped", async () => {
+    harness.deps.stopCloudRun = vi.fn().mockResolvedValue(false);
+
+    await expect(archiveTask(TASK_ID, harness.deps)).rejects.toThrow(
+      "Couldn't stop the task",
+    );
+
+    expect(harness.deps.archive).not.toHaveBeenCalled();
+    expect(harness.ids).not.toContain(TASK_ID);
+  });
+
+  it.each([
+    ["local workspace", { mode: "local" }],
+    ["task without workspace state", null],
+  ])(
+    "checks a %s for a cloud run before archiving",
+    async (_name, workspace) => {
+      harness.deps.getWorkspace = vi.fn().mockResolvedValue(workspace);
+
+      await archiveTask(TASK_ID, harness.deps);
+
+      expect(harness.deps.stopCloudRun).toHaveBeenCalledWith(TASK_ID);
+      expect(harness.deps.archive).toHaveBeenCalledWith(TASK_ID);
+    },
+  );
 });
 
 describe("archiveTasks", () => {

packages/core/src/archive/archiveOrchestration.ts

@@ -36,6 +36,7 @@ export interface ArchiveOrchestrationDeps {
   ): void;
   getFocusedWorktreePath(): string | null | undefined;
   disableFocus(): Promise<void>;
+  stopCloudRun(taskId: string, runId?: string): Promise<boolean>;
   disconnectFromTask(taskId: string): Promise<void>;
   archive(taskId: string): Promise<void>;
   logError(message: string, error: unknown): void;
@@ -58,8 +59,13 @@ export async function archiveTask(
   deps: ArchiveOrchestrationDeps,
   options?: ArchiveTaskOptions,
 ): Promise<void> {
-  const optimistic = options?.optimistic ?? true;
   const workspace = await deps.getWorkspace(taskId);
+  const stopped = await deps.stopCloudRun(taskId);
+  if (!stopped) {
+    throw new Error("Couldn't stop the task. Try again in a moment.");
+  }
+
+  const optimistic = options?.optimistic ?? true;
   const pinnedTaskIds = await deps.getPinnedTaskIds();
   const wasPinned = pinnedTaskIds.includes(taskId);
 

packages/core/src/cloud-task/cloud-task.test.ts

@@ -26,6 +26,7 @@ import { CloudTaskService } from "./cloud-task";
 
 const mockAuthService = {
   authenticatedFetch: vi.fn(),
+  getCloudContext: vi.fn(),
 };
 
 function createJsonResponse(
@@ -114,6 +115,11 @@ describe("CloudTaskService", () => {
       ),
     );
     mockAuthService.authenticatedFetch.mockReset();
+    mockAuthService.getCloudContext.mockReset();
+    mockAuthService.getCloudContext.mockResolvedValue({
+      apiHost: "https://us.posthog.com",
+      teamId: 2,
+    });
     vi.stubGlobal("fetch", fetchRouter);
 
     mockAuthService.authenticatedFetch.mockImplementation(
@@ -2335,6 +2341,42 @@ describe("CloudTaskService", () => {
     ).toBe(false);
   });
 
+  it("stops a cloud run through the run cancel endpoint", async () => {
+    mockNetFetch.mockResolvedValueOnce(
+      createJsonResponse({ id: "run-1", status: "in_progress" }, 202),
+    );
+
+    const result = await service.stop({
+      taskId: "task-1",
+      runId: "run-1",
+    });
+
+    expect(result).toEqual({ success: true, runStatus: "in_progress" });
+    const [url, init] = mockNetFetch.mock.calls[0] as [string, RequestInit];
+    expect(url).toBe(
+      "https://us.posthog.com/api/projects/2/tasks/task-1/runs/run-1/cancel/",
+    );
+    expect(init.method).toBe("POST");
+  });
+
+  it("surfaces the backend error and retryability when a stop fails", async () => {
+    mockNetFetch.mockResolvedValueOnce(
+      createJsonResponse(
+        { error: "Could not reach the run's workflow; try again" },
+        503,
+      ),
+    );
+
+    const result = await service.stop({
+      taskId: "task-1",
+      runId: "run-1",
+    });
+
+    expect(result.success).toBe(false);
+    expect(result.error).toBe("Could not reach the run's workflow; try again");
+    expect(result.retryable).toBe(true);
+  });
+
   it("does not let a stale backend-error count inflate a transport reconnect delay", async () => {
     vi.useFakeTimers();
 

packages/core/src/cloud-task/cloud-task.ts

@@ -19,6 +19,8 @@ import {
   isTerminalStatus,
   type SendCommandInput,
   type SendCommandOutput,
+  type StopInput,
+  type StopOutput,
   type TaskRunStatus,
   type WatchInput,
 } from "./schemas";
@@ -520,6 +522,65 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
     }
   }
 
+  async stop(input: StopInput): Promise<StopOutput> {
+    try {
+      const context = await this.auth.getCloudContext();
+      if (!context) {
+        return { success: false, error: "No active cloud project" };
+      }
+      const url = `${context.apiHost}/api/projects/${context.teamId}/tasks/${input.taskId}/runs/${input.runId}/cancel/`;
+      const response = await this.auth.authenticatedFetch(url, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify(input.reason ? { reason: input.reason } : {}),
+      });
+
+      if (!response.ok) {
+        const errorText = await response.text().catch(() => "");
+        let errorMessage = `Stop failed with status ${response.status}`;
+        try {
+          const errorJson = JSON.parse(errorText) as { error?: unknown };
+          if (typeof errorJson.error === "string" && errorJson.error) {
+            errorMessage = errorJson.error;
+          }
+        } catch {
+          if (errorText) errorMessage = errorText;
+        }
+
+        this.log.warn("Cloud run stop failed", {
+          taskId: input.taskId,
+          runId: input.runId,
+          status: response.status,
+          error: errorMessage,
+        });
+        return {
+          success: false,
+          error: errorMessage,
+          retryable: response.status === 503 || response.status >= 500,
+        };
+      }
+
+      const data = (await response.json()) as { status?: string };
+      this.log.info("Cloud run stop accepted", {
+        taskId: input.taskId,
+        runId: input.runId,
+        runStatus: data.status,
+      });
+      return { success: true, runStatus: data.status };
+    } catch (error) {
+      const errorMessage =
+        error instanceof Error ? error.message : "Unknown error";
+      this.log.error("Cloud run stop error", {
+        taskId: input.taskId,
+        runId: input.runId,
+        error: errorMessage,
+      });
+      return { success: false, error: errorMessage, retryable: true };
+    }
+  }
+
   @preDestroy()
   unwatchAll(): void {
     for (const key of [...this.watchers.keys()]) {

packages/core/src/cloud-task/identifiers.ts

@@ -3,4 +3,5 @@ export const CLOUD_TASK_AUTH = Symbol.for("posthog.core.cloudTaskAuth");
 
 export interface ICloudTaskAuth {
   authenticatedFetch(url: string, init?: RequestInit): Promise<Response>;
+  getCloudContext(): Promise<{ apiHost: string; teamId: number } | null>;
 }

packages/core/src/cloud-task/schemas.ts

@@ -77,3 +77,20 @@ export const sendCommandOutput = z.object({
 });
 
 export type SendCommandOutput = z.infer<typeof sendCommandOutput>;
+
+export const stopInput = z.object({
+  taskId: z.string(),
+  runId: z.string(),
+  reason: z.string().optional(),
+});
+
+export type StopInput = z.infer<typeof stopInput>;
+
+export const stopOutput = z.object({
+  success: z.boolean(),
+  runStatus: z.string().optional(),
+  error: z.string().optional(),
+  retryable: z.boolean().optional(),
+});
+
+export type StopOutput = z.infer<typeof stopOutput>;

packages/core/src/context-menu/context-menu.test.ts

@@ -103,6 +103,22 @@ describe("ContextMenuService.showTaskContextMenu", () => {
     expect(labels(menu.lastItems)).not.toContain("Suspend");
   });
 
+  it("offers Stop task only for a stoppable run", async () => {
+    const running = new FakeContextMenu();
+    const result = makeService(running).showTaskContextMenu({
+      ...baseTask,
+      canStop: true,
+    });
+    await running.shown;
+    findItem(running.lastItems, "Stop task").click();
+    expect(await result).toEqual({ action: { type: "stop" } });
+
+    const idle = new FakeContextMenu();
+    makeService(idle).showTaskContextMenu(baseTask);
+    await idle.shown;
+    expect(labels(idle.lastItems)).not.toContain("Stop task");
+  });
+
   it("hides Add to Command Center when already in it", async () => {
     const inCc = new FakeContextMenu();
     makeService(inCc).showTaskContextMenu({

packages/core/src/context-menu/context-menu.ts

@@ -114,6 +114,7 @@ export class ContextMenuService {
       folderPath,
       isPinned,
       isSuspended,
+      canStop,
       isInCommandCenter,
       hasEmptyCommandCenterCell,
       channels,
@@ -141,6 +142,9 @@ export class ContextMenuService {
     return this.showMenu<TaskAction>([
       this.item(isPinned ? "Unpin" : "Pin", { type: "pin" }),
       this.item("Rename", { type: "rename" }),
+      ...(canStop
+        ? [this.separator(), this.item("Stop task", { type: "stop" as const })]
+        : []),
       ...(worktreePath
         ? [
             this.separator(),

packages/core/src/context-menu/schemas.ts

@@ -6,6 +6,7 @@ export const taskContextMenuInput = z.object({
   folderPath: z.string().optional(),
   isPinned: z.boolean().optional(),
   isSuspended: z.boolean().optional(),
+  canStop: z.boolean().optional(),
   isInCommandCenter: z.boolean().optional(),
   hasEmptyCommandCenterCell: z.boolean().optional(),
   // Top-level desktop_file_system channels available as "File to…" targets.
@@ -45,6 +46,7 @@ const taskAction = z.discriminatedUnion("type", [
   z.object({ type: z.literal("rename") }),
   z.object({ type: z.literal("pin") }),
   z.object({ type: z.literal("suspend") }),
+  z.object({ type: z.literal("stop") }),
   z.object({ type: z.literal("archive") }),
   z.object({ type: z.literal("archive-prior") }),
   z.object({ type: z.literal("delete") }),