feat: hook up pi runtime to ChatThread

Author: jonathanlab

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

@@ -188,7 +188,9 @@ import type {
 } from "@posthog/workspace-server/services/mcp-relay/identifiers";
 import type {
   PI_RPC_CLIENT_FACTORY,
+  PI_RUNTIME_FACTORY,
   PiRpcClientFactory,
+  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";
@@ -356,6 +358,8 @@ export interface MainBindings {
   [AGENT_LOGGER]: RootLogger;
   [PI_RPC_CLIENT_FACTORY]: PiRpcClientFactory;
 
+  [PI_RUNTIME_FACTORY]: PiRuntimeFactory;
+
   // Logger
   [ROOT_LOGGER]: RootLogger;
 

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

@@ -189,6 +189,7 @@ import { onboardingImportModule } from "@posthog/workspace-server/services/onboa
 import { osModule } from "@posthog/workspace-server/services/os/os.module";
 import {
   PI_RPC_CLIENT_FACTORY,
+  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";
@@ -230,6 +231,7 @@ import { workspaceMetadataModule } from "@posthog/workspace-server/services/work
 import ExternalAppsStoreImpl from "electron-store";
 import type { FileWatcherBridge } from "../index";
 import { DesktopPiRpcClientFactory } from "../platform-adapters/desktop-pi-rpc-client-factory";
+import { DesktopPiRuntimeFactory } from "../platform-adapters/desktop-pi-runtime-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";
@@ -325,6 +327,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",
 });
@@ -366,6 +379,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);
@@ -406,12 +420,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) => {
@@ -422,12 +436,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) => {
@@ -705,12 +719,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/platform-adapters/desktop-pi-runtime-factory.test.ts

@@ -0,0 +1,20 @@
+import { PiRuntime } from "@posthog/agent/pi/runtime";
+import type { PiRpcClientFactory } from "@posthog/workspace-server/services/pi-session/identifiers";
+import { describe, expect, it, vi } from "vitest";
+import { DesktopPiRuntimeFactory } from "./desktop-pi-runtime-factory";
+
+describe("DesktopPiRuntimeFactory", () => {
+  it("wraps the host-authenticated RPC client", async () => {
+    const client = { onEvent: vi.fn() };
+    const clientFactory = {
+      create: vi.fn(async () => client),
+    } as unknown as PiRpcClientFactory;
+    const factory = new DesktopPiRuntimeFactory(clientFactory);
+
+    const runtime = await factory.create({ cwd: "/workspace" });
+
+    expect(runtime).toBeInstanceOf(PiRuntime);
+    expect(runtime.client).toBe(client);
+    expect(clientFactory.create).toHaveBeenCalledWith({ cwd: "/workspace" });
+  });
+});

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

@@ -0,0 +1,24 @@
+import { PiRuntime } from "@posthog/agent/pi/runtime";
+import {
+  PI_RPC_CLIENT_FACTORY,
+  type PiRpcClientFactory,
+  type PiRuntimeFactory,
+} from "@posthog/workspace-server/services/pi-session/identifiers";
+import { inject, injectable } from "inversify";
+
+@injectable()
+export class DesktopPiRuntimeFactory implements PiRuntimeFactory {
+  constructor(
+    @inject(PI_RPC_CLIENT_FACTORY)
+    private readonly clientFactory: PiRpcClientFactory,
+  ) {}
+
+  async create(input: {
+    cwd: string;
+    model?: string;
+    sessionFile?: string;
+  }): Promise<PiRuntime> {
+    const client = await this.clientFactory.create(input);
+    return new PiRuntime(client);
+  }
+}

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

@@ -75,6 +75,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,
@@ -291,6 +295,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

@@ -33,7 +33,9 @@ import type { LlmMessage } from "@posthog/core/llm-gateway/schemas";
 import { LOCAL_MCP_WORKSPACE_CLIENT } from "@posthog/core/local-mcp/identifiers";
 import type { LocalMcpWorkspaceClient } from "@posthog/core/local-mcp/localMcpImport";
 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,
@@ -87,6 +89,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,
@@ -294,6 +297,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,12 +1,18 @@
 import "reflect-metadata";
 import { TypedContainer } from "@inversifyjs/strongly-typed";
 import { taskThreadCoreModule } from "@posthog/core/canvas/taskThread.module";
+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,
@@ -37,6 +43,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;
@@ -54,6 +61,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"
@@ -116,7 +128,8 @@
   "author": "PostHog",
   "license": "MIT",
   "scripts": {
-    "build": "node ../../scripts/rimraf.mjs dist && tsup && node build/verify-local-tools-mcp-server.mjs",
+    "build": "node ../../scripts/rimraf.mjs dist && tsup && pnpm build:types && node build/verify-local-tools-mcp-server.mjs",
+    "build:types": "tsc -p tsconfig.build.json",
     "dev": "tsup --watch",
     "test": "vitest run",
     "test:watch": "vitest",
@@ -144,6 +157,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",

packages/agent/src/pi/conversation/toolKind.ts

@@ -0,0 +1,14 @@
+import type { ToolsOptions } from "@earendil-works/pi-coding-agent";
+import type { AgentToolKind } from "@posthog/shared";
+
+export type PiToolName = keyof ToolsOptions;
+
+export const TOOL_KIND_BY_NAME: Record<PiToolName, AgentToolKind> = {
+  read: "read",
+  edit: "edit",
+  write: "edit",
+  bash: "execute",
+  grep: "search",
+  find: "search",
+  ls: "read",
+};