refactor(core): extract portable cloud task engine

Generated-By: PostHog Code
Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22

Author: richardsolomou

apps/web/src/web-container.ts

@@ -27,8 +27,8 @@ import {
 } from "@posthog/core/auth/identifiers";
 import { canvasCoreModule } from "@posthog/core/canvas/canvas.module";
 import { taskThreadCoreModule } from "@posthog/core/canvas/taskThread.module";
+import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task";
 import { cloudTaskModule } from "@posthog/core/cloud-task/cloud-task.module";
-import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task-engine";
 import {
   CLOUD_TASK_AUTH,
   CLOUD_TASK_SERVICE,

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

@@ -1,37 +1,24 @@
+import type { RootLogger, ScopedLogger } from "@posthog/di/logger";
+import type { IAnalytics } from "@posthog/platform/analytics";
 import {
-  ROOT_LOGGER,
-  type RootLogger,
-  type ScopedLogger,
-} from "@posthog/di/logger";
-import {
-  ANALYTICS_SERVICE,
-  type IAnalytics,
-} from "@posthog/platform/analytics";
-import type { StoredLogEntry } from "@posthog/shared";
-import {
+  type CloudTaskPermissionRequestUpdate,
+  isTerminalStatus,
   mcpToolKey,
   posthogToolMeta,
+  type StoredLogEntry,
   serializeError,
+  type TaskRunStatus,
   TypedEventEmitter,
 } from "@posthog/shared";
 import { ANALYTICS_EVENTS } from "@posthog/shared/analytics-events";
-import { inject, injectable, optional, preDestroy } from "inversify";
-import type { CloudTaskPermissionRequestUpdate } from "./cloud-task-types";
-import {
-  CLOUD_TASK_AUTH,
-  type ICloudTaskAuth,
-  MCP_RELAY_EXECUTOR,
-  type McpRelayExecutor,
-} from "./identifiers";
+import type { ICloudTaskAuth, McpRelayExecutor } from "./identifiers";
 import {
   CloudTaskEvent,
   type CloudTaskEvents,
-  isTerminalStatus,
   type SendCommandInput,
   type SendCommandOutput,
   type StopInput,
   type StopOutput,
-  type TaskRunStatus,
   type WatchInput,
 } from "./schemas";
 import { type SseEvent, SseEventParser } from "./sse-parser";
@@ -435,23 +422,45 @@ function sandboxAlivePayload(watcher: { lastSandboxAlive: boolean | null }): {
     : { sandboxAlive: watcher.lastSandboxAlive };
 }
 
-@injectable()
-export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
+export interface CloudTaskEngineDependencies {
+  auth: ICloudTaskAuth;
+  analytics: IAnalytics;
+  logger: RootLogger;
+  mcpRelayExecutor?: McpRelayExecutor | null;
+  streamFetch?: CloudTaskFetch;
+}
+
+export type CloudTaskFetch = (
+  input: string | URL | Request,
+  init?: RequestInit,
+) => Promise<Response>;
+
+export function createCloudTaskEngine(
+  dependencies: CloudTaskEngineDependencies,
+): CloudTaskEngine {
+  return new CloudTaskEngine(dependencies);
+}
+
+export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
   private watchers = new Map<string, WatcherState>();
   private readonly log: ScopedLogger;
-
-  constructor(
-    @inject(CLOUD_TASK_AUTH)
-    private readonly auth: ICloudTaskAuth,
-    @inject(ANALYTICS_SERVICE)
-    private readonly analytics: IAnalytics,
-    @inject(ROOT_LOGGER)
-    logger: RootLogger,
-    @inject(MCP_RELAY_EXECUTOR)
-    @optional()
-    private readonly mcpRelayExecutor: McpRelayExecutor | null = null,
-  ) {
+  private readonly auth: ICloudTaskAuth;
+  private readonly analytics: IAnalytics;
+  private readonly mcpRelayExecutor: McpRelayExecutor | null;
+  private readonly streamFetch: CloudTaskFetch;
+
+  constructor({
+    auth,
+    analytics,
+    logger,
+    mcpRelayExecutor = null,
+    streamFetch = globalThis.fetch.bind(globalThis),
+  }: CloudTaskEngineDependencies) {
     super();
+    this.auth = auth;
+    this.analytics = analytics;
+    this.mcpRelayExecutor = mcpRelayExecutor;
+    this.streamFetch = streamFetch;
     this.log = logger.scope("cloud-task");
   }
 
@@ -770,6 +779,22 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
     void this.bootstrapWatcher(key);
   }
 
+  reconnectIfDisconnected(taskId: string, runId: string): void {
+    const key = watcherKey(taskId, runId);
+    const watcher = this.watchers.get(key);
+    if (
+      !watcher ||
+      watcher.sseAbortController ||
+      watcher.reconnectTimeoutId ||
+      watcher.isBootstrapping ||
+      isTerminalStatus(watcher.lastStatus)
+    ) {
+      return;
+    }
+
+    void this.connectSse(key);
+  }
+
   // Resets a watcher to its pre-bootstrap state so bootstrapWatcher can rebuild it from server truth.
   private resetWatcherForRebootstrap(watcher: WatcherState): void {
     watcher.reconnectAttempts = 0;
@@ -959,7 +984,6 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
     }
   }
 
