pi-provider-litellm 3.2.0 → 3.3.0

package/src/index.ts

@@ -1,5 +1,5 @@
 import { execFileSync } from "node:child_process";
-import { createHash, randomBytes } from "node:crypto";
+import { createHash, randomBytes, timingSafeEqual } from "node:crypto";
 import { existsSync, linkSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";
 import { readFile } from "node:fs/promises";
 import { createServer } from "node:http";
@@ -15,6 +15,7 @@ import type {
   OAuthCredentials,
   ProviderAuth,
 } from "@earendil-works/pi-ai";
+import { Type } from "@earendil-works/pi-ai";
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { getAgentDir, readStoredCredential } from "@earendil-works/pi-coding-agent";
 import { setupLiteLLMCostTracking } from "./cost.js";
@@ -24,6 +25,8 @@ import {
   createMcpToolDefinitions,
   credentialFingerprint,
   McpAccessDeniedError,
+  type McpExposure,
+  parseMcpExposurePolicy,
   reportMcpCatalogOutcome,
   reportMcpPartialDiscovery,
   reportMcpRegistrationFatal,
@@ -61,6 +64,9 @@ const CLI_SSO_EXPIRES_IN_SECONDS = 600;
 const CLI_AUTH_DISCOVERY_PATH = "/.well-known/litellm-cli-auth";
 const PKCE_CALLBACK_TIMEOUT_MS = 10 * 60 * 1000;
 const PKCE_FLOW = "litellm_cli_pkce";
+// Persisted in auth.json; never rename.
+const OIDC_FLOW = "oidc_pkce";
+const OIDC_DISCOVERY_PATH = "/.well-known/openid-configuration";
 const DEFAULT_CLI_JWT_EXPIRATION_HOURS = 24;
 const TOKEN_REFRESH_LEAD_MS = 5 * 60 * 1000;
 const PERMANENT_TOKEN_EXPIRES_AT = Number.MAX_SAFE_INTEGER;
@@ -74,6 +80,7 @@ type RawProviderSettings = {
   headers?: unknown;
   enabled?: unknown;
   allowInsecureHttp?: unknown;
+  oidc?: unknown;
 };
 
 type McpRuntimeAuth = LiteLLMRuntimeAuth & { mcpPauseSource?: string };
@@ -88,6 +95,8 @@ type ProviderDefinition = {
   useGcloudTokenAuth: boolean;
   enableOAuth: boolean;
   allowInsecureHttp: boolean;
+  /** Raw `oidc` setting, validated at login so a bad value never breaks startup. */
+  oidc?: unknown;
 };
 
 function isPlainObject(value: unknown): value is Record<string, unknown> {
@@ -507,6 +516,7 @@ function getProviderDefinitions(settings: Record<string, unknown> | undefined):
     useGcloudTokenAuth: isDefault,
     enableOAuth: isDefault,
     allowInsecureHttp: raw?.allowInsecureHttp === true,
+    oidc: isDefault ? raw?.oidc : undefined,
   });
 
   const definitions = [makeDefinition(PROVIDER_NAME, defaultSettings, true)];
@@ -609,7 +619,9 @@ type PkceToken = {
   teamId?: string;
 };
 
-type PkceTokenResult = { ok: true; token: PkceToken } | { ok: false; transient: boolean; message: string };
+type TokenFailure = { ok: false; transient: boolean; message: string };
+
+type PkceTokenResult = { ok: true; token: PkceToken } | TokenFailure;
 
 function authRequestHeaders(headers?: Record<string, string>, contentType?: string): Headers {
   const result = new Headers(headers);
@@ -642,12 +654,17 @@ function isAuthToken(value: unknown): value is string {
   return typeof value === "string" && /^[\x21-\x7e]+$/.test(value);
 }
 
+/** The OAuth `error` code when it is a plain code; descriptions and other fields are never echoed. */
+function oauthErrorCode(value: unknown): string | undefined {
+  return typeof value === "string" && /^[a-z_]{1,64}$/.test(value) ? value : undefined;
+}
+
 async function readAuthJson(response: Response, signal: AbortSignal | undefined, stage: string): Promise<unknown> {
   try {
     return await response.json();
   } catch {
     if (signal?.aborted) throw signal.reason;
-    throw new Error(`LiteLLM CLI auth ${stage} returned invalid JSON`);
+    throw new Error(`${stage} returned invalid JSON`);
   }
 }
 
@@ -664,7 +681,7 @@ async function discoverPkce(
   });
   if (response.status === 404) return undefined;
   if (!response.ok) throw new Error(`LiteLLM CLI auth discovery failed (HTTP ${response.status})`);
-  const data = await readAuthJson(response, signal, "discovery");
+  const data = await readAuthJson(response, signal, "LiteLLM CLI auth discovery");
   if (!isPlainObject(data) || data.contract_version !== 1)
     throw new Error("LiteLLM CLI auth discovery has unsupported contract version");
   if (!Array.isArray(data.code_challenge_methods_supported) || !data.code_challenge_methods_supported.includes("S256"))
@@ -688,13 +705,14 @@ async function discoverPkce(
   };
 }
 
