diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index bca03305ef..9494377250 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -419,6 +419,12 @@ container.bind(CLOUD_TASK_AUTH).toDynamicValue((ctx) => ({ ctx .get(MAIN_AUTH_SERVICE) .authenticatedFetch(fetch, url, init), + getCloudContext: async () => { + const auth = ctx.get(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); diff --git a/packages/core/src/archive/archiveOrchestration.test.ts b/packages/core/src/archive/archiveOrchestration.test.ts index 1058e19286..0f73a1e780 100644 --- a/packages/core/src/archive/archiveOrchestration.test.ts +++ b/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", () => { diff --git a/packages/core/src/archive/archiveOrchestration.ts b/packages/core/src/archive/archiveOrchestration.ts index b7e4c75ceb..19d7801c1a 100644 --- a/packages/core/src/archive/archiveOrchestration.ts +++ b/packages/core/src/archive/archiveOrchestration.ts @@ -36,6 +36,7 @@ export interface ArchiveOrchestrationDeps { ): void; getFocusedWorktreePath(): string | null | undefined; disableFocus(): Promise; + stopCloudRun(taskId: string, runId?: string): Promise; disconnectFromTask(taskId: string): Promise; archive(taskId: string): Promise; logError(message: string, error: unknown): void; @@ -58,8 +59,13 @@ export async function archiveTask( deps: ArchiveOrchestrationDeps, options?: ArchiveTaskOptions, ): Promise { - 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); diff --git a/packages/core/src/cloud-task/cloud-task.test.ts b/packages/core/src/cloud-task/cloud-task.test.ts index 48c7dac251..c633a95e99 100644 --- a/packages/core/src/cloud-task/cloud-task.test.ts +++ b/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(); diff --git a/packages/core/src/cloud-task/cloud-task.ts b/packages/core/src/cloud-task/cloud-task.ts index 10cbfc416f..296f5573e5 100644 --- a/packages/core/src/cloud-task/cloud-task.ts +++ b/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 { } } + async stop(input: StopInput): Promise { + 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()]) { diff --git a/packages/core/src/cloud-task/identifiers.ts b/packages/core/src/cloud-task/identifiers.ts index 5693714913..51c1586cba 100644 --- a/packages/core/src/cloud-task/identifiers.ts +++ b/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; + getCloudContext(): Promise<{ apiHost: string; teamId: number } | null>; } diff --git a/packages/core/src/cloud-task/schemas.ts b/packages/core/src/cloud-task/schemas.ts index 347a5a6c08..08b0e72dd3 100644 --- a/packages/core/src/cloud-task/schemas.ts +++ b/packages/core/src/cloud-task/schemas.ts @@ -77,3 +77,20 @@ export const sendCommandOutput = z.object({ }); export type SendCommandOutput = z.infer; + +export const stopInput = z.object({ + taskId: z.string(), + runId: z.string(), + reason: z.string().optional(), +}); + +export type StopInput = z.infer; + +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; diff --git a/packages/core/src/context-menu/context-menu.test.ts b/packages/core/src/context-menu/context-menu.test.ts index 036e4cffc4..b06a4a79cb 100644 --- a/packages/core/src/context-menu/context-menu.test.ts +++ b/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({ diff --git a/packages/core/src/context-menu/context-menu.ts b/packages/core/src/context-menu/context-menu.ts index e9e9e74f71..d4437e9488 100644 --- a/packages/core/src/context-menu/context-menu.ts +++ b/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([ 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(), diff --git a/packages/core/src/context-menu/schemas.ts b/packages/core/src/context-menu/schemas.ts index 469b95045c..b6b664ede4 100644 --- a/packages/core/src/context-menu/schemas.ts +++ b/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") }), diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 58b611dccd..a0c12d0dda 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -164,6 +164,7 @@ export interface SessionTrpc { unwatch: TrpcMutation; retry: TrpcMutation; sendCommand: TrpcMutation; + stop: TrpcMutation; onUpdate: TrpcSubscription; }; handoff: { @@ -3227,6 +3228,105 @@ export class SessionService { } } + async stopCloudRun(taskId: string, runId?: string): Promise { + const session = this.d.store.getSessionByTaskId(taskId); + let taskRunId: string; + try { + const client = await this.d.getAuthenticatedClient(); + if (!client) return false; + const task = (await client.getTask(taskId)) as Task; + const latestRun = task.latest_run; + if (!latestRun || latestRun.environment !== "cloud") { + return true; + } + if (isTerminalStatus(latestRun.status)) { + const rendererStillExpectsCompletion = + session?.isCloud === true && + session.taskRunId === latestRun.id && + !isTerminalStatus(session.cloudStatus); + if (!rendererStillExpectsCompletion) { + return true; + } + } + taskRunId = latestRun.id; + if (runId && runId !== taskRunId) { + this.d.log.warn("Refusing to stop a newer cloud run", { + taskId, + requestedRunId: runId, + taskRunId, + }); + return false; + } + } catch (error) { + this.d.log.error("Failed to resolve current cloud run", error); + return false; + } + + const matchingSession = + session?.isCloud && session.taskRunId === taskRunId ? session : undefined; + const previousPromptState = matchingSession + ? { + isPromptPending: matchingSession.isPromptPending, + promptStartedAt: matchingSession.promptStartedAt, + } + : undefined; + if (matchingSession) { + this.d.store.updateSession(matchingSession.taskRunId, { + stopRequested: true, + isPromptPending: false, + promptStartedAt: null, + }); + } + + try { + const result = await this.d.trpc.cloudTask.stop.mutate({ + taskId, + runId: taskRunId, + }); + + if (!result.success) { + if (matchingSession) { + this.d.store.updateSession(matchingSession.taskRunId, { + stopRequested: false, + ...previousPromptState, + }); + } + this.d.log.warn("Cloud run stop failed", { + taskId, + error: result.error, + retryable: result.retryable, + }); + return false; + } + + const durationSeconds = matchingSession + ? Math.round((Date.now() - matchingSession.startedAt) / 1000) + : undefined; + const promptCount = matchingSession?.events.filter( + (e) => "method" in e.message && e.message.method === "session/prompt", + ).length; + this.d.track(ANALYTICS_EVENTS.TASK_RUN_STOPPED, { + task_id: taskId, + execution_type: "cloud", + ...(durationSeconds === undefined + ? {} + : { duration_seconds: durationSeconds }), + ...(promptCount === undefined ? {} : { prompts_sent: promptCount }), + }); + + return true; + } catch (error) { + if (matchingSession) { + this.d.store.updateSession(matchingSession.taskRunId, { + stopRequested: false, + ...previousPromptState, + }); + } + this.d.log.error("Failed to stop cloud run", error); + return false; + } + } + private async getCloudCommandAuth(): Promise<{ apiHost: string; teamId: number; diff --git a/packages/core/src/sessions/sessionServiceStopCloudRun.test.ts b/packages/core/src/sessions/sessionServiceStopCloudRun.test.ts new file mode 100644 index 0000000000..ba1d1a127f --- /dev/null +++ b/packages/core/src/sessions/sessionServiceStopCloudRun.test.ts @@ -0,0 +1,199 @@ +import type { AgentSession } from "@posthog/shared"; +import { describe, expect, it, vi } from "vitest"; +import { SessionService, type SessionServiceDeps } from "./sessionService"; + +function makeSession(overrides: Partial = {}): AgentSession { + return { + taskRunId: "run-1", + taskId: "task-1", + taskTitle: "Test task", + channel: "", + events: [], + startedAt: 1, + status: "connected", + isPromptPending: false, + isCompacting: false, + promptStartedAt: null, + pendingPermissions: new Map(), + pausedDurationMs: 0, + messageQueue: [], + optimisticItems: [], + isCloud: true, + cloudStatus: "in_progress", + ...overrides, + } as AgentSession; +} + +function createHarness( + session: AgentSession | undefined, + stopMutate: ReturnType, + latestRun: { + id: string; + environment: "local" | "cloud"; + status: + | "not_started" + | "queued" + | "in_progress" + | "completed" + | "failed" + | "cancelled"; + } | null = session + ? { + id: session.taskRunId, + environment: session.isCloud ? "cloud" : "local", + status: session.cloudStatus ?? "in_progress", + } + : { id: "run-1", environment: "cloud", status: "in_progress" }, +) { + const updateSession = vi.fn(); + const track = vi.fn(); + const deps = { + store: { + getSessionByTaskId: (taskId: string) => + session?.taskId === taskId ? session : undefined, + updateSession, + }, + log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }, + track, + fetchAuthState: async () => ({ cloudRegion: "us", currentProjectId: 2 }), + getAuthenticatedClient: async () => ({ + getTask: async () => ({ latest_run: latestRun }), + }), + trpc: { + agent: { + onSessionIdleKilled: { + subscribe: () => ({ unsubscribe: vi.fn() }), + }, + }, + cloudTask: { + stop: { mutate: stopMutate }, + }, + }, + } as unknown as SessionServiceDeps; + + return { service: new SessionService(deps), updateSession, track }; +} + +describe("SessionService.stopCloudRun", () => { + it("marks the session stopping and stops the run via the backend", async () => { + const stopMutate = vi.fn().mockResolvedValue({ success: true }); + const { service, updateSession, track } = createHarness( + makeSession(), + stopMutate, + ); + + const result = await service.stopCloudRun("task-1"); + + expect(result).toBe(true); + expect(stopMutate).toHaveBeenCalledWith({ + taskId: "task-1", + runId: "run-1", + }); + expect(updateSession).toHaveBeenCalledWith( + "run-1", + expect.objectContaining({ stopRequested: true }), + ); + expect(track).toHaveBeenCalledWith( + "Task run stopped", + expect.objectContaining({ task_id: "task-1", execution_type: "cloud" }), + ); + }); + + it("treats an already terminal cloud run as stopped", async () => { + const stopMutate = vi.fn(); + const { service } = createHarness( + makeSession({ cloudStatus: "cancelled" }), + stopMutate, + ); + + const result = await service.stopCloudRun("task-1"); + + expect(result).toBe(true); + expect(stopMutate).not.toHaveBeenCalled(); + }); + + it("repairs completion when the backend is terminal but the renderer is stale", async () => { + const stopMutate = vi.fn().mockResolvedValue({ success: true }); + const { service } = createHarness(makeSession(), stopMutate, { + id: "run-1", + environment: "cloud", + status: "cancelled", + }); + + const result = await service.stopCloudRun("task-1"); + + expect(result).toBe(true); + expect(stopMutate).toHaveBeenCalledWith({ + taskId: "task-1", + runId: "run-1", + }); + }); + + it("does not call the backend for a local session", async () => { + const stopMutate = vi.fn(); + const { service } = createHarness( + makeSession({ isCloud: false }), + stopMutate, + ); + + const result = await service.stopCloudRun("task-1"); + + expect(result).toBe(true); + expect(stopMutate).not.toHaveBeenCalled(); + }); + + it("clears the stopping marker when the stop fails", async () => { + const stopMutate = vi + .fn() + .mockResolvedValue({ success: false, error: "boom" }); + const { service, updateSession, track } = createHarness( + makeSession({ isPromptPending: true, promptStartedAt: 123 }), + stopMutate, + ); + + const result = await service.stopCloudRun("task-1"); + + expect(result).toBe(false); + expect(updateSession).toHaveBeenLastCalledWith("run-1", { + stopRequested: false, + isPromptPending: true, + promptStartedAt: 123, + }); + expect(track).not.toHaveBeenCalled(); + }); + + it("stops a sidebar run that has no mounted session", async () => { + const stopMutate = vi.fn().mockResolvedValue({ success: true }); + const { service, updateSession, track } = createHarness( + undefined, + stopMutate, + ); + + const result = await service.stopCloudRun("task-1", "run-1"); + + expect(result).toBe(true); + expect(stopMutate).toHaveBeenCalledWith({ + taskId: "task-1", + runId: "run-1", + }); + expect(updateSession).not.toHaveBeenCalled(); + expect(track).toHaveBeenCalledWith("Task run stopped", { + task_id: "task-1", + execution_type: "cloud", + }); + }); + + it("does not stop a newer run when the sidebar supplied a stale run id", async () => { + const stopMutate = vi.fn().mockResolvedValue({ success: true }); + const { service } = createHarness(undefined, stopMutate, { + id: "run-current", + environment: "cloud", + status: "queued", + }); + + const result = await service.stopCloudRun("task-1", "run-stale"); + + expect(result).toBe(false); + expect(stopMutate).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/sidebar/taskRunning.test.ts b/packages/core/src/sidebar/taskRunning.test.ts index ce08b8c75b..6d322c47f5 100644 --- a/packages/core/src/sidebar/taskRunning.test.ts +++ b/packages/core/src/sidebar/taskRunning.test.ts @@ -67,11 +67,18 @@ describe("isTaskActivelyRunning", () => { expected: true, }, { - name: "queued cloud run (not yet running)", + name: "queued cloud run", environment: "cloud", status: "queued", isGenerating: false, - expected: false, + expected: true, + }, + { + name: "not-started cloud run", + environment: "cloud", + status: "not_started", + isGenerating: false, + expected: true, }, { name: "completed cloud run", diff --git a/packages/core/src/sidebar/taskRunning.ts b/packages/core/src/sidebar/taskRunning.ts index 6348e0d089..73363a36fa 100644 --- a/packages/core/src/sidebar/taskRunning.ts +++ b/packages/core/src/sidebar/taskRunning.ts @@ -1,3 +1,4 @@ +import { isTerminalStatus } from "@posthog/shared/domain-types"; import type { TaskData } from "./sidebarData.types"; // Drives the "Archive running task?" confirmation. Persisted run status is only @@ -7,6 +8,8 @@ import type { TaskData } from "./sidebarData.types"; export function isTaskActivelyRunning(task: TaskData): boolean { if (task.isGenerating) return true; return ( - task.taskRunEnvironment === "cloud" && task.taskRunStatus === "in_progress" + task.taskRunEnvironment === "cloud" && + task.taskRunStatus !== undefined && + !isTerminalStatus(task.taskRunStatus) ); } diff --git a/packages/core/src/tasks/contextMenuActions.test.ts b/packages/core/src/tasks/contextMenuActions.test.ts index 1ee3cc38ab..9e06ca86fe 100644 --- a/packages/core/src/tasks/contextMenuActions.test.ts +++ b/packages/core/src/tasks/contextMenuActions.test.ts @@ -21,6 +21,9 @@ describe("resolveTaskContextMenuIntent", () => { expect(resolveTaskContextMenuIntent({ type: "rename" }, {})).toEqual({ type: "rename", }); + expect(resolveTaskContextMenuIntent({ type: "stop" }, {})).toEqual({ + type: "stop", + }); expect(resolveTaskContextMenuIntent({ type: "delete" }, {})).toEqual({ type: "delete", }); diff --git a/packages/core/src/tasks/contextMenuActions.ts b/packages/core/src/tasks/contextMenuActions.ts index 34df7916fb..a9d12056ea 100644 --- a/packages/core/src/tasks/contextMenuActions.ts +++ b/packages/core/src/tasks/contextMenuActions.ts @@ -7,6 +7,7 @@ export type TaskContextMenuIntent = | { type: "rename" } | { type: "pin" } | { type: "suspend" } + | { type: "stop" } | { type: "restore" } | { type: "archive" } | { type: "archive-prior" } @@ -26,6 +27,8 @@ export function resolveTaskContextMenuIntent( return { type: "pin" }; case "suspend": return flags.isSuspended ? { type: "restore" } : { type: "suspend" }; + case "stop": + return { type: "stop" }; case "archive": return { type: "archive" }; case "archive-prior": diff --git a/packages/host-router/src/routers/cloud-task.router.ts b/packages/host-router/src/routers/cloud-task.router.ts index be0b7dd451..a49cdd8b47 100644 --- a/packages/host-router/src/routers/cloud-task.router.ts +++ b/packages/host-router/src/routers/cloud-task.router.ts @@ -6,6 +6,8 @@ import { retryInput, sendCommandInput, sendCommandOutput, + stopInput, + stopOutput, unwatchInput, watchInput, } from "@posthog/core/cloud-task/schemas"; @@ -43,6 +45,13 @@ export const cloudTaskRouter = router({ .sendCommand(input), ), + stop: publicProcedure + .input(stopInput) + .output(stopOutput) + .mutation(({ ctx, input }) => + ctx.container.get(CLOUD_TASK_SERVICE).stop(input), + ), + onUpdate: publicProcedure .input(onUpdateInput) .subscription(async function* (opts) { diff --git a/packages/shared/src/analytics-events.ts b/packages/shared/src/analytics-events.ts index e3b5394b3a..c1dffa2051 100644 --- a/packages/shared/src/analytics-events.ts +++ b/packages/shared/src/analytics-events.ts @@ -134,6 +134,13 @@ export interface TaskRunCancelledProperties { prompts_sent: number; } +export interface TaskRunStoppedProperties { + task_id: string; + execution_type: ExecutionType; + duration_seconds?: number; + prompts_sent?: number; +} + export interface PromptSentProperties { task_id: string; is_initial: boolean; @@ -1068,6 +1075,7 @@ export const ANALYTICS_EVENTS = { TASK_RUN_STARTED: "Task run started", TASK_RUN_COMPLETED: "Task run completed", TASK_RUN_CANCELLED: "Task run cancelled", + TASK_RUN_STOPPED: "Task run stopped", PROMPT_SENT: "Prompt sent", // Claude Code session import @@ -1228,6 +1236,7 @@ export type EventPropertyMap = { [ANALYTICS_EVENTS.TASK_RUN_STARTED]: TaskRunStartedProperties; [ANALYTICS_EVENTS.TASK_RUN_COMPLETED]: TaskRunCompletedProperties; [ANALYTICS_EVENTS.TASK_RUN_CANCELLED]: TaskRunCancelledProperties; + [ANALYTICS_EVENTS.TASK_RUN_STOPPED]: TaskRunStoppedProperties; [ANALYTICS_EVENTS.PROMPT_SENT]: PromptSentProperties; // Claude Code session import diff --git a/packages/shared/src/sessions.ts b/packages/shared/src/sessions.ts index f938d1b9a4..e2b75ee374 100644 --- a/packages/shared/src/sessions.ts +++ b/packages/shared/src/sessions.ts @@ -87,6 +87,7 @@ export interface AgentSession { initialPrompt?: ContentBlock[]; cloudBranch?: string | null; handoffInProgress?: boolean; + stopRequested?: boolean; optimisticItems: OptimisticItem[]; contextUsed?: number; contextSize?: number; diff --git a/packages/ui/src/features/archive/useArchiveTask.ts b/packages/ui/src/features/archive/useArchiveTask.ts index 9d1e60979a..2f8000d4e5 100644 --- a/packages/ui/src/features/archive/useArchiveTask.ts +++ b/packages/ui/src/features/archive/useArchiveTask.ts @@ -116,6 +116,11 @@ function makeOrchestrationDeps( log.info("Unfocusing workspace before archiving"); await useFocusStore.getState().disableFocus(); }, + stopCloudRun: (taskId, runId) => + resolveService(SESSION_SERVICE).stopCloudRun( + taskId, + runId, + ), disconnectFromTask: (taskId) => resolveService(SESSION_SERVICE).disconnectFromTask( taskId, diff --git a/packages/ui/src/features/sessions/components/StopCloudRunButton.tsx b/packages/ui/src/features/sessions/components/StopCloudRunButton.tsx new file mode 100644 index 0000000000..6836985673 --- /dev/null +++ b/packages/ui/src/features/sessions/components/StopCloudRunButton.tsx @@ -0,0 +1,53 @@ +import { Spinner, Stop } from "@phosphor-icons/react"; +import { isTerminalStatus } from "@posthog/core/cloud-task/schemas"; +import { Button as QuillButton } from "@posthog/quill"; +import { useState } from "react"; +import { shallow } from "zustand/shallow"; +import { useSessionSelector } from "../useSession"; +import { StopCloudRunDialog } from "./StopCloudRunDialog"; + +interface StopCloudRunButtonProps { + taskId: string; +} + +export function StopCloudRunButton({ taskId }: StopCloudRunButtonProps) { + const { isCloud, cloudStatus, stopRequested } = useSessionSelector( + taskId, + (session) => ({ + isCloud: session?.isCloud ?? false, + cloudStatus: session?.cloudStatus ?? null, + stopRequested: session?.stopRequested ?? false, + }), + shallow, + ); + const [confirmOpen, setConfirmOpen] = useState(false); + + if (!isCloud || isTerminalStatus(cloudStatus)) return null; + + return ( + <> +
+ setConfirmOpen(true)} + > + {stopRequested ? ( + + ) : ( + + )} + {stopRequested ? "Stopping..." : "Stop run"} + +
+ + + ); +} diff --git a/packages/ui/src/features/sessions/components/StopCloudRunDialog.tsx b/packages/ui/src/features/sessions/components/StopCloudRunDialog.tsx new file mode 100644 index 0000000000..8e651b4c8a --- /dev/null +++ b/packages/ui/src/features/sessions/components/StopCloudRunDialog.tsx @@ -0,0 +1,75 @@ +import { Stop } from "@phosphor-icons/react"; +import { + SESSION_SERVICE, + type SessionService, +} from "@posthog/core/sessions/sessionService"; +import { useService } from "@posthog/di/react"; +import { GitDialog } from "@posthog/ui/features/git-interaction/components/GitInteractionDialogs"; +import { Text } from "@radix-ui/themes"; +import { useState } from "react"; + +interface StopCloudRunDialogProps { + open: boolean; + taskId: string; + runId?: string; + title: string; + buttonLabel: string; + onOpenChange: (open: boolean) => void; + onStopped?: () => void; +} + +export function StopCloudRunDialog({ + open, + taskId, + runId, + title, + buttonLabel, + onOpenChange, + onStopped, +}: StopCloudRunDialogProps) { + const sessionService = useService(SESSION_SERVICE); + const [isStopping, setIsStopping] = useState(false); + const [error, setError] = useState(null); + + const handleOpenChange = (nextOpen: boolean) => { + if (!nextOpen) { + setError(null); + } + onOpenChange(nextOpen); + }; + + const handleConfirm = async () => { + setIsStopping(true); + setError(null); + try { + const stopped = await sessionService.stopCloudRun(taskId, runId); + if (!stopped) { + setError("Couldn't stop the run. Try again in a moment."); + return; + } + onStopped?.(); + handleOpenChange(false); + } finally { + setIsStopping(false); + } + }; + + return ( + } + title={title} + error={error} + buttonLabel={buttonLabel} + isSubmitting={isStopping} + onSubmit={handleConfirm} + > + + This ends the cloud session and shuts down its sandbox. You can pick the + conversation back up later by sending a new message. To stop only the + current response, press Esc instead. + + + ); +} diff --git a/packages/ui/src/features/sidebar/components/ArchiveRunningTaskDialog.tsx b/packages/ui/src/features/sidebar/components/ArchiveRunningTaskDialog.tsx index 2a2e91ddb1..38e91fe4f9 100644 --- a/packages/ui/src/features/sidebar/components/ArchiveRunningTaskDialog.tsx +++ b/packages/ui/src/features/sidebar/components/ArchiveRunningTaskDialog.tsx @@ -1,24 +1,51 @@ import { Warning } from "@phosphor-icons/react"; -import { AlertDialog, Button, Flex } from "@radix-ui/themes"; +import { AlertDialog, Button, Flex, Text } from "@radix-ui/themes"; +import { useState } from "react"; interface ArchiveRunningTaskDialogProps { open: boolean; taskTitle: string; - onConfirm: () => void; + stopsCloudSandbox: boolean; + onConfirm: () => Promise; onCancel: () => void; } export function ArchiveRunningTaskDialog({ open, taskTitle, + stopsCloudSandbox, onConfirm, onCancel, }: ArchiveRunningTaskDialogProps) { + const [isSubmitting, setIsSubmitting] = useState(false); + const [error, setError] = useState(null); + + const handleConfirm = async (event: React.MouseEvent) => { + event.preventDefault(); + if (isSubmitting) return; + setIsSubmitting(true); + setError(null); + try { + await onConfirm(); + } catch (caughtError) { + setError( + caughtError instanceof Error + ? caughtError.message + : "Couldn't archive the task. Try again in a moment.", + ); + } finally { + setIsSubmitting(false); + } + }; + return ( { - if (!isOpen) onCancel(); + if (!isOpen && !isSubmitting) { + setError(null); + onCancel(); + } }} > @@ -30,17 +57,37 @@ export function ArchiveRunningTaskDialog({ {taskTitle ? `"${taskTitle}"` : "This task"} is still running. - Archiving it now will stop the agent. You can unarchive it later. + {stopsCloudSandbox + ? " Archiving it will stop its cloud run and shut down the sandbox." + : " Archiving it now will stop the agent."}{" "} + You can unarchive it later. + {error ? ( + + {error} + + ) : null} + - - diff --git a/packages/ui/src/features/sidebar/components/SidebarMenu.tsx b/packages/ui/src/features/sidebar/components/SidebarMenu.tsx index 29c274b45f..3be10a1f11 100644 --- a/packages/ui/src/features/sidebar/components/SidebarMenu.tsx +++ b/packages/ui/src/features/sidebar/components/SidebarMenu.tsx @@ -10,6 +10,7 @@ import { import { useCommandCenterStore } from "@posthog/ui/features/command-center/commandCenterStore"; import { useExternalAppAction } from "@posthog/ui/features/external-apps/useExternalAppAction"; import { useFolders } from "@posthog/ui/features/folders/useFolders"; +import { StopCloudRunDialog } from "@posthog/ui/features/sessions/components/StopCloudRunDialog"; import { useArchivingTasksStore } from "@posthog/ui/features/sidebar/archivingTasksStore"; import { useSidebarStore } from "@posthog/ui/features/sidebar/sidebarStore"; import { useTaskSelectionStore } from "@posthog/ui/features/sidebar/taskSelectionStore"; @@ -74,10 +75,10 @@ function SidebarMenuComponent() { activeView: view, }); - const taskMap = new Map(); - for (const task of allTasks) { - taskMap.set(task.id, task); - } + const taskMap = useMemo( + () => new Map(allTasks.map((task) => [task.id, task])), + [allTasks], + ); const commandCenterCells = useCommandCenterStore((s) => s.cells); const assignTaskToCommandCenter = useCommandCenterStore((s) => s.assignTask); @@ -107,6 +108,12 @@ function SidebarMenuComponent() { const [archiveConfirm, setArchiveConfirm] = useState<{ taskId: string; taskTitle: string; + stopsCloudSandbox: boolean; + } | null>(null); + const [stopConfirm, setStopConfirm] = useState<{ + taskId: string; + taskTitle: string; + runId?: string; } | null>(null); // Escape clears any bulk task selection (moved here from the retired @@ -297,6 +304,7 @@ function SidebarMenuComponent() { const taskData = allSidebarTasks.find((t) => t.id === taskId); const task = taskMap.get(taskId) ?? taskData; if (task) { + const runId = taskMap.get(taskId)?.latest_run?.id; const workspace = workspaces[taskId]; const isInCommandCenter = commandCenterCells.some( (id) => id === taskId && taskMap.has(id), @@ -310,9 +318,19 @@ function SidebarMenuComponent() { folderPath: workspace?.folderPath ?? undefined, isPinned, isSuspended: taskData?.isSuspended, + canStop: + taskData?.taskRunEnvironment === "cloud" && + isTaskActivelyRunning(taskData), + runId, isInCommandCenter, hasEmptyCommandCenterCell, onTogglePin: () => handleTaskTogglePin(taskId), + onStop: (stopTaskId, taskTitle, stopRunId) => + setStopConfirm({ + taskId: stopTaskId, + taskTitle, + runId: stopRunId, + }), onArchive: handleTaskArchive, onArchivePrior: handleArchivePrior, onAddToCommandCenter: () => { @@ -333,18 +351,25 @@ function SidebarMenuComponent() { // shows a spinner and ignores clicks/pins/right-clicks until it resolves. // Guards against repeated clicks: a second call while archiving is a no-op. const runArchive = useCallback( - (taskId: string) => { + async (taskId: string) => { const store = useArchivingTasksStore.getState(); - if (store.isArchiving(taskId)) return; + if (store.isArchiving(taskId)) { + return { + success: false, + error: new Error("Task is already archiving"), + }; + } store.startArchiving(taskId); - void archiveTask({ taskId }) - .catch((error) => { - log.error("Failed to archive task", error); - toast.error("Failed to archive task"); - }) - .finally(() => { - useArchivingTasksStore.getState().stopArchiving(taskId); - }); + try { + await archiveTask({ taskId }); + return { success: true as const }; + } catch (error) { + log.error("Failed to archive task", error); + toast.error("Failed to archive task"); + return { success: false as const, error }; + } finally { + useArchivingTasksStore.getState().stopArchiving(taskId); + } }, [archiveTask], ); @@ -354,19 +379,28 @@ function SidebarMenuComponent() { if (useArchivingTasksStore.getState().isArchiving(taskId)) return; const task = allSidebarTasks.find((t) => t.id === taskId); if (task && isTaskActivelyRunning(task)) { - setArchiveConfirm({ taskId, taskTitle: task.title }); + setArchiveConfirm({ + taskId, + taskTitle: task.title, + stopsCloudSandbox: task.taskRunEnvironment === "cloud", + }); return; } - runArchive(taskId); + void runArchive(taskId); }, [allSidebarTasks, runArchive], ); - const handleConfirmArchive = useCallback(() => { + const handleConfirmArchive = useCallback(async () => { if (!archiveConfirm) return; const { taskId } = archiveConfirm; + const result = await runArchive(taskId); + if (!result.success) { + throw result.error instanceof Error + ? result.error + : new Error("Couldn't archive the task. Try again in a moment."); + } setArchiveConfirm(null); - runArchive(taskId); }, [archiveConfirm, runArchive]); const handleTaskTogglePin = useCallback( @@ -486,9 +520,23 @@ function SidebarMenuComponent() { setArchiveConfirm(null)} /> + {stopConfirm ? ( + { + if (!open) setStopConfirm(null); + }} + onStopped={() => toast.success("Stop requested")} + /> + ) : null} ); } diff --git a/packages/ui/src/features/tasks/useTaskContextMenu.ts b/packages/ui/src/features/tasks/useTaskContextMenu.ts index d7c9c29c81..2aa6a32b4e 100644 --- a/packages/ui/src/features/tasks/useTaskContextMenu.ts +++ b/packages/ui/src/features/tasks/useTaskContextMenu.ts @@ -45,9 +45,12 @@ export function useTaskContextMenu() { folderPath?: string; isPinned?: boolean; isSuspended?: boolean; + canStop?: boolean; + runId?: string; isInCommandCenter?: boolean; hasEmptyCommandCenterCell?: boolean; onTogglePin?: () => void; + onStop?: (taskId: string, taskTitle: string, runId?: string) => void; onArchive?: (taskId: string) => void; onArchivePrior?: (taskId: string) => void; onAddToCommandCenter?: () => void; @@ -61,9 +64,12 @@ export function useTaskContextMenu() { folderPath, isPinned, isSuspended, + canStop, + runId, isInCommandCenter, hasEmptyCommandCenterCell, onTogglePin, + onStop, onArchive, onArchivePrior, onAddToCommandCenter, @@ -76,6 +82,7 @@ export function useTaskContextMenu() { folderPath, isPinned, isSuspended, + canStop, isInCommandCenter, hasEmptyCommandCenterCell, channels: channels.map(({ id, name }) => ({ id, name })), @@ -100,6 +107,10 @@ export function useTaskContextMenu() { case "restore": await restoreTask(task.id); break; + case "stop": { + onStop?.(task.id, task.title, runId); + break; + } case "archive": if (onArchive) { onArchive(task.id); diff --git a/packages/ui/src/shell/ContentHeader.tsx b/packages/ui/src/shell/ContentHeader.tsx index c66e3728fa..f52888a09b 100644 --- a/packages/ui/src/shell/ContentHeader.tsx +++ b/packages/ui/src/shell/ContentHeader.tsx @@ -14,6 +14,7 @@ import { BranchSelector } from "@posthog/ui/features/git-interaction/components/ import { CloudGitInteractionHeader } from "@posthog/ui/features/git-interaction/components/CloudGitInteractionHeader"; import { TaskActionsMenu } from "@posthog/ui/features/git-interaction/components/TaskActionsMenu"; import { HandoffConfirmDialog } from "@posthog/ui/features/sessions/components/HandoffConfirmDialog"; +import { StopCloudRunButton } from "@posthog/ui/features/sessions/components/StopCloudRunButton"; import { useHandoffDialogStore } from "@posthog/ui/features/sessions/handoffDialogStore"; import { useSessionCallbacks } from "@posthog/ui/features/sessions/hooks/useSessionCallbacks"; import { useSessionForTask } from "@posthog/ui/features/sessions/useSession"; @@ -197,10 +198,13 @@ export function ContentHeader() { {isCloudTask ? ( - + <> + + + ) : ( )}