diff --git a/src/api-key.ts b/src/api-key.ts new file mode 100644 index 0000000..0ab52f5 --- /dev/null +++ b/src/api-key.ts @@ -0,0 +1,68 @@ +import { existsSync, readFileSync } from "node:fs" +import { homedir } from "node:os" +import { join } from "node:path" + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function stringValue(value: unknown): string | undefined { + return typeof value === "string" ? value : undefined +} + +function defaultAuthPaths(home: string): string[] { + return [ + join(home, ".commandcode", "auth.json"), + join(home, ".pi", "agent", "auth.json"), + join(home, ".omp", "agent", "auth.json"), + ] +} + +function apiKeyFromCredential(value: unknown): string | undefined { + if (!isRecord(value)) return undefined + + if (stringValue(value.type) === "oauth") return stringValue(value.access) + if (stringValue(value.type) === "api") return stringValue(value.key) + return stringValue(value.access) ?? stringValue(value.key) +} + +export function getConfiguredApiKey( + options: { + env?: NodeJS.ProcessEnv + authPaths?: readonly string[] + homeDir?: () => string + } = {}, +): string | undefined { + const env = options.env ?? process.env + if (env.COMMANDCODE_API_KEY) return env.COMMANDCODE_API_KEY + + const home = options.homeDir?.() ?? homedir() + const authPaths = options.authPaths ?? defaultAuthPaths(home) + + for (const authPath of authPaths) { + try { + if (!existsSync(authPath)) continue + const parsed: unknown = JSON.parse(readFileSync(authPath, "utf-8")) + if (!isRecord(parsed)) continue + + const apiKey = stringValue(parsed.apiKey) + if (apiKey) return apiKey + + const commandcode = stringValue(parsed.commandcode) + if (commandcode) return commandcode + + const providerKey = apiKeyFromCredential(parsed.commandcode) + if (providerKey) return providerKey + + const commandCode = stringValue(parsed["command-code"]) + if (commandCode) return commandCode + + const commandCodeKey = apiKeyFromCredential(parsed["command-code"]) + if (commandCodeKey) return commandCodeKey + } catch { + // Ignore malformed or unreadable auth files. + } + } + + return undefined +} diff --git a/src/oauth.ts b/src/oauth.ts index be77391..0a68ac5 100644 --- a/src/oauth.ts +++ b/src/oauth.ts @@ -1,13 +1,13 @@ /** * Command Code OAuth provider for pi's /login flow. * - * Implements a browser-assisted API key retrieval flow: - * 1. Starts a local HTTP server on a Command Code CLI-compatible port - * 2. Opens the Command Code Studio auth page in the browser - * 3. The user authenticates on the Command Code website - * 4. The website POSTs the API key back to the local server - * 5. If browser transfer fails, the user can paste the API key manually - * 6. The API key is stored in pi's auth.json as OAuth credentials + * Implements two API key retrieval flows: + * 1. Browser-assisted login opens Command Code Studio and waits for the + * website to POST the API key back to a local callback server. + * 2. Direct API key login prompts the user to paste a Studio API key. + * + * If browser transfer fails, the user can still paste the API key manually. + * The API key is stored in pi's auth.json as OAuth credentials. * * Since Command Code API keys don't expire, we store them as * OAuth credentials with a far-future expiry. @@ -101,13 +101,35 @@ async function promptForApiKey(callbacks: OAuthLoginCallbacks, message: string) return credentialsFromApiKey(apiKey) } -/** - * Starts the browser-based login flow for Command Code. - * - * Returns OAuth credentials where access == refresh == the user's API key. - * The keys don't expire, so we set a far-future expiry. - */ -export async function login(callbacks: OAuthLoginCallbacks): Promise { +type LoginChoice = { type: "browser" } | { type: "prompt" } | { type: "apiKey"; apiKey: string } + +async function chooseLoginFlow(callbacks: OAuthLoginCallbacks): Promise { + const input = sanitizeApiKey( + await callbacks.onPrompt({ + message: + "Command Code login: press Enter for browser login, type 'key' to paste an API key, or paste the API key directly:", + }), + ) + const normalized = input.toLowerCase() + + if (!input || normalized === "1" || normalized === "b" || normalized === "browser") { + return { type: "browser" } + } + + if ( + normalized === "2" || + normalized === "k" || + normalized === "key" || + normalized === "api" || + normalized === "paste" + ) { + return { type: "prompt" } + } + + return { type: "apiKey", apiKey: input } +} + +async function browserLogin(callbacks: OAuthLoginCallbacks): Promise { let authServer try { authServer = await startAuthServer() @@ -151,6 +173,23 @@ export async function login(callbacks: OAuthLoginCallbacks): Promise { + const choice = await chooseLoginFlow(callbacks) + + if (choice.type === "apiKey") return credentialsFromApiKey(choice.apiKey) + if (choice.type === "prompt") { + return promptForApiKey(callbacks, "Paste your Command Code API key:") + } + + return browserLogin(callbacks) +} + /** * Command Code API keys don't expire, so "refresh" is a no-op. * Returns the same credentials with an updated far-future expiry. diff --git a/tests/test-api-key.ts b/tests/test-api-key.ts new file mode 100644 index 0000000..a497978 --- /dev/null +++ b/tests/test-api-key.ts @@ -0,0 +1,59 @@ +import assert from "node:assert/strict" +import { mkdtemp, rm, writeFile } from "node:fs/promises" +import { tmpdir } from "node:os" +import { join } from "node:path" +import { describe, it } from "node:test" + +import { getConfiguredApiKey } from "../src/api-key.ts" + +async function withAuthFile( + value: unknown, + run: (authPath: string) => Promise, +): Promise { + const directory = await mkdtemp(join(tmpdir(), "pi-commandcode-auth-")) + const authPath = join(directory, "auth.json") + try { + await writeFile(authPath, JSON.stringify(value), "utf-8") + await run(authPath) + } finally { + await rm(directory, { recursive: true, force: true }) + } +} + +describe("getConfiguredApiKey()", () => { + it("prefers the environment variable", () => { + assert.equal( + getConfiguredApiKey({ env: { COMMANDCODE_API_KEY: "env-key" }, authPaths: [] }), + "env-key", + ) + }) + + it("reads pi OAuth and API credentials", async () => { + const cases: readonly { credential: unknown; expected: string }[] = [ + { + credential: { commandcode: { type: "oauth", access: "oauth-key" } }, + expected: "oauth-key", + }, + { credential: { commandcode: { type: "api", key: "api-key" } }, expected: "api-key" }, + { credential: { "command-code": { type: "api", key: "cli-key" } }, expected: "cli-key" }, + { credential: { apiKey: "legacy-key" }, expected: "legacy-key" }, + ] + + for (const testCase of cases) { + await withAuthFile(testCase.credential, async (authPath) => { + assert.equal(getConfiguredApiKey({ env: {}, authPaths: [authPath] }), testCase.expected) + }) + } + }) + + it("ignores malformed files", async () => { + const directory = await mkdtemp(join(tmpdir(), "pi-commandcode-auth-")) + const authPath = join(directory, "auth.json") + try { + await writeFile(authPath, "not json", "utf-8") + assert.equal(getConfiguredApiKey({ env: {}, authPaths: [authPath] }), undefined) + } finally { + await rm(directory, { recursive: true, force: true }) + } + }) +}) diff --git a/tests/test-oauth.ts b/tests/test-oauth.ts index 0ff0fab..00d94a0 100644 --- a/tests/test-oauth.ts +++ b/tests/test-oauth.ts @@ -186,7 +186,7 @@ describe("login()", () => { authUrl = params.url }, onPrompt(_params: { message: string }): Promise { - throw new Error("onPrompt should not be called in browser flow") + return Promise.resolve("") }, } @@ -239,7 +239,7 @@ describe("login()", () => { process.env.COMMANDCODE_AUTH_TIMEOUT_MS = "1" let authUrl = "" - let promptMessage = "" + const promptMessages: string[] = [] try { const result = await login({ @@ -247,13 +247,13 @@ describe("login()", () => { authUrl = params.url }, async onPrompt(params: { message: string }): Promise { - promptMessage = params.message - return "\u001b[200~ user_manualApiKey\n\u001b[201~" + promptMessages.push(params.message) + return promptMessages.length === 1 ? "" : "\u001b[200~ user_manualApiKey\n\u001b[201~" }, }) assert.match(authUrl, /^https:\/\/commandcode\.ai\/studio\/auth\/cli\?/) - assert.match(promptMessage, /Paste your Command Code API key/) + assert.match(promptMessages[1] ?? "", /Paste your Command Code API key/) assert.equal(result.access, "user_manualApiKey") assert.equal(result.refresh, "user_manualApiKey") assert.ok(result.expires > Date.now(), "expiry should be far in the future") @@ -263,6 +263,37 @@ describe("login()", () => { } }) + it("accepts a directly pasted API key", async () => { + let authOpened = false + const result = await login({ + onAuth() { + authOpened = true + }, + onPrompt(): Promise { + return Promise.resolve("user_directApiKey") + }, + }) + + assert.equal(authOpened, false) + assert.equal(result.access, "user_directApiKey") + }) + + it("offers an explicit API key prompt", async () => { + let promptCount = 0 + const result = await login({ + onAuth() { + throw new Error("browser should not open") + }, + onPrompt(): Promise { + promptCount += 1 + return Promise.resolve(promptCount === 1 ? "key" : "user_promptedApiKey") + }, + }) + + assert.equal(result.access, "user_promptedApiKey") + assert.equal(promptCount, 2) + }) + it("rejects on state token mismatch", async () => { let authUrl = "" const callbacks = { @@ -270,7 +301,7 @@ describe("login()", () => { authUrl = params.url }, onPrompt(_params: { message: string }): Promise { - throw new Error("should not prompt") + return Promise.resolve("") }, }