-async function requestPkceToken(
+/** POSTs a token form without following redirects; each flow validates its own success fields. */
+async function postTokenForm(
   endpoint: string,
   form: URLSearchParams,
   signal: AbortSignal | undefined,
-  headers?: Record<string, string>,
-  existingRefreshToken?: string,
-): Promise<PkceTokenResult> {
+  headers: Record<string, string> | undefined,
+  label: string,
+): Promise<{ ok: true; data: Record<string, unknown> } | TokenFailure> {
   let response: Response;
   try {
     response = await fetch(endpoint, {
@@ -706,7 +724,7 @@ async function requestPkceToken(
     });
   } catch {
     if (signal?.aborted) throw signal.reason;
-    return { ok: false, transient: true, message: "LiteLLM token exchange failed (network error)" };
+    return { ok: false, transient: true, message: `${label} token exchange failed (network error)` };
   }
   let data: Record<string, unknown> | undefined;
   try {
@@ -715,23 +733,36 @@ async function requestPkceToken(
   } catch (error) {
     if (signal?.aborted) throw signal.reason;
     if (response.ok && !(error instanceof SyntaxError)) {
-      return { ok: false, transient: true, message: "LiteLLM token exchange failed (network error)" };
+      return { ok: false, transient: true, message: `${label} token exchange failed (network error)` };
     }
   }
   signal?.throwIfAborted();
   if (!response.ok) {
+    const code = oauthErrorCode(data?.error);
     return {
       ok: false,
       transient: response.status === 429 || response.status >= 500,
-      message:
-        data?.error === "invalid_grant"
-          ? "LiteLLM token exchange rejected (invalid_grant)"
-          : `LiteLLM token exchange failed (HTTP ${response.status})`,
+      message: code
+        ? `${label} token exchange rejected (${code})`
+        : `${label} token exchange failed (HTTP ${response.status})`,
     };
   }
-  const expires = typeof data?.expires_in === "number" ? Date.now() + data.expires_in * 1_000 : NaN;
+  return { ok: true, data: data ?? {} };
+}
+
+async function requestPkceToken(
+  endpoint: string,
+  form: URLSearchParams,
+  signal: AbortSignal | undefined,
+  headers?: Record<string, string>,
+  existingRefreshToken?: string,
+): Promise<PkceTokenResult> {
+  const response = await postTokenForm(endpoint, form, signal, headers, "LiteLLM");
+  if (!response.ok) return response;
+  const { data } = response;
+  const expires = typeof data.expires_in === "number" ? Date.now() + data.expires_in * 1_000 : NaN;
   if (
-    !isAuthToken(data?.access_token) ||
+    !isAuthToken(data.access_token) ||
     (!isAuthToken(data.refresh_token) && !isAuthToken(existingRefreshToken)) ||
     typeof data.token_type !== "string" ||
     data.token_type.toLowerCase() !== "bearer" ||
@@ -755,12 +786,26 @@ async function requestPkceToken(
   };
 }
 
-async function loginPkce(
-  interaction: AuthInteraction,
-  baseUrl: string,
-  discovery: CliAuthDiscovery,
-  headers?: Record<string, string>,
-): Promise<OAuthCredential> {
+function constantTimeEqual(value: unknown, expected: string): boolean {
+  if (typeof value !== "string") return false;
+  const left = Buffer.from(value);
+  const right = Buffer.from(expected);
+  return left.length === right.length && timingSafeEqual(left, right);
+}
+
+type LoopbackCallback = { redirectUri: string; state: string; code: () => Promise<string> };
+
+/**
+ * Serves an RFC 8252 loopback redirect on 127.0.0.1 for the lifetime of `run`. `code()` resolves with
+ * the authorization code of the first callback carrying the expected state.
+ */
+async function withLoopbackCallback<T>(
+  signal: AbortSignal | undefined,
+  label: string,
+  ports: readonly number[],
+  run: (callback: LoopbackCallback) => Promise<T>,
+): Promise<T> {
+  const state = randomBytes(32).toString("base64url");
   const server = createServer();
   let settleCallback: ((result: { code?: string; error?: Error }) => void) | undefined;
   server.on("request", (request, response) => {
@@ -780,14 +825,15 @@ async function loginPkce(
       response.writeHead(405, { Allow: "GET" }).end();
       return;
     }
-    if (url.searchParams.get("state") !== state) {
+    if (!constantTimeEqual(url.searchParams.get("state"), state)) {
       response.writeHead(400).end("Invalid OAuth state");
       return;
     }
     const error = url.searchParams.get("error");
     if (error) {
       response.writeHead(400).end("OAuth login failed");
-      settleCallback?.({ error: new Error("LiteLLM PKCE login was denied") });
+      const code = oauthErrorCode(error);
+      settleCallback?.({ error: new Error(`${label} login was denied${code ? ` (${code})` : ""}`) });
       return;
     }
     const code = url.searchParams.get("code");
@@ -802,102 +848,42 @@ async function loginPkce(
     response.end("<!doctype html><title>LiteLLM login complete</title><p>You can close this window.</p>");
     settleCallback?.({ code });
   });
-  const verifier = randomBytes(32).toString("base64url");
-  const state = randomBytes(32).toString("base64url");
   try {
-    await new Promise<void>((resolve, reject) => {
-      const onError = (error: Error) => reject(error);
-      server.once("error", onError);
-      server.listen(0, "127.0.0.1", () => {
-        server.off("error", onError);
-        resolve();
-      });
-    });
-    const address = server.address();
-    if (!address || typeof address === "string") throw new Error("LiteLLM PKCE callback failed to start");
-    const redirectUri = `http://127.0.0.1:${address.port}/callback`;
-    const registrationResponse = await fetch(discovery.registrationEndpoint, {
-      method: "POST",
-      headers: authRequestHeaders(headers, "application/json"),
-      body: JSON.stringify({
-        client_name: "pi-provider-litellm",
-        redirect_uris: [redirectUri],
-        token_endpoint_auth_method: "none",
-        grant_types: ["authorization_code", "refresh_token"],
-        response_types: ["code"],
-      }),
-      redirect: "manual",
-      signal: boundedLoginSignal(interaction.signal),
-    });
-    if (!registrationResponse.ok)
-      throw new Error(`LiteLLM PKCE client registration failed (HTTP ${registrationResponse.status})`);
-    const registration = await readAuthJson(registrationResponse, interaction.signal, "registration");
-    if (
-      !isPlainObject(registration) ||
-      !isAuthToken(registration.client_id) ||
-      !Array.isArray(registration.redirect_uris) ||
-      !registration.redirect_uris.includes(redirectUri)
-    ) {
-      throw new Error("LiteLLM PKCE client registration returned an invalid response");
+    // Try each pre-registered port in order; without any, the OS assigns one (RFC 8252 §7.3).
+    const candidates = ports.length > 0 ? ports : [0];
+    for (const [index, port] of candidates.entries()) {
+      try {
+        await new Promise<void>((resolve, reject) => {
+          const onError = (error: Error) => reject(error);
+          server.once("error", onError);
+          server.listen(port, "127.0.0.1", () => {
+            server.off("error", onError);
+            resolve();
+          });
+        });
+        break;
+      } catch (error) {
+        if (index === candidates.length - 1) throw error;
+      }
     }
-    const authorizationUrl = new URL(discovery.authorizationEndpoint);
-    for (const [key, value] of Object.entries({
-      client_id: registration.client_id,
-      redirect_uri: redirectUri,
-      response_type: "code",
-      resource: discovery.resource,
-      code_challenge: createHash("sha256").update(verifier).digest("base64url"),
-      code_challenge_method: "S256",
-      state,
-    }))
-      authorizationUrl.searchParams.set(key, value);
-    interaction.notify({
-      type: "auth_url",
-      url: authorizationUrl.toString(),
-      instructions: "Open this URL in a browser to sign in with LiteLLM.",
-    });
-    const code = await new Promise<string>((resolve, reject) => {
-      const timeout = setTimeout(() => finish(new Error("LiteLLM PKCE login timed out")), PKCE_CALLBACK_TIMEOUT_MS);
-      const onAbort = () => finish(interaction.signal?.reason ?? new Error("LiteLLM PKCE login cancelled"));
-      const finish = (error?: Error, value?: string) => {
-        clearTimeout(timeout);
-        interaction.signal?.removeEventListener("abort", onAbort);
-        settleCallback = undefined;
-        if (error) reject(error);
-        else resolve(value!);
-      };
-      settleCallback = ({ code, error }) => finish(error, code);
-      interaction.signal?.addEventListener("abort", onAbort, { once: true });
-      if (interaction.signal?.aborted) onAbort();
-    });
-    const result = await requestPkceToken(
-      discovery.tokenEndpoint,
-      new URLSearchParams({
-        grant_type: "authorization_code",
-        code,
-        redirect_uri: redirectUri,
-        client_id: registration.client_id,
-        code_verifier: verifier,
-        resource: discovery.resource,
-      }),
-      interaction.signal,
-      headers,
-    );
-    if (!result.ok) throw new Error(result.message);
-    const credential: PkceCredentials = {
-      type: "oauth",
-      access: result.token.access,
-      refresh: result.token.refresh,
-      expires: result.token.expires,
-      baseUrl,
-      flow: PKCE_FLOW,
-      clientId: registration.client_id,
-      tokenEndpoint: discovery.tokenEndpoint,
-      resource: discovery.resource,
-      userId: result.token.userId,
-      teamId: result.token.teamId,
-    };
-    return { ...credential, type: "oauth" };
+    const address = server.address();
+    if (!address || typeof address === "string") throw new Error(`${label} callback failed to start`);
+    const code = () =>
+      new Promise<string>((resolve, reject) => {
+        const timeout = setTimeout(() => finish(new Error(`${label} login timed out`)), PKCE_CALLBACK_TIMEOUT_MS);
+        const onAbort = () => finish(signal?.reason ?? new Error(`${label} login cancelled`));
+        const finish = (error?: Error, value?: string) => {
+          clearTimeout(timeout);
+          signal?.removeEventListener("abort", onAbort);
+          settleCallback = undefined;
+          if (error) reject(error);
+          else resolve(value!);
+        };
+        settleCallback = ({ code, error }) => finish(error, code);
+        signal?.addEventListener("abort", onAbort, { once: true });
+        if (signal?.aborted) onAbort();
+      });
+    return await run({ redirectUri: `http://127.0.0.1:${address.port}/callback`, state, code });
   } finally {
     await new Promise<void>((resolve) => {
       server.close(() => resolve());
@@ -906,6 +892,277 @@ async function loginPkce(
   }
 }
 
+async function loginPkce(
+  interaction: AuthInteraction,
+  baseUrl: string,
+  discovery: CliAuthDiscovery,
+  headers?: Record<string, string>,
+): Promise<OAuthCredential> {
+  const verifier = randomBytes(32).toString("base64url");
+  return withLoopbackCallback(
+    interaction.signal,
+    "LiteLLM PKCE",
+    [],
+    async ({ redirectUri, state, code: callbackCode }) => {
+      const registrationResponse = await fetch(discovery.registrationEndpoint, {
+        method: "POST",
+        headers: authRequestHeaders(headers, "application/json"),
+        body: JSON.stringify({
+          client_name: "pi-provider-litellm",
+          redirect_uris: [redirectUri],
+          token_endpoint_auth_method: "none",
+          grant_types: ["authorization_code", "refresh_token"],
+          response_types: ["code"],
+        }),
+        redirect: "manual",
+        signal: boundedLoginSignal(interaction.signal),
+      });
+      if (!registrationResponse.ok)
+        throw new Error(`LiteLLM PKCE client registration failed (HTTP ${registrationResponse.status})`);
+      const registration = await readAuthJson(
+        registrationResponse,
+        interaction.signal,
+        "LiteLLM CLI auth registration",
+      );
+      if (
+        !isPlainObject(registration) ||
+        !isAuthToken(registration.client_id) ||
+        !Array.isArray(registration.redirect_uris) ||
+        !registration.redirect_uris.includes(redirectUri)
+      ) {
+        throw new Error("LiteLLM PKCE client registration returned an invalid response");
+      }
+      const authorizationUrl = new URL(discovery.authorizationEndpoint);
+      for (const [key, value] of Object.entries({
+        client_id: registration.client_id,
+        redirect_uri: redirectUri,
+        response_type: "code",
+        resource: discovery.resource,
+        code_challenge: createHash("sha256").update(verifier).digest("base64url"),
+        code_challenge_method: "S256",
+        state,
+      }))
+        authorizationUrl.searchParams.set(key, value);
+      interaction.notify({
+        type: "auth_url",
+        url: authorizationUrl.toString(),
+        instructions: "Open this URL in a browser to sign in with LiteLLM.",
+      });
+      const code = await callbackCode();
+      const result = await requestPkceToken(
+        discovery.tokenEndpoint,
+        new URLSearchParams({
+          grant_type: "authorization_code",
+          code,
+          redirect_uri: redirectUri,
+          client_id: registration.client_id,
+          code_verifier: verifier,
+          resource: discovery.resource,
+        }),
+        interaction.signal,
+        headers,
+      );
+      if (!result.ok) throw new Error(result.message);
+      const credential: PkceCredentials = {
+        type: "oauth",
+        access: result.token.access,
+        refresh: result.token.refresh,
+        expires: result.token.expires,
+        baseUrl,
+        flow: PKCE_FLOW,
+        clientId: registration.client_id,
+        tokenEndpoint: discovery.tokenEndpoint,
+        resource: discovery.resource,
+        userId: result.token.userId,
+        teamId: result.token.teamId,
+      };
+      return { ...credential, type: "oauth" };
+    },
+  );
+}
+
+type OidcConfig = { issuer: string; clientId: string; scope: string; redirectPorts: number[] };
+
+type OidcDiscovery = { issuer: string; authorizationEndpoint: string; tokenEndpoint: string };
+
+type OidcCredentials = OAuthCredentials & {
+  flow: typeof OIDC_FLOW;
+  baseUrl: string;
+  issuer: string;
+  clientId: string;
+  tokenEndpoint: string;
+  subject: string;
+};
+
+type OidcTokenResult =
+  | { ok: true; token: { access: string; refresh: string; expires: number; subject: string } }
+  | TokenFailure;
+
+/** An https URL without credentials or fragment, returned verbatim. */
+function httpsUrl(value: unknown, { allowQuery }: { allowQuery: boolean }): string | undefined {
+  if (typeof value !== "string" || value.includes("#") || (!allowQuery && value.includes("?"))) return undefined;
+  try {
+    const url = new URL(value);
+    return url.protocol === "https:" && !url.username && !url.password ? value : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+function parseOidcConfig(raw: unknown): OidcConfig {
+  const invalid = (field: string, expected: string) =>
+    new Error(`Invalid LiteLLM ${field} setting: expected ${expected}`);
+  if (!isPlainObject(raw)) throw invalid("oidc", "an object");
+  const issuer = httpsUrl(raw.issuer, { allowQuery: false });
+  if (!issuer) throw invalid("oidc.issuer", "an https URL without credentials, query, or fragment");
+  if (!isAuthToken(raw.clientId)) throw invalid("oidc.clientId", "a non-empty string without spaces");
+  const scope = raw.scope ?? "openid";
+  if (typeof scope !== "string" || !scope.split(" ").includes("openid"))
+    throw invalid("oidc.scope", 'a space-separated scope list that includes "openid"');
+  const redirectPorts = raw.redirectPorts ?? [];
+  if (
+    !Array.isArray(redirectPorts) ||
+    !redirectPorts.every((port) => Number.isInteger(port) && port >= 1 && port <= 65535)
+  )
+    throw invalid("oidc.redirectPorts", "an array of port numbers from 1 to 65535");
+  return { issuer, clientId: raw.clientId, scope, redirectPorts };
+}
+
+// IdP requests never carry the proxy's headers: LITELLM_HEADERS and provider headers can hold gateway credentials.
+async function discoverOidc(issuer: string, signal?: AbortSignal): Promise<OidcDiscovery> {
+  const withoutSlash = (value: string) => value.replace(/\/+$/, "");
+  const response = await fetch(`${withoutSlash(issuer)}${OIDC_DISCOVERY_PATH}`, {
+    headers: authRequestHeaders(),
+    redirect: "manual",
+    signal: boundedLoginSignal(signal),
+  });
+  if (!response.ok) throw new Error(`OIDC discovery failed (HTTP ${response.status})`);
+  const data = await readAuthJson(response, signal, "OIDC discovery");
+  if (!isPlainObject(data) || typeof data.issuer !== "string" || withoutSlash(data.issuer) !== withoutSlash(issuer))
+    throw new Error("OIDC discovery issuer does not match the configured issuer");
+  const methods = data.code_challenge_methods_supported;
+  if (methods !== undefined && (!Array.isArray(methods) || !methods.includes("S256")))
+    throw new Error("OIDC discovery does not support S256");
+  const authorizationEndpoint = httpsUrl(data.authorization_endpoint, { allowQuery: true });
+  if (!authorizationEndpoint) throw new Error("OIDC discovery has invalid authorization_endpoint");
+  const tokenEndpoint = httpsUrl(data.token_endpoint, { allowQuery: true });
+  if (!tokenEndpoint) throw new Error("OIDC discovery has invalid token_endpoint");
+  return { issuer: data.issuer, authorizationEndpoint, tokenEndpoint };
+}
+
+function jwtClaims(token: string): Record<string, unknown> | undefined {
+  const [, payload, ...rest] = token.split(".");
+  if (rest.length !== 1 || !payload || !/^[\w-]+$/.test(payload)) return undefined;
+  try {
+    const claims: unknown = JSON.parse(Buffer.from(payload, "base64url").toString("utf8"));
+    return isPlainObject(claims) ? claims : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+/**
+ * Exchanges a grant for an id_token and checks its claims (OIDC Core §3.1.3.7). The signature is
+ * deliberately not verified here: the proxy verifies it against the IdP's JWKS before honouring the
+ * bearer, so these checks only stop Pi from storing a token meant for another client or login.
+ */
+async function requestOidcToken(
+  endpoint: string,
+  form: URLSearchParams,
+  signal: AbortSignal | undefined,
+  expected: { issuer: string; clientId: string; nonce?: string; subject?: string },
+  existingRefreshToken = "",
+): Promise<OidcTokenResult> {
+  const response = await postTokenForm(endpoint, form, signal, undefined, "OIDC");
+  if (!response.ok) return response;
+  const { id_token: idToken, refresh_token: refreshToken } = response.data;
+  const claims = isAuthToken(idToken) ? jwtClaims(idToken) : undefined;
+  if (!isAuthToken(idToken) || !claims)
+    return { ok: false, transient: false, message: "OIDC token response has no valid id_token" };
+  const invalid = (claim: string): TokenFailure => ({
+    ok: false,
+    transient: false,
+    message: `OIDC id_token has invalid ${claim}`,
+  });
+  if (claims.iss !== expected.issuer) return invalid("iss");
+  if (!(Array.isArray(claims.aud) ? claims.aud : [claims.aud]).includes(expected.clientId)) return invalid("aud");
+  // OIDC Core §3.1.3.7 (errata set 2): azp is optional, even with several audiences, but when present it must name us.
+  if (claims.azp !== undefined && claims.azp !== expected.clientId) return invalid("azp");
+  if (typeof claims.sub !== "string" || !claims.sub) return invalid("sub");
+  if (expected.subject !== undefined && claims.sub !== expected.subject) return invalid("sub");
+  const expiresAt = typeof claims.exp === "number" ? Math.floor(claims.exp * 1_000) : NaN;
+  if (!Number.isSafeInteger(expiresAt) || expiresAt <= Date.now()) return invalid("exp");
+  if (expected.nonce !== undefined && !constantTimeEqual(claims.nonce, expected.nonce)) return invalid("nonce");
+  return {
+    ok: true,
+    token: {
+      access: idToken,
+      // The refresh token is optional, and an IdP that does not rotate it keeps the existing one valid.
+      refresh: isAuthToken(refreshToken) ? refreshToken : existingRefreshToken,
+      // Refresh ahead of expiry, but by at most half the remaining lifetime, so a short-lived id_token is
+      // still used rather than treated as already expired.
+      expires: expiresAt - Math.ceil(Math.min(TOKEN_REFRESH_LEAD_MS, (expiresAt - Date.now()) / 2)),
+      subject: claims.sub,
+    },
+  };
+}
+
+async function loginOidc(interaction: AuthInteraction, baseUrl: string, config: OidcConfig): Promise<OAuthCredential> {
+  const discovery = await discoverOidc(config.issuer, interaction.signal);
+  const verifier = randomBytes(32).toString("base64url");
+  const nonce = randomBytes(32).toString("base64url");
+  return withLoopbackCallback(
+    interaction.signal,
+    "OIDC",
+    config.redirectPorts,
+    async ({ redirectUri, state, code }) => {
+      const authorizationUrl = new URL(discovery.authorizationEndpoint);
+      for (const [key, value] of Object.entries({
+        response_type: "code",
+        client_id: config.clientId,
+        redirect_uri: redirectUri,
+        scope: config.scope,
+        state,
+        nonce,
+        code_challenge: createHash("sha256").update(verifier).digest("base64url"),
+        code_challenge_method: "S256",
+      }))
+        authorizationUrl.searchParams.set(key, value);
+      interaction.notify({
+        type: "auth_url",
+        url: authorizationUrl.toString(),
+        instructions: "Open this URL in a browser to sign in with your identity provider.",
+      });
+      const result = await requestOidcToken(
+        discovery.tokenEndpoint,
+        new URLSearchParams({
+          grant_type: "authorization_code",
+          code: await code(),
+          redirect_uri: redirectUri,
+          client_id: config.clientId,
+          code_verifier: verifier,
+        }),
+        interaction.signal,
+        { issuer: discovery.issuer, clientId: config.clientId, nonce },
+      );
+      if (!result.ok) throw new Error(result.message);
+      const credential: OidcCredentials = {
+        type: "oauth",
+        access: result.token.access,
+        refresh: result.token.refresh,
+        expires: result.token.expires,
+        baseUrl,
+        flow: OIDC_FLOW,
+        issuer: discovery.issuer,
+        clientId: config.clientId,
+        tokenEndpoint: discovery.tokenEndpoint,
+        subject: result.token.subject,
+      };
+      return { ...credential, type: "oauth" };
+    },
+  );
+}
+
 async function startCliSso(
   baseUrl: string,
   signal?: AbortSignal,
@@ -1055,8 +1312,12 @@ async function loginWithPastedToken(
 }
 
 async function loginOAuth(interaction: AuthInteraction, definition: ProviderDefinition): Promise<OAuthCredential> {
-  const headers = resolveHeaders(definition);
+  const oidc = definition.oidc === undefined ? undefined : parseOidcConfig(definition.oidc);
   const baseUrl = await promptBaseUrl(interaction, definition);
+  // Direct OIDC sends nothing to the proxy. The IdP is never derived from it either: the proxy's own OAuth
+  // metadata describes LiteLLM's authorization server, not the IdP that signs the JWTs it accepts.
+  if (oidc) return loginOidc(interaction, baseUrl, oidc);
+  const headers = resolveHeaders(definition);
   const discovery = await discoverPkce(baseUrl, interaction.signal, headers);
   if (discovery) return loginPkce(interaction, baseUrl, discovery, headers);
   const cliSso = await startCliSso(baseUrl, interaction.signal, headers);
@@ -1085,6 +1346,24 @@ async function loginOAuth(interaction: AuthInteraction, definition: ProviderDefi
 // still-failing token endpoint (e.g. during an outage or rate limit).
 const pkceTransientRefreshBackoff = new Map<string, number>();
 
+async function refreshWithBackoff(
+  credentials: OAuthCredentials,
+  request: () => Promise<PkceTokenResult | OidcTokenResult>,
+): Promise<OAuthCredentials> {
+  if (Date.now() < credentials.expires) {
+    const backoffUntil = pkceTransientRefreshBackoff.get(credentials.refresh);
+    if (backoffUntil !== undefined && Date.now() < backoffUntil) return credentials;
+  }
+  const result = await request();
+  if (!result.ok && result.transient && Date.now() < credentials.expires) {
+    pkceTransientRefreshBackoff.set(credentials.refresh, Date.now() + PKCE_TRANSIENT_REFRESH_BACKOFF_MS);
+    return credentials;
+  }
+  pkceTransientRefreshBackoff.delete(credentials.refresh);
+  if (!result.ok) throw new Error(`${result.message}; run /login litellm again`);
+  return { ...credentials, ...result.token };
+}
+
 async function refreshLiteLLM(
   credentials: OAuthCredentials,
   definition: ProviderDefinition,
@@ -1092,46 +1371,61 @@ async function refreshLiteLLM(
 ): Promise<OAuthCredentials> {
   signal?.throwIfAborted();
   if (credentials.flow === PKCE_FLOW) {
+    const { baseUrl, clientId } = credentials;
     if (
-      typeof credentials.baseUrl !== "string" ||
-      !isAuthToken(credentials.clientId) ||
+      typeof baseUrl !== "string" ||
+      !isAuthToken(clientId) ||
       !isAuthToken(credentials.access) ||
       !isAuthToken(credentials.refresh) ||
       !Number.isSafeInteger(credentials.expires)
     ) {
       throw new Error("Invalid LiteLLM PKCE credential; run /login litellm again");
     }
-    if (Date.now() < credentials.expires) {
-      const backoffUntil = pkceTransientRefreshBackoff.get(credentials.refresh);
-      if (backoffUntil !== undefined && Date.now() < backoffUntil) return credentials;
+    return refreshWithBackoff(credentials, () => {
+      const root = requireCredentialRoot(normalizeBaseUrl(baseUrl, definition.allowInsecureHttp), definition.name);
+      canonicalIssuer(root);
+      const issuer = new URL(root);
+      const tokenEndpoint = sameOriginUrl(credentials.tokenEndpoint, issuer, "token endpoint");
+      const resource = sameOriginUrl(credentials.resource, issuer, "resource");
+      return requestPkceToken(
+        tokenEndpoint,
+        new URLSearchParams({
+          grant_type: "refresh_token",
+          refresh_token: credentials.refresh,
+          client_id: clientId,
+          resource,
+        }),
+        signal,
+        resolveHeaders(definition),
+        credentials.refresh,
+      );
+    });
+  }
+  // Must precede the `!command` fallthrough below: an IdP refresh token is never executed.
+  if (credentials.flow === OIDC_FLOW) {
+    const { refresh, issuer, clientId, subject } = credentials;
+    const tokenEndpoint = httpsUrl(credentials.tokenEndpoint, { allowQuery: true });
+    if (!isAuthToken(refresh))
+      throw new Error("LiteLLM OIDC credential has no refresh token; run /login litellm again");
+    if (
+      !tokenEndpoint ||
+      typeof issuer !== "string" ||
+      !isAuthToken(clientId) ||
+      typeof subject !== "string" ||
+      !subject ||
+      !Number.isSafeInteger(credentials.expires)
+    ) {
+      throw new Error("Invalid LiteLLM OIDC credential; run /login litellm again");
     }
-    const baseUrl = requireCredentialRoot(
-      normalizeBaseUrl(credentials.baseUrl, definition.allowInsecureHttp),
-      definition.name,
-    );
-    canonicalIssuer(baseUrl);
-    const issuer = new URL(baseUrl);
-    const tokenEndpoint = sameOriginUrl(credentials.tokenEndpoint, issuer, "token endpoint");
-    const resource = sameOriginUrl(credentials.resource, issuer, "resource");
-    const result = await requestPkceToken(
-      tokenEndpoint,
-      new URLSearchParams({
-        grant_type: "refresh_token",
-        refresh_token: credentials.refresh,
-        client_id: credentials.clientId,
-        resource,
-      }),
-      signal,
-      resolveHeaders(definition),
-      credentials.refresh,
+    return refreshWithBackoff(credentials, () =>
+      requestOidcToken(
+        tokenEndpoint,
+        new URLSearchParams({ grant_type: "refresh_token", refresh_token: refresh, client_id: clientId }),
+        signal,
+        { issuer, clientId, subject },
+        refresh,
+      ),
     );
-    if (!result.ok && result.transient && Date.now() < credentials.expires) {
-      pkceTransientRefreshBackoff.set(credentials.refresh, Date.now() + PKCE_TRANSIENT_REFRESH_BACKOFF_MS);
-      return credentials;
-    }
-    pkceTransientRefreshBackoff.delete(credentials.refresh);
-    if (!result.ok) throw new Error(`${result.message}; run /login litellm again`);
-    return { ...credentials, ...result.token };
   }
   if (!credentials.refresh.startsWith("!")) {
     if (credentials.expires < PERMANENT_TOKEN_EXPIRES_AT) {
@@ -1568,6 +1862,7 @@ export default async function (pi: ExtensionAPI): Promise<void> {
   const definitions = getProviderDefinitions(settings);
   const skillsEnabled = isFeatureEnabled(settings, "skills");
   const mcpEnabled = isFeatureEnabled(settings, "mcp");
+  const mcpExposure = parseMcpExposurePolicy(settings?.mcp);
   const providerNames = new Set(definitions.map((definition) => definition.name));
   const oauthRuntimeRoots = new Map<string, { apiKey: string; root: string }>();
   let mcpUI: ExtensionContext["ui"] | undefined;
@@ -1593,6 +1888,7 @@ export default async function (pi: ExtensionAPI): Promise<void> {
     sessionStarted = false;
     pendingMcpMessages.clear();
   });
+  for (const warning of mcpExposure.warnings) notifyMcp(`LiteLLM MCP: ${warning}`);
 
   function discoveryDisabledReason(): string | null {
     if (isOffline()) return `${ENV_OFFLINE}=1`;
@@ -1645,6 +1941,9 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 
   // Each provider owns its own MCP catalog, keyed by provider name.
   const registeredMcpIdentities = new Map<string, string>();
+  // Generated tool names each provider registered in its last successful pass; Pi cannot unregister,
+  // so names that vanish from the catalog are withdrawn by re-registering them hidden.
+  const registeredMcpToolNames = new Map<string, Set<string>>();
   let mcpSeeded = false;
   const mcpRegistrations = new Map<string, Promise<void>>();
   // Pi's registerTool throws only from assertActive(), whose staleness flag is set with `??=` and
@@ -1810,6 +2109,65 @@ export default async function (pi: ExtensionAPI): Promise<void> {
     });
   }
 
+  // Pi before 0.99 reports no `exposure` and declares every registered tool.
+  function piHonorsToolExposure(): boolean {
+    return pi.getAllTools().some((tool) => "exposure" in tool);
+  }
+
+  const WITHDRAWN_MCP_MESSAGE = "This LiteLLM MCP tool is no longer offered by the proxy.";
+  function withdrawMcpTools(names: ReadonlySet<string>): void {
+    for (const name of names) {
+      // Older Pi ignores `exposure`, so the stub stays callable and must fail without calling the proxy.
+      pi.registerTool({
+        name,
+        label: "Withdrawn LiteLLM MCP tool",
+        description: WITHDRAWN_MCP_MESSAGE,
+        parameters: Type.Object({}),
+        exposure: "hidden",
+        async execute() {
+          throw new Error(WITHDRAWN_MCP_MESSAGE);
+        },
+      });
+    }
+    if (piHonorsToolExposure()) pi.setActiveTools(pi.getActiveTools().filter((name) => !names.has(name)));
+  }
+
+  let warnedMcpExposureIgnored = false;
+  let warnedMcpUnreachable = false;
+  // As in Pi's own MCP extension: tools that are not declared to the model are reached through
+  // codemode (`codemode`, `codemode-deferred`) or tool_search (`deferred`). Activate the one the
+  // configured exposure needs, and say so once when neither is available.
+  function activateMcpDiscoveryTools(exposures: readonly McpExposure[]): void {
+    const needsCodemode = exposures.includes("codemode") || exposures.includes("codemode-deferred");
+    const needsToolSearch = exposures.includes("deferred");
+    if (!needsCodemode && !needsToolSearch) return;
+    const tools = pi.getAllTools();
+    // Without exposure support every registered tool is declared, so nothing needs activating.
+    if (!piHonorsToolExposure()) {
+      if (warnedMcpExposureIgnored) return;
+      warnedMcpExposureIgnored = true;
+      notifyMcp(
+        "LiteLLM MCP: this Pi version ignores litellm.mcp.exposure; MCP tools are declared directly to the model.",
+      );
+      return;
+    }
+    const hasCodemode = tools.some((tool) => tool.name === "codemode");
+    const hasToolSearch = tools.some((tool) => tool.name === "tool_search");
+    const active = pi.getActiveTools();
+    const activate = [
+      ...(needsCodemode && hasCodemode && !active.includes("codemode") ? ["codemode"] : []),
+      ...(needsToolSearch && hasToolSearch && !active.includes("tool_search") ? ["tool_search"] : []),
+    ];
+    if (activate.length > 0) pi.setActiveTools([...active, ...activate]);
+    const reachable = [...active, ...activate];
+    if ((hasCodemode && reachable.includes("codemode")) || (hasToolSearch && reachable.includes("tool_search"))) return;
+    if (warnedMcpUnreachable) return;
+    warnedMcpUnreachable = true;
+    notifyMcp(
+      "LiteLLM MCP: tools with codemode or deferred exposure are reachable only through Pi's codemode or tool_search tool, and neither is available; they cannot be called.",
+    );
+  }
+
   async function registerMcpTools(
     definition: ProviderDefinition,
     auth: McpRuntimeAuth,
@@ -1850,6 +2208,7 @@ export default async function (pi: ExtensionAPI): Promise<void> {
           signal,
           notify,
           namespace,
+          mcpExposure.policy,
         );
         signal?.throwIfAborted();
         if (loginGeneration !== mcpLoginGeneration(definition.name)) return;
@@ -1872,6 +2231,27 @@ export default async function (pi: ExtensionAPI): Promise<void> {
         }
         reportMcpRegistrationSuccess(notify);
         reportMcpPartialDiscovery(report.partialFailure, registeredNames, notify);
+        const previousNames = registeredMcpToolNames.get(definition.name) ?? new Set<string>();
+        const currentNames = new Set(registeredNames);
+        // A partial or empty catalog may omit live tools, so it withdraws nothing and keeps the old names.
+        const settled = definitions.length > 0 && !report.partialFailure;
+        const withdrawn = settled ? new Set([...previousNames].filter((n) => !currentNames.has(n))) : new Set<string>();
+        try {
+          if (withdrawn.size > 0) withdrawMcpTools(withdrawn);
+        } catch (error) {
+          if (signal?.aborted) throw signal.reason;
+          mcpRegistrationFatal = true;
+          reportMcpRegistrationFatal(registeredNames.length, definitions.length, error, notify);
+          return;
+        }
+        if (definitions.length > 0) {
+          if (!settled) for (const name of previousNames) currentNames.add(name);
+          registeredMcpToolNames.set(definition.name, currentNames);
+        }
+        if (withdrawn.size > 0 && isVerboseDiscovery()) {
+          notifyMcp(`${label}: withdrew ${withdrawn.size} MCP tools no longer offered by the proxy.`, "info");
+        }
+        activateMcpDiscoveryTools(report.exposures);
         if (isVerboseDiscovery()) {
           notifyMcp(
             `${label}: registered ${registeredNames.length} of ${definitions.length} prepared MCP tools ` +