From e8f8e2aec0221ba833ee68b9f8327c91f549407b Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Fri, 31 Jul 2026 11:23:04 +0200 Subject: [PATCH] fix(pi): preserve cloud model selection --- apps/web/package.json | 1 + apps/web/src/web-host-router.ts | 2 +- packages/agent/package.json | 4 - packages/agent/tsup.config.ts | 1 - .../pi-runtime/piSessionController.test.ts | 37 +++++ .../src/pi-runtime/piSessionController.ts | 23 +++ .../src/task-detail/taskCreationSaga.test.ts | 4 +- .../core/src/task-detail/taskCreationSaga.ts | 4 +- .../posthog-provider}/model-catalog.test.ts | 0 .../posthog-provider}/model-catalog.ts | 14 +- packages/harness/tsup.config.ts | 1 + .../pi-sessions/PiSessionModelControls.tsx | 148 ++++++++++++++++++ .../features/pi-sessions/PiSessionView.tsx | 98 ++++-------- .../pi-sessions/piPendingConfigStore.ts | 60 +++++++ .../features/pi-sessions/usePiModelCatalog.ts | 22 +++ .../task-detail/components/TaskInput.tsx | 29 ++-- packages/workspace-server/package.json | 1 + .../src/services/agent/agent.ts | 2 +- pnpm-lock.yaml | 6 + 19 files changed, 353 insertions(+), 104 deletions(-) rename packages/{agent/src/pi => harness/src/extensions/posthog-provider}/model-catalog.test.ts (100%) rename packages/{agent/src/pi => harness/src/extensions/posthog-provider}/model-catalog.ts (85%) create mode 100644 packages/ui/src/features/pi-sessions/PiSessionModelControls.tsx create mode 100644 packages/ui/src/features/pi-sessions/piPendingConfigStore.ts create mode 100644 packages/ui/src/features/pi-sessions/usePiModelCatalog.ts diff --git a/apps/web/package.json b/apps/web/package.json index bf1eff1d71..dfd6a082f2 100644 --- a/apps/web/package.json +++ b/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:*", diff --git a/apps/web/src/web-host-router.ts b/apps/web/src/web-host-router.ts index 07f148b42d..fb2ac0a73a 100644 --- a/apps/web/src/web-host-router.ts +++ b/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"; diff --git a/packages/agent/package.json b/packages/agent/package.json index c4e12b6365..3c9ebe3832 100644 --- a/packages/agent/package.json +++ b/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" diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index 55bfd03fea..b8167cd239 100644 --- a/packages/agent/tsup.config.ts +++ b/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", diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 897e1799ee..37df1cae8e 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/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(), diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index f2b202350d..f86aedb302 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -35,6 +35,11 @@ export type { export type PiModelSelection = Pick; +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 { 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 { + 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 { const session = await this.getPiSession(taskId); const status = await session.client.getState(); diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index 3fa05dc72a..1d3849110d 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/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, }), ); diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index 461c39fca4..da3f55a9ca 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/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, diff --git a/packages/agent/src/pi/model-catalog.test.ts b/packages/harness/src/extensions/posthog-provider/model-catalog.test.ts similarity index 100% rename from packages/agent/src/pi/model-catalog.test.ts rename to packages/harness/src/extensions/posthog-provider/model-catalog.test.ts diff --git a/packages/agent/src/pi/model-catalog.ts b/packages/harness/src/extensions/posthog-provider/model-catalog.ts similarity index 85% rename from packages/agent/src/pi/model-catalog.ts rename to packages/harness/src/extensions/posthog-provider/model-catalog.ts index 96b6bb45d5..9642b09249 100644 --- a/packages/agent/src/pi/model-catalog.ts +++ b/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, + "provider" +> & { provider: "posthog"; - id: string; name: string; - contextWindow: number; thinkingLevels: ModelThinkingLevel[]; -} +}; export function resolvePosthogPiModelCatalog( gatewayModels: GatewayModel[], diff --git a/packages/harness/tsup.config.ts b/packages/harness/tsup.config.ts index fe45566598..0ebc718d44 100644 --- a/packages/harness/tsup.config.ts +++ b/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", diff --git a/packages/ui/src/features/pi-sessions/PiSessionModelControls.tsx b/packages/ui/src/features/pi-sessions/PiSessionModelControls.tsx new file mode 100644 index 0000000000..bcd2449842 --- /dev/null +++ b/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 ; + } + + const supportsThinking = thinkingLevels.some((level) => level !== "off"); + return ( + + + {currentThinkingLevel && thinkingLevelsLoaded && supportsThinking && ( + + )} + + ); +} diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index d3e2c4317e..0e2ac2c9a5 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -5,10 +5,8 @@ import { } from "@posthog/core/message-editor/content"; import { PI_SESSION_CONTROLLER } from "@posthog/core/pi-runtime/identifiers"; import { - type PiModelSelection, PiOperationError, type PiSessionController, - type PiThinkingLevel, } from "@posthog/core/pi-runtime/piSessionController"; import { toPiContextUsage } from "@posthog/core/pi-runtime/piSessionUsage"; import { useService } from "@posthog/di/react"; @@ -22,7 +20,6 @@ import { Skeleton, } from "@posthog/quill"; import type { AgentConversationEvent } from "@posthog/shared"; -import { isTerminalStatus } from "@posthog/shared/domain-types"; import { useUsageLimitStore } from "@posthog/ui/features/billing/usageLimitStore"; import { PromptInput } from "@posthog/ui/features/message-editor/components/PromptInput"; import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; @@ -44,11 +41,12 @@ import { Box, Flex } from "@radix-ui/themes"; import { type ReactElement, useCallback, useEffect, useRef } from "react"; import { useStore } from "zustand"; import { PiQueuedMessagesDock } from "./PiQueuedMessagesDock"; +import { PiMessagingModeSelector } from "./PiSessionControls"; +import { PiSessionModelControls } from "./PiSessionModelControls"; import { - PiMessagingModeSelector, - PiModelSelector, - PiThinkingLevelSelector, -} from "./PiSessionControls"; + getPiPendingConfig, + usePiPendingConfigStore, +} from "./piPendingConfigStore"; const log = logger.scope("pi-session-view"); @@ -66,6 +64,12 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { (state) => state.sessions[taskId], ); const draftActions = useDraftStore((state) => state.actions); + const pendingConfig = usePiPendingConfigStore((state) => + getPiPendingConfig(state, taskId, taskRunId), + ); + const clearPendingConfig = usePiPendingConfigStore( + (state) => state.clearConfig, + ); const workspace = useWorkspace(taskId); const repoPath = workspace?.worktreePath ?? workspace?.folderPath; const messagingMode = useMessagingModeStore( @@ -152,8 +156,11 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { messagingMode, ); void piSessionController - .submit(taskId, message, isStreaming, messagingMode) + .submit(taskId, message, isStreaming, messagingMode, pendingConfig) .then(() => { + if (action === "prompt" && pendingConfig && taskRunId) { + clearPendingConfig(taskId, taskRunId); + } if (action === "compact") { toast.success("Pi context compacted"); } @@ -168,35 +175,16 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { }, [ handleControllerError, + clearPendingConfig, isStreaming, messagingMode, + pendingConfig, piSessionController, taskId, + taskRunId, ], ); - const setModel = useCallback( - (model: PiModelSelection) => { - void piSessionController - .setModel(taskId, model) - .catch((error) => - handleControllerError(error, "Failed to change Pi model"), - ); - }, - [handleControllerError, piSessionController, taskId], - ); - - const setThinkingLevel = useCallback( - (level: PiThinkingLevel) => { - void piSessionController - .setThinkingLevel(taskId, level) - .catch((error) => - handleControllerError(error, "Failed to change Pi thinking level"), - ); - }, - [handleControllerError, piSessionController, taskId], - ); - const toggleMessagingMode = useCallback(() => { const nextMode = messagingMode === "steer" ? "queue" : "steer"; setMessagingMode(taskId, nextMode); @@ -354,49 +342,12 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { } const controlsPending = status ? isStreaming || isBashRunning : false; - const controlsDisabled = - controlsPending || - isCompacting || - session.connectionState !== "connected" || - (session.cloudStatus !== undefined && - isTerminalStatus(session.cloudStatus)); const hasQueuedMessage = session.queue.steering.length + session.queue.followUp.length > 0; - let modelSelector: ReactElement = ( - - ); - let reasoningSelector: ReactElement | null = ( - - ); let messagingModeToggle: ReactElement = ( ); - if (status && session.modelsLoaded) { - modelSelector = ( - - ); - } - - if (status && session.thinkingLevelsLoaded) { - const supportsThinking = session.thinkingLevels.some( - (level) => level !== "off", - ); - reasoningSelector = supportsThinking ? ( - - ) : null; - } - if (status) { messagingModeToggle = ( + } + reasoningSelector={null} messagingModeToggle={messagingModeToggle} onToggleMessagingMode={toggleMessagingMode} onPromptRecall={handlePromptRecall} diff --git a/packages/ui/src/features/pi-sessions/piPendingConfigStore.ts b/packages/ui/src/features/pi-sessions/piPendingConfigStore.ts new file mode 100644 index 0000000000..8dd822afb6 --- /dev/null +++ b/packages/ui/src/features/pi-sessions/piPendingConfigStore.ts @@ -0,0 +1,60 @@ +import type { + PiModelSelection, + PiThinkingLevel, +} from "@posthog/core/pi-runtime/piSessionController"; +import { electronStorage } from "@posthog/ui/shell/rendererStorage"; +import { create } from "zustand"; +import { persist } from "zustand/middleware"; + +/** Deferred Pi settings selected while a cloud sandbox is inactive. They are applied after its next startup and before the triggering prompt. */ +export interface PiPendingConfig { + model?: PiModelSelection; + thinkingLevel?: PiThinkingLevel; +} + +interface PiPendingConfigState { + configsByRunKey: Record; + setConfig: (taskId: string, runId: string, config: PiPendingConfig) => void; + clearConfig: (taskId: string, runId: string) => void; +} + +function runKey(taskId: string, runId: string): string { + return `${taskId}:${runId}`; +} + +export const usePiPendingConfigStore = create()( + persist( + (set) => ({ + configsByRunKey: {}, + setConfig: (taskId, runId, config) => + set((state) => { + const key = runKey(taskId, runId); + return { + configsByRunKey: { + ...state.configsByRunKey, + [key]: { ...state.configsByRunKey[key], ...config }, + }, + }; + }), + clearConfig: (taskId, runId) => + set((state) => { + const { [runKey(taskId, runId)]: _, ...configsByRunKey } = + state.configsByRunKey; + return { configsByRunKey }; + }), + }), + { + name: "pi-pending-config-storage", + storage: electronStorage, + partialize: (state) => ({ configsByRunKey: state.configsByRunKey }), + }, + ), +); + +export function getPiPendingConfig( + state: PiPendingConfigState, + taskId: string, + runId: string | undefined, +): PiPendingConfig | undefined { + return runId ? state.configsByRunKey[runKey(taskId, runId)] : undefined; +} diff --git a/packages/ui/src/features/pi-sessions/usePiModelCatalog.ts b/packages/ui/src/features/pi-sessions/usePiModelCatalog.ts new file mode 100644 index 0000000000..7dab0bcd9f --- /dev/null +++ b/packages/ui/src/features/pi-sessions/usePiModelCatalog.ts @@ -0,0 +1,22 @@ +import { useHostTRPC } from "@posthog/host-router/react"; +import { getCloudUrlFromRegion } from "@posthog/shared"; +import { useAuthStateValue } from "@posthog/ui/features/auth/store"; +import { useQuery } from "@tanstack/react-query"; +import { useMemo } from "react"; + +export function usePiModelCatalog(enabled: boolean) { + const trpc = useHostTRPC(); + const cloudRegion = useAuthStateValue((state) => state.cloudRegion); + const apiHost = useMemo( + () => (cloudRegion ? getCloudUrlFromRegion(cloudRegion) : null), + [cloudRegion], + ); + + return useQuery({ + ...trpc.agent.getPiModelCatalog.queryOptions({ + apiHost: apiHost ?? "", + region: cloudRegion ?? "us", + }), + enabled: enabled && apiHost !== null, + }); +} diff --git a/packages/ui/src/features/task-detail/components/TaskInput.tsx b/packages/ui/src/features/task-detail/components/TaskInput.tsx index ae2adb4e78..2c0a9a8dc8 100644 --- a/packages/ui/src/features/task-detail/components/TaskInput.tsx +++ b/packages/ui/src/features/task-detail/components/TaskInput.tsx @@ -13,11 +13,7 @@ import { isValidConfigValue } from "@posthog/core/task-detail/configOptions"; import { useServiceOptional } from "@posthog/di/react"; import { useHostTRPC, useHostTRPCClient } from "@posthog/host-router/react"; import { ButtonGroup } from "@posthog/quill"; -import { - type AgentRuntime, - ANALYTICS_EVENTS, - getCloudUrlFromRegion, -} from "@posthog/shared"; +import { type AgentRuntime, ANALYTICS_EVENTS } from "@posthog/shared"; import type { Task } from "@posthog/shared/domain-types"; import { openSettings } from "@posthog/ui/features/settings/hooks/useOpenSettings"; import type { TaskInputReportAssociation } from "@posthog/ui/features/task-detail/stores/taskInputPrefillStore"; @@ -81,6 +77,7 @@ import { PiModelSelector, PiThinkingLevelSelector, } from "../../pi-sessions/PiSessionControls"; +import { usePiModelCatalog } from "../../pi-sessions/usePiModelCatalog"; import { DropZoneOverlay } from "../../sessions/components/DropZoneOverlay"; import { ReasoningLevelSelector } from "../../sessions/components/ReasoningLevelSelector"; import { getCurrentModeFromConfigOptions } from "../../sessions/sessionStore"; @@ -261,17 +258,8 @@ export function TaskInput({ ); const [selectedPiThinkingLevel, setSelectedPiThinkingLevel] = useState(null); - const piApiHost = useMemo( - () => (cloudRegion ? getCloudUrlFromRegion(cloudRegion) : null), - [cloudRegion], - ); - const { data: piModelCatalog = [], isPending: isPiConfigLoading } = useQuery({ - ...trpc.agent.getPiModelCatalog.queryOptions({ - apiHost: piApiHost ?? "", - region: cloudRegion ?? "us", - }), - enabled: runtime === "pi" && piApiHost !== null, - }); + const { data: piModelCatalog = [], isPending: isPiConfigLoading } = + usePiModelCatalog(runtime === "pi"); const [cloudRepoSearchQuery, setCloudRepoSearchQuery] = useState(""); const [isCloudRepoPickerOpen, setIsCloudRepoPickerOpen] = useState(false); const [cloudBranchSearchQuery, setCloudBranchSearchQuery] = useState(""); @@ -779,6 +767,7 @@ export function TaskInput({ ) ? (selectedPiThinkingLevel ?? "high") : piThinkingLevels[0]; + const supportsPiThinking = piThinkingLevels.some((level) => level !== "off"); const currentContextWindow = contextWindowOption?.type === "select" && (contextWindowOption.currentValue === "200k" || @@ -809,7 +798,11 @@ export function TaskInput({ : currentReasoningLevel; const taskModel = runtime === "pi" ? currentPiModel?.id : effectiveModel; const taskReasoningLevel = - runtime === "pi" ? currentPiThinkingLevel : effectiveReasoningLevel; + runtime === "pi" + ? supportsPiThinking + ? currentPiThinkingLevel + : undefined + : effectiveReasoningLevel; useWarmTask({ workspaceMode, @@ -1427,7 +1420,7 @@ export function TaskInput({ } reasoningSelector={ autoresearchDraft ? null : runtime === "pi" ? ( - currentPiThinkingLevel ? ( + currentPiThinkingLevel && supportsPiThinking ? (