@danmademe/pi-provider-litellm 0.2.0 → 0.3.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@danmademe/pi-provider-litellm",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "description": "Pi agent extension for LiteLLM proxy auto-discovery and model configuration",
   "type": "module",
   "main": "./src/index.ts",

package/src/gcloud-token.ts

@@ -0,0 +1,169 @@
+import { existsSync, readFileSync } from 'fs'
+import { join } from 'path'
+
+let cachedToken: string | null = null
+let cachedAt: number = 0
+let inflight: Promise<string | null> | null = null
+export const CACHE_TTL = 50 * 60 * 1000 // 50 minutes in ms
+
+interface AuthorizedUserCredentials {
+  type: 'authorized_user'
+  client_id: string
+  client_secret: string
+  refresh_token: string
+  account?: string
+  universe_domain?: string
+}
+
+interface ServiceAccountCredentials {
+  type: 'service_account'
+}
+
+type GoogleCredentials = AuthorizedUserCredentials | ServiceAccountCredentials
+
+const ADC_FILENAME = 'application_default_credentials.json'
+
+function getAdcPath(): string | null {
+  // 1. GOOGLE_APPLICATION_CREDENTIALS env var (all platforms)
+  const envPath = typeof process !== 'undefined' ? process.env.GOOGLE_APPLICATION_CREDENTIALS : undefined
+  if (envPath) {
+    return envPath
+  }
+
+  // 2. Default ADC locations (Google's official search order)
+  const candidates: string[] = []
+
+  // Linux / macOS: ~/.config/gcloud/
+  const home = typeof process !== 'undefined' ? process.env.HOME : undefined
+  if (home) {
+    candidates.push(join(home, '.config', 'gcloud', ADC_FILENAME))
+  }
+
+  // Windows: %APPDATA%/gcloud/
+  const appData = typeof process !== 'undefined' ? process.env.APPDATA : undefined
+  if (appData) {
+    candidates.push(join(appData, 'gcloud', ADC_FILENAME))
+  }
+
+  for (const path of candidates) {
+    if (existsSync(path)) {
+      return path
+    }
+  }
+
+  return null
+}
+
+function readCredentials(path: string): GoogleCredentials | null {
+  try {
+    const content = readFileSync(path, 'utf-8')
+    return JSON.parse(content) as GoogleCredentials
+  } catch {
+    return null
+  }
+}
+
+async function exchangeRefreshToken(credentials: AuthorizedUserCredentials): Promise<string | null> {
+  const body = new URLSearchParams({
+    grant_type: 'refresh_token',
+    client_id: credentials.client_id,
+    client_secret: credentials.client_secret,
+    refresh_token: credentials.refresh_token,
+  }).toString()
+
+  try {
+    const response = await fetch('https://oauth2.googleapis.com/token', {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
+      body,
+      signal: AbortSignal.timeout(10_000),
+    })
+
+    if (!response.ok) {
+      const text = await response.text()
+      console.warn(`[pi-provider-litellm] Token exchange failed (${response.status}): ${text}`)
+      return null
+    }
+
+    const data = await response.json()
+    return data.access_token || null
+  } catch (error) {
+    console.warn(`[pi-provider-litellm] Token exchange failed: ${error}`)
+    return null
+  }
+}
+
+/**
+ * Gets a Google OAuth access token from the ADC JSON file, cached with a 50-minute TTL.
+ * Concurrent calls share one in-flight request (request coalescing).
+ * Returns null if credentials are not available or the token cannot be fetched.
+ * Logs a warning on failure.
+ */
+export async function getGcloudToken(): Promise<string | null> {
+  // Return cached token if still valid
+  if (cachedToken && (Date.now() - cachedAt) < CACHE_TTL) {
+    return cachedToken
+  }
+
+  // Coalesce concurrent calls: reuse the in-flight promise if one exists
+  if (inflight) {
+    return inflight
+  }
+
+  inflight = (async () => {
+    try {
+      const adcPath = getAdcPath()
+      if (!adcPath) {
+        console.warn(
+          '[pi-provider-litellm] No Google ADC file found. Set GOOGLE_APPLICATION_CREDENTIALS or run `gcloud auth application-default login`.',
+        )
+        return null
+      }
+
+      const credentials = readCredentials(adcPath)
+      if (!credentials) {
+        console.warn(`[pi-provider-litellm] Failed to read ADC file: ${adcPath}`)
+        return null
+      }
+
+      if (credentials.type === 'authorized_user') {
+        const token = await exchangeRefreshToken(credentials)
+        if (token) {
+          cachedToken = token
+          cachedAt = Date.now()
+        }
+        return token
+      }
+
+      if (credentials.type === 'service_account') {
+        console.warn('[pi-provider-litellm] Service account credentials are not yet supported. Use an authorized_user credential or set GOOGLE_APPLICATION_CREDENTIALS to an authorized_user JSON file.')
+        return null
+      }
+
+      // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
+      console.warn(`[pi-provider-litellm] Unknown credential type: ${(credentials as { type: string }).type}`)
+      return null
+    } finally {
+      inflight = null
+    }
+  })()
+
+  return inflight
+}
+
+/**
+ * Pre-warms the token cache by fetching a token in the background.
+ * Safe to call without awaiting — errors are swallowed since getGcloudToken logs them.
+ */
+export function warmGcloudToken(): void {
+  void getGcloudToken()
+}
+
+/**
+ * Resets the token cache. Exported for testing purposes.
+ */
+export function resetTokenCache(): void {
+  cachedToken = null
+  cachedAt = 0
+  inflight = null
+}

