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 c6384e1fa2..927a9defd2 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -1943,6 +1943,48 @@ describe("PiRpcAgentSession", () => { expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe("off"); }); + test.each(["high", "off", "max"] as const)( + "publishes an extension's thinking change to %s before refreshing state", + async (level) => { + const { pi, session } = await createSession(); + onTestFinished(() => session.close()); + const changes: AgentStreamEvent[] = []; + session.subscribe((event) => { + if (event.type === "thinking_option_changed") { + expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe(level); + changes.push(event); + } + }); + const runtime = pi.latestSession(); + runtime.state = { ...runtime.state, thinkingLevel: level }; + runtime.emit({ type: "thinking_level_changed", level }); + + expect(changes).toEqual([ + { type: "thinking_option_changed", provider: "pi", thinkingOptionId: level }, + ]); + expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe(level); + await expect(session.getRuntimeInfo()).resolves.toMatchObject({ thinkingOptionId: level }); + }, + ); + + test.each([undefined, null, "invalid", 123])( + "ignores a thinking change with invalid level %s", + async (level) => { + const { pi, session } = await createSession(); + onTestFinished(() => session.close()); + const initialThinking = session.describePersistence()?.metadata?.thinkingOptionId; + const changes: AgentStreamEvent[] = []; + session.subscribe((event) => { + if (event.type === "thinking_option_changed") changes.push(event); + }); + + pi.latestSession().emit({ type: "thinking_level_changed", level }); + + expect(changes).toEqual([]); + expect(session.describePersistence()?.metadata?.thinkingOptionId).toBe(initialThinking); + }, + ); + test("surfaces state refresh failures after a Pi thinking update", async () => { const { pi, session } = await createSession(); onTestFinished(() => session.close()); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 38a1f9815f..146b56b36f 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -1005,6 +1005,8 @@ function isProcessExitEvent( function isPiAgentSessionEvent(event: PiRuntimeEvent): event is PiAgentSessionEvent { switch (event.type) { + case "thinking_level_changed": + return typeof event.level === "string" && normalizePiThinkingOption(event.level) !== null; case "agent_start": case "turn_start": case "message_start": @@ -2136,6 +2138,15 @@ export class PiRpcAgentSession implements AgentSession { } switch (event.type) { + case "thinking_level_changed": + this.state = { ...this.state, thinkingLevel: event.level }; + this.config.thinkingOptionId = event.level; + this.emit({ + type: "thinking_option_changed", + provider: this.provider, + thinkingOptionId: event.level, + }); + return; case "agent_start": this.activeTurnStarted = true; this.clearNoTurnBuffers(); 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 cfa09daeb6..3e488c028c 100644 --- a/packages/server/src/server/agent/providers/pi/rpc-types.ts +++ b/packages/server/src/server/agent/providers/pi/rpc-types.ts @@ -160,6 +160,7 @@ export type PiAssistantMessageEvent = | { type: "start" | "text_start" | "text_end" | "thinking_start" | "thinking_end" | "done" }; export type PiAgentSessionEvent = + | { type: "thinking_level_changed"; level: PiThinkingLevel } | { type: "agent_start" } | { type: "turn_start" } | { type: "message_start"; message: PiAgentMessage }