diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 16905fcec..18f2b7ddc 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Added + +- Added `dry-balance` CLI command to perform a dry-run OAuth account balancing check across configurable random session IDs, with sample and concurrency options, JSON output, and success/failure summary reporting +- Added `--json` output mode and machine-readable result format to `omp dry-balance` for automated use ### Fixed @@ -9337,4 +9341,4 @@ Initial public release. - Git branch display in footer - Message queueing during streaming responses - OAuth integration for Gmail and Google Calendar access -- HTML export with syntax highlighting and collapsible sections +- HTML export with syntax highlighting and collapsible sections \ No newline at end of file diff --git a/packages/coding-agent/src/cli-commands.ts b/packages/coding-agent/src/cli-commands.ts index 8fa568001..c9ac00741 100644 --- a/packages/coding-agent/src/cli-commands.ts +++ b/packages/coding-agent/src/cli-commands.ts @@ -20,6 +20,7 @@ export const commands: CommandEntry[] = [ { name: "completions", load: () => import("./commands/completions").then(m => m.default) }, { name: "__complete", load: () => import("./commands/complete").then(m => m.default) }, { name: "config", load: () => import("./commands/config").then(m => m.default) }, + { name: "dry-balance", load: () => import("./commands/dry-balance").then(m => m.default) }, { name: "grep", load: () => import("./commands/grep").then(m => m.default) }, { name: "grievances", load: () => import("./commands/grievances").then(m => m.default) }, { name: "install", load: () => import("./commands/install").then(m => m.default) }, diff --git a/packages/coding-agent/src/cli/dry-balance-cli.ts b/packages/coding-agent/src/cli/dry-balance-cli.ts new file mode 100644 index 000000000..03249c705 --- /dev/null +++ b/packages/coding-agent/src/cli/dry-balance-cli.ts @@ -0,0 +1,331 @@ +import type { Api, Model, OAuthAccess } from "@oh-my-pi/pi-ai"; +import { getProjectDir } from "@oh-my-pi/pi-utils"; +import chalk from "chalk"; +import { ModelRegistry } from "../config/model-registry"; +import { + formatModelString, + resolveAllowedModels, + resolveCliModel, + resolveModelRoleValue, + type ModelMatchPreferences, +} from "../config/model-resolver"; +import { Settings } from "../config/settings"; +import { discoverAuthStorage } from "../sdk"; + +const DEFAULT_SAMPLE_COUNT = 100; +const DEFAULT_CONCURRENCY = 32; + +export interface DryBalanceCommandArgs { + model?: string; + flags: { + model?: string; + count?: number; + concurrency?: number; + json?: boolean; + }; +} + +export interface DryBalanceAuthOptions { + baseUrl?: string; + modelId?: string; + signal?: AbortSignal; +} + +export interface DryBalanceAuthStorage { + getOAuthAccess(provider: string, sessionId?: string, options?: DryBalanceAuthOptions): Promise; +} + +export interface DryBalanceModelRegistry { + authStorage: DryBalanceAuthStorage; + getAll(): Model[]; + getAvailable(): Model[]; + getApiKey(model: Model, sessionId?: string): Promise; + getCanonicalVariants(model: Model): Model[]; + resolveCanonicalModel?(model: Model): Model | undefined; + getCanonicalId?(model: Model): 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 DryBalanceSummary { + model: string; + provider: string; + samples: number; + concurrency: number; + success: { + total: number; + accounts: DryBalanceAccountStat[]; + }; + failure: { + total: number; + reasons: DryBalanceFailureStat[]; + }; +} + +interface DryBalanceDependencies { + createRuntime?: () => Promise; + randomSessionId?: () => string; + writeStdout?: (text: string) => void; + writeStderr?: (text: string) => void; + setExitCode?: (code: number) => void; +} + +type DryBalanceAttemptResult = + | { + ok: true; + account: string; + } + | { + ok: false; + reason: 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; +} + +async function createDefaultRuntime(): Promise { + 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; 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, + modelRegistry: DryBalanceModelRegistry, + sessionId: string, +): Promise { + 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" }; + const account = access.email ?? access.accountId ?? access.projectId ?? access.enterpriseUrl ?? "(unknown oauth account)"; + return { ok: true, account }; + } catch (error) { + return { ok: false, reason: error instanceof Error ? error.message : String(error) }; + } +} + +async function mapConcurrent(items: T[], concurrency: number, fn: (item: T) => Promise): Promise { + const results = new Array(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]); + } + }), + ); + return results; +} + + +function sortedStats(map: Map, 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 summarizeResults( + model: Model, + samples: number, + concurrency: number, + results: DryBalanceAttemptResult[], +): DryBalanceSummary { + const accounts = new Map(); + const reasons = new Map(); + 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, + })); + return { + 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, + }, + }; +} + + +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), + ]; + return `${lines.join("\n")}\n`; +} + +export async function runDryBalanceCommand( + command: DryBalanceCommandArgs, + deps: DryBalanceDependencies = {}, +): Promise { + const samples = normalizePositiveInteger("count", command.flags.count, DEFAULT_SAMPLE_COUNT); + const concurrency = 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 runtime = await (deps.createRuntime ?? createDefaultRuntime)(); + 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`); + const sessionIds = Array.from({ length: samples }, () => randomSessionId()); + const results = await mapConcurrent(sessionIds, concurrency, sessionId => + runOneAttempt(model, runtime.modelRegistry, sessionId), + ); + const summary = summarizeResults(model, samples, concurrency, results); + if (command.flags.json) { + writeStdout(`${JSON.stringify(summary, null, 2)}\n`); + } else { + writeStdout(formatDryBalanceText(summary)); + } + if (summary.failure.total > 0) setExitCode(1); + return summary; + } finally { + runtime.close?.(); + } +} diff --git a/packages/coding-agent/src/commands/dry-balance.ts b/packages/coding-agent/src/commands/dry-balance.ts new file mode 100644 index 000000000..2e2c2631e --- /dev/null +++ b/packages/coding-agent/src/commands/dry-balance.ts @@ -0,0 +1,40 @@ +import { Args, Command, Flags } from "@oh-my-pi/pi-utils/cli"; +import { runDryBalanceCommand } from "../cli/dry-balance-cli"; + +export default class DryBalance extends Command { + static description = "Dry-run OAuth account balancing across random session ids"; + + static args = { + model: Args.string({ + description: "Model selector (provider/model or fuzzy id). Defaults to the configured default model.", + required: false, + }), + }; + + static flags = { + model: Flags.string({ description: "Model selector (same syntax as --model on omp)" }), + count: Flags.integer({ description: "Number of random session ids to try", default: 100 }), + concurrency: Flags.integer({ description: "Maximum concurrent credential resolutions", default: 32 }), + json: Flags.boolean({ description: "Output JSON" }), + }; + + static examples = [ + "# Dry-run the configured default model with 100 random session ids\n omp dry-balance", + "# Dry-run a specific model\n omp dry-balance anthropic/claude-sonnet-4-5", + "# Larger run with bounded concurrency\n omp dry-balance --model openai-codex/gpt-5-codex --count 1000 --concurrency 64", + "# Machine-readable output\n omp dry-balance --json", + ]; + + async run(): Promise { + const { args, flags } = await this.parse(DryBalance); + await runDryBalanceCommand({ + model: args.model, + flags: { + model: flags.model, + count: flags.count, + concurrency: flags.concurrency, + json: flags.json, + }, + }); + } +}