pi-provider-litellm 4.1.0 → 4.3.0

package/src/budget.ts

@@ -0,0 +1,561 @@
+import { createHmac, randomBytes } from "node:crypto";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { fetchJson, normalizeBaseUrl } from "./discover.js";
+import type { LiteLLMRuntimeAuth } from "./types.js";
+
+export type BudgetLevelName = "key" | "user" | "team" | "member" | "org";
+export const BUDGET_LEVELS: readonly BudgetLevelName[] = ["key", "user", "team", "member", "org"];
+
+export type BudgetEndpoint = "key" | "userV2" | "user" | "team" | "org";
+
+export interface BudgetLevel {
+  spend: number;
+  maxBudget: number;
+  resetAt?: number;
+}
+export type BudgetLevels = Partial<Record<BudgetLevelName, BudgetLevel>>;
+
+export type BudgetPoll = { ok: true; levels: BudgetLevels; keyPolled: boolean } | { ok: false; reason: string };
+
+type Obj = Record<string, unknown>;
+type Fetched = { status: "ok"; data: unknown } | { status: "denied" } | { status: "transient"; reason: string };
+
+function obj(value: unknown): Obj | undefined {
+  return typeof value === "object" && value !== null && !Array.isArray(value) ? (value as Obj) : undefined;
+}
+
+function str(value: unknown): string | undefined {
+  return typeof value === "string" && value !== "" ? value : undefined;
+}
+
+// A level exists only for a positive, finite limit; anything else means "no budget".
+function level(spend: unknown, maxBudget: unknown, resetAt: unknown): BudgetLevel | undefined {
+  if (typeof maxBudget !== "number" || !Number.isFinite(maxBudget) || maxBudget <= 0) return undefined;
+  const parsedReset = typeof resetAt === "string" ? Date.parse(resetAt) : Number.NaN;
+  return {
+    spend: typeof spend === "number" && Number.isFinite(spend) && spend >= 0 ? spend : 0,
+    maxBudget,
+    ...(Number.isFinite(parsedReset) ? { resetAt: parsedReset } : {}),
+  };
+}
+
+// Levels carried over from the last good poll when their endpoint failed transiently.
+function keep(previous: BudgetLevels, names: readonly BudgetLevelName[]): BudgetLevels {
+  const kept: BudgetLevels = {};
+  for (const name of names) if (previous[name]) kept[name] = previous[name];
+  return kept;
+}
+
+async function get(
+  auth: LiteLLMRuntimeAuth,
+  path: string,
+  endpoint: BudgetEndpoint,
+  denied: Map<BudgetEndpoint, number>,
+  timeoutMs: number,
+): Promise<Fetched> {
+  try {
+    const url = `${normalizeBaseUrl(auth.baseUrl, auth.allowInsecureHttp)}${path}`;
+    const result = await fetchJson<unknown>(url, auth.apiKey, { timeoutMs, headers: auth.headers });
+    if (result.ok) return { status: "ok", data: result.data };
+    // LiteLLM answers 500 when handling the request raises (no database, a credential it cannot look up, a
+    // version bug), which repeats on every call, so it is remembered like a 4xx. 429 and 502-504 stay retryable.
+    if ((result.status >= 400 && result.status < 500 && result.status !== 429) || result.status === 500) {
+      denied.set(endpoint, result.status);
+      return { status: "denied" };
+    }
+    return { status: "transient", reason: `HTTP ${result.status}` };
+  } catch {
+    return { status: "transient", reason: "request failed" };
+  }
+}
+
+type UserPoll = { levels: BudgetLevels; userId?: string; teamIds?: string[]; transient?: true };
+
+async function pollUser(
+  auth: LiteLLMRuntimeAuth,
+  denied: Map<BudgetEndpoint, number>,
+  previous: BudgetLevels,
+  timeoutMs: number,
+): Promise<UserPoll> {
+  const keepUser = (): UserPoll => ({ levels: keep(previous, ["user"]), transient: true });
+  let fetched: Fetched | undefined;
+  let nested = false;
+  if (!denied.has("userV2")) {
+    fetched = await get(auth, "/v2/user/info", "userV2", denied, timeoutMs);
+    if (fetched.status === "denied" && denied.get("userV2") === 404) fetched = undefined;
+  } else if (denied.get("userV2") !== 404) {
+    return { levels: {} };
+  }
+  if (!fetched) {
+    if (denied.has("user")) return { levels: {} };
+    nested = true;
+    fetched = await get(auth, "/user/info", "user", denied, timeoutMs);
+  }
+  if (fetched.status === "denied") return { levels: {} };
+  if (fetched.status === "transient") return keepUser();
+  const data = obj(fetched.data);
+  const body = nested ? obj(data?.user_info) : data;
+  if (!body) return keepUser();
+  const user = level(body.spend, body.max_budget, body.budget_reset_at);
+  // `/v2/user/info` lists team ids; `/user/info` lists team rows.
+  const teams = Array.isArray(data?.teams) ? data.teams.map((team) => str(team) ?? str(obj(team)?.team_id)) : [];
+  return {
+    levels: user ? { user } : {},
+    userId: str(body.user_id),
+    teamIds: teams.filter((team): team is string => team !== undefined),
+  };
+}
+
+type TeamPoll = { levels: BudgetLevels; organizationId?: string; transient?: true };
+
+async function pollTeam(
+  auth: LiteLLMRuntimeAuth,
+  teamId: string,
+  userId: string | undefined,
+  denied: Map<BudgetEndpoint, number>,
+  previous: BudgetLevels,
+  timeoutMs: number,
+): Promise<TeamPoll> {
+  const path = `/team/info?team_id=${encodeURIComponent(teamId)}`;
+  const fetched = await get(auth, path, "team", denied, timeoutMs);
+  if (fetched.status === "denied") return { levels: {} };
+  const body = fetched.status === "ok" ? obj(fetched.data) : undefined;
+  const info = obj(body?.team_info);
+  if (!info) return { levels: keep(previous, ["team", "member"]), transient: true };
+
+  const levels: BudgetLevels = {};
+  const team = level(info.spend, info.max_budget, info.budget_reset_at);
+  if (team) levels.team = team;
+  if (userId) {
+    const memberships = body?.team_memberships;
+    const own = Array.isArray(memberships) ? memberships.map(obj).find((m) => m?.user_id === userId) : undefined;
+    const table = obj(own?.litellm_budget_table) ?? obj(info.team_member_budget_table);
+    const member = level(own?.spend, table?.max_budget, table?.budget_reset_at);
+    if (member) levels.member = member;
+  }
+  return { levels, organizationId: str(info.organization_id) };
+}
+
+async function pollOrg(
+  auth: LiteLLMRuntimeAuth,
+  organizationId: string,
+  denied: Map<BudgetEndpoint, number>,
+  previous: BudgetLevels,
+  timeoutMs: number,
+): Promise<BudgetLevels> {
+  const fetched = await get(auth, "/organization/list", "org", denied, timeoutMs);
+  if (fetched.status === "denied") return {};
+  if (fetched.status === "transient" || !Array.isArray(fetched.data)) return keep(previous, ["org"]);
+  const entry = fetched.data.map(obj).find((o) => o?.organization_id === organizationId);
+  const table = obj(entry?.litellm_budget_table);
+  const org = level(entry?.spend, table?.max_budget, table?.budget_reset_at);
+  return org ? { org } : {};
+}
+
+export async function pollBudget(
+  auth: LiteLLMRuntimeAuth,
+  denied: Map<BudgetEndpoint, number>,
+  previous: BudgetLevels,
+  timeoutMs: number,
+): Promise<BudgetPoll> {
+  const levels: BudgetLevels = {};
+  let userId: string | undefined;
+  let teamId: string | undefined;
+  let organizationId: string | undefined;
+  const keyPolled = !denied.has("key");
+
+  if (keyPolled) {
+    const fetched = await get(auth, "/key/info", "key", denied, timeoutMs);
+    if (fetched.status === "transient") return { ok: false, reason: fetched.reason };
+    if (fetched.status === "ok") {
+      const info = obj(obj(fetched.data)?.info);
+      if (!info) return { ok: false, reason: "malformed response" };
+      const key = level(info.spend, info.max_budget, info.budget_reset_at);
+      if (key) levels.key = key;
+      userId = str(info.user_id);
+      teamId = str(info.team_id);
+      organizationId = str(info.organization_id);
+    }
+  }
+  const keyKnown = keyPolled && !denied.has("key");
+
+  const userPoll: Promise<UserPoll> =
+    userId || !keyKnown ? pollUser(auth, denied, previous, timeoutMs) : Promise.resolve({ levels: {} });
+  // A JWT (LiteLLM's own test: three dot-separated parts), such as an OIDC login's id_token, has no key row, and LiteLLM
+  // charges its user's team: the only one is certain, but with several it depends on the proxy's JWT settings. A
+  // virtual key's user may be in a team the key is not charged to, such as the master key's default user.
+  if (!keyKnown && auth.apiKey.split(".").length === 3) {
+    const user = await userPoll;
+    if (user.transient) Object.assign(levels, keep(previous, ["team", "member", "org"]));
+    else if (user.teamIds?.length === 1) [teamId, userId] = [user.teamIds[0], user.userId];
+  }
+  const [user, team] = await Promise.all([
+    userPoll,
+    teamId && !denied.has("team")
+      ? pollTeam(auth, teamId, userId, denied, previous, timeoutMs)
+      : ({ levels: {} } as TeamPoll),
+  ]);
+  Object.assign(levels, user.levels, team.levels);
+
+  const org = organizationId ?? team.organizationId;
+  if (org && userId && !denied.has("org")) Object.assign(levels, await pollOrg(auth, org, denied, previous, timeoutMs));
+  // The team was the only source of the org id, so a temporary team failure keeps the last org figure too.
+  else if (!org && team.transient) Object.assign(levels, keep(previous, ["org"]));
+  return { ok: true, levels, keyPolled: keyKnown };
+}
+
+export type BudgetDisplay = "all" | "tightest";
+export interface BudgetTheme {
+  fg(color: "dim" | "warning" | "error", text: string): string;
+}
+
+// A header counts when it is a non-blank string holding a finite number >= 0.
+function headerNumber(value: string | undefined): number | undefined {
+  if (value === undefined || value.trim() === "") return undefined;
+  const parsed = Number(value);
+  return Number.isFinite(parsed) && parsed >= 0 ? parsed : undefined;
+}
+
+export function mergeKeyHeaders(
+  key: BudgetLevel | undefined,
+  keyPolled: boolean,
+  headers: Record<string, string>,
+): BudgetLevel | undefined {
+  const spend = headerNumber(headers["x-litellm-key-spend"]);
+  if (spend === undefined) return key;
+  const limit = headerNumber(headers["x-litellm-key-max-budget"]);
+  const maxBudget = limit !== undefined && limit > 0 ? limit : undefined;
+  if (!keyPolled) return maxBudget === undefined ? undefined : { spend, maxBudget };
+  if (key) return { ...key, spend: Math.max(key.spend, spend), maxBudget: maxBudget ?? key.maxBudget };
+  return maxBudget === undefined ? undefined : { spend, maxBudget };
+}
+
+// Locale formatting rounds the decimal value (1.815 to 1.82), where toFixed rounds the binary one (to 1.81).
+const money = (value: number, decimals: number): string =>
+  `$${value.toLocaleString("en-US", { minimumFractionDigits: decimals, maximumFractionDigits: 2 })}`;
+
+export function formatCompactAmount(value: number): string {
+  if (value < 100) return money(value, Number.isInteger(value) ? 0 : 2);
+  if (value < 1000) return `$${Math.round(value)}`;
+  const [scaled, suffix] = value >= 1_000_000 ? [value / 1_000_000, "M"] : [value / 1000, "k"];
+  return `$${scaled < 10 ? scaled.toFixed(1).replace(/\.0$/, "") : Math.round(scaled)}${suffix}`;
+}
+
+export function formatRelativeReset(resetAt: number | undefined, now: number): string | undefined {
+  if (resetAt === undefined || resetAt <= now) return undefined;
+  const ms = resetAt - now;
+  if (ms >= 86_400_000) return `${Math.floor(ms / 86_400_000)}d`;
+  if (ms >= 3_600_000) return `${Math.floor(ms / 3_600_000)}h`;
+  return `${Math.max(1, Math.floor(ms / 60_000))}m`;
+}
+
+const percent = (l: BudgetLevel): number => Math.floor((l.spend / l.maxBudget) * 100);
+
+function colour(l: BudgetLevel): "dim" | "warning" | "error" {
+  const ratio = l.spend / l.maxBudget;
+  return ratio >= 1 ? "error" : ratio >= 0.8 ? "warning" : "dim";
+}
+
+export function formatBudgetStatus(
+  levels: BudgetLevels,
+  display: BudgetDisplay,
+  prefix: string,
+  theme: BudgetTheme,
+  now: number,
+): string | undefined {
+  const set = BUDGET_LEVELS.flatMap((name) => (levels[name] ? [{ name, level: levels[name] }] : []));
+  if (set.length === 0) return undefined;
+  const head = `${theme.fg("dim", prefix)} `;
+  const text = (name: string, l: BudgetLevel) =>
+    `${name} ${formatCompactAmount(l.spend)}/${formatCompactAmount(l.maxBudget)}`;
+  if (display === "all") {
+    return head + set.map(({ name, level: l }) => theme.fg(colour(l), text(name, l))).join(theme.fg("dim", " · "));
+  }
+  // Strict "<" keeps the first level in order on ties.
+  const { name, level: l } = set.reduce((a, b) =>
+    b.level.maxBudget - b.level.spend < a.level.maxBudget - a.level.spend ? b : a,
+  );
+  const reset = formatRelativeReset(l.resetAt, now);
+  return (
+    head +
+    theme.fg(colour(l), `${text(name, l)} (${percent(l)}%)`) +
+    (reset ? theme.fg("dim", ` · resets ${reset}`) : "")
+  );
+}
+
+export function formatBudgetDetails(
+  providerName: string,
+  levels: BudgetLevels,
+  denied: ReadonlyMap<BudgetEndpoint, number>,
+  now: number,
+): string {
+  const lines = [`LiteLLM ("${providerName}") budget`];
+  for (const name of BUDGET_LEVELS) {
+    const l = levels[name];
+    if (!l) continue;
+    const reset = formatRelativeReset(l.resetAt, now);
+    const limit = money(l.maxBudget, Number.isInteger(l.maxBudget) ? 0 : 2);
+    lines.push(
+      `  ${name.padEnd(8)}${money(l.spend, 2)} of ${limit} (${percent(l)}%)${reset ? `, resets in ${reset}` : ""}`,
+    );
+  }
+  const userV2 = denied.get("userV2");
+  const unreadable = [
+    ["key", denied.get("key")],
+    ["user", userV2 !== undefined && userV2 !== 404 ? userV2 : denied.get("user")],
+    ["team", denied.get("team")],
+    ["org", denied.get("org")],
+  ].flatMap(([name, code]) => (code === undefined ? [] : [`${name} (${code})`]));
+  if (lines.length === 1 && unreadable.length === 0) lines.push("  no budgets set");
+  if (unreadable.length > 0) lines.push(`  Not readable with this credential: ${unreadable.join(", ")}`);
+  return lines.join("\n");
+}
+
+export function budgetDisplaySetting(budgetSettings: unknown): { display: BudgetDisplay; warning?: string } {
+  const display = obj(budgetSettings)?.display;
+  if (display === undefined || display === "all" || display === "tightest") return { display: display ?? "all" };
+  return { display: "all", warning: `LiteLLM budget: unknown display ${JSON.stringify(display)}; using "all".` };
+}
+
+export const BUDGET_STATUS_KEY = "litellm-budget";
+const POLL_SETTLE_DELAY_MS = 15_000;
+const POLL_MIN_INTERVAL_MS = 60_000;
+
+export interface BudgetOptions {
+  providers: ReadonlyArray<{ name: string; displayName: string }>;
+  display: BudgetDisplay;
+  displayWarning?: string;
+  resolveAuth(ctx: ExtensionContext, providerName: string): Promise<LiteLLMRuntimeAuth | undefined>;
+  timeoutMs(): number;
+  disabledReason(): string | null;
+  hostOffline(): boolean;
+  missingCredentials(providerName: string): string;
+}
+
+type PollOutcome = "ok" | "no-auth" | { failed: string };
+
+interface ProviderState {
+  digest?: string;
+  levels: BudgetLevels;
+  denied: Map<BudgetEndpoint, number>;
+  keyPolled: boolean;
+  lastPollAt?: number;
+  lastTurnAt?: number;
+  inFlight?: Promise<PollOutcome>;
+}
+
+export function setupLiteLLMBudget(pi: ExtensionAPI, options: BudgetOptions): void {
+  const displayNames = new Map(options.providers.map(({ name, displayName }) => [name, displayName]));
+  const states = new Map<string, ProviderState>();
+  // Keyed like the MCP registration identity in index.ts, because custom headers can carry gateway credentials.
+  const identityKey = randomBytes(32);
+  let shownText: string | undefined;
+  let activeProvider: string | undefined;
+  let latestCtx: ExtensionContext | undefined;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  let generation = 0;
+  let warned = false;
+
+  const stateFor = (name: string): ProviderState => {
+    let state = states.get(name);
+    if (!state) {
+      state = { levels: {}, denied: new Map(), keyPolled: false };
+      states.set(name, state);
+    }
+    return state;
+  };
+  const forget = (state: ProviderState): void => {
+    state.levels = {};
+    state.denied.clear();
+    state.keyPolled = false;
+  };
+  const gated = (): boolean => options.disabledReason() !== null || options.hostOffline();
+
+  // Draws only the active provider's text, and only when it changed, so a context without
+  // setStatus is untouched while there is nothing to show.
+  function render(ctx: ExtensionContext, name: string): void {
+    if (!ctx.hasUI || activeProvider !== name) return;
+    const text = formatBudgetStatus(
+      stateFor(name).levels,
+      options.display,
+      displayNames.get(name)!,
+      ctx.ui.theme,
+      Date.now(),
+    );
+    if (text === shownText) return;
+    ctx.ui.setStatus(BUDGET_STATUS_KEY, text);
+    shownText = text;
+  }
+
+  async function run(ctx: ExtensionContext, name: string, state: ProviderState): Promise<PollOutcome> {
+    const started = generation;
+    state.lastPollAt = Date.now();
+    let auth: LiteLLMRuntimeAuth | undefined;
+    try {
+      auth = await options.resolveAuth(ctx, name);
+    } catch {
+      auth = undefined;
+    }
+    if (!auth) {
+      forget(state);
+      state.digest = undefined;
+      if (started === generation) render(ctx, name);
+      return "no-auth";
+    }
+    const headers = Object.entries(auth.headers ?? {}).sort(([a], [b]) => (a < b ? -1 : 1));
+    const identity = JSON.stringify([auth.baseUrl, auth.apiKey, headers]);
+    const digest = createHmac("sha256", identityKey).update(identity).digest("hex");
+    if (state.digest !== undefined && state.digest !== digest) {
+      forget(state);
+      if (started === generation) render(ctx, name);
+    }
+    state.digest = digest;
+    const result = await pollBudget(auth, state.denied, state.levels, options.timeoutMs());
+    if (result.ok) {
+      const levels = { ...result.levels };
+      if (!result.keyPolled && state.levels.key) levels.key = state.levels.key;
+      state.levels = levels;
+      state.keyPolled = result.keyPolled;
+    }
+    if (started === generation) render(ctx, name);
+    return result.ok ? "ok" : { failed: result.reason };
+  }
+
+  function poll(ctx: ExtensionContext, name: string): Promise<PollOutcome> {
+    const state = stateFor(name);
+    state.inFlight ??= run(ctx, name, state).finally(() => {
+      state.inFlight = undefined;
+    });
+    return state.inFlight;
+  }
+
+  function track(ctx: ExtensionContext, provider: string | undefined): string | undefined {
+    latestCtx = ctx;
+    activeProvider = provider !== undefined && displayNames.has(provider) ? provider : undefined;
+    return activeProvider;
+  }
+
+  function warnOnce(ctx: ExtensionContext): void {
+    if (warned || !options.displayWarning) return;
+    warned = true;
+    if (ctx.hasUI) ctx.ui.notify(options.displayWarning, "warning");
+    else process.stderr.write(`${options.displayWarning}\n`);
+  }
+
+  function schedulePoll(name: string): void {
+    const state = stateFor(name);
+    const now = Date.now();
+    const at = Math.max(state.lastTurnAt! + POLL_SETTLE_DELAY_MS, (state.lastPollAt ?? 0) + POLL_MIN_INTERVAL_MS);
+    timer = setTimeout(() => {
+      timer = undefined;
+      if (latestCtx && activeProvider === name && !gated()) {
+        poll(latestCtx, name).catch(() => {});
+        // A turn that ended within the settle delay may not have its spend written yet: poll once more. Compared
+        // with the planned time, not Date.now(), because a timer can fire a few milliseconds early.
+        if (state.lastTurnAt! + POLL_SETTLE_DELAY_MS > at) schedulePoll(name);
+      }
+    }, at - now);
+    timer.unref();
+  }
+
+  pi.on("session_start", (_event, ctx) => {
+    warnOnce(ctx);
+    if (!ctx.hasUI) return;
+    const name = track(ctx, ctx.model?.provider);
+    if (name && !gated()) void poll(ctx, name);
+  });
+
+  pi.on("model_select", (event, ctx) => {
+    if (!ctx.hasUI) return;
+    const previous = activeProvider;
+    const name = track(ctx, event.model.provider);
+    // A pending turn poll belongs to the provider its turns ran on; a newly selected one is polled below if stale.
+    if (name !== previous) {
+      clearTimeout(timer);
+      timer = undefined;
+    }
+    if (!name) {
+      if (shownText !== undefined) {
+        ctx.ui.setStatus(BUDGET_STATUS_KEY, undefined);
+        shownText = undefined;
+      }
+      return;
+    }
+    render(ctx, name);
+    const state = stateFor(name);
+    if (gated()) return;
+    if (state.lastPollAt === undefined || Date.now() - state.lastPollAt >= POLL_MIN_INTERVAL_MS) void poll(ctx, name);
+    // A turn whose poll was dropped when this provider was switched away from still needs one.
+    const missed = state.lastTurnAt !== undefined && state.lastTurnAt + POLL_SETTLE_DELAY_MS > (state.lastPollAt ?? 0);
+    if (missed && !timer) schedulePoll(name);
+  });
+
+  pi.on("turn_end", (_event, ctx) => {
+    if (!ctx.hasUI) return;
+    const name = track(ctx, ctx.model?.provider);
+    if (!name || gated()) return;
+    stateFor(name).lastTurnAt = Date.now();
+    if (!timer) schedulePoll(name);
+  });
+
+  pi.on("after_provider_response", (event, ctx) => {
+    if (!ctx.hasUI) return;
+    const name = ctx.model?.provider;
+    if (!name || !displayNames.has(name)) return;
+    const state = stateFor(name);
+    const key = mergeKeyHeaders(state.levels.key, state.keyPolled, event.headers ?? {});
+    if (key) state.levels.key = key;
+    else delete state.levels.key;
+    render(ctx, name);
+  });
+
+  pi.on("session_shutdown", (_event, ctx) => {
+    clearTimeout(timer);
+    timer = undefined;
+    generation++;
+    if (shownText !== undefined && ctx?.hasUI) ctx.ui.setStatus(BUDGET_STATUS_KEY, undefined);
+    shownText = undefined;
+    activeProvider = undefined;
+    latestCtx = undefined;
+  });
+
+  pi.registerCommand("litellm-budget", {
+    description: "Show LiteLLM key, user, team, member, and organization budgets",
+    getArgumentCompletions: (prefix) => {
+      const matches = [...displayNames.keys()].filter((name) => name.startsWith(prefix));
+      return matches.length > 0 ? matches.map((name) => ({ value: name, label: name })) : null;
+    },
+    handler: async (args, ctx) => {
+      const notify = (message: string, level: "info" | "warning"): void => {
+        if (ctx.hasUI) ctx.ui.notify(message, level);
+        else process.stderr.write(`${message}\n`);
+      };
+      const disabledReason = options.disabledReason();
+      if (disabledReason) {
+        notify(`LiteLLM: budget check skipped (${disabledReason}).`, "warning");
+        return;
+      }
+      const requested = args.trim();
+      const selected = requested ? options.providers.filter(({ name }) => name === requested) : options.providers;
+      if (selected.length === 0) {
+        const configured = options.providers.map(({ name }) => JSON.stringify(name)).join(", ");
+        notify(`LiteLLM: unknown provider ${JSON.stringify(requested)}; configured: ${configured}.`, "warning");
+        return;
+      }
+      for (const { name } of selected) {
+        const label = `LiteLLM (${JSON.stringify(name)})`;
+        const state = stateFor(name);
+        await state.inFlight?.catch(() => undefined);
+        state.denied.clear();
+        const outcome = await poll(ctx, name);
+        if (outcome === "no-auth") {
+          notify(`${label}: budget check skipped; ${options.missingCredentials(name)}`, "warning");
+        } else if (outcome !== "ok") {
+          notify(`${label}: budget check failed (${outcome.failed}).`, "warning");
+        } else {
+          notify(formatBudgetDetails(name, state.levels, state.denied, Date.now()), "info");
+        }
+      }
+    },
+  });
+}

