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

import assert from "node:assert/strict"
import { test } from "node:test"

import type { OAuthCredentials, OAuthLoginCallbacks } from "@earendil-works/pi-ai"

import {
  credentialsFromApiKey,
  getApiKey,
  login,
  refreshToken,
  sanitizeApiKey,
  validateApiKey,
} from "../../src/auth.ts"

/** Records the callback surface pi supplies so flows can be driven deterministically. */
function createCallbacks(options: {
  select?: string | undefined
  prompt: string | string[]
}): { callbacks: OAuthLoginCallbacks; authUrls: string[]; prompts: string[] } {
  const authUrls: string[] = []
  const prompts: string[] = []
  const answers = Array.isArray(options.prompt) ? [...options.prompt] : [options.prompt]

  return {
    authUrls,
    prompts,
    callbacks: {
      onAuth: (info) => authUrls.push(info.url),
      onDeviceCode: () => {},
      onSelect: async () => options.select,
      onPrompt: async (prompt) => {
        prompts.push(prompt.message)
        const answer = answers.shift()
        if (answer === undefined) throw new Error("No prompt answer configured")
        return answer
      },
    },
  }
}

/** Replaces global fetch for one test and restores it afterwards. */
async function withFetch(stub: typeof fetch, run: () => Promise<void>): Promise<void> {
  const original = globalThis.fetch
  globalThis.fetch = stub
  try {
    await run()
  } finally {
    globalThis.fetch = original
  }
}

test("sanitizeApiKey removes paste markers, control characters and padding", () => {
  const escape = String.fromCharCode(27)

  assert.equal(sanitizeApiKey(`  user_abc${escape}[200~def[201~\n`), "user_abcdef")
  assert.equal(sanitizeApiKey("user_abc\t\r\n"), "user_abc")
})

test("validateApiKey rejects invalid keys and accepts valid ones", async () => {
  await assert.rejects(
    validateApiKey("user_bad", { fetchImpl: async () => new Response("{}", { status: 401 }) }),
    /rejected the API key/,
  )
  await assert.rejects(
    validateApiKey("user_bad", { fetchImpl: async () => new Response("", { status: 500 }) }),
    /\(500\)/,
  )
  await validateApiKey("user_good", {
    fetchImpl: async (input) => {
      assert.equal(String(input), "https://api.commandcode.ai/alpha/whoami")
      return new Response(JSON.stringify({ success: true }), { status: 200 })
    },
  })
})

test("credentialsFromApiKey keeps the key valid for a decade", () => {
  const credentials = credentialsFromApiKey("user_abc")

  assert.equal(credentials.refresh, "user_abc")
  assert.equal(credentials.access, "user_abc")
  assert.ok(credentials.expires > Date.now() + 9 * 365 * 24 * 60 * 60 * 1000)
  assert.equal(getApiKey(credentials), "user_abc")
})

test("refreshToken returns non-expiring credentials unchanged", async () => {
  const refreshed = await refreshToken(credentialsFromApiKey("user_abc") as OAuthCredentials)

  assert.equal(refreshed.access, "user_abc")
  assert.ok(refreshed.expires > Date.now() + 9 * 365 * 24 * 60 * 60 * 1000)
})

test("login with a selected key option prompts and validates the pasted key", async () => {
  const { callbacks, prompts } = createCallbacks({ select: "key", prompt: "user_abc" })

  await withFetch(async () => new Response("{}", { status: 200 }), async () => {
    const credentials = await login(callbacks)
    assert.equal(credentials.access, "user_abc")
  })
  assert.deepEqual(prompts, ["Paste your Command Code API key:"])
})

test("login accepts a pasted key without opening the selector flow", async () => {
  const { callbacks } = createCallbacks({ prompt: "user_abc" })

  await withFetch(async () => new Response("{}", { status: 200 }), async () => {
    const credentials = await login(callbacks)
    assert.equal(credentials.access, "user_abc")
  })
})

test("login rejects a key the account endpoint refuses", async () => {
  const { callbacks } = createCallbacks({ prompt: "user_nope" })

  await withFetch(async () => new Response("{}", { status: 401 }), async () => {
    await assert.rejects(login(callbacks), /rejected the API key/)
  })
})

test("browser login falls back to a pasted key when the callback never arrives", async () => {
  const { callbacks, authUrls, prompts } = createCallbacks({ select: "browser", prompt: "user_abc" })
  process.env.COMMANDCODE_AUTH_TIMEOUT_MS = "30"

  try {
    await withFetch(async () => new Response("{}", { status: 200 }), async () => {
      const credentials = await login(callbacks)
      assert.equal(credentials.access, "user_abc")
    })
  } finally {
    delete process.env.COMMANDCODE_AUTH_TIMEOUT_MS
  }

  assert.equal(authUrls.length, 1)
  assert.match(authUrls[0] ?? "", /^https:\/\/commandcode\.ai\/studio\/auth\/cli\?callback=/)
  assert.match(prompts.at(-1) ?? "", /Automatic transfer timed out/)
})