fix(pi): preserve cloud model selection

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,

packages/agent/src/pi/model-catalog.ts → packages/harness/src/extensions/posthog-provider/model-catalog.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[],

packages/harness/tsup.config.ts

@@ -13,6 +13,7 @@ export default defineConfig({
     "src/extensions/posthog-provider/index.ts",
     "src/extensions/posthog-provider/provider.ts",
     "src/extensions/posthog-provider/models.ts",
+    "src/extensions/posthog-provider/model-catalog.ts",
     "src/extensions/posthog-provider/oauth.ts",
     "src/extensions/posthog-provider/gateway.ts",
     "src/extensions/posthog-provider/gateway-auth.ts",

packages/ui/src/features/pi-sessions/PiSessionModelControls.tsx

@@ -0,0 +1,148 @@
+import type {
+  PiModelSelection,
+  PiSessionController,
+  PiThinkingLevel,
+} from "@posthog/core/pi-runtime/piSessionController";
+import type { PiControllerSessionState } from "@posthog/core/pi-runtime/piSessionStore";
+import { Skeleton } from "@posthog/quill";
+import { isTerminalStatus } from "@posthog/shared/domain-types";
+import { useCallback } from "react";
+import { PiModelSelector, PiThinkingLevelSelector } from "./PiSessionControls";
+import {
+  getPiPendingConfig,
+  usePiPendingConfigStore,
+} from "./piPendingConfigStore";
+import { usePiModelCatalog } from "./usePiModelCatalog";
+
+interface PiSessionModelControlsProps {
+  taskId: string;
+  taskRunId?: string;
+  session: PiControllerSessionState;
+  controller: PiSessionController;
+  isOnline: boolean;
+  onError: (error: unknown, fallback: string) => void;
+}
+
+export function PiSessionModelControls({
+  taskId,
+  taskRunId,
+  session,
+  controller,
+  isOnline,
+  onError,
+}: PiSessionModelControlsProps) {
+  const isCloudSession = session.cloudStatus !== undefined;
+  const isTerminalCloudRun =
+    isCloudSession && isTerminalStatus(session.cloudStatus);
+  const pendingConfig = usePiPendingConfigStore((state) =>
+    getPiPendingConfig(state, taskId, taskRunId),
+  );
+  const setPendingConfig = usePiPendingConfigStore((state) => state.setConfig);
+  const { data: catalog = [], isPending: catalogLoading } =
+    usePiModelCatalog(isCloudSession);
+  const controlsDisabled =
+    session.status?.isStreaming ||
+    session.status?.isCompacting ||
+    session.isBashRunning ||
+    session.connectionState !== "connected";
+  const currentModel = pendingConfig?.model ?? session.status?.model;
+  const models = isCloudSession ? catalog : session.models;
+  const modelsLoaded = isCloudSession ? !catalogLoading : session.modelsLoaded;
+  const catalogModel = catalog.find(
+    (model) =>
+      model.provider === currentModel?.provider && model.id === currentModel.id,
+  );
+  const thinkingLevels = isCloudSession
+    ? (catalogModel?.thinkingLevels ?? [])
+    : session.thinkingLevels;
+  const requestedThinkingLevel =
+    pendingConfig?.thinkingLevel ?? session.status?.thinkingLevel;
+  const currentThinkingLevel =
+    requestedThinkingLevel && thinkingLevels.includes(requestedThinkingLevel)
+      ? requestedThinkingLevel
+      : thinkingLevels[0];
+  const thinkingLevelsLoaded = isCloudSession
+    ? !catalogLoading
+    : session.thinkingLevelsLoaded;
+  const disabled = isTerminalCloudRun ? !isOnline : controlsDisabled;
+  const setModel = useCallback(
+    (model: PiModelSelection) => {
+      if (taskRunId && isTerminalCloudRun) {
+        const nextThinkingLevels =
+          catalog.find(
+            (candidate) =>
+              candidate.provider === model.provider &&
+              candidate.id === model.id,
+          )?.thinkingLevels ?? [];
+        const requestedThinkingLevel =
+          pendingConfig?.thinkingLevel ?? session.status?.thinkingLevel;
+        const thinkingLevel =
+          requestedThinkingLevel &&
+          nextThinkingLevels.includes(requestedThinkingLevel)
+            ? requestedThinkingLevel
+            : nextThinkingLevels[0];
+        setPendingConfig(taskId, taskRunId, { model, thinkingLevel });
+        return;
+      }
+
+      void controller
+        .setModel(taskId, model)
+        .catch((error) => onError(error, "Failed to change Pi model"));
+    },
+    [
+      catalog,
+      controller,
+      isTerminalCloudRun,
+      onError,
+      pendingConfig?.thinkingLevel,
+      session.status?.thinkingLevel,
+      setPendingConfig,
+      taskId,
+      taskRunId,
+    ],
+  );
+  const setThinkingLevel = useCallback(
+    (level: PiThinkingLevel) => {
+      if (taskRunId && isTerminalCloudRun) {
+        setPendingConfig(taskId, taskRunId, { thinkingLevel: level });
+        return;
+      }
+
+      void controller
+        .setThinkingLevel(taskId, level)
+        .catch((error) => onError(error, "Failed to change Pi thinking level"));
+    },
+    [
+      controller,
+      isTerminalCloudRun,
+      onError,
+      setPendingConfig,
+      taskId,
+      taskRunId,
+    ],
+  );
+
+  if (!modelsLoaded) {
+    return <Skeleton className="h-7 w-32 bg-foreground/15" />;
+  }
+
+  const supportsThinking = thinkingLevels.some((level) => level !== "off");
+  return (
+    <span className="flex gap-1">
+      <PiModelSelector
+        models={models}
+        currentModel={currentModel}
+        disabled={disabled}
+        onChange={setModel}
+      />
+      {currentThinkingLevel && thinkingLevelsLoaded && supportsThinking && (
+        <PiThinkingLevelSelector
+          level={currentThinkingLevel}
+          levels={thinkingLevels}
+          disabled={disabled}
+          onChange={setThinkingLevel}
+        />
+      )}
+    </span>
+  );
+}