package/src/discover.ts

@@ -8,6 +8,7 @@ import {
   resolveCatalogProvider,
   routesOnlyThrough,
 } from "./backend-identity.js";
+import { type DiscoveredImageModel, mapFromImageModelInfoGroup } from "./images.js";
 import {
   type CatalogResolution,
   catalogResolution,
@@ -87,7 +88,14 @@ export function normalizeBaseUrl(input: string, allowInsecureHttp = false): stri
   if (url.protocol !== "https:" && !(url.protocol === "http:" && (loopback || allowInsecureHttp))) {
     throw new Error("LiteLLM base URL must use HTTPS except for loopback hosts");
   }
-  return input.replace(/\/+$/, "").replace(/\/v1\/?$/i, "");
+  return trimTrailingSlashes(input).replace(/\/v1\/?$/i, "");
+}
+
+// A loop, not /\/+$/: that regex backtracks quadratically on a long slash run followed by another character.
+export function trimTrailingSlashes(value: string): string {
+  let end = value.length;
+  while (end > 0 && value[end - 1] === "/") end--;
+  return value.slice(0, end);
 }
 
 // Matches both the conventional `anthropic/...` prefix and aliases that
@@ -606,7 +614,7 @@ function loadDiscoveryPublicCatalog(options: DiscoveryOptions): Promise<PublicCa
   );
 }
 
