pi-requesty-discover 2.1.1 → 3.0.1

package/src/health-check.ts

@@ -1,7 +1,7 @@
 import type { ProviderModelConfig } from '@earendil-works/pi-coding-agent'
 import fs from 'node:fs'
 import path from 'node:path'
-import { getEnv, type Env } from './env'
+import { type Env, getEnv } from './env'
 import { formatModelsDiffSummary, type ModelsDiff } from './models-json'
 
 const HEALTH_CHECK_CONCURRENCY = 10
@@ -23,9 +23,11 @@ export type Provider = {
   apiKey: string
 }
 
+export type HealthCheckStatus = 'ok' | 'warning' | 'error'
+
 export type HealthCheckResult = {
   modelId: string
-  ok: boolean
+  status: HealthCheckStatus
   latencyMs: number
   error?: string
 }
@@ -38,6 +40,11 @@ export type HealthCheckProgress = {
   modelId: string
 }
 
+export type HealthCheckLogContext = {
+  providerId: string
+  bannedModels: string[]
+}
+
 export type HealthCheckOptions = {
   concurrency?: number
   timeoutMs?: number
@@ -87,53 +94,73 @@ export async function checkModels(
 }
 
 export function formatHealthSummary(results: HealthCheckResult[]): string {
-  const passed = results.filter(r => r.ok)
-  const failed = results.filter(r => !r.ok)
-
-  if (failed.length === 0) {
-    return `Health check: all ${passed.length} OK.`
-  }
+  const passed = results.filter(r => r.status === 'ok')
+  const warned = results.filter(r => r.status === 'warning')
+  const failed = results.filter(r => r.status === 'error')
 
-  const failedModels = failed.map(r => `- ${r.modelId}`).join('\n')
+  const parts = [`${passed.length} OK`]
+  if (warned.length > 0) parts.push(`${warned.length} warning${warned.length === 1 ? '' : 's'}`)
+  if (failed.length > 0) parts.push(`${failed.length} failed`)
 
-  return `Health check: ${passed.length} OK, ${failed.length} failed:\n${failedModels}\n`
+  return `Health check: ${parts.join(', ')}.`
 }
 
 export function writeHealthCheckLog(
   provider: Provider,
   results: HealthCheckResult[],
   diff: ModelsDiff,
+  context: HealthCheckLogContext,
   envConfig: Env = getEnv(),
 ): void {
-  const passed = results.filter(r => r.ok)
-  const failed = results.filter(r => !r.ok)
+  const passed = results.filter(r => r.status === 'ok')
+  const warned = results.filter(r => r.status === 'warning')
+  const failed = results.filter(r => r.status === 'error')
   const lines = [
     `Requesty health check log`,
     `Timestamp: ${new Date().toISOString()}`,
-    `Provider: ${envConfig.provider_id}`,
+    `Provider: ${context.providerId}`,
     `Base URL: ${provider.baseUrl}`,
     `Total: ${results.length}`,
     `Passed: ${passed.length}`,
+    `Warnings: ${warned.length}`,
     `Failed: ${failed.length}`,
+    `Banned: ${context.bannedModels.length}`,
+    ...context.bannedModels.map(id => `- ${id}`),
     '',
     formatModelsDiffSummary(diff),
     '',
   ]
 
-  if (failed.length === 0) {
+  if (warned.length === 0 && failed.length === 0) {
     lines.push('No failed models.')
   } else {
-    lines.push('Failed models:', '')
-    for (const result of failed) {
-      lines.push(
-        `Model: ${result.modelId}`,
-        `Latency: ${result.latencyMs}ms`,
-        'Error:',
-        result.error || 'Unknown error',
-        '',
-        '---',
-        '',
-      )
+    if (warned.length > 0) {
+      lines.push('Models with warnings:', '')
+      for (const result of warned) {
+        lines.push(
+          `Model: ${result.modelId}`,
+          `Latency: ${result.latencyMs}ms`,
+          'Error:',
+          result.error || 'Unknown error',
+          '',
+          '---',
+          '',
+        )
+      }
+    }
+    if (failed.length > 0) {
+      lines.push('Failed models:', '')
+      for (const result of failed) {
+        lines.push(
+          `Model: ${result.modelId}`,
+          `Latency: ${result.latencyMs}ms`,
+          'Error:',
+          result.error || 'Unknown error',
+          '',
+          '---',
+          '',
+        )
+      }
     }
   }
 
@@ -166,8 +193,10 @@ export async function postChatCompletion(
 
       if (!response.ok) {
         const text = await response.text().catch(() => '')
+        // rate limits are transient - a warning, not an error
+        const status: HealthCheckStatus = response.status === 429 ? 'warning' : 'error'
         return {
-          ok: false,
+          status,
           latencyMs,
           error: `HTTP ${response.status} ${response.statusText}${text ? `: ${text}` : ''}`,
         }
@@ -183,7 +212,7 @@ export async function postChatCompletion(
       }
       const attempts = attempt + 1
       return {
-        ok: false,
+        status: 'error',
         latencyMs: Date.now() - start,
         error: isTimeoutError(err)
           ? `Timed out after ${attempts} attempt(s); per-attempt timeout is ${healthCheckOptions.timeoutMs / 1000}s`
@@ -192,7 +221,7 @@ export async function postChatCompletion(
     }
   }
 
-  return { ok: false, latencyMs: Date.now() - start, error: 'Unknown error' }
+  return { status: 'error', latencyMs: Date.now() - start, error: 'Unknown error' }
 }
 
 async function verifyFirstStreamChunk(body: ReadableStream<Uint8Array>, start: number): Promise<ModelCheckResult> {
@@ -219,14 +248,14 @@ async function verifyFirstStreamChunk(body: ReadableStream<Uint8Array>, start: n
         try {
           const parsed = JSON.parse(payload) as { choices?: unknown }
           if (Array.isArray(parsed.choices) && parsed.choices.length > 0) {
-            return { ok: true, latencyMs: Date.now() - start }
+            return { status: 'ok', latencyMs: Date.now() - start }
           }
         } catch {
           // partial JSON or non-choices chunk — keep reading
         }
       }
     }
-    return { ok: false, latencyMs: Date.now() - start, error: 'Stream ended without content' }
+    return { status: 'error', latencyMs: Date.now() - start, error: 'Stream ended without content' }
   } finally {
     await reader.cancel()
   }
@@ -248,7 +277,7 @@ async function checkModel(
     options,
   )
 
-  if (!basicResult.ok || !model.reasoning || !checkReasoning) {
+  if (basicResult.status !== 'ok' || !model.reasoning || !checkReasoning) {
     return basicResult
   }
 
@@ -276,16 +305,16 @@ async function checkModel(
     options,
   )
 
-  if (!reasoningResult.ok) {
+  if (reasoningResult.status !== 'ok') {
     return {
-      ok: false,
+      status: 'error',
       latencyMs: reasoningResult.latencyMs,
       error: `Reasoning/tool check failed: ${reasoningResult.error}`,
     }
   }
 
   return {
-    ok: true,
+    status: 'ok',
     latencyMs: basicResult.latencyMs + reasoningResult.latencyMs,
   }
 }

package/src/index.ts

@@ -1,306 +1,141 @@
-import {
-  ExtensionAPI,
-  ExtensionCommandContext,
-  ExtensionContext,
-  ProviderModelConfig,
-} from '@earendil-works/pi-coding-agent'
+import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from '@earendil-works/pi-coding-agent'
 import { type Env, getEnv } from './env'
+import { type GetApiKey, getRequestyConfig } from './models-json'
+import { type ApiKeyInfo, fetchApiUsage } from './requesty-api'
+import { RequestyStatusLoader } from './ui/requesty-status-loader'
 import {
-  diffModels,
-  formatModelsDiffSummary,
-  getRequestyConfig,
-  type ModelsDiff,
-  updateModelsJson,
-} from './models-json'
-import { type ApiKeyInfo, discoverModels, fetchApiUsage } from './requesty-api'
-import { checkModels, formatHealthSummary, writeHealthCheckLog } from './health-check'
-import { RequestyStatusLoader } from './ui/requesty-status-loader.ts'
+  type DiscoveryEvaluation,
+  type DiscoveryUi,
+  type RefreshModelsRegistry,
+  evaluateDiscovery,
+  finalizeDiscovery,
+  formatDiscoveryFailure,
+  getArgumentCompletions,
+} from './discovery'
+import { type DiscoverySettings, readDiscoverySettings } from './settings'
+import { runCatching, runCatchingAsync, type Try } from './utils'
 
 const COMMAND_NAME = 'requesty-discover'
-const DRY_RUN_ARG = '--dry-run'
 export const USAGE_STATUS_KEY = 'requesty-usage'
 
-interface AutocompleteItem {
-  value: string
-  label: string
-  description: string
-}
-
-export type NotificationLevel = 'info' | 'warning' | 'error'
-
-export type Notifier = {
-  notify(message: string, level: NotificationLevel): void
-}
-
-export type Confirmer = {
-  confirm(title: string, message: string): Promise<boolean>
-}
-
-export type StatusReporter = {
-  set(message: string): void
-}
-
-type ModelsJsonData = ReturnType<typeof getRequestyConfig>['data']
-
-type DiscoveryEvaluation = {
-  dryRun: boolean
-  modelCount: number
-  failedCount: number
-  passing: ProviderModelConfig[]
-  diff: ModelsDiff
-  healthCheckSummary: string
-  logNote: string
-  data: ModelsJsonData
-}
-
-export type Try<T> = { ok: true; value: T } | { ok: false; error: unknown }
-
 // token suppresses stale writes when turns overlap
 let latestToken: object = {}
 
 // noinspection JSUnusedGlobalSymbols
 export default function (pi: ExtensionAPI) {
   const env = runCatching(() => getEnv())
+  if (!env.ok) {
+    console.error('env loading failed', env.error)
+    return
+  }
+  const settings = runCatching(() => readDiscoverySettings(env.value))
+  if (!settings.ok) {
+    console.error('settings loading failed', settings.error)
+    return
+  }
 
   pi.registerCommand(COMMAND_NAME, {
     description: 'Dynamically discover Requesty models, run health checks, and update the local models.json.',
     getArgumentCompletions,
     handler: async (args, ctx) => {
-      await runDiscoveryWorkflow(ctx, env, args)
+      await runDiscoveryWorkflow(ctx, settings.value, env.value, args)
     },
   })
 
   pi.on('turn_end', (_event, ctx) => {
-    void updateUsageStatus(ctx, env)
+    void updateUsageStatus(ctx, settings.value, env.value)
   })
 
   pi.on('session_start', (_event, ctx) => {
-    complainOnBrokenEnv(ctx, env)
-    void updateUsageStatus(ctx, env)
+    void updateUsageStatus(ctx, settings.value, env.value)
   })
 
   pi.on('model_select', (_event, ctx) => {
-    void updateUsageStatus(ctx, env)
+    void updateUsageStatus(ctx, settings.value, env.value)
   })
 }
 
-export async function runDiscoveryWorkflow(ctx: ExtensionCommandContext, env: Try<Env>, args: string) {
-  complainOnBrokenEnv(ctx, env)
-  if (!env.ok) return
-  // Interactive TUI command; no print/json/rpc path.
+export async function runDiscoveryWorkflow(
+  ctx: ExtensionCommandContext,
+  settings: DiscoverySettings,
+  env: Env,
+  args: string,
+) {
+  const ui = ctx.mode === 'tui' ? createTuiUi(ctx) : createConsoleUi()
   if (ctx.mode !== 'tui') {
+    return runSilentDiscoveryWorkflow(ctx, settings, env, args, ui)
+  } else {
+    return runInteractiveDiscoverWorkflow(ctx, settings, env, args, ui)
+  }
+}
+
+async function runSilentDiscoveryWorkflow(
+  ctx: ExtensionCommandContext,
+  settings: DiscoverySettings,
+  env: Env,
+  args: string,
+  ui: DiscoveryUi,
+): Promise<void> {
+  const evaluationResult = await runCatchingAsync(() =>
+    evaluateDiscovery(args, settings, env, ui, createGetApiKey(ctx)),
+  )
+  if (!evaluationResult.ok) {
+    ui.notify(formatDiscoveryFailure(evaluationResult.error), 'error')
     return
   }
 
-  const notifier = createUiNotifier(ctx)
-  const confirmer = createUiConfirmer(ctx)
+  await finalizeDiscovery(evaluationResult.value, settings, env, ui)
+}
 
-  // Phase A: progress UI only. Loader must close before confirm (Phase B).
-  // Errors are wrapped because ctx.ui.custom resolves via done() and does not reject.
+async function runInteractiveDiscoverWorkflow(
+  ctx: ExtensionCommandContext,
+  settings: DiscoverySettings,
+  env: Env,
+  args: string,
+  ui: DiscoveryUi,
+) {
   const evaluationResult: Try<DiscoveryEvaluation> = await runWithStatusUi(
     ctx,
     'Discovering models...',
-    async status => await runCatchingAsync(() => evaluateDiscovery(args, env.value, status)),
+    async statusUi =>
+      await runCatchingAsync(() => evaluateDiscovery(args, settings, env, statusUi, createGetApiKey(ctx))),
   )
 
   if (!evaluationResult.ok) {
-    notifier.notify(formatDiscoveryFailure(evaluationResult.error), 'error')
+    ui.notify(formatDiscoveryFailure(evaluationResult.error), 'error')
     return
   }
 
-  // Phases B + C: decide, optional write, final notify — outside the loader.
-  await finalizeDiscovery(evaluationResult.value, confirmer, notifier, env.value)
-}
-
-function formatDiscoveryFailure(error: unknown): string {
-  const detail = formatError(error)
-  return `Discovery failed: ${detail}`
+  await finalizeDiscovery(evaluationResult.value, settings, env, ui, createRefreshRegistry(ctx))
+  await updateUsageStatus(ctx, settings, env)
 }
 
 async function runWithStatusUi<T>(
   ctx: ExtensionCommandContext,
   initialMessage: string,
-  fn: (status: StatusReporter) => Promise<T>,
+  fn: (statusUi: DiscoveryUi) => Promise<T>,
 ): Promise<T> {
-  return ctx.ui.custom<T>((_tui, theme, _kb, done) => {
-    const loader = new RequestyStatusLoader(_tui, theme, initialMessage)
-    const status = createLoaderStatusReporter(loader)
+  return ctx.ui.custom<T>((tui, theme, _kb, done) => {
+    const loader = new RequestyStatusLoader(tui, theme, initialMessage)
     void Promise.resolve()
-      .then(() => fn(status))
+      .then(() => fn(createTuiUi(ctx, loader)))
       .then(done)
       .catch(done)
     return loader
   })
 }
 
-async function evaluateDiscovery(args: string, env: Env, status: StatusReporter): Promise<DiscoveryEvaluation> {
-  status.set('Discovering Requesty models...')
-  const dryRun = args.split(' ').includes(DRY_RUN_ARG)
-
-  const { data, provider, existingModelIds } = getRequestyConfig(env)
-  const models = await discoverModels(provider)
-  const modelsMap = new Map(models.map(m => [m.id, m]))
-
-  let diff: ModelsDiff
-  let failedCount = 0
-  let passing: ProviderModelConfig[]
-  let logNote = ''
-  let healthCheckSummary = ''
-
-  if (env.health_check_mode !== 'off') {
-    status.set(`Checking models 0/${models.length}...`)
-    const healthResults = await checkModels(provider, models, env.health_check_mode === 'full', {
-      onProgress: ({ completed, total }) => {
-        status.set(`Checking models ${completed}/${total}...`)
-      },
-    })
-    const sortedResults = healthResults.toSorted((a, b) => a.modelId.localeCompare(b.modelId))
-    failedCount = sortedResults.filter(r => !r.ok).length
-    passing = sortedResults.flatMap(r => {
-      const model = modelsMap.get(r.modelId)
-      return r.ok && model ? [model] : []
-    })
-    diff = diffModels(existingModelIds, passing)
-    healthCheckSummary = formatHealthSummary(sortedResults)
-    writeHealthCheckLog(provider, sortedResults, diff, env)
-    logNote = `Full health check log: ${env.health_check_log_path}\n`
-  } else {
-    passing = models
-    diff = diffModels(existingModelIds, passing)
-  }
-
-  return {
-    dryRun,
-    modelCount: models.length,
-    failedCount,
-    passing,
-    diff,
-    healthCheckSummary,
-    logNote,
-    data,
-  }
-}
-
-async function finalizeDiscovery(
-  evaluation: DiscoveryEvaluation,
-  confirmer: Confirmer,
-  notifier: Notifier,
-  env: Env,
-): Promise<void> {
-  const level = notificationLevel(evaluation)
-  const summary = buildDiscoverySummary(evaluation)
-
-  // Always surface the discovery result first (toast styling), then decide.
-  if (evaluation.dryRun) {
-    notifier.notify(
-      `${summary}
-Dry run: left models.json unchanged.`,
-      level,
-    )
-    return
-  }
-
-  if (evaluation.passing.length === 0) {
-    notifier.notify(
-      `${summary}
-Left models.json unchanged.`,
-      level,
-    )
-    return
-  }
-
-  notifier.notify(summary, level)
-  const { title, message } = buildConfirmPrompt(evaluation)
-  const shouldUpdate = await confirmer.confirm(title, message)
-  if (shouldUpdate) {
-    updateModelsJson(evaluation.data, evaluation.passing, env)
-    notifier.notify('Updated models.json. Run /reload to use the changes.', 'info')
-    return
-  }
-
-  notifier.notify('Left models.json unchanged.', 'info')
-}
-
-function buildDiscoverySummary(evaluation: DiscoveryEvaluation): string {
-  return [
-    `Discovered ${evaluation.modelCount} Requesty model(s).`,
-    evaluation.healthCheckSummary.trimEnd(),
-    formatModelsDiffSummary(evaluation.diff),
-    evaluation.logNote.trimEnd(),
-  ]
-    .filter(part => part.length > 0)
-    .join('\n')
-}
-
-function buildConfirmPrompt(evaluation: DiscoveryEvaluation): { title: string; message: string } {
-  const hasIdChanges = evaluation.diff.added.length > 0 || evaluation.diff.removed.length > 0
-  if (hasIdChanges) {
-    return {
-      title: `Write ${evaluation.passing.length} model(s)?`,
-      message: 'Update models.json with the discovery result above.',
-    }
-  }
-  return {
-    title: 'Refresh models.json?',
-    message: 'No model ID changes. Rewrite the file to refresh metadata anyway?',
-  }
-}
-
-function notificationLevel(evaluation: DiscoveryEvaluation): NotificationLevel {
-  if (evaluation.failedCount === 0) return 'info'
-  if (evaluation.failedCount < evaluation.modelCount) return 'warning'
-  return 'error'
-}
-
-function getArgumentCompletions(prefix: string): AutocompleteItem[] {
-  const options = [
-    {
-      value: DRY_RUN_ARG,
-      label: DRY_RUN_ARG,
-      description: 'Preview discovery without offering to write models.json',
-    },
-  ]
-  if (!prefix) return options
-  return options.filter(o => o.value.toLowerCase().startsWith(prefix.toLowerCase()))
-}
-
-function createUiNotifier(ctx: ExtensionContext): Notifier {
-  return {
-    notify(message: string, level?: NotificationLevel) {
-      const prefixedMessage = `${COMMAND_NAME}: ${message}`
-      ctx.ui.notify(prefixedMessage, level)
-    },
-  }
-}
-
-function createUiConfirmer(ctx: ExtensionContext): Confirmer {
-  return {
-    confirm(title, message) {
-      return ctx.ui.confirm(title, message)
-    },
-  }
-}
-
-function createLoaderStatusReporter(loader: RequestyStatusLoader): StatusReporter {
-  return {
-    set(message: string) {
-      loader.setMessage(message)
-    },
-  }
-}
-
-async function updateUsageStatus(ctx: ExtensionContext, env: Try<Env>): Promise<void> {
-  if (!env.ok) return
+async function updateUsageStatus(ctx: ExtensionContext, settings: DiscoverySettings, env: Env): Promise<void> {
   if (!ctx.hasUI) return // no footer to write to (print/json mode): skip the wasted fetch
   const token: object = {}
   latestToken = token
-  const shouldClear = ctx.model?.provider !== env.value.provider_id
+  const shouldClear = ctx.model?.provider !== settings.providerId
   try {
     if (shouldClear) {
       ctx.ui.setStatus(USAGE_STATUS_KEY, undefined)
       return
     }
-    const info = await fetchUsageStatus(env.value)
+    const info = await fetchUsageStatus(createGetApiKey(ctx), settings, env)
     if (latestToken !== token) return
     ctx.ui.setStatus(USAGE_STATUS_KEY, formatUsageStatus(info))
   } catch {
@@ -311,30 +146,26 @@ async function updateUsageStatus(ctx: ExtensionContext, env: Try<Env>): Promise<
 export function formatUsageStatus(info: ApiKeyInfo): string {
   const spend = `$${info.monthlySpend.toFixed(2)}`
   if (info.monthlyLimit <= 0) {
-    return `Requesty Usage (${info.name}): ${spend} (unlimited)`
+    return `${info.name}: ${spend} (unlimited)`
   }
   const limit = `$${info.monthlyLimit.toFixed(2)}`
   const percent = Math.floor((info.monthlySpend / info.monthlyLimit) * 100)
-  return `Requesty Usage (${info.name}): ${spend}/${limit} (${percent}%)`
+  return `${info.name}: ${spend}/${limit} (${percent}%)`
 }
 
 let lastFetched: { value: ApiKeyInfo; time: Date } | undefined
 
-async function fetchUsageStatus(env: Env): Promise<ApiKeyInfo> {
+async function fetchUsageStatus(getApiKey: GetApiKey, settings: DiscoverySettings, env: Env): Promise<ApiKeyInfo> {
   const now = new Date()
   if (lastFetched?.time && now.getTime() - lastFetched.time.getTime() < 2000) {
     return lastFetched.value
   }
-  const { provider } = getRequestyConfig(env)
+  const { provider } = await getRequestyConfig(getApiKey, settings, env)
   const value = await fetchApiUsage(provider)
   lastFetched = { value, time: new Date() }
   return value
 }
 
-function formatError(error: unknown) {
-  return error instanceof Error ? error.message : String(error)
-}
-
 /**
  * For testing purposes, reset the cache after each run
  */
@@ -342,25 +173,39 @@ export function resetUsageStatusCache(): void {
   lastFetched = undefined
 }
 
-function runCatching<T>(fn: () => T): Try<T> {
-  try {
-    return { ok: true, value: fn() }
-  } catch (error) {
-    return { ok: false, error }
+/** TUI adapter: routes the workflow's UI needs to Pi's ctx.ui. */
+export function createTuiUi(ctx: ExtensionContext, loader?: RequestyStatusLoader): DiscoveryUi {
+  return {
+    notify: (message, level) => ctx.ui.notify(`${COMMAND_NAME}: ${message}`, level),
+    confirm: (title, message) => ctx.ui.confirm(title, message),
+    setStatus: message => loader?.setMessage(message),
   }
 }
 
-async function runCatchingAsync<T>(fn: () => Promise<T>): Promise<Try<T>> {
-  try {
-    return { ok: true, value: await fn() }
-  } catch (error) {
-    return { ok: false, error }
+/**
+ * Console adapter for non-interactive modes. Always confirms: print mode is
+ * non-interactive, so the write proceeds without a confirmation dialog.
+ */
+export function createConsoleUi(): DiscoveryUi {
+  return {
+    notify: (message, level) => console.log(`[${level}] ${message}`),
+    // eslint-disable-next-line @typescript-eslint/require-await -- must match the DiscoveryUi promise signature
+    confirm: async () => true,
+    setStatus: message => console.log(message),
   }
 }
 
-function complainOnBrokenEnv(ctx: ExtensionContext, env: Try<Env>) {
-  if (!env.ok) {
-    const notifier = createUiNotifier(ctx)
-    notifier.notify(`failed to load env: ${formatError(env.error)}`, 'error')
-  }
+/** Re-reads models.json through Pi's model registry so new models are usable without /reload. */
+export function createRefreshRegistry(ctx: ExtensionContext): RefreshModelsRegistry {
+  return () =>
+    ctx.modelRegistry.refresh({ allowNetwork: false }).then(result => {
+      if (result.aborted || result.errors.size > 0) {
+        const detail = [...result.errors.entries()].map(([id, error]) => `${id}: ${error.message}`).join(', ')
+        throw new Error(detail || 'refresh aborted')
+      }
+    })
+}
+
+function createGetApiKey(ctx: ExtensionContext): GetApiKey {
+  return providerId => ctx.modelRegistry.getApiKeyForProvider(providerId)
 }