@0xkahi/pi-qol 0.0.1 → 0.0.3

package/src/extensions/custom-footer/subscription-usage-manager.ts

@@ -0,0 +1,91 @@
+import { AnthropicOauthUsageStrategy } from '../../libs/subscription-usage/strategy/anthropic-oauth-usage.strategy';
+import { OpenAiCodexUsageStrategy } from '../../libs/subscription-usage/strategy/openai-codex-usage.strategy';
+import {
+  type FetchUsageResponse,
+  type RateWindow,
+  SubscriptionUsageApi,
+  type SubscriptionUsageStrategy,
+} from '../../libs/subscription-usage/subscription-usage-api.util';
+import { SUBSCRIPTION_USAGE_TTL_MS } from './constants';
+import type { SupportedProvider } from './types';
+
+export type SubscriptionUsageApiLike = Pick<SubscriptionUsageApi, 'fetchUsage' | 'formatResetDescription'>;
+
+export function resolveSupportedProvider(provider: string | undefined): SupportedProvider | undefined {
+  switch (provider?.toLowerCase()) {
+    case 'anthropic':
+      return 'anthropic';
+    case 'openai-codex':
+    case 'codex':
+    case 'openai':
+    case 'chatgpt':
+      return 'openai-codex';
+    default:
+      return undefined;
+  }
+}
+
+export function pickSubscriptionUsageWindow(windows: readonly RateWindow[]): RateWindow | undefined {
+  let earliestWithReset: RateWindow | undefined;
+
+  for (const window of windows) {
+    if (!window.resetAt) continue;
+    if (!earliestWithReset?.resetAt || window.resetAt.getTime() <= earliestWithReset.resetAt.getTime()) {
+      earliestWithReset = window;
+    }
+  }
+
+  return earliestWithReset ?? windows[0];
+}
+
+export class SubscriptionUsageManager {
+  private cache = new Map<SupportedProvider, FetchUsageResponse>();
+  private inFlight = new Set<SupportedProvider>();
+  private lastFetchAt = new Map<SupportedProvider, number>();
+
+  constructor(
+    private api: SubscriptionUsageApiLike = new SubscriptionUsageApi(),
+    private onUpdate?: () => void,
+    private ttlMs = SUBSCRIPTION_USAGE_TTL_MS,
+  ) {}
+
+  get(provider: SupportedProvider): FetchUsageResponse | undefined {
+    return this.cache.get(provider);
+  }
+
+  ensureFresh(provider: SupportedProvider): FetchUsageResponse | undefined {
+    const last = this.lastFetchAt.get(provider) ?? 0;
+    if (!this.inFlight.has(provider) && Date.now() - last > this.ttlMs) {
+      this.refresh(provider);
+    }
+
+    return this.cache.get(provider);
+  }
+
+  formatResetDescription(date: Date): string {
+    return this.api.formatResetDescription(date);
+  }
+
+  private refresh(provider: SupportedProvider): void {
+    this.inFlight.add(provider);
+    this.lastFetchAt.set(provider, Date.now());
+
+    void this.api
+      .fetchUsage(this.strategyFor(provider))
+      .then(response => {
+        if (!response) return;
+        this.cache.set(provider, response);
+        this.onUpdate?.();
+      })
+      .catch(() => {
+        // Best-effort status: auth/network failures simply hide the segment.
+      })
+      .finally(() => {
+        this.inFlight.delete(provider);
+      });
+  }
+
+  private strategyFor(provider: SupportedProvider): SubscriptionUsageStrategy {
+    return provider === 'anthropic' ? new AnthropicOauthUsageStrategy() : new OpenAiCodexUsageStrategy();
+  }
+}

package/src/extensions/custom-footer/token-stats.ts

