diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index a5e2586df..22e0d8bdc 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -18,6 +18,7 @@ - Fixed the `/mcp` sign-in URL not being clickable when it wraps across lines, by emitting it as a terminal hyperlink with a `Cmd/Ctrl+click to open` line like `/login` ([#10186](https://github.com/earendil-works/pi/issues/10186)). - Fixed codemode `image()` accepting malformed base64 data or unsupported image types, which persisted an invalid image block that made every later provider request fail with HTTP 400 ([#10215](https://github.com/earendil-works/pi/issues/10215)). - Fixed codemode failing to start its script worker from the standalone Windows executable ([#10204](https://github.com/earendil-works/pi/issues/10204)). +- Fixed prompt submission slowing down with session length, because resolving the session's model selection looked up the model catalog once per assistant message ([#10198](https://github.com/earendil-works/pi/issues/10198)). ## [0.99.1] - 2026-09-29 diff --git a/packages/coding-agent/src/core/virtual-models.ts b/packages/coding-agent/src/core/virtual-models.ts index c6b5e6998..269ef6858 100644 --- a/packages/coding-agent/src/core/virtual-models.ts +++ b/packages/coding-agent/src/core/virtual-models.ts @@ -122,26 +122,37 @@ export function findLatestResponse(messages: readonly AgentMessage[]): Assistant * `model_change`, because responses name the physical models it routed to. Otherwise the latest * physical response wins, as in sessions without virtual models. A virtual model that is no longer * registered does not hold, so the selection falls back to the physical model that answered last. + * + * Only the last `model_change` can hold, so this looks up at most one model in the catalog. */ export function getBranchSelection( branch: readonly SessionEntry[], getModel: (provider: string, modelId: string) => Model | undefined, ): { provider: string; modelId: string } | undefined { - const isVirtual = (provider: string, modelId: string) => { - const model = getModel(provider, modelId); - return model !== undefined && isVirtualModel(model); - }; - let selection: { provider: string; modelId: string } | undefined; - for (const entry of branch) { + for (let i = branch.length - 1; i >= 0; i--) { + const entry = branch[i]; if (entry.type === "model_change") { - selection = { provider: entry.provider, modelId: entry.modelId }; - } else if (entry.type === "message" && entry.message.role === "assistant" && !isVirtualModel(entry.message)) { - if (!selection || !isVirtual(selection.provider, selection.modelId)) { - selection = { provider: entry.message.provider, modelId: entry.message.model }; - } + return { provider: entry.provider, modelId: entry.modelId }; + } + if (entry.type === "message" && entry.message.role === "assistant" && !isVirtualModel(entry.message)) { + const response = { provider: entry.message.provider, modelId: entry.message.model }; + const change = findLastModelChange(branch, i); + const model = change && getModel(change.provider, change.modelId); + return change && model && isVirtualModel(model) ? change : response; } } - return selection; + return undefined; +} + +function findLastModelChange( + branch: readonly SessionEntry[], + before: number, +): { provider: string; modelId: string } | undefined { + for (let i = before - 1; i >= 0; i--) { + const entry = branch[i]; + if (entry.type === "model_change") return { provider: entry.provider, modelId: entry.modelId }; + } + return undefined; } /** Latest router state a session branch stores for a virtual model. */ diff --git a/packages/coding-agent/test/virtual-models.test.ts b/packages/coding-agent/test/virtual-models.test.ts index 7d1d0b651..3650c2ccd 100644 --- a/packages/coding-agent/test/virtual-models.test.ts +++ b/packages/coding-agent/test/virtual-models.test.ts @@ -14,7 +14,7 @@ import { AuthStorage } from "../src/core/auth-storage.ts"; import { ModelRuntime } from "../src/core/model-runtime.ts"; import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; -import type { ModelRouteRequest, VirtualModelDefinition } from "../src/core/virtual-models.ts"; +import { getBranchSelection, type ModelRouteRequest, type VirtualModelDefinition } from "../src/core/virtual-models.ts"; import { createTestResourceLoader } from "./utilities.ts"; async function createRuntime(requests: ModelRouteRequest[] = []) { @@ -51,6 +51,68 @@ function assistantFrom(model: Model, text: string): AssistantMessage { return { ...fauxAssistantMessage(text), api: model.api, provider: model.provider, model: model.id }; } +describe("getBranchSelection", () => { + function select(build: (sessionManager: SessionManager) => void, runtime: ModelRuntime) { + const sessionManager = SessionManager.inMemory(); + build(sessionManager); + const lookups: string[] = []; + const selection = getBranchSelection(sessionManager.getBranch(), (provider, modelId) => { + lookups.push(`${provider}/${modelId}`); + return runtime.getModel(provider, modelId); + }); + return { selection, lookups }; + } + + // #10198: the selection must not cost one catalog lookup per assistant message. + it("looks up only the last model_change", async () => { + const { runtime, virtual } = await createRuntime(); + const small = runtime.getModel("faux", "small")!; + const large = runtime.getModel("faux", "large")!; + + const physical = select((sessionManager) => { + sessionManager.appendModelChange(small.provider, small.id); + for (let i = 0; i < 100; i++) sessionManager.appendMessage(assistantFrom(large, "ok")); + }, runtime); + expect(physical.selection).toEqual({ provider: "faux", modelId: "large" }); + expect(physical.lookups).toEqual(["faux/small"]); + + const routed = select((sessionManager) => { + sessionManager.appendModelChange(small.provider, small.id); + sessionManager.appendMessage(assistantFrom(small, "ok")); + sessionManager.appendModelChange(virtual.provider, virtual.id); + for (let i = 0; i < 100; i++) sessionManager.appendMessage(assistantFrom(large, "ok")); + }, runtime); + expect(routed.selection).toEqual({ provider: "router", modelId: "auto" }); + expect(routed.lookups).toEqual(["router/auto"]); + }); + + it("uses the last model_change without responses after it", async () => { + const { runtime, virtual } = await createRuntime(); + const small = runtime.getModel("faux", "small")!; + const { selection, lookups } = select((sessionManager) => { + sessionManager.appendModelChange(virtual.provider, virtual.id); + sessionManager.appendMessage(assistantFrom(small, "ok")); + sessionManager.appendModelChange(small.provider, small.id); + }, runtime); + expect(selection).toEqual({ provider: "faux", modelId: "small" }); + expect(lookups).toEqual([]); + }); + + it("uses the latest physical response without a model_change", async () => { + const { runtime, virtual } = await createRuntime(); + const small = runtime.getModel("faux", "small")!; + const large = runtime.getModel("faux", "large")!; + const { selection, lookups } = select((sessionManager) => { + sessionManager.appendMessage(assistantFrom(small, "ok")); + sessionManager.appendMessage(assistantFrom(large, "ok")); + // Failed routing leaves the virtual model on its message. + sessionManager.appendMessage({ ...assistantFrom(virtual, ""), stopReason: "error" }); + }, runtime); + expect(selection).toEqual({ provider: "faux", modelId: "large" }); + expect(lookups).toEqual([]); + }); +}); + describe("ModelRuntime virtual models", () => { it("lists a virtual model and routes it to a physical model with a clamped thinking level", async () => { const requests: ModelRouteRequest[] = [];