pi-provider-litellm 2.0.6 → 2.1.0

package/README.md

@@ -46,17 +46,17 @@ You'll be prompted for the base URL and API key. Credentials are persisted to `~
 
 #### Enterprise SSO login
 
-If your LiteLLM proxy requires SSO/OAuth authentication (enterprise deployments), you can authenticate via a browser SSO flow and optionally pair the resulting JWT with a stable virtual key:
+If your LiteLLM proxy requires SSO/OAuth authentication (enterprise deployments), you can authenticate through LiteLLM's CLI SSO flow:
 
 1. Run `/login litellm` inside pi and select `Sign in with LiteLLM SSO`
 2. Enter the proxy URL
-3. Your default browser opens the LiteLLM SSO login URL (e.g. `https://litellm.your-domain.com/sso/key/generate`) automatically — the URL is also displayed in case it can't be opened. Authenticate via SSO
-4. Copy your token from the LiteLLM UI and paste it at the prompt (copying a full `Bearer ...` header value is fine — the prefix is stripped automatically)
-5. When prompted to generate a virtual key, press Enter to accept (recommended) or enter `n` to use the JWT directly
+3. Pi displays a verification code and opens the LiteLLM SSO login URL automatically
+4. Authenticate in the browser and confirm the displayed code
+5. If your account belongs to multiple teams, select the team for this session in Pi
 
-When you generate a virtual key, the resulting `sk-...` key is stored as your credential and used for all API requests. If the proxy's key policy attaches an expiry to the generated key, Pi will prompt you to re-authenticate when it nears expiry; otherwise the key is treated as permanent until revoked in LiteLLM.
+Pi polls the proxy until authentication completes and stores the returned short-lived CLI token. Newer proxies advertise the token lifetime; for older proxies, Pi uses `LITELLM_CLI_JWT_EXPIRATION_HOURS` or LiteLLM's 24-hour default. On LiteLLM versions that gate CLI SSO, start the proxy with `EXPERIMENTAL_UI_LOGIN=True`.
 
-When using a JWT directly, the extension reads its `exp` claim and Pi will prompt you to re-authenticate when the token nears expiry. Run `/login litellm` again to refresh.
+Older proxies without `/sso/cli/start` fall back to the legacy browser flow: copy the token from the LiteLLM UI, paste it into Pi, and optionally exchange it for a virtual key. Pi reads JWT expiry claims and prompts you to run `/login litellm` again when re-authentication is required.
 
 ### Option B — environment variables
 
@@ -155,6 +155,7 @@ Treat the configured LiteLLM proxy as trusted: Skills can add instructions to th
 | `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 |
 | `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 and MCP discovery (login, refresh, startup); discovery is silent by default |
 | `LITELLM_MODELS_DEV` | enabled | Set to `0` to disable models.dev metadata enrichment, including its cache and network request; `/v1/models` still uses Pi catalog metadata and defaults |
 
@@ -242,6 +243,8 @@ Opening `/model` refreshes configured provider catalogs in the background using
 | Discovery times out | Increase `LITELLM_DISCOVERY_TIMEOUT_MS` or set `LITELLM_OFFLINE=1` to fall back on cached models |
 | `401 Token expired` | Set `LITELLM_API_KEY_HELPER`. |
 | No models with gcloud auth | Verify `gcloud auth application-default login` has been run or set `GOOGLE_APPLICATION_CREDENTIALS` to an `authorized_user` ADC file |
+| Enterprise SSO waits for token insertion | The proxy returned 404/405 for `/sso/cli/start`, so Pi used the legacy flow — upgrade LiteLLM or paste the UI token |
+| Enterprise CLI SSO start/poll fails | Check the proxy logs and verify `/sso/cli/start` and `/sso/cli/poll/{login_id}` are reachable; only 404/405 falls back to legacy login |
 | Enterprise SSO login shows "virtual key generation failed" | The LiteLLM instance may lack a database (`/key/generate` requires one), your user account may lack key-generation permission, or the request timed out; the JWT is used directly as a fallback |
 | Enterprise SSO token prompt fails with "SSO token is required" | The token field was left empty — paste the token copied from the LiteLLM UI |
 | MCP tools not showing | Verify the proxy exposes `/mcp-rest/tools/list` and open `/model` after fixing the proxy |

package/dist/index.js

@@ -1,6 +1,7 @@
 import { execFileSync } from "node:child_process";
 import { readFile } from "node:fs/promises";
 import { join } from "node:path";
+import { setTimeout as delay } from "node:timers/promises";
 import { getAgentDir } from "@earendil-works/pi-coding-agent";
 import { setupLiteLLMCostTracking } from "./cost.js";
 import { discoverModels, isGpt55Model, isMoonshotModel, normalizeBaseUrl, shouldSuppressReasoningContent, } from "./discover.js";
