pi-commandcode-provider/index.ts

/**
 * Command Code provider for pi.
 *
 * Registers the Provider API catalog as a pi provider and layers pi's native
 * auth, model persistence (`models-store.json`), and OpenAI/Anthropic stream
 * adapters on top of it. Chat requests therefore use the same code paths as
 * built-in providers and need no custom transport.
 *
 * Provider API reference: https://api.commandcode.ai/provider/v1
 */

import type { Api, Model, RefreshModelsContext } from "@earendil-works/pi-ai"
import type { ExtensionAPI, ProviderConfig } from "@earendil-works/pi-coding-agent"

import { discoverApiKey } from "./src/api-key.ts"
import { getApiKey, login, refreshToken } from "./src/auth.ts"
import {
  accountApiBase,
  fetchLiveCatalog,
  getModelsTimeoutMs,
  modelsFromCatalog,
  modelsFromLive,
  modelsUrl,
  PROVIDER_ID,
  providerBaseUrl,
  providerHeaders,
  toProviderModel,
  type CommandCodeModel,
} from "./src/models.ts"
import { normalizeCommandCodeMessage } from "./src/overflow.ts"
import { fetchCommandCodeQuota } from "./src/quota.ts"
import { formatQuota } from "./src/quota-format.ts"

type ProviderModelConfig = NonNullable<ProviderConfig["models"]>[number]

/**
 * pi resolves this env template itself; leaving it unresolved means "not
 * configured", so /login credentials and --api-key keep working.
 */
const API_KEY_ENV_REFERENCE = "$COMMAND_CODE_API_KEY"

/** Models are stored in pi's own catalog cache, which expects full models. */
function toStoredModel(model: CommandCodeModel, apiBase: string): Model<"openai-completions" | "anthropic-messages"> {
  const config = toProviderModel(model, apiBase)
  return {
    ...config,
    api: model.api,
    provider: PROVIDER_ID,
    baseUrl: config.baseUrl ?? apiBase,
  } as Model<"openai-completions" | "anthropic-messages">
}

export default function commandCodeProvider(pi: ExtensionAPI): void {
  const apiBase = providerBaseUrl()
  const catalogUrl = modelsUrl()
  const catalogTimeoutMs = getModelsTimeoutMs()
  const headers = providerHeaders()
  const baseline = modelsFromCatalog()

  pi.registerProvider(PROVIDER_ID, {
    name: "Command Code",
    baseUrl: apiBase,
    apiKey: API_KEY_ENV_REFERENCE,
    api: "openai-completions",
    ...(headers ? { headers } : {}),
    models: baseline.map((model) => toProviderModel(model, apiBase)),
    oauth: {
      name: "Command Code",
      isSubscription: true,
      login,
      refreshToken,
      getApiKey,
    },
    refreshModels: async (context: RefreshModelsContext): Promise<ProviderModelConfig[]> => {
      const stored = (context.stored?.models ?? []).filter(
        (model: Model<Api>) => model.provider === PROVIDER_ID,
      )

      if (context.allowNetwork && !context.signal.aborted) {
        try {
          const live = await fetchLiveCatalog({
            url: catalogUrl,
            timeoutMs: catalogTimeoutMs,
            signal: context.signal,
          })
          const models = modelsFromLive(live)
          await context.publish({
            persist: {
              models: models.map((model) => toStoredModel(model, apiBase)),
              checkedAt: Date.now(),
            },
          })
          return models.map((model) => toProviderModel(model, apiBase))
        } catch {
          // An unreachable endpoint keeps the persisted or generated catalog.
        }
      }

      if (stored.length > 0) return stored as ProviderModelConfig[]
      return baseline.map((model) => toProviderModel(model, apiBase))
    },
  })

  // pi only auto-compacts when it recognizes the overflow wording; Command Code
  // reports context limits in its own phrasing.
  pi.on("message_end", (event, ctx) => {
    if (event.message.role !== "assistant") return
    const normalized = normalizeCommandCodeMessage(event.message, ctx.model?.provider)
    return normalized ? { message: normalized.message } : undefined
  })

  pi.registerCommand("commandcode-quota", {
    description: "Show Command Code account usage and quota",
    handler: async (_args, ctx) => {
      await ctx.waitForIdle()
      const apiKey = discoverApiKey()
      if (!apiKey) {
        ctx.ui.notify(
          "No Command Code API key found. Run /login and select Command Code, or set COMMAND_CODE_API_KEY.",
          "warning",
        )
        return
      }

      const result = await fetchCommandCodeQuota({
        apiKey,
        baseUrl: accountApiBase(apiBase),
        ...(headers ? { headers } : {}),
      })
      if (!result.ok) {
        ctx.ui.notify(result.error, "error")
        return
      }
      ctx.ui.notify(formatQuota(result.quota), "info")
    },
  })
}