pi-commandcode-provider/tests/extension/provider.test.ts

import assert from "node:assert/strict"
import { createServer, type Server } from "node:http"
import type { AddressInfo } from "node:net"
import { test } from "node:test"

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

import commandCodeProvider from "../../index.ts"
import { CATALOG } from "../../src/catalog.ts"

/** Minimal ExtensionAPI stand-in that records what the extension registers. */
function createStubPi(): {
  pi: ExtensionAPI
  provider: { id: string; config: ProviderConfig } | undefined
  commands: Map<string, (args: string, ctx: ExtensionCommandContext) => Promise<void>>
  handlers: Map<string, (event: unknown, ctx: unknown) => unknown>
} {
  const state = {
    pi: undefined as unknown as ExtensionAPI,
    provider: undefined as { id: string; config: ProviderConfig } | undefined,
    commands: new Map<string, (args: string, ctx: ExtensionCommandContext) => Promise<void>>(),
    handlers: new Map<string, (event: unknown, ctx: unknown) => unknown>(),
  }

  state.pi = {
    registerProvider: (id: string, config: ProviderConfig) => {
      state.provider = { id, config }
    },
    registerCommand: (name: string, options: { handler: (args: string, ctx: ExtensionCommandContext) => Promise<void> }) => {
      state.commands.set(name, options.handler)
    },
    on: (event: string, handler: (event: unknown, ctx: unknown) => unknown) => {
      state.handlers.set(event, handler)
    },
  } as unknown as ExtensionAPI

  return state
}

function createRefreshContext(options: {
  allowNetwork: boolean
  stored?: { models: unknown[] }
}): { context: RefreshModelsContext; published: unknown[] } {
  const published: unknown[] = []
  const context = {
    allowNetwork: options.allowNetwork,
    signal: new AbortController().signal,
    stored: options.stored,
    publish: async (publication: unknown) => {
      published.push(publication)
      return true
    },
  } as unknown as RefreshModelsContext
  return { context, published }
}

/** Serves the Provider API catalog from a local port for one test. */
async function withCatalogServer(
  payload: unknown,
  run: (url: string) => Promise<void>,
): Promise<void> {
  const server: Server = createServer((_request, response) => {
    response.writeHead(200, { "content-type": "application/json" })
    response.end(JSON.stringify(payload))
  })
  await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve))
  const { port } = server.address() as AddressInfo

  try {
    await run(`http://127.0.0.1:${port}/provider/v1/models`)
  } finally {
    await new Promise<void>((resolve) => server.close(() => resolve()))
  }
}

test("extension registers the Command Code provider with the generated catalog", () => {
  const stub = createStubPi()
  commandCodeProvider(stub.pi)
  const provider = stub.provider

  assert.equal(provider?.id, "commandcode")
  assert.equal(provider?.config.name, "Command Code")
  assert.equal(provider?.config.baseUrl, "https://api.commandcode.ai/provider/v1")
  assert.equal(provider?.config.apiKey, "$COMMAND_CODE_API_KEY")
  assert.equal(provider?.config.models?.length, CATALOG.length)
  assert.equal(provider?.config.oauth?.name, "Command Code")
  assert.ok(provider?.config.oauth?.isSubscription)

  const claude = provider?.config.models?.find((model) => model.id === "claude-sonnet-4-6")
  assert.equal(claude?.api, "anthropic-messages")
  assert.equal(claude?.baseUrl, "https://api.commandcode.ai/provider")

  const deepseek = provider?.config.models?.find((model) => model.id === "deepseek/deepseek-v4.1-flash")
  assert.equal(deepseek?.api, "openai-completions")
  assert.deepEqual(deepseek?.input, ["text", "image"])
})

