@oh-my-pi/pi-coding-agent 15.9.1 → 15.9.5

package/src/cli/dry-balance-cli.ts

@@ -0,0 +1,823 @@
+import type {
+	Api,
+	AssistantMessage,
+	AssistantMessageEvent,
+	AssistantMessageEventStream,
+	Context,
+	Model,
+	OAuthAccess,
+	OAuthAccessResolution,
+	SimpleStreamOptions,
+} from "@oh-my-pi/pi-ai";
+import { streamSimple } from "@oh-my-pi/pi-ai";
+import { replaceTabs, truncateToWidth } from "@oh-my-pi/pi-tui";
+import { formatDuration, getProjectDir } from "@oh-my-pi/pi-utils";
+import chalk from "chalk";
+import type { CanonicalModelVariant } from "../config/model-equivalence";
+import { type CanonicalModelQueryOptions, ModelRegistry } from "../config/model-registry";
+import {
+	formatModelString,
+	type ModelMatchPreferences,
+	resolveAllowedModels,
+	resolveCliModel,
+	resolveModelRoleValue,
+} from "../config/model-resolver";
+import { Settings } from "../config/settings";
+import dryBalanceBenchPrompt from "../prompts/dry-balance-bench.md" with { type: "text" };
+import { discoverAuthStorage } from "../sdk";
+
+const DEFAULT_SAMPLE_COUNT = 100;
+const DEFAULT_CONCURRENCY = 32;
+const BENCH_MAX_TOKENS = 512;
+const BENCH_RENDER_INTERVAL_MS = 80;
+const BENCH_ACCOUNT_WIDTH = 60;
+const BENCH_ERROR_WIDTH = 110;
+const BENCH_SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"] as const;
+const DRY_BALANCE_BENCH_PROMPT = dryBalanceBenchPrompt.trim();
+
+export interface DryBalanceCommandArgs {
+	model?: string;
+	flags: {
+		model?: string;
+		count?: number;
+		concurrency?: number;
+		json?: boolean;
+		bench?: boolean;
+	};
+}
+
+export interface DryBalanceAuthOptions {
+	baseUrl?: string;
+	modelId?: string;
+	signal?: AbortSignal;
+}
+
+export interface DryBalanceAuthStorage {
+	getOAuthAccess(
+		provider: string,
+		sessionId?: string,
+		options?: DryBalanceAuthOptions,
+	): Promise<OAuthAccess | undefined>;
+	getOAuthAccesses?(provider: string, options?: DryBalanceAuthOptions): Promise<OAuthAccessResolution[]>;
+}
+
+export interface DryBalanceModelRegistry {
+	authStorage: DryBalanceAuthStorage;
+	getAll(): Model<Api>[];
+	getAvailable(): Model<Api>[];
+	getApiKey(model: Model<Api>, sessionId?: string): Promise<string | undefined>;
+	getCanonicalVariants(canonicalId: string, options?: CanonicalModelQueryOptions): CanonicalModelVariant[];
+	resolveCanonicalModel?(canonicalId: string, options?: CanonicalModelQueryOptions): Model<Api> | undefined;
+	getCanonicalId?(model: Model<Api>): string | undefined;
+}
+
+export interface DryBalanceRuntime {
+	modelRegistry: DryBalanceModelRegistry;
+	settings?: Settings;
+	close?: () => void;
+}
+
+export interface DryBalanceAccountStat {
+	account: string;
+	count: number;
+	percent: number;
+}
+
+export interface DryBalanceFailureStat {
+	reason: string;
+	count: number;
+	percent: number;
+}
+
+export interface DryBalanceBenchSuccessResult {
+	ok: true;
+	account: string;
+	ttftMs: number;
+	durationMs: number;
+	outputTokens: number;
+	tokensPerSecond: number;
+}
+
+export interface DryBalanceBenchFailureResult {
+	ok: false;
+	account?: string;
+	error: string;
+}
+
+export type DryBalanceBenchResult = DryBalanceBenchSuccessResult | DryBalanceBenchFailureResult;
+
+export interface DryBalanceBenchSummary {
+	total: number;
+	success: {
+		total: number;
+		averageTtftMs: number | null;
+		averageTokensPerSecond: number | null;
+	};
+	failure: {
+		total: number;
+		reasons: DryBalanceFailureStat[];
+	};
+	results: DryBalanceBenchResult[];
+}
+
+export interface DryBalanceSummary {
+	model: string;
+	provider: string;
+	samples: number;
+	concurrency: number;
+	success: {
+		total: number;
+		accounts: DryBalanceAccountStat[];
+	};
+	failure: {
+		total: number;
+		reasons: DryBalanceFailureStat[];
+	};
+	bench?: DryBalanceBenchSummary;
+}
+
+type DryBalanceStreamSimple = (
+	model: Model<Api>,
+	context: Context,
+	options?: SimpleStreamOptions,
+) => AssistantMessageEventStream;
+
+export interface DryBalanceDependencies {
+	createRuntime?: () => Promise<DryBalanceRuntime>;
+	randomSessionId?: () => string;
+	writeStdout?: (text: string) => void;
+	writeStderr?: (text: string) => void;
+	setExitCode?: (code: number) => void;
+	streamSimple?: DryBalanceStreamSimple;
+	now?: () => number;
+	stdoutIsTTY?: boolean;
+	stderrIsTTY?: boolean;
+}
+
+type DryBalanceAttemptResult =
+	| {
+			ok: true;
+			account: string;
+	  }
+	| {
+			ok: false;
+			reason: string;
+	  };
+
+type DryBalanceBenchProgressStatus =
+	| { state: "waiting" }
+	| { state: "running"; account: string }
+	| { state: "success"; result: DryBalanceBenchSuccessResult }
+	| { state: "failure"; result: DryBalanceBenchFailureResult };
+
+interface DryBalanceBenchProgressSink {
+	markRunning(index: number, account: string): void;
+	complete(index: number, result: DryBalanceBenchResult): void;
+	close(): void;
+}
+
+type DryBalanceBenchTarget =
+	| {
+			ok: true;
+			account: string;
+			accessToken: string;
+	  }
+	| {
+			ok: false;
+			account: string;
+			error: string;
+	  };
+
+function normalizePositiveInteger(name: string, value: number | undefined, fallback: number): number {
+	const resolved = value ?? fallback;
+	if (!Number.isInteger(resolved) || resolved <= 0) {
+		throw new Error(`--${name} must be a positive integer`);
+	}
+	return resolved;
+}
+
+function getErrorMessage(error: unknown): string {
+	if (error instanceof Error && error.message) return error.message;
+	const message = String(error);
+	return message ? message : "Unknown error";
+}
+
+function extractAccount(access: {
+	email?: string;
+	accountId?: string;
+	projectId?: string;
+	enterpriseUrl?: string;
+}): string {
+	return access.email ?? access.accountId ?? access.projectId ?? access.enterpriseUrl ?? "(unknown oauth account)";
+}
+
+function getBenchTargetKey(access: {
+	credentialId?: number;
+	email?: string;
+	accountId?: string;
+	projectId?: string;
+	enterpriseUrl?: string;
+	accessToken?: string;
+}): string {
+	return (
+		access.email ??
+		access.accountId ??
+		access.projectId ??
+		access.enterpriseUrl ??
+		(access.credentialId === undefined ? access.accessToken : `credential:${access.credentialId}`) ??
+		"(unknown oauth account)"
+	);
+}
+
+function sanitizeBenchText(text: string, width: number): string {
+	return truncateToWidth(replaceTabs(text).replace(/\r?\n/g, " "), width);
+}
+
+function formatBenchIndex(index: number, total: number): string {
+	return `#${String(index + 1).padStart(String(total).length, "0")}`;
+}
+
+function formatBenchAccount(account: string | undefined): string {
+	return account ? sanitizeBenchText(account, BENCH_ACCOUNT_WIDTH) : chalk.dim("(no account)");
+}
+
+function formatBenchDuration(ms: number): string {
+	return formatDuration(Math.max(0, Math.round(ms)));
+}
+
+function formatBenchTps(tokensPerSecond: number): string {
+	return `${tokensPerSecond.toFixed(1)}/s`;
+}
+
+function isBenchSuccess(result: DryBalanceBenchResult): result is DryBalanceBenchSuccessResult {
+	return result.ok;
+}
+
+function isBenchFirstTokenEvent(event: AssistantMessageEvent): boolean {
+	switch (event.type) {
+		case "text_delta":
+		case "thinking_delta":
+		case "toolcall_delta":
+			return event.delta.length > 0;
+		case "text_end":
+		case "thinking_end":
+			return event.content.length > 0;
+		default:
+			return false;
+	}
+}
+
+function resolveBenchMaxTokens(model: Model<Api>): number {
+	return Number.isFinite(model.maxTokens) && model.maxTokens > 0
+		? Math.min(BENCH_MAX_TOKENS, model.maxTokens)
+		: BENCH_MAX_TOKENS;
+}
+
+function normalizeBenchMs(value: number): number {
+	return Number.isFinite(value) && value > 0 ? value : 0;
+}
+
+function renderBenchResultLine(index: number, total: number, result: DryBalanceBenchResult): string {
+	const prefix = formatBenchIndex(index, total);
+	if (result.ok) {
+		return `${chalk.green("✓")} ${prefix} ${formatBenchAccount(result.account)} ${chalk.dim("TTFT")} ${formatBenchDuration(
+			result.ttftMs,
+		)} ${chalk.dim("TPS")} ${formatBenchTps(result.tokensPerSecond)}`;
+	}
+	return `${chalk.red("✗")} ${prefix} ${formatBenchAccount(result.account)} ${chalk.red(
+		sanitizeBenchText(result.error, BENCH_ERROR_WIDTH),
+	)}`;
+}
+
+function renderBenchStatusLine(
+	status: DryBalanceBenchProgressStatus,
+	index: number,
+	total: number,
+	frame: number,
+): string {
+	const prefix = formatBenchIndex(index, total);
+	switch (status.state) {
+		case "waiting":
+			return `${chalk.dim("○")} ${prefix} ${chalk.dim("waiting")}`;
+		case "running": {
+			const spinner = BENCH_SPINNER_FRAMES[frame % BENCH_SPINNER_FRAMES.length] ?? "*";
+			return `${chalk.yellow(spinner)} ${prefix} ${formatBenchAccount(status.account)} ${chalk.dim("sending request")}`;
+		}
+		case "success":
+			return renderBenchResultLine(index, total, status.result);
+		case "failure":
+			return renderBenchResultLine(index, total, status.result);
+	}
+}
+
+function createBenchProgressSink(
+	total: number,
+	write: (text: string) => void,
+	interactive: boolean,
+): DryBalanceBenchProgressSink {
+	const statuses: DryBalanceBenchProgressStatus[] = Array.from({ length: total }, () => ({ state: "waiting" }));
+	if (!interactive) {
+		return {
+			markRunning(index, account) {
+				statuses[index] = { state: "running", account };
+				write(`${renderBenchStatusLine(statuses[index], index, total, 0)}\n`);
+			},
+			complete(index, result) {
+				statuses[index] = result.ok ? { state: "success", result } : { state: "failure", result };
+				write(`${renderBenchResultLine(index, total, result)}\n`);
+			},
+			close() {},
+		};
+	}
+
+	let frame = 0;
+	let lineCount = 0;
+	let timer: NodeJS.Timeout | undefined;
+	const render = (): void => {
+		const lines = [
+			chalk.bold("bench requests"),
+			...statuses.map((status, index) => renderBenchStatusLine(status, index, total, frame)),
+		];
+		if (lineCount > 0) write(`\x1b[${lineCount}A`);
+		write(`${lines.map(line => `\x1b[2K${line}`).join("\n")}\n`);
+		lineCount = lines.length;
+	};
+	render();
+	timer = setInterval(() => {
+		frame += 1;
+		render();
+	}, BENCH_RENDER_INTERVAL_MS);
+	timer.unref?.();
+	return {
+		markRunning(index, account) {
+			statuses[index] = { state: "running", account };
+			render();
+		},
+		complete(index, result) {
+			statuses[index] = result.ok ? { state: "success", result } : { state: "failure", result };
+			render();
+		},
+		close() {
+			if (timer) {
+				clearInterval(timer);
+				timer = undefined;
+			}
+			render();
+		},
+	};
+}
+
+async function runBenchRequest(
+	model: Model<Api>,
+	sessionId: string,
+	account: string,
+	accessToken: string,
+	streamFn: DryBalanceStreamSimple,
+	now: () => number,
+): Promise<DryBalanceBenchResult> {
+	const startedAt = now();
+	let firstTokenAt: number | undefined;
+	try {
+		const context: Context = {
+			messages: [
+				{
+					role: "user",
+					content: DRY_BALANCE_BENCH_PROMPT,
+					timestamp: Date.now(),
+					attribution: "user",
+				},
+			],
+		};
+		const stream = streamFn(model, context, {
+			apiKey: accessToken,
+			sessionId,
+			maxTokens: resolveBenchMaxTokens(model),
+			temperature: 0.2,
+			disableReasoning: true,
+			hideThinkingSummary: true,
+		});
+		let message: AssistantMessage | undefined;
+		for await (const event of stream) {
+			if (firstTokenAt === undefined && isBenchFirstTokenEvent(event)) {
+				firstTokenAt = now();
+			}
+			if (event.type === "error") {
+				return { ok: false, account, error: event.error.errorMessage ?? "request failed" };
+			}
+			if (event.type === "done") {
+				message = event.message;
+			}
+		}
+		message ??= await stream.result();
+		if (message.stopReason === "error" || message.errorMessage) {
+			return { ok: false, account, error: message.errorMessage ?? "request failed" };
+		}
+		const durationMs = normalizeBenchMs(message.duration ?? now() - startedAt);
+		const ttftMs = normalizeBenchMs(
+			message.ttft ?? (firstTokenAt === undefined ? durationMs : firstTokenAt - startedAt),
+		);
+		const outputTokens = Number.isFinite(message.usage.output) && message.usage.output > 0 ? message.usage.output : 0;
+		const tokensPerSecond = durationMs > 0 ? (outputTokens * 1000) / durationMs : 0;
+		return {
+			ok: true,
+			account,
+			ttftMs,
+			durationMs,
+			outputTokens,
+			tokensPerSecond,
+		};
+	} catch (error) {
+		return { ok: false, account, error: getErrorMessage(error) };
+	}
+}
+
+async function resolveBenchTargets(
+	model: Model<Api>,
+	authStorage: DryBalanceAuthStorage,
+): Promise<DryBalanceBenchTarget[]> {
+	const resolved = authStorage.getOAuthAccesses
+		? await authStorage.getOAuthAccesses(model.provider, {
+				baseUrl: model.baseUrl,
+				modelId: model.id,
+			})
+		: await authStorage
+				.getOAuthAccess(model.provider, undefined, {
+					baseUrl: model.baseUrl,
+					modelId: model.id,
+				})
+				.then(access => (access ? [{ ok: true as const, ...access }] : []));
+	const targets: DryBalanceBenchTarget[] = [];
+	const seen = new Set<string>();
+	for (const entry of resolved) {
+		const key = getBenchTargetKey(entry);
+		if (seen.has(key)) continue;
+		seen.add(key);
+		const account = extractAccount(entry);
+		if (entry.ok) {
+			targets.push({ ok: true, account, accessToken: entry.accessToken });
+		} else {
+			targets.push({ ok: false, account, error: entry.error });
+		}
+	}
+	return targets;
+}
+
+async function runBenchTargets(
+	model: Model<Api>,
+	targets: DryBalanceBenchTarget[],
+	randomSessionId: () => string,
+	progress: DryBalanceBenchProgressSink | undefined,
+	streamFn: DryBalanceStreamSimple,
+	now: () => number,
+): Promise<DryBalanceBenchResult[]> {
+	return Promise.all(
+		targets.map(async (target, index) => {
+			if (!target.ok) {
+				const result: DryBalanceBenchFailureResult = {
+					ok: false,
+					account: target.account,
+					error: target.error,
+				};
+				progress?.complete(index, result);
+				return result;
+			}
+			progress?.markRunning(index, target.account);
+			const result = await runBenchRequest(
+				model,
+				randomSessionId(),
+				target.account,
+				target.accessToken,
+				streamFn,
+				now,
+			);
+			progress?.complete(index, result);
+			return result;
+		}),
+	);
+}
+
+async function createDefaultRuntime(): Promise<DryBalanceRuntime> {
+	const authStorage = await discoverAuthStorage();
+	try {
+		const settings = await Settings.init({ cwd: getProjectDir() });
+		const modelRegistry = new ModelRegistry(authStorage);
+		return {
+			modelRegistry,
+			settings,
+			close: () => authStorage.close(),
+		};
+	} catch (error) {
+		authStorage.close();
+		throw error;
+	}
+}
+
+async function resolveDryBalanceModel(
+	modelSelector: string | undefined,
+	modelRegistry: DryBalanceModelRegistry,
+	settings: Settings | undefined,
+	randomSessionId: () => string,
+): Promise<{ model: Model<Api>; warning?: string }> {
+	const preferences: ModelMatchPreferences = {
+		usageOrder: settings?.getStorage()?.getModelUsageOrder(),
+	};
+	if (modelSelector) {
+		const resolved = resolveCliModel({
+			cliModel: modelSelector,
+			modelRegistry,
+			preferences,
+		});
+		if (resolved.error) throw new Error(resolved.error);
+		if (!resolved.model) throw new Error(`Model "${modelSelector}" not found`);
+		return { model: resolved.model, warning: resolved.warning };
+	}
+
+	const allowedModels = await resolveAllowedModels(modelRegistry, settings, preferences);
+	if (allowedModels.length === 0) {
+		throw new Error(
+			"No models available. Use --model to select a model or configure enabledModels/default model settings.",
+		);
+	}
+
+	const defaultRoleSpec = resolveModelRoleValue(settings?.getModelRole("default"), allowedModels, {
+		settings,
+		matchPreferences: preferences,
+		modelRegistry,
+	});
+	if (defaultRoleSpec.model) {
+		return { model: defaultRoleSpec.model, warning: defaultRoleSpec.warning };
+	}
+
+	for (const candidate of allowedModels) {
+		const apiKey = await modelRegistry.getApiKey(candidate, randomSessionId());
+		if (apiKey) return { model: candidate };
+	}
+
+	return {
+		model: allowedModels[0],
+		warning:
+			"No allowed model had usable credentials during default resolution; dry-balance will report OAuth failures for the first allowed model.",
+	};
+}
+
+async function runOneAttempt(
+	model: Model<Api>,
+	modelRegistry: DryBalanceModelRegistry,
+	sessionId: string,
+): Promise<DryBalanceAttemptResult> {
+	try {
+		// AuthStorage.getOAuthAccess shares the OAuth credential ranking, refresh,
+		// usage-limit, broker, and session-sticky path used by getApiKey(), while
+		// returning the selected account metadata instead of bearer bytes.
+		const access = await modelRegistry.authStorage.getOAuthAccess(model.provider, sessionId, {
+			baseUrl: model.baseUrl,
+			modelId: model.id,
+		});
+		if (!access) return { ok: false, reason: "no OAuth access resolved" };
+		return { ok: true, account: extractAccount(access) };
+	} catch (error) {
+		return { ok: false, reason: getErrorMessage(error) };
+	}
+}
+
+async function mapConcurrent<T, R>(
+	items: T[],
+	concurrency: number,
+	fn: (item: T, index: number) => Promise<R>,
+): Promise<R[]> {
+	const results = new Array<R>(items.length);
+	let nextIndex = 0;
+	const workerCount = Math.min(concurrency, items.length);
+	await Promise.all(
+		Array.from({ length: workerCount }, async () => {
+			while (true) {
+				const index = nextIndex;
+				nextIndex += 1;
+				if (index >= items.length) return;
+				results[index] = await fn(items[index], index);
+			}
+		}),
+	);
+	return results;
+}
+
+function sortedStats(
+	map: Map<string, number>,
+	samples: number,
+): Array<{ label: string; count: number; percent: number }> {
+	return [...map.entries()]
+		.map(([label, count]) => ({ label, count, percent: (count / samples) * 100 }))
+		.sort((left, right) => right.count - left.count || left.label.localeCompare(right.label));
+}
+
+function summarizeBenchResults(results: DryBalanceBenchResult[]): DryBalanceBenchSummary | undefined {
+	if (results.length === 0) return undefined;
+	const successes = results.filter(isBenchSuccess);
+	const failureReasons = new Map<string, number>();
+	for (const result of results) {
+		if (!result.ok) {
+			failureReasons.set(result.error, (failureReasons.get(result.error) ?? 0) + 1);
+		}
+	}
+	const average = (values: number[]): number | null =>
+		values.length === 0 ? null : values.reduce((sum, value) => sum + value, 0) / values.length;
+	return {
+		total: results.length,
+		success: {
+			total: successes.length,
+			averageTtftMs: average(successes.map(result => result.ttftMs)),
+			averageTokensPerSecond: average(successes.map(result => result.tokensPerSecond)),
+		},
+		failure: {
+			total: results.length - successes.length,
+			reasons: sortedStats(failureReasons, results.length).map(stat => ({
+				reason: stat.label,
+				count: stat.count,
+				percent: stat.percent,
+			})),
+		},
+		results,
+	};
+}
+
+function summarizeResults(
+	model: Model<Api>,
+	samples: number,
+	concurrency: number,
+	results: DryBalanceAttemptResult[],
+): DryBalanceSummary {
+	const accounts = new Map<string, number>();
+	const reasons = new Map<string, number>();
+	for (const result of results) {
+		if (result.ok) {
+			accounts.set(result.account, (accounts.get(result.account) ?? 0) + 1);
+		} else {
+			reasons.set(result.reason, (reasons.get(result.reason) ?? 0) + 1);
+		}
+	}
+	const accountStats: DryBalanceAccountStat[] = sortedStats(accounts, samples).map(stat => ({
+		account: stat.label,
+		count: stat.count,
+		percent: stat.percent,
+	}));
+	const failureStats: DryBalanceFailureStat[] = sortedStats(reasons, samples).map(stat => ({
+		reason: stat.label,
+		count: stat.count,
+		percent: stat.percent,
+	}));
+	const summary: DryBalanceSummary = {
+		model: formatModelString(model),
+		provider: model.provider,
+		samples,
+		concurrency,
+		success: {
+			total: results.filter(result => result.ok).length,
+			accounts: accountStats,
+		},
+		failure: {
+			total: results.filter(result => !result.ok).length,
+			reasons: failureStats,
+		},
+	};
+	return summary;
+}
+
+function formatRows(rows: Array<{ count: number; percent: number; label: string }>): string[] {
+	if (rows.length === 0) return [`  ${chalk.dim("(none)")}`];
+	const maxCountWidth = Math.max(...rows.map(row => row.count.toString().length));
+	return rows.map(row => {
+		const count = row.count.toString().padStart(maxCountWidth);
+		const percent = `${row.percent.toFixed(1)}%`.padStart(6);
+		return `  ${count}  ${percent}  ${row.label}`;
+	});
+}
+
+export function formatDryBalanceText(summary: DryBalanceSummary): string {
+	const accountRows = summary.success.accounts.map(row => ({
+		count: row.count,
+		percent: row.percent,
+		label: row.account,
+	}));
+	const failureRows = summary.failure.reasons.map(row => ({
+		count: row.count,
+		percent: row.percent,
+		label: row.reason,
+	}));
+	const lines = [
+		chalk.bold("dry-balance"),
+		`model: ${summary.model}`,
+		`provider: ${summary.provider}`,
+		`samples: ${summary.samples}`,
+		`concurrency: ${summary.concurrency}`,
+		"",
+		`${chalk.green("success")} ${summary.success.total}`,
+		...formatRows(accountRows),
+		"",
+		`${summary.failure.total > 0 ? chalk.red("failure") : chalk.dim("failure")} ${summary.failure.total}`,
+		...formatRows(failureRows),
+	];
+	if (summary.bench) {
+		const avgTtft =
+			summary.bench.success.averageTtftMs === null ? "-" : formatBenchDuration(summary.bench.success.averageTtftMs);
+		const avgTps =
+			summary.bench.success.averageTokensPerSecond === null
+				? "-"
+				: formatBenchTps(summary.bench.success.averageTokensPerSecond);
+		const benchFailureRows = summary.bench.failure.reasons.map(row => ({
+			count: row.count,
+			percent: row.percent,
+			label: row.reason,
+		}));
+		lines.push(
+			"",
+			chalk.bold("bench"),
+			`requests: ${summary.bench.total}`,
+			`${chalk.green("success")} ${summary.bench.success.total}`,
+			`avg TTFT: ${avgTtft}`,
+			`avg TPS: ${avgTps}`,
+			"",
+			`${summary.bench.failure.total > 0 ? chalk.red("failure") : chalk.dim("failure")} ${summary.bench.failure.total}`,
+			...formatRows(benchFailureRows),
+		);
+	}
+	return `${lines.join("\n")}\n`;
+}
+
+export async function runDryBalanceCommand(
+	command: DryBalanceCommandArgs,
+	deps: DryBalanceDependencies = {},
+): Promise<DryBalanceSummary> {
+	const isBench = command.flags.bench === true;
+	const samples = isBench ? 0 : normalizePositiveInteger("count", command.flags.count, DEFAULT_SAMPLE_COUNT);
+	const concurrency = isBench
+		? 0
+		: Math.min(samples, normalizePositiveInteger("concurrency", command.flags.concurrency, DEFAULT_CONCURRENCY));
+	const randomSessionId = deps.randomSessionId ?? (() => Bun.randomUUIDv7());
+	const writeStdout = deps.writeStdout ?? ((text: string) => process.stdout.write(text));
+	const writeStderr = deps.writeStderr ?? ((text: string) => process.stderr.write(text));
+	const setExitCode =
+		deps.setExitCode ??
+		((code: number) => {
+			process.exitCode = code;
+		});
+	const streamFn = deps.streamSimple ?? streamSimple;
+	const now = deps.now ?? (() => performance.now());
+	const runtime = await (deps.createRuntime ?? createDefaultRuntime)();
+	let progress: DryBalanceBenchProgressSink | undefined;
+	let progressClosed = false;
+	const closeProgress = (): void => {
+		if (progressClosed) return;
+		progressClosed = true;
+		progress?.close();
+	};
+	try {
+		const modelSelector = command.flags.model ?? command.model;
+		const { model, warning } = await resolveDryBalanceModel(
+			modelSelector,
+			runtime.modelRegistry,
+			runtime.settings,
+			randomSessionId,
+		);
+		if (warning) writeStderr(`${chalk.yellow(`Warning: ${warning}`)}\n`);
+		let results: DryBalanceAttemptResult[];
+		let benchResults: DryBalanceBenchResult[] | undefined;
+		let summarySamples = samples;
+		let summaryConcurrency = concurrency;
+		if (isBench) {
+			const targets = await resolveBenchTargets(model, runtime.modelRegistry.authStorage);
+			if (targets.length === 0) throw new Error(`No OAuth accounts resolved for provider ${model.provider}`);
+			summarySamples = targets.length;
+			summaryConcurrency = targets.length;
+			const progressWrite = command.flags.json ? writeStderr : writeStdout;
+			const progressInteractive = command.flags.json
+				? (deps.stderrIsTTY ?? process.stderr.isTTY === true)
+				: (deps.stdoutIsTTY ?? process.stdout.isTTY === true);
+			progress = createBenchProgressSink(targets.length, progressWrite, progressInteractive);
+			benchResults = await runBenchTargets(model, targets, randomSessionId, progress, streamFn, now);
+			results = targets.map(target =>
+				target.ok ? { ok: true, account: target.account } : { ok: false, reason: target.error },
+			);
+		} else {
+			const sessionIds = Array.from({ length: samples }, () => randomSessionId());
+			results = await mapConcurrent(sessionIds, concurrency, sessionId =>
+				runOneAttempt(model, runtime.modelRegistry, sessionId),
+			);
+		}
+		closeProgress();
+		const summary = summarizeResults(model, summarySamples, summaryConcurrency, results);
+		if (benchResults) {
+			const benchSummary = summarizeBenchResults(benchResults);
+			if (benchSummary) summary.bench = benchSummary;
+		}
+		if (command.flags.json) {
+			writeStdout(`${JSON.stringify(summary, null, 2)}\n`);
+		} else {
+			writeStdout(formatDryBalanceText(summary));
+		}
+		if (summary.failure.total > 0 || (summary.bench?.failure.total ?? 0) > 0) setExitCode(1);
+		return summary;
+	} finally {
+		closeProgress();
+		runtime.close?.();
+	}
+}