oh-my-pi/packages/coding-agent/test/model-selector-role-badge-thinking.test.ts

import { beforeAll, describe, expect, test, vi } from "bun:test";
import { stripVTControlCharacters } from "node:util";
import type { Model } from "@oh-my-pi/pi-ai";
import { buildModel } from "@oh-my-pi/pi-catalog/build";
import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
import type { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
import { ModelSelectorComponent } from "@oh-my-pi/pi-coding-agent/modes/components/model-selector";
import { getThemeByName, setThemeInstance } from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
import { AUTO_THINKING, type ConfiguredThinkingLevel } from "@oh-my-pi/pi-coding-agent/thinking";
import type { TUI } from "@oh-my-pi/pi-tui";

function normalizeRenderedText(text: string): string {
	return stripVTControlCharacters(text).replace(/\s+/g, " ").trim();
}

const DEFAULT_RETRY_FALLBACK_ACTION_LABEL = "Set as DEFAULT retry fallback";
const DEFAULT_RETRY_FALLBACK_ACTION = "retryFallback";

type ModelSelectorAction = "modelRole" | typeof DEFAULT_RETRY_FALLBACK_ACTION;
type TestRoleSelectArgs = [
	model: Model,
	role: string | null,
	thinkingLevel?: ConfiguredThinkingLevel,
	selector?: string,
	action?: ModelSelectorAction,
];
type TestRoleSelectCallback = (...args: TestRoleSelectArgs) => void;

function isSelectedMenuLine(line: string): boolean {
	const trimmed = line.trimStart();
	return trimmed.startsWith("❯") || trimmed.startsWith("▸") || trimmed.startsWith(">") || trimmed.startsWith("\uf054");
}

function selectMenuAction(selector: ModelSelectorComponent, label: string): void {
	for (let attempt = 0; attempt < 20; attempt++) {
		const selectedTarget = stripVTControlCharacters(selector.render(220).join("\n"))
			.split("\n")
			.find(line => line.includes(label) && isSelectedMenuLine(line));
		if (selectedTarget) return;
		selector.handleInput("\x1b[B");
	}
	throw new Error(`Menu action not selectable: ${label}`);
}

function createSelector(model: Model, settings: Settings): ModelSelectorComponent {
	const modelRegistry = {
		getAll: () => [model],
		getDiscoverableProviders: () => [],
	} as unknown as ModelRegistry;
	const ui = {
		requestRender: vi.fn(),
	} as unknown as TUI;

	return new ModelSelectorComponent(
		ui,
		model,
		settings,
		modelRegistry,
		[{ model, thinkingLevel: "off" }],
		() => {},
		() => {},
	);
}

function createOllamaCloudModel(id: string): Model {
	return buildModel({
		id,
		name: "DeepSeek V4 Pro",
		api: "ollama-chat",
		provider: "ollama-cloud",
		baseUrl: "https://ollama.com",
		reasoning: true,
		input: ["text"],
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
		contextWindow: 1_000_000,
		maxTokens: 8192,
	});
}
function createContextTestModel(id: string, contextWindow: number): Model {
	return buildModel({
		id,
		name: id,
		api: "ollama-chat",
		baseUrl: "https://example.com",
		reasoning: false,
		provider: "test",
		input: ["text"],
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
		contextWindow,
		maxTokens: 1024,
	});
}

function createScopedSelector(
	models: Model[],
	settings: Settings,
	onSelect: TestRoleSelectCallback,
	options?: { temporaryOnly?: boolean; currentContextTokens?: number },
): ModelSelectorComponent {
	const modelRegistry = {
		getAll: () => models,
		getDiscoverableProviders: () => [],
	} as unknown as ModelRegistry;
	const ui = {
		requestRender: vi.fn(),
	} as unknown as TUI;
	return new ModelSelectorComponent(
		ui,
		undefined,
		settings,
		modelRegistry,
		models.map(model => ({ model })),
		(
			model: Model,
			role: string | null,
			thinkingLevel?: ConfiguredThinkingLevel,
			selector?: string,
			action?: ModelSelectorAction,
		) => onSelect(model, role, thinkingLevel, selector, action),
		() => {},
		options,
	);
}
let testTheme = await getThemeByName("dark");

function installTestTheme(): void {
	if (!testTheme) {
		throw new Error("Failed to load dark theme for ModelSelector tests");
	}
	setThemeInstance(testTheme);
}

describe("ModelSelector role badge thinking display", () => {
	beforeAll(async () => {
		testTheme = await getThemeByName("dark");
		if (!testTheme) {
			throw new Error("Failed to load dark theme for ModelSelector tests");
		}
	});

	test("shows custom roles from cycleOrder/modelRoles and honors built-in metadata overrides", async () => {
		installTestTheme();
		const model = getBundledModel("anthropic", "claude-sonnet-4-5");
		if (!model) throw new Error("Expected bundled model anthropic/claude-sonnet-4-5");

		const settings = Settings.isolated({
			cycleOrder: ["smol", "custom-fast", "default"],
			modelRoles: {
				default: `${model.provider}/${model.id}`,
				"custom-fast": `${model.provider}/${model.id}:low`,
				smol: `${model.provider}/${model.id}`,
			},
			modelTags: {
				smol: { name: "Quick", color: "error" },
			},
		});

		const selector = createSelector(model, settings);
		await Bun.sleep(0);
		installTestTheme();

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("custom-fast (low)");
		expect(rendered).toContain("SMOL (inherit)");

		selector.handleInput("\n");
		installTestTheme();
		const menuRendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(menuRendered).toContain("Set as custom-fast");
		expect(menuRendered).toContain("Set as SMOL (Quick)");
	});

	test("renders the xhigh-ceiling ladder without a speculative max tier (GPT-5.5)", async () => {
		installTestTheme();
		const model = getBundledModel("openai", "gpt-5.5");
		if (!model) throw new Error("Expected bundled model openai/gpt-5.5");

		const selector = createSelector(model, Settings.isolated({}));
		await Bun.sleep(0);
		installTestTheme();

		selector.handleInput("\n");
		selector.handleInput("\n");

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("Thinking for: Default (gpt-5.5)");
		expect(rendered).toContain("low medium high xhigh");
		// gpt-5.5's wire has no max tier; the selector must not invent one.
		expect(rendered).not.toContain("max");
	});

	test("renders max as a real final tier on max-capable models (GPT-5.6)", async () => {
		installTestTheme();
		const model = getBundledModel("openai", "gpt-5.6");
		if (!model) throw new Error("Expected bundled model openai/gpt-5.6");

		const selector = createSelector(model, Settings.isolated({}));
		await Bun.sleep(0);
		installTestTheme();

		selector.handleInput("\n");
		selector.handleInput("\n");

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("Thinking for: Default (gpt-5.6)");
		expect(rendered).toContain("low medium high xhigh max");
	});

	test("reloads DEFAULT(auto) from defaultThinkingLevel", async () => {
		installTestTheme();
		const model = getBundledModel("openai", "gpt-5.5");
		if (!model) throw new Error("Expected bundled model openai/gpt-5.5");

		const settings = Settings.isolated({
			defaultThinkingLevel: AUTO_THINKING,
			modelRoles: {
				default: `${model.provider}/${model.id}`,
			},
		});

		const selector = createSelector(model, settings);
		await Bun.sleep(0);
		installTestTheme();

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("DEFAULT (auto)");
	});

	test("renders DEFAULT (auto) when modelRoles.default carries an explicit :auto suffix", async () => {
		installTestTheme();
		const model = getBundledModel("openai", "gpt-5.5");
		if (!model) throw new Error("Expected bundled model openai/gpt-5.5");

		const settings = Settings.isolated({
			modelRoles: {
				default: `${model.provider}/${model.id}:auto`,
			},
		});

		const selector = createSelector(model, settings);
		await Bun.sleep(0);
		installTestTheme();

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("DEFAULT (auto)");
		expect(rendered).not.toContain("DEFAULT (inherit)");
	});

	test("renders SMOL (auto) when modelRoles.smol carries an explicit :auto suffix", async () => {
		installTestTheme();
		const model = getBundledModel("openai", "gpt-5.5");
		if (!model) throw new Error("Expected bundled model openai/gpt-5.5");

		const settings = Settings.isolated({
			modelRoles: {
				default: `${model.provider}/${model.id}`,
				smol: `${model.provider}/${model.id}:auto`,
			},
		});

		const selector = createSelector(model, settings);
		await Bun.sleep(0);
		installTestTheme();

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("SMOL (auto)");
		expect(rendered).not.toContain("SMOL (inherit)");
	});

	test("shows compact auto badges for unconfigured role defaults", async () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const haiku = createContextTestModel("claude-haiku-4.5", 128_000);
		const codex = createContextTestModel("gpt-5.1-codex", 128_000);

		const selector = createScopedSelector([codex, haiku], settings, () => {});
		await Bun.sleep(0);
		installTestTheme();

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("claude-haiku-4.5");
		expect(rendered).toContain("gpt-5.1-codex");
		expect(rendered).toContain("[SMOL auto]");
		expect(rendered).toContain("[SLOW auto]");
	});

	test("dims and disables models below the current context size in temporary mode", async () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const small = createContextTestModel("a-small", 4096);
		const large = createContextTestModel("b-large", 128_000);
		const selected: string[] = [];
		const selector = createScopedSelector([small, large], settings, model => selected.push(model.id), {
			temporaryOnly: true,
			currentContextTokens: 6000,
		});
		await Bun.sleep(0);
		installTestTheme();

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("a-small");
		expect(rendered).toContain("context>4.1k");

		selector.handleInput("\n");
		expect(selected).toEqual(["b-large"]);
	});

	test("labels temporary picker as session-only and points to role assignment", async () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const model = createContextTestModel("session-model", 128_000);
		const selector = createScopedSelector([model], settings, () => {}, { temporaryOnly: true });
		await Bun.sleep(0);
		installTestTheme();

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("Temporary model selection is session-only");
		expect(rendered).toContain("Alt+M or /model");
		expect(rendered).toContain("default/smol/plan/task/slow/custom roles");
	});

	test("opens over-context default role actions for global configuration", async () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const small = createContextTestModel("only-small", 4096);
		const onSelect = vi.fn();
		const selector = createScopedSelector([small], settings, onSelect, {
			currentContextTokens: 6000,
		});
		await Bun.sleep(0);
		installTestTheme();

		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("only-small");
		expect(rendered).not.toContain("current context 6k > 4.1k limit");

		selector.handleInput("\n");
		const afterOpen = normalizeRenderedText(selector.render(220).join("\n"));
		expect(afterOpen).toContain("Action for: only-small");
		expect(afterOpen).toContain("Set as DEFAULT (Default)");
		expect(afterOpen).not.toContain("context>4.1k");

		selector.handleInput("\n");
		const afterRoleEnter = normalizeRenderedText(selector.render(220).join("\n"));
		expect(afterRoleEnter).toContain("Thinking for: Default (only-small)");
		expect(onSelect).not.toHaveBeenCalled();

		selector.handleInput("\n");
		expect(onSelect.mock.calls[0]?.[0]).toBe(small);
		expect(onSelect.mock.calls[0]?.[1]).toBe("default");
		expect(onSelect.mock.calls[0]?.[3]).toBe("test/only-small");
	});

	test("assigns selected model as default retry fallback without opening thinking options", () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const fallback = createContextTestModel("retry-fallback-model", 128_000);
		const onSelect = vi.fn();
		const selector = createScopedSelector([fallback], settings, onSelect);
		installTestTheme();

		selector.handleInput("\n");
		const menuRendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(menuRendered).toContain("Action for: retry-fallback-model");
		expect(menuRendered).toContain(DEFAULT_RETRY_FALLBACK_ACTION_LABEL);

		selectMenuAction(selector, DEFAULT_RETRY_FALLBACK_ACTION_LABEL);
		selector.handleInput("\n");

		const afterEnter = normalizeRenderedText(selector.render(220).join("\n"));
		expect(afterEnter).not.toContain("Thinking for:");
		expect(onSelect).toHaveBeenCalledTimes(1);
		const call = onSelect.mock.calls[0];
		expect(call?.[0]).toBe(fallback);
		expect(call?.[1]).toBe("default");
		expect(call?.[3]).toBe("test/retry-fallback-model");
		expect(call?.[4]).toBe(DEFAULT_RETRY_FALLBACK_ACTION);
	});

	test("uses cached models for Enter while offline refresh is still pending", () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const cachedModel = createContextTestModel("cached-fast", 128_000);
		const refreshGate = Promise.withResolvers<void>();
		const onSelect = vi.fn();
		const modelRegistry = {
			getAll: () => [cachedModel],
			refresh: vi.fn(() => refreshGate.promise),
			refreshProvider: vi.fn(async () => {}),
			getError: () => undefined,
			getAvailable: () => [cachedModel],
			getDiscoverableProviders: () => [],
		} as unknown as ModelRegistry;
		const ui = {
			requestRender: vi.fn(),
		} as unknown as TUI;

		const selector = new ModelSelectorComponent(
			ui,
			undefined,
			settings,
			modelRegistry,
			[],
			model => onSelect(model.id),
			() => {},
			{ temporaryOnly: true },
		);

		selector.handleInput("\n");
		expect(onSelect).toHaveBeenCalledWith("cached-fast");
		expect(modelRegistry.refresh).toHaveBeenCalledTimes(1);
		refreshGate.resolve();
	});

	test("keeps the highlighted model when a background refresh reorders the list", async () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const modelBb = createContextTestModel("bb-model", 128_000);
		const modelCc = createContextTestModel("cc-model", 128_000);
		const modelAa = createContextTestModel("aa-model", 128_000);
		let availableModels: Model[] = [modelBb, modelCc];
		const refreshGate = Promise.withResolvers<void>();
		const onSelect = vi.fn();
		const modelRegistry = {
			getAll: () => availableModels,
			refresh: vi.fn(() => refreshGate.promise),
			refreshProvider: vi.fn(async () => {}),
			getError: () => undefined,
			getAvailable: () => availableModels,
			getDiscoverableProviders: () => [],
		} as unknown as ModelRegistry;
		const ui = {
			requestRender: vi.fn(),
		} as unknown as TUI;

		const selector = new ModelSelectorComponent(
			ui,
			undefined,
			settings,
			modelRegistry,
			[],
			model => onSelect(model.id),
			() => {},
			{ temporaryOnly: true },
		);

		// Highlight the second entry, then let the pending refresh land a model
		// that sorts ahead of it and shifts every index.
		selector.handleInput("\x1b[B");
		availableModels = [modelAa, modelBb, modelCc];
		refreshGate.resolve();
		await Bun.sleep(0);

		selector.handleInput("\n");
		expect(onSelect).toHaveBeenCalledWith("cc-model");
	});

	test("refreshes Ollama Cloud using provider id instead of tab label", async () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const discoveredModel = createOllamaCloudModel("deepseek-v4-pro");
		let availableModels: Model[] = [];
		const refreshProvider = vi.fn(async (providerId: string) => {
			if (providerId === "ollama-cloud") {
				availableModels = [discoveredModel];
			}
		});
		const modelRegistry = {
			getAll: () => availableModels,
			refresh: vi.fn(async () => {}),
			refreshProvider,
			getError: () => undefined,
			getAvailable: () => availableModels,
			getDiscoverableProviders: () => ["ollama-cloud"],
			getProviderDiscoveryState: () => ({
				provider: "ollama-cloud",
				status: "idle",
				optional: false,
				stale: false,
				models: [],
			}),
		} as unknown as ModelRegistry;
		const ui = {
			requestRender: vi.fn(),
		} as unknown as TUI;

		const selector = new ModelSelectorComponent(
			ui,
			undefined,
			settings,
			modelRegistry,
			[],
			() => {},
			() => {},
		);
		await Bun.sleep(0);
		installTestTheme();

		const initialRendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(initialRendered).toContain("OLLAMA CLOUD");

		selector.handleInput("\t");
		await Bun.sleep(125);
		installTestTheme();

		expect(refreshProvider).toHaveBeenCalledWith("ollama-cloud", "online");
		expect(modelRegistry.refresh).toHaveBeenCalledTimes(1);
		const rendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(rendered).toContain("deepseek-v4-pro");
		expect(rendered).not.toContain("Provider has not been refreshed yet");
	});

	test("switches provider tabs immediately and refreshes in background with spinner animation", async () => {
		installTestTheme();
		const settings = Settings.isolated({});
		const discoveredModel = createOllamaCloudModel("deepseek-v4-pro");
		let availableModels: Model[] = [];
		let resolveRefresh: (() => void) | undefined;
		const refreshProvider = vi.fn(
			(_providerId: string, _strategy?: string) =>
				new Promise<void>(resolve => {
					resolveRefresh = () => {
						availableModels = [discoveredModel];
						resolve();
					};
				}),
		);
		const modelRegistry = {
			getAll: () => availableModels,
			refresh: vi.fn(async () => {}),
			refreshProvider,
			getError: () => undefined,
			getAvailable: () => availableModels,
			getDiscoverableProviders: () => ["ollama-cloud"],
			getProviderDiscoveryState: () => ({
				provider: "ollama-cloud",
				status: "idle",
				optional: false,
				stale: false,
				models: [],
			}),
		} as unknown as ModelRegistry;
		const ui = {
			requestRender: vi.fn(),
		} as unknown as TUI;

		const selector = new ModelSelectorComponent(
			ui,
			undefined,
			settings,
			modelRegistry,
			[],
			() => {},
			() => {},
		);
		await Bun.sleep(0);
		installTestTheme();

		selector.handleInput("\t");

		// Core regression: tab switch must not synchronously enter provider refresh.
		expect(refreshProvider).not.toHaveBeenCalled();

		const immediateRendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(immediateRendered).toContain("Refreshing OLLAMA CLOUD in background");

		await Bun.sleep(5);
		expect(refreshProvider).not.toHaveBeenCalled();
		await Bun.sleep(120);
		expect(refreshProvider).toHaveBeenCalledWith("ollama-cloud", "online");

		const spinnerFrame1 = selector.render(220).join("\n");
		await Bun.sleep(100);
		installTestTheme();
		const spinnerFrame2 = selector.render(220).join("\n");
		expect(normalizeRenderedText(spinnerFrame2)).toContain("Refreshing OLLAMA CLOUD in background");
		expect(spinnerFrame2).not.toEqual(spinnerFrame1);

		resolveRefresh?.();
		await Bun.sleep(10);
		installTestTheme();

		expect(modelRegistry.refresh).toHaveBeenCalledTimes(1);
		const finalRendered = normalizeRenderedText(selector.render(220).join("\n"));
		expect(finalRendered).toContain("deepseek-v4-pro");
		expect(finalRendered).not.toContain("Refreshing OLLAMA CLOUD in background");
	});
});