fix(pi): honor per-model thinking configuration (#4413)

Derive Pi thinking options from each model's thinkingLevelMap, read Pi's effective level back after create, resume, set_model and set_thinking_level, surface failed state reads instead of persisting the requested level, and show the thinking control for single-level models.

Closes #4382

Co-authored-by: therainisme <therainisme@qq.com>
Co-authored-by: trim21 <i@trim21.me>

packages/app/e2e/browser/command-center-agent-controls.spec.ts

@@ -1,4 +1,4 @@
-import { test } from "../support/fixtures";
+import { expect, test } from "../support/fixtures";
 import type { FormPreferences } from "@/create-agent-preferences/preferences";
 import { closeCommandCenter, openCommandCenter } from "../support/helpers/command-center";
 import {
@@ -79,6 +79,27 @@ test.describe("Command Center agent controls", () => {
     }
   });
 
+  test("shows an existing agent's only supported thinking level", async ({ page }) => {
+    const workspace = await seedMockAgentWorkspace({
+      repoPrefix: "agent-controls-single-thinking-",
+      title: "Single thinking level",
+      model: "max-only-thinking-stream",
+      thinkingOptionId: "max",
+    });
+    try {
+      await openAgentRoute(page, {
+        workspaceId: workspace.workspaceId,
+        agentId: workspace.agentId,
+      });
+
+      await expect(
+        page.getByTestId("agent-thinking-selector").filter({ visible: true }),
+      ).toHaveAccessibleName("Select thinking option (Max)");
+    } finally {
+      await workspace.cleanup();
+    }
+  });
+
   test("activates the best-matching row, not the highest-ranked group", async ({ page }) => {
     // "Mode › Load Test" matches "load" at offset 0; every "Model › Mock Load Test › …" row
     // matches at offset 5. Models used to outrank modes by fixed group rank, so Enter picked a

packages/app/e2e/support/helpers/mock-agent.ts

@@ -18,6 +18,7 @@ export interface MockAgentOptions {
   port?: number;
   initialPrompt?: string;
   model?: string;
+  thinkingOptionId?: string;
   modeId?: string;
   featureValues?: Record<string, unknown>;
 }
@@ -43,6 +44,7 @@ export async function seedMockAgentWorkspace(
       title: options.title,
       modeId: options.modeId ?? "load-test",
       model: options.model ?? "e2e-fast-stream",
+      thinkingOptionId: options.thinkingOptionId,
       initialPrompt: options.initialPrompt,
       featureValues: options.featureValues,
     });

packages/app/src/composer/agent-controls/index.tsx

@@ -1788,7 +1788,7 @@ export const AgentControls = memo(function AgentControls({
         onEditAgentProfiles={handleEditAgentProfiles}
         onCreateAgentProfile={profileActions.create}
         onEditAgentProfile={profileActions.edit}
-        thinkingOptions={thinkingOptions.length > 1 ? thinkingOptions : undefined}
+        thinkingOptions={thinkingOptions.length > 0 ? thinkingOptions : undefined}
         selectedThinkingOptionId={modelSelection.selectedThinkingId ?? undefined}
         onSelectThinkingOption={handleSelectThinkingOption}
         features={agent.features}

packages/server/src/server/agent/agent-manager.test.ts

@@ -4767,6 +4767,93 @@ test("later explicit config mutations win over events emitted by earlier mutatio
   expect(manager.getAgent(snapshot.id)?.config.thinkingOptionId).toBe("high");
 });
 
+test("setAgentThinkingOption adopts the session's effective clamped level", async () => {
+  const workdir = mkdtempSync(join(tmpdir(), "agent-manager-thinking-clamp-"));
+  class ClampingSession extends TestAgentSession {
+    async setThinkingOption(): Promise<void> {}
+
+    override async getRuntimeInfo() {
+      return {
+        ...(await super.getRuntimeInfo()),
+        thinkingOptionId: "high",
+      };
+    }
+  }
+  class ClampingClient extends TestAgentClient {
+    override async createSession(config: AgentSessionConfig): Promise<AgentSession> {
+      return new ClampingSession(config);
+    }
+  }
+
+  const storage = new AgentStorage(join(workdir, "agents"), logger);
+  const manager = new AgentManager({
+    clients: { codex: new ClampingClient() },
+    registry: storage,
+    logger,
+    idFactory: () => "00000000-0000-4000-8000-000000000135",
+  });
+  const snapshot = await manager.createAgent(
+    {
+      provider: "codex",
+      cwd: workdir,
+      thinkingOptionId: "low",
+    },
+    undefined,
+    { workspaceId: undefined },
+  );
+
+  await manager.setAgentThinkingOption(snapshot.id, "medium");
+  await manager.flush();
+
+  const agent = manager.getAgent(snapshot.id);
+  expect(agent?.config.thinkingOptionId).toBe("high");
+  expect(agent?.runtimeInfo?.thinkingOptionId).toBe("high");
+  const persisted = await storage.get(snapshot.id);
+  expect(persisted?.config?.thinkingOptionId).toBe("high");
+});
+
+test("setAgentThinkingOption surfaces a failed state read and keeps the previous level", async () => {
+  const workdir = mkdtempSync(join(tmpdir(), "agent-manager-thinking-read-failure-"));
+  class UnreadableStateSession extends TestAgentSession {
+    async setThinkingOption(): Promise<void> {}
+
+    override async getRuntimeInfo(): Promise<never> {
+      throw new Error("state unavailable");
+    }
+  }
+  class UnreadableStateClient extends TestAgentClient {
+    override async createSession(config: AgentSessionConfig): Promise<AgentSession> {
+      return new UnreadableStateSession(config);
+    }
+  }
+
+  const storage = new AgentStorage(join(workdir, "agents"), logger);
+  const manager = new AgentManager({
+    clients: { codex: new UnreadableStateClient() },
+    registry: storage,
+    logger,
+    idFactory: () => "00000000-0000-4000-8000-000000000136",
+  });
+  const snapshot = await manager.createAgent(
+    {
+      provider: "codex",
+      cwd: workdir,
+      thinkingOptionId: "low",
+    },
+    undefined,
+    { workspaceId: undefined },
+  );
+
+  await expect(manager.setAgentThinkingOption(snapshot.id, "medium")).rejects.toThrow(
+    "state unavailable",
+  );
+  await manager.flush();
+
+  expect(manager.getAgent(snapshot.id)?.config.thinkingOptionId).toBe("low");
+  const persisted = await storage.get(snapshot.id);
+  expect(persisted?.config?.thinkingOptionId).toBe("low");
+});
+
 test("session config drift events update state through the stream channel", async () => {
   const workdir = mkdtempSync(join(tmpdir(), "agent-manager-session-config-events-"));
   let capturedSession: TestAgentSession | null = null;

packages/server/src/server/agent/agent-manager.ts

@@ -1934,11 +1934,17 @@ export class AgentManager {
     }
     await this.drainSessionEvents(agentId);
 
-    agent.config.thinkingOptionId = normalizedThinkingOptionId ?? undefined;
+    let effectiveThinkingOptionId = normalizedThinkingOptionId;
+    const runtimeInfo = await agent.session.getRuntimeInfo();
+    if (runtimeInfo.thinkingOptionId !== undefined) {
+      effectiveThinkingOptionId = runtimeInfo.thinkingOptionId;
+    }
+
+    agent.config.thinkingOptionId = effectiveThinkingOptionId ?? undefined;
     if (agent.runtimeInfo) {
       agent.runtimeInfo = {
         ...agent.runtimeInfo,
-        thinkingOptionId: normalizedThinkingOptionId,
+        thinkingOptionId: effectiveThinkingOptionId,
       };
     }
     this.touchUpdatedAt(agent);

packages/server/src/server/agent/providers/mock-load-test-agent.ts

@@ -145,6 +145,18 @@ const MODELS: AgentModelDefinition[] = [
       intervalMs: 5,
     },
   },
+  {
+    provider: MOCK_LOAD_TEST_PROVIDER_ID,
+    id: "max-only-thinking-stream",
+    label: "Max-only thinking stream",
+    description: "Fast realistic stream with one supported thinking level.",
+    thinkingOptions: [{ id: "max", label: "Max", isDefault: true }],
+    defaultThinkingOptionId: "max",
+    metadata: {
+      durationMs: 10_000,
+      intervalMs: 5,
+    },
+  },
   {
     provider: MOCK_LOAD_TEST_PROVIDER_ID,
     id: "e2e-fast-stream",

packages/server/src/server/agent/providers/pi/agent.test.ts

@@ -24,11 +24,35 @@ import {
   transformPiModels,
 } from "./agent.js";
 import { FakePi } from "./test-utils/fake-pi.js";
+import type { PiModel, PiThinkingLevel } from "./rpc-types.js";
 import type { PiUsagePollScheduler } from "./usage-poller.js";
 
 const ONE_BY_ONE_PNG_BASE64 =
   "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII=";
 
+const RESTRICTED_THINKING_MODEL: PiModel = {
+  provider: "kimi-coding",
+  id: "kimi-k3",
+  name: "Kimi K3",
+  reasoning: true,
+  thinkingLevelMap: {
+    off: null,
+    minimal: null,
+    low: "low",
+    medium: null,
+    high: "high",
+    xhigh: null,
+    max: "max",
+  },
+};
+
+interface PiThinkingCatalogCase {
+  name: string;
+  model: PiModel;
+  optionIds: PiThinkingLevel[] | undefined;
+  defaultThinkingOptionId: PiThinkingLevel | undefined;
+}
+
 test("Pi RPC timeout defaults to 60 seconds and accepts an override", () => {
   expect(PiProviderParamsSchema.parse({}).rpcTimeoutMs).toBe(60_000);
   expect(PiProviderParamsSchema.parse({ rpcTimeoutMs: 90_000 }).rpcTimeoutMs).toBe(90_000);
@@ -179,6 +203,51 @@ test("keeps normal Pi agent sessions persisted", async () => {
   await session.close();
 });
 
+test("lets Pi choose the thinking level when none is requested", async () => {
+  const pi = new FakePi();
+  pi.queueSessionSetup((session) => {
+    session.state = {
+      ...session.state,
+      model: RESTRICTED_THINKING_MODEL,
+      thinkingLevel: "max",
+    };
+  });
+  const session = await createClient(pi).createSession(
+    createConfig({ model: "kimi-coding/kimi-k3" }),
+  );
+  onTestFinished(() => session.close());
+
+  expect(pi.recordedLaunches[0]?.thinkingOptionId).toBeUndefined();
+  expect(pi.recordedLaunches[0]?.argv).not.toContain("--thinking");
+  expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe("max");
+  await expect(session.getRuntimeInfo()).resolves.toMatchObject({ thinkingOptionId: "max" });
+});
+
+test.each([
+  { requested: "medium", effective: "high" },
+  { requested: "xhigh", effective: "max" },
+] as const)(
+  "adopts Pi's $effective level when creating with $requested",
+  async ({ requested, effective }) => {
+    const pi = new FakePi();
+    pi.queueSessionSetup((session) => {
+      session.state = {
+        ...session.state,
+        model: RESTRICTED_THINKING_MODEL,
+        thinkingLevel: effective,
+      };
+    });
+    const session = await createClient(pi).createSession(
+      createConfig({ model: "kimi-coding/kimi-k3", thinkingOptionId: requested }),
+    );
+    onTestFinished(() => session.close());
+
+    expect(pi.recordedLaunches[0]?.thinkingOptionId).toBe(requested);
+    expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe(effective);
+    await expect(session.getRuntimeInfo()).resolves.toMatchObject({ thinkingOptionId: effective });
+  },
+);
+
 class SessionEvents {
   private readonly events: AgentStreamEvent[] = [];
   private readonly waiters: Array<{
@@ -1370,6 +1439,32 @@ describe("PiRpcAgentSession", () => {
     ]);
   });
 
+  test("adopts Pi's clamped thinking level when resuming a session", async () => {
+    const pi = new FakePi();
+    pi.queueSessionSetup((session) => {
+      session.state = {
+        ...session.state,
+        model: RESTRICTED_THINKING_MODEL,
+        thinkingLevel: "high",
+      };
+    });
+    const session = await createClient(pi).resumeSession({
+      provider: "pi",
+      sessionId: "pi-session-1",
+      nativeHandle: "/tmp/native-pi-session",
+      metadata: {
+        cwd: "/workspace/project",
+        model: "kimi-coding/kimi-k3",
+        thinkingOptionId: "medium",
+      },
+    });
+    onTestFinished(() => session.close());
+
+    expect(pi.recordedLaunches[0]?.thinkingOptionId).toBe("medium");
+    expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe("high");
+    await expect(session.getRuntimeInfo()).resolves.toMatchObject({ thinkingOptionId: "high" });
+  });
+
   test("reports the persisted Pi entry attached to the submitted message", async () => {
     const pi = new FakePi();
     const client = createClient(pi);
@@ -1436,8 +1531,6 @@ describe("PiRpcAgentSession", () => {
       "pi",
       "--mode",
       "rpc",
-      "--thinking",
-      "medium",
       "--extension",
       actualLaunch.extensionPaths[0],
     ]);
@@ -1507,6 +1600,91 @@ describe("PiRpcAgentSession", () => {
     expect(fakeSession.setThinkingLevelRequests).toEqual(["high"]);
   });
 
+  test.each([
+    { requested: "medium", effective: "high", rpcRequest: "medium" },
+    { requested: "xhigh", effective: "max", rpcRequest: "xhigh" },
+    { requested: null, effective: "high", rpcRequest: "medium" },
+  ] as const)(
+    "stores Pi's $effective thinking level after requesting $requested",
+    async ({ requested, effective, rpcRequest }) => {
+      const pi = new FakePi();
+      pi.queueSessionSetup((session) => {
+        session.state = {
+          ...session.state,
+          model: RESTRICTED_THINKING_MODEL,
+          thinkingLevel: "low",
+        };
+      });
+      const { session } = await createSession(pi);
+      onTestFinished(() => session.close());
+      const fakeSession = pi.latestSession();
+      fakeSession.state = { ...fakeSession.state, thinkingLevel: effective };
+
+      await session.setThinkingOption(requested);
+
+      expect(fakeSession.setThinkingLevelRequests).toEqual([rpcRequest]);
+      expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe(effective);
+      await expect(session.getRuntimeInfo()).resolves.toEqual({
+        provider: "pi",
+        sessionId: "pi-session-1",
+        model: "kimi-coding/kimi-k3",
+        thinkingOptionId: effective,
+        modeId: null,
+      });
+    },
+  );
+
+  test("refreshes Pi's thinking level after changing models", async () => {
+    const { pi, session } = await createSession();
+    onTestFinished(() => session.close());
+    const fakeSession = pi.latestSession();
+    fakeSession.setModelResult = RESTRICTED_THINKING_MODEL;
+    fakeSession.state = { ...fakeSession.state, thinkingLevel: "high" };
+
+    await session.setModel("kimi-coding/kimi-k3");
+
+    expect(fakeSession.setModelRequests).toEqual([{ provider: "kimi-coding", modelId: "kimi-k3" }]);
+    expect(session.describePersistence()?.metadata).toEqual({
+      cwd: "/tmp/paseo-pi-rpc-test",
+      model: "kimi-coding/kimi-k3",
+      thinkingOptionId: "high",
+    });
+    await expect(session.getRuntimeInfo()).resolves.toMatchObject({
+      model: "kimi-coding/kimi-k3",
+      thinkingOptionId: "high",
+    });
+  });
+
+  test("reports updated Pi thinking state instead of a cached selection", async () => {
+    const { pi, session } = await createSession();
+    onTestFinished(() => session.close());
+    const fakeSession = pi.latestSession();
+    fakeSession.state = { ...fakeSession.state, thinkingLevel: "off" };
+
+    await expect(session.getRuntimeInfo()).resolves.toMatchObject({ thinkingOptionId: "off" });
+    expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe("off");
+  });
+
+  test("surfaces state refresh failures after a Pi thinking update", async () => {
+    const { pi, session } = await createSession();
+    onTestFinished(() => session.close());
+    pi.latestSession().getStateError = new Error("Pi state unavailable");
+
+    await expect(session.setThinkingOption("high")).rejects.toThrow("Pi state unavailable");
+    expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe("medium");
+  });
+
+  test("surfaces state refresh failures after a Pi model update", async () => {
+    const { pi, session } = await createSession();
+    onTestFinished(() => session.close());
+    const fakeSession = pi.latestSession();
+    fakeSession.setModelResult = RESTRICTED_THINKING_MODEL;
+    fakeSession.getStateError = new Error("Pi state unavailable");
+
+    await expect(session.setModel("kimi-coding/kimi-k3")).rejects.toThrow("Pi state unavailable");
+    expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe("medium");
+  });
+
   test("materializes image prompts as text hints for text-only Pi models", async () => {
     const { pi, session } = await createSession();
     const fakeSession = pi.latestSession();
@@ -2253,6 +2431,9 @@ describe("PiRpcAgentClient", () => {
       "utf8",
     );
     const pi = new FakePi();
+    pi.queueSessionSetup((session) => {
+      session.state.thinkingLevel = "high";
+    });
     const client = new PiRpcAgentClient({
       logger: pino({ level: "silent" }),
       runtime: pi,
@@ -2324,6 +2505,113 @@ describe("PiRpcAgentClient", () => {
     expect(pi.recordedLaunches[0]).toMatchObject({ cwd: "/workspace/with-extension" });
   });
 
+  test("honors per-model Pi thinking maps and clamps the catalog default upward", async () => {
+    const pi = new FakePi();
+    pi.queueSessionSetup((session) => {
+      session.models = [RESTRICTED_THINKING_MODEL];
+    });
+    const client = createClient(pi);
+
+    const catalog = await client.fetchCatalog({
+      scope: "workspace",
+      cwd: "/workspace/project",
+      force: false,
+    });
+
+    expect(catalog.models).toEqual([
+      {
+        provider: "pi",
+        id: "kimi-coding/kimi-k3",
+        label: "Kimi K3",
+        description: "kimi-coding/kimi-k3",
+        metadata: { provider: "kimi-coding", modelId: "kimi-k3" },
+        thinkingOptions: [
+          { id: "low", label: "Low", description: "Faster reasoning" },
+          { id: "high", label: "High", description: "Deeper reasoning", isDefault: true },
+          { id: "max", label: "Max", description: "Extreme reasoning" },
+        ],
+        defaultThinkingOptionId: "high",
+      },
+    ]);
+  });
+
+  test.each<PiThinkingCatalogCase>([
+    {
+      name: "no thinking map",
+      model: { ...RESTRICTED_THINKING_MODEL, thinkingLevelMap: undefined },
+      optionIds: ["off", "minimal", "low", "medium", "high"],
+      defaultThinkingOptionId: "medium",
+    },
+    {
+      name: "a partial thinking map",
+      model: {
+        ...RESTRICTED_THINKING_MODEL,
+        thinkingLevelMap: { minimal: null, medium: null, xhigh: "extra-high", max: null },
+      },
+      optionIds: ["off", "low", "high", "xhigh"],
+      defaultThinkingOptionId: "high",
+    },
+    {
+      name: "explicitly mapped extended levels",
+      model: {
+        ...RESTRICTED_THINKING_MODEL,
+        thinkingLevelMap: { xhigh: "extra-high", max: "maximum" },
+      },
+      optionIds: ["off", "minimal", "low", "medium", "high", "xhigh", "max"],
+      defaultThinkingOptionId: "medium",
+    },
+    {
+      name: "only lower thinking levels",
+      model: {
+        ...RESTRICTED_THINKING_MODEL,
+        thinkingLevelMap: { medium: null, high: null },
+      },
+      optionIds: ["off", "minimal", "low"],
+      defaultThinkingOptionId: "low",
+    },
+    {
+      name: "a non-reasoning model",
+      model: { ...RESTRICTED_THINKING_MODEL, reasoning: false },
+      optionIds: undefined,
+      defaultThinkingOptionId: undefined,
+    },
+    {
+      name: "no supported thinking levels",
+      model: {
+        ...RESTRICTED_THINKING_MODEL,
+        thinkingLevelMap: {
+          ...RESTRICTED_THINKING_MODEL.thinkingLevelMap,
+          low: null,
+          high: null,
+          max: null,
+        },
+      },
+      optionIds: [],
+      defaultThinkingOptionId: "off",
+    },
+  ])(
+    "resolves Pi thinking options for $name",
+    async ({ model, optionIds, defaultThinkingOptionId }) => {
+      const pi = new FakePi();
+      pi.queueSessionSetup((session) => {
+        session.models = [model];
+      });
+      const catalog = await createClient(pi).fetchCatalog({
+        scope: "workspace",
+        cwd: "/workspace/project",
+        force: false,
+      });
+
+      expect(catalog.models).toHaveLength(1);
+      const thinkingOptions = catalog.models[0]?.thinkingOptions;
+      expect(thinkingOptions?.map((option) => option.id)).toEqual(optionIds);
+      expect(catalog.models[0]?.defaultThinkingOptionId).toBe(defaultThinkingOptionId);
+      expect(
+        thinkingOptions?.filter((option) => option.isDefault).map((option) => option.id),
+      ).toEqual(optionIds?.filter((id) => id === defaultThinkingOptionId));
+    },
+  );
+
   test("lists no draft features without starting a Pi session", async () => {
     const pi = new FakePi();
     const client = createClient(pi);
@@ -2622,8 +2910,6 @@ describe("PiRpcAgentClient", () => {
       "pi",
       "--mode",
       "rpc",
-      "--thinking",
-      "medium",
       "--mcp-config",
       actualLaunch.mcpConfigPath,
       "--extension",
@@ -2704,8 +2990,6 @@ describe("PiRpcAgentClient", () => {
       "pi",
       "--mode",
       "rpc",
-      "--thinking",
-      "medium",
       "--extension",
       actualLaunch.extensionPaths[0],
     ]);

packages/server/src/server/agent/providers/pi/agent.ts

@@ -24,6 +24,7 @@ import {
   type AgentRunOptions,
   type AgentRunResult,
   type AgentRuntimeInfo,
+  type AgentSelectOption,
   type AgentSession,
   type AgentSessionConfig,
   type AgentSlashCommand,
@@ -171,12 +172,11 @@ const PI_THINKING_OPTIONS: ReadonlyArray<{
   id: PiThinkingLevel;
   label: string;
   description: string;
-  isDefault?: boolean;
 }> = [
   { id: "off", label: "Off", description: "No extra reasoning" },
   { id: "minimal", label: "Minimal", description: "Light reasoning" },
   { id: "low", label: "Low", description: "Faster reasoning" },
-  { id: "medium", label: "Medium", description: "Balanced reasoning", isDefault: true },
+  { id: "medium", label: "Medium", description: "Balanced reasoning" },
   { id: "high", label: "High", description: "Deeper reasoning" },
   { id: "xhigh", label: "XHigh", description: "Very deep reasoning" },
   { id: "max", label: "Max", description: "Extreme reasoning" },
@@ -373,21 +373,6 @@ function parseAutoCompactMode(value: string | undefined): AutoCompactMode {
   return "unknown";
 }
 
-function mapThinkingOption(option: (typeof PI_THINKING_OPTIONS)[number]) {
-  const mappedOption = {
-    id: option.id,
-    label: option.label,
-    description: option.description,
-  };
-  if (option.isDefault) {
-    return {
-      ...mappedOption,
-      isDefault: true,
-    };
-  }
-  return mappedOption;
-}
-
 function piModelSupportsImageInput(model: PiModel | null | undefined): boolean {
   return model?.input?.includes("image") === true;
 }
@@ -791,14 +776,6 @@ function isPiRequestAbortError(error: unknown): boolean {
   return /\brequest was aborted\b|\babort(ed)?\b/i.test(toDiagnosticErrorMessage(error));
 }
 
-function resolveThinkingOptionId(
-  cachedThinkingOptionId: string | null,
-  sessionThinkingLevel: PiThinkingLevel,
-): PiThinkingLevel | null {
-  const currentThinking = cachedThinkingOptionId ?? sessionThinkingLevel;
-  return normalizePiThinkingOption(currentThinking);
-}
-
 function modelToId(model: PiModel | null | undefined): string | null {
   return model?.provider && model.id ? `${model.provider}/${model.id}` : null;
 }
@@ -1195,7 +1172,50 @@ function buildExtensionUiResponse(
   return { value: answer };
 }
 
+function resolvePiThinkingConfig(
+  model: PiModel,
+): Pick<AgentModelDefinition, "thinkingOptions" | "defaultThinkingOptionId"> {
+  if (!model.reasoning) {
+    return { thinkingOptions: undefined, defaultThinkingOptionId: undefined };
+  }
+
+  const supportedOptions = PI_THINKING_OPTIONS.filter((option) => {
+    const mapped = model.thinkingLevelMap?.[option.id];
+    if (mapped === null) {
+      return false;
+    }
+    if (option.id === "xhigh" || option.id === "max") {
+      return mapped !== undefined;
+    }
+    return true;
+  });
+  const defaultIndex = PI_THINKING_OPTIONS.findIndex(
+    (option) => option.id === DEFAULT_PI_THINKING_LEVEL,
+  );
+  // Pi clamps upward first, then falls back to the highest remaining lower level.
+  const higherDefault = supportedOptions.find(
+    (option) => PI_THINKING_OPTIONS.indexOf(option) >= defaultIndex,
+  );
+  const defaultOption = higherDefault ?? supportedOptions.at(-1);
+  const defaultThinkingOptionId = defaultOption?.id ?? "off";
+  return {
+    thinkingOptions: supportedOptions.map((option) => {
+      const mappedOption: AgentSelectOption = {
+        id: option.id,
+        label: option.label,
+        description: option.description,
+      };
+      if (option.id === defaultThinkingOptionId) {
+        mappedOption.isDefault = true;
+      }
+      return mappedOption;
+    }),
+    defaultThinkingOptionId,
+  };
+}
+
 function mapPiModel(model: PiModel, provider: AgentProvider): AgentModelDefinition {
+  const { thinkingOptions, defaultThinkingOptionId } = resolvePiThinkingConfig(model);
   return {
     provider,
     id: `${model.provider}/${model.id}`,
@@ -1205,8 +1225,8 @@ function mapPiModel(model: PiModel, provider: AgentProvider): AgentModelDefiniti
       provider: model.provider,
       modelId: model.id,
     },
-    thinkingOptions: model.reasoning ? PI_THINKING_OPTIONS.map(mapThinkingOption) : undefined,
-    defaultThinkingOptionId: model.reasoning ? DEFAULT_PI_THINKING_LEVEL : undefined,
+    thinkingOptions,
+    defaultThinkingOptionId,
   };
 }
 
@@ -1244,7 +1264,6 @@ export class PiRpcAgentSession implements AgentSession {
   private activePromptRequestId: string | null = null;
   private readonly pendingPromptResults = new Map<string, boolean>();
   private readonly pendingSteerSubmissions: PiPendingSteerSubmission[] = [];
-  private lastKnownThinkingOptionId: string | null;
   currentLeafOverrideId: string | null | undefined;
   private readonly capturedUserEntries: PiCapturedEntry[] = [];
   private readonly capturedUserEntriesById = new Map<string, PiCapturedEntry>();
@@ -1266,14 +1285,11 @@ export class PiRpcAgentSession implements AgentSession {
     this.runtimeSession = options.runtimeSession;
     this.config = options.config;
     this.state = options.initialState;
+    this.config.thinkingOptionId = this.state.thinkingLevel;
     this.capabilities = options.capabilities;
     this.provider = PI_PROVIDER;
     this.currentModeId = options.currentModeId ?? null;
     this.cleanup = options.cleanup;
-    this.lastKnownThinkingOptionId =
-      normalizePiThinkingOption(options.config.thinkingOptionId) ??
-      this.state.thinkingLevel ??
-      null;
     this.extensionTimeoutMs = options.extensionTimeoutMs ?? DEFAULT_PI_EXTENSION_RESULT_TIMEOUT_MS;
     this.logger = options.logger;
     this.usagePoller = new PiUsagePoller({
@@ -1469,10 +1485,7 @@ export class PiRpcAgentSession implements AgentSession {
       provider: this.provider,
       sessionId: this.state.sessionId,
       model: modelToId(this.state.model),
-      thinkingOptionId: resolveThinkingOptionId(
-        this.lastKnownThinkingOptionId,
-        this.state.thinkingLevel,
-      ),
+      thinkingOptionId: this.state.thinkingLevel,
       modeId: this.currentModeId,
     };
   }
@@ -1689,22 +1702,14 @@ export class PiRpcAgentSession implements AgentSession {
     }
 
     const model = await this.runtimeSession.setModel(parsedReference.provider, parsedReference.id);
-    this.state = {
-      ...this.state,
-      model,
-    };
+    await this.refreshState();
     this.config.model = `${model.provider}/${model.id}`;
   }
 
   async setThinkingOption(thinkingOptionId: string | null): Promise<void> {
     const thinkingLevel = normalizePiThinkingOption(thinkingOptionId) ?? DEFAULT_PI_THINKING_LEVEL;
     await this.runtimeSession.setThinkingLevel(thinkingLevel);
-    this.lastKnownThinkingOptionId = thinkingLevel;
-    this.config.thinkingOptionId = thinkingLevel;
-    this.state = {
-      ...this.state,
-      thinkingLevel,
-    };
+    await this.refreshState();
   }
 
   private emit(event: AgentStreamEvent): void {
@@ -2496,6 +2501,7 @@ export class PiRpcAgentSession implements AgentSession {
 
   private async refreshState(): Promise<void> {
     this.state = await this.runtimeSession.getState();
+    this.config.thinkingOptionId = this.state.thinkingLevel;
   }
 
   private async refreshAfterTurn(finalUsage: Promise<void>): Promise<void> {
@@ -2542,8 +2548,7 @@ export class PiRpcAgentClient implements AgentClient {
       runtimeSession = await this.runtime.startSession({
         cwd: config.cwd,
         model: config.model,
-        thinkingOptionId:
-          normalizePiThinkingOption(config.thinkingOptionId) ?? DEFAULT_PI_THINKING_LEVEL,
+        thinkingOptionId: normalizePiThinkingOption(config.thinkingOptionId) ?? undefined,
         noSession: config.internal === true,
         env: launchContext?.env,
         mcpConfigPath: mcpConfig?.path,

packages/server/src/server/agent/providers/pi/rpc-types.ts

@@ -74,6 +74,7 @@ export interface PiModel {
   id: string;
   name?: string;
   reasoning?: boolean;
+  thinkingLevelMap?: Record<string, string | null>;
   contextWindow?: number;
   maxTokens?: number;
   api?: string;

packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts

@@ -265,6 +265,7 @@ export class FakePiSession implements PiRuntimeSession {
     if (!this.setModelResult) {
       throw new Error("FakePi setModel requires setModelResult to be scripted");
     }
+    this.state = { ...this.state, model: this.setModelResult };
     return this.setModelResult;
   }