feat: add durable cloud Pi runtime sessions (#3719)

Co-authored-by: posthog[bot] <206114724+posthog[bot]@users.noreply.github.com>

Author: jonathanlab

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

@@ -13,6 +13,10 @@ import {
   type AutoresearchSessionClient,
   type AutoresearchStorageClient,
 } from "@posthog/core/autoresearch/identifiers";
+import {
+  CLOUD_TASK_CLIENT,
+  type CloudTaskClient,
+} from "@posthog/core/cloud-task/cloudTaskClient";
 import {
   CODE_REVIEW_WORKSPACE_CLIENT,
   REVERT_HUNK_SERVICE,
@@ -76,8 +80,10 @@ import {
 import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers";
 import type { PiRunner } from "@posthog/core/pi-runtime/piRunner";
 import {
-  PI_SESSION_CLIENT,
-  type PiSessionClient,
+  LOCAL_PI_SESSION_FACTORY,
+  PI_SESSION_PROVIDER,
+  type PiSessionFactory,
+  type PiSessionProvider,
 } from "@posthog/core/pi-runtime/piSessionController";
 import {
   type BundleLocalSkill,
@@ -295,7 +301,9 @@ export interface RendererBindings {
   [ANALYTICS_TRACKER]: AnalyticsTracker;
   [TASK_CREATION_HOST]: ITaskCreationHost;
   [PI_RUNNER]: PiRunner;
-  [PI_SESSION_CLIENT]: PiSessionClient;
+  [PI_SESSION_PROVIDER]: PiSessionProvider;
+  [LOCAL_PI_SESSION_FACTORY]: PiSessionFactory;
+  [CLOUD_TASK_CLIENT]: CloudTaskClient;
   [TASK_CREATION_EFFECTS]: TaskCreationEffects;
   [RENDERER_TASK_SERVICE]: TaskService;
   [TASK_SERVICE]: TaskService;

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

@@ -2,6 +2,10 @@ import "reflect-metadata";
 import { useDevFlagsStore } from "@features/dev-toolbar/devFlagsStore";
 import { TypedContainer } from "@inversifyjs/strongly-typed";
 import type { TrpcRouter } from "@main/trpc/router";
+import {
+  CLOUD_TASK_CLIENT,
+  type CloudTaskClient,
+} from "@posthog/core/cloud-task/cloudTaskClient";
 import {
   CODE_REVIEW_WORKSPACE_CLIENT,
   REVERT_HUNK_SERVICE,
@@ -35,7 +39,7 @@ import type { LocalMcpWorkspaceClient } from "@posthog/core/local-mcp/localMcpIm
 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 { LOCAL_PI_SESSION_FACTORY } from "@posthog/core/pi-runtime/piSessionController";
 import {
   CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL,
   CLOUD_ARTIFACT_READ_FILE_AS_BASE64,
@@ -89,7 +93,9 @@ 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 { TrpcCloudTaskClient } from "@posthog/host-router/cloud-task-client";
+import { TrpcPiRunner } from "@posthog/host-router/pi-runner";
+import { TrpcPiSessionFactory } from "@posthog/host-router/pi-session-factory";
 import {
   BROWSER_TABS_CLIENT,
   type BrowserTabsClient,
@@ -156,7 +162,6 @@ 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";
 
@@ -297,7 +302,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.bind(LOCAL_PI_SESSION_FACTORY).to(TrpcPiSessionFactory);
+container.bind<CloudTaskClient>(CLOUD_TASK_CLIENT).to(TrpcCloudTaskClient);
 container.load(piRuntimeModule);
 container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects);
 container.bind<TaskService>(RENDERER_TASK_SERVICE).to(TaskService);

apps/web/src/web-container.ts

@@ -29,6 +29,10 @@ import { canvasCoreModule } from "@posthog/core/canvas/canvas.module";
 import { taskThreadCoreModule } from "@posthog/core/canvas/taskThread.module";
 import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task";
 import { cloudTaskModule } from "@posthog/core/cloud-task/cloud-task.module";
+import {
+  CLOUD_TASK_CLIENT,
+  type CloudTaskClient,
+} from "@posthog/core/cloud-task/cloudTaskClient";
 import {
   CLOUD_TASK_AUTH,
   CLOUD_TASK_SERVICE,
@@ -86,10 +90,14 @@ import {
   type GithubConnectClient as OnboardingGithubConnectContract,
 } from "@posthog/core/onboarding/identifiers";
 import { onboardingModule } from "@posthog/core/onboarding/onboarding.module";
+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,
-  type PiSessionClient,
+  LOCAL_PI_SESSION_FACTORY,
+  PI_SESSION_PROVIDER,
+  type PiSessionFactory,
+  type PiSessionProvider,
 } from "@posthog/core/pi-runtime/piSessionController";
 import {
   type BundleLocalSkill,
@@ -161,7 +169,9 @@ import {
   HOST_TRPC_CLIENT,
   type HostTrpcClient,
 } from "@posthog/host-router/client";
-import { TrpcPiSessionClient } from "@posthog/host-router/pi-session-client";
+import { TrpcCloudTaskClient } from "@posthog/host-router/cloud-task-client";
+import { TrpcPiRunner } from "@posthog/host-router/pi-runner";
+import { TrpcPiSessionFactory } from "@posthog/host-router/pi-session-factory";
 import {
   ANALYTICS_SERVICE,
   type IAnalytics,
@@ -314,7 +324,10 @@ import { hostTrpcClient } from "./web-trpc";
 
 interface WebBindings {
   [HOST_TRPC_CLIENT]: HostTrpcClient;
-  [PI_SESSION_CLIENT]: PiSessionClient;
+  [PI_SESSION_PROVIDER]: PiSessionProvider;
+  [LOCAL_PI_SESSION_FACTORY]: PiSessionFactory;
+  [CLOUD_TASK_CLIENT]: CloudTaskClient;
+  [PI_RUNNER]: PiRunner;
   [ROOT_LOGGER]: RootLogger;
   [HOST_LOGGER]: HostLogger;
   [FEATURE_FLAGS]: FeatureFlags;
@@ -395,7 +408,9 @@ export const container = new TypedContainer<WebBindings>({
 // Keystone: the same typed host client the renderer binds — served in-process
 // here (web-trpc.ts) instead of over Electron IPC.
 container.bind(HOST_TRPC_CLIENT).toConstantValue(hostTrpcClient);
-container.bind(PI_SESSION_CLIENT).to(TrpcPiSessionClient);
+container.bind(LOCAL_PI_SESSION_FACTORY).to(TrpcPiSessionFactory);
+container.bind(CLOUD_TASK_CLIENT).to(TrpcCloudTaskClient);
+container.bind(PI_RUNNER).to(TrpcPiRunner);
 container.load(piRuntimeModule);
 
 // Logger: web uses console; electron uses electron-log. Same RootLogger shape.

apps/web/src/web-host-router.ts

@@ -1,3 +1,7 @@
+import { fetchPosthogPiModelCatalog } from "@posthog/agent/pi/model-catalog";
+import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
+import type { AuthService } from "@posthog/core/auth/auth";
+import { AUTH_SERVICE } from "@posthog/core/auth/auth.module";
 import { TEAM_SKILLS_SERVICE } from "@posthog/core/skills/identifiers";
 import type { TeamSkillsService } from "@posthog/core/skills/teamSkillsService";
 import { resolveService } from "@posthog/di/container";
@@ -134,6 +138,19 @@ const agentStubRouter = router({
   }),
   // Called by resetSessionService() on logout/project switch.
   resetAll: publicProcedure.mutation(() => undefined),
+  getPiModelCatalog: publicProcedure
+    .input(
+      z.object({ apiHost: z.string(), region: z.enum(["us", "eu", "dev"]) }),
+    )
+    .query(async ({ input }) => {
+      const auth = resolveService<AuthService>(AUTH_SERVICE);
+      const { accessToken } = await auth.getValidAccessToken();
+      return fetchPosthogPiModelCatalog(
+        getLlmGatewayUrl(input.apiHost),
+        input.region,
+        accessToken,
+      );
+    }),
   // Model/mode/effort options for the task-input preview + cloud run creation
   // (a cloud run requires a model). Real: fetched from the CORS-open PostHog LLM
   // gateway, same logic the desktop main process runs (see web-agent-config.ts).

packages/agent/package.json

@@ -32,6 +32,14 @@
       "types": "./dist/pi/rpc-client.d.ts",
       "import": "./dist/pi/rpc-client.js"
     },
+    "./pi/rpc-transport": {
+      "types": "./dist/pi/rpc-transport.d.ts",
+      "import": "./dist/pi/rpc-transport.js"
+    },
+    "./pi/remote-rpc-client": {
+      "types": "./dist/pi/remote-rpc-client.d.ts",
+      "import": "./dist/pi/remote-rpc-client.js"
+    },
     "./pi/conversation": {
       "types": "./dist/pi/conversation/translatePiConversation.d.ts",
       "import": "./dist/pi/conversation/translatePiConversation.js"
@@ -44,6 +52,10 @@
       "types": "./dist/pi/types.d.ts",
       "import": "./dist/pi/types.js"
     },
+    "./pi/model-catalog": {
+      "types": "./dist/pi/model-catalog.d.ts",
+      "import": "./dist/pi/model-catalog.js"
+    },
     "./pr-url-detector": {
       "types": "./dist/pr-url-detector.d.ts",
       "import": "./dist/pr-url-detector.js"
@@ -114,7 +126,7 @@
     }
   },
   "bin": {
-    "agent-server": "./dist/server/bin.cjs"
+    "agent-server": "./dist/server/bin.js"
   },
   "type": "module",
   "keywords": [