feat(harness): integrate pi runtime into desktop app

Author: jonathanlab

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,
   ],
 

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(),

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",

apps/code/src/main/di/bindings.ts

@@ -180,6 +180,10 @@ import type {
 } from "@posthog/workspace-server/services/local-logs/identifiers";
 import type { MCP_PROXY_AUTH } from "@posthog/workspace-server/services/mcp-proxy/identifiers";
 import type { McpProxyAuth } from "@posthog/workspace-server/services/mcp-proxy/ports";
+import type {
+  PI_RPC_CLIENT_FACTORY,
+  PiRpcClientFactory,
+} from "@posthog/workspace-server/services/pi-session/identifiers";
 import type { PosthogPluginService } from "@posthog/workspace-server/services/posthog-plugin/posthog-plugin";
 import type { ProcessTrackingService } from "@posthog/workspace-server/services/process-tracking/process-tracking";
 import type {
@@ -344,6 +348,7 @@ export interface MainBindings {
   [AGENT_REPO_FILES]: unknown;
   [AGENT_AUTH]: unknown;
   [AGENT_LOGGER]: RootLogger;
+  [PI_RPC_CLIENT_FACTORY]: PiRpcClientFactory;
 
   // Logger
   [ROOT_LOGGER]: RootLogger;

apps/code/src/main/di/container.ts

@@ -184,6 +184,12 @@ 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_RPC_CLIENT_FACTORY,
+  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";
@@ -220,6 +226,7 @@ import { workspaceModule } from "@posthog/workspace-server/services/workspace/wo
 import { workspaceMetadataModule } from "@posthog/workspace-server/services/workspace-metadata/workspace-metadata.module";
 import ExternalAppsStoreImpl from "electron-store";
 import type { FileWatcherBridge } from "../index";
+import { DesktopPiRpcClientFactory } from "../platform-adapters/desktop-pi-rpc-client-factory";
 import { ElectronAppLifecycle } from "../platform-adapters/electron-app-lifecycle";
 import { ElectronAppMeta } from "../platform-adapters/electron-app-meta";
 import { ElectronAppMetrics } from "../platform-adapters/electron-app-metrics";
@@ -356,10 +363,15 @@ 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);
 container.bind(AGENT_AUTH).toService(MAIN_AUTH_SERVICE);
+container
+  .bind(PI_RPC_CLIENT_FACTORY)
+  .to(DesktopPiRpcClientFactory)
+  .inSingletonScope();
 container.bind(AGENT_LOGGER).toConstantValue(logger);
 container.load(osModule);
 container.bind<RootLogger>(ROOT_LOGGER).toConstantValue(logger);
@@ -391,8 +403,12 @@ container.bind(MCP_PROXY_AUTH).toDynamicValue((ctx) => {
 });
 container.load(archiveModule);
 container.bind(ARCHIVE_SESSION_CANCELLER).toDynamicValue((ctx) => ({
-  cancelSessionsByTaskId: (taskId: string) =>
-    ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
+  cancelSessionsByTaskId: async (taskId: string) => {
+    await Promise.all([
+      ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
+      ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
+    ]);
+  },
 }));
 container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({
   stopWatching: async (worktreePath: string) => {
@@ -403,8 +419,12 @@ container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({
 }));
 container.load(suspensionModule);
 container.bind(SUSPENSION_SESSION_CANCELLER).toDynamicValue((ctx) => ({
-  cancelSessionsByTaskId: (taskId: string) =>
-    ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
+  cancelSessionsByTaskId: async (taskId: string) => {
+    await Promise.all([
+      ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
+      ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
+    ]);
+  },
 }));
 container.bind(SUSPENSION_FILE_WATCHER).toDynamicValue((ctx) => ({
   stopWatching: async (worktreePath: string) => {
@@ -675,7 +695,12 @@ container.load(workspaceModule);
 container.bind(WORKSPACE_AGENT).toDynamicValue((ctx): WorkspaceAgent => {
   const agent = ctx.get<AgentService>(AGENT_SERVICE);
   return {
-    cancelSessionsByTaskId: (taskId) => agent.cancelSessionsByTaskId(taskId),
+    cancelSessionsByTaskId: async (taskId) => {
+      await Promise.all([
+        agent.cancelSessionsByTaskId(taskId),
+        ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
+      ]);
+    },
     onAgentFileActivity: (handler) =>
       agent.on(AgentServiceEvent.AgentFileActivity, handler),
   };

apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.test.ts

@@ -0,0 +1,44 @@
+import type { PiRpcClient } from "@posthog/agent/pi/rpc-client";
+import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
+import { getCloudUrlFromRegion } from "@posthog/shared";
+import type { AgentAuth } from "@posthog/workspace-server/services/agent/ports";
+import type { AuthProxyService } from "@posthog/workspace-server/services/auth-proxy/auth-proxy";
+import { describe, expect, it, vi } from "vitest";
+import { DesktopPiRpcClientFactory } from "./desktop-pi-rpc-client-factory";
+
+const createPiRpcClient = vi.hoisted(() => vi.fn());
+
+vi.mock("@posthog/agent/pi/rpc-client", () => ({ createPiRpcClient }));
+
+describe("DesktopPiRpcClientFactory", () => {
+  it("routes Pi through the shared host auth proxy", async () => {
+    const auth = {
+      getOAuthCredentials: vi.fn(async () => ({
+        access: "access-token",
+        refresh: "refresh-token",
+        expires: 1,
+        region: "eu" as const,
+      })),
+    } as unknown as AgentAuth;
+    const authProxy = {
+      start: vi.fn(async () => "http://127.0.0.1:1234"),
+    } as unknown as AuthProxyService;
+    const client = {} as PiRpcClient;
+    createPiRpcClient.mockReturnValue(client);
+    const factory = new DesktopPiRpcClientFactory(auth, authProxy);
+
+    await expect(factory.create({ cwd: "/workspace" })).resolves.toBe(client);
+    expect(authProxy.start).toHaveBeenCalledWith(
+      getLlmGatewayUrl(getCloudUrlFromRegion("eu")),
+    );
+    expect(createPiRpcClient).toHaveBeenCalledWith({
+      cwd: "/workspace",
+      providerOptions: {
+        region: "eu",
+        baseUrl: "http://127.0.0.1:1234",
+        apiKey: "posthog-code-auth-proxy",
+      },
+      env: { POSTHOG_REGION: "eu" },
+    });
+  });
+});

apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.ts

@@ -0,0 +1,59 @@
+import {
+  createPiRpcClient,
+  type PiRpcClient,
+} from "@posthog/agent/pi/rpc-client";
+import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
+import { type CloudRegion, getCloudUrlFromRegion } from "@posthog/shared";
+import { AGENT_AUTH } from "@posthog/workspace-server/services/agent/identifiers";
+import type { AgentAuth } from "@posthog/workspace-server/services/agent/ports";
+import type { AuthProxyService } from "@posthog/workspace-server/services/auth-proxy/auth-proxy";
+import { AUTH_PROXY_SERVICE } from "@posthog/workspace-server/services/auth-proxy/identifiers";
+import type { PiRpcClientFactory } from "@posthog/workspace-server/services/pi-session/identifiers";
+import { inject, injectable } from "inversify";
+
+const PROXY_API_KEY = "posthog-code-auth-proxy";
+
+@injectable()
+export class DesktopPiRpcClientFactory implements PiRpcClientFactory {
+  private proxyRegion?: CloudRegion;
+  private proxyUrlPromise?: Promise<string>;
+
+  constructor(
+    @inject(AGENT_AUTH) private readonly auth: AgentAuth,
+    @inject(AUTH_PROXY_SERVICE)
+    private readonly authProxy: AuthProxyService,
+  ) {}
+
+  async create(input: {
+    cwd: string;
+    model?: string;
+    sessionFile?: string;
+  }): Promise<PiRpcClient> {
+    const credentials = await this.auth.getOAuthCredentials();
+    if (!credentials) {
+      throw new Error("Pi requires PostHog authentication");
+    }
+
+    const baseUrl = await this.getProxyUrl(credentials.region);
+
+    return createPiRpcClient({
+      ...input,
+      providerOptions: {
+        region: credentials.region,
+        baseUrl,
+        apiKey: PROXY_API_KEY,
+      },
+      env: { POSTHOG_REGION: credentials.region },
+    });
+  }
+
+  private getProxyUrl(region: CloudRegion): Promise<string> {
+    if (this.proxyRegion !== region || !this.proxyUrlPromise) {
+      this.proxyRegion = region;
+      const gatewayUrl = getLlmGatewayUrl(getCloudUrlFromRegion(region));
+      this.proxyUrlPromise = this.authProxy.start(gatewayUrl);
+    }
+
+    return this.proxyUrlPromise;
+  }
+}

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,

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;

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<ITaskCreationHost>(TASK_CREATION_HOST).to(TrpcTaskCreationHost);
+container.bind<PiRunner>(PI_RUNNER).to(TrpcPiRunner);
 container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects);
 container.bind<TaskService>(RENDERER_TASK_SERVICE).to(TaskService);
 container.bind<TaskService>(TASK_SERVICE).toService(RENDERER_TASK_SERVICE);