Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion apps/web/src/web-container.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
96 changes: 60 additions & 36 deletions packages/core/src/cloud-task/cloud-task-engine.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -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");
}

Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -959,7 +984,6 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
}
}

@preDestroy()
unwatchAll(): void {
for (const key of [...this.watchers.keys()]) {
this.stopWatcher(key);
Expand Down Expand Up @@ -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,
Expand Down
27 changes: 27 additions & 0 deletions packages/core/src/cloud-task/cloud-task-service.test.ts
Original file line number Diff line number Diff line change
@@ -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");
});
});
2 changes: 1 addition & 1 deletion packages/core/src/cloud-task/cloud-task.module.ts
Original file line number Diff line number Diff line change
@@ -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 }) => {
Expand Down
47 changes: 27 additions & 20 deletions packages/core/src/cloud-task/cloud-task.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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(),
Expand Down Expand Up @@ -86,8 +92,8 @@ async function waitFor(
}
}

describe("CloudTaskService", () => {
let service: CloudTaskService;
describe("CloudTaskEngine", () => {
let service: CloudTaskEngine;

beforeEach(() => {
const scopedLog = {
Expand All @@ -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();
Expand Down Expand Up @@ -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>;
Expand All @@ -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();
Expand Down
35 changes: 35 additions & 0 deletions packages/core/src/cloud-task/cloud-task.ts
Original file line number Diff line number Diff line change
@@ -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();
}
}
22 changes: 7 additions & 15 deletions packages/core/src/cloud-task/schemas.ts
Original file line number Diff line number Diff line change
@@ -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 ---

Expand Down
2 changes: 1 addition & 1 deletion packages/core/src/handoff/handoff.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
23 changes: 23 additions & 0 deletions packages/core/vitest.config.ts
Original file line number Diff line number Diff line change
@@ -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/**"],
},
});
2 changes: 1 addition & 1 deletion packages/host-router/src/routers/cloud-task.router.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
Loading