diff --git a/packages/coding-agent/docs/cli.md b/packages/coding-agent/docs/cli.md index bbb11111a..a044273d1 100644 --- a/packages/coding-agent/docs/cli.md +++ b/packages/coding-agent/docs/cli.md @@ -13,6 +13,7 @@ pi update [target] [options] pi list pi config [options] pi auth [options] +pi auth --continue [payload] pi mcp [options] ``` @@ -303,6 +304,7 @@ Authentication commands require `--provider ` or `--model `. Se | `pi auth check` | Print `ready`, `not_ready`, or `invalid`; exit with status `0`, `1`, or `2`, respectively | | `pi auth print-api-key` | Print the resolved API key | | `pi auth print-bearer-token` | Print a resolved OAuth bearer token | +| `pi auth --continue` | Continue an auth handoff from a base64url JSON payload; prompts when no payload is supplied | | Option | Applies to | Description | |---|---|---| diff --git a/packages/coding-agent/src/cli/auth-command.ts b/packages/coding-agent/src/cli/auth-command.ts index 9ee804602..2ff9ab531 100644 --- a/packages/coding-agent/src/cli/auth-command.ts +++ b/packages/coding-agent/src/cli/auth-command.ts @@ -2,7 +2,7 @@ import type { AuthResult } from "@earendil-works/pi-ai"; import { APP_NAME } from "../config.ts"; import type { Args } from "./args.ts"; -export type AuthCommandKind = "check" | "api_key" | "bearer_token"; +export type AuthCommandKind = "check" | "api_key" | "bearer_token" | "continue"; export interface AuthCommand { kind: AuthCommandKind; @@ -19,10 +19,17 @@ const AUTH_COMMAND_USAGE: Record = { check: `${APP_NAME} auth check --provider [--json] [--credentials] [--no-refresh]`, api_key: `${APP_NAME} auth print-api-key --provider [--model ]`, bearer_token: `${APP_NAME} auth print-bearer-token --provider [--model ] [--min-expiry ]`, + continue: `${APP_NAME} auth --continue [payload]`, }; export function getAuthCommandName(kind: AuthCommandKind): string { - return kind === "check" ? "auth check" : kind === "api_key" ? "auth print-api-key" : "auth print-bearer-token"; + return kind === "check" + ? "auth check" + : kind === "api_key" + ? "auth print-api-key" + : kind === "bearer_token" + ? "auth print-bearer-token" + : "auth --continue"; } export function getAuthCommandUsage(kind: AuthCommandKind): string { @@ -41,8 +48,9 @@ export function printAuthCommandHelp(): void { pi auth print-api-key [--provider ] [--model ] pi auth print-bearer-token [--provider ] [--model ] [--min-expiry ] pi auth check [--provider ] [--model ] [--json] [--credentials] [--no-refresh] + pi auth --continue [payload] -Auth commands require at least one of --provider or --model. Checks refresh expired OAuth credentials by default; --no-refresh prevents this. --credentials emits the credential, or includes it in JSON output.`); +Auth commands require at least one of --provider or --model, except --continue. Checks refresh expired OAuth credentials by default; --no-refresh prevents this. --credentials emits the credential, or includes it in JSON output. --continue accepts a base64url JSON continuation payload; when omitted, Pi prompts for it.`); } export function parseAuthCommand(args: string[]): AuthCommand | undefined { @@ -55,10 +63,12 @@ export function parseAuthCommand(args: string[]): AuthCommand | undefined { ? "api_key" : args[1] === "print-bearer-token" ? "bearer_token" - : undefined; + : args[1] === "--continue" || args[1] === "continue" + ? "continue" + : undefined; if (!kind) { throw new AuthCommandError( - `Unknown auth command "${args[1] ?? ""}". Use "${APP_NAME} auth print-api-key", "${APP_NAME} auth print-bearer-token", or "${APP_NAME} auth check".`, + `Unknown auth command "${args[1] ?? ""}". Use "${APP_NAME} auth print-api-key", "${APP_NAME} auth print-bearer-token", "${APP_NAME} auth check", or "${APP_NAME} auth --continue".`, ); } @@ -69,6 +79,10 @@ export function parseAuthCommand(args: string[]): AuthCommand | undefined { let minExpiryMs: number | undefined; for (let index = 2; index < args.length; index++) { const arg = args[index]; + if (kind === "continue") { + commandArgs.push(arg); + continue; + } if (arg === "--min-expiry") { if (kind !== "bearer_token") throw new AuthCommandError("--min-expiry is only supported by print-bearer-token"); diff --git a/packages/coding-agent/src/cli/auth-continue.ts b/packages/coding-agent/src/cli/auth-continue.ts new file mode 100644 index 000000000..2851872f0 --- /dev/null +++ b/packages/coding-agent/src/cli/auth-continue.ts @@ -0,0 +1,419 @@ +import { Buffer } from "node:buffer"; +import { randomUUID } from "node:crypto"; +import { stdin as input, stdout as output } from "node:process"; +import { createInterface, type Interface } from "node:readline/promises"; +import type { AuthEvent, AuthPrompt, AuthResult } from "@earendil-works/pi-ai"; +import type { McpOAuthState } from "@earendil-works/pi-mcp/oauth"; +import chalk from "chalk"; +import { APP_NAME } from "../config.ts"; +import { ModelRuntime } from "../core/model-runtime.ts"; +import { openBrowser } from "../utils/open-browser.ts"; + +export class AuthContinueError extends Error {} + +type JsonObject = Record; + +interface AuthContinuationBase { + version: 1; + serviceUrl: string; + providerId: string; + continuationId: string; + secret: string; + returnUrl: string; +} + +interface McpContinuationBase { + serverUrl: string; + serverName: string; +} + +export type AuthContinuationPayload = AuthContinuationBase & + ( + | { mcp: McpContinuationBase & { authType: "none" } } + | { mcp: McpContinuationBase & { authType: "bearer" } } + | { + mcp: McpContinuationBase & { + authType: "oauth"; + oauth?: { + scope?: string; + clientName?: string; + clientRegistration?: "dcr" | "cimd"; + authServerMetadataUrl?: string; + }; + }; + } + ); + +type OAuthAuthContinuationPayload = Extract; +type BearerAuthContinuationPayload = Extract; + +function isBearerContinuation(payload: AuthContinuationPayload): payload is BearerAuthContinuationPayload { + return payload.mcp.authType === "bearer"; +} + +function isOAuthContinuation(payload: AuthContinuationPayload): payload is OAuthAuthContinuationPayload { + return payload.mcp.authType === "oauth"; +} + +function asObject(value: unknown, context: string): JsonObject { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new AuthContinueError(`${context} must be a JSON object`); + } + return value as JsonObject; +} + +function requiredString(value: unknown, field: string): string { + if (typeof value !== "string" || value.trim().length === 0) { + throw new AuthContinueError(`${field} must be a non-empty string`); + } + return value.trim(); +} + +function optionalString(value: unknown, field: string): string | undefined { + if (value === undefined) return undefined; + return requiredString(value, field); +} + +function requiredAuthType(value: unknown): AuthContinuationPayload["mcp"]["authType"] { + if (value === "none" || value === "bearer" || value === "oauth") return value; + throw new AuthContinueError('mcp.authType must be "none", "bearer", or "oauth"'); +} + +function optionalRegistration(value: unknown, field: string): "dcr" | "cimd" | undefined { + if (value === undefined) return undefined; + if (value !== "dcr" && value !== "cimd") throw new AuthContinueError(`${field} must be "dcr" or "cimd"`); + return value; +} + +function normalizeContinuationInput(inputValue: string): string { + const trimmed = inputValue.trim(); + if (!trimmed) throw new AuthContinueError("Continuation payload is empty"); + const commandMatch = /^\s*pi\s+auth\s+--continue\s+(.+)$/su.exec(trimmed); + return commandMatch ? commandMatch[1].trim() : trimmed; +} + +function decodeBase64Json(inputValue: string): unknown { + const normalized = normalizeContinuationInput(inputValue); + if (normalized.startsWith("{")) return JSON.parse(normalized); + const base64 = normalized.replace(/-/g, "+").replace(/_/g, "/"); + const padded = base64.padEnd(base64.length + ((4 - (base64.length % 4)) % 4), "="); + return JSON.parse(Buffer.from(padded, "base64").toString("utf8")); +} + +function parseAuthContinuationObject(raw: unknown): AuthContinuationPayload { + const object = asObject(raw, "Continuation payload"); + if (object.version !== 1) throw new AuthContinueError("version must be 1"); + const mcpRaw = asObject(object.mcp, "mcp"); + const base: AuthContinuationBase = { + version: 1, + serviceUrl: requiredString(object.serviceUrl, "serviceUrl"), + providerId: requiredString(object.providerId, "providerId"), + continuationId: requiredString(object.continuationId, "continuationId"), + secret: requiredString(object.secret, "secret"), + returnUrl: requiredString(object.returnUrl, "returnUrl"), + }; + const mcpBase: McpContinuationBase = { + serverUrl: requiredString(mcpRaw.serverUrl, "mcp.serverUrl"), + serverName: requiredString(mcpRaw.serverName, "mcp.serverName"), + }; + const authType = requiredAuthType(mcpRaw.authType); + if (authType === "none") { + if (mcpRaw.oauth !== undefined) throw new AuthContinueError("mcp.oauth is only valid when mcp.authType is oauth"); + return { ...base, mcp: { ...mcpBase, authType } }; + } + if (authType === "bearer") { + if (mcpRaw.oauth !== undefined) throw new AuthContinueError("mcp.oauth is only valid when mcp.authType is oauth"); + return { ...base, mcp: { ...mcpBase, authType } }; + } + const oauthRaw = mcpRaw.oauth === undefined ? undefined : asObject(mcpRaw.oauth, "mcp.oauth"); + return { + ...base, + mcp: { + ...mcpBase, + authType, + ...(oauthRaw + ? { + oauth: { + scope: optionalString(oauthRaw.scope, "mcp.oauth.scope"), + clientName: optionalString(oauthRaw.clientName, "mcp.oauth.clientName"), + clientRegistration: optionalRegistration( + oauthRaw.clientRegistration, + "mcp.oauth.clientRegistration", + ), + authServerMetadataUrl: optionalString( + oauthRaw.authServerMetadataUrl, + "mcp.oauth.authServerMetadataUrl", + ), + }, + } + : {}), + }, + }; +} + +export function decodeAuthContinuationPayload(inputValue: string): AuthContinuationPayload { + try { + return parseAuthContinuationObject(decodeBase64Json(inputValue)); + } catch (error) { + if (error instanceof AuthContinueError) throw error; + throw new AuthContinueError( + `Could not decode continuation payload as base64url JSON: ${error instanceof Error ? error.message : String(error)}`, + ); + } +} + +async function readPayloadFromUser(): Promise { + const rl = createInterface({ input, output }); + try { + return await rl.question("Paste continuation payload: "); + } finally { + rl.close(); + } +} + +function normalizeServiceUrl(value: string): string { + const url = new URL(value); + url.pathname = url.pathname.replace(/\/$/u, ""); + url.search = ""; + url.hash = ""; + return url.href; +} + +function continuationEndpoint( + payload: AuthContinuationPayload, + kind: "claim" | "complete", + serviceUrl: string, +): string { + const url = new URL(`/v1/auth/continuations/${encodeURIComponent(payload.continuationId)}/${kind}`, serviceUrl); + return url.href; +} + +function credentialFromAuth(auth: AuthResult | undefined): string | undefined { + if (auth?.auth.apiKey) return auth.auth.apiKey; + const authorization = Object.entries(auth?.auth.headers ?? {}).find( + ([name]) => name.toLowerCase() === "authorization", + )?.[1]; + return typeof authorization === "string" ? /^Bearer\s+(.+)$/iu.exec(authorization)?.[1] : undefined; +} + +function printAuthEvent(event: AuthEvent): void { + if (event.type === "auth_url") { + console.log(`\nOpen this URL in your browser:\n${event.url}`); + if (event.instructions) console.log(event.instructions); + openBrowser(event.url); + return; + } + if (event.type === "device_code") { + console.log(`\nOpen this URL in your browser:\n${event.verificationUri}`); + console.log(`Enter code: ${event.userCode}`); + return; + } + console.log(event.type === "info" ? event.message : chalk.dim(event.message)); +} + +async function answerAuthPrompt(rl: Interface, prompt: AuthPrompt): Promise { + if (prompt.type === "select") { + console.log(`\n${prompt.message}`); + for (let index = 0; index < prompt.options.length; index++) { + const option = prompt.options[index]; + console.log(` ${index + 1}. ${option.label}${option.description ? ` — ${option.description}` : ""}`); + } + const answer = await rl.question(`Enter number (1-${prompt.options.length}): `); + const selected = prompt.options[Number.parseInt(answer, 10) - 1]; + if (!selected) throw new AuthContinueError("Invalid selection"); + return selected.id; + } + return rl.question(`${prompt.message}${prompt.placeholder ? ` (${prompt.placeholder})` : ""}: `); +} + +async function getContinuationServiceToken(providerId: string, signal: AbortSignal): Promise { + const runtime = await ModelRuntime.create({ allowModelNetwork: false, signal }); + let token = credentialFromAuth(await runtime.getAuth(providerId, { minOAuthValidityMs: 5 * 60_000, signal })); + if (token) return token; + console.log(chalk.dim("Signing in to the continuation service...")); + const rl = createInterface({ input, output }); + try { + await runtime.login( + providerId, + "oauth", + { + signal, + prompt: (prompt) => answerAuthPrompt(rl, prompt), + notify: printAuthEvent, + }, + { getDeviceId: () => randomUUID() }, + ); + } finally { + rl.close(); + } + token = credentialFromAuth(await runtime.getAuth(providerId, { minOAuthValidityMs: 5 * 60_000, signal })); + if (!token) throw new AuthContinueError("Sign-in did not produce a bearer token"); + return token; +} + +async function postContinuation( + endpoint: string, + token: string, + body: JsonObject, + signal: AbortSignal, +): Promise { + const response = await fetch(endpoint, { + method: "POST", + headers: { + accept: "application/json", + "content-type": "application/json", + authorization: `Bearer ${token}`, + }, + body: JSON.stringify(body), + signal, + }); + if (!response.ok) { + const text = await response.text().catch(() => ""); + throw new AuthContinueError( + `Continuation request failed (${response.status})${text ? `: ${text.slice(0, 500)}` : ""}`, + ); + } + const text = await response.text(); + if (!text.trim()) return {}; + return asObject(JSON.parse(text), "Continuation response"); +} + +async function completeMcpContinuation( + payload: AuthContinuationPayload, + token: string, + serviceUrl: string, + body: JsonObject, + signal: AbortSignal, +): Promise { + await postContinuation( + continuationEndpoint(payload, "complete", serviceUrl), + token, + { + continuationId: payload.continuationId, + secret: payload.secret, + mcp: { + serverUrl: payload.mcp.serverUrl, + serverName: payload.mcp.serverName, + authType: payload.mcp.authType, + }, + ...body, + }, + signal, + ); + openBrowser(payload.returnUrl); +} + +async function runMcpBearerContinuation( + payload: BearerAuthContinuationPayload, + token: string, + serviceUrl: string, + signal: AbortSignal, +): Promise { + const rl = createInterface({ input, output }); + let bearerToken: string; + try { + bearerToken = requiredString( + await rl.question(`Paste bearer token for ${payload.mcp.serverName}: `), + "bearer token", + ); + } finally { + rl.close(); + } + await completeMcpContinuation(payload, token, serviceUrl, { bearerToken }, signal); + console.log(chalk.green("MCP bearer token uploaded.")); +} + +async function runMcpOAuthContinuation( + payload: OAuthAuthContinuationPayload, + token: string, + serviceUrl: string, + signal: AbortSignal, +): Promise { + const [{ MemoryOAuthStateStore }, { signInMcpServer, McpSignInCancelledError }] = await Promise.all([ + import("@earendil-works/pi-mcp/oauth"), + import("../extensions/mcp/oauth.ts"), + ]); + const store = new MemoryOAuthStateStore(); + let state: McpOAuthState | undefined; + const captureStore = { + load: () => store.load(), + save: (next: McpOAuthState) => { + state = structuredClone(next); + store.save(next); + }, + }; + const rl = createInterface({ input, output }); + try { + await signInMcpServer({ + serverUrl: payload.mcp.serverUrl, + store: captureStore, + settings: { + scope: payload.mcp.oauth?.scope, + clientName: payload.mcp.oauth?.clientName ?? APP_NAME, + clientRegistration: payload.mcp.oauth?.clientRegistration, + authServerMetadataUrl: payload.mcp.oauth?.authServerMetadataUrl + ? new URL(payload.mcp.oauth.authServerMetadataUrl) + : undefined, + }, + prompt: { + showAuthorizationUrl(url) { + console.log(`\nOpen this URL in your browser:\n${url.href}`); + openBrowser(url.href); + }, + async promptForRedirectUrl(promptSignal) { + try { + return await rl.question( + "If the browser cannot reach the local callback, paste the final redirect URL here (or press Enter to cancel): ", + { signal: promptSignal }, + ); + } catch { + return undefined; + } + }, + }, + signal, + }); + } catch (error) { + if (error instanceof McpSignInCancelledError) throw new AuthContinueError("MCP sign-in cancelled"); + throw error; + } finally { + rl.close(); + } + state ??= store.load(); + if (!state?.tokens) throw new AuthContinueError("MCP sign-in did not produce OAuth tokens"); + await completeMcpContinuation(payload, token, serviceUrl, { authState: state }, signal); + console.log(chalk.green("MCP authentication uploaded.")); +} + +async function runMcpContinuation( + payload: AuthContinuationPayload, + token: string, + serviceUrl: string, + signal: AbortSignal, +): Promise { + if (isBearerContinuation(payload)) { + await runMcpBearerContinuation(payload, token, serviceUrl, signal); + return; + } + if (isOAuthContinuation(payload)) { + await runMcpOAuthContinuation(payload, token, serviceUrl, signal); + return; + } + await completeMcpContinuation(payload, token, serviceUrl, {}, signal); + console.log(chalk.green("MCP connection saved.")); +} + +export async function runAuthContinueCommand(args: string[]): Promise { + const inputPayload = args.length > 0 ? args.join(" ") : await readPayloadFromUser(); + const payload = decodeAuthContinuationPayload(inputPayload); + const serviceUrl = normalizeServiceUrl(payload.serviceUrl); + const signal = AbortSignal.timeout(15 * 60_000); + const token = await getContinuationServiceToken(payload.providerId, signal); + await postContinuation( + continuationEndpoint(payload, "claim", serviceUrl), + token, + { continuationId: payload.continuationId, secret: payload.secret }, + signal, + ); + await runMcpContinuation(payload, token, serviceUrl, signal); +} diff --git a/packages/coding-agent/src/cli/credential-print.ts b/packages/coding-agent/src/cli/credential-print.ts index 304876659..c0d563a6f 100644 --- a/packages/coding-agent/src/cli/credential-print.ts +++ b/packages/coding-agent/src/cli/credential-print.ts @@ -6,7 +6,7 @@ import { AuthCommandError, type AuthCommandKind, getAuthCredential, validateAuth const DEFAULT_BEARER_TOKEN_MIN_EXPIRY_MS = 30 * 60_000; -type CredentialPrintKind = Exclude; +type CredentialPrintKind = Exclude; /** * Resolve one configured provider credential. diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index 0710ed9e6..faa3c78f6 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -26,6 +26,7 @@ import { printAuthCommandHelp, validateAuthCommandArgs, } from "./cli/auth-command.ts"; +import { AuthContinueError, runAuthContinueCommand } from "./cli/auth-continue.ts"; import { resolveCredentialForPrint } from "./cli/credential-print.ts"; import { processFileArguments } from "./cli/file-processor.ts"; import { buildInitialMessage } from "./cli/initial-message.ts"; @@ -147,15 +148,20 @@ async function runAuthCommand(args: string[]): Promise { } if (!command) return false; - const parsed = parseArgs(command.args); - if (parsed.unknownFlags.size > 0) { - const option = parsed.unknownFlags.keys().next().value; - console.error(chalk.red(`Unknown option --${option} for "${getAuthCommandName(command.kind)}".`)); - console.error(chalk.dim(`Use "${APP_NAME} --help" or "${getAuthCommandUsage(command.kind)}".`)); - process.exitCode = 1; - return true; - } try { + if (command.kind === "continue") { + await runAuthContinueCommand(command.args); + return true; + } + + const parsed = parseArgs(command.args); + if (parsed.unknownFlags.size > 0) { + const option = parsed.unknownFlags.keys().next().value; + console.error(chalk.red(`Unknown option --${option} for "${getAuthCommandName(command.kind)}".`)); + console.error(chalk.dim(`Use "${APP_NAME} --help" or "${getAuthCommandUsage(command.kind)}".`)); + process.exitCode = 1; + return true; + } if (parsed.diagnostics.length > 0) { throw new AuthCommandError(parsed.diagnostics.map((diagnostic) => diagnostic.message).join("\n")); } @@ -201,7 +207,10 @@ async function runAuthCommand(args: string[]): Promise { process.stdout.write(`${output}\n`); process.exitCode = result.status === "ready" ? 0 : result.status === "not_ready" ? 1 : 2; } catch (error) { - const message = error instanceof AuthCommandError ? error.message : "Failed to resolve credential"; + const message = + error instanceof AuthCommandError || error instanceof AuthContinueError + ? error.message + : "Failed to resolve credential"; console.error(chalk.red(`Error: ${message}`)); process.exitCode = command.kind === "check" ? 2 : 1; } diff --git a/packages/coding-agent/test/credential-print.test.ts b/packages/coding-agent/test/credential-print.test.ts index a1ca56d9c..21350a5c4 100644 --- a/packages/coding-agent/test/credential-print.test.ts +++ b/packages/coding-agent/test/credential-print.test.ts @@ -2,6 +2,7 @@ import { InMemoryModelsStore } from "@earendil-works/pi-ai"; import { describe, expect, test, vi } from "vitest"; import { parseArgs } from "../src/cli/args.ts"; import { AuthCommandError, isAuthCommandHelp, parseAuthCommand } from "../src/cli/auth-command.ts"; +import { decodeAuthContinuationPayload } from "../src/cli/auth-continue.ts"; import { resolveCredentialForPrint } from "../src/cli/credential-print.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; import { ModelRuntime } from "../src/core/model-runtime.ts"; @@ -112,6 +113,13 @@ describe("credential print commands", () => { noRefresh: false, minExpiryMs: 30 * 60_000, }); + expect(parseAuthCommand(["auth", "--continue", "eyJ2ZXJzaW9uIjoxfQ"])).toEqual({ + kind: "continue", + args: ["eyJ2ZXJzaW9uIjoxfQ"], + json: false, + credentials: false, + noRefresh: false, + }); expect(() => parseAuthCommand(["auth", "print-api-key", "--min-expiry", "30m"])).toThrow( "only supported by print-bearer-token", ); @@ -127,4 +135,24 @@ describe("credential print commands", () => { resolveCredentialForPrint(parseArgs(["--provider", "openai-codex"]), runtime, "api_key"), ).rejects.toThrow("configured with OAuth"); }); + + test("decodes auth continuation payloads", () => { + const payload = { + version: 1, + serviceUrl: "https://auth.example", + providerId: "auth-service", + continuationId: "cont_123", + secret: "one-time-secret", + returnUrl: "https://auth.example/connections?continued=cont_123", + mcp: { + serverUrl: "https://mcp.example/mcp", + serverName: "example", + authType: "oauth", + oauth: { scope: "project:read" }, + }, + }; + const encoded = Buffer.from(JSON.stringify(payload)).toString("base64url"); + expect(decodeAuthContinuationPayload(encoded)).toMatchObject(payload); + expect(decodeAuthContinuationPayload(`pi auth --continue ${encoded}`)).toMatchObject(payload); + }); });