fix(pi): preserve cloud model selection (#4032)

Author: jonathanlab

apps/web/package.json

@@ -19,6 +19,7 @@
     "@posthog/core": "workspace:*",
     "@posthog/di": "workspace:*",
     "@posthog/host-router": "workspace:*",
+    "@posthog/harness": "workspace:*",
     "@posthog/host-trpc": "workspace:*",
     "@posthog/platform": "workspace:*",
     "@posthog/shared": "workspace:*",

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

@@ -1,10 +1,10 @@
-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";
+import { fetchPosthogPiModelCatalog } from "@posthog/harness/extensions/posthog-provider/model-catalog";
 import { analyticsRouter } from "@posthog/host-router/routers/analytics.router";
 import { authRouter } from "@posthog/host-router/routers/auth.router";
 import { canvasDataRouter } from "@posthog/host-router/routers/canvas-data.router";

packages/agent/package.json

@@ -52,10 +52,6 @@
       "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"

packages/agent/tsup.config.ts

@@ -132,7 +132,6 @@ export default defineConfig([
       "src/pi/rpc-client.ts",
       "src/pi/runtime.ts",
       "src/pi/types.ts",
-      "src/pi/model-catalog.ts",
       "src/pi/conversation/translatePiConversation.ts",
       "src/resume.ts",
       "src/types.ts",

packages/core/src/pi-runtime/piSessionController.test.ts

@@ -819,6 +819,43 @@ describe("PiSessionController", () => {
     expect(resumedSession.client.prompt).toHaveBeenCalledWith("continue");
   });
 
+  it("applies deferred Pi config before the first resumed prompt", async () => {
+    const terminalSession = {
+      ...createSession(),
+      resumeRequired: true,
+      taskRunId: "run-1",
+    };
+    const resumedSession = createSession();
+    const provider = {
+      get: vi
+        .fn()
+        .mockResolvedValueOnce(terminalSession)
+        .mockResolvedValue(resumedSession),
+    } as PiSessionProvider;
+    const resumeCloudPiRun = vi.fn(async () => ({ id: "run-1" }));
+    const controller = new PiSessionController(provider, {
+      resumeCloudPiRun,
+    } as unknown as TaskService);
+
+    await controller.connect("task-1");
+    await controller.submit("task-1", "continue", false, "steer", {
+      model: { provider: "posthog", id: "gpt-5.6-terra" },
+      thinkingLevel: "high",
+    });
+
+    expect(resumedSession.client.setModel).toHaveBeenCalledWith(
+      "posthog",
+      "gpt-5.6-terra",
+    );
+    expect(resumedSession.client.setThinkingLevel).toHaveBeenCalledWith("high");
+    expect(resumedSession.client.prompt).toHaveBeenCalledWith("continue");
+    expect(
+      vi.mocked(resumedSession.client.setModel).mock.invocationCallOrder[0],
+    ).toBeLessThan(
+      vi.mocked(resumedSession.client.prompt).mock.invocationCallOrder[0],
+    );
+  });
+
   it("resumes and retries a message when the prior sandbox is gone", async () => {
     const staleSession = {
       ...createSession(),

packages/core/src/pi-runtime/piSessionController.ts

@@ -35,6 +35,11 @@ export type {
 
 export type PiModelSelection = Pick<PiNativeModelInfo, "provider" | "id">;
 
+export interface PiDeferredConfig {
+  model?: PiModelSelection;
+  thinkingLevel?: PiThinkingLevel;
+}
+
 export const PI_SESSION_PROVIDER = Symbol.for("posthog.pi.sessionProvider");
 export const LOCAL_PI_SESSION_FACTORY = Symbol.for(
   "posthog.pi.localSessionFactory",
@@ -301,6 +306,7 @@ export class PiSessionController {
     text: string,
     isStreaming: boolean,
     messagingMode: PiMessagingMode,
+    deferredConfig?: PiDeferredConfig,
   ): Promise<PiSubmitResult> {
     const message = text.trim();
     const action = this.getSubmitAction(message, isStreaming, messagingMode);
@@ -382,6 +388,7 @@ export class PiSessionController {
 
       try {
         const session = await this.getWritablePiSession(taskId);
+        await this.applyDeferredConfig(session, deferredConfig);
         this.markTurnPending(taskId);
         if (session.sendUserMessage && messageId) {
           const taskRunId = this.taskRunIds.get(taskId);
@@ -1088,6 +1095,22 @@ export class PiSessionController {
     });
   }
 
+  private async applyDeferredConfig(
+    session: PiSession,
+    config: PiDeferredConfig | undefined,
+  ): Promise<void> {
+    if (!config) {
+      return;
+    }
+
+    if (config.model) {
+      await session.client.setModel(config.model.provider, config.model.id);
+    }
+    if (config.thinkingLevel) {
+      await session.client.setThinkingLevel(config.thinkingLevel);
+    }
+  }
+
   private async refreshStatus(taskId: string): Promise<void> {
     const session = await this.getPiSession(taskId);
     const status = await session.client.getState();

packages/core/src/task-detail/taskCreationSaga.test.ts

@@ -453,8 +453,8 @@ describe("TaskCreationSaga", () => {
         branch: "main",
         adapter: undefined,
         piRuntime: true,
-        model: undefined,
-        reasoningLevel: undefined,
+        model: "gpt-5.4",
+        reasoningLevel: "high",
         initialPermissionMode: undefined,
       }),
     );

packages/core/src/task-detail/taskCreationSaga.ts

@@ -408,8 +408,8 @@ export class TaskCreationSaga extends Saga<
             branch,
             adapter: cloudAdapter,
             ...(isPiRuntime ? { piRuntime: true } : {}),
-            model: isPiRuntime ? undefined : input.model,
-            reasoningLevel: isPiRuntime ? undefined : input.reasoningLevel,
+            model: input.model,
+            reasoningLevel: input.reasoningLevel,
             contextWindow: isPiRuntime ? undefined : input.contextWindow,
             fastMode: isPiRuntime ? undefined : input.fastMode,
             sandboxEnvironmentId: input.sandboxEnvironmentId,

…kages/agent/src/pi/model-catalog.test.ts …s/posthog-provider/model-catalog.test.tspackages/agent/src/pi/model-catalog.test.ts renamed to packages/harness/src/extensions/posthog-provider/model-catalog.test.ts packages/agent/src/pi/model-catalog.test.ts renamed to packages/harness/src/extensions/posthog-provider/model-catalog.test.ts

@@ -2,20 +2,22 @@ import {
   getSupportedThinkingLevels,
   type ModelThinkingLevel,
 } from "@earendil-works/pi-ai";
+import type { ModelInfo } from "@earendil-works/pi-coding-agent";
+import type { CloudRegion } from "@posthog/shared";
 import {
   fetchPosthogGatewayModels,
   type GatewayModel,
   resolveModelConfigsFromGatewayModels,
-} from "@posthog/harness/extensions/posthog-provider/models";
-import type { CloudRegion } from "@posthog/shared";
+} from "./models";
 
-export interface PiModelCatalogEntry {
+export type PiModelCatalogEntry = Omit<
+  Pick<ModelInfo, "provider" | "id" | "contextWindow">,
+  "provider"
+> & {
   provider: "posthog";
-  id: string;
   name: string;
-  contextWindow: number;
   thinkingLevels: ModelThinkingLevel[];
-}
+};
 
 export function resolvePosthogPiModelCatalog(
   gatewayModels: GatewayModel[],