From ee7949ae25441d67e0780b4e99bf1d4e7c57897c Mon Sep 17 00:00:00 2001 From: Matt Cowger Date: Mon, 21 Sep 2026 02:53:52 -0700 Subject: [PATCH] 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 Co-authored-by: trim21 --- .../command-center-agent-controls.spec.ts | 23 +- .../app/e2e/support/helpers/mock-agent.ts | 2 + .../app/src/composer/agent-controls/index.tsx | 2 +- .../src/server/agent/agent-manager.test.ts | 87 +++++ .../server/src/server/agent/agent-manager.ts | 10 +- .../agent/providers/mock-load-test-agent.ts | 12 + .../server/agent/providers/pi/agent.test.ts | 296 +++++++++++++++++- .../src/server/agent/providers/pi/agent.ts | 101 +++--- .../server/agent/providers/pi/rpc-types.ts | 1 + .../agent/providers/pi/test-utils/fake-pi.ts | 1 + 10 files changed, 477 insertions(+), 58 deletions(-) diff --git a/packages/app/e2e/browser/command-center-agent-controls.spec.ts b/packages/app/e2e/browser/command-center-agent-controls.spec.ts index 0748250302..658e0e1123 100644 --- a/packages/app/e2e/browser/command-center-agent-controls.spec.ts +++ b/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 diff --git a/packages/app/e2e/support/helpers/mock-agent.ts b/packages/app/e2e/support/helpers/mock-agent.ts index d1447700a3..f6176b0f07 100644 --- a/packages/app/e2e/support/helpers/mock-agent.ts +++ b/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; } @@ -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, }); diff --git a/packages/app/src/composer/agent-controls/index.tsx b/packages/app/src/composer/agent-controls/index.tsx index 32941e46e7..f873b42148 100644 --- a/packages/app/src/composer/agent-controls/index.tsx +++ b/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} diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 2abdf65040..c107b3da0b 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/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 {} + + override async getRuntimeInfo() { + return { + ...(await super.getRuntimeInfo()), + thinkingOptionId: "high", + }; + } + } + class ClampingClient extends TestAgentClient { + override async createSession(config: AgentSessionConfig): Promise { + 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 {} + + override async getRuntimeInfo(): Promise { + throw new Error("state unavailable"); + } + } + class UnreadableStateClient extends TestAgentClient { + override async createSession(config: AgentSessionConfig): Promise { + 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; diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 9d3a51feb2..31ea9af5fd 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/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); diff --git a/packages/server/src/server/agent/providers/mock-load-test-agent.ts b/packages/server/src/server/agent/providers/mock-load-test-agent.ts index 315bc728c1..2ebcc0b23e 100644 --- a/packages/server/src/server/agent/providers/mock-load-test-agent.ts +++ b/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", diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index 234fe59cec..52d7c65128 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/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([ + { + 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], ]); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index aaf517422d..aa6117f550 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/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 { + 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(); private readonly pendingSteerSubmissions: PiPendingSteerSubmission[] = []; - private lastKnownThinkingOptionId: string | null; currentLeafOverrideId: string | null | undefined; private readonly capturedUserEntries: PiCapturedEntry[] = []; private readonly capturedUserEntriesById = new Map(); @@ -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 { 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 { this.state = await this.runtimeSession.getState(); + this.config.thinkingOptionId = this.state.thinkingLevel; } private async refreshAfterTurn(finalUsage: Promise): Promise { @@ -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, diff --git a/packages/server/src/server/agent/providers/pi/rpc-types.ts b/packages/server/src/server/agent/providers/pi/rpc-types.ts index f84e0ff69e..a0948ade15 100644 --- a/packages/server/src/server/agent/providers/pi/rpc-types.ts +++ b/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; contextWindow?: number; maxTokens?: number; api?: string; diff --git a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts index b1d369eb85..ce20b4d8ee 100644 --- a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts +++ b/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; }