maximhar
122 lines
4.2 KiB
TypeScript
122 lines
4.2 KiB
TypeScript
import { afterEach, beforeEach, describe, expect, it } from "bun:test";
|
|
import * as path from "node:path";
|
|
import { Agent } from "@oh-my-pi/pi-agent-core";
|
|
import { getBundledModel, type ThinkingLevel } from "@oh-my-pi/pi-ai";
|
|
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
|
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
|
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
|
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
|
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
|
import { TempDir } from "@oh-my-pi/pi-utils";
|
|
|
|
describe("AgentSession role model thinking behavior", () => {
|
|
let tempDir: TempDir;
|
|
let session: AgentSession;
|
|
|
|
beforeEach(() => {
|
|
tempDir = TempDir.createSync("@pi-role-thinking-");
|
|
});
|
|
|
|
afterEach(async () => {
|
|
if (session) {
|
|
await session.dispose();
|
|
}
|
|
tempDir.removeSync();
|
|
});
|
|
|
|
function getAnthropicModelOrThrow(id: string) {
|
|
const model = getBundledModel("anthropic", id);
|
|
if (!model) throw new Error(`Expected anthropic model ${id} to exist`);
|
|
return model;
|
|
}
|
|
|
|
async function createSession(options: {
|
|
initialModelId: string;
|
|
initialThinkingLevel: ThinkingLevel;
|
|
modelRoles: Record<string, string>;
|
|
}) {
|
|
const model = getAnthropicModelOrThrow(options.initialModelId);
|
|
const agent = new Agent({
|
|
initialState: {
|
|
model,
|
|
systemPrompt: "Test",
|
|
tools: [],
|
|
messages: [],
|
|
thinkingLevel: options.initialThinkingLevel,
|
|
},
|
|
});
|
|
const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db"));
|
|
authStorage.setRuntimeApiKey("anthropic", "test-key");
|
|
const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir.path(), "models.yml"));
|
|
|
|
session = new AgentSession({
|
|
agent,
|
|
sessionManager: SessionManager.inMemory(),
|
|
settings: Settings.isolated({ modelRoles: options.modelRoles }),
|
|
modelRegistry,
|
|
});
|
|
}
|
|
|
|
it("re-applies explicit role thinking each time that role is selected", async () => {
|
|
const defaultModel = getAnthropicModelOrThrow("claude-sonnet-4-5");
|
|
const slowModel = getAnthropicModelOrThrow("claude-sonnet-4-6");
|
|
|
|
await createSession({
|
|
initialModelId: defaultModel.id,
|
|
initialThinkingLevel: "high",
|
|
modelRoles: {
|
|
default: `${defaultModel.provider}/${defaultModel.id}`,
|
|
slow: `${slowModel.provider}/${slowModel.id}:off`,
|
|
},
|
|
});
|
|
|
|
const firstSwitch = await session.cycleRoleModels(["default", "slow"]);
|
|
expect(firstSwitch?.role).toBe("slow");
|
|
expect(firstSwitch?.model.id).toBe(slowModel.id);
|
|
expect(firstSwitch?.thinkingLevel).toBe("off");
|
|
expect(session.thinkingLevel).toBe("off");
|
|
|
|
session.setThinkingLevel("high");
|
|
expect(session.thinkingLevel).toBe("high");
|
|
|
|
const secondSwitch = await session.cycleRoleModels(["default", "slow"]);
|
|
expect(secondSwitch?.role).toBe("default");
|
|
expect(secondSwitch?.model.id).toBe(defaultModel.id);
|
|
expect(session.thinkingLevel).toBe("high");
|
|
|
|
const thirdSwitch = await session.cycleRoleModels(["default", "slow"]);
|
|
expect(thirdSwitch?.role).toBe("slow");
|
|
expect(thirdSwitch?.model.id).toBe(slowModel.id);
|
|
expect(thirdSwitch?.thinkingLevel).toBe("off");
|
|
expect(session.thinkingLevel).toBe("off");
|
|
});
|
|
|
|
it("preserves current thinking when switching into default/no-suffix role", async () => {
|
|
const defaultModel = getAnthropicModelOrThrow("claude-sonnet-4-5");
|
|
const slowModel = getAnthropicModelOrThrow("claude-sonnet-4-6");
|
|
|
|
await createSession({
|
|
initialModelId: defaultModel.id,
|
|
initialThinkingLevel: "low",
|
|
modelRoles: {
|
|
default: `${defaultModel.provider}/${defaultModel.id}`,
|
|
slow: `${slowModel.provider}/${slowModel.id}:high`,
|
|
},
|
|
});
|
|
|
|
const toSlow = await session.cycleRoleModels(["default", "slow"]);
|
|
expect(toSlow?.role).toBe("slow");
|
|
expect(toSlow?.thinkingLevel).toBe("high");
|
|
expect(session.thinkingLevel).toBe("high");
|
|
|
|
session.setThinkingLevel("minimal");
|
|
expect(session.thinkingLevel).toBe("minimal");
|
|
|
|
const toDefault = await session.cycleRoleModels(["default", "slow"]);
|
|
expect(toDefault?.role).toBe("default");
|
|
expect(toDefault?.model.id).toBe(defaultModel.id);
|
|
expect(toDefault?.thinkingLevel).toBe("minimal");
|
|
expect(session.thinkingLevel).toBe("minimal");
|
|
});
|
|
});
|