@@ -0,0 +1,144 @@
+import type { ContextUsage, SessionEntry } from '@earendil-works/pi-coding-agent';
+import { crayon } from '../../utils/crayon.util';
+import { clampPercent, renderProgressBar } from './progress-bar';
+import type { CustomFooterColors, CustomFooterDisplay, CustomFooterIcons, FooterTheme, SupportedProvider, UsageTotals } from './types';
+
+export type SubscriptionUsageSegmentInput = {
+  provider: SupportedProvider;
+  responseLabel: string;
+  windowLabel: string;
+  usedPercent: number;
+  resetDescription?: string;
+};
+
+export type BuildStatsLeftOptions = {
+  totals: UsageTotals;
+  display: CustomFooterDisplay;
+  icons: Pick<CustomFooterIcons, 'cache' | 'cacheRead' | 'cacheWrite'>;
+  contextUsage: ContextUsage | undefined;
+  contextWindow: number;
+  usingSubscription: boolean;
+  subscriptionUsageSegment?: string;
+  theme: FooterTheme;
+};
+
+export function formatTokens(count: number): string {
+  if (count < 1000) return count.toString();
+  if (count < 10000) return `${(count / 1000).toFixed(1)}k`;
+  if (count < 1000000) return `${Math.round(count / 1000)}k`;
+  if (count < 10000000) return `${(count / 1000000).toFixed(1)}M`;
+  return `${Math.round(count / 1000000)}M`;
+}
+
+export function calculateUsageTotals(entries: readonly SessionEntry[]): UsageTotals {
+  const totals: UsageTotals = {
+    totalInput: 0,
+    totalOutput: 0,
+    totalCacheRead: 0,
+    totalCacheWrite: 0,
+    totalCost: 0,
+  };
+
+  for (const entry of entries) {
+    if (entry.type !== 'message' || entry.message.role !== 'assistant') continue;
+
+    const usage = entry.message.usage;
+    totals.totalInput += usage.input;
+    totals.totalOutput += usage.output;
+    totals.totalCacheRead += usage.cacheRead;
+    totals.totalCacheWrite += usage.cacheWrite;
+    totals.totalCost += usage.cost.total;
+
+    const latestPromptTokens = usage.input + usage.cacheRead + usage.cacheWrite;
+    totals.latestCacheHitRate = latestPromptTokens > 0 ? (usage.cacheRead / latestPromptTokens) * 100 : undefined;
+  }
+
+  return totals;
+}
+
+export function buildSubscriptionUsageSegment({
+  colors,
+  icons,
+  theme,
+  usage,
+}: {
+  colors: Pick<CustomFooterColors, 'anthropicUsage' | 'codexUsage'>;
+  icons: Pick<CustomFooterIcons, 'refresh'>;
+  theme: FooterTheme;
+  usage: SubscriptionUsageSegmentInput;
+}): string {
+  const providerColor = usage.provider === 'anthropic' ? colors.anthropicUsage : colors.codexUsage;
+  const pct = clampPercent(usage.usedPercent);
+  const roundedPercent = Math.round(pct).toString();
+  const bar = renderProgressBar(pct);
+
+  const resetParts = usage.resetDescription ? [icons.refresh, usage.resetDescription].filter(Boolean).join(' ') : '';
+
+  return [
+    crayon.colorize(usage.responseLabel, { fg: providerColor }),
+    crayon.colorize(usage.windowLabel, { fg: providerColor }),
+    `${crayon.colorize(bar.filled, { fg: providerColor })}${theme.fg('dim', bar.empty)}`,
+    crayon.colorize(`${roundedPercent}%`, { fg: providerColor }),
+    resetParts ? theme.fg('dim', resetParts) : undefined,
+  ]
+    .filter((part): part is string => Boolean(part))
+    .join(' ');
+}
+
+export function buildStatsLeft({
+  contextUsage,
+  contextWindow,
+  display,
+  icons,
+  subscriptionUsageSegment,
+  theme,
+  totals,
+  usingSubscription,
+}: BuildStatsLeftOptions): string {
+  const parts: string[] = [];
+
+  if (display.tokens) {
+    if (totals.totalInput) parts.push(theme.fg('dim', `↑${formatTokens(totals.totalInput)}`));
+    if (totals.totalOutput) parts.push(theme.fg('dim', `↓${formatTokens(totals.totalOutput)}`));
+  }
+
+  if (display.cache && (totals.totalCacheRead > 0 || totals.totalCacheWrite > 0)) {
+    const hitRate = totals.latestCacheHitRate !== undefined ? ` ${totals.latestCacheHitRate.toFixed(1)}%` : '';
+    parts.push(
+      theme.fg(
+        'dim',
+        `${icons.cache}[${icons.cacheRead}${formatTokens(totals.totalCacheRead)} ${icons.cacheWrite}${formatTokens(totals.totalCacheWrite)}${hitRate}]`,
+      ),
+    );
+  }
+
+  if (totals.totalCost || usingSubscription) {
+    parts.push(theme.fg('dim', `$${totals.totalCost.toFixed(3)}${usingSubscription ? ' (sub)' : ''}`));
+  }
+
+  parts.push(buildContextUsageSegment({ contextUsage, contextWindow, theme }));
+
+  if (subscriptionUsageSegment) {
+    parts.push(theme.fg('dim', '•'), subscriptionUsageSegment);
+  }
+
+  return parts.join(' ');
+}
+
+function buildContextUsageSegment({
+  contextUsage,
+  contextWindow,
+  theme,
+}: {
+  contextUsage: ContextUsage | undefined;
+  contextWindow: number;
+  theme: FooterTheme;
+}): string {
+  const contextPercentValue = contextUsage?.percent ?? 0;
+  const contextPercent = contextUsage?.percent !== null ? contextPercentValue.toFixed(1) : '?';
+  const contextPercentDisplay = contextPercent === '?' ? `?/${formatTokens(contextWindow)}` : `${contextPercent}%/${formatTokens(contextWindow)}`;
+
+  if (contextPercentValue > 90) return theme.fg('error', contextPercentDisplay);
+  if (contextPercentValue > 70) return theme.fg('warning', contextPercentDisplay);
+  return theme.fg('dim', contextPercentDisplay);
+}

