pi-provider-litellm 4.0.1 → 4.2.0

package/README.md

@@ -156,7 +156,7 @@ Provider fields:
 | `baseUrl` | `LITELLM_BASE_URL` for `litellm`; required for aliases | LiteLLM proxy URL, with or without `/v1`. Must be a full `http`/`https` URL; a provider with no resolvable base URL exposes no models (see [Model host enforcement](#model-host-enforcement)) |
 | `apiKey` | `LITELLM_API_KEY_HELPER`/`LITELLM_API_KEY` for `litellm`; required for aliases | Pi config value for this provider's key. Use `$ENV_VAR`, `${ENV_VAR}`, `!command`, or a literal key. Escape a literal `$` as `$$`. |
 | `headers` | `$LITELLM_HEADERS` for `litellm`; unset for aliases | JSON string env reference or inline object of request headers |
-| `displayName` | provider name | Label shown in Pi UI |
+| `displayName` | `LITELLM_DISPLAY_NAME` / `"LiteLLM"` for `litellm`; alias name for aliases | Label shown in Pi UI |
 | `enabled` | `true` | Set `false` to skip an alias |
 | `oidc` | unset | `litellm` only. Sign in directly with an OpenID Connect identity provider instead of the LiteLLM-hosted flows; see [Direct OIDC login](#direct-oidc-login) |
 | `allowInsecureHttp` | `false` | Set `true` to permit plaintext HTTP for this provider, for example `http://host.docker.internal`. Credentials and request data will not be encrypted. Loopback HTTP works without this setting. |
@@ -165,7 +165,7 @@ Provider fields:
 
 ### Optional LiteLLM features
 
-LiteLLM Skills and MCP integration are enabled by default. Disable either feature, or choose how MCP tools reach the model, globally in `~/.pi/agent/settings.json`:
+LiteLLM Skills, MCP integration, and the budget footer are enabled by default. Disable any of them, choose how MCP tools reach the model, or choose how budgets are shown, globally in `~/.pi/agent/settings.json`:
 
 ```json
 {
@@ -180,12 +180,16 @@ LiteLLM Skills and MCP integration are enabled by default. Disable either featur
         "*delete*": "hidden",
         "github-search_code": "direct"
       }
+    },
+    "budget": {
+      "enabled": true,
+      "display": "all"
     }
   }
 }
 ```
 
-Setting `skills.enabled` to `false` disables the Skills Gateway management tools, skill fetching, and system-prompt injection. Setting `mcp.enabled` to `false` stops the extension from registering the proxy's MCP server. Restart Pi after changing these settings.
+Setting `skills.enabled` to `false` disables the Skills Gateway management tools, skill fetching, and system-prompt injection. Setting `mcp.enabled` to `false` stops the extension from registering the proxy's MCP server. Setting `budget.enabled` to `false` registers no budget footer, polls, or `/litellm-budget` command; `budget.display` is described under [Budget status](#budget-status). Restart Pi after changing these settings.
 
 `mcp.exposure` and `mcp.toolExposure` are passed to Pi unchanged and mean what they mean in a Pi `mcp.json` server entry: `codemode` (Pi's default when `exposure` is unset; callable from Pi's `codemode` scripts), `deferred` (declared to the model once Pi's `tool_search` loads it), `direct` (declared on every request), or `hidden`. `toolExposure` keys are tool names as LiteLLM's `/mcp` endpoint offers them, which LiteLLM prefixes with the server name, for example `github-search_code`; `*` matches any characters. Pi validates both settings, and an unusable value is reported once.
 
@@ -203,6 +207,43 @@ Every request routed through a configured LiteLLM provider includes Pi's canonic
 `x-litellm-session-id` header. This groups Chat Completions, Responses, and native Messages requests in LiteLLM without
 adding transport-specific fields to request bodies.
 
+## Budget status
+
+Pi's footer shows the LiteLLM budgets set for your credential, its user, its team, and its organization, for the active model's provider:
+
+```
+LiteLLM key $3.10/$10 · user $40/$100 · team $412/$1k · member $20/$50 · org $9.2k/$50k
+```
+
+Each segment is spend and limit for one level: `key`, `user`, `team`, `member` (your own budget inside the team), and `org`, in that order. These are the limits set on each level, and LiteLLM may not enforce every one of them on this key: since LiteLLM 1.95, a user's personal budget applies to a team key only when the operator sets `apply_user_budget_to_team_keys`. A level whose limit is unset, `null`, or `0` is not shown, so with no budgets the footer is empty. A segment is dim, turns to the warning colour from 80 % used, and to the error colour from 100 %. Set `"display": "tightest"` to show only the level with the least money left, with its percentage and reset time; any other `display` value falls back to `"all"` with a warning:
+
+```
+LiteLLM team $412/$1k (41%) · resets 12d
+```
+
+The footer carries only the provider's display name, these labels, and numbers; team names, key aliases, and proxy messages never reach it.
+
+The extension reads `/key/info` first, then `/v2/user/info` (`/user/info` on older proxies) and `/team/info`, then `/organization/list`, with the provider's own credential and custom headers, and only from that provider's proxy. LiteLLM's permissions decide what you can see:
+
+- A team key without a user shows no `user`, `member`, or `org` level.
+- A JWT credential, such as the one a [Direct OIDC login](#direct-oidc-login) presents, has no key row, so `/key/info` cannot be read. It shows the `user` level and, when that user belongs to exactly one team, that team's `team`, `member`, and `org` levels. With several teams, the one LiteLLM charges depends on the proxy's JWT settings, so only `user` shows. A virtual key whose `/key/info` cannot be read shows no team levels, because its user's team need not be the one the key is charged to.
+- Organization budgets need LiteLLM Enterprise.
+- A key whose `allowed_routes` exclude the info routes cannot be polled. It falls back to the `x-litellm-key-spend` and `x-litellm-key-max-budget` headers on every model response, which show only the `key` level, and only when the key has a limit.
+- An endpoint that answers 4xx (other than 429) or 500 is not asked again until the credential or the custom headers change, or you run `/litellm-budget`. LiteLLM answers 500 for errors that repeat on every call, such as a proxy without a database. Other errors keep the last values and are retried at the next refresh. Nothing is reported while polling.
+
+It refreshes when a session starts, when you select a LiteLLM model whose numbers are older than a minute, and after a turn: 15 seconds later, so LiteLLM can record the turn's spend, and never within 60 seconds of the previous poll. Idle sessions do not poll. Between polls the response headers can only raise the key's spend. LiteLLM writes spend to its database every 10 to 60 seconds, and spend resets are applied by a job that runs about every 10 minutes, so spend can lag by up to a minute and a reset by about ten minutes; a reset time already past is not shown.
+
+`/litellm-budget` forgets remembered 4xx and 500 answers, polls every configured LiteLLM provider now (or `/litellm-budget <provider>` for one), and prints a breakdown with two-decimal amounts, percentages, reset times, and the levels this credential cannot read, with their HTTP status:
+
+```
+LiteLLM ("litellm") budget
+  key     $3.10 of $10 (31%), resets in 12h
+  team    $412.30 of $1,000 (41%), resets in 24d
+  Not readable with this credential: org (401)
+```
+
+Automatic polls need Pi's interactive UI and do not run under `LITELLM_OFFLINE=1`, `LITELLM_DISCOVERY_TIMEOUT_MS=0`, or `PI_OFFLINE`; header updates still apply. `/litellm-budget` is stopped by the first two only. The credential and custom headers are kept in memory only as a keyed digest, and proxy responses are reduced to the numbers shown and never logged.
+
 ## Model transport
 
 A `/model/info` route group uses native Anthropic `/v1/messages` only when every deployment explicitly reports Chat mode, identifies a Claude backend through an Anthropic-, Bedrock-, or Vertex-capable adapter, and every declared backend identity resolves the same Anthropic compatibility policy. A `claude-<name>-<major>-<minor>` backend id resolves through its `claude-<name>-<major>` entry when the exact id is absent; an unknown generation stays on Chat Completions. When `supported_endpoints` is a list, it must include `/v1/messages`; an explicit endpoint list without it keeps the group on Chat Completions. A `null` value, which LiteLLM v1.100 and later report for models absent from their model map, counts as absent. An unresolved routing or base-model identity denies native Messages rather than being discarded from the unanimity check. Mixed-generation Claude groups can use Messages when their serializer policies match, with reasoning levels restricted to their common supported set. Groups with incompatible policies, mixed families, or unknown backend evidence remain on Chat Completions; unanimous explicit Responses mode takes precedence. Routes discovered through `/health` never select native Messages: Messages candidates are downgraded to Chat Completions with reasoning controls closed, while correlated Responses detail remains on Responses. Evidence-free `/v1/models` and wholly health-only `/health` groups take transport and presentation metadata from the bounded Pi catalog lookup: `openai-responses` when the catalog says so, otherwise Chat Completions. They expose no speculative reasoning selector and never select native Messages. Catalog identity such as Amazon Bedrock remains separate from the selected wire protocol.
@@ -214,11 +255,12 @@ Native Messages authenticates with `x-api-key`; every transport carries the `x-l
 | Variable | Default | Effect |
 |---|---|---|
 | `LITELLM_API_KEY_HELPER` | unset | Command that prints a fresh LiteLLM bearer token. Takes precedence over `LITELLM_API_KEY`. The extension runs it while resolving request auth, and Pi's per-request auth path is uncached, so rotating/short-lived tokens stay fresh. |
+| `LITELLM_DISPLAY_NAME` | `LiteLLM` | Custom display name for the default `litellm` provider in the Pi UI. |
 | `LITELLM_HEADERS` | unset | JSON object of extra headers sent to LiteLLM provider, discovery, MCP, and Skills Gateway requests. Provider aliases can use it with `"headers": "$LITELLM_HEADERS"`. |
 | `LITELLM_GCLOUD_TOKEN_AUTH` | unset | If set to a non-empty value other than `0`, use Google Application Default Credentials as the LiteLLM bearer token source. This takes precedence over `LITELLM_API_KEY_HELPER` and `LITELLM_API_KEY` when no stored `/login litellm` credential exists. |
 | `GOOGLE_APPLICATION_CREDENTIALS` | Google default ADC path | Optional path to an ADC JSON file used by `LITELLM_GCLOUD_TOKEN_AUTH`. If unset, the extension checks the default gcloud ADC locations. |
 | `LITELLM_OFFLINE` | unset | If `1`, disable all model and MCP discovery, including post-login discovery; use cached models only when their stored canonical proxy root exactly matches the active credential root, including any path prefix. URL-standard host casing and default ports are canonicalized, but paths remain case-sensitive. |
-| `PI_OFFLINE` | unset | Pi documents it as disabling automatic network activity, including model catalog refreshes. Through Pi 1.0.0, any set value (including `0` or an empty string) disables startup model discovery and `/model` network refreshes. Run `/litellm-refresh`, or unset it, to discover; cached models can still be used offline. |
+| `PI_OFFLINE` | unset | Pi documents it as disabling automatic network activity, including model catalog refreshes. Through Pi 1.0.3, any set value (including `0` or an empty string) disables startup model discovery and `/model` network refreshes. Run `/litellm-refresh`, or unset it, to discover; cached models can still be used offline. |
 | `LITELLM_DISCOVERY_TIMEOUT_MS` | `5000` | Background and explicit discovery fetch timeout in ms; `0` disables automatic discovery |
 | `LITELLM_CLI_JWT_EXPIRATION_HOURS` | `24` | CLI SSO token lifetime fallback for older proxies whose poll response omits `expires_in`; mirror a non-default proxy setting locally |
 | `LITELLM_VERBOSE_DISCOVERY` | unset | If `1`, enable progress messages during model discovery (login, refresh, startup), including startup skip reasons and defaulted `/model/info` context windows. Progress messages are off by default |

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-provider-litellm",
-  "version": "4.0.1",
+  "version": "4.2.0",
   "description": "LiteLLM proxy provider extension for Pi",
   "type": "module",
   "license": "MIT",
@@ -67,8 +67,8 @@
   },
   "devDependencies": {
     "@biomejs/biome": "^2.4.15",
-    "@earendil-works/pi-ai": "^1.0.0",
-    "@earendil-works/pi-coding-agent": "^1.0.0",
+    "@earendil-works/pi-ai": "^1.0.3",
+    "@earendil-works/pi-coding-agent": "^1.0.3",
     "@kitlangton/terminal-control": "1.2.1",
     "@types/node": "^26.0.0",
     "@typescript/native-preview": "7.0.0-dev.20260707.2",

package/src/backend-identity.ts

@@ -1,4 +1,4 @@
-export const LITELLM_DISCOVERY_VERSION = 5 as const;
+export const LITELLM_DISCOVERY_VERSION = 6 as const;
 
 export type BackendFamily = "claude" | "deepseek" | "gemini" | "kimi" | "openai";
 

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

@@ -26,7 +26,12 @@ import {
   wireString,
 } from "./model-groups.js";
 import { type ProxyVersion, type ProxyVersionFloor, probeProxyVersion, proxyVersionAtLeast } from "./proxy-version.js";
-import { loadPublicCatalog, type PublicCatalog, type PublicCatalogRecord } from "./public-catalog.js";
+import {
+  loadPublicCatalog,
+  PI_AZURE_PROVIDER,
+  type PublicCatalog,
+  type PublicCatalogRecord,
+} from "./public-catalog.js";
 import { intersectThinkingLevelMaps } from "./thinking-levels.js";
 import type {
   DiscoveredModel,
@@ -443,11 +448,11 @@ function nativeMessagesCatalog(
     : {};
 }
 
-const ADAPTER_CATALOG_PROVIDERS: Readonly<Record<string, BuiltinProvider>> = {
+const ADAPTER_CATALOG_PROVIDERS: Readonly<Record<string, BuiltinProvider | undefined>> = {
   anthropic: "anthropic",
   claude: "anthropic",
-  azure: "azure-openai-responses",
-  azure_ai: "azure-openai-responses",
+  azure: toKnownProvider(PI_AZURE_PROVIDER),
+  azure_ai: toKnownProvider(PI_AZURE_PROVIDER),
   bedrock: "amazon-bedrock",
   bedrock_converse: "amazon-bedrock",
   chatgpt: "openai-codex",
@@ -601,7 +606,7 @@ function loadDiscoveryPublicCatalog(options: DiscoveryOptions): Promise<PublicCa
   );
 }
 
-async function fetchJson<T>(
+export async function fetchJson<T>(
   url: string,
   apiKey: string,
   options: DiscoveryOptions,

package/src/index.ts

@@ -18,6 +18,7 @@ import type { ExtensionAPI, ExtensionContext, McpServerConfig } from "@earendil-
 // A namespace import, so a Pi without `VERSION` reaches the version check instead of failing to link.
 import * as piCodingAgent from "@earendil-works/pi-coding-agent";
 import { getAgentDir, readStoredCredential } from "@earendil-works/pi-coding-agent";
+import { budgetDisplaySetting, setupLiteLLMBudget } from "./budget.js";
 import { setupLiteLLMCostTracking } from "./cost.js";
 import { discoverModels, isGpt6OrNewerModel, isGpt55OrNewerModel, normalizeBaseUrl } from "./discover.js";
 import { getGcloudToken, hasGcloudAdcCredentials, isGcloudTokenAuthEnabled } from "./gcloud-token.js";
@@ -37,6 +38,7 @@ const PROVIDER_NAME = "litellm";
 const SETTINGS_KEY = "litellm";
 const ENV_BASE_URL = "LITELLM_BASE_URL";
 const ENV_API_KEY = "LITELLM_API_KEY";
+const ENV_DISPLAY_NAME = "LITELLM_DISPLAY_NAME";
 const GCLOUD_ADC_SOURCE = "gcloud ADC";
 const ENV_API_KEY_HELPER = "LITELLM_API_KEY_HELPER";
 const ENV_HEADERS = "LITELLM_HEADERS";
@@ -487,7 +489,10 @@ function normalizeProviderSettings(raw: unknown): RawProviderSettings | undefine
   return record;
 }
 
-function isFeatureEnabled(settings: Record<string, unknown> | undefined, feature: "skills" | "mcp"): boolean {
+function isFeatureEnabled(
+  settings: Record<string, unknown> | undefined,
+  feature: "skills" | "mcp" | "budget",
+): boolean {
   const raw = settings?.[feature];
   return !isPlainObject(raw) || raw.enabled !== false;
 }
@@ -496,6 +501,7 @@ function getProviderDefinitions(settings: Record<string, unknown> | undefined):
   const rawProviders = settings?.providers && typeof settings.providers === "object" ? settings.providers : undefined;
   const providerSettings = rawProviders as Record<string, unknown> | undefined;
   const defaultSettings = normalizeProviderSettings(providerSettings?.[PROVIDER_NAME]);
+  const defaultDisplayName = cleanConfig(process.env[ENV_DISPLAY_NAME]) ?? "LiteLLM";
 
   const makeDefinition = (
     name: string,
@@ -503,7 +509,7 @@ function getProviderDefinitions(settings: Record<string, unknown> | undefined):
     isDefault: boolean,
   ): ProviderDefinition => ({
     name,
-    displayName: stringSetting(raw?.displayName) ?? (isDefault ? "LiteLLM" : name),
+    displayName: stringSetting(raw?.displayName) ?? (isDefault ? defaultDisplayName : name),
     baseUrl: stringSetting(raw?.baseUrl),
     apiKeyConfig: stringSetting(raw?.apiKey),
     headers: raw?.headers ?? (isDefault ? `$${ENV_HEADERS}` : undefined),
@@ -1881,6 +1887,7 @@ export default async function (pi: ExtensionAPI): Promise<void> {
   const definitions = getProviderDefinitions(settings);
   const skillsEnabled = isFeatureEnabled(settings, "skills");
   const mcpEnabled = isFeatureEnabled(settings, "mcp");
+  const budgetEnabled = isFeatureEnabled(settings, "budget");
   const providerNames = new Set(definitions.map((definition) => definition.name));
   const oauthRuntimeRoots = new Map<string, { apiKey: string; root: string }>();
   let mcpUI: ExtensionContext["ui"] | undefined;
@@ -2228,6 +2235,21 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 
   setupLiteLLMCostTracking(pi, [...providerNames]);
 
+  if (budgetEnabled) {
+    const { display, warning } = budgetDisplaySetting(settings?.budget);
+    const definitionNamed = (name: string) => definitions.find((definition) => definition.name === name)!;
+    setupLiteLLMBudget(pi, {
+      providers: definitions.map(({ name, displayName }) => ({ name, displayName })),
+      display,
+      displayWarning: warning,
+      resolveAuth: (ctx, name) => getRuntimeAuth(ctx, definitionNamed(name)),
+      timeoutMs: getDiscoveryTimeoutMs,
+      disabledReason: discoveryDisabledReason,
+      hostOffline: isHostOffline,
+      missingCredentials: (name) => missingCredentials(definitionNamed(name)),
+    });
+  }
+
   if (skillsEnabled) {
     for (const tool of createSkillToolDefinitions(resolveDefaultRuntimeAuth)) {
       pi.registerTool(tool);

package/src/public-catalog.ts

@@ -8,6 +8,8 @@ export const MODELS_DEV_URL = "https://models.dev/api.json";
 const MODELS_DEV_CACHE_TTL_MS = 28 * 24 * 60 * 60 * 1000;
 const DEFAULT_TIMEOUT_MS = 5000;
 const KNOWN_PROVIDERS = new Set<string>(getProviders());
+// Pi 1.0.3 renamed its Azure catalog from `azure-openai-responses` to `azure`.
+export const PI_AZURE_PROVIDER = KNOWN_PROVIDERS.has("azure") ? "azure" : "azure-openai-responses";
 
 export interface PublicCatalogRecord {
   source: "models.dev" | "pi-vendor" | "pi-adapter";
@@ -72,8 +74,8 @@ const PROVIDER_ALIASES: Readonly<Record<string, readonly string[]>> = {
 
 const PI_PROVIDER_ALIASES: Readonly<Record<string, readonly string[]>> = {
   "amazon-bedrock": ["amazon-bedrock"],
-  azure: ["azure-openai-responses"],
-  azure_ai: ["azure-openai-responses"],
+  azure: [PI_AZURE_PROVIDER],
+  azure_ai: [PI_AZURE_PROVIDER],
   chatgpt: ["openai-codex"],
   "fireworks-ai": ["fireworks"],
   openai: ["openai"],