trending/pi
1
0
Fork 0
mirror of https://github.com/earendil-works/pi.git synced 2026-10-02 00:35:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(coding-agent): resolve branch model selection with one catalog lookup

Browse Source 
getBranchSelection looked up the model catalog once per assistant message,
which made prompt submission scale with session length. Walk the branch
backward instead: only the last model_change can hold against later
responses, so at most that model needs a lookup.

closes #10198
This commit is contained in:
Armin Ronacher
2026-09-30 14:20:03 +02:00
parent 37c9d0d200
commit a0660b174e
3 changed files with 87 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/CHANGELOG.md
+1
View File
@@ -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
packages/coding-agent/src/core/virtual-models.ts
+23 -12
View File
@@ -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<Api> | 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. */
packages/coding-agent/test/virtual-models.test.ts
+63 -1
View File
@@ -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<string>, 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[] = [];