package/src/extensions/custom-footer/types.ts

@@ -0,0 +1,40 @@
+import type { ExtensionContext } from '@earendil-works/pi-coding-agent';
+import type { TUI } from '@earendil-works/pi-tui';
+import type { ConfigLoader } from '../../config-loader';
+import type { Config } from '../../schemas/config.schema';
+
+export type SupportedProvider = 'anthropic' | 'openai-codex';
+
+export type CustomFooterConfig = Config['custom_footer'];
+export type CustomFooterColors = CustomFooterConfig['colors'];
+export type CustomFooterDisplay = CustomFooterConfig['display'];
+export type CustomFooterIcons = CustomFooterConfig['icons'];
+
+export type FooterTheme = {
+  fg(color: 'dim' | 'error' | 'warning', text: string): string;
+};
+
+export type FooterDataProvider = {
+  getGitBranch(): string | null;
+  getExtensionStatuses(): ReadonlyMap<string, string>;
+  getAvailableProviderCount(): number;
+  onBranchChange(cb: () => void): () => void;
+};
+
+export type CustomFooterComponentDeps = {
+  tui: TUI;
+  theme: FooterTheme;
+  footerData: FooterDataProvider;
+  ctx: ExtensionContext;
+  config: ConfigLoader;
+  getThinkingLevel: () => string;
+};
+
+export type UsageTotals = {
+  totalInput: number;
+  totalOutput: number;
+  totalCacheRead: number;
+  totalCacheWrite: number;
+  totalCost: number;
+  latestCacheHitRate?: number;
+};

package/src/extensions/model-select/constants.ts

@@ -0,0 +1,8 @@
+import { piVimKeyEventId } from '../../constants';
+
+export const COMMAND_NAME = 'select-model';
+export const MAX_VISIBLE_MODELS = 10;
+export const MAX_CONFIG_WARNING_LINES = 4;
+
+// Cross-extension activation hook (for example, pi-vim-keys can emit this to open the picker).
+export const PI_VIM_KEY_EVENT_ID = piVimKeyEventId('model_select');

package/src/extensions/model-select/index.ts