@@ -16,11 +17,15 @@ const ENV_API_KEY = "LITELLM_API_KEY";
 const ENV_API_KEY_HELPER = "LITELLM_API_KEY_HELPER";
 const ENV_HEADERS = "LITELLM_HEADERS";
 const ENV_TIMEOUT = "LITELLM_DISCOVERY_TIMEOUT_MS";
+const ENV_CLI_JWT_EXPIRATION_HOURS = "LITELLM_CLI_JWT_EXPIRATION_HOURS";
 const ENV_OFFLINE = "LITELLM_OFFLINE";
 const ENV_VERBOSE_DISCOVERY = "LITELLM_VERBOSE_DISCOVERY";
 const ENV_MODELS_DEV = "LITELLM_MODELS_DEV";
 const DEFAULT_TIMEOUT_MS = 5000;
 const LOGIN_TIMEOUT_MS = 10_000;
+const CLI_SSO_POLL_INTERVAL_MS = 2_000;
+const CLI_SSO_EXPIRES_IN_SECONDS = 600;
+const DEFAULT_CLI_JWT_EXPIRATION_HOURS = 24;
 const MODELS_DEV_CACHE_FILENAME = "litellm-models-dev.json";
 const TOKEN_REFRESH_LEAD_MS = 5 * 60 * 1000;
 const PERMANENT_TOKEN_EXPIRES_AT = Number.MAX_SAFE_INTEGER;
@@ -125,10 +130,17 @@ function tokenExpiresAt(apiKey, opaqueFallback = PERMANENT_TOKEN_EXPIRES_AT) {
         return opaqueFallback;
     }
 }