-  @preDestroy()
   unwatchAll(): void {
     for (const key of [...this.watchers.keys()]) {
       this.stopWatcher(key);
@@ -1306,7 +1330,7 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
     try {
       // The proxy authenticates with the run-scoped Bearer token; the Django leg uses the session.
       const response = usingProxy
-        ? await fetch(url.toString(), {
+        ? await this.streamFetch(url.toString(), {
             method: "GET",
             headers,
             signal: controller.signal,

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

@@ -0,0 +1,27 @@
+import { describe, expect, it, vi } from "vitest";
+import { CloudTaskService } from "./cloud-task";
+import { CloudTaskEngine } from "./cloud-task-engine";
+
+describe("CloudTaskService", () => {
+  it("preserves the injectable service API as a thin engine wrapper", () => {
+    const scopedLog = {
+      debug: vi.fn(),
+      info: vi.fn(),
+      warn: vi.fn(),
+      error: vi.fn(),
+    };
+    const service = new CloudTaskService(
+      {
+        authenticatedFetch: vi.fn(),
+        getCloudContext: vi.fn(),
+      },
+      { track: vi.fn() } as never,
+      { ...scopedLog, scope: vi.fn(() => scopedLog) },
+    );
+
+    expect(service).toBeInstanceOf(CloudTaskEngine);
+    expect(service.watch).toBeTypeOf("function");
+    expect(service.retry).toBeTypeOf("function");
+    expect(service.unwatchAll).toBeTypeOf("function");
+  });
+});

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

@@ -1,5 +1,5 @@
 import { ContainerModule } from "inversify";
-import { CloudTaskService } from "./cloud-task-engine";
+import { CloudTaskService } from "./cloud-task";
 import { CLOUD_TASK_SERVICE } from "./identifiers";
 
 export const cloudTaskModule = new ContainerModule(({ bind }) => {

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

@@ -5,14 +5,17 @@ const mockNetFetch = vi.hoisted(() => vi.fn());
 const mockStreamFetch = vi.hoisted(() => vi.fn());
 const mockStreamTokenFetch = vi.hoisted(() => vi.fn());
 
-// The service now uses global fetch for BOTH authenticated API calls (JSON)
-// and SSE streaming. The two used to be distinct (net.fetch vs global fetch).
 // Route by URL: /stream_token/ → token mock (read-leg resolution), the stream leg
 // (Django /stream/ or proxy /v1/runs/:run/stream) → stream mock, everything else → API mock.
 // The token mock has a Django-path default so existing fixtures (which never set it) are untouched.
 const fetchRouter = vi.hoisted(() =>
-  vi.fn((input: string | Request, init?: RequestInit) => {
-    const url = typeof input === "string" ? input : input.url;
+  vi.fn((input: string | URL | Request, init?: RequestInit) => {
+    const url =
+      typeof input === "string"
+        ? input
+        : input instanceof URL
+          ? input.toString()
+          : input.url;
     const impl = url.includes("/stream_token/")
       ? mockStreamTokenFetch
       : /\/stream(\/|\?|$)/.test(url)
@@ -22,7 +25,10 @@ const fetchRouter = vi.hoisted(() =>
   }),
 );
 
-import { CloudTaskService } from "./cloud-task-engine";
+import {
+  type CloudTaskEngine,
+  createCloudTaskEngine,
+} from "./cloud-task-engine";
 
 const mockAuthService = {
   authenticatedFetch: vi.fn(),
@@ -86,8 +92,8 @@ async function waitFor(
   }
 }
 
-describe("CloudTaskService", () => {
-  let service: CloudTaskService;
+describe("CloudTaskEngine", () => {
+  let service: CloudTaskEngine;
 
   beforeEach(() => {
     const scopedLog = {
@@ -98,11 +104,12 @@ describe("CloudTaskService", () => {
     };
     const loggerMock = { ...scopedLog, scope: vi.fn(() => scopedLog) };
     const analyticsMock = { track: vi.fn() };
-    service = new CloudTaskService(
-      mockAuthService as never,
-      analyticsMock as never,
-      loggerMock,
-    );
+    service = createCloudTaskEngine({
+      auth: mockAuthService as never,
+      analytics: analyticsMock as never,
+      logger: loggerMock,
+      streamFetch: fetchRouter,
+    });
     mockNetFetch.mockReset();
     mockStreamFetch.mockReset();
     mockStreamTokenFetch.mockReset();
@@ -3077,8 +3084,8 @@ describe("CloudTaskService", () => {
   });
 });
 
-describe("CloudTaskService MCP relay", () => {
-  let relayService: CloudTaskService;
+describe("CloudTaskEngine MCP relay", () => {
+  let relayService: CloudTaskEngine;
   let mcpRelayExecutor: {
     execute: ReturnType<typeof vi.fn>;
     closeRun: ReturnType<typeof vi.fn>;
@@ -3099,12 +3106,12 @@ describe("CloudTaskService MCP relay", () => {
       })),
       closeRun: vi.fn(async () => {}),
     };
-    relayService = new CloudTaskService(
-      mockAuthService as never,
-      analyticsMock as never,
-      loggerMock,
-      mcpRelayExecutor as never,
-    );
+    relayService = createCloudTaskEngine({
+      auth: mockAuthService as never,
+      analytics: analyticsMock as never,
+      logger: loggerMock,
+      mcpRelayExecutor: mcpRelayExecutor as never,
+    });
 
     mockNetFetch.mockReset();
     mockStreamFetch.mockReset();

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

@@ -0,0 +1,35 @@
+import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger";
+import {
+  ANALYTICS_SERVICE,
+  type IAnalytics,
+} from "@posthog/platform/analytics";
+import { inject, injectable, optional, preDestroy } from "inversify";
+import { CloudTaskEngine } from "./cloud-task-engine";
+import {
+  CLOUD_TASK_AUTH,
+  type ICloudTaskAuth,
+  MCP_RELAY_EXECUTOR,
+  type McpRelayExecutor,
+} from "./identifiers";
+
+@injectable()
+export class CloudTaskService extends CloudTaskEngine {
+  constructor(
+    @inject(CLOUD_TASK_AUTH)
+    auth: ICloudTaskAuth,
+    @inject(ANALYTICS_SERVICE)
+    analytics: IAnalytics,
+    @inject(ROOT_LOGGER)
+    logger: RootLogger,
+    @inject(MCP_RELAY_EXECUTOR)
+    @optional()
+    mcpRelayExecutor: McpRelayExecutor | null = null,
+  ) {
+    super({ auth, analytics, logger, mcpRelayExecutor });
+  }
+
+  @preDestroy()
+  override unwatchAll(): void {
+    super.unwatchAll();
+  }
+}

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

@@ -1,20 +1,12 @@
-import type { TaskRunStatus } from "@posthog/shared";
+import type { CloudTaskUpdatePayload } from "@posthog/shared";
 import { z } from "zod";
-import type { CloudTaskUpdatePayload } from "./cloud-task-types";
 
-export type { CloudTaskUpdatePayload, TaskRunStatus };
-
-export const TERMINAL_STATUSES = ["completed", "failed", "cancelled"] as const;
-
-export function isTerminalStatus(
-  status: TaskRunStatus | string | null | undefined,
-): boolean {
-  return (
-    status !== null &&
-    status !== undefined &&
-    TERMINAL_STATUSES.includes(status as (typeof TERMINAL_STATUSES)[number])
-  );
-}
+export {
+  type CloudTaskUpdatePayload,
+  isTerminalStatus,
+  type TaskRunStatus,
+  TERMINAL_STATUSES,
+} from "@posthog/shared";
 
 // --- Events ---
 

packages/core/src/handoff/handoff.ts

@@ -5,7 +5,7 @@ import {
   TypedEventEmitter,
 } from "@posthog/shared";
 import { inject, injectable } from "inversify";
-import type { CloudTaskService } from "../cloud-task/cloud-task-engine";
+import type { CloudTaskService } from "../cloud-task/cloud-task";
 import { CLOUD_TASK_SERVICE } from "../cloud-task/identifiers";
 import { HandoffSaga, type HandoffSagaDeps } from "./handoff-saga";
 import {

packages/core/vitest.config.ts

@@ -0,0 +1,23 @@
+import { defineConfig } from "vitest/config";
+import { trunkTestOptions } from "../../vitest.config.base";
+
+export default defineConfig({
+  oxc: false,
+  esbuild: {
+    tsconfigRaw: {
+      compilerOptions: {
+        experimentalDecorators: true,
+        target: "ES2022",
+        useDefineForClassFields: false,
+        verbatimModuleSyntax: true,
+      },
+    },
+  },
+  test: {
+    globals: true,
+    ...trunkTestOptions,
+    environment: "node",
+    include: ["src/**/*.test.ts", "src/**/*.test.tsx"],
+    exclude: ["**/node_modules/**", "**/dist/**"],
+  },
+});

packages/host-router/src/routers/cloud-task.router.ts

@@ -1,4 +1,4 @@
-import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task-engine";
+import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task";
 import { CLOUD_TASK_SERVICE } from "@posthog/core/cloud-task/identifiers";
 import {
   CloudTaskEvent,