diff --git a/apps/code/electron-builder.ts b/apps/code/electron-builder.ts index d124f10bba..d0fd293eb3 100644 --- a/apps/code/electron-builder.ts +++ b/apps/code/electron-builder.ts @@ -49,6 +49,8 @@ const config: Configuration = { ".vite/build/plugins/posthog/**", ".vite/build/codex-acp/**", ".vite/build/grammars/**", + ".vite/build/rpc-host.js", + ".vite/build/rpc-host.js.map", ...asarUnpackGlobs, ], diff --git a/apps/code/electron.vite.config.ts b/apps/code/electron.vite.config.ts index 2e8172dd17..16bce4a5c2 100644 --- a/apps/code/electron.vite.config.ts +++ b/apps/code/electron.vite.config.ts @@ -22,6 +22,7 @@ import { copyCodexAcpBinaries, copyDrizzleMigrations, copyEnricherGrammars, + copyPiRpcHost, copyPosthogPlugin, fixFilenameCircularRef, getBuildDate, @@ -93,6 +94,7 @@ export default defineConfig(({ mode }) => { autoServicesPlugin(path.join(__dirname, "src/main/services")), fixFilenameCircularRef(), copyClaudeExecutable(), + copyPiRpcHost(), copyPosthogPlugin(isDev), copyDrizzleMigrations(), copyCodexAcpBinaries(), diff --git a/apps/code/package.json b/apps/code/package.json index a0a50e5b15..37961dcd6e 100644 --- a/apps/code/package.json +++ b/apps/code/package.json @@ -152,6 +152,7 @@ "reflect-metadata": "^0.2.2", "semver": "^7.8.1", "shadcn": "^4.1.2", + "superjson": "catalog:", "smol-toml": "^1.6.1", "tailwindcss-scroll-mask": "^0.0.3", "tw-animate-css": "^1.4.0", diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index e620dce39d..c121b765d8 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -184,6 +184,9 @@ import { OAUTH_CALLBACK_SERVER } from "@posthog/workspace-server/services/oauth- import { oauthCallbackModule } from "@posthog/workspace-server/services/oauth-callback/oauth-callback.module"; import { onboardingImportModule } from "@posthog/workspace-server/services/onboarding-import/onboarding-import.module"; import { osModule } from "@posthog/workspace-server/services/os/os.module"; +import { PI_SESSION_SERVICE } from "@posthog/workspace-server/services/pi-session/identifiers"; +import type { PiSessionService } from "@posthog/workspace-server/services/pi-session/pi-session"; +import { piSessionModule } from "@posthog/workspace-server/services/pi-session/pi-session.module"; import { POSTHOG_PLUGIN_SERVICE } from "@posthog/workspace-server/services/posthog-plugin/identifiers"; import { posthogPluginModule } from "@posthog/workspace-server/services/posthog-plugin/posthog-plugin.module"; import { PROCESS_TRACKING_SERVICE } from "@posthog/workspace-server/services/process-tracking/identifiers"; @@ -356,6 +359,7 @@ container .bind(MAIN_DEFAULT_ADDITIONAL_DIRECTORY_REPOSITORY) .toService(DEFAULT_ADDITIONAL_DIRECTORY_REPOSITORY); container.load(agentModule); +container.load(piSessionModule); container.bind(AGENT_SLEEP_COORDINATOR).toService(MAIN_SLEEP_SERVICE); container.bind(AGENT_MCP_APPS).toService(MCP_APPS_SERVICE); container.bind(AGENT_REPO_FILES).toService(MAIN_FS_SERVICE); @@ -391,8 +395,12 @@ container.bind(MCP_PROXY_AUTH).toDynamicValue((ctx) => { }); container.load(archiveModule); container.bind(ARCHIVE_SESSION_CANCELLER).toDynamicValue((ctx) => ({ - cancelSessionsByTaskId: (taskId: string) => - ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), + cancelSessionsByTaskId: async (taskId: string) => { + await Promise.all([ + ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), + ctx.get(PI_SESSION_SERVICE).stop(taskId), + ]); + }, })); container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({ stopWatching: async (worktreePath: string) => { @@ -403,8 +411,12 @@ container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({ })); container.load(suspensionModule); container.bind(SUSPENSION_SESSION_CANCELLER).toDynamicValue((ctx) => ({ - cancelSessionsByTaskId: (taskId: string) => - ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), + cancelSessionsByTaskId: async (taskId: string) => { + await Promise.all([ + ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), + ctx.get(PI_SESSION_SERVICE).stop(taskId), + ]); + }, })); container.bind(SUSPENSION_FILE_WATCHER).toDynamicValue((ctx) => ({ stopWatching: async (worktreePath: string) => { @@ -675,7 +687,12 @@ container.load(workspaceModule); container.bind(WORKSPACE_AGENT).toDynamicValue((ctx): WorkspaceAgent => { const agent = ctx.get(AGENT_SERVICE); return { - cancelSessionsByTaskId: (taskId) => agent.cancelSessionsByTaskId(taskId), + cancelSessionsByTaskId: async (taskId) => { + await Promise.all([ + agent.cancelSessionsByTaskId(taskId), + ctx.get(PI_SESSION_SERVICE).stop(taskId), + ]); + }, onAgentFileActivity: (handler) => agent.on(AgentServiceEvent.AgentFileActivity, handler), }; diff --git a/apps/code/src/main/trpc/router.ts b/apps/code/src/main/trpc/router.ts index c6bd789ffe..aca22ec9af 100644 --- a/apps/code/src/main/trpc/router.ts +++ b/apps/code/src/main/trpc/router.ts @@ -35,6 +35,7 @@ import { notificationRouter } from "@posthog/host-router/routers/notification.ro import { oauthRouter } from "@posthog/host-router/routers/oauth.router"; import { onboardingImportRouter } from "@posthog/host-router/routers/onboarding-import.router"; import { osRouter } from "@posthog/host-router/routers/os.router"; +import { piSessionRouter } from "@posthog/host-router/routers/pi-session.router"; import { processTrackingRouter } from "@posthog/host-router/routers/process-tracking.router"; import { provisioningRouter } from "@posthog/host-router/routers/provisioning.router"; import { secureStoreRouter } from "@posthog/host-router/routers/secure-store.router"; @@ -93,6 +94,7 @@ export const trpcRouter = router({ onboardingImport: onboardingImportRouter, logs: logsRouter, os: osRouter, + piSession: piSessionRouter, processTracking: processTrackingRouter, provisioning: provisioningRouter, sleep: sleepRouter, diff --git a/apps/code/src/renderer/di/bindings.ts b/apps/code/src/renderer/di/bindings.ts index dbdca0c7df..cece5bfd31 100644 --- a/apps/code/src/renderer/di/bindings.ts +++ b/apps/code/src/renderer/di/bindings.ts @@ -65,6 +65,8 @@ import { GITHUB_CONNECT_CLIENT, type GithubConnectClient, } from "@posthog/core/onboarding/identifiers"; +import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers"; +import type { PiRunner } from "@posthog/core/pi-runtime/piRunner"; import { type BundleLocalSkill, CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL, @@ -276,6 +278,7 @@ export interface RendererBindings { [SHELL_PROCESS_READER]: ShellProcessReader; [ANALYTICS_TRACKER]: AnalyticsTracker; [TASK_CREATION_HOST]: ITaskCreationHost; + [PI_RUNNER]: PiRunner; [TASK_CREATION_EFFECTS]: TaskCreationEffects; [RENDERER_TASK_SERVICE]: TaskService; [TASK_SERVICE]: TaskService; diff --git a/apps/code/src/renderer/di/container.ts b/apps/code/src/renderer/di/container.ts index 807f63b398..b691dfc103 100644 --- a/apps/code/src/renderer/di/container.ts +++ b/apps/code/src/renderer/di/container.ts @@ -30,6 +30,8 @@ import { import { LLM_GATEWAY_SERVICE } from "@posthog/core/llm-gateway/identifiers"; import type { LlmGatewayService } from "@posthog/core/llm-gateway/llm-gateway"; import type { LlmMessage } from "@posthog/core/llm-gateway/schemas"; +import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers"; +import type { PiRunner } from "@posthog/core/pi-runtime/piRunner"; import { CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL, CLOUD_ARTIFACT_READ_FILE_AS_BASE64, @@ -149,6 +151,7 @@ import { trpcClient } from "@renderer/trpc"; import { hostTrpcClient } from "@renderer/trpc/client"; import type { TRPCClient } from "@trpc/client"; import { hostLog, logger } from "@utils/logger"; +import { TrpcPiRunner } from "../platform-adapters/trpc-pi-runner"; import type { RendererBindings } from "./bindings"; import { TASK_SERVICE as RENDERER_TASK_SERVICE, TRPC_CLIENT } from "./tokens"; @@ -288,6 +291,7 @@ container // Bind services container.bind(TASK_CREATION_HOST).to(TrpcTaskCreationHost); +container.bind(PI_RUNNER).to(TrpcPiRunner); container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects); container.bind(RENDERER_TASK_SERVICE).to(TaskService); container.bind(TASK_SERVICE).toService(RENDERER_TASK_SERVICE); diff --git a/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts b/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts new file mode 100644 index 0000000000..3f4311f7fb --- /dev/null +++ b/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts @@ -0,0 +1,28 @@ +import type { + PiResumeInput, + PiRunInput, + PiRunner, +} from "@posthog/core/pi-runtime/piRunner"; +import { resolveService } from "@posthog/di/container"; +import { + HOST_TRPC_CLIENT, + type HostTrpcClient, +} from "@posthog/host-router/client"; + +function hostClient(): HostTrpcClient { + return resolveService(HOST_TRPC_CLIENT); +} + +export class TrpcPiRunner implements PiRunner { + async create(input: PiRunInput): Promise { + await hostClient().piSession.start.mutate(input); + } + + resume(input: PiResumeInput): Promise { + return hostClient().piSession.resume.mutate(input); + } + + stop(taskId: string): Promise { + return hostClient().piSession.stop.mutate({ taskId }); + } +} diff --git a/apps/code/src/renderer/trpc/client.ts b/apps/code/src/renderer/trpc/client.ts index 34151a047e..8cc86cdc64 100644 --- a/apps/code/src/renderer/trpc/client.ts +++ b/apps/code/src/renderer/trpc/client.ts @@ -7,14 +7,18 @@ import { createTRPCOptionsProxy, } from "@trpc/tanstack-react-query"; import { queryClient } from "@utils/queryClient"; +import superjson from "superjson"; import type { TrpcRouter } from "../../main/trpc/router"; export const trpcClient = createTRPCClient({ - links: [ipcInstrumentationLink(), ipcLink()], + links: [ + ipcInstrumentationLink(), + ipcLink({ transformer: superjson }), + ], }); export const hostTrpcClient = createTRPCClient({ - links: [ipcLink()], + links: [ipcLink({ transformer: superjson })], }); const context = createTRPCContext(); diff --git a/apps/code/vite-main-plugins.mts b/apps/code/vite-main-plugins.mts index f83bea74b7..8d3d23e295 100644 --- a/apps/code/vite-main-plugins.mts +++ b/apps/code/vite-main-plugins.mts @@ -163,6 +163,29 @@ function copyClaudeSupportAssets(sourcePath: string, destDir: string): void { } } +export function copyPiRpcHost(): Plugin { + return { + name: "copy-pi-rpc-host", + writeBundle() { + const candidates = [ + join(__dirname, "node_modules/@posthog/agent/dist/pi/rpc-host.js"), + join( + __dirname, + "../../node_modules/@posthog/agent/dist/pi/rpc-host.js", + ), + join(__dirname, "../../packages/agent/dist/pi/rpc-host.js"), + ]; + const source = candidates.find((candidate) => existsSync(candidate)); + if (!source) { + throw new Error( + `[copy-pi-rpc-host] Unable to find Pi RPC host, required at runtime by createPiRpcClient. Build @posthog/agent first. Checked:\n ${candidates.join("\n ")}`, + ); + } + copyFileSync(source, join(__dirname, ".vite/build/rpc-host.js")); + }, + }; +} + export function copyClaudeExecutable(): Plugin { return { name: "copy-claude-executable", diff --git a/apps/web/package.json b/apps/web/package.json index d06a50e056..d30e7da97f 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -20,6 +20,7 @@ "@posthog/workspace-client": "workspace:*", "@tanstack/react-query": "^5.100.14", "@trpc/client": "^11.17.0", + "superjson": "catalog:", "@trpc/tanstack-react-query": "^11.17.0", "inversify": "^7.10.6", "react": "19.2.6", diff --git a/apps/web/src/web-trpc.ts b/apps/web/src/web-trpc.ts index c5ffce58ac..b4e7d36a83 100644 --- a/apps/web/src/web-trpc.ts +++ b/apps/web/src/web-trpc.ts @@ -5,6 +5,7 @@ import { httpSubscriptionLink, splitLink, } from "@trpc/client"; +import superjson from "superjson"; // The ENTIRE electron->web transport difference. The renderer builds the same // client with `links: [ipcLink()]`; web swaps in HTTP: httpBatchLink for @@ -16,8 +17,8 @@ export const hostTrpcClient = createTRPCClient({ links: [ splitLink({ condition: (op) => op.type === "subscription", - true: httpSubscriptionLink({ url: API_URL }), - false: httpBatchLink({ url: API_URL }), + true: httpSubscriptionLink({ url: API_URL, transformer: superjson }), + false: httpBatchLink({ url: API_URL, transformer: superjson }), }), ], }); diff --git a/packages/agent/package.json b/packages/agent/package.json index 76f04bf751..2cbd4714e4 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -28,6 +28,10 @@ "types": "./dist/posthog-products.d.ts", "import": "./dist/posthog-products.js" }, + "./pi/rpc-client": { + "types": "./dist/pi/rpc-client.d.ts", + "import": "./dist/pi/rpc-client.js" + }, "./pr-url-detector": { "types": "./dist/pr-url-detector.d.ts", "import": "./dist/pr-url-detector.js" diff --git a/packages/agent/src/pi/rpc-client.test.ts b/packages/agent/src/pi/rpc-client.test.ts index b653484d42..a7d7d4f4d6 100644 --- a/packages/agent/src/pi/rpc-client.test.ts +++ b/packages/agent/src/pi/rpc-client.test.ts @@ -7,7 +7,12 @@ describe("createPiRpcClient", () => { const client = createPiRpcClient({ cwd: "/workspace", model: "claude-opus-4-8", - apiKey: "token", + oauthCredentials: { + access: "access-token", + refresh: "refresh-token", + expires: 0, + region: "us", + }, env: { POSTHOG_REGION: "us" }, }); @@ -17,7 +22,12 @@ describe("createPiRpcClient", () => { cwd: "/workspace", model: "claude-opus-4-8", env: { - POSTHOG_API_KEY: "token", + POSTHOG_OAUTH_CREDENTIALS: JSON.stringify({ + access: "access-token", + refresh: "refresh-token", + expires: 0, + region: "us", + }), POSTHOG_REGION: "us", }, provider: "posthog", diff --git a/packages/agent/src/pi/rpc-client.ts b/packages/agent/src/pi/rpc-client.ts index b30a4162ff..21f2c6afc2 100644 --- a/packages/agent/src/pi/rpc-client.ts +++ b/packages/agent/src/pi/rpc-client.ts @@ -1,37 +1,49 @@ +import type { ChildProcess } from "node:child_process"; import { fileURLToPath } from "node:url"; import { RpcClient, type RpcClientOptions, } from "@earendil-works/pi-coding-agent"; +import type { PosthogOAuthCredentials } from "@posthog/harness"; export type PiRpcClient = RpcClient; +type RpcClientProcessAccess = { + process?: ChildProcess; +}; + +// HACK: Pi RpcClient does not expose its spawned child process or PID publicly. +export function getPiRpcClientProcess( + client: PiRpcClient, +): ChildProcess | null { + return (client as unknown as RpcClientProcessAccess).process ?? null; +} + export type PiRpcClientOptions = Pick< RpcClientOptions, "cwd" | "env" | "model" > & { - /** PostHog token passed only to the isolated RPC host process. */ - apiKey?: string; + sessionFile?: string; + oauthCredentials: PosthogOAuthCredentials; }; -/** - * Create a native Pi RPC client backed by the PostHog harness. - * - * The client owns an isolated child process. Call `start()` before sending - * commands and `stop()` when the run is finished. - */ -export function createPiRpcClient( - options: PiRpcClientOptions = {}, -): PiRpcClient { - const { apiKey, env, ...rpcOptions } = options; +export function createPiRpcClient(options: PiRpcClientOptions): PiRpcClient { + const { env, sessionFile, oauthCredentials, ...rpcOptions } = options; + const args = sessionFile ? ["--session-file", sessionFile] : undefined; + const cliPath = fileURLToPath(new URL("./rpc-host.js", import.meta.url)); + const oauthEnvironment: Record = { + POSTHOG_OAUTH_CREDENTIALS: JSON.stringify(oauthCredentials), + }; + const childEnvironment: Record = { + ...env, + ...oauthEnvironment, + }; return new RpcClient({ ...rpcOptions, - cliPath: fileURLToPath(new URL("./rpc-host.js", import.meta.url)), - env: { - ...env, - ...(apiKey ? { POSTHOG_API_KEY: apiKey } : {}), - }, + args, + cliPath, + env: childEnvironment, provider: "posthog", }); } diff --git a/packages/agent/src/pi/rpc-host.ts b/packages/agent/src/pi/rpc-host.ts index c98d26ce87..10702c91eb 100644 --- a/packages/agent/src/pi/rpc-host.ts +++ b/packages/agent/src/pi/rpc-host.ts @@ -1,16 +1,36 @@ -import { createHarnessRuntime, runRpcMode } from "@posthog/harness"; +import { SessionManager } from "@earendil-works/pi-coding-agent"; +import { + createHarnessRuntime, + parsePosthogOAuthCredentials, + runRpcMode, +} from "@posthog/harness"; + +const SESSION_FILE_ARGUMENT = "--session-file"; +const MODEL_ARGUMENT = "--model"; function argumentValue(name: string): string | undefined { const index = process.argv.indexOf(name); return index === -1 ? undefined : process.argv[index + 1]; } +const sessionFile = argumentValue(SESSION_FILE_ARGUMENT); +const oauthCredentials = parsePosthogOAuthCredentials( + process.env.POSTHOG_OAUTH_CREDENTIALS, +); +if (!oauthCredentials) { + throw new Error("Pi RPC host requires PostHog OAuth credentials"); +} +const cwd = process.cwd(); +const sessionManager = sessionFile + ? SessionManager.open(sessionFile, undefined, cwd) + : undefined; const runtime = await createHarnessRuntime({ - cwd: process.cwd(), - apiKey: process.env.POSTHOG_API_KEY ?? process.env.POSTHOG_PERSONAL_API_KEY, + cwd, + sessionManager, + posthogOAuthCredentials: oauthCredentials, }); -const requestedModel = argumentValue("--model")?.replace(/^posthog\//, ""); +const requestedModel = argumentValue(MODEL_ARGUMENT)?.replace(/^posthog\//, ""); if (requestedModel) { const model = runtime.services.modelRegistry.find("posthog", requestedModel); if (!model) { diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index 09321156bf..7830e651e4 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -87,6 +87,7 @@ const sharedOptions = { "@posthog/shared", "@posthog/git", "@posthog/enricher", + "@posthog/harness", "fflate", ], external: [ @@ -113,7 +114,6 @@ export default defineConfig([ "src/posthog-products.ts", "src/pr-url-detector.ts", "src/pi/rpc-client.ts", - "src/pi/rpc-host.ts", "src/resume.ts", "src/types.ts", "src/adapters/claude/questions/utils.ts", @@ -168,4 +168,16 @@ export default defineConfig([ clean: false, ...sharedOptions, }, + { + entry: { "pi/rpc-host": "src/pi/rpc-host.ts" }, + format: ["esm"], + dts: false, + clean: false, + banner: { + js: 'import { createRequire as __createRequire } from "node:module"; const require = __createRequire(import.meta.url);', + }, + ...sharedOptions, + noExternal: [/^(?!node:)/], + external: [...builtinModules, ...builtinModules.map((m) => `node:${m}`)], + }, ]); diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index 18e3882e88..a7cc4e7cf6 100644 --- a/packages/core/src/auth/auth.ts +++ b/packages/core/src/auth/auth.ts @@ -152,6 +152,22 @@ export class AuthService extends TypedEventEmitter { apiHost: getCloudUrlFromRegion(session.cloudRegion), }; } + async getOAuthCredentials(): Promise<{ + access: string; + refresh: string; + expires: number; + region: CloudRegion; + } | null> { + if (this.tokenOverride) return null; + await this.initialize(); + const session = await this.ensureValidSession(); + return { + access: session.accessToken, + refresh: session.refreshToken, + expires: session.accessTokenExpiresAt, + region: session.cloudRegion, + }; + } async refreshAccessToken(): Promise { const override = this.tokenOverride; if (override) { diff --git a/packages/core/src/inbox/reportTaskCreation.ts b/packages/core/src/inbox/reportTaskCreation.ts index ee85056652..67f5d9cb95 100644 --- a/packages/core/src/inbox/reportTaskCreation.ts +++ b/packages/core/src/inbox/reportTaskCreation.ts @@ -9,7 +9,7 @@ export interface PreviewConfigChoice { /** Minimal shape of a preview-config option we scan for the default model. */ export interface PreviewConfigOption { id?: string; - category?: string; + category?: string | null; type?: string; currentValue?: string | boolean | null; options?: PreviewConfigChoice[]; diff --git a/packages/core/src/pi-runtime/identifiers.ts b/packages/core/src/pi-runtime/identifiers.ts new file mode 100644 index 0000000000..42e0c14ab3 --- /dev/null +++ b/packages/core/src/pi-runtime/identifiers.ts @@ -0,0 +1 @@ +export const PI_RUNNER = Symbol.for("posthog.pi.runner"); diff --git a/packages/core/src/pi-runtime/piRunner.ts b/packages/core/src/pi-runtime/piRunner.ts new file mode 100644 index 0000000000..b174fb95fd --- /dev/null +++ b/packages/core/src/pi-runtime/piRunner.ts @@ -0,0 +1,17 @@ +export interface PiRunInput { + taskId: string; + cwd: string; + prompt: string; + model?: string; +} + +export interface PiResumeInput { + taskId: string; + cwd: string; +} + +export interface PiRunner { + create(input: PiRunInput): Promise; + resume(input: PiResumeInput): Promise; + stop(taskId: string): Promise; +} diff --git a/packages/core/src/task-detail/piTaskCreator.ts b/packages/core/src/task-detail/piTaskCreator.ts new file mode 100644 index 0000000000..bdae81cc6a --- /dev/null +++ b/packages/core/src/task-detail/piTaskCreator.ts @@ -0,0 +1,196 @@ +import { + Saga, + type SagaLogger, + type TaskCreationInput, + type TaskCreationOutput, + type Workspace, +} from "@posthog/shared"; +import type { Task } from "@posthog/shared/domain-types"; +import type { PiRunner } from "../pi-runtime/piRunner"; +import type { TaskCreationApiClient } from "./taskCreationApiClient"; +import type { ITaskCreationHost } from "./taskCreationHost"; +import { resolveTaskRepository } from "./taskRepository"; + +export interface PiTaskCreatorDeps { + posthogClient: TaskCreationApiClient; + host: ITaskCreationHost; + piRunner: PiRunner; + onTaskReady?: (output: TaskCreationOutput) => void; +} + +export class PiTaskCreator extends Saga { + readonly sagaName = "PiTaskCreator"; + + constructor( + private readonly deps: PiTaskCreatorDeps, + logger?: SagaLogger, + ) { + super(logger); + } + + protected async execute( + input: TaskCreationInput, + ): Promise { + if (input.workspaceMode === "cloud") { + throw new Error("Pi tasks are only supported in local workspaces"); + } + + const task = await this.createTask(input); + const repoPath = input.repoPath; + let workspace: Workspace | null = null; + + if (repoPath) { + workspace = await this.createWorkspace(task, repoPath, input); + } else if (input.allowNoRepo) { + workspace = await this.createScratchWorkspace(task); + } + + if (!workspace) { + throw new Error("Pi tasks require a workspace or scratch directory"); + } + + const cwd = workspace.worktreePath ?? workspace.folderPath; + const additionalDirectories = (input.additionalDirectories ?? []).filter( + (path) => path && path !== input.repoPath, + ); + if (additionalDirectories.length > 0) { + await this.step({ + name: "additional_directories", + execute: async () => { + await Promise.all( + additionalDirectories.map((path) => + this.deps.host.addAdditionalDirectory({ taskId: task.id, path }), + ), + ); + return { taskId: task.id, paths: additionalDirectories }; + }, + rollback: async ({ taskId, paths }) => { + await Promise.all( + paths.map((path) => + this.deps.host.removeAdditionalDirectory({ taskId, path }), + ), + ); + }, + }); + } + + await this.step({ + name: "pi_session", + execute: async () => { + await this.deps.piRunner.create({ + taskId: task.id, + cwd, + prompt: input.content ?? "", + model: input.model, + }); + return { taskId: task.id }; + }, + rollback: async ({ taskId }) => this.deps.piRunner.stop(taskId), + }); + + this.deps.onTaskReady?.({ task, workspace }); + return { task, workspace }; + } + + private async createTask(input: TaskCreationInput): Promise { + const repository = await resolveTaskRepository( + input, + this.deps.host, + this.log, + ); + + return this.step({ + name: "task_creation", + execute: async () => + (await this.deps.posthogClient.createTask({ + description: input.content ?? "", + repository: repository ?? undefined, + origin_product: input.signalReportId + ? "signal_report" + : "user_created", + signal_report: input.signalReportId ?? undefined, + channel: input.channelId ?? undefined, + })) as unknown as Task, + rollback: async (task) => this.deps.posthogClient.deleteTask(task.id), + }); + } + + private async createWorkspace( + task: Task, + repoPath: string, + input: TaskCreationInput, + ): Promise { + const folder = await this.deps.host.getFolders().then(async (folders) => { + const existing = folders.find((candidate) => candidate.path === repoPath); + return existing ?? this.deps.host.addFolder({ folderPath: repoPath }); + }); + const workspaceInfo = await this.step({ + name: "workspace_creation", + execute: () => + this.deps.host.createWorkspace({ + taskId: task.id, + mainRepoPath: repoPath, + folderId: folder.id, + folderPath: repoPath, + mode: input.workspaceMode ?? "local", + branch: input.branch ?? undefined, + allowRemoteBranchCheckout: input.allowRemoteBranchCheckout, + reuseExistingWorktree: input.reuseExistingWorktree, + }), + rollback: () => + this.deps.host.deleteWorkspace({ + taskId: task.id, + mainRepoPath: repoPath, + }), + }); + + const workspaceMode = input.workspaceMode ?? "local"; + const worktree = workspaceInfo.worktree; + if (workspaceMode === "worktree" && !worktree) { + throw new Error("Pi worktree creation did not return a worktree"); + } + if (worktree) { + return { + taskId: task.id, + folderId: folder.id, + folderPath: repoPath, + mode: workspaceMode, + worktreePath: worktree.worktreePath, + worktreeName: worktree.worktreeName, + branchName: worktree.branchName, + baseBranch: worktree.baseBranch, + linkedBranch: workspaceInfo.linkedBranch, + createdAt: worktree.createdAt, + }; + } + + return { + taskId: task.id, + folderId: folder.id, + folderPath: repoPath, + mode: "local", + worktreePath: null, + worktreeName: null, + branchName: workspaceInfo.branchName, + baseBranch: input.branch ?? null, + linkedBranch: workspaceInfo.linkedBranch, + createdAt: new Date().toISOString(), + }; + } + + private async createScratchWorkspace(task: Task): Promise { + const folderPath = await this.deps.host.ensureScratchDir(task.id); + return { + taskId: task.id, + folderId: "", + folderPath, + mode: "local", + worktreePath: null, + worktreeName: null, + branchName: null, + baseBranch: null, + linkedBranch: null, + createdAt: new Date().toISOString(), + }; + } +} diff --git a/packages/core/src/task-detail/taskCreationHost.ts b/packages/core/src/task-detail/taskCreationHost.ts index cbf55daa39..d9978341f6 100644 --- a/packages/core/src/task-detail/taskCreationHost.ts +++ b/packages/core/src/task-detail/taskCreationHost.ts @@ -1,6 +1,6 @@ import type { ContentBlock } from "@agentclientprotocol/sdk"; import type { CloudSkillBundleRef } from "@posthog/core/sessions/cloudArtifactIdentifiers"; -import type { Workspace, WorkspaceMode } from "@posthog/shared"; +import type { Workspace, WorkspaceInfo, WorkspaceMode } from "@posthog/shared"; import type { TaskCreationApiClient } from "./taskCreationApiClient"; export interface CloudPromptTransport { @@ -21,16 +21,7 @@ export interface CreateWorkspaceArgs { reuseExistingWorktree?: boolean; } -export interface CreatedWorkspaceInfo { - worktree?: { - worktreePath?: string | null; - worktreeName?: string | null; - branchName?: string | null; - baseBranch?: string | null; - createdAt?: string | null; - } | null; - linkedBranch?: string | null; -} +export type CreatedWorkspaceInfo = WorkspaceInfo; export interface TaskFolderInfo { id: string; @@ -82,6 +73,7 @@ export interface ITaskCreationHost { * task. Returns its absolute path so the agent session can start there. */ ensureScratchDir(taskId: string): Promise; + getTaskRuntime(taskId: string): Promise<"acp" | "pi">; getWorkspace(taskId: string): Promise; createWorkspace(args: CreateWorkspaceArgs): Promise; deleteWorkspace(args: { diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index f3d626053a..8c56f3a362 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/packages/core/src/task-detail/taskCreationSaga.test.ts @@ -10,6 +10,8 @@ const mockHost = vi.hoisted(() => ({ getAuthenticatedClient: vi.fn(), getTaskDirectory: vi.fn(), ensureScratchDir: vi.fn(), + startPiSession: vi.fn(), + stopPiSession: vi.fn(), getWorkspace: vi.fn(), createWorkspace: vi.fn(), deleteWorkspace: vi.fn(), @@ -35,6 +37,7 @@ const mockHost = vi.hoisted(() => ({ linkTaskBranch: vi.fn(), })); +import { PiTaskCreator } from "./piTaskCreator"; import { TaskCreationSaga } from "./taskCreationSaga"; const host = mockHost as unknown as ITaskCreationHost; @@ -338,6 +341,38 @@ describe("TaskCreationSaga", () => { ); }); + it("starts a Pi session without creating an ACP session", async () => { + const createdTask = createTask({ repository: undefined }); + const saga = new PiTaskCreator({ + posthogClient: { + createTask: vi.fn().mockResolvedValue(createdTask), + deleteTask: vi.fn(), + } as never, + host, + piRunner: { + create: mockHost.startPiSession, + stop: mockHost.stopPiSession, + } as never, + }); + + const result = await saga.run({ + content: "Draft a launch email", + workspaceMode: "local", + runtime: "pi", + model: "claude-sonnet", + allowNoRepo: true, + }); + + expect(result.success).toBe(true); + expect(mockHost.startPiSession).toHaveBeenCalledWith({ + taskId: "task-123", + cwd: "/tmp/scratch/task-123", + prompt: "Draft a launch email", + model: "claude-sonnet", + }); + expect(sessionService.connectToTask).not.toHaveBeenCalled(); + }); + it("uploads initial cloud attachments before starting the run", async () => { const createdTask = createTask(); const startedTask = createTask({ latest_run: createRun() }); diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index 1e69d53e6b..19df449923 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -24,6 +24,7 @@ import type { ImportedClaudeCliSession, ITaskCreationHost, } from "./taskCreationHost"; +import { resolveTaskRepository } from "./taskRepository"; export interface TaskCreationDeps { posthogClient: TaskCreationApiClient; @@ -717,27 +718,9 @@ export class TaskCreationSaga extends Saga< input: TaskCreationInput, warmPayload: WarmActivationPayload | null, ): Promise { - let repository = input.repository; - - const repoPathForDetection = input.repoPath; - if (!repository && repoPathForDetection) { - // Detection only fills the org/repo metadata on the task; a transient - // failure (e.g. the workspace-server transport dropping mid-call) must - // not abort task creation, so degrade to an untagged task instead. - const detected = await this.readOnlyStep("repo_detection", () => - this.deps.host - .detectRepo({ directoryPath: repoPathForDetection }) - .catch((error) => { - this.log.warn("Repo detection failed; creating task without one", { - error, - }); - return null; - }), - ); - if (detected) { - repository = `${detected.organization}/${detected.repository}`; - } - } + const repository = await this.readOnlyStep("repo_detection", () => + resolveTaskRepository(input, this.deps.host, this.log), + ); return this.step({ name: "task_creation", diff --git a/packages/core/src/task-detail/taskInput.test.ts b/packages/core/src/task-detail/taskInput.test.ts index f45243d608..22bf0c7398 100644 --- a/packages/core/src/task-detail/taskInput.test.ts +++ b/packages/core/src/task-detail/taskInput.test.ts @@ -21,6 +21,23 @@ describe("prepareTaskInput", () => { }, ); + it("defaults task creation to the ACP runtime", () => { + const input = prepareTaskInput("do the thing", [], { + workspaceMode: "local", + }); + + expect(input.runtime).toBe("acp"); + }); + + it("preserves the selected Pi runtime", () => { + const input = prepareTaskInput("do the thing", [], { + workspaceMode: "local", + runtime: "pi", + }); + + expect(input.runtime).toBe("pi"); + }); + it("drops customInstructions for cloud when none is set", () => { const input = prepareTaskInput("do the thing", [], { workspaceMode: "cloud", diff --git a/packages/core/src/task-detail/taskInput.ts b/packages/core/src/task-detail/taskInput.ts index ae47616776..602d504018 100644 --- a/packages/core/src/task-detail/taskInput.ts +++ b/packages/core/src/task-detail/taskInput.ts @@ -1,6 +1,7 @@ import { buildCloudTaskDescription } from "@posthog/core/editor/cloud-prompt"; import type { Adapter, + AgentRuntime, TaskCreationInput, WorkspaceMode, } from "@posthog/shared"; @@ -17,6 +18,7 @@ export interface PrepareTaskInputOptions { reuseExistingWorktree?: boolean; executionMode?: ExecutionMode; adapter?: Adapter; + runtime?: AgentRuntime; model?: string; reasoningLevel?: string; environmentId?: string | null; @@ -55,6 +57,7 @@ export function prepareTaskInput( reuseExistingWorktree: options.reuseExistingWorktree, executionMode: options.executionMode, adapter: options.adapter, + runtime: options.runtime ?? "acp", model: options.model, reasoningLevel: options.reasoningLevel, environmentId: options.environmentId ?? undefined, diff --git a/packages/core/src/task-detail/taskRepository.ts b/packages/core/src/task-detail/taskRepository.ts new file mode 100644 index 0000000000..ccd22a9bcf --- /dev/null +++ b/packages/core/src/task-detail/taskRepository.ts @@ -0,0 +1,26 @@ +import type { SagaLogger, TaskCreationInput } from "@posthog/shared"; +import type { ITaskCreationHost } from "./taskCreationHost"; + +export async function resolveTaskRepository( + input: TaskCreationInput, + host: Pick, + logger: Pick, +): Promise { + if (input.repository) { + return input.repository; + } + if (!input.repoPath) { + return undefined; + } + + try { + const detected = await host.detectRepo({ directoryPath: input.repoPath }); + if (!detected) { + return undefined; + } + return `${detected.organization}/${detected.repository}`; + } catch (error) { + logger.warn("Repo detection failed; creating task without one", { error }); + return undefined; + } +} diff --git a/packages/core/src/task-detail/taskService.ts b/packages/core/src/task-detail/taskService.ts index 0989046c37..a2002da7ca 100644 --- a/packages/core/src/task-detail/taskService.ts +++ b/packages/core/src/task-detail/taskService.ts @@ -11,7 +11,10 @@ import type { } from "@posthog/shared"; import type { Task } from "@posthog/shared/domain-types"; import { inject, injectable } from "inversify"; +import { PI_RUNNER } from "../pi-runtime/identifiers"; +import type { PiRunner } from "../pi-runtime/piRunner"; import { TASK_CREATION_EFFECTS, TASK_CREATION_HOST } from "./identifiers"; +import { PiTaskCreator } from "./piTaskCreator"; import type { TaskCreationEffects } from "./taskCreationEffects"; import type { ITaskCreationHost } from "./taskCreationHost"; import { TaskCreationSaga } from "./taskCreationSaga"; @@ -42,6 +45,8 @@ export class TaskService { private readonly sessionService: SessionService, @inject(TASK_CREATION_EFFECTS) private readonly effects: TaskCreationEffects, + @inject(PI_RUNNER) + private readonly piRunner: PiRunner, @inject(ROOT_LOGGER) rootLogger: RootLogger, ) { @@ -103,30 +108,35 @@ export class TaskService { } } - const saga = new TaskCreationSaga( - { - posthogClient, - host: this.host, - sessionService: this.sessionService, - track: (event, props) => this.host.track(event, props), - onTaskReady: onTaskReady - ? (output) => { - this.effects.onWorkspaceCreated(output); - this.effects.onCreateSuccess(output, input); - onTaskReady(output); - } - : undefined, - }, - this.log, - ); - - const result = await saga.run(input); + let result: CreateTaskResult; + if (input.runtime === "pi") { + const creator = new PiTaskCreator( + { + posthogClient, + host: this.host, + piRunner: this.piRunner, + onTaskReady, + }, + this.log, + ); + result = await creator.run(input); + } else { + const creator = new TaskCreationSaga( + { + posthogClient, + host: this.host, + sessionService: this.sessionService, + track: (event, props) => this.host.track(event, props), + onTaskReady, + }, + this.log, + ); + result = await creator.run(input); + } if (result.success) { this.effects.onWorkspaceCreated(result.data); - if (!onTaskReady) { - this.effects.onCreateSuccess(result.data, input); - } + this.effects.onCreateSuccess(result.data, input); } return result; @@ -147,6 +157,16 @@ export class TaskService { }; } + let runtime: "acp" | "pi" = "acp"; + try { + runtime = await this.host.getTaskRuntime(taskId); + } catch (error) { + this.log.warn("Task runtime lookup failed; assuming ACP", { + taskId, + error, + }); + } + const existingWorkspace = await this.host.getWorkspace(taskId); if (existingWorkspace) { this.log.info("Workspace already exists, fetching task only", { taskId }); @@ -159,6 +179,13 @@ export class TaskService { task.latest_run = run; } + if (runtime === "pi") { + await this.piRunner.resume({ + taskId, + cwd: existingWorkspace.worktreePath ?? existingWorkspace.folderPath, + }); + } + return { success: true, data: { @@ -176,6 +203,27 @@ export class TaskService { } } + if (runtime === "pi") { + try { + const cwd = await this.host.ensureScratchDir(taskId); + await this.piRunner.resume({ taskId, cwd }); + const task = await posthogClient.getTask(taskId); + return { + success: true, + data: { task: task as unknown as Task, workspace: null }, + }; + } catch (error) { + return { + success: false, + error: + error instanceof Error + ? error.message + : "Failed to resume Pi session", + failedStep: "pi_session", + }; + } + } + const saga = new TaskCreationSaga( { posthogClient, diff --git a/packages/harness/src/extensions/posthog-provider/provider.ts b/packages/harness/src/extensions/posthog-provider/provider.ts index 856b8834bc..3d55de913a 100644 --- a/packages/harness/src/extensions/posthog-provider/provider.ts +++ b/packages/harness/src/extensions/posthog-provider/provider.ts @@ -1,5 +1,6 @@ import type { Api, Model, OAuthCredentials } from "@earendil-works/pi-ai"; import type { + AuthStorage, ProviderConfig, ProviderModelConfig, } from "@earendil-works/pi-coding-agent"; @@ -19,6 +20,28 @@ export interface PosthogProviderOptions { apiKey?: string; } +export type PosthogOAuthCredentials = Pick< + OAuthCredentials, + "access" | "refresh" | "expires" +> & { + region: CloudRegion; +}; + +export function parsePosthogOAuthCredentials( + serialized: string | undefined, +): PosthogOAuthCredentials | null { + return serialized + ? (JSON.parse(serialized) as PosthogOAuthCredentials) + : null; +} + +export function setPosthogOAuthCredentials( + storage: AuthStorage, + credentials: PosthogOAuthCredentials, +): void { + storage.set(POSTHOG_PROVIDER_NAME, { type: "oauth", ...credentials }); +} + /** * Re-routes this provider's already-resolved models to the gateway for the * region baked into `credentials` (set at login time). Called by pi whenever diff --git a/packages/harness/src/index.ts b/packages/harness/src/index.ts index 1e9865e3af..e731d7d3b9 100644 --- a/packages/harness/src/index.ts +++ b/packages/harness/src/index.ts @@ -1,5 +1,10 @@ import type { AgentSessionRuntime } from "@earendil-works/pi-coding-agent"; +export { + type PosthogOAuthCredentials, + parsePosthogOAuthCredentials, + setPosthogOAuthCredentials, +} from "./extensions/posthog-provider/provider"; export { createHarnessRuntime, type HarnessRuntimeOptions, diff --git a/packages/harness/src/runtime.test.ts b/packages/harness/src/runtime.test.ts index 506593ad16..0ce0c8f228 100644 --- a/packages/harness/src/runtime.test.ts +++ b/packages/harness/src/runtime.test.ts @@ -1,4 +1,5 @@ -import { mkdtemp, rm } from "node:fs/promises"; +import { existsSync } from "node:fs"; +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; @@ -56,4 +57,106 @@ describe("createHarnessRuntime", () => { await runtime.dispose(); } }); + + it("keeps desktop-provided OAuth credentials in memory without touching auth.json", async () => { + vi.stubEnv("PI_OFFLINE", "1"); + const pi = await import("@earendil-works/pi-coding-agent"); + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + + const runtime = await createHarnessRuntime({ + agentDir, + cwd, + posthogOAuthCredentials: { + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + region: "us", + }, + sessionManager: pi.SessionManager.inMemory(cwd), + }); + + try { + expect(runtime.services.authStorage.get("posthog")).toMatchObject({ + type: "oauth", + access: "access-token", + refresh: "refresh-token", + }); + expect(existsSync(join(agentDir, "auth.json"))).toBe(false); + } finally { + await runtime.dispose(); + } + }); + + it("seeds the in-memory store from auth.json without writing back to it", async () => { + vi.stubEnv("PI_OFFLINE", "1"); + const pi = await import("@earendil-works/pi-coding-agent"); + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + const authPath = join(agentDir, "auth.json"); + const storedCredentials = { + anthropic: { type: "api_key", key: "anthropic-key" }, + posthog: { + type: "oauth", + access: "stale-access", + refresh: "stale-refresh", + expires: 0, + }, + }; + await writeFile(authPath, JSON.stringify(storedCredentials)); + + const runtime = await createHarnessRuntime({ + agentDir, + cwd, + posthogOAuthCredentials: { + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + region: "us", + }, + sessionManager: pi.SessionManager.inMemory(cwd), + }); + + try { + expect(runtime.services.authStorage.get("anthropic")).toMatchObject({ + type: "api_key", + key: "anthropic-key", + }); + expect(runtime.services.authStorage.get("posthog")).toMatchObject({ + access: "access-token", + refresh: "refresh-token", + }); + expect(JSON.parse(await readFile(authPath, "utf8"))).toEqual( + storedCredentials, + ); + } finally { + await runtime.dispose(); + } + }); + + it("uses file-backed auth storage when no desktop credentials are provided", async () => { + vi.stubEnv("PI_OFFLINE", "1"); + const pi = await import("@earendil-works/pi-coding-agent"); + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + + const runtime = await createHarnessRuntime({ + agentDir, + cwd, + sessionManager: pi.SessionManager.inMemory(cwd), + }); + + try { + runtime.services.authStorage.set("posthog", { + type: "oauth", + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + }); + + expect(existsSync(join(agentDir, "auth.json"))).toBe(true); + } finally { + await runtime.dispose(); + } + }); }); diff --git a/packages/harness/src/runtime.ts b/packages/harness/src/runtime.ts index bbe15307b6..7f7e3a16aa 100644 --- a/packages/harness/src/runtime.ts +++ b/packages/harness/src/runtime.ts @@ -1,17 +1,35 @@ +import { readFileSync } from "node:fs"; import { join } from "node:path"; import type { AgentSessionRuntime, + AuthStorage, CreateAgentSessionFromServicesOptions, CreateAgentSessionRuntimeFactory, CreateAgentSessionServicesOptions, } from "@earendil-works/pi-coding-agent"; import { installHogBrandEnv } from "./extensions/hog-branding/brand-env"; +import { + type PosthogOAuthCredentials, + setPosthogOAuthCredentials, +} from "./extensions/posthog-provider/provider"; import type { HarnessExtensionOptions } from "./extensions/registry"; type PiRuntimeTarget = Parameters[0]; +type AuthStorageSnapshot = Parameters[0]; -export type HarnessRuntimeOptions = HarnessExtensionOptions & - Partial< +function loadAuthStorageSnapshot( + authPath: string, +): AuthStorageSnapshot | undefined { + try { + return JSON.parse(readFileSync(authPath, "utf8")) as AuthStorageSnapshot; + } catch { + return undefined; + } +} + +export type HarnessRuntimeOptions = HarnessExtensionOptions & { + posthogOAuthCredentials?: PosthogOAuthCredentials; +} & Partial< Pick< PiRuntimeTarget, "cwd" | "agentDir" | "sessionManager" | "sessionStartEvent" @@ -35,6 +53,7 @@ export type HarnessRuntimeOptions = HarnessExtensionOptions & export async function createHarnessRuntime( options: HarnessRuntimeOptions = {}, ): Promise { + const { posthogOAuthCredentials, ...runtimeOptions } = options; // Pi reads its application branding when the SDK is first evaluated. Keep // every runtime import below dynamic so this always happens first. installHogBrandEnv(); @@ -45,8 +64,8 @@ export async function createHarnessRuntime( import("./extensions/posthog-provider/models"), ]); - const cwd = options.cwd ?? process.cwd(); - const agentDir = options.agentDir ?? pi.getAgentDir(); + const cwd = runtimeOptions.cwd ?? process.cwd(); + const agentDir = runtimeOptions.agentDir ?? pi.getAgentDir(); const createRuntime: CreateAgentSessionRuntimeFactory = async ({ cwd: runtimeCwd, @@ -54,31 +73,38 @@ export async function createHarnessRuntime( sessionManager, sessionStartEvent, }) => { + const authPath = join(runtimeAgentDir, "auth.json"); const authStorage = - options.authStorage ?? - pi.AuthStorage.create(join(runtimeAgentDir, "auth.json")); + runtimeOptions.authStorage ?? + (posthogOAuthCredentials + ? pi.AuthStorage.inMemory(loadAuthStorageSnapshot(authPath)) + : pi.AuthStorage.create(authPath)); + if (posthogOAuthCredentials) { + setPosthogOAuthCredentials(authStorage, posthogOAuthCredentials); + } const services = await pi.createAgentSessionServices({ - ...options, + ...runtimeOptions, cwd: runtimeCwd, agentDir: runtimeAgentDir, authStorage, resourceLoaderOptions: { - ...options.resourceLoaderOptions, + ...runtimeOptions.resourceLoaderOptions, extensionFactories: [ - ...(options.resourceLoaderOptions?.extensionFactories ?? []), + ...(runtimeOptions.resourceLoaderOptions?.extensionFactories ?? []), ...harnessExtensions(options), ], }, }); const created = await pi.createAgentSessionFromServices({ - ...options, + ...runtimeOptions, services, sessionManager, sessionStartEvent, model: - options.model ?? services.modelRegistry.find("posthog", DEFAULT_MODEL), + runtimeOptions.model ?? + services.modelRegistry.find("posthog", DEFAULT_MODEL), }); return { @@ -97,12 +123,12 @@ export async function createHarnessRuntime( }; const sessionManager = - options.sessionManager ?? pi.SessionManager.create(cwd); + runtimeOptions.sessionManager ?? pi.SessionManager.create(cwd); return pi.createAgentSessionRuntime(createRuntime, { cwd: sessionManager.getCwd(), agentDir, sessionManager, - sessionStartEvent: options.sessionStartEvent, + sessionStartEvent: runtimeOptions.sessionStartEvent, }); } diff --git a/packages/host-router/src/router.ts b/packages/host-router/src/router.ts index 71f787c4c2..a8e116ff9d 100644 --- a/packages/host-router/src/router.ts +++ b/packages/host-router/src/router.ts @@ -34,6 +34,7 @@ import { notificationRouter } from "./routers/notification.router"; import { oauthRouter } from "./routers/oauth.router"; import { onboardingImportRouter } from "./routers/onboarding-import.router"; import { osRouter } from "./routers/os.router"; +import { piSessionRouter } from "./routers/pi-session.router"; import { processTrackingRouter } from "./routers/process-tracking.router"; import { provisioningRouter } from "./routers/provisioning.router"; import { secureStoreRouter } from "./routers/secure-store.router"; @@ -84,6 +85,7 @@ export const hostRouter = router({ oauth: oauthRouter, onboardingImport: onboardingImportRouter, os: osRouter, + piSession: piSessionRouter, processTracking: processTrackingRouter, provisioning: provisioningRouter, secureStore: secureStoreRouter, diff --git a/packages/host-router/src/routers/pi-session.router.ts b/packages/host-router/src/routers/pi-session.router.ts new file mode 100644 index 0000000000..a45b2c86c9 --- /dev/null +++ b/packages/host-router/src/routers/pi-session.router.ts @@ -0,0 +1,69 @@ +import { publicProcedure, router } from "@posthog/host-trpc/trpc"; +import { PI_SESSION_SERVICE } from "@posthog/workspace-server/services/pi-session/identifiers"; +import type { PiSessionService } from "@posthog/workspace-server/services/pi-session/pi-session"; +import { + piSessionEntriesInput, + piSessionPromptInput, + piSessionRuntimeOutput, + piSessionStartOutput, + piSessionTranscriptInput, + resumePiSessionInput, + startPiSessionInput, +} from "@posthog/workspace-server/services/pi-session/schemas"; + +const getService = (container: { get(token: symbol): T }) => + container.get(PI_SESSION_SERVICE); + +export const piSessionRouter = router({ + start: publicProcedure + .input(startPiSessionInput) + .output(piSessionStartOutput) + .mutation(({ ctx, input }) => getService(ctx.container).start(input)), + + resume: publicProcedure + .input(resumePiSessionInput) + .mutation(({ ctx, input }) => getService(ctx.container).resume(input)), + + prompt: publicProcedure + .input(piSessionPromptInput) + .mutation(({ ctx, input }) => + getService(ctx.container).prompt(input.taskId, input.prompt), + ), + + abort: publicProcedure + .input(piSessionTranscriptInput) + .mutation(({ ctx, input }) => + getService(ctx.container).abort(input.taskId), + ), + + stop: publicProcedure + .input(piSessionTranscriptInput) + .mutation(({ ctx, input }) => getService(ctx.container).stop(input.taskId)), + + runtime: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionRuntimeOutput) + .query(({ ctx, input }) => getService(ctx.container).runtime(input.taskId)), + + status: publicProcedure + .input(piSessionTranscriptInput) + .query(({ ctx, input }) => getService(ctx.container).status(input.taskId)), + + entries: publicProcedure + .input(piSessionEntriesInput) + .query(({ ctx, input }) => + getService(ctx.container).entries(input.taskId, input.since), + ), + + onEvent: publicProcedure + .input(piSessionTranscriptInput) + .subscription(async function* (opts) { + const service = getService(opts.ctx.container); + const iterable = service.toIterable("event", { signal: opts.signal }); + for await (const payload of iterable) { + if (payload.taskId === opts.input.taskId) { + yield payload.event; + } + } + }), +}); diff --git a/packages/host-trpc/package.json b/packages/host-trpc/package.json index d567d05c36..e27a5b490a 100644 --- a/packages/host-trpc/package.json +++ b/packages/host-trpc/package.json @@ -15,7 +15,8 @@ "clean": "node ../../scripts/rimraf.mjs .turbo" }, "dependencies": { - "@trpc/server": "catalog:" + "@trpc/server": "catalog:", + "superjson": "catalog:" }, "peerDependencies": { "inversify": "catalog:" diff --git a/packages/host-trpc/src/trpc.ts b/packages/host-trpc/src/trpc.ts index d209a88580..2150e00df9 100644 --- a/packages/host-trpc/src/trpc.ts +++ b/packages/host-trpc/src/trpc.ts @@ -1,8 +1,10 @@ import { initTRPC } from "@trpc/server"; +import superjson from "superjson"; import type { HostContext } from "./context"; const t = initTRPC.context().create({ isServer: true, + transformer: superjson, }); export const router = t.router; diff --git a/packages/shared/src/agent-runtime.ts b/packages/shared/src/agent-runtime.ts new file mode 100644 index 0000000000..a85a4a2ab1 --- /dev/null +++ b/packages/shared/src/agent-runtime.ts @@ -0,0 +1,3 @@ +export const AGENT_RUNTIMES = ["acp", "pi"] as const; + +export type AgentRuntime = (typeof AGENT_RUNTIMES)[number]; diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index a4f453d04d..33d2f564ed 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -1,4 +1,5 @@ export * from "./adapter"; +export * from "./agent-runtime"; export * from "./analytics-events"; export { type ArchivedTask, archivedTaskSchema } from "./archive-domain"; export { withTimeout } from "./async"; diff --git a/packages/shared/src/task-creation-domain.ts b/packages/shared/src/task-creation-domain.ts index 10bb6ca59f..369287d450 100644 --- a/packages/shared/src/task-creation-domain.ts +++ b/packages/shared/src/task-creation-domain.ts @@ -1,4 +1,5 @@ import type { Adapter } from "./adapter"; +import type { AgentRuntime } from "./agent-runtime"; import type { CloudRunSource, PrAuthorshipMode } from "./cloud"; import type { Task } from "./domain-types"; import type { ExecutionMode } from "./exec-types"; @@ -31,6 +32,7 @@ export interface TaskCreationInput { githubUserIntegrationId?: string; executionMode?: ExecutionMode; adapter?: Adapter; + runtime?: AgentRuntime; model?: string; reasoningLevel?: string; environmentId?: string; diff --git a/packages/ui/package.json b/packages/ui/package.json index a93ec5c3a5..ef0cfd7514 100644 --- a/packages/ui/package.json +++ b/packages/ui/package.json @@ -18,6 +18,7 @@ }, "dependencies": { "@agentclientprotocol/sdk": "0.22.1", + "@earendil-works/pi-coding-agent": "catalog:", "@base-ui/react": "^1.3.0", "@codemirror/lang-angular": "^0.1.4", "@codemirror/lang-cpp": "^6.0.3", diff --git a/packages/ui/src/features/canvas/components/ChannelIntro.tsx b/packages/ui/src/features/canvas/components/ChannelIntro.tsx index 33919338f7..4d172c0817 100644 --- a/packages/ui/src/features/canvas/components/ChannelIntro.tsx +++ b/packages/ui/src/features/canvas/components/ChannelIntro.tsx @@ -63,7 +63,7 @@ export function ChannelIntro({
{contextMdState === "created" && ( - + @@ -76,7 +76,7 @@ export function ChannelIntro({ )} {contextMdState === "building" && ( - + diff --git a/packages/ui/src/features/canvas/components/CreateChannelModal.tsx b/packages/ui/src/features/canvas/components/CreateChannelModal.tsx index d2c61fc89b..e1c23747be 100644 --- a/packages/ui/src/features/canvas/components/CreateChannelModal.tsx +++ b/packages/ui/src/features/canvas/components/CreateChannelModal.tsx @@ -187,7 +187,7 @@ export function CreateChannelModal({ )} - + {!isDescribeMode && ( <> diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx new file mode 100644 index 0000000000..f97cdaa3c4 --- /dev/null +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -0,0 +1,196 @@ +import { useHostTRPC, useHostTRPCClient } from "@posthog/host-router/react"; +import { + Button, + Empty, + EmptyDescription, + EmptyHeader, + EmptyTitle, +} from "@posthog/quill"; +import { useQuery } from "@tanstack/react-query"; +import { useSubscription } from "@trpc/tanstack-react-query"; +import { useEffect, useMemo, useState } from "react"; +import { + applyPiEvent, + emptyLiveFeed, + type PiEntries, + PiEntriesSyncer, + type PiEvent, + type PiLiveFeed, + type PiMessage, +} from "./piSessionFeed"; +import { useEnsurePiSession } from "./useEnsurePiSession"; + +interface PiSessionViewProps { + taskId: string; +} + +type PiMessageWithContent = Extract; + +function PiMessageView({ message }: { message: PiMessage }) { + if (message.role === "bashExecution") { + return <>{message.output}; + } + if ( + message.role === "branchSummary" || + message.role === "compactionSummary" + ) { + return <>{message.summary}; + } + if ("content" in message) { + return <>{messageContentText(message)}; + } + return null; +} + +function messageContentText(message: PiMessageWithContent): string { + if (typeof message.content === "string") { + return message.content; + } + + return message.content + .flatMap((part) => (part.type === "text" ? [part.text] : [])) + .join("\n"); +} + +function messageBubbleClass(role: string): string { + if (role === "user") { + return "mb-3 ml-auto max-w-[80%] rounded-lg bg-accent-3 p-3 text-sm"; + } + return "mb-3 max-w-[80%] whitespace-pre-wrap rounded-lg bg-gray-3 p-3 text-sm"; +} + +export function PiSessionView({ taskId }: PiSessionViewProps) { + const trpc = useHostTRPC(); + const client = useHostTRPCClient(); + const { error: ensureError, isSuccess: sessionReady } = + useEnsurePiSession(taskId); + + const [prompt, setPrompt] = useState(""); + const [liveFeed, setLiveFeed] = useState(emptyLiveFeed); + const [syncedEntries, setSyncedEntries] = useState( + undefined, + ); + + const { data: fetchedEntries } = useQuery({ + ...trpc.piSession.entries.queryOptions({ taskId }), + enabled: sessionReady, + }); + const { error: statusError } = useQuery({ + ...trpc.piSession.status.queryOptions({ taskId }), + enabled: sessionReady, + }); + + const syncer = useMemo( + () => + new PiEntriesSyncer( + (since) => client.piSession.entries.query({ taskId, since }), + setSyncedEntries, + ), + [client, taskId], + ); + + useEffect(() => { + syncer.seed(fetchedEntries); + }, [syncer, fetchedEntries]); + + const history = syncedEntries ?? fetchedEntries; + + useSubscription( + trpc.piSession.onEvent.subscriptionOptions( + { taskId }, + { + enabled: sessionReady, + onData: (event: PiEvent) => { + setLiveFeed((feed) => applyPiEvent(feed, event)); + + if (event.type === "agent_settled") { + void syncer.sync().then(() => setLiveFeed(emptyLiveFeed)); + } + }, + }, + ), + ); + + const send = async () => { + const text = prompt.trim(); + if (!text) { + return; + } + await client.piSession.prompt.mutate({ taskId, prompt: text }); + setPrompt(""); + }; + + const sessionError = ensureError ?? statusError; + if (sessionError) { + return ( + + + Pi session failed to start + {sessionError.message} + + + ); + } + + if (!sessionReady) { + return ( + + + Starting Pi session… + + + ); + } + + return ( +
+
+ {history?.entries.map((entry) => { + if (entry.type !== "message") { + return null; + } + + return ( +
+ +
+ ); + })} + {liveFeed.liveMessages.map((message) => ( +
+ +
+ ))} + {liveFeed.streamingMessage ? ( +
+ +
+ ) : null} +
+
+