/** * `omp auth-gateway` command handlers. * * Boots a forward-proxy server that lets less-trusted clients (the macOS * usage widget, robomp containers, …) make provider API calls without ever * seeing the access token. The gateway is itself a broker client and * resolves credentials through the configured broker (via the same * `OMP_AUTH_BROKER_URL` / `auth.broker.url` precedence used elsewhere). * * Sub-verbs: * - `serve [--bind=…]` — boots the gateway against the configured broker. * - `token` / `token --regenerate` — manages the gateway bearer token file. * - `status` — prints the locally-stored gateway token and bind hint. */ import * as crypto from "node:crypto"; import * as fs from "node:fs/promises"; import * as path from "node:path"; import { type Api, AuthBrokerClient, AuthStorage, DEFAULT_AUTH_GATEWAY_BIND, type GeneratedProvider, getBundledModels, getBundledProviders, type Model, RemoteAuthCredentialStore, type SnapshotResponse, startAuthGateway, } from "@oh-my-pi/pi-ai"; import { getConfigRootDir, isEnoent, VERSION } from "@oh-my-pi/pi-utils"; import chalk from "chalk"; import { type AuthBrokerClientConfig, resolveAuthBrokerConfig } from "../session/auth-broker-config"; export type AuthGatewayAction = "serve" | "token" | "status"; export interface AuthGatewayCommandArgs { action: AuthGatewayAction; flags: { json?: boolean; bind?: string; regenerate?: boolean; /** * Disable bearer-token auth on inbound requests. Useful when the gateway * is bound to loopback (the default `127.0.0.1:4000`) and you don't want * to wire token-paste plumbing into every local client. */ noAuth?: boolean; }; } const ACTIONS: readonly AuthGatewayAction[] = ["serve", "token", "status"]; function getTokenFilePath(): string { return path.join(getConfigRootDir(), "auth-gateway.token"); } async function readToken(): Promise { try { const raw = await Bun.file(getTokenFilePath()).text(); const trimmed = raw.trim(); return trimmed.length > 0 ? trimmed : null; } catch (err) { if (isEnoent(err)) return null; throw err; } } async function writeToken(token: string): Promise { const file = getTokenFilePath(); await fs.mkdir(path.dirname(file), { recursive: true, mode: 0o700 }); await fs.writeFile(file, token, { mode: 0o600 }); try { await fs.chmod(file, 0o600); } catch { // Best-effort (e.g. Windows). } } /** * Atomically create the token file, refusing to clobber an existing one. * Returns `true` on success, `false` when the file already existed (so the * caller re-reads it instead of racing another concurrent `ensureToken`). */ async function createTokenExclusive(token: string): Promise { const file = getTokenFilePath(); await fs.mkdir(path.dirname(file), { recursive: true, mode: 0o700 }); try { // `wx` = O_CREAT | O_EXCL — fails with EEXIST if the file is already there. await fs.writeFile(file, token, { flag: "wx", mode: 0o600 }); } catch (err) { if ((err as NodeJS.ErrnoException).code === "EEXIST") return false; throw err; } try { await fs.chmod(file, 0o600); } catch { // Best-effort (e.g. Windows). } return true; } function generateToken(): string { return crypto.randomBytes(32).toString("base64url"); } async function ensureToken(): Promise { const existing = await readToken(); if (existing) return existing; const token = generateToken(); if (await createTokenExclusive(token)) return token; // Another concurrent invocation won the create race; read what they wrote. const fromRace = await readToken(); if (fromRace) return fromRace; // File existed-then-disappeared between EEXIST and read; last resort, write // our generated token unconditionally so callers don't see an empty string. await writeToken(token); return token; } function createBrokerClient(brokerConfig: AuthBrokerClientConfig): AuthBrokerClient { return new AuthBrokerClient({ url: brokerConfig.url, token: brokerConfig.token }); } async function fetchBrokerSnapshot(client: AuthBrokerClient): Promise { const result = await client.fetchSnapshot(); if (result.status !== 200) throw new Error("Auth broker returned no initial snapshot"); return result.snapshot; } async function runServe(flags: AuthGatewayCommandArgs["flags"]): Promise { const brokerConfig = await resolveAuthBrokerConfig(); if (!brokerConfig) { throw new Error( "`omp auth-gateway serve` requires OMP_AUTH_BROKER_URL (or `auth.broker.url`/`auth.broker.token` in config.yml). The gateway is itself a broker client.", ); } const bind = flags.bind ?? DEFAULT_AUTH_GATEWAY_BIND; const gatewayToken = flags.noAuth ? null : await ensureToken(); // Build a broker-backed AuthStorage — same pattern as discoverAuthStorage() // in sdk.ts. The gateway never touches local SQLite. const client = createBrokerClient(brokerConfig); const initialSnapshot = await fetchBrokerSnapshot(client); const store = new RemoteAuthCredentialStore({ client, initialSnapshot }); // Refresh + usage both flow through the store's broker hooks automatically — // `RemoteAuthCredentialStore.refreshOAuthCredential` and `.fetchUsageReports`. // AuthStorage discovers them when no explicit option overrides them, so the // gateway only needs to construct the store and pass it in. const storage = new AuthStorage(store, { sourceLabel: `broker ${brokerConfig.url}`, }); await storage.reload(); // Build the model resolver + catalog from pi-ai's bundled metadata, scoped // to providers we hold credentials for. Format handlers ask `resolveModel` // to translate a client-requested `model` field into a pi-ai `Model` // before dispatch; `listModels` powers `/v1/models`. const snapshot = storage.exportSnapshot(); const providersWithCreds = new Set(); for (const entry of snapshot.credentials) providersWithCreds.add(entry.provider); const modelById = new Map>(); for (const provider of getBundledProviders()) { if (!providersWithCreds.has(provider)) continue; for (const model of getBundledModels(provider as GeneratedProvider)) { // First-write-wins so a canonical model id collisions across providers // stick to the provider listed first by getBundledProviders. if (!modelById.has(model.id)) modelById.set(model.id, model); } } const handle = startAuthGateway({ storage, bind, bearerTokens: gatewayToken ? [gatewayToken] : [], version: VERSION, resolveModel: (id: string) => modelById.get(id), listModels: () => modelById.values(), }); process.stdout.write(`auth-gateway listening on ${handle.url}\n`); if (gatewayToken) { process.stdout.write(`bearer token: ${getTokenFilePath()} (chmod 0600)\n`); } else { process.stdout.write(`auth: disabled (--no-auth) — any client can call this gateway\n`); } process.stdout.write(`upstream broker: ${brokerConfig.url}\n`); const stopped = Promise.withResolvers(); let shutdownStarted = false; const stop = async (signal: NodeJS.Signals): Promise => { if (shutdownStarted) return; shutdownStarted = true; process.stdout.write(`\nReceived ${signal}, shutting down...\n`); let closeError: unknown; try { await handle.close(); } catch (error) { closeError = error; } finally { storage.close(); } if (closeError) { stopped.reject(closeError); } else { stopped.resolve(); } }; const onSigint = (): void => { void stop("SIGINT"); }; const onSigterm = (): void => { void stop("SIGTERM"); }; process.once("SIGINT", onSigint); process.once("SIGTERM", onSigterm); try { await stopped.promise; } finally { process.off("SIGINT", onSigint); process.off("SIGTERM", onSigterm); } } async function runToken(flags: AuthGatewayCommandArgs["flags"]): Promise { if (flags.regenerate) { const next = generateToken(); await writeToken(next); if (flags.json) { process.stdout.write(`${JSON.stringify({ token: next, path: getTokenFilePath() })}\n`); } else { process.stdout.write(`${next}\n`); } return; } const token = await ensureToken(); if (flags.json) { process.stdout.write(`${JSON.stringify({ token, path: getTokenFilePath() })}\n`); } else { process.stdout.write(`${token}\n`); } } async function runStatus(flags: AuthGatewayCommandArgs["flags"]): Promise { const token = await readToken(); const brokerConfig = await resolveAuthBrokerConfig(); const tokenFile = getTokenFilePath(); if (!brokerConfig) { const status = { ready: false, reason: "not_configured", tokenFile, tokenPresent: token !== null, broker: null, brokerConfigured: false, brokerAuthenticated: false, }; if (flags.json) { process.stdout.write(`${JSON.stringify(status)}\n`); } else { process.stdout.write(`${chalk.yellow("No broker configured.")} Set OMP_AUTH_BROKER_URL.\n`); process.stdout.write( `token: ${status.tokenPresent ? chalk.green("present") : chalk.red("missing")} at ${status.tokenFile}\n`, ); } process.exitCode = 1; return; } try { const snapshot = await fetchBrokerSnapshot(createBrokerClient(brokerConfig)); const tokenPresent = token !== null; const status = { ready: tokenPresent, reason: tokenPresent ? null : "token_missing", tokenFile, tokenPresent, broker: brokerConfig.url, brokerConfigured: true, brokerAuthenticated: true, credentialCount: snapshot.credentials.length, }; if (flags.json) { process.stdout.write(`${JSON.stringify(status)}\n`); } else { const brokerLine = `upstream broker: ${brokerConfig.url} (${snapshot.credentials.length} credential${ snapshot.credentials.length === 1 ? "" : "s" })`; process.stdout.write(`${tokenPresent ? chalk.green("ready") : chalk.yellow("not ready")} ${brokerLine}\n`); process.stdout.write( `token: ${tokenPresent ? chalk.green("present") : chalk.red("missing")} at ${status.tokenFile}\n`, ); if (!tokenPresent) { process.stdout.write( "Run `omp auth-gateway token` or `omp auth-gateway serve` to create a bearer token.\n", ); } } if (!tokenPresent) process.exitCode = 1; } catch (error) { const message = error instanceof Error ? error.message : String(error); const status = { ready: false, reason: "broker_unavailable", tokenFile, tokenPresent: token !== null, broker: brokerConfig.url, brokerConfigured: true, brokerAuthenticated: false, error: message, }; if (flags.json) { process.stdout.write(`${JSON.stringify(status)}\n`); } else { process.stdout.write(`${chalk.red("FAILED")} upstream broker: ${brokerConfig.url}: ${message}\n`); process.stdout.write( `token: ${status.tokenPresent ? chalk.green("present") : chalk.red("missing")} at ${status.tokenFile}\n`, ); } process.exitCode = 1; } } export async function runAuthGatewayCommand(cmd: AuthGatewayCommandArgs): Promise { switch (cmd.action) { case "serve": await runServe(cmd.flags); return; case "token": await runToken(cmd.flags); return; case "status": await runStatus(cmd.flags); return; default: { const _exhaustive: never = cmd.action; throw new Error(`Unknown auth-gateway action: ${String(_exhaustive)}`); } } } export { ACTIONS as AUTH_GATEWAY_ACTIONS };