feat: hook up pi runtime to ChatThread

Author: jonathanlab

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_RUNTIME_FACTORY,
+  PiRuntimeFactory,
+} 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_RUNTIME_FACTORY]: PiRuntimeFactory;
 
   // Logger
   [ROOT_LOGGER]: RootLogger;

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

@@ -184,7 +184,10 @@ 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 {
+  PI_RUNTIME_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";
@@ -253,6 +256,7 @@ import {
   TokenCipherPortAdapter,
 } from "../services/auth/port-adapters";
 import { DeepLinkService } from "../services/deep-link/service";
+import { DesktopPiRuntimeFactory } from "../services/desktop-pi-runtime-factory";
 import { DevActionsService } from "../services/dev-actions/service";
 import { DevFlagsService } from "../services/dev-flags/service";
 import { DevLogsService } from "../services/dev-logs/service";
@@ -318,6 +322,17 @@ import {
   WORKTREE_REPOSITORY as MAIN_WORKTREE_REPOSITORY,
 } from "./tokens";
 
+async function cancelTaskSessions(
+  agentService: AgentService,
+  piSessionService: PiSessionService,
+  taskId: string,
+): Promise<void> {
+  await Promise.all([
+    agentService.cancelSessionsByTaskId(taskId),
+    piSessionService.stop(taskId),
+  ]);
+}
+
 export const container = new TypedContainer<MainBindings>({
   defaultScope: "Singleton",
 });
@@ -359,6 +374,7 @@ container
   .bind(MAIN_DEFAULT_ADDITIONAL_DIRECTORY_REPOSITORY)
   .toService(DEFAULT_ADDITIONAL_DIRECTORY_REPOSITORY);
 container.load(agentModule);
+container.bind(PI_RUNTIME_FACTORY).to(DesktopPiRuntimeFactory);
 container.load(piSessionModule);
 container.bind(AGENT_SLEEP_COORDINATOR).toService(MAIN_SLEEP_SERVICE);
 container.bind(AGENT_MCP_APPS).toService(MCP_APPS_SERVICE);
@@ -395,12 +411,12 @@ container.bind(MCP_PROXY_AUTH).toDynamicValue((ctx) => {
 });
 container.load(archiveModule);
 container.bind(ARCHIVE_SESSION_CANCELLER).toDynamicValue((ctx) => ({
-  cancelSessionsByTaskId: async (taskId: string) => {
-    await Promise.all([
-      ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
-      ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
-    ]);
-  },
+  cancelSessionsByTaskId: (taskId: string) =>
+    cancelTaskSessions(
+      ctx.get<AgentService>(AGENT_SERVICE),
+      ctx.get<PiSessionService>(PI_SESSION_SERVICE),
+      taskId,
+    ),
 }));
 container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({
   stopWatching: async (worktreePath: string) => {
@@ -411,12 +427,12 @@ container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({
 }));
 container.load(suspensionModule);
 container.bind(SUSPENSION_SESSION_CANCELLER).toDynamicValue((ctx) => ({
-  cancelSessionsByTaskId: async (taskId: string) => {
-    await Promise.all([
-      ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
-      ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
-    ]);
-  },
+  cancelSessionsByTaskId: (taskId: string) =>
+    cancelTaskSessions(
+      ctx.get<AgentService>(AGENT_SERVICE),
+      ctx.get<PiSessionService>(PI_SESSION_SERVICE),
+      taskId,
+    ),
 }));
 container.bind(SUSPENSION_FILE_WATCHER).toDynamicValue((ctx) => ({
   stopWatching: async (worktreePath: string) => {
@@ -687,12 +703,12 @@ container.load(workspaceModule);
 container.bind(WORKSPACE_AGENT).toDynamicValue((ctx): WorkspaceAgent => {
   const agent = ctx.get<AgentService>(AGENT_SERVICE);
   return {
-    cancelSessionsByTaskId: async (taskId) => {
-      await Promise.all([
-        agent.cancelSessionsByTaskId(taskId),
-        ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
-      ]);
-    },
+    cancelSessionsByTaskId: (taskId) =>
+      cancelTaskSessions(
+        agent,
+        ctx.get<PiSessionService>(PI_SESSION_SERVICE),
+        taskId,
+      ),
     onAgentFileActivity: (handler) =>
       agent.on(AgentServiceEvent.AgentFileActivity, handler),
   };

apps/code/src/main/services/desktop-pi-runtime-factory.test.ts

@@ -0,0 +1,44 @@
+import type { PiRuntime } from "@posthog/agent/pi/runtime";
+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 { DesktopPiRuntimeFactory } from "./desktop-pi-runtime-factory";
+
+const createPiRuntime = vi.hoisted(() => vi.fn());
+
+vi.mock("@posthog/agent/pi/runtime", () => ({ createPiRuntime }));
+
+describe("DesktopPiRuntimeFactory", () => {
+  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 runtime = {} as PiRuntime;
+    createPiRuntime.mockReturnValue(runtime);
+    const factory = new DesktopPiRuntimeFactory(auth, authProxy);
+
+    await expect(factory.create({ cwd: "/workspace" })).resolves.toBe(runtime);
+    expect(authProxy.start).toHaveBeenCalledWith(
+      getLlmGatewayUrl(getCloudUrlFromRegion("eu")),
+    );
+    expect(createPiRuntime).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/services/desktop-pi-runtime-factory.ts

@@ -0,0 +1,56 @@
+import { createPiRuntime, type PiRuntime } from "@posthog/agent/pi/runtime";
+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 { PiRuntimeFactory } from "@posthog/workspace-server/services/pi-session/identifiers";
+import { inject, injectable } from "inversify";
+
+const PROXY_API_KEY = "posthog-code-auth-proxy";
+
+@injectable()
+export class DesktopPiRuntimeFactory implements PiRuntimeFactory {
+  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<PiRuntime> {
+    const credentials = await this.auth.getOAuthCredentials();
+    if (!credentials) {
+      throw new Error("Pi requires PostHog authentication");
+    }
+
+    const baseUrl = await this.getProxyUrl(credentials.region);
+
+    return createPiRuntime({
+      ...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/services/index.ts

@@ -3,5 +3,6 @@
  * This file is auto-generated by vite-plugin-auto-services.ts
  */
 
+import "./desktop-pi-runtime-factory.js";
 import "./dev-toolbar.js";
 import "./settingsStore.js";

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

@@ -67,6 +67,10 @@ import {
 } from "@posthog/core/onboarding/identifiers";
 import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers";
 import type { PiRunner } from "@posthog/core/pi-runtime/piRunner";
+import {
+  PI_SESSION_CLIENT,
+  type PiSessionClient,
+} from "@posthog/core/pi-runtime/piSessionController";
 import {
   type BundleLocalSkill,
   CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL,
@@ -279,6 +283,7 @@ export interface RendererBindings {
   [ANALYTICS_TRACKER]: AnalyticsTracker;
   [TASK_CREATION_HOST]: ITaskCreationHost;
   [PI_RUNNER]: PiRunner;
+  [PI_SESSION_CLIENT]: PiSessionClient;
   [TASK_CREATION_EFFECTS]: TaskCreationEffects;
   [RENDERER_TASK_SERVICE]: TaskService;
   [TASK_SERVICE]: TaskService;

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

@@ -31,7 +31,9 @@ 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 { piRuntimeModule } from "@posthog/core/pi-runtime/pi-runtime.module";
 import type { PiRunner } from "@posthog/core/pi-runtime/piRunner";
+import { PI_SESSION_CLIENT } from "@posthog/core/pi-runtime/piSessionController";
 import {
   CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL,
   CLOUD_ARTIFACT_READ_FILE_AS_BASE64,
@@ -85,6 +87,7 @@ import {
 import { WorkspaceSetupService } from "@posthog/core/workspace/WorkspaceSetupService";
 import { setRootContainer } from "@posthog/di/container";
 import { HOST_TRPC_CLIENT } from "@posthog/host-router/client";
+import { TrpcPiSessionClient } from "@posthog/host-router/pi-session-client";
 import {
   BROWSER_TABS_CLIENT,
   type BrowserTabsClient,
@@ -292,6 +295,8 @@ container
 // Bind services
 container.bind<ITaskCreationHost>(TASK_CREATION_HOST).to(TrpcTaskCreationHost);
 container.bind<PiRunner>(PI_RUNNER).to(TrpcPiRunner);
+container.bind(PI_SESSION_CLIENT).to(TrpcPiSessionClient);
+container.load(piRuntimeModule);
 container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects);
 container.bind<TaskService>(RENDERER_TASK_SERVICE).to(TaskService);
 container.bind<TaskService>(TASK_SERVICE).toService(RENDERER_TASK_SERVICE);

apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts

@@ -3,26 +3,27 @@ import type {
   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";
+import { inject, injectable } from "inversify";
 
-function hostClient(): HostTrpcClient {
-  return resolveService<HostTrpcClient>(HOST_TRPC_CLIENT);
-}
-
+@injectable()
 export class TrpcPiRunner implements PiRunner {
+  constructor(
+    @inject(HOST_TRPC_CLIENT) private readonly hostClient: HostTrpcClient,
+  ) {}
+
   async create(input: PiRunInput): Promise<void> {
-    await hostClient().piSession.start.mutate(input);
+    await this.hostClient.piSession.start.mutate(input);
   }
 
   resume(input: PiResumeInput): Promise<void> {
-    return hostClient().piSession.resume.mutate(input);
+    return this.hostClient.piSession.resume.mutate(input);
   }
 
   stop(taskId: string): Promise<void> {
-    return hostClient().piSession.stop.mutate({ taskId });
+    return this.hostClient.piSession.stop.mutate({ taskId });
   }
 }

apps/web/src/web-container.ts

@@ -1,11 +1,17 @@
 import "reflect-metadata";
 import { TypedContainer } from "@inversifyjs/strongly-typed";
+import { piRuntimeModule } from "@posthog/core/pi-runtime/pi-runtime.module";
+import {
+  PI_SESSION_CLIENT,
+  type PiSessionClient,
+} from "@posthog/core/pi-runtime/piSessionController";
 import { setRootContainer } from "@posthog/di/container";
 import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger";
 import {
   HOST_TRPC_CLIENT,
   type HostTrpcClient,
 } from "@posthog/host-router/client";
+import { TrpcPiSessionClient } from "@posthog/host-router/pi-session-client";
 import { sandboxProxyHtml } from "@posthog/shared/mcp-sandbox-proxy";
 import {
   AUTH_SIDE_EFFECTS,
@@ -36,6 +42,7 @@ import { hostTrpcClient } from "./web-trpc";
 
 interface WebBindings {
   [HOST_TRPC_CLIENT]: HostTrpcClient;
+  [PI_SESSION_CLIENT]: PiSessionClient;
   [ROOT_LOGGER]: RootLogger;
   [FEATURE_FLAGS]: FeatureFlags;
   [ANALYTICS_TRACKER]: AnalyticsTracker;
@@ -53,6 +60,8 @@ export const container = new TypedContainer<WebBindings>({
 
 // Keystone: the same typed host client the renderer binds, over HTTP not IPC.
 container.bind(HOST_TRPC_CLIENT).toConstantValue(hostTrpcClient);
+container.bind(PI_SESSION_CLIENT).to(TrpcPiSessionClient);
+container.load(piRuntimeModule);
 
 // Logger: web uses console; electron uses electron-log. Same RootLogger shape.
 const scoped = (name?: string): RootLogger => ({

packages/agent/package.json

@@ -32,6 +32,18 @@
       "types": "./dist/pi/rpc-client.d.ts",
       "import": "./dist/pi/rpc-client.js"
     },
+    "./pi/conversation": {
+      "types": "./dist/pi/conversation/translatePiConversation.d.ts",
+      "import": "./dist/pi/conversation/translatePiConversation.js"
+    },
+    "./pi/runtime": {
+      "types": "./dist/pi/runtime.d.ts",
+      "import": "./dist/pi/runtime.js"
+    },
+    "./pi/types": {
+      "types": "./dist/pi/types.d.ts",
+      "import": "./dist/pi/types.js"
+    },
     "./pr-url-detector": {
       "types": "./dist/pr-url-detector.d.ts",
       "import": "./dist/pr-url-detector.js"
@@ -140,6 +152,8 @@
     "@agentclientprotocol/sdk": "1.1.0",
     "@anthropic-ai/claude-agent-sdk": "0.3.197",
     "@anthropic-ai/sdk": "0.109.0",
+    "@earendil-works/pi-agent-core": "catalog:",
+    "@earendil-works/pi-ai": "catalog:",
     "@earendil-works/pi-coding-agent": "catalog:",
     "@hono/node-server": "^1.19.9",
     "@openai/codex": "0.140.0",