diff --git a/apps/web/src/web-container.ts b/apps/web/src/web-container.ts index cf17cd493e..3075d139ce 100644 --- a/apps/web/src/web-container.ts +++ b/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, diff --git a/packages/core/src/cloud-task/cloud-task-engine.ts b/packages/core/src/cloud-task/cloud-task-engine.ts index 22b5a1ddab..006b0ab6d2 100644 --- a/packages/core/src/cloud-task/cloud-task-engine.ts +++ b/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 { +export interface CloudTaskEngineDependencies { + auth: ICloudTaskAuth; + analytics: IAnalytics; + logger: RootLogger; + mcpRelayExecutor?: McpRelayExecutor | null; + streamFetch?: CloudTaskFetch; +} + +export type CloudTaskFetch = ( + input: string | URL | Request, + init?: RequestInit, +) => Promise; + +export function createCloudTaskEngine( + dependencies: CloudTaskEngineDependencies, +): CloudTaskEngine { + return new CloudTaskEngine(dependencies); +} + +export class CloudTaskEngine extends TypedEventEmitter { private watchers = new Map(); 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 { 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 { } } - @preDestroy() unwatchAll(): void { for (const key of [...this.watchers.keys()]) { this.stopWatcher(key); @@ -1306,7 +1330,7 @@ export class CloudTaskService extends TypedEventEmitter { 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, diff --git a/packages/core/src/cloud-task/cloud-task-service.test.ts b/packages/core/src/cloud-task/cloud-task-service.test.ts new file mode 100644 index 0000000000..644f9dd7cb --- /dev/null +++ b/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"); + }); +}); diff --git a/packages/core/src/cloud-task/cloud-task.module.ts b/packages/core/src/cloud-task/cloud-task.module.ts index 464011d14f..02f0cc91db 100644 --- a/packages/core/src/cloud-task/cloud-task.module.ts +++ b/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 }) => { diff --git a/packages/core/src/cloud-task/cloud-task.test.ts b/packages/core/src/cloud-task/cloud-task.test.ts index 9002f948ff..4178bd3d78 100644 --- a/packages/core/src/cloud-task/cloud-task.test.ts +++ b/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; closeRun: ReturnType; @@ -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(); diff --git a/packages/core/src/cloud-task/cloud-task.ts b/packages/core/src/cloud-task/cloud-task.ts new file mode 100644 index 0000000000..1e2003d85a --- /dev/null +++ b/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(); + } +} diff --git a/packages/core/src/cloud-task/schemas.ts b/packages/core/src/cloud-task/schemas.ts index d694e52141..b8c03eb202 100644 --- a/packages/core/src/cloud-task/schemas.ts +++ b/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 --- diff --git a/packages/core/src/handoff/handoff.ts b/packages/core/src/handoff/handoff.ts index 2b8d477099..c85c2fcfa8 100644 --- a/packages/core/src/handoff/handoff.ts +++ b/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 { diff --git a/packages/core/vitest.config.ts b/packages/core/vitest.config.ts new file mode 100644 index 0000000000..bed14b24e0 --- /dev/null +++ b/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/**"], + }, +}); diff --git a/packages/host-router/src/routers/cloud-task.router.ts b/packages/host-router/src/routers/cloud-task.router.ts index 4546995404..15d577ce59 100644 --- a/packages/host-router/src/routers/cloud-task.router.ts +++ b/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,