test("refreshModels merges the live catalog and persists it for offline starts", async () => {
  await withCatalogServer(
    {
      object: "list",
      data: [
        {
          id: "deepseek/deepseek-v4.1-flash",
          object: "model",
          name: "DeepSeek V4.1 Flash",
          context_length: 1_000_000,
        },
        { id: "vendor/fresh", object: "model", name: "Fresh Model", context_length: 64_000 },
      ],
    },
    async (url) => {
      process.env.COMMANDCODE_MODELS_URL = url
      try {
        // The endpoint is read when the extension loads, like pi does at startup.
        const stub = createStubPi()
        commandCodeProvider(stub.pi)
        const provider = stub.provider
        const { context, published } = createRefreshContext({ allowNetwork: true })
        const models = await provider?.config.refreshModels?.(context)

        assert.deepEqual(models?.map((model) => model.id), [
          "deepseek/deepseek-v4.1-flash",
          "vendor/fresh",
        ])
        const persisted = published[0] as { persist?: { models: { id: string; provider: string }[] } }
        assert.deepEqual(persisted.persist?.models.map((model) => model.id), [
          "deepseek/deepseek-v4.1-flash",
          "vendor/fresh",
        ])
        assert.equal(persisted.persist?.models[0]?.provider, "commandcode")
      } finally {
        delete process.env.COMMANDCODE_MODELS_URL
      }
    },
  )
})

test("refreshModels keeps the generated catalog when the endpoint is unreachable", async () => {
  process.env.COMMANDCODE_MODELS_URL = "http://127.0.0.1:1/provider/v1/models"
  try {
    const stub = createStubPi()
    commandCodeProvider(stub.pi)
    const { context, published } = createRefreshContext({ allowNetwork: true })
    const models = await stub.provider?.config.refreshModels?.(context)

    assert.equal(models?.length, CATALOG.length)
    assert.deepEqual(published, [])
  } finally {
    delete process.env.COMMANDCODE_MODELS_URL
  }
})

test("refreshModels restores the persisted catalog without network access", async () => {
  const stub = createStubPi()
  commandCodeProvider(stub.pi)
  const provider = stub.provider

  const stored = [{ id: "stored/model", provider: "commandcode", name: "Stored" }]
  const { context } = createRefreshContext({ allowNetwork: false, stored: { models: stored } })
  const models = await provider?.config.refreshModels?.(context)

  assert.deepEqual(models, stored)
})

test("message_end rewrites Command Code context overflow errors", () => {
  const stub = createStubPi()
  commandCodeProvider(stub.pi)

  const result = stub.handlers.get("message_end")?.(
    {
      message: {
        role: "assistant",
        provider: "commandcode",
        stopReason: "error",
        errorMessage: "prompt is too long: 210000 tokens > 200000 maximum",
      },
    },
    { model: { provider: "commandcode" } },
  ) as { message?: { errorMessage?: string } } | undefined

  assert.match(result?.message?.errorMessage ?? "", /^context_length_exceeded: /)
})

test("/commandcode-quota reports the account snapshot", async () => {
  const stub = createStubPi()
  commandCodeProvider(stub.pi)

  const originalFetch = globalThis.fetch
  const originalKey = process.env.COMMAND_CODE_API_KEY
  process.env.COMMAND_CODE_API_KEY = "user_abc"
  globalThis.fetch = (async (input: Parameters<typeof fetch>[0]) => {
    const url = String(input)
    if (url.includes("/alpha/whoami")) {
      return new Response(JSON.stringify({ user: { userName: "shark-cat" } }), { status: 200 })
    }
    if (url.includes("/alpha/billing/credits")) {
      return new Response(
        JSON.stringify({ credits: { monthlyCredits: 12.5, purchasedCredits: 0, freeCredits: 0 } }),
        { status: 200 },
      )
    }
    if (url.includes("/alpha/billing/subscriptions")) {
      return new Response(JSON.stringify({ data: { planId: "go", status: "active" } }), { status: 200 })
    }
    return new Response(JSON.stringify({ totalCost: 1.23, totalCount: 45 }), { status: 200 })
  }) as typeof fetch

  const notifications: { message: string; type?: string }[] = []
  const ctx = {
    waitForIdle: async () => {},
    ui: { notify: (message: string, type?: string) => notifications.push({ message, type }) },
  } as unknown as ExtensionCommandContext

  try {
    await stub.commands.get("commandcode-quota")?.("", ctx)
  } finally {
    globalThis.fetch = originalFetch
    if (originalKey === undefined) delete process.env.COMMAND_CODE_API_KEY
    else process.env.COMMAND_CODE_API_KEY = originalKey
  }

  assert.equal(notifications.length, 1)
  assert.equal(notifications[0]?.type, "info")
  assert.match(notifications[0]?.message ?? "", /Account: shark-cat/)
  assert.match(notifications[0]?.message ?? "", /Credits remaining: 12\.50/)
})