@@ -0,0 +1,119 @@
+import type { Api, Model } from '@earendil-works/pi-ai';
+import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from '@earendil-works/pi-coding-agent';
+import type { ConfigLoader } from '../../config-loader';
+import { COMMAND_NAME, PI_VIM_KEY_EVENT_ID } from './constants';
+import { ModelFormatter } from './model-formatter';
+import { buildModelLists, findExactModel } from './model-lists';
+import { ModelSelectDialog } from './model-select-dialog';
+import type { DialogResult } from './types';
+
+async function applySelectedModel(pi: ExtensionAPI, ctx: ExtensionContext, model: Model<Api>): Promise<void> {
+  const success = await pi.setModel(model);
+  if (success) {
+    ctx.ui.notify(`Model set to ${ModelFormatter.modelLabel(model)}`, 'info');
+  } else {
+    ctx.ui.notify(`No configured auth for ${ModelFormatter.modelLabel(model)}`, 'error');
+  }
+}
+
+async function showModelSelector(pi: ExtensionAPI, args: string, ctx: ExtensionContext, configLoader: ConfigLoader): Promise<void> {
+  // `waitForIdle` only exists on command contexts. When invoked from the event
+  // bus we get a plain ExtensionContext, so fall back to a best-effort guard.
+  if ('waitForIdle' in ctx && typeof ctx.waitForIdle === 'function') {
+    await (ctx as ExtensionCommandContext).waitForIdle();
+  }
+  ctx.modelRegistry.refresh();
+
+  const exactModel = findExactModel(ctx, args);
+  if (exactModel) {
+    await applySelectedModel(pi, ctx, exactModel);
+    return;
+  }
+
+  if (!ctx.hasUI) {
+    ctx.ui.notify('The /select-model picker requires an interactive UI. Pass provider/modelId to select directly.', 'warning');
+    return;
+  }
+
+  const config = configLoader.getModelSelect();
+  const modelLists = await buildModelLists(ctx, config);
+  const registryError = ctx.modelRegistry.getError();
+
+  const selected = await ctx.ui.custom<DialogResult>(
+    (tui, theme, keybindings, done) =>
+      new ModelSelectDialog(tui, theme, keybindings, {
+        currentModel: ctx.model,
+        favouriteItems: modelLists.favouriteItems,
+        favouriteWarnings: modelLists.favouriteWarnings,
+        hasFavouriteSection: config.favourite.length > 0,
+        searchItems: modelLists.searchItems,
+        providerFilter: config.provider_filter,
+        configWarnings: registryError ? [`models.json: ${registryError}`] : [],
+        initialSearch: args.trim(),
+        layout: config.layout,
+        onDone: done,
+      }),
+    config.layout === 'overlay'
+      ? {
+          overlay: true,
+          overlayOptions: {
+            anchor: 'center',
+            width: '85%',
+            margin: 1,
+          },
+        }
+      : undefined,
+  );
+
+  if (selected) {
+    await applySelectedModel(pi, ctx, selected);
+  }
+}
+
+function activateModelSelect(pi: ExtensionAPI, deps: { config: ConfigLoader; initialCtx?: ExtensionContext }) {
+  // Keep a reference to the latest context so the event-bus handler (which gets
+  // no context of its own) can open the modal too. Lazy activation may happen
+  // during session_start, so seed this with the context that triggered it.
+  let latestCtx: ExtensionContext | undefined = deps.initialCtx;
+
+  pi.on('session_start', (_event, ctx) => {
+    latestCtx = ctx;
+  });
+
+  pi.registerCommand(COMMAND_NAME, {
+    description: 'Select/search models with favourites and provider filtering',
+    handler: async (args, ctx) => {
+      latestCtx = ctx;
+      if (!deps.config.isEnabled('model_select')) {
+        ctx.ui.notify('(pi-qol) model_select is disabled', 'warning');
+        return;
+      }
+      await showModelSelector(pi, args, ctx, deps.config);
+    },
+  });
+
+  // Cross-extension activation: another extension can open the modal via
+  //   pi.events.emit(PI_VIM_KEY_EVENT_ID)
+  pi.events.on(PI_VIM_KEY_EVENT_ID, () => {
+    const ctx = latestCtx;
+    if (!ctx || !deps.config.isEnabled('model_select')) {
+      return;
+    }
+    void showModelSelector(pi, '', ctx, deps.config).catch(error => {
+      ctx.ui.notify(`Failed to open model selector: ${error instanceof Error ? error.message : String(error)}`, 'error');
+    });
+  });
+}
+
+export function registerModelSelect(pi: ExtensionAPI, deps: { config: ConfigLoader }) {
+  let modelSelectRegistered = false;
+
+  pi.on('session_start', (_event, ctx) => {
+    if (!deps.config.isEnabled('model_select') || modelSelectRegistered) {
+      return;
+    }
+
+    activateModelSelect(pi, { config: deps.config, initialCtx: ctx });
+    modelSelectRegistered = true;
+  });
+}

package/src/extensions/model-select/model-formatter.ts