-async function generateVirtualKey(baseUrl, userToken, signal, headers) {
-    const boundedSignal = signal
+function configuredCliJwtExpiresAt() {
+    const configured = Number(process.env[ENV_CLI_JWT_EXPIRATION_HOURS]);
+    const hours = Number.isFinite(configured) && configured > 0 ? configured : DEFAULT_CLI_JWT_EXPIRATION_HOURS;
+    return Date.now() + hours * 60 * 60 * 1_000;
+}
+function boundedLoginSignal(signal) {
+    return signal
         ? AbortSignal.any([signal, AbortSignal.timeout(LOGIN_TIMEOUT_MS)])
         : AbortSignal.timeout(LOGIN_TIMEOUT_MS);
+}
+async function generateVirtualKey(baseUrl, userToken, signal, headers) {
     const response = await fetch(`${baseUrl}/key/generate`, {
         method: "POST",
         headers: {
@@ -137,7 +149,7 @@ async function generateVirtualKey(baseUrl, userToken, signal, headers) {
             "Content-Type": "application/json",
         },
         body: JSON.stringify({}),
-        signal: boundedSignal,
+        signal: boundedLoginSignal(signal),
     });
     if (!response.ok) {
         throw new Error(`Virtual key generation failed (${response.status})`);
@@ -400,15 +412,108 @@ async function loginApiKey(interaction) {
         throw new Error("Both base URL and API key are required");
     return { type: "api_key", key, env: { [ENV_BASE_URL]: normalizeBaseUrl(rawBaseUrl) } };
 }
-async function loginOAuth(interaction, headers) {
-    const rawBaseUrl = (await interaction.prompt({
-        type: "text",
-        message: "Enter LiteLLM proxy URL (no trailing /v1):",
-        placeholder: "https://litellm.example.com",
-    })).trim();
-    if (!rawBaseUrl)
-        throw new Error("Base URL is required");
-    const baseUrl = normalizeBaseUrl(rawBaseUrl);
+async function startCliSso(baseUrl, signal, headers) {
+    const response = await fetch(`${baseUrl}/sso/cli/start`, {
+        method: "POST",
+        headers,
+        signal: boundedLoginSignal(signal),
+    });
+    if (response.status === 404 || response.status === 405)
+        return undefined;
+    if (!response.ok)
+        throw new Error(`LiteLLM CLI SSO start failed (HTTP ${response.status})`);
+    const data = (await response.json());
+    if (typeof data.login_id !== "string" ||
+        !data.login_id ||
+        typeof data.poll_secret !== "string" ||
+        !data.poll_secret ||
+        typeof data.user_code !== "string" ||
+        !data.user_code) {
+        throw new Error("LiteLLM CLI SSO start returned an invalid response");
+    }
+    return {
+        loginId: data.login_id,
+        pollSecret: data.poll_secret,
+        userCode: data.user_code,
+        expiresInSeconds: typeof data.expires_in === "number" && Number.isFinite(data.expires_in) && data.expires_in > 0
+            ? data.expires_in
+            : CLI_SSO_EXPIRES_IN_SECONDS,
+    };
+}
+async function waitForNextCliSsoPoll(interaction) {
+    try {
+        await delay(CLI_SSO_POLL_INTERVAL_MS, undefined, interaction.signal ? { signal: interaction.signal } : undefined);
+    }
+    catch (error) {
+        if (interaction.signal?.aborted)
+            throw interaction.signal.reason;
+        throw error;
+    }
+}
+async function pollCliSso(baseUrl, start, interaction, headers) {
+    const deadline = Date.now() + start.expiresInSeconds * 1_000;
+    const pollUrl = new URL(`${baseUrl}/sso/cli/poll/${encodeURIComponent(start.loginId)}`);
+    let teamId;
+    while (Date.now() < deadline) {
+        if (teamId)
+            pollUrl.searchParams.set("team_id", teamId);
+        let response;
+        try {
+            response = await fetch(pollUrl, {
+                headers: { ...headers, "x-litellm-cli-poll-secret": start.pollSecret },
+                signal: boundedLoginSignal(interaction.signal),
+            });
+        }
+        catch {
+            if (interaction.signal?.aborted)
+                throw interaction.signal.reason;
+            await waitForNextCliSsoPoll(interaction);
+            continue;
+        }
+        if (response.status === 429 || response.status >= 500) {
+            await waitForNextCliSsoPoll(interaction);
+            continue;
+        }
+        if (response.status === 400)
+            throw new Error("LiteLLM CLI SSO login expired or is invalid");
+        if (!response.ok)
+            throw new Error(`LiteLLM CLI SSO polling failed (HTTP ${response.status})`);
+        const data = (await response.json());
+        if (data.status === "ready" && typeof data.key === "string" && data.key) {
+            return {
+                access: data.key,
+                expiresInSeconds: typeof data.expires_in === "number" && Number.isFinite(data.expires_in) && data.expires_in > 0
+                    ? data.expires_in
+                    : undefined,
+            };
+        }
+        if (data.status === "ready" && data.requires_team_selection === true && !teamId) {
+            const details = Array.isArray(data.team_details) ? data.team_details : [];
+            const teams = details.flatMap((team) => {
+                if (!team || typeof team !== "object")
+                    return [];
+                const record = team;
+                const id = typeof record.team_id === "string" ? record.team_id : record.id;
+                const alias = record.team_alias;
+                return typeof id === "string" && id ? [{ id, label: typeof alias === "string" && alias ? alias : id }] : [];
+            });
+            if (teams.length === 0 && Array.isArray(data.teams))
+                teams.push(...data.teams.flatMap((id) => (typeof id === "string" && id ? [{ id, label: id }] : [])));
+            if (teams.length === 0)
+                throw new Error("LiteLLM CLI SSO requested team selection without any teams");
+            const selected = await interaction.prompt({ type: "select", message: "Select a LiteLLM team:", options: teams });
+            if (!teams.some((team) => team.id === selected))
+                throw new Error("Invalid LiteLLM team selection");
+            teamId = selected;
+            continue;
+        }
+        if (data.status !== "pending")
+            throw new Error("LiteLLM CLI SSO polling returned an invalid response");
+        await waitForNextCliSsoPoll(interaction);
+    }
+    throw new Error("LiteLLM CLI SSO login expired");
+}
+async function loginWithPastedToken(interaction, baseUrl, headers) {
     interaction.notify({
         type: "auth_url",
         url: `${baseUrl}/sso/key/generate`,
@@ -448,6 +553,35 @@ async function loginOAuth(interaction, headers) {
     }
     return { type: "oauth", access, refresh: "", expires, baseUrl };
 }
+async function loginOAuth(interaction, headers) {
+    const rawBaseUrl = (await interaction.prompt({
+        type: "text",
+        message: "Enter LiteLLM proxy URL (no trailing /v1):",
+        placeholder: "https://litellm.example.com",
+    })).trim();
+    if (!rawBaseUrl)
+        throw new Error("Base URL is required");
+    const baseUrl = normalizeBaseUrl(rawBaseUrl);
+    const cliSso = await startCliSso(baseUrl, interaction.signal, headers);
+    if (!cliSso)
+        return loginWithPastedToken(interaction, baseUrl, headers);
+    interaction.notify({
+        type: "device_code",
+        userCode: cliSso.userCode,
+        verificationUri: `${baseUrl}/sso/key/generate?${new URLSearchParams({ source: "litellm-cli", key: cliSso.loginId })}`,
+        expiresInSeconds: cliSso.expiresInSeconds,
+    });
+    const result = await pollCliSso(baseUrl, cliSso, interaction, headers);
+    return {
+        type: "oauth",
+        access: result.access,
+        refresh: "",
+        expires: result.expiresInSeconds
+            ? Date.now() + result.expiresInSeconds * 1_000
+            : tokenExpiresAt(result.access, configuredCliJwtExpiresAt()),
+        baseUrl,
+    };
+}
 async function refreshLiteLLM(credentials, _signal) {
     if (!credentials.refresh.startsWith("!")) {
         if (credentials.expires < PERMANENT_TOKEN_EXPIRES_AT) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-provider-litellm",
-  "version": "2.0.6",
+  "version": "2.1.0",
   "description": "LiteLLM proxy provider extension for Pi",
   "type": "module",
   "license": "MIT",