-async function fetchJson<T>(
+export async function fetchJson<T>(
   url: string,
   apiKey: string,
   options: DiscoveryOptions,
@@ -1167,9 +1175,12 @@ export async function discoverModels(
   const progress = options.silent ? undefined : options.onProgress;
   progress?.("Querying /model/info endpoint...");
   const infoResult = await fetchJson<ModelInfoResponse>(`${base}/model/info`, apiKey, options);
-  if (infoResult.ok) {
+  // An empty `/model/info` list is treated like an unavailable endpoint: some proxies
+  // answer 200 with no rows for keys that can still list models through `/v1/models`.
+  const infoEntries = infoResult.ok ? (infoResult.data.data ?? []) : [];
+  if (infoEntries.length > 0) {
     const groups = new Map<string, ModelInfoEntry[]>();
-    for (const entry of infoResult.data.data ?? []) {
+    for (const entry of infoEntries) {
       // A route without a readable public name cannot be grouped or addressed.
       const route = wireString(entry.model_name);
       if (!route) continue;
@@ -1294,12 +1305,21 @@ export async function discoverModels(
           "set model_info.max_input_tokens on every deployment in each route",
       );
     }
-    return { source: "model_info", models, ...(proxyVersion ? { proxyVersion } : {}) };
+    // Image routes are a separate operation in Pi; the chat reducer above withholds them.
+    const imageModels = [...groups.entries()]
+      .map(([route, group]) => mapFromImageModelInfoGroup(route, group))
+      .filter((model): model is DiscoveredImageModel => model !== undefined);
+    return {
+      source: "model_info",
+      models,
+      ...(imageModels.length > 0 ? { imageModels } : {}),
+      ...(proxyVersion ? { proxyVersion } : {}),
+    };
   }
-  if (![401, 403, 404].includes(infoResult.status)) {
+  if (!infoResult.ok && ![401, 403, 404].includes(infoResult.status)) {
     throw new Error(`/model/info returned ${infoResult.status}`);
   }
-  progress?.("/model/info unavailable, trying /v1/models...");
+  progress?.("/model/info unavailable or empty, trying /v1/models...");
   const listResult = await fetchJson<ModelsListResponse>(`${base}/v1/models`, apiKey, options);
   if (!listResult.ok) {
     if ([401, 403, 404].includes(listResult.status)) {