mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
This adds experimental support for virtual models. A virtual model is a catalog entry that an extension registers with pi.registerVirtualModel(). It does not talk to a provider itself; for every request it picks a physical model and thinking level. The motivation is to make routing policies pluggable: plan on a strong model and implement on a cheap one, pick a model with a classifier, or move to a larger window when the context grows. Today this needs model switching by hand or hacks in prepareRequest. The selection stays virtual: model_change records it, resume restores it, and each response records the physical model and thinking level that produced it. Router state can be stored on the session branch, so it follows /tree and forks and survives compaction. Context limits, compaction and image sizing use the physical model a request was routed to. examples/extensions/jev-router.ts shows a full router that uses the Jev classifier to pick Sol or Terra for planning and hands off to Luna after the first edit. docs/virtual-models.md describes the API.
269 lines
6.8 KiB
TypeScript
269 lines
6.8 KiB
TypeScript
import { visibleWidth } from "@earendil-works/pi-tui";
|
|
import { beforeAll, describe, expect, it } from "vitest";
|
|
import type { AgentSession } from "../src/core/agent-session.ts";
|
|
import type { ReadonlyFooterDataProvider } from "../src/core/footer-data-provider.ts";
|
|
import { FooterComponent, formatCwdForFooter } from "../src/modes/interactive/components/footer.ts";
|
|
import { initTheme } from "../src/modes/interactive/theme/theme.ts";
|
|
import { stripAnsi } from "../src/utils/ansi.ts";
|
|
|
|
type AssistantUsage = {
|
|
input: number;
|
|
output: number;
|
|
cacheRead: number;
|
|
cacheWrite: number;
|
|
cost: { total: number };
|
|
};
|
|
|
|
function createSession(options: {
|
|
sessionName: string;
|
|
modelId?: string;
|
|
provider?: string;
|
|
reasoning?: boolean;
|
|
thinkingLevel?: string;
|
|
usage?: AssistantUsage;
|
|
branchUsage?: AssistantUsage;
|
|
compactionUsage?: AssistantUsage;
|
|
toolUsage?: AssistantUsage;
|
|
usingSubscription?: boolean;
|
|
routedModel?: { model: { id: string }; thinkingLevel?: string };
|
|
}): AgentSession {
|
|
const usage = options.usage;
|
|
const entries: Array<Record<string, unknown>> = [];
|
|
|
|
if (usage !== undefined) {
|
|
entries.push({
|
|
type: "message",
|
|
message: {
|
|
role: "assistant",
|
|
usage,
|
|
},
|
|
});
|
|
}
|
|
|
|
if (options.branchUsage !== undefined) {
|
|
entries.push({
|
|
type: "branch_summary",
|
|
usage: options.branchUsage,
|
|
});
|
|
}
|
|
|
|
if (options.compactionUsage !== undefined) {
|
|
entries.push({
|
|
type: "compaction",
|
|
usage: options.compactionUsage,
|
|
});
|
|
}
|
|
|
|
if (options.toolUsage !== undefined) {
|
|
entries.push({
|
|
type: "message",
|
|
message: {
|
|
role: "toolResult",
|
|
usage: options.toolUsage,
|
|
},
|
|
});
|
|
}
|
|
|
|
const session = {
|
|
state: {
|
|
model: {
|
|
id: options.modelId ?? "test-model",
|
|
provider: options.provider ?? "test",
|
|
contextWindow: 200_000,
|
|
reasoning: options.reasoning ?? false,
|
|
},
|
|
thinkingLevel: options.thinkingLevel ?? "off",
|
|
},
|
|
sessionManager: {
|
|
getEntries: () => entries,
|
|
getSessionName: () => options.sessionName,
|
|
getCwd: () => "/tmp/project",
|
|
},
|
|
getContextUsage: () => ({ contextWindow: 200_000, percent: 12.3 }),
|
|
routedModel: options.routedModel,
|
|
modelRuntime: {
|
|
isUsingSubscription: () => options.usingSubscription ?? false,
|
|
},
|
|
};
|
|
|
|
return session as unknown as AgentSession;
|
|
}
|
|
|
|
function createFooterData(providerCount: number): ReadonlyFooterDataProvider {
|
|
const provider = {
|
|
getGitBranch: () => "main",
|
|
getExtensionStatuses: () => new Map<string, string>(),
|
|
getAvailableProviderCount: () => providerCount,
|
|
onBranchChange: (callback: () => void) => {
|
|
void callback;
|
|
return () => {};
|
|
},
|
|
};
|
|
|
|
return provider;
|
|
}
|
|
|
|
describe("formatCwdForFooter", () => {
|
|
it("does not abbreviate sibling paths that share the home prefix", () => {
|
|
expect(formatCwdForFooter("/home/user2", "/home/user")).toBe("/home/user2");
|
|
});
|
|
|
|
it("abbreviates the home directory and descendants", () => {
|
|
expect(formatCwdForFooter("/home/user", "/home/user")).toBe("~");
|
|
expect(formatCwdForFooter("/home/user/project", "/home/user")).toBe("~/project");
|
|
});
|
|
});
|
|
|
|
describe("FooterComponent width handling", () => {
|
|
beforeAll(() => {
|
|
initTheme(undefined, false);
|
|
});
|
|
|
|
it("keeps all lines within width for wide session names", () => {
|
|
const width = 93;
|
|
const session = createSession({ sessionName: "한글".repeat(30) });
|
|
const footer = new FooterComponent(session, createFooterData(1));
|
|
|
|
const lines = footer.render(width);
|
|
for (const line of lines) {
|
|
expect(visibleWidth(line)).toBeLessThanOrEqual(width);
|
|
}
|
|
});
|
|
|
|
it("keeps stats line within width for wide model and provider names", () => {
|
|
const width = 60;
|
|
const session = createSession({
|
|
sessionName: "",
|
|
modelId: "模".repeat(30),
|
|
provider: "공급자",
|
|
reasoning: true,
|
|
thinkingLevel: "high",
|
|
usage: {
|
|
input: 12_345,
|
|
output: 6_789,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
cost: { total: 1.234 },
|
|
},
|
|
});
|
|
const footer = new FooterComponent(session, createFooterData(2));
|
|
|
|
const lines = footer.render(width);
|
|
for (const line of lines) {
|
|
expect(visibleWidth(line)).toBeLessThanOrEqual(width);
|
|
}
|
|
});
|
|
|
|
it("shows the physical model a virtual model routed to", () => {
|
|
const session = createSession({
|
|
sessionName: "",
|
|
modelId: "auto",
|
|
reasoning: true,
|
|
thinkingLevel: "high",
|
|
routedModel: { model: { id: "gpt-5.6-luna" }, thinkingLevel: "medium" },
|
|
});
|
|
const footer = new FooterComponent(session, createFooterData(1));
|
|
|
|
const statsLine = stripAnsi(footer.render(120)[1]);
|
|
|
|
expect(statsLine).toContain("auto \u2022 high \u2192 gpt-5.6-luna \u2022 medium");
|
|
});
|
|
|
|
it("includes summary and tool result usage in the total cost", () => {
|
|
const session = createSession({
|
|
sessionName: "",
|
|
usage: {
|
|
input: 100,
|
|
output: 10,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
cost: { total: 0.5 },
|
|
},
|
|
branchUsage: {
|
|
input: 20,
|
|
output: 5,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
cost: { total: 0.25 },
|
|
},
|
|
compactionUsage: {
|
|
input: 5,
|
|
output: 2,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
cost: { total: 0.125 },
|
|
},
|
|
toolUsage: {
|
|
input: 15,
|
|
output: 3,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
cost: { total: 0.375 },
|
|
},
|
|
});
|
|
const footer = new FooterComponent(session, createFooterData(1));
|
|
|
|
const statsLine = stripAnsi(footer.render(120)[1]);
|
|
expect(statsLine).toContain("$1.250");
|
|
});
|
|
|
|
it("shows the latest cache hit rate when cache usage is present", () => {
|
|
const session = createSession({
|
|
sessionName: "",
|
|
usage: {
|
|
input: 100,
|
|
output: 10,
|
|
cacheRead: 50,
|
|
cacheWrite: 50,
|
|
cost: { total: 0.001 },
|
|
},
|
|
});
|
|
const footer = new FooterComponent(session, createFooterData(1));
|
|
|
|
const statsLine = stripAnsi(footer.render(120)[1]);
|
|
expect(statsLine).toContain("CH25.0%");
|
|
});
|
|
|
|
it("marks Kimi Coding costs as subscription estimates", () => {
|
|
const session = createSession({
|
|
sessionName: "",
|
|
provider: "kimi-coding",
|
|
usage: {
|
|
input: 100,
|
|
output: 10,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
cost: { total: 1.234 },
|
|
},
|
|
});
|
|
const footer = new FooterComponent(session, createFooterData(1));
|
|
|
|
expect(stripAnsi(footer.render(120)[1])).toContain("$1.234 (sub)");
|
|
});
|
|
|
|
it("marks explicitly identified subscription auth", () => {
|
|
const session = createSession({ sessionName: "", provider: "anthropic", usingSubscription: true });
|
|
const footer = new FooterComponent(session, createFooterData(1));
|
|
|
|
expect(stripAnsi(footer.render(120)[1])).toContain("$0.000 (sub)");
|
|
});
|
|
|
|
it("does not mark generic OAuth sign-in as a subscription", () => {
|
|
const session = createSession({
|
|
sessionName: "",
|
|
provider: "openrouter",
|
|
usage: {
|
|
input: 100,
|
|
output: 10,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
cost: { total: 1.234 },
|
|
},
|
|
});
|
|
const footer = new FooterComponent(session, createFooterData(1));
|
|
const stats = stripAnsi(footer.render(120)[1]);
|
|
|
|
expect(stats).toContain("$1.234");
|
|
expect(stats).not.toContain("(sub)");
|
|
});
|
|
});
|