package/src/index.ts

@@ -1,17 +1,37 @@
 import type { ExtensionAPI, BeforeAgentStartEvent, BeforeAgentStartEventResult } from '@earendil-works/pi-coding-agent'
 import { resolvePluginConfig, discoverModels, discoverMcpTools, listSkills, buildProviderConfig } from './litellm-api.js'
 import { createMcpToolDefinitions, createSkillToolDefinitions, createSkillsInjector } from './tools.js'
+import { getGcloudToken } from './gcloud-token.js'
+import { loadModelCache, saveModelCache } from './model-cache.js'
 import type { LiteLLMModelInfo, McpTool, PluginConfig } from './types.js'
 
+const LOG = '[pi-provider-litellm]'
+
 export default async function (pi: ExtensionAPI): Promise<void> {
   const config = resolvePluginConfig()
   if (!config) {
+    console.warn(`${LOG} No config found — set LITELLM_URL and LITELLM_KEY (or LITELLM_GCLOUD_TOKEN_AUTH=1)`)
     return
   }
 
-  await discoverAndRegister(pi, config)
+  const isGcloudAuth = !!(process.env.LITELLM_GCLOUD_TOKEN_AUTH &&
+    process.env.LITELLM_GCLOUD_TOKEN_AUTH !== '' &&
+    process.env.LITELLM_GCLOUD_TOKEN_AUTH !== '0')
+
+  // When gcloud token auth is enabled, fetch a live token instead of using the static apiKey
+  const getToken = async (): Promise<string> => {
+    if (isGcloudAuth) {
+      return (await getGcloudToken()) ?? ''
+    }
+    return config.apiKey
+  }
 
-  const injector = createSkillsInjector(config, config.apiKey)
+  // Await discovery so PI blocks until models are registered before resolving
+  // model patterns. Cache is loaded at the top of discoverAndRegister so the
+  // first call returns quickly on subsequent startups.
+  await discoverAndRegister(pi, config, getToken)
+
+  const injector = createSkillsInjector(config, getToken)
   const setupCompleteSessions = new Set<string>()
   pi.on('before_agent_start', async (event: BeforeAgentStartEvent, ctx): Promise<BeforeAgentStartEventResult> => {
     const sessionId = ctx.sessionManager.getSessionFile()
@@ -26,7 +46,7 @@ export default async function (pi: ExtensionAPI): Promise<void> {
   pi.on('session_start', async (_event, _ctx) => {
     setupCompleteSessions.clear()
     injector.clearCache()
-    await discoverAndRegister(pi, config)
+    await discoverAndRegister(pi, config, getToken)
   })
 
   pi.on('session_shutdown', async (_event, _ctx) => {
@@ -34,11 +54,16 @@ export default async function (pi: ExtensionAPI): Promise<void> {
   })
 }
 
-export async function discoverAndRegister(pi: ExtensionAPI, config: PluginConfig): Promise<void> {
-  try {
-    pi.unregisterProvider(config.providerId)
-  } catch {
-    // Provider not yet registered
+export async function discoverAndRegister(pi: ExtensionAPI, config: PluginConfig, getToken: () => Promise<string>): Promise<void> {
+  // Register from cache before live discovery so models are visible immediately.
+  // Must await getToken() — PI rejects an empty apiKey. In gcloud mode this is
+  // one OAuth exchange (~500ms), but subsequent starts within the 50-min TTL
+  // return the cached token instantly. On first-ever run there is no cache file,
+  // so this block is skipped entirely.
+  const cached = loadModelCache(config.providerId)
+  if (cached) {
+    const token = await getToken()
+    pi.registerProvider(config.providerId, buildProviderConfig(config.url, token, cached))
   }
 
   const DISCOVERY_TIMEOUT_MS = 30_000
@@ -51,11 +76,12 @@ export async function discoverAndRegister(pi: ExtensionAPI, config: PluginConfig
   })
 
   try {
+    const token = await getToken()
     const results = await Promise.race([
       Promise.allSettled([
-        discoverModels(config, config.apiKey),
-        discoverMcpTools(config, config.apiKey),
-        listSkills(config, config.apiKey),
+        discoverModels(config, token),
+        discoverMcpTools(config, token),
+        listSkills(config, token),
       ]),
       timeoutPromise,
     ])
@@ -71,19 +97,30 @@ export async function discoverAndRegister(pi: ExtensionAPI, config: PluginConfig
     mcpResult = { status: 'rejected', reason: error as Error }
   }
 
-  if (modelsResult.status === 'fulfilled' && Object.keys(modelsResult.value).length > 0) {
-    const providerConfig = buildProviderConfig(config.url, config.apiKey, modelsResult.value)
-    pi.registerProvider(config.providerId, providerConfig)
+  if (modelsResult.status === 'fulfilled') {
+    const modelCount = Object.keys(modelsResult.value).length
+    if (modelCount > 0) {
+      saveModelCache(config.providerId, modelsResult.value)
+      const token = await getToken()
+      const providerConfig = buildProviderConfig(config.url, token, modelsResult.value)
+      pi.registerProvider(config.providerId, providerConfig)
+    } else {
+      console.warn(`${LOG} No models discovered — check LiteLLM /health endpoint`)
+    }
+  } else {
+    console.error(`${LOG} Model discovery error: ${modelsResult.reason}`)
   }
 
   if (mcpResult.status === 'fulfilled') {
-    const mcpTools = createMcpToolDefinitions(config, config.apiKey, mcpResult.value)
+    const mcpTools = createMcpToolDefinitions(config, getToken, mcpResult.value)
     for (const tool of mcpTools) {
       pi.registerTool(tool)
     }
+  } else {
+    console.warn(`${LOG} MCP tool discovery failed: ${mcpResult.reason}`)
   }
 
-  const skillTools = createSkillToolDefinitions(config, config.apiKey)
+  const skillTools = createSkillToolDefinitions(config, getToken)
   for (const tool of skillTools) {
     pi.registerTool(tool)
   }

package/src/litellm-api.ts

@@ -318,11 +318,17 @@ export function resolvePluginConfig(): PluginConfig | null {
   // Check env vars first
   const envUrl = process.env.LITELLM_URL
   const envKey = process.env.LITELLM_KEY
+  const envGcloudAuth = process.env.LITELLM_GCLOUD_TOKEN_AUTH
 
   if (envUrl && envKey) {
     return { url: envUrl, apiKey: envKey, providerId: process.env.LITELLM_PROVIDER_ID ?? 'litellm' }
   }
 
+  // Allow missing LITELLM_KEY when gcloud token auth is enabled
+  if (envUrl && envGcloudAuth && envGcloudAuth !== '' && envGcloudAuth !== '0') {
+    return { url: envUrl, apiKey: envKey ?? '', providerId: process.env.LITELLM_PROVIDER_ID ?? 'litellm' }
+  }
+
   // Check settings.json
   try {
     const settingsPath = path.join(os.homedir(), '.pi', 'agent', 'settings.json')

package/src/model-cache.ts

@@ -0,0 +1,51 @@
+import os from 'node:os'
+import path from 'node:path'
+import fs from 'node:fs'
+import type { LiteLLMModelInfo } from './types.js'
+
+const CACHE_FILENAME = 'pi-provider-litellm-cache.json'
+
+interface ModelCache {
+  savedAt: number
+  providerId: string
+  models: Record<string, LiteLLMModelInfo>
+}
+
+function getCachePath(): string {
+  return path.join(os.homedir(), '.pi', 'agent', CACHE_FILENAME)
+}
+
+/**
+ * Loads the model cache from disk. Returns null if the file does not exist or
+ * cannot be parsed.
+ */
+export function loadModelCache(providerId: string): Record<string, LiteLLMModelInfo> | null {
+  try {
+    const raw = fs.readFileSync(getCachePath(), 'utf-8')
+    const cache = JSON.parse(raw) as ModelCache
+
+    // Ignore cache for a different provider
+    if (cache.providerId !== providerId) return null
+    if (!cache.models || typeof cache.models !== 'object') return null
+
+    return cache.models
+  } catch {
+    return null
+  }
+}
+
+/**
+ * Saves the discovered models to the cache file on disk.
+ */
+export function saveModelCache(providerId: string, models: Record<string, LiteLLMModelInfo>): void {
+  try {
+    const cache: ModelCache = {
+      savedAt: Date.now(),
+      providerId,
+      models,
+    }
+    fs.writeFileSync(getCachePath(), JSON.stringify(cache, null, 2), 'utf-8')
+  } catch {
+    // Non-fatal — discovery still succeeded, cache will be written next time
+  }
+}

package/src/tools.ts

@@ -114,7 +114,7 @@ export function buildTypeBoxSchema(inputSchema: Record<string, unknown>): TSchem
 
 export function createMcpToolDefinitions(
   config: PluginConfig,
-  token: string,
+  getToken: () => Promise<string>,
   mcpTools: McpTool[],
 ): ToolDefinition[] {
   return mcpTools.map((tool) => {
@@ -141,6 +141,7 @@ export function createMcpToolDefinitions(
           ? (typeof params.args === 'string' ? (() => { try { return JSON.parse(params.args) } catch { return params.args } })() : params.args)
           : params
 
+        const token = await getToken()
         const result = await executeMcpTool(config, token, server, toolName, args as Record<string, unknown>)
         return {
           content: [{ type: 'text', text: result }],
@@ -153,7 +154,7 @@ export function createMcpToolDefinitions(
 
 export function createSkillToolDefinitions(
   config: PluginConfig,
-  token: string,
+  getToken: () => Promise<string>,
 ): ToolDefinition[] {
   return [
     {
@@ -168,6 +169,7 @@ export function createSkillToolDefinitions(
         _onUpdate: AgentToolUpdateCallback<unknown> | undefined,
         _ctx: ExtensionContext,
       ): Promise<AgentToolResult<undefined>> {
+        const token = await getToken()
         const skills = await listSkills(config, token)
         if (!skills.length) {
           return { content: [{ type: 'text', text: 'No skills found.' }], details: undefined }
@@ -195,6 +197,7 @@ export function createSkillToolDefinitions(
         _onUpdate: AgentToolUpdateCallback<unknown> | undefined,
         _ctx: ExtensionContext,
       ): Promise<AgentToolResult<undefined>> {
+        const token = await getToken()
         const skills = await listSkills(config, token)
         const skill = skills.find((s) => s.name === params.name)
         if (!skill) {
@@ -225,6 +228,7 @@ export function createSkillToolDefinitions(
         _onUpdate: AgentToolUpdateCallback<unknown> | undefined,
         _ctx: ExtensionContext,
       ): Promise<AgentToolResult<undefined>> {
+        const token = await getToken()
         const result = await registerSkill(
           config,
           token,
@@ -249,6 +253,7 @@ export function createSkillToolDefinitions(
         _onUpdate: AgentToolUpdateCallback<unknown> | undefined,
         _ctx: ExtensionContext,
       ): Promise<AgentToolResult<undefined>> {
+        const token = await getToken()
         const result = await enableSkill(config, token, params.name)
         return { content: [{ type: 'text', text: result }], details: undefined }
       },
@@ -265,6 +270,7 @@ export function createSkillToolDefinitions(
         _onUpdate: AgentToolUpdateCallback<unknown> | undefined,
         _ctx: ExtensionContext,
       ): Promise<AgentToolResult<undefined>> {
+        const token = await getToken()
         const result = await disableSkill(config, token, params.name)
         return { content: [{ type: 'text', text: result }], details: undefined }
       },
@@ -279,7 +285,7 @@ export interface SkillsInjector {
 
 export function createSkillsInjector(
   config: PluginConfig,
-  token: string,
+  getToken: () => Promise<string>,
 ): SkillsInjector {
   let cache: { skills: Skill[]; timestamp: number } | null = null
   const TTL = 60_000 // 60 seconds
@@ -289,6 +295,7 @@ export function createSkillsInjector(
     if (cache && now - cache.timestamp < TTL) {
       return cache.skills
     }
+    const token = await getToken()
     const skills = await listSkills(config, token)
     cache = { skills, timestamp: now }
     return skills