@@ -0,0 +1,60 @@
+import { type Api, type Model, modelsAreEqual } from '@earendil-works/pi-ai';
+import type { ModelItem } from './types';
+
+export class ModelFormatter {
+  static formatTokenCount(count: number): string {
+    if (count >= 1_000_000) {
+      const millions = count / 1_000_000;
+      return millions % 1 === 0 ? `${millions}M` : `${millions.toFixed(1)}M`;
+    }
+    if (count >= 1_000) {
+      const thousands = count / 1_000;
+      return thousands % 1 === 0 ? `${thousands}K` : `${thousands.toFixed(1)}K`;
+    }
+    return count.toString();
+  }
+
+  static modelLabel(model: Model<Api>): string {
+    return `${model.provider}/${model.id}`;
+  }
+
+  static describeModel(model: Model<Api>): string {
+    const capabilities: string[] = [];
+    if (model.reasoning) {
+      capabilities.push('thinking');
+    }
+    if (model.input.includes('image')) {
+      capabilities.push('images');
+    }
+
+    const capabilityText = capabilities.length > 0 ? ` • ${capabilities.join(', ')}` : '';
+    return `${model.name} • ctx ${ModelFormatter.formatTokenCount(model.contextWindow)} • max ${ModelFormatter.formatTokenCount(model.maxTokens)}${capabilityText}`;
+  }
+
+  static toModelItem(model: Model<Api>): ModelItem {
+    return {
+      model,
+      description: ModelFormatter.describeModel(model),
+      searchText: `${model.provider} ${model.id} ${model.provider}/${model.id} ${model.name}`,
+    };
+  }
+
+  static sortModels(models: Model<Api>[], currentModel: Model<Api> | undefined): Model<Api>[] {
+    return [...models].sort((a, b) => {
+      const aIsCurrent = modelsAreEqual(a, currentModel);
+      const bIsCurrent = modelsAreEqual(b, currentModel);
+      if (aIsCurrent && !bIsCurrent) {
+        return -1;
+      }
+      if (!aIsCurrent && bIsCurrent) {
+        return 1;
+      }
+
+      const provider = a.provider.localeCompare(b.provider);
+      if (provider !== 0) {
+        return provider;
+      }
+      return a.id.localeCompare(b.id);
+    });
+  }
+}

package/src/extensions/model-select/model-lists.ts

@@ -0,0 +1,63 @@
+import type { Api, Model } from '@earendil-works/pi-ai';
+import type { ExtensionContext } from '@earendil-works/pi-coding-agent';
+import type { Config } from '../../schemas/config.schema';
+import { ModelFormatter } from './model-formatter';
+import type { ModelItem, ModelLists } from './types';
+
+export async function buildModelLists(ctx: ExtensionContext, config: Config['model_select']): Promise<ModelLists> {
+  ctx.modelRegistry.refresh();
+
+  const availableModels = ctx.modelRegistry.getAvailable();
+  const providerFilterSet = new Set(config.provider_filter);
+  const searchModels = config.provider_filter.length > 0 ? availableModels.filter(model => providerFilterSet.has(model.provider)) : availableModels;
+  const sortedSearchModels = ModelFormatter.sortModels(searchModels, ctx.model).map(model => ModelFormatter.toModelItem(model));
+
+  const favouriteItems: ModelItem[] = [];
+  const favouriteWarnings: string[] = [];
+  const seenFavouriteModels = new Set<string>();
+
+  for (const favourite of config.favourite) {
+    const model = ctx.modelRegistry.find(favourite.provider, favourite.modelId);
+    const label = `${favourite.provider}/${favourite.modelId}`;
+
+    if (!model) {
+      favouriteWarnings.push(`${label} was not found`);
+      continue;
+    }
+
+    if (!ctx.modelRegistry.hasConfiguredAuth(model)) {
+      favouriteWarnings.push(`${label} has no configured auth`);
+      continue;
+    }
+
+    const key = `${model.provider}\u0000${model.id}`;
+    if (seenFavouriteModels.has(key)) {
+      continue;
+    }
+    seenFavouriteModels.add(key);
+    favouriteItems.push(ModelFormatter.toModelItem(model));
+  }
+
+  return { favouriteItems, favouriteWarnings, searchItems: sortedSearchModels };
+}
+
+export function findExactModel(ctx: ExtensionContext, args: string): Model<Api> | undefined {
+  const trimmed = args.trim();
+  if (!trimmed) {
+    return undefined;
+  }
+
+  const slashIndex = trimmed.indexOf('/');
+  if (slashIndex > 0 && slashIndex < trimmed.length - 1) {
+    const provider = trimmed.slice(0, slashIndex).trim();
+    const modelId = trimmed.slice(slashIndex + 1).trim();
+    return ctx.modelRegistry.find(provider, modelId);
+  }
+
+  const [provider, modelId] = trimmed.split(/\s+/, 2);
+  if (provider && modelId) {
+    return ctx.modelRegistry.find(provider, modelId);
+  }
+
+  return undefined;
+}